CUDA加速火箭联盟强化学习：设备驻留仿真让训练提速百倍

在强化学习（RL）领域，环境仿真的速度往往决定了训练的效率上限。当大多数人还在用CPU跑游戏模拟器时，一位开发者把《火箭联盟》（Rocket League）的物理引擎搬到了GPU上，让数千局游戏并行运行，训练一个超越Nexto水平的智能体只需要几个小时——而原本，这需要数周。

从PureJaxRL到CARL：一次跨越两年的工程冒险

这个项目的灵感来自PureJaxRL——一个用JAX实现完全并行化RL环境的开源项目。受此启发，开发者决定为《火箭联盟》构建一个类似的GPU并行环境。

已有的RocketSim是一个高性能CPU仿真器，已被集成到RLGym中用于RL训练。但CPU的并行能力有限，在训练大规模神经网络时，host-device之间的数据传输会成为瓶颈。于是，开发者花了近两年时间（断断续续），将物理逻辑移植到CUDA上，项目命名为CARL——CUDA-Accelerated Rocket League。

这个项目的核心思路并不复杂，但实现起来极具挑战：把所有仿真状态保留在GPU显存中，训练过程中彻底消除主机与设备之间的数据传输开销。

JAX 与 PureJaxRL：JAX 是 Google 开发的高性能数值计算库，支持自动微分、向量化和即时编译（JIT），并可通过 XLA 将计算图编译到 GPU/TPU。PureJaxRL 是一个基于 JAX 的强化学习库，它利用 jax.vmap 和 jax.lax.scan 等原语，将整个 RL 训练循环（包括环境仿真、策略推理和参数更新）表达为纯函数，从而在 GPU 上并行执行数千个环境实例。与传统的 CPU 环境（如 Gym）不同，PureJaxRL 的环境状态也是 JAX 数组，因此可以自然地驻留在设备上。CARL 借鉴了这一设计，但针对《火箭联盟》的复杂刚体物理和碰撞检测，用 CUDA 手写了定制化的仿真内核，以在 GPU 上高效推进物理状态。

设备驻留仿真：为什么它如此关键？

传统RL训练流程中，环境仿真在CPU上运行，神经网络推理在GPU上进行。每一步都需要将观测数据从CPU传到GPU，再将动作结果传回CPU——这种来回拷贝看似微小，但在数百万步的训练中会累积成巨大的时间开销。

CARL的做法是让整个仿真循环都留在GPU上。数千局游戏的状态同时驻留在显存中，GPU的数千个核心并行推进每一局游戏的物理模拟。这不仅消除了数据传输，还充分利用了GPU的吞吐能力。

这一理念与PureJaxRL一脉相承，但CARL针对《火箭联盟》的复杂物理做了专门优化。开发者提到，他在这个过程中学到了大量GPU编程和性能优化的知识，同时也坦言“还有很多要学”。

Host-Device 数据传输瓶颈：在典型的 RL 训练循环中，环境仿真（如游戏物理步进）在 CPU 上执行，而策略网络的前向/反向传播在 GPU 上完成。每一步都需要将观测（状态）从主机内存拷贝到设备显存，再将动作从显存拷贝回主机内存，随后 CPU 才能推进仿真。这种拷贝虽然单次延迟不高，但训练通常需要数百万到数十亿步，累积的拷贝时间会远超实际计算时间。此外，PCIe 总线的带宽有限，频繁的小批量传输无法充分利用带宽，进一步放大开销。设备驻留仿真的核心思想，就是让仿真状态和计算图都常驻 GPU 显存，使得整个“仿真→推理→动作→仿真”的闭环完全在设备内完成，仅通过内核启动进行同步，从而消除拷贝开销，并利用 GPU 的数千核心并行加速物理步进。

训练速度：从数周到数小时

最令人印象深刻的成果来自一项对比实验。开发者使用著名《火箭联盟》AI机器人Nexto的默认训练配置来训练智能体，结果在短短几个小时内就达到了Nexto的水平。

而Nexto的原始训练，按照开发者所述，花费了数周的墙上时钟时间（wall clock time）。

这个对比并非严格的控制实验——硬件配置、实现细节等可能有所不同——但量级上的差异足以说明设备驻留仿真的潜力。更重要的是，训练出的智能体能够无缝迁移到真实游戏中，说明CARL的物理模拟与实际游戏保持了良好的一致性。

Nexto 与 Rocket League 智能体水平：Nexto 是《火箭联盟》社区中一个著名的基于强化学习的机器人，由开发者 Roland 训练，其水平已达到甚至超越顶级人类玩家（如 Grand Champion 段位）。Nexto 使用 RLGym 和 RocketSim 在 CPU 上进行训练，训练规模巨大，耗费了数周的真实时间。CARL 使用与 Nexto 相同的训练配置（包括奖励函数、超参数和网络结构），但借助 GPU 并行仿真，在数小时内就达到了相当的竞技水平。需要注意的是，这一对比并非严格的控制实验，因为硬件不同、仿真器实现也有差异，但训练时间从数周缩短到数小时，量级上的提升足以证明设备驻留仿真的巨大潜力。

功能完整，但仍有待完善

开发者表示CARL目前已达到“功能完整”状态，但仍有多项工作待推进：

• 内置训练包API：用于针对性训练特定技能（如空中运球、防守等）
• 速度与一致性基准测试：系统性地评估CARL与RocketSim的物理偏差
• 测试与工程最佳实践：完善代码质量和可维护性
• JAX与TensorFlow集成：探索与主流RL框架的对接

开发者计划在后续发布演示视频和更详细的性能对比。

社区反响与意义

这个项目在Reddit上发布后引起了关注。对于RL研究者而言，CARL展示了一条清晰的路径：将环境仿真从CPU迁移到GPU，可以带来数量级的训练加速。这一思路不仅适用于《火箭联盟》，也可以推广到其他需要复杂物理模拟的RL任务中。

当然，完全在GPU上实现复杂物理引擎并非易事。开发者花了两年时间才走到这一步，且物理一致性的验证仍在进行中。但方向已经明确——当仿真成为瓶颈时，让仿真本身跑在GPU上，可能是最直接的解法。

项目代码和更多细节，开发者表示将在完善后陆续公开。