PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Minghao Fu, Tavis Siebert, Eryk Halicki, Randall Balestriero

arXiv:2610.12818v1cs.ROcs.AI

TL;DR

Existing JEPA control pipelines can leave action supervision disconnected from visual representation learning, and latent models may retain information irrelevant to control. LeWAM jointly trains latent prediction and action flow matching with an isotropic Gaussian latent constraint, supporting direct policy execution and planning. It reports stronger goal-reaching results than LeWM in the supplied visual-planning comparisons and competitive behavior-cloning performance, while its authors identify limited data coverage and task- and embodiment-specific training as boundaries.

  • Problem

    Latent visual world models may retain distractors or extraneous observation information, while control often depends on costly online search or separately trained action predictors whose supervision is disconnected from representation learning.

    Source: §1 · p. 2

  • Method

    LeWAM jointly trains shared visual representations, action-conditioned latent dynamics, and a stochastic flow-matching policy, with an isotropic Gaussian constraint on latent states.

    Source: §1 · p. 2

  • Results

    LeWAM averages 93.5% success versus 50.9% for LeWM (CEM) across ten tasks, and at a 100-step goal offset achieves 46.2% mean success versus 10.9% for LeWM across four tasks.

    Source: §4.2 · p. 6

  • Takeaways & Limitations

    The learned policy supports direct control while its latent dynamics support visual planning, and LeWAM performs competitively with established behavioral-cloning baselines on simulated and physical robots.

    Source: §1 · p. 2, §5 · p. 9

    The task suite uses limited-scale expert demonstrations and non-expert interactions, and models are trained or adapted separately for each task and embodiment.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Existing latent visual world models can retain distractors, while control often relies on costly online search or separately trained action predictors whose supervision is disconnected from representation learning. LeWAM jointly trains representation, action-conditioned dynamics, and a stochastic policy, supporting both direct execution and planning; the paper reports gains in simulated goal-reaching tasks and competitive behavior cloning.

1 INTRODUCTION

  • Existing JEPA pipelines may retain temporally persistent distractors or extraneous observation information rather than control-relevant information.The introduction attributes these risks to JEPA representations and action-conditioned dynamics, respectively.Source: §1 · p. 2
  • CEM can require extensive rollouts as action dimensionality or planning horizon grows, and contact transitions complicate local action optimization.Separately trained action predictors can provide direct control or guide search, but training them on frozen JEPA embeddings disconnects action supervision from visual representation learning.Source: §1 · p. 2
  • LeWAM jointly learns visual representations, action-conditioned dynamics, and a stochastic policy using next-latent prediction, action flow matching, and latent-state regularization.Its Mixture-of-Transformers predictor uses clean actions for dynamics prediction and noisy actions for policy learning, so both objectives train the shared representation.Source: §1 · p. 2
  • Across ten heterogeneous simulated environments, LeWAM improves goal-reaching success over LeWM, particularly on contact-rich tasks.The work also extends comparisons to longer goal horizons and reports faster planning; on goal-agnostic completion in simulation and physical robot execution, it performs competitively with established baselines.Source: §1 · pp. 2–3

2 RELATED WORK

Prior work spans latent and generative world models, JEPA representation learning, and action prediction or search for control. LeWAM combines latent dynamics and policy learning in one end-to-end predictor trained jointly on state prediction and action supervision.

World Models

  • World models support planning by simulating candidate action outcomes, using latent rollouts, reward- and value-grounded dynamics, generative latent dynamics, or discrete visual tokens.LeWAM learns latent dynamics and a policy from offline trajectories without reward supervision or pixel reconstruction.Source: §2 · p. 3

Joint-Embedding Predictive Architectures

  • Action-conditioned JEPAs variously use pretrained representations, learn representations and dynamics from scratch, or add action-prediction objectives; LeWAM jointly trains dynamics and policy in one predictor.Both consequence prediction and action supervision shape LeWAM's shared representation.Source: §2 · p. 3

Control with Latent Dynamics

  • Population-based solvers can require many rollouts as action dimension or horizon grows, while contact-mode changes complicate local optimization.GC-IDM post-trains an action predictor for direct control, whereas PRISM uses a Gaussian prior head on frozen features to guide MPPI sampling; LeWAM jointly trains its world model and flow-based action predictor.Source: §2 · p. 3

World Action Models

  • World action models couple world and action prediction, but UVA, UWM, and Cosmos 3 optimize only one objective per sample rather than both jointly on the same pass.LeWAM trains its encoder and predictor end-to-end with latent-state prediction and action flow matching, using clean actions for forward dynamics and noisy actions for stochastic policy modeling.Source: §2 · p. 3

3 METHODOLOGY

LeWAM jointly trains a visual encoder, action-conditioned latent dynamics, and a flow-matching policy from reward-free offline pixel-observation and action trajectories. It supports direct control and planning by sampling action sequences, predicting latent states, and refining candidates against a goal cost.

3 METHODOLOGY

  • LeWAM jointly trains a visual encoder and a predictor for action-conditioned dynamics and a flow-matching policy from reward-free offline trajectories of pixel observations and actions.The model supports both direct control and planning with action-conditioned latent dynamics.Source: §3 · p. 3

3.1 END-TO-END TRAINING WITH LEWAM

  • A shared ResNet-18 encoder with spatial softmax pooling and an MLP processes history, future-target, and goal images into latent representations.The model uses frameskip K for future dynamics targets and predicts contiguous action chunks of length C, where C ≥ K.Source: §3.1 · p. 4
  • The Mixture-of-Transformers predictor uses shared masked self-attention with separate modality projections; the state query reads history latents and clean actions, while noisy action tokens read history and their causal prefix.Goal information conditions noisy action tokens, and history latents attend only to themselves to block indirect clean-action leakage to the action branch.Source: §3.1 · p. 4, Figure 3 · p. 4
  • Training combines next-latent-state prediction and action flow matching with SIGReg, which regularizes latent states toward an isotropic Gaussian distribution.The losses use mean squared error; SIGReg uses random unit directions and the univariate Epps–Pulley statistic against N(0, 1).Source: §3.1 · p. 4

3.2 EFFICIENT CONTROL WITH LEWAM

  • For direct control, the policy samples an action chunk by integrating its learned velocity field from Gaussian noise, executes the first K actions, then resamples after the next observation.The policy can use fixed history and optional goal context.Source: §3.2 · p. 5
  • For planning, LeWAM alternates sampled action chunks with latent-state predictions to build a population of candidate plans over the requested horizon.At each step, the dynamics branch uses the first K actions of each chunk to predict the next candidate state, which is appended to that candidate's history.Source: §3.2 · p. 5, Figure 4 · p. 5
  • LeWAM refines each candidate by gradient descent through latent dynamics against a goal cost, selects the lowest-cost iterate, and executes the first K actions before replanning.A small proximity penalty discourages large changes from sampled plans; sampling and refinement repeat in receding-horizon fashion until the goal is reached.Source: §3.2 · p. 5

4 EXPERIMENTS

The experiments test LeWAM across varied simulated navigation and manipulation tasks, visual-goal planning, and goal-agnostic behavior cloning. They assess planning accuracy and speed, learned representations and dynamics, and control in simulation and on physical robots.

4.2 VISUAL PLANNING WITH LEWAM

  • At a 100-step goal offset, LeWAM reaches 46.2% mean success across four tasks versus 10.9% for LeWM.The sweep uses the same checkpoints and holds the starting frame fixed on Push-T, Cube, Scene, and Puzzle.Source: §4.2 · p. 6, Figure 7 · p. 7
  • Across four visual-goal tasks, LeWAM averages 52.50 plans/s versus 1.62 plans/s for LeWM (CEM), a 32.3× speedup.The reported mechanisms are a learned expressive action-distribution prior and gradient-based optimization.Source: §4.2 · p. 6
  • Across ten simulated tasks, LeWAM averages 93.5% goal-reaching success versus 50.9% for LeWM (CEM), a 42.7-percentage-point gap.LeWAM exceeds LeWM (CEM) on every task, including ToolHang (94.7 versus 35.4) and Transport (83.3 versus 2.8).Source: §4.2 · p. 6, Figure 6 · p. 7

4.3 UNDERSTANDING LEWAM’S DYNAMICS, REPRESENTATIONS, AND ACTIONS

  • Action modeling raises average proprioception, action, and IDM R2 over LeWM by +0.21, +0.20, and +0.19, respectively.The linear probes measure recovery of physical states, proprioception, and actions from LeWAM representations.Source: §4.3 · p. 7
  • Removing SIGReg lowers planning success on all three reported tasks, most on Push-T (25.3 percentage points), followed by Puzzle (5.3) and Scene (4.0).Without SIGReg, latent-only LeWM has R2 = 0.00 on every reported Push-T probe, while LeWAM with either action head retains decodable state and action information.Source: §4.3 · p. 7
  • Across four tasks, direct success averages 86.7% with the jointly trained encoder, versus 79.0% with frozen CLS features and 75.8% with frozen patch features.The additional gain from planning indicates that the learned dynamics also support action refinement.Source: §4.3 · p. 8
  • Flow matching yields larger planning gains than noisy MSE and retains action-distribution distinctions, while LeWAM’s dynamics remain competitive when planned with CEM.On Cube, LeWAM has clearer control-relevant visual predictions and higher planning success than LeWM; on Push-T, planning success is competitive.Source: §4.3 · p. 8

4.4 BEHAVIOR CLONING AND PHYSICAL ROBOT CONTROL

  • In simulated task completion, LeWAM is competitive with Diffusion Policy on two tasks and improves Transport success by 16 percentage points; on physical tasks, it is competitive with both baselines.With the same demonstrations, its encoder trained from scratch achieves a similar average success rate to the frozen V-JEPA 2.1 variant.Source: §4.4 · p. 9

5 CONCLUSION AND LIMITATIONS

LeWAM jointly predicts future latent states and actions, supporting both planning and policy execution; it improves goal-reaching planning over prior JEPA world models and is competitive on behavioral cloning. The authors identify limits in data coverage, generalization, per-task and per-embodiment training, planning hierarchy, and scale.

5 CONCLUSION AND LIMITATIONS

  • LeWAM jointly predicts future latent states and actions and supports control through both planning and policy execution.The conclusion reports improved dynamics-based planning over previous JEPA-based world models and competitive behavioral cloning against established policy baselines.Source: §5 · p. 9
  • The task suite combines expert demonstrations and non-expert interactions at limited scale, constraining coverage of difficult contact transitions and behaviors.The authors propose broader play data and contact-rich exploration to extend that coverage.Source: §5 · p. 10
  • Models are trained or adapted separately for each task and embodiment, leaving multitask robot and language-conditioned training as future work..Source: §5 · p. 10
  • The authors propose a high-level planner to decompose longer tasks into subgoals executed through LeWAM's low-level planning and control.Source: §5 · p. 10
  • The visual-planning model has approximately 17M parameters; larger models and datasets, candidate populations, and refinement budgets are proposed to extend scale.Source: §5 · p. 10
  • Physical deployment requires independent safety mechanisms, and the experiments collect no new human-subject data.Source: ETHICS STATEMENT · p. 10

A IMPLEMENTATION DETAILS

The implementation uses a shared visual encoder and a Mixture-of-Transformers predictor for latent dynamics and action flow matching. Planning samples action sequences from the learned policy and refines them against learned dynamics.

A IMPLEMENTATION DETAILS

  • Dynamics predict a latent state at stride K = 5 from five actions, while flow matching models contiguous chunks of C = 10 actions.The history contains two latents; the history stride is K for goal-reaching and 1 for task-completion, and imagined rollouts retain the two latest latents.Source: A IMPLEMENTATION DETAILS · p. 15
  • The visual encoder maps images to 192-dimensional latents, while a four-block, four-head MoT predictor uses separate state and action readouts.The encoder uses a ResNet-18-style backbone, spatial softmax pooling, and a projection MLP; the predictor has width 192 and dropout 0.1.Source: A IMPLEMENTATION DETAILS · p. 15
  • Optional goal conditioning supplies a goal-image latent and normalized remaining horizon to noisy action tokens through AdaLN-Zero.Goal conditioning is omitted for goal-free task completion.Source: A IMPLEMENTATION DETAILS · p. 15
  • The visual-planning model has approximately 16.9M trainable parameters, and its reported objective configuration uses SIGReg with λsig = 0.04.Simulated action sampling uses eight Euler steps; the MSE dynamics configuration uses an online encoder target and no state noise or state-flow integration.Source: A IMPLEMENTATION DETAILS · p. 15
  • Goal-conditioned runs train for 50 epochs and goal-free task-completion runs for 120, selecting checkpoints by validation action flow-matching loss rather than task success.Both use AdamW and a warmup followed by an epochwise cosine learning-rate schedule.Source: A IMPLEMENTATION DETAILS · p. 15
  • The attention mask assigns visibility by query and blocks clean action tokens from relaying target information to history or action-flow tokens.The figure depicts queries as rows, keys as columns, and distinguishes permitted attention from masked cells.Source: A IMPLEMENTATION DETAILS · p. 16, Figure 10 · p. 16

A.1 LEWAM PLANNING

  • Planning samples candidate sequences from the learned policy, refines them by gradient descent through frozen dynamics, and selects the lowest-cost iterate.Only normalized action plans receive gradients; a proximity penalty discourages large changes from sampled plans, and actions are clipped after updates.Source: §A.1 · p. 17
  • The refinement procedure uses Adam with learning rate 0.05, gradient-norm clipping at 10, amax = 3, and λprox = 0.01.Source: §A.1 · p. 17

A.2 BASELINE IMPLEMENTATIONS

The baselines cover latent-dynamics planning, a post-trained inverse-dynamics policy, and direct goal-conditioned policies. Their implementations retain or adapt released configurations, with shared comparison conditions where specified.

A.2 BASELINE IMPLEMENTATIONS

  • Goal-conditioned comparisons share observations, actions, start-goal pairs, and success predicates, while baseline implementations generally follow their official configurations.The baselines use the released stable-worldmodel codebase and retain its environment wrappers, dataset interface, action bounds, and evaluation loop unless noted otherwise.Source: §A.2 · p. 17
  • LeWM (IDM) freezes a pretrained LeWM checkpoint before post-training a goal-conditioned inverse-dynamics model; LeWM (CEM) uses official LeWM training and CEM evaluation settings.Source: §A.2 · p. 17
  • DINO-WM comparisons include the released patch-token model and a matched CLS-token variant differing only in tokenization.The six added tasks use ten training epochs, while the original four retain four-epoch checkpoints.Source: §A.2 · p. 17
  • GCBC uses a frozen pretrained DINOv2 encoder and executes predicted five-step action chunks without planning at evaluation.Its training duration is 100 epochs, and short-horizon goals use a uniformly later frame from the same episode.Source: §A.2 · pp. 17–18
  • The task-completion simulation baseline is trained on demonstrations conditioned on observation history, while Diffusion Policy predicts 16 actions and executes eight before resampling.The Diffusion Policy inference procedure uses 100 DDPM steps.Source: §A.2 · p. 18
  • The described policy training uses 224 × 224 RGB inputs cropped to 202 × 202, a 90/10 training-validation split, and checkpoint selection by validation denoising loss.It trains for 120 epochs with batch size 64 and evaluates checkpoints over three seeds with 50 episodes per seed.Source: §A.2 · p. 18

B.1 TASKS AND DATASETS

The benchmark combines ten varied simulated tasks and behavior-cloning datasets with goal-reaching and task-completion evaluations. The appendix specifies success criteria, goal horizons, planner sensitivity measurements, and timing comparisons.

B.1 TASKS AND DATASETS

  • The ten-task suite retains four LeWM environments and adds six tasks spanning navigation, object interaction, coupled button states, tool placement, bimanual transfer, and drawers.The tasks vary in action dimensions, camera views, trajectory lengths, and data-collection regimes.Source: §B.1 · p. 18
  • Datasets include noisy or random interactions, scripted expert data, and human or generated behavior-cloning demonstrations across the listed environments.ToolHang uses human demonstrations; Transport and DrawerCleanup use generated demonstrations.Source: §B.1 · pp. 18–19
  • The three behavior-cloning datasets broaden comparison to contact-rich, higher-dimensional behaviors, but task-focused demonstrations may limit diversity of sampled visual goals.Table 3 summarizes views, raw action dimensions, dataset scale, and recorded lengths for the ten tasks.Source: §B.1 · p. 19, Table 3 · p. 19

B.2 GOALS AND SUCCESS CRITERIA

  • Goal-reaching evaluation uses a goal 25 steps ahead and a 50-action budget, except Two-Room, which uses a 100-step offset and 150 actions.The horizon sweep compares recorded goals 25, 50, 75, or 100 steps ahead.Source: §B.2 · p. 19, §B.3 · p. 21
  • The original Cube object-only criterion credited zero action on 42.7% of 25-step episodes; adding an end-effector check reduced that floor to 2.7%.The corrected predicate requires both cube position and end-effector position to be within 0.04 m of their goal-frame positions; Figure 11 illustrates a spurious object-only success.Source: §B.2 · pp. 19–20, Figure 11 · p. 20

B.3 PLANNING AND TIMING PROTOCOLS

  • Planner sensitivity varies refinement updates and sampled-candidate counts on Push-T, measuring success and latency; timing compares 300 × 30 CEM with five LeWAM gradient steps on the same H100 GPU.Figure 1(c) averages throughput across four tasks; Table 6 reports time per planner call.Source: §B.3 · p. 21

C.1 VISUAL PLANNING RESULTS

The visual-planning appendix reports planner sensitivity and component evaluations, compares learned dynamics and action distributions, and documents additional simulation and physical-robot results.

C.1 VISUAL PLANNING RESULTS

  • Success plateaus after five refinement updates, while latency rises; sampling more than eight candidates yields no consistent gain in this Push-T sweep.Source: §C.1 · p. 21, Figure 12 · p. 22

C.2 COMPONENT AND VISUAL-ENCODER ABLATIONS

  • The ablations compare direct policy execution with test-time planning across SIGReg, action-head, and visual-encoder variants, using epoch-30 checkpoints for the component comparisons.Source: §C.2 · p. 22, Table 7 · p. 22, Table 8 · p. 23

C.3 PROBING LATENT REPRESENTATIONS

  • The probing table reports linear-probe R2 for object, proprioceptive, and action information across four tasks, alongside effective rank and inverse-dynamics probes.Source: §C.3 · p. 23, Table 9 · p. 24

C.4 VISUAL RECONSTRUCTIONS AND OPEN-LOOP ROLLOUTS

  • LeWAM exceeds LeWM by 22.2 percentage points on Cube and achieves competitive Push-T success; decoded rollouts show clearer robotarm, gripper, and object detail on 3D tasks.Source: §C.4 · p. 24

C.5 MULTIMODAL ACTION DISTRIBUTIONS

  • On the two-door task, flow matching follows both demonstrated routes, whereas the deterministic MSE head concentrates on the lower route and sometimes collides with a wall.Source: §C.5 · p. 27, Figure 15 · p. 27

C.6 SIMULATION TASK-COMPLETION RESULTS

  • The task-completion table reports per-task mean success rates and standard deviations across three seeds, with 50 episodes per seed.Source: §C.6 · p. 27, Table 11 · p. 27

D PHYSICAL-ROBOT SETUP AND EVALUATION

D.1 PLATFORM AND DATA COLLECTION

The physical-robot evaluation uses teleoperated demonstrations, two-camera observations, matched action-chunk execution, and repeated human-labeled trials.

D.1 PLATFORM AND DATA COLLECTION

  • Successful demonstrations are collected through joint-space teleoperation from a reBot Arm 102 leader to a reBot B601-DM follower, with impedance control and gravity feedforward.Source: §D.1 · p. 27

D.2 MODEL IMPLEMENTATION

  • The robot models use two camera streams; the end-to-end variant trains its ViT-B visual encoder from scratch, while the comparison variant freezes a V-JEPA 2.1-initialized encoder.Source: §D.1 · p. 27, §D.2 · p. 28

D.3 PHYSICAL-ROBOT BASELINE CONFIGURATIONS

  • Physical-robot baselines are adapted to the same cameras and seven-dimensional joint-target actions, returning 36 actions and executing the first 18 before requesting a new chunk.Source: §D.3 · p. 28

D.4 CLOSED-LOOP EVALUATION

  • Each method is evaluated in 30 trials per task, with a human evaluator labeling each episode as success or failure under the teleoperation criteria.Source: §D.4 · p. 29

D.5 REPRESENTATION DIAGNOSTICS AND EVALUATION

Representation diagnostics decode frozen latent states into frames, while the discussion relates action supervision and latent prediction to control-relevant representations.

D.5 REPRESENTATION DIAGNOSTICS AND EVALUATION

  • A post-hoc pixel decoder is trained on frozen encoder outputs; framewise reconstructions decode z_t, and rollout predictions decode latent states advanced under recorded actions.Source: §D.5 · p. 29, Figure 16 · p. 30

E DISCUSSION OF THE ROLE OF JOINT PREDICTION

E.1 REPRESENTATION AND DYNAMICS

  • When visual history changes conditional action velocity, a constant code raises optimal action-regression risk; action supervision favors those distinctions, while SIGReg regularizes the global latent distribution.Source: §E.1 · pp. 31–32

E.2 PLANNING WITH LEARNED ACTION PROPOSALS

The discussion frames contact-rich planning as difficult for broad-prior search and presents learned action proposals with predicted consequences as a heuristic response.

E.2 PLANNING WITH LEARNED ACTION PROPOSALS

  • The figure contrasts smooth rollouts with candidate rollouts that branch at contact, then depicts learned proposals paired with predicted outcomes.Source: Figure 18 · p. 32
  • Contact events can switch dynamics modes, making nearby actions yield distinct outcomes; if sampled plans miss useful contact, a goal cost may provide little ranking signal.Source: §E.2 · p. 32
  • LeWAM combines learned action proposals with predicted consequences to address contact-rich planning under constrained budgets, but the authors describe the solution as heuristic and leave rigorous verification for future work.Source: §E.2 · p. 32, Figure 18 · p. 32

Related papers

Loading 2610.12818v1…