You train a behavioral cloning policy. It works on the demos. You deploy it, and at some point a perturbation pushes it off track. The policy keeps producing actions that look reasonable. They’re just slightly off the manifold of successful executions, and there’s no restoring force.

The well-documented failure mode of pure BC (behavioral cloning) on manipulation is degenerate oscillation: the agent approaches the object, nudges it, backs off, re-approaches from the same angle, and loops. It never discovers it needs to go around.

The standard fixes have costs. Fine-tuning invalidates your safety case. Adding a planner requires a world model, and if your world model is wrong, the planner optimizes against a hallucination. Test-time adaptation reacts to the observation but ignores what the expert would’ve done.

My approach: don’t retrain the policy, don’t replace it, repair it in flight. Train a small recovery head, offline, that learns when the frozen policy is about to fail and how to push its action chunk back toward expert behavior. The training signal is counterfactual regret labels. The transport mechanism is flow matching in DCT-compressed action space.

Glossary (click to expand)
Term Meaning
BC Behavioral Cloning. Learning a policy by supervised regression on expert demonstration data.
MPPI Model Predictive Path Integral control. A sampling-based planner that rolls out candidate actions through a dynamics model and picks the lowest-cost one.
CFR Counterfactual Regret Minimization. A label that says "would the repair have done better than the baseline from this exact state?"
DCT Discrete Cosine Transform. A fixed mathematical basis that compresses an action chunk into a few low-frequency coefficients.
PACE Policy Anchor-Conditioned Execution. Executing the full recovered action chunk before re-planning, instead of replanning every step.
DoF Degrees of Freedom. The dimensionality of the action space (PushT has 2, LIBERO has 7).
SAE Sparse Autoencoder. A neural net that learns a sparse, interpretable representation of latent activations. Used here for drift detection.
Action chunk A sequence of H future actions predicted by the policy in one forward pass (H=8 in this work).
Anchor bank A fixed collection of expert action chunks retrieved by similarity to the current state.
Gate A learned binary classifier that decides whether to intervene on the frozen policy's output.
Flow matching Training a velocity field to transport one distribution toward another. Here: BC chunk toward expert anchor chunk.
Frozen BC Policy observation → chunk action Environment step(action_chunk) next obs Off-manifold drift no restoring force Repair Head gate + flow decoder retrieve Anchor Bank expert demos flow-matched repair
The control loop. The frozen BC policy (blue) runs normally. When it drifts off the success manifold (red), the learned repair head (yellow) retrieves an expert anchor and transports the BC chunk back via flow matching.

The repair manifold

Think of the policy as producing actions on a manifold in action space: the set of chunks that lead to successful task completion. Expert demos live on this manifold. A well-trained BC policy produces actions near it. Under perturbation, the outputs slide off and there’s no restoring force.

A repair layer takes an off-manifold chunk & transports it back, conditioned on what the expert would’ve done from a similar state. Two questions, separable: when to intervene (a gate), and where to go (a decoder). Get either wrong and you either fail to repair or break working plans.

action dim 1 action dim 2 expert anchors BC (in-distribution) drifted BC flow transport
The success manifold (blue band). Expert demos (blue dots) lie on it. BC stays near it in-distribution. Under perturbation, BC drifts off (red). Flow matching transports it back (yellow).

The gate: when to intervene

The gate is the whole game. A repair layer that fires too often is worse than one that fires too rarely. Every intervention is a chance to break a working plan.

My first gate label was “any candidate has positive regret AND baseline fails.” That conflates “there exists a better action” with “the baseline would fail if left alone.” The gate fires whenever a better option exists, even when BC+MPPI (model predictive path integral control) would’ve succeeded. You get regressions on working plans.

The fix: relabel with paired counterfactual rollouts from saved states. Run both the baseline and the repair from the same frozen state. The gate label is TRUE only if the baseline fails AND the repair succeeds.

Saved state s_t freeze env, branch rollout Baseline rollout BC + MPPI from s_t → Success? Repair rollout CFR action from s_t → Success? Intervention Label = (baseline fails) AND (repair succeeds)
Paired counterfactual relabeling. The gate label is TRUE only when the baseline fails and the repair succeeds. This is the only label that produces a conservative, precision-first gate.

The result

Three-seed validation on hard perturbations, 50 episodes per seed:

Seed BC+MPPI strict CFR strict Strict net Legacy net
42 0.08 0.08 0 (2 fix, 2 reg) 0
123 0.02 0.06 +2 (2 fix, 0 reg) -0.02
999 0.02 0.00 -1 (0 fix, 1 reg) -0.02
Total 0.04 0.05 +1 (4 fix, 3 reg) -0.04

Mean net per seed is +0.33, std 1.53. The first 30-episode run gave “+3 fixes, 0 regressions.” Scaling to 3 seeds × 50 episodes showed the variance was larger than the effect. The signal is real but small.

The decoder: two approaches

Approach A: CFR (counterfactual regret) selection. MPPI proposes 256 chunks, a learned dynamics model rolls each forward, keep top-16. A transformer selects one. The world model is in the inference loop.

Approach B: flow-matching transport. Skip MPPI. Take the BC chunk, transport it toward an expert anchor via a learned velocity field in DCT (discrete cosine transform) space. No world model at inference.

Approach A: CFR selection BC + MPPI 256 candidates World Model rollout → top-16 CFR Selector attend over 16 tokens world model in inference loop Approach B: flow matching BC chunk no MPPI Anchor Bank frozen expert demos Flow velocity net transport in DCT space no world model at inference
Approach A (left) needs the world model at inference. Approach B (right) only needs the frozen anchor bank. The world model is the ceiling on A; B sidesteps it.

The world model is the ceiling on A. I compared two dynamics models on multi-step rollout accuracy:

Model h=10 mean px h=50 mean px h=50 p90 px
3-layer MLP 12.8 46.2 80.4
6-layer residual v1 18.7 93.1 161.1
6-layer residual v3 18.5 69.1 107.0

The 3-layer MLP is the best I have, and it’s still 46px off at 50 steps. And the bigger residual is worse: it overfits short-horizon corrections and amplifies them over long rollouts. For world models, capacity is a liability. This pushed me toward Approach B.

Why DCT space

The flow head transports the BC chunk toward the anchor in DCT-edge coefficient space. The 8-step chunk keeps the first 2 timesteps raw (the edge: what you’re about to execute) and projects the rest onto 4 DCT-II basis coefficients. 8 raw values become 6 numbers.

t0 t1 2 Edge steps (raw) k0 k1 k2 k3 4 DCT Coefficients 8 raw values → 6 numbers ● Lower-dimensional ● Inherently smooth ● Exact next action
DCT-edge encoding. First 2 timesteps raw (yellow) for exact immediate action. Remaining 6 projected onto 4 DCT coefficients (blue), high-freq near zero (red). Low-dimensional, smooth by construction, exact edge.

Lower-dimensional: the velocity field has less to predict. Smooth by construction: you can’t produce a jerky chunk from 4 low-frequency coefficients. Exact edge: the immediate next action is kept raw, not smeared by DCT. When you have a strong prior, use it directly.

Both the gate and the flow head start conservative: negative bias init on the gate, zero-init on the velocity head. At initialization the repair does nothing. The model has to learn to deviate. A repair that never fires is safe; one that fires too often is dangerous.

LIBERO results

PushT is 2D. LIBERO is 7-DoF manipulation from images. If the method is real, it survives the transfer. The portable pieces: anchor bank, DCT-edge encoding, flow velocity net, intervention gate. The one adaptation: flow only transports the continuous 6-DoF chunk. The gripper stays fixed from BC.

100-episode evaluation on 3 LIBERO_SPATIAL tasks, PACE (exec_horizon=8, meaning the full recovered chunk is executed before replanning), old 6-epoch BC baseline:

Task BC dense Flow dense Flow minus BC
0 134.59 142.27 +7.68
1 110.58 113.89 +3.31
2 115.26 122.64 +7.38

Flow beats BC on all three tasks by +6.12 avg dense reward. The CaB gate alone adds almost nothing. Flow matching is the key driver.

dense reward BC Flow 134.6 142.3 Task 0 110.6 113.9 Task 1 115.3 122.6 Task 2
LIBERO dense reward, 100 episodes per task, PACE h=8. Flow (yellow) beats BC (gray) on all three.

Success rate is 0/100 for all methods. The BC baseline was only trained for 6 epochs (loss -4.49), so it never solves the task. Dense reward shows flow improves over BC, but neither crosses the success threshold yet. A 50-epoch BC retrain (loss ~-20.75) is in progress. These numbers show flow helps a weak policy move in the right direction; the next step is seeing if that converts to success rate with a trained baseline.

What didn’t work

Original gate label. “Positive regret AND baseline fails” produced over-intervention and regressions. The paired counterfactual label fixed it conceptually, but the effect size is still small.

More data with the wrong label. Scaled 50 to 200 episodes. Offline F1 dropped 0.90 to 0.79. More data compounds label noise.

Bigger dynamics models are worse. The 6-layer residual gets 93px error at h=50 vs 46px for the 3-layer MLP. Capacity is a liability for world models.

Always-on repair. A repair layer with no gate breaks working episodes. The conservative gate is load-bearing.

The frozen policy knows when it’s failing

A raw linear probe on the BC’s 512-dim latent space hits F1 0.98 for failure detection. A sparse-autoencoder (SAE) drift detector gets F1 0.36 to 0.49. The raw latent beats the SAE by a factor of two.

F1 score 0.98 Raw latent probe 0.42 SAE detector
The frozen policy's own representations already know when it's failing. You don't need a fancy drift detector; a linear probe suffices.

The frozen policy’s own representations already know when it’s failing. You just have to read it out. The SAE direction was a detour.

Status

Component Status Notes
PushT CFR Built, evaluated +1 net strict over 180 episodes
Dynamics bottleneck Diagnosed 3-layer MLP is the ceiling; deeper is worse
Flow repair (LIBERO) Trained, evaluated +6.12 dense reward over BC, 0/100 success (BC undertrained)
Drift detection Built, evaluated F1 0.98
BC retraining (50 ep) In progress Crashed before checkpoint save
Multi-seed Not done All results single-seed (42)

The LIBERO +6.12 dense reward is with a 6-epoch BC. The PushT +1 net strict is small. I’m sharing these because they’re directional and worth chasing. The next step is a trained BC baseline and multi-seed runs. If the signal survives those, there’s a paper. If not, the diagnosis (world model is the ceiling, gate label is the lever, DCT flow is the right representation) is still the contribution.

Sharing work in progress. The next experiment might shift the picture. That’s the point.