Discarded Predictions: Rewriting Action-Chunking's Waste as Latent Distributions, Reality-Grounded Supervision, and Lipschitz Geometry

A research note — an idea, not a result. The thread: action chunking throws away a huge amount of information on every inference. If we treat the future latent as a distribution, supervise it against real observations, and then bolt on Lipschitz / contraction geometry, we can upgrade “prediction error” from a scalar loss into a bounded, falsifiable certificate — and, with that, recast control as a constraint-satisfaction problem. Every figure below is a live p5.js demo; drag the sliders and the 3-D scenes as you read.

1. The starting point: action chunking, and what it quietly throws away

Modern VLA policies — ACT, Diffusion Policy, the π0 / π0.5 flow-matching family — mostly run on action chunking: predict a whole future action sequence $a_{t:t+H}$ at once (π0.5 uses $H=50$, about a second), execute only the first $s \le H$ steps, then re-observe and re-predict. Two hard-nosed reasons this became the default:

First, reality itself is the dynamics model. During the open-loop $s$ steps, the policy needs no learned forward model — it isn’t simulating a rollout, it’s actually moving the robot and letting the next real frame report what happened. The error is physical, and physics corrects it for free.

Second, actions are a privileged interface. They’re the one thing the robot can consume, and demonstrations come with action labels, so predicting actions is the most direct, best-grounded supervision you can get.

But precisely because it works so well, the paradigm quietly drops two kinds of information. The first is the unexecuted tail of every chunk — those $H-s$ steps. In a receding-horizon sense, discarding the tail is correct for control: it was computed from the stale observation $o_t$; by the time you reach $t+s$ you already have the fresher $o_{t+s}$, which should replace it. This “waste” is not a bug — drag $s$ below to see why.

The second kind is subtler and bigger: the policy’s implicit belief about the future. To predict a good $a_{t:t+H}$ at all, the model must have implicitly modeled “what the world will look like after these actions.” That belief is never exposed and never supervised — and it’s the genuinely large signal being wasted. The discarded tail, paired with the trajectory reality actually went on to take, is a free self-supervised label: “you predicted this, the world did that.” Scrub time below to harvest those pairs.

2. Make the implicit belief explicit — and ground it in reality

So let’s not throw the belief away. Have a latent forward predictor $g$ roll the current latent $z_t$ and the known actions into a future latent, and supervise it against the latent encoded from the real future observation:

\[\mathcal{L}_{\text{pred}}=\sum_{k=1}^{H} d\Big(\,g_k\big(z_t,\,a_{t:t+k}\big),\ \ \mathrm{sg}\!\big[\mathrm{enc}_{\text{target}}(o_{t+k})\big]\,\Big)\]

with $\mathrm{sg}[\cdot]$ a stop-gradient, $\mathrm{enc}_{\text{target}}$ an EMA target encoder, $d$ usually a cosine distance. This is exactly the self-predictive / latent-consistency machinery from model-based RL (SPR, TD-MPC2, DreamerV3), and it can be bolted straight onto a π0.5-style policy as a co-training head.

Plain version: predict the future in latent space, then grade the prediction against what reality actually became — never against your own guess.

That last clause is the trap you must not step in. If the target is the model’s own rolled-out latent, the objective is self-referential and the gradient drags every latent onto a single constant point — classic representation collapse. The whole JEPA stack of stop-gradient / EMA / variance-covariance regularizers exists to prevent exactly this. Let the target come from the real $o_{t+k}$, and the prediction-vs-reality discrepancy becomes a real gradient that can’t collapse. Flip the toggle below and watch it happen.

This is a training-time auxiliary loss (the target only exists after you’ve executed), good for relabeling already-collected trajectories offline or for online continual learning — it is not a way to speed up a single inference.

3. Treat the latent as a distribution, not a point

A flow / diffusion policy already outputs a distribution, not a point estimate — so the future latent should be one too. The immediate payoff is keeping multimodality. “Go left” and “go right” around an obstacle are both valid; a point estimate, or naive time-averaging across chunks, crushes them into the mean — the straight line through the obstacle, the single most dangerous action, one nobody proposed. Drag the slider from “distribution” to “point estimate” and watch the collision risk climb.

The same point bites at chunk boundaries. Splice two perfectly-valid chunks together naively — a hard cut, or a blind average — and the executed action jumps, which means a spike in acceleration: the kind of thing that snaps a wrist or trips a safety stop. RTC’s fix is to condition the new chunk on the already-committed one (“inpainting”), i.e. match the distribution across the seam, not the point. Widen the blend below and watch the spike vanish.

The catch is honest: distribution matching (KL / optimal transport / score matching) is harder and less stable to train than a consistency loss. No free lunch — but the multimodality is worth paying for.

4. Lipschitz & contraction — making the latent route trustworthy

Here’s the question that’s been hanging over everything: rolling out in latent space compounds error, and there’s no ground truth in latent space to measure that error against. This is the central difficulty of model-based control, and it’s the prerequisite for steps 2–3 to mean anything.

Contraction is the key that pins it down. If the latent dynamics $F(\cdot,a)$ is $\kappa$-Lipschitz with $\kappa<1$ (a contraction), then $k$-step rollout error converges geometrically instead of exploding:

\[\|e_k\|\le\kappa^{k}\,\|e_0\|+\frac{1}{1-\kappa}\,\sup_j\|\delta_j\|\]

Plain version: each step multiplies yesterday’s error by κ and adds fresh noise. If κ<1, old mistakes fade away and the new noise can only pile into a finite puddle of height supδ/(1−κ). If κ≥1, errors feed on themselves and avalanche. Only when the dynamics is contractive is “correct the latent once, then roll” well-posed. Drag κ across 1 below.

Once the latent is a distribution, the natural way to measure “prediction vs reality diverged” is the Wasserstein distance — and Wasserstein and Lipschitz are dual (Kantorovich–Rubinstein):

\[W_1(\mu,\nu)=\sup_{\mathrm{Lip}(g)\le 1}\Big(\mathbb{E}_{\mu}[g]-\mathbb{E}_{\nu}[g]\Big)\]

Plain version: the cost of shovelling one pile of sand onto another equals the best score any “gentle ruler” — a function whose slope never exceeds 1 — can give by rewarding μ and docking ν. Drag the ruler below; it can never beat the true $W_1$, and the best one touches it exactly.

And if the observation map $f$ is $L$-Lipschitz, pushing both distributions through it can’t blow their distance up by more than $L$:

\[W_1(f_{\#}\mu,\,f_{\#}\nu)\ \le\ L\,W_1(\mu,\nu)\]

That’s a strict exchange rate between “divergence I see in observation space” and “distribution drift in latent space” — the error stops being a dimensionless scalar and becomes a bounded, geometric quantity.

The 1-D error bound is easier to believe once you can walk around it. Below, twenty corrected latents are rolled forward by the same map: with $\kappa<1$ they all spiral into one fixed point (the initial spread is forgotten — so it doesn’t matter exactly where you corrected to); push $\kappa$ past 1 and the same map flings them apart. Drag to orbit.

5. The inverse projection — the crux

The genuinely interesting direction is backwards: from a divergence you see in observation space (that “swerve”), infer how much the latent should be corrected. That needs not the forward Lipschitz (an upper bound) but a co-Lipschitz lower bound:

\[\ell\,\|z_1-z_2\|\le\|f(z_1)-f(z_2)\|\quad\Longrightarrow\quad\|z_1-z_2\|\le\tfrac{1}{\ell}\,\|f(z_1)-f(z_2)\|\]

Both directions together is bi-Lipschitz:

\[\ell\,\|z_1-z_2\|\ \le\ \|f(z_1)-f(z_2)\|\ \le\ L\,\|z_1-z_2\|\]

Plain version: f is a fair currency exchange — it never shrinks a distance below ℓ× or grows it above L×. The lower rail ℓ is the one that lets you read a latent gap off an observation gap. Drag a point below; the output stays sandwiched in the ring between the ℓ and L circles, and the inner ring dies as ℓ→0.

I want to be honest: this step is the crux of the whole idea, and the hardest part. Forward Lipschitz is cheap (a dab of spectral normalization). The lower bound $\ell$ is brutal: deep nets collapse information by default (many latent directions map to the same output), and without a lower bound a tiny observation change can correspond to an enormous latent change — the inverse projection just fails. The 3-D scene below is observation space: $f$ squashes one axis by $\ell$, so a fixed little observation gap pulls back to a latent needle of length $d/\ell$. Drag $\ell$ toward 0 and the needle explodes.

Getting bi-Lipschitz essentially needs an invertible architecture — which collides nicely with “treat the latent as a distribution and sample it”: a normalizing flow is bijective by construction, bi-Lipschitz-ish, and gives exact densities. That’s not a coincidence; the constraint is choosing the architecture for us. One more engineering reality: global Lipschitz constants in deep nets are almost always vacuous — what’s usable is a local, per-region certificate, not one global number.

5.5 Why a normalizing flow actually is bi-Lipschitz

I claimed a normalizing flow hands you bi-Lipschitz almost for free. Since that’s the one load-bearing assumption, it deserves a derivation.

Step 1 — bi-Lipschitz is a statement about the Jacobian’s singular values. Zoom in near a point $z$ and the map looks linear, $f(z+\delta)\approx f(z)+J(z)\delta$, with $J$ the Jacobian. The local stretch in direction $\delta$ is $|J\delta|/|\delta|$, and over all directions that ratio sweeps exactly the gap between the smallest and largest singular values of $J$:

\[\sigma_{\min}(J)\ \le\ \frac{\|J\delta\|}{\|\delta\|}\ \le\ \sigma_{\max}(J)\]

So the local upper rail is $L=\sigma_{\max}$ and the local lower rail is $\ell=\sigma_{\min}$; a map is bi-Lipschitz on a region exactly when $\sigma_{\min}\ge\ell>0$ and $\sigma_{\max}\le L$ across it. Geometrically a tiny circle maps to an ellipse with semi-axes $\sigma_{\max},\sigma_{\min}$ — bi-Lipschitz means that ellipse never collapses to a line.

Plain version: the little ellipse you get by pushing a little circle through the map is never a pancake and never a black hole — its fattest and thinnest radii both stay finite and above zero. Drag the probe and watch it.

A coupling-layer flow warps the grid without ever folding it (a bijection), and the unit circle at the probe becomes an ellipse whose axes are exactly the local $\sigma_{\max}$ and $\sigma_{\min}$ — the local $L$ and $\ell$. Drag the probe; raise the layers and the clamp.

Step 2 — invertible ⟺ no fold ⟺ the lower rail survives. $f$ is invertible exactly when $\det J=\prod_i\sigma_i\neq 0$ everywhere — no direction is ever crushed to zero, i.e. $\sigma_{\min}>0$. That is the same statement as “the warped grid never folds onto itself.” And the inverse has Jacobian $J^{-1}$, whose singular values are $1/\sigma_i$, so

\[\mathrm{Lip}(f^{-1})=\frac{1}{\sigma_{\min}(J)}=\frac{1}{\ell}\]

That $1/\ell$ is literally the inverse-projection bound from §5 — the radius of the latent trust region. A flow lets you read a latent gap off an observation gap precisely because it keeps $\sigma_{\min}$ off the floor.

Step 3 — two constructions that guarantee it on purpose.

(a) Coupling layers (RealNVP / NICE). Split the latent in two, $z=(z_a,z_b)$, and update one half conditioned on the other: $z_b’=z_b\odot e^{s(z_a)}+t(z_a)$. The Jacobian is triangular, so $\det=\prod e^{s}=e^{\sum s}$, and the scale $s$ controls the singular values. Clamp $s\in[-c,c]$ and the layer is bi-Lipschitz with rails $[e^{-c},e^{c}]$ — explicit, cheap, exact. That’s the flow in the demo above.

(b) Residual flows (i-ResNet). Take $f(z)=z+g(z)$ and force $\mathrm{Lip}(g)=\kappa<1$ with spectral normalization. The reverse triangle inequality gives

\[(1-\kappa)\,\|\delta\|\ \le\ \|f(z+\delta)-f(z)\|\ \le\ (1+\kappa)\,\|\delta\|\]

so $\ell=1-\kappa$, $L=1+\kappa$ — bi-Lipschitz, and invertible by the Banach fixed-point theorem precisely because $\kappa<1$. Look at what that is: the very same $\kappa<1$ that made the latent rollout converge back in §4 is what makes the flow invertible here. One inequality, third appearance. Push $\kappa$ past 1 below and the grid folds.

Residual flow $f=z+g$: while $\mathrm{Lip}(g)=\kappa<1$ the grid is a clean bijection (rails $[1-\kappa,1+\kappa]$); cross $\kappa=1$ and red folds appear where $\det J$ flips sign — invertibility is lost. The same $\kappa<1$ as the rollout.

Step 4 — stacking, and the density you get for free. Composition is submultiplicative: for $f=f_L\circ\cdots\circ f_1$ we get $\ell\ge\prod_i\ell_i$ and $L\le\prod_i L_i$, so a deep stack stays bi-Lipschitz — the rails just compound (that’s the “layers” slider in the first demo). And the same nonzero Jacobian that buys invertibility also buys an exact density through change-of-variables, $\log p(x)=\log p(z)-\log\lvert\det J\rvert$. The lower bound that makes the inverse projection possible is the very same one that makes the likelihood well-defined — which is why “treat the latent as a distribution” (§3) and “invert observation gaps into latent corrections” (§5) end up wanting the same architecture. The constraint isn’t a nuisance; it’s quietly choosing the model for us.

6. Redefining the problem: from loss to certificate

Put it together and the lower bound pays off. You measure an observation gap $d$ between prediction and reality; the co-Lipschitz bound guarantees the true latent correction lives inside a ball of radius $d/\ell$ around your current latent. Not a guess — a certificate. With a decent $\ell$ it’s tight and you fix the latent without searching; shrink $\ell$ and the ball swallows the space and goes vacuous. Drag both sliders below.

And here’s the upgrade that makes the whole detour worth it. The model’s dynamics asserts that the next latent must land inside a $\kappa$-reachable set. Encode the real observation and check:

• Inside the set → an ordinary belief update, plus a free reality-grounded training target (exactly §2).
• Outside the set (more than the Lipschitz step allows) → this is not a big prediction error to average away. It’s the world doing something the model swore was impossible: out-of-distribution, or an unmodelled disturbance.

Control becomes a constraint-satisfaction / certificate problem: a boundary violation is itself a high-value event — it can drive a surprise / curiosity signal, trigger full replanning, or mark the state as somewhere to go collect data. A plain MSE blurs “small wobble where the model is confident” and “the model was just falsified” into one number; this framing keeps them apart and gives the second one a crisp, actionable meaning. Drag “reality” in and out of the ring.

7. One thread running through all of it

Worth stressing: the single condition of contraction holds up three different things at once.

1. It makes latent-rollout error converge geometrically → it answers “won’t rolling out in latent space drift?” (§4, the 3-D field).
2. It makes the distributional supervision signal bounded and measurable, via the Wasserstein–Lipschitz duality → it makes §2–§3 interpretable.
3. It makes an observation-space divergence invertible into a latent trust region, via bi-Lipschitz → it makes the §5–§6 certificate exist.

Three apparently separate ideas, one shared mathematical premise. That’s also why “betting on which map satisfies which property” is the whole ballgame.

8. The design choices that decide everything

Whether any of this stands up depends on a few not-yet-pinned-down choices:

• Lipschitz of which map, in which metric? Encoder $o\to z$, decoder $z\to o$, dynamics $(z,a)\to z’$, policy $z\to a$ — they mean entirely different things. My bet: bi-Lipschitz encoder + contractive dynamics, and the metric is probably not Euclidean but a learned Riemannian / contraction metric.
• Local vs global bounds. Global is vacuous; everything has to be local, and getting reliable local $L$ and $\ell$ efficiently is the core difficulty.
• The cost of invertibility. Normalizing flows hand you bi-Lipschitz and exact density, but at a price in capacity and training stability, to be balanced against the action objective.
• The benefit is context-dependent. Prediction / consistency objectives sometimes help little in-distribution and more on generalization and long horizons. “The signal is being wasted” is true; “using it is always a net win” needs experiments.

9. Closing

Action chunking wins as the default because it outsources the hardest job — dynamics — to the real world, and the real world speaks in actions. I’m not arguing to overthrow it. I’m arguing to scoop back the implicit future belief it discards: treat it as a distribution, ground it in reality, and harden it with Lipschitz / contraction geometry.

The reward of that route isn’t “save one inference.” It’s a more structured problem definition — prediction error upgraded to a bounded, falsifiable certificate, control recast as constraint satisfaction. The cost is concentrated in one place: the bi-Lipschitz / local-bound bone. Whether that bone can be chewed pretty much decides whether this is an elegant idea on paper or an actual method.

References

Pointers, not a precise bibliography — read these for the real thing.

• ACT / ALOHA — Zhao et al. Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware. arXiv:2304.13705
• Diffusion Policy — Chi et al. arXiv:2303.04137 · π0 arXiv:2410.24164 · π0.5 arXiv:2504.16054
• RTC (real-time / async chunking via inpainting) — Black et al. Real-Time Execution of Action Chunking Flow Policies. arXiv:2506.07339
• Self-predictive / latent consistency — SPR (Schwarzer et al., arXiv:2007.05929) · TD-MPC2 (arXiv:2310.16828) · DreamerV3 (arXiv:2301.04104)
• Predict & plan in latent — V-JEPA 2 / V-JEPA 2-AC (Assran et al., arXiv:2506.09985); JEPA (LeCun, 2022)
• Geometry — Kantorovich–Rubinstein duality (Villani, Optimal Transport) · contraction / incremental stability (Lohmiller & Slotine, Automatica 1998) · normalizing flows (Rezende & Mohamed, arXiv:1505.05770; Papamakarios et al. survey, arXiv:1912.02762)