EDGI: Equivariant Diffusion for Planning with Embodied Agents

Embodied agents operate in a structured world, often solving tasks with spatial, temporal, and permutation symmetries. Most algorithms for planning and modelbased reinforcement learning (MBRL) do not take this rich geometric structure into account, leading to sample inefficiency and poor generalization. We introduce the Equivariant Diffuser for Generating Interactions (EDGI), an algorithm for MBRL and planning that is equivariant with respect to the product of the spatial symmetry group SE(3), the discrete-time translation group Z, and the object permutation group S n . EDGI follows the Diffuser framework by Janner et al. [2022] in treating both learning a world model and planning in it as a conditional generative modeling problem, training a diffusion model on an offline trajectory dataset. We introduce a new SE(3) × Z × S n -equivariant diffusion model that supports multiple representations. We integrate this model in a planning loop, where conditioning and classifier guidance let us softly break the symmetry for specific tasks as needed. On object manipulation and navigation tasks, EDGI is substantially more sample efficient and generalizes better across the symmetry group than non-equivariant models. * Equal contribution, order determined through a game of table tennis † Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc. ‡ Work done during an internship at Qualcomm AI Research 4 This is true in the approximately flat spacetime on Earth, as long as all velocities are much smaller than the speed of light. A machine learning researcher who finds herself close to a black hole may disagree. 37th Conference on Neural Information Processing Systems (NeurIPS 2023).