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

arXiv: https://arxiv.org/abs/2207.08673 · HTML: https://mlsanity.com/paper/2207.08673/

## Key facts

- Cited by 20 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 47 references that are in the map.
- First posted on arXiv 18 Jul 2022.

## Builds on (most cited first, 30 of 33)

- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction](https://mlsanity.com/paper/1906.00949/) (2019)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [An Algorithmic Perspective on Imitation Learning](https://mlsanity.com/paper/1811.06711/) (2018)
- [Safe Learning in Robotics: From Learning-Based Control to Safe Reinforcement Learning](https://mlsanity.com/paper/2108.06266/) (2021)
- [COMBO: Conservative Offline Model-Based Policy Optimization](https://mlsanity.com/paper/2102.08363/) (2021)
- [Learning Hand-Eye Coordination for Robotic Grasping with Deep Learning and Large-Scale Data Collection](https://mlsanity.com/paper/1603.02199/) (2016)
- [Deep Imitation Learning for Complex Manipulation Tasks from Virtual Reality Teleoperation](https://mlsanity.com/paper/1710.04615/) (2017)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Contrastive Learning of Structured World Models](https://mlsanity.com/paper/1911.12247/) (2019)
- [Recovery RL: Safe Reinforcement Learning with Learned Recovery Zones](https://mlsanity.com/paper/2010.15920/) (2020)
- [Deep Spatial Autoencoders for Visuomotor Learning](https://mlsanity.com/paper/1509.06113/) (2015)
- [PLAS: Latent Action Space for Offline Reinforcement Learning](https://mlsanity.com/paper/2011.07213/) (2020)
- [State Representation Learning for Control: An Overview](https://mlsanity.com/paper/1802.04181/) (2018)
- [Learning to be Safe: Deep RL with a Safety Critic](https://mlsanity.com/paper/2010.14603/) (2020)
- [Overcoming Limitations of Mixture Density Networks: A Sampling and Fitting Framework for Multimodal Future Prediction](https://mlsanity.com/paper/1906.03631/) (2019)
- [Uncertainty-Aware Reinforcement Learning for Collision Avoidance](https://mlsanity.com/paper/1702.01182/) (2017)
- [Deep Predictive Policy Training using Reinforcement Learning](https://mlsanity.com/paper/1703.00727/) (2017)
- [Learning to Linearize Under Uncertainty](https://mlsanity.com/paper/1506.03011/) (2015)
- [Guided Uncertainty-Aware Policy Optimization: Combining Learning and Model-Based Strategies for Sample-Efficient Policy Learning](https://mlsanity.com/paper/2005.10872/) (2020)
- [Affordance Learning for End-to-End Visuomotor Robot Control](https://mlsanity.com/paper/1903.04053/) (2019)
- [Equivariant Representation Learning via Class-Pose Decomposition](https://mlsanity.com/paper/2207.03116/) (2022)
- [NewtonianVAE: Proportional Control and Goal Identification from Pixels via Physical Latent Spaces](https://mlsanity.com/paper/2006.01959/) (2020)
- [LS3: Latent Space Safe Sets for Long-Horizon Visuomotor Control of Sparse Reward Iterative Tasks](https://mlsanity.com/paper/2107.04775/) (2021)

## Built on by (most cited first, 20 of 20)

- [Uncertainty-aware Latent Safety Filters for Avoiding Out-of-Distribution Failures](https://mlsanity.com/paper/2505.00779/) (2025)
- [CCIL: Continuity-based Data Augmentation for Corrective Imitation Learning](https://mlsanity.com/paper/2310.12972/) (2023)
- [One ACT Play: Single Demonstration Behavior Cloning with Action Chunking Transformers](https://mlsanity.com/paper/2309.10175/) (2023)
- [SAFE-GIL: SAFEty Guided Imitation Learning for Robotic Systems](https://mlsanity.com/paper/2404.05249/) (2024)
- [RecoveryChaining: Learning Local Recovery Policies for Robust Manipulation](https://mlsanity.com/paper/2410.13979/) (2024)
- [Out-of-Distribution Recovery with Object-Centric Keypoint Inverse Policy for Visuomotor Imitation Learning](https://mlsanity.com/paper/2411.03294/) (2024)
- [SafeMimic: Towards Safe and Autonomous Human-to-Robot Imitation for Mobile Manipulation](https://mlsanity.com/paper/2506.15847/) (2025)
- [Green-VLA: Staged Vision-Language-Action Model for Generalist Robots](https://mlsanity.com/paper/2602.00919/) (2026)
- [ARCADE: Scalable Demonstration Collection and Generation via Augmented Reality for Imitation Learning](https://mlsanity.com/paper/2410.15994/) (2024)
- [Real-Time Operator Takeover for Visuomotor Diffusion Policy Training](https://mlsanity.com/paper/2502.02308/) (2025)
- [Using Non-Expert Data to Robustify Imitation Learning via Offline Reinforcement Learning](https://mlsanity.com/paper/2510.19495/) (2025)
- [Dense Dynamics-Aware Reward Synthesis: Integrating Prior Experience with Demonstrations](https://mlsanity.com/paper/2412.01114/) (2024)
- [Failure Forecasting Boosts Robustness of Sim2Real Rhythmic Insertion Policies](https://mlsanity.com/paper/2507.06519/) (2025)
- [Learning Robust Execution in Robotic Manipulation with Agentic Reinforcement Learning](https://mlsanity.com/paper/2607.13818/) (2026)
- [Bimanual rope manipulation skill synthesis through context dependent correction policy learning from human demonstration](https://mlsanity.com/paper/2209.13850/) (2022)
- [Exploiting Policy Idling for Dexterous Manipulation](https://mlsanity.com/paper/2508.15669/) (2025)
- [Diffusing Trajectory Optimization Problems for Recovery During Multi-Finger Manipulation](https://mlsanity.com/paper/2510.07030/) (2025)
- [DEFENDER: DTW-Based Episode Filtering Using Demonstrations for Enhancing RL Safety](https://mlsanity.com/paper/2305.04727/) (2023)
- [PACT: Self-Evolving Physical Safety Alignment for Diffusion Policies in Embodied Manipulation](https://mlsanity.com/paper/2606.08414/) (2026)
- [UniIntervene++: An Adaptive Intervention Agent for Efficient Real-World Reinforcement Learning](https://mlsanity.com/paper/2610.03620/) (2026)

## Abstract (arXiv)

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.