First-hand information for everyone
Main discovery LEAP (Learned Energy on Action Paths) frames "action learning" as recovering a stochastic action — a trajectory-level objective — from observed trajectories. It anchors an unknown action to a tractable reference process and learns a pathwise correction from comparisons between observed and reference trajectories. (An "action" here means a function that assigns a cost or log density to whole trajectories.)
How LEAP works LEAP introduces a tractable reference path distribution and represents the unknown target action as the reference action plus a correction equal to the path log-density ratio between data and reference. It parameterizes that correction as G_theta and trains it by comparing one observed trajectory with K reference trajectories using a setwise contrastive (InfoNCE) objective. Training is done with minibatches. For each observed trajectory in a minibatch, the algorithm samples K reference trajectories and evaluates G_theta on the observed and reference trajectories. It then updates θ by a stochastic gradient step on the InfoNCE objective (see Alg. 1). After training, the learned correction is centered under the reference and the learned action is S_theta(Φ;c) = -log p_ref(Φ|c) - eG_theta(Φ;c).
Structured correction for shared-diffusion Markov models When data and reference are discrete-time Gaussian Markov processes with the same diffusion, the exact path log-likelihood ratio decomposes over transitions. The correction is a sum over time steps of a velocity-alignment term and a normalization cost per step. In this shared-diffusion case the structured correction equals the exact path log-likelihood ratio for drift differences.
Identification at the population level Under mild conditions (positive densities on the same support and a model class that contains the true correction), every population minimizer of the comparison objective recovers the target log-density ratio up to a context-dependent additive constant. After centering under the reference, the learned action matches the true action up to a context constant that does not affect within-context path rankings or optimization.
Action error controls downstream inference If the learned action differs from the true action by at most ε on a domain of interest, then (1) a δ-suboptimal trajectory under the learned objective is at most 2ε+δ suboptimal under the true objective, and (2) the log density ratio between the learned and true path distributions is bounded by 2ε. Thus action error provides a quantitative bound on later inference and on how much normalized trajectory densities can change.
Double-well experiment — recovering drift and reusing actions Setup: overdamped diffusion with b(x)=x−x^3 and potential U(x)=(x^2−1)^2/4. Trajectories start at x0=−1 and are selected when x_T∈[0.9,1.1] at T=3. Data used 5,000 selected trajectories per seed with diffusion D=0.20 and Δt=0.01. LEAP used a Brownian reference conditioned on the same terminal window. Results: LEAP recovered the unconditioned drift with drift MSE (3.3 ± 2.7) × 10^−4. By comparison, Kramers–Moyal gave 0.8132 ± 0.0030, cubic velocity regression 0.7752 ± 0.0043, and the zero-field baseline 0.0759. Reuse: Keeping the learned action fixed and changing endpoints and horizon, LEAP inferred near-optimal path costs. At T=4 LEAP gave true path cost 0.2533 ± 0.0001 versus analytic 0.2532. At T=16 LEAP gave 0.2506 ± 0.0006 versus analytic 0.2500. By contrast, cubic velocity regression and a geometric path baseline performed worse on these reused-inference tasks.
Nonadditive orbital correction — accumulating weak effects across time Setup: a simulated multi-body orbital system initialized from JPL HORIZONS. An unobserved binary latent z∈{−1,+1} perturbs the solar mass to m_⊙(1+zη), fixed per trajectory. This creates a nonadditive path correction whose evidence accumulates across the whole trajectory. Results: Models that operate on single transitions (LEAP-Local, Transition NCE) stayed near the constant-prediction baseline (NRMSE_ΔG ≈ 1.0013 and 1.0008). A complete-path model (LEAP-Global) that accumulates features across time reached NRMSE_ΔG ≈ 0.228 — reducing about 95% of the squared variation in exact correction differences. LEAP-Global recovered the hidden mass perturbation magnitude as 1.02 × 10^−4 (true value 1.04 × 10^−4). As system size and horizon grew (e.g., 42 bodies with 256 transitions), LEAP-Global continued to accumulate the small effect (NRMSE 0.111) while LEAP-Local stayed near ~1.0015.
Molecular-path experiments — completion, temperature transfer, preferences and constraints Data and training: Alanine dipeptide trajectories were generated by OpenMM Langevin simulations and reduced to backbone angles (ϕ,ψ). Datasets included 169 reactive C7eq→C7ax trajectories at 500 K and 613 at 600 K. LEAP was trained in these reduced coordinates. Completion results: When endpoints and three interior states were observed at 600 K, LEAP achieved circular MSE 0.0931 ± 0.0016. Conditional DDPM scored 0.119 ± 0.012 and conditional flow matching 0.125 ± 0.011. With endpoints only, LEAP gave error 0.6225 while velocity regression and binary NRE gave 0.6609 and 0.6682 respectively. Conditional generators trained with completion masks achieved lower error in the endpoints-only setting. Temperature transfer: Keeping a learned 600 K correction fixed but using the target diffusion/reference, training at 600 K and evaluating at 500 K gave within-context AUROC 0.9579 ± 0.0057 (direct 500 K training gave 0.9403 ± 0.0099). Training at 500 K and evaluating at 600 K gave 0.9719 ± 0.0068 (direct 600 K training 0.9754 ± 0.0012). These results show the learned correction remains useful after temperature changes. Preferences and geometric constraints: Using the same learned 600 K action, increasing a basin-preference weight from 0 to 20 raised the mean preference M for the shown path from 0.0658 to 0.2517 and increased projected transition cost from 0.9389 to 0.9809 nats per active step. Imposing an angular bound |ϕ| ≤ 150° changed the route while preserving the terminal condition. These demonstrate reuse of a learned action with added preferences and constraints without retraining.
Limitations and caveats LEAP requires sufficient overlap between the reference and observed path distributions. Poor overlap makes the correction harder to estimate accurately. The structured correction expression (Eq. 9) is exact only for Gaussian Markov processes with shared diffusion. State-dependent diffusion or projecting dynamics to reduced coordinates can introduce additional likelihood-ratio terms, memory effects, or state-dependent friction. In reduced-coordinate experiments (for example, (ϕ,ψ) for alanine), LEAP therefore learns an effective action in the reduced coordinates rather than the full molecular action.
Takeaway LEAP offers a practical procedure to learn a reusable trajectory-level objective by anchoring an unknown action to a tractable reference and learning a correction via trajectory comparisons. The learned action can be reused for path ranking, constrained inference, transfer across conditions, and steering, within the identified limitations above.