ForestDRLearner : outcome binary and treatement is discret ( 3 values)

i'm building model with ForestDRLearner . I would to have the treatment which minimizes the outcome and in the end to have
client, best_treatment
1, 0
2, 1
3, 2
4, 0
ect ...

how make this final dataset with this code ? what is the best solution ? this code is not quite what I need

X = sampling.drop(columns=['T', 'Y'])
Y = sampling['Y']
T = sampling['T']

X_train, X_test, T_train, T_test, Y_train, Y_test = train_test_split(X, T, Y, test_size=0.2, random_state=123)

model = ForestDRLearner(
model_propensity=XGBClassifier(learning_rate=0.1, max_depth=3, objective="multi:softprob"),
model_regression=XGBClassifier(learning_rate=0.1, max_depth=3, objective="binary:logistic"),
discrete_outcome=True,
random_state=1,
)

model.fit(Y=Y_train, T=T_train, X=X_train, inference="auto")

cate_estimates = model.effect(X_test)
cate_estimates

best_treatment = np.argmin(cate_estimates, axis=1)

results = pd.DataFrame({

'best_treatment': best_treatment

})

Start with the supplied ForestDRLearner.fit and model.effect workflow, checking how the learner represents a binary outcome and three discrete treatments. Reproduce the example and clarify whether the requested client-to-best-treatment table is supported; done when a documented, reproducible approach or a precise limitation is established.