Skip to content

Models API

uplift_bench.models.base

Abstract base for all uplift meta-learners.

The contract is intentionally tiny — fit and predict_uplift. Anything fancier (CV-based hyper-tuning, per-group calibration) is the job of a wrapper, not the base class.

Naming: I use X, t, y as parameters because these are the universal short names in the causal-inference literature. The lint rule that bans single-letter names is per-file-disabled in pyproject for the same reason.

UpliftModel

Bases: ABC

Common interface every meta-learner implements.

Subclasses store fitted state on self; we don't enforce a particular layout because each learner needs different things (S-learner: one model; T-learner: two; X-learner: 3+).

fit(X, t, y) abstractmethod

Train the model. Must mark self._fitted = True before returning.

predict_uplift(X) abstractmethod

Return per-row uplift estimate (same length as X).

uplift_bench.models.factory

Build an UpliftModel from a string name + kwargs.

Used by the Hydra config loader and the CLI. The registry is the single place to add a new meta-learner — every other module imports from here.

make_model(name, **kwargs)

Instantiate a meta-learner by short name.

uplift_bench.models.s_learner

S-learner.

Single model trained on (X, T) → Y. Uplift estimate is the difference of predictions when T is set to 1 vs 0:

tau_hat(X) = mu(X, T=1) - mu(X, T=0)

The simplest meta-learner. Tends to win on data where the treatment effect is small and the model is highly regularised — because in that regime the shared parameter pool acts as a useful prior. Loses badly when the treated and control groups have very different X distributions.

uplift_bench.models.t_learner

T-learner.

Two separate models — one trained on treated rows, one on control rows. Uplift = mu_1(X) - mu_0(X). The opposite trade-off to S-learner: never shares signal across arms, so it overfits when one arm is small but captures arm-specific structure perfectly when both arms are large.

uplift_bench.models.x_learner

X-learner (Künzel, Sekhon, Bickel, Yu 2019).

Two-stage:

Stage 1 — fit mu_0 on control, mu_1 on treated (same as T-learner). Stage 2 — impute counterfactual differences: D_treated = Y_1 - mu_0(X_1) D_control = mu_1(X_0) - Y_0 then fit tau_0 on D_control (X_0) and tau_1 on D_treated (X_1). Combine via propensity weights: tau(X) = e(X) * tau_0(X) + (1 - e(X)) * tau_1(X)

The propensity weights e(X) come from a separate model. We clip them to [0.05, 0.95] — if you don't, X-learner explodes on small datasets where propensity gets close to 0 or 1. Painful lesson learned on RetailHero.

uplift_bench.models.r_learner

R-learner (Nie & Wager 2021).

Fits propensity e(X) and outcome m(X) = E[Y|X] in stage 1 (with cross-fitting to avoid overfitting bias), then minimises the R-loss in stage 2:

L(tau) = sum_i ( (Y_i - m(X_i)) - (T_i - e(X_i)) * tau(X_i) )^2

Equivalent to a weighted regression of the outcome residual on the treatment residual times the candidate tau. We express this as:

target_i = (Y_i - m_hat(X_i)) / (T_i - e_hat(X_i))
weight_i = (T_i - e_hat(X_i))^2

then fit a regressor on (X, target) with sample_weight=weight.

Implementation note: R-learner is the meta-learner where cross-fitting matters most. We use 5-fold by default; less and stage-1 leakage shows up as overoptimistic Qini on training folds.

RLearner

Bases: UpliftModel

uplift_bench.models.dr_learner

DR-learner — Doubly Robust meta-learner (Kennedy 2023).

Same skeleton as R-learner (cross-fit nuisances) but the stage-2 target is the doubly-robust pseudo-outcome:

psi_i = mu_1(X_i) - mu_0(X_i)
      + (T_i / e(X_i))     * (Y_i - mu_1(X_i))
      - ((1-T_i)/(1-e(X_i))) * (Y_i - mu_0(X_i))

This is unbiased for tau(X) when either the outcome model or the propensity model is correctly specified — hence "doubly robust". On the benchmark this is the meta-learner I expect to win on Criteo with a flexible base learner like CatBoost; the propensity scores from CatBoost are usually well-calibrated enough that even when the outcome model is slightly off, the IPW correction recovers the bias.

uplift_bench.models.class_transformation

Class Variable Transformation (Jaskowski & Jaroszewicz 2012).

Trick: under randomised treatment assignment with P(T=1) = 0.5, define

Z = T * Y + (1 - T) * (1 - Y)

i.e. Z = 1 iff (treated and responded) OR (not treated and not responded). Then 2 * P(Z=1|X) - 1 = E[Y|T=1, X] - E[Y|T=0, X] = tau(X).

So a single classifier on (X, Z) recovers an unbiased uplift estimator — without ever fitting two outcome models.

Important assumption: this learner is correct only under randomised treatment with marginal propensity P(T=1) ≈ const (i.e. an RCT, like Hillstrom, Criteo Uplift, RetailHero, MegaFon, all of which are randomised). The marginal-propensity reweighting we apply (1/p_t for treated, 1/(1-p_t) for control) compensates for unbalanced marginal share but does not correct for covariate-conditional propensity e(X). On observational data with non-random treatment assignment use DR-learner or X-learner instead. We refuse to fit when the marginal propensity is outside [0.05, 0.95] as a coarse RCT check, but this does not catch heterogeneous propensity.

Only defined for binary outcomes; calling with continuous outcome raises.

References
  • Jaskowski & Jaroszewicz 2012, "Uplift modeling for clinical trial data" ICML Workshop on Clinical Data Analysis.
  • scikit-uplift ClassTransformation documentation — https://www.uplift-modeling.com/en/latest/api/models/ClassTransformation.html

uplift_bench.models.causal_forest

Causal forest wrapper around econml's CausalForestDML.

We wrap rather than re-export because (a) econml's API is opinionated (treatment must be float, predictions return shape (n, 1, 1)…), and (b) we want a uniform predict_uplift(X) -> 1-D ndarray contract across all seven meta-learners.

Reference: Athey & Wager (2019), "Estimation and Inference of Heterogeneous Treatment Effects Using Random Forests" — JASA.