SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \underset {b} {\arg \min } \mathbb {E} _ {\boldsymbol {x} _ {0: \tau} \sim \pi^ {b}} [ C (\boldsymbol {x} _ {0: \tau}) ] = \underset {b} {\arg \min } \mathbb {E} _ {\boldsymbol {x} _ {0: \tau} \sim \pi^ {b}} \left[ \frac {1}{\lambda} \left(\varphi (\boldsymbol {x} _ {\tau})\right) + \log \frac {\pi^ {b} \left(\boldsymbol {x} _ {0 : \tau}\right)}{\pi^ {0} \left(\boldsymbol {x} _ {0 : \tau}\right)} \right] (24) \\ = \arg \min  _ {b} \mathbb {E} _ {\boldsymbol {x} _ {0: \tau} \sim \pi^ {b}} \left[ - \log \left(\mathbf {1} _ {P} (\boldsymbol {r} _ {\tau})\right) + \log \frac {\pi^ {b} \left(\boldsymbol {x} _ {0 : \tau}\right)}{\pi^ {0} \left(\boldsymbol {x} _ {0 : \tau}\right)} \right] (25) \\ = \arg \min  _ {b} \mathbb {E} _ {\boldsymbol {x} _ {0: \tau} \sim \pi^ {b}} \left[ \log \frac {\pi^ {b} \left(\boldsymbol {x} _ {0 : \tau}\right)}{\mathbf {1} _ {R} \left(\boldsymbol {r} _ {\tau}\right) \cdot \pi \left(\boldsymbol {x} _ {0 : \tau}\right) \cdot \mathbf {1} _ {P} \left(\boldsymbol {r} _ {\tau}\right)} \right] (26) \\ = \underset {b} {\arg \min } \mathbb {E} _ {\boldsymbol {x} _ {0: \tau} \sim \pi^ {b}} \left[ \log \frac {\pi^ {b} \left(\boldsymbol {x} _ {0 : \tau}\right)}{\pi^ {*} \left(\boldsymbol {x} _ {0 : \tau}\right)} \right] (27) \\ = \underset {b} {\arg \min } \mathbb {D} _ {\mathrm {K L}} \left(\pi^ {b} \mid \pi^ {*}\right) (28) \\ \end{array}
$$
where  $\pi^{*}$  is the TP distribution as defined in definition 1.
Algorithm 1: Training Policy  $\pmb{u}_{\theta}$
Input:  $r_0, r_T$ : Initial and target molecular positions,  $U(\cdot)$ : Potential Energy function,  $\gamma$ : Langevin Friction,  $\varphi(\cdot)$ : Terminal cost,  $u_\theta(\cdot,\cdot)$ : Initial parameterized policy,  $N$ : Number of trajectories sampled per update,  $\tau$ : Time horizon,  $\nu$ : Variance of Brownian noise,  $R$ : Control cost matrix,  $\mu$ : Learning rate, dt: Time discretization step
while not converged do  $\begin{array}{rl} & {\triangleright   \mathrm{Generate~trajectories~with~current~policy}  \boldsymbol{u}_{\theta}}\\ & {\lambda \leftarrow \boldsymbol {R}\nu ;}\\ & {n\leftarrow 0;} \end{array}$
while  $n <   N$  do  $\triangleright$  Initialize initial trajectory state  $(\pmb {r}_{n,0},\pmb{v}_{n,0},t)\gets (\pmb {r}_0,\pmb {0},0);$  while  $t <   (\tau /\mathrm{dt})$  do Sample Brownian noise and action  $\varepsilon_{n,t}\sim \mathcal{N}(0,\sqrt{2m\gamma k_B T});$ $\hat{\varepsilon}_{n,t}\sim \mathcal{N}(0,\nu);$ $\pmb{u}_{n,t}\gets \pmb{u}_{\theta}(\pmb{r}_{n,t},t);$  Update positions and velocity  $\pmb{r}_{n,t + 1}\leftarrow \pmb{r}_{n,t} + \pmb{v}_{n,t}\cdot \mathrm{d}t;$ $\pmb{v}_{n,t + 1}\leftarrow \pmb{v}_{n,t} + \left(\frac{-\nabla_{\pmb{r}}U(\pmb{r})}{\pmb{m}} +\pmb{u}_{n,t} - \gamma \pmb {v} + \varepsilon_{n,t} + \hat{\varepsilon}_{n,t}\right)\cdot \mathrm{d}t;$ $t\gets t + 1$  end Determine trajectory cost and gradient  $C_n\gets \frac{1}{\lambda} (\varphi (\pmb {r}_n,\tau) + \sum_{i = 0}^{\tau}\pmb{u}_{n,i}^T\pmb {R}\pmb{u}_{n,i} + \pmb{u}_{n,i}^T\pmb {R}\varepsilon_{n,i});$ $\Delta \theta_{n}\gets \exp (-C_{n}) + \sum_{i = 0}^{\tau}\frac{\partial\pmb{u}_{n,i}}{\partial\theta}\pmb{R}\varepsilon_{n,i};$ $n\gets n + 1$
end
end  $\triangleright$  Determine gradient normalization and perform policy update  $\eta \leftarrow \sum_{i = 0}^{N}\exp (-C_{i})$ $\theta \gets \theta +\frac{\mu}{\eta}\sum_{i = 0}^{N}\Delta \theta_{i};$
C Extension Experimental section
C.1 OpenMM
General setup: We use the Velocity Verlet with Velocity Randomization (VVVR) integrator [Sivak et al., 2013] within OpenMM at a temperature of  $300\mathrm{K}$  with a collision rate of  $1.0~\mathrm{ps}^{-1}$ . All code is implemented in Pytorch and ran on a single GPU (either an NVIDIA RTX3080 or RTX2080).
Alanine Dipeptide: We use the amber 99sb-ildn force field [Lindorff-Larsen et al., 2010] without any solvent, a time-step of 1.0 fs for theVVVR integrator and a cutoff of  $1\mathrm{nm}$  for the Particle Mesh Ewald (PME) method [Essmann et al., 1995]. The policy network for 15000 roll-outs with a 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 a standard deviation of 0.1.
Polyproline Helix: We initialize OpenMM with the amber protein.  $\text{ff14SBonlysc}$  forcefield and gbn2 as the implicit solvent forcefield. The VVVR integrator had a timestep of 2.0 fs and a cutoff of 5 nm for PME. The proposed method was ran for a total of 10,000 fs (resulting in 5,000 policy steps). The policy networks was trained over 500 rollouts with 25 samples each using a learning rate of  $3 \times 10^{-5}$  and a standard deviation of 0.1 for the Brownian motion.
Chignolin: 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. The policy network is trained for 500 roll-outs of 16 samples with a learning rate of  $10^{-4}$  and a standard deviation of 0.05 for the Brownian motion.
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 the 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 conformation has a EPD of  $2.25 \times 10^{-3}$ . All trajectories with an EPD of less than this can thus be considered to transition the molecule within one 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 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  $\phi$  and  $\psi$  are 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 are 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
B:
C:
Figure 5: Visualization of a trajectory sampled with the proposed force prediction method. Left: The sampled trajectory projected on the free energy landscape of Alanine Dipeptide as a 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 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.
higher temperatures we should have a 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 simulation for this set is not limited to 500 femtoseconds, but is instead ran until the target conformation is reached. We consider a trajectory to have reached its target if the following two conditions have 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.
By running the MD simulations until the target is 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 can be found for Alanine Dipeptide [Hooft et al., 2021].