Stochastic Control Policies for Robust Molecular Transition Path Sampling

Jingqian Liu; Yu-Hsiang Wang; Yanru Qu; Ge Liu

cs.LGAug 13, 2026

Stochastic Control Policies for Robust Molecular Transition Path Sampling

Authors: Jingqian Liu, Yu-Hsiang Wang, Yanru Qu, Ge Liu

Organizations: Siebel School of Computing and Data Science, University of Illinois Urbana–Champaign · Department of Electrical and Computer Engineering, University of Illinois Urbana–Champaign

Abstract

Transition path sampling (TPS) aims to efficiently generate rare molecular transition trajectories between metastable states and is essential for understanding biomolecular mechanisms. Beyond traditional molecular dynamics (MD)-based sampling, machine learning has become central to state-of-the-art TPS. One major class of methods learns control forces during explicit MD rollouts. By preserving the underlying molecular dynamics, these methods tend to produce more physically plausible trajectories than endpoint-conditioned generators that construct paths directly. However, rollout-based control methods have been reported to exhibit unstable and strongly seed-dependent performance. We recast rollout-based control as learning a path-space proposal distribution and investigate stochasticity placement as a design choice for improving exploration and optimization robustness. We develop two stochastic policies: FS-TPS, which directly parameterizes a state-dependent Gaussian distribution over the control policy output, and LaS-TPS, which samples a compact latent control variable and decodes it into structured, cross-atom-correlated force variation. We conduct extensive multi-seed experiments on three biomolecular systems of increasing size: alanine dipeptide, chignolin, and BBL, a fast-folding protein. Stochastic policies consistently improve transition success and path quality over deterministic-policy baselines while substantially reducing sensitivity to random initialization.

Figures & tables

Figure 1: Training and inference performance of TPS-DPS, FS-TPS, and LaS-TPS on alanine dipeptide transition path sampling, averaged over nine random seeds. Top: RMSD, THP, and ETS as a function of training rollout, showing the evolution of each metric during training. Bottom: Final RMSD, THP, and ETS at inference, with individual seed results overlaid as points and error bars denoting standard deviation across seeds.

Figure 2: Comparison of sampled transition paths on the alanine dipeptide free-energy surface for TPS-DPS, FS-TPS, and LaS-TPS for a policy trained starting from the same random seed (72831). The stars show the two saddle regions.

Method	RMSD (Å $\downarrow$ )	THP (% $\uparrow$ )	ETS (kJ/mol $\downarrow$ )
Deterministic force control
TPS-DPS	$0.43\pm 0.22$	$45.67\pm 28.26$	$29.00\pm 11.75$
Force-space stochastic control
FS-TPS ( $\sigma$ -scale $=1.0$ )	$0.26\pm 0.05$	$69.79\pm 8.84$	$29.64\pm 9.76$
$\sigma$ -scale $=0$ (mean force)	$0.25\pm 0.04$	$68.16\pm 12.16$	$31.64\pm 10.36$
$\sigma$ -scale $=0.5$	$0.26\pm 0.06$	$67.97\pm 12.19$	$31.35\pm 8.01$

Table 1: Ablation study of the proposed stochastic biasing policies on alanine dipeptide. Unless otherwise indicated, results are reported as mean

\pm

standard deviation across the same nine random seeds. ETS is computed over runs with valid transition paths. Arrows indicate the preferred direction of each metric.

Method	Internal $r_{\mathrm{eff}}$	Internal EV@1	Excess $R_{\mathrm{cross}}$	RMS noise/mean	$r_{\mathrm{eff}}(z)/d_{z}$	Global $r_{\mathrm{eff}}$
Alanine dipeptide
FS-TPS	$52.56\pm 3.00$	$2.82\%\pm 0.16\%$	$0.00\pm 0.00$	$0.20\pm 0.05$	/	/
LaS-TPS	$1.20\pm 0.04$	$95.86\%\pm 1.02\%$	$0.58\pm 0.02$	$0.40\pm 0.20$	$0.98\pm 0.00$ ( $d_{z}=10$ )	$1.07\pm 0.02$
Chignolin
FS-TPS	$30.23\pm 5.12$	$12.30\%\pm 2.19\%$	$0.00\pm 0.00$	$0.05\pm 0.02$	/	/
LaS-TPS	$1.00\pm 0.00$	$100.00\%\pm 0.00\%$	$0.75\pm 0.01$	$1.34\pm 0.06$	$0.99\pm 0.00$ ( $d_{z}=24$ )	$14.18\pm 1.33$

Table 2: The internal force covariance

\Sigma_{\mathrm{int}}(x)

is characterized by its effective rank, top-mode explained variance (EV@1), excess cross-atom correlation relative to an atom-permuted null, and RMS stochastic amplitude relative to the mean force. For LaS-TPS,

r_{\mathrm{eff}}(z)/d_{z}

measures latent-space rank utilization, while global

r_{\mathrm{eff}}

measures whether the dominant force direction is shared across states. Values are mean

\pm

standard deviation over independently trained seeds.

Figure 3: Signed atom-pair directional correlations of the stochastic control for alanine dipeptide (left) and chignolin (right). The left lower and right upper triangles show FS-TPS and LaS-TPS, respectively. Positive and negative values indicate aligned and opposing force fluctuations. Results are averaged over seven training seeds.

$\beta$	RMSD (Å) $\downarrow$	THP (%) $\uparrow$	ETS (kJ/mol) $\downarrow$
Alanine dipeptide
$10^{-4}$	$0.25\pm 0.03$	$72.58\pm 5.01$	$21.6\pm 5.5$
$10^{-3}$	$0.23\pm 0.05$	$73.69\pm 6.99$	$22.4\pm 3.5$
$10^{-2}$	$0.27\pm 0.07$	$75.22\pm 13.17$	$21.9\pm 7.2$
Chignolin
$10^{-4}$	$1.21\pm 0.18$	$20.31\pm 14.26$	$-699.6\pm 123.2$

Table 4: Ablation study of the KL coefficient

\beta

for LaS-TPS. Results are reported as mean

\pm

standard deviation across random seeds. Arrows indicate the preferred direction of each metric. Blue cells indicate the best result within each molecular system.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

System	Method	Layer dimensions
Alanine dipeptide	TPS-DPS	$(\mathrm{input},128,256,256,256,128,\mathrm{output})$
	LaS-TPS	Encoder: $(\mathrm{input},64,128,64)$ ; latent heads: $2\times(64,10)$ ; decoder: $(10,64,128,64,\mathrm{output})$
Chignolin	TPS-DPS	$(\mathrm{input},512,1024,2048,1024,512,\mathrm{output})$
	LaS-TPS	Encoder: $(\mathrm{input},256,512,256)$ ; latent heads: $2\times(256,24)$ ; decoder: $(24,256,512,256,\mathrm{output})$
BBL	TPS-DPS	$(\mathrm{input},512,1024,2048,1024,512,\mathrm{output})$
	LaS-TPS	Encoder: $(\mathrm{input},1024,512,256)$ ; latent heads: $2\times(256,128)$ ; decoder: $(128,256,512,1024,\mathrm{output})$

Appendix

Table 5: Network architectures used for TPS-DPS and LaS-TPS.

Figure 4: Sampled transition paths on the alanine dipeptide free-energy surface for Doob’s Lagrangian.

x_{\mathcal{B},t}^{(m)}=\operatorname{Align}\!\left(x_{t}^{(m)},x_{\mathcal{B}}\right),\qquad\xi_{t}^{(m)}=\left(x_{t}^{(m)},x_{\mathcal{B},t}^{(m)}\right),

Appendix

Algorithm 1 Training with FS-TPS (colored in yellow) and LaS-TPS (colored in purple)

\left(x_{0}^{(m)},v_{0}^{(m)}\right)\sim\rho_{\mathcal{A}}.

Appendix

Algorithm 2 Inference with FS-TPS (colored in yellow) and LaS-TPS (colored in purple)

Explore similar work

Transition Path Sampling Using Koopman Operators and Exit-Time Optimal Control
Oct 7, 2026Boya Hou, Shane Wang, Siddharth Ambekar +2Stochastic Optimal Control Optimal Control
University of Illinois Urbana-Champaign
Correlation-Free Transition Path Sampling through Shooting Point Generation Guided by Committor Learning
Sep 17, 2026Maximilian Negedly, Sebastian Falkner, Alessandro Coretti +1Conditional Generative Modeling
Faculty of Physics, University of Vienna, 1090 Vienna, Austria · Vienna Doctoral School in Physics, University of Vienna, 1090 Vienna, Austria · Institute of Physics, University of Augsburg, 86159 Augsburg, Germany
Rare Event Analysis via Stochastic Optimal Control
Apr 14, 2026Yuanqi Du, Jiajun He, Dinghuai Zhang +2Rare-Event Estimation Stochastic Optimal Control
Microsoft Research New England · Cornell University · University of Cambridge +1

Oct 7, 2026J/K move · Enter open · S save

Transition Path Sampling Using Koopman Operators and Exit-Time Optimal Control

Boya Hou, Shane Wang, Siddharth Ambekar, Maxim Raginsky+1

University of Illinois Urbana-Champaign

Stochastic Optimal Control Optimal Control

Open paper AskPDF arXiv

Sampling transitions between metastable states is a central problem in dynamical systems theory and molecular dynamics in particular. A key challenge is the existence of high free-energy barriers that separate the states, making transitions extremely rare. Recent machine learning-based methods cast transition path sampling (TPS) as an optimal stochastic control (OSC) problem over a fixed time horizon, and parameterize the drift bias via a neural network trained by simulation-in-the-loop, requiring repeated biased rollouts. To address computational and performance guarantee issues of these models, we propose a new approach for the problem based on Koopman operators. Because Koopman operators are linear, their leading eigenfunctions reveal the metastable sets and provide an estimate of the committor function with no transition path information required. Furthermore, we formulate TPS as an OSC problem up to an exit time. Our time horizon is the first hitting time of the target set, and our running cost penalizes time spent in nonreactive regions by encoding the estimated committor function. We derive the optimal controller in closed form and approximate it in a reproducing kernel Hilbert space (RKHS). This reduces the problem of constructing the optimal controller to solving a single equality-constrained quadratic program, whose solution can be characterized by a linear Karush-Kuhn-Tucker (KKT) system. On the two-channel double well and alanine dipeptide, our controller increases the fraction of trajectories reaching the target from 0% to 99.8% within 1000 steps, and from 0% to 93% within 1ps, respectively.

Abstract

Explore similar work

Transition Path Sampling Using Koopman Operators and Exit-Time Optimal Control

Correlation-Free Transition Path Sampling through Shooting Point Generation Guided by Committor Learning

Rare Event Analysis via Stochastic Optimal Control