Learning Domain Invariant Representations in Goal-conditioned Block MDPs

1 Introduction . Deep Reinforcement Learning ( RL ) has achieved remarkable success in solving high-dimensional Markov Decision Processes ( MDPs ) problems , e.g. , Alpha Zero Silver et al . [ 2017 ] for Go , DQN Mnih et al . [ 2015 ] for Atari games and SAC Haarnoja et al . [ 2018 ] for locomotion control . However , current RL algorithms requires massive amounts of trial and error to learn Silver et al . [ 2017 ] , Mnih et al . [ 2015 ] , Haarnoja et al . [ 2018 ] . They also tend to overfit to specific environments and often fail to generalize beyond the environment they were trained on Packer et al . [ 2018 ] . Unfortunately , this characteristic limits the applicability of RL algorithms for many real world applications . Deployed RL agents , e.g . robots in the field , will often face environment changes in their input such as different backgrounds , lighting conditions or object shapes Julian et al . [ 2020 ] . Many of these changes are often spurious and unrelated to the underlying task , e.g . control . However , RL agents trained without experiencing these changes are sensitive to the changes and often perform poorly in practice Julian et al . [ 2020 ] , Zhang et al . [ 2020a , b ] . In our work , we seek to tackle changing , diverse problems with goal-conditioned RL agents . Goalconditioned Reinforcement Learning is a popular research topic as its formulation and method is practical for many robot learning problems Marcin et al . [ 2017 ] , Eysenbach et al . [ 2020 ] . In goalconditioned MDPs , the agent has to achieve a desired goal state g which is sampled from a prior distribution . The agent should be able to achieve not only the training goals but also new test-time goals . Moreover , in practice , goal-conditioned RL agents often receive high-dimensional inputs for both observations and goals Paster et al . [ 2020 ] , Péré et al . [ 2018 ] . Thus , it is important to ensure that 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the behaviour of goal-conditioned RL agents is invariant to any irrelevant environmental changes in the input at test time . Previous work Zhang et al . [ 2020a ] tries to address these problems via model bisimulation metric Ferns et al . [ 2011 ] . These methods aim to acquire a minimal representation which is invariant to irrelevant environment factors . However , as goal-conditioned MDPs are a family of MDPs indexed by the goals , it is inefficient for these methods to acquire the model bisimulation representation for every possible goal , especially in high-dimensional continuous goal spaces ( such as images ) . In our work , we instead choose to optimize a surrogate objective to learn the invariant policy . Our main contributions are : 1 . We formulate the Goal-conditioned Block MDPs ( GBMDPs ) to study domain generalization in the goal-conditioned reinforcement learning setting ( Section 2 ) , and propose a general theory characterizing how well a policy generalizes to unseen environments ( Section 3.1 ) . 2 . We propose a theoretically-motivated algorithm based on optimizing a surrogate objective , perfect alignment , with aligned data ( Section 3.2 ) . We then describe a practical implementation based on Skew-Fit Pong et al . [ 2020 ] to achieve the objective ( Section 3.3 ) . 3 . Empirically , our experiments for a sawyer arm robot simulation with visual observations and goals demonstrates that our proposed method achieves state-of-the-art performance compared to data augmentation and bisimulation baselines at generalizing to unseen test environments in goal-conditioned tasks ( Section 4 ) . 2 Problem Formulation In this section , we formulate the domain invariant learning problem as solving Goal-conditioned Block MDPs ( GBMDPs ) . This extends previous work on learning invariances Zhang et al . [ 2020a ] , Du et al . [ 2019 ] to the goal-conditioned setting Kaelbling [ 1993 ] , Schaul et al . [ 2015 ] , Marcin et al . [ 2017 ] . We consider a family of Goal-conditioned Block MDP environments ME = { ( S , A , X e , T e , G , ) |e 2 E } where e stands for the environment index . Each environment consists of shared state space S , shared action space A , observation space X e , transition dynamic T e , shared goal space G ⇢ S and the discount factor . Moreover , we assume that ME follows the generalized Block structure Zhang et al . [ 2020a ] . The observation xe 2 X e is determined by state s 2 S and the environmental factor be 2 Be , i.e. , xe ( s , be ) ( Figure 6 ( c ) ) . For brevity , we use xet ( s ) to denote the observation for domain e at state s and step t. We may also omit t as xe ( s ) if we do not emphasize on the step t or the exact environmental factor bet . The transition function is thus consists of state transition p ( st+1|st , at ) ( also p ( s0 ) ) , environmental factor transition qe ( bet+1|b e t ) . In our work , we assume the state transition is nearly deterministic , i.e. , 8s , a , entropy H ( p ( st+1|st , at ) ) , H ( p ( s0 ) ) ⌧ 1 , which is quite common in most RL benchmarks and applications Mnih et al . [ 2015 ] , Greg et al . [ 2016 ] , Pong et al . [ 2018 ] . Most importantly , X E = [ e2EX e satisfies the disjoint property Du et al . [ 2019 ] , i.e. , each observation x 2 X E uniquely determines its underlying state s. Thus , the observation space X E can be partitioned into disjoint blocks X ( s ) , s 2 S . This assumption prevents the partial observation problem . The objective function in GBMDP is to learn a goal-conditioned policy ⇡ ( a|xe , g ) that maximizes the discounted state density function J ( ⇡ ) Eysenbach et al . [ 2020 ] across all domains e 2 E . In our theoretical analysis , we do not assume the exact form of g to the policy . One can regard ⇡ ( ·|xe , g ) as a group of RL policies indexed by the goal state g. J ( ⇡ ) = Ee⇠E , g⇠G , ⇡ `` ( 1 ) 1X t=0 tpe⇡ ( st = g|g ) # = Ee⇠E [ Je ( ⇡ ) ] ( 1 ) pe⇡ ( st = g|g ) denotes the probability of achieving goal g under policy ⇡ ( ·|xe , g ) at step t in domain e. Besides , e ⇠ E and g ⇠ G refers to uniform samples from each set . As pe⇡ is defined over state space , it may differs among environments since policy ⇡ takes xe as input . Fortunately , in a GBMDP , there exist optimal policies ⇡G ( ·|xe , g ) which are invariant over all environments , i.e. , ⇡G ( a|xe ( s ) , g ) = ⇡G ( a|xe 0 ( s ) , g ) , 8a 2 A , s 2 S , e , e0 2 E . During training , the agent has access to training environments { ei } Ni=1 = Etrain ⇢ E with their environment indices . However , we do not assume that Etrain is i.i.d sampled from E . Thus , we want the goal-conditioned RL agent to acquire the ability to neglect the spurious and unrelated environmental factor be and capture the underlying invariant state information . This setup is adopted in many recent works such as in Zhang et al . [ 2020a ] and in domain generalization Koh et al . [ 2020 ] , Arjovsky et al . [ 2019 ] for supervised learning . 3 Method . In this section , we propose a novel learning algorithm to solve GBMDPs . First , we propose a general theory to characterize how well a policy ⇡ generalizes to unseen test environments after training on Etrain . Then , we introduce perfect alignment as a surrogate objective for learning . This objective is supported by the generalization theory . Finally , we propose a practical method to acquire perfect alignment . 3.1 Domain Generalization Theory for GBMDP . In a seminal work , Ben-David et al . Ben-David et al . [ 2010 ] shows it is possible to bound the error of a classifier trained on a source domain on a target domain with a different data distribution . Follow-up work extends the theory to the domain generalization setting Sicilia et al . [ 2021 ] , Albuquerque et al . [ 2019 ] . In GBMDP , we can also derive similar theory to characterize the generalization from training environments Etrain to target test environment t. The theory relies on the Total Variation Distance DTV Wikipedia [ 2021 ] of two policies ⇡1 , ⇡2 with input ( xe , g ) , which is defined as follows . DTV ( ⇡1 ( ·|x e , g ) k ⇡2 ( ·|x e , g ) ) = sup A02 ( A ) |⇡1 ( A 0 |xe , g ) ⇡2 ( A 0 |xe , g ) | In the following statements , we denote ⇢ ( x , g ) as some joint distributions of goals and observations that g ⇠ G and x is determined by ⇢ ( x|g ) . Additionally , we use ⇢e⇡ ( xe|g ) to denote the discounted occupancy measure of xe in environment e under policy ⇡ ( ·|xe , g ) and refer ⇢e⇡ ( xe ) as the marginal distribution . Furthermore , we denote ✏⇢ ( x , g ) ( ⇡1 k ⇡2 ) as the average DTV between ⇡1 and ⇡2 , i.e. , ✏⇢ ( x , g ) ( ⇡1 k ⇡2 ) = E⇢ ( x , g ) [ DTV ( ⇡1 ( ·|x , g ) k ⇡2 ( ·|x , g ) ) ] . This quantity is crucial in our theory as it can characterize the performance gap between two policies ( see Appendix C ) . Then , similar to the famous H H-divergence Ben-David et al . [ 2010 ] , Sicilia et al . [ 2021 ] in domain adaptation theory , we define ⇧ ⇧-divergence of two joint distributions ⇢ ( x , g ) and ⇢ ( x , g ) 0 in terms of the policy class ⇧ : d⇧ ⇧ ( ⇢ ( x , g ) , ⇢ ( x , g ) 0 ) = sup ⇡ , ⇡02⇧ |✏⇢ ( x , g ) ( ⇡ k ⇡0 ) ✏⇢ ( x , g ) 0 ( ⇡ k ⇡0 ) | On one hand , d⇧ ⇧ is a distance metric which reflects the distance between two distributions w.r.t function class ⇧ . On the other hand , if we fix these two distributions , it also reveals the quality of the function class ⇧ , i.e. , smaller d⇧ ⇧ means more invariance to the distribution change . Finally , we state the following Proposition in which ⇡G is some optimal and invariant policy . Proposition 1 ( Informal ) . For any ⇡ 2 ⇧ , we consider the occupancy measure { ⇢ei⇡ ( xei , g ) } Ni=1 for training environments and ⇢t⇡G ( x t , g ) for the target environment . For simplicity , we use ✏ei as the abbreviation of ✏⇢ ei ⇡ ( x ei , g ) , ✏t as ✏⇢ t ⇡G ( xt , g ) and = maxei , e0i2Etrain d⇧ ⇧ ( ⇢ ei ⇡ ( x ei , g ) , ⇢ e0i ⇡ ( xe 0 i , g ) ) . Let = 1 N NX i=1 ✏ei ( ⇡⇤ k ⇡G ) + ✏ t ( ⇡⇤ k ⇡G ) , ⇡ ⇤ = argmin ⇡02⇧ NX i=1 ✏ei ( ⇡0 k ⇡G ) Then , we have J t ( ⇡G ) J t ( ⇡ ) 1 N NX i=1 ✏ei ( ⇡ k ⇡G ) + + + min ⇢ ( x , g ) 2B d⇧ ⇧ ( ⇢ ( x , g ) , ⇢ t ⇡G ( x t , g ) ) ( 2 ) where B is a characteristic set of joint distributions determined by Etrain and policy class ⇧ . The formal statement and the proof are shown in Appendix C.2 . Generally speaking , the first term of the right hand side in Eq . ( 2 ) quantifies the performance of ⇡ in the N training environments . quantifies the optimality of the policy class ⇧ over all environments . reflects how the policy class ⇧ can reflect the difference among { ⇢ei⇡ ( xei , g ) , ei 2 Etrain } , which should be small if the policy class is invariant . The last term characterizes the distance between training environment and target environment and will be small if the training environments are diversely distributed . Many works on domain generalization of supervised learning Ben-David et al . [ 2010 ] , Liu et al . [ 2019 ] , Sicilia et al . [ 2021 ] , Albuquerque et al . [ 2019 ] , Akuzawa et al . [ 2019 ] spend much effort in discussing the trade-offs among different terms similar to the ones in Eq . ( 2 ) , e.g. , minimizing may increase Akuzawa et al . [ 2019 ] , and in developing sophisticated techniques to optimize the bound , e.g . distribution matching Louizos et al . [ 2016 ] , Li et al . [ 2018 ] , Jin et al . [ 2020 ] or adversarial learning Liu et al . [ 2019 ] . Different from their perspectives , in GBMDPs , we propose a simple but effective criteria to minimize the bound . From now on , we only consider the policy class ⇧ = ⇧ = { w ( ( x ) , g ) , 8w } . Usually , will be referred as an encoder which maps x 2 X E to some latent representation z = ( x ) . We will also use the notation z ( s ) = ( x ( s ) ) if we do not emphasize on the specific environment . Definition 1 ( Perfect Alignment ) . An encoder is called a perfect alignment encoder w.r.t environment set E if 8e , e0 2 E and 8s , s0 2 S , ( xe ( s ) ) = ( xe 0 ( s0 ) ) if and only if s = s0 . As illustrated in Figure 5 , an encoder is in perfect alignment if it maps two observations of the same underlying state s to the same latent encoding z ( s ) while also preventing meaningless embedding , i.e. , mapping observations of different states to the same z . We believe perfect alignment plays an important role in domain generalization for goal-conditioned RL agents . Specifically , it can minimize the bound of Eq . ( 2 ) as follows . Proposition 2 ( Informal ) . If the encoder is a perfect alignment over Etrain , then J t ( ⇡G ) J t ( ⇡ ) 1 N NX i=1 ✏ei ( ⇡ k ⇡G ) | { z } ( E ) + ✏t ( ⇡⇤ k ⇡G ) + d⇧ ⇧ ( ⇢̃ ( x , g ) , ⇢ t ⇡G ( x t , g ) ) | { z } ( t ) ( 3 ) where ⇢̃ ( x , g ) and ⇡⇤ are defined in Proposition 1 ( also Appendix C ) . In Appendix C.3 , we formally prove Proposition 2 when is a ( ⌘ , ) -perfect alignment , i.e. , is only near perfect alignment . The proof shows that the generalization error bound is minimized on the R.H.S of Eq . ( 3 ) when asymptotically becomes an exact perfect alignment encoder . Therefore , in our following method , we aim to learn a perfect alignment encoder via aligned sampling ( Section 3.2 ) . For the remaining terms in the R.H.S of Eq . ( 3 ) , we find it hard to quantify them task agnostically , as similar difficulties also exist in the domain generalization theory of supervised learning Sicilia et al . [ 2021 ] . Fortunately , we can derive upper bounds for the remaining terms under certain assumptions and we observe that these upper bounds are significantly reduced via our method in the experiments ( Section 4 ) . The ( E ) term represents how well the learnt policy ⇡ approximates the optimal invariant policy on the training environments and is reduced to almost zero via RL ( Table 1 ) . For the ( t ) term , we show that an upper bound of ( t ) is proportion to the invariant quality of on the target environment . Moreover , we find that learning a perfect alignment encoder over Etrain empirically improves the invariant quality over other unseen environments ( t ) ( Figure 4 ) . Thus , this ( t ) term upperbound is reduced by learning perfect alignment . Please refer to Appendix C.4 for more details . Based on the theory we derived in this subsection , we adopt perfect alignment as the heuristic to address GBMDPs in our work . In the following subsections , we propose a practical method to acquire a perfect alignment encoder over the training environments .

Summary: **High-Level Summary**  This paper studies a new representation learning approach for goal-conditioned RL for the particular case of Block MDPs (Du et al., 2019). The goal is to learn a goal-conditioned policy that generalizes across a set of Block MDPs given access to a train set of Block MDPs. It is assumed that latent dynamics are nearly deterministic which is critical. Formally, the proposed representation learning approach samples sequence of actions and executes them in different sampled environments. An encoder is trained to map observations at the same time step along these two trajectories to a similar value (line 173) while separating two randomly chosen observations (line 186). It is argued that learning this representation yields invariant representation which helps optimize the generalization error. Experimental evidence supports the claim.