Frozen world models skip per-step backprop: WARP-JEPA swaps planning for learned proposals, speeding long-horizon episodes 9.6–18.8×
Synopsis
Hierarchical Joint-Embedding Predictive Architectures (H-JEPA) previously planned by running gradient descent on actions at test time, and the iterative solver consumed 96.6%–97.7% of episode runtime on long-horizon tasks. WARP-JEPA keeps the world models frozen and replaces online optimization with learned multi-candidate trajectory proposers, scoring candidates with the frozen predictor inside the generation loop. Across 30 frozen two-, three- and four-level models on four benchmarks, it speeds long-horizon episode inference by 9.6–18.8× (2.9–3.0× on Cube) and raises task success by up to 33.3 percentage points.
Interpretation
The planning cost of hierarchical world models moves from repeated test-time gradient iterations to a single pass of candidate proposal and scoring, while the world model itself stays frozen. H-JEPA previously optimized sampled action sequences online with AdamW at every level, and the nested solver loops ran thousands of sequential forward and backward passes. Timing two-, three- and four-level H-JEPA models on four benchmarks shows the iterative solver accounts for 96.6%–97.7% of episode runtime on long-horizon tasks and 79.9% on the short-horizon Cube.
Each level gets a trained proposer that emits several candidate trajectories at once, and the frozen world model scores them in one batched rollout to pick the lowest-cost plan, with no gradient iterations at inference. Proposers are trained on a joint objective combining model-predicted rollout cost, which points toward the goal, with best-of-N behavior cloning, which pulls candidates back toward demonstrated actions. On Push-T L2, the variant without behavior regularization drifts off the demonstrated action manifold and reaches 42.7% success, while the regularized MM reaches 48.0% against 42.0% for online gradient descent.
Top and subordinate levels use different architectures: subordinate levels predict the two-step bridge in parallel with residual MLPs, while the top level uses a one-layer causal Transformer that feeds the imagined state from the frozen predictor back into each macro-action query. The top level must compose up to 8–12 interdependent macro-actions, and an MLP predicting them in parallel from the start and goal alone leaves later actions blind to the intermediate states earlier actions induce. On Push-T L2 and FourRoom L2, which have long top-level horizons, the closed-loop TM beats the one-shot MM by 8.8 and 8.0 percentage points respectively, whereas on Cube, whose top horizon is 2, MM is 6.7 points higher.
Perspective
Evaluation is confined to 30 frozen H-JEPA models on Push-T, FourRoom, Visual AntMaze and OGBench Cube at depths of two, three and four levels, with success averaged over three model seeds and timing run sequentially on a single H200. On Push-T L4, H-JEPA's 40-step top-level stride aliases contact dynamics and every planner fails, so that setting does not support the method's effectiveness.
The architectural split between top and subordinate levels pays off on long top-level horizons, but on Cube, whose top horizon is 2, MM is 6.7 points above TM and the gap is not significant; end-of-episode unconditioned bypass solves still run 10 gradient iterations, covering 9%–33% of subordinate solves and 75% of level-1 solves on Cube; about 3 seconds per AntMaze episode goes to encoding and simulation outside the solvers, bounding further speed-ups.