🤖 AI

JaxMARL: Multi-Agent RL Environments and Algorithms in JAX

이 논문은 JAX를 활용하여 GPU 가속화된 대규모 병렬 다중 에이전트 강화 학습 환경과 알고리즘을 제공함으로써 기존 방식보다 상당한 속도 향상을 달est하고 StarCraft Multi-Agent Challenge의 유연한 엔진 프리 재구현을 제공하는 오픈 소스 Python 라이브러리인 JaxMARL을 소개한다.

원저자: Alexander Rutherford, Benjamin Ellis, Matteo Gallici, Jonathan Cook, Andrei Lupu, Gardar Ingvarsson, Timon Willi, Ravi Hammond, Akbir Khan, Christian Schroeder de Witt, Alexandra Souly, Saptarashmi Ba

게시일 2026-07-07

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Alexander Rutherford, Benjamin Ellis, Matteo Gallici, Jonathan Cook, Andrei Lupu, Gardar Ingvarsson, Timon Willi, Ravi Hammond, Akbir Khan, Christian Schroeder de Witt, Alexandra Souly, Saptarashmi Bandyopadhyay, Mikayel Samvelyan, Minqi Jiang, Robert Tjarko Lange, Shimon Whiteson, Bruno Lacerda, Nick Hawes, Tim Rocktaschel, Chris Lu, Jakob Nicolaus Foerster

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신은 로봇 팀이 서로 협력하는 법을 가르치려 한다고 상상해 보세요. 인공지능의 세계에서는 이를 **다중 에이전트 강화학습(Multi-Agent Reinforcement Learning, MARL)**이라고 부릅니다. 보통 이 로봇들을 가르치기 위해 연구자들은 수천 번의 시뮬레이션을 실행해야 합니다.

전통적인 방식을 설명하자면, 한 명의 느린 코치가 경기장을 직접 뛰어다니며 한 번에 한 명의 선수에게만 지시를 내리며 축구팀을 가르치는 것과 같습니다. 작동은 하지만, 시간이 너무 오래 걸립니다. 만약 100가지의 서로 다른 코칭 전략을 테스트하고 싶다면, 결과가 나오기를 기다리는 데만 몇 달을 보낼 수도 있습니다. 이것이 바로 논문에서 설명하는 "병목 현상"입니다. 시뮬레이션을 처리하기 위해 사용되는 컴퓨터(CPU)가 팀 단위의 에이전트를 효과적으로 훈련하는 데 필요한 방대한 연습량을 감당하기에는 너무 느리기 때문입니다.

JaxMARL의 등장.

저자들은 JaxMARL이라는 새로운 도구를 만들었습니다. 여러분은 이 도구를 '단 한 명의 느린 코치'를 '10,000대의 드론을 거느린 슈퍼 오케스트라 지휘자'로 업그레이드한 것으로 생각할 수 있습니다.

그들이 어떻게 이 일을 해냈는지, 그리고 무엇을 발견했는지를 다음과 같이 쉽게 나누어 설명합니다.

1. 속도 향상 (The "Drone Fleet")

표준 컴퓨터 프로세서에서 하나씩 시뮬레이션을 실행하는 대신, JaxMARL은 컴퓨터가 강력한 그래픽 카드(GPU)—게이밍용 칩과 동일한 것—를 사용하여 수학적 계산을 수행할 수 있게 해주는 JAX라는 특별한 프로그래밍 라이브러리를 사용합니다.

비유: 만약 벽 10,000개를 칠해야 한다고 가정해 봅시다. 기존 방식(CPU)은 한 명의 화가가 벽 하나를 칠하고, 말리고, 그다음 벽을 칠하는 것과 같습니다. JaxMARL 방식은 10,000개의 벽을 동시에 분사할 수 있는 기계를 사용하는 것과 같습니다.
결과: 논문은 이 시스템이 단일 훈련 실행 시 14배 더 빠르다고 주장합니다. 하지만 연구자들이 흔히 하는 것처럼 여러 실험을 동시에 실행할 때는 최대 12,500배 더 빠릅니다. 이는 몇 주가 걸리던 과정을 몇 시간 또는 몇 분으로 단축합니다.

2. 새로운 놀이터 (Environments)

이 AI 에이전트들을 훈련시키려면 "게임"이나 환경이 필요합니다. 이 논문은 이 초고속 시스템에서 실행되도록 다시 작성된 다양한 게임 라이브러리를 소개합니다.

SMAX (StarCraft 대체제): AI 팀을 훈련시키는 데 가장 인기 있는 게임 중 하나는 스타크래프트 II를 기반으로 합니다. 하지만 원본 게임은 AI를 훈련시키기 위해 전체 3D 그래픽 엔진을 구동해야 하므로 무겁고 느립니다.
- 해결책: 저자들은 SMAX를 만들었습니다. 이것을 게임의 "스켈레톤(뼈대) 버전"이라고 생각하면 됩니다. 스타크래프트의 모든 규칙과 전략은 유지하되, 화려한 3D 그래픽은 제거했습니다. 단순히 GPU에서 로직만 실행되기 때문에, 원본 게임 엔진보다 40,000배 더 빠릅니다.
STORM (Grid World): 그들은 또한 STORM이라는 새로운 게임 세트를 구축했습니다. 플레이어들이 그리드 위를 움직이며 코인을 모으는 보드 게임을 상상해 보세요. 하지만 규칙은 그들이 얼마나 잘 협력하거나 경쟁하는지를 테스트하도록 설계되었습니다. 이는 연구자들이 에이전트가 어떻게 협력하거나 서로를 속이는 법을 배우는지 연구하는 데 도움을 줍니다.

3. 코치들 (Algorithms)

빠른 놀이터를 갖는 것만으로는 충분하지 않습니다. 좋은 훈련 방법이 필요합니다. 논문은 또한 이 새로운 빠른 시스템에서 작동하도록 몇 가지 유명한 "코칭 전략"(PPO 및 QMIX와 같은 알고리즘)을 다시 작성했습니다. 그들은 이 새로운 빠른 코치들이 기존의 느린 코치들과 똑같이 똑똑한 결과를 만들어낸다는 것을 검증했습니다.

4. 왜 이것이 중요한가 (The "Evaluation Crisis")

논문은 이 분야의 문제점을 지적합니다. 훈련이 너무 느리기 때문에, 연구자들은 종로 종종 자신의 아이디어를 한두 개의 게임에서만 테스트하곤 합니다. 이것은 마치 요리사가 새로운 레시피를 오직 한 종류의 파스타에 대해서만 테스트하는 것과 같습니다. 만약 그 레시피가 스파게티에는 작동하지만 펜네에는 실패한다면, 요리사는 알 수 없게 됩니다.

JaxMARL은 매우 빠르기 때문에, 연구자들은 이제 단 하나의 게임을 테스트하는 데 걸렸던 시간 동안 수십 가지의 서로 다른 게임에서 자신의 아이디어를 테스트할 수 있습니다. 이는 AI가 단순히 특정 게임 하나를 "암기"하는 것이 아니라, 실제로 똑똑하고 범용적인지를 확인하는 데 도움을 줍니다.

요 요약

요약하자면, JaxMARL은 연구자들이 현대적인 그래픽 카드의 막대한 힘을 사용하여 AI 에이전트 팀을 훈련시킬 수 있도록 해주는 무료 오픈 소스 툴박스입니다. 이는 느리고 무거운 비디오 게임 엔진을 가볍고 번개처럼 빠른 버전으로 대체하여, 과학자들이 컴퓨터가 따라잡기를 기다리며 정체되지 않고 AI가 협력하고, 경쟁하며, 복잡한 문제를 해결하는 더 나은 방법을 찾아낼 수 있도록 돕습니다.

기술 요약: JaxMARL

문제 정의

강화 학습(RL) 연구는 문제를 정의하고 알고리즘 비교를 가능하게 하는 벤치마크에 크게 의존합니다. 그러나 전통적인 RL 환경은 일반적으로 CPU에서 실행되며, 이는 확장성을 저해하는 중요한 병목 현상을 생성합니다. 이러한 한계는 다중 에이전트 강화 학습(MARL)에서 더욱 두드러지는데, 여러 에이전트를 동시에 훈련시키고, 비정상성(non-stationarity)을 처리하며, 분산된 부분 관측 가능성(decentralized partial observability)을 관리해야 하는 계산적 부담이 가중되기 때문입니다. 결과적으로 실험 실행은 느리고 비용이 많이 들며 범위가 제한적입니다.

나아가, 이 분야는 "평가 위기"에 직면해 있습니다. 최근의 분석에 따르면 많은 MARL 논문들이 단 하나 또는 두 개의 환경에서만 테스트를 수행하며, 이는 잠재적인 과적합, 편향된 비교, 그리고 불분명한 발전 지표로 이어집니다. 하드웨어 가속(예: GPU)은 대규모 병렬 처리를 가능하게 함으로써 해결책을 제시하지만, 기존 구현체들은 종종 비-Python 코드베이스(예: C++)를 요구하여 Python이 표준인 광범위한 연구 커뮤니티의 접근성을 떨어뜨립니다. 비록 JAX가 단일 에이전트 RL(예: PureJaxRL)을 가속화하는 데 성공했지만, 다중 에이전트 시나리오에는 널리 채택되지 않았습니다.

방법론

저자들은 GPU 가속이 가능한 대규모 병렬 MARL 환경 및 알고리즘을 제공하기 위해 JAX 프레임워크를 활용하는 오픈 소스 Python 기반 라이브러리인 JaxMARL을 제시합니다.

핵심 기술 접근 방식

JAX 통합: 이 라이브러리는 JAX의 vmap 연산을 사용하여 입력 차원에 따라 함수를 벡터화함으로써, 단일 GPU에서 수천 개의 환경 인스턴스가 병렬로 실행될 수 있도록 합니다. 또한 성능 최적화를 위해 Just-In-Time (JIT) 컴파일을 채택합니다.
공동 배치(Co-location): 환경 역학(dynamics)과 에이전트 훈련 파이프라인을 모두 JAX로 구현함으로써, CPU와 GPU 메모리 간의 데이터 복사로 발생하는 지연 시간을 제거합니다.
인터페이스 설계: PettingZoo와 Gymnax에서 영감을 받은 단순한 인터페이스를 설계하였으며, 명확한 단일 파일 구현을 제공하는 "CleanRL" 철학을 준수합니다.

환경 제품군

JaxMARL은 연속 및 이산 역학, 가변적인 에이전트 수, 그리고 협력적, 경쟁적, 혼합 인센티브 설정을 아우르는 다양한 환경 세트를 포함합니다:

새로운 환경:
- SMAX: StarCraft Multi-Agent Challenge (SMAC/SMACv2)의 JAX 기반 근사 재구현체입니다. 무거운 StarCraft II 게임 엔진에 의존하는 원본과 달리, SMAX는 경량 역학을 사용하며 게임 엔진의 필요성을 제거하고 경쟁적 셀프 플레이와 높은 커스터마이징을 지원합니다.
- STORM (Spatial-Temporal Representations of Matrix Games): 매트릭스 게임(예: 죄수의 딜레마)을 그리드 월드 시나리오로 확장한 것입니다. 사회적 딜레마와 상대방 형성(opponent shaping)을 연구하기 위해 시간적으로 확장된 행동, 부분 관측 가능성 및 확률성을 도입합니다.
기존 환경: Hanabi, Overcooked, Multi-Agent Brax (MABrax), Multi-Agent Particle Environment (MPE), Coin Game, 그리고 Switch Riddle의 JAX 재구현체입니다.

알고리즘 구현

라이브러리는 JAX의 JIT 및 vmap 기능과 호환되는 인기 있는 MARL 베이스라인들을 구현합니다:

정책 경사(Policy Gradient): 파라미터 공유를 지원하는 피드포워드 및 RNN 정책을 포함하는 Independent PPO (IPPO)와 Multi-Agent PPO (MAPPO).
Q-Learning: PyMARL 코드베이스와의 일관성을 보장하도록 구현된 Independent Q-Learning (IQL), Value Decomposition Networks (VDN), 그리고 QMIX.

주요 기여

최초의 오픈 소스 JAX MARL 라이브러리: JaxMARL은 GPU 기반의 효율성과 광범위한 인기 MARL 환경 및 베이스라인 알고리즘 지원을 결합한 최초의 라이브러리입니다.
새로운 벤치마크 제품군: SMAX(SMAC의 더 빠르고 유연한 대안)와 STORM(매트릭스 게임의 그리드 월드 확장)의 도입은 MARL 연구를 위한 가용 도구를 확장합니다.
알고리즘 구현: IPPO, MAPPO, QMIX, VDN, IQL의 JAX 네이티브 버전을 제공합니다.
포괄적인 벤치마킹: 저자들은 기존 CPU 기반 저장소(예: PyMARL, SMAC)와 비교하여 구현의 속도 및 정확성을 검증하는 철저한 벤치마크를 제공합니다.
평가 권장 사항: 다양한 연구 설정(CTDE, Zero-Shot Coordination, General-Sum, Continuous Actions)에 대한 표준 최소 환경 세트를 제안하며, 편향을 완화하기 위해 사분위수 평균(IQM) 통계치를 사용하는 대규모 평가 스크립트를 제공합니다.

결과

저자들은 상당한 성능 향상과 정확성 동등성을 입증했습니다:

속도:
- 실제 시간(wall-clock time) 기준으로, JAX 기반 훈련 파이프라인은 단일 훈련 실행 시 기존 방식(예: Py-MARL)보다 약 14배 빠릅니다.
- 여러 훈련 실행을 벡터화할 경우, 속도 향상은 최대 12,500배에 달합니다 (예: 1,024개의 시드를 약 198초 만에 훈련, PyMARL의 단일 실행이 한 시간 이상 걸리는 것과 대조적임).
- SMAX 훈련은 단일 NVIDIA 2080에서 벡터화되었을 때 동등한 SMAC 구현체보다 40,000배 빠른 것으로 보고되었습니다.
- 무작위 행동에 대한 환경 스텝 속도(초당 스텝 수)는 CPU 대응 모델보다 수십 배 높습니다 (예: MPE Simple Spread는 1만 개의 환경에서 초당 $4.0 \times 10^7$ 스텝에 도달함).
성능:
- 다중 환경 비교(SMAX, Overcooked, MPE)에서 PPO 기반 방법(IPPO/MAPPO)이 Q-Learning 방법보다 명확한 우위를 보였습니다.
- Hanabi 및 MPE에서의 IPPO 및 MAPPO와 같은 알고리즘의 훈련 곡선과 최종 점수는 이전 연구(예: PyTorch/C++ 구현체)의 결과와 일치하거나 이를 상회하여, JAX 구현의 정확성을 검증했습니다.
정확성: 저자들은 환경 전이(transitions)와 보상이 엄격한 허용 오차( $1 \times 10^{-4}$ ) 내에서 원본 구현(예: Overcooked-AI, PettingZoo, Hanabi Learning Environment)과 일치함을 확인했습니다.

의의

본 논문은 JaxMARL이 MARL 연구의 결정적인 병목 현상을 해결한다고 주장합니다:

평가 위기 완화: 훈련 및 평가의 계산 비용을 획기적으로 낮춤으로써, 라이브러리는 연구자들이 몇 개의 환경이 아닌 광범위한 환경 세트에서 알고리즘을 테스트할 수 있게 하여 과적합 위험을 줄이고 더 견고한 결론을 도출할 수 있게 합니다.
하드웨어 가속의 민주화: Python 기반의 사용하기 쉬운 인터페이스를 제공함으로써, 복잡한 C++ 코드를 작성할 수 없거나 작성하고 싶지 않은 연구자들도 GPU 가속 MARL을 사용할 수 있게 합니다.
새로운 연구 방향 제시: 라이브러리의 속도와 유연성(특히 SMAX를 통해)은 이전에 계산적으로 불가능했던 셀프 플레이, 메타 러닝, 대규모 인구 기반 훈련(population-based training)과 같은 연구 분야의 잠재력을 열어줍니다.

저자들은 가변적인 에이전트 수를 효율적으로 처리하는 문제나 JAX에서의 오프 폴리시(off-policy) 방법의 특정 한계와 같은 과제가 여전히 남아 있지만, JaxMARL이 MARL 연구를 더 빠르고, 철저하며, 접근 가능하게 만드는 중요한 진전이라고 결론짓습니다.