Iterative improvement of free energy landscape reconstructions with optimal protocols derived using automatic differentiation

Free energy landscapes encode the kinetics, intermediates, and transition states that govern stochastic 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 energy landscape. Here, harnessing automatic differentiation, 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 in JAX MD; (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.