pufferlib

PufferLib - 高性能強化学習

概要

PufferLib は、高速な並列環境シミュレーションと学習のために設計された高性能強化学習ライブラリです。最適化されたベクトル化、ネイティブのマルチエージェントサポート、効率的な PPO 実装（PuffeRL）により、毎秒数百万ステップの学習を実現します。このライブラリは 20 以上の環境からなる Ocean スイートと、Gymnasium、PettingZoo、および特化した RL フレームワークとのシームレスな統合を提供します。

このスキルを使用するタイミング

次のような場合にこのスキルを使用します:

• 任意の環境（単一エージェントまたはマルチエージェント）で PPO を使って RL エージェントを学習させるとき
• PufferEnv API を使って カスタム環境を作成するとき
• 並列環境シミュレーション（ベクトル化）の パフォーマンスを最適化するとき
• Gymnasium、PettingZoo、Atari、Procgen などの既存環境を 統合するとき
• CNN、LSTM、またはカスタムアーキテクチャで ポリシーを開発するとき
• より速い実験のために RL を毎秒数百万ステップまで スケールさせるとき
• ネイティブのマルチエージェント環境サポートを使って マルチエージェント RL を行うとき

主な機能

1. 高性能学習（PuffeRL）

PuffeRL は、毎秒 100 万〜400 万ステップを達成する、PufferLib の最適化された PPO+LSTM 学習アルゴリズムです。

クイックスタートの学習:

# CLI training
puffer train procgen-coinrun --train.device cuda --train.learning-rate 3e-4

# Distributed training
torchrun --nproc_per_node=4 train.py

Python の学習ループ:

import pufferlib
from pufferlib import PuffeRL

# Create vectorized environment
env = pufferlib.make('procgen-coinrun', num_envs=256)

# Create trainer
trainer = PuffeRL(
    env=env,
    policy=my_policy,
    device='cuda',
    learning_rate=3e-4,
    batch_size=32768
)

# Training loop
for iteration in range(num_iterations):
    trainer.evaluate()  # Collect rollouts
    trainer.train()     # Train on batch
    trainer.mean_and_log()  # Log results

包括的な学習ガイダンスについては、references/training.md を参照してください。内容は以下のとおりです:

• 学習ワークフローの全体とコマンドラインオプション
• Protein を使ったハイパーパラメータのチューニング
• 分散マルチ GPU / マルチノード学習
• ロガーとの統合（Weights & Biases、Neptune）
• チェックポイントと学習の再開
• パフォーマンス最適化のヒント
• カリキュラム学習のパターン

2. 環境開発（PufferEnv）

PufferEnv API を使って、カスタムの高性能環境を作成します。

基本的な環境の構造:

import numpy as np
from pufferlib import PufferEnv

class MyEnvironment(PufferEnv):
    def __init__(self, buf=None):
        super().__init__(buf)

        # Define spaces
        self.observation_space = self.make_space((4,))
        self.action_space = self.make_discrete(4)

        self.reset()

    def reset(self):
        # Reset state and return initial observation
        return np.zeros(4, dtype=np.float32)

    def step(self, action):
        # Execute action, compute reward, check done
        obs = self._get_observation()
        reward = self._compute_reward()
        done = self._is_done()
        info = {}

        return obs, reward, done, info

テンプレートスクリプトの利用: scripts/env_template.py には、単一エージェントおよびマルチエージェントの環境テンプレートが完全な形で用意されており、以下の例が含まれています:

• さまざまな観測空間の種類（ベクトル、画像、辞書）
• 行動空間のバリエーション（離散、連続、マルチディスクリート）
• マルチエージェント環境の構造
• テストユーティリティ

環境開発の詳細については、references/environments.md を参照してください。内容は以下のとおりです:

• PufferEnv API の詳細とインプレース操作のパターン
• 観測空間と行動空間の定義
• マルチエージェント環境の作成
• Ocean スイート（20 以上の既製環境）
• パフォーマンス最適化（Python から C へのワークフロー）
• 環境ラッパーとベストプラクティス
• デバッグと検証の手法

3. ベクトル化とパフォーマンス

最適化された並列シミュレーションで最大のスループットを達成します。

ベクトル化のセットアップ:

import pufferlib

# Automatic vectorization
env = pufferlib.make('environment_name', num_envs=256, num_workers=8)

# Performance benchmarks:
# - Pure Python envs: 100k-500k SPS
# - C-based envs: 100M+ SPS
# - With training: 400k-4M total SPS

主な最適化:

• ゼロコピーの観測受け渡しのための共有メモリバッファ
• パイプやキューの代わりとなるビジーウェイトフラグ
• 非同期返却のための余剰環境
• ワーカーごとの複数環境

ベクトル化の最適化については、references/vectorization.md を参照してください。内容は以下のとおりです:

• アーキテクチャとパフォーマンス特性
• ワーカーとバッチサイズの設定
• シリアル、マルチプロセシング、非同期モードの比較
• 共有メモリとゼロコピーのパターン
• 大規模向けの階層的ベクトル化
• マルチエージェントのベクトル化戦略
• パフォーマンスのプロファイリングとトラブルシューティング

4. ポリシー開発

標準の PyTorch モジュールとして、任意のユーティリティを使ってポリシーを構築します。

基本的なポリシーの構造:

import torch.nn as nn
from pufferlib.pytorch import layer_init

class Policy(nn.Module):
    def __init__(self, observation_space, action_space):
        super().__init__()

        # Encoder
        self.encoder = nn.Sequential(
            layer_init(nn.Linear(obs_dim, 256)),
            nn.ReLU(),
            layer_init(nn.Linear(256, 256)),
            nn.ReLU()
        )

        # Actor and critic heads
        self.actor = layer_init(nn.Linear(256, num_actions), std=0.01)
        self.critic = layer_init(nn.Linear(256, 1), std=1.0)

    def forward(self, observations):
        features = self.encoder(observations)
        return self.actor(features), self.critic(features)

ポリシー開発の詳細については、references/policies.md を参照してください。内容は以下のとおりです:

• 画像観測向けの CNN ポリシー
• 最適化された LSTM を使った再帰ポリシー（推論が 3 倍高速）
• 複雑な観測のためのマルチ入力ポリシー
• 連続行動ポリシー
• マルチエージェントポリシー（パラメータの共有と独立）
• 高度なアーキテクチャ（アテンション、残差）
• 観測の正規化と勾配クリッピング
• ポリシーのデバッグとテスト

5. 環境の統合

人気の RL フレームワークの環境をシームレスに統合します。

Gymnasium 統合:

import gymnasium as gym
import pufferlib

# Wrap Gymnasium environment
gym_env = gym.make('CartPole-v1')
env = pufferlib.emulate(gym_env, num_envs=256)

# Or use make directly
env = pufferlib.make('gym-CartPole-v1', num_envs=256)

PettingZoo マルチエージェント:

# Multi-agent environment
env = pufferlib.make('pettingzoo-knights-archers-zombies', num_envs=128)

対応フレームワーク:

• Gymnasium / OpenAI Gym
• PettingZoo（parallel および AEC）
• Atari（ALE）
• Procgen
• NetHack / MiniHack
• Minigrid
• Neural MMO
• Crafter
• GPUDrive
• MicroRTS
• Griddly
• その他多数...

統合の詳細については、以下の内容について references/integration.md を参照してください:

• 各フレームワークの完全な統合例
• カスタムラッパー（観測、報酬、フレームスタッキング、アクションリピート）
• 空間のフラット化と復元
• 環境の登録
• 互換性パターン
• パフォーマンスに関する考慮事項
• 統合のデバッグ

クイックスタートワークフロー

既存の環境でトレーニングする場合

1. Ocean スイートまたは互換フレームワークから環境を選択します
2. 出発点として scripts/train_template.py を使用します
3. タスクに合わせてハイパーパラメータを設定します
4. CLI または Python スクリプトでトレーニングを実行します
5. Weights & Biases または Neptune で監視します
6. 最適化については references/training.md を参照してください

カスタム環境を作成する場合

1. scripts/env_template.py から始めます
2. 観測空間とアクション空間を定義します
3. reset() および step() メソッドを実装します
4. ローカルで環境をテストします
5. pufferlib.emulate() または make() でベクトル化します
6. 高度なパターンについては references/environments.md を参照してください
7. 必要に応じて references/vectorization.md で最適化します

ポリシー開発の場合

1. 観測内容に基づいてアーキテクチャを選択します:
   • ベクトル観測 → MLP ポリシー
   • 画像観測 → CNN ポリシー
   • 逐次的なタスク → LSTM ポリシー
   • 複雑な観測 → マルチ入力ポリシー
2. 適切な重み初期化のために layer_init を使用します
3. references/policies.md のパターンに従います
4. 本格的なトレーニングの前に環境でテストします

パフォーマンス最適化の場合

1. 現在のスループット（毎秒のステップ数）をプロファイリングします
2. ベクトル化の設定（num_envs、num_workers）を確認します
3. 環境コードを最適化します（インプレース演算、numpy のベクトル化）
4. 重要な経路については C による実装を検討します
5. 体系的な最適化には references/vectorization.md を使用します

リソース

scripts/

train_template.py - 以下を含む完全なトレーニングスクリプトのテンプレート:

• 環境の作成と設定
• ポリシーの初期化
• ロガーの統合（WandB、Neptune）
• チェックポイント保存を伴うトレーニングループ
• コマンドライン引数の解析
• マルチ GPU 分散トレーニングの設定

env_template.py - 環境実装のテンプレート:

• 単一エージェントの PufferEnv の例（グリッドワールド）
• マルチエージェントの PufferEnv の例（協調ナビゲーション）
• 複数の観測・アクション空間のパターン
• テストユーティリティ

references/

training.md - 包括的なトレーニングガイド:

• トレーニングのワークフローとCLIオプション
• ハイパーパラメータの設定
• 分散トレーニング（マルチGPU、マルチノード）
• モニタリングとロギング
• チェックポイント
• Protein のハイパーパラメータチューニング
• パフォーマンス最適化
• 一般的なトレーニングパターン
• トラブルシューティング

environments.md - 環境開発ガイド:

• PufferEnv API と特性
• 観測空間とアクション空間
• マルチエージェント環境
• Ocean スイートの環境
• カスタム環境の開発ワークフロー
• Python から C への最適化パス
• サードパーティ環境の統合
• ラッパーとベストプラクティス
• デバッグ

vectorization.md - ベクトル化の最適化:

• アーキテクチャと主要な最適化
• ベクトル化モード（シリアル、マルチプロセッシング、非同期）
• ワーカーとバッチの設定
• 共有メモリとゼロコピーパターン
• 高度なベクトル化（階層型、カスタム）
• マルチエージェントのベクトル化
• パフォーマンスのモニタリングとプロファイリング
• トラブルシューティングとベストプラクティス

policies.md - ポリシーアーキテクチャガイド:

• 基本的なポリシー構造
• 画像向けの CNN ポリシー
• 最適化を施した LSTM ポリシー
• マルチインプットポリシー
• 連続アクションポリシー
• マルチエージェントポリシー
• 高度なアーキテクチャ（アテンション、残差）
• 観測の処理と展開（unflattening）
• 初期化と正規化
• デバッグとテスト

integration.md - フレームワーク統合ガイド:

• Gymnasium との統合
• PettingZoo との統合（parallel と AEC）
• サードパーティ環境（Procgen、NetHack、Minigrid など）
• カスタムラッパー（観測、報酬、フレームスタッキングなど）
• 空間の変換と展開（unflattening）
• 環境の登録
• 互換性パターン
• パフォーマンス上の考慮事項
• 統合のデバッグ

成功のためのヒント

1. シンプルに始める: カスタム環境を作成する前に、Ocean 環境または Gymnasium 統合から始めてください

2. 早期にプロファイルする: ボトルネックを特定するため、最初から毎秒のステップ数を計測してください

3. テンプレートを使う: scripts/train_template.py と scripts/env_template.py が確かな出発点を提供します

4. 必要に応じてリファレンスを読む: 各リファレンスファイルは自己完結しており、特定の機能に焦点を当てています

5. 段階的に最適化する: まず Python で実装し、プロファイルを取り、必要に応じて重要な箇所を C で最適化してください

6. ベクトル化を活用する: 高いスループットを達成するうえで、PufferLib のベクトル化が鍵となります

7. トレーニングを監視する: WandB または Neptune を使って実験を追跡し、問題を早期に特定してください

8. 環境をテストする: トレーニングを拡大する前に、環境のロジックを検証してください

9. 既存の環境を確認する: Ocean スイートには 20 以上の既製環境が用意されています

10. 適切な初期化を使う: ポリシーには必ず pufferlib.pytorch の layer_init を使用してください

一般的なユースケース

標準ベンチマークでのトレーニング

# Atari
env = pufferlib.make('atari-pong', num_envs=256)

# Procgen
env = pufferlib.make('procgen-coinrun', num_envs=256)

# Minigrid
env = pufferlib.make('minigrid-empty-8x8', num_envs=256)

マルチエージェント学習

# PettingZoo
env = pufferlib.make('pettingzoo-pistonball', num_envs=128)

# Shared policy for all agents
policy = create_policy(env.observation_space, env.action_space)
trainer = PuffeRL(env=env, policy=policy)

カスタムタスクの開発

# Create custom environment
class MyTask(PufferEnv):
    # ... implement environment ...

# Vectorize and train
env = pufferlib.emulate(MyTask, num_envs=256)
trainer = PuffeRL(env=env, policy=my_policy)

高性能最適化

# Maximize throughput
env = pufferlib.make(
    'my-env',
    num_envs=1024,      # Large batch
    num_workers=16,     # Many workers
    envs_per_worker=64  # Optimize per worker
)

インストール

pip install pufferlib

ドキュメント

• 公式ドキュメント: https://puffer.ai/docs.html
• GitHub: https://github.com/PufferAI/PufferLib
• Discord: コミュニティによるサポートを利用できます