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

Summary

Abstract Summary

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.

Key Points

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.

Preliminaries

We aim to minimize covariate shift by optimizing the KL-divergence between the trajectories visited by the trained robot policy $p(\xi/\pi_{\theta})$ and those visited by the noised expert $p(\xi/\pi_{\theta^{*}},\psi)$. In DART [ 6 ] , the authors minimize covariate shift by optimizing the KL-divergence with respect to the sufficient statistics $\psi$ of the noise distribution.

Policy-Calibrated DAgger

We introduce Policy-Calibrated DAgger, a method that predicts data augmentation noise parameters that make the expert demonstrations cover states closer to the final robot’s distribution. A key difference between our method and DART [ 6 ] is that we do not inject noise into the expert’s control loop—instead, we compute the sufficient statistics offline and in a simulator.

This places the burden of corrections on the policy instead of the expert and allows the expert to provide clean demonstrations, uninterrupted.

Policy-Calibrated DAgger — Policy-Calibrated DAgger: Offline Calibrated Noise Injection for Imitation Learning — 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.

Experiments

We conduct a series of experiments to validate our method for data-efficient interactive imitation learning. Our primary research question is how our proposed augmentations to intervention data affect policy optimization given the same number of expert interventions. Our augmentations are based on a diffusion model’s [ 43 ] denoising process.

Sources and Demonstrations

Figures are reproduced from Wang et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.