Skip to content

Cluster Execution

Dependency-aware launchers live in notebooks/experiments/scripts/. They validate environment names and scheduler targets, route notebook arguments to the correct stage, and pass only the algorithms selected for each job. The examples assume a remote checkout at ~/Projects/omsh, with ~/scripts/omsh linked to the launcher directory.

Supported Environment Keys

congestion, gathering, markov_stag_hunt, pursuit, bertrand, chicken

An explicit notebook directory or .py path is also accepted by the per-environment launchers. A misspelled bare key is rejected. The curated suite launchers (pipeline.sh, pipeline_all_envs*.sh) cover the five original paper keys; Congestion runs through pipeline_env*.sh and env*.sh.

The four primary algorithm names are ippo, ippo_lag, cpo, and omsh_monotone. omsh_credible is explicitly selectable for diagnostic ablations with algorithms=omsh_credible; the two shield modes cannot share one job.

Slurm: One Complete Environment

bash notebooks/experiments/scripts/pipeline_env.sh \
  gathering training num_runs=3

The launcher submits WM, attaches OM with afterok:<wm-job-id>, then attaches the four primary algorithm jobs with afterok:<om-job-id>.

For a configurable environment:

bash notebooks/experiments/scripts/pipeline_env.sh \
  pursuit a100 n_agents=3 n_steps=1000000 num_runs=3

For coordinated multi-role shielding, om_focal_indices creates one OM job per role. The jobs run in sequence because they share a timing file, and the algorithm jobs wait for the last one:

bash notebooks/experiments/scripts/pipeline_env.sh congestion amd24 \
  n_agents=2 om_focal_indices=0,1 shielded_agent_ids=player_0,player_1

Argument routing is deliberate:

Keys Reach
n_agents, model_variant, reward_game_variant Every stage
n_steps, pretrain_max_steps, wm_* (including wm_variant), reward_inference_* World model only
n_levels, om_*, agent_updates, opp_updates, imagination_*, agent_rollouts, opp_rollouts, opponent_rollout_horizon Opponent model only
venv, html Cluster runner
Everything else, such as num_runs Algorithm stage only

The launcher rejects caller-provided run_*, run_algorithms, and just_vis, because it owns the complete algorithm selection.

Pipelines use notebook defaults for the algorithm stage

pipeline_env.sh passes no experiment_tag, no campaign seeds, and no frozen-policy evaluation unless you add them. Use it to train the world and opponent models; use corrected_final_campaign.sh (or pass the full campaign protocol) for algorithm runs you intend to report.

Slurm: Suites and Campaigns

The curated partition assignment is launched with:

bash notebooks/experiments/scripts/pipeline.sh num_runs=3

The full six-configuration matrix uses one selected partition. It includes every curated environment once, with Pursuit at both two and three agents, which makes 36 jobs:

bash notebooks/experiments/scripts/pipeline_all_envs.sh a100 num_runs=5

On cpucluster, use bundles=4 to give every primary algorithm its own job:

bash notebooks/experiments/scripts/pipeline_all_envs.sh \
  amd48 bundles=4 num_runs=5

The published algorithm protocol for the five original environments is encoded in one script. It submits algorithm jobs only, with fixed seeds, budgets, and 1,000-episode frozen evaluation:

PURSUIT_MAX_RISK=0.16 bash notebooks/experiments/scripts/corrected_final_campaign.sh \
  amd48 <experiment_tag>

Always pass the tag: it defaults to the August campaign's paper_final_corrected_2026_08_v2. PURSUIT_MAX_RISK defaults to the product-Bernoulli value 0.11; a successor-world-model campaign must pass 0.16. ONLY_ENVIRONMENTS="bertrand chicken" submits a subset, so environments whose opponent model is ready can start early. Congestion is run separately through pipeline_env.sh with its own protocol (seed offset 5300, evaluation offset 4,200,000, both roles).

Imperial PBS Pro (CX3)

The full one-environment PBS pipeline is:

bash notebooks/experiments/scripts/pipeline_env_hpc.sh \
  gathering num_runs=3

It submits GPU WM → GPU OM → four CPU algorithm jobs with native qsub afterok dependencies. The defaults are 16 CPUs/64 GB/1 GPU for WM and OM, 8 CPUs/32 GB/no GPU per algorithm, and 71:59:59 walltime. Override them with the OMSH_HPC_WM_*, OMSH_HPC_OM_*, and OMSH_HPC_ALGORITHM_* variables.

Leave the optional PBS queue unset to use Imperial routing, or provide it as the second positional value. The full six-configuration matrix is:

bash notebooks/experiments/scripts/pipeline_all_envs_hpc.sh num_runs=5

It creates 36 PBS jobs.

Single-core OmSh campaigns on CX3

For algorithm-only reruns whose timings are not reported, CX3 runs one single-core job per environment from a workstation:

ssh cx3 "PATIENCE=512 bash -s" < notebooks/experiments/scripts/cx3_campaign_patience.sh

cx3_campaign_patience.sh and cx3_patience_screen.sh encode the campaign seeds, budgets, and evaluation streams. They need inputs already present on CX3: the campaign's wm/, om/, and shield/transition_shield_bundle*.pkl for each environment, copied with modification times preserved so the shield cache's staleness check passes. Three CX3 details matter:

  • CPU-only JAX. The CX3 virtual environment carries the JAX CUDA plugin. CPU jobs must set JAX_PLATFORMS=cpu and CUDA_VISIBLE_DEVICES="", or they stall after setup while memory climbs to the limit. The submitters refuse a runner without these settings.
  • Start-up races. Each job's runner updates the shared checkout and environment. Simultaneous starts race on the Git ref lock, so the submitters stagger starts with qsub -a.
  • Separate storage. CX3 cannot see the Department of Computing volumes. Copy results back (for example with tar over ssh) before analysis.

Existing Artifacts and Visualization

When compatible WM/OM artifacts already exist, env.sh (Slurm) and env_hpc.sh (PBS) submit only algorithm jobs. vis.sh performs CPU visualization reruns with just_vis=True and defaults to amd48.

Use DRY_RUN=1 to inspect commands without submitting. Dry-run output prints arguments with printf %q, so a \, inside player_0\,player_1 is display quoting, not part of the value. PROJECT_PREFIX changes the checkout path below $HOME/Projects; JOB_SCRIPT, SUBMIT_CMD, DEFAULT_VENV, and SUBMIT_DELAY_SECONDS control deployment details. Shared cluster checkouts must serialize their git pull --ff-only and uv sync updates before sibling jobs execute notebooks.

Memory

Without these variables, Slurm jobs receive the partition's default memory (4 GB on the CPU partitions). Set the stage that needs more:

Variable Stage
OMSH_SLURM_WM_MEMORY World model (sized to the graph)
OMSH_SLURM_OM_MEMORY Opponent model
OMSH_SLURM_ALGORITHM_MEMORY Every algorithm job
OMSH_SLURM_OMSH_MEMORY Algorithm jobs that include an OmSh condition

Graph-sized stages need about 64 GB for Markov Stag Hunt, 32 GB for Gathering, 16 GB for Pursuit, and 8 GB otherwise. Markov Stag Hunt's OM stage also needs 64 GB. The ablation experiment notebooks need at least 16 GB for their end-of-run plotting.

A marimo kernel that dies does not end its job. The job can stay RUNNING indefinitely while its afterok dependents wait. Check the log for kernel was stopped unexpectedly rather than trusting the queue state.

Results Archive

Experiment results are copied to shared storage at /vol/bitbucket/oja24/omsh/results/<experiment>-<yyyy-mm-dd>/. Each folder mirrors export-relative paths and has a README.txt recording what ran, where, with which seeds and protocol, and the headline outcome. This path is specific to the maintainer's Department of Computing account; see the results archive note.

See the maintainer launcher contract for bundle composition, job-ID parsing, exact resource variables, and validation rules, and the successor campaign note for a worked end-to-end campaign.