Skip to content

Pursuit IPPO Experiment Helper

src/experiments/pursuit_experiment.py retains run_pursuit_ippo_experiment(...) as the single-algorithm IPPO entrypoint. The four active Pursuit notebooks use run_pursuit_experiments(...) so they can select any of the 15 supported variants; see Pursuit Multi-Algorithm Notebooks for that suite contract.

Shared notebook helper

Direct callers pass an explicit resolved config:

config = apply_cli_overrides({...})
experiment = run_pursuit_ippo_experiment(spec, config)

The helper is responsible for:

  • constructing the wrapped PursuitEvasion environment,
  • validating notebook config keys and applying CLI overrides that Marimo parsed,
  • mapping those knobs onto an IPPO config,
  • running all seeds,
  • collecting wrapper histories for plotting,
  • writing one checkpoint bundle per run,
  • writing runs.npy, nash.npy, nash_report.json, and related result artifacts as soon as the algorithm finishes,
  • returning runs, nash_runs, nash_summary, and save metadata for notebook cells.

The suite notebooks themselves:

  • define the experiment spec,
  • define the full user-facing config in one place,
  • show the resolved runtime config in the notebook output,
  • render the NashConv summary,
  • plot EWMA success,
  • display the artifact paths returned by the helper.

IPPO config mapping

The direct helper accepts shared orchestration and flat IPPO knobs, including:

  • num_env_steps
  • num_runs
  • progress_chunks
  • calculate_conv
  • calculate_alpharank
  • calculate_satisfaction
  • calculate_conv_quick
  • eval_episodes
  • br_train_iters
  • br_batch_steps
  • satisfaction_eval_episodes
  • satisfaction_delta
  • nashconv_delta
  • nashconv_auto_resume
  • nashconv_auto_resume_threshold
  • nashconv_auto_resume_step_fraction
  • resume_partial_runs
  • save_dir
  • num_envs
  • num_steps
  • num_minibatches
  • update_epochs
  • lr
  • anneal_lr
  • max_grad_norm
  • activation
  • gamma
  • gae_lambda
  • clip_eps
  • vf_coef
  • ent_coef

The helper resolves IPPO training config separately:

  • num_steps=min(256, target_chunk_size) when the notebook leaves num_steps=None,
  • total_timesteps is rounded up to a whole number of PPO updates,
  • num_minibatches is reduced automatically if the requested value does not divide the batch size.

progress_chunks controls how often the notebook progress bar redraws, not the PPO rollout length itself. The trainer still emits metrics every update, but the shared experiment helper only forwards progress to tqdm at milestone boundaries. This keeps notebook memory use bounded while avoiding output spam.

Checkpoint bundle format

Each run writes a pickle bundle under the experiment checkpoints/ directory.

Bundle fields:

  • agent_ids
  • run_seed
  • config
  • train_states

train_states is serialized with flax.serialization.to_state_dict(...), so the bundle is stable to pickle without storing live TrainState objects directly.

IPPO NashConv path

Notebook verification runs directly from saved IPPO train states.

approximate_nashconv_bernoulli_ippo(...):

  • evaluates the baseline policy profile directly from saved IPPO train states,
  • trains one best response at a time by warm-starting IPPO from the snapshot,
  • freezes non-target agents with trainable_agent_ids={target_agent},
  • reuses the existing NashConv payload shape so save_nashconv_report(...) stays unchanged.

BR defaults intentionally use:

  • NUM_ENVS=1
  • NUM_MINIBATCHES=1
  • TOTAL_TIMESTEPS = br_train_iters * br_batch_steps
  • NUM_STEPS = br_batch_steps

That keeps BR training easy to reason about and matches the notebook-facing br_train_iters / br_batch_steps controls.