Learning Parsimonious Dynamics for Generalization in Reinforcement Learning

TL;DR

This paper introduces a model that learns simple, interpretable latent dynamics in reinforcement learning, inspired by human reasoning, to improve generalization and planning in new environments.

Contribution

It proposes a novel approach to learn parsimonious latent dynamics using locally linear transformations and a variational objective, enhancing generalization in RL.

Findings

Improved generalization in unseen environments.

Effective planning using learned latent dynamics.

Parsimonious models require fewer transformations.

Abstract

Humans are skillful navigators: We aptly maneuver through new places, realize when we are back at a location we have seen before, and can even conceive of shortcuts that go through parts of our environments we have never visited. Current methods in model-based reinforcement learning on the other hand struggle with generalizing about environment dynamics out of the training distribution. We argue that two principles can help bridge this gap: latent learning and parsimonious dynamics. Humans tend to think about environment dynamics in simple terms -- we reason about trajectories not in reference to what we expect to see along a path, but rather in an abstract latent space, containing information about the places' spatial coordinates. Moreover, we assume that moving around in novel parts of our environment works the same way as in parts we are familiar with. These two principles work together in tandem: it is in the latent space that the dynamics show parsimonious characteristics. We develop a model that learns such parsimonious dynamics. Using a variational objective, our model is trained to reconstruct experienced transitions in a latent space using locally linear transformations, while encouraged to invoke as few distinct transformations as possible. Using our framework, we demonstrate the utility of learning parsimonious latent dynamics models in a range of policy learning and planning tasks.