It's Not Just More Demos: Counterfactual Action Sensitivity Coverage for Data-Efficient Robust Robot Imitation

Technical Summary: Counterfactual Nuisance Action Sensitivity Coverage for Data-Efficient Robust Robot Imitation

Problem Statement
Visuomotor imitation learning has achieved significant success in manipulation tasks, yet trained policies often remain brittle when deployed under task-preserving visual nuisances (e.g., lighting changes, background shifts, distractor objects, or surface texture variations). While these nuisances do not alter the underlying task state or the required expert action, they cause substantial performance degradation in policies trained on "clean" data. The robotics community has traditionally addressed this by scaling data collection or using domain randomization. However, blindly increasing data diversity is inefficient; it often covers variations the policy already handles while missing specific visual conditions that induce brittle responses. The core challenge is identifying, for a specific trained policy, which additional data points are most informative for repairing its robustness without requiring costly online rollouts or expert intervention.

Methodology: Counterfactual Nuisance Behaviour Cloning (CFNBC)
The authors propose CFNBC, an offline data-selection framework designed to repair visuomotor policies by targeting their specific fragility modes. The pipeline operates in three stages:

1. Generation of Paired Counterfactuals: Starting with a nominal policy ($\pi_{\theta_0}$) trained on clean demonstrations, the method generates paired observations for the same task state: one clean ($s^c_t$) and one with a visual nuisance ($s^n_t$). Crucially, the expert action $a^*$ remains invariant between the pair.
2. Action Drift Measurement: The policy's predicted actions for the clean and nuisance inputs are compared in a normalized action space. The action drift ($d_i$) is defined as the magnitude of the change in the policy's predicted action between the paired inputs. High drift indicates that the policy is sensitive to task-irrelevant visual changes. Additionally, a response feature ($r_i$) captures the type of action change induced, allowing for the comparison of different failure modes.
3. Response-Guided Repair Selection: Instead of simply selecting the examples with the highest drift (which may be redundant), CFNBC selects a compact repair set ($R_K$) that covers diverse high-sensitivity response modes. Using a greedy optimization approach, the algorithm maximizes the coverage of the candidate response space under a fixed budget $K$. This ensures the selected data addresses a wide variety of fragility patterns rather than just the most extreme outliers.
4. Fine-Tuning: The selected nuisance observations are paired with their original expert actions (inherited from the clean demonstrations) to form a counterfactual repair dataset. The nominal policy is fine-tuned on a mixture of the original clean data and this selected repair set. This process requires no new expert labels, rollout success labels, or rewards.

Key Contributions

• Formulation of Robustness as Data Selection: The paper frames the repair of visuomotor policies not as a problem of generating more data, but as selecting the most informative counterfactual examples for a specific policy.
• Action Drift Signal: The introduction of "action drift" as an offline, policy-specific sensitivity signal. This metric ranks nuisance conditions by fragility without requiring environment rollouts or success labels.
• Response-Guided Selection: A novel selection strategy that prioritizes diversity in policy response modes over raw drift magnitude, ensuring the repair set covers a broad spectrum of failure behaviors.
• Empirical Validation: Demonstration that this targeted approach significantly outperforms random data selection and high-drift-only selection under low-data budgets.

Experimental Results
The framework was evaluated on two simulated manipulation tasks: MuJoCo bimanual cube transfer and SimplerEnv cube stacking, using Action Chunking with Transformers (ACT) as the base architecture.

• Brittleness Confirmation: Nominal policies achieved high success (~90-96%) in clean domains but suffered severe degradation (success rates dropping to 0-44%) under task-preserving visual nuisances.
• Drift as a Proxy: Action drift strongly correlated with rollout failure rates (Spearman correlations of $\rho=0.95$ and $\rho=0.88$), validating it as a reliable offline indicator of fragility.
• Data Efficiency: In the low-budget regime (20–30 selected candidates), response-guided repair substantially outperformed matched-budget random selection.
  • In Cube Transfer, the proposed method achieved a 0.96 success rate across all nuisances with 20 candidates, whereas random selection required ~100–500 samples to reach similar performance.
  • In Cube Stacking, the method improved all-nuisance success by +0.44 over the nominal policy with 30 candidates, compared to +0.28 for random selection.
• Generalization: The repaired policies maintained high performance on clean data, indicating that robustness was gained without sacrificing nominal behavior. While performance on held-out nuisance conditions varied, the method consistently improved transfer within the tested candidate pool.

Significance and Claims
The paper argues for a "data-centric" view of robustness repair. It claims that the most useful data for repairing a policy are not necessarily the most numerous, visually diverse, or obviously difficult examples. Instead, the most effective data are those that cover the fragile response modes of the current policy. By using counterfactual generation as an offline audit rather than unstructured augmentation, CFNBC enables the construction of compact, high-impact repair sets that significantly improve robustness with minimal data overhead, avoiding the costs and safety risks associated with online failure rollouts.