From Action Labels to Sets: Rethinking Action Supervision for Imitation Learning from Corrective Feedback

1. Problem Statement

The paper addresses a fundamental limitation in Interactive Imitation Learning (IIL) and Behavior Cloning (BC): the brittleness of pointwise supervision.

• The Issue: Traditional BC treats human demonstrations or corrective feedback as exact, pointwise action targets ($a^*$). While effective with perfect data, this approach fails when human feedback is noisy, suboptimal, or relative (indicating direction but not magnitude).
• The Consequence: When using expressive policy classes (e.g., Energy-Based Models (EBMs) or Diffusion Models), pointwise supervision causes the policy to overfit to specific, potentially noisy labels. This creates sharp local minima around the provided actions, steering the policy away from the true underlying optimal behavior.
• Limitations of Alternatives:
  • Pairwise Preference Learning: While more robust than pointwise labels, pairwise comparisons (e.g., "Action A is better than Action B") only constrain the relative ordering of two actions. In high-dimensional continuous spaces, this leaves the rest of the action space largely unconstrained, allowing the policy to assign high probability to suboptimal regions while still satisfying the preference constraint.
• The Gap: There is a need for a supervision objective that is more robust than pointwise BC (to handle noise/suboptimality) but more informative than pairwise comparisons (to constrain the action space effectively).

2. Methodology: CLIC

The authors propose CLIC (Contrastive policy Learning from Interactive Corrections), a framework that shifts from pointwise labels to set-valued supervision.

Core Concept: Desired Action Sets

Instead of treating a human correction as a single target action, CLIC interprets it as defining a Desired Action Set ($\hat{A}$). This set contains all actions consistent with the feedback while excluding the executed (suboptimal) action.

• Absolute Corrections: If a human provides a new action $a_h$, the desired set is a region around $a_h$ (e.g., a ball or polytope).
• Relative Corrections: If a human provides a direction $h$ to improve the current action $a_r$, the desired set is a cone or half-space in that direction, without committing to a specific magnitude.

Key Components

A. Construction of Desired Action Sets
The paper defines two specific geometric constructions for these sets:

1. Polytope Sets (CLIC-Half): Constructed by intersecting multiple "half-space" constraints derived from data augmentation.
   • A single correction $(a_r, a_h)$ is augmented into multiple contrastive pairs $(a^-_i, a^+_i)$.
   • Each pair defines a half-space where actions closer to $a^+_i$ are preferred.
   • The intersection of these half-spaces forms a polytope (approaching a cone).
   • Parameters: Directional certainty ($\alpha$) and magnitude certainty ($\epsilon$) control the shape and size of the set.
2. Circular Sets (CLIC-Circular): Specifically for absolute corrections.
   • Defines a ball centered at the human action $a_h$ with a radius proportional to the distance between the robot's action and the human's action.
   • This captures the intuition that the optimal action lies within a neighborhood of the correction.

B. Aggregation and Convergence

• Iterative Refinement: A single correction is often insufficient to pinpoint the optimal action. CLIC aggregates multiple corrections over time.
• Intersection: The overall desired action set for a state is the intersection of all individual sets derived from the feedback buffer ($\hat{A}_{D_s} = \bigcap \hat{A}_i$).
• Convergence: Theoretically, as more feedback is aggregated, the intersection of these sets converges to the set of optimal actions ($A^*_s$), provided the teacher provides feedback on suboptimal actions and the sets always include at least one optimal action.

C. Policy Shaping Objective
To train the policy $\pi_\theta$ to align with these sets, CLIC minimizes the KL divergence between the current policy and a target distribution ($\pi_{target}$):
$$\ell_{KL} = \mathbb{E} [ KL(\pi_{target}(\cdot|s) \parallel \pi_\theta(\cdot|s)) ]$$

• Target Distribution: Defined via Bayes' rule using an observation model.
  • Observation Model: A soft membership function (sigmoid) that assigns high probability to actions inside the desired set and low probability to those outside.
  • Prior: Uses the current policy as the prior ($\pi_\theta$) rather than a uniform prior. This creates a "policy-weighted Bayes loss," ensuring updates are conservative and incremental, preventing the policy from drifting too far from previously learned behaviors.
• Implicit vs. Explicit:
  • Implicit (EBM): The loss is optimized using contrastive sampling (InfoNCE-style), where the target distribution weights samples based on their set membership.
  • Explicit (Gaussian): For unimodal tasks, the method simplifies to a hinge loss ensuring the policy mean lies within the polytope.

3. Key Contributions

1. Set-Valued Supervision Perspective: Proposes a novel formulation where corrective feedback defines a region of acceptable actions rather than a single point, bridging the gap between pointwise BC and pairwise preference learning.
2. CLIC Framework: Introduces a practical algorithm that constructs and aggregates these sets to iteratively refine the policy. It supports both absolute and relative feedback and handles multi-modal behaviors naturally via implicit policies (EBMs).
3. Theoretical Analysis: Provides convergence conditions showing that aggregated set-valued feedback can recover the optimal action set, even when individual feedback signals are noisy or relative.
4. Comprehensive Evaluation: Extensive validation in simulation and on real robots (KUKA iiwa) demonstrating robustness across diverse feedback types.

4. Experimental Results

Simulation Experiments:

• Accurate Feedback: CLIC performs competitively with state-of-the-art (SOTA) methods (e.g., Diffusion Policy, IBC), proving that set-valued supervision does not sacrifice performance when data is perfect.
• Noisy Feedback: CLIC significantly outperforms baselines (IBC, Diffusion Policy, PVP). While pointwise BC methods overfit to noise and fail to converge, CLIC's set-based approach absorbs the noise, maintaining high success rates.
• Relative & Partial Feedback:
  • Relative Corrections: CLIC-Half and CLIC-Explicit successfully learn from directional feedback where pointwise BC methods fail completely (due to the "on-policy" problem where past corrections become suboptimal).
  • Partial Feedback: In multi-robot tasks where humans only correct one arm, CLIC-Half excels by constraining only the corrected dimensions, whereas baselines trying to imitate the full (partially suboptimal) action vector fail.

Real-Robot Experiments:

• Insert-T (Multi-modal): CLIC-Half achieved 80% success on hard tasks, significantly outperforming Diffusion Policy (30%) and IBC (10%). This highlights CLIC's ability to handle complex, multi-modal decision-making without overfitting.
• Ball Catching (Dynamic/Partial): Successfully learned a high-frequency task using relative and partial feedback, achieving a 100% success rate within two attempts.
• Water Pouring (Precision): Demonstrated the ability to switch between absolute and relative corrections during training, refining the policy from coarse to fine control.

Ablation Studies:

• Confirmed that the policy-weighted prior is crucial for stability.
• Showed that the hyperparameters controlling the set geometry ($\alpha, \epsilon$) must be tuned to match the uncertainty of the feedback signal.
• Visualized that as the set radius shrinks to zero, CLIC degenerates into overfitting IBC, validating the necessity of the "set" formulation.

5. Significance and Impact

• Robustness to Human Error: The primary significance is making imitation learning robust to the inherent imperfections of human teachers (fatigue, teleoperation limits, cognitive load).
• Unified Framework: CLIC unifies learning from absolute demonstrations, relative corrections, and partial feedback under a single mathematical framework.
• Scalability to Complex Tasks: By avoiding the overfitting of expressive models to noisy labels, CLIC enables the use of powerful implicit policies (EBMs) in real-world, interactive settings where data is rarely perfect.
• Practical Applicability: The method reduces the cognitive load on human teachers (who can provide rough directions rather than exact trajectories) and accelerates the learning process for complex robotic tasks.

In summary, CLIC rethinks the fundamental unit of supervision in imitation learning from a single point to a set, providing a mathematically sound and empirically robust solution for learning from imperfect human feedback.