Symmetry-Aware Steering of Equivariant Diffusion Policies: Benefits and Limits

TL;DR

This paper demonstrates that symmetry-aware diffusion policies can be effectively steered using reinforcement learning, improving sample efficiency and policy performance, especially in tasks with geometric symmetries, while also discussing the limits of strict equivariance.

Contribution

The paper provides a theoretical foundation for equivariant diffusion policies and introduces a symmetry-aware steering framework that enhances RL performance with geometric symmetries.

Findings

Symmetry-aware RL improves sample efficiency in diffusion policies.

Exploiting symmetry prevents value divergence during training.

Strict equivariance has practical limits under symmetry-breaking conditions.

Abstract

Equivariant diffusion policies (EDPs) combine the generative expressivity of diffusion models with the strong generalization and sample efficiency afforded by geometric symmetries. While steering these policies with reinforcement learning (RL) offers a promising mechanism for fine-tuning beyond demonstration data, directly applying standard (non-equivariant) RL can be sample-inefficient and unstable, as it ignores the symmetries that EDPs are designed to exploit. In this paper, we theoretically establish that the diffusion process of an EDP is equivariant, which in turn induces a group-invariant latent-noise MDP that is well-suited for equivariant diffusion steering. Building on this theory, we introduce a principled symmetry-aware steering framework and compare standard, equivariant, and approximately equivariant RL strategies through comprehensive experiments across tasks with varying degrees of symmetry. While we identify the practical boundaries of strict equivariance under symmetry breaking, we show that exploiting symmetry during the steering process yields substantial benefits-enhancing sample efficiency, preventing value divergence, and achieving strong policy improvements even when EDPs are trained from extremely limited demonstrations.