Reinforcement Learning for Control of Non-Markovian Cellular Population Dynamics

Response to Reviewer JJJP

We are pleased to read your positive review of our submission. We found your comments and criticisms helpful in further strengthening the manuscript. Below, we provide our responses to your specific questions. Please let us know if this addresses your concerns or you would like to discuss more. **Weaknesses:** - Great point: due to the highly non-Markovian nature of the model, we also considered the use of recurrent architectures. Interestingly, our initial experiments showed that using RNN/LSTMs actually gave the same or worse performance, but took much longer to train than vanilla DQN. - As a result, we decided to focus on using DQN and FQF given their relative simplicity and robust performance, however we do plan to further study the use of RNNs in future work, in addition to the ideas mentioned in our other responses. As time permits, we can continue to run additional experiments in the same setting with RNNs and report back if there are positive results. - As noted in other reviewer responses which raised similar concerns, we agree that another family of RL algorithms should be studied. Thus, we examined the performance of PPO and SAC in this setting. The main results of these experiments are expressed in the general comment at the top of the review thread. **Regarding feasibility:** - In the experimental setting, measuring the population growth rate of cancer or bacterial cells in vitro periodically is straightforward and performed routinely. Thus, we chose this as our experimental observable (unlike subpopulation size used in prior work, which is harder to measure). As a result, data collection for general model parameterization and testing would not be a major challenge. Since this is where our method would initially be tested, we focused on a realistic experimental setting. Similarly, administration of the drug would require coordination (and potential automation) but there are no major technical hurdles that we can foresee. Implementing precise drug administration and obtaining cell growth rates in the clinical setting at scale would be more challenging, but possible with adequate resources. - The main technical limitation in real-world application in the clinic would be obtaining growth rate measurements at short intervals from real patients. Your intuition on various $\Delta$ values driving the optimization procedure to different regimes is interesting, and merits further experimentation in future work. Interestingly, the results shown in Fig. 6 (appendix) indicate that values of $\Delta \leq 0.2$ all yield similar performances and thus may not be a significant limitation in practice. - Furthermore, as it is not practical to obtain large amounts of patient data before using the model to train from scratch, our vision would be to obtain initial experimental data to calibrate the population dynamics environment, train the agent in simulation, then finetune the model on real patient data before or during clinical implementation (essentially sim2real, as suggested). This approach is significantly less expensive (especially compared to measuring subpopulation fractions) in both the required data collection and compute, and also allows for personalized treatment models for individual patients. Thanks again for your review and we would be happy to discuss further as needed.

Response to Reviewer 9huZ

We are excited to see that you found the problem and model development exposition clear, and that you appreciate the novelty of the model and of our combination of insights from OC and RL. Furthermore, we are glad you recognize the clinical significance of the proposed problem. **Methodology:** Indeed, we agree that the primary goal of this paper was not to develop completely new deep RL techniques. Instead, the main goal was to instead apply existing techniques systematically to a novel, highly-non-Markovian environment, with significant potential real-world impact. The application of RL and OC as presented provides new insights to an otherwise intractable model. As you point out, the model itself and insights gained are part of the main contributions. Furthermore, we believe our approach of using control theory to simplify the RL problem is novel in the field, and we show that combining insights from both can be powerful. As such, we maintain that our work is a strong fit at ICLR, despite not presenting a specific new deep RL algorithm. Please see our response below regarding the timescales used in our simulations, and let us know if you would like to discuss further. **Questions:** 1. We appreciate your concern for the biological plausibility of our model in addition to considering the mathematical and computational aspects of our work. You are correct that individual mammalian cells typically divide roughly every 24 hours; but in a given population, divisions occur asynchronously, so population growth appears much more continuous. Furthermore, cell death and drug uptake can occur at any time and in many cases are not directly connected to division timing. Thus, drug pulsing can affect cell physiology on smaller timescales than the cell cycle. The main real-world bottleneck to a small $\Delta$ parameter is being able to precisely administer the drug at the appropriate time. In the future we hope to collaborate with clinical groups to validate our computational approach in real-world settings, but the focus of the current work is to first establish a computational framework. 2. Given that adaptive living systems such as cancerous or bacterial cell populations can dynamically adapt their memory capacity over time, Fig. 4 was aimed at testing the RL model’s robustness to changes in memory strength which could occur as a result of several biological mechanisms such as mutations and changes in gene expression. We have updated the manuscript to make the justification of this experiment more clear, along with additional references giving experimental justification. 3. To choose the specific method ultimately employed (FQF with NoisyNets), we swept over different algorithmic choices (with and without PER), specifically: DQN, DDQN, C51, FQF, FQF+dueling, FQF+noisy and chose the model with the absolute best evaluation performance. This model was FQF+noisy (without PER). To add clarity to this point, we are re-running several experiments with finetuned parameters to compare as time permits. However, we would like to emphasize that the main point is to achieve the best evaluation reward, measured as the lowest simultaneous cell counts across $\mu \in [0.6,\ 0.7,\ 0.8,\ 0.9,\ 1.0]$. Thank you for reviewing our work. We would be happy to continue discussing during the remaining rebuttal period.

Response to Reviewer yE9z

**Presentation:** We are glad to hear that you found the paper well-explained, well-motivated, and relevant. We also appreciate your input on how to further improve readability. As suggested, we have re-done the figure referencing (by fixing a ref to Fig 1 and moving Fig 4 further up) to improve the readability – please let us know if you see any other areas where the readability or referencing can be improved. Additionally, we added a gray dashed line indicating the optimal solution from the memoryless case in Fig 2 (right). Thank you for this suggestion, we agree this provides a better intuition of how the RL solutions compare to that of OC. **Minor Comments:** Thank you for carefully reading the manuscript to find these points. We have updated the reference to Fig 1, rather than Fig 2 as you pointed out. Regarding SGD vs. gradient ascent, we only meant SGD is performed on the parameters of $Q_\theta$ based on the Bellman residual (Eq 17). But you are correct about RL in general performing ascent on the return objective. If you think it will help to add a clarifying sentence along these lines, we can do so. **Specific Questions:** 1. Our intuition was that continuous action algorithms would have a more difficult time learning bang-bang control for two reasons: (1) The corresponding bang-bang action values are at the extreme limits of the output policies (thus they tend to saturate gradients at the output layer) and (2) Since neural net representations of the policy implicitly assume a smoothness with respect to state and in some algorithms (SAC) give bonuses for higher entropy, a deterministic policy with such extremes w.r.t. observation would be hard to learn. Based on our recent experiments with SAC and PPO, our basic intuition seems to hold up, and it appears that such algorithms do have a harder time learning in these environments. 2. We appreciate your concern for ensuring that our RL approach is robust to potentially different or unobserved biological dynamics. As summarized in the general response, our new robustness experiment shows the RL agent performs well across a range of the most important switching parameters, even when they were not seen during training. Moreover, the generalization of our bang-bang derivation to allow for unique dose-response functions for each rate ensures our use of the binary action space for the RL agent is well-founded. Taken together, these results support the notion that our discrete action RL approach will be successful even if our model does not precisely capture the true underlying biological dynamics. 3. While the objective of maximizing net death rate is indeed the one used in defining the reward function, the goal is to (equivalently) minimize or suppress the total population. To accomplish this, the agent must ensure that resistant cells are periodically driven back to the susceptible state (by withholding treatment), such that the overall population can continue to be reduced. If we mentioned at some point ‘maintaining the stability’, this is not entirely accurate, and we apologize for any confusion. 4. We foresee this being an important problem for future work, which requires more computational resources and potentially real-world data. In the future, we would like to use known techniques such as domain randomization and recurrence (to have a better internal state representation in a contextual setting) to allow a very general RL policy to be useful in deployment. To avoid overfitting, we foresee the training of an agent in simulation on a vast range over all physically plausible environments, to make the final policy robust rather than overfitting on a specific parameter setting. Thank you for your careful review, and we look forward to discussing further as needs be.

Thank you

We are glad to hear that our rebuttal addressed all of your questions and concerns. We believe that your feedback did indeed help us improve our manuscript. Please let us know if you have any other questions or concerns, and we would be happy to discuss for the remainder of the rebuttal period.

Response to Reviewer mj4C

**Presentation:** We are glad to hear that you found our submission well-written, and we appreciate your suggestions to further improve readability. We agree that the presentation of our manuscript would be strengthened by adding the structure that you proposed, in particular by better connecting the OC insights with the RL results. To this end, we have added several lines in the introduction outlining the paper’s structure with the proper emphasis you noted (end of Sec. 1). Furthermore, we believe the addition of continuous action experiments (PPO and SAC, discussed in the general response) successfully ties the two sections together, as they highlight the importance of simplifying the control problem with a discrete action space. **Experiments:** We are glad to see that you found our experiments extensive and that our work bridges some of the gaps between real-life problems and RL. As we mentioned in the general response, to further strengthen our results, we ran similar experiments with PPO and SAC. Our results indicate that these continuous action algorithms struggle to learn an optimal policy, while attempting to reach a bang-bang control. We added this figure to the paper (new Fig. 4) to further motivate the need for the discrete action space. These experiments independently validate the use of the binary action space derived by OC. Additionally, we previously included all experimental details, including how long the RL agent was trained ($3 \times 10^5$ steps) in the appendix, but we moved this to the main text for clarity. Please let us know if you believe any other experimental details should be added or moved to the main text. **Contribution:** We thoroughly agree that testing on real experimental data would strengthen the paper. However, given the high cost of such real-world data, this work instead develops a principled model (cf Sec 2.1 for references with similar models in other work) demonstrating the applicability of RL to this domain. A goal of this work is to show a strong proof of concept, which will motivate experimental works on adaptive therapy. Furthermore, through Figs. 5 and 7 we show the generalization abilities of deep RL, which offers flexibility even when real-world dynamics differ from simulations. This work paves the way for offline-online finetuning methods when real-world data becomes available in the future. **Specific Questions:** Thanks for catching this. We have now added it to the caption. Each line represents the mean over 10 trajectories, with the standard error indicated by the shaded region. As mentioned above, we just finished running additional experiments to test robustness. The results are summarized above, but please let us know if you have further questions.

Update: New Experiments

The experiments mentioned above recently finished and we are able to report that the overall message agrees with our prior intuition, summarized as: (1) Continuous action algorithms struggle to learn bang-bang control and (2) the deep RL agent is robust to changes in switching rates within an expected range of perturbation. We elaborate further on these conclusions below. We invite the reviewers to check the updated manuscript, where Figure 2 and 7 relate to points (1) and (2) respectively. **Continuous action algorithms** - We have run both PPO and SAC on a fixed memory strength ($\mu=0.9$) environment, with ~100 days of compute each in hyperparameter search, on par with the finetuning compute timescales from DDQN and FQF. - The best model from these sweeps (i.e. the checkpoint with the highest evaluation reward) was then evaluated, and the results are depicted in Fig 2. This figure illustrates an interesting phenomenon where the continuous action algorithms are able to find the “high dose” constant application period at the beginning of the episode, but as the episode continues, both PPO and SAC attempt a short-lived bang-bang policy before defaulting to the average action (0.5 dose) for the remainder. - Because of their inability to reach a bang-bang policy, PPO and SAC underperform DDQN by a significant margin. The algorithm and its corresponding total population of cells remaining at the end of the episode are DDQN (560), SAC (785), PPO (790), where lower is better. - We believe continuous action algorithms struggle because of the saturating gradient effect at the limits (bang-bang) of action space. Additionally, continuous action algorithms that give bonuses for policy entropy (e.g. SAC) make learning these deterministic policies even more difficult. **Parameter Robustness** - To test the robustness of our trained agent against perturbations or shifts in the underlying model parameters (besides $\mu$), we chose to sweep over the switching rate constants $\alpha$ and $\delta$, which determine the rate of change from the susceptible state to the resistant state and vice versa (cf. Fig 1). We chose these rates because they allow for an intuitive control over the difficulty of the control problem. For the parameter regime which we have focused on ($\kappa_S^{max},\kappa_R^{max}>0$, $\kappa_S^{min},\kappa_R^{min}<0$), population reduction is always possible. - However, the difficulty of the control problem can be tuned by changing $\alpha$ and $\delta$, thus reducing the space of pulsing frequencies which successfully lead to population reduction. For example, increasing $\alpha$ makes the control problem harder because cells more quickly transition to the resistant state, thus evading drug treatment. - We systematically swept over a grid of $\alpha$ and $\delta$ values in the range $\pm 20$ % outside of the (**fixed**) rates used for training. Despite only being trained on fixed rates, the model generalizes well across the range of values tested. - The expected increase in difficulty of control is indicated by a diagonal red line in Fig. 7 (increased $\alpha$ and decreased $\delta$) and the agent is indeed seen to struggle more in the most difficult regime. However, the agent’s performance remains nearly fixed along anti-diagonals, showing invariance to the exact rates (only changes to relative rates affect agent performance). - Beyond these new figures (now Fig’s 2 & 7), we have also updated Fig. 3 (right) with a visualization of the OC solution. We have also made other changes to the manuscript, highlighted with red text. One additional change we’d like to highlight is a generalization of the OC derivation. This generalization guarantees “bang-bang” optimal control when the dose-response functions are unique for each rate, whereas we previously assumed the same functional form throughout. We are finalizing the individual reviewer responses and they will be posted shortly. Thank you to all the reviewers for suggesting these experiments, as we believe they further enhance the main messages of the paper.