iKCE Diagnostic Experiment

Empirical anchor for the RSS-W 2026 workshop paper "Imagined Rollouts Are Kinematic, Not Dynamic: A Diagnosis of Long-Horizon World-Model Failure." Computes the imagined-rollout kinematic-consistency error (iKCE) and the conditioning-perturbation dose-response curve on DreamerV3, plus the post-hoc statistical analyses (flatness test, horizon-emergence test) and TikZ figure pipeline that ship with the paper.

Interpretation matters. Low iKCE is not good. Low iKCE means the world model is imagining kinematically — the headline finding the paper is arguing for. The signature of dynamic imagination is iKCE > 0 that grows systematically and correlates with physical drivers. See src/ikce_diagnostic/metrics/ikce.py and docs/INTERPRETATION.md.

Quick start

1. Create the conda environment

The DMC stack (mujoco, dm_control) and DreamerV3 port together pull in several heavy native dependencies. Use a dedicated conda env:

conda env create -f environment.yml
conda activate ikce-diagnostic

2. Install PyTorch matching your GPU

pip install torch                                                # CPU only
pip install torch --index-url https://download.pytorch.org/whl/cu121   # RTX 30/40-series
pip install torch --index-url https://download.pytorch.org/whl/cu128   # RTX 50-series (Blackwell)

3. Install the project + the extras you need

pip install -e ".[dev]"                # core + tests, no MuJoCo
pip install -e ".[dev,dmc]"            # adds dm_control + MuJoCo
pip install 'dreamerv3-torch @ git+https://github.com/NM512/dreamerv3-torch.git@6ef8646'
                                       # the WM checkpoint host (see docs/CHECKPOINTS.md)

4. Fetch the pre-trained checkpoints

The five trained DreamerV3 walker-walk checkpoints (~1.1 GB total) are hosted on the Hugging Face Hub rather than in git, and fetched into checkpoints/ (SHA-256-verified against the manifest in docs/CHECKPOINTS.md):

hf auth login                     # one-time; needs a (free) Hugging Face account
bash scripts/fetch_artifacts.sh   # → checkpoints/<run>/latest.pt (+ config, metrics)

This populates walker_walk (default imag_horizon=15), walker_walk_h64 (actor-horizon ablation), walker_walk_dr (domain-randomized friction), and walker_walk_seed{1,2} (seed replicates) — the checkpoints the sweep configs expect. To retrain from scratch instead, follow the per-checkpoint commands in docs/CHECKPOINTS.md (~14 h per checkpoint on an RTX-class GPU).

5. Run the experiments

# Unit + integration tests
pytest

# Sanity check on synthetic physics (no MuJoCo needed):
python scripts/sanity_check_ikce.py --backend synthetic --horizon 64

# Sanity check on the real walker (requires the [dmc] extra):
python scripts/sanity_check_ikce.py --backend dmc --domain walker --task walk

# Empirically determine the regime boundary μ (writes
# results/regime_boundary_friction.csv; the headline figure's dashed
# line is anchored here).
python scripts/determine_regime_boundary.py configs/walker_walk_physics_policy.yaml

# Full perturbation sweep (one per run config; ~2 h total on RTX 5090
# for all eight runs: physics/WM × identity/gait × h=15/h=64 actor):
python scripts/run_perturbation_sweep.py configs/walker_walk_wm_h64.yaml --deterministic
# ...repeat for the other run configs.

# Post-hoc statistical analyses (no rollouts; reads saved per-step
# iKCE from the sweep above).
python scripts/flatness_test.py
python scripts/horizon_sweep_analysis.py

# Refresh tex/data/ from results/ so the PGFPlots figures pick up the
# latest numbers.
python scripts/prepare_figure_data.py

# Local figure-compile harness (mimics the IEEEtran double-column
# layout of the workshop paper).
cd tex && pdflatex test_double_column.tex

Repository layout

ikce-diagnostic-experiment/
├── README.md
├── pyproject.toml / environment.yml
├── src/ikce_diagnostic/
│   ├── types.py                   ← RolloutOutput, KinematicSpec
│   ├── config.py                  ← YAML → SweepConfig
│   ├── sweep.py                   ← sweep runner
│   ├── metrics/
│   │   ├── ikce.py                ← iKCE definition (paper Eq. 1)
│   │   └── extrapolation.py       ← constant_velocity, constant_acceleration
│   ├── rollouts/
│   │   ├── base.py                ← RolloutBackend protocol
│   │   ├── synthetic.py           ← analytical physics
│   │   ├── dmc_physics.py         ← MuJoCo ground truth
│   │   ├── dreamer.py             ← legacy DreamerV3 wrapper (kept for reference)
│   │   └── dreamer_nm512.py       ← active adapter for the NM512 DreamerV3 port
│   ├── perturbations/
│   │   ├── base.py                ← PerturbationSpec, registry
│   │   ├── dmc.py                 ← friction / gravity / init_velocity / joint_noise
│   │   └── synthetic.py           ← test/CI counterpart
│   └── analysis/
│       └── aggregate.py           ← bootstrap CIs, CSV I/O
├── scripts/
│   ├── test_dreamer_rollout.py    ← end-to-end smoke test
│   ├── sanity_check_ikce.py       ← metric-implementation sanity
│   ├── run_perturbation_sweep.py  ← sweep driver (one YAML config per run)
│   ├── determine_regime_boundary.py ← empirical μ=0.20 procedure
│   ├── flatness_test.py           ← H2 falsifiable slope test (Appendix)
│   ├── horizon_sweep_analysis.py  ← horizon-emergence reanalysis (Appendix)
│   ├── dump_perstep_csv.py        ← per-step CSVs for Fig. 4 / Fig. 5
│   ├── prepare_figure_data.py     ← refreshes tex/data/ from results/
│   └── make_headline_figure.py    ← legacy matplotlib export; superseded by PGFPlots
├── configs/                       ← one YAML per run (physics/WM × view × actor-h)
├── tex/
│   ├── styles.tex                 ← shared PGFPlots styling
│   ├── headline_template.tex      ← parameterised macro for headline-style figures
│   ├── headline_identity_h64.tex  ← Fig. 1 source
│   ├── regime_curve.tex           ← Fig. 2 source
│   ├── headline_identity_a64.tex  ← Fig. 3 source (actor-horizon ablation)
│   ├── perstep_compare_h64.tex    ← Fig. 4 source (per-step decomposition)
│   ├── perstep_actor_ablation.tex ← Fig. 5 source
│   ├── headline_gait_h64.tex      ← Fig. 6 source (gait DOFs view)
│   ├── horizon_sweep.tex          ← Fig. 7 source (horizon-emergence test)
│   ├── test_double_column.tex     ← local compile harness (IEEEtran double-column)
│   └── data/                      ← committed CSVs the figure sources consume
├── tests/                         ← pytest suite
├── docs/
│   ├── CHECKPOINTS.md             ← trained WM provenance (h=15 default + h=64 ablation)
│   ├── INITIAL_iKCE.md            ← experimental design + measured numbers
│   ├── KINEMATIC_STATE.md         ← observation-slicing convention
│   ├── KINEMATIC_EXTRAPOLATION.md ← null-model spec
│   ├── PERTURBATIONS.md           ← perturbation-axis design
│   ├── PHYSICAL_REGIMES.md        ← what "regime boundary" means here
│   ├── PHASE_4_PROTOCOL.md        ← sweep-protocol notes
│   ├── INTERPRETATION.md          ← reads-before-using guide
│   └── CAVEATS.md                 ← things a reviewer might reasonably ask
└── results/                       ← gitignored, generated by sweep + analyses

What runs what

Artefact in the paper	Generated by
Table I, Fig. 1, Fig. 2 (main paper)	run_perturbation_sweep.py + determine_regime_boundary.py + prepare_figure_data.py
Fig. 3 (actor-horizon ablation)	sweep using walker_walk_*_a64.yaml configs against the walker_walk_h64 checkpoint
Fig. 4 (per-step decomposition)	dump_perstep_csv.py
Fig. 5 (per-step actor ablation)	dump_perstep_csv.py over the h=15 + h=64 runs
Fig. 6 (gait DOFs view)	sweep using walker_walk_*_gait*.yaml configs
Fig. 7 (horizon-emergence test)	horizon_sweep_analysis.py
Appendix flatness regression	flatness_test.py
Regime-boundary plot	determine_regime_boundary.py

All figure sources in tex/*.tex read from tex/data/*.csv via \datadir/tex/data/...; \datadir is set in the paper preamble (use . for an Overleaf project with the paper at root + tex/ as a subdirectory). prepare_figure_data.py is the bridge that keeps tex/data/ in sync with the latest sweep outputs in results/.

Non-negotiable constraints

• iKCE follows the definition in the paper (§III-A) exactly. The implementation in src/ikce_diagnostic/metrics/ikce.py is the authoritative version. Any change to it must be discussed first.
• Kinematic extrapolation is chosen once per experimental run and held constant. SweepConfig.extrapolation is the single source of truth.
• All rollouts and sweeps are reproducible from one YAML config. SweepRunner writes the config snapshot next to its outputs so reviewers can rerun without consulting configs/.
• Full sweeps are reported, not curated subsets. The sweep runner refuses to silently drop any cell; missing files are logged as warnings and excluded from aggregation with explicit row counts in the CSV's n column.
• Sanity check must pass before any sweep result is trusted. The unit tests partially encode this via tests/test_ikce.py::test_ikce_zero_on_ballistic_under_constant_acceleration and tests/test_ikce.py::test_ikce_spikes_on_bouncing_ball; the production check still requires running scripts/sanity_check_ikce.py --backend dmc.

Caveats

Anything that surprised us during the experiments — perturbation timing, motion-magnitude dilution at long horizons, the "peak vs elevation" framing of the physics-side response — lives in docs/CAVEATS.md. Read it before quoting numbers from this repo.