Generalization in Visual Reinforcement Learning with the Reward Sequence Distribution

TL;DR

This paper introduces RSD-OA, a reward sequence distribution method conditioned on actions, which improves generalization in visual reinforcement learning by focusing on task-relevant information and ignoring visual distractions.

Contribution

The paper proposes RSD-OA, a novel reward sequence distribution approach that is invariant to visual distractions and captures long-term task-relevant information for better generalization.

Findings

RSD-OA significantly outperforms state-of-the-art methods on DeepMind Control tasks with visual distractions.

The approach improves generalization to unseen environments in VRL.

RSD-OA effectively isolates task-relevant information from visual noise.

Abstract

Generalization in partially observed markov decision processes (POMDPs) is critical for successful applications of visual reinforcement learning (VRL) in real scenarios. A widely used idea is to learn task-relevant representations that encode task-relevant information of common features in POMDPs, i.e., rewards and transition dynamics. As transition dynamics in the latent state space -- which are task-relevant and invariant to visual distractions -- are unknown to the agents, existing methods alternatively use transition dynamics in the observation space to extract task-relevant information in transition dynamics. However, such transition dynamics in the observation space involve task-irrelevant visual distractions, degrading the generalization performance of VRL methods. To tackle this problem, we propose the reward sequence distribution conditioned on the starting observation and the predefined subsequent action sequence (RSD-OA). The appealing features of RSD-OA include that: (1) RSD-OA is invariant to visual distractions, as it is conditioned on the predefined subsequent action sequence without task-irrelevant information from transition dynamics, and (2) the reward sequence captures long-term task-relevant information in both rewards and transition dynamics. Experiments demonstrate that our representation learning approach based on RSD-OA significantly improves the generalization performance on unseen environments, outperforming several state-of-the-arts on DeepMind Control tasks with visual distractions.