SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Constrained Policy Optimization with Explicit Behavior Density for Offline Reinforcement Learning
Jing Zhang
jzhanggy@connect.ust.hk
Kuaishou Technologies
zhangchi08@kuaishou.com
Wenjia Wang*
HKUST (GZ) and HKUST
wenjiawang@ust.hk
Bing-Yi Jing*
SUSTech
jingby@sustech.edu.cn
Due to the inability to interact with the environment, offline reinforcement learning (RL) methods face the challenge of estimating the Out-of-Distribution (OOD) points. Existing methods for addressing this issue either control policy to exclude the OOD action or make the  $Q$  function pessimistic. However, these methods can be overly conservative or fail to identify OOD areas accurately. To overcome this problem, we propose a Constrained Policy optimization with Explicit Behavior density (CPED) method that utilizes a flow-GAN model to explicitly estimate the density of behavior policy. By estimating the explicit density, CPED can accurately identify the safe region and enable optimization within the region, resulting in less conservative learning policies. We further provide theoretical results for both the flow-GAN estimator and performance guarantee for CPED by showing that CPED can find the optimal  $Q$ -function value. Empirically, CPED outperforms existing alternatives on various standard offline reinforcement learning tasks, yielding higher expected returns.
As a form of active learning, reinforcement learning (RL) has achieved great empirical success in both simulation tasks and industrial applications [1, 2, 3, 4, 5, 6, 7]. The great success of RL is largely due to sufficient information exploration and the online training paradigm that the agent has plenty of interactions with the environment. However, the merits of RL methods are limited when it is difficult or costly to sequentially interact with the environment. In many real-world scenarios such as driverless or intelligent diagnostics [8], online interaction is unacceptable. Thus, vanilla RL methods may possibly fail under such circumstances.
One recently developed approach to mitigating the constrained online-interaction is offline RL [8, 9, 10]. Motivated by various data-driven techniques, the offline RL leverages prior experiences to train a policy without interacting with the environment. In offline RL, a critical challenge is distribution shift (also called "extrapolation error" in literature). Specifically, vanilla RL methods
*Corresponding authors: Bing-Yi Jing and Wenjia Wang.
using the Bellman equation and the  $Q$ -function approximation probably fail to deliver good estimates for Out-of-Distribution (OOD) points. To this end, two types of offline RL solutions have been proposed. One type is the  $Q$ -function constraint. By adding constraints on the  $Q$ -function, this type of solution provides a pessimistic estimation of the  $Q$ -function and prevents the  $Q$ -value from being too large. The constraints include ensembling multiple  $Q$ -functions [11, 12], penalizing  $Q$ -function with high uncertainty [13, 14], among others. Recently, implicit Q-learning [15] and mildly conservative Q-learning [16] are suggested. However, the uncertainty of the  $Q$ -function or a tight lower bound of  $Q$ -value is difficult to estimate. Thus the performance is not guaranteed neither theoretically nor practically.
Another type of solution is policy control, where the learning policy is controlled over the support of the behavior policy. Thus the OOD points will not be visited and the distribution shift problem is alleviated [17, 18, 19]. It is noted that previous policy control methods limit the safe areas within the close neighborhoods of the training dataset by distribution distance metrics such as KL divergence, maximum mean discrepancy (MMD), and Wasserstein Distance [20, 17, 18]. Such policy control approaches often lead to an over-conservative and suboptimal learning policy, and the optimal policy value may not be reached. For example, if some state-action points are within the training distribution but unobserved in the training dataset, they may be considered OOD points in previous policy control approaches. Nevertheless, these points are safe when the estimated  $Q$ -function can well generalize on this area [12, 21], and the optimal policy likely lies in the unseen area of the data space. In this sense, a global estimator of the conditional action data space is required to define the safe areas and examine whether an unseen point is safe. To this end, Wu et al. [22] imposes the density-based constraints and proposes the SPOT method. However, the density estimator in SPOT approximates the lower bound of the behavior policy density, and policy control methods with more accurate density estimators are worth further discussing. As in the field of policy control methods for offline RL problems, the idea of ensuring the consistency of the support of learned policy and that of the behavior policy is considered the most desirable approach to tackle distribution shift. Therefore, a direct estimation density of the behavior policy is desired for guarantee support consistency.
Recently, generative models have shown strong capabilities in diverse fields, in which variational autoencoder (VAE) and generative adversarial network (GAN) [23, 24, 25, 26] are popular tools utilized for sampling and other related applications. In RL areas, both VAE and GAN have been introduced in RL methods [19, 16, 27, 28]. However, as random sample generators, VAE and GAN cannot be directly applied to estimate behavior densities in offline RL, nor can they provide accurate density estimates. Wang et al. [29], Yang et al. [30], Chen et al. [31], Ghasemipour et al. [32] have shown that VAE has shortcomings in estimating complex distributions, especially multimodal distributions or complex behavior policy distributions, and discuss that typically-used VAEs may not align well with the behavior dataset. As VAE is prone to cover mode, this phenomenon diminishes the VAE's capacity to accurately estimate data distributions, particularly in scenarios involving multimodal data distributions, such as behavior policy distributions. Compared to VAE, GAN has more advantages in learning complex data distributions. Recently, various GAN variants have been suggested. Among them, the Flow-GAN [33] model combines the normalizing flow model [34, 35] and GAN, and provides an explicit density estimator. The Flow-GAN model is further renowned for its robustness and tuning flexibility and has garnered increasing attention.
In this work, we propose a novel approach called Constrained Policy optimization with Explicit Behavior Density  $(\mathrm{CPED}^2)$  for offline RL. The CPED suggests a strong density-based constraint and employs an explicit and accurate density estimator to identify the safe areas. Specifically, CPED utilizes the Flow-GAN model as the density estimator, whose accuracy is theoretically guaranteed. By leveraging this capability, CPED can accurately identify the feasible region, which includes both observed and unobserved but safe points (illustrated in Figure 1 c). By expanding the feasible region and enabling reasonable exploration within it, CPED exhibits reduced risk-aversion and demonstrates an enhanced ability to generate superior policies.
To summarize, the merits of the proposed CPED include the following.
- By introducing the Flow-GAN model, the CPED directly tackles the density estimation problem and proposes a novel policy control scheme for solving offline RL tasks.
$^{2}$ Code available at https://github.com/evalarzj/cped
- The CPED does not behave over-conservatively and allows exploration in the safe region, especially high-density areas, to find the optimal policy.
- Practically, the effectiveness of CPED is fully demonstrated empirically. Extensive experiments on competitive tasks show that CPED substantially outperforms state-of-the-art methods.
- From the theoretical perspective, we show that GAN with a hybrid loss can accurately estimate density, which itself contributes to the GAN community. Furthermore, we prove the policy learned from CPED can access the optimal value function, theoretically ensuring the efficiency of CPED.
Figure 1: (a): The ground truth safe area in offline RL optimization, and the updates of policies and  $Q$ -functions are done within the green area. The blue points are collected behavior data  $\mathcal{D}$ , and the red point denotes the optimal policy given the states. (b): In previous approaches, the exploration of the policy takes place in a small neighborhood of points in  $\mathcal{D}$  (the orange circles). (c): The CPED relaxes the exploration area and constructs the feasible region (pink areas), which includes the unobserved but safe points (black point).
Value Function Constraint. Taking OOD actions can lead to overestimated value function and introduce bootstrap error. Value function constraint methods aim to construct pessimistic value functions or reduce the uncertainty of the value function to avoid OOD actions. Following [11, 12, 36], a pessimistic value function can be obtained by ensembling multiple value functions. CQL [13] proposes a regularization term so that the estimated  $Q$ -function low-bounds its true value. TQC [37] follows the distributional RL method and learns the percentile of value function distribution. IQL [15] combines expectile regression with exponential weighted advantage function to avoid extreme value function. MCQ [16] tries to use a special bellman operator to separate the estimation of the target value function for the in-distribution point and OOD point. The value function shows great uncertainty when taking OOD points. UWAC [14] estimates the uncertainty of the value function using the dropout technique. Urpi et al. [38] applies the tail risk measurement for the  $Q$ -function, in which the value function with high uncertainty is penalized, and the learned value function is forced to be pessimistic. However, finding the lower bound of the value function or estimating the uncertainty is not easy. The error from estimating the value function and finding the lower bound hinders the performance of value function constraint methods. Recently, a series of in-sample learning methods [39, 40, 41, 42, 43] have been proposed to optimize the  $Q$ -value function or the policy value function  $V$  solely utilizing in-sample data. As a result, the extrapolation error, induced by out-of-distribution (OOD) data points in the  $Q$ -value function, can be mitigated.
Policy Control Method. To prevent the learning of policies that take out-of-distribution (OOD) actions, the policy control method incorporates a penalty term during policy learning and restricts the distribution distance between the learned policy and the behavior policy. The BEAR method [17] uses the MMD distance to constrain the difference between the behavior policy and the learned policy. [18, 44, 20, 45, 46] measures the distribution shift between learning policy and behavior policy with KL divergence and Wasserstein distance. DOGE [46] learns a distance function to control the learning policy staying close to the behavior policy. Besides the explicit distance-based policy control methods, some implicit control methods [19, 47, 22], which generate the learning policy action based on the estimated conditional probability distributions, are suggested. However, explicit distance control methods suffer from the generalization problem since the distance is estimated from the limited training samples. On the other hand, implicit control methods tend to generate over-conservative
policies, where the generated actions are too close to the behavior policy. Consequently, the learned policy by these methods discourages exploration and remains conservative.
Behavior Policy Density Estimation In offline RL setting, determining the action space  $\mathcal{A}$  under given states is crucial. This is equivalent to learning the distribution of the behavior policy. However, traditional density estimation methods [48, 49, 50, 51, 52, 53] often fail in this scenario of density estimation due to the curse of high dimensionality. Recently, deep neural network-based density estimation methods have proven successful in learning the distribution of the behavior policy. Conditional variational autoencoders (CVAE) have been utilized as both behavior policy samplers [19, 17, 47, 16, 54] and density estimators [22]. Generative adversarial networks (GAN) are also introduced in offline RL scenario [55, 56]. Singh et al. [57] applies the generative normalized flow model to learn an invertible map from the noisy prior to the behavior policy action space and do behavior policy sampling by the trained model. EMaQ [32] uses the Autoregressive Generative Model instead of CVAE as a sampler. Additionally, the diffusion model has been applied to estimate the density of the behavior policy in both model-free RL [31] and model-based RL [58]. Specifically, SPOT [22] approaches the lower bound of the behavior density with ELBO derived from CVAE. However, to the best of our knowledge, most of the aforementioned methods either estimate the behavior density implicitly or lack the ability to provide accurate estimations. In our proposed method, the CPED finds an exact density function using the Flow model [34, 35] and GAN [23, 24, 25, 59, 60].
3 Probability Controlled Offline RL Framework
3.1 Offline RL Problem Settings
In RL setting, the dynamic system is described in terms of a Markov decision process (MDP), which can be defined by a set of tuples  $\{S, \mathcal{A}, T, p_0, r, \gamma\}$ . Here,  $S$  and  $\mathcal{A}$  denote the state space and action space, respectively;  $T$  is a conditional transition kernel, and  $T(s'|s, a)$  is the conditional transition probability between states, describing the dynamics of the entire system;  $p_0$  is the distribution of the initial states;  $r(s, a)$  is a deterministic reward function; and  $\gamma \in (0, 1)$  is a scalar discount factor.
The goal of RL is to find a distribution of actions  $\pi(a|s)$  (named as policy) that yields the highest returns within a trajectory, i.e., maximizing the expected cumulative discounted returns. We denote this expected discounted return starting from  $(s,a)$  under policy  $\pi$  as  $Q^{\pi}(s,a)$ . Under the Q-learning framework, the optimal expected discounted return  $Q^{*}(s,a)$  can be obtained by minimizing the  $L_{2}$  norm of Bellman residuals:
$$
\mathbb {E} _ {s ^ {\prime} \sim T (s ^ {\prime} | s, a), a ^ {\prime} \sim \pi (a | s)} [ Q (s, a) - (\mathcal {B} Q) (s, a) ] ^ {2}, \tag {1}
$$
where  $\mathcal{B}$  is the Bellman operator defined as
$$
(\mathcal {B} Q) (s, a) := r (s, a) + \gamma \mathbb {E} _ {s ^ {\prime \prime} \sim T (s ^ {\prime} | s, a)} [ \max  _ {a ^ {\prime}} Q (s ^ {\prime}, a ^ {\prime}) ].
$$
In the context of deep RL, various solutions are provided to estimate the  $Q$ -function [61, 62, 63]. When the action space is continuous, the learning process is divided into actor ( $Q$  value function) training and critic (learning policy) training [64]. The critic network is learned from Eq.1 without "maximum" in the Bellman operator. The actor is formulated with various forms [65, 66, 67, 1] such as maximizing the  $Q$ -value (DDPG), the weighted log of action density (policy gradient), and clipped weighted action density ratio (PPO).
In the offline RL, the learning policy cannot interact with the dynamic system yet only learns from a training dataset  $\mathcal{D}(s,a,s',r)$ . Dataset  $\mathcal{D}$  contains many triples  $(s,a,s',r)$  that can be viewed as independent. Assume that the dataset  $\mathcal{D}$  is generated by a behavior policy  $\pi_{\beta}$ . We need to redefine the MDP corresponding to  $\mathcal{D}$  under the offline RL settings.
Definition 3.1 (Offline MDP). Given a dataset  $\mathcal{D}$  with continuous state and action space. Let  $\pi_{\beta}$  denote the behavior policy that generated  $\mathcal{D}$ . The MDP of dataset  $\mathcal{D}$  is defined by  $\mathcal{M}_{\mathcal{D}} := \{\mathcal{S}, \mathcal{A}, T_{\mathcal{D}}, p_{0_{\mathcal{D}}}, r, \gamma\}$ , where  $\mathcal{S} = \{S, \mathbb{P}_{\mathcal{D}}\}$  is the measured state space related to  $\mathcal{D}$  with probability measure  $\mathbb{P}_{\mathcal{D}}$ ,  $\mathbb{P}_{\mathcal{D}}$  is the marginal probability of each state in state space of  $\mathcal{D}$ ,  $\mathcal{A} = \{\mathcal{A}, \pi_{\beta}\}$  is the measured action space related to  $\mathcal{D}$  with conditional probability measure  $\pi_{\beta}(a|s)$ ,  $T_{\mathcal{D}}(s'|s, a)$  is the transition probability between states, describing the dynamics of  $\mathcal{D}$ , and  $p_{0_{\mathcal{D}}}$  is the distribution of the initial states of  $\mathcal{D}$ .
In offline RL, we can obtain the optimal  $Q$ -function by minimizing Eq.1. However, due to the distribution shift problem [17, 8], the Bellman residuals are biased during the training phase. Specifically, at  $k$ -th iteration, since  $\pi_k$  is trained to maximize the  $Q$ -function, it is possible that  $\pi_k$  visits some OOD points with high estimated  $Q$ -function values. Nevertheless, we cannot evaluate the rewards of the OOD points under offline scenarios. Consequently the Bellman residuals at the OOD points may dominate the loss, and the learned  $Q$ -function value is misleading. Therefore, the optimization of offline RL must occur in the bounded safe state-action space  $\mathcal{S} \times \mathcal{A}$ , such that the estimated  $Q$ -function is trustworthy.
3.2 Behavior Policy Estimation by GAN with Specific Density
In offline RL, the dataset  $\mathcal{D}$  is generated by the behavior policy, whose density can be estimated via inverse RL [68]. One widely used approach in inverse RL, maximum entropy inverse RL (MaxEnt IRL, [69]), models the density via a Boltzmann distribution
$$
p _ {\theta} (\tau) = Z ^ {- 1} \exp (- c _ {\theta} (\tau)), \tag {2}
$$
where the energy is given by the cost function  $c_{\theta}(\tau) = \sum_{t=0}^{H} c_{\theta}(s_t, a_t)$ ,  $\tau = (s_0, a_0, s_1, a_1, \dots, s_H, a_H)$  is a trajectory, and  $Z$  is a scale factor. The cost function  $c_{\theta}(\tau) = \sum_{t=0}^{H} c_{\theta}(s_t, a_t)$  is a parameterized cost function that can be learned by maximizing the likelihood of the trajectories in  $\mathcal{D}$ . And the scale factor  $Z$  is the sum of exponential cost  $\exp(-c_{\theta}(\tau))$  of all the trajectories that belong to the product space of the state and action of  $\mathcal{D}$ . However, estimating  $Z$  is difficult because it is scarcely possible to traverse the entire action space  $\mathcal{A}$ . An alternative way to estimate  $Z$  is via guided cost learning [70], where it has been shown that learning  $p_{\theta}(\tau)$  is equivalent to training the dataset  $\mathcal{D}$  with a GAN [70, 71, 72, 73].
Originally, the GAN is a random sample generator for high-dimensional data, and it consists of two models: a generator  $G$  and a discriminator  $D$ . The generator is trained to generate random samples that are close to those drawn from the underlying true distribution, and the discriminator tries to classify whether the input is generated sample or an actual one from the training dataset. The GAN cannot provide an explicit density estimation, which is an indispensable ingredient when controlling the probability of actions and specifying an safe area. Thus, we cannot directly apply the MaxEnt IRL to our CPED. To overcome the difficulty, we follow the approach in [33] and consider GAN with flow model as a generator (we call if Flow-GAN hereafter as in [33]).
Flow-based generative models have shown their power in the fields of image processing and natural language processing [34, 35]. The normalizing flow model allows us to start from a simple prior noise distribution and obtain an explicit estimate of the density function after a series of non-linear invertible transformations.
Specifically, let  $f_{\theta}:\mathbb{R}^{A}\mapsto \mathbb{R}^{A}$  be a nonlinear invertible transformation, where  $A$  is the dimension of action space  $\mathcal{A}$ , and  $\theta$  is parameter. By change of variables, the probability density on  $\mathcal{D}$  can be obtained by
$$
p _ {\theta} (\tau) = p _ {H} (f (\tau)) \left| \det  \frac {\partial f _ {\theta} (\tau)}{\partial \tau} \right|, \tag {3}
$$
where  $\tau \in \mathcal{D}$  is a trajectory. Since  $f_{\theta}$  is invertible, a generator can be constructed by  $G_{\theta} = f_{\theta}^{-1}$ . In practice, both generator and discriminator in GAN are modeled as deep neural networks. Two specific neural network structures, NICE [34] and Real_NVP [35], are proposed to ensure that  $G_{\theta}$  (or  $f_{\theta}$ ) is invertible. The Jacobian matrices of the mapping functions  $f$  constructed by these two methods are both triangular matrices, guaranteeing efficient computation of  $|\operatorname*{det}\frac{\partial f(\tau)}{\partial\tau}|$ .
Directly training a Flow-GAN by original loss function in GAN may lead to poor log-likelihoods, and GAN is prone to mode collapse and producing less diverse distributions [33]. Thus, following the approach in [33], we adopt a hybrid loss function containing a maximum likelihood loss and a GAN structure loss as the optimization objective for estimating the probability measure of behavior policy:
$$
\min  _ {\theta} \max  _ {\phi} \mathcal {L} _ {\mathcal {G}} \left(G _ {\theta}, D _ {\phi}\right) - \lambda \mathbb {E} _ {\tau \sim P _ {d a t a}} [ \log \left(p _ {\theta} (\tau)\right) ] \tag {4}
$$
where  $\mathcal{L}_{\mathcal{G}}(G_{\theta},D_{\phi})$  can be any min-max loss function of GAN,  $D_{\phi}$  is a discriminator indexed by parameter  $\phi$ , and  $\lambda >0$  is a hyperparameter balancing the adversary learning process and the maximum likelihood process.
Remark 3.1. The reason for not simply using MLE is that MLE tends to cover all the modes of distribution (even though some modes have small probabilities), and MLE is not robust against model misspecification [74], which often occurs in high dimensional scenarios.
In the following proposition, we show that the estimated density of the behavior policy using the hybrid loss Eq.4 is equivalent to that using MaxEnt IRL. The proof of Proposition 3.1 is provided in Appendix B.
Proposition 3.1. For offline dataset  $\mathcal{D}$  generated by behavior policy  $\pi_{\beta}$ , the learned likelihood function  $L^{\pi_{\beta}}$ , using GAN with hybrid loss in Eq.4 is equivalent to that trained by MaxEnt IRL. If the generator of GAN can give a specific likelihood function  $p_{\theta}^{G}(\tau)$ , then
$$
L _ {\theta} ^ {\pi_ {\beta}} (\tau) = C Z ^ {- 1} \exp (- c _ {\theta} (\tau)) \propto p _ {\theta} ^ {G} (\tau) \tag {5}
$$
where  $C$  is a constant related to  $\mathcal{D}$
3.3 Constrained Policy Optimization with Explicit Behavior Density Algorithm (CPED)
In this work, we propose an algorithm for probabilistically controlling the learned policy, CPED. In CPED, we first utilize the Flow-GAN and obtain a probability estimate of the trajectory  $p_{\theta}(\tau)$ . The optimize objective follows the hybrid loss in Eq.4:
$$
\min  _ {\theta} \max  _ {\phi} \mathcal {L} _ {\mathcal {G}} \left(G _ {\theta}, D _ {\phi}\right) - \lambda \mathbb {E} _ {(s, a) \sim D} \left[ \log \left(L _ {\theta} ^ {\pi_ {\beta}} (s, a)\right)\right). \tag {6}
$$
Since the episode  $\tau$  is a sequential collection of states  $s$  and actions  $a$ , we rewrite the input of  $L_{\theta}^{\pi_{\beta}}(\cdot)$  with  $(s, a)$  instead of  $\tau$  in Eq.6 and afterwards. From Proposition 3.1, the solution to Eq.6 is equivalent to the behavior estimator obtained from MaxEnt IRL, thus providing an effective estimation of the behavior policy.
Following the learning paradigm in classic RL methods, the policy learning process is divided into two parts: the actor training and critic training. With the estimated function  $L_{\theta}^{\pi_{\beta}}(\cdot)$  approximating the density of behavior policy, we can naturally propose a density-based constraint in policy learning, and obtain the safe area according to the estimated density. The points with low density values are excluded and the updated policy cannot are exceed the safe region. The actor learning is formulated as: