{
  "skill": "alterlab-pufferlib",
  "evals": [
    {
      "id": "vectorized-ppo-training-throughput",
      "prompt": "I'm training a PPO agent on Procgen coinrun and standard training is way too slow. I want to push hundreds of parallel environments and hit millions of steps per second on my GPU. How do I set this up?",
      "expected_output": "Invokes alterlab-pufferlib: builds a vectorized environment with pufferlib.vector.make(env_creator, num_envs=256, num_workers=...) passing an env-constructor callable, trains with the PuffeRL (PPO+LSTM) trainer on device='cuda' using the evaluate()/train()/mean_and_log() loop, and tunes num_envs/num_workers/batch_size to maximize steps-per-second throughput.",
      "assertions": [
        { "type": "should_trigger", "value": true },
        { "type": "output_contains", "value": "PuffeRL" },
        { "type": "behavior", "value": "Uses pufferlib.vector.make with many parallel envs and the PuffeRL trainer to maximize steps-per-second, tuning num_envs/num_workers." }
      ]
    },
    {
      "id": "custom-pufferenv-implementation",
      "prompt": "I need to build my own high-performance grid-world environment for RL. How do I implement it so it works with PufferLib's fast training and vectorization?",
      "expected_output": "Invokes alterlab-pufferlib: implements a custom environment by subclassing PufferEnv, defining single_observation_space and single_action_space plus num_agents BEFORE calling super().__init__(buf), and implementing reset(seed) returning (obs, info-list) and step(action) returning (obs, rewards, terminals, truncations, info) with in-place buffer operations; suggests starting from scripts/env_template.py.",
      "assertions": [
        { "type": "should_trigger", "value": true },
        { "type": "output_contains", "value": "PufferEnv" },
        { "type": "behavior", "value": "Subclasses PufferEnv, defines observation/action spaces before super().__init__(buf), and implements reset/step with the PufferEnv signature." }
      ]
    },
    {
      "id": "multi-agent-pettingzoo-integration",
      "prompt": "I have a cooperative multi-agent PettingZoo environment (knights_archers_zombies) and I want to train all agents with a shared policy at scale. How do I wrap and vectorize it in PufferLib?",
      "expected_output": "Invokes alterlab-pufferlib: wraps the PettingZoo parallel_env in pufferlib.emulation.PettingZooPufferEnv via an env_creator callable, vectorizes with pufferlib.vector.make(env_creator, num_envs=...), and trains a shared policy built from env.single_observation_space / env.single_action_space using the PuffeRL trainer for native multi-agent RL.",
      "assertions": [
        { "type": "should_trigger", "value": true },
        { "type": "output_contains", "value": "PettingZooPufferEnv" },
        { "type": "behavior", "value": "Wraps the PettingZoo env with PettingZooPufferEnv, vectorizes it, and trains a shared multi-agent policy." }
      ]
    },
    {
      "id": "lstm-cnn-policy-development",
      "prompt": "My environment gives image observations and the task is partially observable, so I think I need a recurrent CNN policy. How do I write the policy module for PufferLib training?",
      "expected_output": "Invokes alterlab-pufferlib: builds the policy as a standard torch.nn.Module with a CNN encoder for image observations feeding an optimized LSTM for the partially-observable/recurrent component, with separate actor and critic heads, using layer_init from pufferlib.pytorch for proper weight initialization (actor std=0.01, critic std=1.0); points to references/policies.md.",
      "assertions": [
        { "type": "should_trigger", "value": true },
        { "type": "output_contains", "value": "layer_init" },
        { "type": "behavior", "value": "Builds a CNN+LSTM torch.nn.Module policy with actor/critic heads using layer_init from pufferlib.pytorch." }
      ]
    },
    {
      "id": "near-miss-alterlab-stable-baselines3",
      "prompt": "I just want to quickly prototype a standard PPO agent on CartPole-v1 using a well-documented, off-the-shelf implementation with sensible defaults. I don't care about squeezing out maximum throughput, I just want something that works in a few lines.",
      "expected_output": "Does NOT invoke this skill; defers to alterlab-stable-baselines3. The user explicitly wants quick prototyping with a standard, well-documented off-the-shelf algorithm rather than high-throughput vectorized scaling, which is stable-baselines3's territory rather than PufferLib's.",
      "assertions": [
        { "type": "should_not_trigger", "value": true },
        { "type": "output_contains", "value": "alterlab-stable-baselines3" }
      ]
    },
    {
      "id": "near-miss-alterlab-pytorch-lightning",
      "prompt": "I'm training a supervised image classifier with a CNN and I want clean multi-GPU training, checkpointing, and logging boilerplate handled for me. How should I structure the training loop?",
      "expected_output": "Does NOT invoke this skill; defers to alterlab-pytorch-lightning. The user is doing supervised deep learning (image classification), not reinforcement learning with environments and policies, so the general training-loop organization belongs to pytorch-lightning rather than PufferLib's RL-specific tooling.",
      "assertions": [
        { "type": "should_not_trigger", "value": true },
        { "type": "output_contains", "value": "alterlab-pytorch-lightning" }
      ]
    }
  ]
}

# PufferLib Environments Guide

## Overview

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.

## 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
- **Native spaces must be flat**: `single_observation_space` must be a `Box`, and `single_action_space` a `Discrete`, `MultiDiscrete`, or `Box`. A `Dict` (or other structured) observation space is **not** allowed for a native PufferEnv — the base class raises `APIUsageError`. If you need Dict/structured observations, build a Gymnasium env and use `pufferlib.emulation.GymnasiumPufferEnv`, which flattens them for you.

### Creating a PufferEnv

```python
import numpy as np
import gymnasium
import pufferlib
from pufferlib import PufferEnv

class MyEnvironment(PufferEnv):
    def __init__(self, buf=None):
        # Define spaces and num_agents BEFORE super().__init__(buf).
        # Native obs space MUST be a Box (no Dict). Here: an 84x84x3 image.
        self.single_observation_space = gymnasium.spaces.Box(
            low=0, high=255, shape=(84, 84, 3), dtype=np.uint8)
        self.single_action_space = gymnasium.spaces.Discrete(4)  # 4 discrete actions
        self.num_agents = 1

        super().__init__(buf)

    def reset(self, seed=None):
        """Reset environment to initial state. Returns (obs, info-list)."""
        # Reset internal state
        self.agent_pos = np.array([0, 0])
        self.step_count = 0

        # Return initial observation and an (empty) info LIST (asserted by PufferEnv)
        obs = np.zeros((84, 84, 3), dtype=np.uint8)
        return obs, []

    def step(self, action):
        """Execute one environment step. Returns 5 values."""
        # Update state based on action
        self.step_count += 1

        # Calculate reward
        rewards = self._compute_reward()

        # Check termination vs. truncation
        terminals = self._is_done()
        truncations = self.step_count >= 1000

        # Generate observation
        obs = self._get_observation()

        # Additional info (list of dicts; PufferEnv asserts info is a list)
        info = [{'episode': {'r': rewards, 'l': self.step_count}}] if truncations else []

        return obs, rewards, terminals, truncations, info

    def _compute_reward(self):
        """Compute reward for current state."""
        return 1.0

    def _get_observation(self):
        """Generate observation from current state (a Box-shaped array)."""
        return np.random.randint(0, 256, (84, 84, 3), dtype=np.uint8)
```

### Observation Spaces

Define spaces directly with `gymnasium.spaces` and assign them to
`self.single_observation_space` / `self.single_action_space` before calling
`super().__init__(buf)`.

#### Native observation spaces (Box only)

A native PufferEnv observation space must be a `Box`. Encode discrete or
heterogeneous state into a single Box vector/array:

```python
import gymnasium

# Flat continuous/feature vector
self.single_observation_space = gymnasium.spaces.Box(
    low=-np.inf, high=np.inf, shape=(10,), dtype=np.float32)

# Image observation
self.single_observation_space = gymnasium.spaces.Box(
    low=0, high=255, shape=(84, 84, 3), dtype=np.uint8)
```

#### Dict / structured observations (emulation path only)

`Dict` observation spaces are **not** valid for a native PufferEnv. To use them,
build a Gymnasium env and wrap it with `pufferlib.emulation.GymnasiumPufferEnv`,
which flattens the Dict to a Box for you (recover the structure in the policy via
`pufferlib.pytorch.nativize_tensor`):

```python
import gymnasium
# Inside a standard gymnasium.Env (NOT a PufferEnv):
observation_space = gymnasium.spaces.Dict({
    'image': gymnasium.spaces.Box(low=0, high=255, shape=(84, 84, 3), dtype=np.uint8),
    'vector': gymnasium.spaces.Box(low=-np.inf, high=np.inf, shape=(10,), dtype=np.float32),
})
```

### Action Spaces

```python
import gymnasium

# Discrete actions
self.single_action_space = gymnasium.spaces.Discrete(4)  # 4 actions: 0, 1, 2, 3

# Continuous actions
self.single_action_space = gymnasium.spaces.Box(
    low=-1.0, high=1.0, shape=(3,), dtype=np.float32)  # 3D continuous action

# Multi-discrete actions
self.single_action_space = gymnasium.spaces.MultiDiscrete([3, 3])  # Two 3-way discrete choices
```

## Multi-Agent Environments

PufferLib has native multi-agent support, treating single-agent and multi-agent environments uniformly: set `num_agents > 1` and use the **same** array-based API as a single-agent env. There is no `{agent_id: ...}` dict / `'__all__'` convention in native PufferEnv — that belongs to PettingZoo (use `pufferlib.emulation.PettingZooPufferEnv` for those). The native obs/reward/terminal/truncation buffers are simply sized along a leading agent axis, and `info` is a list of dicts.

### Multi-Agent PufferEnv

```python
import gymnasium
import numpy as np

class MultiAgentEnv(PufferEnv):
    def __init__(self, num_agents=4, buf=None):
        # Define spaces and num_agents BEFORE super().__init__(buf).
        # single_*_space is PER AGENT and must be a Box (obs) / Discrete-or-Box (action).
        self.num_agents = num_agents
        self.single_observation_space = gymnasium.spaces.Box(
            low=-np.inf, high=np.inf, shape=(14,), dtype=np.float32)  # e.g. pos+vel+global
        self.single_action_space = gymnasium.spaces.Discrete(5)

        super().__init__(buf)

    def reset(self, seed=None):
        """Reset all agents. Returns (obs, info-list)."""
        # obs is shaped (num_agents, *single_observation_space.shape)
        obs = np.zeros((self.num_agents, 14), dtype=np.float32)
        return obs, []

    def step(self, actions):
        """Step all agents. `actions` is an array of shape (num_agents, ...).
        Returns (obs, rewards, terminals, truncations, info-list) where the first
        four are arrays of length num_agents and info is a list of dicts."""
        obs = self._get_observations()                 # (num_agents, 14)
        rewards = self._compute_rewards()               # (num_agents,)
        terminals = self._terminated()                  # (num_agents,) bool
        truncations = self._truncated()                 # (num_agents,) bool
        info = []
        return obs, rewards, terminals, truncations, info
```

## Environments: Ocean vs. third-party bindings

PufferLib ships two distinct collections — keep them straight:

- **Ocean** (`pufferlib.ocean`) is PufferLib's own suite of fast, mostly native-C
  environments. It includes `breakout`, `pong`, `snake`, `enduro`, `freeway`,
  `connect4`, `go`, `g2048`, `tetris`, `pacman`, `nmmo3`, `moba`, `drive`,
  `rware`, `trash_pickup`, `cartpole`, `grid`, and many more (40+). These are the
  high-throughput envs the SPS benchmarks refer to.
- **Third-party bindings** (`pufferlib.environments.*`) wrap external suites:
  `atari`, `procgen`, `nethack`, `minihack`, `minigrid`, `crafter`, `craftax`,
  `butterfly` (PettingZoo), `magent`, `nmmo`, `gpudrive`, `microrts`, `griddly`,
  `pokemon_red`, `mujoco`, `dm_control`, `vizdoom`, and others. So Atari, Procgen,
  and NetHack are **bindings, not Ocean envs**.

### Using these environments

```python
import functools
import pufferlib.vector
import pufferlib.ocean

# Ocean (native): resolve a constructor with env_creator. The name is prefixed
# 'puffer_' (matching the CLI `puffer train puffer_breakout`).
breakout = pufferlib.ocean.environment.env_creator('puffer_breakout')
env = pufferlib.vector.make(breakout, num_envs=256)

# Bind constructor kwargs with functools.partial (PufferLib has no string registry).
snake = pufferlib.ocean.environment.env_creator('puffer_snake')
env = pufferlib.vector.make(functools.partial(snake, num_agents=4), num_envs=128)

# Third-party bindings live under pufferlib.environments.<name> and expose an
# `env_creator`. They are not native, so pass an explicit backend.
import pufferlib.environments.atari as atari
env = pufferlib.vector.make(
    atari.env_creator('BreakoutNoFrameskip-v4'),
    num_envs=128, backend=pufferlib.vector.Multiprocessing)
```

## 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

```c
// my_env.c
#include <Python.h>
#include <numpy/arrayobject.h>

// 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);
}
```

## Third-Party Environment Integration

### Gymnasium Environments

```python
import gymnasium as gym
import pufferlib.emulation
import pufferlib.vector

# Wrap a Gymnasium env in a GymnasiumPufferEnv, then vectorize. Wrapped
# (non-native) envs require an explicit backend (default backend=PufferEnv is
# native-only).
def env_creator():
    return pufferlib.emulation.GymnasiumPufferEnv(
        env_creator=lambda: gym.make('CartPole-v1'))

env = pufferlib.vector.make(
    env_creator, num_envs=256, backend=pufferlib.vector.Multiprocessing)
```

### PettingZoo Environments

```python
from pettingzoo.butterfly import pistonball_v6
import pufferlib.emulation
import pufferlib.vector

# Wrap a PettingZoo env in a PettingZooPufferEnv, then vectorize.
def env_creator():
    return pufferlib.emulation.PettingZooPufferEnv(
        env_creator=lambda: pistonball_v6.parallel_env())

env = pufferlib.vector.make(
    env_creator, num_envs=128, backend=pufferlib.vector.Multiprocessing)
```

### Custom Wrappers

A PufferEnv defines `single_observation_space` / `single_action_space` — assigning
`observation_space` / `action_space` is rejected by the base class. Mirror the
wrapped env's *single* spaces, set `num_agents` before `super().__init__`, and keep
the native 5-tuple step signature.

```python
class CustomWrapper(pufferlib.PufferEnv):
    """Wrapper to modify a native PufferEnv's behavior."""

    def __init__(self, base_env, buf=None):
        self.base_env = base_env
        # Copy the per-agent spaces (NOT observation_space/action_space).
        self.single_observation_space = base_env.single_observation_space
        self.single_action_space = base_env.single_action_space
        self.num_agents = base_env.num_agents
        super().__init__(buf)

    def reset(self, seed=None):
        obs, info = self.base_env.reset(seed)
        return self._process_obs(obs), info

    def step(self, action):
        action = self._process_action(action)
        obs, rewards, terminals, truncations, info = self.base_env.step(action)
        obs = self._process_obs(obs)
        rewards = self._process_reward(rewards)
        return obs, rewards, terminals, truncations, 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, terminal, truncation, info
```

### Episode Termination

PufferLib keeps termination and truncation separate (the 4th return value).
A time limit is a *truncation*; reaching a goal/failure is a *termination*.

```python
def step(self, action):
    # ... environment logic ...

    terminal = self._check_success() or self._check_failure()  # task-ending
    truncation = self.step_count >= self.max_steps             # time limit
    info = [{'success': self._check_success()}] if (terminal or truncation) else []

    return obs, reward, terminal, truncation, 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. Note: PufferEnv exposes single_*_space (per agent);
# observation_space/action_space are the joint (batched) spaces set by the base class.
def step(self, action):
    obs, reward, terminal, truncation, info = self._step_impl(action)

    assert self.single_observation_space.contains(obs), "Invalid observation"
    assert np.all(np.isfinite(reward)), "Non-finite reward"

    return obs, reward, terminal, truncation, 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, terminal, truncation, info = self._step_impl(action)

    if np.any(terminal) or np.any(truncation):
        logger.info(f"Episode finished: reward={self.total_reward}")

    return obs, reward, terminal, truncation, info
```

# PufferLib Integration Guide

> Targets PufferLib **3.0.x**.
>
> **Critical rule for this whole guide:** any env wrapped via the emulation layer
> (`GymnasiumPufferEnv` / `PettingZooPufferEnv`) is *not* a native PufferEnv, so
> `pufferlib.vector.make` requires an explicit `backend=pufferlib.vector.Multiprocessing`
> (or `Serial`). The default `backend=PufferEnv` is native-only and raises
> `APIUsageError` on a wrapped env. The examples below omit it for brevity only where
> noted — add it in real code.

## Overview

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. Many third-party suites also ship ready-made bindings under `pufferlib.environments.*` (e.g. `atari`, `procgen`, `nethack`, `minigrid`, `crafter`), each exposing an `env_creator`.

## Gymnasium Integration

### Basic Gymnasium Environments

```python
import gymnasium as gym
import pufferlib.emulation
import pufferlib.vector

# PufferLib has no string registry -- pass an environment constructor
# (callable) to `pufferlib.vector.make`. There is no top-level
# `pufferlib.emulate`/`pufferlib.make`; wrap Gymnasium envs with
# `pufferlib.emulation.GymnasiumPufferEnv`, then vectorize.

# Method 1: Wrap a single Gymnasium env, then vectorize.
# Wrapped envs require an explicit backend (default backend=PufferEnv is native-only).
def cartpole_creator():
    return pufferlib.emulation.GymnasiumPufferEnv(
        env_creator=lambda: gym.make('CartPole-v1'))

env = pufferlib.vector.make(
    cartpole_creator, num_envs=256, backend=pufferlib.vector.Multiprocessing)

# Method 2: Same pattern with an inline creator
env = pufferlib.vector.make(
    lambda: pufferlib.emulation.GymnasiumPufferEnv(
        env_creator=lambda: gym.make('CartPole-v1')),
    num_envs=256, backend=pufferlib.vector.Multiprocessing,
)

# 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 (MyGymEnv is a gym.Env, so emulate it)
env = pufferlib.vector.make(
    lambda: pufferlib.emulation.GymnasiumPufferEnv(env_creator=MyGymEnv),
    num_envs=128, backend=pufferlib.vector.Multiprocessing,
)
```

### Atari Environments

The simplest path is PufferLib's bundled Atari binding, which handles the standard
preprocessing for you:

```python
import pufferlib.environments.atari as atari
import pufferlib.vector

# The binding takes an ALE ROM name (e.g. 'breakout', 'pong') and returns an
# already-emulated env, so still pass a non-native backend.
env = pufferlib.vector.make(
    atari.env_creator('breakout'),
    num_envs=256, backend=pufferlib.vector.Multiprocessing)
```

You can also build the Gymnasium env yourself and emulate it:

```python
import functools
import gymnasium as gym
from gymnasium.wrappers import AtariPreprocessing, FrameStack
import pufferlib.emulation
import pufferlib.vector

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

def atari_creator():
    return pufferlib.emulation.GymnasiumPufferEnv(env_creator=make_atari_env)

env = pufferlib.vector.make(
    atari_creator, num_envs=256, backend=pufferlib.vector.Multiprocessing)
```

### Complex Observation Spaces

```python
import numpy as np
import gymnasium as gym
from gymnasium.spaces import Dict, Box, Discrete
import pufferlib.emulation
import pufferlib.vector

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, {}

# GymnasiumPufferEnv flattens Dict spaces to a Box for the trainer; recover the
# structure in the policy via pufferlib.pytorch.nativize_tensor (see policies.md).
env = pufferlib.vector.make(
    lambda: pufferlib.emulation.GymnasiumPufferEnv(env_creator=ComplexObsEnv),
    num_envs=128, backend=pufferlib.vector.Multiprocessing,
)
```

## PettingZoo Integration

### Parallel Environments

```python
from pettingzoo.butterfly import pistonball_v6
import pufferlib.emulation
import pufferlib.vector

# Multi-agent envs are wrapped with PettingZooPufferEnv, then vectorized.
# PufferLib has no string registry -- pass the constructor (callable).
def pistonball_creator():
    return pufferlib.emulation.PettingZooPufferEnv(
        env_creator=lambda: pistonball_v6.parallel_env())

env = pufferlib.vector.make(
    pistonball_creator, num_envs=128, backend=pufferlib.vector.Multiprocessing)
```

### AEC (Agent Environment Cycle) Environments

```python
from pettingzoo.classic import chess_v5
from pettingzoo.utils import aec_to_parallel
import pufferlib.emulation
import pufferlib.vector

# Convert AEC to a parallel env, wrap with PettingZooPufferEnv, then vectorize.
# PufferLib has no string registry -- pass the constructor (callable).
def chess_creator():
    return pufferlib.emulation.PettingZooPufferEnv(
        env_creator=lambda: aec_to_parallel(chess_v5.env()))

env = pufferlib.vector.make(
    chess_creator, num_envs=64, backend=pufferlib.vector.Multiprocessing)
```

### Multi-Agent Training

```python
import pufferlib.emulation
import pufferlib.vector
from pufferlib.pufferl import PuffeRL, load_config

# Create multi-agent environment. `kaz_creator` is a placeholder that wraps a
# PettingZoo env (e.g. knights_archers_zombies) with PettingZooPufferEnv.
env = pufferlib.vector.make(
    kaz_creator, num_envs=128, backend=pufferlib.vector.Multiprocessing)

# One shared policy serves all agents. PufferLib emulates PettingZoo as a single
# flat agent axis, so build the policy from the PER-AGENT spaces.
policy = create_policy(env.single_observation_space, env.single_action_space)

# Train: PuffeRL(config_dict, vecenv, policy). Loop on global_step.
args = load_config('puffer_breakout')          # or your own env config
config = {**args['train'], 'env': 'multiagent'}
trainer = PuffeRL(config, env, policy)

while trainer.global_step < config['total_timesteps']:
    trainer.evaluate()   # rollouts are batched arrays, not {agent_id: ...} dicts
    trainer.train()
    trainer.mean_and_log()
```

## Third-Party Environments

PufferLib ships bindings for these suites under `pufferlib.environments.<name>`,
each exposing an `env_creator(name, ...)` (install the underlying package first).
They are emulated, so pass `backend=pufferlib.vector.Multiprocessing`. The
`make_*` names in the GPUDrive/MicroRTS/Griddly snippets below are placeholders
for whichever creator you obtain (binding `env_creator` or your own callable).

### Procgen

```python
import pufferlib.environments.procgen as procgen
import pufferlib.vector

env = pufferlib.vector.make(
    procgen.env_creator('coinrun'),
    num_envs=256, backend=pufferlib.vector.Multiprocessing)
```

### NetHack / MiniHack

```python
import pufferlib.environments.nethack as nethack
import pufferlib.vector

env = pufferlib.vector.make(
    nethack.env_creator('nethack'),
    num_envs=128, backend=pufferlib.vector.Multiprocessing)
```

### Minigrid

```python
import pufferlib.environments.minigrid as minigrid
import pufferlib.vector

env = pufferlib.vector.make(
    minigrid.env_creator('MiniGrid-Empty-8x8-v0'),
    num_envs=256, backend=pufferlib.vector.Multiprocessing)
```

### Neural MMO

```python
import functools
import pufferlib.vector

# Large-scale multi-agent environment.
# Pass a constructor callable -- there is no string registry. `make_neuralmmo`
# is a placeholder for the env creator you provide; bind kwargs with partial.
env = pufferlib.vector.make(
    functools.partial(
        make_neuralmmo,
        num_agents=128,  # Agents per environment
        map_size=128,
    ),
    num_envs=64,
)
```

### Crafter

```python
import pufferlib.vector

# Open-ended crafting environment.
# Pass a constructor callable -- there is no string registry. `make_crafter`
# is a placeholder for the env creator you provide.
env = pufferlib.vector.make(make_crafter, num_envs=128)
```

### GPUDrive

```python
import functools
import pufferlib.vector

# GPU-accelerated driving simulator.
# Pass a constructor callable -- there is no string registry. `make_gpudrive`
# is a placeholder for the env creator you provide; bind kwargs with partial.
env = pufferlib.vector.make(
    functools.partial(make_gpudrive, num_vehicles=8),
    num_envs=1024,  # Can handle many environments on GPU
)
```

### MicroRTS

```python
import functools
import pufferlib.vector

# Real-time strategy game.
# Pass a constructor callable -- there is no string registry. `make_microrts`
# is a placeholder for the env creator you provide; bind kwargs with partial.
env = pufferlib.vector.make(
    functools.partial(make_microrts, map_size=16, max_steps=2000),
    num_envs=128,
)
```

### Griddly

```python
import pufferlib.vector

# Grid-based games.
# Pass a constructor callable -- there is no string registry. make_griddly_*
# are placeholders for the env creators you provide.
env = pufferlib.vector.make(make_griddly_clusters, num_envs=256)
env = pufferlib.vector.make(make_griddly_sokoban, num_envs=256)
```

## Custom Wrappers

There is **no `pufferlib.Wrapper`** base class. For emulated envs, wrap at the
**Gymnasium** level (subclass `gymnasium.Wrapper`) *before* passing the env to
`GymnasiumPufferEnv` — that way you keep the standard Gymnasium 5-tuple step
`(obs, reward, terminated, truncated, info)` and PufferLib handles the rest.
PufferLib also ships ready wrappers (`pufferlib.ResizeObservation`,
`pufferlib.ClipAction`, `pufferlib.EpisodeStats`). For frame stacking / action
repeat, prefer the standard `gymnasium.wrappers` (e.g. `FrameStack`).

```python
import numpy as np
import gymnasium as gym
import pufferlib.emulation
import pufferlib.vector

class RewardShaping(gym.Wrapper):
    """Add a shaped reward at the Gymnasium level."""

    def __init__(self, env, shaping_fn):
        super().__init__(env)
        self.shaping_fn = shaping_fn

    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        reward = reward + self.shaping_fn(obs, action)
        return obs, reward, terminated, truncated, info

def proximity_shaping(obs, action):
    goal_pos = np.array([10, 10])
    distance = np.linalg.norm(goal_pos - obs[:2])
    return -0.1 * distance

# Apply wrappers inside the creator, then emulate + vectorize.
def shaped_creator():
    base = gym.make('CartPole-v1')
    base = RewardShaping(base, proximity_shaping)
    return pufferlib.emulation.GymnasiumPufferEnv(env_creator=lambda: base)

env = pufferlib.vector.make(
    shaped_creator, num_envs=128, backend=pufferlib.vector.Multiprocessing)
```

## 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
```

### Recovering structure in the policy

There is no `unflatten_observations` helper. Use
`pufferlib.pytorch.nativize_dtype` (once, from `env.emulated`) plus
`pufferlib.pytorch.nativize_tensor` per batch to recover the structured dict:

```python
import pufferlib.pytorch

class PolicyWithDictObs(nn.Module):
    def __init__(self, env):
        super().__init__()
        self.dtype = pufferlib.pytorch.nativize_dtype(env.emulated)
        # ... encoders / heads ...

    def encode_observations(self, flat_observations, state=None):
        obs = pufferlib.pytorch.nativize_tensor(flat_observations, self.dtype)
        image_features = self.image_encoder(obs['image'].float() / 255.0)
        vector_features = self.vector_encoder(obs['vector'])
        # ...
```

## Environment Lookup

### Mapping Names to Constructors

PufferLib has no string registry and no `pufferlib.register`/`pufferlib.make`.
If you want name-based lookup, keep your own mapping of names to environment
constructors (callables) and pass the resolved constructor to
`pufferlib.vector.make`:

```python
import functools
import pufferlib.vector
from my_package.envs import MyEnvironment

# Your own registry: name -> constructor (callable)
ENV_REGISTRY = {
    'my-custom-env': functools.partial(MyEnvironment, param1='value1'),
}

# Resolve the name to a constructor, then vectorize
env_creator = ENV_REGISTRY['my-custom-env']
env = pufferlib.vector.make(env_creator, num_envs=256)
```

### Sharing Constructors

To make an environment easy for others to use, simply expose its constructor
(a `PufferEnv` subclass, or a function/`functools.partial` that returns one)
from your package and document the kwargs:

```python
# In my_package/envs.py
def make_my_env(default_param='default_value'):
    return MyEnvironment(param1=default_param)
```

## Compatibility Patterns

### Gymnasium to PufferLib

```python
import gymnasium as gym
import pufferlib.emulation
import pufferlib.vector

# 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

# Wrap with GymnasiumPufferEnv, then vectorize (pass a constructor callable)
env = pufferlib.vector.make(
    lambda: pufferlib.emulation.GymnasiumPufferEnv(env_creator=GymEnv),
    num_envs=128,
)
```

### PettingZoo to PufferLib

```python
from pettingzoo import ParallelEnv
import pufferlib.emulation
import pufferlib.vector

# 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

# Wrap with PettingZooPufferEnv, then vectorize (pass a constructor callable)
env = pufferlib.vector.make(
    lambda: pufferlib.emulation.PettingZooPufferEnv(env_creator=PZEnv),
    num_envs=128,
)
```

### Legacy Gym (v0.21) to PufferLib

```python
import gym  # Old gym
import pufferlib.emulation
import pufferlib.vector

# 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

# Wrap with GymnasiumPufferEnv, then vectorize (pass a constructor callable)
env = pufferlib.vector.make(
    lambda: pufferlib.emulation.GymnasiumPufferEnv(env_creator=LegacyEnv),
    num_envs=128,
)
```

## Performance Considerations

### Efficient Integration

```python
import gymnasium as gym
import pufferlib.emulation
import pufferlib.vector

# Fast: a native PufferEnv constructor (no emulation layer).
# `make_coinrun` is a placeholder for your PufferEnv creator.
env = pufferlib.vector.make(make_coinrun, num_envs=256)

# Slower: Generic Gymnasium wrapper (emulated -> needs a non-native backend)
env = pufferlib.vector.make(
    lambda: pufferlib.emulation.GymnasiumPufferEnv(
        env_creator=lambda: gym.make('CartPole-v1')),
    num_envs=256, backend=pufferlib.vector.Multiprocessing,
)

# Slowest: Nested wrappers add overhead
def nested_creator():
    gym_env = gym.make('CartPole-v1')
    gym_env = SomeWrapper(gym_env)
    gym_env = AnotherWrapper(gym_env)
    return pufferlib.emulation.GymnasiumPufferEnv(env_creator=lambda: gym_env)

env = pufferlib.vector.make(
    nested_creator, num_envs=256, backend=pufferlib.vector.Multiprocessing)
```

### Minimize Wrapper Overhead

```python
import pufferlib.emulation
import pufferlib.vector

# BAD: Too many wrappers
def bad_creator():
    env = gym.make('CartPole-v1')
    env = Wrapper1(env)
    env = Wrapper2(env)
    env = Wrapper3(env)
    return pufferlib.emulation.GymnasiumPufferEnv(env_creator=lambda: env)

env = pufferlib.vector.make(
    bad_creator, num_envs=256, backend=pufferlib.vector.Multiprocessing)

# GOOD: Combine wrapper logic
class CombinedWrapper(gym.Wrapper):
    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        # Apply all transformations at once
        obs = self._transform_obs(obs)
        reward = self._transform_reward(reward)
        return obs, reward, terminated, truncated, info

def good_creator():
    env = CombinedWrapper(gym.make('CartPole-v1'))
    return pufferlib.emulation.GymnasiumPufferEnv(env_creator=lambda: env)

env = pufferlib.vector.make(
    good_creator, num_envs=256, backend=pufferlib.vector.Multiprocessing)
```

## Debugging Integration

### Verify Environment Compatibility

```python
import numpy as np

def test_environment(env, num_steps=100):
    """Smoke-test a single PufferEnv before vectorizing (5-tuple native API)."""
    obs, info = env.reset()
    assert isinstance(info, list), "PufferEnv reset must return an info LIST"

    for _ in range(num_steps):
        # One action per agent.
        action = np.array([env.single_action_space.sample()
                           for _ in range(env.num_agents)])
        obs, reward, terminal, truncation, info = env.step(action)

        assert np.all(np.isfinite(np.asarray(reward, dtype=float))), "Non-finite reward"
        assert isinstance(info, list), "info must be a list of dicts"

        if np.any(terminal) or np.any(truncation):
            obs, info = 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.emulation
import pufferlib.vector
import numpy as np

gym_env = gym.make('CartPole-v1')
puffer_env = pufferlib.vector.make(
    lambda: pufferlib.emulation.GymnasiumPufferEnv(
        env_creator=lambda: gym.make('CartPole-v1')),
    num_envs=1, backend=pufferlib.vector.Multiprocessing,
)

# Both reset() calls return (obs, info). vecenv.step returns a 5-tuple:
# (obs, rewards, terminals, truncations, infos).
gym_env.reset(seed=42)
puffer_obs, _ = puffer_env.reset(seed=42)

for _ in range(100):
    action = gym_env.action_space.sample()

    gym_obs, gym_reward, gym_term, gym_trunc, gym_info = gym_env.step(action)
    puffer_obs, puffer_reward, puffer_term, puffer_trunc, *_ = \
        puffer_env.step(np.array([action]))

    # Compare outputs (accounting for the leading batch/agent dimension)
    assert np.allclose(gym_obs, puffer_obs[0])
    assert gym_reward == puffer_reward[0]
```

# PufferLib Policies Guide

## Overview

PufferLib policies are standard PyTorch modules — the library does not enforce a base class. There are two requirements for use with the PuffeRL trainer (3.0.x):

1. PuffeRL calls `policy.forward_eval(observations, state)` during rollouts and expects it to return `(logits, value)`. Provide `forward_eval` (and usually an identical `forward`).
2. Structure the network as `encode_observations(obs, state)` + `decode_actions(hidden)` so it can be wrapped by `pufferlib.models.LSTMWrapper` without rewrites. `pufferlib.models.Default` is the canonical reference implementation.

The example policies below show the network bodies; in practice wrap each one in this `encode_observations` / `decode_actions` / `forward_eval` interface (as `pufferlib.models.Default` does).

## Policy Architecture

### Basic Policy Structure

```python
import torch
import torch.nn as nn
from pufferlib.pytorch import layer_init

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
```

### Layer Initialization

PufferLib provides `layer_init` for proper weight initialization:

```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))
```

## CNN Policies

For image-based observations:

```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
```

### Efficient CNN Architecture

```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 via `pufferlib.models.LSTMWrapper` (note: it lives in `pufferlib.models`, not `pufferlib.pytorch`). Wrap a policy that exposes `encode_observations` / `decode_actions` and a `hidden_size` attribute, and set `use_rnn=True` in the trainer config:

```python
from pufferlib.models 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 (dict / structured observations)

When an emulated env exposes a Dict observation space, PufferLib flattens it to a
byte tensor. Recover the structured tensors with `pufferlib.pytorch.nativize_dtype`
(once, from `env.emulated`) plus `pufferlib.pytorch.nativize_tensor` per batch —
there is no `unflatten_observations` helper.

```python
import pufferlib.pytorch

class PolicyWithDictObs(nn.Module):
    def __init__(self, env):
        super().__init__()

        # Build the native dtype description once from the emulated env.
        self.dtype = pufferlib.pytorch.nativize_dtype(env.emulated)

        self.encoders = nn.ModuleDict({
            'image': self._make_image_encoder(),
            'vector': self._make_vector_encoder(),
        })
        # ... actor / critic heads ...

    def encode_observations(self, flat_observations, state=None):
        # Recover the structured observation dict from the flat byte tensor.
        obs = pufferlib.pytorch.nativize_tensor(flat_observations, self.dtype)

        image_features = self.encoders['image'](obs['image'].float() / 255.0)
        vector_features = self.encoders['vector'](obs['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
# The trainer clips gradients for you; set the norm via the config dict
# (not a PuffeRL kwarg). PuffeRL(config, vecenv, policy) reads config['max_grad_norm'].
config['max_grad_norm'] = 1.5  # default in pufferlib/config/default.ini
trainer = PuffeRL(config, vecenv, policy)
```

### Model Compilation

```python
# torch.compile is also config-driven: set config['compile'] = True (and
# optionally config['compile_mode']). The trainer compiles the policy internally,
# so you do not call torch.compile yourself.
config['compile'] = True
trainer = PuffeRL(config, vecenv, policy)
```

## Debugging Policies

### Check Output Shapes

```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)

    # Forward pass
    logits, value = policy(obs)

    # Check shapes
    assert logits.shape == (batch_size, policy.action_space.n)
    assert value.shape == (batch_size, 1)

    print("✓ Policy shapes correct")
```

### Verify Gradients

```python
def check_gradients(policy, observation_space):
    """Check that gradients flow properly."""
    obs = torch.randn(1, *observation_space.shape, requires_grad=True)

    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 Training Guide

> Targets PufferLib **3.0.x**. The trainer is config-driven: `PuffeRL.__init__(self, config, vecenv, policy, logger=None)` where `config` is a **dict** (typically `args['train']` from `load_config`), not flat keyword args. The dev `4.0` branch differs.

## Overview

PuffeRL is PufferLib's high-performance training algorithm based on CleanRL's PPO with optional LSTMs, enhanced with research improvements. It achieves high training throughput through optimized vectorization and an efficient implementation.

## Training Workflow

### Basic Training Loop

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()
```

### CLI Training

Quick start via the `puffer` CLI. `env_name` resolves to a config section in `pufferlib/config/*.ini` plus its registered Ocean env. Config keys are overridden with `--<section>.<key>`:

```bash
# Basic training (Ocean Breakout)
puffer train puffer_breakout --train.device cuda

# Override config keys (note real key names: learning-rate, total-timesteps)
puffer train puffer_breakout \
    --train.device cuda \
    --train.learning-rate 0.015 \
    --train.total-timesteps 10_000_000 \
    --vec.num-envs 256
```

### Python Training Script

The supported high-level entry point is `pufferl.train(env_name, ...)`, which builds the vecenv, policy, logger, and config for you:

```python
import pufferlib.pufferl as pufferl

pufferl.train('puffer_breakout')                  # all defaults from config/*.ini
pufferl.train('puffer_breakout', vecenv=my_vec, policy=my_policy)  # override pieces
```

To drive the trainer manually, build the config dict yourself:

```python
import pufferlib.vector
from pufferlib.pufferl import PuffeRL, load_config

# Native PufferEnv -> default backend=PufferEnv. Wrapped (Gymnasium/PettingZoo)
# envs require backend=pufferlib.vector.Multiprocessing (or Serial).
vecenv = pufferlib.vector.make(env_creator, num_envs=256)

# load_config returns a nested args dict; PuffeRL takes the 'train' section (a dict).
args = load_config('puffer_breakout')
config = {**args['train'], 'env': 'puffer_breakout'}
config['device'] = 'cuda'

trainer = PuffeRL(config, vecenv, my_policy)  # (config, vecenv, policy, logger=None)

# Training loop is driven by global_step, not a fixed iteration count.
while trainer.global_step < config['total_timesteps']:
    trainer.evaluate()      # Collect rollouts
    trainer.train()         # Train on batch
    trainer.mean_and_log()  # Aggregate + log
```

## Key Training Parameters

Keys and defaults below are from `pufferlib/config/default.ini` (3.0.x); per-env `.ini` files override them. Use the exact key names — they differ from generic CleanRL/PPO conventions (e.g. `update_epochs`, not `n_epochs`; `minibatch_size`/`bptt_horizon`, not a single `num_steps`).

### Core Hyperparameters (`[train]`)

- **learning_rate**: Optimizer LR (default: `0.015`; note PufferLib's default `optimizer = muon`)
- **batch_size**: Timesteps per training batch (default: `auto`, derived from `bptt_horizon` x agents)
- **minibatch_size**: Minibatch size (default: `8192`)
- **max_minibatch_size**: Gradient accumulation above this size (default: `32768`)
- **bptt_horizon**: BPTT / rollout horizon per segment (default: `64`)
- **update_epochs**: Training epochs per batch (default: `1`)
- **total_timesteps**: Total env steps to train for (default: `10_000_000`)

### Vectorization (`[vec]`)

- **num_envs**: Parallel environments (default: `2`)
- **num_workers**: Vectorization workers (default: `auto`)
- **backend**: `Multiprocessing` (CLI default) / `Serial` / `Ray` / `PufferEnv` (native)

### PPO Parameters (`[train]`)

- **gamma**: Discount factor (default: `0.995`)
- **gae_lambda**: GAE lambda (default: `0.90`)
- **clip_coef**: PPO clipping coefficient (default: `0.2`)
- **ent_coef**: Entropy coefficient (default: `0.001`)
- **vf_coef**: Value loss coefficient (default: `2.0`)
- **vf_clip_coef**: Value clipping coefficient (default: `0.2`)
- **max_grad_norm**: Max gradient norm (default: `1.5`)

### Performance Parameters (`[train]`)

- **device**: Computing device (`cuda` / `cpu`, default `cuda`)
- **compile**: Use torch.compile (default: `False`; `compile_mode = max-autotune-no-cudagraphs`)
- **cpu_offload**: Keep observations on CPU to save GPU memory (default: `False`)

## Distributed Training

### Multi-GPU Training

Use torchrun for distributed training across multiple GPUs:

```bash
torchrun --nproc_per_node=4 train.py \
    --train.device cuda \
    --train.batch-size 131072
```

### Multi-Node Training

For distributed training across multiple nodes:

```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
```

## Monitoring and Logging

### Logger Integration

Loggers live in `pufferlib.pufferl` (`WandbLogger`, `NeptuneLogger`, `NoLogger`), **not** the top-level package. Each takes the nested `args` config dict and reads keys from it (it does not take `project=`/`name=` kwargs).

The intended path is to let `pufferl.train` pick the logger from config flags:

```bash
# Weights & Biases: set --wandb plus the wandb_project / wandb_group keys
puffer train puffer_breakout --train.wandb True --train.wandb-project my_project

# Neptune: set --neptune plus neptune_name / neptune_project
puffer train puffer_breakout --train.neptune True --train.neptune-project my_project
```

To construct one manually, pass the `args` dict (with the relevant keys populated):

```python
from pufferlib.pufferl import WandbLogger, NeptuneLogger, NoLogger

# args is the nested dict from load_config(); WandbLogger reads
# args['wandb_project'], args['wandb_group'], args['tag'], args['no_model_upload'].
logger = WandbLogger(args)          # or NeptuneLogger(args), or NoLogger(args)
trainer = PuffeRL(config, vecenv, policy, logger=logger)
```

### 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

Checkpointing is **auto-managed**: PuffeRL calls `save_checkpoint()` every `config['checkpoint_interval']` epochs and on completion. The method takes **no path argument** — it writes `model_*.pt` and `trainer_state.pt` under `config['data_dir']/<env>_<run_id>/`.

```python
# Manual save (no path arg). Returns the model path written.
model_path = trainer.save_checkpoint()
```

### Resuming / loading

There is no `trainer.load_checkpoint(path)`. Loading happens at policy-build time via config keys, then training resumes from the loaded weights:

```bash
# Resume from the most recent local checkpoint (top-level flag, not in the [train] section)
puffer train puffer_breakout --load-model-path latest

# Or load a specific file, or a logged W&B/Neptune run with --load-id <run_id> --wandb
puffer train puffer_breakout --load-model-path experiments/puffer_breakout_000400.pt
```

## Hyperparameter Tuning with Protein

Protein is PufferLib's Pareto-genetic / GP sweep method (`pufferlib.sweep.Protein`). You do **not** instantiate it directly with a `search_space=` kwarg — sweeps are config-driven via the `[sweep]` section and run with `pufferl.sweep()`. Sweeps **require** a wandb or neptune logger.

In the config (`.ini`), `[sweep]` selects the method and metric, and `[sweep.<section>.<key>]` blocks declare the search space:

```ini
[sweep]
method = Protein
metric = score
goal = maximize

[sweep.train.learning_rate]
distribution = log_normal
min = 1e-4
max = 1e-2

[sweep.vec.num_envs]
distribution = uniform_pow2
min = 1
max = 16
```

Run it:

```bash
puffer sweep puffer_breakout --train.wandb True --train.wandb-project my_sweep
```

```python
import pufferlib.pufferl as pufferl
pufferl.sweep(env_name='puffer_breakout')   # args must enable wandb or neptune
```

## 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
import functools
import pufferlib.vector

# Start with easy tasks, gradually increase difficulty.
# Bind constructor kwargs with functools.partial -- there is no string
# registry, so `MyEnv` here is your own PufferEnv class / creator callable.
difficulty_levels = [0.1, 0.3, 0.5, 0.7, 1.0]

for difficulty in difficulty_levels:
    env = pufferlib.vector.make(
        functools.partial(MyEnv, difficulty=difficulty), num_envs=256)
    trainer = PuffeRL(config, env, policy)  # (config dict, vecenv, policy)

    while trainer.global_step < steps_per_level:
        trainer.evaluate()
        trainer.train()
        trainer.mean_and_log()
```

### Reward Shaping

```python
# A PufferEnv step returns a 5-tuple (obs, rewards, terminals, truncations, info).
class RewardShapedEnv(pufferlib.PufferEnv):
    def step(self, actions):
        obs, rewards, terminals, truncations, info = super().step(actions)

        # Add shaped rewards
        rewards = rewards + 0.1 * proximity_bonus

        return obs, rewards, terminals, truncations, info
```

### Multi-Stage Training

`config` is the dict passed at construction; PuffeRL reads most values once at
init, so change a stage's learning rate by rebuilding the trainer with an
updated config (or use the built-in `anneal_lr` / `min_lr_ratio` config keys).

```python
stages = [
    {'learning_rate': 1e-3, 'steps': 1_000_000},   # Exploration
    {'learning_rate': 3e-4, 'steps': 5_000_000},   # Main training
    {'learning_rate': 1e-4, 'steps': 2_000_000},   # Fine-tuning
]

for stage in stages:
    config = {**config, 'learning_rate': stage['learning_rate']}
    trainer = PuffeRL(config, vecenv, policy)
    target = trainer.global_step + stage['steps']
    while trainer.global_step < target:
        trainer.evaluate()
        trainer.train()
        trainer.mean_and_log()
```

## Troubleshooting

### Low Performance

- 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 Instability

- Reduce learning_rate
- Decrease batch_size
- Increase num_envs for more diverse samples
- Add entropy coefficient for more exploration
- Check reward scaling

### Memory Issues

- Reduce batch_size or num_envs
- Use gradient accumulation
- Disable compile mode if causing OOM
- Check for memory leaks in custom environments

# PufferLib Vectorization Guide

## Overview

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.

## Vectorization Architecture

### Key Optimizations

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

### 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

```python
import pufferlib.vector

# PufferLib has no string registry -- pass an environment constructor
# (callable) to `pufferlib.vector.make`. `env_creator` is a placeholder for
# the callable that builds your env (a PufferEnv class, or a function /
# functools.partial that returns one).

# Native PufferEnv (default backend=PufferEnv)
env = pufferlib.vector.make(env_creator, num_envs=256)

# With explicit configuration. envs-per-worker is derived as num_envs // num_workers.
env = pufferlib.vector.make(
    env_creator,
    num_envs=256,
    num_workers=8,
    backend=pufferlib.vector.Multiprocessing,
)
```

### Choosing a backend

The backends live in `pufferlib.vector` (`Serial`, `Multiprocessing`, `Ray`) — there is **no `pufferlib.vectorization` module**. In normal use you do not instantiate them directly; you pass the class as the `backend=` argument to `pufferlib.vector.make`:

```python
import pufferlib.vector

# Serial — single process, easy debugging
env = pufferlib.vector.make(
    MyEnvironment, num_envs=16, backend=pufferlib.vector.Serial)

# Multiprocessing — parallel workers (required for wrapped Gymnasium/PettingZoo envs)
env = pufferlib.vector.make(
    MyEnvironment, num_envs=256, num_workers=8,
    backend=pufferlib.vector.Multiprocessing)

# Native (default) — a native PufferEnv that handles per-process batching itself
env = pufferlib.vector.make(MyNativePufferEnv, num_envs=256)  # backend=PufferEnv
```

## Vectorization Modes

`pufferlib.vector.make` accepts only these extra kwargs: `num_workers`, `batch_size`, `zero_copy`, `overwork`, `backend` (plus `num_envs`, `seed`, `env_args`, `env_kwargs`). Passing anything else (e.g. `envs_per_worker`, `mode`, `surplus_envs`) raises `APIUsageError`.

### Serial

Best for debugging and lightweight environments — all envs run in the main process:

```python
import pufferlib.vector

env = pufferlib.vector.make(
    env_creator, num_envs=16, backend=pufferlib.vector.Serial)
```

**When to use:** development/debugging, very fast envs, small env counts, single-threaded profiling.

### Multiprocessing

Best for most production use cases and required for wrapped Gymnasium/PettingZoo envs:

```python
import pufferlib.vector

env = pufferlib.vector.make(
    env_creator, num_envs=256, num_workers=8,
    backend=pufferlib.vector.Multiprocessing)
```

**When to use:** production training, CPU-intensive envs, large-scale parallel simulation.

### Asynchronous batching (surplus envs)

There is no `mode='async'` flag. Async-style behavior comes from setting
`batch_size` smaller than `num_envs`: PufferLib simulates the surplus
environments and returns the first `batch_size` agents that are ready, which
improves GPU utilization with variable step times.

```python
env = pufferlib.vector.make(
    env_creator, num_envs=256, batch_size=128, num_workers=8,
    backend=pufferlib.vector.Multiprocessing)  # 128 surplus envs hide stragglers
```

## 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 (derived, not a kwarg)

PufferLib computes envs-per-worker as `num_envs // num_workers` — there is no
`envs_per_worker` argument. Tune the ratio by choosing `num_envs` and
`num_workers`:

```python
# Fast envs -> pack many per worker (high num_envs / num_workers ratio)
env = pufferlib.vector.make(env_creator, num_envs=512, num_workers=8,    # 64 each
    backend=pufferlib.vector.Multiprocessing)

# Slow envs -> fewer per worker
env = pufferlib.vector.make(env_creator, num_envs=128, num_workers=8,    # 16 each
    backend=pufferlib.vector.Multiprocessing)
```

### 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
import gymnasium

class OptimizedEnv(PufferEnv):
    def __init__(self, buf=None):
        # Define spaces and num_agents BEFORE super().__init__(buf).
        # Native obs space must be a Box.
        self.single_observation_space = gymnasium.spaces.Box(
            low=0, high=255, shape=(84, 84, 3), dtype=np.uint8)
        self.single_action_space = gymnasium.spaces.Discrete(4)
        self.num_agents = 1

        # super().__init__ calls set_buffers, which allocates the shared buffers
        # self.observations / self.rewards / self.terminals / self.truncations.
        super().__init__(buf)

    def reset(self, seed=None):
        # Write directly into the shared observation buffer (no allocation).
        self._render_to_buffer(self.observations)
        return self.observations, []

    def step(self, action):
        self._update_state(action)
        self._render_to_buffer(self.observations)  # in-place into shared buffer

        # Write reward/flags into the shared buffers, then return them.
        self.rewards[:] = self._reward()
        self.terminals[:] = self._terminated()
        self.truncations[:] = self._truncated()
        return self.observations, self.rewards, self.terminals, self.truncations, []
```

### 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, terminal, truncation, info

# GOOD: In-place operations
def step(self, action):
    self.state += action  # In-place
    return obs, reward, terminal, truncation, info
```

## Advanced Vectorization

### Wrapping a vecenv

There is no public `VectorEnv` base class to subclass. To add custom behavior,
wrap the object returned by `pufferlib.vector.make` and delegate `reset`/`step`:

```python
import pufferlib.vector

class CustomVecWrapper:
    """Illustrative wrapper around a PufferLib vecenv."""

    def __init__(self, vecenv):
        self.vecenv = vecenv
        self.num_envs = vecenv.num_envs
        self.single_observation_space = vecenv.single_observation_space
        self.single_action_space = vecenv.single_action_space

    def reset(self, seed=0):
        return self.vecenv.reset(seed)

    def step(self, actions):
        return self.vecenv.step(actions)  # delegate, then post-process as needed

vecenv = pufferlib.vector.make(
    env_creator, num_envs=256, backend=pufferlib.vector.Multiprocessing)
vecenv = CustomVecWrapper(vecenv)
```

### Large-scale parallelism

PufferLib already packs `num_envs // num_workers` environments into each worker
process, so a single `make` call covers the "hierarchical" case — you do not
nest backends. Scale by raising `num_envs` and `num_workers`:

```python
# 256 envs across 8 workers => 32 envs per worker, handled internally.
env = pufferlib.vector.make(
    env_creator, num_envs=256, num_workers=8,
    backend=pufferlib.vector.Multiprocessing)
```

## Multi-Agent Vectorization

### Native Multi-Agent Support

PufferLib treats multi-agent environments as first-class citizens:

```python
import functools
import pufferlib.vector

# Multi-agent environment automatically vectorized.
# PufferLib has no string registry -- pass an environment constructor
# (callable). For PettingZoo envs, wrap them with
# pufferlib.emulation.PettingZooPufferEnv inside the creator. `MultiAgentEnv`
# here is your own multi-agent PufferEnv class / creator callable; bind
# constructor kwargs with functools.partial.
env = pufferlib.vector.make(
    functools.partial(MultiAgentEnv, num_agents=4),
    num_envs=128,
)

# Observations: {agent_id: [batch_obs]} for each agent
# Actions: {agent_id: [batch_actions]} for each agent
# Rewards: {agent_id: [batch_rewards]} for each agent
```

Multi-agent vectorization is handled internally — you do not implement a custom
vectorizer. Pass your multi-agent PufferEnv (or a PettingZooPufferEnv-wrapped
env) to `pufferlib.vector.make` and PufferLib batches across agents and envs.

## 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 MonitoredVecEnv:
    """Wraps a PufferLib vecenv to log throughput (no base class to subclass)."""

    def __init__(self, vecenv):
        self.vecenv = vecenv
        self.num_envs = vecenv.num_envs
        self.step_times = []
        self.step_count = 0

    def reset(self, seed=0):
        return self.vecenv.reset(seed)

    def step(self, actions):
        start = time.perf_counter()
        result = self.vecenv.step(actions)
        self.step_times.append(time.perf_counter() - start)
        self.step_count += 1

        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
import pufferlib.vector

# Reduce number of environments (lowers per-worker env count too)
env = pufferlib.vector.make(
    env_creator, num_envs=128, num_workers=8,   # was num_envs=256
    backend=pufferlib.vector.Multiprocessing)

# Use the Serial backend for debugging (single process)
env = pufferlib.vector.make(
    env_creator, num_envs=16, backend=pufferlib.vector.Serial)
```

### 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

Tune `num_envs` and `num_workers` (envs-per-worker is the derived ratio):

```python
# Start conservative, then scale up iteratively.
vec_kwargs = dict(num_envs=64, num_workers=4)      # 16 envs/worker
vec_kwargs = dict(num_envs=256, num_workers=8)     # 32 envs/worker

# Monitor and adjust:
#   sps below target  -> raise num_envs and/or num_workers
#   memory too high   -> lower num_envs (and consider cpu_offload in training)
```

### 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, terminal, truncation, info
```

### Testing

```python
import numpy as np
import pufferlib.vector

# Verify the Multiprocessing backend matches Serial for the same seed.
serial_env = pufferlib.vector.make(
    env_creator, num_envs=4, seed=42, backend=pufferlib.vector.Serial)
vec_env = pufferlib.vector.make(
    env_creator, num_envs=4, num_workers=2, seed=42,
    backend=pufferlib.vector.Multiprocessing)

serial_obs, _ = serial_env.reset(seed=42)
vec_obs, _ = vec_env.reset(seed=42)

assert np.allclose(serial_obs, vec_obs), "Vectorization mismatch!"
```

#!/usr/bin/env python3
"""
PufferLib 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.
"""

import numpy as np
import gymnasium
from pufferlib import PufferEnv


class MyEnvironment(PufferEnv):
    """
    Custom PufferLib environment template.

    This is a simple grid world example. Customize it for your specific task.
    """

    def __init__(self, buf=None, grid_size=10, max_steps=1000):
        """
        Initialize environment.

        Args:
            buf: Shared memory buffer (managed by PufferLib)
            grid_size: Size of the grid world
            max_steps: Maximum steps per episode
        """
        self.grid_size = grid_size
        self.max_steps = max_steps

        # Define spaces and num_agents BEFORE calling super().__init__(buf)

        # Define observation space.
        # IMPORTANT: a native PufferEnv obs space MUST be a Box (Dict is rejected
        # by the base class). For Dict/structured observations, build a Gymnasium
        # env and wrap it with pufferlib.emulation.GymnasiumPufferEnv instead.
        #
        # Option 1: Flat vector observation
        self.single_observation_space = gymnasium.spaces.Box(
            low=-np.inf, high=np.inf, shape=(4,), dtype=np.float32)  # [x, y, goal_x, goal_y]

        # Option 2: Image observation (still a Box)
        # self.single_observation_space = gymnasium.spaces.Box(
        #     low=0, high=255, shape=(grid_size, grid_size, 3), dtype=np.uint8)

        # Define action space
        # Option 1: Discrete actions
        self.single_action_space = gymnasium.spaces.Discrete(4)  # 0: up, 1: right, 2: down, 3: left

        # Option 2: Continuous actions
        # self.single_action_space = gymnasium.spaces.Box(
        #     low=-1.0, high=1.0, shape=(2,), dtype=np.float32)  # [dx, dy]

        # Option 3: Multi-discrete actions
        # self.single_action_space = gymnasium.spaces.MultiDiscrete([3, 3])  # Two 3-way choices

        self.num_agents = 1

        # Initialize state
        self.agent_pos = None
        self.goal_pos = None
        self.step_count = 0

        super().__init__(buf)

    def reset(self, seed=None):
        """
        Reset environment to initial state.

        Returns:
            observation: Initial observation
            info: List of info dicts (one per agent); empty list here
        """
        # 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 and an (empty) info list
        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
            terminal: Whether episode terminated (goal/failure)
            truncation: Whether episode was truncated (time limit)
            info: List of info dicts (one per agent)
        """
        self.step_count += 1

        # Execute action
        self._apply_action(action)

        # Compute reward
        reward = self._compute_reward()

        # Episode terminates on goal; truncates on the step limit
        terminal = self._is_terminated()
        truncation = self.step_count >= self.max_steps

        # Get new observation
        observation = self._get_observation()

        # Additional info (list of dicts, one per agent)
        info = []
        if terminal or truncation:
            # Include episode statistics when episode ends
            info.append({'episode': {
                'r': reward,
                'l': self.step_count
            }})

        return observation, reward, terminal, truncation, 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_terminated(self):
        """Check if episode terminated (goal reached). Timeout is a truncation, handled in step()."""
        distance = np.linalg.norm(self.agent_pos - self.goal_pos)
        goal_reached = distance < 0.5

        return goal_reached

    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 (NATIVE PufferEnv API).

    Example: cooperative navigation where each agent reaches its own goal.

    Native multi-agent uses the SAME array-based API as single-agent: set
    num_agents > 1, make single_*_space PER AGENT, and have reset/step return
    arrays whose leading dimension is num_agents. There is NO {agent_id: ...}
    dict / '__all__' convention here -- that belongs to PettingZoo (use
    pufferlib.emulation.PettingZooPufferEnv for those). obs space must be a Box.
    """

    def __init__(self, buf=None, num_agents=4, grid_size=10, max_steps=1000):
        # Define spaces and num_agents BEFORE calling super().__init__(buf)
        self.num_agents = num_agents
        self.grid_size = grid_size
        self.max_steps = max_steps

        # Per-agent observation: [pos(2), goal(2), other agent positions(2*(n-1))]
        obs_dim = 2 + 2 + 2 * (num_agents - 1)
        self.single_observation_space = gymnasium.spaces.Box(
            low=-np.inf, high=np.inf, shape=(obs_dim,), dtype=np.float32)

        # Per-agent action space
        self.single_action_space = gymnasium.spaces.Discrete(5)  # 4 directions + stay

        # Initialize state
        self.agent_positions = None
        self.goal_positions = None
        self.step_count = 0

        super().__init__(buf)

    def reset(self, seed=None):
        """Reset all agents. Returns (obs, info-list)."""
        self.agent_positions = np.random.rand(self.num_agents, 2) * self.grid_size
        self.goal_positions = np.random.rand(self.num_agents, 2) * self.grid_size
        self.step_count = 0

        return self._get_obs(), []

    def step(self, actions):
        """
        Step all agents.

        Args:
            actions: array of shape (num_agents,) of discrete actions

        Returns:
            observations: (num_agents, obs_dim) float32 array
            rewards: (num_agents,) float32 array
            terminals: (num_agents,) bool array
            truncations: (num_agents,) bool array
            info: list of dicts (PufferEnv asserts info is a list)
        """
        self.step_count += 1

        actions = np.asarray(actions).reshape(self.num_agents)
        for agent_idx in range(self.num_agents):
            self._apply_action(agent_idx, int(actions[agent_idx]))

        observations = self._get_obs()
        rewards = self._compute_rewards()
        terminals = self._terminated()
        truncations = np.full(
            self.num_agents, self.step_count >= self.max_steps, dtype=bool)

        return observations, rewards, terminals, truncations, []

    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 _compute_rewards(self):
        """Per-agent reward: negative normalized distance to each agent's goal."""
        distances = np.linalg.norm(
            self.agent_positions - self.goal_positions, axis=1)
        return (-distances / self.grid_size).astype(np.float32)

    def _terminated(self):
        """Per-agent termination: True once close to its goal."""
        distances = np.linalg.norm(
            self.agent_positions - self.goal_positions, axis=1)
        return distances < 0.5

    def _get_obs(self):
        """Build the (num_agents, obs_dim) observation array."""
        obs = np.zeros(
            (self.num_agents, self.single_observation_space.shape[0]),
            dtype=np.float32)
        for i in range(self.num_agents):
            others = np.concatenate([
                self.agent_positions[j] for j in range(self.num_agents) if j != i
            ]) if self.num_agents > 1 else np.zeros(0, dtype=np.float32)
            obs[i] = np.concatenate([
                self.agent_positions[i], self.goal_positions[i], others
            ]).astype(np.float32)
        return obs


def test_environment():
    """Test environment to verify it works correctly."""
    print("Testing single-agent environment...")
    env = MyEnvironment()

    obs, info = env.reset()
    print(f"Initial observation shape: {obs.shape}")

    for step in range(10):
        action = env.single_action_space.sample()
        obs, reward, terminal, truncation, info = env.step(action)

        print(f"Step {step}: reward={reward:.3f}, terminal={terminal}, truncation={truncation}")

        if terminal or truncation:
            obs, info = env.reset()
            print("Episode finished, resetting...")

    print("\nTesting multi-agent environment...")
    multi_env = MultiAgentEnvironment(num_agents=4)

    obs, info = multi_env.reset()
    print(f"Obs shape (num_agents, obs_dim): {obs.shape}")

    for step in range(10):
        # One action per agent: array of shape (num_agents,)
        actions = np.array([
            multi_env.single_action_space.sample()
            for _ in range(multi_env.num_agents)
        ])
        obs, rewards, terminals, truncations, info = multi_env.step(actions)

        print(f"Step {step}: mean_reward={rewards.mean():.3f}")

        if terminals.all() or truncations.any():
            obs, info = multi_env.reset()
            print("Episode finished, resetting...")

    print("\nEnvironment tests passed.")


if __name__ == '__main__':
    test_environment()

#!/usr/bin/env python3
"""
PufferLib Training Template (targets PufferLib 3.0.x)

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.

Key 3.0 API facts this template respects:
- pufferlib.vector.make(env_creator, num_envs=..., backend=...). The default
  backend is PufferEnv (native vectorization); wrapped Gymnasium/PettingZoo
  envs require backend=pufferlib.vector.Multiprocessing (or Serial).
- PuffeRL(config, vecenv, policy, logger=None) -- the first arg is a config
  DICT (not flat kwargs), the env arg is the vecenv, and the loop is driven by
  trainer.global_step < config['total_timesteps'].
- Loggers (WandbLogger/NeptuneLogger/NoLogger) live in pufferlib.pufferl and
  take the nested args config dict, not project=/name= kwargs.
"""

import argparse
import numpy as np
import torch
import torch.nn as nn
import pufferlib
import pufferlib.vector
from pufferlib.pufferl import PuffeRL, NoLogger, WandbLogger, NeptuneLogger
from pufferlib.pytorch import layer_init


class Policy(nn.Module):
    """Example policy network for a flat-vector, discrete-action env.

    PufferLib does not enforce a base class, but PuffeRL calls
    `policy.forward_eval(obs, state)` and expects it to return (logits, values).
    Structuring the network as encode_observations + decode_actions (as below)
    lets you wrap it with pufferlib.models.LSTMWrapper later without rewrites.
    The constructor takes the (driver) env so it can read single_observation_space
    / single_action_space, mirroring pufferlib.models.Default.
    """

    def __init__(self, env, hidden_size=256):
        super().__init__()
        self.hidden_size = hidden_size

        obs_dim = int(np.prod(env.single_observation_space.shape))
        num_actions = env.single_action_space.n

        self.encoder = nn.Sequential(
            layer_init(nn.Linear(obs_dim, hidden_size)),
            nn.ReLU(),
            layer_init(nn.Linear(hidden_size, hidden_size)),
            nn.ReLU(),
        )
        self.actor = layer_init(nn.Linear(hidden_size, num_actions), std=0.01)
        self.critic = layer_init(nn.Linear(hidden_size, 1), std=1.0)

    def encode_observations(self, observations, state=None):
        return self.encoder(observations.float())

    def decode_actions(self, hidden):
        return self.actor(hidden), self.critic(hidden)

    def forward_eval(self, observations, state=None):
        hidden = self.encode_observations(observations, state)
        return self.decode_actions(hidden)

    def forward(self, observations, state=None):
        return self.forward_eval(observations, state)


def resolve_env_creator(env_name):
    """Map a CLI env name to an environment constructor (callable).

    Ocean environments resolve by name in PufferLib 3.0:

        from pufferlib.ocean import env_creator
        return env_creator(f'puffer_{env_name}')   # e.g. 'puffer_breakout'

    For your own environments, maintain a mapping, e.g.:

        registry = {'my_task': MyTask, 'my_task_4p': functools.partial(MyTask, num_agents=4)}
        return registry[env_name]
    """
    raise NotImplementedError(
        f"Define a constructor for {env_name!r} (see resolve_env_creator).")


def make_env():
    """Create vectorized environment. Customize this for your task.

    PufferLib has no string registry -- pass an environment constructor
    (callable) to `pufferlib.vector.make`.
    """
    # Option 1: A PufferEnv constructor (e.g. an Ocean environment).
    #   `env_creator` is a placeholder: replace it with the callable that
    #   constructs your environment, e.g. pufferlib.ocean.env_creator('puffer_breakout').
    env_creator = None  # TODO: replace with your environment constructor (callable)
    return pufferlib.vector.make(env_creator, num_envs=256)

    # Option 2: A Gymnasium environment wrapped for PufferLib. Wrapped (non-native)
    # envs require an explicit backend -- the default backend=PufferEnv is native-only.
    # import gymnasium as gym
    # import pufferlib.emulation
    # gym_creator = lambda: pufferlib.emulation.GymnasiumPufferEnv(
    #     env_creator=lambda: gym.make('CartPole-v1'))
    # return pufferlib.vector.make(
    #     gym_creator, num_envs=256, backend=pufferlib.vector.Multiprocessing)

    # Option 3: A custom PufferEnv.
    # from my_envs import MyEnvironment
    # return pufferlib.vector.make(MyEnvironment, num_envs=256)


def create_policy(env):
    """Create policy network. PufferLib policies read single_*_space off the env."""
    return Policy(env, hidden_size=256)


def train(args):
    """Main training function."""
    # Set random seeds
    torch.manual_seed(args.seed)

    # Create environment.
    # `pufferlib.vector.make` takes an environment constructor (callable).
    # Resolve the CLI name to a constructor here; `resolve_env_creator` is a
    # placeholder you implement (Ocean: pufferlib.ocean.env_creator('puffer_<name>');
    # your own envs: a dict of PufferEnv classes / functools.partial).
    print(f"Creating environment with {args.num_envs} parallel environments...")
    env_creator = resolve_env_creator(args.env_name)
    env = pufferlib.vector.make(
        env_creator,
        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:
        args.device = 'cpu'
        print("Using CPU")

    # Build the config DICT PuffeRL expects. In real use, prefer
    # `args = pufferlib.pufferl.load_config(env_name)` and pass `args['train']`
    # so every required key is filled from pufferlib/config/*.ini. Here we set
    # the keys this template exposes explicitly.
    config = {
        'env': args.env_name,
        'seed': args.seed,
        'torch_deterministic': True,
        'cpu_offload': False,
        'device': args.device,
        'total_timesteps': args.total_timesteps,
        'learning_rate': args.learning_rate,
        'anneal_lr': True,
        'min_lr_ratio': 0.0,
        'gamma': args.gamma,
        'gae_lambda': args.gae_lambda,
        'update_epochs': args.update_epochs,
        'clip_coef': args.clip_coef,
        'ent_coef': args.ent_coef,
        'vf_coef': args.vf_coef,
        'vf_clip_coef': 0.2,
        'max_grad_norm': args.max_grad_norm,
        'batch_size': args.batch_size,
        'minibatch_size': args.minibatch_size,
        'max_minibatch_size': args.minibatch_size,
        'bptt_horizon': args.bptt_horizon,
        'compile': args.compile,
        'use_rnn': False,
        'data_dir': args.checkpoint_dir,
        'checkpoint_interval': args.checkpoint_interval,
    }

    # Logger. The 3.0 loggers read keys from a config dict; the simplest path is
    # NoLogger. Populate wandb_*/neptune_* keys in `config` if you wire those up.
    if args.logger == 'wandb':
        logger = WandbLogger(config)
    elif args.logger == 'neptune':
        logger = NeptuneLogger(config)
    else:
        logger = NoLogger(config)

    # Create trainer: positional (config, vecenv, policy, logger).
    print("Creating trainer...")
    trainer = PuffeRL(config, env, policy, logger)

    # Training loop is driven by global_step, not a fixed iteration count.
    print(f"Training to {config['total_timesteps']:,} timesteps...")
    while trainer.global_step < config['total_timesteps']:
        trainer.evaluate()      # Collect rollouts
        trainer.train()         # Train on batch
        trainer.mean_and_log()  # Aggregate + log (PuffeRL also auto-checkpoints
                                # every config['checkpoint_interval'] epochs)

    print("Training complete!")

    # Final checkpoint (save_checkpoint takes no path; it writes under data_dir).
    final_path = trainer.save_checkpoint()
    print(f"Saved final model to {final_path}")


def main():
    parser = argparse.ArgumentParser(description='PufferLib Training')

    # Environment
    parser.add_argument('--env-name', type=str, default='my-env',
                        help='Env key resolved to a constructor by '
                             'resolve_env_creator (PufferLib has no string registry)')
    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 (defaults mirror pufferlib/config/default.ini, 3.0.x)
    parser.add_argument('--total-timesteps', type=int, default=10_000_000,
                        help='Total environment steps to train for')
    parser.add_argument('--learning-rate', type=float, default=0.015,
                        help='Learning rate (PufferLib default optimizer is muon)')
    parser.add_argument('--batch-size', type=int, default=32768,
                        help='Timesteps per training batch')
    parser.add_argument('--minibatch-size', type=int, default=8192,
                        help='Minibatch size (also used as max for grad accumulation)')
    parser.add_argument('--bptt-horizon', type=int, default=64,
                        help='BPTT / rollout horizon per segment')
    parser.add_argument('--update-epochs', type=int, default=1,
                        help='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.995,
                        help='Discount factor')
    parser.add_argument('--gae-lambda', type=float, default=0.90,
                        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.001,
                        help='Entropy coefficient')
    parser.add_argument('--vf-coef', type=float, default=2.0,
                        help='Value function coefficient')
    parser.add_argument('--max-grad-norm', type=float, default=1.5,
                        help='Maximum gradient norm')

    # Logging
    parser.add_argument('--logger', type=str, default='none',
                        choices=['wandb', 'neptune', 'none'],
                        help='Logger to use (set wandb_*/neptune_* config keys to wire up)')

    # Checkpointing (PuffeRL auto-saves under data_dir every checkpoint_interval epochs)
    parser.add_argument('--checkpoint-dir', type=str, default='experiments',
                        help='data_dir where checkpoints are written')
    parser.add_argument('--checkpoint-interval', type=int, default=200,
                        help='Checkpoint save frequency (epochs)')

    # 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)


if __name__ == '__main__':
    main()