LeWAM: End-to-End World and Action Modeling with JEPAs

Minghao Fu; Tavis Siebert; Eryk Halicki; Randall Balestriero

cs.ROOct 8, 2026

LeWAM: End-to-End World and Action Modeling with JEPAs

Authors: Minghao Fu, Tavis Siebert, Eryk Halicki, Randall Balestriero

Organizations: UC San Diego · ETH Zurich · University of British Columbia · Brown University

Abstract

Joint-embedding predictive architectures (JEPAs) learn compact latent world models without pixel reconstruction. Existing JEPA pipelines typically train visual representations and action-conditioned dynamics before adding a control policy or solver. We introduce LeWAM, an end-to-end JEPA for control with: (1) a joint training objective over next latent state prediction and action flow matching, and (2) an isotropic Gaussian constraint on the latent state, enabling a single model to support direct policy execution and efficient planning by evaluating sampled actions with its learned latent dynamics. Compared with prior JEPAs in visual planning, LeWAM improves success rate from 28.6% to 89.7% on contact-rich tasks and from 10.9% to 46.2% on long-horizon tasks, and plans 32.3 times faster. On behavior cloning tasks in simulation and physical robots, LeWAM is also competitive with established methods. Interpretability analyses show that LeWAM learns control-relevant representations and action distributions effectively guiding test-time planning and control.

Figures & tables

Panel of Figure 1 — Figure 1: LeWAM Training. A shared encoder encodes observed frames with SIGReg ( Balestriero and LeCun, 2025 ) regularization. The Mixture-of-Transformers ( Liang et al., 2025 ) learns latent-state prediction from clean actions and action flow matching from noisy actions with goal information. Results Overview. The mean of (a) planning success on contact-rich tasks; (b) planning success on a 100-step goal offset; (c) plans/s versus CEM ( Rubinstein, 1999 ) ; (d) action chunks per TFLOP versus $\pi_{0.5}$ ( Physical Intelligence et al., 2025 ) . Full results are in Section 4 .

Figure 2: Positioning of LeWAM. The columns place LeWAM alongside related approaches to world-model training, visual representation learning, world action modeling, and action search. LeWAM jointly trains latent dynamics and a policy through a shared encoder, predicts in latent space without reconstruction, and refines learned action proposals through its dynamics. The final column summarizes its evaluation on contact-rich tasks and longer goal horizons.

Figure 3: MoT Attention. Filled cells permit attention, while white cells are masked. The state query reads history and clean actions, and noisy action tokens read state history and their own causal prefix. Small grids show causal attention. Arrows mark the state branch and action velocity branch.

Figure 4: Planning with LeWAM. Each MoT predictor stage samples $N$ action chunks from Gaussian noise and uses the first $K$ actions in each chunk, denoted ${\bm{A}}_{t,K}$ , to advance the latent dynamics from step $t$ to $t+K$ . Five-point curves depict action chunks (we use $K=5$ ), and stacked curves denote $N$ sampled plans of length $L$ . Step 1 refines all plans by gradient descent through the frozen LeWAM dynamics using the goal-state cost. Step 2 selects the lowest-cost refined iterate as the final plan, shown as ${\bm{A}}^{\star}$ . Appendix A.1 gives the full procedure in a pseudo-algorithm.

Figure 5: Simulation Environments. Symbols beneath each frame denote overlapping task attributes, with their meanings shown in the legend to the left.

Figure 6: Visual Planning across Ten Simulated Environments. Most panels report the success rate under a goal offset of

L=25

environment steps and permit

2L=50

actions across 3 seeds. For Two-Room , we use 100/150 action steps following LeWM setting; For LeWM (CEM), we report 10 seeds for stability. Details are present in Appendix B.2 . Error bars denote one standard deviation; indicates models with end-to-end training. Linear axes start at 50% for the first five panels.

Figure 7: Visual Planning Success across Different Horizons. The task panels show mean success rates across 3 seeds. Each sweep starts at the first trajectory frame but with different goal offsets.

Environment	Model	SIGReg	Eff. rank	Object $\uparrow$	Proprio $\uparrow$	Action (IDM) $\uparrow$
Push-T	LeWAM	Yes	$42.90$	$0.90$	$\mathbf{0.86}$	$\mathbf{0.40}(\mathbf{0.75})$
	LeWAM	No	$10.94$	$0.88$	$0.85$	$0.38(0.73)$
	LeWAM (MSE)	No	$6.96$	$0.66$	$\mathbf{0.86}$	$0.36\,(0.74)$
	LeWM	Yes	$94.20$	$\mathbf{0.95}$	$0.80$	$0.29(0.71)$
	LeWM	No	$16.23$	$0.00$	$0.00$	$0.00\,(0.00)$

Table 1: Linear probing and the role of SIGReg. All linear probes report

R^{2}

score using ridge regression. Object denotes the average

R^{2}

over all object ground truth information and Proprio the ground truth proprioceptive state. Action predicts the action chunk from latent

{\bm{z}}_{t}

; parenthesized IDM scores use both

{\bm{z}}_{t},{\bm{z}}_{t+K}

where reported. Eff. rank denotes effective rank ( Roy and Vetterli, 2007 ) .

Figure 8: (a) Planning gain (percentage points) from SIGReg. (b) Direct success and planning gain for visual encoders, averaged over four tasks. (c) Planning gain for flow (circles) and MSE (squares) action heads. (d) Task-completion success for LeWAM and Diffusion Policy.

Figure 9: Paired scene and wrist view examples from three physical-robot tasks.

Method	Pick and Place T	Push-T	Tape Stack	Average	GFLOPs
Diffusion Policy ( Chi et al., 2023 )	$19/30$	$16/30$	0/30	38.9%	$270.0$
$\pi_{0.5}$ ( Physical Intelligence et al., 2025 )	$14/30$	3/30	9/30	28.9%	$3603.6$
LeWAM (Frozen V-JEPA 2.1 Encoder)	$15/30$	$12/30$	5/30	35.6%	$270.2$
LeWAM	$16/30$	$8/30$	6/30	33.3%	$270.2$

Table 2: Physical robot evaluation. Successful trials out of 30 on Pick and Place T, Push-T, and Tape Stack. GFLOPs count one inference for a full action chunk, including all denoising steps. One full LeWAM action chunk uses 270.2 GFLOPs versus 3603.6 GFLOPs for

\pi_{0.5}

. Figure 1 (d) expresses the reciprocal as action chunks per TFLOP, a

13.3\times

difference in operation-count efficiency.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 10: Token-Level MoT Attention. Expansion of Figure 3 for two history latents, one learned state query, five clean actions, and ten interpolated action tokens. Rows are queries; columns are keys. Filled cells permit attention; white cells are masked. Blue and brown mark latent state and action velocity prediction queries; gray marks context. Wider gaps separate token groups. Optional goal conditioning

{\bm{c}}_{t}

is omitted.

\Delta

is the history stride, and action indices remain at the raw environment rate.

Task	Collection regime	Views	Dim.	Traj.	Mean frames	Total frames	Hz	Hours
Two-Room	Noisy heuristic	1	2	10,000	92	920,809	10	25.58
Reacher	Uniform random	1	2	10,000	201	2,010,000	25	22.33
Push-T	Non-expert	1	2	18,685	125	2,336,736	10	64.91
Cube	Scripted expert	1	5	10,000	201	2,010,000	20	27.92
PointMaze-Large	Scripted navigation	1	2	1,000	1,000	1,000,000	10	27.78
Scene	Scripted expert	1	5	10,024	200	2,004,800	20	27.84

Appendix

Table 3: Simulation Benchmark Datasets. Ten tasks in Figure 5 order. Views gives camera-stream count; Dim. is the raw action dimension. Statistics cover complete source datasets before splits. Total hours are computed before rounding per-task values.

Figure 11: A Degenerate Cube Success under Goal Reaching. Each row shows a policy rollout (left, frames over time) beside its goal frame (far right, after the divider). Throughout the execution, the gripper never reaches the goal end-effector pose, yet the cube stays within the $0.04$ m object-only threshold of the goal, so object-position-only predicate scores the episode a spurious success. Our end-effector term (Table 4 ) rejects these degenerate successes by additionally requiring the gripper to reach the goal pose.

Environment	Raw action dim.	25-step goal reaching		Task Completion
		Zero-action	Random	Zero-action	Random
Two-Room	2	$7.3\pm 0.9$	$25.3\pm 5.0$	$0.0\pm 0.0$	$2.7\pm 2.5$
Reacher	2	$3.3\pm 1.9$	$12.0\pm 1.6$	$0.0\pm 0.0$	$12.0\pm 4.9$
Push-T	2	$0.0\pm 0.0$	$1.3\pm 0.9$	$0.0\pm 0.0$	$0.0\pm 0.0$
Cube	5	$42.7\pm 6.6$	$44.0\pm 1.6$	$4.7\pm 2.5$	$6.0\pm 1.6$
Cube ( $+$ EEF)	5	$2.7\pm 0.9$	$4.7\pm 0.9$	$4.7\pm 2.5$	$6.0\pm 1.6$

Appendix

Table 4: Action Dimensions and Policy-Floor Controls. Rows cover all ten environments in Figure 5 , plus a Cube (

+

EEF) row isolating the corrected end-effector predicate. Raw action dimensionality is reported beside zero-action and random-policy success rates (%) for short-horizon goal reaching and full-trajectory task completion, averaged over three seeds. These blind controls measure baseline success from zero actions and random actions. Cube has the same task-completion values in both rows because the end-effector correction applies only to short-horizon goal reaching. DrawerCleanup exposes 24 raw coordinates; the policy predicts 22 active dimensions after removing two constant joints.

Task	Method	Goal horizon (steps)
		25	50	75	100
Push-T	LeWAM	$94.7\pm 3.1$	$60.7\pm 3.1$	$44.0\pm 3.5$	$43.3\pm 4.6$
	LeWM (CEM)	$68.7\pm 3.1$	$12.0\pm 4.0$	$3.3\pm 2.3$	$2.7\pm 1.2$
Cube	LeWAM	$100.0\pm 0.0$	$88.0\pm 2.0$	$80.7\pm 2.3$	$82.0\pm 2.0$
	LeWM (CEM)	$100.0\pm 0.0$	$0.0\pm 0.0$	$4.0\pm 3.5$	$4.0\pm 3.5$
Scene	LeWAM	$82.7\pm 3.1$	$73.3\pm 8.1$	$65.3\pm 3.1$	$56.0\pm 4.0$

Appendix

Table 5: Goal-Reaching Success across Horizons. Success rate (%), mean

\pm

standard deviation. Measured values for Figure 7 . The Cube horizon sweep is distinct from the sampled-start 25-step evaluation in Figure 6 .

Task	Action dim.	ms/plan	plans/s	ms/plan	plans/s	Speedup
		LeWM (CEM)		LeWAM
Scene	$5$	$705$	$1.42$	$18.8$	$53.19$	$37.5\times$
Puzzle	$5$	$731$	$1.37$	$16.3$	$61.35$	$44.8\times$
Cube	$5$	$771$	$1.30$	$21.1$	$47.39$	$36.5\times$
Push-T	$2$	$414$	$2.42$	$20.8$	$48.08$	$19.9\times$
Task mean	–	–	$1.62$	–	$52.50$	$32.3\times$

Appendix

Table 6: Planning Speed across Four Visual-Goal Tasks. Each plan is one planner call. LeWM (CEM) ( Maes et al., 2026b ) uses

300\times 30

CEM; LeWAM uses five gradient steps. Action dimensions are raw environment action dimensions. Throughput is the reciprocal of time per plan. The mean row averages per-task throughput, and its ratio compares those two means.

Figure 12: Planner Sensitivity on Push-T . Goal-reaching success (mean and standard deviation across three seeds, 50 episodes each) and per-plan latency versus (left) Adam refinement updates with 32 candidate plans and (right) sampled candidate plans with five refinement updates.

Environment	Variant	Direct	Planning	Gain (pp)
Push-T	LeWAM	$82.7\pm 8.1$	$97.3\pm 3.1$	$+14.7$
	w/o SIGReg	$69.3\pm 10.1$	$72.0\pm 3.5$	$+2.7$
	MSE action head	$90.0\pm 4.0$	$95.3\pm 1.2$	$+5.3$
Scene	LeWAM	$92.0\pm 5.3$	$97.3\pm 1.2$	$+5.3$
	w/o SIGReg	$91.3\pm 2.3$	$93.3\pm 3.1$	$+2.0$
	MSE action head	$89.3\pm 4.2$	$90.0\pm 2.0$	$+0.7$

Appendix

Table 7: SIGReg and Action-Head Ablations. Direct and planning success rates (%) using epoch-30 checkpoints for all variants. Gain is planning minus direct success in percentage points. Standard deviations are reported where available; none is inferred for gains.

Visual encoder	Mode	Push-T	Cube	Scene	Puzzle	Average
LeWAM (end-to-end)	Direct	$85.3\pm 5.0$	$100.0\pm 0.0$	$88.0\pm 6.9$	$73.3\pm 5.8$	$86.7$
	Planning	$98.7\pm 1.2$	$100.0\pm 0.0$	$93.3\pm 3.1$	$74.0\pm 12.5$	$91.5$
	Gain (pp)	$+13.4$	$0.0$	$+5.3$	$+0.7$	$+4.9$
frozen DINOv2 CLS	Direct	$72.7\pm 7.6$	$98.7\pm 1.2$	$82.0\pm 4.0$	$62.7\pm 7.0$	$79.0$
	Planning	$78.7\pm 6.4$	$99.3\pm 1.2$	$91.3\pm 1.2$	$61.3\pm 13.0$	$82.7$
	Gain (pp)	$+6.0$	$+0.6$	$+9.3$	$-1.4$	$+3.6$

Appendix

Table 8: Visual Encoders for Direct Control and Planning. Success rate (%) with the same goal-conditioned evaluation protocol for all encoder variants. Direct executes the policy without search; planning uses test-time action search. Gain is planning minus direct success in percentage points. Averages are equal-weight means over the four tasks; gains use unrounded means.

Environment	Model	SIGReg	Eff. rank	Object $\uparrow$	Proprio $\uparrow$	Action (IDM) $\uparrow$
Push-T	LeWAM	Yes	$42.90$	$0.90$	$\mathbf{0.86}$	$\mathbf{0.40}\,(\mathbf{0.75})$
	LeWAM	No	$10.94$	$0.88$	$0.85$	$0.38\,(0.73)$
	LeWAM (MSE)	No	$6.96$	$0.66$	$\mathbf{0.86}$	$0.36\,(0.74)$
	LeWM	Yes	$94.20$	$\mathbf{0.95}$	$0.80$	$0.29\,(0.71)$
	LeWM	No	$16.23$	$0.00$	$0.00$	$0.00\,(0.00)$
Cube	LeWAM	Yes	$46.60$	$\mathbf{0.99}$	$0.67$	$\mathbf{0.37}\,(0.66)$

Appendix

Table 9: Effective Rank and Linear Probing. Results on Push-T , Cube , Scene , and Puzzle , with the Object, Proprio, and Action columns reporting linear-probe

R^{2}

. Rows compare models with and without SIGReg; parenthesized values use both

{\bm{z}}_{t}

and

{\bm{z}}_{t+K}

for inverse-dynamics probing.

Model	Push-T	Cube
LeWM (CEM)	$\mathbf{92.4\pm 1.2}$	$63.8\pm 2.2$
LeWAM (CEM)	$84.0\pm 1.3$	$\mathbf{86.0\pm 3.2}$

Appendix

Table 10: CEM Planning with Learned Dynamics. Goal-reaching success rates (%), reported as mean

\pm

standard deviation. LeWAM is evaluated with its action prediction head removed.

Figure 13: Aligned Visual Sequences on Push-T and Cube . Each task shows recorded ground truth, framewise reconstructions for LeWM and LeWAM, and decoded open-loop predictions for both models. Columns use the same relative frame indices; predictions follow the recorded actions.

Figure 14: Aligned Visual Sequences on Scene and Puzzle . Rows show recorded ground truth, LeWM and LeWAM framewise reconstructions, and their decoded open-loop predictions. Columns use the same relative frame indices and the same recorded action sequence for both prediction rows.

Panel of Figure 15 — Figure 15: Two-Door Demonstrations and Action-Head Rollouts. Expert demonstrations (a), executed trajectories from the flow-matching action head (b), and executed trajectories from the deterministic MSE action head (c). Black segments mark walls, and the green cross marks the goal. Blue and orange indicate upper- and lower-door trajectories; dashed red paths end in collision.

Method	DrawerCleanup	Transport	ToolHang	Average
Diffusion Policy ( Chi et al., 2023 )	$\mathbf{54.7\pm 1.2}$	$45.3\pm 9.0$	$\mathbf{84.7\pm 3.1}$	61.6
LeWAM	$52.7\pm 5.0$	$\mathbf{61.3\pm 9.9}$	$80.0\pm 3.5$	$\mathbf{64.7}$

Appendix

Table 11: Task completion. Mean success rates (%)

\pm

standard deviation across three evaluation seeds, with 50 episodes per seed. Bold marks the better result.

Task	Episodes	Frames	Duration	Task Description
Pick and Place T	200	52,297	29 min	”Pick up the T block and place it onto the outline.”
Push-T	100	61,996	34 min	”Push the T block onto the outline.”
Tape Stack	100	23,726	13 min	”Stack the green roll of tape on the white roll of tape.”

Appendix

Table 12: Datasets for the physical-robot tasks in Section 4.4 .

Panel of Figure 16 — Figure 16: Matched visual reconstructions and recorded-action latent rollouts for the scene camera. Within each model, rows show the ground-truth (GT) video, independent frame reconstructions, and decoded autoregressive predictions.

World models
Method	FDM	Policy	Joint FDM and Policy	End-to-end from scratch
DINO-WM ( Zhou et al., 2024 )	✓	✗	✗	✗
V-JEPA 2-AC ( Assran et al., 2025 )	✓	✗	✗	✗
LeWM (CEM) ( Maes et al., 2026b )	✓	✗	✗	✓
World action models
UVA ( Li et al., 2025 )	✓	✓	✗	✗

Appendix

Table 13: Modeling and Control Capabilities. FDM denotes forward dynamics conditioned on clean actions; Policy denotes learned action prediction. Entries describe reported training recipes or the goal-reaching configurations used here.

Figure 18: Action Proposals and Predicted Consequences. (a) Smooth rollouts; (b) candidate rollouts branch at contact; (c) LeWAM pairs learned action proposals with predicted outcomes. Dashed links denote search; the solid link denotes joint prediction.

World Model Learning

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

See all 399 papers

JEPA-WAM: Learning Vision-Language-Action Policies with Joint-Embedding World Modeling
Aug 10, 2026Yihan Lin, Jiawei He, Shifeng Bao +6Joint-Embedding Predictive Architecture Robotic Manipulation
School of Information, Renmin University of China, Beijing, China · XYZ Embodied AI, Beijing, China · Key Laboratory of Data Engineering and Knowledge Engineering, Beijing, China +3
LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC
Oct 8, 2026Shashank Hegde, Alexander Popov, Elie Aljalbout +1World Action Models Latent World Models
NVIDIA
Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding
Jun 30, 2026Zhenghao Zhang, Yuanxiang Wang, Zhenyu Guan +11World Models World Model-Based Planning
School of Computer Science and Technology, University of Chinese Academy of Sciences, Beijing · Institute of Information Engineering, Chinese Academy of Sciences, Beijing · School of Computer Science and Technology, Harbin Institute of Technology, Weihai +2

Aug 10, 2026J/K move · Enter open · S save

JEPA-WAM: Learning Vision-Language-Action Policies with Joint-Embedding World Modeling

Yihan Lin, Jiawei He, Shifeng Bao, Chen Zhao+5

School of Information, Renmin University of China, Beijing, China · XYZ Embodied AI, Beijing, China · Key Laboratory of Data Engineering and Knowledge Engineering, Beijing, China+3

Joint-Embedding Predictive Architecture Robotic Manipulation Vision-Language-Action Models Latent World Models Embodied World Models

Open paper AskPDF arXiv

Robust robot control benefits from explicitly modeling state transitions, but video-generation world action models (WAMs) introduce substantial deployment cost. Existing latent WAMs avoid explicit future generation, but often compress predictive representations or separate predictive modeling from the representations used for action generation. We introduce JEPA-WAM, a latent WAM built in a pretrained V-JEPA space, which couples latent transition prediction with continuous action generation through a shared predictor. JEPA-WAM predicts a spatially structured joint current-future target that captures task-shared visual temporal structure between current and future observations, while preserving dense patch-level correspondence. Through the shared predictor, transition supervision directly shapes the backbone, from which dedicated representations are extracted for action prediction. The same design can also be instantiated in pretrained VLA policies while preserving their original perception and action pathways. On LIBERO-Plus, JEPA-WAM achieves 79.2%, the best result without large-scale robot-policy pretraining, while its pretrained

π_{0.5}

instantiation reaches 86.3%, achieving the best overall performance. Experiments on RoboTwin 2.0 and real-world bimanual manipulation further demonstrate strong generalization under visual and spatial shifts.

Abstract

World Model Learning

Explore similar work

JEPA-WAM: Learning Vision-Language-Action Policies with Joint-Embedding World Modeling

LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC

Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding