Skip to content

Getting Started

The project requires Python 3.11 or newer and uses uv for dependency management and command execution.

Install Dependencies

From the repository root:

uv sync

The dependency set includes GPU-oriented JAX and PyTorch packages, Marimo, plotting libraries, and the documentation toolchain. A first sync can be large. The gridworld substrates are native code and do not install Melting Pot or Lab2D.

Preview The Documentation

The docs are built with Zensical. The Markdown source lives in docs/, and generated HTML is written to site/.

uv run zensical serve

The preview server listens on localhost:8000 by default and reloads when docs change.

To build the static site:

uv run zensical build --strict

The generated site/ directory is ignored by git and should not be committed.

Run A Notebook

Experiment entrypoints are Marimo notebooks stored as Python files. Open one with:

uv run marimo edit notebooks/matrix/stag_hunt/SustainedStag_2_mo.py

or serve it as a read-only app:

uv run marimo run notebooks/matrix/stag_hunt/SustainedStag_2_mo.py

For non-interactive scheduler runs, use the cluster launchers instead of keeping a notebook server alive.

Most training notebooks define a visible config block near the top. Common knobs include:

  • num_env_steps: total training steps.
  • seed: base seed for repeated runs.
  • num_runs: number of seeds to evaluate.
  • run_*: algorithm toggles.
  • calculate_conv: whether to estimate NashConv.
  • calculate_alpharank: whether to compute AlphaRank.
  • calculate_satisfaction: whether to run final satisfaction rollouts.
  • satisfaction_eval_episodes: number of final-satisfaction episodes.
  • nashconv_delta: overall failure probability for joint NashConv rollout bounds.
  • satisfaction_delta: overall failure probability for satisfaction bounds.
  • nashconv_auto_resume: whether a high-NashConv run receives one bounded continuation-training pass.
  • resume_partial_runs: whether validated completed seeds are reused after an interrupted invocation.
  • save_dir: export directory for reports, checkpoints, and plots.

Some top-level config values can be overridden from the CLI through the notebook helper:

uv run marimo run notebooks/gridworlds/territory/RoleClaim_2_mo.py -- \
  --num-env-steps=100000 --num-runs=1 --run-ippo=true

Nested algorithm blocks such as ippo, iql, and nashq are usually edited in the notebook.

A Small First Run

For a lightweight end-to-end run, use the matrix smoke tool. Its defaults run one Sustained Stag seed with IPPO, IQL, and Nash-Q over 4,096 environment steps:

JAX_PLATFORMS=cpu uv run python scripts/run_matrix_smoke.py

It writes ordinary experiment artifacts under exports/matrix/stag_hunt/sustained_stag/2/. Use --help to select other Stag Hunt objectives or change the smoke budgets.

The smoke tool defaults to --clean and removes that suite directory before running. Use it only when those exports are disposable; pass --no-clean to preserve unrelated files, while remembering that the selected algorithm files are still rewritten.

For a full configurable notebook, open:

uv run marimo edit notebooks/matrix/stag_hunt/SustainedStag_2_mo.py

For gridworld inspection before training, use one of the play notebooks:

uv run marimo edit notebooks/gridworlds/territory/play_labelled_mo.py

Play notebooks are useful for checking action maps, labels, temporal atoms, and whether a formula describes the behavior you intend.