Baselines

In addition to kpe(), the package provides two baseline estimators of the personalization effect: train_eval() and papd().

The examples below run on the bundled toy dataset (see load_kpe_toy), the same one used in Getting started:

from kpe.data import load_kpe_toy

d = load_kpe_toy()
X, A, Y, propensity = d["X"], d["A"], d["Y"], d["propensity"]

train_eval()

from kpe import train_eval

res = train_eval(
    X, A, Y, propensity,
    n_shuffles=100,
    contextual_policy="policy_tree",
    best_arm="ips",
    reward_model="random_forest",
    random_state=0,
)

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

Output. A KPEResult object. policy_stability and overall_stability are 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()

from kpe import papd

res = papd(
    X, A, Y, propensity,
    n_shuffles=100,
    contextual_policy="policy_tree",
    best_arm="ips",
    random_state=0,
)

Inputs. Same 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.