Title: Finding Counterfactually Optimal Action Sequences in Continuous State Spaces
Venue: NeurIPS 2023
Accepted: yes
Keywords: Counterfactual reasoning

Markov decision process

Structural causal model

A* search
Abstract: Whenever a clinician reflects on the efficacy of a sequence of treatment decisions for a patient, they may try to identify critical time steps where, had they made different decisions, the patient's health would have improved. While recent methods at the intersection of causal inference and reinforcement learning promise to aid human experts, as the clinician above, to *retrospectively* analyze sequential decision making processes, they have focused on environments with finitely many discrete states. However, in many practical applications, the state of the environment is inherently continuous in nature. In this paper, we aim to fill this gap. We start by formally characterizing a sequence of discrete actions and continuous states using finite horizon Markov decision processes and a broad class of bijective structural causal models. Building upon this characterization, we formalize the problem of finding counterfactually optimal action sequences and show that, in general, we cannot expect to solve it in polynomial time. Then, we develop a search method based on the A* algorithm that, under a natural form of Lipschitz continuity of the environment’s dynamics, is guaranteed to return the optimal solution to the problem. Experiments on real clinical data show that our method is very efficient in practice, and it has the potential to offer interesting insights for sequential decision making tasks.

Finding Counterfactually Optimal Action Sequences in Continuous State Spaces

Abstract

Whenever a clinician reflects on the efficacy of a sequence of treatment decisions for a patient, they may try to identify critical time steps where, had they made different decisions, the patient's health would have improved. While recent methods at the intersection of causal inference and reinforcement learning promise to aid human experts, as the clinician above, to retrospectively analyze sequential decision making processes, they have focused on environments with finitely many discrete states. However, in many practical applications, the state of the environment is inherently continuous in nature. In this paper, we aim to fill this gap. We start by formally characterizing a sequence of discrete actions and continuous states using finite horizon Markov decision processes and a broad class of bijective structural causal models. Building upon this characterization, we formalize the problem of finding counterfactually optimal action sequences and show that, in general, we cannot expect to solve it in polynomial time. Then, we develop a search method based on the A ∗ algorithm that, under a natural form of Lipschitz continuity of the environment's dynamics, is guaranteed to return the optimal solution to the problem. Experiments on real clinical data show that our method is very efficient in practice, and it has the potential to offer interesting insights for sequential decision making tasks.

1 Introduction

Had the chess player moved the king one round later, would they have avoided losing the game? Had the physician administered antibiotics one day earlier, would the patient have recovered? The process of mentally simulating alternative worlds where events of the past play out differently than they did in reality is known as counterfactual reasoning [1]. Thoughts of this type are a common by-product of human decisions and they are tightly connected to the way we attribute causality and responsibility to events and others' actions [2]. The last decade has seen a rapid development of reinforcement learning agents, presenting (close to) human-level performance in a variety of sequential decision making tasks, such as gaming [3, 4], autonomous driving [5] and clinical decision support [6, 7]. In conjunction with the substantial progress made in the field of causal inference [8, 9], this has led to a growing interest in machine learning methods that employ elements of counterfactual reasoning to improve or to retrospectively analyze decisions in sequential settings [10-15].

In the context of reinforcement learning, sequential decision making is typically modeled using Markov Decision Processes (MDPs) [16]. Here, we consider MDPs with a finite horizon where each episode ( i.e. , each sequence of decisions) consists of a finite number of time steps. As an example, consider a clinician treating a patient in an intensive care unit (ICU). At each time step, the clinician observes the current state of the environment ( e.g. , the patient's vital signs) and they choose among a set of potential actions ( e.g. , standardized dosages of a drug). Consequently, the chosen action causes the environment to transition (stochastically) into a new state, and the clinician earns a reward ( e.g. ,

satisfaction inversely proportional to the patient's severity). The process repeats until the horizon is met and the goal of the clinician is to maximize the total reward.

In this work, our goal is to aid the retrospective analysis of individual episodes as the example above. For each episode, we aim to find an action sequence that differs slightly from the one taken in reality but, under the circumstances of that particular episode, would have led to a higher counterfactual reward. In our example above, assume that the patient's condition does not improve after a certain period of time. A counterfactually optimal action sequence could highlight to the clinician a small set of time steps in the treatment process where, had they administered different drug dosages, the patient's severity would have been lower. In turn, a manual inspection of those time steps could provide insights to the clinician about potential ways to improve their treatment policy.

To infer how a particular episode would have evolved under a different action sequence than the one taken in reality, one needs to represent the stochastic state transitions of the environment using a structural causal model (SCM) [8, 17]. This has been a key aspect of a line of work at the intersection of counterfactual reasoning and reinforcement learning, which has focused on methods to either design better policies using offline data [10, 12] or to retrospectively analyze individual episodes [11, 13]. Therein, the work most closely related to ours is by Tsirtsis et al. [13], which introduces a method to compute counterfactually optimal action sequences in MDPs with discrete states and actions using a Gumbel-Max SCM to model the environment dynamics [11]. However, in many practical applications, such as in critical care, the state of the environment is inherently continuous in nature [18]. In our work, we aim to fill this gap by designing a method to compute counterfactually optimal action sequences in MDPs with continuous states and discrete actions. Refer to Appendix A for a discussion of further related work and to Pearl [8] for an overview of the broad field of causality.

Our contributions. We start by formally characterizing sequential decision making processes with continuous states and discrete actions using finite horizon MDPs and a general class of bijective SCMs [19]. Notably, this class of SCMs includes multiple models introduced in the causal discovery literature [20-25]. Building on this characterization, we make the following contributions:

(i) We formalize the problem of finding a counterfactually optimal action sequence for a particular episode in environments with continuous states under the constraint that it differs from the observed action sequence in at most k actions.

(ii) We show that the above problem is NP-hard using a novel reduction from the classic partition problem [26]. This is in contrast with the computational complexity of the problem in environments with discrete states, which allows for polynomial time algorithms [13].

(iii) We develop a search method based on the A ∗ algorithm that, under a natural form of Lipschitz continuity of the environment's dynamics, is guaranteed to return the optimal solution to the problem upon termination.

Finally, we evaluate the performance and the qualitative insights of our method by performing a series of experiments using real patient data from critical care. 1

2 A causal model of sequential decision making processes

At each time step t ∈ [ T -1] := { 0 , 1 , . . . , T -1 } , where T is a time horizon, the decision making process is characterized by a D -dimensional vector state s t ∈ S = R D , an action a t ∈ A , where A is a finite set of N actions, and a reward R ( s t , a t ) ∈ R associated with each pair of states and actions. Moreover, given an episode of the decision making process, τ = { ( s t , a t ) } T -1 t =0 , the process's outcome o ( τ ) = ∑ t R ( s t , a t ) is given by the sum of the rewards. In the remainder, we will denote the elements of a vector s t as s t, 1 , . . . , s t,D . 2

Further, we characterize the dynamics of the decision making process using the framework of structural causal models (SCMs). In general, an SCM is consisted of four parts: (i) a set of endogenous variables (ii) a set of exogenous (noise) variables (iii) a set of structural equations assigning values to the endogenous variables, and (iv) a set of prior distributions characterizing the exogenous variables [8]. In our setting, the endogenous variables of the SCM C are the random variables representing the states S 0 , . . . , S T -1 and the actions A 0 , . . . , A T -1 . The action A t at time

step t is chosen based on the observed state S t and is given by a structural (policy) equation

where Z t ∈ Z is a vector-valued noise variable, to allow some level of stochasticity in the choice of the action, and its prior distribution P C ( Z t ) is characterized by a density function f C Z t . Similarly, the state S t +1 in the next time step is given by a structural (transition) equation

where U t ∈ U is a vector-valued noise variable with its prior distribution P C ( U t ) having a density function f C U t , and we refer to the function g S as the transition mechanism . Note that, in Eq. 2, the noise variables { U t } T -1 t =0 are mutually independent and, keeping the sequence of actions fixed, they are the only source of stochasticity in the dynamics of the environment. In other words, a sampled sequence of noise values { u t } T -1 t =0 and a fixed sequence of actions { a t } T -1 t =0 result into a single (deterministic) sequence of states { s t } T -1 t =0 . This implicitly assumes that the state transitions are stationary and there are no unobserved confounders. Figure 4 in Appendix B depicts the causal graph G corresponding to the SCM C defined above.

The above representation of sequential decision making using an SCM C is a more general reformulation of a Markov decision process, where a (stochastic) policy π ( a | s ) is entailed by Eq. 1, and the transition distribution ( i.e. , the conditional distribution of S t +1 | S t , A t ) is entailed by Eq. 2. Specifically, the conditional density function of S t +1 | S t , A t is given by

where do ( A t = a t ) denotes a (hard) intervention on the variable A t , whose value is set to a t . 3 Here, the first equality holds because S t +1 and A t are d-separated by S t in the sub-graph obtained from G after removing all outgoing edges of A t 4 and the second equality follows from Eq. 2.

Moreover, as argued elsewhere [11, 13], by using an SCM to represent sequential decision making, instead of a classic MDP, we can answer counterfactual questions. More specifically, assume that, at time step t , we observed the state S t = s t , we took action A t = a t and the next state was S t +1 = s t +1 . Retrospectively, we would like to know the probability that the state S t +1 would have been s ′ if, at time step t , we had been in a state s , and we had taken an action a , (generally) different from s t , a t . Using the SCM C , we can characterize this by a counterfactual transition density function

where f C | S t +1 = s t +1 , S t = s t ,A t = a t U t is the posterior distribution of the noise variable U t with support such that s t +1 = g S ( s t , a t , u ) .

In what follows, we will assume that the transition mechanism g S is continuous with respect to its last argument and the SCM C satisfies the following form of Lipschitz-continuity:

Definition 1. An SCM C is Lipschitz-continuous iff the transition mechanism g S and the reward R are Lipschitz-continuous with respect to their first argument, i.e. , for each a ∈ A , u ∈ U , there exists a Lipschitz constant K a, u ∈ R + such that, for any s , s ′ ∈ S , ‖ g S ( s , a, u ) -g S ( s ′ , a, u ) ‖ ≤ K a, u ‖ s -s ′ ‖ , and, for each a ∈ A , there exists a Lipschitz constant C a ∈ R + such that, for any s , s ′ ∈ S , | R ( s , a ) -R ( s ′ , a ) | ≤ C a ‖ s -s ′ ‖ . In both cases, ‖·‖ denotes the Euclidean distance.

Note that, although they are not phrased in causal terms, similar Lipschitz continuity assumptions for the environment dynamics are common in prior work analyzing the theoretical guarantees of reinforcement learning algorithms [27-35]. Moreover, for practical applications ( e.g. , in healthcare), this is a relatively mild assumption to make. Consider two patients whose vitals s and s ′ are similar

at a certain point in time, they receive the same treatment a , and every unobserved factor u that may affect their health is also the same. Intuitively, Definition 1 implies that their vitals will also evolve similarly in the immediate future, i.e. , the values g S ( s , a, u ) and g S ( s ′ , a, u ) will not differ dramatically. In this context, it is worth mentioning that, when the transition mechanism g S is modeled by a neural network, it is possible to control its Lipschitz constant during training, and penalizing high values can be seen as a regularization method [36, 37].

Further, we will focus on bijective SCMs [19], a fairly broad class of SCMs, which subsumes multiple models studied in the causal discovery literature, such as additive noise models [20], post-nonlinear causal models [21], location-scale noise models [22] and more complex models with neural network components [23-25].

Definition 2. An SCM C is bijective iff the transition mechanism g S is bijective with respect to its last argument, i.e. , there is a well-defined inverse function g -1 S : S ×A×S → U such that, for every combination of s t +1 , s t , a t , u t with s t +1 = g S ( s t , a t , u t ) , it holds that u t = g -1 S ( s t , a t , s t +1 ) .

Importantly, bijective SCMs allow for a more concise characterization of the counterfactual transition density given in Eq. 4. More specifically, after observing an event S t +1 = s t +1 , S t = s t , A t = a t , the value u t of the noise variable U t can only be such that u t = g -1 S ( s t , a t , s t +1 ) , i.e. , the posterior distribution of U t is a point mass and its density is given by

where 1 [ · ] denotes the indicator function. Then, for a given episode τ of the decision making process, we have that the (non-stationary) counterfactual transition density is given by

Since this density is also a point mass, the resulting counterfactual dynamics are purely deterministic. That means, under a bijective SCM , the answer to the question ' What would have been the state at time t +1 , had we been at state s and taken action a at time t , given that, in reality, we were at s t , we took a t and the environment transitioned to s t +1 ? ' is just given by s ′ = g S ( s , a, g -1 S ( s t , a t , s t +1 ) ) .

glyph[negationslash]

On the counterfactual identifiability of bijective SCMs. Very recently, Nasr-Esfahany and Kiciman [38] have shown that bijective SCMs are in general not counterfactually identifiable when the exogenous variable U t is multi-dimensional. In other words, even with access to an infinite amount of triplets ( s t , a t , s t +1 ) sampled from the true SCM C , it is always possible to find an SCM Mglyph[negationslash] = C with transition mechanism h S and distributions P M ( U t ) that entails the same transition distributions as C ( i.e. , it fits the observational data perfectly), but leads to different counterfactual predictions. Although our subsequent algorithmic results do not require the SCM C to be counterfactually identifiable, the subclass of bijective SCMs we will use in our experiments in Section 5 is counterfactually identifiable. The defining attribute of this subclass, which we refer to as element-wise bijective SCMs , is that the transition mechanism g S can be decoupled into D independent mechanisms g S,i such that S t +1 ,i = g S,i ( S t , A t , U t,i ) for i ∈ { 1 , . . . , D } . This implies S t +1 ,i ⊥ ⊥ U t,j | U t,i , S t , A t for j = i , however, U t,i , U t,j do not need to be independent. Informally, we have the following identifiability result (refer to Appendix C for a formal version of the theorem along with its proof, which follows a similar reasoning to proofs found in related work [12, 19]):

Theorem 3 (Informal) . Let C and M be two element-wise bijective SCMs such that their entailed transition distributions for S t +1 given any value of S t , A t are always identical. Then, all their counterfactual predictions based on an observed transition ( s t , a t , s t +1 ) will also be identical.

On the assumption of no unobserved confounding. The assumption that there are no hidden confounders is a frequent assumption made by work at the intersection of counterfactual reasoning and reinforcement learning [10-13] and, more broadly, in the causal inference literature [39-43]. That said, there is growing interest in developing off-policy methods for partially observable MPDs (POMDPs) that are robust to certain types of confounding [44-46], and in learning dynamic treatment regimes in sequential settings with non-Markovian structure [47, 48]. Moreover, there is a line of work focusing on the identification of counterfactual quantities in non-sequential confounded environments [49-51]. In that context, we consider the computation of (approximately) optimal counterfactual action sequences under confounding as a very interesting direction for future work.

3 Problem statement

Let τ be an observed episode of a decision making process whose dynamics are characterized by a Lipschitz-continuous bijective SCM. To characterize the counterfactual outcome that any alternative action sequence would have achieved under the circumstances of the particular episode, we build upon the formulation of Section 2, and we define a non-stationary counterfactual MDP M + = ( S + , A , F + τ,t , R + , T ) with deterministic transitions. Here, S + = S × [ T -1] is an enhanced state space such that each s + ∈ S + is a pair ( s , l ) indicating that the counterfactual episode would have been at state s ∈ S with l action changes already performed. Accordingly, R + is a reward function which takes the form R + (( s , l ) , a ) = R ( s , a ) for all ( s , l ) ∈ S + , a ∈ A , i.e. , it does not change depending on the number of action changes already performed. Finally, the time-dependent transition function F + τ,t : S + ×A→S + is defined as glyph[negationslash]

Intuitively, here we set the transition function according to the point mass of the counterfactual transition density given in Eq. 6, and we use the second coordinate to keep track of the changes that have been performed in comparison to the observed action sequence up to the time step t .

Now, given the initial state s 0 of the episode τ and any counterfactual action sequence { a ′ t } T -1 t =0 , we can compute the corresponding counterfactual episode τ ′ = { ( s ′ t , l t ) , a ′ t } T -1 t =0 . Its sequence of states is given recursively by

Then, similarly as in Tsirtsis et al. [13], our ultimate goal is to find the counterfactual action sequence { a ′ t } T -1 t =0 that, starting from the observed initial state s 0 , maximizes the counterfactual outcome subject to a constraint on the number of counterfactual actions that can differ from the observed ones, i.e. , glyph[negationslash]

where a 0 , . . . , 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: 5

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 + ( τ ∗ ) 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 τ , let V τ ( 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, glyph[negationslash]

(a) Search graph

(b) Heuristic function computation

S=s'

=

0

R(s')

S=So

R(s'）

s=S0

l=0

R(s,)

l=k

t=0

R(s,

t=T

S=

R(s",)

R(s",

Observed action

sequence

t=0t=1

S÷

ES

R²,la,t+ 1

R2,l,t

Then, it is easy to see that the quantity V τ ( s , l, t ) can be given by the recursive function

where ( s a , l a ) = F + τ,t (( s , l ) , a ) . In the base case of l = k ( i.e. , all allowed action changes are already performed), we have V τ ( s , k, t ) = R ( s , a t ) + V τ ( s a t , l a t , t +1) for all s ∈ S and t < T -1 , and V τ ( s , k, T -1) = R ( s , a T -1 ) for t = T -1 . Lastly, when t = T -1 and l < k , we have V τ ( s , l, T -1) = max a ∈A R ( s , a ) for all s ∈ S .

Given the optimal substructure of the problem, one may be tempted to employ a typical dynamic programming approach to compute the values V τ ( 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 ∈ S , l ∈ [ k ] and t ∈ [ T ] . Every node v = ( s , l, t ) with l < k and t < T -1 has |A| outgoing edges, each one associated with an action a ∈ 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 + τ,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 ∅ , k, T ) which we call the goal node , and it has zero outgoing edges itself. Note that, the exact value of s ∅ 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 , . . . , v i , . . . , 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 + ( τ ) = ∑ T -1 t =0 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 |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

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 + ˆ V τ ( s i , l i , t ) , where ˆ V τ 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 τ ( 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 τ ( s i , l i , t ) by ˆ V τ ( s i , l i , t ) . Next, we will look for a heuristic function that satisfies consistency 6 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 ˆ V τ ( s , l, t ) for a finite set of points such that l ∈ [ k ] , t ∈ [ T -1] , s ∈ S † ⊂ S , where S † 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 C , we show that these computed values of ˆ V τ are valid upper bounds of the corresponding values V τ ( s , l, t ) and we expand the definition of the heuristic function ˆ V τ over all s ∈ 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 ˆ V τ , we exploit the observation that the values V τ ( s , l, t ) satisfy a form of Lipschitz-continuity, as stated in the following Lemma.

Lemma 5. Let u t = g -1 S ( s t , a t , s t +1 ) , K u t = max a ∈A K a, u t , C = max a ∈A C a and the sequence L 0 , . . . , L T -1 ∈ R + be such that L T -1 = C and L t = C + L t +1 K u t for t ∈ [ T -2] . Then, it holds that | V τ ( s , l, t ) -V τ ( s ′ , l, t ) | ≤ L t ‖ s -s ′ ‖ , for all t ∈ [ T -1] , l ∈ [ k ] and s , s ′ ∈ S .

Based on this observation, our algorithm proceeds in a bottom-up fashion and computes valid upper bounds of the values V τ ( s , l, t ) for all l ∈ [ k ] , t ∈ [ T -1] and s in the anchor set S † . To get the

intuition, assume that, for a given t , the values ˆ V τ ( s , l, t +1) are already computed for all s ∈ S † , l ∈ [ k ] , and they are indeed valid upper bounds of the corresponding V τ ( s , l, t + 1) . Then, let ( s a , l a ) = F + τ,t (( s , l ) , a ) for some s ∈ S † and l ∈ [ k ] . Since s a itself may not belong to the finite anchor set S † , the algorithm uses the values ˆ V τ ( s † , l a , t +1) of all anchors s † ∈ S † in combination with their distance to s a , and it sets the value of ˆ V τ ( s , l, t ) in way that it is also guaranteed to be a (maximally tight) upper bound of V τ ( 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 ∈ S † , l ∈ [ k ] , t ∈ [ T -1] , it holds that ˆ V τ ( s , l, t ) ≥ V τ ( s , l, t ) , where ˆ V τ ( s , l, t ) are the values of the heuristic function computed by Algorithm 1.

Next, we use the values ˆ V τ ( s , l, t ) computed by Algorithm 1 to expand the definition of ˆ V τ over the entire domain as follows. For some s ∈ S , a ∈ A , let ( s a , l a ) = F + τ,t (( s , l ) , a ) , then, we have that

where A ′ = { a t } for l = k and A ′ = A for l < k . Finally, the following theorem shows that the resulting heuristic function ˆ V τ 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 ˆ V τ ( s , l, t ) ≥ R ( s , a ) + ˆ V τ ( 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 ∅ , k, T ) , it holds that ˆ V τ ( s , l, T -1) ≥ R ( s , a ) + ˆ V τ ( s ∅ , k, T ) .

Kick-starting the heuristic function computation with Monte Carlo anchor sets. For any s glyph[negationslash]∈ S † , whenever we compute ˆ V τ ( s , l, t ) using Eq. 11, the resulting value is set based on the value ˆ V τ ( s † , l a , t +1) of some anchor s † , increased by a penalty term L t +1 ‖ s † -s a ‖ . Intuitively, this allows us to think of the heuristic function ˆ V τ as an upper bound of the function V τ 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 , . . . , s T -1 } and all unique states { s ′ 0 , . . . , s ′ T -1 } resulting by M randomly sampled counterfactual action sequences a ′ 0 , . . . , a ′ T -1 . Specifically, for each action sequence, we first sample a number k ′ of actions to be changed and what those actions are going to be, both uniformly at random from { 1 , . . . , k } and A k ′ , respectively. Then, we sample from { 0 , . . . , T -1 } the k ′ time steps where the changes take place, with each time step t having a probability L t / ∑ t ′ L t ′ 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 .

0.8

0.9

1.0

1.1

1.2

1.3

L

h

1

.

5

2

.

0

2

.

5

EBF

0

20

40

60

A

∗

average runtime (s)

(a) Efficiency vs.

L

h

500

1000

1500

2000

2500

3000

M

1

.

9

2

.

0

2

.

1

2

.

2

EBF

10

20

A

∗

average runtime (s)

(b) Efficiency vs.

M

0

1

2

3

4

5

6

# of action changes,

k

1

2

3

EBF

0

200

400

600

A

∗

average runtime (s)

(c) Efficiency vs.

k

To form our state space S = 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 ˜ 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 ∈ 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 ∈ A . Following related work [6, 55, 57], we consider an action space A that consists of 25 actions, which correspond to 5 × 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 C whose transition mechanism takes a location-scale form g S ( S t , A t , U t ) = h ( S t , A t ) + φ ( S t , A t ) glyph[circledot] U t , where h, φ : S × A → R ˜ D , and glyph[circledot] 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 φ and enforce their Lipschitz constants to be L h and L φ , respectively. This results in a Lipschitz continuous SCM C with K a, u = L h + L φ max i | u i | . Further, we assume that the noise variable 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 φ 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 φ = 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 † , 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 ≥ 1 such that the number of nodes expanded by A ∗ is equal to 1 + b + b 2 + · · · + 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

0

1

2

3

4

5

6

# of action changes,

k

0%

5%

10%

Average c/f improvement

(a) Average c/f improvement vs. k

0%

10%

20%

30%

40%

Counterfactual improvement

0

1000

2000

Number of patients

20%

30%

0

50

100

!

(b) Distribution of c/f improvement

(c) Observed vs. c/f episode

2 . 1 12 ≈ 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.

Next, we investigate to what extent the counterfactual action sequences generated by our method would have led the patients in our dataset to better outcomes. For each patient, we measure their counterfactual improvement-the relative decrease in cumulative SOFA score between the counterfactual and the observed episode. Figures 3a and 3b summarize the results, which show that: (i) the average counterfactual improvement shows a diminishing increase as k increases; (ii) the median counterfactual improvement is only 5% , indicating that, the treatment choices made by the clinicians for most of the patients were close to optimal, even with the benefit of hindsight; and (iii) there are 176 patients for whom our method suggests that a different sequence of actions would have led to an outcome that is at least 15% better. That said, we view patients at the tail of the distribution as 'interesting cases' that should be deferred to domain experts for closer inspection, and we present one such example in Fig. 3c. In this example, our method suggests that, had the patient received an early higher dosage of intravenous fluids while some of the later administered fluids where replaced by vasopressors, their SOFA score would have been lower across time. Although we present this case as purely anecdotal, the counterfactual episode is plausible, since there are indications of decreased mortality when intravenous fluids are administered at the early stages of a septic shock [61].

6 Conclusions

In this paper, we have introduced the problem of finding counterfactually optimal action sequences in sequential decision making processes with continuous state dynamics. We showed that the problem is NP-hard and, to tackle it, we introduced a search method based on the A ∗ algorithm that is guaranteed to find the optimal solution, with the caveat that its runtime can vary depending on the problem instance. Lastly, using real clinical data, we have found that our method is very efficient in practice, and it has the potential to offer interesting insights to domain experts by highlighting episodes and time-steps of interest for further inspection.

Our work opens up many interesting avenues for future work. For example, it would be interesting to develop algorithms with approximation guarantees that run in polynomial time, at the expense of not achieving strict counterfactual optimality. Moreover, since the practicality of methods like ours relies on the assumption that the SCM describing the environment is accurate, it would be interesting to develop methods to learn SCMs that align with human domain knowledge. Finally, it would be interesting to validate our method using real datasets from other applications and carry out user studies in which the counterfactual action sequences found by our method are systematically evaluated by the human experts ( e.g. , clinicians) who took the observed actions.

Acknowledgements. Tsirtsis and Gomez-Rodriguez acknowledge support from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (grant agreement No. 945719).

References

Ruth MJ Byrne. Counterfactual thought. Annual review of psychology , 67:135-157, 2016.

David A Lagnado, Tobias Gerstenberg, and Ro'i Zultan. Causal responsibility and counterfactuals. Cognitive science , 37(6):1036-1073, 2013.

David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature , 529(7587):484-489, 2016.

Łukasz Kaiser, Mohammad Babaeizadeh, Piotr Miłos, Bła˙ zej Osi´ nski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In International Conference on Learning Representations , 2020.

B Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey. IEEE Transactions on Intelligent Transportation Systems , 23(6):4909-4926, 2021.

Matthieu Komorowski, Leo A Celi, Omar Badawi, Anthony C Gordon, and A Aldo Faisal. The artificial intelligence clinician learns optimal treatment strategies for sepsis in intensive care. Nature medicine , 24(11):1716-1720, 2018.

Chao Yu, Jiming Liu, Shamim Nemati, and Guosheng Yin. Reinforcement learning in healthcare: A survey. ACM Computing Surveys (CSUR) , 55(1):1-36, 2021.

Judea Pearl. Causality . Cambridge university press, 2009.

Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: foundations and learning algorithms . The MIT Press, 2017.

Lars Buesing, Theophane Weber, Yori Zwols, Nicolas Heess, Sebastien Racaniere, Arthur Guez, and Jean-Baptiste Lespiau. Woulda, coulda, shoulda: Counterfactually-guided policy search. In International Conference on Learning Representations , 2019.

Michael Oberst and David Sontag. Counterfactual off-policy evaluation with gumbel-max structural causal models. In International Conference on Machine Learning , pages 4881-4890. PMLR, 2019.

Chaochao Lu, Biwei Huang, Ke Wang, José Miguel Hernández-Lobato, Kun Zhang, and Bernhard Schölkopf. Sample-efficient reinforcement learning via counterfactual-based data augmentation. arXiv preprint arXiv:2012.09092 , 2020.

Stratis Tsirtsis, Abir De, and Manuel Gomez-Rodriguez. Counterfactual explanations in sequential decision making under uncertainty. Advances in Neural Information Processing Systems , 34:30127-30139, 2021.

Jonathan Richens, Rory Beard, and Daniel H Thompson. Counterfactual harm. In Advances in Neural Information Processing Systems , volume 35, pages 36350-36365, 2022.

Kimia Noorbakhsh and Manuel Gomez-Rodriguez. Counterfactual temporal point processes. Advances in Neural Information Processing Systems , 35:24810-24823, 2022.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Elias Bareinboim, Juan D Correa, Duligur Ibeling, and Thomas Icard. On pearl's hierarchy and the foundations of causal inference. In Probabilistic and causal inference: the works of judea pearl , pages 507-556. 2022.

Malcolm Elliott and Alysia Coventry. Critical care: the eight vital signs of patient monitoring. British Journal of Nursing , 21(10):621-625, 2012.

Arash Nasr-Esfahany, Mohammad Alizadeh, and Devavrat Shah. Counterfactual identifiability of bijective causal models. In Proceedings of the 40th International Conference on Machine Learning , ICML'23. JMLR.org, 2023.

Patrik Hoyer, Dominik Janzing, Joris M Mooij, Jonas Peters, and Bernhard Schölkopf. Nonlinear causal discovery with additive noise models. Advances in neural information processing systems , 21, 2008.

Kun Zhang and Aapo Hyvärinen. On the identifiability of the post-nonlinear causal model. In Proceedings of the Twenty-Fifth Conference on Uncertainty in Artificial Intelligence , UAI '09, page 647-655, Arlington, Virginia, USA, 2009. AUAI Press.

Alexander Immer, Christoph Schultheiss, Julia E Vogt, Bernhard Schölkopf, Peter Bühlmann, and Alexander Marx. On the identifiability and estimation of causal location-scale noise models. In International Conference on Machine Learning , pages 14316-14332. PMLR, 2023.

Nick Pawlowski, Daniel Coelho de Castro, and Ben Glocker. Deep structural causal models for tractable counterfactual inference. Advances in Neural Information Processing Systems , 33:857-869, 2020.

Ilyes Khemakhem, Ricardo Monti, Robert Leech, and Aapo Hyvarinen. Causal autoregressive flows. In International conference on artificial intelligence and statistics , pages 3520-3528. PMLR, 2021.

Pedro Sanchez and Sotirios A. Tsaftaris. Diffusion causal models for counterfactual estimation. In First Conference on Causal Learning and Reasoning , 2022.

Richard M Karp. Reducibility among combinatorial problems. In Complexity of computer computations , pages 85-103. Springer, 1972.

Dimitri Bertsekas. Convergence of discretization procedures in dynamic programming. IEEE Transactions on Automatic Control , 20(3):415-419, 1975.

Karl Hinderer. Lipschitz continuity of value functions in markovian decision processes. Mathematical Methods of Operations Research , 62:3-22, 2005.

Emmanuel Rachelson and Michail G Lagoudakis. On the locality of action domination in sequential decision making. 2010.

Jason Pazis and Ronald Parr. Pac optimal exploration in continuous space markov decision processes. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 27, pages 774-781, 2013.

Matteo Pirotta, Marcello Restelli, and Luca Bascetta. Policy gradient in lipschitz markov decision processes. Machine Learning , 100:255-283, 2015.

Felix Berkenkamp, Matteo Turchetta, Angela Schoellig, and Andreas Krause. Safe model-based reinforcement learning with stability guarantees. Advances in neural information processing systems , 30, 2017.

Kavosh Asadi, Dipendra Misra, and Michael Littman. Lipschitz continuity in model-based reinforcement learning. In International Conference on Machine Learning , pages 264-273. PMLR, 2018.

Ahmed Touati, Adrien Ali Taiga, and Marc G Bellemare. Zooming for efficient model-free reinforcement learning in metric spaces. arXiv preprint arXiv:2003.04069 , 2020.

Omer Gottesman, Kavosh Asadi, Cameron Allen, Sam Lobel, George Konidaris, and Michael Littman. Coarse-grained smoothness for rl in metric spaces. arXiv preprint arXiv:2110.12276 , 2021.

Cem Anil, James Lucas, and Roger Grosse. Sorting out lipschitz function approximation. In International Conference on Machine Learning , pages 291-301. PMLR, 2019.

Henry Gouk, Eibe Frank, Bernhard Pfahringer, and Michael J Cree. Regularisation of neural networks by enforcing lipschitz continuity. Machine Learning , 110:393-416, 2021.

Arash Nasr-Esfahany and Emre Kiciman. Counterfactual (non-)identifiability of learned structural causal models. arXiv preprint arXiv:2301.09031 , 2023.

Scott Sussex, Caroline Uhler, and Andreas Krause. Near-optimal multi-perturbation experimental design for causal structure learning. Advances in Neural Information Processing Systems , 34:777-788, 2021.

Ioana Bica, Daniel Jarrett, and Mihaela van der Schaar. Invariant causal imitation learning for generalizable policies. Advances in Neural Information Processing Systems , 34:3952-3964, 2021.

Chris Cundy, Aditya Grover, and Stefano Ermon. Bcd nets: Scalable variational approaches for bayesian causal discovery. Advances in Neural Information Processing Systems , 34:7095-7110, 2021.

Virginia Aglietti, Neil Dhir, Javier González, and Theodoros Damoulas. Dynamic causal bayesian optimization. Advances in Neural Information Processing Systems , 34:10549-10560, 2021.

Xinshi Chen, Haoran Sun, Caleb Ellington, Eric Xing, and Le Song. Multi-task learning of orderconsistent causal graphs. Advances in Neural Information Processing Systems , 34:11083-11095, 2021.

Guy Tennenholtz, Uri Shalit, and Shie Mannor. Off-policy evaluation in partially observable environments. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pages 10276-10283, 2020.

Hongseok Namkoong, Ramtin Keramati, Steve Yadlowsky, and Emma Brunskill. Off-policy policy evaluation for sequential decisions under unobserved confounding. Advances in Neural Information Processing Systems , 33:18819-18831, 2020.

Lingxiao Wang, Zhuoran Yang, and Zhaoran Wang. Provably efficient causal reinforcement learning with confounded observational data. Advances in Neural Information Processing Systems , 34:21164-21175, 2021.

Junzhe Zhang and Elias Bareinboim. Near-optimal reinforcement learning in dynamic treatment regimes. Advances in Neural Information Processing Systems , 32, 2019.

Junzhe Zhang and Elias Bareinboim. Designing optimal dynamic treatment regimes: A causal reinforcement learning approach. In International Conference on Machine Learning , pages 11012-11022. PMLR, 2020.

Ilya Shpitser and Eli Sherman. Identification of personalized effects associated with causal pathways. In Uncertainty in artificial intelligence: proceedings of the... conference. Conference on Uncertainty in Artificial Intelligence , volume 2018. NIH Public Access, 2018.

Juan Correa, Sanghack Lee, and Elias Bareinboim. Nested counterfactual identification from arbitrary surrogate experiments. Advances in Neural Information Processing Systems , 34:68566867, 2021.

Junzhe Zhang, Jin Tian, and Elias Bareinboim. Partial counterfactual identification from observational and experimental data. In International Conference on Machine Learning , pages 26548-26558. PMLR, 2022.

Stuart Russel, Peter Norvig, et al. Artificial intelligence: a modern approach , volume 256. Pearson Education Limited London, 2013.

Judea Pearl. Heuristics: intelligent search strategies for computer problem solving . AddisonWesley Longman Publishing Co., Inc., 1984.

Alistair EW Johnson, Tom J Pollard, Lu Shen, Li-wei H Lehman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G Mark. Mimic-iii, a freely accessible critical care database. Scientific data , 3(1):1-9, 2016.

Aniruddh Raghu, Matthieu Komorowski, Imran Ahmed, Leo Celi, Peter Szolovits, and Marzyeh Ghassemi. Deep reinforcement learning for sepsis treatment. arXiv preprint arXiv:1711.09602 , 2017.

Siqi Liu, Kay Choong See, Kee Yuan Ngiam, Leo Anthony Celi, Xingzhi Sun, and Mengling Feng. Reinforcement learning for clinical decision support in critical care: comprehensive review. Journal of medical Internet research , 22(7):e18477, 2020.

Taylor W Killian, Haoran Zhang, Jayakumar Subramanian, Mehdi Fatemi, and Marzyeh Ghassemi. An empirical study of representation learning for reinforcement learning in healthcare. arXiv preprint arXiv:2011.11235 , 2020.

Mervyn Singer, Clifford S Deutschman, Christopher Warren Seymour, Manu Shankar-Hari, Djillali Annane, Michael Bauer, Rinaldo Bellomo, Gordon R Bernard, Jean-Daniel Chiche, Craig M Coopersmith, et al. The third international consensus definitions for sepsis and septic shock (sepsis-3). Jama , 315(8):801-810, 2016.

Simon Lambden, Pierre Francois Laterre, Mitchell M Levy, and Bruno Francois. The sofa score-development, utility and challenges of accurate assessment in clinical trials. Critical Care , 23(1):1-9, 2019.

Flavio Lopes Ferreira, Daliana Peres Bota, Annette Bross, Christian Mélot, and Jean-Louis Vincent. Serial evaluation of the sofa score to predict outcome in critically ill patients. Jama , 286(14):1754-1758, 2001.

Jason Waechter, Anand Kumar, Stephen E Lapinsky, John Marshall, Peter Dodek, Yaseen Arabi, Joseph E Parrillo, R Phillip Dellinger, Allan Garland, Cooperative Antimicrobial Therapy of Septic Shock Database Research Group, et al. Interaction between fluids and vasoactive agents on mortality in septic shock: a multicenter, observational study. Critical care medicine , 42(10):2158-2168, 2014.

Lihong Li and Michael L Littman. Lazy approximation for solving continuous finite-horizon mdps. In AAAI , volume 5, pages 1175-1180, 2005.

Richard Dearden, Nicholas Meuleau, Richard Washington, and Zhengzhu Feng. Dynamic programming for structured continuous markov decision problems. In Twentieth Conference on Uncertainty in Artificial Intelligene (UAI-04) , 2004.

Prashan Madumal, Tim Miller, Liz Sonenberg, and Frank Vetere. Explainable reinforcement learning through a causal lens. In Proceedings of the AAAI conference on artificial intelligence , volume 34, pages 2493-2500, 2020.

Ioana Bica, Daniel Jarrett, Alihan Hüyük, and Mihaela van der Schaar. Learning" what-if" explanations for sequential decision-making. arXiv preprint arXiv:2007.13531 , 2020.

Amir-Hossein Karimi, Julius Von Kügelgen, Bernhard Schölkopf, and Isabel Valera. Algorithmic recourse under imperfect causal knowledge: a probabilistic approach. Advances in neural information processing systems , 33:265-277, 2020.

Stratis Tsirtsis and Manuel Gomez-Rodriguez. Decisions, counterfactual explanations and strategic behavior. Advances in Neural Information Processing Systems , 33:16749-16760, 2020.

A Further related work

Counterfactual reasoning and reinforcement learning. As mentioned in Section 1, there is a closely related line of work [10-13] that focuses on the development of machine learning methods that employ elements of counterfactual reasoning to improve or to retrospectively analyze decisions in sequential settings. Buesing et al. [10] use SCMs to express the transition dynamics in Partially Observable MDPs (POMDPs), and they propose a method to efficiently compute a policy based on counterfactual realizations of logged episodes. Lu et al. [12] adopt a similar modeling framework, and they propose a counterfactual data augmentation approach to speed up standard Q-learning. Oberst and Sontag [11] introduce the Gumbel-Max SCM to express the dynamics of an arbitrary discrete POMPD, and they develop a method for counterfactual off-policy evaluation to identify episodes where a given alternative policy would have achieved a higher reward. However, none of these works aims to find an action sequence, close to the observed sequence of a particular episode, that is counterfactually optimal.

Planning in continuous-state MDPs. Our work has additional connections to pieces of work that aim to approximate the optimal value function in an MDP with continuous states and a finite horizon [27, 62, 63]. Therein, the work most closely related to ours is the one by Bertsekas [27]. It shows that, under a Lipschitz continuity assumption on the transition dynamics, a value function computed via value iteration in an MDP with discretized states, converges to the optimal value function of the original (continous-state) MDP as the discretization becomes finer. Although some of the proof mechanics of this work are similar to ours, the contributions are orthogonal, as we do not employ any form of discretization, and we leverage the Lipschitz continuity assumption to compute optimal action sequences in continuous states using the A ∗ algorithm.

Counterfactual reasoning and explainability. Our work has ties to pieces of work that use forms of counterfactual reasoning as a tool towards learning explainable machine learning models. For example, Madumal et al. [64] propose to express the action selection process of a reinforcement learning agent as a causal graph, and they use it to generate explanations for the agent's chosen actions. Bica et al. [65] introduce an inverse reinforcement learning approach to learn an interpretable reward function from expert demonstrated behavior that is expressed in terms of preferences over potential (counterfactual) outcomes. Moreover, our work is broadly related to the work on counterfactual explanations (in classification) that aims to find a minimally different set of features that would have led to a different outcome, in settings where single decisions are taken [66, 67].

B Supporting table and figure

S

S

S

t-1

t+1

At

U

Z,

U

t-

S

†

C Counterfactual identifiability of element-wise bijective SCMs

In this section, we show that element-wise bijective SCMs, a subclass of bijective SCMs which we formally define next, are counterfactually identifiable.

Definition 8. An SCM C is element-wise bijective iff it is bijective and there exist functions g S,i : R × A × R → R with i ∈ { 1 , . . . , D } such that, for every combination of s t +1 , s t , a t , u t with s t +1 = g S ( s t , a t , u t ) , it holds that s t +1 ,i = g S,i ( s t , a t , u t,i ) for i ∈ { 1 , . . . , D } .

Under our assumption that the transition mechanism g S is continuous with respect to its third argument, it is easy to see that, for any element-wise bijective SCM, the functions g S,i are always strictly monotonic functions of the respective u t,i . Based on this observation, we have the following theorem of counterfactual identifiability:

Theorem 9. Let C and M be two element-wise bijective SCMs with transition mechanisms g S and h S , respectively, and, for any observed transition ( s t , a t , s t +1 ) , let u t = g -1 S ( s t , a t , s t +1 ) and ˜ u t = h -1 S ( s t , a t , s t +1 ) . Moreover, given any s ∈ S , a ∈ A , let s ′ = g S ( s , a, u t ) and s ′′ = h S ( s , a, ˜ u t ) . If P C ( S t +1 | S t = s , A t = a ) = P M ( S t +1 | S t = s , A t = a ) for all s ∈ S , a ∈ A , it must hold that s ′ = s ′′ .

Proof. We prove the theorem by induction, starting by establishing the base case s ′ 1 = s ′′ 1 . Without loss of generality, assume that both g S, 1 and h S, 1 are strictly increasing with respect to their third argument. Since the two SCMs entail the same transition distributions, we have that

where ( ∗ ) holds because both SCMs are element-wise bijective, and ( ∗∗ ) holds because g S, 1 and h S, 1 are increasing with respect to their third argument. Similarly, we have that

where in ( glyph[star] ) , ( † ) we have used the monotonicity of g S and h S , and ( glyph[star]glyph[star] ) follows from the previous result. The last equality implies that s ′ 1 and s ′′ 1 correspond to the same quantile of the distribution for S t +1 , 1 | S t = s , A t = a . Therefore, it is easy to see that s ′ 1 = s ′′ 1 since the opposite would be in contradiction to g S, 1 being bijective. Note that, we can reach that conclusion irrespective of the direction of monotonicity of g S, 1 and h S, 1 , since any change in the direction of the inequalities happening at step ( ∗∗ ) is reverted at steps ( glyph[star] ) and ( † ) .

Now, starting from the inductive hypothesis that s ′ i = s ′′ i for all i ∈ { 1 , . . . , n } with n < D , we show the inductive step, i.e. , s ′ n +1 = s ′′ n +1 . Again, without loss of generality, assume that both g S,n +1 and h S,n +1 are strictly increasing with respect to their last argument. Note that, the two SCMs entail the same transition distributions, i.e. , the same joint distributions for S t +1 | S t , A t . Following from the law of total probability, they also entail the same conditional distributions for S t +1 ,n +1 | S t +1 , ≤ n , S t , A t , where we use the notation x ≤ n to refer to a vector that contains the first n elements of a D -dimensional vector x . Therefore, we have that

where for the first equality we have used the inductive hypothesis, ( ∗ ) holds because both SCMs are element-wise bijective, and ( ∗∗ ) holds because g S,n +1 and h S,n +1 are increasing with respect to their third argument. Similarly, we get that

where in ( glyph[star] ) , ( † ) we have used the invertibility and monotonicity of g S and h S , and ( glyph[star]glyph[star] ) follows from the previous result. With the same argument as in the base case, the last equality implies that s ′ n +1 = s ′′ n +1 . That concludes the proof.

D Proofs

D.1 Proof of Theorem 4

We prove the hardness of our problem as defined in Eq. 9 by performing a reduction from the partition problem [26], which is known to be NP-Complete. In the partition problem, we are given a multiset of B positive integers V = { v 1 , . . . , v B } and the goal is to decide whether there is a partition of V into two subsets V 1 , V 2 with V 1 ∩ V 2 = ∅ and V 1 ∪ V 2 = V , such that their sums are equal, i.e. , ∑ v i ∈V 1 v i = ∑ v j ∈V 2 v j .

Consider an instance of our problem where S = U = R 2 , A contains 2 actions a diff , a null and the horizon is T = B +1 . Let C be an element-wise bijective SCM with arbitrary prior distributions P C ( U t ) such that their support is on R 2 and a transition mechanism g S such that

Moreover, assume that the reward function is given by

where sum ( V ) is the sum of all elements ∑ B i =1 v i . Note that, the SCM C defined above is Lipschitzcontinuous as suggested by the following Lemma (refer to Appendix D.1.1 for a proof).

Lemma10. The SCM C defined by Equations 12, 13 is Lipschitz-continuous according to Definition 1.

Now, assume that the counterfactual action sequence can differ in an arbitrary number of actions from the action sequence in the observed episode τ , i.e. , k = T and, let the observed action sequence be such that a t = a null for t ∈ { 0 , . . . , T -1 } . Lastly, let the initial observed state be s 0 = [0 , v 1 ] , the observed states { s t } T -2 t =1 be such that s t = [ ∑ t i =1 v i , v t +1 ] for t ∈ { 1 , . . . , T -2 } and the last observed state be s T -1 = [ sum ( V ) , 0] . Then, it is easy to see that the noise variables U t have posterior distributions with a point mass on the respective values

Note that, for all t ∈ { 1 , . . . , T -2 } , we have 0 ≤ s t, 1 < sum ( V ) and s t, 2 ≥ 1 , hence the immediate reward according to Eq. 13 is equal to 0 . Consequently, the outcome of the observed episode τ is o + ( τ ) = R ( s T -1 , a null ) = -max(0 , sum ( V ) 2 ) -max(0 , -sum ( V ) 2 ) = -sum ( V ) 2 .

Next, we will characterize the counterfactual outcome o ( τ ′ ) of a counterfactual episode τ ′ with a sequence of states { s ′ t } T -1 t =0 resulting from an alternative sequence of actions { a ′ t } T -1 t =0 . Let D ′ t , N ′ t denote the set of time steps until time t , where the actions taken in a counterfactual episode τ ′ are a diff and a null respectively. Formally, D ′ t = { t ′ ∈ { 0 , . . . , t } : a ′ t ′ = a diff } , N ′ t = { t ′ ∈ { 0 , . . . , t } : a ′ t ′ = a null } . Then, as an intermediate result, we get the following Lemma (refer to Appendix D.1.2 for a proof).

Following from that, we get that 0 ≤ s ′ t, 1 ≤ sum ( V ) for all t ∈ { 1 , . . . , T -1 } . Moreover, we can observe that the transition mechanism given in Eq. 12 is such that g S, 2 ( S t , A t , U t, 2 ) = U t, 2 for all t ∈ { 0 , . . . , T -2 } , independently of S T and A t . Therefore, it holds that s ′ t, 2 = u t -1 , 2 ≥ 1 for t ∈ { 1 , . . . , T -2 } , and s ′ 0 , 2 = s 0 , 2 = v 1 ≥ 1 . As a direct consequence, it is easy to see that R ( s ′ t , a ′ t ) = 0 for all t ∈ { 0 , . . . , T -2 } , and the counterfactual outcome is given by

In addition to that, we have that u T -2 , 2 = 0 , hence

Now, we will show that, if we can find the action sequence { a ∗ t } T -1 t =0 that gives the optimal counterfactual outcome o + ( τ ∗ ) for the aforementioned instance in polynomial time, then we can make a decision about the corresponding instance of the partition problem, also in polynomial time. To this end, let { s ∗ t } T -1 t =0 be the sequence of states in the optimal counterfactual realization and, let D ∗ T -2 = { t ∈ { 0 , . . . , T -2 } : a ∗ t = a diff } , N ∗ T -2 = { t ′ ∈ { 0 , . . . , T -2 } : a ∗ t ′ = a null } .

From Eq. 14, we get that the optimal counterfactual outcome is o + ( τ ∗ ) = R ( s ∗ T -1 , a ∗ T -1 ) , and it is easy to see that the reward function given in Eq. 13 is always less or equal than zero. If o ( τ ∗ ) = 0 , it has to hold that

where ( ∗ ) follows from Eq. 15. As a consequence, the subsets V 1 = { v i : i -1 ∈ N ∗ T -2 } and V 2 = { v i : i -1 ∈ D ∗ T -2 } partition V and their sums are equal.

On the other hand, if o + ( τ ∗ ) < 0 , as we will show, there is no partition of V into two sets with equal sums. For the sake of contradiction, assume there are two sets V 1 , V 2 that partition V , with sum ( V 1 ) = sum ( V 2 ) = sum ( V ) / 2 , and let N ′ T -2 = { t ∈ { 0 , . . . , T -2 } : v t +1 ∈ V 1 } and D ′ T -2 = { t ∈ { 0 , . . . , T -2 } : v t +1 ∈ V 2 } . Then, consider the counterfactual episode τ ′ with an action sequence { a ′ t } T -1 t =0 such that its elements take values a null and a diff based on the sets N ′ T -2 , D ′ T -2 respectively, with a ′ T -1 taking an arbitrary value. It is easy to see that which is a contradiction. This step concludes the reduction and, therefore, the problem given in Eq. 9 cannot be solved in polynomial time, unless P = NP .

D.1.1 Proof of Lemma 10

It is easy to see that, for all u ∈ U and for all s , s ′ ∈ S , the function g S ( S t , a null , u ) satisfies ‖ g S ( s , a null , u ) -g S ( s ′ , a null , u ) ‖ ≤ ‖ s -s ′ ‖ , and therefore K a null , u = 1 satisfies Definition 1. For the case of A t = a diff, we have that

and therefore ‖ g S ( s , a diff , u ) -g S ( s ′ , a diff , u ) ‖ 2 ≤ ( s 1 -s ′ 1 ) 2 +( s 2 -s ′ 2 ) 2 +2 | s 1 -s ′ 1 || s 2 -s ′ 2 | . We also have that

By combining these, we get

Hence, we can easily see that K a diff , u = √ 2 satisfies Definition 1.

Next, we need to show that, for all a ∈ A there exists a C a ∈ R + such that, for all s , s ′ ∈ S , it holds | R ( s , a ) -R ( s ′ , a ) | ≤ C a ‖ s -s ′ ‖ . Note that, to show that a function of the form max(0 , f ( s )) with f : R 2 → R is Lipschitz continuous, it suffices to show that f ( s ) is Lipschitz continuous, since the function max(0 , x ) with x ∈ R has a Lipschitz constant equal to 1 .

We start by showing that the function f ( s ) = s 1 -α -s 2 · α is Lipschitz continuous, where α = sum ( V ) / 2 is a positive constant. For an arbitrary pair s , s ′ ∈ S , we have that

We also have that

By combining these, we get

Hence, we arrive to | f ( s ) -f ( s ′ ) | ≤ √ 1 + α ‖ s -s ′ ‖ , and the function f is Lipschitz continuous. It is easy to see that the function φ ( s ) = α -s 1 -s 2 · α is also Lipschitz continuous with the proof being almost identical. As a direct consequence, the reward function given in Equation 13 satisfies

Definition 1 with C a null = C a diff = 2 √ 1 + sum ( V ) 2 . This concludes the proof of the lemma.

D.1.2 Proof of Lemma 11

We will prove the lemma by induction. For the base case of t = 1 , we distinguish between the cases a ′ 0 = a diff and a ′ 0 = a null . In the first case, we have s ′ 1 , 1 = u 0 , 1 + s 0 , 1 -s 0 , 2 = v 1 +0 -v 1 = 0 and N ′ 0 = ∅ and, therefore, the statement holds. In the second case, we have s ′ 1 , 1 = u 0 , 1 + s 0 , 1 = v 1 +0 = v 1 , N ′ 0 = { 0 } and ∑ t ′ ∈N ′ 0 v t ′ +1 = v 1 . Therefore, the statement also holds.

For the inductive step ( t > 1 ), we assume that s ′ t -1 , 1 = ∑ t ′ ∈N ′ t -2 v t ′ +1 and we will show that s ′ t, 1 = ∑ t ′ ∈N ′ t -1 v t ′ +1 . Again, we distinguish between the cases a ′ t -1 = a diff and a ′ t -1 = a null . However, note that, in both cases, s ′ t -1 , 2 = u t -2 , 2 +0 = v t . Therefore, in the case of a ′ t -1 = a diff , we get

where the last equation holds because a ′ t -1 = a diff and, therefore, N ′ t -1 = N ′ t -2 . In the case of a ′ t -1 = a null, we get

where the last equation holds because a ′ t -1 = a null and, therefore, N ′ t -1 = N ′ t -2 ∪ { t -1 } .

D.2 Proof of Lemma 5

We will prove the proposition by induction, starting from the base case, where t = T -1 . First, Let t = T -1 and l = k . It is easy to see that, if the process is at a state s ∈ S in the last

time step with no action changes left, the best reward that can be achieved is R ( s , a T -1 ) , as already discussed after Eq. 10. Therefore, it holds that | V τ ( s , k, T -1) -V τ ( s ′ , k, T -1) | = | R ( s , a T -1 ) -R ( s ′ , a T -1 ) | ≤ C a T -1 ‖ s -s ′ ‖ ≤ C ‖ s -s ′ ‖ , where the last step holds because C = max a ∈A C a . Now, consider the case of t = T -1 with l taking an arbitrary value in { 0 , . . . , k -1 } . Let s , s ′ be two states in S and a ∗ be the action that gives the maximum immediate reward at state s , that is, a ∗ = argmax a ∈A { R ( s , a ) } . Then, we get

where ( ∗ ) follows from the fact that R ( s ′ , a ∗ ) ≤ max a ∈A { R ( s ′ , a ) } . Therefore, for any l ∈ { 0 , . . . , k } and s , s ′ ∈ S , it holds that | V τ ( s , l, T -1) -V τ ( s ′ , l, T -1) | ≤ L T -1 ‖ s -s ′ ‖ , where L T -1 = C .

Now, we will proceed to the induction step. Let t < T -1 , l < k and, as an inductive hypothesis, assume that L t +1 ∈ R + as defined in Lemma 5 is such that, for all l ∈ { 0 , . . . , k } and s , s ′ ∈ S , it holds that | V τ ( s , l, t +1) -V τ ( s ′ , l, t +1) | ≤ L t +1 ‖ s -s ′ ‖ . Additionally, let ( s a , l a ) , ( s ′ a , l a ) denote the enhanced states that follow from ( s , l ) , ( s ′ , l ) after taking an action a , i.e. , ( s a , l a ) = F + τ,t (( s , l ) , a ) and ( s ′ a , l a ) = F + τ,t (( s ′ , l ) , a ) . Lastly, let a ∗ be the action that maximizes the future total reward starting from state s , i.e. , a ∗ = argmax a ∈A { R ( s , a ) + V τ ( s a , l a , t +1) } . Then, we have that

In the above, ( ∗ ) holds due to R ( s ′ , a ∗ ) + V τ ( s ′ a ∗ , l a ∗ , t +1) ≤ max a ∈A { R ( s ′ , a ) + V τ ( s ′ a , l a , t +1) } , ( ∗∗ ) follows from the inductive hypothesis, and ( ∗ ∗ ∗ ) holds because C = max a ∈A C a and K u t = max a ∈A K a, u t . It is easy to see that, similar arguments hold for the simple case of l = k , therefore, we omit the details. This concludes the inductive step and the proof of Lemma 5.

D.3 Proof of Proposition 6

We will prove the proposition by induction, starting from the base case, where t = T -1 . If t = T -1 , the algorithm initializes ˆ V τ ( s , l, T -1) to max a ∈A R ( s , a ) for all s ∈ S † , l ∈ { 0 , . . . , k -1 } and ˆ V τ ( s , k, T -1) to R ( s , a T -1 ) . It is easy to see that those values are optimal, as already discussed after Eq. 10. Therefore, the base case ˆ V τ ( s , l, T -1) ≥ V τ ( s , l, T -1) follows trivially.

Now, we will proceed to the induction step. Let t < T -1 and, as an inductive hypothesis, assume that ˆ V τ ( s , l, t +1) ≥ V τ ( s , l, t +1) for all s ∈ S † , l ∈ { 0 , . . . , k } . Our goal is to show that ˆ V τ ( s , l, t ) ≥ V τ ( s , l, t ) for all s ∈ S † , l ∈ { 0 , . . . , k } . First, let l < k . For a given point s ∈ S † , Algorithm 1 finds the next state s a that would have occurred by taking each action a , i.e. , ( s a , l a ) = F + τ,t (( s , l ) , a ) , and it computes the associated value V a = min s † ∈S † { ˆ V τ ( s † , l a , t +1) + L t +1 ‖ s † -s a ‖} . Then, it

simply sets ˆ V τ ( s , l, t ) equal to max a ∈A { R ( s , a ) + V a } . We have that

where ( ∗ ) follows from the inductive hypothesis, and ( ∗∗ ) is a consequence of Lemma 5. Then, we get

Additionally, when l = k , we have ˆ V τ ( s , k, t ) = R ( s , a t ) + min s † ∈S † { ˆ V τ ( s † , k, t + 1) + L t +1 ‖ s † -s a t ‖} and V τ ( s , k, t ) = R ( s , a t ) + V τ ( s a t , k, t + 1) . Therefore, the proof for ˆ V τ ( s , k, t ) ≥ V τ ( s , k, t ) is almost identical.

D.4 Proof of Theorem 7

We start from the case where t = T -1 . Let v = ( s , l, T -1) and, consider an edge associated with action a ∗ connecting v to the goal node v T = ( s ∅ , k, T ) that carries a reward R ( s , a ∗ ) . Then, we have

and the base case holds.

For the more general case, where t < T -1 , we first establish the following intermediate result, whose proof is given in Appendix D.4.1.

Lemma 12. For every s , s ′ ∈ S , l ∈ { 0 , . . . , k } , t ∈ { 0 , . . . , T -1 } , it holds that | ˆ V τ ( s , l, t ) -ˆ V τ ( s ′ , l, t ) | ≤ L t ‖ s -s ′ ‖ , where L t is as defined in Lemma 5.

That said, consider an edge associated with an action a ∗ connecting node v = ( s , l, t ) to node v a ∗ = ( s a ∗ , l a ∗ , t +1) . Then, we have

That concludes the proof and, therefore, the heuristic function ˆ V τ is consistent.

D.4.1 Proof of Lemma 12

Without loss of generality, we will assume that l < k , since the proof for the case of l = k is similar and more straightforward. We start from the case where t = T -1 and, for two states s , s ′ ∈ S we have

where the last inequality follows from Lemma 5.

Now, consider the case t < T -1 , and let ( s a , l a ) denote the enhanced state that follows from ( s , l ) after taking an action a at time t , i.e. , ( s a , l a ) = F + τ,t (( s , l ) , a ) . Then, we have

Let a ∗ be the action a ∈ A that maximizes the first part of the above subtraction, i.e. ,

Then, Eq. 16 implies that

Now, let ˜ s be the s † ∈ S † that minimizes the second part of the above subtraction, i.e. ,

As a consequence and in combination with Eq. 17, we get

where in ( ∗ ) we use the triangle inequality and the fact that the SCM C is Lipschitz-continuous, and ( ∗∗ ) follows from Lemma 5.

Algorithm 2: Graph search via A ∗

glyph[negationslash]

E A* algorithm

Algorithm 2 summarizes the step-by-step process followed by the A ∗ algorithm. Therein, we represent each node v by an object with 4 attributes: (i) the 'tuple' ( s , l, t ) of the node, (ii) the total reward 'rwd' of the path that has led the search from the root node v 0 to the node v , (iii) the parent node 'par' from which the search arrived to v , and (iv) the action 'act' associated with the edge connecting the node v with its parent. In addition to the queue of nodes to visit, the algorithm maintains a set of explored nodes, and it adds a new node to the queue only if it has not been previously explored. The algorithm terminates when the goal node is chosen to be visited, i.e. , the 'tuple' attribute of v has the format ( ∗ , ∗ , T ) , where ∗ denotes arbitrary values. Once the goal node v T has been visited, the algorithm reconstructs the action sequence that led from the root node v 0 to the goal node and returns it as output.

10000

15000

20000

25000

30000

Size of anchor set,

|S

†

|

1

.

9

2

.

0

2

.

1

EBF

(a) MC-Lipschitz

10000

15000

20000

25000

30000

Size of anchor set,

|S

†

|

1

.

9

2

.

0

2

.

1

EBF

(b) MC-Uniform

10000

15000

20000

25000

30000

Size of anchor set,

|S

†

|

3

.

00

3

.

05

3

.

10

EBF

(c) Facility-Location

F Experimental evaluation of anchor set selection strategies

In this section, we benchmark the anchor set selection strategy presented in Section 4 against two alternative competitive strategies using the sepsis management dataset and the same experimental setup as in Section 5. More specifically, we consider the following anchor set selection strategies:

(i) MC-Lipschitz : This is the strategy described in depth in Section 4, based on Monte Carlo simulations of counterfactual episodes under randomly sampled counterfactual action sequences. Notably, the time steps where each counterfactual action sequence differs from the observed one are sampled proportionally to the respective Lipschitz constant L t of the SCM's transition mechanism. To ensure a fair comparison with other strategies, instead of controlling the number of sampled action sequences M , we fix the desired size of the anchor set S † , and we repeatedly sample counterfactual action sequences until the specified size is met.

(ii) MC-Uniform : This strategy is a variant of the previous strategy where we sample the time steps where each counterfactual action sequence differs from the observed one uniformly at random, rather than biasing the sampling towards time steps with higher Lipschitz constants L t .

(iii) Facility-Location : Under this strategy, the anchor set is the solution to a minimax facility location problem defined using the observed available data. Let S o be the union of all state vectors observed in all episodes τ in a given dataset. Then, we choose an anchor set S † ⊂ S o of fixed size |S † | = b , such that the maximum distance of any point in S o to its closest point in S † is minimized. Here, the rationale is that counterfactual states resulting from counterfactual action sequences for one observed episode are likely to be close to the observed states of some other episode in the data. Formally,

Although the above problem is known to be NP-Complete, we find a solution using the farthest-point clustering algorithm, which is known to have an approximation factor equal to 2 and runs in polynomial time. The algorithm starts by adding one point from S o to S † at random. Then, it proceeds iteratively and, at each iteration, it adds to S † the point from S o that is the furthest from all points already in S † , i.e. , S † = S † ∪ s , where s = max s ′ ∈S o { min s † ∈S † || s ′ -s † || } . The algorithm terminates after b iterations.

Results. We compare the computational efficiency of our method under each of the above anchor set selection strategies for various values of the size of the anchor set |S † | . Figure 5 summarizes the results. We observe that the Facility-Location selection strategy performs rather poorly compared to the other two strategies, achieving an effective branching factor (EBF) higher than 3 . In contrast, the MC-Lipschitz and MC-Uniform strategies achieve an EBF close to 2 , which decreases rapidly as the size of the anchor set increases. Among these two strategies, the MC-Lipschitz strategy, which we use in our experiments in Section 5, achieves the lowest EBF.

G Additional details on the experimental setup

G.1 Features and actions in the sepsis management dataset

As mentioned in Section 5, our state space is S = R D , where D = 13 is the number of features. We distinguish between three types of features: (i) demographic features, whose values remain constant across time, (ii) contextual features, for which we maintain their observed (and potentially varying) values throughout all counterfactual episodes and, (iii) time-varying features, whose counterfactual values are given by the SCM C . The list of features is as follows:

Gender (demographic)

Age (demographic)

Re-admission (demographic)

Mechanical ventilation (contextual)

PaO 2 (time-varying)

FiO 2 (time-varying)

Platelet count (time-varying)

Glasgow Coma Scale (time-varying)

Bilirubin (time-varying)

Mean arterial blood pressure (time-varying)

Urine output (time-varying)

Creatinine (time-varying)

SOFA score (time-varying)

To define our set of actions A we follow related work [6, 55, 57], and we consider 25 actions corresponding to 5 × 5 levels of administered vasopressors and intravenous fluids. Specifically, for both vasopressors and fluids, we find all non-zero values appearing in the data, and we divide them into 4 intervals based on the quartiles of the observed values. Then, we set the 5 levels to be the median values of the 4 intervals and 0 . Table G.1 shows the resulting values of vasopressors and fluids.

G.2 Additional details on the network architecture & training

We represent the location and scale functions h and φ of the SCM C using neural networks with 1 hidden layer, 200 hidden units and tanh activation functions. The mapping from a state s and an action a to the hidden vector z takes the form z = tanh ( W s s + W a a ) , where a is a 2 -D vector representation of the respective action. The mapping from the hidden vector z to the network's output is done via a fully connected layer with weights W z . To enforce a network to have a Lipschitz constant L with respect to the state input, we apply spectral normalization to the weight matrices W s and W z , so that their spectral norms are ‖ W s ‖ 2 = ‖ W z ‖ 2 = 1 . Additionally, we add 2 intermediate layers between the input and the hidden layer and between the hidden layer and the output layer, each one multiplying its respective input by a constant √ L . Since it is known that the tanh activation function has a Lipschitz constant of 1 , it is easy to see that, by function composition, the resulting network is guaranteed to be Lipschitz continuous with respect to its state input with constant L . Note

0.01

0.02

0.05

0.1

0.2

0.5

1.0

1.1

1.2

1.5

L

h

0.01

0.02

0.05

0.1

0.2

0.5

1.0

1.1

1.2

1.5

L

φ

-

20%

-

10%

that, since the matrix W a is not normalized, the network's Lipschitz constant with respect to the action input can be arbitrary.

To train the SCM C , for each sample, we compute the negative log-likelihood of the observed transition under the SCM's current parameter values ( i.e. , network weight matrices & covariance matrix of the multivariate Gaussian prior), and we use that as a loss. Subsequently, we optimize those parameters using the Adam optimizer with a learning rate of 0 . 001 , a batch size of 256 , and we train the model for 100 epochs.

We train the model under multiple values of the Lipschitz constants L h , L φ of the location and scale networks, and we evaluate the log-likelihood of the data under each model using 5 -fold crossvalidation. Specifically, for each configuration of L h and L φ , we randomly split the dataset into a training and a validation set (with a size ratio 4 -to1 ), we train the corresponding SCM using the training set, and we evaluate the log-likelihood of the validation set based on the trained SCM. This results in the log-likelihood always being measured on a different set of data points than the one used for training. For each configuration of L h and L φ , we repeat the aforementioned procedure 5 times and we report the average log-likelihood achieved on the validation set. In addition, we train an unconstrained model without spectral normalization, which can have an arbitrary Lipschitz constant.

Figure 6 shows the decrease in log-likelihood of the respective constrained model as a percentage of the log-likelihood achieved by the unconstrained model, under various values of the Lipschitz constants L h , L φ . We observe that, the model's performance is predominantly affected by the Lipschitz constant of the location network L h , and its effect is more pronounced when L h takes values smaller than 1 . Additionally, we can see that the scale network's Lipschitz constant L φ has a milder effect on performance, especially when L h is greater or equal than 1 . Since we are interested in constraining the overall Lipschitz constant of the SCM C , in our experiments in Section 5, we set L h = 1 and L φ = 0 . 1 , which achieves a log-likelihood only 6% lower to that of the best model trained without any Lipschitz constraint.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors present a method for finding c/f (optimal) action sequences in sequential decision making problems with uncertainty, with the novelty that they consider continuous state dynamics. They apply their method to the interesting setting of sepsis treatment at the end of the paper and therein demonstrate impressive results for their method (compared to the actual, observed, action sequences).\n\nSTRENGTHS:\nPlease see the Questions section for the full review.\n\nWEAKNESSES:\nPlease see the Questions section for the full review.\n\nQUESTIONS:\n## Abstract\n\n- I re-read the abstract multiple times and I unfortunately struggle with the first half. Consider changing the sentence construction or just simplify the first half. It is very unclear what you are trying to say and I fear that the positioning of the paper suffers because of it. One suggestion could be to anchor the paper with one or two real examples and then take the justification from there.\n- The first paragraph of the introduction does a much better job of framing your paper, I would suggest adapting that approach for the first half of the abstract.\n\n## Introduction \n\n- Spell out ICU before you use it; it is common-enough term in the West but may be unfamiliar to a non-native speaker (or indeed anyone who does not frequently use English).\n- Excellent second paragraph.\n- Line 49: an SCM does not have the concept of state or transitions, merely a distribution P() over the endogenous and exogenous variables in the model (as defined originally by Pearl). Consequently how you have chosen to introduce it here does not quite follows. Consider re-phrasing. Moreover, SCMs are, as originally defined, state-less representations of a causal environment.\n- It would be helpful to the reader and yourselves if you could show us (the readers) a quadrant of work which has considered MDPs with: (continuous states, discrete actions), (continuous states, continuous action) … and so on — so that we understand what gap you are filling with this work (certainly you spell it out, but I would argue that a quadrant would be more useful and visually more powerful). \n- Line 60: please explain what a bijective SCM is. It is not enough to simply cite the original paper it if it is going to be key to your method.\n\n##\tA Causal Model of Sequential Decision Making Processes \n\n- Line 76: presumably $a_t \\in \\mathcal{A} = \\mathbb{Z}_+$?\n- How come you are not considered the discounted cumulative reward?\n- I think you should differentiate your setting more than what you are currently doing. At present it sounds as if eq (2) is part of the standard SCM definition. Perhaps introduce a definition paragraph for your SCM? That way the reader is clear that your are formalising a new concept and there is now ambiguity between yours and Pearl’s definition.\n- I would like to see a much longer discussion on why you make the causal sufficiency assumption (Line 97) in this work (no unobserved confounders). It is a common assumption as you say, but merely saying that other works make the same assumption, is not good enough justification to make it a tenable assumption. It is a huge assumption and highly unrealistic in most realistic scenarions. It makes life computationally far easier but trades-off usability down the line. Consequently, please discuss why you make this assumption. \n- Line 104: the do() express does not have to constitute a hard (atomic) intervention, it can also be a soft intervention (and others). Please consider or at least mention these settings too and why you settle for hard interventions.\n- Definition 1: is this definition somehow different from the standard one? Perhaps you ought to also give it a source since this is common fare but it may differ from the one readers are used to (and would also elucidate on the matter if it the standard definition or not).\n- The paragraph at the bottom of page 3 is great, very interesting application of Lipschitz continuity on a real problem.\n- Define the indicator function in equation 5.\n- Lines 149-161: I wonder if it may be not better for you to formally have a small lemma here, possibly even a theorem, explicating the results for which you have a proof. It currently reads as if it was merely a by-product of your method, whereas I imagine it is rather more important than that?\n\n## Problem statement\n\n- I would recommend revising the structure of your paper given that the problem statement appears on page four at conference with a page limit of nine. The paper will read far better if you frame the problem early on, allowing the reader to understand your angle of attack from the very start.\n- Line 163: you already introduced the episode on line 78.\n- Where is practically relevant to know that eq 9 is NP hard? A trivial question no doubt, but I would be interested to hear what the authors have in mind with regards to this theorem (is it really a theorem?)\n\n## Finding the optimal counterfactual action sequence via A* search\n\n- Line 210: I think a diagram here would be very helpful, explaining algorithm 1; a lot of page six could gainfully be explained with a figure rather than paragraphs of (rather dense) text.\n\n## Experiments using clinical sepsis management data\n\n- Looking at figure 2(a) I would have thought that as k increased we would have seen a much larger increase in the y-axis but this does not seem to be the case, the increase is fairly modest. Can you explain? \n- According to panel 2(b), 50% of patients see less than or equal to ~5% c/f improvement? Is that correctly read?\n- I enjoyed this section and the example is good but it is long and arduous, and we have no evidence to suggest that your method works for more (even synthetic) settings. It is an old trope of reviewing that we (reviewers) always want more experiments but in this case I think it is a fair request. Simply because I am struggling to understand how your ideas generalise across different settings and applications (even synthetic) nor do we have a comparison against other methods (are there any to which you can gainfully compare against?)\n\nLIMITATIONS:\nPlease see the Questions section for the full review.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper tackles the problem of finding counterfactual action sequences in sequential decision making problems. The main difference to previous methods is that this paper regards a continuous state space instead of a discrete one, which renders previous solution methods infeasible. The authors formalize the problem as an SCM and show that the solution in general is NP hard. As a solution, they present a new method based on the A-star algorithm, which yields for many problems an efficient solution. They evaluate the method on clinical data.\n\nSTRENGTHS:\nThe paper is well-written and easy to follow.\n\nThe problem of efficiently finding counterfactual action sequences for the continuous state space seems new to me and relevant to the community.\n\nEven if the SCM formulation follows closely the work of Tsirtsis et al. (2021), a likewise dynamic programming approach would require enumerating all possible action sequences. From my point of view, the main contribution lies therefore in the A-star-based search method (and the anchor set selection), which requires further technical details such as bijectivity and Lipschitz continuity of the SCM for the heuristics. Therefore, the technical contribution seems good to me.\n\nThe method is evaluated with regard to the influence of different parameters.\n\nWEAKNESSES:\nThe worse case complexity of the method is the same as for brute-force search. The efficiency of the method depends on the number of samples of the Monte Carlo samples in the anchor set. It is not clear to me, which conditions the problem must have in order for the method to work effiently and when it falls back to worse case complexity.\n\nMinor (typos):\nline 190 comma before \"such that\"\n\nQUESTIONS:\nHow important are the assumptions that the SCM is bijective and Lipschitz-continuous? You need this assumption for the bounds in your algorithm. Can you still approach the problem if one condition is violated? How much do you think do the conditions restrict the field of applications of your method?\n\nAny thoughts about the conditions of the problem under which the method works efficiently?\n\nLIMITATIONS:\nLimitations and directions for future work are adequately addressed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles the problem of finding optimal action sequences in domains with continuous state spaces with counterfactual reasoning.  They propose an A* based search approach and show the efficacy of the approach on a clinical sepsis management problem.\n\nSTRENGTHS:\n* The paper focuses on a very significant research problem\n* The experimental evaluation on a  real-world clinical dataset demonstrates the strong utility of the approach\n\nWEAKNESSES:\n* I have some concerns regarding the causal d-separation assumption (see Q1 & Q2).\n\nQUESTIONS:\n1. Graphical implication of an intervention on a variable $X$ (ie.  $do(X=x)$ ) is that all the incoming edges to $X$ are removed. In other words, as we are intervening on X, the influence from the $PA_X$  (parent nodes of X) on the node $X$ becomes zero. Refer to paragraph after Def 3.2.1 in Chapter 3 of Pearl. In Equation 3, it seems that all the arrows pointing out of $A_t$ are removed to adjust for $do(A_t = a_t)$ (as mentioned in line 106). So, the d-separation assumed does not follow from the do-calculus. \n2. Similar d-separation assumptions are made in Eq. 4 and 6.  \n\t1. Equation 1. suggests that two variables have a parent-child relationship: $S_t \\rightarrow A_t$.  \n\t2. Equation 2 suggests that their variable form a collider: $S_t \\rightarrow S_{t+1} \\leftarrow A_{t}$.\n\t3. On performing do-operation on $A_t$  would remove the parent-child relationship. But still maintain the collider  $S_t \\rightarrow S_{t+1} \\leftarrow A_{t}$. So, $S_{t+1}$ and $A_t$ are not d-separated by $S_t$. \n3. What are the implications of modifying the d-separation assumption on the proposed algorithms?\n4. How was the generated counterfactual action sequence evaluated? The SOFA score requires vital signs for calculations, how were the vital signs obtained?\n5. Figure 1 needs a legend. It is not clear what the pink & green lines represent.\n6. Did any physician analyze the counterfactual action sequences? Was any physian involved at any stage of the process?\n7. I recommend authors provide a reference table for notations in the appendix, it was really difficult to keep track of all the notations used in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies a question that is very natural for a sequential decision maker to ask itself: \"how could I best improve the return I obtained in a trajectory by only changing a fixed (k) number of actions from the sequence of actions I just executed while keeping the rest of the actions fixed\". Authors study this question in a context where the state space is continuous but the action space is discrete. The paper generalizes the result of previous work which studied a similar question, but in the context of discrete states. The setting is very nicely motivated and the problem is quite relevant and important.\n\nThis formulation seems to be novel, general, and useful. However, the paper shows that answering this question is in fact an NP-hard, one that we can only solve in the worst case using exhaustive search. Therefore, the paper goes on to find an admissible heuristic to be used in conjunction with the A* algorithm to reduce the complexity of the search space.\n\nTo propose a consistent/admissible heuristic, one needs to find reasonable upper bounds for the value of unseen states which are found using counterfactual reasonings. The paper proposes to get this done using the notion of Lipschitz continuity where intuitively, the difference between the value of the two states is upper bounded by a constant times the norm of the difference of the two states. This may be a stringent requirement generally, but it seems to make sense in the domain that is of interest in the paper, namely in medical trials. Experiments show that using this heuristic does result in significantly reducing the average runtime of the search.\n\nOverall, this is a fair contribution in the intersection of causal reasoning and planning/reinforcement learning.\n\nSTRENGTHS:\nThe competency that the paper attempts to develop, namely being able to answer counterfactual questions in the context of planning and reinforcement learning, is quite interesting. Formulating this problem and showing that the original problem is NP-hard also sheds some light on the complexity of asking such questions. Some of the assumptions made, such as using a bijective causal model does limit the scope of the result, but I still find the formulation very interesting.\n\nWEAKNESSES:\nWe know that Lipchitz continuity can be a pretty lose upper bound for the function value. Especially in the context of this paper, when data is sparse, the upper bounds derived in Lemma 4 can be pretty bad. Moreover, recently there have been more efforts to develop Lipschitz-like tools that are more conducive to RL algorithms by ensuring that the new tool are Coarse-Grained. See for example \"Coarse-Grained Smoothness for Reinforcement Learning in Metric Spaces\", by Gottesman et al, 2023 or \"Zooming for Efficient Model-Free Reinforcement Learning in Metric Spaces\" by Touati et al 2020. This paper though relies on more classical notions of smoothness (vanilla Lipchitz continuity) which is not the most effective in the context of RL.\n\n\nAnother point is that under the assumptions made in this paper (namely Lipschitz reward and transition dynamics), the value function itself becomes Lipschitz, so when learning the value function one can restrict the set of feasible value functions to be Lipschitz also. Is there any reason why this is not leveraged? In the same vein, we clearly know that the value function is always between (1-\\gamma) R_max and -(1-\\gamma) R_max. Can you elaborate if the bounds that are obtained by the Lipschitz interpolation would be able to provide much tighter bounds or that the bounds become vacuous? Can you also elaborate on the amount of data needed in your medical domains before these bound become non vacuous?\n\nIn practice, we also need to make sure that we do not underestimate the Lipschitz constant, because otherwise the heuristic used in A* will no longer be admissible. The estimate also needs to be tight enough so as to make sure that it is effective when used as a heuristic. But I am not sure how this trade-off will be maintained in problems where the Lipschitz constant is unknown.\n\nQUESTIONS:\nI would like to see a comparison with the discrete case, in particular the \"Woulda, coulda, shoulda: Counterfactually-guided policy search\" paper of Buesing et al, 2018. I do understand that this paper is dealing with the continuous state space. But one natural baseline to compare against could for example be the paper mentioned above where one can simply discretize the continuous state space and then apply their algorithm to compare against.\n\nHow important is the kick-starting part with the monte carlo anchoring? In particular, it would have been nice to see ablations where the monte-carlo anchoring is performed more or less to understand its effect on performance.\n\nLIMITATIONS:\nAuthors maintain that they will release code for their experiments if accepted, but they seem to have access to confidential patient data. I hope and trust that authors will they take sufficient precautions to ensure the anonymity of patients whose information is used in these experiments.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the problem of finding a counterfactual action sequence to maximize the outcome of a trajectory in an MDP characterized by a structural causal model. The focus is on continuous state spaces under a set of Lipschitz constraints. The authors show that this problem is NP hard and propose an algorithm based on A* search, which they test empirical on a dataset on clinical decision making.\n\nSTRENGTHS:\n- The paper studies a relevant well-motivated problem that will be of interest to the NeurIPS community\n- The authors do a good job in clarifying and formalizing the problem\n- Showing that the proposed problem is NP-hard is a valuable insight to guide future theoretical and empirical work\n- The proposed algorithm is a natural choice and it's formulation is sound\n- The empirical evaluation is performed on a realistic dataset, and the results are promising\n\nWEAKNESSES:\n- The discussion of related work is too limited. I think there is a lot of related work about related problems in SCM (outside RL) that should be mentioned. I'm not very familiar with this literature myself, which made it difficult to judge the potential impact of the contribution the present paper makes.\n\n- Theorem 3 (NP-hardness) is a key contribution of the paper. However, it is only stated in the paper and almost not discussed. It would be valuable to give a high level picture of the reduction used to prove the result, as well as a discussion of the implications of this result and how it results in the choice of algorithm later in the paper.\n\n- The empirical results are difficult to interpret due to a lack of baselines. While the efficiency results in Figure 1 and performance results in Figure 2 look as expected qualitatively, there is no point of comparison.\n   - For efficiency it would be useful to compare to some naive baselines, like a search based method without a specific heuristic.\n   - The counterfactual performance improvements in Figure 2 are pretty small which makes me wonder if the dataset is a good enough benchmark to evaluate the present method. Maybe the trajectories in the dataset are too close to optimal?\n   - It could be interesting to compare to an approach the discretizes the environment and uses prior methods for discrete state spaces.\n\nQUESTIONS:\n- What could be baselines to compare the empirical results to?\n   - What impact does the chosen heuristic have on the performance of the search algorithm?\n   - How much worse would be a method that discretizes the environment?\n\n- How can we be confident that the present dataset is a good benchmark to evaluate the approach?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}