Back to the Manifold: Recovering from Out-of-Distribution States

TL;DR

This paper introduces a recovery policy for robots that uses data density approximation and a learned latent space to bring the robot back to in-distribution states after external perturbations, improving task success.

Contribution

The paper proposes a novel recovery policy framework that addresses deployment-time distributional shifts by leveraging data density and equivariant mappings.

Findings

Recovery policy successfully restores robot states after perturbations.

Behavioral cloning fails under distributional shift without recovery.

Real robotic experiments validate the effectiveness of the proposed method.

Abstract

Learning from previously collected datasets of expert data offers the promise of acquiring robotic policies without unsafe and costly online explorations. However, a major challenge is a distributional shift between the states in the training dataset and the ones visited by the learned policy at the test time. While prior works mainly studied the distribution shift caused by the policy during the offline training, the problem of recovering from out-of-distribution states at the deployment time is not very well studied yet. We alleviate the distributional shift at the deployment time by introducing a recovery policy that brings the agent back to the training manifold whenever it steps out of the in-distribution states, e.g., due to an external perturbation. The recovery policy relies on an approximation of the training data density and a learned equivariant mapping that maps visual observations into a latent space in which translations correspond to the robot actions. We demonstrate the effectiveness of the proposed method through several manipulation experiments on a real robotic platform. Our results show that the recovery policy enables the agent to complete tasks while the behavioral cloning alone fails because of the distributional shift problem.