Technical Summary: IMWM (Intuition Models Complement World Models for Latent Planning)

1. Problem Statement

The paper addresses a fundamental limitation in model-based reinforcement learning (RL) for control from raw pixels: the reliance on a learned latent world model (WM) paired with a finite-budget sample-based planner (e.g., Cross-Entropy Method, CEM). While improving the accuracy of the world model is the standard approach to enhancing planning performance, the authors argue that this instinct is insufficient.

The core problem identified is that search, not prediction, can be the binding constraint. Even if a planner possesses a "perfect" world model (operationalized as an oracle that uses true environment dynamics for rollout), it may still fail to reach a goal if the finite search budget prevents it from proposing a successful action sequence within the vast action space. The paper posits that the bottleneck lies in the proposal volume—the ability of the planner to land on goal-reaching candidates—rather than the accuracy of the forward predictor.

2. Diagnosis and Theoretical Basis

Before proposing a solution, the authors provide both experimental and theoretical evidence that search is the failure point:

• Experimental Diagnosis: They replaced the learned forward predictor in a standard latent planner with an idealized rollout of the true environment dynamics (Oracle Dynamics). Despite having perfect dynamics, the planner still failed significantly on tasks like OGBench-Cube and Two-Room. Analysis of failed episodes revealed that in the vast majority of cases (e.g., 98.9% on Two-Room), the CEM population contained zero goal-reaching candidates at any replanning step.
• Theoretical Bound: The authors prove (Theorem A.1) that under a finite budget of black-box queries against any latent cost surface, the probability of success is bounded by the volume of the action space covered by the queries. This bound is independent of how the predictor was trained, confirming that a better predictor cannot overcome a search bottleneck if the proposal distribution does not cover the success manifold.

3. Methodology: IMWM

To address the search bottleneck without requiring a better world model, the authors propose IMWM (Intuition Model + World Model). This framework pairs a frozen latent world model with a separate Intuition Model trained on demonstrations. The intuition model acts as a guide to recognize promising actions, complementing rather than replacing the world model.

3.1 The Intuition Model

The intuition model ($D_{\psi}^{inv}$) is a contrastive scalar scorer trained on demonstration windows using an InfoNCE-style objective. It takes a tuple of (start latent, goal latent, action chunk) and outputs a compatibility score.

• Inverse Conditioning: Unlike the world model (which predicts $s_{t+1}$ given $s_t, a_t$), the intuition model scores actions given start and goal evidence.
• Separate Latent Space: It operates in a distinct latent space ($z_I$) decoupled from the world model's forward-prediction space ($z_L$), preventing the compounding of rollout errors.
• Training: It is trained to distinguish demonstration action chunks from negative samples within the same task context.

3.2 Three Lightweight Collaboration Components

The intuition model interacts with the planner through three specific mechanisms:

1. Retrieval Initialization:

   • Before the cost function evaluates candidates, the planner retrieves an action chunk from a fixed bank of demonstrations.
   • The retrieval is based on cosine similarity between the current (start, goal) latents and the stored demonstration keys in the intuition encoder's latent space.
   • The CEM proposal distribution is initialized as an isotropic Gaussian centered on this retrieved action chunk, biasing the search toward demonstration-supported regions.

2. Hybrid Cost:

   • The planner minimizes a composite cost function $J$ that combines the intuition score and the world-model rollout error.
   • $J = \alpha_{inv} \cdot zS(-D_{\psi}^{inv}) + \beta_{eff} \cdot zS(\text{MSE}(\text{rollout}, \text{goal}))$.
   • Both terms are z-scored across the candidate set to handle incommensurate scales. The weights $\alpha$ and $\beta$ determine the reliance on intuition versus the world model.

3. Reliability Gate:

   • Since intuition is not universally useful (e.g., if the demonstration bank lacks relevant examples), a gate decides per task/cell how much to trust the intuition model.
   • The gate computes two diagnostics from the retrieved chunk: a score margin ($r_{inv}$) against neutral chunks and a lag-1 autocorrelation ($r_{lag}$).
   • Based on frozen thresholds, the gate selects one of three pre-specified "recipes":
     • Trusted Hybrid: High reliance on both intuition and world model.
     • Intuition-Dominant: High reliance on intuition, low on world model.
     • Forward-Only Fallback: Intuition is disabled; the planner reverts to the standard world-model-only configuration.
   • Crucially, this selection is automatic based on diagnostics, not hand-tuned per task.

4. Key Contributions

1. Diagnosis of Finite-CEM Latent Planning: The paper demonstrates experimentally and theoretically that even with an oracle-dynamics world model, finite-budget planners fail due to search limitations (proposal volume) rather than prediction errors.
2. The IMWM Method: A novel architecture that integrates a contrastive intuition model with a world model via Retrieval Initialization, Hybrid Cost, and a Reliability Gate.
3. Empirical Gains: Significant improvements over world-model-only baselines on four pixel-based goal-reaching tasks without per-task tuning of the core algorithm.

5. Experimental Results

The method was evaluated on a 48-cell grid (4 tasks $\times$ 12 seeds) using the stable-worldmodel platform. The tasks were: Two-Room, Reacher, Push-T, and OGBench-Cube.

• Two-Room: IMWM achieved 99.2% success rate, a +11.5 percentage point (pp) improvement over the baseline (87.7%).
• OGBench-Cube: IMWM achieved 94.7% success rate, a massive +28.5 pp improvement over the baseline (66.2%).
• Push-T: IMWM achieved 92.7%, a +2.8 pp improvement (baseline 89.8%).
• Reacher: IMWM achieved 83.8%, a near-tie with the baseline (83.2%, +0.7 pp). In this case, the Reliability Gate correctly routed to the "forward-only fallback" recipe, as the intuition signal was not beneficial.

Ablation Studies:

• Retrieval Initialization: Turning on retrieval alone significantly boosted performance on search-limited tasks (e.g., OGBench-Cube from 66.2% to 88.3%) but hurt performance on tasks where the world model was already strong (Reacher), validating the need for the gate.
• Hybrid Cost: The full configuration (both intuition and world model costs) outperformed using either term alone.
• Reliability Gate: No single fixed recipe worked best across all tasks. The adaptive gate successfully selected the optimal recipe for each task automatically, matching the performance of a hand-tuned best-case scenario.

6. Significance and Claims

The paper claims that the primary bottleneck in latent world-model planning is often the search strategy, not the model's predictive accuracy. By introducing an "intuition" mechanism derived from demonstrations, IMWM effectively concentrates the search budget on promising regions of the action space, bypassing the volume limitations of blind sampling.

The authors emphasize that:

• The improvements are not due to a better world model or per-task hyperparameter tuning. The algorithmic family and trained artifacts are fixed; only the gate's recipe selection varies based on pre-computed diagnostics.
• The approach is complementary: The intuition model does not replace the world model but guides it.
• Limitations: The reliability gate is hand-designed with fixed thresholds (though interpretable), the results are limited to a single benchmark family of pixel-based tasks, and the method currently requires successful demonstrations to build the intuition bank. The authors do not claim transfer to real-robot control or broader benchmarks without further work.

In conclusion, IMWM demonstrates that pairing a world model with a retrieval-based intuition model can resolve the "search failure" mode of latent planning, achieving state-of-the-art performance on specific goal-reaching tasks where pure world-model planning struggles.