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: This paper introduces a Conservative Offline Model-Based policy Optimization RL algorithm, called COMBO, that learns a pessimistic model by enabling a lower bound optimization on the policy performance without requiring uncertainty quantification. COMBO employs an actor-critic method to learn the value function on both offline and synthetic datasets. The idea is to penalize Q-values learned from state-action pairs that are out-of-support of the offline dataset.