LIBERO를 MuJoCo Warp로 포팅: AMD GPU에서 수행한 130개 로봇 조작 태스크

LIBERO는 로봇 조작 벤치마크입니다. 이 벤치마크는 다섯 가지 스위트(공간, 객체, 목표, scene10, scene90)에 걸쳐 총 130개의 태스크로 구성되어 있으며, 각 태스크마다 인간의 시연과 언어적 목표가 주어집니다. LIBERO는 언어 조건부 모방 학습(language-conditioned imitation)의 표준 테스트베드이며, 일반적으로 CPU MuJoCo를 사용하는 robosuite에서 한 번에 하나의 환경으로 실행됩니다.

저는 이 모든 것을 MuJoCo Warp로 포팅하여 RX 9070 XT($700, 16 GB, RDNA4)에서 실행했습니다. 물리 시뮬레이션은 2,048개의 월드에서 초당 17,137 env-steps를 수행합니다. CPU robosuite는 초당 38을 기록했습니다. BC(Behavioral Cloning) 트랜스포머를 50 에폭 동안 학습시키면 Warp에서는 42.5%의 성능을 얻었고, CPU에서는 50%였습니다.

GPU가 중요한 이유: 행동 복제(behavioral cloning)는 롤아웃(rollouts)이 필요합니다. 정책을 실행하여 실패하는 지점을 확인하고, 그곳에서 레이블이나 시연을 생성해야 합니다. CPU 환경에서는 한 번에 하나의 환경만 처리할 수 있어, 단일 스위트(10개 태스크)의 시연을 생성하는 데 8~9시간이 걸렸습니다. 하지만 GPU를 사용하면 몇 분 만에 가능하여, 단순히 한 번 실행하는 것이 아니라 반복적인 개선이 가능합니다.

AMD에서 Warp를 사용하는 것이 흥미로운 부분입니다. Warp는 NVIDIA의 GPU 시뮬레이션 프레임워크이며, AMD HIP/ROCm 포트는 최근에 이루어졌습니다(Tomas Thoresen, Strix Halo). RDNA4에서 이것을 작동시키고, Warp가 gfx1201용 HIP 커널을 컴파일하며, JAX가 rocm:0을 인식하고, PyTorch가 cuda를 인식하도록 만든 것이 가능하게 했습니다. 렌더러가 가장 까다로운 부분이었습니다. BC 정책은 픽셀의 고정된 함수인데, Warp의 레이 트레이서(ray tracer)는 MuJoCo의 CPU 렌더러와 다릅니다. 제가 처음 Warp로 평가했을 때 점수는 0%였습니다. 네 가지 수정 사항을 거쳐 42.5%까지 올릴 수 있었습니다: 세로 반전(vertical flip), 그림자 상수 0.3 -> 0.0, 밝기 1.15배 증가, 그리고 큐브맵 샘플링(테이블의 나무 무늬가 평평하게 렌더링되던 문제). 밝기만 조정하는 것만으로도 12.5점만큼의 가치가 있었습니다.

모든 것이 ROCm 기반의 공개 소스에서 구축되었으며, 예제와 21.6 MB 크기의 BC 체크포인트는 리포지토리에 포함되어 있습니다. 제가 완료하지 못한 부분은 다음과 같습니다: Warp 경로는 롤아웃을 수집하지만, 학습 자체는 여전히 PyTorch에서 오프라인으로 진행됩니다. 저는 여름 동안 AMD의 R9700 및 일부 Instinct 카드를 다루었기 때문에, 다음 목표는 인-루프 비전(in-loop vision)입니다.

작성 글: https://amohan.dev/blog/2026/libero-warp-mjx-rdna4/
코드: https://github.com/poad42/libero_mjx