alterlab-pufferlib

PufferLib - 高性能強化学習

概要

PufferLib は、高速な並列環境シミュレーションとトレーニングを目的として設計された、高性能な強化学習ライブラリです。最適化されたベクトル化、ネイティブなマルチエージェント対応、効率的な PPO 実装（PuffeRL）により、毎秒数百万ステップのトレーニングを実現します。このライブラリには、20 以上の環境を含む Ocean スイートと、Gymnasium、PettingZoo、各種 RL フレームワークとのシームレスな統合機能が用意されています。

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

次の場合にこのスキルを使用します：

• PPO を使って、あらゆる環境（単一エージェントまたはマルチエージェント）で RL エージェントをトレーニングする場合
• PufferEnv API を使用してカスタム環境を作成する場合
• 並列環境シミュレーション（ベクトル化）のパフォーマンスを最適化する場合
• Gymnasium、PettingZoo、Atari、Procgen などの既存環境を統合する場合
• CNN、LSTM、またはカスタムアーキテクチャを使ってポリシーを開発する場合
• 実験を高速化するため、RL を毎秒数百万ステップにスケールする場合
• ネイティブなマルチエージェント環境対応を活用して、マルチエージェント RL を行う場合

対象外

シナリオ	代わりに使用
標準的で十分にドキュメント化された PPO／SAC／DQN 実装を使った、手軽な単一エージェントのプロトタイプ作成	alterlab-stable-baselines3
チェックポイント保存やマルチ GPU に対応した、教師あり深層学習のトレーニングループ（分類、回帰）	alterlab-pytorch-lightning
ルールベースのエージェントを使い、学習を行わない社会システムのエージェントベースシミュレーション	alterlab-abm-mesa

主な機能

1. 高性能トレーニング（PuffeRL）

PuffeRL は、高スループットのトレーニング向けに構築された、PufferLib の最適化済み PPO トレーナーです（CleanRL をベースとし、models.LSTMWrapper によるオプションの LSTM に対応）。

推奨方法 — CLI／高水準ヘルパー。 トレーナーを手作業で配線するのではなく、設定（pufferlib/config/ 内の .ini）を使ってトレーニングを実行します：

# CLI: env name resolves to a config in pufferlib/config/ + an Ocean env
puffer train puffer_breakout --train.device cuda --train.learning-rate 0.015
# Resume from a checkpoint (top-level flag, not a [train] key)
puffer train puffer_breakout --load-model-path latest

import pufferlib.pufferl as pufferl
# train(env_name, args=None, vecenv=None, policy=None, logger=None)
pufferl.train('puffer_breakout')

手動ループ。 PuffeRL(config, vecenv, policy, logger=None) — 最初の引数は設定辞書（フラットなキーワード引数ではありません）、環境の引数名は vecenv で、ループは global_step によって進行します。3 つのループメソッドは実際に利用できます：evaluate()、train()、mean_and_log()。

import pufferlib.vector
from pufferlib.pufferl import PuffeRL, load_config

# Native PufferEnv -> default backend=PufferEnv. For wrapped (Gymnasium/
# PettingZoo) envs you MUST pass backend=pufferlib.vector.Multiprocessing.
vecenv = pufferlib.vector.make(MyPufferEnv, num_envs=256)

# load_config returns a nested args dict (sections: 'train', 'vec', 'env', ...)
# with defaults from pufferlib/config/*.ini. PuffeRL takes the 'train' section.
args = load_config('puffer_breakout')
config = {**args['train'], 'env': 'puffer_breakout'}
config['device'] = 'cuda'

trainer = PuffeRL(config, vecenv, my_policy)
while trainer.global_step < config['total_timesteps']:
    trainer.evaluate()      # Collect rollouts
    trainer.train()         # Train on batch
    trainer.mean_and_log()  # Aggregate + log

トレーニングの詳細なガイダンスについては、references/training.md を参照してください：

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

2. 環境開発（PufferEnv）

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

基本的な環境構造：

import numpy as np
import gymnasium
from pufferlib import PufferEnv

class MyEnvironment(PufferEnv):
    def __init__(self, buf=None):
        # Define spaces BEFORE calling super().__init__(buf)
        self.single_observation_space = gymnasium.spaces.Box(
            low=-np.inf, high=np.inf, shape=(4,), dtype=np.float32)
        self.single_action_space = gymnasium.spaces.Discrete(4)
        self.num_agents = 1

        super().__init__(buf)

    def reset(self, seed=None):
        # Reset state and return (observation, info-list)
        obs = self._get_observation()
        return obs, []

    def step(self, action):
        # Execute action, compute reward, check termination/truncation
        obs = self._get_observation()
        rewards = self._compute_reward()
        terminals = self._is_done()
        truncations = self._is_truncated()
        info = []

        return obs, rewards, terminals, truncations, info

テンプレートスクリプトを使用する： scripts/env_template.py には、次の例を含む、単一エージェント環境とマルチエージェント環境の完全なテンプレートが用意されています：

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

環境開発の詳細については、references/environments.md を参照してください：

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

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

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

ベクトル化の設定：

import pufferlib.vector

# Pass an env-constructor callable. Default backend=PufferEnv is native-only;
# for wrapped (Gymnasium/PettingZoo) envs add backend=pufferlib.vector.Multiprocessing.
env = pufferlib.vector.make(env_creator, num_envs=256, num_workers=8)

# Performance benchmarks (PufferLib's published figures; vary by env/hardware):
# - 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.emulation
import pufferlib.vector

# Wrap a Gymnasium env in a GymnasiumPufferEnv, then vectorize.
# Wrapped (non-native) envs require an explicit backend (Serial or Multiprocessing);
# the default backend=PufferEnv is only for native PufferEnvs.
def env_creator():
    return pufferlib.emulation.GymnasiumPufferEnv(
        env_creator=lambda: gym.make('CartPole-v1'))

env = pufferlib.vector.make(
    env_creator, num_envs=256, backend=pufferlib.vector.Multiprocessing)

PettingZoo のマルチエージェント：

import pufferlib.emulation
import pufferlib.vector
from pettingzoo.butterfly import knights_archers_zombies_v10

# Wrap a PettingZoo env in a PettingZooPufferEnv, then vectorize.
def env_creator():
    return pufferlib.emulation.PettingZooPufferEnv(
        env_creator=lambda: knights_archers_zombies_v10.parallel_env())

env = pufferlib.vector.make(
    env_creator, num_envs=128, backend=pufferlib.vector.Multiprocessing)

対応フレームワーク：

• Gymnasium / OpenAI Gym
• PettingZoo（並列および 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.emulation.GymnasiumPufferEnv でラップし、pufferlib.vector.make() でベクトル化します
6. 高度なパターンについては references/environments.md を参照してください
7. 必要に応じて references/vectorization.md を使って最適化します

ポリシーを開発する場合

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

パフォーマンスを最適化する場合

1. 現在のスループット（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 suite の環境
• カスタム環境の開発ワークフロー
• Python から C への最適化パス
• サードパーティ環境の統合
• ラッパーとベストプラクティス
• デバッグ

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

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

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

• 基本的なポリシー構造
• 画像用 CNN ポリシー
• 最適化を施した LSTM ポリシー
• 複数入力ポリシー
• 連続行動ポリシー
• マルチエージェントポリシー
• 高度なアーキテクチャ（アテンション、残差）
• 観測の処理とアンフラット化
• 初期化と正規化
• デバッグとテスト

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

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

成功のためのヒント

1. シンプルに始める：カスタム環境を作成する前に、Ocean の環境または Gymnasium の統合から始めましょう

2. 早い段階でプロファイリングする：最初から 1 秒あたりのステップ数を測定し、ボトルネックを特定しましょう

3. テンプレートを使う：scripts/train_template.py と scripts/env_template.py は、確かな出発点になります

4. 必要に応じてリファレンスを読む：各リファレンスファイルはそれぞれ独立しており、特定の機能に焦点を当てています

5. 段階的に最適化する：まず Python で始めてプロファイリングし、必要に応じて重要な処理パスを C で最適化しましょう

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

7. トレーニングをモニタリングする：WandB または Neptune を使用して実験を追跡し、問題を早期に特定しましょう

8. 環境をテストする：トレーニングを大規模化する前に、環境のロジックを検証しましょう

9. 既存の環境を確認する：Ocean suite には、事前構築済みの環境が 20 以上用意されています

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

一般的なユースケース

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

import pufferlib.vector

# Atari (pass an env-constructor callable)
env = pufferlib.vector.make(make_pong_env, num_envs=256)

# Procgen
env = pufferlib.vector.make(make_coinrun_env, num_envs=256)

# Minigrid
env = pufferlib.vector.make(make_minigrid_env, num_envs=256)

マルチエージェント学習

import pufferlib.vector

# PettingZoo, wrapped via PettingZooPufferEnv (needs an explicit backend)
env = pufferlib.vector.make(
    make_pistonball_env, num_envs=128, backend=pufferlib.vector.Multiprocessing)

# One shared policy serves all agents (single_observation_space / single_action_space
# are per-agent). Pass config (dict), vecenv, policy positionally to PuffeRL.
policy = create_policy(env.single_observation_space, env.single_action_space)
trainer = PuffeRL(config, env, policy)

カスタムタスクの開発

import pufferlib.vector

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

# Native PufferEnv -> default backend=PufferEnv is fine here.
env = pufferlib.vector.make(MyTask, num_envs=256)
trainer = PuffeRL(config, env, my_policy)  # config is a dict (see Training above)

高性能な最適化

import pufferlib.vector

# Maximize throughput (pass an env-constructor callable)
env = pufferlib.vector.make(
    my_env_creator,     # env constructor callable
    num_envs=1024,      # Large batch
    num_workers=16,     # Many workers
    backend=pufferlib.vector.Multiprocessing,
)

インストール

# Pin the 3.0 line — the config-dict trainer API and import paths in this skill
# target it. PyPI ships only an sdist, so a C compiler is needed to build it.
uv pip install "pufferlib==3.0.*"

PufferLib 5.0 は別の製品です。 上流のデフォルトブランチ（5.0）は、PufferTank を使ってソースからビルドする C/CUDA トレーナーです（./build.sh ENV、./puffer train）。pip パッケージはなく、サードパーティ（Gymnasium/PettingZoo 形式）の統合は削除されており、CPU での評価はできますが CPU でのトレーニングはできません。そのため、このスキルに記載されている Python API（PuffeRL、pufferlib.vector.make、pufferlib.emulation）が利用できるとは限りません。ユーザーが 5.0 を使用している場合は、その旨を伝え、3.0 のコードに合わせるのではなく、5.0 のドキュメントを案内してください。

ドキュメント

• 3.0 のソース（このスキルに対応）：https://github.com/PufferAI/PufferLib/tree/3.0
• 現在の上流ドキュメント（3.0 ではなく 5.0 を説明）：https://puffer.ai/docs.html
• GitHub：https://github.com/PufferAI/PufferLib
• Discord：コミュニティサポートを利用できます