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

1 Introduction . 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 can not 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 . 2 Problem Setup . 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 can not 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 . 3 Latent Dynamics Mixture . 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 . 3.1 Policy Network . 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 . 3.2 Latent Dynamics Network . 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 .

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.