Multi-Group Equivariant Augmentation for Reinforcement Learning in Robot Manipulation

TL;DR

This paper introduces Multi-Group Equivariance Augmentation (MEA), a novel data augmentation method leveraging non-isometric symmetries in visuomotor reinforcement learning, improving sampling efficiency in robotic manipulation tasks.

Contribution

The work extends symmetry-based augmentation to non-isometric transformations, formulates a new POMDP incorporating these symmetries, and demonstrates improved efficiency in simulation and real robots.

Findings

Enhanced sampling efficiency in robotic manipulation tasks.

Effective integration of MEA with offline reinforcement learning.

Successful real-robot experiments across multiple domains.

Abstract

Sampling efficiency is critical for deploying visuomotor learning in real-world robotic manipulation. While task symmetry has emerged as a promising inductive bias to improve efficiency, most prior work is limited to isometric symmetries -- applying the same group transformation to all task objects across all timesteps. In this work, we explore non-isometric symmetries, applying multiple independent group transformations across spatial and temporal dimensions to relax these constraints. We introduce a novel formulation of the partially observable Markov decision process (POMDP) that incorporates the non-isometric symmetry structures, and propose a simple yet effective data augmentation method, Multi-Group Equivariance Augmentation (MEA). We integrate MEA with offline reinforcement learning to enhance sampling efficiency, and introduce a voxel-based visual representation that preserves translational equivariance. Extensive simulation and real-robot experiments across two manipulation domains demonstrate the effectiveness of our approach.