API Reference
Core Package
Environment registration plus the full export surface of so101_nexus.
so101_nexus provides shared types, configuration classes, scene objects, observation components, asset helpers, and the environment registry. Backends live under so101_nexus.mujoco and so101_nexus.warp and are registered by importing them.
import so101_nexusEnvironment registration
Environment IDs are registered with Gymnasium at import time. Importing the top-level so101_nexus package alone registers nothing.
import so101_nexus.mujoco # registers the MuJoCo ids
import so101_nexus.warp # registers the Warp ids (so101-nexus[warp])| Package | Import | Registers |
|---|---|---|
so101-nexus | import so101_nexus.mujoco | Single-env MuJoCo*-v1 ids |
so101-nexus[warp] | import so101_nexus.warp | Batched GPU-parallel Warp*-v1 ids |
MuJoCo environments are constructed with gym.make():
import gymnasium as gym
import so101_nexus.mujoco
env = gym.make("MuJoCoPickLift-v1", render_mode="human")Warp environments are natively batched vector environments, so they are constructed with gym.make_vec():
import gymnasium as gym
import so101_nexus.warp
envs = gym.make_vec("WarpTouch-v1", num_envs=4096, device="cuda")
obs, info = envs.reset(seed=0)Warp observations, actions, and rewards are torch tensors on device, and autoreset is same-step: the done step already returns the post-reset observation.
Both gym.make and gym.make_vec accept config= and an episode-length keyword:
from so101_nexus import CubeObject, PickConfig
config = PickConfig(objects=[CubeObject(color="blue")])
env = gym.make("MuJoCoPickLift-v1", config=config, max_episode_steps=512)See Configuration Classes for the config surface and Environments for the id list, step caps, and observation widths.
all_registered_env_ids()
from so101_nexus.env_ids import all_registered_env_ids
ids = all_registered_env_ids()Returns the SO101-Nexus environment IDs for whichever backends are currently imported: six MuJoCo*-v1 ids with the MuJoCo backend imported, plus six Warp*-v1 ids with the Warp backend imported.
Exports by Category
Type Aliases
| Name | Definition |
|---|---|
ColorName | Literal["red", "orange", "yellow", "green", "blue", "purple", "black", "white", "gray"] |
ColorConfig | ColorName | list[ColorName] |
ControlMode | Literal["pd_joint_pos", "pd_joint_delta_pos", "pd_joint_target_delta_pos", "pd_ee_pose", "pd_ee_delta_pose"] (see Control Modes) |
ObsMode | Literal["state", "visual"] |
MoveDirection | Literal["up", "down", "left", "right", "forward", "backward"] |
YcbModelId | Literal union of 10 YCB object IDs (see Scene Objects and Assets) |
Constants
| Name | Type | Description |
|---|---|---|
SO101_JOINT_NAMES | tuple[str, ...] | Ordered joint names for the SO-101 arm |
DIRECTION_VECTORS | dict[MoveDirection, tuple] | Maps direction names to unit vectors |
COLOR_MAP | dict[str, list[float]] | Maps color names to RGBA values |
YCB_OBJECTS | dict[str, str] | Maps YCB model IDs to human-readable names |
ROBOT_CAMERA_PRESETS | dict[str, RobotCameraPreset] | Named camera mounting presets for SO-100 and SO-101 |
ASSETS_DIR | Path | Root directory for bundled assets |
SO101_DIR | Path | Asset directory for the SO-101 arm (URDF/XML retained for teleop calibration; the MuJoCo backend loads the menagerie MJCF) |
Configuration Classes
All configuration classes have sensible defaults. See Configuration Classes for full reference.
| Class | Purpose |
|---|---|
RenderConfig | Render camera resolution and view selection (visualization only, not observations) |
RobotConfig | Rest pose, grasp force threshold, velocity threshold, end-effector solver knobs |
RobotCameraPreset | Named preset combining camera mounting parameters |
RewardConfig | Reward component weights with a compute() method |
EnvironmentConfig | Base environment parameters (spawn region, obs mode, episode length) |
PickConfig | Extends EnvironmentConfig for pick/lift tasks |
PickAndPlaceConfig | Extends EnvironmentConfig for pick-and-place tasks |
StackCubeConfig | Extends EnvironmentConfig for stack-cube tasks |
TouchConfig | Extends PickConfig for touch-an-object tasks |
LookAtConfig | Extends EnvironmentConfig for look-at tasks |
MoveConfig | Extends EnvironmentConfig for directional move tasks |
Observation Components
See Observations for conceptual overview and Configuration Classes for full parameter reference.
| Class | Type | Description |
|---|---|---|
Observation | Abstract base | Base class for all observation components |
JointPositions | State (6-dim) | Current joint angles |
JointVelocities | State (6-dim) | Current joint angular velocities (rad/s) |
JointEfforts | State (6-dim) | Actuator force on each joint (N*m) |
GripperContactForce | State (3-dim) | World-frame resultant contact force on the fingers (N) |
EndEffectorPose | State (7-dim) | TCP position + quaternion orientation |
TargetOffset | State (3-dim) | Vector to the goal: goal minus object in manipulation tasks, goal minus TCP elsewhere |
GazeDirection | State (3-dim) | Unit vector from the wrist camera toward the target object |
GazeState | State (1-dim) | Binary in-frame flag (target object inside the wrist camera FOV) |
GraspState | State (1-dim) | Binary grasp flag (two-sided, opposing-normal contact) |
ObjectPose | State (7-dim) | Object position + quaternion orientation |
ObjectVelocity | State (6-dim) | Object linear + angular velocity |
ObjectOffset | State (3-dim) | Vector from gripper to object |
TargetPosition | State (3-dim) | Absolute goal position |
WristCamera | Camera | RGB image from wrist-mounted camera |
OverheadCamera | Camera | RGB image from stationary overhead camera |
CameraObservation | Camera (abstract) | Base class for camera observation components (WristCamera, OverheadCamera) |
Scene Objects
See Scene Objects and Assets for full reference.
| Class | Description |
|---|---|
SceneObject | Abstract base for all scene objects |
CubeObject | A colored cube with configurable size and mass |
YCBObject | A YCB benchmark object loaded from mesh assets |
MeshObject | A custom mesh object loaded from collision and visual mesh files |
Functions
Color
| Function | Signature | Description |
|---|---|---|
sample_color | (colors: ColorConfig, rng: Generator | None = None) -> list[float] | Resolve a ColorConfig to an RGBA list, sampling uniformly if given a list of colors |
Reward and observation helpers
| Function | Signature | Description |
|---|---|---|
privileged_state_feature_names | (observations: Sequence[Observation] | None) -> list[str] | Per-dimension names for the non-camera state observation vector |
observations_from_feature_names | (names: Sequence[str]) -> list[Observation] | Inverse of the above: rebuild the state component layout a recording declares |
relabel_environment_state | (env_state, recorded_names, observations, *, dt, episode_index=None) -> ndarray | Re-lay a recorded observation.environment_state matrix onto a newer component layout, reconstructing JointVelocities by finite difference |
reach_progress | (distance, *, scale) -> float | ndarray | tensor | Tanh-shaped progress in [0, 1]: 1 at distance 0, decaying toward 0 |
orientation_progress | (cos_similarity) -> float | ndarray | tensor | Maps cosine similarity in [-1, 1] to a reward in [0, 1] |
lift_progress | (height, *, scale, grasped) -> float | ndarray | tensor | Tanh-shaped lift progress in [0, 1], zero unless grasped |
simple_reward | (*, progress, completion_bonus, success) -> float | ndarray | tensor | Single-objective task reward; shaping fills [0, 1 - completion_bonus] via progress |
direction_to_object | (camera_pos, object_pos) -> ndarray | tensor | Unit vector from a camera toward an object |
gaze_cosine | (gaze_axis, direction) -> float | ndarray | tensor | Clamped cosine between a camera's optical axis and a direction |
gaze_angle_rad | (cosine) -> float | ndarray | tensor | Gaze angle in radians for a gaze_cosine result |
object_in_view | (angle_rad, half_fov_rad) -> bool | ndarray | tensor | Whether a gaze bearing falls inside a camera's half field of view |
Asset Paths
See Scene Objects and Assets for full reference.
| Function | Returns | Description |
|---|---|---|
get_so101_simulation_dir() | Path | Path to SO-101 SO101/ assets (URDF/XML, retained for teleop calibration) |
get_so101_mujoco_model_dir() | Path | Directory of the vendored menagerie MJCF (loaded by the MuJoCo backend) |
get_so101_mujoco_model_path() | Path | Path to the menagerie so101.xml (loaded by the MuJoCo backend) |
YCB Asset Management
See Scene Objects and Assets for full reference.
| Function | Description |
|---|---|
ensure_ycb_assets(model_id) | Download and cache YCB assets, return cache path |
get_ycb_mesh_dir(model_id) | Path to mesh directory for a YCB model |
get_ycb_collision_parts(model_id) | Convex collision parts with their mass fractions |
get_ycb_collision_meshes(model_id) | Paths of the convex collision parts |
get_ycb_collision_mesh(model_id) | Path of the first convex collision part |
get_ycb_visual_mesh(model_id) | Path to visual.obj |
get_ycb_texture_file(model_id) | Expected path to the optional texture.png |
get_mujoco_ycb_rest_pose(verts) | Rest orientation and spawn height for MuJoCo |
Environment Registry
| Function | Description |
|---|---|
all_registered_env_ids() | All registered Gymnasium environment IDs for the imported backends |
Processor APIs
The optional so101_nexus.processors subpackage provides LeRobot-style action and observation processor steps and a Gym wrapper that emits LeRobot-shaped observations. It requires the teleop extra. See LeRobot Processors for the full reference and Installation for the extras matrix.