Action Jets: Learning the Local Shape of the Action-conditioned Data Manifold

Anton Arapin; Chahyon Ku; Bernadette Bucher

cs.CVOct 13, 2026

Action Jets: Learning the Local Shape of the Action-conditioned Data Manifold

Authors: Anton Arapin, Chahyon Ku, Bernadette Bucher

Organizations: University of Michigan

Abstract

World models are a powerful means of reasoning and planning, but many of these action-conditioned predictive models are high dimensional and computationally expensive. These challenges limit the number of action candidates that can be evaluated by a planner. We show that this cost is avoidable for planners over short horizons, which only require local information with respect to the current state. For a fixed state, the frames reachable by varying a low-dimensional action (a camera movement or an end-effector command) lie on a low-dimensional manifold. We estimate the local shape of that manifold directly via a network trained to estimate how each image region moves per unit of each action axis supervised by data collected from samples by the generator. A correction head fit to trajectory data estimates the shape of this manifold region up to a third-order Taylor polynomial. We call this network estimating the Taylor polynomial describing the manifold shape an Action Jet and compute its radius of validity in closed form. By traversing this approximation of the local manifold shape, we only have to perform inference once per state for any number of candidate actions. Scoring sixteen candidates from a single extraction is over 300x faster than with either source generator, and this gap widens with the number of candidates a planner evaluates. We measure against state-of-the-art action-conditioned generators, IRASim and WorldGym on BridgeData V2, a real 7-DoF manipulation benchmark, and FrameCrafter on DL3DV scenes, where we chart from the novel-view generators DepthSplat and LagerNVS. Because only the tangent comes from the generator and every other term from the data, the Action Jet also predicts held-out frames more accurately than IRASim, the generator that supplied its tangent, at short horizons and at over 100x less compute.

Figures & tables

Panel of Figure 1 — Figure 1: A chart of the action-conditioned data manifold. (a) For a fixed state $x$ , varying the action traces a curved surface in latent space, the manifold $\mathcal{M}_{x}$ , on which transitions recorded from $x$ lie. (b) Probing the frozen generator $T$ at $\pm\delta$ (blue) gives the tangent plane at the current latent $z_{t}$ . While $T$ ’s own response bends away from the manifold, the Action Jet (orange) fits the Taylor polynomial to the recorded transitions in the data and so follows $\mathcal{M}_{x}$ up to the radius $r_{1}^{*}$ (green) it computes from its own coefficients. (c) One token of the chart is marked on the tangent plane in (b). The geometry encoder writes the fields of all $A$ action axes in one orthonormal frame $Q$ with coefficients $G$ .

Panel of Figure 2 — Figure 2: Method overview. Fill color marks what supervises each part: blue, the generator’s measured Jacobians; orange, recorded transitions; green, no network at all. In column 1, $z^{\pm\delta_{i}}=\varphi_{x}(\pm\delta_{i}e_{i})$ is the encoded render of each probe (Eq. 2 ).

	PSNR $\uparrow$ (gap 1/2/3/4)				LPIPS $\downarrow$ (gap 1/2/3/4)
tangent only	30.43	26.48	23.66	22.80	.0241	.0486	.0720	.0926
free-form Transformer head	34.45	31.19	28.28	27.26	.0153	.0267	.0391	.0511
3-jet correction head (ours)	31.80	29.05	26.92	26.27	.0183	.0305	.0403	.0501

Table 2: Kubric

4

-frame horizon prediction through the shared fine-tuned decoder. The free-form head re-runs a Transformer network for every candidate action, while the 3-jet answers any candidate by arithmetic on one extraction. The 3-jet reaches LPIPS parity at gaps 3–4 and trails by

1.0

–

2.6

dB in PSNR, a gap that shrinks with the horizon. The 3-jet dominates the 2-jet at matched decoder, and decoder fine-tuning is a strong rendering lever applied to all variants (App. Table 23 ).

	PSNR $\uparrow$				LPIPS $\downarrow$				policy
gap $g$	1	2	3	4	1	2	3	4	ADE $\downarrow$	FDE $\downarrow$
copy-frame (control)	24.48	21.48	20.04	19.19	.051	.084	.105	.121	1.577	1.994
WorldGym	23.80	21.10	19.70	18.86	.059	.093	.118	.137	1.448	1.783
IRASim (source generator)	23.34	22.37	21.64	21.12	.101	.116	.127	.136	0.846	1.035
Action Jet	24.21	22.07	20.67	19.67	.074	.110	.144	.177	1.474	1.992
$+$ fine-tuned decoder	24.57	22.42	21.07	20.11	.073	.111	.143	.170	1.299	1.663

Table 3: Open-loop rollouts on BridgeData V2 (

n{=}441

), and Octo policy consistency (

H{=}4

, a perfect simulator scores

0

). The Action Jet basis is charted from IRASim. Action Jet leads on short horizons against its generator.

Table 7

Table 8

Appendix figures & tables25 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Architecture. (A1) probing; (A2) the geometry encoder, emitting $Q$ and $G$ ; (A3) the correction head, whose coefficient fields depend on the state and a noise draw but never on the action. Everything in the red region is arithmetic on tensors already in memory.

Figure 6: Extracting and tracing a local chart. A held-out state is encoded once and the correction head’s coefficient fields are extracted once. The six branches then read that one extraction along each action axis. Each branch shows a surface of the local action manifold $\mathcal{M}_{x}$ and a corresponding decoded frame, predicted from the step along the labeled axis. Reading a surface render: $x$ and $y$ are commanded action in the plane spanned by the labeled axis and one companion axis, the latter present only so that the patch is a surface rather than a curve. The dashed circle is the query radius for that axis for that episode, and red dashed circle is the validity radius. The dashed line on the tangent plane is the departure from the tangent basis plane. The flat gray sheet is the first-order chart $z_{t}+QGa$ , which is the zero-height plane in this embedding and so is flat by construction. The blue sheet adds the correction head, $z_{t}+QGa+\tilde{r}(a,\varepsilon)$ . The two touch at the origin but the blue sheet leaves it with a nonzero slope. Every vertex of both sheets is arithmetic on the single extraction, no network is re-run anywhere in the figure, and only the displayed frames are decoded.

gap	1	2	3	4	6	8	10	12
IRASim basis	99.0	93.8	76.0	66.7	53.1	39.6	31.2	27.1
WorldGym basis	100.0	90.6	72.9	60.4	41.7	26.0	21.9	22.9
median $\|a\|/r_{1}^{*}$ (IRASim)	0.23	0.42	0.59	0.76	0.96	1.16	1.26	1.37

Appendix

Table 9: Real BridgeData V2 transitions inside

r_{1}^{*}

for both charts by gap. 96 states per gap. The two charts are built from different generators and use different action conventions, yet their coverage curves track each other closely.

$\|a\|/r_{1}^{*}$	0.22	0.43	0.63	0.83	1.02	1.26	1.70	2.24
jet (LPIPS)	0.039	0.063	0.080	0.084	0.101	0.130	0.154	0.192
copy	0.047	0.093	0.130	0.115	0.154	0.160	0.173	0.154
decoder floor	0.016	0.017	0.016	0.015	0.017	0.017	0.017	0.015
margin	$+0.008$	$+0.030$	$+0.051$	$+0.031$	$\mathbf{+0.053}$	$+0.030$	$+0.020$	$-0.038$
relative latent error
full jet	0.937	0.885	0.908	0.917	0.932	0.984	1.162	1.767

Appendix

Table 10: IRASim-basis chart on real Bridge frames, binned by

|a|/r_{1}^{*}

, median over 1152 transitions. Margin is copy minus jet, where positive means the chart beats predicting no motion. The margin peaks at the radius and reaches zero at

1.89\,r_{1}^{*}

. The lower block is an independent confirmation in latent space only.

Figure 7: The radius of Action Jet (geometry encoder of which was trained on IRASim generator) measured against real frames with no generator in the loop. Left: relative latent error of the full jet against its first-order truncation, with the stochastic fill level. Center: LPIPS against the real frame, with the decoder round-trip floor and the copy baseline. Right: share of real transitions inside $r_{1}^{*}$ by gap. Dashed vertical line is $r_{1}^{*}$ .

Figure 8: Qualitative results and comparison between WorldGym, IRASim and Action Jet. IRASim often hallucinates or performs slightly wrong motion, while WorldGym under-commands and can show wrong movement.

Figure 9: Advantage over the copy baseline for two Bridge charts built from different generators, each against its own $r_{1}^{*}$ . The charts have nearly orthogonal tangent fields and translation gains differing by $10\times$ , yet both stop beating the baseline at about twice their own predicted radius.

Figure 10: The WorldGym-basis chart, same three panels as Fig. 7 : latent error of the full jet against its first-order truncation, LPIPS against the real frame with the decoder floor and the copy baseline, and coverage inside $r_{1}^{*}$ by gap. This is the negative control of App. B.5.3 , a chart that anti-selects downstream and does not anchor on real transitions, and its radius is nonetheless calibrated. Being well calibrated about your own limits and charting the right object are separate properties, which is why we measure both.

probe axes	probe motion (LPIPS)	decoder floor	verdict
translations $t_{x},t_{y},t_{z}$	0.015–0.109	0.12–0.15	below the floor
rotations $r_{x},r_{y},r_{z}$	0.22–0.85	0.12–0.15	$1.5$ – $7\times$ above
Bridge, for contrast	floor $0.016$ ; every axis moves the image		measurable

Appendix

Table 11: Did the DL3DV probe move the picture? LPIPS between each probe render and the generator’s own render at the base action, against the decoder round-trip floor on the same frames. A decoded measurement below its own instrument’s noise floor carries no information, which is what rules the translation axes out of the pixel-space radius measurement on this benchmark (Fig. 11 ). The agreement results of Tables 8 and 13 compare latent fields directly and never pass through the decoder, so this floor does not limit them. Range over the six ladder intervals,

\{0.25\ldots 1.5\}\,r_{1}^{*}

n=64

states.

Figure 11: DL3DV degradation across the probe ladder. The chart-to-generator gap is what the ratio of $0.74$ – $0.85$ quoted above measures, and it is stable across the whole ladder. Left: all six axes on the shared normalized abscissa, log scale, with $r_{1}^{*}$ marked, with the flat translation lines measured in Table 11 . Center and right: the same data against the physical step, split because $0.5$ chart units and $0.5^{\circ}$ are not the same picture.

axis	$r_{1}^{*}$ (chart)	generator-implied	ratio	$\rho$ per state	$r_{1}^{*}/p_{95}(\|a\|)$
$t_{x}$	0.900	0.446	0.50	$+0.40$	0.7
$t_{y}$	0.646	0.283	0.44	$+0.24$	1.5
$t_{z}$	1.403	0.617	0.44	$+0.28$	1.7
$r_{x}$	0.480	0.250	0.52	$-0.02$	18.4
$r_{y}$	0.446	0.220	0.49	$-0.27$	6.8
$r_{z}$	0.426	0.217	0.51	$-0.20$	26.7

Appendix

Table 12: DL3DV: the learned chart’s

r_{1}^{*}

against the radius implied by the generator’s own probe responses, per axis,

n=64

states. The last column gives

r_{1}^{*}

as a multiple of the 95th percentile of the recorded action magnitude. Radii are for the DepthSplat chart, mean over the 64 ladder states, so they differ from Table 20 , which reports medians.

axis	Action Jets $\cos(F_{A},F_{B})$	generators $\cos(\widehat{J}_{A},\widehat{J}_{B})$	null
$t_{x}$	0.77 [0.71, 0.80]	0.72 [0.62, 0.79]	- 0.00
$t_{y}$	0.81 [0.74, 0.85]	0.73 [0.63, 0.79]	$0.00$
$t_{z}$	0.77 [0.70, 0.81]	0.68 [0.56, 0.75]	$0.00$
$r_{x}$	0.84 [0.78, 0.87]	0.69 [0.64, 0.73]	$0.00$
$r_{y}$	0.82 [0.79, 0.85]	0.84 [0.80, 0.87]	- 0.00
$r_{z}$	0.79 [0.72, 0.84]	0.86 [0.82, 0.89]	$0.00$

Appendix

Table 13: Cross-generator tangent agreement on DL3DV, median [IQR] over 256 states, probed at

0.25\,r_{1}^{*}

per axis. Action Jets column compares the two Action Jet charts learned from DepthSplat and LagerNVS probes. Generators column compares the DepthSplat and LagerNVS generators that supervised them, both of which are deterministic (App. B.5.1 ).

pair	$t_{x}$	$t_{y}$	$t_{z}$	$r_{x}$	$r_{y}$	$r_{z}$	mean	$n$
generator $\leftrightarrow$ generator
DL3DV: DepthSplat $\leftrightarrow$ LagerNVS	.81	.82	.81	.88	.89	.85	.84	148
Bridge: IRASim $\leftrightarrow$ WorldGym (span)	–	–	–	–	–	–	.036	148
generator $\leftrightarrow$ probed baseline
DL3DV: LagerNVS $\leftrightarrow$ FrameCrafter	.67	.60	.67	.78	.79	.69	.70	148
DL3DV: DepthSplat $\leftrightarrow$ FrameCrafter	.59	.55	.60	.72	.75	.63	.64	148

Appendix

Table 14: Robustness check for charts agreement: the same comparisons at a

1\sigma

probe, band-square window on scenes, 148 states. FrameCrafter is probed but not charted. Conclusions are unchanged, the absolute values are not comparable with Table 13 , because the two ladders step the axes by different physical amounts.

axis	Action Jets	generators	$\rho_{\mathrm{IRASim}}\geq$	$\rho_{\mathrm{WG}}$	ceiling $\geq$	$\cos(s_{A},s_{B})\leq$
$t_{x}$	$-0.02$	$-0.02$	0.48	0.87	0.65	0.04
$t_{y}$	- 0.06	- 0.03	0.51	0.87	0.67	0.05
$t_{z}$	- 0.08	- 0.05	0.40	0.89	0.59	0.08
$r_{x}$	$-0.09$	$-0.04$	0.19	0.81	0.39	0.11
$r_{y}$	- 0.09	- 0.04	0.23	0.84	0.44	0.10
$r_{z}$	- 0.12	- 0.05	0.08	0.79	0.25	0.21

Appendix

Table 15: The two generators for BridgeData V2 do not share a tangent signal.

\rho_{\mathrm{WG}}

is measured from a seed replicate, while

\rho_{\mathrm{IRASim}}

is the lower bound

\cos(F,\widehat{J})^{2}

of Eq. 9 , since IRASim is a diffusion model too slow to re-render. The final column is the implied upper bound on the cosine between the two generators’ noise-free responses.

	IRASim Action Jet	WorldGym Action Jet
held-out fit to own generator’s fields	$0.513$	$0.542$
real-transition anchoring (cos, mid-band, $n{=}48$ )	$\mathbf{+0.249}$	$\mathbf{-0.011}$
rolled-action control	$-0.014$	$+0.088$
anchoring gap (true $-$ control)	$\mathbf{+0.263}$	$\mathbf{-0.099}$
least-squares gain $\kappa^{*}$ (median)	$0.74$	$-0.10$
tangent vs. copy floor (PSNR, dB)	$13.49$ vs. $12.94$	$12.84$ vs. $12.94$

Appendix

Table 16: Fit to the generator cannot tell the two Action Jets apart, so a distillation metric would call them equally good, but anchoring against real held-out transitions separates them at once. The rolled-action control pairs each state’s fields with another state’s action, so it measures the anchoring any two arrays would show. The least-squares gain

\kappa^{*}

is the scalar that best maps the chart’s predicted displacement onto the real one, where negative means the chart’s best linear fit to the data points backwards.

rung pair	$t_{x}$	$t_{y}$	$t_{z}$	$r_{x}$	$r_{y}$	$r_{z}$	mean
$0.25\!\leftrightarrow\!0.5$	.67	.78	.73	.83	.33	.83	.69
$0.5\!\leftrightarrow\!1$	.52	.67	.62	.78	.22	.74	.59
$1\!\leftrightarrow\!2$	.36	.53	.49	.37	.25	.49	.41
$2\!\leftrightarrow\!3$	.44	.59	.57	.46	.42	.54	.50
$3\!\leftrightarrow\!6$	.26	.31	.34	.23	.40	.24	.29
Action Jet trained at $\delta{=}1\sigma$ , held-out cosine at probe scale $0.25/0.5/1/2/3/6$

Appendix

Table 17: Full curvature ladder (DepthSplat, DL3DV, 148 states): per-axis cosine between per-unit fields at consecutive probe magnitudes, and an Action Jet trained at

\delta{=}1\sigma

evaluated against every interval. Curvature tracks image flow per unit action, while

r_{y}

(pan), the largest-flow axis, is most curved, and the Action Jet peaks sharply at the displacement it was trained on and collapses off-scale.

control	artifact it would expose	measured
random-field floor	cosine inflated by dimensionality or the shared VAE	$0.001$
seed-paired repeat, same generator	disagreement that is only sampling noise	$0.573$
split-half reliability ceiling $\sqrt{\rho}$	a stochastic generator penalized for its noise	$0.89$ – $0.94$
per-state $6{\times}6$ mixing fit	disagreement recoverable by a linear relabelling	$0.12$ vs. $0.78$
kinematic convention certification	generators compared in incompatible action frames	$+1.0000$ vs. $-0.72$
rolled-state null	an Action Jet emitting a generic, state-independent field	$0.00$ – $0.09$

Appendix

Table 18: Controls, and the artifact each one rules out. Every agreement measurement in the paper is reported against a floor and, where the generator is stochastic, a ceiling.

Action Jet	window	$t_{x}$	$t_{y}$	$t_{z}$	$r_{x}$	$r_{y}$	$r_{z}$	mean	null
DL3DV / DepthSplat	full	.80	.82	.79	.86	.83	.84	.82	.01
DL3DV / LagerNVS	square	.73	.78	.73	.82	.79	.78	.77	.00
DL3DV / DepthSplat	square	.69	.73	.69	.78	.75	.74	.73	.00
Bridge / IRASim	–	.65	.68	.60	.41	.44	.30	.51	.04
Bridge / WorldGym	–	.55	.56	.62	.51	.54	.47	.54	.09

Appendix

Table 19: Geometry-encoder held-out field cosine against the Action Jet’s own generator, with rolled-state nulls at right. DL3DV Action Jets score about

1.5\times

Bridge Action Jets, tracking generator determinism and conditioning tightness. Encoder latency is

8.3

–

15.9

ms per state.

Figure 12: DL3DV forward prediction and geometric consistency qualitative results. Held-out scenes from DL3DV Benchmark dataset at target frame $t+4$ compared to Action Jet and FrameCrafter predictions. The overlay marks the MASt3R correspondences measured in Table 1 . Action Jet may blur high-frequency details or leave slight artifacts at repeated patterns, while the diffusion baseline shows camera-motion and reprojection errors.

Figure 13: Additional DL3DV forward prediction and geometric consistency qualitative results, continued from Fig. 12 .

Figure 14: Additional qualitative results for BridgeData V2 open-loop rollouts. Held-out episodes under the recorded actions at gap $4$ against the real frames and the source generators. These are the rollouts scored in Table 3 .

	DL3DV (LagerNVS chart)		DL3DV (DepthSplat chart)		Bridge (IRASim chart)
axis	$r_{1}^{*}$	fill	$r_{1}^{*}$	fill	$r_{1}^{*}$ / fill
$t_{x}$	0.80 [0.68,0.92]	1.2	0.88 [0.73,1.07]	1.4	1.89 [1.65,2.14] / 0.71
$t_{y}$	0.53 [0.45,0.64]	1.3	0.55 [0.50,0.68]	1.3	2.10 [1.83,2.44] / 0.30
$t_{z}$	1.16 [1.02,1.34]	1.1	1.21 [1.08,1.54]	1.4	2.13 [1.79,2.39] / 0.97
$r_{x}$	0.47 [0.45,0.50]	8.7	0.48 [0.44,0.52]	8.5	3.33 [2.80,4.04] / 0.18
$r_{y}$	0.44 [0.42,0.46]	3.5	0.45 [0.41,0.47]	3.7	3.02 [2.44,3.68] / 0.10

Appendix

Table 20: Self-reported validity radii

r_{1}^{*}

(median [IQR]) and stochastic fill floors per unit radius, in normalized chart units where

1.0

is one probe displacement. The normalization is what makes the two embodiments comparable.

method	$g$	$n$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	tail LPIPS $\downarrow$	$\Delta$ LPIPS vs. ours [95% CI]
copy-frame (control)	1	370	18.63	0.506	.122	.220	$[-0.045,-0.035]$
	2	370	16.02	0.372	.205	.347	$[-0.008,+0.005]$
	3	370	14.83	0.313	.267	.422	$[+0.011,+0.025]$
	4	370	14.11	0.282	.315	.471	$[+0.025,+0.039]$
DepthSplat	1	370	25.64	0.862	.047	.076	$[-0.119,-0.110]$
	2	370	23.17	0.800	.075	.125	$[-0.138,-0.127]$

Appendix

Table 21: All methods, causal tier (two trailing context frames), common 256-pixel window. The copy control and the difficulty strata use exactly our states. Tail LPIPS averages the hardest third of states, ranked by how much the copy control misses.

\Delta

LPIPS is paired (row minus ours) on shared states; negative means better than ours. An interpolation tier (six surrounding context views) is not reported: its context includes the target frame at

g{\geq}2

	PSNR $\uparrow$				LPIPS $\downarrow$
gap $g$	1	2	3	4	1	2	3	4
LagerNVS generator
single extraction, FT decoder	21.2	19.7	18.5	17.6	.139	.172	.209	.243
composed (AR), FT decoder	21.2	18.6	16.8	15.5	.139	.217	.284	.344
composed (AR), frozen decoder	21.0	18.0	16.1	14.8	.158	.273	.388	.478
DepthSplat generator

Appendix

Table 22: Ablations of our chart on the

n{=}100

states per gap on all four gaps. Composing four single-step queries autoregressively is worse than one extrapolated query at

g{\geq}2

(paired

\Delta

LPIPS

+0.045

+0.103

, every CI excluding zero): re-encoding each decoded frame compounds error faster than the calibrated jet loses accuracy in extrapolation. The fine-tuned decoder’s gain grows under composition (LPIPS

+0.016

g{=}1

+0.153

g{=}4

, frozen minus fine-tuned), because every step re-encodes the previous decode. At

g{=}1

both extraction modes coincide exactly, a pipeline-identity check. The full-window row is not decoder-matched and is shown for feasibility only.

	decoder	gap 1		gap 2		gap 3		gap 4
free-form	frozen	33.54	.0325	30.32	.0638	27.53	.0942	26.24	.1249
free-form	FT	34.45	.0153	31.19	.0267	28.28	.0391	27.26	.0511
free-form+pixFT	frozen	32.45	.0285	28.73	.0558	25.96	.0884	24.54	.1275
free-form+pixFT	FT	33.49	.0165	30.17	.0274	27.72	.0381	26.77	.0484
2-jet+pixFT	frozen	28.69	.0332	26.09	.0680	23.65	.1033	22.76	.1373
2-jet+pixFT	FT	30.46	.0213	27.93	.0350	25.86	.0468	25.33	.0575

Appendix

Table 23: Kubric full ablation ladder. Correction-head ladder

\times

decoder (PSNR / LPIPS per gap). “+pixFT” marks training with optional adversarial loss. The shared decoder was fine-tuned on {3-jet, free-form+pixFT} latents; the out-of-mix plain free-form model rendering best through it evidences that the shared renderer plays no favorites.

Efficient Inference for World Action Models

Efficient inference for world action models cuts the cost of predicting action-conditioned futures for responsive control and planning.

See all 79 papers

Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding
Jun 30, 2026Zhenghao Zhang, Yuanxiang Wang, Zhenyu Guan +11World Models World Model-Based Planning
School of Computer Science and Technology, University of Chinese Academy of Sciences, Beijing · Institute of Information Engineering, Chinese Academy of Sciences, Beijing · School of Computer Science and Technology, Harbin Institute of Technology, Weihai +2
Fast LeWorldModel
Jun 24, 2026Yuntian Gao, Xiangyu XuLatent Dynamics Modeling World Model-Based Planning
Xi’an Jiaotong University
ActionRoPE: Action as Coordinates for Spatially Consistent Control in Planar-Camera Game World Models
Oct 19, 2026Zhiyang Deng, Yuzhe Wu, Zijun Lin +1Action-Conditioned Video Generation World Models
Tencent · National University of Singapore · Nanyang Technological University

Jun 30, 2026J/K move · Enter open · S save

Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding

Zhenghao Zhang, Yuanxiang Wang, Zhenyu Guan, Yujia Yang+10

School of Computer Science and Technology, University of Chinese Academy of Sciences, Beijing · Institute of Information Engineering, Chinese Academy of Sciences, Beijing · School of Computer Science and Technology, Harbin Institute of Technology, Weihai+2

World Models World Model-Based Planning Joint-Embedding Predictive Architecture Action-Conditioned World Models Model-Based Planning Latent World Models

Open paper AskPDF arXiv

Learning visual world models for planning requires compact latent dynamics that remain sensitive to actions, yet reconstruction-free joint-embedding objectives can collapse to action-insensitive representations. We propose Delta-JEPA, an end-to-end reconstruction-free world model that augments latent forward prediction with a Latent Difference Action Decoder (LDAD). Unlike inverse decoders that infer actions from concatenated endpoint embeddings, LDAD reconstructs the executed action from the latent displacement between consecutive observations. This displacement-level supervision directly regularizes transition geometry: adjacent embeddings cannot collapse without losing action information, and different actions are encouraged to induce distinguishable latent changes for rollout-based planning. Delta-JEPA uses only latent prediction and action reconstruction, avoiding pixel reconstruction and distribution-matching regularizers. Across four visual continuous-control tasks, Delta-JEPA improves planning over JEPA-based and representation-learning world model baselines. Ablations show that displacement-based action decoding is consistently more effective than endpoint concatenation, and action-sensitivity analyses show clearer action-conditioned latent responses. These results indicate that supervising latent differences is a simple and effective mechanism for collapse-resistant and action-sensitive world model learning.

Abstract

Efficient Inference for World Action Models

Explore similar work

Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding

Fast LeWorldModel

ActionRoPE: Action as Coordinates for Spatially Consistent Control in Planar-Camera Game World Models