Action Learning for Scientific Simulations

Sarah Liaw; Yuanqi Du; Yilun Du

cs.LGOct 22, 2026

Action Learning for Scientific Simulations

Authors: Sarah Liaw, Yuanqi Du, Yilun Du

Organizations: Kempner Institute · Harvard University · Microsoft Research New England

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: Local dynamics learning versus action learning. Local models learn short-range transition laws that induce a path likelihood, whereas LEAP learns an action from complete trajectories while retaining structured local dynamics when appropriate.

Figure 2: LEAP. Given data trajectories and a tractable reference process, LEAP learns a centered path-space correction and combines it with the reference action to obtain $S_{\theta}(\Phi;c)$ .

Figure 3: Downstream trajectory inference with the learned action. Once $S_{\theta}(\Phi;c)$ is learned, the same action supports representative path inference, mechanism steering (tilting), constrained inference and sampling.

Panel of Figure 6 — Figure 6: Reuse of a fixed molecular action. All panels use the same $600$ K LEAP action. (a) Completion from endpoints and three interior observations. (b) Increasing the basin preference raises $M$ from $0.23$ to $0.81$ for the shown path. (c) Adding $|\phi|\leq 150^{\circ}$ changes the route while preserving the terminal condition. Diamonds and squares are endpoints, dots are saved frames and the background is empirical free energy.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

F(\Phi;c)=S_{\theta}(\Phi;c)-\sum_{j}\gamma_{j}M_{j}(\Phi)+\sum_{\ell}\lambda_{\ell}C_{\ell}(\Phi)

Appendix

Algorithm 2 Path optimization with a learned LEAP action

	Double well	Orbital	Alanine
Training paths	$5{,}000$	$1{,}016$	$130$ / $463$
Reference paths $K$	$8$	$4$	$8$
Training budget	$500$ epochs	$100$ epochs	$1{,}200$ updates
Batch size	$64$	$8$	$32$
Learning rate	$10^{-3}$	$3\times 10^{-4}$	$3\times 10^{-4}$
Training seeds	$3$	$3$	$3$

Appendix

Table 4: Training settings for the principal experiments. Molecular counts correspond to the

500

K and

600

K training sets. Orbital model inputs are described in § E.2 .

Method	Definition
Kramers-Moyal	Binned mean local increment
Cubic velocity regression	Cubic least-squares fit to local increments
Zero field	$b(x)=0$
Geometric path	Fixed waiting and linear crossing path
Analytic drift	Exact drift and action
LEAP	Complete path comparison objective

Appendix

Table 5: Double well methods and controls.

	Drift MSE $\downarrow$	True path cost $\downarrow$
Method		$T=4$	$T=16$
Kramers-Moyal	$0.8132{\scriptscriptstyle\pm 0.0030}$	—	—
Cubic velocity reg.	$0.7752{\scriptscriptstyle\pm 0.0043}$	$0.3189{\scriptscriptstyle\pm 0.0008}$	$0.6126{\scriptscriptstyle\pm 0.0079}$
Zero field	$0.0759$	$0.2637$	$0.4454$
Geometric path	—	$0.2637$	$0.2637$
Analytic drift	$0$	$0.2532$	$0.2500$

Appendix

Table 6: Physical recovery and path inference. Drift error and path cost under new endpoints and horizons.

Reference diffusion
Setting	LEAP	Transition NCE	Gaussian MLE
$D_{\mathrm{ref}}=0.06$	$0.31{\scriptscriptstyle\pm 0.13}$	$0.24{\scriptscriptstyle\pm 0.13}$	$0.347{\scriptscriptstyle\pm 0.014}$
$D_{\mathrm{ref}}=0.12$	$0.58{\scriptscriptstyle\pm 0.16}$	$0.64{\scriptscriptstyle\pm 0.12}$	$0.644{\scriptscriptstyle\pm 0.084}$
$D_{\mathrm{ref}}=0.20$	$0.66{\scriptscriptstyle\pm 0.14}$	$0.71{\scriptscriptstyle\pm 0.11}$	$0.729{\scriptscriptstyle\pm 0.084}$
$D_{\mathrm{ref}}=0.32$	$0.795{\scriptscriptstyle\pm 0.089}$	$0.77{\scriptscriptstyle\pm 0.10}$	$0.855{\scriptscriptstyle\pm 0.062}$
$D_{\mathrm{ref}}=0.48$	$0.939{\scriptscriptstyle\pm 0.040}$	$0.837{\scriptscriptstyle\pm 0.080}$	$0.931{\scriptscriptstyle\pm 0.034}$

Appendix

Table 7: Reference choice. Correction NRMSE

\downarrow

. The diffusion sweep uses

b_{\mathrm{ref}}=0

; the drift sweep uses

D_{\mathrm{ref}}=0.20

. Contrastive models use

K=4

Setting	LEAP	Transition NCE	Gaussian MLE
$\Delta t=0.02$ , $N=48$	$0.828{\scriptscriptstyle\pm 0.096}$	$0.769{\scriptscriptstyle\pm 0.081}$	$0.730{\scriptscriptstyle\pm 0.033}$
$\Delta t=0.04$ , $N=24$	$0.66{\scriptscriptstyle\pm 0.14}$	$0.71{\scriptscriptstyle\pm 0.11}$	$0.729{\scriptscriptstyle\pm 0.084}$
$\Delta t=0.08$ , $N=12$	$0.67{\scriptscriptstyle\pm 0.15}$	$0.671{\scriptscriptstyle\pm 0.070}$	$0.66{\scriptscriptstyle\pm 0.13}$

Appendix

Table 8: Temporal resolution. Correction NRMSE

\downarrow

, mean

\pm

sample SD at fixed

T=0.96

Model	Input	$\mathrm{NRMSE}_{\Delta G}\downarrow$
LEAP-Local	One transition	$1.0013{\scriptscriptstyle\pm 0.0015}$
Transition NCE	One transition	$1.0008{\scriptscriptstyle\pm 0.0013}$
LEAP-Global	Complete path + accumulated features	$\mathbf{0.228}{\scriptscriptstyle\mathbf{\pm 0.011}}$
LEAP-Hybrid	Local + complete path	$0.232{\scriptscriptstyle\pm 0.014}$
Gaussian MLE	State + time	$11.92{\scriptscriptstyle\pm 0.66}$
History Gaussian MLE	Causal history	$10.162{\scriptscriptstyle\pm 0.093}$

Appendix

Table 9: Orbital controls. Correction NRMSE

\downarrow

on the

31

-body benchmark; a constant prediction gives

1

. Bold is the lowest mean within each control block.

Training data
Setting	Training	$\mathrm{NRMSE}_{\Delta G}\downarrow$
$K=4$	validation stop	$0.2208{\scriptscriptstyle\pm 0.0056}$
$K=4$	$275$ epochs	$0.2207{\scriptscriptstyle\pm 0.019}$
$K=16$	validation stop	$0.147{\scriptscriptstyle\pm 0.017}$
$K=32$	validation stop	$0.138{\scriptscriptstyle\pm 0.032}$
Training paths	LEAP-Global	Binary NRE

Appendix

Table 10: Comparison and data budgets. Correction NRMSE

\downarrow

. Increasing

K

changes the number of reference paths per observed trajectory; the lower panel increases the number of training paths. Separate runs from Table 9 .

31 bodies
Transitions	Path coords.	LEAP-Local	LEAP-Global
32	$6{,}138$	$0.9994{\scriptscriptstyle\pm 1.2\times 10^{-3}}$	$0.2584{\scriptscriptstyle\pm 0.0097}$
64	$12{,}090$	$0.9985{\scriptscriptstyle\pm 7\times 10^{-4}}$	$0.127{\scriptscriptstyle\pm 0.015}$
128	$23{,}994$	$1.0052{\scriptscriptstyle\pm 0.0066}$	$0.113{\scriptscriptstyle\pm 0.027}$
256	$47{,}802$	$0.9965{\scriptscriptstyle\pm 7\times 10^{-4}}$	$0.117{\scriptscriptstyle\pm 0.055}$
42 bodies

Appendix

Table 11: Scaling with system size and path length. Correction NRMSE

\downarrow

. Each configuration uses its exact mixture correction and

K=4

. Separate runs from Table 9 .

Objective difference	Error $\downarrow$	AUROC $\uparrow$	Spearman $\uparrow$
Denoising	$1.0000{\scriptscriptstyle\pm 2.3\times 10^{-6}}$	$0.4999{\scriptscriptstyle\pm 2\times 10^{-4}}$	—
Variational	$0.9883{\scriptscriptstyle\pm 1\times 10^{-4}}$	$0.4999{\scriptscriptstyle\pm 2\times 10^{-4}}$	$0.0066{\scriptscriptstyle\pm 0.0055}$

Appendix

Table 12: Trajectory diffusion baseline. GeoTDM objective differences on the orbital benchmark with

31

bodies. The denoising objective is evaluated using correction NRMSE, while the variational objective is evaluated using affine relative error.

Method	Sparse completion	Endpoints only
Reference action	$0.1031{\scriptscriptstyle\pm 0.0016}$	$0.8183{\scriptscriptstyle\pm 0.0042}$
Untrained action	$0.0947{\scriptscriptstyle\pm 0.0000}$	$0.7953{\scriptscriptstyle\pm 0.0027}$
Velocity regression	$0.0900{\scriptscriptstyle\pm 0.0003}$	$0.6609{\scriptscriptstyle\pm 0.0032}$
Binary NRE	$0.0897{\scriptscriptstyle\pm 0.0002}$	$0.6682{\scriptscriptstyle\pm 0.0118}$
Conditional DDPM	$0.119{\scriptscriptstyle\pm 0.012}$	$0.409{\scriptscriptstyle\pm 0.034}$
Conditional flow matching	$0.125{\scriptscriptstyle\pm 0.011}$	$0.350{\scriptscriptstyle\pm 0.012}$

Appendix

Table 13: Molecular completion at

600

K. Circular MSE

\downarrow

on unobserved frames for the same observation patterns and durations.

Train	Test	LEAP	Binary NRE	Velocity	Untrained
$500$	$500$	$0.9403{\scriptscriptstyle\pm 0.0099}$	$0.9337{\scriptscriptstyle\pm 0.0030}$	$0.9427{\scriptscriptstyle\pm 0.0033}$	$0.8385{\scriptscriptstyle\pm 0.0050}$
$600$	$500$	$\mathbf{0.9579}{\scriptscriptstyle\mathbf{\pm 0.0057}}$	$0.9489{\scriptscriptstyle\pm 0.0051}$	$0.9437{\scriptscriptstyle\pm 0.0022}$	$0.8385{\scriptscriptstyle\pm 0.0050}$
$500$	$600$	$\mathbf{0.9719}{\scriptscriptstyle\mathbf{\pm 0.0068}}$	$0.9386{\scriptscriptstyle\pm 0.0026}$	$0.9533{\scriptscriptstyle\pm 0.0025}$	$0.8847{\scriptscriptstyle\pm 0.0028}$
$600$	$600$	$0.9754{\scriptscriptstyle\pm 0.0012}$	$0.9692{\scriptscriptstyle\pm 0.0007}$	$0.9559{\scriptscriptstyle\pm 0.0029}$	$0.8847{\scriptscriptstyle\pm 0.0028}$

Appendix

Table 14: Temperature transfer. Within-context AUROC

\uparrow

. Bold marks LEAP evaluated at a temperature different from its training temperature.

	$\lambda=0$	$\lambda=5$	$\lambda=20$	$\lambda=50$
Preference $M$ $\uparrow$	$0.0658{\scriptscriptstyle\pm 0.0059}$	$0.1095{\scriptscriptstyle\pm 0.0088}$	$0.2517{\scriptscriptstyle\pm 0.0045}$	$0.3641{\scriptscriptstyle\pm 0.0116}$
Transition cost $\downarrow$	$0.9389{\scriptscriptstyle\pm 0.0545}$	$0.9159{\scriptscriptstyle\pm 0.0228}$	$0.9809{\scriptscriptstyle\pm 0.0231}$	$1.3774{\scriptscriptstyle\pm 0.0252}$

Appendix

Table 15: Preference reuse with a fixed LEAP action. Increasing

\lambda

raises the basin preference while changing the unguided projected transition cost in nats per active step.

Learning Transferable Policies from Action-free Time Series Through Dynamical Embeddings
Oct 2, 2026Niklas Emonds, Georgia KoppeTransfer Learning in RL Reinforcement Learning
Hector Institute for AI in Psychiatry (HITKIP) & Department of Psychiatry and Psychotherapy, Central Institute of Mental Health (CIMH), Medical Faculty Mannheim, Heidelberg University · Interdisciplinary Center for Scientific Computing (IWR), Faculty of Mathematics and Computer Science, Heidelberg University · Hertie Institute for AI in Brain Health, University of Tübingen, Germany
LeapBot-JAC: Joint Action Chains for World Action Models
Oct 9, 2026Chenhao Zhu, Baibing Liu, Pei Liu +5Robot Policy Learning Action Chunking
Leapmotor · The Hong Kong University of Science and Technology (Guangzhou) · The Hong Kong University of Science and Technology
Langevin-Informed Transfer Learning: Replacing Target Samples by Black-Box Feedback
Oct 1, 2026Vladimir R. Kostic, Karim Lounici, Hélène Halconruy +3Latent Dynamics Modeling Transfer Learning
CSML, Istituto Italiano di Tecnologia · University of Novi Sad · CMAP-Ecole Polytechnique +5

Oct 2, 2026J/K move · Enter open · S save

Learning Transferable Policies from Action-free Time Series Through Dynamical Embeddings

Niklas Emonds, Georgia Koppe

Hector Institute for AI in Psychiatry (HITKIP) & Department of Psychiatry and Psychotherapy, Central Institute of Mental Health (CIMH), Medical Faculty Mannheim, Heidelberg University · Interdisciplinary Center for Scientific Computing (IWR), Faculty of Mathematics and Computer Science, Heidelberg University · Hertie Institute for AI in Brain Health, University of Tübingen, Germany

Transfer Learning in RL Reinforcement Learning Model-Based RL Representation Learning for RL Offline RL Hierarchical RL

Open paper AskPDF arXiv

Learning control from action-free recordings is challenging because intervention effects are unobserved and policies may exploit errors in reconstructed dynamics. We present a hierarchical model-based reinforcement learning framework that uses shared structure across related systems to learn system-specific control policies from action-free recordings. A hierarchical dynamical system reconstruction model captures shared dynamics and individual variation through low-dimensional embeddings. These embeddings are then reused to parameterize shared policy and value networks, linking differences in reconstructed dynamics to differences in control. Policies are trained entirely via simulation under an explicit intervention model with additive latent perturbations. Piecewise-linear recurrent neural networks enable mechanistic analyses of the controlled dynamics, while decoder-based constraints make the immediate effects of interventions interpretable in observation space and permit interventions on one modality while protecting another from direct manipulation. On Lorenz-63 and double-pendulum systems, hierarchical policies improve transfer over independently trained policies. On Lorenz-63, they also achieve a higher mean reward than repeated planning with the same reconstructed models, perform comparably to methods trained with controlled interactions, and generalize to systems absent from policy training after embedding inference alone. Applications to neural-behavioral recordings demonstrate suppression of predicted movement under constrained neural perturbations. Together, these findings show how shared dynamical representations support transferable control and mechanistic hypothesis generation from action-free recordings.

Abstract

Explore similar work

Learning Transferable Policies from Action-free Time Series Through Dynamical Embeddings

LeapBot-JAC: Joint Action Chains for World Action Models

Langevin-Informed Transfer Learning: Replacing Target Samples by Black-Box Feedback