Doubly Robust Fusion of Many Treatments for Policy Learning

We sincerely appreciate your insightful and constructive comments. Below, we have provided our detailed, point-by-point responses.

1. Behavior under assumption violation

Our 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.

In 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

Remark 3.10, which gives the gap tolerance condition; please also see our response to Reviewer UwD5, point 2, for additional context.

2. Additional experiments and details

We 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.

We will move the experimental setup tables to the Appendix.

• Results are averaged over 200 runs; variance is reported in the revised Table S5 below.
• Policy tree (Zhou et al., 2023) is applied with default tuning.
• Fused lasso uses extended BIC (line 159) for model selection, a standard criterion for recovering sparse structure.
• 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.

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.

---

3. Related work on complex treatments

We will incorporate the recommended references into the revised manuscript to better position our work within the broader literature.

4. Presentation

We revised Section 3.2 to begin with a summary of the key results without technical details, followed by the necessary mathematical intuition and formalism.

5. Clarification on DGP assumptions

Our 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.

On the second bullet point under weaknesses, the sentence after "Second–" appears incomplete. Could you kindly clarify the intended point?

6. Motivation for calibration weighting and fused lasso over end-to-end approaches

We 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.

While 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.

7. Comparison with overlap-based approaches for continuous treatments

Thank 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.

In 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.

We sincerely appreciate your insightful and constructive comments. Below, we have provided our detailed, point-by-point responses.

1. Necessity of recovering group structure

When 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.

By 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.

We 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.

2. Approximate equality of $\mu_a(X)$ and gap tolerance

We 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.

In 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.

We explicitly provide the gap tolerance required for recovery in Remark 3.10. For reference, the condition is:

$$\min\_{b\neq b^\prime} \| \boldsymbol{\beta}\_b^* -\boldsymbol{\beta}\_{b^\prime}^* \|\_{\infty} / c > \lambda\_n \gg \sqrt{n\log(n)}/K,$$

where $\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.

3. When is group recovery necessary for policy learning?

We 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.

However, 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.

We sincerely appreciate your insightful and constructive comments. Below, we have provided our detailed, point-by-point responses.

1. Simulations for validation of the double robustness

We 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$.

Since 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.

Note 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.

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.

---

2. Real-world experiments

The 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.

To 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.

3. Simulation with more treatments

We 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.

4. Completeness assumption

To 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.

5. Presentation

We 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.

While $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.

We will revise Assumption 3.2 based on your suggestion.

6. Assumption 2.1 (Identification)

We replaced “SUTVA” with the consistency and removed $P(X = x) > 0$.

7. Title terminology

We revised the title and text to use “many treatments” instead of “high-dimensional treatments,” following conventions in the literature (Ma et al., 2022, 2023).

8. Assumption 3.1 (Convergence of calibration weight)

For every treatment $a\in\mathcal{A}$, calibration weighting is an optimization problem and can be solved using the method of Lagrange multipliers:

$$L\_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).$$

Minimizing $L\_a(w\_1,\ldots,w\_n)$ gives:

$$\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})]},$$

where the function $\rho\_{\gamma}(x)$ for different $\gamma$ values are summarized in Table S4 below, and $\widehat{\lambda}$ solves the equation

\sum\_{i:A\_i=a}\rho\_{\gamma} $ \lambda^\top(X\_i-\bar{X}) $ (X\_i-\bar{X})=0.

Therefore, $\hat{\lambda}$ is an M-estimator and, under standard regularity conditions for M-estimators (Boos and Stefanski, 2013), it is root-n consistent.

Table S4: $\rho\_{\gamma}(x)$ for Cressie-Read family.

---

9. Bounded assumption

The bounded outcome and covariates assumption is only required for the results in Section 3.3 and not for those in Section 3.2.

10. Generality of results

Propositions 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.

11. Related literature

We will incorporate the recommended references into the revised manuscript to better situate our work within the broader literature.

We sincerely appreciate your insightful and constructive comments. Below, we have provided our detailed, point-by-point responses.

1. Experiments for increasing $K$ and fixed $n$

We 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).

---

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.

---

Table S2: Simulation results for $K=48$.

---

2. Convergence rate of the weights

In 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.

In 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.

3. Synthetic combinations

We 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).

While 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.

4. Arms with few or even zero units

If 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.

5. Iterative learning of weights for fused treatments

It 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.

In 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.

In the policy learning phase, we indeed re-estimate the weights for the fused treatments when learning the individualized treatment rule.