RobertTLange/gymnax

해결하는 문제

Gymnax는 느리고 CPU 기반의 강화학습(RL) 환경 처리의 성능 저하 문제를 해결합니다. 연구자들은 가속기(GPU/TPU)에서 환경 전이와 정책을 실행하여 고처리량 실험과 대규모 벡터화를 달성할 수 있습니다.

작동 방식

Gymnax는 JAX를 사용하여 고전적인 RL 환경(예: OpenAI Gym의 환경)을 재구현합니다. 이를 통해 jit(즉시 컴파일), vmap(벡터화 매핑), pmap(병렬 매핑)과 같은 JAX의 기본 기능을 환경 API에 도입합니다. 이로 인해 환경과 정책을 동일한 가속기에서 실행할 수 있어, 빠른 병렬 롤아웃과 난수 시드 및 하이퍼파라미터와 같은 환경 설정에 대한 함수형 제어가 가능해집니다.

대상 사용자

메타-RL, 진화적 최적화, 또는 Anakin 하위 아키텍처와 같은 분산 아키텍처를 다루는 강화학습 연구자 및 개발자에게 적합합니다.

주요 특징

• 가속된 환경: 고전적 제어, bsuite, MinAtar, 다양한 메타-RL 작업의 구현 포함.
• JAX 통합: jit, vmap, pmap에 대한 완전한 지원으로 난수 키 또는 환경 파라미터를 기반으로 대규모 벡터화 가능.
• 미분 가능한 전이: 연속 제어 환경(예: Pendulum, MountainCarContinuous)에서 JAX 기울기의 선택적 지원.
• 빠른 롤아웃: lax.scan을 지원하여 전체 에피소드 루프를 하나의 가속기 연산으로 컴파일 가능.
• 자동 리셋: 종료 또는 단절된 단계 후 자동 리셋 기능 구현.