# Response properties in phaseless auxiliary field quantum Monte Carlo

Ankit Mahajan<EMAIL_ADDRESS>Department of Chemistry, Columbia
University, New York, NY 10027, USA  Department of Chemistry, University of
Colorado, Boulder, CO 80302, USA  Jo S. Kurian  Department of Chemistry,
University of Colorado, Boulder, CO 80302, USA  Joonho Lee  Department of
Chemistry, Columbia University, New York, NY 10027, USA  Department of
Chemistry and Chemical Biology, Harvard University, Cambridge, MA 02138, USA
David R. Reichman  Department of Chemistry, Columbia University, New York, NY
10027, USA  Sandeep Sharma<EMAIL_ADDRESS>Department of Chemistry,
University of Colorado, Boulder, CO 80302, USA

###### Abstract

We present a method for calculating first-order response properties in
phaseless auxiliary field quantum Monte Carlo (AFQMC) through the application
of automatic differentiation (AD). Biases and statistical efficiency of the
resulting estimators are discussed. Our approach demonstrates that AD enables
the calculation of reduced density matrices (RDMs) with the same computational
cost scaling as energy calculations, accompanied by a cost prefactor of less
than four in our numerical calculations. We investigate the role of self-
consistency and trial orbital choice in property calculations.

##  1 Introduction

The calculation of energy gradients is of great importance in electronic
structure theory. Derivatives of the energy with respect to external
perturbations can be used to evaluate many properties of interest. Nuclear
gradients are used for geometry optimization and molecular dynamics
simulations. Parameter derivatives are useful in wave function and basis set
optimizations. Throughout the history of quantum chemistry, the development of
efficient methods for evaluating energy gradients has thus been a focal point
due to their central importance.Pulay (1969); Pople et al. (1979); Handy and
Schaefer III (1984); Salter et al. (1989); Helgaker et al. (2012) This
development often involved meticulous derivation of analytic expressions and
manual coding of implementations. The development of efficient algorithms has
led to cost-effective gradient evaluations, comparable in cost to the
evaluation of the energy. In recent years, automatic or algorithmic
differentiation (AD) has gained recognition as a powerful tool for calculating
derivatives of functions defined by computer programs, notably in machine
learning. By applying the chain rule to elementary steps in a computation, AD
can calculate numerically exact derivatives, considerably reducing the
implementation effort.Griewank and Walther (2008) The reverse-mode AD is
particularly remarkable, enabling the evaluation of first-order derivatives
with respect to an arbitrary number of parameters at the same cost scaling as
the underlying function. As a result, it becomes an appealing choice for
calculating all nuclear forces or reduced density matrices in quantum
chemistry. Various studies have explored the utility of AD in conventional
electronic structure methods.Ekström et al. (2010); Bast et al. (2011); Abbott
et al. (2021); Zhang and Chan (2022); Kasim et al. (2022) Reference 12
showcased the effectiveness of AD in differentiating tensor network-based
functions.

Quantum Monte Carlo (QMC) approaches comprise a scalable and accurate suite of
alternatives to traditional quantum chemistry methods.Kalos et al. (1974);
Ceperley et al. (1977); Ceperley and Alder (1980); Nightingale and Umrigar
(1998); Foulkes et al. (2001); Becca and Sorella (2017) Among these, the
auxiliary field quantum Monte Carlo (AFQMC) approach stands out as a powerful
tool for investigating correlated electronic systems.Zhang et al. (1995,
1997); Zhang and Krakauer (2003) It has been used to study a wide variety of
systems, including solids and molecules with a range of correlation
strengths.Al-Saidi et al. (2006); Suewattana et al. (2007); Hao et al. (2018);
Lee et al. (2020); Malone et al. (2022); Shee et al. (2019) However, the
evaluation of energy derivatives has posed challenges for projection QMC
methods. Mixed estimators are usually used in energy calculations with these
methods, which are not accurate for general properties or nuclear gradients.
An incorrect application of the Hellman-Feynman theorem leads to this mixed
estimator that is easy to evaluate but leads to an uncontrolled bias for
operators that do not commute with the Hamiltonian. Although various
approximations can address this shortcoming,Badinski et al. (2010); Casalegno
et al. (2003); Filippi and Umrigar (2000); Moroni et al. (2014); Assaraf and
Caffarel (2003); Motta and Zhang (2017); Thomas et al. (2015) many challenges
still remain due to uncontrolled biases or large variances. It is worth
mentioning that within the context of variational Monte Carlo (VMC), AD was
first employed by Sorella and Capriotti to calculate nuclear gradients.Sorella
and Capriotti (2010) Additionally, there have been studies of AD in diffusion
Monte Carlo (DMC)Poole (2014) and lattice VMCZhang et al. (2019), following an
approach similar to that presented here.

The problem of calculating derivatives of complicated stochastic functions
plays a crucial role in various domains ranging from sensitivity analysis in
financeGlasserman (2004) and queueing theory in managementGlynn (1990) to
optimizations in reinforcement learning.Bengio et al. (2013) An effective
approach for formulating low-variance estimators involves utilizing correlated
sampling, by employing common random numbers for example, to achieve small
cost differences with high accuracy. When Monte Carlo simulations involve
discrete random variables, like branching in QMC calculations, biases can
arise in derivative estimators, and addressing these errors is an active area
of research.Mohamed et al. (2020) In this paper, we leverage some of this
technology to compute first-order response properties in AFQMC. Specifically,
we demonstrate that reverse-mode AD enables the calculation of Reduced Density
Matrices (RDMs) with the same cost scaling as energy calculations.

This paper is organized as follows: we first review relevant parts of the
AFQMC algorithm (Section 2.1) and existing approaches for calculating
properties (Section 2.2). Then we discuss use of AD in AFQMC, including
various contributions to the derivatives as well as biases in our method
(Section 2.3). We then present illustrative numerical results to demonstrate
the accuracy and statistical efficiency of our method (Section 3.1). Finally,
we compare AD-AFQMC dipole moments of some molecules with other quantum
chemistry methods and experiments (Section 3).

##  2 Theory

###  2.1 Review of AFQMC

In this section, we present an outline of the AFQMC algorithm, focusing on the
aspects relevant to gradient calculations. We refer the reader to recent
review articles Motta and Zhang (2018); Shi and Zhang (2021); Lee et al.
(2022) for a more comprehensive description. In AFQMC, the ground state is
represented as a weighted sum of non-orthogonal Slater determinants

$|\Psi_{0}\rangle=\sum_{i}w_{i}\frac{|\phi_{i}\rangle}{\langle\psi_{T}|\phi_{i}\rangle},$
(1)

where $w_{i}$ are weights, $|\phi_{i}\rangle$ are walker Slater determinants
with complex orbitals, and $|\psi_{T}\rangle$ is the trial state. This state
is obtained by applying an exponential form of projector onto the trial state
as

$e^{-\tau\hat{H}}|\psi_{T}\rangle\xrightarrow{\tau\rightarrow\infty}|\Psi_{0}\rangle,$
(2)

where $\tau$ is the imaginary time, and we assume
$\langle\psi_{T}|\psi_{0}\rangle\neq 0$. We use the quantum chemistry
Hamiltonian given by

$\hat{H}=\sum_{ij}h_{ij}\hat{a}_{i}^{\dagger}\hat{a}_{j}+\frac{1}{2}\sum_{\gamma}\left(\sum_{ij}L^{\gamma}_{ij}\hat{a}_{i}^{\dagger}\hat{a}_{j}\right)^{2},$
(3)

where $h_{ij}$ are one-electron integrals and $L^{\gamma}_{ij}$ are Cholesky
decomposed two-electron integrals in an orthonormal orbital basis. To sample
the action of the propagator, the Hubbard-Stratonovic transform is used to
write the short-time propagator as

$e^{-\Delta\tau\hat{H}}=\int d\mathbf{x}p(\mathbf{x})\hat{B}(\mathbf{x}),$ (4)

where $\mathbf{x}$ is the vector of auxiliary fields, $p(\mathbf{x})$ is the
standard normal distribution of the auxiliary fields, and
$\hat{B}(\mathbf{x})$ is a complex propagator given by the exponential of a
one-body operator. Due to Thouless’ theorem, $\hat{B}(\mathbf{x})$ acts on a
Slater determinant $|\phi\rangle$ as

$\hat{B}(\mathbf{x})|\phi\rangle=|\phi(\mathbf{x})\rangle,$ (5)

where $|\phi(\mathbf{x})\rangle$ is another Slater determinant obtained by a
linear transformation of the orbitals in $|\phi\rangle$. In the hybrid
approximation, importance sampling is achieved by shifting the auxiliary
fields with a force bias given as

$\bar{x}_{\gamma}=-\sqrt{\Delta\tau}\frac{\langle\psi_{T}|\sum_{ij}L^{\gamma}_{ij}\hat{a}_{i}^{\dagger}\hat{a}_{j}|\phi\rangle}{\langle\psi_{T}|\phi\rangle},$
(6)

and with the weights propagated as

$w_{i}(\mathbf{x})=\bigg{|}\frac{\langle\psi_{T}|\phi(\mathbf{x})\rangle}{\langle\psi_{T}|\phi\rangle}e^{\mathbf{x}.\bar{\mathbf{x}}-\frac{\bar{\mathbf{x}}^{2}}{2}}\bigg{|}w_{i}.$
(7)

Force bias is not enough, by itself, to control the large fluctuations
stemming from the sign or phase problem. The phaseless constraintZhang and
Krakauer (2003) can be used to overcome the phase problem at the expense of a
systematic bias in the sampled wave function by performing a cosine projection
as

$w_{i}(\mathbf{x})=\text{max}(0,\cos(\Delta\theta)),w_{i}(\mathbf{x})$ (8)

where

$\Delta\theta=\text{Arg}\left(\frac{\langle\psi_{T}|\phi(\mathbf{x})\rangle}{\langle\psi_{T}|\phi\rangle}\right).$
(9)

The size of this phaseless bias is dictated by the accuracy of the trial
state.

For statistical efficiency, we periodically perform reconfigurations of the
walker population. We use the stochastic reconfiguration (SR) methodBuonaura
and Sorella (1998) where walkers with small weights are removed, and those
with large weights are duplicated. Figure 1 illustrates this procedure
schematically. It has the desirable property that when the walker weights are
uniform, the walker population remains unchanged. It is important to note that
SR is a discontinuous function, as small changes in weights lead to different
reconfiguration instances. These discontinuities are a source of bias in our
AD gradient calculations, as described in section 2.3.

Figure 1: A schematic showing the stochastic reconfiguration (SR) procedure
with four walkers. $\xi$ is drawn from a uniform distribution on
[0,$\frac{1}{4}$] and $\bar{w}_{i}=\frac{w_{i}}{\sum_{i}w_{i}}$ denote
normalized weights. In this instance, SR leads to the reconfiguration
$\left\\{\phi_{1},\phi_{2},\phi_{3},\phi_{4}\right\\}\rightarrow\left\\{\phi_{1},\phi_{3},\phi_{3},\phi_{4}\right\\}$.

###  2.2 Calculating properties in AFQMC

Given an observable $\hat{O}$, the simplest way to estimate its expectation
value in AFQMC is to use the mixed estimator given by

$\langle\hat{O}\rangle_{\text{mixed}}=\frac{\sum_{\phi}w_{\phi}\frac{\langle\psi_{T}|\hat{O}|\phi\rangle}{\langle\psi_{T}|\phi\rangle}}{\sum_{\phi}w_{\phi}}.$
(10)

If $\hat{O}$ does not commute with the Hamiltonian, this estimator has a bias
due to its non-variational nature. One way to mitigate this error is to use
the extrapolated estimatorWhitlock et al. (1979); Ceperley and Kalos (1986);
Rothstein (2013)

$\langle\hat{O}\rangle_{\text{extrapolated}}=2\langle\hat{O}\rangle_{\text{mixed}}-\frac{\langle\psi_{T}|\hat{O}|\psi_{T}\rangle}{\langle\psi_{T}|\psi_{T}\rangle}.$
(11)

The simplicity of evaluating this estimator comes with a crucial caveat; its
accuracy heavily relies on the trial state employed. For instance, achieving
convergence of dipole moments in reference 49 required the use of a large
multi-Slater trial state.

A more common approach to calculating properties in AFQMC is using
backpropagation,Zhang et al. (1995) which overcomes the non-variationality
issue by propagating the trial state backward as

$\langle\hat{O}\rangle_{\text{bp}}=\frac{\langle\psi_{T}|e^{-\tau_{\text{bp}}\hat{H}}\hat{O}e^{-\tau\hat{H}}|\psi_{T}\rangle}{\langle\psi_{T}|e^{-(\tau_{\text{bp}}+\tau)\hat{H}}|\psi_{T}\rangle}.$
(12)

A direct simulation of two independent random walks results in a highly noisy
and computationally intensive estimator. To control this noise, importance
sampling is employed to correlate the two walks. However, in certain systems,
this procedure has been observed to introduce a significant phaseless bias.
This bias can be reduced by undoing the phaseless approximation in the
backpropagation, albeit at the expense of noisier estimates.Motta and Zhang
(2017) Backpropagation has also been used to evaluate nuclear gradients based
on the Hellman-Feynman theorem.Motta and Zhang (2018)

With the goal of bringing accuracy and computational cost of property
calculations on par with that of calculating the energy, we pursue a response
approach in this paper. Consider a perturbed Hamiltonian given by

$\hat{H}(\lambda)=\hat{H}+\lambda\hat{O},$ (13)

where $\lambda$ is a real parameter. The response estimator can be written as

$\langle\hat{O}\rangle_{\text{response}}=\frac{dE(\lambda)}{d\lambda}\bigg{|}_{\lambda=0},$
(14)

where $E(\lambda)$ is the AFQMC energy calculated with the perturbed
Hamiltonian $\hat{H}(\lambda)$. For a variationally optimized wave function,
this estimator is equivalent to the expectation value by virtue of the
Hellman-Feynman theorem. For non-variational methods like coupled cluster
theory, response properties are found to be very accurate and are routinely
used. One way to evaluate energy derivatives is using the method of finite
differences. However, as AFQMC is a stochastic method, it is necessary to
correlate the finite difference energy calculations to minimize the noise in
the energy difference. Such correlated sampling methods have been used in
AFQMCShee et al. (2017) as well as other QMC methodsAssaraf et al. (2011);
Filippi and Umrigar (2000). Using a common stream of random numbers is an
intuitively appealing strategy to maintain two Monte Carlo runs close to each
other, but technical challenges arise in ensuring that the runs remain
coherent. A promising approach to address this challenge was recently reported
in reference 53. A drawback of the correlated sampling finite difference
method is the requirement for separate energy calculations for each
perturbation, making the calculation of reduced density matrices and all
nuclear forces computationally expensive, with the cost scaling linearly with
the number of perturbations. In the next section, we will explore how reverse-
mode automatic differentiation utilizes the common random number strategy,
while also allowing calculations of responses to multiple perturbations at a
cost comparable to that of a single energy calculation.

###  2.3 Derivatives of AFQMC energy

We can write the AFQMC energy expression making the dependence on the
perturbation explicit as

$\begin{split}E(\hat{H}_{\lambda},\psi_{T}(\lambda),\left\\{w_{i}(\lambda)\right\\},&\left\\{\phi_{i}(\lambda)\right\\})\\\
&=\frac{\sum_{i}w_{i}(\lambda)\frac{\langle\psi_{T}(\lambda)|\hat{H}(\lambda)|\phi_{i}(\lambda)\rangle}{\langle\psi_{T}(\lambda)|\phi_{i}(\lambda)\rangle}}{\sum_{i}w_{i}(\lambda)}.\end{split}$
(15)

Differentiating this expression with respect to $\lambda$ gives

$\begin{split}\frac{dE}{d\lambda}=&\frac{\sum_{i}w_{i}(\lambda)\frac{\langle\psi_{T}(\lambda)|\hat{O}|\phi_{i}(\lambda)\rangle}{\langle\psi_{T}(\lambda)|\phi_{i}(\lambda)\rangle}}{\sum_{i}w_{i}(\lambda)}+\frac{\partial
E}{\partial\psi_{T}(\lambda)}\frac{d\psi_{T}(\lambda)}{d\lambda}\\\
&+\sum_{i}\left(\frac{\partial E}{\partial
w_{i}(\lambda)}\frac{dw_{i}(\lambda)}{d\lambda}+\frac{\partial
E}{\partial\phi_{i}(\lambda)}\frac{d\phi_{i}(\lambda)}{d\lambda}\right).\end{split}$
(16)

The first term is the mixed estimator 10. The second term contains the
explicit effect of the relaxation of the trial state in response to the
perturbation on the energy. It also appears in the response calculation of
many conventional quantum chemistry methods where it is evaluated (implicitly)
using the coupled perturbed Hartree Fock (CPHF) equations.Gerratt and Mills
(1968) The final term contains contributions from the derivatives of weights
and walkers, which, in turn, are influenced by the response of the trial
state. The dependence of weights and walkers on the perturbation is
accumulated through many steps of AFQMC propagation and local energy
measurements. One can evaluate their derivatives by repeated application of
the chain rule. AD offers an efficient and convenient way to accomplish this
task.

All these contributions can be evaluated using AD on an energy function that
incorporates the effects of the perturbation on the final energy. If
stochastic reconfiguration is not performed, the AFQMC energy can be written
as the integral

$E(\lambda)=\int d\mathbf{x}\mathcal{N}(\mathbf{x})f(\mathbf{x},\lambda),$
(17)

where $\mathbf{x}$ is the collection of normally distributed auxiliary fields
for all time steps and walkers, and $f$ denotes the AFQMC procedure, including
propagation and measurement. For the sake of notational convenience, we have
suppressed the dependence of $f$ on other quantities that are not pertinent to
the current discussion. The derivative of this integral can be obtained as

$\frac{dE}{d\lambda}=\int d\mathbf{x}\ \mathcal{N}(\mathbf{x})\frac{\partial
f}{\partial\lambda}(\mathbf{x}.\lambda),$ (18)

Note that the phaseless constraint (8) is not differentiable at the origin,
but it is continuous and differentiable almost everywhere. The switching of
derivative and integral is allowed in this case because all the operations
performed in the evaluation of $f$ are continuous and differentiable almost
everywhere (within the region of interest).Glasserman (2004) As a result, (18)
provides a means to sample the derivative through Monte Carlo sampling, akin
to the energy, and is known as a pathwise estimator. This estimator exhibits
desirable properties of low variance, thanks to correlated sampling, and
avoids numerical issues stemming from finite difference step sizes. This
pathwise estimator is precisely the one that is evaluated by performing
automatic differentiation on the energy function. Upon introducing SR, the
energy can be written as

$E(\lambda)=\int d\mathbf{x}d\boldsymbol{\xi}\
\mathcal{N}(\mathbf{x})\mathcal{U}(\boldsymbol{\xi})\tilde{f}(\mathbf{x},\mathbf{\xi},\lambda),$
(19)

where $\mathbf{\xi}$ are uniformly distributed variables. Due to the discrete
nature of SR, the function $\tilde{f}$ is not continuous and the pathwise
estimator is biased. Several methods have been developed to address biases
arising from discontinuities, score function or likelihood ratio estimators
being a popular choice.Glynn (1990); Glasserman (2004); Schulman et al. (2015)
However, in preliminary numerical experiments, we observed this method to be
quite noisy. Hence, we defer a detailed study of these bias mitigation
techniques to future work. In this paper, we use the biased pathwise estimator
to sample the derivative of the energy function. We present a numerical
analysis of this bias in section 3.

###  2.4 Implementation details

In AFQMC energy calculations, a sequence of phaseless imaginary time
propagation steps is executed before reconfiguration and energy measurement.
At the end of this block of steps, the walker weights and local energies are
averaged to yield estimates of the energy along with its associated stochastic
error. To obtain energy derivative samples, AD is performed on the averaged
energy from a sufficiently large number of such blocks. The number of blocks
to be included in the AD calculation is dictated by the time required for the
convergence of walker and weight derivatives (16). We present a numerical
investigation of this convergence in section 3.1.

AD can be used in two modes: forward and reverse. In forward-mode AD, the
derivative of the energy function is evaluated by propagating the perturbation
through the energy function. This approach calculates the response to one
perturbation at a time, resulting in a computational cost scaling comparable
to correlated finite difference methods. However, it offers the advantage of
having no finite difference errors and requires low memory, making it suitable
for evaluating a few properties in large systems. For example, one can
evaluate the dipole moment of a system by using the perturbation given by

$\hat{H}(\boldsymbol{\varepsilon})=\hat{H}+\boldsymbol{\varepsilon}.\hat{\mathbf{d}},$
(20)

where $\hat{\mathbf{d}}$ is the dipole operator, and calculating the
derivative with respect to the electric field $\boldsymbol{\varepsilon}$,
requiring at most three forward-mode AD calculations. In reverse-mode AD,
gradients are computed by propagating changes in energy backward through to
the perturbations. Remarkably, it shares the same computational cost scaling
as the energy function itself when calculating gradients with respect to an
arbitrary number of perturbations. Depending on the specific implementation,
the cost is usually between 3 to 5 times that of the energy evaluation. This
property makes reverse-mode AD well-suited for efficiently calculating many
properties simultaneously. For example, one can evaluate one and two-particle
reduced density matrices (RDM) by using the following perturbed Hamiltonian

$\hat{H}(\lambda)=\hat{H}+\sum_{ij}\lambda_{ij}\hat{a}_{i}^{\dagger}\hat{a}_{j}+\sum_{iajb}\gamma_{ipjq}\hat{a}_{i}^{\dagger}\hat{a}_{j}^{\dagger}\hat{a}_{q}\hat{a}_{p},$
(21)

and evaluating derivatives with respect to $\lambda_{ij}$ and $\gamma_{ipjq}$,
respectively.

Reverse-mode AD achieves its impressive computational cost scaling by storing
intermediates during the evaluation of the energy function resulting in a
linear increase in memory usage with the number of AFQMC time steps to be
differentiated. Since the derivatives of weights and walkers need to be
converged with respect to the propagation time, this leads to a substantial
memory cost for reverse-mode AD with increasing system size. We use
checkpointing,Griewank and Walther (2008) which recalculates some
intermediates on the backward pass instead of storing them, to reduce this
cost. This reduction in memory usage comes at the expense of an increased
computational cost prefactor.

For capturing the orbital relaxation of a HF trial state, the AFQMC energy
function can be written to include a call to a HF solver. We use a fixed and
conservatively large number of SCF iterations to ensure that the derivatives
are fully converged. Additionally, issues related to degeneracies in reverse-
mode AD of eigendecompositions can be dealt with by ignoring rotations in
degenerate spaces.Zhang and Chan (2022); Liao et al. (2019)

##  3 Results

In this section, we present the results of our response ph-AFQMC calculations,
referred to as AD-AFQMC, and analyze both systematic and stochastic errors.
Our analysis focuses on the accuracy of 1-RDM and dipole moment calculations
using AFQMC with a HF trial state. We used PySCFSun et al. (2018) to obtain
molecular integrals and to perform quantum chemistry wave function
calculations. The near-exact Density Matrix Renormalization Group (DMRG)
calculations on hydrogen chains were done using the block2 code.Zhai and Chan
(2021) We used QMCPACKKim et al. (2018) to do backpropagation AFQMC
calculation. The code used to perform AD-AFQMC calculations is available in a
public repositorydqm , with the JaxBradbury et al. (2018) library used for AD.
Reverse-mode AD was employed for 1-RDM calculations, while forward-mode AD was
used for dipole moments (except for self-consistent calculations). Input and
output files for all calculations can be accessed from a public repository.afq
We used a time step of 0.01 a.u. in all ph-AFQMC calculations. Cholesky
decompositions were calculated up to a threshold error of $10^{-5}$. Errors
due to these approximations can be controlled systematically and we estimate
them to be smaller than the statistical errors in the results presented here.

###  3.1 Illustrative results

Figure 2: Dipole moment of NH3 in the cc-pVDZ basis (top) and one electron
energy of a H20 chain in the STO-6G basis (bottom) as a function of
propagation time.

First, we analyze various contributions to observables and their convergence
with imaginary time. In the top panel of figure 2, we show the convergence of
the AD-AFQMC dipole moment of the ammonia molecule in the cc-pVDZ basis at
equilibrium geometry. The x-axis represents imaginary time, over which the AD
is performed, starting with equilibrated walkers. After around 8 a.u., the
dipole moment converges to a value outside the stochastic error bar of the
finite difference AFQMC estimator calculated using common random numbers. This
discrepancy arises from the bias due to SR in sampling AD derivatives, leading
to a bias of 0.003(1) a.u in this case. Since this is a relatively small
system, we can perform AD-AFQMC calculations without SR and this value matches
with the finite difference calculation within stochastic error. However,
without SR, more sampling effort is required to achieve the same level of
stochastic error. One can perform SR less frequently to control this bias, but
we do not explore this approach further since the bias is relatively small. We
also show the converged value of the AD-AFQMC dipole moment (with SR) without
including trial orbital relaxation, illustrating that including orbital
relaxation brings the dipole moment closer to the exact value.

In the bottom panel of figure 2, we conduct a similar analysis for the one-
electron energy of a twenty-atom linear hydrogen chain with an interatomic
separation of 2.4 a.u. in the minimal STO-6G basis. We use an unrestricted HF
(UHF) trial state for the AFQMC calculations. In this case, the SR bias of AD-
AFQMC is 0.001(1) a.u. The orbital relaxation contribution is substantial and
increases the accuracy of the estimator significantly.

###  3.2 Hydrogen chains

Figure 3: Relative error in the one-electron energy (left) and 1-RDM (right)
for a hydrogen chain (STO-6G basis) with an inter-hydrogen separation of
$d=1.6$ Bohr (top) and $d=2.4$ Bohr (bottom).

In figure 3, we present a comparison of the accuracy of AD-AFQMC one-electron
energies and 1-RDMs with coupled cluster methods for hydrogen chains. Near-
exact reference benchmark results for this one-dimensional system were
obtained using DMRG. Unrestricted coupled cluster calculations were performed
on top of UHF references. We did not include orbital relaxation in CC
calculations, as we found it not to significantly impact the results. All
AFQMC calculations were done with UHF trial states and with SR. 1-RDM errors
were calculated using the Frobenius norm. We consider two bond lengths ($d$)
of 1.6 a.u. and 2.4 a.u.

For $d=1.6$ a.u., out of the CC methods CCSD(T) performs very well, as one
would expect for a system close to equilibrium. We note that for the same bond
length, the chains become more correlated with increasing length, explaining
the increasing error in both the one-electron energy and 1-RDM for the CC
methods. AD-AFQMC one-electron energy systematic errors are smaller than
CCSD(T) for all chain lengths, except for H10. This discrepancy is due to the
proximity of $d=1.6$ a.u. to the Coulson-Fischer point of this chain, beyond
which the UHF solution becomes more stable. The AD-AFQMC 1-RDM error is also
smaller than CCSD for all chains except H10 and smaller than CCSD(T) beyond
H50.

For $d=2.4$ a.u., the CC methods are less accurate due to the increased
correlation in the system. In this case, restricted methods have convergence
issues, so we only present the results for unrestricted observables. For both
one-electron energies and 1-RDMs, AD-AFQMC errors are significantly smaller
compared to UCCSD and UCCSD(T), demonstrating the efficacy of AD-AFQMC in
providing accurate estimates of properties across a range of correlation
strengths.

Figure 4: Stochastic error in the occupation number in an atomic orbital in
the middle of a hydrogen chain (STO-6G basis).

Analyzing the system-size scaling of stochastic errors in AD-AFQMC estimates
is crucial to assess their applicability in large systems. For a local
observable, it is desirable that the stochastic error not grow with system
size, since the observable itself does not scale. In figure 4, we show the
stochastic error in the occupation of a hydrogen 1$s$ STO-6G atomic orbital in
the middle of a hydrogen chain for $d=1.6$ a.u. and $d=2.4$ a.u. for the same
amount of sampling effort. Remarkably, it is observed that the stochastic
error does not increase with system size asymptotically. While the stochastic
error in energy increases as $\sqrt{N_{\text{Hydrogens}}}$, the noise in the
energy derivative does not increase. This behavior has previously been
reported for correlated sampling estimators in 52.

Figure 5: Computational time scaling of AD-AFQMC calculations for hydrogen
chains of different lengths with the same sampling effort. Reverse mode AD
times correspond to the 1-RDM calculations shown in figure 3, while forward
mode AD times correspond to the occupation number calculations in figure 4.

In figure 5, we show the cost (wall time, including compilation time) of
forward and reverse mode AD-AFQMC 1-RDM calculations relative to that of
energy. As expected, the reverse mode AD costs are 3 to 4 times as expensive
as the energy calculation for all chain lengths, while the forward mode AD
costs are 2 to 3 times as expensive. We also show the size scaling of energy
calculation cost, which is seen to be between quadratic and cubic. We note
that computational cost scaling of the backpropagation algorithm for the
calculation of properties and forces is also similarly favorable.Motta and
Zhang (2017, 2018)

###  3.3 Dipole moments

Table 1: Dipole moments (in a.u.) of molecules at equilibrium geometries using the aug-cc-pVTZ basis set. BP-AFQMC refers to estimates obtained using backpropagation in AFQMC with partial path restoration. Molecule |  | MP2 |  | CCSD |  | BP-AFQMC |  | AD-AFQMC |  | CCSD(T) |  | Experiment
---|---|---|---|---|---|---|---|---|---|---|---|---
H2O |  | 0.7298 |  | 0.7335 |  | 0.707(4) |  | 0.720(2) |  | 0.7247 |  | 0.730(2)Lide (2004)
NH3 |  | 0.5996 |  | 0.6015 |  | 0.578(6) |  | 0.592(2) |  | 0.5938 |  | 0.581(1)Shimizu (1970)
HCl |  | 0.4389 |  | 0.4318 |  | 0.400(5) |  | 0.429(1) |  | 0.4273 |  | 0.430Lovas et al. (2005)
HBr |  | 0.3379 |  | 0.3289 |  | 0.33(1) |  | 0.329(2) |  | 0.3245 |  | 0.325Lovas et al. (2005)
CO |  | 0.1050 |  | 0.0199 |  | 0.12(1) |  | 0.019(4) |  | 0.0429 |  | 0.048(1)Meerts et al. (1977)
CH2O |  | 0.9375 |  | 0.9666 |  | 0.82(1) |  | 0.965(5) |  | 0.9389 |  | 0.918(1)Theulé et al. (2003)
C4H5N |  | 0.7354 |  | 0.7241 |  | - |  | 0.730(9) |  | 0.7255 |  | 0.72(2)Nelson et al. (1967)

In table 1, we show a comparison of dipole moments calculated using AD-AFQMC
with those calculated using backpropagation (BP-AFQMC) and with conventional
quantum chemistry methods. We use the aug-cc-pVTZ basis which is large enough
to allow meaningful comparison with experimental values. The equilibrium
geometries used here are provided in reference 61. For pyrrole, we used the
geometry reported in reference 68. We include orbital relaxation contributions
and use the frozen-core approximation in all methods. Orbital relaxation
contributions are not relevant for BP-AFQMC.

Since these molecules are largely single reference, CCSD(T) values are
expected to be very accurate. We see that AD-AFQMC dipole moments are in
better agreement with CCSD(T) values than CCSD in most cases. They are also
close to experimental dipole moments with some residual basis set error. This
is an encouraging sign considering that AFQMC has a less steep cost scaling
compared to the CC methods. Results obtained with BP-AFQMC using partial path
restoration are significantly less accurate than AD-AFQMC for CO and CH2O.
While the sampling effort in BP and AD-AFQMC is not directly comparable, we
used a similar allocation of computational resources for both sets of
calculations. Note that one can use better trial states to obtain improved BP
estimates as demonstrated in reference 33. Details about the convergence of
BP-AFQMC are provided in the supporting information.

We see greater errors in both AD-AFQMC and CCSD dipole moments of CO and CH2O
molecules. To mitigate this error, we use the self-consistent AFQMC
approach.Qin et al. (2016); Qin (2023) We calculate the 1-RDM using reverse-
mode AD-AFQMC and use the corresponding natural orbitals in the trial state of
the subsequent AFQMC calculation. This procedure is repeated until the desired
convergence threshold. Note that we include orbital relaxation contributions
in the zeroth iteration to obtain an accurate 1-RDM initially, but not in the
subsequent iterations. The results of this procedure for CO and CH2O are shown
in figure 6. In both cases, the self-consistent iterations can be seen to
improve the dipole accuracy with convergence achieved in less than five steps.
Availability of 1-RDM from AD-AFQMC can thus enable systematic improvement in
AFQMC properties. We also show the results of a single AD-AFQMC calculation
using DFT-B3LYP as the trial state without orbital relaxation. These dipole
moments show better agreement with the CCSD(T) value suggesting that these
orbitals are optimal for AFQMC calculations. Thus it may be possible to devise
orbital optimization schemes other than the one used in self-consistent AFQMC
to obtain better trial states for property calculations.

Figure 6: Convergence of the dipole moment of CO (top panel) and CH2O (bottom
panel) with the number of self-consistent AFQMC iterations using the aug-
ccpVTZ basis. Note that self-consistency performs worse than using B3LYP
orbitals.

##  4 Conclusion

In this work, we investigated the utility of automatic differentiation (AD)
for calculating response properties in auxiliary field quantum Monte Carlo
(AFQMC) with the aim of achieving accuracy and computational cost comparable
with energy calculations. We studied various contributions to energy
derivatives and their convergence behavior. Our analysis demonstrated the
accuracy of AD-AFQMC for computing one-particle reduced density matrices
(1-RDM) in hydrogen chains, achieved at a computational cost no more than four
times that of the energy calculation. We found that the stochastic error of
the AD-AFQMC estimator for a local orbital density does not increase with
system size. Finally, dipole moments of several molecules calculated using AD-
AFQMC were found to be in good agreement with CCSD(T) and experiments. We were
able to further improve the accuracy of these dipole moments by employing
self-consistent AFQMC calculations, and alternatively by using trial states
from DFT.

A way to improve AD-AFQMC involves addressing the bias due to stochastic
reconfiguration. While for the systems studied in this paper, the SR bias was
found to be small, it is possible that in systems where there are large
fluctuations in the weights, this bias can be significant. We plan to explore
bias mitigation techniques in future work. Another avenue for improvement lies
in using more accurate trial states. In this work, we used a single
determinant trial state, but it is possible to use multi-determinant trial
states. Furthermore, to enable calculations on much larger systems,
advancements in AD implementations are necessary. Developing more
sophisticated checkpointing schemes may help manage the memory usage of
reverse-mode AD, allowing for more extensive calculations without excessive
memory constraints. Overall, these potential enhancements will contribute to
the broader applicability of AD-AFQMC.

Despite these challenges, we believe that AD-AFQMC presents a viable and
promising approach for calculating properties in AFQMC. A natural extension of
this work is the calculation of nuclear gradients using AD-AFQMC, which can be
accomplished by combining our code with differentiable molecular integral
evaluators. The low cost-scaling and high-accuracy of AD-AFQMC open up
possibilities for extremely accurate molecular dynamics simulations.

## Data availability

Code used for AD-AFQMC calculations is available in a public GitHub repository
at reference 59, and the input and output files for these calculations are
also available at reference 61.

## Acknowledgements

We would like to thank Matthew Foulkes, Shiwei Zhang, and Miguel Morales for
helpful discussions. AM and JSK were supported by NSF Career award
CHE-2145209. SS was supported by a grant from the Camille and Henry Dreyfus
Foundation. DRR acknowledges funding support from NSF CHE–2245592. This work
utilized resources from the University of Colorado Boulder Research Computing
Group, which is supported by the National Science Foundation (Award Nos.
ACI-1532235 and ACI-1532236), the University of Colorado Boulder, and Colorado
State University. Some calculations were also performed at the Advanced
Research Computing at Hopkins (ARCH) core facility through allocation
CHE230028 from the Advanced Cyberinfrastructure Coordination Ecosystem:
Services and Support (ACCESS) program, which is supported by National Science
Foundation grants #2138259, #2138286, #2138307, #2137603, and #2138296.

## References

  * Pulay (1969) Pulay, P. Ab initio calculation of force constants and equilibrium geometries in polyatomic molecules: I. Theory. _Molecular Physics_ 1969, _17_ , 197–204.
  * Pople et al. (1979) Pople, J.; Krishnan, R.; Schlegel, H.; Binkley, J. S. Derivative studies in hartree-fock and møller-plesset theories. _International Journal of Quantum Chemistry_ 1979, _16_ , 225–241.
  * Handy and Schaefer III (1984) Handy, N. C.; Schaefer III, H. F. On the evaluation of analytic energy derivatives for correlated wave functions. _The Journal of chemical physics_ 1984, _81_ , 5031–5033.
  * Salter et al. (1989) Salter, E.; Trucks, G. W.; Bartlett, R. J. Analytic energy derivatives in many-body methods. I. First derivatives. _The Journal of Chemical Physics_ 1989, _90_ , 1752–1766.
  * Helgaker et al. (2012) Helgaker, T.; Coriani, S.; Jørgensen, P.; Kristensen, K.; Olsen, J.; Ruud, K. Recent advances in wave function-based methods of molecular-property calculations. _Chemical reviews_ 2012, _112_ , 543–631.
  * Griewank and Walther (2008) Griewank, A.; Walther, A. _Evaluating derivatives: principles and techniques of algorithmic differentiation_ ; SIAM, 2008.
  * Ekström et al. (2010) Ekström, U.; Visscher, L.; Bast, R.; Thorvaldsen, A. J.; Ruud, K. Arbitrary-order density functional response theory from automatic differentiation. _Journal of chemical theory and computation_ 2010, _6_ , 1971–1980.
  * Bast et al. (2011) Bast, R.; Ekström, U.; Gao, B.; Helgaker, T.; Ruud, K.; Thorvaldsen, A. J. The ab initio calculation of molecular electric, magnetic and geometric properties. _Physical Chemistry Chemical Physics_ 2011, _13_ , 2627–2651.
  * Abbott et al. (2021) Abbott, A. S.; Abbott, B. Z.; Turney, J. M.; Schaefer III, H. F. Arbitrary-order derivatives of quantum chemical methods via automatic differentiation. _The journal of physical chemistry letters_ 2021, _12_ , 3232–3239.
  * Zhang and Chan (2022) Zhang, X.; Chan, G. K.-L. Differentiable quantum chemistry with PySCF for molecules and materials at the mean-field level and beyond. _The Journal of Chemical Physics_ 2022, _157_ , 204801.
  * Kasim et al. (2022) Kasim, M. F.; Lehtola, S.; Vinko, S. M. DQC: A Python program package for differentiable quantum chemistry. _The Journal of chemical physics_ 2022, _156_.
  * Liao et al. (2019) Liao, H.-J.; Liu, J.-G.; Wang, L.; Xiang, T. Differentiable programming tensor networks. _Physical Review X_ 2019, _9_ , 031041.
  * Kalos et al. (1974) Kalos, M. H.; Levesque, D.; Verlet, L. Helium at zero temperature with hard-sphere and other forces. _Physical Review A_ 1974, _9_ , 2178.
  * Ceperley et al. (1977) Ceperley, D.; Chester, G. V.; Kalos, M. H. Monte Carlo simulation of a many-fermion study. _Physical Review B_ 1977, _16_ , 3081\.
  * Ceperley and Alder (1980) Ceperley, D. M.; Alder, B. J. Ground state of the electron gas by a stochastic method. _Physical review letters_ 1980, _45_ , 566.
  * Nightingale and Umrigar (1998) Nightingale, M. P.; Umrigar, C. J. _Quantum Monte Carlo methods in physics and chemistry_ ; Springer Science & Business Media, 1998.
  * Foulkes et al. (2001) Foulkes, W.; Mitas, L.; Needs, R.; Rajagopal, G. Quantum Monte Carlo simulations of solids. _Reviews of Modern Physics_ 2001, _73_ , 33.
  * Becca and Sorella (2017) Becca, F.; Sorella, S. _Quantum Monte Carlo approaches for correlated systems_ ; Cambridge University Press, 2017.
  * Zhang et al. (1995) Zhang, S.; Carlson, J.; Gubernatis, J. E. Constrained path quantum Monte Carlo method for fermion ground states. _Physical review letters_ 1995, _74_ , 3652.
  * Zhang et al. (1997) Zhang, S.; Carlson, J.; Gubernatis, J. E. Constrained path Monte Carlo method for fermion ground states. _Physical Review B_ 1997, _55_ , 7464\.
  * Zhang and Krakauer (2003) Zhang, S.; Krakauer, H. Quantum Monte Carlo method using phase-free random walks with Slater determinants. _Physical review letters_ 2003, _90_ , 136401.
  * Al-Saidi et al. (2006) Al-Saidi, W.; Zhang, S.; Krakauer, H. Auxiliary-field quantum Monte Carlo calculations of molecular systems with a Gaussian basis. _The Journal of chemical physics_ 2006, _124_ , 224101.
  * Suewattana et al. (2007) Suewattana, M.; Purwanto, W.; Zhang, S.; Krakauer, H.; Walter, E. J. Phaseless auxiliary-field quantum Monte Carlo calculations with plane waves and pseudopotentials: Applications to atoms and molecules. _Physical Review B_ 2007, _75_ , 245123.
  * Hao et al. (2018) Hao, H.; Shee, J.; Upadhyay, S.; Ataca, C.; Jordan, K. D.; Rubenstein, B. M. Accurate predictions of electron binding energies of dipole-bound anions via quantum Monte Carlo methods. _The journal of physical chemistry letters_ 2018, _9_ , 6185–6190.
  * Lee et al. (2020) Lee, J.; Malone, F. D.; Morales, M. A. Utilizing essential symmetry breaking in auxiliary-field quantum Monte Carlo: Application to the spin gaps of the C36 fullerene and an iron porphyrin model complex. _Journal of chemical theory and computation_ 2020, _16_ , 3019–3027.
  * Malone et al. (2022) Malone, F. D.; Mahajan, A.; Spencer, J. S.; Lee, J. ipie: A Python-Based Auxiliary-Field Quantum Monte Carlo Program with Flexibility and Efficiency on CPUs and GPUs. _Journal of Chemical Theory and Computation_ 2022, _19_ , 109–121.
  * Shee et al. (2019) Shee, J.; Arthur, E. J.; Zhang, S.; Reichman, D. R.; Friesner, R. A. Singlet–triplet energy gaps of organic biradicals and polyacenes with auxiliary-field quantum Monte Carlo. _Journal of chemical theory and computation_ 2019, _15_ , 4924–4932.
  * Badinski et al. (2010) Badinski, A.; Haynes, P.; Trail, J.; Needs, R. Methods for calculating forces within quantum Monte Carlo simulations. _Journal of Physics: Condensed Matter_ 2010, _22_ , 074202.
  * Casalegno et al. (2003) Casalegno, M.; Mella, M.; Rappe, A. M. Computing accurate forces in quantum Monte Carlo using Pulay’s corrections and energy minimization. _The Journal of Chemical Physics_ 2003, _118_ , 7193–7201.
  * Filippi and Umrigar (2000) Filippi, C.; Umrigar, C. Correlated sampling in quantum Monte Carlo: A route to forces. _Physical Review B_ 2000, _61_ , R16291.
  * Moroni et al. (2014) Moroni, S.; Saccani, S.; Filippi, C. Practical schemes for accurate forces in quantum Monte Carlo. _Journal of chemical theory and computation_ 2014, _10_ , 4823–4829.
  * Assaraf and Caffarel (2003) Assaraf, R.; Caffarel, M. Zero-variance zero-bias principle for observables in quantum Monte Carlo: Application to forces. _The Journal of Chemical Physics_ 2003, _119_ , 10536–10552.
  * Motta and Zhang (2017) Motta, M.; Zhang, S. Computation of ground-state properties in molecular systems: Back-propagation with auxiliary-field quantum Monte Carlo. _Journal of chemical theory and computation_ 2017, _13_ , 5367–5378.
  * Thomas et al. (2015) Thomas, R. E.; Opalka, D.; Overy, C.; Knowles, P. J.; Alavi, A.; Booth, G. H. Analytic nuclear forces and molecular properties from full configuration interaction quantum Monte Carlo. _The Journal of Chemical Physics_ 2015, _143_.
  * Sorella and Capriotti (2010) Sorella, S.; Capriotti, L. Algorithmic differentiation and the calculation of forces by quantum Monte Carlo. _The Journal of chemical physics_ 2010, _133_.
  * Poole (2014) Poole, T. Calculating derivatives within quantum Monte Carlo. 2014,
  * Zhang et al. (2019) Zhang, S.-X.; Wan, Z.-Q.; Yao, H. Automatic differentiable Monte Carlo: theory and application. _arXiv preprint arXiv:1911.09117_ 2019,
  * Glasserman (2004) Glasserman, P. _Monte Carlo methods in financial engineering_ ; Springer, 2004; Vol. 53.
  * Glynn (1990) Glynn, P. W. Likelihood ratio gradient estimation for stochastic systems. _Communications of the ACM_ 1990, _33_ , 75–84.
  * Bengio et al. (2013) Bengio, Y.; Léonard, N.; Courville, A. Estimating or propagating gradients through stochastic neurons for conditional computation. _arXiv preprint arXiv:1308.3432_ 2013,
  * Mohamed et al. (2020) Mohamed, S.; Rosca, M.; Figurnov, M.; Mnih, A. Monte carlo gradient estimation in machine learning. _The Journal of Machine Learning Research_ 2020, _21_ , 5183–5244.
  * Motta and Zhang (2018) Motta, M.; Zhang, S. Ab initio computations of molecular systems by the auxiliary-field quantum Monte Carlo method. _Wiley Interdisciplinary Reviews: Computational Molecular Science_ 2018, _8_ , e1364.
  * Shi and Zhang (2021) Shi, H.; Zhang, S. Some recent developments in auxiliary-field quantum Monte Carlo for real materials. _The Journal of Chemical Physics_ 2021, _154_ , 024107.
  * Lee et al. (2022) Lee, J.; Pham, H. Q.; Reichman, D. R. Twenty years of auxiliary-field quantum Monte Carlo in quantum chemistry: An overview and assessment on main group chemistry and bond-breaking. _Journal of Chemical Theory and Computation_ 2022, _18_ , 7024–7042.
  * Buonaura and Sorella (1998) Buonaura, M. C.; Sorella, S. Numerical study of the two-dimensional Heisenberg model using a Green function Monte Carlo technique with a fixed number of walkers. _Physical Review B_ 1998, _57_ , 11446.
  * Whitlock et al. (1979) Whitlock, P.; Ceperley, D.; Chester, G.; Kalos, M. Properties of liquid and solid He 4. _Physical Review B_ 1979, _19_ , 5598.
  * Ceperley and Kalos (1986) Ceperley, D. M.; Kalos, M. _Monte Carlo methods in statistical physics_ ; Springer, 1986; pp 145–194.
  * Rothstein (2013) Rothstein, S. M. A survey on pure sampling in quantum Monte Carlo methods. _Canadian Journal of Chemistry_ 2013, _91_ , 505–510.
  * Mahajan et al. (2022) Mahajan, A.; Lee, J.; Sharma, S. Selected configuration interaction wave functions in phaseless auxiliary field quantum Monte Carlo. _The Journal of Chemical Physics_ 2022, _156_.
  * Motta and Zhang (2018) Motta, M.; Zhang, S. Communication: Calculation of interatomic forces and optimization of molecular geometry with auxiliary-field quantum Monte Carlo. _The Journal of Chemical Physics_ 2018, _148_.
  * Shee et al. (2017) Shee, J.; Zhang, S.; Reichman, D. R.; Friesner, R. A. Chemical transformations approaching chemical accuracy via correlated sampling in auxiliary-field quantum Monte Carlo. _Journal of Chemical Theory and Computation_ 2017, _13_ , 2667–2680.
  * Assaraf et al. (2011) Assaraf, R.; Caffarel, M.; Kollias, A. Chaotic versus nonchaotic stochastic dynamics in monte carlo simulations: A route for accurate energy differences in n-body systems. _Physical Review Letters_ 2011, _106_ , 150601\.
  * Chen et al. (2023) Chen, S.; Yang, Y.; Morales, M.; Zhang, S. Algorithm for branching and population control in correlated sampling. _arXiv preprint arXiv:2307.15203_ 2023,
  * Gerratt and Mills (1968) Gerratt, J.; Mills, I. M. Force constants and dipole-moment derivatives of molecules from perturbed Hartree–Fock calculations. I. _The Journal of Chemical Physics_ 1968, _49_ , 1719–1729.
  * Schulman et al. (2015) Schulman, J.; Heess, N.; Weber, T.; Abbeel, P. Gradient estimation using stochastic computation graphs. _Advances in neural information processing systems_ 2015, _28_.
  * Sun et al. (2018) Sun, Q.; Berkelbach, T. C.; Blunt, N. S.; Booth, G. H.; Guo, S.; Li, Z.; Liu, J.; McClain, J. D.; Sayfutyarova, E. R.; Sharma, S.; Wouters, S.; Chan, K.-L. G. PySCF: the Python-based simulations of chemistry framework. _WIREs Comput. Mol. Sci._ 2018, _8_ , e1340.
  * Zhai and Chan (2021) Zhai, H.; Chan, G. K. Low communication high performance ab initio density matrix renormalization group algorithms. _The Journal of Chemical Physics_ 2021, _154_.
  * Kim et al. (2018) Kim, J.; Baczewski, A. D.; Beaudet, T. D.; Benali, A.; Bennett, M. C.; Berrill, M. A.; Blunt, N. S.; Borda, E. J. L.; Casula, M.; Ceperley, D. M., et al. QMCPACK: an open source ab initio quantum Monte Carlo package for the electronic structure of atoms, molecules and solids. _Journal of Physics: Condensed Matter_ 2018, _30_ , 195901.
  * (59) https://github.com/ankit76/ad_afqmc/.
  * Bradbury et al. (2018) Bradbury, J.; Frostig, R.; Hawkins, P.; Johnson, M. J.; Leary, C.; Maclaurin, D.; Necula, G.; Paszke, A.; VanderPlas, J.; Wanderman-Milne, S.; Zhang, Q. JAX: composable transformations of Python+NumPy programs. http://github.com/google/jax, 2018.
  * (61) https://github.com/ankit76/ad_afqmc_data.
  * Lide (2004) Lide, D. R. _CRC handbook of chemistry and physics_ ; CRC press, 2004; Vol. 85.
  * Shimizu (1970) Shimizu, F. Stark Spectroscopy of NH3 $\nu$ 2 Band by 10-$\mu$ CO2 and N2O Lasers. _The Journal of Chemical Physics_ 1970, _52_ , 3572–3576.
  * Lovas et al. (2005) Lovas, F.; Tiemann, E.; Coursey, J.; Kotochigova, S.; Chang, J.; Olsen, K.; Dragoset, R. Diatomic spectral database, NIST standard reference database 114\. 2005.
  * Meerts et al. (1977) Meerts, W.; De Leeuw, F.; Dymanus, A. Electric and magnetic properties of carbon monoxide by molecular-beam electric-resonance spectroscopy. _Chemical Physics_ 1977, _22_ , 319–324.
  * Theulé et al. (2003) Theulé, P.; Callegari, A.; Rizzo, T. R.; Muenter, J. Fluorescence detected microwave Stark effect measurements in excited vibrational states of H 2 CO. _The Journal of chemical physics_ 2003, _119_ , 8910–8915.
  * Nelson et al. (1967) Nelson, R. D.; Lide, D. R.; Maryott, A. A., et al. _Selected values of electric dipole moments for molecules in the gas phase_ ; US National Bureau of Standards Washington, DC, 1967; Vol. 10.
  * Loos et al. (2020) Loos, P.-F.; Lipparini, F.; Boggio-Pasqua, M.; Scemama, A.; Jacquemin, D. A mountaineering strategy to excited states: Highly accurate energies and benchmarks for medium sized molecules. _Journal of Chemical Theory and Computation_ 2020, _16_ , 1711–1741.
  * Qin et al. (2016) Qin, M.; Shi, H.; Zhang, S. Coupling quantum Monte Carlo and independent-particle calculations: Self-consistent constraint for the sign problem based on the density or the density matrix. _Physical Review B_ 2016, _94_ , 235119.
  * Qin (2023) Qin, M. Self-consistent optimization of the trial wave function within the constrained path auxiliary field quantum Monte Carlo method using mixed estimators. _Physical Review B_ 2023, _107_ , 235124.


# Discovering Intrinsic Reward with Contrastive Random Walk

Zixuan Pan
Department of Computer Science Engineering
University of Michigan
Ann Arbor, USA
<EMAIL_ADDRESS>Zihao Wei
Department of Computer Science Engineering
University of Michigan
Ann Arbor, USA
<EMAIL_ADDRESS>
Yidong Huang
Department of Computer Science Engineering
University of Michigan
Ann Arbor, USA
<EMAIL_ADDRESS>
Aditya Gupta
Department of Electrical and Computer Engineering
University of Michigan
Ann Arbor, USA
<EMAIL_ADDRESS>

https://github.com/Zi-hao-Wei/Contrastive-Random-Walk-as-Intrinsic-Rewards
Corresponding author: Zihao Wei<EMAIL_ADDRESS>
###### Abstract

The aim of this paper is to demonstrate the efficacy of using Contrastive
Random Walk as a curiosity method to achieve faster convergence to the optimal
policy. Contrastive Random Walk defines the transition matrix of a random walk
with the help of neural networks. It learns a meaningful state representation
with a closed loop. The loss of Contrastive Random Walk serves as an intrinsic
reward and is added to the environment reward. Our method works well in non-
tabular sparse reward scenarios, in the sense that our method receives the
highest reward within the same iterations compared to other methods.
Meanwhile, Contrastive Random walk is more robust. The performance doesn’t
change much with different random initialization of environments. We also find
that adaptive restart and appropriate temperature are crucial to the
performance of Contrastive Random Walk.

##  1 Introduction

Reinforcement learning (RL) learns policies that maximize the reward given by
the environment. Methods considering only the environment rewards work well
when rewards are dense. However, when the rewards are sparse, exploration has
been proven important to make faster convergence of finding optimal policy. In
the past few decades, dozens of methods have been proposed to encourage
exploration by introducing an exploration bonus as an intrinsic rewards, which
decides based on the chosen time of a certain state[3, 1, 15].

With the advance of RL, people are interested in applying it to more
complicated scenarios, where the observations are some high-dimensional
signals (e.g. video frame)[9]. To understand the signals and better define
states, instead of human definition, features extracted by neural networks are
applied to represent states[9, 17]. It provokes two notable changes in the
definition of states: states become high-dimensional and states are no longer
invariant to observations. These changes make counting the visits to certain
states unavailable.

To deal with the problem that traditional exploration rewards can’t be used in
complicated environments, researchers propose the concept of curiosity and use
neural network loss as rewards[15]. Neural networks are good at understanding
high-dimensional signals and have astonishing generalization ability[13],
which perfectly fits the scenario. The loss of neural networks can be viewed
as a pseudo-count, that is lower when the state is frequently visited and
higher for un-visited states.

However, two problems remain with many existing curiosity methods. On one
hand, many of the methods lack mathematical foundations. The intuition of why
these methods work is unclear. On the other hand, neural networks in most
curiosity methods neural networks only serve as function fitters and are not
extracting meaningful features. We argue that these neural networks can also
help build a better state encoder if properly designed. Good state
representations can thus further leverage the generalization ability of neural
networks and implicitly create a good clustering of observations. Such
clusters can well define curiosity for similar observations, i.e. intrinsic
rewards will not decrease when seeing exactly the same observations, but will
also decrease for similar states. More meaningful representations will be
better at defining what is "similar".

To tackle these two problems, we proposed applying Contrastive Random Walk
(CRW) as a curiosity method[8]. CRW can extract meaningful representations
from time series by creating a closed-loop for the series. In RL, experience
is naturally a time series and hence CRW could also be able to extract good
state representations. Meanwhile, we provide a mathematical proof that the
loss of CRW is a good approximation to the information gain, which is a good
measurement of how visited a state is.

We tested our method on two environments: "MountainCar-v0" and "Montezuma’s
Revenge". Both environments have sparse environment rewards and thus good
exploration strategies are critical to ensure a faster convergence of policy.
Meanwhile, both environments use neural networks to get state representations.
MountainCar-v0 is a simple environment where the agent’s task is to reach the
top of the mountain within less time. The observation is some attributes of
the car, including speed and position. While Montezuma’s Revenge is a more
complicated environment… The observation is just video frames. Our method has
obtained better performance than all the other baselines, including the
curiosity method, count-based method, and randomized strategy. Meanwhile, our
methods are less subject to randomness of environments and randomness of
policy that can give similar performance on all kinds of random initialization
compared to other methods.

##  2 Preliminary and Notations

###  2.1 Reinforcement Learning

Given a Markov decision process (MDP), which can be described as
$\mathcal{M}=\langle\mathcal{S},\mathcal{A},P,E,\gamma\rangle$, where
$\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space,
$P\left(s_{t+1}\mid s_{t},a_{t}\right)$ is the probability distribution over
transitions, $E\left(s_{t},a_{t}\right)$ is the environment reward function,
and $\gamma$ is the discount factor. We will focus on infinite-horizon MDPs.

###  2.2 Intrinsic Reward

We define an exploration bonus based on the current state $I(s_{t})$ and add
it to the reward $E\left(s_{t},a_{t}\right)$ given by the environment, getting
a total reward $R\left(s_{t},a_{t}\right)=E\left(s_{t},a_{t}\right)+I(s_{t})$.
The goal of the addition of exploration bonus is to embed the new reward to
existing reinforcement learning algorithms and help to make a better policy.

###  2.3 Non-tabular State Representation

The non-tabular environment of reinforcement learning can be viewed as an
infinite-horizon stochastic dynamic programming, where states are some high-
dimensional representations extracted by neural networks, denoted as
$s=\phi(o)$, $s=[s_{1},s_{2},...,s_{n}]^{T}\in\mathbb{R}^{n}$, where $o$ is
the observation (video frame). A prominent characteristic of non-tabular
states is that the state representations are variant. As a neural network is
applied and updated with experience, state representation of observations will
change (i.e. $s$ will be updated with time). performance for count-based
methods.

##  3 Related Works

###  3.1 Curiosity Driven Intrinsic Reward

In order to make count-based methods also work for non-tabular state
representations, discrete "counts" are changed into some density and
continuous functions[3]. However, this approximation also incurs performance
drop for count based methods. Therefore, curiosity methods are proposed to
create a new kind of exploration bonus for non-tabular states.

Curiosity driven intrinsic rewards encourages the agent to explore the
environment by learning an intrinsic reward $I(s_{t})$ besides the environment
reward[5, 15, 4, 16, 6]. The intrinsic reward is given by neural network loss.

$I(s_{t})=L_{nn}(s_{t})+C$

where C is a constant.

It applies the fact that if an input has been seen, loss of the neural network
will decrease, indicating a lower intrinsic reward. Neural network loss style
intrinsic reward also applies the generalization ability that makes it
suitable for continuous state representation[13]. The loss will drop for
exactly the same input $s_{t}$, but also states with similar feature to
$s_{t}$, which makes it appropriate for non-tabular tasks.

There are mainly two types of curiosity methods: action dependent methods and
state-only methods. State-only methods only use state representations from
observations as inputs to the neural network[5, 16, 10]. Action dependent
methods take both actions and state representations as input into neural
networks[15, 6].

###  3.2 Information gain as Measurement of Intrinsic Reward

The information gain or mutual information is defined as:

$IG(X;Y)=\mathbb{E}_{p(x,y)}\log\frac{p(x,y)}{p(x)p(y)}=H(X)-H(X\mid Y)$

It is widely used in Reinforcement learning as a measurement of intrinsic
reward[3, 6].

For curiosity methods, information gain is made larger for unvisited states,
while lower for frequently visited states to encourage the agent exploring the
environment.

There exist various definitions of information gain in Reinforcement
Learning[17, 20]. We will use the state-only transition information, which has
been proved to be a good indicator for state representation[17].

$IG(s_{t+k};s_{t})=H(s_{t+k})-H(s_{t+k}|s_{t})$

In non-tabular scenario, we can change it to

$IG(s_{t+k};s_{t})=\Sigma_{i=1}^{n}H(s_{t+k,i})-H(s_{t+k,i}|s_{t,i})$

. And the intrinsic reward for $s_{t+k}$ is thus

$I(s_{t+k})=IG(s_{t+k};s_{t})$

###  3.3 State Representation Learning

There are a bunch of previous works on how to extract better state
representations[19, 12, 17, 7]. Better representations have been proved to be
crucial to make RL algorithms perform better in many scenarios. Our work is
most similar to [7], where the information gain is used to evaluate how good a
representation is. The goal of this paper is to maximize information gain
between inputs and outputs of the state encoder. Our method is also arguably
increasing such information gain.

##  4 Contrastive Random Walk

Figure 1: Overview of Contrastive Random Walk. A backtracing process is
appended to the forward path. States are connected with affinity matrices.

Inspired by previous work, we view each experience as a reversible process[8].
Given an experience starting with $s_{1}$, ending with $s_{t}$, we also
obtained another experience starting with $s_{t}$, ending with $s_{1}$. We
concatenate the reverse experience to the original experience as shown in
Figure1. We give each state an embedding, denoting the encoder as $\phi$.

$A_{t}^{t+1}(i,j)=\operatorname{softmax}\left(f(s_{t}s_{t+1}^{\top})\right)_{ij}=\frac{\exp\left(d\left(\mathbf{s}_{t}^{i},\mathbf{s}_{t+1}^{j}\right)/\tau\right)}{\sum_{l=1}^{N}\exp\left(d\left(\mathbf{s}_{t}^{i},\mathbf{s}_{t+1}^{l}\right)/\tau\right)}$

where
$d\left(s_{1},s_{2}\right)=\left\langle\left(s_{1}\right),\left(s_{2}\right)\right\rangle$
is the similarity between embeddings. $\tau$ is the temperature, which is a
constant in softmax. By using the softmax, we make this affinity matrix also a
transition matrix. Then we can define a continuous transition from t to t+k.

$\bar{A}_{t}^{t+k}=\prod_{i=0}^{k-1}A_{t+i}^{t+i+1}=P\left(s_{t+k}\mid
s_{t}\right)$

Intuitively, it just predicts the state of t+k given state of t. Our goal is
defined as:

$\mathcal{L}_{cyc}^{k}=\mathcal{L}_{CE}\left(\bar{A}_{t}^{t+k}\bar{A}_{t+k}^{t},I\right)=-\sum_{i=1}^{N}\log
P\left(s_{t+2k,i}\mid s_{t,i}\right)$

When the objective function is minimized, each affinity matrix will be a one-
hot matrix containing only 0 and 1, with one 1 in each row and column. This is
to give a one-by one map from previous state the current state. Each node in
the previous state embedding aims to maximize the similarity with another node
in the current state, and minimize the similarity with all the other nodes in
the current state, which can be viewed as a contrastive learning with n-1
negative samples.

For the intrinsic reward at each state, we use

$I(s_{t})=\mathcal{L}_{cyc}^{k}+C$

starting from state $s_{t}$ and ends with $s_{t+k}$. A simple choice of $k$
will be 1 for all states if we use the replay buffer and Deep Deterministic
Policy Gradient algorithm. If the batch contains the entire experience, we can
set $k>1$. It can provide implicitly supervision to intermediate states.

###  4.1 Adaptive Restart

Since the environment we use is quite simple, the convergence of CRW is often
too fast, and results in the intrinsic reward approaching zero with more
iterations. Therefore, we applied the strategy of adaptive restart in
optimization to settle the problem[14]. The parameters of CRW encoder are
reset after the average intrinsic reward is smaller than a threshold. A psuedo
code of our method can be found in Algorithm 1.

Algorithm 1 model.forward()

Input: $x$, threshold, temperature

#x size: $B\times T\times N$

A = einsum("bti,btj$\rightarrow$ btij" x[:,:,:-1], x[:,:,1:]) / temperature

#affinity matrix of forward path with size $B\times T\times N\times N$

A_ = A.transpose(-2,-1)

for i in range(T):

AA[i] = bmm(A[:,i],A_[:,i])

AA = multiply(AA)

# Multiply all elements in AA

loss = sum(-log(diag(AA,dim=-2,-1)))

# Compute Cross Entropy loss.

If loss $\leq$ threshold:

reset_param()

# Adaptive restart

###  4.2 Approximation to Information Gain

For the contrastive random walk,

$\displaystyle
IG(s_{t+2k};s_{t})=\Sigma_{i=1}^{n}H(s_{t+k,i})-H(s_{t+2k,i}|s_{t,i})$
$\displaystyle=\Sigma_{i=1}^{n}\mathbb{E}_{s_{t+2k}\sim\rho^{\prime}}[log(s_{t+2k,i})]+\mathbb{E}_{p(s_{t+2k})\sim\rho^{\prime},s_{t}\sim\rho}[log(p(s_{t+2k,i}|s_{t,i}))]$

where $\rho^{\prime}$ is the discounted state distribution for backtracing
states, and $\rho$ is the discounted state distribution for forward states.

Because the backtracing gives exact the same states of forward path,

$\mathbb{E}_{s_{t+2k}\sim\rho^{\prime}}[log(p(s_{t+2k,i}))]=\mathbb{E}_{s_{t+2k}\sim\rho}[log(p(s_{t,i}))]$

is not determined by the contrastive random walk, we just replace it by a
constant $\mathrm{C}$.

Contrastive random walk optimizes the second term, such that:

$\mathbb{E}_{p(s_{t+2k})\sim\rho^{\prime},s_{t}\sim\rho}[log(p(s_{t+2k,i}|s_{t,i}))]=\frac{\Sigma_{m=1}^{M}p_{m}(s_{t+2k,i}|s_{t,i})log(p_{m}(s_{t+2k,i}|s_{t,i}))}{M}$

Denote $P=\max p_{m}$, and when $P$ is close to 1, the above equation can be
approximated by:

$\mathbb{E}_{p(s_{t+2k})\sim\rho^{\prime},s_{t}\sim\rho}[log(p(s_{t+2k,i}|s_{t,i}))]\approx\frac{log(P(s_{t+2k,i}|s_{t,i}))}{M}$

Remove the constant $M$, we get

$\mathbb{E}_{p(s_{t+2k})\sim\rho^{\prime},s_{t}\sim\rho}[log(p(s_{t+2k,i}|s_{t,i}))]\approx
log(P(s_{t+2k,i}|s_{t,i}))$

This approximation is sufficient as the goal of contrastive random walk is to
maximize $P$.

Thus, we can rewrite the information gain as:

$\displaystyle
IG(s_{t+2k};s_{t})=\mathbb{E}_{s_{t+2k}\sim\rho}[log(p(s_{t,i}))]+L_{cyc}^{k}$

Then, by optimizing the information gain of $s_{t+2k}$, we are also optimizing
the information gain of $s_{t}$, as they share the same state representation.

$\displaystyle I(s_{t})=I(s_{t+2k})=IG(s_{t+2k};s_{t})$

##  5 Experiments

###  5.1 Training Settings

####  5.1.1 Environments

Because of the lack of computational resources, we use a simple environment
called "MountainCar-v0" [11], to test our methods and other baselines. In the
MountainCar-v0 environments(shown in Figure 2), a car is originally located in
the valley and the agents are forced to learn a policy which can drive up the
car from the bottom to the mountain on the right. However, since the car’s
engine is not strong enough, the only way to succeed is to drive back and
forth to build up momentum. The entire environment has an observation of 2
dimension $(o_{1},o_{2})$, denoting the position of the car on the X-axis and
the velocity of the car; three action space $(a_{1},a_{2},a_{3})$, standing
for accelerating right, left, or not. In order to encourage the agents to
arrive at the destination as fast as possible, a reward $r$ of -1 is given at
each step when the car hasn’t arrived, and the total steps in one episode are
200, meaning the reward is ranging from -200 to 0.

Figure 2: Overview of the MountainCar-V0 environment.

###  5.2 Implementation Details

To test the effectiveness of the exploration bonus we proposed, we combine our
CRW Intrisic Reward with a Deep Q Network(DQN) backbone. Here we use a three-
layer MLP with the default hidden dimensions (256,128,64) with 0.3 dropout
ratio as our state encoder $\phi$ and Q-network. We compare our methods with
several exploring baselines, which are Epsilon-greedy, Upper Confidence Bound
(UCB), and Random Network Distillation(RND). These baseline models will use
the same DQN as backbone. Q value is updated with the function

$\tilde{Q}(s_{t},a_{t})\longleftarrow\left(1-\beta_{k}\right)\tilde{Q}(s_{t},a_{t})+\beta_{k}\left(E(s_{t},a_{t})+I(s_{t})+\alpha\max_{v}\tilde{Q}(s^{\prime},v)\right)$

The target and predictor network of RND will share the same architecture of
state encoder of our CRW. The main hyper-parameters and their default value
used in the experiments are, temperature $\tau=1$, DQN learning rate
$lr_{1}=1e-3$, state encoder learning rate $lr_{2}=1e-5$, restart running loss
threshold $s=5e-3$, batch size $B=128$, replay buffer size $M=100000$, episode
num $N=1000$. And the entire experiment is carried out on one NVIDIA RTX 3080
GPU. More details can be found in the code.

###  5.3 Results

Here we report the performance of our CRW and other baselines in Figure 3. Our
methods reach a final episode running reward of -110.607, which greatly
surpasses all the other baselines, where epsilon greedy achieves -139.57, UCB
achieves -121.037 and RND achieves -118.297. Meanwhile, we also observe that
our methods begin to have rewards earlier than other networks. This phenomenon
shows that our method wastes less efforts on testing a similar states, which
benefits the convergence speed.

(a) Epsilon Greedy

(b) Upper Confidence Bound

(c) Random Network Distillation

(d) Contrastive Random Walk

Figure 3: The running episode reward of the training process of CRW and
different baselines.

###  5.4 Ablation Study

Here we show how each hyperparameter will contribute to the final results of
our methods in Figure 4. The hyper-parameter tuned including feature dimension
of the state encoder, temperature of crw, batch size during training and
learning rate. Except for the parameter of being tuned, other parameters are
set as default.

(a) Tuning feature dimension in state encoder

(b) Tuning crw temperature

(c) Tuning learning rate of the state encoder

(d) Tuning restart running loss ratio

Figure 4: The running episode reward of the training process of CRW and
different baselines.

###  5.5 Discussion

As shown in the results, temperature is an important factor to improve model
performance. We find that the reward at 1000 iteration is linearly related to
the choice of temperature. Generally, a small temperature with adaptive
restart works better. Also, we found that adaptive restart is essential for
CRW. When the restart threshold is set smaller, the performance of CRW may
drop significantly.

##  6 Experiments on more complicated environments

Besides the “MountainCar-v0" [11] environment, we also tried to implement our
algorithm on the classic atari game “MontezumaRevenge-v0" [2] which has a
sparse reward. As a traditional arcade game 5, the observation space of the
game is a RGB image with hieght 210 pixels and width 160 pixels. The aim for
the agent is to go through different rooms and get as many points as possible.
Former work [5] showed that with the help of intrinsic rewards, the agent can
get out of the first room and gain points while normal RL algorithms fail to
get a single point.

Figure 5: Overview of the MontezumaRevenge-v0 environment.

###  6.1 Implementation Details

For the agent network, we used an actor-critic architecture with cnn state
encoder. We use a PPO [18] training method to update the networks, which is a
same architecture suggested in the RND model [5]. To get the experience, we
implement 64 workers using python multiprocessing. We also set the PPO
learning rate $lr=1e-5$. For the CRW model that we proposed, we used a
slightly different mean of encoding.

The pre-processing of the observations starts with compressing the image into
a 160*160 large image. Getting the compressed image, we split the image into
16 patches, each of 64*64 large. The overlapping region between neighboring
patches is used to avoid missing information near the edges. After getting the
16 patches of images, we use a 4-layer convolution network with leakyrelu as
activation layer to get the embedding of each patch. With these 16 embedding,
we implement the CRW model and get the intrinsic reward. To further extract
property of the moving characters, we use a localized method, which also
extract 16 smaller patches around the moving character to produce intrinsic
reward. We trained the network on a NVIDIA RTX 3080 GPU. More implementation
details can be found in the code.

###  6.2 Results and qualitative analysis

Because of the lack of computational resource, we just tried the
implementation for 300 episodes and investigated the change of intrinsic
reward. Because 300 episodes were not enough for the agent to get out of the
first room, the implementation was in fact a purely exploration problem. As is
shown in Figure 6, we can see that as the episode number increase, the
intrinsic reward given by our CRW model first increase and then decrease,
showing that the exploring first arise curiosity of the world, and then after
many states are visited, the model lose the curiosity, thus the loss will drop
to a lower value. This implies that our model can do the job of helping the
algorithm to explore more unvisited states.

Figure 6: Result of implementation in MentonzumaRevenge-v0.

##  7 Conclusion

We demonstrated a new intrinsic reward with Contrastive Random Walk that
achieves better performance than other methods on sparse-reward non-tabular
scenarios. This method has a more solid mathematical foundation related to
information gain, and whose neural network learns meaningful representations
of states. These two advantages ensures our method also has the potential to
provide good results in more complicated environments. Our success also
provokes thinking on the importance of learning a good representation win the
meantime of exploring the environment. It also encourages us to explore the
mathematical proof of a faster conversion using this method as future work.

## References

  * [1] Peter Auer.  Using confidence bounds for exploitation-exploration trade-offs.  Journal of Machine Learning Research, 3(Nov):397–422, 2002.
  * [2] M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling.  The arcade learning environment: An evaluation platform for general agents.  Journal of Artificial Intelligence Research, 47:253–279, jun 2013\.
  * [3] Marc Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos.  Unifying count-based exploration and intrinsic motivation.  Advances in neural information processing systems, 29, 2016.
  * [4] Yuri Burda, Harri Edwards, Deepak Pathak, Amos Storkey, Trevor Darrell, and Alexei A Efros.  Large-scale study of curiosity-driven learning.  arXiv preprint arXiv:1808.04355, 2018.
  * [5] Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov.  Exploration by random network distillation.  arXiv preprint arXiv:1810.12894, 2018.
  * [6] Ildefons Magrans de Abril and Ryota Kanai.  Curiosity-driven reinforcement learning with homeostatic regulation.  In 2018 international joint conference on neural networks (ijcnn), pages 1–6. IEEE, 2018.
  * [7] R Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon, Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio.  Learning deep representations by mutual information estimation and maximization.  arXiv preprint arXiv:1808.06670, 2018.
  * [8] Allan Jabri, Andrew Owens, and Alexei Efros.  Space-time correspondence as a contrastive random walk.  Advances in neural information processing systems, 33:19545–19560, 2020.
  * [9] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al.  Human-level control through deep reinforcement learning.  nature, 518(7540):529–533, 2015.
  * [10] Shakir Mohamed and Danilo Jimenez Rezende.  Variational information maximisation for intrinsically motivated reinforcement learning.  Advances in neural information processing systems, 28, 2015.
  * [11] Andrew William Moore.  Efficient memory-based learning for robot control.  Technical report, University of Cambridge, 1990.
  * [12] Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine.  Near-optimal representation learning for hierarchical reinforcement learning.  arXiv preprint arXiv:1810.01257, 2018.
  * [13] Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro.  Towards understanding the role of over-parametrization in generalization of neural networks.  arXiv preprint arXiv:1805.12076, 2018.
  * [14] Brendan O’donoghue and Emmanuel Candes.  Adaptive restart for accelerated gradient schemes.  Foundations of computational mathematics, 15(3):715–732, 2015.
  * [15] Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell.  Curiosity-driven exploration by self-supervised prediction.  In International conference on machine learning, pages 2778–2787. PMLR, 2017.
  * [16] Deepak Pathak, Dhiraj Gandhi, and Abhinav Gupta.  Self-supervised exploration via disagreement.  In International conference on machine learning, pages 5062–5071. PMLR, 2019.
  * [17] Kate Rakelly, Abhishek Gupta, Carlos Florensa, and Sergey Levine.  Which mutual-information representation learning objectives are sufficient for control?, 2021.
  * [18] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms, 2017.
  * [19] Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin.  Decoupling representation learning from reinforcement learning.  In International Conference on Machine Learning, pages 9870–9879. PMLR, 2021.
  * [20] Stas Tiomkin and Naftali Tishby.  A unified bellman equation for causal information and value in markov decision processes.  arXiv preprint arXiv:1703.01585, 2017.