Technical Summary: Difference-Aware Retrieval Policies for Imitation Learning (DARP)

Problem Statement

Parametric imitation learning via Behavior Cloning (BC) often suffers from poor generalization to out-of-distribution (OOD) states. While BC minimizes supervised risk on expert states, it fails to control variance in low-density regions of the state space encountered during closed-loop rollouts. Small errors accumulate (covariate shift), driving the agent into states where the learned policy produces high-variance, unstable behavior, frequently leading to task failure. Existing solutions to mitigate compounding errors often require additional assumptions, such as simulators, interactive expert feedback, or large quantities of sub-optimal data, moving beyond the standard BC regime of learning solely from expert state-action pairs.

Methodology

The paper introduces Difference-Aware Retrieval Policies for Imitation Learning (DARP), a semi-parametric retrieval-based approach that addresses these limitations without requiring additional data collection or supervision.

Theoretical Foundation

The authors first establish that enforcing local consistency in policy learning corresponds to manifold regularization (Laplacian smoothing). They define a Manifold Regularized Imitation Learning (MRIL) objective that adds a smoothness penalty to the standard BC loss, penalizing high-frequency variations between neighboring expert states. Theoretical analysis (Theorem 1) demonstrates that MRIL reduces estimator variance, guarantees local Lipschitz continuity, and improves policy stability compared to vanilla BC.

Implicit Regularization via Architecture

Rather than explicitly optimizing the regularized objective (which requires tuning a hyperparameter $\lambda$), the paper proposes Implicit Manifold Regularized Imitation Learning (iMRIL). This approach achieves the benefits of Laplacian smoothing by modifying the policy architecture rather than the loss function.

• Mechanism: Instead of learning a direct mapping $f(s)$, the policy predicts actions based on a neighborhood of retrieved $k$-nearest neighbors (k-NN) from the training dataset.
• Equivalence: The authors prove (Theorem 2) that iMRIL is equivalent to a parameter-free Laplacian regularization with an effective smoothing strength of $\lambda \approx 1$. The neighbor aggregation step acts as a fixed low-pass spectral filter on the k-NN graph, suppressing high-frequency variance while preserving low-frequency modes.

DARP Algorithm

DARP instantiates iMRIL with a specific parameterization designed to capture local geometry:

1. Retrieval: For a query state $s_q$, the algorithm retrieves $k$-nearest neighbors $(s^*_i, a^*_i)$ from the expert dataset $D^*$ based on a distance metric.
2. Difference-Aware Parameterization: The core innovation is the input representation for the neural network $f_\theta$. Instead of simply processing the neighbor state, $f_\theta$ takes a tuple consisting of:
   • The neighbor state $s^*_i$
   • The neighbor's expert action $a^*_i$
   • The difference vector $\Delta s_i = s^*_i - s_q$
     This allows the model to learn how expert actions vary under local perturbations relative to the query state.
3. Aggregation: The network produces candidate actions $a'_i = f_\theta(s^*_i, a^*_i, \Delta s_i)$. These candidates are aggregated via a permutation-invariant function $g_\psi$ (e.g., simple averaging, Set Transformer, or a generative model head) to produce the final action prediction $\hat{a}_q$.
4. Training: The model is trained using the standard imitation learning objective (minimizing the discrepancy between the aggregated prediction and the expert action), requiring no additional loss terms or hyperparameters beyond standard BC.

Key Contributions

1. Novel Algorithm: The introduction of DARP, which reparameterizes imitation learning in terms of relative differences between query states and retrieved neighbors, rather than direct state-to-action mappings.
2. Theoretical Insight: A formal proof that the proposed architecture implicitly performs Laplacian smoothing, providing variance reduction and stability guarantees without explicit regularization terms.
3. Semi-Parametric Design: A method that combines the robustness of non-parametric local methods (retrieval) with the stability of parametric global learning, operating strictly within the standard BC regime (no extra data or feedback).
4. Architectural Analysis: Identification of critical components, specifically the necessity of difference vectors and permutation-invariant aggregation, through comprehensive ablation studies.

Experimental Results

The authors evaluate DARP across continuous control (MuJoCo), robotic manipulation (Robosuite, RoboCasa), and high-dimensional visual tasks.

• Performance Gains: DARP consistently outperforms standard Behavior Cloning and other baselines (including Locally Weighted Regression and explicit regularization methods). Improvements range from 15% to 46% across diverse domains.
  • In MuJoCo locomotion tasks, DARP achieved scores significantly higher than BC (e.g., ~3545 vs. ~2313 on Hopper).
  • In Robosuite and RoboCasa manipulation tasks, DARP doubled the success rate of BC in real-world furniture assembly scenarios (92% vs. 44%).
• Generalization: DARP maintains performance improvements when using high-dimensional visual features (R3M embeddings) and multimodal action distributions (e.g., Gaussian Mixture Models, Diffusion Models).
• Ablation Findings:
  • Difference Vectors: Removing the difference vector ($\Delta s_i$) significantly degrades performance, confirming that learning relative changes is crucial.
  • Permutation Invariance: Using a permutation-dependent aggregator drastically reduces success rates.
  • Neighbor Selection: Random neighbor selection performs poorly, validating the importance of meaningful distance metrics.
• Divergence Analysis: In rollouts where BC fails and DARP succeeds, the point of divergence occurs when the query state is OOD (low state likelihood) but the difference vector relative to neighbors remains in-distribution (high delta likelihood). This supports the hypothesis that DARP's reparameterization improves prediction stability in OOD regions.

Significance and Claims

The paper claims that DARP offers a general, scalable architecture for imitation learning that naturally extends to complex representations and action distributions. By reparameterizing the problem to focus on local neighborhood structures and relative differences, DARP achieves the robustness of local methods and the stability of regularized global policies without the computational or data overhead of existing advanced methods. The authors assert that this approach provides a parameter-free form of smoothing that adapts to the local density and geometry of the dataset, significantly reducing the variance of BC policies and enhancing robustness to distribution shift while adhering strictly to the standard behavior cloning assumptions.