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: The authors present an empirical study of several uncertainty quantification heuristics applied to model learning in offline model-based reinforcement learning. Specifically, they consider the basic architecture of MOPO, in which an uncertainty based state-action penalty function is applied on top of the standard reward to construct a pessimistic MDP. Within this set up, the authors perform a empirical study of several different uncertainty penalties, exploring their correlation to the prediction error of the model, as well as in terms of their ability to detect individual transitions with high-percentile prediction errors. Finally, the authors perform Bayesian optimization over the choice of the uncertainty penalty as well as other hyperparameters such as number of ensemble elements, planning horizon, penalty weights, etc. They present these results showing that the optimal choice of penalty and penalty weight can vary significantly, not only between environments, but also within an environment as a function of offline dataset.