Skip to content

PettingZoo IPPO Trainer

src/rl/ippo.py now targets PettingZoo ParallelEnv instances directly.

Supported entrypoint

Use:

train = make_train(config, env_factory)
out = train(seed=0)
  • config is the PPO hyperparameter dictionary.
  • env_factory(seed) must return a fresh PettingZoo ParallelEnv.

The trainer no longer depends on jaxmarl.make(...), ENV_NAME, or ENV_KWARGS.

The returned trainer callable now also accepts optional keyword arguments:

out = train(
    seed=0,
    initial_train_states=...,
    trainable_agent_ids={"player_0"},
    on_update=callback,
    capture_env_histories=True,
)
  • initial_train_states: warm-start all agents from a prior IPPO snapshot.
  • trainable_agent_ids: update only that subset; all other agents act as frozen opponents.
  • on_update(update_idx, total_updates, metric): called once per PPO update with scalar metrics.
  • capture_env_histories=True: includes env_histories, a per-env list of serialized wrapper histories.

Current assumptions

  • PettingZoo ParallelEnv only.
  • gymnasium.spaces.Box observations only.
  • gymnasium.spaces.Discrete actions only.
  • One independent actor-critic and optimizer state per agent.
  • Whole-environment autoreset on any agent termination or truncation.

That autoreset rule matches the current experiment environments and reward wrappers, which treat the episode as finished once any agent is done.

Rollout vs update split

Environment interaction is Python-side because PettingZoo envs are not pure JAX functions.

  • Rollouts are collected by stepping NUM_ENVS live env instances in Python.
  • PPO updates stay in JAX/Flax and are JIT-compiled per agent.
  • GAE, PPO clipping, and optimizer scheduling remain close to the earlier implementation.

Metrics

Episode metrics are now trainer-owned rather than coming from a JAXMARL log wrapper.

Returned metrics include:

  • episode_return_mean
  • episode_length_mean
  • completed_episodes
  • PPO loss scalars
  • per_agent_episode_return_mean

If no episode finishes during an update, the return/length means are recorded as NaN.

When trainable_agent_ids is narrower than the full agent set, PPO loss scalars are aggregated over the trainable agents only. Episode-return metrics still cover all agents.

Warm-starting and frozen-opponent BR runs

The warm-start path restores the serialized train-state contents into a fresh TrainState template built from the current config. That keeps the active optimizer definition tied to the current run config while preserving learned params, step, and optimizer state from the snapshot.

This is what the IPPO NashConv verifier uses for approximate best responses:

  • clone the saved train states,
  • rebuild a fresh IPPO trainer with BR-specific rollout sizes,
  • set trainable_agent_ids={target_agent},
  • train only the deviating agent while opponents stay fixed at the snapshot.

Captured histories

When capture_env_histories=True, the trainer walks the wrapper chain of each live env and serializes to_dicts() or history_by_agent if available. This keeps notebook plotting compatible with the existing plot_ewma_success_all_agents(...) utilities.

Follow-up candidates

  • Optional shared-parameter mode for homogeneous-agent experiments.
  • Shared checkpoint load helpers for restoring serialized notebook bundles.
  • Support for partial-agent death or AEC envs if future environments need that behavior.