Add ManiSkill3 environments (Franka Panda + SIMPLER Bridge)

[qhua360]

Summary

Adds ManiSkill3 (SAPIEN) manipulation tasks as standard SWM environments via a single robot-agnostic wrapper. Two stationary-arm families:

• Franka Panda table-top manipulation (swm/MS*) — PickCube, PushCube, PullCube, PokeCube, StackCube, LiftPegUpright, PegInsertionSide, PlugCharger, PickSingleYCB, RollBall.
• SIMPLER / real2sim Bridge digital twins (WidowX, swm/Simpler*) — carrot-on-plate, spoon-on-towel, stack-cube, eggplant-in-basket.

Includes data collection (collect_maniskill.py) and world-model eval support (LeWM training data config + goal-conditioned MPC eval config). World.evaluate + SWM's policies/solvers also handle policy eval via the task's native success detector (mapped onto terminated).

Design

• ManiSkillWrapper(gym.Wrapper) mirrors FetchWrapper: wraps gym.make(..., num_envs=1), debatches torch→numpy, maps info['success'] → terminated, lifts env_name/proprio/state/instruction into info, renders (H,W,3) uint8. No robot-specific branching.
• TASK_SPECS is the single extension point — one line per task/embodiment.
• Goal-conditioned eval: _set_state/_set_goal_state restore/compare the flat ManiSkill sim state (info['state']), mirroring the PushT env, so SWM's MPC eval (eval_wm.py) works.
• Opt-in maniskill dependency group (GPU-only — needs CUDA+Vulkan; excluded from uv sync --all-extras/[all], installed via uv sync --group maniskill), lazily imported so import stable_worldmodel stays CPU-importable.
• Assets auto-download on first use (MS_SKIP_ASSET_DOWNLOAD_PROMPT=1).

Factors of Variation

The env declares a variation_space of visual distribution-shift knobs (SIMPLER-aligned), applied in _apply_variations using ManiSkill's own SAPIEN idioms (cf. digital_twins/so100_arm/grasp_cube.py) then scene.update_render()+get_obs() so the frame reflects them; sampled values flow into info['variation.*']. Each was verified on an A100 to change the rendered frame via reset(options={'variation_values': ...}):

Factor	Mechanism	Δ pixels
light.intensity	scene.set_ambient_light	17074
camera.angle_delta	render-camera set_local_pose	15300
object.color	cube material set_base_color	102 (small object)
rendering.transparent_arm	robot-link material alpha	5129

Deferred (documented): table.color/background.color — those surfaces are textured, so set_base_color is a verified no-op; they need a texture swap / greenscreen. Distractors and object-pose-as-settable-factor are follow-ups (object pose is already seed-randomized per episode).

Success rate

Demo replay. No trained model is bundled, so scripts/examples/maniskill_demo_replay.py replays ManiSkill's official demonstrations through the wrapper — restoring each demo's initial env_state then replaying its recorded actions, reporting success_rate:

PickCube-v1 (motionplanning demos, pd_joint_pos): success_rate = 20/20 = 100.0%

This confirms the env + success→terminated wiring yields and detects real successes (vs. ~0% from a random policy). (Reproduction uses the initial env_state, not the seed — batched-GPU demos aren't seed-reproducible; the absolute-joint demos replay deterministically.)

World-model MPC. A LeWM world model trained on a collected PickCube dataset and evaluated with goal-conditioned MPC reached 2% (1/50) goal-reaching success on swm/MSPickCube-v0 (random-policy data, 30 epochs, goal threshold 2.0 on the flat sim state) — a real but low baseline:

python scripts/data/collect_maniskill.py
python scripts/train/lewm.py data=maniskill output_model_name=lewm_mspickcube wandb.enabled=false
python scripts/plan/eval_wm.py --config-name maniskill

Creating an eval dataset

World-model MPC eval is wired (env callables + collect/train/eval configs), but — exactly like PushT — it runs on a prepared dataset, not raw World.collect output. The dataset must be in the benchmark layout: per-row episode_idx/step_idx + ep_len/ep_offset, with pixels (HWC), action, proprio, and state (the flat ManiSkill sim state used by _set_state) — the same format as the provided pusht_expert_train.h5. World.collect output isn't natively in that layout (Lance hides the index columns; HDF5 omits them), so the eval dataset is produced/provided separately. The trained checkpoint's config.json must be the model block (save_pretrained(config_key='model')). Real Open-X datasets (DROID / BridgeData V2) are for training policies, not as MPC-eval datasets (they carry no ManiSkill sim state to restore).

Verification (A100, CUDA 12.8, NumPy 2)

• pytest tests/envs/test_maniskill.py → 5 passed (registry + variation-space CPU tests + 2 GPU rollouts). On CPU/CI without ManiSkill the GPU tests skip.
• FoV per-factor pixel-diffs above; render() frames visually confirmed (real Panda/WidowX scenes).
• collect→train→eval verified on an H100 (LeWM goal-reaching 2% / 1-of-50 on swm/MSPickCube-v0).

Notes

• Panda tasks default to native joint control; pass control_mode='pd_ee_delta_pose' for a uniform 7-D EE action.
• Google Robot SIMPLER tasks aren't ported to ManiSkill3 yet (out of scope).

🤖 Generated with Claude Code