COMBO: Conservative Offline Model-Based Policy Optimization

1 Introduction . Offline reinforcement learning ( offline RL ) [ 30 , 34 ] refers to the setting where policies are trained using static , previously collected datasets . This presents an attractive paradigm for data reuse and safe policy learning in many applications , such as healthcare [ 62 ] , autonomous driving [ 65 ] , robotics [ 25 , 48 ] , and personalized recommendation systems [ 59 ] . Recent studies have observed that RL algorithms originally developed for the online or interactive paradigm perform poorly in the offline case [ 14 , 28 , 26 ] . This is primarily attributed to the distribution shift that arises over the course of learning between the offline dataset and the learned policy . Thus , development of algorithms specialized for offline RL is of paramount importance to benefit from the offline data available in aformentioned applications . In this work , we develop a principled model-based offline RL algorithm that matches or exceeds the performance of prior offline RL algorithms in benchmark tasks . A major paradigm for algorithm design in offline RL is to incorporate conservatism or regularization into online RL algorithms . Model-free offline RL algorithms [ 15 , 28 , 63 , 21 , 29 , 27 ] directly incorporate conservatism into the policy or value function training and do not require learning a dynamics model . However , model-free algorithms learn only on the states in the offline dataset , which can lead to overly conservative algorithms . In contrast , model-based algorithms [ 26 , 67 ] learn a pessimistic dynamics model , which in turn induces a conservative estimate of the value function . By generating and training on additional synthetic data , model-based algorithms have the potential for 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . broader generalization and solving new tasks using the offline dataset [ 67 ] . However , these methods rely on some sort of strong assumption about uncertainty estimation , typically assuming access to a model error oracle that can estimate upper bounds on model error for any state-action tuple . In practice , such methods use more heuristic uncertainty estimation methods , which can be difficult or unreliable for complex datasets or deep network models . It then remains an open question as to whether we can formulate principled model-based offline RL algorithms with concrete theoretical guarantees on performance without assuming access to an uncertainty or model error oracle . In this work , we propose precisely such a method , by eschewing direct uncertainty estimation , which we argue is not necessary for offline RL . Our main contribution is the development of conservative offline model-based policy optimization ( COMBO ) , a new model-based algorithm for offline RL . COMBO learns a dynamics model using the offline dataset . Subsequently , it employs an actor-critic method where the value function is learned using both the offline dataset as well as synthetically generated data from the model , similar to Dyna [ 57 ] and a number of recent methods [ 20 , 67 , 7 , 48 ] . However , in contrast to Dyna , COMBO learns a conservative critic function by penalizing the value function in state-action tuples that are not in the support of the offline dataset , obtained by simulating the learned model . We theoretically show that for any policy , the Q-function learned by COMBO is a lower-bound on the true Q-function . While the approach of optimizing a performance lower-bound is similar in spirit to prior model-based algorithms [ 26 , 67 ] , COMBO crucially does not assume access to a model error or uncertainty oracle . In addition , we show theoretically that the Q-function learned by COMBO is less conservative than model-free counterparts such as CQL [ 29 ] , and quantify conditions under which the this lower bound is tighter than the one derived in CQL . This is illustrated through an example in Figure 1 . Following prior works [ 31 ] , we show that COMBO enjoys a safe policy improvement guarantee . By interpolating model-free and model-based components , this guarantee can utilize the best of both guarantees in certain cases . Finally , in our experiments , we find that COMBO achieves the best performance on tasks that require out-of-distribution generalization and outperforms previous latent-space offline model-based RL methods on image-based robotic manipulation benchmarks . We also test COMBO on commonly studied benchmarks for offline RL and find that COMBO generally performs well on the benchmarks , achieving the highest score in 9 out of 12 MuJoCo domains from the D4RL [ 12 ] benchmark suite . 2 Preliminaries . Markov Decision Processes and Offline RL . We study RL in the framework of Markov decision processes ( MDPs ) specified by the tupleM = ( S , A , T , r , µ0 , γ ) . S , A denote the state and action spaces . T ( s′|s , a ) and r ( s , a ) ∈ [ −Rmax , Rmax ] represent the dynamics and reward function respectively . µ0 ( s ) denotes the initial state distribution , and γ ∈ ( 0 , 1 ) denotes the discount factor . We denote the discounted state visitation distribution of a policy π using dπM ( s ) : = ( 1− γ ) ∑∞ t=0 γ tP ( st = s|π ) , where P ( st = s|π ) is the probability of reaching state s at time t by rolling out π in M. Similarly , we denote the state-action visitation distribution with dπM ( s , a ) : = d π M ( s ) π ( a|s ) . The goal of RL is to learn a policy that maximizes the return , or long term cumulative rewards : maxπ J ( M , π ) : = 11−γE ( s , a ) ∼dπM ( s , a ) [ r ( s , a ) ] . Offline RL is the setting where we have access only to a fixed dataset D = { ( s , a , r , s′ ) } , which consists of transition tuples from trajectories collected using a behavior policy πβ . In other words , the dataset D is sampled from dπβ ( s , a ) : = dπβ ( s ) πβ ( a|s ) . We defineM as the empirical MDP induced by the dataset D and d ( s , a ) as sampled-based version of dπβ ( s , a ) . In the offline setting , the goal is to find the best possible policy using the fixed offline dataset . Model-Free Offline RL Algorithms . One class of approaches for solving MDPs involves the use of dynamic programming and actor-critic schemes [ 56 , 5 ] , which do not explicitly require the learning of a dynamics model . To capture the long term behavior of a policy without a model , we define the action value function as Qπ ( s , a ) : = E [ ∑∞ t=0 γ t r ( st , at ) | s0 = s , a0 = a ] , where future actions are sampled from π ( ·|s ) and state transitions happen according to the MDP dynamics . Consider the following Bellman operator : BπQ ( s , a ) : = r ( s , a ) + γE s′∼T ( ·|s , a ) , a′∼π ( ·|s′ ) [ Q ( s′ , a′ ) ] , and its sample based counterpart : B̂πQ ( s , a ) : = r ( s , a ) + γQ ( s′ , a′ ) , associated with a single transition ( s , a , s′ ) and a′ ∼ π ( ·|s′ ) . The action-value function satisfies the Bellman consistency criterion given by BπQπ ( s , a ) = Qπ ( s , a ) ∀ ( s , a ) . When given an offline dataset D , standard approximate dynamic programming ( ADP ) and actor-critic methods use this criterion to alternate between policy evaluation [ 40 ] and policy improvement . A number of prior works have observed that such a direct extension of ADP and actor-critic schemes to offline RL leads to poor results due to distribution shift over the course of learning and over-estimation bias in the Q function [ 14 , 28 , 63 ] . To address these drawbacks , prior works have proposed a number of modifications aimed towards regularizing the policy or value function ( see Section 6 ) . In this work , we primarily focus on CQL [ 29 ] , which alternates between : Policy Evaluation : The Q function associated with the current policy π is approximated conservatively by repeating the following optimization : Qk+1←argmin Q β ( Es∼D , a∼µ ( ·|s ) [ Q ( s , a ) ] −Es , a∼D [ Q ( s , a ) ] ) + 1 2 Es , a , s′∼D [ ( Q ( s , a ) −B̂πQk ( s , a ) ) 2 ] , ( 1 ) where µ ( ·|s ) is a wide sampling distribution such as the uniform distribution over action bounds . CQL effectively penalizes the Q function at states in the dataset for actions not observed in the dataset . This enables a conservative estimation of the value function for any policy [ 29 ] , mitigating the challenges of over-estimation bias and distribution shift . Policy Improvement : After approximating the Q function as Q̂π , the policy is improved as π ← argmaxπ′Es∼D , a∼π′ ( ·|s ) [ Q̂π ( s , a ) ] . Actor-critic methods with parameterized policies and Q functions approximate argmax and argmin in above equations with a few gradient descent steps . Model-Based Offline RL Algorithms . A second class of algorithms for solving MDPs involve the learning of the dynamics function , and using the learned model to aid policy search . Using the given dataset D , a dynamics model T̂ is typically trained using maximum likelihood estimation as : minT̂ E ( s , a , s′ ) ∼D [ log T̂ ( s′|s , a ) ] . A reward model r̂ ( s , a ) can also be learned similarly if it is un- known . Once a model has been learned , we can construct the learned MDP M̂ = ( S , A , T̂ , r̂ , µ0 , γ ) , which has the same state and action spaces , but uses the learned dynamics and reward function . Subsequently , any policy learning or planning algorithm can be used to recover the optimal policy in the model as π̂ = argmaxπ J ( M̂ , π ) . This straightforward approach is known to fail in the offline RL setting , both in theory and practice , due to distribution shift and model-bias [ 51 , 26 ] . In order to overcome these challenges , offline model-based algorithms like MOReL [ 26 ] and MOPO [ 67 ] use uncertainty quantification to construct a lower bound for policy performance and optimize this lower bound by assuming a model error oracle u ( s , a ) . By using an uncertainty estimation algorithm like bootstrap ensembles [ 43 , 4 , 37 ] , we can estimate u ( s , a ) . By constructing and optimizing such a lower bound , offline model-based RL algorithms avoid the aforementioned pitfalls like model-bias and distribution shift . While any RL or planning algorithm can be used to learn the optimal policy for M̂ , we focus specifically on MBPO [ 20 , 57 ] which was used in MOPO . MBPO follows the standard structure of actor-critic algorithms , but in each iteration uses an augmented dataset D ∪Dmodel for policy evaluation . Here , D is the offline dataset and Dmodel is a dataset obtained by simulating the current policy using the learned dynamics model . Specifically , at each iteration , MBPO performs k-step rollouts using T̂ starting from state s ∈ D with a particular rollout policy µ ( a|s ) , adds the model-generated data to Dmodel , and optimizes the policy with a batch of data sampled from D ∪Dmodel where each datapoint in the batch is drawn from D with probability f ∈ [ 0 , 1 ] and Dmodel with probability 1− f .

Summary: The paper presents a new approach for model-based offline RL. Here, the model is not used to estimate the return directly using roll-outs, but to generate synthetic data that is then used to train the Q-function. In a way, the method is a combination of CQL and MOPO. What is new is the idea of using this synthetic data to, penalize extrapolation instead of using uncertainty estimation. This modification performs much better on the selected benchmarks. As benchmarks a selection of D4RL benchmarks is used (3 benchmarks) and two vision based benchmarks.