RFC: CausalML v1.0 roadmap

[jeongyoonlee]

CausalML (0.17.0) is mid-transition to sklearn-idiomatic estimators (BaseEstimator meta-learners, save/load, native polars/pandas). We'd like to define the road to v1.0 and get community input before locking it in.

What v1.0 is: a stabilize + expand release positioning CausalML as the production-scale library for uplift / CATE estimation. It will ship through intermediate milestones (each released as a 0.x minor); 1.0 GA is the API freeze at the end — not a single big-bang.

Where we see CausalML in the ecosystem: strongest on uplift + policy optimization (Qini/AUUC/RATE, multi-treatment, value optimization) with a clean meta-learner API and SHAP interpretability. Current gaps include: no distributed/out-of-core compute, no standard benchmark-dataset loaders or PEHE/policy-risk metrics, a fragmented neural-net story, and heavy serial bootstrap-CI paths.

M1 — Foundation & stability

• sklearn conformance (check_estimator for all learners); move black out of runtime deps
• serialization hygiene + a pickle-size regression guard in CI
• slim the bootstrap-ensemble memory footprint
• decision on fit() argument order (CausalML is not compatible with Scikit-Learn Pipeline #854) — the one intentional breaking change, with a deprecation shim
• unify the tree base: port UpliftTreeClassifier / UpliftRandomForestClassifier onto the shared _tree/ Cython kernel that CausalTreeRegressor / CausalRandomForestRegressor already use — introduce an UpliftClassificationCriterion alongside the existing causal criteria, delete the duplicated node / builder / splitter / bagging code in uplift.pyx (~2.3 KLoC), and let both families inherit BaseDecisionTree. Prerequisite for the check_estimator and serialization-hygiene bullets above; public class names and fit/predict signatures preserved via thin subclasses, with deprecation warnings for any drift.

M2 — Scale-out (production headline)

• profile + optimize the estimators' own hot paths (uplift/causal-tree Cython, per-group meta-learner fitting)
• parallelize bootstrap-CI paths and respect the active joblib backend, so Ray works as a drop-in via register_ray() (optional causalml[ray] extra)
• out-of-core batch scoring of fitted learners on large populations (Ray Data) — the real large-scale win
• GPU passthrough for boosting nuisances and NN batching

M3 — Benchmarks & evaluation

• sklearn-style fetch_* loaders (cached) for standard datasets: IHDP, Lalonde, Twins, Jobs, ACIC, Criteo-uplift, Hillstrom (subject to licensing)
• ground-truth metrics: PEHE / √PEHE, ε_ATE, policy risk — complementing AUUC/Qini/RATE
• a reproducible estimator × dataset leaderboard in the docs

M4 — Modern NN & uncertainty

• consolidate the TF/PyTorch/JAX DragonNet backends onto one primary backend with friendly optional-import errors
• conformal ITE prediction intervals over any learner
• a unified uncertainty API reconciling analytic, bootstrap, and conformal intervals
• SOTA architectures: TARNet, the CATENets family (SNet/FlexTENet/OffsetNet), deep R/DR learners

1.0 GA

API freeze, deprecation policy + semver commitment, migration guide, and a "choosing an estimator" guide.

Explicit non-goals for v1.0

• survival/censored-outcome HTE (low demand)
• continuous/dose-response treatments
• head-to-head DML parity with EconML (we prefer DoWhy interop)
• in-context causal foundation models (experimental, post-1.0)

Open questions we'd love feedback on

1. Primary neural-net backend — PyTorch or JAX?
2. fit() argument order (CausalML is not compatible with Scikit-Learn Pipeline #854): keep fit(X, treatment, y) or move to (X, y, treatment) for sklearn alignment — how painful is the break for you?
3. Which benchmark datasets and metrics matter most to your work?
4. Is Ray the right distributed backend, or would you prefer Dask/Spark?
5. Biggest scale/performance pain points with CausalML today?

Feedback, disagreement, and additional use cases all welcome. If there's appetite, we'll convert agreed items into tracked issues per milestone.

[jeongyoonlee]

Adding one item to M1 that I missed in the first draft: unifying the underlying tree base for the two tree families.

Today CausalTreeRegressor / CausalRandomForestRegressor already sit on the shared causalml/inference/tree/_tree/ Cython kernel (Tree, TreeBuilder, Splitter, RegressionCriterion) that was vendored from sklearn 1.5.2 in #806. UpliftTreeClassifier / UpliftRandomForestClassifier don't — they live in a single 2,742-line uplift.pyx with their own Python-object node struct, their own recursive builder, their own bagging forest, and their own KL / ED / Chi² / DDP / IT / CIT criteria. They also don't inherit from BaseEstimator, which is why save / load / _is_fitted are monkey-patched in tree/__init__.py and why check_estimator conformance (M1) is currently blocked for these classes.

Scope of the proposed refactor:

• extract BaseUpliftDecisionTree(BaseDecisionTree) analogous to today's BaseCausalDecisionTree
• add UpliftClassificationCriterion alongside CausalRegressionCriterion, keeping the existing impurity math (KL, ED, Chi², DDP, IT, CIT) as pluggable criterion classes (~440 LoC preserved)
• port the builder / node / splitter usage to the shared kernel; delete the duplicated ~2.3 KLoC from uplift.pyx
• switch UpliftRandomForestClassifier from its own bagging loop to sklearn's ForestClassifier
• preserve public class names, arg names, and defaults; add deprecation warnings for any inevitable drift (e.g. the monkey-patched save / load becoming real methods)

Why M1 and not M2: it's a prerequisite for two M1 bullets already in the RFC — check_estimator conformance and serialization hygiene — and it changes the profiling surface for M2's "uplift/causal-tree Cython hot paths" bullet (one surface instead of two). Doing it after 1.0 GA becomes a breaking-change hazard.

Explicit non-goals: no change to which uplift-splitting criteria are supported; no honesty-semantics change unless separately RFC'd; no new estimator types.

Open questions:

1. Ship in one PR, or one PR per criterion family (KL / ED / Chi² / …)?
2. Deprecation window for the monkey-patched save / load semantics — one minor, or fold it into the M1 serialization-hygiene work?
3. Anyone relying on the current DecisionTree Python class as a public API surface (vs. only using it via UpliftTreeClassifier)? If so, we need a compatibility shim.