A package for evaluating multi-armed treatment rules via the Multi-Armed Qini - a generalization of the Qini curve to multiple costly treatment arms.
The latest release of the package can be installed through CRAN:
install.packages("maq")
The development version can be installed with:
devtools::install_github("grf-labs/maq", subdir = "r-package/maq")
(Installing from source requires a compiler that implements C++11 or later)
Python bindings are here.
library(maq) # Fit a CATE estimator on a training sample. n <- 3000 p <- 5 X <- matrix(runif(n * p), n, p) W <- as.factor(sample(c("0", "1", "2"), n, replace = TRUE)) Y <- X[, 1] + X[, 2] * (W == "1") + 1.5 * X[, 3] * (W == "2") + rnorm(n) train <- sample(1:n, n/2) tau.forest <- grf::multi_arm_causal_forest(X[train, ], Y[train], W[train]) # Predict CATEs on held out evaluation data. test <- -train tau.hat <- predict(tau.forest, X[test, ], drop = TRUE)$predictions # Assume costs equal a unit's pre-treatment covariate - the following are a toy example. cost <- cbind(X[test, 4] / 4, X[test, 5]) # Fit an evaluation forest to compute doubly robust scores on the test set. eval.forest <- grf::multi_arm_causal_forest(X[test, ], Y[test], W[test]) DR.scores <- grf::get_scores(eval.forest, drop = TRUE) # Fit a Qini curve on evaluation data, using 200 bootstrap replicates for confidence intervals. ma.qini <- maq(tau.hat, cost, DR.scores, R = 200) # Plot the Qini curve. plot(ma.qini) legend("topleft", c("All arms", "95% CI"), lty = c(1, 3)) # Get an estimate of gain at a given spend per unit along with standard errors. average_gain(ma.qini, spend = 0.2) # Get the treatment allocation matrix at a given spend per unit. pi.mat <- predict(ma.qini, spend = 0.2) # If the treatment randomization probabilities are known, then an alternative to # evaluation via AIPW scores is to use inverse-propensity weighting (IPW). W.hat <- rep(1/3, 3) IPW.scores <- get_ipw_scores(Y[test], W[test], W.hat) mq.ipw <- maq(tau.hat, cost, IPW.scores)
Let τ̂( ⋅ ) be an estimated CATE function where the k-th element measures the conditional average treatment effect E[Yi(k) − Yi(0)|Xi] for a given unit Xi for one of k = 1, …, K treatment arms, where k = 0 is a control arm. Let C( ⋅ ) be some known cost function that quantifies the cost of assigning a given treatment arm to the i-th unit. maq
delivers estimates of the Qini curve
Q(B) = E[⟨πB(Xi), τ(Xi)⟩],
which is the expected gain, at any budget constraint B, when assigning treatment using the policy πB that optimally selects (using the given functions τ̂( ⋅ ) and C( ⋅ )) which arm to assign to which unit such that the average incurred cost is less than or equal to B. The policy πB is a solution to a linear program: maq
computes a solution path for these treatment allocations over increasing budget levels B via an algorithm that leverages the multiple-choice knapsack structure of this problem. See the algorithm reference for an overview.
Erik Sverdrup, Han Wu, Susan Athey, and Stefan Wager. Qini Curves for Multi-Armed Treatment Rules. 2023. [arxiv]