AI Research 技术调研

MuJoCo Playground: An Open-Source Framework for GPU-Accelerated Robot Learning and Sim-to-Real Transfer

13分钟阅读

一句话定位

基于 MJX（MuJoCo 的 JAX/GPU 分支）+ Madrona 批渲染器的全开源机器人学习框架，把大规模并行 MuJoCo 强化学习收进单卡：pip install playground 即可在几分钟内训练策略，覆盖四足 / 人形 / 灵巧手 / 机械臂，并在 6 个真实平台上演示了从状态和像素输入的零样本 sim-to-real（RSS 2025 Outstanding Demo Paper Award）。

背景与定位

这是一篇 sim-infra（仿真基础设施） 论文，不是模型工作——它提供的是环境套件 + 训练管线 + 批渲染集成，属于 GPU 并行仿真这一 paradigm。

paradigm 定位：机器人 RL 的主流 sim-to-real 配方是「建数字孪生环境 → 设计 reward → 仿真训练 → 部署上机」，其中 reward 设计需人为迭代（步骤 2、3 反复循环），因此 time-to-robot（时间到机器人） 是关键，而它取决于仿真器的速度。传统上四足 / 灵巧手策略需要多机数天训练；GPU 并行把「agent-环境」循环全放在设备上、利用宽 batch on-policy RL，可缩短到单卡分钟级。

与闭源阵营的对照是本文的核心动机：目前绝大多数 sim-to-real 成果建立在 NVIDIA PhysX（闭源）之上，经由 IsaacGym / Isaac Lab 落地，研究者无法为特定任务扩展仿真器。开源 GPU 物理引擎（MJX、Brax、Warp、Taichi）此前只有很少的机器人环境利用。本文用 MuJoCo 生态补上这块，并用 Madrona 打通端到端视觉 RL。相关工作可对比 cosmos-predict（生成式世界模型做数据/仿真）、以及依赖 IsaacGym/Isaac Lab 的人形/操作栈 asap、groot-n1、homie、gmt。

模型架构

本框架没有单一「模型」，其架构是一个 物理引擎 + 批渲染器 + 训练环境 的集成栈：

  • 物理后端：MuJoCo XLA（MJX），MuJoCo 的 JAX 分支，跑在 GPU 上、可在设备内直接训练。README 更新披露：论文之后已同时支持 MuJoCo Warp（MJWarp） 后端（--impl warp），用以缓解 MJX 的静态形状限制。
  • 批渲染：集成 Madrona——一个 GPU 上的 entity-component-system（ECS），提供两个后端：CUDA 写的软件批光线追踪器（本文实验所用，支持复杂光照、阴影、纹理、材质）与 Vulkan 光栅化器。Madrona 通过底层 JAX primitive 连接其 init/render 函数，可与 jit、vmap 无缝组合。支持逐实例定制，用于视觉属性的域随机化（几何尺寸、颜色、光照、相机位姿）。
  • 环境三大类：① DM Control Suite（[61] 的大多数任务在 MJX 中重实现，作为入门任务，部分 XML 改了 solver iterations/line-search/contact 参数以提速）；② Locomotion；③ Manipulation。资产直接复用 MuJoCo Menagerie。
  • 视觉观测格式：64×64×3，取当前帧 + 前两帧、转灰度后转置喂给 CNN（PandaPickCubeCartesian）。
  • 策略网络：MLP。Locomotion 默认 policy (128,128,128,128)；Go1 joystick/handstand 用 (512,256,128) 的 policy 与 value，value 读 privileged state（非对称 actor-critic）。视觉任务用 CNN。

支持的机体（sim 中可用，★=已演示真机零样本迁移）：

  • 四足：Unitree Go1★（joystick / handstand / footstand / fall-recovery）、Boston Dynamics Spot、Google Barkour。
  • 人形：Berkeley Humanoid★（每腿 6 DoF）、Unitree G1★（29 DoF）、Booster T1★（23 DoF）、Unitree H1、Robotis OP3。
  • 灵巧手：LEAP hand★（7 cm 立方体 SE(3) 重定向）。
  • 机械臂：Franka Emika Panda + Robotiq★（非抓取式 yoga block 重定向；视觉 pick-cube★）、Aloha 双臂 peg-insertion。

数据

这是 RL 框架而非数据集工作，没有离线演示数据——「数据」来自 GPU 并行仿真中在线生成的 transition。真正决定 sim-to-real 的是各任务的 域随机化（domain randomization） 配方，均在附录逐任务给出具体设置：

  • Go1 四足：随机化传感器噪声、动力学属性、任务不确定性；先在受限指令范围的平地训练（5 min），再在粗糙地形放宽范围 finetune。
  • 人形（Berkeley/G1/T1）：沿用四足的域随机化 + finetune 策略。
  • LEAP 手内重定向：随机化机器人参数、立方体质量与摩擦，加传感器噪声；用渐进 curriculum 逐步增大位姿噪声与动作正则。
  • Franka 非抓取：随机化初始位姿、关节位置/速度，注入动作与观测的随机延迟（模拟硬件延迟），curriculum 逐步扩大 block 位移/朝向范围。
  • Franka 视觉 pick-cube：随机化光照、阴影、相机位姿、物体颜色 + 随机亮度后处理 + 高达 250 ms 的随机抓取延迟；训练时把动作维降到 3 维（Y、Z、离散夹爪）提采样效率。

规模上以「环境步」计：DM Control PPO 训练 60M 步（部分 100M）、SAC 5M 步；Locomotion PPO 100M 步（Go1 joystick 200M）；视觉 PPO 1M 步；均在单/多 GPU 上分钟级完成（详见训练方法与 Infra）。

训练方法

  • 算法：主用 on-policy PPO（Brax 实现），另提供 SAC（Brax）；也支持 torch 系的 RSL-RL PPO——论文结果主要用 Brax PPO/SAC。Brax 与 RSL-RL 在 Go1 上达到相近奖励与步态、wallclock 相近（Fig 8）。
  • 关键超参（Brax PPO，Locomotion 默认，Table XV）：num_timesteps 100M、num_envs 8192、batch size 256、unroll length 20、num_minibatches 32、updates_per_batch 4、discount 0.97、lr 3e-4、entropy cost 1e-2、max_grad_norm 1.0，policy (128,128,128,128)。Go1JoystickFlat/Rough：200M 步、policy/value (512,256,128)、value 读 privileged state（Table XVI）。
  • Brax PPO（DM Control，Table XII）：60M 步（Acrobot/Swimmer/WalkerRun 用 100M）、num_envs 2048、batch 1024、unroll 30、discount 0.995、lr 1e-3、entropy 1e-2。
  • Brax PPO（视觉，Table XIII）：madrona_backend=True、1M 步、num_envs 1024、eval_envs 1024、batch 256、unroll 10、discount 0.97、lr 5e-4、entropy 5e-3。
  • Brax SAC（Table XIV）：5M 步（部分任务 10M）、num_envs 128、batch 512、grad_updates_per_step 8、max_replay 1048576×4、min_replay 8192、Q 网络 layer norm、discount 0.99、lr 1e-3。
  • 视觉策略端到端训练：借 Madrona 批渲染直接在仿真中从像素训练，不需要 teacher-student 蒸馏（对比传统「先训 proprioceptive 教师再蒸馏到视觉学生」的两段式）。
  • 各任务训练时长：Go1 平地 <5 min（2×RTX4090）；Berkeley Humanoid <15 min、Unitree G1 / Booster T1 <30 min（2×RTX4090）；LEAP 立方体重定向 30 min（2×RTX4090）；Franka yoga block 10 min（16×A100）；Franka 视觉 pick-cube 10 min（单 RTX4090 + i9-14900KF）；DM Control 多数任务单卡 <10 min。

Infra（训练 / 推理工程）

训练侧：

  • 硬件覆盖 NVIDIA RTX 4090、A100、H100 及多卡拓扑。以 LeapCubeReorient（接触密集）为例，固定超参（8192 num envs）下：1×4090 ≈ 2080 s，8×H100 ≈ 670 s（约 3× 加速）；更大拓扑能把训练时间缩到 1/3。
  • RL 吞吐（单 A100，Brax PPO，Locomotion，Table VII）：多数环境达 2–4×10⁵ env steps/s，最高如 Go1JoystickFlatTerrain 417,451、SpotJoystickGaitTracking 407,572 steps/s（DM Control Suite 的 FingerSpin/WalkerRun 在 Table IV，仅 246,791/141,581 steps/s）。
  • 像素吞吐（MJX+Madrona）：CartpoleBalance ≈ 403,000 steps/s，PandaPickCubeCartesian ≈ 37,000 steps/s（Franka 物理比 Cartpole 贵 >20×，故对分辨率不敏感）。
  • 训练瓶颈迁移：在 PPO 训练回路中，物理 + 渲染 + 推理合计只占 Cartpole 9% / Franka 43% 的总训练时间，其余耗在 CNN 策略更新上——瓶颈从「采数据」转到「处理数据」。Table XI：Cartpole 中 >90% 负担在网络更新（policy update fraction 0.91）；Franka 为 0.57。Table X 原始 FPS：Cartpole 纯物理 1.37×10⁶ → 加像素 4.03×10⁵ → 加推理 3.41×10⁵ → 加训练 3.13×10⁴；Panda 对应 6.40×10⁴ → 3.69×10⁴ → 3.60×10⁴ → 1.56×10⁴（RTX4090，5 次平均）。
  • 并行/精度：全流程 JAX（jit/vmap），单设备并行数千环境；JIT 编译 1–3 min。README 提示 Ampere（RTX 30/40）默认 TF32 matmul 会损 RL 稳定性/复现，需 export JAX_DEFAULT_MATMUL_PRECISION=highest 才与 float32 一致。Madrona 与 Maniskill3、IsaacLab 的批渲染吞吐做过粗对比（Fig 23，数据源自 [60]），结论是「有竞争力」。
  • 计算由 Lambda Labs 云 credit 与 The AI Institute 赞助。

推理 / 部署侧（边缘）：

  • Locomotion：ros2-control + C++ 实时；策略经 ONNX Runtime 以 50 Hz 推理，配 model-based estimator 在硬件通信频率 500–2000 Hz 运行提供线速度观测；Unitree SDK1/2、Berkeley EtherCAT master 抽象为统一 sensor/actuator 接口（同接口也接 Gazebo 做 sim-to-sim 验证）。
  • LEAP 手：策略 20 Hz，单 Intel RealSense D415 从上方估计立方体位姿（预训练检测器）。
  • Franka 非抓取：200 Hz 直接力矩控制（7 关节输出电机力矩，夹爪闭合），学控力矩而非位置 → 平滑柔顺、迁移更好。
  • Franka 视觉 pick-cube：15 Hz，RTX 3090 做推理，64×64 RGB 方裁降采样后经 C++ ONNX ROS 节点出笛卡尔增量 + 夹爪指令，D435 相机。
  • 团队在 不到 8 周内 把状态与视觉策略部署到 6 个真实平台。

评测 benchmark

sim-to-real 真机结果（一手数据）：

  • LEAP 手内立方体重定向（10 次试验，记连续成功旋转数）：median 3.5、mean 7.1（Table I）。主要失败模式：立方体卡在指-掌间隙致策略停滞；偶发食指-拇指互锁（低成本硬件形变）。
  • Franka 非抓取 yoga block 重定向（35 次真机试验，Table II）：真机成功率 85.7 ± 12.2%，位置误差 5.28 ± 3.26 cm、朝向误差 3.32 ± 1.59°（median 分别 100% / 1.95 cm / 1.72°）；成功判据为位置 3 cm 内、朝向 10° 内。训练于 16×A100。
  • Franka 视觉 pick-cube（从像素）：12 次真机试验 100% 成功，对光照与轻微相机抖动鲁棒。
  • 四足 Go1：joystick / handstand / footstand / fall-recovery 四个策略均鲁棒零样本迁移到室内外（草地、水泥、不同摩擦地面），能扛外部扰动。
  • 人形：Berkeley Humanoid 在从硬地到软滑地面稳健跟踪速度指令；G1 与 Booster T1 零样本稳定行走与转向。
  • 仿真基线对照：Brax PPO vs RSL-RL PPO 奖励曲线相近（Fig 8）；Madrona 批渲染吞吐对 IsaacLab / Maniskill3 具竞争力（Fig 23，作者明说是「粗略对比」，因仿真器差异难做完全受控 benchmark）。

创新点与影响

论文自述三大贡献：① 用 MJX 建成一套覆盖四足 / 人形 / 灵巧手 / 机械臂的完整机器人环境套件，并演示跨平台 sim-to-real；② 集成开源 Madrona 批渲染器，实现单卡端到端视觉策略训练、并在操作任务上零样本迁移；③ 提供含 notebook、超参、训练曲线的完整可复现管线。

影响与定位：这是首个「完全开源（Apache-2.0）+ 覆盖广机体 + 真机零样本验证」的 MJX 机器人学习套件，是对闭源 PhysX/IsaacGym/Isaac Lab 路线的开源替代；把「大规模并行 MuJoCo RL」的门槛压到 pip install + 单张消费级 GPU + 一个 Colab notebook。视觉 RL 无需蒸馏直接从像素端到端训练是相对以往两段式管线的方法学改变。获 RSS 2025 Outstanding Demo Paper Award。

作者自陈局限：① 继承 MJX/JAX 约束——JIT 编译慢（1–3 min）；接触计算量按「场景中可能接触对数」而非「实际激活接触数」增长（JAX 静态形状所致），可由 Warp/Taichi 更灵活框架缓解（正在开发，README 已上 MJWarp 后端）；② 基于 Madrona 的视觉训练仍处早期阶段（目前仅在 CartpoleBalance 完整演示端到端像素 RL）。

原始链接

一手源存档（sources/）

关系图谱

反向链接