Difference-Aware Retrieval Policies for Imitation Learning

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

Difference-Aware Retrieval Policies for Imitation Learning

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

Published as a conference paper at ICLR 2026

D IFFERENCE -AWARE R ETRIEVAL P OLICIES FOR I MITATION L EARNING Quinn Pfeifer1 , Ethan Pronovost1 , Paarth Shah2 , Khimya Khetarpal3,4 , Siddhartha Srinivasa1 , Abhishek Gupta1,2 1 Paul G. Allen School of Computer Science & Engineering, University of Washington 2 Toyota Research Institute 3 Google DeepMind

arXiv:2606.09758v1 [cs.RO] 8 Jun 2026

4 Mila

A BSTRACT 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 taskspecific 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. Code and demos are available at https://weirdlabuw.github.io/darp-site/.

I NTRODUCTION Imitation Learning

Difference-Aware Retrieval Policies

…

Figure 1: Overview of DARP: Unlike standard BC (left), DARP (right) utilizes a retrieval-based reparameterization 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 at inference-time, while DARP utilizes it to perform retrieval to find a local neighborhood of expert state-action pairs around each query point sq .

Published as a conference paper at ICLR 2026

Imitation learning via behavior cloning (BC) (Pomerleau, 1991) has enabled robots to learn complex, dexterous behaviors from expert demonstrations (Zhao et al., 2023; Chi et 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 errors accumulate during rollouts, driving the agent into states not well represented in the demonstration data (Spencer et al., 2021; Ross et al., 2011). In such out-of-distribution regions, BC policies are highly unstable, producing unreliable 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., 2015; Ke et 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 smoothness, 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., 2023; Chi et al., 2024) attempt to compress the entire demonstration dataset into a single parametric function, which is typically brittle under distribution shift. “Local” methods (Pari et al., 2022; Mansimov & Cho, 2018; Salzberg & 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-Aware 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 a feedforward pass on a parametric function, DARP (Fig. 1) first retrieves a set of k neighbors from the training corpus and then predicts k actions conditioned on tuples of (neighbor state, associated action, and difference from the query state). These neighbor-informed predictions are then aggregated in a permutationinvariant manner to produce a 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 on the retrieved neighbors). We show that doing so reduces variance without requiring any additional assumptions beyond those made for standard Behavior Cloning – we need no additional data, online supervision, or 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 a parameter-free 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 a fixed low-pass spectral filter that suppresses high-frequency 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

Published as a conference paper at ICLR 2026

DARP is a general, scalable architecture that naturally extends to 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-based algorithms for sequential decision-making problems.

D IFFERENCE -AWARE R ETRIEVAL P OLICIES FOR I MITATION L EARNING

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. In this section, we thoroughly derive theoretical guarantees from first-principles. Readers primarily interested in the practical experimental results may skip to Section 2.5 for a self-contained summary of DARP before proceeding to Section 3 for empirical results. As a warmup, we define the problem setting (Section 2.1) and discuss a variant of regularized imitation learning (Section 2.2) that imposes additional structure from the data for improvements in variance, generalization, and stability. In Section 2.3, we 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

P RELIMINARIES : B EHAVIOR C LONING FOR I MITATION L EARNING

We operate in the typical imitation learning setting, formalized by a finite-horizon Markov Decision Process (MDP), M = {S , A , P0 }, where S is the state space, A is the action space, and P0 is the initial state distribution. A policy maps a state to a distribution of actions fθ : S → ∆A so as to maximize task-relevant objectives (for brevity, we drop the subscript θ when discussing the general function class f ). We assume access to expert human-provided demonstrations D ∗ as a collection of state-action pairs: D ∗ = {(s∗j , a∗j )}. We use the notation s∗ and a∗ specifically to denote states and actions belonging to the expert dataset. The behavior cloning (Pomerleau, 1991) algorithm learns a policy fθ from this dataset by casting imitation as a typical supervised learning problem — arg maxθ E(s∗ ,a∗ )∼D ∗ [log( fθ (a∗ | s∗ ))]. While the distribution class of fθ 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 simplicity1 . 2.2

WARM - UP : N EIGHBOR M ANIFOLD R EGULARIZED I MITATION L EARNING

While behavior cloning minimizes only the supervised imitation loss over expert states drawn from D ∗ , 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-distribution regions where a BC policy may act arbitrarily, especially for overparameterized neural networks (Ross & Bagnell, 2010). 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 a 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-regularized loss (LMRIL ), where the standard imitation learning objective (LBC ) is combined with an additional smoothness penalty (LS ) enforcing predictions to respect the geometry of the dataset rather than relying solely on pointwise supervision. 1We show that this can be relaxed in Section 2.4

Published as a conference paper at ICLR 2026

" LMRIL ( f ) = E(s∗ ,a∗ )∼D ∗ ℓ f (s∗ ), a∗ +λ Es∗ ∼D ∗ | {z } | supervised risk(LBC )

∑

∗

wi (s )

i∈Nk (s∗ )

2 f (s ) − f (s∗i ) 2 ∗

# ,

(1)

}

smoothness regularizer(LS )

where ℓ( f (s∗ ), a∗ ) is the supervised imitation loss, Nk (s∗ ) are the k-nearest neighbors of s∗ from the expert dataset, and weights wi (s∗ ) are normalized kernel weights based on the state differences the ∗ ∗ ∥s −s ∥ — wi (s∗ ) ∝ K∆ i h . As we discuss briefly below (and in detail in Appendix A.1.1), this corresponds to a 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 (LMRIL ) improves over vanilla BC (LBC )). Let f ∗ : S → A be the true, underlying expert policy, assumed to be C2 -smooth on a compact state space S . Let f : S → A denote the learned policy estimator. Consider two estimators trained on expert demonstrations: 1. Vanilla BC: a global supervised model minimizing LBC ( f ) = E(s∗ ,a∗ )∼D ∗ [ℓ( f (s∗ ), a∗ )]. 2. MRIL: a neighbor-based estimator minimizing " LMRIL ( f ) = LBC ( f ) + λ Es∗ ∼D ∗

∑

i∈Nk (s∗ )

wi (s∗ ) f (s∗ ) − f (s∗i )

2 2

# ,

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. (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: In a closed-loop rollout, the deviation recursion ∆t+1 ≤ Ls ∆t + La ∥ f (st ) − f ∗ (st∗ )∥2

accumulates error linearly for vanilla BC, but sublinearly for MRIL, since the smoothness regularizer enforces ∥ f (s∗ ) − f (s′∗ )∥ = O(∥s∗ − s′∗ ∥) 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.1, but provide a brief sketch. The key idea of the proof is to first show that the smoothness regularizer directly corresponds to a graph Laplacian penalty on a graph constructed by a k-nearest neighbor (k-NN) affinity matrix defined by the kernel wi . 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 is 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, it shrinks the local Lipschitz constant of f along directions where the data density p(s) is high, ensuring that small changes in state lead to small, consistent changes in the predicted action. As a result, the policy generalizes more reliably on in-distribution (ID) states and extrapolates in a structured manner on new out-of-distribution (OOD) states in the neighborhood. 2 L and L are the Lipschitz constants of the environment transition dynamics with respect to state and action, s a respectively.

Published as a conference paper at ICLR 2026

2.3

I MPLICIT M ANIFOLD R EGULARIZATION VIA I N -C ONTEXT A RCHITECTURES

While our MRIL objective does amortize local learning to provide improvements over vanilla BC, there are two notable drawbacks. First, it requires a hyperparameter λ that must be tuned to balance supervised accuracy and smoothness. Second, the requirement to optimize a modified, regularized objective rather than a standard BC objective may modify the optimization landscape in adverse ways. This raises a natural question: can we obtain the same benefits conferred by MRIL (Eq. 1), by modifying the policy architecture rather than modifying the objective? In this section, we introduce a retrieval-based change in policy architecture that leads to an implicit manifold regularization effect (iMRIL), despite using a standard imitation objective. With iMRIL, we can obtain the benefits of Laplacian smoothing (from MRIL) by training on a standard BC objective (as shown in Fig. 2), without introducing λ as an additional hyperparameter for training. We then build on this algorithm to develop a practical instantiation of this method (DARP) in Section 2.4. iMRIL architecture: The high-level idea behind iMRIL is simple – we propose moving the neighborhood aggregation (averaging) operation from the objective (as in Eq. 1) to the architecture itself. So instead of learning a standard feedforward predictor f (s) that is trained BC iMRIL against a neighborhood regularized smoothness objective (Eq. 1), we propose embedding the Figure 2: iMRIL implicitly achieves Laplacian structure of neighborhood aggregation directly smoothing, which reduces variance and enforces local into the parameterization of the action predic- consistency, whereas the lack of smoothness constraint tor fˆ itself, while maintaining the objective as on standard BC allows for arbitrarily jagged function standard imitation learning. iMRIL learns the approximations. parameters of a per-state predictor fθ such that an action predictor explicitly parameterized via neighborhood-aggregation fˆ(s∗ ) = 1k ∑i∈Nk (s∗ ) fθ (s∗i ) across nearest neighbor states from the training set {s∗i }i∈Nk (s∗ ) generates accurate predictions of the corresponding expert action a∗ . With this parameterization, iMRIL optimizes at training time:

arg min E(s∗ ,a∗ )∼D ∗ θ

! 1 ∗ ∗ fθ (si ) −a k i∈N∑(s∗ ) 2 k | {z }

(2)

fˆ(s∗ )

At deployment time, inference can be performed on a new state sq simply by retrieving the k-NN of sq from the training set and performing neighborhood aggregation â = fˆ(sq ) = 1k ∑i∈Nk (sq ) fθ (s∗i ). As we show in Section 2.4, the particular parameterization of f is of crucial importance and plays a significant role in the empirical performance of iMRIL— leading to the development of DARP. Intuitively, we are parameterizing the action predictor fˆ as an aggregation of predictions at neighbor states from the training data f (s∗i ), and then learning f . Supervising the post-aggregation function implicitly prevents any f predictions from being arbitrarily non-smooth, conferring the benefits noted in 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.1, we state our main result and a 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 in Section 2.2), with eigenpairs {(µ j , u j )}nj=1 , where 0 = µ1 ≤ µ2 ≤ · · · ≤ µn ≤ 2.

The minimizers of the explicit MRIL objective (Section 2.2) and the implicit iMRIL objective (Section 2.3) have the following closed form expansions 5

Published as a conference paper at ICLR 2026

1 ⟨a∗ , u j ⟩ u j 1 + λ µ j j=1

fˆiMRIL = ∑ (1 − µ j ) ⟨ f , u j ⟩ u j

fMRIL = ∑

j=1

iMRIL ’s neighbor aggregation step applies the fixed spectral filter φiMRIL (µ) = 1 − µ to the graph Laplacian L, preserving low-frequency modes and suppressing high-frequency modes. The congruence between fˆiMRIL and fMRIL shows that iMRIL is equivalent to a built-in form of Laplacian 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. The explicit regularizer admits a spectral solution by diagonalizing the k-NN Laplacian, yielding a filter of the form (1+λ µ)−1 on each eigenmode. The implicit objective can be expressed as neighbor aggregation fˆ = S f with S = D−1 A, the random-walk matrix, which has the same eigenvectors and applies the fixed filter 1 − µ. Intuitively, both act as low-pass filters on the graph: modes with small eigenvalues (smooth variation across the data manifold) are largely preserved, while modes with large eigenvalues (rapid, high-variance fluctuations between neighbors) are strongly damped. Thus iMRIL implicitly performs Laplacian smoothing, reducing variance and enforcing local consistency without needing to tune λ . Note that the implicit Laplacian smoothing view does not replace the need to learn a policy; rather, it constrains the class of functions that can be represented after aggregation. The neighbor-conditioned network fθ 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

D IFFERENCE -AWARE R ETRIEVAL P OLICIES : A P RACTICAL I NSTANTIATION OF I MRIL FOR I MITATION L EARNING

Given the conceptual framework of iMRIL, we instantiate a practical algorithm for large-scale imitation learning. We build on the objective outlined in Eq. 2 and instantiate a careful choice of (1) parameterization, (2) neighbor aggregation that leads to strong empirical performance. 2.4.1

D IFFERENCE - BASED PARAMETERIZATION OF fθ

The objective described in Eq. 2 leaves the parameterization and input representations of fθ open to broad interpretation. We make the observation that the neighborhood aggregation should learn how expert actions vary under local perturbations. This suggests that fθ should use knowledge of differences between a query state and a neighbor state to adaptively propose locally adapted actions for each neighbor. In Difference-Aware Retrieval Policies (DARP), instead of simply parameterizing fθ by fθ (s∗i ), we provide additional context about the optimal neighbor action a∗i , as well as the difference between the query state and the neighbor state ∆si = s∗i − sq ; a predictor fθ predicts an action candidate a′i for a query state sq and a neighbor (s∗i , a∗i ) using the difference information as a′i = fθ (s∗i , a∗i , ∆si = s∗i − sq ).

Let Nk (sq ) be the index set of the k-nearest neighbors retrieved according to some distance function d(sq , s∗i ).3 For generating predictions with DARP, we can then perform neighborhood aggregation (as outlined in Section 2.3) to predict an action for any query state sq âq = fDARP (sq ) =

1 a′ k i∈N∑(s ) i k

1 fθ (s∗i , a∗i , ∆si = s∗i − sq ). k i∈N∑(s ) k

(3)

(4)

3 In our work we use the Euclidean distance in a pre-trained embedding space, although other neighborhood functions are also applicable. We refer the reader to Appendix Section A.2.1 for a thorough discussion of design decisions in constructing neighborhood sets via retrieval.

Published as a conference paper at ICLR 2026

At training time, this can be used to define a straightforward imitation learning objective from the expert dataset D ∗ : h i fDARP (s∗q ) − a∗q

arg min E(s∗q ,a∗q )∼D ∗ θ

(5)

where we optimize for the parameters of the predictor fθ , minimizing the discrepancy between the predicted action âq and optimal action a∗q . Given the simplicity of the objective, any parameterization can be used for fθ , in our case, standard feedforward or convolutional neural networks. As we show in Section 3, this difference-based parameterization is crucial for performance. At inference time, we generate actions to execute by retrieving k-NN and performing inference through the neighborhood aggregation operation defined in Eq. 3. 2.4.2

G OING B EYOND L INEAR AGGREGATION

While the process of neighborhood aggregation thus far has been restricted to averaging over neighborhood predictions âq = 1k ∑i∈Nk (sq ) a′i , this is a special case of a broader class of permutation-invariant aggregation functions gψ ({a′i }i∈Nk (sq ) ). For instance, gψ could be parameterized with more expressive set-compliant neural models like the set transformer (Lee et al., 2019) or DeepSets (Zaheer et al., 2017). This suggests a generalization of the prediction model in Eq. 3 as âq = gψ ({ fθ (s∗i , a∗i , ∆si = s∗i − sq )}i∈Nk (sq ) ). Besides benefits in expressivity, generalizing from a simple averaging operation to a parametric aggregation model gψ 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 to the L2 -regression objective defined in Eq. 5. Rather than predicting âq directly, DARP can predict the parameters α of an action distribution p(aq ; α) – for instance the means, covariances, and weights for a 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 L2 -regression: arg max E(s∗q ,a∗q )∼D ∗ log p(a∗q ; αθ (s∗q )) , where αθ (s∗q ) = gψ { fθ (s∗i , a∗i , ∆si = s∗i − s∗q )}i∈Nk (s∗q ) θ

(6)

Inference for a query state sq can be performed by sampling âq ∼ p(· ; αθ (sq )), constructing αθ (sq ) = ∗ ∗ ∗ gψ { fθ (si , ai , ∆si = si − sq )}i∈Nk (sq ) from a set of neighbors retrieved at test time. We refer readers to Appendix Section A.3 for detailed training pseudocode. 2.5

S UMMARY

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 over standard BC (Theorem 1). Rather than adding an 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, we propose the following algorithm: at training time, for each expert state-action pair (s∗q , a∗q ), DARP retrieves the k-nearest neighbor states from D ∗ , computes difference vectors ∆si = s∗i − s∗q from each neighbor to the query state, passes each neighbor tuple (s∗i , a∗i , ∆si ) through a network fθ to produce candidate actions, and finally aggregates these candidates via a permutation-invariant function gψ to predict the final action. This is trained only with the standard imitation learning objective.

E XPERIMENTAL E VALUATION

Next, we evaluate DARP in order to 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 to DARP’s performance gains? We conduct experiments across multiple domains using low-dimensional state representations, high-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. 7

Published as a conference paper at ICLR 2026

3.1

BASELINE C OMPARISONS AND TASK D ESCRIPTIONS

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 on a flat plane. These tasks include: Hopper (single-legged hopping robot), Walker (bipedal humanoid), Ant (quadruped), and HalfCheetah (biped). Robosuite Tasks: The Robosuite (Zhu et al., 2020) tasks all entail a 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 a thin, needle-like tool and insert it into a 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 a 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 a variety of baselines and ablations: (1) R&P (Sridhar et 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 on retrieved neighbors, (3) BC: refers to standard parametric behavior cloning, (4) REGENT (Sridhar et al., 2025): refers to a transformer-based in-context learning method conditioned on retrieved neighbors, (5) MRIL: refers to the explicitly smoothed version of DARP outlined in Section 2.2.

3.2

C AN DARP CONSISTENTLY OUTPERFORM STANDARD BEHAVIOR CLONING ? (Q1)

In this experiment, we evaluate DARP’s core hypothesis on tasks with low-dimensional state representations, where the distance metrics between states are well-defined and interpretable. This evaluation spans locomotion tasks from MuJoCo and robotic manipulation tasks from Robosuite and RoboCasa with data generated with MimicGen (Mandlekar et al., 2023). In these experiments, the aggregation function g is implemented as a simple average of all neighbor action predictions a′ .

Method R&P (Sridhar et al., 2025) LWR (Pari et al., 2022) BC REGENT (Sridhar et al., 2025) MRIL DARP DARP Set Transformer

Hopper

Ant

Walker

HalfCheetah

711.82 ± 85.63 1703.78 ± 245.95 2313.65 ± 203.75 1819.39 ± 186.24 2793.63 ± 156.41 3545.57 ± 3.54 2965.86 ± 103.08

-305.97 ± 76.42 846.59 ± 216.06 2376.20 ± 339.43 -302.10 ± 146.67 3869.08 ± 241.00 4383.28 ± 266.37 4063.79 ± 218.80

419.18 ± 50.21 1484.91 ± 356.54 2658.40 ± 274.08 507.01 ± 76.10 4370.96 ± 168.13 4894.01 ± 75.12 4752.42 ± 109.23

-178.64 ± 29.75 1945.82 ± 567.26 1063.23 ± 371.08 169.85 ± 63.10 701.58 ± 195.08 5515.41 ± 841.33 3417.85 ± 764.57

Table 1: Both DARP and DARP Set Transformer outperform other approaches across all domains. Performance Comparison of DARP vs. BC and other baselines across MuJoCo Environments Using LowDimensional State. Scores reported are averaged across 100 independent trials with 95% confidence intervals. The parametric policies utilized Multi-Layer Perceptrons, for results with diffusion policies, see Section 8.

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 1 and Table 2). 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 a 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. 8

Published as a conference paper at ICLR 2026

Figure 3: The “DrawerClose” Figure 4: The “Threading” Ro- Figure 5: Real-world FurniRoboCasa task in which the bosuite task in which the robotic tureBench square table assembly robotic arm is tasked with clos- arm is tasked with inserting the task. The robot is tasked with ing the open drawer. needle implement into the small picking up a table leg and screwhole. ing it into a hole in the corner of the tabletop. RoboCasa

Robosuite

Real

Method

Drawer

Door

Stove

Stack

Thrd.

Peg

Sq. Table

BC DARP

54 85

29 45

28 43

47 72

37 63

46 62

44 92

Table 2: Comparing DARP against BC across all three evaluation environments using low-dimensional state

features. Scores are success percentages (RoboCasa & Robosuite: 100 trials; Real: 50 trials). DARP consistently outperforms BC across all tasks and environments, more than doubling the score of BC in the real world. RoboCasa and Robosuite policies utilized Multi-Layer Perceptrons, while the real-world results used diffusion policy, see Section 8.

3.3

C AN DARP HANDLE MORE COMPLEX STATE REPRESENTATION AND ACTION DISTRIBUTIONS ? (Q2)

High-Dimensional Visual Input Representations. To test the applicability of DARP beyond the regime of Method Stack Thrd. Peg compact, low-dimensional states, we evaluate DARP on simulated robotic manipulation tasks using R3M image BC 44 38 17 embeddings (Nair et al., 2022). This tests whether the DARP 75 76 52 neighbor-based approach remains effective when states are represented as high-dimensional feature vectors ex- Table 3: Success rates (%) on vision-based tracted from visual observations (see Table 3). Observe Robosuite tasks. that, not only does DARP outperform standard BC, the average improvement, ∼ 35%, is actually higher than the average improvement on Robosuite tasks in low-dimensional state (∼ 22%). Empirically, this means that DARP was better at adapting to complex, high-dimensional state representations than standard BC. 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

H OW DO DIFFERENT ARCHITECTURAL COMPONENTS CONTRIBUTE TO 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 (1) standard DARP; (2) DARP, but without including the neighbor actions; (3) an ensemble of 10 BC agents; (4) DARP, but we choose random neighbors as opposed to using a distance metric; (5) DARP, but we take the L2 norm of the distance vector; (6) BC baseline, which is just the query state 9

Published as a conference paper at ICLR 2026

Model Architecture Ablation

0.85

0.72

Success Rate

0.67 0.59

0.58

0.53

0.47

0.42 0.19

0.00 (s∗i , a∗i , s∗i − sq ) (s∗i , s∗i − sq )

BC Ensem.

Rando k (s∗i , a∗i , ||s∗i − sq ||2 )

(s∗i , a∗i , sq )

Perm. Dep.

Figure 6: Distance vectors and permutation invariance contribute heavily to DARP’s success. Exploration of how the performance of a DARP agent is impacted as various changes are made to the core architecture demonstrates that DARP success is most attributed to the distance vectors (s∗i , a∗i , s∗i − sq ). Success rate is averaged across 100 trials on the Robosuite Stack environment with 95% confidence intervals.

sq ; (7) DARP, but include just the query state rather than the distance vector between the query state and neighbor states; (8) DARP, but using a permutation-dependent (so not permutation-invariant) aggregator to combine all a′ s. We report in Figure 6 the results of this systematic ablation. The ablation study reveals that distance vectors and permutation invariance are crucial for DARP’s success, while neighbor actions have a more modest impact. Random neighbor selection performs poorly, confirming that meaningful distance metrics informing neighbor selection are crucial. The permutation-invariant aggregation function g proves critical, as permutation-dependent alternatives significantly degrade performance.

State Likelihood

Accum. Reward

BC DARP SoD

State Likelihood vs. Time

Delta Likelihood vs. Time Delta Likelihood

Cumulative Reward vs. Time

Step

200

Step

200

Step

200

Figure 7: 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 < τs (OOD), but the delta likelihood is > τ∆ (in distribution).

Divergence Analysis: To better understand DARP’s success over standard BC, we analyze the point of divergence in rollouts in 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 a significantly different reward. We define τs and τ∆ as the 1st percentile of likelihoods of the training set (that is, 1% of the deltas seen at training time are less likely than τ∆ ). In all six different rollouts across two different tasks (the Robosuite Stack task and the MuJoCo Hopper task), the query state at the SoD has a state likelihood of < τs but a delta likelihood of > τ∆ . This result bolsters our hypothesis that DARP gains occur partly due to improved prediction on slightly out-of-distribution states due to reparameterization in terms of difference vectors to neighbors. (see Figure 7 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. 10

Published as a conference paper at ICLR 2026

R ELATED W ORK

Non-Parametric Imitation Learning Methods: Non-parametric IL algorithms demonstrate surprising 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 et al. (2025) uses mixtures of Gaussian processes to 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. FlowRetrieval Lin et al. (2024), STRAP Memmel et al. (2025), and Behavior Retrieval Du et al. (2023) perform non-parametric retrieval 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 (s∗i , a∗i , s∗i − sq ) tuples and learning a semi-parametric 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. Smoothness-Constrained Policy Learning: Much recent literature has explored explicit smoothness constraints to improve policy stability and robustness. L2C2 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 a similar methodology for model-based RL models with Lipschitz constraints. CCIL Ke et al. (2024a) extends these ideas to generate synthetic corrective labels for imitation learning using a Lipschitz-constrained dynamics model. This has also been scaled up to humanoid controllers Chen et al. (2024) to reduce shakiness on deployment. DARP differs from these methods by enforcing smoothness implicitly through an architecture change while using standard imitation learning objectives. In-Context Learning Methods: Recent work has explored non-parametric retrieval from the perspective of in-context imitation learning. REGENT Sridhar et al. (2025) investigates retrievalaugmented generalization by incorporating retrieved states, actions, and rewards into a causal transformer, while DPT Lee et al. (2023) uses supervised pretraining for transformers to predict actions given query states and in-context datasets, effectively learning how to explore. Other incontext architectures include ICRT Fu et al. (2024), Instant Policy Vosylius & Johns (2025), and KAT Di Palo & Johns (2024). These methods aim to quickly adapt to new tasks and environments, whereas DARP focuses on accomplishing higher performance and stability on standard imitation learning.

C ONCLUSION

We introduced Difference-Aware Retrieval Policies for Imitation Learning (DARP) (DARP), a 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 state-to-action mappings. We prove that our method implicitly achieves Laplacian smoothing. Our experimental evaluation across diverse domains, including continuous control and robotic manipulation, validates three key hypotheses. First, DARP consistently outperforms standard behavior cloning when using low-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-invariant aggregation are crucial components to our algorithm.

R EPRODUCIBILITY S TATEMENT

A link to supplementary source code is provided. This codebase contains all code used to train and evaluate our models. It also contains policy and environment configuration files to generate all results seen in this paper. We provide all data used in MuJoCo 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 a transformer (REGENT, Set Transformer), which are non-deterministic and may differ slightly from reported numbers. 11

Published as a conference paper at ICLR 2026

ACKNOWLEDGMENTS

This work was funded by the Toyota Research Institute under the University 2.0 program, along with grants from the National Science Foundation NRI (#2132848), DARPA RACER (#HR001121-C-0171), and the Office of Naval Research (#N00014-24-S-B001 and #2022-016-01 UW). We gratefully acknowledge gifts from Amazon, Collaborative Robotics, Cruise, the Research Scholarship from the Mary Gates Endowment for Students, and others.

R EFERENCES Kavosh Asadi, Dipendra Misra, and Michael Littman. Lipschitz continuity in model-based reinforcement learning. In International Conference on Machine Learning. PMLR, PMLR, 2018. Mikhail Belkin and Partha Niyogi. Towards a theoretical foundation for laplacian-based manifold methods. Journal of Computer and System Sciences, 74(8):1289–1308, 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. π0 : A vision-languageaction flow model for general robot control, 2024. URL https://arxiv.org/abs/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 lipschitz-constrained policies. arxiv preprint arXiv:2410.11825, 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.1090/cbms/092. Michael Jae-Yoon Chung, Maxwell Forbes, Maya Cakmak, and Rajesh PN Rao. Accelerating imitation learning through crowdsourcing. In 2014 IEEE International Conference on Robotics 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 arXiv:2408.15980, 2024. Liyiming Ke, Yunchu Zhang, Abhay Deshpande, Siddhartha Srinivasa, and Abhishek Gupta. Ccil: Continuity-based data augmentation for corrective imitation learning. In The Twelfth International Conference on Learning Representations, 2024a. URL https://openreview.net/forum? id=LQ6LQ8f4y8. Liyiming Ke, Yunchu Zhang, Abhay Deshpande, Siddhartha S. Srinivasa, and Abhishek Gupta. CCIL: continuity-based data augmentation for corrective imitation learning. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024b. URL https://openreview.net/forum?id=LQ6LQ8f4y8. 12

Published as a conference paper at ICLR 2026

Taisuke Kobayashi. L2c2: Locally lipschitz continuous constraint towards stable and smooth reinforcement learning. In 2022 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 4032–4039. IEEE, IEEE, October 2022. Jonathan Lee et 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, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invariant 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 learning: Tutorial, review, and perspectives on open problems, 2020. URL https://arxiv.org/abs/2005. 01643. Li-Heng Lin, Yuchen Cui, Amber Xie, Tianyu Hua, and Dorsa Sadigh. Flowretrieval: Flowguided 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://proceedings.mlr.press/v270/lin25a.html. Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In International Conference on Learning Representations (ICLR) 2023, 2023. URL https://arxiv.org/abs/2210.02747. Originally appeared as arXiv preprint arXiv:2210.02747. 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. In 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 as a 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 Learning Research, pp. 892–909. PMLR, 2022. URL https://proceedings.mlr.press/v205/nair23a.html. 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 Robotics: Science and Systems (RSS), 2024. 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.2022. XVIII-052. Emmanuel Pignat and Sylvain Calinon. Bayesian gaussian mixture model for robotic policy imitation. IEEE Robotics and Automation Letters, 2019. URL https://arxiv.org/pdf/1904. 10716.pdf. Preprint version. Dean Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural Comput., 3(1):88–97, 1991. doi: 10.1162/NECO.1991.3.1.88. URL https://doi.org/10. 1162/neco.1991.3.1.88. 13

Published as a conference paper at ICLR 2026

Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings of 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 on Artificial Intelligence and Statistics, AISTATS 2011, Fort Lauderdale, USA, April 11-13, 2011, volume 15 of JMLR Proceedings, pp. 627–635. JMLR.org, 2011. URL http://proceedings.mlr. press/v15/ross11a/ross11a.pdf. 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/978-1-4612-2660-4\_33. URL https://link.springer. com/chapter/10.1007/978-1-4612-2660-4_33. Jonathan C. Spencer, Sanjiban Choudhury, Arun Venkatraman, Brian D. Ziebart, and J. Andrew Bagnell. Feedback in imitation learning: The three regimes of covariate shift. CoRR, abs/2102.02872, 2021. URL https://arxiv.org/abs/2102.02872. Kaustubh Sridhar, Souradeep Dutta, Dinesh Jayaraman, and Insup Lee. REGENT: A retrievalaugmented 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 IEEE/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 TwentyNinth AAAI Conference on Artificial Intelligence, January 25-30, 2015, Austin, Texas, USA, pp. 3024–3030. AAAI Press, 2015. doi: 10.1609/AAAI.V29I1.9590. URL https://doi.org/ 10.1609/aaai.v29i1.9590. Jan Ole von Hartz, Adrian Röfer, Joschka Boedecker, and Abhinav Valada. The unreasonable effectiveness of discrete-time gaussian process mixtures for robot policy learning. arXiv preprint arXiv:2505.03296, 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. Garnett (eds.), Advances in Neural Information Processing Systems 30, pp. 3391–3401. Curran Associates, Inc., 2017. URL http://papers.nips.cc/paper/ 6931-deep-sets.pdf. Tony Z. Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual manipulation with low-cost hardware. In Proceedings of Robotics: Science and Systems (RSS), 2023. Dengyong Zhou, Olivier Bousquet, Thomas Navin Lal, Jason Weston, and Bernhard Schölkopf. Learning with local and global consistency. In Advances in Neural Information Processing Systems (NeurIPS), volume 16, 2003. Yuke Zhu, Josiah Wong, Ajay Mandlekar, Roberto Martín-Martín, Abhishek Joshi, Soroush Nasiriany, Yifeng Zhu, and Kevin Lin. robosuite: A modular simulation framework and benchmark for robot learning. In arXiv preprint arXiv:2009.12293, 2020.

Published as a conference paper at ICLR 2026

A PPENDIX

A.1

L EMMAS AND P ROOFS

A.1.1

P ROOF OF T HEOREM 1

To prove Theorem 1, we first start with a well-known result (Chung, 1997; Zhou et al., 2003; Belkin & Niyogi, 2008). Lemma 1 (Smoothness regularizer as k-NN graph Laplacian penalty). Let {s∗1 , . . . , s∗n } be the expert states with corresponding predicted actions f (s∗i ) ∈ Rda . For each i, let Nk (s∗i ) denote the indices of the k-nearest neighbors of s∗i (excluding i). Define asymmetric weights w j (s∗i ), if j ∈ Nk (s∗i ), W̃i j = 0, otherwise, and construct a symmetric affinity matrix Wi j = 12 W̃i j + W̃ ji . Let D be the degree matrix with Dii = ∑ j Wi j , and define the k-NN graph Laplacian L = D −W .

Then the smoothness regularizer can be written as the quadratic form LS ( f ) =

1 n w j (s∗i ) ∥ f (s∗i ) − f (s∗j )∥2 ∝ Tr F ⊤ LF , ∑ ∑ n i=1 j∈N (s∗ ) k

where F = [ f (s∗1 ), f (s∗2 ), . . . , f (s∗n )]⊤ ∈ Rn×da . Equivalently, in the scalar case, LS ( f ) ∝ f ⊤ L f . Corollary 1 (Continuum limit of smoothness regularizer). Assume states {s∗i }ni=1 are sampled i.i.d. from a smooth density p(s∗ ) supported on an m-dimensional C2 manifold M ⊂ Rd . Let W be the symmetrized k-NN affinity matrix constructed from a kernel K∆ with bandwidth h, and let L = D −W be the graph Laplacian.

If n → ∞, h → 0, and nhm+2 → ∞, then the normalized quadratic form converges to the weighted Dirichlet energy: 1 Tr F ⊤ LF −→ CK n2 hm+2

Z M

∥∇M f (s)∥22 p(s)2 dvol(s),

where CK > 0 is a constant depending only on the kernel K∆ . Theorem 1 (Manifold Regularized BC (LMRIL ) improves over vanilla BC (LBC )). Let f ∗ : S → A be the true, underlying expert policy, assumed to be C2 -smooth on a compact state space S . Let f : S → A denote the learned policy estimator. Consider two estimators trained on expert demonstrations: 1. Vanilla BC: a global supervised model minimizing LBC ( f ) = E(s∗ ,a∗ )∼D ∗ [ℓ( f (s∗ ), a∗ )]. 2. MRIL: a neighbor-based estimator minimizing " LMRIL ( f ) = LBC ( f ) + λ Es∗ ∼D ∗

∑

i∈Nk (s∗ )

wi (s∗ ) f (s∗ ) − f (s∗i )

2 2

# ,

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: In a closed-loop rollout, the deviation recursion ∆t+1 ≤ Ls ∆t + La ∥ f (st ) − f ∗ (st∗ )∥4

The Laplacian penalty in LMRIL is

∑ Wi j ∥ f (s∗i ) − f (s∗j )∥2 = 2 f ⊤ L f , i, j

where L = D −W is the graph Laplacian, and W is the k-NN affinity matrix. By Lemma 1, 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 a quadratic penalty is equivalent to Tikhonov regularization with respect to the graph Laplacian norm ∥ f ∥2L = f ⊤ L f . In statistical learning theory, adding a Tikhonov penalty strictly reduces the variance of the estimator compared to 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 1): for i.i.d. samples {s∗i } from density p on a smooth manifold M , the normalized penalty converges to Z M

∥∇ f (s)∥2 p(s)2 dvol(s).

This is the weighted Dirichlet energy of f on M . If this integral is finite, f belongs to the Sobolev space H 1 (M , p2 ), and in particular f is locally Lipschitz almost everywhere with ∥ f (s∗ ) − f (s′∗ )∥ ≤ C∥s∗ − s′∗ ∥

for p-a.e. neighbor pairs s∗ , s′∗ .

Therefore minimizers of LMRIL have uniformly bounded local Lipschitz constants along high-density regions of the state space. By contrast, minimizers of vanilla BC have no 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 ∆t = ∥st − st∗ ∥ denote the deviation at time t. For Lipschitz dynamics T ,

Decompose the action error:

∆t+1 ≤ Ls ∆t + La ∥ f (st ) − f ∗ (st∗ )∥.

∥ f (st ) − f ∗ (st∗ )∥ ≤ ∥ f (st ) − f ∗ (st )∥ + ∥ f ∗ (st ) − f ∗ (st∗ )∥. For vanilla BC, the first term ∥ f (st ) − f ∗ (st )∥ is only minimized on the empirical distribution PS ; off-distribution, it may be O(1) regardless of ∆t . The second term satisfies ∥ f ∗ (st ) − f ∗ (st∗ )∥ = O(∆t ) by smoothness of f ∗ . Hence the recursion can take the form ∆t+1 ≤ Ls ∆t + La O(1) + O(∆t ) ,

which accumulates linearly in t.

For MRIL, the Laplacian penalty enforces ∥ f (s∗ ) − f (s′∗ )∥ ≤ C∥s∗ − s′∗ ∥

for neighbor pairs (s∗ , s′∗ ),

4 L and L are the Lipschitz constants of the environment transition dynamics with respect to state and action, s a respectively.

Published as a conference paper at ICLR 2026

as shown in part (ii). Thus ∥ f (st ) − f ∗ (st )∥ = O(h2 ) (where h is the radius of the kernel bandwidth around the point st ) by local-linear regression error bounds, and ∥ f ∗ (st ) − f ∗ (st∗ )∥ = O(∆t ) by smoothness of f ∗ . Combining these, ∆t+1 ≤ Ls ∆t + La O(h2 ) + O(∆t ) . Since the constant multiplying ∆t is strictly smaller under the smoothness constraint, the cumulative error grows strictly slower than in vanilla BC. In particular, error growth is sublinear in the rollout horizon when h is small, whereas it is linear for vanilla BC.

Conclusion. Claims (i)–(iii) 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 a Gaussian kernel ∥s∗ −s∗ ∥2 , wi (s∗ ) ∝ exp − 2h2i ∑ wi (s∗ ) = 1, i∈Nk (s∗ )

with bandwidth h set to the median distance to the k-th nearest neighbor across the dataset. This choice is standard in manifold regularization (Belkin & Niyogi, 2008; Zhou et al., 2003) and ensures that the graph Laplacian penalty converges to the Dirichlet energy in the continuum limit. In practice, we found this default to be stable across tasks, though other kernels (e.g., uniform k-NN or exponential decay) yield qualitatively similar results. A.1.2

P ROOF OF T HEOREM 2

We begin with the required Lemmas establishing the spectral form of explicit Laplacian regularization and neighbor aggregation. Lemma 2 (Spectral form of explicit Laplacian regularization). Let L be the symmetric normalized graph Laplacian with eigenpairs {(µ j , u j )}nj=1 , where 0 = µ1 ≤ µ2 ≤ · · · ≤ µn ≤ 2. The minimizer of the penalized objective Lλ ( f ) = ∥ f − a∗ ∥2 + λ f ⊤ L f has the closed-form expansion n

1 ⟨a∗ , u j ⟩ u j . 1 + λ µj j=1

fλ = ∑

Thus λ directly determines the spectral filter φλ (µ) = (1 + λ µ)−1 applied to each Laplacian mode. Proof. Diagonalize L = UΛU ⊤ with U = [u1 , . . . , un ] orthogonal and Λ = diag(µ1 , . . . , µn ). Write f = Uc, a∗ = Ub in this basis. The objective becomes n

∥Uc −Ub∥2 + λ c⊤ Λc = ∑ (c j − b j )2 + λ µ j c2j . j=1

Minimizing each term yields c j = 1+λ1 µ j b j . Transforming back gives the stated expansion. Lemma 3 (Spectral form of neighbor aggregation). Let S = D−1 A be the random-walk matrix of the k-NN graph, with adjacency A and degree D. For any prediction vector f , the neighbor-averaged prediction is fˆ = S f . In the Laplacian eigenbasis, this corresponds to the spectral filter n

fˆ = ∑ (1 − µ j ) ⟨ f , u j ⟩ u j , j=1

i.e. φDARP (µ) = 1 − µ. Proof. By definition, L = I − D−1/2 AD−1/2 and S = D−1 A = I − Lrw where Lrw = D−1 L is the random-walk Laplacian. Since Lrw and L share the same spectrum up to similarity transform, the eigenbasis {u j } diagonalizes S. Thus for each mode u j , Su j = (1 − µ j )u j , 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 in Section 2.2), with eigenpairs {(µ j , u j )}nj=1 , where 0 = µ1 ≤ µ2 ≤ · · · ≤ µn ≤ 2.

The minimizers of the explicit MRIL objective (Section 2.2) and the implicit iMRIL objective (Section 2.3) have the following closed form expansions n

1 ⟨a∗ , u j ⟩ u j 1 + λ µ j j=1

fˆiMRIL = ∑ (1 − µ j ) ⟨ f , u j ⟩ u j

fMRIL = ∑

j=1

On the other hand, Lemma 2 shows that explicit Laplacian regularization with parameter λ yields the spectral filter φλ (µ) = (1 + λ µ)−1 . Both filters downweight high-frequency modes (µ ≫ 0) while preserving low-frequency modes (µ ≈ 0). The key difference is that φλ (µ) requires tuning λ , whereas φDARP (µ) is parameter-free. To see the equivalence, note that for small µ, φDARP (µ) = 1 − µ ≈ (1 + µ)−1 = φλ =1 (µ) up to O(µ 2 ) terms. Thus DARP can be interpreted as performing Laplacian smoothing with an effective regularization weight of order λ ≈ 1 in normalized units. Moreover, for large µ, φDARP (µ) damps high-frequency modes even more strongly by driving them toward zero, providing a sharper low-pass effect than explicit regularization. Therefore, DARP ’s aggregation step is mathematically equivalent to implicit Laplacian regularization with fixed spectral filter φDARP , eliminating the need to tune λ explicitly. DARP can therefore be viewed as a form of locally adaptive implicit regularization: rather than introducing an explicit global weight λ , its neighbor aggregation step enforces smoothness automatically through the graph structure. The effective regularization strength varies with local degree Spectral Filters: Explicit Laplacian vs. DARP

Spectral filter φ (µ)

φDARP (µ) = 1 − µ φλ =0.5 (µ) φλ =1 (µ) φλ =2 (µ)

0.5

−0.5

0.2

0.4

0.6

0.8 1 1.2 Eigenvalue µ

1.4

1.6

1.8

Figure 8: 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 (µ) = 1 − µ, which suppresses high-frequency modes more aggressively than any fixed explicit λ . Figure 8 illustrates this comparison, showing how DARP achieves sharper low-pass filtering without the need for hyperparameter tuning.

A.2.1

R ETRIEVAL

DARP BC

DARP Success Rate vs. Lookback

1 Success Rate

DARP Success Rate vs. Number of Neighbors

Success Rate

A DDITIONAL E XPERIMENTAL D ETAILS

2000

DARP BC

Lookback

DARP Success Rate vs. Decay

1 Success Rate

A.2

DARP BC

Decay

-5

Figure 9: DARP performance analysis as retrieval hyperparameters are swept: (left) observe that the perfor-

mance of a DARP model is poor when using few neighbors, reaches a global optimum when retrieving about 500 neighbors, and plateaus just above BC’s success rate as k goes to the size of the dataset; (center) observe that the performance of a DARP model generally slightly improves as more history is considered, and only performs worse than BC when very little or no history is considered; (right) observe that the performance of a DARP model is sensitive to how much weight is applied to older observations when performing retrieval. Intuitively, if this decay is too high, DARP performance is nearly identical to having little to no lookback, performing worse than BC. Success rate is measured out of 50 trials on the Robosuite Stack environment. 95% confidence intervals are included.

The selection of the distance function d(sq , s∗i ) 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 we are working with sequences of states and incorporates history in our distance calculation. Suppose we have a query trajectory Sq = (. . . , sq,−1 , sq,0 ) where sq,0 is the current query state sq . Now suppose we want to calculate d(sq , s∗i ), where s∗i is some state from the expert dataset. We first find the trajectory this state is from—call this S∗j —and the index of s∗i in this trajectory—call this t. Thus, s∗i can be rewritten as s∗j,t . Given some lookback parameter ℓ which denotes how many past states we want to consider, we get: ℓ−1

d(sq , s∗i ) = ∑ ∥sq,−n − s∗j,t−n ∥ n=0

This is simply the accumulation of Euclidean distances of the current and last ℓ − 1 states from both the query trajectory and the source trajectory, assuming valid indices. Of course, in practice, we generally want to put more emphasis on more recent states, as we want them to be more influential in the selection of neighbors. Thus, given some rate of exponential decay γ ≥ 0, we have ℓ−1

d(sq , s∗i ) = ∑ ∥sq,−n − s∗j,t−n ∥ · e−γn n=0

See Figure 9 for an experimental analysis on how the success rate in an environment changes as these parameters are swept. A.2.2

C AN DARP HANDLE TASKS REQUIRING THE REPRESENTATION OF MULTI - MODAL ACTION DISTRIBUTIONS ?

We test DARP’s ability to handle complex action distributions by evaluating on the Push-T task, as described in (Chi et al., 2024), which requires representing multi-modal action distributions (see Figure 10 for a visualization). For this experiment, DARP employs a Set Transformer head that predicts parameters of a Gaussian Mixture Model. We note that DARP with a GMM head is able to handle multi-modal distributions effectively, showing a 22% improvement over BC on the Push-T 19

Published as a conference paper at ICLR 2026

Method

Score

BC DARP

48 ± 8 70 ± 8

Table 4: Push-T Results. Aver-

aged over 100 trials, DARP outperforms BC.

Figure 10: Push-T Environment. The goal is to control the blue circle to push the T-shaped block.

task (Q2) (see Table 4). This demonstrates that DARP can be further adapted to multi-modal action distribution modeling requirements. A.2.3

C AN DARP HANDLE DISCONTINUOUS ENVIRONMENTS WHERE NEARBY STATES MAY REQUIRE OPPOSING ACTIONS ?

A key concern for neighbor-based approaches is performance in environments with strong discontinuities, 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 11 for a visualization). Even in this deliberately challenging discontinuous environment, DARP achieves a 57% success rate compared to BC’s 25%. (Q3) (see Table 5) This suggests that the distance vectors and permutation-invariant aggregation help the model distinguish between appropriate and inappropriate neighbors, even when spatial proximity doesn’t guarantee action similarity.

Succ. (%)

BC DARP

25 57

Table 5: Long maze results. Averaged over 100 trials, DARP significantly outperforms BC.

Figure 11: Long maze environment. The goal is to move a force-actuated ball from the green start to the red destination. A.2.4

Method

C AN DARP R ECOVER F ROM BC E RROR ?

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 a fork of the environment and begin rolling out a DARP agent in that clone of the environment. The results (seen in Fig. 12) 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 error and can perform well in the slightly out-of-distribution states that a failing BC agent drifts into. 20

Published as a conference paper at ICLR 2026

Forked DARP Scores (Hopper)

Forked DARP Scores (Stack)

140

3600

Score

BC DARP

1000 0 190

Step

BC DARP

Step

400

Figure 12: In two different tasks (the Robosuite Stack task and the MuJoCo Hopper task), we rollout a BC agent and create a fork of the environment every k steps (in this case, k = 10). Observe that, even as BC nears the end of its failing rollout, DARP is able to scale highly, and is only prevented from doing so about halfway through the Stack rollout and about 80% through the Hopper rollout. A.2.5

C OMPARISON WITH CCIL Environment

Method

Improvement over BC (%)

Hopper

CCIL DARP

32.6% 53.2%

Walker

CCIL DARP

84.1% 84.1%

Ant

CCIL DARP

12.6% 84.5%

HalfCheetah

CCIL DARP

5.4% 418.7%

Table 6: Relative improvement over BC compared to CCIL. Comparing DARP against CCIL on standard MuJoCo benchmarks. DARP achieves higher or equal relative gains across all tasks. As shown in Table 6, 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 a particularly large margin on HalfCheetah (418.7% vs 5.4% improvement). A.2.6

D ISTANCE M ETRIC S ENSITIVITY Method

Success Rate (%)

BC DARP (Cosine Similarity) DARP (Euclidean Distance)

47 70 75

Table 7: Distance Metric Sensitivity. Comparison of 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 to choose nearby neighbors, it is natural to consider alternative metrics, like cosine similarity, especially in high-dimensional embeddings such as R3M. We find (as shown in Table 7) that DARP performs similarly – only 5% worse – when using cosine similarity rather than Euclidean distance. This suggests DARP is robust to the choice of distance metric used. 21

Published as a conference paper at ICLR 2026

A.2.7

DARP IN C OMBINATION W ITH D IFFUSION P OLICY

Method

Stack

Threading

Peg Insertion

BC (MLP) DARP (MLP)

34 62

12 38

28 40

Diffusion DARP w/ Diffusion

52 82

44 72

54 66

Table 8: 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 are state-based. Note that the number of demonstrations used is much less than that used in Table 2, so baseline BC and DARP numbers do not match.

While all reported experiments are performed with an MLP backbone, diffusion policy (Chi et al., 2024) has proven to be a state-of-the-art model class for imitation learning, particularly for manipulation tasks. Table 8 reveals that DARP is not mutually exclusive with diffusion and can be combined for an even more performant model. Using the DARP architecture with a diffusion backbone outperforms 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)

5.610 0.559

0.15700 0.00437

Diffusion DARP (MLP)

Table 9: 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 ≈ 10× slower in training and ≈ 36× slower during inference. Additionally, we empirically find that DARP with an MLP backbone is much faster than standard diffusion, particularly in inference – see Table 9.

A.2.8

DARP T RAINED ON H UMAN D EMONSTRATIONS

Method

Success Rate (%)

BC DARP

45 60

Table 10: Results with human demonstrations. Comparison of success rates when training on human data rather than data collected by RL policies. Results are on the Robosuite Stack task.

The expert demonstrations used to 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 on expert demonstrations collected by humans. Indeed, when trained on human demonstrations on the Robosuite Stack task, DARP outperforms standard BC by 15% (see Table 10). 22

Published as a conference paper at ICLR 2026

A.2.9

C HOICE OF R ETRIEVAL H YPERPARAMETERS Environment

Method

Score (Mean)

Hopper

BC DARP

507.49 805.69

Stack

BC DARP

0.12 0.31

Table 11: 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 a performance gain in comparison to BC. Figure 9 reveals that there are selections of retrieval parameters (for example, a 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 11. A.3

P SEUDOCODE

We provide pseudocode of the DARP algorithm, see Algorithm 1. Algorithm 1 Difference-Aware Retrieval Policies 1: Input: Expert demonstrations D ∗ = {(s∗j , a∗j )}, number of neighbors k 2: Initialize: f parameters θ 3: if g is parametric then 4: Initialize: g parameters ψ 5: end if 6: // Training Loop 7: while not converged do 8: Sample batch of query data (s∗q , a∗q ) ∼ D ∗ 9: for each query pair (s∗q , a∗q ) in batch do 10: // Find k-Nearest Neighbors from the entire dataset D ∗ 11: Nk (s∗q ) ← arg min -k j d(s∗q , s∗j ) 12: // Compute Neighbor-based Predictions 13: for each neighbor index i ∈ Nk (s∗q ) do 14: a′i ← fθ (s∗i , a∗i , s∗i − s∗q ) 15: end for 16: // Aggregate Predictions 17: if g is parametric then 18: âq ← gψ ({a′i }i∈Nk (s∗q ) ) 19: else 20: âq ← g({a′i }i∈Nk (s∗q ) ) 21: end if 22: end for 23: // Update Parameters based on the batch loss 24: L ← ∑(s∗q ,a∗q )∈batch ∥âq − a∗q ∥2 25: // Gradient descent step 26: θ ← θ − α∇θ L 27: if g is parametric then 28: ψ ← ψ − α∇ψ L 29: end if 30: end while 31: Output: Trained parameters θ and, if applicable, ψ

Published as a conference paper at ICLR 2026

A.4

RUNTIME A NALYSIS Environment

Train (s/epoch)

Test (s/step)

Hopper (1,000 datapoints)

BC 100 200 300 400 500

0.102 0.126 0.128 0.130 0.131 0.130

0.00127 0.00413 0.00418 0.00420 0.00421 0.00419

Stack (4,200 datapoints)

BC 100 250 500 750 1000 1500 2000

0.431 0.556 0.539 0.559 0.576 0.591 0.682 0.758

0.00139 0.00437 0.00442 0.00437 0.00433 0.00478 0.00452 0.00451

Table 12: Runtime comparison across environments. Training and testing speeds (in seconds) for Behavior Cloning (BC) and varying values of k on Hopper and Stack datasets.

As shown in Table 12, we analyze the computational overhead of DARP at both training-time and inference-time. Our analysis indicates that computation cost scales sub-linearly as k increases, and that inference time is tractable for real-time robotic application. For example, on the Robosuite Stack task at k = 500, inference takes approximately 0.00437 seconds per step on our hardware, corresponding to a control frequency higher than 230 Hz.

Record · ID 267643 · SHA-256 148bf30e5da329e3
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.