# Probing reaction channels via reinforcement learning

Senwei Liang  Applied Mathematics and Computational Research Division,
Lawrence Berkeley National Laboratory, Berkeley, CA 94720, USA  Aditya N.
Singh  Department of Chemistry, University of California Berkeley, Berkeley,
CA 94720, US  Chemical Sciences Division, Lawrence Berkeley National
Laboratory, Berkeley, CA 94720, USA  Yuanran Zhu  Applied Mathematics and
Computational Research Division, Lawrence Berkeley National Laboratory,
Berkeley, CA 94720, USA  David T. Limmer  Department of Chemistry, University
of California Berkeley, Berkeley, CA 94720, US  Chemical Sciences Division,
Lawrence Berkeley National Laboratory, Berkeley, CA 94720, USA  Materials
Science Division, Lawrence Berkeley National Laboratory, Berkeley, CA 94720,
USA  Kavli Energy Nanoscience Institute at Berkeley, Berkeley, CA 94720, USA
Chao Yang  Applied Mathematics and Computational Research Division, Lawrence
Berkeley National Laboratory, Berkeley, CA 94720, USA

###### Abstract

We propose a reinforcement learning based method to identify important
configurations that connect reactant and product states along chemical
reaction paths. By shooting multiple trajectories from these configurations,
we can generate an ensemble of configurations that concentrate on the
transition path ensemble. This configuration ensemble can be effectively
employed in a neural network-based partial differential equation solver to
obtain an approximation solution of a restricted Backward Kolmogorov equation,
even when the dimension of the problem is very high. The resulting solution,
known as the committor function, encodes mechanistic information for the
reaction and can in turn be used to evaluate reaction rates.

##  1 Introduction

The study of rare reactive events is a fundamental topic within the field of
chemical physics.[1] These reactions are dynamical processes that can be
characterized by the transition of a molecular system from one collection of
meta-stable atomic configurations, i.e., a reactant, to another, i.e., a
product. This metastability arises from the features of the potential energy
surface, where metastable states are configurations near the local minimizer
separated by high energy saddle points. When the saddle points lie along the
typical transition paths between reactants and products, they are referred to
as transition states. The direct numerical studies through molecular dynamics
(MD) simulations are typically prohibitive computationally because these
reactions are rare relative to the timescales of the typical thermal
fluctuations of the system[2].

A primary quantity of interest that has been employed to infer a mechanistic
understanding of reactive events is the committor, a function that maps the
phase space of the system to the probability of reacting.[3, 4] This function
encodes the ideal reaction coordinate, and the computation of this function
can be framed as a multidimensional optimization problem.[5, 6, 7] The high-
dimensionality of this problem poses difficulties in its inference, however a
multitude of methods have been developed over the past decade that have made
the computation of this function tractable. Some notable examples include the
string method[5, 6], diffusion maps[8, 9, 10, 11] and neural networks[12, 13,
14, 15, 16, 17, 18, 19, 20, 21].

This optimization of the committor can be framed as two separate but
intertwined problems – (1) finding configurations with high reactive densities
and (2) fitting a nonlinear ansatz to solve the optimization problem on those
configurations to compute the committor. To solve the first problem, one needs
a measure of the reactive density that depends on the committor itself, and to
solve the second problem, one needs to access those configurations. One way to
solve the first problem is to obtain an ensemble of reactive trajectories via
Transition Path Sampling [22, 23, 3, 24], a Monte-Carlo based method that
enables generation of new reactive trajectories from old ones. While this
method provides access to the ideal set of configurations to compute the
committor on, it often suffers from low acceptance rates resulting in long
decorrelation times between subsequent trajectories, and further refinements
based on approximations of the committor are required to enable efficient
sampling. Beyond Transition Path Sampling, a wide range of methods have been
developed in the past two decades to access rare but important configurations
that encode information of reactive events. Such examples include but are not
limited to metadynamics[25, 26], weighted ensemble method[27, 28] and
variational enhanced sampling[29]. These methods rely on applying a carefully
chosen bias potential to the original Hamiltonian along some low-rank ansatz
of the reaction coordinate, often referred to as an order parameter to
generate more atomic configurations near the transition state. While these
methods can obtain accurate estimates of thermodynamic and kinetic
observables, they strongly hinge on the overlap between the order parameter
and the reaction coordinate. As such, the ensembles of configurations are
often limited by the choice of the order parameter.

In this paper, we use a reinforcement learning (RL) based method to obtain an
ensemble of configurations with high reactive densities. This is done through
a two-step method in which we first find saddle points along the reactive
probability density, which we refer to as connective configurations. These
configurations can be thought of as the maxima of reactive probability along
each reactive channel of the reactive process. This optimization procedure
allows for the discovery of saddle points along each reactive channel when
multiple reactive pathways exist. The optimization proceeds through an RL
algorithm where an agent moves from one state to another by taking an action
in order to achieve a certain goal. Each action is associated with a reward,
and which action to take depends on a policy function designed to guide the
agent toward the goal. In the context of searching for connective
configurations, a state is an atomic configuration, an action simply moves the
agent from one configuration to another according to a policy that is updated
(or learned) over time. The optimal policy, which is obtained after performing
several RL episodes with each episode consisting of a sequence of actions,
would allow us (the agent) to move from any arbitrary configuration toward a
connective configuration. In our RL algorithm, the reward function, which we
will describe in detail in section 3, is chosen as a proxy to the true
objective function to be maximized. While such RL based methods have been
previously used to identify transition states [30] and perform transition path
sampling [31, 32, 33] or cloning [34, 35], the proposed method has a different
objective and yields different quantities of interest. Once we have obtained
these configurations, we perform shooting operations from these points to
obtain a set of configurations along a relatively small subspace in the
configuration space, which we refer to as reaction channels. The notion of
reaction channels is similar to the concept of transition tubes introduced in
Ref. [7]. Both are expected to concentrate on the reactive path ensemble.

Finally, once we obtain these configurations within the reaction channels, we
train a feed-forward neural network (NN) to solve the backward-Kolmogorov
equation (BKE). While a range of methods have been proposed to approximate a
committor using an NN, the method demonstrated in this work is unique in that
it solves the exact BKE rather than the variational form[15, 13, 16, 14] or
the Feynman-Kac form[17, 18]. This form of optimization is more accurate,
however similar to other methods, it is strongly sensitive to the
configurations that it is trained on. We find that training on samples
obtained from RL based method proposed in this work outperforms optimization
on samples generated from a grid-based or high-temperature based sampling
method by at least one order of magnitude. Beyond the accuracy of the
committor, the fidelity of this method is also demonstrated through accurate
estimates of reactive rates, that are often non-trivial to converge.

The rest of the paper is organized as follows. In Section 2, we introduce
terminologies and concepts that are relevant to the approach we take to study
chemical reactions. We outline the basic scheme we use with some detail in
Section 3. Section 3.1 is devoted to the details of the RL algorithm we use to
identify connective configurations. In particular, we discuss how to define
the reward function and design an effective policy. We discuss how to use NN
to obtain the values of the committor function at selected configurations
generated within reaction channels in Section 3.2. We demonstrate the
effectiveness of our approach using a few examples in Section 4. We conclude
the paper with some additional perspectives in Section 5. Some of the
computational details and discussions are provided in the appendix.

##  2 Preliminaries

In this section, we introduce terminologies and concepts needed to present our
algorithms in subsequent sections. Specifically, we define a transition path
and reactive trajectory associated with an overdamped Langevin dynamics, the
committor function associated with transition paths, and describe how the
reaction rate constant can be calculated from the committor function. We then
show how the committor function can be computed by using a feed-forward neural
network.

###  2.1 Transition path and reactive trajectory

Let $\bm{x}(t)=(x_{1}(t),x_{2}(t),...,x_{d}(t))\in\mathbb{R}^{d}$ be the
coordinates of a system that satisfy an overdamped Langevin dynamics defined
by

$\displaystyle\gamma_{i}\dot{x_{i}}(t)=-\frac{\partial V(\bm{x}(t))}{\partial
x_{i}}+\xi_{i}(t),\ \ i=1,\cdots,d,$ (1)

where $V:\Omega\to\mathbb{R}$ is a potential energy function, $\gamma_{i}$ is
the friction coefficient for $x_{i}$, $\xi_{i}(t)$ is white noise with mean
$\langle\xi_{i}(t)\rangle$ and variance
$\langle\xi_{i}(t)\xi_{j}(t^{\prime})\rangle=2\beta^{-1}\gamma_{i}\delta(t-t^{\prime})\delta_{ij}$
and $\beta$ is the inverse of the product of temperature and Boltzmann’s
constant.

Let $A$ and $B$ be two disjoint metastable regions of interest in the
configuration space $\Omega\subset\mathbb{R}^{d}$. They correspond to regions
surrounding two distinct local minima of the potential energy $V(\bm{x})$. We
are interested in trajectories that start in $A$ and terminate in $B$. Along
such a trajectory, $\bm{x}(t)$ must escape out of one metastable region and
cross over a transition region $\Omega\setminus(A\cup B)$ before reaching
another metastable region. Such a trajectory is often referred to as a
transition path. For chemical systems, a transition region is associated with
a chemical reaction. A transition path is also referred to as a reactive
trajectory.

The probability density of $\bm{x}$ follows the Boltzmann-Gibbs distribution
$p(\bm{x})=\exp(-\beta V(\bm{x}))/Z$, where $Z=\int_{\Omega}\exp(-\beta
V(\bm{x})){\rm d}\bm{x}$ is a normalization constant or the partition
function. The probability of observing $\bm{x}$ in a transition region
relative to the probability of observing $\bm{x}$ in a metastable region is
very low. As a result, a transition path or a reactive trajectory is a rare
event and generally requires a very long period of simulation time to observe.

Committor function. Reactive trajectories are not unique. In fact, a chemical
reaction is often characterized by an ensemble of reactive trajectories. Along
each reactive trajectory, we are particularly interested in configurations
that are equally likely to evolve (or commit) to either one of the metastable
regions $A$ and $B$. Such likelihood can be characterized by what is known as
a committor function. To give a precise definition of a committor function
$q(\bm{x})$, let $\tau_{D}(\bm{x})$ be the first hitting time of region $D$
when the dynamics is initiated from $\bm{x}$, i.e. $\tau_{D}(\bm{x})$
represents the time it takes for the chemical system to first enter region $D$
when the dynamics starts from $\bm{x}$. A committor function, denoted by
$q(\bm{x})$, is defined as the probability that a trajectory $\bm{x}(t)$,
starting from a point $\bm{x}(0)=\bm{x}$ within a given set $\Omega$, reaches
$B$ before it reaches $A$, namely, $q(\bm{x})={\rm
Prob}\left(\tau_{B}<\tau_{A}\mid\bm{x}(0)=\bm{x}\right)$. It is well known
that $q(\bm{x})$ satisfies the backward Kolmogorov equation (BKE)

$\displaystyle\mathcal{L}(q)$
$\displaystyle:=\sum_{i=1}^{d}\left(-\gamma_{i}^{-1}\frac{\partial
V(\bm{x})}{\partial x_{i}}\frac{\partial q(\bm{x})}{\partial
x_{i}}+\gamma_{i}^{-1}\beta^{-1}\frac{\partial^{2}q(\bm{x})}{\partial
x_{i}^{2}}\right)=0,\ \ $ (2) $\displaystyle\bm{x}$
$\displaystyle\in\Omega\setminus(A\cup B),\quad q(\bm{x})=0,\ \bm{x}\in
A,\quad q(\bm{x})=1,\ \bm{x}\in B.$

When dimension of $\Omega$ is low (2 or 3), this partial differential equation
can be solved numerically by standard methods such as the finite difference or
finite element method. However, when the dimension of $\Omega$ is high, it is
not practical to use these methods to solve for $q(\bm{x})$. We will discuss
an alternative approach that uses a feed-forward neural network to solve for
$q(\bm{x})$ in Section 2.2.

Probability of being reactive. We use $\rho(\bm{x})$ to denote the probability
of observing a reactive trajectory crossing $\bm{x}$. It follows from the
Markovian property of the underlying overdamped Langevin dynamics that
$\rho(\bm{x})$ can be written as the product of $1-q(\bm{x})$, which describes
the probability of the trajectory arriving at $\bm{x}$ from $A$ without going
through $B$ (under the assumption of time reversibility at the statistical
equilibrium [36]), $p(\bm{x})$, which represents the probability of being at
$\bm{x}$, and $q(\bm{x})$, which describes the probability of reaching $B$
first after leaving $\bm{x}$, i.e.,

$\displaystyle\rho(\bm{x})=(1-q(\bm{x}))q(\bm{x})p(\bm{x}).$ (3)

A configuration $\bm{x}$ that has a high reactive density $\rho(\bm{x})$ is of
particular interest because it marks a transition region along a reactive
trajectory associated with an overdamped Langevin dynamics. However, the
definition of $q(\bm{x})$ given in (3) requires the committor function
$q(\bm{x})$ to be known in advance. Because $q(\bm{x})$ is generally difficult
to calculate for an arbitrary $\bm{x}$, it is not easy to calculate
$\rho(\bm{x})$ in practice.

Reaction rate. The committor function is used in Ref. [7] to introduce the
notion of the current or flux associated with a reactive trajectory. At a
configuration $\bm{x}$ along a reactive trajectory, the flux $\bm{J}(\bm{x})$
across $\bm{x}$ is defined as

$J_{i}(\bm{x})=Z^{-1}e^{-\beta
V(\bm{x})}\gamma_{i}^{-1}\beta^{-1}\frac{\partial{q(\bm{x})}}{\partial
x_{i}},$ (4)

where $\bm{J}(\bm{x})=(J_{1}(\bm{x}),\cdots,J_{d}(\bm{x}))$ and $Z$ is again
the partition function.

If $\bm{J}(\bm{x})$ is known for all $\bm{x}$ along a dividing surface $S$
that separates $A$ and $B$, we can use it to evaluate the reaction rate
constant $\kappa$ via

$\displaystyle\kappa=\int_{S}n_{S}(\bm{x})\bm{J}(\bm{x})d\sigma_{S}(\bm{x}),$
(5)

where $n_{S}(\bm{x})$ is normal vector of $S$.

Note that the main contribution to the integral (5) comes from configurations
$\bm{x}$ along the dividing surface that has a relatively large magnitude of
the flux $\bm{J}(\bm{x})$. Because these configurations typically occupy a
small area $S^{\prime}$ on $S$ for rare events [7, 36], we can focus on this
area and approximate (5) by

$\kappa\approx\int_{S^{\prime}}n_{S^{\prime}}(\bm{x})\bm{J}(\bm{x})d\sigma_{S}(\bm{x}).$
(6)

The area $S^{\prime}$ can be defined by the intersection of $S$ and the so-
called reaction channel to be defined in Sec. 3.

We should note that as long as $S^{\prime}$ can be easily identified and
$\bm{J}(\bm{x})$ can be efficiently evaluated for all $\bm{x}\in S^{\prime}$,
the formula (6) provides a more practical way to compute the rate constant
compared to a brute force approach in which the rate constant is computed
according to the alternative definition

$\kappa=\lim_{T\to\infty}\frac{N_{T}}{T},$ (7)

where $N_{T}$ is the number of trajectory segments that leave $A$ and enter
$B$ within the time interval $[0,T]$ [7]. In the latter approach, $\kappa$ can
be approximated by a direct simulation, i.e., we can generate a sufficiently
long Langevin trajectory from a random starting point and count $N_{T}$.
However, when the system contains a high barrier, it becomes extremely rare to
observe a reactive trajectory even with a long time simulation.

If the dividing surface is not accessible, an alternative method is to
calculate the integral of the flux on the transition region, which can be done
as follows:

$\displaystyle\kappa=\int_{\Omega\setminus(A\cup B)}Z^{-1}e^{-\beta
V(\bm{x})}\beta^{-1}\sum_{i=1}^{d}\gamma_{i}^{-1}\left(\frac{\partial{q(\bm{x})}}{\partial
x_{i}}\right)^{2}d\bm{x}\,,$ (8)

for the systems considered here.

###  2.2 Solving committor function via deep learning

In recent year, deep learning has demonstrated remarkable progress in various
domains of scientific exploration [37, 38, 39], thanks to the exceptional
approximation and generalization capabilities of neural networks (NNs) [40].
In particular, deep learning has emerged as a powerful tool for solving a wide
range of partial differential equations (PDEs), including those formulated in
high-dimensional spaces where conventional solvers like finite difference and
finite element methods suffer from the curse of dimensionality [41].
Furthermore, NN-based solvers can be easily adapted to solve PDEs defined on
an irregular domain. These two key attributes of an NN-based PDE solver make
it highly advantageous for the specific problem we aim to address in this
study. As we will show in Section 4, an NN-based PDE solver enables us to
solve a 66-dimensional BKE associated with a 22-atom molecule within a
reaction channel that consists of an ensemble of configurations not uniformly
distributed or organized in a regular domain in the configuration space. The
NN-based PDE solver overcomes these complexities and enables us to effectively
handle this challenging scenario.

Typically in an NN-based PDE solver, the approximation to the solution of the
PDE is represented as an NN $q(\bm{x};\bm{\theta})$ parameterized by a set of
weights and biases denoted by a vector $\bm{\theta}$. The network takes
$\bm{x}$ as the input and generates $q(\bm{x};\bm{\theta})$ as the output. The
NN parameters are determined in an iterative training procedure that minimizes
a loss function with respect to $\bm{\theta}$ for different choices of the
input $\bm{x}$. To use deep learning to solve the BKE (2) on $\Omega$, we
define the loss function as

$\displaystyle\|\mathcal{L}(q(\bm{x};\bm{\theta}))\|^{2}_{L^{2}(\Omega)}+\ell\|q(\bm{x};\bm{\theta})\|^{2}_{L^{2}(A)}+\ell\|q(\bm{x};\bm{\theta})-1\|^{2}_{L^{2}(B)},$
(9)

where $\ell$ is a penalty coefficient used to impose the boundary constraints.
In practice, the $L^{2}$-norm in (9) is evaluated by summing the loss on the
data points sampled randomly and uniformly in $\Omega$ and $\mathcal{L}$ is
computed by auto-differentiation using advanced deep learning frameworks
(e.g., Pytorch [42]). The NN-based optimization (9) can be conducted by
stochastic gradient descent (SGD), such as Adam [43], a variant of SGD based
on momentum. Similarly, when solving BKE (2) on $K$ sub-domains
$\Omega_{1},\cdots,\Omega_{K}\subset\Omega$, the PDE becomes
$\mathcal{L}(q)=0$ for $\bm{x}\in\cup_{s=1}^{K}\Omega_{s}$ along with the
boundary conditions, which is referred to as restricted BKE. We can define the
NN-optimization problem by

$\displaystyle\min_{\bm{\theta}}\sum_{s=1}^{K}\|\mathcal{L}(q(\bm{x};\bm{\theta}))\|^{2}_{L^{2}(\Omega_{s})}+\ell\|q(\bm{x};\bm{\theta})\|^{2}_{L^{2}(A)}+\ell\|q(\bm{x};\bm{\theta})-1\|^{2}_{L^{2}(B)}.$
(10)

When we have data points $\\{\bm{x}^{i,s}\\}_{i=1}^{N_{s}}\subset\Omega_{s}$,
$s=1,\cdots,K$, $\\{\hat{\bm{x}}^{i}\\}_{i=1}^{N_{A}}\subset A$ and
$\\{\tilde{\bm{x}}^{i}\\}_{i=1}^{N_{B}}\subset B$, the loss function in (10)
can be evaluated as

$\displaystyle\frac{1}{\sum_{s=1}^{K}N_{s}}\sum_{s=1}^{K}\sum_{i=1}^{N_{s}}(\mathcal{L}q(\bm{x}^{i,s};\bm{\theta}))^{2}+\frac{\ell}{N_{A}}\sum_{i=1}^{N_{A}}q(\hat{\bm{x}}^{i};\bm{\theta})^{2}+\frac{\ell}{N_{B}}\sum_{i=1}^{N_{B}}(q(\tilde{\bm{x}}^{i};\bm{\theta})-1)^{2}.$
(11)

##  3 Methodology

As we indicated in Sec. 2.1, configurations $\bm{x}$ with high reactive
density $\rho(\bm{x})$ are of interest because they mark a transition region
in which reactive trajectories are more likely to be observed. Intuitively, if
we shoot a trajectory from a configuration $\bm{x}$ with a high reactive
density $\rho(\bm{x})$ and initiate it with a random momentum, it is likely
that the trajectory will stay within a region where reactive trajectories pass
through. Even though such a trajectory may not be part of a reactive
trajectory, the configurations along such a trajectory occupy a small subspace
that is likely to contain several reactive trajectories. We will refer to the
subspace formed by these configurations as a reactive channel. Because
configurations within such a subspace are likely to have high reactive flux
$\bm{J}(\bm{x})$, we will focus on these configurations, and solve a
restricted BKE within a reactive channel using the neural network technique
discussed in Sec. 2.2 to obtain an approximate committor function and and its
gradient at configurations within the channel. With these, one can
approximately calculate the reaction rate constant by evaluating (6).

We should note that the concept of reaction channel introduced here is similar
in spirit to the notion of transition tube introduced in Ref. [7]. A
transition tube is defined to be an ensemble of regions on non-intersection
dividing surfaces between two metastable regions $A$ and $B$ that have
localized flux. Because a transition tube is characterized by configurations
with relatively high reactive flux which depends on the unknown committor
function, it is not easy to identify directly. Although the central curve
within the transition tube can be approximated by the minimum energy path
which can be computed by the string method [5], defining the region of the
tube is still not trivial.

Because a reaction channel is generated by shooting trajectories from a single
configuration, it is relatively easy to produce as long as we can select a
proper configuration to shoot from. Ideally, that configuration should be the
one that has a high reactive density $\rho(\bm{x})$. However, because
$\rho(\bm{x})$ is defined in terms of the committor function, it is not easy
to identify configurations with high $\rho(\bm{x})$ directly because that
would require solving the original BKE (2). In the following, we will present
a reinforcement learning-based technique to identify configurations that are
likely to have a high $\rho(\bm{x})$, and we will refer to these
configurations as connective configurations.

To create reaction channels, we start by performing a shooting procedure from
connective configurations. Within each reaction channel, we can determine the
committor function on each configuration by solving a restricted BKE using a
neural network. Using the NN solution and its gradient, we can then calculate
the reactive flux for every configuration within the reaction channels. As we
will see in the next section, the reaction channel generated by shooting
trajectories from connective configurations is likely to contain
configurations with relatively high reactive flux. This is sufficient to
provide a good estimation of statistics, such as rate, even if not all
configurations within the channel have high reactive flux.

###  3.1 Seeking connective configurations via reinforcement learning

In this section, we show how to use a reinforcement learning (RL) method to
identify connective configurations. Our basic strategy is to treat each
configuration as a state $\bm{x}^{t}$ and train an agent to take a sequence of
actions $\\{\bm{a}^{0},\bm{a}^{1},\cdots,\bm{a}^{n}\\}$ to move from an
arbitrary state $\bm{x}^{0}$ to $\bm{x}^{1}$, $\bm{x}^{2},\cdots$ successively
through the operation $\bm{x}^{t+1}=\bm{x}^{t}+\bm{a}^{t}$, for
$t=0,1,\cdots,n-1$, until it ultimately reaches a desired state $\bm{x}^{n}$
which corresponds to a connective configuration. The sequence of state-action
pairs $\\{(\bm{x}^{t},\bm{a}^{t})\\}$, with $t=0,1,2,..,n-1$ is an instance of
a policy $\pi$ the agent follows, which is initially not optimal. However,
over a multi-episode learning process, the policy is gradually improved based
on the feedback the agent receives from the environment, which consists of
configurations not being visited, through a policy gradient. An optimal policy
allows the agent to move from an arbitrary state to the desired state
efficiently.

In an RL algorithm, an action that an agent takes at a particular state
$\bm{x}$ is associated with a reward $r(\bm{x},\bm{a})$ that measures the
effectiveness of that state action pair. The policy an agent follows at a
particular state $\bm{x}$ is often designed to maximize not just the reward
$r(\bm{x},\bm{a})$, but the expectation of a sequence of discounted future
rewards, i.e.,

$\mathbb{E}_{\tau\sim\pi}\left[R(\tau)|\bm{x}^{0}=\bm{x},\bm{a}^{0}=\bm{a}\right],$
(12)

where $\tau$ is an instance of a policy $\pi$ that is specified by a sequence
of state action pairs
$\tau:=\\{(\bm{x}^{0},\bm{a}^{0}),\cdots,(\bm{x}^{t},\bm{a}^{t}),\cdots\\}$,
and

$R(\tau):=\sum_{t=1}^{\infty}\eta^{t}r(\bm{x}^{t},\bm{a}^{t}),$

for some discount factor $0<\eta\leq 1$. Such expectation of discounted future
rewards is often referred to as a $Q$-value function or $Q$-function in short
and denoted by $Q^{\pi}(\bm{x},\bm{a})$.

We use $A(\bm{x})$ to denote the action that maximizes
$Q^{\pi}(\bm{x},\bm{a})$ for a given policy $\pi$, i.e.,

$\displaystyle A(\bm{x})=\arg\max_{\bm{a}}Q^{\pi}(\bm{x},\bm{a}).$ (13)

It is well known that the optimal $Q$-function
$Q^{*}(\bm{x},\bm{a}):=\max_{\pi}Q^{\pi}(\bm{x},\bm{a})$ satisfies the Bellman
equation [44]

$Q^{*}(\bm{x},\bm{a})={\mathbb{E}}_{\bm{x}^{\prime}}\left[r(\bm{x},\bm{a})+\eta\max_{\bm{a}^{\prime}}Q^{*}\left(\bm{x}^{\prime},\bm{a}^{\prime}\right)\right],$
(14)

where the expectation is taken with respect to the conditional probability of
the new state $\bm{x}^{\prime}$ given $\bm{x}$ and $\bm{a}$. Associated with
this optimal $Q$-function is the optimal action

$\displaystyle A^{*}(\bm{x})=\arg\max_{\bm{a}}Q^{*}(\bm{x},\bm{a}).$ (15)

Because the state space that contains $\bm{x}$ and the action space that
contains $\bm{a}$ in our problem are not discrete and because the analytical
form of the optimal $Q$-function is generally unknown, it is difficult to
solve the Bellman equation (14) or the optimization problem (15) directly.
Finding the optimal $Q$-function is often an intractable problem, and
approximate solutions are commonly used instead. Several methods such as the
deep deterministic policy gradient [45] (DDPG) method, Twin Delayed DDPG [46]
(TD3) method have been developed to solve the problem approximately. In these
methods, $Q(\bm{x},\bm{a})$ and $A(\bm{x})$ are represented by neural networks
parameterized by $\Psi$ and $\Phi$ respectively. These neural networks are
trained by a set of data $\mathbb{P}$ that contains a collection of
$(\bm{x},\bm{a},r(\bm{x},\bm{a}),\bm{x}^{\prime})$, where $\bm{x}^{\prime}$
denotes a new state reached by the agent after taking the action $\bm{a}$.
Roughly speaking, the parameters $\Psi$ and $\Phi$ are optimized in an
alternate fashion by maximizing the objectives derived from (15) and the
Bellman equation. We have included a diagram in Appendix A to provide a visual
illustration of the Q-learning training process. To be specific, the training
process is used to solve the following optimization problems alternately

$\max_{\Phi}\underset{\left(\bm{x},\bm{a},r,\bm{x}^{\prime}\right)\sim\mathbb{P}}{\mathbb{E}}Q(\bm{x},A(\bm{x};\Phi);\Psi)$
(16)

and

$\min_{\Psi}\underset{\left(\bm{x},\bm{a},r,\bm{x}^{\prime}\right)\sim\mathbb{P}}{\mathbb{E}}\left[Q(\bm{x},\bm{a};\Psi)-\left(r(\bm{x},\bm{a})+\eta
Q\left(\bm{x}^{\prime},A(\bm{x}^{\prime};\Phi);\Psi\right)\right)\right]^{2}.$
(17)

Reward. Note that the solution of the second problem (17) depends on how the
reward function $r(\bm{x},\bm{a})$ is defined. Because our ultimate goal is to
identify connective configurations that are expected to have a high reactive
density $\rho(\bm{x})$, ideally, we would like to use $\rho(\bm{x}^{\prime})$,
where $\bm{x}^{\prime}=\bm{x}+\bm{a}$, as the reward function. The difficulty
is that, as we indicated earlier, $\rho(\bm{x})$ is defined in terms of the
committor function $q(\bm{x})$ which is unknown in general. Therefore, setting
$r(\bm{x},\bm{a})$ to the exact $\rho(\bm{x}^{\prime})$ is not practical.

However, as $q(\bm{x})$ is defined as the probability of a trajectory starting
from a point $\bm{x}(0)=\bm{x}$ and reaching $B$ before reaching $A$, we can
estimate this probability numerically by shooting trajectories from $\bm{x}$,
and performing statistical analysis of these trajectories. To be specific, we
propose to use a shooting procedure to estimate $q(\bm{x})$ by counting the
number of trajectories originating from $\bm{x}$ with a random momentum and
terminating in one of the metastable regions $A$ or $B$ within a fixed number
of time steps.

We shoot $N$ trajectories from $\bm{x}$ with a random momentum or force. Let
$N_{A}$ be the number of trajectories reaching $A$ first rather than $B$
within a fixed number of time steps $T$, and $N_{B}$ be the number of
trajectories reaching $B$ first rather than $A$. Typically, $T$ is far smaller
than the time scale required to observe a reactive trajectory. Clearly,
$N_{A}+N_{B}\leq N$ since some trajectories may hover around the starting
configuration for a long time and never reach either $A$ or $B$ within $T$
time interval. We can view $N_{A}/(N_{A}+N_{B})$ as an approximation to
$q(\bm{x})$ and $N_{B}/(N_{A}+N_{B})$ as an approximation to $1-q(\bm{x})$.
Consequently, we can use

$\hat{\rho}(\bm{x})=\frac{N_{A}N_{B}}{(N_{A}+N_{B})^{2}}p(\bm{x})$ (18)

as a proxy for $\rho(\bm{x})$. As a result, the reward function associated
with the state action pair $(\bm{x},\bm{a})$ can be defined as

$\displaystyle
r(\bm{x},\bm{a}):=\frac{N_{A}N_{B}}{(N_{A}+N_{B})^{2}}\cdot\frac{\exp(-\beta
V(\bm{x}^{\prime}))}{Z},$ (19)

where $\bm{x}^{\prime}=\bm{x}+\bm{a}$ is a new state. In practice, we can
ignore the constant $Z$ in (19).

Once the neural networks are properly trained, the optimal action to be taken
at each $\bm{x}$ satisfies

$Q^{*}(\bm{x},A(\bm{x}))=\max_{\bm{a}}Q^{*}(\bm{x},\bm{a}).$

RL is a multi-episode iterative learning process. In each episode, a random
configuration $\bm{x}^{0}$ is chosen to start the learning process. A neural
network that represents the action function $A(\bm{x})$ takes the
configuration as the input and generates an action $\bm{a}^{0}$ as the output.
Taking such an action yields a new configuration $\bm{x}^{1}$ for which a
reward $r$ can be obtained by shooting several trajectories from $\bm{x}^{1}$
and evaluating (19). This process can be repeated several times until we
generate a sequence of state action pairs $\\{(\bm{x}^{t},\bm{a}^{t})\\}$.
Figure 1 gives a schematic illustration of the process of generating a
sequence of state action pairs within a single episode.

Figure 1:  A schematic description of the proposed RL action-reward feedback
loop. In each episode, a sequence of states (configurations)
$\\{\bm{x}^{0},\cdots,\bm{x}^{t},\cdots\\}$ is generated iteratively. Each
iteration consists of the following stages. ① Initialization: a configuration
$\bm{x}^{0}$ is randomly drawn from the configuration set; ② Evolution: the
parameterized policy yields the action $\bm{a}^{t}$ after observing the state
$\bm{x}^{t}$. After taking the action, the agent moves to a new configuration
$\bm{x}^{t+1}{=}\bm{x}^{t}+\bm{a}^{t}$; ③ Interaction: we shoot multiple
trajectories from $\bm{x}^{t+1}$, and evaluate its reward
$r(\bm{x}^{t},\bm{a}^{t})$. The quadruples
$(\bm{x}^{t},\bm{a}^{t},r(\bm{x}^{t},\bm{a}^{t}),\bm{x}^{t+1})$ is used in
Q-learning to update the policy.

The generated sequence of state action pairs, together with the rewards
evaluated for states form the training data set $\mathbb{P}$ that we use to
optimize the neural network representions of the Q-function $Q(\bm{x},\bm{a})$
and action $A(\bm{x})$ by solving the minimization problems (16) and (17) in
an alternate fashion.

We use the TD3 method to perform the optimization of $A(\bm{x})$ and
$Q(\bm{x},\bm{a})$. This variant of the DDPG method uses a variety of
techniques to improve the stability of the training process and mitigate the
risk of potentially over-estimating the $Q$-value function. In particular, TD3
uses the moving average of parameters associated with multiple neural networks
to solve (16) and (17). Furthermore, TD3 uses a “delayed” policy update, where
the policy network (which is used to determine the optimal action) is only
updated after a certain number of Q-network updates. This delayed updating
scheme helps to stabilize the training process and reduces the likelihood of
the policy network training being stuck in an undesirable local minimum.

The main steps of a multi-episode RL algorithm for seeking the optimal policy
that allows us to quickly identify connection configurations from an arbitrary
starting configuration is shown in Algorithm 1.

Algorithm 1 An RL algorithm for seeking connective configurations

Input: Random parameter initialization $\Psi_{1}$ and $\Psi_{2}$ of critic
networks and $\Phi$ of actor network.

Output: Action network $A(\cdot;\Phi)$.

1:Initialize target networks $\Psi^{\prime}_{1}\leftarrow\Psi_{1}$,
$\Psi^{\prime}_{2}\leftarrow\Psi_{2}$ and $\Phi^{\prime}\leftarrow\Phi$

2:Replay buffer $\mathbb{P}\leftarrow\emptyset$

3:for episode from 1 to the maximal number of episodes do

4: $\bm{x}^{0}\sim{\rm Uniform}(\Omega)$ $\triangleright$ Start a new episode

5: for $t$ from $0$ to $L-1$ do

6: Obtain a new action $\bm{a}^{t}=A(\bm{x}^{t};\Phi)$ and a new state
$\bm{x}^{t+1}=\bm{x}^{t}+\bm{a}^{t}$ $\triangleright$ New state update rule

7: if $p(\bm{x}^{t+1})$ is smaller than a threshold then Break

8: end if

9: Compute the reward $r(\bm{x}^{t},\bm{a}^{t})$ using (19) $\triangleright$
Compute the reward by shooting

10: Append $(\bm{x}^{t},\bm{a}^{t},r(\bm{x}^{t},\bm{a}^{t}),\bm{x}^{t+1})$ to
$\mathbb{P}$

11: end for

12: for $tt$ from $0$ to $t$ do

13: Sample a mini-batch $(\bm{x},\bm{a},r,\bm{x}^{\prime})$ of size $B$ from
$\mathbb{P}$

14: $y\leftarrow
r+\eta\min\\{Q(\bm{x}^{\prime},A(\bm{x}^{\prime};\Phi^{\prime});\cdot)):\Psi_{1}^{\prime},\Psi_{2}^{\prime}\\}$

15: Update $\Psi_{i}$ with the loss
$\frac{1}{B}\sum(Q(\bm{x},\bm{a};\Psi_{i})-y)^{2}$ for $i=1,2$

16: if $tt$ mod $policy\\_delay=0$ then

17: Update $\Phi$ with the loss $\frac{1}{B}\sum-Q(s,A(s;\Phi);\Psi_{1})$

18: Update target networks:
$\Phi^{\prime}\leftarrow\tau\Phi+(1-\tau)\Phi^{\prime}$ and
$\Psi^{\prime}_{i}\leftarrow\tau\Psi_{i}+(1-\tau)\Psi_{i}^{\prime}$ for
$i=1,2$

19: end if

20: end for

21:end for

###  3.2 Generating reaction channels and computing reaction rate constant

After performing several episodes of RL using Algorithm 1, we obtain an
optimal action function $A(\cdot;\Phi^{*})$ parameterized by $\Phi^{*}$. Such
a function yields the policy we follow at each configuration $\bm{x}$ to
quickly move towards a connective configuration. Figure 2(a) shows how such a
policy (represented as the vector field) looks like for a simple potential
energy surface. The value of $A(\bm{x})$, which is a vector with two
components, is plotted as an arrow for each $\bm{x}$ uniformly sampled in
$[-2.0,2.0]\times[-1.5,2.0]$. We see the arrows point to two configurations
marked by crosses. These correspond to two connective configurations.

Once we identify a connective configuration, we perform additional shooting
operations from that configuration to generate multiple trajectories
originating from the connective configuration. The configurations generated
along these trajectories are considered as samples within a reactive channel
between $A$ and $B$. If multiple connective configurations are identified,
each one of them can be used to generate a distinct reactive channel.

The configurations generated within reactive channels can be used to compute
the committor function $q(\bm{x})$ by solving a restricted BKE on these
configurations using a deep neural network as we described in Section 2.2. The
utilization of NN-based PDE solver is often associated with an implicit bias
towards fitting smooth functions that exhibit fast decay in the frequency
domain. [47] Consequently, This bias can make it challenging for NN models to
capture drastic changes in the committor function. However, by choosing an
appropriate training dataset, one can mitigate this bias. [48] In Appendix
D.2, we demonstrate the advantage of using configurations within reaction
channels to train the NN designed to solve the restricted BKE.

The neural network not only returns $q(\bm{x})$ for each $\bm{x}$ within all
reactive channels, but also its gradient $\nabla q(\bm{x})$. This will allow
us to compute the reactive flux at each $\bm{x}$ within the reaction channel.
As a result, by using (6), we can calculate the rate constant.

##  4 Numerical results

In this section, we present several numerical experiments that demonstrate the
effectiveness of the RL method introduced in Sec. 3 for identifying connective
configurations and using them to generate configurations that characterize
reactive channels. Our experiments were performed on model potentials (triple-
well and rugged Muller potentials) as well as the Alanine dipeptide molecule
in vacuum. While we demonstrate the RL results with initializations uniformly
sampled from $\Omega$ in this section, we have the flexibility to relax this
constraint by considering initializations from metastable states (see Appendix
D.1).

###  4.1 Potential with multiple reaction channels

We consider the triple-well potential defined by

$\displaystyle\begin{split}V(x_{1},x_{2})=3e^{-x_{1}^{2}-\left(x_{2}-\frac{1}{3}\right)^{2}}-3e^{-x_{1}^{2}-\left(x_{2}-\frac{5}{3}\right)^{2}}-5e^{-(x-1)^{2}-x_{2}^{2}}-5e^{-(x_{1}+1)^{2}-x_{2}^{2}}+0.2x_{1}^{4}+0.2(x_{2}-\frac{1}{3})^{4}.\end{split}$
(20)

We focus on the domain $\Omega=[-2,2]\times[-1.2,2]$. Figure 2(a) shows this
potential as a color-mapped image. The two meta-stable regions $A$ and $B$ are
defined by

$\displaystyle A=\\{\bm{x}:V(x_{1},x_{2})<-2\text{ and }x_{1}\leq-0.1\\},\quad
B=\\{\bm{x}:V(x_{1},x_{2})<-2\text{ and }x_{1}\geq 0.1\\}.$

We can see from Figure 2(a) that there are two transition paths from $A$ to
$B$. The top transition path goes through the third well in the top part of
the image and two transition states between the third well and $A$ ($B$). The
bottom transition path goes directly from $A$ to $B$ and crosses the
transition state. Because the potential function is symmetric with respect to
$x_{1}=0$, all configurations along $\\{\bm{x}:x_{1}=0\\}$ have an equal
probability of reaching either $A$ or $B$. As a result, $\\{\bm{x}:x_{1}=0\\}$
represents the half-isocommittor surface $\\{\bm{x}:q(\bm{x})=0.5\\}$. We
experimented with both a relatively low-temperature regime $\beta=6.67$ and a
relatively high-temperature regime $\beta=1.67$. We discuss the results for
$\beta=6.67$ here, which is more challenging for studying rare events. We
report a similar observation for the $\beta=1.67$ case in Appendix B.1. By
default, the friction coefficient is set as 1.

Identifying reaction channels. In our experiments, the initial configuration
for each RL episode is randomly sampled from a uniform distribution of
configurations in $\Omega$. The reward $r(\bm{x},\bm{a})$ (19) is obtained by
shooting $N=50$ trajectories from $\bm{x}$. The maximal number of evolution
steps is set to $L=15$. The maximum number of episodes used in RL is set to
1000. While we initially set a large number of episodes for the RL training
process, it is worth noting that the convergence of RL does not require an
extensive number of episodes. Additional discussion can be found in Appendix
B.3.

Figure 2(a) shows the learned action $A(\bm{x};\Phi^{*})$ as a vector field
that represents an optimal policy. Two attractors can be seen from this policy
field. They correspond to two connective configurations located in two
different transition paths.

From each of the identified connective configuration, we shoot $50$
trajectories by simulating the overdamped Langevin dynamics (1) using the
Euler-Maruyama scheme. We choose a uniform time step size $\Delta t=5\times
10^{-3}$, and propagate the solution from $T=0.0$ to $T=2.0$. The total number
of configurations generated along these trajectories is 40,000. These
configurations lie either in the metastable region $A$ or $B$ or two
transition regions between $A$ and $B$ as shown in Fig. 2(b). These two
transition regions correspond to the two reactive channels associated with
this potential energy surface.

Solving BKE. We then solve the restricted BKE in the identified reaction
channels by using the NN-based solver discussed in Sec. 2.2. The loss function
(11) is optimized by the Adam optimizer [43]. The hyperparameters used in the
NN are listed in Appendix C. Figure 2(d) shows the contour plots of the NN
solution (colored contour lines) and the potential (dotted contour lines). The
$0.5$-level set of NN solution marked in the figure almost coincides with the
true half-isocommittor $\\{\bm{x}:x_{1}=0\\}$. As a reference, we also used
the finite difference method (FDM) to solve the BKE on the entire domain
$\Omega$. The absolute difference between the NN and FDM solutions on a
$100\times 100$ uniform mesh is shown in Figure 2(e). We observe small errors
in the NN solution in the regions that contain a large number of sampled
configurations (such as the top transition path) and relatively large errors
in regions where configurations are sparsely sampled (such as the region near
$(0.0,0.5)$).

Rate estimation. The NN approximation to the committor function and its
gradient on configurations with two reactive channels are used to calculate
the reaction rate by integrating (5) over the dividing sub-surfaces identified
with each reaction channel based on the data. We compare the computed rate
with the one obtained from a direct dynamics simulation (see Sec. 2.1) and
that computed from the committor function obtained from the finite difference
solution of the KBE on the entire domain $\Omega$ in Table 1. We list the
computed rates for both $\beta=1.67$ and $\beta=6.67$. In the direct dynamics
simulation, we generate a long trajectory by time evolving the solution to the
overdampled Langevin equation to $T=2\times 10^{7}$. In the FDM calculation,
the rate is computed with numerical integration of (5) on the entire line
segment $\\{\bm{x}:x_{1}=0,-1.2\leq x_{2}\leq 2\\}$. From these numerical
experiments, we find that the generated configurations mainly cross the line
segment $\\{\bm{x}:x_{1}=0,-1.0\leq x_{2}\leq 2.2\\}$ for $\beta=1.67$ and
line segments $\\{\bm{x}:x_{1}=0,-0.8\leq x_{2}\leq 0.0,0.8\leq x_{2}\leq
1.8\\}$ for $\beta=6.67$. We then compute the rates using the NN solutions on
these segments. As we can see, the NN solution on reaction channels gives
comparable rates as the one obtained by the FDM and a direct simulation.
Finally, we validate the rate calculation with our NN solution using the
formula (8).

(a) Learned action function (vector field) and the connective configurations
it reveals (crosses).

(b) Sampled configurations generated by shooting from two identified
connective configurations.

(c) The FDM solution of (2) on $\Omega$.

(d) The NN solution of (2).

(e) The difference between (c) and (d).

Figure 2:  The results obtained from running the RL algorithm for Triple-well potential at inverse temperature $\beta=6.67$. (a) The action function (vector field) learned by our proposed RL. The action function reveals two connective configurations (crosses). (b) The configurations generated by shooting trajectories initiated at the identified connective configurations. (c-e) The FDM (c) and NN (d) solutions of the BKE (2) and their difference (e). Table 1: Reaction rate of the triple-well potential under various temperatures. “N/A” means that no reactive trajectory is observed in the long trajectory of time $T=2\times 10^{7}$.  Methods |  | $\beta$
---|---|---
| 1.67 |  | 6.67
Direct simulation |  | $2.18\times 10^{-2}$ |  | N/A
Finite difference (Eqn. (6)) |  | $2.16\times 10^{-2}$ |  | $7.43\times 10^{-8}$
Our NN (Eqn. (6)) |  | $2.00\times 10^{-2}$ |  | $7.11\times 10^{-8}$
Our NN (Eqn. (8)) |  | $1.99\times 10^{-2}$ |  | $7.32\times 10^{-8}$

(a) Learned action function (vector field) and the connective configurations
it reveals (crosses).

(b) Sampled configurations generated by shooting trajectories from the
identified connective configuration.

(c) The FDM solution of (2) on $\Omega$.

(d) The NN solution of (2).

(e) The difference between (c) and (d).

Figure 3:  The results obtained from running the RL algorithm for the Rugged
Muller potential at the inverse temperature $\beta=0.25$. (a) The action
function (vector field) learned by the RL algorithm. It reveals one connective
configuration (cross). (b) The configurations generated by shooting
trajectories from the identified connective configurations. (c-e) The FDM (c)
and NN (d) solutions of the BKE (2) and their difference (e).

###  4.2 Potential with rough landscape

In the second example, we consider the rugged Muller potential on the domain
$\Omega=[-1.5,1]\times[-0.5,2]$. The potential function is defined by

$\displaystyle V\left(x_{1},x_{2}\right)=$
$\displaystyle\sum_{i=1}^{4}D_{i}\exp\left[a_{i}\left(x_{1}-X_{i}\right)^{2}+b_{i}\left(x_{1}-X_{i}\right)\left(x_{2}-Y_{i}\right)+c_{i}\left(x_{2}-Y_{i}\right)^{2}\right]+\gamma\sin\left(2k\pi
x_{1}\right)\sin\left(2k\pi x_{2}\right).$ (21)

Here the parameters $\gamma$ and $k$ control the roughness of the landscape,
which are set to $9$ and $5$, respectively. Other model parameters
($a_{i},b_{i},c_{i},X_{i},Y_{i},D_{i}$, $i=1,\cdots,4$) are exactly the same
as the ones used in Refs. [14, 15]. By default, the friction coefficient is
set as 1.

Identifying reaction channels. We discuss the numerical results in the low
temperature regime $\beta=0.25$ (Fig. 3) here and refer readers to Appendix
B.2 for results obtained for $\beta=0.1$. In this example, the reward in the
RL algorithm is calculated by shooting $N=20$ trajectories up to $T=0.25$. The
step size of used in each trajectory is set to $5\times 10^{-5}$. Each RL
episode consists of $L=20$ steps (actions). We ran the RL algorithm for $1000$
episodes. Figure 3(a) shows that the learned policy points to a single
connective configuration. From that configuration, we shoot $50$ trajectories
using the Euler-Maruyama scheme. These trajectories contain 100,000
configurations that lie in the metastable regions $A$ and $B$ as well as the
transition region in between (Fig. 3(b)). The latter is viewed as the reactive
channel for this particular potential energy surface. The hyperparameter
setting for the NNs used in the RL algorithm is listed in Appendix C.

Solving BKE. We use the configurations contained in the reactive channel to
solve the BKE via a NN. Figures 3(c) and 3(d) show the contour plots of the
solution to the BKE obtained from both the FDM and the NN, respectively. The
difference between the two is also shown in Figure 3(e). From these plots, we
observe that the NN solution agrees well with the FDM solution. In particular,
the NN solution captures drastical changes near the point $(-0.8,0.6)$.

Rate estimation. The computed committor functions obtained by the FDM and the
NN approach are used to compute the reaction rate under different $\beta$
values. The computed rates are compared with those obtained from direct
numerical simulations of the corresponding overdamped Langevin dynamics in
Table 2. In direction numerical simulations, we ran a long trajectory until
$T=1.5\times 10^{4}$ with a time step size of $10^{-5}$. When $\beta=0.25$, no
reactive trajectory is observed. When the committor function is obtained from
the FDM, the rate is computed from the numerical integration of (5) on the
dividing surface $\\{x_{1}=0.0,-0.5\leq x_{2}\leq 2.0\\}$. When using an NN to
compute the commitor function within the reactive channel, the reaction rate
is computed from the numerical integration of (5) along the line segments
$\\{x_{1}=-0.5,0.0\leq x_{2}\leq 0.8\\}$ for $\beta=0.1$. When $\beta=0.25$,
we calculate the rate by integrating (5) along the line segment
$\\{x_{1}=-0.5,0.3\leq x_{2}\leq 0.65\\}$. We observe that the rates computed
by all three methods are comparable when $\beta=0.1$. When $\beta=0.25$, the
rates obtained from the NN and FDM approximation of the committor function
have the same magnitude.

Table 2: Reaction rate of the rugged Muller potential under various temperatures. “N/A” means that no reactive trajectory is observed in the long trajectory of time $T=1.5\times 10^{4}$. Methods |  | $\beta$
---|---|---
| 0.1 |  | 0.25
Direct simulation |  | $4.75\times 10^{-3}$ |  | N/A
Finite difference (Eqn. (6)) |  | $4.31\times 10^{-3}$ |  | $1.78\times 10^{-10}$
Our NN (Eqn. (6)) |  | $4.36\times 10^{-3}$ |  | $6.27\times 10^{-10}$

###  4.3 Alanine dipeptide in vacuum

In this example, we show how the RL algorithm introduced above can be used to
identify a reaction channel of an Alanine dipeptide (ADP) molecule in vacuum
that corresponds to its isomerization process. In our numerical experiment, we
set the temperature to $300$ K. The ADP molecule contains $22$ atoms (see Fig.
4(e)). Therefore, the dimension of the configuration space is 66. The
isomerization process is principally described by two the dihedral angles
$(\phi,\psi)\in[-180^{\circ},180^{\circ}]^{2}$ of a subset of atoms (indexed
by $4,6,8,14$ and $6,8,14,16$, respectively.) With a slight abuse of notation,
we use $\phi(\bm{x})$ and $\psi(\bm{x})$ to donate the mapping from a
configuration $\bm{x}$ to the two specific torsion angles. Figure 4(a) shows
the potential energy landscape in the $(\phi,\psi)$-space. The plot is
constructed as follows. We first generate a long trajectory at a relatively
high temperature ($1200$ K). We denote the set of configurations along this
trajectory by $\mathbb{S}$. The potential energy for each configuration in
$\mathbb{S}$ is stored. Next, we discretize $(\phi,\psi)$ in
$(-180^{\circ},180^{\circ}]\times(-180^{\circ},180^{\circ}]$ by generating a
$100\times 100$ uniform grid. For each $(\phi_{i},\psi_{j})$ pair on the grid,
we define a neighborhood

$C(\phi_{i},\psi_{j})=\\{\bm{x}\in\mathbb{S}:|\phi(\bm{x})-\phi_{i}|\leq
5^{\circ},|\psi(\bm{x})-\psi_{j}|\leq 5^{\circ}\\}.$

If $C(\phi_{i},\psi_{j})\neq\emptyset$, we set the potential energy associated
with $(\phi_{i},\psi_{j})$ to $V_{\min}$, where

$V_{\min}=\min_{\bm{x}\in C(\phi_{i},\psi_{j})}V(\bm{x}).$

Otherwise, the potential energy of $(\phi_{i},\psi_{j})$ is undefined and
colored by white in Figure 4. The two metastable regions $A$ (solid box in
Fig. 4(a)) and $B$ (dotted box) are defined by
$A=\\{\bm{x}:-150^{\circ}\leq\phi(\bm{x})\leq-65^{\circ},0^{\circ}\leq\psi(\bm{x})\leq
150^{\circ}\\},B=\\{\bm{x}:30^{\circ}\leq\phi(\bm{x})\leq
100^{\circ},-150^{\circ}\leq\psi(\bm{x})\leq 0^{\circ}\\}$. Figure 4(b) shows
the snapshots of one trajectory initiated at a configuration near $A$ of
$T=1\times 10^{8}$ fs when the temperature is $300$ K and we can see that the
no reactive trajectory is observed.

Identifying reaction channels. Our RL algorithm aims to find a connective
configuration in the $(\phi,\psi)$-space. Subsequently, we use this
configuration to generate additional configurations that bridge two metastable
states. We should note that, for the ADP system, the actions taken in the RL
algorithm are defined in a low-dimensional space specified by $(\phi,\psi)$
whereas the shooting procedure used to evaluate the reward takes place in the
66-dimensional configuration space. To address this disparity in
dimensionality, it is necessary to establish a one-to-one mapping between each
$(\phi,\psi)$ pair and a configuration in the phase space. To this end, we
first construct a configuration set $\mathbb{P}$ by generating trajectories at
high temperature $1200$ K. We then map a given torsion coordinate in
$(\phi,\psi)$ back to the configuration space by choosing a configuration from
$\mathbb{P}$ with lowest potential energy. We simulate the Langevin dynamics
with a step size of $2$ fs and friction $10\text{ ps}^{-1}$ using the package
Openmm Python API [49]. The reward for each action is computed by shooting 10
trajectories of $T=2\times 10^{3}$ fs with kinetic initialization randomly
sampled from the Boltzmann-Gibbs distribution. Each RL episode consists of
$L=10$ steps (actions). Figure 4(c) shows the learned action function reveals
$2$ different connective configurations, i.e., $(-10^{\circ},-62^{\circ})$ and
$(139^{\circ},-120^{\circ})$. However, our primary interest is in
$(-10^{\circ},-62^{\circ})$ configuration, as it has been extensively studied
in the existing literature [50, 14, 20]. We shoot 100 trajectories of
$T=2\times 10^{3}$ fs from this connective configuration to generate the
additional configurations that bridge two metastable states as shown in Fig.
4(d). We also visualize the change of molecular structures from the connective
configuration to two metastable states in Fig. 4(e).

Solving BKE. Our next objective is to solve the BKE (2) in a 66-dimensional
space. The obtained numerical solution is used to generate the plot of the
approximate committor function in $(\phi,\psi)$. Such an approximation is then
compared with the approximation obtained from the diffusion map (DM) [8, 51]
method. Note that it is not possible to use traditional PDE solvers, such as
finite difference and finite element to solve (2) because they suffer from the
curse of dimensionality, i.e., their computational cost increase exponentially
with respect to dimension of the problem [41]. The DM method is another
sample-based method that allows for the solution of the BKE to be approximated
on an arbitrary set of configurations $\\{\bm{x}^{i}\\}$. However, it is
important to note that DM may not always produce an accurate approximation of
the derivatives at configurations near the boundary [52, 53]. This can result
in less accurate DM solutions to BKE near the boundaries of a fixed domain.

We use the dataset $\\{\bm{x}^{i}\\}$ identified by our RL method as shown in
Fig. 4(d) and apply DM and NN method to get the solutions of the BKE. The
solutions are presented in Fig. 5. We observe that the half-isocommittor
region of the solution, i.e., the set of configurations on which the committor
function value is close to 0.5, obtained from the DM method occupies a
relatively large area in the $(\phi,\psi)$ plane, whereas the half-
isocommittor region defined by the NN solution is confined to a small area
defined by $[-25^{\circ},25^{\circ}]\times[-80^{\circ},25^{\circ}]$. Our
solution is found to be more consistent with the results presented in Fig. 1
Panel 2 of Ref. [50].

Rate estimation. To estimate the reaction rate, we use formula (8) and
approximate the evaluation of the integral using the generated configurations
$\\{\bm{x}^{i}\\}$ obtained through the shooting method. This approach yields
the approximation formula:

$\displaystyle\kappa=\frac{\int_{\Omega}e^{-\beta
V(\bm{x})}\beta^{-1}\left\|\nabla
q(\bm{x})\right\|^{2}d\bm{x}}{\int_{\Omega}e^{-\beta
V(\bm{x})}d\bm{x}}\approx\frac{\sum_{i=1}^{N}e^{-\beta
V(\bm{x}_{i})}\beta^{-1}\left\|\nabla
q(\bm{x}_{i})\right\|^{2}}{\sum_{i=1}^{N}e^{-\beta V(\bm{x}_{i})}}.$ (22)

Finally, evaluating (22) using the approximated solution of KBE, we obtain a
rate of $1.58\times 10^{-5}\text{ ps}^{-1}$. This calculated estimate is
comparable to the approximated value $4.54\times 10^{-5}\text{ ps}^{-1}$
(reported in Figure 5 of Ref. [20]).

(a) Potential energy in $(\phi,\psi)$ and two metastable regions.

(b) Configurations generated from Langevin dynamics at temperature $300$K

(c) Learned action and identified CCs

(d) Reactive channel generated from ${\rm CC}_{1}$

(e) Change of configuration from from ${\rm CC}_{1}$ to metastable states.

Figure 4: Alanine dipeptide in vacuum under temperature $300$ K. (a) The
potential energy landscape visualized in $(\phi,\psi)$-space and two
metastable states (solid and dotted rectangle). (b) configurations generated
from overdamped Langevin dynamis at temperature 300 K. (c) The action function
learned by the proposed RL algorithm. The action function reveals two
connective configurations (CC’s). (d) Configurations generated by shooting
from the identified connective configurations. (e) The change of molecule
configuration from the identified connective configuration to metasable
states.

(a) Diffusion map.

(b) NN.

Figure 5: A comparison of solutions of the BKE (2) obtained from a diffusion
map and NN respectively at temperature 300K.

##  5 Conclusion

We presented a novel RL-based approach for identifying and characterizing an
ensemble of configurations where reactive trajectories are likely to be found.
The optimized action function returned from the RL algorithm reveals
connective configurations that have high reactive probabilities. One of the
key elements of the RL algorithm is the proper construction of a reward
function that serves as a surrogate for measuring the reactive probability of
a configuration, which is normally defined in terms of the value of it
committor function. Because the exact committor function is generally unknown
in advance, we employ a randomized shooting procedure to estimate its value at
an arbitrary configuration. Our results demonstrate the effectiveness of this
approach across multiple model problems of different sizes. Using the
identified connective configurations, we generate trajectories directed
towards metastable regions. The configurations along these trajectories are
utilized to define reactive channels on which a restricted BKE is solved by a
NN-based PDE solver. The solution yields an approximate committor function
evaluated within these channels. This committor function can then be used to
estimate reaction rates.

## Acknowledgement

This material is based upon work supported by the U.S. Department of Energy,
Office of Science, Office of Advanced Scientific Computing Research and Office
of Basic Energy Science, Scientific Discovery through Advanced Computing
(SciDAC) program under Contract No. DE-AC02-05CH11231. This work used the
computational resources of the National Energy Research Scientific Computing
(NERSC) center under NERSC Award ASCR-ERCAP m1027 for 2023, which is supported
by the Office of Science of the U.S. Department of Energy under Contract No.
DE-AC02-05CH11231.

## References

  * [1] Peters, B.  _Reaction rate theory and rare events_ (Elsevier, 2017).
  * [2] Chandler, D.  Statistical mechanics of isomerization dynamics in liquids and the transition state approximation.  _The Journal of Chemical Physics_ 68, 2959–2970 (1978).
  * [3] Bolhuis, P. G., Chandler, D., Dellago, C. & Geissler, P. L.  Transition path sampling: Throwing ropes over rough mountain passes, in the dark.  _Annual review of physical chemistry_ 53, 291–318 (2002).
  * [4] Geissler, P. L., Dellago, C. & Chandler, D.  Kinetic pathways of ion pair dissociation in water.  _The Journal of Physical Chemistry B_ 103, 3706–3710 (1999).
  * [5] Weinan, E., Ren, W. & Vanden-Eijnden, E.  String method for the study of rare events.  _Physical Review B_ 66, 052301 (2002).
  * [6] Vanden-Eijnden, E. _et al._ Towards a theory of transition paths.  _Journal of statistical physics_ 123, 503–523 (2006).
  * [7] Vanden-Eijnden, E. _et al._ Transition-path theory and path-finding algorithms for the study of rare events.  _Annual review of physical chemistry_ 61, 391–420 (2010).
  * [8] Coifman, R. R., Kevrekidis, I. G., Lafon, S., Maggioni, M. & Nadler, B.  Diffusion maps, reduction coordinates, and low dimensional representation of stochastic systems.  _Multiscale Modeling & Simulation_ 7, 842–864 (2008).
  * [9] Thiede, E. H., Giannakis, D., Dinner, A. R. & Weare, J.  Galerkin approximation of dynamical quantities using trajectory data.  _The Journal of chemical physics_ 150, 244111 (2019).
  * [10] Evans, L., Cameron, M. K. & Tiwary, P.  Computing committors via mahalanobis diffusion maps with enhanced sampling data.  _The Journal of Chemical Physics_ 157, 214107 (2022).
  * [11] Trstanova, Z., Leimkuhler, B. & Lelièvre, T.  Local and global perspectives on diffusion maps in the analysis of molecular systems.  _Proceedings of the Royal Society A_ 476, 20190036 (2020).
  * [12] Ma, A. & Dinner, A. R.  Automatic method for identifying reaction coordinates in complex systems.  _The Journal of Physical Chemistry B_ 109, 6769–6779 (2005).
  * [13] Rotskoff, G. M., Mitchell, A. R. & Vanden-Eijnden, E.  Active importance sampling for variational objectives dominated by rare events: Consequences for optimization and generalization.  In _Mathematical and Scientific Machine Learning_ , 757–780 (PMLR, 2022).
  * [14] Li, Q., Lin, B. & Ren, W.  Computing committor functions for the study of rare events using deep learning.  _The Journal of Chemical Physics_ 151, 054112 (2019).
  * [15] Khoo, Y., Lu, J. & Ying, L.  Solving for high-dimensional committor functions using artificial neural networks.  _Research in the Mathematical Sciences_ 6, 1–13 (2019).
  * [16] Hasyim, M. R., Batton, C. H. & Mandadapu, K. K.  Supervised learning and the finite-temperature string method for computing committor functions and reaction rates.  _arXiv preprint arXiv:2107.13522_ (2021).
  * [17] Strahan, J., Finkel, J., Dinner, A. R. & Weare, J.  Forecasting using neural networks and short-trajectory data.  _arXiv preprint arXiv:2208.01717_ (2022).
  * [18] Strahan, J., Guo, S. C., Lorpaiboon, C., Dinner, A. R. & Weare, J.  Inexact iterative numerical linear algebra for neural network-based spectral estimation and rare-event prediction.  _arXiv preprint arXiv:2303.12534_ (2023).
  * [19] Jung, H. _et al._ Machine-guided path sampling to discover mechanisms of molecular self-organization.  _Nature Computational Science_ 1–12 (2023).
  * [20] Singh, A. N. & Limmer, D. T.  Variational deep learning of equilibrium transition path ensembles.  _arXiv preprint arXiv:2302.14857_ (2023).
  * [21] Zhu, Y., Tang, Y.-H. & Kim, C.  Learning stochastic dynamics with statistics-informed neural network.  _Journal of Computational Physics_ 474, 111819 (2023).
  * [22] Dellago, C., Bolhuis, P. G. & Geissler, P. L.  Transition path sampling.  _Advances in chemical physics_ 123, 1–78 (2002).
  * [23] Bolhuis, P. G. & Swenson, D. W.  Transition path sampling as markov chain monte carlo of trajectories: Recent algorithms, software, applications, and future outlook.  _Advanced Theory and Simulations_ 4, 2000237 (2021).
  * [24] Dellago, C., Bolhuis, P. G. & Geissler, P. L.  Transition path sampling methods.  In _Computer Simulations in Condensed Matter Systems: From Materials to Chemical Biology Volume 1_ , 349–391 (Springer, 2006).
  * [25] Barducci, A., Bussi, G. & Parrinello, M.  Well-tempered metadynamics: a smoothly converging and tunable free-energy method.  _Physical review letters_ 100, 020603 (2008).
  * [26] Barducci, A., Bonomi, M. & Parrinello, M.  Metadynamics.  _Wiley Interdisciplinary Reviews: Computational Molecular Science_ 1, 826–843 (2011).
  * [27] Huber, G. A. & Kim, S.  Weighted-ensemble brownian dynamics simulations for protein association reactions.  _Biophysical journal_ 70, 97–110 (1996).
  * [28] Zwier, M. C. _et al._ Westpa: An interoperable, highly scalable software package for weighted ensemble simulation and analysis.  _Journal of chemical theory and computation_ 11, 800–809 (2015).
  * [29] Valsson, O. & Parrinello, M.  Variational approach to enhanced sampling and free energy calculations.  _Physical review letters_ 113, 090601 (2014).
  * [30] Zhang, J. _et al._ Deep reinforcement learning of transition states.  _Physical Chemistry Chemical Physics_ 23, 6888–6895 (2021).
  * [31] Das, A., Rose, D. C., Garrahan, J. P. & Limmer, D. T.  Reinforcement learning of rare diffusive dynamics.  _The Journal of Chemical Physics_ 155, 134105 (2021).
  * [32] Das, A., Kuznets-Speck, B. & Limmer, D. T.  Direct evaluation of rare events in active matter from variational path sampling.  _Physical Review Letters_ 128, 028005 (2022).
  * [33] Lelièvre, T., Robin, G., Sekkat, I., Stoltz, G. & Cardoso, G. V.  Generative methods for sampling transition paths in molecular dynamics.  _arXiv preprint arXiv:2205.02818_ (2022).
  * [34] Das, A. & Limmer, D. T.  Variational control forces for enhanced sampling of nonequilibrium molecular dynamics simulations.  _The Journal of chemical physics_ 151, 244123 (2019).
  * [35] Rose, D. C., Mair, J. F. & Garrahan, J. P.  A reinforcement learning approach to rare trajectory sampling.  _New Journal of Physics_ 23, 013013 (2021).
  * [36] Metzner, P., Schütte, C. & Vanden-Eijnden, E.  Illustration of transition path theory on a collection of simple examples.  _The Journal of chemical physics_ 125, 084110 (2006).
  * [37] Zhang, X., Landsness, E. C., Culver, J. P. & Anastasio, M. A.  Identifying functional brain networks from spatial-temporal wide-field calcium imaging data via a recurrent autoencoder.  In _Neural Imaging and Sensing 2022_ , PC1194612 (SPIE, 2022).
  * [38] Nan, Y. & Ji, H.  Deep learning for handling kernel/model uncertainty in image deconvolution.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2388–2397 (2020).
  * [39] Lin, B., Li, Q. & Ren, W.  A data driven method for computing quasipotentials.  In _Mathematical and Scientific Machine Learning_ , 652–670 (PMLR, 2022).
  * [40] Shen, Z., Yang, H. & Zhang, S.  Deep Network with Approximation Error Being Reciprocal of Width to Power of Square Root of Depth.  _Neural Computation_ (2021).
  * [41] Weinan, E., Han, J. & Jentzen, A.  Algorithms for solving high dimensional pdes: From nonlinear monte carlo to machine learning.  _Nonlinearity_ 35, 278 (2021).
  * [42] Paszke, A. _et al._ Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ 32 (2019).
  * [43] Kingma, D. P. & Ba, J.  Adam: A method for stochastic optimization.  In Bengio, Y. & LeCun, Y. (eds.) _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_ (2015).
  * [44] Sutton, R. S. & Barto, A. G.  _Reinforcement learning: An introduction_ (MIT press, 2018).
  * [45] Silver, D. _et al._ Deterministic policy gradient algorithms.  In _International conference on machine learning_ , 387–395 (PMLR, 2014).
  * [46] Fujimoto, S., Hoof, H. & Meger, D.  Addressing function approximation error in actor-critic methods.  In _International conference on machine learning_ , 1587–1596 (PMLR, 2018).
  * [47] Xu, Z.-Q. J., Zhang, Y., Luo, T., Xiao, Y. & Ma, Z.  Frequency Principle: Fourier Analysis Sheds Light on Deep Neural Networks.  _Communications in Computational Physics_ 28, 1746–1767 (2020).
  * [48] Liang, S., Huang, Z. & Zhang, H.  Stiffness-aware neural network for learning hamiltonian systems.  In _International Conference on Learning Representations_ (2022).
  * [49] Eastman, P. _et al._ Openmm 7: Rapid development of high performance algorithms for molecular dynamics.  _PLoS computational biology_ 13, e1005659 (2017).
  * [50] Kikutsuji, T. _et al._ Explaining reaction coordinates of alanine dipeptide isomerization obtained from deep neural networks using explainable artificial intelligence (xai).  _The Journal of Chemical Physics_ 156, 154108 (2022).
  * [51] Ko, T. _et al._ Using diffusion maps to analyze reaction dynamics for a hydrogen combustion benchmark dataset.  _arXiv preprint arXiv:2304.09296_ (2023).
  * [52] Liang, S., Jiang, S. W., Harlim, J. & Yang, H.  Solving pdes on unknown manifolds with machine learning.  _arXiv preprint arXiv:2106.06682_ (2021).
  * [53] Jiang, S. W. & Harlim, J.  Ghost point diffusion maps for solving elliptic pdes on manifolds with classical boundary conditions.  _Communications on Pure and Applied Mathematics_ 76, 337–405 (2023).

##  Appendix A Q-learning algorithm

Figure 6 gives a schematic description of how the action function and the
$Q$-function are optimized in the RL algorithm used to identify connective
configurations.

Figure 6:  A schematic illustration of how the Q-learning is performed for a
continuous action.

##  Appendix B Numerical results for high temperature regime

###  B.1 Triple-well potential

This section presents the the results obtained from running the RL algorithm
to identify connective configurations for a triple-well potential with an
inverse temperature of $\beta=1.67$. The reward function defined in Equation
(19) was evaluated by shooting 50 trajectories of that evolve up to time
$T=0.75$. The maximum number of time steps taken in each trajectory is set to
$L=20$. The RL procedure was run for a total of 1000 episodes.

Figure 7(a) shows the action function $A(x;\Phi)$ produced at the end of the
RL run. We mark two distinct connectivity configurations identified by
crosses. They are nearly identical to the connective configurations we found
when the trajectories were generated using $\beta=6.67$. In Figure 7(b), we
plot configurations generated by shooting trajectories from two connective
configurations. We observe that the configurations generated from running
trajectories using $\beta=1.67$ cover a wider region of the configuration
space that includes the local maximum near $(0.0,0.5)$, as well as some
configurations outside the pre-defined domain $\Omega=[-2,2]\times[-1.2,2]$.
Figures 7(c-e) show that the neural network solution to BKE is comparable to
solution obtained from the finite difference method (FDM).

(a) Learned action function and the connective configurations it reveals.

(b) Sampled configurations generated by shooting trajectories from the two
identified connective configurations.

(c) The FDM solution of (2) on $\Omega$

(d) The NN solution of (2).

(e) The difference between (c) and (d).

Figure 7: The results obtained from running the RL algorithm for Triple-well
potential at inverse temperature $\beta=1.67$. (a) The action function learned
by our proposed RL. The action function reveals two connective configurations
(crosses). (b) The configurations generated by shooting trajectories initiated
at the identified connective configurations. (c-e) The FDM (c) and NN (d)
solutions of the BKE (2) and their difference (e).

###  B.2 Rugged Muller potential

Figure 8 shows the results obtained from applying the previously presented RL
algorithm to the Rugged Muller potential at an inverse temperature of
$\beta=0.1$. The reward function defined in Equation (19) was computed based
on shooting $N=10$ trajectories that evolve up to time $T=0.05$. The maximum
number of time steps taken in each trajectory was set to $L=20$. The RL
algorithm was run for a total of 1000 episodes. We can observe that the
configurations generated by shooting trajectories from the identified
connective configurations with $\beta=0.1$ (see Fig. 8(b)) cover a wider area
compared to that obtained from performing RL and generating trajectories at
$\beta=0.25$. The committor function appears to be more stable near the region
around $(-0.8,0.6)$. The error is relatively small in the region covered by
the sampled configurations.

(a) Learned action function and the connective configurations it reveals.

(b) Sampled configurations generated by shooting from identified connective
configurations.

(c) The FDM solution of (2).

(d) The NN solution of (2).

(e) The difference between (d) and (d).

Figure 8: The results obtained from running the RL algorithm for the Rugged
Muller potential at the inverse temperature $\beta=0.1$. (a) The action
function learned by the RL algorithm. It reveals two connective configurations
(crosses). (b) The configurations generated by shooting trajectories from the
identified connective configurations. (c-e) The FDM (c) and NN (d) solutions
of the BKE (2) and their difference (e).

(a) Muller: $\beta=0.1$

(b) Muller: $\beta=0.25$

(c) Tri-well: $\beta=1.67$

(d) Tri-well: $\beta=6.67$

Figure 9: The moving average of the reward as a function of the episode number
during the RL procedure. Here the value of the reward is linearly scaled to
the interval $[0,1]$ by min-max normalization.

###  B.3 Monitoring RL progress

Figure 9 shows the reward as a function of the episode in the RL algorithm. A
maximum of 1000 episodes are allowed. However, the largest rewards are
achieved around episode 500. This suggests that the RL process is efficient
and reaches a high level of performance before reaching the limit on the
allowed number of episodes.

##  Appendix C Implementation details of NN solutions of the BKE

In this paper, we approximate the solution of the BKE by a fully connected
neural network (FNN), which can be viewed as a composition of $L$ simple
nonlinear functions, i.e.,
$\phi(\mathbf{x};\bm{\theta}):=\sigma_{2}\circ\mathbf{a}^{\top}\mathbf{h}_{L}\circ\mathbf{h}_{L-1}\circ\cdots\circ\mathbf{h}_{1}(\mathbf{x})$.
Here,
$\mathbf{h}_{\ell}(\mathbf{x})=\sigma\left(\mathbf{W}_{\ell}\mathbf{x}+\mathbf{b}_{\ell}\right)$
with $\mathbf{W}_{\ell}\in\mathbb{R}^{N_{\ell}\times N_{\ell-1}}$,
$\mathbf{b}_{\ell}\in\mathbb{R}^{N_{\ell}}$ for $\ell=1,\dots,L$,
$\mathbf{a}\in\mathbb{R}^{N_{L}}$, $\sigma$ is the tanh function, and
$\sigma_{2}$ is a sigmoid function such that the range of output is $[0,1]$.
We use an FNN with $L=2$ and a uniform width $m$, i.e., $N_{\ell}=m$ for all
$\ell\neq 0$.

The FNN is optimized by Adam [43]. In Adam, we use an initial learning rate of
0.001 for $T$ iterations. The learning rate is then adjusted in each iteration
by a factor of $0.5(\cos(\frac{\pi t}{T})+1)$, where $t$ is the current
iteration number. We set the batch size to be the total number of training
points. The hyperparameter setting for each numerical example is listed in
Table 3.

Table 3: Hyperparameter setting for solving BKE with NN. Note that we use two different boundary coefficients for boundary conditions of $A$ and $B$ in the Alanine dipeptide case.  Example | Triple-well |  | Rugged Muller |  | Alanine dipeptide
---|---|---|---|---|---
$\beta$ | 1.67 | 6.67 |  | 0.1 | 0.25 |  | 1/2.5 kJ-1mol
Width $m$ | 50 | 50 |  | 100 | 100 |  | 200
Boundary coefficient $\ell$ | 10 | 100 |  | 10,000 | 100,000 |  | 1 million & 2 million
Iteration | 30,000 | 30,000 |  | 100,000 | 30,000 |  | 100,000

##  Appendix D Additional notes on the RL method

###  D.1 RL initialization from stable states

In the numerical results presented in section B, the initialization of the
reinforcement learning (RL) algorithm is performed by randomly sampling from
$\Omega$ uniformly. In a revised approach, we revise the initialization scheme
by sampling the initial configurations from the meta-stable regions. To show
the impact of this change, we retrain the RL model and plot the final
configurations obtained from RL procedure using 1,000 different
initializations sampled from the meta-stable region in Figures 10.

With this initialization scheme, the RL method still successfully identifies
two regions with a high reactive probability for the triple-well potential.
Similarly, in the case of the Rugged Muller potential, the RL algorithm
converges closely to the configurations depicted in Figures 3(a) and 8(a).
Furthermore, in the Alanine dipeptide case, the RL approach reveals reactive
regions similar to those shown in Figure 4(c). These findings demonstrate the
robustness of the RL model with respect to different initialization
strategies.

(a) Triple-well $\beta=1.67$ (b) Triple-well $\beta=6.67$ (c) Rugged Muller
$\beta=0.1$ (d) Rugged Muller $\beta=0.25$ (e) Alanine dipeptide

Figure 10: The learned action produced by the RL algorithm where the initial
configurations are randomly sampled from the meta-stable region uniformly.
Each “+” represents the configuration generated in the final configuration of
RL episode.

###  D.2 Solving the BKE using NN with different training datasets

The utilization of NN-based optimization methods is often associated with an
implicit bias toward fitting smooth functions that exhibit fast decay in the
frequency domain. [47] Consequently, training NN models that can be used to
approximate the committor function can be challenging when we attempt to
capture drastic changes in the committor function. Such implicit bias can be
mitigated by using an appropriate training dataset. [48] By carefully
selecting the training dataset to provide the necessary samples that encompass
the desired variations in the committor function, we can overcome the
limitations posed by the implicit bias of NN-based optimization. Here we
compare NN solutions of the BKE trained on various datasets. These datasets
consist of configurations sampled uniformly on $\Omega$, configurations
obtained from overdamped Langevin dynamics ran at a higher temperature, and
configurations generated by shooting from CC’s, as depicted in Figures 11 and
13 respectively. We evaluate the performance of the NN on two examples (a
triple-well potential and a Rugged Muller potential) at a low temperature.
Figures 12 and 14 show the mean square error (MSE) of the training loss and
the NN solutions obtained from different training datasets.

From these results, it is evident that the NN solution trained on the dataset
generated by shooting trajectories from CC’s achieves a lower MSE and yields a
more accurate approximation.

(a) Uniformly sampled in $\Omega$.

(b) Overdamped Langevin ($\beta=1.67$)

(c) Shooting from CC

Figure 11: Different training datasets used for NN solutions of the BKE for
the triple well potential with $\beta=6.67$. (a) Uniform samples from
$\Omega$. (b) Samples along a overdamped Langevin dynamics trajectory with
$\beta=1.67$. (c) Configurations generated by shooting trajectories from the
identified connective configurations.

(a) MSE of the NN loss function,

(b) Uniformly sampled in $\Omega$. (c) Overdamped Langevin ($\beta=1.67$)

(d) Shooting from CC

Figure 12: (a) Mean square error of the training loss at configurations
sampled from the transition region of the triple well potential. (b-d) NN
solutions obtained from different training datasets.

(a) Uniformly sampled in $\Omega$.

(b) Overdamped Langevin ($\beta=0.05$)

(c) Shooting from CC

Figure 13: Different training datasets used for NN solutions of the BKE for
the Rugged Muller potential with $\beta=0.25$. (a) Uniformly sampled in
$\Omega$. (b) Sample along a overdamped Langevin dynamics trajectory with
$\beta=0.05$. (c) Configurations generated by shooting from the identified
connective configurations.

(a) MSE for BKE

(b) Uniform dataset

(c) Overdamped Langevin ($\beta=0.05$) (d) Shooting from CC

Figure 14: (a) Mean square error of the training loss at configurations
sampled from the transition region of the rugged Muller-Brown potential. (b-d)
NN solutions obtained from different training datasets.