# Improving Robustness to Out-of-Distribution States in Imitation Learning via Deep Koopman-Boosted Diffusion Policy

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 37 of its 49 references that are in the map.
- First posted on arXiv 1 Nov 2025.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)
- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Deep Inside Convolutional Networks: Visualising Image Classification Models and Saliency Maps](https://mlsanity.com/paper/1312.6034/) (2013)
- [Shortcut Learning in Deep Neural Networks](https://mlsanity.com/paper/2004.07780/) (2020)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [RLBench: The Robot Learning Benchmark & Learning Environment](https://mlsanity.com/paper/1909.12271/) (2019)
- [Deep learning for universal linear embeddings of nonlinear dynamics](https://mlsanity.com/paper/1712.09707/) (2017)
- [Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion](https://mlsanity.com/paper/2310.02279/) (2023)
- [Implicit Behavioral Cloning](https://mlsanity.com/paper/2109.00137/) (2021)
- [Understanding Diffusion Models: A Unified Perspective](https://mlsanity.com/paper/2208.11970/) (2022)
- [Causal Confusion in Imitation Learning](https://mlsanity.com/paper/1905.11979/) (2019)
- [Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets](https://mlsanity.com/paper/2504.02792/) (2025)
- [Unified Video Action Model](https://mlsanity.com/paper/2503.00200/) (2025)
- [Consistency Policy: Accelerated Visuomotor Policies via Consistency Distillation](https://mlsanity.com/paper/2405.07503/) (2024)
- [RVT-2: Learning Precise Manipulation from Few Demonstrations](https://mlsanity.com/paper/2406.08545/) (2024)
- [Learning to Act without Actions](https://mlsanity.com/paper/2312.10812/) (2023)
- [Steering Your Generalists: Improving Robotic Foundation Models via Value Guidance](https://mlsanity.com/paper/2410.13816/) (2024)
- [Imitating Latent Policies from Observation](https://mlsanity.com/paper/1805.07914/) (2018)
- [One-Step Diffusion Policy: Fast Visuomotor Policies via Diffusion Distillation](https://mlsanity.com/paper/2410.21257/) (2024)
- [Decomposing the Generalization Gap in Imitation Learning for Visual Robotic Manipulation](https://mlsanity.com/paper/2307.03659/) (2023)
- [Fighting Copycat Agents in Behavioral Cloning from Observation Histories](https://mlsanity.com/paper/2010.14876/) (2020)
- [ManiFlow: A General Robot Manipulation Policy via Consistency Flow Training](https://mlsanity.com/paper/2509.01819/) (2025)
- [A Survey on Causal Reinforcement Learning](https://mlsanity.com/paper/2302.05209/) (2023)
- [PerAct2: Benchmarking and Learning for Robotic Bimanual Manipulation Tasks](https://mlsanity.com/paper/2407.00278/) (2024)
- [Diff-DAgger: Uncertainty Estimation with Diffusion Policy for Robotic Manipulation](https://mlsanity.com/paper/2410.14868/) (2024)
- [Deep Koopman Operator with Control for Nonlinear Systems](https://mlsanity.com/paper/2202.08004/) (2022)

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

- [Dreaming the Unseen: World Model-regularized Diffusion Policy for Out-of-Distribution Robustness](https://mlsanity.com/paper/2603.21017/) (2026)

## Abstract (arXiv)

Integrating generative models with action chunking has shown significant promise in imitation learning for robotic manipulation. However, the existing diffusion-based paradigm often struggles to capture strong temporal dependencies across multiple steps, particularly when incorporating proprioceptive input. This limitation can lead to task failures, where the policy overfits to proprioceptive cues at the expense of capturing the visually derived features of the task. To overcome this challenge, we propose the Deep Koopman-boosted Dual-branch Diffusion Policy (D3P) algorithm. D3P introduces a dual-branch architecture to decouple the roles of different sensory modality combinations. The visual branch encodes the visual observations to indicate task progression, while the fused branch integrates both visual and proprioceptive inputs for precise manipulation. Within this architecture, when the robot fails to accomplish intermediate goals, such as grasping a drawer handle, the policy can dynamically switch to execute action chunks generated by the visual branch, allowing recovery to previously observed states and facilitating retrial of the task. To further enhance visual representation learning, we incorporate a Deep Koopman Operator module that captures structured temporal dynamics from visual inputs. During inference, we use the test-time loss of the generative model as a confidence signal to guide the aggregation of the temporally overlapping predicted action chunks, thereby enhancing the reliability of policy execution. In simulation experiments across six RLBench tabletop tasks, D3P outperforms the state-of-the-art diffusion policy by an average of 14.6%. On three real-world robotic manipulation tasks, it achieves a 15.0% improvement. Code: https://github.com/dianyeHuang/D3P.