Beckmann World Models: Direct Terminal-Map Learning for One-Step Action-Conditioned Prediction

Adam Lee; Shobhit Agarwal; June Pyo Suh

cs.CVOct 16, 2026

Beckmann World Models: Direct Terminal-Map Learning for One-Step Action-Conditioned Prediction

Authors: Adam Lee, Shobhit Agarwal, June Pyo Suh

Organizations: UC Berkeley · Stanford University · Independent Researcher

Abstract

Action-conditioned world models support planning and policy evaluation by predicting future states. However, learning an accurate one-step predictor can require substantial training computation. We introduce Beckmann World Models (BWM), which learn a direct map from noise, visual history, and actions to future frames without a pretrained generative teacher. We adapt the terminal-map conservation law of Beckmann Transport Models to action-conditioned video prediction, training directly from noise-data pairs without conditional populations of generated samples required by DriftWorld. However, transport consistency on interpolated inputs does not ensure accurate prediction from the pure noise used at inference. We address this gap with endpoint-aligned training that directly supervises the source prediction, emphasizes motion and spatial structure, and teaches recovery from generated histories. All components train the same map and retain one network call per future chunk. On PushT and Robomimic Can, the final model reduces 64-frame prediction MSE by 68% and 57%, respectively, relative to same-backbone regression baseline. For policy evaluation, BWM achieves the lowest mean IoU error among the evaluated models, 0.0100, at 24.8 ms per four-frame PushT prediction. On Robomimic Can, the BWM stage trained with motion supervision and generated histories achieves 3.3 times higher throughput and 78% lower peak memory than local one-step DriftWorld. In a latent-space extension to RT-1 and Bridge-V2, equal-update controls show that feature supervision under generated histories improves perceptual prediction. These findings show how terminal-map learning can combine efficient training with accurate one-step visual prediction.

Figures & tables

Figure 1: Conditional terminal-map prediction and training. (a) Observed history and aligned actions condition a one-call prediction. The endpoint image shows a representative prediction on Robomimic Can. (b) Pair-based training uses $x_{s}=(1-s)z+sy$ and the target boundary $y$ . Source anchoring supervises $T_{\theta}(z,h,a)$ at $s=0$ . (c) Target frames define motion masks and within-chunk differences. Prediction quality is evaluated in Section 4 .

Figure 2: Source anchoring improves prediction from initial noise. Measured Can endpoint MSE on 90 fixed windows and four shared source draws per window. The stabilized map has low error at the target boundary (

s=1

) but much higher error at the source (

s=0

). Source anchoring reduces source MSE by

71\%

. The continuation with motion supervision and generated-history recovery reduces it further. Points denote measured input locations. Connecting lines guide the eye. These checkpoints precede spatial training.

Figure 3: Endpoint supervision. (a) A target-defined mask weights changed pixels, normalized per example. The notation $\langle\cdot\rangle$ denotes averaging over tensor coordinates. (b) Temporal differences are compared within a future chunk. Color shows channel-averaged signed RGB differences on a shared scale. (c) A frozen VGG16 encoder matches three levels of normalized features, detaching only target features. (d) A detached generated prefix conditions a new prediction using actions and targets from the dataset. Fixed PushT crops and saved Can rollout frames illustrate the losses, rather than a captured training batch. Appendix A gives exact definitions and schedules.

Figure 4: Can prediction quality and training resources. Left: evaluated checkpoints against cumulative training-loop time, counting all preceding BWM stages. The BWM curve ends at the checkpoint trained with motion supervision and generated-history recovery, before spatial supervision. Right: complete-update profiling on two A800 GPUs, with one conditioning window per GPU. Throughput counts these primary windows, including the cost of auxiliary recovery predictions. Peak memory is allocated memory per GPU. Both models use one call at inference.

Figure 5: Policy-score accuracy versus inference speed on PushT. Error is the mean absolute difference between model and simulator policy-level mean IoUs. Latency is for four predicted frames. Training histories and sampler calls are given in Table 1 . Timings exclude policy execution and reward estimation.

World model	Training	Calls	ms $\downarrow$	IoU error $\downarrow$	$r$ $\uparrow$
DriftWorld	Local	1	21.5	.0306	.8565
GPC	Local	12	140.7	.5127	.3659
AVDC	Local	100	1866.2	.3951	-.0801
DriftWorld	Released	1	21.5	.0203	.9255
GPC	Released	12	139.1	.0860	.8438
BWM	Local	1	24.8	.0100	.9280

Table 1: PushT policy evaluation and inference cost. Error and correlation use seven policy means, each from 300 paired trials. Calls and latency refer to four predicted frames. Local baselines are five-epoch fits.

	PushT			Robomimic Can
Model	MSE	Moving MSE	LPIPS	MSE	Moving MSE	LPIPS
DriftWorld	.0314	.4093	.0790	.0104	.0769	.0682
MSE U-Net	.0303	.4008	.0804	.0051	.0544	.0104
GPC	.0630	1.4241	.1609	.0276	.3103	.0589
AVDC	.0684	1.1840	.1860	.0400	.2644	.1194
BWM	.0098	.1247	.0280	.0022	.0242	.0056

Table 2: RGB prediction on episode-held-out splits. MSE, moving-region MSE, and LPIPS for 60 autoregressively predicted frames after the initial context. All metrics are lower-is-better. BWM is the final spatial-stage model. Baselines are local five-epoch fits. These rows compare different total training budgets.

Figure 6: Autoregressive RGB predictions. Reference and predicted frames at positions 24, 44, and 64 in two selected evaluation sequences, with the same actions and times across rows. BWM uses the final spatial continuation. Local baselines use their five-epoch fits. Appendix B.7 gives sample selection and additional baseline rows.

Figure 7: Native-video prediction with generated-history feature supervision. Reference, control, and recovery-refined predictions at matched positions in selected RT-1 and Bridge-V2 episodes. References are VAE reconstructions. Each eight-frame prediction window starts from four frames. Separated columns span several such windows. Residual gripper blur remains visible.

Dataset	Model	PSNR $\uparrow$	LPIPS $\downarrow$	FID $\downarrow$	FVD $\downarrow$
RT-1	Parent	27.35	.0917	19.77	167.50
	Control	27.34	.0933	20.56	172.89
	$+$ recovery	27.27	.0612	8.05	104.91
Bridge-V2	Parent	25.00	.1113	31.86	175.53
	Control	25.01	.1106	30.98	171.56
	$+$ recovery	24.95	.0841	21.34	117.20

Table 3: Feature supervision under generated histories. Full-data BWM parents and matched-length continuations. RT-1 adds 1,000 updates. Bridge-V2 adds 500. The control retains clean-history feature drifting. Recovery adds it at generated-history endpoints. Both continuations use the same parent and one training seed.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Protocol	Metric	PushT	Robomimic Can
64 frames	MSE $\downarrow$	0.0098 [0.0086, 0.0111]	0.0022 [0.0016, 0.0028]
	Moving MSE $\downarrow$	0.1247 [0.1098, 0.1401]	0.0242 [0.0177, 0.0319]
	LPIPS $\downarrow$	0.0280 [0.0242, 0.0320]	0.0056 [0.0045, 0.0068]
	SSIM $\uparrow$	0.9754 [0.9731, 0.9776]	0.9843 [0.9812, 0.9871]
	PSNR $\uparrow$	31.21 [30.82, 31.61]	34.75 [33.79, 35.69]
Full episode	MSE $\downarrow$	0.0154 [0.0137, 0.0171]	0.0082 [0.0057, 0.0112]

Appendix

Table 4: Final BWM prediction metrics and episode-bootstrap intervals. Each cell is a mean with its 95% interval. The spatial stage has 30 epochs.

Recipe	Mean s	95th percentile s	Windows $\mathrm{s}^{-1}$	Peak GiB
Local DriftWorld	1.841	1.843	1.09	36.07
Source-anchored map	.496	.509	4.03	6.12
Motion and recovery	.553	.739	3.62	7.77

Appendix

Table 5: Complete-update cost for the Can RGB recipes. Equal batches of training windows in full precision.

	Released GPC		AVDC-derived
Policy epoch	Model	Simulator	Model	Simulator
50	0.5105	0.6190	0.1800	0.6103
100	0.5100	0.5929	0.1770	0.5922
150	0.4939	0.6029	0.1829	0.6042
200	0.4862	0.5630	0.1768	0.5609
250	0.4798	0.5536	0.1857	0.5507

Appendix

Table 6: PushT baseline policy estimates. World-model and simulator mean IoUs for each frozen policy. Each entry averages 300 trials.

Checkpoint	Policy error $\downarrow$	Pearson $r$ $\uparrow$	LPIPS $\downarrow$
Spatial, 5 epochs	.0098	.9462	.0407
Spatial, 30 epochs (main)	.0100	.9280	.0280
Learned-direction alternative	.0355	.9743	.0501

Appendix

Table 7: Separate PushT policy and visual checkpoints. Seven frozen policies and 300 paired trials each. LPIPS uses AR64. Entries are descriptive stage results, not independently selected best metrics of one model.

Task	Checkpoint	MSE	Moving MSE	LPIPS
PushT	Released DriftWorld	.0085	.1182	.0163
	Released GPC	.0192	.2446	.0324
	Final BWM	.0098	.1247	.0280
Can	Released DriftWorld	.0525	.0662	.0676
	Final BWM	.0022	.0242	.0056

Appendix

Table 8: Released references on the RGB evaluator. All errors are lower-is-better. These are not equal-training-data comparisons.

Figure 8: Additional RGB baseline comparisons. The same two evaluation sequences as Figure 6 , with local GPC and AVDC-derived rows included. BWM uses its final spatial continuation.

Dataset	Arm	SSIM $\uparrow$	PSNR $\uparrow$	LPIPS $\downarrow$	FID $\downarrow$	FVD $\downarrow$
RT-1	Parent	0.8840	27.35	0.0917	19.77	167.50
	Control	0.8836	27.34	0.0933	20.56	172.89
	Recovery	0.8861	27.27	0.0612	8.05	104.91
Bridge-V2	Parent	0.8555	25.00	0.1113	31.86	175.53
	Control	0.8556	25.01	0.1106	30.98	171.56
	Recovery	0.8580	24.95	0.0841	21.34	117.20

Appendix

Table 9: Complete native continuation metrics. Same cohorts and models as Table 3 .

Dataset	Method	Calls per frame	PSNR $\uparrow$	LPIPS $\downarrow$	FID $\downarrow$	FVD $\downarrow$
RT-1	DriftWorld	1	22.42	0.1295	27.48	237.45
	MSE U-Net	1	23.03	0.1788	62.41	419.16
	GPC-derived	3	20.69	0.2937	124.18	884.52
	AVDC-derived	100	20.27	0.1751	94.87	719.10
	BWM spatial	1	23.15	0.1659	75.53	535.59
Bridge-V2	DriftWorld	1	21.92	0.1042	14.18	188.13

Appendix

Table 10: Thirty-epoch common-subset native comparison. The same training episode subsets and evaluation cohorts are used within each dataset. Batches and optimizer recipes differ. These models precede BWM’s full-data and recovery continuations.

Figure 9: Real robotics video prediction. RT-1 drawer manipulation and Bridge-V2 gripper motion from the final recovery-refined checkpoints. Columns show separated positions in each displayed video (positions 13, 49, and 85 for RT-1 and positions 4, 25, and 46 for Bridge-V2). Predictions are generated within the conditioned windows described in Section 4.6 . References are VAE reconstructions.

Split	Model	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	FID $\downarrow$	FVD $\downarrow$
Validation	Parent	25.00	.8555	.11125	31.86	175.53
	Paired-feature child	24.95	.8577	.10642	31.36	188.04
Test	Parent	25.53	.8631	.10439	32.69	160.06
	Paired-feature child	25.47	.8651	.09964	31.77	172.55

Appendix

Table 11: Paired-feature variant: validation and independent test. Bridge-V2 parent versus the separately selected 1,000-update EMA. This is not the 500-update recovery-population variant in Table 3 . FID and FVD are retained evaluator outputs.

Model	0	.25	.50	.75	1
Stabilized BTM	.002831	.002027	.001437	.000810	.000105
$+$ source	.000823	.001429	.001178	.000687	.000096
$+$ motion and recovery	.000503	.000868	.000747	.000498	.000025
Learned direction	.001170	.002241	.002476	.002513	.000112

Appendix

Table 12: Endpoint MSE across input locations on Can. The source is

s=0

and the target boundary is

s=1

Task	Model	MSE $\downarrow$	Moving MSE $\downarrow$	LPIPS $\downarrow$
PushT	Local DriftWorld	.0314	.4093	.0790
	MSE U-Net	.0303	.4008	.0804
	Fresh BWM	.0211	.2832	.0980
Can	Local DriftWorld	.0104	.0769	.0682
	MSE U-Net	.0051	.0544	.0104
	Fresh BWM	.0048	.0475	.0154

Appendix

Table 13: Fresh generators at matched clean-window exposure. 64-frame evaluation after five capped epochs. BWM has zero inherited updates. Recovery targets, optimizer, and compute differ. The predeclared final checkpoint is used, separately from Table 2 .

World Model Learning

World model learning builds predictive models of how environments change under actions to support planning, simulation, and decision-making.

See all 414 papers

Learning Visual Feature-Based World Models via Residual Latent Action
May 8, 2026Xinyu Zhang, Zhengtong Xu, Yutian Tao +3Latent Action Learning World Model Learning
Rutgers University · Purdue University · University of Wisconsin-Madison
From World Models to World Action Models: A Concise Tutorial for Robotics
Jul 1, 2026Xiaoxiong Zhang, Xiong Zeng, Wei ZhangWorld Model Learning World Models for Robotics
School of Automation and Intelligent Manufacturing, Southern University of Science and Technology, Shenzhen, China · LimX Dynamics, Shenzhen, China
Video World Models Need Predictable Latents
Oct 26, 2026Mattia Scardecchia, Sherry YangVideo World Models World Model Learning
New York University

May 8, 2026J/K move · Enter open · S save

Learning Visual Feature-Based World Models via Residual Latent Action

Xinyu Zhang, Zhengtong Xu, Yutian Tao, Yeping Wang+2

Rutgers University · Purdue University · University of Wisconsin-Madison

Latent Action Learning World Model Learning Action-Conditioned World Models Visual RL Video World Models

Open paper AskPDF arXiv

World models predict future transitions from observations and actions. Existing works predominantly focus on image generation only. Visual feature-based world models, on the other hand, predict future visual features instead of raw video pixels, offering a promising alternative that is more efficient and less prone to hallucination. However, current feature-based approaches rely on direct regression, which leads to blurry or collapsed predictions in complex interactions, while generative modeling in high-dimensional feature spaces still remains challenging. In this work, we discover that a new type of latent action representation, which we refer to as Residual Latent Action (RLA), can be easily learned from DINO residuals. We also show that RLA is predictive, generalizable, and encodes temporal progression. Building on RLA, we propose RLA World Model (RLA-WM), which predicts RLA values via flow matching. RLA-WM outperforms both state-of-the-art feature-based and video-diffusion world models on simulation and real-world datasets, while being orders of magnitude faster than video diffusion. Furthermore, we develop two robot learning techniques that use RLA-WM to improve policy learning. The first one is a minimalist world action model with RLA that learns from actionless videos, and improves VLA on LIBERO and real robot. The second one is a visual RL framework trained entirely inside a world model learned from offline videos only, using a video-aligned reward and no online interactions. Project page: https://mlzxy.github.io/rla-wm

Abstract

World Model Learning

Explore similar work

Learning Visual Feature-Based World Models via Residual Latent Action

From World Models to World Action Models: A Concise Tutorial for Robotics

Video World Models Need Predictable Latents