Skip to content

Marimo Notebooks

Runnable notebooks are reactive marimo applications stored as Python files. They work in the editor and in non-interactive script/cluster execution.

Notebook Families

Name Role
train_wm_mo.py World-model data, training, evaluation, and graph augmentation.
train_om_mo.py Opponent dataset preparation and imagined-level training.
experiments_mo.py Baseline and shielded RL experiments.
play_mo.py Interactive policy/environment inspection.
true_dynamics_mo.py Exact transition graph inspection.

Generic reports and viewers live under notebooks/generic/; conceptual graph notebooks live under notebooks/diagrams/; reusable non-UI helpers live under notebooks/utils/. notebooks/ablations/ contains the maintained WM/OM strength studies, while notebooks/benchmarks/ contains environment-only benchmark notebooks rather than the canonical WM → OM → experiment pipeline.

Interactive Use

uv run marimo edit notebooks/matrix/chicken/experiments_mo.py

Marimo derives a dependency graph from cell definitions. mo.stop(...) skips an expensive stage and its dependent cells, which is how disabled algorithms avoid loading their artifacts.

Script Use and CLI Overrides

uv run python notebooks/matrix/chicken/experiments_mo.py \
  --run-ippo=true \
  --num-runs=1 \
  --timesteps=64 \
  --rollouts=64

apply_cli_overrides(...) maps kebab-case CLI names to snake-case config keys, coerces simple booleans/numbers/None, and raises on unknown keys. Keys are case-insensitive. In notebooks where an early cell needs n_agents, that cell reads get_cli_overrides() and the complete config is validated later.

All five algorithm run_* switches and the separate run_exact_anytime_diagnostic switch default to false. Select only the stages a scripted experiment should run. The cluster launchers own the algorithm flags and set only the learners assigned to each job; the exact-anytime switch is an offline certificate audit rather than a sixth learner.

Within run_ippo_shielded, shield_action_mode and shield_opponent_mode select condition variants. experiment_tag isolates durable stage outputs and seed_offset reserves a disjoint trainer-seed block. See Experiments for the condition matrix and tagged-report workflow.

Stable Working Directory

Runnable notebooks call set_cwd_to_project_root() in their setup cell. Paths such as exports/... are therefore stable whether execution starts from the repository root, the notebook directory, marimo, or a cluster runner.

Timing and Outputs

Notebooks use perf_counter() and the shared helpers in notebooks/utils/timings.py. The canonical files are environment-root timings.json and timings_hardware.json; locked, atomic writes allow several algorithm jobs to merge their stage records safely.

See Configuration and CLI for common keys and Artifacts for output paths.