Source-linked AI summary
Making Latent Evolution Explicit: Operator-Structured Transitions for World Action Models
Xiaoxiao Lu, Yunlong Dong, Jiahao Shi, Ye Yuan
TL;DR
Latent WAM transitions commonly use Transformer predictors whose inductive structure emphasizes token interaction rather than explicit temporal evolution. LEON introduces context-modulated operator propagation and additive forcing in learned observable coordinates, and it improves or preserves closed-loop performance across distinct WAM policy couplings. The results support treating transition realization as a consequential architectural choice.
Problem
Latent WAMs increasingly predict in latent space, but their Transformer-based transitions do not explicitly represent temporal, action-conditioned evolution.
Method
LEON models latent evolution in learned observable coordinates using context-modulated operators with shared basis components and a complementary additive forcing path.
Results
Across VLA-JEPA and LaWAM, LEON improves or preserves closed-loop performance, including 99.05% average LIBERO success and 80.6% aggregate LIBERO-Plus success.
Takeaways & Limitations
Transition realization is a first-class architectural dimension in latent WAMs alongside predictive representation and prediction–policy coupling.
Takeaways & Limitations
Controlled results do not imply that learned WAM representations follow corresponding physical dynamics or admit an exact Koopman description.
Abstract
from arXiv · showhide
World Action Models (WAMs) augment robot policies by predicting how task-relevant scene states may evolve under interaction. Recent WAMs increasingly perform such prediction in latent representation spaces, avoiding full appearance-level generation while preserving control-relevant information. Yet latent transitions are commonly realized with Transformer-based predictors whose inductive structure is centered on token interaction rather than temporal evolution. We study transition realization as an architectural choice distinct from predictive representation and prediction-policy coupling. We introduce the Latent Evolution Operator Network (LEON), which models latent evolution in a learned observable space through context-modulated operator-based propagation and additive forcing. Grounded in the controlled Koopman generator view of evolution, LEON organizes context-dependent transition variation around a shared evolution-operator structure while retaining a complementary path for additive change. Controlled dynamical systems verify the resulting evolution-specific inductive bias and the complementary roles of operator propagation and forcing. Across two WAM formulations that integrate latent prediction into the policy differently, LEON improves closed-loop performance and robustness while remaining effective under full transition replacement. These results establish transition realization as a consequential architectural choice in latent WAMs.
1 Introduction
Latent transition realization is a distinct WAM architectural choice: LEON makes evolution structure explicit through context-modulated operators and additive forcing, improving or preserving closed-loop performance across policy couplings.
- Controlled dynamical systems test LEON's evolution-specific inductive bias through extrapolation, nonlinear dynamics, and interventions on operator and forcing branches.Structured models outperform a Transformer baseline beyond the training range, with the advantage persisting for nonlinear pendulum dynamics.
- VLA-JEPA + LEON reaches 99.05% average success on LIBERO, improving the corresponding Transformer realization by 1.85 percentage points.
- On LIBERO-Plus, LEON achieves 80.6% aggregate success versus 79.5% for VLA-JEPA, while LaWAM replacement preserves near-baseline RoboTwin performance at 84.13% versus 84.50%.
- Latent transition realization is distinct from predictive representation and prediction–policy coupling, making temporal evolution an explicit modeling object.
- LEON maps visual latents into learned observables and uses context-modulated operator propagation plus additive forcing to model evolution.The operator basis is shared across contexts, while context-dependent coefficients modulate its composition.
2 Background and Related Work
WAMs use latent prediction to provide predictive context without full image generation, but commonly realize transitions with generic Transformer token interactions rather than explicit evolution structure.
- WAMs augment robot policies with predicted future scene information, while latent prediction avoids reconstructing full visual observations and appearance-level redundancy.
- Latent WAMs differ in whether future prediction shapes policy representations during training or is directly consumed by the action generator at inference.VLA-JEPA exemplifies the representation-mediated design.
- Transformer-based predictors jointly transform current-state and conditioning tokens to produce future latents through attention and feed-forward mixing.
- Generic token transformations do not explicitly single out temporal, action-conditioned state evolution as a distinct modeling object.The future latent represents how the current state evolves under action and semantic context, not merely another prediction target.
- Locally linear latent models and deep Koopman methods make temporal evolution more tractable through direct propagation or linear operator evolution in learned coordinates.
- Controlled Koopman formulations allow controls to modulate evolution operators, whereas LEON applies this principle to learned semantic and action-related WAM context.
3 Latent Evolution Operator Network
LEON makes latent transition realization explicitly evolution-structured: learned context modulates a shared operator family in observable space, while additive forcing captures complementary change. The design adapts controlled Koopman generator principles to latent WAMs and preserves the existing predictive representations and policy couplings.
- Transition realization: LEON treats transition realization as a distinct architectural choice from predictive representation and prediction–policy coupling.Its intervention changes the form of Tθ while preserving the predictive representation, future target, conditioning pathway, and downstream policy interface.
- Learned observables and context: LEON maps visual latents into learned observable coordinates and constructs context from semantic and action-related conditioning.The resulting context modulates how the observable representation evolves.
- Controlled Koopman basis: The controlled Koopman generator decomposes evolution into shared drift and control-induced components, motivating context-dependent operator modulation in LEON.In latent WAMs, learned transition context replaces explicit physical control while retaining the shared-basis modulation principle.
- Context-modulated evolution: LEON combines context-modulated operator propagation with additive forcing to separate multiplicative state evolution from flexible context-dependent change.The operator acts directly on the current observable, whereas forcing adds change not constrained to be multiplicative.
- Structured evolution operator: Context-dependent operator variation lies in a compact shared low-rank subspace, with fixed basis operators and a context-independent baseline.The factor width r upper-bounds the number of independent operator directions available for transition variation.
- Instantiation across couplings: LEON composes observable-space updates with a learned readout and is instantiated in both representation-mediated VLA-JEPA and policy-facing LaWAM.Both systems replace the Transformer transition while preserving their predictive representation and prediction–policy coupling.
4 Experiments
Replacing Transformer-based transition realizations with LEON improves or preserves closed-loop WAM performance across distinct prediction–policy couplings. Controlled studies further show stronger extrapolation and that operator propagation and additive forcing contribute differently to structured transitions.
- WAM performance: VLA-JEPA + LEON achieves the highest average success among methods listed in the LIBERO comparison.It reports the strongest performance on Goal while remaining near ceiling on Spatial and Object.
- WAM performance: 84.13% versus 84.50% aggregate success shows that full transition replacement preserves near-baseline LaWAM performance under matched four-task training.LEON improves clean Hanging Mug success from 51% to 55% while retaining strong performance on other tasks.
- Prediction–policy coupling: LEON remains effective when predicted future latents directly condition action generation, not only when they shape policy-relevant representations during training.This conclusion follows from the LaWAM replacement result and the distinct VLA-JEPA coupling.
- Controlled dynamics: 0.0554 versus 3.0315 energy MAE shows Adaptive Koopman substantially improves long-horizon physical fidelity over Fixed Koopman under amplitude shift.The corresponding energy-balance residual decreases from 0.0223 to 0.0018.
- Controlled dynamics: Removing operator propagation raises OOD rollout RMSE to 10.95× the full-model level, versus 3.03× after removing additive forcing.The intervention identifies operator propagation as the dominant contribution and forcing as complementary.
5 Discussion and Conclusion
The paper frames transition realization as a distinct architectural choice in latent WAMs and proposes LEON’s evolution-specific operator structure. LEON remains effective across prediction–policy couplings, while controlled studies clarify its intended inductive behavior and scope.
- Transition realization is a distinct architectural dimension in latent World Action Models, alongside predictive representation and prediction–policy coupling.
- LEON organizes latent transitions in learned observable coordinates using context-modulated operator propagation and complementary additive forcing.
- 97.2% to 99.05%: full transition replacement with LEON increases average LIBERO success in VLA-JEPA.On policy-facing LaWAM, replacement preserves near-baseline matched four-task performance at 84.13% versus 84.50%.
- Controlled studies indicate that explicit evolution structure supports unchanged-law extrapolation and nonlinear dynamics, while weakening across qualitative regime changes.Operator propagation provides the primary contribution, with additive forcing contributing complementarily.
- The controlled results do not establish that learned WAM representations follow physical dynamics or admit an exact Koopman description.