PaperlayerView exact revision on arXiv

Source-linked AI summary

Action Learning for Scientific Simulations

Sarah Liaw, Yuanqi Du, Yilun Du

arXiv:2610.31585v1cs.LG

TL;DR

Trajectory data often do not provide the known action needed to evaluate and optimize complete paths. LEAP learns a stochastic action by comparing observed trajectories with a tractable reference process. Across physical, orbital, and molecular settings, the learned action recovers dynamics, captures temporal dependencies, and supports reuse across inference tasks and conditions.

  • Problem

    Trajectory observations often lack a known action for assigning probabilities or costs to complete paths, and existing learned methods do not directly provide a reusable path objective.

    Source: §1 · p. 2

  • Method

    LEAP learns a differentiable stochastic action by combining a tractable reference action with a learned correction for differences between observed and reference path distributions.

    Source: §2 · p. 3, §2.1 · p. 4

  • Results

    Across the double well, orbital benchmark, and molecular coordinates, the learned action recovers physical drift, captures nonadditive interactions, and supports diverse inference tasks and conditions.

    Source: §5 · p. 10

  • Takeaways & Limitations

    The same learned action can be reused with new path preferences, including intermediate-basin rewards and geometric constraints, without retraining.

    Source: §1 · p. 2

    LEAP requires sufficient overlap between reference and observed path distributions; poor overlap makes the learned correction harder to estimate accurately.

    Source: §5 · p. 10

Abstract

from arXiv · show

Understanding and controlling how systems evolve is central to science and engineering, where many problems require reasoning about complete trajectories. Existing approaches typically learn local dynamics or generate trajectories, but objectives constructed from individual transitions can miss dependencies across time, while generative models do not necessarily provide a tractable objective for evaluating and optimizing arbitrary paths. We study $\textit{action learning}$, the problem of recovering a stochastic action from trajectory observations, yielding a common objective for reasoning about dynamical systems. We introduce $\textbf{LEAP}$ ($\textbf{L}$earned $\textbf{E}$nergy on $\textbf{A}$ction $\textbf{P}$aths), which anchors the unknown action to a tractable reference process and learns the remaining correction by comparing observed trajectories with reference trajectories. This construction yields a differentiable objective on path space that can capture both local dynamical structure and nonadditive dependencies across time. Once learned, the same action can be flexibly reused for path ranking and inference, mechanism steering, constrained optimization and sampling. Across stochastic physical systems, many-body orbital dynamics and molecular simulations, LEAP recovers underlying dynamics, captures dependencies across time and enables transfer across inference objectives and physical conditions. We show that $\textit{action learning}$ is a promising framework for turning trajectory observations into reusable objectives for scientific inference and control.

1 FROM TRAJECTORY SAMPLES TO PATH INFERENCE

Reasoning about transitions and constraints requires objectives over complete paths, but observed trajectories often arrive without a known action. LEAP learns a reusable stochastic action from trajectory samples, complementing methods that learn local dynamics or generate trajectories.

1 FROM TRAJECTORY SAMPLES TO PATH INFERENCE

  • Trajectory observations often lack the known action required for path inference, while local learned methods do not directly provide a reusable path objective.Objectives over complete paths support inference of representative transitions and preferences over trajectories.Source: §1 · pp. 1–2
  • Local predictors can misestimate underlying dynamics on globally selected trajectories, and accumulated transition errors can distort path optimization.Locally plausible segments can still lead to poor path-level minima.Source: §1 · p. 2
  • LEAP learns a stochastic action from complete trajectories by anchoring it to a tractable reference and learning a correction that can capture nonadditive dependencies.The learned action can be reused with new path preferences, including rewards for intermediate basins or geometric penalties, without retraining.Source: §1 · p. 2
  • The evaluations span controlled physical systems, a 31-body orbital benchmark, and alanine dipeptide molecular benchmarks, testing recovery, temporal dependence, and objective reuse.The double-well and orbital systems provide exact actions for evaluation.Source: §1 · p. 2

1.1 RELATED WORK

  • Known-model inference uses available dynamics or actions [Kappen, 2005], while learned evolution/generation methods generally lack reusable trajectory objectives [Chen et al., 2019; Song et al., 2021].LEAP instead targets the stochastic action on complete paths for optimization, control, and sampling.Source: §1.1 · pp. 2–3

2 LEARNING ACTION FUNCTIONALS FROM TRAJECTORY SAMPLES

LEAP learns a stochastic action from complete trajectory samples by anchoring it to a tractable reference process and learning a correction. The resulting path-space objective supports inference tasks beyond training.

2 LEARNING ACTION FUNCTIONALS FROM TRAJECTORY SAMPLES

  • LEAP learns a differentiable stochastic action from trajectory samples, yielding a reusable objective for evaluating and optimizing complete paths.The action does not require the underlying drift, potential, or simulator likelihood.Source: §2 · p. 3

2.1 TARGET ACTION AND TRAINING OBJECTIVE

  • LEAP combines a reference action with a correction learned by comparing observed and reference trajectories; at the population solution, it recovers the data action up to a context-dependent constant.The comparison objective identifies the trajectory-varying part of the correction; the remaining constant does not affect path rankings or within-context optimization.Source: §2.1 · p. 4

2.2 REFERENCE PROCESS

  • The reference sets the baseline path measure and inductive bias, requiring a tractable density, efficient sampling, and sufficient probability mass over data trajectories.A discrete-time Gaussian Markov reference has analytic transitions parameterized by drift and diffusion; with shared diffusion, drift differences give the correction a dynamical interpretation.Source: §2.2 · p. 4

2.3 LOCAL AND GLOBAL ACTION CORRECTIONS

  • LEAP-Local sums transition corrections when density ratios factorize; LEAP-Global models whole paths when shared trajectory information breaks that factorization, and LEAP-Hybrid combines both.When local and global terms are learned jointly, only their sum is identified, so interpreting the local term as a physical drift correction requires additional restrictions.Source: §2.3 · p. 5

2.4 DOWNSTREAM TRAJECTORY INFERENCE

  • The learned action can rank candidate paths, infer representative trajectories, steer mechanisms, impose constraints, and define trajectory distributions for sampling.Mechanism preferences and constraint penalties can be combined to weight plausible paths.Source: §2.4 · pp. 5–6

3 WHY LEAP RECOVERS AN ACTION

The paper gives guarantees for recovering an action from a comparison objective, identifying structured drift corrections, and bounding downstream inference error. The guarantees depend on support, model, and trajectory-selection conditions.

3 WHY LEAP RECOVERS AN ACTION

  • Under shared positive support, conditionally independent candidates, and a score class containing the target correction, every population minimizer of the comparison objective recovers the action.The supplied passage states this result but does not include the equation that completes its formal statement.Source: §3 · p. 6
  • For Euler–Maruyama dynamics with shared diffusion, time step, and initial distribution, the structured correction captures drift changes; recovering unconditioned drift from terminally selected paths requires variation across terminal conditions.The recovery applies at states and times visited by the process and uses a structured drift parameterization.Source: §3 · pp. 6–7
  • The same action error bounds both the inferred trajectory’s objective value and the change in its normalized trajectory density.The bounds continue to hold after adding the same preference or constraint functional to both objectives.Source: §3 · p. 7

4 EXPERIMENTS

Across the experiments, LEAP recovers physical dynamics from selected paths, captures trajectory-wide dependencies, and reuses a learned action across inference tasks and physical conditions. The results also test reference choice, accumulated features, and scaling.

4.1 RECOVERING DYNAMICS UNDER TERMINAL SELECTION

  • Drift MSE is (3.3±2.7)×10^-4 for LEAP, versus 0.8132±0.0030 for Kramers-Moyal and 0.7752±0.0043 for cubic regression on terminally selected double-well paths.Although cubic regression can represent the exact drift, it fits the selected process’s Doob h-transformed drift.Source: §4.1 · p. 7
  • With the learned potential fixed, LEAP matches analytic path costs at T = 4 and T = 16 after changing the endpoint and horizon, while a closer reference reduces correction error from 0.66 to 0.223.The T = 4 costs are 0.2533 ± 0.0001 versus 0.2532 analytically; at T = 16 they are 0.2506 ± 0.0006 versus 0.2500.Source: §4.1 · pp. 7–8

4.2 RECOVERING A NONADDITIVE TRAJECTORY ACTION

  • At 42 bodies and 256 transitions, LEAP-Global reaches NRMSE 0.111, while LEAP-Local remains at 1.0015 across paths with 64,764 coordinates.The global model accumulates the small hidden-mass effect across transitions as the system and trajectories grow.Source: §4.2 · p. 8
  • LEAP-Global reaches NRMSE 0.228 while local methods stay near constant prediction, and the exact orbital correction cannot be represented as a sum of separate transition corrections.The complete-path model follows the exact log cosh A dependence across held-out trajectories.Source: §4.2 · p. 8, Table 2 · p. 9, Figure 6 · p. 9
  • Adding accumulated features lowers binary NRE NRMSE from 1.00000 ± 3 × 10^-5 to 0.175 ± 0.011 and gives 0.228 ± 0.015 with categorical NRE.Both complete-path ratio objectives improve; the binary result also shows action learning is not tied exactly to the categorical objective.Source: §4.2 · p. 8, Table 3 · p. 9

4.3 REUSING LEARNED ACTIONS ACROSS CONDITIONS AND OBJECTIVES

  • For molecular path completion, LEAP has circular MSE 0.0931± 0.0016 with three interior observations, but conditional generators perform better with endpoints only.LEAP uses the same action for both observation patterns, while it is trained only on complete trajectories.Source: §4.3 · p. 9
  • The fixed learned correction remains useful after temperature changes, and the same 600 K action supports stronger basin preferences and angular constraints without retraining.The constrained path stays within |ϕ| ≤150° while reaching the same endpoint.Source: §4.3 · p. 9

5 CONCLUSION

The paper presents action learning as a reusable objective for trajectory data and reports its use across physical recovery, orbital dynamics, and molecular inference. It also identifies overlap, diffusion, and reduced-coordinate boundaries.

5 CONCLUSION

  • The same learned action supports drift recovery, nonadditive orbital corrections, molecular completion, temperature transfer, preferences, and constraints without retraining.The paper frames path-space objectives as a common interface between trajectory data and downstream inference, including reduced coordinates without closed-form dynamics.Source: §5 · pp. 9–10, Figure 7 · p. 10
  • LEAP requires sufficient overlap between reference and observed path distributions; poor overlap makes accurate correction estimation harder.The structured correction is exact for Gaussian Markov processes with shared diffusion, while state-dependent diffusion requires additional likelihood-ratio terms.Source: §5 · p. 10
  • In alanine experiments, LEAP learns an effective action in (ϕ, ψ), not the full molecular action, because reduced coordinates can introduce memory and state-dependent friction or diffusion..Source: §5 · p. 10

IMPACT STATEMENT

The impact statement describes learned trajectory objectives as potentially useful for analyzing and steering dynamical systems, while cautioning that deployment needs validation and attention to data.

IMPACT STATEMENT

  • Learned trajectory objectives may help analyze and steer dynamical systems in science, robotics, planning, and other sequential domains.The passage also notes that such objectives could steer systems toward harmful outcomes, so deployment requires domain validation and careful attention to data.Source: IMPACT STATEMENT · p. 10

APPENDIX

  • The appendix outlines additional materials on theory, shared evaluation procedures, and experiments spanning double-well, orbital, and molecular settings.Its sections cover action recovery and inference theory, reference choice and training details, orbital comparisons and scaling, and molecular validation.Source: APPENDIX · p. 14

A THEORY AND DERIVATIONS

  • The theory and derivations section includes population action recovery, structured Gaussian corrections, drift interpretation, matched trajectory selection, nonadditive actions, and downstream stability.The notation fixes a context and trajectory space and defines data and reference densities on the same base measure.Source: A THEORY AND DERIVATIONS · p. 14

A.1 ACTION RECOVERY FROM TRAJECTORY COMPARISONS

The appendix establishes when trajectory comparisons recover an action, when structured corrections identify drift, and how conditioning or latent regimes affect path objectives. It also bounds the effect of reusing an imperfect learned action.

A.1 ACTION RECOVERY FROM TRAJECTORY COMPARISONS

  • At the population optimum, trajectory comparisons recover the data action up to a context-dependent constant for every K ≥ 1, under the stated support, independence, and model-class conditions.The constant does not affect action differences between paths in the same context.Source: §A.1 · p. 14

A.2 STRUCTURED CORRECTION AND DRIFT IDENTIFICATION

  • Under shared-diffusion assumptions, matching unconditioned trajectory distributions identifies the drift correction on visited state-time pairs.For structured LEAP, this yields uθ*=u* on the support of observed states and times.Source: §A.2 · p. 15

A.3 MATCHED CONDITIONING AND SELECTED TRAJECTORIES

  • Matched conditioning preserves action differences between paths with the same condition, but selected-path local regression estimates a conditioned drift rather than necessarily the unconditioned drift.With both endpoints fixed, constant-drift likelihood ratios are identical across connecting paths, so those trajectories alone cannot distinguish drift values.Source: §A.3 · p. 16

A.4 NONADDITIVE ACTION FROM A PERSISTENT LATENT REGIME

  • Marginalizing a latent regime held fixed over a trajectory can induce nonadditive actions that cannot be represented by additive local transition costs.The result applies where the relevant trajectory terms can vary independently.Source: §A.4 · pp. 16–17

A.5 DOWNSTREAM STABILITY UNDER ACTION REUSE

  • If the learned action has error at most ε and an inferred path is δ-suboptimal for its objective, its true-objective suboptimality is at most 2ε + δ.The bound continues to hold after adding the same preferences or constraints to both objectives.Source: §A.5 · p. 17

B EXPERIMENTAL DETAILS

This section specifies the experimental metrics, inference procedure, and training resources used across the paper. It also describes uncertainty reporting and the principal experiment settings.

B.1 METRICS AND UNCERTAINTY

  • Because action is identified up to a context-dependent constant, action correction error compares within-context correction differences, scoring exact differences 0 and constant predictions 1; trajectory discrimination uses within-context AUROC, and numerical action differences use NRMSEΔG.For the double well, b*(x) = x −x3; this cost is used only for evaluation, while each method infers with its own objective.Source: §B.1 · p. 18

B EXPERIMENTAL DETAILS

  • Experiments use at least three independent seeds, report table entries as sample mean ± sample standard deviation, and run training jobs on one NVIDIA GPU with up to 8 CPU cores and 64 GB of memory.The reference study runs on CPUs.Source: §B.1 · p. 18, B EXPERIMENTAL DETAILS · p. 18

B.2 PATH INFERENCE

  • Path optimization updates free trajectory coordinates using the objective, projects onto hard constraints when required, and can redistribute knots before the next optimizer step.The algorithm returns the optimized path.Source: §B.2 · p. 18

B.3 TRAINING AND MODEL SELECTION

  • The principal experiments use the settings summarized in Table 4, with additional double-well architecture and optimizer details and validation stopping for orbital experiments.The double-well potential network has three 512-wide hidden layers, GELU activations, and LayerNorm; orbital training uses validation stopping with patience 20.Source: §B.3 · p. 19, Table 4 · p. 19

C DOUBLE WELL EXPERIMENTS

The double-well experiments use terminally selected trajectories to evaluate drift recovery and path inference under changed boundary conditions. The long-horizon analysis relates the optimal action to the potential barrier.

C.1 TRAJECTORY CONSTRUCTION AND BASELINES

  • Training uses 5,000 trajectories per seed selected by xT ∈ [0.9, 1.1], with Brownian motion at the same diffusion and terminal condition as the reference.LEAP parameterizes drift through a scalar potential; local baselines use the same selected trajectories, whose terminal selection changes their regression target.Source: §C.1 · p. 19

C.2 PATH INFERENCE UNDER NEW BOUNDARY CONDITIONS

  • Path inference keeps the learned potential fixed, changes the endpoint to xT = 0, and evaluates horizons T ∈ {1, 4, 8, 16}.The free path coordinates are optimized from four fixed linear and waiting initializations, then evaluated with the ground-truth cost.Source: §C.2 · pp. 19–20

C.3 LONG-HORIZON BARRIER SCALE

  • As the horizon grows, the optimal action approaches the barrier value, and the reverse-gradient path approaches this bound.At finite horizons, optimized discrete paths are evaluated under the true action.Source: §C.3 · p. 20

D REFERENCE CHOICE AND ROBUSTNESS

The experiments test reference sensitivity, robustness, and whether action recovery requires modeling dependencies across time. In the orbital setting, a trajectory-wide correction substantially outperforms transition-level estimators, while structured features allow different ratio-estimation objectives to recover the correction.

D REFERENCE CHOICE AND ROBUSTNESS

  • Adding physical drift to the reference decreases LEAP correction NRMSE from 0.66 at aref = 0 to 0.223 at aref = 0.5 and 0.079 at aref = 1.Across three temporal grids, estimators reach similar correction error; increasing hidden integration substeps from 2 to 8 changes LEAP AUROC from 0.9171 to 0.9216.Source: D REFERENCE CHOICE AND ROBUSTNESS · p. 20, Table 7 · p. 21, Table 8 · p. 21
  • Increasing training trajectories from 32 to 256 decreases LEAP NRMSE from 0.87 to 0.66, while increasing K from 1 to 4 gives the largest reference-budget improvement.Under the default setting, lowest-predicted-action selection yields regret of 0.0075 nats for LEAP versus 0.0230 nats for Gaussian MLE.Source: D REFERENCE CHOICE AND ROBUSTNESS · pp. 20–21

E ORBITAL EXPERIMENTS

  • The orbital benchmark uses 31 bodies, 186-dimensional states, and 32 transitions, with held-out initial contexts evaluated after chronological splitting.Training resimulates data and reference trajectories from 254 training initial states each epoch.Source: E ORBITAL EXPERIMENTS · p. 21

E.1 EXACT TRAJECTORY ACTION

  • A solar-mass regime is sampled once per trajectory and remains fixed across transitions; marginalizing it combines likelihoods over whole trajectories into a nonadditive correction.The reference uses nominal mass, while the data use perturbed mass regimes.Source: §E.1 · p. 21

E.2 PATH STRUCTURE AND ESTIMATOR CONTROLS

  • Transition-level models remain at NRMSE 1, while LEAP-Global reaches 0.228 on the 31-body orbital benchmark.The exact correction mixes over a latent regime shared across the trajectory, highlighting the importance of dependencies across time.Source: §E.2 · p. 22, Table 9 · p. 22
  • Binary and categorical ratio estimation reach NRMSE 0.175 and 0.228 with structured features, while binary estimation with a generic path encoder stays near NRMSE 1.Both objectives can recover the correction when paired with the same Transformer and accumulated physical features.Source: §E.2 · p. 22

E.3 COMPARISON AND DATA BUDGETS

Orbital experiments show that more reference paths and training trajectories improve correction estimates, and that LEAP-Global maintains low error in larger systems and longer paths. A trajectory-diffusion baseline, by contrast, does not track exact orbital corrections.

E.3 COMPARISON AND DATA BUDGETS

  • Increasing K from 4 to 16 decreases LEAP-Global NRMSE from 0.2208 to 0.147, while increasing training paths from 4,064 to 8,128 decreases it from 0.149 to 0.110.Extending the K = 4 model to 275 epochs leaves NRMSE at 0.2207.Source: §E.3 · p. 22

E.4 LONGER PATHS AND LARGER SYSTEMS

  • At 42 bodies and 256 transitions, LEAP-Global reaches NRMSE 0.111, while LEAP-Local remains at 1.0015.Each path contains 64,764 coordinates; LEAP-Global maintains low correction error as state dimension and trajectory length increase.Source: §E.4 · p. 23, Table 11 · p. 23

E.5 TRAJECTORY DIFFUSION BASELINE

  • GeoTDM gives chance AUROC (0.4999), and its variational estimate has Spearman correlation 0.0066 with the exact correction.Trajectories with larger exact corrections do not systematically receive larger GeoTDM estimates.Source: §E.5 · p. 23, Table 12 · p. 23

F.1 MOLECULAR DATA AND MODELS

The molecular experiments evaluate trajectory completion, temperature transfer, and reuse of a learned action for preferences and constraints. LEAP performs well among action-based methods, transfers across temperatures with the target reference process, and supports changing inference objectives while the action stays fixed.

F.1 MOLECULAR DATA AND MODELS

  • Molecular paths use backbone angles sampled every 0.1 ps, with at most 201 frames; reference paths share the reactive event and use wrapped Gaussian transition densities.The data comprise train/validation/test trajectories at 500 K and 600 K, with simulation replicas assigned to disjoint splits.Source: §F.1 · p. 23

F.2 COMPLETION CONTROLS

  • With endpoints and three interior observations, LEAP achieves circular MSE 0.0931 versus 0.119 for conditional DDPM and 0.125 for conditional flow matching.With endpoints only, LEAP has the lowest error among action-based methods (0.6225 versus 0.6609 for velocity regression and 0.6682 for binary NRE), while conditional generators achieve lower reconstruction error.Source: §F.2 · p. 24, Table 13 · p. 24

F.3 TEMPERATURE TRANSFER

  • Training at 600 K gives AUROC 0.9579 at 500 K, while training at 500 K gives 0.9719 at 600 K.The correction transfers across temperatures when combined with the target reference process.Source: §F.3 · p. 24, Table 14 · p. 24

F.4 PREFERENCE AND CONSTRAINT REUSE

  • Increasing preference weight λ from 0 to 20 raises mean preference from 0.0658 to 0.2517 while transition cost changes from 0.9389 to 0.9809.At λ = 50, preference reaches 0.3641 with transition cost 1.3774; reconstructed trajectories also preserve molecular geometry at both temperatures.Source: §F.4 · p. 25, Table 15 · p. 25

Related papers

Loading 2610.31585v1…