Doubly Robust Fusion of Many Treatments for Policy Learning
Authors: Ke Zhu
Jianing Chu
Ilya Lipkovich
Wenyu Ye
Shu Yang
Venue: ICML 2025 poster (2025)
Keywords: calibration weighting
individualized treatment rule
fused Lasso
Primary area: general_machine_learning->causality
TL;DR: We propose a robust treatment fusion method for high-dimensional individualized treatment recommendations, addressing data sparsity and covariate imbalance to improve patient outcomes with theoretical guarantees and practical applications.

Abstract: Individualized treatment rules/recommendations (ITRs) aim to improve patient outcomes by tailoring treatments to the characteristics of each individual. However, in high-dimensional treatment settings, existing methods face significant challenges due to data sparsity within treatment groups and highly unbalanced covariate distributions across groups. To address these challenges, we propose a novel calibration-weighted treatment fusion procedure that robustly balances covariates across treatment groups and fuses similar treatments using a penalized working model. The fusion procedure ensures the recovery of latent treatment group structures when either the calibration model or the outcome model is correctly specified. In the fused treatment space, practitioners can seamlessly apply state-of-the-art ITR learning methods with the flexibility to utilize a subset of covariates, thereby achieving robustness while addressing practical concerns such as fairness. We establish theoretical guarantees, including consistency, the oracle property of treatment fusion, and regret bounds when integrated with multi-armed ITR learning methods such as policy trees. Simulation studies show superior group recovery and policy value compared to existing approaches. We illustrate the practical utility of our method using EHR-derived data from patients with Chronic Lymphocytic Leukemia and Small Lymphocytic Lymphoma.

Doubly Robust Fusion of Many Treatments for Policy Learning
Ke Zhu 12*
Jianing Chu
Ilya Lipkovich Wenyu Ye
Shu Yang'
Abstract
1. Introduction
In precision medicine, individualized treatment rules (ITRs) aim to optimize patient outcomes by tailoring treatment recommendations individuals based their character- istics. This personalization is essential because treatment effects can vary across individuals. Developing a rule that recommends the most effective treatment for each person requires accounting for these variations in a systematic and data-driven manner. For settings with two or multiple treat- ments, numerous machine learning approaches have been developed for estimating ITRs, commonly referred ITR learning or policy learning. These methods can be broadly classified into two categories. The first involves modeling treatment outcomes, as in Q-learning (Watkins & Dayan, 1992; Qian & Murphy, 2011 I; Song et al., 2015), A-learning (Murphy, 2003; Shi et al., 2018), or D-learning (Qi & Liu, 2018), where the focus on estimating the expected out- come under each treatment to derive the optimal rule. The second class of methods directly optimizes the value func- tion, which measures the expected outcome under given decision rule (Zhang et al., 2012; Zhao et al., 2012; Athey & Wager, 2021). In these methods, inverse propensity score weighting (IPW) or augmented IPW (AIPW) estimators are employed to evaluate the value function, and the optimal rule is identified by maximizing the value function over a class of decision functions, such linear ITRs (Zhang et al., 2012; Zhao et al., 2012) or tree-based ITRs (Zhang et al., 2015; Laber & Zhao, 2015; Athey & Wager, 2021).
Individualized treatment rules/recommendationss (ITRs) aim to improve patient outcomes by tai- loring treatments to the characteristics of each individual. However, when there are many treat- ment groups, existing methods face significant challenges due to data sparsity within treatment groups and highly unbalanced covariate distribu- tions across groups. To address these challenges, we propose novel calibration walibration-weighted treat- ment fusion procedure that robustly balances co- variates across treatment groups and fuses similar treatments using a penalized working model. The fusion procedure ensures the recovery of latent treatment group structures when either the calibra- tion model or the outcome model is correctly spec- ified. In the fused treatment space, practitioners can seamlessly apply state-of-the-art ITR learning methods with the flexibility to utilize subset of covariates, thereby achieving robustness while ad- dressing practical concerns such as fairness. We establish theoretical guarantees, including consis- tency, the oracle property of treatment fusion, and regret bounds when integrated with multi-armed ITR learning methods such as policy trees. Simu- lation studies show superior group recovery and policy value compared to existing approaches. We illustrate the practical utility of our method using nationwide electronic health ecord-derived de- identified database containing data from patients with Chronic Lymphocytic Leukemia and Small Lymphocytic Lymphoma.
However, these approaches encounter significant challenges when the number of treatment levels becomes large (Rashid et al., 2021). When treatment levels are numerous, data often sparse within each treatment group, making it diffi- cult to estimate treatment effects accurately. Additionally, covariate shifts across treatment groups exacerbate insta- bility, particularly when balancing methods like IPW are used. This instability is further amplified in underrepre- sented treatment groups, where propensity scores are small and highly variable.
"qqual contribution 'Department of Statistics, North Carolina State University, Raleigh, NC 27695, U.S.A. 'Department of Bio- statistics and Bioinformatics. Duke University, Durham, 27710 U.S.A. "Amazon (This work was done prior joining Amazon) "Eli Lilly& & Company, Indianapolis, IN 46285, U.S.A.. Correspon- dence to: Shu Yang <syang24 ncsu..duuu
A key insight address these challenges lies in recogniz- ing that many treatments share commonalities (Ma et al., 2022; 2023). For example, different pharmaceutical com- panies may develop treatments targeting the same disease mechanisms or symptoms, resulting in similar effects across
Proceedings of the n2nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

Doubly Robust Fusion of Many Treatments for Policy Learning
treatments. By grouping such treatments into clusters, we can reduce the effective dimensionality of the treatment space. Treatments within the same group can be treated as equivalent, enabling the application of efficient multi-armed ITR learning methods to the grouped treatments.
tural assumptions. Gao al. (2024) used a low-rank tensor with block structure to fuse treatments. In contrast, we assume group structure among treatment levels and ap- ply ciaibration-weigheee fused lasso. These approaches are complementary, targeting different structural assumptions.
Despite its potential, the task of treatment fusion introduces its own challenges due to sparse data and unbalanced covari- ate distributions. Sparse data within treatment groups neces- sitates the use of simple linear working models to prevent overfitting. However, these models risk misspecification, which can lead to biased fusion results. Additionally, severe covariate shifts across groups make traditional balancing methods like IPW unreliable, especially for treatments with very small sample sizes.
Continuous treatments. Other methods focus on continu- ous treatments (e.g.. Chemoozukov et al., 2019; Cai et al., 2021). Cai et al. (2021) also explores action grouping in a continuous setting. These approaches differ methodologi- cally from ours, which focuses on discrete actions.
Complex treatments. Recent work studies complex treat- ment types such as images, text, or chemical structures (Kad- dour et al., 2021; Nilforoshan et al., 2023; Schweisthal et al., 2023; Marmarelis al., 2024). Schweisthal et al. (2023) address large treatment spaces induced by high-dimensional continuous variables using neural networks and constrained optimization, focusing on regions with sufficient overlap. In contrast, our method targets the entire population, is easy to implement, compatible with existing algorithms, and inter- pretable.
To overcome these difficulties, we propose a novel pro- cedure called calibration helibrationweiighted weighted treatment fusion. This method uses calibration weighting (Lee et al., 2023; Wu & Yang, 2023) to robustly balance covariates across treatment groups, addressing the instability of traditional balancing techniques. The calibrated weights are then used in a pe- nalized working model with fused Lasso (Tibshirani et al., 2005) to group treatments based on their effects. Our pro- cedure is doubly robust, meaning that the true latent group structure can be recovered as long as either the calibra- tion weighting model or the outcome model is correctly specified. This robustness significantly enhances both the efficiency and reliability of treatment fusion compared to existing methods.
2. Preliminaries
We consider a K-armed setting where the treatment
A (1,...,,,..
Let X E X' S RP denote a vector of covariates, and Y E I denote the observed outcome of interest. We assume that larger values of Y are preferred by convention. The observed data (Y;, A., X,) are assumed to be independent and identically distributed. The potential outcomes Y (a). A, represent the outcomes that would be observed if subject received treatment a. The following standard assumptions in causal inference are made (Rubin, 1978).
After performing data-driven fusion, practitioners can trans- parently review the grouping results and scamlessly apply state-of-the-art multi-armed ITR learning methods on the grouped treatments, such as policy trees (Zhou et al., 2023), to align with their application contexts. This approach not only robustly reduces the dimensionality of the treatment space but also ensures flexibility, interpretability, and im- proved policy learning outcomes. By addressing the chal- lenges of many treatments, our method provides robust and practical framework for advancing precision medicine.
Assumption 2.1 (Identification). (i) Consistency: Y =
(()). (ii) Unconfoundedness: Y(a) X, Va E A. (iii) Positivity: P(A a X = x) 1 for all a E A.
An individualized treatment rule (ITR) is a decision func- tion d(.) X + A, which maps the covariate space to the treatment space. For any arbitrary ITR d(-), the correspond- ing potential outcome is defined as Y (dd()), which would be observed if randomly chosen individual were assigned treatment according to d(.), i.e., A d(X). The value function under d(-) is then defined as the expectation of Y (d(())) i.e., V(d) == (d((())))) Let the propensity score be
1.1. Related work
In this paper, we study settings with large action space induced by single discrete treatment variable with many levels (Saito & Joachims. 2022; Saito et al.. 2023; Peng et al., 2023; Sachdeva et al., 2024; Aouali et al., 2024). Related work considers policy evaluation or learning under alternative treatment structures.
(x) = P(A - a X
Combination treatments. Some studies consider treat- ments formed by combinations of multiple variables (Liang et al., 2018 8; Agarwal et al., 2023; Xu et al., 2024a;b). Agar- wal et al. (2023) propose synthetic combinations that impute counterfactuals via low-rank matrix completion under struc-
and the outcome mean function be
[a(x) = ({Y(a) X=
Under Assumption 2.1, the value function (dd) can be iden- tified using observed data through inverse propensity score
2

---

Doubly Robust Fusion of Many Treatments for Policy Learning
(X)]- Suppose D is class of ITRs of interest, such as linear ITRs (Zhang et al., 2012; Zhao et al., 2012; Cheng & Yang, 2024) or bree- ITRs (Zhang et al., 2015; Laber & Zhao, 2015; Athey & Wager, 2021). The optimal ITR is defined as d² (X) := ar V (dd). The complexity of 'D in- creases exponentially with K. Consequently, existing lit- erature often assumes that K is fixed (Zhou et al., 2023). However, in practice, the treatment dimension K may be high (Rashid et al., 2021), making the task of learning d' (X) significantly more challenging.
weighting (IPW): V(d) E
ances covariates across treatments. Since (a( = M (1) implies that the linear projections of Y(a) and ((a') are equal, we employ penalized linear working model to per- form treatment fusion, achieving double robustness when combined with calibration weighting. After treatment group- ing, any state-of-the- ITR learning method, such as policy trees, can be adopted, allowing for flexible outcome model- ing and superior policy learning.
3. Methodology
3.1. Calibration- Treatment Fusion
To address this challenge, a key insight is that certain treat- ments, such as those in drug development targeting similar disease symptoms and mechanisms, may yield compara- ble or identical outcomes (Ma et al., 2022; 2023). This observation suggests group structure.
Since the number of groups M and the group structure - are both unknown, we employ fused Lasso to jointly determine M and the partition. Specifically, we consider the following working linear model:
Definition 2.2 (Oracle group structure). A U g
Y =Mo(X)+++ + Σ 1(A a)x +E,
where 97, 's are disjoint sets satisfying: (i) (X)) ne (X) for a, a' E G and (ii) (X) # (X) for E 9% a' E Gi with b DV.
aEA
Σ I(A I(A= aa 00 aEA
s.t.
Remark 2.3. We define the oracle group structure by exact ) to ensure identifiability and interpretabil- ity. While seemingly restrictive, this serves as a natural basis for grouping similar treatments and enables formal guaran- tees. In practice, the fused lasso penalization we use allows for grouping treatments with approximately equal effects by tolerating small differences due to sampling variability. The explicit gap tolerance required for recovery is provided in Remark 3.11. Overall, when no exact group structure exists, fusion trades variance for bias, and this trade-off can improve performance when data are limited.
equality of (X))
where the redundant function Mx(X) represents the main ef- feet of treatments, and X Sa captures the interaction effect between treatment a and the covariates. The sum-to-zero constraint on the interaction terms ensures the identifiability of the regression function. The main effect function Mx(X) can be estimated using weighted parametric or nonparamet- ric regression models. As this is not the focus of our paper, we assume it has been accurately estimated and define the transformed outcome Y Y My
To estimate and group Su 's, we consider the following opti- mization problem, which imposes pairwise fusion penalty on each pair of treatment-speciffe parameters:
Motivated by this latent structure, we can first learn a group mapping 6 A + B :=
M), and subsequently learn the grouped ITR (.): - B by using established multi-armed policy learning methods. Since M is smaller than K after grouping, learning the grouped optimal ITR d := argmax, V (d) becomes more efficient, where DE is a class of ITRs d= After obtaining de, for any X such that dx(((] b. we define d( d(X) as randomly selecting one such that g
1 Σ n
Σ
c (V..x.
min
aE i:A,ma
Σ IECCO'KK
Phn (1<aa Sa. Sarll)
(I)
where C(:s-) is a prespecified loss function that measures the goodness of fit, || 1/2 denotes the & norm of vector, Px, is a penalty function that encourages the fusion of ?," into groups, and A, is the tuning parameter, which can be selected by multiple model selection criteria, such as Bayesian information criterion (BIC) (Schwarz, 1978) or extended BIC (EBIC) (Chen & Chen, 2008).
Therefore, the primary objective of this paper is to recover the true group structure ,M, GE through data-driven fusion. This task is particularly challenging due to the sparsity of data within treatment groups caused by a large K, which hinders the accurate estimation of Ha (X) using complex models beyond linear ones. Additionally, some treatment groups have very low propensity scores, leading high vari- ability in conventional inverse propensity score balancing methods. Both model misspecification and covariate shifts can introduce bias, resulting in poor treatment fusion and suboptimal policy learning. To address these challenges, we propose calibration weighting method that robustly bal-
However, the applicability of the objective function (1) is limited to scenarios where either (i) the true outcome func- tion Pa (x) is linear or (ii) the covariate distributions are identical across all treatment groups In cases where these conditions do not hold, the estimated Ç, 's obtained from
3

---

Doubly Robust Fusion of Many Treatments for Policy Learning
(I) may deviate from the true group structure 1 In observational studies, variations in covariate distributions across different treatment groups often prevent these groups from accurately representing the entire population. To miti- gate this issue, a preliminary step involves reweighting the samples within each treatment group so that the weighted samples better reflect the overall population. To achieve this, we propose the following calibration weighting approach.
Algorithm 1: Calibration-Weightedd Treatment Fusion
Input: Data {(X;, A,, Y))
for a = 1,..., K do
Solve calibration weights úi by optimizing:
min wiii:A==a
Σ
byy (wa),
t.A==
S.t.
Σ EAi=o
w;Y, =
Σ i.Ama
wi 11
For each treatment group a, aim to assign weights {w; A; = a) to calibrate the covariate distribution of the group to match the overall sample mean X. This is achieved by solving the following optimization problem for each a E A::
Solve ? by weighted fused Lasso:
min <=(C) (...)))) T ERKR
+ (L((5)+PP((())
min Mi m
Σ
by (we),
i.A.aa
S.t.
Σ A=A.=a
= = X,
Σ =A
=,=1,
Ln(S) 1
2n
Σ
E
Th (i, X1 Su)
aEd s:A.=a
Pa(() =
Σ
PAn (llSa - S1r
where h, (w) quantifies the discrepancy between the cali- bration weights and the uniform distribution na -1, with na denoting the sample size of treatment group a. The function (w) can chosen from the Cressie and Read family of discrepancies (Cressic & Read, 1984). defined as:
Egaca' &K
Forming groups 8(a) = ((a') if C. = 2a..
Output: Group mapping ó + B.
Σ by (w;) =:A
Σ (nly + (m,a 1}. #AA=a
structure is equivalent recovering the projection vectors. In Section 3.2.2, we establish the convergence of the oracle estimator for the projection vectors under doubly robust and regularity conditions (Theorem 3.8). In Section 3.2.3, we show that the oracle estimator is a local minimizer of the objective function in Algorithm (Theorem 3.1.). Taken together. these results imply that Algorithm consistently recovers the oracle group structure. Technical clarifications are provided in remarks and can be skipped by readers less interested in such details.
For example, minimizing Ees,=o (w: equivalent to maximizing EEA,= log(w;), leading to the maxi- mum empirical log-likelihood objective function. Min- imizing EeA, (o is equivalent to maximizing - A. Wy log(w;), leading to the maximum empirical exponential likelihood or entropy.
Let Wi denote the calibrated weights solved by (2) for in- dividual i. Using these weights, the calibrated objective
function becomes:
1.2.1. REPRESENTATION OF ORACLE GROUP STRUCTURE
min Σ n
Σ aEA i:A=m
wiL
For a E A, let (a) Y(a) - MX denote the trans- formed potential outcome. We project Ya onto the linear space spanned by X and denote the projection vector by
Σ K
Psn (115a - -arlll)
(3)
C argmin E {f (TT)
x
(4)
CERP
The entire procedure is summarized in Algorithm I, consid- ering the least squares loss function as an example.
where X includes the intercept term Solving (4) yields:
E
x (xa) =0
(5)
3.2. Double Robustness of Treatment Fusion
Theory roadmap. This section provides theoretical guar- antees for Algorithm 1. In Section 3.2.1, we represent the oracle group structure U 1 via the linear projection of potential outcomes onto the covariate space, without requir- ing the linear model to be correctly specified. Under the completeness Assumption 3. recovering the oracle group
We define the projection residual by
:(a) = {((a) x C..
(6)
By (5) and (6), we have
y(a) = X S + E(a), E EX e(a)) =0
(7)

---

Doubly Robust Fusion of Many Treatments for Policy Learning
From the above derivation, condition (7) does not assume linear relationship between ((a) and X; it holds solely due to the projection (4) and the definition of e(a). Con- dition (7) alone is generally insufficient to guarantee the consistency of the estimated projection vector derived from the unweighted working linear model (1). Ma et al. (2022) assumes a linear relationship between (a) and X, which
3.2.2. CONSISTENCY ORACLE ESTIMATOR
We establish the convergence of Ç to S" under the assump- tions stated below. Let Ci C2. C3, and C4 denote positive constants. We allow K, M, and p to grow with n, omitting their dependence on n for notational simplicity. We write On >> bn to denote b2- = =(an).
essentially imposes a stronger condition:
Assumption 3.2 (Convergence of calibration weight). Vi= 1,..., n, 0. w + Op(1/ and C1 S 3 S C2-
V(a) = X 5h + e(a), B(e(a) x) =
(8)
Assumption 3.3 (Doubly robust model assumption). One of the following conditions holds: (i) (Correct calibration weighting) W. = 1/TAA, (X.), or (ii) (Correct outcome
In the following, we show that by using calibration ealbration-weighted treatment fusion (3), the consistency results hold if either the calibration weighting (2) is correctly specified or the outcome model (8) is correctly specified, but not necessarily both. Consequently, our approach provides a more robust fusion method against model misspecification.
model) ({s(a) X = ({Y - X C² [X] = 0 for any a E A.
Remark 3.4. Assumption 3.2 requires the /r--convereence of the working weights W; to bounded limits w/,, which are not necessarily the true inverse propensity scores. This typically holds under mild conditions for posited parametric models for weighting, such as the entropy balancing method (see Section A. for details). The convergence rate require- ment for the weights in the fusion stage may be relaxed by using undersmoothed estimators or advanced doubly robust methods (Chambaz et al., 2012; Ertefaic et al., 2023; Bruns- Smith et al., 2025). Assumption 3.3 requires only that either the calibration weighting model or the outcome model is correctly specified, highlighting the double robustness of our results.
Since X includes the intercept term, and from (7), we have E{E(a)} = 0. Therefore,
a. E 9h + Pa(X) = Mar (X) = S = -a (9)
To ensure the reverse direction of (9) holds, we impose the following assumption.
Assumption 3.1 (Completeness). For any function h(-). if
(X)(X)= = 0, then h(X)= = 0 almost surely.
Under Assumption 3.1, we have
Assumption 3.5 (Regularity condition for X). For any 1..... P. n Σx= X3 I C3. For any b E B, Amin (ENA.ES; X.X,)) /Nmin 2 C2. where (-in(-) de- note the smallest eigenvalue of matrix, and Nmin == minser Σi= 94} is the smallest sample size across groups
C - Co
Ma(X) (xx))=a(X]] (X),
which ensures that identifying S, recovers the oracle group- ing.
We denote the group-shared projection vector as B = S% Va E Gi. The transformed potential outcome can then be expressed as:
Assumption 3.6 Sub-Gaussian error). For any a E A, e(a) = (s),...... En(()) has sub-Gaussian tails, that is, 30. > 0, for any b E R and t > 0, PP E(a)
8(a) - XT + =((a)= = x 3 + E(a).
(10)
E(b E(a) > [[b]l2t) I 2 pexp(-f/200
Remark 3.7. Assumptions 3.5 and 3.6 are typical regularity conditions used in high-dimensional statistics (Wainwright, 2019). Notably, Assumption 3.6 only requires that b Te e(a) concentrates around its expectation. If the outcome model is misspecified and there is covariate shift across treatments, E(b {aa} may not equal zero when b, = T(A, = a) a)j,;. leading to bias. However, by using calibration weighting with b, = I(A, = a a)w; Xtj. we can robustly eliminate this bias if either the calibration weighting model or the outcome model is correctly specified, as shown in Lemma A.I.
If the true group structure 9 is known, the data within each group can be pooled to estimate the group-shared pro- jection vector as
B,
min p. Elr 2n
Σ i=1 aE
2 (f:-x, Bm)
Then, the oracle estimator for the projection vector Ç* (Si T ....., Cx I can be obtained by expanding Br such nor that c = (GIT CorT ..... )T, where 2.2 = G for all 2or
nor
Theorem 3.8 (Consistency of C Suppose Assumptions 2.1, 3.2, 3.3, 3.5, and 3.6 hold. If Mp/n o(1) and Vpn log(n)/ //min = o(1). then for some constant C > O. with probability at least 1 - 2M in (where in + 0 as -or
G3. In practice, since the true group structure is un- known, the estimated projection vector Č is obtained using Algorithm 1. Define the objective function as QC) = T.n (<) + Pc()).
n - n-))),, we have 1|5 - -5*ll S CCvp log(no/
5

---

Doubly Robust Fusion of Many Treatments for Policy Learning
Remark 3.9. To ensure the consistency of ç it is required that v pn log(n) < Nmin S nr/M, which implies that the number of groups must satisfy M = o Vvo/tplog(n)3).
ploying policy trees as the specific policy class, and review its theoretical results. CAIPWL. involves three main steps. First, it estimates the nuisance functions
Tb (z) = Σ (a()); REC
[b(x)) = pa pa(x)) Va E gh+
3.2.3. ORACLE PROPERTY OF 2
Next, we establish the oracle property of Č. To encourage the grouping of similar projection vectors and reduce bias introduced by the penalty, we require the penalty function (+).()) (+) := (x)( to have a sharp derivative near O. More- over, the &o distances between the projection vectors of two different groups must be sufficiently large to ensure they can be separated. Thus, we impose the following regularity condition, which is commonly used in high-dimensional statistics (c.g., Ma & Huang, 2017).
using L-fold cross- fittingitinn. Next, to evaluate the value of a policy d8, an augmented IPW (AIPW) estimator is used:
1 n n i=1
Y, -BB,
-=(()) (X:)
V(d2)
IBB - de (x.))
$ n,
(X.)
(1))
Finally, (id is optimized over specified policy class D8, such as decision tree, obtain the ITR estimator d The detailed procedure is outlined in Algorithm 2. Note that PB. p0(()) (X,) and RB. ==1(1) 11)(())) (X.) denote the nuisance functions estimated using T. folds of data, excluding the 1(i)-th fold containing the i-th unit.
Assumption 3.10 (Penalty function). The penalty function (+.,()) = (-) is symmetric about 0, satisfies Px., (0) = 0, is differentiable near 0 with pt) continuous except at finitely many? t and p' p0(0-) = 1, and becomes constant for t 2 2 for some e > 0. Additionally, mingyy [|3; 18/1100/ /e> >> An >> On + +pp0// p0n/Kmin + n /nlog(n)/ Kmin- where On. == (Vpnlog(n)//N /iinn C' is the constant in Theorem 3.8, and Kiin == min [9&] is the smallest number of treatments across groups
Algorithm 2: Cross-Fitted AIPW Policy Learning
Input: Data {(x,, A,, Y=) Group mapping 8. Mapping the treatment into groups B, = =(A,).
Remark 3.11. When (i) the covariate dimension p and the number of groups M are fixed, and (ii) Nuun 7/nn/ M and Kmnn 1/2 K/M where 17; 122 E (0, 1] are fixed con- stants, the term on + /n/ Kmin + /n log(n)/ Kmin is of order O(v log(n)/ K). In this simplified case, it suf- fices to assume that min648 [|3, Bi llo/c > An >> vn log(n)/K.
Split the data into L folds.
for ( = 1,..., L do
for = 1,............ 1,..., M do
Fit t. (xx using rest L - 1 folds.
Fit Ar (z)) using rest L -1 folds.
Compute the estimated value of policy d by (11).
Solving d = argmax,#<ps Vdd5).
Theorem 3.12 (Oracle property of 2) Suppose the con-
Output: Optimal policy dd5
ditions in Theorem 3.8 and Assumption 3.10 are satisfied. If Kp/n = =(1), there exists local minimizer Č of the objective functionn Qn( such that )
Consider D = D in Algorithm 2, where depth-D trees serve as candidate policies dB(()) : XX - B. A depth- tree maps a covariate vector X = (X1 , XP) E X into an action E B by traversing D branch layers followed by a final layer of leaf nodes. Each branch node splits on covariate X at a threshold 1, directing traversal to the left child if X < l and to the right child otherwise. The traversal ends at leaf node, assigned one of M actions in B, partitioning X' into up to 2D disjoint regions, each associated with an action b. Figure 1 shows a depth-5 tree learned from a real data application.
Combining Theorems 3.8 and 3.12, we demonstrate that minimizing (C) facilitates the recovery of S*, which in- dicates the group structure UM under the completeness Assumption 3.1.
3.3. Multi-armed Policy Learning
After fusing treatments, those within the same group can be treated as identical. This enables our doubly robust treatment fusion procedure to seamlessly integrate with any state-of-the-a thett-tttrrtttt multi-armed ITR learning method to identify the optimal grouped ITR dB. X + B, where B = (1,2, M} represents the set of group indices.
Notably, the covariates used for node splits may be a smaller subset of those used for treatment fusion in Algorithm and for estimating nuisance functions (xx and (x(X). While treatment fusion and nuisance function estimation employ a richer covariate set to ensure correct model specification, the decision function d prioritizes a smaller, interpretable subset of covariates that are actionable and relevant for
As a concrete example, we combine the doubly robust treat- ment fusion with the Cross-Fitte Frots-FittedAaae Augmented IPW Learning (CAIPWI.) approach proposed by Zhou et al. (2023), em-
6

---

Doubly Robust Fusion of Many Treatments for Policy Learning
Table 1. Summary of definitions for the CAIPWI. and policy tree.
Description
Notation
Definition
Hamming Distance on { C X
Hdd di (adl (e, DE, (m)HL)
n "i= (df(,)) + di (zi)].
e-Hamming Covering Number on {rr)/=1 C.
The smallest number L of policies di. ..... .d. in DB, such
that Va E DB, 3d E DD, ld di (x))= S.
e-Hamming Covering Number
Ner(e,D9)
Nup(Nn(e, DB, (m:)))) m 2 1, (. C C) Jo Vlg (e,
Entropy Integral
x(DB)
Policy value of d8
I (X)) Y- + (as(x)(
xu(X)
Worst-case variance in evaluating the difference between two policies in DB
V.
MfEED (( (dg))}
Regret Bound
R )
(XX - BYY (d (x))).
Decision Tree Class
D
Set of all depth- D trees.
tree
decision-making This separation highlights the robustness and flexibility of our procedure, ensuring the resulting policy is both interpretable and practical for implementation.
Pe (X) all 5(a) = 5(a'). For the ds learned from Algo- rithm 2, we have R(d0)= = = (N(DB)/VV.))
Proposition 3.18 (Regret bound of policy tree). Un- der Assumptions 2.1. 3.13, and 3.15, suppose that na (x( = (X) for all 5(a) = 8(d'). For the d learned from Algorithm 2 with DE = D tree we have R(d3) A. 1D. -I/
We provide the Vn-reerett bound for CAIPWL, relying on the rate doubly robust model, policy class complexity, and bounded outcome and covariate assumptions. Table I is summary of definitions relevant the theoretical result.
log 2D log M +
Assumption 3.13 (Rate doubly robust model assumption). For any b B, ..... L, A, (XX and it, (X) satisfy:
E ( xx)) 2 + O, B B1, (X) (x))) 0,
4. Numerical Experiments
4.1. Synthetic Scenarios
E ( (()))=((x))* E (1p (x) [x]) =o o)n
o(n
We considered M = 4 treatment groups, with the structure summarized in Table 2. Each group comprises | = 4 treatments sharing identical outcome mean functions, as detailed in Table 4. We considered covariate distributions with shifts and varying sample sizes across treatment groups, as shown in Table 3. The covariance matrices are:
Assumption 3.14 (Complexity of the policy class). 40 <
< 1. N. ) < C₁ exp((2e for some constants C1,C, 0,0 < w < 0.5.
Assumption 3.15 (Bounded outcome and covariate). For all E A, (a)) is bounded, and X is bounded.
Remark 3.16. Assumption 3.13 is weaker than the standard doubly robust model assumption, which requires either the estimator of Th (X) or (X) to be Vii-consistent. Instead, Assumption 3.13 permits a trade-off between their accu- racies, requiring only that the product of their error terms scales as o Modern machine learning methods offer effective estimators for these quantities. Assumption 3. 14 requires the logarithm of the policy class covering number to grow at low-order polynomial rate with 1/e, condi- tion satisfied by the finite-depth trees considered here (Zhou et al., 2023). Assumption 3.1,, required only for the results in Section 3.3, is a standard regularity condition in the policy leaming literature.
1
-0.25
1
- -0.3
Σ1
-00.5
0.25
Σg = 0.3
-0.3
1
Table 2. Group structure.
Group Treatment (1223344
2
3
4
(5,6,788)
19,10,11.12) [13,14,
As a baseline, we implemented the CAIPWL method (Zhou et al., 2023) without calibration weighting fusion, using the default tuning in the R package po l i cytree to learn a depth-3 3 policy tree. We further implemented the fusion step both with and without calibration weighting, using CAIPWL to learn the corresponding optimal policy trees. In the fusion step, fused lasso uses extended BIC (Chen & Chen, 2008)
Proposition 3.17 (Regret bound of CAIPWL). Under As- sumptions 2.1, 3.13, 3.14, and 3,15, suppose that (X)) =
7

---

Doubly Robust Fusion of Many Treatments for Policy Learning
0.26, whereas fusion with calibration weighting achieved a significantly higher average ARI of 0.96. For the "fusion policy tree" approach, the poor quality of fusion led to lower average testing value compared to the "CW + fusion + policy tree" method. However, both approaches still outper- formed the baseline. In contrast, the method proposed by Ma et al. (2022) suffered from misspecified outcome mean functions, yielding results that were even worse than the baseline. We perform additional simulations with increased K and fixed n, and under misspecified weighting model, deferring the details to Section B.
Table Covariate distribution.
Treatment
Covariate
Sample Size
X1 (eraoul1
(1.59,13)
(X2, X,) X1 = N((1, -1)". Ex)
150
(Xg, X,) 1 X 0 11)-
X1 Bermoulti(04)
(2.6.10.14)
(X,, [|
125
Xx,X,]) X = ~ N((-1.1)), Eg)
X1 ~ (er5
(3.7.11115
(X2, X,)]]|x==11 -
100
X1 (e
(4.8.12.16)
=
75
(Xg, = 1((-,,))
4.2. Real Data Application
Table 4. Outcome mean functions for each group.
We illustrate the proposed methods through an application to data of patients with Chronic Lymphocytic Leukemia (CLL) and Small Lymphocytic Lymphoma (SLL) from the nationwide Flatiron Health electronic health record-derived database. The Flatiron Health database is a longitudinal database, comprising de-identified patient-level structured and unstructured data, curated via technology-enables ab- straction (Ma et al., 2020; Birnbaum et al., 2020). During the study period, the de-identified data originated from approxi- mately 280 US cancer clinics (~~800 sites of care; primarily community oncology settings). The data are de-identified and subject to obligations to prevent re-identification and protect patient confidentiality. CL.I. and SLL are slow- growing, indolent hematologic malignancies that primarily affect lymphocytes. In CLL, cancer cells are mainly found in the blood and bone marrow, while in SLL,, they are mostly located in the lymph nodes. The relative 5-year survival rate following an initial CLL diagnosis is estimated to be 88. 1% (source: SEER Cancer Statistics).
[x[X]
[0.7 + +0.1X; 0.3%, 0.223 0.4sign( X + +33, 1. -25]]
2
3exp(0.5 + 0.1X) +0115,- 0.3X3 + 0.55agmm X3 3Xx- 0 225)
3
Mexp + 0.1X, - -0.15x,-- 0.3.X-- + 0 + +3X,- - 2.5)
3exp(00.6 + 0.1X, + 022, -..1., -01x3- 0.7aign[X3 Xs- 23
for model selection. Treatments are grouped if the Euclidean distance between their fused lasso estimates is less than 0.25. Additionally, for comparison, we implemented the method proposed by Ma et al. (2022), where the group structure learned using fused Lasso without calibration weighting, and linear working models are used for both treatment fusion and policy learning.
We used the adjusted Rand index (ARI) (Gates & Ahn, 2017) to assess the quality of the fusion by comparing it to the true underlying group structure presented in Table 2. The ARI measures the similarity between two clusterings while ac- counting for random chance. It evaluates how pairs of items are grouped together separated in both clusterings, adjust- ing for the possibility that some agreement could occur by chance. ARI produces a score between -1 and I, where in- dicates perfect agreement, 0 suggests no better than random chance, and negative values indicate conse-thhn-randdon fusion. We then generated a large dataset following the same distribution as the sample dataset and used the average outcome mean function over the entire dataset as the test value. We conducted 200 replications of the learning-testing procedure, with the average results summarized in Table 5.
The dataset includes 10,346 patients who received first line of therapy (LOT), with details provided in Table 6. The primary outcome is patient overall survival status (1 for survival, O for death), along with 10 covariates: race, region, PayerBin, SES Index (2015-2019), gender, ECOG score, Rai stage, lymphadenopathy, age at the start of first LOT, and the time from diagnosis to the initiation of first LOT.
Table Sample size for each treatment.
Treatment
Number of patients
Table S. Simulation results for K = 16.
cBTКi mono
3392
Method
ARI
Number of groups
Value
policy tree (baseline)
16
8.77 (0:08)
AntiCD20 + Chemotherapy Only
1726
fusion + policy tree
0.26 (0.14)
10.725 (1.93)
8.78 (0.09)
AntiCD20 mono
1230
CW fusion policy tree
0.96 (0.06)
4.335 (0.60)
8.89 (0.11)
BCL2i + AntiCD20 Only
463
Ma et al. (2022)
0.26 (0.14)
10.725 (1.93)
8.51 (0.12)
cBKTi + AntiCD20 Only
408
CW= Calibration Weighting. ARI (Adjusted Rand Index for fusion quality) and
policy value: higher is better. Oraele number of groups 44 Nunibers pareutheses are Monte standard errors. Results are avenged over 200 runs.
Chemotherapy Only
215
Other
412
Total
10346
We observed that, due to heterogeneity in the covariate dis- tribution across different treatment groups. fusion without calibration weighting resulted in an average ART of only
We implemented the proposed "CW + fusion + policy tree"
8

---

Doubly Robust Fusion of Many Treatments for Policy Learning
Yes
timex=1.87
No
Yes
dmme-1.00
No
Yos
time--36.99
No
ly lymohadengpathy
age--84
Jalse-
ralstage__ 0
EC0G_0-1-0
Ves
Yes
Yes
Yes
tm.3--0.3
um.ex
agec-72
agex-822
timex 35, 58
tmme- 15.31
timer-
timex
Yes
Yes
No
Group 3
Group 5
Group
Group
Group
Group 5
Group
Group S
Group
Group
Group
Group 2
Group
Group
Group S
Group 2
2
Group
<BTK moro
ArRiDD00 mono
Chemothenyyy Only
ArrCDD0
Othe
dEKTI AniCO20 Only
BCL3I AMRCD22
Orly
Chemotwrapy Only
Figure The fusion results and the learned optimal policy tree assigns patients to grouped treatments based on covariate splits. time the time from diagnosis to the first LOT; age is the age at the start the first LOT; lymphadenopathy sarpppttt__llsse__aepppppaeeeeee sffflseelll indicates no lymph node swelling, and 1ymphadenopa taaaathyaallllll wal.aaleee0 indicates its presenee; raistage._ee 0=1 indicates Rai stage (low risk), and rais sisstaeee-0 indicates stages 1-IV (intermediate to high risk); ECOG_0-11 indicates an ECOG score of 0 or (good functional status), and ECOG_0-1=1 indicates score of 2-5 (diminished functional status).
method described in Algorithms and 2. All 10 covariates were included in the calibration weighting and the estima- tion of nuisance functions in CAIPWL. The algorithm offers flexibility by enabling a subset of covariates to be used for learning the ITR, in contrast to methods that require the same covariates for weighting, modeling nuisance functions, and ITR learning. Certain confounders, such as race, region, and proxies of social status like PayerBin and SES Index (2015-2019), were excluded from ITR learning avoid their use in treatment assignment. As result, the remaining six covariates were used for fused Lasso and ITR.
fused Lasso to recover the latent group structure of treat- ments, providing theoretical guarantees of consistency and double robustness. Practitioners can seamlessly integrate the fusion results with state-of-the-a ITR learning methods, such as policy trees, which offer Vn-regret bounds, enabling flexible and interpretable decision-making, Simulation stud- ies highlight the superiority of our method over baseline and competing approaches. Additionally, we demonstrate its practical utility through real data application, yielding clinically relevant insights.
In extreme cases where some treatment arms have few or no observations, our method may become unstable due to the lack of information, unless additional structural assump- tions, such as the combinatorial structure in Agarwal et al. (2023), are imposed. This instability underscores the in- herent difficulty of the problem and suggests directions for methodological improvement. Currently, the fusion step is performed once; an alternative is to adopt an iterative pro- cedure that alternates between treatment fusion and weight estimation to enhance stability.
The fusion results and the learned optimal policy tree are presented in Figure 1. The following insights can drawn: (i) Two monotherapies are grouped together in Group 1, reflecting their similar mechanisms of action and treatment intensity. (ii) Combination therapies are assigned to dis- tinct groups, while chemotherapy-only forms its own sep- arate group. (iii) Older patients or those with shorter time since diagnosis are more likely recommended to Group 5 sheemothrrapy-only)) likely due to limited treatment toler- ance and the need for immediate intervention. (iv) Relatively younger patients or those with longer time since diagnosis tend to be recommended to Groups 2, 3, 4 (combination therapies), likely due to their better functional status to toler- ate aggressive treatments, and longer diagnostic times often indicate a chronic disease course requiring more targeted interventions to manage progression. Our findings provide valuable insights for guiding future individualized treatment strategies while ensuring their practical feasibility.
Our method can be extended in several directions. First, while this study focuses on single data source, future work could enhance the procedure by integrating multi-source data or generalizing learned ITRs to other target popula- tions (Mo et al., 2021; Chu et al., 2023; Wu & Yang, 2023; Zhang et al., 2024; Carranza & Athey, 2024). Second, be- yond ITR estimation, providing inference and uncertainty quantification is critical (Liang et al., 2022; Ghosh et al., 2023; Cheng & Yang, 2024), especially in high-stakes con- texts like medicine, with conformal prediction offering promising approach (Osama et al., 2020; Taufiq et al., 2022; Zhang et al., 2023). Finally, nigh-dimensional treatments also occur in heterogeneous treatment effect (HTE) estima- tion (Goplerud et al., 2025), and extending our procedure to HTE frameworks would enhance its applicability.
5. Conclusion
This paper introduces a calibration -eiirattoo-weighted treatment fu- sion procedure to address the challenges of many treatments in ITR learning. By leveraging treatment similarities and ro- bustly balancing covariates, our method employs weighted
9

---

Doubly Robust Fusion of Many Treatments for Policy Learning
Acknowledgment
continuous exposure. Electronic Journal of Statistics, 6: 1059-1099, 2012.
We thank the anonymous reviewers and meta-reviewers of ICML 2025 for their valuable feedback, which led to greatly improved manuscript. Yang was partially supported by the National Science Foundation grant SES 2242776 and the National Institutes of Health grants IR01ES03165 and IR01HL169347.
Chen, J. and Chen, Z. Extended bayesian information criteria for model selection with large model spaces. Biometrika, 95(3):759-777 7711, 2008.
Cheng, Y. and Yang, S. Inference for optimal linear treatment regimes in personalized decision making. In Kiyavash, N. and Mooij, J. M. (eds.), Proceedings of the Fortieth Conference on Uncertainty in Artificial Intelligence, volume 244 of Proceedings of Machine Learning Research, pp. 718-735. PMLR, 15-19 Jul 2024. URLhttps: wwwphhttsss/pprcceeeeee http///proeeeeggge/eeeee/e mlr press/ v2 44 / 2444/chegg2433 4bb html.
Impact Statement
This paper presents work whose goal advance the field of policy learning. There are many potential societal conse- quences of our work. none which we feel must be specifi- cally highlighted here.
Chernozhukov, V., Demirer, M., Lewis, G., and Syrgkanis, V. Semi Semi-parametrit efficient policy learning with continu- ous actions. Advances in Neural Information Processing Systems, 32, 2019.
References
Agarwal, A., Agarwal, A., and Vijaykumar, S. Synthetic combinations: A causal inference framework for combi- natorial interventions. Advances in Neural Information 2023.
Chu, J., Lu, W., and Yang, S. Targeted optimal treatment regime learning using summary statisties. Biometrika, 110(4))913-331 2023.
Processing Systems,
Aouali, I., Brunel, V.-E., Rohde, D., and Korba, A. Bayesian off off-policy evaluation and learning for large action spaces. arXiv preprint arXiv:2402. / /4664, 2024.
Cressie, N. and Read, T. R. Multinomial goodness-of-fi tests. Journal of the Royal Statistical Society: Series (Methodological), 46(3):440 464, 1984.
Athey, S. and Wager, S. Policy learning with observational data. Econometrica, 89(1):133-161, 2021.
Ertefaie, A., Hejazi, N. S., and van der Laan, M. Nonpara- estimators based on the highly adaptive lasso. Biometrics, 79(2):1029-104) 2023.
metric inverse- peerrseproobbbittyyeeeeee
Birnbaum, B., Nussbaum, N., Seidl Rathk-Ra, K., Agrawal, M., Estevez, M., Estola, E., Haimson, J., He, L., Lar- son, P., and Richardson, P. Model-assistee cohort selec- tion with bias analysis for generating large-scale cohorts from the EHR for oncology research. arXiv preprint arXiv:2001. 00765,, 2020.
Gao, C., Zhang, Z.,, and Yang, S. Causal customer churn analysis with low-rank tensor block hazard model. In In- ternational Conference on Machine Learning, pp. 14920- 14953. PMLR, 2024.
Boos, D. D. and Stefanski, L. A. Essential statistical infer- ence: theory and methods, volume 120. Springer Science & Business Media, 2013.
Gates, A. J. and Ahn, Y.-Y. The impact of random models on clustering similarity. Journal of Machine Learning Research, 18(87):1-22 2017.
Bruns-Smith, D., Dukes, O., Feller, A., and Ogburn, E. L. Augmented balancing weights as linear regression. Jour- nal of the Royal Statistical Society Series B: Statistical Methodology, pp- qkaf019, 2025.
Ghosh, T., Ma, Y., Song, R., and Zhong, P. Flexible in- ference of optimal individualized treatment strategy in covariate adjusted randomization with multiple covari- ates. Electronic Journal of Statistics, 17(1):1344-1370 2023.
Cai, H., Shi, C., Song, R., and Lu, W. Deep jump learning for off-policy evaluation in continuous treatment settings. Advances in Neural Information Processing Systems, 34: 15285-15300, 2021.
Goplerud, M., Imai, K., and Pashley, N. E. Estimating het- erogeneous causal effects of nigh-dimensional treatments: Application to conjoint analysis. The Annals of Applied Statistics, press, 2025.
Carranza, A. G. and Athey, S. Robust offline policy learning with observational data from multiple sources. arXiv preprint arXiv:2410.08537 2024.
Kaddour, J., Zhu, Y., Liu, Q., Kusner, M. J., and Silva, R. Causal effect inference for structured treatments. Ad- vances in Neural Information Processing Systems, 34: 24841-24854, 2021.
Chambaz, A., Neuvial, P., and van der Laan, M. Estimation of paramarie variable importance measure of
10

---

Doubly Robust Fusion of Many Treatments for Policy Learning
Laber, E. B. and Zhao, Y-Q. Tree-based methods for individ- ualized treatment regimes. Biometrika, 102(3):501-514, 2015.
Osama, M., Zachariah, D., and Stoica, Learning robust de- cision policies from observational data. Advances in Neu- ral Information Processing Systems, 33:18205-18214, 2020.
Lee, D., Yang, S., Dong, L., Wang, X., Zeng, D., and Cai, J. Improving trial generalizability using observational studies. Biometrics, 79(2):1213-1225, 2023.
Peng. J., Zou, H., Liu, Li, S., Jiang, Y., Pei, J., and Cui, P. Offline policy evaluation in large action spaces via outcome-oriented action grouping. Proceedings of the ACM Web Conference 2023, pp. 1220-1230, 2023.
Liang, M., Ye, T., and Fu, H. Estimating individualized optimal combination therapies through outcome weighted deep learning algorithms. Statistics in medicine, 37(27): 3869-3886, 2018.
Qi, Z. and Liu, Y. D-learning to estimate optimal individual treatment rules. Electronic Journal of Statistics, 12(2), 2018.
Liang, M., Choi, Y-G.., Ning, Y. Smith, M. A., and Zhao, Y.-Q. Estimation and inference on nigh-dimensional in- dividualized treatment rule in observational data using split-and-poleed de-correlated score. Journal of Machine Learning Research, 23(262): 1 -6, 2022.
Qian, M. and Murphy, S. A. Performance guarantees for individualized treatment rules. The Annals of Statistics, 2011.
Rashid, N. U., Luckett, D. J., Chen, J., Lawson, M. T., Wang, L., Zhang, Y., Laber, B., Liu, Y., Yeh, J., Zeng, D., al. High-dimensional precision medicine from patient- derived xenografts. Journal of the American Statistical Association, 116(535): (1::5):1140-11 1154, 2021.
Ma, H., Zeng, D., and Liu, Y. Learning individualized treat- ment rules with many treatments: A supervised clustering approach using adaptive fusion. Advances in Neural In- formation Processing Systems, 35:15956-15969, 2022.
Ma, H., Zeng, D., and Liu, Y. Learning optimal group- structured individualized treatment rules with many treat- ments. Journal of Machine Learning Research, 24(102), 1-48, 2023.
Rubin, D. B. Bayesian inference for causal effects: The role randomization. The Annals of Statistics, 6(1):34-58, 1978.
Sachdeva, N., Wang, L., Liang, D., Kallus, N., and McAuley, J. Off-policy evaluation for large action spaces via policy convolution. In Proceedings of the ACM Web Conference 2024, pp. 2576-3585, 2024.
Ma, S. and Huang, J. A concave pairwise fusion approach to subgroup analysis. Journal of the American Statistical Association, 12(517):410-422) 2017.
Ma, X., Long, L., Moon, S., Adamson, B. J., and Baxi, S. S. Comparison of population characteristics in real-world clinical oncology databases in the US: Flatiron health, SEER, and NPCR. Medrxiv, 2020.
Saito, Y. and Joachims, T. Off-policy evaluation for large ac- tion spaces via embeddings. In International Conference on Machine Learning, pp. 19089-19122. PMLR, 2022.
Saito, Y., Ren, Q., and Joachims, T. Off-policy evaluation for large action spaces via conjunct effect modeling. In international conference on Machine learning, pp. 29734- 29759. PMLR, 2023.
Marmarelis, M. G., Morstatter, F., Galstyan, A., and Ver Steeg, G. Policy learning for localized interventions from observational data. In International Conference Artificial Intelligence and Statistics, pp. 4456-4464. PMLR, 2024.
Schwarz, G. Estimating the dimension of a model. The annals of statistics, pp. 461-464, 1978.
Mo, W., Qi, Z., and Liu, Y. Learning optimal distribution- ally robust individualized treatment rules. Journal of the American Statistical Association, 1 116(534)659-674 2021.
Schweisthal, J., Frauen, D... Melnychuk, V., and Feuerriegel, S. Reliable off-policy learning for dosage combinations, Advances in Neural Information Processing Systems, 36: 67900-67924, 2023.
Murphy, S. A. Optimal dynamic treatment regimes. Jour- nal of the Royal Statisticall Society: Series B (Statistical Methodology), 65(2):331-355 2003.
Shi, C., Fan, A., Song, R., and Lu, W. nigh-dimensional learning for optimal dynamic treatment regimes. Annals of statistics, 46(3):925, 2018.
Nilforoshan, H., Moor, M., Roohani, Y., Chen, Y., Šurina, A., Yasunaga, M., Oblak, S., and Leskovec, J. Zero- shot causal learning. Advances in Neural Information Processing Systems, 36:6862-6901, 6901, 2023.
Song, R., Wang, W., Zeng, D., and Kosorok, M. R. Penal- ized Q-learning for dynamic treatment regimens. Statis- tica Sinica, 25(3):901, 2015.

---

Doubly Robust Fusion of Many Treatments for Policy Learning
Taufiq, M. F., J.-F, Cornish, R., Teh, Y. W., and Doucet, A. Conformal off-policy prediction in contextual bandits. Advances in Neural Information Processing Systems, 35: 91512-31524 2022.
Tibshirani, R., Saunders, M., Rosset, S., Zhu, J., and Knight, K. Sparsity and smoothness via the fused lasso. Jour- nal of the Royal Statistical Society Series B: Statistical Methodology, 67(1):91-108, 2005.
Wainwright, M. J. High-dimensional Statistics: A Non- asymptotic Viewpoint, volume 48. Cambridge University Press, 2019.
Watkins, C. J. and Dayan, P. Q-learning. Machine Learning, 8(3-4):279-299 1992.
Wooldridge, M. Introductory Econometrics: A Modern Approach. South- Weeter Cengage Learning, Mason, OH, 5th edition, 2012. ISBN 978-1-111-5310411.
Wu, L. and Yang, S. Transfer learning of individualized treatment rules from experimental to real-world data. Journal of Computational and Graphical Statistics, 32 (3):1036-1045, 2023.
Xu, Q. Cao, X., Chen, G., Zeng, H., Fu, H., and Qu, A. Multi-label residual weighted learning for individualized combination treatment rule. Electronic Journal of Statis- fics, 18(1):1517-154) 2024a.
Xu, Q., Fu, H., and Qu, A. Optimal individualized treatment rule for combination treatments under budget constraints. Journal of the Royal Statistical Society Series B: Statisti- cal Methodology, 86(3):714-745 20244b
Zhang, B., Tsiatis, A. A., Laber, F. B., and Davidian, M. A robust method for estimating optimal treatment regimes. Biometrics, 6814)1010-1 2012.
Zhang, Y., Laber, F. B., Tsiatis, A., and Davidian, M. Using decision lists to construct interpretable and parsimonious treatment regimes. Biometrics, 71(4):895-904, 2015.
Zhang, Y., Shi, C., and Luo, S. Conformal off-policy pre- diction. In International Conference on Artificial Intelli gence and Statistics, pp. 2751-2766 PMLR, 2023.
Zhang, Y., Huang, M., and Imai, K. Minimax regret estima- tion for generalizing heterogencous treatment effects with multisite data. arXiv preprint arXiv:2412.111388 2024.
Zhao, Y., Zeng. D... Rush, A. J., and Kosorok. M. R. Es- timating individualized treatment rules using outcome weighted learning. Journal of the American Statistical Association, 107(499): 100/499))100--10 1011110611110 1118,1111111 2012.
Zhou, Z. Athey. S., and Wager, S. Offline multi-action policy learning: Generalization and optimization. Opera- tions Research, 71(1):148-188) 2023.
12

---

Doubly Robust Fusion of Many Treatments for Policy Learning
A. Proof
A.1. Rationale for Assumption 3.2
For every treatment a E A, calibration weighting is an optimization problem and can be solved using the method of Lagrange multipliers:
(now,))t1 -
La Ta(w1,.... Wn) Σ
NAT Σ
wi(X; X) ng 1 E Wi
+(++1)
=A,=0
=1,=0
=A,=0
Minimizing La(w....., wp)) gives:
P (X, -X1
dq = w(Xii\iii)
Σ (XTT(X XX]
where the function (1) for different 7 values are summarized in Table 7, and Ā solves the equation
Σ EAA=a
(x, X) = O.
Therefore, Ā an M-estimator and, under standard regularity conditions for M-estimators (Boos & Stefanski, 2013), is root-n consistent.
Table 7. py(x)) for Cressie-Read family.
7
hy(w)
(x
In(nur) nw In(ww)
- x) exp(x)
0
(nw) -1
(1 + 12)
7(7 +
A.2. Lemma A.1
Lemma A.1. Suppose Assumptions 2.1 and 3.3 hold. For any i = 1,...,
n, = 1.....p and a E A.
E[I(A, a)w Xije((a) = 0.
Proof of Lemma A.I. By the unconfoundedness in Assumption 2.1, we have
[[[(A, = a) X.) = E [I(A a) (()e((a) x;}]
E[[A a)w; Xx X,, x.
(i) (correct calibration weighting) in Assumption 3.3 holds, i.e., [T(A, a)w; X:) 1, we have
B[E[I(A; = a) X.)((,)((e((a)) x.]] = EX x.)] = [()
By the definitions of the projection vector and projection residual:
S% argmin E ((()) (a)) x and =(a)== i(a) x c..
(12)
CER
the result follows from
E (x {(a)) = E X [xT(r(1) - x c::]] 0.
Note that (xTT((a) does not require linear model between ((a) and X; holds solely due the projection (12).
13

---

Doubly Robust Fusion of Many Treatments for Policy Learning
If (ii) (correct outcome model) in Assumption 3.3 holds, i.e., B(ee(a)
x 0, we have
[[B( = a)w; [.)xX,)B(s((A) X. = 0.
Note that the above equation does not require calibration weighting to be correct; it holds solely due to E ($(a) XX] 0. condition stronger than E [{)) = 0. This condition is also referred to as the exogeneity assumption or the zero conditional mean condition (Wooldridge, 2012).
A.3. Proof of Theorem 3.8
Proof of Theorem 3.8. We rewrite B, to relate it to the potential outcomes
BB = min
Σ
II(A; = a)i; (%: - x3 (.))
9,, ERP 2n
i=1 aESi
I[A; a) (Vi(a) X, 3ym))
min
Sn,
2n
i=1 RECi
T[A; = a) (x) 3 + =((a) x._ Bm)²
min 3,, ERP
2n
1=1 NEG;
where the second equation is due to I(A = a)Y; I(A; (a) and the third equation is due to (10). Then, the least squares estimation leads to
B 3 =
= a(A==));;xi;;
T(A: = a) (i:((a)
aEg;
aEg;
Twe
We have
max DEB
|13 -,1100 S max BEB
[
#x [6]
| (e
(13)
examine
||T || ooo and |[ |[Taeo/looo respectively.
Step (Bound [|r 1 oo)) By Assumption 3.2, ]e 0, with probability least in (where in + 0 n + 20), W 2 w
- wel 2 Ci - c = C₁ for any i. Thus, we have
| Tx Amax (T (T)xp)) Amin (Taxpp) S C.C, Nmin
where the last inequality is due to Assumptions 3.5. Then, we have
VD N.C.
max
|
1oo 1lo I /p ||I 12 S
(14)
BEB
bEB
Step 2 (Bound [[T 1ooo). We first bound
I [w =
T(A, a) (a)
1=1 aEg;
By Assumptions 3.2 and 3.5, we have
Σ
Σ
I(A, a)(w; < nChCs.
1=1 REg.
14

---

Doubly Robust Fusion of Many Treatments for Policy Learning
Combined with Lemma A.I and Assumption 3.6, for any -I...., and b B, for L >
have
P
Σ a
I(A, - x));
Vn03C3l S 22exp(-12/203
Then, have
max
[ [e
VnC2Cat - 2MMpep((-P/22))
BEB
Letting 2Mpexp(-t(2oo) = 2Mp/n, we have that with probability at least 1 2M/p/n.
max BEB
[ustslls S 2CC3Cao. Iogg())
(15)
Finally, by Assumption 3.2,
ppn
max bEB
|| [xx00]/>>0|1 (e( - we(l) 0m Op
(16)
Nunin
The result follows from (13), (14), (15), and (16).
A.4. Proof of Theorem 3.12
Proof of Theorem 3.12. Recall that the true group structure is U h- Gr The space of ('s that have the true group structure is defined as
zor = { (CT
<k) E Kr vb e B, Va, d' E 96, S. = Sa/
Note that 5 oor
zor, Define the mapping
T: RK Kp + z", SHC,
where Č .Si....,cii _kk and E.a = =Eore; Sa./19%). Va E g.. Define the neighbor of projection vector Ç* as
o = E RKP |ls 10 S )n)
Note that Vg θ, have T(C) O. Define the event Fg (C O). By Theorem 3.8, have ((E)) 2 1 2Mp/n 4. The result follows from the following two statements, each of which will prove separately.
Statement On event Ej, for all such that T(C) + "" have (I((C) ())
Statement 2 There exists an event Eg such that P(E2) Kp/n - lus and on E2, for all Ç θ, we have Qx(C) (C)))
Proof of Statement 1. We examine L, (T(()) Ln(C") and Pa(T(()) P.((""), respectively.
Step 1.1 (Examine L,). By definition, restricted to z cor ca0
C is the unique minimizer of Ln (C), that is, for all Ç θ such
that T()) 7
L, (T(()) (n
(17)
Step 1.2 (Examine P.). For any Č = (ST ......... ".. E 2or no (including the case Č = C -or
_Tk)TT
Pr(C) = Σ Ps. (1 - -carll
IgaKa'EK
treatments belong to different groups, i.c., a E Gš and a' E Gú. with * b/ by Assumption 3.10, we have
11. 11
2 || 12. - -arloo
2 || c-2 Co - -2115 - -''
c)n 2án
cdn/2.
15

---

Doubly Robust Fusion of Many Treatments for Policy Learning
Then, Ph.. (1)5, So (1) is constant by Assumption 3.10. If treatments belong to the same group, i.e., a, gb have Sa Sa' || = 0, thus pA.. (lISa - Sor |lI)) = 0. Overall, we have that Px(() is a constant. Since T({) E 2 n θ and
C E z no on Ei, we have
(T(T(())= = P.
(18)
The Statement follows from (17) and (18).
Proof of Statement 2: Let Ç := T(). We examine (C) Ln (n(() and Pa(C) - P,(C), respectively.
Step 2.1 (Examine L,). By Taylor expansion. there is < & such that Ç= = EC + (1 E)C E O, we have
(S) Ln(() 1
Σ
I(A = a) Y 2.) X Sa - -x, 3.)
12
AEA
1(4 a)
() (Y((a) -
-X, [, (x (x7(.)
Σ
[ a) x7 Sa +c((a) x, { (xi Sa - S.)
71 EE
T(4: a) x (S% -2.) + xxe((a)) (Sa -3a)
n
AEA
Let
1
Σ
(Sa I(A: a) ( -2a) + X.Ei(a)
Va
n
Then, we have (S) - Ln(() = - -nnes Vo (Sa - Sa). Since to = Enweg 111 Va E 9r, by some algebra, we
have
(C) Ln(() --E Σ
T (Sa -aa)
bEB DEgS
Eweg; Saa
Σ
Un Sn-
T
196
BEB REG
U (Sa Sar)
Σ ΣΣΣ
19
BEB BEGg d'Esi
(Va vaa) (Sa Sar)
=
2/98
dEB NES; d'Es;
(va va/) T (Sa Saa))
1961
dEB 0, EG a<aa
10 Vo loc 115a Sar
Σ
1981
By Assumptions 3.5 and 3.2, we have
T(A x)x (S"- - (.))
<
T(A: = -AA
|1: - 2.1
n
= Or(p
Following a similar derivation as in Step 2 of the proof of Theorem 3.8, there exists an event Fgg such that PP(22) 2 1 Kp/n lns and on Ez, we have
I[A = a)o,
O(v log(n)).
16

---

Doubly Robust Fusion of Many Treatments for Policy Learning
Thus, have
1,0 - -22 120 I 2 max Vallzo = Oz (po, + v vnlog(n))
aeA
and
Op (pon + Vnlog(n)
(x( -
Σ E BEB a,a' a.' EDI aaca
ll5a So 1l
(19)
19
Step 2.2 (Examine P.). Following similar arguments in Step 1.2 of the proof of Statement if treatments belong
to different groups, i.e,, a E gr and a' E G. with b & U, by Assumption 3.10, we have IlS.. Sor >> cdn/2 and [ca 1ar >> c2n/2, thus,
Ps., (llsa Sailli) PA. (1 - -arll)) O.
If treatments belong the same group, i.c. a E Gh. we have |[a - Sarll O, thus psa <l|Sa - - 0. However, since IlC. - -S * they are the only terms contributing Pn(C) Pn(C). Thus, we have
Pa(() -v(()) = Σ Σ
Paa (lISa - 11
hEB a, ED2 aacae
PA. (ISa. Sarll1) ||5. - Sarll.
Σ
(20)
155a Sarll
bEB a,a' a.a' Elg aacaa
We have l|5a Sw S 1K.. Calli + 115a 12 p((0+) = 1, and An >> p / Kmin + (on))/)))/ Kmins we have
2Φn By Assumption 3.10, that PA. (-) Anp(-).
Pan (1[.a - S 2
pon + vn log(n)
(21)
[[5a l1
Kmin
The Statement follows from (19) (20), and (21).
B. Additional Simulation Results
B.1. Simulations for increasing K and fixed
We keep the sample size fixed our proposed method outperforms other baselines in terms of fusion quality and policy value (see Tables 8 and 9). The number of recovered groups increases slightly, expected.
1800 and increase the number of treatments K from 16 32 and 48. In such regime,
Table 8. Simulation results for = 32
Method
ARI
Number of groups
Value
policy tree (baseline)
32
8.66 (0.18)
fusion policy tree
0.21 (0.09)
16.74 (3.64)
8.63 (0.18)
CW fusion policy tree
0.85 10)
5.68 (1.72)
8.80 (0.21)
Ma et al. (2022)
0.21 (0.09)
16.74 (3.64)
8.52 (0.12)
CW = Calibration Weighting. ARI (Adjusted Rand Index for fusion quality) and policy value: higher is better. Oracle number of groups 4. Numbers in parentheses are Monte Carlo standard errors. Results are averaged over 200 runs.
B.2. Simulations under misspecified weighting model
We additionally considered a scenario where the outcome mean functions are linear (i.e., correctly specified), but the weighting model misspecified. Specifically, in calibration weighting, we excluded X, and used only X₂ and X3. The
17

---

Doubly Robust Fusion of Many Treatments for Policy Learning
Table 9, Simulation results for 48
Method
ARI
Number of groups
Value
policy tree (baseline)
/
48
8.49 (0.20)
fusion policy tree
0.16 (0.08)
21.00 (5.58)
8.40 (0.30)
CW fusion policy tree
0.74 10)
7.35 (2.31)
8.52 (0.23)
Ma et al. (2022)
0.16 (0.08)
21.00 (5.58)
8.41 (0.12)
CW Calibration Weighting. ARI (Adjusted Rand Index for fusion quality) and policy value: higher is better. Oracle number of groups 4. Numbers in parentheses are Monte Carlo standard errors. Results are averaged over 200 runs.
outcome mean functions were set follows:
Yi 2.5 0.5X, 1,5x2 -3,,
Y/ X1 -2X, - 255x,,
Y/ = 2 0.5X + +X2 2X3;
Y1 = -1 X1 - -X2 X3.
Table 10. Simulation results under misspecified weighting model
Method
ARI
Number of groups
Value
policy tree (baseline)
16
6.35 (0.06)
fusion policy tree
0.88 (0. 13)
5.42 (1.42)
6.41 (0.04)
CW fusion policy tree
0.96 (0.06)
4.46 (0.66)
6.43 (0.02)
Ma et al. (2022)
0.88 (0.13)
5,42 (1.42)
6.39 (0.00)
CW Calibration Weighting. ART (Adjusted Rand Index fusion quality) and
policy value: higher is better. Oracle number of groups 4. Numbers in parentheses are Monte Carlo standard errors. Results are averaged over 200 runs.
Table 10 presents the results. As the outcome models correctly specified, both the fusion (based on linear model) and the CW fusion approaches achieved strong ARI scores, illustrating the double robustness property of CW fusion. Ma et al. (2022) also performed well, as their method relies on the same linearity assumption, which holds in this setting, Nevertheless, our proposed method consistently achieved the best overall performance. Note that the simulation in Section already showed the advantage of the proposed method when the outcome model is misspecified while the weighting model is correct.
18


Reviews:
[{"review_id": "KtZdeXLudw", "reviewer": "Reviewer_c8W6", "summary": "The authors study causal inference with many treatments and sparse data within each group, which makes estimating treatment effects challenging. They observe that many treatments share commonalities and can be clustered to reduce the effective dimensionality of the treatment space. Within each group, treatments are treated as equivalent, allowing the application of efficient multi-armed ITR learning methods. The authors propose calibration-weighted treatment fusion, a technique that leverages balancing weights in a regression model to estimate the group structure of treatments in a doubly robust manner. Using these fused groups, they propose AIPW estimators for policy learning.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your insightful and constructive comments. Below, we have provided our detailed, point-by-point responses.\n\n**1. Experiments for increasing $K$ and fixed $n$**\n\nWe keep the sample size fixed at $n=1800$ and increase the number of treatments $K$ from 16 to 32 and 48. In such regime, our proposed method outperforms other baselines in terms of fusion quality and policy value. The number of recovered groups increases slightly, as expected. Our method may become unstable when the number of treatments approaches the sample size, leaving very few observations per arm. In such regimes, iterative weight learning, as suggested by the reviewer, may help stabilize performance (see point 5).\n\n---\n**Table S1**: Simulation results for $K=32$; CW = Calibration Weighting; ARI (Adjusted Rand Index for fusion quality) and policy value: higher is better; oracle number of groups = 4.\n\n|Method|ARI|Number of Groups|Value|\n|-|-|-|-|\n|policy tree (baseline)|/|32|8.66 (0.18)|\n|fusion + policy tree|0.21 (0.09)|16.74 (3.64)|8.63 (0.18)|\n|CW + fusion + policy tree|**0.85** (0.10)|**5.68** (1.72)|**8.80** (0.21)|\n|Ma et al. (2022)|0.21 (0.09)|16.74 (3.64)|8.52 (0.12)|\n---\n**Table S2**: Simulation results for $K=48$.\n\n|Method|ARI|Number of Groups|Value|\n|-|-|-|-|\n|policy tree (baseline)|/|48|8.49 (0.20)|\n|fusion + policy tree|0.16 (0.08)|21.00 (5.58)|8.40 (0.30)|\n|CW + fusion + policy tree|**0.74** (0.10)|**7.35** (2.31)|**8.52** (0.23)|\n|Ma et al. (2022)|0.16 (0.08)|21.00 (5.58)|8.41 (0.12)|\n---\n\n**2. Convergence rate of the weights**\n\nIn the fusion phase, Assumption 3.1 posits parametric-rate estimation of the weights to ensure consistency of the estimated grouping under outcome model misspecification. Specifically, if the weight model is correctly specified and estimated at the parametric rate, our method can tolerate full misspecification of the linear working outcome model, which is an appealing feature in settings with many treatment arms and limited per-arm data. In contrast, allowing slower convergence for the weight estimator (e.g., $o(n^{-1/4})$) would require the development of doubly robust estimators and the estimation of an additional nuisance model, specifically the true conditional mean function of the outcome, at a rate of at least $O(n^{-1/4})$. Such doubly robust estimators are appealing when there is a large sample size per treatment arm. However, in our setting with limited data per arm, more flexible models may result in overfitting and unstable performance.\n\nIn the policy learning phase, Assumption 3.1 is no longer required for the weight model. Instead, we adopt Assumption 3.13, which imposes a weaker rate double robustness condition, as the reviewer suggested. This relaxation is justified by the increased sample size within each fused group, enabling the use of more flexible models for both the weight and outcome components.\n\n**3. Synthetic combinations**\n\nWe clarify that our setting involves $K$ treatment levels without combinatorial structure, while [2] considers $2^p$ combinations from $p$ binary treatments. To reflect this distinction, we will revise our title to “Doubly Robust Fusion of Many Treatments for Policy Learning,” following Reviewer UeTD’s suggestion and literature conventions (Ma et al., 2022, 2023).\n\nWhile our method uses calibration-weighted fused lasso to reduce the action space for policy learning, synthetic combinations aim to impute counterfactuals under structural assumptions. These are complementary strategies, and we added a discussion in the revised manuscript.\n\n**4. Arms with few or even zero units**\n\nIf some treatment arms have very few or even zero observations, there is little information for those arms, and our method might not be stable in such cases unless additional structural assumptions across the $K$ treatments are imposed, such as the combinatorial structure used in [2]. The instability of learned balancing weights in these challenging scenarios highlights the difficulty of the problem and motivates potential methodological refinements. For example, one possible solution is to conduct iterative weighting and fusion, as discussed in the next point.\n\n**5. Iterative learning of weights for fused treatments**\n\nIt is beneficial to estimate weights for the fused treatments, as there might be more data in each treatment group than before fusion, which is also our main motivation for performing fusion. See also point 2 regarding the convergence rate of the weights for the fusion and policy learning phases. \n\nIn the fusion phase, while our current approach performs a single round of fusion, we appreciate the reviewer’s suggestion of iterative learning of weights to refine the preliminary fusion and discussed this extension to handle challenging cases with few observations in some treatment arms in the manuscript.\n\nIn the policy learning phase, we indeed re-estimate the weights for the fused treatments when learning the individualized treatment rule.", "claims_and_evidence": "Theoretical claims are supported by theorems and proofs. The authors provide formal guarantees for consistent estimation of the counterfactual mean parameters of treatments and recovery of fused treatment groups. They then provide regret bounds for a Cross-Fitted AIPW Policy Learning algorithm that leverages their fused treatment approach.", "methods_and_evaluation_criteria": "The authors evaluate their approach in synthetic and real-world datasets against reasonable benchmarks. Their method performs well in the simulations and the real-world dataset highlights the practical applicability and interpretability of the method in policy learning.\n\nIt would be beneficial to provide experiments investigating the performance of the method as the number of treatments increases for a fixed sample size, with the number of observations per treatment thus decreasing. Does the method become unstable? What regimes would the method be expected to work well in, and where might it not perform well?", "theoretical_claims": "The theoretical claims seem correct but I need not check the proofs. \n\nAssumption 3.1 requires parametric-rate estimation of the weights. Is this necessary? Since the estimator is doubly robust with respect to the outcome and weight models, I would expect that a weaker condition, such as faster than n^{-1/4} convergence of nuisances, might suffice, e.g., as in [1].\n\n\n[1] Chambaz, Antoine, Pierre Neuvial, and Mark J. van der Laan. \"Estimation of a non-parametric variable importance measure of a continuous exposure.\" Electronic journal of statistics 6 (2012): 1059.", "experimental_designs_or_analyses": "The experimental design looked well thought out. However, as mentioned in Methods And Evaluation Criteria, it would be useful to see how the method performs in a more diverse range of scenarios characterizing by data sparsity and number of treatments.", "relation_to_broader_scientific_literature": "Related literature is discussed.", "essential_references_not_discussed": "Some relevant work worth mentioning is [2] and the references therein on synthetic combinations, which provide a synthetic control-like framework for estimating counterfactual means using low-rank matrix completion when the number of possible treatments is large or combinatorially explosive.\n\n\n\n[2] Agarwal, Abhineet, Anish Agarwal, and Suhas Vijaykumar. \"Synthetic combinations: A causal inference framework for combinatorial interventions.\" Advances in Neural Information Processing Systems 36 (2023): 19195-19216.", "other_strengths_and_weaknesses": "Strength: The approach tackles a notable problem in causal inference with many treatments, where treatment groups may be sparse and have insufficient overlap. The idea of grouping similar treatments to reduce the dimensionality of the treatment space provides a practical solution. A nice feature of the proposed method is that once these groups are learned using their method, standard causal inference and policy learning methods can then be used with these fused treatments.\n\nWeaknesses: The method requires estimation of balancing weights within each treatment arm. Since each treatment arm may contain very few observations, it seems that this could be quite challenging to estimate well. For example, in some settings, as considered in [2], some treatment arms may have very few or even zero observations. In such cases, I imagine the authors' technique may not apply or may perform poorly.", "other_comments_or_suggestions": "See other parts for comments.", "questions_for_authors": "When learning the balancing weights, would it be beneficial to use weights for the fused treatments, such that more data can be used to learn them? For instance, one could use an iterative approach that first learns a preliminary fused treatment and then estimates the calibrated weights for the fused treatments. This may allow the application of this method when treatment arms have very few observations such that learning weights is infeasible.", "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": "", "extra_scores": {}}, {"review_id": "vp1ZzeKmIo", "reviewer": "Reviewer_UeTD", "summary": "The paper proposes to learn an optimal individualized treatment regime in settings with a large number of treatment actions. To handle large action space, the author(s) proposed a novel doubly robust method for grouping treatments into a few categories to enable the application of existing methodologies developed with few treatment options to handle these settings. They further validate their proposed method both theoretically and empirically. Specifically, in theory:\n\n1.  They established the doubly robustness of their group procedure, requiring the correct specification of either the propensity score or the outcome regression model. \n2. They proved the oracle property of their estimated $\\widehat{\\zeta}$.\n3. They derived the regret bound of the resulting estimated optimal treatment regime. \n\nEmpirically, they conducted simulations and real data analyses to demonstrate the superior finite sample performance of their proposal.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your insightful and constructive comments. Below, we have provided our detailed, point-by-point responses.\n\n**1. Simulations for validation of the double robustness**\n\nWe additionally considered a scenario where the outcome mean functions are linear (correctly specified), while the weighting model is misspecified. Specifically, for calibration weighting, we only used $X_2$ and $X_3$, omitting $X_1$.\n\nSince the outcome models are correctly specified, both the *fusion* and *CW + fusion* approaches achieved strong ARI scores, demonstrating the double robustness property. Ma et al. (2022) also performed well since the linearity assumption they rely on holds in this setting. However, our proposed method still achieved the best performance.\n\nNote that the simulation in Section 4.1 already showed the advantage of the proposed method when the outcome model is misspecified while the weighting model is correct.\n\n**Table S3**: Simulation results under misspecified weighting model; CW = Calibration Weighting; ARI (Adjusted Rand Index for fusion quality) and policy value: higher is better; oracle number of groups = 4.\n\n|Method|ARI|Number of Groups|Value|\n|-|-|-|-|\n|policy tree (baseline)|/|16|6.35 (0.06)|\n|fusion + policy tree|0.88 (0.13)|5.42 (1.42)|6.41 (0.04)|\n|CW + fusion + policy tree|**0.96** (0.06)|**4.46** (0.66)|**6.43** (0.02)|\n|Ma et al. (2022)|0.88 (0.13)|5.42 (1.42)|6.39 (0.00)|\n---\n\n**2. Real-world experiments**\n\nThe real-world dataset motivates our treatment fusion approach due to practical challenges such as limited sample size per treatment arm. We acknowledge that the grouping from 7 to 5 is determined by the underlying structure of the real data.\n\nTo further illustrate the effectiveness of our method in larger treatment spaces, our simulation study demonstrates its ability to reduce 16 treatment arms to 4 groups. We also consider scenarios with more treatments in the next response point.\n\n**3. Simulation with more treatments**\n\nWe increased the number of treatments to 32 and 48 while keeping the total sample size fixed. We refer to Tables S1 and S2 in the response for Reviewer c8W6, point 1, for results. The proposed method continues to outperform other methods.\n\n**4. Completeness assumption**\n\nTo guarantee recovery of the oracle grouping, we additionally imposed the completeness assumption: for any function $h(\\cdot)$, if $\\mathbb{E}\\\\{X h(X)\\\\}=0$, then $h(X)=0$ almost surely.\n\n**5. Presentation**\n\nWe will streamline the manuscript by adding plain-language explanations (e.g., after Theorem 3.8) and moving heavy notation and technical content to the appendix to enhance accessibility.\n\nWhile $w^*$ equals the inverse propensity score when the model is correct, it is not in general, as we allow model misspecification; thus, the second part of Assumption 3.1 is retained.\n\nWe will revise Assumption 3.2 based on your suggestion.\n\n**6. Assumption 2.1 (Identification)**\n\nWe replaced “SUTVA” with the consistency and removed $P(X = x) > 0$.\n\n**7. Title terminology**\n\nWe revised the title and text to use “many treatments” instead of “high-dimensional treatments,” following conventions in the literature (Ma et al., 2022, 2023). \n\n**8. Assumption 3.1 (Convergence of calibration weight)**\n\nFor every treatment $a\\in\\mathcal{A}$, calibration weighting is an optimization problem and can be solved using the method of Lagrange multipliers:\n$$\nL\\_a(w\\_1,\\ldots,w\\_n)=\\sum_{i:A\\_i=a}\\\\{\\gamma(\\gamma+1)\\\\}^{-1}\\\\{(n\\_a w\\_i)^{\\gamma+1}-1\\\\}-n\\lambda^\\top \\sum\\_{i:A\\_i=a}w\\_i(X_i-\\bar{X})+n\\varphi\\left(1-\\sum_{i:A\\_i=a} w_i\\right).\n$$\nMinimizing $L\\_a(w\\_1,\\ldots,w\\_n)$ gives:\n$$\n\\hat{w}\\_i=w(X_i;\\hat\\lambda)=\\frac{\\rho\\_{\\gamma}[\\hat\\lambda^\\top(X_i-\\bar{X})]}{\\sum\\_{j:A_j=a} \\rho\\_{\\gamma}[\\hat\\lambda^\\top(X_j-\\bar{X})]},\n$$\nwhere the function $\\rho\\_{\\gamma}(x)$ for different $\\gamma$ values are summarized in Table S4 below, and $\\widehat{\\lambda}$ solves the equation\n$$\n\\sum\\_{i:A\\_i=a}\\rho\\_{\\gamma}\\[\\lambda^\\top(X\\_i-\\bar{X})\\](X\\_i-\\bar{X})=0.\n$$\nTherefore, $\\hat{\\lambda}$ is an M-estimator and, under standard regularity conditions for M-estimators (Boos and Stefanski, 2013), it is root-n consistent.\n\n**Table S4**: $\\rho\\_{\\gamma}(x)$ for Cressie-Read family.\n\n|$\\gamma$|$h\\_{\\gamma}(w)$|$\\rho\\_{\\gamma}(x)$|\n|-|-|-|\n|$-1$|$-\\ln(nw)$|$(1 - x)^{-1}$|\n|$0$|$nw\\ln(nw)$|$\\exp(x)$|\n|$\\gamma$|$\\frac{(nw)^{\\gamma+1} - 1}{\\gamma(\\gamma+1)}$|$(1 + \\gamma x)^{1/\\gamma}$|\n---\n\n**9. Bounded assumption**\n\nThe bounded outcome and covariates assumption is only required for the results in Section 3.3 and not for those in Section 3.2.\n\n**10. Generality of results**\n\nPropositions 3.16 and 3.17 do not rely on the policy class being restricted to trees. We revised the section title to “Multi-armed Policy Learning” to reflect the broader applicability.\n\n**11. Related literature**\n\nWe will incorporate the recommended references into the revised manuscript to better situate our work within the broader literature.", "claims_and_evidence": "The double robustness, oracle property, and regret bounds are theoretically verified. Numerically, the authors conducted experiments to evaluate the performance of their proposed treatment regime and the (ARI of their grouping method. To further strengthen the claims, it would be beneficial for the authors to conduct additional simulations to empirically validate the double robustness property.", "methods_and_evaluation_criteria": "The paper utilizes a real-world dataset with 7 treatment options, which partially motivates the problem, as the scenario extends beyond the typical two-arm case commonly studied in the literature. However, 7 is not an extraordinarily large number. In fact, the analysis groups these treatments into 5 categories. I would not consider the reduction from 7 to 5 as substantial enough to lead to significant enhancements.", "theoretical_claims": "The theories appear reasonable to me, though I did not check the proofs in detail.", "experimental_designs_or_analyses": "As I mentioned earlier, the analysis is theoretically sound and partially supported by numerical experiments. However, one potential limitation of the experimental design is that the current simulation considers only 16 treatment groups. To better address scenarios with truly large action spaces, I would suggest increasing the number of treatment groups to 50 or even 100 in simulations.", "relation_to_broader_scientific_literature": "The paper makes valuable contributions that extend beyond the existing literature. Scientifically, large action spaces have been relatively underexplored in the context of individualized (dynamic) treatment regimes. Methodologically, the proposed approach is novel and supported by theoretical guarantees, including the proven double robustness property. A particularly appealing feature of the method is that it does not rely on linear model assumptions. Instead, it uses linear models solely as working models to facilitate the application of fused lasso for grouping treatments. This represents a significant advancement over prior work, such as Ma et al. (2022).", "essential_references_not_discussed": "I do not think any essential references are missing. However, there is an extensive literature on estimating (dynamic) treatment regimes, and many related papers exist. For further details, please refer to my comments in the \"Other Comments & Suggestions\" section.", "other_strengths_and_weaknesses": "I have discussed the strengths of the paper in the \"Summary\" and \"Relation To Broader Scientific Literature\" Sections. Some potential weakness are given below:\n\n1. As I mentioned, one strength of the paper lies in that they do not rely on the linearity assumption. However, this comes at a price for their grouping method. Specifically, their identified groups might be different from the oracle groups. This is because when the oracle groups are the same, it ensures that their identified groups are asymptotically the same, based on the equation above (9). However, the converse is not true. I was wondering if the author(s) can impose some completeness-type assumptions to guarantee the reverse also holds. \n\n2. Another potential limitation lies in the presentation. The paper relies heavily on notations and mathematical expressions, which could be streamlined. It would be ideal if the authors could reduce the number of notations and replace unnecessary formulas with words. For example, the discussion following Theorem 3.8 is intended to explain the theorem, but it itself involves several mathematical notations, making it less accessible. Similarly, Assumption 3.9 is highly technical, and Remark 3.10, which aims to clarify it, remains quite technical as well. Some of the assumptions and notations could be simplified. For instance, $w_i^*$ essentially represents the inverse propensity score, so the second part of Assumption 3.1 is automatically satisfied due to the positivity assumption and does not need to be explicitly imposed. Likewise, Assumption 3.2 fundamentally requires the correct specification of the propensity score, which could be stated more straightforwardly.", "other_comments_or_suggestions": "1. Assumption 2.1(i) Should be the consistency assumption, instead of SUTVA. SUTVA requires the potential outcome depends on its own treatment value, which differs from the consistency which requires the observed outcome to equal the potential outcome.\n\n2. I will not include $\\mathbb{P}(X=x)>0$ in Assumption 2.1(iii). It is not the typical positivity assumption referred to in the literature.  \n\n3. Page 2, K is probably not the treatment dimension, but the treatment cardinality. Similarly, the title refers to \"high-dimensional treatments,\" but it seems more appropriate to use \"high-cardinality treatments.\"\n\n4. The first part of Assumption 3.1 is more like a claim than an assumption. Would it be possible to provide a proof for this rather than assuming it holds?\n\n5. You imposed the bounded outcome and covariates only in Section 3.3. Does this imply that the theories in Section 3.2 do not rely on this assumption?\n\n6. Propositions 3.16 and 3.17 seem quite general. Do they truly depend on specifying the policy class as trees? If not, I suggest revising the title of this section and placing less emphasis on policy trees to better reflect the generality of the results.\n\n7. As I mention earlier, there is a large literature on individualized (dynamic) treatment regimes. The author(s) might want to include more relevant papers. For instance:\n\n    * Q-learning introduced in the introduction is adapted from the well-known Q-learning algorithm by Watkins https://link.springer.com/article/10.1007/BF00992698. There are also follow-up works on dynamic treatment regimes; see e.g., https://www3.stat.sinica.edu.tw/sstest/j25n3/j25n35/j25n35.html. \n    * In addition to the work by Zhang et al. (2012), outcome-weighted learning https://pmc.ncbi.nlm.nih.gov/articles/PMC3636816/ also belongs to the value-search method. Linear ITRS have been explored thoroughly in the literature, including in Zhang et al. (2012), the outcome-weighted learning paper, and e.g., https://pmc.ncbi.nlm.nih.gov/articles/PMC5966293/. \n    * Similarly, tree-type methods have been explored very early by https://pubmed.ncbi.nlm.nih.gov/26893526/ and  https://arxiv.org/abs/1504.07715. It is unclear why the authors cite Athey and Wager’s paper as the sole representative of this line of research. \n    * Finally, https://proceedings.neurips.cc/paper_files/paper/2021/file/816b112c6105b3ebd537828a39af4818-Paper.pdf also studies the grouping of actions, but in a continuous action space.", "questions_for_authors": "Please refer to the section above.", "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": "", "extra_scores": {}}, {"review_id": "CUxae6fBFx", "reviewer": "Reviewer_UwD5", "summary": "This paper addresses the problem of learning individualized treatment rules (ITR) in the high-dimensional setting when the cardinality of the action space is large. This setting can be challenging given the search space for ITR grows exponentially with the cardinality of the action space. The paper proposes a calibration-weighted fusion approach that robustly groups treatments into a smaller number of groups thereby reducing the effective dimensionality of the action space. The method first uses calibration weighting to re-weight observations so that the covariate distributions in each treatment group resembles the overall population, and then use fused Lasso on a linear model to group treatments whose outcome are similar conditioned on covariates.\n\nThe approach is doubly robust: correct group recovery is guaranteed if either the calibration weighting model is correctly specified, or the linear outcome model is correctly specified. After forming these fused groups, the user can then employ standard multi-armed policy-learning methods. The paper shows that the final learned policy is near-optimal with regret bounds. It also provides simulations and a real-world example on CLL and SLL to demonstrate the benefits of the approach in terms of group recovery and policy value optimization.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your insightful and constructive comments. Below, we have provided our detailed, point-by-point responses.\n\n**1. Necessity of recovering group structure**\n\nWhen there are many treatments and the sample size is limited, directly performing policy learning presents several challenges: (i) the outcome model cannot be too complex or flexible due to limited data per arm, (ii) covariate shifts across treatment groups exacerbate instability, especially when using balancing methods like IPW, and this instability is amplified for underrepresented arms with small and highly variable propensity scores, and (iii) the theoretical guarantees (e.g., regret bounds) of well-established methods such as multi-arm policy trees (Zhou et al., 2023) typically require the number of treatment levels to be fixed.\n\nBy using our doubly robust fusion approach followed by policy learning over the fused treatments, we address these issues: (i) more complex and flexible outcome models can be used due to increased sample sizes within fused groups, (ii) weights can be more stably estimated across fused treatments, and (iii) by reducing the number of treatments, existing multi-arm policy learning methods with theoretical guarantees can be directly applied.\n\nWe acknowledge that there are alternative approaches for policy learning without explicit group recovery, such as end-to-end methods mentioned by Reviewer 7HRd. However, we emphasize that our approach (i) avoids potential overfitting issues when sample sizes are small, (ii) is easy to implement using fused lasso and calibration weighting, (iii) is readily compatible with state-of-the-art multi-arm policy learning algorithms, and (iv) offers interpretability, which is especially desirable in high-stakes domains such as medicine.\n\n**2. Approximate equality of $\\mu_a(X)$ and gap tolerance**\n\nWe define treatment groups based on equality of the expected outcome functions to ensure identifiability and interpretability of the fused structure, which is central to our theoretical analysis. While exact equality may seem restrictive, it serves as a natural starting point for grouping treatments with similar effects and enables formal guarantees.\n\nIn practice, our method does tolerate small differences in $\\mu_a$ due to sampling variability and regularization. The fused lasso penalization encourages the fusion of approximately equal effects, allowing treatments with similar but not identical outcomes to be grouped together.\n\nWe explicitly provide the gap tolerance required for recovery in Remark 3.10. For reference, the condition is:\n$$\n\\min\\_{b\\neq b^\\prime} \\| \\boldsymbol{\\beta}\\_b^* -\\boldsymbol{\\beta}\\_{b^\\prime}^* \\|\\_{\\infty} / c > \\lambda\\_n \\gg \\sqrt{n\\log(n)}/K,\n$$\nwhere $\\boldsymbol{\\beta}\\_b^*$ and $\\boldsymbol{\\beta}\\_{b^\\prime}^*$ are the group-shared projection vectors from groups $b$ and $b^\\prime$, $c$ is a constant, $\\lambda_n$ is the penalization level, $n$ is the sample size, and $K$ is the number of treatments.\n\n**3. When is group recovery necessary for policy learning?**\n\nWe agree that pre-grouping is not strictly necessary for all policy learning scenarios. In the example the reviewer described, if the primary goal is simply to identify the single best action with the highest expected outcome, then fully recovering the group structure among suboptimal actions may not be essential.\n\nHowever, in many practical settings, especially in personalized decision-making, the goal is to learn individualized treatment rules rather than selecting a single best action for the entire population. In such cases, identifying meaningful treatment groupings can improve estimation stability, interpretability, and computational tractability, particularly when the action space is large and the sample size is limited. Grouping also allows us to leverage shared structure across similar treatments to reduce variance and enhance overall performance.", "claims_and_evidence": "The paper focuses on the problem of partitioning high dimensional treatments into groups such that treatments within each group have identical conditional outcomes. For this purpose they give a calibration weighting approach that is doubly robust for group recovery, i.e., it can recover the correct grouping if either the calibration or the linear outcome model is correct. Once treatments are grouped, standard multi-armed policy learning methods can be applied to the grouped space with a near-optimal policy regret bound.\n\nThe theoretical claims are supported by proofs. But I am a bit skeptical why is it necessary to recover such groups for policy learning policy. For example, there can be scenarios where exact recovery of groups is more costly than directly doing policy learning, or there may be other techniques that does policy learning without full recovery of the groups.", "methods_and_evaluation_criteria": "The paper offers a valuable approach for merging many treatments into fewer groups, thus enabling more effective policy learning techniques for high dimensional settings. The method’s double-robustness is appealing, and the real data application underscores its practicality and interpretability.", "theoretical_claims": "The paper shows double robustness of the proposed method; the group structure is recovered if either the calibration weighting model is correct or the linear outcome model is correct (Theorem 3.8, 3.11). Minimizing the weighted fused Lasso objective recovers the “oracle” solution that merges arms that share identical $μ_a(x)$. If the grouping is correct, the final policy learning with an M-armed approach achieves an $O(\\sqrt{n})$ regret bound (Propositions 3.16, 3.17).", "experimental_designs_or_analyses": "The paper provides empirical results in both synthetic scenarios and real data application. \n\nIn the synthetic scenarios, it tests the proposed method on a simulation that involves 16 treatments that form 4 true groups, each group having 4 arms that share the same outcome function and the covariate distribution differs across arms. It shows that using calibration weight, fusion and policy tree obtains the best policy value compared with baselines. \n\nIn the real data scenario, it tests on a dataset with 10,346 patients with CLL/SLL with 8 treatment arms and show that the their method with calibration weighting, fusion and policy tree yield 5 groups that implies clinical patterns in the dataset.", "relation_to_broader_scientific_literature": "The paper extends standard policy learning to handle the high dimensional case with a large number of arms through identifying group structures in among the arms that share the same outcome model. The paper combines calibration weight with fused Lasso for a doubly robust approach.", "essential_references_not_discussed": "n.a.", "other_strengths_and_weaknesses": "Strengths:\n- The paper extends standard policy learning to handle the high dimensional case with a large number of arms through identifying group structures in among the arms that share the same outcome model. The paper combines calibration weight with fused Lasso for a doubly robust approach. \n\nWeaknesses: \n- It seems restrictive that the grouping method only works for groups of action that shares the same outcome function. And it is unclear why one necessarily need to identify all such groups for policy learning purposes.", "other_comments_or_suggestions": "n.a.", "questions_for_authors": "- The paper defines the grouping structure for different actions $a\\in A$ when the expected outcome $\\mu_a(X)$ are the same. Why does it have to be equality and not approximate equality? And what should be the gap tolerance between different expected outcomes for the algorithm to work?\n- Is the grouping necessarily needed for policy learning purposes and is it always the best thing to do to identify this group structure before optimizing? Are there scenarios when, e.g. one of the actions has very high expected outcome, and the other actions have some complicated group structures, but finding this best action doesn't necessarily require identifying the group structure within the suboptimal actions?", "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": "", "extra_scores": {}}, {"review_id": "3S2fHFUsbx", "reviewer": "Reviewer_7HRd", "summary": "The paper proposes a new methodology for policy learning in scenarios where many treatments are available. In such settings, naive methods may fail due to e.g., lack of overlap. The proposed methodology assumes that treatments can be assigned into \"groups\" that have the same treatment effect and learn this grouping data-driven. Then, standard policy learning is applied on top of the reduced treatment space.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your insightful and constructive comments. Below, we have provided our detailed, point-by-point responses.\n\n**1. Behavior under assumption violation**\n\nOur method assumes a ground-truth group structure to establish theoretical guarantees such as consistency of group recovery. This enables identifiability and recovery bounds. While potentially restrictive, it provides a principled foundation for analysis.\n\nIn practice, if no exact group structure exists but some treatments have similar effects, the fused lasso encourages grouping of approximately equal effects. If all treatments differ beyond a certain gap, our method will not fuse them. Due to space limit, we refer to\n\nRemark 3.10, which gives the gap tolerance condition; please also see our response to Reviewer UwD5, point 2, for additional context.\n\n**2. Additional experiments and details**\n\nWe have added additional simulations for more treatments and under misspecified weight models; see Tables S1 and S2 in the response for Reviewer c8W6 and Table S3 in the response for Reviewer UeTD.\n\nWe will move the experimental setup tables to the Appendix.\n\n- Results are averaged over 200 runs; variance is reported in the revised Table S5 below.  \n- Policy tree (Zhou et al., 2023) is applied with default tuning.  \n- Fused lasso uses extended BIC (line 159) for model selection, a standard criterion for recovering sparse structure.\n- Overall, when no group structure exists, fusion trades variance for bias: it increases effective sample size but may introduce bias if outcome functions differ. Such trade-off is well-studied in multi-source learning and data fusion literature, and can improve overall performance when data are limited.\n\n**Table S5**: Simulation results; CW = Calibration Weighting; ARI (Adjusted Rand Index for fusion quality) and policy value: higher is better; oracle number of groups = 4.\n\n|Method|ARI|Number of Groups|Value|\n|-|-|-|-|\n|policy tree (baseline)|/|16|8.77 (0.08)|\n|fusion + policy tree|0.26 (0.14)|10.725 (1.93)|8.78 (0.09)|\n|CW + fusion + policy tree|**0.96** (0.06)|**4.335** (0.60)|**8.89** (0.11)|\n|Ma et al. (2022)|0.26 (0.14)|10.725 (1.93)|8.51 (0.12)|\n---\n\n**3. Related work on complex treatments**\n\nWe will incorporate the recommended references into the revised manuscript to better position our work within the broader literature.\n\n**4. Presentation**\n\nWe revised Section 3.2 to begin with a summary of the key results without technical details, followed by the necessary mathematical intuition and formalism.\n\n**5. Clarification on DGP assumptions**\n\nOur method does not assume a parametric model on the outcome function. Instead, we rely on projections onto a linear working model without assuming the model is correctly specified. This allows our method to remain robust to outcome model misspecification. \n\nOn the second bullet point under weaknesses, the sentence after \"Second–\" appears incomplete. Could you kindly clarify the intended point?\n\n**6. Motivation for calibration weighting and fused lasso over end-to-end approaches**\n\nWe adopt a linear fused lasso in the fusion step because of data sparsity prior to grouping. More flexible models like neural networks may overfit in this setting. To mitigate potential model misspecification from using a simple working model, we incorporate calibration weighting, which offers a doubly robust guarantee.\n\nWhile end-to-end approaches such as neural networks are possible alternatives, they often require larger sample sizes to avoid overfitting and may lack interpretability. In contrast, our method is (i) easy to implement using fused lasso and calibration weighting, (ii) compatible with existing multi-arm policy learning algorithms, and (iii) interpretable, which is especially important in high-stakes applications such as medicine.\n\n**7. Comparison with overlap-based approaches for continuous treatments**\n\nThank you for highlighting the well-established approach that focuses on treatments with sufficient data (e.g., overlap regions) and the related reference. We will include them in the related work and discussion. This approach is well-suited for continuous, multi-dimensional treatments, such as dosage combinations, and leverages neural networks and constrained optimization to model joint dose-response surfaces. Its strength lies in capturing drug-drug interactions and ensuring reliable off-policy evaluation in such challenging settings.\n\nIn contrast, our doubly robust fusion method targets discrete treatments with many levels. In this setting, we are able to utilize information from all regions, not just those with strong overlap. For example, new treatments with few observations but potentially strong effects can be fused and compared to standard-of-care treatments with sufficient historical data. Rather than excluding underrepresented treatments, our method allows for their inclusion through meaningful fusion, enabling a more comprehensive and interpretable evaluation.", "claims_and_evidence": "Yes and no. Both theoretical and empirical evidence are provided. However, the theory relies on strong assumptions, and the experimental section could be improved (more details below).", "methods_and_evaluation_criteria": "The authors choose reasonable metrics and datasets for evaluation. Benchmarking on synthetic datasets is standard in causal inference as ground-truth is not available for real-world data.", "theoretical_claims": "The authors provide a comprehensive theoretical analysis and show consistency of their approach. However, this relies und (potentially strong) assumptions, such as assuming a ground-truth group-structure within the treatments (see more in weaknesses below).", "experimental_designs_or_analyses": "The experimental section is relatively weak. The proposed method is only benchmarked on a single simulated dataset. The insights on the real-world dataset are nice but do not compare baselines/ ablations (which I am fine with since evaluation on real-world data is difficult). Most tables in Sec. 4 do not show results but details regarding experimental setups (which could also be moved to the Appendix). Furthermore, many details are not reported.\n- Over how many runs are averaged? There is no variance reported in Table 5\n- How are hyperparameters chosen? Did the authors use hyperparameter tuning?\n- How were the lasso parameters chosen? The proposed method seems to find the correct sparseness structure, but could this be due to a different lasso penalty parameter?\n- How sensitive is the proposed method with respect to violations of assumptions (e.g., what if there is no sparse group structure in the DGP?)\n\nIn summary, I think the experimental section needs to be significantly improved.", "relation_to_broader_scientific_literature": "The key contribution to learn a sparse group structure to reduce the treatment space for policy learning seems to be novel. However, there are a few related ideas in the literature on \"complex\" treatments that have not been acknowledged. Examples:\nhttps://proceedings.neurips.cc/paper/2021/file/d02e9bdc27a894e882fa0c9055c99722-Paper.pdf\nhttps://proceedings.neurips.cc/paper_files/paper/2023/hash/15ddb1773510075ef44981cdb204330b-Abstract-Conference.html\nhttps://dl.acm.org/doi/10.1145/3459637.3482349\nhttps://proceedings.neurips.cc/paper_files/paper/2023/hash/d69103d7895f4e2083f24b664003d386-Abstract-Conference.html", "essential_references_not_discussed": "I am not aware of any existing paper that proposes to learn a group treatment structure.", "other_strengths_and_weaknesses": "Strenghts:\n- The general idea to \"group\" multiple treatments into a single one and then apply policy learning on a reduced treatment space seems appealing, especially in settings with many treatments in which we often do not have sufficient overlap to do proper policy learning.\n\nWeaknesses:\n- The writing could be improved, especially in Sec. 3.2. The authors start introducing various quantities to outline a sketch of proof for their theoretical results. In my opinion, it would help the flow of the paper to first state the main results and then potentially provide some intuition later on.\n- The proposed methodologies relies on strong assumptions. First, parametric structure is imposed on the DGP (linear interactions between treatments and covariates). Second-", "other_comments_or_suggestions": "No further comments", "questions_for_authors": "In summary, I like the main of idea of the paper. My two major concerns are:\n- the experimental section (see above)\n- The assumptions imposed and violations thereof. Could the authors elaborate on this point and outline what would happen if, e.g., the proposed method would try to enforce a group structure that does not exist within the DGP? \n\nAnother question: I am not sure if I understand the (rather complicated) motivation of using calibration weights and a linear lasso approach to learn the group structure. Why not simply train, e.g., a neural net to learn a group structure and a policy that maximizes policy value (in an end-to-end manner)?\n\nFinally, another well-established approach to deal with large treatment spaces is to only focus on the treatments for which sufficient data is available (e.g., overlap regions; see an example here: https://proceedings.neurips.cc/paper_files/paper/2023/hash/d69103d7895f4e2083f24b664003d386-Abstract-Conference.html). How does the proposed approach compare with such methods?", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: I concur with the reviewers' positive view of the submission and think the paper would make a nice addition. Several points for revision and improvement were brought up in the discussion (thanks to all for participating actively) and I fully expect the authors to incorporate these into the final paper. A couple points were left unresponded that can also be addressed (such as discussion of completeness and oracle grouping).
Decision: Accept (poster)