PIPS: Path Integral Stochastic Optimal Control for Path Sampling in Molecular Dynamics
Authors: Lars Holdijk
Yuanqi Du
Ferry Hooft
Priyank Jaini
Bernd Ensing
Max Welling
Venue: Submitted to ICLR 2023 (2023)

Abstract: We consider the problem of Sampling Transition Paths: Given two metastable conformational states of a molecular system, \eg\ a folded and unfolded protein, we aim to sample the most likely transition path between the two states. Sampling such a transition path is computationally expensive due to the existence of high free energy barriers between the two states. To circumvent this, previous work has focused on simplifying the trajectories to occur along specific molecular descriptors called Collective Variables (CVs). However, finding CVs is non trivial and requires chemical intuition. For larger molecules, where intuition is not sufficient, using these CV-based methods biases the transition along possibly irrelevant dimensions. In this work, we propose a method for sampling transition paths that considers the entire geometry of the molecules. We achieve this by relating the problem to recent works on the Schr\"odinger bridge problem and stochastic optimal control. Using this relation, we construct a path integral method that incorporates important characteristics of molecular systems such as second-order dynamics and invariance to rotations and translations. We demonstrate our method on commonly studied protein structures like Alanine Dipeptide, and also consider larger proteins such as Polyproline and Chignolin.

Under review as a conference paper at ICLR 2023
PIPS: PATH INTEGRAL STOCHASTIC OPTIMAL CON- TROL FOR PATH SAMPLING IN MOLECULAR DYNAMICS
Anonymous authors Paper under double-blind review
ABSTRACT
We consider the problem of Sampling Transition Paths: Given metastable conformational states molecular system, e.g. a folded and unfolded protein, we aim to sample the most likely transition path between the two states. Sampling such a transition path is computationally expensive due to the existence of high free energy barriers between the two states. To circumvent this, previous work has focused on simplifying the trajectories to occur along specific molecular descriptors called Collective Variables (CVs). However, finding CVs is non trivial and requires chemical intuition. For larger molecules, where intuition is not sufficient, using these CV-based methods biases the transition along possibly irrelevant dimensions. In this work, we propose method for sampling transition paths that considers the entire geometry of the molecules. We achieve this by relating the problem to recent works the Schrödinger bridge problem and stochastic optimal control. Using this relation, we construct a path integral method that incorporates important characteristics of molecular systems such second-order dynamics and invariance to rotations and translations. We demonstrate our method on commonly studied protein structures like Alanine Dipeptide, and also consider larger proteins such as Polyproline and Chignolin.
INTRODUCTION
Modeling non-equilibrium systems in natural sciences involves analyzing dynamical behaviour that occur with very low probability known as rare events, i.e. particular instances of the dynamical system that are atypical. The kinetics of many important molecular processes, such as phase transitions, protein folding, conformational changes, and chemical reactions, all dominated by these rare events. One way to sample these rare events is to follow the time evolution of the underlying dynamical system using Molecular Dynamic (MD) simulations until a reasonable number of events have been observed. However, this is highly inefficient computationally due to the large time-scales involved in MD simulations, which are typically related to the presence of high energy or entropy barriers between the metastable states. Thus, the main problem is: How can we efficiently sample trajectories between metastable states that give rise to these rare but interesting transition events?
Numerous enhanced sampling methods such as steered MD (Jarzynski, 1997), umbrella sampling (Tor- rie and Valleau, 1977), constrained MD (Carter al., 1989), transition path sampling (Dellago and Bolhuis, 2009), and many more, have been developed to deal with the problem of rare events in molecular simulation. Most of these methods bias the dynamical system with well-chosen geometric descriptors of the transition (analogous to lower dimensional features), called collective variables (CVs), that allow the system to overcome high-energy transition barriers and sample these rare events. The performance of these enhanced sampling techniques is critically dependent on the choice of these CVs. However, choosing appropriate CVs for all but the simplest molecular systems is fraught with difficulty, as it relies on human intuition, insights about the molecular system, and trial and error.
A key alternative to sampling these rare transition paths is to model an alternate dynamical system that allows sampling these rare trajectories in optimal manner (Ahamed et al., 2006; Jack, 2020; Todorov, 2009) or by learning an optimal RL policy for such a transition system Rose et al. (2021).
In this paper, consider the problem of sampling rare transition paths by developing alternative dynamical system using path integral stochastic optimal control (Kappen, 2005; 2007; Kappen and Ruiz, 2016; Theodorou et al., 2010). Our method models this altermative dynamics of the system

---

Under review as a conference paper at ICLR 2023
by applying externl control policy each of the atoms in the molecule. We learn the external control policy such that it minimizes the amount of external work needed to overcome the lowest energy barrier and transition the molecular system from initial meta-stable state to a final one. The method does not require any knowledge of CVs to sample these rare trajectories. Furthermore, draw connections between sampling rare transition paths and the Schrödinger bridge problem (Schrödinger, 1931; 1932). Subsequently, we show that stochastic optimal control is well suited to solving these problems by extending the work of Kappen and Ruiz (2016) for molecular systems by incorporating Hamiltonian dynamics and equivariance constraints in our path integral SOC method.
Our main contributions this paper are:
We demonstrate the equivalence between the problem of sampling transition paths, the Schrödinger bridge problem, and path integral stochastic optimal control (SOC) ($2).
We develop PIPS, path integral SOC method that incorporates second order Hamiltonian dynamics with clear physical interpretations of the system ($3).
In contrast to earlier work, PIPS does not require any knowledge CVs, which is important for modeling large and complex molecular transitions for which CVs are unknown ($2-3),
Due to considering second order Hamiltonian dynamics, PIPS seamlessly integrates with common molecular dynamics frameworks such as OpenMM (Eastman et al., 2017).
We demonstrate the efficacy of PIPS on conformational transitions in three molecular systems of varying complexity, namely Alanine Dipeptide, Polyproline, and Chignolin ($4).
PRELIMINARIES AND PROBLEM SETUP
Consider a system evolving over time where π (xx) the distribution of states æ and I. (, (2;;1)) a Markovian transition kernel. The distribution of trajectories generated by such a system is given by:
ग (x(T)) ( II Te (=
where (T) defines trajectory of states of length r discretized over time into an ordered sequence of states (x)) = (200.11, ..,..
The problem of sampling transition paths involves sampling trajectories from this distribution, T (z)()), with the boundary condition that the initial state 2y and terminal state x, are drawn from pre-specified marginal distributions Tu and Tr, respectively. These marginal distributions describe the stable states of the molecular system located at the local minimas of the free energy surface e.g. these stable states can be reactants and products of chemical reactions, native and unfolded states of protein. Thus, these marginal distributions defining the stable states can viewed as Dirac delta distributions. Unfortunately, these stable states are often separated by high free energy barriers making the trajectories, (T)), sampled starting from 29 to terminate in the target state 8, unlikely.
In this paper, we construct sampling approach that generates trajectories that still likely under the distribution т (2(T)) while also adhering to the boundary conditions by crossing the high free energy barrier by incorporating relevant inductive biases of the system. Formally, find alternate dynamical system 7 (I(T)) with marginals To and म, that is as close to π (I(T)) as possible, i.e.
# (x(T)) = arg min a ((((r)))D(moo.)))
DKL A (x((T) ([((T))](((((11) (I(())))
(2)
where D(#o, πT) is the space of path measures with marginals uo and This problem of learning an alternative dynamical system also known as the Schrödinger Bridge Problem (SBP) (Schrödinger, 1931; 1932). We, thus, take inspiration from recent computational advances for solving SBP (Vargas et al., 202 De Bortoli et al., 2021) to develop our solution in $3 to solve the problem of sampling transition paths that can efficiently cross the high free energy barriers. Additionally, in this work, propose an alternative approach to solving SBP using path integral stochastic optimal control that lends itself well to modelling the chemical nature of our problem.
In the next section, we will set the stage for this novel approach by first relating the problem of sampling transition paths a path integral stochastic optimal control problem. Subsequently, we will
2

---

Under review as a conference paper at ICLR 2023
establish equivalence between learning alternative dynamical system for sampling transition paths, Schrödinger bridge problem, and stochastic optimal control.
2.1 SAMPLING TRANSITION PATHS THROUGH STOCHASTIC OPTIMAL CONTROL
The original dynamics of the system, as given in eq. (I), can be reformulated as stochastic process:
dæ; = f t) df + G( t) de t [0,
(3)
where f. 12 x 12+ + Ed and G Hd R + dxd are deterministic functions representing the drift and volatility of the system. The stochastic process E, is a Brownian motion with variance v.
As we stated before, the system dynamics in Equation (3) is insufficient for sampling molecular transition paths as they do not adhere to the boundary conditions imposed by the problem. We, thus, add external bias potential (or control) u( I) IRd x R+ the system that pushes the molecule over the transition state barriers. We can write the dynamics of this new system as follows:
dæ. f dt G( (u(zz, 1) dt + de.),
t E [0, r]
(4)
Given a trajectory (r)) (zo. ... a E R generated through the SDE in eq. (4), we define the cost of this trajectory under control following Kappen (2007); Theodorou et al. (2010) as:
C(x(T), v.) = (x(xz)) + Σ u(x, t)a Ru(x;, t) + u(x,, t)= Re;
(5)
t=0
where yp denotes the terminal cost, À is a constant and R is the cost of taking action u in the current state and is given as weight matrix for quadratic control cost. The goal then becomes find the optimal control u" that minimizes the expected cost in Equation (5):
2
arg min E,,ee (C(x(T), Ee)]
(6)
u
where the expectation is taken over trajectories 7 sampled using the SDE under control u. Before proceeding further, couple of remarks are in order:
We note that the control u(x,,t) t) in Equation (4) does not operate directly the system dynamics but is controlled through the same control matrix G as the Brownian motion. This formulation is highly crucial for our method, PIPS, incorporate system specific second order Hamiltonian dynamics as we will show in Section 3.
Remark 1.
The last term in the cost function in eq. (5) relating the Brownian motion and the control is unusual and devoid of clear intuition. However. this term plays an important role when relating the cost to a KL--diverencce which we will establish next. Additionally, as discussed in Thijssen and Kappen (2015), the additional cost vanishes under expectation (Ter,,, 1)z Ræ,] 0) and thus, does not influence the optimal control u given by eq. (6)
Remark 2.
Relation to sampling transition paths: Interestingly, the objective in Equation (6) is exactly related to the problem of sampling transition paths given in Equation (2). As Kappen and Ruiz (2016) establish, Equation (4) defines a probability distribution Tu (x(T)) over trajectories x(T) through:
Tu (2(()) = II (a 1 Ex) t
with μ. - x= 8(),,)) + G(æ,, (u 8) di) and Σ, - G(@a, (FG
For different u, these distributions are related through the Girsanov Theorem (Cameron and Martin, 1944). As shown in appendix B, we make the common assumption that the control cost R and the variance of the Brownian motion are inversely correlated XR V, we can obtain:
Tu (x()))
1 I /bpuls
log
Σ t=
u(xi t)T Ru( + u( ( Re:
(8)
To (T))
3

---

Under review as a conference paper at ICLR 2023
where To (z())) denotes the distribution over trajectories with control i.e. u 0 (Śeq. (I)). This assumption of relating the control cost and the variance of the Brownian motion is a common trait of control problems referred to as Path Integral Stochastic Optimal Control (Kappen, 2005).
We observe that the right-hand side of eq. (8) can also be found in the definition of the control cost in eq. (5), including the additional cost term related to the Brownian noise. As Kappen and Ruiz (2016)
show, we can thus use eq. (8) to rewrite the objective in Equation (6) as:
Tu arg min (r) [ + IKLL (Tu ( To (z(()))
(9)
This objective is an approximation of the Schrodinger Bridge formulation in Equation (2) where the constraints on the marginal distributions are replaced by a regularization term in the form of the terminal cost. Therefore, when the terminal cost dominates the KL--diergence term above, it enforces the target boundary constraints of the problem. Before we discuss an algorithm to learn this optimal policy in Equation (9) next, end this part with remark:
Remark 3. This connection between the Schrödinger Bridge Problem and stochastic optimal control has been previously established (Chen al., 2016; Pavon et al., 2021). However, through the formulations in Equations (2) and (9), we also establish the equivalence between sampling transition paths, Schrödinger bridge problem, and stochastic optimal control. This allows us utilize solutions for finding the optimal control in Equation (9) for the aforementioned problems.
Optimal Control Policy: Kappen and Ruiz (2016) introduced the Path Integral Cross Entropy (PICE)
method for solving Equation (9). The PICE method derives an explicit expression for the optimal policy and distribution Tu- when x = v RR given by:
Tu- n)(x, t) Tu (Ir(T) exp(- exp(-C(x((r),)) u))
(10)
where 712, t) (r( lexp( ( (xx is the normalization constant. This establishes the optimal distribution Tu- as reweighing of any distribution induced by an arbitrary control u. Similar to importance sampling, depending on the choice of the proposal distribution wus the estimator variance can greatly differ. Thus, the objective is to find the that best approximates u².
PICE, subsequently, achieves this by minimizing the KL-divergence between the optimal controlled distribution nu- and a parameterized distribution Tug using gradient descent as follows:
(TT (uo)) 00
Ous 00
E Br exp((C((x(T), uo)) Σ 1)
(Re:
)]
(11)
Similar to the optimal control in eq. (10), the gradient used to minimize the KI KI--divergence is found by reweighing for each sampled trajectory, (TT). the gradient of the control policy uy by the cost of said trajectory. Algorithm in the appendix provides a method for finding this gradient and training the policy ug. Thus, PICE provides an iterative gradient descent method to learn a parameterized policy us and subsequently distribution over paths x(r). We can then use this learned control, us. to approximate the solution for sampling transition paths as well as the Schrödinger bridge problem.
In this section, we up our main problem of sampling transition paths and established its relationship to both the Schrödinger bridge problem and stochastic optimal control. Subsequently, we discussed an iterative gradient descent based method for solving the optimal control problem. In the next section, we will extend this iterative algorithm to consider the entire geometry of the molecular system by incorporating Hamiltonian dynamics using an augmented state space It, and symmetries by learning a policy network ug.
PATH INTEGRAL OPTIMAL CONTROL FOR SAMPLING TRANSITION PATHS
We consider a molecule consisting of n atoms with an initial and final configuration To E 33 and T, E 3x i.. are given vector defining the 3D positions of each atom in the molecule. Thus, direct method to sample transition paths r(T) for this problem is to learn a control us acting directly on the positions of the molecule using the iterative gradient descent method discussed in Section 2.
4

---

Under review as a conference paper at ICLR 2023
However, the collective behaviour of the atoms and molecules are governed by classical molecular dynamics i.e. Newtonian equations of motion:
dr = o(t) de, and, de = (()) di
(12)
where v(t) E R3x" is the velocity and (t) E R is acceleration given by a(t) = vrU( where (r) is the potential energy of the system and m is the mass. The potential energy of system is defined by a parameterized sum of pairwise empirical potential functions, such as harmonic bonds, angle potentials, inter-molecular electrostatic and Van der Waals potentials. In our work, we compute this potential energy using the OpenMM framework (Eastman et al., 2017). Therefore, in light of Equation (12), we need to adapt the dynamical system defined in Equation (4) to incorporate these molecular dynamics.
Incorporating second order dynamics: Formally, we incorporate the second order dynamics of the system defined above by considering an augmented state space: Let 2o == (To. va) E R x 3 be the initial configuration of the system defining the initial positions and velocities of each atom and Fr == (r,, 1T) be the final configuration. We, thus, model the dynamical system in Equation (4) as:
dr
Vi
dt +
O3n
(u(a,,, 1) dt - +de;),
EE [0,-]
(13)
dv
- rr, Urt))
Tan
doe
f(
G(
Due to the choice of G( t) in Equation (13) above, the additional bias force, ux t). applied to the system only influences the acceleration and velocity of the atoms and does not act directly on the positions of the atoms. dr, is solely influenced by the velocity Vr. thus conforming to the classical molecular dynamics of the system given in Equation (12).
Unfortunately, this new dynamical system in Equation (13) leads to a singular covariance matrix, Eg in eq. (7) due to the choice of G. However, due to the conditional independence of Tr+1 given (Te, ve). we are able to factorize the distribution in Equation (7) which circumvents the singularity of the covariance matrix. Due to space constraint, we provide details and derivations in Appendix B.1.
Remark 4. We note here that second order dynamics have been considered before for stochastic optimal control by Kappen (2007) for a synthetic spring experiment in one dimension and SBP by Vargas et al. (2021a) for modelling motion. Our formulation of incorporating second-order dynamics here is distinct and more practical than these previous works. Courtsey of eq. (13), we have a clear physical interpretation of the control as an external physical force by limiting it act linearly on the velocity U. This is interesting for downstream applications of the sampled transition paths such as reconstructing free-energy surfaces. Additionally, it also simplisies incorporating the control with MD) simulation software like OpenMM which we will discuss in detail in section 4.
Invariance to rotations and translations: Secondly, the molecules in consideration are invariant w.r.t. translations and 3D rotations i.e. the molecular orientations achieved along transition path need to incorporate this equivariance w.r.t. the SE(3) group. For this purpose, we need to make the terminal cost function, (x(,),, ). in Equation (5) to be equivariant. We enforce this by defining the terminal cost as the exponentiated pairwise distance between atoms which is commonly used distance metric (Shi al., 2021) that is invariant to rotations and translations i.c. pp(rt) exp Ei, (dis(re) -di)))d)rr))) (r+))
where dey (r.) l1(r,); - (re))l12-
Physics inspired policy network (ug)): The main learnable component of our PIPS method (as described by eq. (13) for sampling transition paths is the policy network ug. Following the discussion above and formalized in Equation (13), we can interpret the control ua as an additive bias force applied to the system. In this work, we consider two different design approaches to modelling ugg In our first approach, we model ug as a neural network that predicts the bias force on the system in which case the velocity evolves as du = Uv,, + us.:) dt. Alternatively, in our second approach, we model ug network predicting the bias potential energy. In this case, the corresponding force, F(r:) ,applied to the system is calculated by backpropagating through the network, F(r:)= == Vr, "g,-- The change in velocity then given by du = (V,, (r.)) + F(r:)) dt. Additionally, ue uo can be implemented using recent advances in physics inspired equivariant neural networks (Cohen and Welling, 2016; Satorras et al., 2021) that take into account the SE(3) symmetry of the system. We provide details for training the control network us in Appendix A.
5

---

Under review as a conference paper at ICLR 2023
Temp.
EPD (4)
THP (t)
ETP (4)
fs
K
nm x 10
%
kJ mol
Force Prediction
500
300
2.07
41.1 %
0.68
Energy Prediction
500
300
1.25
89.2 %
-5.21
MD w. fixed timescale
500
300
7.92
0%
500
1500
7.47
0%
500
4500
6.33
0%
500
9000
6.82
1.7
1019.83
MD w/ fixed timescale
34810
1500
1.88
100%
551.51
48683
4500
2.01
100%
1647.35
Table I: Benchmark scores for the proposed method and extended MD baselines. From-left-to-rightt Time-horizon т representing the trajectory length (note that we take one policy step every fs), simulation temperature, Expected Pairwise distance (EPD), Target Hit Percentage (THP), and Energy Transition Point (ETP). ETP can only be calculate when a trajectory reaches the target. All metrics are averaged over 1000 trajectories except for MD w/ fixed timescale which is ran only for 10 trajectories.
EXPERIMENTS
We evaluate our path integral stochastic optimal control method for sampling transition paths with three different molecular systems, namely (i) Alanine Dipeptide, small amino acid with well- studied transition paths, (ii) Polyproline, a small protein with two distinct conformations with different helix orientations, and (iii) Chignolin, an artificial mini-protein studied to understand the folding process of proteins. We begin by detailing the experimental setup below.
Molecular Dynamics Simulation: As we discussed in section 3, we use the OpenMM framework to simulate the molecular dynamics following Equation (13). Crucially, by considering the second order dynamics, the control acts linearly on the molecular potential function in this formulation of the molecular dynamics. This allows us implement the resulting control a bias potential that is acting on the system addition to the molecular potential. At every step of the Molecular Dynamics this bias potential is calculated using our Py Py'oorch implementation of the control and then passed to OpenMM as a custom external force. Implementing the control this way thus allows us to use the optimized configuration capabilities of OpenMM, such as forcefield definitions (the potential function description) and integrators (for the dineediscettzzation of our dynamics). We report the molecule specific OpenMM configuration in appendix C. Generally, we run our simulations 300 K.
Policy Network, u(x,, t): We implement the policy network as a 6 layer MLP with ReLU activation for all our experiments below. The width of the layers of the policy network is dependent on the number of atoms in the molecule under consideration. We implement all code in Pytorch. We ran the experiments on single GPU (either an NVIDIA RTX3080 or RTX2080), Our code, including full stand-alone notebook re-implemennttion, is available here: https://giie com/pips 4anonyouus /pips anonymous.
4.1 ALANINE DIPEPTIDE
Alanine Dipeptide an extensively studied molecule (Tobias and Brooks III, 1992; Rossky and Karplus, 1979; Head-Gordon et al., 1991; Swenson et al., 2018) for developing and testing enhanced sampling methods due to ready availability of its two CVs (ø, tb). The conformation transition for Alanine Dipeptide can thus be understood in terms of these two dihedral angles 6 and t as displayed in Figure IA. Prior work has, thus. focused on transforming from the initial configuration (see Figure IA) to the final configuration (Figure IE) by rotating these CVs. As we discussed previously, major advantage our method is that we do not require the knowledge of CVs sample a transition path. However, in our experiment for Alanine Dipeptide, will use these CVs to compare the quality the trajectory sampled by our method.
Setup: For our experiment, we consider both the design choices for the policy network, us(#; f). discussed in Section 3 i.e. directly predicting the force and predicting the energy. We trained the policy networks for 15,000 roll-outs with time horizon of 500 fs each consisting of 16 samples. A gradient update was made to the policy network after each roll-out with a learning rate of 10 -5 The Brownian motion has standard deviation of 0.1.

---

Under review as a conference paper at ICLR 2023
A:
B:
D:
c:
D
Figure Visualization of trajectory sampled with the proposed method. Left: The sampled trajectory projected on the free energy landscape of Alanine Dipeptide as function of two CVs Right: Conformations along the sampled trajectory: A) starting conformation showing the CV dihedral angles, B-D) intermediate conformations with C being the highest energy point on the trajectory, and E) final conformation, which closely aligns with the target conformation. Bottom: Potential energy during transition. Letters represent the same configurations in the transition.
We compare our method MD simulations with extended time- horizon and increased system temperatures to sample transition paths. To our knowledge, there are no fixed quantitative metrics in the literature to compare different methods that sample transition paths. Thus, we introduce here three metrics to evaluate the quality of transition paths: (i) Expected Pairwise Distance (EPD) measures the euclidean distance between the final conformation in the trajectory and the target conformation, reflecting the goal of the transition to end in the target state, (ii) Target Hit Percentage (THP) assures that the final configuration is also close in terms of CVs by measuring the percentage of trajectories correctly transforming these CVs, and (iii) Energy Transition Point (ETP) which evaluates the capacity of each method to find transition paths that cross the high-energy barrier low point by taking the maximum potential energy of the molecule along the trajectory. A good trajectory will be one that passes through the minimal high-energy barrier and ETP aims to measure this. We provide more details in Appendix C.2.1.
Baseline and evaluation metries:
We first visualize the trajectory generated by the energy prediction policy in Figure 1 and defer the visualization for the force prediction policy to Appendix C.2.2. The trajectory Figure demonstrates that the control policy transforms the molecule from the initial position (A) to the final position (E) by transitioning over the barrier with the least energy at (C). Interestingly, the trajectory follows the expected transitions in the CVs without them being explicitly specified c.g. the transition path visualized on the left in Figure shows that the molecule first rotates the dihedral angle associated with CV 6 in (A B), then gradually rotates along both to and ø in (B C-> D), and finally rotates &o in (D >> E) to reach the final configuration. As expected, we observe that the potential energy goes up during the transition until it reaches the top of the energy barrier (C). After this point, the molecule settles down in its new low-energy state.
Results:
Next, in Table predicting policy networks with MD simulations on the metrics introduced before. We find that the trajectories generated by both the policy networks outperform the MD baselines, but the more physics-aligned energy predicting policy performs best under our metrics. This policy network consistently reaches the target conformation both in terms of full geometry and the CVs orientation. Furthermore, our policy network generates these trajectories in significantly shorter time than temperature enhanced MD simulations without a fixed timescale. When we do limit MD to run for the same timescale as the proposed method, we found that, in contrast to the proposed method, temperature enhanced MD simulations are unable to generate successful trajectories.
compare the performance of the trajectories sampled using the force and energy
4.2 POLYPROLINE HELIX
Polyproline is helix- shaped protein structure that consists of repeating proline residues. Polyproline helix can form different conformations namely Polyproline-I (PP-I) and Polyproline-Il helix
7

---

Under review as a conference paper at ICLR 2023
Figure 2: Visualization the Polyproline transformation from PP-II to PP-L. From-top- crom-top-to-botiee tep-tt-bottom stages of the transition, s CVs, ø CVs, CVs, and Potential Energy. For the CVs multiple instances of the same dihedral angles can be found in a single molecule. Stars indicate target CV states. Colored bonds represent the bonds involved in the CV.
(PP-II) (Moradi et al., 2009; 2010). These conformations can be distinguished by their respective helix rotation. PP-I forms a compact night-hander helix due its peptide bonds having cis-isomers while PP-II has trans-isomer peptide bonds and forms a left-handed helix. Furthermore, the backbone of the polyproline helix also contains bonds and dihedral-angles as the w, & and & CVs respectively. Polyproline can have varying lengths due its repeated structure. In our experiment, consider the polyproline trimer with proline residues transitioning from PP-II to PP-I.
different dihedral angles. We will refer to these peptide
Setup: The policy network trained over 500 rollouts with 25 samples each using learning rate 3 x 10 and a standard deviation of 0.1 for the Brownian motion.
Results: We visualize the transformation of the three collective variables (w, c, vi) as well as the corresponding potential energy of the conformation in Figure 2 for a sampled transition path from our trained policy network. The CV admits the biggest change for the transition from PP-I going from 180º to 0% We observe that the transition path sampled by our method aligns with the expected changes CVs in spite of our method not containing any knowledge about these CVs. Figure shows that the peptide bonds transition from trans-isomer to a cis-isomer state at steps 450 and 3,000. We notice the biggest changes CVs at these steps in Figure 2. We also note that in addition to the change in the peptide bonds, the final conformation differs from the initial in one of the vi-dihedral angles. Technically, PP-I has diddihedral angles similar to PP-II, but as a result of the inherent noise of MD target conformation successfully learned to sample transition paths terminating in a similar perturbed state. This indicates that our proposed method is resilient target states not having minimal-energy configuration.
sampled with slight rotation here well. Interestingly.
method
4.3 CHIGNOLIN
Chignolin is small 8-hairpin protein constructed artificially study protein folding mecha- nisms (Honda et al., 2004; Seibert et al., 2005). Due to its small size, its folding process is easier to study than larger scale proteins while being similar enough to shed light on this complex process. In contrast Alanine Dipeptide and Polyproline, there is no agreement the transition mechanism de- scribing the folding Chignolin. Both the CVs involved (Satoh et al., 2006; Paissoni and Camilloni, 2021), as well as the sequence of steps (Harada and Kitao, 2011 Satoh et al., 2006; Suenaga et al., 2007; Enemark et al., 2012) describing the folding process have multiple different interpretations. Thus, methods that do not require prior knowledge of CVs particularly useful to study this protein.
Setup: We sample transition paths between the folded and unfolded state of the Chignolin protein using total time horizon of 5000 fs. Note that the typical folding time of Chignolin recorded to be
8

---

Under review as a conference paper at ICLR 2023
Figure 3: Visualization of the Chignolin folding process. Top: 5 stages of the folding process, Middle: Pairwise distance wrt to the target conformation of the molecule, Bottom: Potential Energy.
0.6 (Lindorff-Larsee) et al., 2011). The policy network is trained over 500 rollouts of 16 samples with learning rate of x 10-4 and standard deviation of 0.05 for the Brownian motion.
Results: In Figure 3, we visualize the transition of Chignolin at different timesteps during the transition path. We observe that transition the protein from its low energy unfolded state to the folded conformation, the proposed method guides the protein into a region of higher energy. This increase is initially more steep (0-+1500 than in the later stages. Additionally, most of the finer- finer-grained folding (2500-44000) occurs with high potential energy before settling into the lower- energy folded state. We notice that for the restricted folding time we use our experiments (5000 fs vs 0.6 ps), the molecule does not end at the final configuration but reaches close to it shown by the plot on pairwise distance. Furthermore, the learned policy network is able transition through the high energy transition barrier in this restricted time. We do not encounter this for molecules with shorter natural transition time (as illustrated by the potential energy Alanine Dipeptide in fig. I),
5 DISCUSSIONS, LIMITATIONS, AND FUTURE WORK
In this work, we proposed path integral stochastic optimal control method for the problem of sampling rare transition paths for molecular systems that incorporates the Hamiltonian dynamics and equivariance of the system. In passing, we showed an equivalence between the problem of sampling transition paths, stochastic optimal control, and the Schrödinger bridge problem. We empirically tested our method on three different molecular systems of varying sizes and demonstrated that it was able to sample transition paths on the full geometry of the system without biasing along CVs.
One observed limitation of the proposed method is that for molecules with long natural transition times, we observe the transitions to not converge to the configuration of minimal energy after crossing the high-energy transition barrier. We hypothesize that this is due to the method operating reduced time horizon (e.g. 5000 instead of 0.6 js in the case of Chignolin). or due to the terminal control cost function requiring the velocity to be zero the end of the transition. Nevertheless, we note that the method is successful in transitioning the molecules over the high energy barriers as exemplified by the known CVs changing appropriately.
There are many exciting directions for future work. Most importantly, we are excited to see how research from other machine learning fields can be used develop, possibly more efficient, methods for sampling trajectories between molecular conformations. Given the vast literature on Stochastic Optimal Control theory, Schrodinger Bridge samplers, and other topics such as Covariance Control (Yin et al., 2021; Hotz and Skelton, 1987) and Reinforcement Learning (Das et al., 2021) we hope that our path-integral based method can serve as a starting point for machine learning based solutions for the problem introduced in our work. Following this, these approaches and their sampled trajectories, could be used for solving related problems in chemistry. For example, our experiments showed that the molecules transitioned along the CVs correctly in spite of not having any information about the CVs. It will be interesting to see if we can infer these CVs from the learned policy and dynamics of the systems. Lastly, method can have implications for training diffusion models within fixed time-scale by additionally learning the control policy to transform one distribution into another. First explorations of this line of thinking presented in (Vargas al., 2021b; Zhang and Chen, 2021).

---

Under review as a conference paper at ICLR 2023
REFERENCES
Christopher Jarzynski. Nonequilibrium equality for free energy differences, Physical Review Letters, 78(14),,690, 1997.
Glenn Torric and John Valleau. Nonphysical sampling distributions in Monte Carlo free-energy estimation: Umbrella sampling, Journal of Computational Physics, 23(2):187-199, 1977.
EA Carter, Giovanni Ciccotti, James T Hynes, and Raymond Kapral. Constrained reaction coordinate dynamics for the simulation of rare events. Chemical Physics Letters, 156(5):472-477 1989.
Christoph Dellago and Peter G Bolhuis. Transition path sampling and other advanced simulation techniques for rare events. Advanced computer simulation approaches for soft matter sciences III, pages 167-233, 2009.
TP Imthias Ahamed, Vivek S Borkar, and S Juneja. Adaptive importance sampling technique for Markov chains using stochastic approximation. Operations Research, 54(3):489 -044, 2006.
Robert L Jack. Ergodicity and large deviations in physical systems with stochastic dynamics. The European Physical Journal B, 93(4):1-22 2020.
Emanuel Todorov. Efficient computation of optimal actions. Proceedings of the national academy of sciences, (00028):114781111 11483, 2009.
Dominic C Rose, Jamie Mair, and Juan P Garrahan. A reinforcement learning approach to rare trajectory sampling. New Journal of Physies, 23(1 J:013013, 2021. 2021.
Hilbert Kappen. Path integrals and symmetry breaking for optimal control theory. Journal of statistical mechanics: theory and experiment, 2005 1011 2005.
Hilbert J Kappen. An introduction to stochastic control theory, path integrals and reinforcement learning. In AIP conference proceedings, volume 887, pages 149-181. American Institute of Physics, 2007.
Hilbert Johan Kappen and Hans Christian Ruiz. Adaptive importance sampling for control and inference. Journal of Statistical Physics, 164(5):24441126, 2016.
Evangelos Theodorou, Jonas Buchli, and Stefan Schaal. A generalized path integral control approach to reinforcement learning. The Journal of Machine Learning Research, 11:3137-3181, 2010.
Erwin Schrödinger. Über die unkehrung der naturgesetze. Verlag der Akademie der Wissenschaften in Kommission bei Walter De Gruyter u.... 1931.
Erwin Schrödinger. Sur la théorie relativiste de l'électron et l'interprétation de la mécanique quantique. In Annales de l'institut Henri Poincaré, volume 2, pages 269-310, 1932.
Peter Fastman, Jason Swails, John D Chodera, Robert T McGibbon, Yutong Zhao, Kyle Beauchamp, Lee-Ping Wang, Andrew C Simmonett, Matthew P Harrigan, Chaya D Stern, et al. OpenMM 7: Rapid development high performance algorithms for molecular dynamics. PLoS computational biology, (((7)e0000655) 2017.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy, 23(9): 134, 2021a.
Valentin De Bortoli, James Thornoon, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. Advances Neural Information Processing Systems, 34, 2021.
Sep Thijssen and H. J. Kappen. Path integral control and state-dependent feedback. feedbaa Phys. Rev. 91:032104, Mar 2015. doi: 10. (tw.///hyseeE..1.022222 URL httpt://t/// aps org/ 91 032104.
doi 2oi/10.10 100.1113//Phssee
Robert H Cameron and William Martin. Transformations of weiner integrals under translations. Annals of Mathematics, pages 386-396, 1944.
10

---

Under review as a conference paper at ICLR 2023
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. On the relation between optimal transport and Schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, 169(2):671-66)) (-12), 2016.
Michele Pavon, Giulio Trigila, and Esteban G Tabak. The Data-Driven Schrödinger Bridge. Commu- nications on Pure and Applied Mathematics, 74(7):1545-1573, 2021.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In International Conference on Machine Learning, pages 9558 9558-9568. PMLR, 2021.
Taco Cohen and Max Welling, Group equivariant convolutional networks. In International conference on machine learning, pages 2990-2999, PMLR, 2016.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. arXiv preprint arXiv:2102.09844 2021.
Douglas J Tobias and Charles L Brooks III. Conformational equilibrium in the alanine dipeptide in the gas phase and aqueous solution: A comparison of theoretical results. The Journal of Physical Chemistry, 96(9):3864-388 3870, 1992.
Peter Rossky and Martin Karplus. Solvation. molecular dynamics study of dipeptide in water. Journal of the American Chemical Society, 101(8)19133-133, 1979.
Teresa Head- Gordon, Martin Head-Gordon, Michael J Frisch, Charles L Brooks III, and John A Pople. Theoretical study of blocked glycine and alanine peptide analogs. Journal of the American chemical sociery, 113(16):5989 5997 1991.
David WH Swenson,
Prinz, Frank Noe, John D Chodera, and Peter G Bolhuis. Open- PathSampling: A Python framework for path sampling simulations. 1. Basics, Journal of chemical theory and computation, 15(2):813-836, 2018
Mahmoud Moradi, Volodymyr Babin, Christopher Roland, Thomas A Darden, and Celeste Sagui. Conformations and free energy landscapes of polyproline peptides. Proceedings of the National Academy of Sciences, 106(49):20744-20751, 2009.
Mahmoud Moradi, Volodymyr Babin, Christopher Roland, and Celeste Sagui. A classical molecular dynamics investigation of the free energy and structure of short polyproline conformers. The Journal of chemical physics, 133(12))096614 2010.
Shinya Honda, Kazuhiko Yamasaki, Yoshito Sawada, and Hisayuki Morii. 10 residue folded peptide designed by segment statisties. Structure, 1208):1507-1518, 2004.
M Marvin Seibert, Alexandra Patriksson, Berk Hess, and David Van Der Spoel. Reproducible polypeptide folding and structure prediction using molecular dynamics simulations. Journal of molecular biology, 354(1):173-183, 2005.
Daisuke Satoh, Kentaro Shimizu, Shugo Nakamura, and Tohru Terada. Folding free-energy landscape of 10-residue mini-protein, chignolin. FEBS letters,
2006.
Cristina Paissoni and Carlo Camilloni. How determine accurate conformational ensembles by metadynamics metainference: a chignolin study case. Frontiers in molecular biosciences, 8:491, 2021.
Ryuhei Harada and Akio Kitao. Exploring the folding free energy landscape 8-hairpin minipro- tein, chignolin, using multiscale free energy landscape calculation method. The Journal of Physical Chemistry B, (1((2))8800-888) 8812, 2011.
Atsushi Suenaga, Tetsu Narumi, Noriyuki Futatsugi, Ryoko Yanai, Yousuke Ohno, Noriaki Okimoto, and Makoto Taiji. Folding dynamics of 10-residue 8-hairpin peptide chignolin. Chemistry-An Asian Journal, 2(5):591-598 2007.
Søren Enemark, Nicholas A Kurniawan, and Raj Rajagopalan. 3-Hairpin forms by rolling up from C-terminal: Topological guidance of early folding dynamics. Scientific Reports, 2(1):166, 2012.
11

---

Under review as a conference paper at ICLR 2023
Kresten Lindorff-Larsee Stefano Piana, Ron O Dror, and David E Shaw. How fast-folding proteins fold. Science, 534(605))511-522,, 2011.
Ji Yin, Zhiyuan Zhang, Evangelos Theodorou, and Panagiotis Tsiotras. Improving model predictive path integral using covariance steering, arXiv preprint arXiv:2109.22147, 2021.
Anthony Hotz and Robert E Skelton. Covariance control theory. International Journal of Control, 46 (1):13-32, 1987.
Avishek Das, Dominic C Rose, Juan Garrahan, and David T Limmer. Reinforcement learning of rare diffusive dynamics. The Journal of Chemical Physics, 155(13):134105, 2021.
Francisco Vargas, Andrius Ovsianas, David Fernandes, Mark Girolami, Neil Lawrence, and Niko- las Nüsken. Bayesian Learning via Neural Schrö odinger-Föllmer Flows. arXiv preprini arXiv:2111.10510. 2021b.
Qinsheng Zhang and Yongxin Chen. Path Integral Sampler: a stochastic control approach for sampling. arXiv preprint arXiv:2111.15141, 2021.
David A Sivak, John Chodera, and Gavin Crooks. Using nonequilibrium fluctuation theorems to understand and correct errors in equilibrium and nonequilibrium simulations of discrete Langevin dynamics. Physical Review X, 3(1):01 1007, 2013.
Kresten Lindorff-Larsen, Stefano Piana, Kim Palmo, Paul Maragakis, John L Klepeis, Ron O Dror, and David E Shaw. Improved side-chain torsion potentials for the Amber f19SB protein force field. Proteins: Structure, Function, and Bioinformatics, 1950-199,, 2010.
Uirich Essmann, Lalith Perera, Max L. Berkowitz, Tom Darden, Hsing Lee, and Lee G Pedersen. A smooth particle mesh Ewald method. The Journal of chemical physics, 103(19):8577-853,, 1995.
Ferry Hooft, Alberto Pérez de Alba Ortíz, and Bernd Ensing. Discovering collective variables of molecular transitions via genetic algorithms and neural networks. Journal chemical theory and computation, 17(4):2294-235 2306, 2021.
12

---

Under review as a conference paper at ICLR 2023
A ALGORITHMS
A.I LEARNING
Algorithm 1: Training Policy Up
Input: Po. Pri Initial and target molecular positions,
U(): Potential Energy function,
(p()): : Terminal cost,
Le (. Initial parameterized policy,
N: Number of trujectories sampled per update,
Time horizon,
Variance Brownian noise,
Control cost matrix,
μ. Learning rate,
dt: Time discretization step
while not converged do
4 Generate trajecturies with current policy
Ry:
0:
while < do
b Initialize initial trajectory state
(rn.0, Un,0, f) <- (ro. 0, 0);
while t (T/ dt) do
b Sample Brownian noise and action
En.t (0,,v));
un.t t uu((n,,,
& Update positions and velocity
Tn,++1 + Pn,t Un,t df;
Un.t11 f Un.t - (Vr... (rna)) + + un.t + En.e) dt;
end
& Determine trajectory cost and gradient
Cn (p((n))) Ei u Ru + uh Rena);
AP - exp(-C.)) 2iz 90
Oun.i
end
b Determine gradient normalization and perform policy update
17 EN
0 - 0 EN
EN A0;;
end
A.2 SAMPLING
Algorithm 2: Sampling using parameterized control ug
Input: ro: Initial molecular positions,
U(): Potential Energy function,
uo(: ) Trained parameterized policy,
T: Time horizon,
dt: Time discretization step
Initialize initial trajectory state
(r.. Cs 1) +- -(ro, 0,0);
while (r/df)
b Determine action
u. + us uo(r.,t), t);
& Update positions and velocity
re+1 rt
dt;
0++1 vt (V,UU(r) ue) de;
f 2t- t
end
13

---

Under review as a conference paper at ICLR 2023
B STOCHASTIC OPTIMAL CONTROL
In this section we expand on Section 2.1. Specifically, we expand on the derivation the Stochastic Optimal Control (SOC) objective terms of KL-divergence (appendix B.I) and the derivation of the iterative gradient descent procedure (appendix B.2), Note that the derivations presented here are rephrasing of those given in (Kappen and Ruiz, 2016) using notation similar the remainder of the paper. One difference to prior work can however be found in the relation between the distribution over uncontrolled and controlled dynamics.
Let us start by restating the objective of Path Integral Stochastic Optimal Control. Given control u and the Brownian motion e, with variance v, Equation (4) defines trajectory =<T) (Zo..... 27)
Rexd We can define the cost for said trajectory as
C(æ(T), u, u. = (v((x))) Σ Ju(er t)7 Ru( t) + u( t) Re:
(14)
where yp denotes the terminal cost, a constant and R defines a weighted control cost.
This is restatement of Equation (5), included here to make future reference easier. We note number of important observations.
1. Following eq. (4), observe that the control acts linearly the dynamics of the system.
2. The cost of the control itself is quadratic, weighted by the matrix R.
3. Under expectation the final term vanished; Eg (1 t)² REi = 0
The first two observations are what classify the current control problem in the family of Path Integral Stochastic Optimal Control (Kappen, 2005) and are requirement to be able to derive the explicit expression for the optimal control policy given in eq. (10). The third observation, while unusual in the context of SOC, needed rewrite the SOC objective terms of the KL-divergence as we will see next. Additionally, if we restate the SOC-objective
u = arg min E,,,, (C(x((T), u, Ec)]
(15)
we observe that observation shows that the additional cost does change the optimal control u². Lastly, we note that the family of Path Integral Stochastic Optimal Control problems assumes that
Rv. This assumption is needed both for rewriting the SOC objective KL-diieegeenc and to find an explicit expression for the solution.
B.1 SOC OBJECTIVE AS A EEEDIVERGENCE
As noted in the main body of the paper, an adjustment needs to be made to Equation (7) due to incorporation of the second-order dynamics. For this purpose restate eq. (13) as
ri+1
r: + (-Vr.((r)))
(u( () de:). (16)
01+1
Uit1
222
di +
Isn
f
G(
with E [0, r].
observe here that given r. and v. re+1 and we+1 are conditionally independent. As such we can derive factorized probability distribution Tue (x(T)) over trajectories xr()) as
Tu (xx(T) = क. (I(T)) - (x(T))
(17)
with
= (x(()) II (n
(18)
x (2(())
N Ip Et).
(19)
14

---

Under review as a conference paper at ICLR 2023
Here, the normal distribution describing the transition probabilities for the velocity compo-
nent similar to eq. (4) with μ, = Vs + (I s) dt + Go(ag, (u(a,, s) dt) and Σ, = (T, vy (Z., s). With fr and G, denote the components of f and G acting the velocity, respectively - (r,) and Taa-
Similarly, Equation (3) defines probability distribution no (3((T)) where now u = 0:
<o(2(()) = I (x((T) कह. (z()))
(20)
with n. (x(()) - =b (I((T) and
To (x(())
II N Ma E.)
(21)
with P. u, + fa s) df and Σ. 2.
Because we are only interested in the relation between Tu (x((T) and To (z(()))), the same analysis as (Kappen and Ruiz, 2016) applies with n. (IT()) and my (ax(7) cancelling out. Following Girsanov (Cameron and Martin, 1944), we get:
u G G.u. G.uilfe + 0z-1)
Tu (x(()))) = mo( exp
2
Ex
u G G.u
GaulGrl(us + +x))
=o( exp
2
Ex
u G. Gru ui G. Gi&:
=m(( exp
+
2
Ex
E,
1=0
= mo(x(T)) exp
u. G. G + u N.G.E,
TT(x(TT) exp
ui
T 1 Ut + ui v v -1)
1
mo(x(T))exx
Σ 1=0
' in
Ru, + ut Re:
(22)
2
where we use the assumption x = Ry in the last step. We use shorthand notation to simplify
W u(x,, t), G. = Go t), and f. f f.(x,,t) t). From here we can obtain the relation in Equation (8).
Now, as again show in (Kappen and Ruiz, 2016), can use this relation to rewrite the cost in eq. (14)
as
Tu (II(T)
C(æ(T), sx) - (((())) + log
(23)
No (z)
and thus, the distribution over trajectories under optimal control u² can now be defined as
Tu- arg min Ex) [C(x(T), u. Et)
1 (F(0)) + +log arg min (T(
Tu (x(())))
To (x((T)
Tu
Tu (u(x(())
arg min Ree(r) ~u T ( + + Ber 1og
To (x((T)
au
arg min Eg E [((
+ IDKL The (2((T)) || (x(())))
(24)
Th
This objective is an approximation of the Schrodinger Bridge formulation Equation (2) where the constraints the marginal distributions replaced by regularization term in the form of the terminal cost. When the expected terminal cost dominates the KL-divergence the found distribution should be similar.
15

---

Under review as a conference paper at ICLR 2023
B.2 ITERATIVE GRADIENT DESCENT
As mentioned earlier, the specific control problem we are considering here (linear acting control and weighted quadratic control cost) is known as Path Integral Stochastic Optimal Control. Work on this control problem has established that under the additional assumption that À = Rv there exists an explicit solution describing the optimal control u². While there are number of different papers establishing this result (Kappen, 2005; Theodorou et al., 2010; Kappen, 2007), we note that (Kappen and Ruiz, 2016) is most in line with our work. As such, we refer the interested reader to this work to find the proof for the following statement that defines the distribution over optimal trajectories as reweighing of the distributions over trajectories under no control:
T = To (z())) exp(-
A ((xx))
(25)
where 1) Ex cxp( X (xx))] is the normalization constant. Given the previously established relation (eq. (8)) between To and Tu. we can equivalently express the optimal control u" in terms of any control u using importance sampling
1 no (I((T)
Tu (I(()) exp( 1 ((((e)))
Tu
1) Tu (I(T))
1
1
Tu (x((T) exp(- (x((==)))))))
" exp
Ei=o 2 uu Ru + uf REi)
{
1
nu (z( expl v -
1
1
Σ
7
1
u Ru: un RE,)
T
7)
2
1=0
Tu (z(()) exp(- C((((T)) u, E)))
(26)
7
With an explicit expression for the optimal control policy given, the PICE method aims to find parameterized distribution Tug- that is close the optimal control in terms of KL-divergence
Tug arg min DK (TTu (x(()) Tu (T)
(27)
Using the explicit expression for the optimal control, the KL-divergence is given follows:
DKL Tu- (I(()) || Tep (I()))
x - Eu,,.
[log Tus
1
Eu (t) G. G:us(t)
G.uo(t) + 2x x1-1)
log log -E
exp
2
E,
1 u00)
G. Gruo(!) Grua(t) (1) (((++++))) +.)))
- EE,,-
2
Σ,
t=1
Em...
(t) Rus(t)= - uo( (uu")) ut( ((tT)Ree Re:
2
1=0
1 Enr [e- (((())
us un(")) Rus(!) - us (1) [u((() - us un(") Re: (28)
2
We use shorthand notation to simplify u(t) u (1, t), G. G( t), and f. f(x,, t). Line 1 we discard the constant term E.,. log u--] Line 2 we make use of the established relation between Tu and To for any control Line 3 we discard the constant term E,,. log тo] and note that the expectation is over trajectories sampled from म... Line 4 we rewrite using the assumption that Rv. Line we use Equation (26) rewrite the distribution over arbitrary control u using.
16

---

Under review as a conference paper at ICLR 2023
We minimize this explicit expression using Gradient Descent, to do this, derive the gradient of the KL-divergence
0D (T (I(()) (20(T))
00
- 1i le- (((((())
Dre (t)
RRu - Ru(t) Re;)
(29)
00
t=0
Finally, we note that the expectation is over trajectories of any distribution Tu. and as such, this distribution can also be chosen to be equal to the parameterized distribution Tu Tuo This gives us the final gradient
0DK Tu- (ar(T)) Teee (2r(T))
00
Esuu
[e -((((r)).uue))
(u) - -Ru (t) REi)
Sur (4)
(30)
1)
00
t=
1 Exoo (-C(0(7))uu.sz))
@uo(!)
Rec
(31)
00
C EXTENSION EXPERIMENTAL SECTION
C.I
OPENMM
General setup: We use the Velocity Verlet with Velocity Randomization (VVVR) integrator (Sivak al., 2013) within OpenMM temperature of 300 K with collision rate of 1.0 1.0
Alanine Dipeptide: We use the amber 99sb-ildn force field (Lindorff-Larre et al., 2010) without any solvent, time-step of 1.0 fs for the VVVR integrator and a cutoff of 1 nm for the Particle Mesh Ewald (PME) method (Essmann et al., 1995).
Polyproline Helix: We initialize OpeNMM with the amber protein. ff44BBolyssscc forcefield and gbn2 as the implicit solvent forcefield. The VVVR integrator had timestep of 2.0 fs and cutoff of 5 nm for PME. The proposed method was ran for a total 10.000 fs (resulting in 5,000 policy steps).
To sample transition paths between the folded and unfolded state of the Chignolin protein, we initialize OpenMM using the same forcefield and VVVR integrator as for Polyproline with the exception that we sample a new force from our policy network every 1.0 fs.. We do this 5000 times for each rollout for a total time horizon of 5000 fs.
Chignolin:
C.2 ALANINE DIPEPTIDE
C.2.1 DISCUSSION BASELINES AND EVALUATION METRICS
Metrics Three different metrics are used for the comparison covering multiple desiderata for the sampled transition trajectories. For each metric we report the score over 1000 trajectories with exception of the Molecular Dynamics without fixed timescale baseline which is only ran until 10 trajectories are successfully generated.
Expected Pairwise Distance (EPD) The EPD measures the similarity between the final conformation in the trajectory and the target conformation taking into account the full 3D geometry of the molecule. Note that the expected pairwise distance for uncontrolled MD with the target as the starting conforma- tion has a EPD of 2.25 10-3 3 All trajectories with an EPD of less than this can thus be considered to transition the molecule within standard deviation of the target distribution.
Target Hit Percentage (THP): The second metric under which we evaluate the proposed Transition Path Sampler measures the similarity of the final and target conformation in terms of the collective
17

---

Under review as a conference paper at ICLR 2023
A:
B:
C:
Figure 4: Visualization a trajectory sampled with the proposed force prediction method. Left: The sampled trajectory projected on the free energy landscape of Alanine Dipeptide as function of two CVs Right: Conformations along the sampled trajectory: A) starting conformation showing the CV dihedral angles, B-D) intermediate conformations with D being the highest energy point the trajectory, and E) final conformation, which closely aligns with the target conformation. Bottom: Potential energy during transition. Letters represent the same configurations in the transition.
variables. The THP measures the percentage of generated trajectories/paths that reach the target state. As such, higher hit percentages are preferred. We determine a trajectory to have hit the target in CV space when e and both within 0.75 of the target.
Energy Transition Point (ETP): The final metric looks at the potential energy of the transition point the conformation in the trajectory with the highest potential energy. This directly evaluates the capability of the method to find the transition path that crosses the boundary at the lowest saddle point.
Baselines We compare the proposed Transition Path Sampling method with extended Molecular Dynamics simulation using different time-scales and temperature points. As discussed earlier, there currently no other methods available for Transition Path Sampling using the full 3D geometry of the molecules.
Molecular Dynamics with fixed timescale: This set of baselines is limited to the same timescale as the proposed Transition Path Sampler, 500 femtoseconds, but uses varying temperatures. With higher temperatures should have higher probability of crossing the barrier and hitting the target configuration.
Molecule Dynamics without fixed timescales: In contrast to the other set of baselines, the MD simu- lation for this set is not limited to 500 femtoseconds, but is instead ran until the target conformation reached. We consider a trajectory to have reached its target if the following been met: 1) the current conformation classifies as having hit the target under the conditions of the metric described above and 2) the current conformation is within one standard deviation of the target distributions mean.
conditions have
By running the MD simulations until the target reached we aim to gain intuition into the speed-up that it achieved by the fixed timescale of the proposed Transition Path Sampler.
C.2.2 ADDITIONAL RESULTS: VISUALIZATION FORCE PREDICTION
We observe that the force predicting policy has learned a different trajectory then the energy predicting model presented in the main body of the paper. While different, both of the trajectories pass the high energy barrier in a locally low point. Previous work on finding transition path has also observed that multiple viable paths be found for Alanine Dipeptide (Hooft al., 2021).
18


Reviews:
[{"review_id": "qhnJyVtZsK_", "reviewer": "Reviewer_RDYQ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a method for finding the most likely transition between two states in high-dimensional configuration spaces (for chemical compounds, typically). It does so on building on recent advances in the Schrödinger bridge problem (SBP) and stochastic optimal control (SOC).\n\nIt claims to show equivalence between sampling transition paths (the most likely one) and these two other equivalent problems (SBP and SOC).\n\nThe main novelty is the use of second order Hamiltonian dynamics, i.e. sec 3, in particular Eq (13) (which seems to me like a rather classic trick to assemble position and velocity to make a Dynamical System become Markov-like, under time discretization).\n\nAnother novelty, possibly (I am not so aware of the literature) is the very use of SOC / RL-style approaches to transition path sampling. But, this does not seem to be altogether new: people are discussing the use of RL for path sampling, which is very close to the proposed approach (apart from finite vs infinite time horizon).\n\nThe proposed approach is tested on 3 classic transition-paths small molecules, with some success (but no comparison with other similar approaches, as Reinforcement Learning).", "summary_of_the_review": "Overall, this paper is rather sound and rather clear, but lacks novelty, by its own account, and additionally does not focus on ML at all.\n\nHence, I lean to reject, quite strongly, until further arguments are given about the originality of the approach and its merit for the community.", "strength_and_weaknesses": "Strengths:\n- the notations are overall clear, the content is understandable and concisely reported.\n- the experiments are abundantly commented (in terms of commenting the results)\n\nWeaknesses:\n- Machine Learning, or Learning Representation, seems quite lacking from the paper. Neural networks are mentioned quickly, only on page 5. \n- mostly, the lack of novelty, in several points. I detail here:\n- I recall that Eq (2) is the Schrödinger Bridge Problem (SBP). Then Eq (9), which is claimed to establish equivalence between sampling transition paths, Schrödinger bridge problem, and stochastic optimal control, however, has the term of terminal cost inside it, making it different from (2). Paper says that:\n     > Therefore, when the terminal cost dominates the KL-divergence term above, it enforces the target boundary constraints of the problem.\n\n     But in practice, in most problems of finding good transition paths, the initial and final configurations' energies are quite similar, and it's not obvious that the \"terminal cost\" dominates (I realized afterwards I probably inferred incorrectly what is the terminal cost.. it should be defined early and intuitively in the paper. If its expression depends on how we parameterize the configuration... I do not understand well how it can relate with weights of paths\n\n- algorithmic novelty: here again, PIPS does not seem to significantly differ from PICE (Hilbert Johan Kappen and Hans Christian Ruiz. Adaptive importance sampling for control and inference. Journal of Statistical Physics, 162(5):1244–1266, 2016). Actually the code submitted contains the file PICE.py, which performs Algorithm A.1 (page 13). Indeed, the authors are very honest about it:\n     > we note that (Kappen and Ruiz, 2016) is most in line with our work\n- on the theory side, the paper seems to be strongly inspired (and it is not hidden at all, on the contrary, which is very much honorable) by works from Kappen:  Kappen, 2005; 2007; Kappen and Ruiz, 2016. All equations/theory/results stated up to end of sec. 2 are attributed to other papers, and mostly these 2 papers.\n- On the interpretation side, similarly, Das 2021 and Rose 2021 (including Garrahan) are cited, but not exploited. By very quickly looking at these references, it seems like a good deal of the intuitive physics (that is I think only partially explained in the paper), are explained over there. For instance, papers says:\n   > Remark 2. The last term in the cost function in eq. (5) relating the Brownian motion and the control is unusual and devoid of a clear intuition.\n    I think intuition comes from considering, where could this external force we apply to the molecule, come from, in the original (no external force) setup ? I think that the answer is the noise: it can \"conspire\" to correlate to be equal or close to the control policy, .. if one is very lucky (hence the rare transitions are rare !)\n- as I said in the summary, the main (only?) novelty is the use of second order Hamiltonian dynamics, i.e. sec 3, which is in 1 page long, in particular Eq (13). However, this seems to me like a rather classic trick to assemble position and velocity, in order to make a Dynamical System become Markov-like, under time discretization. However, I recognize that if it's the needed step to apply new techniques (SOC, SBP, PICE) to transition path sampling, well then, good job. (this could become a strength if I am convinced this is very much new and insightful compared to previous works).\n- even sec. 3 is inspired from refs:\n     > We, thus, take inspiration from recent computational advances for solving SBP (Vargas et al., 2021a; De Bortoli et al., 2021) to develop our solution in §3 to solve the problem of sampling transition paths that can efficiently cross the high free energy barriers\n\n    but I did not have time to check these references.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "op4RAuaOshh", "reviewer": "Reviewer_XQgJ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a control algorithm to explore transition paths between reactant and product basins in molecular dynamics simulations. The work is strictly based on the theory of Knappen and Ruiz (2016), and implements it for molecular systems.\n", "summary_of_the_review": "The approach seems to generate a reasonable path in energy landscape. The proposed approach can be extraordinarily aggressive and may not work well with temperature or to produce potential of mean force. It may be useful for people trying to find transition paths but will likely require downstream refinement to be practically useful.  ", "strength_and_weaknesses": "\nThe loss fucntion that is optimized is a combination of two:\n\nDistance between desired and actual descriptors at the end of the trajectory: This part of the loss is unarguable. The trajectories minimizing this part are obviously transition paths.\n\nA penalty term for the magnitude of the control: This is less justified. \nThere is a very strong assumption underneath that the weight of the penalty is directly proportional to the noise. In a theoretical investigation of control this is fine. In a case when we one does not seek physically ideal bias and control strength e.g. when first part of the loss is much more important that the second one, this is also fine. But in case of transition paths and potentially free energy estimates - as proposed in the paper outline - such simplification does not stand and is likely too simple to produce realiable estimates of the potential of the means force.\n\nThe gradient update is the weighted average of the trajectories, reinforcing those that had low loss and forgetting the one with high loss. This will always produce a path connecting A and B. But in what sense is this a transition path? Are tehre guarantees that the underlying physics are respected? What would be the result if one optimized just MSE of the (actual - desired) descriptors in the final point?\n\nAppendix A1 There is possibly a typo in \\Delta \\Theta_n -> .... line. Is the exponential weight with the control gradient or added to?\n\nThe authors claim they cannot compare the method to any other. I think the method potentially could be compared to differentiable simulation of protein folding, using codes like TorchMD or JAXMD.\n\nThe simulations construct what are essentially transition paths, but are not called so. If this method produces better TS paths then the second part of the loss helps in some way.\n\nIt might be easier then to explain the difference between figure 1 and 4 if they were plotted together. Illustrative molecules can be omitted for on of the results as there is not much insight gained from them being just printed.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "YR798Br1C7p", "reviewer": "Reviewer_5HjD", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work considers the problem of sampling transition paths between two metastable states of a molecular system. This problem is difficult as the energy barrier between the states may be large, making it computationally expensive for traditional MD simulation. This work first relates this problem to literature on optimal control and Schrodinger bridges then provides a parameterization in terms of second-order dynamics to learn likely transition paths. This method is demonstrated on three systems using two slightly different parameterizations, either directly modelling the control force, or indirectly modelling the force as $\\nabla E$ and directly parameterizing energy. ", "summary_of_the_review": "I found this paper interesting and timely with the increased interest in diffusion-based models. For me I would like to see more experimental validation and rigor in terms of quantitative comparisons and experimental setup. ", "strength_and_weaknesses": "Strengths:\n\n- Interesting application of optimal control\n- Neat parameterization in terms of second-order dynamics that allows easy integration with OpenMM.\n- Well explained experiments which are understandable by non MD experts\n\nWeaknesses:\n\n- Limited theoretical novelty as the theory is established in prior work.\n- Limited quantitative experiments on a single (simple) molecule. As a non-expert in MD I can’t evaluate how often good CVs are available and how much they help in this path sampling problem. Do the authors think an experiment with MD using knowledge of the CVs be useful? It would help me understand how this baseline method works. Presumably when the CVs are less accurate (such as in chignolin?) CV-based optimizations would not work as well as PIPS? Could this be confirmed experimentally?\n- Another concern here is overfitting. How are we sure that the parameters are not simply overfit to each molecule making it difficult to apply to new molecule simulations? It would be helpful to state how hyperparameters were tuned and on what data.\n\nComments:\n\nI don’t think I understand remark 3. Is this a claim that equations (2) and (9) are the first to establish the equivalence between sampling transition paths and the SB problem which hasn’t been previously established? I don’t understand how these equations are meant to establish this and would be very surprised if this is the first time this equivalence has been established. \n\nThe notation in the “Physics inspired policy network” is a bit confusing to me with bold $u_\\theta$ and regular $u_\\theta$ having very different (and unspecified) dimensions. Instead of defining F(r_t) could we say instead that bold $u_\\theta$ is either directly parameterized (force prediction) or parameterized as $\\nabla_{r_t} E_\\theta$ (energy prediction)?\n\nI understand the current evaluation is on the quality of the paths. I would also be curious as to the quality in terms of the optimization, i.e. what is the value of $E_{\\tau, \\epsilon_t} C(x(\\tau), u, \\epsilon_t)$ for the force vs. energy models? \n\n$u_\\theta$ with some invariances is mentioned, but the end network is an MLP. Was a more sophisticated network tried? Perhaps it is not helpful for these single-molecule studies.\n\nQuestion: How does Temperature come in to play? Its slightly odd to me that MD without a fixed timescale is evaluated on a different temperature than PIPS. Could PIPS be evaluated on the same temperatures or could the authors explain why 300K is “better”?\n\nMinor remarks:\n\nRemark 3? is this a novelty claim?\n\nmisplaced comma last paragraph of page 5.\n\nCourtsey —> Curtesy Remark 4\n\nIt would be useful to mention Table 1 is on Alanine Dipeptide in the caption. Also it slightly strange that it comes before figure 1, but is referenced afterwards.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject