CausalTree regression with 'max_leaf_nodes=xxx' doesn't work

[lmaors]

Describe the bug
Causaltree regression with 'max_leaf_nodes=xxx' doesn't work in synthetic dataset, the ate score (ctree_1) is similar to the score of random

Screenshots

Environment (please complete the following information):

• Python Version: 3.7
• Causalml Version: newest

Additional context

y, X, w, tau, b, e = synthetic_data(mode=2, n=1000000, p=20, sigma=5.0)

df = pd.DataFrame(X)
feature_names = [f'feature_{i}' for i in range(X.shape[1])]
df.columns = feature_names
df['outcome'] = y
df['treatment'] = w
df['treatment_effect'] = tau

df_train, df_test = train_test_split(df, test_size=0.2, random_state=111)
n_test = df_test.shape[0]
n_train = df_train.shape[0]

tree1 = CausalTreeRegressor(criterion='causal_mse',
                            control_name=0,
                            min_samples_leaf=200,
                            leaves_groups_cnt=True)
tree1.fit(X=df_train[feature_names].values,
          treatment=df_train['treatment'].values,
          y=df_train['outcome'].values
          )
tree2 = CausalTreeRegressor(criterion='causal_mse',
                            max_leaf_nodes=512,
                            control_name=0,
                            min_samples_leaf=1,
                            leaves_groups_cnt=True)
tree2.fit(X=df_train[feature_names].values,
          treatment=df_train['treatment'].values,
          y=df_train['outcome'].values
          )

tree1_ite_pred = tree1.predict(df_test[feature_names].values)
tree2_ite_pred = tree2.predict(df_test[feature_names].values)

df_result = pd.DataFrame(
    {
        'ctree_1': tree1_ite_pred,
        'ctree_2': tree2_ite_pred,
        'outcome': df_test['outcome'],
        'is_treated': df_test['treatment'],
        'treatment_effect': df_test['treatment_effect']
    }
)
df_result = df_result.reset_index(drop=True)

stat_columns = ['treatment_effect', 'is_treated', 'outcome',
                'ctree_2','ctree_1'
               ]
plot_qini(df_result[stat_columns],
          outcome_col='outcome',
          treatment_col='is_treated',
          treatment_effect_col='treatment_effect',
         )