PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Action Prediction: Learning Search-Free Planning in Latent Space

Guanchen Du, Jinhao Du, Yanpeng Sun, Shan Zhang, Na Zhao, Zechao Li

arXiv:2610.23615v1cs.ROcs.AI

TL;DR

Latent world model planning typically requires costly test-time optimization, and action prediction alone does not ensure that goal representations are action-relevant or that predicted actions produce consistent latent transitions. HIJP trains hindsight-conditioned inverse dynamics together with a latent predictor, achieving 86.8% overall success and 1.9–16.5× planning speedups over LeWM across four visual control tasks.

  • Problem

    Latent world models support visual planning but typically require costly test-time search, while action prediction alone does not ensure action-relevant goal representations or consistent predicted transitions.

    Source: §1 · pp. 1–2

  • Method

    HIJP relabels future observations as goals, lets inverse-action gradients shape goal representations, and aligns predicted-action transitions with encoded successor observations.

    Source: §1 · p. 2

  • Results

    86.8% overall success, 4.9 percentage points above the strongest baseline, and 1.9–16.5× planning speedups over LeWM with matched replanning counts across four visual control tasks.

    Source: §1 · p. 2, §4.3 · p. 7

  • Takeaways & Limitations

    HIJP generates multi-step action sequences by alternating inverse dynamics and latent prediction without online action search.

    Source: §1 · p. 2

    The evaluation uses tasks with limited temporal span, fixed planning horizons, and fixed replanning schedules; extending HIJP to substantially longer-horizon tasks remains future work.

    Source: §6 · p. 10

Abstract

from arXiv · show

Latent world models learn action-conditioned state transitions in a compact representation space, providing a predictive model for visual planning and control. However, planning with these models typically requires costly test-time optimization over action sequences. We propose to replace this online search with a goal-conditioned inverse model that directly predicts actions from current and goal observations. A key challenge is that action prediction alone neither ensures that goal representations encode action-relevant information nor that predicted actions are consistent with the learned latent prediction model. We introduce Hindsight Inverse JEPA Planning (HIJP), a search-free framework that addresses both aspects during training. HIJP relabels future observations along offline trajectories as goals, providing action supervision across multiple temporal horizons. By retaining gradients through the goal representations, this supervision directly shapes the visual representation. HIJP further feeds predicted actions into the latent predictor model and aligns the resulting next-state representations with encoded successor observations, enforcing action-transition consistency. At inference time, the inverse dynamics model and latent predictor models are applied alternately to generate multi-step action sequences without online search. Across four visual control tasks, HIJP achieves 86.8% overall success and a 1.9-16.5x planning speedup over search-based planning.

1 INTRODUCTION

Latent world model planning often relies on costly test-time search, while action prediction alone leaves goal representations and predicted transitions underconstrained. HIJP trains goal-conditioned action prediction and its induced latent transition together, achieving strong success and planning-efficiency results across four visual control tasks.

1 INTRODUCTION

  • Latent world models support visual planning, but commonly use CEM or MPPI to repeatedly generate and evaluate candidate action sequences at each replanning step.These optimizers require substantial test-time computation.Source: §1 · p. 1
  • Two gaps motivate HIJP: action supervision should shape goal-side visual features, and predicted actions should be constrained to produce the intended latent transitions.Offline trajectories contain multiple future goals at different horizons, while separate action and latent-prediction objectives do not directly constrain their planning-time composition.Source: §1 · pp. 1–2
  • HIJP relabels future observations as goals, retains gradients through goal representations, and trains predicted actions to induce transitions aligned with encoded successor observations.At test time, the inverse dynamics model and latent predictor alternate to generate multi-step action sequences without online search.Source: §1 · p. 2
  • 86.8% overall success, 4.9 percentage points above the strongest baseline; with matched replanning calls, HIJP achieves 1.9–16.5× planning speedups over LeWM.The evaluation covers four visual control tasks and four configurations of goal temporal offsets and execution budgets.Source: §1 · p. 2

2 RELATED WORK

Related work develops latent world models, action-supervised representations, and methods that amortize planning computation. HIJP combines goal-conditioned action prediction with forward latent prediction and couples the two during training.

Latent World Models

  • Latent world models use pretrained visual features, end-to-end representation learning, or large-scale video pretraining; systematic studies examine architectures, objectives, and planning algorithms.Gao & Xu (2026) reduce forward-rollout costs with parallel prediction, while Wang et al. (2026a) use inverse dynamics and goal-conditioned supervision but retain CEM.Source: §2 · p. 2

Inverse Dynamics and Goal-Conditioned Action Learning

Learned and Amortized Planning

Consistency and Goal-Aware Planning Objectives

  • Prior consistency objectives connect model prediction with planning through action-recovery costs or multi-step rollout supervision [Seo et al., 2026; Bai & Xiong, 2026].HIJP aligns the successor latent predicted under a generated action with the encoding of the next observation, allowing state-prediction errors to supervise action learning directly.Source: §2 · p. 3

3 METHODOLOGY

HIJP jointly learns a visual encoder, latent predictor, and goal-conditioned inverse dynamics model from offline visual trajectories. Training couples hindsight-based action supervision with action-transition consistency, matching the inverse-to-predictor composition used for search-free planning.

3.1 PROBLEM FORMULATION

  • The task is to produce actions that reach a goal observation within a fixed interaction budget, using recorded observations and actions without reward annotations.Planning repeatedly composes an inverse model that predicts an action from current and goal latents with a predictor that estimates the action-induced latent state.Source: §3.1 · p. 3

3.2 LATENT WORLD MODEL

  • A visual encoder maps observations to latent states, and an action-conditioned predictor learns latent transitions from demonstrated actions.The predictor is trained on demonstrated transitions; a regularization term discourages representational collapse, and the same model accepts the shorter context used by the consistency objective.Source: §3.2 · pp. 3–4

3.3 HINDSIGHT-CONDITIONED INVERSE DYNAMICS

  • HIJP samples future observations at varying temporal offsets as goals, then predicts the demonstrated immediate action at the anchor observation.The target is the immediate action, not the intervening action sequence, so each window supplies supervision under goals at different horizons.Source: §3.3 · p. 4
  • Gradients from inverse-action supervision flow through both current and future-goal latents, training the visual features used to represent each observation.The inverse model is an MLP operating on the concatenated latent pair.Source: §3.3 · p. 4

3.4 ACTION TRANSITION CONSISTENCY

  • The action-transition consistency objective passes the predicted action through the latent predictor and aligns its successor with the encoded one-step observation.It updates the inverse head, action encoder, predictor, and anchor-side visual representation, coupling action generation with its induced transition even when the hindsight goal is farther away.Source: §3.4 · p. 5

3.5 JOINT OBJECTIVE AND SEARCH-FREE PLANNING

  • Joint training combines latent prediction, representation regularization, hindsight-conditioned action prediction, and action-transition consistency; inference alternates inverse action prediction and latent prediction to form multi-step plans without online search.The models are jointly optimized, and each replanning call uses one inverse-to-latent prediction rollout rather than sampling or iteratively refining candidate action sequences.Source: §3.5 · p. 5

4 EXPERIMENTS

Across four visual-control tasks, the experiments assess goal-reaching, planning efficiency, ablations, latent-space progress, and action smoothness. HIJP achieves the highest aggregate success and faster planning than LeWM.

4.1 EXPERIMENT SETUP

  • The evaluation covers four continuous-control tasks, four goal-horizon and interaction-budget tiers, and eight baselines under a shared protocol.HIJP plans with a single H = 5 amortized rollout and no test-time search.Source: §4.1 · pp. 5–6

4.2 PLANNING PERFORMANCE

  • 86.8% overall success rate gives HIJP the highest aggregate result, exceeding the strongest baseline by 4.9 percentage points.The largest reported gains are on challenging manipulation tiers.Source: Table 1 · p. 6, §4.2 · p. 7

4.3 PLANNING EFFICIENCY

  • 1.9–16.5× speedups over LeWM under matched replanning counts show reduced planning time across all evaluated settings.Average speedups are 6.2× on TwoRoom, 7.1× on Reacher, 15.0× on PushT, and 2.8× on OGCube.Source: §4.3 · p. 7, Table 2 · p. 7

4.4 ABLATIONS

  • Removing Linv reduces average success from 82.5% to 45.7%; removing SIGReg reduces it to 47.0%, with near-collapse on Reacher.The ablations support the importance of action supervision and representation regularization.Source: Figure 2 · p. 8, Table 3 · p. 8
  • Restricting goals to each window’s final observation lowers average success to 71.7%, while stopping gradients through zg lowers it to 78.5%.These ablations test hindsight goal sampling and gradient flow through goal representations.Source: Figure 2 · p. 8, Table 3 · p. 8
  • Removing Lcyc lowers average success to 76.0%, blocking its predictor updates to 78.3%, and replacing the inverse-guided rollout with CEM to 69.3%.The effects vary across tasks; the CEM comparison keeps the trained representation fixed.Source: §4.4 · p. 8, Table 3 · p. 8

4.5 LATENT-SPACE PROGRESS ANALYSIS

  • HIJP reaches a lower final normalized latent distance than its same-representation CEM ablation on all four tasks.On PushT, HIJP continues reducing distance after CEM largely plateaus; on Reacher, its distance fluctuates less after the initial decrease.Source: §4.5 · p. 9

4.6 SMOOTHNESS OF EXECUTED ACTIONS

  • HIJP’s action jerk is 14.7–26.4× lower than its CEM ablation and 15.2–24.5× lower than LeWM across four tasks.The smoothness comparison is on the extreme tier.Source: §4.6 · p. 9, Table 4 · p. 9

5 CONCLUSION

HIJP combines hindsight goal relabeling, action supervision, and latent-transition consistency in a search-free visual-control framework. Across four tasks it achieves the highest aggregate goal-reaching performance among evaluated methods while reducing planning time relative to a search-based world-model planner.

5 CONCLUSION

  • HIJP relabels future observations as goals, trains goal representations with inverse-action gradients, and constrains predicted-action latent transitions with a consistency objective.At planning time, inverse and latent-predictor models alternate to construct multi-step plans without online search.Source: §5 · p. 9
  • Across four visual-control tasks and multiple evaluation configurations, HIJP achieves the highest aggregate goal-reaching performance among evaluated methods while reducing planning time versus a search-based world-model planner.Under the most challenging setting, its executed actions are smoother than those produced by CEM using the same representations.Source: §5 · p. 9

6 LIMITATIONS

The evaluation is limited to tasks with a limited temporal span and fixed planning horizons and replanning schedules. The experiments reuse existing LeWM offline datasets without modification.

6 LIMITATIONS

  • The evaluation uses tasks with a limited temporal span and fixed planning horizons and replanning schedules.Extending HIJP to substantially longer-horizon tasks remains future work.Source: §6 · p. 10

A DATASET DETAILS AND PROTOCOL

A.1 DATASETS

The four datasets differ in behavior policies and action distributions, and share a common visual-data format while varying in action dimensions and episode lengths.

A.1 DATASETS

  • Behavior policies range from goal-directed navigation to near-uniform random and expert policies, producing distinct action distributions across tasks.Reacher actions are essentially uniform, Tworoom actions saturate at ±1, and PushT and OGCube actions are concentrated and structured.Source: §A.1 · p. 13, Figure 4 · p. 14
  • Each dataset uses a shared HDF5 schema with RGB observations, continuous actions, and episode boundaries that prevent sampled windows from crossing episodes.The data loader assigns 90% of valid trajectory windows to training and 10% to validation, replacing NaN action entries with zeros during training.Source: §A.1 · pp. 13–14
  • All observations are 224 × 224 RGB frames; action spaces are 2D for Tworoom, Reacher, and PushT, and 5D for OGCube.Reacher and OGCube episodes have fixed lengths of 201 steps, while Tworoom and PushT episode lengths vary with task completion.Source: §A.1 · pp. 13–14
  • Actions and other non-pixel inputs are standardized per dimension using dataset statistics, while the figure shows raw action distributions before normalization.The same z-score transformation is applied across all four tasks.Source: §A.1 · p. 14, Figure 4 · p. 14

A.2 EVALUATION PROTOCOL

The evaluation varies goal horizon and interaction budget across four tiers and uses simulator success criteria. Shared CEM settings, temporal discretization, and preprocessing define the planning and data-flow protocol.

A.2 EVALUATION PROTOCOL

  • Episodes use offline-trajectory start states and goals K steps ahead; success requires simulator-reported completion within budget, not a latent-distance threshold.Four tiers vary goal offset and step budget, from easy at K=10 to extreme at K=75; larger K means a longer-horizon goal.Source: §A.2 · pp. 14–15, Table 6 · p. 15, Table 7 · p. 15

A.3 CEM PLANNING PROTOCOL

  • CEM searches a factorized Gaussian over H=5 model steps, sampling 300 candidates for 30 iterations and retaining 30 elites.Each sequence covers 25 environment steps; each replanning call evaluates 45,000 latent transitions in parallel on the GPU, which is not equivalent to a wall-clock slowdown.Source: §A.3 · p. 15
  • CEM comparisons match HIJP’s horizon, action chunking, interaction budget, and replanning cadence; the controlled ablation changes only the test-time planner.The controlled CEM ablation uses HIJP’s encoder and forward model while disabling its inverse-dynamics action generator.Source: §A.3 · pp. 15–16

B IMPLEMENTATION DETAILS

  • The implementation shares architecture and temporal settings across all four tasks, with images normalized and resized and continuous actions standardized for training.Predicted action chunks are inverse-transformed to the environment action space at execution.Source: B IMPLEMENTATION DETAILS · p. 17, §B.1 · p. 17

B.1 TEMPORAL DISCRETIZATION AND DATA FLOW

  • The world model groups five consecutive atomic actions into one model-level action chunk, so each latent transition spans five environment steps.Raw action dimensions are 2 for Tworoom, Reacher, and PushT and 5 for OGCube.Source: §B.1 · p. 17
  • Training windows contain temporally subsampled observations aligned with action chunks, with five raw environment steps between consecutive observations.The first and last observations are separated by (W −1)k raw steps.Source: §B.1 · p. 17

B.2 ARCHITECTURE

HIJP combines a visual encoder, action encoder, latent predictor, and hindsight-conditioned inverse head. Joint training couples prediction, representation regularization, inverse supervision, and action-transition consistency; inference alternates inverse prediction and latent rollout.

B.2 ARCHITECTURE

  • The architecture combines a visual encoder, action encoder, autoregressive latent predictor, and inverse-dynamics head that maps current and future goal latents to an action chunk.The visual encoder produces 192-dimensional latents, and the predictor uses recent latent-action pairs as context.Source: §B.2 · pp. 17–18

B.3 TRAINING

  • Joint training combines next-latent prediction, SIGReg anti-collapse regularization, hindsight-conditioned inverse dynamics, and action-transition consistency.The prediction term uses a causal context of length three, while SIGReg is applied to latents from all 16 frames in the sampled window.Source: §B.3 · pp. 18–19
  • Hindsight training samples anchor-goal pairs at varied future offsets, exposing the inverse head to multiple temporal distances without rewards or task labels.The supervised target is the demonstrated immediate action chunk at the anchor, not the action sequence between anchor and goal.Source: §B.3 · pp. 18–19
  • For each hindsight pair, the predicted action is fed to the latent predictor and its next-latent output is regressed toward the actual successor latent.The target is detached; gradients update the inverse head, action encoder, predictor, visual encoder, and anchor-side representation.Source: §B.3 · p. 19

B.4 PLANNING AND RECEDING-HORIZON EXECUTION

  • At each rollout step, the inverse head predicts an action chunk toward the goal latent and the latent predictor advances the imagined state.An H-step plan requires H evaluations of each model, with no candidate population or iterative test-time action optimization.Source: §B.4 · pp. 19–20
  • Planning uses shared settings across tasks and a receding-horizon loop that executes the predicted sequence before re-encoding the resulting observation and replanning.The target remains fixed within an episode; execution ends on task success or when the raw-step budget is exhausted.Source: §B.4 · p. 20

C LOSS WEIGHT ABLATION

The ablation varies the inverse-dynamics and cycle-consistency loss weights while evaluating checkpoints under the main protocol. Default weights yield the best overall average among the tested settings.

C LOSS WEIGHT ABLATION

  • 82.5% overall average success is achieved by the default weights λinv=1.0 and λcyc=0.1, the highest among tested settings.The comparison reports success rates for loss-weight ablations; Overall Avg. averages all 12 task–tier combinations.Source: Table 12 · p. 21, C LOSS WEIGHT ABLATION · p. 21
  • The ablation varies λinv and λcyc separately while fixing the other weight, and omits TwoRoom because its success is near 100% and cannot distinguish settings.Models are evaluated using the same protocol as the main experiments.Source: C LOSS WEIGHT ABLATION · p. 21
  • Raising λinv to 2.0 improves PushT and matches default OGCube performance, but reduces Reacher performance and lowers the overall average to 81.8%.Changing λcyc to 0.01 or 0.2 yields overall averages of 81.2% and 78.0%, respectively.Source: C LOSS WEIGHT ABLATION · p. 21

D ACTION SMOOTHNESS EXPERIMENT

The action-smoothness study compares HIJP with CEM-based planners using action jerk across four environments and four difficulty tiers. HIJP consistently produces lower jerk than its CEM ablation and LeWM.

D ACTION SMOOTHNESS EXPERIMENT

  • Methods were evaluated closed-loop on the same start–goal instances, using 50 episodes per task with random seed 42.The four tiers use goal-offset and step-budget pairs listed in Table 7; sampling depends on the dataset, seed, goal offset, and episode count.Source: D ACTION SMOOTHNESS EXPERIMENT · p. 22
  • HIJP directly executes its amortized plan, while its CEM ablation uses the same learned representation and predictive model but replaces inverse-guided actions with CEM planning.LeWM provides another reference by performing CEM planning over its own forward model.Source: D ACTION SMOOTHNESS EXPERIMENT · p. 22
  • On the extreme tier (K = 75), the CEM ablation produces 14.7×–26.4× higher jerk than HIJP, which shares its learned representation and predictive model.The reported controlled comparison attributes the smoothness advantage primarily to amortized inverse-guided planning rather than different learned representations.Source: Table 13 · p. 22, D ACTION SMOOTHNESS EXPERIMENT · p. 23
  • Successful episodes are scored only over the action prefix through the first goal-reaching step, while unsuccessful episodes contribute their full budget-length sequence.This truncation measures executed controls without including post-success idling.Source: D ACTION SMOOTHNESS EXPERIMENT · p. 22
  • Action jerk is the mean ℓ2 norm of the executed action sequence's discrete second difference, with lower values indicating smoother controls.Constant or linearly varying actions have zero second difference, while oscillatory or abruptly changing commands produce larger values.Source: D ACTION SMOOTHNESS EXPERIMENT · pp. 22–23
  • For each method, environment, and tier, jerk is computed per episode and averaged across episodes; ratios above one indicate smoother actions from HIJP than the corresponding CEM-based planner.All methods use the same normalized action space, making their jerk values directly comparable.Source: D ACTION SMOOTHNESS EXPERIMENT · p. 23

E VISUALIZING PLANNED TRAJECTORIES

The paper visualizes HIJP rollouts as pixel-space filmstrips across four environments. The figures show four difficulty tiers per environment, with the goal observation displayed separately.

E VISUALIZING PLANNED TRAJECTORIES

  • Each environment figure shows Easy, Default, Hard, and Extreme tiers, defined by goal offsets of 10, 25, 50, and 75 steps and budgets of 25, 50, 100, and 100 steps.Each tier's rollout uses 16 uniformly sampled frames arranged in two rows of eight.Source: E VISUALIZING PLANNED TRAJECTORIES · p. 23
  • The goal observation appears separately on the right with a red outline, and each tier label indicates whether the rollout reaches the goal within its budget.The filmstrip frames span 0% to 100% of the rollout.Source: E VISUALIZING PLANNED TRAJECTORIES · p. 23
  • Figures 6–9 present representative HIJP trajectories in Push-T, Reacher, TwoRoom, and OGCube, respectively.Each figure visualizes planner rollouts from a start observation toward a goal observation.Source: E VISUALIZING PLANNED TRAJECTORIES · p. 23, Figure 6 · p. 24, Figure 7 · p. 25
Loading 2610.23615v1…