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

📄 arXiv: 2605.31481v1 📥 PDF

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

分类: cs.RO

发布日期: 2026-05-29

🔗 代码/项目: GITHUB


💡 一句话要点

BARD:用于GPU加速机器人学习的PyTorch批量可微刚体动力学库

🎯 匹配领域: 支柱一:机器人控制 (Robot Control) 支柱二:RL算法与架构 (RL & Architecture) 支柱八:物理动画 (Physics-based Animation)

关键词: 刚体动力学 GPU加速 机器人学习 自动微分 PyTorch Featherstone算法 强化学习

📋 核心要点

  1. 现有机器人控制依赖于CPU绑定的Pinocchio等库,成为GPU训练管道的吞吐量瓶颈,限制了大规模强化学习的应用。
  2. BARD通过分层缓存、无矩阵乘法关节变换和水平并行传播等优化手段,实现了高效的批量GPU刚体动力学计算和自动微分。
  3. 实验表明,BARD在吞吐量上显著优于Pinocchio,并成功应用于机器人系统辨识和强化学习任务,验证了其性能和可微性。

📝 摘要(中文)

本文提出了BARD(Batched Articulated Rigid-body Dynamics),一个自包含的PyTorch实现的Featherstone刚体动力学算法库,专门为批量GPU评估和自动微分优化。BARD通过三个关键设计实现高效性:分层惰性求值缓存,避免冗余的树遍历;通过预计算的Rodrigues常数实现无矩阵乘法的关节变换;以及减少顺序操作到树深度批量步骤的水平并行传播。在五个机器人模型(7-23自由度)上,BARD在数值上与Pinocchio匹配,同时在NVIDIA H200上,批量大小为4096时,正向运动学吞吐量高达64倍,雅可比矩阵吞吐量高达63倍。通过对7自由度机械臂的基于梯度的系统辨识验证了可微性,在5%的扭矩噪声下,连杆质量的平均误差恢复到1.24%。BARD被集成到Isaac Lab AMP训练管道中,用于具有4096个并行环境的11自由度脊柱四足机器人,其在环动力学计算速度比Pinocchio快8.5倍,比ADAM快2.0倍。BARD已开源。

🔬 方法详解

问题定义:论文旨在解决机器人控制中,大规模强化学习对动力学计算的需求与现有CPU绑定库(如Pinocchio)造成的GPU训练瓶颈之间的矛盾。现有方法在GPU加速方面存在不足,限制了并行计算能力,影响了训练效率。

核心思路:论文的核心思路是利用PyTorch构建一个自包含、可微分的刚体动力学库BARD,并针对GPU架构进行优化,实现高效的批量计算和自动微分。通过优化算法和数据结构,减少计算冗余,提高并行度,从而加速机器人学习过程。

技术框架:BARD的整体框架基于Featherstone算法,包括正向运动学、逆动力学和雅可比矩阵计算等核心模块。其主要流程包括:1)构建机器人模型的树状结构;2)利用分层缓存避免重复计算;3)使用预计算的Rodrigues常数进行关节变换;4)通过水平并行传播进行动力学计算;5)利用PyTorch的自动微分功能计算梯度。

关键创新:BARD的关键创新在于针对GPU架构的优化设计:1)分层惰性求值缓存:避免冗余的树遍历,减少重复计算;2)无矩阵乘法的关节变换:通过预计算的Rodrigues常数,避免耗时的矩阵乘法;3)水平并行传播:将顺序操作减少到树深度级别的批量步骤,提高并行度。

关键设计:BARD的关键设计包括:1)分层缓存的层级结构和缓存策略;2)Rodrigues常数的预计算方法和存储格式;3)水平并行传播的调度策略和数据依赖关系;4)损失函数的设计,例如在系统辨识任务中使用的均方误差损失函数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

BARD在五个机器人模型上,正向运动学吞吐量最高提升64倍,雅可比矩阵计算吞吐量最高提升63倍(批量大小为4096,NVIDIA H200)。在7自由度机械臂的系统辨识实验中,在5%扭矩噪声下,连杆质量的平均误差恢复到1.24%。在11自由度四足机器人的强化学习训练中,BARD比Pinocchio快8.5倍,比ADAM快2.0倍。

🎯 应用场景

BARD可广泛应用于机器人控制、强化学习、系统辨识、运动规划等领域。其高效的GPU加速能力和自动微分功能,能够加速机器人算法的开发和训练,提高机器人性能。尤其适用于需要大量并行环境的强化学习任务,例如复杂地形下的四足机器人运动控制。

📄 摘要(原文)

As robot control shifts toward large-scale reinforcement learning with in-loop dynamics computation, the community's reliance on CPU-bound libraries such as Pinocchio creates a throughput bottleneck in GPU-based training pipelines. We present BARD (Batched Articulated Rigid-body Dynamics), a self-contained PyTorch implementation of Featherstone's rigid-body dynamics algorithms, optimized for batched GPU evaluation and automatic differentiation. Three design choices make this efficient: a tiered lazy-evaluation cache that avoids redundant tree traversals, matmul-free joint transforms via pre-computed Rodrigues constants, and level-parallel propagation that reduces sequential operations to tree-depth batched steps. On five robot models (7-23 DOFs), BARD matches Pinocchio numerically while reaching up to 64x higher throughput for Forward Kinematics and 63x for Jacobians at batch size 4096 on an NVIDIA H200. We validate differentiability through gradient-based system identification on a 7-DOF manipulator, recovering link masses to 1.24% mean error under 5% torque noise, and integrate BARD into an Isaac Lab AMP training pipeline for an 11-DOF spined quadruped with 4096 parallel environments, where it is 8.5x faster than Pinocchio and 2.0x faster than ADAM for in-loop dynamics. BARD is open-sourced at: https://github.com/YueWang996/bard-pytorch-dynamics.