Skip to content

Pursuit 3-Agent Opponent-Model Memory

This note records a resolved memory incident in the older deterministic graph path. It does not resolve the current stochastic graph limit: with 0.1 slip, the August 2026 build enumerated 415,886,464 successor edges and was OOM-killed at 205 GiB before atomic publication. The default graph builder now rejects that materialization using a fanout-aware preflight; compact or procedural transition storage is required before three-agent OM training is viable.

On 2026-04-24, Pursuit train_om_mo.py runs with n_agents=3 failed for jobs 234679 and 234840. The generated marimo notebooks had no useful traceback because the Python kernel was killed externally. The backing scheduler logs were /homes/oja24/logs/job_234679.out and /homes/oja24/logs/job_234840.out; both reported:

The Python kernel ... died unexpectedly.
Detected 1 oom_kill event ... Some of the step tasks have been OOM Killed.

At the time of the incident, the 3-agent export had no saved OM artifacts under exports/gridworlds/pursuit/3/om, so those jobs had not reached the final save cell. That observation describes the failed run, not a current artifact contract; later exports may contain a completed iop_stack.pt.

The main memory pressure in that older deterministic build came from prepare_om_training_inputs and the graph adapter path:

  • exports/gridworlds/pursuit/3/wm/D.pkl has 1,000,000 transitions and loads to about 3.25 GiB RSS in a local probe.
  • exports/gridworlds/pursuit/3/wm/env_transition_graph.pkl has 85,184 nodes and 10,648,000 edges. Loading it alone reached about 15.9 GiB RSS.
  • Constructing GraphWorldModelAdapter over that graph climbed past 24 GiB RSS before the probe was stopped.

Because train_om_mo.py loads the replay buffer, materializes OM state arrays, loads the full env graph, builds a validation adapter inside prepare_om_training_inputs, then later builds another adapter for rollout PI, the peak memory can exceed common 24-32 GiB job allocations before the policy improvement phase starts.

Considered fixes

  • Avoid constructing GraphWorldModelAdapter twice. Return the validation adapter from prepare_om_training_inputs, or split graph validation into a lightweight state lookup path.
  • Release the raw replay buffer and temporary states_np as soon as OMBuffer and graph indices are built.
  • Consider loading/building an indexed or compact graph representation for OM rollouts instead of keeping the dataclass edge list plus duplicated adapter indices in memory.
  • If no code change is made yet, request a substantially larger memory allocation for 3-agent Pursuit OM jobs.

Resolution implemented

  • prepare_om_training_inputs now returns the validated GraphWorldModelAdapter as prepared.dyn, while keeping prepared.graph for existing callers.
  • notebooks/gridworlds/pursuit/train_om_mo.py uses prepared.dyn instead of constructing a second adapter.
  • Large temporary replay/state/action arrays are released once the OMBuffer and graph-state indices have been built.
  • GraphWorldModelAdapter now bounds its legal-mask cache with an LRU cap. This only changes memoization lifetime; legal masks are deterministic and recomputed exactly after eviction.
  • The same prepared.dyn reuse pattern was propagated to the other train_om_mo.py notebooks so they avoid the duplicate adapter construction too.
  • The patch does not move OMBuffer storage to CPU; notebook runs still pass the selected device into prepare_om_training_inputs. The memory win here comes from removing duplicate graph-adapter construction and releasing temporary CPU arrays earlier.
  • The default legal-mask cache cap is intentionally conservative at 100k entries. Raw masks are small, but Python tuple, NumPy array, and OrderedDict overhead dominates; raising the cap to 1M could plausibly cost hundreds of MiB in a path already failing from RAM pressure. The cap is exactness-preserving because evicted masks are recomputed deterministically.