Skip to content

johnsutor/so101-nexus

Repository files navigation

SO-101 Arm

SO101-Nexus: SO-101 robot learning, from demos to policies

License Python Docs Tests GitHub release Open In Colab Discord

Beta: APIs may change between releases. Feedback and bug reports are welcome.

SO101-Nexus is an end-to-end Python library for taking an SO-101 robot from demonstrations to a trained policy. It combines physical leader-arm teleoperation, LeRobot-compatible dataset recording, Gymnasium/MuJoCo manipulation environments, and training/evaluation hooks in one installable package.

For full documentation, visit so101-nexus.com/docs.

Demo Rollouts

Recorded MuJoCo teleoperation datasets are available on Hugging Face:

TaskDatasetExample Rollout
PickLift johnsutor/MuJoCoPickLift-v1 Example Rollout
PickAndPlace johnsutor/MuJoCoPickAndPlace-v1 Example Rollout

Why

There are useful SO-101 tools, but few packages connect teleoperation, LeRobot datasets, environments, and training loops in one workflow, all using simulations. SO101-Nexus is built around the record -> clone -> reinforce path: collect demonstrations, replay and evaluate in matching SO-101 environments, bootstrap with imitation learning, then fine-tune with RL.

MuJoCo is the default backend. An optional MuJoCo Warp backend (so101-nexus[warp]) adds GPU-parallel, batched environments for large-scale RL.

What You Get

  • Teleoperation recorder: drive a simulated follower with a physical SO-100 or SO-101 leader arm.
  • LeRobot dataset output: save demonstrations with SO follower state/action units and wrist/overhead camera fields.
  • Gymnasium environments: run SO-101 MuJoCo tasks for touch, look-at, move, pick-lift, and pick-and-place.
  • Configurable curricula: swap objects, add distractors, randomize colors, tune rewards, and choose observation components.
  • Training and evaluation hooks: start with the PPO baseline, LeRobot processors, and policy adapters for real-policy evaluation.
  • GPU-parallel Warp backend (optional, experimental): batched Warp*-v1 vector environments for large-scale RL, installed with so101-nexus[warp].

Installation

pip install so101-nexus

From source

git clone https://github.com/johnsutor/so101-nexus.git
cd so101-nexus
uv sync

Start with the Workflow

Record demonstrations

uvx --from "so101-nexus[teleop]" so101-nexus teleop \
    --leader-port /dev/ttyACM0

See the teleoperation docs for hardware setup, camera fields, environment customization, and Hub upload.

Run an environment

import gymnasium as gym
import so101_nexus.mujoco  # noqa: F401

env = gym.make("MuJoCoPickLift-v1", render_mode="rgb_array")
obs, info = env.reset()

for _ in range(256):
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    if terminated or truncated:
        obs, info = env.reset()

env.close()

See the environment reference for all task IDs.

Run the GPU-parallel Warp backend

Experimental: The Warp backend's API and physics may change between minor releases while the MuJoCo backend is stable. See Stability and versioning.

Install the optional extra and create a batched vector environment:

pip install "so101-nexus[warp]"
import gymnasium as gym
import so101_nexus.warp  # noqa: F401

envs = gym.make_vec("WarpTouch-v1", num_envs=4096, device="cuda")
obs, info = envs.reset(seed=0)
obs, reward, terminated, truncated, info = envs.step(envs.action_space.sample())
envs.close()

Train a policy

The default workflow is demo-seeded: behavior cloning from teleoperation demonstrations, then PPO fine-tuning on the GPU-parallel Warp backend. Train end to end in your browser:

Open In Colab

Or run it locally with examples/bc_ppo_warp.py. For the full record -> clone -> reinforce walkthrough, see the Workflow docs.

Prefer a from-scratch baseline instead? SO101-Nexus also ships a CleanRL-style PPO baseline for Gymnasium environments (no demonstration seeding). See Training with PPO for the command-line workflow and tuning notes, or train a strong policy end to end in your browser:

Open In Colab

Roadmap

  • MuJoCo environments for the SO-101 arm
  • SO-101 tasks: Touch, LookAt, Move, PickLift, PickAndPlace
  • Physical leader-arm teleop recorder for LeRobot datasets
  • MuJoCo Warp backend for GPU-parallel throughput
  • Stronger training baselines and exemplars for every environment
  • Integration with the LeRobot Hub

Development

git clone https://github.com/johnsutor/so101-nexus.git
cd so101-nexus
uv sync

make test       # run all tests
make format     # format code
make lint       # lint code

License

This repository's source code is available under the Apache-2.0 License.

About

Full-stack robot learning for the SO-101 arm: teleoperation, imitation learning, and RL in MuJoCo.

Topics

Resources

License

Contributing

Security policy

Stars

20 stars

Watchers

0 watching

Forks

Packages

 
 
 

Contributors