Methods¶

This vignette summarizes the K-Fold Personalization Estimator.

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

Target¶

Given tuples \((X_i, A_i, Y_i)\) with known or estimated propensities \(p(A_i \mid X_i)\), the personalization effect is defined as

\[ \psi = \mathbb{E}\!\left[\,Y \mid \pi^\star(X)\,\right] - \max_{a}\,\mathbb{E}\!\left[\,Y \mid A=a\,\right], \]

where \(\pi^\star\) is the optimal personalized policy and the right-hand term is the expected outcome under the best single treatment applied to every individual.

Algorithm 2: per-shuffle cross-fit (kpe2 shared-nuisance variant)¶

Partition the \(n\) rows into \(K =\) n_folds folds. Choose \(m =\) n_nuisance_folds folds on which to fit the nuisances (default \(m = K - 1\)); the evaluation block size is \(b = K - m\) and must divide \(K\), so the folds tile into \(G = K / b\) non-overlapping evaluation blocks. For each block \(g\):

Fit all nuisances — the reward model \(m(x, a)\), the contextual policy \(\pi(x)\), the best single arm \(\pi_0\), and (when is_rct=False) the propensity \(\hat p(a \mid x)\) — on the \(m\) folds outside block \(g\).
Evaluate the AIPW influence function (Eq. S17) on the \(b\) folds of block \(g\):

\[ \psi_i = m(X, \pi(X)) - m(X, \pi_0) + \frac{\mathbf{1}\{A=\pi(X),\; \pi_0 \neq \pi(X)\} - \mathbf{1}\{A=\pi_0,\; \pi_0 \neq \pi(X)\}} {p(A \mid X)}\,\bigl(Y - m(X, A)\bigr). \]

Concatenating the per-block \(\psi\) vectors across all \(G\) blocks yields a length-\(n\) vector in which each sample is evaluated exactly once. The shuffle-level summary is \(\hat\psi_s = \text{mean}\) and \(\hat\sigma_s = \text{std}\).

Unlike the original Algorithm 2 — which fits the policy and reward models on disjoint subsets of the non-evaluation folds — kpe2 fits every nuisance on the same \(m\) folds. The default \(m = K - 1\) gives \(b = 1\) and \(G = K\): leave-one-fold-out, fitting every nuisance on \(K - 1\) folds and evaluating on the remaining one, rotated \(K\) times (with \(K = 6\), that is 5/6 of the data to fit and 1/6 to evaluate). When is_rct=True (the default), \(p(A \mid X)\) above is the known design propensity and nothing is fit for it; when is_rct=False, the fitted \(\hat p(a \mid x)\) is used in both the influence evaluation and the best-arm comparison.

Algorithm 1: aggregation across shuffles¶

Algorithm 2 is repeated for \(S\) random shuffles. The per-shuffle \((\hat\psi_s, \hat\sigma_s)\) pairs are aggregated via

\[ \hat\psi = \frac{1}{S}\sum_s \hat\psi_s, \qquad \widehat{\operatorname{Var}}(\hat\psi) = \frac{1}{S}\sum_s\!\left(\hat\sigma_s^2 + (\hat\psi_s - \hat\psi)^2\right)\!\Big/\,n. \]

The first term \(\hat\sigma_s^2\) accounts for within-shuffle influence variance; the second term \((\hat\psi_s - \hat\psi)^2\) accounts for across-shuffle variability introduced by sample splitting.

The one-sided upper-tail t-statistic is \(t = \hat\psi / \sqrt{\widehat{\operatorname{Var}}(\hat\psi)}\). The reported p_value is the corresponding upper-tail probability under \(t_{n-1}\). The reported confidence_interval is the 95% Wald interval \(\hat\psi \pm 1.96\,\sqrt{\widehat{\operatorname{Var}}(\hat\psi)}\).

Stability diagnostic¶

policy_stability is the fraction of samples for which every shuffle's contextual policy predicted the same action. Its range is \([0, 1]\). overall_stability is the analogous quantity for the best-arm policy.