pufferlib

name	pufferlib
description	Best practices and guidance for developing PufferLib reinforcement learning environments. Use when creating RL environments, implementing observation/action spaces, setting up vectorization, training policies, or integrating with PufferLib's API.
allowed-tools	Read, Grep, Glob, Bash, Write, Edit
argument-hint	["topic"]

PufferLib Development Guide

You are helping develop a high-performance reinforcement learning environment using PufferLib. Follow these best practices and patterns from the official PufferLib documentation.

Environment Development

Native PufferEnv (Recommended for Performance)

For maximum performance, implement pufferlib.PufferEnv directly with zero-copy buffer sharing:

import gymnasium
import pufferlib

class MyPufferEnv(pufferlib.PufferEnv):
    def __init__(self, config_param=1.0, buf=None, seed=0):
        # Define spaces BEFORE calling super().__init__
        self.single_observation_space = gymnasium.spaces.Box(
            low=-1, high=1, shape=(obs_dim,), dtype=np.float32
        )
        self.single_action_space = gymnasium.spaces.Discrete(num_actions)
        self.num_agents = 1  # or more for multi-agent
        super().__init__(buf)

        # Custom initialization AFTER super().__init__
        self.config_param = config_param

    def reset(self, seed=0):
        # Write directly to shared buffer (zero-copy)
        self.observations[:] = initial_obs
        return self.observations, []

    def step(self, actions):
        # Process actions and write to buffers directly
        self.observations[:] = next_obs
        self.rewards[:] = reward_values
        self.terminals[:] = done_flags
        self.truncations[:] = truncation_flags

        infos = [{'metric': value} for _ in range(self.num_agents)]
        return self.observations, self.rewards, self.terminals, self.truncations, infos

    def close(self):
        pass

Key Points:

• Define single_observation_space, single_action_space, and num_agents BEFORE super().__init__
• Use buf=None and seed=0 parameters for compatibility with vectorization
• Write directly to self.observations, self.rewards, self.terminals, self.truncations buffers
• Return empty list [] for infos if no agent-specific info needed

Gymnasium Wrapper (For Existing Environments)

Wrap existing Gymnasium environments for PufferLib compatibility:

import gymnasium
import pufferlib.emulation

gym_env = YourGymnasiumEnv()
puffer_env = pufferlib.emulation.GymnasiumPufferEnv(gym_env)

# Check emulation status
print(f"Observation emulated: {puffer_env.is_obs_emulated}")
print(f"Action emulated: {puffer_env.is_atn_emulated}")
print(f"Num agents: {puffer_env.num_agents}")

Vectorization

Serial Backend (Development/Debugging)

Use for debugging - runs environments sequentially:

import pufferlib.vector

vecenv = pufferlib.vector.make(
    MyPufferEnv,
    num_envs=4,
    backend=pufferlib.vector.Serial,
    env_kwargs={'config_param': 2.0}
)

Multiprocessing Backend (Production)

Use for CPU-bound environments with async send/recv pattern:

vecenv = pufferlib.vector.make(
    MyPufferEnv,
    num_envs=8,           # Total environments
    num_workers=4,        # Parallel processes
    batch_size=4,         # Envs per recv call
    backend=pufferlib.vector.Multiprocessing,
    env_kwargs={'config_param': 2.0}
)

# Async API for maximum throughput
vecenv.async_reset()
obs, rewards, terminals, truncations, infos, env_ids, masks = vecenv.recv()

for step in range(num_steps):
    actions = policy(obs)
    vecenv.send(actions)
    obs, rewards, terminals, truncations, infos, env_ids, masks = vecenv.recv()

vecenv.close()

Heterogeneous Configurations

Pass different arguments to each environment:

vecenv = pufferlib.vector.make(
    [MyPufferEnv, MyPufferEnv, MyPufferEnv],
    num_envs=3,
    backend=pufferlib.vector.Serial,
    env_args=[[1.0], [2.0], [3.0]],
    env_kwargs=[
        {'max_steps': 100},
        {'max_steps': 150},
        {'max_steps': 200}
    ]
)

Policy Definition

Simple MLP Policy

import torch
import pufferlib

class SimplePolicy(torch.nn.Module):
    def __init__(self, env):
        super().__init__()
        obs_shape = env.single_observation_space.shape[0]
        num_actions = env.single_action_space.n

        self.net = torch.nn.Sequential(
            pufferlib.pytorch.layer_init(torch.nn.Linear(obs_shape, 128)),
            torch.nn.ReLU(),
            pufferlib.pytorch.layer_init(torch.nn.Linear(128, 128)),
            torch.nn.ReLU(),
        )
        self.action_head = torch.nn.Linear(128, num_actions)
        self.value_head = torch.nn.Linear(128, 1)

    def forward(self, observations, state=None):
        hidden = self.net(observations)
        logits = self.action_head(hidden)
        values = self.value_head(hidden)
        return logits, values

Research Policy with LayerNorm

class ResearchPolicy(torch.nn.Module):
    def __init__(self, obs_dim, action_dim, hidden_size=256):
        super().__init__()
        self.encoder = torch.nn.Sequential(
            pufferlib.pytorch.layer_init(torch.nn.Linear(obs_dim, hidden_size)),
            torch.nn.LayerNorm(hidden_size),
            torch.nn.ReLU(),
            pufferlib.pytorch.layer_init(torch.nn.Linear(hidden_size, hidden_size)),
        )
        self.action_head = pufferlib.pytorch.layer_init(
            torch.nn.Linear(hidden_size, action_dim), std=0.01
        )
        self.value_head = pufferlib.pytorch.layer_init(
            torch.nn.Linear(hidden_size, 1), std=1.0
        )

    def forward(self, observations, state=None):
        hidden = self.encoder(observations)
        return self.action_head(hidden), self.value_head(hidden).squeeze(-1)

Use pufferlib.pytorch.layer_init() for proper weight initialization.

Training Configuration

Default PPO Hyperparameters

train_config = {
    'device': 'cuda',
    'seed': 42,
    'batch_size': 4096,
    'bptt_horizon': 32,
    'minibatch_size': 512,
    'update_epochs': 4,
    'learning_rate': 3e-4,
    'adam_beta1': 0.9,
    'adam_beta2': 0.999,
    'adam_eps': 1e-8,
    'optimizer': 'adam',
    'clip_coef': 0.1,
    'vf_coef': 0.5,
    'ent_coef': 0.01,
    'max_grad_norm': 1.0,
    'total_timesteps': 10_000_000,
    'compile': False,          # torch.compile for 2x speedup
    'compile_mode': 'default',
    'precision': 'float32',    # or 'bfloat16' for mixed precision
    'use_rnn': False,          # LSTM for partial observability
}

Training Loop

from pufferlib import pufferl

trainer = pufferl.PuffeRL(train_config, vecenv, policy)

best_return = -float('inf')
while trainer.epoch < trainer.total_epochs:
    trainer.evaluate()  # Collect experience
    logs = trainer.train()  # Update policy

    if trainer.epoch % 50 == 0:
        avg_return = logs.get('return', 0)
        if avg_return > best_return:
            best_return = avg_return
            torch.save({
                'epoch': trainer.epoch,
                'policy_state_dict': policy.state_dict(),
                'optimizer_state_dict': trainer.optimizer.state_dict(),
                'return': avg_return,
            }, f'best_policy.pt')

        print(f"Epoch {trainer.epoch}: Return={avg_return:.2f}, SPS={logs.get('SPS', 0):.0f}")

trainer.close()

High-Level Training API

from pufferlib import pufferl

# Load config, customize, and train
args = pufferl.load_config('default')
args['train']['total_timesteps'] = 10_000_000
args['train']['learning_rate'] = 3e-4

vecenv = pufferl.load_env(env_name, args)
policy = pufferl.load_policy(args, vecenv, env_name)
trainer = pufferl.PuffeRL(args['train'], vecenv, policy)

Best Practices

Observation Space Design

1. Use flat arrays - PufferLib works best with Box spaces
2. Normalize observations - Keep values in [-1, 1] or [0, 1] range
3. Fixed dimensions - Avoid variable-length observations
4. Float32 dtype - Standard for neural network inputs

Action Space Design

1. Discrete actions - Use gymnasium.spaces.Discrete(n)
2. Multi-discrete - Use gymnasium.spaces.MultiDiscrete([n1, n2, ...])
3. Continuous - Use gymnasium.spaces.Box (requires different policy head)

Performance Optimization

1. Native PufferEnv - Use instead of Gymnasium wrapper when possible
2. Zero-copy buffers - Write directly to self.observations, etc.
3. Multiprocessing - Use for CPU-bound environments
4. Batch size - Larger batches improve GPU utilization
5. torch.compile - Enable compile: True for 2x speedup (PyTorch 2.0+)
6. Mixed precision - Use precision: 'bfloat16' for faster training

Debugging

1. Serial backend first - Debug with pufferlib.vector.Serial
2. Single environment - Test with num_envs=1 initially
3. Print observations - Verify shape and values
4. Check rewards - Ensure rewards are reasonable scale
5. Monitor terminals - Verify episodes end correctly

Project-Specific Notes

For this Yomi Hustle project:

• Package manager: Always use uv, never pip
• Testing: Run uv run pytest tests/ -v for all tests
• GDScript parity: Python must match original Godot behavior
• Work streams: A (core), B (fighter), C (states), D (integration)

Common Patterns

Environment with Episode Tracking

def reset(self, seed=0):
    self.episode_step = 0
    self.episode_return = 0.0
    self.observations[:] = self._get_initial_obs()
    return self.observations, []

def step(self, actions):
    self.episode_step += 1
    reward = self._compute_reward(actions)
    self.episode_return += reward

    done = self._check_terminal()
    truncated = self.episode_step >= self.max_steps

    self.observations[:] = self._get_obs()
    self.rewards[:] = reward
    self.terminals[:] = done
    self.truncations[:] = truncated

    infos = [{'episode_return': self.episode_return}] if done or truncated else []
    return self.observations, self.rewards, self.terminals, self.truncations, infos

Multi-Agent Environment

class MultiAgentEnv(pufferlib.PufferEnv):
    def __init__(self, num_agents=2, buf=None, seed=0):
        self.single_observation_space = gymnasium.spaces.Box(...)
        self.single_action_space = gymnasium.spaces.Discrete(...)
        self.num_agents = num_agents  # Set before super().__init__
        super().__init__(buf)

    def step(self, actions):
        # actions shape: (num_agents,)
        for i, action in enumerate(actions):
            self._process_agent_action(i, action)

        # Fill buffers for all agents
        for i in range(self.num_agents):
            self.observations[i] = self._get_agent_obs(i)
            self.rewards[i] = self._get_agent_reward(i)
            self.terminals[i] = self._is_agent_done(i)
            self.truncations[i] = False

        return self.observations, self.rewards, self.terminals, self.truncations, []