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

B EYOND ACTION P REDICTION : L EARNING S EARCH F REE P LANNING IN L ATENT S PACE Guanchen Du1 , Jinhao Du2 , Shan Zhang4 , Na Zhao3 , 1 3

Yanpeng Sun3 , Zechao Li1

Nanjing University of Science and Technology Singapore University of Technology and Design

2

University of Oxford Adelaide University

4

arXiv:2610.23615v1 [cs.RO] 16 Oct 2026

Project page: https://hijp.demohub.live

A BSTRACT 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 goalconditioned 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.5× planning speedup over search-based planning.

1

I NTRODUCTION

Latent world models encode high-dimensional visual observations into compact latent representations and learn an action-conditioned latent predictor from offline trajectories, providing a foundation for visual planning and control (Zhou et al., 2025; Maes et al., 2026; Zhang et al., 2025). However, planning with these models typically requires optimizers such as CEM (Rubinstein, 1999) or MPPI (Williams et al., 2017) to repeatedly generate and evaluate candidate action sequences at each replanning step, which requires substantial test-time computation (Pinneri et al., 2021). Learning a direct mapping from the current state and goal to actions therefore provides a promising alternative for improving the planning efficiency of latent world models (Sun et al., 2026; Nguyen et al., 2026). Recent work has introduced action supervision to improve latent world models’ control capabilities. SMWM learns action-relevant representations through inverse dynamics supervision (Ivashkov et al., 2026), while Delta-JEPA (Zhang et al., 2026) recovers actions from latent state differences, encouraging representations to retain action-related state changes. GC-IDM (Nguyen et al., 2026) further learns a goal-conditioned action mapping on frozen representations, while INTACT (Sun et al., 2026)incorporates local and future-goal action supervision into end-to-end training, supporting search-free recursive control. These approaches demonstrate the value of action supervision for goal-conditioned control. However, accurate action prediction alone does not necessarily provide all the supervision needed for search-free planning. We identify two aspects that remain underconstrained. First, goal-conditioned action supervision should not only train an inverse model to predict actions, but also help shape visual representations 1

to capture how actions relate current states to future goals. Offline trajectories naturally contain such relations at multiple temporal scales: for the same current state and demonstrated action, different future observations define goals at different temporal horizons. However, existing training schemes (Sun et al., 2026; Wang et al., 2026a) exploit only a subset of these relations, for example by conditioning action prediction on a fixed future endpoint. More generally, treating the goal representation only as a conditioning target limits the extent to which action supervision can shape goal-side visual features. Second, accurate action prediction alone does not ensure that the predicted action induces the correct transition under the learned latent predictor. The inverse dynamics model and the latent predictor model are typically supervised separately, with the latter conditioned on demonstrated actions (Ivashkov et al., 2026; Sun et al., 2026). During search-free rollout planning, however, these two mappings are applied sequentially: the inverse dynamics model predicts an action toward the goal, and the latent predictor model uses this predicted action to advance the latent state for subsequent action generation. The resulting inverse-to-predictor composition is therefore directly involved in planning, yet is not explicitly constrained by separate action and latent prediction objectives. This motivates extending action supervision beyond action prediction itself, using it both to shape goal-conditioned representations and to couple predicted actions with their induced latent transitions. Based on this principle, we propose Hindsight Inverse JEPA Planning (HIJP), which jointly trains a visual encoder, a goal-conditioned inverse dynamics model, and a latent predictor model end-to-end. HIJP relabels future observations within each trajectory window as goals, exposing the inverse objective to diverse state–goal relations across temporal horizons. Gradients are retained through the goal representations, allowing inverse-action supervision to directly shape the visual features used for goal-conditioned action prediction. We further feed the predicted actions into the latent predictor and align the resulting successor latents with the encodings of the next observations. This actiontransition consistency objective jointly updates the inverse and predictor models, directly coupling action generation with its induced latent transition. In this way, HIJP trains the same inverse-topredictor mapping that is repeatedly applied during planning. At test time, the inverse dynamics module and latent predictor model alternate to generate multi-step action sequences, which are executed through closed-loop replanning without online action search. We evaluate HIJP on four visual control tasks under four configurations of goal temporal offsets and execution budgets. HIJP achieves an overall success rate of 86.8%, outperforming the strongest baseline by 4.9 percentage points. Ablations on hindsight goal sampling, training objectives, and gradient routing examine how goal-side representation learning and inverse - predictor coupling contribute to control performance. Beyond success rates, we also evaluate efficiency and the temporal smoothness of executed actions. With matched numbers of replanning calls, HIJP achieves a 1.9–16.5× planning speedup over LeWM. Under the most challenging setting, its executed actions exhibit lower jerk than those produced by CEM using the same representations, indicating smoother executed controls.

2

R ELATED WORK

Latent World Models Latent world models learn action-conditioned changes in representation space to support goal-directed planning from offline visual trajectories. Existing approaches span pretrained visual features, end-to-end representation learning, and large-scale video pretraining (Zhou et al., 2025; Maes et al., 2026; Assran et al., 2025; Bo et al., 2026). Systematic studies further investigate how architectures, training objectives, and planning algorithms affect control performance (Terver et al., 2025). Gao & Xu (2026) reduce forward-rollout costs through parallel prediction conditioned on action prefixes. Wang et al. (2026a) use inverse dynamics and goal-conditioned action supervision to improve planning representations, while retaining CEM for control. HIJP builds on latent forward dynamics and focuses on learning a goal-conditioned action mapping that can be directly deployed for control. Inverse Dynamics and Goal-Conditioned Action Learning Inverse dynamics and actionsupervised objectives encourage representations to retain action-relevant state changes. These objectives can also be combined with physical-state alignment to improve representations for control (Ivashkov et al., 2026; Zhang et al., 2026; Liu et al., 2026). Hindsight learning relabels achieved states along trajectories as goals, providing supervision for goal-conditioned policies and inverse dy2

namics learning (Andrychowicz et al., 2018; Ghosh et al., 2019; Sun et al., 2019). Goal-conditioned action-sequence prediction has also been used for visual planning (Paster et al., 2020). Related theoretical studies analyze the benefits of inverse dynamics through bias–variance trade-offs and learning-problem complexity (Schäfer et al., 2026; Morin et al., 2026). Within latent world models, GC-IDM(Nguyen et al., 2026) learns a goal- and horizon-conditioned action mapping on frozen representations, while INTACT (Sun et al., 2026) jointly learns local and goal-conditioned action prediction to support search-free recursive control. HIJP follows this direction by enriching supervision through future-goal relabeling within trajectory windows. Learned and Amortized Planning Planning computation can be shifted to training by learning action priors, latent subgoals, or plan-update rules (Wang et al., 2026b; Cheng et al., 2026; Masip et al., 2026; Sommer & Schilling, 2026; Tang et al., 2025). Another line of work first predicts future states or generates state trajectories and then recovers actions through inverse dynamics (Tian et al., 2025; Xie et al., 2025; Zeng et al., 2026; Ajay et al., 2022; Huang et al., 2026). Rakhimov et al. (2026) jointly train state and action flows, allowing a single model to support multiple control inference modes. HIJP alternates goal-conditioned action prediction with forward latent-state prediction to construct multi-step plans without online action search. Consistency and Goal-Aware Planning Objectives Consistency objectives also connect model prediction with planning. Seo et al. (2026) incorporate action-recovery errors from predicted transitions into the online planning cost, while Bai & Xiong (2026) combine temporal-progress learning with multi-step rollout consistency supervision. HIJP focuses on the composed output of the goalconditioned action module and the forward model. During training, it aligns the successor latent state predicted under the generated action with the encoding of the next observation in the recorded trajectory, allowing state-prediction errors to supervise action learning directly.

3

M ETHODOLOGY

Our goal is to extend action supervision beyond action prediction itself, such that it also shapes goal-conditioned visual representations and constrains the latent transitions induced by predicted actions. HIJP jointly learns a visual encoder fθ , a latent predictor model gϕ , and a goal-conditioned inverse dynamics model hξ . Given offline visual trajectories, the inverse dynamics model learns to predict actions toward future goals, while the predictor predicts action-conditioned latent transitions. We couple these components during training to supervise the local inverse-to-predictor composition repeatedly used in search-free planning. Figure 1 provides an overview. 3.1

P ROBLEM F ORMULATION

Let D = {τ (i) }N i=1 be a dataset of trajectories, where τ = (o0 , a0 , o1 , . . . , aT −1 , oT ).

(1)

Here ot is a visual observation and at is the action taken at time t. Given a current observation ot and a goal observation og , the task is to produce actions that reach the goal within a fixed interaction budget. Training uses recorded observations and actions, with future observations in each trajectory serving as hindsight goals, where reward annotations are not required. Our planning procedure repeatedly composes two mappings. An inverse model predicts an action from the current latent state and goal latent state, and a latent predictor model predicts the latent state induced by that action. This composition motivates training the models together: action prediction must identify a useful action, while the resulting transition must agree with the observed dynamics. 3.2

L ATENT W ORLD M ODEL

Following the prior work (Maes et al., 2026), a visual encoder fθ : O → Rd maps observations to latent states zt = fθ (ot ), and an action-conditioned latent predictor model gϕ predicts the next latent state from a causal context of latent states and actions: ẑt+1 = gϕ (zt−m+1:t , at−m+1:t ) , 3

zt+1 = fθ (ot+1 ).

(2)

Initial observation

Trajectories

Goal observation

…

𝑎1

𝑎0

Goal-conditioned rollout

𝑜0

𝑜1

𝑜2

Hindsight Inverse Dynamics

Latent World Model

𝑜𝑡 Action-Transition Consistency

𝑎ƶ 𝑡 Latent Predictor Model

Latent Predictor Model

Inverse Dynamics Model

𝑧ƶ𝑡+1 ℒpred

𝑧ƶ𝑡+1 ℒcyc

ℒinv

𝑧𝑡

𝑧𝑡+1

𝑧𝑡

𝑧𝑡+𝜅

Encoder

Encoder

Encoder

Encoder

𝑧𝑡

𝑧𝑡+1

Encoder

Encoder

𝑜𝑡

𝑜𝑡+1

𝑎𝑡 𝑜𝑡

𝑎𝑡

𝑜𝑡+1

𝑜𝑡

𝑜𝑡+𝜅

Figure 1: Overview of the HIJP training framework. (A) The latent world model predicts the next latent state from the current state and action. (B) The hindsight-conditioned inverse dynamics model infers the action from current and future latent states. (C) The inferred action is reused to enforce action–transition consistency. All modules are jointly optimized using the prediction, regularization, inverse-dynamics, and action transition consistency losses. We train the latent prediction model on demonstrated transitions using Lpred = MSE(ẑt+1 , zt+1 ).

(3)

We also apply Lsigreg to the encoded observations in each training window to discourage representational collapse. We use a ViT-Tiny encoder trained from scratch and a causal Transformer model with context length m = 3. An action encoder supplies the latent prediction model with action embeddings. Architecture and optimization details are given in Appendix B. The same model also accepts the shorter context used for the one-step consistency objective below. The latent prediction model learns how demonstrated actions transform latent states. Our next objective is complementary: to learn how actions can be inferred from current and goal representations, while allowing this action supervision to shape the representations themselves. 3.3

H INDSIGHT- CONDITIONED I NVERSE DYNAMICS

A trajectory contains multiple valid future goals for a given observation. To expose these relations during training, we sample a window of W observations, an anchor i and a future offset κ: i ∼ Unif{0, . . . , W − 2},

κ | i ∼ Unif{1, . . . , W − 1 − i}.

(4)

Each sampled pair therefore defines a goal-conditioned action-learning problem at a potentially different temporal horizon. Unlike using only the endpoint of a trajectory window, this construction exposes the inverse objective to diverse relations between current observations and future goals using the same offline data. The observation at i + κ becomes the hindsight goal. The goal-conditioned inverse model predicts the demonstrated action at the anchor: Linv = MAE(âi , ai ).

âi = hξ (zi , zi+κ ),

(5)

Although the sampled goal may be several steps ahead, the target is the immediate action ai , not the intervening action sequence. Each training window therefore supplies action supervision under goals at different temporal offsets. Gradients from Linv flow through both zi and zi+κ . Consequently, the inverse objective trains not only hξ , but also the visual features used to represent the current observation and the future goal. We implement hξ as an MLP operating on the concatenated latent pair. 4

3.4

ACTION T RANSITION C ONSISTENCY

The inverse model can predict an action close to the demonstrated action without ensuring that this prediction produces the observed transition under the learned predictor model. This distinction matters at planning time, where the model receives predicted rather than demonstrated actions. We therefore pass the predicted action through gϕ : z̃i+1 = gϕ (zi , âi ) ,

âi = hξ (zi , zi+κ ),

(6)

and align the result with the encoded one-step successor: Lcyc = MSE(z̃i+1 , sg[zi+1 ]) ,

(7)

where sg denotes stop-gradient. This objective uses the one-step successor zi+1 even when the hindsight goal zi+κ is farther away. It allows Lcyc to update the inverse head, action encoder, latent prediction model, and the anchor-side visual representation. Thus, the consistency objective does not merely use the prediction model to provide an error signal to the inverse dynamics model; it jointly adapts action generation and the latent predictor around their composed prediction. This directly supervises the inverse-to-latent predictor transition that is repeatedly composed during search-free planning. 3.5

J OINT O BJECTIVE AND S EARCH - FREE P LANNING

Joint training. The full training objective combines latent prediction, representation regularization, hindsight-conditioned action prediction, and action-transition consistency: L = Lpred + λreg Lsigreg + λinv Linv + λcyc Lcyc .

(8)

All components are optimized jointly, and all terms are averaged over their corresponding batch, temporal, sampled-pair, and feature dimensions before weighting. For all tasks, the loss weights are set to λreg = 0.09, λinv = 1.0, and λcyc = 0.05 for TwoRoom and 0.10 for the remaining tasks. Search-free planning. At planning time, we use the inverse dynamics model and latent prediction models in the same order in which they are coupled during training. At each replanning point, the latest observation and fixed goal observation are encoded as zt = fθ (ot ) and zg = fθ (og ). Starting from ẑ0 = zt , we construct an H-step plan by alternating goal-conditioned action prediction and latent prediction:
aj = hξ (ẑj , zg ), ẑj+1 = gϕ ẑℓj :j , aℓj :j , ℓj = max(0, j −m+1), j = 0, . . . , H −1. (9) The inverse model therefore predicts each action from the current imagined state and fixed goal, while the latent prediction model advances the imagined state under that predicted action. We use H = 5, execute the resulting action sequence, observe the environment again, and repeat the procedure in a receding-horizon manner. Each replanning call requires only one inverse-to-latent prediction rollout, without sampling or iteratively refining candidate action sequences.

4

E XPERIMENTS

We conduct experiments on four tasks to answer the following key questions: (i) whether our method achieves strong goal-reaching performance, particularly for long-horizon goals; (ii) how much inference speed is gained by replacing the CEM search at test time with a single amortized rollout; (iii) how hindsight goal relabeling, inverse-action supervision, gradient routing, and predicted-action transition consistency contribute to control performance; and (iv) how HIJP and CEM differ in latent-space goal progress and executed-action smoothness. 4.1

E XPERIMENT SETUP

Benchmarks and Tasks We evaluate on four goal-conditioned control benchmarks with distinct dynamics: TwoRoom (2D navigation), Reacher (DMControl reacher), PushT (planar pushing), and OGCube (OGBench cube manipulation). All tasks use 224×224 RGB observations, continuous actions, and observation-specified goals. The dataset information is provided in Appendix A. We use the offline datasets released by LeWM (Maes et al., 2026) without modification. To vary difficulty, 5

Table 1: Success rates (%) across four tasks and four difficulty tiers. Avg. denotes the arithmetic mean across the four difficulty tiers; Overall Avg. averages all 16 task–tier combinations with equal weight. The best and second-best results in each row are shown in bold and underlined text. Reference Task

TwoRoom

Reacher

PushT

Baselines

Ours

Random

GCBC

PLDM

DINO-WM

LeWM

SMWM

∆-JEPA

INTACT

HIJP

Easy Default Hard Extreme

66% 24% 8% 2%

100% 100% 100% 100%

100% 88% 90% 96%

100% 100% 98% 74%

96% 88% 62% 50%

96% 96% 86% 72%

76% 38% 48% 24%

100% 100% 96% 90%

100% 100% 100% 98%

Avg.

25.0%

100.0%

93.5%

93.0%

74.0%

87.5%

46.5%

96.5%

99.5%

Easy Default Hard Extreme

12% 10% 6% 6%

4% 0% 0% 0%

70% 72% 72% 84%

46% 78% 98% 94%

40% 76% 96% 84%

34% 46% 62% 58%

14% 20% 52% 40%

44% 100% 100% 96%

48% 92% 92% 94%

Tier

Avg.

8.5%

1.0%

74.5%

79.0%

74.0%

50.0%

31.5%

85.0%

81.5%

Easy Default Hard Extreme

14% 2% 2% 0%

94% 74% 40% 4%

94% 42% 22% 0%

98% 74% 56% 28%

78% 88% 58% 12%

80% 56% 30% 16%

72% 64% 36% 0%

98% 80% 54% 6%

90% 80% 66% 42%

Avg.

4.5%

53.0%

39.5%

64.0%

59.0%

45.5%

43.0%

59.5%

69.5%

Easy Default Hard Extreme

68% 44% 26% 40%

98% 90% 58% 76%

74% 56% 28% 50%

92% 78% 62% 72%

98% 82% 56% 70%

98% 88% 70% 86%

88% 74% 54% 74%

100% 100% 70% 76%

98% 96% 96% 96%

Avg.

44.5%

80.5%

52.0%

76.0%

76.5%

85.5%

72.5%

86.5%

96.5%

Overall Avg.

20.6%

58.6%

64.9%

78.0%

70.9%

67.1%

48.4%

81.9%

86.8%

OGCube

we set each goal to the environment state K ∈ {10, 25, 50, 75} steps ahead of the start and assign a corresponding interaction budget B ∈ {25, 50, 100, 100}, defining the easy, default (LeWM’s(Maes et al., 2026) original goal-setting horizon and interaction budget), hard, and extreme tiers, respectively. All methods are evaluated in the original simulator on the same start–goal instances, with 50 episodes per task–tier combination. An episode counts as successful only when the simulator’s native success condition is met; we use no latent-space distance threshold. We report success rate as the primary metric and planning time for the efficiency study, with additional details in Appendix A.2. Baselines We compare against eight baselines: Random, GCBC (Ghosh et al., 2021), DINOWM (Zhou et al., 2025), PLDM (Sobal et al., 2025), LeWM (Maes et al., 2026), SMWM (Ivashkov et al., 2026), delta-JEPA (Zhang et al., 2026) and INTACT(Sun et al., 2026). Random provides a chance-level reference; GCBC and INTACT provide direct goal-conditioned controllers, with INTACT representing the closest end-to-end search-free baseline. The remaining methods are latent world-model planners that rely on CEM- or MPPI-based search at test time. All learned methods are trained on the same offline datasets and evaluated under the same protocol. All CEM-based comparisons use the shared search and execution protocol specified in Appendix A.3. Implementation Details HIJP combines a latent world model with an amortized inverse head trained end-to-end. At planning time, it produces an H = 5 step latent plan with a single amortized rollout, executes the corresponding actions in a receding-horizon manner, and then replans from the latest observation. Unlike the CEM-based world-model baselines, HIJP uses no test-time search. Architecture and optimization details are provided in Appendix B. 4.2

P LANNING P ERFORMANCE

We first examine how HIJP’s goal-reaching performance compares with the baselines. As detailed in Section 4.1, we test HIJP and all baseline methods on each task-tier combination, using 50 episodes per pair, and report success rates in Table 1. 6

Table 2: Total wall-clock planning time (seconds) for LeWM and HIJP. Speedup denotes LeWM’s planning time divided by HIJP’s. Avg. reports the arithmetic mean across difficulty tiers, including the per-tier speedup ratios. Lower planning times are shown in bold. Task

TwoRoom

Reacher

PushT

OGCube

Difficulty tier

Method

Avg.

Easy

Default

Hard

Extreme

LeWM HIJP

69.8 13.5

124.9 20.8

254.7 36.2

253.9 37.4

175.9 27.0

Speedup

5.2×

6.0×

7.0×

6.8×

6.2×

LeWM HIJP

153.6 32.6

296.0 45.4

566.6 65.3

581.9 69.0

399.5 53.1

Speedup

4.7×

6.5×

8.7×

8.4×

7.1×

LeWM HIJP

183.0 14.9

368.5 24.3

682.3 41.4

669.6 41.6

475.8 30.6

Speedup

12.3×

15.2×

16.5×

16.1×

15.0×

LeWM HIJP

110.5 59.5

180.9 70.5

298.2 87.4

297.7 90.8

221.8 77.0

Speedup

1.9×

2.6×

3.4×

3.3×

2.8×

Table 1 shows that HIJP achieves the highest overall success rate of 86.8%, exceeding the strongest baseline by 4.9 percentage points. The gains are most pronounced in challenging manipulation settings. On PushT, HIJP achieves 66% and 42% success in the hard and extreme tiers, respectively, outperforming the best baseline in each tier by 8 and 14 percentage points. On OGCube, it maintains 96% success in both tiers, with corresponding gains of 26 and 10 percentage points. On TwoRoom and Reacher, HIJP remains competitive, ranking second in average success on both tasks. The stronger gains on manipulation are consistent with our design to extend action supervision beyond action prediction itself, shaping goal-conditioned visual representations and constraining the latent transitions induced by predicted actions. We examine the contributions of these two design choices through the ablations in Section 4.4. 4.3

P LANNING E FFICIENCY

To assess planning efficiency, we compare HIJP with LeWM, a CEM-based world-model planner, under matched numbers of replanning calls. Following Section 4.1, we report total wall-clock planning time for each task–tier combination, with all episodes processed as a single vectorized batch on an RTX 4090 (Table 2). The result shows that HIJP reduces planning time across all evaluated settings, achieving 1.9–16.5× speedups over LeWM under matched replanning counts. The average speedups are 6.2×, 7.1×, 15.0×, and 2.8× on TwoRoom, Reacher, PushT, and OGCube, respectively. Together with the results in Table 1, these findings show that HIJP achieves substantial planning speedups over LeWM while attaining higher average success rates on all four tasks. 4.4

A BLATIONS

To systematically evaluate the contribution of each component, we conduct ablation studies on the training objectives, gradient propagation, hindsight goal relabeling, and the test-time planner. For the training objectives, we remove one term at a time from Eq. 8: the inverse loss Linv , the transitionconsistency loss Lcyc , and the SIGReg regularizer Lsigreg . We further evaluate two gradient-related variants. In sg[zg ], we stop gradients through the goal representation zg . In the frozen-gϕ variant, we retain the latent predictor gϕ in the forward computation but block parameter updates from Lcyc while retaining those from Lpred. To evaluate the contribution of hindsight goal relabeling, we introduce a fixed-goal variant that always uses the latent representation of the final observation in each training window as the goal. Specifically, we sample an anchor step t uniformly from {0, . . . , T − 2}, set zg = zT −1 , and predict the one-step action at from (zt , zg ). 7

For the planner ablation, we keep the fully trained representation fixed and replace our amortized rollout with CEM search, thereby isolating the effect of the planning mechanism. All other settings follow Section 4.1. We report results on Reacher, PushT, and OGCube in Table 3. We exclude TwoRoom because all variants achieve near-100% success rates on this task. Furthermore, we examine sensitivity to the inverse and transition-consistency loss weights in Appendix C.

Table 3: Success rates for HIJP and its ablation variants. Avg. reports the arithmetic mean across the four difficulty tiers for each task, and Overall Avg. averages all 12 task–tier combinations. The best and second-best results in each row are shown in bold and underlined text. Full model Task

Reacher

PushT

OGCube

Action Consistency

Planner

Tier

HIJP

w/o Linv

Training objectives w/o Lsigreg

Endpoint goal

sg[zg ]

w/o Lcyc

Frozen gϕ

CEM

Easy Default Hard Extreme

48% 92% 92% 94%

36% 72% 90% 86%

2% 0% 0% 0%

36% 66% 72% 82%

40% 88% 88% 92%

34% 78% 84% 94%

40% 84% 80% 84%

40% 70% 98% 94%

Avg.

81.5%

71.0%

0.5%

64.0%

77.0%

72.5%

72.0%

75.5%

Easy Default Hard Extreme

90% 80% 66% 42%

36% 6% 0% 0%

74% 58% 34% 20%

76% 74% 50% 44%

90% 80% 48% 38%

84% 78% 54% 22%

86% 72% 64% 46%

78% 74% 44% 24%

Avg.

69.5%

10.5%

46.5%

61.0%

64.0%

59.5%

67.0%

55.0%

Easy Default Hard Extreme

98% 96% 96% 96%

80% 56% 46% 40%

98% 96% 90% 92%

92% 88% 88% 92%

94% 96% 94% 94%

98% 96% 94% 96%

98% 96% 94% 96%

96% 80% 58% 76%

Avg.

96.5%

55.5%

94.0%

90.0%

94.5%

96.0%

96.0%

77.5%

Overall Avg.

82.5%

45.7%

47.0%

71.7%

78.5%

76.0%

78.3%

69.3%

HIJP (ours)

PushT norm. latent distance to goal

Goal supervision

CEM (HIJP ablation)

LeWM (baseline)

TwoRoom

Reacher

OGCube

1.00 0.75 0.50 0.25 0.00

0

20

40

60

environment step

80

0

20

40

60

80

0

environment step

20

40

60

environment step

80

0

20

40

60

80

environment step

Figure 2: Normalized latent distance to the goal for HIJP, its CEM ablation, and LeWM across four tasks on the extreme tier. Lines show means, and shaded regions indicate pointwise 95% confidence intervals. HIJP and its CEM ablation share the same latent space, while LeWM uses its own encoder. The results show that removing Linv causes the largest overall decline, reducing average success from 82.5% to 45.7%. Removing SIGReg also reduces average success substantially, to 47.0%, with near-collapse on Reacher. These results support the importance of action supervision and representation regularization for HIJP’s control performance. The two goal-supervision ablations examine different parts of our design choice. Restricting goals to the final observation of each window reduces average success to 71.7%, while stopping gradients through zg reduces it to 78.5%. This supports using multiple future observations as goals and allowing action-related gradients to shape goal representations. The action consistency ablations examine the contribution of the transition-consistency objective. Removing Lcyc lowers average success to 76.0%, with the clearest decreases on Reacher and PushT and little change on OGCube. Blocking predictor updates from Lcyc lowers average success to 78.3%, primarily affecting Reacher. These results support the transition-consistency objective and its updates to the latent predictor, with gains varying across tasks. Finally, replacing the inverseguided rollout with CEM reduces average success to 69.3% under the matched evaluation protocol, supporting the effectiveness of HIJP’s learned planning mechanism. 8

4.5

L ATENT-S PACE P ROGRESS A NALYSIS

HIJP shows its largest gains on long-horizon goals. We next examine how the latent distance to the goal evolves during execution. We compare HIJP, its CEM ablation (Section 4.4), and LeWM. The CEM ablation shares HIJP’s learned representation and replaces HIJP’s amortized inverse-guided rollout with CEM search, isolating the contribution of the planner. On the extreme tier (K=75), we execute all methods in the original simulator on the same start–goal instances, following each method’s replanning schedule, and track the normalized latent distance to the goal, dt = ∥zt − zg ∥/∥z0 − zg ∥, where zt and zg are the encoded observation and goal. After an episode succeeds, we hold its last recorded distance constant for the remaining steps. Figure 2 reports the mean and pointwise 95% confidence intervals. Figure 2 shows that HIJP reaches a lower final normalized latent distance than its CEM ablation on all four tasks. On PushT, the methods exhibit similar initial progress, but HIJP continues to reduce the distance later in the episode as CEM largely plateaus. On TwoRoom and OGCube, HIJP reaches a lower plateau, while on Reacher it exhibits fewer fluctuations after the initial decrease. Compared with its CEM ablation, HIJP’s executed trajectories end closer to the goal in the shared latent representation.

4.6

S MOOTHNESS OF E XECUTED ACTIONS

We additionally evaluate the temporal smoothness of executed actions. We quan- Table 4: Action jerk on the extreme tier. Ratio denotes tify this with action jerk, the mean ℓ2 each comparator’s jerk divided by HIJP’s. The lowest norm of the second finite difference P of the jerk in each environment is shown in bold. 1 executed action sequence, T −2 t ∥at − HIJP CEM (ablation) LeWM 2at−1 + at−2 ∥2 (lower is smoother), following the convention of Nguyen et al. Environment Jerk ↓ Jerk ↓ Ratio ↑ Jerk ↓ Ratio ↑ 0.048 0.697 14.7× 0.724 15.2× (2026). To attribute any difference to the PushT 0.146 2.503 17.1× 2.558 17.5× planner rather than the representation, we TwoRoom 0.046 1.222 26.4× 1.134 24.5× compare HIJP against its CEM ablation, Reacher OGCube 0.099 2.065 20.9× 2.304 23.3× which shares the same learned representation and differs only in the planning rule, and include LeWM as a reference. Table 4 reports the extreme tier (K=75); the metric definition and full per-tier results are in Appendix D. The results show that HIJP produces smoother executed actions across all four tasks, with action jerk 14.7–26.4× lower than its CEM ablation and 15.2–24.5× lower than LeWM. These results indicate that HIJP’s amortized rollout improves the temporal smoothness of executed controls in addition to its goal-reaching performance and planning efficiency.

5

C ONCLUSION

We introduced Hindsight Inverse JEPA Planning (HIJP), a search-free framework for goalconditioned visual control. HIJP samples future observations as hindsight goals to enrich inversedynamics supervision and allows training gradients to update goal representations. It further feeds predicted actions into a latent predictor model and constrains the resulting latent transitions through a consistency objective. At planning time, the inverse and latent predictor models alternate to construct multi-step action plans without online search. Across four visual control tasks and multiple evaluation configurations, HIJP achieves the highest aggregate goal-reaching performance among the evaluated methods while reducing planning time compared to a search-based world-model planner. Under the most challenging setting, its executed actions are also smoother than those produced by CEM using the same representations. Together, these results suggest that aligning goal-conditioned action learning with the latent transitions used during planning provides an effective route toward efficient, search-free visual control. 9

6

L IMITATIONS

Our evaluation focuses on tasks with a limited temporal span, using fixed planning horizons and replanning schedules. Extending HIJP to substantially longer-horizon tasks remains an important direction. Future work will explore hierarchical subgoal generation, adaptive planning horizons, and dynamic replanning schedules. These extensions aim to decompose longer tasks into shorter goal-conditioned segments and adapt planning and execution intervals to the requirements of each segment.

R EFERENCES Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657, 2022. Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay, 2018. URL https://arxiv.org/abs/1707.01495. Mido Assran, Adrien Bardes, David Fan, Quentin Garrido, Russell Howes, Mojtaba, Komeili, Matthew Muckley, Ammar Rizvi, Claire Roberts, Koustuv Sinha, Artem Zholus, Sergio Arnaud, Abha Gejji, Ada Martin, Francois Robert Hogan, Daniel Dugas, Piotr Bojanowski, Vasil Khalidov, Patrick Labatut, Francisco Massa, Marc Szafraniec, Kapil Krishnakumar, Yong Li, Xiaodong Ma, Sarath Chandar, Franziska Meier, Yann LeCun, Michael Rabbat, and Nicolas Ballas. V-jepa 2: Self-supervised video models enable understanding, prediction and planning, 2025. URL https://arxiv.org/abs/2506.09985. Jiaxin Bai and Jiaxuan Xiong. Temporal-distance jepa: Plan-aware representation learning for latent world model predictive control. arXiv preprint arXiv:2607.25337, 2026. Randall Balestriero and Yann LeCun. Lejepa: Provable and scalable self-supervised learning without the heuristics. arXiv preprint arXiv:2511.08544, 2025. Weihao Bo, Shan Zhang, Yanpeng Sun, Jie Liu, Yongke Yao, Jinhao Du, Wei He, Kai Zou, Zechao Li, and Jingdong Wang. Diagram-mmu: A multi-modal benchmark for scientific diagrams. In European Conference on Computer Vision, pp. 524–543. Springer, 2026. Letian Cheng, Qi Zhang, and Yisen Wang. Sage: Subgoal-conditioned action generation for latent world model planning. arXiv preprint arXiv:2607.17973, 2026. Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL https: //openreview.net/forum?id=YicbFdNTTy. Yuntian Gao and Xiangyu Xu. Fast leworldmodel. arXiv preprint arXiv:2606.26217, 2026. Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Devin, Benjamin Eysenbach, and Sergey Levine. Learning to reach goals via iterated supervised learning. arXiv preprint arXiv:1912.06088, 2019. Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Manon Devin, Benjamin Eysenbach, and Sergey Levine. Learning to reach goals via iterated supervised learning. In International Conference on Learning Representations, 2021. URL https://openreview.net/ forum?id=rALA0Xo6yNJ. Hsiang-Wei Huang, Jianxu Shangguan, Junbin Lu, and Jenq-Neng Hwang. Leflow: Generative latent flow planning for world models. arXiv preprint arXiv:2608.24855, 2026. Petr Ivashkov, Randall Balestriero, and Bernhard Schölkopf. Sensorimotor world models: Perception for action via inverse dynamics, 2026. URL https://arxiv.org/abs/2606.20104. 10

Muyuan Liu, Yue Huang, Zheng Liang, and Xiang Gao. Toward physically grounded jepa world models for goal-conditioned robotic planning, 2026. URL https://arxiv.org/abs/ 2609.03565. Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun, and Randall Balestriero. Leworldmodel: Stable end-to-end joint-embedding predictive architecture from pixels. arXiv preprint arXiv:2603.19312, 2026. Sergi Masip, Jonathan Swinnen, Yutong Hu, Renaud Detry, and Tinne Tuytelaars. Ff-jepa: Longhorizon planning in world models with latent planners. arXiv preprint arXiv:2606.09311, 2026. Sacha Morin, Moonsub Byeon, Alexia Jolicoeur-Martineau, and Sébastien Lachapelle. On the sample efficiency of inverse dynamics models for semi-supervised imitation learning. arXiv preprint arXiv:2602.02762, 2026. Hoang Nguyen, Xiaohao Xu, and Xiaonan Huang. Latent geometry beyond search: Amortizing planning in world models, 2026. URL https://arxiv.org/abs/2605.08732. Keiran Paster, Sheila A McIlraith, and Jimmy Ba. Planning from pixels using inverse dynamics models. arXiv preprint arXiv:2012.02419, 2020. William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 4195–4205, October 2023. Cristina Pinneri, Shambhuraj Sawant, Sebastian Blaes, Jan Achterhold, Joerg Stueckler, Michal Rolinek, and Georg Martius. Sample-efficient cross-entropy method for real-time planning. In Conference on Robot Learning, pp. 1049–1065. PMLR, 2021. Ruslan Rakhimov, George Bredis, Yuriy Maksyuta, and Daniil Gavrilov. Qantara: Bridge-flow training for multi-paradigm jepa control. arXiv preprint arXiv:2607.04978, 2026. Reuven Rubinstein. The cross-entropy method for combinatorial and continuous optimization. Methodology and computing in applied probability, 1(2):127–190, 1999. Lukas Schäfer, Pallavi Choudhury, Abdelhak Lemkhenter, Chris Lovett, Somjit Nath, Luis França, Matheus Ribeiro Furtado de Mendonça, Alex Lamb, Riashat Islam, Siddhartha Sen, et al. When does predictive inverse dynamics outperform behavior cloning? arXiv preprint arXiv:2601.21718, 2026. Gawon Seo, Dongwon Kim, and Suha Kwak. Acid: Action consistency via inverse dynamics for planning with world models, 2026. URL https://arxiv.org/abs/2607.02403. Vlad Sobal, Wancong Zhang, Kyunghyun Cho, Randall Balestriero, Tim G. J. Rudner, and Yann LeCun. Stress-testing offline reward-free reinforcement learning: A case for planning with latent dynamics models. In 7th Robot Learning Workshop: Towards Robots with Human-Level Abilities, 2025. URL https://openreview.net/forum?id=jON7H6A9UU. Armin Sommer and Jannik Schilling. Reinforced planning with latent world models. arXiv preprint arXiv:2608.18669, 2026. Hao Sun, Zhizhong Li, Xiaotong Liu, Bolei Zhou, and Dahua Lin. Policy continuation with hindsight inverse dynamics. Advances in Neural Information Processing Systems, 32, 2019. Junhan Sun, Hao Zhao, and Guofeng Zhang. Intact: Isomorphic intent-to-action learning for searchfree world models, 2026. URL https://arxiv.org/abs/2607.26056. Wei Tang, Yanpeng Sun, Shan Zhang, Weihao Bo, Xiaofan Li, Piotr Koniusz, Wei Li, Na Zhao, and Zechao Li. Artemis: Structured visual reasoning for perception policy learning. arXiv preprint arXiv:2512.01988, 2025. Basile Terver, Tsung-Yen Yang, Jean Ponce, Adrien Bardes, and Yann LeCun. What drives success in physical planning with joint-embedding predictive world models? arXiv preprint arXiv:2512.24497, 2025. 11

Yang Tian, Sizhe Yang, Jia Zeng, Ping Wang, Dahua Lin, Hao Dong, and Jiangmiao Pang. Predictive inverse dynamics models are scalable learners for robotic manipulation. In International Conference on Learning Representations, volume 2025, pp. 92033–92052, 2025. Jiawei Wang, Ke Rui, Yushen Zuo, Yichun Feng, and Minglei Li. Decision-metric alignment in latent world models: Diagnostics and action-conditioned objectives for mpc planning. arXiv preprint arXiv:2608.18746, 2026a. Yuhai Wang, Jiawei Xia, Rongxuan Zhou, Xiao Hu, Yongliang Shi, Jing Du, and Yang Ye. Prism: Prior-guided imagination sampling in world models. arXiv preprint arXiv:2606.07974, 2026b. Grady Williams, Nolan Wagener, Brian Goldfain, Paul Drews, James M Rehg, Byron Boots, and Evangelos A Theodorou. Information theoretic mpc for model-based reinforcement learning. In 2017 IEEE international conference on robotics and automation (ICRA), pp. 1714–1721. IEEE, 2017. Amber Xie, Oleh Rybkin, Dorsa Sadigh, and Chelsea Finn. Latent diffusion planning for imitation learning. In Forty-second International Conference on Machine Learning, 2025. URL https: //openreview.net/forum?id=vhACnRfuYh. Yuxiang Zeng, Kun Xie, Yang Wang, Jigang Wen, Xiaocan Li, Guangxing Zhang, Gaogang Xie, and Jiannong Cao. Dta-gt: Direction- and topology-aware graph transformer for neural network representation learning, 2026. URL https://arxiv.org/abs/2610.16822. Zhenghao Zhang, Yuanxiang Wang, Zhenyu Guan, Yujia Yang, Bingkang Shi, Tianyu Zong, Hongzhu Yi, Guoqing Chao, Xingchen Chen, Tiankun Yang, Chenxi Bao, Tao Yu, Jingjing Zhou, and Jungang Xu. Delta-jepa: Learning action-sensitive world models via latent difference decoding, 2026. URL https://arxiv.org/abs/2606.31232. Zhiwei Zhang, Hui Zhang, Kaihong Huang, Chenghao Shi, and Huimin Lu. Efficient image-goal navigation with representative latent world model, 2025. URL https://arxiv.org/abs/ 2511.11011. Gaoyue Zhou, Hengkai Pan, Yann LeCun, and Lerrel Pinto. DINO-WM: World models on pretrained visual features enable zero-shot planning. In Forty-second International Conference on Machine Learning, 2025. URL https://openreview.net/forum?id=D5RNACOZEI.

12

A PPENDIX A

DATASET D ETAILS AND P ROTOCOL

We evaluate on four continuous-control benchmarks: Tworoom, Reacher, PushT, and OGBenchCube (hereafter OGCube), which together span 2D navigation, low-dimensional continuous control, planar object pushing, and robotic cube manipulation; Figure 3 shows a representative trajectory from each. We do not collect any trajectories ourselves: we reuse the full-scale offline datasets released by LeWM (Maes et al., 2026) without modification. A.1

DATASETS

The four datasets are gathered under different behavior policies: a goal-directed navigation policy for Tworoom, a near-uniform random policy for Reacher, and expert policies for PushT and OGCube. This heterogeneity is directly visible in their action distributions (Figure 4): Reacher actions are essentially uniform over [−1, 1], Tworoom actions saturate at the ±1 boundaries, whereas PushT and OGCube produce concentrated, structured actions. Each dataset is stored as a single HDF5 file under a shared schema: RGB observations pixels of shape (N, 224, 224, 3) in uint8, a continuous action array of shape (N, A), and per-episode boundaries encoded by ep offset and ep len, so that a sampled window never straddles two episodes. Table 5 summarizes the resulting statistics.

Tworoom

Reacher

PushT

OGCube

Figure 3: Task overview. Each row shows one representative trajectory per benchmark, sampled uniformly in time from its start (leftmost) to its end (rightmost).

Table 5: Statistics of the four datasets (full versions, as used in all experiments). Frames denotes the total number of transitions; Traj. length reports the minimum, mean, and maximum episode lengths. Dataset

Simulator

Tworoom Reacher PushT OGCube

2D navigation DMControl Reacher PushT OGBench Cube

Action dim A

#Episodes

Traj. length (min/mean/max)

#Frames

2 2 2 5

10,000 10,000 18,685 10,000

31 / 92 / 101 201 / 201 / 201 49 / 125 / 246 201 / 201 / 201

920,809 2,010,000 2,336,736 2,010,000

Observations and actions. All observations are 224 × 224 RGB frames, which we feed to the encoder at full resolution. The action space is a 2D command for Tworoom, Reacher, and PushT, 13

Tworoom

Reacher

goal-directed nav | dim=2, std=0.869

random policy | dim=2, std=0.577

a0 a1

8

a0 a1

density

6 4 2 Uniform

0

PushT

OGCube

expert | dim=2, std=0.214

expert | dim=5, std=0.421

a0 a1

8

a0 a1 a2 a3 a4

density

6 4 2 0

1.0

0.5

0.0

action value

0.5

1.0

1.0

0.5

0.0

action value

0.5

1.0

Figure 4: Per-dimension marginal action distributions for each dataset. The dashed line marks the density of a uniform distribution over [−1, 1]. and a 5D vector (end-effector pose and gripper) for OGCube. Reacher and OGCube episodes have a fixed length of 201 steps, whereas Tworoom and PushT episodes vary in length because they terminate upon task completion. Action normalization. Before training, every non-pixel input—the actions, and (where used) proprioception and state—is standardized per dimension by its dataset statistics, i.e., a z-score transform to zero mean and unit variance, applied identically to all four tasks. Figure 4 plots the corresponding raw (pre-normalization) action distributions, which differ markedly across tasks because of their different behavior policies. Training samples and monitoring split. Following the data-loading and window-level splitting procedure of the LeWM implementation (Maes et al., 2026), we enumerate all valid windows within each trajectory and randomly assign 90% to training and 10% to validation. Each window contains W observations sampled at intervals of k raw environment steps, together with aligned action chunks. The loader requires W k consecutive raw time-step records, while the first and last sampled observations are separated by (W − 1)k environment steps. We replace NaN action entries with zeros during training rather than dropping the window. A.2

E VALUATION P ROTOCOL

We follow the dataset-based evaluation procedure of the LeWM implementation (Maes et al., 2026). For each evaluation episode, we sample a start state from an offline trajectory and define the goal as the environment state K steps ahead (at index start + K). The agent acts for at most B steps. We count an episode as successful only if the simulator reports task completion within this budget Table 6; latent-space distance is not used to determine success. 14

Table 6: Environment-defined success criteria. Task

Success condition

Tworoom Reacher

Agent-to-goal position distance < 16 in environment coordinates. Absolute error of every joint angle < 0.05 radians relative to the target configuration. Euclidean error over agent and block positions < 20, and wrapped block orientation error < π/9. Target cube-to-goal position distance ≤ 0.04 m.

PushT OGCube

Table 7: Difficulty tiers used in evaluation. Each tier is fully specified by the goal offset K (how many steps ahead of the start the goal is placed) and the step budget B (the maximum number of environment steps allowed). Tier Goal offset K Step budget B easy default hard extreme

10 25 50 75

25 50 100 100

By varying the goal offset K and the step budget B, we obtain four difficulty tiers (Table 7; Figure 5), ranging from easy (K=10) to extreme (K=75): a larger K yields a longer-horizon goal. The extreme tier is the most challenging setting and is where our method attains its largest margin over the baselines. A.3

CEM P LANNING P ROTOCOL

CEM plays several roles in our experiments. It is used as the test-time optimizer for the CEM-based world-model baselines, as the computational reference in the planning-efficiency experiment, and as a controlled planner ablation in which HIJP’s amortized action generation is replaced by search over the same learned representation and forward dynamics. We therefore use one shared CEM configuration throughout all experiments. At every replanning point, CEM maintains a factorized Gaussian distribution over action sequences of length H = 5 model steps. Each model-level action contains a chunk of k = 5 environment actions, so every candidate sequence covers Hk = 25 environment steps. The sampling distribution is initialized with zero mean and sampling scale 1.0. At each optimization iteration, CEM samples N = 300 candidate action sequences and rolls each candidate forward through the corresponding latent dynamics model. Candidates are ranked by the squared distance between their terminal predicted latent and the encoded goal, 2

Jgoal (a0:H−1 ) = ∥ẑH (a0:H−1 ) − zg ∥2 .

(10)

The best 30 candidates are retained as elites and are used to refit the sampling mean and scale. This procedure is repeated for I = 30 iterations. Consequently, every CEM replanning call evaluates N I = 300 × 30 = 9,000

(11)

candidate action sequences. Since every sequence contains H = 5 latent transitions, this corresponds to approximately N IH = 300 × 30 × 5 = 45,000 (12) latent transition predictions per environment and replanning call. These operations are evaluated in parallel on the GPU; therefore, the sequence count should not be interpreted as an equivalent wall-clock slowdown. All CEM comparisons use the same planning horizon, action chunking, interaction budget, and replanning cadence as HIJP. Both planners generate H = 5 model-level actions and execute R = 5 model steps, equivalent to Rk = 25 environment steps, before observing the environment and 15

Start

Goal (harder ) default hard K=25 K=50

easy K=10

extreme K=75

Tworoom

Reacher

PushT

OGCube

Figure 5: Visualization of the four difficulty tiers. For each task (row), we fix a start observation (leftmost column) and show the goal observation placed K steps ahead of it, one per tier. A larger offset K moves the goal farther along the trajectory and demands a longer planning horizon.

Table 8: CEM configuration shared across all CEM-based comparisons. Hyperparameter

Value

Planning horizon H Action-chunk length k Execution stride R Environment steps per replan Rk Candidate sequences N Optimization iterations I Elite sequences Initial sampling scale Sequence evaluations per replan Latent transitions per replan

5 model steps 5 environment steps 5 model steps 25 300 30 30 1.0 9,000 45,000

replanning. Under the interaction budgets B ∈ {25, 50, 100, 100}, the maximum numbers of replanning calls are therefore 1, 2, 4, and 4 for the easy, default, hard, and extreme tiers, respectively. For the main LeWM baseline, CEM searches using LeWM’s own encoder and forward model. In the controlled CEM ablation, it instead searches using HIJP’s encoder and forward model while completely disabling the inverse-dynamics action generator. This controlled comparison changes only the test-time planning mechanism and therefore isolates the effect of amortized inverse-guided planning. 16

B

I MPLEMENTATION D ETAILS

This appendix provides a concrete instantiation of the method. We first specify the temporal discretization and the architecture of each component, then describe the training procedure and testtime planning protocol. Unless stated otherwise, the architecture and temporal settings are shared across all four tasks. B.1

T EMPORAL D ISCRETIZATION AND DATA F LOW

The environments expose atomic actions us ∈ Rda at raw environment step s. The world model operates at a coarser temporal resolution by grouping k consecutive atomic actions into one modellevel action chunk, at = (utk , utk+1 , . . . , utk+k−1 ) ∈ Rkda . We use k = 5 for every task. Consequently, one transition in the learned latent dynamics corresponds to five raw environment steps. The raw action dimension is da = 2 for Tworoom, Reacher, and PushT, and da = 5 for OGCube. The effective model-level action dimension is therefore da = kda ∈ {10, 25}. A training sample contains W temporally subsampled observations, together with their aligned model-level action chunks. k = 5 raw environment steps separate consecutive observations. Thus, the first and final observations in a window are separated by (W − 1)k raw environment steps. Before being passed to the model, images are converted to floating-point, normalized using ImageNet statistics, and resized to 224 × 224. Following the LeWM preprocessing procedure, continuous actions are standardized per dimension using statistics computed from the training dataset. At execution time, predicted action chunks are mapped back to the environment action space using the inverse transformation. B.2

A RCHITECTURE

The implementation contains a visual encoder fθ , an action encoder ψω , a latent predictor gϕ , and a hindsight-conditioned inverse-dynamics head hξ . The action encoder is absorbed into the notation of gϕ in the main text and is made explicit here. Visual encoder fθ . We instantiate fθ as a ViT-Tiny(Dosovitskiy et al., 2021) trained from scratch, without a pretrained backbone or an external visual representation. It contains 12 Transformer layers, 3 attention heads, and a hidden width of 192. Each 224 × 224 observation is divided into patches of size 14 × 14. We take the final[CLS] token and pass it through a projection MLP 192 → 2048 → 192. Batch normalization is applied at the hidden layer of the projection head. The resulting vector is the latent representation zt = fθ (ot ) ∈ R192 . Action encoder ψω . The flattened action chunk at ∈ Rda is embedded in two stages. A pointwise linear projection, implemented as a 1 × 1 one-dimensional convolution, first maps da → 10. A twolayer MLP then maps 10 → 768 → 192, with a SiLU activation between the two linear layers. This produces an action embedding with the same width as the visual latent, eat = ψω (at ) ∈ R192 . latent predictor gϕ . The predictor is a six-layer autoregressive Transformer with causal selfattention. It operates at width 192, uses 16 attention heads with head dimension 64, an MLP width of 2048, attention and MLP dropout of 0.1, and learned positional embeddings. The predictor receives a sequence of latent tokens and a time-aligned sequence of action embeddings. Actions condition every Transformer block through AdaLN-Zero(Peebles & Xie, 2023). In particular, each action embedding produces the shift, scale, and residual-gating parameters for both the attention and feed-forward sublayers; the modulation layer is initialized to zero. The predictor uses at most the m = 3 most recent latent-action pairs. For a context ending at time t, its last output is interpreted as the prediction of the next latent,ẑt+1 = gϕ (zt−m+1:t , ψω (at−m+1:t )) ,where the context is shortened near the beginning of a sequence. The Transformer output is mapped to the latent space by a second projection MLP 192 → 2048 → 192 with batch normalization. 17

Table 9: Architecture used in all experiments. Component

Configuration

Latent width Visual encoder fθ

d = 192 ViT-Tiny; 12 layers; 3 heads; width 192; patch size 14; 224 × 224 inputs; trained from scratch MLP 192 → 2048 → 192 with batch normalization 1 × 1 convolution da → 10, followed by MLP 10 → 768 → 192 with SiLU Causal Transformer; 6 layers; 16 heads; head dimension 64; MLP width 2048; dropout 0.1; AdaLN-Zero conditioning m = 3 latent–action pairs MLP 192 → 2048 → 192 with batch normalization Three hidden layers of width 512, each with LayerNorm and GELU k = 5 raw environment steps per model step

Visual projection Action encoder ψω latent predictor gϕ Predictor context Predictor projection Inverse head hξ Frameskip

Inverse-dynamics head hξ . future goal latent,

The inverse head receives the concatenation of a current latent and a

hξ : R2d → Rda , âi = hξ (zi , zi+κ ). (13) It consists of three hidden layers of width 512. Each hidden layer applies a linear transformation, LayerNorm, and GELU activation. A final linear layer maps to the model-level action dimension da . B.3

T RAINING

All four modules are optimized jointly. Training combines next-latent prediction, anti-collapse representation regularization, hindsight-conditioned inverse dynamics, and action transition consistency. latent prediction loss. Although every sampled window contains W = 16 model steps, the prediction term uses a causal context of length m = 3. Specifically, the predictor consumes z0:m−1 and the corresponding action embeddings and predicts the shifted latent sequence z1:m . Because causal attention is used, the output at position t only depends on positions up to t. Denoting the three predicted latents by ẑ1:m , we compute B m−1

Lpred =

2 1 X X (b) (b) ẑt+1 − zt+1 . Bmd 2 t=0

(14)

b=1

The target latent in this term remains attached to the computational graph, for which we do not employ a momentum target encoder. Anti-collapse regularization. We apply the Sketch Isotropic Gaussian Regularizer (SIGReg; Balestriero & LeCun, 2025) to the latents from all W frames, including frames that are not used by Equation 14. SIGReg measures the discrepancy between random one-dimensional projections of the latent distribution and a standard Gaussian reference distribution. We use 17 integration knots and 1024 random projections. Applying this objective over the complete window increases the number and temporal diversity of latent samples contributing to the regularizer. Hindsight-conditioned inverse dynamics. For each training window, we independently sample Na anchor–goal pairs. For the n-th pair, the anchor is sampled uniformly from all positions with an available successor, in ∼ Unif{0, . . . , W − 2}. (15) Conditioned on in , we sample a future offset uniformly from the remaining part of the window, κn ∼ Unif{1, . . . , W − 1 − in }.

(16)

The resulting goal index is in + κn . Sampling is performed independently for each pair and each element of the minibatch, so anchors may repeat. This procedure exposes the inverse head to goal conditions at multiple temporal distances while requiring no rewards or task labels. 18

Table 10: Task-specific training hyperparameters for HIJP. Task

Batch

Learning rate

Na

λreg

λinv

λcyc

Epochs

Scheduler

TwoRoom Reacher PushT OGCube

48 48 48 32

3 × 10−5 5 × 10−5 5 × 10−5 5 × 10−5

4 4 4 4

0.09 0.09 0.09 0.09

1.0 1.0 1.0 1.0

0.05 0.10 0.10 0.10

10 10 10 10

LWC LWC LWC LWC

For every sampled pair, the inverse head predicts the demonstrated action chunk executed at the anchor, âin = hξ (zin , zin +κn ). (17) Importantly, the target is the immediate action chunk ain , not the sequence of actions between the anchor and the sampled goal. The inverse loss is the element-wise mean absolute error N

B

Linv =

a XX 1 (b) (b) âin − ain . BNa da 1 n=1

(18)

b=1

Action transition consistency loss. The action consistency term couples the inverse head to the learned latent prediction. For each hindsight pair, we embed the predicted action and perform a single predictor step from the anchor latent, z̃in +1 = gϕ (zin , ψω (âin )) . (19) We then regress the prediction toward the actual next latent using B Na h i 2 1 XX (b) (b) , z̃in +1 − sg zin +1 Lcyc = BNa d 2 n=1

(20)

b=1

where sg denotes stop-gradient. Gradients from this loss update hξ , ψω , gϕ , and fθ , as well as the anchor-side representation through zin , while the next-latent target is detached. The single-step context used by this term is supported by the same autoregressive predictor used for the length-m forward objective. Joint objective.

The complete objective is L = Lpred + λreg Lsigreg + λinv Linv + λcyc Lcyc . (21) All terms are averaged over the batch and their corresponding temporal, anchor, action, or latent dimensions before the weights are applied. Optimization. We use a single AdamW optimizer for fθ , ψω , gϕ , and hξ , with weight decay 10−3 . Training uses bf16 mixed precision and a global gradient-norm clipping threshold of 1.0. The learning rate follows a linear-warmup cosine-decay schedule, starting from zero. With ten training epochs, the configured warmup is rounded to one epoch, followed by cosine decay to zero. Training and validation use the LeWM-based window-level split described in the dataset appendix A, with a 90%/10% ratio and random seed 3072. Algorithm 1 summarizes the complete training procedure, and Table 10 reports the task-specific training parameters. B.4

P LANNING AND R ECEDING -H ORIZON E XECUTION

At each replanning point, the current observation and target observation are encoded as z0 = fθ (ot ), zg = fθ (og ).

(22)

Starting from z0 , we construct an H-step action sequence through the autoregressive latent rollout introduced in the main text. At rollout step j, the inverse head first predicts the action chunk that is consistent with moving from the current imagined latent toward the goal latent, aj = hξ (zj , zg ). (23) 19

Algorithm 1 Joint training with hindsight-conditioned inverse dynamics Require: Offline dataset D; window length W ; predictor context length m; number of hindsight pairs Na 1: for each training minibatch do 2: Sample a minibatch of windows (o0:W −1 , a0:W −2 ) ∼ D 3: Encode observations and actions: zt ← fθ (ot ) and eat ← ψω (at ) 4: Predict the next latents: ẑ1:m ← gϕ (z0:m−1 , ea0:m−1 ) 5: Lpred ← MSE(ẑ1:m , z1:m ) 6: Lsigreg ← SIGReg(z0:W −1 ) 7: for n = 1, . . . , Na do 8: Sample in ∼ Unif{0, . . . , W − 2} 9: Sample κn ∼ Unif{1, . . . , W − 1 − in } 10: Predict the demonstrated action: âin ← hξ (zin , zin +κn ) 11: Predict the successor latent: z̃in +1 ← gϕ (zin , ψω (âin )) 12: end for   Na a 13: Linv ← MAE {âin }N n=1 , {ain }n=1   Na a 14: Lcyc ← MSE {z̃in +1 }N , {sg[z ]} in +1 n=1 n=1 15: L ← Lpred + λreg Lsigreg + λinv Linv + λcyc Lcyc 16: Update (θ, ω, ϕ, ξ) with AdamW 17: end for

The action chunk is embedded and passed to the latent predictor, which advances the imagined state,
zj+1 = gϕ zmax(0,j−m+1):j , ψω (amax(0,j−m+1):j ) . (24) The available context grows from one to m latent-action pairs during the first rollout steps and is subsequently truncated to the most recent m = 3 pairs. Thus, constructing an H-step plan requires H evaluations of the inverse head and H autoregressive evaluations of the latent predictor. No population of candidate action sequences or iterative test-time action optimization is required. We use H = 5 model steps for every task. Since each model-level action contains k = 5 atomic actions, one predicted plan covers Hk = 25 raw environment steps. The controller operates in a receding-horizon loop with execution stride R = 5 model steps. It therefore executes the complete predicted action sequence, re-encodes the resulting observation, and constructs a new plan. The target observation remains unchanged within an episode but is encoded together with the current observation at each replanning point. Execution terminates when the task-specific success condition is satisfied, or the allotted budget of raw environment steps is exhausted. Algorithm 2 gives the complete procedure. Table 11 collects the planning settings shared across the four tasks. Keeping these settings unchanged ensures that differences across tasks do not arise from task-specific planning budgets or model-rollout depths. Table 11: Planning settings shared across all four tasks. Hyperparameter

Value

Planning horizon H Action-chunk length k Raw planning horizon Hk Execution stride R Predictor context m

C

5 model steps 5 raw steps 25 environment steps 5 model steps 3 latent–action pairs

L OSS W EIGHT A BLATION

To examine the sensitivity of HIJP to the weighting coefficients of its training objectives, we conduct ablations on the inverse dynamics loss weight λinv and the cycle-consistency loss weight λcyc . 20

Algorithm 2 Amortized latent planning Require: Trained fθ , ψω , gϕ , and hξ ; goal observation og ; planning horizon H; execution stride R; action-chunk length k; predictor context length m; raw-step budget B 1: n ← 0 2: while n < B and the task has not succeeded do 3: Encode the current and goal observations: ẑ0 ← fθ (ot ), zg ← fθ (og ) 4: Initialize Z ← [ẑ0 ] and A ← [ ] 5: for j = 0, . . . , H − 1 do 6: Predict an action chunk: aj ← hξ (ẑj , zg ) 7: Append aj to A 8: Predict the next imagined latent: ẑj+1 ← gϕ (tailm (Z), ψω (tailm (A))) 9: Append ẑj+1 to Z 10: end for    11: q ← min R, H, B−n k 12: Execute the first q action chunks in A, terminating early if the task succeeds 13: Let s be the number of raw environment steps actually executed 14: n←n+s 15: end while

Table 12: Success rates (%) for the loss-weight ablations. Avg. reports the arithmetic mean across the four difficulty tiers, and Overall Avg. averages all 12 task–tier combinations. Bold column headers indicate the default weights. In the Overall Avg. row, the best and second-best results within each weight group are shown in bold and underlined text. λinv (λcyc = 0.1) Task

Reacher

PushT

λcyc (λinv = 1.0)

Tier

0

0.1

1.0

2.0

0

0.01

0.1

0.2

Easy Default Hard Extreme

36% 72% 90% 86%

42% 88% 86% 98%

48% 92% 92% 94%

44% 78% 90% 90%

34% 78% 84% 94%

52% 86% 90% 90%

48% 92% 92% 94%

44% 64% 78% 92%

Avg.

71.0% 78.5% 81.5% 75.5% 72.5% 79.5% 81.5% 69.5%

Easy Default Hard Extreme

36% 6% 0% 0%

Avg.

10.5% 65.0% 69.5% 73.5% 59.5% 69.0% 69.5% 71.5%

Easy Default OGCube Hard Extreme Avg. Overall Avg.

80% 56% 46% 40%

90% 78% 56% 36% 94% 94% 90% 96%

90% 80% 66% 42% 98% 96% 96% 96%

92% 82% 62% 58% 100% 94% 96% 96%

84% 78% 54% 22% 98% 96% 94% 96%

90% 72% 56% 58% 96% 94% 96% 94%

90% 80% 66% 42% 98% 96% 96% 96%

92% 82% 68% 44% 92% 92% 94% 94%

55.5% 93.5% 96.5% 96.5% 96.0% 95.0% 96.5% 93.0% 45.7% 79.0% 82.5% 81.8% 76.0% 81.2% 82.5% 78.0%

Specifically, we vary λinv ∈ {0, 0.1, 1.0, 2.0} while fixing λcyc = 0.1, and vary λcyc ∈ {0, 0.01, 0.1, 0.2} while fixing λinv = 1.0. We train a separate model for each weight setting and evaluate the resulting checkpoint using the same evaluation protocol as in the main experiments. All training and evaluation settings follow Section 4.1. We report the results on Reacher, PushT, and OGCube in Table 12. We omit TwoRoom because its performance is already saturated under our evaluation protocol, with near 100% success rates, making it uninformative for distinguishing different loss-weight choices. The default setting, λinv = 1.0 and λcyc = 0.1, achieves the highest overall average success rate of 82.5%. Increasing λinv to 2.0 improves performance on PushT and matches the default performance on OGCube, but reduces performance on Reacher, resulting in a slightly lower overall average of 81.8%. Similarly, decreasing λcyc to 0.01 or increasing it to 0.2 yields overall averages of 81.2% and 78.0%, respectively. These results indicate that the default weights provide the best overall balance across the three evaluated environments. 21

D

ACTION S MOOTHNESS E XPERIMENT

This appendix provides the complete evaluation protocol, metric definition, and per-tier results for the action-smoothness study in Section 4.6. We evaluate HIJP, its CEM ablation, and LeWM on the four domains described in Appendix A. For each domain and difficulty tier, all methods are evaluated closed-loop on the same set of start– goal instances. Instance sampling depends only on the dataset, random seed, goal offset, and number of evaluation episodes, and is therefore independent of the evaluated method. We use 50 episodes per task with random seed 42. The four difficulty tiers are defined by the goal-offset and step-budget pairs (K, B) listed in Table 7. Each rollout is allowed to run for the corresponding step budget, and we record the action executed at every environment step. HIJP directly executes the amortized plan produced by its inverse-guided controller and performs no test-time search. The CEM ablation uses the same learned representation and predictive model as HIJP, but replaces inverse-guided action generation with CEM planning. LeWM, which performs CEM planning over its own forward model, is included as an additional reference. All methods use the same control frequency and normalized action space. Table 13: Action jerk across four environments and four difficulty tiers. Avg. reports the arithmetic mean across tiers, with average ratios computed from the per-tier ratios. The lowest jerk in each row is shown in bold. HIJP CEM (ablation) LeWM Environment

Tier

Jerk ↓

Jerk ↓

Ratio ↑

Jerk ↓

Ratio ↑

PushT

easy default hard extreme Avg.

0.038 0.040 0.048 0.048 0.043

0.439 0.450 0.611 0.697 0.549

11.5× 11.2× 12.8× 14.7× 12.6×

0.451 0.407 0.595 0.724 0.544

11.8× 10.1× 12.4× 15.2× 12.4×

TwoRoom

easy default hard extreme Avg.

0.082 0.084 0.120 0.146 0.108

1.838 2.216 2.325 2.503 2.220

22.4× 26.3× 19.4× 17.1× 21.3×

1.774 2.191 2.428 2.558 2.238

21.6× 26.0× 20.2× 17.5× 21.3×

Reacher

easy default hard extreme Avg.

0.040 0.046 0.046 0.046 0.045

1.179 1.174 1.200 1.222 1.194

29.5× 25.6× 26.3× 26.4× 27.0×

1.099 1.114 1.134 1.134 1.120

27.5× 24.3× 24.8× 24.5× 25.3×

OGCube

easy default hard extreme Avg.

0.126 0.073 0.088 0.099 0.097

1.596 1.843 1.848 2.065 1.838

12.7× 25.2× 21.0× 20.9× 19.9×

1.517 1.886 2.159 2.304 1.967

12.1× 25.8× 24.5× 23.3× 21.4×

If an episode reaches its goal before exhausting the step budget, we truncate its action sequence at the first success step and compute smoothness only over the executed prefix. This avoids either rewarding or penalizing post-success idling and measures the controls actually required to complete the task. Unsuccessful episodes contribute their full budget-length action sequences. Let a1:T = (a1 , . . . , aT ) denote the executed action sequence of an episode. We define action jerk as the mean ℓ2 norm of the discrete second difference:

Jerk(a1:T ) =

T 1 X ∥at − 2at−1 + at−2 ∥2 . T − 2 t=3

22

(25)

Lower values indicate smoother controls. A temporally constant or linearly varying action sequence has zero second difference, whereas oscillatory or abruptly changing commands produce a large value. For each method, environment, and difficulty tier, we compute Eq. 25 independently for every episode and report the mean across episodes. Let J m denote this episode-averaged jerk for method m. We additionally report RCEM =

J CEM , J HIJP

RLeWM =

J LeWM . J HIJP

(26)

A ratio greater than one indicates that HIJP produces smoother actions than the corresponding CEMbased planner. Because all methods are evaluated in the same normalized action space, their jerk values are directly comparable. Table 13 reports the results across all four difficulty tiers. HIJP consistently achieves lower action jerk than both its CEM ablation and LeWM. On the extreme tier (K = 75), the CEM ablation produces 14.7×–26.4× higher jerk than HIJP, despite sharing the same learned representation and predictive model. This controlled comparison shows that the smoothness advantage primarily arises from HIJP’s amortized inverse-guided planning rather than from a difference in learned representations. LeWM exhibits a similar level of jerk to the CEM ablation, suggesting that independently optimizing successive plans with CEM can introduce substantial temporal discontinuities into the executed controls.

E

V ISUALIZING P LANNED T RAJECTORIES

For each environment, we run our planner (Algorithm 2) from a start observation toward a goal observation and record the resulting trajectory. Figures 6–9 visualize these rollouts as pixel-space filmstrips. Each figure corresponds to one environment and contains four difficulty tiers: Easy, Default, Hard, and Extreme. The tiers are defined by goal offsets of 10, 25, 50, and 75 environment steps, with corresponding planning budgets of 25, 50, 100, and 100 steps, respectively. For each tier, we show one rollout using 16 frames sampled uniformly over its execution, from 0% to 100%, and arrange them in two rows of eight frames. The goal observation appears separately on the right, outlined in red. The label beside each tier indicates whether the rollout reaches the goal within the planning budget.

23

0%

7%

13%

Push-T

20%

27%

33%

40%

47%

Goal

Easy success

53%

60%

67%

73%

80%

87%

93%

100%

Default success

Hard success

Extreme unsuccess

Figure 6: Representative trajectories produced by HIJP in the Push-T environment.

24

0%

7%

13%

Reacher

20%

27%

33%

40%

47%

Goal

Easy success

53%

60%

67%

73%

80%

87%

93%

100%

Default success

Hard success

Extreme unsuccess

Figure 7: Representative trajectories produced by HIJP in the Reacher environment.

25

0%

7%

13%

Two-Room 20%

27%

33%

40%

47%

Goal

Easy success

53%

60%

67%

73%

80%

87%

93%

100%

Default success

Hard success

Extreme unsuccess

Figure 8: Representative trajectories produced by HIJP in the TwoRoom environment.

26

0%

7%

13%

OGCube

20%

27%

33%

40%

47%

Goal

Easy success

53%

60%

67%

73%

80%

87%

93%

100%

Default success

Hard success

Extreme unsuccess

Figure 9: Representative trajectories produced by HIJP in the OGCube environment.

27

Related documents

• SAGE: Subgoal-Conditioned Action Generation for Latent World Model Planning#1118656
• LeWAM: End-to-End World and Action Modeling with JEPAs#1296825
• WAVE-0: Bridging VLA Policies and WAMs with Lightweight Latent Future Prediction#1344069
• Making Latent Evolution Explicit: Operator-Structured Transitions for World Action Models#1035606
• Staircase Policy: Streaming Inference for World-Action Models with Large Action Chunks#1131313