Iterative improvement of free energy landscape reconstructions with optimal protocols derived from differentiable simulations

TL;DR

This work tackles the challenge of reconstructing unknown free energy landscapes from nonequilibrium trajectories by introducing an iterative scheme that alternates between landscape estimation via the Hummer–Szabo framework and automatic-differentiation-driven minimum-dissipation protocol optimization. By updating driving protocols based on a differentiable Brownian dynamics model, the method progressively improves free energy reconstructions for 1D and 2D control without prior landscape knowledge, reducing bias and variance relative to naive linear driving. The authors demonstrate successful recovery of diverse landscapes (bistable symmetric, asymmetric, and triple-well) across near- and far-from-equilibrium regimes, with rapid convergence (often 1–4 iterations) for many cases. They also discuss degeneracy in the loss landscape, scalability to higher dimensions, and future directions toward experimental realization and alternative loss objectives for robust, high-dimensional control. Overall, the approach broadens the toolkit for extracting full free-energy landscapes from unknown systems using differentiable simulations and iterative protocol optimization.

Abstract

Free energy landscapes encode the kinetics, intermediates, and transition states that govern molecular processes and are thus a key target of single biomolecule research. Typical approaches to deriving optimal, error-minimizing, non-equilibrium driving protocols for estimating these landscapes require a priori knowledge of the landscape. Here, we present an alternative: an iterative algorithm for optimizing full free energy landscape reconstructions which can be used alongside experiments on unknown landscapes. Our approach (i) takes experimental or simulated trajectory data; (ii) reconstructs an `approximate' energy landscape; (iii) derives optimal control protocols from low-dimensional differentiable Brownian dynamics simulations on the candidate landscape using automatic differentiation; (iv) re-runs the experiment or simulation using the updated protocol; and (v) iterates until convergence. Using this approach, we recover known benchmarks from the literature and probe far-from-equilibrium regimes for symmetric, asymmetric, and triple-well energy landscapes under both 1- and 2-dimensional control. Our control protocols -- derived with no a priori knowledge of the energy landscape -- yield substantially reduced variance and bias in free energy landscape reconstructions compared to naive linear protocols.

Figures (10)

• Figure 1: Schematic of our iterative scheme for reconstructing a priori unknown free energy landscapes. (a) Nonequilibrium barrier crossing trajectories are collected, either through simulations or experiments. Shown here is a hypothetical laser optical tweezers (LOT) setup measuring the unfolding of a DNA hairpin. Reproduced from Ref. Woodside2016 with permission. (b) For each barrier crossing trajectory, external work is measured as a function of time. The process is stochastic, as seen by the width of the work distributions. (c) Using Equation \ref{['eq:hummerszabo']}, the trajectories are used to reconstruct an estimated underlying free energy landscape as a function of a molecular coordinate. (d) Brownian simulations are performed in JAX MD and automatically differentiated as described in section \ref{['sec:opt_barrier_crossing']} and reference engel_optimal_2023 to minimize loss, which here is the dissipated work. This yields optimal control protocols for the time-dependent stiffness and position of the external harmonic forcing potential. Using these new control protocols, the experiments or simulations are performed again and the process repeats from (a) until the landscape estimate converges.
• Figure 2: Performance of the iterative landscape reconstruction scheme as quantified by the landscape bias, Equation \ref{['eq:bias']}, for asymmetric, bistable symmetric, and triple well landscapes; various barrier heights; and for slower (left panel, (a)) and faster (right panel, (a)) driving. The results of using the iterative procedure for 1D (orange) and 2D (blue) control are compared to using the a priori knowledge of the landscape directly for 1D (purple) and 2D (red) control. Also shown is the landscape bias resulting from using a naive linear protocol (green). Landscapes were generated with 1000 trajectories. For 1D control, all simulations used a default $k_s = 0.4$ pN/nm throughout the entire simulation.
• Figure 3: Iterative scheme of 2D protocol landscape reconstructions and optimized protocols for 1 ms simulations with 80$\, k_{\textup{B}}\textup{T}$ barrier height. Each landscape had 10 total iterations, with the darkest line representing the protocol for the final iteration. 2D protocols optimized with the iterated scheme recovers near-perfect landscape reconstruction performance of 2D protocols optimized with gradients of the true underlying landscape, despite differing in form from the protocols that used knowledge of the 'true' landscape (shown in red), indicating degeneracy in the loss landscape.
• Figure 4: Standard deviation of free energy landscape reconstructions. Individual landscapes are constructed with $300µs$ simulations and an $25\, k_{\textup{B}}\textup{T}$ barrier height, with 1000 trajectories per landscape. For every protocol type we generated 100 independent landscapes; $\sigma\left(\Delta\hat{F}\right)$ denotes the empirical standard deviation at each position bin, computed across those landscapes.
• Figure 5: Histograms of work distributions for 300 µs simulations at near-equilibrium 20$\, k_{\textup{B}}\textup{T}$ (left) and far-from-equilibrium 40$\, k_{\textup{B}}\textup{T}$ (right) barrier heights. Far-from equilibrium, AD optimized 1D protocols (True Position, green) increase skewness, despite decreasing average work. The increased positive skewness decreases landscape reconstruction accuracy, explaining the poor performance of 1D protocols (True Position) in Fig. \ref{['fig:bias-barrier']}.