curl -fsSL https://raw.githubusercontent.com/PufferAI/PufferTank/refs/heads/5.0/install.sh | bash

git clone --filter=blob:none --single-branch --branch 5.0 https://github.com/pufferai/puffertank
cd puffertank
./docker.sh test

Installation Test

./build.sh breakout --cu
./puffer train
./puffer eval latest

./build.sh breakout --cpu
./breakout

Remote Rendering

# remote
sudo apt install tigervnc-standalone-server
vncpasswd
printf '#!/bin/sh\nexec sleep infinity\n' > ~/.vnc/xstartup
chmod +x ~/.vnc/xstartup
vncserver :1 -geometry 1920x1080 -localhost yes

# local
sudo apt install tigervnc-viewer
ssh -L 5901:127.0.0.1:5901 user@host
vncviewer localhost:5901

# remote
DISPLAY=:1 ./puffer eval latest
DISPLAY=:1 ./seethestars

# PuffeRL
./build.sh ENV
./puffer [train|eval|sweep] [--section.key=value]

# Standalone
./build.sh ENV --cpu
./ENV [latest|your_model.bin] [--section.key=value]

# Example: Train and render
./build.sh breakout --cu
./puffer train --env.frameskip=3
./puffer eval latest

# Example: Multi-GPU sweep + visualization
./build.sh breakout
./puffer sweep --sweep.gpus=4
./build.sh cache_data # Builds resources/constellation/experiments.ini from logs/
./build.sh constellation
./seethestars

# Example: Headless environment dev
./build.sh breakout --cpu --debug
./breakout --headless

# Extra flags
--debug    # -O0 -g; ASan/UBSan on --cpu
--profile  # Build kernel benchmark
--cu       # Trainer build flag for CUDA native environments
--float    # For numerics testing vs. default bf16
--web      # emscripten web build (install separately)

The PufferLib 5.0 native backend is ~10,000 lines of CUDA C. ~5,000 is devoted to training and infrastructure. The rest is sweeps, visualization, and logging. Fork the project and edit directly. There is no prebuilt package blackboxing functionality. Everything is written as simply and transparently as possible. If you're new to low-level dev, it's much easier than you think. Give the environment tutorial a try!

Memory Management: Tensors in PufferLib are just structs with a shape and a data pointer. Every tensor registers its size with an allocator at startup time. After all tensors are registered, the allocator sums up the sizes and does a single allocation of contiguous memory. There are separate allocators for weights, gradients, and activations. No tensors are created or reallocated afterwards. Static memory improves performance, simplifies cudagraph tracing, and cleans up profile timelines. Since weights and gradients are contiguous, we can apply updates to them in a single kernel without looping over tensors.

Tracing: Cudagraphs capture and replay GPU operations in order to reduce kernel launch overhead. On the first epoch of each run, PufferLib traces both the rollout forward pass and the entire train minibatch + loss + policy update. We bloat memory a bit by tracing separate cudagraphs for each step during rollouts in order to avoid an extra data copy for intermediate graph buffers.

Vectorization: Environment instances are chunked into buffers, each of which is associated with a rollout worker on a separate CUDA stream. Within each buffer, environment execution is parallelized with OMP threading. Rollout workers are independent of each other but each process the same number of environment steps per epoch. Each buffer asynchronously queues data transfers to/from the GPU and uses pinned memory.

Kernels: The main consideration for performance is fusing small elementwise operations to reduce memory bandwidth. For most of our kernels, the efficiency of the compute load is secondary. The MinGRU kernels are load-bearing and have received more attention to performance. The learnable workload of MinGRU is a set of linear layers implemented as cuBLAS matmuls. None of our kernels include complex operations with tensor cores.

Algorithm: PufferLib implements a PPO-variant with improvements based on our own research. It supports both synchronous and 1-epoch asynchronous training as in CleanBa.

PufferNet: Our default model architecture combines MinGRU, a linear recurrence that is parallelizable over the time dimension, with highway nets, a fancy residual that replaces expensive normalization layers. It's equal to or better than an LSTM on every environment we've tested and much faster.

Sweeps: The fundamental unit of compute in PufferLib is a hyperparameter sweep, not a single experiment. Protein is a tuning algorithm based on our own research. It combines Gaussian processes with a simple genetic algorithm over the Pareto-frontier defined by wall-clock experiment time and score.

Ocean environments are written in C. Mostly very simple C. Like first 2 weeks of an intro systems course C. Compile standalone builds with --cpu. We recommend --debug during development to catch indexing and overflow bugs with the address sanitizer. C isn't much fun without it. Observations, actions, rewards, and terminals are each allocated as big chunks of memory that are contiguous across all (usually thousands) of environment instances.

The Minimal environment is a clean template with comments that walk you through our API. Read and understand it first. To create your own environment, first copy and rename all the files. PufferLib explicitly looks for your_env.h inside of ocean/your_env. The .cu file included is a more advanced example of a custom encoder, which most environments will not require.

Here's a checklist of common bugs if your env is not training:

• Zero or incorrect observations/actions: Ensure you have defined your observation and action metadata (space/size/type) correctly in your binding.c. Note that observation buffers are not zeroed for you every step, so if you are only writing information to specific indicies (i.e. one-hots), be sure you are zeroing them. One memset is usually the quickest.
• NaN losses: If your losses NaN after 1 epoch, your environment almost certainly has data corruption. You are writing to memory that does not match the obs/action sizes you have defined, or you have defined them using the wrong types. If your losses NaN after a longer time, it may be hyperparameters.
• Incorrect or missing resets: Your environment should handle its own resets internally. For envs that never reset, it is often useful to respawn agents if they are stuck (e.g. no reward for 500 steps). If every rollout has the same number of timesteps, you should stagger your environments on init.
• Not zeroing observations: If you are writing over your observation buffers but only write to some elements each step, If you don't zero out rewards, they will retain their value from the previous step. Ditto for terminals. Zero them at the start of puf_step to be safe. A single memset will do it for multiagent envs. If you are not setting every element of every obs (i.e. one-hots), make sure to clear those too.
• Manually inspect data scale: You want observations and rewards to be roughly in the range of -1 to 1. 0 to 1 is fine. -5 to 5 is kind of fine. Randomly having one obs with magnitude 1000 is not.

Can I run PufferLib without an Nvidia GPU? As of 5.0, we have a solid --cpu eval mode but no CPU training option. A 10c/hr cloud GPU will outperform your CPU by 10x.

Why are the docs so short? Our API surface is tiny and fully covered by our example envs.

I'm new to RL. How do I contribute? Start by building a simple new environment and getting it to train. I review environment PRs from new contributors live on stream.

Where did all the Python/third-party stuff go? It was all 100x+ slower than PufferLib is now. We do plan on hooking the C/C++ for Atari and maybe ProcGen into our low-level interface when we have time.

Why is it called PufferLib? Would you have rathered yet another minimal tech logo? Here, have a pufferfish 🐡