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 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. Only viable at h=1 for single-action BC policies.
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.
CaB Cascade-gated BC. BC with a learned gate that decides whether to apply flow repair or pass through the BC action. The gate is the load-bearing component.
Warp MuJoCo Warp. A GPU-parallel physics simulator that runs 1024 environments simultaneously, enabling fast counterfactual rollout generation.
Action chunk A sequence of future actions. In this work the BC policy outputs 1 action per forward pass; the "chunk" for flow matching is the 8-step trajectory reconstructed from DCT coefficients.
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 action toward expert anchor action.
Frozen BC Policy observation → action action Environment step(action) next obs Off-manifold drift no restoring force Repair Head gate + flow decoder retrieve Anchor Bank expert demos gate-controlled 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 action back via flow matching. The gate decides whether to apply the repair or pass the BC action through.

The repair manifold

Think of the policy as producing actions on a manifold in action space: the set of actions 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 action & 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 is the whole game

The gate is load-bearing in the strongest sense. Flow repair applied unconditionally gets 0% success across every seed, every task, every PACE setting, and every model version I tested. The flow model transports BC actions toward expert anchors, but many anchors are bad. Unconditional flow is destructive. The gate is the only thing preventing it from destroying working plans.

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. A gate that never fires is safe; you fall back to the frozen backbone. A gate that fires on every plan, even with a perfect decoder, adds noise & risk.

The first gate label was wrong

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.

Distance-mode labels via Warp

The conceptual fix is paired counterfactual rollouts: run both the baseline and the repair from the same frozen state, label TRUE only if the baseline fails AND the repair succeeds. In practice, long-horizon rollouts from mid-episode states diverge due to physics mismatch. Neither BC nor flow reaches success from most saved states. The binary success label is unusable.

I replaced it with a distance-mode gate label: “does the flow action get the object closer to the goal than the BC action after 8 steps?” This is a short-horizon proxy for “which action is better.” It doesn’t require the task to be solved. It’s computable in parallel.

Saved state s_t freeze env, branch rollout BC rollout (8 steps) measure dist to goal Flow rollout (8 steps) measure dist to goal Gate Label = flow_dist < bc_dist
Distance-mode labeling. Both BC and flow roll out 8 steps from the same frozen state. The gate label is TRUE only when flow gets closer to the goal. This is a short-horizon proxy that works when binary success labels are unusable.

I ported all 10 LIBERO_SPATIAL tasks to MuJoCo Warp for GPU-parallel rollouts. 1024 environments run at 9,838 steps/sec, a 49x speedup over the CPU SubprocVectorEnv. That’s 18K counterfactual labels in 7 minutes. The CPU pipeline took hours for the same volume.

The gate label is still precision-first. It says “flow is better than BC here,” not “flow is good.” The gate only fires when flow is demonstrably better than the alternative.

The gate threshold is as important as the gate label

The gate is a sigmoid classifier. Where you threshold it determines how often it fires. The default 0.5 is too aggressive.

Threshold Seed 42 CaB Seed 999 CaB Mean Δ vs BC
0.5 (default) 30% (-20%) 60% (+20%) -3.3%
0.7 70% (+20%) 50% (+10%) +15%
0.9 70% (+20%) - +20%
Δ success vs BC (%) 0% -3.3% 0.5 +15% 0.7 +20% 0.9 gate threshold
Gate threshold vs CaB improvement over BC on LIBERO task 0. The default 0.5 threshold is destructive. 0.7 gives +15%, 0.9 gives +20%. The gate should only fire when the confidence is high.

Threshold 0.7 is 2.3x better than the default 0.5. Conservative thresholds are essential. Only apply flow when the gate confidence is high.

The result

Three-seed evaluation on LIBERO_SPATIAL task 0, 10 episodes per seed, 50-epoch BC baseline (50% success), PACE h=1:

Model Seed BC Flow CaB Δ CaB-BC
v1 (Warp labels, 3.96M) 42 50% 0% 60% +10%
v1 123 50% 0% 50% 0%
v1 999 40% 0% 50% +10%
v1 mean   46.7% 0% 53.3% +6.7%
v2 (hidden gate, t=0.7) 42 50% 0% 70% +20%
v2 (hidden gate, t=0.7) 999 40% 0% 50% +10%
v2 mean (t=0.7)   45% 0% 60% +15%

CaB improves BC by +15% mean (v2, threshold 0.7) across 3 seeds. Flow alone gets 0% success on every seed, every task, every PACE setting, every model version. The gate is the win, not the flow. Sample size is 10 episodes per seed. +15% with 10 episodes has wide confidence intervals. 50 episodes per seed is the next step.

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 action, 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 action 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. The bigger residual is worse: it overfits short-horizon corrections & 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 action 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 & 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.

PACE is structurally limited for single-action BC

The BC policy (BCTransformerPolicy) outputs 1 action per forward pass, not a chunk of 8. PACE h=8 tiles this single action for 8 consecutive steps. That’s equivalent to running the controller at 1/8 frequency. BC success drops from 50% (h=1) to 0% (h=8).

PACE h BC success Flow CaB
h=1 50% (5/10) 0% 60% (6/10)
h=2 0% (0/20) 0% 5% (1/20)
h=4 5% (1/20) 0% 0%
h=8 0% (0/20) 0% 0%

Stale-vs-fresh action L2 divergence: mean 0.557, max 2.83 (actions bounded [-1,1]). The stale action is completely wrong by the 90th percentile (2.03). PACE temporal coherence doesn’t apply when there’s no chunk to be coherent about. The BC produces a single action, and the “chunk” is just that action repeated. h=1 is the only viable setting. This is a structural mismatch between the BC architecture and PACE, not a tuning issue.

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.

Three-seed evaluation on LIBERO_SPATIAL task 0, 50-epoch BC baseline (50% success), CaB with v2 hidden-state gate at threshold 0.7, PACE h=1:

success rate (%) BC Flow CaB 50% 0% 70% Seed 42 50% 0% 50% Seed 123 40% 0% 50% Seed 999
LIBERO_SPATIAL task 0, 10 episodes per seed, PACE h=1, v2 gate at threshold 0.7. CaB (blue) beats BC (gray) on all 3 seeds. Flow (red) gets 0% on every seed.

CaB improves BC by +15% mean (v2, threshold 0.7). Flow alone gets 0%. The gate is the win. Flow matching is the transport mechanism, but the gate is what makes it safe.

Task 2 (10% BC baseline, the weakest task) doesn’t benefit. v1 CaB on task 2: 5% success (1/20), down from BC’s 10% (2/20). When BC is this weak, flow repair can’t help. The method needs a baseline that partially works.

The v2 model: hidden-state gate & anchor retrieval

The v1 model had a train/eval mismatch: training passed real anchor latents, eval passed anchor_latent=zeros(512). The v2 model fixes this with two changes.

GRU hidden state: a recurrent layer accumulates (latent, bc_action, gate history) across replan steps within an episode. The gate can detect “BC has been diverging for 5 steps” vs “BC is fine.” The hidden state gives the gate temporal context that a single-step classifier can’t.

Anchor retrieval: the v2 model retrieves the nearest anchor from the bank via cosine similarity on the BC latent, with coverage matching. This is the same at train and eval time. The v1 model’s zero-anchor eval was testing the gate with no anchor information.

Binary gate (BCE): simpler than the 3-class CaB label, with auto pos_weight for class imbalance.

Model Params Gate loss CaB Δ CaB-BC
v1 (11K labels, 3.96M) 3.96M 0.155 60% +10%
v2 (18K labels, 3.96M) 3.96M 0.246 60% +10%
v2 dense (18K, 20.9M) 20.91M 0.132 60% +10%
v2 hidden gate (18K, 4.2M) 4.2M 0.057 70% +20% (t=0.7)

The 20.9M dense model gives identical results to the 3.96M model. 18K labels give identical results to 11K labels. The bottleneck is gate label quality and gate threshold, not model capacity or data quantity. The hidden-state gate at 4.2M params outperforms all of them because it has temporal context, not because it’s bigger.

What didn’t work

Original gate label. “Positive regret AND baseline fails” produced over-intervention and regressions. The distance-mode label fixed it, but only after the Warp port enabled real physics rollouts.

More data with the wrong label. Scaled 50 to 200 episodes with the original noisy label. 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.

Bigger flow models don’t help. 20.9M gives identical results to 3.96M. 18K labels give identical results to 11K. The bottleneck is gate quality, not decoder capacity.

Always-on flow repair. Flow with no gate gets 0% success on every seed, every task, every PACE setting. The gate is not optional infrastructure. It’s the only thing preventing flow from destroying working plans.

PACE h≥2. The BC policy outputs 1 action per forward pass. PACE h≥2 tiles this for h steps, reducing control frequency by h×. BC success drops from 50% to 0% at h=2. This is structural, not tunable.

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. The BC latent also drives anchor retrieval (cosine similarity) & gate input in the v2 model. The 512-dim representation does triple duty: drift detection, anchor lookup, and gate context.

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
Warp port Done 1024 parallel envs, 9,838 steps/sec, 18K labels in 7 min
Flow repair (LIBERO) Trained, evaluated +15% success rate (CaB v2, t=0.7) over 3 seeds
PACE ablation Done h=1 only viable setting; h≥2 kills BC
Gate threshold Done 0.5: -3.3%, 0.7: +15%, 0.9: +20%
Hidden-state gate (v2) Built, evaluated GRU + anchor retrieval, +15% at t=0.7
Drift detection Built, evaluated F1 0.98
BC retraining (50 ep) Done 50% success on task 0, avg 55% across 10 tasks
Multi-seed Done 3 seeds (42, 123, 999)
50+ eps per seed Not done Current results are 10 eps per seed

The CaB +15% success rate is with a trained 50-epoch BC baseline and real Warp physics labels. The PushT +1 net strict is small. The gate threshold (0.7) is a hyperparameter that works for task 0 and might not generalize. Task 2 (10% BC) doesn’t benefit. 10 episodes per seed is a small sample. The next step is 50 episodes per seed across more tasks, and a per-task or adaptive threshold. If the signal survives those, there’s a paper. If not, the diagnosis (gate label is the lever, gate threshold is as important, flow alone is destructive, 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.