Environments

SO101-Nexus ships six Gymnasium-compatible manipulation tasks for the SO-101 arm. Importing so101_nexus.mujoco registers the single-environment MuJoCo*-v1 IDs; importing so101_nexus.warp registers the GPU-batched Warp*-v1 counterparts, which share the same task definitions and step caps.

For a runnable loop, see the quickstart. For the observation components themselves, see Observations; for action spaces, see Control modes.

Tasks

State dim is the width of the default state observation. It changes if you override observations on the config. all_registered_env_ids() returns these six IDs for each backend you have imported.

Success conditions

Every task publishes a binary success flag in the info dict.

PickLift succeeds when the object is grasped and its height above the spawn pose exceeds lift_threshold (default 0.05 m). Both conditions must hold on the same step.

PickAndPlace succeeds when the cube is placed, settled, and released: its XY distance to the target disc is within goal_thresh (default 0.025 m) and it has come back down to table height, its linear and angular speeds are below the static thresholds, and is_grasped < 0.5. Letting go is mandatory. The arm's own motion is not part of the predicate: is_robot_static is reported in info as a diagnostic only, so a policy may still be retreating when the episode succeeds.

StackCube succeeds when cube A's horizontal offset from cube B is within sqrt(2) * cube_half_size + stack_alignment_margin and its height above cube B is within stack_alignment_margin of exactly 2 * cube_half_size (resting flush on cube B's top face), cube A is static, the robot is static, and cube A has been released (is_grasped < 0.5). Unlike PickAndPlace, the robot-static check is required here.

Touch succeeds when the TCP comes within the target object's bounding radius plus touch_margin (default 0.03 m) of the object center, so the threshold scales with object size.

LookAt succeeds when the angle between the wrist camera's optical axis and the direction to the target is at most half the camera's vertical field of view (fov_deg / 2, or the live cam_fovy / 2 when fov_deg is None). This is exactly the GazeState observation.

Move succeeds when the TCP's displacement along the move direction reaches at least target_distance - success_threshold. Perpendicular drift is ignored.

Rewards

Each step returns a scalar roughly in [-1, 1], where 1.0 means the task just finished. It is a weighted sum of reaching, grasping, task_objective, and a one-time completion_bonus, and those four weights must sum to 1.0. The component functions live in so101_nexus.rewards; the weights are set through RewardConfig.

Simple tasks pay the current level. Touch, LookAt, and Move each have exactly one thing to get right, so the reward is how close the robot is right now: low when far away, rising smoothly toward 1.0 as it closes in, plus the bonus on success. Distance terms use tanh shaping rather than a hard cutoff, so reward rises gradually instead of jumping.

Multi-step tasks pay the change in progress. PickLift, PickAndPlace, and StackCube use potential-based shaping: outside the completion bonus, the reward is how much better this step is than the last, not how good the current state looks. Once the robot stops improving, that term drops to exactly zero. This removes the incentive to freeze in a good-looking pose (holding a lifted object for the rest of the episode, or hovering a cube above the target) instead of finishing, because dwelling pays nothing.

Because those terms measure a change, reward can go slightly negative when progress is genuinely lost, such as dropping an object that was already grasped. That is intentional and symmetric with the credit earned for making the same progress in the first place. Success always pays the full 1.0.

Two optional penalties, action_delta_penalty (jerky motion) and energy_penalty (high effort), default to 0.0. See Configs for every reward field and its default.

Task details

Training

PickLift, Touch, LookAt, and Move have solved vanilla-PPO baselines on the Warp backend. PickAndPlace is solved by the demo-seeded BC-then-PPO recipe, not by vanilla PPO. StackCube has no baseline yet. Commands, hyperparameters, and per-seed results are in Training.