Optimal stopping for two-strike cancer therapy

Problem formulation

Background

Cancers present a remarkable capability of developing resistance to diverse treatments. This leads often leads to first-line therapy failure due to the emergence of and selection for resistance phenotypes. To address this issue, two-strike extinction therapy proposes the use of two treatments sequentially: the first reduces the tumor size, and the second pushes the now-vulnerable cancer cell population to extinction before it recovers from the first treatment. A schematic visualization of this protocol can be seen in the figure below.

The effectiveness of two-strike extinction therapy is highly dependent on the timing of the second strike. However, optimal timing varies across patients and is difficult to be assessed in real time. Previous work has shown that the ideal intervention time is often near the time of minimal tumor size following the first line of treatment, which can only be known retrospectively. This project formulates the second strike timing as an optimal stopping problem and uses machine learning to optimize two-strike cancer therapy.

---

Simulated dataset

The models are trained using an in silico dataset generated using SLiM following the protocol described in this repository.

Each simulated patient is characterized by a unique set of biological and treatment parameters, including

• Baseline mutation rate
• First-line treatment dosage
• Second-line treatment dosage
• Cross-resistance parameters
• Additional simulation parameters recorded in the manifest file

For every patient, the tumor is first simulated under continuous first-line therapy to obtain a baseline trajectory under a fixed seed. The tumor population is observed at 20 evenly spaced scan times, representing the times at which a clinician could choose to administer the second treatment.

To estimate the outcome of switching at each candidate scan, the simulation is branched from the baseline trajectory at every scan time. Beginning from the same tumor state, the second treatment is initiated and 100 independent stochastic simulations are performed. These replicates capture the inherent randomness of mutation, selection, and extinction, allowing the probability of successful treatment to be estimated empirically.

The average outcome over these 100 stochastic replicates defines the expected consequence of switching at a particular scan. A schematic representation of this approach is shown in the figure below.

These summaries form the target values used throughout the optimal stopping framework. The extinction probabilities are used to characterize the expected success of switching at each scan, while the individual replicate outcomes are used to construct reward targets and train the backward dynamic programming algorithm.

---

Data pipeline

The machine learning models are not trained directly from the SLiM output.

Instead, sim_data_processor.py converts the raw simulation logs into several processed datasets that are used during model training.

SLiM logs
      │
      ▼
sim_data_processor.py
      │
      ├── simulation_statistics.csv
      ├── extinction_by_lag.csv
      ├── model_expected_outcomes.json
      └── model_dataset.json

---

Raw simulation output

Input

The processor expects

• A directory containing all SLiM simulation logs
• The simulation manifest
• An output directory for recovered baseline trajectories

The log directory should have the structure

seed_<seed>/
    log_<switch_scan>_<replicate>.csv

Example

seed_2040/
    log_1_1.csv
    log_1_2.csv
    ...
    log_20_100.csv

Each log represents one stochastic realization after switching treatments at a specific scan.

---

Processing

For every patient, the processor

1. Reconstructs the baseline trajectory,
2. Extracts the final population size of every replicate,
3. Estimates extinction probability at every switch time,
4. Computes summary statistics,
5. Creates model-ready datasets,
6. Filters out patients that reach extinction before treatment 1 starts.

---

Processed datasets

simulation_statistics.csv

One row per patient containing all simulation parameters.

seed	base_mut_rate	dose_drug_1_mono	dose_drug_2_combo	...
2040	0.001	1.25	0.8	...
2041	0.0005	1.00	0.8	...

---

extinction_by_lag.csv

One row per patient and switch time.

seed	lag	n_reps	extinction_prob	mean_final_N	min_final_N	max_final_N
2040	1	100	0.01	814	0	1422
2041	2	100	0.08	401	0	918
...	...	...	...	...	...	...

This file summarizes the expected outcome of switching at each candidate scan.

---

expected_outcomes.json

Contains expected extinction statistics for every patient and every switch time.

{
    seed: {
        lag: {
            "extinction_prob": float,
            "avg_final_n": float
        }
    }
}

Example

{
    2040: {
        1: {
            "extinction_prob": 0.02,
            "avg_final_n": 812.5
        },
        2: {
            "extinction_prob": 0.08,
            "avg_final_n": 391.4
        }
    }
}

These values are used when training the reward and continuation value regressors.

---

model_dataset.json

Primary dataset used throughout training.

patient
│
├── baseline trajectory
├── mutation rate
├── drug 1 dose
│
└── switch outcomes
      │
      ├── lag 1
      │      ├── replicate 1
      │      ├── replicate 2
      │      └── ...
      │
      ├── lag 2
      └── ...

Each replicate stores

• Extinction status: binary indicator of whether extinction is reached for the given simulation)
• Time to extinction: if extinction is reached, represents the number of cycles applying drug 2 until extinction)
• Proportional population change: computed as $(N_{final} - N_t)/N_t$. Indicates the change in population from the switch point to the end divided by the population at the switch point.

Example

{
    2040: {

        "baseline": [
            1530,
            1492,
            1401,
            1270,
            ...
        ],

        "mut_rate": 0.001,

        "drug_1_dose": 1.25,

        "switches": {

            1: {

                1: {
                    "extinct": 0,
                    "time_to_extinction": null,
                    "percent_change": -0.12
                },

                2: {
                    "extinct": 1,
                    "time_to_extinction": 4.0,
                    "percent_change": -1.0
                }
            },

            2: {
                ...
            }
        }
    }
}

The baseline trajectory provides the observations available before switching, while the nested switches dictionary stores the outcome of every simulated intervention time. These data are used to construct training examples for the stopping classifiers and to compute reward and continuation-value targets during backward induction.

---

Machine learning pipeline

Additional pre-processing

Patients that never reach a minimal best case extinction probability are removed before training. This is governed by the parameter min_extinction_prob in config.py. The standard value used is 0.05.

State representation

During training, the processed dataset is converted into differently-sized state spaces based on how many scans have been collected. Since a different model is trained to predict switch probabilities, reward, and continuation values at each point, they can have varying-sized inputs.

In our model, we assume that we have a scan at every cycle in our simulation. For this reason, we don't include the timing of the scans in the state representation, given it's information is already encoded positionally.

For scan t, the input state has dimension t + 1 and the following form:

[
    drug_1_dose,
    first_scan,
    ...                 # [all_scans_up_to_time_t]
    current_scan
]

---

Model architecture

The optimal stopping policy is approximated using three groups of MLPs with backward induction. For all of these, training and architectural features such as number of hidden layers, batch sizes, and learning rate can be edited through the config.py file. Model classes are in the model_building/models.py script.

1. Stopping classifiers

At each scan, predict Switch or Continue. A sigmoid activation is used in the output layer so that the output is the probability of switch prediction .

---

2. Reward regressors

Estimate the immediate reward obtained by switching at the current scan.

The reward function for patient $i$ switching treatments at time $t$ is defined as:

$$R_i(t) = \frac{1}{N}\sum_{j=1}^N r_{i,j}(t)+\lambda t.$$

Where $r_{i,j}$ represents the reward for each future trajectory $j$ out of the $N$ simulated dynamics. $r_{i,j}$ is defined as:

$$ r_{i,j}(t)= \begin{cases} -\beta, & \text{if extinction does not occur},\\ \beta + \frac{T_{\max,i}(t)-T_{i,j}(t)}{T_{\max,i}(t)}, & \text{if extinction occurs}. \end{cases} $$

Variables and parameters explanation:

• $N$: the number of simulated future trajectories following a switch at any cycle.
• $\lambda$: fixed reward given when the model waits for another scan (benefit of getting more information). Can be edited under continuation_bonus in config.py.
• $\beta$: benefit for reaching / penalty for not reaching extinction. Can be edited under extinction_bonus in config.py.
• $T_{\max,i}(t)$: across the simulated trajectories switching at time $t$ that reached extinction, $T_{\max,i}(t)$ represents the one with the longest time to extinction after drug 2 is introduced.
• $T_{i,j}(t)$: the time it takes for simulation $j$ for patient $i$ switching treatments at time $t$ to reach exticntion.

---

3. Continuation value regressors

Estimate the expected future reward obtained by waiting.

The continuation value function is defined as:

$$C_i(t,X_i) = \sum_{k=t+1}^T\left[ \left( \prod_{j=t}^{k-1}(1-p_j) \right) \times p_k \times R_i(k) \right]$$

Where $p_j$ is defined as:

$$\mathbb{P}(\text{switch} | X_j) \approx \text{MLP}(X_j)$$

The MLP used here is the classifier that aims to predict the probability of switching treatments at time $j$. $X_j$ represents the covariates available by time $j$ (i.e. observed cell counts at all previous time points and the drug 1 dosage).

---

Decision rule

The policy compares the immediate reward of switching against the expected continuation value of waiting. At every scan, the decision is as follows:

$$\text{Switch if } \begin{cases} \hat{R}_i(t)\geq \hat{C}_i(t,X_i) \text{ and } p_t > \gamma,\\ \Delta N_{i,t} > 0 \text{ and } \hat{R}_i(t) \geq \eta \beta \end{cases}$$

Variables and parameters explanation:

• $\hat{R}_i(t)$ and $\hat{C}_i(t,X_i)$ are the estimated reward and continuation values for patient $i$ at timee $t$ by the appropriate MLP regressors.
• $\gamma$ is a minimal switch probability threshold for making a choice of switch. This can be edited under decision_threshold in config.py.
• $\Delta N_{i,t}$ represents the change in population size (cell count) from time $t-1$ to time $t$.
• $\eta$ is a value between 0 and 1 that represents the minimal percentage of the reward given to reaching extinction ($\beta$). Jointly, $\eta\beta$ defines a minimal value the reward must reach for a stop to be predicted when the population is increasing. $\eta$ can be edited under reward_stopping_pct in config.py.

---

Repository structure

config.py                  OSTConfig dataclass - every hyperparameter lives here
sim_data_processor.py     load data directly from simulation results and process it to be in the models' format  (longest run time: approximately 1.5 hour running on CPU for 80 patients with trajectories with 20 cycles and 100 potential outcomes at each switch)
data_loading.py             load pre-processed data / filter / split / normalize
model_building/
    models.py                StopClassifier, MLPRegressor (network definitions)
    factory.py               centralized control for building / saving/loading models
    state.py                   build_state, build_scan_dataset, build_labels (builds model-appropriate data formats)
    reward.py                   immediate_reward, continuation_value (defines value functions for switches)
    training.py                  train_backward, train_reward_continuation (run training loops for classifiers and regressors)
    prediction.py                 predict_switch_scan (runs prediction for new patients)
results.py                  save predictions/performance to disk as tidy CSVs
cross_validation.py         k-fold CV orchestration
visualizations/
    data_exploration.py       pre-modeling plots (extinction distributions, trajectories)
    performance_plots.py       classifier/regression performance by scan, CV variance
    prediction_plots.py         switch-scan / extinction-prob outcome plots
main.py                     single train/test run (with defined train/test split)
run_cross_validation.py     full CV run + saved plots (edit number of folds in config.py)
sweep.py                    hyperparameter grid search over CV runs
param_optimization_bayesian.py  hyperparameter Bayesian search over CV runs
true_reward_test.py         sanity check script. evaluate model training only classifiers and assuming true reward function is known.

Running

Command line prompts

pip install -r requirements.txt

# create model-ready data from SLiM simulation output
python sim_data_processor.py

# single train/test split
python main.py

# k-fold cross validation with plots 
python run_cross_validation.py

# grid-search hyperparameter optimization (edit build_sweep_configs() in sweep.py first)
python sweep.py

# bayesian hyperparameter optimization (edit suggest_config() in param_optimization_bayesiann.py to match desired values)
python param_optimization_bayesian.py 

Bash scripts are also provided for training (train_model_job.sh), cross-validation performance evaluation (cross_validation_job.sh), and Bayesian hyperparameter optimization (optimize_hyperparams_job.sh).

Best practices

Before running, edit OSTConfig.data_dir in config.py (or pass a config with a different data_dir into main()/run_cross_validation.main()) to point at your simulation_statistics.csv, extinction_by_lag.csv, model_expected_outcomes.json, and model_ready_dataset_sim1_2.json.

Output

Every run writes to output/<run_id>/, where run_id is derived from the config (hidden dim, layers, bonuses, thresholds, etc.), so different hyperparameter settings never overwrite each other:

output/<run_id>/
    config.json                    snapshot of every hyperparameter used
    models/                         binary_classifiers.pt, reward_mlp.pt, continuation_mlp.pt
    results/raw_model_output/        classifier_performance.csv, reward_performance.csv,
                                       continuation_performance.csv, raw_predictions.json
    results/visualizations/          all saved plots
    cross_validation/                 (reserved for CV-specific artifacts)

Possible modifications

• New hyperparameters: add a field to OSTConfig in config.py. It will automatically show up in config.json snapshots and can be swept in sweep.py.
• New model type: add an entry to MODEL_TYPES in model_building/factory.py with its constructor. Nowhere else needs to change.
• New plot: add a function to the relevant file in visualizations/ that takes a DataFrame (or the raw dict, for prediction plots) and returns a Figure. Wire it into run_cross_validation.py's figs dict.
• Areas for improvement ideas: making sim_data_processor.py more computationally efficient. Testing different reward and continuation value functions. Exploring new ways of estimating $P(\text{switch} | X_j)$ other than the current MLP classifier.