DOCSEARCHbyAlgolia
Submit Docs

PufferLib Documentation

repository·4.0·Indexed 24 days ago

https://github.com/pufferai/pufferlib

A fast reinforcement learning library designed for training small, superhuman models. PufferLib provides integrated learning algorithms, hyperparameter tuning, and simulation methods. It includes specialized environment support for NetHack (via fast-nle), Overcooked (with configurable layouts and reward systems), and Craftax (featuring a native C implementation transition for high-performance step execution).

Tokens
8.9K
Snippets
13
Records
56
Agent score
91%

What's inside PufferLib

  1. Overview of PufferLib

    4.0
    PufferLib is a high-performance reinforcement learning library designed to train small, superhuman models rapidly. It includes custom learning algorithms, hyperparameter tuning, and simulation methods developed through internal research. The library is open-source and optimized for speed and usability in reinforcement learning workflows.

  2. Understand the Craftax Ocean Environment Implementation

    4.0

    The ocean/craftax/ environment is a full Craftax Ocean implementation with a symbolic observation size of 8268 and an action count of 43.

    Current Architecture:

    • Reset: The c_reset function is implemented natively in C. It generates the full 9-floor world state (including maps, item maps, mob maps, light maps, ladders, chest flags, and scalar reset fields) and encodes them into a flat symbolic observation.
    • Step: The c_step function is currently a Python/JAX proxy. It acquires the Python GIL, calls the JAX Craftax-Symbolic-v1 implementation, and copies the resulting float32 observation, reward, terminal flag, and terminal achievement log into PufferLib-owned buffers.
    • State Alignment: After a native reset, the first delegated step performs an internal proxy reset to ensure the JAX-owned step state remains aligned with the native reset observation.

    Note on Performance: Because c_step still allocates through Python/JAX and serializes on the GIL, training throughput is currently limited and not yet optimized for high-performance targets.

  4. Verify Craftax parity with the stress battery

    4.0

    To ensure correctness between JAX and native implementations, use the stress battery in tests/craftax_parity_stress.py. The battery executes the following configurations:

    • 64 seeds $\times$ 10,000 steps with mixed policy.
    • 16 seeds $\times$ 30,000 steps with descend policy.
    • 32 seeds $\times$ 5,000 steps with suicide policy.
    • 16 seeds $\times$ 5,000 steps with combat policy.

    All stress cases use atol=1e-5 for observations and rewards and require exact terminal matching.

  5. Watch a NetHack policy

    4.0

    To watch a trained policy, first build the demo binary using the --fast flag, then run the ./nethack binary.

    To play specific weights, set the NH_WEIGHTS environment variable to the path of your checkpoint file.

    Binary Arguments:

    • [steps]: Number of steps to run.
    • [ms_per_frame]: Milliseconds per frame. Setting this to 0 runs the environment headless.

    Environment Variables:

    • NH_WEIGHTS: Path to the weights file (e.g., checkpoints/nethack/<run>/<step>.bin).
    • NH_SEED: Set this to replay a specific run.

    Example usage:

    ./build.sh nethack --fast
./nethack
NH_WEIGHTS=checkpoints/nethack/<run>/<step>.bin ./nethack

  6. Understand the Craftax Native Step Integration Roadmap

    4.0

    The Craftax environment is transitioning from a Python/JAX proxy-based step execution to a native C implementation (c_step). This transition is being done via standalone subsystem ports that are verified against JAX for parity before being integrated.

    Current Status (as of 2026-04-18):

    • Completed: Native reset (PRNG, noise, world gen), simple subsystems, medium subsystems, crafting/placement subsystems, and the do_action subsystem.
    • Remaining Tasks:
      • Native ports for update_mobs and spawn_mobs.
      • Native bookkeeping for rewards, terminal states, timesteps, light-levels, RNG, and achievement-deltas.
      • Integration of all standalone ports into a single native c_step behind an explicit switch.
      • Removal of the Python/JAX proxy from the normal step path.
      • Restoring production vector sizes in config/ocean/craftax.ini.

    Note: Until the integration is complete, the live Ocean environment still delegates steps to the Python/JAX proxy, and rendering remains a no-op.

  8. Use deterministic action policies in the Craftax parity harness

    4.0

    The standalone parity harness (tests/craftax_parity.py) supports several deterministic action policies for testing beyond uniform random exploration:

    • uniform: Original random action stream.
    • combat: Biases toward DO, arrows, fireballs, and iceballs when meaningful; otherwise moves toward live mobs.
    • descend: Uses mirrored state to push toward down ladders, clear blocked levels via combat, and exercise placement/crafting.
    • suicide: Steers into lava, water, mob-occupied, or projectile-heavy danger.
    • boss: Warm up with downward navigation, then repeatedly attempts descent while routing toward ladders.
    • mixed: Round-robins the above policies every 500 steps.

    When a divergence is detected, the harness reports the policy, seed, step, action, reward delta, terminal delta, first symbolic-observation field, suspected subsystem, and the last 10 actions.

  9. Build and Run the Overcooked Environment

    4.0

    To use the Overcooked environment, you must first build it using the provided setup script, then you can run tests or standalone demos.

    Build

    python setup.py build_overcooked --inplace

    Run Tests

    python pufferlib/ocean/overcooked/overcooked.py

    Run Standalone Demos

    Run the ./overcooked executable followed by the layout name:

    ./overcooked cramped_room
./overcooked asymmetric_advantages
./overcooked forced_coordination
./overcooked coordination_ring
./overcooked counter_circuit
    # Build the environment
python setup.py build_overcooked --inplace

# Run standalone test
python pufferlib/ocean/overcooked/overcooked.py

# Run standalone demo with specific layout
./overcooked cramped_room
./overcooked asymmetric_advantages
./overcooked forced_coordination
./overcooked coordination_ring
./overcooked counter_circuit

  10. Setup the NetHack environment

    4.0

    To use the NetHack environment, install the current package in editable mode and run the build script to clone and build the fast-nle vendor and the training backend.

    Note: Run these commands from the repository root. The engine expects data at vendor/fast-nle/build/dat. You can override this location by setting the NETHACKDIR environment variable.

    pip install -e .
./build.sh nethack

  11. Configure the Overcooked Environment

    4.0

    Initialize the Overcooked environment using the Python wrapper. You can specify the number of parallel environments, the layout, the number of agents, and various reward configurations to customize the training signal.

    Key configuration parameters include:

    • num_envs: Number of parallel environments.
    • layout: The name of the map layout (e.g., cramped_room, asymmetric_advantages).
    • num_agents: Number of agents per environment.
    • render_mode: Set to enable rendering.
    • reward_*: Various keys to adjust the weight of main and intermediate rewards (e.g., reward_dish_served_whole_team, reward_pot_started, reward_ingredient_picked).
    env = Overcooked(
    num_envs=1,                          # Number of parallel environments
    layout="cramped_room",               # Layout name (see Available Layouts)
    num_agents=2,                        # Agents per environment
    render_mode=None,                    # Set to enable rendering
    log_interval=128,                    # Steps between log aggregation
    grid_size=32,                        # Render tile size in pixels

    # Reward configuration (from config/ocean/overcooked.ini)
    reward_dish_served_whole_team=1.0,   # Shared reward for correct dish
    reward_dish_served_agent=0.0,        # Bonus for serving agent
    reward_pot_started=0.15,             # Starting correct recipe
    reward_ingredient_added=0.15,        # Adding onion to pot
    reward_ingredient_picked=0.05,       # Picking up ingredient
    reward_soup_plated=0.20,             # Plating cooked soup
    reward_wrong_dish_served=0.0,        # Serving incorrect dish
    reward_step_penalty=0.0,             # Per-step penalty
)

  12. Configure Craftax Environment via INI

    4.0

    The environment configuration is managed via config/ocean/craftax.ini.

    Warning: Currently, this configuration uses a small proxy-friendly vector size to accommodate the Python/JAX step proxy. Once the step logic is fully ported to native C, this vector size should be increased.