Bilinear World Models: Learning Representations with Structured Dynamics for Efficient Control

Antonio Pariente; Ignacio Boero; Nikolai Matni; Alejandro Ribeiro

cs.ROSep 28, 2026

Bilinear World Models: Learning Representations with Structured Dynamics for Efficient Control

Authors: Antonio Pariente, Ignacio Boero, Nikolai Matni, Alejandro Ribeiro

Organizations: University of Pennsylvania

Abstract

World models jointly learn latent representations and dynamics that predict how high-dimensional observations evolve under actions. In this work, we propose a JEPA-style world model in which, rather than learning arbitrary latent dynamics, we restrict them to follow a bilinear parameterization. This structure enables efficient planning and control while shifting the modeling burden onto the encoder, encouraging richer representations that expose the controllable geometry of the system. In particular, this structured parameterization allows us to structurally enforce action recoverability, thereby preventing representation collapse by construction. Although prescribing a bilinear parametrization may appear restrictive, we show that a broad class of nonlinear dynamical systems admits a transformation under which the dynamics become bilinear. Empirically, we show across standard 2D and 3D control tasks that representations with bilinear-parameterized dynamics can be learned directly from high-dimensional observations, reducing planning time by nearly three orders of magnitude while retaining or even improving control accuracy. We also propose more demanding regimes of longer-horizon planning and real-time control, and demonstrate that our method succeeds in both, moving JEPA-style world models beyond short-horizon offline planning.

Figures & tables

Figure 1: Planning success and time across four benchmarks. Solid bars report success rate over 100 evaluation episodes, each defined by a start–goal pair with a goal offset of 25 steps; hatched bars report the average total wall-clock planning time per episode on a log scale. Our Bilinear World Model matches or improves success across all environments while reducing planning time by two to three orders of magnitude.

Panel of Figure 2 — Figure 2: Geometry of the learned latent representations. (Top) Variance explained by each principal component of the latent embeddings, with the action dimension indicated by the red dashed line. (Middle) Comparison between selected physical state variables and the corresponding projections of the latent embeddings onto their principal components. In TwoRoom, PC1 and PC2 are shown against the agent’s horizontal and vertical positions. In PushT, Reacher, and Cube, principal-component projections are shown alongside angular state variables represented through their sine and cosine coordinates. (Bottom) Visualizations of the corresponding environments.

Panel of Figure 3 — Figure 3: Trajectories and action smoothness in TwoRoom. (Left) Average jerk, measured from changes between consecutive actions, for each world model. (Right) Representative planned trajectories toward the goal. Bilinear World Models produce lower-jerk action sequences which result in smoother trajectories.

Panel of Figure 4 — Figure 4: Performance across increasing task horizons. Each panel reports success rate (solid bars) and planning time (hatched bars) for goal offsets of 25, 50, and 75 steps, denoted as @25, @50, and @75. Bilinear World Models exhibit a smaller degradation as the horizon increases while maintaining substantially lower planning times.

Figure 5: Control performance on moving-goal tasks. Success rate on M-TwoRoom, M-PushT, M-Reacher, and M-Cube under a 20 Hz control loop. The target follows a random walk, requiring repeated replanning under controller latency. SMWM and LeWorld use iCEM ( Pinneri et al., 2021 ) , while Bilinear World Models use the same Gauss–Newton controller as in the other experiments. The Random baseline applies uniformly sampled actions.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Task	$d_{a}$	$H$	$K$	$n_{b}$	$K_{c}$	iters	$n_{\text{cand}}$	$\rho$
PushT	2	5	50	5	10	40	16	0.5
Reacher	2	3	30	3	6	20	64	2.0
TwoRoom	2	5	50	3	6	40	0	—
Cube	5	5	125	3	15	20	0	—

Appendix

Table 1: Gauss–Newton planner settings for the reported cells.

H

: planning horizon in action blocks;

m=k\,d_{a}

with block size

k=5

;

K=Hm

raw unknowns,

K_{c}

coherent unknowns;

n_{b}

: time blocks refreshed per iteration. Shared across all tasks:

\lambda_{g}=1

c=0.01

\mu=10^{-4}

\varepsilon=10^{-3}

, line search

\{1,\tfrac{1}{2},\tfrac{1}{4},\tfrac{1}{8}\}

n_{\text{starts}}=1

, tolerance

10^{-6}

Panel of Figure 6 — Figure 6: Training dynamics of Bilinear World Models. The latent forward-prediction loss decreases throughout training, while the minimum singular value $\sigma_{\min}(R)$ remains bounded away from zero. Since $M(z)=Q(z)R$ with $Q(z)$ orthonormal, $\sigma_{\min}(R)=\sigma_{\min}(M(z))$ for every latent state $z$ . The learned action map therefore remains full column rank throughout training.

Panel of Figure 8 — Figure 8: Controller performance across world models. Success rate and planning time are reported for CEM, iCEM, and Gauss–Newton (GN) applied to Bilinear World Models, SMWM, and LeWorld. GN is particularly effective when paired with Bilinear World Models, where it can exploit the explicit structure of the latent dynamics. CEM remains applicable across all architectures but incurs substantially higher planning cost, while iCEM reduces computation at the expense of degraded performance on the more challenging tasks.

Panel of Figure 10 — Figure 10: PCA visualization of the learned latent representations. Latent states are projected onto their first three principal components. TwoRoom, PushT, and Reacher exhibit approximately low-dimensional surface structure, while Cube displays a richer geometry that is not fully captured by a three-dimensional projection.

Explore similar work

Sep 28, 2026cs.LG

LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models

Compact JEPA world models enable efficient latent-space planning, but low-dimensional representation trained under reward-free self-supervision must encode both action-conditioned dynamics and predictable visual context. This competition can entangle controllable state with high-rank nuisance appearance and degrade planning as scenes become more complex. We introduce LRC-JEPA, a lightweight end-to-end world model that routes information into a compact predictive latent

\mathbf{z}

and learned-query residual-context embeddings

\mathbf{u}

. Only

\mathbf{z}

is propagated by the dynamics model and used for planning, while

\mathbf{u}

captures temporally persistent information for cross-attention reconstruction; a differentiable residual connection encourages the latent to retain complementary dynamic content. Under explicit assumptions, we show that the resulting representation is sufficient, minimal, nuisance-invariant, and disentangled. Across four simulated control environments, LRC-JEPA improves average planning success over a parameter-matched JEPA baseline by 9 percentage points and matches or exceeds substantially larger pretrained models. On the real-world Bridge-v2 set, its 5.5M-parameter active encoder outperforms DINO-WM (22.1M) and V-JEPA2 (303.9M) encoders while also enabling faster planning. Physical-state probes, reconstruction interventions, and ablations confirm the effectiveness of LRC-JEPA's representation disentanglement.

Luzhe Huang, Lei Chu, Jingyi Liang +1

Independent researcher

Latent World Models Low-Latency Latent Planning S-Jepa World

Ask PDF

Aug 6, 2026cs.CV

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

We propose PhyLatent, a dynamics-relevant training objective for Joint-Embedding Predictive Architecture (JEPA) world models. Our key observation is that preventing global latent collapse does not necessarily ensure that the learned representation preserves physically meaningful state and action relationships. We identify three failure modes: sensitivity to appearance changes that leave the physical state unchanged, insufficient separation of distinct physical states, and insufficient separation of different action-conditioned futures. We refer to these as Physical Invariance Collapse, Physical Distinguishability Collapse, and Counterfactual Dynamics Collapse, respectively. PhyLatent targets these failures through three coordinated training pathways, implemented with static visual invariance, physical state grounding, future representation alignment, counterfactual branch separation, and latent denoising. On OGBench-Cube, PhyLatent reduces the three collapse rates by 43.9%, 27.9%, and 47.1%, respectively, while improving model predictive control (MPC) success by 12.0 percentage points (17.2% relative). Across four visual-control tasks, average planning success increases by 6.62 percentage points (8.3% relative). These results show that global non-collapse alone is insufficient for learning a reliable JEPA world-model state space, and that explicitly preserving dynamics-relevant structure can improve closed-loop planning.

Xi Zeng, Haojie Ren, Ziying Song +2

School of Mechanical and Aerospace Engineering, Nanyang Technological University, Singapore · The University of Sheﬃeld, Sheﬃeld, United Kingdom · School of Artiﬁcial Intelligence (School of Software), Yanshan University, Qinhuangdao, China

Latent World Models Joint-Embedding Predictive Architectures World Models Inference Dynamics Latent Variable

Ask PDF

Jul 24, 2026cs.LG

On the Identifiability of Controlled World Models

World model serves as a promising tool to infer environment dynamics under high-dimensional observations and candidate actions. Recently, LeCun's JEPA provides a compelling framework for learning such models in representation space. Its action-conditioned extension plays a central role in visual control and latent-space planning, but leaves a fundamental question: can it recover the controlled dynamics from nonlinear observations? This paper presents a joint identifiability condition for controlled world models with Gaussian latent states, which consists of two coupled components: (1) representation identifiability and (2) transition identifiability. The former depends on the spectral separation property while the latter is related to non-degenerate variation of conditional action. We prove that when this condition holds, minimizing the LeJEPA-style predictive objective can recover both latent states and controlled dynamics in the sense of orthogonal transformation. We further prove that the upper bound of transition prediction error is inversely proportional to the spectral separation margin. We also characterize an attainable amplification of counterfactual prediction error that scales inversely with the weakest conditional action-excitation margin. The theoretical predictions are empirically supported across four nonlinear observation settings.

Xiangteng Zhang, Yang Guan, Bo Zhang +3

School of Vehicle and Mobility, Tsinghua University · 2Didi Voyager Labs, Didi Autonomous Driving · School of Computing and Data Science, The University of Hong Kong +2

Latent World Models Controllability World Models Identifiability

Ask PDF