minimax: Efficient Baselines for Autocurricula in JAX

要約

教師なし環境設計 (UED) は、目に見えない環境にゼロショットで移行できるように堅牢な意思決定エージェントをトレーニングするための自動カリキュラム学習の形式です。
このようなオートカリキュラムは、RL コミュニティから大きな関心を集めています。
ただし、CPU ロールアウトと GPU モデルの更新に基づく UED の実験には、多くの場合、数週間のトレーニングが必要でした。
このコンピューティング要件は、この分野の急速なイノベーションにとって大きな障害となります。
この作業では、高速化されたハードウェアでの UED トレーニング用の Minimax ライブラリを導入します。
JAX を使用して完全にテンソル化された環境とオートカリキュラム アルゴリズムを実装することで、minimax ではトレーニング ループ全体をハードウェア アクセラレーション用にコンパイルできるようになります。
迅速な実験のためのペトリ皿を提供するために、minimax には、手続き的に生成された環境でオートカリキュラムを実行するための再利用可能な抽象化に加えて、MiniGrid に基づくテンソル化されたグリッド ワールドが含まれています。
これらのコンポーネントを使用して、minimax は新しい並列化バリアントを含む強力な UED ベースラインを提供し、同じバッチ サイズでトレーニングする場合、以前の実装と比較して所要時間で 120$\times$ 以上の高速化を達成します。
Minimax ライブラリは、Apache 2.0 ライセンスに基づいて https://github.com/facebookresearch/minimax で入手できます。

要約(オリジナル)

Unsupervised environment design (UED) is a form of automatic curriculum learning for training robust decision-making agents to zero-shot transfer into unseen environments. Such autocurricula have received much interest from the RL community. However, UED experiments, based on CPU rollouts and GPU model updates, have often required several weeks of training. This compute requirement is a major obstacle to rapid innovation for the field. This work introduces the minimax library for UED training on accelerated hardware. Using JAX to implement fully-tensorized environments and autocurriculum algorithms, minimax allows the entire training loop to be compiled for hardware acceleration. To provide a petri dish for rapid experimentation, minimax includes a tensorized grid-world based on MiniGrid, in addition to reusable abstractions for conducting autocurricula in procedurally-generated environments. With these components, minimax provides strong UED baselines, including new parallelized variants, which achieve over 120$\times$ speedups in wall time compared to previous implementations when training with equal batch sizes. The minimax library is available under the Apache 2.0 license at https://github.com/facebookresearch/minimax.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google