Technical Summary: PRISM (PRior-guided Imagination Sampling in world Models)

1. Problem Statement

In continuous control, learned world models provide agents with "physical intuition" to simulate future states. However, the effectiveness of model-based planning (e.g., Model Predictive Control, MPC) critically depends on how candidate actions are generated. Existing planners typically initialize their sampling distribution with a zero-mean, fixed-variance Gaussian (uninformed initialization) or use expert demonstrations only to set a mean, discarding the expert's state-conditioned confidence.

Current approaches to obtaining a robust action prior often rely on architectural bloat, such as independent visual encoders or large-scale Vision-Language Models (VLMs). This creates redundancy, as the world model's own learned representations inherently encode the agent's "action intuition" (which actions drive the state toward a goal). The core challenge is to extract this implicit intuition from the world model itself to guide the planner's sampling process without adding significant computational overhead or requiring separate, heavy prior models.

2. Methodology: PRISM

PRISM is a task-agnostic framework that integrates a learned action prior directly into the planner's sampling distribution using a standard JEPA-style (Joint-Embedding Predictive Architecture) latent world model.

2.1 Architecture

• Shared Backbone: PRISM utilizes a single frozen JEPA encoder ($h_\psi$) that embeds both the current observation ($o_t$) and the goal observation ($o_g$).
• Physical Intuition: The frozen predictor ($f_\theta$) rolls out action sequences in the embedding space to score them against the goal (visual cost).
• Action Intuition (The Prior): A lightweight, trainable Multi-Layer Perceptron (MLP) head ($g_\phi$, $\approx$1M parameters) is attached directly to the frozen encoder. It maps the current and goal embeddings to a Gaussian distribution over the upcoming $H$-step action sequence: $g_\phi([z_t, z_g]) \to (\mu_p, \sigma_p)$.
  • This head is trained offline on the same dataset used for the world model using a $\beta$-NLL loss to stabilize variance prediction.
  • It operates in the StandardScaler-normalized action space.

2.2 Precision-Weighted Fusion (Product of Gaussians)

At planning time, PRISM fuses the planner's default initialization ($N(\mu_\pi, \sigma_\pi^2)$) with the learned prior ($N(\mu_p, (s\sigma_p)^2)$) using a closed-form Product of Gaussians (PoG) update:

• Fused Variance: $\sigma_{fused}^2 = (\tau_\pi + \tau_p)^{-1}$, where $\tau$ represents precision ($1/\sigma^2$).
• Fused Mean: $\mu_{fused} = \sigma_{fused}^2 (\tau_\pi \mu_\pi + \tau_p \mu_p)$.
• Scaling: A scalar $s$ rescales the prior's standard deviation (set to 1 in experiments).

2.3 Integration with MPC

The fused distribution initializes a standard MPPI (Model Predictive Path Integral) loop. Crucially, PRISM maintains the fixed-covariance property of MPPI (updating only the mean, not the variance, across iterations). This ensures that the prior's confidence (precision) persists throughout the entire optimization horizon, rather than being overwritten by the iterative refinement process as happens in CEM (Cross-Entropy Method).

2.4 Graceful Degradation

The fusion mechanism is inherently robust:

• High Confidence: If the prior is confident ($\sigma_p$ is small), its precision dominates, narrowing the search around $\mu_p$.
• Low Confidence: If the prior is uncertain (e.g., out-of-distribution states where $\sigma_p$ is large), its precision vanishes ($\tau_p \to 0$), and the initialization automatically reverts to the vanilla, prior-free MPPI. This requires no learned gates or tuning.

3. Key Contributions

1. Architectural Simplicity: Demonstrates that action priors can be extracted directly from a world model's frozen encoder using a lightweight MLP, eliminating the need for separate, heavy perceptual systems or VLMs.
2. Parameter-Free Integration: Introduces a closed-form, precision-weighted fusion mechanism that requires no inference-time optimization or hyperparameter tuning beyond a single scale factor.
3. Robustness via Precision: Unlike simple warm-starting (which discards variance information), PRISM preserves the prior's per-state confidence throughout the planning horizon, enabling graceful degradation when the prior is unreliable.
4. Sample Efficiency: Achieves significant performance gains with drastically reduced sampling budgets, making it suitable for real-time control on constrained hardware.

4. Experimental Results

The authors evaluated PRISM on two simulated visual manipulation tasks (PushT and Cube) and validated it on two real-robot platforms (Franka PushT and ARX X5 Cube).

• Performance Gains:
  • PushT: PRISM-MPPI improved success rates by 35 percentage points over vanilla MPPI (89% vs. 54% at $K=128$) and by 23–35 pp across various sample budgets.
  • Cube: PRISM-MPPI improved success rates by 32 percentage points over vanilla MPPI (79% vs. 47% at $K=128$).
  • Sample Efficiency: PRISM with a small sample budget ($K=32$) outperformed vanilla MPPI with a large budget ($K=128$) by +25 pp on PushT and +35 pp on Cube.
• Ablation Studies:
  • Variance vs. Mean: Using only the prior's mean (warm-start) improved performance but failed to scale with sample budget and sometimes degraded. Incorporating the variance (full PoG fusion) was essential for robust scaling.
  • Planner Choice (MPPI vs. CEM): PRISM-MPPI (fixed covariance) significantly outperformed PRISM-CEM (adaptive covariance) at low sample budgets ($K=32$), where CEM's variance refitting collapsed due to biased elites.
  • Encoder Agnosticism: The fusion mechanism provided gains even when swapping the task-trained encoder for a frozen DINOv2 encoder, though absolute performance depended on the encoder's suitability for the task.
• Real-Robot Deployment: PRISM successfully transferred to real robots (Franka and ARX X5) with success rates of 35% and 45% respectively, demonstrating zero-latency overhead in the control loop.
• Comparison to BC/DP: While Behavior Cloning (BC) and Diffusion Policies (DP) showed promise, they generally underperformed planning methods. Notably, a 19.3M-parameter Diffusion Policy collapsed on the multimodal PushT task (41% SR) compared to PRISM's 89%, highlighting PRISM's robustness in multimodal settings despite using only a 1.0M parameter head.

5. Significance and Claims

The paper claims that PRISM represents a shift in how planners utilize learned priors. By arguing that the "architectural bloat" of separate prior models is unnecessary, PRISM leverages the inherent structure of the world model's latent space.

The significance lies in:

• Efficiency: It achieves superior planning performance with a fraction of the candidate samples, enabling higher real-time control frequencies or operation on edge devices.
• Robustness: The precision-weighted fusion ensures that the planner never relies on a misleading prior; it automatically falls back to a safe, uninformed baseline when the prior is uncertain.
• Simplicity: It offers a "parameter-free" (in terms of inference-time optimization) and closed-form solution that integrates seamlessly with existing MPPI pipelines without retraining the world model or adding complex inference-time networks.

The authors conclude that PRISM exemplifies a broader principle: using precision arithmetic to integrate lightweight, learned priors into a planner's initialization can yield substantial gains in continuous control without compromising safety or computational efficiency.