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)
env.action_space.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()reset(seed=0) controls the environment state. action_space.seed(0) controls
the random actions from action_space.sample(). Use both calls to reproduce this rollout.
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
Replace the environment ID in gym.make() to choose another task.
Environments lists all IDs, episode lengths, observations, rewards, and success conditions.
Use a different object
Pass a config to change what is on the table. Here the default cube becomes a YCB gelatin box:
import gymnasium as gym
import so101_nexus.mujoco
from so101_nexus import PickConfig, YCBObject
config = PickConfig(objects=YCBObject(model_id="009_gelatin_box"))
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 tasks as 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.8, 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 the backends you imported:
from so101_nexus.env_ids import all_registered_env_ids
print(all_registered_env_ids())