Improving Generalization in Meta-RL with Imaginary Tasks from Latent Dynamics Mixture
Authors: Suyoung Lee
Sae-Young Chung
Venue: NeurIPS 2021 Poster (2021)
Keywords: Reinforcement Learning
Meta-learning
Generalization

Abstract: The generalization ability of most meta-reinforcement learning (meta-RL) methods is largely limited to test tasks that are sampled from the same distribution used to sample training tasks. To overcome the limitation, we propose Latent Dynamics Mixture (LDM) that trains a reinforcement learning agent with imaginary tasks generated from mixtures of learned latent dynamics. By training a policy on mixture tasks along with original training tasks, LDM allows the agent to prepare for unseen test tasks during training and prevents the agent from overfitting the training tasks. LDM significantly outperforms standard meta-RL methods in test returns on the gridworld navigation and MuJoCo tasks where we strictly separate the training task distribution and the test task distribution.

Improving Generalization in Meta-RL with Imaginary Tasks from Latent Dynamics Mixture
Suyoung Lee"
Sau-Young Chung KAIST schungCkaist. ac.
KAIST
suyoung. 10kaist. ac.
Abstract
The generalization ability of most meta-reinforcement learning (meta-RL) methods is largely limited to test tasks that are sampled from the same distribution used to sample training tasks. To overcome the limitation, propose Latent Dynamics Mixture (LDM) that trains a reinforcement learning agent with imaginary tasks generated from mixtures of learned latent dynamics. By training policy mixture tasks along with original training tasks, LDM allows the agent to prepare for unseen test tasks during training and prevents the agent from overfitting the training tasks. LDM significantly outperforms standard meta-RL methods in test returns the gridworld navigation and MuloCo tasks where we strictly separate the training task distribution and the test task distribution.
Introduction
Overfitting and lack of generalization ability have been raised as the most critical problems of deep reinforcement learning (RL) [5] 30
47 52]. Numbers of meta-reinforcemen learning (meta-RL) methods have proposed solutions to the problems by meta peeta-triining a policy that easily adapts unseen but similar tasks. Meta-RI. trains agent in multiple sample tasks construct an inductive bias over the shared structure across tasks. Most meta-RI. works evaluate their agents on test tasks that sampled from the same distribution used sample training tasks. Therefore, the
34
38
vulnerability of meta-RL to test-time distribution shift is hardly revealed [2] 26 29]30]
One major category of meta-RL is gradient-based meta-RL that learns initialization of model such that few steps of policy gradient are sufficient to attain good performance
new task
99, 36. 39 40 (55). Most of these methods require many test-time rollouts for adaptation that may be costly in real environments. Moreover, the networks are composed of feedforward networks that make online adaptation within a rollout difficult.
Another major eategory of meta-RI. is context-based meta-R RI. that tries to learn the tasks' structures by utilizing recurrent or memory- augmeneed models 6
13 23
25
32
35| 49
[66] A context-based
meta-RI. agent encodes its collected experience into a context. The policy conditioned on the context is trained to maximize the return. These methods have difficulties generalizing unseen out-of- distribution (OOD) tasks mainly because of two reasons. (1) The process of encoding unseen task dynamics into a context is not well generalized. (2) Even if the unseen dynamics is well encoded, the policy that has never been trained conditioned on the unseen context cannot interpret the context to output optimal actions.
We propose Latent Dynamics Mixture (LDM), novel meta-RL method that overcomes the aforemen- tioned limitations and generalizes to strictly unseen test tasks without any additional test-time updates. LDM based on variational Bayes-adaptive meta-RI. that meta-learns approximate inference on latent belief distribution over multiple reward and transition dynamics [56]. We generate imaginary
Corresponding author.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

LDM
variBAD
RL*
RL2 Oracle
30
Mean Renure
(a) Musin
(b) Mses
(c) Mean returns achieved for each task at the N-th episode.
Figure 1: The gridworld example for the problem setup. (a) Training MDPs M goal is located at one of the 18 shaded states. During training, the agent has to navigate to discover the unknown goal position randomly sampled the beginning of each task. (b) Test MDPs Missr' test goal is located at one of the 27 shaded states disjoint to the set of training goals in (a). The agent does not have access to the tasks in Must during training. The solid lines and dashed lines represent examples of optimal paths with and without the knowledge of the true goals, respectively, o and x symbols represent the initial state and the hidden goal state, respectively.
tasks using mixtures of training tasks' meta-learned latent beliefs. By providing the agent with imaginary tasks during training, the agent can train its context encoder and policy given the context for unseen tasks that may appear during testing. Since LDM prepares for the test during training, does not require additional gradient adaptation during testing.
For example, let there be four types of training tasks, each of which must move east, north, west, and south. By mixing the two tasks moving east and north, we may create a new task of moving northeast. By mixing the training tasks in different weights, we may create tasks with goals in any direction.
We evaluate LDM and other ml. methods on the gridworld navigation task and MuloCo meta-RI. tasks, where we completely separate the distributions training tasks and test tasks. We show that LDM, without any prior knowledge on the distribution of test tasks during training, achieves superior test returns compared to other meta-RL. methods.
2 Problem Setup
Our work is motivated by the meta-learing setting of variBAD [56], therefore follow most of the
problem setup and notations except for a key difference that the test and training task distributions
strictly disjoint in our setup. Markov decision process (MDP) M (S,A, R, T, To, 7, H) consists of a set of states S, a set of actions A. a reward function R(re+ 18,, 81+-1 ), transition function T (st+11/8, a,), initial state distribution (5))) discount factor 7 and time horizon H.
During meta-training, a task (or a batch of tasks) sampled following p(M) over the set of MDPs M at every iteration. Each MDP M, (S, AA Tk, To. 7, H) has individual reward function R& (e.g., goal location) and transition function T& (e.g., amount of friction), while sharing some general structures. We assume that the agent does not have access to the task index k which determines the MDP. At meta-test, standard meta-RL methods evaluate agents on tasks sampled from the same distribution that is used sample the training tasks. To evaluate the generalization ability of agents in environments unseen during training, we split M into two strictly disjoint training and test sets of MDPs, i.e., M -Mrriin U Miss and Minin M 0. The agent does not have any prior information about Misst and cannot interact in Miss during training.
Since the MDP is initially unknown, the best the agent can do is update its belief (R, about the environment according to its experience T; = {so, ag: F1, $1, 01 , St.. According to the Bayesian RL formulation, the agent's 's belief about the reward and transition dynamics timestep t can be formalized as a posterior over the MDP given the agent's 's trajectory, by RR T) p(R, T)r:). By augmenting the belief to the state, Bayes-Adaptive MDP (BAMDP) can be constructed [7]. The agent's 's goal in a BAMDP is to maximize the expected return while exploring the environment by minimizing the uncertainty about the initially unknown MDP.
The inference and posterior update problem in a BAMDP can be solved by combining meta-learning and approximate variational inference [56). An inference model encodes the experience into low-

---

Policy Network
Normal Workers W03 in MC
Latent Dynamics Mixture
Mixture Worker
Encador
1P in
Policy
Enc p
Pol.
02
X (/
Decoder St-1--
Pt.
Dynamics Network
Dec.
Eacodor-u
Encodar p
mmmm (B
26.
by
no.
(ml]Fa) mi
tytni Aelm)
1
40g
bg
9
Decoder
Enc-p
Pol.
5/61
P02
TJ+1
Enc-v - ml
(A]
Pol.
Der.
Dby
3/61
Intur dropout
Figure 2: Imaginary task generation from latent dynamics mixture. We train a normal workers and a mixture worker in parallel. Each normal worker W trains a policy network and a latent dynamics network on its sampled MDP M E Mirain A mixture latent model rit, is generated as a weighted sum of normal workers' latent models m,
(6) All workers share single policy network and a single latent dynamics network. We feed this mixture belief to the latent dynamies network's
learned decoder to generate new reward f and construct an imaginary task M.
dimensional stochastic latent variable m to represent the posterior belief over the MDPs The reward and transition dynamics can be formulated as shared functions across MDPs: R( [tc ae St+1; m) and T(s++1|52 Qe m.. Then the problem of computing the posterior p(R, T).;) becomes inferring the posterior pTm])==)) over m. By conditioning the policy on the posterior pm [xx). an approximately Bayes-optimal policy can be obtained.
Refer to Figure T for the gridworld navigation example that is the same as the example used in variBAD (56) except for the increased number of cells and that the task set M is divided into disjoint Murin and MMes. A Bayes--ptimal agent for a task in Mirain first assigns a uniform prior to the goal states of Murin (Figure Lap and then explores these states until it discovers a goal state as the dashed path in Figure Ial If this agent, trained to solve the tasks in Musin only, is put in to solve task in Must without any prior knowledge (Figure Ib], the best policy an agent can take first is to maintain its initial belief learned in Murain and explore the goal states of M Once the agent realizes that there no goal states in Mirainn it could start exploring the states that are not visited (i.e., .M- - Miain), and discover an unseen goal state in Meeu However, it is unlikely that the agent trained only in Musin will encode its experience into beliefs for unseen tasks accurately and explore the goal states of Miest efficiently conditioned on the unseen context without any prior knowledge or test-time adaptation.
3 Latent Dynamics Mixture
Our work aims to train an agent that prepares for unseen test tasks during training as in Figure [b We provide the agent during training with imaginary tasks created from mixtures of training tasks' latent beliefs. By training the agent to solve the imaginary tasks, the agent learns to encode unseen dynamics and to produce optimal policy given the beliefs for tasks not only in Musin but also for more general tasks that may appear during testing.
Refer Figure 2 for an overview of the entire process. We train n normal workers W(L) ,.... (n) and mixture worker W in parallel. For the convenience of explanation, we first focus on the case with only one mixture worker. At the beginning of every iteration, we sample MDPs M(1) ...., (n) from Musin and assign each MDP M to each normal worker Wii), All normal and mixture workers share a single policy network and single latent dynamics network. Normal workers train the shared policy network and latent dynamics network using true rewards from the sampled MDPs. A mixture We use the terms context, latent belief, and latent (dynamics) model interchangeably denote
3

---

worker trains the policy network with imaginary rewards from the learned decoder's output given mixture beliefs.
3.1 Policy Network
Any type of recurrent network that can encode the past trajectory into a belief state b, is sufficient for the policy network. We use an RL 6 type of policy network (Figure 2. Each normal worker W trains a recurrent Encoder-p (parameterized by @p) and a feedforward policy network (parameterized by 4 to maximize the return for its assigned MDP M (1), A mixture worker trains the same policy network to maximize the return in an imaginary task M where the imaginary reward fr is from the decoder given the mixture model ri, Any online RL algorithm can be used to train the policy network. We use A2C for the gridworld and PPO 37 for MuJoCo tasks to optimize op and t end-to-end.
3.2 Latent Dynamics Network
We use the same network structure and training methods of VAE introduced in variBAD [56] for latent dynamics network (Figure 2). The only difference is that the policy network and the latent dynamics network do not share an encoder. Therefore Encoder-v (parameterized by Φv) of the latent dynamics network does not need to output the context necessary for the poliey but only needs to encode the MDP dynamics into a low-dimensional stochastic latent embedding m. The latent dynamics model m changes over time the agent explores MDP (denoted as m,,) but converges as the agent collects sufficient information to infer the dynamics of the current MDP. The latent dynamics network is not involved in the action selection of workers. We store trajectories from Muain in buffer and use samples from the buffer to train the latent dynamics network offline. Each normal worker trains the latent dynamics network decode the entire trajectory. including the future, to allow inference about unseen future transitions. In this work, we focus on MDPs where only reward dynamics varies 9
17 26 and only train the reward decoder (parameterized by ²n) as in [56]. The parameters or and On are optimized end-to-end to maximize the ELBO (20 using a reparameterization trick.
10 12
3.3 Imaginary Task Generation from Latent Dynamies Mixture
While training the policy network of the normal workers W ,..., (in) in parallel, we generate an imaginary latent model mn as a randomly weighted sum of the latent models m, ,...,m, m.. normal workers.
(x) of the
mi
m
and (1) (",,..., (m) ~ ~3. Dirichlet(1.......
(1)
1=1
(0's are random mixture weights multiplied each latent model m, At the beginning of every iteration when the normal workers are assigned to new MDPs, we also sample new mixture weights fixed for that iteration. There are many distributions suitable for sampling mixture weights, and we use the Dirichlet distribution in Equation I 3 is hyperparameter that controls the mixture's degree of extrapolation. The sum of mixture weights equals regardless of the hyperparameter 3. If 1, all (''s are bounded between 0 and 1. Then the mixture model becomes convex combination of the training models. If B > 1, the resulting mixture model may express extrapolated dynamics of training tasks. We find f = 1.0 suits best for most of our experiments among (0.5, 1.0, 1.5, 2.0, 2.5) that we tried. Refer to extrapolation results in Section 5.1.2 where / greater than I can be effective.
mixture worker interacts with an MDP sampled from Minain but we replace the environment reward r, with the imaginary reward ft to construct a mixture task M. mixture worker trains the policy network to maximize the return for the imaginary task M.. We expect the imaginary task MM to share some common structures with the training tasks because the mixture task is generated using the decoder that is trained to fit the training tasks' reward dynamies. On the other hand, the decoder can generate unseen rewards because we feed the decoder unseen mixture beliefs. The mixture worker only trains the policy network but not the decoder with imaginary dynamics of M.

---

Dropout of state and action input for the decoder As the latent dynamics network is trained for the tasks in Mrain we find that the decoder easily overfits the state and action observations, ignoring the latent model m. Returning to the gridworld example, if we train the decoder with the tasks in Minin (Figure Ta), and feed the decoder one of the goal states in Miest, the decoder refers to the next state input 1,41 only and always returns zero rewards regardless of the latent model m (Figure Fa). We apply dropout of rate Párop to all inputs of the decoder except the latent model m. It forces the decoder to refer to the latent model when predicting the reward and generate general mixture tasks. Refer to Appendix
D for ablations on dropout.
Training the decoder with a single-step regression loss is generally less complex than training the policy network with multi-step policy gradient loss. Therefore the decoder can be stably trained even with input dropout, generalizing better than the meta-trained policy. Refer to Appendix for empirical results on the test-time generalization ability of the latent dynamics network with dropout.
3.4 Implementation Details
Multiple episodes of the same task in one iteration Following the setting of variBAD [56. we define an iteration sequence of N episodes of the same task and train the agent to act Bayes- optimally within the N rollout episodes (i.e., H+ = N x H steps). After every N episodes, new tasks are sampled from Murain for the normal workers, and new mixture weights are sampled for the mixture worker. Then we can compare our method to other meta-RL methods that are designed to maximize the return after rollouts of many episodes.
Multiple mixture workers We may train more than one mixture worker at the same time by sampling different sets of mixture weights for different mixture workers. Increasing the ratio of mixture workers to normal workers may help the agent generalize to unseen tasks faster, but the normal workers may require more iterations to learn optimal policies in Mininn- We train 14 normal workers and ñ = 2 mixture workers in parallel unless otherwise stated. Refer to Appendix D for empirical analysis on the ratio of workers.
4 Related Work
RettaRReinfor Learning Although context-based meta-RL methods require large amount of data for meta-training, they can learn within the task and make online adaptations. RI." [6] is the most simple. yet effective context-based based model-free meta-RL method that utilizes recurrent network to encode the experience into a policy. PEARL (35) integrates an off-policy meta-RL method with online probabilistic filtering of latent task variables to achieve high meta-training efficiency. MAML 9 learns an initialization such that a few steps of policy gradient is enough for the agent to adapt to a new test task. E-MAML and ProMP [36] 40) extend MAML by proposing exploration strategies for collecting rollouts for adaptation.
Bayesian Reinforcement Learning Bayesian RL quantifies the uncertainty or the posterior belief over the MDPs using past experience. Conditioning on the environment's uncertainty, a Bayes- optimal policy can set the optimal balance between exploration and exploitation to maximize the return during training [1 [T 19 28 In a adyestad MDP (BAMDP), where the agent augments the state space of MDP with its belief, it is almost impossible to find the optimal policy due to the unknown parameterization and the intractable belief update. VariBAD (56] proposes an approximate but tractable solution that combines meta-learning and approximate variational inference. However, it restricted to settings where the training and test task distributions are almost the same. Furthermore, the learned latent model is only used as additional information for the policy. LDM uses the learned latent model more actively by creating mixture tasks to train the policy for more general test tasks out of training task distribution.
Curriculum, Goal and Task Generation The idea of generating new tasks for RL not new. Florensa et al. [10] proposes generative adversarial training to generate goals. Gupta et al. [12] proposes an automatic task design process based on mutual information. SimOpt 3 learns the randomization of simulation parameters based on few real-world rollouts. POET and enhanced
POET 45 46 generate the terrain for 2D walker given access to the environment parameters.
Dream to Control
[4] solves long-horizon task using latent-imagination. BIRD [54] learns from
5

---

imaginary trajectories by maximizing the mutual information between imaginary and real trajectories. Chandak et al. 2 trains an agent to forecast the future tasks non- statiiony MDPs. Most of these works require prior knowledge control of the environment parameters, or the pool of the generated tasks is restricted to the training task distribution.
aata-nngmennttion for Reinforcement Learning Many image augmentation techniques such as random convolution, random shift, 12- 12-regularization, dropout, batch normalization and noise injection are shown to improve the generalization of RL (S) 16 22 24
33 50) Mixreg [4 that applies the
idea of mix-up [53] in RL, generates new training data as a convex interpolation of input observation and output reward. LDM can be thought of as data-augmentation method in way that it generates a mixture task using the data from training tasks. However, LDM generates a new task in the latent space with the latent model that fully encodes the MDPs' dynamics. Instead of the pre-defined augmentation techniques with heuristics, we generate mixture tasks using the learned decoder, which contains the shared structure of the MDPs.
Out-of-diste distrbbiinttnnnn Meta-Reinforcement Learning Some recent works aim to generalize RL to OOD tasks [8]. MIER (30 relabels past trajectories in buffer during test to generate synthetic training data suited for the test MDP. FLAP [31] Iearns a shared linear representation of the policy. To adapt to new task FLAP only needs to predict a set of linear weights for fast adaptation. MetaGenRL 21 meta-learns an RL objective to train a randomly initialized policy on a test task. Most of these methods require experience from the test for additional training updates for the network. AdMRL [26] performs adversarial virtual training with varying rewards. AdMRL assumes known reward space and parameterization, while LDM meta-learns.
5 Experiments
We evaluate LDM and other meta-RL methods on the gridworld example (Figure D and three MuJoCo meta-RL tasks 42]. We slightly modify the standard MuJoCo tasks by splitting the task space M into disjoint Muain and Mees. We use RL 2 [6] and variBAD [56] as baselines representing the context-based based meta-RL methods. LDM without mixture training and the latent dynamics network reduces to RL LDM without mixture training reduces to variBAD if the policy and latent dynamics networks share an encoder. We use E-MAML 40 and ProMP (36 as baselines representing the gradient-based meta-RL. methods. We implement -baed Mixreg [44 to evaluate the difference between generating mixture tasks in the latent space and the observation space. Refer to Appendix for the details of implementations and hyperparameters.
All methods are trained on tasks sampled from Munin uniformly random, except for the oracle methods that are trained on tasks in the entire task set M = Mirain U Miest- Note that the oracle performance is only for reference and can not be achieved by any non-oracle methods theoretically. Non-oracle methods without prior knowledge on Mest require additional exploration during testing to experience the changes from Mysin- whereas the oracle agent can exploit the knowledge of the test distribution. In the gridworld example, the oracle agent can search for goals in Mies before navigating the outermost goal states of Musin. Therefore the main focus should be on the relative improvement of LDM compared to non-oracle methods, bridging the gap between the oracle and the non-oracle methods. We report mean results using 8 random seeds and apply a moving average of window size 5 for all main experiments (4 seeds for ablations). Shaded areas indicate standard deviations for all plots.
5.1 Gridworld Navigation
Experimental setup We use the gridworld navigation task introduced in Figure I The agent is allowed 5 actions: up, down, left, right, and stay. The reward is for reaching or staying at the hidden goal and -0.1 for all other transitions. Each episode lasts for H = 30 steps. All baselines are given N = 4 rollout episodes for fixed task except for ProMP and E-MAML that are given N 20 rollouts. Such choice of N follows from the reference implementations of the baselines. The time horizon has been carefully set so that the agent can not visit all states within the first episode but can visit them within two episodes. If a rollout is over, the agent is reset to the origin. The optimal policy is to search for the hidden goal and stay at the goal or return to the goal as quickly as possible. After N episodes, new task is sampled from Muain uniformly at random. The reward decoder

---

25
25
Mean Return of Maran
20
Mean Retum of MMust
20
reached Mast
25
Number goals
20
15
15
10
10
5
0.4 0.6 Timestops
0.8
0.2
Timesteps
Timesteps
LDM
Mixreg
RL2
RL Oracle
variBAD
variBAD
Oracle
ProMP
E-MAML
Figure 3: Results for the gridworld task evaluated at the N-th episode in terms of the mean returns in Minin and Muesa and the number of tasks Miest in which the agent fails to reach the goal.
(a)
(b)
[cl
(d)
(e)
(f)
0.0
0.5
Parp 0.7
0.7
0.9
0.3
0.2
0.1
Mixture Reward Madituue
00.0
1.0
0.0
Figure 4: Examples of mixture tasks generated by LDM. First row: Mixture weights () multiplied
the latent models (m the end of each training task (Equation 1). When the same training tasks are sampled multiple times, we plot the sum of weights. Second Tow (reward map): decoder output for each next state conditioned on the mixture weights from the first row. A cross mark denotes the state with the maximum mixture reward.
requires only the latent model and the next state as input for this task. We apply dropout with the rate Pdrop 0.7 the next state input the reward decoder.
Results Except for ProMP and E-MAML that use feedforward networks, all the baseline methods achieve the same optimal performance in Miain at the N-th episode (Figure 3). However, our method outperforms the non-oracle baselines in Mist Out of the 27 tasks in Mestt our agent succeeds to visit the goals in 19 tasks on average. Although LDM still does not reach 8 test goals on average, considering that there no prior information on the test distribution and that the optimal policy for Miruin (dashed path in Figure Ta does not visit most of the goal states in Miest, the improvement of LDM over RI_2 and VA significant. RL_2 achieves a high test return initially, but as the policy overfits the training tasks, its test return decays. VariBAD achieves high returns in Mirain but fails to generalize most of the tasks in Miest- Mixreg performs better than RL 2 but does match LDM.
Refer to Figure 8a for example of LDM's test-time behavior for a task in Miea that was introduced in Figure Ta The agent searches for goal in Mirain first. Once it explores all the goal states in Mirain- it starts to search for the goals in Mess. From the second episode, the agent directly heads to the goal based on the context made during the first episode. Note that the initial prior is nonzero even for the goals in Mess due to the dropout applied to the decoder's state input.
5.1.1 Tasks generated by LDM
We present an empirical analysis in Figure4[to verify that LDM indeed generates meaningful new tasks that help to solve the test tasks. Without prior knowledge or interaction with Missts it is impossible to create exactly the same reward maps of Mies- But the reward maps in Figure 4b. are sufficient to induce exploration toward the goal states of Miea- Due to the dropout applied to the next state input of the decoder, the decoder assigns high rewards some goals that belong to
7

---

Must Note that the decoder does not generalize without dropout therefore assigns high rewards only training goals (Figure 41).
5.1.2 Extrapolation ability of LDM
To demonstrate that LDM is not restricted to target tasks inside the convex hull of training tasks, we design gridworld-extrapolattio task as in Figure This task is similar to the previous gridworld task in Section [.1] except that we shift the outermost cells of Marain inside to construct extrapolation tasks M (Figure 5c). Refer to Figure 6 for the final return for each task when we train LDM (Párop 0.5) with different values of 3. For small values of 3 1.0 and 3 1.5, LDM focuses on training the interpolation tasks in Miest As increases, returns for tasks in Misat decrease, but the returns for tasks in M1ss2 increase.
(a) Masin
(b) Mieu
(c) MMs2
Figure 5: Gridwoord-ecttraooattioo task.
30
RL2
B 1.0
β 1.5
2.0
B 2.5
20
Mean Retum
10
1
3
2
Figure 6: Returns of each task for different extrapolation level on the gridworld-extrapolltion task. Mean return for each task M at the N-th episode, mean of 4 random seeds.
5.2 MuJoCo
Experimental setup We evaluate our agent and baselines on three standard MuJoCo domains (Ant-direction, Ant-goal, and lllffchhetah--velocit) to verify how effective LDM is for continuous control tasks. We evaluate the agents on fixed set of evaluation tasks Meval C MMeas for each domain ensure the result not affected by the sampling of evaluation tasks (Table I For MuJoCo tasks we report the results of PEARL [35]. MQL [8). and MIER (30) with number of rollouts per iteration (N) equal 3, 11 and 3 respectively. ProMP and E-MAML are given N 20 and all other methods are given N 2 rollouts (following the reference implementations). H = 200 steps for each episode of all MuJoCo tasks. Because PEARL, MQL, and MIER use an off-policy RL algorithm, their performance converges with much less training data. Therefore, for the three baselines, we report the converged asymptotic performance after million steps of environment interaction. We set LDM's Pärop 0.5 for all MuloCo tasks.
Ant-direction We construct the Ant-direction task based on the example discussed the introduc- tion. 0 denotes the angle of the target direction from the origin. A major component of the reward is the dot product of the target direction and the agent's velocity.
Ant-goal We construct the Ant-goal task similar to the gridworld task, where the agent has to reach the unknown goal position. The training task distribution Mirinn (shaded region of Figure &&c is continuous, unlike the gridworld and Ant-direction tasks. and denote the radius and angle of the

---

Ant-direction
Ant-goal
Half-cheetah-velooe
1250
-100
Mean return in M Miaa
1000
-200
-100
750
330
-200
4400
250
-500
-400
0
-600
-500
0.5
1.0
1.5
2.0
0.5
1.0 Timesteps
1.5
2,0
1
2
5
Timesteps
1ce
108
Timesteps
167
LDM
Mixreg
E-MAML
MQL
Oracle
Oracle
Figure 7: Mean returns the N-th episodes in Meval of three MuJoCo tasks.
Table 1: Set of training, test and evaluation tasks of Mujoco tasks. & E 1,,1,
Ant-direction
Ant-goal
Half-cheeta--velocity
0
0
Minann
90" x k
[0.0, 1.0 3.05
10 360%)
[0.0, 0.5) U 3.50
Must
90
k + 45
[1.0, 2.5)
060)
[0.5, 3.0)
Meval
+ 45
1.75
90 k
{0.75, 1.25, 1.75, 2.25, 2.75)
goal from the origin, respectively. A major component of the reward is the negative value of the taxicab distance between the agent and the goal. We set 3 because LDM with goals near the origin mostly due to the symmetry of training tasks. Refer to Appendix C for a detailed analysis regarding the choice of B. One may argue that the training and test tasks are disjoint because the agent passes through the goal states of Miea momentarily while learning to solve M However, the task inference with the unseen rewards from Miest cannot be learned during training. The policy to stay at the goals of Miest until the end of the time horizon is also not learned.
generates
Halff-chettah-velocity We train the Half-cheetah agent to match the target velocity sampled from the two distributions at extremes and test for target velocities in between. This task is relatively easier than the previous Ant tasks due to the reduced dimension of the task space. Therefore we train more mixture workers than the other tasks, training in parallel. Refer to Appendix [oror for additional results with different numbers of mixture workers. The target velocity is where the velocity is measured environment steps). A major component of the reward is the negative value the difference between the target velocity and the agent's velocity. Similar Ant-Goal, although the test velocities in Miss are achieved momentarily in the process of reaching a target velocity in [3.0, 3.5), the agent does not learn to maintain the target velocities in Misa during training Murin-
12 normal workers and ń 4 mixture workers
change of position per second (or 20
Results Refer to Figure 7 for the test returns in Meval- LDM outperforms the non-oracle baselines the N-th episode. VariBAD oracle achieves the best test returns for all MuloCo tests, revealing the strength of variBAD's task inference ability for tasks seen during training. However, the non-oracle variBAD has difficulty generalizing Mistt- The performance of Mixreg is lower than that of RL which reveals the limitation of mixing in complex observation space. In all MuJoCo tasks, the agent can infer the true task dynamics using the reward every timestep, even the first step of an episode. Therefore maddient-based methods with feedforward networks also make progress, unlike the gridworld task. Because PEARL is designed collect data for the first two episodes, achieves low returns before the N-th episode. Even after accumulating context, the policy of PEARL not prepared for the unseen contexts Ant-direction and Half-cheetah--vllote and MIER require additional buffer-bass based re-labeling or re-training for testing after collecting some rollouts of the test task. LDM can prepare in advance during training so that we do not require collection of test rollouts and extra buffer-based training during testing.
Note that MQL

---

LDM
50 100 150 200 variBAD Oracle
1-0
Ce 50
60
90
120
Ep500o3
5p63624
x:position
x-positon
x-position
(a) Sample trajectory of LDM for 4 rollout episodes. The red shade denotes decoder output for each state conditioned on online context.
(b) Ant-drection. 4 trajectories from Meval-
LDM
RL2
variBAD Oracle
LDM
RL2
variBAD Oracle
y-positen
Postbon
20
3
32- 2
3-2
0
50 100 150 200 Environment
50 100 150 200
x-position
Environment
Environimntt
(c) Ant-goal, 4 trajectories from Meval.
(d) half-chhetah--veloiity trajectories from Mout-
Figure 8: Example trajectories of the agents in Meval- We illustrate the behavior at the N-th episode as colored paths. The targets Meval are indicated as dashed lines or cross marks
Because LDM is not trained in Mist. it can not solve the target tasks optimally from the beginning of each task the Ant-direction and Ant-goal tasks, unlike variBAD Oracle (Figure 3b and Sc). However, LDM reaches the target direction goal close enough after sufficiently small amount of exploration, unlike RL For the Half-chhetabh-ellecii task, LDM generates well-separated policies compared RL² (Figure 8d). Refer Appendix [ for additional experimental results the test returns at the first rollout episode, the training returns, and sample trajectories on training tasks.
Conclusion
We propose Latent Dynamics Mixture to improve the generalization ability of meta-RL by training policy with generated mixture tasks. Our method outperforms bascline meta-RL methods in experi- ments with strictly disjoint training and test task distributions, even reaching the oracle performance some tasks. Because our latent dynamics network and the task generation process are independent of the policy network, we expect LDM to make orthogonal contributions when combined with not only RL but most of other meta-RL. methods.
We believe that our work can be a starting point for many interesting future works. For example, in- stead of heuristic weight-sampling for generating the mixture, may incorporate OOD generation techniques in the latent space. Another extension is to train the latent dynamics network to decode only the reward but also the state transition dynamics. Then we can generate a purely imaginary mixture task without additional environment interaction.
Acknowledgments
This work was supported by the National Research Foundation of Korea (NRF) grant funded by the Korca government (MSIT) 1022RI1A220005111 We thank Dong Hoon Lee, Minguk Jang, and Dong Geun Shin for constructive feedback and discussions.
References
[1] A. R. Cassandra, I.. Kaelbling, and M. Littman. Acting optimally in partially observable stochastic domains. In National Conference Artificial Intelligence, 1994.
10

---

[2] Y. Chandak, G. Theocharous, S. Shankar, M. White, S. Mahadevan, and P. Thomas. Optimizing for the future in non-stationary mdps. In International Conference on Machine Learning, (ICML), pages 1 11:14414-1425, 2020.
[3] Y. Chebotar, A. Handa, V. Makoviychuk, M. Macklin, J. Issac, N. D. Ratliff, and D. Fox. Closing the sim-to-real loop: Adapting simulation randomization with real world experience. In International Conference on Robotics and Automation, (ICRA), pages 8973-8979, 2019.
[4] B. Chen, W. Deng, and H. Shen. Virtual class enhanced discriminative embedding learning. In Advances Neural Information Processing Systems (NeurlPS), pages 31:1942-14 1952, 2018.
[5] K. Cobbe, O. Klimov, C. Hesse, T. Kim, and J. Schulman. Quantifying generalization in reinforcement learning. In International Conference on Machine Learning, (ICML., pages 97:1282-1289, 2019.
[6] Y. Duan, Schulman, X. Chen, P. L. Bartlett, I. Sutskever, and P. Abbeel. RL:3: Fast re- inforcement learning via slow reinforcement learning. arXiv preprint arXiv: 161 1611.02779, 2016.
[7] M. O. Duff. Optimal Learning: Computational Procedures for Bayes-adaptive Markov Decision Processes. University of Massachusetts at Amherst, 2002.
[8] R. Fakoor, P. Chaudhari, S. Soatto, and A. J. Smola. Meta-q-learning, In International Conference on Learning Representations, (ICLR), 2020.
[9] C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning, (ICML), pages 70:1 1t:1126-1113 1135, 2017.
[10] C. Florensa, D. Held, X. Geng, and P. Abbeel. Automatic goal generation for reinforcement learning agents. In International Conference on Machine Learning (ICML), pages 80: 1515- 1528, 2018
[11] M. Ghavamzadeh, Mannor, J. Pineau, and A. Tamar. Bayesian reinforcement learning: A survey, Found. Trends Mach. Learn., pages 855-6)3359-48,, 2015,
[12] A. Gupta, B. Eysenbach, C. Finn, and S. Levine. Unsupervised meta-learning for reinforcement learning. arXiv preprint arXiv: 2806.04640 2018
[13] A. Gupta, R. Mendonca, Liu, P. Abbeel, and S. Levine. Meta-reinforcement learning of structured exploration strategies. In Advances in Neural Information Processing Systems (NeurIPS), pages 31:5302-5311, 2018.
[14] D. Hafner, P. Lillicrap. J. Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, (ICLR), 2020.
[15] J. Humplik, A. Galashov, L. Hasenclever, P. A. Ortega, Y. W. Teh, and N. Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv: 1905.06424, 2019.
[16] M. Igl, K. Ciosek, Li, S. Tschiatschek, C. Zhang. Devlin, and K. Hofmann. Generalization in reinforcement learning with selective noise injection and information bottleneck. In Neural Information Processing Systems (NeurIPS), pages 02:11397-1999 2019.
[17) A. Jabri, K. Hsu, B. Eysenbach, A. Gupta, S. Levine, and C. Finn. Unsupervised curricula for visual meta-reiinforeement learning. In Advances in Neural Information Processing Systems (NeurIPS), volume 32, 2019.
[18] Kaddour, S. Sæeuudssoon, and M. P. Deisenroth. Probabilistic active meta-learning. In Advances in Neural Information Processing Systems (NeurIPS), pages 33:20813 3-20822, 2020.
[19] L. P. Kaelbling, M. L. Littman, and A. R. Cassandra. Planning and acting in partially observable stochastic domains. Artificial Intelligence, 101:99-134, 1998.
[20] D. Kingma and M. Welling. Auto-encoding variational bayes. International Conference on Learning Representations (ICLR), 2014.
11

---

[21] L. Kirsch, S. van Steenkiste, and J. Schmidhuber. Improving generalization in meta reinforce- ment learning using learned objectives. In International Conference on Learning Representa- lions, (ICIR), 2020.
[22] M. Laskin, K. Lee, A. Stooke, L. Pinto, P. Abbeel, and A. Srinivas. Reinforcement learning with augmented data. In Neural Information Processing Systems (NeurIPS), pages 33:19884- 19895, 2020.
[23] A. X. Lee, A. Nagabandi, P. Abbeel, and S. Levine. Stochastic latent actor-critie: Deep reinforcement learning with a latent variable model. In Advances in Neural Information Processing Systems (NeurIPS), pages 33:741-752, 2020.
[24] K. Lee, K. Lee, J. Shin, and H. Lee. Network randomization: A simple technique for generaliza- tion deep reinforcement learning. In International Conference on Learning Representations, (ICLR), 2020.
[25] K. Lee, Y. Seo, S. Lee, H. Lee, and J. Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In International Conference on Machine Learning. (ICML), pages 11:57777-5766, 2020.
[26] Z. Lin, G. Thomas, G. Yang, and T. Ma. Model-based adversarial meta-reinforceement learnin In Advances in Neural Information Processing Systems (NeurIPS), pages 33:10161-10173, 2020.
[27] S. Liu, A. J. Davison, and E. Johns. Self-supervised generalisation with meta auxiliary learning. In Neural Information Processing Systems (NeurIPS), pages 32:1679-1689, 2019.
[28] J. Martin. Bayesian Decision Problems and Markov Chains. Wiley, 1967.
[29] B. Mehta, T. Deleu, C. Raparthy, C. J. Pal, and L. Paull. Curriculum in gradient-based meta-reinforcement learning. arXiv preprint 2rXiv:2002.07956, 2020.
[30] R. Mendonca, X. Geng, C. Finn, and S. Levine. Meta-reinforcemeen learning robust to distribu- tional shift via model identification and experience relabeling, arXiv preprint arXiv: 2006.07178. 2020.
[31] M. Peng, B. Zhu, and J. Jiao. Linear representation meta-reinforcemmnt learning for instant adaptation. arXiv preprint arXiv: 2101.04750, 2021.
[32] R. Raileanu, M. Goldstein, A. Szlam, and R. Fergus. Fast adaptation to new environments via policy-dynamies value functions. In International Conference on Machine Learning, (ICML), pages 119:7920-7931. 2020.
[33] R. Raileanu, M. Goldstein, D. Yarats, I. Kostrikov, and R. Fergus. Automatic data augmentation for generalization in deep reinforcement learning, arXiv preprint arXiv:2006. 11...22866 /2862, 2020.
[34] A. Rajeswaran, K. Lowrey, E. Todorov, and S. M. Kakade. Towards generalization and simplicity continuous control. Advances in Neural Information Processing Systems (NeurIPS), pages 30:6550-6561, 2017.
[35] K. Rakelly, A. Zhou, C. Finn, S. Levine, and D. Quillen. Efficient off- off-policy meta-reinforcemmen learning via probabilistic context variables. In International Conference on Machine Learning, (ICML), pages 97:5331 2019.
[36] J. Rothfuss, D. Lee, I. Clavera, T. Asfour, and P. Abbeel. Promp: Proximal meta-policy search. In International Conference on Learning Representations, (ICLR), 2019.
[37] Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprini arXiv: 1707.06347, 2017.
[38] X. Song, Y. Jiang, S. Tu, Y. Du, and B. Neyshabur. Observational overfitting in reinforcement learning. In International Conference on Learning Representations, (ICLR), 2020.
[39] Y. Song, A. Mavalankar, W. Sun, and S. Gao. Provably efficient model-based policy adaptation.
In International Conference on Machine Learning (ICML), pages 119:9088-900 9999988888888,,,8888888899888888 2020.
12

---

[40] B. C. Stadie, G. Yang, R. Houthooft, X. Chen, Y. Duan, Y. Wu, P. Abbeel, and I. Sutskever. The importance of sampling in metu-reinforreemen learning. In Advances in Neural Information Processing Systems (NeurIPS), pages (19280099290 2018.
[41] J. Tobin, R. Fong. A. Ray, J. Schneider, W. Zaremba, and P. Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In International Conference on Intelligent Robots and Systems (IROS), pages 23-30, 2017.
[42] E. Todorov, T. Erez, and Y. Tassa. Mujoco: A physics engine for model-based control. In International Conference on Intelligent Robots and Systems (IROS), pages 5026-5033, 2012.
[43] J. X. Wang, Z. Kurth-Nelson, D. Tirumala, H. Soyer, J. Z. Leibo, R. Munos, C. Blundell, D. Kumaran, and M. Botvinick. Learning to reinforcement learn. In Annual Meeting of the Cognitive Science Community (CogSci), 2016.
[44] K. Wang, B. Kang, J. Shao, and J. Feng. Improving generalization in reinforcement learning with mixture regularization. In Advances Neural Information Processing Systems (NeurIPS), pages 33:7968-7978, 2020.
[45] R. Wang. J. Lehman, J. Clune, and K. O. Stanley. Paired open-ended trailblazer (POET): endlessly generating increasingly complex and diverse learning environments and their solutions. arXiv preprint arXiv: 1901.01753, 2019.
[46] R. Wang, J. Lehman, A. Rawal, J. Zhi, Y. J. Clune, and K. Stanley. Enhanced poet: Open- ended reinforcement learning through unbounded invention of learning challenges and their solutions. In International Conference on Machine Learning. (ICML), pages 11::9940-9551, 2020.
[47] S. Whiteson, B. Tanner, M. E. Taylor; and P. Stone. Protecting against evaluation overfitting in empirical reinforcement learning. In IFEE Symposium on Adaptive Dynamic Programming and Reinforcement Learning (ADPRL.), pages 120-127, 2011.
[48] A. Xie, J. Harrison, and C. Finn. Deep reinforcement learning amidst lifelong non sottsttitooaitt. arXiv preprint arXiv: 2006. 10701, 2020.
[49] J. Yang, B. K. Petersen, H. Zha, and D. Faissol. Single episode policy transfer in reinforcement learning. In International Conference on Learning Representations, (ICLR), 2020.
[50] D. Yarats, I. Kostrikov, and Fergus. Image augmentation all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations, (ICLR), 2021.
[51] M. Yin, G. Tucker, M. Zhou, Levine, and C. Finn. Meta-learning without memorization. In International Conference on Learning Representations, (ICLR), 2020.
[52] C. Zhang, O. Vinyals, R. Munos, and S. Bengio. A study overfitting in deep reinforcement learning. arXiv preprint arXiv: 1804.06893, 2018
[53] H. Zhang, M. Cissé, Y. N. Dauphin, and D. Lopez-Paz. mixup: Beyond empirical risk mini- mization. In International Conference on Learning Representations, (ICLR), 2018.
[54] G. Zhu, M. Zhang, H. Lee, and Zhang, Bridging imagination and reality for model-based deep reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), pages 33:8993- 90069 2020.
[55] L. M. Zintgraf, Shiarlis, Kurin, K. Hofmann, and Whiteson. Fast context adaptation via meta-learning. In International Conference on Machine Learning, (ICML), pages 97:7693-7702. 2019.
[56] L. M. Zintgraf, K. Shiarlis, M. Igl, S. Schulze, Y. Gal, K. Hofmann, and Whiteson. Varibad: A very good method for bayes-adaptive deep RL via meta-learning. In International Conference on Learning Representations, (ICLR), 2020.
13

---

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] In the conclusion section
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(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? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (cither in the supplementary material or as a URL)? [Yes] as a URL in Appendix
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] in Appendix
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [Yes] in Appendix
4. 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) source codes of the baselines
(b) Did you mention the license of the assets? [Yes] the supplementary material
(c) Did you include any new assets either in the supplementary material or as a URL?
Yes] Links the baseline codes in Appendix
(d) Did you discuss whether and how consent obtained from people whose data you'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]
14


Reviews:
[{"review_id": "izkNOqaM9Dm", "reviewer": "Reviewer_BvaU", "summary": "This paper focuses on the field of meta-reinforcement learning (meta-RL), where the authors propose a method to augment the tasks for meta-reinforcement by applying a mixture of learned reward functions. Specifically, the paper builds on the framework of variBad [1], which uses a VAE to learn a latent dynamics and reward model by encoding the data and decoding the next states and rewards. The authors propose a method that uses a mixture of the encoded latent variables from different tasks during the training time to obtain new latent encodings corresponding to imaginary tasks, and then applies these latent encodings to the reward decoder to obtain new reward functions to train the policy on these imaginary tasks.\n\nThe authors then evaluate the proposed method on grid world tasks and several MuJoCo locomotion tasks, where the authors specifically emphasize the performance of adapting to out-of-distribution tasks. The empirical results demonstrate that the proposed method outperforms several baselines.\n\n\nReferences\n\n[1] Zintgraf, L., et al. \"VariBAD: a very good method for Bayes-adaptive deep RL via meta-learning.\" Proceedings of ICLR 2020 (2020).\n", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall I think this paper presents an interesting idea on task augmentation for metal-reinforcement learning. The paper is well-written and the authors include a lot of experiment details and ablation studies in the paper. However, there are also many limitations in this paper that need to be addressed.\n\nPros:\n\nThe main idea is simple and intuitive. Using a mixture of latent codes to augment the task is a natural idea and it is easy to believe that it should work well.\n\nThe paper is well written. The main idea is presented in a way that is easy to understand, and the whole paper is structured in an organized way.\n\nThe authors provide detailed analysis of the experiments, with many informative behavior visualizations and ablation studies about the proposed method.\n\nCons:\n\nSeveral important baseline comparisons are missing. The main advantage of the proposed method is the capability of adapting well to out-of-distribution tasks, where the training tasks and test tasks come from non-overlapping distributions. As the authors mentioned in related works, several prior works such as MQL [1], MIER [2] and FLAP [3] have attempted to address the out-of-distribution problem of test time tasks. Since this paper addresses the same problem, it would be important to compare the proposed method to these baselines.\n\nThe scope of experiments in this paper is rather limited. Besides simple grid world environments, the authors only include 3 MuJoCo locomotion environments, which I don’t believe is enough. In order to add more environments, the authors could consider out-of-distribution environments proposed in [1, 2], and also MetaWorld [4]. \n\nIn summary, being an empirical paper, this paper does not include enough experiments to convince me about the performance of the proposed method. Therefore, it is hard to recommend acceptance of this paper at the current state.\n\n\nReferences\n\n[1] Fakoor, Rasool, et al. \"Meta-q-learning.\" arXiv preprint arXiv:1910.00125 (2019).\n\n[2] Mendonca, Russell, et al. \"Meta-reinforcement learning robust to distributional shift via model identification and experience relabeling.\" arXiv preprint arXiv:2006.07178 (2020).\n\n[3] Peng, Matt, Banghua Zhu, and Jiantao Jiao. \"Linear representation meta-reinforcement learning for instant adaptation.\" arXiv preprint arXiv:2101.04750 (2021).\n\n[4] Yu, Tianhe, et al. \"Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning.\" Conference on Robot Learning. PMLR, 2020.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See main review for limitations and N/A for social impact.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "haNIEaCSXfD", "reviewer": "Reviewer_Ncyq", "summary": "This paper tackles out-of-distribution generalization in meta-reinforcement learning. The proposed method, latent dynamics mixture (LDM), extends meta-RL methods (e.g. variBAD) that support inference and generation of a reward function by interpolating between inferred latents of training tasks to generate new reward functions. The intuition is that \"filling in\" areas of the task space not supported by the training task distribution corresponds to shifting towards the test task distribution. ", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Originality\nThe technical contribution is very limited: the delta over variBAD consists entirely of obtaining a convex combination of latent vectors for decoding.\n\n## Quality\nThe experiments are reasonable, but there is a lack of comparison to methods that also attempt to tackle out-of-distribution generalization, aside from MixReg-based RL2. All other comparisons are made to meta-RL methods that assume iid training and test distributions. Besides quantitative results, informative qualitative analysis that investigates the impact of the proposed method are presented. \n\n## Clarity\nThe paper is well-written. \n\n## Significance\nWhile this is a cleanly executed paper, I am not sure it will have a significant impact. The technical contribution is undisputably incremental, and the resulting improvements marginal. Given that LDM implicitly assumes that test tasks are \"interpolations\" of training tasks (this is reflected in the design of the experiments) and that the latent space is regularized by the variational inference training procedure, it's also entirely unsurprising that it works. \n\n## Summary\nIn essence, my main concern is that the contributions are not sufficient to warrant a NeurIPS publication. However, I look forward to discussing this point with everyone.\n\n## Post-rebuttal update\nThe authors have clarified via ablations that dropout on the decoder input and separate encoders are important for LDM. They have also expanded their experiments by comparing to meta-learning methods that address out-of-distribution generalization, and on more environments. I have improved my score to a 6.\n\nI still maintain similar concerns to Reviewer 1a3x regarding an implicit \"convex hull\"/\"filling in\" assumption underpinning LDM. While the authors have acknowledged this point in their rebuttal, I encourage them to develop this notion further in future revisions. It would make the paper stronger.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "This is adequately done.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "W1Mabdd2nIF", "reviewer": "Reviewer_1a3x", "summary": "This work studies the partial observability problem of training a policy to adapt to new, unseen out-of-distribution tasks at test time. The authors focus on the setting where the dynamics across all tasks are the same but reward changes. They propose to train an ensemble of latent dynamics models, policies, and reward decoders and train an additional mixture version with dropout. During training, they replace the real reward with an imaginary reward generated by the mixture decoder to create new imaginary tasks. They show that in gridworld and MuJoCo environments, their method generalizes better than other Meta-RL methods. \n", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "To really evaluate these results requires more than 4 seeds. The variance seems quite high in Figure 3. Given the importance of dropout to prevent memorization of train tasks, another good baseline to check would be RL^2 and/or VariBAD with dropout as well. Perhaps that is all the regularization needed to get better generalization performance? You mention that LDM test time behavior is to first try all the train goal points before exploring for new test goal points. Given that the other recurrence-based baselines also have > 0 reward at test time, do they do something similar?\n\nOverall it seems like a nice idea but I’m not convinced by the results. More analysis of the types of problems that LDM would work well vs. not well for would also be appreciated. All the examples chosen are of tasks with more “continuous” context spaces. I assume LDM fails in tasks with only discrete structure? What types of test tasks would this fail at? It seems like there is some notion of a convex hull over training environments that this captures. Requiring a strict separation of train and test seems like a limitation? Does it need to be strict? Does the method not work if it is not?\n\nI am currently giving this work a borderline reject which I am happy to raise to an accept contingent on some analysis and discussion of limitations and assumptions being made about the structure of the environment and results with 10 seeds. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "I_Nl3IRik0G", "reviewer": "Reviewer_yAYE", "summary": "The authors attempt to solve a problem in context-based RL, where a model tries to learn and identify the context within which it is behaving. This context can be defined by differences between MDPs, such as differences in the transition dynamics, and the reward function; the agent will make use of such a context in its policy. \n\nThis work builds heavily on the variBAD, as is mentioned in the paper. As its main contribution, it introduces a so-called Latent Dynamics Mixture model where during training, latent context values that are seen over the sampled training tasks are used (additively) to generate new, 'imaginary' tasks, to augment the training. \n\nThis method aims to improve on the problem of generalising to unseen tasks that are generated from an unseen task distribution. \n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Originality \n- This paper builds on variBAD, but I believe has sufficient novelty. I haven't seen similar work otherwise. \n\n\n### Quality\n- I think that experiments that the authors had run, do support their claim that generalisation to unseen tasks from an unseen distribution works better with their models. But some questions do arise. See (Misc. Comments for some questions I have). \n\n- Are 4 seeds really enough for any meaningful statistical significance?\n \n### Clarity\n- I think that this paper is written sufficiently well. I would have liked to have seen the equation of the loss they are using, written in terms of their notation. It would help with my understanding of what is going on, and the interaction between the different components of this. \n\n- That being said, Figure 2 was quite a nice summary of their solution.  \n\n- I appreciate the inclusion of the code, and what appears like sufficient hyperparameter details of how to reproduce this work in the Appendices. \n\n- Line 40: I think that the language here could be clarified a little. When I first read this, I thought that the policy would be able to learn new 'actions' for its policy. Perhaps more clarity about which environment this is in, and more precise descriptions would be useful. \n\n- Line 249: Repeated definition of H+ is repeated as in Line 161. Also, I'm not certain H+ is described or used anywhere substantial. \n\n- I enjoyed Figure 6(b) and (c), as they are nice visualisations the different ways the methods behaved. \n\n### Significance\n- I think that this work is quite interesting. This could be a nice solution for OOD problems in RL as discussed in the conclusion of the paper. \n\n  But, some questions arise from this work (see Misc. Comments below) that I think should be addressed first. \n\n### Misc. comments\n- Generally speaking, I am confused about the use of the word \"Dynamics\" in LDM. Is this referring to the dynamics of the MDP? Or is it referring to some dynamics of the latent context variable? \n\n-  I am also confused as to how exactly this method would help generalise to new tasks. As I understand it, a main novelty is the generation of the imaginary tasks, which are generated as a convex combination of the latent context values from the training tasks. Now, if the claim is that methods that don't use such a generation of imaginary tasks do not generalise well because the latent values are generated only over $M_{train}$, why is there an expectation that a convex sum of such $m$ would result in meaningfully new tasks?\n\n- In fact, from Figure 4, and the description given in the paragraph starting at Line 268, it seems like the dropout plays a key role in generalising to unseen goals. So do the imaginary tasks do what is claimed?\n\n   From the previous submission to ICML 2021, the authors had added Figure 4 to address the following issue:\n   \n   \"We included additional empirical results regarding the tasks generated by LDM in Figure 4. We added a paragraph describing that LDM indeed generates meaningful new tasks that help to solve unseen test tasks.\"\n\n   I'm not certain the new figure answers this, due to the effects of dropout. \n\n- What is the significance of encoders for the policy network and latent dynamics network not being shared? If I understand correctly, the policy uses a latent embedding too yes? Couldn't it use the embedding learnt from Encoder-v? Why are they separated? Could there be a comparison where the encoders are shared, as an ablation?\n\n### Conclusion\nI think that this work is interesting. The clarifications of the questions above would help my confidence in the work. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I think that some limitations are mentioned (for example the fact that dropout is important for the imaginary tasks to actually look at unseen goals). I think some more discussion of why this is would have been good. See Main Review for more details about this. \n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes to learn an ensemble of latent dynamics models and provide a mixture of them as input to a reward decoder, which allows the entire framework to generate imaginary tasks (i.e., imaginary reward functions). By training the agent on the generated tasks, the agent can generalize better to out-of-distribution tasks on grid-world and MuJoCo environments. The idea of generating imaginary reward functions using a mixture of latent dynamics model is novel, though the overall framework builds upon VariBad. Although there was an initial concern around the significance of the results (e.g., lack of baselines, challenging environments), the additional results during the rebuttal period addressed most of the questions raised by the reviewers. However, the question about how the proposed framework can extrapolate so well and how the mixture hyperparameters ($\beta$) affects interpolation/extrapolation performance remains the same, which may need to be discussed in the main paper rather than in the appendix. Given that the main result is still solid and consistent, I'd recommend acceptance for this paper.
Decision: Accept (Poster)