Bootstrap State Representation using Style Transfer for Better Generalization in Deep Reinforcement Learning

TL;DR

This paper introduces Thinker, a style transfer-based bootstrapping method that enhances deep reinforcement learning agents' ability to generalize by removing confounding features through trajectory clustering and style translation.

Contribution

It presents a novel unsupervised style transfer approach for trajectory augmentation, improving generalization in deep RL beyond existing data augmentation techniques.

Findings

Thinker improves generalization in Procgen environments.

Trajectory style transfer reduces overfitting to training environments.

Method is applicable across various RL settings.

Abstract

Deep Reinforcement Learning (RL) agents often overfit the training environment, leading to poor generalization performance. In this paper, we propose Thinker, a bootstrapping method to remove adversarial effects of confounding features from the observation in an unsupervised way, and thus, it improves RL agents' generalization. Thinker first clusters experience trajectories into several clusters. These trajectories are then bootstrapped by applying a style transfer generator, which translates the trajectories from one cluster's style to another while maintaining the content of the observations. The bootstrapped trajectories are then used for policy learning. Thinker has wide applicability among many RL settings. Experimental results reveal that Thinker leads to better generalization capability in the Procgen benchmark environments compared to base algorithms and several data augmentation techniques.