π€ AI Summary
This work addresses the low sample efficiency and poor generalization of reinforcement learning in high-dimensional, nonlinear legged robot control by proposing a novel approach that integrates morphological symmetry with Koopman-based dynamic modeling. The method employs an autoencoder to learn the systemβs Koopman operator, whose predictions are leveraged as privileged observations for the critic. Concurrently, group-theoretic symmetry is embedded as an inductive bias within the policy network architecture to enable equivariant policy learning. By uniquely combining symmetry priors with Koopman-based forecasting, the approach significantly enhances both training efficiency and cross-environment transferability. Empirical results demonstrate substantially reduced convergence times and improved final rewards across a range of complex bipedal locomotion tasks.
π Abstract
Reinforcement learning (RL) algorithms classically suffer from poor sample efficiency. In robotics, a recent line of work has emerged addressing this problem by encoding physics priors in the learning process. However, most of these approaches are validated on well-defined, low-dimensional benchmark systems rather than high-dimensional robots with complex nonlinear dynamics. In this paper, we introduce \textit{SKooP (Symmetric Koopman Predictions)}, an approach combining the advantages of morphological symmetries with those of a Koopman model learned via autoencoder to enhance policy learning. SKooP learns a Koopman model of the system dynamics alongside the policy. The resulting Koopman predictions are used as privileged observations for the critic, allowing the agent to learn based on smoother, more informative features. We also incorporate group symmetries into the actor, critic, encoder and decoder networks to produce a highly equivariant policy. The SKooP approach is validated via in-depth analysis of the learned Koopman models and symmetric policies to showcase how each of these influences the agent's performance. We also show that the learned policies are transferable to different simulation environments. Our results show that SKooP consistently reduces convergence time and increases the learned reward for multiple challenging bipedal locomotion tasks on a quadruped robot. Project page: https://evelyd.github.io/SymmetricKoopmanPredictions/