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

Action Jets teaser figure — One encoder pass and one coefficient extraction per state. Every candidate action after that is arithmetic on tensors already in memory.

Abstract

World models are a powerful means of reasoning and planning for a variety of applications including robotics. However, many of these world models such as video generation models are high dimensional and computationally expensive, which limits 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 need 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 describes the shape of this manifold region up to a third-order Taylor polynomial. We call the 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. We are over 300× faster at sixteen candidates compared to the source generators, 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 DepthSplat, LagerNVS and FrameCrafter on DL3DV scenes. Because only the first-order term comes from the generator and the rest 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 roughly 10⁴× less compute.

1.8ms

per candidate action, with the network run twice per state no matter how many candidates are scored

0.80

tangent agreement between charts distilled from two unrelated generators, against a permutation null of 0.00

99%

of single-step robot transitions fall inside the chart's own closed-form validity radius

+1.2dB

over the generator that supplied the tangent, at the next frame, at ~10⁴× less compute

Charting the action-conditioned data manifold

Latent image states live in the space of a frozen image autoencoder $(E, D)$ : $z = E (x)$ . We split $z$ into $N$ tokens, each a $2 \times 2$ block of latent pixels, so every token is a vector in $R^{d}$ . For an image state $x$ , let $M_{x}$ be the set of latent states reachable from $x$ by actions in a neighbourhood $U$ of zero. This is a smooth surface of dimension at most $A$ , fixed by the data distribution including the dynamics. Write $ψ_{x} : U \to M_{x}$ for the map that sends an action to the latent state it reaches in the data (its mean, where transitions are stochastic), so that $ψ_{x} (0) = z_{t} : = E (x)$ . It is a property of the dataset, not of any model.

A generator $T$ trained on that data defines $φ_{x} (a) = E (T (x, a))$ , and where it has learned the data, $φ_{x} \approx ψ_{x}$ : an approximate chart of the action-conditioned data manifold. The local shape of $M_{x}$ is its Taylor expansion,

ψ_{x} (a) = z_{t} + i \sum a_{i} J_{i} + \frac{1}{2} i, j \sum a_{i} a_{j} H_{ij} + \frac{1}{6} i, j, k \sum a_{i} a_{j} a_{k} C_{ij k} + \dots

whose first-order fields $J_{i} = \partial ψ_{x} / \partial a_{i}$ span the tangent space. An Action Jet is a network that emits a truncation of that expansion from the current frame, with the first order measured from a generator and everything above it fit to data.

**Method overview.** Fill colour marks what supervises each part: the generator's measured Jacobians, recorded transitions, or no network at all. The generator is queried only once, offline.

1. Probe the generator for the action Jacobian

For each generator, state and axis $i$ we render at $\pm δ_{i}$ along that axis alone, encode both renders, and difference them into the field

J_{i} = \frac{φ _{x} ( δ _{i} e _{i} ) - φ _{x} ( - δ _{i} e _{i} )}{2 δ _{i}} \in R^{N \times d}

$J_{i}$ is the estimated action Jacobian for axis $i$ , and it says how each token moves per unit of that axis. Both renders show the same room, table and robot, and subtracting removes all of it, so the measurement records the generator's response to the action rather than the quality of its rendering. For stochastic generators we average over samples. The probe needs no gradients, so any model that maps a state and an action to an image can be probed, including multi-step diffusion samplers.

2. Estimate the first-order term with the geometry encoder

Probing costs two generator passes per axis per state, which a planner scoring many candidates cannot afford. We therefore train a geometry encoder that predicts the action Jacobian from the current frame in one forward pass, supervised by the probed fields.

A ViT reads $z_{t}$ and one context frame $z_{t - 1}$ as $N$ tokens. For each token $n$ it emits a feature vector $h_{n}$ , used later by the correction head, a basis $B_{n} \in R^{d \times m}$ , and coefficients $G_{n} \in R^{m \times A}$ . Taking $Q_{n}$ to be the $Q$ factor of the reduced QR decomposition of $B_{n}$ makes $Q_{n}^{⊤} Q_{n} = I_{m}$ hold by construction, and token $n$ 's motion per unit of axis $i$ is predicted as $F_{i, n} = Q_{n} G_{n} e_{i}$ . All $A$ axes share that one rank- $m$ frame, which makes the output a tangent space rather than a set of unrelated fields. The encoder minimises

L_{enc} = i = 1 \sum A \frac{∥ F _{i} - J _{i} ∥ ^{2}}{∥ J _{i} ∥ ^{2}} + λ_{cov} i = 1 \sum A \frac{∥( I - Q Q ^{⊤} ) J _{i} ∥}{∥ J _{i} ∥}

where the norms are taken over all tokens and $(I - Q Q^{⊤})$ acts token-wise. Normalising each axis by its own energy keeps the axes that cause large image motion from dominating the gradient; the second term prices the part of each probed field that lies outside the predicted frame. Stacking tokens gives the encoder's estimate of $J_{i}$ , and with it the first-order truncation

ψ_{x}^{(1)} (a) = z_{t} + i = 1 \sum A a_{i} F_{i}

3. Fit the higher orders to the data

The first-order chart is a plane and $M_{x}$ is a curved surface. A correction head adds the second- and third-order terms. Unlike the geometry encoder it is not supervised by the generator: it is trained on recorded transitions, so every term above first order is fit to the data itself. For each token it reads the encoder feature $h_{n}$ and a noise draw $ε_{n}$ and emits the coefficient fields of

r (a, ε) = i \sum a_{i} L_{i} + \frac{1}{2} i, j \sum a_{i} a_{j} H_{ij} + \frac{1}{6} i, j, k \sum a_{i} a_{j} a_{k} C_{ij k} + i \sum a_{i} U_{i}^{(1)} (ε) + i, j \sum a_{i} a_{j} U_{ij}^{(2)} (ε)

where every field lies in $R^{N \times d}$ . $H_{ij}$ and $C_{ij k}$ are the head's estimates of the second- and third-order coefficients, $L_{i}$ corrects the first-order term, and $U^{(1)}, U^{(2)}$ are noise-conditioned fill fields for disoccluded content, which is a discontinuous function of the action and has no polynomial form. The prediction for action $a$ is

z^(a, ε) = ψ_{x}^{(1)} (a) + r ~ (a, ε) = z_{t} + i \sum a_{i} F_{i} + r ~ (a, ε)

with $r ~ (a, ε) : = r (a, ε) - r (0, ε)$ , which guarantees $z^(0, ε) = z_{t}$ . The head is trained with the encoder frozen.

Every coefficient field is computed from $h$ and $ε$ , and never from $a$ . That is the property the whole method rests on: the fields are extracted once per state, and each candidate action after that costs only a contraction against tensors already in memory. It is also why an Action Jet can be more accurate than the generator it came from. The generator's contribution ends at first order, supplying the directions in which the image can move at each state, including along single axes the dataset never isolates. Recorded data sets every coefficient above it.

4. Query candidate futures inside the radius of validity

The estimate does not hold for actions that move the state too far. A camera moved somewhere with no visual overlap with the current scene cannot have its new view predicted from the old one, and computationally, a Taylor polynomial stops being accurate far from its expansion point. Because the prediction is an explicit polynomial in $a$ , that boundary can be computed from the chart's own coefficients rather than found by a sweep.

Let $u$ be a unit direction in normalised action units. The predicted displacement from $z_{t}$ at step length $r$ is exactly a cubic,

P (r) = E_{ε} [z^(r u, ε)] - z_{t} = α r + β r^{2} + γ r^{3}

with $α = \sum_{i} u_{i} (F_{i} + L_{i} + E_{ε} [U_{i}^{(1)}])$ , $β = \frac{1}{2} H [u, u] = \frac{1}{2} \sum_{i, j} u_{i} u_{j} H_{ij}$ and $γ = \frac{1}{6} C [u, u, u]$ . Writing $p_{1} = ∥ α ∥$ , $p_{2} = ∥ β ∥$ and $p_{3} = ∥ γ ∥$ , the radius of validity is the positive root of $p_{2} r^{2} + p_{3} r^{3} = p_{1} r$ ,

r_{1}^{*} (u) = \frac{- p _{2} + p _{2}^{2} + 4 p _{1} p _{3}}{2 p _{3}}

past which the nonlinear terms can be as large as the linear one. Inside $r_{1}^{*}$ any number of further candidate actions can be simulated from the current state; outside it, a new chart has to be extracted. The whole quantity costs six evaluations of the correction head and no rendering. It is a property of the Action Jet alone, with no generator and no ground-truth frame anywhere in it, so it is available per state at inference time.

Sliding past the radius. One held-out episode, one extracted chart, and a slide along a single action axis (t_x) with the chart's own radius

r_{1}^{*}

marked on the surface. Inside it the decoded frame stays sharp. As the slide passes

r_{1}^{*}

, the nonlinear terms overtake the linear one and the image loses its sharpness. No network is re-run along the slide.

On BridgeData V2 we check it against real held-out frames with no generator in the loop. 99% of single-step transitions and 67% of four-step ones fall inside $r_{1}^{*}$ , and the chart stops beating the do-nothing baseline at $1.89 r_{1}^{*}$ . Repeating the whole measurement on a second chart built from a different generator, with nearly orthogonal tangent fields and translation gains a factor of ten apart, puts the crossing at $2.04 r_{1}^{*}$ . The two agree about where they stop working to within 8%.

One extraction, any number of candidates

Screening candidate actions with a world model has one shape: every candidate costs a full forward simulation. An action jet pays two network evaluations per state regardless of K, because the action never enters the network. Sixteen candidates cost 1.4× one and sixty-four cost 3.1×, so the cost per candidate falls as breadth grows. Ranking happens in latent space, so the decoder runs once rather than K times.

world model	NFE / state	K=1	16	64	ms / cand
IRASim	25K	5714	91431	—	5478
WorldGym	40K	1240	17501	—	1094
Action Jet, composed	8	87.8	604.3	2347.9	36.7
Action Jet, single	2	37.9	54.0	117.5	1.8

Measured cost in milliseconds of evaluating K candidate four-step futures at one state, identical hardware.

Breadth search on three held-out robot states. A goal-conditioned policy proposes candidate action chunks; the action jet predicts every outcome from a single extraction, and the selector keeps the one whose predicted future is closest to the target.

Kubric: one factor at a time

Kubric scenes are synthetic, so every factor of variation is declared and can be moved on its own: the camera's azimuth, the light's azimuth and the object's position. The generator we chart here is a feed-forward next-frame transformer, which is differentiable, so the field the probe estimates is also available exactly by automatic differentiation. The probed fields match it at cosine 0.87 for the camera, 0.90 for the light and 0.99 for the object. All three axes share one rank-12 frame, and the effective rank of the probed energy in that frame is 3.9, against three true factors. In the grid below each row is a different scene, and each column drives one factor through the 3-jet prediction; the last column moves all three at once.

Camera rotation

Light azimuth

Object translation

Combined

3-jet predictions on held-out Kubric scenes. One scene per row. Each column drives one factor, and the last drives all three together.

Predicting future frames

Only the first order comes from the generator, so an Action Jet is not capped by the model it was distilled from. We check that on two domains that fail in opposite directions: a robot arm, where consecutive frames barely differ and doing nothing is a strong baseline, and indoor scenes, where the camera moves far enough that a prediction has to invent pixels.

BridgeData V2

Under the recorded actions, the chart predicts the next frame at 24.57 dB against IRASim's 23.34, so it beats the generator that supplied its tangent by 1.2 dB at roughly 100× less compute. By gap 2 the two are level, 22.42 against 22.37, and from gap 3 IRASim pulls ahead by about 0.6-1.0 dB, which is the price of composing steps rather than a limit of the chart. Fine-tuning the decoder on real pixels is worth a further 0.3–0.4 dB at every gap.

DL3DV

On indoor scenes the comparison that matters is against the other method that has to synthesize pixels from two causal frames. Against FrameCrafter, a 25-step diffusion sampler, the chart predicts at 20.5 dB to its 17.4 at gap 1 and 16.8 to 13.7 at gap 4, a margin of 3.0 to 3.3 dB held across the horizon, with LPIPS .159 against .221 narrowing to .279 against .429. The copy-frame control wins LPIPS at gap 1, .136 against .159, and is passed from gap 2 onward. DepthSplat and LagerNVS sit above both, but they re-render real pixels from posed views rather than synthesizing them, so they are an upper bound rather than a competitor. Swapping which of them supplied the tangent moves the result by at most 0.006 LPIPS and 0.08 dB at every gap.

Pixel metrics cannot separate a frame that moved the camera correctly from one that blurred toward the mean, and on this domain the copy control exploits exactly that. So we score the geometry instead: pair each prediction with the real photograph at the same target pose, match with MASt3R, recover the relative camera by PnP, and note that its ground truth is the identity. A photograph against itself gives the instrument's noise, 0.12° and 0.4 px. At gap 4 the chart cuts the control's rotation error by 54% and its reprojection error by 69%, 0.79° and 5.9 px against 1.73° and 19.1 px, and is below the control on 97% of states. It is the only generative method below the control at all: FrameCrafter's frames imply a camera farther from the target than the input frame itself, 2.16° and 51.1 px, below the control on 23% of states. A model can win LPIPS by predicting nothing. It cannot win this.

DL3DV predictions and geometric consistency — **DL3DV predictions** against the real photograph at the same target pose, with the MASt3R correspondences the reprojection numbers are computed from. Click to enlarge.

Two generators, one manifold

An extracted chart could describe the dataset's manifold or only the generator we probed. The comparison that separates them is well posed because of where supervision enters: the geometry encoder sees only generator probes, so the tangent field is the one part of the model that the objection “both were trained on the same real frames” cannot reach.

On DL3DV, two charts distilled from architecturally unrelated novel-view generators — DepthSplat, explicit Gaussian splatting on estimated depth, and LagerNVS, a transformer initialised from VGGT — agree at mean per-axis cosine 0.80 against a state-permutation null of 0.00. Their generators agree at 0.75, and because both generators are deterministic that 0.75 is an unattenuated ceiling rather than a noise-limited one. On four of six axes the two learned bases agree with each other more than the generators they came from do.

	t_x	t_y	t_z	r_x	r_y	r_z	mean
DL3DV — DepthSplat and LagerNVS, deterministic generators
learned charts	.77	.81	.77	.84	.82	.79	.80
generators	.72	.73	.68	.69	.84	.86	.75
permutation null	.00	.00	.00	.00	.00	.00	.00
BridgeData V2 — IRASim and WorldGym
learned charts	−.02	.06	.08	−.09	.09	.12	.04
generators	−.02	.03	.05	−.04	.04	.05	.02

Tangent agreement between learned charts built from different generators. Where the generators agree, the learned bases recover the shared structure and agree at least as well. Where the generators disagree, on Bridge, the learned charts inherit that disagreement rather than manufacturing consensus — which is what makes the DL3DV result evidence about the data rather than about shared architecture.

Tracing the manifold

Each clip below walks the demonstrated action out from one extracted chart, on a different held-out episode. The surfaces are measured patches of $M_{x}$ , four tangent pairs at once: the flat grey sheet is the first-order chart alone, the coloured sheet adds the correction head, and the marker tracks the current step. The frames beside them are decoded from that same single extraction. No network is re-run anywhere in the sweep.

Episode 1

Episode 2

Episode 3