Joke¶

The Jester joke-ratings dataset is encoded as a contextual bandit with 10 actions (joke IDs), 100-dimensional random-projection contexts, and a rating outcome normalized to [0, 5.0].

Data¶

The code assumes that you have created a KPE_DATA_DIR env variable which points to the directory you are storing the RCT_joke_data.csv` (n = 48 445).

The RCT_joke_data.csv is a preprocessed RCT form of the Jester dataset [https://goldberg.berkeley.edu/jester-data/]. See details here: https://datadryad.org/dataset/doi:10.5061/dryad.bg79cnpp7

Prerequisites¶

Install the package from GitHub; all learner backends (scikit-learn for lasso, the per-arm linear policy) ship as dependencies — see Getting started for the full learner/backend table. This vignette also requires RCT_joke_data.csv and the KPE_DATA_DIR environment variable pointing to the directory that contains it (see Data 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

import os
import numpy as np
import pandas as pd

# Processed Jester arrays (columns X0..X99, A, Y), exported once from
# RCT_joke_data.pkl to a CSV 
df = pd.read_csv(os.path.join(os.environ["KPE_DATA_DIR"], "RCT_joke_data.csv"))
X = df[[f"X{i}" for i in range(100)]].to_numpy(dtype=float)
A = df["A"].to_numpy(dtype=int)
Y = df["Y"].to_numpy(dtype=float)
n = X.shape[0]

unique, counts = np.unique(A, return_counts=True)
p = counts / n
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="linear",
    best_arm="best_arm_erm",
    reward_model="lasso",
    n_jobs=5,
    random_state=12314,
)
print(result.summary())

The research code uses a custom learner named joker_lasso_flat; this package implements a per-arm linear regression (contextual_policy = "linear") combined with a Lasso reward model (reward_model = "lasso") as the closest in-package equivalent. The joker_* learners are not exported by this package.

Results¶

Running the above call on the authoritative pickle under the default leave-one-fold-out cross-fitting (n_nuisance_folds = n_folds - 1, so 5/6 of the data fits every nuisance) yields:

quantity	value
`psi`	`0.1676`
`std_error`	`0.0224`
95% confidence interval	`[0.124, 0.211]`
z-statistic	`7.50`
one-sided p-value	`3.4e-14`
`policy_stability`	`0.274`

For reference, this two-fold shared-nuisance scheme used here slightly improves over the original results report in Li and Brunskill (Science, 2026). The sign of psi and the reject-at-0.05 decision are unchanged.

Multi-arm learner choice¶

The Joke dataset has 10 actions. Among the built-in contextual-policy strings:

"linear" fits a separate OLS regression per arm and returns argmax predictions; supports any number of actions.
"policy_tree" (alias "dr_econml_2") wraps econml.policy.DRPolicyTree; supports any number of actions.
"policy_forest" (alias "dr_econml") wraps econml.policy.DRPolicyForest; supports any number of actions.

Notes¶

Joke ratings are normalized to the range shown above prior to this call; psi is on that normalized outcome scale.