Difference-Aware Retrieval Policies for Imitation Learning

• URL: http://arxiv.org/abs/2606.09758v1
• Date: Mon, 08 Jun 2026 17:18:19 GMT
• Title: Difference-Aware Retrieval Policies for Imitation Learning
• Abstract summary: Parametric imitation learning via behavior cloning can suffer from poor generalization to out-of-distribution states.<n>We show that reusing the training data during inference via a semi-parametric retrieval-based imitation learning approach can alleviate this challenge.<n>We present Difference-Aware Retrieval Policies for Imitation Learning (DARP), a semi-parametric retrieval-based imitation learning approach.
• Score: 8.909659724448819
• License: http://creativecommons.org/licenses/by/4.0/
• 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. Code and demos are available at https://weirdlabuw.github.io/darp-site/.

Related papers

• Scalable Data Attribution via Forward-Only Test-Time Inference [3.5466521714943138]
  Data attribution seeks to trace model behavior back to the training examples that shaped it.<n>We propose a data attribution method that preserves the same first-order counterfactual target.<n>Our method provides a theoretical framework for practical, real-time data attribution in large pretrained models.
  arXiv  Detail & Related papers  (2025-11-25T00:11:39Z)
• A Dual Approach to Imitation Learning from Observations with Offline Datasets [19.856363985916644]
  Demonstrations are an effective alternative to task specification for learning agents in settings where designing a reward function is difficult. We derive DILO, an algorithm that can leverage arbitrary suboptimal data to learn imitating policies without requiring expert actions.
  arXiv  Detail & Related papers  (2024-06-13T04:39:42Z)
• Learning a Diffusion Model Policy from Rewards via Q-Score Matching [93.0191910132874]
  We present a theoretical framework linking the structure of diffusion model policies to a learned Q-function.<n>We propose a new policy update method from this theory, which we denote Q-score matching.
  arXiv  Detail & Related papers  (2023-12-18T23:31:01Z)
• Get Back Here: Robust Imitation by Return-to-Distribution Planning [43.26690674765619]
  We consider the Imitation Learning (IL) setup where expert data are not collected on the actual deployment environment but on a different version. To address the resulting distribution shift, we combine behavior cloning (BC) with a planner that is tasked to bring the agent back to states visited by the expert whenever the agent deviates from the demonstration distribution. The resulting algorithm, POIR, can be trained offline, and leverages online interactions to efficiently fine-tune its planner to improve performance over time.
  arXiv  Detail & Related papers  (2023-05-02T13:19:08Z)
• Let Offline RL Flow: Training Conservative Agents in the Latent Space of Normalizing Flows [58.762959061522736]
  offline reinforcement learning aims to train a policy on a pre-recorded and fixed dataset without any additional environment interactions. We build upon recent works on learning policies in latent action spaces and use a special form of Normalizing Flows for constructing a generative model. We evaluate our method on various locomotion and navigation tasks, demonstrating that our approach outperforms recently proposed algorithms.
  arXiv  Detail & Related papers  (2022-11-20T21:57:10Z)
• Value-Consistent Representation Learning for Data-Efficient Reinforcement Learning [105.70602423944148]
  We propose a novel method, called value-consistent representation learning (VCR), to learn representations that are directly related to decision-making. Instead of aligning this imagined state with a real state returned by the environment, VCR applies a $Q$-value head on both states and obtains two distributions of action values. It has been demonstrated that our methods achieve new state-of-the-art performance for search-free RL algorithms.
  arXiv  Detail & Related papers  (2022-06-25T03:02:25Z)
• Accounting for the Sequential Nature of States to Learn Features for Reinforcement Learning [2.0646127669654826]
  We investigate the properties of data that cause popular representation learning approaches to fail. In particular, we find that in environments where states do not significantly overlap, variational autoencoders (VAEs) fail to learn useful features. We demonstrate this failure in a simple gridworld domain, and then provide a solution in the form of metric learning.
  arXiv  Detail & Related papers  (2022-05-12T10:20:43Z)
• Causal Imitation Learning under Temporally Correlated Noise [39.4969161422156]
  We develop algorithms for imitation learning from policy data corrupted by temporally correlated noise in expert actions. In particular, we present two techniques, one of a generative-modeling flavor (DoubIL) that can utilize access to a simulator, and one of a game-theoretic flavor (ResiduIL) that can be run entirely offline.
  arXiv  Detail & Related papers  (2022-02-02T22:33:08Z)
• Curriculum Offline Imitation Learning [72.1015201041391]
  offline reinforcement learning tasks require the agent to learn from a pre-collected dataset with no further interactions with the environment. We propose textitCurriculum Offline Learning (COIL), which utilizes an experience picking strategy for imitating from adaptive neighboring policies with a higher return. On continuous control benchmarks, we compare COIL against both imitation-based and RL-based methods, showing that it not only avoids just learning a mediocre behavior on mixed datasets but is also even competitive with state-of-the-art offline RL methods.
  arXiv  Detail & Related papers  (2021-11-03T08:02:48Z)
• Object-Aware Regularization for Addressing Causal Confusion in Imitation Learning [131.1852444489217]
  This paper presents Object-aware REgularizatiOn (OREO), a technique that regularizes an imitation policy in an object-aware manner. Our main idea is to encourage a policy to uniformly attend to all semantic objects, in order to prevent the policy from exploiting nuisance variables strongly correlated with expert actions.
  arXiv  Detail & Related papers  (2021-10-27T01:56:23Z)
• Graph Sampling Based Deep Metric Learning for Generalizable Person Re-Identification [114.56752624945142]
  We argue that the most popular random sampling method, the well-known PK sampler, is not informative and efficient for deep metric learning. We propose an efficient mini batch sampling method called Graph Sampling (GS) for large-scale metric learning.
  arXiv  Detail & Related papers  (2021-04-04T06:44:15Z)
• Strictly Batch Imitation Learning by Energy-based Distribution Matching [104.33286163090179]
  Consider learning a policy purely on the basis of demonstrated behavior -- that is, with no access to reinforcement signals, no knowledge of transition dynamics, and no further interaction with the environment. One solution is simply to retrofit existing algorithms for apprenticeship learning to work in the offline setting. But such an approach leans heavily on off-policy evaluation or offline model estimation, and can be indirect and inefficient. We argue that a good solution should be able to explicitly parameterize a policy, implicitly learn from rollout dynamics, and operate in an entirely offline fashion.
  arXiv  Detail & Related papers  (2020-06-25T03:27:59Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.