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の完全サポートにより、乱数キーまたは環境パラメータ across で大規模なベクトル化が可能。
• 微分可能な遷移: 連続制御環境（例：Pendulum、MountainCarContinuous）でJAX勾配のオプトインサポート。
• 高速ロールアウト: lax.scanをサポートし、エピソード全体のループを1つのアクセラレータ操作にコンパイル可能。
• 自動リセット: 終了または切断ステップ後に自動リセットを実装。