Improving Generalization in Meta-RL with Imaginary Tasks from Latent Dynamics Mixture
Year: 2021

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.

heading: 1 Introduction
text: Overfitting and lack of generalization ability have been raised as the most critical problems of deep reinforcement learning (RL) [5, 8, 30, 34, 38, 47, 52]. Numbers of meta-reinforcement learning (meta-RL) methods have proposed solutions to the problems by meta-training a policy that easily adapts to unseen but similar tasks. Meta-RL trains an agent in multiple sample tasks to construct an inductive bias over the shared structure across tasks. Most meta-RL works evaluate their agents on test tasks that are sampled from the same distribution used to sample training tasks. Therefore, the vulnerability of meta-RL to test-time distribution shift is hardly revealed [12, 26, 29, 30].
One major category of meta-RL is gradient-based meta-RL that learns an initialization of a model such that few steps of policy gradient are sufficient to attain good performance in a new task [9, 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 category of meta-RL is context-based meta-RL that tries to learn the tasks’ structures by utilizing recurrent or memory-augmented models [6, 13, 23, 25, 32, 35, 49, 56]. A context-based meta-RL 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 to unseen out-ofdistribution (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), a novel meta-RL method that overcomes the aforementioned limitations and generalizes to strictly unseen test tasks without any additional test-time updates. LDM is based on variational Bayes-adaptive meta-RL that meta-learns approximate inference on a latent belief distribution over multiple reward and transition dynamics [56]. We generate imaginary
∗Corresponding author.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
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, it 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 of 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 meta-RL methods on the gridworld navigation task and MuJoCo meta-RL tasks, where we completely separate the distributions of 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.

heading: 2 Problem Setup
text: Our work is motivated by the meta-learning setting of variBAD [56], therefore we follow most of the problem setup and notations except for a key difference that the test and training task distributions are strictly disjoint in our setup. A Markov decision process (MDP) M = (S,A, R, T, T0, γ,H) consists of a set of states S, a set of actions A, a reward function R(rt+1|st, at, st+1), a transition function T (st+1|st, at), an initial state distribution T0(s0), a discount factor γ and a time horizon H . During meta-training, a task (or a batch of tasks) is sampled following p(M) over the set of MDPs M at every iteration. Each MDP Mk = (S,A, Rk, Tk, T0, γ,H) has individual reward function Rk (e.g., goal location) and transition function Tk (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 p that is used to sample the training tasks. To evaluate the generalization ability of agents in environments unseen during training, we splitM into two strictly disjoint training and test sets of MDPs, i.e., M = Mtrain ∪Mtest andMtrain ∩Mtest = ∅. The agent does not have any prior information aboutMtest and cannot interact inMtest during training. Since the MDP is initially unknown, the best the agent can do is to update its belief bt(R, T ) about the environment according to its experience τ:t = {s0, a0, r1, s1, a1, r2, . . . , st}. According to the Bayesian RL formulation, the agent’s belief about the reward and transition dynamics at timestep t can be formalized as a posterior over the MDP given the agent’s trajectory, bt(R, T ) = p(R, T |τ:t). By augmenting the belief to the state, a Bayes-Adaptive MDP (BAMDP) can be constructed [7]. The agent’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 a low-
dimensional stochastic latent variable m to represent the posterior belief over the MDPs.2 The reward and transition dynamics can be formulated as shared functions across MDPs: R(rt+1|st, at, st+1;m) and T (st+1|st, at;m). Then the problem of computing the posterior p(R, T |τ:t) becomes inferring the posterior p(m|τ:t) overm. By conditioning the policy on the posterior p(m|τ:t), an approximately Bayes-optimal policy can be obtained.
Refer to Figure 1 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 setM is divided into disjoint Mtrain andMtest. A Bayes-optimal agent for a task inMtrain first assigns a uniform prior to the goal states ofMtrain (Figure 1a) and then explores these states until it discovers a goal state as the dashed path in Figure 1a. If this agent, trained to solve the tasks inMtrain only, is put in to solve a task inMtest without any prior knowledge (Figure 1b), the best policy an agent can take first is to maintain its initial belief learned inMtrain and explore the goal states ofMtrain. Once the agent realizes that there are no goal states inMtrain, it could start exploring the states that are not visited (i.e.,M−Mtrain), and discover an unseen goal state inMtest. However, it is unlikely that the agent trained only inMtrain will encode its experience into beliefs for unseen tasks accurately and explore the goal states ofMtest efficiently conditioned on the unseen context without any prior knowledge or test-time adaptation.

heading: 3 Latent Dynamics Mixture
text: Our work aims to train an agent that prepares for unseen test tasks during training as in Figure 1b. 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 inMtrain but also for more general tasks that may appear during testing.
Refer to Figure 2 for an overview of the entire process. We train n normal workers W (1), . . . ,W (n)
and a 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 n MDPs M (1), . . . ,M (n) fromMtrain and assign each MDPM (i) to each normal workerW (i). All normal and mixture workers share a single policy network and a single latent dynamics network. Normal workers train the shared policy network and latent dynamics network using true rewards from the sampled MDPs. A mixture
2We use the terms context, latent belief, and latent (dynamics) model interchangeably to denote m.
worker trains the policy network with imaginary rewards from the learned decoder’s output given mixture beliefs.

heading: 3.1 Policy Network
text: Any type of recurrent network that can encode the past trajectory into a belief state bt is sufficient for the policy network. We use an RL2 [6] type of policy network (Figure 2). Each normal worker W (i) trains a recurrent Encoder-p (parameterized by φp) and a feedforward policy network (parameterized by ψ) to maximize the return for its assigned MDP M (i). A mixture worker trains the same policy network to maximize the return in an imaginary task M̂ where the imaginary reward r̂t is from the decoder given the mixture model m̂t. 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 φp and ψ end-to-end.

heading: 3.2 Latent Dynamics Network
text: We use the same network structure and training methods of VAE introduced in variBAD [56] for our 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 policy but only needs to encode the MDP dynamics into a low-dimensional stochastic latent embedding m. The latent dynamics model m changes over time as the agent explores an MDP (denoted as mt), 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 fromMtrain in a buffer and use samples from the buffer to train the latent dynamics network offline. Each normal worker trains the latent dynamics network to 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, 10, 12, 13, 17, 26] and only train the reward decoder (parameterized by θR) as in [56]. The parameters φv and θR are optimized end-to-end to maximize the ELBO [20] using a reparameterization trick.

heading: 3.3 Imaginary Task Generation from Latent Dynamics Mixture
text: While training the policy network of the normal workers W (1), . . . ,W (n) in parallel, we generate an imaginary latent model m̂t as a randomly weighted sum of the latent models m (1) t , . . . ,m (n) t of the normal workers.
m̂t = n∑ i=1 α(i)m (i) t and α (1), . . . , α(n) ∼ β · Dirichlet(1,. . . ,1)− β − 1 n . (1)
α(i)’s are random mixture weights multiplied to each latent model m(i)t . 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 1. β is a hyperparameter that controls the mixture’s degree of extrapolation. The sum of mixture weights equals 1 regardless of the hyperparameter β. If β = 1, all α(i)’s are bounded between 0 and 1. Then the mixture model becomes a convex combination of the training models. If β > 1, the resulting mixture model may express extrapolated dynamics of training tasks. We find β = 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 1 can be effective.
A mixture worker interacts with an MDP sampled from Mtrain, but we replace the environment reward rt with the imaginary reward r̂t to construct a mixture task M̂ . A mixture worker trains the policy network to maximize the return for the imaginary task M̂ . We expect the imaginary task M̂ 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 dynamics. 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 inMtrain, 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 Mtrain (Figure 1a), and feed the decoder one of the goal states inMtest, the decoder refers to the next state input si+1 only and always returns zero rewards regardless of the latent model m (Figure 4a). We apply dropout of rate pdrop 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 E for ablations on dropout.
Training the decoder with a single-step regression loss is generally less complex than training the policy network with a 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 B for empirical results on the test-time generalization ability of the latent dynamics network with dropout.

heading: 3.4 Implementation Details
text: Multiple episodes of the same task in one iteration Following the setting of variBAD [56], we define an iteration as a sequence of N episodes of the same task and train the agent to act Bayesoptimally within the N rollout episodes (i.e., H+ = N ×H steps). After every N episodes, new tasks are sampled fromMtrain 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 inMtrain. We train n = 14 normal workers and n̂ = 2 mixture workers in parallel unless otherwise stated. Refer to Appendix D for empirical analysis on the ratio of workers.

heading: 4 Related Work
text: Meta-Reinforcement Learning Although context-based meta-RL methods require a large amount of data for meta-training, they can learn within the task and make online adaptations. RL2 [6] is the most simple, yet effective context-based model-free meta-RL method that utilizes a 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 Bayesoptimal policy can set the optimal balance between exploration and exploitation to maximize the return during training [1, 11, 19, 28]. In a Bayes-adaptive 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 is 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 is 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 a few real-world rollouts. POET and enhanced POET [45, 46] generate the terrain for a 2D walker given access to the environment parameters. Dream to Control [14] solves a long-horizon task using a latent-imagination. BIRD [54] learns from
imaginary trajectories by maximizing the mutual information between imaginary and real trajectories. Chandak et al. [2] trains an agent to forecast the future tasks in non-stationary MDPs. Most of these works require prior knowledge or control of the environment parameters, or the pool of the generated tasks is restricted to the training task distribution.
Data-augmentation for Reinforcement Learning Many image augmentation techniques such as random convolution, random shift, l2-regularization, dropout, batch normalization and noise injection are shown to improve the generalization of RL [5, 16, 22, 24, 33, 50]. Mixreg [44] 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 a data-augmentation method in a 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-distribution Meta-Reinforcement Learning Some recent works aim to generalize RL to OOD tasks [8]. MIER [30] relabels past trajectories in a buffer during a test to generate synthetic training data suited for the test MDP. FLAP [31] learns a shared linear representation of the policy. To adapt to a 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 or updates for the network. AdMRL [26] performs adversarial virtual training with varying rewards. AdMRL assumes known reward space and parameterization, while LDM meta-learns.

heading: 5 Experiments
text: We evaluate LDM and other meta-RL methods on the gridworld example (Figure 1) and three MuJoCo meta-RL tasks [42]. We slightly modify the standard MuJoCo tasks by splitting the task spaceM into disjointMtrain andMtest. We use RL2 [6] and variBAD [56] as baselines representing the context-based meta-RL methods. LDM without mixture training and the latent dynamics network reduces to RL2. 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 RL2-based Mixreg [44] to evaluate the difference between generating mixture tasks in the latent space and the observation space. Refer to Appendix A for the details of implementations and hyperparameters.
All methods are trained on tasks sampled fromMtrain uniformly at random, except for the oracle methods that are trained on tasks in the entire task setM =Mtrain ∪Mtest. 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 onMtest require additional exploration during testing to experience the changes fromMtrain, whereas the oracle agent can exploit the knowledge of the test distribution. In the gridworld example, the oracle agent can search for goals inMtest before navigating the outermost goal states ofMtrain. 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.

heading: 5.1 Gridworld Navigation
text: Experimental setup We use the gridworld navigation task introduced in Figure 1. The agent is allowed 5 actions: up, down, left, right, and stay. The reward is 1 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 a 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, a new task is sampled fromMtrain uniformly at random. The reward decoder
requires only the latent model and the next state as input for this task. We apply dropout with the rate pdrop = 0.7 to the next state input of the reward decoder.
Results Except for ProMP and E-MAML that use feedforward networks, all the baseline methods achieve the same optimal performance inMtrain at the N -th episode (Figure 3). However, our method outperforms the non-oracle baselines inMtest. Out of the 27 tasks inMtest, 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 is no prior information on the test distribution and that the optimal policy for Mtrain (dashed path in Figure 1a) does not visit most of the goal states inMtest, the improvement of LDM over RL2 and variBAD is significant. RL2 achieves a high test return initially, but as the policy overfits the training tasks, its test return decays. VariBAD achieves high returns inMtrain, but fails to generalize in most of the tasks inMtest. Mixreg performs better than RL2, but does not match LDM. Refer to Figure 8a for example of LDM’s test-time behavior for a task inMtest that was introduced in Figure 1a. The agent searches for a goal inMtrain first. Once it explores all the goal states inMtrain, it starts to search for the goals inMtest. 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 inMtest due to the dropout applied to the decoder’s state input.

heading: 5.1.1 Tasks generated by LDM
text: We present an empirical analysis in Figure 4 to verify that LDM indeed generates meaningful new tasks that help to solve the test tasks. Without prior knowledge or interaction with Mtest, it is impossible to create exactly the same reward maps ofMtest. But the reward maps in Figure 4c,d,e are sufficient to induce exploration toward the goal states ofMtest. Due to the dropout applied to the next state input of the decoder, the decoder assigns high rewards to some goals that belong to
Mtest. Note that the decoder does not generalize without dropout therefore assigns high rewards only to training goals (Figure 4a).

heading: 5.1.2 Extrapolation ability of LDM
text: To demonstrate that LDM is not restricted to target tasks inside the convex hull of training tasks, we design gridworld-extrapolation task as in Figure 5. This task is similar to the previous gridworld task in Section 5.1, except that we shift the outermost cells ofMtrain inside to construct extrapolation tasksMtest2 (Figure 5c). Refer to Figure 6 for the final return for each task when we train LDM (pdrop = 0.5) with different values of β. For small values of β = 1.0 and β = 1.5, LDM focuses on training the interpolation tasks inMtest1. As β increases, the returns for tasks inMtest1 decrease, but the returns for tasks inMtest2 increase.

heading: 5.2 MuJoCo
text: Experimental setup We evaluate our agent and baselines on three standard MuJoCo domains (Ant-direction, Ant-goal, and Half-cheetah-velocity) to verify how effective LDM is for continuous control tasks. We evaluate the agents on a fixed set of evaluation tasks Meval ⊂ Mtest for each domain to ensure the test result is not affected by the sampling of evaluation tasks (Table 1). For MuJoCo tasks we report the results of PEARL [35], MQL [8], and MIER [30] with number of rollouts per iteration (N ) equal to 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 5 million steps of environment interaction. We set LDM’s pdrop = 0.5 for all MuJoCo tasks.
Ant-direction We construct the Ant-direction task based on the example discussed in the introduction. θ 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 distributionMtrain (shaded region of Figure 8c) is continuous, unlike the gridworld and Ant-direction tasks. r and θ denote the radius and angle of the
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 β = 2 because LDM with β = 1 generates goals near the origin mostly due to the symmetry of training tasks. Refer to Appendix C for a detailed analysis regarding the choice of β. One may argue that the training and test tasks are not disjoint because the agent passes through the goal states ofMtest momentarily while learning to solveMtrain. However, the task inference with the unseen rewards fromMtest cannot be learned during training. The policy to stay at the goals ofMtest until the end of the time horizon is also not learned.
Half-cheetah-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 n = 12 normal workers and n̂ = 4 mixture workers in parallel. Refer to Appendix D for additional results with different numbers of mixture workers. The target velocity is v, where the velocity is measured as a change of position per second (or 20 environment steps). A major component of the reward is the negative value of the difference between the target velocity and the agent’s velocity. Similar to Ant-Goal, although the test velocities inMtest 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 inMtest during training inMtrain.
Results Refer to Figure 7 for the test returns inMeval. LDM outperforms the non-oracle baselines at the N -th episode. VariBAD oracle achieves the best test returns for all MuJoCo tests, revealing the strength of variBAD’s task inference ability for tasks seen during training. However, the non-oracle variBAD has difficulty generalizing to Mtest. The performance of Mixreg is lower than that of RL2, 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 at every timestep, even at the first step of an episode. Therefore gradient-based methods with feedforward networks also make progress, unlike the gridworld task. Because PEARL is designed to collect data for the first two episodes, it achieves low returns before the N -th episode. Even after accumulating context, the policy of PEARL is not prepared for the unseen contexts in Ant-direction and Half-cheetah-velocity. Note that MQL and MIER require additional buffer-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.
Episode 1
t = 0
Episode 1
t = 30
Episode 2
t = 60
Episode 3
t = 90
Episode 4
t = 120
(a) Sample trajectory of LDM for 4 rollout episodes. The red shade denotes decoder output for each state conditioned on the online context.
-3 -2 -1 0 1 2 3 x-position
-3 -2 -1 0 1 2 3 ypo sit io n
LDM
-3 -2 -1 0 1 2 3 x-position
RL2
-3 -2 -1 0 1 2 3 x-position
variBAD Oracle
(b) Ant-drection, 4 trajectories fromMeval.
Because LDM is not trained inMtest, it can not solve the target tasks optimally from the beginning of each task in the Ant-direction and Ant-goal tasks, unlike variBAD Oracle (Figure 8b and 8c). However, LDM reaches the target direction or goal close enough after a sufficiently small amount of exploration, unlike RL2. For the Half-cheetah-velocity task, LDM generates well-separated policies compared to RL2 (Figure 8d). Refer to Appendix F for additional experimental results on the test returns at the first rollout episode, the training returns, and sample trajectories on training tasks.

heading: 6 Conclusion
text: We propose Latent Dynamics Mixture to improve the generalization ability of meta-RL by training policy with generated mixture tasks. Our method outperforms baseline meta-RL methods in experiments with strictly disjoint training and test task distributions, even reaching the oracle performance in 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 RL2 but most of other meta-RL methods.
We believe that our work can be a starting point for many interesting future works. For example, instead of a heuristic weight-sampling for generating the mixture, we may incorporate OOD generation techniques in the latent space. Another extension is to train the latent dynamics network to decode not only the reward but also the state transition dynamics. Then we can generate a purely imaginary mixture task without additional environment interaction.

heading: Acknowledgments
text: This work was supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) [2021R1A2C2007518]. We thank Dong Hoon Lee, Minguk Jang, and Dong Geun Shin for constructive feedback and discussions.

Review summaries:
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.   This 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.   This method aims to improve on the problem of generalising to unseen tasks that are generated from an unseen task distribution.  

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. 

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.  The 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.   References  [1] Zintgraf, L., et al. "VariBAD: a very good method for Bayes-adaptive deep RL via meta-learning." Proceedings of ICLR 2020 (2020). 

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.