Difference-Aware Retrieval Policies for Imitation Learning
Authors: Quinn Pfeifer
Ethan Pronovost
Paarth Shah
Khimya Khetarpal
Siddhartha Srinivasa
Abhishek Gupta
Venue: ICLR 2026 Poster (2026)
Keywords: behavior cloning
robotics
imitation learning
nearest neighbor
retrieval
RAG
retrieval-augmented generation
Primary area: applications to robotics, autonomy, planning
TL;DR: DARP addresses behavior cloning's instability by conditioning action predictions on difference vectors between query states and neighbors

Abstract: Parametric imitation learning via behavior cloning can suffer from poor generalization to out-of-distribution states due to compounding errors during deployment. We show that reusing the training data during inference via a semi-parametric retrieval-based imitation learning approach can alleviate this challenge. We present Difference-Aware Retrieval Policies for Imitation Learning (DARP), a semi-parametric retrieval-based imitation learning approach that addresses this limitation by reparameterizing the imitation learning problem in terms of local neighborhood structure rather than direct state-to-action mappings. Instead of learning a global policy, DARP trains a model to predict actions based on k-nearest neighbors from expert demonstrations, their corresponding actions, and the relative distance vectors between neighbor states and query states. DARP requires no additional assumptions beyond those made for standard behavior cloning -- it does not require additional data collection, online expert feedback, or task-specific knowledge. We demonstrate consistent performance improvements of 15-46% over standard behavior cloning across diverse domains, including continuous control and robotic manipulation, and across different representations, including high-dimensional visual features.

Published as a conference paper at ICLR 2026
DIFFERENCE-AWARE RETRIEVAL POLICIES FOR IMITA- TION LEARNING
Quinn Pfeifer', Ethan Pronovost', Paarth Shah", Khimya Khetarpal 34 Siddhartha Srinivasa Abhishek Gupta
Paul G. Allen School of Computer Science & Engineering, University of Washington
Toyota Research Institute
Google DeepMind
4 Mila
ABSTRACT
Parametric imitation learning via behavior cloning can suffer from poor general- ization to out-of-distribution states due to compounding errors during deployment. We show that reusing the training data during inference via a semi-parametric retrieval-basee imitation learning approach can alleviate this challenge. We present Difference- Aare Retrieval Policies for Imitation Learning (DARP), a semi- parametric retrieval-based imitation learning approach that addresses this limitation by reparameterizing the imitation learning problem in terms of local neighbor- hood structure rather than direct state-to-action mappings. Instead of learning a global policy, DARP trains model to predict actions based on k-nearest neighbors from expert demonstrations, their corresponding actions, and the relative distance vectors between neighbor states and query states. DARP requires no additional assumptions beyond those made for standard behavior cloning it does not require additional data collection, online expert feedback, or task-specific knowledge. We demonstrate consistent performance improvements of 15-46% over standard be- havior cloning across diverse domains, including continuous control and robotic manipulation, and across different representations, including high-dimensional visual features.
INTRODUCTION
Imitation Learning
Difference-Aware Retrieval Policies
ff
S
fo
fo
ag
ag
Figure 1: Overview of DARP: Unlike standard BC (left), DARP (right) utilizes retrieval-hased reparameteri- zation centered around difference vectors between query states and retrieved neighbors. In standard behavior cloning, the dataset of expert state-action pairs is used only for training and is discarded inference- timee, while DARP utilizes to performs retrieval to find neighborhood of expert state-action pairs around each query point Sq.
Imitation learning via behavior cloning (BC) (Pomerleau, 1991) has enabled robots to learn complex, dexterous behaviors from expert demonstrations (Zhao et al. 2023 Chi al.) 2024 Black et al.) 2024) Chung et al.) 2014). Yet despite its simplicity, BC often proves brittle in practice, especially for long-horizon tasks (Ross et al. 2011). The core issue is covariate shift: small emors accumulate

---

Published as a conference paper at ICLR 2026
during rollouts, driving the agent into states not well represented in the demonstration data (Spencer jet at. [2021; Ross et al., (011)). In such -iit----istiitttion regions, BC policies are highly unstable, producing unrelable and high-variance behavior that frequently leads to failure.
This problem is well recognized, and many approaches have been proposed to mitigate compounding error (Ross et al.] 2011 Venkatraman et al., 2013 Ke al. 2024b Levine et al. 2020). However, these typically go beyond the standard BC assumptions, requiring simulators, interactive experts, large quantities of sub-optimal data, or strong task-specific structure. By contrast, our goal is to remain in the pure BC regime: learn only from expert state-action pairs, with no additional supervision or feedback. The central question is thus: can we reduce the variance of BC policies using only the original demonstration dataset?
From a statistical standpoint, BC minimizes only the supervised risk on expert states. This controls bias on the training distribution, but leaves variance unchecked: in low-density regions of the state space (which are often encountered during closed-loop rollouts), the learned policy can oscillate arbitrarily. A natural remedy is to enforce smoothhess, so that nearby states yield similar predicted actions (Kobayashi] 2022; Asadi et al. 2018 Ke et al., 2024a; Chen et al.. 2024). This discourages spurious fluctuations and improves rollout stability. Several approaches to encourage smoothness
have been explored, see related work in Section 4
Although sometimes effective, each has drawbacks: augmentation does not guarantee consistency, global priors can blur distinct behaviors, temporal penalties only act along time (not space), and explicit graph regularizers require tuning extra smoothness hyperparameters.
A complementary line of work contrasts global and local learning. "Global" supervised models (Black et al. 2024 Zhao et al.) [023] Chi et al.] (022) attempt compress the entire demonstration dataset into a single parametric function, which is typically brittle under distribution shift. "Local" methods (Pari et al. 2022 Mansimov & Chol 2018, (aalzergg& & Aha 1994) instead adapt predictions
to the structure of the dataset itself, consulting neighborhoods of similar states and generating outputs from non-parametric or semi-parametric operations on the training distribution of expert behavior. This locality offers robustness since it avoids reliance on a single parametric function, but it also has limitations: its effectiveness inherently depends on the distance metric, naive averaging of neighborhood can blur distinct actions and struggle to represent multimodality, and treating neighbors only in terms of their absolute states can limit generalization.
We introduce Difference-A wwaeee ware Retrieval Policies for Imitation Learning (DARP), which combines the robustness of local methods with the stability of regularized global policy learning. At inference time, rather than predicting actions to execute only from the current query state via feedforward pass on a parametric function, DARP (Fig. [ first retrieves a set of k neighbors from the training corpus and then predicts k actions conditioned tuples of (neighbor state, associated action, and difference from the query state). These
predictions are then aggregated in a permutation- invariant manner to produce single robust action prediction. This design both grounds predictions in observed data (due to non-parametric retrieval) and implicitly enforces local consistency (due to parametric action prediction, conditional the retrieved neighbors). We show that doing reduces variance without requiring any additional assumptions beyond those made for standard Behavior Cloning - we need no additional data, online supervision, or task task-specific knowledge. In spectral terms, this form of neighbor aggregation approximates a Laplacian filter on the k-NN (k-nearest neighbor) graph of expert states, providing parameter- form of smoothing that adapts to the local density and geometry of the dataset.
We provide both theoretical and empirical evidence that while operating under the same requirements as behavior cloning, DARP improves performance considerably by reducing variance and enhancing robustness to distribution shift. Our analysis formalizes the connection to Laplacian regularization, showing that DARP implicitly applies fixed low-pass spectral filter that suppresses fighffrry variance. Empirically, on imitation learning evaluations, DARP achieves 15-46% gains over typical behavior cloning across continuous control (MuJoCo), robotic manipulation (Robosuite, Robocasa), and high-dimensional visual imitation tasks (Robosuite with image state). We demonstrate that DARP is general, scalable architecture that naturally extends image-based domains, with rich policy classes like transformers and Gaussian mixture models. We perform a careful set of ablations to highlight the importance of our particular choice of representation and architecture, providing general- purpose insights into retrieval-basee algorithms for sequential decision maninng problems.
2

---

Published as a conference paper at ICLR 2026
2 SIFFERENCE-AWARE RETRIEVAL POLICIES FOR IMITATION LEARNING
In this work, we instantiate a new class of imitation learning methods that get the best of both "global" parametric learning methods and "local" learning methods. We propose a new architecture and simple training objective that allows for learning under the same requirements as typical behavior cloning. while providing significant improvements both theoretically and empirically. we thoroughly derive theoretical guarantees from first-principles. Readers primarily interested in the practical experimental results may skip to Section 2.3 for a self-contained summary of DARP before proceeding to Section B for empirical results.
this section,
As a warmup, we define the problem setting (Section 2.1 and discuss variant of regularized imitation learning SSection2.22 that imposes additional structure from the data for improvements in variance, generalization, and stability. In Section 2.3 then show how the benefits of explicitly regularized learning can be implicitly accomplished by modifying policy architecture rather than the objective. Finally, in Section 2.4 we introduce our practical algorithm DARP, which realizes these benefits through a semi-parametric retrieval augmented architecture that can be generally applied to imitation learning with modern neural networks and generative modeling tools.
2.1 PRELIMINARIES: BEHAVIOR CLONING FOR IMITATION LEARNING
We operate in the typical imitation learning setting, formalized by finite-horizon Markov Decision Process (MDP), 1 = { , .Pi). where y is the state space, ad is the action space, and Pb is the initial state distribution. A policy maps state distribution of actions Tg : As so as to maximize task-relevant objectives. We assume access to expert human-provide
demonstrations
9* as collection of state-action pairs: 9' - -[5r),a))) We use the notation s' and a' specifically to denote states and actions belonging the expert dataset. The behavior cloning (Pomerleau 1991) algorithm learns a policy nig from this dataset by casting imitation as a typical supervised learning problem arg max "")= [To[[Te (a"
s))]] While the distribution class of ne can be
an arbitrary complex generative model (Lipman et al 2023: Chi et al. 2024), we will start with a
Gaussian parameterization for the sake of simplicity
2.2 WARM-UP: NEIGHBOR MANIFOLD REGULARIZED IMITATION LEARNING
While behavior cloning minimizes only the supervised imitation loss over expert states drawn from Dr, such an objective alone does not control how the policy behaves on states that deviate from the manifold of expert states. In practice, accumulating errors lead the agent to out-of-distribttion regions where a BC policy may act arbitrarily, especially for overparameterized neural networks (Ross& & (eggell((210))
To mitigate this, we note that behavior cloning enforces function evaluations only at the training states, but it does not explicitly take into account the relationship between states (and their corresponding actions) in neighborhood, thereby ignoring the underlying data manifold. To incorporate this information into policy learning, let us consider a modified objective that introduces a regularization term that explicitly encourages local consistency of predictions: nearby states in the expert dataset should be mapped to similar actions. This intuition leads to the following neighborhood-regulan loss (MMRIL)) ), where the standard imitation learning objective (LBc) is combined with additional smoothness penalty (LS) enforcing predictions to respect the geometry of the dataset rather than relying solely on pointwise supervision.
LMRR.(A) = =B R(s0))~-- [((s(s)a)] +2 Eyugi
Σ (ES)
wa(s) ((()))
(1)
supervised risk(.2nc)
smoothness regularizer(
where E(f(s),a) is the supervised imitation loss, Ms) are the k-nearest neighbors of from the expert dataset, and the weights (s) are normalized kernel weights based on the state differences
[1g wi(s) Kg h
[+l- As we discuss briefly below (and in detail in Appendix A.1.1), this corresponds
We show that this can he relaxed in Section 2.4
3

---

Published as a conference paper at ICLR 2026
to form of manifold regularization or Laplacian smoothing, where the policy is penalized for high- frequency variation across the neighborhood of expert states. This manifold regularization provably leads to improvements in policy variance, stability, and generalization.
Theorem 1 (Manifold Regularized BC (mMrH) improves over vanilla BC (LRC). Let I - of be the true, underlying expert policy, assumed to C--smotth on a compact state space f. Let f: + A denote the learned policy estimator. Consider two estimators trained on expert demonstrations:
1. Vanilla BC: a global supervised model minimizing
LBC(f) = (isa)=go(((((ssaa)))
2. MRIL: a neighbor- estimator minimizing
LMRIL() = LBC(f) 2 2
Σ (())) (eE))
where wi(s) are the kernel weights defined above and
Then, under the smoothness assumption on f. the following hold:
(i) Variance reduction: The Laplacian penalty in MRIL acts as a data-dependent Tikhonov regularizer, yielding smaller estimator variance than vanilla BC.
(ii) Smoothness guarantee: Minimizers of LMRR, satisfy a uniform bound on the local Lipschitz constant of f. whereas vanilla BC admits interpolants with arbitrarily large Lipschitz constants between training states.
(iii) Policy stability: In closed-loop rollout, the deviation recursion
A,+1 S L, (s)
accumulates error linearly for vanilla BC, but sublinearly for MRIL, since the smoothness regularizer enforces ([f(s) O[[|s- s'1) for neighbors s.s.
This suggests that MRIL enjoys strictly better generalization and stability guarantees than BC.
Proof sketch. We defer the detailed proof to Appendix A.1.I but provide brief sketch. The key idea of the proof first show that the smoothness regularizer directly corresponds to graph Laplacian penalty graph constructed by k-nearest neighbor (k-NN) affinity matrix defined by the kernel w;. Next, we show that as the number of samples tends to infinity, this graph Laplacian penalty converges to the weighted Dirichlet energy (Belkin & Niyogi 2008 Zhou et al. 2003). Minimizing this Dirichlet energy (1) ensures that the learned f is locally Lipschitz almost everywhere, ensuring smoothness and, in turn, policy stability, and (2) corresponds to Tikhonov regularization, thereby reducing estimator variance, while keeping the bias controlled.
Intuitively, the smoothness regularizer not merely penalizing pairwise disagreements between neighbors, but is driving the learned policy to be smooth with respect to the underlying data manifold. In particular, shrinks the local Lipschitz constant
along directions where the data density p(s) high, ensuring that small changes in state lead small, consistent changes in the predicted action. As result, the policy generalizes more reliably on in-distribution structured manner on new diittriiistiiittionn (OOD) states in the neighborhood.
(ID) states and extrapolates in a
2.3 IMPLICIT MANIFOLD REGULARIZATION VIA IN-CONTEXT ARCHITECTRRES
While MRIL objective does amortize local learning to provide improvements vanilla BC, there are two notable drawbacks. First, the presence of a hyperparameter 2 that must be tuned balance supervised accuracy and smoothness. Second, the requirement optimize modified, regularized objective rather than standard BC objective may modify the optimization landscape in
2,, and T.p are the I. Lipschitz constants of the environment transition dynamics with respect to state and action, respectively. We denote xT(x,) the action predicted by the agent from the time f, St-
4

---

Published as a conference paper at ICLR 2026
adverse ways. This raises natural question: we obtain the same benefits conferred by MRIL (Eq. []. by modifying the policy architecture rather than modifying the objective?
In this section, introduce retrieval-based change in policy architecture that leads implicit manifold regularization effect (iMRIL), despite using a standard imitation objective. With iMRLL we obtain the benefits of Laplacian smoothing (from MRIL) by training standard BC objective (as shown in Fig. 2 without introducing À as an additional hyperparameter for training. We then build this algorithm to develop practical instantiation of this method (DARP) in Section 2.4
iMRIL architecture: The high-level idea be- hind iMRIL is simple we propose moving the neighborhood aggregation (averaging) op- eration from the objective (as in Eq. 1 to the architecture itself. So instead of learning a stan- dard feedforward predictor f(s) that is trained against a neighborhood regularized smoothness objective (Eq. D. we propose embedding the structure of neighborhood aggregation directly into the parameterization of the action predic- tor ff itself, while maintaining the objective standard imitation learning. iMRIL learns the
BC iMRIL
Figure 2: iMRIL implicitly achieves Laplacian smoothing, which reduces variance and enforces local consistency, whereas the lack of smoothness constraint standard BC allows for arbitrarily jagged function approximations.
parameters of per-state predictor fo such that action predictor explicitly parameterized via neighborhoodd-ggrree
f(5*) (ist fo(;) across nearest neighbor states from the training set {s/ ((e)) generates accurate predictions of the corresponding expert action a", With this parameterization, iMRIL optimizes at training time:
Σ fo(si) (E.
-a"
k
((*))
At deployment time, inference can be performed on new state sq simply by retrieving the k-NN of 5q from the training set and performing neighborhood aggregation à fs,) Kis. fe(s) As we show in Section 2.4 the particular parameterization of f is of crucial importance and plays significant role in the empirical performance of iMRIL- leading the development of DARP.
Intuitively, we parameterizing the action predictor f as aggregation of predictions at neighbor states from the training data ) and then learning f. Supervising the past-aggregation function implicitly prevents any f predictions from being arbitrarily non-smooth, conferring the benefits noted Section 2.2
We prove a direct equivalence of iMRIL to the Laplacian regularization in Section 2.2
Equivalence between iMRIL and MRIL: While we defer a full proof of formal equivalence between MRIL and iMRIL to the Appendix Section A.I state our main result and proof sketch to this effect here.
Theorem 2 (iMRIL is parameter-free Laplacian regularization for BC (MRIL)). Consider the symmetric normalized k-NN graph Laplacian L (defined Section 2.2 with eigenpairs ((1j:
where 0 μ1 4 ... Mn 2.
The minimizers of the explicit MRIL objective (Section 2.2 and the implicit iMRIL objective (Sec- tion 2.3 have the following closed form expansions
MMRIL-
-E (a',u))uj
Σ J-1 1+224
PIRIL (1-
iMRIL 's neighbor aggregation step applies the fixed spectral filter OMRRIL (u) = ( 4 the graph Laplacian L, preserving low-frequency modes and suppressing high-frequency modes. The con- gruence between fiMRU. and fMRL shows that iMRIL. is equivalent built-in in form of Laplacian
5

---

Published as a conference paper at ICLR 2026
smoothing (MRIL) with effective À = 1 in normalized units. Unlike explicit regularization, this implicit filter requires no additional hyperparameter tuning.
Proof sketch. We defer full details to Appendix Section A.1.2 spectral solution by diagonalizing the k-NN Laplacian, yielding a filter of the form (1 + 24) on each eigenmode. The implicit objective can be expressed neighbor aggregation Sf with S 5 D A, the random-walk matrix, which has the same eigenvectors and applies the fixed filter - μ. Intuitively, both act low-pass filters the graph: modes with small eigenvalues (smooth variation across the data manifold) are largely preserved, while modes with large eigenvalues (rapid, high-variane fluctuations between neighbors) are strongly damped. Thus iMRIL implicitly performs Laplacian smoothing, reducing variance and enforcing local consistency without needing to tume
The explicit regularizer admits a
Note that the implicit Laplacian smoothing view does not replace the need to learn a policy; rather, constrains the class of functions that can be represented after aggregation. The neighbor-conditions conditioned network fe learns how expert actions vary under local perturbations, proposing locally adapted actions for each neighbor. The aggregation operator then enforces variance reduction by smoothing these proposals across the neighborhood. In this way, learning provides accuracy by correcting local bias, while aggregation provides stability by controlling variance.
2.4 DIFFERENCE- AWARE RETRIEVAL POLICIES: A PRACTICAL INSTANTIATION OF IMRIL FOR IMITATION LEARNING
Given the conceptual framework of iMRIL, we instantiate a practical algorithm for large-scale imitation learning. We build on the objective outlined in Eq. and instantiate a careful choice of (1) parameterization, (2) neighbor aggregation that leads to strong empirical performance.
2.4.1 DIFFERENCE- BASED PARAMETERIZATION OF fo
The objective described in Eq. 2 leaves the parameterizaaion and input representations of fe open to broad interpretation. We make the observation that the neighborhood aggregation should learn how expert actions vary under local perturbations. This suggests that fo should use knowledge of differences between query state and a neighbor state to adaptively propose locally adapted actions for each neighbor. Difference-A ware Retrieval Policies (DARP), instead of simply parameterizing Se by fe(s)), ), we provide additional context about the optimal neighbor action a', as well as the difference between the query state and the neighbor state As = sf sgi a predictor fo predicts an action candidate d; for query state sy and neighbor (s/ a using the difference information = fe; a As; S Sq)-
Let (sg be the index set of the k-nearest neighbors retrieved according to some distance function dss,.,;) For generating predictions with DARP, can then perform neighborhood aggregation (as outlined Section 2.3) to predict an action for any query state Sq
ay = fDARP((q)
Σ df ((.
(3)
Σ
fo(s) af As; = S7 sq)-
At training time, this can be used to define straightforward imitation learning objective from the expert dataset 9*:
arg min E (xgpa)))0* [|la ay [2]
where we optimize for the parameters of the predictor fe. minimizing the discrepancy between the predicted action ag and optimal action ag. Given the simplicity of the objective, any parameterization can be used for fe. in our case, standard feedforward or convolutional neural networks. As we show
In our work we use the Euclidean distance in a pre- gre-trained embedding space, although other neighborhood functions are also applicable. We refer the reader to Appendix Section A.2. for a thorough discussion of design decisions in constructing neighborhood sets via retrieval.

---

Published as a conference paper at ICLR 2026
in Section 3]
this difference-based generate actions to exccute by retrieving k-NN and performing inference through the neighborhood aggregation operation defined in Eq.
parpeeaattttiiitttooooon parameterization is crucial for performance. At inference time,
2.4.2 GOING BEYOND LINEAR AGGREGATION
While the process of neighborhood aggregation thus far has been restricted to averaging neighbor- hood predictions ag Eie. Lie. d/ this is a special case of broader class of permutation-invan -nvvariat aggregation functions 5 ((ts))))) For instance, &y could be parameterized with more expres- sive set-compliant neural models like the set transformer (Lee al. 2019) DeepSets (Zaheer et al.) 2017). This suggests a generalization of the prediction model in Eq. 3 as àg gy( (s) a As; si (e)))e Besides benefits expressivity, generalizing from simple averaging operation parametric aggregation model gy allows for the representation of richer action distributions (e.g Gaussian mixture models (Pignat & Calinon 2019) or diffusion models (Chi et al.
2024)) than
the Gaussian distribution that is implicit the 12 regresssnon objective defined in Eq. Rather than predicting åg directly. DARP can predict the parameters & of an action distribution p(ag: a) for instance the means, covariances, and weights for Gaussian mixture model, or the score function for a diffusion model. This allows DARP to perform maximum likelihood training of multimodal action distributions rather than just unimodal (2-daall2--rr Iz-regression:
arg max E (0 logplag; (ogg)) where as (5g) 8v (ffo(s).a,AAs - (4))ic.(((ss)
(6)
Inference for a query state Sy can be performed by sampling from p(ayi ag (sy)), constructing dg (5₂) = gv (sf af s - Sy) sy)lic.((.))) from a set of neighbors retrieved at test time. We refer readers Appendix Section A.3 for detailed training pseudocode.
2.5 SUMMARY
We show that encouraging local consistency when fitting a behavior cloning policy corresponds to manifold regularization (Laplacian smoothing) and provably improves variance, smoothness, and stability standard BC (Theorem TT. Rather than adding explicit regularization term (and its attendant hyperparameter )) into our loss function, we show that we can induce smoothing implicitly by building neighborhood aggregation into the policy architecture itself (Theorem 2 Thus, propose the following algorithm: at training time, for each expert state-action pair (Ng. ag). DARP retrieves the k-nearest neighbor states from 9*, computes difference vectors As; 5 - Sg from each neighbor to the query state, passes each neighbor tuple (s) az As;) through a network fo to produce candidate actions, and finally aggregates these candidates via a permutation-invarit predict the final action. This is trained only with the standard imitation learning objective.
function gy to
3 EXPERIMENTAL EVALUATION
Next, evaluate DARP in order answer three key questions: Q1: Can DARP consistently outperform standard behavior cloning?, Q2: Can DARP handle more complex state representation and action distributions?, Q3: How do different architectural components contribute DARP's performance gains? We conduct experiments across multiple domains using low-dimensional state representations, nigh-dimensional image features, and diverse action representations. Our evaluation includes continuous control tasks (MuJoCo), robotic manipulation (Robosuite), and specially designed discontinuous environments that stress-test the neighbor- based approach.
3.1 BASELINE COMPARISONS AND TASK DESCRIPTIONS
MuJoCo Tasks:
The MuJoCo Todorov et al. 2012 Fu et al. 2020) tasks entail controlling various
legged figures in multiple embodiments to achieve forward locomotion a flat plane. These tasks include: Hopper (single-legged hopping robot), Walker (bipedal humanoid), Ant (quadruped), and HalfCheetah (biped).
7

---

Published as a conference paper at ICLR 2026
Robosuite Tasks:
The Robosuite (Zhu et al. 2020) tasks all entail single robotic arm manipulating objects. In the Stack task, the goal is to put a smaller cube on top of a larger one. In the Threading task, the goal is to manipulate thin, needle-like tool and insert it into small ring. In the Square Peg task, the goal is to manipulate a square wooden block with a hole in the center and place it onto a square peg.
RoboCasa Tasks: The RoboCasa (Nasiriany et al. 2024) tasks all entail a single robotic arm manipulating objects in randomized kitchen setting. In the Drawer task, the goal is to close an open drawer. In the Door task, the goal is to close an open microwave oven door. In the Stove task, the goal is to twist a knob to turn off a stove burner.
Baseline Comparisons: We compare DARP against variety of baselines and ablations: (I)
R&P (Sridhar al. 2025):: refers to directly taking the action corresponding to the nearest neighbor, (2) LWR (Pari et al., 2022): refers to performing locally weighted regression retrieved neighbors,
(3) BC: refers to standard parametric behavior cloning, (4) REGENT (Sridhar al. 2025); : refers
to transformer-based in-context learning method conditioned on retrieved neighbors, (5) MRIL:
refers to the explicitly smoothed version DARP outlined Section 2.2
3.2 CAN DARP CONSISTENTLY OUTPERFORM STANDARD BEHAVIOR CLONING? (Q1)
In this experiment, we evaluate DARP's core hypothesis on tasks with low-dimensional state represen- tations, where the distance metrics between states well-defined and interpretable. This evaluation spans locomotion tasks from MuloCo and robotic manipulation tasks from Robosuite and RoboCasa with data generated with MimicGen (Mandlekar al., 2023). In these experiments. the aggregation function is implemented as a simple average of all neighbor action predictions d'.
Method
Hopper
Ant
Walker
HalfCheetth
R&P (Sridhar et al.
2025
711.82 # 85.63
-305.97 = 76.42
419.18 + 50.21
-178. 64 = 29.75
LWR (Part et al.
2022
1703.78 + 245.95
846.59 +: 216.06
1484.91 +: 356.54
1945.82 a: 567.26
BC
2313.65 + 203.75
2376.20 339.43
2658 40 + 274.08
1063.23 # 371.08
REGENT Sridhhar et al.
2025
1819.39 + 186.24
-302.10 146.67
507.01 + 76.10
169.85 + 63.10
MRIL
2793.63 # 156.41
3869.08 + 241.00
4370.96 + 168.13
701.58 # 195.08
DARP
3545.57 # 3.54
4383.28 # 266.37
4894.01 # 75.12
5515.41 # 841.33
DARP Set Transformer
2965.86 # 103.08
4063.79 + 218.80
4752.42 # 109.23
3417.85 a 764.57
Table 1: Both DARP and DARP Set Transformer outperform other approaches across all domains. Performance Comparison of DARP BC and other baselines across MuJoCo Environments Using Low- Dimensional State. Scores reported are averaged across 100 independent trials with 95% confidence intervals.
We find that DARP demonstrates substantial improvements over standard behavior cloning across all tested environments. We observe performance gains ranging from 15-25% points in robotic manipulation tasks and significant score improvements in locomotion tasks (see Table Table 3 and T2b We observe that purely non-parametric methods (R&P and LWR) perform poorly on these tasks, and while MRIL is nearly always able to get score higher than vanilla BC, the highest scores on this suite of tasks are always achieved by our DARP architecture.
Given the changes introduced for the practical instantiation in Section 2.4 we evaluate whether DARP scales up to higher- dimensional input representations such as images.
Method
Stack
Thrd.
Pog
Method
Drawer
Door
Stove
R&P
38
11
31
BC
54
29
28
LWR
21
39
30
DARP
85
45
43
BC
47
37
46
DARP
72
63
62
Table Comparing across RoboCasa Environments using low-dimensional state features. Scores are listed as success percentage. DARP outperforms BC.
Table 3: Comparing across Robosuite Environments using low-dim state features.
8

---

Published as a conference paper at ICLR 2026
3.3 CAN DARP HANDLE MORE COMPLEX STATE REPRESENTATION AND ACTION DISTRIBUTIONS? (Q2)
ffihh-Dimeniionl Visual Input Representations. To test the applicability of DARP beyond the regime of compact, low-dimensional states, evaluate DARP on simulated robotic manipulation tasks using R3M image embeddings (Nair et al. 2022). This tests whether the neighbor-hased approach remains effective when states are represented as nigh-dimensional feature vectors ex- tracted from visual observations (see Table Observe that, not only does DARP outperform standard BC, the average improvement, 35%, is actually higher than the average improvement Robosuite tasks in low-dimensional state (~~ 22%).. Empirically, this means that DARP was better at adapting to complex, figh-dimensional state representations than standard BC.
Method
Stack
Thrd. Peg
BC
44
38
17
DARP
75
76
52
Table 4: Success rates (%) on vision-based Robosuite tasks.
Multi-modal Action Distributions. We show that DARP can solve complex multimodal imitation learning tasks such as the Push-T environment over 20% better than behavior cloning. We defer details to Appendix A.2.2
3.4 How DO DIFFERENT ARCHITECTURAL COMPONENTS CONTRIBUTE DARP'S PERFORMANCE GAINS? (Q3)
Ablation Study: To understand which components of the DARP architecture contribute most to its performance gains, we conduct a comprehensive ablation study examining each design choice, namely (I) standard DARP; (2) DARP, but without including the neighbor actions; (3) an ensemble of 10 BC agents; (4) DARP, but we choose random neighbors opposed to using distance metrie; (5) DARP, but we take the L2 norm of the distance vector; (6) BC baseline, which is just the query state sy; (7) DARP, but include just the query state rather than the distance vector between the query state (so not cormutation-nnvaiinne
and neighbor states; (8) DARP, but using a
aggregator combine all d's. We report Figure 3 the results of this systematic ablation.
The ablation study reveals that distance vectors and permuta- tion invariance are crucial for DARP's success, while neighbor actions have more modest im- pact. Random neighbor selec- tion performs poorly, confirm- ing that meaningful distance met- rics informing neighbor selection are crucial. The permutation- invariant aggregation function g proves critical, as permutation- dependent alternatives signifi- cantly degrade performance.
Model Architecture Ablation
0.85
Success Rate
0.59
0.55
0.47
0-42
0.19
0.00
11.0 (4)
BC Been
Radoo (g-4).l0 -xll)
(6 (g:4))4)) Prr Dep.
Figure 3: Distance vectors and permutation invariance contribute heavily to DARP's success. Exploration of how the performance of DARP agent is impacted various changes architecture demonstrates that DARP success is most attributed the distance vectors 16. sg). Success rate is averaged across trials on the Robosuite Stack environment with 95% confidence intervals.
made the
Divergence Analysis: To better understand DARP's success over standard BC, we analyze the point of divergence in rollouts which the latter fails but the former succeeds. We identify the "step of divergence" as the point at which DARP and BC begin to receive significantly different reward. We define 7, and va as the percentile of likelihoods of the training set (that is, 1 of the deltas seen at training time are less likely than T.). In all six different rollouts across two different tasks (the Robosuite Stack task and the MuloCo Hopper task), the query state the SoD has state likelihood of but delta likelihood of > This result bolsters our hypothesis that DARP gains occur partly due to improved prediction on slightly out-of out-of-distribution states due to reparameterization in terms of difference vectors to neighbors. (see Figure for plots of reward drift, SoDs, and state and delta likelihood for one task.) See Appendix A.2.4 for additional experiments regarding DARP robustness.

---

Published as a conference paper at ICLR 2026
Stap
Step
200
Scp
200
Figure 4: Cumulative rewards for BC and DARP on the Robosuite stack task illustrate initially identical rollouts that diverge as BC fails the task and DARP succeeds. A vertical dashed line indicates the step in which the two diverge, labeled "SoD". At the SoD, the state likelihood is < 5 (OOD), but the delta likelihood is TA (in distribution).
RELATED WORK
Non-Paree Mon-Parametric Imitation Learning Methods: Non- Non-parametri IL algorithms demonstrate surpris- ing performance by leveraging local structure. VINN Pari et al. (2022) explores locally weighted regression for imitation, showing surprising results in image embedding spaces, and MiDiGaP von Hartz al. (2255) uses mixtures of Gaussian processes model multimodal trajectories and achieve rapid generalization. SEABO Lyu et al. (2024) uses retrieval methods to perform offline RL by rewarding transitions close to neighbors to form a reward function. FlowRetrieve Lin et al. (2024), STRAP Memmel et al. (2025), and Behavior Retrieval Du et al. (2023) perform non- con--arametric re- trieval and finetuning from large unlabeled datasets, enabling generalization through test-time training. DARP differs from the above in its unique parameterization of retrieved states into (sf ,a, - 8g) tuples and learning a semi-paramee parametaaetticc policy rather than relying purely on non-parametric aggregation or test-time training. This provides the variance reduction of local methods and generalization of parametric policies.
Much recent literature has explored explicit smoothness constraints to improve policy stability and robustness. L.2C2 1.2C2 Kobayashi (2022) considers model-free RL under local Lipschitz continuity constraints, achieving smoothness and noise robustness without sacrificing expressiveness, while Asadi et al. (2018) proposed similar methodology for model-based RL models with Lipschitz constraints. CCIL, synthetic corrective labels for imitation learning using Lipschitz-constraieed dynamics model. This has also been scaled up to humanoid controllers [Cen et al. (2024) to reduce shakiness on deployment. DARP differs from these methods by enforcing smoothness implicitly through architecture change while using standard imitation learning objectives.
Cooothhnes-Constrannee Policy Learning:
al. (2024a) extends these ideas to generate
SonnrollersChhen
In-Context Learning Methods: Recent work has explored non-parametric retrieval from the perspective in-context imitation learning. REGENT et al. (2025) investigates retrieval- augmented generalization by incorporating retrieved states, actions, and rewards into causal transformer, while DPT et al.] (2023) uses supervised pretraining for transformers to predict actions given query states and in-context datasets, effectively learning how explore. Other in- context architectures include ICRT Fu et al. (2024), Instant Policy Vosylius & Johns (2025), and KAT Di Palo & Johns 2024). These methods aim quickly adapt new tasks and environments, whereas DARP focuses on accomplishing higher performance and stability on standard imitation learning.
5 CONCLUSION
We introduced Difference-Awar Retrieval Policies for Imitation Learning (DARP) (DARP), nearest- neighbor-based algorithm that reparameterizes the imitation learning problem in terms of relative differences between query states and their nearest neighbors, rather than learning direct mappings. We prove that our method implicitly achieves Laplacian smoothing. Our experimental evaluation across diverse domains, including continuous control and robotic manipula- tion, validates three key hypotheses. First, DARP consistently outperforms standard behavior cloning when using oiw-dimensional state representation. Second, DARP maintains performance across different state representations, action distribution modeling requirements, and task complexities, with improvements ranging from 15-46% across tested scenarios. Third, architectural ablations reveal that distance vectors and permutation-invarian aggregation are crucial components to our algorithm.
10

---

Published as a conference paper at ICLR 2026
6 REPRODUCIBILITY STATEMENT
link to supplementary source code is provided. This codebase contains code used to train and evaluate models. It also contains policy and environment configuration files to generate all results seen in this paper. We provide all data used in MuloCo experiments and provide scripts to generate expert demonstrations for Robosuite tasks via MimicGen. We also provide all code necessary to transform between different modalities, such as low-dimensional state representation to images to R3M features. Results will be identical to those in the paper on NVIDIA L40 and L40s GPUs, with the exception of results that require the use of transformer (REGENT, Set Transformer), which are non-deterministic and may differ slightly from reported numbers.
7 ACKNOWLEDGMENTS
This work was (partially) funded by grants from the National Science Foundation NRI (#2132848), DARPA RACER and #2022-016-01 UW). We gratefully acknowledge gifts from Amazon, Collaborative Robotics, Cruise, the Toyota Research Institute under the University 2.0 program, the Research Scholarship from the Mary Gates Endowment for Students, and others.
and the Office of Naval Research (#N00014-24-S) --B001
REFERENCES
Kavosh Asadi, Dipendra Misra, and Michael Littman. Lipschitz continuity in model-based reinforce- ment learning. In International Conference on Machine Learning. PMLR, PMLR, 2018.
Mikhail Belkin and Partha Niyogi. Towards theoretical foundation for laplacian basedd manifold methods. Journal of Computer and System Sciences, (448))1229-1338 2008.
Kevin Black, Noah Brown, Danny Driess, Adnan Esmail, Michael Equi, Chelsea Finn, Niccolo Fusai, Lachy Groom, Karol Hausman, Brian Ichter, Szymon Jakubczak, Tim Jones, Liyiming Ke, Sergey Levine, Adrian Li-Bell, Mohith Mothukuri, Suraj Nair, Karl Pertsch, Lucy Xiaoyang Shi, James Tanner, Quan Vuong. Anna Walling, Haohuan Wang, and Ury Zhilinsky. Mo. A vision-languaae- action flow model for general robot control, 2024. htRp:ttp:: //t::////iivv. orq /bb////b//// 2410 24164
Zixuan Chen, Xialin He, Yen-Jen Wang. Qiayuan Liao, Yanjie Ze, Zhongyu Li, S. Shankar Sastry, Jiajun Wu, Koushil Sreenath, Saurabh Gupta, and Xue Bin Peng. Learning smooth humanoid locomotion through lipschit--onnsraained policies. arxiv preprint arXiv:2410.1 1/825 2024.
Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal of Robotics Research, 2024.
Fan R. K. Chung. Spectral Graph Theory, volume 92 of CBMS Regional Conference Series in Mathematics. American Mathematical Society, 1997. doi: 10.1099/cbms0092.
Michael Yoon Chung, Maxwell Forbes, Maya Cakmak, and Rajesh PN Rao. Accelerating imitation learning through crowdsourcing. 2014 IEEE International Conference on Roboties and Automation (ICRA), pp. 4777 4784. IEEE, 2014.
Norman Di Palo and Edward Johns. Keypoint action tokens enable in-context imitation learning in robotics. In Proceedings of Robotics: Science and Systems (RSS), 2024.
Maximilian Du, Suraj Nair, Dorsa Sadigh, and Chelsea Finn. Behavior retrieval: Few-shot imitation learning by querying unlabeled datasets. In Proceedings of Robotics: Science and Systems (RSS), 2023.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. 2020.
Letian Fu, Huang Huang, Gaurav Datta, Lawrence Yunliang Chen, William Chung- Ho Panitch, Fangchen Liu, Hui Li, and Ken Goldberg. In-context imitation learning via next-token prediction. arXiv preprint 2408 55980, 2024.
11

---

Published as a conference paper at ICLR 2026
Liyiming Ke, Yunchu Zhang, Abhay Deshpande, Siddhartha Srinivasa, and Abhishek Gupta. Ceil: Continuity- data augmentation for corrective imitation learning. In The Tweifth International Conference Learning Representations, 2024a.
id=1.961.8f449
Liyiming Ke, Yunchu Zhang, Abhay Deshpande. Siddhartha S. Srinivasa, and Abhishek Gupta. CCIL: continuity- based data augmentation for corrective imitation learning. In The Tweifth Inter- national Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. net/ forum wwtoorumiiiigge T
OpenReview.net, 2024b. httpthtss////ppeeeeeee
Taisuke Kobayashi. L2c2: Locally lipschitz continuous constraint towards stable and smooth reinforcement learning. In 2022 IEEE/RSJ International Conference Intelligent Robots and Systems (IROS), pp. 4032-4039, IEEE, IEEE, October 2022.
Jonathan Lee al. Supervised pretraining can learn in-context reinforcement learning. In Advances in Neural Information Processing Systems, volume 36, pp. 43057-43083, 2023.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set trans- former: A framework for attention permutation-invarian neural networks. In Proceedings of the 36th International Conference on Machine Learning, pp. 3744 3753, 2019.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement Jeammi Tutorial, review, and perspectives on open problems, 2020. URL 01643
Li-Heng Lin, Yuchen Cui, Amber Xie, Tianyu Hua, and Dorsa Sadigh. Flowretrieval: Flow- guided data retrieval for few-shot imitation learning. In Pulkit Agrawal, Oliver Kroemer, and Wolfram Burgard (eds.), Conference on Robot Learning, 6-9 November 2024, Munich, Germany, volume 270 of Proceedings of Machine Learning Research, pp. 4084-4099, PMLR, 2024. URL https::/proceedingg www.nng.mmm....mmmm..m.......eme mlr www.spessseee..... wttpss//voo/Liin2a..o. html
Yaron Lipman, Ricky Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In International Conference Learning Representations (ICLR) 2023, 2023, URL htt http::://arriiv.orr//a//////2 as arXiv preprint arXiv:22 10.00777..
wwtp:oorg/bbbb///o// 02747 Originally appeared
Jiafei Lyu, Xiaoteng Ma. Le Wan, Runze Liu, Xiu Li, and Zongqing Lu. Seabo: A simple search-based method for offline imitation learning. In International Conference on Learning Representations, 2024.
Ajay Mandlekar, Soroush Nasiriany, Bowen Wen, Iretiayo Akinola, Yashraj Narang, Linxi Fan, Yuke Zhu, and Dieter Fox. Mimicgen: A data generation system for scalable robot learning using human demonstrations. 7th Annual Conference on Robot Learning, 2023.
Elman Mansimov and Kyunghyun Cho. Simple nearest neighbor policy method for continuous control tasks. In International Conference on Learning Representations (ICLR) 2018, 2018. Under review conference paper at ICLR 2018
Marius Memmel, Jacob Berg, Bingqing Chen, Abhishek Gupta, and Jonathan Francis. Strap: Robot sub-trajectory retrieval for augmented policy learning. In The Thirteenth International Conference on Learning Representations, 2025.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3M: A universal visual representation for robot manipulation. In Karen Liu, Dana Kulic, and Jeffrey Ichnowski (eds.), Conference on Robot Learning. CoRL 2022, 14-18 December 2022, Auckland, New Zealand, volume 205 of Proceedings of Machine Leurning Research, pp. 892-909. PMLR, 2022. URL htttp::/pprocee.nggs ml www..prres/v200/raaa
Soroush Nasiriany, Abhiram Maddukuri, Lance Zhang, Adeet Parikh, Aaron Lo, Abhishek Joshi, Ajay Mandlekar, and Yuke Zhu. Robocasa: Large-scale simulation of everyday tasks for generalist robots. In Roboties: Science and Systems (RSS), 2024.
12

---

Published as a conference paper at ICLR 2026
Jyothish Pari, Nur Muhammad Shafiullah, Sridhar Pandian Arunachalam, and Lerrel Pinto. The surprising effectiveness of representation learning for visual imitation. In Robotics: Science and Systems (RSS). Robotics: Science and Systems Foundation, June 2022. doi: 10.15607/RSS.20222.
XVIII-052.
Emmanuel Pignat and Sylvain Calinon. Bayesian gaussian mixture model for robotic policy imita- tion. IEEE Robotics and Automation Letters, 2019. URL
10716. pdf Preprint version.
Dean Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural Compur., 3(1):88 97, 1991. doi: 10.1162/NEC0.999 1991.39 1991.3.1.88 3.8.8.. URL httppttts////o.rrr//
1162/neco. 1991. 3.1.88
Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings the thirteenth international conference on artificial intelligence and statistics, pp. 661-668. JMLR Workshop and Conference Proceedings, 2010.
Stéphane Ross, Geoffrey J. Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Geoffrey J. Gordon, David B. Dunson, and Miroslav Dudík (eds.), Proceedings of the Fourteenth International Conference Artificial Intelligence and Statistics, AISTATS 2011, Fort Lauderdale, USA, April 1-13, 201 volume 15 of JMLR Proceedings, pp. 627-635. JMLR.org, 2011. URL http::/pprocedddngggee mlr
press/v15 10445/oosslee wtt.ssssaaarssssssseeesessssssss ross
Steven L. Salzberg and David W. Aha. Learning to catch: Applying nearest neighbor algorithms to dynamic control tasks. In P. Cheeseman and R. W. Oldford (eds.), Selecting Models from Data: Artificial Intelligence and Statistics IV, volume 89 of Lecture Notes in Statistics, pp. 321-328. Springer, 1994. doi: 10.1007/9788 1111779788---612-220 1---1---1666----- /00 1007000100000880000
URL https: //tps:///ink... workkssprnnggg nger 33
Som/chapterr/00
Jonathan C. Spencer, Sanjiban Choudhury, Arun Venkatraman, Brian D. Ziebart, and J. Andrew Bag- nell. Feedback in imitation learning: The three regimes of covariate shift. CoRR, abs/2102.02872, 2021. URL hatpl:://///// arxiv. org /aabs/210.... 02872
Kaustubh Sridhar, Souradeep Dutta, Dinesh Jayaraman, and Insup Lee. REGENT: A retrieval- augmented generalist agent that can act in-context in new environments. In The Thirteenth International Conference on Learning Representations (ICLR), 2025. Oral presentation.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEF/RSJ international conference on intelligent robots and systems, pp. 5026-5033. IEEE, 2012.
Arun Venkatraman, Martial Hebert, and J. Andrew Bagnell. Improving multi-step prediction of learned time series models. In Blai Bonet and Sven Koenig (eds.). Proceedings of the Twenty- Ninth AAAI Conference on Artificial Intelligence, January 25-30, 2015, Austin, Texas, USA, pp.
5024-3030 AAAI Press, 2015. doi: 10. 10.1609/AAALVV291 1990. URL pettps://doi.orge
To. 1 609 aaat v2..1.11..... 9590
Jan Ole von Hartz, Adrian Röfer, Joschka Boedecker, and Abhinav Valada. The unreasonable effectiveness discrete-time gaussian process mixtures for robot policy learning. arXiv preprint arXiv:2505.03226. 2025.
Vitalis Vosylius and Edward Johns. Instant policy: In-context imitation learning via graph diffusion. In Proceedings of the International Conference on Learning Representations (ICLR), 2025.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan R Salakhutdinov, and Alexander J Smola. Deep sets. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gamett (eds.), Advances in Neural Information Processing Systems 30, pp. 3391-3401. Curran Associates, Inc., 2017. URL http: http:::/pppprs.niii nips cc/paper/ 6931 deep- s.ff
Tony Z. Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine- fine-grained bimanual manipulation with low-cost hardware. In Proceedings of Robotics: Science and Systems (RSS), 2023.
13

---

Published as a conference paper at ICLR 2026
Dengyong Zhou, Olivier Bousquet, Thomas Navin Lal, Jason Weston, and Bernhard Schölkopf. Learning with local and global consistency. Advances in Neural Information Processing Systems (NeurIPS), volume 16, 2003.
Yuke Zhu, Josiah Wong, Ajay Mandlekar, Roberto Martin-Martin, Abhishek Joshi, Soroush Nasiriany, Yifeng Zhu, and Kevin Lin. robosuite: A modular simulation framework and benchmark for robot learning. In arXiv preprint (tXXx::0099122990 2020.
14

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 LEMMAS AND PROOFS
A.1.1 PROOF THEOREM
To prove Theorem we first start with a well-known result (Chung 1997 Zhou et al. 2003
Belkin
& Niyogi 2008).
Lemma (Smoothness regularizer as k-NN graph Laplacian penalty). Let (s be the expert states with corresponding predicted actions f(;; E Rd, For each i, let (s,,) denote the indices of the k-nearest neighbors of s; (excluding i). Define asymmetric weights
(w)(s)), FEEM)ss) 1.
Wo
otherwise.
and construct symmetric affinity matrix
Wij (Wi WR).
Let be the degree matrix with D = Σ, Wij, and define the k-NN graph Laplacian L = D
Then the smoothness regularizer can be written the quadratic
Ls(f) = - Σ
(y(ss)
TF(FLLF).
(E
where F =
E Raxd, Equivalently, in the scalar case,
Ls() f
Corollary 1 (Continuum limit of smoothness regularizer). Assume states {si)n are sampled i..d. from smooth density p(s) supported on m-dimensional C² manifold M Rd. Let W be the symmetrized k-NN affinity matrix constructed from a kernel Kg with bandwidth h, and let I. 1.= 1) W be the graph applccian.
If o oo - O, and nlr + oo, then the normalized quadratic form converges to the weighted Dirichlet energy:
1 TT(F LF) - CK / I (
where Cx O is constant depending only the kernel Ks-
Theorem 1 (Manifold Regularized BC (EMRI)) improves over vanilla BC (LBC)). Let f" I - a be the true, underlying expert policy, assumed to -m-oooth on a compact state space P. Let f: + denote the learned policy estimator. Consider two estimators trained expert demonstrations:
1. Vanilla BC: a global supervised model minimizing
Lec(f)
E(say- [(()(s)
2. MRIL: a brighhborbbsedd estimator minimizing
LMRUL(F) = -BC(S) + 2 2E
Σ (s()) - /(s))
(ee))(s))
where wi(s) are the kernel weights defined above and À 0.
Then, under the smoothness assumption on f. the following hold:
(i) Variance reduction: The Laplacian penalty in MRIL acts as a data-dependent Tikhonov regularizer, yielding smaller estimator variance than vanilla BC.
15

---

Published as a conference paper at ICLR 2026
(ii) Smoothness guarantee: Minimizers of LMRIL satisfy a uniform bound on the local Lipschitz constant of f. whereas vanilla BC admits interpolants with arbitrarily large Lipschitz constants between training states.
(iii) Policy stability: closed-loop rollout, the deviation recursion
A++1 S L, L (s))
accumulates error linearly for vanilla BC, but sublinearly for MRIL, since the smoothness regularizer enforces (lf(s) f(')ll = O[lss s for neighbors s,3.
This suggests that MRIL enjoys strictly better generalization and stability guarantees than BC.
Proof. We prove each claim in turn.
(i) Variance reduction. The Laplacian penalty in LMRII is
Σ Wi ((s)) = 2 lf,
where L= D W is the graph Laplacian, and W is the k-NN affinity matrix. By Lemmal I this equals the empirical Dirichlet energy of f on the k-NN graph. It is well known (Zhou et al. 2003 Belkin & Niyogi 2008) that such quadratic penalty is equivalent Tikhonov regularization with respect to the graph Laplacian norm || f Lf. In statistical learning theory, adding Tikhonov penalty strictly reduces the variance of the estimator compared the unregularized solution while keeping the bias term controlled. Thus MRIL enjoys smaller estimator variance than vanilla BC, which uses no such penalty.
(ii) Smoothness guarantee. Consider the continuum limit (Corollary I for i.i.d. samples {s;} from density on a smooth manifold M, the normalized penalty converges to
[V][(s)] p(s)
dvs))dvs(ss). dvol(s).
This is the weighted Dirichlet energy of f on M. If this integral is finite, f belongs to the Soboley space H (MM, )²). and particular ff locally Lipschitz almost everywhere with
I - f(s')) s Clss- for p-a.e. neighbor pairs s.s'.
Therefore minimizers of .MRH. have uniformly bounded local Lipschitz constants along high-density regions of the state space. By contrast, minimizers of vanilla BC have such constraint: any oscillatory interpolant that matches the training data exactly yields the same supervised risk, so arbitrarily large Lipschitz constants are possible.
(iii) Policy stability.
Let A, = ||s, - st denote the deviation at time 1. For Lipschitz dynamics T,
A+1 L,A, <=ll -s)))
Decompose the action error:
[z(s) +1((s) (16)))-
For vanilla BC, the first term [x((s)) [sx)]] is only minimized on the empirical distribution py; off-distribution, may be ((1) regardless of A, The second term satisfies [[f(s)) -(())) 0(4) by smoothness of f. Hence the recursion can take the form
A L, La
which accumulates linearly in
For MRIL, the Laplacian penalty enforces
(T(s - x(5)|| S C|ls- for neighbor pairs (s,5'),
L, and l.p, are the I. Lipschitz constants of the environment transition dynamics with respect to state and action, respectively. We denote sT(x,) the action predicted by the agent from the time St-
16

---

Published as a conference paper at ICLR 2026
shown in part (ii). Thus | = O(h²) (where h is the radius of the kernel band- width around the point 5, by local-linear regression error bounds, and [[f((s) f(,)) = 0(4, by smoothness of f. Combining these,
+ A+t1 S +, (..
Since the constant multiplying is strietly smaller under the smoothness constraint, the cumulative егror grows strictly slower than in vanilla BC. In particular, error growth is sublinear in the rollout horizon when is small, whereas it is linear for vanilla BC
Conclusion. Claims (i)- establish that MRIL yields lower variance, uniform smoothness control, and sublinear rollout error accumulation compared to vanilla BC, completing the proof.
Kernel choice. For the IC smoothness regularizer, we adopt Gaussian kernel
ws(s) oc cxp
||s-
Σ wi(s) = 1, (E
with bandwidth set to the median distance to the k-th nearest neighbor across the dataset. This choice is standard in manifold regularization (Belkin Niyogil 2008 Zhou al. (2003) and ensures that the graph Laplacian penalty converges to the Dirichlet energy in the continuum limit. In practice, we found this default be stable across tasks, though other kernels (e.g.. uniform k-NN exponential decay) yield qualitatively similar results.
A.1.2 PROOF THEOREM 2
We begin with the required Lemmas establishing the spectral form of explicit Laplacian regularization and neighbor aggregation.
Lemma 2 (Spectral form explicit Laplacian regularization). Let I. be the symmetric normalized
graph Laplacian with eigenpairs ((4), where Mi S U2 S S Un 2. The minimizer of the penalized objective
Lx() -1 LS
has the closed-form expansion
fa. Σ j=1
(a" (a',,)) uj.
I + +22
Thus directly determines the spectral filter da (1) (1 -+) applied to each Laplacian mode.
Proof. Diagonalize L. T with = (u1,..., Un] orthogonal and .... Write f Uc, a" = Ub in this basis. The objective becomes
Uc UUc-Ubl2 + 2c Ac
AC = Σ
Minimizing each term yields cj = T+211 bj. Transforming back gives the stated expansion.
Lemma 3 (Spectral form of neighbor aggregation). Let S D be the random-walk matrix of the K-NN graph, with adjacency A and degree D. For any prediction vector f. the neighbor-average prediction is f Sf. In the Laplacian eigenbasis, this corresponds to the spectral filter filter
f Σ (1))(
i.e. (DARP(")) μ.
Proof. By definition, L - and S D- = 1 Lyw where L,w D is the random-walk Laplacian. Since Law and L share the same spectrum up similarity transform, the eigenbasis { diagonalizes S. Thus for each mode uj. Su, = - 4y)uj, yielding the claimed spectral filter.
17

---

Published as a conference paper at ICLR 2026
Theorem 2 (iMRIL is parameter-free Laplacian regularization for BC (MRIL)). Consider the symmetric normalized k-NN graph Laplacian L (defined Section 2.2), with eigenpairs (H4),)))
where μi S 12 - ... Um -
The minimizers of the explicit MRIL objective (Section 2.2 and the implicit iMRIL objective (Sec-
tion 2.3 have the following closed form expansions
fMRI. Σ + 2,1 (a',u)) "j
AIRRI. -1
iMRIL 's neighbor aggregation step applies the fixed spectral filter OIRRIL 1 - to the graph Laplacian L, preserving low-freqeency modes and suppressing high-frequency modes. The con- gruence between fIMRU. and fARu, shows that iMRIL is equivalent to built-in form of Laplacian smoothing (MRIL) with effective à =s 1 in normalized units. Unlike explicit regularization, this implicit filter requires no additional hyperparameter tuning.
Proof. From Lemma the neighbor aggregation operator S= D-11 'A acts on Laplacian eigenmodes "
= (
where U4 are the normalized Laplacian cigenvalues. Thus in the graph Fourier basis, neighbor aggregation corresponds to multiplying each mode by the fixed spectral filter @DARP (1) 1 1-u
On the other hand, Lemma 2 shows that explicit Laplacian regularization with parameter yields the spectral filter 6x (u) (1+ + 24)-1. Both filters downweight high- fiehfrruyency modes (μ >> 0) while preserving low-frrquee cww-frequency modes (μu ss 0). The key difference is that Φ2 (u) requires tuning A, whereas (H) is parameter-free.
To see the equivalence, note that for small μ4,
(D))P( = ~ Thus DARP can be interpreted as performing Laplacian smoothing with an effective regularization weight of order 2. = 1 in normalized units. Moreover, for large μ. ΦDARP(H) damps geh-frrquency modes even more strongly by driving them toward zero, providing a sharper low-pass effect than explicit regularization.
()) up to O(u²) terms.
Therefore, DARP aggregation step mathematically equivalent to implicit Laplacian regularization with fixed spectral filter ΦDARP, eliminating the need to tune 2 explicitly.
DARP can therefore be viewed as a form of locally adaptive implicit regularization: rather than introducing explicit global weight 2, its neighbor aggregation step enforces smoothness auto- matically through the graph structure. The effective regularization strength varies with local degree
Spectral Filters: Explicit Laplacian vs. DARP
QDARP (M)
(2-0,5(4)
Spectral filter (())
(2-1(14)
0.5
0
-0.5
0.2
0.6
0.8
1.2
1.8
Figenvalue
Figure 5: DARP achieves sharper low-pass filtering
18

---

Published as a conference paper at ICLR 2026
and neighborhood geometry, adapting to the density of the expert demonstrations. Spectrally, this corresponds to the fixed filter (DARP(H)) = - which suppresses high- frequency modes more aggressively than any fixed explicit Å. Figure Saarre/)jiluuuue illlsttrtes this comparison, showing how DARP achieves sharper low-pass filtering without the need for hyperparameter tuning.
A.2
ADDITIONAL EXPERIMENTAL DETAILS
A.2.1
RETRIEVAL
DARP Success Rate vs. Number of Neighhoss
DARP Success Rate vs. Lookback
DARP Success Rate vs. Decay
Success Rate
Success Rate
Success Rate
DARP
DARP BC
DARF
BC
BC
Lookback
Decay
5
2000
50
Figure 6: DARP performance analysis as retrieval hyperparameters swept: (left) observe that the perfor- mance of DARP model poor when using few neighbors, reaches global optimum when retrieving about neighbors, and plateaus just above BC's success rate k goes the size the performance of DARP model generally slightly improves as more history is considered, and only performs worse than BC when very little no history is considered; (right) observe that the performance of a DARP model is sensitive how much weight applied older observations when performing retrieval. Intuitively, if decay is too high. DARP performance is nearly identical having little to lookback, performing worse than BC. Success rate measured out 50 trials on the Robosuite Stack environment. 95% confidence intervals are included.
dataset; (center) observe that
The selection of the distance function s((),s))) to select k neighbors is crucial. While we find that simple Euclidean distance between states can work, in our experiments, we use a slightly modified algorithm that takes advantage of the fact that are working with sequences of states and incorporates history in our distance calculation.
where Sg.0 the current query state Sg. Now suppose we want to calculate d((sy ). where s/ is some state from the expert dataset. We first find the trajectory this state is from-call this and the index of s/ in this trajectory-call this Thus, sf can be rewritten Sj.- Given some lookback parameter l which denotes how many past states want to consider, we get:
Suppose have query trajectory Sg (....-99-1::88
d(sg:s7)
Σ ||s si,- ==
This is simply the accumulation of Euclidean distances of the current and last - states from both the query trajectory and the source trajectory, assuming valid indices. Of course, in practice, we generally want to put more emphasis more recent states, want them to be more inffuential in the selection of neighbors. Thus, given some rate of exponential decay y 2 O, we have
E-1 Σ R=0
d(sg. s1)
See Figure 6 for experimental analysis how the success rate in environment changes these parameters are swept.
A.2.2 CAN DARP HANDLE TASKS REQUIRING THE REPRESENTATION OF MULTI-MODAL ACTION DISTRIBUTIONS?
We test DARP's ability handle complex action distributions by evaluating the Push-T task, as described in (Chi et al. 2024). which requires representing multi-mode action distributions (see Figure / for visualization). For this experiment, DARP employs Set Transformer head that predicts parameters of a Gaussian Mixture Model. We note that DARP with GMM head is able to handle multi-modal distributions effectively, showing 22% improvement BC the Push-T
19

---

Published as a conference paper at ICLR 2026
Method
Score
BC
48 = 8
DARP
70 8
Table 5: Push-T Results. Aver- aged 100 trials, DARP out- performs BC.
Figure 7: Push-T Environment. The goal is control the blue cir- to push the T-shaped block.
task (Q2) (see Table 5. This demonstrates that DARP can be further adapted to multi-modal action distribution modeling requirements.
A.2.3 CAN DARP HANDLE DISCONTINUOUS ENVIRONMENTS WHERE NEARBY STATES REQUIRE OPPOSING ACTIONS?
key concern for neighbor-based approaches is performance in environments with strong discontinu- ities, where states that are close in Euclidean distance may require drastically different actions. To address this concern, we design a stress test using a modified version of D4RL's Umaze environment (see Figure & for visualization).
Even in this deliberately challenging discontinuous environment, DARP achieves a 57 success rate compared to BC's 25%. (Q3) (see Table 6 This suggests that the distance vectors and permutation-innarian aggregation help the model distinguish between appropriate and inappropriate neighbors, even when spatial proximity doesn't 'tt guarantee action similarity.
Method
Succ. (%)
BC
25
DARP
57
Table 6: Long maze re- sults. Averaged over 100 tri- als, DARP significantly out- performs BC.
Figure 8: Long maze environ- ment. The goal is move a force-actuated ball from the green start to the red destina- tion.
A.2.4 CAN DARP RECOVER FROM BC ERROR?
In order to analyze DARP's robustness to accumulated error, we roll out a BC agent in an environment in which we know it will fail, but every k steps, we create fork of the environment and begin rolling out a DARP agent in that clone of the environment. The results (seen in Fig. 9 show that, even as BC approaches failure and drifts away from the support of expert demonstrations, DARP is able to recover and score very highly. This suggests that DARP indeed has superior robustness to accumulation of enror and can perform well in the slightly -itt---distribution states that failing BC agent drifts into.
20

---

Published as a conference paper at ICLR 2026
Forked DARP Scores (Stack)
Forked DARP Scores (Hopper)
140
3600
m
BC DARP
Score
Score
BC
DARP
1000
Step
Step
190
400
Figure In two different tasks (the Robosuite Stack task and the MuJoCo Hopper task), we rollout BC agent and create fork of the environment every k steps (in this case, k = 10). Observe that, even as BC nears the end its failing rollout, DARP able to scale highly, and is only prevented from doing so about halfway through the Stack rollout and about 80% through the Hopper rollout.
Environment
Method
Improvement over BC (%)
Hopper
CCIL
32.6%
DARP
53.2% %
Walker
CCIL
84.1%
DARP
84.1 %
Ant
CCIL
12.6%
DARP
84.5 %
HalfCheetah
CCIL
5.4%
DARP
418.7%
Table 7: Relative improvement over BC compared to CCIL. Comparing DARP against CCIL standard MuloCo benchmarks. DARP achieves higher or equal relative gains across all tasks.
A.2.5 COMPARISON WITH CCIL.
As shown in Table 7 we evaluate our method against CCIL, a baseline that explicitly induces smoothness. We use the reported scores in the CCIL paper, and compare percent improvement over BC. Observe that DARP outperforms CCIL significantly on three out of four environments, with particularly large margin on HalfCheetah (418.77 8.7% vs 5.4% improvement).
A.2.6 DISTANCE METRIC SENSITIVITY
Method
Success Rate (%
BC
47
DARP (Cosine Similarity)
DARP (Euclidean Distance)
75
Table 8: Distance Metric Sensitivity. Comparison success rates using R3M features. DARP with Euclidean distance and cosine similarity perform similarly, beating BC by 28% and 23% respectively. Results are on the Robosuite Stack task.
While all experiments performed in this paper use Euclidean distance choose nearby neighbors, is natural to consider alternative metrics, like cosine similarity, especially in high-dimensional embeddings such as R3M. We find (as shown in Table that DARP performs similarly only 5% worse when using cosine similarity rather than Euclidean distance. This suggests DARP robust to the choice of distance metric used.
21

---

Published as a conference paper at ICLR 2026
A.2.7 DARP IN COMBINATION WITH DIFFUSION POLICY
Method
Stack
Threading
Peg Insertion
BC (MLP)
34
12
28
DARP (MLP)
62
38
40
Diffusion
52
44
54
DARP w/ Diffusion
82
72
Table 9: DARP in combination with diffusion policy. DARP provides significant improvements when applied to both standard MLP policies and Diffusion policies. Results are on the Robosuite Stack, Threading, and Peg Insertion tasks as success rate across 50 trials. States state-based. Note that the number of demonstrations used is much less than that used in Table so baseline BC and DARP numbers do match.
While all reported experiments are performed with an MLP backbone, diffusion policy (Chi et al. 2024) has proven to be state-of-fe the-ttr-rt model class for imitation learning, particularly for manipula- tion tasks. Table[rrevee
that DARP not mutually exclusive with diffusion and can be combined for even more performant model. Using the DARP architecture with diffusion backbone outper- forms DARP with an MLP backbone by 20%, 34%, and 26% success rate increase for all three tasks respectively, beating standard BC by a total 48%, 60%, and 38% for the three tasks.
Method
Training (s/epoch)
Inference (s/step)
Diffusion
5.610
0.15700
DARP (MLP)
0.559
0.00437
Table 10: Computational efficiency of DARP with an MLP backbone and diffusion policy. Comparison of runtime costs between DARP (MLP) and diffusion policy. Diffusion policy is significantly more expensive, being R 10x slower in training and x 36x slower during inference.
Additionally, we empirically find that DARP with an MLP backbone much faster than standard diffusion, particularly in inference see Table 10
A.2.8 DARP TRAINED ON HUMAN DEMONSTRATIONS
Method
Success Rate (%)
BC
45
DARP
60
Table 11: Results with human demonstrations. Comparison of success rates when training human data rather than data collected by RL policies. Results are on the Robosuite Stack task.
The expert demonstrations used train models for the MuJoCo and Robosuite environments are collected by an optimal Reinforcement Learning policy. It is crucial to ensure DARP maintains a performance gain in comparison to BC when trained expert demonstrations collected by humans. Indced, when trained on human demonstrations on the Robosuite Stack task, DARP outperforms standard BC by 15% (see Table 11)
22

---

Published as a conference paper at ICLR 2026
A.2.9
CHOICE OF RETRIEVAL HYPERPARAMETERS
Environment
Method
Score (Mean)
Hopper
BC
507.49
DARP
805.69
Stack
BC
0.12
DARP
0.31
Table 12: Performance comparison when validation loss is used to select training epochs and retrieval hyperparameters. Mean scores on Hopper and Stack environments. Observe that DARP maintains performance gain in comparison to BC.
Figure revealss that there are selections of retrieval parameters (for example, very low number of neighbors) which cause DARP to perform worse than standard BC. However, we find that choosing retrieval hyperparameters that minimize validation loss is an effective strategy to find performant settings, see Table I2
A.3 PSEUDOCODE
We provide pseudocode of the DARP algorithm, Algorithm
Algorithm Difference-Aware Retrieval Policies
1: Input: Expert demonstrations 2 {($),47)) number of neighbors k
2: Initialize: / parameters 0
3: if is parametric then
4: Initialize: g parameters y
5: end if
6: // Training Loop
7: while not converged do
8: Sample batch of query data (sg) "
9:
for each query pair (gg:ag) in batch
10:
// Find K-Nearest Neighbors from the entire dataset ø* M(ssg) & arg -kgmink)(((s.s5))
11:
12:
// Compute Neighbor-based Predictions
13:
for each neighbor index i (G(s))))) do 5q)
14:
d f fo (s) end for
15:
// Aggregate Predictions if g is parametric then àg + 8v(ldhe.n)))))
17:
else
20:
àg + 8(fa)tie.((on))
21:
end if
22:
end for
23:
// Update Parameters based on the batch loss
L + Elyyay 1aallll,,0-0 agll
25:
// Gradient descent step
26:
0 0 -Vo
27:
if is parametric then
28:
+ v aV L
29:
end if
30:
end while
31: Output: Trained parameters 0 and, if applicable, w
23

---

Published as a conference paper at ICLR 2026
A.4
RUNTIME ANALYSIS
Environment
k
Train (s/epoch)
Test (s/step)
Hopper (1,000 datapoints)
BC
0.102
0.00127
100
0.126
0.00413
200
0.128
0.00418
300
0.130
0.00420
400
0.131
0.00421
500
0.130
0.00419
Stack (4,200 datapoints)
BC
0.431
0.00139
100
0,556
0.00437
250
0.539
0,00442
500
0.559
0.00437
750
0.576
0.00433
1000
0,591
0.00478
1500
0.682
0.00452
2000
0.758
0.00451
Table 13: Runtime comparison across environments. Training and testing speeds (in seconds) for Behavior
Cloning (BC) and varying values of k on Hopper and Stack datasets.
analyze the computational overhead of DARP at both training- time and inference-time. Our analysis indicates that computation cost scales sub-linearly as increases, and that inference time is tractable for real-time robotic application. For example, the Robosuite Stack task at k = 500, inference takes approximately 0.00437 seconds per step on our hardware, corresponding to control frequency higher than 230 Hz.
As shown in Table 13
24


Reviews:
[{"review_id": "PhZeie3tPj", "reviewer": "Reviewer_T2hb", "summary": "This paper introduces DARP, a retrieval-based imitation learning method that addresses covariate shift in behavior cloning by reparameterizing the problem using local neighborhood structure. The key innovation is conditioning action predictions on k-nearest neighbors along with their actions and crucially, the difference vectors between neighbor states and query states. The authors provide theoretical analysis showing DARP implicitly performs Laplacian smoothing and demonstrate 15-46% performance improvements across continuous control and robotic manipulation tasks.", "questions": "- What is the time complexity for training and inference? How does this scale with dataset size and k?\n- How sensitive is DARP to the choice of distance metric, especially in high-dimensional spaces? Have you experimented with learned metrics?\n-  Does the Laplacian smoothing analysis extend to the full DARP algorithm with difference vectors and parametric aggregation functions?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The use of difference vectors (s*_i - s_q) rather than just neighbor states is creative and well-motivated. The ablation in Figure 5 convincingly shows this is crucial for performance.\n- The connection to Laplacian smoothing provides intuition for why the method works and bridges local and global learning paradigms effectively.\n- The experiments span MuJoCo locomotion, robotic manipulation, visual observations, and even deliberately discontinuous environments, showing broad applicability.", "weaknesses": "-  The paper doesn't discuss the computational cost of k-NN retrieval at every training and inference step. For k=500 (as suggested in Figure 8), this could be prohibitive for large datasets or real-time applications.\n- While the paper compares to BC and some retrieval methods, it lacks comparison with other smoothness-inducing approaches mentioned in related work (L2C2, CCIL). The MRIL baseline helps but isn't a published method.\n- The jump from iMRIL (which uses simple averaging) to DARP (with difference vectors and learned aggregation) isn't theoretically justified. Does the Laplacian smoothing interpretation still hold?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7oBzaxRwis", "reviewer": "Reviewer_vWKv", "summary": "This paper presents DARP, a method to reduce variances in behavior cloning methods. It conditions the action generation using the retrieved nearest neighbors and aggregates the predictions from individual neighbors to output the final action. The experiments show the improvements over prior methods based on retrieved neighbors and BC on MuJoCo and robosuite tasks.", "questions": "- The step of divergence analysis can be considered as slight evidence that DARP handles OOD better, but how often can we detect the divergence? Also, it seems like this analysis is only possible with deterministic policies. How can we apply the same analysis to the push-T task, where DARP uses a GMM head?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper bridges the non-parametric retrieval-based method and the parametric BC to improve the policy performance.\n- The paper provides good analysis, both theoretical and experimental, to show the effect of DARP.", "weaknesses": "- The main experiment result is based on MuJoCo tasks, which have a lot more neighbors than manipulation, so it is easier to show the improvements. For manipulation tasks, it is unclear which BC architecture it is employing. If it is an MLP, it is not a well-performing architecture for manipulation; it will need to compare with other BC architectures, like diffusion policy, to show the real improved performance.\n- As shown in Appendix A2.1, the performance is significantly affected by the distance measure, e.g.,  the amount of look-back. This can show that the retrieval mechanism outweighs the architectural changes proposed in the paper. \n- The motivation of the paper is about improving generalization. However, none of the experiments evaluate the generalization performance of DARP.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "h6NK4UcfDj", "reviewer": "Reviewer_kF51", "summary": "This paper proposes DARP, a new semi-parametric BC method that integrates non-parametric retrieval with parametric prediction. DARP retrieves a set of nearest-neighbor expert demonstrations for a given query state and conditions the policy’s action prediction on both the neighbor states and their difference vectors relative to the query. The retrieved neighbors are aggregated in a permutation-invariant manner, enabling the model to enforce implicit local smoothness without explicit regularization. The authors show theoretically that DARP approximates a Laplacian smoothing operation over the expert k-NN graph, promoting manifold-consistent policies. Empirical results demonstrate the improved stability, generalization, and robustness across various continuous-control tasks.", "questions": "1.  Have the authors analyzed the trade-off between inference latency and performance?\n\n1. In the proof of Theorem 1 (iii), what does $r$ represent? There seems a loss of definition.\n\n1. Can the method be extended to other distance metrics?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The core idea of reparameterizing the policy in terms of relative differences to known expert state-actions is highly intuitive. It anchors the policy's predictions to the ground-truth data manifold, providing a strong and data-centric form of 'regularization'.\n\n1. The authors provide a strong theoretical motivation for their architectural choice. The connection drawn between neighbor aggregation (iMRIL) and explicit manifold regularization (MRIL) via spectral graph theory (Theorem 2) offers a compelling, first-principles explanation for why this method should reduce variance and improve stability, linking it to low-pass filtering on the graph Laplacian.\n\n1. The experimental results are good, and the ablation is impressive and sufficient.", "weaknesses": "Overall, the proposed method is inspiring and technically sound, and the paper is easy to read. I only have concerns on the inference cost and the distance metric:\n\n1. The method could suffer from large computational/storage overhead at inference time. For every single decision step, the policy must perform a k-NN search over the entire $N$-point expert dataset and then perform $k$ network forward passes. Also, the method requires storing a large expert dataset or its embeddings, which increases memory and storage requirements by orders of magnitude compared to a standalone policy network. \n\n1. The paper assumes that Euclidean distance (or a simple temporal extension with decay) suffices to identify meaningful neighbors, which strong in high-dimensional or structured state spaces, where Euclidean proximity does not necessarily imply functional similarity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nr09YNzrXm", "reviewer": "Reviewer_iFBA", "summary": "Authors propose a way to combine retrieval based IL with standard BC that will turns out to be equivalent to manifold smoothed BC. Authors add both theoretical analysis and experimental validation in terms of standard MuJoCo IL experiments and ablation studies.", "questions": "- How would you contrast the proposed method to the retrieval IL method presented in: Federico Malato, Ville Hautamäki, \"Online Adaptation for Enhancing Imitation Learning Policies\", CoG 2024?\n- In Section 2.4.2 you show an interesting way to go beyond simple averaging, but in experiments I see no improvement in using it. Why is that? I would suggest not to bold DARP Set Transformer numbers in Table 1 as all of those numbers are worse than plain DARP.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- Interesting, theoretically motivated and practical offline IL method. \n- MuJoCo and robosuite experimental results are quite good. \n- Narrative is well written and easy to follow.", "weaknesses": "- Even though narrative is well written, the mathematical exposition leaves room for improvement. For example looking at page 3, we see a lot of symbols that have not been defined: s^*, \\pi(s_t), L, ... \n- Even though I understand that focus is in offline IL, it would be good obtain results from IL methods that have access to environment, such as adversarial methods (GAIL, AIRL etc). Can the proposed method bridge the gap between BC and GAIL? \n- It appears that expert demos have been obtained from an optimal RL policy, or have I been mistaken? In the \"What Matters for Adversarial Imitation Learning?\" -paper authors note that there is a difference when using expert demos from optimal policy and expert demos from human expert. It would be good to test with such demos, one option is to use ALE environment and with Atari-HEAD expert demo dataset.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)