📊 statistics

Counterfactual Q Learning via the Linear Buckley James Method for Longitudinal Survival Data

This paper proposes a Counterfactual Buckley-James Q-Learning framework that integrates the Buckley-James imputation method with reinforcement learning to estimate optimal dynamic treatment regimes for maximizing survival time in the presence of censored longitudinal data.

Original authors: Jeongjin Lee, Jong-Min Kim

Published 2026-07-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jeongjin Lee, Jong-Min Kim

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a coach trying to build the ultimate training schedule for a team of athletes. Your goal is simple: keep them in the game as long as possible. However, there's a catch. Some athletes quit the team early (they drop out), some get injured and stop playing before the season ends, and some just leave the stadium before the final whistle blows. In statistics, this is called censoring. You know they were playing up to a certain point, but you don't know how long they would have lasted if they had stayed.

This paper introduces a new coaching strategy called Counterfactual Buckley-James Q-Learning. Here is how it works, broken down into simple parts:

1. The Problem: The "Ghost" Athletes

In traditional medical studies (like the Cox model mentioned in the paper), when an athlete leaves early, statisticians often treat them as if they never existed or try to guess their future based on a rigid rule (like "everyone slows down at the same rate"). This can lead to bad advice. If you don't know how long an athlete would have played, you might pick the wrong training plan for the next season.

2. The Solution: The "Crystal Ball" Imputation

The authors use a clever trick called the Buckley-James method. Think of this as a crystal ball that doesn't predict the future magically, but rather uses math to make a very educated guess.

How it works: If an athlete leaves early, the method looks at everyone else who is similar (same age, same injury history, same starting stats) and asks, "Based on the people who stayed, how much longer would this person have likely played?"
The Result: It fills in the missing "ghost" time with a calculated estimate. Now, instead of having a broken dataset with holes, the coach has a complete picture of every athlete's potential lifespan.

3. The Strategy: Learning by Rewinding (Q-Learning)

Once the data is "filled in," the paper uses a technique called Q-Learning. Imagine playing a video game where you want to find the path to the highest score.

The Game: The "game" is a patient's life over time.
The Moves: At each checkpoint (like a doctor's visit), the patient gets a treatment (Medicine A or Medicine B).
The Reward: The "score" is how long the patient survives.
The Trick: The algorithm works backwards. It starts at the end of the game, looks at who survived the longest, and then rewinds to figure out: "If I had chosen Medicine A at the start, would that player have ended up with a higher score?"

By combining the "Crystal Ball" (Buckley-James) with the "Backwards Rewind" (Q-Learning), the system learns the best sequence of moves to keep the player in the game the longest.

4. The Proof: The Simulation Race

The authors tested this new method against the old standard (the Cox model) in a giant computer simulation.

The Setup: They created 1,000 fake patients with different body types and tumor sizes. They gave them treatments and then "censored" 50% of them (made them leave the study early).
The Race: They asked, "Who can figure out the best treatment plan?"
The Winner: The new Buckley-James Q-Learning method was like a master strategist. It correctly identified the best treatment about 97% of the time with large groups of people. The old method (Cox) was like a novice coach, getting it right only about 77% of the time. The old method struggled because it couldn't handle the "ghost" data as well.

5. The Real-World Test: The HIV Dataset

The authors tried their method on real data from a famous HIV study (ACTG175).

The Challenge: This data was messy. 75% of the patients were "censored" (they left the study before the event happened), and some data was missing.
The Finding: When they used their method to fill in the missing times, the results clearly showed that a combination therapy (using two drugs together) kept patients alive longer than using just one drug. The "imputed" curves (the filled-in guesses) were smoother and more confident than the jagged lines of the raw, incomplete data.

The Bottom Line

This paper claims that by using a mathematical "fill-in-the-blanks" technique (Buckley-James) combined with a smart, backward-looking learning algorithm (Q-Learning), doctors can make much better decisions about which treatments to give patients, even when a lot of the data is incomplete. It turns a messy, half-finished puzzle into a clear picture of what works best.

Technical Summary: Counterfactual Q-Learning via the Linear Buckley–James Method for Longitudinal Survival Data

Problem Statement
Personalized healthcare increasingly relies on Dynamic Treatment Regimes (DTRs) to optimize sequential treatment decisions based on evolving patient conditions. However, a significant challenge in applying reinforcement learning (specifically Q-Learning) to survival analysis is the prevalence of censored data. In clinical trials, survival times are often incomplete due to patient dropout or study termination (right-censoring). Standard Q-Learning frameworks struggle with this incompleteness because they require precise reward estimates (survival times) to learn optimal policies. While the Cox proportional hazards model is a common tool for survival analysis, it relies on the proportional hazards assumption and offers hazard ratios that can be counterintuitive for direct outcome comparison. Furthermore, existing methods for handling censoring in DTRs, such as inverse probability weighting or weighted least squares, may not adequately facilitate direct comparisons of imputed survival times across different treatment strategies under a counterfactual framework.

Methodology
The authors propose the Linear Buckley–James Q-Learning (BJ-Q) framework, which integrates the Buckley–James (BJ) imputation method with reinforcement learning to address censored survival data.

Framework Structure: The method operates within a longitudinal setting with $K$ decision points. At each stage $k$ , the goal is to estimate the Q-function, defined as the expected cumulative survival time from stage $k$ onward given the patient's history and a specific treatment.
Buckley–James Imputation: To handle right-censoring, the framework employs the Buckley–James estimator within an Accelerated Failure Time (AFT) model. For censored observations, the unobserved survival time is imputed using the conditional expectation of the survival time given the observed data and the censoring status.
- The imputed survival time $Y^*_{i,k}$ is calculated as $\delta_{i,k}T_{i,k} + (1-\delta_{i,k})E[T_{i,k} | T_{i,k} > C_i, H_{i,k}]$ .
- The conditional expectation is approximated using the Kaplan–Meier estimator of the residual distribution.
- Regression coefficients are estimated iteratively using a Nelder–Mead simplex algorithm to solve the modified estimating equations, ensuring convergence despite the non-monotonicity of the objective function.
Recursive Q-Learning: The framework utilizes backward induction.
- Final Stage ( $K$ ): The Q-function is initialized using the Buckley–James imputed survival times.
- Intermediate Stages ( $k < K$ ): Pseudo-outcomes are constructed by summing the immediate imputed survival time at stage $k$ and the maximum estimated Q-value from stage $k+1$ . A linear regression model is then fitted to these pseudo-outcomes to estimate the Q-function parameters for stage $k$ .
Counterfactual Decision Making: To determine the optimal treatment, the method evaluates the Q-function under hypothetical treatment assignments (counterfactuals). The optimal decision rule selects the treatment that maximizes the estimated Q-value, effectively comparing potential outcomes across different strategies.

Key Contributions

Integration of BJ and Q-Learning: The paper introduces a novel framework that combines the robustness of the Buckley–James imputation method with the sequential decision-making capabilities of Q-Learning, specifically tailored for censored survival data.
Direct Comparison of Strategies: Unlike some existing methods, this approach explicitly accounts for censored observations while ensuring that imputed survival times allow for robust, direct comparisons across different treatment regimes under a counterfactual framework.
Handling Treatment Effect Heterogeneity: The method is designed to capture and leverage treatment effect heterogeneity (e.g., how tumor size modifies treatment efficacy) to identify personalized optimal strategies.
Extension to Double Censoring: The methodology notes that the Buckley–James component can be extended to handle double censoring (both left and right), broadening its applicability to more complex data scenarios.

Results
The study evaluates the BJ-Q framework through extensive simulation studies and a real-world application:

Simulation Studies:
- Single-Stage DTR: In simulations with 50% right-censoring and 50% missingness, BJ-Q consistently produced Q-value estimates closely aligned with true oracle values across sample sizes ( $n=100, 500, 1000$ ). In contrast, Cox-based Q-Learning (Cox-Q) demonstrated persistent bias, particularly in capturing treatment effect heterogeneity.
- Decision Accuracy: BJ-Q achieved significantly higher treatment decision accuracy than Cox-Q. For $n=1000$ , the median accuracy for BJ-Q was 0.973, compared to 0.770 for Cox-Q.
- Two-Stage DTR: In multi-stage settings, BJ-Q maintained robust performance with cumulative decision accuracies reaching 0.900 at $n=1000$ . Cox-Q showed limited improvement with increasing sample size, with cumulative accuracy stagnating around 0.506, highlighting its inability to accurately recover total survival under censoring in complex longitudinal settings.
Real-World Application (ACTG175): The method was applied to the ACTG175 HIV clinical trial dataset, which features a high censoring rate (~75%). The analysis compared monotherapy (ddI) versus combination therapy (AZT + ddI). The BJ-Q framework successfully imputed censored survival times, revealing that combination therapy consistently yielded superior survival outcomes compared to monotherapy. The Kaplan–Meier curves based on imputed data showed tighter confidence bands and clearer separation between treatment groups than observed data alone.

Significance and Claims
The authors claim that the BJ-Q framework provides a principled and practical tool for optimizing survival outcomes in complex clinical settings where data is incomplete. The significance of the work lies in its ability to:

Overcome Censoring Bias: By imputing censored times via conditional expectations within an AFT model, the method avoids the proportional hazards assumptions and interpretational difficulties of the Cox model.
Enable Personalized DTRs: It facilitates the identification of optimal dynamic treatment regimes that maximize expected survival time, even when outcomes are censored.
Provide Robustness: The method demonstrates superior performance over Cox-based alternatives in both single and multi-stage settings, particularly in scenarios with treatment effect heterogeneity.

The authors maintain a modest stance regarding theoretical guarantees, noting that while the method is expected to converge to the optimal decision as sample size increases, the validity depends on the behavior of Kaplan–Meier residuals and may require additional regularity conditions regarding the censoring mechanism. They acknowledge that the current implementation relies on the Kaplan–Meier estimator for residual approximation and suggest that future work could explore nonparametric machine learning methods (e.g., random survival forests) to address potential biases from informative censoring. Additionally, they propose extending the framework to high-dimensional covariate spaces using penalized methods.