github/scientific-agent-skills
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update PufferLib versioned workflows

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Timothy Kassis
2026-07-23 17:10:07 -07:00
parent 27ea7f4d63
commit f515922468
17 changed files with 3512 additions and 3410 deletions
Download Patch File Download Diff File Expand all files Collapse all files

647
skills/pufferlib/SKILL.md
View File

@@ -1,433 +1,332 @@
---
name: pufferlib
description: High-performance reinforcement learning framework optimized for speed and scale. Use when you need fast parallel training, vectorized environments, multi-agent systems, or integration with game environments (Atari, Procgen, NetHack). Achieves 2-10x speedups over standard implementations. For quick prototyping or standard algorithm implementations with extensive documentation, use stable-baselines3 instead.
license: MIT license
metadata: {"version": "1.0", "skill-author": "K-Dense Inc."}
description: Version-aware guidance for PufferLib reinforcement-learning environments, vectorization, policies, PuffeRL training, evaluation, and safe checkpoint review. Use when adapting Gymnasium/PettingZoo environments to published PufferLib 3.0.0 or working with the redesigned native 4.0 source line.
license: MIT
compatibility: Bundled CLIs require Python 3.10+ and use only the standard library. Published pufferlib 3.0.0 supports Python >=3.9 but ships as a native-code source archive; current 4.0 source requires Python >=3.10, Torch >=2.9, and an audited CPU/CUDA toolchain. Network, GPU, native builds, environment plug-ins, assets, checkpoints, and external logging are never required by the bundled CLIs.
allowed-tools:
- Read
- Bash
- Grep
- Python
metadata:
version: "1.1"
skill-author: "K-Dense Inc."
last-reviewed: "2026-07-23"
---
# PufferLib - High-Performance Reinforcement Learning
# PufferLib
## Overview
Use PufferLib with an explicit version profile. Upstream currently has two
incompatible surfaces:
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.
| Profile | Status on 2026-07-23 | Main use |
|---|---|---|
| `pufferlib==3.0.0` | Latest stable PyPI release, published 2025-06-23 | Python/Gymnasium/PettingZoo emulation, `pufferlib.vector`, Torch PuffeRL |
| source `4.0` | Upstream default branch; not the latest stable PyPI artifact | Native C Ocean environments, native CUDA trainer, optional Torch fallback |
## When to Use This Skill
Do not combine 3.0 imports with 4.0 config/CLI examples. The 4.0 redesign
removed the 3.0 `emulation`, `vector`, and `pytorch` modules from the current
package tree.
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
## Safe defaults
## Core Capabilities
1. Start with bundled synthetic, CPU-only, network-free tools.
2. Do not import an arbitrary environment by dotted path. Bundled tools accept
only allowlisted built-ins and slug identifiers.
3. Do not install or execute an unreviewed environment package, native
extension, ROM, map, checkpoint, or pickle file.
4. Verify official source, immutable revision, licenses, checksums or
attestations, and build hooks. Sandbox native builds and first execution.
5. Cap steps, environments, agents, workers, threads, buffers, memory, disk,
render size, and wall time.
6. Keep training and evaluation environments/seeds separate.
7. Default logging to local/none. External logging requires explicit opt-in,
disclosure acknowledgment, and separate artifact-upload approval.
8. Never pass W&B or Neptune credentials via CLI, INI, JSON, tags, run names, or
logger configuration. Never print them.
9. Never dump all environment variables or recursively search for `.env`.
10. Hash checkpoint bytes before trusted, sandboxed loading; metadata inspection
is not proof of safety.
### 1. High-Performance Training (PuffeRL)
## First local checks
PuffeRL is PufferLib's optimized PPO+LSTM training algorithm achieving 1M-4M steps/second.
All bundled CLIs are dependency-free and emit strict JSON:
**Quick start training:**
```bash
# CLI training
puffer train procgen-coinrun --train.device cuda --train.learning-rate 3e-4
# Distributed training
torchrun --nproc_per_node=4 train.py
python3 scripts/env_template.py --help
python3 scripts/env_contract_validator.py
python3 scripts/benchmark_vectorization.py --backend serial
python3 scripts/train_template.py
python3 scripts/validate_plan.py
python3 scripts/repro_plan.py
```
**Python training loop:**
```python
import pufferlib
from pufferlib import PuffeRL
Defaults are synthetic, deterministic, bounded, local, CPU-only, no-network,
and dry-run where training would otherwise occur.
# Create vectorized environment
env = pufferlib.make('procgen-coinrun', num_envs=256)
## Installation and provenance
# Create trainer
trainer = PuffeRL(
env=env,
policy=my_policy,
device='cuda',
learning_rate=3e-4,
batch_size=32768
)
### Published 3.0.0
# Training loop
for iteration in range(num_iterations):
trainer.evaluate() # Collect rollouts
trainer.train() # Train on batch
trainer.mean_and_log() # Log results
PyPI supplies only `pufferlib-3.0.0.tar.gz`:
```text
sha256: 7df3a3e3f5f894d78d2a1f5374097890aec01473183e748abefe4f3faa10eaa9
Requires-Python: >=3.9
```
**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
After source/build review, create a pinned uv project:
### 2. Environment Development (PufferEnv)
Create custom high-performance environments with the PufferEnv API.
**Basic environment structure:**
```python
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
```bash
uv venv --python 3.11
uv add --exact --no-sync "pufferlib==3.0.0"
uv lock
uv sync --frozen
```
**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
Commit `pyproject.toml` and `uv.lock`; verify the archive digest and every
resolved dependency. The source build can compile native code and fetch build
assets, so resolve/build in a sandbox without credentials or sensitive mounts.
The uploaded metadata does not pin Torch or CUDA; do not claim a supported CUDA
matrix that PyPI does not declare.
**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
### Current 4.0 source
### 3. Vectorization and Performance
The reviewed branch head on 2026-07-23 was:
Achieve maximum throughput with optimized parallel simulation.
**Vectorization setup:**
```python
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
```text
25647630e1b15330bb3153a5a0d3ff8d234c3acf
```
**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
Pin the commit, not branch `4.0`:
**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:**
```python
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)
return self.actor(features), self.critic(features)
```bash
uv add --no-sync \
"pufferlib @ git+https://github.com/PufferAI/PufferLib.git@25647630e1b15330bb3153a5a0d3ff8d234c3acf"
uv lock
```
**For complete policy development**, read `references/policies.md` for:
- CNN policies for image observations
- Recurrent policies with optimized LSTM (3x faster inference)
- Multi-input policies for complex observations
- Continuous action policies
- Multi-agent policies (shared vs independent parameters)
- Advanced architectures (attention, residual)
- Observation normalization and gradient clipping
- Policy debugging and testing
The current package declares Python `>=3.10` and Torch `>=2.9`. Upstream
PufferTank currently uses Ubuntu 24.04, Python 3.12, and an NVIDIA CUDA
13.0.2/cuDNN development image with the `cu130` Torch index, but does not pin
the exact Torch wheel or all system packages. Treat it as a reference, not a
complete lock. Never execute a remote installer directly from a pipe.
### 5. Environment Integration
Read `references/training.md` before any installation or build.
Seamlessly integrate environments from popular RL frameworks.
## Environment workflow
### 1. Validate the contract
Gymnasium reset returns `(observation, info)`. Step returns:
**Gymnasium integration:**
```python
import gymnasium as gym
import pufferlib
# Wrap Gymnasium environment
gym_env = gym.make('CartPole-v1')
env = pufferlib.emulate(gym_env, num_envs=256)
# Or use make directly
env = pufferlib.make('gym-CartPole-v1', num_envs=256)
(observation, reward, terminated, truncated, info)
```
**PettingZoo multi-agent:**
```python
# Multi-agent environment
env = pufferlib.make('pettingzoo-knights-archers-zombies', num_envs=128)
Validate spaces, shapes, dtypes, finite rewards, booleans, reset-before-step,
reset-after-end, seeding, and cleanup. `terminated` is an MDP terminal;
`truncated` is an external cutoff such as a time limit. Preserve the distinction
for bootstrapping and metrics.
```bash
python3 scripts/env_contract_validator.py \
--steps 64 --episodes 8 --seed 42
```
**Supported frameworks:**
- Gymnasium / OpenAI Gym
- PettingZoo (parallel and AEC)
- Atari (ALE)
- Procgen
- NetHack / MiniHack
- Minigrid
- Neural MMO
- Crafter
- GPUDrive
- MicroRTS
- Griddly
- And more...
### 2. Adapt only after review
**For integration details**, read `references/integration.md` for:
- Complete integration examples for each framework
- Custom wrappers (observation, reward, frame stacking, action repeat)
- Space flattening and unflattening
- Environment registration
- Compatibility patterns
- Performance considerations
- Integration debugging
Published 3.0 uses explicit wrappers:
## Quick Start Workflow
### For Training Existing Environments
1. Choose environment from Ocean suite or compatible framework
2. Use `scripts/train_template.py` as starting point
3. Configure hyperparameters for your task
4. Run training with CLI or Python script
5. Monitor with Weights & Biases or Neptune
6. Refer to `references/training.md` for optimization
### For Creating Custom Environments
1. Start with `scripts/env_template.py`
2. Define observation and action spaces
3. Implement `reset()` and `step()` methods
4. Test environment locally
5. Vectorize with `pufferlib.emulate()` or `make()`
6. Refer to `references/environments.md` for advanced patterns
7. Optimize with `references/vectorization.md` if needed
### For Policy Development
1. Choose architecture based on observations:
- Vector observations → MLP policy
- Image observations → CNN policy
- Sequential tasks → LSTM policy
- Complex observations → Multi-input policy
2. Use `layer_init` for proper weight initialization
3. Follow patterns in `references/policies.md`
4. Test with environment before full training
### For Performance Optimization
1. Profile current throughput (steps per second)
2. Check vectorization configuration (num_envs, num_workers)
3. Optimize environment code (in-place ops, numpy vectorization)
4. Consider C implementation for critical paths
5. Use `references/vectorization.md` for systematic optimization
## Resources
### scripts/
**train_template.py** - Complete training script template with:
- Environment creation and configuration
- Policy initialization
- Logger integration (WandB, Neptune)
- Training loop with checkpointing
- Command-line argument parsing
- Multi-GPU distributed training setup
**env_template.py** - Environment implementation templates:
- Single-agent PufferEnv example (grid world)
- Multi-agent PufferEnv example (cooperative navigation)
- Multiple observation/action space patterns
- Testing utilities
### references/
**training.md** - Comprehensive training guide:
- Training workflow and CLI options
- Hyperparameter configuration
- Distributed training (multi-GPU, multi-node)
- Monitoring and logging
- Checkpointing
- Protein hyperparameter tuning
- Performance optimization
- Common training patterns
- Troubleshooting
**environments.md** - Environment development guide:
- PufferEnv API and characteristics
- Observation and action spaces
- Multi-agent environments
- Ocean suite environments
- Custom environment development workflow
- Python to C optimization path
- Third-party environment integration
- Wrappers and best practices
- Debugging
**vectorization.md** - Vectorization optimization:
- Architecture and key optimizations
- Vectorization modes (serial, multiprocessing, async)
- Worker and batch configuration
- Shared memory and zero-copy patterns
- Advanced vectorization (hierarchical, custom)
- Multi-agent vectorization
- Performance monitoring and profiling
- Troubleshooting and best practices
**policies.md** - Policy architecture guide:
- Basic policy structure
- CNN policies for images
- LSTM policies with optimization
- Multi-input policies
- Continuous action policies
- Multi-agent policies
- Advanced architectures (attention, residual)
- Observation processing and unflattening
- Initialization and normalization
- Debugging and testing
**integration.md** - Framework integration guide:
- Gymnasium integration
- PettingZoo integration (parallel and AEC)
- Third-party environments (Procgen, NetHack, Minigrid, etc.)
- Custom wrappers (observation, reward, frame stacking, etc.)
- Space conversion and unflattening
- Environment registration
- Compatibility patterns
- Performance considerations
- Debugging integration
## Tips for Success
1. **Start simple**: Begin with Ocean environments or Gymnasium integration before creating custom environments
2. **Profile early**: Measure steps per second from the start to identify bottlenecks
3. **Use templates**: `scripts/train_template.py` and `scripts/env_template.py` provide solid starting points
4. **Read references as needed**: Each reference file is self-contained and focused on a specific capability
5. **Optimize progressively**: Start with Python, profile, then optimize critical paths with C if needed
6. **Leverage vectorization**: PufferLib's vectorization is key to achieving high throughput
7. **Monitor training**: Use WandB or Neptune to track experiments and identify issues early
8. **Test environments**: Validate environment logic before scaling up training
9. **Check existing environments**: Ocean suite provides 20+ pre-built environments
10. **Use proper initialization**: Always use `layer_init` from `pufferlib.pytorch` for policies
## Common Use Cases
### Training on Standard Benchmarks
```python
# Atari
env = pufferlib.make('atari-pong', num_envs=256)
import pufferlib.emulation
# Procgen
env = pufferlib.make('procgen-coinrun', num_envs=256)
# Minigrid
env = pufferlib.make('minigrid-empty-8x8', num_envs=256)
wrapped = pufferlib.emulation.GymnasiumPufferEnv(reviewed_gymnasium_instance)
```
### Multi-Agent Learning
```python
# PettingZoo
env = pufferlib.make('pettingzoo-pistonball', num_envs=128)
For a reviewed PettingZoo Parallel environment:
# Shared policy for all agents
policy = create_policy(env.observation_space, env.action_space)
trainer = PuffeRL(env=env, policy=policy)
```python
wrapped = pufferlib.emulation.PettingZooPufferEnv(reviewed_parallel_instance)
```
### Custom Task Development
```python
# Create custom environment
class MyTask(PufferEnv):
# ... implement environment ...
There is no supported 3.0 `pufferlib.emulate(...)` shortcut matching the old
skill. Read `references/environments.md` and `references/integration.md`.
# Vectorize and train
env = pufferlib.emulate(MyTask, num_envs=256)
trainer = PuffeRL(env=env, policy=my_policy)
```
### 3. Native environments
Published 3.0 `PufferEnv` requires
`single_observation_space`, `single_action_space`, and `num_agents` before
`super().__init__(buf)`. It uses in-place vector buffers and returns separate
terminal/truncation arrays plus a list of info dictionaries.
Current 4.0 uses C bindings. Start from upstream `ocean/squared` (single-agent)
or `ocean/target` (multi-agent), build one environment in local/sanitized mode,
and verify every buffer size/type/index before optimization.
## Vectorization workflow
Published 3.0:
### High-Performance Optimization
```python
# Maximize throughput
env = pufferlib.make(
'my-env',
num_envs=1024, # Large batch
num_workers=16, # Many workers
envs_per_worker=64 # Optimize per worker
import pufferlib.vector
vecenv = pufferlib.vector.make(
reviewed_creator,
backend=pufferlib.vector.Serial,
num_envs=4,
seed=42,
)
```
## Installation
Move to `Multiprocessing` only after serial traces pass. Record
`num_envs`, `num_workers`, `batch_size`, zero-copy mode, start method, agent
count, masks, and actual returned shapes. For multi-agent environments, batch
length is based on agent slots, not necessarily `num_envs`.
```bash
uv pip install pufferlib
Current 4.0 config instead uses:
```ini
[vec]
total_agents = 4096
num_buffers = 2
num_threads = 16
```
## Documentation
Read `references/vectorization.md`. Benchmark fixed work with warmup and at least
three repeats; report simulation and end-to-end training SPS separately. The
bundled benchmark measures only its synthetic harness.
- Official docs: https://puffer.ai/docs.html
- GitHub: https://github.com/PufferAI/PufferLib
- Discord: Community support available
## Policy workflow
Published 3.0 policies are Torch modules sized from
`single_observation_space`/`single_action_space`. Stable recurrent composition
uses `encode_observations` and `decode_actions`; structured emulation uses
`pufferlib.pytorch.nativize_dtype` and `nativize_tensor`.
Current 4.0 Torch fallback composes:
```python
pufferlib.models.Policy(encoder=encoder, decoder=decoder, network=network)
```
It provides MLP, MinGRU, LSTM, and GRU network choices; `--slowly` selects this
fallback instead of the native backend. Check output/state shapes, masks,
finite values, gradients, and eager-versus-compiled behavior. See
`references/policies.md`.
## Training and evaluation
Published 3.0 trainer import:
```python
from pufferlib import pufferl
trainer = pufferl.PuffeRL(train_config, vecenv, policy)
```
Current 4.0 CLI:
```bash
puffer train ENV_NAME
puffer eval ENV_NAME --load-model-path EXACT_TRUSTED_PATH
puffer sweep ENV_NAME
```
Generate a plan instead of launching by default:
```bash
python3 scripts/train_template.py \
--profile pypi-3.0.0 \
--environment synthetic \
--device cpu \
--total-timesteps 10000
```
Validate a custom strict-JSON plan:
```bash
python3 scripts/validate_plan.py --root . --config plan.json
```
The schema rejects secret-bearing keys, unbounded resources, dotted environment
paths, invalid vector divisibility, mixed-version options, and coupled
train/eval seeds. See `references/training.md`.
## Logging
PufferLib 3.0 exposes W&B and Neptune; current 4.0 CLI exposes W&B. Both are
optional external services. They may transmit configuration, metrics, source
metadata, hardware telemetry, output, and approved artifacts, with privacy,
retention, access-control, and cost implications.
- W&B credential: named environment variable `WANDB_API_KEY`.
- Neptune credential: named environment variable `NEPTUNE_API_TOKEN`.
- Never put values in arguments/config/logs.
- Sanitize config keys before logging.
- Keep source/model upload off unless explicitly approved.
The planner requires both:
```bash
python3 scripts/train_template.py \
--logger wandb \
--enable-external-logging \
--acknowledge-external-disclosure
```
It reports only the required variable name and never reads its value.
## Checkpoint workflow
PufferLib 3.0 and the 4.0 Torch fallback use Torch serialization; current native
4.0 writes opaque `.bin` weights. PyTorch warns that untrusted models are
programs and that `torch.load` uses unpickling.
```bash
python3 scripts/inspect_checkpoint.py checkpoint.pt \
--root . \
--expected-sha256 0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef
```
The inspector hashes and classifies only. It does not call `torch.load`, import
pickle/Torch, inspect archive members, or extract files. Verify source, license,
architecture, environment revision, sidecar metadata, and checksum before any
sandboxed load. Never use `latest` in a reproducible evaluation.
## Bundled files
### Scripts
- `scripts/env_template.py` — deterministic synthetic Gymnasium-style template.
- `scripts/env_contract_validator.py` — bounded contract and seed checks.
- `scripts/benchmark_vectorization.py` — capped serial/spawn synthetic benchmark.
- `scripts/train_template.py` — non-executing 3.0/4.0 training-plan generator.
- `scripts/validate_plan.py` — strict config/resource/security validator.
- `scripts/inspect_checkpoint.py` — metadata/hash inspection without deserialization.
- `scripts/repro_plan.py` — separate-seed evaluation and benchmark plan.
### References
- `references/environments.md` — Gymnasium, stable PufferEnv, emulation, native C.
- `references/vectorization.md` — backends, shapes, start methods, benchmarks.
- `references/policies.md` — stable/current policy contracts and state safety.
- `references/training.md` — installs, config, CLI, PuffeRL, eval, logs, checkpoints.
- `references/integration.md` — migration matrix, third-party and credential safety.
## Dated upstream sources
- [PyPI pufferlib 3.0.0](https://pypi.org/project/pufferlib/3.0.0/) —
released 2025-06-23; checked 2026-07-23.
- [PyPI 3.0.0 metadata](https://pypi.org/pypi/pufferlib/3.0.0/json) —
digest/dependencies; checked 2026-07-23.
- [PufferLib official docs](https://puffer.ai/docs.html) — current 4.0 docs;
checked 2026-07-23.
- [PufferLib source](https://github.com/PufferAI/PufferLib) — default branch and
implementation; checked 2026-07-23.
- [PufferTank 4.0 Dockerfile](https://github.com/PufferAI/PufferTank/blob/4.0/puffertank.dockerfile)
— CUDA/Python reference; checked 2026-07-23.
- [PufferLib 2.0 paper](https://openreview.net/forum?id=qRyteMTgn0) —
Reinforcement Learning Journal, 2025; use only for its stated benchmarks.
- [PufferLib compatibility paper](https://arxiv.org/abs/2406.12905) —
submitted 2024-06-18; describes an earlier API/performance profile.

662
skills/pufferlib/references/environments.md
View File

@@ -1,508 +1,260 @@
# PufferLib Environments Guide
# Environment Contracts and Native Environments
## Overview
Research snapshot: **2026-07-23**.
PufferLib provides the PufferEnv API for creating high-performance custom environments, and the Ocean suite containing 20+ pre-built environments. Environments support both single-agent and multi-agent scenarios with native vectorization.
## Start from the Gymnasium contract
## PufferEnv API
### Core Characteristics
PufferEnv is designed for performance through in-place operations:
- Observations, actions, and rewards are initialized from a shared buffer object
- All operations happen in-place to avoid creating and copying arrays
- Native support for both single-agent and multi-agent environments
- Flat observation/action spaces for efficient vectorization
### Creating a PufferEnv
A current single-agent Gymnasium environment defines `observation_space` and
`action_space`, then implements:
```python
def reset(self, *, seed=None, options=None):
super().reset(seed=seed)
return observation, info
def step(self, action):
return observation, reward, terminated, truncated, info
```
Contract requirements:
- `observation` must be contained in `observation_space` after reset and every
step, with the documented shape and dtype.
- `action` must be contained in `action_space`.
- `reward` is a finite scalar for ordinary single-agent tasks.
- `terminated` means the task's MDP reached a terminal state.
- `truncated` means an external limit ended the episode, commonly a time limit.
- `info` is a dictionary; never hide the only termination signal in it.
- Call `reset()` after either `terminated` or `truncated`.
- Seed the environment through `reset(seed=...)`. Seed the action space
separately when sampled actions must be reproducible.
- Always call `close()`.
Do not collapse `terminated` and `truncated` during learning. A time-limit
truncation can still permit value bootstrapping; a true terminal state does not.
Run the local contract tool before involving PufferLib:
```bash
python3 scripts/env_contract_validator.py
```
It validates only the bundled synthetic environment. It intentionally has no
module-path option, so it cannot dynamically import an untrusted package.
## Published PufferLib 3.0.0 native contract
For a native Python `PufferEnv`, assign these attributes **before** calling
`super().__init__(buf)`:
```python
import gymnasium
import numpy as np
import pufferlib
from pufferlib import PufferEnv
class MyEnvironment(PufferEnv):
def __init__(self, buf=None):
class ReviewedEnv(pufferlib.PufferEnv):
def __init__(self, buf=None, seed=0):
self.single_observation_space = gymnasium.spaces.Box(
low=-1.0, high=1.0, shape=(4,), dtype=np.float32
)
self.single_action_space = gymnasium.spaces.Discrete(3)
self.num_agents = 2
super().__init__(buf)
# Define observation and action spaces
self.observation_space = self.make_space({
'image': (84, 84, 3),
'vector': (10,)
})
self.action_space = self.make_discrete(4) # 4 discrete actions
# Initialize state
self.reset()
def reset(self):
"""Reset environment to initial state."""
# Reset internal state
self.agent_pos = np.array([0, 0])
self.step_count = 0
# Return initial observation
obs = {
'image': np.zeros((84, 84, 3), dtype=np.uint8),
'vector': np.zeros(10, dtype=np.float32)
}
return obs
def step(self, action):
"""Execute one environment step."""
# Update state based on action
self.step_count += 1
# Calculate reward
reward = self._compute_reward()
# Check if episode is done
done = self.step_count >= 1000
# Generate observation
obs = self._get_observation()
# Additional info
info = {'episode': {'r': reward, 'l': self.step_count}} if done else {}
return obs, reward, done, info
def _compute_reward(self):
"""Compute reward for current state."""
return 1.0
def _get_observation(self):
"""Generate observation from current state."""
return {
'image': np.random.randint(0, 256, (84, 84, 3), dtype=np.uint8),
'vector': np.random.randn(10).astype(np.float32)
}
```
### Observation Spaces
The stable base accepts a Box observation space and Discrete, MultiDiscrete, or
Box action space. It allocates or attaches:
#### Discrete Spaces
- `observations`
- `actions`
- `rewards`
- `terminals`
- `truncations`
- `masks`
Native methods operate on those buffers:
```python
# Single discrete value
self.observation_space = self.make_discrete(10) # Values 0-9
def reset(self, seed=None):
# update self.observations in place
return self.observations, []
# Dict with discrete values
self.observation_space = self.make_space({
'position': (1,), # Continuous
'type': self.make_discrete(5) # Discrete
})
def step(self, actions):
# update all buffers in place
return (
self.observations,
self.rewards,
self.terminals,
self.truncations,
[],
)
```
#### Continuous Spaces
The `infos` value for native Puffer environments is a list of dictionaries.
PufferLib's native interface expects vector rows for agents, even when there is
one agent. Native environments handle their own resets; clear rewards,
terminals, truncations, masks, and partially written observations explicitly.
Never leave a previous step's buffer values in place.
### Native shape checklist
For `A = num_agents` and single observation shape `S`:
- observations: `(A, *S)`
- rewards: `(A,)`
- terminals: `(A,)`
- truncations: `(A,)`
- masks: `(A,)`
- actions: joint shape derived from the single action space and `A`
Validate the exact allocated action shape rather than assuming `(A,)`, especially
for MultiDiscrete and Box actions.
## Stable Gymnasium and PettingZoo adaptation
PufferLib 3.0 uses explicit adapters:
```python
# Box space (continuous)
self.observation_space = self.make_space({
'image': (84, 84, 3), # Image
'vector': (10,), # Vector
'scalar': (1,) # Single value
})
import pufferlib.emulation
wrapped = pufferlib.emulation.GymnasiumPufferEnv(reviewed_gymnasium_instance)
```
#### Multi-Discrete Spaces
or:
```python
# Multiple discrete values
self.observation_space = self.make_multi_discrete([3, 5, 2]) # 3 values, 5 values, 2 values
wrapped = pufferlib.emulation.PettingZooPufferEnv(reviewed_parallel_instance)
```
### Action Spaces
There is no supported 3.0 `pufferlib.emulate(...)` convenience function matching
the old skill examples. Pass either an `env` instance or an `env_creator`
callable according to the class signature; do not pass both.
```python
# Discrete actions
self.action_space = self.make_discrete(4) # 4 actions: 0, 1, 2, 3
The Gymnasium adapter:
# Continuous actions
self.action_space = self.make_space((3,)) # 3D continuous action
- maps structured observation/action spaces to flat arrays;
- checks the first observation and action against the original spaces;
- returns separate terminal and truncation values;
- requires reset before step and reset after episode end.
# Multi-discrete actions
self.action_space = self.make_multi_discrete([3, 3]) # Two 3-way discrete choices
```
The PettingZoo adapter:
## Multi-Agent Environments
- targets the Parallel API;
- uses `possible_agents` as the fixed slot set;
- pads missing agents and exposes masks;
- canonicalizes per-agent spaces and flattened buffers.
PufferLib has native multi-agent support, treating single-agent and multi-agent environments uniformly.
Validate heterogeneous-agent spaces before use. The adapter derives its single
spaces from the first possible agent, so environments with incompatible spaces
need an explicit reviewed transformation.
### Multi-Agent PufferEnv
### Structured spaces
```python
class MultiAgentEnv(PufferEnv):
def __init__(self, num_agents=4, buf=None):
super().__init__(buf)
Stable emulation supports Box, Discrete, MultiDiscrete, Tuple, and Dict patterns
through a packed NumPy dtype. This is byte-layout conversion, not semantic
feature engineering. Check:
self.num_agents = num_agents
- deterministic Dict key order;
- leaf shape and dtype;
- finite numeric values;
- lossless action reconstruction;
- policy-side unflattening;
- padding/mask handling for variable populations.
# Per-agent observation space
self.single_observation_space = self.make_space({
'position': (2,),
'velocity': (2,),
'global': (10,)
})
## Current 4.0 Ocean contract
# Per-agent action space
self.single_action_space = self.make_discrete(5)
The 4.0 default branch focuses on first-party C environments. It no longer
provides the 3.0 Python emulation/vector modules. The official starting points
are:
self.reset()
- `ocean/squared`: commented single-agent template
- `ocean/target`: commented multi-agent template
def reset(self):
"""Reset all agents."""
self.agents = {f'agent_{i}': Agent(i) for i in range(self.num_agents)}
# Return observations for all agents
return {
agent_id: self._get_obs(agent)
for agent_id, agent in self.agents.items()
}
def step(self, actions):
"""Step all agents."""
# actions is a dict: {agent_id: action}
observations = {}
rewards = {}
dones = {}
infos = {}
for agent_id, action in actions.items():
agent = self.agents[agent_id]
# Update agent
agent.update(action)
# Generate results
observations[agent_id] = self._get_obs(agent)
rewards[agent_id] = self._compute_reward(agent)
dones[agent_id] = agent.is_done()
infos[agent_id] = {}
# Check for global done condition
dones['__all__'] = all(dones.values())
return observations, rewards, dones, infos
```
## Ocean Environment Suite
PufferLib provides the Ocean suite with 20+ pre-built environments:
### Available Environments
#### Arcade Games
- **Atari**: Classic Atari 2600 games via Arcade Learning Environment
- **Procgen**: Procedurally generated games for generalization testing
#### Grid-Based
- **Minigrid**: Partially observable gridworld environments
- **Crafter**: Open-ended survival crafting game
- **NetHack**: Classic roguelike dungeon crawler
- **MiniHack**: Simplified NetHack variants
#### Multi-Agent
- **PettingZoo**: Multi-agent environment suite (including Butterfly)
- **MAgent**: Large-scale multi-agent scenarios
- **Neural MMO**: Massively multi-agent survival game
#### Specialized
- **Pokemon Red**: Classic Pokemon game environment
- **GPUDrive**: High-performance driving simulator
- **Griddly**: Grid-based game engine
- **MicroRTS**: Real-time strategy game
### Using Ocean Environments
```python
import pufferlib
# Make environment
env = pufferlib.make('procgen-coinrun', num_envs=256)
# With custom configuration
env = pufferlib.make(
'atari-pong',
num_envs=128,
frameskip=4,
framestack=4
)
# Multi-agent environment
env = pufferlib.make('pettingzoo-knights-archers-zombies', num_agents=4)
```
## Custom Environment Development
### Development Workflow
1. **Prototype in Python**: Start with pure Python PufferEnv
2. **Optimize Critical Paths**: Identify bottlenecks
3. **Implement in C**: Rewrite performance-critical code in C
4. **Create Bindings**: Use Python C API
5. **Compile**: Build as extension module
6. **Register**: Add to Ocean suite
### Performance Benchmarks
- **Pure Python**: 100k-500k steps/second
- **C Implementation**: 100M+ steps/second
- **Training with Python env**: ~400k total SPS
- **Training with C env**: ~4M total SPS
### Python Optimization Tips
```python
# Use NumPy operations instead of Python loops
# Bad
for i in range(len(array)):
array[i] = array[i] * 2
# Good
array *= 2
# Pre-allocate arrays instead of appending
# Bad
observations = []
for i in range(n):
observations.append(generate_obs())
# Good
observations = np.empty((n, obs_shape), dtype=np.float32)
for i in range(n):
observations[i] = generate_obs()
# Use in-place operations
# Bad
new_state = state + delta
# Good
state += delta
```
### C Extension Example
A binding defines compile-time metadata such as:
```c
// my_env.c
#include <Python.h>
#include <numpy/arrayobject.h>
#define OBS_SIZE 121
#define NUM_ATNS 1
#define ACT_SIZES {5}
#define OBS_TENSOR_T ByteTensor
// Fast environment step implementation
static PyObject* fast_step(PyObject* self, PyObject* args) {
PyArrayObject* state;
int action;
if (!PyArg_ParseTuple(args, "O!i", &PyArray_Type, &state, &action)) {
return NULL;
}
// High-performance C implementation
// ...
return Py_BuildValue("Ofi", obs, reward, done);
}
static PyMethodDef methods[] = {
{"fast_step", fast_step, METH_VARARGS, "Fast environment step"},
{NULL, NULL, 0, NULL}
};
static struct PyModuleDef module = {
PyModuleDef_HEAD_INIT,
"my_env_c",
NULL,
-1,
methods
};
PyMODINIT_FUNC PyInit_my_env_c(void) {
import_array();
return PyModule_Create(&module);
}
#define Env Squared
#include "vecenv.h"
```
## Third-Party Environment Integration
The environment struct must include pointers for observations, actions,
rewards, and terminals, plus `num_agents` and a log struct. It implements
`c_reset`, `c_step`, `c_render`, and `c_close`; `binding.c` supplies `my_init`
and `my_log`.
### Gymnasium Environments
Security and correctness rules:
```python
import gymnasium as gym
import pufferlib
1. Treat the C environment and every linked library as native code.
2. Verify repository/commit, license, asset rights, and checksums before build.
3. Build only the selected environment in a disposable container or VM.
4. Start with the local/address-sanitizer build described by upstream.
5. Match `OBS_SIZE`, tensor dtype, action branch count/sizes, and actual writes.
6. Bounds-check every index and allocation; use checked arithmetic for sizes.
7. Initialize every output element each step. Reset reward/terminal buffers
before early returns.
8. Use an environment-owned RNG seeded per instance; do not use global RNG
state for reproducibility.
9. Free only memory owned by the environment. Do not free framework buffers.
10. Fuzz reset/step/action boundaries before optimization.
# Wrap Gymnasium environment
gym_env = gym.make('CartPole-v1')
puffer_env = pufferlib.emulate(gym_env, num_envs=256)
`c_step` may reset immediately after marking a terminal. Record this autoreset
behavior when interpreting terminal observations.
# Or use make directly
env = pufferlib.make('gym-CartPole-v1', num_envs=256)
```
## Environment provenance
### PettingZoo Environments
An environment package may execute arbitrary Python/native code and may fetch
assets at import, build, reset, or render time. Before execution:
```python
from pettingzoo.butterfly import pistonball_v6
import pufferlib
- use the official repository and immutable revision;
- inspect package/build scripts and transitive dependencies;
- verify artifact hashes or attestations;
- review license compatibility for code, datasets, media, ROMs, maps, and model
opponents separately;
- reject unlicensed ROMs or “accept ROM license” automation without proof of
rights;
- disable network and credentials in the first-run sandbox;
- cap disk, memory, processes, threads, episode length, agents, and render size;
- do not load bundled checkpoints or pickle files during environment import.
# Wrap PettingZoo environment
pz_env = pistonball_v6.env()
puffer_env = pufferlib.emulate(pz_env, num_envs=128)
An entry in Ocean/config is not a blanket security, quality, or licensing
approval.
# Or use make directly
env = pufferlib.make('pettingzoo-pistonball', num_envs=128)
```
## Testing ladder
### Custom Wrappers
1. Built-in synthetic contract validator.
2. One environment, one seed, serial, tens of steps.
3. Boundary actions and intentionally invalid actions.
4. Termination and time-limit truncation tests.
5. Same-seed trace comparison.
6. Independent-seed diversity check.
7. Structured-space round trip.
8. Multi-agent join/leave and mask tests.
9. Serial versus vectorized trace equivalence where ordering permits.
10. Bounded throughput benchmark only after correctness passes.
```python
class CustomWrapper(pufferlib.PufferEnv):
"""Wrapper to modify environment behavior."""
## Sources
def __init__(self, base_env, buf=None):
super().__init__(buf)
self.base_env = base_env
self.observation_space = base_env.observation_space
self.action_space = base_env.action_space
def reset(self):
obs = self.base_env.reset()
# Modify observation
return self._process_obs(obs)
def step(self, action):
# Modify action
modified_action = self._process_action(action)
obs, reward, done, info = self.base_env.step(modified_action)
# Modify outputs
obs = self._process_obs(obs)
reward = self._process_reward(reward)
return obs, reward, done, info
```
## Environment Best Practices
### State Management
```python
# Store minimal state, compute on demand
class EfficientEnv(PufferEnv):
def __init__(self, buf=None):
super().__init__(buf)
self.agent_pos = np.zeros(2) # Minimal state
def _get_observation(self):
# Compute full observation on demand
observation = np.zeros((84, 84, 3), dtype=np.uint8)
self._render_scene(observation, self.agent_pos)
return observation
```
### Reward Scaling
```python
# Normalize rewards to reasonable range
def step(self, action):
# ... environment logic ...
# Scale large rewards
raw_reward = compute_raw_reward()
reward = np.clip(raw_reward / 100.0, -10, 10)
return obs, reward, done, info
```
### Episode Termination
```python
def step(self, action):
# ... environment logic ...
# Multiple termination conditions
timeout = self.step_count >= self.max_steps
success = self._check_success()
failure = self._check_failure()
done = timeout or success or failure
info = {
'TimeLimit.truncated': timeout,
'success': success
}
return obs, reward, done, info
```
### Memory Efficiency
```python
# Reuse buffers instead of allocating new ones
class MemoryEfficientEnv(PufferEnv):
def __init__(self, buf=None):
super().__init__(buf)
# Pre-allocate observation buffer
self._obs_buffer = np.zeros((84, 84, 3), dtype=np.uint8)
def _get_observation(self):
# Reuse buffer, modify in place
self._render_scene(self._obs_buffer)
return self._obs_buffer # Return view, not copy
```
## Debugging Environments
### Validation Checks
```python
# Add assertions to catch bugs
def step(self, action):
assert self.action_space.contains(action), f"Invalid action: {action}"
obs, reward, done, info = self._step_impl(action)
assert self.observation_space.contains(obs), "Invalid observation"
assert np.isfinite(reward), "Non-finite reward"
return obs, reward, done, info
```
### Rendering
```python
class DebuggableEnv(PufferEnv):
def __init__(self, buf=None, render_mode=None):
super().__init__(buf)
self.render_mode = render_mode
def render(self):
"""Render environment for debugging."""
if self.render_mode == 'human':
# Display to screen
self._display_scene()
elif self.render_mode == 'rgb_array':
# Return image
return self._render_to_array()
```
### Logging
```python
import logging
logger = logging.getLogger(__name__)
def step(self, action):
logger.debug(f"Step {self.step_count}: action={action}")
obs, reward, done, info = self._step_impl(action)
if done:
logger.info(f"Episode finished: reward={self.total_reward}")
return obs, reward, done, info
```
- [Gymnasium Env API](https://gymnasium.farama.org/api/env/) — current reset,
step, spaces, and seeding contract; accessed 2026-07-23.
- [Gymnasium terminated/truncated explanation](https://farama.org/Gymnasium-Terminated-Truncated-Step-API)
— published 2023-10-27; accessed 2026-07-23.
- [PufferLib 3.0 core environment source](https://github.com/PufferAI/PufferLib/blob/3.0/pufferlib/pufferlib.py)
— stable native contract; accessed 2026-07-23.
- [PufferLib 3.0 emulation source](https://github.com/PufferAI/PufferLib/blob/3.0/pufferlib/emulation.py)
— stable adapters; accessed 2026-07-23.
- [PufferLib 3.0 Gymnasium example](https://github.com/PufferAI/PufferLib/blob/3.0/examples/gymnasium_env.py)
— stable example; accessed 2026-07-23.
- [PufferLib 3.0 PettingZoo example](https://github.com/PufferAI/PufferLib/blob/3.0/examples/pettingzoo_env.py)
— stable example; accessed 2026-07-23.
- [PufferLib 4.0 Squared template](https://github.com/PufferAI/PufferLib/tree/4.0/ocean/squared)
— current single-agent native template; accessed 2026-07-23.
- [PufferLib 4.0 Target template](https://github.com/PufferAI/PufferLib/tree/4.0/ocean/target)
— current multi-agent native template; accessed 2026-07-23.
- [PufferLib Ocean](https://puffer.ai/ocean.html) — current first-party
collection; accessed 2026-07-23.

779
skills/pufferlib/references/integration.md
View File

@@ -1,621 +1,192 @@
# PufferLib Integration Guide
# Integration, Security, and Migration Guide
## Overview
Research snapshot: **2026-07-23**.
PufferLib provides an emulation layer that enables seamless integration with popular RL frameworks including Gymnasium, OpenAI Gym, PettingZoo, and many specialized environment libraries. The emulation layer flattens observation and action spaces for efficient vectorization while maintaining compatibility.
## Compatibility matrix
## Gymnasium Integration
| Need | Published `pufferlib==3.0.0` | Current `4.0` source |
|---|---|---|
| Gymnasium instance adaptation | `pufferlib.emulation.GymnasiumPufferEnv` | Removed from current source |
| PettingZoo Parallel adaptation | `pufferlib.emulation.PettingZooPufferEnv` | Removed from current source |
| Python vector backends | `pufferlib.vector` | Removed from current source |
| Native Python `PufferEnv` | Supported | Replaced by current C/Ocean interface |
| Trainer | `pufferlib.pufferl.PuffeRL` | Native backend or `pufferlib.torch_pufferl.PuffeRL` |
| External logging | W&B and Neptune | W&B in current CLI |
| Primary config | merged INI sections | different INI schema |
| Checkpoints | Torch state dict plus trainer state | native `.bin`; Torch fallback state dict |
### Basic Gymnasium Environments
Pin a profile. Do not import from a floating branch or blend examples across
columns.
## Correct stable adaptation patterns
### Gymnasium
```python
import gymnasium as gym
import pufferlib
import gymnasium
import pufferlib.emulation
import pufferlib.vector
# Method 1: Direct wrapping
gym_env = gym.make('CartPole-v1')
puffer_env = pufferlib.emulate(gym_env, num_envs=256)
# Method 2: Using make
env = pufferlib.make('gym-CartPole-v1', num_envs=256)
def make_env():
raw = gymnasium.make("CartPole-v1")
return pufferlib.emulation.GymnasiumPufferEnv(raw)
# Method 3: Custom Gymnasium environment
class MyGymEnv(gym.Env):
def __init__(self):
self.observation_space = gym.spaces.Box(low=-1, high=1, shape=(4,))
self.action_space = gym.spaces.Discrete(2)
def reset(self, seed=None, options=None):
super().reset(seed=seed)
return self.observation_space.sample(), {}
def step(self, action):
obs = self.observation_space.sample()
reward = 1.0
terminated = False
truncated = False
info = {}
return obs, reward, terminated, truncated, info
# Wrap custom environment
puffer_env = pufferlib.emulate(MyGymEnv, num_envs=128)
```
### Atari Environments
```python
import gymnasium as gym
from gymnasium.wrappers import AtariPreprocessing, FrameStack
import pufferlib
# Standard Atari setup
def make_atari_env(env_name='ALE/Pong-v5'):
env = gym.make(env_name)
env = AtariPreprocessing(env, frame_skip=4)
env = FrameStack(env, num_stack=4)
return env
# Vectorize with PufferLib
env = pufferlib.emulate(make_atari_env, num_envs=256)
# Or use built-in
env = pufferlib.make('atari-pong', num_envs=256, frameskip=4, framestack=4)
```
### Complex Observation Spaces
```python
import gymnasium as gym
from gymnasium.spaces import Dict, Box, Discrete
import pufferlib
class ComplexObsEnv(gym.Env):
def __init__(self):
# Dict observation space
self.observation_space = Dict({
'image': Box(low=0, high=255, shape=(84, 84, 3), dtype=np.uint8),
'vector': Box(low=-np.inf, high=np.inf, shape=(10,), dtype=np.float32),
'discrete': Discrete(5)
})
self.action_space = Discrete(4)
def reset(self, seed=None, options=None):
return {
'image': np.zeros((84, 84, 3), dtype=np.uint8),
'vector': np.zeros(10, dtype=np.float32),
'discrete': 0
}, {}
def step(self, action):
obs = {
'image': np.random.randint(0, 256, (84, 84, 3), dtype=np.uint8),
'vector': np.random.randn(10).astype(np.float32),
'discrete': np.random.randint(0, 5)
}
return obs, 1.0, False, False, {}
# PufferLib automatically flattens and unflattens complex spaces
env = pufferlib.emulate(ComplexObsEnv, num_envs=128)
```
## PettingZoo Integration
### Parallel Environments
```python
from pettingzoo.butterfly import pistonball_v6
import pufferlib
# Wrap PettingZoo parallel environment
pz_env = pistonball_v6.parallel_env()
puffer_env = pufferlib.emulate(pz_env, num_envs=128)
# Or use make directly
env = pufferlib.make('pettingzoo-pistonball', num_envs=128)
```
### AEC (Agent Environment Cycle) Environments
```python
from pettingzoo.classic import chess_v5
import pufferlib
# Wrap AEC environment (PufferLib handles conversion to parallel)
aec_env = chess_v5.env()
puffer_env = pufferlib.emulate(aec_env, num_envs=64)
# Works with any PettingZoo AEC environment
env = pufferlib.make('pettingzoo-chess', num_envs=64)
```
### Multi-Agent Training
```python
import pufferlib
from pufferlib import PuffeRL
# Create multi-agent environment
env = pufferlib.make('pettingzoo-knights-archers-zombies', num_envs=128)
# Shared policy for all agents
policy = create_policy(env.observation_space, env.action_space)
# Train
trainer = PuffeRL(env=env, policy=policy)
for iteration in range(num_iterations):
# Observations are dicts: {agent_id: batch_obs}
rollout = trainer.evaluate()
# Train on multi-agent data
trainer.train()
trainer.mean_and_log()
```
## Third-Party Environments
### Procgen
```python
import pufferlib
# Procgen environments
env = pufferlib.make('procgen-coinrun', num_envs=256, distribution_mode='easy')
# Custom configuration
env = pufferlib.make(
'procgen-coinrun',
num_envs=256,
num_levels=200, # Number of unique levels
start_level=0, # Starting level seed
distribution_mode='hard'
vecenv = pufferlib.vector.make(
make_env,
backend=pufferlib.vector.Serial,
num_envs=2,
seed=42,
)
try:
observations, infos = vecenv.reset(seed=42)
actions = vecenv.action_space.sample()
observations, rewards, terminals, truncations, infos = vecenv.step(actions)
finally:
vecenv.close()
```
### NetHack
This example is an API pattern, not authorization to install or execute
`CartPole-v1` or another plug-in. Review the exact environment and dependencies
first.
### PettingZoo
Use a reviewed Parallel environment instance:
```python
import pufferlib
# NetHack Learning Environment
env = pufferlib.make('nethack', num_envs=128)
# MiniHack variants
env = pufferlib.make('minihack-corridor', num_envs=128)
env = pufferlib.make('minihack-room', num_envs=128)
wrapped = pufferlib.emulation.PettingZooPufferEnv(reviewed_parallel_env)
```
### Minigrid
```python
import pufferlib
# Minigrid environments
env = pufferlib.make('minigrid-empty-8x8', num_envs=256)
env = pufferlib.make('minigrid-doorkey-8x8', num_envs=256)
env = pufferlib.make('minigrid-multiroom', num_envs=256)
```
### Neural MMO
```python
import pufferlib
# Large-scale multi-agent environment
env = pufferlib.make(
'neuralmmo',
num_envs=64,
num_agents=128, # Agents per environment
map_size=128
)
```
### Crafter
```python
import pufferlib
# Open-ended crafting environment
env = pufferlib.make('crafter', num_envs=128)
```
### GPUDrive
```python
import pufferlib
# GPU-accelerated driving simulator
env = pufferlib.make(
'gpudrive',
num_envs=1024, # Can handle many environments on GPU
num_vehicles=8
)
```
### MicroRTS
```python
import pufferlib
# Real-time strategy game
env = pufferlib.make(
'microrts',
num_envs=128,
map_size=16,
max_steps=2000
)
```
### Griddly
```python
import pufferlib
# Grid-based games
env = pufferlib.make('griddly-clusters', num_envs=256)
env = pufferlib.make('griddly-sokoban', num_envs=256)
```
## Custom Wrappers
### Observation Wrappers
```python
import numpy as np
import pufferlib
from pufferlib import PufferEnv
class NormalizeObservations(pufferlib.Wrapper):
"""Normalize observations to zero mean and unit variance."""
def __init__(self, env):
super().__init__(env)
self.obs_mean = np.zeros(env.observation_space.shape)
self.obs_std = np.ones(env.observation_space.shape)
self.count = 0
def reset(self):
obs = self.env.reset()
return self._normalize(obs)
def step(self, action):
obs, reward, done, info = self.env.step(action)
return self._normalize(obs), reward, done, info
def _normalize(self, obs):
# Update running statistics
self.count += 1
delta = obs - self.obs_mean
self.obs_mean += delta / self.count
self.obs_std = np.sqrt(((self.count - 1) * self.obs_std ** 2 + delta * (obs - self.obs_mean)) / self.count)
# Normalize
return (obs - self.obs_mean) / (self.obs_std + 1e-8)
```
### Reward Wrappers
```python
class RewardShaping(pufferlib.Wrapper):
"""Add shaped rewards to environment."""
def __init__(self, env, shaping_fn):
super().__init__(env)
self.shaping_fn = shaping_fn
def step(self, action):
obs, reward, done, info = self.env.step(action)
# Add shaped reward
shaped_reward = reward + self.shaping_fn(obs, action)
return obs, shaped_reward, done, info
# Usage
def proximity_shaping(obs, action):
"""Reward agent for getting closer to goal."""
goal_pos = np.array([10, 10])
agent_pos = obs[:2]
distance = np.linalg.norm(goal_pos - agent_pos)
return -0.1 * distance
env = pufferlib.make('myenv', num_envs=128)
env = RewardShaping(env, proximity_shaping)
```
### Frame Stacking
```python
class FrameStack(pufferlib.Wrapper):
"""Stack frames for temporal context."""
def __init__(self, env, num_stack=4):
super().__init__(env)
self.num_stack = num_stack
self.frames = None
def reset(self):
obs = self.env.reset()
# Initialize frame stack
self.frames = np.repeat(obs[np.newaxis], self.num_stack, axis=0)
return self._get_obs()
def step(self, action):
obs, reward, done, info = self.env.step(action)
# Update frame stack
self.frames = np.roll(self.frames, shift=-1, axis=0)
self.frames[-1] = obs
if done:
self.frames = None
return self._get_obs(), reward, done, info
def _get_obs(self):
return self.frames
```
### Action Repeat
```python
class ActionRepeat(pufferlib.Wrapper):
"""Repeat actions for multiple steps."""
def __init__(self, env, repeat=4):
super().__init__(env)
self.repeat = repeat
def step(self, action):
total_reward = 0.0
done = False
for _ in range(self.repeat):
obs, reward, done, info = self.env.step(action)
total_reward += reward
if done:
break
return obs, total_reward, done, info
```
## Space Conversion
### Flattening Spaces
PufferLib automatically flattens complex observation/action spaces:
```python
from gymnasium.spaces import Dict, Box, Discrete
import pufferlib
# Complex space
original_space = Dict({
'image': Box(0, 255, (84, 84, 3), dtype=np.uint8),
'vector': Box(-np.inf, np.inf, (10,), dtype=np.float32),
'discrete': Discrete(5)
})
# Automatically flattened by PufferLib
# Observations are presented as flat arrays for efficient processing
# But can be unflattened when needed for policy processing
```
### Unflattening for Policies
```python
from pufferlib.pytorch import unflatten_observations
class PolicyWithUnflatten(nn.Module):
def __init__(self, observation_space, action_space):
super().__init__()
self.observation_space = observation_space
# ... policy architecture ...
def forward(self, flat_observations):
# Unflatten to original structure
observations = unflatten_observations(
flat_observations,
self.observation_space
)
# Now observations is a dict with 'image', 'vector', 'discrete'
image_features = self.image_encoder(observations['image'])
vector_features = self.vector_encoder(observations['vector'])
# ...
```
## Environment Registration
### Registering Custom Environments
```python
import pufferlib
# Register environment for easy access
pufferlib.register(
id='my-custom-env',
entry_point='my_package.envs:MyEnvironment',
kwargs={'param1': 'value1'}
)
# Now can use with make
env = pufferlib.make('my-custom-env', num_envs=256)
```
### Registering in Ocean Suite
To add your environment to Ocean:
```python
# In ocean/environment.py
OCEAN_REGISTRY = {
'my-env': {
'entry_point': 'my_package.envs:MyEnvironment',
'kwargs': {
'default_param': 'default_value'
}
}
}
```
## Compatibility Patterns
### Gymnasium to PufferLib
```python
import gymnasium as gym
import pufferlib
# Standard Gymnasium environment
class GymEnv(gym.Env):
def reset(self, seed=None, options=None):
return observation, info
def step(self, action):
return observation, reward, terminated, truncated, info
# Convert to PufferEnv
puffer_env = pufferlib.emulate(GymEnv, num_envs=128)
```
### PettingZoo to PufferLib
```python
from pettingzoo import ParallelEnv
import pufferlib
# PettingZoo parallel environment
class PZEnv(ParallelEnv):
def reset(self, seed=None, options=None):
return {agent: obs for agent, obs in ...}, {agent: info for agent in ...}
def step(self, actions):
return observations, rewards, terminations, truncations, infos
# Convert to PufferEnv
puffer_env = pufferlib.emulate(PZEnv, num_envs=128)
```
### Legacy Gym (v0.21) to PufferLib
```python
import gym # Old gym
import pufferlib
# Legacy gym environment (returns done instead of terminated/truncated)
class LegacyEnv(gym.Env):
def reset(self):
return observation
def step(self, action):
return observation, reward, done, info
# PufferLib handles legacy format automatically
puffer_env = pufferlib.emulate(LegacyEnv, num_envs=128)
```
## Performance Considerations
### Efficient Integration
```python
# Fast: Use built-in integrations when available
env = pufferlib.make('procgen-coinrun', num_envs=256)
# Slower: Generic wrapper (still fast, but overhead)
import gymnasium as gym
gym_env = gym.make('CartPole-v1')
env = pufferlib.emulate(gym_env, num_envs=256)
# Slowest: Nested wrappers add overhead
import gymnasium as gym
gym_env = gym.make('CartPole-v1')
gym_env = SomeWrapper(gym_env)
gym_env = AnotherWrapper(gym_env)
env = pufferlib.emulate(gym_env, num_envs=256)
```
### Minimize Wrapper Overhead
```python
# BAD: Too many wrappers
env = gym.make('CartPole-v1')
env = Wrapper1(env)
env = Wrapper2(env)
env = Wrapper3(env)
puffer_env = pufferlib.emulate(env, num_envs=256)
# GOOD: Combine wrapper logic
class CombinedWrapper(gym.Wrapper):
def step(self, action):
obs, reward, done, truncated, info = self.env.step(action)
# Apply all transformations at once
obs = self._transform_obs(obs)
reward = self._transform_reward(reward)
return obs, reward, done, truncated, info
env = gym.make('CartPole-v1')
env = CombinedWrapper(env)
puffer_env = pufferlib.emulate(env, num_envs=256)
```
## Debugging Integration
### Verify Environment Compatibility
```python
def test_environment(env, num_steps=100):
"""Test environment for common issues."""
# Test reset
obs = env.reset()
assert env.observation_space.contains(obs), "Invalid initial observation"
# Test steps
for _ in range(num_steps):
action = env.action_space.sample()
obs, reward, done, info = env.step(action)
assert env.observation_space.contains(obs), "Invalid observation"
assert isinstance(reward, (int, float)), "Invalid reward type"
assert isinstance(done, bool), "Invalid done type"
assert isinstance(info, dict), "Invalid info type"
if done:
obs = env.reset()
print("✓ Environment passed compatibility test")
# Test before vectorizing
test_environment(MyEnvironment())
```
### Compare Outputs
```python
# Verify PufferLib emulation matches original
import gymnasium as gym
import pufferlib
import numpy as np
gym_env = gym.make('CartPole-v1')
puffer_env = pufferlib.emulate(lambda: gym.make('CartPole-v1'), num_envs=1)
# Test with same seed
gym_env.reset(seed=42)
puffer_obs = puffer_env.reset()
for _ in range(100):
action = gym_env.action_space.sample()
gym_obs, gym_reward, gym_done, gym_truncated, gym_info = gym_env.step(action)
puffer_obs, puffer_reward, puffer_done, puffer_info = puffer_env.step(np.array([action]))
# Compare outputs (accounting for batch dimension)
assert np.allclose(gym_obs, puffer_obs[0])
assert gym_reward == puffer_reward[0]
assert gym_done == puffer_done[0]
```
The stable source does not document automatic AEC-to-Parallel conversion in
this adapter. Convert explicitly with PettingZoo's supported utilities only
when the environment's turn semantics permit it, then test action ordering,
dead-agent handling, masks, and termination/truncation dictionaries.
### Native stable environment
Subclass `pufferlib.PufferEnv`, define `single_observation_space`,
`single_action_space`, and `num_agents` before `super().__init__`, then update
the provided arrays in place. Native Puffer environments are already vector
interfaces; do not return Gym's scalar four-tuple.
## Unsupported shortcuts from the old skill
Remove or migrate these historical patterns:
| Historical pattern | Current guidance |
|---|---|
| `pufferlib.make("name", ...)` | Stable: import an audited creator and use `pufferlib.vector.make`; 4.0: build/configure a named native environment |
| `pufferlib.emulate(...)` | Stable: instantiate `GymnasiumPufferEnv` or `PettingZooPufferEnv` explicitly |
| `pufferlib.vectorization.Serial` | Stable module is `pufferlib.vector.Serial` |
| `from pufferlib import PuffeRL` | Stable trainer is `pufferlib.pufferl.PuffeRL`; 4.0 fallback is in `torch_pufferl` |
| define native `observation_space`/`action_space` | Stable native class requires `single_observation_space`/`single_action_space` before `super()` |
| return `(obs, reward, done, info)` | Return separate termination and truncation values |
| native multi-agent dictionaries and `dones["__all__"]` | Use stable vector buffers or a reviewed PettingZoo Parallel adapter |
| arbitrary dotted `entry_point` registration | Import an audited callable directly; bundled tools reject dotted paths |
| top-level `WandbLogger`/`NeptuneLogger` | Stable logger classes live in `pufferlib.pufferl`; prefer the CLI and sanitized config |
| assume Atari/Procgen/NetHack names exist everywhere | Verify the chosen version's config/source and install the separately reviewed environment |
## Migrating 3.0 to 4.0
This is a redesign, not a drop-in upgrade:
1. Preserve the 3.0 lock, source digest, config, checkpoint hashes, and baseline
evaluation before changing anything.
2. Inventory use of `emulation`, `vector`, `PufferEnv`, third-party environments,
policy wrappers, INI keys, logger flags, and Torch checkpoints.
3. Decide whether the application should stay on published 3.0.0 or port to a
native 4.0 C environment. The current docs say the Python/third-party layer
was removed from 4.0.
4. Port environment logic to the reviewed Squared/Target C binding contract.
5. Recreate configuration using 4.0 `[vec]`, `[policy]`, `[torch]`, and `[train]`
keys. Do not mechanically rename old keys.
6. Rebuild policy composition around 4.0 encoder/decoder/network modules or the
native backend.
7. Treat old `.pt` and new `.bin` files as incompatible unless an official,
tested converter says otherwise. Do not improvise binary conversion.
8. Re-run contract, same-seed trace, throughput, and held-out learning
baselines. Attribute behavior changes; do not compare headline SPS alone.
The default branch contains some stale 3.0-style examples even though the
corresponding modules are absent. Prefer current implementation and docs over
those copied examples.
## Third-party environments and native code
Environment extras can pull old Gym versions, native libraries, renderers,
emulators, datasets, model opponents, and ROM tooling. A package name in a
PufferLib optional extra is not a security or license endorsement.
Before install/import/build:
1. Identify the official repository and immutable revision.
2. Read build/install hooks and all network downloads.
3. Verify licenses for code and assets separately.
4. Verify hashes/attestations; record missing provenance.
5. Use a disposable sandbox without credentials, home-directory mounts, or
network after required artifacts are staged.
6. Cap processes, threads, memory, disk, render resolution, agents, and steps.
7. Do not execute bundled native extensions, ROMs, checkpoints, or pickle files
until separately trusted.
For Atari and similar systems, the user must supply legally obtained assets.
Never download ROM sets or auto-accept a license on the user's behalf.
## Logging integration
External tracking is disabled by default. The stable logger implementations can
log the full argument mapping and can upload model artifacts. Therefore:
- sanitize arguments before logger construction;
- keep `WANDB_API_KEY` and `NEPTUNE_API_TOKEN` only in an approved environment
injection or secret manager;
- never add credential keys to nested INI/JSON/config objects;
- do not pass a token on the command line;
- disable model/source upload unless explicitly approved;
- review project visibility, retention, residency, access controls, and cost;
- use vendor offline/disabled modes only after confirming what is written
locally and how later sync behaves.
Do not print all environment variables or recursively discover `.env` files.
Checking whether one explicitly named credential variable exists can be
acceptable; reading or logging its value is not.
## Integration acceptance test
For each reviewed environment/profile:
1. Create one instance without network or GPU.
2. Validate spaces and reset return.
3. Step a fixed action trace until both ordinary and episode-end paths run.
4. Verify terminated/truncated semantics and final observation behavior.
5. Close and confirm no child processes/resources remain.
6. Run stable Serial or one 4.0 local native instance.
7. Compare a same-seed trace.
8. Scale to two workers/threads with small caps.
9. Run policy shape and finite-value checks.
10. Run held-out evaluation with logging still disabled.
Only then consider GPU training, external logging, or larger parallelism.
## Sources
- [PufferLib PyPI 3.0.0](https://pypi.org/project/pufferlib/3.0.0/) —
published 2025-06-23; accessed 2026-07-23.
- [PufferLib 3.0 emulation source](https://github.com/PufferAI/PufferLib/blob/3.0/pufferlib/emulation.py)
— stable adapters; accessed 2026-07-23.
- [PufferLib 3.0 vector source](https://github.com/PufferAI/PufferLib/blob/3.0/pufferlib/vector.py)
— stable vector API; accessed 2026-07-23.
- [PufferLib 3.0 trainer source](https://github.com/PufferAI/PufferLib/blob/3.0/pufferlib/pufferl.py)
— stable logger/checkpoint behavior; accessed 2026-07-23.
- [PufferLib 4.0 package tree](https://github.com/PufferAI/PufferLib/tree/4.0/pufferlib)
— current modules; accessed 2026-07-23.
- [PufferLib 4.0 docs](https://puffer.ai/docs.html) — current architecture and
removal note; accessed 2026-07-23.
- [PufferLib releases](https://github.com/PufferAI/PufferLib/releases) —
checked for source releases on 2026-07-23.
- [Gymnasium Env API](https://gymnasium.farama.org/api/env/) — current
single-agent contract; accessed 2026-07-23.
- [PyTorch security policy](https://github.com/pytorch/pytorch/security) —
model/native-package safety; accessed 2026-07-23.

762
skills/pufferlib/references/policies.md
View File

@@ -1,653 +1,181 @@
# PufferLib Policies Guide
# Policies and Model Contracts
## Overview
Research snapshot: **2026-07-23**. Policy APIs changed substantially between
published PufferLib 3.0.0 and current 4.0 source.
PufferLib policies are standard PyTorch modules with optional utilities for observation processing and LSTM integration. The framework provides default architectures and tools while allowing full flexibility in policy design.
## Published 3.0.0
## Policy Architecture
PufferLib 3.0 policies are ordinary `torch.nn.Module` objects. The environment
exposes `single_observation_space` and `single_action_space`; size heads from
those single-agent spaces, not from the batched spaces.
### Basic Policy Structure
### Minimal feed-forward policy
Build an `nn.Module` with an encoder sized from
`env.single_observation_space.shape`, an action head sized from
`env.single_action_space`, and a one-value critic head. The official stable
example defines a rollout method named `forward_eval(observations, state=None)`
and makes the normal `forward` method use the same contract. Here, `forward_eval`
is a PufferLib/PyTorch method name; it does **not** invoke Python's dangerous
`eval()` builtin.
For a discrete action space, the first output contains action logits and the
second is the value estimate. Preserve the leading agent-batch dimension.
### Recurrent composition
The stable `pufferlib.models.LSTMWrapper` expects a base policy with:
```python
import torch
import torch.nn as nn
from pufferlib.pytorch import layer_init
def encode_observations(self, observations, state=None):
...
class BasicPolicy(nn.Module):
def __init__(self, observation_space, action_space):
super().__init__()
self.observation_space = observation_space
self.action_space = action_space
# Encoder network
self.encoder = nn.Sequential(
layer_init(nn.Linear(observation_space.shape[0], 256)),
nn.ReLU(),
layer_init(nn.Linear(256, 256)),
nn.ReLU()
)
# Policy head (actor)
self.actor = layer_init(nn.Linear(256, action_space.n), std=0.01)
# Value head (critic)
self.critic = layer_init(nn.Linear(256, 1), std=1.0)
def forward(self, observations):
"""Forward pass through policy."""
# Encode observations
features = self.encoder(observations)
# Get action logits and value
logits = self.actor(features)
value = self.critic(features)
return logits, value
def get_action(self, observations, deterministic=False):
"""Sample action from policy."""
logits, value = self.forward(observations)
if deterministic:
action = logits.argmax(dim=-1)
else:
dist = torch.distributions.Categorical(logits=logits)
action = dist.sample()
return action, value
def decode_actions(self, hidden):
...
```
### Layer Initialization
The wrapper uses an `LSTMCell` during rollout inference and an `LSTM` over
time-batched data during training. Do not manually reshape recurrent state
without checking the source's batch/time convention. Reset hidden state on
actual terminations and truncations according to the trainer's mask behavior.
PufferLib provides `layer_init` for proper weight initialization:
### Structured observations
Stable emulation flattens `Dict` and `Tuple` spaces into a homogeneous array.
The byte layout is described by `env.emulated`. In policy setup:
```python
from pufferlib.pytorch import layer_init
# Default orthogonal initialization
layer = layer_init(nn.Linear(256, 256))
# Custom standard deviation
actor_head = layer_init(nn.Linear(256, num_actions), std=0.01)
critic_head = layer_init(nn.Linear(256, 1), std=1.0)
# Works with any layer type
conv = layer_init(nn.Conv2d(3, 32, kernel_size=8, stride=4))
native_dtype = pufferlib.pytorch.nativize_dtype(env.emulated)
```
## CNN Policies
For image-based observations:
In the forward pass:
```python
class CNNPolicy(nn.Module):
def __init__(self, observation_space, action_space):
super().__init__()
# CNN encoder for images
self.encoder = nn.Sequential(
layer_init(nn.Conv2d(3, 32, kernel_size=8, stride=4)),
nn.ReLU(),
layer_init(nn.Conv2d(32, 64, kernel_size=4, stride=2)),
nn.ReLU(),
layer_init(nn.Conv2d(64, 64, kernel_size=3, stride=1)),
nn.ReLU(),
nn.Flatten(),
layer_init(nn.Linear(64 * 7 * 7, 512)),
nn.ReLU()
)
self.actor = layer_init(nn.Linear(512, action_space.n), std=0.01)
self.critic = layer_init(nn.Linear(512, 1), std=1.0)
def forward(self, observations):
# Normalize pixel values
x = observations.float() / 255.0
features = self.encoder(x)
logits = self.actor(features)
value = self.critic(features)
return logits, value
structured = pufferlib.pytorch.nativize_tensor(observations, native_dtype)
```
### Efficient CNN Architecture
Keep the original flattened dtype. Constructing a new float tensor before
unflattening can destroy the packed representation. Validate every recovered
leaf shape and dtype before training.
### Action spaces
The 3.0 source handles:
- `Discrete`: one categorical logits tensor.
- `MultiDiscrete`: one logits tensor per action branch.
- `Box`: a Normal distribution path for continuous actions.
Do not infer support from the 2024 paper's limitations section; that paper
describes an earlier release. Test clipping/scaling against the environment's
actual `Box.low`, `Box.high`, shape, and dtype. A `tanh` output is not a general
substitute for affine mapping to arbitrary bounds.
### Stable model utilities
Useful 3.0 symbols include:
- `pufferlib.pytorch.layer_init`
- `pufferlib.pytorch.nativize_dtype`
- `pufferlib.pytorch.nativize_tensor`
- `pufferlib.models.Default`
- `pufferlib.models.LSTMWrapper`
- `pufferlib.models.Convolutional`
- `pufferlib.models.ProcgenResnet`
Inspect the exact 3.0 source before copying signatures. Do not use top-level
`from pufferlib import PuffeRL`; the trainer is
`pufferlib.pufferl.PuffeRL`.
## Current 4.0 source
The current PyTorch fallback composes a policy from three modules:
```python
class EfficientCNN(nn.Module):
"""Optimized CNN for Atari-style games."""
def __init__(self, observation_space, action_space):
super().__init__()
in_channels = observation_space.shape[0] # Typically 4 for framestack
self.network = nn.Sequential(
layer_init(nn.Conv2d(in_channels, 32, 8, stride=4)),
nn.ReLU(),
layer_init(nn.Conv2d(32, 64, 4, stride=2)),
nn.ReLU(),
layer_init(nn.Conv2d(64, 64, 3, stride=1)),
nn.ReLU(),
nn.Flatten()
)
# Calculate feature size
with torch.no_grad():
sample = torch.zeros(1, *observation_space.shape)
n_features = self.network(sample).shape[1]
self.fc = layer_init(nn.Linear(n_features, 512))
self.actor = layer_init(nn.Linear(512, action_space.n), std=0.01)
self.critic = layer_init(nn.Linear(512, 1), std=1.0)
def forward(self, x):
x = x.float() / 255.0
x = self.network(x)
x = torch.relu(self.fc(x))
return self.actor(x), self.critic(x)
```
## Recurrent Policies (LSTM)
PufferLib provides optimized LSTM integration with automatic recurrence handling:
```python
from pufferlib.pytorch import LSTMWrapper
class RecurrentPolicy(nn.Module):
def __init__(self, observation_space, action_space, hidden_size=256):
super().__init__()
# Observation encoder
self.encoder = nn.Sequential(
layer_init(nn.Linear(observation_space.shape[0], 128)),
nn.ReLU()
)
# LSTM layer
self.lstm = nn.LSTM(128, hidden_size, num_layers=1)
# Policy and value heads
self.actor = layer_init(nn.Linear(hidden_size, action_space.n), std=0.01)
self.critic = layer_init(nn.Linear(hidden_size, 1), std=1.0)
# Hidden state
self.hidden_size = hidden_size
def forward(self, observations, state=None):
"""
Args:
observations: (batch, obs_dim)
state: Optional (h, c) tuple for LSTM
Returns:
logits, value, new_state
"""
batch_size = observations.shape[0]
# Encode observations
features = self.encoder(observations)
# Initialize hidden state if needed
if state is None:
h = torch.zeros(1, batch_size, self.hidden_size, device=features.device)
c = torch.zeros(1, batch_size, self.hidden_size, device=features.device)
state = (h, c)
# LSTM forward
features = features.unsqueeze(0) # Add sequence dimension
lstm_out, new_state = self.lstm(features, state)
lstm_out = lstm_out.squeeze(0)
# Get outputs
logits = self.actor(lstm_out)
value = self.critic(lstm_out)
return logits, value, new_state
```
### LSTM Optimization
PufferLib's LSTM optimization uses LSTMCell during rollouts and LSTM during training for up to 3x faster inference:
```python
class OptimizedLSTMPolicy(nn.Module):
def __init__(self, observation_space, action_space, hidden_size=256):
super().__init__()
self.encoder = nn.Sequential(
layer_init(nn.Linear(observation_space.shape[0], 128)),
nn.ReLU()
)
# Use LSTMCell for step-by-step inference
self.lstm_cell = nn.LSTMCell(128, hidden_size)
# Use LSTM for batch training
self.lstm = nn.LSTM(128, hidden_size, num_layers=1)
self.actor = layer_init(nn.Linear(hidden_size, action_space.n), std=0.01)
self.critic = layer_init(nn.Linear(hidden_size, 1), std=1.0)
self.hidden_size = hidden_size
def encode_observations(self, observations, state):
"""Fast inference using LSTMCell."""
features = self.encoder(observations)
if state is None:
h = torch.zeros(observations.shape[0], self.hidden_size, device=features.device)
c = torch.zeros(observations.shape[0], self.hidden_size, device=features.device)
else:
h, c = state
# Step-by-step with LSTMCell (faster for inference)
h, c = self.lstm_cell(features, (h, c))
logits = self.actor(h)
value = self.critic(h)
return logits, value, (h, c)
def decode_actions(self, observations, actions, state):
"""Batch training using LSTM."""
seq_len, batch_size = observations.shape[:2]
# Reshape for LSTM
obs_flat = observations.reshape(seq_len * batch_size, -1)
features = self.encoder(obs_flat)
features = features.reshape(seq_len, batch_size, -1)
if state is None:
h = torch.zeros(1, batch_size, self.hidden_size, device=features.device)
c = torch.zeros(1, batch_size, self.hidden_size, device=features.device)
state = (h, c)
# Batch processing with LSTM (faster for training)
lstm_out, new_state = self.lstm(features, state)
# Flatten back
lstm_out = lstm_out.reshape(seq_len * batch_size, -1)
logits = self.actor(lstm_out)
value = self.critic(lstm_out)
return logits, value, new_state
```
## Multi-Input Policies
For environments with multiple observation types:
```python
class MultiInputPolicy(nn.Module):
def __init__(self, observation_space, action_space):
super().__init__()
# Separate encoders for different observation types
self.image_encoder = nn.Sequential(
layer_init(nn.Conv2d(3, 32, 8, stride=4)),
nn.ReLU(),
layer_init(nn.Conv2d(32, 64, 4, stride=2)),
nn.ReLU(),
nn.Flatten()
)
self.vector_encoder = nn.Sequential(
layer_init(nn.Linear(observation_space['vector'].shape[0], 128)),
nn.ReLU()
)
# Combine features
combined_size = 64 * 9 * 9 + 128 # Image features + vector features
self.combiner = nn.Sequential(
layer_init(nn.Linear(combined_size, 512)),
nn.ReLU()
)
self.actor = layer_init(nn.Linear(512, action_space.n), std=0.01)
self.critic = layer_init(nn.Linear(512, 1), std=1.0)
def forward(self, observations):
# Process each observation type
image_features = self.image_encoder(observations['image'].float() / 255.0)
vector_features = self.vector_encoder(observations['vector'])
# Combine
combined = torch.cat([image_features, vector_features], dim=-1)
features = self.combiner(combined)
return self.actor(features), self.critic(features)
```
## Continuous Action Policies
For continuous control tasks:
```python
class ContinuousPolicy(nn.Module):
def __init__(self, observation_space, action_space):
super().__init__()
self.encoder = nn.Sequential(
layer_init(nn.Linear(observation_space.shape[0], 256)),
nn.ReLU(),
layer_init(nn.Linear(256, 256)),
nn.ReLU()
)
# Mean of action distribution
self.actor_mean = layer_init(nn.Linear(256, action_space.shape[0]), std=0.01)
# Log std of action distribution
self.actor_logstd = nn.Parameter(torch.zeros(1, action_space.shape[0]))
# Value head
self.critic = layer_init(nn.Linear(256, 1), std=1.0)
def forward(self, observations):
features = self.encoder(observations)
action_mean = self.actor_mean(features)
action_std = torch.exp(self.actor_logstd)
value = self.critic(features)
return action_mean, action_std, value
def get_action(self, observations, deterministic=False):
action_mean, action_std, value = self.forward(observations)
if deterministic:
return action_mean, value
else:
dist = torch.distributions.Normal(action_mean, action_std)
action = dist.sample()
return torch.tanh(action), value # Bound actions to [-1, 1]
```
## Observation Processing
PufferLib provides utilities for unflattening observations:
```python
from pufferlib.pytorch import unflatten_observations
class PolicyWithUnflatten(nn.Module):
def __init__(self, observation_space, action_space):
super().__init__()
self.observation_space = observation_space
# Define encoders for each observation component
self.encoders = nn.ModuleDict({
'image': self._make_image_encoder(),
'vector': self._make_vector_encoder()
})
# ... rest of policy ...
def forward(self, flat_observations):
# Unflatten observations into structured format
observations = unflatten_observations(
flat_observations,
self.observation_space
)
# Process each component
image_features = self.encoders['image'](observations['image'])
vector_features = self.encoders['vector'](observations['vector'])
# Combine and continue...
```
## Multi-Agent Policies
### Shared Parameters
All agents use the same policy:
```python
class SharedMultiAgentPolicy(nn.Module):
def __init__(self, observation_space, action_space, num_agents):
super().__init__()
self.num_agents = num_agents
# Single policy shared across all agents
self.encoder = nn.Sequential(
layer_init(nn.Linear(observation_space.shape[0], 256)),
nn.ReLU()
)
self.actor = layer_init(nn.Linear(256, action_space.n), std=0.01)
self.critic = layer_init(nn.Linear(256, 1), std=1.0)
def forward(self, observations):
"""
Args:
observations: (batch * num_agents, obs_dim)
Returns:
logits: (batch * num_agents, num_actions)
values: (batch * num_agents, 1)
"""
features = self.encoder(observations)
return self.actor(features), self.critic(features)
```
### Independent Parameters
Each agent has its own policy:
```python
class IndependentMultiAgentPolicy(nn.Module):
def __init__(self, observation_space, action_space, num_agents):
super().__init__()
self.num_agents = num_agents
# Separate policy for each agent
self.policies = nn.ModuleList([
self._make_policy(observation_space, action_space)
for _ in range(num_agents)
])
def _make_policy(self, observation_space, action_space):
return nn.Sequential(
layer_init(nn.Linear(observation_space.shape[0], 256)),
nn.ReLU(),
layer_init(nn.Linear(256, 256)),
nn.ReLU()
)
def forward(self, observations, agent_ids):
"""
Args:
observations: (batch, obs_dim)
agent_ids: (batch,) which agent each obs belongs to
"""
outputs = []
for agent_id in range(self.num_agents):
mask = agent_ids == agent_id
if mask.any():
agent_obs = observations[mask]
agent_out = self.policies[agent_id](agent_obs)
outputs.append(agent_out)
return torch.cat(outputs, dim=0)
```
## Advanced Architectures
### Attention-Based Policy
```python
class AttentionPolicy(nn.Module):
def __init__(self, observation_space, action_space, d_model=256, nhead=8):
super().__init__()
self.encoder = layer_init(nn.Linear(observation_space.shape[0], d_model))
self.attention = nn.MultiheadAttention(d_model, nhead, batch_first=True)
self.actor = layer_init(nn.Linear(d_model, action_space.n), std=0.01)
self.critic = layer_init(nn.Linear(d_model, 1), std=1.0)
def forward(self, observations):
# Encode
features = self.encoder(observations)
# Self-attention
features = features.unsqueeze(1) # Add sequence dimension
attn_out, _ = self.attention(features, features, features)
attn_out = attn_out.squeeze(1)
return self.actor(attn_out), self.critic(attn_out)
```
### Residual Policy
```python
class ResidualBlock(nn.Module):
def __init__(self, dim):
super().__init__()
self.block = nn.Sequential(
layer_init(nn.Linear(dim, dim)),
nn.ReLU(),
layer_init(nn.Linear(dim, dim))
)
def forward(self, x):
return x + self.block(x)
class ResidualPolicy(nn.Module):
def __init__(self, observation_space, action_space, num_blocks=4):
super().__init__()
dim = 256
self.encoder = layer_init(nn.Linear(observation_space.shape[0], dim))
self.blocks = nn.Sequential(
*[ResidualBlock(dim) for _ in range(num_blocks)]
)
self.actor = layer_init(nn.Linear(dim, action_space.n), std=0.01)
self.critic = layer_init(nn.Linear(dim, 1), std=1.0)
def forward(self, observations):
x = torch.relu(self.encoder(observations))
x = self.blocks(x)
return self.actor(x), self.critic(x)
```
## Policy Best Practices
### Initialization
```python
# Always use layer_init for proper initialization
good_layer = layer_init(nn.Linear(256, 256))
# Use small std for actor head (more stable early training)
actor = layer_init(nn.Linear(256, num_actions), std=0.01)
# Use std=1.0 for critic head
critic = layer_init(nn.Linear(256, 1), std=1.0)
```
### Observation Normalization
```python
class NormalizedPolicy(nn.Module):
def __init__(self, observation_space, action_space):
super().__init__()
# Running statistics for normalization
self.obs_mean = nn.Parameter(torch.zeros(observation_space.shape[0]), requires_grad=False)
self.obs_std = nn.Parameter(torch.ones(observation_space.shape[0]), requires_grad=False)
# ... rest of policy ...
def forward(self, observations):
# Normalize observations
normalized_obs = (observations - self.obs_mean) / (self.obs_std + 1e-8)
# Continue with normalized observations
return self.policy(normalized_obs)
def update_normalization(self, observations):
"""Update running statistics."""
self.obs_mean.data = observations.mean(dim=0)
self.obs_std.data = observations.std(dim=0)
```
### Gradient Clipping
```python
# PufferLib trainer handles gradient clipping automatically
trainer = PuffeRL(
env=env,
policy=policy,
max_grad_norm=0.5 # Clip gradients to this norm
policy = pufferlib.models.Policy(
encoder=encoder,
decoder=decoder,
network=network,
)
```
### Model Compilation
The source contract is:
```python
# Enable torch.compile for faster training (PyTorch 2.0+)
policy = MyPolicy(observation_space, action_space)
- `Policy.initial_state(batch_size, device)`
- `Policy.forward_eval(x, state)` for rollout inference
- `Policy.forward(x)` for time-batched training
- encoder maps observations to hidden vectors
- recurrent/network module maps hidden vectors and state
- decoder maps hidden vectors to action logits and values
# Compile the model
policy = torch.compile(policy, mode='reduce-overhead')
Current built-ins include `DefaultEncoder`, `DefaultDecoder`, `MLP`, `MinGRU`,
`LSTM`, `GRU`, `NatureEncoder`, and `ImpalaEncoder`. INI config selects the
Torch fallback components:
# Use with trainer
trainer = PuffeRL(env=env, policy=policy, compile=True)
```ini
[torch]
network = MinGRU
encoder = DefaultEncoder
decoder = DefaultDecoder
[policy]
hidden_size = 128
num_layers = 4
```
## Debugging Policies
The default 4.0 backend is the native implementation, not this Torch fallback.
The CLI flag `--slowly` selects the fallback.
### Check Output Shapes
## Shape and numerical checks
```python
def test_policy_shapes(policy, observation_space, batch_size=32):
"""Verify policy output shapes."""
# Create dummy observations
obs = torch.randn(batch_size, *observation_space.shape)
Run these checks before a long job:
# Forward pass
logits, value = policy(obs)
1. Reset the reviewed environment and record observation shape/dtype/range.
2. Run one policy inference under `torch.no_grad()`.
3. For discrete actions, require logits shape
`(agent_batch, action_space.n)`.
4. Require values to represent one scalar per active agent.
5. For `MultiDiscrete`, verify branch count and each branch width.
6. For recurrent policies, verify state batch matches active agent rows and
that masks reset state at episode boundaries.
7. Reject NaN/Infinity in observations, logits, values, losses, and gradients.
8. Confirm inactive/padded multi-agent rows do not contribute to loss.
9. Run backward once and verify finite, non-missing gradients.
10. Compare eager and compiled outputs before enabling compilation.
# Check shapes
assert logits.shape == (batch_size, policy.action_space.n)
assert value.shape == (batch_size, 1)
`torch.compile` and reduced precision can alter performance and numerical
behavior. Record PyTorch, CUDA, compiler mode, precision, and deterministic
settings. Do not claim determinism solely because seeds are fixed.
print("✓ Policy shapes correct")
```
## Checkpoint-safe policy workflow
### Verify Gradients
- Save weights/state dictionaries, architecture config, environment revision,
package lock, seed, and checksum separately.
- Do not serialize arbitrary policy objects.
- Never call `torch.load` on an untrusted file. PufferLib 3.0 and the 4.0 Torch
fallback use `torch.load` for model paths; provenance review is therefore a
precondition, not an optional cleanup.
- Inspect metadata first with `scripts/inspect_checkpoint.py`; it never imports
Torch or deserializes.
- Verify an expected SHA-256 and license before loading.
- If business requirements force inspection of an untrusted model, isolate the
operation in a disposable sandbox with no credentials, network, host mounts,
or sensitive data. PyTorch warns that models are programs and that even
inspection tools may execute model code.
```python
def check_gradients(policy, observation_space):
"""Check that gradients flow properly."""
obs = torch.randn(1, *observation_space.shape, requires_grad=True)
## Sources
logits, value = policy(obs)
# Backward pass
loss = logits.sum() + value.sum()
loss.backward()
# Check gradients exist
for name, param in policy.named_parameters():
if param.grad is None:
print(f"⚠ No gradient for {name}")
elif torch.isnan(param.grad).any():
print(f"⚠ NaN gradient for {name}")
else:
print(f"✓ Gradient OK for {name}")
```
- [PufferLib 3.0 policy example](https://github.com/PufferAI/PufferLib/blob/3.0/examples/pufferl.py)
— stable example; accessed 2026-07-23.
- [PufferLib 3.0 PyTorch utilities](https://github.com/PufferAI/PufferLib/blob/3.0/pufferlib/pytorch.py)
— stable implementation; accessed 2026-07-23.
- [PufferLib 3.0 models](https://github.com/PufferAI/PufferLib/blob/3.0/pufferlib/models.py)
— stable model classes; accessed 2026-07-23.
- [PufferLib 4.0 models](https://github.com/PufferAI/PufferLib/blob/4.0/pufferlib/models.py)
— current source model contract; accessed 2026-07-23.
- [PufferLib 4.0 Torch trainer](https://github.com/PufferAI/PufferLib/blob/4.0/pufferlib/torch_pufferl.py)
— current fallback and checkpoint loading; accessed 2026-07-23.
- [PyTorch security policy](https://github.com/pytorch/pytorch/security) —
untrusted-model guidance; accessed 2026-07-23.
- [PyTorch `torch.load` documentation](https://docs.pytorch.org/docs/stable/generated/torch.load.html)
— deserialization warning; accessed 2026-07-23.

549
skills/pufferlib/references/training.md
View File

@@ -1,360 +1,287 @@
# PufferLib Training Guide
# Training, Evaluation, Configuration, and Logging
## Overview
Research snapshot: **2026-07-23**.
PuffeRL is PufferLib's high-performance training algorithm based on CleanRL's PPO with LSTMs, enhanced with proprietary research improvements. It achieves training at millions of steps per second through optimized vectorization and efficient implementation.
## Choose a version profile first
## Training Workflow
### Published stable package
### Basic Training Loop
PyPI's latest stable `pufferlib` release is **3.0.0**, published
**2025-06-23**. It declares Python `>=3.9` and is distributed only as a
60.7 MB source archive:
The PuffeRL trainer provides three core methods:
```python
# Collect environment interactions
rollout_data = trainer.evaluate()
# Train on collected batch
train_metrics = trainer.train()
# Aggregate and log results
trainer.mean_and_log()
```text
pufferlib-3.0.0.tar.gz
sha256: 7df3a3e3f5f894d78d2a1f5374097890aec01473183e748abefe4f3faa10eaa9
```
### CLI Training
The uploaded metadata depends on NumPy `<2.0`, Gym `<=0.23`, Gymnasium
`<=0.29.1`, PettingZoo `<=1.24.1`, Shimmy, Torch, Neptune, W&B, and other
packages without a complete transitive lock. It does not declare a CUDA
version or a minimum Torch version. Do not invent compatibility guarantees.
Quick start training via command line:
### Current source line
The upstream default branch is `4.0`; its `pyproject.toml` says version `4.0.0`,
Python `>=3.10`, and Torch `>=2.9`. As of the research date, this source line
is not the latest stable PyPI artifact.
The current PufferTank Dockerfile uses:
- Ubuntu 24.04
- NVIDIA CUDA `13.0.2` cuDNN development image
- Python 3.12
- the CUDA 13.0 PyTorch wheel index
- Nsight Systems `2025.6.3`
The Dockerfile does **not** pin an exact Torch wheel, uv version, PufferLib
commit, or every apt package. It is an upstream convenience environment, not a
complete reproducibility lock.
## Reproducible uv workflow
Do not use an unpinned `uv pip install pufferlib`. Work in a disposable,
project-specific environment and commit `pyproject.toml` plus `uv.lock`.
For the published profile, after reviewing the source archive and build:
```bash
# Basic training
puffer train environment_name --train.device cuda --train.learning-rate 0.001
# Custom configuration
puffer train environment_name \
--train.device cuda \
--train.batch-size 32768 \
--train.learning-rate 0.0003 \
--train.num-iterations 10000
uv venv --python 3.11
uv add --exact --no-sync "pufferlib==3.0.0"
uv lock
uv sync --frozen
```
### Python Training Script
Confirm the lock records the published SHA-256 above and review every resolved
dependency. The 3.0.0 source build can compile native code and may fetch build
assets. Resolve and build in a sandbox with no credentials or sensitive mounts.
Do not treat a successful resolver run as a security review.
```python
import pufferlib
from pufferlib import PuffeRL
# Initialize environment
env = pufferlib.make('environment_name', num_envs=256)
# Create trainer
trainer = PuffeRL(
env=env,
policy=my_policy,
device='cuda',
learning_rate=3e-4,
batch_size=32768,
n_epochs=4,
gamma=0.99,
gae_lambda=0.95,
clip_coef=0.2,
ent_coef=0.01,
vf_coef=0.5,
max_grad_norm=0.5
)
# Training loop
for iteration in range(num_iterations):
# Collect rollouts
rollout_data = trainer.evaluate()
# Train on batch
train_metrics = trainer.train()
# Log results
trainer.mean_and_log()
```
## Key Training Parameters
### Core Hyperparameters
- **learning_rate**: Learning rate for optimizer (default: 3e-4)
- **batch_size**: Number of timesteps per training batch (default: 32768)
- **n_epochs**: Number of training epochs per batch (default: 4)
- **num_envs**: Number of parallel environments (default: 256)
- **num_steps**: Steps per environment per rollout (default: 128)
### PPO Parameters
- **gamma**: Discount factor (default: 0.99)
- **gae_lambda**: Lambda for GAE calculation (default: 0.95)
- **clip_coef**: PPO clipping coefficient (default: 0.2)
- **ent_coef**: Entropy coefficient for exploration (default: 0.01)
- **vf_coef**: Value function loss coefficient (default: 0.5)
- **max_grad_norm**: Maximum gradient norm for clipping (default: 0.5)
### Performance Parameters
- **device**: Computing device ('cuda' or 'cpu')
- **compile**: Use torch.compile for faster training (default: True)
- **num_workers**: Number of vectorization workers (default: auto)
## Distributed Training
### Multi-GPU Training
Use torchrun for distributed training across multiple GPUs:
For 4.0 source work, pin an immutable revision rather than branch `4.0`:
```bash
torchrun --nproc_per_node=4 train.py \
--train.device cuda \
--train.batch-size 131072
uv add --no-sync \
"pufferlib @ git+https://github.com/PufferAI/PufferLib.git@25647630e1b15330bb3153a5a0d3ff8d234c3acf"
uv lock
```
### Multi-Node Training
The commit above is the reviewed 4.0 branch head on 2026-07-23. Re-review before
updating it. Native training still requires an audited build of a specific
environment; uv locking does not lock compilers, CUDA, NCCL, cuDNN, Raylib, or
system libraries.
For distributed training across multiple nodes:
Never run remote install scripts directly from a pipe. Download, inspect, pin,
verify, and execute only in an appropriate sandbox.
## Published 3.0.0 training
### CLI
The 3.0 console entry point is `puffer = pufferlib.pufferl:main`:
```bash
# On main node (rank 0)
torchrun --nproc_per_node=8 \
--nnodes=4 \
--node_rank=0 \
--master_addr=MASTER_IP \
--master_port=29500 \
train.py
# On worker nodes (rank 1, 2, 3)
torchrun --nproc_per_node=8 \
--nnodes=4 \
--node_rank=NODE_RANK \
--master_addr=MASTER_IP \
--master_port=29500 \
train.py
puffer train ENV_NAME [OPTIONS]
puffer eval ENV_NAME [OPTIONS]
puffer sweep ENV_NAME [OPTIONS]
puffer autotune ENV_NAME [OPTIONS]
puffer profile ENV_NAME [OPTIONS]
puffer export ENV_NAME [OPTIONS]
```
## Monitoring and Logging
Environment, vector, policy, recurrent, training, and sweep values come from
INI sections. Overrides use section-qualified flags:
### Logger Integration
PufferLib supports multiple logging backends:
#### Weights & Biases
```python
from pufferlib import WandbLogger
logger = WandbLogger(
project='my_project',
entity='my_team',
name='experiment_name',
config=trainer_config
)
trainer = PuffeRL(env, policy, logger=logger)
```bash
puffer train puffer_breakout \
--train.device cpu \
--train.total-timesteps 100000 \
--vec.backend Serial \
--vec.num-envs 2
```
#### Neptune
Run `puffer train ENV_NAME --help` against the exact locked environment because
available options are generated from merged INI files.
### Python API
The stable trainer is `pufferlib.pufferl.PuffeRL`, not a top-level
`pufferlib.PuffeRL`:
```python
from pufferlib import NeptuneLogger
from pufferlib import pufferl
logger = NeptuneLogger(
project='my_team/my_project',
name='experiment_name',
api_token='YOUR_TOKEN'
)
args = pufferl.load_config("puffer_breakout")
vecenv = pufferl.load_env("puffer_breakout", args)
policy = pufferl.load_policy(args, vecenv, "puffer_breakout")
trainer = pufferl.PuffeRL(args["train"], vecenv, policy)
trainer = PuffeRL(env, policy, logger=logger)
```
#### No Logger
```python
from pufferlib import NoLogger
trainer = PuffeRL(env, policy, logger=NoLogger())
```
### Key Metrics
Training logs include:
- **Performance Metrics**:
- Steps per second (SPS)
- Training throughput
- Wall-clock time per iteration
- **Learning Metrics**:
- Episode rewards (mean, min, max)
- Episode lengths
- Value function loss
- Policy loss
- Entropy
- Explained variance
- Clipfrac
- **Environment Metrics**:
- Environment-specific rewards
- Success rates
- Custom metrics
### Terminal Dashboard
PufferLib provides a real-time terminal dashboard showing:
- Training progress
- Current SPS
- Episode statistics
- Loss values
- GPU utilization
## Checkpointing
### Saving Checkpoints
```python
# Save checkpoint
trainer.save_checkpoint('checkpoint.pt')
# Save with additional metadata
trainer.save_checkpoint(
'checkpoint.pt',
metadata={'iteration': iteration, 'best_reward': best_reward}
)
```
### Loading Checkpoints
```python
# Load checkpoint
trainer.load_checkpoint('checkpoint.pt')
# Resume training
for iteration in range(resume_iteration, num_iterations):
trainer.evaluate()
trainer.train()
trainer.mean_and_log()
```
## Hyperparameter Tuning with Protein
The Protein system enables automatic hyperparameter and reward tuning:
```python
from pufferlib import Protein
# Define search space
search_space = {
'learning_rate': [1e-4, 3e-4, 1e-3],
'batch_size': [16384, 32768, 65536],
'ent_coef': [0.001, 0.01, 0.1],
'clip_coef': [0.1, 0.2, 0.3]
}
# Run hyperparameter search
protein = Protein(
env_name='environment_name',
search_space=search_space,
num_trials=100,
metric='mean_reward'
)
best_config = protein.optimize()
```
## Performance Optimization Tips
### Maximizing Throughput
1. **Batch Size**: Increase batch_size to fully utilize GPU
2. **Num Envs**: Balance between CPU and GPU utilization
3. **Compile**: Enable torch.compile for 10-20% speedup
4. **Workers**: Adjust num_workers based on environment complexity
5. **Device**: Always use 'cuda' for neural network training
### Environment Speed
- Pure Python environments: ~100k-500k SPS
- C-based environments: ~4M SPS
- With training overhead: ~1M-4M total SPS
### Memory Management
- Reduce batch_size if running out of GPU memory
- Decrease num_envs if running out of CPU memory
- Use gradient accumulation for large effective batch sizes
## Common Training Patterns
### Curriculum Learning
```python
# Start with easy tasks, gradually increase difficulty
difficulty_levels = [0.1, 0.3, 0.5, 0.7, 1.0]
for difficulty in difficulty_levels:
env = pufferlib.make('environment_name', difficulty=difficulty)
trainer = PuffeRL(env, policy)
for iteration in range(iterations_per_level):
try:
while trainer.epoch < trainer.total_epochs:
trainer.evaluate()
trainer.train()
trainer.mean_and_log()
finally:
trainer.close()
```
### Reward Shaping
The exact public methods include `evaluate`, `train`, `mean_and_log`,
`save_checkpoint`, `print_dashboard`, and `close`. Use the CLI when possible;
the Python trainer is a relatively low-level implementation surface.
```python
# Wrap environment with custom reward shaping
class RewardShapedEnv(pufferlib.PufferEnv):
def step(self, actions):
obs, rewards, dones, infos = super().step(actions)
### Stable configuration checks
# Add shaped rewards
shaped_rewards = rewards + 0.1 * proximity_bonus
- Make rollout/batch relationships explicit; do not rely on `auto` in a
published experiment.
- Record environment, vector, policy, recurrent, and train sections verbatim.
- Fix `seed` in both `[vec]` and `[train]`, then run multiple independent seeds.
- Record `torch_deterministic`, precision, compile settings, optimizer, horizon,
minibatch, and total timesteps.
- Keep evaluation seeds, instances, and metrics separate from training.
return obs, shaped_rewards, dones, infos
## Current 4.0 training
Build one audited environment, then use:
```bash
puffer train breakout
puffer eval breakout --load-model-path checkpoints/.../weights.bin
puffer sweep breakout
puffer match breakout \
--load-model-path trusted-a.bin \
--load-enemy-model-path trusted-b.bin
```
### Multi-Stage Training
Current modes are `train`, `eval`, `sweep`, `paretosweep`, and `match`.
Native training is the default. `--slowly` selects the Torch fallback.
Configuration uses sections such as:
```python
# Train in multiple stages with different configurations
stages = [
{'learning_rate': 1e-3, 'iterations': 1000}, # Exploration
{'learning_rate': 3e-4, 'iterations': 5000}, # Main training
{'learning_rate': 1e-4, 'iterations': 2000} # Fine-tuning
]
```ini
[vec]
total_agents = 4096
num_buffers = 2
num_threads = 16
for stage in stages:
trainer.learning_rate = stage['learning_rate']
for iteration in range(stage['iterations']):
trainer.evaluate()
trainer.train()
trainer.mean_and_log()
[train]
total_timesteps = 10_000_000
minibatch_size = 8192
horizon = 64
[torch]
network = MinGRU
encoder = DefaultEncoder
decoder = DefaultDecoder
```
## Troubleshooting
Current source validates that `minibatch_size` is divisible by `horizon` and
does not exceed `horizon * total_agents`. Multi-GPU launch uses spawn. Do a
small CPU/local build and contract test before CUDA training.
### Low Performance
## Held-out evaluation
- Check environment is vectorized correctly
- Verify GPU utilization with `nvidia-smi`
- Increase batch_size to saturate GPU
- Enable compile mode
- Profile with `torch.profiler`
Training rollouts are not evaluation. For every reported result:
### Training Instability
1. Freeze one checkpoint-selection rule before inspecting held-out scores.
2. Construct fresh evaluation environment instances.
3. Use evaluation seeds disjoint from training seeds.
4. Disable optimizer updates, exploration noise unless explicitly measuring it,
curriculum updates, normalization-stat updates, and reward shaping used only
for training.
5. Report deterministic and stochastic policy protocols separately.
6. Run enough episodes for uncertainty; report per-seed results and aggregate
intervals, not only a best run.
7. Preserve terminated versus truncated semantics in return/length accounting.
8. Record wrappers, frame skip, autoreset mode, opponent pool, policy state
reset, and rendering state.
- Reduce learning_rate
- Decrease batch_size
- Increase num_envs for more diverse samples
- Add entropy coefficient for more exploration
- Check reward scaling
Generate a starting plan:
### Memory Issues
```bash
python3 scripts/repro_plan.py --environment synthetic
```
- Reduce batch_size or num_envs
- Use gradient accumulation
- Disable compile mode if causing OOM
- Check for memory leaks in custom environments
## Checkpoints
PufferLib 3.0 saves a policy `state_dict` with `torch.save` and a separate
trainer state containing optimizer state, global step, epoch, and run ID. Its
loading paths call `torch.load`. The 4.0 native backend writes `.bin` weight
files; the 4.0 Torch fallback also uses `torch.save`/`torch.load`.
Rules:
- Never load an untrusted checkpoint, even to “inspect” it.
- Record SHA-256, size, source URL, immutable revision, license, environment,
policy architecture, package lock, and training config in a strict JSON
sidecar.
- Do not use `latest` in a reproducible run; resolve and record the exact path
and digest.
- Do not auto-download a run artifact by ID.
- Test restore and evaluation in a disposable environment before a long resume.
- A model-only checkpoint is not a bitwise resume; optimizer, scheduler,
normalizer, RNG, environment, and recurrent state may also matter.
Safe metadata inspection:
```bash
python3 scripts/inspect_checkpoint.py trusted/model.pt \
--expected-sha256 EXPECTED_DIGEST
```
The helper hashes and classifies bytes only. It never imports Torch, invokes
pickle, opens archive members, or extracts files.
## External logging
Local logging is the default. W&B and Neptune are optional network services
that may transmit configuration, metrics, source metadata, hardware telemetry,
stdout/stderr, and explicitly uploaded checkpoints/artifacts. They can create
storage, seat, compute, or retention costs and are subject to vendor privacy,
access, and retention policies.
Credential rules:
- W&B: use the named environment variable `WANDB_API_KEY` or an approved secret
manager.
- Neptune: use `NEPTUNE_API_TOKEN` or an approved secret manager.
- Never pass either secret as a CLI argument, INI/JSON value, logger config,
tag, run name, or chat/tool input.
- Never print the value or include it in a broad environment dump.
- Do not recursively search for `.env` files. If policy permits a local secret
file, read only the explicitly named key from the explicitly named file.
- Sanitize configuration before logging; reject keys containing token, secret,
password, credential, authorization, private key, or API key.
- Disable checkpoint/source upload unless separately approved.
PufferLib 3.0 supports both `--wandb` and `--neptune`; its sweep mode requires
one. Current 4.0 source exposes W&B but no Neptune CLI integration. In either
profile, require explicit logging opt-in and disclosure acknowledgment. The
bundled training planner enforces this without reading credential values:
```bash
python3 scripts/train_template.py \
--logger wandb \
--enable-external-logging \
--acknowledge-external-disclosure
```
## Sources
- [PyPI: pufferlib 3.0.0](https://pypi.org/project/pufferlib/3.0.0/) —
released 2025-06-23; accessed 2026-07-23.
- [PyPI 3.0.0 JSON metadata](https://pypi.org/pypi/pufferlib/3.0.0/json) —
package requirements and digest; accessed 2026-07-23.
- [PufferLib 3.0 trainer](https://github.com/PufferAI/PufferLib/blob/3.0/pufferlib/pufferl.py)
— stable CLI, logger, and checkpoint source; accessed 2026-07-23.
- [PufferLib 3.0 default config](https://github.com/PufferAI/PufferLib/blob/3.0/pufferlib/config/default.ini)
— stable parameters; accessed 2026-07-23.
- [PufferLib 4.0 docs](https://puffer.ai/docs.html) — current CLI and
architecture; accessed 2026-07-23.
- [PufferLib 4.0 trainer](https://github.com/PufferAI/PufferLib/blob/4.0/pufferlib/pufferl.py)
— current modes/config/checkpoints; accessed 2026-07-23.
- [PufferLib 4.0 package metadata](https://github.com/PufferAI/PufferLib/blob/4.0/pyproject.toml)
— current Python/Torch requirements; accessed 2026-07-23.
- [PufferTank 4.0 Dockerfile](https://github.com/PufferAI/PufferTank/blob/4.0/puffertank.dockerfile)
— CUDA/Python reference environment; accessed 2026-07-23.
- [Neptune Run API](https://docs.neptune.ai/run) — token and offline-mode
guidance; accessed 2026-07-23.
- [W&B documentation](https://docs.wandb.ai/) — logging and credential
guidance; accessed 2026-07-23.

689
skills/pufferlib/references/vectorization.md
View File

@@ -1,557 +1,210 @@
# PufferLib Vectorization Guide
# Vectorization and Throughput
## Overview
Research snapshot: **2026-07-23**. PufferLib's published 3.0.0 package and
current 4.0 source line expose different vectorization systems. Never mix their
configuration names.
PufferLib's vectorization system enables high-performance parallel environment simulation, achieving millions of steps per second through optimized implementation inspired by EnvPool. The system supports both synchronous and asynchronous vectorization with minimal overhead.
## Version split
## Vectorization Architecture
| Profile | Vectorization surface | Use for |
|---|---|---|
| PyPI `pufferlib==3.0.0` | `pufferlib.vector.make`; `Serial`, `Multiprocessing`, optional `Ray`, and native `PufferEnv` backends | Published Python/Gymnasium/PettingZoo compatibility workflows |
| Source `4.0` at a pinned commit | Native C vector interface configured by `[vec] total_agents`, `num_buffers`, and `num_threads` | Current Ocean/native trainer source |
### Key Optimizations
The 4.0 package directory no longer contains the 3.0 `vector.py`,
`emulation.py`, or `pytorch.py` modules. Treat old examples importing those
modules as 3.0 examples, even if a stale copy remains under the 4.0 `examples/`
tree.
1. **Shared Memory Buffer**: Single unified buffer across all environments (unlike Gymnasium's per-environment buffers)
2. **Busy-Wait Flags**: Workers busy-wait on unlocked flags rather than using pipes/queues
3. **Zero-Copy Batching**: Contiguous worker subsets return observations without copying
4. **Surplus Environments**: Simulates more environments than batch size for async returns
5. **Multiple Envs per Worker**: Optimizes performance for lightweight environments
## Published 3.0.0 API
### Performance Characteristics
- **Pure Python environments**: 100k-500k SPS
- **C-based environments**: 100M+ SPS
- **With training**: 400k-4M total SPS
- **Vectorization overhead**: <5% with optimal configuration
## Creating Vectorized Environments
### Basic Vectorization
The source signature is:
```python
import pufferlib
# Automatic vectorization
env = pufferlib.make('environment_name', num_envs=256)
# With explicit configuration
env = pufferlib.make(
'environment_name',
num_envs=256,
num_workers=8,
envs_per_worker=32
pufferlib.vector.make(
env_creator_or_creators,
env_args=None,
env_kwargs=None,
backend=pufferlib.PufferEnv,
num_envs=1,
seed=0,
**kwargs,
)
```
### Manual Vectorization
Select a backend explicitly during development:
```python
from pufferlib import PufferEnv
from pufferlib.vectorization import Serial, Multiprocessing
import pufferlib.vector
# Serial vectorization (single process)
vec_env = Serial(
env_creator=lambda: MyEnvironment(),
num_envs=16
serial = pufferlib.vector.make(
reviewed_env_creator,
backend=pufferlib.vector.Serial,
num_envs=4,
seed=42,
)
# Multiprocessing vectorization
vec_env = Multiprocessing(
env_creator=lambda: MyEnvironment(),
num_envs=256,
num_workers=8
parallel = pufferlib.vector.make(
reviewed_env_creator,
backend=pufferlib.vector.Multiprocessing,
num_envs=16,
num_workers=4,
batch_size=8,
zero_copy=True,
seed=42,
)
```
## Vectorization Modes
Do not replace `reviewed_env_creator` with a dotted import string. Import the
audited callable directly in trusted code. Environment construction executes
package code and may initialize native libraries.
### Serial Vectorization
### Stable backends
Best for debugging and lightweight environments:
- `PufferEnv` is the default native backend. `vector.make` requires
`num_envs=1` for this backend because that one native environment can manage
many agents internally.
- `Serial` runs multiple environment instances in the caller process. Use it
for contract debugging and deterministic comparisons.
- `Multiprocessing` uses worker processes and shared arrays. It supports the
synchronous `reset`/`step` facade and asynchronous `async_reset`/`recv` plus
`send`/`recv`.
- `Ray` is an optional 3.0 backend and requires the package's pinned Ray extra.
It introduces a separate distributed runtime and is not a safe local default.
### Shape semantics
Do not assume `num_envs == returned batch length`.
- A single environment advertises `num_agents`,
`single_observation_space`, and `single_action_space`.
- The full stable batch contains agent slots. For fixed-population environments,
serial batch length is normally `num_envs * num_agents`.
- `vecenv.agents_per_batch` is the number of agent rows returned by one receive.
- Observations have leading agent-batch dimension; rewards, terminals,
truncations, agent IDs, and masks have matching leading length.
- A synchronous call returns
`(observations, rewards, terminals, truncations, infos)`.
- The asynchronous `recv()` additionally returns `agent_ids` and `masks`.
- Structured observations/actions are flattened by emulation. Preserve the
recorded dtype metadata and unflatten in the policy; do not cast arbitrary
byte views to float first.
Validate actual shapes and dtypes at reset and the first step. In multi-agent
workflows, use masks to exclude padded or inactive slots from loss and metrics.
### Multiprocessing constraints
The 3.0 source validates these relationships:
1. `num_envs` must be divisible by `num_workers`.
2. `batch_size` defaults to `num_envs`.
3. `batch_size` must be divisible by `num_envs / num_workers`.
4. With zero-copy enabled, `num_envs` must be divisible by `batch_size`.
5. Physical-core oversubscription is rejected unless `overwork=True`; do not
bypass this for benchmark headline numbers.
Always call `close()` in `finally`. Keep constructors top-level and serializable.
Protect process creation with `if __name__ == "__main__":`.
### Start methods
The stable API does not expose a `start_method` argument. The process context is
therefore affected by Python and platform defaults. Set an application-wide
method before constructing workers if your program requires one:
```python
from pufferlib.vectorization import Serial
import multiprocessing as mp
vec_env = Serial(
env_creator=env_creator_fn,
num_envs=16
)
# All environments run in main process
# No multiprocessing overhead
# Easier debugging with standard tools
if __name__ == "__main__":
mp.set_start_method("spawn")
main()
```
**When to use:**
- Development and debugging
- Very fast environments (< 1μs per step)
- Small number of environments (< 32)
- Single-threaded profiling
Prefer `spawn` when CUDA, threads, or non-fork-safe native libraries may already
be initialized. Do not call `set_start_method(..., force=True)` inside a library.
Record the effective method in benchmark output. `forkserver` can also be
appropriate when available and tested. Never compare results that silently use
different methods.
### Multiprocessing Vectorization
### Seeding
Best for most production use cases:
- Pass one integer seed to `vector.make`; the stable implementation derives
per-environment seeds.
- `reset(seed=base_seed)` similarly offsets seeds across serial environments.
- Seed action-space sampling separately when random actions are part of a test.
- A deterministic seed does not guarantee bitwise deterministic GPU training or
deterministic third-party simulators.
- Recreate workers and environments for independent replicates; do not treat
adjacent episodes from one long run as independent seeds.
```python
from pufferlib.vectorization import Multiprocessing
## Current 4.0 source
vec_env = Multiprocessing(
env_creator=env_creator_fn,
num_envs=256,
num_workers=8,
envs_per_worker=32
)
The current default branch uses native C environments and a different vector
layout:
# Parallel execution across workers
# True parallelism for CPU-bound environments
# Scales to hundreds of environments
```ini
[vec]
total_agents = 4096
num_buffers = 2
num_threads = 16
```
**When to use:**
- Production training
- CPU-intensive environments
- Large-scale parallel simulation
- Maximizing throughput
Environment instances are grouped into buffers. Native execution uses OpenMP
threads inside those buffers; rollout workers coordinate buffer transfers and,
for GPU training, pinned memory and CUDA streams. The default trainer backend is
native; `--slowly` selects the PyTorch fallback. Multi-GPU launch code explicitly
uses a `spawn` multiprocessing context.
### Async Vectorization
Build and test one audited Ocean environment at a time. The C binding defines
observation size/type and action branches. A mismatch can cause memory
corruption rather than a friendly Python shape error.
For environments with variable step times:
## Benchmark methodology
```python
vec_env = Multiprocessing(
env_creator=env_creator_fn,
num_envs=256,
num_workers=8,
mode='async',
surplus_envs=32 # Simulate extra environments
)
Throughput is not a library constant. Report enough detail to reproduce it:
# Returns batches as soon as ready
# Better GPU utilization
# Handles variable environment speeds
1. Pin source/package, Python, dependencies, compiler, and environment revision.
2. Record CPU model/core topology, GPU/driver/CUDA, OS, precision, backend,
start method, env count, agent count, workers/threads, buffers, and batch.
3. State whether timing includes construction, reset, policy inference,
host-device transfer, learning, rendering, logging, and checkpoint I/O.
4. Warm up separately; use a fixed number of **agent steps**, not only wall time.
5. Run at least three independent repeats; report all samples plus median and
spread. Report failures and memory use.
6. Validate equivalent observations, actions, reset/autoreset behavior, frame
skip, and policy workload before comparing backends.
7. Distinguish simulation SPS from end-to-end training SPS.
The 2024 compatibility paper benchmarked PufferLib 1.x-style vectorization on an
i9-14900K/RTX 4090 desktop and an i7-10750H/RTX 3070 laptop. The 2025 PufferLib
2.0 paper reports a different Ocean/training system. Those results are scoped to
their listed hardware and workloads; they are not expected values for 3.0 or
4.0.
Run the bundled bounded harness first:
```bash
python3 scripts/benchmark_vectorization.py --backend serial
python3 scripts/benchmark_vectorization.py \
--backend multiprocessing --start-method spawn \
--envs 8 --workers 2 --steps-per-env 2000
```
**When to use:**
- Variable environment step times
- Maximizing GPU utilization
- Network-based environments
- External simulators
## Optimizing Vectorization Performance
### Worker Configuration
```python
import multiprocessing
# Calculate optimal workers
num_cpus = multiprocessing.cpu_count()
# Conservative (leave headroom for training)
num_workers = num_cpus - 2
# Aggressive (maximize environment throughput)
num_workers = num_cpus
# With hyperthreading
num_workers = num_cpus // 2 # Physical cores only
```
### Envs Per Worker
```python
# Fast environments (< 10μs per step)
envs_per_worker = 64 # More envs per worker
# Medium environments (10-100μs per step)
envs_per_worker = 32 # Balanced
# Slow environments (> 100μs per step)
envs_per_worker = 16 # Fewer envs per worker
# Calculate from target batch size
batch_size = 32768
num_workers = 8
envs_per_worker = batch_size // num_workers
```
### Batch Size Tuning
```python
# Small batch (< 8k): Good for fast iteration
batch_size = 4096
num_envs = 256
steps_per_env = batch_size // num_envs # 16 steps
# Medium batch (8k-32k): Good balance
batch_size = 16384
num_envs = 512
steps_per_env = 32
# Large batch (> 32k): Maximum throughput
batch_size = 65536
num_envs = 1024
steps_per_env = 64
```
## Shared Memory Optimization
### Buffer Management
PufferLib uses shared memory for zero-copy observation passing:
```python
import numpy as np
from multiprocessing import shared_memory
class OptimizedEnv(PufferEnv):
def __init__(self, buf=None):
super().__init__(buf)
# Environment will use provided shared buffer
self.observation_space = self.make_space({'obs': (84, 84, 3)})
# Observations written directly to shared memory
self._obs_buffer = None
def reset(self):
# Write to shared memory in-place
if self._obs_buffer is None:
self._obs_buffer = np.zeros((84, 84, 3), dtype=np.uint8)
self._render_to_buffer(self._obs_buffer)
return {'obs': self._obs_buffer}
def step(self, action):
# In-place updates only
self._update_state(action)
self._render_to_buffer(self._obs_buffer)
return {'obs': self._obs_buffer}, reward, done, info
```
### Zero-Copy Patterns
```python
# BAD: Creates copies
def get_observation(self):
obs = np.zeros((84, 84, 3))
# ... fill obs ...
return obs.copy() # Unnecessary copy!
# GOOD: Reuses buffer
def get_observation(self):
# Use pre-allocated buffer
self._render_to_buffer(self._obs_buffer)
return self._obs_buffer # No copy
# BAD: Allocates new arrays
def step(self, action):
new_state = self.state + action # Allocates
self.state = new_state
return obs, reward, done, info
# GOOD: In-place operations
def step(self, action):
self.state += action # In-place
return obs, reward, done, info
```
## Advanced Vectorization
### Custom Vectorization
```python
from pufferlib.vectorization import VectorEnv
class CustomVectorEnv(VectorEnv):
"""Custom vectorization implementation."""
def __init__(self, env_creator, num_envs, **kwargs):
super().__init__()
self.envs = [env_creator() for _ in range(num_envs)]
self.num_envs = num_envs
def reset(self):
"""Reset all environments."""
observations = [env.reset() for env in self.envs]
return self._stack_obs(observations)
def step(self, actions):
"""Step all environments."""
results = [env.step(action) for env, action in zip(self.envs, actions)]
obs, rewards, dones, infos = zip(*results)
return (
self._stack_obs(obs),
np.array(rewards),
np.array(dones),
list(infos)
)
def _stack_obs(self, observations):
"""Stack observations into batch."""
return np.stack(observations, axis=0)
```
### Hierarchical Vectorization
For very large-scale parallelism:
```python
# Outer: Multiprocessing vectorization (8 workers)
# Inner: Each worker runs serial vectorization (32 envs)
# Total: 256 parallel environments
def create_serial_vec_env():
return Serial(
env_creator=lambda: MyEnvironment(),
num_envs=32
)
outer_vec_env = Multiprocessing(
env_creator=create_serial_vec_env,
num_envs=8, # 8 serial vec envs
num_workers=8
)
# Total environments: 8 * 32 = 256
```
## Multi-Agent Vectorization
### Native Multi-Agent Support
PufferLib treats multi-agent environments as first-class citizens:
```python
# Multi-agent environment automatically vectorized
env = pufferlib.make(
'pettingzoo-knights-archers-zombies',
num_envs=128,
num_agents=4
)
# Observations: {agent_id: [batch_obs]} for each agent
# Actions: {agent_id: [batch_actions]} for each agent
# Rewards: {agent_id: [batch_rewards]} for each agent
```
### Custom Multi-Agent Vectorization
```python
class MultiAgentVectorEnv(VectorEnv):
def step(self, actions):
"""
Args:
actions: Dict of {agent_id: [batch_actions]}
Returns:
observations: Dict of {agent_id: [batch_obs]}
rewards: Dict of {agent_id: [batch_rewards]}
dones: Dict of {agent_id: [batch_dones]}
infos: List of dicts
"""
# Distribute actions to environments
env_actions = self._distribute_actions(actions)
# Step each environment
results = [env.step(act) for env, act in zip(self.envs, env_actions)]
# Collect and batch results
return self._batch_results(results)
```
## Performance Monitoring
### Profiling Vectorization
```python
import time
def profile_vectorization(vec_env, num_steps=10000):
"""Profile vectorization performance."""
start = time.time()
vec_env.reset()
for _ in range(num_steps):
actions = vec_env.action_space.sample()
vec_env.step(actions)
elapsed = time.time() - start
sps = (num_steps * vec_env.num_envs) / elapsed
print(f"Steps per second: {sps:,.0f}")
print(f"Time per step: {elapsed/num_steps*1000:.2f}ms")
return sps
```
### Bottleneck Analysis
```python
import cProfile
import pstats
def analyze_bottlenecks(vec_env):
"""Identify vectorization bottlenecks."""
profiler = cProfile.Profile()
profiler.enable()
vec_env.reset()
for _ in range(1000):
actions = vec_env.action_space.sample()
vec_env.step(actions)
profiler.disable()
stats = pstats.Stats(profiler)
stats.sort_stats('cumulative')
stats.print_stats(20)
```
### Real-Time Monitoring
```python
class MonitoredVectorEnv(VectorEnv):
"""Vector environment with performance monitoring."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.step_times = []
self.step_count = 0
def step(self, actions):
start = time.perf_counter()
result = super().step(actions)
elapsed = time.perf_counter() - start
self.step_times.append(elapsed)
self.step_count += 1
# Log every 1000 steps
if self.step_count % 1000 == 0:
mean_time = np.mean(self.step_times[-1000:])
sps = self.num_envs / mean_time
print(f"SPS: {sps:,.0f} | Step time: {mean_time*1000:.2f}ms")
return result
```
## Troubleshooting
### Low Throughput
```python
# Check configuration
print(f"Num envs: {vec_env.num_envs}")
print(f"Num workers: {vec_env.num_workers}")
print(f"Envs per worker: {vec_env.num_envs // vec_env.num_workers}")
# Profile single environment
single_env = MyEnvironment()
single_sps = profile_single_env(single_env)
print(f"Single env SPS: {single_sps:,.0f}")
# Compare vectorized
vec_sps = profile_vectorization(vec_env)
print(f"Vectorized SPS: {vec_sps:,.0f}")
print(f"Speedup: {vec_sps / single_sps:.1f}x")
```
### Memory Issues
```python
# Reduce number of environments
num_envs = 128 # Instead of 256
# Reduce envs per worker
envs_per_worker = 16 # Instead of 32
# Use Serial mode for debugging
vec_env = Serial(env_creator, num_envs=16)
```
### Synchronization Problems
```python
# Ensure thread-safe operations
import threading
class ThreadSafeEnv(PufferEnv):
def __init__(self, buf=None):
super().__init__(buf)
self.lock = threading.Lock()
def step(self, action):
with self.lock:
return super().step(action)
```
## Best Practices
### Configuration Guidelines
```python
# Start conservative
config = {
'num_envs': 64,
'num_workers': 4,
'envs_per_worker': 16
}
# Scale up iteratively
config = {
'num_envs': 256, # 4x increase
'num_workers': 8, # 2x increase
'envs_per_worker': 32 # 2x increase
}
# Monitor and adjust
if sps < target_sps:
# Try increasing num_envs or num_workers
pass
if memory_usage > threshold:
# Reduce num_envs or envs_per_worker
pass
```
### Environment Design
```python
# Minimize per-step allocations
class EfficientEnv(PufferEnv):
def __init__(self, buf=None):
super().__init__(buf)
# Pre-allocate all buffers
self._obs = np.zeros((84, 84, 3), dtype=np.uint8)
self._state = np.zeros(10, dtype=np.float32)
def step(self, action):
# Use pre-allocated buffers
self._update_state_inplace(action)
self._render_to_obs()
return self._obs, reward, done, info
```
### Testing
```python
# Test vectorization matches serial
serial_env = Serial(env_creator, num_envs=4)
vec_env = Multiprocessing(env_creator, num_envs=4, num_workers=2)
# Run parallel and verify results match
serial_env.seed(42)
vec_env.seed(42)
serial_obs = serial_env.reset()
vec_obs = vec_env.reset()
assert np.allclose(serial_obs, vec_obs), "Vectorization mismatch!"
```
It benchmarks only the bundled synthetic environment, never imports PufferLib,
and cannot substantiate an upstream PufferLib performance claim.
## Sources
- [PufferLib 3.0 vector source](https://github.com/PufferAI/PufferLib/blob/3.0/pufferlib/vector.py)
— stable API source; accessed 2026-07-23.
- [PufferLib 3.0 vectorization example](https://github.com/PufferAI/PufferLib/blob/3.0/examples/vectorization.py)
— stable usage example; accessed 2026-07-23.
- [PufferLib 4.0 documentation](https://puffer.ai/docs.html) — current native
architecture and CLI; accessed 2026-07-23.
- [PufferLib 4.0 trainer source](https://github.com/PufferAI/PufferLib/blob/4.0/pufferlib/pufferl.py)
— current config and spawn behavior; accessed 2026-07-23.
- [PufferLib compatibility paper](https://arxiv.org/abs/2406.12905) — submitted
2024-06-18.
- [PufferLib 2.0 paper](https://openreview.net/forum?id=qRyteMTgn0) —
Reinforcement Learning Journal, 2025.

1
skills/pufferlib/scripts/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Local, dependency-free helpers for the PufferLib skill."""

199
skills/pufferlib/scripts/_common.py Normal file
View File

@@ -0,0 +1,199 @@
#!/usr/bin/env python3
"""Shared safety and strict-JSON helpers for bundled PufferLib CLIs."""
from __future__ import annotations
import json
import math
import re
from pathlib import Path
from typing import Any
MAX_JSON_BYTES = 1_048_576
MAX_STEPS = 1_000_000_000
MAX_ENVS = 65_536
MAX_WORKERS = 256
MAX_EVAL_EPISODES = 10_000
STABLE_SDIST_SHA256 = (
"7df3a3e3f5f894d78d2a1f5374097890aec01473183e748abefe4f3faa10eaa9"
)
SOURCE_4_COMMIT = "25647630e1b15330bb3153a5a0d3ff8d234c3acf"
LOGGER_CREDENTIAL_ENV = {
"none": None,
"wandb": "WANDB_API_KEY",
"neptune": "NEPTUNE_API_TOKEN",
}
_SLUG = re.compile(r"^[a-z0-9][a-z0-9_-]{0,63}$")
_SHA256 = re.compile(r"^[0-9a-f]{64}$")
_SECRET_KEY = re.compile(
r"(?:^|_)(?:api_?key|token|secret|password|credential|"
r"private_?key|authorization)(?:$|_)",
re.IGNORECASE,
)
class UserInputError(ValueError):
"""Raised for bounded, user-correctable input errors."""
def bounded_int(value: str | int, *, name: str, minimum: int, maximum: int) -> int:
"""Parse an integer while rejecting bools and out-of-range values."""
if isinstance(value, bool):
raise UserInputError(f"{name} must be an integer, not bool")
try:
parsed = int(value)
except (TypeError, ValueError) as exc:
raise UserInputError(f"{name} must be an integer") from exc
if not minimum <= parsed <= maximum:
raise UserInputError(f"{name} must be between {minimum} and {maximum}")
return parsed
def validate_slug(value: str, *, name: str = "name") -> str:
"""Accept a local identifier, never a dotted import path."""
if not isinstance(value, str) or not _SLUG.fullmatch(value):
raise UserInputError(
f"{name} must match {_SLUG.pattern!r}; dotted import paths are not accepted"
)
return value
def validate_sha256(value: str, *, name: str = "sha256") -> str:
"""Validate a lowercase SHA-256 digest."""
if not isinstance(value, str) or not _SHA256.fullmatch(value):
raise UserInputError(f"{name} must be 64 lowercase hexadecimal characters")
return value
def _reject_constant(value: str) -> None:
raise UserInputError(f"non-finite JSON number is not allowed: {value}")
def _reject_duplicate_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
result: dict[str, Any] = {}
for key, value in pairs:
if key in result:
raise UserInputError(f"duplicate JSON key: {key}")
result[key] = value
return result
def _assert_finite_json(value: Any, path: str = "$") -> None:
if isinstance(value, float) and not math.isfinite(value):
raise UserInputError(f"{path} contains a non-finite number")
if isinstance(value, dict):
for key, item in value.items():
_assert_finite_json(item, f"{path}.{key}")
elif isinstance(value, list):
for index, item in enumerate(value):
_assert_finite_json(item, f"{path}[{index}]")
def strict_json_loads(text: str) -> Any:
"""Load JSON with duplicate-key and non-finite-number rejection."""
try:
value = json.loads(
text,
object_pairs_hook=_reject_duplicate_keys,
parse_constant=_reject_constant,
)
except json.JSONDecodeError as exc:
raise UserInputError(
f"invalid JSON at line {exc.lineno}, column {exc.colno}: {exc.msg}"
) from exc
_assert_finite_json(value)
return value
def strict_json_dumps(value: Any, *, pretty: bool = True) -> str:
"""Serialize deterministic JSON and reject NaN or Infinity."""
_assert_finite_json(value)
return json.dumps(
value,
allow_nan=False,
ensure_ascii=True,
indent=2 if pretty else None,
separators=None if pretty else (",", ":"),
sort_keys=True,
)
def emit_json(value: Any, *, pretty: bool = True) -> None:
print(strict_json_dumps(value, pretty=pretty))
def resolve_local_path(
value: str | Path,
*,
root: str | Path,
must_exist: bool = True,
reject_symlink: bool = True,
) -> Path:
"""Resolve a path beneath an explicit root without directory traversal."""
root_path = Path(root).expanduser().resolve(strict=True)
raw_path = Path(value).expanduser()
candidate = raw_path if raw_path.is_absolute() else root_path / raw_path
if reject_symlink and candidate.is_symlink():
raise UserInputError(f"symlinks are not accepted: {candidate}")
try:
resolved = candidate.resolve(strict=must_exist)
except OSError as exc:
raise UserInputError(f"cannot resolve path: {candidate}") from exc
try:
resolved.relative_to(root_path)
except ValueError as exc:
raise UserInputError(f"path escapes root {root_path}: {candidate}") from exc
return resolved
def load_json_object(
path: str | Path,
*,
root: str | Path,
max_bytes: int = MAX_JSON_BYTES,
) -> dict[str, Any]:
"""Read one explicitly named, bounded UTF-8 JSON object."""
resolved = resolve_local_path(path, root=root, must_exist=True)
size = resolved.stat().st_size
if size > max_bytes:
raise UserInputError(f"JSON file exceeds {max_bytes} bytes: {resolved}")
try:
text = resolved.read_text(encoding="utf-8")
except (OSError, UnicodeError) as exc:
raise UserInputError(f"cannot read UTF-8 JSON file: {resolved}") from exc
value = strict_json_loads(text)
if not isinstance(value, dict):
raise UserInputError("top-level JSON value must be an object")
return value
def secret_key_paths(value: Any, path: str = "$") -> list[str]:
"""Return key paths that look like credential-bearing configuration."""
matches: list[str] = []
if isinstance(value, dict):
for key, item in value.items():
normalized = re.sub(r"[^a-z0-9]+", "_", str(key).lower()).strip("_")
child_path = f"{path}.{key}"
if _SECRET_KEY.search(normalized):
matches.append(child_path)
matches.extend(secret_key_paths(item, child_path))
elif isinstance(value, list):
for index, item in enumerate(value):
matches.extend(secret_key_paths(item, f"{path}[{index}]"))
return matches
def require_keys(
mapping: dict[str, Any],
*,
allowed: set[str],
required: set[str],
path: str,
) -> list[str]:
"""Return schema errors for missing and unknown mapping keys."""
errors = [f"{path}.{key} is required" for key in sorted(required - mapping.keys())]
errors.extend(f"{path}.{key} is not allowed" for key in sorted(mapping.keys() - allowed))
return errors

244
skills/pufferlib/scripts/benchmark_vectorization.py Normal file
View File

@@ -0,0 +1,244 @@
#!/usr/bin/env python3
"""Bounded synthetic vectorization benchmark with no PufferLib import."""
from __future__ import annotations
import argparse
import multiprocessing as mp
import os
import platform
import statistics
import time
from typing import Any
try:
from ._common import UserInputError, bounded_int, emit_json
from .env_template import SyntheticGymEnv
except ImportError: # Direct script execution.
from _common import UserInputError, bounded_int, emit_json
from env_template import SyntheticGymEnv
def _run_partition(payload: tuple[int, int, int, int, int]) -> dict[str, Any]:
"""Run an allowlisted synthetic partition in one process."""
start_index, num_envs, steps_per_env, seed, max_steps = payload
agent_steps = 0
checksum = 0.0
resets = 0
for offset in range(num_envs):
env_index = start_index + offset
env = SyntheticGymEnv(max_steps=max_steps)
env.reset(seed=seed + env_index)
env_resets = 0
for step_index in range(steps_per_env):
action = (seed + 17 * env_index + step_index) % env.action_space.n
observation, reward, terminated, truncated, _ = env.step(action)
checksum += reward + observation[0] * 1e-6
agent_steps += 1
if terminated or truncated:
env_resets += 1
resets += 1
env.reset(seed=seed + env_index + env_resets * 1_000_003)
env.close()
return {"agent_steps": agent_steps, "checksum": checksum, "resets": resets}
def _partitions(
*,
num_envs: int,
workers: int,
steps_per_env: int,
seed: int,
max_steps: int,
) -> list[tuple[int, int, int, int, int]]:
parts: list[tuple[int, int, int, int, int]] = []
base, remainder = divmod(num_envs, workers)
start = 0
for worker_index in range(workers):
count = base + int(worker_index < remainder)
if count:
parts.append((start, count, steps_per_env, seed, max_steps))
start += count
return parts
def _summarize(samples: list[float]) -> dict[str, float]:
ordered = sorted(samples)
def percentile(fraction: float) -> float:
index = round((len(ordered) - 1) * fraction)
return ordered[index]
return {
"maximum": max(samples),
"median": statistics.median(samples),
"minimum": min(samples),
"p10": percentile(0.10),
"p90": percentile(0.90),
}
def benchmark(
*,
backend: str,
num_envs: int,
workers: int,
steps_per_env: int,
repeats: int,
warmup_steps: int,
seed: int,
max_steps: int,
start_method: str,
) -> dict[str, Any]:
"""Measure a fixed synthetic workload; never claim upstream PufferLib SPS."""
workers = min(workers, num_envs)
payloads = _partitions(
num_envs=num_envs,
workers=workers,
steps_per_env=steps_per_env,
seed=seed,
max_steps=max_steps,
)
warmup_payloads = _partitions(
num_envs=num_envs,
workers=workers,
steps_per_env=warmup_steps,
seed=seed,
max_steps=max_steps,
)
elapsed_samples: list[float] = []
throughput_samples: list[float] = []
checksums: list[float] = []
observed_steps: list[int] = []
pool: Any = None
try:
if backend == "multiprocessing":
context = mp.get_context(start_method)
pool = context.Pool(processes=workers)
if warmup_steps:
pool.map(_run_partition, warmup_payloads)
elif warmup_steps:
for payload in warmup_payloads:
_run_partition(payload)
for repeat_index in range(repeats):
adjusted = [
(start, count, steps, run_seed + repeat_index, limit)
for start, count, steps, run_seed, limit in payloads
]
started = time.perf_counter()
if backend == "multiprocessing":
results = pool.map(_run_partition, adjusted)
else:
results = [_run_partition(payload) for payload in adjusted]
elapsed = time.perf_counter() - started
agent_steps = sum(int(item["agent_steps"]) for item in results)
checksum = sum(float(item["checksum"]) for item in results)
elapsed_samples.append(elapsed)
throughput_samples.append(agent_steps / elapsed)
observed_steps.append(agent_steps)
checksums.append(checksum)
finally:
if pool is not None:
pool.close()
pool.join()
return {
"backend": backend,
"benchmark": "bundled-synthetic-harness",
"checksums": checksums,
"elapsed_seconds": _summarize(elapsed_samples),
"environment_construction_in_timed_region": True,
"hardware": {
"logical_cpus": os.cpu_count(),
"platform": platform.platform(),
"python": platform.python_version(),
},
"network_used": False,
"parameters": {
"max_steps": max_steps,
"num_envs": num_envs,
"repeats": repeats,
"seed": seed,
"start_method": start_method if backend == "multiprocessing" else None,
"steps_per_env": steps_per_env,
"warmup_steps": warmup_steps,
"workers": workers,
},
"steps_per_second": _summarize(throughput_samples),
"total_agent_steps_per_repeat": observed_steps,
"warning": (
"This measures the bundled synthetic harness, not PufferLib, an Ocean "
"environment, training throughput, or cross-machine performance."
),
}
def build_parser() -> argparse.ArgumentParser:
safe_methods = [method for method in ("spawn", "forkserver") if method in mp.get_all_start_methods()]
parser = argparse.ArgumentParser(
description=(
"Benchmark only the built-in synthetic environment. Defaults are CPU-only, "
"network-free, and intentionally small."
)
)
parser.add_argument(
"--backend", choices=["serial", "multiprocessing"], default="serial"
)
parser.add_argument("--envs", type=int, default=4, help="1..128")
parser.add_argument("--workers", type=int, default=2, help="1..32")
parser.add_argument("--steps-per-env", type=int, default=1_000, help="1..100000")
parser.add_argument("--repeats", type=int, default=3, help="1..7")
parser.add_argument("--warmup-steps", type=int, default=32, help="0..1000")
parser.add_argument("--max-steps", type=int, default=32, help="1..10000")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument(
"--start-method",
choices=safe_methods,
default="spawn" if "spawn" in safe_methods else safe_methods[0],
help="fork is intentionally unavailable",
)
parser.add_argument("--compact", action="store_true")
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
try:
num_envs = bounded_int(args.envs, name="envs", minimum=1, maximum=128)
workers = bounded_int(args.workers, name="workers", minimum=1, maximum=32)
steps = bounded_int(
args.steps_per_env,
name="steps_per_env",
minimum=1,
maximum=100_000,
)
repeats = bounded_int(args.repeats, name="repeats", minimum=1, maximum=7)
warmup = bounded_int(
args.warmup_steps, name="warmup_steps", minimum=0, maximum=1_000
)
max_steps = bounded_int(
args.max_steps, name="max_steps", minimum=1, maximum=10_000
)
report = benchmark(
backend=args.backend,
num_envs=num_envs,
workers=workers,
steps_per_env=steps,
repeats=repeats,
warmup_steps=warmup,
seed=args.seed,
max_steps=max_steps,
start_method=args.start_method,
)
except (UserInputError, ValueError, RuntimeError) as exc:
parser.error(str(exc))
emit_json(report, pretty=not args.compact)
return 0
if __name__ == "__main__":
raise SystemExit(main())

198
skills/pufferlib/scripts/env_contract_validator.py Normal file
View File

@@ -0,0 +1,198 @@
#!/usr/bin/env python3
"""Validate a built-in synthetic environment without importing plug-ins."""
from __future__ import annotations
import argparse
import math
import random
from typing import Any
try:
from ._common import UserInputError, bounded_int, emit_json
from .env_template import SyntheticGymEnv
except ImportError: # Direct script execution.
from _common import UserInputError, bounded_int, emit_json
from env_template import SyntheticGymEnv
def _require(condition: bool, message: str, errors: list[str]) -> None:
if not condition:
errors.append(message)
def _validate_reset(result: Any, env: SyntheticGymEnv, errors: list[str]) -> Any:
_require(
isinstance(result, tuple) and len(result) == 2,
"reset() must return (observation, info)",
errors,
)
if not isinstance(result, tuple) or len(result) != 2:
return None
observation, info = result
_require(
env.observation_space.contains(observation),
"reset observation is outside observation_space",
errors,
)
_require(isinstance(info, dict), "reset info must be a dict", errors)
return observation
def _validate_step(result: Any, env: SyntheticGymEnv, errors: list[str]) -> tuple[bool, bool]:
_require(
isinstance(result, tuple) and len(result) == 5,
"step() must return (observation, reward, terminated, truncated, info)",
errors,
)
if not isinstance(result, tuple) or len(result) != 5:
return False, False
observation, reward, terminated, truncated, info = result
_require(
env.observation_space.contains(observation),
"step observation is outside observation_space",
errors,
)
_require(
isinstance(reward, (int, float))
and not isinstance(reward, bool)
and math.isfinite(float(reward)),
"reward must be a finite scalar",
errors,
)
_require(type(terminated) is bool, "terminated must be bool", errors)
_require(type(truncated) is bool, "truncated must be bool", errors)
_require(isinstance(info, dict), "step info must be a dict", errors)
_require(
not (terminated and truncated),
"synthetic environment must not terminate and truncate simultaneously",
errors,
)
return bool(terminated), bool(truncated)
def _check_determinism(seed: int, max_steps: int, errors: list[str]) -> None:
env_a = SyntheticGymEnv(max_steps=max_steps)
env_b = SyntheticGymEnv(max_steps=max_steps)
first_a = env_a.reset(seed=seed)
first_b = env_b.reset(seed=seed)
_require(first_a == first_b, "same reset seed produced different results", errors)
actions = [0, 2, 1, 2, 0, 1]
for action in actions:
result_a = env_a.step(action)
result_b = env_b.step(action)
_require(result_a == result_b, "same action trace produced different results", errors)
if result_a[2] or result_a[3]:
break
env_a.close()
env_b.close()
def validate_synthetic(
*,
seed: int,
steps: int,
episodes: int,
max_steps: int,
) -> dict[str, Any]:
"""Run bounded API, space, reset, termination, and determinism checks."""
errors: list[str] = []
env = SyntheticGymEnv(max_steps=max_steps)
_require(hasattr(env, "observation_space"), "missing observation_space", errors)
_require(hasattr(env, "action_space"), "missing action_space", errors)
action_rng = random.Random(seed + 1)
observation = _validate_reset(env.reset(seed=seed), env, errors)
total_steps = 0
completed_episodes = 0
terminations = 0
truncations = 0
while total_steps < steps and completed_episodes < episodes:
action = env.action_space.sample(action_rng)
_require(env.action_space.contains(action), "sampled action is invalid", errors)
terminated, truncated = _validate_step(env.step(action), env, errors)
total_steps += 1
if terminated or truncated:
completed_episodes += 1
terminations += int(terminated)
truncations += int(truncated)
observation = _validate_reset(
env.reset(seed=seed + completed_episodes), env, errors
)
_require(
observation is None or env.observation_space.contains(observation),
"final observation is invalid",
errors,
)
_check_determinism(seed, max_steps, errors)
env.close()
return {
"checks": {
"deterministic_seed": "passed" if not errors else "see errors",
"reset_two_tuple": True,
"spaces": True,
"step_five_tuple": True,
},
"contract": "gymnasium",
"environment": "synthetic",
"errors": errors,
"network_used": False,
"observed": {
"episodes": completed_episodes,
"steps": total_steps,
"terminations": terminations,
"truncations": truncations,
},
"seed": seed,
"status": "passed" if not errors else "failed",
"vector_shape_expectation": {
"actions": ["num_envs"],
"observations": ["num_envs", 4],
"rewards": ["num_envs"],
"terminations": ["num_envs"],
"truncations": ["num_envs"],
},
}
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description=(
"Validate only the allowlisted built-in synthetic environment; "
"external modules and dotted import paths are not supported."
)
)
parser.add_argument("--environment", choices=["synthetic"], default="synthetic")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--steps", type=int, default=64, help="1..10000")
parser.add_argument("--episodes", type=int, default=8, help="1..100")
parser.add_argument("--max-steps", type=int, default=16, help="1..10000")
parser.add_argument("--compact", action="store_true")
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
try:
steps = bounded_int(args.steps, name="steps", minimum=1, maximum=10_000)
episodes = bounded_int(args.episodes, name="episodes", minimum=1, maximum=100)
max_steps = bounded_int(
args.max_steps, name="max_steps", minimum=1, maximum=10_000
)
report = validate_synthetic(
seed=args.seed,
steps=steps,
episodes=episodes,
max_steps=max_steps,
)
except (UserInputError, ValueError, RuntimeError) as exc:
parser.error(str(exc))
emit_json(report, pretty=not args.compact)
return 0 if report["status"] == "passed" else 1
if __name__ == "__main__":
raise SystemExit(main())

488
skills/pufferlib/scripts/env_template.py
View File

@@ -1,340 +1,210 @@
#!/usr/bin/env python3
"""
PufferLib Environment Template
"""Dependency-free synthetic Gymnasium-style environment template.
This template provides a starting point for creating custom PufferEnv environments.
Customize the observation space, action space, and environment logic for your task.
This module is intentionally local and synthetic. It does not import PufferLib,
Gymnasium, environment plug-ins, native extensions, or ROMs. Port the contract
to a separately reviewed Gymnasium or PufferLib environment only after the
validator passes.
"""
import numpy as np
import pufferlib
from pufferlib import PufferEnv
from __future__ import annotations
import argparse
import math
import random
from dataclasses import dataclass
from typing import Any
try:
from ._common import UserInputError, bounded_int, emit_json
except ImportError: # Direct script execution.
from _common import UserInputError, bounded_int, emit_json
class MyEnvironment(PufferEnv):
"""
Custom PufferLib environment template.
@dataclass(frozen=True)
class DiscreteSpace:
"""Minimal stand-in for a discrete action space."""
This is a simple grid world example. Customize it for your specific task.
"""
n: int
def __init__(self, buf=None, grid_size=10, max_steps=1000):
"""
Initialize environment.
@property
def shape(self) -> tuple[int, ...]:
return ()
Args:
buf: Shared memory buffer (managed by PufferLib)
grid_size: Size of the grid world
max_steps: Maximum steps per episode
"""
super().__init__(buf)
def contains(self, value: Any) -> bool:
return isinstance(value, int) and not isinstance(value, bool) and 0 <= value < self.n
self.grid_size = grid_size
def sample(self, rng: random.Random) -> int:
return rng.randrange(self.n)
@dataclass(frozen=True)
class BoxSpace:
"""Minimal one-dimensional finite box used by the synthetic environment."""
low: float
high: float
shape: tuple[int, ...]
dtype: str = "float32"
def contains(self, value: Any) -> bool:
if len(self.shape) != 1 or not isinstance(value, (list, tuple)):
return False
if len(value) != self.shape[0]:
return False
for item in value:
if isinstance(item, bool) or not isinstance(item, (int, float)):
return False
if not math.isfinite(float(item)) or not self.low <= float(item) <= self.high:
return False
return True
class SyntheticGymEnv:
"""Small deterministic environment implementing Gymnasium's five-tuple API."""
metadata = {"render_modes": []}
def __init__(self, *, max_steps: int = 16) -> None:
if not 1 <= max_steps <= 10_000:
raise UserInputError("max_steps must be between 1 and 10000")
self.max_steps = max_steps
self.observation_space = BoxSpace(-1.0, 1.0, (4,))
self.action_space = DiscreteSpace(3)
self._rng = random.Random()
self._initialized = False
self._done = False
self._position = 0.0
self._target = 0.75
self._step_count = 0
self._last_action = 0.0
# Define observation space
# Option 1: Flat vector observation
self.observation_space = self.make_space((4,)) # [x, y, goal_x, goal_y]
def _observation(self) -> list[float]:
return [
float(self._position),
float(self._target),
float(self._step_count / self.max_steps),
float(self._last_action),
]
# Option 2: Dict observation with multiple components
# self.observation_space = self.make_space({
# 'position': (2,),
# 'goal': (2,),
# 'grid': (grid_size, grid_size)
# })
def reset(
self,
*,
seed: int | None = None,
options: dict[str, Any] | None = None,
) -> tuple[list[float], dict[str, Any]]:
"""Reset state and return ``(observation, info)``."""
if seed is not None:
self._rng.seed(seed)
if options is not None and set(options) - {"position", "target"}:
raise UserInputError("reset options may contain only position and target")
# Option 3: Image observation
# self.observation_space = self.make_space((grid_size, grid_size, 3))
self._position = self._rng.uniform(-0.5, 0.5)
self._target = self._rng.choice((-0.75, 0.75))
if options:
self._position = float(options.get("position", self._position))
self._target = float(options.get("target", self._target))
if not -1.0 <= self._position <= 1.0 or not -1.0 <= self._target <= 1.0:
raise UserInputError("position and target options must be within [-1, 1]")
# Define action space
# Option 1: Discrete actions
self.action_space = self.make_discrete(4) # 0: up, 1: right, 2: down, 3: left
self._step_count = 0
self._last_action = 0.0
self._initialized = True
self._done = False
observation = self._observation()
return observation, {"seed": seed, "synthetic": True}
# Option 2: Continuous actions
# self.action_space = self.make_space((2,)) # [dx, dy]
def step(
self, action: int
) -> tuple[list[float], float, bool, bool, dict[str, Any]]:
"""Advance one step and return the Gymnasium five-tuple."""
if not self._initialized:
raise RuntimeError("reset() must be called before step()")
if self._done:
raise RuntimeError("reset() must be called after termination or truncation")
if not self.action_space.contains(action):
raise ValueError(f"action {action!r} is outside the action space")
# Option 3: Multi-discrete actions
# self.action_space = self.make_multi_discrete([3, 3]) # Two 3-way choices
movement = (-0.125, 0.0, 0.125)[action]
self._position = max(-1.0, min(1.0, self._position + movement))
self._last_action = movement / 0.125
self._step_count += 1
# Initialize state
self.agent_pos = None
self.goal_pos = None
self.step_count = 0
self.reset()
def reset(self):
"""
Reset environment to initial state.
Returns:
observation: Initial observation
"""
# Reset state
self.agent_pos = np.array([0, 0], dtype=np.float32)
self.goal_pos = np.array([self.grid_size - 1, self.grid_size - 1], dtype=np.float32)
self.step_count = 0
# Return initial observation
return self._get_observation()
def step(self, action):
"""
Execute one environment step.
Args:
action: Action to take
Returns:
observation: New observation
reward: Reward for this step
done: Whether episode is complete
info: Additional information
"""
self.step_count += 1
# Execute action
self._apply_action(action)
# Compute reward
reward = self._compute_reward()
# Check if episode is done
done = self._is_done()
# Get new observation
observation = self._get_observation()
# Additional info
info = {}
if done:
# Include episode statistics when episode ends
info['episode'] = {
'r': reward,
'l': self.step_count
}
return observation, reward, done, info
def _apply_action(self, action):
"""Apply action to update environment state."""
# Discrete actions: 0=up, 1=right, 2=down, 3=left
if action == 0: # Up
self.agent_pos[1] = min(self.agent_pos[1] + 1, self.grid_size - 1)
elif action == 1: # Right
self.agent_pos[0] = min(self.agent_pos[0] + 1, self.grid_size - 1)
elif action == 2: # Down
self.agent_pos[1] = max(self.agent_pos[1] - 1, 0)
elif action == 3: # Left
self.agent_pos[0] = max(self.agent_pos[0] - 1, 0)
def _compute_reward(self):
"""Compute reward for current state."""
# Distance to goal
distance = np.linalg.norm(self.agent_pos - self.goal_pos)
# Reward shaping: negative distance + bonus for reaching goal
reward = -distance / self.grid_size
# Goal reached
if distance < 0.5:
reward += 10.0
return reward
def _is_done(self):
"""Check if episode is complete."""
# Episode ends if goal reached or max steps exceeded
distance = np.linalg.norm(self.agent_pos - self.goal_pos)
goal_reached = distance < 0.5
timeout = self.step_count >= self.max_steps
return goal_reached or timeout
def _get_observation(self):
"""Generate observation from current state."""
# Return flat vector observation
observation = np.concatenate([
self.agent_pos,
self.goal_pos
]).astype(np.float32)
return observation
class MultiAgentEnvironment(PufferEnv):
"""
Multi-agent environment template.
Example: Cooperative navigation task where agents must reach individual goals.
"""
def __init__(self, buf=None, num_agents=4, grid_size=10, max_steps=1000):
super().__init__(buf)
self.num_agents = num_agents
self.grid_size = grid_size
self.max_steps = max_steps
# Per-agent observation space
self.single_observation_space = self.make_space({
'position': (2,),
'goal': (2,),
'others': (2 * (num_agents - 1),) # Positions of other agents
})
# Per-agent action space
self.single_action_space = self.make_discrete(5) # 4 directions + stay
# Initialize state
self.agent_positions = None
self.goal_positions = None
self.step_count = 0
self.reset()
def reset(self):
"""Reset all agents."""
# Random initial positions
self.agent_positions = np.random.rand(self.num_agents, 2) * self.grid_size
# Random goal positions
self.goal_positions = np.random.rand(self.num_agents, 2) * self.grid_size
self.step_count = 0
# Return observations for all agents
return {
f'agent_{i}': self._get_obs(i)
for i in range(self.num_agents)
distance = abs(self._target - self._position)
terminated = distance <= 0.0625
truncated = self._step_count >= self.max_steps and not terminated
reward = 1.0 if terminated else -distance
self._done = terminated or truncated
info = {
"distance": float(distance),
"episode_step": self._step_count,
}
return self._observation(), float(reward), terminated, truncated, info
def step(self, actions):
"""
Step all agents.
def close(self) -> None:
self._initialized = False
self._done = True
Args:
actions: Dict of {agent_id: action}
Returns:
observations: Dict of {agent_id: observation}
rewards: Dict of {agent_id: reward}
dones: Dict of {agent_id: done}
infos: Dict of {agent_id: info}
"""
self.step_count += 1
def run_demo(*, seed: int, steps: int, max_steps: int) -> dict[str, Any]:
"""Run a bounded deterministic rollout for documentation and smoke tests."""
env = SyntheticGymEnv(max_steps=max_steps)
action_rng = random.Random(seed + 1)
observation, _ = env.reset(seed=seed)
total_reward = 0.0
resets = 0
terminated_count = 0
truncated_count = 0
observations = {}
rewards = {}
dones = {}
infos = {}
for index in range(steps):
action = env.action_space.sample(action_rng)
observation, reward, terminated, truncated, _ = env.step(action)
total_reward += reward
terminated_count += int(terminated)
truncated_count += int(truncated)
if terminated or truncated:
resets += 1
observation, _ = env.reset(seed=seed + resets + index + 1)
# Update all agents
for agent_id, action in actions.items():
agent_idx = int(agent_id.split('_')[1])
env.close()
return {
"environment": "synthetic",
"last_observation": observation,
"network_used": False,
"resets": resets,
"seed": seed,
"steps": steps,
"terminated": terminated_count,
"total_reward": total_reward,
"truncated": truncated_count,
}
# Apply action
self._apply_action(agent_idx, action)
# Generate outputs
observations[agent_id] = self._get_obs(agent_idx)
rewards[agent_id] = self._compute_reward(agent_idx)
dones[agent_id] = self._is_done(agent_idx)
infos[agent_id] = {}
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description="Run the dependency-free synthetic environment template."
)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--steps", type=int, default=16, help="1..10000")
parser.add_argument("--max-steps", type=int, default=16, help="1..10000")
parser.add_argument("--compact", action="store_true", help="Emit compact JSON")
return parser
# Global done condition
dones['__all__'] = all(dones.values()) or self.step_count >= self.max_steps
return observations, rewards, dones, infos
def _apply_action(self, agent_idx, action):
"""Apply action for specific agent."""
if action == 0: # Up
self.agent_positions[agent_idx, 1] += 1
elif action == 1: # Right
self.agent_positions[agent_idx, 0] += 1
elif action == 2: # Down
self.agent_positions[agent_idx, 1] -= 1
elif action == 3: # Left
self.agent_positions[agent_idx, 0] -= 1
# action == 4: Stay
# Clip to grid bounds
self.agent_positions[agent_idx] = np.clip(
self.agent_positions[agent_idx],
0,
self.grid_size - 1
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
try:
steps = bounded_int(args.steps, name="steps", minimum=1, maximum=10_000)
max_steps = bounded_int(
args.max_steps, name="max_steps", minimum=1, maximum=10_000
)
def _compute_reward(self, agent_idx):
"""Compute reward for specific agent."""
distance = np.linalg.norm(
self.agent_positions[agent_idx] - self.goal_positions[agent_idx]
)
return -distance / self.grid_size
def _is_done(self, agent_idx):
"""Check if specific agent is done."""
distance = np.linalg.norm(
self.agent_positions[agent_idx] - self.goal_positions[agent_idx]
)
return distance < 0.5
def _get_obs(self, agent_idx):
"""Get observation for specific agent."""
# Get positions of other agents
other_positions = np.concatenate([
self.agent_positions[i]
for i in range(self.num_agents)
if i != agent_idx
])
return {
'position': self.agent_positions[agent_idx].astype(np.float32),
'goal': self.goal_positions[agent_idx].astype(np.float32),
'others': other_positions.astype(np.float32)
}
result = run_demo(seed=args.seed, steps=steps, max_steps=max_steps)
except (UserInputError, ValueError, RuntimeError) as exc:
parser.error(str(exc))
emit_json(result, pretty=not args.compact)
return 0
def test_environment():
"""Test environment to verify it works correctly."""
print("Testing single-agent environment...")
env = MyEnvironment()
obs = env.reset()
print(f"Initial observation shape: {obs.shape}")
for step in range(10):
action = env.action_space.sample()
obs, reward, done, info = env.step(action)
print(f"Step {step}: reward={reward:.3f}, done={done}")
if done:
obs = env.reset()
print("Episode finished, resetting...")
print("\nTesting multi-agent environment...")
multi_env = MultiAgentEnvironment(num_agents=4)
obs = multi_env.reset()
print(f"Number of agents: {len(obs)}")
for step in range(10):
actions = {
agent_id: multi_env.single_action_space.sample()
for agent_id in obs.keys()
}
obs, rewards, dones, infos = multi_env.step(actions)
print(f"Step {step}: mean_reward={np.mean(list(rewards.values())):.3f}")
if dones.get('__all__', False):
obs = multi_env.reset()
print("Episode finished, resetting...")
print("\n✓ Environment tests passed!")
if __name__ == '__main__':
test_environment()
if __name__ == "__main__":
raise SystemExit(main())

225
skills/pufferlib/scripts/inspect_checkpoint.py Normal file
View File

@@ -0,0 +1,225 @@
#!/usr/bin/env python3
"""Inspect checkpoint file metadata without deserializing checkpoint contents."""
from __future__ import annotations
import argparse
import hashlib
import os
import stat
from pathlib import Path
from typing import Any, BinaryIO
try:
from ._common import (
UserInputError,
bounded_int,
emit_json,
load_json_object,
resolve_local_path,
secret_key_paths,
validate_sha256,
)
except ImportError: # Direct script execution.
from _common import (
UserInputError,
bounded_int,
emit_json,
load_json_object,
resolve_local_path,
secret_key_paths,
validate_sha256,
)
_SIDECAR_FIELDS = {
"created_at",
"environment",
"format",
"framework",
"framework_version",
"license",
"notes",
"parent_sha256",
"policy",
"schema_version",
"seed",
"sha256",
"source_commit",
"source_url",
"training_steps",
}
def _open_regular_no_follow(path: Path) -> tuple[BinaryIO, os.stat_result]:
flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
try:
descriptor = os.open(path, flags)
except OSError as exc:
raise UserInputError(f"cannot open checkpoint safely: {path}") from exc
try:
file_stat = os.fstat(descriptor)
if not stat.S_ISREG(file_stat.st_mode):
raise UserInputError("checkpoint must be a regular file")
return os.fdopen(descriptor, "rb"), file_stat
except Exception:
os.close(descriptor)
raise
def _detect_format(prefix: bytes, suffix: str) -> dict[str, Any]:
if prefix.startswith(b"PK\x03\x04"):
return {
"family": "zip-container",
"risk": "may contain a torch.save pickle payload; not opened",
}
if prefix.startswith(b"\x80"):
protocol = prefix[1] if len(prefix) > 1 else None
return {
"family": "pickle-like",
"pickle_protocol_byte": protocol,
"risk": "unsafe to deserialize unless provenance is trusted",
}
if suffix.lower() == ".bin":
return {
"family": "opaque-bin",
"risk": "could be PufferLib native weights or another binary format",
}
return {
"family": "opaque",
"risk": "format not identified; no deserialization attempted",
}
def _hash_and_prefix(handle: BinaryIO, *, chunk_bytes: int = 1_048_576) -> tuple[str, bytes]:
digest = hashlib.sha256()
prefix = b""
while True:
chunk = handle.read(chunk_bytes)
if not chunk:
break
if not prefix:
prefix = chunk[:16]
digest.update(chunk)
return digest.hexdigest(), prefix
def _safe_sidecar(
metadata_path: str | None, *, root: str | Path
) -> tuple[dict[str, Any] | None, list[str]]:
if metadata_path is None:
return None, []
raw = load_json_object(metadata_path, root=root, max_bytes=262_144)
secret_paths = secret_key_paths(raw)
if secret_paths:
raise UserInputError(
"sidecar contains credential-bearing keys: " + ", ".join(secret_paths)
)
safe = {key: raw[key] for key in sorted(raw) if key in _SIDECAR_FIELDS}
unknown = sorted(set(raw) - _SIDECAR_FIELDS)
return safe, unknown
def inspect_checkpoint(
checkpoint_path: str,
*,
root: str | Path,
metadata_path: str | None,
expected_sha256: str | None,
max_bytes: int,
) -> dict[str, Any]:
"""Hash and classify one local regular file without importing torch or pickle."""
resolved = resolve_local_path(
checkpoint_path,
root=root,
must_exist=True,
reject_symlink=True,
)
handle, file_stat = _open_regular_no_follow(resolved)
with handle:
if file_stat.st_size > max_bytes:
raise UserInputError(
f"checkpoint is {file_stat.st_size} bytes; cap is {max_bytes}"
)
digest, prefix = _hash_and_prefix(handle)
expected = validate_sha256(expected_sha256) if expected_sha256 else None
sidecar, unknown_fields = _safe_sidecar(metadata_path, root=root)
sidecar_digest = sidecar.get("sha256") if sidecar else None
if sidecar_digest is not None:
validate_sha256(sidecar_digest, name="sidecar sha256")
return {
"checkpoint": {
"format_detection": _detect_format(prefix, resolved.suffix),
"name": resolved.name,
"sha256": digest,
"size_bytes": file_stat.st_size,
},
"deserialized": False,
"expected_sha256_matches": None if expected is None else digest == expected,
"metadata": sidecar,
"metadata_sha256_matches": (
None if sidecar_digest is None else digest == sidecar_digest
),
"network_used": False,
"sidecar_unknown_fields_omitted": unknown_fields,
"warning": (
"Inspection does not establish trust. Verify source, license, signature or "
"attestation, and checksum before sandboxed loading."
),
}
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description=(
"Hash and classify one local checkpoint without torch.load, pickle, "
"archive extraction, dynamic imports, or network access."
)
)
parser.add_argument("checkpoint", help="Checkpoint path beneath --root")
parser.add_argument("--root", default=".", help="Allowed local path root")
parser.add_argument("--metadata", help="Explicit strict-JSON sidecar beneath --root")
parser.add_argument("--expected-sha256")
parser.add_argument(
"--max-bytes",
type=int,
default=2_147_483_648,
help="1..68719476736",
)
parser.add_argument("--compact", action="store_true")
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
try:
max_bytes = bounded_int(
args.max_bytes,
name="max_bytes",
minimum=1,
maximum=68_719_476_736,
)
report = inspect_checkpoint(
args.checkpoint,
root=args.root,
metadata_path=args.metadata,
expected_sha256=args.expected_sha256,
max_bytes=max_bytes,
)
except (UserInputError, OSError, ValueError) as exc:
report = {
"deserialized": False,
"errors": [str(exc)],
"network_used": False,
"status": "invalid",
}
emit_json(report, pretty=not args.compact)
return 1
emit_json(report, pretty=not args.compact)
return 0
if __name__ == "__main__":
raise SystemExit(main())

177
skills/pufferlib/scripts/repro_plan.py Normal file
View File

@@ -0,0 +1,177 @@
#!/usr/bin/env python3
"""Generate a bounded reproducibility and held-out evaluation plan."""
from __future__ import annotations
import argparse
from typing import Any
try:
from ._common import (
SOURCE_4_COMMIT,
STABLE_SDIST_SHA256,
UserInputError,
bounded_int,
emit_json,
validate_slug,
)
from .validate_plan import PROFILES
except ImportError: # Direct script execution.
from _common import (
SOURCE_4_COMMIT,
STABLE_SDIST_SHA256,
UserInputError,
bounded_int,
emit_json,
validate_slug,
)
from validate_plan import PROFILES
def generate_plan(
*,
profile: str,
environment: str,
base_seed: int,
replicates: int,
eval_episodes: int,
benchmark_repeats: int,
) -> dict[str, Any]:
"""Create non-overlapping train/eval seeds and reporting requirements."""
train_seeds = [base_seed + index for index in range(replicates)]
eval_seed_base = base_seed + 1_000_000
eval_seeds = [eval_seed_base + index for index in range(replicates)]
if set(train_seeds) & set(eval_seeds):
raise UserInputError("training and evaluation seeds overlap")
if profile == "pypi-3.0.0":
upstream = {
"artifact": "pufferlib-3.0.0.tar.gz",
"package": "pufferlib==3.0.0",
"python": ">=3.9",
"sha256": STABLE_SDIST_SHA256,
"warning": (
"PyPI provides only an sdist. Its build can download and compile native "
"dependencies; audit and sandbox the build before installation."
),
}
else:
upstream = {
"commit": SOURCE_4_COMMIT,
"package": "pufferlib source 4.0",
"python": ">=3.10",
"torch": ">=2.9",
"warning": (
"The 4.0 default branch is not the latest stable PyPI artifact. Pin the "
"commit and use an audited CUDA/CPU build environment."
),
}
return {
"benchmarking": {
"aggregate": ["median", "p10", "p90"],
"exclude_setup": False,
"fixed_workload": True,
"record": [
"CPU model and logical/physical cores",
"GPU model, driver, CUDA, and precision when applicable",
"OS, Python, PufferLib, NumPy, Gymnasium, and PyTorch versions",
"backend, start method, workers, envs, buffers, and batch size",
"warmup, repeats, wall time, agent steps, and reset count",
],
"repeats": benchmark_repeats,
"warning": "Do not compare SPS across changed workloads or hardware.",
},
"environment": {
"name": environment,
"record": [
"source URL and immutable revision",
"license and asset/ROM rights",
"environment and wrapper configuration",
"observation/action spaces and dtypes",
"termination, truncation, autoreset, and frame-skip semantics",
],
},
"evaluation": {
"checkpoint_selected_without_eval_feedback": True,
"deterministic_policy_pass": True,
"episodes_per_seed": eval_episodes,
"learning_disabled": True,
"report_per_seed_and_aggregate": True,
"seeds": eval_seeds,
"separate_environment_instances": True,
"stochastic_policy_pass": True,
},
"network_used": False,
"profile": profile,
"provenance": {
"checkpoint_sha256_required": True,
"dependency_lock_required": True,
"record_git_diff": True,
"record_source_commit": True,
"upstream": upstream,
},
"schema_version": 1,
"training": {
"determinism_limitations_recorded": True,
"seeds": train_seeds,
},
}
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description=(
"Emit a local reproducibility/evaluation plan; no packages, checkpoints, "
"environments, GPU, or network services are opened."
)
)
parser.add_argument("--profile", choices=PROFILES, default="pypi-3.0.0")
parser.add_argument("--environment", default="synthetic")
parser.add_argument("--base-seed", type=int, default=42)
parser.add_argument("--replicates", type=int, default=3, help="1..32")
parser.add_argument("--eval-episodes", type=int, default=100, help="1..10000")
parser.add_argument("--benchmark-repeats", type=int, default=5, help="3..20")
parser.add_argument("--compact", action="store_true")
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
try:
environment = validate_slug(args.environment, name="environment")
base_seed = bounded_int(
args.base_seed, name="base_seed", minimum=0, maximum=2**32 - 1_000_033
)
replicates = bounded_int(
args.replicates, name="replicates", minimum=1, maximum=32
)
eval_episodes = bounded_int(
args.eval_episodes,
name="eval_episodes",
minimum=1,
maximum=10_000,
)
repeats = bounded_int(
args.benchmark_repeats,
name="benchmark_repeats",
minimum=3,
maximum=20,
)
plan = generate_plan(
profile=args.profile,
environment=environment,
base_seed=base_seed,
replicates=replicates,
eval_episodes=eval_episodes,
benchmark_repeats=repeats,
)
except (UserInputError, ValueError) as exc:
parser.error(str(exc))
emit_json(plan, pretty=not args.compact)
return 0
if __name__ == "__main__":
raise SystemExit(main())

487
skills/pufferlib/scripts/train_template.py
View File

@@ -1,239 +1,282 @@
#!/usr/bin/env python3
"""
PufferLib Training Template
"""Safe PufferLib training-plan template.
This template provides a complete training script for reinforcement learning
with PufferLib. Customize the environment, policy, and training configuration
as needed for your use case.
This script never imports PufferLib, starts training, loads checkpoints, uses a
GPU, or contacts an external logger. It emits a validated argv preview for a
human to review in an appropriately sandboxed, pinned environment.
"""
from __future__ import annotations
import argparse
import torch
import torch.nn as nn
import pufferlib
from pufferlib import PuffeRL
from pufferlib.pytorch import layer_init
import copy
from typing import Any
try:
from ._common import LOGGER_CREDENTIAL_ENV, UserInputError, bounded_int, emit_json, validate_slug
from .validate_plan import PROFILES, default_plan, validate_plan
except ImportError: # Direct script execution.
from _common import LOGGER_CREDENTIAL_ENV, UserInputError, bounded_int, emit_json, validate_slug
from validate_plan import PROFILES, default_plan, validate_plan
class Policy(nn.Module):
"""Example policy network."""
def _command_preview(plan: dict[str, Any]) -> list[str]:
"""Build argv without shell interpolation or execution."""
environment = plan["environment"]["name"]
if environment == "synthetic":
return []
def __init__(self, observation_space, action_space, hidden_size=256):
super().__init__()
self.observation_space = observation_space
self.action_space = action_space
# Encoder network
self.encoder = nn.Sequential(
layer_init(nn.Linear(observation_space.shape[0], hidden_size)),
nn.ReLU(),
layer_init(nn.Linear(hidden_size, hidden_size)),
nn.ReLU()
)
# Policy head (actor)
self.actor = layer_init(nn.Linear(hidden_size, action_space.n), std=0.01)
# Value head (critic)
self.critic = layer_init(nn.Linear(hidden_size, 1), std=1.0)
def forward(self, observations):
"""Forward pass through policy."""
features = self.encoder(observations)
logits = self.actor(features)
value = self.critic(features)
return logits, value
def make_env():
"""Create environment. Customize this for your task."""
# Option 1: Use Ocean environment
return pufferlib.make('procgen-coinrun', num_envs=256)
# Option 2: Use Gymnasium environment
# return pufferlib.make('gym-CartPole-v1', num_envs=256)
# Option 3: Use custom environment
# from my_envs import MyEnvironment
# return pufferlib.emulate(MyEnvironment, num_envs=256)
def create_policy(env):
"""Create policy network."""
return Policy(
observation_space=env.observation_space,
action_space=env.action_space,
hidden_size=256
)
def train(args):
"""Main training function."""
# Set random seeds
torch.manual_seed(args.seed)
# Create environment
print(f"Creating environment with {args.num_envs} parallel environments...")
env = pufferlib.make(
args.env_name,
num_envs=args.num_envs,
num_workers=args.num_workers
)
# Create policy
print("Initializing policy...")
policy = create_policy(env)
if args.device == 'cuda' and torch.cuda.is_available():
policy = policy.cuda()
print(f"Using GPU: {torch.cuda.get_device_name(0)}")
else:
print("Using CPU")
# Create logger
if args.logger == 'wandb':
from pufferlib import WandbLogger
logger = WandbLogger(
project=args.project,
name=args.exp_name,
config=vars(args)
)
elif args.logger == 'neptune':
from pufferlib import NeptuneLogger
logger = NeptuneLogger(
project=args.project,
name=args.exp_name,
api_token=args.neptune_token
training = plan["training"]
vector = plan["vectorization"]
command = [
"puffer",
"train",
environment,
"--train.total-timesteps",
str(training["total_timesteps"]),
"--train.seed",
str(training["seed"]),
]
if plan["profile"] == "pypi-3.0.0":
backend_name = {
"serial": "Serial",
"multiprocessing": "Multiprocessing",
"native": "PufferEnv",
}[vector["backend"]]
command.extend(
[
"--train.device",
training["device"],
"--vec.backend",
backend_name,
"--vec.num-envs",
str(vector["num_envs"]),
"--vec.num-workers",
str(vector["num_workers"]),
"--vec.batch-size",
str(vector["batch_size"]),
]
)
else:
from pufferlib import NoLogger
logger = NoLogger()
command.extend(
[
"--vec.total-agents",
str(vector["total_agents"]),
"--vec.num-buffers",
str(vector["num_buffers"]),
"--vec.num-threads",
str(vector["num_threads"]),
]
)
if vector["backend"] == "torch":
command.append("--slowly")
# Create trainer
print("Creating trainer...")
trainer = PuffeRL(
env=env,
policy=policy,
device=args.device,
learning_rate=args.learning_rate,
batch_size=args.batch_size,
n_epochs=args.n_epochs,
gamma=args.gamma,
gae_lambda=args.gae_lambda,
clip_coef=args.clip_coef,
ent_coef=args.ent_coef,
vf_coef=args.vf_coef,
max_grad_norm=args.max_grad_norm,
logger=logger,
compile=args.compile
logger = plan["logging"]["backend"]
if logger != "none":
command.append(f"--{logger}")
if not plan["logging"]["upload_checkpoints"] and plan["profile"] == "pypi-3.0.0":
command.append("--no-model-upload")
return command
def make_plan(args: argparse.Namespace) -> dict[str, Any]:
plan = default_plan(args.profile)
environment = validate_slug(args.environment, name="environment")
plan["environment"]["name"] = environment
plan["environment"]["adapter"] = args.adapter
plan["environment"]["provenance_verified"] = (
environment == "synthetic" or args.provenance_verified
)
# Training loop
print(f"Starting training for {args.num_iterations} iterations...")
for iteration in range(1, args.num_iterations + 1):
# Collect rollouts
rollout_data = trainer.evaluate()
# Train on batch
train_metrics = trainer.train()
# Log results
trainer.mean_and_log()
# Save checkpoint
if iteration % args.save_freq == 0:
checkpoint_path = f"{args.checkpoint_dir}/checkpoint_{iteration}.pt"
trainer.save_checkpoint(checkpoint_path)
print(f"Saved checkpoint to {checkpoint_path}")
# Print progress
if iteration % args.log_freq == 0:
mean_reward = rollout_data.get('mean_reward', 0)
sps = rollout_data.get('sps', 0)
print(f"Iteration {iteration}/{args.num_iterations} | "
f"Mean Reward: {mean_reward:.2f} | "
f"SPS: {sps:,.0f}")
print("Training complete!")
# Save final model
final_path = f"{args.checkpoint_dir}/final_model.pt"
trainer.save_checkpoint(final_path)
print(f"Saved final model to {final_path}")
plan["training"].update(
{
"device": args.device,
"horizon": bounded_int(
args.horizon, name="horizon", minimum=1, maximum=65_536
),
"minibatch_size": bounded_int(
args.minibatch_size,
name="minibatch_size",
minimum=1,
maximum=16_777_216,
),
"seed": bounded_int(
args.seed, name="seed", minimum=0, maximum=2**32 - 1
),
"total_timesteps": bounded_int(
args.total_timesteps,
name="total_timesteps",
minimum=1,
maximum=1_000_000_000,
),
}
)
plan["evaluation"].update(
{
"deterministic": args.deterministic_eval,
"episodes": bounded_int(
args.eval_episodes,
name="eval_episodes",
minimum=1,
maximum=10_000,
),
"seed": bounded_int(
args.eval_seed, name="eval_seed", minimum=0, maximum=2**32 - 1
),
"separate": True,
}
)
plan["logging"].update(
{
"backend": args.logger,
"disclosure_ack": args.acknowledge_external_disclosure,
"external_opt_in": args.enable_external_logging,
"upload_checkpoints": args.upload_checkpoints,
}
)
if args.profile == "pypi-3.0.0":
backend = args.backend or "serial"
plan["vectorization"].update(
{
"backend": backend,
"batch_size": bounded_int(
args.batch_size,
name="batch_size",
minimum=1,
maximum=65_536,
),
"num_envs": bounded_int(
args.num_envs, name="num_envs", minimum=1, maximum=65_536
),
"num_workers": bounded_int(
args.num_workers,
name="num_workers",
minimum=1,
maximum=256,
),
"start_method": args.start_method,
"zero_copy": args.zero_copy,
}
)
else:
backend = args.backend or "torch"
plan["vectorization"].update(
{
"backend": backend,
"num_buffers": bounded_int(
args.num_buffers, name="num_buffers", minimum=1, maximum=256
),
"num_threads": bounded_int(
args.num_threads, name="num_threads", minimum=1, maximum=256
),
"start_method": "spawn",
"total_agents": bounded_int(
args.total_agents,
name="total_agents",
minimum=1,
maximum=65_536,
),
}
)
plan["checkpoint"]["format"] = (
"state_dict" if backend == "torch" else "native-bin"
)
return plan
def main():
parser = argparse.ArgumentParser(description='PufferLib Training')
# Environment
parser.add_argument('--env-name', type=str, default='procgen-coinrun',
help='Environment name')
parser.add_argument('--num-envs', type=int, default=256,
help='Number of parallel environments')
parser.add_argument('--num-workers', type=int, default=8,
help='Number of vectorization workers')
# Training
parser.add_argument('--num-iterations', type=int, default=10000,
help='Number of training iterations')
parser.add_argument('--learning-rate', type=float, default=3e-4,
help='Learning rate')
parser.add_argument('--batch-size', type=int, default=32768,
help='Batch size for training')
parser.add_argument('--n-epochs', type=int, default=4,
help='Number of training epochs per batch')
parser.add_argument('--device', type=str, default='cuda',
choices=['cuda', 'cpu'], help='Device to use')
# PPO Parameters
parser.add_argument('--gamma', type=float, default=0.99,
help='Discount factor')
parser.add_argument('--gae-lambda', type=float, default=0.95,
help='GAE lambda')
parser.add_argument('--clip-coef', type=float, default=0.2,
help='PPO clipping coefficient')
parser.add_argument('--ent-coef', type=float, default=0.01,
help='Entropy coefficient')
parser.add_argument('--vf-coef', type=float, default=0.5,
help='Value function coefficient')
parser.add_argument('--max-grad-norm', type=float, default=0.5,
help='Maximum gradient norm')
# Logging
parser.add_argument('--logger', type=str, default='none',
choices=['wandb', 'neptune', 'none'],
help='Logger to use')
parser.add_argument('--project', type=str, default='pufferlib-training',
help='Project name for logging')
parser.add_argument('--exp-name', type=str, default='experiment',
help='Experiment name')
parser.add_argument('--neptune-token', type=str, default=None,
help='Neptune API token')
parser.add_argument('--log-freq', type=int, default=10,
help='Logging frequency (iterations)')
# Checkpointing
parser.add_argument('--checkpoint-dir', type=str, default='checkpoints',
help='Directory to save checkpoints')
parser.add_argument('--save-freq', type=int, default=100,
help='Checkpoint save frequency (iterations)')
# Misc
parser.add_argument('--seed', type=int, default=42,
help='Random seed')
parser.add_argument('--compile', action='store_true',
help='Use torch.compile for faster training')
args = parser.parse_args()
# Create checkpoint directory
import os
os.makedirs(args.checkpoint_dir, exist_ok=True)
# Run training
train(args)
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description=(
"Create a bounded, local PufferLib dry-run plan. This command never "
"executes the generated argv."
)
)
parser.add_argument("--profile", choices=PROFILES, default="pypi-3.0.0")
parser.add_argument("--environment", default="synthetic")
parser.add_argument(
"--adapter",
choices=["synthetic", "gymnasium", "pettingzoo", "native-ocean"],
default="synthetic",
)
parser.add_argument(
"--provenance-verified",
action="store_true",
help="Explicitly attest review for a non-synthetic environment",
)
parser.add_argument(
"--backend",
choices=["serial", "multiprocessing", "native", "torch"],
default=None,
)
parser.add_argument("--num-envs", type=int, default=4)
parser.add_argument("--num-workers", type=int, default=1)
parser.add_argument("--batch-size", type=int, default=4)
parser.add_argument("--zero-copy", action="store_true")
parser.add_argument(
"--start-method", choices=["spawn", "forkserver"], default="spawn"
)
parser.add_argument("--total-agents", type=int, default=64)
parser.add_argument("--num-buffers", type=int, default=2)
parser.add_argument("--num-threads", type=int, default=1)
parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
parser.add_argument("--total-timesteps", type=int, default=10_000)
parser.add_argument("--horizon", type=int, default=16)
parser.add_argument("--minibatch-size", type=int, default=256)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--eval-seed", type=int, default=1_000_042)
parser.add_argument("--eval-episodes", type=int, default=10)
parser.add_argument(
"--stochastic-eval",
dest="deterministic_eval",
action="store_false",
default=True,
)
parser.add_argument("--logger", choices=["none", "wandb", "neptune"], default="none")
parser.add_argument("--enable-external-logging", action="store_true")
parser.add_argument("--acknowledge-external-disclosure", action="store_true")
parser.add_argument("--upload-checkpoints", action="store_true")
parser.add_argument("--compact", action="store_true")
return parser
if __name__ == '__main__':
main()
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
try:
plan = make_plan(args)
errors = validate_plan(plan)
logger = plan["logging"]["backend"]
report = {
"command_preview": _command_preview(plan) if not errors else [],
"credential": {
"environment_variable": LOGGER_CREDENTIAL_ENV.get(logger),
"value_read_or_logged": False,
},
"dry_run": True,
"errors": errors,
"external_logging_disclosure": (
"External services may receive configuration, metrics, source metadata, "
"hardware telemetry, and explicitly enabled artifacts; review vendor "
"privacy, retention, access, and pricing before use."
if logger != "none"
else None
),
"network_used": False,
"plan": copy.deepcopy(plan),
"status": "valid" if not errors else "invalid",
}
except (UserInputError, KeyError, ValueError) as exc:
report = {
"command_preview": [],
"dry_run": True,
"errors": [str(exc)],
"network_used": False,
"plan": None,
"status": "invalid",
}
emit_json(report, pretty=not args.compact)
return 0 if report["status"] == "valid" else 1
if __name__ == "__main__":
raise SystemExit(main())

570
skills/pufferlib/scripts/validate_plan.py Normal file
View File

@@ -0,0 +1,570 @@
#!/usr/bin/env python3
"""Strict, dependency-free validator for PufferLib training plans."""
from __future__ import annotations
import argparse
import copy
import re
from typing import Any
try:
from ._common import (
LOGGER_CREDENTIAL_ENV,
MAX_ENVS,
MAX_EVAL_EPISODES,
MAX_STEPS,
MAX_WORKERS,
SOURCE_4_COMMIT,
STABLE_SDIST_SHA256,
UserInputError,
emit_json,
load_json_object,
require_keys,
secret_key_paths,
validate_slug,
)
except ImportError: # Direct script execution.
from _common import (
LOGGER_CREDENTIAL_ENV,
MAX_ENVS,
MAX_EVAL_EPISODES,
MAX_STEPS,
MAX_WORKERS,
SOURCE_4_COMMIT,
STABLE_SDIST_SHA256,
UserInputError,
emit_json,
load_json_object,
require_keys,
secret_key_paths,
validate_slug,
)
PROFILES = ("pypi-3.0.0", "source-4.0")
_COMMIT = re.compile(r"^[0-9a-f]{40}$")
def default_plan(profile: str = "pypi-3.0.0") -> dict[str, Any]:
"""Return a bounded local plan that never starts training."""
if profile not in PROFILES:
raise UserInputError(f"profile must be one of {PROFILES}")
common: dict[str, Any] = {
"schema_version": 1,
"profile": profile,
"environment": {
"adapter": "synthetic",
"name": "synthetic",
"provenance_verified": True,
},
"training": {
"device": "cpu",
"horizon": 16,
"minibatch_size": 256,
"seed": 42,
"total_timesteps": 10_000,
},
"evaluation": {
"deterministic": True,
"episodes": 10,
"seed": 1_000_042,
"separate": True,
},
"logging": {
"backend": "none",
"disclosure_ack": False,
"external_opt_in": False,
"upload_checkpoints": False,
},
"checkpoint": {
"format": "state_dict",
"trusted_only": True,
},
}
if profile == "pypi-3.0.0":
common["package"] = {
"name": "pufferlib",
"sha256": STABLE_SDIST_SHA256,
"version": "3.0.0",
}
common["vectorization"] = {
"backend": "serial",
"batch_size": 4,
"num_envs": 4,
"num_workers": 1,
"start_method": "spawn",
"zero_copy": False,
}
else:
common["package"] = {
"commit": SOURCE_4_COMMIT,
"name": "pufferlib",
"version": "4.0.0-source",
}
common["vectorization"] = {
"backend": "torch",
"num_buffers": 2,
"num_threads": 1,
"start_method": "spawn",
"total_agents": 64,
}
return common
def _mapping(value: Any, path: str, errors: list[str]) -> dict[str, Any]:
if not isinstance(value, dict):
errors.append(f"{path} must be an object")
return {}
return value
def _integer(
value: Any,
*,
path: str,
minimum: int,
maximum: int,
errors: list[str],
) -> int | None:
if isinstance(value, bool) or not isinstance(value, int):
errors.append(f"{path} must be an integer")
return None
if not minimum <= value <= maximum:
errors.append(f"{path} must be between {minimum} and {maximum}")
return None
return value
def _boolean(value: Any, *, path: str, errors: list[str]) -> bool | None:
if type(value) is not bool:
errors.append(f"{path} must be boolean")
return None
return value
def _validate_package(
package: dict[str, Any], profile: str, errors: list[str]
) -> None:
if profile == "pypi-3.0.0":
errors.extend(
require_keys(
package,
allowed={"name", "version", "sha256"},
required={"name", "version", "sha256"},
path="$.package",
)
)
if package.get("version") != "3.0.0":
errors.append("$.package.version must equal 3.0.0")
if package.get("sha256") != STABLE_SDIST_SHA256:
errors.append("$.package.sha256 must match the published 3.0.0 sdist")
elif profile == "source-4.0":
errors.extend(
require_keys(
package,
allowed={"name", "version", "commit"},
required={"name", "version", "commit"},
path="$.package",
)
)
commit = package.get("commit")
if not isinstance(commit, str) or not _COMMIT.fullmatch(commit):
errors.append("$.package.commit must be a pinned 40-character commit")
if package.get("version") != "4.0.0-source":
errors.append("$.package.version must equal 4.0.0-source")
if package.get("name") != "pufferlib":
errors.append("$.package.name must equal pufferlib")
def _validate_environment(environment: dict[str, Any], errors: list[str]) -> None:
errors.extend(
require_keys(
environment,
allowed={"adapter", "name", "provenance_verified"},
required={"adapter", "name", "provenance_verified"},
path="$.environment",
)
)
try:
validate_slug(environment.get("name"), name="environment.name")
except UserInputError as exc:
errors.append(str(exc))
if environment.get("adapter") not in {
"synthetic",
"gymnasium",
"pettingzoo",
"native-ocean",
}:
errors.append("$.environment.adapter is not allowlisted")
if environment.get("name") == "synthetic" and environment.get("adapter") != "synthetic":
errors.append("the synthetic environment must use the synthetic adapter")
if environment.get("name") != "synthetic" and environment.get("adapter") == "synthetic":
errors.append("a non-synthetic environment cannot use the synthetic adapter")
if environment.get("provenance_verified") is not True:
errors.append("$.environment.provenance_verified must be true")
def _validate_training(training: dict[str, Any], errors: list[str]) -> None:
errors.extend(
require_keys(
training,
allowed={"device", "horizon", "minibatch_size", "seed", "total_timesteps"},
required={"device", "horizon", "minibatch_size", "seed", "total_timesteps"},
path="$.training",
)
)
if training.get("device") not in {"cpu", "cuda"}:
errors.append("$.training.device must be cpu or cuda")
_integer(
training.get("seed"),
path="$.training.seed",
minimum=0,
maximum=2**32 - 1,
errors=errors,
)
_integer(
training.get("total_timesteps"),
path="$.training.total_timesteps",
minimum=1,
maximum=MAX_STEPS,
errors=errors,
)
horizon = _integer(
training.get("horizon"),
path="$.training.horizon",
minimum=1,
maximum=65_536,
errors=errors,
)
minibatch = _integer(
training.get("minibatch_size"),
path="$.training.minibatch_size",
minimum=1,
maximum=16_777_216,
errors=errors,
)
if horizon and minibatch and minibatch % horizon:
errors.append("$.training.minibatch_size must be divisible by horizon")
def _validate_vectorization(
vectorization: dict[str, Any],
profile: str,
training: dict[str, Any],
errors: list[str],
) -> None:
if profile == "pypi-3.0.0":
allowed = {
"backend",
"batch_size",
"num_envs",
"num_workers",
"start_method",
"zero_copy",
}
errors.extend(
require_keys(
vectorization,
allowed=allowed,
required=allowed,
path="$.vectorization",
)
)
backend = vectorization.get("backend")
if backend not in {
"serial",
"multiprocessing",
"native",
}:
errors.append("$.vectorization.backend is invalid for PufferLib 3.0.0")
num_envs = _integer(
vectorization.get("num_envs"),
path="$.vectorization.num_envs",
minimum=1,
maximum=MAX_ENVS,
errors=errors,
)
workers = _integer(
vectorization.get("num_workers"),
path="$.vectorization.num_workers",
minimum=1,
maximum=MAX_WORKERS,
errors=errors,
)
batch = _integer(
vectorization.get("batch_size"),
path="$.vectorization.batch_size",
minimum=1,
maximum=MAX_ENVS,
errors=errors,
)
zero_copy = _boolean(
vectorization.get("zero_copy"),
path="$.vectorization.zero_copy",
errors=errors,
)
if vectorization.get("start_method") not in {"spawn", "forkserver"}:
errors.append("$.vectorization.start_method must be spawn or forkserver")
if num_envs and workers and num_envs % workers:
errors.append("$.vectorization.num_envs must be divisible by num_workers")
if num_envs and batch and batch > num_envs:
errors.append("$.vectorization.batch_size cannot exceed num_envs")
if num_envs and batch and zero_copy and num_envs % batch:
errors.append(
"$.vectorization.num_envs must be divisible by batch_size "
"when zero_copy is true"
)
if num_envs and workers and batch:
envs_per_worker = num_envs // workers if num_envs % workers == 0 else 0
if envs_per_worker and batch % envs_per_worker:
errors.append(
"$.vectorization.batch_size must be divisible by envs_per_worker"
)
if backend == "native" and num_envs != 1:
errors.append(
"$.vectorization.num_envs must equal 1 for the stable native backend"
)
else:
allowed = {
"backend",
"num_buffers",
"num_threads",
"start_method",
"total_agents",
}
errors.extend(
require_keys(
vectorization,
allowed=allowed,
required=allowed,
path="$.vectorization",
)
)
if vectorization.get("backend") not in {"native", "torch"}:
errors.append("$.vectorization.backend must be native or torch")
agents = _integer(
vectorization.get("total_agents"),
path="$.vectorization.total_agents",
minimum=1,
maximum=MAX_ENVS,
errors=errors,
)
buffers = _integer(
vectorization.get("num_buffers"),
path="$.vectorization.num_buffers",
minimum=1,
maximum=256,
errors=errors,
)
_integer(
vectorization.get("num_threads"),
path="$.vectorization.num_threads",
minimum=1,
maximum=MAX_WORKERS,
errors=errors,
)
if vectorization.get("start_method") != "spawn":
errors.append("$.vectorization.start_method must be spawn for 4.0")
if agents and buffers and agents % buffers:
errors.append("$.vectorization.total_agents must be divisible by num_buffers")
horizon = training.get("horizon")
minibatch = training.get("minibatch_size")
if (
isinstance(agents, int)
and isinstance(horizon, int)
and isinstance(minibatch, int)
and minibatch > agents * horizon
):
errors.append(
"$.training.minibatch_size cannot exceed total_agents * horizon"
)
def _validate_evaluation(
evaluation: dict[str, Any], training: dict[str, Any], errors: list[str]
) -> None:
allowed = {"deterministic", "episodes", "seed", "separate"}
errors.extend(
require_keys(
evaluation,
allowed=allowed,
required=allowed,
path="$.evaluation",
)
)
_boolean(evaluation.get("deterministic"), path="$.evaluation.deterministic", errors=errors)
_boolean(evaluation.get("separate"), path="$.evaluation.separate", errors=errors)
_integer(
evaluation.get("episodes"),
path="$.evaluation.episodes",
minimum=1,
maximum=MAX_EVAL_EPISODES,
errors=errors,
)
_integer(
evaluation.get("seed"),
path="$.evaluation.seed",
minimum=0,
maximum=2**32 - 1,
errors=errors,
)
if evaluation.get("separate") is not True:
errors.append("$.evaluation.separate must be true")
if evaluation.get("seed") == training.get("seed"):
errors.append("training and evaluation seeds must differ")
def _validate_logging(
logging: dict[str, Any], profile: str, errors: list[str]
) -> None:
allowed = {
"backend",
"disclosure_ack",
"external_opt_in",
"upload_checkpoints",
}
errors.extend(
require_keys(logging, allowed=allowed, required=allowed, path="$.logging")
)
backend = logging.get("backend")
if backend not in LOGGER_CREDENTIAL_ENV:
errors.append("$.logging.backend must be none, wandb, or neptune")
return
if profile == "source-4.0" and backend == "neptune":
errors.append("Neptune is not a current source-4.0 integration")
opt_in = _boolean(
logging.get("external_opt_in"),
path="$.logging.external_opt_in",
errors=errors,
)
disclosure = _boolean(
logging.get("disclosure_ack"),
path="$.logging.disclosure_ack",
errors=errors,
)
upload = _boolean(
logging.get("upload_checkpoints"),
path="$.logging.upload_checkpoints",
errors=errors,
)
if backend == "none" and any(value is True for value in (opt_in, disclosure, upload)):
errors.append("external logging flags must be false when backend is none")
if backend != "none" and (opt_in is not True or disclosure is not True):
errors.append(
"external logging requires external_opt_in=true and disclosure_ack=true"
)
def _validate_checkpoint(checkpoint: dict[str, Any], errors: list[str]) -> None:
allowed = {"format", "trusted_only"}
errors.extend(
require_keys(
checkpoint,
allowed=allowed,
required=allowed,
path="$.checkpoint",
)
)
if checkpoint.get("format") not in {"state_dict", "native-bin", "opaque"}:
errors.append("$.checkpoint.format is invalid")
if checkpoint.get("trusted_only") is not True:
errors.append("$.checkpoint.trusted_only must be true")
def validate_plan(plan: Any) -> list[str]:
"""Return all deterministic schema and safety errors."""
errors: list[str] = []
if not isinstance(plan, dict):
return ["$ must be an object"]
secret_paths = secret_key_paths(plan)
if secret_paths:
errors.append(
"credential-bearing keys are forbidden in plans: " + ", ".join(secret_paths)
)
top_allowed = {
"checkpoint",
"environment",
"evaluation",
"logging",
"package",
"profile",
"schema_version",
"training",
"vectorization",
}
errors.extend(
require_keys(
plan,
allowed=top_allowed,
required=top_allowed,
path="$",
)
)
if plan.get("schema_version") != 1:
errors.append("$.schema_version must equal 1")
profile = plan.get("profile")
if profile not in PROFILES:
errors.append(f"$.profile must be one of {PROFILES}")
return errors
package = _mapping(plan.get("package"), "$.package", errors)
environment = _mapping(plan.get("environment"), "$.environment", errors)
training = _mapping(plan.get("training"), "$.training", errors)
vectorization = _mapping(plan.get("vectorization"), "$.vectorization", errors)
evaluation = _mapping(plan.get("evaluation"), "$.evaluation", errors)
logging = _mapping(plan.get("logging"), "$.logging", errors)
checkpoint = _mapping(plan.get("checkpoint"), "$.checkpoint", errors)
_validate_package(package, profile, errors)
_validate_environment(environment, errors)
_validate_training(training, errors)
_validate_vectorization(vectorization, profile, training, errors)
_validate_evaluation(evaluation, training, errors)
_validate_logging(logging, profile, errors)
_validate_checkpoint(checkpoint, errors)
return errors
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description=(
"Validate a strict local JSON plan. With no --config, validate the "
"bounded built-in dry-run plan."
)
)
parser.add_argument("--config", help="Explicit JSON file beneath --root")
parser.add_argument("--root", default=".", help="Allowed local path root")
parser.add_argument("--profile", choices=PROFILES, default="pypi-3.0.0")
parser.add_argument("--compact", action="store_true")
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
try:
plan = (
load_json_object(args.config, root=args.root)
if args.config
else default_plan(args.profile)
)
errors = validate_plan(plan)
report = {
"errors": errors,
"network_used": False,
"plan": copy.deepcopy(plan),
"status": "valid" if not errors else "invalid",
}
except UserInputError as exc:
report = {
"errors": [str(exc)],
"network_used": False,
"plan": None,
"status": "invalid",
}
emit_json(report, pretty=not args.compact)
return 0 if report["status"] == "valid" else 1
if __name__ == "__main__":
raise SystemExit(main())

1
skills/pufferlib/tests/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Tests for the bundled PufferLib safety tools."""

244
skills/pufferlib/tests/test_scripts.py Normal file
View File

@@ -0,0 +1,244 @@
#!/usr/bin/env python3
"""Dependency-free synthetic tests for all bundled PufferLib CLIs."""
from __future__ import annotations
import hashlib
import json
import os
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
SKILL_DIR = Path(__file__).resolve().parents[1]
SCRIPTS_DIR = SKILL_DIR / "scripts"
sys.path.insert(0, str(SCRIPTS_DIR))
import env_template # noqa: E402
import validate_plan # noqa: E402
def run_script(name: str, *arguments: str) -> subprocess.CompletedProcess[str]:
environment = os.environ.copy()
environment["PYTHONDONTWRITEBYTECODE"] = "1"
return subprocess.run(
[sys.executable, str(SCRIPTS_DIR / name), *arguments],
check=False,
capture_output=True,
env=environment,
text=True,
timeout=30,
)
class SyntheticEnvironmentTests(unittest.TestCase):
def test_seeded_trace_is_deterministic(self) -> None:
first = env_template.SyntheticGymEnv(max_steps=8)
second = env_template.SyntheticGymEnv(max_steps=8)
self.assertEqual(first.reset(seed=7), second.reset(seed=7))
for action in (0, 2, 1, 0):
self.assertEqual(first.step(action), second.step(action))
def test_contract_validator(self) -> None:
result = run_script(
"env_contract_validator.py",
"--steps",
"24",
"--episodes",
"4",
"--compact",
)
self.assertEqual(result.returncode, 0, result.stderr)
report = json.loads(result.stdout)
self.assertEqual(report["status"], "passed")
self.assertFalse(report["network_used"])
class CliTests(unittest.TestCase):
def test_all_help_commands_are_dependency_free(self) -> None:
names = (
"env_template.py",
"env_contract_validator.py",
"benchmark_vectorization.py",
"train_template.py",
"validate_plan.py",
"inspect_checkpoint.py",
"repro_plan.py",
)
for name in names:
with self.subTest(name=name):
result = run_script(name, "--help")
self.assertEqual(result.returncode, 0, result.stderr)
self.assertIn("usage:", result.stdout.lower())
def test_serial_benchmark_is_bounded_json(self) -> None:
result = run_script(
"benchmark_vectorization.py",
"--envs",
"2",
"--steps-per-env",
"8",
"--repeats",
"1",
"--warmup-steps",
"0",
"--compact",
)
self.assertEqual(result.returncode, 0, result.stderr)
report = json.loads(result.stdout)
self.assertEqual(report["total_agent_steps_per_repeat"], [16])
self.assertFalse(report["network_used"])
def test_train_template_defaults_to_dry_run(self) -> None:
result = run_script("train_template.py", "--compact")
self.assertEqual(result.returncode, 0, result.stderr)
report = json.loads(result.stdout)
self.assertTrue(report["dry_run"])
self.assertEqual(report["command_preview"], [])
self.assertEqual(report["plan"]["logging"]["backend"], "none")
def test_source_profile_uses_torch_dry_run(self) -> None:
result = run_script(
"train_template.py", "--profile", "source-4.0", "--compact"
)
self.assertEqual(result.returncode, 0, result.stderr)
report = json.loads(result.stdout)
self.assertEqual(report["plan"]["vectorization"]["backend"], "torch")
self.assertEqual(report["plan"]["checkpoint"]["format"], "state_dict")
def test_external_logger_requires_two_opt_ins(self) -> None:
result = run_script("train_template.py", "--logger", "wandb", "--compact")
self.assertEqual(result.returncode, 1)
report = json.loads(result.stdout)
self.assertEqual(report["status"], "invalid")
self.assertTrue(report["errors"])
def test_external_logger_never_reads_credential_value(self) -> None:
result = run_script(
"train_template.py",
"--logger",
"wandb",
"--enable-external-logging",
"--acknowledge-external-disclosure",
"--compact",
)
self.assertEqual(result.returncode, 0, result.stderr)
report = json.loads(result.stdout)
self.assertEqual(report["credential"]["environment_variable"], "WANDB_API_KEY")
self.assertFalse(report["credential"]["value_read_or_logged"])
def test_source_profile_rejects_neptune(self) -> None:
result = run_script(
"train_template.py",
"--profile",
"source-4.0",
"--logger",
"neptune",
"--enable-external-logging",
"--acknowledge-external-disclosure",
"--compact",
)
self.assertEqual(result.returncode, 1)
self.assertIn("not a current source-4.0 integration", result.stdout)
def test_repro_plan_separates_seeds(self) -> None:
result = run_script(
"repro_plan.py", "--replicates", "2", "--eval-episodes", "3", "--compact"
)
self.assertEqual(result.returncode, 0, result.stderr)
report = json.loads(result.stdout)
self.assertTrue(
set(report["training"]["seeds"]).isdisjoint(
report["evaluation"]["seeds"]
)
)
class PlanValidationTests(unittest.TestCase):
def test_default_plans_validate(self) -> None:
for profile in validate_plan.PROFILES:
with self.subTest(profile=profile):
self.assertEqual(validate_plan.validate_plan(validate_plan.default_plan(profile)), [])
def test_secret_key_is_rejected(self) -> None:
plan = validate_plan.default_plan()
plan["logging"]["api_token"] = None
errors = validate_plan.validate_plan(plan)
self.assertTrue(any("credential-bearing" in error for error in errors))
def test_external_environment_requires_attested_provenance(self) -> None:
result = run_script(
"train_template.py",
"--environment",
"reviewed-env",
"--adapter",
"gymnasium",
"--compact",
)
self.assertEqual(result.returncode, 1)
report = json.loads(result.stdout)
self.assertTrue(any("provenance_verified" in error for error in report["errors"]))
def test_duplicate_json_key_is_rejected(self) -> None:
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "plan.json").write_text(
'{"schema_version":1,"schema_version":1}', encoding="utf-8"
)
result = run_script(
"validate_plan.py",
"--root",
str(root),
"--config",
"plan.json",
"--compact",
)
self.assertEqual(result.returncode, 1)
self.assertIn("duplicate JSON key", result.stdout)
class CheckpointInspectorTests(unittest.TestCase):
def test_inspector_hashes_without_deserialization(self) -> None:
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
payload = b"\x80\x04synthetic-not-executable"
checkpoint = root / "sample.pt"
checkpoint.write_bytes(payload)
digest = hashlib.sha256(payload).hexdigest()
sidecar = root / "sample.metadata.json"
sidecar.write_text(
json.dumps(
{
"schema_version": 1,
"format": "state_dict",
"sha256": digest,
"seed": 42,
}
),
encoding="utf-8",
)
result = run_script(
"inspect_checkpoint.py",
"sample.pt",
"--root",
str(root),
"--metadata",
"sample.metadata.json",
"--expected-sha256",
digest,
"--compact",
)
self.assertEqual(result.returncode, 0, result.stderr)
report = json.loads(result.stdout)
self.assertFalse(report["deserialized"])
self.assertTrue(report["expected_sha256_matches"])
self.assertTrue(report["metadata_sha256_matches"])
self.assertEqual(
report["checkpoint"]["format_detection"]["family"], "pickle-like"
)
if __name__ == "__main__":
unittest.main()