SlowGuess/NeurIPS_2023_OCR huggingface dataset:
ReDS: Offline Reinforcement Learning With Heteroskedastic Datasets via Support Constraints
Anikait Singh $^{1,*}$ , Aviral Kumar $^{1,*}$ , Quan Vuong $^{2}$ , Yevgen Chebotar $^{2}$ , Sergey Levine $^{1}$
$^{1}$ UC Berkeley,  $^{2}$ Google DeepMind (*Equal contribution)
asap7772@berkeley.edu
Offline reinforcement learning (RL) learns policies entirely from static datasets. Practical applications of offline RL will inevitably require learning from datasets where the variability of demonstrated behaviors changes non-uniformly across the state space. For example, at a red light, nearly all human drivers behave similarly by stopping, but when merging onto a highway, some drivers merge quickly, efficiently, and safely, while many hesitate or merge dangerously. Both theoretically and empirically, we show that typical offline RL methods, which are based on distribution constraints fail to learn from data with such non-uniform variability, due to the requirement to stay close to the behavior policy to the same extent across the state space. Ideally, the learned policy should be free to choose per state how closely to follow the behavior policy to maximize long-term return, as long as the learned policy stays within the support of the behavior policy. To instantiate this principle, we reweight the data distribution in conservative Q-learning (CQL) to obtain an approximate support constraint formulation. The reweighted distribution is a mixture of the current policy and an additional policy trained to mine poor actions that are likely under the behavior policy. Our method, CQL (ReDS), is theoretically motivated, and improves performance across a wide range of offline RL problems in games, navigation, and pixel-based manipulation.
Recent advances in offline RL [39, 36] hint at exciting possibilities in learning high-performing policies, entirely from offline datasets, without requiring dangerous [19] or expensive [25] active interaction. Analogously to the importance of data diversity in supervised learning [9], the practical benefits of offline RL depend heavily on the coverage of behavior in the offline datasets [35]. Intuitively, the dataset must illustrate the consequences of a diverse range of behaviors, so that an offline RL method can determine what behaviors lead to high returns, ideally returns that are significantly higher than the best single behavior in the dataset.
One easy option to attain this kind of coverage is to combine many realistic sources of data, but doing so can lead to the variety of demonstrated behaviors varying in highly non-uniform ways across the state space, i.e. the dataset is heteroskedastic. For example, a driving dataset might show very high variability in driving habits, with some drivers being timid and some more aggressive, but remain remarkably consistent in "critical" states (e.g., human drivers are extremely unlikely to swerve in an empty road or drive off a bridge). A good offline RL algorithm should combine the best parts of each behavior in the dataset – e.g., in the above example, the algorithm should produce a policy that is as good as the best human in each situation, which would be better than any human driver overall. At the same time, the learned policy should not attempt to extrapolate to novel actions in subset of the state space where the distribution of demonstrated behaviors is narrow (e.g., the algorithm should not attempt to drive off a bridge). How effectively can current offline RL methods selectively choose on a per-state basis how closely to stick to the behavior policy?
Most existing methods [32, 33, 28, 27, 53, 17, 23] constrain the learned policy to stay close to the behavior policy with so-called "distribution constraints". Using a combination of empirical and theoretical evidence, we first show that distribution constraints are insufficient when the heteroskedasticity of the demonstrated behaviors varies non-uniformly across states, because the strength of the constraint is state-agnostic, and may be overly conservative at some states even when it is not conservative enough at other states. We also devise a measure of heteroskedasticity that enables us to determine if certain offline datasets would be challenging for distribution constraints.
Our second contribution is a simple observation: distribution constraints against a reweighted version of the behavior policy give rise to support constraints. That is, the return-maximization optimization process can freely choose per-state how much the learned policy should stay close to the behavior policy, so long as the learned policy remains within the data support. We show that it is convenient to instantiate this insight on top of conservative Q-learning (CQL) [33], a recent offline RL method. The new method, CQL (ReDS), changes minimally the form of regularization, design decisions employed by CQL and inherits existing hyper-parameter values. CQL (ReDS) attains better performance than recent distribution constraints methods on a variety of tasks with more heteroskedastic distributions.
The goal in offline RL is to find the optimal policy in a Markov decision process (MDP) specified by the tuple  $\mathcal{M} = (\mathcal{S},\mathcal{A},T,r,\mu_0,\gamma)$ .  $\mathcal{S},\mathcal{A}$  denote the state and action spaces.  $T(\mathbf{s}'|\mathbf{s},\mathbf{a})$  and  $r(\mathbf{s},\mathbf{a})$  represent the dynamics and reward function.  $\mu_0(s)$  denotes the initial state distribution.  $\gamma \in (0,1)$  denotes the discount factor. We wish to learn a policy that maximizes return, denoted by  $J(\pi) \coloneqq \frac{1}{1 - \gamma}\mathbb{E}_{(\mathbf{s}_t,\mathbf{a}_t)\sim \pi}[\sum_t\gamma^t r(\mathbf{s}_t,\mathbf{a}_t)]$ . We must find this policy while only having access to an offline dataset of transitions collected using a behavior policy  $\pi_{\beta}$ ,  $\mathcal{D} = \{(\mathbf{s},\mathbf{a},r,\mathbf{s}')\}$ .
Offline RL via distributional constraints. Most offline RL algorithms regularize the learned policy  $\pi$  from querying the target Q-function on unseen actions [17, 30], either implicitly or explicitly. For our theoretical analysis, we will abstract the behavior of distributional constraint offline RL algorithms into a generic formulation following Kumar et al. [33]. As shown in Equation 1, we consider the problem where we must maximize the return of the learned policy  $\pi$  (in the empirical MDP)  $\widehat{J} (\pi)$ , while also penalizing the divergence from  $\pi_{\beta}$ :
$$
\max  _ {\pi} \mathbb {E} _ {\mathbf {s} \sim \widehat {d} ^ {\pi}} \left[ \widehat {J} (\pi) - \alpha D (\pi , \pi_ {\beta}) (\mathbf {s}) \right], \tag {1}
$$
where  $D$  denotes a divergence between the learned policy  $\pi$  and the behavior policy  $\pi_{\beta}$  at state  $\mathbf{s}$ .
Conservative Q-learning. [33] enforces the distributional constraint on the policy implicitly. To see why this is the case, consider the CQL objective, which consists of two terms:
$$
\min  _ {\theta} \alpha \underbrace {\left(\mathbb {E} _ {\mathbf {s} \sim \mathcal {D} , \mathbf {a} \sim \pi} \left[ Q _ {\theta} (\mathbf {s} , \mathbf {a}) \right] - \mathbb {E} _ {\mathbf {s} , \mathbf {a} \sim \mathcal {D}} \left[ Q _ {\theta} (\mathbf {s} , \mathbf {a}) \right]\right)} _ {\mathcal {R} (\theta)} + \frac {1}{2} \mathbb {E} _ {\mathbf {s}, \mathbf {a}, \mathbf {s} ^ {\prime} \sim \mathcal {D}} \left[ \left(Q _ {\theta} (\mathbf {s}, \mathbf {a}) - \mathcal {B} ^ {\pi} \bar {Q} (\mathbf {s}, \mathbf {a})\right) ^ {2} \right], \tag {2}
$$
where  $\mathcal{B}^{\pi}\bar{Q} (\mathbf{s},\mathbf{a})$  is the Bellman backup operator applied to a delayed target Q-network,  $\bar{Q}$ :  $\mathcal{B}^{\pi}\bar{Q} (\mathbf{s},\mathbf{a})\coloneqq r(\mathbf{s},\mathbf{a}) + \gamma \mathbb{E}_{\mathbf{a}^{\prime}\sim \pi (\mathbf{a}^{\prime}|\mathbf{s}^{\prime})}[\bar{Q} (\mathbf{s}^{\prime},\mathbf{a}^{\prime})]$ . The second term (in blue) is the standard TD error [40, 18, 22]. The first term  $\mathcal{R}(\theta)$  (in red) attempts to prevent overestimation in the Q-values for out-of-distribution (OOD) actions by minimizing the Q-values under a distribution  $\mu (\mathbf{a}|\mathbf{s})$ , which is automatically chosen to pick actions with high Q-values  $Q_{\theta}(\mathbf{s},\mathbf{a})$ , and counterbalances by maximizing the Q-values of the actions in the dataset. Kumar et al. [33] show that Equation 2 gives rise to a pessimistic Q-function that modifies the optimal Q function by the ratios of densities,  $\pi (\mathbf{a}|\mathbf{s}) / \pi_{\beta}(\mathbf{a}|\mathbf{s})$  at a given state-action pair  $(s,a)$ . Formally, the Q-function obtained after one iteration is given by:
$$
Q _ {\theta} (\mathbf {s}, \mathbf {a}) := \mathcal {B} ^ {\pi} \bar {Q} (\mathbf {s}, \mathbf {a}) - \alpha \left[ \frac {\pi (\mathbf {a} | \mathbf {s})}{\pi_ {\beta} (\mathbf {a} | \mathbf {s})} - 1 \right]. \tag {3}
$$
The Q function is unchanged only if the density of the learned policy  $\pi$  matches that of the behavior policy  $\pi_{\beta}$ . Otherwise, for state-action pairs where  $\pi (\mathbf{a}|\mathbf{s}) < \pi_{\beta}(\mathbf{a}|\mathbf{s})$ , Eq. 3 increases their Q values and encourages the policy  $\pi$  to assign more mass to the action. Vice versa, if  $\pi (\mathbf{a}|\mathbf{s}) > \pi_{\beta}(\mathbf{a}|\mathbf{s})$ , Eq. 3 encourages the policy  $\pi$  to assign smaller density to the action  $\mathbf{a}$ . In Eq. 3,  $\alpha$  is a constant for every state, and hence the value function learned by CQL is altered by the ratio of action probabilities to the same extent at all possible state-action pairs. As we will discuss in the next section, this can be sub-optimal when the learnt policy should stay close to the behavior policy in some states, but not others. We elaborate on this intuition in the next section.
3 Why Distribution Constraints Fail with Heteroskedastic Data
In statistics, heteroskedasticity is typically used to refer to the condition when the standard deviation in a given random variable varies non-uniformly over time (see for example, Cao et al. [6]). We call a offline dataset heteroskedastic when the variability of the behavior differs in different regions of the state space: for instance, if for certain regions of the state space, the observed behaviors in the dataset assign the most probability mass to a few actions, but in other regions, the observed behaviors are more diverse. Realistic offline datasets are often heteroskedastic as they are typically generated by multiple policies, each with its own characteristics, under different conditions. E.g., driving datasets come from multiple humans [11], and many robotic datasets are collected by multiple teleoperators [10], resulting in systematic variability in different regions of the state space.
3.1 A Didactic Example
To understand why distribution constraints are insufficient with heteroskedastic data, we present a didactic example. Motivated by the driving scenario, we consider a maze navigation task shown in Fig. 1. The task is to navigate from the position labeled as "Start" to the position labeled as "Goal" using five actions at every possible state (L:  $\leftarrow$ , R:  $\rightarrow$ , U:  $\uparrow$ , D:  $\downarrow$ , No: No Op), while making sure that the executed actions do not hit the walls of the grid.
Dataset construction. To collect a heteroskedastic dataset, we consider a mixture of several behavior policies that attain a uniform occupancy over different states in the maze. However, the dataset action distributions differ significantly in different states. The induced action distribution is heavily biased to move towards the goal in the narrow hallways (e.g., the behavior policy moves upwards at state A). In contrast, the action distribution is quite diverse in the wider rooms. In these rooms, the behavior policy often selects actions that do not immediately move the agent towards
the goal (e.g., the behavior policy at state B), because doing so does not generally hit the walls as the rooms are wider, and hence the agent is not penalized. Whereas, the agent must take utmost precaution to not hit the walls in the narrow hallways. More details are in Appendix B.
Figure 1: Failure mode of distribution constraints. In this navigation task, an offline RL algorithm must find a path from the start state to the goal state as indicated in (a). The offline dataset provided exhibits non-uniform coverage at different states, e.g., in the state marked as "B" located in a wide room has more uniform action distribution, whereas the states in the narrow hallways exhibit a more narrow action distribution. This is akin to how the behavior of human drivers varies in certain locations ("B"), but is very similar in other situations ("A"). To perform well, an algorithm must stay close to the data in the hallways ("A"), but deviate significantly from the data in the rooms ("B"), where the data supports many different behaviors (most are not good). AWR and CQL get stuck because they stay too close to the bad behavior policy in the rooms, e.g. the left and right arrows near State B in Fig (b) and (c). Our method, CQL (ReDS), learns to ignore the bad behavior action in state B and prioritizes the good action, indicated by the downward arrow near B in (d).
Representative distribution constraint algorithms such as AWR [44, 43] and CQL [33] fail to perform the task, as shown in Figure 1. To ensure fair comparison, we tune each method to its best evaluation performance using online rollouts. The visualization in Figure 1 demonstrates that these two algorithms fail to learn reasonable policies because the learned policies match the random behavior of the dataset actions too closely in the wider rooms, and therefore are unable to make progress towards the Goal position. This is a direct consequence of enforcing too strong of a constraint on the learned policy to stay close to the behaviors in the dataset. Therefore, we also evaluated the performance of CQL and AWR in this example, with lower amounts of conservatism (Appendix B) and found that utilizing a lower amount of conservatism suffers from the opposite failure mode: it is unable to prevent the policies from hitting the walls in the narrow hallways. This means that conservatism prevents the algorithm from making progress in the regions where the behavior in the dataset is more diverse, whereas not being conservative enough hurts performance in regions where the behaviors in the dataset agree with each other. The method we propose in this paper to tackle this challenge, indicated as "CQL (ReDS)", effectively traverses the maze,  $80\%$  of the time.
3.2 Challenges with Distribution Constraints
Having seen that distribution constraints can fail in certain scenarios, we now formally characterize when offline RL datasets is heteroskedastic, and why distribution constraints may be ineffective in
Figure 2: Empirically computing  $C_{\mathrm{diff}}^{\pi}$  with three datasets: uniform (top), mixed (middle) and skewed (bottom) on a gridworld. We also visualize  $D(\pi, \pi_{\beta})(s)$  across states in the maze as the colors on different cells, a histogram of  $D(\pi, \pi_{\beta})(s)$  to visualize variation in this quantity and the performance of running standard CQL. Top: The uniform distribution leads to low  $C_{\mathrm{diff}}^{\pi}$ , uniform  $D(\pi, \pi_{\beta})(s)$ , and highest success. Middle: The mixed distribution leads to medium  $C_{\mathrm{diff}}^{\pi}$ , less uniformly distributed  $D(\pi, \pi_{\beta})(s)$ , and a drop in task success. Bottom: The skewed distribution leads to a high  $C_{\mathrm{diff}}^{\pi}$ , non-uniform  $D(\pi, \pi_{\beta})(s)$ , and poor performance.
such scenarios. Similar to how standard analyses utilize concentrability coefficient [46], which upper bounds the ratio of state-action visitation under a policy  $d^{\pi}(\mathbf{s},\mathbf{a})$  and the dataset distribution  $\mu$ , i.e.,  $\max_{\mathbf{s},\mathbf{a}}d^{\pi}(\mathbf{s},\mathbf{a}) / \mu (\mathbf{s},\mathbf{a})\leq C^{\pi}$ , we introduce a new metric called differential concentrability, which measures dataset heteroskedasticity (i.e., the variability in the dataset behavior across different states).
Definition 3.1 (Differential concentrability.). Given a divergence  $D$  over the action space, the differential concentrability of a given policy  $\pi$  with respect to the behavioral policy  $\pi_{\beta}$  is given by:
$$
C _ {\text {d i f f}} ^ {\pi} = \underset {\mathbf {s} _ {1}, \mathbf {s} _ {2} \sim d ^ {\pi}} {\mathbb {E}} \left[ \left(\sqrt {\frac {D (\pi , \pi_ {\beta}) (\mathbf {s} _ {1})}{\mu (\mathbf {s} _ {1})}} - \sqrt {\frac {D (\pi , \pi_ {\beta}) (\mathbf {s} _ {2})}{\mu (\mathbf {s} _ {2})}}\right) ^ {2} \right]. \tag {4}
$$
Eq. 4 measures the variation in the divergence between a given policy  $\pi(\mathbf{a}|\mathbf{s})$  and the behavior policy  $\pi_{\beta}(\mathbf{a}|\mathbf{s})$  weighted inversely by the density of these states in the offline dataset (i.e.,  $\mu(\mathbf{s})$  in the denominator). For simplicity, let us revisit the navigation example from Section 3.1 and first consider a scenario where  $\mu(\mathbf{s}) = \mathrm{Unif}(S)$ . For any given policy  $\pi$ , if there are states where  $\pi$  chooses actions that lie on the fringe of the data distribution (e.g., in the wider rooms), as well as states where the policy  $\pi$  chooses actions at the mode of the data distribution (e.g., as in the narrow passages), then  $C_{\mathrm{diff}}^{\pi}$  would be large any policy  $\pi$  that we learn. Crucially,  $C_{\mathrm{diff}}^{\pi}$  would be small even if the learned policy  $\pi$  deviates significantly from the behavior policy  $\pi_{\beta}$ , such that  $D(\pi, \pi_{\beta})(\mathbf{s})$  is large, but  $|D(\pi, \pi_{\beta})(\mathbf{s}_1) - D(\pi, \pi_{\beta})(\mathbf{s}_2)|$  is small, indicating the dataset is not heteroskedastic.
Connection between variability in the action distribution and high  $C_{\mathrm{diff}}^{\pi}$ . Consider a simpler formula where we remove the counts  $n(\mathbf{s})$  from the expression of differential concentrability and set  $\pi$  in  $C_{\mathrm{diff}}^{\pi}$  to be the uniform distribution over actions. Then, we can show that the value of  $C_{\mathrm{diff}}^{\pi}$  is exactly equal to twice the variance of  $D(\pi, \pi_{\beta})(\mathbf{s})$  across states. Therefore, we will demonstrate in Section 5 that arbitrary policy checkpoints  $\pi$  learned by offline RL algorithms generally attain a low value of the variance in  $D(\pi, \pi_{\beta})(\mathbf{s})$  on offline datasets from non-heteroskedastic sources, such as those covered in the D4RL [13] benchmark. Of course, we cannot always exclude the counts of states  $n(\mathbf{s})$ , however, we note that in high-dimensional state spaces, such as those in our experiments, each state in the offline data is likely to be unique, thus validating the condition that  $n(\mathbf{s}) = 1$ . That said, we do compute the exact value of  $C_{\mathrm{diff}}^{\pi}$  (with  $n(\mathbf{s})$ ) in a didactic gridworld maze shown in Figure 2. In this case, we find that our definition of  $C_{\mathrm{diff}}^{\pi}$  is actually able to reflect the intuitive notion of heteroskedasticity.
We now use the definition of differential concentrability to bound both the improvement and deprovement of  $\pi$  w.r.t.  $\pi_{\beta}$  for distribution constraint algorithms using the framework of safe policy
improvement [37, 33]. We show that when  $C_{\mathrm{diff}}^{\pi}$  is large, then constraints (Eq. 1) may not improve significantly over  $\pi_{\beta}$ , even for the best value for the weight  $\alpha$  (proof in Appendix C):
Theorem 3.2 (Informal; Limited policy improvement via distributional constraints.). W.h.p.  $\geq 1 - \delta$  for any prescribed level of safety  $\zeta$ , the maximum possible policy improvement over choices of  $\alpha$ ,  $\max_{\alpha} [J(\pi_{\alpha}) - J(\pi_{\beta})] \leq \zeta^{+}$ , where  $\zeta^{+}$  is given by:
$$
\zeta^ {+} := \max  _ {\alpha} \frac {h ^ {*} (\alpha)}{(1 - \gamma) ^ {2}} s. t. \frac {c _ {1} \sqrt {\log \frac {| S | | \mathcal {A} |}{\delta}}}{(1 - \gamma) ^ {2}} \frac {\sqrt {C _ {d i f f} ^ {\pi_ {\alpha}}}}{| \mathcal {D} |} - \frac {\alpha \mathbb {E} _ {\mathbf {s} \sim \widehat {d} ^ {\pi_ {\alpha}}} [ D (\pi_ {\alpha} , \pi_ {\beta}) (\mathbf {s}) ]}{1 - \gamma} \leq \zeta , \tag {5}
$$
where  $h^*$  is a monotonically decreasing function of  $\alpha$ , and  $h(0) = \mathcal{O}(1)$ .
Theorem 3.2 quantifies the fundamental tradeoff with distribution constraints: to satisfy a given  $\zeta$ -safety constraint in problems with larger  $C_{\mathrm{diff}}^{\pi}$ , we would need a larger  $\alpha$ . Since the maximum policy improvement  $\zeta^{+}$  is upper bounded by  $h^{*}(\alpha)$ , the policy may not necessarily improve over the behavior policy if  $\alpha$  is large. On the flip side, if we choose to fix the value of  $\alpha$  to be small in hopes to attain more improvement in problems where  $C_{\mathrm{diff}}^{\pi}$  is high for all policies, we would end up compromising on the safety guarantee as  $\zeta$  needs to be large for a small  $\alpha$  and large  $C_{\mathrm{diff}}^{\pi}$ . Thus, in this case, the policy may not improve over the behavior policy reliably.
Note that a larger value of  $C_{\mathrm{diff}}^{\pi}$  need not imply large  $\mathbb{E}_{\mathbf{s} \sim \widehat{d}^{\pi}}[D(\pi, \pi_{\beta})(\mathbf{s})]$  because the latter does not involve  $\mu(\mathbf{s})$ .  $C_{\mathrm{diff}}^{\pi}$  also measures the dispersion of  $D(\pi, \pi_{\beta})(\mathbf{s})$ , while the latter performs a mean over states. In addition, Theorem 3.2 characterizes the maximum possible improvement with an oracle selection of  $\alpha$ , though it is not feasible in practice. Thus, when  $C_{\mathrm{diff}}^{\pi}$  is large, distribution constraint algorithms could either not safely improve over  $\pi_{\beta}$  or would attain only a limited improvement with any possible value of  $\alpha$ . Finally, we remark that complementing [32, 39] that discuss failure modes of distribution constraints with high-entropy behavior policies, Theorem 3.2 quantifies when this would be the case: this happens when  $C_{\mathrm{diff}}^{\pi}$  is large.
4 Support Constraints As Reweighted Distribution Constraints
Thus far, we have seen that distribution constraints can be ineffective with heteroskedastic datasets. If we can impose the distribution constraint such that the constraint strength can be modulated per state, then in principle, we can alleviate the issue raised in Theorem 3.2 and Section 3.1.