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

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-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), 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 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.

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=(\mathcal{S},\mathcal{A},R,T,T_{0},\gamma,H)$ consists of a set of states $\mathcal{S}$, a set of actions $\mathcal{A}$, a reward function $R(r_{t+1}|s_{t},a_{t},s_{t+1})$, a transition function $T(s_{t+1}|s_{t},a_{t})$, an initial state distribution $T_{0}(s_{0})$, a discount factor $\gamma$ 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 $\mathcal{M}$ at every iteration. Each MDP $M_{k}=(\mathcal{S},\mathcal{A},R_{k},T_{k},T_{0},\gamma,H)$ has individual reward function $R_{k}$ (e.g., goal location) and transition function $T_{k}$ (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 split $\mathcal{M}$ into two strictly disjoint training and test sets of MDPs, i.e., $\mathcal{M}=\mathcal{M}_{\textrm{train}}\cup\mathcal{M}_{\textrm{test}}$ and $\mathcal{M}_{\textrm{train}}\cap\mathcal{M}_{\textrm{test}}=\emptyset$. The agent does not have any prior information about $\mathcal{M}_{\textrm{test}}$ and cannot interact in $\mathcal{M}_{\textrm{test}}$ during training.

Since the MDP is initially unknown, the best the agent can do is to update its belief $b_{t}(R,T)$ about the environment according to its experience $\tau_{:t}=\{s_{0},a_{0},r_{1},s_{1},a_{1},r_{2},\ldots,s_{t}\}$. 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, $b_{t}(R,T)=p(R,T|\tau_{: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.¹¹1We use the terms context, latent belief, and latent (dynamics) model interchangeably to denote $m$. The reward and transition dynamics can be formulated as shared functions across MDPs: $R(r_{t+1}|s_{t},a_{t},s_{t+1};m)$ and $T(s_{t+1}|s_{t},a_{t};m)$. Then the problem of computing the posterior $p(R,T|\tau_{:t})$ becomes inferring the posterior $p(m|\tau_{:t})$ over $m$. By conditioning the policy on the posterior $p(m|\tau_{: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 set $\mathcal{M}$ is divided into disjoint $\mathcal{M}_{\textrm{train}}$ and $\mathcal{M}_{\textrm{test}}$. A Bayes-optimal agent for a task in $\mathcal{M}_{\textrm{train}}$ first assigns a uniform prior to the goal states of $\mathcal{M}_{\textrm{train}}$ (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 in $\mathcal{M}_{\textrm{train}}$ only, is put in to solve a task in $\mathcal{M}_{\textrm{test}}$ without any prior knowledge (Figure 1b), the best policy an agent can take first is to maintain its initial belief learned in $\mathcal{M}_{\textrm{train}}$ and explore the goal states of $\mathcal{M}_{\textrm{train}}$. Once the agent realizes that there are no goal states in $\mathcal{M}_{\textrm{train}}$, it could start exploring the states that are not visited (i.e., $\mathcal{M}-\mathcal{M}_{\textrm{train}}$), and discover an unseen goal state in $\mathcal{M}_{\textrm{test}}$. However, it is unlikely that the agent trained only in $\mathcal{M}_{\textrm{train}}$ will encode its experience into beliefs for unseen tasks accurately and explore the goal states of $\mathcal{M}_{\textrm{test}}$ efficiently conditioned on the unseen context without any prior knowledge or test-time adaptation.

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 in $\mathcal{M}_{\textrm{train}}$ 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)},\ldots,W^{(n)}$ and a mixture worker $\hat{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)},\ldots,M^{(n)}$ from $\mathcal{M}_{\textrm{train}}$ and assign each MDP $M^{(i)}$ to each normal worker $W^{(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 worker trains the policy network with imaginary rewards from the learned decoder’s output given mixture beliefs.

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 space $\mathcal{M}$ into disjoint $\mathcal{M}_{\textrm{train}}$ and $\mathcal{M}_{\textrm{test}}$. We use RL² [6] and variBAD [56] as baselines representing the context-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 RL²-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 from $\mathcal{M}_{\textrm{train}}$ uniformly at random, except for the oracle methods that are trained on tasks in the entire task set $\mathcal{M}=\mathcal{M}_{\textrm{train}}\cup\mathcal{M}_{\textrm{test}}$. 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 $\mathcal{M}_{\textrm{test}}$ require additional exploration during testing to experience the changes from $\mathcal{M}_{\textrm{train}}$, whereas the oracle agent can exploit the knowledge of the test distribution. In the gridworld example, the oracle agent can search for goals in $\mathcal{M}_{\textrm{test}}$ before navigating the outermost goal states of $\mathcal{M}_{\textrm{train}}$. 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.

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

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.

1. 1.

   For all authors…

   1. (a)

      Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes]

   2. (b)

      Did you describe the limitations of your work? [Yes] In the conclusion section

   3. (c)

      Did you discuss any potential negative societal impacts of your work? [N/A]

   4. (d)

      Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]

2. 2.

   If you are including theoretical results…

   1. (a)

      Did you state the full set of assumptions of all theoretical results? [N/A]

   2. (b)

      Did you include complete proofs of all theoretical results? [N/A]

3. 3.

   If you ran experiments…

   1. (a)

      Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplementary material or as a URL)? [Yes] as a URL in Appendix

   2. (b)

      Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] in Appendix

   3. (c)

      Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes]

   4. (d)

      Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] in Appendix

4. 4.

   If you are using existing assets (e.g., code, data, models) or curating/releasing new assets…

   1. (a)

      If your work uses existing assets, did you cite the creators? [Yes] source codes of the baselines

   2. (b)

      Did you mention the license of the assets? [Yes] in the supplementary material

   3. (c)

      Did you include any new assets either in the supplementary material or as a URL? [Yes] Links to the baseline codes in Appendix

   4. (d)

      Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [N/A]

   5. (e)

      Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]

5. 5.

   If you used crowdsourcing or conducted research with human subjects…

   1. (a)

      Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]

   2. (b)

      Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]

   3. (c)

      Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]