Transformer-Based Reinforcement Learning for Phase I Dose Allocation

Shijie Yuan

stat.MLOct 30, 2026

Transformer-Based Reinforcement Learning for Phase I Dose Allocation

Authors: Shijie Yuan

Organizations: The University of Texas at Austin

Abstract

Phase I dose-finding trials must learn from sparse toxicity observations while limiting exposure to overly toxic doses. We develop a single reinforcement-learning policy for trials with different numbers of ordered dose levels. Each dose is represented as a token containing its normalized position, accumulated patient count, observed dose-limiting toxicities (DLTs), and current-dose indicator. A Transformer encoder summarizes the dose tokens, trial-conditioned cross-attention produces a fixed-dimensional trial representation, and a three-action policy chooses de-escalation, staying, or escalation. The final maximum tolerated dose (MTD) decision is deliberately kept outside the learned policy: after the planned sample size, weighted isotonic regression and a prespecified safety cap select the dose closest to the target toxicity probability. In the simulations, toxicity-penalized policies reduce DLT burden while unpenalized policies improve selection accuracy at the cost of greater toxicity. These results demonstrate a controllable safety--accuracy trade-off but do not establish superiority over conventional dose-finding designs.

Figures & tables

Method	PCS	POS	PCA	POA	False stop	Mean DLT	Unsafe fraction	Total reward
RL, $\lambda_{\mathrm{DLT}}=0.05$ , lr $10^{-4}$	$0.561\pm 0.017$	$0.172\pm 0.025$	$0.323\pm 0.026$	$0.160\pm 0.023$	$0.035\pm 0.020$	$6.51\pm 0.57$	$0.177\pm 0.027$	$-0.130\pm 0.016$
RL, $\lambda_{\mathrm{DLT}}=0.05$ , lr $3\times 10^{-5}$	$0.568\pm 0.017$	$0.155\pm 0.020$	$0.323\pm 0.036$	$0.160\pm 0.038$	$0.021\pm 0.013$	$6.97\pm 1.22$	$0.191\pm 0.055$	$-0.140\pm 0.051$
BOIN	$0.588$	$0.220$	$0.381$	$0.256$	$0.008$	$8.23$	$0.280$	$-0.172$
i3+3	$0.589$	$0.229$	$0.386$	$0.272$	$0.007$	$8.42$	$0.293$	$-0.175$
mTPI-2	$0.591$	$0.209$	$0.383$	$0.248$	$0.008$	$8.15$	$0.271$	$-0.162$
CRM	$0.596$	$0.211$	$0.387$	$0.267$	$0.008$	$8.39$	$0.285$	$-0.161$

Table 1: Primary five-seed v2-selection comparison with a DLT penalty. POS is the probability of selecting above the true MTD; PCA and POA are the proportions of patients allocated at and above the true MTD, respectively. RL values are mean

\pm

standard deviation across seeds; benchmark-design values are single matched simulation summaries. All methods use the same final selection rule.

Setting	PCS	POS	PCA	POA	False stop	Mean DLT	Unsafe fraction
$\lambda_{\mathrm{DLT}}=0.05$ , lr $10^{-4}$	$0.561\pm 0.017$	$0.172\pm 0.025$	$0.323\pm 0.026$	$0.160\pm 0.023$	$0.035\pm 0.020$	$6.51\pm 0.57$	$0.177\pm 0.027$
$\lambda_{\mathrm{DLT}}=0.05$ , lr $3\times 10^{-5}$	$0.568\pm 0.017$	$0.155\pm 0.020$	$0.323\pm 0.036$	$0.160\pm 0.038$	$0.021\pm 0.013$	$6.97\pm 1.22$	$0.191\pm 0.055$
$\lambda_{\mathrm{DLT}}=0$ , lr $10^{-4}$	$0.595\pm 0.003$	$0.229\pm 0.018$	$0.352\pm 0.032$	$0.324\pm 0.071$	$0.004\pm 0.003$	$9.77\pm 1.17$	$0.358\pm 0.069$
$\lambda_{\mathrm{DLT}}=0$ , lr $3\times 10^{-5}$	$0.605\pm 0.010$	$0.186\pm 0.022$	$0.377\pm 0.030$	$0.287\pm 0.054$	$0.008\pm 0.004$	$9.06\pm 0.74$	$0.316\pm 0.057$

Table 2: Sensitivity of RL operating characteristics to the DLT penalty. POS is the probability of selecting above the true MTD; PCA and POA are the proportions of patients allocated at and above the true MTD, respectively. Values are mean

\pm

standard deviation across five seeds.

Setting	PCS	False stop	Mean DLT	Unsafe fraction	Total reward
Baseline: batch 256, entropy 0.01, $\lambda=0.05$	$0.563\pm 0.022$	$0.026\pm 0.022$	$6.63\pm 0.77$	$0.181\pm 0.035$	$-0.132\pm 0.018$
Batch 1024, $\lambda=0.05$	$0.551\pm 0.013$	$0.044\pm 0.006$	$5.80\pm 0.31$	$0.149\pm 0.020$	$-0.110\pm 0.007$
Entropy 0.003, $\lambda=0.05$	$0.562\pm 0.021$	$0.041\pm 0.026$	$6.36\pm 0.81$	$0.169\pm 0.050$	$-0.116\pm 0.005$
Entropy 0.03, $\lambda=0.05$	$0.552\pm 0.007$	$0.030\pm 0.014$	$6.33\pm 0.39$	$0.167\pm 0.031$	$-0.133\pm 0.014$
Baseline: batch 256, entropy 0.01, $\lambda=0$	$0.595\pm 0.004$	$0.003\pm 0.003$	$9.54\pm 0.27$	$0.343\pm 0.015$	$-0.226\pm 0.019$
Batch 1024, $\lambda=0$	$0.597\pm 0.006$	$0.006\pm 0.005$	$10.04\pm 0.79$	$0.350\pm 0.024$	$-0.245\pm 0.036$

Table 3: Optimization sensitivity sweep. Values are mean

\pm

standard deviation across three seeds. The sweep used the original rule-based selection settings for its comparator simulations; only RL rows are shown.

Transformer-Based RL

Transformer-based reinforcement learning uses attention-based sequence models to learn, plan, and control from experience across tasks.

See all 55 papers

Adaptive Chemotherapy Control under Tumor Heterogeneity via Reinforcement Learning
Sep 10, 2026Bereket Sitotaw Kidane, Md Samiul Haque Motayed, Shuo WangRL Control Reinforcement Learning
Department of Mechanical and Aerospace Engineering, University of Texas at Arlington, Arlington, TX 76019, USA
Recurrent Deep Reinforcement Learning for Chemotherapy Control under Partial Observability
May 4, 2026Firas Mohamed Elamine Kiram, Imane Youkana, Rachida Saouli +2Reinforcement Learning Healthcare
Dept. of Information Engineering University of Padova Padova, Italy · LINFI Lab, Dept. of Computer Science University of Biskra Biskra, Algeria
CaresAI at CT-DEB26: Detecting Dosing Errors In Clinical Trials Using Domain-Specific Transformer Embeddings and Classification Models
Jun 29, 2026Leon Hamnett, Favour Igwezeke, Joseph Itopa Abubakar +1Clinical Decision Support Clinical NLP
CaresAI, Australia · Faculty of Pharmaceutical Sciences, Nsukka Enugu, Nigeria · Menzies School of Health Research, Australia

Sep 10, 2026J/K move · Enter open · S save

Adaptive Chemotherapy Control under Tumor Heterogeneity via Reinforcement Learning

Bereket Sitotaw Kidane, Md Samiul Haque Motayed, Shuo Wang

Department of Mechanical and Aerospace Engineering, University of Texas at Arlington, Arlington, TX 76019, USA

RL Control Reinforcement Learning Deep Q-Learning Oncology Generalization in RL

Open paper AskPDF arXiv

Designing effective chemotherapy regimens is hindered by tumor heterogeneity and drug resistance, which complicate the deployment of patient-specific model-based optimal control across diverse populations. We develop and compare closed-loop deep reinforcement learning (DRL) dosing policies with continuous (TD3) and discrete (DQN) action spaces trained on a high-dimensional heterogeneous tumor model. The DRL policies are benchmarked against a Pontryagin's Maximum Principle (PMP)-derived open-loop benchmark. We assess generalization under parametric heterogeneity using a 100-patient virtual cohort with plus or minus 10 percent uniform perturbations in growth and drug-sensitivity parameters. Across this cohort, TD3 achieves higher average tumor reduction, while DQN yields tighter inter-patient dosing consistency, revealing a clear efficacy-consistency trade-off in this study. Our simulations assume full observation of all tumor subpopulations; translation to sparse and noisy clinical measurements will require partial-observability formulations and/or state estimation. Overall, the results show that simulation-trained DRL can learn state-dependent feedback dosing policies that complement open-loop optimal control benchmarks.

Abstract

Transformer-Based RL

Explore similar work

Adaptive Chemotherapy Control under Tumor Heterogeneity via Reinforcement Learning

Recurrent Deep Reinforcement Learning for Chemotherapy Control under Partial Observability

CaresAI at CT-DEB26: Detecting Dosing Errors In Clinical Trials Using Domain-Specific Transformer Embeddings and Classification Models