araffin/sbx

何を解決するか

SBX は Jax フレームワークを使用して Stable-Baselines3 (SB3) のプロトタイプ実装です。パフォーマンス最適化された信頼性の高い強化学習 (RL) アルゴリズムのセットを提供し、SB3 API と互換性があるため、PyTorch ベースの SB3 と Jax ベースの SBX の間で高レベルのコード構造を変更せずに切り替えることができます。

どう動くか

SB3 アルゴリズムを Jax で再実装することで、Jax のハードウェア加速と関数型プログラミングモデルの利点を活かしています。PPO、DQN、SAC、TD3、TQC などの主要な RL アルゴリズムを実装しており、DroQ や SimBa といった特別な設定もサポートしています。SB3 と API 互換性を維持しているため、RL Zoo などの既存ツールに統合して学習や評価が可能です。

対象ユーザー

Jax のパフォーマンスの利点を活かしつつ、Stable-Baselines3 と同じ使いやすさの API を維持したい強化学習の研究者や開発者。

特徴

• SB3 API 互換性: RL Zoo やその他の SB3 互換ツールとシームレスに統合可能。
• 詳細なアルゴリズムセット: Soft Actor-Critic (SAC)、Proximal Policy Optimization (PPO)、Deep Q Network (DQN) など幅広いアルゴリズムを含む。
• Jax ベース: 計算効率の向上を実現する Jax で構築。
• 高度な RL 技術: オフポリシー アルゴリズムのパラメータリセットや、SimBa などの特別なアーキテクチャをサポート。