Learning Transferable Policies from Action-free Time Series Through Dynamical Embeddings
Authors: Niklas Emonds, Georgia Koppe
Organizations: 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
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.
Figures & tables
Figure 1: Framework overview. a : A hierarchical DSR model is trained on multiple related dynamical systems (“subjects”) and used as world model for MBRL . Subject-specific parameters φ(j) are obtained by projecting shared parameters Pφ onto low-dimensional subject embeddings S(j) . These embeddings are reused to condition the policy πψ and value networks Qθ by modulating group-level parameters Pψ and Pθ . b : For unseen subjects, learning only a new embedding via DSR fine-tuning enables insertion into the control model. c : The high-dimensional dynamics are controlled only in the directions that are decoded ( z~t ). To underactuate the problem and for application-specific constraints, actions are constrained to be orthogonal to the target modalities and target dimensions. d : In the Lorenz-63 benchmark, the actions only affect the y− and z− directions.
Figure 2: Controlling the Lorenz-63 benchmark. a : Reward histograms for randomly initialized trajectories across all subjects for the simulated system (left) and the true system (right) using hierarchically (top) and independently trained models (bottom). Black line denotes uncontrolled baseline. b : Representative Lorenz–63 trajectories with reward =−0.020 generated by the world model under SAC control during uncontrolled and controlled phases (gray shading), with the value function and occupied LSRs shown below. While x2 and x3 are directly controlled and rapidly approach the target (dashed red line), x1 converges to a nearby FP . c : Corresponding trajectories of the true system, demonstrating sim-to-real transfer, with a reward of −0.03 . d, e : For the same subject, the most frequently occupied subregions (indices 3 and 5) are those admitting non-virtual FPs (red crosses) close to the target. f : During extended control, the trajectory approaches the target while action-dependent FPs are progressively shifted toward it.
Lorenz-63
Double-pendulum
World Model
True Environment
World Model
True Environment
ALRNN -SAC, Hierarchical
−0.047±0.006
−0.040±0.004
−47.93±10.57
−9.74±2.42
ALRNN -SAC, Independent
−0.0232±4⋅10−4
−0.171±0.017
−51.00±18.47
−16.82±7.78
ALRNN -CEM-MPC, Independent
- 0.02090±4⋅10−5
−0.04424±8⋅10−5
−15.35±0.71
−6.34±0.55
Online-TD-MPC, Independent
−40±2
−0.0328±9⋅10−4
−3547±323
−27.94±7.21
Online-DreamerV3, Independent
−0.034±0.002
−0.047±0.001
−6.40±0.78
−5.87±0.30
Table 1: Control performance . The reward is reported with mean and standard deviation across seeds.
Figure 3: Reconstruction and stabilization of the double-pendulum. a : Colors indicate mass ratios m1/m2 and length ratios l1/l2 , respectively. Analogous to the Lorenz-63 System, subject vectors cluster according to system parameters. b : Example trajectory of the true double-pendulum under control, with the target position indicated by the dashed red line. The 6 variables are sine and cosine of the joint angles as well as the angular velocities. Despite its inherent instability the controlled system converges to the target.
Figure 4: Reconstruction and model-based control of the LINK data. a : The first principal component of the subject vectors shows a linear relationship with recording days, indicating that the learned subject representations capture structure across sessions. b : Finger positions of the index and MRS finger groups, neural activity, and control energy for one representative session. Each new goal position, presented to the monkey and provided to the model as external input, triggers a movement in the recorded (red) and simulated (blue) trajectories. Under the learned policy (orange), the simulated finger positions stay at rest. Neural activity follows a quasi-periodic dynamics in the uncontrolled model and changes under control. Control energy is the squared norm of the decoded action. Controlled trajectories are reinitialized from the recorded state every 200 time steps, which produces the visible jumps. c : Action magnitude over time for 0 to 12 excluded channels (color, top) and median reward as a function of the number of excluded channels (bottom). Actions are forced to zero on randomly selected channels. Action magnitude first increases to compensate and is then suppressed entirely once 8 channels are excluded, causing a drop in reward.
Figure 5: Behavioral control through neural stimulation. a, left: Predicted trajectories without control. The model receives the goal position cue (black dashed) and go cue (cyan) as external inputs and generates trajectories (solid) that reproduce the empirical test data (dashed). Hand (green) and eye (blue) positions are shown separately. In successful trials, the hand follows the instructed target, whereas eye position is not consistently target-aligned in either the model or the data. center: Two-dimensional trajectories without (top) and with (bottom) control. The policy suppresses hand movement and, to a lesser extent, eye movement despite external inputs driving both toward the target. Color gradients indicate time. right: Neural stimulation remains sparse and selective, with no statistically significant changes in mean firing rate ( t=−1.33 , p=0.18 ) or Fano factor ( t=1.04 , p=0.30 ). b: Value function, control energy, and latent states for a representative trial. Actions become effective once the goal-position cue is presented and steer the system into a high-value region of latent space associated with trajectories that remain insensitive to the subsequent go cue.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Subject vectors cluster according to Lorenz-63 parameters. PCA projection of the subject vectors. Colors indicate constant values of σ , while vertical and horizontal separation reflect β and ρ . The subject vectors condition SAC training and convey information about subject-specific dynamics.
Figure 7: Limited controllability under single dimension actuation. Example trajectory of the true Lorenz-63 system with actions applied exclusively to the third state dimension. The dashed line indicates the target state.
Figure 8: Mechanisms underlying Lorenz-63 control. a : Number of LSRs required to cover 90% of time steps for uncontrolled (blue) and controlled (red) trajectories of each subject. Control strongly reduces region occupancy. Subject 54 (cf. figure 2 ) is highlighted. b : Minimal distance to the target of the non-virtual FPs computed from sampled actions in each occupied LSR (cf. figure 2(e) ). c : Eigenvalues of the Jacobian at the FPs , which depends only on the LSR , indicate stability along most directions. The unit circle is shown in gray.
Figure 9: Robust control of heterogeneous nonlinear dynamics. Trajectories of all 64 Lorenz-63 systems under the learned policy. Dashed lines indicate the target state, and gray shading marks the action phase.
Figure 10: Generalization through fine-tuned subject vectors. On unseen subjects, fine-tuned subject vectors achieved a reward of =−0.036±0.002 , compared with −0.046±0.005 for the unchanged vectors of the nearest training subjects. a : Fine-tuned subject vectors inserted into the fixed SAC model control the true, unseen systems without further optimization (cf. figure 2(c) ). b : PCA projection of the subject vectors of the fine-tuned (colored) and original Lorenz-63 systems. Shifted parameter regimes produce corresponding shifts in the PCA embedding (cf. figure 6 ).
Figure 11: Control mechanisms in multimodal neural data. a : Occupancy of LSRs during control. b: Minimal distance of non-virtual FPs to the target in each occupied LSR . The regions occupied during control are those that admit a non-virtual FP close to the target, indicating that the policy steers the system into regions of latent space whose linear dynamics favor reaching the goal (cf. figure 2(d) - 2(e) ).
Figure 12: Trajectories of all LINK data subjects.
Figure 13: Reconstruction of all Lorenz-63 subjects. Ground-truth test trajectories (blue) and reconstructions (orange).
Figure 14: Reconstruction of 16 double-pendulum subjects. a: Ground truth (blue) and reconstructed (orange) trajectories of the θ1,θ2,ω1 , and ω2 agree closely up to the training sequence length of 30 time steps (red dashed line). b: Energy remains largely constant at first and fluctuates increasingly over time.
Figure 15: Reconstruction of neural spike rates. Empirical (left) and reconstructed (right) spike rates for a representative trial. Across all test trials, the world model achieves a variance-normalized mean squared error of 0.17±0.06 .
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 action learning, the problem of recovering a stochastic action from trajectory observations, yielding a common objective for reasoning about dynamical systems. We introduce LEAP (Learned Energy on Action Paths), 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 action learning is a promising framework for turning trajectory observations into reusable objectives for scientific inference and control.
Sarah Liaw, Yuanqi Du, Yilun Du
Kempner Institute · Harvard University · Microsoft Research New England
Training reinforcement learning (RL) policies from scratch is costly: it requires careful reward and environment design, extensive tuning, and substantial computation. Yet many control problems already have a functional but suboptimal policy available as a baseline. This paper proposes a method for embedding such a baseline into the RL training process, simultaneously improving training efficiency relative to from-scratch methods and producing a learning policy that outperforms the baseline. At each step, the method arbitrates between the baseline policy and a trainable learning policy, initially relying strongly on the baseline policy and then progressively transferring agency to the learning policy. By the end of training, the learning policy is a standalone neural network that operates without baseline policy support. The paper formalizes what it means for the baseline policy to be functional: under this policy, the agent reaches a goal set and remains there with high probability. The proposed arbitration mechanism is designed to exploit this property during training, yielding high goal-reaching rates right from the beginning of training. A theoretical analysis provides a formal interpretation of this behavior under stated assumptions and extends it to the final baseline-free regime, where explicit lower bounds are derived for the goal-reaching probability of the standalone learning policy. Empirical results on continuous-control benchmarks show that the proposed method achieves returns that match or exceed those of competitive approaches, while maintaining the highest goal-reaching rates throughout training among the compared methods -- including in the final stage, where the learning policy operates without any baseline support.
Anton Bolychev, Georgiy Malaniya, Sinan Ibrahim +1
Center for Engineering Systems and Sciences · Central University; Center for Engineering Systems and Sciences; Sirius University of Science and Technology
Animal brains flexibly and efficiently achieve many behavioral tasks with a single neural network. A core goal in modern neuroscience is to map the mechanisms of the brain's flexibility onto the dynamics underlying neural populations. However, identifying task-specific dynamical rules from limited, noisy, and high-dimensional experimental neural recordings remains a major challenge, as experimental data often provide only partial access to brain states and dynamical mechanisms. While recurrent neural networks (RNNs) directly constrained neural data have been effective in inferring underlying dynamical mechanisms, they are typically limited to single-task domains and struggle to generalize across behavioral conditions. Here, we introduce JEDI, a hierarchical model that captures neural dynamics across tasks and contexts by learning a shared embedding space over RNN weights. This model recapitulates individual samples of neural dynamics while scaling to arbitrarily large and complex datasets, uncovering shared structure across conditions in a single, unified model. Using simulated RNN datasets, we demonstrate that JEDI accurately learns robust, generalizable, condition-specific embeddings. By reverse-engineering the weights learned by JEDI, we show that it recovers ground truth fixed point structures and unveils key features of the underlying neural dynamics in the eigenspectra. Finally, we apply JEDI to motor cortex recordings during monkey reaching to extract mechanistic insight into the neural dynamics of motor control. Our work shows that joint learning of contextual embeddings and recurrent weights provides scalable and generalizable inference of brain dynamics from recordings alone.
Anirudh Jamkhandi, Ali Korojy, Olivier Codol +2
Mila - Quebec AI Institute · Université de Montréal