기술 요약: frax: JAX 기반의 고속 로봇 운동학 및 동역학

문제 제기

로봇 제어, 계획, 학습 분야에서 고성능, 사용 편의성, 그리고 다양한 하드웨어 아키텍처 (CPU, GPU, TPU) 와의 호환성을 동시에 갖춘 강체 동역학 라이브러리에 대한 절실한 필요가 존재합니다. 기존 솔루션들은 상당한 트레이드오프에 직면해 있습니다:

• CPU 최적화 라이브러리: Pinocchio 와 MuJoCo 와 같은 라이브러리는 CPU 에서의 저지연 실행에 탁월하지만, 가속기에서 직접 실행될 수 없습니다. 이들의 Python 바인딩은 종종 오버헤드를 유발하고 자동 미분 (AD) 워크플로우와의 통합을 복잡하게 만듭니다.
• GPU 최적화 라이브러리: MJX 와 BRAX 와 같은 라이브러리는 GPU 병렬성을 활용하여 높은 처리량의 시뮬레이션을 수행하지만, CPU 에서 단일 로봇 제어를 사용할 때 높은 오버헤드를 겪는 경우가 많아, 실시간 킬로헤르츠 제어 루프에는 적합하지 않습니다.
• AD 호환성: C++ 기반 동역학 라이브러리를 Python 기반의 머신러닝 또는 계획 워크플로우에 통합할 때, 언어 경계를 가로지르는 자동 미분 적용의 어려움으로 인해 마찰이 자주 발생합니다.

성능이나 사용 편의성을 희생하지 않으면서 여러 아키텍처를 대상으로 하는 통합 프레임워크는 부재합니다.

방법론

저자들은 frax를 소개합니다. 이는 JAX를 기반으로 구축된 순수 Python 라이브러리로, 로봇 운동학 및 동역학을 위한 고성능 인터페이스를 제공합니다. 핵심 방법론적 혁신은 대부분의 C++ 라이브러리에서 사용되는 전통적인 재귀적 트리 탐색 방법 (예: 재귀적 뉴턴 - 오일러 알고리즘) 과는 구별되는, 강체 동역학에 대한 완전 벡터화 접근법입니다.

주요 기술 구성 요소:

1. 벡터화된 공식화: 재귀적 루프 대신, frax 는 조상 마스크 (ancestor mask)($U \in \{0, 1\}^{n \times n}$) 를 사용하여 운동학 트리 구조를 인코딩합니다. 이를 통해 동역학 계산을 행렬 곱셈과 배열 브로드캐스팅을 통해 수행할 수 있습니다.
   • 역동역학 (RNEA): 마스크 $U$를 사용하여 트리 위아래로 힘과 가속도를 합산함으로써 순차적 종속성을 피하고 계산합니다.
   • 질량 행렬 (CRBA): 복합 관성량의 축적을 벡터화하여 계산합니다.
2. 복잡도 트레이드오프: 이 접근법은 RNEA 와 같은 알고리즘의 이론적 계산 복잡도를 $O(n)$에서 $O(n^2)$로, CRBA 를 $O(nd)$에서$O(n^2)$로 (여기서$n$은 자유도) 의도적으로 증가시킵니다. 그러나 저자들은 현대 하드웨어 (XLA 컴파일을 통한) 에서의 세밀한 병렬성으로 인한 성능 향상과 루프 오버헤드 제거가 중복 계산보다 훨씬 더 중요하다고 주장합니다.
3. JAX 통합: JAX 를 활용함으로써 frax 는 다음을 네이티브로 지원합니다:
   • Just-In-Time (JIT) 컴파일: CPU, GPU, TPU 에서 최적화된 실행을 위해.
   • 자동 미분 (AD): 긴 언롤링 루프를 추적하는 비용 없이 최적화 기반 제어를 위한 효율적인 기울기 계산을 가능하게 함.
   • 야코비안 - 벡터 곱 (JVPs): 전체 야코비안을 계산하지 않고도 제어 제약 조건 (예: 제어 장벽 함수) 에 대한 리 미분 (Lie derivatives) 을 효율적으로 평가할 수 있게 함.

주요 기여

1. 통합 아키텍처: frax 는 CPU, GPU, TPU 에서 효율적으로 실행되는 단일 코드베이스를 제공하여 저지연 제어와 고처리량 시뮬레이션 간의 격차를 해소합니다.
2. 성능:
   • CPU 에서: 킬로헤르츠 제어 속도에 적합한 마이크로초 수준의 계산 시간을 달성하여, Pinocchio 와 MuJoCo 의 Python 바인딩보다 2~5 배 빠르며 최적화된 C++ 구현체의 원시 성능에 근접합니다.
   • GPU/TPU 에서: 수천 개의 병렬 인스턴스로 확장되어 초당 1 억 개 이상의 동역학 평가를 달성하며, MJX 및 BRAX 와 같은 전용 GPU 시뮬레이터의 성능과 맞먹거나 능가합니다.
3. 사용 편의성: 순수 Python 라이브러리로서 프로토타이핑을 단순화하고 Python 생태계와 원활하게 통합하며, 높은 성능을 유지하면서 C++ 바인딩의 필요성을 제거합니다.
4. 오픈 소스: 이 라이브러리는 오픈 소스 소프트웨어로 출시되었으며, Franka Panda 와 Unitree G1 과 같은 특정 로봇에 대한 지원과 조정된 충돌 모델을 포함합니다.

결과

저자들은 Franka Panda 매니퓰레이터와 Unitree G1 휴머노이드에서 frax 를 검증했습니다.

• 지연 시간: CPU 에서 frax 는 역운동학 (미분 IK) 시간에서 Franka 기준 4.09 µs, G1 기준 6.71 µs를 달성했으며, 역동역학 (작업 공간 제어) 시간에서는 Franka 기준 7.03 µs, G1 기준 24.6 µs를 달성했습니다.
• 비교:
  • Pinocchio 와 MuJoCo 의 Python API 와 비교할 때, frax 는 JIT 컴파일과 컴파일러 퓨전으로 인해 2~5 배 더 빠릅니다.
  • Pinocchio/MuJoCo 의 C++ 인터페이스와 비교할 때, frax 는 원시 성능에 근접하지만 Python 의 유연성과 AD 기능을 유지합니다.
  • MJX 와 BRAX 와 비교할 때, frax 는 GPU 에서 동등하거나 약간 더 빠른 성능을 보였지만, JIT 컴파일 시간을 612 초에서 12 초로 크게 단축하여 더 빠른 반복 개발을 가능하게 했습니다.
• 처리량: GPU 에서 이 라이브러리는 대용량 배치 (4096) 에서 초당 1 억 개 이상의 평가 처리량을 보여주어 대규모 병렬 학습 환경에 적합함을 나타냈습니다.

중요성 및 주장

이 논문은 frax가 로봇 제어 및 학습 워크플로우를 통합하는 방향으로 중요한 진전을 이루었다고 주장합니다. 그 주요 중요성은 동일한 코드베이스 내에서 CPU 에서의 실시간 단일 로봇 제어를 가능하게 하는 동시에 가속기에서의 대규모 병렬 시뮬레이션을 지원한다는 점에 있습니다.

저자들은 frax 가 연구자들이 다음을 가능하게 한다고 강조합니다:

• 하드웨어 배포에 필요한 성능을 희생하지 않고 자동 미분 (예: 비선형 궤적 최적화 또는 MPC 용) 을 사용하여 Python 에서 제어기를 프로토타이핑할 수 있습니다.
• 안전이 중요한 제약 조건을 위해 전체 야코비안 계산을 효율적인 JVP 로 대체할 수 있습니다.
• 알고리즘을 다시 작성하지 않고도 서로 다른 하드웨어 백엔드 (실행용 CPU, 학습용 GPU/TPU) 에 동일한 로직을 배포할 수 있습니다.

이 논문은 frax 가 Pinocchio 와 같은 포괄적인 라이브러리 (예: 전역 IK 솔버 및 전체 시뮬레이션 환경 부재) 보다 현재는 더 최소화된 형태이며, C++ 구현체가 여전히 약간 더 빠른 제어 속도를 제공할 수 있음을 인정하며 겸손한 입장을 유지합니다. 그러나 frax 를 현대 로봇 제어와 머신러닝의 교차점을 위해 특별히 설계된 유연하고 고성능의 기반으로서 포지셔닝합니다.