Skip to content
/ maq Public

Qini curves for multi-armed treatment rules

License

Notifications You must be signed in to change notification settings

grf-labs/maq

Folders and files

NameName
Last commit message
Last commit date

Latest commit

d62bdc3 · Apr 15, 2025
Nov 11, 2024
Jun 21, 2023
Apr 23, 2024
May 16, 2024
Apr 5, 2025
Apr 15, 2025
Oct 12, 2023
Jun 23, 2023
Apr 14, 2025
Jun 12, 2024
Apr 11, 2025
Nov 3, 2024
Apr 14, 2025
Nov 3, 2024

Repository files navigation

Multi-Armed Qini

CRANstatus Build Status

A package for policy evaluation using generalized Qini curves: Evaluate data-driven treatment targeting rules for one or more treatment arms over different budget constraints in experimental or observational settings under unconfoundedness.

Installation

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.

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 is 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 [ Y i ( k ) Y i ( 0 ) | X i ] for a given unit X i 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 ( X i ) ,   τ ( X i ) ] ,

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. Journal of Computational and Graphical Statistics, forthcoming. [paper, arxiv]