SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hao Sun\*, Alihan Huyuk, Daniel Jarrett, Mihaela van der Schaar  
Department of Applied Mathematics and Theoretical Physics  
University of Cambridge
Learning controllers with offline data in decision-making systems is an essential area of research due to its potential to reduce the risk of applications in real-world systems. However, in responsibility-sensitive settings such as healthcare, decision accountability is of paramount importance, yet has not been adequately addressed by the literature. This paper introduces the Accountable Offline Controller (AOC) that employs the offline dataset as the Decision Corpus and performs accountable control based on a tailored selection of examples, referred to as the Corpus Subset. AOC operates effectively in low-data scenarios, can be extended to the strictly offline imitation setting, and displays qualities of both conservation and adaptability. We assess AOC's performance in both simulated and real-world healthcare scenarios, emphasizing its capability to manage offline control tasks with high levels of performance while maintaining accountability.
In recent years, offline control that uses pre-collected data to generate control policies has gained attention due to its potential to reduce the costs and risks associated with applying control algorithms in real-world systems [1], which are especially advantageous in situations where real-time feedback is challenging or expensive to obtain [2-5]. However, existing literature has primarily focused on enhancing learning performance, leaving the need for accountable and reliable control policies in offline settings largely unaddressed, particularly for high-stakes, responsibility-sensitive applications.
However, in many critical real-world applications such as healthcare, the challenge is beyond enhancing policy performance. It requires the decisions made by learned policies to be transparent, traceable, and justifiable. Yet those essential properties, summarized as Accountability, are left largely unaddressed by existing literature.
In our context, we use Accountability to indicate the existence of a supportive basis for decision-making. For instance, in tumor treatment using high-risk options like radiotherapy and chemotherapy, the treatment decisions should be based on the successful outcomes experienced by previous patients who share similar conditions and were given the same medication. Another concrete illustrative example is the allocation decisions of ventilator machines. The decision to allocate a ventilator should be accountable, in the way that it juxtaposes the potential consequences of both utilization and non-utilization and provides a reasonable decision on those bases. In those examples, the ability to refer to existing cases that support current decisions can enhance reliability and facilitate reasoning or debugging of the policy. To advance offline control towards real-world responsibility-sensitive applications, five properties are desirable: (P1) controllable conservation: this ensures the policy learning performance by avoiding aggressive extrapolation. (P2) accountability: as underscored by our prior examples, there is a need for a clear basis upon which decisions are made. (P3) suitability
*hs789@cam.ac.uk
for low-data regimes: given the frequent scarcity of high-stake decision data, it's essential to have methods that perform well with limited data. (P4) adaptability to user specification: this ensures the policy can adjust to changes, like evolving clinical guidelines, allowing for tailored solutions. (P5) flexibility in Strictly Offline Imitation: this property ensures a broader applicability across various scenarios and data availability.
To embody all these properties, we need to venture beyond the current scope of literature focused on conservative offline learning. In our work:
1. Methodologically, we introduce the formal definitions and necessary concepts in accountable decision-making. We propose the Accountable Offline Controller (AOC), which makes decisions according to a decomposition on the basis of the representative existing decision examples.
2. Theoretically, we prove the existence and uniqueness of the decomposition under mild conditions, guiding the design of our algorithms.
3. Practically, we introduce an efficient algorithm that takes all the aforementioned desired properties into consideration and circumvented the computational difficulty.
4. Empirically, we verify and highlight the desired properties of AOC on a variety of offline control tasks, including five simulated continuous control tasks and one real-world healthcare dataset.
POMDP We develop our work under the general Partially Observable Markov Decision Process (POMDP) setting, denoted as a tuple  $(\mathcal{X},\omega ,\mathcal{O},\mathcal{A},\mathcal{T},\mathcal{R},\gamma ,\rho_0)$ , where  $\mathcal{X}\subseteq \mathbb{R}^{d_x}$  denotes the underlying  $d_{x}$  dimensional state space,  $\omega :\mathcal{X}\mapsto \mathcal{O}$  is the emission function that maps the underlying state into a  $d_{o}$  dimensional observation  $\mathcal{O}\subseteq \mathbb{R}^{d_o};\mathcal{A}\subseteq \mathbb{R}^{d_a}$  is a  $d_{a}$  dimensional action space, transition dynamics  $\mathcal{T}:\mathcal{X}\times \mathcal{A}\mapsto \mathcal{X}$  controls the underlying transition between states given actions; and the reward function  $\mathcal{R}:\mathcal{X}\mapsto \mathbb{R}$  maps state into a reward scalar; we use  $\gamma$  to denote the discount factor and  $\rho_0$  the initial state distribution. Note that the MDP is the case when  $\omega$  is an identical mapping.
Offline Control Specifically, our work considers the offline learning problem: a logged dataset  $\mathcal{D} = \{o_t^i,a_t^i,r_t^i,o_{t + 1}^i\}_{t = 1,\dots ,T}^{i = 1,\dots ,N}$  containing  $N$  trajectories with length  $T$  is collected by rolling out some behavior policies in the POMDP, where  $x_0^i\sim \rho_0$  sampled from the initial state distribution and  $o_0^i = \omega (x_0^i),a_t^i\sim \pi_b^i$  sampled from unknown behavior policies  $\pi_b^i$ . We denote the observational transition history until time  $t$  as  $h_t$ , such that  $h_t = (o_{< t},a_{< t},r_{< t})\in \mathcal{H}\subseteq \mathbb{R}^{(d_o + d_a + 1)\cdot (t - 1)}$ .
Learning Objective A policy  $\pi : \mathcal{O} \times \mathcal{H} \mapsto \mathcal{A}$  is a function of observational transition history. The learning objective is to find  $\pi$  that maximizes the expected cumulative return in the POMDP:
$$
\max  _ {\pi} \mathbb {E} _ {\tau \sim \rho_ {0}, \pi , \mathcal {T}} \left[ \gamma^ {t} r _ {t} \right] \tag {1}
$$
where  $\tau \sim \rho_0, \pi, \mathcal{T}$  means the trajectories  $\tau$  is generated by sampling initial state  $x_0$  from  $\rho_0$ , and action  $a$  from  $\pi$ , such that:  $\tau = \{x_0, o_0 = \omega(x_0), a_0 = \pi(o_0, h_0 = \phi), x_1 = \mathcal{T}(x_0, a_0), r_0 = \mathcal{R}(x_1), \ldots\}$ . A learned belief variable  $b_t \in \mathbb{R}^{d_b}$  can be introduced to represent the underlying state  $x_t$ .
3 Accountable Control with Decision Corpus
3.1 Method Sketh: A Roadmap
The high-level core idea of our work is to introduce an example-based accountable framework for offline decision-making, such that the decision basis can be clear and transparent.
To achieve this, a naive approach would be to leverage the insights of Nearest Neighbors: for each action, this involves finding the most similar transitions in the offline dataset, and estimating the corresponding outcomes. Nonetheless, a pivotal hurdle arises in defining similarity, particularly when taking into account the intricate nature of trajectories, given both the observation space heterogeneity and the inherent temporal structure in decision-making. Compounding such a challenge, another difficulty arises in identifying the most representative examples and integrating the pivotal principle
of conservation, which is widely acknowledged to be essential for the success of offline policy learning.
Our proposed method seeks to address those challenges. We start by introducing the basic definitions to support a formal discussion of accountability: in Definition 3.1, we introduce the concept of Decision Corpus.
To address the similarity challenge, we showcase that a nice linear property (Property 3.2) generally exists (Remarks 3.3 & 3.4) when working in the belief space (Definition 3.5). This subsequently leads to a theoretical bound for estimation error (Proposition 3.8);
To address the representative challenge while obeying the principle of conservation, we underscore those examples that span the convex hull (Definition 3.6) and introduce the related optimization objective (Definition 3.7 & 3.9). In a nutshell, the intuition is to use a minimal set of representative training examples to encapsulate test-time decisions. Under mild conditions, we show the solution would exist and be unique (Proposition 3.10).
3.2 Understanding Decision-Making with a Subset of Offline Data
In order to perform accountable decision-making, we construct our approach upon the foundation of the example-based explanation framework [6, 7], which motivates us to define the offline decision dataset  $\mathcal{D}$  as the Decision Corpus, and introduce the concept of Corpus Subset that is composed of a representative subset of the decision corpus. These corpus subsets will be utilized for comprehending control decisions. Formally:
Definition 3.1 (Corpus Subset). A Corpus Subset  $\mathcal{C}$  is defined as a subset of the decision corpus  $\mathcal{D}$ , indexed by  $[C] := \{1,2,\dots,C\}$  — the natural numbers between 1 and  $C$ .
$$
\mathcal {C} = \left\{\left(o ^ {c}, a ^ {c}, h ^ {c}, v ^ {c}\right) \in \mathcal {D} \mid c \in [ C ] \right\}. \tag {2}
$$
In the equation above, the transition tuple  $(o^c, a^c, h^c, v^c)$  can be obtained from the decision corpus. We absorb both superscripts  $(i)$  and subscripts  $(t)$  to the index  $(c)$  for the conciseness of notions. Later in our work, we use subscripts  $t$  to denote the control time step. The variable  $v^c \in \mathcal{V} \subseteq \mathbb{R}$  represents a performance metric that is defined based on the user-specified decision objective. e.g.,  $v^c$  may correspond to an ongoing cumulative return, immediate reward, or a risk-sensitive measure.
Our objective is to understand the decision-making process for control time rollouts with the decision corpus. To do this, a naive approach is to reconstruct the control time examples with the corpus subset  $(o_{t},a_{t},h_{t}) = \sum_{c = 1}^{C}w^{c}(o^{c},a^{c},h^{c})$ , where  $w^{c}\in [0,1]$  and  $\sum_{c = 1}^{C}w^{c} = 1$ . Then the corresponding performance measures  $v^{c}$  can be used to estimate  $v_{t}|a_{t}$  — the control time return of a given decision  $a_{t}$ . Nonetheless, there are critical issues associated with this straightforward method:
1. Defining a distance metric for the joint observation-action-history space presents a challenge. This is not only due to the fact that various historical observations are defined over different spaces but also because the Euclidean distance fails to capture the structural information effectively.
2. The conversion from the observation-action-history space to the performance measure does not necessarily correspond to a linear mapping:  $v_{t} \neq \sum_{c=1}^{C} w^{c} v^{c}$ .
To address those difficulties, we propose an alternative approach called Accountable Offline Controller (AOC) that executes traceable offline control leveraging certain examples as decision corpus.
3.3 Linear Belief Learning for Accountable Offline Control
To alleviate the above problems, we first introduce a specific type of belief function  $\pmb{b} : \mathcal{O} \times \mathcal{A} \times \mathcal{H} \mapsto \mathcal{B} \subseteq \mathbb{R}^{d_b}$ , which maps a joint observation-action-history input onto a  $d_b$ -dimensional belief space.
Property 3.2 (Belief Space Linearity). The performance measure function  $V: \mathcal{O} \times \mathcal{A} \times \mathcal{H} \mapsto \mathcal{V}$  can always be decomposed as  $V = l \circ b$ , where  $b: \mathcal{O} \times \mathcal{A} \times \mathcal{H} \mapsto \mathcal{B} \subseteq \mathbb{R}^{d_b}$  maps the joint observation-action-history space to a  $d_b$  dimensional belief variable  $b_t = b(o_t, a_t, h_t)$  and  $l: \mathcal{B} \mapsto \mathcal{V}$  is a linear function that maps the belief  $b_t$  to an output  $l(b_t) \in \mathcal{V}$ .
Remark 3.3. This property is often the case with prevailing neural network approximators, where the belief state is the last activated layer before the final linear output layer.
Remark 3.4. The belief mapping function maps a  $d_{o} + d_{a} + (d_{o} + d_{a} + 1) \cdot (t - 1)$  dimensional varied-length input into a fixed-length  $d_{b}$  dimensional output. Usually, we have  $1 = d_v < d_b \ll d_o + d_a + (d_o + d_a + 1) \cdot (t - 1)$ . It is important to highlight that calculating the distance in the belief space is feasible, whereas it is intractable in the joint original space.
Based on Property 3.2, the linear relationship between the belief variable and the performance measure makes belief space a better choice for example-based decision understanding — this is because linearly decomposing the belief space also breaks down the performance measure, hence can guide action selection towards optimizing the performance metric. To be specific, we have
$$
v \left(a _ {t} \mid o _ {t}, h _ {t}\right) = \boldsymbol {l} \circ \boldsymbol {b} \left(o _ {t}, a _ {t}, h _ {t}\right) = \boldsymbol {l} \circ \sum_ {c = 1} ^ {C} w ^ {c} \cdot \boldsymbol {b} \left(o ^ {c}, a ^ {c}, h ^ {c}\right) = \sum_ {c = 1} ^ {C} w ^ {c} \cdot \boldsymbol {l} \circ \boldsymbol {b} \left(o ^ {c}, a ^ {c}, h ^ {c}\right) = \sum_ {c = 1} ^ {C} w ^ {c} \cdot v ^ {c} \tag {3}
$$
Although we have formally defined of the Corpus Subset, its composition has not been explicitly specified. We will now present the fundamental principles and techniques employed to create a practical Corpus Subset.
3.4 Selection of Corpus Subset
As suggested by Equation (3), if a belief state can be expressed as a linear combination of examples from the offline dataset, the same weights applied to the linear decomposition of the belief space can be utilized to decompose and synthesize the outcome. More formally, we define the convex hull spanned by the image of the Corpus Subset under the mapping function  $\pmb{b}$ :
Definition 3.5 (Belief Corpus Subset). A Belief Corpus Subset  $\mathbf{b}(\mathcal{C})$  is defined by applying the belief function  $\mathbf{b}$  to the corpus subset  $\mathcal{C}$ , with the corresponding set of values denoted as  $v(\mathcal{C})$ :
$$
\boldsymbol {b} (\mathcal {C}), v (\mathcal {C}) = \left\{b ^ {c} = \boldsymbol {b} \left(o ^ {c}, a ^ {c}, h ^ {c}\right), v ^ {c} \mid \left(o ^ {c}, a ^ {c}, h ^ {c}, v ^ {c}\right) \in \mathcal {C} \right\} \subset \mathcal {B} \times \mathcal {V} \tag {4}
$$
Definition 3.6 (Belief Corpus Convex Hull). The Belief Corpus Convex Hull spanned by a corpus subset  $\mathcal{C}$  with the belief corpus  $b(\mathcal{C})$  is the convex set
$$
\mathcal {C B} (\mathcal {C}) = \left\{\sum_ {c = 1} ^ {C} w ^ {c} b ^ {c} \mid w ^ {c} \in [ 0, 1 ], b ^ {c} \in \boldsymbol {b} (\mathcal {C}), \forall c \in [ C ], \sum_ {c = 1} ^ {C} w ^ {c} = 1 \right\} \tag {5}
$$
Note that an accurate belief corpus subset decomposition for a specific control time belief value  $b_{t}$  may be unattainable when  $b_{t} \notin \mathcal{CB}(\mathcal{C})$ . In such situations, the optimal course of action is to identify the element  $\tilde{b}_{t}$  within  $\mathcal{CB}(\mathcal{C})$  that provides the closest approximation to  $b_{t}$ . Denoting the norm in the belief space as  $||\cdot||_{\mathcal{B}}$ , this equates to minimizing a corpus residual that is defined as the distance from the control time belief to the nearest point within the belief corpus convex hull:
Definition 3.7 (Belief Corpus Residual). The Belief Corpus Residual given a control time belief variable  $b_{t}$  and a belief corpus subset  $\mathcal{C}$  is given as
$$
r _ {\mathcal {C}} \left(b _ {t}\right) = \min  _ {\tilde {b} _ {t} \in \mathcal {C B} (\mathcal {C})} \left\| b _ {t} - \tilde {b} _ {t} \right\| _ {\mathcal {B}} \tag {6}
$$
In summary, when  $b_{t}$  resides within  $\mathcal{CB}(\mathcal{C})$ , it can be decomposed using the belief corpus subset according to  $b_{t} = \sum_{c=1}^{C} w^{c} b^{c}$ . Here, the weights can be understood as representing the similarity and importance in reconstructing the control time belief variable through the belief corpus subset; otherwise, the optimizer of Equation (6) can be employed as a proxy for  $b_{t}$ . Under such circumstances, the belief corpus residual acts as a gauge for the degree of extrapolation. Given those definitions, we use the following proposition to inform the choice of an appropriate corpus subset.
Proposition 3.8 (Estimation Error Bound for  $v(a_{t})$ , cf. Crabbé et al. [7]). Consider the belief variable  $b_{t} = \pmb{b}(o_{t}, a_{t}, h_{t})$  and  $\hat{b}$  the optimizer of Equation (6), the estimated value residual between  $l(b_{t})$  and  $l(\hat{b})$  is controlled by the corpus residual:
$$
\left| \left| \boldsymbol {l} (\hat {b}) - \boldsymbol {l} \left(b _ {t}\right) \right| \right| _ {\mathcal {V}} \leq \left| \left| \boldsymbol {l} \right| \right| _ {\mathrm {o p}} \cdot r _ {\mathcal {C}} \left(b _ {t}\right) \tag {7}
$$
where  $||\cdot ||_{\mathcal{V}}$  is a norm on  $\mathcal{V}$  and  $||l||_{\mathrm{op}} = \inf \left\{\lambda \in \mathbb{R}^{+}:||l(b)||_{\mathcal{V}}\leq \lambda ||b||_{\mathcal{B}}\right\}$  is the operator norm for the linear mapping.
Figure 1: We illustrate the concepts of Corpus Residual and Minimal Hull in a 2-dim example. The belief function  $\pmb{b}$  maps trajectories in the offline dataset (colors of blue and green denote that they may come from different behavior policies) onto the belief space. The control trajectory with different sampled actions at time  $t$  can also be mapped to the belief space. Some of those sampled action candidates (e.g.,  $a_{t}^{(1)}, a_{t}^{(2)}$ ) can be well-supported by a corpus subset, while in another case  $(a_{t}^{(3)})$  the corpus residual manifests the extrapolation error. Among the three plotted beliefs generated by candidate actions,  $b_{t}^{(3)}$  is out of the convex hulls that can be spanned by any decision corpus,  $b_{t}^{(2)}$  has the minimal corpus (green shaded area),  $b_{t}^{(1)}, b_{t}^{(2)}$  are also in a larger corpus hull (blue shaded area) but it is not ideal.
Proposition 3.8 indicates that minimizing the belief corpus residual also minimizes the estimation error. In order to approximate the belief variables in control time accurately, we propose to use the corpus subset that spans the minimum convex hull in the belief space.
Definition 3.9 (Minimal Hull and Minimal Corpus Subset). Denoting the corpus subsets whose belief convex hull contains a belief variable  $b_{t}$  by
$$
\mathcal {C} (b _ {t}) = \left\{\mathcal {C} \subseteq \mathcal {D} \middle | b _ {t} \in \mathcal {C B} (\mathcal {C}) \right\}
$$
Among those subsets, the one that contains maximally  $d_b + 1$  examples and has the smallest hypervolume forms the minimal hull is called the minimal corpus subset (w.r.t.  $b_t$ ), denoted by  $\tilde{\mathcal{C}}(b_t)$ :
$$
\tilde {\mathcal {C}} (b _ {t}) := \arg \min _ {\mathcal {C}} \int_ {\mathcal {C B} (\mathcal {C})} d V
$$
Proposition 3.10 (Existence and Uniqueness). Consider the belief variable  $b_{t} = \pmb{b}(o_{t}, a_{t}, h_{t})$ , if  $r_{\mathcal{C}}(b_{t}) \leq 0$  holds for  $\mathcal{C} = \mathcal{D}$ , then  $\tilde{\mathcal{C}}(b_{t})$  exists and the decomposition on the corresponding minimal convex hull exists and is unique.
In this work, we use the corpus subset that constructs the minimal (belief corpus convex) hull to represent the control time belief variable  $b_{t}$ . Those concepts are illustrated in Figure 1.
3.5 Accountable Offline Control with Belief Corpus
Given the linear relationship between the variables  $v$  and  $b$ , for a specific action  $a_{t}$ , we can estimate the corresponding  $\hat{v}(a_{t}) = l(\hat{b}) = \sum_{c=1}^{C} w^{c} v^{c}$ , where  $o_{t}$  and  $h_{t}$  are the control time observation and transition history respectively. Applying an arg-max type of operator to the above value estimator, we are able to select the action with the highest value that is supported by the belief corpus:
$$
\pi_ {\mathrm {A O C}} \left(o _ {t}, h _ {t}\right) = \arg \max  _ {a _ {t}} \hat {v} \left(a _ {t}\right), \tag {8}
$$
Flexibility of the Measurement Function: Plug-and-Play. As discussed earlier, the choice of measurement function  $\mathbf{V}$  depends on the objective of policy learning. For the low-dimensional tasks, we can instantiate  $\mathbf{V}$  using the Monte Carlo estimation of the cumulative return:  $v^{c} = \sum_{t = t_{c}}^{T}\gamma^{t - t_{c}}r_{t}$ , as the most straightforward choice in control tasks; for higher-dimensional tasks, the measurement metric can be the learned Q-value function; for other applications like when safety-constraints are also taken into consideration, the measurement function can be adapted accordingly.
3.6 Optimization Procedure
The optimization process in our proposed AOC method involves three main steps: optimizing the belief function and linear mapping, reconstructing the belief during control time rollout, and
Algorithm 1 The Accountable Offline Controller
Bached Dataset  $\mathcal{D} = \{o_t^i, a_t^i, r_t^i, o_{t+1}^i\}_{t=1,\dots,T}^{i=1,\dots,N}$ , control time observation  $o_t$ , control history  $h_t$ . Hyperparameters:  $K$  for the number of candidate actions,  $\epsilon$  for the threshold.
Accountable action  $a_{t}$  for control.
1. Optimize belief function  $\pmb{b}$  and linear mapping  $l$  according to Equation (9).
2. Sample  $K$  candidate actions uniformly in the action space:  $a_{t}^{k}\sim U(\mathcal{A}),k\in [K]$
# 3. Find the minimal hulls  $\tilde{\mathcal{C}}(b_t)$  for each action, based on  $\pmb{b}$  and  $b_t = \pmb{b}(o_t, a_t^k, h_t)$ .
# 4. Minimize corpus residuals for each action candidate  $a_{t}^{k}$  according to Equation (10).
# 5. Estimate the value of action  $a_t$  by  $\hat{v}(a_t | o_t, h_t) = \sum_{c=1}^{C} w^c v^c$ .
# 6. Outputs the candidate action that has the highest estimated value according to Equation (11).
performing sampling-based optimization on the action space. In this section, we detail each of these steps along with corresponding optimization objectives. First, the belief function  $\pmb{b}$  and the linear mapping  $l$  are learned by minimizing the difference between the predicted performance measure and the actual value  $v$  from the dataset  $\mathcal{D}$ :
$$
\boldsymbol {b}, \boldsymbol {l} = \arg \min  _ {\boldsymbol {b}, \boldsymbol {l}} \mathbb {E} _ {(o, a, h, v) \in \mathcal {D}} (v - \boldsymbol {l} \circ \boldsymbol {b} (o, a, h)) ^ {2} \tag {9}
$$
Next, with the observation  $o_t$  and trace  $h_t$ , any candidate action  $a_t$  corresponds to a belief  $b_t = \pmb{b}(o_t, a_t, h_t)$ , which can be reconstructed by the minimal corpus subset  $\bar{\mathcal{C}}(b_t)$ , according to:
$$
w ^ {c} = \arg \min  _ {w ^ {c}} | | \sum_ {c = 1} ^ {C} w ^ {c} \boldsymbol {b} \left(o ^ {c}, a ^ {c}, h ^ {c}\right) - b _ {t} | |, \text {s . t .} \sum_ {c = 1} ^ {C} w ^ {c} = 1, \left(o ^ {c}, a ^ {c}, h ^ {c}\right) \in \tilde {\mathcal {C}} (b _ {t}) \tag {10}
$$
Then the value of action  $a_{t}$  can be estimated by  $\hat{v}(a_{t}|o_{t},h_{t}) = \sum_{c=1}^{C}w^{c}v^{c}$ . Finally, we perform a sampling-based optimization on the action space, with the corpus residual taken into consideration for the pursuance of conservation. For each control step, we uniformly sample  $K$  actions in the action space  $a^{(k)}t\sim U(\mathcal{A})$  and then search for the action that corresponds to the maximal performance measure while constraining the belief corpus residual to be smaller than a threshold  $\epsilon \geq 0$ :