Building a one-step
world model

Diffusion models learn to reverse a noise-corruption process. Sampling follows a sequence of model evaluations from noise toward data. A world model repeats this process whenever it needs another future chunk, multiplying the cost across time and candidate action sequences.

Faster samplers and consistency models shorten this path. Consistency training can be standalone or distilled; a pretrained teacher is not a requirement of every few-step approach. The broader question is whether we can train a useful one-call generator directly.

Drifting models take that route by evolving the generated distribution during training, using attraction toward data and repulsion among generated samples. DriftWorld brings this idea to action-conditioned video. It is an important one-step comparison: our speed advantage is over iterative diffusion, not over other one-call models.

Beckmann transport gives us a different organizing principle: learn the destination associated with an autonomous trajectory. We use that principle to train a conditional map, then add supervision aimed at the endpoint a world model actually produces.

Fix the context $c=(h,a)$, where $h$ is the observed history and $a$ is a proposed action sequence. For an autonomous flow, the terminal map is constant along a trajectory. Its differential form and data boundary condition are:

For RGB video, our map is a residual U-Net. It receives the future tensor being transported, four history frames, and action embeddings. The inference rule is simply:

One call predicts four future frames on PushT or two on Robomimic Can, at 96 × 96 resolution. We append the prediction to the history and repeat. The network has no interpolation-time input; this is a time-independent endpoint map.

We sample a straight interpolation between Gaussian source $z$ and recorded future $y$. A Jacobian–vector product (JVP) measures how the endpoint map changes along the source-to-target chord, holding history and actions fixed:

$\operatorname{sg}$ means stop-gradient. The target is detached, so optimization regresses the map toward a stopped update; it does not backpropagate through a squared JVP using second derivatives. The chord is a sampling device, not the actual autonomous characteristic, and individual sampled JVPs need not be zero.

The last two terms anchor the data boundary and the neighborhood near data. Squared norms here use coordinate means; the near-data term is averaged over the whole batch. A detached coefficient $\lambda$ balances transport and anchor gradient norms on the last spatial output layer. Implementation ↗

Transport geometry describes a map over the future tensor. World modeling adds two demands: the prediction must correspond to the supplied actions, and the model must remain useful when its history contains earlier predictions. We addressed these through the training recipe.

A correct boundary value $T(y,c)\approx y$ does not directly supervise $T(z,c)$, the query used at inference. Source loss does. Its coefficient ramps to one over 500 updates, outside the transport-loss balancing rule.

There is a tradeoff: squared error rewards the conditional mean and penalizes variation across source draws. At fixed context and target, the error decomposes into mean error plus sampling variance:

Lower error therefore does not prove that a generator preserves all plausible futures. Conditional diversity remains an open evaluation question.

Most pixels in a manipulation scene can stay almost unchanged. We add a target-motion-weighted loss, a within-chunk temporal-difference loss, and a penalty for predictions outside the normalized RGB range:

The moving-region denominator has a floor so sparse motion does not amplify the loss excessively. Every fourth update also uses short generated-history recovery: earlier predictions are detached, and only the last prediction receives gradients. Its weight ramps to 0.25 over 1,000 updates.

The recorded future remains the training target, so this is a local robustness regularizer, not supervision for arbitrary counterfactual histories. We retain unclipped RGB feedback during evaluation. The working continuation also uses a quarter of the earlier learning rate; the observed gains belong to the complete recipe.

Pixel error alone gives a limited description of object structure. We add a loss between spatial VGG16 feature maps of the source prediction and target, using relu1_2, relu2_2, and relu3_3:

Features are normalized across channels at each location, then channel-squared differences are summed and averaged across locations, frames, and examples. The feature extractor is frozen, but gradients flow through its input into the prediction. The weight ramps to 0.05 over 1,000 updates.

This is output supervision using pretrained features; it is not the learned LPIPS evaluation metric. VGG is absent at inference. The RGB generator is trained without a pretrained generative teacher.

The most useful experiments were not all wins. Two branches tested whether adding a learned transport direction or propagating gradients through connected predictions would improve the existing map.

We added a head to predict the chord direction and gradually substituted it into the JVP. On Can, 64-frame MSE rose from 0.0068 to 0.0203 and LPIPS from 0.0225 to 0.0854. PushT improved in this continuation, so the evidence argues against adopting it as a reliable default—not against every learned-direction method.

The Can map was already deteriorating before the learned tangent was enabled. The auxiliary head shared decoder features and a clipping group with the endpoint model. That makes auxiliary-task interference a plausible explanation; the logs do not establish its share of the damage or rule out optimizer-restart effects. [audit]

The attached variant backpropagated through two connected predictions at the full learning rate. The detached recovery recipe used a lower learning rate and performed better on the final 64-frame metrics for both tasks. Because both the gradient path and optimization recipe changed, this comparison does not isolate detachment alone.

Can also exposed a gap between short and long rollouts. The source model’s 64-frame MSE was 0.0068, but full-episode MSE was 0.0940, with a worst-episode value of 0.737. The robust recipe reduced full-episode MSE to 0.0112. Good short-window prediction had not been enough to make repeated feedback reliable.

Spatial supervision subsequently improved 64-frame quality, but not every number moved together: Can’s full-episode MSE was 0.0112 for the robust model and 0.0114 after five spatial epochs. We kept the full-episode measurements visible while extending training.

The completed spatial continuations produce the strongest final RGB prediction metrics in our local comparison. Relative to spatial epoch five, 64-frame MSE falls from 0.0155 to 0.0098 on PushT and from 0.0028 to 0.0022 on Can. The improvements require no additional inference steps.

“30 epochs” refers to the spatial stage, not the model’s whole training history. Including inherited training, the final models use 301,028 updates on PushT and 45,000 on Can. Update counts also do not equal compute: batch sizes and per-update costs can differ.

For context, a released long-trained DriftWorld checkpoint on PushT records MSE 0.0085 and LPIPS 0.0163 after about 1.18 million updates—better than our final values. Its exclusion of these evaluation episodes is unverified, so we show it as a reference rather than a controlled ranking. The evidence supports an improving one-step recipe, not a universal state-of-the-art claim.

Averages hide where a model loses the scene. The viewer below synchronizes all methods on the same saved episode. It comes from the earlier spatial-epoch-five comparison; Figure 1 contains the new epoch-30 predictions.

Even at epoch 30, some PushT predictions become visibly blurred or diverge from the recorded block configuration later in the rollout. The examples below make those remaining errors easy to inspect. We cannot attribute their cause from appearance alone.

On a dedicated A800, the spatial-epoch-five RGB model takes 20.6 ms per PushT chunk and 23.9 ms per Can chunk. Accounting for four versus two output frames gives 193.9 and 83.8 predicted frames per second, respectively.

The measured gains are approximately 6.8× / 2.9× over GPC and 101× / 99× over AVDC for PushT / Can. DriftWorld and the regression U-Net are already one-call models with similar latency. For those comparisons, the question is prediction quality and training behavior, not fewer sampling steps. [timings]

PushT planning uses a frozen pose reader to score the last predicted frame for candidate action sequences. The planner needs to order nearby alternatives correctly. That is a different requirement from reducing average error along recorded trajectories.

The final BWM model scores 0.721 on GPC-RANK, versus 0.709 for the local DriftWorld baseline. The paired difference is +0.012 with a 95% interval of [−0.047, +0.071]. This does not establish a planning improvement.

In a 64-window probe, swapping actions while holding real history fixed increases the spatial model’s moving-region error by 11.0× on PushT and 6.2× on Can. The predictions respond to the action input, rather than relying only on visual history.

A useful next target is candidate-ranking error: whether the model prefers the action that really produces the better outcome. Neither a sharp frame nor a large action-swap penalty answers that question on its own.

We also transferred the recipe to 256-pixel robot video. The native model uses SD3 VAE latents with 16 channels at 32 × 32 resolution and a DriftWorld-style native U-Net: 175.4M parameters for Bridge-V2 and 160.3M for RT-1.

Each call predicts one latent frame, conditioned on four context frames and actions. The spatial and range losses operate through the decoded endpoint; the temporal term is zero for a one-frame output. The pretrained VAE remains part of this system.

The completed subset runs use 2,078 eligible training episodes on Bridge and 2,041 on RT-1. Unlike the RGB continuations, these runs schedule map, source, robustness, and spatial stages within a single 30-epoch subset budget.

Generator time is 16.3 ms per Bridge frame and 13.1 ms per RT-1 frame; VAE decoding adds about 26 ms per frame. The latest run table also records approximately 2.8 and 2.3 training GPU-hours for the subset runs, excluding evaluation.

At the September 29, 04:16 UTC source snapshot, full-data continuations were still training and matched-subset baselines were queued. There are no completed full-data results to report. Snapshot ↗ · Latest run table ↗

The practical contribution is a direct conditional map whose output is supervised where it is used: at the Gaussian source, under imperfect history, and in spatial feature space. The strongest evidence is the completed improvement in video prediction while retaining one-call generation.

The immediate open questions are whether these gains persist in matched native comparisons, whether useful conditional variation survives source supervision, and whether better endpoint structure improves action selection. The current measurements leave all three questions open.