Technical Summary: SAGE – Subgoal-Conditioned Action Generation for Latent World Model Planning

Problem Statement

Latent world models have established themselves as efficient planning paradigms by learning action-conditioned predictive dynamics in a compact representation space, allowing agents to simulate and evaluate candidate action sequences without pixel-level reconstruction. However, as the planning horizon extends, the efficacy of these models becomes increasingly constrained by the quality of action proposals. Under a fixed candidate budget, initializing action searches with a generic random distribution fails to provide adequate coverage of the exponentially expanding action space. Consequently, goal-directed sequences are rarely sampled, and the world model is not exposed to high-quality candidate futures for evaluation. This "proposal gap" severely limits long-horizon planning performance, even when the underlying predictive dynamics are accurate.

Methodology: SAGE Framework

To address the proposal gap, the authors introduce SAGE (Subgoal-Conditioned Action Generation), a prior-conditioned planning framework. SAGE replaces random proposal initialization with structured guidance derived from predicted latent subgoals. The framework operates by decomposing long-horizon planning into two tightly coupled, learned decisions while keeping the underlying world model (specifically LeWM) frozen.

1. Goal-Conditioned Subgoal Generator

The first component is a variable-length latent subgoal generator. Given the current observation history, a distant goal, and a requested future duration ($\tau$), this module predicts a reachable intermediate latent state (subgoal) at time $t+\tau$.

• Architecture: A four-layer Transformer decoder that takes history tokens, current state, far-goal tokens, and temporal embeddings ($\delta, \tau$) as input.
• Mechanism: It outputs a residual correction to the far-goal latent, producing a local target $\hat{z}_{t+\tau}$ that is reachable within the specific duration $\tau$.
• Function: This generator decomposes distant goals into reachable intermediate states at varying temporal scales, capturing both fine-grained local dynamics (short $\tau$) and higher-level task progress (long $\tau$).

2. Subgoal-Conditioned Action Generator

The second component maps the current history, the far goal, and the predicted local subgoal to a distribution of action options.

• Architecture: A three-layer Transformer decoder using learned action queries to attend to the conditioning memory (history, state, subgoal, far goal).
• Output: The network parameterizes a trajectory-level Gaussian Mixture Model (GMM), $q_\phi(a_{t:t+\tau-1} | \dots) = \sum \pi_m \mathcal{N}(a; \mu_m, \Sigma_m)$.
• Function: Instead of sampling from a generic distribution, the planner samples action sequences conditioned on the specific local subgoal. This ensures the proposed actions are structurally aligned with the immediate, reachable target.

3. Planning and Refinement Loop

The planning process integrates these learned components with the frozen world model:

1. Decomposition: A long-horizon query $H$ is decomposed into a schedule of durations $(\tau_1, \dots, \tau_R)$.
2. Subgoal Prediction: At each stage, the subgoal generator predicts a local target $\hat{z}_{t+\tau}$ based on the remaining horizon.
3. Proposal Generation: The action generator produces $K$ candidate action sequences conditioned on this local target.
4. Evaluation & Refinement: The frozen LeWM predicts the latent futures of these candidates and evaluates their cost against the local subgoal. The Cross-Entropy Method (CEM) then refines the elite set of candidates over 30 rounds.
5. Execution: The agent executes the leading segment of the selected sequence, observes the new state, and replans.

Key Contributions

1. Multi-Horizon Latent Subgoal Generator: A novel module that decomposes distant goals into reachable intermediate states at different temporal scales, effectively bridging the gap between current state and distant objectives.
2. Subgoal-Conditioned Action Generator: A mechanism that generates action candidates conditioned on predicted subgoals, replacing generic random initialization with structured, task-aware priors.
3. Empirical Validation of Structured Proposals: Demonstration that coupling variable-length subgoal generation with conditioned action generation significantly improves long-horizon planning while preserving short-horizon performance, without modifying the frozen world model's predictive dynamics.

Experimental Results

The method was evaluated on PushT (requiring sustained planar contact) and OGBench Cube (requiring arm motion, grasping, and placement) using a controlled setup with a frozen LeWM checkpoint, identical candidate budgets (300), and CEM refinement (30 rounds).

• Long-Horizon Performance: As the planning horizon increases, SAGE demonstrates substantial gains over baselines (standard LeWM and PRISM).
  • At a target offset of H=150:
    • PushT: Success rate improved from 12.7% (LeWM) to 64.7% (SAGE).
    • OGBench Cube: Success rate improved from 26.7% (LeWM) to 67.3% (SAGE).
• Short-Horizon Performance: SAGE maintains strong performance on short horizons (e.g., 94.0% on PushT at H=25), indicating that the structured priors do not degrade performance when the goal is near.
• Ablation Studies:
  • Generator-Only: Using generated local targets with a standard Gaussian proposal (without the learned action generator) still yielded significant improvements (e.g., 58.7% on PushT at H=150 vs. 12.7%), confirming the value of local target decomposition.
  • SAGE Prior Top: Executing the highest-weight option from the learned prior without LeWM ranking or CEM refinement resulted in poor long-horizon performance (16.0% at H=150), highlighting the necessity of world-model-guided search and refinement.
• Temporal Schedules: Experiments with different duration schedules (e.g., varying the mix of 15, 20, and 25-step commitments) revealed that both the frequency of replanning and the ordering of durations significantly impact success rates.

Significance and Claims

The paper claims that SAGE offers a practical pathway to extend the planning capabilities of existing latent world models without the computational cost of retraining their predictive dynamics. By improving the structure of action proposals—specifically by decomposing distant goals into reachable local targets and conditioning action generation on these targets—the framework effectively mitigates the proposal quality bottleneck inherent in long-horizon planning.

The authors emphasize that their approach separates the "proposal" mechanism from the "prediction" mechanism. While the world model remains a fixed simulator for evaluating futures, SAGE ensures that the candidates presented to the simulator are plausible and aligned with the task structure. This modular improvement allows for substantial gains in complex, long-horizon robotic tasks while retaining the efficiency and robustness of the underlying world model.