Title: Accelerating Model-Based Reinforcement Learning Using Equivariance
Primary area: reinforcement learning
Keywords: Reinforcement Learning, Model-based Reinforcement Learning, World models, Equivariance

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.

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