# What Do You Need for Compositional Generalization in Diffusion Planning?

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for offline RL, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 60 of its 79 references that are in the map.
- First posted on arXiv 23 May 2025.

## Placed in

- Offline RL and imitation learning from data / [Diffusion and flow policies for offline RL](https://mlsanity.com/offlinerl/diffusion-and-flow-policies-for-offline-rl/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Decoupled Weight Decay Regularization](https://mlsanity.com/paper/1711.05101/) (2017)
- [π0.5: a Vision-Language-Action Model with Open-World Generalization](https://mlsanity.com/paper/2504.16054/) (2025)
- [Make-A-Video: Text-to-Video Generation without Text-Video Data](https://mlsanity.com/paper/2209.14792/) (2022)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [RePaint: Inpainting using Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2201.09865/) (2022)
- [GR00T N1: An Open Foundation Model for Generalist Humanoid Robots](https://mlsanity.com/paper/2503.14734/) (2025)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion](https://mlsanity.com/paper/2407.01392/) (2024)
- [VideoCrafter1: Open Diffusion Models for High-Quality Video Generation](https://mlsanity.com/paper/2310.19512/) (2023)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [robosuite: A Modular Simulation Framework and Benchmark for Robot Learning](https://mlsanity.com/paper/2009.12293/) (2020)
- [Making Convolutional Networks Shift-Invariant Again](https://mlsanity.com/paper/1904.11486/) (2019)
- [Causal Confusion in Imitation Learning](https://mlsanity.com/paper/1905.11979/) (2019)
- [Demystifying CLIP Data](https://mlsanity.com/paper/2309.16671/) (2023)
- [Imitating Human Behaviour with Diffusion Models](https://mlsanity.com/paper/2301.10677/) (2023)
- [OGBench: Benchmarking Offline Goal-Conditioned RL](https://mlsanity.com/paper/2410.20092/) (2024)
- [RvS: What is Essential for Offline RL via Supervised Learning?](https://mlsanity.com/paper/2112.10751/) (2021)
- [Generalization in diffusion models arises from geometry-adaptive harmonic representations](https://mlsanity.com/paper/2310.02557/) (2023)
- [Deep Reinforcement Learning and the Deadly Triad](https://mlsanity.com/paper/1812.02648/) (2018)
- [History-Guided Video Diffusion](https://mlsanity.com/paper/2502.06764/) (2025)

## Abstract (arXiv)

In policy learning, stitching and compositional generalization refer to the extent to which the policy is able to piece together sub-trajectories of data it is trained on to generate new and diverse behaviours. While stitching has been identified as a significant strength of offline reinforcement learning, recent generative behavioural cloning (BC) methods have also shown proficiency at stitching. However, the main factors behind this are poorly understood, hindering the development of new algorithms that can reliably stitch by design. Focusing on diffusion planners trained via generative behavioural cloning, and without resorting to dynamic programming or TD-learning, we find three properties are key enablers for composition: shift equivariance, local receptive fields, and inference choices. We use these properties to explain architecture, data, and inference choices in existing generative BC methods based on diffusion planning including replanning frequency, data augmentation, and data scaling. Our experiments show that while local receptive fields are more important than shift equivariance in creating a diffusion planner capable of composition, both are crucial. Using findings from our experiments, we develop a new architecture for diffusion planners called Eq-Net, that is simple, produces diverse trajectories competitive with more computationally expensive methods such as replanning or scaling data, and can be guided to enable generalization in goal-conditioned settings. We show that Eq-Net exhibits significant compositional generalization in a variety of navigation and manipulation tasks designed to test planning diversity.