💻 computer science

Batched Differentiable Rigid Body Dynamics in PyTorch for GPU-Accelerated Robot Learning

本文介绍了 BARD，这是一个自包含、针对 GPU 优化的 PyTorch 库，用于处理批处理可微刚体动力学，它在吞吐量和训练速度上显著优于像 Pinocchio 这样受限于 CPU 的现有解决方案，同时保持了数值精度并实现了有效的基于梯度的系统辨识。

原作者： Yue Wang, Yanran Xu, Wenbo Wu, Chuanhang Qiu, Zhaoxing Li

发布于 2026-06-01

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Yue Wang, Yanran Xu, Wenbo Wu, Chuanhang Qiu, Zhaoxing Li

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在试图教会一个机器人走路、跑步或捡起物体。为了高效地完成这项任务，机器人的大脑（AI）需要不断地问自己：“如果我这样移动我的腿，我的平衡会发生什么变化？”这需要一种复杂的数学运算，叫做刚体动力学（rigid-body dynamics）。

长期以来，进行这种数学运算就像是用一把小茶匙试图填满一个游泳池。科学家们使用的标准工具（例如一个名为 Pinocchio 的库）是为单个“大脑”（CPU）设计的，一次只能处理一个计算。当研究人员尝试同时训练成千上万个机器人时，这些旧工具就成了瓶颈。这就像是试图让一辆法拉利在拖拉机的车队中穿行。

这篇论文介绍了 bard（批量关节刚体动力学，Batched Articulated Rigid-body Dynamics），这是一个专门为 PyTorch（一种流行的 AI 框架）构建的新工具，旨在解决这个“交通拥堵”问题。

以下是通过简单的类比来解释 bard 是如何工作的：

1. “偷懒”的厨师（分层延迟计算/Tiered Lazy Evaluation）

想象一位正在为 4,000 名宾客准备盛大宴会的厨师。

旧的方法： 厨师为每一位宾客烹饪每一道菜，即使宾客 A 只想要汤，而宾客 B 只想要沙拉。这浪费了大量的时间和精力。
bard 的方法： 这里的厨师是以一种聪明的方式在“偷懒”。他们只烹饪实际被要求的食物。如果 AI 只需要知道机器人的手在哪里（正向运动学），厨师就不会费力去计算机器人肌肉内部的力量。如果它只需要知道力量，厨师就会跳过手的位姿计算。通过跳过不必要的工作，这节省了大量时间。

2. “预混”的油漆（无矩阵乘法变换/Matmatmul-Free Transforms）

在旧工具中，每当机器人的关节移动时，计算机都必须执行复杂的两个大型数字网格（矩阵）相乘，以确定新的位置。进行成千上万次这样的操作非常缓慢。

bard 的方法： 作者意识到，机器人的关节“形状”永远不会改变，改变的只是角度。因此，他们在表演开始前就预先混合好了“油漆”（常量）。与其在每次关节移动时都重新混合油漆，不如只需加入一点点“角度”（正弦和余弦）到预混好的基底中。这把一个缓慢、沉重的混合过程变成了一个快速、简单的搅拌过程。

讨论 3. “流水线” vs. “单人工作者”（层级并行传播/Level-Parallel Propagation）

机器人就像树一样构建（一个带有手臂和腿的躯干）。旧工具从树干到树梢，一次处理一个分支，必须等待前一步完成后才能开始下一步。

bard 的方法： 想象一条流水线。bard 不是由一名工人完成整棵树，而是将所有处于相同高度（深度）的分支组合在一起。它同时计算四足机器人四条腿的运动，而不是一个接一个地计算。它在一个巨大的批次中处理整个“层级”的树结构，从而充分利用 GPU 的强大算力。

4. 结果：速度与精度

研究人员在 NVIDIA H200 GPU（一种非常强大的计算芯片）上测试了 bard，同时运行了 4,096 个机器人模拟。

速度： 在计算机器人手部位置方面，bard 比旧的标准工具快了 64 倍。在计算机器人的内部力量时，它仍然显著更快（在实际训练场景中快了高达 8.5 倍）。
精度： 尽管速度如此之快，但 bard 与旧的、缓慢的工具一样精确。其数学误差微乎其微，几乎可以忽略不计（就像是一粒沙子与一座大山之间的差别）。
现实测试： 他们使用 bard 来教一个 11 腿机器人（一个带有脊椎的四足机器人）如何通过一种称为强化学习的方法进行移动。由于 bard 如此之快，训练过程比使用旧工具快了 8.5 倍。

5. 为什么这很重要

这篇论文表明，bard 允许研究人员在单个 GPU 上运行复杂的机器人训练，而以前这需要一个庞大的计算机集群。它作为一个“即插即用”的替代方案，这意味着科学家们可以在不重写整个代码的情况下，直接将旧工具替换为 bard。

简而言之，bard 将机器人物理学中沉重、缓慢的数学运算进行了优化，使其能够以现代 AI 的速度运行，从而让机器人学习得更快、更高效。

技术摘要：基于 PyTorch 的批处理可微刚体动力学 (bard)

问题陈述
随着机器人控制向大规模强化学习（RL）及循环内动力学计算转型，该领域面临着关键的吞吐量瓶颈。当前的训练流水线（如 Isaac Lab 和 MuJoCo Playground 中的流水线）依赖于 PyTorch 进行批处理 GPU 执行。然而，标准的动力学库 Pinocchio 是受限于 CPU 且串行的。这种架构导致每次评估都会产生线性的 CPU–GPU 数据传输开销，并使 GPU 核心在串行树遍历期间处于闲置状态。虽然现有的 GPU 模拟器（如 Brax 和 MuJoCo MJX）具有高吞吐量，但它们绑定于非 PyTorch 生态系统（JAX），或者作为单体式模拟器存在，无法提供独立的、可微的算子（例如雅可比矩阵或质量矩阵），而这些算子对于自定义损失函数或物理信息奖励至关重要。相反，现有的兼容 PyTorch 的可微库（如 ADAM）虽然牺牲了 GPU 吞吐量以换取后端通用性，但未能充分利用内核融合，也无法避免破坏 PyTorch 编译器图追踪的数据依赖型控制流。

方法论
作者提出了 bard（Batched Articulated Rigid-body Dynamics，批处理关节刚体动力学），这是一个纯 PyTorch 库，实现了 Featherstone 刚体动力学算法（包括 FK、RNEA、CRBA、ABA 和雅可比矩阵）。该库旨在进行批处理 GPU 评估，并完全兼容 torch.compile。其性能提升通过四种特定的设计策略实现：

分层延迟评估（Tiered Lazy Evaluation）： bard 并没有在每次调用时计算所有运动学量，而是将运动学状态分为三个由有效性标志保护的层级：(1) 本地父子变换，(2) 空间伴随矩阵（仅在动力学算法中需要），以及 (3) 世界坐标系位姿（仅在 FK/雅可比查询中需要）。这确保了仅调用 RNEA 的 RL 循环可以避免世界坐标系传播的计算成本。
无矩阵乘法的关节变换（Matrix-Multiplication-Free Joint Transforms）： 为了消除在 GPU 上启动数千个小型 $4 \times 4$ 矩阵乘法带来的开销，该库在模型构建阶段预计算了罗德里格斯（Rodrigues）旋转常数。在运行时，父子变换通过仅使用 $\sin(q_i)$ 和 $\cos(q_i)$ 的标量-矩阵乘法来恢复，从而允许 PyTorch 编译器将所有关节和批处理元素中的操作融合进单个内核中。
层并行树传播（Level-Parallel Tree Propagation）： 虽然 Featherstone 算法本质上是串行的，但 bard 将树遍历重新组织为层并行的批处理操作。通过将相同树深度的节点进行分组，该库将 $n_b$ （物体数量）个串行步骤减少为 $d$ （树深度）个批处理步骤。虽然由于数据依赖关系，反向传递（力累积）仍然是串行的，但所有 $B$ 个环境都在每个节点处同时进行处理。
直接雅可比列构造（Direct Jacobian Column Construction）： 对于单自由度（single-DOF）关节，该库绕过了构建完整的 $6 \times 6$ 伴随矩阵的过程。相反，它利用预计算的偏移量和可用的世界坐标系位姿直接计算雅可比列，用高效的叉积和矩阵-向量乘法取代了伴随矩阵的构建。

该库采用批处理优先（batch-first）的内存布局，并连续存储转置后的空间变换，以避免缓存行浪费。它同时支持 eager 模式（跳过梯度追踪以获得最大吞吐量）和 functional 模式（全自动微分），并通过使用静态控制流和常量索引张量，实现了与 torch.compile 的完全兼容。

核心贡献

bard： 一个开源的、纯 PyTorch 的动力学库，支持 Featherstone 算法的批处理 GPU 评估。
系统级技术： 在 GPU 上运行递归树计算的新方法，包括运动学缓存、预分配工作空间和层并行传播。
验证： 在基于梯度的系统辨识任务以及包含 4,096 个并行环境的大规模 RL 训练中进行了全面测试。

实验结果
作者在 NVIDIA H200 GPU 上，针对五种机器人模型（自由度从 7 到 23 不等）对 bard 与 Pinocchio 和 ADAM 进行了对比评估：

数值精度： bard 在浮点舍入误差范围内与 Pinocchio 保持一致。对于 FK 和雅可比矩阵，最大差异接近机器极小值（ $\approx 10^{-16}$ ）；对于 RNEA 和 CRBA，差异低于 $10^{-13}$ 。由于 ABA 三阶段结构中的误差累积，其显示出略大的差异（最高达 $10^{-10}$ ），但在实际范围内仍属正常。
吞吐量： 在批处理大小为 4,096 时，bard 在正向运动学（FK）方面比 Pinocchio 高出高达 64 倍，在雅可比矩阵方面高出 63 倍。对于动力学算法，加速比在 3 倍至 7 倍之间。
系统辨识： 在一个带有 5% 转矩噪声的 7 自由度机械臂的梯度系统辨理任务中，bard 成功恢复了连杆质量，平均误差为 1.24%。
RL 集成： 在集成到 Isaac Lab AMP 训练流水线（用于 11 自由度四足机器人，4,096 个并行环境）中时，bard 的吞吐量比 Pinocchio 高 8.5 倍，比 ADAM 高 2.0 倍。

意义与主张
论文声称 bard 通过为 PyTorch 提供高性能、模块化的原生动力学库，填补了机器人生态系统中的关键空白。它使得在单个 GPU 上进行带有循环内动力学的大规模 RL 训练成为可能，而此前这一任务受限于 CPU 瓶颈或生态系统不兼容。作者指出，在大型批处理规模下，运动学主导的工作负载其吞吐量优势最为显著。他们同时也提醒，由于寄存器溢出问题，torch.compile 在低端 GPU（如 NVIDIA L4）上处理动力学密集型工作负载时可能会降低性能，因此建议在这种情况下使用 eager 模式。这项工作证明，通过重新组织串行算法以实现并行 GPU 执行，可以在不牺牲数值精度或可微性的前提下，获得显著的性能提升。