Revisiting Design Choices in Offline Model Based Reinforcement Learning

1 INTRODUCTION . In offline ( or batch ) reinforcement learning ( RL ) ( Ernst et al. , 2005 ; Levine et al. , 2020 ) , the goal is to leverage offline datasets of transitions in an environment to train a policy that transfers to an online task . This could have vast implications for using RL in real-world settings , as agents can make use of ever-increasing amounts of data without the need for an accurate simulator , while also avoiding expensive and potentially even unsafe exploration in the environment . Model-based reinforcement learning ( MBRL ) has recently shown promise in this paradigm , obtaining state-of-the-art performance on offline RL benchmarks ( Kidambi et al. , 2020 ; Yu et al. , 2021 ) , improving upon powerful model-free approaches ( e.g . Kumar et al . ( 2020 ) ) . MBRL works by training a dynamics model from the offline data , then optimizing a policy using imaginary rollouts from the model . This allows the agent to learn from on-policy experience , as the model is agnostic to the policy used to generate data , making it possible to achieve high returns using data collected from even a random policy . Furthermore , recent work has demonstrated the utility of world models beyond maximizing return , such as generalizing to unseen variations in an environment ( Ball et al. , 2021 ) , transferring to new tasks ( Yu et al. , 2020 ) , and learning with safety constraints ( Argenson & Dulac-Arnold , 2021 ) . Therefore , the case for MBRL in offline RL is clear : not only does it represent state-of-the-art in terms of performance , but it also provides the opportunity to maximize the signal in the offline data to generalize onto tasks beyond those encoded by the behavior policy . This is crucial for offline RL to be useful for real-world tasks ( Dulac-Arnold et al. , 2021 ) , where there will inevitably be differences between the data and desired task . However , a common failure mode of MBRL is when policies exploit the model in parts of the state-action space where the model is inaccurate . Thus , naïve application of MBRL to offline data can result in suboptimal performance . To prevent this , concurrent works ( Yu et al. , 2020 ; Kidambi et al. , 2020 ) have approached the problem by training a policy in a pessimistic MDP ( P-MDP ) . The P-MDP lower bounds the true MDP , and discourages the policy from regions where there is large discrepancy between the true and learned dynamics ; this often provides a theoretical guarantee of improvement over cloning the behavior policy that generated the offline data . This is made practically possible by adding a penalty correlated with the uncertainty in the dynamics model . However , while these recent successes are similar in principle , in practice they differ in a series of design choices . First and foremost , they make use of different heuristics to measure model uncertainty , in some cases deviating from simpler metrics which are more consistent with the theory . In this paper we conduct a rigorous investigation into a series of these design choices . We begin by focusing on the choice of uncertainty metric , comparing both recent state-of-the-art offline approaches ( Kidambi et al. , 2020 ; Yu et al. , 2020 ; Rafailov et al. , 2020 ) with additional metrics used in the online setting ( Ball et al. , 2020 ; Pan et al. , 2020 ; Cowen-Rivers et al. , 2020 ) . We also explore the interaction with a series of other hyperparameters , such as the number of models and imaginary rollout length . Interestingly , the relationship between these variables and model uncertainty varies significantly depending on the choice of uncertainty penalty . Furthermore , we compare these uncertainty heuristics under new evaluation protocols that , for the first time , capture the specific covariate shift induced by model-based RL . This allows us to assess calibration to model exploitation in MBRL , observing that some existing penalties are surprisingly successful at capturing the errors in predicted dynamics , as seen in Fig . 1a ( see App . D for details ) . Then , using the insights gained from sections 4 and 5 , we then achieve a 43 % gain over a previously grid-searched method by using a single hyperparameter value across all environments . To further test these insights , we then fine-tune over our identified key variables , modeled jointly using a powerful Bayesian Optimization algorithm ( Wan et al. , 2021 ) . We find that , when properly tuned , the simpler uncertainty measures can provide state-of-the-art results in continuous control offline benchmarks , and that the chosen optimal hyperparameters continue to align with our analysis . Finally , we rigorously confirm the aggregate statistical significance of our results using the RLiable framework ( Agarwal et al. , 2021 ) in Fig . 1b , showing that the improvements over existing methods are significant ( see App . H for details ) . This work is intended to benefit both researchers and practitioners in offline RL . Our main findings include : • Longer horizon rollouts with larger penalties can improve existing methods . Contrary to common wisdom , conducting significantly longer rollouts inside the model , coupled with larger uncertainty penalties , typically improves performance . • Penalties that use canonical forms of uncertainty estimation achieve better correlation with OOD measures . The deep ensembles approach of Lakshminarayanan et al . ( 2017 ) often outperforms the penalty from state-of-the-art methods ( Yu et al. , 2020 ; Kidambi et al. , 2020 ) . We observe that the ensemble standard deviation is statistically strikingly similar to that used in Kidambi et al . ( 2020 ) , but has improved correlation and scaling behavior . • Uncertainty is more correlated with dynamics error than distribution shift . We find that successful penalties measure the discrepancy in dynamics , and can in fact assign high certainty to data far away from the offline data . 2 RELATED WORK . Two recent works concurrently demonstrated the effectiveness of model-based reinforcement learning ( MBRL ) in the offline setting . MOPO ( Yu et al. , 2020 ) follows the successful online RL algorithm MBPO ( Janner et al. , 2019 ) but trains inside a conservative MDP , penalizing the reward based on the maximum aleatoric uncertainty over the ensemble members . MOReL ( Kidambi et al. , 2020 ) achieves even stronger performance , penalizing the rewards by a penalty based on the maximum pair-wise difference in ensemble member predictions . For pixel-based tasks , LOMPO ( Rafailov et al. , 2020 ) also proposed a novel penalty , using the variance of ensemble log-likelihoods . Outside the offline setting , probabilistic dynamics models leveraging uncertainty have underpinned a series of successes ( Chua et al. , 2018 ; Kurutach et al. , 2018 ; Buckman et al. , 2018 ; Pan et al. , 2020 ; Pacchiano et al. , 2021 ) . Uncertainty can also be measured in MBRL without the use of neural networks ( Deisenroth & Rasmussen , 2011 ) , although these methods tend to be harder to scale and thus lack widespread use . Effective hyperparameter selection in RL has been shown to be crucial to the success of popular algorithms ( Engstrom et al. , 2020 ; Andrychowicz et al. , 2021 ) . This becomes even more challenging in MBRL with additional hyperparameters/design-choices for the dynamics model . Recent work has shown that carefully optimizing these hyperparameters for online MBRL can significantly improve performance , with the tuned agent breaking the MuJoCo simulator ( Zhang et al. , 2021 ) . In contrast , we focus on the offline setting , and investigate parameters specifically related to uncertainty estimation . Previous work studied the impact of hyperparameters in offline RL ( Paine et al. , 2020 ) , finding offline RL algorithms to be brittle to hyperparameter choices . However , unlike our work they only consider model-free approaches , whereas we specifically investigate model-based offline algorithms . Abbas et al . ( 2020 ) investigates the impact of different uncertainty estimation methods in online MBRL ; they too find penalizing with combined aleatoric and epistemic uncertainty improves performance . Our work also relates to the rich literature on deep ensembles ( Lakshminarayanan et al. , 2017 ) , which train multiple deep neural networks with different initializations and dataset orderings , and generally outperform variational Bayes methods ( Mackay , 1992 ; Blundell et al. , 2015 ) . Achieving effective calibration with neural networks is notoriously difficult ( Guo et al. , 2017 ; Kuleshov et al. , 2018 ; Maddox et al. , 2019 ) , and furthermore we require calibration under co-variate shift ( Ovadia et al. , 2019 ) , as the policy learned in the model will likely deviate from the behavior policy that generated the offline data . Recent work has highlighted this issue in offline RL ( Kumar et al. , 2020 ; Yu et al. , 2021 ) and has reported superior performance when eschewing model uncertainty entirely , and instead performing “ conservative '' Q-updates . However , it is unclear if this improvement is due to poor uncertainty calibration , implementation details , or a limitation in the pessimistic-MDP formulation . 3 BACKGROUND . All of the methods we investigate in this paper model the environment as a Markov Decision Process ( MDP ) , defined as a tuple M = ( S , A , P , R , ρ0 , γ ) , where S and A denote the state and action spaces respectively , P ( s′|s , a ) the transition dynamics , R ( s , a ) the reward function , ρ0 the initial state distribution , and γ ∈ ( 0 , 1 ) the discount factor . The goal is to optimize a policy π ( a|s ) that maximizes the expected discounted return Eπ , P , ρ0 [ ∑∞ t=0 γ tR ( st , at ) ] . In offline RL , the policy is not deployed in the environment until test time . Instead , the algorithm only has access to a static dataset Denv = { ( sj , aj , rj , sj+1 ) } Jj=1 , collected by one or more behavioral policies πb . Following the notation in Yu et al . ( 2020 ) we refer to the distribution from which Denv was sampled as the behavioral distribution . The canonical approach in offline MBRL is to train an ensemble of N probabilistic dynamics models ( Nix & Weigend , 1994 ) . These usually learn to predict both the next state st+1 and reward rt from a state-action pair , and are trained on Denv using supervised learning . Concretely , each of the N models output a Gaussian P̂ iφ ( st+1 , rt|st , at ) = N ( µiφ ( st , at ) , Σiφ ( st , at ) ) parameterized by φ . The resulting learned dynamics model P̂ and reward model R̂ define a model MDP M̂ = ( S , A , P̂ , R̂ , ρ0 , γ ) . To train the policy , we use k-step rollouts inside M̂ to generate trajectories ( Sutton , 1991 ) . To prevent policy exploitation in a model , a pessimistic MDP ( P-MDP ) is constructed by lower bounding the true-expected return , ηM ( π ) , using some error between the true and estimated models . For instance , in Yu et al . ( 2020 ) , the authors show that a lower bound on the return can be established by penalizing the reward by a measure that corresponds to estimated model error : ηM ( π ) ≥ E ( s , a ) ∼ρπ P̂ [ R ( s , a ) − γ|Gπ M̂ ( s , a ) | ] ( 1 ) where ρπ P̂ represents transitioning under the dynamics model P̂ and policy π . Several potential choices for |Gπ M̂ ( s , a ) | are proposed , including an upper bound based on the total variation distance between the learned and true dynamics . However , for their practical algorithm , the authors elect to use a heuristic based on impressive empirical results . Concurrent to MOPO , MOReL ( Kidambi et al. , 2020 ) in theory constructs a P-MDP by augmenting a standard MDP with a negative valued absorbing state that is transitioned to when total variation distance between true and learned dynamics is exceeded . They show that a policy learned in this P-MDP exceeds simple behavior cloning . However , while dynamics-based total variation distance has desirable theoretical properties , the practical algorithm relies on another heuristic to approximate this quantity . This motivates the study of penalties used , as well as other under-used candidates , and their overall effectiveness .

Summary: Offline RL has seen an explosion of interest in the past year or so, with model-based RL methods being the state of the art (or close to it) on commonly studied benchmark tasks (like D4RL and RL Unplugged). This paper takes a retrospective view on model-based offline RL to benchmark/evaluate different design choices in recent papers. Based on experimental evaluations, the paper makes a number of interesting observations and suggestions for improvements such as longer rollout horizons and penalties more closer to OOD detection rather than uncertainty estimation, among others.