JobCorps¶

This dataset is a semi-synthetic dataset built from the JobCorps dataset on job training (n = 10 214; binary treatment; outcome = week-4-year earnings rate, denominated in dollars per week).

See details here: https://datadryad.org/dataset/doi:10.5061/dryad.bg79cnpp7

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
Schochet, Peter Z., Burghardt, John, and McConnell, Sheena. Replication data for: Does Job Corps Work? Impact Findings from the National Job Corps Study. Nashville, TN: American Economic Association [publisher], 2008. Ann Arbor, MI: Inter-university Consortium for Political and Social Research [distributor], 2019-10-12. https://doi.org/10.3886/E113269V1

Prerequisites¶

Install the package from GitHub; all learner backends (scikit-learn for random_forest, econml for policy_tree) ship as dependencies — see Getting started for the full learner/backend table.

This vignette also requires the JobCorps CSV and the KPE_DATA_DIR environment variable pointing to the directory that contains it (see Preprocessing below).

This dataset can be accessed by cloning the repository and looking in the "data" directory.

git clone https://github.com/StanfordAI4HI/kpe-py.git

Preprocessing¶

The covariates are one-hot encoded and the age-group columns are retained as the policy_features subset passed to the contextual policy.

The code assumes that you have created a KPE_DATA_DIR env variable which points to the directory you are storing the provided data file: jobcorps_cleaned_nonan_full_with_offset.csv.

This dataset is a semi-synthetic dataset built from the JobCorps dataset on job training (n = 10 214; binary treatment; outcome = week-4-year earnings rate, denominated in dollars per week).

See details here: https://datadryad.org/dataset/doi:10.5061/dryad.bg79cnpp7

import os
import numpy as np
import pandas as pd
import sklearn
from sklearn.preprocessing import OneHotEncoder

df = pd.read_csv(os.path.join(os.environ["KPE_DATA_DIR"],
                              "jobcorps_cleaned_nonan_full_with_offset.csv"))
Y = df["EARNY4"].values
A = df["TREATMNT"].astype(int).values
feature_columns = [c for c in df.columns if c not in ("EARNY4", "TREATMNT")]
X_raw = df[feature_columns].values

encoded_df = pd.DataFrame()
age_groups = []
for i, col_name in enumerate(feature_columns):
    enc = OneHotEncoder(sparse_output=False)
    encoded_col = enc.fit_transform(X_raw[:, i].reshape(-1, 1))
    names = [f"{col_name}_{c}" for c in enc.categories_[0]]
    encoded_df = pd.concat(
        [encoded_df, pd.DataFrame(encoded_col, columns=names)], axis=1)
    if col_name == "AGEGROUP":
        age_groups = names

X = np.nan_to_num(encoded_df.values, nan=0.0)
age_group_index = [encoded_df.columns.get_loc(g) for g in age_groups]

unique, counts = np.unique(A, return_counts=True)
p = counts / len(A)
pa_x = np.array([dict(zip(unique, p))[a] for a in A], dtype=float)

Run¶

from kpe import kpe

result = kpe(
    X=X, A=A, Y=Y, propensity=pa_x,
    n_shuffles=100,
    n_folds=6,
    contextual_policy="dr_econml_2",
    best_arm="best_arm_erm",
    reward_model="random_forest",
    policy_features=age_group_index,
    n_jobs=5,
    random_state=0,
)
print(result.summary())

Results¶

Note that the resulting estimates are in dollars per week.

quantity	value
`psi` ($/week)	`+10.6706`
`std_error`	`2.7222`
95% confidence interval	`[+5.334, +16.007]`
one-sided p-value	`4.5e-05`
`policy_stability`	`0.231`
best-arm stability	`0.840`

The confidence interval excludes zero at the 5% level. policy_stability = 0.231 records the fraction of samples for which every shuffle predicted the same contextual-policy action.

Note that relative to the results reported in Li and Brunskill 2026 (see their Algorithms 1 and 2), this code base uses a two-fold shared-nuisance approach in which every nuisance is trained on 5/6 of the data. This improves the point estimate and the stability, but the sign and reject-at-0.05 decision are unchanged.

Baselines on the same inputs¶

from kpe import train_eval, papd

common = dict(
    X=X, A=A, Y=Y, propensity=pa_x,
    n_shuffles=100, n_folds=6,
    contextual_policy="dr_econml_2",
    best_arm="best_arm_erm",
    reward_model="random_forest",
    policy_features=age_group_index,
    n_jobs=5, random_state=0,
)
res_trev = train_eval(**common)
res_papd = papd(**common)

Both baselines return annualized psi and std_error on the same scale as kpe().

Notes¶

The JobCorps CSV stores earnings in dollars per week. Figure 3 multiplies the estimator's output by 52 for presentation; this vignette reports the annualized scale for direct comparison with the paper.
policy_features = age_group_index restricts the contextual policy to the AGEGROUP one-hot columns.