Skip to content

ctde_v0.probe_message_ablation

probe_message_ablation

STANDALONE CPU probe — does the LPAC GNN backbone actually USE neighbour messages?

No training. Loads a run's config.json + model.eqx (exactly like render.py), drives ONE rollout per seed with the base actor, samples states every --sample-stride steps, and — at each sampled state — ablates the comm graph / a neighbour's belief and measures how much the per-agent belief z and the goal distribution move. The 5th return value of actor(obs, adj, dist=dist, h=h, inference=True) is feat = z (nets.py:874); that is the vector EVERY head reads, so it is what we watch.

Three modes (--mode):

  • graphoff — set adj all-False (every agent isolated) and recompute. Report per-agent rel_dz = ‖z1_i − z0_i‖ / ‖z0_i‖ and the goal-argmax-flip fraction. THE decisive number: mean rel_dz ≈ 0 ⇒ the GNN IGNORES neighbour messages (the belief is a pure per-agent CNN readout; message passing is dead weight).

  • edge — for every PRESENT directed edge i←j remove just that one edge (adj.at[i,j].set(False)) and measure the receiver i's rel_dz, goal-KL and argmax-flip → a per-edge influence heatmap (who listens to whom, and how hard).

  • perturb — inject Gaussian(0, sigma) into ONE agent's belief channels (known + *frontier) and measure Δz / Δgoal at its in-range neighbours → does a corrupted local belief propagate through the graph to the agents that hear it.

Outputs (into --out as a directory): per-mode heatmap PNG(s), histogram PNG(s), and summary.json with an explicit verdict string.

JAX_PLATFORMS=cpu PYTHONPATH=.:../../../FiedlerValueEstimation     /Users/bijanmehr/Project.Zymera/zymera_lab/.venv/bin/python -m ctde_v0.probe_message_ablation         --run-dir runs/warmab/warm_s0 --out probes/msg_warm_s0 --mode graphoff         --seeds 0,1,2 --sample-stride 10