Allowing Categorical Features in T/R Learners

[portchester1989]

Is your feature request related to a problem? Please describe.
When I have a categorical features in part of the dataset, it errors out even if I try to use XGBoost model eith enable_categorical_feature on.

Describe the solution you'd like
Adding CatBoost model, so that we can handle categorical feature.

Describe alternatives you've considered
using XGBRRegressor or XGBTRegressor with enable_categorical_feature on.

Additional context

[jbbqqf]

Hi! Triaging older issues — I think this is largely solvable today and may be a candidate for closure.

The T/R-learners are model-agnostic — you can pass any sklearn-like estimator that supports categorical features natively:

• XGBRegressor/XGBClassifier with enable_categorical=True (note: the kwarg is enable_categorical, not enable_categorical_feature — the issue's reported name has an extra suffix that XGBoost doesn't recognize, which would explain the silent fall-through). Pass pd.Categorical columns and it works.
• LGBMRegressor/LGBMClassifier with the categorical_feature constructor / fit kwarg.
• CatBoostRegressor/CatBoostClassifier with cat_features=[...].

The wrapper invocation would look like BaseTRegressor(learner=XGBRegressor(enable_categorical=True)) and pass a DataFrame with dtype="category" columns through fit. If a concrete reproducer with enable_categorical=True (not enable_categorical_feature) still errors out, that'd be a different bug worth a fresh issue with the traceback.

(Disclosure: I drafted this comment with help from Claude Code while triaging stale issues; the references above were verified manually against the current master branch and against XGBoost's documented categorical-feature support.)

[wangkang85]

I ran into the same issue. The problem is that the T-Learner calls convert_pd_to_np() and it fails when the column type is categorical. So setting enable_categorical=True doesn't work. Thank you for looking into it though.

[aman-coder03]

hey, found the actual bug here..
the issue is in convert_pd_to_np() in causalml/inference/meta/utils.py. it calls .to_numpy() on the dataframe before passing it to the learner, and that strips all dtype info including category
so by the time xgboost sees the data, there are no categorical columns left and enable_categorical=True does nothing

the fix has two parts...

1. utils.py skip conversion when the dataframe has categorical columns...

def convert_pd_to_np(*args):
    def _convert(obj):
        if isinstance(obj, pd.DataFrame) and any(
            pd.api.types.is_categorical_dtype(obj[c]) for c in obj.columns
        ):
            return obj  # pass through so learners can handle categoricals natively
        return obj.to_numpy() if hasattr(obj, "to_numpy") else obj
    output = [_convert(obj) for obj in args]
    return output if len(output) > 1 else output[0]

2. tlearner.py stop converting X to numpy in fit, predict, fit_predict, and estimate_ate...
   only treatment and y need to be numpy (for masking and arithmetic). X should be left alone so categorical dtypes survive all the way to the learner

# before
X, treatment, y = convert_pd_to_np(X, treatment, y)

# after
treatment, y = convert_pd_to_np(treatment, y)

dataframes without any categorical columns still get converted to numpy like before, so nothing breaks for existing users.
one note on XGBoost 3.x, string-valued categoricals don't work even with this fix, you need integer-coded ones...

# works
pd.Categorical(np.random.choice([0, 1, 2], size=n))

# breaks in XGBoost 3.x
pd.Categorical(np.random.choice(['a', 'b', 'c'], size=n))

after the fix..

from xgboost import XGBRegressor
from causalml.inference.meta import BaseTRegressor

learner = BaseTRegressor(learner=XGBRegressor(enable_categorical=True))
# categorical columns must have dtype="category" with integer codes
learner.fit(X, treatment, y)

opening a PR with these changes + tests for both T and R learners

[jeongyoonlee]

resolved by #901