Model-based reinforcement learning for biological sequence design

The ability to design biological structures such as DNA or proteins would have considerable medical and industrial impact . Doing so presents a challenging black-box optimization problem characterized by the large-batch , low round setting due to the need for labor-intensive wet lab evaluations . In response , we propose using reinforcement learning ( RL ) based on proximal-policy optimization ( PPO ) for biological sequence design . RL provides a flexible framework for optimization generative sequence models to achieve specific criteria , such as diversity among the high-quality sequences discovered . We propose a model-based variant of PPO , DyNA PPO , to improve sample efficiency , where the policy for a new round is trained offline using a simulator fit on functional measurements from prior rounds . To accommodate the growing number of observations across rounds , the simulator model is automatically selected at each round from a pool of diverse models of varying capacity . On the tasks of designing DNA transcription factor binding sites , designing antimicrobial proteins , and optimizing the energy of Ising models based on protein structure , we find that DyNA PPO performs significantly better than existing methods in settings in which modeling is feasible , while still not performing worse in situations in which a reliable model can not be learned . 1 INTRODUCTION . Driven by real-world obstacles in health and disease requiring new drugs , treatments , and assays , the goal of biological sequence design is to identify new discrete sequences x which optimize some oracle , typically an experimentally-measured functional property f ( x ) . This is a difficult black-box optimization problem over a combinatorially large search space in which function evaluation relies on slow and expensive wet-lab experiments . The setting induces unusual constraints in black-box optimization and reinforcement learning : large synchronous batches with few rounds total . The current gold standard for biomolecular design is directed evolution , which was recently recognized with a Nobel prize ( Arnold , 1998 ) and is a form of randomized local search . Despite its impact , directed evolution is sample inefficient and relies on greedy hillclimbing to the optimal sequences . Recent work has demonstrated that machine-learning-guided optimization ( Section 3 ) can find better sequences faster . ∗Work done as an intern at Google . Reinforcement learning ( RL ) provides a flexible framework for black-box optimization that can harness modern deep generative sequence models . This paper proposes a simple method for improving the sample efficiency of policy gradient methods such as PPO ( Schulman et al. , 2017 ) for black-box optimization by using surrogate models that are trained online to approximate f ( x ) . Our method updates the policy ’ s parameters using sequences x generated by the current policy πθ ( x ) , but evaluated using a learned surrogate f ′ ( x ) , instead of the true , but unknown , oracle reward function f ( x ) . We learn the parameters of the reward model , w , simultaneously with the parameters of the policy . This is similar to other model-based RL methods , but simpler , since in the context of sequence optimization , the state-transition model is deterministic and known . Initially the learned reward model , f ′ ( x ) , is unreliable , so we rely entirely on f ( x ) to assess sequences and update the policy . This allows a graceful fallback to PPO when the model is not effective . Over time , the reward model becomes more reliable and can be used as a cheap surrogate , similar to Bayesian optimization methods ( Shahriari et al. , 2015 ) . We show empirically that cross-validation is an effective heuristic for assessing the model quality , which is simpler than the inference required by Bayesian optimization . We rigorously evaluate our method on three in-silico sequence design tasks that draw on experimental data to construct functions f ( x ) characteristic of real-world design problems : optimizing binding affinity of DNA sequences of length 8 ( search space size 48 ) ; optimizing anti-microbial peptide sequences ( search space size 2050 ) , and optimizing binary sequences where f ( x ) is defined by the energy of an Ising model for protein structure ( search space size 2050 ) . These do not rely on wet lab experiments , and thus allow for large-scale benchmarking across a range of methods . We show that our DyNA PPO method achieves higher cumulative reward for a given budget ( measured in terms of number of calls to f ( x ) ) than existing methods , such as standard PPO , various forms of the cross-entropy method , Bayesian optimization , and evolutionary search . In summary , our contributions are as follows : • We provide a model-based RL algorthm , DyNA PPO , and demonstrate its effectiveness in performing sample efficient batched black-box function optimization . • We address model bias by quantifying the reliability and automatically selecting models of appropriate complexity via cross-validation . • We propose a visitation-based exploration bonus and show that it is more effective than entropy-regularization in identifying multiple local optima . • We present a new optimization task for benchmarking methods for biological sequence design based on protein energy Ising models . 2 METHODS . Let f ( x ) be the function that we want to optimize and x ∈ V T a sequence of length T over a vocabulary V such as DNA nucleotides ( |V | = 4 ) or amino acids ( |V | = 20 ) . We assume N experimental rounds and that B sequences can be measured per round . Let Dn = { ( x , f ( x ) ) } be the data acquired in round n with |Dn| = B . For simplicity , we assume that the sequence length T is constant , but our approach based on generating sequences autoregressively easily generalizes to variable-length sequences . 2.1 MARKOV DECISION PROCESS . We formulate the design of a single sequence x as a Markov decision process M = ( S , A , p , r ) with state space S , action space A , transition function p , and reward function r. The state space S = ∪t=1 ... TV t is the set of all possible sequence prefixes and A corresponds to the vocabulary V . A sequence is generated left to right . At time step t , the state st = a0 , ... , at−1 corresponds to the t last tokens and the action at ∈ A to the next token . The transition function p ( st + 1|st ) = stat is deterministic and corresponds to appending at to st . The reward r ( st , at ) is zero except at the last step T , where it corresponds to the functional measurement f ( sT−1 ) . For generating variable-length sequences , we extend the vocabulary by a special end-of-sequence token and terminate sequence generation when this token is selected . Algorithm 1 : DyNA PPO 1 : Input : Number of experiment rounds N 2 : Input : Number of model-based training rounds M 3 : Input : Set of candidate models S = { f ′ } 4 : Input : Minimum model score τ for model-based training 5 : Input : Policy πθ with initial parameters θ 6 : for n = 1 , 2 , ... N do 7 : Collect samples Dn = { x , f ( x ) } using policy πθ 8 : Train policy πθ on Dn 9 : Fit candidate models f ′ ∈ S on ⋃n i=1Di and compute their score by cross-validation 10 : Select the subset of models S′ ⊆ S with a score ≥ τ 11 : if S ′ 6= ∅ then 12 : for m = 1 , 2 , ... M do 13 : Sample a batch of sequences x from πθ and observe the reward f ′′ ( x ) = 1|S′| ∑ f ′∈S′ f ′ ( x ) 14 : Update πθ on { x , f ′′ ( x ) } 15 : end for 16 : end if 17 : end for 2.2 POLICY OPTIMIZATION . We train a policy πθ ( at|st ) to optimize the expected sum of rewards : E [ R ( s1 : t ) |s0 , θ ] = ∑ st ∑ at πθ ( at|st ) r ( st , at ) . ( 1 ) We use proximal policy optimization ( PPO ) with KL trust-region constraint ( Schulman et al. , 2017 ) , which we have found to be more stable and sample efficient than REINFORCE ( Williams , 1992 ) . We have also considered off-policy deep Q-learning ( DQN ) ( Mnih et al. , 2015 ) , and categorical distributional deep Q-learning ( CatDQN ) ( Bellemare et al. , 2017 ) , which are in principle more sampleefficient than on-policy learning using PPO since they can reuse samples multiple times . However , they performed worse than PPO in our experiments ( Appendix C ) . We implement algorithms using the TF-Agents RL library ( Guadarrama et al. , 2018 ) . We employ autoregressive models with one fully-connected layer as policy and value networks since they are faster to train and outperformed recurrent networks in our experiments . At time step t , the network takes as input the W last characters at−W , ... , at−1 that are one-hot encoded , where the context window size W is a hyper-parameter . To provide the network with information about the current position of the context window , it also receives the time step t , which is embedded using a sinusoidal positional encoding ( Vaswani et al. , 2017 ) , and concatenated with the one-hot characters . The policy network outputs a distribution πθ ( at|st ) over next the token at . The value network V ( st ) , which approximates the expected future reward for being in state st , is used as a baseline to reduce the variance of stochastic estimates of equation 1 ( Schulman et al. , 2017 ) . 2.3 MODEL-BASED POLICY OPTIMIZATION . Model-based RL learns a model of the environment that is used as a simulator to provide additional pseudo-observations . While model-free RL has been successful in domains where interaction with the environment is cheap , such as those where the environment is defined by a software program , its high sample complexity may be unrealistic for biological sequence design . In model-based RL , the MDP M = ( S , A , p , r ) is approximated by a model M′ = ( S , A , p′ , r′ ) with the same state space S and action space A asM ( Sutton & Barto , 2018 , Ch . 8 ) . Since the transition function p is deterministic in our case , only the reward function r ( st , at ) needs to be approximated by r′ ( st , at ) . Since r ( sT , aT ) is non-zero at the last step T and then corresponds to f ( x ) with x == sT−1 , the problem reduces to approximating f ( x ) . This can be done by supervised regression by fitting a regressor f ′ ( x ) on the data ∪n′ < =nDn′ collected so far . We then use the resulting model to collect additional observations ( x , f ′ ( x ) ) and update the policy in a simulation phase , instead of only using observations ( x , f ( x ) ) from the the true environment , which are expensive to collect . We call our method DyNA PPO since it is similar to the DYNA architecture ( Sutton ( 1991 ) ; Peng et al . ( 2018 ) ) and since can be used for DNA sequence design . Model-based RL provides the promise of improved sample efficiency when the model is accurate , but it can reduce performance if insufficient data are available for training a trustworthy model . In this case , the policy is prone to exploit regions where the model is inaccurate ( Janner et al. , 2019 ) . To reap the benefit of model-based RL when the model is accurate and avoid reduced performance when it is not , we ( i ) automatically select the model from a set of candidate models of varying complexity , ( ii ) only use the selected model if it is accurate , and iii ) stop model-based training as soon the the model uncertainty increases by a certain threshold . After each round of experiment , we fit a set of candidate models on all available data to estimate f ( x ) via supervised regression . We quantify the accuracy of each candidate model by the R2 score , which we estimate by five-fold cross-validation . See Appendix G for a discussion of different data splitting strategies to select models using crossvalidation . If the R2 score of all candidate model is below a pre-specified threshold τ , we do not perform model-based training in that round . Otherwise , we build an ensemble model that includes all models with a score greater or equal than τ , and use the average prediction as reward for training the policy . We considered τ as a tunable hyper-parameter , were we found τ = 0.5 to be optimal for all problems ( see Figure 14 . By ignoring the model if it is inaccurate , we aim to prevent the policy from exploiting deficiencies of the model ( Janner et al. , 2019 ) . We perform up to M model-based optimization rounds ( see Algorithm 1 ) and stop as soon as the model uncertainty increased by a certain factor relative to the model uncertainty at the first round ( m = 1 ) . This is motivated by our observation that the model uncertainty is strongly correlated with the unknown model error , and prevents from training the policy with inaccurate model predictions ( see Figure 12 , 13 ) as soon as the model starts to explore regions on which the model was not trained on . For models , we consider nearest neighbor regression , Bayesian ridge regression , random forests , gradient boosting trees , Gaussian processes , and ensemble of deep neural networks . Within each model family , we additionally use cross-validation for tuning hyper-parameters , such as the number of trees , tree depth , kernels and kernel parameters , or the number of hidden layers and units ( see Appendix A.7 for details ) . By testing and optimizing the hyper-parameters of different models automatically , the model capacity can dynamically increase as data becomes available . In Bayesian optimization , non-parametric models such as Gaussian processes are popular regressors , and they also automatically grow model capacity as more data arrives ( Shahriari et al. , 2015 ) . However , with Bayesian optimization there is no opportunity to ignore the regressor entirely if it is unreliable . Furthermore , Bayesian optimization relies on performing ( approximate ) Bayesian inference , which in practice is sensitive to the choice of hyper-parameter ( Snoek et al. , 2012 ) . Overall , our method combines the positive attributes of both generative and discriminative approaches to sequence design . Our experiments do not compare to prior work on model-based RL , since these methods primarily focus on estimating a dynamics model for state transitions .

Summary: This paper apply a model-based RL algorithm, DyNA-PPO for designing biological sequences. By being model-based, this algorithm is sample efficiency compared to model-free RL algorithms. This advantage is attractive and important in the context of biological sequence design since the designed is constrained to be done in the large batch / low round settings. To further improves model efficiency, the authors reduce learning bias by quantifying the reliability and automatically selecting models of appropriate complexity via cross validation. To encourage diversity in the target distribution they also penalize the reward using a visitation-based strategy.