Technical Summary: Prismatic World Model (PRISM-WM)

1. Problem Statement

Model-Based Reinforcement Learning (MBRL) and planning in robotic domains are fundamentally limited by the hybrid nature of physical dynamics. Real-world systems, particularly in legged locomotion and manipulation, involve continuous motion punctuated by discrete events (e.g., contacts, impacts, sticking vs. sliding). These interactions create non-smooth, highly non-linear dynamics with sharp discontinuities in the underlying vector field.

Conventional latent world models typically employ monolithic neural networks that enforce global continuity. While theoretically capable of universal approximation, these models practically learn an averaged representation that over-smooths distinct dynamic modes. For trajectory optimization planners, this over-smoothing induces compounding errors during long-horizon lookaheads, particularly at physical boundaries. This results in:

• Spectral blurring across contact boundaries, causing significant gradient mismatch.
• Compounding drift in high-dimensional configuration spaces, rendering long-term planning unreliable.
• Policy collapse in contact-rich environments where standard models fail to distinguish between stable and unstable regimes.

2. Methodology: The Prismatic World Model (PRISM-WM)

To address these limitations, the authors introduce PRISM-WM, a structured architecture designed to decompose complex hybrid dynamics into composable primitives. The model is built upon a context-aware Mixture-of-Experts (MoE) framework with a novel latent orthogonalization objective.

Core Architecture

• Context-Aware Gating: A gating network ($G$) receives a context vector $c_t$ (comprising state $z_t$, action $a_t$, and optionally task embedding $e_{task}$) and produces softmax-normalized weights $w_t$. This mechanism implicitly identifies the active physical regime (e.g., flight vs. stance) without explicit supervision.
• Specialized Experts: The state is routed to a set of $K$ expert networks $\{E_k\}$. Each expert specializes in learning the continuous flow within a specific local dynamic regime, avoiding interference from conflicting modes.
• Residual Prediction: The next latent state is predicted as a weighted combination of expert outputs added to the current state:
  $$\hat{z}_{t+1} = z_t + \sum_{k=1}^{K} w_{t,k} \cdot v_k(x_t)$$
  where $v_k$ represents the residual dynamics predicted by the $k$-th expert.

Key Innovation: Latent Orthogonalization

To prevent mode collapse (where experts learn redundant features) and ensure diversity, PRISM-WM introduces a differentiable Gram-Schmidt orthogonalization layer.

• After each expert outputs a penultimate-layer vector $u_k$, the model applies a stabilized Gram-Schmidt process to generate orthogonal basis vectors $v_k$.
• This enforces the expert representations to lie on a Stiefel manifold, ensuring they span a maximal-rank, non-redundant subspace for the residual dynamics.
• While physical contact modes are unordered, this orthogonalization ensures the experts capture distinct components of the residual manifold, allowing the gating network to reconstruct unordered physical modes via a weighted sum.

Integration with Planning

PRISM-WM is integrated into the TD-MPC planning framework. By factorizing complex dynamics into non-redundant experts, the model reduces the one-step model residual from $\epsilon$ to approximately $\epsilon/K\alpha$. This structural decomposition mitigates compounding drift, enabling high-fidelity long-horizon rollouts and stable gradient propagation for policy learning.

3. Key Contributions

1. Architectural Decomposition: Introduction of PRISM-WM, which decomposes hybrid dynamics using a context-aware MoE architecture coupled with latent orthogonalization, moving away from monolithic global approximations.
2. Orthogonalization Objective: A novel training objective using differentiable Gram-Schmidt orthogonalization to enforce expert diversity and prevent mode collapse, structurally aligning the latent space with discrete physical transitions.
3. Theoretical and Empirical Validation: Demonstration that this decomposition reduces compounding errors in high-dimensional spaces and improves the geometry of the loss landscape for both planning (TD-MPC) and direct policy learning (PWM).
4. Zero-Shot Transfer Capability: Evidence that the orthogonal basis allows the model to generalize across simulators with different contact physics (e.g., soft-contact IsaacGym to stiff-contact MuJoCo) by isolating specific dynamic modes (e.g., friction vs. vertical support).

4. Experimental Results

The authors evaluated PRISM-WM on continuous control benchmarks, including high-dimensional humanoids, multi-task settings, and zero-shot sim-to-sim transfer.

• High-Dimensional Locomotion (DiffRL & Humanoid-Bench): PRISM-WM achieved higher sample efficiency and asymptotic performance compared to monolithic baselines (TD-MPC2, DreamerV3, PWM, SAC). Notably, while monolithic models experienced policy collapse in bipedal locomotion tasks (Humanoid, SNU-Humanoid) due to averaging non-smooth contact events, PRISM-WM maintained stable optimization.
• Multi-Task Generalization: On a 30-task DeepMind Control Suite benchmark, PRISM-WM achieved a mean normalized score of 0.531, a 23.5% improvement over the monolithic TD-MPC2 baseline (0.430). It outperformed the baseline on 26 of 30 tasks, effectively mitigating negative transfer between heterogeneous dynamics.
• Zero-Shot Sim-to-Sim Transfer: In transferring from IsaacGym to MuJoCo, the monolithic MLP baseline suffered a 47.8% performance drop with a 12% success rate. PRISM-WM reduced this drop to 11.8% with an 88% success rate. The orthogonalization allowed the model to adjust weights for specific experts (e.g., horizontal friction) without compromising the global dynamics manifold.
• Stability and Efficiency: Ablation studies confirmed that orthogonalization prevents reward estimation degradation seen in unconstrained MoE models. Inference efficiency tests on an NVIDIA RTX 4090 showed that PRISM-WM (K=4) achieved 7,712 FPS, slightly outperforming the monolithic baseline (7,499 FPS) with minimal memory overhead.

5. Significance and Claims

The paper claims that PRISM-WM provides a high-fidelity substrate for trajectory optimization algorithms, acting as a foundational model for model-based agents. By decomposing complex physics into orthogonal primitives, the model:

• Mitigates long-term drift in planning horizons, enabling reliable long-horizon reasoning in contact-rich environments.
• Stabilizes gradient propagation for policy learning, transforming fractured loss landscapes into functionally smooth convex blends.
• Enables robust decision-making in hybrid systems where standard world models fail due to over-smoothing.

The authors conclude that decomposing complex physics into orthogonal primitives is "fundamentally essential for mastering contact-rich control." While the current framework requires manual tuning of the expert count $K$, the work establishes a structural foundation for building general-purpose, physically-grounded world models for future embodied agents.