araffin/sbx

解決的問題

SBX 是使用 Jax 框架實現的 Stable-Baselines3 (SB3) 概念驗證版本。它提供一組性能最佳化且可靠的強化學習（RL）演算法，與 SB3 API 相容，讓使用者可以在不改變高階程式碼結構的情況下，無縫切換 PyTorch 基於的 SB3 與 Jax 基於的 SBX。

如何運作

透過將 SB3 演算法重寫為 Jax，此專案利用 Jax 的硬體加速與函數式程式設計模型。實作了多種關鍵的 RL 演算法，包括 PPO、DQN、SAC、TD3 和 TQC，以及 DroQ 和 SimBa 等特殊設定。它維持與 SB3 的 API 相容性，因此可整合至現有的工具（如 RL Zoo）中，用於訓練與評估。

適用對象

希望在保持 Stable-Baselines3 一致使用者友善 API 的同時，獲得 Jax 性能優勢的強化學習研究人員與開發者。

主要亮點

• SB3 API 相容性：可無縫整合至 RL Zoo 及其他 SB3 相容工具中。
• 完整的演算法套件：包含 Soft Actor-Critic (SAC)、Proximal Policy Optimization (PPO)、Deep Q Network (DQN) 等多種演算法。
• Jax 驅動：基於 Jax 建構，實現更高的計算效率。
• 進階強化學習技術：支援離策略演算法的參數重設以及 SimBa 等特殊架構。