alterlab-pufferlib

PufferLib - 高性能强化学习

概述

PufferLib 是一个高性能强化学习库，旨在实现快速的并行环境模拟与训练。它通过优化的向量化、原生多智能体支持和高效的 PPO 实现（PuffeRL），实现每秒数百万步的训练速度。该库提供包含 20 多个环境的 Ocean 套件，并可与 Gymnasium、PettingZoo 及专业 RL 框架无缝集成。

何时使用此技能

在以下情况下使用此技能：

• 使用 PPO 训练 RL Agent，适用于任何环境（单智能体或多智能体）
• 使用 PufferEnv API 创建自定义环境
• 优化性能，实现并行环境模拟（向量化）
• 集成现有环境，包括 Gymnasium、PettingZoo、Atari、Procgen 等
• 使用 CNN、LSTM 或自定义架构开发策略
• 扩展 RL，实现每秒数百万步的训练速度，加快实验
• 多智能体 RL，使用原生多智能体环境支持

不适用的场景

场景	改用
使用标准、文档完善的 PPO/SAC/DQN 实现快速构建单智能体原型	alterlab-stable-baselines3
使用检查点和多 GPU 进行监督式深度学习训练循环（分类、回归）	alterlab-pytorch-lightning
使用基于规则的 Agent 模拟社会系统，且不涉及学习	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 驱动。三个循环方法均可用：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，了解：

• 架构和性能特征
• 工作进程和批次大小配置
• 串行、多进程和异步模式
• 共享内存和零拷贝模式
• 面向大规模场景的分层向量化
• 多 Agent 向量化策略
• 性能分析与故障排查

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 倍）
• 用于复杂观测的多输入策略
• 连续动作策略
• 多 Agent 策略（共享参数与独立参数）
• 高级架构（注意力、残差）
• 观测归一化和梯度裁剪
• 策略调试与测试

5. 环境集成

无缝集成来自热门强化学习框架的环境。

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 多 Agent：

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（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.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. 分析当前吞吐量（每秒步数）
2. 检查向量化配置（num_envs、num_workers）
3. 优化环境代码（原地操作、NumPy 向量化）
4. 对关键路径考虑使用 C 实现
5. 使用 references/vectorization.md 进行系统性优化

资源

scripts/

train_template.py —— 完整的训练脚本模板，包含：

• 环境创建与配置
• 策略初始化
• 日志记录器集成（WandB、Neptune）
• 带检查点保存的训练循环
• 命令行参数解析
• 多 GPU 分布式训练设置

env_template.py —— 环境实现模板：

• 单 Agent PufferEnv 示例（网格世界）
• 多 Agent PufferEnv 示例（协作导航）
• 多种观测／动作空间模式
• 测试工具

references/

training.md - 完整的训练指南：

• 训练工作流和 CLI 选项
• 超参数配置
• 分布式训练（多 GPU、多节点）
• 监控和日志记录
• 检查点
• Protein 超参数调优
• 性能优化
• 常见训练模式
• 故障排除

environments.md - 环境开发指南：

• PufferEnv API 和特性
• 观测和动作空间
• 多 Agent 环境
• Ocean 套件环境
• 自定义环境开发工作流
• 从 Python 到 C 的优化路径
• 第三方环境集成
• 包装器和最佳实践
• 调试

vectorization.md - 向量化优化：

• 架构和关键优化
• 向量化模式（串行、多进程、异步）
• Worker 和批次配置
• 共享内存和零拷贝模式
• 高级向量化（分层、自定义）
• 多 Agent 向量化
• 性能监控和分析
• 故障排除和最佳实践

policies.md - 策略架构指南：

• 基本策略结构
• 用于图像的 CNN 策略
• 带优化的 LSTM 策略
• 多输入策略
• 连续动作策略
• 多 Agent 策略
• 高级架构（注意力、残差）
• 观测处理和反展平
• 初始化和归一化
• 调试和测试

integration.md - 框架集成指南：

• Gymnasium 集成
• PettingZoo 集成（parallel 和 AEC）
• 第三方环境（Procgen、NetHack、Minigrid 等）
• 自定义包装器（观测、奖励、帧堆叠等）
• 空间转换和反展平
• 环境注册
• 兼容性模式
• 性能注意事项
• 集成调试

成功秘诀

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

常见用例

在标准基准上训练

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)

多 Agent 学习

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）是一个 C/CUDA 训练器，通过 PufferTank 从源代码构建（./build.sh ENV、./puffer train）。它没有 pip 软件包，已移除第三方（Gymnasium/PettingZoo 风格）集成，并支持 CPU 评估但不支持 CPU 训练，因此不要假设此技能中的 Python API（PuffeRL、pufferlib.vector.make、pufferlib.emulation）在其中存在。如果用户使用的是 5.0，请明确告知，并引导他们查看该版本的文档，而不是改写 3.0 代码。

文档

• 3.0 源代码（与此技能匹配）：https://github.com/PufferAI/PufferLib/tree/3.0
• 当前上游文档（介绍 5.0，而非 3.0）：https://puffer.ai/docs.html
• GitHub：https://github.com/PufferAI/PufferLib
• Discord：可获取社区支持