💻 computer science

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

The paper proposes MoRE (Mode Redirection), a method that distills a temporary mode classifier's guidance into policy weights via a short "uncloning" step to permanently suppress unsafe or undesired behavior modes in behavior-cloned policies without incurring inference-time overhead, thereby significantly improving deployment success rates across diverse robotic tasks.

Original authors: Hao Wang, Jiuzhou Lei, Dayou Li, Bangya Liu, Minghui Zheng, Manling Li, Ruohan Zhang, Zhiwen Fan

Published 2026-06-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Hao Wang, Jiuzhou Lei, Dayou Li, Bangya Liu, Minghui Zheng, Manling Li, Ruohan Zhang, Zhiwen Fan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you teach a robot to do a job, like handing you a knife, by showing it hundreds of videos of people doing that task. Because the videos come from different people in different kitchens, the robot learns multiple ways to do the job.

Sometimes, it learns the safe way (holding the knife by the handle). But because it learned from everything, it also learns the dangerous way (holding the knife by the blade). When you send this robot into a real home, it might get confused and choose the dangerous way, even though it knows how to do the job safely.

This paper introduces a solution called MoRE (Mode Redirection). Think of it as a "behavioral unlearning" tool that fixes the robot's bad habits without needing to re-teach it from scratch or slow it down while it works.

Here is how it works, using simple analogies:

The Problem: The "Jack-of-All-Trades" Robot

When robots are trained on big, messy datasets, they become like a student who studied for a test using every possible textbook. They know the answer, but they might choose the wrong method to get there.

The Issue: If you try to fix this by throwing away the "bad" videos and retraining the robot, it takes forever and costs a lot of money.
The Alternative: Some people try to add a "traffic cop" to stand next to the robot while it works, shouting "No! Do it this way!" This works, but it slows the robot down because the traffic cop has to think and shout every single second.

The Solution: MoRE (The "Internal Compass")

MoRE is different. Instead of adding a traffic cop or retraining the whole robot, it performs a quick "surgery" on the robot's brain (its software weights) to permanently change its habits.

Here is the step-by-step process:

The Temporary Coach: First, MoRE builds a tiny, temporary "classifier." Think of this as a coach who can look at the robot's current actions and say, "Oh, you are about to do the blade-first move. That's the bad mode."
The Gentle Nudge: The robot tries to do a task. When the coach sees it starting to drift toward a bad habit (like the blade-first move), the coach gives a gentle "nudge" (a mathematical signal) to the robot's brain. This signal says, "Hey, steer away from that path."
The Safety Net (The Retain Loss): The robot also needs to keep its skills. If we only tell it "don't do the bad thing," it might forget how to do the task at all. So, MoRE also tells the robot, "While you are fixing the bad habit, make sure you keep doing the good habit perfectly." This is like telling a student, "Stop cheating, but keep studying hard for the right answer."
The "Uncloning": Once the robot has learned to avoid the bad habits and stick to the good ones, the temporary coach is fired and thrown away. The robot is now a standalone expert that naturally avoids the bad moves.

Why This is a Big Deal

The paper claims MoRE is superior because:

No Speed Bumps: Unlike the "traffic cop" method, MoRE doesn't add any extra steps when the robot is actually working. It runs just as fast as before.
No Re-training: You don't need the original videos or weeks of training. You just need a few examples of the "good" vs. "bad" ways to do the task.
It Works Everywhere: The authors tested this on simulated robots (in video games) and real robots (physical arms moving knives and bottles). In almost every case, the robots became much safer and more successful at following the specific rules you wanted, without losing their ability to do the job.

The Bottom Line

MoRE takes a robot that is confused by having too many options and "distills" a clear preference directly into its brain. It's like taking a student who knows how to drive but keeps speeding, and giving them a quick lesson that rewires their instinct to always drive safely, without needing a driving instructor to sit in the passenger seat forever.

Key Results from the Paper:

On average, MoRE improved the robot's success rate by 44% compared to the unedited robot.
It performed almost as well as if they had retrained the robot from scratch with only "good" data, but much faster.
It works on different types of robot brains (from simple ones to advanced AI models) and for different tasks (moving objects, walking, handing things over).

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

Problem Statement

Behavior cloning (BC) enables the training of capable visuomotor policies from large, diverse demonstration datasets. However, these datasets often aggregate demonstrations from different operators and contexts, resulting in policies that learn multiple "behavior modes" for the same task. While all modes may successfully complete the task, some are unsafe, undesirable, or incompatible with specific deployment constraints (e.g., passing a knife blade-first versus handle-first).

Existing solutions face significant trade-offs:

Filtered Retraining: Removing undesired demonstrations and retraining from scratch provides direct control but requires access to original data and incurs high computational costs for every new target mode.
Inference-Time Steering: Methods that re-rank actions, guide diffusion sampling, or intervene on hidden states during deployment avoid retraining but introduce substantial inference-time overhead, requiring external modules or sampling loops that complicate low-latency control.

The paper identifies a gap for a method that can suppress undesired modes in an already-trained policy without requiring the original dataset and without adding inference-time overhead.

Methodology: MoRE (Mode Redirection)

The authors propose MoRE, a post-hoc policy-editing framework that distills a mode-redirection signal directly into the policy weights. The process involves three main stages:

1. Problem Formulation

Given a mixed-mode policy $\pi_{\theta_0}$ trained on a dataset $D_{BC}$ partitioned into modes $D_k$ , the goal is to edit the policy to a new parameter set $\theta^*$ such that closed-loop rollouts fall within a desired deployment set $S \subseteq \{1, \dots, K\}$ . The editing data $D_{edit}$ can consist of original demonstrations or closed-loop rollouts from the mixed policy, labeled with behavior modes.

2. Differentiable Mode-Redirection Editor

MoRE utilizes a temporary, lightweight behavior-mode classifier $g_\phi$ trained on features extracted from the original policy.

Feature Extraction: For Vision-Language-Action (VLA) policies, features are mean-pooled hidden states; for Diffusion Policies, features include the observation condition and the predicted clean action chunk.
Freezing: The classifier parameters $\phi$ are frozen after training. During the editing phase, the policy parameters $\theta$ are updated while the classifier remains fixed. This allows gradients from the classifier to flow back through the policy-dependent features into the policy weights.

3. Optimization with Retain and Redirect Losses

The editing objective combines two loss terms to balance behavior suppression with task competence:
$L_{MoRE}(\theta) = \underbrace{\mathbb{E}_{x \sim D_{des}} [L_{BC}(x; \theta)]}_{\text{Retain}} + \gamma \underbrace{\mathbb{E}_{x \sim D_{undes}^\theta} [L_{red}(x; \theta, \phi, S)]}_{\text{Redirect}}$

Retain Loss: Applied to samples in the desired mode set ( $D_{des}$ ), this term uses the original imitation loss (e.g., noise prediction or velocity matching) to preserve task competence and prevent the policy from forgetting how to complete the task.
Redirect Loss: Applied to undesired-mode samples ( $D_{undes}$ ), this term uses a subset-probability loss to push the classifier's predicted probability mass toward the desired set $S$ .
Source-Mode Masking: To ensure stability, the redirect loss is only applied to undesired samples where the source mode's probability is below a threshold $\tau$ (e.g., 0.5). This ensures redirection occurs only in state-action regions shared across modes, avoiding disruption to mode-specific execution branches where the policy has already committed to a trajectory.

After editing, the classifier is discarded. The resulting policy $\pi_{\theta^*}$ runs through the original inference path with zero additional overhead.

Key Contributions

Formulation of Behavior Uncloning: The paper defines a new policy-editing setting focused on suppressing deployment-undesired modes in pre-trained policies, evaluated via a "Deployment Success Rate" (SR) metric that requires both task completion and adherence to an acceptable behavior mode.
The MoRE Framework: A method that distills classifier-guided redirection into policy weights using a retain loss to maintain competence. It supports single- and multi-target edits, works with demonstrations or rollouts, and removes the classifier post-editing.
Comprehensive Validation: MoRE is validated across Diffusion Policy and $\pi_{0.5}$ VLA backbones, binary and multi-way mode control, simulated and real-robot deployments, and both manipulation and quadruped navigation tasks.

Experimental Results

Experiments were conducted on eight tasks (four simulated, four real-robot) involving tasks like knife handover, multi-location placement, and quadruped navigation.

Performance Gains: MoRE improved the average Deployment Success Rate (SR) by 44 percentage points over the original mixed-mode policy. For example, on the Push-T task, SR increased from 50.0% to 98.3%.
Comparison to Baselines:
- Vs. Preference Learning (CPL): MoRE achieved stronger mode shifts and higher SR. CPL provided only partial shifts and sometimes reduced Task Completion Rate (TCR).
- Vs. Inference-Time Steering (DynaGuide): While DynaGuide improved mode alignment on some tasks, it incurred an 8.47x inference-time slowdown. MoRE maintained the original inference speed.
- Vs. Filtered Retraining: MoRE approached the performance of policies retrained from scratch on filtered data without requiring access to the original dataset or full retraining.
Generalization: The method successfully transferred across different policy backbones (Diffusion Policy to $\pi_{0.5}$ VLA), different embodiment types (arm manipulation to quadruped), and from simulation to real-world deployment.
Data Efficiency: MoRE remained effective even when using only closed-loop rollouts from the mixed policy as editing data, though original demonstrations yielded slightly better results.

Significance and Claims

The paper positions MoRE as a practical solution for the "deployment tension" between dataset diversity and safety/acceptability constraints. By distilling redirection signals into the weights, MoRE enables the creation of standalone, safe policies that do not require external steering modules or retraining from scratch.

The authors explicitly state that MoRE should be viewed as a mode-control editing method rather than a certified safety layer. Its effectiveness relies on specific assumptions:

The availability of mode-labeled editing data (or reliably inferable rollouts).
A classifier interface where features expose mode information before the trajectory fully commits to a source mode.

The work suggests that targeted behaviors can be redirected directly in model weights, offering a path to safer, more controllable robot deployment without the computational or latency costs of inference-time intervention. Future work directions identified include handling noisier labels, discovering hierarchical mode taxonomies, and addressing stronger distribution shifts.