google-deepmind/mctx

해결하는 문제

심층 신경망과 탐색 알고리즘을 결합할 때 종종 C++과 같은 빠른 컴파일 언어로 구현해야 하며, 연구자들이 사용하거나 수정하기 어려울 수 있습니다. Mctx는 Python에서 사용 가능하면서도 실행 속도와 연구자 친화성 사이의 격차를 메우는 고성능 JAX 네이티브 몬테카를로 트리 탐색(MCTS) 구현을 제공합니다.

작동 방식

Mctx는 JAX를 사용하여 MCTS 알고리즘(AlphaZero, MuZero, Gumbel MuZero 포함)을 구현합니다. JIT 컴파일을 활용하고 병렬로 여러 입력을 처리하여 하드웨어 가속기의 효율을 극대화합니다. 사용자는 루트 상태의 표현 함수나 환경 동역학을 위한 재귀 함수와 같은 학습된 구성 요소를 제공하며, 라이브러리는 이를 사용해 탐색 트리를 구성하고 행동을 제안합니다.

대상 사용자

컴파일된 코드의 성능을 원하면서도 Python 생태계를 벗어나지 않고 싶은 탐색 기반 강화학습 에이전트를 연구하는 AI 연구자를 대상으로 합니다.

주요 특징

• JAX 네이티브: JIT 컴파일을 완전히 지원하여 계산 속도를 크게 향상시킵니다.
• 병렬 탐색: 병렬로 여러 입력을 처리하여 가속기 사용을 최적화합니다.
• 알고리즘 지원: AlphaZero, MuZero, Gumbel MuZero의 구현이 포함되어 있습니다.
• 사용자 정의 가능: 새로운 탐색 기반 아이디어 탐색을 가능하게 하기 위해 매우 유연하게 구성되어 있습니다.