Skip to content

Add rlearner_score (R-loss) for CATE model evaluation, EconML RScorer parity#935

Open
aman-coder03 wants to merge 3 commits into
uber:masterfrom
aman-coder03:feature/r-loss-cate-scoring
Open

Add rlearner_score (R-loss) for CATE model evaluation, EconML RScorer parity#935
aman-coder03 wants to merge 3 commits into
uber:masterfrom
aman-coder03:feature/r-loss-cate-scoring

Conversation

@aman-coder03

Copy link
Copy Markdown
Contributor

Proposed changes

2nd PR following #930 per @jeongyoonlee's scoping notes on #922
Adds rlearner_score() as a standalone R-loss (Nie & Wager, 2021) scoring function for CATE model evaluation, giving EconML RScorer parity. also factors the cross-fitting logic out of BaseRLearner.fit() so the R-learner's own training loop and the new standalone scorer share the same nuisance-fitting code instead of duplicating it.

what's included,

  • compute_r_residuals() in causalml/propensity.py, cross-fitted (stratified on treatment) outcome/treatment residuals (y - m_hat(X), w - e_hat(X)). lives in propensity.py rather than inference.meta or metrics.cate_scoring to avoid a real circular import between causalml.metrics and causalml.inference.meta, propensity.py is the only module with zero dependency on either. hit an actual ImportError tracing this down when it was first placed in cate_scoring.py
  • rlearner_score() in causalml/metrics/cate_scoring.py, scores fitted CATE models via R-loss = mean[(y_residual - w_residual * tau_hat(X))**2]. named rlearner_score rather than r_score per review feedback on Add CATE model selection via surrogate metrics (R-score, DR-loss) #922 (too close to rate_score). shares bootstrap CI logic with dr_score()/plug_in_t_score() via a new _bootstrap_loss_ci() helper. no p_value column, same reasoning as Add DR and plug-in T scoring metrics for CATE model evaluation #930
  • BaseRLearner.fit() refactored to call compute_r_residuals() instead of its own inline cross_val_predict block. scoped to BaseRLearner.fit() only, BaseRClassifier.fit() and XGBRRegressor.fit() keep their existing inline logic for now
  • 10 new tests in test_cate_scoring.py, mirroring the existing DR/plug-in-T test shapes

worth flagging for review...

  • switches BaseRLearner.fit()'s outcome-model cross-fitting from plain KFold to StratifiedKFold on treatment, for consistency with the propensity/DR nuisance splitting elsewhere (and the Add DR and plug-in T scoring metrics for CATE model evaluation #930 rare-treatment fix). model_mu estimates E[Y|X] unconditionally so it was never exposed to the empty-arm crash itself, this is a consistency change, not a bugfix, but fold assignment for existing R-learner fits will differ run to run vs. before given the same random_state
  • removed self.cv = KFold(...) from BaseRLearner.fit(), dead code once compute_r_residuals() builds its own splitter internally, and nothing else in the class reads self.cv

scoped out/follow-ups...

  • BaseRClassifier.fit() and XGBRRegressor.fit() still use inline cross_val_predict rather than compute_r_residuals()
  • compute_r_residuals() does X = np.asarray(X) internally, which is fine for metrics.cate_scoring callers but means BaseRLearner.fit() now coerces X to numpy for this one step, diverging from the polars/pandas-native contract the rest of fit() follows (collect_if_lazy, filter_mask, etc.). not fixed here since it'd need either a native-X rewrite of compute_r_residuals() or branching logic, flagging rather than blocking on it
  • docs/methodology.rst and docs/refs.bib updates (extending the "CATE Model Evaluation" section from Add DR and plug-in T scoring metrics for CATE model evaluation #930, adding a citation for the R-learner paper) aren't in this PR yet, will add here or as a doc-only follow-up

closes the remaining scope from #922

Types of changes

What types of changes does your code introduce to CausalML?
Put an x in the boxes that apply

  • Bugfix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Documentation Update (if none of the other choices apply)

Checklist

Put an x in the boxes that apply. You can also fill these out after creating the PR. If you're unsure about any of them, don't hesitate to ask. We're here to help! This is simply a reminder of what we are going to look for before merging your code.

  • I have read the CONTRIBUTING doc
  • I have signed the CLA
  • Lint and unit tests pass locally with my changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have added necessary documentation (if appropriate)
  • Any dependent changes have been merged and published in downstream modules

Further comments

If this is a relatively large or complex change, kick off the discussion by explaining why you chose the solution you did and what alternatives you considered, etc. This PR template is adopted from appium.

@jeongyoonlee jeongyoonlee left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

BaseRLearner.fit() routes the outcome cross-fit through compute_r_residuals() but drops w_residual (_) and passes neither p nor propensity_learner — so compute_r_residuals cross-fits a full propensity model (default ElasticNet, per fold) on every fit and discards it. fit() already computed p above; this is wasted work: fit() is ~4× slower even when the caller supplies p (0.8s → 3.3s at n=8k, p=20). Add an outcome-only path to compute_r_residuals (skip propensity when w_residual isn't needed) and call that here.

@aman-coder03 aman-coder03 requested a review from jeongyoonlee July 9, 2026 09:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants