so101-nexus 0.4.11

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

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:

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, pick-and-place, and stack-cube.
• 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] (NVIDIA/CUDA only).
• ROCm training (optional): install PyTorch from AMD's ROCm 7.2 wheel index instead of CUDA with so101-nexus[rocm], for behavior cloning and PPO on the MuJoCo backend.

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 on an AMD GPU (ROCm)

so101-nexus[rocm] installs PyTorch from the ROCm 7.2 wheel index instead of the default CUDA build, for behavior-cloning and PPO training on the (CPU-simulated) MuJoCo backend on Linux x86_64. It targets the train extra's torch dependency, not the GPU-parallel Warp backend: Warp is built on NVIDIA Warp, which has no ROCm/AMD support and always requires a CUDA GPU.

uv sync --extra train --extra rocm --no-default-groups

--no-default-groups skips the dev dependency group, which pins lerobot<0.6 (and therefore torch<2.11) for the test suite, a version range incompatible with the ROCm 7.2 torch build. uv rejects combining rocm with teleop, dev, or test for the same reason.

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:

Or run it locally with examples/bc_ppo_warp.py. Defaults target WarpPickLift-v1; the harder pick-and-place task (carry the object onto a goal disc, lower it, and hold still) is the same script with a few opt-in flags (--env-id WarpPickAndPlace-v1 --demo-repo johnsutor/MuJoCoPickAndPlace-v1 --success-bonus 50 --total-timesteps 160000000 --anneal-timesteps 80000000 --lr-min-frac 0.1, all off/at their WarpPickLift-v1-safe default otherwise); validated across 3 seeds at best_success=0.86 mean (std=0.08), see the module docstring for the full sweep. 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:

Roadmap

• MuJoCo environments for the SO-101 arm
• SO-101 tasks: Touch, LookAt, Move, PickLift, PickAndPlace, StackCube
• 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.