Add rlearner_score (R-loss) for CATE model evaluation, EconML RScorer parity#935
Open
aman-coder03 wants to merge 3 commits into
Open
Add rlearner_score (R-loss) for CATE model evaluation, EconML RScorer parity#935aman-coder03 wants to merge 3 commits into
rlearner_score (R-loss) for CATE model evaluation, EconML RScorer parity#935aman-coder03 wants to merge 3 commits into
Conversation
jeongyoonlee
requested changes
Jul 8, 2026
jeongyoonlee
left a comment
Collaborator
There was a problem hiding this comment.
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.
10 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 EconMLRScorerparity. also factors the cross-fitting logic out ofBaseRLearner.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()incausalml/propensity.py, cross-fitted (stratified on treatment) outcome/treatment residuals(y - m_hat(X), w - e_hat(X)). lives inpropensity.pyrather thaninference.metaormetrics.cate_scoringto avoid a real circular import betweencausalml.metricsandcausalml.inference.meta,propensity.pyis the only module with zero dependency on either. hit an actualImportErrortracing this down when it was first placed incate_scoring.pyrlearner_score()incausalml/metrics/cate_scoring.py, scores fitted CATE models via R-loss =mean[(y_residual - w_residual * tau_hat(X))**2]. namedrlearner_scorerather thanr_scoreper review feedback on Add CATE model selection via surrogate metrics (R-score, DR-loss) #922 (too close torate_score). shares bootstrap CI logic withdr_score()/plug_in_t_score()via a new_bootstrap_loss_ci()helper. nop_valuecolumn, same reasoning as Add DR and plug-in T scoring metrics for CATE model evaluation #930BaseRLearner.fit()refactored to callcompute_r_residuals()instead of its own inlinecross_val_predictblock. scoped toBaseRLearner.fit()only,BaseRClassifier.fit()andXGBRRegressor.fit()keep their existing inline logic for nowtest_cate_scoring.py, mirroring the existing DR/plug-in-T test shapesworth flagging for review...
BaseRLearner.fit()'s outcome-model cross-fitting from plainKFoldtoStratifiedKFoldon 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_muestimatesE[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 samerandom_stateself.cv = KFold(...)fromBaseRLearner.fit(), dead code oncecompute_r_residuals()builds its own splitter internally, and nothing else in the class readsself.cvscoped out/follow-ups...
BaseRClassifier.fit()andXGBRRegressor.fit()still use inlinecross_val_predictrather thancompute_r_residuals()compute_r_residuals()doesX = np.asarray(X)internally, which is fine formetrics.cate_scoringcallers but meansBaseRLearner.fit()now coercesXto numpy for this one step, diverging from the polars/pandas-native contract the rest offit()follows (collect_if_lazy,filter_mask, etc.). not fixed here since it'd need either a native-X rewrite ofcompute_r_residuals()or branching logic, flagging rather than blocking on itdocs/methodology.rstanddocs/refs.bibupdates (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-upcloses the remaining scope from #922
Types of changes
What types of changes does your code introduce to CausalML?
Put an
xin the boxes that applyChecklist
Put an
xin 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.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.