/pufferlib

name: pufferlib
description: This skill should be used when working with reinforcement learning tasks including high-performance RL training, custom environment development, vectorized parallel simulation, multi-agent systems, or integration with existing RL environments (Gymnasium, PettingZoo, Atari, Procgen, etc.). Use this skill for implementing PPO training, creating PufferEnv environments, optimizing RL performance, or developing policies with CNNs/LSTMs.

PufferLib - High-Performance Reinforcement Learning

Overview

PufferLib is a high-performance reinforcement learning library designed for fast parallel environment simulation and training. It achieves training at millions of steps per second through optimized vectorization, native multi-agent support, and efficient PPO implementation (PuffeRL). The library provides the Ocean suite of 20+ environments and seamless integration with Gymnasium, PettingZoo, and specialized RL frameworks.

When to Use This Skill

Use this skill when:

• **Training RL agents** with PPO on any environment (single or multi-agent)
• **Creating custom environments** using the PufferEnv API
• **Optimizing performance** for parallel environment simulation (vectorization)
• **Integrating existing environments** from Gymnasium, PettingZoo, Atari, Procgen, etc.
• **Developing policies** with CNN, LSTM, or custom architectures
• **Scaling RL** to millions of steps per second for faster experimentation
• **Multi-agent RL** with native multi-agent environment support

Core Capabilities

1. High-Performance Training (PuffeRL)

PuffeRL is PufferLib's optimized PPO+LSTM training algorithm achieving 1M-4M steps/second.

**Quick start training:**

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

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

**Python training loop:**

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

**For comprehensive training guidance**, read `references/training.md` for:

• Complete training workflow and CLI options
• Hyperparameter tuning with Protein
• Distributed multi-GPU/multi-node training
• Logger integration (Weights & Biases, Neptune)
• Checkpointing and resume training
• Performance optimization tips
• Curriculum learning patterns

2. Environment Development (PufferEnv)

Create custom high-performance environments with the PufferEnv API.

**Basic environment structure:**

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

**Use the template script:** `scripts/env_template.py` provides complete single-agent and multi-agent environment templates with examples of:

• Different observation space types (vector, image, dict)
• Action space variations (discrete, continuous, multi-discrete)
• Multi-agent environment structure
• Testing utilities

**For complete environment development**, read `references/environments.md` for:

• PufferEnv API details and in-place operation patterns
• Observation and action space definitions
• Multi-agent environment creation
• Ocean suite (20+ pre-built environments)
• Performance optimization (Python to C workflow)
• Environment wrappers and best practices
• Debugging and validation techniques

3. Vectorization and Performance

Achieve maximum throughput with optimized parallel simulation.

**Vectorization setup:**

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

**Key optimizations:**

• Shared memory buffers for zero-copy observation passing
• Busy-wait flags instead of pipes/queues
• Surplus environments for async returns
• Multiple environments per worker

**For vectorization optimization**, read `references/vectorization.md` for:

• Architecture and performance characteristics
• Worker and batch size configuration
• Serial vs multiprocessing vs async modes
• Shared memory and zero-copy patterns
• Hierarchical vectorization for large scale
• Multi-agent vectorization strategies
• Performance profiling and troubleshooting

4. Policy Development

Build policies as standard PyTorch modules with optional utilities.

**Basic policy structure:**

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