kpe-py — K-Fold Personalization Estimator¶

kpe-py is the Python implementation; the R version is kpe-r. The installed and imported module is kpe (from kpe import kpe) — "kpe-py" is the project/repository name.

Released under the MIT License — provided "as is", without warranty of any kind; the authors are not liable for any claim or damages.

A Python package for estimating the personalization effect — the difference in expected outcome between the best personalized policy and the best single treatment applied to every individual — from observational or experimental data. The package implements a 2 fold variant of the K-Fold Personalization Estimator of Li and Brunskill (2026) together with two baseline estimators (TrainEval and PAPD) under a common argument signature and return type.

References:¶

Zhaoqi Li, Emma Brunskill, A statistical test for the benefits of personalizing interventions. Science 393, eaeb9506 (2026). DOI: 10.1126/science.aeb9506
2-fold variant detailed here

Install¶

Installed from GitHub — not yet on PyPI (planned):

pip install "git+https://github.com/StanfordAI4HI/kpe-py.git"

Example usage¶

from kpe import kpe
from kpe.data import load_kpe_toy

d = load_kpe_toy()
res = kpe(d["X"], d["A"], d["Y"], d["propensity"],
          n_shuffles=100, n_folds=6, random_state=0)
print(res.summary())

See Getting started for a short example on the bundled kpe_toy dataset, and the JobCorps and Joke vignettes for dataset-specific analyses.

Estimators¶

All three estimators share the same argument signature and return a KPEResult dataclass.

`kpe()`¶

Inputs. X is an (n, d) covariate matrix. A is a length-n integer vector of observed actions in {0, ..., K-1}. Y is a length-n numeric vector of observed outcomes. propensity is a scalar, length-n vector, or (n, K) matrix giving p(A | X) for the observed action; it is required when is_rct=True (the default) and ignored when is_rct=False (the propensity is then fit on the nuisance folds). Additional keyword arguments include n_shuffles, n_folds, n_nuisance_folds, is_rct, contextual_policy, best_arm, reward_model, propensity_model, policy_features, n_jobs, random_state, and verbose.

Output. A KPEResult object.

Computation. The K-Fold Personalization Estimator of Li and Brunskill (2026), in the two-fold shared-nuisance variant: for each of n_shuffles random partitions of the data into n_folds folds, all nuisances — the reward model, contextual policy, best-arm learner, and (when is_rct=False) the propensity — are fit on a shared set of n_nuisance_folds folds, and the AIPW influence function of Eq. S17 is evaluated on the held-out block of the remaining n_folds - n_nuisance_folds folds. The blocks tile the folds without overlap, so each sample is evaluated exactly once per shuffle. The default n_nuisance_folds = n_folds - 1 gives leave-one-fold-out: with n_folds = 6, every nuisance is fit on 5/6 of the data and evaluated on the held-out 1/6, rotated six times. The Algorithm-1 variance decomposition combines the per-shuffle (psi_hat, sigma_hat) pairs into a point estimate, standard error, one-sided upper-tail t-test p-value, and 95% Wald confidence interval.

`train_eval()`¶

Inputs. Same signature as kpe(). The n_folds argument is accepted for symmetry and is ignored.

Output. A KPEResult object. policy_stability and overall_stability are returned as NaN because each shuffle evaluates a different randomly-chosen held-out half, so per-sample policy agreement across shuffles is not well-defined.

Computation. On each of n_shuffles random 50/50 splits of the data, the reward, contextual-policy, and best-arm learners are estimated on one half and the AIPW influence function of Eq. S17 is evaluated on the other half. The Algorithm-1 variance decomposition is applied to the per-shuffle (psi_hat, sigma_hat) pairs using n / 2 as the effective sample size for the t-test denominator.

`papd()`¶

Inputs. Same signature as kpe(). The n_folds and reward_model arguments are accepted for symmetry and are ignored; papd() always uses 3 folds and does not fit a reward model.

Output. A KPEResult object.

Computation. Implementation of the Population Average Prescriptive Difference of Imai and Li (2023), specialized to the case where one policy class is restricted to providing a single best arm for all individuals, as described in Li and Brunskill (2026, Eq. S18). On each of n_shuffles random 3-fold partitions, the contextual policy and best-arm learner are trained on two folds and the test-fold-centered inverse-propensity statistic

\[ \psi_i = \frac{\mathbf{1}\{A_i = \pi(X_i),\; \pi_0 \neq \pi(X_i)\} - \mathbf{1}\{A_i = \pi_0,\; \pi_0 \neq \pi(X_i)\}} {p(A_i \mid X_i)}\,\bigl(Y_i - \bar Y_{\text{test}}\bigr) \]

is computed on the held-out fold. The Algorithm-1 variance decomposition is applied to the per-shuffle (psi_hat, sigma_hat) pairs.

Sibling R package¶

An R implementation with matching argument names and return fields is available at StanfordAI4HI/kpe-r (documentation: https://stanfordai4hi.github.io/kpe-r/).

Citation¶

Li, Z. and Brunskill, E. (2026). A Statistical Test for the Benefits of Personalizing Interventions.

License¶

kpe-py is released under the MIT License — provided "as is", without warranty of any kind; the authors are not liable for any claim or damages.