Prismatic World Model: Learning Compositional Dynamics for Planning in Hybrid Systems
This addresses the challenge of unreliable long-horizon planning for robotic systems with hybrid dynamics, offering a foundational model for next-generation agents, though it is incremental as it builds on existing Mixture-of-Experts frameworks.
The paper tackled the problem of model-based planning in robotic domains with hybrid dynamics, where conventional models over-smooth discrete events, by introducing the Prismatic World Model (PRISM-WM) that decomposes dynamics into composable primitives, resulting in significantly reduced rollout drift and superior performance in high-dimensional benchmarks.
Model-based planning in robotic domains is fundamentally 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, inevitably over-smoothing the distinct dynamic modes (e.g., sticking vs. sliding, flight vs. stance). For a planner, this smoothing results in catastrophic 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 leverages 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, effectively preventing mode collapse. By accurately modeling the sharp mode transitions in system dynamics, PRISM-WM significantly reduces rollout drift. Extensive experiments on challenging continuous control benchmarks, including high-dimensional humanoids and diverse multi-task settings, demonstrate that PRISM-WM provides a superior high-fidelity substrate for trajectory optimization algorithms (e.g., TD-MPC), proving its potential as a powerful foundational model for next-generation model-based agents.