Nonzero imbalance.penalty causes trees to split less (and tuning to fail)

gistfile1.R

library(grf)

set.seed(123)

# Very simple scenario, but already adversarial
n <- 200
p <- 2
X <- matrix(rnorm(n * p), n, p)
W <- rbinom(p=0.5, size=1, n=n)
eta <-  1
Tau <- 0.3 * X[,1]
Y <- eta + 1/2*(2*W - 1) * Tau + 0.1 * rnorm(n)

# This is without tuning
rf_notune <- causal_forest(X, Y, W,
                            tune.parameters = FALSE,
                            seed=123)
tau_notune <- predict(rf_notune)$predictions
infeasible_mse_notune <- mean((Tau - tau_notune)^2)
print("Infeasible MSE(no tune)")
print(infeasible_mse_notune)

# The parameters for min.node.size, mtry, and alpha are
# the ones that would be chosen by tuning, if we manually
# forced imbalance.penalty to zero.
rf_partial_manual_tune <- causal_forest(X, Y, W,
                                    tune.parameters = FALSE,
                                    min.node.size=4,
                                    mtry=2,
                                    seed=123,
                                    alpha=0.1,
                                    imbalance.penalty = 0)
tau_partial_manual_tuned <- predict(rf_partial_manual_tune)$predictions
infeasible_mse_partial <- mean((Tau - tau_partial_manual_tuned)^2)
print("Infeasible MSE(partial manual tune)")
print(infeasible_mse_partial)

# Now watch what happens when I turn imbalance.penalty on by a little bit.
rf_bad <- causal_forest(X, Y, W,
                        tune.parameters = FALSE,
                        min.node.size=4,
                        mtry=2,
                        seed=123,
                        alpha=0.1,
                        imbalance.penalty = 0.00000000000000001)
tau_bad <- predict(rf_bad)$predictions
infeasible_mse_bad <- mean((Tau - tau_bad)^2)
print("Infeasible MSE(bad)")
print(infeasible_mse_bad)