Policy-Calibrated DAgger: Offline Calibrated Noise Injection for Imitation Learning

Jenny Wang; George Kantor

cs.ROSep 24, 2026

Policy-Calibrated DAgger: Offline Calibrated Noise Injection for Imitation Learning

Authors: Jenny Wang, George Kantor

Organizations: Carnegie Mellon University, Pittsburgh PA, USA

Abstract

Policies trained with imitation learning can accumulate errors over time, causing the robot to drift outside the training distribution. Existing methods mitigate this covariate shift by collecting additional data where the policy fails or is likely to fail. The first places the robot in unsafe conditions and the second requires choosing an appropriate noise distribution to collect new expert demonstrations under that noise. We propose Policy-Calibrated DAgger, a method that makes use of the properties of recent generative policies to estimate the policy's noise offline by using its own predicted action distribution. We measure a diffusion policy's spread of predicted actions at observations along the expert trajectory and measure its closed-loop error relative to a recorded trajectory. To address issues with measuring error in a multimodal action space, we guide the policy towards the trajectory during closed-loop control through partial denoising, and use properties of a diffusion model to unnormalize the measured error as if we did not guide it. We experiment in a scenario where a robot is tasked to reach an engine lever in a cluttered and narrow environment and show results in a 3D photorealistic simulator and a 2D planar reacher environment. We show that our method surpasses policies trained with dataset aggregation without noising and matches the performance of the best noise level in hindsight, without requiring a sweep over noise levels.

Figures & tables

Fig. 1 : Our method estimates policy error while tracking a trajectory by making use of the probabilistic properties of generative policies. By sampling along pre-recorded expert data and making use of the model’s noise schedule to undo the normalizing effects of partial denoising towards the expert distribution in closed-loop control in simulation, we can estimate the policy’s noise with respect to the expert data. We can sample targeted recovery trajectories within this policy-calibrated noise distribution to reduce covariate shift between the policy’s visited states and the expert data.

Fig. 2 : Our method augments the data collected during a DAgger-style data collection loop. The policy at data collection round $k$ is trained on the base dataset $\mathcal{D}_{0}$ , aggregated intervention data $\mathcal{D}_{1}\ldots\mathcal{D}_{k}$ and noised intervention data $\hat{\mathcal{D}}_{1}\ldots\hat{\mathcal{D}}_{k}$ . To generate policy-calibrated noised intervention data: (a) We sample the policy open loop at every other timestep along the expert trajectory to estimate $\hat{\Sigma}_{t}$ . (b) We sample blended trajectories between the policy and expert with partial denoising to estimate $W$ . (c) We combine (a) and (b) to compute the policy noise per timestep $\Sigma^{\alpha}_{t}$ . (d) We average the per-timestep $\Sigma^{\alpha}_{t}$ across all data to obtain the pooled $\bar{\Sigma}^{\alpha}$ . In this example, the per-timestep ellipsoidal noise estimates average to become more isotropic. We sample states and corrective action pairs within this local area around the expert intervention.

Fig. 3 : Empirically, performing partial denoising at lower blend ratios $\beta$ results in action chunks closer to the sampled expert chunk. The unguided policy $\beta=1.0$ may sometimes travel in a different direction than the sampled expert chunk.

Fig. 4 : Our fitted curve for estimating $W$ on the planar reacher environment. As blend ratio $\beta\rightarrow 1$ , noise variance fraction $\sigma_{k_{sw}}^{2}\rightarrow 1$ and thus our estimate $x(\sigma_{k_{sw}}^{2})\rightarrow W$ . We measure error in terms of median step size measured from the demonstration data.


Method	$K{=}1$	$K{=}2$	$K{=}3$	$K{=}4$	$K{=}5$
BC (no aggregation)	68.0 $\pm$ 0.2 (1)
HG-DAgger [ 7 ]	79.9 $\pm$ 0.5 (5)	81.9 $\pm$ 0.7 (5)	83.3 $\pm$ 1.5 (5)	85.3 $\pm$ 0.7 (5)	87.3 $\pm$ 1.0 (5)
Fixed iso. noise ( $\sigma{=}2$ ) [ 28 ]	81.9 $\pm$ 0.8 (5)	84.8 $\pm$ 0.7 (5)	85.1 $\pm$ 0.3 (5)	88.6 $\pm$ 0.4 (5)	89.1 $\pm$ 0.4 (5)
Fixed iso. noise ( $\sigma{=}4$ )	82.7 $\pm$ 0.9 (5)	86.1 $\pm$ 0.7 (5)	86.3 $\pm$ 0.5 (5)	87.7 $\pm$ 0.6 (5)	89.2 $\pm$ 0.8 (5)
Fixed iso. noise ( $\sigma{=}8$ )	83.9 $\pm$ 0.8 (5)	85.1 $\pm$ 0.8 (5)	85.1 $\pm$ 0.8 (5)	86.1 $\pm$ 0.8 (5)	87.1 $\pm$ 0.8 (5)

TABLE I : Mean

\pm

SE success rates on a 100-scenario planar reacher environment (3 evaluations per scenario). Most cells are averaged across 5 DAgger lineages (5) . BC is trained on

\mathcal{D}_{0}

(500 demonstrations). To prevent sequential optimization errors, all policy variants are fine-tuned from BC’s 75k checkpoint for each round

K

(175k total steps). We sweep over isotropic noise levels and calibrated noise estimated before training for the round. Darker colors denote higher success. Our method matches the performance of the best noise level in the hyperparameter sweep without requiring extra expert data.


Method	$K{=}1$	$K{=}2$
BC (no aggregation)	58.0 $\pm$ 1.2 (1)
HG-DAgger [ 7 ]	57.4 $\pm$ 5.6 (3)	58.1 $\pm$ 4.0 (3)
Pooled $\bar{\Sigma}^{\alpha}$ (Ours)	68.8 $\pm$ 4.2 (3)	63.4 $\pm$ 4.8 (3)

TABLE II : Mean

\pm

SE success rates on a 100-scenario photorealistic engine lever-reaching environment (3 evaluations per scenario). Following Table I ’s protocol, policy variants are fine-tuned from BC’s 75k training step checkpoint to a total of 95k training steps. BC is trained on 500 demonstrations.

Fig. 5 : We evaluate our method in two cluttered goal-reaching tasks with narrow passageways, where the robot often must untangle its links to make progress towards the goal. The 3D photorealistic physics simulation recreates our real mockup. We run ablations on the 3-link planar robot.

Explore similar work

WM-DAgger: Enabling Efficient Data Aggregation for Imitation Learning with World Models
Apr 13, 2026Anlan Yu, Zaishu Chen, Peili Song +6Robot Imitation Learning Data-Efficient Robot Learning
Peking University, Beijing, China · JD Logistics, China · Nankai University, Tianjin, China +4
BlenDAgger: Blended Shared Control for Interactive Imitation Learning
Sep 29, 2026Cailyn Smith, Geoffrey Sun, Henny Admoni +1Robotic Control Robotic Manipulation
The Robotics Institute, School of Computer Science Carnegie Mellon University Pittsburgh, USA · School of Computer Science Carnegie Mellon University Pittsburgh, USA
How to Mitigate the Distribution Shift Problem in Robotics Control: A Robust and Adaptive Approach Based on Offline to Online Imitation Learning
May 25, 2026Hyung-Suk Yoon, Seung-Woo SeoDistribution Shift Robustness Imitation Learning
Department of Electronic and Computer Engineering, Seoul National University, Seoul, South Korea.

Apr 13, 2026J/K move · Enter open · S save

WM-DAgger: Enabling Efficient Data Aggregation for Imitation Learning with World Models

Anlan Yu, Zaishu Chen, Peili Song, Zhiqing Hong+5

Peking University, Beijing, China · JD Logistics, China · Nankai University, Tianjin, China+4

Robot Imitation Learning Data-Efficient Robot Learning Interactive Imitation Learning Synthetic Data Generation for Robotics World Models for Robotics

Open paper AskPDF arXiv

Imitation learning is a powerful paradigm for training robotic policies, yet its performance is limited by compounding errors: minor policy inaccuracies could drive robots into unseen out-of-distribution (OOD) states in the training set, where the policy could generate even bigger errors, leading to eventual failures. While the Data Aggregation (DAgger) framework tries to address this issue, its reliance on continuous human involvement severely limits scalability. In this paper, we propose WM-DAgger, an efficient data aggregation framework that leverages World Models to synthesize OOD recovery data without requiring human involvement. Specifically, we focus on manipulation tasks with an eye-in-hand robotic arm and only few-shot demonstrations. To avoid synthesizing misleading data and overcome the hallucination issues inherent to World Models, our framework introduces two key mechanisms: (1) a Corrective Action Synthesis Module that generates task-oriented recovery actions to prevent misleading supervision, and (2) a Consistency-Guided Filtering Module that discards physically implausible trajectories by anchoring terminal synthesized frames to corresponding real frames in expert demonstrations. We extensively validate WM-DAgger on multiple real-world robotic tasks. Results that our method significantly improves success rates, achieving a 93.3% success rate in soft bag pushing with only five demonstrations. The source code is publicly available at https://github.com/czs12354-xxdbd/WM-Dagger.

Abstract

Explore similar work

WM-DAgger: Enabling Efficient Data Aggregation for Imitation Learning with World Models

BlenDAgger: Blended Shared Control for Interactive Imitation Learning

How to Mitigate the Distribution Shift Problem in Robotics Control: A Robust and Adaptive Approach Based on Offline to Online Imitation Learning