Decomposed Mutual Information Optimization for Generalized Context in Meta-Reinforcement Learning
Anonymous Author(s)
AffiliationAddress
Abstract
Adapting to the changes in transition dynamics is essential in robotic applications.
By learning a conditional policy with a compact context, context-aware metareinforcement learning provides a flexible way to adjust behavior according todynamics changes.
However, in real-world applications, the agent may encountercomplex dynamics changes.
Multiple confounders can influence the transitiondynamics, making it challenging to infer accurate context for decision-making.
This paper addresses such a challenge by D ec O mposed M utual IN formation
O ptimization (DOMINO) for context learning, which explicitly learns a disentangled context to maximize the mutual information between the context and historicaltrajectories, while minimizing the state transition prediction error.
Our theoreticalanalysis shows that DOMINO can overcome the underestimation of the mutualinformation caused by multi-confounded challenges via learning disentangledcontext and reduce the demand for the number of samples collected in variousenvironments.
Extensive experiments show that the context learned by DOMINObenefits both model-based and model-free reinforcement learning algorithms fordynamics generalization in terms of sample efficiency and performance in unseenenvironments.
Open-sourced code and videos are released on our anonymoushomepage.
Introduction
Dynamics generalization in deep reinforcement learning (RL) investigates the problem of training a
RL agent in a few kinds of environments and adapting across unseen system dynamics or structures,such as different physical parameters or robot mythologies.
Meta-Reinforcement Learning (MetaRL) has been proposed to tackle the problem by training on a range of tasks and fast adaptingto a new task with the learned prior knowledge.
However, training in meta-RL requires orders ofmagnitudes more samples than single-task RL since the agent not only has to learn to infer the changeof environment but also has to learn the corresponding policies.
Context-aware meta-RL methodstake a step further and show promising potential to capture local dynamics explicitly by learningan additional context vector from historical trajectories [1, 2, 3, 4].
The historical trajectories aresampled from the joint distribution of multiple confounders, which are the key factors that causethe dynamics changes.
Accordingly, if multiple confounders affect the dynamics simultaneously,the state transition distribution will become highly multi-modal, leading to challenges in extractingaccurate context.
Recent advanced context-aware meta-RL methods [5, 6, 7, 8] further improve meta-RL via contrastivelearning, which optimizes the InfoNCE bound [9] of the mutual information in essence.
Thesemethods show a promising improvement in entangled context learning, which performs well insingle confounded environments.
However, as demonstrated in Figure 1a, in real-world situations for
[Figure]
robotic applications with partially unspecified dynamics, the transition dynamics can be influencedby multiple confounders simultaneously, such as mass changes, damping, friction, or more extremesituations with malfunctional modules like a crippled leg.
For example, when a transportation robotis working in the wild, the load will dynamically change as the task progresses, while the humidityand roughness of the road also may vary.
Moreover, some works also construct a confounder set forunsupervised RL environment generalization[10, 11, 12, 13, 14, 15].
Such changeable environmentsbring 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 thenumber of confounders increases, InfoNCE will be a loose bound of mutual information (MI) withthe samples in limited seen environments, which is called MI underestimation [16].
To tackle thisproblem, we propose a D ec O mposed M utual IN formation
O ptimization (DOMINO) frameworkfor context learning in meta-RL.
The context encoder aims to embed the past state-action pairs intodisentangled context vectors and is optimized by maximizing the mutual information between thedisentangled context vectors and historical trajectories while minimizing the state transition predictionerror.
DOMINO decomposes the full MI optimization problem into a summation of N smaller MIoptimization problems by learning disentangled context.
We then theoretically prove that DOMINOcould alleviate the underestimation bias of the InfoNCE and reduce the demand for the samplescollected in various environments [17, 18].
Last, with the learned disentangled context, we furtherdevelop the context-aware model-based and model-free algorithms to learn the context-conditionedpolicy and illustrate that DOMINO can consistently improve generalization performance in bothways to overcome the challenge of multi-confounded dynamics.
Extensive experiments demonstrate that DOMINO benefits meta-RL on both the generalizationperformance in unseen environments and sample efficiency during the training process under thechallenging multi-confounded setting.
For example, as show in Figure 1b, it achieves 1.5 times performance improvement to T-MCL [3] in the Cheetah domain and 2.6 times performance improvementto T-MCL in the Crippled-Ant domain.
Visualization of the learned context demonstrates that thedisentangled context generated by DOMINO under different environments could be more clearlydistinguished in the embedding space, which indicates its advantage to extract high-quality contextualinformation from the environment.
Related Work2.1 Meta-Reinforcement Learning
Meta-RL extends the framework of meta-learning [19, 20] to reinforcement learning, aiming tolearn an adaptive policy being able to generalize to unseen tasks.
Specifically, meta-RL methodslearn the policy based on the prior knowledge discovered from various training environments andreuse the policy to fast adapt to unseen testing environments after zero or few shots.
Gradientbased meta-RL algorithms [21, 22, 23, 24] learn a model initialization and adapt the parameterswith few policy gradient updates in new dynamics.
Context-based meta-RL algorithms [1, 2, 3, 4]learn contextual information to capture local dynamics explicitly and show great potential to tacklegeneralization tasks in complicated environments.
Many model-free context-based methods areproposed to learn a policy conditioned on the latent context that can adapt with off-policy data byleveraging context information and is trained by maximizing the expected return.
PEARL [1] adapts to a new environment by inferring latent context variables from a small number of trajectories.
Recentadvanced methods further improve the quality of contextual representation leveraging contrastivelearning [5, 6, 7, 8].
Unlike the model-free methods mentioned above, context-aware world modelsare proposed to learn the dynamics with confounders directly.
CaDM [25] learns a global model thatgeneralizes across tasks by training a latent context to capture the local dynamics.
T-MCL [4] furtherenhances this method by combining multiple-choice learning with context-aware world model andachieves state-of-the-art results on the dynamics generalization tasks.
RIA [26] introduces a relationalintervention approach to T-MCL [4] which also attempts to maximize the MI between context vectorand historical trajectories.
However, existing context-based approaches focus on learning entangled contexts.
In a multiconfounding environment, learning entangled contexts requires orders of magnitude higher samplesto capture accurate dynamics information.
This is due to the highly multi-modal state transitionscaused by the combined effects of multiple confounding factors.
To tackle this challenge, DOMINOlearns disentangled context and reduces the amount of demand for diverse samples by decomposedmutual information optimization and thus improves the generalization of the policy to overcome theadaptation 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 varioustasks such as computer vision [27, 28], natural language processing [29, 18], and RL [30], exploitingnoise-contrastive estimation (NCE) [31], InfoNCE [9] and variational objectives [32].
InfoNCEhas gained recent interest with respect to variational approaches due to its lower variance [33] andsuperior 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 [16] first scaffoldsthe total MI estimation into a sequence of smaller estimation problems.
In this paper, since theconfounders in the real world are commonly independent, we simplify the complexity of mutualinformation decomposition and eliminate the need to learn conditional mutual information as asub-term, assuming that multiple confounders are independent of each other.
Preliminaries
We consider standard RL framework where an agent optimizes a specified reward function throughinteracting with an environment.
Formally, we formulate our problem as a Markov decision process(MDP)
[34], which is defined as a tuple ( S , A , p, r, γ, ρ 0 ) .
Here, S is the state space, A is theaction space, p ( s ′ | s, a ) is the transition dynamics, r ( s, a ) is the reward function, ρ 0 is the initialstate distribution, and γ ∈ [0 , 1) is the discount factor.
The goal of RL is to obtain a policy,mapping from states to actions, that maximizes the expected return.
In order to address the problemof generalization, we further consider the distribution of MDPs, where the transition dynamicsp ˜ u
( s ′ | s, a ) varies according to multiple confounders ˜
u
= { u 0 , u 1 , . . .
, u N } .
The confounders can becontinuous random variables, like the mass, damping, random disturbance force, or discrete randomvariables, such as one of the robot’s leg is crippled.
We assume that the true transition dynamicsmodel is unknown, but the state transition data can be sampled by taking actions in the environment.
Given a set of training setting sampled from p (˜ u train ) , 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 test-time, the policy should adapt to the new MDP under the testsetting ˜ u test drawn from p (˜ u test ) .
Our goal is to learn a policy to maximizing the expected return R train condition on the context c which is encoded from the sequences of current state action pairs{ s τ , a τ , s τ +1 } tτ =
t − H in several training scenarios and enable it to perform well and achieve a highexpected return R test in test scenarios never seen before.
[Equation]
Decomposed Mutual Information Optimization for Context Learning
In this section, we first provide a theoretical analysis to show why the multi-confounded environmentsare more challenging.
We find that when the number of confounders increases, the InfoNCE will be a
[Figure]
𝑠
[Figure]
𝑎
[Figure]
Damping
Crippled Leg ⋯
[Figure]
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 contextby decomposed MI optimization.
We theoretically illustrate that the decomposed MI optimizationcan alleviate the underestimation of MI and reduce the demand for the number of samples.
Thedisentangled context c = { c 0 , c 1 , . . .
, c N } is embedded by the context encoder with parameter φ fromthe past state-action pairs τ ∗ =
{ s ∗ l , a ∗ l } t − l =
− H in current episode.
DOMINO explicitly maximizes the
MI between the context c and the historical trajectories T = (cid:8)
τ i (cid:9) M i =1 ( τ i = (cid:8) s il , a il (cid:9) T l =0 ) collectedbased on the combination of multiple confounders
˜ u
= { u 0 , u 1 , . . .
, u N } as same as the currentconfounder setting, while minimizing the state transition prediction error conditioned on the learnedcontext.
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 MIoptimization into smaller ones to alleviate the underestimation of the mutual information and reducethe demand for the number of samples collected in various environments.
4.1 InfoNCE Bound for Mutual Information Optimization
InfoNCE bound I NCE ( x ; y ) is a lower bound of the mutual information
I ( x ; y ) , where NCE stands for
Noise-Contrastive Estimation, is a type of contrastive loss function used for self-supervised learning.
InfoNCE is obtained by comparing pairs sampled from the joint distribution x, y 1 ∼ p ( x, y ) ( y 1 iscalled the positive example) to pairs x, y i built using a set of negative examples, y 2: K ∼ p ( y 2: K )
=
144 (cid:81)
Kk
[Equation]
p (
[Equation] [Equation]
where ψ is a function assigning a similarity score to x, y pairs and K denotes the number of samples.
Through discriminating naturally the paired positive instances from the randomly paired negativeinstances, it is proved to bring universal performance gains in various domains, such as computervision and natural language processing.
Lemma 1 I NCE ( X ; Y | K ) ≤ I ( x ; y ) ≤ log K is a necessary condition for I NCE ( X ; Y | K ) to bea tight bound of I ( x ; y ) .
(see proof in Appendix A)
Some previous context-aware methods learn an entangled context c by maximizing the mutualinformation between the context c 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 I NCE ( c ; T ) , which can beviewed as a contrastive estimation [9] of I NCE ( c ; T ) , and obtain promising improvement in singleconfounded environment.
However, according to Lemma 1, the I NCE ( c ; T ) may be loose if thetrue mutual information I ( c ; T ) is larger than log K , which is called underestimation of the mutualinformation.
Therefore, to make the InfoNCE bound to be a tight bound of I ( c ; T ) , the minimumnumber of samples is e I ( c ; T ) .
In real-world robotic control tasks, the dynamics of the robot is commonly influenced by multipleconfounders ˜ u
= { u 0 , . . .
, u i , . .
. , u j , . .
, u N } simultaneously, under the assumption that the confounders are independent (such as mass and damping), the mutual information between the historicaltrajectories T and the context c can be derived as
[Equation]
As the number of confounders increases, the lower bound of I ( c ; T ) will become larger, and thenecessary condition for I NCE ( T ; c | K ) to be a tight bound of I ( c ; T ) will become more difficultto satisfy.
Since I ( c ; T ) ≥ (cid:80) Ni =
I ( u i ; T ) , to let the necessary consition satisfied, the amount ofdata K must be larger than e (cid:80) Ni
=
I ( u i ; T ) according to Lemma 1.
Thus the demand for data increasessignificantly.
Since the confounders are commonly independent in real-world, can we relax thiscondition by learning disentangled context vectors instead of entangled context intuitively?
4.2 Decomposed MI Optimization
If the context vectors c = { c 0 , c 1 , . . .
, c N } can be independent, then we can ease this problem byapplying the chain rule on MI to decompose the total MI into a sum of potentially smaller MI terms,i.e.,
[Equation]
Theorem
If the context vectors { c 0 , c 1 , . . . , c N } can be independent, then the necessary conditionfor I NCE ( c ; T ) to be a tight bound can be relaxed to I ( c ; T ) ≤
N log K
=
log
K N .
Thus, the needof the number of samples can be reduced from K ≥
e
I ( c ; T ) to K ≥ e 1 N I ( c ; T ) .
Inspired by Theorem 1, we intuitively learn disentangled context vectors and maximize the mutual information between the historical trajectories T and the context vectors { c 0 , . .
. , c N } whileminimizing the I NCE between the context vectors, i.e., to maximize the L NCE
[Equation]
[Equation]
where the I NCE ( c i ; T ) can be obtained with the positive trajectory τ + and negative trajectories
[Equation]
i.e.,
[Equation]
The context c = { c 0 , . . .
, c i , . . .
, c N }
( c i ∈ R m ) is encoded from the past state action pairs τ ∗ incurrent episode by g φ ( · ) .
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 m byh w ( · ) .
The critic function ψ ( · , · ) measures the cosine similarity between inputs by dot product afternormalization.
Under the assumption that the setting of confounders will not change in one episode, to obtain the
I NCE ( c i ; c j ) , we use c + j sampled from same episode like c i as the positive example, and use c − jsampled from different episode as the negative example.
Thus the I NCE ( c i ; c j ) can be derived as
[Equation]
NCE
[Equation]
i
[Equation]
Then, the future state s t +1 can be predicted with the the current state s t , action a t and the disentangledcontext vectors { c 0 t , . . .
, c N t } by the state transition prediction network f ϕ ( · ) .
We aim to minimizethe prediction loss, which is equal to maximizing
[Equation]
where B is the training set, and H is the prediction horizon.
The whole framework of DOMINO isdemonstrated in Figure 2 and the overall objective function of DOMINO is
[Equation]
4.3 Combine DOMINO with Downstream RL Methods
Combination with Model-based RL.
With DOMINO we can learn the context encoder andthe context-aware world model together.
First, the past state-action pairs are encoded into thedisentangled context vectors by the context encoder.
According to the learned context, the transitionprediction network predicts the future states of different actions.
In particular, we use the crossentropy method (CEM) [35], a typical neural model predictive control (MPC) [36] method, toselect actions, in which several candidate action sequences are iteratively sampled from a candidatedistribution, which is adjusted based on best-performing action samples.
The optimal action sequencea t : t + T .
= { a t , . . .
, a t + T } can be obtained by
[Equation]
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.1.
As for the adaptation process, thepolicy and context encoder zero-shot adapts to the unseen confounders setting u test , and we use thesame adaptive planning method as T-CML[3], which selects the most accurate prediction head over arecent experience condition on the inferred context.
The details is introduced in Appendix D.
Combination with Model-free RL.
Previous works show that a policy learned by model-freemethod can be more robust to dynamics changes when it takes the contextual information as anadditional input [37, 38, 39].
Motivated by this, we investigate whether the context encoder learnedby DOMINO can be used as a plug-and-play module to improve the final generalization performanceof model-free RL methods.
We concatenate the disentangled context encoded by a pre-trainedcontext encoder from DOMINO and the current state-action pairs, and learn a conditional policyπ ( a t | s t , c 0 , . . .
, c N ) .
We use the Proximal Policy Optimization (PPO) method to train the agent [40],which learns the policy by maximizing
[Equation]
11) where ˆ A t is the estimation of the advantage function at timestep t .
We provide detailed pseudo-codein the Appendix B.2.
Experiments
In this section, we evaluate the performance of our DOMINO method to answer the followingquestions: (1) Can DOMINO help the model-based RL methods overcome the multi-confoundedchallenges in dynamics generalization (see comparison with ablation in Figure 3 and Figure 4)?
[Figure]
Can the context encoder learned by DOMINO be used as a plug-and-play module to improve thegeneralization abilities of model-free RL methods in multi-confounded environments (see comparisonwith ablation in Table 1 and Table 2)?
Can the proposed decomposed MI optimization benefit theforward prediction of the world model?
(see Figure 5) (4) Does the disentangled context extract moremeaningful contextual information than entangled context (see Figure 6)?
5.1 Setups
We demonstrate the effectiveness of our proposed method on 8 benchmarks, which contain 6 typicalrobotic control tasks based on the MuJoCo physics engine [41] and 2 classical control tasks (CartPoleand Pendulum) from OpenAI Gym [42].
Different from previous works, all the environmentsare influenced by multiple confounders simultaneously.
In our experiments, we modify multipleenvironment parameters at 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 2 more difficultenvironments (Crippled Ant and Crippled HalfCheetah) affected by both continuous and discreteconfounders.
As for Cartpole, we modify the push force f
and the pole length l simultaneously.
As for Pendulum, we scale the pendulum mass by scale factor m and modify the pendulum lengthl .
As for Hopper, Half-cheetah, Ant, and Slimhumanoid, we scale the mass of every rigid link byscale factor m , and scale damping of every joint by scale factor d .
As for Crippled Ant and Crippled
Half-cheetah, we scale the mass of every rigid link by scale factor m , scale damping of every joint byscale factor d , and randomly select one leg, and make it crippled.
The detailed settings are illustratedin Appendix C (Table 3).
We implement these environments based on the publicly available codeprovide by [43, 4], and we also open-source the code of the multiple-confounded environments 1 .
For both training and testing phase, we sample the confounders at the beginning of each episode.
During training, we randomly select a combination of confounders from a 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] as the key baseline in comparison with model-based methods,which achieves the state-of-the-art results in zero-shot dynamics generalization tasks.
We alsoconsider an ablation version of DOMINO as a baseline (denoted as MINO) to show the effectivenessof the decomposed MI optimization, which optimizes the MI and predicts the future states with anentangled context without decomposition.
Results.
Figure 3 shows the average return during the learning process in the training environments.
The results illustrate that DOMINO learns the policy more efficiently and stably than the baselines.
Figure 4 shows the generalization performance tested in the unseen environments.
The results showthat DOMINO surpasses T-MCL in terms of the generalization performance and the learning sampleefficiency.
This demonstrates that the disentangled context improves the context-aware world model.
Especially, the performance gain becomes much more significant in more complex environments
[Figure] [Figure]
.g.
, long-horizon, sensitive, and high-dimensional domains like Cripple-Ant, Ant, and Hopper).
Forexample, DOMINO achieves about 2.6 times performance improvement to T-MCL in Cripple-Ant,which is one of the most difficult environment in the experiments, whose leg will randomly becrippled, and its mass will be changed in testing.
More details are shown in the Appendix D.
5.3 Comparison with Model-free Methods
We also verify whether the learned disentangled context is useful for improving the generalizationperformance 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 as conditional input(PPO+DOMINO), is compared with several context-conditional policies [45, 39].
Specifically, weconsider combining the PPO with the context learned by T-MCL (PPO+T-MCL), which learns thecontext encoder via a context-aware world model and achieves the state-of-the-art performance ondynamics generalization.
We also consider PEARL [45], which learns probabilistic context variableby maximizing the expected returns.
We further develop an ablation version of DOMINO, whichoptimize the MI with entangled context (PPO+MINO) as a baseline to illustrate the effectiveness ofthe decomposed MI optimization.
We provide more detailed explanations in Appendix D and morecomparison experiments with other baselines are shown in Appendix E.
Results .
Table 1 and Table 2 show the performance of various model-free RL methods on bothtraining and test environments.
PPO+DOMINO shows superior performance and shows bettergeneralization performances than previous conditional policy methods, implying that
the proposed
DOMINO method can extract contextual information more effectively than both the context learnedby 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, especiallyin 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 thatcompared to PEARL, the context learned by T-MCL and DOMINO has better performance whichimplies that the state transition perdition can help to extract contextual information more effectively.
5.4 Disentangled Context Analysis
Prediction errors .
To show that our method indeed helps with a transition prediction, we comparebaseline methods with DOMINO in terms of prediction error across 8 environments with varyingmultiple confounders.
As shown in Figure 5, our model demonstrates superior prediction performance,which indicates that the learned context capture better contextual information compare with entangledcontext (see more results in Appendix E.1).
Visualization .
We visualize the disentangled context learned by DOMINO via t-SNE [46] andcompare it with the entangled context learned by T-MCL.
We run the learned policies under 5randomly sampled setups of multiple confounders and collect 200 trajectories for each setting.
Further, we encode the collected trajectories into context in embedding space and visualize via
The transition dynamics will change in both training and test environments in every episode
[Table] [Table]
[Figure]
t-SNE.
As shown in Figure 6, we find that the disentangled context vectors encoded from trajectoriescollected under different confounder settings could be more clearly distinguished in the embeddingspace than the entangled context learned by T-MCL.
This indicates that DOMINO extracts highquality task-specific information from the environment compared with T-MCL.
We provide morevisualization results based on both t-SNE [46] and PCA [47] in Appendix F.2.
[Figure]
, leg = 1 ), Setup2 ( m = 1 . 25 , leg = 0 ), Setup3 ( m = 0 .
85 , leg = 0 ), Setup4 ( m = 1 . 0 , leg = 2 ).
Conclusion
In this paper, we propose a decomposed mutual information optimization (DOMINO) frameworkto learn the generalized context for zero-shot dynamics generalization.
The disentangled context islearned by maximizing the mutual information between the context and historical trajectories whileminimizing the state transition prediction error.
By decomposing the whole mutual informationoptimization problem into smaller ones, DOMINO can reduce the need for samples collectedin various environments and overcome the underestimation of the mutual information in multiconfounded environments.
Extensive experiments illustrate that DOMINO benefits the generalizationperformance in unseen environments with both model-based RL and model-free RL.
We believe ourwork can lay the foundation of dynamics generalization in complex environments.
Limitations and Negative Social Impact.
DOMINO sets the number of disentangled context vectorsas a hyper-parameter as same as the number of confounders in the environments.
Based on ourframework, capturing the number of confounders automatically, could be future works.
We believethat DOMINO will not cause any negative social impact.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions 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]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Section4.2.(b)
Did you include complete proofs of all theoretical results?
[Yes] See Section 4.2 andAppendix A.
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 a URL)?
[Yes] See Section 5.
and Appendix D.
(b) Did you specify all the training details (e
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 experiments multiple times)?
[Yes] See Section 5.(d)
Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix D.
If you are using existing assets (e
g., code, data, models) or curating/releasing new 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 new assets either in the supplemental material or as a URL?
[Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
[Table]
information or offensive content?
[N/A]
[Table]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
[Table]
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation?
[N/A]
[Table]
A Derivations
A.1 Proof of Lemma 1
According to the Barber and Agakov’s variational lower bound [48], the mutual information I ( x ; y )between x and y can be bounded as follows:
[Equation]
where q is an arbitrary distribution.
Specifically, q ( y | x ) is defined by independently sampling a set ofexamples { y 1 , . . .
, y K } from a proposal distribution π ( y ) and then choosing y from { y 1 , . . .
, y K }in proportion to the importance weights w y =
[Equation]
ψ ( x,yk ) , where ψ is a function that takes x andy and outputs a scalar.
According to the section 2.3 in [9], by setting the proposal distribution asthe marginal distribution π ( y ) ≡ p ( y ) , the unnormalized density of y given a specific set of samplesy 2: K = { y 2 , . . .
, y K } and x is:
[Equation]
) where K denotes the numbers of samples.
According to the equation 3 of section 2 in [49], theexpectation of q ( y | x, y 2: K ) with respect to resampling of the alternatives y 2: K from p ( y ) producesa normalized density:
[Equation]
With Equation 14 and Jensen’s inequality applied in Equation 12, we have
[Equation]
It is obviously that
[Equation]
x
[Equation]
x [Equation] 2 e ψ ( x,yk ) ≤ 1 , thus we have [Equation]
[Equation]
With Equation 16, we have
[Equation]
[Equation]
Therefore, we have
[Equation]
If I ( x, y ) > log K , then I ( x, y ) > log K ≥ I NCE ( x ; y | ψ, K ) , and I NCE will be a loose bound.
Thus, I NCE ≤ I ( x, y ) ≤ log K is the necessary condition for I NCE to be a tight bound of I ( x, y ) .
A.2 Detailed derivation of Theorem 1
As the number of confounders increases, although the true mutual information I ( c ; T ) does notincrease, the necessary condition of I NCE to be a tight lower bound of I NCE becomes more difficultto satisfy, and the demand of data increases significantly.
As for an entangled context, the necessary condition of the InfoNCE lower bound I NCE ( c ; T ) to bea tight bound is
[Equation]
Since I ( c ; T ) ≥ (cid:80) Ni =
I ( u i ; T ) , to let the above condition satisfied, the amount of data K mustsatisfy
[Equation] [Equation]
Therefore, if the number of confounders increases, then the demand for data will grow exponentially.
When data is not rich enough, the nesseray condition may not be satisfied.
The InfoNCE lower bound
I NCE ( c ; T ) may be loose, that is I NCE ( c ; T ) may be much smaller than the true mutual information
I ( c ; T ) , thus the MI optimization based on I NCE ( c ; T ) will be severely affected.
I NCE ( c i ; T ) is the lower bound of I ( c i ; T ) and the necessary condition of I NCE ( c i ; T ) to be atight bound of I ( c i ; T ) is
[Equation]
As for disentangled context c = { c 1 , c 2 , · · · , c N } , we then derive the necessary condition of I ( c, T )to be a tight lower bound of I ( c, T ) :
With the assumption that the contexts { c 1 , c 2 , · · · , c N } are independent to each other, then I ( c ; T )could be derived as (cid:80) I
( c i ; T ) .
Therefore, under the confounder independent assumption, let
I NCE ( c ; T ) be a tight bound is only necessary to let every I NCE ( c i ; T ) to be a tight bound.
If every I NCE ( c i ; T )( i = 1 , 2 , . . . , N ) is a tight bound, then we have [Equation]
[Equation]
under the confounder independent assumption, we have
[Equation]
Thus, the necessary condition of I NCE ( c ; T ) to be a tight bound of I ( c ; T ) could be relaxed to
[Equation]
Therefore, by decomposing the MI estimation under the confounder independent assumption, thedemand of the amount K of data could be reduced from K ≥
e I ( c ; T ) to K ≥
e 1 N I ( c ; T ) .
Andwith I ( c ; T ) ≥ (cid:80)
Ni
=
I ( u i ; T ) , specificly, the the amount K of data could be reduced from K ≥e (cid:80) Ni =
I ( u i ; T ) to K ≥ e 1 N (cid:80) Ni =
I ( u i ; T ) .
B Pseudo-codeB.1 Combination with model-based methods
We provide the pseudo-code of DOMINO combined with model-based methods.
Firstly, the past stateaction pairs are encoded into the disentangled context vectors by the context encoder.
According to thelearned context, the transition prediction network predicts the future states of different actions.
Then,the context encoder is optimized by maximizing the mutual information between the disentangledcontext vectors and historical trajectories while minimizing the state transition prediction error.
Inparticular, we use the cross entropy method (CEM) [35], a typical neural model predictive control(MPC)
[36] method, to select actions, in which several candidate action sequences are iterativelysampled from a candidate distribution, which is adjusted based on best-performing action samples.
Algorithm
Inputs : learning rate α ,maximum number of iteration P , batch size B , the number of past observations H past , maxium rollout step max _ step and the number of future observations H future .
Initialize parameters of prediction network ϕ , context encoder φ .
Initialize replay buffer D ← ∅ .
for P iterations do // C OLLECT TRAINING SAMPLES step =
V =while step ≤ max _ step do
Sample u V ∼ p
u train ( u ) .
V = V +for t = 1 to TaskHorizon do step = step +
Get context latent vectors c t 0 , c t 1 , . .
c t N = g
( τ t ; φ ) , τ t = { s l , a l } t − 1 l = t − H past
Collect samples { ( s t , a t , s t +1 , r t , τ t ) } from the environment using the planning algorithm based on CEM with the context vectors end for
Update D u V ← D u V ∪ { ( s t , a t , s t +1 , r t , τ t ) } end while //
U PDATE DYNAMICS MODELS AND ENCODER
Initialize batch B ← ∅ .
for i = 1 to B do sample V ∗ from [0, V max ]
Sample { s t , a t , s t +1 , r t , τ t } from D u V∗
Sample positive trajectories τ + from D u V∗
Sample negative trajectories (cid:8) τ − k (cid:9) K k =2 from D u V !
= V∗
Get context latent vectors c t 0 , c t 1 , . .
c t N = g
( τ t ; φ )
Update B ← B ∪ { ( s t , a t , s t +1 , r t , τ t ) } end for
L pred ← E τ ∗ ∼B (cid:104) − 1 H (cid:80) t
+ H future − 1 λ = t log f ϕ ( s λ +1 | s λ , a λ , ( c 0 λ , . . .
, c N λ ))
(cid:105)
L NCE ← (cid:80) Ni I NCE ( c i ; T ) − (cid:80) Nj (cid:80)
Ni =0 ,i ̸
=
j I NCE ( c i ; c j )
Update φ ← φ − α ∇ φ L NCE
Update φ ← φ − α ∇ φ L pred
Update ϕ ← ϕ − α ∇ ϕ L pred end for
B.2 Combination with model-free methods
We provide the pseudo-code for the combination between DOMINO and the model-free method,which uses the context encoder learned by DOMINO as a plug-and-play module to extract accuratecontext.
We concatenate the disentangled context encoded by a pre-trained context encoder from
DOMINO and the current state-action pairs, and learn a conditional policy π ( a t | s t , c 0 , . .
. , c N ) .
Wechoose the Proximal Policy Optimization (PPO) method [40] to train the agents.
Algorithm
Inputs :Maximum number of iteration P , number of actor updates M , number of critic updates B , the KL regular coefficient λ , scaling coefficient α and the learning rate β .
Initialize parameters of policy network θ , value network ξ , and context encoder φ .
for P iterations do
Encode disentangled context vectors c t 0 , c t 1 , . .
c t N by the learned context encoder g φ ( · )
Run policy π θ for T timesteps, collecting {{ s t , c t 0 , c t 1 , . .
c t N , a t , r t }} Tt =
Estimate advantages ˆ A t = (cid:80) t ′ >t γ t ′ − t r t ′ − V ξ ( s t , c t 0 , c t 1 , . . . ,
c t N ) π old ← π θ for M updates do
[Equation]
θ
[Equation]
end for for B updates do
[Equation]
, c t 1 , . . . ,
c t N ))
ξ ← ξ − β ∇ ξ L BL end for if KL[ π old | π θ ] > β high KL target then λ ← αλ else if KL[ π old | π θ ] < β low KL target then λ ← λ/α end if end for
C Details about the testing environments
[Figure]
For CartPole environments, we use open-source implementation of CartPoleSwingUp-v2 2 , whichis the modified version of original CartPole environments from OpenAI Gym.
The objective of
CartPole task is to swing up the pole by moving a cart and keep the pole upright.
For our experiments,we modify the push force f
and the pole length l simultaneously.
As for Pendulum, we scale thependulum mass by scale factor m and modify the pendulum length l .
For Pendulum environments,we use the open-source implementation of from the OpenAI Gym.
The objective of Pendulum is toswing up the pole and keep the pole upright within 200 timesteps.
We scale the pendulum mass byscale factor m and modify the pendulum length l .
As for Hopper, Half-cheetah, Ant, and Slimhumanoid, we use the environments from MuJoCophysics engine 3 , and scale the mass of every rigid link by scale factor m , and scale damping of
[Table]
[Table] [Table]
[Figure]
[Table]
every joint by scale factor d .
As for Crippled Ant and Crippled Half-cheetah, we scale the massof every rigid link by scale factor m , scale damping of every joint by scale factor d , and randomlyselect one leg, and make it crippled.
The objectives of these tasks are to move forward as fast aspossible while minimizing the action cost.
The detailed settings are illustrated in Table 3.
We providethe pyTorch-style pseudo-code for multi-confounded environments in Listing 1 and Listing 2.
Weimplement these environments based on the publicly available code provide by [43, 4], and we alsoopen-source the code of the multiple-confounded environments 4 .
For both the training and testingphase, we sample the confounders at the beginning of each episode.
During training, we randomlyselect a combination of confounders from a training set.
At test time, we evaluate each algorithm in unseen environments with confounders outside the training range.
We also provide the PyTorch-stylepseudo-code for the dynamics change based on Mujoco engine.
D Implementation details
D.1 Combination with Model-based RL
The context encoder is modeled as multi-layer perceptrons (MLPs) with 3 hidden layers and Noutput heads which are single-layer MLPs.
Every disentangled context vector is produced as a10-dimensional vector by the 3 hidden layers and a specific output head.
Then, the disentangledcontext vectors are used as the additional input to the prediction network, i.e., the input is givenas a concatenation of state, action, and context vector.
We use H past = 10 for the number ofpast observations and H future = 5 for the number of future observations.
The prediction networkis modeled as multi-layer perceptrons (MLPs) with 4 hidden layers of 200 units each and Swishactivations.
For each prediction head, the mean and variance are parameterized by a single linearlayer that takes the output vector of the backbone network as an input.
To train the prediction network,we collect 10 trajectories with 200 timesteps from environments using the MPC controller and trainthe model for 50 epochs at every iteration.
We train the prediction network for 10 iterations for everyexperiment.
We evaluate trained models on environments over 8 random seeds every iteration toreport the testing performance.
The Adam optimizer [50] is used with a learning rate 1 × 10 − 4 .
Forplanning, we use the cross entropy method (CEM) with 200 candidate actions for all the environments.
The horizon of MPC is set as 30.
D.2 Combination with Model-free RL
We train the model-free agents for 5 million timesteps on OpenAI-Gym and MuJoCo environments(i.e.
, Hopper, Half-cheetah, Ant, Crippled Half-cheetah, Crippled Half-Ant, Slim-Humanoid) and0.5 million timesteps on CartPole and Pendulum.
The trained agents are evaluated every 10,000timesteps over 5 random seeds.
We use a discount factor γ = 0 .
99 , a generalized advantage estimator[51] parameter λ = 0 .
95 and an entropy bonus of 0.01 for exploration.
In every iteration, the agentrollouts 200 timesteps in the environments with the learned policy, and then it will be trained for 8epochs with 4 mini-batches.
The Adam optimizer is used with the learning rate 5 × 10 − 4 .
D.3 Details of InfoNCE
We provide detailed pseudocode for the calculation of InfoNCE bound.
Specifically, the temperatureτ is set as 0.004 to the calculation of I ( c i , T ) and is set as 0.1 to calculation of I ( c i , c j ) .
[Table]
D.4 Details of the adaptive planning used in adaption process
The prediction model has 3 output head head 0 , head 1 , head 2 which are used for selecting actions by planning.
The adaptive planning method selects the most accurate prediction head over a recent experience.
Given N past transitions, we select the prediction head h ∗ by
[Equation]
where ℓ is the mean square error function.
All the hyper-parameter is set as same as T-MCL[3].
E Additional Results
E.1 Prediction Error
As shown in Figure 7, DOMINO has a smaller prediction error compared to T-MCL and its ablationversion MINO (optimize MI with entangled context), indicating that the learned context can effectivelyhelp predict the future state more accurately, which is the key to the performance of the model-basedplanning.
[Figure]
E.2 Comparison with RIA
We compare DOMINO with RIA and TMCL together under multi-confounded setting, the environments including Hopper- m - d , Halfcheetah- m - d , Slim-humanoid- m - d and Pendulum- m - l .
As shownin Figure 8, DOMINO achieves better generalization performance than RIA and TMCL, especially incomplex environments like Halfcheetah- m - d and Slim-humanoid- m - d .
E.3 Sensitivity Analysis of the hyper-parameter N
We compare the performance of DOMINO with different hyper-parameter N , which is equal or notequal to the number of confounders in the environment.
In this experiment, the confounder is thedamping, mass, and a crippled leg (number of confounders is 3), and we compare the performanceof DOMINO with different hyper-parameter N = 1 , 2 , 3 .
As shown in Figure 9, even though thehyper-parameter N is not equal to the ground truth value of the confounder number, DOMINO alsobenefits the context learning.
In practice, under a conservative setup of hyper-parameter N , DOMINOcan also benefit the context learning compared to the baselines like TMCL.
[Figure] [Figure]
F Visualization
F.1 Verifying whether the contexts is disentangled
We add an additional experiment to show that the context vectors inferred by DOMINO are disentangled well.
We vary only one of the confounders and observe the changes of N disentangledvectors.
In this experiment, we set up two different confounders: mass m and damping d .
Under the
DOMINO framework, the context encoder inferred two disentangled context vectors: context 0 andcontext 1.
As shown in Figure 11 and Figure 10, the context 1 is more related to damping.
Whenthe confounders are set as the same mass but different damping, the visualization result of context1 under different settings are separated clearly from each other, while under the same damping butdifferent mass settings, the visualization result of context 1 is much more blurred from each other.
Similarly, context 0 is more related to mass.
When the confounders are set to the same dampingbut different mass, the visualization result of context 0 under different settings is separated clearlyfrom each other, while under the same mass but different damping settings, the visualization result ofcontext 0 is less different from each other.
F.2 Visualization of the whole context
Visualization .
We visualize the whole context which is a the concatenation of the disentangledcontexts learned by DOMINO via t-SNE [46] and compare it with the entangled context learnedby T-MCL.
We run the learned policies under 5 randomly sampled setups of multiple confounders
[Figure] [Figure]
and collect 200 trajectories for each setting.
Further, we encode the collected trajectories intocontext in embedding space and visualize via t-SNE [46] and PCA [47].
As shown in Figure 12 and
Figure 13, we find that the disentangled context vectors encoded from trajectories collected underdifferent confounder settings could be more clearly distinguished in the embedding space than theentangled context learned by T-MCL.
This indicates that DOMINO extracts high-quality task-specificinformation from the environment compared with T-MCL.
Accordingly, the policy conditioned onthe disentangled context is more likely to get a higher expected return on dynamics generalizationtasks, which is consistent with our prior empirical findings.
[Figure] [Figure]