Exploration Implies Data Augmentation: Reachability and Generalisation in Contextual MDPs

TL;DR

This paper introduces Explore-Go, an exploration method that enhances generalisation in zero-shot policy transfer for contextual MDPs by increasing state coverage and accuracy, applicable to various RL algorithms.

Contribution

The paper proposes Explore-Go, a simple exploration phase that improves generalisation in contextual MDPs by balancing coverage and value function accuracy.

Findings

Explore-Go significantly improves generalisation across multiple environments.

Combining Explore-Go with existing RL algorithms yields better transfer performance.

Increased exploration coverage enhances policy robustness in partially observable settings.

Abstract

In the zero-shot policy transfer (ZSPT) setting for contextual Markov decision processes (MDP), agents train on a fixed set of contexts and must generalise to new ones. Recent work has argued and demonstrated that increased exploration can improve this generalisation, by training on more states in the training contexts. In this paper, we demonstrate that training on more states can indeed improve generalisation, but can come at a cost of reducing the accuracy of the learned value function which should not benefit generalisation. We hypothesise and demonstrate that using exploration to increase the agent's coverage while also increasing the accuracy improves generalisation even more. Inspired by this, we propose a method Explore-Go that implements an exploration phase at the beginning of each episode, which can be combined with existing on- and off-policy RL algorithms and significantly improves generalisation even in partially observable MDPs. We demonstrate the effectiveness of Explore-Go when combined with several popular algorithms and show an increase in generalisation performance across several environments. With this, we hope to provide practitioners with a simple modification that can improve the generalisation of their agents.