Learning Abstract World Models with a Group-Structured Latent Space

Thomas Delliaux, Nguyen-Khanh Vu, Vincent Francois-Lavet, Elise van der Pol, Emmanuel Rachelson

Keywords: representation learning, reinforcement learning, geometric prior, abstract world model, model-based reinforcement learning

Abstract: Learning meaningful abstract models of Markov Decision Processes (MDPs) is crucial for improving generalization from limited data. In this work, we show how geometric priors can be imposed on the low-dimensional representation manifold of a learned transition model. We incorporate known symmetric structures via appropriate choices of the latent space and the associated group actions, which encode prior knowledge about invariances in the environment. In addition, our framework allows the embedding of additional unstructured information alongside these symmetries. We show experimentally that this leads to better predictions of the latent transition model than fully unstructured approaches, as well as quicker learning on downstream RL tasks, in environments with rotational and translational features, including in first-person views of 3D environments. Additionally, our experiments show that this leads to simpler and more disentangled representations.

Confirmation: I understand that authors of each paper submitted to EWRL may be asked to review 2-3 other submissions to EWRL.

Serve As Reviewer: ~Thomas_Delliaux1

Track: Regular Track: unpublished work

Submission Number: 116