stable-baselines3

Stable Baselines3

概要

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

主な機能

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): observation, reward, terminated, truncated, info を返す
• render(): 可視化（任意）
• close(): リソースのクリーンアップ

主な制約:

• 画像の観測は範囲 [0, 255] の np.uint8 でなければなりません
• 可能な場合はチャネルファースト形式（channels, height, width）を使用してください
• 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

# 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)

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

API の違い:

• reset() は観測のみを返します（info は 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
)

ビデオ録画:

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))

マルチ入力ポリシー（辞書形式の観測）:

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

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

Hindsight Experience Replay:

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: ベクトル化環境の使用方法とラッパー

インストール

# Basic installation
uv pip install stable-baselines3

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