Skip to main content
PufferLib provides a comprehensive wrapper system that unifies single-agent and multi-agent environments from Gymnasium and PettingZoo into a common vectorized interface. This allows you to write training code once and use it across diverse environments.

Wrapper architecture

PufferLib’s wrapper system consists of three main components:
  1. Emulation wrappers - Convert Gymnasium and PettingZoo environments to the PufferEnv interface
  2. Preprocessing wrappers - Apply transformations like observation resizing and episode statistics
  3. Environment-specific wrappers - Handle unique requirements of specific environments

How wrappers integrate

The typical wrapper stack for an environment looks like this:
import pufferlib
import pufferlib.emulation
import gymnasium

# 1. Create base environment
env = gymnasium.make('CartPole-v1')

# 2. Apply preprocessing wrappers
env = pufferlib.EpisodeStats(env)

# 3. Convert to PufferEnv interface
env = pufferlib.emulation.GymnasiumPufferEnv(env=env)
All environments are converted to use:
  • Vectorized interface - Batch dimension for observations, rewards, terminals
  • In-place updates - Zero-copy buffer management
  • Unified API - Same interface for single-agent and multi-agent environments

Available environment wrappers

PufferLib includes 36 environment integrations out of the box:

Classic environments

  • Atari - Arcade Learning Environment games (Breakout, Pong, etc.)
  • Classic Control - CartPole, MountainCar, Acrobot, Pendulum
  • Box2D - LunarLander, BipedalWalker, CarRacing
  • MuJoCo - Ant, HalfCheetah, Hopper, Humanoid, Walker2D

Procedural generation

  • Procgen - 16 procedurally generated games
  • Crafter - Survival crafting environment
  • Craftax - JAX-accelerated crafting environment
  • MiniGrid - Partially observable gridworld environments
  • Griddly - Grid-based game engine environments

Roguelike games

  • NetHack - The classic roguelike dungeon crawler
  • MiniHack - Simplified NetHack tasks
  • Pokémon Red - Gameboy game environment
  • Link’s Awaken - Zelda game environment

DeepMind environments

  • DM-Control - MuJoCo-based continuous control suite
  • DM-Lab - 3D navigation and puzzle-solving
  • BSuite - Behavior suite for RL agents

Multi-agent environments

  • Butterfly - Cooperative PettingZoo environments
  • MAgent - Large-scale multi-agent battle simulations
  • SMAC - StarCraft Multi-Agent Challenge
  • Neural MMO - Massively multi-agent survival game
  • OpenSpiel - Game theory and multi-agent games
  • Slime Volleyball - Competitive two-player game
  • MicroRTS - Real-time strategy game
  • CoGames - Cooperative game environments

Retro and simulation

  • Stable-Retro - Classic console games via emulation
  • VizDoom - 3D FPS environments
  • MineRL - Minecraft learning environment
  • Kinetix - Character animation and control
  • GPUDrive - GPU-accelerated driving simulator
  • ManiSkill - Robotic manipulation tasks

Custom environments

  • NMMO - Neural MMO with custom wrappers
  • Trade Sim - Economic trading simulation
  • Tribal Village - Multi-agent village simulation
  • Metta - Custom environment framework
  • GVGAI - General Video Game AI framework
  • Test - Mock environments for testing

Preprocessing wrappers

PufferLib provides several built-in preprocessing wrappers:

ResizeObservation

Efficiently downscales observations using NumPy slicing:
import pufferlib

env = pufferlib.ResizeObservation(env, downscale=2)
This is much faster than OpenCV-based resizing used in other libraries.

EpisodeStats

Tracks episode returns and lengths:
env = pufferlib.EpisodeStats(env)
# Automatically adds episode_return and episode_length to info dict

ClipAction

Clips continuous actions to valid ranges:
env = pufferlib.ClipAction(env)

MultiagentEpisodeStats

Tracks per-agent statistics in multi-agent environments:
env = pufferlib.MultiagentEpisodeStats(env)

MeanOverAgents

Averages info dictionaries across agents:
env = pufferlib.MeanOverAgents(env)

Space emulation

PufferLib automatically handles complex observation and action spaces by “emulating” them as flat arrays:
  • Dict spaces - Flattened into single array
  • Tuple spaces - Concatenated into single array
  • Discrete/MultiDiscrete - Converted to integer arrays
This happens transparently in GymnasiumPufferEnv and PettingZooPufferEnv. See Gymnasium integration and PettingZoo integration for details.

Creating environment instances

Each environment module provides a make() function and env_creator() function:
from pufferlib.environments import atari

# Direct creation
env = atari.make('breakout')

# Creator function (for parallel environments)
env_creator = atari.env_creator('breakout')
env = env_creator()
The env_creator() returns a functools.partial that can be called multiple times to create identical environments for vectorization.

Buffer management

PufferLib environments use pre-allocated NumPy buffers for zero-copy performance:
import numpy as np
import pufferlib.emulation

# Create custom buffers
buf = {
    'observations': np.zeros((num_agents, *obs_shape), dtype=obs_dtype),
    'rewards': np.zeros(num_agents, dtype=np.float32),
    'terminals': np.zeros(num_agents, dtype=bool),
    'truncations': np.zeros(num_agents, dtype=bool),
    'masks': np.ones(num_agents, dtype=bool),
    'actions': np.zeros((num_agents, *atn_shape), dtype=atn_dtype),
}

env = pufferlib.emulation.GymnasiumPufferEnv(env=base_env, buf=buf)
This allows multiple environments to share the same memory buffers, eliminating copies during vectorization.

Next steps

Gymnasium integration

Wrap single-agent Gymnasium environments

PettingZoo integration

Wrap multi-agent PettingZoo environments

Custom wrappers

Create your own environment wrappers

Ocean environments

Browse all Ocean environments

Build docs developers (and LLMs) love