Merge MultiModelWrapper and SeparateModel into a single dispatcher#1039
Open
kbattocchi wants to merge 2 commits into
Open
Merge MultiModelWrapper and SeparateModel into a single dispatcher#1039kbattocchi wants to merge 2 commits into
kbattocchi wants to merge 2 commits into
Conversation
MultiModelWrapper is rewritten to dispatch rows to per-category models based
on the treatment block trailing the feature matrix. It now supports three
encodings via the new `encoding` keyword:
- `'drop_first'` (default): the last `K - 1` columns are a drop-first
one-hot encoding; the all-zero row is the control category. This matches
the encoding used throughout EconML's discrete-treatment estimators
(e.g. via `OneHotEncoder(drop='first')` in DRLearner), so
MultiModelWrapper can now actually be used as a first-stage model for
them.
- `'full'`: the last `K` columns are a full one-hot encoding (column 0
is the control). This matches what DROrthoForest's nuisance estimator
feeds to `model_Y` after re-adding the dropped column.
- `'label'`: the last column is an integer category index in
`{0, ..., K - 1}`. This makes MultiModelWrapper a strict superset of
SeparateModel.
A single positional model combined with `n_categories=K` clones the model
K times, which is convenient when all arms share a hyperparameter search.
The previous MultiModelWrapper implementation was untested, not exported
anywhere else in the package, advertised by DROrthoForest's docstrings,
and could not actually work with DROrthoForest because it did not account
for the extra one-hot column that nuisance_estimator_generator stitches on.
The DROrthoForest docstrings are updated to recommend the
`encoding='full'` variant explicitly, and a TODO is added to
nuisance_estimator_generator noting that we could later harmonize with the
rest of EconML by feeding the drop-first encoding directly.
SeparateModel is now deprecated via the existing `@deprecated` decorator;
`MultiModelWrapper(..., encoding='label')` is a drop-in replacement. The
existing test in test_model_selection.py that exercised SeparateModel is
migrated to MultiModelWrapper.
Tests cover all three encodings, validation paths, sample_weight
forwarding, the single-model + `n_categories` form, an end-to-end
LinearDRLearner integration with 3 treatment categories (which the old
wrapper could not handle), and the SeparateModel deprecation warning.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Keith Battocchi <kebatt@microsoft.com>
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR modernizes MultiModelWrapper to support multiple discrete-treatment encodings, improves its documentation, updates dependent references/tests, and deprecates SeparateModel in favor of the enhanced wrapper.
Changes:
- Redesign
MultiModelWrapperto supportdrop_first,full, andlabeltreatment encodings; add stronger validation and cloning behavior. - Deprecate
SeparateModelwith guidance to migrate toMultiModelWrapper(..., encoding='label'). - Update ORF docs and add/adjust tests to cover new wrapper behavior and integrations.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 7 comments.
| File | Description |
|---|---|
| econml/utilities.py | Reworks MultiModelWrapper API/behavior; deprecates and reintroduces SeparateModel with warnings. |
| econml/tests/test_utilities.py | Adds comprehensive unit tests for MultiModelWrapper encodings/validation and deprecation warning test. |
| econml/tests/test_model_selection.py | Migrates model selection tests from SeparateModel to MultiModelWrapper. |
| econml/orf/_ortho_forest.py | Clarifies that ORF uses full one-hot and documents MultiModelWrapper(…, encoding='full'). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…odelWrapper Addresses two issues raised in the PR review: - `_split` previously called `np.asarray(Xt)` unconditionally, which collapses `scipy.sparse` matrices to 0-d object arrays and breaks the subsequent column slicing. `_split` now routes both `'drop_first'` and `'full'` through `inverse_onehot` on the trailing `K - 1` columns (the `'full'` case's leading control column simply drops out, leaving the drop-first form), which works on sparse inputs natively via matmul; the `'label'` case densifies only the single trailing column. - The new positional-only constructor was a breaking change from the prior `MultiModelWrapper(model_list=[...])` and `MultiModelWrapper([...])` forms (`clone(list)` would fail). Both spellings now work again, emit a `FutureWarning`, and are unpacked into the existing positional path. Passing both `*models` and `model_list=` together raises `ValueError`. Tests cover sparse inputs across all three encodings (`drop_first`, `full`, `label`) and the two backward-compat constructor spellings (plus the conflict case). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Keith Battocchi <kebatt@microsoft.com>
fverac
approved these changes
Jun 17, 2026
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.
MultiModelWrapper is rewritten to dispatch rows to per-category models based on the treatment block trailing the feature matrix. It now supports three encodings via the new
encodingkeyword:'drop_first'(default): the lastK - 1columns are a drop-first one-hot encoding; the all-zero row is the control category. This matches the encoding used throughout EconML's discrete-treatment estimators (e.g. viaOneHotEncoder(drop='first')in DRLearner), so MultiModelWrapper can now actually be used as a first-stage model for them.'full': the lastKcolumns are a full one-hot encoding (column 0 is the control). This matches what DROrthoForest's nuisance estimator feeds tomodel_Yafter re-adding the dropped column.'label': the last column is an integer category index in{0, ..., K - 1}. This makes MultiModelWrapper a strict superset of SeparateModel.A single positional model combined with
n_categories=Kclones the model K times, which is convenient when all arms share a hyperparameter search.The previous MultiModelWrapper implementation was untested, not exported anywhere else in the package, advertised by DROrthoForest's docstrings, and could not actually work with DROrthoForest because it did not account for the extra one-hot column that nuisance_estimator_generator stitches on. The DROrthoForest docstrings are updated to recommend the
encoding='full'variant explicitly, and a TODO is added to nuisance_estimator_generator noting that we could later harmonize with the rest of EconML by feeding the drop-first encoding directly.SeparateModel is now deprecated via the existing
@deprecateddecorator;MultiModelWrapper(..., encoding='label')is a drop-in replacement. The existing test in test_model_selection.py that exercised SeparateModel is migrated to MultiModelWrapper.Tests cover all three encodings, validation paths, sample_weight forwarding, the single-model +
n_categoriesform, an end-to-end LinearDRLearner integration with 3 treatment categories (which the old wrapper could not handle), and the SeparateModel deprecation warning.