Metadata-Version: 2.4
Name: so101-nexus
Version: 0.4.12
Summary: End-to-end simulation and training stack for the SO-101 arm: record demonstrations, clone behavior, and reinforce with GPU-parallel MuJoCo / MuJoCo Warp environments, built on LeRobot
License-Expression: Apache-2.0
License-File: LICENSE.md
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Typing :: Typed
Requires-Python: >=3.12
Requires-Dist: gymnasium<2,>=1.0.0
Requires-Dist: huggingface-hub<2
Requires-Dist: mujoco<4,>=3.1.3
Requires-Dist: numpy<3
Requires-Dist: scipy<2
Requires-Dist: trimesh<5
Requires-Dist: tyro<2,>=0.9.0
Provides-Extra: molmoact
Requires-Dist: accelerate; extra == 'molmoact'
Requires-Dist: huggingface-hub; extra == 'molmoact'
Requires-Dist: pillow; extra == 'molmoact'
Requires-Dist: transformers>=5.8; extra == 'molmoact'
Provides-Extra: rocm
Requires-Dist: pytorch-triton-rocm; (sys_platform == 'linux' and platform_machine == 'x86_64') and extra == 'rocm'
Requires-Dist: torch; (sys_platform == 'linux' and platform_machine == 'x86_64') and extra == 'rocm'
Requires-Dist: triton-rocm; (sys_platform == 'linux' and platform_machine == 'x86_64') and extra == 'rocm'
Provides-Extra: teleop
Requires-Dist: gradio>=5.0.0; extra == 'teleop'
Requires-Dist: lerobot[feetech]<0.6,>=0.5.0; extra == 'teleop'
Requires-Dist: opencv-python>=4.8.0; extra == 'teleop'
Requires-Dist: pillow; extra == 'teleop'
Requires-Dist: plotly>=6.0.0; extra == 'teleop'
Provides-Extra: train
Requires-Dist: tensorboard>=2.0.0; extra == 'train'
Requires-Dist: torch; extra == 'train'
Requires-Dist: wandb[media]>=0.16.0; extra == 'train'
Provides-Extra: viz
Requires-Dist: pillow; extra == 'viz'
Provides-Extra: warp
Requires-Dist: mujoco-warp<3.10,>=3.9.0.1; extra == 'warp'
Requires-Dist: torch; extra == 'warp'
Description-Content-Type: text/markdown

<div align="center">

<img src="https://raw.githubusercontent.com/johnsutor/so101-nexus/main/assets/so101.png" width="250" alt="SO-101 Arm">

<h3 align="center">
    <p>SO101-Nexus: SO-101 robot learning, from demos to policies</p>
</h3>

<p align="center">
    <a href="https://github.com/johnsutor/so101-nexus/blob/main/LICENSE.md"><img alt="License" src="https://img.shields.io/github/license/johnsutor/so101-nexus.svg?color=blue"></a>
    <a href="https://www.python.org/downloads/"><img alt="Python" src="https://img.shields.io/badge/python-3.12%2B-blue"></a>
    <a href="https://so101-nexus.com/docs"><img alt="Docs" src="https://img.shields.io/badge/docs-so101--nexus.com-blue"></a>
    <a href="https://github.com/johnsutor/so101-nexus/actions"><img alt="Tests" src="https://img.shields.io/github/actions/workflow/status/johnsutor/so101-nexus/ci.yml?label=tests"></a>
    <a href="https://github.com/johnsutor/so101-nexus/releases"><img alt="GitHub release" src="https://img.shields.io/github/release/johnsutor/so101-nexus.svg"></a>
    <a href="https://colab.research.google.com/github/johnsutor/so101-nexus/blob/main/examples/bc_ppo_warp_colab.ipynb"><img alt="Open In Colab" src="https://colab.research.google.com/assets/colab-badge.svg"></a>
    <a href="https://discord.gg/37kKRXDh8"><img alt="Discord" src="https://img.shields.io/badge/Discord-Join_Us-5865F2?style=flat&logo=discord&logoColor=white"></a>
</p>

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

</div>

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](https://so101-nexus.com/docs).

## Demo Rollouts

<div align="center">
  <video controls muted playsinline width="720" aria-label="MuJoCo PickAndPlace teleoperation rollout">
    <source src="https://raw.githubusercontent.com/johnsutor/so101-nexus/main/docs/public/videos/pick-it-up.mp4" type="video/mp4">
    Open the <a href="https://huggingface.co/spaces/lerobot/visualize_dataset?path=%2Fjohnsutor%2FMuJoCoPickAndPlace-v1%2Fepisode_0">PickAndPlace episode viewer</a> instead.
  </video>
  <p>
    <a href="https://huggingface.co/datasets/johnsutor/MuJoCoPickAndPlace-v1">PickAndPlace dataset</a>
    ·
    <a href="https://huggingface.co/spaces/lerobot/visualize_dataset?path=%2Fjohnsutor%2FMuJoCoPickAndPlace-v1%2Fepisode_0">PickAndPlace viewer</a>
  </p>
</div>

Recorded MuJoCo teleoperation datasets are available on Hugging Face:

<table align="center">
  <thead>
    <tr><th>Task</th><th>Dataset</th><th>Example Rollout</th></tr>
  </thead>
  <tbody>
    <tr>
      <td>PickLift</td>
      <td><a href="https://huggingface.co/datasets/johnsutor/MuJoCoPickLift-v1">johnsutor/MuJoCoPickLift-v1</a></td>
      <td><a href="https://huggingface.co/spaces/lerobot/visualize_dataset?path=%2Fjohnsutor%2FMuJoCoPickLift-v1%2Fepisode_0">Example Rollout</a></td>
    </tr>
    <tr>
      <td>PickAndPlace</td>
      <td><a href="https://huggingface.co/datasets/johnsutor/MuJoCoPickAndPlace-v1">johnsutor/MuJoCoPickAndPlace-v1</a></td>
      <td><a href="https://huggingface.co/spaces/lerobot/visualize_dataset?path=%2Fjohnsutor%2FMuJoCoPickAndPlace-v1%2Fepisode_0">Example Rollout</a></td>
    </tr>
  </tbody>
</table>

## 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

```bash
pip install so101-nexus
```

### From source

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

## Start with the Workflow

### Record demonstrations

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

See the [teleoperation docs](https://so101-nexus.com/docs/teleoperation/overview) for hardware setup, camera fields, environment customization, and Hub upload.

### Run an environment

```python
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](https://so101-nexus.com/docs/environments) 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](https://so101-nexus.com/docs/getting-started/stability).

Install the optional extra and create a batched vector environment:

```bash
pip install "so101-nexus[warp]"
```

```python
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](https://download.pytorch.org/whl/rocm7.2) 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.

```bash
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:

[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/johnsutor/so101-nexus/blob/main/examples/bc_ppo_warp_colab.ipynb)

Or run it locally with [`examples/bc_ppo_warp.py`](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](https://so101-nexus.com/docs/workflow/overview).

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](https://so101-nexus.com/docs/training/ppo) for the command-line workflow and tuning notes, or train a strong policy end to end in your browser:

[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/johnsutor/so101-nexus/blob/main/examples/ppo_warp_colab.ipynb)

## Roadmap

- [x] MuJoCo environments for the SO-101 arm
- [x] SO-101 tasks: Touch, LookAt, Move, PickLift, PickAndPlace, StackCube
- [x] Physical leader-arm teleop recorder for LeRobot datasets
- [x] MuJoCo Warp backend for GPU-parallel throughput
- [x] Stronger training baselines and exemplars for every environment
- [ ] Integration with the [LeRobot Hub](https://huggingface.co/docs/lerobot/en/envhub)

## Development

```bash
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](LICENSE.md).
