Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: **High-Level Summary**

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

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

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

**Strength**

1. Studies goal-conditioned RL, and invariant representation which are important topics in modern RL. The proposed approach is novel to my knowledge.
2. Experimental evidence supports the claim.

**Weakness**

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

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

Some technical questions and concerns are raised below. 

**Concerns and Questions**

1. Why is the learned state encoding a vector when a Block MDP has finite states?

2. 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}$?

3. What is the exact formalism for near-deterministic dynamics?

4. (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.

5. Where is "characteristic set of joint distributions " defined? (Line 120) 

**Presentation**

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

2. Consider adding pseudocode early on. It is easy to get lost since many things are introduced suddenly such as the Skew-Fit algorithm section. 

3. Line: 77 "_You_ can regard" seems non-standard. Consider using "One can regard".

4. Line 150, space before the comma in "as the ones in training , it"

5. Line 160 has a letter "s" in $\mathcal{E}_{train}s$ that doesn't parse

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

7. Line 168: Grammar in "stored in a aligned"
limitations_and_societal_impact: See main review for limitations.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: 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.

.

main_review: Decision
--------

The technical novelty and reported performance benefit of the algorithm
seems very compelling. I have, however, some reservations about the
perfect alignment property and the way it is enforced through the
regularizer. Based on these reservations and some issues with the
empirical investigation, I am rating the paper at a 5 ("Marginally below
the acceptance threshold"). If, after discussing with the authors and
other reviewers, I am able to clear my reservations then I will be
willing to increase the score marginally. Below, I have outlined my
concerns more specifically and I hope it will help the authors improve
their paper further.

Strengths
---------

-   The theoretical analysis proposed in this paper seems novel. In
    particular, the combination of goal-conditioned RL and block MDPs,
    and the resultant domain generalization theory, is interesting.
    Although the proof of lemma 1 seems similar to largely follow the
    technical machinery in most theoretical RL papers, proposition 1 is
    still insightful. Perfect alignment is somewhat motivated from this
    bound, but this reasoning is loose due to the challenging (t) term
    in Equation 3.
-   Overall problem is well described and motivated.
-   The empirical results on the mulitworld environment suggest that the
    proposed algorithm does benefit from improved generalization. The
    choice in baselines seems comprehensive, and the ablation study does
    demonstrate the importance of the particular combination of
    techniques for perfectly aligned Skew-Fit.

Weaknesses
----------

-   Perfect alignment seems at odds with fundamental properties of the
    reinforcement learning problem. In a particular environment $e$, two
    states $s$ and $s'$ could have the same policy
    $\pi(a | s) = \pi( a | s')$ or, given an arbitrary policy, the two
    states could have equal value $\pi$, $V^\pi(s) = V^\pi(s')$. If the
    encoder is perfectly aligned, and we use a linear predictor $w$ on
    top of the encoded state $z$, then I do not think you could
    represent equivalences of value/policy between different states.
-   Empirical concerns: I understand that using a self-supervised
    goal-conditioned RL formulation requires a certain environment setup
    (i.e. goal formulation and semantic observations). I think this does
    not disqualify some classic control environments like mountaincar,
    which would help contextualize your results.
-   Ablation on regularization hyperparameters. Your main novelty is the
    introduction of a regularizer to enforce perfect alignment. With the
    addition of these regularizers however, you also introduce more
    degrees of freedom in the form of the hyperparameters controlling
    the regularization strength. I think there should be an ablation
    study for these hyperparameters to ensure that the performance gains
    is not attributed to the extra degrees of freedom instead of the
    property of perfect alignment.

Detailed Comments
-----------------

-   Importance of aligned sampling vs observation sampling is unclear.
    The discussion in Section 3 is brief and the importance of splitting
    the data collection to only partially aligned sampling is not
    investigated in the experiments.

Minor Comments
--------------

-   Line 71: The disjoint property suggests that the environment $b^e$
    is not actually necessary in constructing the observation. What is
    the role of $b^e$ in this case?
-   What is Figure 5 supposed to show? There is no clustering of the
    points but I am unsure if that is to be expected. It seems that the
    puck position may be the same on left and the same on the right.
    This could be better visualized through a plot of latent distance
    and puck distance

Post Rebuttal
--------------

In light of the discussion below, I have increased my score from a 5 -> 6.
limitations_and_societal_impact: The authors adequately addressed the limitations and potential negative societal impact.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
main_review: Strength: 

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

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

Limitations and inquiries

1. 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?


2. 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?


3. 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?
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_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: 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). 

main_review: * key strengths: 
1. both theoretical analyses and practical algorithm (with empirical results) are provided.  
2. alignment sampling proposed in Section 3.2 is a novel idea and interesting.  

* key weakness:  
1 The theoretical parts (Section 3.2 and 3.3) are a bit hard to follow.  
1-1.  too many symbols are used. It would be more clear if the table list of all symbols is provided.  
1-2.  I am not sure if the following assumption is really valid:  
1-2-1. lines 143-144: "Besides, under mild assumptions, if (E) ! 0
144 then % ! 0".  
1-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".  
For 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.  
2. 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."  
I 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.)  

* Minor comments:  
1. Typos?:  
line 145:  an goal-conditioned -> a goal-conditioned  
line 149:  criteria -> criterion  
line 330: is formally formulated ->  is formulated  
line 369: mdps -> {MDP}s  

2. Reference style is not consistent (e.g., the abbreviated style of conference name is used in some references, and not in the others).  

---Edit after reading the author response and the other reviews:------

I would like to improve my score, as the author has basically adequately addressed my concerns: WR->WA

I still have some concerns about the assumptions made in the theoretical analysis. 
For example, I saw the author's response regarding the term (t), but was not sufficiently convinced. 

I also think that the presentation of the theoretical analysis section needs to be improved (e.g., as suggested by Reviewer S2fL).


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_title: null

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)