Interpretability and sensitivity checks of the RSPA revision (ESM Sections S11 and S16)

[Ed-Mehrez]

Adds the checks behind two sections of the electronic supplementary material (ESM) of the revised manuscript (RSPA-2026-0421): Section S11 ("Reading the policy off the representation", and the closing sentences of "A Note on Interpretability") and Section S16 ("Accuracy along the learned policy and sensitivity of SKVI").

The first three commits are the S11 checks Preston reviewed and approved. The next three add S16, the third of them answering Preston's review of the S16 commits. The newest (1cebcf5) makes the comments of both modules self-contained, following Ludger's suggestion (see Code comments below).

Contents

• koopmanrl_utils/skvi_policy_checks.py (S11): trains SKVI with the tuned configurations in configurations/ and reads its deployed Gibbs policy off the Koopman tensor.
  • Linear system: learned quadratic form vs the discounted Riccati matrix; mean-action gain vs the discounted LQR gain.
  • Double well: value function pruned (odd terms removed; diagonal quadratic only); returns on paired rollouts and total-variation distance to the full policy; gain of the mean action near the origin; closed-form Gaussian mean vs the grid policy.
  • For S16, train_skvi takes optional arguments (state-dictionary order, action-cost scale, transitions appended to the identification data, and whether they go to the tensor, the value fit or both) and quadratic_cost an optional action-cost scale. The defaults are unchanged, and the trained tensor and value weights are bitwise identical to before (checked on seed 100 for both environments).
• koopmanrl_utils/skvi_sensitivity_checks.py (S16, new):
  • --study accuracy (Table S6, all four benchmarks): one-step error of the tensor identified from random-agent data, along the trained SKVI policy and on fresh random-agent data, with the scale of each dictionary function fixed from the random-agent targets for both distributions. Also the persistence baseline, the error of the predicted state in state units, the exact conditional mean for the double well, the one-step Jacobian of the fitted model at the target against the environment's one-step map, and the distance of the closed loop and of the uncontrolled system from the target.
  • --study sensitivity (Tables S7 and S8, double well and Lorenz): SKVI retrained with other dictionary orders, with the action cost weighted by the time step, and after re-identifying the tensor on transitions of the trained policy, compared with no control and with the LQR controller of the linearisation on the same episodes. Only SKVI is retrained. Two Lorenz variants separate the two effects of re-identification (refit the tensor only, or add the policy's states to the value fit only), and --variants runs a subset with the same numbers as a full run.
  • --summarize_only prints the medians quoted in Tables S6 to S8, including explicit failure counts. Results go to an untracked skvi_sensitivity_checks_results/, with the settings of each run. The summary refuses to mix runs made with different settings, a run never replaces a result file made with other settings, and unknown variant names are rejected.
• Tests: tests/test_skvi_policy_checks.py (S11) and tests/test_skvi_sensitivity_checks.py (S16): the sampler equals the package's pis, the prediction equals K_, the double-well conditional mean checked through the environment's own noise, the error metrics, episode cost sums and the near-target fraction on a deterministic test system, the appended identification data, the local Jacobian on the linear system, the LQR reference, the action-cost scale and both guards. 14 tests, about 4 seconds.
• koopmanrl_utils/AGENTS.md lists both modules; .gitignore ignores both output directories.

Code comments
The comments describe each computation in full, so the code can be read without the paper: the normalised one-step error is written out, the R·dt action-cost scaling is explained, and the comments say why appended transitions are rounded to single precision and why the gain uses the generator state left by the previous evaluation. Where the code implements a theoretical statement, the docstring gives the formula and cites the statement by number and name: action_fields (Proposition S11.1, the policy is an exponential family over the action dictionary) and policy_reading (Corollary S11.2, the Gaussian closed form). Each module docstring ends with a single reference to the paper's supplementary material, by section title. Tables, figures and section numbers are no longer named in the code, and the summary prints "Accuracy" and "Sensitivity" in place of table numbers. Apart from those two printed labels, 1cebcf5 changes no code: with docstrings removed, the syntax trees of both modules are otherwise identical.

Running it

uv run -m koopmanrl_utils.skvi_policy_checks                                # S11: both environments, seeds 100-107
uv run -m koopmanrl_utils.skvi_sensitivity_checks --study accuracy          # S16, Table S6
uv run -m koopmanrl_utils.skvi_sensitivity_checks --study sensitivity       # S16, Tables S7 and S8
uv run -m koopmanrl_utils.skvi_sensitivity_checks --summarize_only

No new dependencies.

Verification. The S16 tables were first produced by exploratory scripts that skvi_sensitivity_checks.py consolidates. A full run of the module on all eight seeds agrees exactly with the saved results of those scripts for every reported number of both studies (1,624 values), and --summarize_only prints the medians of Tables S6 to S8. The re-identification ablation, run on the same eight seeds, gives a median Lorenz cost of 97,000 with the tensor refitted only and 844,000 with the policy states added to the value fit only, against 241,000 for the tuned SKVI and 48,000 for full re-identification. The exception is the model Jacobians in the sensitivity study, which the module evaluates in double precision with a smaller step (differences below 1e-7). An independent review by Codex (GPT, extra-high reasoning) found no error in the metrics, the conditional means, the paired noise, the LQR reference or the cost weighting. Its seven requests on clarity, provenance and tests are addressed in the S16 commits. Ruff, ruff format, isort, vulture and the hygiene hooks pass.

🤖 Generated with Claude Code

[Ed-Mehrez]

@Pdbz199 @ludgerpaehler could you both take a look at the new commit (34c6d01)? It adds koopmanrl_utils/skvi_sensitivity_checks.py with tests, the code behind the new ESM Section S16 (accuracy along the learned policy and sensitivity of SKVI).

• Preston: it follows the structure of the S11 module you approved. The only change to the S11 code is optional arguments to train_skvi and quadratic_cost; with the defaults it gives bitwise-identical tensors and value weights.
• Ludger: it retrains SKVI on all four benchmarks with the tuned configurations in configurations/, so a check of the fluid-flow and Lorenz settings would help.

The updated description has the details and the verification. We won't get this merged before the resubmission, so the data accessibility statement will point to this PR and commit. Thanks!

[Ed-Mehrez]

@Pdbz199 @ludgerpaehler a quick bump in case this got buried. There are two new commits since Preston's approval: 34c6d01 adds the S16 checks (koopmanrl_utils/skvi_sensitivity_checks.py and its tests), and 12a9edf makes the failure criterion in its summary explicit.

A full run of the module on all eight seeds reproduces every number in the S16 tables exactly. The manuscript's data accessibility statement now cites this PR at 12a9edf, so it would help to hear soon if anything needs changing. A short look is enough; the description lists what to check.

[Ed-Mehrez]

@ludgerpaehler following your point about paper references in code comments, 1cebcf5 rewrites the comments of both modules so that each one describes its computation in full:

• the normalised one-step error is written out, with its target and the two choices of scale, instead of pointing at an equation of the paper;
• the R·dt action-cost scaling, the single-precision rounding of appended transitions and the generator order in gain_near_target are explained where they occur, instead of referring to the runs reported in the supplementary material;
• functions, arguments and constants describe their study rather than naming a table, and the summary prints "Accuracy" and "Sensitivity" instead of table numbers.

Following the middle path you suggested, the two places where the code implements a theoretical statement cite it next to the formula, by number and name: action_fields (Proposition S11.1, the policy is an exponential family over the action dictionary) and policy_reading (Corollary S11.2, the Gaussian closed form). Each module docstring ends with one reference to the supplementary material, by section title.

No computation changed. With docstrings removed, the syntax trees are identical apart from the two printed labels, and the 14 tests and the pre-commit hooks pass. Once you have tagged the release with a Zenodo DOI, we can cite that in the data accessibility statement instead of a commit hash.