SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yao Liu<sup>1</sup>, Pratik Chaudhari<sup>1,2</sup>, Rasool Fakoor<sup>1</sup>  
<sup>1</sup>Amazon Web Services, <sup>2</sup>University of Pennsylvania {yaoliuai, prtic, fakoor}@amazon.com
The main challenge of offline reinforcement learning, where data is limited, arises from a sequence of counterfactual reasoning dilemmas within the realm of potential actions: What if we were to choose a different course of action? These circumstances frequently give rise to extrapolation errors, which tend to accumulate exponentially with the problem horizon. Hence, it becomes crucial to acknowledge that not all decision steps are equally important to the final outcome, and to budget the number of counterfactual decisions a policy makes in order to control the extrapolation. Contrary to existing approaches that use regularization on either the policy or value function, we propose an approach to explicitly bound the amount of out-of-distribution actions during training. Specifically, our method utilizes dynamic programming to decide where to extrapolate and where not to, with an upper bound on the decisions different from behavior policy. It balances between the potential for improvement from taking out-of-distribution actions and the risk of making errors due to extrapolation. Theoretically, we justify our method by the constrained optimality of the fixed point solution to our  $Q$  updating rules. Empirically, we show that the overall performance of our method is better than the state-of-the-art offline RL methods on tasks in the widely-used D4RL benchmarks.
One of the primary hurdles in reinforcement learning (RL), or online RL, is its reliance on interacting with an environment in order to learn [47]. This can be a significant barrier to applying RL to real-world problems, where it may not be feasible or safe to interact with the environment directly [48]. In contrast, batch or offline RL [6, 27, 30] provides a more suitable framework for effectively learning policies by leveraging previously collected data to learn a policy. Notably, offline RL relies on a fixed but limited dataset comprising previously collected data from an unknown policy(s) and lacks the ability to continue interacting with the environment to gather additional samples.
These limitations in offline RL give rise to various challenges, with one notable issue being the occurrence of extrapolation problems resulting from the scarcity of training data [8, 12]. In order to overcome this challenge, previous approaches to offline RL primarily focus on constraining the gap between the behavioral policy and the learned policy [4, 11, 12, 25, 40, 51], or limiting the disparity between the state-action values of logged actions in the data and extrapolated actions [8, 22, 26, 50, 53]. Alternatively, other approaches utilize a learned model that avoids making predictions outside the distribution of the dataset [21, 35, 55]. All these offline RL methods involve a sequence of counterfactual reasoning problems within the realm of potential actions, explicitly or implicitly. The question of "what if" we were to choose a different course of action than the behavior policy often arises, leading to extrapolation errors. Hence the difficulty of the problem increases as we plan for a longer horizon and involves more counterfactual decisions. For offline RL algorithms, it is difficult to find a balance between the potential for improvement from taking out-of-distribution actions and the risk of making errors due to extrapolation.
A critical aspect of offline RL which is relatively unexplored pertains to the varying importance of different actions at each step in determining the final outcome. Not all actions are created equal, and some have more influence on the outcome than others. Thus, it is not always essential to consider alternative actions or counterfactuals at every step than the one dictated by the behavior policy. Rather, the emphasis should be placed on the actions that wield the most impact, warranting careful decision-making at those steps. This limits the number of counterfactual decisions we make and assigns them to the most needed states/steps, to increase the return on the "extrapolation investment".
Unlike existing approaches that employ explicit or implicit regularization on either the policy, value functions, or both, we introduce a novel and effective algorithm, called Budgeting Counterfactual for Offline RL (BCOL), that explicitly constrains the level of extrapolation during training. Specifically, our approach leverages dynamic programming to determine when and where extrapolation should occur, with a budget on the deviation from the behavior policy. This enables us to strike a balance between the potential for improvement achieved by exploring out-of-distribution actions and the inherent risks associated with extrapolation errors. Conceptually, such a method is different from regularized offline RL by allowing non-uniform constraints while keeping strict bounds on counterfactual decisions.
Our contribution. We propose a novel algorithm, BCOL, making only a few but important counterfactual decisions in offline RL. This idea extends the space of current offline RL methods that extensively focus on regularized methods. We demonstrate that the fixed point resulting from our  $Q$ -value updating rule corresponds to the optimal  $Q$ -value function, under the constraints that the policy deviates from the behavior policy within a budget. We conduct thorough evaluations on a diverse array of tasks from the widely-utilized D4RL benchmarks [9]. In terms of overall performance, our approach exhibited a favorable comparison to existing state-of-the-art methods.
2 Problem Settings
We study the RL problem in the context of the infinite horizon, discounted Markov Decision Process (MDP) [41]. An MDP is defined as a tuple  $\mathcal{M} = < S, \mathcal{A}, r, P_0, P, \gamma>$  where  $\mathcal{S}$  is a state space,  $\mathcal{A}$  is an action space, and both of these spaces can be infinite or continuous in nature.  $r: S \times \mathcal{A} \to \mathbb{R}$  is the reward function.  $P_0$  is a distribution over  $\mathcal{S}$  that the initial states are drawn from.  $P: S \times \mathcal{A} \to \Delta(S)$  maps a state-action pair to a distribution over state space.  $\gamma$  is the discount factor over future reward. The goal in MDP is to maximize the expectation of discounted future reward  $v^{\pi} = \mathbb{E}\left[\sum_{t=0}^{\infty} \gamma^{t} r_{t} \mid a_{t} \sim \pi\right]$ . An important function in MDP is the state-action value function, known as the  $Q$  function, which represents the expected future reward, given an initial state or state-action pair.  $Q^{\pi}(s, a) := \mathbb{E}\left[\sum_{t=0}^{\infty} \gamma^{t} r_{t} \mid s_{0} = s, a_{0} = a\right]$ .
The primary emphasis of this paper is on model-free learning algorithms within the context of offline RL. The goal is to learn a target policy  $\pi$  that maximizes the expected future reward using a fixed dataset  $\mathcal{D}$ . The dataset  $\mathcal{D}$  consists of transitions  $(s, a, r, s', a')$  where  $s$  can be drawn from any fixed distribution,  $a \sim \mu(\cdot | s)$ ,  $r = r(s, a)$ ,  $s' \sim P(\cdot | s, a)$ ,  $a' \sim \mu(\cdot | s')$ . Here  $\mu$  denotes an unknown behavior policy that was utilized in the past to collect dataset  $\mathcal{D}$ .
In offline RL, the task of learning an optimal policy poses a greater challenge compared to online RL. This is primarily due to the fact that the target policy can be different from the behavior policy, leading to the introduction of out-of-distribution actions. This challenge is further amplified by the use of function approximation and the problem horizon as the extrapolation errors from fitted  $Q$  function tend to accumulate over time-steps [8, 12]. In particular, it will results in the  $Q(s, a)$  values increasing dramatically for out-of-distribution state and action pairs. Consequently, this results in the learning of a policy that is likely to perform poorly and risky at the deployment time. Many previous works on offline RL aim to address these issues by some form of regularization, which might hurt the flexibility of making certain key decisions at each step (which we will return to later). However, we take a different approach in this paper where we emphasize on cautiousness of making counterfactual decisions and, hence propose a method that only deviates from behavior policy for a few steps.
As mentioned earlier, prior studies in offline RL mainly address the issue of extrapolation through two types of regularization terms: policy regularization or value regularization. Policy regularization
typically involves utilizing a divergence metric between current candidate policy and behavior policy while value regularization is commonly defined as the positive difference between  $Q(s, \hat{a})^1$  and  $Q(s, a)$ , where  $a \sim \mu(\cdot | s)$ . In these works, the regularization terms were uniformly applied to all samples in the dataset, enforcing a flat penalization on the distance between  $\pi(\cdot | s)$  and  $\mu(\cdot | s)$ , or flat penalization of overestimating  $Q(s, \hat{a})$  for all  $s$ . The utilization of a uniform regularization approach can present a problem. When strong regularization is enforced, the divergence bound from behavior policy over all states are small, hence the policy may not be able to improve significantly. In such cases, the resulting policy can only be as good as the behavior policy itself. Conversely, due to the fact that trajectory distribution mismatch increases exponentially with respect to the problem horizon [8, 12, 34, 36], if the bound is not small enough, it causes significant issues such as large extrapolation errors that lead to a risky and poor-performing policy. Therefore, applying regularization uniformly without considering the specific characteristics of each state can pose challenges and limitations in offline RL.
This work emphasizes the importance of making counterfactual decisions, i.e., decisions that are different from the decisions that would have been made by the behavior policy. Building upon the notion that not all decision steps carry equal importance, we introduce a novel offline RL algorithm that incorporates only a limited number of counterfactual decisions. In other words, the algorithm follows the behavior policy  $\mu$  most of the time, but it only makes counterfactual decisions (i.e. performing offline RL) in certain states. Thus intuitively it maximizes the potential improvement from counterfactual decisions while keeping the extrapolation error under control<sup>2</sup>. We refer to this number by the budget of counterfactual, denoted as  $B$ . The grid world MDP in Figure 1 shows a simple example of why using a budget of counterfactuals is an effective constraint.
Here, the goal is to find the shortest path from the start to the goal state without passing through the fail state, given a dataset of previously recorded trajectories from different and unknown behavior policies. There are two types of trajectories in the dataset, marked as yellow and blue in Figure 1. While yellow trajectories are the most common, the shortest path requires stitching together these two types of trajectories and taking a counterfactual decision (shown by dashed green line). A budget of counterfactuals ensures that the agent only takes the most beneficial counterfactual decision and prevents it from taking unnecessary counterfactual actions.
In order to spend the budget of counterfactual strategically, a dynamic programming algorithm is utilized to plan and assess the potential policy improvement resulting from actions taken outside of the data distribution. Thus, the choice of policy on each step needs to balance between the immediate gain from taking greedy actions in the current step and the benefit from future counterfactual decisions. Naturally, the policy also depends on how many counterfactual decisions it can take before it exceeds the upper bound  $B$ . We define the budget of counterfactual decisions at a time step  $t$  as follows:
$$
b _ {t + 1} = b _ {t} - \mathbb {1} \left\{\pi (\cdot | s _ {t}, b _ {t}) \neq \mu (\cdot | s _ {t}) \right\}, \quad b _ {0} = B, \tag {1}
$$
where  $\mathbb{1}$  is an indicator function. The policies studied in this paper take the current budget  $b_{t}$  as an input, which  $b_{t}$  is the initial budget  $B$  subtracted by the number of counterfactual steps taken before time step  $t$  as shown in (1). Given that, the goal of our method is to solve the following constrained policy optimization:
$$
\max  _ {\pi} \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r _ {t} \mid s _ {0} \sim P _ {0}, \pi \right] s. t. b _ {t} \geq 0, \forall t \geq 0, \forall \{(s _ {t}, b _ {t}, a _ {t}) \} _ {t = 0} ^ {\infty} \in \mathcal {E} _ {\mathcal {M}} (\pi) \tag {2}
$$
$\mathcal{E}_{\mathcal{M}}(\pi)$  is the support set of trajectory distribution introduced by the MDP and policy  $\pi$ . Without additional statements, later we only consider state-action-budget trajectories in  $\mathcal{E}_{\mathcal{M}}(\pi)$  and drop this
There are different approaches to obtaining  $\hat{a}$  [5, 8]. Generally, it approximates  $\mathrm{argmax}_a Q(s,\cdot)$  or is sampled from the current policy.
Technically, extrapolation refers to out-of-distribution actions. A counterfactual decision policy may still overlap with the behavior policy and thus not truly extrapolate. Whether or not actions from counterfactual policy lead to extrapolation is unknown when we only observe logged action rather than  $\mu$ . In this paper, we refer to this "possible extrapolation" as extrapolation as well, for the simplicity of discussion.
requirement for the ease of notation. Note that  $b_{t}$  in (2) is a function of the initial budget  $b_{0} = B$ , policies, and states before step  $t$ . To provide more clarity, we expand the constraint in (2) as follows:
$$
\sum_ {t = 0} ^ {\infty} \mathbb {1} \left\{\pi (\cdot | s _ {t}, b _ {t}) \neq \mu (\cdot | s _ {t}) \right\} \leq B \tag {3}
$$
In order to develop our offline RL algorithm that maximizes (2), we introduce a new Bellman operator that plans on backup future values as well as the number of counterfactual decisions.
Definition 1 (Counterfactual-Budgeting Bellman Operator).
$$
\mathcal {T} _ {\mathrm {C B}} Q (s, b, a) := r (s, a) + \gamma \underset {s ^ {\prime}} {\mathbb {E}} \left[ V _ {Q} \left(s ^ {\prime}, b\right) \right] \tag {4}
$$
$$
\text {w h e r e} V _ {Q} (s ^ {\prime}, b) := \left\{ \begin{array}{l l} \max  \{\max  _ {a ^ {\prime}} Q (s ^ {\prime}, b - 1, a ^ {\prime}), \mathbb {E} _ {a ^ {\prime} \sim \mu} Q (s ^ {\prime}, b, a ^ {\prime}) \} & b > 0 \\ \mathbb {E} _ {a ^ {\prime} \sim \mu} Q (s ^ {\prime}, b, a ^ {\prime}) & b = 0 \end{array} \right.
$$
The Bellman operator  $\mathcal{T}_{\mathrm{CB}}$  updates the  $Q$  values by taking the maximum value between two terms. The first term refers to the case where a counterfactual decision is made by maximizing  $Q$  values over all possible actions in the next state. This leads to a decrease in the counterfactual budget. The second term involves following the behavior policy in the next state (similar to SARSA [43]) while keeping the counterfactual budget intact. By selecting the optimal backup from the two cases, the  $Q$  value for a given budget  $b$  strikes a balance between the benefits of counterfactual decisions (i.e.  $\max_{a'} Q(s', b - 1, a')$ ) in the next step and in further future. Since  $b \leq B$ , there are at most  $B$  backup steps taking the maximum  $\max_{a'} Q(s', b - 1, a')$  in the backup path. This intuitively upper bounds the amount of extrapolation that the  $Q$  function can take. We can recover the standard Bellman operator on state-action  $Q$  functions by considering  $b \equiv \infty$  in  $Q(s, b, a)$  and the rule  $\infty - 1 = \infty$ . Since  $b = b - 1 = \infty$ , the first max operator in  $V_Q$  always prefers the former term and gives us the standard Bellman operator.
It is straightforward to show the counterfactual-budgeting Bellman operator is a  $\gamma$ -contraction. Based on that, we prove the fixed point is the optimal  $Q$  function constrained by a counterfactual budget.
Theorem 2. There exists a unique fixed point of  $\mathcal{T}_{\mathrm{CB}}$ , and it is
$$
Q ^ {\star} (s, b, a) := \max  _ {\pi} \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r _ {t} \mid s _ {0} = s, a _ {0} = a, b _ {1} = b, \pi \right], \text {s . t .} b _ {t} \geq 0, \forall t \geq 1 \tag {5}
$$
This theorem (proved in the appendix) indicates that the fixed point iteration of  $\mathcal{T}_{\mathrm{CB}}$  will lead to the optimal value function with the upper bound of counterfactual decisions. Thus it motivates us to minimize a temporal difference (TD) error of the counterfactual-budgeting Bellman operator, towards the goal of only improving behavior policy over a limited number of but important decision steps.
3.1 Algorithm
In this section, we derive a practical offline algorithm for our approach when the counterfactual-budgeting Bellman operator (4) is estimated via function approximation. First, we need to replace the expectations in (4) by one-sample estimation to evaluate the Bellman backup with the fixed dataset  $\mathcal{D}$ . Next, considering that the action space is continuous, we approximate the  $\max_{a}$  operator in (4) by maximizing over actions sampled from a policy network that is trained to maximize the current  $Q$  function. Note that these design choices are commonly employed in previous offline RL algorithms [8, 12, 25, 26]. As a result, we obtain a sampled-based counterfactual-budgeting Bellman operator denoted as  $\widehat{\mathcal{T}}_{\mathrm{CB}}$ .
Definition 3 (Approximate Counterfactual-Budgeting Bellman Operator).
$$
\forall (s, a), \widehat {\mathcal {T}} _ {\mathrm {C B}} Q _ {\theta} (s, b, a) := r (s, a) + \gamma \left\{ \begin{array}{l l} \max  \{\underset {\bar {a} \in \{a _ {k} \} _ {k = 1} ^ {m}} {\max } Q _ {\theta} (s ^ {\prime}, b - 1, \bar {a}), Q _ {\theta} (s ^ {\prime}, b, a ^ {\prime}) \} & b > 0 \\ Q _ {\theta} (s ^ {\prime}, b, a ^ {\prime}) & b = 0 \end{array} \right.
$$
Here  $Q^{\star}(s,b,a)$  is defined as given  $b_{1} = b$  rather than  $b_{0}$  because the action  $a$  is already given and the future (optimal) value should be independent to which distribution  $a$  is drawn from.
Algorithm 1 BCOL Training	Algorithm 2 BCOL Inference
1: Input: θ0, φ0, T, B, D	1: Input: θ, φ, B, μ, M
2: for t = 0 to T - 1 do	2: b0← B, s0 ~ M
3: θt+1← θt - αq∇θLQ(θ, θt; φt, B, D)	3: for t = 0 to trajectory ends do
4: φt+1← Φt - αp∇φLπ(φ; θt, B, D)	4: π, bt+1 ← Select(φ, μ; st, bt, Qθ)
5: end for	5: at ~ π(·), rt, st+1 ~ M
6: Return θT, φT	6: end for
where  $(s', a')$  sampled from  $\mathcal{D}$  and  $\{a_k\}_{k=1}^m$  are sampled from  $\pi$ . Although this operator requires additional input  $b$  to the  $Q$  functions compared with the standard Bellman operator, it does not add any extra data requirement. We can augment the  $(s, a, s', a')$  tuple from an offline dataset with all  $0 \leq b \leq B$  where  $B$  is the maximum number of counterfactual decisions that we consider.
To learn the  $Q$  function, we minimize the least square TD error introduced by  $\widehat{\mathcal{T}}_{\mathrm{CB}}$ , using dataset  $\mathcal{D}$  consisting of  $(s,a,s^{\prime},a^{\prime})$  sampled from behavior policy  $\mu$ . Subsequently, we update the policy by making it more likely to choose actions with higher Q-values. The objective functions for the  $Q$  function and the policy  $\pi$  are formally defined as follows, with  $\theta$  representing the parameters of the  $Q$  function and  $\phi$  representing the parameters of the policy networks:
$$
\mathcal {L} _ {Q} (\theta , \bar {\theta}; \phi , B, \mathcal {D}) := \sum_ {b = 0} ^ {B} \underset {(s, a, s ^ {\prime}, a ^ {\prime}) \sim \mathcal {D}} {\mathbb {E}} \left[ \left(Q _ {\theta} (s, b, a) - \widehat {\mathcal {T}} _ {\mathrm {C B}} Q _ {\bar {\theta}} (s, b, a)\right) ^ {2} \right] \tag {6}
$$
$$
\mathcal {L} _ {\pi} (\phi ; \theta , B, \mathcal {D}) := - \sum_ {b = 0} ^ {B} \underset {s \sim \mathcal {D}, a \sim \pi_ {\phi} (\cdot | s, b)} {\mathbb {E}} Q _ {\theta} (s, b, a) \tag {7}
$$
where  $\bar{\theta}$  denotes delayed target network. Algorithm 1 describes the training process of our method<sup>1</sup>. The counterfactual-budgeting Bellman operator has an interesting property that  $Q$  is monotonically increased with  $b$ , for any  $Q$  such that  $Q = \widehat{T}_{\mathrm{CB}}Q$ .
$$
Q (s, b, a) \geq Q (s, b ^ {\prime}, a), \quad \forall s, a, b > b ^ {\prime}. \tag {8}
$$
The proof is straightforward since  $Q$  with a larger budget maximizes the value over more action sequences. This property will be followed asymptotically. However, with a limited number of iterations, and function approximations, the gap is not always positive. Intuitively, enforcing the monotonic gap might reduce the search space and potentially accelerate the convergence. It is important to note that this modification does not alter the fixed point solution since  $Q^{\star}$  always has a monotonic gap. Such a regularizer does not enforce any behavior/pessimistic constraint to the policy. It only constrains the Q-values for the same action with different budgets, not the Q-values for different actions. This regularizer will help the dynamics programming over counterfactual actions finding a self-consistent solution.
To implement it, we introduce the following penalty term and add it to the  $\mathcal{L}_Q$ . This results in the revised form of  $\mathcal{L}_Q$  as follows:
$$
\mathcal {L} _ {Q} (\theta , \bar {\theta}; \phi , B, \mathcal {D}) + \omega \sum_ {b = 0} ^ {B - 1} \underset {s \sim \mathcal {D}, a \sim \pi_ {\phi} (\cdot | s, b)} {\mathbb {E}} \left[ \left(\max  \left\{Q _ {\theta} (s, b, a) - Q _ {\theta} (s, b + 1, a), 0 \right\}\right) ^ {2} \right] \tag {9}
$$
This penalty term minimizes the gap for the actions sampled from  $\pi_{\phi}(\cdot |s,b)$ . As  $\pi_{\phi}(\cdot |s,b)$  maximizes  $Q_{\theta}(s,b,\cdot)$ , which is supposed to be smaller than  $Q_{\theta}(s,b + 1,\cdot)$ , it can be viewed as a more efficient sampler of actions, compared to uniformly sampling, in order to ensure the constraints in Equation 8.
Putting these together, we have Algorithm 1: Budgeting Counterfactual for Offline reinforcement Learning (BCOL). It uses an offline actor-critic style algorithm to learn the  $Q$  function and a policy maximizing the learned  $Q$  values. However, as the problem is to find the best policy under the constraints, greedy policy  $\pi_{\phi}$  itself is not enough for inference. The action selection method during inference is implicitly included in the definition of  $T_{\mathrm{CB}}$ . At test time, the policy needs to look ahead based on the current budget  $b_{t}$  and the  $Q$  values of taking counterfactual actions v.s. taking the
<sup>1</sup>Note  $\mathcal{L}_{\pi}$  is not exact but abstract in the sense that gradient to the  $\phi$  as a distribution parameter cannot be calculated directly and need the so-called policy gradient estimator. We clarify our implementation in Section 3.3.
behavior policy. Based on that decision, we updated the new budget as well. We define this step as an operator below.
$$
\mathsf {S e l e c t} (\pi , \mu ; s, b, Q) := \left\{ \begin{array}{l} (\mu (\cdot | s), b) \quad \text {i f} \quad \underset {\bar {a} \sim \pi (s, b)} {\mathbb {E}} Q (s, b - 1, \bar {a}) \leq \underset {\bar {a} \sim \mu (s)} {\mathbb {E}} Q (s, b, \bar {a}) \text {o r} b = 0 \\ (\pi (\cdot | s, b), b - 1) \quad \text {o . w .} \end{array} \right.
$$
The complete inference time procedure is described in Algorithm 2. It also takes the learned policy and  $Q$  function parameters as well as an approximate behavior policy  $\widehat{\mu}$ . In case of unknown behavior policy,  $\widehat{\mu}$  can be learned from behavior cloning or any other imitation learning algorithm. Algorithm 2 starts with an initial counterfactual budget  $B$ , takes action each time according to the condition in Select, and update the budget  $b_{t}$  if the action is not drawn from  $\widehat{\mu}$ .