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

LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models
Sep 28, 2026Luzhe Huang, Lei Chu, Jingyi Liang +1Latent World Models Low-Latency Latent Planning
Independent researcher
PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models
Aug 6, 2026Xi Zeng, Haojie Ren, Ziying Song +2Latent World Models Joint-Embedding Predictive Architectures
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
On the Identifiability of Controlled World Models
Jul 24, 2026Xiangteng Zhang, Yang Guan, Bo Zhang +3Latent World Models Controllability
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

Sep 28, 2026, cs.LGJ/K move · Enter open · S save

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

Luzhe Huang, Lei Chu, Jingyi Liang, Yuhuan Zhao

Independent researcher

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

On the Identifiability of Controlled World Models