alterlab-stable-baselines3

Stable Baselines3

概要

Stable Baselines3（SB3）は、強化学習アルゴリズムの信頼性の高い実装を提供する、PyTorch ベースのライブラリです。このスキルでは、SB3 の統一 API を使用して、RL エージェントのトレーニング、カスタム環境の作成、コールバックの実装、トレーニングワークフローの最適化を行うための包括的なガイダンスを提供します。

このスキルを使用する場面

ユーザーが次のことを希望する場合に、このスキルを使用します。

• Gymnasium 環境で、単一エージェントの RL ポリシー（PPO、A2C、SAC、TD3、DDPG、DQN、HER）をトレーニングして評価する。
• 研究用シミュレーターをカスタムの gymnasium.Env としてラップし、check_env で検証する。
• コールバックを通じて、評価、チェックポイント保存、早期停止、カスタムロギングを追加する。
• ベクトル化環境を使用してトレーニングを高速化する。または、トレーニング済みエージェントを再現可能な方法で保存、読み込み、レポートする。

該当しないケース

シナリオ	代わりに使用
1 秒あたり数百万ステップの処理、ネイティブなマルチエージェント環境、または PufferEnv／Atari 規模のスループット	alterlab-pufferlib
教師ありディープラーニングのトレーニングループ（LightningModule、マルチ GPU 戦略）	alterlab-pytorch-lightning
逐次的な意思決定プロセスを伴わない、パラメーターのブラックボックス最適化または多目的最適化	alterlab-pymoo
ポリシーを学習せずにキューやリソース競合をシミュレーションする	alterlab-simpy

主な機能

1. RL エージェントのトレーニング

基本的なトレーニングパターン：

import gymnasium as gym
from stable_baselines3 import PPO

# Create environment
env = gym.make("CartPole-v1")

# Initialize agent
model = PPO("MlpPolicy", env, verbose=1)

# Train the agent
model.learn(total_timesteps=10000)

# Save the model
model.save("ppo_cartpole")

# Load the model (without prior instantiation)
model = PPO.load("ppo_cartpole", env=env)

重要な注意事項：

• total_timesteps は下限です。バッチ収集のため、実際のトレーニング量がこれを超える場合があります。
• model.load() は既存のインスタンスではなく、静的メソッドとして使用します。
• スペースを節約するため、リプレイバッファはモデルとともに保存されません。

アルゴリズムの選択： アルゴリズムの特性と選択方法の詳細については、references/algorithms.md を使用してください。簡易リファレンス：

• PPO/A2C：汎用的で、すべての種類のアクション空間に対応し、マルチプロセッシングに適しています。
• SAC/TD3：連続制御向けで、オフポリシー方式によりサンプル効率に優れています。
• DQN：離散アクション向けのオフポリシー方式です。
• HER：目標条件付きタスク向けです。

ベストプラクティスを含むトレーニングテンプレート全体については、scripts/train_rl_agent.py を参照してください。

2. カスタム環境

要件： カスタム環境は gymnasium.Env を継承し、次を実装する必要があります。

• __init__()：action_space と observation_space を定義します。
• reset(seed, options)：初期観測値と info 辞書を返します。
• step(action)：観測値、報酬、terminated、truncated、info を返します。
• render()：可視化します（任意）。
• close()：リソースを解放します。

主な制約：

• 画像の観測値は、範囲 [0, 255] の np.uint8 である必要があります。
• 可能な場合は、チャンネルファースト形式（チャンネル、高さ、幅）を使用します。
• SB3 は画像を 255 で割って自動的に正規化します。
• 事前に正規化済みの場合は、policy_kwargs で normalize_images=False を設定します。
• SB3 は start!=0 の Discrete または MultiDiscrete 空間に対応していません。

検証：

from stable_baselines3.common.env_checker import check_env

check_env(env, warn=True)

カスタム環境のテンプレート全体については、scripts/custom_env_template.py を参照してください。詳細なガイダンスについては、references/custom_environments.md を参照してください。

3. ベクトル化環境

目的： ベクトル化環境は、複数の環境インスタンスを並列に実行してトレーニングを高速化し、特定のラッパー（フレームスタッキング、正規化）を利用できるようにします。

種類：

• DummyVecEnv：現在のプロセス内で順次実行します（軽量な環境向け）。
• SubprocVecEnv：複数のプロセス間で並列実行します（計算負荷の高い環境向け）。

クイックセットアップ：

from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.vec_env import SubprocVecEnv

# Create 4 parallel environments
env = make_vec_env("CartPole-v1", n_envs=4, vec_env_cls=SubprocVecEnv)

model = PPO("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=25000)

SubprocVecEnv を使用する場合、呼び出し元を if __name__ == "__main__": で囲んでください（macOS／Windows の spawn では必須です）。

オフポリシー最適化： オフポリシーアルゴリズム（SAC、TD3、DQN）で複数の環境を使用する場合は、gradient_steps=-1 を設定して環境ステップごとに1回の勾配更新を行い、実時間とサンプル効率のバランスを取ります。

API の違い：

• reset() は観測値のみを返します（情報は vec_env.reset_infos で取得できます）。
• step() は5要素のタプルではなく、4要素のタプル (obs, rewards, dones, infos) を返します。
• エピソード終了後、環境は自動的にリセットされます。
• 終了時の観測値は infos[env_idx]["terminal_observation"] で取得できます。

ラッパーと高度な使い方について詳しくは、references/vectorized_envs.md を参照してください。

4. 監視と制御のためのコールバック

目的： コールバックを使うと、コアアルゴリズムを変更することなく、メトリクスの監視、チェックポイントの保存、早期停止の実装、カスタムのトレーニングロジックの追加ができます。

一般的なコールバック：

• EvalCallback：定期的に評価し、最良のモデルを保存します。
• CheckpointCallback：一定間隔でモデルのチェックポイントを保存します。
• StopTrainingOnRewardThreshold：目標報酬に達した時点で停止します。
• ProgressBarCallback：時間情報とともにトレーニングの進捗を表示します。

カスタムコールバックの構造：

from stable_baselines3.common.callbacks import BaseCallback

class CustomCallback(BaseCallback):
    def _on_training_start(self):
        # Called before first rollout
        pass

    def _on_step(self):
        # Called after each environment step
        # Return False to stop training
        return True

    def _on_rollout_end(self):
        # Called at end of rollout
        pass

利用可能な属性：

• self.model：RL アルゴリズムのインスタンス
• self.num_timesteps：環境ステップの合計数
• self.training_env：トレーニング環境

コールバックの連結：

from stable_baselines3.common.callbacks import CallbackList

callback = CallbackList([eval_callback, checkpoint_callback, custom_callback])
model.learn(total_timesteps=10000, callback=callback)

コールバックの包括的なドキュメントについては、references/callbacks.md を参照してください。

5. モデルの永続化と確認

保存と読み込み：

# Save model
model.save("model_name")

# Save normalization statistics (if using VecNormalize)
vec_env.save("vec_normalize.pkl")

# Load model
model = PPO.load("model_name", env=env)

# Load normalization statistics
vec_env = VecNormalize.load("vec_normalize.pkl", vec_env)

パラメーターへのアクセス：

# Get parameters
params = model.get_parameters()

# Set parameters
model.set_parameters(params)

# Access PyTorch state dict
state_dict = model.policy.state_dict()

6. 評価と記録

評価：

from stable_baselines3.common.evaluation import evaluate_policy

mean_reward, std_reward = evaluate_policy(
    model,
    env,
    n_eval_episodes=10,
    deterministic=True
)

動画の記録（ラップ済み環境は render_mode="rgb_array" を指定して作成する必要があります。また、moviepy がインストールされている必要があります。例：uv pip install "gymnasium[other]"）：

from stable_baselines3.common.vec_env import VecVideoRecorder

# Wrap environment with video recorder
env = VecVideoRecorder(
    env,
    "videos/",
    record_video_trigger=lambda x: x % 2000 == 0,
    video_length=200
)

評価と記録の完全なテンプレートについては、scripts/evaluate_agent.py を参照してください。

7. 高度な機能

学習率スケジュール：

def linear_schedule(initial_value):
    def func(progress_remaining):
        # progress_remaining goes from 1 to 0
        return progress_remaining * initial_value
    return func

model = PPO("MlpPolicy", env, learning_rate=linear_schedule(0.001))

# Built-in equivalent (SB3 >= 2.7; replaces the deprecated get_linear_fn/constant_fn/get_schedule_fn)
from stable_baselines3.common.utils import LinearSchedule
model = PPO("MlpPolicy", env, learning_rate=LinearSchedule(1e-3, 0.0, 1.0))

N ステップリターン（オフポリシー、SB3 >= 2.7）： SAC、TD3、DDPG、DQN は n_steps= を受け付け、N ステップリターンからブートストラップできます（NStepReplayBuffer によって実現）。

マルチ入力ポリシー（Dict 観測）：

model = PPO("MultiInputPolicy", env, verbose=1)

観測値が辞書の場合（画像とセンサーデータを組み合わせる場合など）に使用します。

後知恵経験再生：

from stable_baselines3 import SAC, HerReplayBuffer

model = SAC(
    "MultiInputPolicy",
    env,
    replay_buffer_class=HerReplayBuffer,
    replay_buffer_kwargs=dict(
        n_sampled_goal=4,
        goal_selection_strategy="future",
    ),
)

TensorBoard 連携：

model = PPO("MlpPolicy", env, tensorboard_log="./tensorboard/")
model.learn(total_timesteps=10000)

ワークフローのガイダンス

新しい RL プロジェクトの開始：

1. 問題を定義する：観測空間、行動空間、報酬の構造を特定します。
2. アルゴリズムを選択する：選定の指針として references/algorithms.md を使用します。
3. 環境を作成／適応する：必要に応じて scripts/custom_env_template.py を使用します。
4. 環境を検証する：トレーニング前に必ず check_env() を実行します。
5. トレーニングを設定する：開始用テンプレートとして scripts/train_rl_agent.py を使用します。
6. 監視を追加する：評価とチェックポイント保存のためのコールバックを実装します。
7. パフォーマンスを最適化する：高速化のためにベクトル化環境を検討します。
8. 評価して反復する：評価には scripts/evaluate_agent.py を使用します。

よくある問題：

• メモリーエラー：オフポリシーアルゴリズムの buffer_size を小さくするか、並列環境の数を減らします。
• トレーニングが遅い：並列環境には SubprocVecEnv の使用を検討します。
• トレーニングが不安定：別のアルゴリズムを試す、ハイパーパラメーターを調整する、または報酬のスケーリングを確認します。
• インポートエラー：stable_baselines3 がインストールされていることを確認します：uv pip install "stable-baselines3[extra]"

リソース

scripts/

• train_rl_agent.py：ベストプラクティスを取り入れた、完全なトレーニングスクリプトのテンプレート
• evaluate_agent.py：エージェントの評価と動画記録のテンプレート
• custom_env_template.py：カスタム Gym 環境のテンプレート

references/

• algorithms.md：アルゴリズムの詳細な比較と選択ガイド
• custom_environments.md：カスタム環境の包括的な作成ガイド
• callbacks.md：コールバックシステムの完全なリファレンス
• vectorized_envs.md：ベクトル化環境の使用方法とラッパー

インストール

Python 3.10 以降と PyTorch >= 2.8 が必要です（SB3 2.9 では下限が 2.3 から引き上げられました）。[extra] オプションを指定すると、TensorBoard、OpenCV、ale-py（Atari）、pygame-ce、tqdm/rich（プログレスバー）が追加されます。また、2.9 以降では pandas と matplotlib も追加されます。これらはコア依存関係ではなくなり、results_plotter／load_results に必要です。"PongNoFrameskip-v4" や "ALE/Pong-v5" などの Atari ID は、ale_py をインポートすると登録されます（import ale_py; gym.register_envs(ale_py)）。

# Basic installation
uv pip install stable-baselines3

# With extra dependencies (Tensorboard, Atari, etc.)
uv pip install "stable-baselines3[extra]"

AlterLab Academic Skills スイートの一部です。