Learning Generalizable Representations for Reinforcement Learning via Adaptive Meta-learner of Behavioral Similarities

1 INTRODUCTION . Designing effective reinforcement learning algorithms to learn control from high-dimensional visual observations is crucial and has attracted more attentions in the reinforcement learning ( RL ) community ( Yarats et al. , 2021 ; Laskin et al. , 2020a ; Schwarzer et al. , 2021 ; Lesort et al. , 2018 ) . To learn a policy efficiently from high-dimensional observations , prior approaches first learned an encoder to map high-dimensional observations , e.g. , images , to low-dimensional representations , and subsequently trained a policy from low-dimensional representations to actions based on various RL algorithms . Therefore , how to learn low-dimensional representations which are able to provide semantic abstraction for high-dimensional observations plays a key role . Early works on deep reinforcement learning train encoders based on various reconstruction losses ( Lange & Riedmiller , 2010 ; Watter et al. , 2015 ; Wahlström et al. , 2015 ; Higgins et al. , 2017 ) , which aim to enforce that the learned low-dimensional representations can be decoded to reconstruct the original high-dimensional observations . Promising results have been achieved in some application domains , such as playing video games , simulated tasks , etc . However , the policy learned on representations of states trained with reconstruction losses may not generalize well to complex environments , which are even though semantically similar to the source environments . The reason is that a reconstruction loss is computed over all the pixels , which results in all the details of the high-dimensional images being preserved in low-dimensional representations . However , some of the observed details , e.g. , complex background objects in an image , are task-irrelevant and highly environment-dependent . Encoding such details in representation learning makes the learned representation less effective for a downstream reinforcement learning task of interest and less generalizable to new environments . To address the aforementioned issue , some data augmentation techniques have been proposed to make the learned representations more robust ( Yarats et al. , 2021 ; Lee et al. , 2020b ; Laskin et al. , 2020b ) . However , these approaches rarely consider the properties of a Markov Decision Process ( MDP ) , such as conditional transition probabilities between states given an action , in the representation learning procedure . Recent research ( Zhang et al. , 2020 ; Agarwal et al. , 2021 ; Castro , 2020 ; Ferns et al. , 2004 ) has shown that the bisimulation metric and its variants are potentially effective to be exploited to learn a more generalizable reinforcement learning agent across semantically similar environments . The bisimulation metric measures the “ behavioral similarity ” between two states based on two terms : 1 ) reward difference that considers the difference in immediate task-specific reward signals between two states , and 2 ) dynamics difference that considers the similarity of the long-term behaviors between two states . A general idea of these approaches is to learn an encoder to map observations to a latent space such that the distance or similarity between two states in the latent space approximates their bisimulation metric . In this way , the learned representations ( in the latent space ) are task-relevant and invariant to task-irrelevant details . However , manually specifying a form of distance , e.g. , the L1 norm as used in ( Zhang et al. , 2020 ) , in the latent space may limit the approximation precision for the bisimulation metric and potentially lose parts of the state features that are useful for policy learning . Moreover , existing approaches rarely explore how to learn a linear combination of reward and dynamics differences in the bisimulation metric , which may vary in different tasks or environments . We propose a novel framework for learning generalizable state representation in RL , called Adaptive Meta-learner of Behavioral Similarities ( AMBS ) . In this framework we design a network with two encoders that map the high-dimensional observations to two decomposed representations regarding the two properties of MDP , rewards and dynamics . For the purpose of learning behavioral similarity in state representations , we introduce a pair of meta-learners that learn similarities in order to measure the reward and dynamics similarity between two states over the corresponding decomposed state representations , respectively . The meta-learners are self-learned by approximating the corresponding terms in the bisimulation metric , containing the reward difference and the dynamics difference . Then the meta-learners update the state representations according to their behavioral distance to the other state representations . Unlike previous approaches with a hand-crafted form of distance/similarity evaluating state encoding in L1 space , our learned similarities measure two state representations via a neural architecture , where more side information can be preserved for further policy learning . Our experiments also observe that a smaller loss can be obtained by using the learned similarity . This demonstrates that the proposed meta-learners can overcome the approximation precision issue and provide more stable gradients for robust learning of state representations . Moreover , we explore the impact between reward and dynamics terms in the bisimulation metric . We propose a learning-based adaptive strategy to balance the effect of reward and dynamics in different tasks or environments by introducing a learnable importance parameter , which is jointly learned with the state-action value function . Finally , we use a simple but effective data augmentation strategy to accelerate the RL procedure and help learning more robust state representations . The main contributions of our work are 3-fold : 1 ) we propose a meta-learner-based framework to learn task-relevant and environment-details-invariant state representations ; 2 ) we propose a network architecture that decomposes each state into two different types of representations for measuring the similarities in terms of reward and dynamics , respectively , and design a learnable adaptive strategy to balance them to estimate the bisimulation metric between states ; 3 ) we verify our proposed framework on extensive experiments and demonstrate new state-of-the-art results on backgrounddistraction DeepMind control suite ( Tassa et al. , 2018 ; Zhang et al. , 2018 ; Stone et al. , 2021 ) and other visual-based RL tasks . 2 RELATED WORK . Representation learning for reinforcement learning from pixels Various existing deep RL methods have been proposed to address the sample-efficiency and generalization problem of RL algorithms learning from pixel observation . In end-to-end deep RL , neural networks learn representations implicitly by optimizing some RL objective ( Mnih et al. , 2015 ; Espeholt et al. , 2018 ) . Watter et al . ( 2015 ) and Wahlström et al . ( 2015 ) proposed to learn a encoder with training the dynamics model jointly to produce observation representations . Lee et al . ( 2020a ) , Hafner et al . ( 2019 ) , Hafner et al . ( 2020 ) and Zhang et al . ( 2019 ) aimed to learn an environment dynamics model with reconstruction loss to compact pixels to latent representation . Gelada et al . ( 2019 ) proposed to learn representations by predicting the dynamics model along with the reward , and analyzed its theoretical connection to the bisimulation metric . Agarwal et al . ( 2021 ) learns representations by state distances based on policy distribution . Kemertas & Aumentado-Armstrong ( 2021 ) modified the learning of bisimulation metric with intrinsic rewards and inverse dynamics regularization . Castro et al . ( 2021 ) replaced the Wasserstein distance in bisimulation metric with a parametrized metric . Data Augmentation in reinforcement learning Laskin et al . ( 2020b ) and Ye et al . ( 2020b ) explored the various data augmentation techniques from computer vision , which are limited to the transformation on input images . Cobbe et al . ( 2019 ) applied data augmentation on domain transfer . Ye et al . ( 2020a ) studied data augmentation on game environments for zero-shot generalization . RandConv ( Lee et al. , 2020b ) proposes a randomized convolutional neural network to generate randomized observations in order to perform data augmentation . A recent work DrQ ( Yarats et al. , 2021 ) suggests to perform random crop on image observations and provides regularized formulation for updating Q-function . Raileanu et al . ( 2021 ) proposed to use data augmentation for regularizing on-policy actor-critic RL objectives . Pitis et al . ( 2020 ) proposed a counterfactual data augmentation technique by swapping observed trajectory pairs . 3 PRELIMINARIES . We define an environment as a Markov Decision Process ( MDP ) described by a tuple M = ( S , A , P , R , γ ) , where S is the high-dimensional state space ( e.g. , images ) , A is the action space , P ( s′|s , a ) is the transition dynamics model that captures the probability of transitioning to next state s′ ∈ S given current state s ∈ S and action a ∈ A , R is the reward function yielding a reward signal r = R ( s , a ) ∈ R , and γ ∈ [ 0 , 1 ) is the discounting factor . The goal of reinforcement learning is to learn a policy π ( a|s ) that maximizes the expected cumulative rewards : maxπ E [ ∑ t γ tr ( st , at ) |at ∼ π ( ·|st ) , s′t ∼ P ( s′|s , a ) ] . In the scope of this paper , we do not consider partial observability , and use use stacked consecutive frames as the fully observed states . Soft Actor-Critic ( SAC ) is a widely used off-policy model-free reinforcement learning algorithm ( Haarnoja et al. , 2018 ) . SAC aims to maximize the entropy objective ( Ziebart , 2010 ) which is the reinforcement learning objective augmented with an entropy term J ( π ) =∑ t E [ r ( st , at ) + αH ( π ( ·|st ) ) ] . In order to maximize the objective , SAC learns a state-value function Qθ ( s , a ) , a stochastic policy πψ ( a|s ) and the temperature α , where θ and ψ are the parameters of Qθ and πψ , respectively . The Q-function is trained by minimizing the squared soft Bellman residual JQ ( θ ) = E ( st , at ) ∼D [ 1 2 ( Qθ ( st , at ) − ( r ( st , at ) + γEst+1∼P ( ·|st , at ) [ Vθ̄ ( st+1 ) ] ) ) 2 ] , ( 1 ) where D is the dataset or replay buffer storing the transitions , and Vθ̄ is the value function parameterized by θ̄ . The parameters ψ of policy is learned by maximizing the following objective Jπ ( ψ ) = Est∼D [ Eat∼πψ [ α log ( πψ ( a|s ) ) −Qθ ( st , at ) ] ] . ( 2 ) Bisimulation Metric defines a pseudometric d : S × S → R,1 where d quantifies the behavioral similarity of two discrete states ( Ferns et al. , 2004 ) . An extension of bisimulation metric to both continuous and discrete state spaces has also been developed ( Ferns et al. , 2011 ) . An variant of the bisimulation metric proposed in Castro ( 2020 ) defines a metric with respect to a policy π , which is known as the on-policy bisimulation metric . It removes the requirement of matching actions in the dynamics model but focuses on the policy π , which is able to better capture behavioral similarity for a specific task . Because of this property , in this paper , we focus on the π-bisimulation metric ( Castro , 2020 ) : Fπ ( d ) ( si , sj ) = ( 1− c ) |Rπsi −R π sj |+ cW1 ( d ) ( P π si , P π sj ) . ( 3 ) where Rπs : = ∑ a π ( a|s ) R ( s , a ) , Pπs : = ∑ a π ( a|s ) P ( ·|s , a ) , mapping Fπ : met → met , met denotes the set of all pseudometrics on state space S and W1 ( d ) ( · , · ) denotes the 1-Wasserstein distance given the pseudometric d. Then Fπ has a least fixed point denoted by d∗ . Deep Bisimulation 1If the pseudometric d of two states is 0 , then the two states belong to an equivalence relation . for Control ( DBC ) ( Zhang et al. , 2020 ) proposed to learn latent representations such that the L1 distances in the latent space equal to the π-bisimulation metric in the state space : argmin ϕ ( ∥ϕ ( si ) − ϕ ( sj ) ∥1 − |Rπsi −R π sj | − γW2 ( d ) ( P π si , P π sj ) ) 2 , where ϕ is the state encoder , and W2 ( d ) ( · , · ) denotes the 2-Wasserstein distance . DBC is combined with the reinforcement learning algorithm SAC where ϕ ( s ) is the input for SAC .

Summary: The paper studies the problem of learning invariant visual representations for reinforcement learning. Specifically, it builds on prior work Deep Bisimulation for Control (DBC) (Zhang et al), which learns an image representation predictive of rewards and dynamics and uses it with SAC. This paper makes a few modifications on top of DBC, specifically (1) instead of a single state embedding, it learns two separate ones (one for rewards and one for dynamics), (2) instead of using L1 distance in the embedding space to measure distance it instead uses a learned MLP which predicts distance, (3) using gradients from SAC it learns to dynamically adjust the balance between the dynamics weight and reward weight, and (4) adds image augmentation to the training. Experiments indicate on Distracting Control that on 2/4 tasks it matches DrQ and DBC and on the other 2 tasks outperforms them. Ablations suggest that all components are important, particularly (1) and (2).