[논문 리뷰] Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation

C. Daniel Freeman, Erik Frey외 4명|arXiv (Cornell University)|2021. 06. 24.

Context-Aware Activity Recognition Systems인용 수 61

한 줄 요약

Brax는 JAX로 만들어진 오픈소스 미분 가능 강체 물리 엔진으로, 가속기에서 실행되어 대규모 병렬 환경 시뮬레이션과 번들 RL 알고리즘을 통한 장치(on-device) 학습을 가능하게 합니다.

ABSTRACT

We present Brax, an open source library for rigid body simulation with a focus on performance and parallelism on accelerators, written in JAX. We present results on a suite of tasks inspired by the existing reinforcement learning literature, but remade in our engine. Additionally, we provide reimplementations of PPO, SAC, ES, and direct policy optimization in JAX that compile alongside our environments, allowing the learning algorithm and the environment processing to occur on the same device, and to scale seamlessly on accelerators. Finally, we include notebooks that facilitate training of performant policies on common OpenAI Gym MuJoCo-like tasks in minutes.

연구 동기 및 목표

• 가속기에서 물리와 학습을 같은 장소에 colocating 하여 RL 연구의 속도와 접근성을 높인다.
• 미분 가능하고 오픈 소스 엔진을 제공하여 기울기 기반 RL 방법을 가능하게 한다.
• 프로토버프(ProtoBuf) 기반 환경 명세와 Gym 유사 인터페이스를 제공하여 확장 가능한 과제를 지원한다.
• Brax의 확장성 및 성능을 가속기에서의 실행에 걸쳐 MuJoCo와 비교 평가한다.
• 미분 가능 물리 엔진의 개선을 위한 한계점 및 향후 연구 방향을 제시한다.

제안 방법

• QP 데이터 프리미티브와 벡터화 변환을 이용한 최대 좌표 미요소rigid-body 다이나믹스 구현.
• 대규모 환경 수천 개를 가속기에서 실행하기 위해 JAX를 이용한 자동 벡터화, 디바이스-병렬성, JIT 및 자동 미분 활용.
• 물리적 단계들을 Brax_system에 적용되는 2차 아이솔렉틱 오일러 적분기와 함께 병렬 변환(조인트, 액추에이터, 충돌체)로 정의.
• ProtoBuf 기반 시스템 명세와 물리 시나리오를 구축하고 해결하기 위한 gym 유사 Env를 제공.
• PPO, SAC, ES, APG를 JAX에 번들링하여 가속기에서 번들 تشغيل 및 디그리언트 기반 업데이트를 위한 온디바이스 롤아웃으로 실행.
• SAC를 사용하여 평가할 때 MuJoCo 유사 Ant, Humanoid, Halfcheetah, Grasp, Fetch 등의 환경군을 시연한다.

실험 결과

연구 질문

• RQ1 가속기에서 실행되는 미분 가능 물리 엔진이 환경 시뮬레이션과 학습을 같은 공간에 배치함으로써 RL 학습 속도를 크게 높일 수 있는가?
• RQ2 Brax의 TPU/GPU에서의 성능과 확장성은 MuJoCo와 같은 전통적인 CPU 기반 엔진과 비교해 어떤가? 일반 RL 벤치마크에서?
• RQ3 Bra x에서 구현된 PPO, SAC, ES, APG의 학습 알고리즘에 대한 미분 가능성의 효과는 어떤가?
• RQ4 대규모의 differentiable rigid-body 시뮬레이션에 필요한 실용적 한계점 및 튜닝은 무엇인가?
• RQ5 Prot oBuf 환경 명세와 gym 유사 인터페이스가 다양한 물리 기반 과제 구축에 얼마나 효과적인가?

주요 결과

• 일부 환경에서 Brax는 단일 가속기에서 초당 수백만 개의 시뮬레이션 스텝을 달성한다.
• 4x2 TPUv3에서 Brax는 가속기 간 분산 계산으로 초당 수억 개의 환경 스텝까지 확장한다.
• PPO 및 SAC 구현이 Brax에 맞춰 컴파일되어 온 디바이스 롤아웃 및 그래디언트 업데이트를 최소한의 데이터 전송 오버헤드로 가능하게 하며 비가속 기준 대비 훨씬 빠른 학습을 달성한다.
• SAC로 평가된 유사 하이퍼파라미터에서 Brax 환경은 MuJoCo 상대 학습 궤적 및 보상과 질적으로 일치한다.
• 비교적 표준 PPO 구현에 비해 Brax의 PPO로의 최적화를 통해 Ant에서의 보행 학습 시간이 약 10초 정도로 단축된다(약 30분의 표준 PPO 구현에 비해).
• Brax는 모멘텀 보존과 에너지/각운동량 거동 측면에서 다른 엔진에 비해 양호한 성능을 보인다.