Stochastic Control Policies for Robust Molecular Transition Path Sampling

# Stochastic Control Policies for Robust Molecular Transition Path Sampling
Source: [https://arxiv.org/html/2608.13800](https://arxiv.org/html/2608.13800)
###### 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\.

## Introduction

Molecular transitions between metastable states underlie processes such as protein folding, conformational switching, and chemical reactions\. Transition path sampling \(TPS\) aims to generate reactive trajectories connecting metastable states, providing mechanistic insight into how biomolecular systems evolve between long\-lived conformations\([5](https://arxiv.org/html/2608.13800#bib.bib10);[3](https://arxiv.org/html/2608.13800#bib.bib11)\)\. Efficiently sampling such trajectories is challenging because the relevant free\-energy barriers are rarely crossed on accessible simulation timescales\. Although carefully designed collective variables \(CVs\) can facilitate this exploration, identifying informative CVs often requires substantial domain knowledge\. This limitation becomes more severe for protein folding, which involves high\-dimensional and coordinated atomic motions that may not admit an obvious low\-dimensional representation\.

Existing TPS approaches can be broadly divided into three families\. Traditional molecular dynamics \(MD\)\-based methods\([10](https://arxiv.org/html/2608.13800#bib.bib35)\)generate trajectories through molecular force fields, sometimes aided by externally applied steering forces\([15](https://arxiv.org/html/2608.13800#bib.bib7)\)or biased sampling approaches\([43](https://arxiv.org/html/2608.13800#bib.bib36);[21](https://arxiv.org/html/2608.13800#bib.bib37)\)\. Learning\-guided MD\-rollout methods instead train a neural policy to predict an instantaneous bias from the current molecular configuration\. The bias is combined with the physical force, and the system is advanced by the MD integrator, allowing the complete path to emerge sequentially from controlled dynamics\. PIPS\([14](https://arxiv.org/html/2608.13800#bib.bib6)\)introduced this policy\-learning formulation; TPS\-DPS improved off\-policy training with a replay buffer and a log\-variance path\-measure objective\([39](https://arxiv.org/html/2608.13800#bib.bib1)\)\. The third family, including Doob’s Lagrangian\([6](https://arxiv.org/html/2608.13800#bib.bib4)\)and FAS\([30](https://arxiv.org/html/2608.13800#bib.bib3)\), directly constructs complete endpoint\-conditioned paths without MD rollouts\. Among these families, learning\-guided MD rollouts remain competitive because they preserve explicit temporal propagation and tend to generate more physically plausible paths than direct endpoint\-conditioned generators\. Evidence reported in the Appendix shows that Doob’s Lagrangian paths collapse toward a geometrical connection in the CV space for alanine dipeptide, consistent with its failure independently reported elsewhere\([19](https://arxiv.org/html/2608.13800#bib.bib39)\)\. FAS produces less favorable low\-energy paths on the larger benchmark, chignolin\. However, rollout\-based methods also have a clear limitation: a poor initial replay buffer can repeatedly expose the policy to a narrow region of path space, reinforce unsuccessful control strategies, and produce strong sensitivity to network initialization\.

Prior work has attempted to address this instability through trust\-region regularization\([2](https://arxiv.org/html/2608.13800#bib.bib2)\), but achieved only limited performance gains\. In this work, we explore a different perspective: introducing stochasticity into policy learning\. This connects to a classical result in stochastic optimal control, where the noise level governs whether the optimal control averages over multiple solutions or collapses onto a single one\([18](https://arxiv.org/html/2608.13800#bib.bib45);[41](https://arxiv.org/html/2608.13800#bib.bib46)\)\. Our motivation also comes from reinforcement learning, which shares with MD\-rollout methods a sequential control structure, state\-dependent actions, trajectory\-based interaction, and replay\-buffer\-based off\-policy optimization, despite using a different learning objective\. Previous studies have shown that action\-space perturbations, parameter\-space noise, and entropy\-regularized stochastic policies can promote exploration\([24](https://arxiv.org/html/2608.13800#bib.bib28);[34](https://arxiv.org/html/2608.13800#bib.bib29);[9](https://arxiv.org/html/2608.13800#bib.bib30);[11](https://arxiv.org/html/2608.13800#bib.bib31)\), grounded in the broader connection between maximum\-entropy control and probabilistic inference\([23](https://arxiv.org/html/2608.13800#bib.bib44)\)\. Inspired by these approaches, we first replace the deterministic bias\-force predictor with a state\-dependent Gaussian policy that samples the applied control during each MD rollout, yielding force\-space stochastic TPS \(FS\-TPS\)\. On alanine dipeptide, FS\-TPS improves RMSD, transition success, and path diversity by39\.5%39\.5\\%,52\.8%52\.8\\%, and22\.2%22\.2\\%, respectively, relative to the deterministic baseline\. These gains remain consistent on the larger chignolin and BBL systems\. More importantly, evaluations over seven to nine random seeds show substantially reduced variability, demonstrating greater robustness to network initialization\.

These results further motivated us to examine where stochasticity should be introduced in higher\-dimensional molecular systems\. Independent fluctuations in Cartesian force components may be insufficient to represent the coordinated atomic motions involved in molecular transitions\. We therefore hypothesize that stochasticity should instead be introduced through a compact latent representation, allowing a shared decoder to transform latent perturbations into correlated variations of the full atom\-wise bias force\. Based on this idea, we propose latent\-space stochastic TPS \(LaS\-TPS\), which samples a state\-dependent latent variable within a bottleneck architecture and deterministically decodes it into the applied control\. Under a first\-order local approximation, we show that the induced force covariance is generally non\-diagonal and state\-dependent, with rank upper\-bounded by the latent dimension\. Our empirical characterization further shows that LaS\-TPS concentrates stochastic variation into low\-effective\-rank, cross\-atom\-correlated force modes, supporting the intended structured parameterization\. Using the same set of random initializations, LaS\-TPS reduces RMSD and transition\-state energy by41\.9%41\.9\\%and25\.4%25\.4\\%, respectively, while improving transition success and mode coverage by58\.9%58\.9\\%and77\.8%77\.8\\%on alanine dipeptide relative to the deterministic baseline\. It also retains improvements in transition success and energetic path quality on chignolin and BBL, with the gains evaluated over seven to nine random seeds\.

Our main contributions are as follows:

- •We revisit learning\-guided MD rollouts from a path\-space proposal perspective and introduce stochastic control as a mechanism for improving exploration and reducing sensitivity to early replay\-buffer trajectories\.
- •We propose FS\-TPS, which replaces the deterministic bias predictor with a state\-dependent Gaussian control policy and uses entropy regularization to maintain exploration during off\-policy training\. Across alanine dipeptide, chignolin, and BBL, FS\-TPS improves transition success and reduces sensitivity to network initialization over seven to nine random seeds\.
- •We propose LaS\-TPS, which introduces stochasticity through a compact bottleneck representation\. Under local linearization, we show that the induced force covariance is generally correlated and state\-dependent, with rank upper\-bounded by the latent dimension, and empirically characterize these structured force variations\. Across the three molecular systems, LaS\-TPS achieves the strongest overall transition\-success and energetic performance while maintaining robustness over random seeds\.

## Related work

##### Monte Carlo path sampling and learned proposals\.

Transition path sampling was originally formulated as Monte Carlo sampling in trajectory space, where shooting and shifting moves generate new reactive trajectories from existing ones\([5](https://arxiv.org/html/2608.13800#bib.bib10);[3](https://arxiv.org/html/2608.13800#bib.bib11)\)\. The efficiency of these samplers depends critically on the quality of the shooting points, which motivated a line of work that replaces hand\-designed moves with learned proposal distributions\. Boltzmann generators\([28](https://arxiv.org/html/2608.13800#bib.bib12)\)and their conditional variants\([8](https://arxiv.org/html/2608.13800#bib.bib13)\)learn to sample equilibrium configurations directly, and have been used to propose shooting points or MCMC moves in path space\([33](https://arxiv.org/html/2608.13800#bib.bib14)\)\. Related approaches guide path sampling with learned reaction coordinates\([16](https://arxiv.org/html/2608.13800#bib.bib15)\), allocate simulation effort adaptively through reinforcement learning\([47](https://arxiv.org/html/2608.13800#bib.bib16)\), or employ diffusion models as enhanced\-sampling proposals for rare events and free\-energy estimation\([46](https://arxiv.org/html/2608.13800#bib.bib17)\)\.

##### Rollout\-based learned control\.

A second line of work replaces the Monte Carlo machinery with a learned biasing force that steers the dynamics toward the target basin during explicit molecular\-dynamics rollouts\. The reactive path ensemble was first cast as a variational control problem in this form, with a time\-dependent control force reweighting the unbiased path measure\([40](https://arxiv.org/html/2608.13800#bib.bib18)\)\. PIPS\([14](https://arxiv.org/html/2608.13800#bib.bib6)\)introduced a collective\-variable\-free policy\-learning perspective, but its KL\-based objective can suffer from mode collapse, capturing only a subset of the reactive path modes\. TPS\-DPS\([39](https://arxiv.org/html/2608.13800#bib.bib1)\)formulates transition path generation as a controlled diffusion process and optimizes the policy with a log\-variance objective\([29](https://arxiv.org/html/2608.13800#bib.bib19)\), alleviating mode collapse but remaining sensitive to random seeds and policy initialization\. TR\-LV\([2](https://arxiv.org/html/2608.13800#bib.bib2)\)adds trust\-region constraints to log\-variance optimization, substantially improving stability while yielding only moderate gains in transition\-path metrics\. Throughout this line, paths are generated by sequential MD rather than imposed by endpoint construction, preserving a closer connection to physically time\-ordered dynamics\. All of these methods, however, learn a*deterministic*map from configuration to control force: the policy commits to a single control action per configuration\.

##### Endpoint\-conditioned and path\-space generative models\.

A complementary family constructs entire trajectories at once, enforcing the two metastable endpoints by construction rather than reaching the target basin through simulation\. One group formulates TPS through Doob’shh\-transform, which characterizes stochastic processes conditioned on rare events: Doob’s Lagrangian\([6](https://arxiv.org/html/2608.13800#bib.bib4)\)learns an endpoint\-conditioned path distribution through a simulation\-free variational objective, and subsequent work scales this formulation to longer transition paths with a sequence\-to\-sequence architecture and fixed\-window attention\([22](https://arxiv.org/html/2608.13800#bib.bib5)\)\. A second group generates the trajectory with a diffusion or flow model, by minimizing the Onsager–Machlup action\([35](https://arxiv.org/html/2608.13800#bib.bib21)\), learning normalizing flows over paths\([27](https://arxiv.org/html/2608.13800#bib.bib22)\), or training directly on trajectory data\([32](https://arxiv.org/html/2608.13800#bib.bib23);[44](https://arxiv.org/html/2608.13800#bib.bib24);[45](https://arxiv.org/html/2608.13800#bib.bib25)\)\. FAS\([30](https://arxiv.org/html/2608.13800#bib.bib3)\)extends adjoint\-based sampling\([48](https://arxiv.org/html/2608.13800#bib.bib26);[37](https://arxiv.org/html/2608.13800#bib.bib27);[12](https://arxiv.org/html/2608.13800#bib.bib40)\)to function spaces, sampling endpoint\-constrained path functions with a terminal cost combining potential energy and path smoothness\.

## Methodology

### Log\-variance objective in learning\-guided transition path sampling

TPS\-DPS\([39](https://arxiv.org/html/2608.13800#bib.bib1)\)trains a biasing policy by minimizing a log\-variance objective between the path distribution induced by biased molecular dynamics and the target transition\-path distribution\. Because this objective can be evaluated on trajectories stored in a replay buffer, it supports off\-policy training and improves sample efficiency\. A learnable control variate estimates the unknown normalization constant and reduces the variance of the gradient estimator\. We refer readers to the original TPS\-DPS paper\([39](https://arxiv.org/html/2608.13800#bib.bib1)\)for the complete derivation and training procedure\.

The deterministic TPS\-DPS policy takes the current molecular configuration and target conformation as input and produces the bias applied during molecular\-dynamics rollouts\. Two output parameterizations are considered\. The*force\-based*parameterization directly predicts the Cartesian bias force\. In contrast, the*scale\-based*parameterization predicts atom\-wise scaling coefficients, which are multiplied by the displacement toward the aligned target conformation to construct the bias force\. The latter introduces a target\-directed inductive bias and is more scalable for larger molecular systems\. Following the original system\-specific setup, we use direct force outputs for alanine dipeptide and scale outputs for the larger protein systems\. All proposed stochastic policy variants use the same output parameterization as their corresponding deterministic TPS\-DPS baseline to ensure a controlled comparison\.

### Force\-Space Distributional Control \(FS\-TPS\)

For FS\-TPS, we relax the assumption that the policy is deterministic by modeling the policy as a conditional probability distribution over admissible bias forces\. Given the current configurationxtx\_\{t\}and target conformationxℬx\_\{\\mathcal\{B\}\}, the policy defines a diagonal Gaussian distribution

πθ​\(𝐮t∣xt,xℬ\)=𝒩⁡\(𝝁θ​\(xt,xℬ\),diag⁡\(𝝈θ2​\(xt,xℬ\)\)\),\\pi\_\{\\theta\}\(\\mathbf\{u\}\_\{t\}\\mid x\_\{t\},x\_\{\\mathcal\{B\}\}\)=\\mathcal\{N\}\\left\(\\boldsymbol\{\\mu\}\_\{\\theta\}\(x\_\{t\},x\_\{\\mathcal\{B\}\}\),\\operatorname\{diag\}\\\!\\left\(\\boldsymbol\{\\sigma\}\_\{\\theta\}^\{2\}\(x\_\{t\},x\_\{\\mathcal\{B\}\}\)\\right\)\\right\),\(1\)
where𝝁θ​\(xt,xℬ\)\\boldsymbol\{\\mu\}\_\{\\theta\}\(x\_\{t\},x\_\{\\mathcal\{B\}\}\)and𝝈θ​\(xt,xℬ\)\\boldsymbol\{\\sigma\}\_\{\\theta\}\(x\_\{t\},x\_\{\\mathcal\{B\}\}\)denote the predicted mean and standard deviation of the bias force, respectively\. A bias force is sampled via the reparameterization trick\([20](https://arxiv.org/html/2608.13800#bib.bib33);[36](https://arxiv.org/html/2608.13800#bib.bib34)\)\.

𝐮t=𝝁θ​\(xt,xℬ\)\+𝝈θ​\(xt,xℬ\)⊙ϵ,ϵ∼𝒩⁡\(𝟎,𝐈\),\\mathbf\{u\}\_\{t\}=\\boldsymbol\{\\mu\}\_\{\\theta\}\(x\_\{t\},x\_\{\\mathcal\{B\}\}\)\+\\boldsymbol\{\\sigma\}\_\{\\theta\}\(x\_\{t\},x\_\{\\mathcal\{B\}\}\)\\odot\\boldsymbol\{\\epsilon\},\\qquad\\boldsymbol\{\\epsilon\}\\sim\\mathcal\{N\}\(\\mathbf\{0\},\\mathbf\{I\}\),\(2\)
#### Entropy\-Regularized Path\-Measure Objective

We regularize the controller by maximizing the conditional differential entropy of the bias\-force policy\. For the diagonal Gaussian policy, the entropy at molecular statextx\_\{t\}is

ℋ\(πθ\(⋅∣xt,xℬ\)\)=12∑i=1dulog\(2πeσθ,i2\(xt,xℬ\)\),\\mathcal\{H\}\\left\(\\pi\_\{\\theta\}\(\\cdot\\mid x\_\{t\},x\_\{\\mathcal\{B\}\}\)\\right\)=\\frac\{1\}\{2\}\\sum\_\{i=1\}^\{d\_\{u\}\}\\log\\left\(2\\pi e\\,\\sigma\_\{\\theta,i\}^\{2\}\(x\_\{t\},x\_\{\\mathcal\{B\}\}\)\\right\),\(3\)wheredud\_\{u\}denotes the dimensionality of the bias force\. We define the trajectory\-averaged policy entropy as

ℋ¯\(πθ\)=𝔼τ\[1T∑t=0T−1ℋ\(πθ\(⋅∣xt,xℬ\)\)\]\.\\overline\{\\mathcal\{H\}\}\(\\pi\_\{\\theta\}\)=\\mathbb\{E\}\_\{\\tau\}\\left\[\\frac\{1\}\{T\}\\sum\_\{t=0\}^\{T\-1\}\\mathcal\{H\}\\left\(\\pi\_\{\\theta\}\(\\cdot\\mid x\_\{t\},x\_\{\\mathcal\{B\}\}\)\\right\)\\right\]\.\(4\)
The final training objective combines path\-measure matching \(ℒLV\\mathcal\{L\}\_\{\\mathrm\{LV\}\}\) with entropy regularization,

ℒFS=ℒLV−λℋ​ℋ¯​\(πθ\)\.\\mathcal\{L\}\_\{\\mathrm\{FS\}\}=\\mathcal\{L\}\_\{\\mathrm\{LV\}\}\-\\lambda\_\{\\mathcal\{H\}\}\\overline\{\\mathcal\{H\}\}\(\\pi\_\{\\theta\}\)\.\(5\)
This objective encourages the induced path measure to approach the target transition\-path measure while preventing premature collapse of the state\-dependent bias\-force distribution\. To account for variations in thermal fluctuations across simulation conditions, we use a temperature\-dependent coefficient with a tunable hyperparameterλℋref\\lambda\_\{\\mathcal\{H\}\}^\{\\mathrm\{ref\}\}, whereTrefT\_\{\\mathrm\{ref\}\}is the final temperature of the annealing schedule when one is used, and the training temperature otherwise \(see Appendix for details\)\.

λℋ​\(T\)=λℋref​TTref,\\lambda\_\{\\mathcal\{H\}\}\(T\)=\\lambda\_\{\\mathcal\{H\}\}^\{\\mathrm\{ref\}\}\\frac\{T\}\{T\_\{\\mathrm\{ref\}\}\},\(6\)

### Latent\-Space Stochastic Control \(LaS\-TPS\)

LaS\-TPS represents uncertainty through a compact stochastic control variable\. It adopts a variational information bottleneck \(VIB\)\([1](https://arxiv.org/html/2608.13800#bib.bib8)\)architecture, in which an input is mapped to a distribution over latent representations and a sample from this distribution is used to perform the downstream task\. Given the current configurationxtx\_\{t\}and target conformationxℬx\_\{\\mathcal\{B\}\}, the encoder predicts a diagonal Gaussian distribution

qθe​\(𝐳t∣xt,xℬ\)=𝒩⁡\(𝝁θe​\(xt,xℬ\),diag⁡\[𝝈θe2​\(xt,xℬ\)\]\)\.q\_\{\\theta\_\{\\mathrm\{e\}\}\}\\left\(\\mathbf\{z\}\_\{t\}\\mid x\_\{t\},x\_\{\\mathcal\{B\}\}\\right\)=\\mathcal\{N\}\\\!\\left\(\\boldsymbol\{\\mu\}\_\{\\theta\_\{\\mathrm\{e\}\}\}\(x\_\{t\},x\_\{\\mathcal\{B\}\}\),\\operatorname\{diag\}\\\!\\left\[\\boldsymbol\{\\sigma\}\_\{\\theta\_\{\\mathrm\{e\}\}\}^\{2\}\(x\_\{t\},x\_\{\\mathcal\{B\}\}\)\\right\]\\right\)\.\(7\)At each molecular\-dynamics step, a latent control variable is sampled using the reparameterization trick and decoded into the bias force:

ϵt∼𝒩⁡\(𝟎,𝐈\),𝐳t=𝝁θe\+𝝈θe⊙ϵt,𝐮t=gθd​\(𝐳t\)\.\\boldsymbol\{\\epsilon\}\_\{t\}\\sim\\mathcal\{N\}\(\\mathbf\{0\},\\mathbf\{I\}\),\\qquad\\mathbf\{z\}\_\{t\}=\\boldsymbol\{\\mu\}\_\{\\theta\_\{\\mathrm\{e\}\}\}\+\\boldsymbol\{\\sigma\}\_\{\\theta\_\{\\mathrm\{e\}\}\}\\odot\\boldsymbol\{\\epsilon\}\_\{t\},\\qquad\\mathbf\{u\}\_\{t\}=g\_\{\\theta\_\{\\mathrm\{d\}\}\}\(\\mathbf\{z\}\_\{t\}\)\.\(8\)
The latent distribution is regularized toward a standard Gaussian prior\. The resulting LaS\-TPS objective is

ℒLaS​\(θ,log⁡Z\)=ℒLV​\(θ,log⁡Z\)\+\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{LaS\}\}\(\\theta,\\log Z\)=\\mathcal\{L\}\_\{\\mathrm\{LV\}\}\(\\theta,\\log Z\)\+\(9\)βKLK⁡\(L\+1\)​∑k=1K∑t=0LDKL​\(qθe​\(𝐳t∣xt\(k\),xℬ\)∥𝒩⁡\(𝟎,𝐈\)\)\.\\displaystyle\\frac\{\\beta\_\{\\mathrm\{KL\}\}\}\{K\(L\+1\)\}\\sum\_\{k=1\}^\{K\}\\sum\_\{t=0\}^\{L\}D\_\{\\mathrm\{KL\}\}\\\!\\left\(q\_\{\\theta\_\{\\mathrm\{e\}\}\}\\left\(\\mathbf\{z\}\_\{t\}\\mid x\_\{t\}^\{\(k\)\},x\_\{\\mathcal\{B\}\}\\right\)\\,\\middle\\\|\\,\\mathcal\{N\}\(\\mathbf\{0\},\\mathbf\{I\}\)\\right\)\.where coefficientβKL\\beta\_\{\\mathrm\{KL\}\}controls the strength of the latent information bottleneck \(see Appendix for details\)\.

##### Replay\-buffer sampling\.

For LaS\-TPS, we use a mildly prioritized replay strategy in which approximately5%5\\%of the training samples are drawn from the30%30\\%of replay\-buffer trajectories with the lowest RMSD, while the remaining samples follow the standard sampling procedure\. An ablation study shows that it provides only a marginal improvement\. Given its limited effect, we do not discuss it further in the main experimental section\. Full implementation details and ablation results are provided in the Appendix\.

### Proposition: Structured control covariance induced by latent stochasticity\.

In LaS\-TPS, stochasticity is introduced in a low\-dimensional latent space and subsequently mapped by a shared decoder to the full atom\-wise control force\. Consequently, a single latent perturbation may simultaneously modify multiple atoms and Cartesian components, allowing the policy to represent correlated control patterns such as coupled atomic motion, compensating force directions, coordinated rotations, and collective backbone displacement\. The following proposition characterizes the local covariance structure induced by this parameterization\.

Let the latent space

z\|x,xℬ∼𝒩⁡\(μz​\(x,xℬ\),Σz​\(x,xℬ\)\),Σz​\(x,xℬ\)=diag⁡\(σz2​\(x,xℬ\)\),\\begin\{split\}z\\mid x,x\_\{\\mathcal\{B\}\}\\sim\\mathcal\{N\}\\\!\\left\(\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\),\\Sigma\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\right\),\\\\ \\qquad\\Sigma\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)=\\operatorname\{diag\}\\\!\\left\(\\sigma\_\{z\}^\{2\}\(x,x\_\{\\mathcal\{B\}\}\)\\right\),\\end\{split\}\(10\)and the control force be given by

u=gθ​\(z\),u=g\_\{\\theta\}\(z\),\(11\)wherez∈ℝdzz\\in\\mathbb\{R\}^\{d\_\{z\}\}andu∈ℝduu\\in\\mathbb\{R\}^\{d\_\{u\}\}\. Define

Jz​\(x,xℬ\)=∂gθ​\(z\)∂z\|z=μz​\(x,xℬ\)\.J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)=\\left\.\\frac\{\\partial g\_\{\\theta\}\(z\)\}\{\\partial z\}\\right\|\_\{z=\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\}\.\(12\)Under a first\-order approximation aroundz=μz​\(x,xℬ\)z=\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\),

𝔼\[u∣x,xℬ\]\\displaystyle\\mathbb\{E\}\[u\\mid x,x\_\{\\mathcal\{B\}\}\]≈gθ​\(μz​\(x,xℬ\)\),\\displaystyle\\approx g\_\{\\theta\}\\\!\\left\(\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\right\),\(13\)Cov⁡\(u∣x,xℬ\)\\displaystyle\\operatorname\{Cov\}\(u\\mid x,x\_\{\\mathcal\{B\}\}\)≈Jz​\(x,xℬ\)​Σz​\(x,xℬ\)​Jz​\(x,xℬ\)⊤\.\\displaystyle\\approx J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\Sigma\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)^\{\\top\}\.\(14\)Therefore, the induced force covariance is generally non\-diagonal and state\-dependent, satisfies

rank⁡\(Cov⁡\(u∣x,xℬ\)\)≤dz,\\operatorname\{rank\}\\\!\\left\(\\operatorname\{Cov\}\(u\\mid x,x\_\{\\mathcal\{B\}\}\)\\right\)\\leq d\_\{z\},\(15\)and restricts the stochastic force variation to the column space ofJz​\(x,xℬ\)J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\. The proof is shown in the Appendix\.

##### Empirical characterization of conditional force fluctuations\.

For each system, we construct a shared bank of fixed molecular configurations spanning the transition from reactant to product \(S=200S=200for alanine dipeptide andS=224S=224for chignolin and BBL\)\. For each fixed statexx, the molecular coordinates are held constant and the trained policy is queriedK=1024K=1024times, producing samples from the conditional bias\-force distributionpθ​\(f∣x,xℬ\)p\_\{\\theta\}\(f\\mid x,x\_\{\\mathcal\{B\}\}\)\. Since the target conformationxℬx\_\{\\mathcal\{B\}\}is held fixed throughout, we omit it from the notation below\. The resulting variation reflects stochastic policy outputs at one exact configuration, rather than variation across molecular states\. After mean\-centering the sampled forces, we estimate

Σ⁡\(x\)=1K−1​∑k=1K\(fk−f¯\)​\(fk−f¯\)⊤,\\Sigma\(x\)=\\frac\{1\}\{K\-1\}\\sum\_\{k=1\}^\{K\}\\left\(f\_\{k\}\-\\bar\{f\}\\right\)\\left\(f\_\{k\}\-\\bar\{f\}\\right\)^\{\\top\},and project out rigid\-body translations and rotations to obtain the internal covarianceΣint​\(x\)\\Sigma\_\{\\mathrm\{int\}\}\(x\)\. We summarize the structure ofΣint​\(x\)\\Sigma\_\{\\mathrm\{int\}\}\(x\)for each fixed state and average these summary statistics over the state bank to obtain one value per trained checkpoint\. We additionally characterize latent\-rank utilization and the variation of the leading local force direction across states\. Full definitions and implementation details are provided in the Appendix\.

## Experiments

### Experimental Settings

We follow the TPS\-DPS implementation frame and evaluation protocol\([39](https://arxiv.org/html/2608.13800#bib.bib1)\)on three systems of increasing size: alanine dipeptide \(22 atoms\), transitioning from C5 to C7ax, and the fast\-folding proteins chignolin \(166 atoms\) and BBL \(711 atoms\)\([25](https://arxiv.org/html/2608.13800#bib.bib32)\), transitioning from unfolded to folded states\. All paths are generated on the fly through biased MD using the TPS\-DPS simulation settings\([7](https://arxiv.org/html/2608.13800#bib.bib41);[26](https://arxiv.org/html/2608.13800#bib.bib42);[4](https://arxiv.org/html/2608.13800#bib.bib43)\)\. Training, sampling, and evaluation are performed independently for each random seed, and each inference run generates 64 paths\. We report heavy\-atom final\-state RMSD after Kabsch alignment\([17](https://arxiv.org/html/2608.13800#bib.bib9)\), target hit percentage \(THP\), and transition\-state energy \(ETS\), the maximum potential energy among states in a transition path\. For alanine dipeptide, we additionally report mode coverage \(also referred to as path diversity\) over successful paths to assess whether they visit the two predefined saddle regions\. Coverage is0%0\\%,50%50\\%, or100%100\\%when none, one, or both regions are covered, respectively\. Further details are provided in the Appendix\.

### FS\-TPS improves training exploration and the learned mean control

![Refer to caption](https://arxiv.org/html/2608.13800v1/comparison_rollout.png)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[1](https://arxiv.org/html/2608.13800#Sx4.F1)shows that FS\-TPS separates from the deterministic baseline early in training, most clearly through a faster increase in THP\. Across nine seeds on alanine dipeptide, FS\-TPS improves THP from45\.67±28\.26%45\.67\\pm 28\.26\\%to69\.79±8\.84%69\.79\\pm 8\.84\\%and reduces RMSD from0\.43±0\.220\.43\\pm 0\.22Å to0\.26±0\.050\.26\\pm 0\.05Å\. The seed\-level results in Figure[2](https://arxiv.org/html/2608.13800#Sx4.F2)show how a run that remained nearly unsuccessful under TPS\-DPS is recovered by FS\-TPS\. Consistent improvements in THP are also observed on chignolin and BBL, which is summarized in Table[3](https://arxiv.org/html/2608.13800#Sx4.T3)\.

We found entropy regularization is essential to this behavior\. Removing it largely degrades the performance\. By directly discouraging premature variance collapse, the entropy term maintains stochastic exploration while the replay buffer is being constructed, as shown in Table[1](https://arxiv.org/html/2608.13800#Sx4.T1)\.

Table[1](https://arxiv.org/html/2608.13800#Sx4.T1)also shows stochastic sampling is not required at inference for FS\-TPS\. Scaling the learned standard deviation from00to1\.51\.5produces little change in RMSD or THP\. Only an extreme scale of1010substantially degrades ETS\. These results suggest that force\-space stochasticity acts primarily during training by enriching early rollouts and constructing a more successful replay buffer, thereby improving the learned mean control rather than relying on inference\-time noise\.

![Refer to caption](https://arxiv.org/html/2608.13800v1/tps-phipsi.png)

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\.We further examine whether simply adding temperature\-scaled Gaussian noise to a deterministic policy can reproduce the THP achieved by FS\-TPS\. As reported in the Appendix, it fails to do so\. These results indicate that the observed improvement cannot be attributed to additional randomness alone; rather, it arises from the learned, state\-dependent force distribution sustained by entropy regularization\.

Table 1:Ablation study of the proposed stochastic biasing policies on alanine dipeptide\. Unless otherwise indicated, results are reported as mean±\\pmstandard deviation across the same nine random seeds\. ETS is computed over runs with valid transition paths\. Arrows indicate the preferred direction of each metric\.MethodRMSD \(Å↓\\downarrow\)THP \(%↑\\uparrow\)ETS \(kJ/mol↓\\downarrow\)Deterministic force controlTPS\-DPS0\.43±0\.220\.43\\pm 0\.2245\.67±28\.2645\.67\\pm 28\.2629\.00±11\.7529\.00\\pm 11\.75Force\-space stochastic controlFS\-TPS \(σ\\sigma\-scale=1\.0=1\.0\)0\.26±0\.050\.26\\pm 0\.0569\.79±8\.8469\.79\\pm 8\.8429\.64±9\.7629\.64\\pm 9\.76σ\\sigma\-scale=0=0\(mean force\)0\.25±0\.040\.25\\pm 0\.0468\.16±12\.1668\.16\\pm 12\.1631\.64±10\.3631\.64\\pm 10\.36σ\\sigma\-scale=0\.5=0\.50\.26±0\.060\.26\\pm 0\.0667\.97±12\.1967\.97\\pm 12\.1931\.35±8\.0131\.35\\pm 8\.01σ\\sigma\-scale=1\.5=1\.50\.25±0\.050\.25\\pm 0\.0570\.90±9\.2670\.90\\pm 9\.2632\.21±9\.6432\.21\\pm 9\.64σ\\sigma\-scale=10=100\.26±0\.050\.26\\pm 0\.0567\.39±11\.1967\.39\\pm 11\.1949\.93±7\.4449\.93\\pm 7\.44w/o entropy regularization0\.54±0\.220\.54\\pm 0\.2229\.77±26\.3329\.77\\pm 26\.3343\.68±38\.4843\.68\\pm 38\.48Latent\-space stochastic controlLaS\-TPS \(σ\\sigma\-scale=1\.0=1\.0\)0\.25±0\.030\.25\\pm 0\.0372\.58±5\.0172\.58\\pm 5\.0121\.63±5\.4921\.63\\pm 5\.49σ\\sigma\-scale=0=0\(mean latent\)1\.01±0\.211\.01\\pm 0\.210\.98±1\.160\.98\\pm 1\.1623\.63±7\.5323\.63\\pm 7\.53σ\\sigma\-scale=0\.5=0\.50\.73±0\.160\.73\\pm 0\.168\.59±5\.488\.59\\pm 5\.4816\.57±5\.3916\.57\\pm 5\.39σ\\sigma\-scale=1\.5=1\.50\.16±0\.020\.16\\pm 0\.0288\.87±4\.9088\.87\\pm 4\.9038\.26±7\.3938\.26\\pm 7\.39w/o KL regularization0\.25±0\.070\.25\\pm 0\.0773\.44±16\.6273\.44\\pm 16\.6229\.40±13\.2129\.40\\pm 13\.21Deterministic bottleneck0\.51±0\.200\.51\\pm 0\.2037\.13±29\.7437\.13\\pm 29\.7429\.44±11\.7829\.44\\pm 11\.78

### From FS\-TPS to LaS\-TPS: Characterizing Sampled\-Force Fluctuations

Table 2:The internal force covarianceΣint​\(x\)\\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,reff​\(z\)/dzr\_\{\\mathrm\{eff\}\}\(z\)/d\_\{z\}measures latent\-space rank utilization, while globalreffr\_\{\\mathrm\{eff\}\}measures whether the dominant force direction is shared across states\. Values are mean±\\pmstandard deviation over independently trained seeds\.MethodInternalreffr\_\{\\mathrm\{eff\}\}Internal EV@1ExcessRcrossR\_\{\\mathrm\{cross\}\}RMS noise/meanreff​\(z\)/dzr\_\{\\mathrm\{eff\}\}\(z\)/d\_\{z\}Globalreffr\_\{\\mathrm\{eff\}\}Alanine dipeptideFS\-TPS52\.56±3\.0052\.56\\pm 3\.002\.82%±0\.16%2\.82\\%\\pm 0\.16\\%0\.00±0\.000\.00\\pm 0\.000\.20±0\.050\.20\\pm 0\.05//LaS\-TPS1\.20±0\.041\.20\\pm 0\.0495\.86%±1\.02%95\.86\\%\\pm 1\.02\\%0\.58±0\.020\.58\\pm 0\.020\.40±0\.200\.40\\pm 0\.200\.98±0\.000\.98\\pm 0\.00\(dz=10d\_\{z\}=10\)1\.07±0\.021\.07\\pm 0\.02ChignolinFS\-TPS30\.23±5\.1230\.23\\pm 5\.1212\.30%±2\.19%12\.30\\%\\pm 2\.19\\%0\.00±0\.000\.00\\pm 0\.000\.05±0\.020\.05\\pm 0\.02//LaS\-TPS1\.00±0\.001\.00\\pm 0\.00100\.00%±0\.00%100\.00\\%\\pm 0\.00\\%0\.75±0\.010\.75\\pm 0\.011\.34±0\.061\.34\\pm 0\.060\.99±0\.000\.99\\pm 0\.00\(dz=24d\_\{z\}=24\)14\.18±1\.3314\.18\\pm 1\.33BBLFS\-TPS139\.09±6\.57139\.09\\pm 6\.574\.01%±0\.34%4\.01\\%\\pm 0\.34\\%−0\.0001±0\.0002\-0\.0001\\pm 0\.00020\.110±0\.0250\.110\\pm 0\.025//LaS\-TPS1\.204±0\.0941\.204\\pm 0\.09495\.66%±2\.01%95\.66\\%\\pm 2\.01\\%0\.785±0\.0310\.785\\pm 0\.0311\.133±0\.1081\.133\\pm 0\.1080\.939±0\.00020\.939\\pm 0\.0002\(dz=128d\_\{z\}=128\)7\.96±0\.657\.96\\pm 0\.65

As established in the methodology section, a latent bottleneck allows a shared decoder to transform compact stochastic perturbations into non\-diagonal, state\-dependent force covariance with rank bounded by the latent dimension\. We use scalable MLPs with latent dimensions of 10, 24, and 128 for alanine dipeptide, chignolin, and BBL, respectively\.

To test the predicted structure, we characterize the conditional force distributions over the shared transition\-spanning state banks described in the methodology section\. Local statistics are computed independently at each fixed configuration fromK=1024K=1024policy samples and then summarized over states and trained seeds\. Because the entropy\-based effective rank is continuous, it may take non\-integer values\. As shown in Table[2](https://arxiv.org/html/2608.13800#Sx4.T2), FS\-TPS distributes its stochastic variation over many force directions, with internal effective ranks of52\.5652\.56,30\.2330\.23, and139\.09139\.09for alanine dipeptide, chignolin, and BBL, respectively\. In contrast, LaS\-TPS produces nearly one\-dimensional local fluctuations, with effective ranks close to one and a leading mode that explains almost all internal\-force variance\. These values are far below the corresponding latent\-dimensional upper bounds and do not arise from negligible stochastic amplitudes\.

The near\-unit output rank is not caused by collapse of the latent distribution\. We compute the effective rank of the latent covariance and normalize it by the nominal latent dimension\. The resulting utilization ratios are close to one for all three systems, indicating that most latent dimensions remain active\. Thus, the encoder does not reduce the latent distribution to a single active coordinate; instead, the decoder maps broadly utilized latent variation into an almost one\-dimensional local force subspace\. We further distinguish local dimensionality from variation of the dominant direction across molecular states\. The global effective rank is computed by aggregating the sign\-invariant outer products of the leading local covariance directions across the state bank\. Alanine dipeptide has a global effective rank close to one, indicating that a similar dominant direction is reused across configurations\. In contrast, chignolin and BBL have substantially larger global effective ranks\. Their fluctuations are therefore locally low\-rank, while the dominant direction changes across molecular configurations rather than collapsing globally to a single shared force mode\.

In Table[2](https://arxiv.org/html/2608.13800#Sx4.T2), we next quantify whether these fluctuations are coordinated across atoms\. For each atom pair, we measure the normalized magnitude of the corresponding off\-diagonal covariance block and subtract the same statistic under an atom\-permuted null, yielding the excess cross\-atom correlationRcrossR\_\{\\mathrm\{cross\}\}\. FS\-TPS remains indistinguishable from the null across all three systems, whereas LaS\-TPS produces substantial excess correlation\. We also report the RMS stochastic amplitude relative to the mean force, which is consistently larger for LaS\-TPS than for FS\-TPS\. Thus, LaS\-TPS concentrates substantial force variation into a small number of dominant directions rather than producing weak fluctuations with an apparently low\-dimensional covariance\. Figure[3](https://arxiv.org/html/2608.13800#Sx4.F3)visualizes the signed directional correlations for alanine dipeptide and chignolin, revealing strong cross\-atom coordination under LaS\-TPS\. For chignolin, these correlations occur both among nearby atoms and across several long\-range atom pairs, showing that a single latent perturbation can coordinate spatially separated molecular regions\.

Overall, latent\-space stochasticity is not merely a lower\-dimensional source of independent noise\. Instead, the decoder reorganizes latent variation into low\-effective\-rank, cross\-atom\-correlated force fluctuations whose dominant direction remains configuration dependent in the larger molecular systems\.

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\.Table 3:Comparison of traditional MD and learning\-guided transition path sampling methods on alanine dipeptide, chignolin, and BBL\. Blue shading indicates the best reported result for each metric, whereas orange shading indicates the second\-best result among the methods included in this table\.†denotes results reported by Blessing et al\. Results for alanine dipeptide are averaged over the same nine random seeds\. Results for chignolin and BBL are averaged over the same seven random seeds\.For BBL, the TPS\-DPS statistics are computed over five successful runs; two additional runs terminated during early rollouts because of numerical instability\.−\-indicates that implementation is not accessible while//indicates that certain values are not applicable\.Alanine Dipeptide \(22 atoms\)Chignolin \(166 atoms\)BBL \(711 atoms\)MethodRMSD\(Å↓\\downarrow\)THP\(%↑\\uparrow\)ETS\(kJ/mol↓\\downarrow\)Mode cov\.\(%↑\\uparrow\)RMSD\(Å↓\\downarrow\)THP\(%↑\\uparrow\)ETS\(kJ/mol↓\\downarrow\)RMSD\(Å↓\\downarrow\)THP\(%↑\\uparrow\)ETS\(kJ/mol↓\\downarrow\)Traditional MDUMD†1\.19±0\.031\.19\\pm 0\.036\.256\.25812\.47±148\.80812\.47\\pm 148\.80007\.23±0\.937\.23\\pm 0\.931\.561\.56388\.17388\.1718\.48±0\.6318\.48\\pm 0\.630\.000\.00/SMD†0\.56±0\.270\.56\\pm 0\.2754\.6954\.6978\.40±12\.7678\.40\\pm 12\.76𝟏𝟎𝟎\\bm\{100\}0\.85±0\.24\\bm\{0\.85\\pm 0\.24\}34\.3834\.38179\.52±138\.87179\.52\\pm 138\.872\.97±0\.332\.97\\pm 0\.337\.817\.81−1738\.57±386\.81\-1738\.57\\pm 386\.81Learning\-guided transition path samplingPIPS†0\.66±0\.150\.66\\pm 0\.1543\.7543\.7528\.17±10\.86\\bm\{28\.17\\pm 10\.86\}50504\.66±0\.174\.66\\pm 0\.170\.000\.00/17\.92±0\.2917\.92\\pm 0\.290\.000\.00/TPS\-DPS0\.43±0\.220\.43\\pm 0\.2245\.67±28\.2645\.67\\pm 28\.2629\.00±11\.7529\.00\\pm 11\.7550\.00±43\.3050\.00\\pm 43\.301\.22±0\.221\.22\\pm 0\.2227\.47±16\.8827\.47\\pm 16\.88−530\.85±171\.38\-530\.85\\pm 171\.382\.29±0\.252\.29\\pm 0\.2528\.76±13\.1928\.76\\pm 13\.19−2002\.00±1892\.27\-2002\.00\\pm 1892\.27TR\-LV†0\.29±0\.030\.29\\pm 0\.0361\.25±4\.0561\.25\\pm 4\.0549\.11±5\.8449\.11\\pm 5\.84–0\.90±0\.01\\bm\{0\.90\\pm 0\.01\}43\.95±5\.64\\bm\{43\.95\\pm 5\.64\}−303\.98±28\.65\-303\.98\\pm 28\.65–––FS\-TPS \(Ours\)0\.26±0\.05\\bm\{0\.26\\pm 0\.05\}69\.79±8\.84\\bm\{69\.79\\pm 8\.84\}29\.64±9\.7629\.64\\pm 9\.7661\.11±22\.0561\.11\\pm 22\.051\.16±0\.151\.16\\pm 0\.1531\.25±8\.0831\.25\\pm 8\.08−541\.29±185\.49\\bm\{\-541\.29\\pm 185\.49\}1\.81±0\.20\\bm\{1\.81\\pm 0\.20\}35\.30±12\.80\\bm\{35\.30\\pm 12\.80\}−2282\.71±718\.30\\bm\{\-2282\.71\\pm 718\.30\}LaS\-TPS \(Ours\)0\.25±0\.03\\bm\{0\.25\\pm 0\.03\}72\.58±5\.01\\bm\{72\.58\\pm 5\.01\}21\.63±5\.49\\bm\{21\.63\\pm 5\.49\}88\.89±22\.05\\bm\{88\.89\\pm 22\.05\}1\.03±0\.051\.03\\pm 0\.0538\.84±10\.63\\bm\{38\.84\\pm 10\.63\}−629\.00±76\.88\\bm\{\-629\.00\\pm 76\.88\}2\.33±0\.20\\bm\{2\.33\\pm 0\.20\}36\.39±14\.21\\bm\{36\.39\\pm 14\.21\}−2818\.29±518\.64\\bm\{\-2818\.29\\pm 518\.64\}

Table 4:Ablation study of the KL coefficientβ\\betafor LaS\-TPS\. Results are reported as mean±\\pmstandard deviation across random seeds\. Arrows indicate the preferred direction of each metric\. Blue cells indicate the best result within each molecular system\.β\\betaRMSD \(Å\)↓\\downarrowTHP \(%\)↑\\uparrowETS \(kJ/mol\)↓\\downarrowAlanine dipeptide10−410^\{\-4\}0\.25±0\.030\.25\\pm 0\.0372\.58±5\.0172\.58\\pm 5\.0121\.6±5\.521\.6\\pm 5\.510−310^\{\-3\}0\.23±0\.050\.23\\pm 0\.0573\.69±6\.9973\.69\\pm 6\.9922\.4±3\.522\.4\\pm 3\.510−210^\{\-2\}0\.27±0\.070\.27\\pm 0\.0775\.22±13\.1775\.22\\pm 13\.1721\.9±7\.221\.9\\pm 7\.2Chignolin10−410^\{\-4\}1\.21±0\.181\.21\\pm 0\.1820\.31±14\.2620\.31\\pm 14\.26−699\.6±123\.2\-699\.6\\pm 123\.210−310^\{\-3\}1\.03±0\.051\.03\\pm 0\.0538\.84±10\.6338\.84\\pm 10\.63−629\.0±76\.9\-629\.0\\pm 76\.910−210^\{\-2\}0\.99±0\.120\.99\\pm 0\.1246\.43±12\.6846\.43\\pm 12\.68−585\.6±119\.9\-585\.6\\pm 119\.9BBL10−710^\{\-7\}2\.14±0\.292\.14\\pm 0\.2925\.89±11\.7125\.89\\pm 11\.71−3069\.1±333\.2\-3069\.1\\pm 333\.210−610^\{\-6\}2\.33±0\.202\.33\\pm 0\.2036\.39±14\.2136\.39\\pm 14\.21−2818\.3±518\.6\-2818\.3\\pm 518\.6
### LaS\-TPS Achieves the Strongest Overall Performance among MD\-Rollout Methods

Across the three molecular systems, LaS\-TPS achieves the strongest overall performance among the MD\-rollout\-based methods while remaining robust across random seeds, as shown in Table[3](https://arxiv.org/html/2608.13800#Sx4.T3)\. Its advantage is particularly pronounced in energetic path quality: LaS\-TPS obtains the lowest ETS among the rollout\-based methods on alanine dipeptide, chignolin, and BBL, while also maintaining competitive RMSD and transition success\. On BBL, both FS\-TPS and LaS\-TPS alleviate the numerical instability observed in training of TPS\-DPS\. In the deterministic baseline, occasional excessively large predicted bias forces cause the MD simulation to diverge during early rollouts, whereas no such failures are observed for either stochastic policy\. This improved stability is consistent with the stochastic parameterizations and their associated regularization producing less brittle control outputs and reducing the likelihood of extreme bias\-force predictions\.

These gains are not attributable to greater model capacity\. Across the three systems, LaS\-TPS uses substantially fewer trainable parameters than the corresponding baseline MLPs, amounting to only20\.8%20\.8\\%,13\.4%13\.4\\%, and71\.5%71\.5\\%of the baseline parameter counts for alanine dipeptide, chignolin, and BBL, respectively\. Complete architecture specifications are provided in the Appendix\. They also cannot be explained by the bottleneck architecture alone\. On alanine dipeptide, replacing the stochastic latent variable with a deterministic bottleneck reduces THP below the TPS\-DPS baseline and substantially increases RMSD, indicating that simply inserting a low\-dimensional hidden layer does not recover the benefit of LaS\-TPS\. Similarly, directly decoding the latent mean at inference yields poor endpoint accuracy and almost no successful transitions\.

KL regularization is also important for controlling the learned latent distribution\. Results in Table[1](https://arxiv.org/html/2608.13800#Sx4.T1)show removing the KL term preserves endpoint performance on alanine dipeptide but increases variability and degrades energetic path quality\. Moreover, varying the KL coefficient reveals a consistent trade\-off between transition success and ETS on the larger systems, as shown in Table[4](https://arxiv.org/html/2608.13800#Sx4.T4), consistent with the rate–distortion trade\-off characterized by the information bottleneck framework\([42](https://arxiv.org/html/2608.13800#bib.bib38);[1](https://arxiv.org/html/2608.13800#bib.bib8);[13](https://arxiv.org/html/2608.13800#bib.bib20)\)\. A stronger KL penalty compresses the configuration\-dependent latent representation toward the shared prior, providing greater regularization and often improving target\-reaching robustness\. A weaker penalty retains more state\-specific information, which can enable lower\-energy control strategies but may reduce transition success\. The KL coefficient controls not only latent regularization but also the balance between reliably reaching the target basin and generating energetically favorable transition paths\.

Because the conditional force covariance consistently exhibits an effective rank close to one, we further tested whether a one\-dimensional latent space would be sufficient\. For alanine dipeptide, we reduced the latent dimension from1010to11and evaluated four random seeds\. One of them showed clear degradation, with a THP of21\.9%21\.9\\%and an ETS of60\.3​kJ/mol60\.3\\,\\mathrm\{kJ/mol\}; full results are provided in the Appendix\. These results suggest that, although a scalar latent variable can be sufficient when training succeeds, the higher\-dimensional latent representation improves optimization robustness and reduces sensitivity to initialization\.

## Conclusion and Limitations

We introduced two stochastic\-control formulations for MD\-rollout\-based TPS\. FS\-TPS directly models a state\-dependent distribution in force space, whereas LaS\-TPS injects stochasticity through a compact latent representation that is decoded into the full atom\-wise bias force\. Across three molecular systems, both approaches improve robustness to random initialization and transition\-path quality relative to deterministic control, with LaS\-TPS providing the strongest overall balance of transition success and energetic quality\. Our analysis further shows that the location of stochasticity substantially changes the resulting control distribution\. While FS\-TPS produces broadly distributed, approximately factorized force fluctuations, LaS\-TPS organizes latent variation into low\-effective\-rank, cross\-atom\-correlated force patterns whose dominant directions remain configuration dependent in the larger systems\. These results suggest that high\-dimensional molecular control may benefit not only from introducing stochasticity, but also from learning how stochastic variation is structured across atoms and force components\.

However, we found there is always a trade\-off between THP and ETS\. Within LaS\-TPS, varying the KL coefficient changes the balance between THP and ETS\. A related trade\-off appears more broadly across different TPS formulations\. MD\-rollout\-based methods must discover the target basin through sequential simulation and therefore do not guarantee endpoint satisfaction, whereas endpoint\-conditioned generators achieve high transition success by construction but may produce paths with less favorable energetic or dynamical properties\. Developing objectives that jointly balance transition probability, energetic quality, and path diversity therefore remains an open challenge\. We hope that LaS\-TPS will motivate further work on compact stochastic representations that transform high\-dimensional molecular uncertainty into coordinated and scalable control strategies\.

## References

- Alemiet al\.\(2016\)A\. A\. Alemi, I\. Fischer, J\. V\. Dillon, and K\. MurphyDeep variational information bottleneck\.arXiv preprint arXiv:1612\.00410\.Cited by:[Latent\-Space Stochastic Control \(LaS\-TPS\)](https://arxiv.org/html/2608.13800#Sx3.SSx3.p1.1),[LaS\-TPS Achieves the Strongest Overall Performance among MD\-Rollout Methods](https://arxiv.org/html/2608.13800#Sx4.SSx4.p3.1)\.
- Blessinget al\.\(2026\)D\. Blessing, J\. Berner, L\. Richter, C\. Domingo i Enrich, Y\. Du, A\. Vahdat, and G\. NeumannTrust region constrained measure transport in path space for stochastic optimal control and inference\.Advances in Neural Information Processing Systems38,pp\. 165462–165510\.Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p3.1),[Rollout\-based learned control\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px2.p1.1)\.
- Bolhuiset al\.\(2002\)P\. G\. Bolhuis, D\. Chandler, C\. Dellago, and P\. L\. GeisslerTRANSITION path sampling: throwing ropes over rough mountain passes, in the dark\.Annual Review of Physical Chemistry53\(Volume 53, 2002\),pp\. 291–318\.External Links:[Document](https://dx.doi.org/https%3A//doi.org/10.1146/annurev.physchem.53.082301.113146),[Link](https://www.annualreviews.org/content/journals/10.1146/annurev.physchem.53.082301.113146),ISSN 1545\-1593Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p1.1),[Monte Carlo path sampling and learned proposals\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px1.p1.1)\.
- Bussi and Parrinello \(2007\)G\. Bussi and M\. ParrinelloAccurate sampling using langevin dynamics\.Physical Review E75\(5\),pp\. 056707\.Cited by:[Experimental Settings](https://arxiv.org/html/2608.13800#Sx4.SSx1.p1.1)\.
- Dellagoet al\.\(1998\)C\. Dellago, P\. G\. Bolhuis, F\. S\. Csajka, and D\. ChandlerTransition path sampling and the calculation of rate constants\.The Journal of Chemical Physics108\(5\),pp\. 1964–1977\.External Links:ISSN 0021\-9606,[Document](https://dx.doi.org/10.1063/1.475562),[Link](https://doi.org/10.1063/1.475562),https://pubs\.aip\.org/aip/jcp/article\-pdf/108/5/1964/19283374/1964\_1\_online\.pdfCited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p1.1),[Monte Carlo path sampling and learned proposals\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px1.p1.1)\.
- Duet al\.\(2024\)Y\. Du, M\. Plainer, R\. Brekelmans, C\. Duan, F\. Noe, C\. P\. Gomes, A\. Aspuru\-Guzik, and K\. NeklyudovDoob’s lagrangian: a sample\-efficient variational approach to transition path sampling\.Advances in Neural Information Processing Systems37,pp\. 65791–65822\.Cited by:[Appendix A](https://arxiv.org/html/2608.13800#A1.SSx8.p1.1),[Introduction](https://arxiv.org/html/2608.13800#Sx1.p2.1),[Endpoint\-conditioned and path\-space generative models\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px3.p1.1)\.
- Eastmanet al\.\(2023\)P\. Eastman, R\. Galvelis, R\. P\. Peláez, C\. R\. Abreu, S\. E\. Farr, E\. Gallicchio, A\. Gorenko, M\. M\. Henry, F\. Hu, J\. Huang,et al\.OpenMM 8: molecular dynamics simulation with machine learning potentials\.The Journal of Physical Chemistry B128\(1\),pp\. 109–116\.Cited by:[Experimental Settings](https://arxiv.org/html/2608.13800#Sx4.SSx1.p1.1)\.
- Falkneret al\.\(2023\)S\. Falkner, A\. Coretti, S\. Romano, P\. L\. Geissler, and C\. DellagoConditioning boltzmann generators for rare event sampling\.Machine Learning: Science and Technology4\(3\),pp\. 035050\.External Links:[Document](https://dx.doi.org/10.1088/2632-2153/acf55c),[Link](https://doi.org/10.1088/2632-2153/acf55c)Cited by:[Monte Carlo path sampling and learned proposals\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px1.p1.1)\.
- Fortunatoet al\.\(2018\)M\. Fortunato, M\. G\. Azar, B\. Piot, J\. Menick, I\. Osband, A\. Graves, V\. Mnih, R\. Munos, D\. Hassabis, O\. Pietquin, C\. Blundell, and S\. LeggNoisy networks for exploration\.InInternational Conference on Learning Representations,Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p3.1)\.
- Frenkel and Smit \(2002\)D\. Frenkel and B\. SmitUnderstanding molecular simulation: from algorithms to applications\.2nd edition,Academic Press\.Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p2.1)\.
- Haarnojaet al\.\(2018\)T\. Haarnoja, A\. Zhou, P\. Abbeel, and S\. LevineSoft actor\-critic: off\-policy maximum entropy deep reinforcement learning with a stochastic actor\.InProceedings of the 35th International Conference on Machine Learning,J\. Dy and A\. Krause \(Eds\.\),Proceedings of Machine Learning Research, Vol\.80,pp\. 1861–1870\.Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p3.1)\.
- Havenset al\.\(2025\)A\. Havens, B\. K\. Miller, B\. Yan, C\. Domingo\-Enrich, A\. Sriram, B\. Wood, D\. Levine, B\. Hu, B\. Amos, B\. Karrer, X\. Fu, G\. Liu, and R\. T\. Q\. ChenAdjoint sampling: highly scalable diffusion samplers via adjoint matching\.arXiv preprint arXiv:2504\.11713\.Cited by:[Endpoint\-conditioned and path\-space generative models\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px3.p1.1)\.
- Higginset al\.\(2017\)I\. Higgins, L\. Matthey, A\. Pal, C\. Burgess, X\. Glorot, M\. Botvinick, S\. Mohamed, and A\. LerchnerBeta\-VAE: learning basic visual concepts with a constrained variational framework\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=Sy2fzU9gl)Cited by:[LaS\-TPS Achieves the Strongest Overall Performance among MD\-Rollout Methods](https://arxiv.org/html/2608.13800#Sx4.SSx4.p3.1)\.
- Holdijket al\.\(2023\)L\. Holdijk, Y\. Du, F\. Hooft, P\. Jaini, B\. Ensing, and M\. WellingStochastic optimal control for collective variable free sampling of molecular transition paths\.Advances in Neural Information Processing Systems36,pp\. 79540–79556\.Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p2.1),[Rollout\-based learned control\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px2.p1.1)\.
- Izrailevet al\.\(1999\)S\. Izrailev, S\. Stepaniants, B\. Isralewitz, D\. Kosztin, H\. Lu, F\. Molnar, W\. Wriggers, and K\. SchultenSteered molecular dynamics\.InComputational Molecular Dynamics: Challenges, Methods, Ideas: Proceedings of the 2nd International Symposium on Algorithms for Macromolecular Modelling, Berlin, May 21–24, 1997,pp\. 39–65\.Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p2.1)\.
- Junget al\.\(2023\)H\. Jung, R\. Covino, A\. Arjun, C\. Leitold, C\. Dellago, P\. G\. Bolhuis, and G\. HummerMachine\-guided path sampling to discover mechanisms of molecular self\-organization\.Nature Computational Science3\(4\),pp\. 334–345\.External Links:[Document](https://dx.doi.org/10.1038/s43588-023-00428-z),[Link](https://doi.org/10.1038/s43588-023-00428-z)Cited by:[Monte Carlo path sampling and learned proposals\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px1.p1.1)\.
- Kabsch \(1976\)W\. KabschA solution for the best rotation to relate two sets of vectors\.Acta Crystallographica Section A32\(5\),pp\. 922–923\.External Links:[Document](https://dx.doi.org/10.1107/S0567739476001873),[Link](https://doi.org/10.1107/S0567739476001873)Cited by:[Appendix A](https://arxiv.org/html/2608.13800#A1.SSx3.SSS0.Px1.p1.1),[Experimental Settings](https://arxiv.org/html/2608.13800#Sx4.SSx1.p1.1)\.
- Kappen \(2005\)H\. J\. KappenPath integrals and symmetry breaking for optimal control theory\.Journal of Statistical Mechanics: Theory and Experiment2005\(11\),pp\. P11011\.Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p3.1)\.
- Karczewskiet al\.\(2026\)R\. Karczewski, M\. Heinonen, A\. Pouplin, S\. Hauberg, and V\. GargThe spacetime of diffusion models: an information geometry perspective\.InInternational Conference on Learning Representations,Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p2.1)\.
- Kingma and Welling \(2014\)D\. P\. Kingma and M\. WellingAuto\-encoding variational bayes\.InInternational Conference on Learning Representations,Cited by:[Force\-Space Distributional Control \(FS\-TPS\)](https://arxiv.org/html/2608.13800#Sx3.SSx2.p2.1)\.
- Laio and Parrinello \(2002\)A\. Laio and M\. ParrinelloEscaping free\-energy minima\.Proceedings of the National Academy of Sciences99\(20\),pp\. 12562–12566\.Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p2.1)\.
- Leeet al\.\(2025\)J\. Lee, M\. Plainer, Y\. Du, L\. Holdijk, R\. Brekelmans, D\. Beaini, and K\. NeklyudovScaling deep learning solutions for transition path sampling\.InFrontiers in Probabilistic Inference: Learning meets Sampling,Cited by:[Table 7](https://arxiv.org/html/2608.13800#A1.T7),[Endpoint\-conditioned and path\-space generative models\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px3.p1.1)\.
- Levine \(2018\)S\. LevineReinforcement learning and control as probabilistic inference: tutorial and review\.arXiv preprint arXiv:1805\.00909\.Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p3.1)\.
- Lillicrapet al\.\(2016\)T\. P\. Lillicrap, J\. J\. Hunt, A\. Pritzel, N\. Heess, T\. Erez, Y\. Tassa, D\. Silver, and D\. WierstraContinuous control with deep reinforcement learning\.InInternational Conference on Learning Representations,Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p3.1)\.
- Lindorff\-Larsenet al\.\(2011\)K\. Lindorff\-Larsen, S\. Piana, R\. O\. Dror, and D\. E\. ShawHow fast\-folding proteins fold\.Science334\(6055\),pp\. 517–520\.External Links:[Document](https://dx.doi.org/10.1126/science.1208351)Cited by:[Experimental Settings](https://arxiv.org/html/2608.13800#Sx4.SSx1.p1.1)\.
- Lindorff\-Larsenet al\.\(2010\)K\. Lindorff\-Larsen, S\. Piana, K\. Palmo, P\. Maragakis, J\. L\. Klepeis, R\. O\. Dror, and D\. E\. ShawImproved side\-chain torsion potentials for the amber ff99sb protein force field\.Proteins78,pp\. 1950–1958\.Cited by:[Experimental Settings](https://arxiv.org/html/2608.13800#Sx4.SSx1.p1.1)\.
- Liuet al\.\(2022\)T\. Liu, W\. Gao, Z\. Wang, and C\. WangPathFlow: a normalizing flow generator that finds transition paths\.InProceedings of the Thirty\-Eighth Conference on Uncertainty in Artificial Intelligence,J\. Cussens and K\. Zhang \(Eds\.\),Proceedings of Machine Learning Research, Vol\.180,pp\. 1232–1242\.External Links:[Link](https://proceedings.mlr.press/v180/liu22b.html)Cited by:[Endpoint\-conditioned and path\-space generative models\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px3.p1.1)\.
- Noéet al\.\(2019\)F\. Noé, S\. Olsson, J\. Köhler, and H\. WuBoltzmann generators: sampling equilibrium states of many\-body systems with deep learning\.Science365\(6457\),pp\. eaaw1147\.External Links:[Document](https://dx.doi.org/10.1126/science.aaw1147),[Link](https://www.science.org/doi/abs/10.1126/science.aaw1147),https://www\.science\.org/doi/pdf/10\.1126/science\.aaw1147Cited by:[Monte Carlo path sampling and learned proposals\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px1.p1.1)\.
- Nüsken and Richter \(2020\)N\. Nüsken and L\. RichterSolving high\-dimensional hamilton\-jacobi\-bellman pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space\.arXiv preprint arXiv:2005\.05409\.Cited by:[Rollout\-based learned control\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px2.p1.1)\.
- Parket al\.\(2025\)B\. Park, J\. Lee, and G\. LiuFunctional adjoint sampler: scalable sampling on infinite dimensional spaces\.arXiv preprint arXiv:2511\.06239\.Cited by:[Appendix A](https://arxiv.org/html/2608.13800#A1.SSx8.p1.1),[Table 7](https://arxiv.org/html/2608.13800#A1.T7),[Introduction](https://arxiv.org/html/2608.13800#Sx1.p2.1),[Endpoint\-conditioned and path\-space generative models\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px3.p1.1)\.
- Pérez\-Hernándezet al\.\(2013\)G\. Pérez\-Hernández, F\. Paul, T\. Giorgino, G\. de Fabritiis, and F\. NoéIdentification of slow molecular order parameters for markov model construction\.The Journal of Chemical Physics139\(1\),pp\. 015102\.External Links:[Document](https://dx.doi.org/10.1063/1.4811489)Cited by:[Appendix A](https://arxiv.org/html/2608.13800#A1.SSx3.SSS0.Px2.p1.1)\.
- Petersenet al\.\(2024\)M\. Petersen, G\. Roig, and R\. CovinoDynamicsDiffusion: generating and rare event sampling of molecular dynamic trajectories using diffusion models\.External Links:[Link](https://openreview.net/forum?id=kKXIYUi8ff)Cited by:[Endpoint\-conditioned and path\-space generative models\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px3.p1.1)\.
- Plaineret al\.\(2024\)M\. Plainer, H\. Stärk, C\. Bunne, and S\. GünnemannTransition path sampling with boltzmann generator\-based mcmc moves\.External Links:2312\.05340,[Link](https://arxiv.org/abs/2312.05340)Cited by:[Monte Carlo path sampling and learned proposals\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px1.p1.1)\.
- Plappertet al\.\(2018\)M\. Plappert, R\. Houthooft, P\. Dhariwal, S\. Sidor, R\. Y\. Chen, X\. Chen, T\. Asfour, P\. Abbeel, and M\. AndrychowiczParameter space noise for exploration\.InInternational Conference on Learning Representations,Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p3.1)\.
- Rajaet al\.\(2025\)S\. Raja, M\. Šípka, M\. Psenka, T\. Kreiman, M\. Pavelka, and A\. S\. KrishnapriyanAction\-minimization meets generative modeling: efficient transition path sampling with the onsager\-machlup functional\.External Links:2504\.18506,[Link](https://arxiv.org/abs/2504.18506)Cited by:[Endpoint\-conditioned and path\-space generative models\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px3.p1.1)\.
- Rezendeet al\.\(2014\)D\. J\. Rezende, S\. Mohamed, and D\. WierstraStochastic backpropagation and approximate inference in deep generative models\.InInternational Conference on Machine Learning,pp\. 1278–1286\.Cited by:[Force\-Space Distributional Control \(FS\-TPS\)](https://arxiv.org/html/2608.13800#Sx3.SSx2.p2.1)\.
- Richter and Berner \(2024\)L\. Richter and J\. BernerImproved sampling via learned diffusions\.External Links:2307\.01198,[Link](https://arxiv.org/abs/2307.01198)Cited by:[Endpoint\-conditioned and path\-space generative models\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px3.p1.1)\.
- Roy and Vetterli \(2007\)O\. Roy and M\. VetterliThe effective rank: a measure of effective dimensionality\.InProceedings of the 15th European Signal Processing Conference \(EUSIPCO\),pp\. 606–610\.Cited by:[Appendix A](https://arxiv.org/html/2608.13800#A1.SSx5.SSS0.Px3.p1.2)\.
- Seonget al\.\(2025\)K\. Seong, S\. Park, S\. Kim, W\. Y\. Kim, and S\. AhnTransition path sampling with improved off\-policy training of diffusion path samplers\.InInternational Conference on Learning Representations,Vol\.2025,pp\. 93040–93062\.Cited by:[Appendix A](https://arxiv.org/html/2608.13800#A1.SSx3.p1.1),[Introduction](https://arxiv.org/html/2608.13800#Sx1.p2.1),[Rollout\-based learned control\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px2.p1.1),[Log\-variance objective in learning\-guided transition path sampling](https://arxiv.org/html/2608.13800#Sx3.SSx1.p1.1),[Experimental Settings](https://arxiv.org/html/2608.13800#Sx4.SSx1.p1.1)\.
- Singh and Limmer \(2023\)A\. N\. Singh and D\. T\. LimmerVariational deep learning of equilibrium transition path ensembles\.The Journal of Chemical Physics159\(2\)\.External Links:ISSN 1089\-7690,[Link](http://dx.doi.org/10.1063/5.0150278),[Document](https://dx.doi.org/10.1063/5.0150278)Cited by:[Rollout\-based learned control\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px2.p1.1)\.
- Theodorouet al\.\(2010\)E\. Theodorou, J\. Buchli, and S\. SchaalA generalized path integral control approach to reinforcement learning\.Journal of Machine Learning Research11\(104\),pp\. 3137–3181\.Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p3.1)\.
- Tishbyet al\.\(1999\)N\. Tishby, F\. C\. Pereira, and W\. BialekThe information bottleneck method\.InProceedings of the 37th Annual Allerton Conference on Communication, Control and Computing,pp\. 368–377\.Cited by:[LaS\-TPS Achieves the Strongest Overall Performance among MD\-Rollout Methods](https://arxiv.org/html/2608.13800#Sx4.SSx4.p3.1)\.
- Torrie and Valleau \(1977\)G\. M\. Torrie and J\. P\. ValleauNonphysical sampling distributions in monte carlo free\-energy estimation: umbrella sampling\.Journal of Computational Physics23\(2\),pp\. 187–199\.Cited by:[Introduction](https://arxiv.org/html/2608.13800#Sx1.p2.1)\.
- Triplett and Lu \(2023\)L\. Triplett and J\. LuDiffusion methods for generating transition paths\.External Links:2309\.10276,[Link](https://arxiv.org/abs/2309.10276)Cited by:[Endpoint\-conditioned and path\-space generative models\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px3.p1.1)\.
- Wanget al\.\(2024\)H\. Wang, Y\. Qiu, Y\. Wang, R\. Brekelmans, and Y\. DuGeneralized flow matching for transition dynamics modeling\.External Links:2410\.15128,[Link](https://arxiv.org/abs/2410.15128)Cited by:[Endpoint\-conditioned and path\-space generative models\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px3.p1.1)\.
- Xieet al\.\(2026\)Y\. Xie, L\. Winkler, L\. Sun, S\. Lewis, A\. E\. Foster, J\. J\. Luna, T\. Hempel, M\. Gastegger, Y\. Chen, I\. Zaporozhets, C\. Clementi, C\. M\. Bishop, and F\. NoéEnhanced diffusion sampling: efficient rare event sampling and free energy calculation with diffusion models\.External Links:2602\.16634,[Link](https://arxiv.org/abs/2602.16634)Cited by:[Monte Carlo path sampling and learned proposals\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px1.p1.1)\.
- Yanget al\.\(2024\)D\. T\. Yang, A\. M\. Goldberg, and L\. T\. ChongRare\-event sampling using a reinforcement learning\-based weighted ensemble method\.bioRxiv\.External Links:[Document](https://dx.doi.org/10.1101/2024.10.09.617475),[Link](https://www.biorxiv.org/content/early/2024/10/11/2024.10.09.617475),https://www\.biorxiv\.org/content/early/2024/10/11/2024\.10\.09\.617475\.full\.pdfCited by:[Monte Carlo path sampling and learned proposals\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px1.p1.1)\.
- Zhang and Chen \(2022\)Q\. Zhang and Y\. ChenPath integral sampler: a stochastic control approach for sampling\.External Links:2111\.15141,[Link](https://arxiv.org/abs/2111.15141)Cited by:[Endpoint\-conditioned and path\-space generative models\.](https://arxiv.org/html/2608.13800#Sx2.SS0.SSS0.Px3.p1.1)\.

Appendix for Stochastic Control Policies for Robust Molecular Transition Path Sampling

## Appendix AAppendix

### Proof of the Structured\-Covariance Proposition

Let

u=gθ\(z\),z∣x,xℬ∼𝒩\(μz\(x,xℬ\),Σz\(x,xℬ\)\)\.u=g\_\{\\theta\}\(z\),\\qquad z\\mid x,x\_\{\\mathcal\{B\}\}\\sim\\mathcal\{N\}\\\!\\left\(\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\),\\Sigma\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\right\)\.A first\-order Taylor expansion ofgθg\_\{\\theta\}around the state\-dependent latent meanμz​\(x,xℬ\)\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)gives

u≈gθ​\(μz​\(x,xℬ\)\)\+Jz​\(x,xℬ\)​\(z−μz​\(x,xℬ\)\),u\\approx g\_\{\\theta\}\\\!\\left\(\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\right\)\+J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\left\(z\-\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\right\),\(16\)where

Jz​\(x,xℬ\)=∂gθ​\(z\)∂z\|z=μz​\(x,xℬ\)\.J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)=\\left\.\\frac\{\\partial g\_\{\\theta\}\(z\)\}\{\\partial z\}\\right\|\_\{z=\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\}\.\(17\)Since

𝔼\[z−μz\(x,xℬ\)∣x,xℬ\]=0,Cov\(z∣x,xℬ\)=Σz\(x,xℬ\),\\mathbb\{E\}\\\!\\left\[z\-\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\mid x,x\_\{\\mathcal\{B\}\}\\right\]=0,\\qquad\\operatorname\{Cov\}\(z\\mid x,x\_\{\\mathcal\{B\}\}\)=\\Sigma\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\),\(18\)the conditional mean and covariance are approximately

𝔼\[u∣x,xℬ\]≈gθ\(μz\(x,xℬ\)\),\\mathbb\{E\}\[u\\mid x,x\_\{\\mathcal\{B\}\}\]\\approx g\_\{\\theta\}\\\!\\left\(\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\right\),\(19\)and

Cov⁡\(u∣x,xℬ\)≈Jz​\(x,xℬ\)​Σz​\(x,xℬ\)​Jz​\(x,xℬ\)⊤\.\\operatorname\{Cov\}\(u\\mid x,x\_\{\\mathcal\{B\}\}\)\\approx J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\Sigma\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)^\{\\top\}\.\(20\)Moreover,

u−𝔼\[u∣x,xℬ\]≈Jz\(x,xℬ\)\(z−μz\(x,xℬ\)\)∈Col\(Jz\(x,xℬ\)\),u\-\\mathbb\{E\}\[u\\mid x,x\_\{\\mathcal\{B\}\}\]\\approx J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\left\(z\-\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\right\)\\in\\operatorname\{Col\}\\\!\\left\(J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\right\),\(21\)and

rank⁡\(Jz​\(x,xℬ\)​Σz​\(x,xℬ\)​Jz​\(x,xℬ\)⊤\)≤rank⁡\(Jz​\(x,xℬ\)\)≤dz\.\\operatorname\{rank\}\\\!\\left\(J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\Sigma\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)^\{\\top\}\\right\)\\leq\\operatorname\{rank\}\\\!\\left\(J\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\\right\)\\leq d\_\{z\}\.\(22\)AlthoughΣz​\(x,xℬ\)\\Sigma\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)is diagonal, the covariance entries satisfy

Cov\(ui,uj∣x,xℬ\)≈∑k=1dzJz,i​k\(x,xℬ\)Jz,j​k\(x,xℬ\)σz,k2\(x,xℬ\),\\operatorname\{Cov\}\(u\_\{i\},u\_\{j\}\\mid x,x\_\{\\mathcal\{B\}\}\)\\approx\\sum\_\{k=1\}^\{d\_\{z\}\}J\_\{z,ik\}\(x,x\_\{\\mathcal\{B\}\}\)J\_\{z,jk\}\(x,x\_\{\\mathcal\{B\}\}\)\\sigma\_\{z,k\}^\{2\}\(x,x\_\{\\mathcal\{B\}\}\),\(23\)which are generally nonzero fori≠ji\\neq j\. Although the decoder does not explicitly depend on the molecular state, its Jacobian is evaluated at the state\-dependent latent meanμz​\(x,xℬ\)\\mu\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\)\. Consequently, together with the state dependence ofΣz​\(x,xℬ\)\\Sigma\_\{z\}\(x,x\_\{\\mathcal\{B\}\}\), the induced covariance and its associated low\-dimensional subspace may vary with\(x,xℬ\)\(x,x\_\{\\mathcal\{B\}\}\)\.

### Algorithm and implementation details of FS\-TPS and LaS\-TPS

The training and inference algorithms for FS\-TPS and LaS\-TPS are provided inAlgorithm 1andAlgorithm 2\. In the algorithms,xℬ,t\(m\)=Align⁡\(xt\(m\),xℬ\)x\_\{\\mathcal\{B\},t\}^\{\(m\)\}=\\operatorname\{Align\}\(x\_\{t\}^\{\(m\)\},x\_\{\\mathcal\{B\}\}\)denotes the target conformationxℬx\_\{\\mathcal\{B\}\}re\-expressed in the reference frame of the current configurationxt\(m\)x\_\{t\}^\{\(m\)\}, andξt\(m\)=\(xt\(m\),xℬ,t\(m\)\)\\xi\_\{t\}^\{\(m\)\}=\(x\_\{t\}^\{\(m\)\},x\_\{\\mathcal\{B\},t\}^\{\(m\)\}\)is the common policy input used by both methods\. The default sampled\-policy evaluation usesα=1\\alpha=1, corresponding to theσ\\sigma\-scale factor reported in Table 1\. For FS\-TPS, settingα=0\\alpha=0gives the conditional mean\-force policy\. For LaS\-TPS, it gives the decoded\-latent\-mean policy\.

##### Detailed settings

For the baselines, we use the default hyperparameter settings from TPS\-DPS\. For alanine dipeptide, we use a batch size of 16, 1000 steps per path, an annealing temperature from 600K to 300K, 1000 rollouts, and a buffer size of 1000\. For chignolin, we use a batch size of 4, 5000 steps per path, an annealing temperature from 600K to 300K, 100 rollouts, and a buffer size of 200\. For BBL protein, we use a batch size of 2, 5000 steps per path, and an annealing temperature from 600K to 400K\.

For FS\-TPS,βH\\beta\_\{\\mathrm\{H\}\}is set to1×10−51\\times 10^\{\-5\}for alanine dipeptide,5×10−65\\times 10^\{\-6\}for chignolin, and5×10−55\\times 10^\{\-5\}for BBL protein\. For LaS\-TPS, we evaluateβKL∈\{10−2,10−3,10−4\}\\beta\_\{\\mathrm\{KL\}\}\\in\\\{10^\{\-2\},10^\{\-3\},10^\{\-4\}\\\}for alanine dipeptide and chignolin, andβKL∈\{10−6,10−7\}\\beta\_\{\\mathrm\{KL\}\}\\in\\\{10^\{\-6\},10^\{\-7\}\\\}for BBL protein\. The architecture details are included in Table[5](https://arxiv.org/html/2608.13800#A1.T5)\. We use random seeds2,50,137,4829,44219,72831,560237,910244,88310212,50,137,4829,44219,72831,560237,910244,8831021for alanine dipeptide and2,42,50,137,4829,44219,728312,42,50,137,4829,44219,72831for chignolin and BBL protein\.

##### Replay buffer sampling

For LaS\-TPS, at each training step, we assemble a minibatch by drawing a small fraction of trajectories from a low\-RMSD pool and the rest uniformly from the full replay buffer\. Concretely, after each rollout, every trajectory’s closest approach to the target state is identified by a relaxed Gaussian indicator with bandwidthσ\\sigma, evaluated over the trajectory after Kabsch alignment\. The RMSD at that frame is stored alongside the trajectory in the buffer\. During training,NbiasedN\_\{\\text\{biased\}\}trajectories are sampled uniformly from the bottompp\-fraction of the buffer ranked by RMSD \(the “low pool”\), and the remaining trajectories are sampled uniformly from the full buffer as usual\. For example, for alanine dipeptide, we usedσ=0\.1\\sigma=0\.1,p=0\.3p=0\.3,Nbatch=16N\_\{\\text\{batch\}\}=16, andNbiased=1N\_\{\\text\{biased\}\}=1, giving the reported∼1/16≈6%\\sim 1/16\\approx 6\\%biased fraction\.

Table 5:Network architectures used for TPS\-DPS and LaS\-TPS\.SystemMethodLayer dimensionsAlanine dipeptideTPS\-DPS\(input,128,256,256,256,128,output\)\(\\mathrm\{input\},128,256,256,256,128,\\mathrm\{output\}\)LaS\-TPSEncoder:\(input,64,128,64\)\(\\mathrm\{input\},64,128,64\);latent heads:2×\(64,10\)2\\times\(64,10\);decoder:\(10,64,128,64,output\)\(10,64,128,64,\\mathrm\{output\}\)ChignolinTPS\-DPS\(input,512,1024,2048,1024,512,output\)\(\\mathrm\{input\},512,1024,2048,1024,512,\\mathrm\{output\}\)LaS\-TPSEncoder:\(input,256,512,256\)\(\\mathrm\{input\},256,512,256\);latent heads:2×\(256,24\)2\\times\(256,24\);decoder:\(24,256,512,256,output\)\(24,256,512,256,\\mathrm\{output\}\)BBLTPS\-DPS\(input,512,1024,2048,1024,512,output\)\(\\mathrm\{input\},512,1024,2048,1024,512,\\mathrm\{output\}\)LaS\-TPSEncoder:\(input,1024,512,256\)\(\\mathrm\{input\},1024,512,256\);latent heads:2×\(256,128\)2\\times\(256,128\);decoder:\(128,256,512,1024,output\)\(128,256,512,1024,\\mathrm\{output\}\)

### Evaluation Metrics

Following evaluations in TPS\-DPS\([39](https://arxiv.org/html/2608.13800#bib.bib1)\), for each sampled transition path, we roll out the trajectory under the learned dynamics and record the atomic positions at every step, from which the corresponding forces and potential energies are obtained by re\-evaluating the system’s force field at each configuration\. All metrics are always computed over an ensemble of 64 sampled paths\.

##### RMSD

Structural accuracy is quantified by the root\-mean\-square deviation \(RMSD\) between the final configuration of each sampled path and the target basin’s reference structure, computed over heavy atoms only after optimal rigid\-body superposition via the Kabsch algorithm\([17](https://arxiv.org/html/2608.13800#bib.bib9)\)\.

##### THP

Path success is defined through a target\-hit criterion evaluated in a low\-dimensional collective\-variable space specific to each system\. For alanine dipeptide, the backbone dihedral angles\(ϕ,ψ\)\(\\phi,\\psi\)of the final configuration are compared with those of the target state, and a path is counted as successful if the periodic angular distance satisfies\(Δ​ϕ\)2\+\(Δ​ψ\)2<0\.752\(\\Delta\\phi\)^\{2\}\+\(\\Delta\\psi\)^\{2\}<0\.75^\{2\}\. For the larger proteins chignolin and BBL, the final configuration is projected onto the leading two time\-lagged independent components \(TIC1, TIC2\) obtained from a TICA model\([31](https://arxiv.org/html/2608.13800#bib.bib48)\)trained on backbone\-torsion features of the folded reference structure\. The same distance criterion,\(Δ​TIC1\)2\+\(Δ​TIC2\)2<0\.752\(\\Delta\\mathrm\{TIC\}\_\{1\}\)^\{2\}\+\(\\Delta\\mathrm\{TIC\}\_\{2\}\)^\{2\}<0\.75^\{2\}, is applied in this reduced space\. The target\-hit percentage \(THP\) is the fraction of sampled paths satisfying this criterion, expressed as a percentage\.

##### ETS

For paths classified as successful, we additionally report the energy at the transition state \(ETS\), defined as the maximum potential energy attained along the trajectory\.

##### Mode coverage

For alanine dipeptide specifically, whose two\-dimensional\(ϕ,ψ\)\(\\phi,\\psi\)landscape admits two known, well\-separated first\-order saddle points connecting the C5 and C7ax basins, we further assess whether sampling explores both transition mechanisms rather than collapsing onto a single dominant pathway\. For each successful path, we determine whether its trajectory passes within a fixed radius \(0\.30\.3rad\) of each saddle point in dihedral space at any point along its length; paths that come within this radius of more than one saddle are treated as ambiguous and excluded from this analysis\.

### Computation of stochastic\-force structure metrics

### Per\-state stochastic\-force metrics

Five per\-state metrics including internal effective rank, internal top\-mode explained variance, excess cross\-atom covariance, RMS stochastic amplitude relative to the mean force, and latent effective\-rank utilization, are computed from the same per\-state stochastic\-force ensemble\. The procedures are identical for alanine dipeptide, chignolin, and BBL; only the number of atomsNN, latent dimensiondzd\_\{z\}, and number of evaluated statesSSdiffer across systems\. We useS=200S=200states for alanine dipeptide andS=224S=224states for chignolin and BBL\.

Each metric is first computed independently at every fixed state\. For each trained seed, the resulting per\-state values are then averaged over the state bank\. We finally report the mean and standard deviation across independently trained seeds\.

Throughout this section, the target conformationxℬx\_\{\\mathcal\{B\}\}is held fixed across the entire state bank and is therefore omitted from the notation; all quantities are written as functions of the configurationxxalone\.

##### Per\-state force ensemble\.

For each fixed molecular configurationxx, we query the trained policyKMC=1024K\_\{\\mathrm\{MC\}\}=1024times while holdingxxunchanged, producing applied bias\-force samples

\{𝐟\(k\)​\(x\)\}k=1KMC,𝐟\(k\)​\(x\)∈ℝ3​N\.\\left\\\{\\mathbf\{f\}^\{\(k\)\}\(x\)\\right\\\}\_\{k=1\}^\{K\_\{\\mathrm\{MC\}\}\},\\qquad\\mathbf\{f\}^\{\(k\)\}\(x\)\\in\\mathbb\{R\}^\{3N\}\.For FS\-TPS, each sample is drawn directly from the predicted force\-space Gaussian\. For LaS\-TPS, a fresh latent variable is sampled and subsequently decoded into the applied bias force\. Let

𝐟¯​\(x\)=1KMC​∑k=1KMC𝐟\(k\)​\(x\)\\overline\{\\mathbf\{f\}\}\(x\)=\\frac\{1\}\{K\_\{\\mathrm\{MC\}\}\}\\sum\_\{k=1\}^\{K\_\{\\mathrm\{MC\}\}\}\\mathbf\{f\}^\{\(k\)\}\(x\)denote the sample mean\. The per\-state force covariance is

𝚺⁡\(x\)=1KMC−1​∑k=1KMC\(𝐟\(k\)−𝐟¯\)​\(𝐟\(k\)−𝐟¯\)⊤\.\\boldsymbol\{\\Sigma\}\(x\)=\\frac\{1\}\{K\_\{\\mathrm\{MC\}\}\-1\}\\sum\_\{k=1\}^\{K\_\{\\mathrm\{MC\}\}\}\\left\(\\mathbf\{f\}^\{\(k\)\}\-\\overline\{\\mathbf\{f\}\}\\right\)\\left\(\\mathbf\{f\}^\{\(k\)\}\-\\overline\{\\mathbf\{f\}\}\\right\)^\{\\\!\\top\}\.

##### Removal of rigid\-body components\.

For the three internal\-force metrics below, we remove the six\-dimensional rigid\-body subspace consisting of three translational and three mass\-weighted rotational modes\. Let𝐐⁡\(x\)\\mathbf\{Q\}\(x\)contain the corresponding basis vectors\. The projector onto the internal subspace is

𝐏int​\(x\)=𝐈−𝐐​\(𝐐⊤​𝐐\)−1​𝐐⊤,\\mathbf\{P\}\_\{\\mathrm\{int\}\}\(x\)=\\mathbf\{I\}\-\\mathbf\{Q\}\\left\(\\mathbf\{Q\}^\{\\top\}\\mathbf\{Q\}\\right\)^\{\-1\}\\mathbf\{Q\}^\{\\top\},and the projected covariance is

𝚺int​\(x\)=𝐏int​𝚺​\(x\)​𝐏int⊤\.\\boldsymbol\{\\Sigma\}\_\{\\mathrm\{int\}\}\(x\)=\\mathbf\{P\}\_\{\\mathrm\{int\}\}\\boldsymbol\{\\Sigma\}\(x\)\\mathbf\{P\}\_\{\\mathrm\{int\}\}^\{\\top\}\.Thus, the reported internal\-force statistics measure only genuine molecular deformations rather than global translation or rotation\. For numerical stability, the covariance eigenvalues are computed as the squared singular values of the projected, centered sample matrix divided byKMC−1K\_\{\\mathrm\{MC\}\}\-1, rather than by explicitly forming and diagonalizing the covariance matrix\.

##### Internal effective rank\.

Let\{λi​\(x\)\}\\\{\\lambda\_\{i\}\(x\)\\\}denote the nonzero eigenvalues of𝚺int​\(x\)\\boldsymbol\{\\Sigma\}\_\{\\mathrm\{int\}\}\(x\), and define

pi​\(x\)=λi​\(x\)∑jλj​\(x\)\.p\_\{i\}\(x\)=\\frac\{\\lambda\_\{i\}\(x\)\}\{\\sum\_\{j\}\\lambda\_\{j\}\(x\)\}\.The internal effective rank is the exponential of the spectral Shannon entropy\([38](https://arxiv.org/html/2608.13800#bib.bib47)\),

reffint\(x\)=exp\[−∑ipi\(x\)logpi\(x\)\]\.r\_\{\\mathrm\{eff\}\}^\{\\mathrm\{int\}\}\(x\)=\\exp\\left\[\-\\sum\_\{i\}p\_\{i\}\(x\)\\log p\_\{i\}\(x\)\\right\]\.This quantity is a continuous measure of the effective number of equally weighted stochastic\-force directions\. It equalsqqwhen exactlyqqnonzero eigenvalues have equal magnitude\.

##### Internal top\-mode explained variance\.

The fraction of internal stochastic\-force variance captured by the dominant mode is

EV​@​1​\(x\)=λ1​\(x\)∑iλi​\(x\),\\mathrm\{EV@1\}\(x\)=\\frac\{\\lambda\_\{1\}\(x\)\}\{\\sum\_\{i\}\\lambda\_\{i\}\(x\)\},whereλ1\\lambda\_\{1\}is the largest eigenvalue\. We report this value as a percentage\. Values near100%100\\%indicate that the stochastic force is concentrated along a single internal direction, whereas smaller values indicate a broader spectrum\.

##### Excess cross\-atom correlation\.

We partition𝚺int\\boldsymbol\{\\Sigma\}\_\{\\mathrm\{int\}\}intoN×NN\\times Nblocks𝚺a​b∈ℝ3×3\\boldsymbol\{\\Sigma\}\_\{ab\}\\in\\mathbb\{R\}^\{3\\times 3\}, where each block describes the covariance between the Cartesian force components of atomsaaandbb\. The cross\-atom covariance fraction is

Rcross​\(x\)=∑a≠b‖𝚺a​b​\(x\)‖F2∑a,b‖𝚺a​b​\(x\)‖F2\.R\_\{\\mathrm\{cross\}\}\(x\)=\\frac\{\\sum\_\{a\\neq b\}\\left\\\|\\boldsymbol\{\\Sigma\}\_\{ab\}\(x\)\\right\\\|\_\{F\}^\{2\}\}\{\\sum\_\{a,b\}\\left\\\|\\boldsymbol\{\\Sigma\}\_\{ab\}\(x\)\\right\\\|\_\{F\}^\{2\}\}\.Because a finite number of Monte Carlo samples produces nonzero off\-diagonal covariance even for independent atom\-wise noise, we estimate a permutation null\. Specifically, theKMCK\_\{\\mathrm\{MC\}\}sample indices are independently permuted for each atom, which preserves each atom’s marginal covariance while destroying cross\-atom coupling\. Repeating this procedureB=50B=50times gives

Rcrossexcess​\(x\)=Rcrossreal​\(x\)−1B​∑ℓ=1BRcrossnull,ℓ​\(x\)\.R\_\{\\mathrm\{cross\}\}^\{\\mathrm\{excess\}\}\(x\)=R\_\{\\mathrm\{cross\}\}^\{\\mathrm\{real\}\}\(x\)\-\\frac\{1\}\{B\}\\sum\_\{\\ell=1\}^\{B\}R\_\{\\mathrm\{cross\}\}^\{\\mathrm\{null\},\\ell\}\(x\)\.A value near zero therefore indicates no cross\-atom structure beyond the finite\-sample background\.

##### RMS stochastic amplitude relative to the mean force\.

This metric is computed from the full, unprojected covariance𝚺⁡\(x\)\\boldsymbol\{\\Sigma\}\(x\)\. The total stochastic variance istr⁡𝚺⁡\(x\)\\operatorname\{tr\}\\boldsymbol\{\\Sigma\}\(x\), while the deterministic force magnitude is‖𝐟¯​\(x\)‖22\\\|\\overline\{\\mathbf\{f\}\}\(x\)\\\|\_\{2\}^\{2\}\. We define

Rnoise/mean​\(x\)=tr⁡𝚺⁡\(x\)‖𝐟¯​\(x\)‖22\.R\_\{\\mathrm\{noise/mean\}\}\(x\)=\\sqrt\{\\frac\{\\operatorname\{tr\}\\boldsymbol\{\\Sigma\}\(x\)\}\{\\\|\\overline\{\\mathbf\{f\}\}\(x\)\\\|\_\{2\}^\{2\}\}\}\.Values below one indicate that the stochastic component is smaller in RMS magnitude than the mean bias force, whereas values above one indicate that stochastic variation is the dominant contribution\.

##### Latent effective\-rank utilization\.

For LaS\-TPS, we additionally verify that low\-rank force covariance is not caused by collapse of the sampled latent distribution\. At each state, the sameKMC=1024K\_\{\\mathrm\{MC\}\}=1024latent samples are used to construct the empirical latent covariance𝚺z​\(x\)\\boldsymbol\{\\Sigma\}\_\{z\}\(x\)\. Its effective rank is computed using the same entropy\-based definition above and normalized by the latent dimension:

Rz​\(x\)=reff​\(𝚺z​\(x\)\)dz\.R\_\{z\}\(x\)=\\frac\{r\_\{\\mathrm\{eff\}\}\\\!\\left\(\\boldsymbol\{\\Sigma\}\_\{z\}\(x\)\\right\)\}\{d\_\{z\}\}\.Values close to one indicate that the sampled latent noise remains approximately full\-rank\. Consequently, a low\-rank force covariance together withRz≈1R\_\{z\}\\approx 1reflects the structure induced by the nonlinear decoder, rather than latent\-space sampling collapse\.

### Global effective rank across states

Unlike the five metrics above, the global effective rank is computed jointly across all states in the state bank rather than independently at each state\. For LaS\-TPS, it measures whether the dominant internal stochastic\-force direction is shared across molecular states or varies along the transition path\. This metric is not reported for FS\-TPS because its force covariance is generally not characterized by a single dominant mode\.

For each statexsx\_\{s\}in the state bank, we obtain the unit\-normalized leading right singular vector

𝐯1​\(xs\)∈ℝ3​N\\mathbf\{v\}\_\{1\}\(x\_\{s\}\)\\in\\mathbb\{R\}^\{3N\}from the rigid\-body\-projected, mean\-centered force sample matrix described above\. This vector represents the dominant internal stochastic\-force direction at that state\.

Before comparing directions across states, the dominant force direction at each state is expressed in the common target reference frame using the Kabsch rotation that aligns the corresponding configurationxsx\_\{s\}to the target conformationxℬx\_\{\\mathcal\{B\}\}\.

We construct the cross\-state directional matrix by assigning equal weight to every state:

𝐌=1S​∑s=1S𝐯1​\(xs\)​𝐯1​\(xs\)⊤\.\\mathbf\{M\}=\\frac\{1\}\{S\}\\sum\_\{s=1\}^\{S\}\\mathbf\{v\}\_\{1\}\(x\_\{s\}\)\\mathbf\{v\}\_\{1\}\(x\_\{s\}\)^\{\\top\}\.Let\{ηi\}\\\{\\eta\_\{i\}\\\}denote the eigenvalues of𝐌\\mathbf\{M\}, and define

qi=ηi∑jηj\.q\_\{i\}=\\frac\{\\eta\_\{i\}\}\{\\sum\_\{j\}\\eta\_\{j\}\}\.The global effective rank uses the same effective rank equation in former sections\. In practice, the spectrum is computed from theS×SS\\times SGram matrix of the aligned dominant modes, which has the same nonzero eigenvalues as𝐌\\mathbf\{M\}and avoids explicitly constructing a3​N×3​N3N\\times 3Nmatrix\.

A global effective rank near one indicates that the dominant stochastic\-force direction is approximately shared across all states\. Larger values indicate that the locally dominant direction changes across the state bank, even when the force covariance at each individual state is approximately low\-rank\. The metric is computed separately for each trained seed and then summarized by the mean and standard deviation across seeds\.

### Additional results and ablation studies

##### Temperature\-scaled Gaussian noise baseline

For alanine dipeptide, we perturb the deterministic policy by directly adding temperature\-scaled Gaussian noise to the predicted bias force:

𝐮~θ​\(x,T\)=𝐮θ​\(x\)\+σ​TTref​ϵ,ϵ∼𝒩⁡\(𝟎,𝐈\),\\widetilde\{\\mathbf\{u\}\}\_\{\\theta\}\(x,T\)=\\mathbf\{u\}\_\{\\theta\}\(x\)\+\\sigma\\sqrt\{\\frac\{T\}\{T\_\{\\mathrm\{ref\}\}\}\}\\,\\boldsymbol\{\\epsilon\},\\qquad\\boldsymbol\{\\epsilon\}\\sim\\mathcal\{N\}\(\\mathbf\{0\},\\mathbf\{I\}\),whereσ=10−3\\sigma=10^\{\-3\}andTref=300​KT\_\{\\mathrm\{ref\}\}=300\\,\\mathrm\{K\}\. Across nine seeds, this baseline achieves a mean THP of only43\.6%43\.6\\%and remains strongly seed\-sensitive, with THP ranging from3\.1%3\.1\\%to70\.3%70\.3\\%, indicating that simply adding Gaussian noise to a deterministic policy does not reproduce the performance of FS\-TPS\.

##### Ablation of prioritized replay sampling

For alanine dipeptide, prioritized replay yields an average RMSD of0\.25±0\.030\.25\\pm 0\.03Å, THP of72\.58±5\.01%72\.58\\pm 5\.01\\%, and ETS of21\.63±5\.4921\.63\\pm 5\.49kJ/mol across nine seeds, compared with0\.25±0\.040\.25\\pm 0\.04Å,71\.30±8\.24%71\.30\\pm 8\.24\\%, and18\.54±7\.0618\.54\\pm 7\.06kJ/mol without prioritized sampling\. \(All results are reported across the same set of nine random seeds\.\) Thus, it provides only a marginal increase in THP, while leaving RMSD essentially unchanged and not improving ETS\.

##### Effect of latent dimension\.

Table[6](https://arxiv.org/html/2608.13800#A1.T6)shows that reducing the latent dimension to one increases sensitivity to initialization for alanine dipeptide\.

Table 6:Effect of the latent dimension on LaS\-TPS for alanine dipeptide\. Individual\-seed results and mean±\\pmstandard deviation over four seeds are reported\. The seed exhibiting degraded performance is highlighted in red\.LatentdimensionSeedRMSD\(Å↓\\downarrow\)THP\(%↑\\uparrow\)ETS\(kJ/mol↓\\downarrow\)dz=1d\_\{z\}=120\.2384\.412\.00500\.2187\.510\.991370\.2078\.120\.6148290\.7021\.960\.29Mean±\\pmstd\.0\.34±0\.240\.34\\pm 0\.2467\.98±30\.9667\.98\\pm 30\.9625\.97±23\.2825\.97\\pm 23\.28dz=10d\_\{z\}=1020\.2871\.932\.70500\.3065\.616\.901370\.2868\.816\.0048290\.2570\.319\.10Mean±\\pmstd\.0\.28±0\.020\.28\\pm 0\.0269\.15±2\.6869\.15\\pm 2\.6821\.18±7\.7921\.18\\pm 7\.79

### Results of Doob’s Lagrangian and FAS

We report the results of Doob’s Lagrangian\([6](https://arxiv.org/html/2608.13800#bib.bib4)\)and FAS\([30](https://arxiv.org/html/2608.13800#bib.bib3)\)to conduct a thorough comparison between MD\-rollout TPS methods and endpoint\-conditioned path generators\. The results are shown in Figure[4](https://arxiv.org/html/2608.13800#A1.F4)and Table[7](https://arxiv.org/html/2608.13800#A1.T7)\. For Doob’s Lagrangian, all 64 paths collapse onto a single nearly straight trajectory in the\(ϕ,ψ\)\(\\phi,\\psi\)collective\-variable space\.

![Refer to caption](https://arxiv.org/html/2608.13800v1/doobs.png)

Figure 4:Sampled transition paths on the alanine dipeptide free\-energy surface for Doob’s Lagrangian\.Table 7:Endpoint\-conditioned path generators on alanine dipeptide and chignolin\. These methods explicitly constrain the trajectory endpoints; therefore, their RMSD and THP are determined primarily by the endpoint conditions and are not directly comparable with those of unconstrained transition path sampling methods\. Chignolin results of Doob’s Lagrangian and all results for FAS are taken from Lee et al\.\([22](https://arxiv.org/html/2608.13800#bib.bib5)\)and Park et al\.\([30](https://arxiv.org/html/2608.13800#bib.bib3)\)\. No corresponding results were reported for BBL\.MethodRMSD\(Å↓\\downarrow\)THP\(%↑\\uparrow\)ETS\(kJ/mol↓\\downarrow\)Alanine Dipeptide \(22 atoms\)Doob’sEndpoint\-constrained−17\.20±0\.01\-17\.20\\pm 0\.01FASRMSD≈0\\approx 0, THP≈100%\\approx 100\\%−29\.46±3\.87\-29\.46\\pm 3\.87Chignolin \(166 atoms\)Doob’sEndpoint\-constrained3828\.38±0\.103828\.38\\pm 0\.10FASRMSD≈0\\approx 0, THP≈100%\\approx 100\\%−360\.85±53\.71\-360\.85\\pm 53\.71

1

Input:Number of rollouts

II; training steps per rollout

JJ; replay\-buffer batch size

KK; number of paths per rollout

MM; target conformation

xℬx\_\{\\mathcal\{B\}\}; annealing schedule

λ1=λstart\>⋯\>λI=λend\\lambda\_\{1\}=\\lambda\_\{\\mathrm\{start\}\}\>\\cdots\>\\lambda\_\{I\}=\\lambda\_\{\\mathrm\{end\}\};force\-entropy coefficientβH\\beta\_\{\\mathrm\{H\}\}; andlatent KL coefficientβKL\\beta\_\{\\mathrm\{KL\}\}

2

Initialize an empty replay buffer

𝒟^\\widehat\{\\mathcal\{D\}\}, policy parameters

θ\\theta, and scalar parameter

ww3

4

5

for*i=1,…,Ii=1,\\ldots,I*do6

7

Generate

MMpaths

\{x0:L\(m\)\}m=1M\\\{x\_\{0:L\}^\{\(m\)\}\\\}\_\{m=1\}^\{M\}from biased MD simulations at annealing temperature

λi\\lambda\_\{i\}\. At each MD step

tt, construct

xℬ,t\(m\)=Align⁡\(xt\(m\),xℬ\),ξt\(m\)=\(xt\(m\),xℬ,t\(m\)\),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\),and apply one of the following stochastic control policies:8

9

10

11

FS\-TPS12

13

Predict the conditional force distribution:

πθ​\(𝐮t∣ξt\(m\)\)\\displaystyle\\pi\_\{\\theta\}\\\!\\left\(\\mathbf\{u\}\_\{t\}\\mid\\xi\_\{t\}^\{\(m\)\}\\right\)=𝒩\(𝝁uθ\(ξt\(m\)\),\\displaystyle=\\mathcal\{N\}\\\!\\Bigl\(\\boldsymbol\{\\mu\}^\{u\}\_\{\\theta\}\(\\xi\_\{t\}^\{\(m\)\}\),OPENdiag⁡\[𝝈θu​\(ξt\(m\)\)2\]\)\.\\displaystyle\\operatorname\{diag\}\\\!\\bigl\[\\boldsymbol\{\\sigma\}^\{u\}\_\{\\theta\}\(\\xi\_\{t\}^\{\(m\)\}\)^\{2\}\\bigr\]\\Bigr\)\.14

Draw fresh Gaussian noise:

ϵt\(m\)∼𝒩⁡\(𝟎,𝐈\),\\boldsymbol\{\\epsilon\}\_\{t\}^\{\(m\)\}\\sim\\mathcal\{N\}\(\\mathbf\{0\},\\mathbf\{I\}\),15

and sample the bias force:

𝐮t\(m\)\\displaystyle\\mathbf\{u\}\_\{t\}^\{\(m\)\}=𝝁θu​\(ξt\(m\)\)\\displaystyle=\\boldsymbol\{\\mu\}^\{u\}\_\{\\theta\}\(\\xi\_\{t\}^\{\(m\)\}\)\+𝝈uθ\(ξt\(m\)\)⊙ϵt\(m\)\.\\displaystyle\+\\boldsymbol\{\\sigma\}^\{u\}\_\{\\theta\}\(\\xi\_\{t\}^\{\(m\)\}\)\\odot\\boldsymbol\{\\epsilon\}\_\{t\}^\{\(m\)\}\.16

17

LaS\-TPS18

19

Predict the conditional latent distribution:

qθ​\(𝐳t∣ξt\(m\)\)\\displaystyle q\_\{\\theta\}\\\!\\left\(\\mathbf\{z\}\_\{t\}\\mid\\xi\_\{t\}^\{\(m\)\}\\right\)=𝒩\(𝝁zθ\(ξt\(m\)\),\\displaystyle=\\mathcal\{N\}\\\!\\Bigl\(\\boldsymbol\{\\mu\}^\{z\}\_\{\\theta\}\(\\xi\_\{t\}^\{\(m\)\}\),OPENdiag⁡\[𝝈θz​\(ξt\(m\)\)2\]\)\.\\displaystyle\\operatorname\{diag\}\\\!\\bigl\[\\boldsymbol\{\\sigma\}^\{z\}\_\{\\theta\}\(\\xi\_\{t\}^\{\(m\)\}\)^\{2\}\\bigr\]\\Bigr\)\.20

Draw fresh Gaussian noise:

ϵt\(m\)∼𝒩⁡\(𝟎,𝐈\),\\boldsymbol\{\\epsilon\}\_\{t\}^\{\(m\)\}\\sim\\mathcal\{N\}\(\\mathbf\{0\},\\mathbf\{I\}\),21

and sample the latent control:

𝐳t\(m\)\\displaystyle\\mathbf\{z\}\_\{t\}^\{\(m\)\}=𝝁θz​\(ξt\(m\)\)\\displaystyle=\\boldsymbol\{\\mu\}^\{z\}\_\{\\theta\}\(\\xi\_\{t\}^\{\(m\)\}\)\+𝝈zθ\(ξt\(m\)\)⊙ϵt\(m\)\.\\displaystyle\+\\boldsymbol\{\\sigma\}^\{z\}\_\{\\theta\}\(\\xi\_\{t\}^\{\(m\)\}\)\\odot\\boldsymbol\{\\epsilon\}\_\{t\}^\{\(m\)\}\.22

Decode the sampled latent variable:

𝐮t\(m\)=gθ​\(𝐳t\(m\)\)\.\\mathbf\{u\}\_\{t\}^\{\(m\)\}=g\_\{\\theta\}\\\!\\left\(\\mathbf\{z\}\_\{t\}^\{\(m\)\}\\right\)\.23

24

25

Update the replay buffer:

𝒟^←𝒟^∪\{x0:L\(m\)\}m=1M\.\\widehat\{\\mathcal\{D\}\}\\leftarrow\\widehat\{\\mathcal\{D\}\}\\cup\\\{x\_\{0:L\}^\{\(m\)\}\\\}\_\{m=1\}^\{M\}\.26

for*j=1,…,Jj=1,\\ldots,J*do27

28

Sample a batch of

KKpaths

\{x0:L\(k\)\}k=1K\\\{x\_\{0:L\}^\{\(k\)\}\\\}\_\{k=1\}^\{K\}from

𝒟^\\widehat\{\\mathcal\{D\}\}29

30

For each sampled path and step, construct

xℬ,t\(k\)=Align⁡\(xt\(k\),xℬ\),ξt\(k\)=\(xt\(k\),xℬ,t\(k\)\)\.x\_\{\\mathcal\{B\},t\}^\{\(k\)\}=\\operatorname\{Align\}\\\!\\left\(x\_\{t\}^\{\(k\)\},x\_\{\\mathcal\{B\}\}\\right\),\\qquad\\xi\_\{t\}^\{\(k\)\}=\\left\(x\_\{t\}^\{\(k\)\},x\_\{\\mathcal\{B\},t\}^\{\(k\)\}\\right\)\.31

Form the common log\-variance objective:

ℒ\(θ,w\)=1K∑k=1K\(logp0\(x0:L\(k\)\)𝟏ℬ\(xL\(k\)\)p𝐮θ\(x0:L\(k\)\)−w\)2\.\\mathcal\{L\}\(\\theta,w\)=\\frac\{1\}\{K\}\\sum\_\{k=1\}^\{K\}\\left\(\\log\\frac\{p\_\{0\}\\\!\\left\(x\_\{0:L\}^\{\(k\)\}\\right\)\\mathbf\{1\}\_\{\\mathcal\{B\}\}\\\!\\left\(x\_\{L\}^\{\(k\)\}\\right\)\}\{p\_\{\\mathbf\{u\}\_\{\\theta\}\}\\\!\\left\(x\_\{0:L\}^\{\(k\)\}\\right\)\}\-w\\right\)^\{2\}\.32

Augment the common objective with the regularizer corresponding to the selected stochastic\-control policy:33

34

ℒFS\(θ,w\)=ℒ\(θ,w\)−βiK⁡\(L\+1\)∑k=1K∑t=0Lℋ\[πθ\(⋅∣ξt\(k\)\)\],βi=βHλiλend\.\\mathcal\{L\}\_\{\\mathrm\{FS\}\}\(\\theta,w\)=\\mathcal\{L\}\(\\theta,w\)\-\\frac\{\\beta\_\{i\}\}\{K\(L\+1\)\}\\sum\_\{k=1\}^\{K\}\\sum\_\{t=0\}^\{L\}\\mathcal\{H\}\\\!\\left\[\\pi\_\{\\theta\}\\\!\\left\(\\cdot\\mid\\xi\_\{t\}^\{\(k\)\}\\right\)\\right\],\\qquad\\beta\_\{i\}=\\beta\_\{\\mathrm\{H\}\}\\frac\{\\lambda\_\{i\}\}\{\\lambda\_\{\\mathrm\{end\}\}\}\.35

36

ℒLaS​\(θ,w\)=ℒ⁡\(θ,w\)\+βKLK⁡\(L\+1\)​∑k=1K∑t=0LDKL​\(qθ​\(𝐳t∣ξt\(k\)\)∥𝒩⁡\(𝟎,𝐈\)\)\.\\mathcal\{L\}\_\{\\mathrm\{LaS\}\}\(\\theta,w\)=\\mathcal\{L\}\(\\theta,w\)\+\\frac\{\\beta\_\{\\mathrm\{KL\}\}\}\{K\(L\+1\)\}\\sum\_\{k=1\}^\{K\}\\sum\_\{t=0\}^\{L\}D\_\{\\mathrm\{KL\}\}\\\!\\left\(q\_\{\\theta\}\\\!\\left\(\\mathbf\{z\}\_\{t\}\\mid\\xi\_\{t\}^\{\(k\)\}\\right\)\\,\\middle\\\|\\,\\mathcal\{N\}\(\\mathbf\{0\},\\mathbf\{I\}\)\\right\)\.37

38

Update

θ\\thetaand

wwusing either

∇θ,wℒFS\\nabla\_\{\\theta,w\}\\mathcal\{L\}\_\{\\mathrm\{FS\}\}or

∇θ,wℒLaS\\nabla\_\{\\theta,w\}\\mathcal\{L\}\_\{\\mathrm\{LaS\}\}39

40

end for41

end for42

43

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

Input:Trained policy parameters

θ⋆\\theta^\{\\star\}; number of sampled paths

MevalM\_\{\\mathrm\{eval\}\}; path length

LL; reactant\-state initialization distribution

ρ𝒜\\rho\_\{\\mathcal\{A\}\}; target conformation

xℬx\_\{\\mathcal\{B\}\}; evaluation temperature

λend\\lambda\_\{\\mathrm\{end\}\}; stochastic scale

α\\alpha\(default

α=1\\alpha=1\)

2

Output:Sampled transition paths

𝒯=\{τ\(m\)\}m=1Meval\\mathcal\{T\}=\\\{\\tau^\{\(m\)\}\\\}\_\{m=1\}^\{M\_\{\\mathrm\{eval\}\}\}
3

Initialize an empty path set

𝒯←∅\\mathcal\{T\}\\leftarrow\\varnothing4

5

6

for*m=1,…,Mevalm=1,\\ldots,M\_\{\\mathrm\{eval\}\}*do7

8

Initialize the molecular configuration and velocity:

\(x0\(m\),v0\(m\)\)∼ρ𝒜\.\\left\(x\_\{0\}^\{\(m\)\},v\_\{0\}^\{\(m\)\}\\right\)\\sim\\rho\_\{\\mathcal\{A\}\}\.9

for*t=0,…,L−1t=0,\\ldots,L\-1*do10

11

Construct the policy input:

xℬ,t\(m\)=Align⁡\(xt\(m\),xℬ\),ξt\(m\)=\(xt\(m\),xℬ,t\(m\)\)\.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\)\.At the current configuration, apply one of the following stochastic control policies:12

13

14

15

16

FS\-TPS17

18

Predict the conditional force distribution:

πθ⋆​\(𝐮t∣ξt\(m\)\)\\displaystyle\\pi\_\{\\theta^\{\\star\}\}\\\!\\left\(\\mathbf\{u\}\_\{t\}\\mid\\xi\_\{t\}^\{\(m\)\}\\right\)=𝒩\(𝝁uθ⋆\(ξt\(m\)\),\\displaystyle=\\mathcal\{N\}\\\!\\Bigl\(\\boldsymbol\{\\mu\}^\{u\}\_\{\\theta^\{\\star\}\}\(\\xi\_\{t\}^\{\(m\)\}\),OPENdiag⁡\[𝝈θ⋆u​\(ξt\(m\)\)2\]\)\.\\displaystyle\\operatorname\{diag\}\\\!\\bigl\[\\boldsymbol\{\\sigma\}^\{u\}\_\{\\theta^\{\\star\}\}\(\\xi\_\{t\}^\{\(m\)\}\)^\{2\}\\bigr\]\\Bigr\)\.19

Draw fresh Gaussian noise:

ϵt\(m\)∼𝒩⁡\(𝟎,𝐈\),\\boldsymbol\{\\epsilon\}\_\{t\}^\{\(m\)\}\\sim\\mathcal\{N\}\(\\mathbf\{0\},\\mathbf\{I\}\),20

and sample the bias force:

𝐮t\(m\)\\displaystyle\\mathbf\{u\}\_\{t\}^\{\(m\)\}=𝝁θ⋆u​\(ξt\(m\)\)\\displaystyle=\\boldsymbol\{\\mu\}^\{u\}\_\{\\theta^\{\\star\}\}\(\\xi\_\{t\}^\{\(m\)\}\)\+α𝝈uθ⋆\(ξt\(m\)\)⊙ϵt\(m\)\.\\displaystyle\+\\alpha\\,\\boldsymbol\{\\sigma\}^\{u\}\_\{\\theta^\{\\star\}\}\(\\xi\_\{t\}^\{\(m\)\}\)\\odot\\boldsymbol\{\\epsilon\}\_\{t\}^\{\(m\)\}\.21

22

LaS\-TPS23

24

Predict the conditional latent distribution:

qθ⋆​\(𝐳t∣ξt\(m\)\)\\displaystyle q\_\{\\theta^\{\\star\}\}\\\!\\left\(\\mathbf\{z\}\_\{t\}\\mid\\xi\_\{t\}^\{\(m\)\}\\right\)=𝒩\(𝝁zθ⋆\(ξt\(m\)\),\\displaystyle=\\mathcal\{N\}\\\!\\Bigl\(\\boldsymbol\{\\mu\}^\{z\}\_\{\\theta^\{\\star\}\}\(\\xi\_\{t\}^\{\(m\)\}\),OPENdiag⁡\[𝝈θ⋆z​\(ξt\(m\)\)2\]\)\.\\displaystyle\\operatorname\{diag\}\\\!\\bigl\[\\boldsymbol\{\\sigma\}^\{z\}\_\{\\theta^\{\\star\}\}\(\\xi\_\{t\}^\{\(m\)\}\)^\{2\}\\bigr\]\\Bigr\)\.25

Draw fresh Gaussian noise:

ϵt\(m\)∼𝒩⁡\(𝟎,𝐈\),\\boldsymbol\{\\epsilon\}\_\{t\}^\{\(m\)\}\\sim\\mathcal\{N\}\(\\mathbf\{0\},\\mathbf\{I\}\),26

and sample the latent control:

𝐳t\(m\)\\displaystyle\\mathbf\{z\}\_\{t\}^\{\(m\)\}=𝝁θ⋆z​\(ξt\(m\)\)\\displaystyle=\\boldsymbol\{\\mu\}^\{z\}\_\{\\theta^\{\\star\}\}\(\\xi\_\{t\}^\{\(m\)\}\)\+α𝝈zθ⋆\(ξt\(m\)\)⊙ϵt\(m\)\.\\displaystyle\+\\alpha\\,\\boldsymbol\{\\sigma\}^\{z\}\_\{\\theta^\{\\star\}\}\(\\xi\_\{t\}^\{\(m\)\}\)\\odot\\boldsymbol\{\\epsilon\}\_\{t\}^\{\(m\)\}\.27

Decode the sampled latent variable into the bias force:

𝐮t\(m\)=gθ⋆​\(𝐳t\(m\)\)\.\\mathbf\{u\}\_\{t\}^\{\(m\)\}=g\_\{\\theta^\{\\star\}\}\\\!\\left\(\\mathbf\{z\}\_\{t\}^\{\(m\)\}\\right\)\.28

29

30

31

Combine the learned bias with the molecular force:

𝐅t,total\(m\)=𝐅MD​\(xt\(m\)\)\+𝐮t\(m\)\.\\mathbf\{F\}\_\{t,\\mathrm\{total\}\}^\{\(m\)\}=\\mathbf\{F\}\_\{\\mathrm\{MD\}\}\\\!\\left\(x\_\{t\}^\{\(m\)\}\\right\)\+\\mathbf\{u\}\_\{t\}^\{\(m\)\}\.32

Advance the system by one biased\-MD step:

\(xt\+1\(m\),vt\+1\(m\)\)=MDStepλend⁡\(xt\(m\),vt\(m\),𝐅t,total\(m\)\)\.\\left\(x\_\{t\+1\}^\{\(m\)\},v\_\{t\+1\}^\{\(m\)\}\\right\)=\\operatorname\{MDStep\}\_\{\\lambda\_\{\\mathrm\{end\}\}\}\\\!\\left\(x\_\{t\}^\{\(m\)\},v\_\{t\}^\{\(m\)\},\\mathbf\{F\}\_\{t,\\mathrm\{total\}\}^\{\(m\)\}\\right\)\.33

end for34

35

Construct the sampled path:

τ\(m\)=x0:L\(m\),\\tau^\{\(m\)\}=x\_\{0:L\}^\{\(m\)\},36

and update the path set:

𝒯←𝒯∪\{τ\(m\)\}\.\\mathcal\{T\}\\leftarrow\\mathcal\{T\}\\cup\\\{\\tau^\{\(m\)\}\\\}\.37

end for38

39

return*𝒯\\mathcal\{T\}*40

41

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