Bilinear World Models: Learning Representations with Structured Dynamics for Efficient Control
Abstract Overview
This paper proposes a JEPA-style world model that constrains latent dynamics to follow a bilinear parameterization, with a normalized action map that supports efficient planning and control. The design shifts more of the modeling burden onto the encoder, organizing latent representations around controllable aspects of the system. Grounded in Koopman operator theory, the authors argue that action-affine nonlinear dynamical systems admit bilinear latent dynamics with action recoverability. Empirically, the approach is evaluated on four 2D and 3D control benchmarks from high-dimensional observations, as well as on extended-horizon and real-time moving-goal settings.
Novelty
The primary novelty is restricting latent dynamics in a JEPA-style world model to a normalized bilinear parameterization, rather than employing an unconstrained neural network predictor. This structured action map structurally enforces action recoverability to prevent representation collapse by construction and enables efficient sensitivity-based planning via a Gauss-Newton controller.
Results
Across the TwoRoom, PushT, Reacher, and Cube benchmarks, the method matches or exceeds baseline success rates (100%, 90%, 75%, and 99%, respectively) while reducing planning time by two to three orders of magnitude. It also shows greater resilience over longer planning horizons and outperforms baselines in all four moving-goal tasks under a 20 Hz real-time replanning loop with an average 27-percentage-point gain.
Key Points
- The model constrains latent dynamics to a normalized bilinear parameterization, enforcing closed-form action recoverability and preventing representation collapse by construction.
- Theoretical analysis connects action-affine nonlinear systems to action-recoverable bilinear Koopman embeddings, motivating the imposed latent structure.
- A sensitivity-based Gauss-Newton planner leverages the differentiable dynamics to achieve high control success while cutting planning times from minutes to seconds in offline and real-time regimes.