Pufferlib V2

PufferLib v2

Use this skill for codebases pinned to PufferLib 2.x (pufferlib>=2.0,<3.0). PufferLib 2.0 was released on PyPI on December 8, 2024 and introduced the Ocean/native-environment generation described by PufferAI as the 1M SPS release.

[!IMPORTANT] This is not the v1 API. Do not write from pufferlib import PuffeRL; import the training module with from pufferlib import pufferl or use the puffer CLI. If the repository already uses trainer.evaluate(), trainer.train(), or trainer.mean_and_log(), inspect imports and pinning before editing because that pattern may be v1 or early-v2 code.

Version Selection

• Prefer pufferlib-v2 only for projects pinned to pufferlib>=2.0,<3.0, reproducing 2024/early-2025 experiments, or maintaining PufferLib 2.x environments.
• Prefer pufferlib-v3 for pufferlib>=3.0,<4.0 projects or when the code uses newer pufferlib.pufferl config helpers.
• Prefer the legacy pufferlib skill only for pufferlib==1.0.0 code.

Installation

python -m venv .venv
source .venv/bin/activate
pip install "pufferlib>=2.0,<3.0"

For environments with native graphics, Atari, NetHack, Procgen, Neural MMO, or CUDA-sensitive dependencies, prefer the upstream PufferTank container or the repository's pinned environment. Keep Python, PyTorch, CUDA, Gymnasium, and PettingZoo versions synchronized with the project lockfile.

Primary Workflow

1. Confirm the version:

   python - <<'PY'
   import pufferlib, importlib.metadata
   print(importlib.metadata.version('pufferlib'))
   print(pufferlib.__file__)
   PY
   

2. Prefer CLI training for supported Ocean/built-in environments:

   puffer train ENV_NAME --help
   puffer train ENV_NAME
   puffer eval ENV_NAME --help
   

3. For Python training, prefer the module-level pufferl API and local examples in the target repository over the old top-level PuffeRL import.
4. Keep vectorized environment settings explicit: number of workers, environments per worker / environment count, batch size, backend, and async/sync mode.
5. Validate with a tiny run before launching a long experiment.

Native Environment and Vectorization Guidance

PufferLib 2.x centers on high-throughput simulation:

• Use native PufferEnv/Ocean environments when performance matters. Native envs can write observations directly into buffers used by vectorization/training.
• Keep Gymnasium/PettingZoo support as a compatibility layer. Wrap existing environments rather than rewriting them first.
• Use a VecEnv-style interface by default. Async operation extends the synchronous step/reset workflow, so implement synchronous logic first, then add send/recv-style async only if needed.
• Start with a conservative worker count. Increase workers/envs per worker only after profiling CPU utilization and policy-forward time.

Training Pattern

Use this shape for v2 maintenance work; adapt names from the installed package and repository examples rather than assuming v1 signatures:

from pufferlib import pufferl

# Prefer repository-provided config/env/policy helpers when available.
config = load_or_create_config()
vecenv = make_vectorized_env(config)
policy = make_policy(vecenv, config)

trainer = pufferl.PuffeRL(config, vecenv, policy)
while not done_training(trainer, config):
    trainer.evaluate()
    logs = trainer.train()
    if logs:
        handle_logs(logs)

checkpoint = trainer.close()

When migrating old v1 code:

• Replace from pufferlib import PuffeRL with from pufferlib import pufferl.
• Replace pufferlib.make(...) calls with the v2 vectorization/environment construction used by the project.
• Ensure policies accept the state argument expected by the trainer when recurrent state or async environment metadata is enabled.
• Treat mean_and_log() as trainer-internal unless local examples call it directly.

Environment Development Checklist

• Define observation/action spaces before reset/step data are consumed.
• Return consistent Gymnasium-style obs, reward, terminated, truncated, info data in wrappers, or the native PufferEnv equivalent in native envs.
• For multi-agent envs, keep agent ordering stable and document how dead/missing agents are padded or masked.
• Avoid per-step allocation in performance-critical environments. Allocate buffers at initialization and write in place.
• Add a smoke test that resets, steps random actions for at least 100 transitions, and checks shapes/dtypes.

Performance Checklist

• Profile environment step time separately from policy forward/backward time.
• Use async vectorization when policy inference and environment stepping can overlap.
• Keep observation copies to a minimum; prefer native buffers for Ocean/C environments.
• Tune num_workers, env count, rollout horizon, batch size, and minibatch size together.
• Compare SPS after warmup, not during first-epoch compilation/build overhead.

Common Failure Modes

• ImportError for PuffeRL: old v1 import; use from pufferlib import pufferl.
• Shape mismatch on first batch: wrapper spaces do not match reset/step output; print single_observation_space, single_action_space, and first observation dtype/shape.
• Hanging multiprocessing run: test the same env with a serial backend, then reduce workers to one and re-enable multiprocessing.
• Low SPS with fast native env: check policy forward time, Python callbacks/logging, info payload size, and whether async mode is enabled.

Sources to Check When Updating

• PyPI release history for exact 2.x versions.
• PufferAI blog post “PufferLib 2.0: Reinforcement Learning at 1,000,000 Steps/Second”.
• The target repository's lockfile and training examples; v2 projects often carried local helper wrappers.