A UGMENTED W ORLD M ODELS F ACILITATE Z ERO -S HOT D YNAMICS G ENERALIZATION F ROM A S INGLE O FFLINE E NVIRONMENT
A BSTRACT
Reinforcement learning from large-scale ofﬂine datasets provides us with the ability to learn policies without potentially unsafe or impractical exploration.
Signiﬁcant progress has been made in the past few years in dealing with the challenge of correcting for differing behavior between the data collection and learned policies.
However, little attention has been paid to potentially changing dynamics when transferring a policy to the online setting, where performance can be up to 90% reduced for existing methods.
In this paper we address this problem with Augmented
World Models (AugWM). We augment a learned dynamics model with simple transformations that seek to capture potential changes in physical properties of the robot, leading to more robust policies.
We not only train our policy in this new setting, but also provide it with the sampled augmentation as a context, allowing it to adapt to changes in the environment.
At test time we learn the context in a self-supervised fashion by approximating the augmentation which corresponds to the new environment.
We rigorously evaluate our approach on over 100 different changed dynamics settings, and show that this simple approach can signiﬁcantly improve the zero-shot generalization of a recent state-of-the-art baseline, often achieving successful policies where the baseline fails.
I NTRODUCTION
Ofﬂine reinforcement learning (RL) describes the problem setting where RL agents learn policies solely from previously collected experience without further interaction with the environment (12; 29).
This could have tremendous implications for real world problems (10), with the potential to leverage rich datasets of past experience where exploration is either not feasible (e g. a Mars Rover) or unsafe (e.g. in medical settings).
As such, interest in ofﬂine RL has surged in recent times.
This work focuses on model-based ofﬂine RL, which has achieved state-of-the-art performance through the use of uncertainty penalized updates (45; 23).
However, existing work only addresses the issue of transferring from different behavior policies in the same environment , ignoring any possibility of distribution shift. Consider the case where it is expensive to collect data, and we have access to a single dataset from a robot.
Using existing methods we would be unable to make any changes that impact the dynamics, such as using a newer model of the robot or deploying it in a different room.
A related setting is the Sim2Real problem which considers transferring an agent from a simulated environment to the real world.
A popular recent approach is domain randomization (42; 20), the process of randomizing non-essential regions of the observation space to make agents robust to ‘observational overﬁtting’ (41).
Indeed, methods seeking to generalize to novel dynamics have also shown promise (34), by randomizing physical properties such as the mass of the agent.
A signiﬁcant limitation of these approaches is the requirement for a simulator, which may not be available. In this work we take inspiration from Sim2Real to generalize solely from an ofﬂine dataset, in a learned simulator or World Model (WM).
We therefore describe our problem setting as follows: an [Figure] agent must learn to generalize to unseen test-time dynamics whilst having access to ofﬂine data from only a single environment ; we call this “dynamics generalization from a single ofﬂine environment”.
In this paper we concentrate on the zero-shot performance of our agents to unseen dynamics, as it may not be practical (nor safe) to perform multiple rollouts at test time.
To tackle this problem, we propose a novel form of data augmentation: rather than augment observations, we focus on augmenting the dynamics .
We ﬁrst learn a world model of the environment, and then augment the transition function at policy training time, making the agent train under different imagined dynamics.
In addition, our agent is given access to the augmentation itself as part of the observations, allowing it to consider the context of modiﬁed dynamics.
At test time we propose a simple, yet surprisingly effective, self-supervised approach to learning an agent’s augmentation context.
We learn a linear dynamics model which is then used to approximate the dynamics augmentation induced by the modiﬁed environment.
This context is then given to the agent, allowing it to adapt on the ﬂy to the new dynamics within a single episode (i.e. zero-shot).
We show that our approach is capable of training agents that can vastly outperform existing Ofﬂine RL methods on the “dynamics generalization from a single ofﬂine environment” problem.
We also note that this approach does not require access to environment rewards at test time.
This facilitates application to Sim2Real problems whereby test time rewards may not be available.
Our contributions are twofold: 1) As far as we are aware, we are the ﬁrst to propose dynamics augmentation for model based RL , allowing us to generalize to changing dynamics despite only training on a single setting.
We do this without access to any environment parameters or prior knowledge.
We propose a simple self-supervised context adaptation reward-free algorithm, which allows our policy to use information from interactions in the environment to vary its behavior in a single episode, increasing zero-shot performance.
We believe both of these approaches are not only novel, but offer signiﬁcant improvement v.s. state-of-the-art methods, improving generalization and providing a promising approach for using ofﬂine RL in the real world.
R ELATED W ORK
In this work we focus on Model Based RL (MBRL).
A key challenge in MBRL is that an inaccurate model can be exploited by the policy, leading to behaviors that fail to transfer to the real environment.
As such, a swathe of recent works have made use of model ensembles to improve robustness (26; 7; 9; 5; 21).
With increased accuracy, MBRL has recently been shown to be competitive with model free methods in continuous control (14; 21; 7) and games (39; 22).
We make use of the ensemble of probabilistic dynamics models, ﬁrst introduced in (7), and subsequently used in (21).
In this paper we focus on Model-Based ofﬂine RL, where MOPO (45) and MOReL (23) have recently demonstrated the effectiveness of learned dynamics models, using model uncertainty to constrain policy optimization.
We build upon this approach for zero-shot dynamics generalization from ofﬂine data.
There have also been successes in off policy methods for ofﬂine RL (44; 25; 12) and context based approaches (1), although these works only consider tasks within the support of the ofﬂine dataset.
Finally, MBOP (4) addresses the problem of goal-conditioned zero-shot transfer from ofﬂine datasets.
However, their goal-conditioning relies on unchanged dynamics in the test environment.
In online RL, recent work has achieved strong dynamics generalization with a learned model (40).
However, this required training under varied dynamics, assigning different experiences to models.
In addition, this work used MPC whereas we train a policy inside the model, which is signiﬁcantly faster at deployment time.
Also related is (8; 30), where the model is trained to quickly to adapt to new dynamics P ( s (cid:48) | s, a ) , however both these works place more emphasis on model-adaption rather than zero-shot policy performance.
Furthermore, access to an underlying task distribution is required, something we do not have in our ofﬂine setting.
Also similar to our work is the recently proposed Policy Adaptation during Deployment (PAD, (18)) approach.
Our approach differs in that we learn a context, whereas PAD uses a auxiliary objective to adapt its features.
In addition, PAD considers the online model free setting, while our method is ofﬂine and model based .
Sim2Real is the setting where an agent trained in a simulator must transfer to the real world.
A common approach to solve this problem is through domain randomization (42; 20), whereby parameters in the simulator are varied during training.
This has shown to be effective for dynamics generalization (2; 3; 34; 46; 47; 32), but requires access to a simulator which we do not have.
Another form of domain randomization, data augmentation, has proved to be effective for training RL policies (27; 28; 24; 36), resulting in improved efﬁciency and generalization.
So far, these works have focused on online model free methods, and used data augmentation on the state space, reducing observational overﬁtting (41).
In contrast, we focus on ofﬂine MBRL and instead augment the dynamics .
We are not the ﬁrst to propose data augmentation in the MBRL setting, (35) proposed Counterfactual data augmentation for improving performance in the context of locally factored tasks.
Approaches to ensuring adversarial robustness can include data augmentations that assist with out-of-domain generalization, as opposed to observational overﬁtting.
In (43) this is done without a simulator and from a single source of data, however they only work on supervised learning problems and require an adversary to be learned, adding computational complexity.
P RELIMINARIES
We consider a Markov Decision Process (MDP), deﬁned as a tuple M =
( S , A , P, R, ρ 0 , γ ) , where
S and A denote the state space and action space respectively, P ( s (cid:48) | s, a ) the transition dynamics,
R ( s, a ) the reward function, ρ 0 the initial state distribution, and γ ∈ (0 , 1) the discount factor.
The goal in RL is to optimize a policy π ( a | s ) that maximizes the expected discounted return E π,P,ρ
[ (cid:80) ∞ t =
γ t R ( s t , a t )] .
The value function V π ( s ) :
= E π,P [ (cid:80) ∞ t =
γ t R ( s t , a t ) | s 0 = s ] gives the expected discounted return under π when starting from state s .
In ofﬂine RL , the policy is not deployed in the environment until test time.
Instead, the algorithm only has access to a static dataset
D env = { ( s, a, r, s (cid:48) ) } , collected by one or more behavioral policies π b .
We borrow notation from (45) and refer to the distribution from which D env was sampled as the behavioral distribution .
When training a model, we follow MBPO (21) and MOPO (45) and train an ensemble of N probabilistic dynamics models (31).
Each of the models learns to predict both the next state s (cid:48) and reward r from a state-action pair, using D env in a supervised fashion.
Concretely, each of the N models output a Gaussian (cid:98) P i ( s t +1 , r t | s t , a t ) =
N ( µ ( s t , a t ) , Σ( s t , a t )) .
The resulting dynamics model (cid:98) P deﬁnes a model MDP (cid:99)
M = ( S , A , (cid:98) P, (cid:98) R, ρ 0 , γ ) , where (cid:98) R refers to the learned reward model.
To train the policy, we use k step rollouts inside (cid:99) M , adding experience to a replay buffer D (cid:99) env to learn an action-value function and a policy, using Soft Actor Critic (SAC (15)).
SAC alternates between a soft policy evaluation step, which estimates Q π ( s, a ) =
E π [ (cid:80) ∞ t =
γ t ( R ( s t , a t )+ α H ( π ( ·| s k )))
| s
=
s, a 0 = a ] using Bellman backups (where α is a temperature parameter for policy entropy H ), and a policy improvement which learns a policy by minimizing the expected KL divergence J π ( φ, D )
=
E s t ∼D (cid:99) env [ D KL ( π (cid:107) exp { 1 α { Q π − V π }} )] .
Note that the SAC algorithm is unchanged from the modelfree setting, aside from the environment being a learned model, and the rollout horizon k being truncated.
Perhaps surprisingly, this approach alone produces strong results in the ofﬂine setting.
However MOPO (45) and MORel (23) show improvement in performance by penalizing rewards in regions of the state space where the ensemble of probabilistic models is less certain.
This implicitly reduces reliance on samples which deviate beyond the support of D env .
While MOPO and MORel have addressed the issue of training a policy in D env , and transferring to the true environment M , they only consider where the data in D env is actually drawn from P .
However, sometimes this may not be sufﬁcient for deployment.
For example, a robot could fail to walk when learning from data that was collected by a different version of the robot (with different mass), or if the same robot collected data but in a different room to deployment (with varied friction).
It is this setting, where dynamics may vary at test time, that is the focus of our work.
To learn successfully we propose a novel approach to training robust context-dependent policies.
[Figure] [Figure]
A UGMENTED W ORLD M ODELS WITH S ELF S UPERVISED P OLICY A DAPTATION
In this section we introduce our algorithm: Augmented World Models (AugWM).
We ﬁrst discuss our training procedure (Fig. 2) before moving onto our self-supervised approach to online context learning (Fig. 3).
4. A UGMENTED W ORLD M ODELS
Rather than seeking to transfer our policy from (cid:99) M to M , we instead wish to transfer from (cid:99) M to
M (cid:63) , where M (cid:63) =
( S , A , P (cid:63) , R (cid:63) , ρ 0 , γ ) is an unseen environment with different dynamics .
Thus, our emphasis shifts to producing experience inside the world model such that our agent is able to generalize to unseen, out of distribution dynamics.
We approach this problem by training our policy in an Augmented World Model .
Formally, we denote an augmentation as z t ∼ Z , z ∈ R |S| , which is sampled at each timestep to produce an augmentation operator T z .
T z is applied to ( s, a, r, s (cid:48) ) tuples from the dataset D , and when used with tuples sampled from D (cid:100) env indirectly induces an augmented distribution (cid:98) P z .
In principle, we wish to produce augmentations such that the true modiﬁed environment dynamics P (cid:63) lies in the support of the distribution of augmented dynamics, i.e.
[Equation]
for all s, a , some small (cid:15) > 0 , and suitable distance/divergence metric D .
We consider several augmentations, beginning with existing works before moving to new approaches which speciﬁcally target the problem of dynamics generalization.
We begin with Random Amplitude Scaling as in (27), which we refer to as RAD.
RAD scales both s t and s t +1 as follows:
[Equation]
for z ∼ Unif([ a, b ] |S| ) .
Given that our focus is on changing dynamics (vs. observational overﬁtting), we also propose to scale only the next state, i.e., Random Amplitude Nextstate Scaling (RANS):
[Equation]
for z ∼ Unif([ a, b ] |S| ) .
Note that while RANS is more focused on augmenting dynamics than RAS, it still suffers from a dependence on the magnitude of s t +1 .
As such, we further propose a more targeted augmentation, which we call Dynamics Amplitude Sampling (DAS).
Rather than directly scale the state, DAS scales the change in the state which we denote with δ t =
s t +1 − s t , as follows:
[Equation]
for z ∼ Unif([ a, b ] |S| ) .
The full training procedure is shown in Algorithm 1.
One crucial addition to our method is the use of context.
Concretely, when we are optimizing the policy using a batch of data, we concatenate the next state with the augmentation vector z .
This allows our policy to be informed of the speciﬁc augmentation that was applied to the environment and thus behave accordingly.
However, at test time we do not know z , so what can we use?
Next we propose a solution to this problem, learning the context on the ﬂy.
4. S ELF -S UPERVISED P OLICY S ELECTION
In the meta-learning literature there have been many recent successes making use of a learned context to adapt a policy at test time to a new environment (38; 48), typically using a blackbox model with a latent state.
Crucially, these approaches require several episodes to adapt at test time, making them unfeasible in our zero-shot setting.
What makes our setting unique is we explicitly know what the context represents: a linear transformation of s t , or δ t .
Using this insight, we are able to learn an effective context on the ﬂy at test time.
Concretely, we observe that from a state s t drawn from M (cid:63) , we can sample an action a t ∼ π and then compute an approximate (cid:98) s t +1 using our model (cid:98) P .
With (cid:98) s t +1 , we have a sample estimate of the state change under (cid:98) P , i.e. (cid:98) δ t =
(cid:98) s t +1 − s t .
We can make this approximation of the next state without interacting with the environment, but once we do take the action a t in the environment, we then receive the true next state s t +1 and can store the true difference δ t =
s t +1 − s t .
Using the DAS augmentation, we can approximate z as δ t / ˆ δ t .
Input: Initial state s 1 , horizon H , policy π , world model (cid:98) P , initial context ˆ z = 1 |S|
Initialize: Linear model f ψ , dataset D
=
Ø, return R 0 = 0 .
for step = 1 , 2 ,...
H do
Select action: a t ∼ π ( s t , (cid:98)
z t ) Take action: s t +1 ∼ P (cid:63) ( s t ,a t )
Update return: R t +1 =
R t + R (cid:63) ( s t ,a t ,s t +1 )
Update Dataset: D ∪ ( s t ,δ t ) .
Update Linear model by minimizing L MSE ( ψ, D ) .
Predict new context using (cid:98) z t = δ t / ˆ δ t .
endReturn: R T
[Figure] This however is retrospective; we can only approximate z having already seen the next state, by which time our agent has already acted.
Furthermore, we believe under changed dynamics the true z likely depends on s , thus we cannot use a previous z for future timesteps.
Therefore, we learn a forward dynamics model using the data collected during the test rollout.
After h timesteps in the environment, we have the following dataset: D = { ( s 1 , s 2 ) , . . .
, ( s h − 1 , s h ) } .
This allows us to learn a simple dynamics model f ψ : ( s t ) (cid:55)→ δ t =
s t +1 − s t , by minimizing the mean squared error L MSE ( ψ, D ) .
Notably, since we never actually plan with this model, it does not need to be as accurate as a typical dynamics model in MBRL.
Instead, it is crucial that the model learns quickly enough such that we can use it in a zero-shot evaluation. Thus, we choose to use a linear model for f . To show the effectiveness of this, in Fig. 4 we show the mean R-squared of linear models learned on the ﬂy during evaluation rollouts.
We observe that in less than 100 timesteps the linear model achieves high accuracy on the test data.
Subsequently, equipped with f ψ , we can approximate δ t , and predict the augmentation as (cid:98) z t = δ t / ˆ δ t .
We then provide the agent with (cid:98) z t to compute action a t .
The full procedure is shown in Algorithm 2.
E XPERIMENTS
In our experiments we aim to investigate the effectiveness of our approach for zero-shot dynamics generalization from a single ofﬂine dataset.
To assess this, we will answer a series of questions, beginning with a question on the necessity of our method:
Do we really need to develop methods speciﬁcally for dynamics generalization?
To answer this, we train MOPO using ofﬂine data from a single environment, and test it under changed dynamics.
We consider the HalfCheetah environment from the OpenAI Gym (6), using ofﬂine data from D4RL (11).
We train a MOPO agent using the mixed dataset, using our own implementation of the algorithm (but using the same hyperparameters as the original authors).
To test the trained policy, we vary both the mass of the agent and damping coefﬁcient by a multiplicative factor.
In this work we consider a grid of the following values for HalfCheetah: { 0 . 25 , 0 . 5 , 0 .
75 , 1 . 0 , 1 . 25 , 1 . 5 , 1 . 75 } and { 0 . 5 , 0 .
75 , 1 . 0 , 1 . 25 , 1 . 5 } for Walker2d, representing a signiﬁcant out-of-distribution shift.
[Figure]
The results (Fig. 5, left), show that MOPO performance is clearly impacted by changing dynamics.
We see in the central cell, that performance for our version of MOPO matches the author results (45), and in some cases we even see small gains (e.g. mass = 0 .
75 , damping = 1 . 0 ).
However, on the top left we see dramatically weaker performance, often below 1 k, indicating the robot is failing to achieve locomotion.
Before evaluating AugWM, we ﬁrst test whether training with the “correct” augmentation improves generalization performance.
In short, we ask:
Is augmenting dynamics even worthwhile?
To answer this, we train SAC for 1 × 10 5 steps and save the states visited in the ‘true’ environment.
We then use these starting states to train a policy using an ofﬂine MBPO 2 approach with AugWM.
However, instead of sampling z t ∼ Z , we provide the actual z = δ (cid:63) / ˆ δ as we have access to the ‘true’ and ‘modiﬁed’ environments; we refer to this as an oracle version of our method, and is designed to assess the viability of our approach.
Note that we do not augment the ‘true’ environment rewards.
We consider two baselines: a) ofﬂine MBPO in the ‘true’ environment; b) online SAC in the ‘modiﬁed’ environment. We train MBPO until convergence, and SAC for 1 × 10 5 steps.
As shown Fig.
5 ( Right ), when provided with the true z , AugWM outperforms both baselines.
The SAC result is surprising: the baseline agent was trained directly on the ‘modiﬁed’ environment for the same number of steps as the policy that generated our oracle starting states.
One explanation is the greater exploration induced by the ‘easier’ dynamics of the ‘true’ environment.
This validates our approach; if we augment the dynamics (cid:98) P from a model correctly, we can generalize to unseen dynamics.
Neither the starting states nor rewards need to be from the ‘modiﬁed’ environment.
Our next question is a simple one:
Which augmentation strategy is most effective?
To test this, we train as in Algorithm 1, without context, to isolate the effectiveness of the training process.
We use the HalfCheetah mixed dataset and train a MOPO agent, augmenting either both s and s (cid:48)
(RAD), just s (cid:48) (RANS) or just δ (DAS); the results are shown in Fig.
As we see, the RAD augmentation fails to improve dynamics generalization, actually leading to worse performance [Figure]
overall.
RANS does improve performance on unseen dynamics, as we are inﬂuencing the dynamics , not just the observation.
However, DAS clearly provides the strongest performance.
As a result [Figure] we use DAS for AugWM.
Our ﬁnal algorithm design question is as follows:
Does training with context improve performance?
[Figure]
To answer this question we return to the HalfCheetah mixed setting from Fig. 6, taking the policy trained with DAS.
We now train two additional agents: 1) Default Context: at train time the agent is provided with the DAS augmentation z as context, at test time it is provided with a vector of ones, 1 |S| ; 2)
Learned Context: trained as in 1) [Figure] but context is learned online using Algorithm 2.
The results are shown in in Fig.
We observe that training with context (orange) improves performance on average, while adapting the context on the ﬂy (green) leads to further gains (+80 on average).
These methods combine to produce our AugWM algorithm.
We are now ready for the ﬁnal question:
Can Augmented World Models improve zero-shot generalization?
[Table]
To answer this question we perform a rigorous analysis, using multiple benchmarks from the D4RL dataset (11).
Namely, we consider the random, medium, mixed and med-expert datasets for both
Walker2d and HalfCheetah.
In each setting, we compare AugWM against base MOPO on zero-shot performance, training entirely on the data provided, but not seeing the test environment until evaluation.
The results are shown as a change v.s. MOPO, averaged over one dimension in Fig. 10, and as a total return number averaged over both dimensions in Table 1.
For additional implementation details (e.g., hyperparameters) see
Appendix B. AugWM provides statistically signiﬁcant improvements in zero-shot performance v.s. MOPO in many cases, achieving successful policies where MOPO fails.
By now we have provided signiﬁcant evidence that AugWM can signiﬁcantly improve performance for HalfCheetah and Walker2d with varied mass and damping.
However, this is only a small subset of possible dynamics changes.
We next consider several signiﬁcantly harder settings.
We test increased dimensionality, using the Ant environment from MOPO (45), and also consider varied dynamics changes (Ant with crippled legs, HalfCheetah with varied physical properties from (19)).
We show the mean results over each of these factors of variation in Table 2, where once again AugWM provides a non-trivial improvement over a strong baseline.
For more details see Appendix B. [Table] Finally, we note that dynamics may change during an episode; consider a robot that suffers a motor fault, reducing the power delivered to its joints.
Evidently the underlying dynamics have been altered, and being robust to such changes when only training from a single dataset of ofﬂine experience is challenging.
To illustrate this, we perform a 1500 step rollout in the HalfCheetah environment, starting with ofﬂine dynamics (mass/damping = 1 ), before changing to mass = 0 .
75 , damping = {+
+}
.after 500 steps; performance is shown in Fig.8.
Observe that after 500 steps, MOPO performance is dramatically reduced.
This is because the agent continues to apply the same force and thus falls forward with lighter mass.
For our AugWM agent, performance initially drops, then when the new context is learned we achieve higher performance than before, making use of the lighter torso.
[Figure]
Discussion We believe that our experiments provide signiﬁcant support to the claim that training with AugWM improves zero-shot dynamics generalization.
In a broad set of commonly used datasets, and with a wide range of out-of-distribution dynamics, our algorithm learns good policies where a state-of-the-art baseline fails.
This is due to a number of novel contributions: 1) using dynamics augmentation rather than observation augmentation ; 2) training and testing with a context-based policy.
Regarding limitations, we note that training inside the WM with context generally takes longer to converge (Appendix B).
Furthermore, in more nonlinear settings such as the HalfCheetah modiﬁed body part setting we saw a reduced performance for the learned context. This could be because the dynamics changes are out of the distribution of DAS augmentations (violating Eqn. 1), or due to the difﬁculty of modeling the task with a linear model. We note that linear models have achieved success in planning (13) and meta learning (33), and are effective in our case due to their data efﬁciency, but can be replaced by more ﬂexible models to deal with different augmentations.
C ONCLUSION AND F UTURE W ORK
In this paper we propose Augmented World Models (AugWM), which we show sufﬁciently simulates changes in dynamics such that agents can generalize in a zero-shot manner.
We believe that we are the ﬁrst to propose this problem setting, and our results show a signiﬁcant improvement over existing state-of-the-art methods which ignore this problem.
A promising line of future work would be to meta-train a policy over AugWM such that it can quickly adapt to new dynamics in the few-shot setting.
There is evidence that data augmentation can improve robustness in meta-learning (37), and could extend to strong performance in out-of-distribution tasks.
We also wish to consider varying goals at test time, and other potential sources of non-stationarity which may impact policy performance.
It may also be possible to extend AugWM to pixel-based tasks, which have received a great deal of recent attention (17; 16).
We believe that our transition based augmentations will be applicable to a latent representation, as commonly used in state-of-the-art vision MBRL approaches.
Thus we think that extending our work to this setting, while a considerable feat of engineering, should not require signiﬁcant methodological changes.
A PPENDIX A A DDITIONAL E XPERIMENTS
Here we show the more granular experimental results, presented in condensed form in Table 1.
In
Fig.
10 we show the improvement vs. MOPO for each dimension of mass (green) or damping (blue), for HalfCheetah (top) and Walker2d (bottom)
[Figure] [Figure]
In this section we show the performance for Augmented World Models with different training ranges for the DAS augmentation ( z train in Table 4).
We train with adaptive context on the HalfCheetah mixed dataset, and present the results in Fig.
As we see, [0 .
75 , 1 . 25] and [0 . 5 , 1 . 5] perform the best.
Based on this, we use [0 . 5 , 1 . 5] for our experiments as we believe this helps us sample a wider set of dynamics, helping us generalize better across all environments and data sets.
B I MPLEMENTATION D ETAILS B. H YPERPARAMETERS
Our algorithm is based on MOPO (45) with values for the rollout length h and penalty coefﬁcient λ shown in Table 3.
AugWM speciﬁc hyperparameters are listed in Table 4.
For each evaluation rollout, we clear the buffer of stored true modiﬁed environment transitions to measure zero-shot performance.
We adapt using the context after a set number of steps, k , in the environment to train the linear model.
The two ranges used for the context z during training and test time are different.
At test time, the estimated context is clipped to remain within the given bounds.
[Figure]
[Table] [Table]
B. D4RL DATASET
We evaluate our method on D4RL (11) datasets based on the MuJoCo continuous control tasks (halfcheetah and walker2d).
The four dataset types we evaluate on are:
• random: roll out a randomly initialized policy for 1M steps.
• medium: partially train a policy using SAC, then roll it out for 1M steps.
• mixed: train a policy using SAC until a certain (environment-speciﬁc) performance threshold is reached, and take the replay buffer as the batch.
• medium-expert: combine 1M samples of rollouts from a fully-trained policy with another 1M samples of rollouts from a partially trained policy or a random policy.
This gives us a total of 8 experiments.
B. A NT E NVIRONMENT
For the Ant experiments, we follow the Ant Changed Direction approach in MOPO (45).
Since this ofﬂine dataset is not provided in the authors’ code, nor is it in the standard D4RL library (11), we were required to generate our own ofﬂine Ant dataset.
Since the authors’ did not outline certain details in their experiment, we found the following was required to match their performance with our codebase: 1) Training our SAC policy for 1 × 10 6 timesteps in the Ant environment provided by the authors’ code in (45); 2) relabelling each reward in the buffer using the new direction, without the living reward; 3) training a world model over this ofﬂine dataset; 4) training a policy in the world model, adding in living reward post-hoc; 5) evaluating the policy with the living reward.
B. H ALF C HEETAH M ODIFIED A GENT
We use the modiﬁed HalfCheetah environments from (19).
In each setting one body part of the agent is changed, from following set: { Foot, Leg, Thigh, Torso, Head } .
The body part can either be “Big” or “Small”, where Big bodyparts involve scaling the mass and width of the limb by 1.25 and Small bodyparts are scaled by 0.75.
In Table 2 we show the mean over each of these ﬁve body parts, for agents trained on each of the D4RL datasets, repeated for ﬁve seeds.