DOMINO: Decomposed Mutual Information Optimization for Generalized Context in Meta-Reinforcement Learning
Authors: Yao Mu
Yuzheng Zhuang
Fei Ni
Bin Wang
Jianyu Chen
Jianye HAO
Ping Luo
Venue: NeurIPS 2022 Accept (2022)
Keywords: Reinforcement Learning
Dynamics Generalization
Meta Reinforcement Learning

Abstract: Adapting to the changes in transition dynamics is essential in robotic applications. By learning a conditional policy with a compact context, context-aware meta-reinforcement learning provides a flexible way to adjust behavior according to dynamics changes. However, in real-world applications, the agent may encounter complex dynamics changes. Multiple confounders can influence the transition dynamics, making it challenging to infer accurate context for decision-making. This paper addresses such a challenge by decomposed mutual information optimization (DOMINO) for context learning, which explicitly learns a disentangled context to maximize the mutual information between the context and historical trajectories while minimizing the state transition prediction error. Our theoretical analysis shows that DOMINO can overcome the underestimation of the mutual information caused by multi-confounded challenges via learning disentangled context and reduce the demand for the number of samples collected in various environments. Extensive experiments show that the context learned by DOMINO benefits both model-based and model-free reinforcement learning algorithms for dynamics generalization in terms of sample efficiency and performance in unseen environments. 

Decomposed Mutual Information Optimization for Generalized Context in Meta-Reinforrementt Learning
Yao Mu
Yuzheng Zhuang Huawei Noah's 's Ark Lab
Fei Ni Tianjin University
The University of Hong Kong muyao0connect. .ku.
zhuangyuzhengghuuvviiccc com
fei_mi0tju edu cn
Bin Wang
Jianyu Chen
Huawei Noah's Ark Lab wangbin1580huawei. com
Tsinghua University
j jianyuchen@tsinghua. edu. cn
Jianye Hao
Ping Luo
Huawei Noah's Ark Lab boojianye@haaeei.comm com
The University of Hong Kong pluoêcs hku hk
Abstract
Adapting to the changes in transition dynamics is essential in robotic applications. By learning conditional policy with compact context, context-aware meta- reinforcement learning provides a flexible way to adjust behavior according to dynamics changes. However, in real-world applications, the agent may encounter complex dynamics changes. Multiple confounders can influence the transition dynamics, making it challenging to infer accurate context for decision- This paper addresses such challenge by DecOmposed Mutual INformation Optimization (DOMINO) for context learning, which explicitly learns a disentan- gled context to maximize the mutual information between the context and historical trajectories, while minimizing the state transition prediction error. Our theoretical analysis shows that DOMINO can overcome the underestimation of the mutual information caused by multi-confounded challenges via learning disentangled context and reduce the demand for the number of samples collected in various environments. Extensive experiments show that the context learned by DOMINO benefits both model- and model-free reinforcement learning algorithms for dynamics generalization in terms of sample efficiency and performance in unseen environments. Open-sourced code is released on homepage.
Introduction
Dynamics generalization deep reinforcement learning (RL) investigates the problem of training RL agent in a few kinds of environments and adapting across unseen system dynamics structures, such different physical parameters or robot mythologies. RL) has been proposed to tackle the problem by training a range of tasks, and fast adapting to new task with the learned prior knowledge. However, training in meta-RL requires orders of magnitudes more samples than single-task RL since the agent not only has to learn to infer the change of environment also has to learn the corresponding policies. Context-aware meta-RL. methods take step further and show promising potential capture local dynamics explicitly by learning an
RettaReeiforremmen Learning (Meta-
Ping Luo is the corresponding author. Yao Mu and Fei Ni conducted this work during the internship in Huawei Noah's 's Ark Lab.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

10000
Mac
Maa
5000
TIMI
2200
Graple
Symnieea Mazo
Man--Daaan Dewee-Daaannnnneeeooen Leg
Cpppicts
Clisial
Coppli Boiad
Happost
Canpole
(a) Example the multi-confounded environments
(b) Performance comparison
Figure 1: Generalization with complex dynamics changes. The transition dynamics of the robot may simultaneously influenced by multiple confounders, such as mass (G) ). length of leg (= ), or a crippled leg (<). ). In real-world situations, all the possible confounders may change simultaneously, which bring challenges to robotic dynamics generalization. DOMINO addresses such problem by decomposed MI optimization and achieves the state-of-the-art performance.
The historical trajectories are sampled from the joint distribution of multiple confounders, which are the key factors that cause the dynamics changes. Accordingly, if multiple confounders affect the dynamics simultaneously, the state transition distribution will become highly multi-modal, leading to challenges in extracting accurate context.
additional context vector from historical trajectories [1-4]
Recent advanced context-aware meta-RL methods [5-8] further improve meta-RL via contrastive learning, which optimizes the InfoNCE bound [9] of the mutual information in essence. These methods show a promising improvement in entangled context learning, which performs well in single confounded environments. However, demonstrated in Figure 1a, in real-world situations for robotic applications with partially unspecified dynamics, the transition dynamics can be influenced by multiple confounders simultaneously, such mass changes, damping, friction, malfunctional modules like a crippled leg. For example, when a transportation robot is working in the wild, the load will dynamically change the task progresses, while the humidity and roughness of the road also may vary. Moreover, some works also construct a confounder set for unsupervised RL environment generalization| (00115]. RIA [16] also constructs confounder sets with multiple confounders for unsupervised dynamics generalization. Such changeable environments bring great challenges to the robot for capturing contextual information, which motivates our study.
Contribution. In this paper, we give a theoretical analysis which demonstrates that when the number of confounders increases, InfoNCE will be loose bound of mutual information (MI) with the samples in limited seen environments, which is called MI underestimation [17]. To tackle this problem, we propose DecOmposed Mutual INformation Optimization (DOMINO) framework for context learning in meta-RL.. The context encoder aims to embed the past state-action pairs into disentangled context vectors and is optimized by maximizing the mutual information between the disentangled context vectors and historical trajectories while minimizing the state transition prediction error. DOMINO decomposes the full MI optimization problem into a summation of N smaller MI optimization problems by learning disentangled context. We then theoretically prove that DOMINO could alleviate the underestimation bias of the InfoNCE and reduce the demand for the samples collected in various environments [18, 19]. Last, with the learned disentangled context, we further develop the context-aware model-based and model-free algorithms learn the context-conditioned policy and illustrate that DOMINO can consistently improve generalization performance in both ways to overcome the challenge of multi-confounded dynamics.
Extensive experiments demonstrate that DOMINO benefits meta-RL on both the generalization performance in unseen environments and sample efficiency during the training process under the challenging multi-confounded setting. For example, as show in Figure 1b, it achieves 1.5 times per- formance improvement to T-MCL [3] the Cheetah domain and 2.6 times performance improvement to T-MCL in the Crippled-Ani domain. Visualization of the learned context demonstrates that the disentangled context generated by DOMINO under different environments could be more clearly distinguished in the embedding space, which indicates its advantage to extract high-quality contextual information from the environment.
Related Work
2.1 Retta-Reinforremen Learning
Meta-RI. extends the framework of meta-learning [20, 21] to reinforcement learning, aiming to learn adaptive policy being able to generalize to unseen tasks. Specifically, meta-RL methods

---

learn the policy based on the prior knowledge discovered from various training environments and reuse the policy to fast adapt to unseen testing environments after zero or few shots. Gradient- based meta-RI. algorithms [22-25] learn a model initialization and adapt the parameters with few policy gradient updates in new dynamics. Context-based meta-RI. algorithms [1-4] learn contextual information to capture local dynamics explicitly and show great potential to tackle generalization tasks in complicated environments. Many model-free context- based methods are proposed to learn a policy conditioned on the latent context that can adapt with off-policy data by leveraging context information and is trained by maximizing the expected return. PEARL [1] adapts to a new environment by inferring latent context variables from small number of trajectories. Recent advanced methods further improve the quality of contextual representation leveraging contrastive learning [5-8]. Unlike the model- -free methods mentioned above, context-avee -were world models proposed to learn the dynamics with confounders directly. CaDM [26] learns a global model that generalizes across tasks by training latent context capture the local dynamics. T-MCL [4] combines multiple-choice learning with context-awan world model and achieves state-of-the-arr results on the dynamics generalization tasks. RIA [16] further expands this method into unsupervised setting without environment label by intervention, and enhances the context learning via MI optimization.
However, existing context-based approaches focus on learning entangled context, in which each trajectory is encoded into only one context vector. In a multi confoundono environment, learning entangled contexts requires orders of magnitude higher samples to capture accurate dynamics infor- mation. To tackle this challenge, different from RIA [16] and T-MCL [4] DOMINO infers several disentangled context vectors from a single trajectory and divides the whole MI optimization into the summation of smaller ones. The proposed decomposed MI optimization reduces the amount of demand for diverse samples and thus improves the generalization of the policy to overcome the
adaptation problem in multi-confounded unseen environments.
2.2 Mutual Information Optimization for Representation Learning
Representation learning based on mutual information (MI) maximization has been applied in various
tasks such as computer vision [27, 28], natural language processing [29, 19], and RL [30], exploiting noise- coisr-oonraattvve estimation (NCE) [31], InfoNCE [9] and variational objectives [32]. InfoNCE has gained recent interest with respect to variational approaches due to its lower variance [33] and superior performance in downstream tasks. However, InfoNCE may underestimate the true MI, given that it is limited by the number of samples. To tackle this problem, DEMI [17] first scaffolds the total MI estimation into a sequence of smaller estimation problems. In this paper, since the confounders in the real world are commonly independent, we simplify the complexity of mutual information decomposition and eliminate the need to learn conditional mutual information sub-term, assuming that multiple confounders are independent of each other.
3 Preliminaries
We consider standard RL framework where an agent optimizes a specified reward function through interacting with an environment. Formally, we formulate our problem as a Markov decision process (MDP) [34], which is defined as a tuple (S, p. 7/ Po). Here, S is the state space, .A is the action space, (8") ls, a) is the transition dynamics, (s,a) is the reward function, Po is the initial state distribution, and 7 E [0, 1) is the discount factor. In order to address the problem of generalization, we further consider the distribution of MDPs, where the transition dynamics pa (s'1s, a) varies according to multiple confounders ū = {uo, u1.... u... The confounders can be continuous random variables, like the mass, damping, random disturbance force. discrete random variables, such as one of the robot's leg is crippled. We assume that the true transition dynamics model is unknown, but the state transition data can be sampled by taking actions in the environment. Given set of training setting sampled from p( Marai)) the meta-training process learns a policy ((s, c) that adapts to the task at hand by conditioning on the embedding of the history of past transitions. which we refer as context c. At temt- the policy should adapt to the new MDP under the test setting Ühstt drawn from p( (ess))
Our goal is to learn a policy to maximizing the expected return Rarain condition on the context e which is encoded from the sequences of current state action pairs {sT, ar 8r+1 34 ) / in several training
3

---

Length
Damping
Crippled Leg
Past (5, pairs
Sc+1
Ly
Loss
Positive sampled trajectory
Negutive sumpled trajectory
Sct1
Context Encoder
P4
Disentangled Context
Prediction
max INCE (ci7)
(st, at)
Net
Embedding Distance
Decomposed MI Optimization
Minimize Transition Prediction Error
The context encoder embeds the past state-action pairs (+..) into disentangled context vectors. The disentangled context vectors (O) are Iearned via the decomposed mutual information optimization errr Minimize transition prediction error: With the current state-action pair and the learned context vectors, future state can be predicted by the prediction network. The gradient of the prediction error will be used update both the context encoder and the prediction network. Decomposed MI Optimization: we optimize the MI between the learned context and the historical trajectories under the confounder setting via maximizing the InfoNCE bound, which aims minimize the embedding distance between the positive sampled trajectories (-.-) and the context, while maximizing the embedding distance between the negative sampled trajectories (*.)) and the context.
Figure 2:
The overall framework of DOMINO
while minimizing the state transition prediction
scenarios and enable it to perform well and achieve a high expected return Riest in test scenarios never seen before.
Σ' E'rr(si, an) }} a ~ m(s..)) # {"train" "test")
R Ba)p((a)
(I)
1=
Decomposed Mutual Information Optimization for Context Learning
In this section, we first provide theoretical analysis to show why the malti-conffonded environments are more challenging. We find that when the number of confounders increases, the InfoNCE will be loose bound of MI with the samples in limited seen environments, resulting in the underestimation of MI. To solve such a problem, we develop the DOMINO framework to learn disentangled context by decomposed MI optimization. We theoretically illustrate that the decomposed MI optimization can alleviate the underestimation of MI and reduce the demand for the number of samples. The disentangled context e = {co, Ci,.. CN embedded by the context encoder with parameter from the past state-action pairs 7* = { aj 1 in current episode. DOMINO explicitly maximizes the
M (r' = {sf, =f T collected MI between the context e and the historical trajectories T = {-
based on the combination of multiple confounders ă = {ug. U1..... uN } as same as the current confounder setting, while minimizing the state transition prediction error conditioned on the learned context. We solve the MI optimization problem by optimizing the InfoNCE lower bound on MI [9], which can be viewed as a contrastive method for the MI optimization, and decompose the full MI optimization into smaller ones to alleviate the underestimation of the mutual information and reduce the demand for the number samples collected various environments.
4.1 InfoNCE Bound for Mutual Information Optimization
InfoNCE bound [NCE(2;)) y) is a lower bound of the mutual information 7 7(x; v), where NCE stands for Coise-Contrastive Estimation, is type of contrastive loss function used for self-supervised learning. InfoNCE is obtained by comparing pairs sampled from the joint distribution z, V1 p(z, y) (y1 is

---

called the positive example) to pairs 2, built using set of negative examples, Y2:K p(y2.K)
=
p(vx): 1 (x; 8) 2 (NCE( 4. K) = E' log
(2)
===1
where & is function assigning a similarity score to æ, y pairs and K denotes the number of samples. Through discriminating naturally the paired positive instances from the randomly paired negative instances, it is proved to bring universal performance gains in various domains, such as computer vision and natural language processing.
Lemma INCE(X; Y; K) < [x; < log K is a necessary condition for INce(;YY K) to be light bound of (2; g). (see proof Appendix A)
Some previous context-aware methods leam an entangled context e by maximizing the mutual information between the context e embedded from the past state-action pairs in the current episode, and the historical trajectories T collected under the same confounder setting as the current episode. They solve this problem by maximizing the InfoNCE lower bound on INCE T), which can be viewed as a contrastive estimation [9] of INCE T), and obtain promising improvement in single- confounded environment. However, according to Lemma I, the INCE (c; T) may be loose if the true mutual information I(c; T) is larger than log K, which is called underestimation of the mutual information. Therefore, to make the InfoNCE bound to tight bound of (c; T), the minimum number of samples e((T) commonly influenced by multiple confounders ū = {Mo, ... M4...... Myy ........ ................................ uN simultaneously, under the assumption that the confounders are independent (such as mass and damping), the mutual information between the historical trajectories T and the context can be derived
real-world robotic control tasks, the dynamics of the robot
I(e; 7) = Ep(rict log P(T I c) p(T)
{
[[[[T] a)p(ū c)dd
Ep(r,C) log
p(T)
(3)
2 Ep(r.c)p(El)
log P(T (T) (1) P(T)
I(ū; T)
By Luy
Σ (u: 2=0
N
As the number of confounders increases, the lower bound of I(c; T) will become larger, and the necessary condition for (T;c K) to be tight bound of I [c;; will become more difficult to satisfy. Since T(c:T)) EN I (w;; T), ) let the necessary consition satisfied, the amount of data K must be larger than e (u,;T) according to Lemma 1. Thus the demand for data increases significantly. Since the confounders commonly independent in real-world, can relax this condition by learning disentangled context vectors instead entangled context intuitively?
4.2 Decomposed MI Optimization
If the context vectors {co, Ci...., CNN can be independent, then we can ease this problem by applying the chain rule on MI to decompose the total MI into a sum of smaller MI terms, i.e.,
INCE K)
(INcEle; T K)} S log K
(4)
Theorem If the context vectors {Co, CI....,cN, CN be independent. then the necessary condition for (xCE(c;7 T) to be a tight bound can be relaxed to I[c; T) S N log K log KN.. Thus, the need of the number of samples can be reduced from K 2 ((eT to K > Me:))
Inspired by Theorem 1, we intuitively leamn disentangled context vectors and maximize the mu- tual information between the historical trajectories T and the context vectors (co...., CN while minimizing the INCE between the context vectors, i.e., to maximize the LNCE
LNCE(P, w)
INCE(CI;
INCE(er;
(5)
1=0
where the INCE(C;;) T) can be obtained with the positive trajectory / and negative trajectories (Tx Vx=grie.
5

---

(ercha(rtt) K (E (er (ct(ssm((r!)))
INCE (ex; T) = E
log
(6)
The context c = (co..........................(.. Ci CC....,,.N ,CN (C: E Rm)) is encoded from the past state action pairs T" in current episode by (+).). Both the positive trajectory (collected in same setting of the confounders) and negative trajectory (collected in different setting of the confounders) are encoded to R by fww (-). The critic functiin +) measures the cosine similarity between inputs by dot product after normalization. Under the assumption that the setting of confounders will not change in one episode, obtain the [NCE(C;; cj), we use oj sampled from same episode like c; as the positive example, and use ej sampled from different episode as the negative example. Thus the (NCE(CE e₁) can be derived as
e (((()))) e (cr
INCE (c):e)) = Ei log
K
+ (c(rsst)
K
Then, the future state St-1 can be predicted with the the current state 8t, action and the disentangled context vectors (Ca...... CN, } by the state transition prediction network fo(.). We aim to minimize the prediction loss, which is equal to maximizing
++
(Pr 2) = E.
1
Σ
log fo (#x+1 | 8A, ax. (aaoo....,,,)))) , (si, (ii,ai) 1i-tH
t-1
(8)
A=(
where B is the training set, and H is the prediction horizon. The whole framework of DOMINO is demonstrated in Figure 2 and the overall objective function of DOMINO is
C(p, w. Φ) = Live (1p, Φ) + LNCE((,, w)
(9)
4.3 Combine DOMINO with Downstream RL Methods
Combination with Model-based RL. With DOMINO we can learn the context encoder and the context-aware world model together. First, the past state-action pairs are encoded into the disentangled context vectors by the context encoder. According to the learned context, the transition prediction network predicts the future states of different actions. particular, we use the cross entropy method (CEM) [35], typical neural model predictive control (MPC) [36] method, to select actions, which several candidate action sequences are iteratively sampled from candidate distribution, which is adjusted based on best-performing action samples. The optimal action sequence ac++r= {a....., ) can be obtained by
LT Σ
argmax
(r), P St, Oey Ctor-r., City
(10)
Then, we use the mean value of adjusted candidate distribution as action and re-plan at every timestep. We provide detailed algorithm pseudo-code in the Appendix B.I. As for the adaptation process, the policy and context encoder zero-shot adapts to the unseen confounders setting utest, and we use the same adaptive planning method [TCMLL(3), which selects the most accurate prediction head over recent experience condition on the inferred context. The details is introduced in Appendix D.4
Combination with Model-free RL. Previous works show that policy learned by model-free method can be more robust to dynamics changes when it takes the contextual information as an additional input [37-39]. Motivated by this, investigate whether the context encoder learned by DOMINO model-freee
can be used as a plug-and-plat module to improve the final generalization performance RL methods. We concatenate the disentangled context encoded by pre-trained context encoder from DOMINO and the current state-action pairs, and learn a conditional policy æ (at 8tt, co...., CN). We the Proximal Policy Optimization (PPO) method train the agent [40], which learns the policy by maximizing
π (at St, Cto Ctg) A, SKL (- tto......, CtN) ( 8t, Cto Caxy
Th
(11)
Teaa (at
88, Ctg.
)
where A, the estimation of the advantage function timestep 1. We provide detailed podedo-code in the Appendix B.2.

---

Cripple-Halfcheeetah--d
Ha Halfcheetah-m-d
Cripple-Ani-m-d
Ant-mtd
12500
3000
4000
Average Return
Awerage fteumm
10000
Average teeurn
Averge Retum
3000
7500
2000
2000
5000
1000
0-
0.0
0.4
08
1.2
1.6
0.0
0.4
0.8
1.2
1.6
2.0
0.0
0.4
0.8
1.2
2.0
0.0
0.4
0.8
16
2.0
Timesteps
1e5
Timestips
1e5
Timesteps
1e5
Timesteps
1e5
Slimhumansid
Hopper
Cartpole
Pendulum m
Average Naturn
2000
Avenge Return
1000
Average Returrn
400-
-1000
300
1000
200
Dlg-DGMING
100
-3000
alg-MINO MCL
-1000
400
0
4000
1.6
0.2
0.6
1.0
Timesneps
1c5
Timesteps
Timesteps
165
Timesteps
le5
Figure 3: The average returns of the model-based methods in training environments (over seeds).
5 Experiments
In this section, we evaluate the performance of our DOMINO method to answer the following questions: (1) Can DOMINO help the model-based RL methods overcome the multi-confounded challenges in dynamics generalization (see comparison with ablation in Figure and Figure 5)? (2) Can the context encoder learned by DOMINO be used plug-and-play module to improve the generalization abilities of model-free RL methods in multi-confounded environments (see comparison with ablation in Table and Table 2)? (3) Can the proposed decomposed MI optimization benefit the forward prediction of the world model? (see Figure 6) (4) Does the disentangled context extract more meaningful contextual information than entangled context (see Figure 7)?
5.1 Setups
We demonstrate the effectiveness of our proposed method on 8 benchmarks, which contain 6 typical robotic control tasks based on the MuloCo physics engine [41] and 2 classical control tasks (CartPole and Pendulum) from OpenAl Gym [42]. Different from previous works, the environments are influenced by multiple confounders simultaneously. In our experiments, we modify multiple environment parameters the same time (e.g., mass, length, damping, push force, and crippled leg) that characterize the transition dynamics. The robotic control tasks contain 4 environments (Hopper, HalfCheetah, Ant, SlimHumanoid) affected by multiple continuous confounders and more difficult environments (Crippled Ant and Crippled HalfCheetah) affected by both continuous and discrete confounders. The detailed settings are illustrated in Appendix C (Table 3). We implement these the
environments based the publicly available code provide by [43, 4], and we also open- code of the multiple-confe confounde e For both training and testing phase, we sample the confounders at the beginning of each episode. During training, we randomly select combination of confounders from training set. At test time, we evaluate each algorithm in unseen environments with confounders outside the training range.
5.2 Comparison with Model-based Methods
Baselines. We consider T-MCL [3] and RIA[16] as the key baselines comparison with model- based methods, which achieve the state-of of offtheeart results in zero-shot dynamics generalization tasks. Since RIA doesn't 't has a adaptive planning process, we provide the DOMINO and T-MCL without adaptive planning to fair compare to the RIA. We also consider ablation version of DOMINO as baseline (denoted as MINO) to show the effectiveness of the decomposed MI optimization, which optimizes the MI and predicts the future states with an entangled context without decomposition.
Figure shows the generalization performance tested in the unseen environments. The results show that DOMINO surpasses T-MCL in terms of the generalization performance and the learning sample
7

---

DOHINO
Results. As shown in Figure 9, DOMINO achieves better generalization performance than RIA and TMCL even without the adaptive planning, especially in complex environments like Halfcheetah-m-- and fiin-humaniid-m-de Figure shows the average return during the learning process in the training envi- ronments. The results illustrate that DOMINO learns the policy more efficiently than T-MCL and MINO.
RIA
TMC
4000
2000
Ho fofchhetat-m- Sim-Humandd--me Pendulum-m-1
Figure 4: Comparison w/o adaptive planning
Chiple--fellccheethh---d
Cripple-Ant-m--d
4200
10000
3000
Average Maturn
8000
15500
5000
6000
2000
2000
Neergg
Avenage
4000
1000
1000
2000
0.0
0.4
Timesteps
le5
165
Timesteps
Timesteps
le5
Slimhumanoid-m-d
Hopper-m-d
Cartpole-fi
Pendulum-m-
1200
Avetage Netum
1000
Retum
400
Averaze Ratum
11000
1000
Averege
300
-150
800
DOMIN al--MINO
600
alg
MCL
0.0
0.4
0.8
1.2
166
2.0
0.0
0.2
0.6 Timesteps
0.4
0.8
10
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
04
0.6
0.8
1.0
Timesteps
le5
le5
Timestaps
le5
Timestaps
1e5
Figure 5: Comparison with model-based methods of the generalization performance (average return) in unseen multi-confoudded environments (over 8 seeds).
efficiency. This demonstrates that the disentangled context improves the context-awan world model. Especially, the performance gain becomes much more significant more complex environments (e.g., long-horizon and high-dimensional domains like Cripple-Ant, Ant, and Hopper). For example, DOMINO achieves about 2.6 times improvement to T-MCL. Cripple-Ant-m-d, which is one the most difficult environment, whose leg will randomly be crippled, and its mass and damping will be changed in testing. More details are shown in Appendix D.
5.3 Comparison with Model-free Methods
We also verify whether the learned disentangled context is useful for improving the generalization performance of model-free RL methods. Similar to [3, 44], we use the Proximal Policy Optimization (PPO [40]) method to train the agents.
Baselines. Our proposed method, which takes the context learned by DOMINO conditional input (PPO+DOMINO), compared with several context-conditional policies [45, 39]. Specifically, we consider combining the PPO with the context learned by T-MCL (PPO+T-MCL), which learns the context encoder via context-aware world model and achieves the state-of-the-art performance on dynamics generalization. We also consider PEARL [45]. which learns probabilistic context variable by maximizing the expected returns. We further develop an ablation version of DOMINO, which optimize the MI with entangled context (PPO+MINO) the decomposed MI optimization. We provide more detailed explanations in Appendix D.
bascline illustrate the effectiveness of
Results. Table and Table 2 show the performance of various model-free RL methods on both training and test environments. PPO+DOMINO shows superior performance and shows better generalization performances than previous conditional policy methods, implying that the proposed DOMINO method can extract contextual information more effectively than both the context learned by the model-based method (PPO+T-MCL) and the context learned by the model-free method (PEARL). Furthermore, PPO+DOMINO shows an obvious advantage over PPO+MINO, especially in complex environments, such as HalfCheetah, Ant, and Hopper, which implies that the decomposed

---

MI optimization improves the context learning significantly. Additionally, the results also show that compared PEARL, the context learned by T-MCL and DOMINO has better performance which implies that the state transition perdition can help to extract contextual information more effectively.
Table Comparison with model-free methods in the multi-oonfounded environments (over seeds). The transition dynamics will change in both training and test environments in every episode.
Carpplle-f-4
Pendulum-m--
Ant-ma-
Train
Test
Train
Test
Train
Test
PEARL
197_12
1751:37
-1265
-12931134
153163
73:125
PPO+T-MCI.
220-27
182+-25
-558+188 184
-579++125
176+-82
173+388
PPO+MINO
267=188
234::46
-4971162
-5264:219
194495
184 = 47
PPO+DOMINO
299 23
283.1.8
-405 139
446.1116
227:186
2161:32
Halfcheetah-m-dd
Slimhumnnoid-m-d
Hopper-m-d
Train
Test
Train
Test
Train
Test
PEARL
1802:1773
530:1-270
6947::3541
3697:1-2674
934:12422
8741:366
PPO-T-MCL
2032+688
674+:395
6157=1435
4136++1528
937=252
896#2388
PPO+MINO
1973+5563
824+-498
6179+1123
4275+1134
1109+3449
964+3323
PPO+DOMINO
2472:1803
1034_476
7825
525811039
1409
1137=335
5,4 Disentangled Context Analysis
Prediction errors. To show that our method indeed helps with transition prediction, we compare baseline methods with DOMINO in terms of prediction error across environments with varying multiple confounders. As shown in Figure 6, our model demonstrates superior prediction performance, which indicates that the learned context capture better contextual information compare with entangled context (see more results Appendix E.1).
Visualization. We visualize concatenation of the disentangled context vectors learned by DOMINO via t-SNE [46] and compare it with the entangled context learned by T-MCL. As shown in Figure 7, we find that the disentangled context vectors encoded from trajectories collected under different confounder settings could be more clearly distinguished in the embedding space than the entangled context learned by T-MCL. This indicates that DOMINO extracts high-quality task-specific informa- tion from the environment compared with T-MCL. We provide more visualization results based on both I-SNE [46] and PCA [47] in Appendix F.2.
DOMBOO
TCM
(a) Visualization in Hopper-m-d.
(b) Visualization in Cripple-Ant-ma
Figure 7: t-SNE visualization of context vectors extracted from trajectories collected various environments. Embedded points from environments with the same confounders have the same color.
Hopper setups: SetupO (m 0.75, d = 1.5). Setup/ (m 1.5,d = 1.25), Setup2 (m 0.5, d = 1.25), Setup3 (m 1.25, = 0.75), Setup4 (m 1.0, d 1.5); Cripple-Ant setups: SetupO (m 1.15, leg 3), Setupl
(m 0.75, blg 1), Setup2 (m 1.25, leg 0), Setup3 (m 0.85, leg 0), Setup4 (m 1.0, leg 2)(20 trajectories per setup).
6 Conclusion
In this paper, we propose a decomposed mutual information optimization (DOMINO) framework to learn the generalized context for zero-shot dynamics generalization. The disentangled context is learned by maximizing the mutual information between the context and historical trajectories while minimizing the state transition prediction error. By decomposing the whole mutual information

---

Table 2: Comparison with model-free methods in more
Crtpp--nntm-d
0:23
difficult multi-confounded environments (over seeds).
Cripple Sepplle-Attrig m-
Cripple Halfcheetah m-dd
Tran
Test
Train
Test
PEARL
182:173
1:11
2538
1028:1-445
0.05
PPO+T-MCL
187265
1091:36
23684726
1006:1-33
2.2
0.
1.0
PPO+MINO
2062164
1131:34
24931166
1197:14244
Timeetopp
Timettgge
165
1324-27
2503--658
Figure 6: The testing prediction error.
optimization problem into smaller ones, DOMINO can reduce the need for samples collected in various environments and overcome the underestimation of the mutual information in multi- confounded environments. Extensive experiments illustrate that DOMINO benefits the generalization performance in unseen environments with both model-based RL and model-free RL. For future work, an effective combination of DOMINO and RIA [16], which expands the decomposed MI optimization to relational intervention approach proposed by RIA could become stronger baseline for unsupervised dynamics generalization. We believe our work can lay the foundation of dynamics generalization in complex environments.
Limitations and Negative Social Impact. DOMINO sets the number of disentangled context vectors aaper-parameter equal the number of confounders in the environments, and capturing the number of confounders automatically could be future works. We believe that DOMINO will not cause any negative social impact.
Acknowledgments and Disclosure of Funding
The authors would like thank the anonymous reviewers for their valuable comments and helpful suggestions. The work is supported by Huawei Noah's 's Ark Lab; Ping Luo is supported by the General Research Fund of HK No.27208720. No. 17212120, and No. 17.0720.
References
[1] Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforeement learning via probabilistic context variables. In Proceedings of the 36th International Conference on Machine Learning, ICML 2019, pages 5331-5340. 2019.
[2] Rasool Fakoor, Pratik Chaudhari, Stefano Soatto, and Alexander J. Smola. Meta-g-learning. International Conference on Learning Representations, ICLR 2020. OpenReview.net, 2020.
[3] Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In International Conference on Machine Learning, 2020.
[4] Younggyo Seo, Kimin Lee, Ignasi Clavera, Thanard Kurutach, Jinwoo Shin, and Pieter Abbcel. Trajectory- wise multiple choice learning for dynamics generalization in reinforcement learning. arXiv preprint arXiv:2010.13303. 2020.
[5] Haotian Fu, Hongyao Tang, Jianye Hao, Chen Chen, Xidong Feng, Dong Li, and Wulong Liu. Towards effective context for meta-reinforcement learning: an approach based on contrastive learning. arXiv preprint 2rXiv:2009.1180 13891, 2020.
[6] Bernie Wang. Simon Xu, Kurt Keutzer, Yang Gao, and Bichen Wu. Improving context-based meta-reinforccement learning with self-supervised trajectory contrastive learning. arXiv preprint arXiv:2103.06386 2021.
[7] Lanqing Li, Yuanhao Huang, Mingzhe Chen, Siteng Luo, Dijun Luo, and Junzhou Huang. Prov- ably improved basse offline meta-rl with attention and contrastive learning. arXiv preprint arXiv:2102. 10774, 2021.
[8] Tong Sang. Hongyao Tang. Yi Ma. Jianye Hao, Yan Zheng, Zhaopeng Meng. Boyan Li, and Zhen Wang. Pandr: Fast adaptation to new environments from offline experiences via decoupling policy and environment representations. arXiv preprint arXhh:2004.02877. 2022.
10

---

[9] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, arXiv preprint 2Xiv::807.03748. 2018.
[10] Michael Dennis, Natasha Jaques, Eugene Vinitsky, Alexandre Bayen, Stuart Russell, Andrew Critch, and Sergey Levine. Emergent complexity and zero-shot transfer via unsupervised environment design. Advances neural information processing systems, 3049-13061, 2020.
[11] Minqi Jiang, Edward Grefenstette, and Tim Rocktäschel, Prioritized level replay. International Conference on Machine Learning, pages 4940-4950. PMLR, 2021.
[12] Minqi Jiang, Michael Demnis, Jack Parker-Holder, Jakob Foerster, Edward Grefenstette, and Tim Rock- täschel. Keplay-guided adversarial environment design. Advances in Neural Information Processing Systems, 1111888881977 2021.
[13] Jack Parker- Mingi Jiang, Michael Dennis, Mikayel Samvelyan, Jakob Foerster, Edward Grefen- stette, and Tim Rocktäschel. Evolving curricula with regret- environment design. arXiv preprint arXiv:2203.0130 2022.
[14] Roberta Raileanu and Rob Fergus. Decoupling value and policy for generalization reinforcement learning. In International Conference on Machine Learning, pages 8787-8798. PMLR, 2021.
[15] Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning, arxiv. arXiv preprint arXiv: 1912.01888,,, 2019.
[16] Jiaxian Guo, Mingming Gong, and Dacheng Tao. A relational intervention approach for unsupervised dynamics generalization in model-based reinforcement learning. In International Conference on Learning Representattonss 2022.
[17] Alessandro Sordoni, Nouha Dziri, Hannes Schulz, Geoff Gordon, Philip Bachman, and Remi Tachet Des Combes. Decomposed mutual information estimation for contrastive representation learning. In International Conference on Machine Learning, pages 9859-9869. PMLR, 2021.
[18] David McAllester and Karl Stratos. Formal limitations on the measurement of mutual information. In International Conference on Artificall Intelligence and Statistics, pages 875-884 PMLR, 2020.
[19] Karl Stratos. Mutual information maximization for simple and accurate pf -freeoh--peech induction. arXiv preprint 1rXiv:1804007844 2018.
[20] Juergen Schmidhuber. Evolutionary principles in self-referential learning. 1987.
[21] Sehastian Thrun and Lorien Y. Pratt. Learning to learn. In Springer US, 1998.
[22] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Proceedings the 34th International Conference on Machine Learning, ICML 2017, pages 1126-1135, 2017.
[23] Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. Promp: Proximal meta- policy search. In 7th International Conference on Learning Representations, ICLR 2019, 2019.
[24] H. Liu, R. Socher, and Caiming Xiong. Taming maml: Efficient unbiased meta-reinforcement learning. In ICML, 2019.
[25] A. Gupta, R. Mendonca, Yuxuan Liu, P. Abbeel, and S. Levine. Meta-reinforreemenl learning of structured exploration strategies. In NewrlPS, 2018.
[26] Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-awand dynamics model for generalization in model-based reinforcement learning. CoRR, abs/2005.06800 068000 2020.
[27] Cean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, al. Bootstrap your own latent: new approach to self-supervised learning. arXiv preprint NXNi:2006.0073:. 2020.
[28] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsu- pervised learning of visual features by contrasting cluster assignments. arXiv preprint wrXiv:2006.09882 2020.
[29] Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representations in vector space, arXiv preprint 101,33130,,378,, 2013.

---

[30] Bogdan Mazoure. Remi Tachet des Combes, Thang Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. 2020.
[31] Michael U Gutmann and Aapo Hyvärinen. Noise-contrastive estimation of unnormalized statistical models, with applications natural image statisties. Journal of Machine Learning Research, 13:307-360 -311, 2012.
[32] R Devon Hjelm, Alex Fedorov, Samuel Lavoie Marchildon, Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. 2019.
[33] Jiaming Song and Stefano Ermon. Understanding the limitations of variational mutual information estimators, 2019.
[34] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT Press, 2018.
[35] Zdravko Botev, Dirk Kroese, Reuven Y Rubinstein, and Pierre L'Ecuyer. The cross-entropy method for optimization. In Handbook of statistics. Elsevier, 2013.
[36] Carlos E Garcia, David M Prett, and Manfred Morari. Model predictive control: theory and practice- survey, Automatica, 25(3):335-344 1989.
[37] Wenhao Yu, Jie Tan, C Karen Liu, and Greg Turk. Preparing for the unknown: Learning a universal policy with online system identification. In RSS, 2018.
[38] Charles Packer, Katelyn Gao, Jemej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization in deep reinforcement learning. arXiv preprint arXiv: 12282, 2018.
[39] Wenxuan Zhou, Lerrel Pinto, and Abhinav Gupta. Environment probing interaction policies. In ICLR, 2019.
[40] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06344 2017.
[41] Emanuel Todorov, Tom Frez, and Yuval Tassa. Mujoco: physics engine for model-based control. In IROS, 2012.
[42] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint 1606.014 2016.
[43] Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing. Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real-world environments through meta-reiiforreemen lcarning. In ICLR, 2019.
[44] Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. AAAI, 2018
[45] Kate Rakelly, Aurick Zhou Deirdre Quillen, Chelsea Finn, and Sergey Levine. Efficient off-policy meta- rentfooeffrre learning via probabilistic context variables. In ICML, 2019.
[46] Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne, Journal of Machine Learning Research, 255x)):559--2655 2008
[47] Ian T Jolliffe. Principal component analysis for special types of data. Springer, 2002.
[48] David Barber and Felix Agakov. The im algorithm: A variational approach to information maximization. page 201-208, 2003.
[49] Chris Cremer, Quaid Morris, and David Duvenaud. Reinterpreting weightt autoencoders. 2017.
[50] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
[51] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbecl. High-dimensional continuous control using generalized advantage estimation. In ICLR, 2016.
[52] Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta-learning. arXiv preprint arXiv:1970008345 2019.
12

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [Yes]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Section 4.2.
(b) Did you include complete proofs of all theoretical results? [Yes] See Section 4.2 and Appendix A.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as URL)? [Yes] See Section and Appendix D.
(b) Did you specify all the training details (c.g.. data splits, hyperparameters, how they were chosen)? [Yes] See Appendix D.
(c) Did you report error bars (e.g.. with respect to the random seed after running experi- ments multiple times)? [Yes] See Section
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] See Appendix D.
4. If you using existing assets (e.g., code, data, models) or curating/releasing
assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any assets either in the supplemental material URL? [Yes]
(d) Did you discuss whether and how consent obtained from people whose data you're rou're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
13


Reviews:
[{"review_id": "cn5mV-9jun", "reviewer": "Reviewer_rYrP", "summary": "This paper studies a contextual reinforcement learning (RL) setting where the environment dynamics are parameterized by independent factors, which the authors refer to as “confounders.” In each episode, the underlying factors can vary. They present a method for contextual meta-reinforcement learning (RL) called DOMINO, which learns to encode the RL agent’s current trajectory into a set of independent context vectors. These independent context vectors can then be used as inputs to the transition model in model-based RL (MBRL) and as an input to the policy in model-free RL, thereby providing the agent with an inferred context for the underlying environment factors in any given episode. Importantly, their method assumes the underlying environment factors are similarly independent. The main contributions of the paper are the method, DOMINO, for learning independent context vectors from the trajectory and their analysis and experimental results demonstrating the favorable properties of this method (including improved empirical performance against baselines learning entangled context vectors), when the underlying independence assumptions are valid.", "questions": "- Could the authors elaborate on their choice of T-MCL as the sole baseline in their MBRL experiments? Further details on this choice would benefit the clarity of the experiments section.\n- Further, since the goal of the method is to perform efficient meta-test adaptation via context vectors, why was PEARL chosen over Varibad, which has been shown to provide much more efficient within-episode adaptation compared to PEARL?\n- It seems that the number of context vectors N must be set to the number of environment factors. Is this understanding correct?\n- This paper assumes the environment confounders impact the transition function, but not the reward function. How do the authors view the role of reward generalization in their work? Can DOMINO be expected to work in settings where the environment confounders also impact the reward function?", "limitations": "The core assumption of this work also acts as its primary limitation: The environment factors of variation are assumed to be independent, and their number known a priori. The authors should make an effort to emphasize this limitation and to what extent they believe such an assumption of independence may be applicable in practice.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n- The paper provides a simple method for improving context-aware meta RL in an environment with multiple independent factors of variation that impact the transition dynamics. The method itself is clearly described. This seems to be the first method to directly exploit an explicit assumption of independence among the underlying environment factors of variation.\n- The method performs well against sensible baselines. Importantly the method performs well against an ablation that does not learn disentangled context vectors.\n\nWeaknesses\n\n- The reported results in the Table 1 and 2 have high overlap between the authors’ DOMINO and MINO methods and the baselines. The signficance of these results could be made clearer by reporting the results of a Welch t-test between the proposed method and the baselines.\n- Similarly, the performance comparison plot in Figure 1b should have error bars. It should also state what method of averaging was used for the plotted values\n- The paper can benefit from a full pass to improve the clarity of the writing. There are numerous missing details about basic figures, such as what measure of uncertainty is represented by the error bars for each plot and table. There are also several ambiguous phrasings and sentences with confusing wording. For example\n- A key aspect of this paper is the analysis of InfoNCE as a “loose bound” of the mutual information. However, the authors never define whether this bound is an upper or lower bound. While this detail can be inferred from context, I think it is important to make this point clearer to the reader. Relatedly, the definition of “MI underestimation” in L45 is unclear.\n- Given that the independence assumption is core to this work, it is unclear how significant this setting will be in practice and for future work.\n- Moreover, it seems important for the experiments to assess how valid such an independence assumption is in practice, and crucially, what is the price in performance one might expect to pay for making this assumption. An experiment assessing the performance of DOMINO and MINO on a more complex environment whose underlying factors of variation are not mutually independent would improve this paper by providing a more complete picture of the effectiveness of this method.\n- There seems to be an underlying assumption that the N independent context vectors aim to encode information about the underlying factors of variation in the environment. However, this connection is actually never explicitly made in the writing, making the jump from discussing MI in terms of environment factors to context vectors (4.1 to 4.2) unclear.\n- It seems that DOMINO requires setting the number of context vectors N equal to the number of environment factors of variation. In general, we may not know this value exactly. Adding a sensitivity analysis to how dependent the performance is on setting N to this exact value would provide important information on how applicable this method is in practice.\n\nMinor comments:\n\n- L22: “mythologies” should be “morphologies”.\n- L47-48: “First the context encoder embeds the past state-action pairs into disentangled context vectors” is an inaccurate description, as it must first be optimized to do so (as next described in L48-49).\n- This paper could consider citing related work in unsupervised environment design [1,2,3,4] and more generally, RL work in procedurally-generated environments [5,6]. These works are deeply related as they effectively perform meta-RL over a space of environment variations with an implicitly learned context. Ignoring this line of work seems like a significant oversight.\n\nReferences\n\n[1] Dennis et al, 2020. Emergent Complexity and Zero-shot Transfer via Unsupervised Environment Design.\n\n[2] Jiang et al, 2021. Prioritized Level Replay.\n\n[3] Jiang et al, 2021. Replay-Guided Adversarial Environment Design.\n\n[4] Parker-Holder et al, 2022. Evolving Curricula with Regret-Based Environment Design.\n\n[5] Raileanu et al, 2021. Decoupling Value and Policy for Generalization in Reinforcement Learning.\n\n[6] Cobbe et al, 2019. Leveraging Procedural Generation to Benchmark Reinforcement Learning.", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "cmSlk6WSCB8", "reviewer": "Reviewer_RkbU", "summary": "This paper proposes a decomposed mutual information method to learn disentangled context information, which can generalize reinforcement learning algorithms into unseen environments. The experimental experiments demonstrate that the proposed method can achieve better performance than the previous methods.", "questions": "Please refer to the \"Weakness\" listed above.", "limitations": "Please refer to the \"Weakness\" listed above.", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The writing of this paper is pretty well, and the idea of it is easy to follow.\n2. The figures in this paper are very clear and very well.\n3. The extensive experiments show the effectiveness of the proposed method.\n\nWeakness: \n1. Based on the title, I assume that this study focuses on the meta-reinforcement learning problem. The conventional meta-reinforcement learning methods include an adaptation process, but this paper makes no mention of this process. Additionally, the paper states that it intends to train a general context-encoder to solve the adaptation problem, indicating that the paper's context is the dynamics generalization in reinforcement learning (this paper also mentions it in line 84), which is in contrast to the title of the paper, which refers to meta-reinforcement learning.\n\n2. The second problem of this paper is the novelty. The paper aims to maximize the mutual information between contexts extracted from historical information and the historical trajectories. However, this paper does not make clear the relationship with [1,2,3] which also attempt to maximize the MI between context vector and historical trajectories. Furthermore, this work does not compare the performance with [3] and even does not acknowledge it, despite the fact that [3] focuses on a similar problem to this paper. As a result of the missing contribution and experimental comparisons with [1,2,3], I believe this paper's uniqueness is somewhat limited.\n\n3. The number of learned context vectors $c$ is set as the number of environments in the study, which is the primary hyperparameter of the suggested technique. However, in a real-world setting, the number of environments is not available, making it unfair to compare it to the baseline TMCL, which doesn't rely on such prior information. This increases my concerns about the technical soundness of this paper.\n\nIn conclusion, while the writing and experimental results are excellent, this paper suffers from the aforementioned clarity and novelty issues. If the authors address my concerns in their response, I will consider raising my score.\n\n\n------------------------------------------------- After Rebuttal ------------------------------------------------\n\nI think that the additional experimental results and discussion in the revision resolve my concerns about the clarity problem of the submission, so I increase my score from 4 to 6 accordingly.\n\nMinors: I believe that RIA considers context information and constructs confounder sets with multiple confounders, so I believe that RIA should be discussed in the introduction's confounder discussion (Line 42).\n\n[1] Haotian Fu, Hongyao Tang, Jianye Hao, Chen Chen, Xidong Feng, Dong Li, and Wulong Liu. Towards effective context for meta-reinforcement learning: an approach based on contrastive learning. \n\n[2] Li, L., Huang, Y., Chen, M., Luo, S., Luo, D., & Huang, J. (2021). Provably Improved Context-Based Offline Meta-RL with Attention and Contrastive Learning. arXiv preprint arXiv:2102.10774.\n\n[3] Guo J, Gong M, Tao D. A Relational Intervention Approach for Unsupervised Dynamics Generalization in Model-Based Reinforcement Learning[C]//International Conference on Learning Representations. 2022.", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "_bOgSB4vaC", "reviewer": "Reviewer_oNNM", "summary": "This paper addresses the problem of learning generalizable context in RL. In particular, it suggests learning disentangled context representation of each confounding in the environment using the proposed model, DOMINO, which optimizes decomposed MI objectives. It adopts the contrastive learning method when learning the disentangled context representation, regarding trajectories sampled from the setting of the same confounding as positive pair and of different confounding as negative pair. The authors also provide a theoretical basis for how optimizing their decomposed MI objective can make $I_{NCE}$ a tighter lower bound by alleviating the underestimation of MI. By learning policy conditioning on the learned context vector, DOMINO can achieve higher generalization performance compared to both model-based and model-free baselines.", "questions": "- In Appendix A,\n  - How does the first equility ($q(y | x, y_{2:K}) = p(y) K w_y$) hold?\n  - What is $w_y$?\n  - What does $E$ mean in $I_{NCE}(x;y | E, K)$?\n- In Theorem 1,\n  - Even if the number of confounders increases, the true mutual information $I(c; \\mathcal T)$ does not. Therefore it shows inconsistency to regard $I_{NCE}(c; \\mathcal T)$ and $I_{NCE}(c_i; \\mathcal T)$ as having the same upper bound.\n  - The following inequality also holds in the same setting. $\\sum I_{NCE}(c_i;\\mathcal T | K) = I_{NCE}(c;\\mathcal T| K) \\le log K$ \n- In Figure 6,\n  - DOMINO encodes the context vector of each confounder mass and damping. However, Figure 6 does not show how effectively each confounder was encoded.\n  - For example, Setup0 and Setup4 have the same damping condition and Setup0 and Setup3 have different mass and damping conditions. Then shouldn't Setup0 and Setup4 be located closer than Setup0 and Setup3?", "limitations": "Yes, the authors adequately addressed the limitations and potential negative social impact of their work.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nThe paper is well written and clear to understand. Using contrastive loss when learning disentangled representation of each confounding is novel and intuitive. And it is intriguing to get an idea of sampling negative pairs from different episodes. The experiments are comprehensive and the results are impressive.\n\n\nWeaknesses:\n\nHowever, the proof of Lemma 1 and Theorem 1 lacks mathematical rigor. Also, there is some missing specific information about notations in the proof, thereby undermining the clarity and soundness of the paper (e.g., $w_y$ and $E$). Visualization of the learned context embeddings does not show how effectively each confounding is encoded.\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "3aa4rQl67Fj", "reviewer": "Reviewer_1VJF", "summary": "This paper tackles the problem of generalization in MDPs where the dynamics changes are assumed to be caused by multiple independent factors, denoted as context. The proposed framework (DOMINO) learns a context encoder that maps trajectories to a latent context via decomposed mutual information using noise-contrastive estimation (InfoNCE). The authors combine DOMINO with model-free and model-based RL algorithms, and perform experiments in classic environments, as well as in the Mujoco benchmark, in settings where multiple confounders change simultaneously. Additionally, qualitative visualizations of the latent context vectors are presented using t-SNE.", "questions": "Furthermore, I have the following questions and constructive criticisms:\n\n- “Our goal is to learn a generalized context encoder, which is learned in training process by maximizing the expected rewards R_train in seen environments, and zero-shot generalized to unseen environments for a high expected rewards R_test” \nThis sentence is hard to follow. The objective is also not very clearly defined in Eq. 1. For instance, how is the context $c$ generated? How do the confounders $u$ affect the expectation?\n\n- In Section 4.2, it is not clear the difference between the two context encoders, $g_\\phi$ and $h_w$? Importantly, why is $h_w$ necessary? Notice that $h_w$ does not appear in the pseudo-code of the algorithms in the appendix.\n\n- It is not clear what are the parameters being optimized through the loss functions in Eq. 5, 8, and 9. I suggest specifying which parameters ($\\psi$, $\\phi$, etc.) are involved in the gradients of each loss function.\n\n- “The critic function $\\psi(·, ·)$ measures the similarity between inputs by dot product.” Why was dot product chosen as the similarity measure? Can you elaborate on this decision?\n\n- “The results illustrate that DOMINO learns the policy more efficiently and stably than the baselines.” \nIt is not possible to infer that DOMINO learns more “stably” than the baselines. In fact, sometimes it shows more variance (see Fig. 3 - Halfcheetah-m-d) than the baselines.\n\n- Why does MINO perform better than T-MCL? It seems the improvements of MINO (in comparison to T-MCL) are more important than the improvements of DOMINO (in comparison to MINO).\n\n- In Figure 6, it is not clear what each point represents. Given a trajectory, the context encoder outputs N different context vectors. In this figure, did you concatenate all vectors before applying t-SNE? \n\n- The key characteristic of the proposed method is the fact that it should learn disentangled and independent latent contexts. However, the visualizations using t-SNE were not able to show that. An important result would be to show that indeed each latent context captures one of the confounders of the environment. This could be shown by varying only one of the confounders and observing whether only one of the latent contexts changes accordingly.\n\n- The authors state in the checklist that they have included the code, data, and instructions needed to reproduce the main experimental results in Appendix D. However, only pseudocode of the algorithms and code for the environment were made available. The actual code of the proposed methods is not available.\n\n- In the paper’s abstract, it is said that the open-sourced code and videos are released on their anonymous homepage. However, there are no videos on this page, and only code for the environments is available, not the algorithms/training code.\n", "limitations": "The paper could benefit from a discussion regarding the assumption of independent confounders. For instance, how difficult it would be to adapt the algorithm to the case where we have co-related confounders?", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The idea of capturing the different confounders that may affect the dynamics of the MDP into different latent contexts is novel and interesting.\n- The experimental results show that the proposed method can, in general, achieve better performance than the state-of-the-art.\n\nWeaknesses:\n- The paper needs improvement regarding the clarity of the mathematical definitions, such as the objective functions.\n- It is not clear whether the improvements are because of the decomposed mutual information framework, or because of other algorithmic improvements (see below).\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept