Leveraging Factored Action Spaces for Efficient Offline Reinforcement Learning in Healthcare
Authors: Shengpu Tang
Maggie Makar
Michael Sjoding
Finale Doshi-Velez
Jenna Wiens
Venue: NeurIPS 2022 Accept (2022)
Keywords: reinforcement learning
offline rl
action space factorization
bias-variance trade-off
domain knowledge
healthcare
sepsis

Abstract: Many reinforcement learning (RL) applications have combinatorial action spaces, where each action is a composition of sub-actions. A standard RL approach ignores this inherent factorization structure, resulting in a potential failure to make meaningful inferences about rarely observed sub-action combinations; this is particularly problematic for offline settings, where data may be limited. In this work, we propose a form of linear Q-function decomposition induced by factored action spaces. We study the theoretical properties of our approach, identifying scenarios where it is guaranteed to lead to zero bias when used to approximate the Q-function. Outside the regimes with theoretical guarantees, we show that our approach can still be useful because it leads to better sample efficiency without necessarily sacrificing policy optimality, allowing us to achieve a better bias-variance trade-off. Across several offline RL problems using simulators and real-world datasets motivated by healthcare, we demonstrate that incorporating factored action spaces into value-based RL can result in better-performing policies. Our approach can help an agent make more accurate inferences within underexplored regions of the state-action space when applying RL to observational datasets. 

Leveraging Factored Action Spaces for Efficient Offline Reinforcement Learning in Healthcare
Shengpu Tang'
Maggie Makar Michael W. Sjoding? Finale Doshi-Veler. Jenna Wiens
Division Computer Science Engineering, University of Michigan, Ann Arbor, MI, USA Division of Pulmonary and Critical Care, Michigan Medicine, Ann Arbor, MI, USA SSEAS, Harvard University, Cambridge, MA, USA
Correspondence to: hinp://wwiiiii@mmmhm....
Reviewed on OpenReview: http:://pppanveewwwooo//////iddt
Abstract
Many reinforcement learning (RL) applications have combinatorial action spaces, where each action is a composition of sub-actions. standard RL approach ignores this inherent factorization structure, resulting in potential failure to make meaningful inferences about rarely observed sub-action combinations; this is particularly problematic for offline settings, where data may be limited. In this work, we propose a form of linear Q-function decomposition induced by factored action spaces. We study the theoretical properties of our approach, identifying scenarios where it is guaranteed to lead to zero bias when used to approximate the Q-function. Outside the regimes with theoretical guarantees, show that our approach still useful because it leads to better sample efficiency without necessarily sacrificing policy optimality, allowing to achieve better bias- variance trade-off. Across several offline RL. problems using simulators and real-world datasets motivated by healthcare, we demonstrate that incorporating factored action spaces into value- based RL can result in better-pen entter-performing policies. Our approach can help an agent make more accurate inferences within underexplored regions of the state-action space when applying RL to observational datasets.
Introduction
In many real-world decision- makinn problems, the action space exhibits an inherent combinatorial structure. For example, in healthcare, an action may correspond to a combination of drugs and treat- ments. When applying reinforcement learning (RL) to these tasks, past work [1-4] typically considers each combination distinct action, resulting in exponentially large action space (Figure la). This inefficient as fails leverage any potential independence among dimensions of the action space.
This type of factorization structure in action space could be incorporated when designing the architec- ture function approximators for RL (Figure lb). Similar ideas have been used the past, primarily to improve online exploration [5, 6]. or to handle multiple agents [7-11] or multiple rewards [12]. However, the applicability of this approach has not been systematically studied, especially in offline settings and when the MDP presents no additional structure (e.g.. when the state space cannot be explicitly factorized).
In this work, we develop an approach for offline RL with factored action spaces by learning linearly decomposable Q-functions. First, study the theoretical properties of this approach, investigating the sufficient and necessary conditions for it to lead to an unbiased estimate of the Q-function (i.e., zero approximation error). Even when the linear decomposition is biased, note that approach
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

(a)
combinatorial action space
(b)
factored action space
select Q(x;a)
Σ Q(x,a)
Figure 1: Illustration of Q-network architectures, which take the state s input and output Q(s, a) for a selected action. In this example, the action space A consists of D = 3 binary sub-action spaces { [ and {O, J. (a) Learning with the combinatorial action space requires 2 output heads (exponential in D), one for each combination of sub-actions. (b) Incorporating the linear Q decomposition for the factored action space requires
6 output heads (linear in D).
leads to reduction of variance, which in turn leads to improvement in sample efficiency. Lastly, we show that when sub-actions exhibit certain structures (e.g.. when two sub-actions "reinforce" their independent effects), the linear approximation, though biased, can still lead the optimal policy. We test our approach in offline RI. domains using a simulator [13] and a real clinical dataset [2], where domain knowledge ahout the relationship among actions suggests our proposed factorization approach is applicable. Empirically, our approach outperforms a non-factored baseline when the sample size is limited, even when the theoretical assumptions (around the validity of a linear decomposition) are perfectly satisfied. Qualitatively, in the real-data experiment, our approach learns policies that better capture the effect of less frequently observed treatment combinations.
Our work provides both theoretical insights and empirical evidence for RL practitioners consider this simple linear decomposition for value-based RL approaches. Our contribution complements many popular offline RL methods focused on distribution shift (e.g., BCQ [14]) and goes beyond pessimism- methods by leveraging domain knowledge. Compatible with any algorithm that has Q-function component, we expect our approach will lead to gains for offline RL problems with combinatorial action spaces where data are limited and when domain knowledge can be used to check the validity of theoretical assumptions.
Problem Setup
We consider Markov decision processes (MDPs) defined by a tuple M =
(S, (S, A,p., (o. 7), where S and A are the state and action spaces, p(s'ls, a) and r(s, are the transition and instantancous reward functions, (o(s) is the initial state distribution, and 7 E [0, 1] is the discount factor. A probability distribution
probabilistic policy π (als) specifies mapping from each state actions. For a deterministic policy, TTs) refers to the action with ([ass) = 1. The state-value function defined as V (s) E.,lE. [Σt=1 > 81 =s] The action-value function, Q (s, a). is defined by further restricting the action taken from the starting state. The goal of RI. is to find policy π = arg max, Ee-po! [VT [V((s) (or an approximation) that has the maximum expected performance.
2.1 Factored Action Spaces
While the standard MDP definition abstracts away the underlying structure within the action space A. in this paper, we explicitly express factored action space Cartesian product of D sub-action spaces, A 8 As A₁ x x Ap. We use a E AA denote each action, which can be written as a vector of sub-actions a = [ai...., .........p anl, with each O, E Aa. In general, a sub-action space can be discrete or continuous, and the cardinalities of discrete sub-action spaces are not required to be the same. For clarity of analysis and illustration, we consider discrete sub-action spaces in this paper.
2.2 Linear Decomposition of Q Function
The traditional factored MDP literature almost exclusively considers state space factorization [15]. In contrast, here we capitalize on action space factorization to parameterize value functions. Specifically, approach considers linear decomposition of the Q function, illustrated in Figure lb:
Qa (s, a) Eb ga(s, az).
(1)
Each component (a(8, aa) in the summation allowed to condition on the full state space s and only one sub-action ag. While similar forms of decomposition have been used in past work, there are key

---

differences in how the summation components are parameterizzed. In the multi-agent RL literature, each component qa(sa, 0.z) can only condition on the corresponding state space of the d-th agent [e.g.. 8, 9]. The decomposition in Eqn. (1) also differs from related form of decomposition considered by Juozapaitis et al. [12] where each component (a(8, a) can condition on the full action a. To the best of our knowledge, we are the first to consider this specific form of Q-function decomposition backed by both theoretical rigor and empirical evidence; in addition, we are the first to apply this idea to offline RL. We discuss other related work in Section 5.
3 Theoretical Analyses
In this section, we study the theoretical properties of the linear Q-function decomposition induced by factored action spaces. We first present sufficient and necessary conditions for our approach to yield unbiased estimates, and then analyze settings in which our approach can reduce variance without sacrificing policy performance when the conditions are violated. Finally, we discuss how domain knowledge may be used check the validity of these conditions, providing examples in healthcare.
3.1 Sufficient Conditions for Zero Bias
If we consider the total return of D MDPs running in parallel, where each MDP is defined by their respective state space Sy and action space Aa. then the desired linear decomposition holds for the MDP defined by the joint state space 8 D Sa and joint action space 2 A. (formally discussed in Appendix I. However, this relies on an explicit, known state space factorization, limiting its applicability. In contrast, we now present a generalization that forgoes the explicit factorization of the state space by making use of state abstractions. Intuitively, the MDP should have some implicit factorization, such that it is homomorphic to D parallel MDPs. It is, however, not a requirement that this factorization is known, as long as it exists.
Theorem Given an MDP defined by S,A.p. p. and a policy π :S + A(A), where A = X d As is factored action space with D sub-action spaces, if there exists D unique corresponding state abstractions Φb = (o1...., where da : S + Zd. 2d = oa(s). 2/2 = &a(s'), such that for all a, s' the following holds:
(()
Ila pa(=zlza,
a) = =E (dlaad)
(3)
T(als) = =Ild- Ta(aalza)
(4)
for some Pa : Za x As + A(Z)), ra : Za x Ag + R. and Na Za - A(Aa), then the Q-function of policy π can be expressed as Q* = Ed= qa(s, ad).
In Appendix B.2, we present an induction proof of Theorem 1. Since every assumption is used in key steps of the proof, we conjecture that the sufficient conditions cannot be relaxed in general. Consequently, if the sufficient conditions are satisfied, then using Eqn. (1) to parameterize the Q-function leads to zero approximation error and results in an unbiased estimator. Note that this does not require knowledge of Φ. To highlight the significance of Theorem 1, we present the following example, in which the state space cannot be explicitly factored, yet the linear decomposition exists (additional examples probing the sufficient conditions can be found in Appendix C).
Example troo-dimensional chains with abstractions). The factored action space shown in Figure 2a, A = Ag x Ay, is the composition of two binary sub-action spaces: A. = {+,, + leading the agent to move left or right, and Ay (L. 1} leading the agent to move down or up. Thus, A consists of [Uxx, ay] leads the agent to move diagonally.
four actions, where each action
Consider the MDP in Figure 2b with action space A. The state space S {so.o, s0,1, $1,0, $1,13 contains 5 different states; subscripts indicate the abstract state vector under &b = [0x, Øy (c.g. $0,1 and 50,1 are two different raw states but are identical under the abstraction, (550.1)
((50.1 = [20,?, 22. There does not exist an explicit state space factorization such that S = S2 Sy. One can check that Eqns. (2) and (3) are satisfied by comparing the raw transitions and rewards against the abstracted version (e.g., action / from $o.9 moves both > (under dx) and 1 (under Φy)) to $1.1 and receives the sum of the two rewards, 1 + 1 2). For Eqn. (4) to hold, the policy must take the same action from 80.1 and 80.1- In Figure 2c, we show the linear decomposition of the Q-function for one such policy where Theorem applies, under which the evolution of the MDP can seen as two chain MDPs running in parallel (also in Figure 2b).
3

---

(a)
exactly the Ag x Ay
(b)
Up
(Left, Upl
|Right, Upl
Abstract MDP in y-direction
Fu
Right
Ag
Original MDP
Down
Left Down
(Right, Down)
RM
Roa
(c)
[0x By]
Gx(50,0, (x)
(50.0. (y)
(" a)
0.9
1.8
0.9
1.9
0.9
1.9
Abstract MDP x-direction
4x
Figure 2: (a) The composition of sub-action spaces Av and Ay results in A Ax Ay depicted by outgoing arrows exiting the corners of each state (denoted by [).. The corner from which the action exits encodes the direction. (b) An MDP with 5 states and actions of the factored action space A. For example, action /= |+, 1 from 50.0 moves the agent both right (-+) and up (f), to $1,1. Under abstractions [ dy], this MDP can be mapped to two abstract MDPs (with action spaces Ag and Ay. respectively). The abstract state spaces are 2, {20,?, 21.?) and Ey = { 22,1 respectively, where indicates the coordinate ignored by the abstraction. $1,1 is an absorbing state whose outgoing transition arrows are not shown. Taking action ~= [ from $0.0 leads to $0,1 with probability and $ with probability (1 - p) (denoted in green). Actions taken by a determinisitic policy are denoted by bold blue arrows. takes the same action y= [+) from 50, and 80, Nonzero rewards are denoted in red. (c) Linear decomposition of Q" for 80,0 with respect the factored
action space (7 = 0.9). Similar decompositions for other states also exist (omitted for space).
3.2 Necessary Conditions for Zero Bias
In Appendix B.5, derive necessary condition for the linear parameterization to be unbiased. Unfortunately, the condition is not verifiable unless the exact MDP parameters are known; this highlights the non-trivial nature of the problem. One may naturally question whether the sufficient conditions (which are arguably more verifiable practice) must hold (i.e., are necessary) for the linear parameterization to be unbiased. Perhaps surprisingly. none of the conditions necessary. We state the following propositions and provide justifications through a set of counterexamples below and in Appendix C.
Proposition 2. There exists an MDP M and a policy for which QM decomposes as Eqn. (1) but the transition function p of M does not satisfy Eqn. (2).
Proposition 3. There exists an MDP M and a policy for which QM decomposes as Equ. (1) but the reward function r of M does satisfy Eqn. (3).
Proposition 4. There exists an MDP M and a policy π for which QM decomposes as Eqn. (1) but the policy æ does satisfy Eqn. (4).
Example 2 (Modified two-dimensional chains). Figure 3, all conditions in Theorem are violated, yet for each state, there exists linear decomposition of Q-values (see Appendix C).
Figure 3: This MDP is similar to Example (except it does not have state $0.1) and we consider the same abstractions ø = (dx, 0y]. The Q-function and decomposition as in the previous example. However, none of the conditions in Theorem I are satisfied. (a) The transition function does not satisfy Eqn. (2) because action >= [-+, 1] from 80, does not move right (- under dx) 81,1 and instead moves back to state s0 (b) The reward function does not satisfy Eqn. (3) as the reward of (1 7) for action /= [-+, t) from $0,1 is not the sum of +1 (- from 20,? under Φx) and 0 (f from 2?,1 under dy ). (c) The policy does not satisfy Eqn. (4) takes different sub-actions from 2g.? under dx (K from 80.0 specifies t-, whereas / from 80, specifies ->)-
Pin
F1.2
Therefore, while Theorem 1 imposes rather stringent set of assumptions on the MDP structure (transitions, rewards) and the policy, violations of these conditions do not preclude the linear parame- terization of the Q-function from being unbiased estimator.

---

3.3 How Does Bias Affect Policy Learning?
When the sufficient conditions do not hold perfectly, using the linear parameterization in Eqn. (1) to fit the Q-function may incur nonzero approximation eтrоr (bias). This can affect the performance of the learned policy: in Appendix B.3, we derive error bounds based on the extent of bias relative to the sufficient conditions in Theorem 1. Despite this bias, our approach always leads to reduction in the variance of the estimator. This gives us an opportunity to achieve better bias-variance trade-off, especially given limited historical data in the offline setting. In addition, as we will demonstrate, biased Q-values do not always result in suboptimal policy performance, and we identify the characteristics of problems where this occurs under our proposed linear decomposition.
3.3.1 Bias- Variance Trade-off
While the amount of bias incurred depends on the problem structure, the benefit of variance reduction is immediate. Intuitively, to learn the Q-function of tabular MDP with state space S and action space A d=1 Ad, the linear parameterization reduces the number of free parameters from 1=1 Aal) to |S S |Aal D + 1) (see Appendix B.4). This reduces the
[ = I - hypothesis class from exponential in D to linear in D. To analyze variance reduction, compare the bounds on Rademacher complexity [16-18] of the Q-function approximator using the factored action space with that of the full combinatorial action space (formally discussed in Appendix B.6).
Proposition 5. Using the linear Q-function decomposition for the factored action space in Eqm. (1) has a smaller lower bound on the empirical Rademacher complexity compared to learning the Q-function in the combinatorial action space.
Proposition 5 shows that our linear Q-function parameterization leads to a smaller function space, which implies a varianee estimator. Hence, our factored-action approach can make more efficient use of limited samples, leading to an interesting bias-variance trade-off that is especially beneficial for offline settings with limited data.
3.3.2 Bias A Suboptimal Performance
Even in the presence of bias, an inaccurate Q-function may still correctly identify the value- maximizing action (Proposition 6). While this statement is generally true, this section, we identify when this occurs specifically given our linear decomposition based on factored action spaces. To focus analysis on the most interesting aspects unique to our approach, we consider a bandit setting; extensions to the sequential RL setting are possible by applying induction similar to the proof for the main theorems (Appendices B.1 and B.2).
Proposition 6. There exists an MDP with the optimal Q* and its approximation Q parameterized in the form of Eqn. (1), such that Q + Q" and yet arg max, ((a) arg max, (a).
Justification. Consider a 1-step bandit problem with a single state and the same action space as before, A Ax x Ay- Taking an action terminate immediately. Since there are no transitions, the Q-values of any policy are simply the immediate reward from each action, Q(a) r(a). We assume the reward function defined as in Figure 4a (Appendix B.7 describes a procedure to standardize an arbitrary reward function).
[ax. ay] leads the agent to move diagonally and
(a)
(b)
Figure (a) two-dimensional bandit problem with action space A.. Rewards are denoted for each arm. (b) Learning using the linear Q decomposition approach corresponds to system linear equations that relates the reward of each arm. The parameter Timeast is dropped in our linear approximation, leading to omitted-variable bias. (c) Solving the system results in an approximate value function Q, which does not equal to the true value function Q" unless f = 0.
1
+Tawn
Tun
Toge Town
Yage
(c)
True Value Function
Linear Approximation
[2)
0
(Q())
(*(S)
(((s)
Q
1 +
(2)())
+
5

---

Applying our approach amounts to solving for the parameters Tet. laag:: "cown Tup of the linear system in Figure 4b, while dropping the interaction term I m'maa resulting in a form of vaitted-variable bias [19). Solving the system gives the approximate value function where the interaction term , appears in the approximation Q for all arms (Figure 4c, details in Appendix B.8).
Note that Q Q" only when 3 = 0. i.e., there is no interaction between the two sub-actions We first consider the family of problems with a and 8 E [-4, 4]. In Figure Sa, we measure the value approximation error RMSE(Q". Q). as well as the suboptimality V - -V* = maxa Q*(a) (arg max, ((a)) of the greedy policy defined by Q as compared to n.. As expected, when B = 0, Q is unbiased and has zero approximation error. When / # O, Q is biased and RMSE > 0; however, for 3 2 -1, Q corresponds to a policy that correctly identifies the optimal action.
We further investigate this phenomenon considering both a, 8 E [-4, (to show all regions with interesting trends), measuring RMSE and suboptimality in the same way as above. As shown in Figure 5b, the approximation error is zero only when 3 regardless of a. However, in Figure Sc, for a wide range of a and B settings, suboptimality is zero; this suggests that in those regions, even in the presence of bias (non-zero approximation error), our approach leads to an approximate value function that correctly identifies the optimal action. The irregular contour outlines multiple regions where this happens; one key region is when the two sub-actions affect the reward in the same direction (i.e.,e a 2 0) and their interaction effects also affect the reward in the same direction (i.e., 3 2 0).
(a)
(b)
RMSE(2'. Q11
(c)
Suboptimallyy V -
6 1.0
10
38
Unblased
S.bpptrratt
34
0.5
Opimal
33
44
9
Figure 5: (a) The approximation error and policy suboptimality of our approach for the bandit problem in Figure 4a, for different settings of 33 when a = 1. The Q-value approximation is unbiased only when B = 0, but the corresponding approximate policy is optimal for wider range of [ > -1. (b-c) The approximation error and policy suboptimality of our approach for the bandit problem in Figure 4a, for different settings of and B. The Q-value approximation is unbiased only when B = O, but the corresponding approximate policy is optimal for a wide range of a and 3 values. The highlighted region of zero suboptimality corresponds to a 2 0 and 8 > O.
3.4 Practical Considerations: Are these Assumptions too Strong?
Based on our theoretical analysis, strong assumptions (Section 3.1) on the problem structure (though not necessary, Section 3.2) are the only known way to guarantee the unbiasedness of our proposed linear approximation. It is thus crucial to understand the applicability (and inapplicability) of our approach in real-world scenarios. Exploring to what extent these assumptions hold in practice is especially important for safety-critical domains such as healthcare where incorrect actions (treatments) can have devastating consequences. Fortunately, RL tasks for healthcare are often equipped with significant domain knowledge, which serves as a better guide to inform the algorithm design than heuristics- driven reasoning alone [20, 5, 9].
Oftentimes, when clinicians treat conditions using multiple medications at the same time (giving rise to the factored action space), is because each medication has different "mechanism of action," " resulting in negligible or limited interactions. For example, several classes of medications are used in the management of chronic heart failure, and each has a unique and incremental benefits on patient outcomes [21]. Problems such as this satisfy the sufficient conditions in Section 3.1 in spite of a non-factorized state space. Moreover, any small interactions would have a bounded effect on RL policy performance (according to Appendix B.3).
Similarly, in the management of sepsis (which we consider in Section 4.2), fluids and vasopressors affect blood pressure to correct hypotension via different mechanisms [22]. Fluid infusion increases "preload" by increasing the blood return to the heart to make sure the heart has enough blood to pump out [23]. In contrast, common vasopressors (c.g., norepinephrine) increase "inotropy" by

---

stimulating the heart muscle and increase peripheral vascular resistance to maintain perfusion to organs [24, 25]. Therefore, while the two treatments may appear to operate on the same part of the state space (e.g., they both increase blood pressure), in general they are not expected to interfere with each other. Recently, there has also been evidence suggesting that their combination can better correct hypotension [26]. which places this problem approximately in the regime discussed in Section 3.3.2.
In offline settings with limited historical data, the benefits of a reduction in variance can outweigh any potential small bias incurred in the scenarios above and lead overall performance improvement (Section 1. However, our approach is not suitable if the interaction is counter to the effect of the sub-actions (e.g., two drugs that raise blood pressure individually, but when combined lead to a decrease). In such scenarios, the resulting bias will likely lead to suboptimal performance (Sec- tion 3.3.2). Nevertheless, many drug-drut interactions are known and predictable [27-30]. In such cases, one can either explicitly encode the interaction terms or resort back to a combinatorial action space (Appendix B.9). While focus on healthcare, there other domains in which significant domain knowledge regarding the interactions among sub-actions is available, e.g., cooperative multi- agent games in finance where there is higher payoff if agents cooperate (positive interaction effects) or intelligent tutoring systems that teach basic arithmetic operations as well as fractions (which are distinct but related skills). For these problems, this knowledge can and should be leveraged.
4 Experimental Evaluations
We apply our approach to two offline RL problems from healthcare: a simulated and a real-data problem, both having an action space that is composed of several sub-action spaces. These problems correspond to settings discussed in Section 3.4 where we expect our proposed approach to perform well. In the following experiments, we compare our proposed approach (Figure lb), which makes assumptions regarding the effect of sub-actionn common baseline that considers combinatorial action space (Figure la).
combination with other sub-actions, against a
4.1 Simulated Domain: Sepsis Simulator
Rationale. First, we apply our approach to a simulated domain modeled after the physiology of patients with sepsis [13]. Although the policies are learned "offline," simulated setting allows us to evaluate the learned policies in an "online" fashion without requiring offline policy evaluation (OPE).
Setup. Following prior work [31], a state is represented by feature vector x(s) E 10, that uses a one-hot encoding for each underlying variable (diabetes status, heart rate, blood pressure, oxygen concentration, glucose; all of which
discrete). The action space is composed of 3 binary
treatments: antibiotics, vasopressors, and mechanical ventilation, such that A = Aahx X Avaso x Ame- with Aahx Avaso = Amy {0, 1} and |A| 2 8. Each treatment affects certain vital signs and may raise or lower their values with pre-specified probabilities (precise definition in [31]). A patient is discharged alive when all vitals are normal and all treatments have been withdrawn; death occurs if 3 or more vitals are abnormal. Rewards are sparse and only assigned at the end of each cpisode (+1 1 for survival and - for death), after which the system transitions into the respective absorbing state. Episodes are truncated at maximum length of 20 following [13] (where no terminal reward is assigned). Here, the MDP partly satisfies the sufficient conditions outlined in Section 3. For example, oxygen saturation (which can be seen as state abstraction) is only affected by mechanical ventilation, whereas heart rate only affected by antibiotics. However, blood pressure is affected by both antibiotics and vasopressors, meaning the effects of these two sub-actions are not independent.
Offline learning. First, we generated datasets with different sample sizes following different behavior policies. We ran fitted Q-iteration for up to 50 iterations using a neural network function approximator, selecting the early-stopping iteration based on ground-truth policy performance. Each setting of sample size and behavior policy was repeated 10 times with different random seeds. Additional details are described in Appendix D.I.
Figure 6 compares median performance of the proposed approach vs. the baseline over the 10 runs (error bars interquartile ranges). We considered behavior policies that take the optimal action with probability and select randomly among non-optimal actions with probability 1 p.
Results.
How does sample size affect performance? We first look at uniformly random behavior policy (p ===10 1/|A|| = 0.125, Figure 6 center). As expected, larger sample sizes (i.e., more training episodes) lead to better policy performance for both the baseline and proposed approaches. For smaller
7

---

Performance (policy value)
(c 0.1)
p (c 0.50
P
0.0
0.4
10a
10a
10
10
10
10
10a
10a
10
10-
102
10*
10a
102
10a
Sample size (training episodes)
Baseline
Proposed
Figure 6: Performance the sepsis simulator across sample sizes and behavior policies. Plots display the performance over 10 runs, with the trend lines showing medians and error bars showing interquartile ranges. P is the probability of taking the optimal action under the behavior policies used to generate offline datasets from the simulator. The left two plots show two e-greedy policies (p 0.125; conversion: p (1 -- E) + e/ e/]A)); the middle plot shows a uniformly random policy (p = 15);; the right two plots show two policies that undersample the optimal action, < 0. 0.:;;; from left right, p decreases. Across different data distributions, our proposed approach outperforms the baseline at small sample sizes, and closely matches baseline performance at large sample sizes. Dashed lines denote the value of the optimal policy, which equals 0.736.
sample sizes (< 5000), the proposed approach consistently outperforms the baseline. As sample size increases further, the performance gap shrinks and eventually the baseline overtakes our proposed approach. This is because variance decreases with increasing sample size but the bias incurred by the factored approximation does not change. Once there are enough samples, reductions in variance are longer advantageous and the incurred bias dominates the performance. Overall, this shows that our approach is promising especially for datasets with limited sample size.
How does behavior policy affect performance? As anneal the behavior policy closer to the optimal policy (p > 0.1 Figure 6 left two), we reduce the randomness in the behavior policy and limit the amount of exploration possible at the same sample size. The same overall trend largely holds. On the other hand, when the probability of taking the optimal action is less than random (p < 0.125, Figure 6 right two), the proposed approach achieves better performance than the baseline with an even larger gap for limited sample sizes (S 10²). Without observing the optimal actions (p = 0), the baseline performs relatively poorly, even for large sample sizes. In comparison, accounts for relationships among actions to some extent and is thus able to better generalize to the unobserved and underexplored optimal actions, thereby outperforming the baseline.
approach
Takeaways. challenging situation where our theoretical assumptions do not perfectly hold, our proposed approach matches outperforms the baseline, especially for smaller sample sizes.
4.2 Real Healthcare Data: Sepsis Treatment in MIMIC-III
Rationale. We apply our method to real-world example learning optimal sepsis treatment policies for patients in the intensive care unit. Acknowledging the challenging nature of OPE for quantitative comparisons [32, 33], here we qualitatively inspect the learned policies using clinical domain knowledge.
Setup. Originally introduced by [2], we use the improved formulations of this task as per [34] and [35]. After applying the specified inclusion and exclusion criteria to the de-identified MIMIC-III database [36), we obtained cohort 19,287 patients and performed 70/15/15 split for training, validation and testing. For each patient, their data include 10 time-invariant demographic and contextual features and a 33-dimensional time series collected at 4h intervals, consisting of measurements from up to 24h before until up 48h after sepsis onset. We used recurrent neural network (RNN) with long short-term memory (LSTM) cells to create an approximate information state [37] to summarize the history into de-dimensional embedding vector. A terminal reward of 100 is assigned for 48h survival and otherwise. Intermediate rewards are all 0. y for learning is 0.99 and for evaluation is 1. Actions pertain treatment decisions in each 4h interval, representing total volume of intravenous (IV) fluids and amount of vasopressors administered, resulting in 5 x factored action space.
Offline learning. After learning the state representations, apply variants of discrete-action batch- constrained Q-learning (BCQ) [38, 14], where the baseline uses the combinatorial action space and the proposed approach incorporates the linear decomposition induced by the factored action space.

---

98-
Eatimated Pakcy
96
Policy
Baseline BCQ
Factored BCQ
Clinician
94
Baseline
Test WIS
90.44 # 2.44 178.32 + 11.42
91.62 = 2.12 178.32 + 11.96
90.29 + 0.51
Propesed
Test ESS
2894
92-
Observed
% agreement with clinician
62.42%
62.37%
57.16%
200
Eflctivv Sample Size
Figure 7: Left- Pareto frontiers of validation performance for the candidate policies (all points plotted in Figure 16). The shaded region does not meet the ESS cutoff of 2 200. The red circles indicate the selected models (based on best validation WIS) for baseline and proposed (both have a BCQ threshold of T = 0.5). Right- Performance on test set, + standard errors from 100 bootstraps.
(a)
Clinician
Baselra BCQ
Factore BCQ
75003
State Vector Hiterogannily
NVuud (n.4M
1 10m0r1s 1
1400
103
167
203
277
NN 0000 (muths
0
0
0
Nu (m you
0
10
10083
2,103
0.175
0.103
0.121
0.
2364
175
1044
6.157
0:164
01103
0.170
854
826
190
22235
538
2508
2200:
253
1801
154
10099
0.108
0.111
0.155
106
48
85
75
0929
01:66
0
52
8.192 0.150
0.1.3
0.144
0.797
V1.4
ViA
SINCONE 37:000 13.44
(gglagrren)
Lug
Vascprase (ughigraan
Figure& 8: (a) Qualitative comparison of policies. (b) Per-action state heterogeneity, measured as the standard deviation of all state embeddings from which a particular action is observed in the dataset, averaged over state embedding dimensions. Actions with higher IV fluid doses exhibited greater heterogeneity in the observed states from which those actions were taken (by the clinician policy).
The -entworrks were trained for maximum of 10, 000 iterations, with checkpoints saved every 100 iterations. We performed model selection [31] over the saved checkpoints (candidate policies) by evaluating policy performance using the validation set with OPE. Specifically, we estimated policy value using weighted importance sampling (WIS) and measured effective sample size (ESS), where the behavior policy was estimated using kt nearest neighbors in the embedding space. Following previous work [39], the final policies were selected by maximizing validation WIS with ESS of 2 200 (we consider other thresholds in Appendix D.3), for which we report results on the test set.
Results. We visualize the validation performance over all candidate policies. Figure 7-1eft shows that the performance Pareto frontier (in terms of WIS and ESS) of the proposed approach generally dominates the baseline.
Quantitative comparisons. Evaluating the final selected policies the test (Figure 7-right) shows that the proposed factored BCQ achieves a higher policy value (estimated using WIS) than baseline BCQ at the same level of ESS. In addition, both policies have a similar level of agreement with the clinician policy, comparable to the average agreement among clinicians.
Qualitative comparisons. In Figure 8a, we compare the distributions of recommended actions by the clinician behavior policy, baseline BCQ and factored BCQ, as evaluated on the test set. While overall the policies look rather similar, in that the most frequently recommended action corresponds to low doses of IV fluids <500mL with no vasopressors, there are notable differences for key parts of the action space. In particular, baseline BCQ almost never recommends higher doses of IV fluids >500 ml, either alone or in combination with vasopressors, whereas both clinician and factored BCQ recommend IV fluids >500 mL more frequently. These actions typically used for critically ill patients, for whom the Surviving Sepsis Campaign guidelines recommends up to >1 of fluids [40]. We hypothesize that this difference is due higher level of heterogeneity in the patient states for which actions with high IV fluid doses were observed, compared to the remaining actions with lower doses of IV fluids. To further understand this phenomenon, we measure the per-action state heterogeneity in the test set by computing, for each action, the standard deviation (averaged over the embedding dimensions) of all RNN state embeddings from which that action is taken according to the behavior policy. As shown in Figure 8b. actions with higher IV fluids generally have larger standard deviations, supporting our hypothesis. The larger heterogeneity combined with lower sample sizes makes it difficult for baseline BCQ to correctly infer the effects of these actions, as it does not leverage the relationship among actions. contrast, our approach leverages the factored action space and can thus make better inferences about these actions.

---

Takeaways. Applied to real clinical data, our proposed approach outperforms the baseline quantita- tively and recommends treatments that align better with clinical knowledge. While promising, these results are based in part on OPE, which has many issues [32, 33J. We stress that further investigation and close collaboration with clinicians are essential before such RI. algorithms are used in practice.
Related Work
For many years, the factored RL literature focused exclusively on state space factorization [15, 41-43]. More recently, interest in action space factorization has grown, as RI. is applied in increasingly more complex planning and control tasks. In particular, researchers have previously considered the model-based setting with known MDP factorizations in which both state and action spaces factorized [44-47]. For model-free approaches, others have studied methods for factored actions with a policy component (i.e.. policy-based or actor-critic) [48, 20, 49-51J. In contrast, our work considers value-based methods as those have been the most successful in offline RL [52].
Among prior work with value-based component (e.g., Q-network), the majority pertains to multi- agent [8-10, 53, 54] or multi-objective [51] problems that impose known, explicit assumptions on the state space or the reward function. Notably, Son et al. [10] established theoretical conditions for factored optimal actions (called "IndviduaalGGlobalMaay
for multi-agent RL and motivated subsequent works [55, 56]; their result differs from our contribution
the unbiasedness of factored Q-functions (instead of actions) for single-agent RL. In the online setting for single-agent deep RL, Sharma al. [20] and Tavakoli et al. [5] incorporated factored action spaces into Q-network architecture designs, but did not provide a formal justification for the linear decomposition. Others have empirically compared various "mixing" functions to combine the values of sub-actions [20, 9]. In contrast, while our work only considers the linear decomposition function, we examine its theoretical properties and provide justifications for using this approach in practical problems, especially in offline settings. Our linear Q-decomposition is related to that of Swaminathan et al. [57] who also applied a linearity assumption for off- off-policy evaluation, but for combinatorial contextual bandits rather than RL. Our insights on the bias-variance trade-off is also related to concurrent work by Saito and Joachims [58] who proposed efficient off- aff-policy evaluation for bandit problems with large (but not necessarily factored) action spaces. In appendix Table 1, we further outline the differences of our work compared to the existing literature.
Finally, the sufficient conditions we establish are related to, but different from, those identified by Van Seijen et al. [59] and Juozapaitis al. [12] who considered reward decompositions in the absence of factored actions. Related, Metz et al. [60] proposed an approach that sequentially predicts values for discrete dimensions of a transformed continuous action space. but assume an a priori ordering of action dimensions, which we do not; Pierrot et al. [50] studied a different form of action space factorization where sub-actions are sequentially selected in an autoregressive manner. Complementary to our work, Tavakoli et al. [6] proposed to organize the sub-actions and interactions as hypergraph and linearly combining the values; our theoretical results on the linear decomposition nonetheless apply to their setting where the sub-action interactions are explicitly identified and encoded.
6 Conclusion
To better leverage factored action spaces in RL, we developed an approach to learning policies that incorporates a simple linear decomposition of the Q-function. We theoretically analyze the sufficient and necessary conditions for this parameterization to yield unbiased estimates, study its effect on variance reduction, and identify scenarios when any resulting bias does not lead to suboptimal performance. We also note how domain knowledge may be used inform the applicability of our approach in practice, for problems where any possible bias is negligible or does not affect optimality. Through empirical experiments on two offline RL problems involving a simulator and real clinical data, we demonstrate the advantage of our approach especially in settings with limited sample sizes. We provide further discussions on limitations, ethical considerations and societal impacts in Appendix A. Though motivated by healthcare, our approach could apply more broadly to scale RL to other applications (e.g., education) involving combinatorial action spaces where domain knowledge may be used to verify the theoretical conditions. Future work should consider the theoretical implications of linear Q decompositions when combined with other offline RL-specific algorithms [52]. Given the challenging nature of identifying the best treatments from offline data, our proposed approach may also be combined with other RI. techniques that do not aim to identify the single best action (c.g., learning dead-ends [61] or set-valued policies [34]).
10

---

Acknowledgments
This work was supported by the National Science Foundation (NSF; award 115-1553146 to JW; award IIS-2007076 to FDV; award IIS-2153083 to MM) and the National Library of Medicine of the National Institutes of Health (NLM; grant ROILM013323 to JW and MWS). The views and those of the authors and should not be interpreted necessarily representing the official policies, either expressed or implied, of the National Science Foundation, nor of the National Institutes of Health. This work supported, in part, by computational resources and services provided by Advanced Research Computing, a division of Information and Technology Services (TTS) at the University of Michigan, Ann Arhor. The authors would like thank Adith Swaminathan, Tabish Rashid, and members of the MLD3 group for helpful discussions regarding this work, well as the reviewers for constructive feedback.
conclusions in this document
Data and Code Availability
The code for all experiments is available at
The sepsis simulator is based on prior work with public implementation at https://github.com com/ clinicalm//gumbel-mne max-sem. The MIMIC-II database used in the real-data experiments of this paper www.o//tttrrgooooone/mmmmiiiiiii .4/. The
is publicly available through the PhysioNet website: cohort definition, extraction and preprocessing code are based on prior work with publicly available implementation at https//jiitaa.cmmmmirrsottiiniii
References
[1] Damien Ernst, Guy-Bart Stan, Jorge Goncalves, and Louis Wehenkel. Clinical data based optimal STI strategies for HIV: a reinforcement learning approach. In Proceedings of the 45th IEEE Conference on Decision and Control, pages 667-677 677. IEEE, 2006.
[2] Matthieu Komorowski, Leo A Celi, Omar Badawi, Anthony C Gordon, and A Aldo Faisal. The artificial intelligence clinician learns optimal treatment strategies for sepsis intensive care. Nature medicine, 24 (11):1716-17) 17200 2018.
[3] Niranjani Prasad, Li Fang Cheng. Corey Chivers, Michael Draugelis, and Barbara E Engelhardt. A reinforcement learning approach weaning of mechanical ventilation in intensive care units. In Conference on Uncertainty in Artificial Intelligence (UAI), 2017.
[4] Sonali Parbhoo, Mario Wieser, Volker Roth, and Finale Doshi-Velez. Transfer learning from well-curated to less-resourced populations with HIV. In Machine Learning for Healthcare Conference, pages 589-606 -699. PMLR, 2020.
Arash Tavakoli, Fahio Pardo, and Petar Kormushev. Action branching architectures for deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018
[6] Arash Tavakoli, Mchdi Fatemi, and Petar Kormushev. Learning to represent action values as a hypergraph the action vertices. In International Conference Learning Representations, 2021. URL https: //openteviie..ee/fommmmkk/K.. s64FiXTv.
[7] Craig Boutilier. Planning, learning and coordination multiagent decision processes. In TARK, volume 96, pages 195-210. Citeseer, 1996.
[8] Peter Sunehag, Guy Lever, Audrunas Gruslys, Wojeiech Marian Czarnecki, Vinicius 7ambaldi, Max Jader- berg, Mare Lanctot, Nicolas Sonnerat, Joel 7. Leibo, Karl Tuyls, and Thore Graepel. Value-decomposition networks for cooperative multi-agent learning based on team reward. In Proceedings of the 17th Interna- tional Conference on Autonomous Agents and MulliAgent Systems, AAMAS, 2018.
[9] Tabish Rashid, Mikayel Samvelyan. Christian Schroeder, Gregory Farquhar, Jakob Foerster, and Shimon Whiteson. QMIX: Monotonic value function factorisation for deep multi-agent reinforcement learning. In International Conference on Machine Learning, pages 4295 4304. PMLR, 2018.
[10] Kyunghwan Son, Daewoo Kim, Wan Ju Kang, David Earl Hostallero, and Yung Yi. QTRAN: Learning to factorize with transformation for cooperative multi agent reinforcement learning. In International conference machine learning, pages 5887-5896. PMLR, 2019.

---

[11] Ming Zhou, Yong Chen, Ying Wen, Yaodong Yang, Yufeng Su, Weinan Zhang. Dell Zhang, and Jun Wang. Factorized Q-leariing for large-scale multi-agent systems. In Proceedings of the First International Conference on Distributed Artificial Intelligence, pages 1-7, 2019.
[12] Zoe Juozapaitis, Anurag Koul, Alan Fern. Martin Erwig, and Finale Doshi Velez. Explainable reinforce- ment learning via reward decomposition. In IJCAI/ECA/ Workshop on Explainahle Artificial Intelligence, 2019.
[13] Michael Oberst and David Sontag. Counterfactual off-policy evaluation with Gumbel-max structural causal models. In International Conference Machine Learning, pages 4881-4890, 2019.
[14] Scott Fujimoto. David Meger, and Doina Precup. Off-policc deep reinforcement learning without explo- ration. In International Conference on Machine Learning, pages 2052-2062. PMI.R, 2019.
[15] Daphne Koller and Ronald Parr. Computing factored value functions for policies in structured MDPs. In Proceedingy of the 16th International Joint Conference on Artificial Intelligence-Voluee Volume 2, pages 1332-1339, 1999.
[16] Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar, Foundations of machine learning. MIT press, 2018.
[17] Yaqi Duan, Chi Jin, and Zhiyuan Risk bounds and Rademacher complexity batch reinforcement learning. arXiv preprint arXiv:2103. 13883, 2021.
[18] Maggie Makar, Ben Packer, Dan Moldovan, Davis Blalock, Yoni Halpem, and Alexander D' Amour. Causally motivated shortcut removal using auxiliary labels. In International Conference Artificial Intelligence and Statistics, pages 739-766. PMLR, 2022.
[19] Jeffiey M Wooldridge. Introductory econometrics: A modern approach. Cengage learning, 2015.
[20] Sahil Sharma, Aravind Suresh, Rahul Ramesh, and Balaraman Ravindran. Learning factor policies and action-value functions: Factored action space representations for deep reinforcement learning. arXiv preprint arXiv:1705.07266 072699 2017.
[21] Michel Komajda, Michael Boehm, Jeffrey Borer, Ian Ford. Luigi Tavazzi, Matthieu Pannaux, and Karl Swedberg. Incremental benefit of drug therapies for chronic heart failure with reduced ejection fraction: a network meta-analysis. European journal of heart failure, 20(9):1315-1322, 2018.
[22] Jeffrey Gouts and Michael Matthay. Sepsis: pathophysiology and clinical management. BMJ, 353, 2016.
[23] Laurent Guérin, Jean-Louis Teboul, Romain Persichini, Martin Dres, Christian Richard, and Xavier Monnet. Effects of passive leg raising and volume expansion on mean systemic pressure and venous return in shock in humans. Critical Care, 19(1):1-9, 9, 2015.
[24] Olfa Hamzaoui, Jean-Frangoi Georger, Xavier Monnet, Hatem Ksouri, Julien Maizel, Christian Richard, and Jean-Louis Teboul. Early administration of norepinephrine increases cardiac preload and cardiac output in septic patients with life-threatening hypotension. Critical care, 14(4):1-9, 2010.
[25] Xavier Monnet, Julien Jabot, Julien Maizel, Christian Richard, and Jean-Louis Teboul. Norepinephrine increases cardiac preload and reduces preload dependency assessed by passive leg raising in septic shock patients. Critical care medicine, (9())6689 669-4444, 2011.
[26] Olfa Hamzaoui. Combining fluids and vasopressors: A magic potion? Journal of Intensive Medicine, 2021. ISSN 2667-100X. doi: https://do..org/100 jott://o...011....mmm.mmmmm 2021000000000000000....
[27] Teijo I Saari, Kari Laine, Mikko Neuvonen. Pertti J Neuvonen, and Klaus T Olkkola. Effect of voricona- zole and fluconazole the pharmacokinetics of intravenous fentanyl. European journal clinical pharmacology, 64:25533- 2008
[28] Pamela L Smithburger. Sandra Kane- Gill, and Amy Seybert. Drug drug interactions in the medical intensive care unit: an assessment of frequency, severity and the medications involved. International Journal of Pharmacy Practice, 20(6):402-408, 2012.
[29] WebMD. Drugs interaction checker- find interactions between medications. https//www.webmndcomm imeracttionchheckkeffefaall....u
[30] Epic. Lexicomp and medi-span for epic- drug reference information. http::/wwww.wotteeskwwwwccc.../ souttoon/eetcoomm//ouuuup//gggge referencee.
12

---

[31] Shengpu Tang and Jenna Wiens. Model selection for offline reinforcement learning: Practical considerations for healtheare settings. In Machine Learning for Healthcare Conference, pages 2-35. PMLR. 2021.
[32] Omer Gottesman, Fredrik Johansson, Joshua Meier, Jack Dent, Donghun Lee, Srivatsan Srinivasan, Linying Zhang, Yi Ding, David Wihl, Xuefeng Peng. et al. Evaluating reinforcement leamning algorithms in observational health settings. arXiv preprint arXiv::805, 1228,, 2018.
[33] Omer Gottesman, Fredrik Johansson, Matthieu Komorowski, Aldo Faisal, David Sontag, Finale Doshi- Velez, and Leo Anthony Celi. Guidelines for reinforcement learning in healthcare. Nature medicine, 25(1): 16-18, 2019.
[34] Shengpu Tang. Aditya Modi, Michael Sjoding, and Jenna Wiens, Clinician-in in-the-loop decision making: Reinforcement learning with near-optimal set-valued policies. In International Conference on Machine Learning, pages 9387-9396. PMLR, 2020.
[35] Taylor W Killian, Haoran Zhang, Jayakumar Subramanian, Mehdi Fatemi, and Marzyeh Ghassemi. An empirical study of representation learning for reinforcement learning in healthcare. In Machine Learning for Health NeurIPS Workshop, 2020.
[36] Alistair EW Johnson, Tom Pollard, Lu Shen, Li-wei Lchman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Lco Anthony Celi, and Roger G Mark. MIMIC-III, a freely accessible critical care database. Scientific data, 3(1):1-9, 2016.
[37] Jayakumar Subramanian and Aditya Mahajan. Approximate information state for partially observed systems. In 2019 IEEE 58th Conference on Dectsion and Control (CDC), pages 1629-1636. IEEE, 2019.
[38] Scott Fujimoto, Edoardo Conti, Mohammad Chavamzadeh, and Joelle Pineau. Benchmarking batch deep reinforcement learning algorithms. arXiv preprint arXiv:1910.01708, 2019.
[39] Yao Liu and Emma Brunskill. Avoiding overfitting the importance weights in offline policy optimization, 2022. URL https://openrevve..ee//oounm Didaarmmmmddal.lXSSSsrstii
[40] Laura Evans, Andrew Rhodes, Waleed Alhazzani, Massimo Antonelli, Craig Coopersmith, Craig French, Flávia R Machado, Lauralyn Mcintyre, Marlies Ostermann, Hallie C Prescott, et al. Surviving sepsis campaign: international guidelines for management of sepsis and septic shock 2021. Intensive care medicine, 47(11):1181 11241777777 2021.
[41] Carlos Guestrin, Daphne Koller, Ronald Parr, and Shohha Venkataraman. Efficient solution algorithms for factored MDPs. Journal of Artificial Intelligence Research, 19:399-468, 2003.
[42] Alexander L Strehl, Carlos Diuk, and Michael L Littman. Efficient structure learning in factored-state MDPs. 2007.
[43] Karina Valdivia Delgado, Scott Sanner, and Leliane Nunes De Barros. Efficient solutions factored MDPs with imprecise transition probabilities. Artificial Intelligence, 75)9-10):1498-1555 201
[44] Aswin Raghavan, Saket Joshi, Alan Fern, Prasad Tadepalli, and Roni Khardon. Planning factored action spaces with symbolic dynamic programming. In Twenty-Sixth AAAI Conference on Artificial Intelligence, 2012.
[45] Aswin Raghavan, Roni Khardon, Alan Fern, and Prasad Tadepalli. Symbolic opportunistic policy iteration for factored-action mdps. Advances in Neural Information Processing Systems, 26, 2013.
[46] Ian Oshand and Benjamin Van Roy. Near-optimal reinforcement learning in factored MDPs. Advances in Neural Information Processing Systems, 27:604-612, 2014.
[47] Yangyi Lu, Amirhossein Meisami, and Ambuj Tewari. Causal Markov decision processes: Learning good interventions efficiently. arXiv preprint arXiv:2102.00 07663, 2021.
[48] Brian Sallans and Geoffrey E Hinton. Reinforcement learning with factored states and actions. The Journal of Machine Learning Research, 5:1063-1088, 2004.
[49] Tom Van de Wiele, David Warde-Farley, Andriy Mnih, and Volodymyr Mnih. Q-learning in enormous action spaces via amortized approximate maximization. arXiv preprint vrXiv:2001.08//6 2020.
[50] Thomas Pierrot, Valentin Macé, Jean- Baptistt Sevestre, Louis Monier, Alexandre Laterre, Nicolas Perrin, Karim Beguir, and Olivier Sigaud. Factored action spaces in deep reinforcement learning, 2021. URL
13

---

[51] Thomas Spooner. Nelson Vadori, and Sumitra Ganesh. Factored policy gradients: Leveraging structure for efficient learning in MOMDPs. In Advances in Neural Information Prucessing Systems, 2021. URL
[52] Sergey Levine. Aviral Kumar, George Tucker, and Justin Fu, Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint (.Xix:2005.06442 2020.
[53] Laetitia Matignon, Guillaume Laurent, and Nadine Le Fort-Piat. Independent reinforcement leamners in cooperative Markov games: a survey regarding coordination problems. The Knowledge Engineering Review, 27(1):1-31, 2012.
[54] Ardi Tampuu, Tambet Matiisen, Dorian Kodelja, llya Kuzovkin, Kristjan Korjus, Juhan Aru, Jaan Aru, and Raul Vicente. Multiagent cooperation and competition with deep reinforeement learning. PLOS ONE, 12(4); 1-15, 04 2017. doi: 10. 1/..3all www.oaal.oooeee..e....e 01 72395. URL http:://doi.org 1/ooaa/aoooo 013395..
[55] Jianhao Wang, Zhizhou Ren, Terry Liu, Yang Yu, and Chongjie Zhang. QPLEX: Duplex ducling multi- agent Q-learning. In International Conference on Learning Representations, 2021. URL https///pcerreview.
[56] Jianhao Wang. Zhizhou Ren, Beining Han, Jianing Ye, and Chongjic Zhang. Towards understanding cooperative multi-agent Q-learning with value factorization. Advances Neural Information Prucessing Systems, 34:29142 42-221555 2021.
[57] Adith Swaminathan. Akshay Krishnamurthy. Alekh Agarwal, Miro Dudik, John Langford, Damien Jose. and Imed Zitouni. Off-policy evaluation for slate recommendation. Advances in Neural Information Processing Systems, 30, 2017.
[58] Yuta Saito and Thorsten Joachims. Off-policy evaluation for large action spaces via embeddings. arXiv preprint arXiv:2202..0000 16 2022.
[59] Harm Van Seijen. Mehdi Fatemi, Joshua Romoff, Romain Laroche, Tavian Barnes, and Jeffrey Tsang- Hybrid reward architecture for reinforcement learning. In Advances in Neural Information Processing Systems, 2017.
[60] Luke Metz, Julian Ibarz, Navdeep Jaitly, and James Davidson. Discrete sequential prediction of continuous actions for deep RI, 2018. URL
[61] Mehdi Fatemi, Taylor W. Killian, Jayakumar Subramanian, and Marzyeh Ghassemi. Medical dead-ends and learning to identify high-risk states and treatments. In Advances in Neural Information Processing Systems, 2021. URL
[62] Jinglin Chen and Nan Jiang. Information-theoreti considerations in batch reinforcement learning. In International Conference on Machine Learning, pages 1042-1051, 2019.
[63] Tengyang Xie and Nan Jiang. Batch value-function approximation with only realizability. International Conference on Machine Learning, pages 11 14404-1143.. PMIR, 2021.
[64] Wenhao Zhan, Baihe Huang, Audrey Huang, Nan Jiang, and Jason D Lee. Offline reinforcement learning with realizability and single-policy concentrability. arXiv preprint arXiv:2202. 04664, 2022.
[65] Lihong Li, Thomas Walsh, and Michael L. Littman. Towards a unified theory of state abstraction for MDPs. Proceedings of the 9th International Symposium on Artificial Intelligence and Mathematics, 4, 2006.
[66] Satinder P Singh and Richard C Yee. An upper bound the loss from approximate optimal-value functions. Machine Learning, 16(3):277-333, 1994.
[67] Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. Machine learning, 49(2):209-322, 2002. 2002.
[68] Michail Lagoudakis and Ronald Parr. Least-squares policy iteration. The Journal of Machine Learning Research, 4:1 107-1149, 2003.
[69] Pranjal Awasthi, Natalie Frank, and Mehryar Mohri, On the Rademacher complexity of linear hypothesis sets. arXiv preprint arXiv:2007.1 10455, 2020.
[70] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations. 2015.
[71] Volodymyr Maih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement leaming. Nature, 518(7540):529-555 533, 2015.
14

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] Section 3.4, Appendix A
(c) Did you discuss any potential negative societal impacts of your work? [Yes] Appx A
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] Sec Appx B
(b) Did you include complete proofs of all theoretical results? [Yes] Appendix
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material or as a URL)? [Yes] see Data Code Availability
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Sec 4.1.2.2, Appendix 1.1-D.2
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes] Fig
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [No] our experiments only involved small neural networks
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes] MIMIC-III by Johnson et al. [36], sepsisSim by Oberst and Sontag [13]
(b) Did you mention the license of the assets? [No] please refer to the website of the creator MIMIC-III
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you'l using/curating? [No] please refer to the website of the creator of MIMIC-III wwwpae/tt.rrg/ooee///imiiiiiiiie
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [Yes] Sec 4.2, MIMIC-III de-identified
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
15


Reviews:
[{"review_id": "suXPC19mRu7", "reviewer": "Reviewer_jzqP", "summary": "This paper proposes to decompose Q-function according to factorized action space. It conducts discussions on the bias-variance trade-off of the decomposed Q-function. Specifically, this paper provides sufficient conditions for zero bias and proves these conditions are not necessary. Furthermore, it shows that this form of Q-function decomposition leads to low variance because of the smaller lower bound on the empirical Rademacher complexity. Finally, it discusses what kinds of tasks can fulfill the sufficient conditions. The experiments on offline healthcare tasks demonstrate that the decomposed Q-function has the potential to help models improve their performance.", "questions": "The title of the article points out that this is an efficient offline RL approach and the sample efficiency is verified in the paper. However, I am still wondering whether the running time and computational cost of factored BCQ are significantly reduced compared to the baseline BCQ?\nThe comparison of the models in the experimental evaluation part may be insufficient. Whether the validity of linear decomposition of the Q-function is applicable to other models of offline RL, such as CQL[1], PQI[2], REM[3]?\nIn Figure 8a, why does factored BCQ almost never recommend IV fluids 1L-2L, which is inconsistent with clinician behavior policy?\n[1] Kumar, A., Zhou, A., Tucker, G., & Levine, S. (2020). Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33, 1179-1191.\n[2] Liu, Y., Swaminathan, A., Agarwal, A., & Brunskill, E. (2020). Provably good batch off-policy reinforcement learning without great exploration. Advances in neural information processing systems, 33, 1264-1274.\n[3] Agarwal, R., Schuurmans, D., & Norouzi, M. (2020, November). An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning (pp. 104-114). PMLR.\n", "limitations": "The authors have discussed the limitations of their proposed approach in the manuscript, which may cause risky results in healthcare.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: The paper is well-structured and easy to follow. This paper studies the guarantee of the unbiasedness of Q-function decomposition, which may inspire the researchers to focus on safe decomposition and improve performance in healthcare tasks. Besides, the proofs of bias-variance trade-off are provided in detail. On the sufficiency and necessity of these conditions, it validates them either by theoretical analysis or examples. \nWeakness: \n1. According to Figure 1, Q-function decomposition can help reduce some calculation steps in the model. However, this paper does not demonstrate the optimization of computational efficiency, especially when it comes to experiments.\n2. According to Section 3.4, it seems that the sufficient conditions are very hard to satisfy which may lead the model rather impractical. \n3. The lack of comparison to more algorithms makes the effectiveness of the decomposition somewhat less convincing.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "UcwnP5M8dvd", "reviewer": "Reviewer_yW3d", "summary": "This study improved and extended the standard reinforcement learning (RL) to factored action space. A form of linear Q-function decomposition was proposed to handle factored action space. The novel method has been analyzed from a theoretical perspective, where the authors discussed the sufficient and necessary conditions for unbiases and studied its effect on variance reduction. The authors also demonstrated the proposed method through empirical experiments based on a simulation study and real-world MIMIC data analysis. The novel method provides both theoretical insights and empirical evidence for RL practitioners to consider this simple linear decomposition approach to factored action space. ", "questions": "1.\tBased on experimental evaluation (Figure 7), I was unable to observe significant improvement (in terms of clinical, real-world applicability) when compared with baseline BCQ and factored BCQ (90.44±2.44 vs 91.62±2.12).Therefore, it might be hard to conclude that the proposed factored BCQ outperforms the baseline.\n\n2.\tBased on Figure 6, the authors also claim better performance for a small sample size. I have concerns on this claim due to the lack of sufficient evidence. The variation is large for a small sample size, and the selection of samples would affect the result. Also, I couldn’t find the authors’ discussions on why the proposed method would get higher performance for small sample size, especially from the theoretical perspective. \n\n3.\tThe authors state that domain knowledge may be used to inform the applicability of the approach in practice. However, very little information in the main text address on how domain knowledge can be involved. This is an important question for real-world applications, thus would like to suggest the authors to provide more details for justifications. \n\n4.\tIf the authors claim that the method works well only for a limited sample size, I would suggest the authors do some analyses to determine the sample size cut-off (or give a quantitative/analytical formula to calculate). When the sample size is large enough, the empirical research shows that the proposed method is less useful and effective as compared to the baseline. Therefore, knowing the sample size cut-off is essential. \n\n5.\tIf possible, adding another real-world use case would be ideal to justify the generalizability and applicability of the proposed method. \n", "limitations": "1.\tThe performance improvement brought by the proposed method seems marginal. I suggest the authors add some other performance indicators and provide additional discussions.\n\n2.\tThere is a lack of interpretability for the proposed method. The authors are suggested to elaborate on how to interpret the results at the individual level and how this can be applied to clinical practice.\n\n3.\tThere are some preconditions required for using this proposed method. I would suggest the author provide some discussions and give more guidelines for future researchers who would like to use your method. \n\n4.\tPlease also add more on the importance of your method and how your method can address the unmet need of current literature. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength: The research question is novel, and the authors found the unmet need of factored action spaces in the RL approach. Solving this problem will facilitate the use of offline RL in a broader application, especially in healthcare. The methodology is well-explained with details. This paper is complete with theoretical insights and empirical experiments. \n\nWeakness: The proposed method showed marginal improvement compared to the existing method. I also have some conservations on the applicability of the method (where many preconditions are required). The improvement has not addressed the key limitations of offline RL (i.e., how to improve exploration, discover high-reward regions and solve distributional shift problems). \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "I2jq73um9F4", "reviewer": "Reviewer_cdHc", "summary": "This paper considers the use of a factored Q-function representation for reinforcement learning that takes advantage of a factored action space in certain \"appropriate\" MDPs, where \"appropriate\" is characterized both theoretically and qualitatively/practically in the paper. The authors show empirically that their factored representation can help in low-data healthcare settings, even when it is not a perfect match for the environment. It is argued that this is because it allows for a better bias-variance trade-off and provides a structural basis to generalize outside of the empirical distribution.", "questions": "A useful tactic in cases where we have an \"approximately correct\" structure, as you have in your experiments is to use a residual (even LSTMs / Resnets are examples of this, where identity is the approximately correct structure). So here, you could have $Q^\\pi(s, \\textbf{a}) = \\sum q_d(s, a_d) + \\mathcal{R}(s, a)$, where the residual function $\\mathcal{R}$ can capture interactions between the components (and can be heavily regularized in cases of low data). It may be interesting for you to consider this approach in your applications. Not really a question, I know. \n\nI guess my most relevant question is: have you looked at the FMDP literature at all? I don't see it mentioned anywhere in the paper. I wonder if your Theorem can be related to some of their representation theorems by a simple transform of the state space (that entangles the previously disentangled state features). Could be interesting / worth looking into. ", "limitations": "I think this paper adequately addressed the limitations and potential negative societal impact of their work. The authors discuss the stringent conditions of their representations at length, and their experiments are performed in environments that *do not* satisfy their sufficient conditions. ", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This is a good paper. At first I was put off by the specificity of the representation---it really did seem to me that the factored action space representation used here was too specific to be of any practical use, but I felt that the authors did well to combat this and demonstrated applicability in healthcare at the very least. The presentation is clear and well done, and the development is logically consistent and flows well. While I found many of the theoretical bits \"obvious\" (e.g., props 2-4), I also felt that their presence adds to the overall depth of the paper as more or less a complete treatment of this \"factored action space\" representation of the Q-function. The fact that the abstraction used in Theorem 1 can be implicit is what makes this Theorem, and the paper as a whole, interesting, and sufficiently differentiated from, say, the literature on factored MDPs (but I'm not an expert in FMDPs). The empirical portions are well done---experiments are highly relevant, and the connection between the experiments and the method is plain---and the discussion is insightful.\n\nI have seen, but not read in any detail, many works in multi-agent RL (and otherwise), which use factored action spaces (there are 69 results for \"factored action spaces\" on Scholar). I am trusting the author's lit review on the novelty point, and I hope that at least one of the other reviewers is sufficiently familiar with the literature to give a good opinion on novelty. To me, it seems sufficiently novel, but I lack broad knowledge of other work on factored action spaces.\n\nI'm going to give this paper an \"Accept\" rating, since I don't see easy ways to improve it, and it feels complete. That being said, I do think the description for \"Weak Accept\" (which is pretty strong?) is more accurate, insofar as this paper is rather specific.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "EUdNNEpPy-K", "reviewer": "Reviewer_CUKv", "summary": "This paper proposes the factored action spaces for offline RL in the applications of healthcare AI. The authors aim to leverage the factored action space in the form of linear Q-function decomposition to improve the sample efficiency compared with the baselines with combinational action spaces. The theoretical guarantees show that the linear Q-function decomposition can lead to zero bias under mild conditions. And even if the assumptions are violated, the method can still achieve good performances (nearly policy optimality). Empirical analysis verifies the proposed claims and effectiveness of the approach on simulated and real-world healthcare RL benchmarks. ", "questions": "#### **1. About the autoregressive factorization**\n\nCan authors give some analysis on both theoretical and empirical views on the scenarios where the actions are entirely in the autoregressive factorization? I notice that the theory part can be extended to this, and the empirical evaluation includes cases where a few actions are not independent. However, the entire autoregressive factorization, which is very common in RL with sub-actions, is not analyzed in the current version. In autoregressive factorization, the state would be augmented with the previously executed actions. Any discussion or analysis would be highly appreciated. \n\n#### **2. About the state space**\n\nIn the current version, each sub-action is based on the full state space. I think that why not leverage the local factorization or factored state space, where the state-action space is entirely factorized, and the state space of each sub-action would be much smaller. I think the sample efficiency would also be improved with the factored state-(sub)action space. **Please note that this is only a discussion. There is no need to run any experiment on this during the rebuttal phase.**\n\n ### **References**\n\n[1] Pitis, Silviu, Elliot Creager, and Animesh Garg. \"Counterfactual data augmentation using locally factored dynamics.\" Advances in Neural Information Processing Systems 33 (2020): 3976-3990.\n\n[2] Huang, Biwei, et al. \"Action-sufficient state representation learning for control with structural constraints.\" ICML 2022.\n", "limitations": "The limitations are more related to the algorithmic contribution and empirical evaluation (listed in the weakness section). I will increase my score if the authors give justifications during the rebuttal phase. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## *Strengths*\n\n#### **1. Technical soundness and significance** \n\nThe paper proposes the factored action space, where the action space is expressed as a Cartesian product of a few sub-actions spaces. Though similar ideas and methods have been proposed in the past few years, this paper gives a detailed and comprehensive analysis in both theoretical and empirical manners. The theory parts are complete and technically sound, and evaluation with practical consideration also verifies the theoretical claims. With the analyzed theorems in this paper, factored approaches (either in state-action space or action space) in RL would be more broadly applicable. \n\n#### **2. Presentation** \n\nThough the presentation can be further improved for better readability (see the weaknesses points below), the overall logical flow is clear, and the concrete examples (Fig. 2-Fig. 3) are helpful for readers to understand the theorems. \n\n## *Weakness*\n\n#### **1. The contributions on the algorithmic aspects are vague** \n\nAs the authors mentioned in the related work section, several works have been proposed to explore the benefits of factored action space for both model-based and model-free RL. Yes, this work offers the solution for leveraging factored actions with value-based methods in \noffline RL. However, from the current version, I cannot tell the significant contributions of the algorithmic aspects of learning the factored space. The Cartesian product for action space has been used extensively in previous works (i.e., [1]). It would be better to explicitly list the contributions to the algorithmic aspects in the revised version. A table with comparisons of all related approaches would be helpful. \n\n#### **2. About evaluation** \n\nMore complicated benchmarks (e.g., Mujoco control tasks, DOTA2, StarCraft, etc.) have been tested in other works using factored action spaces. Can this framework also be applicable for other commonly-used RL benchmarks? **Please note that running on these benchmarks during the rebuttal phase is not a must, but any discussion or analysis would be highly appreciated.**\n\n#### **3. Possible directions to improve the writing for better readability** \n\n#### -> 3.1 Giving algorithmic frameworks into the main paper\n\nThe authors can consider adding the algorithmic frameworks as algorithm pseudo-code or figures. The framework can explain the pipelines of learning or exploiting the factored spaces. Moving some justification contents in Section 3.3.2 into the appendix can save room for this (since I feel like the propositions in 3.3.2 is already very clear). \n\n#### -> 3.2 Adding one background section on factored MDP and factored action space\n\nThe related work section in the appendix briefly introduces factored MDP and action space in RL. I think it is better to briefly give some formal definitions in the background or preliminary sections in the main paper. \n\n\n### **References**\n\n[1] PIERROT, Thomas, et al. \"Factored Action Spaces in Deep Reinforcement Learning.\" (2020).\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept