Q-learning with censored data

Yair Goldberg; Michael R. Kosorok

Source-linked AI summary

Q-learning with censored data

Yair Goldberg, Michael R. Kosorok

arXiv:1205.6659v1May 30, 2012math.ST

TL;DR

The paper addresses multistage treatment decisions with flexible numbers of stages and censored survival outcomes. It develops censored Q-learning with inverse-probability-of-censoring weighting, finite-sample error bounds, and a fixed-stage auxiliary formulation. The method yields individualized treatment strategies and converges to the optimal policy under the stated Q-function approximation condition.

Problem
The paper addresses how to learn treatment regimes that maximize survival when treatment stages vary across patients and survival outcomes are censored.
Method
The authors develop a censored-Q-learning algorithm using inverse-probability-of-censoring weighting and an auxiliary fixed-stage problem to handle flexible stage counts.
Results
The method provides finite-sample generalization-error bounds, and expected survival under learned policies converges to that of the optimal policy when optimal Q-functions lie in the approximation spaces.
Takeaways & Limitations
The censored-Q-learning algorithm can identify patient-tailored treatment strategies in multistage trials with flexible numbers of stages.
Takeaways & Limitations
The method assumes censoring is independent of observed trajectories, and inverse-probability weighting may be inefficient when many trajectories are censored.

Abstract

from arXiv · show

We develop methodology for a multistage decision problem with flexible number of stages in which the rewards are survival times that are subject to censoring. We present a novel Q-learning algorithm that is adjusted for censored data and allows a flexible number of stages. We provide finite sample bounds on the generalization error of the policy learned by the algorithm, and show that when the optimal Q-function belongs to the approximation space, the expected survival time for policies obtained by the algorithm converges to that of the optimal policy. We simulate a multistage clinical trial with flexible number of stages and apply the proposed censored-Q-learning algorithm to find individualized treatment regimens. The methodology presented in this paper has implications in the design of personalized medicine trials in cancer and in other life-threatening diseases.

1. Introduction.

The paper addresses individualized treatment decisions for survival outcomes when censoring and patient-dependent numbers and timing of treatment stages complicate Q-learning. It develops a censored-Q-learning methodology with theoretical guarantees and a flexible-stage formulation.

Motivation: The problem concerns treatment regimes that maximize survival when patients may receive different numbers and timings of treatment lines and outcomes may be censored.These features arise naturally in cancer care, where disease progression and treatment tolerance determine subsequent therapy.
Challenges: Flexible stages and censoring make backward recursion difficult because the final stage is patient-dependent and censored trajectories may not reveal complete outcomes.These issues prevent direct definition of the last-stage optimization and complicate use of incomplete trajectories.
Contributions: The proposed Q-learning algorithm uses inverse-probability-of-censoring weighting to account for censored observations.The method targets policy learning for sequential-multiple-assignment randomized trials with flexible stage counts.
Contributions: The paper provides finite-sample generalization-error bounds and proves universal consistency when optimal Q-functions lie in the algorithm’s approximation spaces.The generalization error measures the average difference in expected survival time between the optimal and learned regimes.
Contributions: The authors develop an auxiliary fixed-stage problem whose results can be translated back to the original problem with flexible numbers and timings of stages.This auxiliary construction enables backward recursion while preserving conclusions for the original decision problem.

2. Q-learning.

Q-learning uses backward recursion to estimate policies without modeling the full process dynamics, but censoring and patient-dependent stage counts require a modified formulation. The paper positions its approach against existing fixed-stage or uncensored methods.

Reinforcement learning: Reinforcement learning estimates action–consequence relationships from trajectories and selects a policy that approximates the most desirable outcome.In medical applications, stages correspond to clinical decision points, actions to treatments, and rewards to patient outcomes.
Q-learning: A policy is a sequence of decision rules mapping patient history and prior actions to treatments, with the objective of maximizing expected cumulative reward.The optimal policy is characterized recursively through the Bellman equation.
Q-learning: Q-learning estimates Q-functions backward from the final stage and computes the policy without requiring the full transition and reward dynamics.This approach reduces memory and computational requirements relative to learning the complete process model.
Challenges: Flexible stage counts make the recursion’s base step unclear, while censoring produces incomplete trajectories and unknown rewards at censoring times.These problems persist even when the nominal multistage structure is fixed if censored patients do not reach all stages.
Existing approaches: Prior censored-Q-learning work handled censoring with support vector regression for a fixed number of stages but did not evaluate theoretical properties.The paper states that existing approaches did not simultaneously address censoring and flexible stage counts.

3. Preliminaries.

The paper formalizes multistage patient trajectories, policies, rewards, and censoring, then targets truncated expected survival time under policies. The setup permits patient-specific stage counts within a finite maximum.

Decision-process setup: The maximum number of decision points is finite, but individual patients can have different numbers of stages because failure can terminate treatment.States include patient covariates and accumulated reward, while actions come from a finite discrete action space.
Censoring assumptions: Censoring is modeled through a variable C that is independent of covariates and failure time, supported on [0,τ], with survival probability bounded away from zero at τ.The censoring indicators record whether censoring has occurred before each subsequent decision point.
Observed trajectories: Failure and censoring alter the usual full-length trajectory because treatment can stop before the maximum decision point and censoring can leave later states and rewards unobserved.The resulting observed trajectory may be shorter than the nominal 2T+1 sequence.
Policy objective: A policy is a sequence of deterministic decision rules, and the objective is to maximize expected truncated-by-τ survival time under the policy-induced trajectory distribution.Truncation reflects the assumption that censoring occurs by τ and the observed maximum survival time is therefore bounded by τ.

4. The auxiliary problem.

The auxiliary problem converts variable-length, censored trajectories into fixed-length modified trajectories with bounded cumulative rewards. Its value functions preserve the original policy values, allowing fixed-stage Q-learning theory to transfer back.

Trajectory construction: The auxiliary construction modifies trajectories so they have fixed length T and modified cumulative rewards no greater than τ.After failure or truncation, later states are completed and subsequent actions are drawn uniformly from the action space.
Auxiliary process: The auxiliary distribution and exploration policy are defined from the original process, with the modified policy agreeing with the original on nonterminal states.At states created by the modification, actions are selected uniformly.
Value functions: The auxiliary value function represents expected truncated remaining survival time from a decision state when a policy is followed thereafter.The corresponding Q-function conditions additionally on the action selected at the current stage.
Value preservation: Lemma 4.1 establishes that auxiliary-policy values equal the expected truncated-by-τ survival times of corresponding policies in the original problem.The result also relates the optimal values after maximizing over the respective policy classes.
Value preservation: The proof uses a one-to-one correspondence between original and modified trajectory sets together with the construction of completed trajectories.This correspondence supports translating results from the fixed-stage auxiliary problem to the original flexible-stage problem.

5. The censored-Q-learning algorithm.

The paper maps the flexible-stage problem to an auxiliary fixed-stage problem, then estimates censored Q-functions by backward recursion and defines the policy through their action-wise maxima.

5. The censored-Q-learning algorithm: The algorithm maps the flexible-stage decision problem to an auxiliary fixed-stage problem, enabling backward recursion for Q-function estimation.The policy is obtained after approximating the Q-functions in the auxiliary problem.
5. The censored-Q-learning algorithm: The estimated policy chooses each action by maximizing the corresponding approximated Q-function over available actions.The construction sets the terminal Q-function to zero and recursively estimates earlier-stage functions.
5. The censored-Q-learning algorithm: Censoring makes the ideal empirical recursion infeasible because some rewards are unknown, so the method incorporates the censoring survival function into estimation.The censoring survival function is estimated with a Kaplan–Meier estimator.

6. Theoretical results.

The theoretical analysis derives finite-sample guarantees for the learned policy under bounded, entropy-controlled approximation spaces. When the optimal Q-functions lie in those spaces, the policy’s expected truncated survival-time gap converges to zero.

6. Theoretical results: When every optimal Q-function belongs to its corresponding approximation space, the algorithm is universally consistent.The consistency result follows by making the approximation error term vanish under this condition.
6. Theoretical results: The guarantees assume bounded approximation functions and a uniform entropy bound controlling the complexity of the Q-function spaces.Uniform entropy supports exponential control of differences between empirical and true expectations involving the Kaplan–Meier estimator.
6. Theoretical results: The analysis proves a finite-sample bound on the difference between the expected truncated survival times of the optimal policy and the learned policy.The proof maps the original problem to an auxiliary one and controls empirical-process, approximation, and censoring-related errors.
6. Theoretical results: The finite-sample bound is not directly suitable for structural risk minimization because it is too loose and contains model-dependent constants that are unspecified.The authors note that the bound can still be used to derive asymptotic rates.
6. Theoretical results: The proof handles censoring through concentration bounds for the Kaplan–Meier estimator and assumes the censoring survival function remains bounded away from zero.These ingredients control the random censoring-weighting terms in the error analysis.

7. Simulation study.

The simulation evaluates censored-Q-learning in a flexible-stage cancer-trial setting, comparing learned policies with fixed treatments and the optimal policy under varying censoring. Results show decreasing variance with larger samples, a gap from the optimal policy under linear Q-function approximation, and substantial bias from ignoring censoring in one setting.

Simulation design: The simulation models a three-year cancer trial in which treatment stages begin when tumor size reaches a critical level, allowing patients to receive different numbers of treatments.Patients receive either aggressive treatment A or less aggressive treatment B, and stages end at tumor progression, failure, or trial completion.
Algorithm and evaluation: The algorithm estimates the censoring survival function, computes Q-functions backward from the final stage, and derives the policy from the estimated Q-functions.The implementation uses trajectories generated under the simulated treatment and disease dynamics.
Simulation results: Estimated expected-survival variance decreases with larger sample sizes and is modestly smaller under lower censoring percentages.This pattern is reported across sample sizes from 50 to 3200 and censoring levels from 0% to 30%.
Simulation results: The learned policy achieves expected survival times a little above 17 months, compared with 17.85 months for the optimal policy.The gap is attributed to estimating Q-functions with linear functions while the optimal Q-function is nonlinear.
Simulation results: Ignoring censored observations produces substantial bias when censoring follows an exponential distribution, although the difference is moderate under uniform censoring.The comparison includes weighting-based estimation, deleting censored trajectories, and deleting the last stage of censored trajectories.

8. Summary.

The paper studies censored survival rewards in flexible-stage multistage decisions, proposes an adjusted Q-learning algorithm, and evaluates its performance theoretically and by simulation. It also identifies independence-of-censoring and efficiency limitations.

8. Summary.: The proposed Q-learning algorithm addresses multistage survival decisions with flexible stages and censoring, with performance demonstrated through simulations.The framework targets real-world multistage decision problems in which survival-time rewards are censored.
8. Summary.: The analysis assumes censoring is independent of observed trajectories, leaving covariate-dependent censoring as an unresolved extension.The authors identify relaxing this assumption as a challenge.
8. Summary.: Under average 50% exponential censoring, Figure 6 compares expected survival times across data-set sizes for the proposed and two censoring-ignoring algorithms.Curves summarize means over 400 simulation repetitions.
8. Summary.: The algorithm uses inverse-probability-of-censoring weighting, but may become inefficient when the percentage of censored trajectories is large.Developing a more efficient method remains an open question.

APPENDIX: SUPPLEMENTARY PROOFS

The supplementary proofs establish entropy bounds for function classes used in the theoretical analysis. They combine covering-number results for transformed, multiplied, squared, and maximized function classes with empirical-process arguments.

APPENDIX: SUPPLEMENTARY PROOFS: The appendix proves an exponential bound on the difference between empirical and true expectations as a function of class uniform entropy.This result is identified as Lemma A.6 and supports the theoretical analysis.
APPENDIX: SUPPLEMENTARY PROOFS: A composition lemma bounds covering numbers for functions formed by applying a Lipschitz map to products of uniformly bounded function classes.The proof works with finitely discrete probability measures and L2(P) distances.
APPENDIX: SUPPLEMENTARY PROOFS: The appendix derives covering-number consequences for reciprocal monotone classes, maxima over actions, and squared function classes.These results are stated in Corollaries A.2, A.3, and A.5.
APPENDIX: SUPPLEMENTARY PROOFS: Applying the preceding lemmas yields entropy bounds for the constructed class under the assumed uniform entropy conditions on the Q-function spaces.The resulting constants depend on the entropy parameters and related bounds.
APPENDIX: SUPPLEMENTARY PROOFS: The appendix concludes the entropy result for every finitely discrete probability measure and invokes an empirical-process theorem for the second assertion.The proof uses a universal covering-number bound for a monotone class.

SUPPLEMENTARY MATERIAL

The supplementary package provides code and datasets accompanying the paper.

SUPPLEMENTARY MATERIAL: Code and datasets are available in the supplementary ZIP archive, with usage details provided in README.pdf.The archive is identified by DOI 10.1214/12-AOS968SUPP.