Domain Generalization for Robust Model-Based Offline Reinforcement Learning
Abstract
Existing offline reinforcement learning (RL) algorithms typically assume that training data is either: 1) generated by a known policy, or 2) of entirely unknown origin.
We consider multi-demonstrator offline RL, a middle ground where we know which demonstrators generated each dataset, but make no assumptions about the underlying policies of the demonstrators.
This is the most natural setting when collecting data from multiple human operators, yet remains unexplored.
Since different demonstrators induce different data distributions, we show that this can be naturally framed as a domain generalization problem, with each demonstrator corresponding to a different domain.
Specifically, we propose DomainInvariant Model-based Offline RL (DIMORL), where we apply Risk Extrapolation (REx) [15] to the process of learning dynamics and rewards models.
Our results show that models trained with REx exhibit improved domain generalization performance when compared with the natural baseline of pooling all demonstrators’ data.
We observe that the resulting models frequently enable the learning of superior policies in the offline model-based RL setting, can improve the stability of the policy learning process, and potentially enable increased exploration.
Introduction
In the standard online reinforcement learning (RL) paradigm, agents often require millions of interactions with the environment in order to learn an optimal policy for completing a task.
In many settings, however, the process of exploration in the real-world is undesirable, impractical or unsafe [17, 19].
It would therefore be preferable to enable learning from demonstrations , rather than direct interaction with the environment.
Datasets targeting activity recognition from ego-centric videos [8] highlight the prevalence and ease of collecting demonstrations in our setting. Demonstrations are provided by one or more demonstrators , each executing their own behavioural policy, π B –their method of achieving the task.
Various approaches to performing offline RL have been proposed [17, 19], however, to the best of our knowledge, none of these have exploited information regarding the collection of data from multiple demonstrators.
A fundamental challenge in off-policy RL is policy-induced distributional shift ; this is especially problematic in offline RL [17, 28].
Standard off-policy methods often fail when trained using offline data [10].
Value-based offline RL algorithms have addressed this by encouraging the learning agent to stay close to the behaviour policy that generated the training data. Such a constraint can (provably) limit performance, however Buckman et al.
[4], and other model-based approaches, which do not suffer from this weakness, have been found to outperform value-based methods [28, 14].
Rather than constraining the learning agent’s exploration, the aim of our work is to increase the robustness of environment models 1 to distributional shifts by applying Risk Extrapolation (REx) [15] during their training.
Our hope is that this will enable learning policies to explore more of the state-action space without incurring significant training instability.
REx assumes that training data from multiple domains (in our case, demonstrators) is available.
While the demonstrator that generated each training record is known, no knowledge about the demonstrator’s policy is required.
Even if distributional shifts more extreme than those observed at training time are encountered, REx aims to achieve similar risks on out-of-distribution domains by encouraging the training losses/risks across the domains present in the training data to be equal [15].
The primary contribution of our work is a practical algorithm for performing Domain-Invariant
Model-based Offline RL (DIMORL) using multi-demonstrator datasets.
We empirically show that environment models trained with REx exhibit improved average and worst-case performance on out-of-distribution datasets.
Furthermore, in a majority of cases, policies trained offline using these models attain higher average returns.
We demonstrate that environment models trained with REx can enable more stable offline model-based policy learning, and potentially support increased exploration.
We present hypotheses for the benefits observed, and share the supporting evidence gathered thus far.
Further experimentation and analysis are needed to draw firm conclusions.
Background and Related Work
Offline Reinforcement Learning In offline RL, an agent π is trained using only a fixed dataset D of transitions ( { s, a, r, s ′ } )
[10], or trajectories ( { s 1 , a 1 , r 1 , ..., s T , a T , r T } )
[1], without any further interaction with the environment [17, 19].
The data may also include meta-data about the policies used to collect the data.
Our aim is to learn the optimal policy, π ∗ ( a | s ) , that maximises the expected sum of discounted rewards π ∗
= arg max π E π
[ (cid:80) ∞ t {+=
+}
γ t r ( s t , a t )] , using discount factor γ ∈ (0 , 1] .
Distributional Shift While the transition distribution of a Markov-Decision Process (MDP) is independent of the policy being followed, the state and state-action visitation distributions induced by behavioural policy π B , d π B ( s ) and d π B ( s, a ) respectively, are not. The data collected for offline policy training is typically assumed to be composed of iid samples drawn from d π B ( s, a ) , and is likely to cover only a fraction of the total state-action space.
In order to learn optimal policies, we would like our learning algorithm to generalize to other areas of this space; departing from the support of the training data in order to learn good behaviours that are not exhibited in the static dataset [28, 27].
The policy being learned, π off , would therefore induce new state and state-action visitation distributions, d π off ( s ) and d π off ( s, a ) .
Errors made by the learned environment models, such as those arising from poor generalization performance under distributional shifts, can result in model exploitation : the policy being trained learns to take advantage of model errors when optimizing the reward, leading to poor performance in the real environment [5, 20, 13, 7, 17]. However, in more extreme cases, model exploitation can cause significant instability in training, preventing a policy from being learned [16, 18].
Offline Model-Based Reinforcement Learning Kidambi et al.
[14] and Yu et al.
[28] proposed the model-based offline RL algorithms MoREL and MOPO respectively.
Both of these methods are conservative in the face of model uncertainty.
However, they do not explicitly penalize deviations from the demonstrator’s distribution, thus avoiding performance limitations demonstrated by Buckman et al.
Both MoREL and MOPO work by modifying a learned MDP to penalize uncertain transitions.
MoREL’s pessimistic MDP adjusts P ( s ′ | s, a ) to transition to a low-reward terminal state when the dynamics of ( s, a ) are uncertain, as measured by disagreement within an ensemble of learned dynamics models [14].
In MOPO, Yu et al.
[28] instead penalize uncertain transitions directly in the reward function, and measure uncertainty as the maximum standard deviation across members of an ensemble.
MOPO extends MBPO [13], which utilizes learned environment models to generate short roll-outs of length h that are employed to update the learning policy using the soft-actor critic (SAC) policy gradient algorithm [11, 12].
Methods
Our method stems from the observation that different demonstrators in an offline RL setting correspond to different domains in a domain generalization setting.
We propose the Domain-Invariant
Model-based Offline RL (DIMORL) algorithm, which supplements model-based offline RL techniques by applying Risk Extrapolation (REx) [15] when training environment models.
3.1 Risk Extrapolation (REx)
The Risk Extrapolation (REx) domain generalization method aims to enforce strict equality of risks (i.e., losses) across training domains [15].
This helps guarantee that a model trained with REx will be invariant; i.e., it should also achieve (roughly) the same loss on test domains that are in the affine span of the training domains.
This is demonstrated in Fig.
1, which shows that an environment model trained with REx achieves more consistent performance when evaluating log-likelihoods across a range of datasets that were each generated by a different held-out demonstrator (i.e., demonstrators that did not contribute to the model’s training data). We use the simple Variance-REx (V-REx) algorithm, which penalizes the variance of the training risks:
[Equation]
where R e is the risk on the e -th domain, and β controls the strength of the variance regularizer.
The robustness of REx to multiple forms of distributional shift [15] make it a strong candidate for our investigations, however other domain generalization techniques could be explored [2, 21, 26, 23, 30].
3.2 Domain-Invariant Model-based Offline RL (DIMORL) Implementation
In this work, we implement a version of DIMORL that extends the MOPO algorithm [28] by applying REx during environment model training.
Experiments are performed using β ∈ { 0 , 5 , 10 , 20 } , λ ∈ { 0 , 1 , 5 } and {+h ∈ { 5 , 10 } , except where otherwise stated. When β = 
 the+} MOPO {+algorithm is recovered, and when additionally λ = 
 the MBPO algorithm is recovered, providing two baselines against which our results can be compared.+}
Ensembles of environment models, p θ,ϕ ( s t +1 , r t | s t , a t ) , are trained against multi-demonstrator datasets, D E =
(cid:83) Me =1 D e , where the outputs of model i parameterize a Gaussian distribution with diagonal covariance matrix: N ( s t +1 , r t ; µ iθ ( s, a ) , Σ iϕ ( s, a )) .
Individual risks, R e , are calculated for each demonstrator’s data, D e , and REx used to reduce the variance of the risks. Thus, we take advantage of knowing which demonstrator generated each record, but assume no knowledge of any demonstrator’s policy.
See Appendix A for further details. [Table] [Table] [Table]
Experiments
As a proxy for real-world demonstrators [Table] we trained a selection of policies online against the OpenAI
Gym MuJoCo Hopper, HalfCheetah and Walker2d environments [3, 25] using the soft actor-critic (SAC) algorithm [11, 12].
Subsets comprising five of these pseudo-demonstrators were used to generate three multi-demonstrator datasets for each environment: Novice , Mixed , and Expert .
See
Appendix B for details of the SAC pseudo-demonstrators and multi-demonstrator datasets produced.
4.1 Domain Generalization Performance of Environment Models
We found that environment models trained with REx achieved higher average and worst-case outof-distribution (OOD) log-likelihoods for each environment and multi-demonstrator dataset.
Across all environments, OOD performance was worst for the Experienced datasets, and was most greatly improved by the incorporation of REx to model training.
This can be seen for the HalfCheetah environment in Figure 1.
The results for the other environments can be found in Appendix C. The average and worst case performance generally continued to improve as the REx penalty coefficient was increased from 5 to 20.
While this might encourage the use of larger coefficients, we did not find the domain generalization performance of the models to be well correlated with the average returns of policies trained using them. This is discussed further in Appendix C.1.
4.2 Offline Agent Training
Policies learned using environment models trained with REx obtained the highest average return for two-thirds of multi-demonstrator datasets across the MuJoCo environments investigated.
Of these, the largest REx penalty coefficient investigated ( β = 20 ) yielded the highest return in all but one setting.
Table 1 provides an overview of the hyperparameter settings used to achieve the highest performing policies.
The full results are discussed in Appendix D.1.
We additionally found that environment models trained with REx were more robust to noisy initial state distributions; highlighting the benefits that improved domain generalization performance can bring to policy training (see Appendix F.1).
We further attempted to learn policies using roll-out starting locations sampled randomly from datasets not used to train the environment models, with the expectation that environment models trained with REx would exhibit improved robustness to such distributional shifts.
However, no good policies were learned under any settings investigated (see Appending D.2), likely due to the distributional shifts encountered being simply too extreme.
4.3 Increased Policy Training Stability
In an attempt to pinpoint the sources of the benefits brought by environment models trained with
REx, we analyzed the agent training performance and found that, for the HalfCheetah environment, the Q-functions of soft actor-critic policies were often less prone to becoming degenerate when using
REx penalty coefficients β ∈ { 10 , 20 } .
This is demonstrated in Fig. We present two hypotheses regarding the source of the Q-function instability, and why REx may be yielding improvements.
Degenerate Predictions: The magnitude of the evaluation returns in Figure 2 highlight that the learned reward models can make degenerate predictions.
This could lead to instability in Q-function training.
However, the returns shown were obtained over 1000 step episodes, whereas roll-outs of h ∈ { 5 , 10 } steps were typically generated for policy training.
The reward models would therefore need to make degenerate predictions within 5-10 steps to negatively impact Q-function learning.
We have observed cases of this (see Appendix E), however it does not appear to be a necessary condition for Q-function instability to occur.
Note, rather than the issue lying solely with the reward models, it is likely that errors in the dynamics models will also contribute by leading the agent into unnatural states.
Further analysis of the trajectories produced in the learned environment is necessary.
Environment Model Complexity: The learned environment model may be less well-behaved and more complex than the real environment.
For instance, a small change in the current state could have a large and unpredictable impact on the predicted next state.
If we assume the neural network used has sufficient capacity and is provided with sufficient data to capture the real environment, the function approximation component of the classic deadly triad [24] would theoretically not be an issue.
However, the environment model state representation may fail to capture all quirks in the model, resulting in function approximation continuing to be problematic.
4.4 Increased Exploration During Policy Training
We hypothesize that the improved domain generalization performance of environment models trained with REx may enable increased exploration of the state-action space during policy training, yielding more optimal policies. Our hope is that policy training will continue to remain stable, and model exploitation will be minimized.
To visualize the extent of exploration, we use PCA to project transition records into two-dimensions.
Initial experiments potentially indicate increased exploration (see Appendix F.2), however further work is necessary to ensure our observations cannot simply be explained by the fraction of the original data’s variance that is captured by the projection.
Conclusions
We have highlighted and exploited the opportunity to take advantage of knowledge regarding which demonstrators generated offline datasets, without assuming any information about each demonstrator’s policy.
We presented a practical algorithm for Domain-Invariant Model-based Offline RL (DIMORL), and have illustrated that this method may improve the stability of policy learning, enable increased exploration, and yield more optimal policies in many settings.
Further experimentation and analysis is needed to validate our findings, and to fully take advantage of the algorithm’s potential.
The flexibility of our approach enables the combination of many alternative domain generalization techniques and model-based RL algorithms, presenting exciting avenues for future investigation.
A DIMORL Implementation Details
The MOPO algorithm [28] and codebase 2 is used as the basis for our practical implementation of DIMORL.
Our modified version is outlined in Algorithm 1.
We train ensembles of environment models, p θ,ϕ ( s t +1 , r t | s t , a t ) , against multi-demonstrator datasets D E
= (cid:83) Me =1 D e , where the outputs of the models parameterise a Gaussian distribution with diagonal covariance matrix: N ( s t +1 , r t ; µ iθ ( s, a ) , Σ iϕ ( s, a )) .
Individual negative log-likelihood values, R e nll , are calculated for each demonstrator’s data, D e .
These are then used to calculate the MOPO V-REx loss given by Equation 2, where β is the REx penalty coefficient.
R wr is an ℓ 2 weight regularisation penalty, and R vb is a loss term relating to the learning of variance bounds.
Chua et al.
[6] highlight that outside of the training distribution the predicted variance can assume arbitrary values, and can both collapse to zero or explode to infinity (in contrast to models like GPs where variance values are better behaved).
Chua et al.
[6] found that bounding the output variance such that it cannot exceed the range seen in the training data was beneficial, and so include the R vb penalty.
The weight regularisation and variance bounding terms also appear in the loss function used by MOPO to train environment models [28].
[Equation]
{ p
[Equation]
ϕ ( s t +1 , r t | s t , a t ) =
N ( µ iθ ( s, a ) , Σ iϕ ( s, a )) }
Ni =1 with REx penalty coefficient β 2: Initialise policy π off and empty replay buffer D model ← ∅ 3: for epoch 1 , 2 , . . .
do 4: for 1 , 2 , . . .
, b (in parallel) do 5:
Sample state s 1 from D E for the initialisation of the roll-out 6: for j = 1 , 2 , . .
, h do 7:
Sample an action a j ∼ π ( s j ) 8:
Randomly pick dynamics ˆ T from { ˆ T i } Ni =1 and sample s j +1 , r j ∼ ˆ T ( s j , a j ) 9:
Compute ˜
r j = r j − λ max Ni =1 || Σ iϕ ( s j , a j ) || F 10:
Add sample ( s j , a j , ˜ r j , s j +1 ) to D model 11:
Drawing samples from D E (cid:83) D model , use SAC to update π off
Environment model training was split into two phases:
ERM: No REx penalty is used (i e., β = 0 ) until the original MOPO termination condition is reached: the loss on a held-out evaluation dataset does not decrease by more than 1 % for any model in the ensemble for five consecutive epochs [28].
REx: Training was then allowed to continue for the same number of epochs completed in the ERM phase, with the REx penalty coefficient now set to a user defined value.
We investigated REx penalty coefficients β ∈ { 0 , 5 , 10 , 20 } in our work.
B SAC Demonstrators and Multi-Demonstrator Datasets
To proxy for real-world demonstrators (whether humans, or pre-existing control policies), we train a selection of demonstrator policies, π e , using online RL algorithms and the OpenAI Gym MuJoCo
Environment
[Table]
Hopper, HalfCheetah and Walker2d environments.
The sole purpose of these policies is to generate individual demonstrator datasets, D e , for use in training environment models.
Each dataset comprises N e transition tuples of the form: ( s e , a e , s ′ e , r e , e ) .
The unique demonstrator identifier e is used to identify individual domains during the training of dynamics models with Risk Extrapolation (REx).
The identifier is simply a number [Table] and provides no information about the policy followed by the demonstrator.
Independently sampled evaluation datasets are also created for each demonstrator.
Collections of M individual datasets are combined to produce multi-demonstrator datasets: D E
= { ( s ei , a ei , s ′ ei , r ei , e ) N e i =1 } Me =1 .
A selection of soft actor-critic (SAC) [11, 12] policies are trained online. To increase the diversity of the demonstrators, we use two different implementations of the SAC algorithm:
Softlearning , the official SAC implementation [11, 12]; and D3RLPY , which implements a version of SAC with delayed policy updates [22]. Policies are trained using the default hyperparameters provided in their respective repositories. As shown in Table 2, snapshots are taken at regular intervals during training to emulate demonstrators (with distinct policies) of varying degrees of skill.
Table 3 shows the individual pseudo-demonstrator policies obtained.
To further increase data diversity, additional demonstrator datasets are created using a random policy, π N .
We select individual subsets of five pseudo-demonstrators (i.e., five distinct policies) to produce three multi-demonstrator datasets for each MuJoCo environment: Novice , Mixed , and Experienced .
As their names indicate, the Novice and Experienced datasets comprise transition records generated using pseudo-demonstrators trained online for minimal and extended number of training steps respectively, while Mixed uses a combination of each.
Each pseudo-demonstrator contributes 20, transition records, and so each multi-demonstrator dataset consists of a total of 100,000 records.
In future work, we would look to train pseudo-demonstrators using a greater variety of online RL algorithms to further increase the diversity of policies used.
[Table]
C Environment Model OOD Performance
Figure 3 shows the average and worst-case log-likelihoods and mean-squared errors (MSEs) obtained for environment models trained on all MuJoCo environments and multi-demonstrator datasets.
The performance metrics were calculated using datasets generated by demonstrators that did not contribute to the multi-demonstrator dataset used to train the model.
Across all environments and multi-demonstrator datasets, environment models trained using REx achieved improved average and worst-case log-likelihoods.
The results were more mixed for the MSE values–small increases in
MSE were sometimes observed, as can be seen for the Walker2d Experienced dataset in Figure 3f.
The discrepancy likely arises given it was the variance of log-likelihoods across training domains that was minimised during model training.
Future experiments could evaluate the impact of updating the training procedure to minimise the variance of the training MSEs, or both the log-likelihoods and MSEs simultaneously.
Increasing the REx penalty coefficient leads to higher log-likelihood values in the majority of cases.
One exception is the peak in worst-case performance that all Hopper multi-demonstrator datasets exhibit when β = 10 .
This indicates that the REx penalty coefficient is still a hyperparameter that needs to be appropriately tuned.
C.1 Correlation between Environment Model OOD Performance and Policy Returns
As demonstrated by Figure 4, no correlation is observed between the average OOD log-likelihood of environment models and the average evaluation return of the policies trained using them.
It should, however, be noted that the distribution of roll-out starting locations differs across the experiments (given that starting locations are drawn from the dataset used to train the environment model), which will have had an influence on the observed results.
D Policy Training Results D.1 ID Roll-out Starting Locations
The results of experiments run against multi-demonstrator datasets generated using the HalfCheetah,
Hopper and Walker2d environments are shown in Table 4, Table 5 and Table 6 respectively.
In all of these experiments, during policy training, roll-out starting locations were drawn from the same dataset used to train the environment model.
Across all environments, a policy trained offline with DIMORL obtained a higher average return for the Novice multi-demonstrator dataset than any of the individual demonstrators that were used to [Figure] [Figure] generate the dataset.
The same consistency was not observed for the Mixed or Experienced datasets.
This reaffirms the importance of action diversity in offline model-based methods, whereas modelfree techniques are potentially better suited to datasets collected from less-diverse, more experienced demonstrators, as was highlighted by Yu et al.
We also note the lack of consistency in the results regarding the optimal MOPO penalty coefficient.
Yu et al.
[29] highlight the poor calibration between the MOPO penalty and actual model error, which is likely at least partly responsible for the inconsistency observed. [Table]
D.2 OOD Roll-out Starting Locations
We refer to roll-out starting locations as being OOD if they are sampled from any dataset other than the one used to train the environment model that is used in policy training.
The results of experiments run against multi-demonstrator datasets generated using the HalfCheetah and Hopper environments are shown in Table 7 and Table 8 respectively.
Roll-out starting locations for these experiments were drawn from the following datasets
RAND:
Transitions were generated by taking random actions in the environment.
D4RL-MR: The D4RL medium-replay datasets for each of the environments [9].
Experiments were not run for environment models trained on the HalfCheetah Experienced multidemonstrator dataset given that (when using MOPO penalty coefficient λ = 0 ) no policies were learned when drawing starting locations from the same dataset used to train the models.
If policies cannot be learned using in-distribution starting locations, then we do not expect to be able to learn them using OOD starting locations.
Further, we are yet to run similar experiments for the Walker2D multi-demonstrator datasets, but we anticipate similar results to those seen for the HalfCheetah and
Walker2D datasets (i.e., an inability to learn a policy).
Additionally, experiments are still to be run for MOPO penalty coefficients λ > 0 .
E Degenerate Reward Predictions
Figure 5 demonstrates that abnormally large reward predictions can occur within the first 10 steps of episodes generated using environment models trained without REx.
Out of the 60 episodes shown
[Table]
[Table]
[Table]
[Table]
in Figure 5, six contained reward predictions with an absolute value exceeding one million.
If these episodes are excluded then the largest absolute reward prediction was 12.29.
No episodes generated absolute reward predictions greater than 12.29 in the first 5 steps, while 2 episodes did within the first 10 steps.
Note that our aim in this experiment was merely to show the possibility of abnormally large reward predictions in the early stages of episodes, not prove that the problem is common/widespread.
F Additional Experiments F.1 Noisy Roll-out Starting Locations
To further assess the potential benefits of environment models trained with REx, we analysed the impact of adding noise to the roll-out starting locations used during offline policy training.
Note that the environment models were trained on the original multi-demonstrator datasets (which have no noise), and starting locations were still sampled from the same datasets used to train the environment models.
The only change was that noise with standard deviation σ ∈ { 0 . 00 , 0 . 01 , 0 .
[Table] was added to the datasets before starting locations were sampled from them.
Figure 6 shows that environment models trained with REx yielded policies with higher average returns across all noise levels for both the HalfCheetah and Hopper environments.
For the Walker2D environment, environment models trained with REx penalty coefficients β ∈ { 10 , 20 } yielded policies with higher returns.
For both HalfCheetah and Hopper environments, the highest policy evaluation return was obtained when σ = 0 .
01 –suggesting that a small amount of noise was beneficial.
We hypothesise that the distribution of noisy starting locations is more diverse, enabling greater exploration.
In the case of the Hopper environment, it was the environment model trained without REx that obtained the highest average return for σ = 0 .
01 , whereas the returns for the equivalent HalfCheetah and Walker2D environment models decreased significantly.
This may be thanks to the greater simplicity of the Hopper environment. [Figure]
F.2 Analysis of Training Exploration using PCA
In an attempt to determine whether environment models trained with REx supported increased exploration of the state-action space during policy training, we investigated the use of PCA to project transition records into two-dimensions, such that they could be visualised.
For experiments using the Gym MuJoCo HalfCheetah environment, we sampled 100,000 transition records from the model pool every 100,000 training steps, excluding the zeroth and final training steps.
Figure 7 shows the projections for a set of six policy training experiments, using the Mixed dataset and β ∈ { 0 , 10 } .
The projection matrix was trained on the complete collection of sampled transition records for the six experiments.
Note that, unlike all other experiments presented in this paper, the environment models used in these experiments received one additional epoch of training prior to being used to learn a policy.
This is the default behaviour of the MOPO codebase when utilising pre-trained environment models [28].
The average return for the experiments that did not use an environment model trained with REx was -420 ± 50, while for the REx experiments it was 8128 ± 1053.
The projected records appear to indicate that increased exploration of the state-action space took place when using an environment model trained with REx.
It is important to recognise, however, that the projections for REx experiments have higher explained variance, which may partially or fully explain the observed difference.
The results will of course be sensitive to the datasets used to learn the projection matrix.
We believe that further investigation of this method is warranted.
training appears to indicate that increased exploration of the state-action space took place when using an environment model trained with REx ( β = 10 ). It is important to recognize, however, that the projections for the REx experiments have higher explained variance, which may partially or fully explain the observed differences. The projection matrix was trained on the complete collection of sampled transition records for the six experiments shown.
MOPO penalty coefficient λ = 5 and roll-out length h = 10 were used.