Multi-Armed Qini

A package for evaluating multi-armed treatment rules via the Multi-Armed Qini - a generalization of the Qini curve to multiple costly treatment arms.

Installation

The latest release of the package can be installed through CRAN:

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.

Usage Example

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)

Details

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.

References

Erik Sverdrup, Han Wu, Susan Athey, and Stefan Wager. Qini Curves for Multi-Armed Treatment Rules. 2023. [arxiv]