A Decision-Margin Theorem for When Learned-Model Planning Stays Correct
Yueyi Li — Washington University in St. Louis
This repository contains the full experimental code and data for the paper. A model-based agent that scores actions by imagined reward alone is blind to commitments that permanently destroy reachable value (stepping into lava, sealing a corridor, spending a non-renewable resource). We price such a commitment at decision time with a destroyed reachable mass read directly from a world model's own predicted reachability, prove a decision-margin theorem (with a localization lemma) characterizing exactly when planning on a learned model still makes the same decision as planning on the exact model, and give a reachability-consistency training objective that repairs the failure case the theorem identifies. All claims are validated in controlled, deterministic gridworlds — three custom-built environments plus a native Farama MiniGrid environment — using a matched reversible/irreversible twin design for causal identification.
pip install -r requirements.txt
Python 3.11. See requirements.txt for exact pinned versions. Everything in
this repository runs on CPU; no GPU is required or used.
core/ shared exact-model primitives (the single destroyed_mass
definition, exact planners, MDP class) -- everything else
imports from here, never redefines it
environments/ the four environments and their reversible/irreversible
twin constructors
worldmodels/ the learned world model (base) and the decision-aware
variant (adds the reachability-consistency training
objective + the no-oracle access counter, CountedMDP)
experiments/ the scripts that reproduce every result in the paper
verification/ standalone scripts for the paper's specific verification
claims (currently: Assumption 1's exhaustive check)
figures/ scripts that turn data/results/*.json into the paper's
figures, plus the exact CSV data behind each one
data/results/ JSON (and a diagnostic PDF each) written by the experiments/
scripts -- already included, so figures can be regenerated
without re-running anything
run.sh one-click reproduction of the whole suite
Later directories import primitives from earlier ones — in particular,
destroyed_mass, the exact planners, and the MDP class are defined once,
in core/exact_mdp.py; every other module imports it and asserts the import
resolved to that exact module (see experiments/common_eval.py), so a
regression there cannot be masked by a silent redefinition elsewhere.
| module | environment | irreversibility mechanism |
|---|---|---|
environments/lava_corridor.py |
abstract lava corridor (Section 4's exact-model construction) | absorbing state |
environments/doorkey_lava.py |
DoorKey-Lava | absorbing state (lava) |
environments/sokoban_barrier.py |
Sokoban-barrier | environment-state (box seals a corridor chokepoint) |
environments/resource_depletion.py |
Resource-depletion | monotone (non-renewable resource) |
environments/minigrid_lavagap.py |
native MiniGrid lava-gap | absorbing state, on the real Farama MiniGrid engine |
Each environment builds a matched twin: identical layout, reward, and
horizon, differing only in whether the designated commitment is reversible.
minigrid_lavagap.py is the only one built on a real simulator (Farama
minigrid 3.1.0) rather than an enumerable custom grid; its exact ground-truth
reachability is obtained by offline breadth-first enumeration, stepping
the real environment, never by the test-time planner.
Figure 1 (the schematic of destroyed reachable mass and the localization argument) is a hand-drawn diagram, not a data-driven plot; there is no script for it in this repository.
For Figures 2–6, the plotting scripts and the exact data they consume are both included, so they can be regenerated without re-running any experiment:
python figures/make_figures.py
reads data/results/*.json and writes:
| output | paper figure |
|---|---|
figures/fig_lambda_threshold.pdf |
Figure 2 — recovery threshold (Theorem 1) |
figures/fig_margin_crossing.pdf |
Figure 3 — decision-flip boundary (Theorem 2) |
figures/fig_error_decomposition.pdf |
Figure 4 — score-error decomposition |
figures/fig_localization.pdf |
Figure 5 — localization of the cost-gap error (Lemma 1) |
figures/fig_repair.pdf |
Figure 6 — decision-aware repair (Proposition 1) |
Runtime: a few seconds.
Table 1 (separation and collapse across the four environments) is read
directly from the mechanism_per_seed entries of
data/results/embodied_family_results.json and
data/results/minigrid_results.json — no separate script reproduces it.
figures/export_data.py and figures/export_margin_multienv.py re-derive the
CSVs under figures/data/ from the same JSON files; their output is already
included, so running them is optional (a few seconds each) and only useful to
verify the derivation independently.
Each of these re-runs the world-model training and closed-loop evaluation
that produced the corresponding data/results/*.json (already included, so
this is optional and only needed to verify the numbers independently).
Runtimes are approximate, single CPU core, no GPU.
python experiments/run_embodied_family.py # Env1/2/3, a few minutes
python experiments/run_minigrid.py # Env4 (native MiniGrid), ~3 minutes
python experiments/run_localization.py # Figure 5 source, ~3 minutes
python experiments/run_repair.py # Figure 6 source, ~3 minutes
Each writes its own data/results/*.json and a diagnostic PDF; a seed-by-seed
diff against the included files is the most direct way to check
reproduction. bash run.sh runs all four in sequence, then verification and
figure regeneration, with a final pass/fail summary
(MRC_QUICK=1 bash run.sh skips straight to verification + figures using the
results already in the repo, ~1 minute).
Assumption 1 (reachability contraction), exhaustive check. The paper states this is verified by exhaustive enumeration on the Sokoban-barrier environment (the least obvious of the four, since pushing the box changes environment state rather than agent state). This is a pure, deterministic graph check: build the full transition graph — including the box's own state variable, not just agent position — and confirm the successor's reachable set is a subset of the predecessor's, for every state and action, on both twins.
python verification/verify_sokoban_monotonicity.py
Runs in under a second; prints a verdict and any counterexample found, and
writes verification/verify_sokoban_monotonicity_results.json.
No-oracle verification at decision time. CountedMDP
(worldmodels/decision_aware.py) wraps the true environment in a counter
that increments on any access to its transition or reward function. Every
closed-loop rollout in experiments/ is preceded by a probe that asserts the
counter registers zero true-environment accesses while the planner makes
its decision — the true dynamics are consulted only by the offline
enumeration that builds ground-truth reachability, never by the test-time
planner. experiments/run_minigrid.py runs this check explicitly and prints
its result (verify_cheat_free) separately from the main experiment output;
experiments/common_eval.py's run_closed_loop enforces the same assertion
for Env1/2/3.
Affine margin cross-validation (dense decision-flip boundary evidence). For a fixed seed and decision state, the exact and learned combined scores are both affine in the planning weight λ (the reward and reachability terms do not themselves depend on λ), so the decision-flip boundary can be evaluated on an arbitrarily fine λ grid by exact interpolation rather than by re-running training at each grid point.
python figures/export_margin_multienv.py
does this for all four environments and, critically, cross-validates the interpolated values against every λ point that was actually measured in the original (coarser) run — the reported max deviation is at floating-point noise level for all four environments, which is what licenses treating the fine grid as exact. It prints, per environment: number of decisions evaluated, number of sign flips, number of margin-preservation violations (expected: zero), and the cross-validation error.
@article{li2026pricing,
title = {Pricing Irreversible Commitments in World Models: A
Decision-Margin Theorem for When Learned-Model Planning
Stays Correct},
author = {Li, Yueyi},
journal = {Journal of Artificial Intelligence Research},
year = {2026},
note = {To appear. Update volume/pages/DOI once assigned.}
}See also CITATION.cff.
MIT — see LICENSE.