SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\underset {a _ {0} ^ {\prime}, \dots , a _ {T - 1} ^ {\prime}} {\text {m a x i m i z e}} o ^ {+} \left(\tau^ {\prime}\right) \quad \text {s u b j e c t t o} \quad s _ {0} ^ {\prime} = s _ {0} \text {a n d} \sum_ {t = 0} ^ {T - 1} \mathbf {1} \left[ a _ {t} \neq a _ {t} ^ {\prime} \right] \leq k, \tag {9}
$$
where  $a_0, \ldots, a_{T-1}$  are the observed actions. Unfortunately, using a reduction from the classic partition problem [26], the following theorem shows that we cannot hope to find the optimal action sequence in polynomial time:
Theorem 4. The problem defined by Eq. 9. is NP-Hard.
The proof of the theorem relies on a reduction from the partition problem [26], which is known to be NP-complete, to our problem, defined in Eq. 9. At a high-level, we map any instance of the partition problem to an instance of our problem, taking special care to construct a reward function and an observed action sequence, such that the optimal counterfactual outcome  $o^{+}(\tau^{*})$  takes a specific value if and only if there exists a valid partition for the original instance. The hardness result of Theorem 4 motivates our subsequent focus on the design of a method that always finds the optimal solution to our problem at the expense of a potentially higher runtime for some problem instances.
4 Finding the optimal counterfactual action sequence via A* search
To deal with the increased computational complexity of the problem, we develop an optimal search method based on the classic  $A^{*}$  algorithm [52], which we have found to be very efficient in practice. Our starting point is the observation that, the problem of Eq. 9 presents an optimal substructure, i.e., its optimal solution can be constructed by combining optimal solutions to smaller sub-problems. For an observed episode  $\tau$ , let  $V_{\tau}(s,l,t)$  be the maximum counterfactual reward that could have been achieved in a counterfactual episode where, at time  $t$ , the process is at a (counterfactual) state  $s$ , and there are so far  $l$  actions that have been different in comparison with the observed action sequence. Formally,
$$
V _ {\tau} (\boldsymbol {s}, l, t) = \max  _ {a _ {t ^ {\prime}} ^ {\prime}, \dots , a _ {T - 1} ^ {\prime}} \sum_ {t ^ {\prime} = t} ^ {T - 1} R \left(\boldsymbol {s} _ {t ^ {\prime}} ^ {\prime}, a _ {t ^ {\prime}} ^ {\prime}\right) \quad \text {s u b j e c t t o} \quad \boldsymbol {s} _ {t} ^ {\prime} = \boldsymbol {s} \text {a n d} \sum_ {t ^ {\prime} = t} ^ {T - 1} \mathbf {1} \left[ a _ {t ^ {\prime}} \neq a _ {t ^ {\prime}} ^ {\prime} \right] \leq k - l.
$$
5The supporting proofs of all Theorems, Lemmas and Propositions can be found in Appendix D.
(a) Search graph
(b) Heuristic function computation
Figure 1: Main components of our search method based on the  $\mathbf{A}^*$  algorithm. Panel (a) shows the search graph for a problem instance with  $|\mathcal{A}| = 2$ . Here, each box represents a node  $v = (s,l,t)$  of the graph, and each edge represents a counterfactual transition. Next to each edge, we include the action  $a\in \mathcal{A}$  causing the transition and the associated reward. Panel (b) shows the heuristic function computation, where the two axes represent a (continuous) state space  $\mathcal{S} = \mathbb{R}^2$  and the two levels on the z-axis correspond to differences in the (integer) values  $(l,t)$  and  $(l_a,t + 1)$ . Here, the blue squares correspond to the finite states in the anchor set  $\mathcal{S}_{\dagger}$  and  $(s_{a},l_{a}) = F_{\tau ,t}^{+}((s,l),a)$ .
Then, it is easy to see that the quantity  $V_{\tau}(s,l,t)$  can be given by the recursive function
$$
V _ {\tau} (\boldsymbol {s}, l, t) = \max  _ {a \in \mathcal {A}} \left\{R (\boldsymbol {s}, a) + V _ {\tau} \left(\boldsymbol {s} _ {a}, l _ {a}, t + 1\right) \right\}, \quad \text {f o r a l l} \boldsymbol {s} \in \mathcal {S}, l <   k \text {a n d} t <   T - 1, \tag {10}
$$
where  $(\pmb{s}_a, l_a) = F_{\tau, t}^+ ((\pmb{s}, l), a)$ . In the base case of  $l = k$  (i.e., all allowed action changes are already performed), we have  $V_{\tau}(\pmb{s}, k, t) = R(\pmb{s}, a_t) + V_{\tau}(\pmb{s}_{a_t}, l_{a_t}, t + 1)$  for all  $\pmb{s} \in S$  and  $t < T - 1$ , and  $V_{\tau}(\pmb{s}, k, T - 1) = R(\pmb{s}, a_{T - 1})$  for  $t = T - 1$ . Lastly, when  $t = T - 1$  and  $l < k$ , we have  $V_{\tau}(\pmb{s}, l, T - 1) = \max_{a \in \mathcal{A}} R(\pmb{s}, a)$  for all  $\pmb{s} \in S$ .
Given the optimal substructure of the problem, one may be tempted to employ a typical dynamic programming approach to compute the values  $V_{\tau}(s,l,t)$  in a bottom-up fashion. However, the complexity of the problem lies in the fact that, the states  $s$  are real-valued vectors whose exact values depend on the entire action sequence that led to them. Hence, to enumerate all the possible values that  $s$  might take, one has to enumerate all possible action sequences in the search space, which is equivalent to solving our problem with a brute force search. In what follows, we present our proposed method to find optimal solutions using the  $A^{*}$  algorithm, with the caveat that its runtime varies depending on the problem instance, and it can be equal to that of a brute force search in the worst case.
Casting the problem as graph search. We represent the solution space of our problem as a graph, where each node  $v$  corresponds to a tuple  $(s, l, t)$  with  $s \in S$ ,  $l \in [k]$  and  $t \in [T]$ . Every node  $v = (s, l, t)$  with  $l < k$  and  $t < T - 1$  has  $|\mathcal{A}|$  outgoing edges, each one associated with an action  $a \in \mathcal{A}$ , carrying a reward  $R(s, a)$ , and leading to a node  $v_{a} = (s_{a}, l_{a}, t + 1)$  such that  $(s_{a}, l_{a}) = F_{\tau, t}^{+}((s, l), a)$ . In the case of  $l = k$ , the node  $v$  has exactly one edge corresponding to the observed action  $a_{t}$  at time  $t$ . Lastly, when  $t = T - 1$ , the outgoing edge(s) lead(s) to a common node  $v_{T} = (s_{\emptyset}, k, T)$  which we call the goal node, and it has zero outgoing edges itself. Note that, the exact value of  $s_{\emptyset}$  is irrelevant, and we only include it for notational completeness.
Let  $s_0$  be the initial state of the observed episode. Then, it is easy to notice that, starting from the root node  $v_0 = (s_0, 0, 0)$ , the first elements of each node  $v_i$  on a path  $v_0, \ldots, v_i, \ldots, v_T$  form a sequence of counterfactual states, and the edges that connect those nodes are such that the corresponding counterfactual action sequence differs from the observed one in at most  $k$  actions. That said, the counterfactual outcome  $o^+(\tau) = \sum_{t=0}^{T-1} R(s_t', a_t')$  is expressed as the sum of the rewards associated with each edge in the path, and the problem defined by Eq. 9 is equivalent to finding the path of maximum total reward that starts from  $v_0$  and ends in  $v_T$ . Figure 1a illustrates the search graph for a simple instance of our problem. Unfortunately, since the states  $s$  are vectors of real values, even enumerating all the graph's nodes requires time exponential in the number of actions  $|\mathcal{A}|$ , which makes classic algorithms that search over the entire graph non-practical.
To address this challenge, we resort to the  $A^{*}$  algorithm, which performs a more efficient search over the graph by preferentially exploring only parts of it where we have prior information that they are
Algorithm 1: It computes upper bounds  $\hat{V}_{\tau}(s,l,t)$  for the values  $V_{\tau}(s,l,t)$
Input: States  $\mathcal{S}$ , actions  $\mathcal{A}$ , observed action sequence  $\{a_{t}\}_{t=0}^{t=T-1}$ , horizon  $T$ , transition function  $F_{\tau,t}^{+}$ , reward function  $R$ , constraint  $k$ , anchor set  $\mathcal{S}_{\dagger}$ .  
Initialize:  $\hat{V}_{\tau}(\boldsymbol{s},l,T-1) \leftarrow \max_{a \in \mathcal{A}} R(\boldsymbol{s},a)$ ,  $s \in \mathcal{S}_{\dagger}$ ,  $l = 0, \dots, k-1$ .  
 $\hat{V}_{\tau}(\boldsymbol{s},k,T-1) \leftarrow R(\boldsymbol{s},a_{t})$ ,  $s \in \mathcal{S}_{\dagger}$ .  
for  $t = T-2, \dots, 0$  do  
for  $l = k, \dots, 0$  do  
available_actions)  $\leftarrow a_{t}$  if  $l = k$  else  $\mathcal{A}$   
for  $\boldsymbol{s} \in \mathcal{S}_{\dagger}$  do  
bounds  $\leftarrow \emptyset$   
for  $a \in available_actions$  do  
 $s_{a},l_{a} \leftarrow F_{\tau,t}^{+}((s,l),a)$ ; /* Get the min bound for  $V_{\tau}(s_{a},l_{a},t+1)$ $V_{a} \leftarrow \min_{s_{\dagger} \in S_{\dagger}} \{\hat{V}_{\tau}(s_{\dagger},l_{a},t+1) + L_{t+1} \| s_{\dagger} - s_{a} \|\};$  based on  $\mathcal{S}_{\dagger} * /$  bounds  $\leftarrow$  bounds  $\cup$ $\{R(s,a) + V_{a}\}$   
end  
 $\hat{V}_{\tau}(s,l,t) \leftarrow \max(bounds);$  /* Get the max bound over the actions */  
end  
return  $\hat{V}_{\tau}(s,l,t)$  for all  $s \in S_{\dagger}$ ,  $l \in [k]$ ,  $t \in [T-1]$
more likely to lead to paths of higher total reward. Concretely, the algorithm proceeds iteratively and maintains a queue of nodes to visit, initialized to contain only the root node  $v_{0}$ . Then, at each step, it selects one node from the queue, and it retrieves all its children nodes in the graph which are subsequently added to the queue. It terminates when the node being visited is the goal node  $v_{T}$ . Refer to Algorithm 2 in Appendix E for a pseudocode implementation of the  $A^{*}$  algorithm.
The key element of the  $A^{*}$  algorithm is the criterion based on which it selects which node from the queue to visit next. Let  $v_{i} = (s_{i},l_{i},t)$  be a candidate node in the queue and  $r_{v_i}$  be the total reward of the path that the algorithm has followed so far to reach from  $v_{0}$  to  $v_{i}$ . Then, the  $A^{*}$  algorithm visits next the node  $v_{i}$  that maximizes the sum  $r_{v_i} + \hat{V}_{\tau}(s_i,l_i,t)$ , where  $\hat{V}_{\tau}$  is a heuristic function that aims to estimate the maximum reward that can be achieved via any path starting from  $v_{i} = (s_{i},l_{i},t)$  and ending in the goal node  $v_{T}$ , i.e., it gives an estimate for the quantity  $V_{\tau}(s_i,l_i,t)$ . Intuitively, the heuristic function can be thought of as an "eye into the future" of the graph search, that guides the algorithm towards nodes that are more likely to lead to the optimal solution and the algorithm's performance depends on the quality of the approximation of  $V_{\tau}(s_i,l_i,t)$  by  $\hat{V}_{\tau}(s_i,l_i,t)$ . Next, we will look for a heuristic function that satisfies consistency<sup>6</sup> to guarantee that the  $A^{*}$  algorithm as described above returns the optimal solution upon termination [52].
Computing a consistent heuristic function. We first propose an algorithm that computes the function's values  $\hat{V}_{\tau}(s,l,t)$  for a finite set of points such that  $l\in [k],t\in [T - 1],s\in S_{\dagger}\subset S,$  where  $S_{\dagger}$  is a pre-defined finite set of states—an anchor set—whose construction we discuss later. Then, based on the Lipschitz-continuity of the SCM  $\mathcal{C}$ , we show that these computed values of  $\hat{V}_{\tau}$  are valid upper bounds of the corresponding values  $V_{\tau}(s,l,t)$  and we expand the definition of the heuristic function  $\hat{V}_{\tau}$  over all  $s\in S$  by expressing it in terms of those upper bounds. Finally, we prove that the function resulting from the aforementioned procedure is consistent.
To compute the upper bounds  $\hat{V}_{\tau}$ , we exploit the observation that the values  $V_{\tau}(s,l,t)$  satisfy a form of Lipschitz-continuity, as stated in the following Lemma.
Lemma 5. Let  $\pmb{u}_t = g_S^{-1}(\pmb{s}_t, a_t, \pmb{s}_{t+1})$ ,  $K_{\pmb{u}_t} = \max_{a \in \mathcal{A}} K_{a,\pmb{u}_t}$ ,  $C = \max_{a \in \mathcal{A}} C_a$  and the sequence  $L_0, \ldots, L_{T-1} \in \mathbb{R}_+$  be such that  $L_{T-1} = C$  and  $L_t = C + L_{t+1} K_{\pmb{u}_t}$  for  $t \in [T-2]$ . Then, it holds that  $|V_\tau(\pmb{s}, l, t) - V_\tau(\pmb{s}', l, t)| \leq L_t \| \pmb{s} - \pmb{s}' \|$ , for all  $t \in [T-1]$ ,  $l \in [k]$  and  $\pmb{s}, \pmb{s}' \in S$ .
Based on this observation, our algorithm proceeds in a bottom-up fashion and computes valid upper bounds of the values  $V_{\tau}(s,l,t)$  for all  $l \in [k]$ ,  $t \in [T - 1]$  and  $s$  in the anchor set  $S_{\dagger}$ . To get the
A heuristic function  $\hat{V}_{\tau}$  is consistent iff, for nodes  $v = (s,l,t)$ ,  $v_{a} = (s_{a},l_{a},t + 1)$  connected with an edge associated with action  $a$ , it satisfies  $\hat{V}_{\tau}(s,l,t) \geq R(s,a) + \hat{V}_{\tau}(s_{a},l_{a},t + 1)$  [53].
intuition, assume that, for a given  $t$ , the values  $\hat{V}_{\tau}(s,l,t + 1)$  are already computed for all  $s \in S_{\dagger}$ ,  $l \in [k]$ , and they are indeed valid upper bounds of the corresponding  $V_{\tau}(s,l,t + 1)$ . Then, let  $(s_a,l_a) = F_{\tau ,t}^{+}((s,l),a)$  for some  $s \in S_{\dagger}$  and  $l \in [k]$ . Since  $s_a$  itself may not belong to the finite anchor set  $S_{\dagger}$ , the algorithm uses the values  $\hat{V}_{\tau}(s_{\dagger},l_a,t + 1)$  of all anchors  $s_{\dagger} \in S_{\dagger}$  in combination with their distance to  $s_a$ , and it sets the value of  $\hat{V}_{\tau}(s,l,t)$  in way that it is also guaranteed to be a (maximally tight) upper bound of  $V_{\tau}(s,l,t)$ . Figure 1b illustrates the above operation. Algorithm 1 summarizes the overall procedure, which is guaranteed to return upper bounds, as shown by the following proposition:
Proposition 6. For all  $s \in S_{\dagger}$ ,  $l \in [k]$ ,  $t \in [T - 1]$ , it holds that  $\hat{V}_{\tau}(s, l, t) \geq V_{\tau}(s, l, t)$ , where  $\hat{V}_{\tau}(s, l, t)$  are the values of the heuristic function computed by Algorithm 1.
Next, we use the values  $\hat{V}_{\tau}(s,l,t)$  computed by Algorithm 1 to expand the definition of  $\hat{V}_{\tau}$  over the entire domain as follows. For some  $s\in S$ ,  $a\in \mathcal{A}$ , let  $(s_a,l_a) = F_{\tau ,t}^{+}((s,l),a)$ , then, we have that
$$
\hat {V} _ {\tau} (\boldsymbol {s}, l, t) = \left\{ \begin{array}{l l} 0 & t = T \\ \max  _ {a \in \mathcal {A} ^ {\prime}} R (\boldsymbol {s}, a) & t = T - 1 \\ \max  _ {a \in \mathcal {A} ^ {\prime}} \left\{R (\boldsymbol {s}, a) + \min  _ {\boldsymbol {s} _ {\dagger} \in \mathcal {S} _ {\dagger}} \left\{\hat {V} _ {\tau} \left(\boldsymbol {s} _ {\dagger}, l _ {a}, t + 1\right) + L _ {t + 1} \| \boldsymbol {s} _ {\dagger} - \boldsymbol {s} _ {a} \| \right\} \right\} & \text {o t h e r w i s e ,} \end{array} \right. \tag {11}
$$
where  $\mathcal{A}' = \{a_t\}$  for  $l = k$  and  $\mathcal{A}' = \mathcal{A}$  for  $l < k$ . Finally, the following theorem shows that the resulting heuristic function  $\hat{V}_{\tau}$  is consistent:
Theorem 7. For any nodes  $v = (s, l, t)$ ,  $v_{a} = (s_{a}, l_{a}, t + 1)$  with  $t < T - 1$  connected with an edge associated with action  $a$ , it holds that  $\hat{V}_{\tau}(s, l, t) \geq R(s, a) + \hat{V}_{\tau}(s_{a}, l_{a}, t + 1)$ . Moreover, for any node  $v = (s, l, T - 1)$  and edge connecting it to the goal node  $v_{T} = (s_{\emptyset}, k, T)$ , it holds that  $\hat{V}_{\tau}(s, l, T - 1) \geq R(s, a) + \hat{V}_{\tau}(s_{\emptyset}, k, T)$ .
Kick-starting the heuristic function computation with Monte Carlo anchor sets. For any  $s \notin S_{\dagger}$ , whenever we compute  $\hat{V}_{\tau}(s,l,t)$  using Eq. 11, the resulting value is set based on the value  $\hat{V}_{\tau}(s_{\dagger},l_a,t + 1)$  of some anchor  $s_{\dagger}$ , increased by a penalty term  $L_{t + 1}\| s_{\dagger} - s_a\|$ . Intuitively, this allows us to think of the heuristic function  $\hat{V}_{\tau}$  as an upper bound of the function  $V_{\tau}$  whose looseness depends on the magnitude of the penalty terms encountered during the execution of Algorithm 1 and each subsequent evaluation of Eq. 11. To speed up the  $A^{*}$  algorithm, note that, ideally, one would want all penalty terms to be zero, i.e., an anchor set that includes all the states  $s$  of the nodes  $v = (s,l,t)$  that are going to appear in the search graph. However, as discussed in the beginning of Sec. 4, an enumeration of those states requires a runtime exponential in the number of actions.
To address this issue, we introduce a Monte Carlo simulation technique that adds to the anchor set the observed states  $\{s_0,\dots ,s_{T - 1}\}$  and all unique states  $\{s_0^{\prime},\ldots ,s_{T - 1}^{\prime}\}$  resulting by  $M$  randomly sampled counterfactual action sequences  $a_0^\prime ,\ldots ,a_{T - 1}^\prime$  . Specifically, for each action sequence, we first sample a number  $k^{\prime}$  of actions to be changed and what those actions are going to be, both uniformly at random from  $\{1,\dots ,k\}$  and  $\mathcal{A}^{k^{\prime}}$  , respectively. Then, we sample from  $\{0,\dots ,T - 1\}$  the  $k^{\prime}$  time steps where the changes take place, with each time step  $t$  having a probability  $L_{t} / \sum_{t^{\prime}}L_{t^{\prime}}$  to be selected. This biases the sampling towards earlier time steps, where the penalty terms are larger due to the higher Lipschitz constants. As we will see in the next section, this approach works well in practice, and it allows us to control the runtime of the  $A^{*}$  algorithm by appropriately adjusting the number of samples  $M$  . We experiment with additional anchor set selection strategies in Appendix F.
5 Experiments using clinical sepsis management data
Experimental setup. To evaluate our method, we use real patient data from MIMIC-III [54], a freely accessible critical care dataset commonly used in reinforcement learning for healthcare [6, 55-57]. We follow the preprocessing steps of Komorowski et al. [6] to identify a cohort of 20,926 patients treated for sepsis [58]. Each patient record contains vital signs and administered treatment information in time steps of 4-hour intervals. As an additional preprocessing step, we discard patient records whose associated time horizon  $T$  is shorter than 10, resulting in a final dataset of 15,992 patients with horizons between 10 and 20.
(a) Efficiency vs.  $L_{h}$
(b) Efficiency vs.  $M$
Figure 2: Computational efficiency of our method under different configurations, as measured by the effective branching factor (pink-left axis) and the runtime of the  $A^{*}$  algorithm (green-right axis). In Panel (a), we set  $M = 2000$  and  $k = 3$ . In Panel (b), we set  $L_{h} = 1.0$  and  $k = 3$ . In Panel (c), we set  $L_{h} = 1.0$  and  $M = 2000$ . In all panels, we set  $L_{\phi} = 0.1$  and error bars indicate  $95\%$  confidence intervals over 200 executions of the  $A^{*}$  algorithm for 200 patients with horizon  $T = 12$ .
(c) Efficiency vs.  $k$
To form our state space  $\mathcal{S} = \mathbb{R}^D$ , we use  $D = 13$  features. Four of these features are demographic or contextual and thus we always set their counterfactual values to the observed ones. The remaining  $\tilde{D} = 9$  features are time-varying and include the SOFA score [59]—a standardized score of organ failure rate—along with eight vital signs that are required for its calculation. Since SOFA scores positively correlate with patient mortality [60], we assume that each  $s \in S$  gives a reward  $R(s)$  equal to the negation of its SOFA value. Here, it is easy to see that this reward function is just a projection of  $s$ , therefore, it is Lipschitz continuous with constant  $C_a = 1$  for all  $a \in \mathcal{A}$ . Following related work [6, 55, 57], we consider an action space  $\mathcal{A}$  that consists of 25 actions, which correspond to  $5 \times 5$  levels of administered vasopressors and intravenous fluids. Refer to Appendix G for additional details on the features and actions.
To model the transition dynamics of the time-varying features, we consider an SCM  $\mathcal{C}$  whose transition mechanism takes a location-scale form  $g_{S}(\pmb{S}_{t}, A_{t}, \pmb{U}_{t}) = h(\pmb{S}_{t}, A_{t}) + \phi(\pmb{S}_{t}, A_{t}) \odot \pmb{U}_{t}$ , where  $h, \phi: \mathcal{S} \times \mathcal{A} \to \mathbb{R}^{\tilde{D}}$ , and  $\odot$  denotes the element-wise multiplication [22, 24]. Notably, this model is element-wise bijective and hence it is counterfactually identifiable, as shown in Section 2. Moreover, we use neural networks to model the location and scale functions  $h$  and  $\phi$  and enforce their Lipschitz constants to be  $L_{h}$  and  $L_{\phi}$ , respectively. This results in a Lipschitz continuous SCM  $\mathcal{C}$  with  $K_{a,\pmb{u}} = L_{h} + L_{\phi} \max_{i} |u_{i}|$ . Further, we assume that the noise variable  $\pmb{U}_{t}$  follows a multivariate Gaussian distribution with zero mean and allow its covariance matrix to be a (trainable) parameter.
We jointly train the weights of the networks  $h$  and  $\phi$  and the covariance matrix of the noise prior on the observed patient transitions using stochastic gradient descent with the negative log-likelihood of each transition as a loss. In our experiments, if not specified otherwise, we use an SCM with Lipschitz constants  $L_{h} = 1.0$ ,  $L_{\phi} = 0.1$  that achieves a log-likelihood only  $6\%$  lower to that of the best model trained without any Lipschitz constraint. Refer to Appendix G for additional details on the network architectures, the training procedure and the way we enforce Lipschitz continuity. $^{7}$
Results. We start by evaluating the computational efficiency of our method against (i) the Lipschitz constant of the location network  $L_{h}$ , (ii) the number of Monte Carlo samples  $M$  used to generate the anchor set  $S_{\dagger}$ , and (iii) the number of actions  $k$  that can differ from the observed ones. We measure efficiency using running time and the effective branching factor (EBF) [52]. The EBF is defined as a real number  $b \geq 1$  such that the number of nodes expanded by  $A^{*}$  is equal to  $1 + b + b^{2} + \dots + b^{T}$ , where  $T$  is the horizon, and values close to 1 indicate that the heuristic function is the most efficient in guiding the search. Figure 2 summarizes the results, which show that our method maintains overall a fairly low running time that decreases with the number of Monte Carlo samples  $M$  used for the generation of the anchor set and increases with the Lipschitz constant  $L_{h}$  and the number of action changes  $k$ . That may not come as a surprise since, as  $L_{h}$  increases, the heuristic function becomes more loose, and as  $k$  increases, the size of the search space increases exponentially. To put things in perspective, for a problem instance with  $L_{h} = 1.0$ ,  $k = 3$  and horizon  $T = 12$ , the  $A^{*}$  search led by our heuristic function is effectively equivalent to an exhaustive search over a full tree with
<sup>7</sup>All experiments were performed using an internal cluster of machines equipped with 16 Intel(R) Xeon(R) 3.20GHz CPU cores, 512GBs of memory and 2 NVIDIA A40 48GB GPUs.
(a) Average  $c / f$  improvement vs.  $k$
Figure 3: Retrospective analysis of patients' episodes. Panel (a) shows the average counterfactual improvement as a function of  $k$  for a set of 200 patients with horizon  $T = 12$ , where error bars indicate  $95\%$  confidence intervals. Panel (b) shows the distribution of counterfactual improvement across all patients for  $k = 3$ , where the dashed vertical line indicates the median. Panel (c) shows the observed (solid) and counterfactual (dashed) SOFA score across time for a patient who presents a  $19.9\%$  counterfactual improvement when  $k = 3$ . Upward (downward) arrows indicate action changes that suggest a higher (lower) dosage of vasopressors (V) and fluids (F). In all panels, we set  $M = 2000$ .
(b) Distribution of c/f improvement
(c) Observed vs. c/f episode
$2.1^{12} \approx 7,355$  leaves while the corresponding search space of our problem consists of more than 3 million action sequences—more than 3 million paths to reach from the root node to the goal node.