RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability

TL;DR

RePo introduces a resilient visual model-based RL method that regularizes the latent representation to be predictive of dynamics and reward while reducing sensitivity to irrelevant visual variations, enabling better performance in dynamic environments.

Contribution

The paper proposes a novel regularization approach for learning resilient latent representations in visual model-based RL, including a test-time adaptation method for distribution shifts.

Findings

Enhanced robustness to visual distractors in simulation benchmarks

Effective test-time adaptation to diverse environments

Successful application to real-world egocentric navigation

Abstract

Visual model-based RL methods typically encode image observations into low-dimensional representations in a manner that does not eliminate redundant information. This leaves them susceptible to spurious variations -- changes in task-irrelevant components such as background distractors or lighting conditions. In this paper, we propose a visual model-based RL method that learns a latent representation resilient to such spurious variations. Our training objective encourages the representation to be maximally predictive of dynamics and reward, while constraining the information flow from the observation to the latent representation. We demonstrate that this objective significantly bolsters the resilience of visual model-based RL methods to visual distractors, allowing them to operate in dynamic environments. We then show that while the learned encoder is resilient to spirious variations, it is not invariant under significant distribution shift. To address this, we propose a simple reward-free alignment procedure that enables test time adaptation of the encoder. This allows for quick adaptation to widely differing environments without having to relearn the dynamics and policy. Our effort is a step towards making model-based RL a practical and useful tool for dynamic, diverse domains. We show its effectiveness in simulation benchmarks with significant spurious variations as well as a real-world egocentric navigation task with noisy TVs in the background. Videos and code at https://zchuning.github.io/repo-website/.