Skip to content

Marimo Experiment Notebooks

The experiment notebooks under notebooks/*/*/experiments_mo.py are marimo notebooks, not Jupyter notebooks. They should use marimo-native control flow and avoid carrying over IPython magic patterns like %%ifcell and %time.

Execution Control

  • Keep run switches in the notebook config dictionary rather than standalone locals.
  • Apply CLI overrides with notebooks.marimo_cli.apply_cli_overrides(...) so the same notebook works in interactive marimo runs and scripted/HPC runs.
  • When an early cell needs one dimension (for example n_agents) before the complete experiment config exists, read it with get_cli_overrides() and defer apply_cli_overrides(...) until the full config cell. Marimo executes cells with generated frame filenames, so source-wide key discovery is not reliable from an early partial-config call; validating there can incorrectly reject later run_*, num_runs, or n_steps arguments.
  • Standard experiment flags are:
  • run_ippo
  • run_ippo_lag
  • run_icpo
  • run_ippo_shielded
  • run_ippo_true_shielded
  • run_exact_anytime_diagnostic
  • just_vis
  • shield_debug_mode
  • shield_trace_limit
  • shield_missing_coverage_log_limit
  • shield_patience
  • All algorithm run_* flags and the independent exact-anytime diagnostic flag default to False in every standard environment notebook. Scripted and cluster runs select only the stages they intend to execute; omitted algorithms and diagnostics stay disabled. The exact-anytime flag builds and audits a certificate without initializing an RL trainer or consuming the policy-training step budget.

This makes script-mode invocations predictable, for example:

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

Why mo.stop(...)

Use mo.stop(...) at the top of expensive cells to replace old notebook-magics that conditionally skipped cells.

  • It is supported by marimo.
  • It stops the current cell cleanly.
  • It also prevents descendant cells from running when they depend on skipped outputs, which is exactly what we want for multi-stage experiment setup like shield construction followed by shielded training.

This is preferred to reviving IPython-specific helpers in src.core, because those helpers only made sense inside Jupyter magic execution and do not map cleanly onto marimo's dependency model.

Timing Sidecars

Replace %time with time.perf_counter() and store durations in a shared root sidecar named timings.json. Store hardware descriptions separately in timings_hardware.json.

  • The canonical file lives at the environment export root, not under wm/, om/, or experiments/.
  • For example, exports/gridworlds/gathering/4/timings.json is shared by the gathering world-model, opponent-model, and experiment notebooks.
  • The timing structure is a dictionary keyed by stage name.
  • Each stage entry contains total_seconds and runs.
  • runs stores seconds, a hardware_id, and an optional run_idx, label, or caller-provided metadata.
  • hardware_id resolves through the sibling timings_hardware.json, whose hardware map deduplicates full descriptions across stages and runs.

Timings now use JSON by default.

  • notebooks/utils/timings.py treats root timings.json as canonical.
  • Saves use per-file advisory locks, merge stage keys with the existing file, and replace files atomically after flushing them. This prevents different experiment processes from dropping each other's stage records or leaving a partially written JSON document.
  • The loader reads only the canonical root JSON file. Other filenames and pickle timing payloads are not accepted.

Hardware allocation semantics

The hardware sidecar describes resources available to the timed process:

  • RAM is the smallest finite limit found while walking the process's cgroup hierarchy. This handles cgroup v2 memory.max and cgroup v1 memory.limit_in_bytes, including the layouts used by Slurm and PBS.
  • When no finite cgroup limit is visible, Slurm SLURM_MEM_PER_NODE/SLURM_MEM_PER_CPU is the fallback. Per-CPU memory is multiplied by the allocated job CPU count. Total physical node RAM is never substituted, because it overstates a cluster job's allocation.
  • CPU availability comes from sched_getaffinity(0) where supported, with scheduler counts and os.cpu_count() only as fallbacks. /proc/cpuinfo is filtered to the accessible logical CPU IDs so physical cores and SMT threads remain distinct.
  • CPU clock ranges normally come from Linux cpufreq and record clock_source: linux_cpufreq. When ACPI CPPC exposes a higher autonomous boost capability, the maximum is derived from nominal_freq * highest_perf / nominal_perf and the source is linux_cpufreq_cppc. KVM guests without these physical interfaces may use a fixed /proc/cpuinfo clock only when every visible vCPU advertises hypervisor and tsc_known_freq and all values agree within 0.1 MHz; that source is proc_cpuinfo_virtual_nominal. Varying readings remain unavailable.
  • A virtual machine's processor name is the guest-visible model and may be generic rather than the physical host CPU SKU.
  • Visible NVIDIA GPUs are filtered by CUDA_VISIBLE_DEVICES or NVIDIA_VISIBLE_DEVICES. The record includes model, VRAM, memory-bus width, maximum memory clock, and derived peak memory bandwidth when NVML exposes them.
  • RAM and VRAM are capacity/allocation metadata, not process peak usage. Use scheduler accounting such as MaxRSS for observed peak memory.

Shared Helper

Common experiment-orchestration helpers live in notebooks/utils/experiments.py.

  • should_run_stage(...) centralizes the flag and not just_vis rule.
  • record_timing(...) updates the in-memory timing structure used across cells.
  • summarize_timings(...) gives a lightweight text summary for notebook output.

record_timing(...) records the current notebook run, not an all-time history.

  • For indexed experiment stages, run_idx=0 resets that stage and later run indices append to the same fresh batch. A normal three-seed IPPO run should therefore leave three entries, not six after two executions.
  • For single-shot stages without run_idx, such as wm_train_model, om_train_policy, component_om, and shield setup, each call replaces the previous entry for that stage.
  • This keeps the shared root sidecar useful across WM, OM, and experiment notebooks while preventing stale repeated executions from inflating totals.

Shared timing file helpers live in notebooks/utils/timings.py.

  • load_shared_timings(...) loads the canonical root timing sidecar.
  • load_shared_timing_hardware(...) loads the deduplicated hardware map.
  • save_shared_timings(...) always writes the canonical timings.json file.
  • timings_filename(...), timings_hardware_filename(...), and shared_timings_root(...) keep notebook path handling consistent.

If the notebook patterns grow beyond simple orchestration, a future refactor could move more of the repeated training-stage assembly into shared functions. For now, the repo keeps the RL setup explicit inside each notebook so environment-specific changes stay easy to inspect.

Shielded experiment setup is now one of those shared paths. The notebooks should use notebooks/utils/shielding.py for bundle loading/caching and fresh shielded env construction instead of inlining old one-off shield code.

World-model notebooks should record the stable stages:

  • component_wm
  • wm_legal_graph
  • wm_collect_rollouts
  • wm_train_model
  • wm_build_env_transition_graph
  • wm_log_predictions
  • wm_evaluate_report

Opponent-model notebooks should record the stable stages:

  • component_om
  • om_prepare_inputs
  • om_fit_level0
  • om_train_policy
  • om_save_artifacts

save_shared_timings(...) now also synthesizes component_wm and component_om as the sum of the stable wm_* / om_* stage totals whenever those stage families are present. That keeps the detailed breakdown while making the end-to-end component time obvious in the same sidecar. component_wm is always refreshed from the current wm_* stage sum because WM notebooks do not record a separate explicit component timer.

For opponent-model notebooks, prefer recording an explicit component_om end-to-end timing at notebook completion as well.

  • This captures the real wall-clock runtime even if one of the fine-grained om_* stage timings is missing in a given notebook version.
  • refresh_component_totals(...) treats that explicit component entry as authoritative when it is larger than the synthesized om_* stage sum, and otherwise falls back to the stage rollup.
  • This means a completed notebook run can still report a truthful component_om total even when the detailed stage family is incomplete; otherwise the detailed stage sum is used.

Timing summaries should render entries as:

  • label when a labeled entry is present
  • run {run_idx} for run-indexed entries
  • entry {ordinal} for unlabeled entries that are neither runs nor setups

When multiple training stages share the same helper, import that helper in the notebook's shared imports cell and return it as a marimo dependency instead of repeating local imports inside each stage cell.

  • This keeps the notebook's dependency graph explicit.
  • It also avoids cell-by-cell drift when a shared helper name changes.

Artifact Dependencies

Baseline experiment stages should not eagerly validate shield-training artifacts.

  • run_ippo, run_ippo_lag, and run_icpo only need the environment factory.
  • run_ippo_shielded depends on the learned shield artifacts built from wm/env_transition_graph.pkl and om/iop_stack.pt.
  • run_ippo_true_shielded depends on the cached true-transition graph under true_shield/, not on wm/world_model.pt.
  • run_exact_anytime_diagnostic builds from the serialized exact graph and writes below true_shield/exact_anytime/; it does not require learned WM/OM artifacts or a trained policy. Matrix graphs are simulator-stepped, whereas stochastic gridworld certificates remain graph-relative until their separate successor generators are exhaustively tied to the simulator semantics.

The shared environment-setup cell in experiments_mo.py should therefore only build and return create_env(...).

  • Do not eagerly call load_world_model_mlp(...) there just to check for wm/world_model.pt.
  • Let the shielded setup cells fail locally when their own required exports are missing.

This keeps plain PPO/ICPO experiment runs usable before any world-model or opponent-model training has happened, while preserving explicit artifact errors for the shielded paths that actually need them.

Figure Exports

Experiment notebooks should save both cumulative and episodic reward/safety variants when world-model pretraining is part of the workflow.

  • cumulative.png and episodic.png keep the existing WM-aware view by passing wm_hist=..., so the shielded run is plotted on the global timeline that includes the world-model pretraining prefix.
  • cumulative_post.png and episodic_post.png omit wm_hist, so the same metrics start at the policy-training phase without the world-model timesteps on the x-axis.
  • All four exports should set truncate_to_smallest=True so algorithms are compared on a shared post-processing horizon instead of letting one longer run stretch the x-axis.

Marimo also treats names assigned inside a cell as notebook-level definitions for dependency tracking, so short-lived locals still need to be unique across sibling cells.

  • Avoid reusing names like setup_start in multiple cells of the same notebook.
  • Prefer stage-scoped names such as ippo_shielded_setup_start and ippo_true_shielded_setup_start.

First Cell Layout

When a marimo notebook needs import marimo as mo, keep that import inside the first meaningful code cell instead of dedicating a standalone cell just to returning mo.

  • For title cells, import mo directly inside the markdown-rendering cell.
  • This keeps the notebook header visually important without adding an extra dependency-only cell at the top.

CLI Parameters

Older .ipynb notebooks used papermill-style parameter cells. In marimo notebooks, prefer CLI-backed defaults instead.

  • Use notebooks.marimo_cli.apply_cli_overrides(...) for notebook CLI config, even when there is only one tunable value like n_steps or n_agents.
  • apply_cli_overrides(...) is backed by mo.cli_args(), normalizes --kebab-case to snake_case, lowercases keys so --N_AGENTS and --n-agents match the same parameter, and coerces simple literals like integers, floats, booleans, and None.
  • Unknown CLI args now raise immediately instead of being silently ignored.
  • apply_cli_overrides(...) validates against notebook-declared CLI keys, so an early scalar parameter cell can coexist with a later full experiment config cell in the same notebook.

Example:

from notebooks.marimo_cli import apply_cli_overrides

cli_config = apply_cli_overrides(
    {
        "n_steps": 500_000,
        "max_risk": 0.2,
    }
)

n_steps = int(cli_config["n_steps"])
max_risk = float(cli_config["max_risk"])

World Model Training Notebooks

For train_wm_mo.py marimo notebooks, keep the trained world model flowing through downstream cells as model instead of reloading the checkpoint into a second variable like model_1.

  • This matches marimo's dependency model better because the sample-logging and evaluation cells depend directly on the output of the training cell.
  • It also avoids notebook drift where later cells silently inspect a freshly reloaded checkpoint rather than the in-memory model that was just trained in the current run.
  • If a future workflow genuinely needs checkpoint-only evaluation, add that as an explicit alternate path with a clearly named flag rather than reusing model_1 as an implicit fallback.

When downstream cells need to work both after a fresh training run and in a later notebook session, use an explicit eval_model value.

  • eval_model should be model when the training cell has already produced an in-memory model.
  • Otherwise, rebuild the model with the notebook's INPUT_DIM, N_CELLS, N_CHANNELS, and REWARD_DIM, then load world_model.pt.
  • Point report and sample-logging cells at eval_model so notebook reruns still work even if only the saved checkpoint is available.