Forest Kernel Energy Balancing for Causal Inference

forestBalance estimates average treatment effects (ATE) by combining multivariate random forests with kernel energy balancing. A joint forest model of covariates, treatment, and outcome defines a proximity kernel that characterizes the confounding structure and emphasizes similarity of observations in terms of confounding. Distributional balancing weights are then obtained via a closed-form kernel energy distance solution. By construction, these balancing weights aim to balance the joint distribution of confounders specifically.

The method is described in:

De, S. and Huling, J.D. (2025). Data adaptive covariate balancing for causal effect estimation for high dimensional data. arXiv:2512.18069.

Installation

# Install from GitHub
devtools::install_github("jaredhuling/forestBalance")

Quick start

# library(forestBalance)
# Simulate observational data with nonlinear confounding (true ATE = 0)
set.seed(123)
dat <- simulate_data(n = 500, p = 10, ate = 0)
# Estimate ATE with forest kernel energy balancing
fit <- forest_balance(dat$X, dat$A, dat$Y)
fit
#> Forest Kernel Energy Balancing
#> -------------------------------------------------- 
#>   n = 500  (n_treated = 173, n_control = 327)
#>   Trees: 1000
#>   Solver: direct
#>   ATE estimate: 0.0455
#>   ESS: treated = 105/173 (61%)   control = 232/327 (71%)
#> -------------------------------------------------- 
#> Use summary() for covariate balance details.

How it works

The method proceeds in three steps:

1. Joint forest model: A grf::multi_regression_forest is fit on covariates