Learning Domain Invariant Representations in Goal-conditioned Block MDPs
Year: 2022

Abstract:
Deep Reinforcement Learning (RL) is successful in solving many complex Markov Decision Processes (MDPs) problems. However, agents often face unanticipated environmental changes after deployment in the real world. These changes are often spurious and unrelated to the underlying problem, such as background shifts for visual input agents. Unfortunately, deep RL agents are usually sensitive to these changes and fail to act robustly against them. This resembles the problem of domain generalization in supervised learning. In this work, we study this problem for goalconditioned RL agents. We propose a theoretical framework in the Block MDP setting that characterizes the generalizability of goal-conditioned policies to new environments. Under this framework, we develop a practical method PA-SkewFit that enhances domain generalization. The empirical evaluation shows that our goal-conditioned RL agent can perform well in various unseen test environments, improving by 50% over baselines.

heading: 1 Introduction
text: 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.

heading: 3 Method
text: 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.

heading: 3.1 Domain Generalization Theory for GBMDP
text: 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.

heading: 3.2 Learning Domain Invariant via Aligned Sampling
text: First, we discuss about the if condition on perfect alignment encoder , i.e., 8s, (xe(s)) = (xe 0 (s)). The proposed method is based on aligned sampling. In contrast, most RL algorithms use observation-dependent sampling from the environment, e.g., ✏-greedy or Gaussian distribution policies Haarnoja et al. [2018], Fujimoto et al. [2018], Pong et al. [2020], Mnih et al. [2015]. However,
with observation-dependent sampling, occupancy measures ⇢e⇡(s), 8e 2 Etrain will be different. Thus, simply aligning the latent representation of these observations will fail to produce a perfect alignment encoder .
Thus, we propose a novel strategy for data collection called aligned sampling. First, we randomly select a trajectory (e.g., from replay buffer etc.), denoted as {xe0, a0, xe1, a1, . . . , xeT } from environment e. The set of corresponding states along this trajectory are denoted as {set (a0:t)}Tt=0. Second, we take the same action sequence a0:T in another domain e0 to get another trajectory {xe 0 0 , a0, x e0 1 , a1, . . . , x e0 T } (so as {s e0 t (a0:t)} T t=0). We refer to the data collected by aligned sampling from all training environments as aligned data. These aligned observations {xeit (a0:t)}, 8ei 2 Etrain are stored in an aligned buffer Ralign corresponding to the aligned action sequence a0:t.
Under the definition of GBMDP, we have 8t 2 [0 : T ], s 2 S, ⇢(set (a0:t)) = ⇢(se 0
t (a0:t)), i.e., the same state distribution. Therefore, we can use MMD loss Gretton et al. [2008] to match distribution of (xe(s)) for the aligned data. More specifically, in each iteration, we sample a mini-batch of B aligned observations of every training environment ei 2 Etrain from Ralign, i.e., Balign = {xei(s ei t (a b 0:t)), 8ei 2 Etrain} B b=1. Then we use the following loss as a computationally efficient approximation of the MMD metric Zhao and Meng [2015], Louizos et al. [2016].
LMMD( ) = Ee,e0⇠Etrain,Balign⇠Ralign [k 1
B
BX
b=1
( (xe(set (a b 0:t)))
1
B
BX
b=1
( (xe 0 (se 0
t (a b 0:t)))) k 2 2]
where is a random expansion function.
In Figure 2(a), we illustrate the intuition of the above approach. When the transition is nearly deterministic, the entropy for ⇢(set (a0:t)) is much smaller, i.e., H(⇢(set (a0:t))) ⌧ H(⇢e⇡(st)). Thus, ⇢(set (a0:t)) can be regarded as small patches in S . We use the MMD loss LMMD to match the latent representation { (xe(s)), s ⇠ ⇢(set (a0:t))}, 8e 2 Etrain together. As a consequence, we should achieve an encoder that is more aligned. We discuss the theoretical property of LMMD in detail in Appendix C.5.
However, simply minimizing LMMD may violate the only if condition for perfect alignment. For example, a trivial solution for LMMD = 0 is mapping all observations to some constant latent. To ensure that (xe(s)) = (xe 0 (s0)) only if s = s0, we additionally use the difference loss LDIFF as follows.
LDIFF( ) = Ee⇠Etrain,xe,x̃e2Re k (xe) (x̃e) k22 where Re refers to the replay buffer of environment e. Clearly, minimizing LDIFF encourages dispersed latent representations over all states s 2 S .
We refer to the combination ↵MMDLMMD + ↵DIFFLDIFF as our perfect alignment loss LPA. Note that LPA resembles contrastive learning Chen et al. [2020], Laskin et al. [2020a]. Namely, observations of aligned data from Ralign are positive pairs while observations sampled randomly from a big replay
buffer are negative pairs. We match the latent embedding of positive pairs via the MMD loss while separating negative pairs via the difference loss. As discussed in Section 3.1, we believe this latent representation will improve generalization to unseen target environments.

heading: 3.3 Perfect Alignment for Skew-Fit
text: In Section 4, we will train goal-conditioned RL agents with perfect alignment encoder using the Skew-Fit algorithm Pong et al. [2020]. Skew-Fit is typically designed for visual-input agents which learn a goal-conditioned policy via purely self-supervised learning.
First, Skew-Fit trains a -VAE with observations collected online to acquire a compact and meaningful latent representation for each state, i.e., z(s) from the image observations x(s). Then, Skew-Fit optimizes a SAC Haarnoja et al. [2018] agent in the goal-conditioned setting over the latent embedding of the image observation and goal, ⇡(a|z, g). The reward function is the negative of l2 distance between the two latent representation z(s) and z(g), i.e., r(s, g) = k z(s) z(g) k2. Furthermore, to improve sample efficiency, Skew-Fit proposes skewed sampling for goal-conditioned exploration.
In our algorithm, Perfect Alignment for Skew-Fit (PA-SF), the encoder is optimized via both -VAE losses as Pong et al. [2020], Nair et al. [2018] and LPA loss to ensure meaningful and perfectly aligned latent representation.
L( , D) = LRECON(xe, x̂e) + DKL(q (z|x e) k p(z)) + ↵MMDL MMD + ↵DIFFL DIFF (4)
In addition, we use both aligned sampling and observation-dependent sampling. Aligned sampling provides aligned data but hurts sample-efficiency while observation-dependent sampling is exploration-efficient but fails to ensure alignment. In practice, we find that collecting a small portion (15% of all data collected) of aligned data in Ralign is enough for perfect alignment via LPA.
Additionally, inspired by Louizos et al. [2016], we also change the -VAE structure to what is shown in Figure 2(b), since in GBMDP data are collected from N training environments and thus, the identity Gaussian distribution is no longer a proper fit for prior. The encoder maps xe(s) to some latent representation z(s) while the decoder D takes both z(s) and the environment index e as input to reconstruct x̂e(s). Note that by using both LPA and LRECON, we require static environmental factors in Etrain (unnecessary for testing environments) for a stable optimization. In future work, we will address the limit from -VAE by training two latent representations simultaneously to stabilize the optimization for generality.

heading: 4 Experiments
text: In this section, we conduct experiments to evaluate our PA-SF algorithms. The experiments are based on multiworld Pong et al. [2018]. Our empirical analysis tries to answer the following questions: (1) How well does PA-SF perform in solving GBMDP problems? (2) How does each component proposed in Section 3 contribute to the performance?

heading: 4.1 Comparative Evaluation
text: In this subsection, we aim to answer the question (1) by comparing our proposed PA-SF method with vanilla Skew-Fit and several other baselines that attempt to acquire invariant policies for RL agents.
Baselines Current methods for obtaining robust policies can be characterized into two categories: (1) data augmentation and (2) model bisimulation. 1. Data Augmentation. Recent work Stone et al. [2021] tries to use data augmentation to prevent
the RL agents from distractions. We implement the most widely accepted data augmentation methods RAD Laskin et al. [2020b] upon Skew-Fit (Skew-Fit + RAD) as a baseline. Note that our PA-SF method does not use any data augmentation and is parallel with this kind of techniques.
2. Model Bisimulation Ferns et al. [2011]. These methods utilize bisimulation metrics to learn a minimal but sufficient representation which will neglect irrelevant features of Block MDPs. We include MISA Zhang et al. [2020a] and DBC Zhang et al. [2020b] in our comparison as they are the most successful implementations for high-dimensional tasks. Moreover, in the goal-conditioned setting, we use an oracle state-goal distance k s g k2 as rewards for these two algorithms in GBMDP. In contrast, our PA-SF method does not have such information.
Environments We evaluate PA-SF and all baselines on a set of GBMDP tasks based on multiworld benchmark Pong et al. [2018], which is widely used to evaluate the performance of visual input goal-conditioned algorithms. We use the following four basic tasks Nair et al. [2018], Pong et al. [2020]: Reach, Door, Pickup and Push. In GBMDP, we create different environments with various backgrounds, desk surfaces, and object appearances. During testing, we also create environments with unseen video backgrounds to mimic environmental factor transitions qe(bet+1|bet ). This makes policy generalization more challenging. Please refer to Appendix E for a full description of our experiment setup and implementation details of the baselines and our algorithm.
Results In Table 1, we show the final average performance of each algorithm on unseen test environments Etest. The corresponding learning curves are shown in Figure 3. This metric shows the generalizability of each RL agent. All these agents are allowed to collect data from Etrain (N = 3) with static environmental factors. Our PA-SF achieves SOTA performance on all tasks. On testing environments, we achieve a relative reduction around 40% to 65% of the corresponding metrics over vanilla Skew-Fit w.r.t the optimal metric possible (Oracle Skew-Fit). Oracle Skew-Fit refers to the performance of a Skew-Fit algorithm trained directly on the single environment (and not simultaneously on all Etrain).
Other invariant policy learning methods perform sluggishly on all tasks. For DBC and MISA, we hypothesize that they struggle for goal-conditioned problems since the model bisimulation metric is defined for a single MDP. In GBMDPs, this means acquiring a set of encoders g that achieves model bisimulation for every possible g and is thus inefficient for learning. By design, our method is not susceptible to this issue as the perfect alignment is a universal invariant representation for all goals. Data augmentation via RAD provides marginal improvement over the vanilla Skew-Fit. Nevertheless, we believe developing adequate data augmentation techniques for GBMDPs is an important research problem and is orthogonal with our method.
Additionally, we also show the performance of PA-SF on the training environments in Table 1. PA-SF is still comparable and as sample-efficient as Skew-Fit in the training environments. This supports the claim that the (E) term in the R.H.S of Eq. (3) is reduced to almost zero via RL training in practice.

heading: 4.2 Design Evaluation
text: In this subsection, we conduct comprehensive analysis on the design of PA-SF to interpret how well it carries out the theoretical framework discussed in Section 3.1 and Section 3.2.
To begin with, we show the learning curves in Figure 4 of different ablations of PA-SF in the Door environment during both training and testing. To understand the roles of LDIFF and LMMD, PA-SF (w/o D) excludes LDIFF and PA-SF (w/o MD) excludes both losses1. Noticing that PA-SF (w/o MD) is equivalent to the Skew-Fit algorithm with our proposed VAE structure (Figure 2(b)). We also add PA-SF (w/o AS) which excludes aligned sampling.
Err( ) = Ee⇠E,s⇠S  k (xe(s)) (xe0(s)) k2
k (xe(s)) k2
In general, the smaller Err( ) is, the closer the encoder is to perfect alignment over the environments E. We first focus on the discussion about training performance.
1. achieves the if condition of perfect alignment over Etrain via LMMD as the LER value of PA-SF and PA-SF (w/o D) is almost 0. While without MMD loss, PA-SF (w/o MD) and Skew-Fit struggle with large LER value despite achieving good training performance. Furthermore, the comparison between PA-SF and PA-SF (w/o AS) demonstrates the importance of using aligned data in the MMD loss (Otherwise, the matching is inherently erroneous). 2. The only if condition, i.e., (xe(s)) = (xe 0 (s0)) only if s = s0, is also achieved empirically by
visualizing the reconstruction of the VAE (Figure 9 in Appendix D) and we believe this is satisfied by both the difference loss and the reconstruction loss. Under the only if condition, the SAC Haarnoja et al. [2018] trained on the latent space achieves the optimal performance. In contrast, PA-SF (w/o AS) fails to learn well on the training environments as its latent representation is mixed over different states.
Second, we focus on the generalization performance on target domains t, i.e., term (t) in Eq. (3). We observe the following: 1. As shown by the learning curve of test environments, the target domain performance of different
ablations match that of the LER metric: SkewFit, PA-SF (w/o AS) > PA-SF (w/o MD) > PA-SF (w/o D) > PA-SF. During training, these ablations have almost the same performance, except PA-SF (w/o AS). This indicates that the increased test performance indeed comes from the improved representation quality of the encoder , i.e., more aligned. This supports our claim at the end of Section 3.1 and the upper bound analysis on the (t) term in Appendix C.4, that the increased invariant property of produces better domain generalization performance. 2. In test environment ablations, the LER is reduced significantly on methods with LMMD. This supports our claim that a perfect alignment encoder on training environments also improves the encoder’s invariant property on unseen environments. In addition, by encouraging dispersed latent representation, the difference loss LDIFF also plays a role in reducing LER during testing. This supports the necessity of both losses for generalization.
1A single LDIFF is not useful here.
We observe the similar results in other tasks as well (Appendix D). Here, we also visualize the latent space by t-SNE plot to illustrate the perfect alignment on task Push. Dots in training environments are matched perfectly and the corresponding test environment dot is approximately near as expected.

heading: 5 Related Work
text: Goal-conditioned RL: Goal-conditioned RL Kaelbling [1993], Schaul et al. [2015] removes the need for complicated reward shaping by only rewarding agents for reaching a desired set of goal states. RIG Nair et al. [2018] is the seminal work for visual-input, reward-free learning in goal-conditioned MDPs. Skew-Fit Pong et al. [2020] improves over RIG Nair et al. [2018] in training efficiency by ensuring the behavioural goals used to explore are diverse and have wide state coverage. However, Skew-Fit has its own limitation in understanding the semantic meaning of the goal-conditioned task. To acquire more meaningful goal’s and observation’s latent representation, several approaches apply inductive biases or seek human feedback. ROLL Wang et al. [2020] applies object extraction methods under strong assumptions, while WSC Lee et al. [2020] uses weak binary labeled data as the reward function. Others explore the same goal-conditioned RL problem via hindsight experience replay Marcin et al. [2017], Ren et al. [2019], Ghosh et al. [2019], unsupervised reward learning Péré et al. [2018], inverse dynamics models Paster et al. [2020], C-learning Eysenbach et al. [2020], goal generation Florensa et al. [2018], Nair and Finn [2019], Pitis et al. [2020], goal-conditioned forward models Nair et al. [2020], and hierarchical RL Li et al., Nachum et al. [2018], Zhang et al. [2020c], Hou et al. [2020]. Our study focus on learning goal-conditioned policies that is invariant of spurious environmental factors. We aim to learn a policy that can generalize to visual goals in unseen test environments.
Learning Invariants in RL: Robustness to domain shifts is crucial for real-world applications of RL. Zhang et al. [2020a,b], Gelada et al. [2019] implement the model-bisimulation framework Ferns et al. [2011] to acquire a minimal but sufficient representation for solving the MDP problem. However, model-bisimulation for high-dimension problems typically requires domain-invariant and dense rewards. These assumptions do not hold in GBMDPs. Contrastive Metric Embeddings (CME) Agarwal et al. [2021] instead uses ⇡⇤-bisimulation metric but it also requires extra information of the optimal policy. Another line of work tries to address these issues via self-supervised learning. Stone et al. [2021] tests multiple data augmentation methods including RAD Laskin et al. [2020b] and DrQ Kostrikov et al. [2020] to boost the robustness of the representation as well as the policy. Our work can also apply data augmentation in practice. However, we find that RAD is not very helpful in the Skew-Fit framework. Additionally, Hansen et al. [2020], Bodnar et al. [2020] use self-supervised correction during real-world adaptation like sim2real transfer but these methods are incompatible for domain generalization.

heading: 6 Conclusion
text: In this paper, we study the problem of learning invariant policies in Goal-conditioned RL agents. The problem is formulated as a GBMDP, which is an extension of Goal-conditioned MDPs and Block MDPs where we want the agent’s policy to generalize to unseen test environments after training on several training environments.
As supported by the generalization bound for GBMDP, we propose a simple but effective heuristic, i.e., perfect alignment which we can minimize the bound asymptotically and benefit the generalization. To learn a perfect alignment encoder, we propose a practical method based on aligned sampling. The method resembles contrastive learning: matching latent representation of aligned data via MMD loss and dispersing the whole latent representations via the DIFF loss. Finally, we propose a practical implementation Perfect Alignment for Skew-Fit (PA-SF) by adding the perfect alignment loss to Skew-Fit and changing the VAE structure to handle GBMDPs.
The empirical evaluation shows that our method is the SOTA algorithm and achieves a remarkable increase in test environments’ performance over other methods. We also compare our algorithm with several ablations and analyze the representation quantitatively. The results support our claims in the theoretical analysis that perfect alignment criteria is effective and that we can effectively optimize the criteria with our proposed method. We believe the perfect alignment criteria will enable applications in diverse problem settings and offers interesting directions for future work, such as extensions to other goal-conditioned learning frameworks Eysenbach et al. [2020], Paster et al. [2020].

heading: Acknowledgements
text: We are grateful for the feedback from anonymous reviewers. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute <http://www.vectorinstitute.ai/partners>.

heading: Andrychowicz Marcin, Wolski Filip, Ray Alex, Schneider Jonas, Fong Rachel, Welinder Peter,
text: McGrew Bob, Tobin Josh, Abbeel Pieter, and Zaremba Wojciech. Hindsight experience replay. In Advances in Neural Information Processing Systems, 2017.

heading: Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. C-learning: Learning to achieve
text: goals via recursive classification. arXiv preprint arXiv:2011.08909, 2020.
Keiran Paster, Sheila A McIlraith, and Jimmy Ba. Planning from pixels using inverse dynamics models. arXiv preprint arXiv:2012.02419, 2020.

heading: Alexandre Péré, Sébastien Forestier, Olivier Sigaud, and Pierre-Yves Oudeyer. Unsupervised learning
text: of goal spaces for intrinsically motivated goal exploration. In International Conference on Learning Representations, 2018.
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing, 40(6):1662–1714, 2011.

heading: Vitchyr Pong, Murtaza Dalal, Steven Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine. Skew-fit:
text: State-covering self-supervised reinforcement learning. In International Conference on Machine Learning, pages 7783–7792. PMLR, 2020.

heading: Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford.
text: Provably efficient RL with rich observations via latent state decoding. In International Conference on Machine Learning, pages 1665–1674. PMLR, 2019.
Leslie Pack Kaelbling. Learning to achieve goals. In IJCAI, pages 1094–1099. Citeseer, 1993.
Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function approximators. In International conference on machine learning, pages 1312–1320. PMLR, 2015.

heading: Brockman Greg, Cheung Vicki, Pettersson Ludwig, Schneider Jonas, Schulman John, Tang Jie, and

heading: Zaremba Wojciech. Openai gym, 2016.
text: Vitchyr Pong, Murtaza Dalal, Steven Lin, and Ashvin Nair. multiworld. https://github.com/ vitchyr/multiworld, 2018.

heading: Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal-
text: subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Sara Beery, et al. Wilds: A benchmark of in-the-wild distribution shifts. arXiv preprint arXiv:2012.07421, 2020.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893, 2019.

heading: Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman
text: Vaughan. A theory of learning from different domains. Machine learning, 79(1):151–175, 2010.
Anthony Sicilia, Xingchen Zhao, and Seong Jae Hwang. Domain adversarial neural networks for domain generalization: When it works and how to improve. arXiv preprint arXiv:2102.03924, 2021.

heading: Isabela Albuquerque, João Monteiro, Mohammad Darvishi, Tiago H Falk, and Ioannis Mitliagkas.
text: Generalizing to unseen domains via distribution matching. arXiv preprint arXiv:1911.00804, 2019.
Wikipedia. Total variation distance of probability measures — Wikipedia, the free encyclopedia. http://en.wikipedia.org/w/index.php?title=Total%20variation% 20distance%20of%20probability%20measures&oldid=1024556616, 2021. [Online; accessed 24-May-2021].
Hong Liu, Mingsheng Long, Jianmin Wang, and Michael Jordan. Transferable adversarial training: A general approach to adapting deep classifiers. In International Conference on Machine Learning, pages 4013–4022. PMLR, 2019.
Kei Akuzawa, Yusuke Iwasawa, and Yutaka Matsuo. Adversarial invariant feature learning with accuracy constraint for domain generalization. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 315–331. Springer, 2019.
Christos Louizos, Kevin Swersky, Yujia Li, Max Welling, and Richard S Zemel. The variational fair autoencoder. In ICLR, 2016.

heading: Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao.
text: Deep domain generalization via conditional invariant adversarial networks. In Proceedings of the European Conference on Computer Vision (ECCV), pages 624–639, 2018.
Xin Jin, Cuiling Lan, Wenjun Zeng, and Zhibo Chen. Feature alignment and restoration for domain generalization and adaptation. arXiv preprint arXiv:2006.12009, 2020.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actorcritic methods. In International Conference on Machine Learning, pages 1587–1596, 2018.

heading: Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A
text: kernel method for the two-sample problem. Journal of Machine Learning Research, 1:1–10, 2008.

heading: Ji Zhao and Deyu Meng. Fastmmd: Ensemble of circular discrepancy for efficient two-sample test.
text: Neural computation, 27(6):1345–1372, 2015.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597–1607. PMLR, 2020.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning, pages 5639–5650. PMLR, 2020a.

heading: Ashvin Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual
text: reinforcement learning with imagined goals. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pages 9209–9220, 2018.
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite–a challenging benchmark for reinforcement learning from pixels. arXiv preprint arXiv:2101.02722, 2021.

heading: Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas.
text: Reinforcement learning with augmented data. arXiv preprint arXiv:2004.14990, 2020b.

heading: Yufei Wang, Gautham Narayan Narasimhan, Xingyu Lin, Brian Okorn, and David Held. Roll: Visual
text: self-supervised reinforcement learning with object reasoning. arXiv preprint arXiv:2011.06777, 2020.

heading: Lisa Lee, Benjamin Eysenbach, Ruslan Salakhutdinov, Chelsea Finn, et al. Weakly-supervised
text: reinforcement learning for controllable behavior. arXiv preprint arXiv:2004.02860, 2020.
Zhizhou Ren, Kefan Dong, Yuan Zhou, Qiang Liu, and Jian Peng. Exploration via hindsight goal generation. Advances in Neural Information Processing Systems, 32:13485–13496, 2019.

heading: Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Devin, Benjamin Eysenbach, and
text: Sergey Levine. Learning to reach goals via iterated supervised learning. arXiv e-prints, pages arXiv–1912, 2019.
Carlos Florensa, David Held, Xinyang Geng, and Pieter Abbeel. Automatic goal generation for reinforcement learning agents. In International conference on machine learning, pages 1515–1528. PMLR, 2018.
Suraj Nair and Chelsea Finn. Hierarchical foresight: Self-supervised learning of long-horizon tasks via visual subgoal generation. In International Conference on Learning Representations, 2019.
Silviu Pitis, Harris Chan, Stephen Zhao, Bradly Stadie, and Jimmy Ba. Maximum entropy gain exploration for long horizon multi-goal reinforcement learning. In International Conference on Machine Learning, pages 7750–7761. PMLR, 2020.
Suraj Nair, Silvio Savarese, and Chelsea Finn. Goal-aware prediction: Learning to model what matters. In International Conference on Machine Learning, pages 7207–7219. PMLR, 2020.
Siyuan Li, Lulu Zheng, Jianhao Wang, and Chongjie Zhang. Learning subgoal representation with slow dynamics.
Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine. Near-optimal representation learning for hierarchical reinforcement learning. In International Conference on Learning Representations, 2018.
Lunjun Zhang, Ge Yang, and Bradly C Stadie. World model as a graph: Learning latent landmarks for planning. arXiv preprint arXiv:2011.12491, 2020c.
Zhimin Hou, Jiajun Fei, Yuelin Deng, and Jing Xu. Data-efficient hierarchical reinforcement learning for robotic assembly control applications. IEEE Transactions on Industrial Electronics, 2020.

heading: Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. Deepmdp:
text: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, pages 2170–2179. PMLR, 2019.

heading: Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Marc G Bellemare. Contrastive
text: behavioral similarity embeddings for generalization in reinforcement learning. arXiv preprint arXiv:2101.05265, 2021.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004.13649, 2020.
Nicklas Hansen, Yu Sun, Pieter Abbeel, Alexei A Efros, Lerrel Pinto, and Xiaolong Wang. Selfsupervised policy adaptation during deployment. arXiv preprint arXiv:2007.04309, 2020.

heading: Cristian Bodnar, Karol Hausman, Gabriel Dulac-Arnold, and Rico Jonschkowski. A geometric
text: perspective on self-supervised policy adaptation. arXiv preprint arXiv:2011.07318, 2020.

heading: John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region
text: policy optimization. In International conference on machine learning, pages 1889–1897. PMLR, 2015.
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8(3-4):279–292, 1992.

Review summaries:
Summary -------  This paper proposes a pair of regularizers to the Skew-Fit algorithm to ensure that the encoder maps two states to the same encoding if and only if the states are the same (whether the states are in different environments or not). The problem setting is goal-conditioned RL but the paper introduces the block MDP structure to better analyze generalization. The perfect alignment property is motivated from a derived inequality measuring the performance difference of two policies with ideas based on the $\mathcal{H}\Delta \mathcal{H}$ divergence. The proposed algorithm is found to be most performant on multiworld amongst the baselines: SkewFit, MISA and DBC. The performance benefit is attributed to the enforcement of the perfect alignment property.  . 

In the paper under review, state abstraction in Goal-conditioned RL (Block MDPs) is discussed.  Theoretical analyses to motivate the use of perfect alignment encoders are provided.  Then, a practical method (including the trajectory sampling method and encoder loss) for building the perfect alignment encoder is proposed.  Experimental results demonstrate that, in robotic arm manipulation benchmarks, the proposed method has a better generalization ability than previous methods (e.g., bi-simulation metric-based method).  

This paper aims to learn a domain invariant representation of observations in the goal-conditioned block MDP (GBMDP). To this end, the authors propose to minimize the deviation between embedded visitation trajectories with a fixed action sequence across different environments. The intention is to learn the latent mapping between states and observations. The authors also characterize the optimality gap between an arbitrary policy and the optimal invariant policy. Their analysis shows that the key to optimal and invariant policy is (i) the optimality of the policy at each environment, (ii) the deviation of the visitation measures across different environments, and (iii) the difference between the test environment and training environments.  The authors conduct extensive experiments on the multi-world environment and show that the proposed PA-SF method outperforms several baselines.

**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.