Quickstart
Run a MuJoCo task with random actions, swap the task or the object, then scale to thousands of GPU-parallel worlds.
Run your first environment
import gymnasium as gym
import so101_nexus.mujoco # registers the MuJoCo env ids
env = gym.make("MuJoCoPickLift-v1", render_mode="human")
obs, info = env.reset(seed=0)
for _ in range(1000):
action = env.action_space.sample()
obs, reward, terminated, truncated, info = env.step(action)
if terminated or truncated:
obs, info = env.reset()
env.close()Importing so101_nexus.mujoco is what registers the ids with Gymnasium, so it must run before gym.make(). Use render_mode="human" for a live viewer window, or render_mode="rgb_array" to capture frames programmatically.
Try another task
Every task uses the same Gymnasium API, so swapping the id is the only change.
| Environment ID | Task |
|---|---|
MuJoCoPickLift-v1 | Pick up an object and lift it |
MuJoCoPickAndPlace-v1 | Carry an object onto a goal and release it |
MuJoCoStackCube-v1 | Stack cube A on top of cube B |
MuJoCoTouch-v1 | Touch an object on the table |
MuJoCoLookAt-v1 | Aim the wrist camera at a target |
MuJoCoMove-v1 | Move the TCP in a cardinal direction |
Environments documents each task's episode length, observation layout, reward, and success condition.
Use a different object
Pass a config to change what is on the table. Here the default cube becomes a YCB banana:
import gymnasium as gym
import so101_nexus.mujoco
from so101_nexus import PickConfig, YCBObject
config = PickConfig(objects=YCBObject(model_id="011_banana"))
env = gym.make("MuJoCoPickLift-v1", config=config, render_mode="human")
obs, info = env.reset(seed=0)
env.close()YCB assets download automatically on first use and are cached locally.
Scale up on GPU
The optional Warp backend registers the same six tasks as Warp*-v1 batched vector environments:
import gymnasium as gym
import so101_nexus.warp # requires so101-nexus[warp]
envs = gym.make_vec("WarpPickLift-v1", num_envs=4096, device="cuda")
obs, info = envs.reset(seed=0)
envs.close()The Warp backend is experimental. It needs an NVIDIA GPU with CUDA >= 12.4, and its physics differs from MuJoCo, so a policy may need re-tuning across backends. See Backends and Stability and versioning.
all_registered_env_ids() reports the ids for whichever backends you imported, six per backend:
from so101_nexus.env_ids import all_registered_env_ids
print(all_registered_env_ids())