Learning Domain Invariant Representations in Goal-conditioned Block MDPs
Authors: Beining Han
Chongyi Zheng
Harris Chan
Keiran Paster
Michael R. Zhang
Jimmy Ba
Venue: NeurIPS 2021 Poster (2021)
Keywords: goal-conditioned reinforcement learning
learning invariants
representation learning
domain generalization
block mdp
deep reinforcement learning

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 policies 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 goal-conditioned 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.

Learning Domain Invariant Representations in Goal-conditioned Block MDPs
Beining Han
Chongyi Zheng Carnegie Mellon University
IIIS, Tsinghua University boundinghan@@mmii.. com
chongyiz@andrew cmu. edu
Harris Chan
Keiran Paster Michael R. Zhang University of Toronto & Vector Institute {hchan, keirp, michael, jba]-cs. toronto.
Jimmy Ba
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 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 goal- conditioned RL agents. We propose theoretical framework in the Block MDP setting that characterizes the generalizability of goal-conditioned policies to new environments. Under this framework, develop practical method PA-SkewFir 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.
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. [017)] for Go, DQN
Mnih 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 [iiver et al. [2017). Mnih
al. 2015]. Haarnoja al. 2018 They also tend overfit specific environments and often fail
to generalize beyond the environment they were trained on Packer et al. (018). Unfortunately, this
characteristic limits the applicability of RL algorithms for many real world applications. Deployed
RL agents, c.g. robots in the field, will often face environment changes in their input such as different backgrounds, lighting conditions or object shapes Julian 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 the changes and often perform poorly practice J Julian et al. 2220 Zhang et al. [2020abb] bl.
In our work, seek to tackle changing, diverse problems with goal-conditioned RL agents. Goal- conditioned Reinforcement Learning is a popular research topic as its formulation and method is practical for many robot learning problems Marcin al. (2017). Eysenbach et al. (200). In goal- conditioned MDPs, the agent has to achieve a desired goal state which 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 important ensure that
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

the behaviour of conditonn 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. [211]. 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 D. and propose a general theory characterizing how well policy generalizes to unseen environments (Section B.J])
2. We propose theoretically-nne dbeoretically-motivatedd motivvteed algorithm based on optimizing a surrogate objective, perfect alignment, with aligned data (Section 3.2). We then describe practical implementation based on Skew-Fit Ping et al. [2020] to achieve the objective (Section B.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-arr performance compared to data augmentation and bisimulation baselines at generalizing to unseen test environments in coal-conditioned tasks (Section
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 condiiondi setting Kaclbling 1993 Schaul et al.
2015
Marcin et al. 2017]]
Sh-2
Sh-1
We consider family of Goal-conditioned Block MDP environments M {(S, A, X T g. E E where e stands for the envi- ronment index. Each environment consists of shared state space S, shared action space A. observation space X', transition dynamic T., shared goal space G C S and the discount factor 7.
Bess
Moreover, assume that M follows the generalized Block struc- ture Zhang et al. [2020a]. The observation xn E X is determined by state s E S and the environmental factor 6t E B", i... (s, 6*) (Figure ((C). For brevity, we use (s)s) to denote the observa- tion for domain e at state 8 and step We may also omit as (s) if we do not emphasize on the step t or the exact environ- mental factor b9. The transition function is thus consists of state transition p( P(8 |tc, ae) (also p(8o)), environmental factor transition o (15++116) In our work, we assume the state transition nearly deterministic, i.e., Vs, a, entropy H(p( (8,,))))) 4x)), H(p(s))) & 1, which quite common in most RL. benchmarks and applications Mmih et al. 2015). Greg et al. 2016). Pong et al. 2018 Most
Figure 1: Graphical model for Goal-conditioned Block MDPs (GBMDPs) setting. The agent takes in the goal g and observation It, which is produced by the domain invariant state 8t and environ- mental state b,, and acts with action ar. Note that by may have temporal dependence indicated by the dashed edge.
importantly, X t' satisfies the disjoint property Du et al.
[019]. i.c., each observation E X uniquely determines its underlying state Thus, the observation space X can be partitioned into disjoint blocks X'(s), s E S. This assumption prevents the partial observation problem.
The objective function in GBMDP is to learn a goal-conditioned policy {a)2"" g) that maximizes the discounted state density function ((T) Eysenbach et al. [2020) across all domains e E E. In our theoretical analysis, we do not assume the exact form of y to the policy. One can regard π (-1aaf .g) as
a group of RL. policies indexed by the goal state g.
J(T) = BeNE (1 (1-7)
Σ t=0
( 96g) = [e-o((((TT)
pf = glg) denotes the probability of achieving goal under policy π (-1xr .g) at step / in domain Besides, ~ E and g G refers to uniform samples from each set. As pf is defined over

---

state space, it may differs among environments since policy т takes x" as input. Fortunately, in GBMDP, there exist optimal policies TT(-12"., .9) which are invariant over all environments, i.e. To (c(a) = Ta(a| (s), Va A, 8 E S e, e' E E
During training, the agent has access to training environments {e:) Erain C E with their environment indices. However, we do not assume that &rain is i.id sampled from E. Thus, we want the coal-condii RL agent to acquire the ability to neglect the spurious and unrelated environmental factor be and capture the underlying invariant state information. This setup adopted in many recent works such as in Zhang et al. [2200a] and in domain generalization Koh et al. 22200]
Arjowsky et al. 2019 for supervised learning.
3 Method
In this section, we propose a novel learning algorithm to solve GBMDPs. First, we propose general theory characterize how well policy и generalizes unseen test environments after training on Euain Then, we introduce perfect aligument as surrogate objective for learning. This objective is supported by the generalization theory. Finally, we propose a practical method acquire perfect alignment.
3.1 Domain Generalization Theory for GBMDP
In seminal work, Ben-David et al. Ben-David et al. 2010 shows it is possible to bound the error of a classifier trained on 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. [202T), J, Albugqerquu et al. 2019]. In GBMDP, we can also derive similar theory to characterize the generalization from training environments Egain to target test environment t. The theory relies the Total Variation Distance Drv Wikipedia [021] of two policies Tı, T2 with input (g).)) which is defined as follows.
m2(-1= g))
sup A'Ea(A)
|T] (A'la" aa0 .9) (A'(aa"
In the following statements, we denote p(x. g) as some joint distributions of goals and observations that g ~ G and is determined by p(x g). Additionally, we use Pf (z"1g) to denote the discounted occupancy measure of x² in environment under policy x(-1=", .)) and refer p=(x") as the marginal distribution. Furthermore, we denote (l(x.a)
T2) as the average Drv between æ1 and 22, i.e., co(2.a) (T1 || π2) p (Div(TI (-12z.) 9) || T2(- (r, 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 Hiv--iiveeeeece BBn--avvdd er al. (2010), Sicilia et al. 2021 in domain
adaptation theory, we define ITATI-divrreenee of two joint distributions p(z, 9) and p(a, g) in terms of the policy class IT:
dnsu(e(E, 9), p(x, sup (e (l) mI' e((oo) (m 1|)
n,त'Ell
On one hand, duAm is a distance metric which reflects the distance between two distributions w.r.t function class II. On the other hand, if we fix these two distributions, it also reveals the quality of the function class II, i.e., smaller dдan means more invariance to the distribution change. Finally, we state the following Proposition in which TG is some optimal and invariant policy.
Proposition 1 (Informal). For any E II, consider the occupancy measure (pf ((t),,))1 for training environments and P'zo (x* g) for the target environment. For simplicity, we use e as the abbreviation of et (e*) 1),, el as la' a) and = maxe, dmAn (I .g,, Px (x4 Let
N
N
(
|| TG) + +x(** ma), n² =ra min n'ell
( ma)
N 1=1
Then, we have
('(Te)
MΣ E ( (T ~
+
+X+0+
min d g), g) Pno (x' .))) (2)
where B is a characteristic set of joint distributions determined by Evain and policy class II.
3

---

The formal statement and the proof are shown in Appendix Generally speaking, the first term the N training environments,
of the right hand side in Eq. 2 quantifies the performance quantifies the optimality of the policy class IT over all environments. ő reflects how the policy class ( 9),e, e, E Emain] which should be small if the policy
II can reflect the difference among 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 [en--avit et al. 2010 Liu et al. [2019] Sicilia al. [20TT] Albuquerque al. [2019] Akuzawa al. [2019] spend much effort
in discussing the trade-off's among different terms similar to the ones in Eq.
e.g., minimizing d
may increase A/Akuzawa al. [2019]. and in developing sophisticated techniques to optimize the
bound, e.g. distribution matching Louizos et al. [2016]. [i et al. (2018) Uin et al. (2200) or adversarial
learning Liu al. 2019
Different from their perspectives, in GBMDPs, we propose simple but effective criteria to minimize the bound. From now on, we only consider the policy class 11 11g {w(((x), g), Yw). Usually, & will be referred as an encoder which maps I E X to some latent representation. will also use the notation z(s) ((x(s)) if we do not emphasize on the specific environment.
(xx). We
Definition 1 (Perfect Alignment). An encoder is called a perfect alignment encoder 4 w.r. environ- ment set E Ve, e' E E and Vs, s' E S. (zx = (xx" (s')) if and only s'.
As illustrated in Figure 5 an encoder is in perfect alignment if it maps two observations of the same underlying state s the same latent encoding =s)) while also preventing meaningless embedding, iee, mapping observations of different states to the same 2. 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 Emaim then
1 N
J'(TG) - J'()) S
Σ (x ma)
+ TT) + (aa g). Pro (x ,g)
(3)
(E)
where p(x, g) and x² are defined in Proposition 7 aalso Appendix
C
In Appendix we formally prove Proposition when a (17, 1 -perfecct alignment, i.e., + is only near periect 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 Iearn a perfect alignment encoder via aligned sampling (Section 3.2)
For the remaining terms in the H.H.S of Eq. B we find hard to quantify them task agnostically, as similar difficulties also exist in the domain generalization theory of supervised learning (Sicilia et al. (02T]. 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 - 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 RI. (Table D.. For the (4) 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 perfect alignment encoder over Euain 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.
3.2 Learning Domain Invariant via Aligned Sampling
First, we discuss about the if condition on perfect alignment encoder 4 i.c., Vs, (s) = (x= (s)). The proposed method is based aligned sumpling. In contrast, most RI. algorithms use oeserratioon-eepeddent sampling from the environment, e.g., e-greedy or Gaussian distribution
policies Haarojja et al. 2018). () Fujimoto et al. 2018 Pong et al. (2020) Mnih et al. 2015). However,

---

Domain
Domsin
should L MMD
DIFF
M
S
1OOO
PE(A)
Encoder
Dacodar
concat
S
D
(a) Illustration of Aligned Sampling
(b) Overall structure
Figure 2: (a): Illustration of Aligned Sampling. Square represents the whole state space S, gray represents the distribution P (s) in two different environments. Small colored areas are the aligned
state distribution generated by aligned sampling in Section 3.2 (b): Overall VAE structure in our
PA-SF, Encoder maps x to the latent emhedding z and decoder reconstructs the observations with
and index T.MDD and I, DIFI denote the losses in Section 3.2
with observation- simply aligning the latent representation of these observations will fail to produce a perfect alignment encoder 4.
sampling, occupancy measures & (s), Ve E Evain will be different. Thus,
Thus, propose a novel strategy for data collection called aligned sampling. First, we randomly select trajectory (e.g., from replay buffer etc.), denoted as {z 00. 2f. 01,..... [TT] 2T) from envi- ronment e. The set of corresponding states along this trajectory are denoted as { (ao:) ( Second, we take the same action sequence ag:r another domain to get another trajectory (26 0o, 2I (11, (u.,...,.f] (so {sî (ao=)){=o)) We refer to the data collected by aligned sampling from all training environments as aligned data. These aligned observations { (da:)); Ve, &main are stored in an aligned buffer Ralign corresponding the aligned action sequence dg:t-
Under the definition of GBMDP, we have Vt E [0 Ti. E S. ((s)(a))))) = (( (ov:))) ).. i...
the same state distribution. Therefore, can use MMD loss Gretton al. [2008 to match distribution of (s) for the aligned data. More specifically, in each iteration, we sample a mini-batch of B aligned observations of every training environment e, Erain from Ralign i.c.,
the following loss computationally efficient approximation of the MMD metric Zhao and Meng (2015. Louizos et al. 2016).
Balign f (sf: (ab..)) Ve, Emain 1 Then
B
B (b(x (s (db)))) 18)
L MMD ( = Ee,e' Bai Raige
[I]
B
B b=1
where &b is random expansion function.
In Figure ((a) we illustrate the intuition of the above approach. When the transition is nearly deterministic, the entropy for p(sf (ao:)) is much smaller, i.e., H(p(s (ana))) < (s:)). Thus, p(sf (ag:)) can regarded as small patches in S. We use the MMD loss L MMD to match the latent representation (s)), s p(s, (4o Ve Eiin together. As consequence, achieve an encoder & that is more aligned. We discuss the theoretical property of I MMD in detail in Appendix
However, simply minimizing L MMD may violate the only if condition for perfect alignment. For example, a trivial solution for MMD 0 is mapping all observations to some constant latent. To ensure that (s)) (xx" (s')) only if s', we additionally use the difference loss L DIFF as follows.
DIFF ( -a--Ean 5 eR || (xxx) ( 11
where RD refers to the replay buffer of environment Clearly, minimizing DIFF encourages dispersed latent representations over all states s E S.
We refer to the combination O.MDD. MMD + opirr /, DIFF as our perfect alignment loss PA Note that T.PA resembles contrastive learning Chen et al. 2020), Laskin al. (2020a) Namely, observations of aligned data from Ralign positive pairs while observations sampled randomly from big replay
5

---

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 Section 3.1 we believe this latent representation will improve generalization to unseen target environments.
3.3 Perfect Alignment for Skew-Fit
In Section will train goal- conditiodd RL agents with perfect alignment encoder using the Skew-Fit a algorithm Pong et al. (222)]. Skew-Fit is typically designed for visual-input agents which learn coal-conditioned policy via purely self-super learning.
First, Skew-Fit trains a 3-VAE with observations collected online to acquire a compact and meaningful latent representation for each state, i.e., =(s) from the image observations =(s). Then, Skew-Fit optimizes a SAC Haarnoja et al. 2018 agent the goal goal-conditioned setting over the latent embedding of the image observation and goal. π (Taa2,, (, The reward function is the negative of l2 distance
between the two latent representation 2 (ss) and =(g), i.c., r(s; g) = - || ((s) - =(g) ) ll2. Furthermore, to improve sample efficiincy, Skew-Fit proposes skewed sampling for goal-conditioned exploration.
In our algorithm. Perfect Alignment for Skew-Fir (PA-SF), the encoder + is optimized via both 3- 3-VAE losses as Pong et al. [2020] Nair et al. 2018 and PA loss to ensure meaningful and perfectly aligned latent representation.
_ D) = RECON (2"", a²r)) + 3Dx (go(()=))) || p(=)) + CMMDD MMD + CDIITL DIFF
(4)
In addition, we use both aligned sampling and observation-depenie pling provides aligned data but hurts sumple-efficiency while dbservation-dependeee exploration-effiee 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 LAA
sampling. Aligned sam-
sampling is
Additionally, inspired by Louizos et al. [016]] we also change the 3-VAE structure to what 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 x²(s) some latent representation z(s) while the decoder D takes both 2(s) and the environment index input to reconstruct (s). Note that by using both PA and RECON we require static environmental factors in Euain (unnecessary for testing environments) for stable optimization. In future work, we will address the limit from 3-VAE by training two latent representations simultaneously to stabilize the optimization for generality.
4 Experiments
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 answer the following questions: (I) How well does PA-SF perform in solving GBMDP problems? (2) How does each component proposed in Section 3 contribute the performance?
4.1 Comparative Evaluation
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 RI. 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 al. 2021 tries use data augmentation prevent the RI. agents from distractions. We implement the most widely accepted data augmentation methods R
Laskin al. (0200) upon Skew-Fit -Stew-Fit+ + RÁD) baseline. Note that PA-SF method does not use any data augmentation and is parallel with this kind of techniques.
2. Model Bisimulation Ferns et al. (201T]. These methods utilize bisimulation metrics to learn minimal but sufficient representation which will neglect irrelevant features of Block MDPs. We include MISA Zhang et al. (2020a) and DBC Zhang et al. (20205) 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 || g lle as rewards for these two algorithms in GBMDP. In contrast, our PA-SF method does not have such information.

---

Table 1: Evaluation of PA-SF and baselines on four control tasks. We report the mean and one standard deviation on each task (lower metric is better).
Reach
Door
Push
Pickup
Algorithm
Hand distance (35K)
Angle difference (150K)
Puck distance (400K)
Object distance (280K)
Skew-Fit
0.111 + 0.010
0.194 + 0.018
0.086 + 0.004
0.037 + 0.006
Skew-Fit + RAD
0.105 = 0.010
0.162 = 0.030
0.082 = 0.008
0.040 + 0.004
Test Avg
MISA
0.239 = 0.0142
0.255 = 0.027
0.099 = 0.006
0.043 + 0.004
DBC
0.185 + 0.037
0.320 + 0.033
0.095 = 0.006
0.045 + 0.002
PA-SF(OUrS)
0.076 + 0.005
0.106 + 0.015
0.069 + 0.005
0.028 + 0.004
PA-SF(Ours)
0.067 + 0.005
0.058 + 0.074
0.060 + 0.005
0.020
0.008
Train Avg
Oracle Skew-Fit
0.055 + 0.010
0.057 + 0.012
0.054 + 0.006
0.020 + 0.006
Dear
Push
6.44
0.39
0640
8.05
Ro35
Tokal Tressress 30
Trasleni DC Sloew-Fit
Oracle Texetton
MISA
Figure 3: Learning curve of all algorithms average across test environments for each task. All curves show the mean and one standard deviation (a half for Pickup to show clearly) of 7 seeds.
Environments We evaluate PA-SF and all baselines set of GBMDP tasks based on multiworld benchmark Pong et al. 2018), which is widely used to evaluate the performance visual input goal-conditioned algorithms. We use the following four basic tasks Nair et al. 2018). Pong al. (220)]: Reach, Door, Pickup and Push. In GBMDP, we create different environments with various backgrounds, desk surfaces, and object appearances. During testing, also create environments with unseen video backgrounds to mimic environmental factor transitions (b, (bf). This makes policy generalization more challenging. Please refer Appendix D for full description of our experiment setup and implementation details of the baselines and our algorithm.
Results In Table I show the final average performance of each algorithm unseen test environments Eest. The corresponding learning curves are shown in Figure This metric shows the generalizability of each RL agent. All these agents are allowed to collect data from Evain (N 3) with static environmental factors. Our PA-SF achieves SOTA performance on all tasks. On testing environments, achieve 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 Skew-Fit algorithm trained directly on the single environment (and not simultaneously on all Emmin).
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 single MDP. In GBMDPs, this means acquiring set of encoders & that achieves model bisimulation for every possible and is thus inefficient for learning. By design, our method is not susceptible to this issue as the perfect alignment is 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 important research problem and orthogonal with our method.
Additionally, we also show the performance of PA-SF on the training environments in Table | PA-SF 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. B is reduced to almost zero via RL training in practice.
7

---

4.2 Design Evaluation
In this subsection, we conduct comprehensive analysis on the design PA-SF to interpret how well it carries out the theoretical framework discussed in Section B.1 and Section R.
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 T, and 7, 7,MMD, PA-SF (w/o D) excludes L D and PA-SF (w/o MD) excludes both losse:] Noticing that PA-SF (w/o MD) is equivalent to the Skew-Fir Fit algorithm with our proposed VAE structure (Figure 2 PA-SF (w/o AS) which excludes aligned sampling.
We also add
Door Trahing Eneronments Perfermance
Door Testing Envronments Perfomarce
4.3
6.0
ToTM Trrateoe.00 30
Tetal Tirattratt
PA-57
PA-5F (nth
PA-SF
PA-5F 04/5
Stoua-72
Onerie StHe-TT
Figure 4: Ablation of PA-SF and visualization of the latent representation via LER metric. All curves represent the mean and one standard deviation across seeds. Additionally, we also quantify the quality of the latent representation (x) in Figure 4 via the metric Latent Error Rate (LER). LER is defined as the average over environment set E E {Eirains Etest) as follows:
((x (s) ll2 (Txx(s)) (s)) 1l
Err(+)= = EE-
In general, the smaller Emr( is, the closer the encoder + is to perfect alignment over the environ- ments E. We first focus on the discussion about training performance.
1. 4 achieves the if condition of perfect alignment over Eiain via MMD as the LER value of PA-SF and PA-SF (w/o D) is almost O. 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 the MMD loss (Otherwise, the matching is inherently erroneous).
2. The only if condition, i.c., (s)( 4(x (s")) only if s'. is also achieved empirically by visualizing the reconstruction of the VAE (Figure D 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 learn well on the training environments as its latent representation mixed over different states.
Second, we focus the generalization performance target domains t, i..., term (t) in Eq. 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 the encoder 4 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. test environment ablations, the LER reduced significantly on methods with MMMDD This supports our claim that 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 L DIFF also plays role in reducing LER during testing. This supports the necessity both losses for generalization.
A single 7 DIFF is not useful here.
8

---

Train Env
Train Env
:: Dim
Train Env
Train Env
Train
Test Env
Test Env
Test Env
-15
-30
I-SNE Dim
Figure 5: t-SNE visualization of the latent space & (x*) trained with PA-SF for three environments:
training and testing of Push well instances visualization.
We observe the similar results other tasks well (Appendix Db. 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.
5 Related Work
Goal-conditioned RL: Goal-conditioned RL. Kaelbling [993], 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- reward- frree learning in goal-conditioned MDPs. Skew- Fit Pong et al. [2020] improves over RIG Nair et al. [0188] training efficiency by ensuring the behavioural goals used to explore 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 '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 LSC et al. 20200]
weak binary labeled data the
reward function. Others explore the same goal-condittoned RL problem via hindsight experience
replay Marcin al. [2077] Ren et [2019]. Ghosh et 2019), unsupervised reward learning Péré
et al. (2018). inverse dynamics models Paster et al. 2020 C-learming Eysenbach et al.) [2020] goal
generation Florensa al. [018). Nair and Finn (2019) [9] Pitis et al.] [2020). goal-conditonned forward
models] Nair et al. 2020J. and hierarchical RL Li et al. Nachum et al. [0188) Zhang et al. [2020c),
Hou et al. 2020 Our study focus on learning goal-conditioned policies that 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 al. 2020a Gelada al.
2019 implement the model-bisimulation framework Ferns et al. 2011 to acquire minimal but sufficient representation for solving the MDP problem. However, mnoeelbisimulation for digh--imension problems typically requires Nomain-invarianl and dense rewards. These assumptions do not hold in GBMDPs. Contrastive Metric Embeddings (CME) Agarwal et al. (2021) instead uses /² -iisimullaioon metric but also requires extra information the optimal policy. Another line of work tries to address these issues via self- suff-supe learning. [Stone] al. (202T) 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
practice. However, we find that RAD is not very helpful in Skew-Fit framework. Additionally, [Hansen et al. 22200] Bodnar et al. 2020) use self-supervise correction during weal-world adaptation like sim2real transfer but these methods are incompatible for domain generalization.

---

6 Conclusion
In this paper, we study the problem of learning invariant policies in Goal-conditioeed RL agents. The problem is formulated as GBMDP, which is an extension of Goal coal-condiiioeed MDPs and Block MDPs where we want the agent's 's policy to generalize to unseen test environments after training on several training environments.
As supported by the generalization bound for GBMDP, propose a simple but effective heuristic, i.e, perfect alignment which we can minimize the bound asymptotically and benefit the generalization. To learn perfect alignment encoder, we propose practical method based 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 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 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 coal-conditioned learning frameworks Eysenbach et al. [2220] Paster et al. [2020]
Acknowledgements
We are grateful for the feedback from anonymous reviewers. Resources used 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://wwww.ecoorrssstuuuuaaaaat
References
David Silver, Julian Schrittieeser, Karen Simonyan, Aj Antonoglou, Ioannis abd Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap. Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of go without human knowledge. Nature, 550, 2017.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg strovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan umaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518, 2015.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1861-1870, Stockholmsmässan, Stockholm Sweden, 10-15 Jul 2018. PMLR.
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization in deep reinforcement learning. CoRR, abs/1810.12221 2018. URL http://arxiv. wwwpl.org/gggoo wwtg/abs///1i/ 12282
Ryan Julian, Benjamin Swanson, Gauray Sukhatme, Sergey Levine, Chelsea Finn, and Karol Hausman. Efficient adaptation for end-to-end vision-based robotic manipulation. arXiv preprint arXiv:2004. 10/90, 2020.
Amy Zhang, Clare Lyle. Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block MDPs. In International Conference on Machine Learning, pages 1214-11224. PMLR, 2020a.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006.10742. 2020b.
10

---

6 Conclusion
In this paper, we study the problem of learning invariant policies in Goal-conditioeed RL agents. The problem is formulated as GBMDP, which is an extension of Goal coal-condiiioeed MDPs and Block MDPs where we want the agent's 's policy to generalize to unseen test environments after training on several training environments.
As supported by the generalization bound for GBMDP, propose a simple but effective heuristic, i.e, perfect alignment which we can minimize the bound asymptotically and benefit the generalization. To learn perfect alignment encoder, we propose practical method based 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 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 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 coal-conditioned learning frameworks Eysenbach et al. [2220] Paster et al. [2020]
Acknowledgements
We are grateful for the feedback from anonymous reviewers. Resources used 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://wwww.ecoorrssstuuuuaaaaat
References
David Silver, Julian Schrittieeser, Karen Simonyan, Aj Antonoglou, Ioannis abd Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap. Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of go without human knowledge. Nature, 550, 2017.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg strovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan umaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518, 2015.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1861-1870, Stockholmsmässan, Stockholm Sweden, 10-15 Jul 2018. PMLR.
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization in deep reinforcement learning. CoRR, abs/1810.12221 2018. URL http://arxiv. wwwpl.org/gggoo wwtg/abs///1i/ 12282
Ryan Julian, Benjamin Swanson, Gauray Sukhatme, Sergey Levine, Chelsea Finn, and Karol Hausman. Efficient adaptation for end-to-end vision-based robotic manipulation. arXiv preprint arXiv:2004. 10/90, 2020.
Amy Zhang, Clare Lyle. Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block MDPs. In International Conference on Machine Learning, pages 1214-11224. PMLR, 2020a.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006.10742. 2020b.
10

---

6 Conclusion
In this paper, we study the problem of learning invariant policies in Goal-conditioeed RL agents. The problem is formulated as GBMDP, which is an extension of Goal coal-condiiioeed MDPs and Block MDPs where we want the agent's 's policy to generalize to unseen test environments after training on several training environments.
As supported by the generalization bound for GBMDP, propose a simple but effective heuristic, i.e, perfect alignment which we can minimize the bound asymptotically and benefit the generalization. To learn perfect alignment encoder, we propose practical method based 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 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 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 coal-conditioned learning frameworks Eysenbach et al. [2220] Paster et al. [2020]
Acknowledgements
We are grateful for the feedback from anonymous reviewers. Resources used 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://wwww.ecoorrssstuuuuaaaaat
References
David Silver, Julian Schrittieeser, Karen Simonyan, Aj Antonoglou, Ioannis abd Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap. Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of go without human knowledge. Nature, 550, 2017.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg strovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan umaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518, 2015.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1861-1870, Stockholmsmässan, Stockholm Sweden, 10-15 Jul 2018. PMLR.
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization in deep reinforcement learning. CoRR, abs/1810.12221 2018. URL http://arxiv. wwwpl.org/gggoo wwtg/abs///1i/ 12282
Ryan Julian, Benjamin Swanson, Gauray Sukhatme, Sergey Levine, Chelsea Finn, and Karol Hausman. Efficient adaptation for end-to-end vision-based robotic manipulation. arXiv preprint arXiv:2004. 10/90, 2020.
Amy Zhang, Clare Lyle. Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block MDPs. In International Conference on Machine Learning, pages 1214-11224. PMLR, 2020a.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006.10742. 2020b.
10

---

6 Conclusion
In this paper, we study the problem of learning invariant policies in Goal-conditioeed RL agents. The problem is formulated as GBMDP, which is an extension of Goal coal-condiiioeed MDPs and Block MDPs where we want the agent's 's policy to generalize to unseen test environments after training on several training environments.
As supported by the generalization bound for GBMDP, propose a simple but effective heuristic, i.e, perfect alignment which we can minimize the bound asymptotically and benefit the generalization. To learn perfect alignment encoder, we propose practical method based 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 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 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 coal-conditioned learning frameworks Eysenbach et al. [2220] Paster et al. [2020]
Acknowledgements
We are grateful for the feedback from anonymous reviewers. Resources used 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://wwww.ecoorrssstuuuuaaaaat
References
David Silver, Julian Schrittieeser, Karen Simonyan, Aj Antonoglou, Ioannis abd Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap. Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of go without human knowledge. Nature, 550, 2017.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg strovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan umaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518, 2015.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1861-1870, Stockholmsmässan, Stockholm Sweden, 10-15 Jul 2018. PMLR.
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization in deep reinforcement learning. CoRR, abs/1810.12221 2018. URL http://arxiv. wwwpl.org/gggoo wwtg/abs///1i/ 12282
Ryan Julian, Benjamin Swanson, Gauray Sukhatme, Sergey Levine, Chelsea Finn, and Karol Hausman. Efficient adaptation for end-to-end vision-based robotic manipulation. arXiv preprint arXiv:2004. 10/90, 2020.
Amy Zhang, Clare Lyle. Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block MDPs. In International Conference on Machine Learning, pages 1214-11224. PMLR, 2020a.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006.10742. 2020b.
10


Reviews:
[{"review_id": "wE32zz3xGAe", "reviewer": "Reviewer_KeFP", "summary": "In the paper under review, state abstraction in Goal-conditioned RL (Block MDPs) is discussed. \nTheoretical analyses to motivate the use of perfect alignment encoders are provided. \nThen, a practical method (including the trajectory sampling method and encoder loss) for building the perfect alignment encoder is proposed. \nExperimental 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). \n", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "* key strengths: \n1. both theoretical analyses and practical algorithm (with empirical results) are provided.  \n2. alignment sampling proposed in Section 3.2 is a novel idea and interesting.  \n\n* key weakness:  \n1 The theoretical parts (Section 3.2 and 3.3) are a bit hard to follow.  \n1-1.  too many symbols are used. It would be more clear if the table list of all symbols is provided.  \n1-2.  I am not sure if the following assumption is really valid:  \n1-2-1. lines 143-144: \"Besides, under mild assumptions, if (E) ! 0\n144 then % ! 0\".  \n1-2-2. lines 148--149: \"Though it is hard to conduct task-agnostic analysis on (t) term, we believe that perfect alignment is still an adequate criteria in minimizing (t) term\".  \nFor the assumption made at lines 148--149, I'd like to know (t) is really reduced by using a perfect alignment encoder (PA-SF) in the robotic arm control benchmarks.  \n2. also,   lines 145--146 (and lines 631--636 in appendix):  \"Theoretically speaking, when # is a perfect alignment encoder, an goal-conditioned RL policy trained 146 over the encoded space {z(s)}s2S will minimize (E) to 0.\"  \nI think this statement is not valid if multiple optimal policies \\pi_G exist. (say that there are two optimal policies \\pi_{G,1} and \\pi_{G,2} and that \\pi is converged to \\pi_{G,1}, \\epsilon^{e_i}(\\pi || \\pi_{G, 2}) is still can be greater than zero.)  \n\n* Minor comments:  \n1. Typos?:  \nline 145:  an goal-conditioned -> a goal-conditioned  \nline 149:  criteria -> criterion  \nline 330: is formally formulated ->  is formulated  \nline 369: mdps -> {MDP}s  \n\n2. Reference style is not consistent (e.g., the abbreviated style of conference name is used in some references, and not in the others).  \n\n---Edit after reading the author response and the other reviews:------\n\nI would like to improve my score, as the author has basically adequately addressed my concerns: WR->WA\n\nI still have some concerns about the assumptions made in the theoretical analysis. \nFor example, I saw the author's response regarding the term (t), but was not sufficiently convinced. \n\nI also think that the presentation of the theoretical analysis section needs to be improved (e.g., as suggested by Reviewer S2fL).\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes. # Discussion about societal impact is not contained in the paper, but I think it is not really necessary for this research. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "SXe0B_cWMtG", "reviewer": "Reviewer_S2fL", "summary": "**High-Level Summary**\n\nThis 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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Technical Summary**\n\nThis paper studies goal-conditioned RL in a set of environments where the goal is to learn a goal-conditioned policy that reaches the goal in different environments. The setting assumes that these environments share a latent state space, action space, latent transition dynamics, and goal states which are a subset of state space. Further, a set of environment factors specific to each environment evolve independently of actions and are combined with the current latent state to generate an environment-specific observation. It is assumed that the latent state transition dynamics is nearly deterministic though the exact assumption is not formally stated in the paper. It can be shown that an optimal policy $\\pi_G$ exists that reaches the goal optimally and only depends on the latent state. \n\nThe first set of results analyze the generalization performance of a policy $\\pi$ compared to $\\pi_G$. This is similar to domain-generalization results. A set of policy is considered that factorizes as $\\pi(x) = \\pi(\\phi(x))$. It is argued that if $\\phi$ is a perfectly aligned decoder that maps two observations to the same value only if they are from the same state, then a policy trained on top of it will minimize the generalization error. This is not mathematically proven in the paper. Finally, the learned decoder is combined with the skew-fit algorithm that trains a soft actor-critic on a reward function that is given by the distance between the encoding of current state and goal and additionally trains the encoder using VAE. The two losses, VAE and the proposed loss are additively combined. \n\n**Strength**\n\n1. Studies goal-conditioned RL, and invariant representation which are important topics in modern RL. The proposed approach is novel to my knowledge.\n2. Experimental evidence supports the claim.\n\n**Weakness**\n\n1. There are two strong environment assumptions: the latent dynamics are assumed to be nearly deterministic and the latent state dynamics are independent of the environment. E.g., this task cannot be used for training a navigation bot in a set of houses where the house layout can be different.\n\n2. Weak theoretical support. While prop 1 and prop 2 are helpful, the discussion in lines 148-154 is hand-wavy. A rigorous mathematical argument would have shown that sub-optimality is based on \"misalignment\" of the decoder and back it up with matching bounds. Also, I am not sure how prop 3 helps. More importantly, the behaviour of the proposed loss $L^{PA}$ is not studied theoretically. Questions such as what is its optimal solution? Does the optimal solution learn the right encoder? Does the empirical minimization achieve this optimal solution? are not addressed.\n\nSome technical questions and concerns are raised below. \n\n**Concerns and Questions**\n\n1. Why is the learned state encoding a vector when a Block MDP has finite states?\n\n2. Why is MMD loss being minimized in $L^{MMD}$ but not in $L^{DIFF}$? What happens if one simply minimizes the L2 distance between encodings in $L^{MMD}$ as in $L^{DIFF}$?\n\n3. What is the exact formalism for near-deterministic dynamics?\n\n4. (Important) If one assumes deterministic dynamics and goals are a subset of state space, then can't one learn open-loop policies which will trivially generalize to all environments? Is the non-trivially coming from things being slightly deterministic? What happens if one continues to use open-loop policies in experiments? Formally, find an open-loop policy that reaches a given goal in the training environment. This should generalize to the test set whenever things are deterministic but also perhaps when things are near-deterministic.\n\n5. Where is \"characteristic set of joint distributions \" defined? (Line 120) \n\n**Presentation**\n\n1. Consider using condensed notations.  E.g., $x^e(s)$ can simply be written as $x$ whenever the environment or state information is not required which is often the case. \n\n2. Consider adding pseudocode early on. It is easy to get lost since many things are introduced suddenly such as the Skew-Fit algorithm section. \n\n3. Line: 77 \"_You_ can regard\" seems non-standard. Consider using \"One can regard\".\n\n4. Line 150, space before the comma in \"as the ones in training , it\"\n\n5. Line 160 has a letter \"s\" in $\\mathcal{E}_{train}s$ that doesn't parse\n\n6. Line 163: $x_T$ should be replaced by $x^{e}_T$. This is an example where the $e$ can be dropped for the whole episode without any problem.\n\n7. Line 168: Grammar in \"stored in a aligned\"", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See main review for limitations. ", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "NkRB8hHIY9F", "reviewer": "Reviewer_6X8Q", "summary": "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength: \n\n1. Learning invariant representation that can generalize across different environments is an important challenge in the study of RL. The authors propose a novel algorithm that tackles such a challenge and outperforms several strong baselines on the multi-world environment.\n\n2. The authors propose an analysis of the optimality gap in  GBMDP and identify several important factors of being optimal in solving GBMDP. The authors further propose their algorithm based on their theoretical findings.\n\nLimitations and inquiries\n\n1. In Proposition 1, is the bound tight? It seems that $\\delta$ measures the discrepancy of observation visitations across different environments, which can hardly go to zero due to different observations generated by the corresponding environments. Under what scenario will the term $\\delta$ become zero?\n\n\n2. In the aligned sampling part, it seems to me that the visitation also depends on the initial states. Hence, for different visitation measures to attain the same latent state distribution under the same action sequence, one should also ensure that the initial state distributions are the same. I wonder how to ensure that the initial observation distributions share the same latent state across different environments?\n\n\n3. Despite the fact that the test environment in this paper is widely accepted, I wonder if swapping only the background grant sufficient challenge in learning to transfer. In fact, with only changes in the background, it is hard to tell if the performance is due to the trained VAE that filters out the background or if the agent learns a correspondence between the pixel input and a latent state. For instance, what happens if one changes the locations of the camera in the training and test environments? Will the agent learn the latent state corresponding to the position of the robot arm in such a scenario?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See Limitations and inquiries in the Main review. I do not have concerns regarding the societal impact of this work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "BwE8d2cqnmx", "reviewer": "Reviewer_WvNf", "summary": "Summary\n-------\n\nThis paper proposes a pair of regularizers to the Skew-Fit algorithm to\nensure that the encoder maps two states to the same encoding if and only\nif the states are the same (whether the states are in different\nenvironments or not). The problem setting is goal-conditioned RL but the\npaper introduces the block MDP structure to better analyze\ngeneralization. The perfect alignment property is motivated from a\nderived inequality measuring the performance difference of two policies\nwith ideas based on the $\\mathcal{H}\\Delta \\mathcal{H}$ divergence. The\nproposed algorithm is found to be most performant on multiworld amongst\nthe baselines: SkewFit, MISA and DBC. The performance benefit is\nattributed to the enforcement of the perfect alignment property.\n\n.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Decision\n--------\n\nThe technical novelty and reported performance benefit of the algorithm\nseems very compelling. I have, however, some reservations about the\nperfect alignment property and the way it is enforced through the\nregularizer. Based on these reservations and some issues with the\nempirical investigation, I am rating the paper at a 5 (\"Marginally below\nthe acceptance threshold\"). If, after discussing with the authors and\nother reviewers, I am able to clear my reservations then I will be\nwilling to increase the score marginally. Below, I have outlined my\nconcerns more specifically and I hope it will help the authors improve\ntheir paper further.\n\nStrengths\n---------\n\n-   The theoretical analysis proposed in this paper seems novel. In\n    particular, the combination of goal-conditioned RL and block MDPs,\n    and the resultant domain generalization theory, is interesting.\n    Although the proof of lemma 1 seems similar to largely follow the\n    technical machinery in most theoretical RL papers, proposition 1 is\n    still insightful. Perfect alignment is somewhat motivated from this\n    bound, but this reasoning is loose due to the challenging (t) term\n    in Equation 3.\n-   Overall problem is well described and motivated.\n-   The empirical results on the mulitworld environment suggest that the\n    proposed algorithm does benefit from improved generalization. The\n    choice in baselines seems comprehensive, and the ablation study does\n    demonstrate the importance of the particular combination of\n    techniques for perfectly aligned Skew-Fit.\n\nWeaknesses\n----------\n\n-   Perfect alignment seems at odds with fundamental properties of the\n    reinforcement learning problem. In a particular environment $e$, two\n    states $s$ and $s'$ could have the same policy\n    $\\pi(a | s) = \\pi( a | s')$ or, given an arbitrary policy, the two\n    states could have equal value $\\pi$, $V^\\pi(s) = V^\\pi(s')$. If the\n    encoder is perfectly aligned, and we use a linear predictor $w$ on\n    top of the encoded state $z$, then I do not think you could\n    represent equivalences of value/policy between different states.\n-   Empirical concerns: I understand that using a self-supervised\n    goal-conditioned RL formulation requires a certain environment setup\n    (i.e. goal formulation and semantic observations). I think this does\n    not disqualify some classic control environments like mountaincar,\n    which would help contextualize your results.\n-   Ablation on regularization hyperparameters. Your main novelty is the\n    introduction of a regularizer to enforce perfect alignment. With the\n    addition of these regularizers however, you also introduce more\n    degrees of freedom in the form of the hyperparameters controlling\n    the regularization strength. I think there should be an ablation\n    study for these hyperparameters to ensure that the performance gains\n    is not attributed to the extra degrees of freedom instead of the\n    property of perfect alignment.\n\nDetailed Comments\n-----------------\n\n-   Importance of aligned sampling vs observation sampling is unclear.\n    The discussion in Section 3 is brief and the importance of splitting\n    the data collection to only partially aligned sampling is not\n    investigated in the experiments.\n\nMinor Comments\n--------------\n\n-   Line 71: The disjoint property suggests that the environment $b^e$\n    is not actually necessary in constructing the observation. What is\n    the role of $b^e$ in this case?\n-   What is Figure 5 supposed to show? There is no clustering of the\n    points but I am unsure if that is to be expected. It seems that the\n    puck position may be the same on left and the same on the right.\n    This could be better visualized through a plot of latent distance\n    and puck distance\n\nPost Rebuttal\n--------------\n\nIn light of the discussion below, I have increased my score from a 5 -> 6.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations and potential negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper focuses on an interesting problem, goal-conditioned policy generalization in block MDPs (a kind of POMDPs where the current state is uniquely identifiable from the current observation, even though a state can emit many different observations). This is a mostly theoretical work, but its experiments are convincing. The reviewers have closely examined the paper's theory and, on the one hand, didn't find errors in it but, on the other hand, found theory gaps in it that could be explained either by the paper being not fully clear or not sufficiently rigorous. Please see the discussions with reviewers WvNf and S2fL, especially reviewer S2fL's "Thanks for your response" comment. 

Nonetheless, due to this work being one of the early ones on this topic and being likely to become a stepping stone for further exploration of this area, the metareviewer recommends acceptance, trusting that the authors incorporate the points that came up in the discussion into the final version (it's hard to think of a reason not to do this). 
Decision: Accept (Poster)