statspai.policy_learning

policy_learning

Policy Learning: Optimal treatment assignment from heterogeneous effects.

Learns an interpretable treatment assignment policy that maximises the expected welfare (value) of the population. Given estimated CATE, finds the optimal tree-based policy: "who should be treated?"

Components

• PolicyTree : Optimal depth-limited decision tree for treatment assignment (Athey & Wager 2021).
• policy_value : Evaluate the expected value of a treatment policy using doubly robust scores.

References

Athey, S. & Wager, S. (2021). Policy Learning with Observational Data. Econometrica, 89(1), 133-161. [@athey2021matrix]

Zhou, Z., Athey, S., & Wager, S. (2023). Offline Multi-Action Policy Learning: Generalization and Optimization. Operations Research, 71(1), 148-183. [@zhou2023offline]

PolicyTree

Optimal depth-limited policy tree.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str		required
treat	str		required
covariates	list of str		required
policy_covariates	list of str		None
max_depth	int		2
min_leaf_size	int		25
n_folds	int		5
alpha	float		0.05
random_state	int		42

Examples:

Fit the estimator directly, then route fresh covariates through the learned rule with :meth:predict:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> treat = rng.integers(0, 2, n)
>>> y = 1.0 + 2.0 * (x1 > 0) * treat + 0.5 * x2 + rng.normal(0, 1, n)
>>> df = pd.DataFrame({"y": y, "treat": treat, "x1": x1, "x2": x2})
>>> tree = sp.PolicyTree(data=df, y="y", treat="treat",
...                      covariates=["x1", "x2"], max_depth=2,
...                      min_leaf_size=30, n_folds=3, random_state=0)
>>> res = tree.fit()
>>> bool(res["value_gain_in_sample"] >= 0)  # by construction
True
>>> rec = tree.predict(np.array([[1.5, 0.0], [-1.5, 0.0]]))
>>> int(rec.shape[0])
2
>>> bool(set(int(v) for v in rec) <= {0, 1})
True

fit

fit() -> Dict[str, Any]

Learn the optimal policy tree.

predict

predict(X_new: ndarray) -> ndarray

Predict treatment assignment for new data.

Parameters:

Name	Type	Description	Default
X_new	ndarray(n, p)	Policy covariates for new observations.	required

Returns:

Type	Description
ndarray(n)	Binary treatment recommendations (0 or 1).

PolicyTreeResult

Bases: dict, ResultProtocolMixin

Result of :func:policy_tree.

Inherits from :class:dict so the legacy result['policy'] API keeps working (and isinstance(result, dict) is still True), while also exposing rich attribute access plus methods:

• :attr:value_policy_se — influence-function SE of the policy value, computed from the AIPW scores :math:\Gamma_i and the policy's assignments. With value_evaluation='cross_fit' (default) each unit is assigned by a tree grown without it, so the SE describes the reported value; with 'in_sample' it ignores the search and the value is optimistic.
• :meth:summary / :meth:plot_tree / :meth:to_latex / :meth:cite that match the Stata / R reporting idioms.
• :meth:to_excel for publication exports.

The tree attribute holds the fitted :class:PolicyTree instance so :meth:PolicyTree.predict is reachable downstream.

Examples:

Produced by :func:policy_tree; the legacy dict API still works alongside attribute access and the rich reporting methods:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> treat = rng.integers(0, 2, n)
>>> y = 1.0 + 2.0 * (x1 > 0) * treat + 0.5 * x2 + rng.normal(0, 1, n)
>>> df = pd.DataFrame({"y": y, "treat": treat, "x1": x1, "x2": x2})
>>> res = sp.policy_tree(df, y="y", treat="treat",
...                      covariates=["x1", "x2"], max_depth=2,
...                      min_leaf_size=30, n_folds=3, random_state=0)
>>> isinstance(res, sp.PolicyTreeResult)
True
>>> isinstance(res, dict)               # legacy result['policy'] still works
True
>>> res["n_obs"]
300
>>> bool(0.0 <= res.fraction_treated <= 1.0)
True
>>> "begin{table}" in res.to_latex()        # LaTeX table output
True

plot_tree

plot_tree(ax: Any = None, figsize: Tuple[float, float] = (8.0, 5.0), node_color: str = '#e8f0fe') -> Tuple[Any, Any]

Draw the policy tree as a labeled hierarchical diagram.

Each split node shows feature ≤ threshold; each leaf shows TREAT / DON'T TREAT plus the leaf value (mean AIPW score). Requires matplotlib.

to_latex

to_latex(caption: Optional[str] = None, label: str = 'tab:policy_tree') -> str

Render a publication-style summary table (LaTeX).

to_excel

to_excel(path: str, digits: int = 6) -> str

Write a single-sheet Excel summary.

Six decimals to match every other to_excel in the package: a spreadsheet is a data-interchange target that gets sorted, charted and recomputed, so it keeps numeric headroom rather than the display precision used by the presentation exports.

direct_method

direct_method(X: ndarray, A: ndarray, R: ndarray, pi_target: Any, n_actions: Optional[int] = None, alpha: float = 0.05, q_hat: Optional[ndarray] = None) -> OPEResult

Direct outcome regression (plug-in Q-model) OPE.

V = mean_i sum_a pi_e(a|X_i) Q(X_i, a). Q is q_hat when given (an (n, n_actions) matrix of reward-model predictions), otherwise a random forest on (X, one-hot A). A 1-D pi_target is a deterministic action vector. With the same Q this is obp's DirectMethod.

Examples:

>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 80
>>> X = rng.normal(size=(n, 3))
>>> A = rng.integers(0, 3, size=n)
>>> R = rng.normal(size=n)
>>> pi_target = rng.dirichlet(np.ones(3), size=n)
>>> res = sp.direct_method(X, A, R, pi_target)
>>> float(res.value)

doubly_robust

doubly_robust(X: ndarray, A: ndarray, R: ndarray, pi_target: Any, pi_behavior: Optional[ndarray] = None, n_actions: Optional[int] = None, clip: float = 50.0, alpha: float = 0.05, q_hat: Optional[ndarray] = None, n_folds: int = 1) -> OPEResult

Doubly-robust OPE (Dudik et al. 2011).

V = mean_i [ sum_a pi_e(a|X_i) Q(X_i, a) + w_i (R_i - Q(X_i, A_i)) ] with w_i = min(pi_e(A_i|X_i) / pi_b(A_i|X_i), clip). Q is q_hat when given, else a random forest. With n_folds > 1 the forest is cross-fitted, so that each observation's Q comes from a forest that did not see it; the default n_folds=1 fits and predicts on the same rows. A 1-D pi_target is a deterministic action vector. With the same Q and clip = lambda_ this is obp's DoublyRobust.

Notes

If pi_behavior is None and the internal behavior-policy logistic regression fails, propensities fall back to uniform 1/K; a ConvergenceWarning is emitted and diagnostics['propensity_fallback'] is set to True.

Examples:

>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 80
>>> X = rng.normal(size=(n, 3))
>>> A = rng.integers(0, 3, size=n)
>>> R = rng.normal(size=n)
>>> pi_target = rng.dirichlet(np.ones(3), size=n)
>>> res = sp.doubly_robust(X, A, R, pi_target)
>>> v = float(res.value)

ips

ips(X: ndarray, A: ndarray, R: ndarray, pi_target: Any, pi_behavior: Optional[ndarray] = None, clip: float = 50.0, alpha: float = 0.05) -> OPEResult

Inverse propensity score OPE.

V = mean_i w_i R_i with w_i = min(pi_e(A_i|X_i) / pi_b(A_i|X_i), clip); pi_target 1-D is pi_e(A_i|X_i), 2-D is (n, K) probabilities. clip is obp's lambda_ (np.inf for no cap). SE sd(w R)/sqrt(n).

.. versionchanged:: 1.30.0 clip caps the weight only; pi_b is no longer floored at 1/clip.

Notes

If pi_behavior is None and the internal behavior-policy logistic regression fails, propensities fall back to uniform 1/K; a ConvergenceWarning is emitted and diagnostics['propensity_fallback'] is set to True.

Examples:

>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 80
>>> X = rng.normal(size=(n, 3))
>>> A = rng.integers(0, 3, size=n)
>>> R = rng.normal(size=n)
>>> pi_target = rng.dirichlet(np.ones(3), size=n)
>>> res = sp.ips(X, A, R, pi_target)
>>> v = float(res.value)

snips

snips(X: ndarray, A: ndarray, R: ndarray, pi_target: Any, pi_behavior: Optional[ndarray] = None, clip: float = 50.0, alpha: float = 0.05) -> OPEResult

Self-normalised IPS (bias-reduction for large IS weights).

V = sum_i w_i R_i / sum_i w_i (obp's SelfNormalizedInverseProbabilityWeighting), weights as in :func:ips. SE is the delta-method sd(w (R - V)) / (mean(w) sqrt(n)).

.. versionchanged:: 1.30.0 The SE divided by sum(w) where mean(w) belongs, so it was too small by a factor of about sqrt(n) (0.00039 against a bootstrap 0.0174 at n = 2000). clip no longer floors pi_b.

Notes

If pi_behavior is None and the internal behavior-policy logistic regression fails, propensities fall back to uniform 1/K; a ConvergenceWarning is emitted and diagnostics['propensity_fallback'] is set to True.

Examples:

>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 80
>>> X = rng.normal(size=(n, 3))
>>> A = rng.integers(0, 3, size=n)
>>> R = rng.normal(size=n)
>>> pi_target = rng.dirichlet(np.ones(3), size=n)
>>> res = sp.snips(X, A, R, pi_target)
>>> v = float(res.value)

policy_value

policy_value(scores: ndarray, policy: ndarray) -> float

Evaluate the value of a treatment policy relative to treating no one.

Returns mean(scores * policy). With scores the doubly robust treatment-effect scores Gamma_1 - Gamma_0 (grf get_scores; policytree's reward matrix is [Gamma_0, Gamma_1]), this is mean(Gamma[i, pi_i]) - mean(Gamma[i, 0]): the AIPW estimate of E[Y(pi)] - E[Y(0)]. It equals mean(policy) times the AIPW ATE of the units the policy treats (grf average_treatment_effect(subset = policy == 1)). No standard error is returned; :func:policy_tree reports one.

Parameters:

Name	Type	Description	Default
scores	ndarray(n)	Doubly robust scores (AIPW pseudo-outcomes for treatment). Positive scores indicate the individual benefits from treatment. Must be one-dimensional.	required
policy	ndarray(n)	Policy recommendations (0 or 1; a probability of treatment gives the value of the stochastic policy). Same length as scores, or a scalar for a constant policy.	required

Returns:

Type	Description
float	Estimated value gain of the policy over treating no one.

Examples:

>>> import numpy as np
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> scores = rng.normal(0.3, 1.0, size=n)  # DR gains from treating

Treat-everyone vs an oracle policy (treat only positive-gain units):

>>> policy_all = np.ones(n, dtype=int)
>>> policy_oracle = (scores > 0).astype(int)
>>> round(float(sp.policy_value(scores, policy_all)), 2)
0.29
>>> round(float(sp.policy_value(scores, policy_oracle)), 2)
0.54

policy_targeting

policy_targeting(cate: Any, *, budget: Optional[int] = None, frac: Optional[float] = None, min_effect: float = 0.0) -> Dict[str, Any]

Rank-and-treat policy under a budget constraint.

Parameters:

Name	Type	Description	Default
cate	array-like, CausalResult, or fitted CATE model	Per-unit effect estimates — a raw array, a metalearner() / tarnet() result, or a fitted causal_forest() model (training-sample effects are used).	required
budget	int	Maximum number of units that can be treated. Mutually exclusive with frac.	None
frac	float	Maximum fraction of units that can be treated (in (0, 1]).	None
min_effect	float	Never treat a unit whose predicted effect is at or below this threshold, even with budget left over — treating predicted non-responders wastes budget and can do harm.	0.0

Returns:

Type	Description
dict	policy (0/1 array in input order), n_treated, threshold (smallest predicted effect among the treated), expected_gain under the policy, and the expected_gain_* baselines (treat_all, random at the same budget), plus a one-row summary DataFrame. Gains are sums of predicted effects — validate against sp.policy_value with doubly robust scores before deployment.

Examples:

>>> import numpy as np
>>> import statspai as sp
>>> tau = np.array([2.0, 1.0, 0.5, -0.5, -2.0])
>>> out = sp.policy_targeting(tau, budget=2)
>>> out["policy"].tolist()
[1, 1, 0, 0, 0]
>>> out["expected_gain"]
3.0
>>> out["expected_gain_treat_all"]
1.0
>>> # budget larger than the number of positive effects: the
>>> # min_effect guard stops at 3 treated units
>>> sp.policy_targeting(tau, budget=5)["n_treated"]
3