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

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

cs.ROOct 16, 2026

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

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

Organizations: Nanjing University of Science and Technology · University of Oxford · Singapore University of Technology and Design · Adelaide University

Abstract

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.

Figures & tables

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.

		Reference	Baselines							Ours
Task	Tier	Random	GCBC	PLDM	DINO-WM	LeWM	SMWM	$\Delta$ -JEPA	INTACT	HIJP
TwoRoom	Easy	66%	100%	100%	100%	96%	96%	76%	100%	100%
	Default	24%	100%	88%	100%	88%	96%	38%	100%	100%
	Hard	8%	100%	90%	98%	62%	86%	48%	96%	100%
	Extreme	2%	100%	96%	74%	50%	72%	24%	90%	98%
	Avg.	25.0%	100.0%	93.5%	93.0%	74.0%	87.5%	46.5%	96.5%	99.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 .

Task	Method	Difficulty tier				Avg.
Task	Method	Easy	Default	Hard	Extreme	Avg.
TwoRoom	LeWM	69.8	124.9	254.7	253.9	175.9
	HIJP	13.5	20.8	36.2	37.4	27.0
	Speedup	$5.2\times$	$6.0\times$	$7.0\times$	$6.8\times$	$6.2\times$
Reacher	LeWM	153.6	296.0	566.6	581.9	399.5
	HIJP	32.6	45.4	65.3	69.0	53.1

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 .

		Full model	Training objectives		Goal supervision		Action Consistency		Planner
Task	Tier	HIJP	w/o $\mathcal{L}_{\mathrm{inv}}$	w/o $\mathcal{L}_{\mathrm{sigreg}}$	Endpoint goal	$\operatorname{sg}[z_{g}]$	w/o $\mathcal{L}_{\mathrm{cyc}}$	Frozen $g_{\phi}$	CEM
Reacher	Easy	48%	36%	2%	36%	40%	34%	40%	40%
	Default	92%	72%	0%	66%	88%	78%	84%	70%
	Hard	92%	90%	0%	72%	88%	84%	80%	98%
	Extreme	94%	86%	0%	82%	92%	94%	84%	94%
	Avg.	81.5%	71.0%	0.5%	64.0%	77.0%	72.5%	72.0%	75.5%

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 .

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.

	HIJP	CEM (ablation)		LeWM
Environment	Jerk $\downarrow$	Jerk $\downarrow$	Ratio $\uparrow$	Jerk $\downarrow$	Ratio $\uparrow$
PushT	0.048	0.697	14.7 $\times$	0.724	15.2 $\times$
TwoRoom	0.146	2.503	17.1 $\times$	2.558	17.5 $\times$
Reacher	0.046	1.222	26.4 $\times$	1.134	24.5 $\times$
OGCube	0.099	2.065	20.9 $\times$	2.304	23.3 $\times$

Table 4: Action jerk on the extreme tier. Ratio denotes each comparator’s jerk divided by HIJP’s. The lowest jerk in each environment is shown in bold .

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

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

Dataset	Simulator	Action dim $A$	#Episodes	Traj. length (min/mean/max)	#Frames
Tworoom	$2$ D navigation	$2$	$10{,}000$	$31\,/\,92\,/\,101$	$920{,}809$
Reacher	DMControl Reacher	$2$	$10{,}000$	$201\,/\,201\,/\,201$	$2{,}010{,}000$
PushT	PushT	$2$	$18{,}685$	$49\,/\,125\,/\,246$	$2{,}336{,}736$
OGCube	OGBench Cube	$5$	$10{,}000$	$201\,/\,201\,/\,201$	$2{,}010{,}000$

Appendix

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.

Figure 4: Per-dimension marginal action distributions for each dataset. The dashed line marks the density of a uniform distribution over

[-1,1]

Task	Success condition
Tworoom	Agent-to-goal position distance $<16$ in environment coordinates.
Reacher	Absolute error of every joint angle $<0.05$ radians relative to the target configuration.
PushT	Euclidean error over agent and block positions $<20$ , and wrapped block orientation error $<\pi/9$ .
OGCube	Target cube-to-goal position distance $\leq 0.04\,\mathrm{m}$ .

Appendix

Table 6: Environment-defined success criteria.

Tier	Goal offset $K$	Step budget $B$
easy	$10$	$25$
default	$25$	$50$
hard	$50$	$100$
extreme	$75$	$100$

Appendix

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).

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.

Hyperparameter	Value
Planning horizon $H$	$5$ model steps
Action-chunk length $k$	$5$ environment steps
Execution stride $R$	$5$ model steps
Environment steps per replan $Rk$	$25$
Candidate sequences $N$	$300$
Optimization iterations $I$	$30$

Appendix

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

Component	Configuration
Latent width	$d=192$
Visual encoder $f_{\theta}$	ViT-Tiny; $12$ layers; $3$ heads; width $192$ ; patch size $14$ ; $224\times 224$ inputs; trained from scratch
Visual projection	MLP $192\!\rightarrow\!2048\!\rightarrow\!192$ with batch normalization
Action encoder $\psi_{\omega}$	$1\times 1$ convolution $d_{\mathbf{a}}\!\rightarrow\!10$ , followed by MLP $10\!\rightarrow\!768\!\rightarrow\!192$ with SiLU
latent predictor $g_{\phi}$	Causal Transformer; $6$ layers; $16$ heads; head dimension $64$ ; MLP width $2048$ ; dropout $0.1$ ; AdaLN-Zero conditioning
Predictor context	$m=3$ latent–action pairs

Appendix

Table 9: Architecture used in all experiments.

Task	Batch	Learning rate	$N_{a}$	$\lambda_{\mathrm{reg}}$	$\lambda_{\mathrm{inv}}$	$\lambda_{\mathrm{cyc}}$	Epochs	Scheduler
TwoRoom	48	$3\times 10^{-5}$	4	0.09	1.0	0.05	10	LWC
Reacher	48	$5\times 10^{-5}$	4	0.09	1.0	0.10	10	LWC
PushT	48	$5\times 10^{-5}$	4	0.09	1.0	0.10	10	LWC
OGCube	32	$5\times 10^{-5}$	4	0.09	1.0	0.10	10	LWC

Appendix

Table 10: Task-specific training hyperparameters for HIJP.

Hyperparameter	Value
Planning horizon $H$	$5$ model steps
Action-chunk length $k$	$5$ raw steps
Raw planning horizon $Hk$	$25$ environment steps
Execution stride $R$	$5$ model steps
Predictor context $m$	$3$ latent–action pairs

Appendix

Table 11: Planning settings shared across all four tasks.

\hat{z}_{j+1}\leftarrow g_{\phi}\!\left(\operatorname{tail}_{m}(\mathcal{Z}),\psi_{\omega}\!\left(\operatorname{tail}_{m}(\mathcal{A})\right)\right)

Appendix

Algorithm 2 Amortized latent planning

		$\lambda_{\mathrm{inv}}$ ( $\lambda_{\mathrm{cyc}}=0.1$ )				$\lambda_{\mathrm{cyc}}$ ( $\lambda_{\mathrm{inv}}=1.0$ )
Task	Tier	0	0.1	1.0	2.0	0	0.01	0.1	0.2
Reacher	Easy	36%	42%	48%	44%	34%	52%	48%	44%
	Default	72%	88%	92%	78%	78%	86%	92%	64%
	Hard	90%	86%	92%	90%	84%	90%	92%	78%
	Extreme	86%	98%	94%	90%	94%	90%	94%	92%
	Avg.	71.0%	78.5%	81.5%	75.5%	72.5%	79.5%	81.5%	69.5%

Appendix

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 .

		HIJP	CEM (ablation)		LeWM
Environment	Tier	Jerk $\downarrow$	Jerk $\downarrow$	Ratio $\uparrow$	Jerk $\downarrow$	Ratio $\uparrow$
PushT	easy	0.038	0.439	11.5 $\times$	0.451	11.8 $\times$
	default	0.040	0.450	11.2 $\times$	0.407	10.1 $\times$
	hard	0.048	0.611	12.8 $\times$	0.595	12.4 $\times$
	extreme	0.048	0.697	14.7 $\times$	0.724	15.2 $\times$
	Avg.	0.043	0.549	12.6 $\times$	0.544	12.4 $\times$

Appendix

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 .

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

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

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

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

Latent-Space Planning

See all 31 papers

H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning
Oct 5, 2026Wancong Zhang, Basile Terver, Michael Rabbat +2World Models for Robotics World Models
NYU · Advanced Machine Intelligence · INRIA Paris +1
FF-JEPA: Long-Horizon Planning in World Models with Latent Planners
Jun 8, 2026Sergi Masip, Jonathan Swinnen, Yutong Hu +2Hierarchical Planning Long-Horizon Robotic Manipulation
PSI, ESAT, KU Leuven · RAM, MECH, KU Leuven
Fast LeWorldModel
Jun 24, 2026Yuntian Gao, Xiangyu XuLatent Dynamics Modeling World Model-Based Planning
Xi’an Jiaotong University

Oct 5, 2026J/K move · Enter open · S save

H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning

Wancong Zhang, Basile Terver, Michael Rabbat, Yann LeCun+1

NYU · Advanced Machine Intelligence · INRIA Paris+1

World Models for Robotics World Models World Model-Based Planning Joint-Embedding Predictive Architecture Long-Horizon Planning Latent-Space Planning

Open paper AskPDF arXiv

Long-horizon planning with latent world models requires reasoning across timescales and levels of abstraction. Existing task-agnostic JEPA world models predict and plan at a single timescale or with multiple horizons in one shared latent space. We introduce H-JEPA, an end-to-end recipe for training a hierarchy of action-conditioned JEPAs in which each level predicts farther ahead in its own learned latent space. Planning proceeds top-down: the top level optimizes progress toward the goal, and each level's predictions become subgoals for the planner below it. When factors in the data evolve at separated timescales, higher levels discard fast, unpredictable detail and retain slower task-relevant state. Across four simulated navigation and manipulation environments, hierarchical planning improves over a flat JEPA; on Visual AntMaze, a three-level hierarchy raises success from 18% to 73% using less planner compute. Ablations attribute these gains to both temporal decomposition and higher-level goal representations. With inverse-dynamics supervision, the approach extends to diverse real-robot videos from DROID, where hierarchy improves offline planning fidelity at lower planner compute.

Abstract

Latent-Space Planning

Explore similar work

H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning

FF-JEPA: Long-Horizon Planning in World Models with Latent Planners

Fast LeWorldModel