Accelerating Model-Based Reinforcement Learning Using Equivariance

Ahmed Agha, Haojie Huang, Dian Wang, Christopher Amato, Robert Platt

Keywords: Reinforcement Learning, Model-based Reinforcement Learning, World models, Equivariance

TL;DR: We combine model-based reinforcement learning with equivariant neural networks to achieve better generalization and solving tasks with fewer environment interactions.

Abstract: Model-based reinforcement learning (MBRL) is a promising approach for learning effective policies in a data-efficient manner by using learned dynamics models to generate synthetic rollouts for actor-critic trianing, thereby reducing the reliance on costly environment interactions. However, when the learned dynamics model is inaccurate, these synthetic rollouts can introduce bias and deteriorate performance. Fortunately, many domains exhibit symmetries that can serve as powerful inductive biases, enabling the learned models to generalize beyond their training data. In this work, we exploit these inherent symmetries in MBRL and formally define equivariant MBRL for POMDPs. Building on this formulation, we introduce EquiDreamer, a framework that integrates symmetry into both world modeling and policy learning through an equivariant latent dynamics architecture. Experiments on visual continuous control tasks demonstrate that our equivariant MBRL method outperforms both model-based and model-free baselines, achieving strong results with substantially fewer environment interactions.

Primary Area: reinforcement learning

Submission Number: 9962