# CAUSALRLSTACK: adaptive balancing of deep representation and causal effect estimation with application to HIV-related health data

## TL;DR

This paper introduces CAUSALRLSTACK, a new framework that improves individualized causal effect estimation in health data by combining modular components and reinforcement learning.

## Contribution

The novel contribution is a modular framework that adaptively balances representation learning and causal effect estimation using reinforcement learning.

## Key findings

- CAUSALRLSTACK outperformed six state-of-the-art models on HIV datasets with higher accuracy, F1-Score, and AUC-ROC.
- The framework achieved the lowest predictive uncertainty, showing robust performance in treatment effect estimation.
- The modular design and adaptive weighting strategy enhance generalizability across diverse populations.

## Abstract

Estimating individualized causal effects plays a vital role in data-driven decision-making, especially in high-risk domains such as public health. However, current causal inference models often lack flexibility and generalizability due to the tight coupling between representation learning and effect estimation. This study aims to develop a modular and adaptive framework to enhance the analysis of individualized causal effects in complex health data.

We propose CAUSALRLSTACK, a modular framework designed to separate representation learning from causal effect estimation. In practice, the model uses a memory-augmented Transformer (TITAN) to capture complex, individualized representations. It is further paired with a doubly robust estimator(DRLearner) to improve the treatment effect estimation. A reinforcement learning agent adjusts how much each component contributes by assigning instance-specific weights. This adaptive weighting process improves the model’s ability to generalize across different populations. Input features are derived from causal graphs, automatically chosen between an expert-defined graph and one discovered from data. To evaluate performance, we applied the framework to two publicly available HIV datasets that reflect community-level testing behavior and post-intervention clinical outcomes.

CAUSALRLSTACK outperforms six state-of-the-art causal inference models across both datasets, achieving the highest accuracy (0.861 and 0.855), F1-Score (0.845 and 0.839), and AUC-ROC (0.897 and 0.892). It also achieves the lowest predictive uncertainty (0.093 and 0.092), indicating robust performance in estimating treatment effects.

The proposed framework offers a flexible and effective solution for individualized causal inference. Its modular architecture and reinforcement learning-based weighting strategy enable adaptive, data-driven estimation across diverse populations. Strong experimental results demonstrate the potential of the framework to advance individualized causal inference in health data and provide a practical basis for designing personalized intervention strategies in HIV and broader public health domains.