Prismatic World Model: Learning Compositional Dynamics for Planning in Hybrid Systems

TL;DR

The paper introduces PRISM-WM, a structured world model that decomposes hybrid system dynamics into primitives using a Mixture-of-Experts framework, improving long-horizon planning accuracy.

Contribution

It proposes a novel architecture with a gating mechanism and expert diversity objective to better model hybrid dynamics for planning tasks.

Findings

PRISM-WM reduces rollout drift in hybrid dynamics.

It enhances trajectory optimization in high-dimensional control tasks.

The model improves planning reliability over monolithic neural networks.

Abstract

Model-based planning in robotic domains is challenged by the hybrid nature of physical dynamics, where continuous motion is punctuated by discrete events such as contacts and impacts. Conventional latent world models typically employ monolithic neural networks that enforce global continuity, which over-smooths distinct dynamic modes (e.g., sticking vs. sliding, flight vs. stance). For a planner, this smoothing results in compounding errors during long-horizon lookaheads, rendering the search process unreliable at physical boundaries. To address this, we introduce the Prismatic World Model (PRISM-WM), a structured architecture designed to decompose complex hybrid dynamics into composable primitives. PRISM-WM uses a context-aware Mixture-of-Experts (MoE) framework where a gating mechanism implicitly identifies the current physical mode, and specialized experts predict the associated transition dynamics. We further introduce a latent orthogonalization objective to ensure expert diversity, preventing mode collapse. By modeling the mode transitions in system dynamics, PRISM-WM reduces rollout drift. Experiments on continuous control benchmarks, including high-dimensional humanoids and multi-task settings, demonstrate that PRISM-WM provides a high-fidelity substrate for trajectory optimization algorithms (e.g., TD-MPC), indicating its potential as a foundational model for model-based agents.