Address Copilot PR review comments

- Improve docstring completeness with detailed parameter descriptions
- Clarify usage of treatment/y parameters for classification metrics
- Explain return_components output format with probability details
- Separate DR learner imports for better readability
- Maintain API consistency with other meta-learners

Addresses comments in #844

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: jeongyoonlee

causalml/inference/meta/drlearner.py

@@ -531,11 +531,21 @@ def predict(
 
         Args:
             X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series, optional): a treatment vector
-            y (np.array or pd.Series, optional): an outcome vector
-            verbose (bool, optional): whether to output progress logs
+            treatment (np.array or pd.Series, optional): a treatment vector. Used for computing
+                classification metrics when y is also provided.
+            y (np.array or pd.Series, optional): an outcome vector. Used for computing
+                classification metrics when treatment is also provided.
+            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1). Currently not used in prediction but kept for API consistency.
+            return_components (bool, optional): whether to return outcome probabilities for treatment and control
+                groups separately. Defaults to False.
+            verbose (bool, optional): whether to output progress logs. Defaults to True.
         Returns:
             (numpy.ndarray): Predictions of treatment effects.
+            If return_components is True, also returns:
+                - dict: Predicted probabilities for the control group (yhat_cs).
+                - dict: Predicted probabilities for the treatment group (yhat_ts).
         """
         X, treatment, y = convert_pd_to_np(X, treatment, y)
 

tests/test_meta_learners.py

@@ -30,7 +30,9 @@
     XGBRRegressor,
 )
 from causalml.inference.meta import TMLELearner
-from causalml.inference.meta import BaseDRLearner, BaseDRRegressor, BaseDRClassifier
+from causalml.inference.meta import BaseDRLearner
+from causalml.inference.meta import BaseDRRegressor
+from causalml.inference.meta import BaseDRClassifier
 from causalml.metrics import ape, auuc_score
 
 from .const import RANDOM_SEED, N_SAMPLE, ERROR_THRESHOLD, CONTROL_NAME, CONVERSION