Our research can be divided into several topic areas.

1. Theory of machine learning through the lens of symmetry including generalization and approximation error for equivariant neural networks, study of approximate and extrinsic symmetry in deep learning, symmetries of neural network parameters spaces and their effects on optimization, structure of equivariant neural networks.

2. design of equivariant neural networks for scientific and engineering applications including fluid dynamics, radar signal processing, astrophysical simulation, radar signal processing, and materials and thin films.

3. Equivariant methods for deep reinforcement learning with an emphasis on robotic learning and perception. Our goal are sample efficient algorithms for real world robot learning.

4. Design models with can utilize approximate and unknown symmetry. These include symmetry discovery methods, relaxed equivariant methods, applications of equivariant models in cases of extrinsic symmetry, and methods which can learn group actions in cases of latent symmetry.

Selected Publications

2023

NeurIPS 2023

Modeling Dynamics over Meshes with Gauge Equivariant Nonlinear Message Passing

Jung Yeon Park, Lawson L.S. Wong, Robin Walters

Conference on Neural Information Processing Systems 2023

Abstract PDF Code Webpage

Abstract: Data over non-Euclidean manifolds, often discretized as surface meshes, naturally arise in computer graphics and biological and physical systems. In particular, solutions to partial differential equations (PDEs) over manifolds depend critically on the underlying geometry. While graph neural networks have been successfully applied to PDEs, they do not incorporate surface geometry and do not consider local gauge symmetries of the manifold. Alternatively, recent works on gauge equivariant convolutional and attentional architectures on meshes leverage the underlying geometry but underperform in modeling surface PDEs with complex nonlinear dynamics. To address these issues, we introduce a new gauge equivariant architecture using nonlinear message passing. Our novel architecture achieves higher performance than either convolutional or attentional networks on domains with highly complex and nonlinear dynamics. However, similar to the non-mesh case, design trade-offs favor convolutional, attentional, or message passing networks for different tasks; we investigate in which circumstances our message passing method provides the most benefit.

2023

NeurIPS 2023

A General Theory of Correct, Incorrect, and Extrinsic Equivariance

Dian Wang, Xupeng Zhu, Jung Yeon Park, Mingxi Jia, Guanang Su, Robert Platt, Robin Walters

Conference on Neural Information Processing Systems 2023

Abstract PDF Code

Abstract: Although equivariant machine learning has proven effective at many tasks, success depends heavily on the assumption that the ground truth function is symmetric over the entire domain matching the symmetry in an equivariant neural network. A missing piece in the equivariant learning literature is the analysis of equivariant networks when symmetry exists only partially in the domain. In this work, we present a general theory for such a situation. We propose pointwise definitions of correct, incorrect, and extrinsic equivariance, which allow us to quantify continu- ously the degree of each type of equivariance a function displays. We then study the impact of various degrees of incorrect or extrinsic symmetry on model error. We prove error lower bounds for invariant or equivariant networks in classification or regression settings with partially incorrect symmetry. We also analyze the poten- tially harmful effects of extrinsic equivariance. Experiments validate these results in three different environments.

2023

ICLR 2023

The Surprising Effectiveness of Equivariant Models in Domains with Latent Symmetry

Dian Wang, Jung Yeon Park, Neel Sortur, Lawson L.S. Wong, Robin Walters*, Robert Platt*

International Conference on Learning Representations 2023

Abstract PDF Webpage

Abstract: Extensive work has demonstrated that equivariant neural networks can significantly improve sample efficiency and generalization by enforcing an inductive bias in the network architecture. These applications typically assume that the domain symmetry is fully described by explicit transformations of the model inputs and outputs. However, many real-life applications contain only latent or partial symmetries which cannot be easily described by simple transformations of the input. In these cases, it is necessary to learn symmetry in the environment instead of imposing it mathematically on the network architecture. We discover, surprisingly, that imposing equivariance constraints that do not exactly match the domain symmetry is very helpful in learning the true symmetry in the environment. We differentiate between extrinsic and incorrect symmetry constraints and show that while imposing incorrect symmetry can impede the model's performance, imposing extrinsic symmetry can actually improve performance. We demonstrate that an equivariant model can significantly outperform non-equivariant methods on domains with latent symmetries both in supervised learning and in reinforcement learning for robotic manipulation and control problems.

2023

ICRA 2023

SEIL: Simulation-augmented Equivariant Imitation Learning

Mingxi Jia*, Dian Wang*, Guanang Su, David Klee, Xupeng Zhu, Robin Walters, Robert Platt

International Conference on Robotics and Automation 2023

Abstract PDF Webpage

Abstract: In robotic manipulation, acquiring samples is extremely expensive because it often requires interacting with the real world. Traditional image-level data augmentation has shown the potential to improve sample efficiency in various machine learning tasks. However, image-level data augmentation is insufficient for an imitation learning agent to learn good manipulation policies in a reasonable amount of demonstrations. We propose Simulation-augmented Equivariant Imitation Learning (SEIL), a method that combines a novel data augmentation strategy of supplementing expert trajectories with simulated transitions and an equivariant model that exploits the O(2) symmetry in robotic manipulation. Experimental evaluations demonstrate that our method can learn non-trivial manipulation tasks within ten demonstrations and outperforms the baselines with a significant margin.

2023

ICRA 2022

Edge Grasp Network: Graph-Based SE(3)-invariant Approach to Grasp Detection

Haojie Huang, Dian Wang, Xupend Zhu, Robin Walters, Robert Platt

International Conference on Robotics and Automation 2023

Abstract PDF Webpage

Abstract: Given point cloud input, the problem of 6-DoF grasp pose detection is to identify a set of hand poses in SE(3) from which an object can be successfully grasped. This important problem has many practical applications. Here we propose a novel method and neural network model that enables better grasp success rates relative to what is available in the literature. The method takes standard point cloud data as input and works well with single-view point clouds observed from arbitrary viewing directions.

2023

Under Review

Fast and Expressive Gesture Recognition using a Combination-Homomorphic Electromyogram Encoder

Niklas Smedemark-Margulies*, Yunus Bicer*, Elifnur Sunger, Tales Imbiriba, Eugene Tunik, Deniz Erdogmus, Mathew Yarossi, Robin Walters

Under Review

Abstract PDF Code Dataset

We study the task of gesture recognition from electromyography (EMG), with the goal of enabling expressive human-computer interaction at high accuracy, while minimizing the time required for new subjects to provide calibration data. To fulfill these goals, we define combination gestures consisting of a direction component and a modifier component. New subjects only demonstrate the single component gestures and we seek to extrapolate from these to all possible single or combination gestures. We extrapolate to unseen combination gestures by combining the feature vectors of real single gestures to produce synthetic training data. This strategy allows us to provide a large and flexible gesture vocabulary, while not requiring new subjects to demonstrate combinatorially many example gestures. We pre-train an encoder and a combination operator using self-supervision, so that we can produce useful synthetic training data for unseen test subjects. To evaluate the proposed method, we collect a real-world EMG dataset, and measure the effect of augmented supervision against two baselines: a partially-supervised model trained with only single gesture data from the unseen subject, and a fully-supervised model trained with real single and real combination gesture data from the unseen subject. We find that the proposed method provides a dramatic improvement over the partially-supervised model, and achieves a useful classification accuracy that in some cases approaches the performance of the fully-supervised model.

2022

ICML 2022

Learning Symmetric Embedding Networks for Equivariant World Models

Jung Yeon Park, Ondrej Biza, Linfeng Zhao, Jan-Willem van de Meent, Robin Walters

We propose learning symmetric embedding networks (SENs) for input spaces where we do not know the effects of symmetry transformations, to a feature space where the transformation is known. This network can be combined with downstream task-specific equivariant networks and trained end-to-end in latent space. SENs can extend the applicability of equivariant networks to a wider variety of domains.

Abstract PDF Code

Abstract: Incorporating symmetries can lead to highly data-efficient and generalizable models by defining equivalence classes of data samples related by transformations. However, characterizing how transformations act on input data is often difficult, limiting the applicability of equivariant models. We propose learning symmetric embedding networks (SENs) that encode an input space (e.g. images), where we do not know the effect of transformations (e.g. rotations), to a feature space that transforms in a known manner under these operations. This network can be trained end-to-end with an equivariant task network to learn an explicitly symmetric representation. We validate this approach in the context of equivariant transition models with 3 distinct forms of symmetry. Our experiments demonstrate that SENs facilitate the application of equivariant networks to data with complex symmetry representations. Moreover, doing so can yield improvements in accuracy and generalization relative to both fully-equivariant and non-equivariant baselines.

2022

PMLR Vol. 197

Image to Icosahedral Projection for SO(3) Object Reasoning from Single-View Images

David M. Klee, Ondrej Biza, Robert Platt, Robin Walters

We develop a hybrid equivariant model that incorporates SO(2) and SO(3) equivariant convolution layers to improve 3D reasoning from 2D images. Our method outperforms baselines on shape classification and pose prediction tasks, especially in the low-data regime.

Abstract PDF Code Webpage

Abstract: Reasoning about 3D objects based on 2D images is challenging due to variations in appearance caused by viewing the object from different orientations. Tasks such as object classification are invariant to 3D rotations and other such as pose estimation are equivariant. However, imposing equivariance as a model constraint is typically not possible with 2D image input because we do not have an a priori model of how the image changes under out-of-plane object rotations. The only SO(3)-equivariant models that currently exist require point cloud or voxel input rather than 2D images. In this paper, we propose a novel architecture based on icosahedral group convolutions that reasons in SO(3) by learning a projection of the input image onto an icosahedron. The resulting model is approximately equivariant to rotation in SO(3). We apply this model to object pose estimation and shape classification tasks and find that it outperforms reasonable baselines.

2022

RLDM 2022

Understanding the Mechanism behind Data Augmentation's Success on Image-based RL

David M. Klee, Robin Walters, Robert Platt

Data augmentation is known to provide substantial benefits for image-based reinforcement learning (RL) but the mechanism is not clear. We show that data augmentation increases both the equivariance and invariance of the convolutional encoder, e.g. the feature map is spatially smooth. We show that a simple Gaussian blur operation can achieve the same effect for some of the tested environments.

Abstract PDF Poster

Abstract: Reinforcement learning for continuous control tasks is challenging with image observations, due to the representation learning problem. A series of recent work has shown that augmenting the observations via random shifts during training significantly improves performance, even matching state-based methods. However, it is not well-understood why augmentation is so beneficial; since the method uses a nearly-shift equivariant convolutional encoder, shifting the input should have little impact on what features are learned. In this work, we investigate why random shifts are useful augmentations for image-based RL and show that it increases both the shift-equivariance and shift-invariance of the encoder. In other words, the visual features learned exhibit spatial continuity, which we show can be partially achieved using dropout. We hypothesize that the spatial continuity of the visual encoding simplifies learning for the subsequent linear layers in the actor-critic networks.

2022

ICLR 2022

SO(2)-Equivariant Reinforcement Learning

Dian Wang, Robin Walters, Robert Platt

International Conference on Learning Representations 2022

Abstract PDF Code Webpage

Abstract: Equivariant neural networks enforce symmetry within the structure of their convolutional layers, resulting in a substantial improvement in sample efficiency when learning an equivariant or invariant function. Such models are applicable to robotic manipulation learning which can often be formulated as a rotationally symmetric problem. This paper studies equivariant model architectures in the context of Q-learning and actor-critic reinforcement learning. We identify equivariant and invariant characteristics of the optimal Q-function and the optimal policy and propose equivariant DQN and SAC algorithms that leverage this structure. We present experiments that demonstrate that our equivariant versions of DQN and SAC can be significantly more sample efficient than competing algorithms on an important class of robotic manipulation problems.

2022

RSS 2022

Equivariant Transporter Network

Haojie Huang, Dian Wang, Robin Walters, Robert Platt

Robotics: Science and Systems 2022

Abstract PDF Code Webpage

Abstract: Transporter Net is a recently proposed framework for pick and place that is able to learn good manipulation policies from a very few expert demonstrations [35]. A key reason why Transporter Net is so sample efficient is that the model incorporates rotational equivariance into the pick-conditioned place module, i.e. the model immediately generalizes learned pick-place knowledge to objects presented in different pick orientations. This paper proposes a novel version of Transporter Net that is equivariant to both pick and place orientation. As a result, our model immediately generalizes pick-place knowledge to different place orientations in addition to generalizing the pick orientation as before. Ultimately, our new model is more sample efficient and achieves better pick and place success rates than the baseline Transporter Net model.

2022

CoRL 2022

On-Robot Learning With Equivariant Models

Dian Wang, Mingxi Jia, Xupeng Zhu, Robin Walters, Robert Platt

Conference on Robot Learning 2022

Abstract PDF Code Webpage

Abstract: Recently, equivariant neural network models have been shown to improve sample efficiency for tasks in computer vision and reinforcement learning. This paper explores this idea in the context of on-robot policy learning in which a policy must be learned entirely on a physical robotic system without reference to a model, a simulator, or an offline dataset. We focus on applications of Equivariant SAC to robotic manipulation and explore a number of variations of the algorithm. Ultimately, we demonstrate the ability to learn several non-trivial manipulation tasks completely through on-robot experiences in less than an hour or two of wall clock time.

2021

CoRL 2021

Equivariant Q Learning in Spatial Action Spaces

Dian Wang, Robin Walters, Xupeng Zhu, Robert Platt

Conference on Robot Learning 2021

Abstract PDF Code Webpage

Abstract: Recently, a variety of new equivariant neural network model architectures have been proposed that generalize better over rotational and reflectional symmetries than standard models. These models are relevant to robotics because many robotics problems can be expressed in a rotationally symmetric way. This paper focuses on equivariance over a visual state space and a spatial action space - the setting where the robot action space includes a subset of SE(2). In this situation, we know a priori that rotations and translations in the state image should result in the same rotations and translations in the spatial action dimensions of the optimal policy. Therefore, we can use equivariant model architectures to make Q learning more sample efficient. This paper identifies when the optimal Q function is equivariant and proposes Q network architectures for this setting. We show experimentally that this approach outperforms standard methods in a set of challenging manipulation problems.