Skip to content
Bud's Wiki

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

TL;DR

Section titled “TL;DR”

SAGE is a prior-conditioned planner that improves long-horizon planning through a frozen latent world model by replacing random action-proposal initialization with structured guidance from predicted intermediate subgoals. At each planning stage, a goal-conditioned generator predicts the next reachable latent subgoal for a specified duration; that subgoal then conditions the generation of candidate action sequences, which are evaluated and refined by the frozen world model before execution. Subgoals of varying durations are used as multi-scale priors, balancing fine-grained local control against long-horizon progress. On PushT and OGBench Cube, coupling subgoal decomposition with prior-conditioned action generation lifts success rates dramatically at long horizons — PushT 12.7% → 64.7% and OGBench Cube 26.7% → 67.3% at target offset 150.

Key claims

Section titled “Key claims”
  • Latent world models used as planners degrade at long horizons because a fixed candidate budget must cover an exponentially growing action space, so proposal quality — not simulation fidelity — becomes the bottleneck [Abstract].
  • Replacing random action-proposal initialization with a goal-conditioned latent-subgoal prior improves the quality of the candidate futures the world model gets to evaluate [Abstract].
  • The subgoal generator is conditioned on a target duration, and using subgoals of varying durations as priors captures semantic information across temporal scales — trading fine-grained local control against higher-level long-horizon progress [Abstract].
  • Prior-conditioned proposals still go through the frozen world model for evaluation and refinement before execution — the world model is not retrained, only the planning stack changes [Abstract].
  • On PushT with target offset 150, SAGE lifts success rate from 12.7% (baseline) to 64.7% [Abstract].
  • On OGBench Cube with target offset 150, SAGE lifts success rate from 26.7% to 67.3% [Abstract].
  • Long-horizon gains do not come at the cost of short-horizon performance — the paper reports “strong short-horizon performance” is preserved [Abstract].

Method

Section titled “Method”

SAGE plugs into an existing latent-world-model MPC loop and changes only the proposal step. The world model itself — action-conditioned latent predictor trained on offline data — is frozen. At each planning stage, given the current latent state and a task goal, a goal-conditioned generator predicts the next reachable latent subgoal along the way to the ultimate goal, parameterized by a duration hyperparameter. This subgoal then conditions a second generator that produces candidate action sequences aimed at reaching that subgoal, rather than sampling action sequences from a task-agnostic prior. The candidate sequences are rolled out through the frozen world model, scored against the subgoal (and downstream goal), and the highest-scoring candidate is executed for one chunk before the loop repeats. To handle horizons that a single fixed-duration subgoal cannot bridge, SAGE runs the subgoal-and-actions machinery at multiple durations in parallel, giving the planner both short-range refinement candidates and long-range progress candidates in the same proposal batch. Only the proposal distribution changes; the world model’s evaluation role is unchanged.

Results

Section titled “Results”
  • PushT, target offset 150: success 12.7% → 64.7% [Abstract]. This is a 5× lift over the random-proposal baseline at the horizon where planning proposal quality dominates.
  • OGBench Cube, target offset 150: success 26.7% → 67.3% [Abstract]. Similar magnitude of gain on a distinct manipulation benchmark.
  • Short-horizon performance is preserved on both benchmarks [Abstract]. Specific short-horizon numbers, ablations on subgoal duration diversity, and comparisons against a subgoal-conditioned baseline without prior-conditioned actions (i.e., isolating the two contributions) are in the paper body — not in the fetched abstract; open to re-read once tables are surfaced.

Why it’s interesting

Section titled “Why it’s interesting”

SAGE is a concrete answer to one of the recurring shape-of-the-solution questions in the World Foundation Models cluster: given a frozen latent world model, where should the design effort go to make it a useful planner? The wiki has been accumulating deployment patterns that leave the WFM frozen and add machinery on the outside — Lifting Embodied World Models for Planning and Control wraps a learned high-level action interface (image-space waypoints) around a frozen PEVA and searches CEM over the lower-dim waypoint space; μ₀: A Scalable 3D Interaction-Trace World Model (µ₀) exposes 3D interaction traces as the interface between a frozen trace-WM and small embodiment-specific action experts; World Pilot: Steering Vision-Language-Action Models with World-Action Priors uses a pretrained WAM to steer a VLA through two latent+action paths. SAGE fits the same “frozen WFM + external lifter” family but changes a different lever: rather than reshaping the action space (Lifted-WM’s waypoints) or the interface modality (µ₀’s traces), it reshapes the proposal distribution by injecting learned intermediate subgoal priors at multiple time-scales, and leaves both the action space and the WFM interface untouched. Complements AdaJEPA: An Adaptive Latent World Model (AdaJEPA) on the opposite axis — AdaJEPA keeps the proposal random but adapts the model at test time; SAGE keeps the model frozen but structures the proposals. Together they carve out the two natural degrees of freedom around a frozen latent WFM MPC loop.

See also

Section titled “See also”