Skip to content

Repository files navigation

RDR

The decoder of a latent world model is trained on latents anchored to observations: encoder outputs, and the one-step predictions teacher forcing makes from them. At deployment it is applied to the model's own free-running rollout, which by then has run hundreds of steps since the last observation.

We propose Rollout-Decoded Reconstruction (RDR), a single loss term that free-runs the model during training, decodes every rollout latent, and penalizes reconstruction error against ground truth. RDR adds no parameters and requires no architectural changes, and setting its weight to zero recovers the standard objective exactly, so every comparison is a one-flag A/B at an identical parameter count.

On the chaotic Kuramoto–Sivashinsky equation, RDR raises valid prediction time from 3.87 to 6.97 time units, a 1.80× improvement at an identical 193,568 parameters, confirmed on seeds never used in selection and in 10 of 10 preregistered configurations at ratios of 1.71–2.50×. The improvement grows with latent width, from 1.41× at dimension 16 to 2.50× at dimension 48, while the standard objective loses horizon as capacity is added.

In model-based control, ensembles of RDR world models reach useful behavior in fewer optimizer steps, winning 20 of 20 paired episodes at every reduced-data rung on two tasks, and are robust to a planner–training rollout mismatch that costs the standard objective more return in 4 of 4 measured rows. Both results come from the same single term.

RDR is an objective rather than an architecture: it composes with any world model that retains a decoder.

Results

Forecasting, chaotic KS at L = 22, three seeds per arm, capacity-matched.

posterior-only +RDR ratio
VPT, canonical horizon 3.87 ± 0.23 tu 6.97 ± 0.42 tu 1.80×
VPT, long horizon (1300 steps) 3.77 tu 6.90 tu 1.83×
trainable parameters 193,568 193,568 identical
preregistered configurations won 10 of 10 1.71–2.50×

Measured at fresh seeds 10–12, none used during selection. The preregistered gate, locked before any sweep number existed, required 5.77 tu; the primary arm clears it at 6.47 ± 0.40 with every seed above the bar.

Capacity scaling. The standard objective loses horizon as the latent widens; RDR converts the same capacity into horizon.

latent posterior-only +RDR ratio
16 3.00 ± 0.36 4.20 ± 0.27 1.41 ± 0.10
24 2.77 ± 0.06 4.73 ± 0.15 1.71 ± 0.09
32 2.60 ± 0.27 5.90 ± 0.46 2.30 ± 0.43
48 2.27 ± 0.06 5.67 ± 0.55 2.50 ± 0.22

The posterior arm falls at every rung in 3 of 3 seeds; all 36 per-seed paired ratios exceed 1. The bracket is descriptive rather than preregistered, its top step holds in only 2 of 3 seeds, and the posterior arm drops below the persistence baseline at latent 48, so part of that last ratio is baseline collapse.

Control, CEM-MPC over ensembles of ten, 20 paired episodes on pendulum and cartpole. RDR wins 20 of 20 at every reduced-data rung under fixed epochs. Matching optimizer steps collapses the margins from +745/+1072/+396 to +3/+13/−18, which places the advantage in optimization rather than data volume. Under a planner–training rollout mismatch the posterior arm loses more return in 4 of 4 rows (+9.0/+18.1/+17.1/+12.7 against RDR's +1.8/+3.6/+12.8/+10.6); under matched conditions the arms are equivalent.

The absolute scale. λ_max ≈ 0.043, so one Lyapunov time is 23.26 tu.

tier VPT (tu) Λt
persistence baseline 2.40 0.10
latent, posterior-only 3.87 0.17
latent + RDR, split head 6.47 0.28
latent + RDR, shared decoder 6.97 0.30
observation-space pushforward, matched budget 7.00 0.30
published latent-ROM band, symmetry-reduced 33–47 1.4–2.0
full-state reservoir (ESN), this work 73.0 3.14

On this fully observed 64-dimensional system an observation-space predictor reaches parity at matched budget, which bounds the latent bottleneck rather than the objective. The RDR contrast is within-latent throughout. The settings that force a latent, partial observability, pixel observations, or planning over a compact state, are also the settings where an observation-space predictor is unavailable.

The full protocol, every ablation, and the caveats that discipline each number are in the paper.

Acknowledgements

Thank you to Google (GCP) and E3A Healthcare for providing the resources to run this project!

About

Reference implementation for the paper "Rollout-Decoded Reconstruction for Latent World Models": one loss term that trains a latent world model's decoder on its own free-running rollout, at zero added parameters.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages