Behavior Uncloning: Distilling Mode Redirection into Policy Weights without Inference-Time Steering

Behavior-cloned policies often learn multiple behavior modes from demonstration datasets, including modes that are unsafe or otherwise undesired at deployment. This paper introduces MoRE (Mode Redirection), a post-hoc editing framework that distills a differentiable mode-redirection signal into policy weights to steer behavior toward desired modes without adding inference-time overhead.

The gist

MoRE distills the redirection signal from a temporary mode classifier into the policy weights to steer behavior, allowing the standalone policy to suppress unwanted modes with zero inference-time overhead.

Problem Formulation and Goal

Behavior cloning trains a policy by supervised imitation of demonstrated observation-action pairs, resulting in mixed-mode policies that can learn multiple successful strategies. The challenge is deploying these policies safely, as some learned modes might be unsafe at deployment. MoRE formulates behavior uncloning as an efficient policy-editing setting for suppressing deployment-undesired modes in already-trained robot policies, where a rollout is successful only when it both completes the task and follows an acceptable behavior mode. The goal is to find a new policy weight set, denoted as θ⋆, such that πθ⋆ produces closed-loop rollouts whose behavior modes fall within a desired deployment set S.

Differentiable Mode-Redirection Editor

MoRE introduces a differentiable mode-redirection editor. First, it trains a K-way behavior-mode classifier (gϕ) on features produced by the original mixed policy to predict the mode label m(x). During editing, this classifier is frozen, and its output provides a redirection signal. The key mechanism involves recomputing the policy-dependent features under the current policy as rθ(x), which allows gradients from the frozen classifier to flow through rθ(x) into the trainable policy weights θ. This process utilizes a unified subset-probability redirection loss, Lred, which redirects probability mass toward the desired deployment set S: Lred(x; θ, ϕ, S) = − log P[i∈S] exp z i.

Optimization for Balancing Task Competence and Behavior Uncloning

The optimization objective balances suppressing undesired modes with preserving task competence through a retain loss. The total MoRE loss is formulated as: LMoRE(θ) = Ex∼Ddes [LBC(x; θ)] z retain + γ Ex∼DMθ undes [Lred(x; θ, ϕ, S)] z redirect. The retain term uses the policy’s original imitation loss to keep desired-mode behavior close to the original policy. The redirect loss is applied only to undesired-mode samples whose source-mode probability is below a threshold τ: Mθ = [xt ∈ Dundes: qϕ(m(xt) rθ(xt)) < τ]. This masking ensures redirection occurs only in regions shared across behavior modes, focusing the edit where the policy can still change the subsequent behavior mode while avoiding disruption to later mode-specific execution.

Validation and Results

MoRE was validated across eight simulated and real-world tasks, including Diffusion Policy and π0.5 VLA backbones, binary and multi-way mode control settings, simulated and real-robot deployments. The primary finding is that MoRE improves the average deployment success rate (SR) by 44 percentage points over the original mixed-mode policy across eight simulated and real-world tasks. MoRE achieves the strongest SR among compared adaptation and steering baselines while approaching filtered-data retraining references with no inference-time overhead, demonstrating its capability to transfer across different policy backbones, mode counts, embodiments, and deployment settings. The method supports single- and multi-target edits by treating any completed rollout whose mode lies in the target set S as mode-aligned when evaluating SR.

Limitations

MoRE is not a certified safety layer; its effectiveness depends on several assumptions, including the availability of mode-labeled editing data or closed-loop rollouts whose modes can be reliably inferred, and a classifier interface that exposes mode information before trajectories have fully committed to a source mode. Future work should focus on extending behavior uncloning to noisier labels, automatically discovered or hierarchical mode taxonomies, longer-horizon tasks, and stronger distribution shifts.

How it works

1. The process begins with a mixed-mode policy trained via behavior cloning (LBC). The dataset is partitioned into demonstration sets Dk based on mode labels k=1 to K.

2. MoRE trains a K-way behavior-mode classifier (gϕ) on cached features from the original policy πθ0 to predict the mode label m(x). This classifier is then frozen.

3. During editing, MoRE recomputes policy-dependent features rθ(x) under the current policy πθ to generate a differentiable redirection signal that flows into the trainable policy weights θ.

As a fastidious and diligent researcher, I have analyzed the MoRE (Mode Redirection) framework presented in this paper. To improve existing AI systems—specifically reinforcement learning policies trained via behavior cloning or imitation learning—the following improvements can be implemented:

Here are the specific improvements and what the resulting AI system can do:

1. A significantly improved robot policy that maintains high task competence while suppressing unsafe or undesired execution modes during real-world deployment, without requiring any additional inference-time computational overhead.

2. The ability to selectively enforce desired behaviors (e.g., pass the knife handle-first instead of accidentally learning pass the knife blade-first) by distilling a lightweight mode classifier directly into the policy weights during a post-hoc editing step.

3. A system that achieves performance gains comparable to full dataset retraining (filtered retraining reference) but without requiring access to the original, potentially massive, demonstration datasets for every target mode.

4. The capability to perform multi-target behavior redirection, allowing a policy to suppress several undesired modes simultaneously while explicitly preserving a set of desired behaviors (e.g., maintaining competence in both left-hand grip and right-hand grip modes).

5. Improved generalization across diverse robotic backbones (Diffusion Policy, π0.5 VLA) and varied task categories (manipulation, quadruped navigation), ensuring the mode redirection mechanism remains effective regardless of the underlying model architecture or physical embodiment.

6. Deployment of complex policies in real-world settings with zero added latency compared to the original policy's inference path, as the correction is baked into the weights rather than requiring external steering modules or verifiers that increase control loop costs by up to 8.47x (as seen in baselines).

7. A robust mechanism for editing policies using only closed-loop rollouts from a mixed-mode policy, proving that original demonstrations are not strictly necessary for mode correction, thereby enabling adaptation from existing learned behavior without full retraining.

Behavior-cloned policies often learn multiple behavior modes from demonstration datasets, including modes that are unsafe or otherwise undesired at deployment. For example, a policy trained on diverse handover demonstrations may learn to pass a knife blade-first. Standard remedies such as data curation and inference-time steering either require access to the original demonstrations for full retraining or add substantial inference-time overhead. To address this gap, we propose MoRE(Mode Redirection), which redirects policy rollouts toward desired behavior modes through a short "uncloning" step. Specifically, MoRE distills the redirection signal from a temporary mode classifier into the policy weights to steer behavior. A retain loss balances this edit by preserving desired-mode competence, allowing the standalone policy to suppress unwanted modes with zero inference-time overhead. Across eight simulated and real-world tasks, MoRE improves the average deployment success rate (SR) by 44 percentage points over the original mixed-mode policy. Among all compared adaptation and steering baselines, MoRE achieves the strongest SR and approaches the filtered-data retraining reference, while preserving task competence and inference speed. MoRE also generalizes across robot policy backbones, including Diffusion Policy and the Pi0.5 VLA, diverse task categories, and real-world deployments.