feat(env): add mujoco-warp support#1096
Conversation
|
@rainstormstudio
|
Extract shared GPU simulation lifecycle, CUDA graph management, rendering, metrics, and auto-reset into a new `MuJoCoWarpEnv` base class. Move CartPole and CubePick task-specific logic into `tasks/cartpole.py` and `tasks/cubepick.py` respectively. Update `__init__.py` registration and entry points to reference the renamed classes.
Introduce `MuJoCoWarpOffloadEnv` that runs the underlying environment in a subprocess so GPU memory can be freed between rollout phases. Extend the base `MuJoCoWarpEnv` with state serialisation hooks and register the offload path in `get_env_cls()` when `enable_offload` is set in config.
envs Move the condition-check for rendering images into a shared `_should_render_obs` / `_maybe_add_render_to_obs` helper, and add a new `render_for_policy` flag under `video_cfg` so observations can be rendered independently of video saving. Update the cartpole and cubepick tasks to use the helper. Add a full OpenVLA PPO config for cubepick with `render_for_policy: True` on both train and eval.
@rainstormstudio Hello, could you please fix this problem? |
This PR adds MuJoCo-Warp environments to RLinf, providing GPU-accelerated physics simulation using the mujoco-warp.
Description
Two envirionments are currently included:
mujoco_warp_cartpole_env.py): classic cart-pole balancing with GPU-accelerated simulationmujoco_warp_cubepick_env.py): Franka Emika Panda robot arm picking a cubeMotivation and Context
Support GPU-accelerated physics simulation within RLinf with mujoco-warp.
How has this been tested?
Roadmap
mujoco_warp_envbase for general useTypes of changes
Checklist: