RobertTLange/gymnax

What it solves

gymnax는 CPU 기반 강화 학습(RL) 환경의 느린 속도라는 병목 현상을 해결합니다. JAX로 클래식 RL 환경을 재구현함으로써, 연구자들이 환경과 정책을 동일한 가속기(GPU/TPU)에서 실행할 수 있게 하여 CPU와 GPU 사이의 통신 오버헤드를 제거하고 대규모 벡터화를 가능하게 합니다.

How it works

클래식 OpenAI Gym API를 JAX 생태계로 가져옵니다. JAX로 작성되었기 때문에 환경은 완전히 기능적이며 jit (just-in-time compilation), vmap (vectorized mapping), pmap (parallel mapping)과 같은 JAX Primitives을 지원합니다. 이를 통해 수천 개의 환경 인스턴스에서 동시에 고처리량 rollout을 수행할 수 있으며, lax.scan을 사용하여 전체 에피소드 루프를 단일 가속기 연산으로 컴파일할 수 있습니다.

How it works

클래식 OpenAI Gym API를 JAX 생태계로 가져| (Note: The user requested translation, translation logic follows below)

Who it’s for

고처리량 실험이 필요한 RL 연구자 및 개발자, 다양한 환경 설정이 병렬화되어야 하는 meta-RL을 연구하는 분들, 그리고 진화 최적화 전략을 사용하는 사용자들을 위해 설계되었습니다.

Highlights

• High-Throughput Vectorization: vmap and pmap을 활용하여 대규모 병렬 rollout을 수행합니다.
• Accelerator-Native: 환경과 정책을 모두 GPU/TPU에서 실행하여 Anakin sub-architecture와 같은 아키텍처를 지원합니다.
• Broad Environment Support: 환경 지원 범위가가 넓습니다. 클래식 컨트롤, bsuite, MinAtar, 그리고 다양한 meta-RL 태스크를 지원합니다.
• Functional Control: 난수 시드와 하이퍼파라미터를 명시적으로 제어하여 가속화된 병렬화를 가능하게 합니다.
• Visualization Tools: 에이전트의 동작을 GIF 애니메이션로 생성하는 Visualizer를 포함합니다.