C ONSISTENT M ETA -R EINFORCEMENT L EARNING VIA M ODEL I DENTIFICATION AND E XPERIENCE R ELABELING
Anonymous authors
Paper under double-blind review
A BSTRACT
Reinforcement learning algorithms can acquire policies for complex tasks automatically, however the number of samples required to learn a diverse set of skills can be prohibitively large.
While meta-reinforcement learning has enabled agents to leverage prior experience to adapt quickly to new tasks, the performance of these methods depends crucially on how close the new task is to the previously experienced tasks.
Current approaches are either not able to extrapolate well, or can do so at the expense of requiring extremely large amounts of data due to on-policy training.
In this work, we present model identiﬁcation and experience relabeling (MIER), a meta-reinforcement learning algorithm that is both efﬁcient and extrapolates well when faced with out-of-distribution tasks at test time based on a simple insight: we recognize that dynamics models can be adapted efﬁciently and consistently with off-policy data, even if policies and value functions cannot.
These dynamics models can then be used to continue training policies for out-ofdistribution tasks without using meta-reinforcement learning at all, by generating synthetic experience for the new task.
I NTRODUCTION
Meta-reinforcement learning algorithms can enable acquisition of new tasks from just a small number of samples by leveraging experience from previous related tasks (Duan et al., 2016;
Wang et al., 2016; Finn et al., 2017).
However, the performance of these methods on new tasks depends crucially on how close the tasks are to the meta-training task distribution.
Meta-learned models can adapt quickly to tasks that are similar to those seen during training, but can lose much of their beneﬁt when adapting to tasks that are too far away from the meta-training distribution.
This places a signiﬁcant burden on the user to carefully construct meta-training task distributions that sufﬁciently cover the kinds of tasks that may be encountered at test time.
Many meta-RL methods either utilize a variant of model-agnostic meta-learning (MAML)
(Finn et al., 2017; Rothfuss et al., 2018;
Nagabandi et al., 2018), or an inference-based formulation with recurrent (Duan et al., 2016; Wang et al., 2016), attentional (Mishra et al., 2017), or variational (Rakelly et al., 2019) methods.
The latter class of methods generally fails to handle out-of-distribution tasks, because the model cannot adequately deal with out-of-distribution inputs corresponding to the new task.
Most of the former class of methods, where gradient-based adaptation corresponds to a well-deﬁned and consistent learning process, require on-policy samples, resulting in high sample complexity for meta-training in order to produce efﬁcient adaptation.
In this paper, we make use of a simple insight to develop model identiﬁcation and experience relabeling (MIER), a meta RL algorithm that is both efﬁcient and which extrapolates effectively when faced with out-of-distribution tasks at meta-test time: we recognize that dynamics and reward models can be adapted consistently, using MAML-style update rules with off-policy data, even if policies and value functions cannot.
These models can then be used to train new policies for out-of-distribution tasks without using meta-reinforcement learning at all, by generating synthetic experience for the new task.
To be able to quickly learn dynamics models, we reformulate the meta-reinforcement learning problem as one of MDP identiﬁcation .
We use this approach since identifying a task involves determining its transition dynamics and reward function, which is exactly the information that a model of the task needs to represent.
Speciﬁcally, we use a supervised meta-learning method that optimizes for a dynamics model initialization such that conditioned on a MDP context descriptor , prediction error on a validation batch of data sampled from the task is minimized, after updating only the context descriptor with a few gradient steps.
Effective model training requires the validation batch to contain data corresponding to optimal behavior for the tasks, which we obtain by training a universal policy conditioned on the context descriptor.
Note that since our formulation ensures that the context descriptor contains sufﬁcient information about the task, the policy does not need to be meta-trained or adapted, and can hence be learned with simple and efﬁcient off-policy RL algorithms, without needing to handle the complexity of meta-reinforcement learning.
At test time, given out-of-distribution tasks, the adapted context descriptor would indeed be out of distribution and thus our context-conditioned model and policy might not perform well for the test task.
However, since our method uses gradient descent which is a consistent learning method for adaptation, we can continue to improve our model using more gradient updates.
We then leverage all of the experience collected from other tasks during meta-training, by using the learned model to relabel the next state and reward information, thus obtaining synthetic data to continue training the policy.
Our main contribution is an off-policy meta-RL algorithm that is sample efﬁcient, stable and which extrapolates well to out-of-distribution tasks.
By formulating the meta-adaptation as MDP identiﬁcation , we are able to transform the meta-RL problem into a supervised meta-learning problem and thus beneﬁt from the stability and consistency of supervised learning methods.
The consistency of our model also enables us to continue improving our policy without collecting extra data by relabeling data collected from other tasks, thus allowing us to efﬁciently adapt to out-of-distribution tasks.
R ELATED W ORK
Meta-reinforcement learning algorithms extend the framework of meta-learning (Schmidhuber, 1987;
Thrun & Pratt, 1998; Naik & Mammone, 1992) to the reinforcement learning setting.
Current metareinforcement learning algorithms can be roughly categorized into 3 large categories, model-free context based methods, model-free gradient based methods and model based methods.
Model-free, context based methods such as Duan et al.
; Wang et al.
; Mishra et al.
Rakelly et al.
(2019) often encode the experience during adaptation into a latent context variable, and the policy is conditioned on the task speciﬁc context to adapt to a given task.
The context encoding process is often done via a recurrent network (Duan et al., 2016;
Wang et al., 2016), via an attention mechanism (Mishra et al., 2017) or via amortized probabilistic inference (Rakelly et al., 2019).
While many of these methods can efﬁciently summarize experiences during adaptation, it is often hard for them to encode any universal learning algorithm such as gradient descent.
This is because given a ﬁxed distribution of training tasks, the most efﬁcient way of adapting to a particular task is to directly infer it from the task distribution, rather than applying any universal learning algorithm.
While efﬁcient in handling in distribution tasks, the lack of universal learning algorithm in these methods implies that they would be able to adapt to out of distribution tasks.
Our method, on the other hand, is able to handle out of distribution tasks through continual adaptation with gradient descent.
Model-free gradient based meta-RL methods, such as Finn et al.
; Rothfuss et al.
Zintgraf et al.
; Rusu et al.
Liu et al.
(2019), implement gradient descent as the adaptation process.
The parameters are often optimized such that the model achieves good performance after only a few steps of gradient descent.
Since gradient descent is a universal learning algorithm, with a large capacity model, these methods are consistent in the way that they would guarantee to improve even on out-of-distribution tasks.
However, most of these methods are based on on-policy
RL algorithm, which means they are rather sample inefﬁcient at training time and requires more data to be collected when adapting to new tasks.
Our method mitigate this problem by applying an off-policy RL algorithm on top of a gradient based meta-learning algorithm to achieve better sample efﬁciency at training time.
At test time, our method reuses the experiences collected during training to enable fast adaptation without collecting a lot more data.
Model based meta-RL methods such as Nagabandi et al.
; Sæmundsson et al.
(2018), a rapidly adapting model instead of a policy is learned.
Since learning a fast adapting model is a supervised meta-learning problem, a broader range of meta-learning algorithms can be applied.
If an universal learning algorithm is used to represent the adaptation process, the model learning can then be made consistent such as in Nagabandi et al.
At test time, when the model is adapted to a particular task, standard planning techniques such as model predictive control is often applied to select actions.
While possible to be consistent, the model based meta-RL methods often suffer in long horizon tasks due to the accumulated error of model while planning long time ahead into the future.
Our method does not suffer from this problem since we use our model to train a model-free policy, and only execute the policy during test time.
This approach is similar to methods presented in Sutton (1991);
Janner et al.
P RELIMINARIES 3. R EINFORCEMENT L EARNING
Formally, a reinforcement learning problem is deﬁned by a Markov decision process (MDP).
We adopt the standard deﬁnition of a MDP T =
( S , A , p, µ 0 , r, γ ) , where S is the state space,
A is the action space, p ( s (cid:48) | s , a ) is the unknown transition probability of landing on next state s (cid:48)
at the next time step when an agent takes action a at state s , µ 0 ( s ) is the initial state distribution, r ( s , a ) is the reward function, and γ ∈ (0 , 1) is the discount factor.
An agent acts according to some policy π ( a | s ) and the objective of learning is to maximize the expected rewards E s t , a t ∼ π
[ (cid:80) t γ t r ( s t , a t )] .
3. M ODEL B ASED R EINFORCEMENT L EARNING
In model based reinforcement learning, a model ˆ p ( s (cid:48) , r | s , a ) that predicts the reward and next state from current state and action is trained using standard supervised learning approaches.
The model is then used to generate data to train a policy.
Speciﬁcally, given a model, we perform the following optimization to obtain the policy: arg max π E s t , a t ∼ π, ˆ p
[ (cid:80) t r ( s t , a t )] .
Note here the expectation is taken with respect to the model distribution of states instead of the true environment distribution.
3. M ETA -R EINFORCEMENT L EARNING
In meta-reinforcement learning, we represent various of similar tasks with a distribution of MDPs ρ ( T ) , where each speciﬁc task is a sample drawn from the distribution.
Given a speciﬁc task T , the agent is allowed to collect little amount of data D ( T ) adapt , adapt the policy to form π T according to the data.
The objective of this problem is to maximize the expected rewards of the adapted policy E
T ∼ ρ ( T ) , s t , a t ∼ π T [ (cid:80) t γ t r ( s t , a t )] .
3. S UPERVISED M ETA -L EARNING AND M ODEL A GNOSTIC M ETA L EARNING
We brieﬂy introduce the supervised meta-learning problem and the model agnostic meta-learning approach, which is an important foundation of our work.
In supervised meta-learning, we also have a distribution of tasks ρ ( T ) similar to the meta-RL setup, except that the task T is now a pair of input and output random variables ( X T , Y T ) .
Given a small dataset D ( T ) adapt sampled from a speciﬁc task
T , the objective is to built a model that performs well on the evaluation data D ( T ) eval sampled from the same task.
If we denote our model as f ( X ; θ ) , the adaptation process as Ξ( θ, D ( T ) adapt ) and our loss function as L , the objective can be written as:
[Equation]
Model agnostic meta-learning (Finn et al., 2017) is an approach to solve the supervised meta-learning problem.
Speciﬁcally, the model f
( X ; θ ) is represented as a neural network, and the adaptation process is represented as few steps of gradient descent.
For simplicity of notation, we only write out the one step of gradient descent:
[Equation]
Note that because Ξ MAML is the standard gradient descent, it does not contain any trainable parameters and hence the trainable parameters are all contained in the network weights θ .
Therefore, the training process of MAML can be summarized as optimizing the loss of the model after few steps of gradient descent on data from the new task.
C ONSISTENT M ETA -O PTIMIZATION BY R ELABELING E XPERIENCE
INPUT: task distribution ρ ( T ) , model training steps N model , policy training steps N policy , learning rate α OUTPUT: policy parameter ψ , model parameter θ , model context φ
Randomly initialize policy parameter ψ , model parameter θ , model con- text φ
Initialize multitask replay buffer R ( T ) ← ∅ while θ , φ , ψ not converged do Sample task T ∼ ρ ( T )
Collect data batch D ( T ) adapt before adaptation using π ψ and φ
Compute adapted context φ T =
Ξ MAML ( θ, φ, D ( T ) adapt )
Collect data batch D ( T ) eval after adaptation using π and φ T
Add data to replay buffer R ( T ) ← R ( T ) ∪ D ( T ) adapt ∪ D ( T ) eval for i = 1 to N model do Sample task in replay buffer T ∼ R
Sample two batches of data D ( T ) adapt , D ( T ) eval ∼ R ( T )
Train model ( θ, φ ) ← ( θ, φ ) − α ∇ θ,φ J ˆ p ( θ, φ, D ( T ) adapt , D ( T ) eval ) end for i = 1 to N policy do Sample task in replay buffer T ∼ R
Sample two batches of data D ( T ) adapt , D ( T ) eval ∼ R ( T )
Compute adapted context φ T =
Ξ MAML ( θ, φ, D ( T ) adapt )
Train policy ψ ← ψ − α ∇ ψ J π ( ψ, D ( T ) eval ) end end
In this section, we lay the metatraining and test time adaptation process of MIER.
We ﬁrst reformulate the meta-RL problem into a meta-model identiﬁcation problem, where we train a fast adapting model to rapidly identify transition dynamics and reward function for a given new task.
We parameterize the model with a latent context descriptor to contain all task speciﬁc information aquired during adaptation.
We then train a universal policy that condition on context descriptor to solve all tasks sampled from the training distribution.
At test time, given a potentially out of distribution task, we continual to improve our policy by using our model to relabel data from previously experienced tasks.
The consistency of our model adaptation process and policy improvement ensures that the policy will continue to improve even the task is out-ofdistribution.
We provide the detailed explain of each part of our method in the following sections, and we lay out the pseudo code for meta-training time in Algorithm 1 and test time in Algorithm 2.
4. M ETA -M ODEL I DENTIFICATION WITH L ATENT C ONTEXT
In a meta-RL problem, tasks are drawn from a distribution of MDPs, and an agent adapts to each of the tasks by observing some data for the tasks.
The adaptation process is usually performed as either an implicit or explicit inference of some task representations.
That is, the agent adapts to a task by identifying it from other tasks in the same distribution.
In a distribution of MDPs, the only varying factors are the dynamics p and the reward function r .
Therefore, a sufﬁcient condition for identifying the task is to learn the transition dynamics and the reward function, and this is exactly what model based RL methods do.
Hence, we can naturally formulate the meta-task identiﬁcation problem as model based RL and solve it with supervised meta-learning methods.
Speciﬁcally, we choose the MAML method for its simplicity and consistency.
Unlike the standard supervised MAML formulation, we condition our model on a latent context vector, and we only change the context vector when adapting to new tasks.
With this formulation, we restrict all the task speciﬁc information learned during adaptation to the context vector, and thus allowing the policy to condition on it.
Let us denote our model as ˆ p ( s (cid:48) , r | s , a ; θ, φ ) , where θ is the neural network parameters and φ is the latent context vector that is passed in as input to the network.
One step of gradient adaptation can be written as follows:
[Equation]
We use the log likelihood as our objective for the probabilistic model.
Then we evaluate the adapted context vector φ T and minimize its loss on the evaluation dataset to learn the model.
Speciﬁcally, we minimize the model meta-loss function J ˆ p ( θ, φ, D ( T ) adapt , D ( T ) eval ) to obtain the optimal parameter θ and context vector initialization φ .
[Equation]
The main difference between our method and previously proposed context based meta-RL methods (Rakelly et al., 2019; Duan et al., 2016) is that we use gradient descent to adapt the context.
Therefore, given a model with large enough capacity, the adaptation process is consistent, in the sense that it will continue to improve regardless of which task it is presented with.
In-distribution tasks will of course have much faster adaptation, since the meta-training process explicitly optimizes for this objective, but the model for out-of-distribution tasks will still adapt to the task given enough samples and gradient steps, since the adaptation process corresponds to a well-deﬁned and convergent learning process.
4. P OLICY O PTIMIZATION WITH L ATENT C ONTEXT
Given the latent context variable from the adapted model, the meta-RL problem can be effectively reduced to a standard RL problem, as the task speciﬁc information has been all encoded in the context variable.
Therefore we can apply any model-free RL algorithm to obtain a policy, as long as we condition the policy on the latent MDP descriptor context.
Speciﬁcally, we choose the soft actor critic (SAC) method (Haarnoja et al., 2018) for its sample efﬁciency and stability.
Let us parameterize our policy π ψ by a parameter ψ .
The soft actor critic method maintains an estimate of the values for the current policy Q π ψ ( s , a , φ T ) =
E s t , a t ∼ π ψ
[ (cid:80) t γ t r ( s t , a t ) | s
=
s, a 0 = a, T ] and V π ψ ( s , φ T ) =
E s t , a t ∼ π ψ
[ (cid:80) t γ t r ( s t , a t ) | s
= s, T ] through Bellman backup, and improves the policy by minimizing the KL divergence between the policy and the exponential advantage J π ( ψ, D ) =
E D [ D KL ( π ψ || exp
{ Q π ψ − V π ψ } )] over some dataset D .
Note that we condition our value functions Q π ψ , V π ψ and policy π ψ all on the adapted task speciﬁc context vector φ T so that the policy and value functions can access full information about the task.
4. C ONTINUAL A DAPTATION AND E XTRAPOLATION THROUGH D ATA R ELABELING Algorithm
MIER
INPUT: test task ˆ T , multitask replay buffer R ( T ) , model training steps N model , policy training steps N policy , policy parameter ψ , model parameter θ , model context φ , learning rate α
OUTPUT: policy parameter ψ , model parameter θ , model context φ
Collect data batch D (ˆ T ) adapt from test task ˆ T using π ψ and φ
Compute adapted context φ ˆ T =
Ξ MAML ( θ, φ, D (ˆ T ) adapt ) while θ , ψ not converged do for i = 1 to N model do
Train model θ ← θ − α ∇ θ E D
(ˆ T ) adapt [ − log ˆ p ( s (cid:48) , r | s , a ; θ, φ ˆ T )] endfor i = 1 to N policy do Sample task in replay buffer T ∼ R
Sample data batch D ( T ) ∼ R ( T )
Relabel data ˆ D (ˆ T ) ← Relabel ( D ( T ) , θ, φ ˆ T )
Train policy ψ ← ψ − α ∇ ψ J π ( ψ, ˆ D (ˆ T ) ) endend
At meta-test time, when the model is given an unseen task
T , it will ﬁrst sample a small batch of data and obtain the latent context φ T by running the gradient descent adaptation process on the context variable.
If the new task is in-distribution, we can directly feed φ T to the policy.
However, in many cases, we might encounter tasks that are out of distribution.
If the task T is out of the meta-training distribution, then the obtained latent context φ T will also be out of distribution for the policy.
Even though our model is adapted with a consistent method, the improvement guarantee is only for the model – the latent context φ T that might be obtained for an outof-distribution task might not immediately produce an effective policy.
However, with an improved model, we can generate new data for the task without sampling from the environment to continue training the policy, providing a sample-efﬁcient (but computationally more costly) adaptation method for out-of-distribution tasks.
When using data generated from a learned model to train a policy, a common caveat is that the model’s predicted trajectory often diverges from the trajectory under the real dynamics in long horizon tasks, due to accumulated error.
In the special case of meta-RL, we argue that this issue can be avoided by the multi-task setup in meta-RL.
Although we have never seen this particular new task, we have seen many other tasks that share the same state space and action space, and we have collected a lot of experience from these tasks.
Using the adapted model and policy, we can relabel the transitions ( s , a , s (cid:48) , r ) from any other tasks by sampling a new actions with our adapted policy, and sample next states and rewards with the adapted model.
The relabeling process can be written as Relabel ( D , θ, φ T ) =
{ ( s , a , s (cid:48) , r ) | ( s , a ) ∈ D ; ( s (cid:48) , r ) ∼ ˆ p ( s (cid:48) , r | s , a ; θ, φ T ) } .
We use these relabeled transitions to continue training the policy.
By doing so, we effectively transferred experiences from other tasks into this new task, and thus avoided the divergence problem of sampling long horizon trajectories as we are only doing one step prediction.
This relabeling scheme is similar to the Dyna algorithm introduced by Sutton (1991).
By relabeling experiences from tasks seen before, we can efﬁciently reuse the large amount of data collected from other tasks to improve our policy on this new task.
E XPERIMENTAL E VALUATION
We aim to answer the following questions in our experiments: (1) Can MIER meta-train efﬁciently on standard meta-RL benchmarks without requiring large amounts of data?
(2) How does MIER compare to prior meta-learning approaches for extrapolation to meta-test tasks with out-of-distribution (a)
reward functions and (b) dynamics?
(3) How important is cross task relabeling in leveraging the model to learn an effective policy efﬁciently?
To answer these questions, we ﬁrst compare the sample efﬁciency of MIER to existing methods on several standard meta-RL benchmark environments.
We then test MIER on a set of environment with out of distribution meta-test tasks to analyze extrapolation performance.
We further run an ablation of our method without cross task labelling, and show that this isn’t able to adapt the policy as effectively as our method.
5. S AMPLE E FFICIENCY ON M ETA -RL B ENCHMARKS
We evaluate MIER on two standard meta-RL benchmark environments based on the MuJoCo physics simulator (Todorov et al., 2012), shown in Figure 1.
These environments were introduced by Finn et al.
(2017), and later used by the state-of-the-art context based meta-RL method Rakelly et al.
They include half cheetah robot running at different velocities (HalfCheetah-Vel) and a humanoid robot running in different directions (Humanoid-Dir).
For all environments, we follow the exact setting in Rakelly et al.
For comparison, we include performance of the PEARL algorithm (Rakelly et al., 2019) and MAML algorithm
(Finn et al., 2017) as our baselines, and we average the performance of all algorithms across 3 random seeds.
We plot the test time performance after adaptation versus number of steps in the environment in
Figure 2.
We see that our method is able to achieve performance comparable to the state of the art method in meta-RL in terms of sample efﬁciency (PEARL).
This is due to the ability of our method to perform supervised meta-learning for the model, and off-policy reinforcement learning for the policy, which results in a stable and efﬁcient optimization procedure during meta-training.
Note that here, off-policy refers to the meta-training process.
However, as we discuss in Section 4.3, our method can also perform off-policy adaptation via cross task relabeling.
We evaluate this in the next section.
5. E XTRAPOLATION TO O UT - OF -D ISTRIBUTION T ASKS
We evaluate the extrapolation capability of MIER separately on environments with varying reward functions, and those with different dynamics.
In order to ensure fair comparisons, we train all algorithms to convergence before testing on new tasks.
Hence, MAML is trained for 80 millions steps, while PEARL and MIER only require about a million steps.
Unless stated otherwise, we adapt with MIER using cross task relabeling in all of these experiments.
5.2. A DAPTATION TO O UT - OF -D ISTRIBUTION R EWARDS
For testing extrapolation ability to tasks with different rewards, we take the HalfCheetah-Vel environment from the standard benchmarks and restrict the meta-training task distribution.
This is required since, in the original environments, the meta-training distribution covers almost the entire feasible task space, leaving little room for extrapolation.
Instead, if we limit the training distribution, we can easily ﬁnd out-of-distribution tasks to test the extrapolation capability of our method.
The training tasks consist of target velocities at which the cheetah needs to run sampled uniformly from 0 to 1.
We then evaluate the ability of the meta-learned polices to adapt in order to run at new speeds chosen at 0.5 intervals from 0.5 to 2.5, using only 2 trajectories sampled from the environment.
From Figure 3, we see that MIER attains much higher return than other methods for target velocities that are substantially outside of the meta-training distribution.
Eventually, performance degrades for all methods, since the state space distribution differs too much from the meta-training tasks, making cross task relabeling ineffective.
Nonetheless, this experiment shows that MIER’s capacity to adapt to out-of-distribution tasks substantially exceeds that of prior meta-RL methods.
5.2. A DAPTATION TO O UT - OF -D ISTRIBUTION D YNAMICS
To study adaptation to out-of-distribution dynamics, we take the HalfCheetah environment, and randomly negate the control of 3 joints out of the total 6 joints of the robot.
During meta-training, we only select joints from the ﬁrst 5 joints to negate, and during meta-test time, we select 2 joints from the ﬁrst 5 and always negate the 6th joint.
This ensures that the agent never sees the negation of the
6th joint during meta-training time, and hence needs to extrapolate well to get good return on the test tasks.
In total, we have 10 training tasks and 10 test tasks.
We compare the return of all methods after adapting on 2 new trajectories of the validation task.
From the plot of average performance across test tasks, in Figure 4, we see that MIER achieves superior performance compared to existing methods.
We are able to extrapolate so efﬁciently since we are able to leverage off-policy data collected during meta-training.
This in turn is enabled by our ability to learn a model quickly using minimal data, which is the consistent meta-training objective.
To further highlight the importance of relabeling data from other tasks during adaptation, we run an ablation where the policy is learned only using the sampled data.
Figure 4 shows performance of the policy plotted against the number of steps where we relabel exsiting data and continue to train the policy.
We see that without the beneﬁt of off-policy data, the policy performance deteriorates rapidly due to overﬁtting to the small batch of data sampled from the test task.
C ONCLUSION
In this paper, we introduce a consistent and sample efﬁcient meta-RL algorithm by reformulating the meta-RL problem as model identiﬁcation.
Our algorithm can adapt to new tasks by determining the parameters of the model, which predicts the reward and future transitions.
This allows us to perform meta-learning via supervised learning of the model, which is more stable and sample efﬁcient.
More importantly, this provides us with a consistent adaptation procedure, where adaptation is performed via gradient descent.
This means that, even for out-of-distribution meta-test tasks, our method eventually learns the right model.
This model then allows us to relabel past experience to reuse it for the new task, running off-policy reinforcement learning to acquire the policy for the new task.
Experiment results show that our method achieves superior performance compared to existing methods, especially on out-of-distribution tasks that requires extrapolation, where consistent adaptation of the model followed by relabeling and off-policy reinforcement learning substantially outperforms prior methods that adapt to the new task directly Appendices A I MPLEMENTATION D ETAILS The implementation of the algorithm largely builds on the Model Based Policy Optimization code ( Janner et al. , and the environments are standard meta-RL benchmarks, with slight modiﬁcations for the extrapolation tasks as described.
The context vector is a variable that is appended to the states and actions when sent to the model, and appended to the states when passed to the policy.
The hyper-parameters were kept mostly ﬁxed across all experiments.
One major change is that environments with changing dynamics had models predicting both next state and reward, while the environments with only changing reward only had prediction for reward.
Some key hyperparameters (common across experiments):
Number of fast adapt steps for model meta-training :2.
Fast adapt learning rate for the model : 0. 3.
Batch size for model adaptation :4.
Reward-dynamics model architecture : 200-200-2005.
Dimension of context vector :6.
Regression weight on post-update contexts : 0.
Number of tasks sampled per training epoch :8.
Number of data points sampled from a task per training epoch :9.
Number of SAC training steps per training epoch :10.
Learning rate for Q functions in SAC : 0.
Q network hidden layers : 256-