rjaf: Regularized Joint Assignment Forest with Treatment Arm Clustering

Learning optimal assignment of treatments is an important problem in economics, public health, and related fields, particularly when faced with a variety of treatment strategies. The problem arises, for example, in settings where randomized controlled trials (RCT) are conducted to evaluate various behavioral science-informed interventions aimed at fostering behavior change (Milkman, Gromet, et al., 2021). Such interventions have been studied across diverse domains, including encouraging gym attendance and increasing vaccine uptake for influenza or COVID-19 (Dai et al., 2021; Milkman, Gromet, et al., 2021; Milkman, Patel, et al., 2021; Milkman et al., 2022). While most studies focus on identifying interventions that perform best on average, this approach often overlooks effect heterogeneity. Ignoring heterogeneity can be a missed opportunity to tailor interventions for maximum effectiveness and may even exacerbate disparities (Bryan et al., 2021). Subject-specific covariates, such as sociodemographics can be harnessed to identify which interventions work best for different segments of the population, allowing for more impactful intervention assignments. The rjaf package provides a user-friendly implementation of the regularized joint assignment forest (RJAF) (Ladhania et al., 2023), a regularized forest-type assignment algorithm based on greedy recursive partitioning (Athey et al., 2019) that shrinks effect estimates across treatment arms. The algorithm is augmented by outcome residualization to reduce baseline variation, and employs a clustering scheme (Hartigan & Wong, 1979) that combines treatment arms with consistently similar outcomes. Personalized treatment learning is achieved by optimizing a regularized empirical analogue of the expected outcome. The integration of R (R Core Team, 2024) and C++ (Stroustrup, 2013) substantially boosts computational efficiency in tree partitioning and aggregating. It is especially suitable in RCT settings with numerous treatment arms and constrained sample sizes, making it a powerful tool for learning personalized intervention strategies.