技术摘要：UniLab：一种超越 GPU 主导范式的机器人强化学习异构架构

1. 问题陈述

当前基于仿真的机器人强化学习（RL）训练日益围绕GPU 主导范式组织，其中物理仿真、轨迹收集（rollout collection）和策略学习均在单一的 GPU 中心路径上执行（例如 Isaac Gym、Isaac Lab、MuJoCo Playground）。虽然这种方法显著提高了训练速度，但它依赖于一个默认假设，即高效训练必须要求物理仿真驻留在 GPU 上。

作者认为，这一假设混淆了硬件部署与系统效率。在仿真主导的机器人控制中，关键瓶颈不一定在于哪个处理器运行物理仿真，而在于仿真吞吐量、策略学习和运行时同步是否构成了一个高效的全闭环。以 GPU 为中心的模型将高吞吐量实验绑定到一组狭窄的软件环境（主要是 NVIDIA CUDA），并且可能并非所有工作负载的最优设计，特别是那些涉及复杂接触动力学或需要跨平台可移植性的工作负载。

2. 方法论：UniLab 架构

本文提出了UniLab，一种异构训练架构，它将CPU 并行仿真与GPU 策略更新解耦。UniLab 不强制将物理仿真置于 GPU 上，而是将训练系统视为一个闭环组织问题，通过优化数据移动、缓冲和同步来最小化开销。

核心组件

• CPU 端仿真：UniLab 利用在 CPU 上运行的批处理刚体物理后端。目前它在统一的运行时契约下支持两种后端：
  • MuJoCoUni：一个 CPU 批处理的 MuJoCo 运行时。
  • MotrixSim：一个映射到相同任务契约的物理和渲染栈。
    这些后端并行执行高吞吐量环境步进和数据生成（轨迹或状态转移）。
• GPU 端学习：策略和价值网络在 GPU（或其他加速器，如 AMD ROCm、Intel XPU 或 Apple Metal）上进行更新。
• 统一运行时：一个中央运行时协调 CPU 收集器与 GPU 学习器之间的数据路径。它管理：
  • 数据移动：使用固定内存（pinned memory）和双缓冲策略进行异步主机到设备（H2D）传输。
  • 缓冲：共享回放缓冲区（用于离线策略方法）或环形缓冲区（用于在线策略方法）。
  • 同步：参数同步（权重更新）以及在线策略算法的同步点。

执行流程与重叠

UniLab 支持同步和松散耦合的执行模式：

• 在线策略（PPO/APPO）：使用同步的轨迹收集/更新循环，或异步公式（APPO），其中收集器将固定时域的轨迹写入共享环形缓冲区，同时继续步进下一个轨迹，从而允许 CPU 收集与 GPU 学习在挂钟时间上重叠。
• 离线策略（SAC/TD3/FlashSAC）：收集器将状态转移批次插入共享回放缓冲区。UniLab 采用**“传输前采样”（sample-before-transfer）*流水线：CPU 收集器从回放缓冲区采样，打包批次，并在学习器请求该批次之前*启动异步 H2D 传输。这使得 CPU 打包和数据传输与 GPU 当前的学习器更新重叠，从而掩盖了延迟。

可移植性

该架构将仿真器后端与学习器加速器后端解耦。这使得 UniLab 能够在多样化的硬件配置上运行，包括：

• NVIDIA CUDA (Linux)
• Apple macOS (Metal/MLX)
• AMD ROCm
• Intel XPU

3. 主要贡献

1. 系统框架：本文将高效的机器人 RL 训练重新定义为仿真 - 学习闭环的系统组织问题，而不仅仅是 GPU 驻留物理仿真的结果。它挑战了高吞吐量训练必须依赖 GPU 物理仿真的必要性。
2. 异构训练架构：UniLab 通过统一运行时连接 CPU 批处理物理后端、GPU 学习器和数据缓冲。它被实现为一个完整、可扩展的系统，支持在单一框架内运行 PPO、SAC、FlashSAC、TD3 和 APPO。
3. 端到端证据：作者提供了实证证据，表明在特定场景下，异构执行可以优于 GPU 驻留仿真，同时也证明了跨平台的可行性。

4. 实验结果

实验在受控的单 GPU/单 CPU 工作站（NVIDIA RTX 4090, AMD Ryzen 9 9950X3D）上进行，涵盖了各种机器人形态（四足机器人、人形机器人、灵巧手）和任务（移动、运动跟踪、操作）。

• 吞吐量：CPU 批处理仿真（MuJoCoUni 和 MotrixSim）展示了足以支撑 GPU 学习的吞吐量，在涉及复杂接触和灵巧操作的任务中，其吞吐量往往匹配甚至超过 GPU 驻留仿真。
• 训练效率：与相同的硬件上的 GPU 驻留基线（如 MjLab、Isaac Gym、Holosoma）相比，UniLab 实现了3–10 倍的端到端挂钟训练时间改进。
  • 例如，在 G1 翻转跟踪（G1 Flip Tracking）任务中，UniLab (PPO) 完成训练耗时约 111 分钟，而 MjLab 耗时约 120 分钟；而在 G1 平地行走（G1 Walk Flat）任务中，UniLab (FastSAC) 耗时 3.0 分钟，而 Holosoma 耗时 18.3 分钟。
• 增益机制：主要增益源于将学习器与收集器解耦。在 GPU 驻留系统中，收集器和学习器经常争夺相同的加速器资源。UniLab 的 CPU 仿真/打包与 GPU 学习的异步重叠消除了这种竞争。
• 跨平台性能：UniLab 成功在 Apple M 系列芯片、AMD GPU 和 Intel XPU 上训练了模型，证明了高效训练并非 NVIDIA CUDA 生态系统所独有。

5. 意义与主张

本文主张，GPU 仿真是实现高效训练的有效路径，但并非必要路径。

• 拓宽设计空间：UniLab 作为一个系统反例，表明高效机器人 RL 训练的设计空间比当前以 GPU 为中心的默认设置所暗示的更为广阔。
• 降低依赖性：通过支持 CPU 仿真/GPU 学习，UniLab 减少了对 NVIDIA CUDA 软件栈的依赖，使其能够在 Apple macOS、AMD ROCm 和 Intel XPU 上执行。
• 范围与局限性：作者明确指出，其主张涉及仿真主导环境下的端到端训练效率。他们承认，在以下情况下，GPU 驻留仿真可能仍然是首选：
  • 仿真器吞吐量不再是主导瓶颈。
  • 工作负载利用极端规模、富含加速器的配置。
  • 专用 GPU 原生栈的优势超过了异构执行的益处。
  • 任务主要由渲染、感知或基于视觉的学习主导，而非刚体物理。

本文得出结论：高效训练取决于高吞吐量、协调良好的仿真 - 学习执行，而 UniLab 提供了一种实用的系统组织方式，在不强制要求 GPU 驻留物理仿真的情况下实现这一目标。