What Do You Need for Compositional Generalization in Diffusion Planning?

TL;DR

This paper investigates the key properties enabling compositional generalization in diffusion-based policy planning, introduces a new architecture called Eq-Net, and demonstrates its effectiveness across navigation and manipulation tasks.

Contribution

It identifies shift equivariance, local receptive fields, and inference choices as crucial for composition, and proposes Eq-Net, a simple architecture that achieves strong generalization without complex re-planning.

Findings

Local receptive fields are more important than shift equivariance for composition.

Eq-Net achieves diverse, goal-conditioned planning comparable to more complex methods.

The properties enable effective architecture, data, and inference design for diffusion planners.

Abstract

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.