💻 computer science

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

이 논문은 수치적 정확도를 유지하면서 효과적인 경사 기반 시스템 식별을 가능하게 하는 동시에, Pinocchio와 같은 기존의 CPU 중심 솔루션보다 처리량과 학습 속도 면에서 크게 뛰어난, 배치형 미분 가능 강체 역학을 위한 자립형 GPU 최적화 PyTorch 라이브러리인 BARD를 소개한다.

원저자: Yue Wang, Yanran Xu, Wenbo Wu, Chuanhang Qiu, Zhaoxing Li

게시일 2026-06-01

📖 3 분 읽기☕ 가벼운 읽기

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)에서 작동하도록 설계되었습니다. 연구자들이 수천 대의 로봇을 동시에 훈련시키기 위해 현대적이고 매우 빠른 그래픽 카드(GPU)를 사용하려 했을 때, 이 오래된 도구들은 병목 현상이 되었습니다. 그것은 마치 트랙터들이 가득한 교통 체증 속에서 페라리를 몰려고 하는 것과 같았습니다.

이 논문은 이 교통 체증을 해결하기 위해 PyTorch(대중적인 AI 프레임워크)를 위해 특별히 제작된 새로운 도구인 bard(Batched Articulated Rigid-body Dynamics)를 소개합니다.

bard가 어떻게 작동하는지 쉬운 비유를 통해 설명하겠습니다.

1. "게으른" 요리사 (계층적 지연 평가 - Tiered Lazy Evaluation)

4,000명의 손님을 위한 거대한 연회를 준비하는 요리사를 상상해 보세요.

기존 방식: 요리사는 손님 A는 스프만 원하고 손님 B는 샐러드만 원하더라도, 모든 손님을 위해 모든 요리를 만듭니다. 이는 엄청난 시간과 에너지를 낭비합니다.
bard 방식: 요리사는 똑똑하게 "게으릅니다". 실제로 요청된 것만 요리합니다. 만약 AI가 로봇의 손 위치(순역학, Forward Kinematics)만 알아야 한다면, 요리사는 로봇의 근육 내부의 힘을 계산하는 데 신경 쓰지 않습니다. 만약 힘(forces)만 필요하다면, 요리사는 손의 위치 계산을 건너뜁니다. 불필요한 작업을 건너뜀으로써 엄청난 시간을 절약합니다.

2. "미리 섞어둔" 페인트 (Matmul-Free Transforms)

기존 도구에서는 로봇의 관절이 움직일 때마다 컴퓨터가 새로운 위치를 파악하기 위해 두 개의 커다란 숫자 격자(행렬)를 곱하는 복잡한 연산을 수행해야 했습니다. 이를 수천 번 반복하는 것은 느립니다.

bard 방식: 저자들은 로봇 관절의 "형태"는 변하지 않고 오직 각도만 변한다는 사실을 깨달았습니다. 그래서 공연이 시작되기도 전에 "페인트"를 미리 섞어두었습니다. 관절이 움직일 때마다 매번 페인트를 섞는 대신, 미리 섞어둔 베이스에 약간의 "각도"(사인과 코사인)를 더하기만 하면 됩니다. 이는 느리고 무거운 혼합 과정을 빠르고 간단한 젓기로 바꿔줍니다.

3. "조립 라인" vs "단일 작업자" (Level-Parallel Propagation)

로봇은 나무 구조(몸통에서 팔다리로 이어지는 구조)처럼 만들어집니다. 기존 도구들은 나무의 밑동부터 끝까지, 한 번에 하나의 가지씩 계산했습니다. 즉, 이전 단계가 끝날 때까지 다음 단계를 시작하지 못하고 기다렸습니다.

bard 방식: 조립 라인을 상상해 보세요. 한 명의 작업자가 전체 나무를 다 하는 대신, bard는 같은 높이(깊이)에 있는 모든 가지를 그룹화합니다. 4족 보행 로봇의 네 다리 움직임을 하나씩 순서대로 계산하는 것이 아니라, 동시에 계산합니다. 이는 GPU의 강력한 성능을 활용하여 전체 "레벨"을 한 번에 거대한 배치(batch)로 처리합니다 man합니다.

4. 결과: 속도와 정확도

연구진은 4,096대의 로봇 시뮬레이션을 동시에 실행하는 매우 강력한 컴퓨터 칩인 NVIDIA H200 GPU에서 bard를 테스트했습니다.

속도: 로봇의 손 위치를 파악하는 데 있어, bard는 기존 표준보다 64배 더 빨랐습니다. 로봇의 내부 힘을 계산하는 데 있어서도 실제 훈련 시나리오에서 최대 8.5배까지 훨씬 더 빨랐습니다.
정확도: 이토록 빠름에도 불구하고, bard는 기존의 느린 도구들과 동일한 정확도를 유지합니다. 수학적 오차는 거의 제로에 가까울 정도로 미미했습니다(모래 한 알과 산의 차이만큼이나 작습니다).
실제 테스트: 연구진은 bard를 사용하여 11개의 부위(척추가 있는 4족 보행 로봇)를 가진 로봇에게 강화 학습(Reinforcement Learning) 방식으로 움직이는 법을 가르쳤습니다. bard가 매우 빨랐기 때문에, 훈련 과정은 기존 도구를 사용할 때보다 8.5배 더 빨랐습니다.

5. 이것이 왜 중요한가

이 논문은 bard를 통해 연구자들이 과거에 거대한 컴퓨터 클러스터가 필요했던 복잡한 로봇 훈련을 단 하나의 GPU에서 실행할 수 있음을 보여줍니다. 이는 "드롭인(drop-in)" 교체 모델, 즉 과학자들이 전체 코드를 다시 작성할 필요 없이 기존 도구를 bard로 바로 바꿀 수 있음을 의미합니다.

요약하자면, bard는 로봇 물리 계산의 무겁고 느린 수학적 과정을 현대적 AI의 속도에 맞춰 최적화하여, 로봇이 훨씬 더 빠르고 효율적으로 학습할 수 있게 해줍니다.

기술 요약: PyTorch 기반 배치형 미분 가능 강체 역학 (bard)

문제 정의
로봇 제어가 루프 내 역학 계산을 포함하는 대규모 강화 학습(RL)으로 전환됨에 따라, 분야는 심각한 처리량 병목 현상에 직면해 있습니다. Isaac Lab 및 MuJoCo Playground와 같은 현재의 훈련 파이프라인은 배치형 GPU 실행을 위해 PyTorch에 의존합니다. 그러나 표준 역학 라이브러리인 Pinocchio는 CPU에 종속적이며 순차적입니다. 이러한 구조는 매 평가마다 선형적인 CPU-GPU 데이터 전송 오버헤드를 강제하며, 순차적인 트리 순회 과정 동안 GPU 코어를 유휴 상태로 방치합니다. Brax나 MuJoCo MJX와 같은 기존 GPU 시뮬레이터들은 높은 처리량을 제공하지만, 비(非) PyTorch 생태계(JAX)에 묶여 있거나, 커스텀 손실 함수 또는 물리 정보 기반 보상에 필요한 독립적인 미분 가능 연산자(예: 자코비안 또는 질량 행렬)를 노출하지 않는 단일 구조의 시뮬레이터로 기능합니다. 반대로, ADAM과 같이 PyTorch와 호환되는 기존 미분 가능 라이브러리들은 백엔드 범용성을 위해 GPU 처리량을 희생하며, 커널 퓨전(kernel fusion)을 활용하지 못하거나 PyTorch 컴파일러의 그래프 트레이싱을 깨뜨리는 데이터 의존적 제어 흐름을 피하지 못합니다.

방법론
저자들은 Featherstone의 강체 역학 알고리즘(FK, RNEA, CRBA, ABA 및 자코비안)을 구현한 순수 PyTorch 라이브러리인 bard(Batched Articulated Rigid-body Dynamics)를 제시합니다. 이 라이며지는 배치형 GPU 평가와 torch.compile과의 완전한 호환성을 위해 설계되었습니다. 성능 이득은 네 가지 구체적인 설계 전략을 통해 달성됩니다:

계층적 지연 평가 (Tiered Lazy Evaluation): 모든 운동학적 양을 매 호출 시마다 계산하는 대신, bard는 유효성 플래그(validity flags)에 의해 보호되는 세 가지 계층으로 운동학적 상태를 분리합니다: (1) 로컬 부모-자식 변환, (2) 공간 접착(spatial adjoints, 역학 알고리즘에만 필요), (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$ 접착 행렬을 구축하는 과정을 건너뜁니다. 대신, 사전 계산된 오프셋과 가용 월드 프레임 포즈를 사용하여 자코비안 컬럼을 직접 계산하며, 이를 통해 접착 구축을 효율적인 외적(cross product) 및 행렬-벡터 곱셈으로 대체합니다.

라이브러리는 배치 우선(batch-first) 메모리 레이아웃을 활용하며, 캐시 라인 낭비를 피하기 위해 전치된 공간 변환(spatial transforms)을 연속적으로 저장합니다. 또한 에이거 모드(eager mode, 최대 처리량을 위해 그래디언트 추적을 건너뜀)와 기능 모드(functional mode, 완전한 자동 미분)를 모두 지원하며, 정적 제어 흐름과 상수 인덱스 텐서를 사용하여 torch.compile과 완전히 호환됩니다.

주요 기여

bard: Featherstone 알고리즘의 배치형 GPU 평가를 가능하게 하는 오픈 소스 순수 PyTorch 역학 라이브러리.
시스템 수준 기술: 키네마틱스 캐싱, 사전 할당된 워크스페이스, 레벨 병렬 전파를 포함하여 GPU에서 재귀적 트리 계산을 실행하기 위한 새로운 방법론.
검증: 그래디언트 기반 시스템 식별 및 4,096개의 병렬 환경을 이용한 대규모 RL 훈련에 대한 종합적인 테스트.

실험 결과
저자들은 NVIDIA H200 GPU에서 7~23 DOF 범위의 5가지 로봇 모델을 대상으로 Pinocchio 및 ADAM과 bard를 비교 평가했습니다:

수치 정확도: bard는 부동 소수점 반올림 오차 범위 내에서 Pinocchio와 일치합니다. 최대 불일치는 FK 및 자코비안의 경우 머신 엡실론( $\approx 10^{-16}$ ) 근처였으며, RNEA 및 CRBA의 경우 $10^{-13}$ 미만이었습니다. ABA는 3-패스 구조에서의 오차 누적으로 인해 약간 더 큰 불일치(최대 $10^{-10}$ )를 보였으나 실용적인 범위 내에 있었습니다.
처리량: 배치 크기 4,096에서 bard는 순전파 운동학(Forward Kinematics)에 대해 Pinocchio보다 최대 64배, 자코비안에 대해 63배 높은 처리량을 달성했습니다. 역학 알고리즘의 경우, 속도 향상은 3배에서 7배 사이였습니다.
시스템 식별: 5% 토크 노이즈가 있는 7-DOF 매니퓰레이터에 대한 그래디언트 기반 식별 작업에서, bard는 링크 질량을 평균 오차 **1.24%**로 성공적으로 복구했습니다.
RL 통합: 4,096개의 병렬 환경을 가진 11-DOF 사족보행 로봇의 Isaac Lab AMP 훈련 파이프라인에 통합되었을 때, bard는 Pinocchio보다 8.5배, ADAM보다 2.0배 높은 처리량을 달ach했습니다.

의의 및 주장
본 논문은 bard가 PyTorch를 위한 고성능, 모듈형 역학 라이브러리를 제공함으로써 로보틱스 생태계의 중요한 공백을 메운다고 주장합니다. 이는 이전에 CPU 병목 현상이나 생태계 불일치로 인해 어려움을 겪었던, 단일 GPU 상에서의 대규모 루프 내 역학을 포함한 RL 훈련을 가능하게 합니다. 저자들은 배치 크기가 큰 키네마틱스 중심의 워크로드에서 처리량 이점이 가장 두드之处를 보인다고 언급했습니다. 또한, torch.compile이 레지스터 스필링(register spilling)으로 인해 저사양 GPU(예: NVIDIA L4)의 역학 집약적 워크로드에서 성능을 저하시킬 수 있음을 경고하며, 이러한 경우에는 에이거 모드를 사용할 것을 권장합니다. 이 연구는 직렬 알고리즘을 병렬 GPU 실행으로 재구성하는 것이 수치적 정확성이나 미분 가능성을 희생하지 않고도 상당한 성능 이득을 얻을 수 있음을 입증합니다.