rwm-reproduction

Model card — RWM reproduction checkpoints

Independent reproduction of the proprioceptive dynamics model of Li, Krause & Hutter, Robotic World Model (arXiv:2501.10100) and Uncertainty-Aware RWM (arXiv:2504.16680). Not affiliated with, endorsed by, or reviewed by the original authors.

Code, evidence and the full claim record: https://github.com/joyjeet-singh/rwm

Paper: Right Ranking, Wrong Scale: Reproducing the Robotic World Model's Dynamics Model and Testing the Uncertainty of Its Follow-up, arXiv:2611.02699.

Read this before using the σ output

These models' predicted standard deviation is not a usable uncertainty estimate. It is not a matter of degree. Measured against realised error on held-out episodes:

Those are the ALEATORIC term, at one forecast step. The quantity the follow-up's method actually penalises rewards with is the EPISTEMIC one, and on the released 5-member checkpoint it is 33.4× [28.7, 39.0] out at h = 100 — the horizon the method's own imagination rollouts run to — with 4.61% coverage at ±1σ. Our ensemble-5 arms reach 10.5× [9.0, 11.5] on the same measurement. Better, and not calibrated.

The cause is structural, not a training accident: the state loss is squared error on a reparameterised sample with no log-σ term, so its optimum is σ = 0, and the bound term that should oppose it cancels algebraically. The corrected-objective-2500 checkpoint uses the reference's unused gaussian_nll branch, which reverses the mechanism and still does not produce a usable estimate (10.9× overconfident). It is released as the corrected-objective artifact, not as a calibrated one.

If you need to rank which predictions will be worse, the σ output carries some signal — but treat that as directional, not established. Converting per-dimension sign counts to P-values against an independent-trials null overstates the evidence badly: the 45 state dimensions are physically coupled and share a forecast-depth trend, so a σ that grows with depth correlates with any trajectory's error. Under a permutation test over whole trajectories, no such count in our paper survives multiplicity correction.

If you need an interval, the raw σ will not give you one. Do not use it for risk-gating, safety margins, or anything that treats σ as a scale.

There is a candidate remedy, and on these checkpoints the evidence is mixed. A single multiplier does not work, because the miscalibration grows with forecast horizon. One multiplier per horizon, fitted on one episode and scored on the other, restored ±1σ coverage on the released reference checkpoint to within 10 points of nominal on 12 of 12 epistemic cells, though at two episodes no single cell is resolvable, where a single global multiplier managed 2; but that checkpoint trained on both episodes. On these ensemble-5 checkpoints, which never saw them, their own multipliers manage 17 of 36 epistemic cells, on two episodes only. Refit it on your own data and check it rather than copying our constants.

The ensemble-5 checkpoints, and what they are for

Ten of the thirteen checkpoints here run at ensemble size 1, where the epistemic term -- the disagreement across ensemble members -- is identically zero by construction. Nothing about ensemble disagreement can be reproduced from one of them alone; the block above says what to do with five of them.

The three autoregressive-ens5-* checkpoints carry it directly. On them:

• ensemble disagreement correlates with realised error in 12 of 12 seed-horizon cells more strongly than the forecast step index does, which is the property our paper argues makes it a usable ranking signal;
• the epistemic term is 10.5x smaller than the realised error at h=100 -- the horizon the method's own imagination rollouts run to -- with 8.19% coverage at +-1 sigma against a calibrated 68.27%, and 13.0x with 6.30% at the 368-step open-loop diagnostic horizon. Better than the released reference checkpoint, and still not an interval;
• it is input-dependent, CoV 0.379-0.395 across a batch.

The pre-registered rule governing this replication returned DOES NOT GENERALISE, on its second condition: the paired difference against the index excludes zero at 1 of 4 horizons, not a majority. The direction replicated everywhere; the separation did not, on the four independent trajectories our held-out arena has. Treat the ranking property as supported on these checkpoints and established only on the released reference one.

Scoring five of these together: the independent-ensemble result

If you download the ensemble-size-1 autoregressive checkpoints, score them as an ensemble. The paper's §5.3 takes seeds 0-4 of Arm A -- five separately initialised, separately trained models, sharing no parameters and no recurrent state -- and scores their disagreement the way the method scores its own. That is the contrast the released five-member checkpoint cannot provide, because its five members share one GRU trunk and one hidden state: 89.15% of each member's state-pathway parameters are numerically identical to every other member's.

• The independent ensemble is 2.03x better calibrated than the shared-trunk autoregressive-ens5-* arms at h=100, against a pre-registered minimum detectable effect of 1.45x. Coverage is +7.12 points higher, against an MDE of 2.26 points. The rule (M-44, committed before the runs) returns MECHANISM SUPPORTED.
• It is still not an interval. 5.2x overconfident at h=100 with 15.31% coverage at +-1 sigma, against a calibrated 68.27%. Building the ensemble properly is worth doing and is not sufficient.
• What the gain is made of. sigma larger by 1.65x -- 71% of the improvement at h=100 -- and the rest ordinary ensembling accuracy. At the 368-step diagnostic horizon the split reverses: 57% of it is the ensemble simply predicting better.
• What it does not isolate. Five independent models differ from five shared heads in initialisation, in data ordering AND in capacity: 3,570,820 state-pathway parameters against 1,024,132, a factor of 3.49. The comparison bounds the trunk-sharing effect rather than isolating it.

Checkpoints

autoregressive-10k-seed0

Autoregressive training — the arm the base paper's claim is about. Seed 0 of three at 10,000 iterations; scores 0.3894 normalised error at a 368-step horizon on held-out episodes (arm mean 0.3582 ± 0.0283 over three seeds).

• source: runs/armA_seed0_10k/weights_10000.pt
• size: 5,683,412 bytes
• sha256: 4dbb6871f69e6b7a3e7014fb0ae7effb8366860f6bd0c84af717c10ea64c6625
• σ calibration, measured at iteration 2,500 (this arm; not re-measured at 10,000): 52× overconfident, coverage 11.67% at ±1σ (h=1)

autoregressive-10k-seed1

Autoregressive training — the arm the base paper's claim is about. Seed 1 of three at 10,000 iterations; scores 0.3509 normalised error at a 368-step horizon on held-out episodes (arm mean 0.3582 ± 0.0283 over three seeds).

• source: runs/armA_seed1_10k/weights_10000.pt
• size: 5,683,412 bytes
• sha256: c88b20a136c7364475e1bf74b3919b33de1ebad9293de521cfe259ed8474ee6e
• σ calibration, measured at iteration 2,500 (this arm; not re-measured at 10,000): 52× overconfident, coverage 11.67% at ±1σ (h=1)

autoregressive-10k-seed2

Autoregressive training — the arm the base paper's claim is about. Seed 2 of three at 10,000 iterations; scores 0.3341 normalised error at a 368-step horizon on held-out episodes (arm mean 0.3582 ± 0.0283 over three seeds).

• source: runs/armA_seed2_10k/weights_10000.pt
• size: 5,683,412 bytes
• sha256: d8e3f43edf8a6f76c25e34d814a7a7c4c35779bbb7ec90e76c5f49a9ec2a192c
• σ calibration, measured at iteration 2,500 (this arm; not re-measured at 10,000): 52× overconfident, coverage 11.67% at ±1σ (h=1)

teacher-forced-10k-seed0

Teacher forcing — the comparison arm. Released so the central claim can be checked rather than taken on trust. Seed 0 of three at 10,000 iterations; scores 1.9710 normalised error at a 368-step horizon on held-out episodes (arm mean 1.6497 ± 0.2858 over three seeds).

• source: runs/armB_seed0_10k/weights_10000.pt
• size: 5,683,412 bytes
• sha256: a8e20cf9ef3f0ed0380ba031559fdd497f9e2a0a3a681ce0de39820b4e66f0c1
• σ calibration, measured at iteration 2,500 (this arm; not re-measured at 10,000): 315× overconfident, coverage 12.96% at ±1σ (h=1)

teacher-forced-10k-seed1

Teacher forcing — the comparison arm. Released so the central claim can be checked rather than taken on trust. Seed 1 of three at 10,000 iterations; scores 1.5540 normalised error at a 368-step horizon on held-out episodes (arm mean 1.6497 ± 0.2858 over three seeds).

• source: runs/armB_seed1_10k/weights_10000.pt
• size: 5,683,412 bytes
• sha256: b473bd70d5afdcbb1065e7ef49e32e85d4941719ec1bf290e4a91a09b8165754
• σ calibration, measured at iteration 2,500 (this arm; not re-measured at 10,000): 315× overconfident, coverage 12.96% at ±1σ (h=1)

teacher-forced-10k-seed2

Teacher forcing — the comparison arm. Released so the central claim can be checked rather than taken on trust. Seed 2 of three at 10,000 iterations; scores 1.4241 normalised error at a 368-step horizon on held-out episodes (arm mean 1.6497 ± 0.2858 over three seeds).

• source: runs/armB_seed2_10k/weights_10000.pt
• size: 5,683,412 bytes
• sha256: 2431bc1ba09229d305d364abb53d422ef79d39a7299f8c3692f1fb743c8240a2
• σ calibration, measured at iteration 2,500 (this arm; not re-measured at 10,000): 315× overconfident, coverage 12.96% at ±1σ (h=1)

autoregressive-2500

Arm A at the paper's stated iteration count, for comparison with the released checkpoint.

• source: runs/armA_seed0/weights_2500.pt
• size: 5,683,374 bytes
• sha256: bc4e0dfbcff28b994e3acabbae2d6f1d833331dc5bfad5cc9aa8f410385bd4b0
• σ calibration, measured at iteration 2,500: 52× overconfident, coverage 11.67% at ±1σ (h=1)

corrected-objective-2500

Trained with the reference's unused gaussian_nll branch. This is the CORRECTED-OBJECTIVE artifact, not a calibrated one — see the limitation above.

• source: runs/armA_seed0_nll/weights_2500.pt
• size: 5,683,374 bytes
• sha256: a1a339b27b077712b0a87b4df4718437933db3e6a370b8b08f941fab30129367
• σ calibration, measured at iteration 2,500: 11× overconfident, coverage 42.78% at ±1σ (h=1)

autoregressive-ens5-seed0

Ensemble size 5. The only checkpoints here with a non-zero epistemic term; the others have exactly zero by construction. See the ensemble-5 block below.

• source: runs/armA_seed0_ens5/weights_2500.pt
• size: 8,020,484 bytes
• sha256: 923e3a0b935cff3cb611cf6b35eea83701df4a78f02a22e3f10a6a2e945b9496
• σ calibration (aleatoric), measured on THIS arm at iteration 2,500: 7.7× overconfident with 12.22% coverage at ±1σ (h=1), and 42.4× with 2.44% at h=100
• σ calibration (epistemic — the quantity the method penalises with): 2.0× at h=1 and 10.3× at h=100, coverage 8.62% at ±1σ

autoregressive-ens5-seed1

Ensemble size 5, seed 1.

• source: runs/armA_seed1_ens5/weights_2500.pt
• size: 8,020,484 bytes
• sha256: e756ca16fe2cd4e1615ca27b47d9bb519dd0d4da6dfe263758b1d5bbdaabeaaa
• σ calibration (aleatoric), measured on THIS arm at iteration 2,500: 9.7× overconfident with 11.11% coverage at ±1σ (h=1), and 43.6× with 2.12% at h=100
• σ calibration (epistemic — the quantity the method penalises with): 2.2× at h=1 and 10.7× at h=100, coverage 7.98% at ±1σ

autoregressive-ens5-seed2

Ensemble size 5, seed 2.

• source: runs/armA_seed2_ens5/weights_2500.pt
• size: 8,020,484 bytes
• sha256: 2688b16030a57b4b4201bc844a9e1b3607df8734e7b62b1f7b7efb11ed5d9652
• σ calibration (aleatoric), measured on THIS arm at iteration 2,500: 8.5× overconfident with 8.33% coverage at ±1σ (h=1), and 39.8× with 2.51% at h=100
• σ calibration (epistemic — the quantity the method penalises with): 2.3× at h=1 and 10.5× at h=100, coverage 7.98% at ±1σ

autoregressive-ens1-seed3

Arm A at ensemble size 1, seed 3. Trained for the independent-ensemble test: seeds 0-4 of this arm are scored together as a five-model ensemble that shares nothing (§5.3). Not part of the three-seed headline. No per-arm σ calibration is quoted below: task1_calibration.py measures seed 0 only, and pasting seed 0's figure onto this checkpoint is exactly the mis-attribution this card was corrected for elsewhere. What was measured on this checkpoint is its contribution to the ensemble above.

• source: runs/armA_seed3/weights_2500.pt
• size: 5,683,374 bytes
• sha256: e44e15f5fc31556c7d16647c96edf91bb0d5a4b8b8ad72078ad841a3e60ad260

autoregressive-ens1-seed4

Arm A at ensemble size 1, seed 4. The second of the two seeds added for the independent-ensemble test, and the same caveat applies: no individually measured σ calibration exists for it, only its contribution to the five-model ensemble above.

• source: runs/armA_seed4/weights_2500.pt
• size: 5,683,374 bytes
• sha256: b5e0fa1123cb4ae935a3b925fb0801e765cfd3ac25f8f1cf0b653582eb69d24a

The result these support

Relative-L1 error, the reference's own metric, at a 368-step horizon on held-out episodes, over three training seeds (standard deviation with ddof=1). That is the horizon the paper's pre-registered rule names; at h=100, the method's own rollout length, the same comparison gives 2.58×.

Autoregressive training is better by a factor of 4.61×. For reference the hold-last floor — predicting that nothing changes — scores 0.9930 in the same cell, so teacher forcing is worse than making no prediction at all.

Every 10,000-iteration checkpoint was cross-checked against the 2,500-iteration run at the same seed: 90,000 logged values compared, 0 differing.

What these are

• Architecture. GRU trunk, ensemble size 1, mean head plus bounded log-σ head, auxiliary contact and termination heads. Rebuilt from scratch, verified against the reference at 0.000e+00 on losses and gradients across 7 terms and 106 parameter tensors before training.
• Data. The released ANYmal D dataset: 10,000 rows at 50 Hz, ten 20-second episodes. Trained on 7,687 episode-respecting windows from eight episodes; two held out.
• Action convention. Row t holds the action that produced state t; these checkpoints are trained and evaluated with that pairing. Do not score them with the reference evaluation script's pairing, which is one step stale: on their own training episodes' 16 trajectories it raises the autoregressive-10k checkpoints' error by +124.6% [+75.8, +201.9] at 1 step and +209.3% [+167.3, +257.6] at 100 steps (three-seed mean, relative-L1). On the two held-out episodes' 4 trajectories the change is smaller and, at 1 step, not resolved: -1.0% [-6.1, +27.7], and +16.6% [+1.9, +72.6] at 100 steps (results/n5_stale_armA_ci.json). The forecasts depend on the action (rule X2, pre-registered: RESPONDS TO THE ACTION; results/action_sensitivity.json), but every action in the training data comes from one policy, so their accuracy under actions that policy would not take is untested.
• Normalisation. States are normalised with the reference's stored mean and std. Actions are not normalised, matching the reference.

Intended use

Reproduction, verification and further study of the claims in the two papers above. These are CPU-trained research artifacts on one dataset, one gait and one terrain. They are not intended for deployment on hardware.

Limitations

• Most of these run at ensemble size 1, against the reference's 5. On those the epistemic term is identically zero and the σ discussed above is aleatoric. The ens5 arms and the five-seed independent ensemble cover the epistemic term.
• One gait, one terrain, one command distribution. Generalisation here means across velocity commands only.
• Long-horizon claims rest on 4 independent 400-step trajectories in the held-out arena. That is the binding statistical constraint.
• The 10k checkpoints are 3 seeds per arm. The headline's ± is a standard deviation over 3 values.
• No policy learning. Dynamics model only.

Licence and attribution

Apache 2.0. Upstream: robotic_world_model_lite (Apache 2.0) and rsl_rl_rwm (BSD 3-Clause, ETH Zurich and NVIDIA); neither is redistributed here.

Citation

Cite the paper, arXiv:2611.02699. CITATION.cff in the repository has the entry.

All checkpoints, in safetensors

Every saved model checkpoint of this reproduction is also here in safetensors format: 272 files from 127 training runs, 1,550 MB, under safetensors/<run>/weights_<iteration>.safetensors, a layout that mirrors runs/ in the code repository. The .pt files in checkpoints/ above are unchanged. Each run keeps every checkpoint it saved: the 500- and 2,500-iteration weights, and the intermediate ones of the 10,000-iteration runs.

safetensors_manifest.json lists, for every file, its sha256, the source .pt and that file's sha256, the iteration, seed and labels (arm, or architecture, regime and spec, and for the fold runs the held-out episode pair), the parameter count, and the results artifact in the code repository that reports its run. A copy is results/hf_safetensors_manifest.json. Each file was read back after conversion and every tensor compared bitwise with the original.

from safetensors import safe_open
from safetensors.torch import load_file

path = "safetensors/armA_seed0/weights_2500.safetensors"
state_dict = load_file(path)           # the model's state dict, as saved
with safe_open(path, "pt") as f:
    labels = f.metadata()              # iter, arm or arch/regime/spec, seed, source .pt
# model.load_state_dict(state_dict)    # with the model built as in the code repository

Not included: the rolling_* resume snapshots, which are training state rather than evaluated models, and the upstream released checkpoint, which belongs to its authors and is not redistributed. The architecture baselines are this reproduction's own implementations.

Follow-up studies' checkpoints, in safetensors

Two later studies started from this reproduction's models: one compressing them, one fine-tuning their magnitude-pruned copies. Their saved checkpoints are on the Hub too: 165 files from 105 training runs, 560 MB, under the same layout.

• Pruning fine-tune study (ft_u*): 30% or 50% of the weights in the GRU and linear layers set to zero by global magnitude, then fine-tuned for 500 iterations with those weights held at zero. Metadata level is the fraction removed, and source_run names the run it started from.
• Compression study, fine-tuned (e3ft_*_w<width>): hidden units removed down to width 192 or 128, chosen by an L2-norm importance score, then fine-tuned for 500 iterations. The width-256 runs are controls, fine-tuned the same way with nothing removed. The kept units are listed in each file's kept_units metadata.
• Compression study, from scratch (scratch_w<width>_ho<ab>_seed<s>): models of width 192 or 128 trained from scratch with episodes a and b held out. Their checkpoints are at 500, 2,500 and 3,000 iterations.

fold (metadata) is the pair of held-out episodes. These runs' manifest entries carry a study field. Their source_pt and results_artifact paths are relative to that study's own folder, which is not part of this repository. Each file was converted, read back and compared tensor by tensor with its .pt, as above. The rolling_* resume snapshots are again not included.