Learning Dynamics of RNNs in Closed-Loop Environments

TL;DR

The paper develops a mathematical theory for the learning dynamics of linear RNNs operating in closed-loop environments, showing that closed-loop training follows distinct trajectories from open-loop training due to the coupled agent–environment spectrum of the system matrix $\mathbf{P}$. It demonstrates a fundamental trade-off between short-term policy improvement and long-term stability, which organizes learning into three stages and a sequence of spectral transitions, including a dominant complex-conjugate pair and a slower real mode. By deriving a low-dimensional effective model $\mathbf{P}_{\mathrm{eff}}$ under rank-1 connectivity, the authors connect high-dimensional RNN dynamics to tractable analytic predictions that capture stage transitions, zig-zag optimization, and the emergence of internal world models. The results generalize to multi-frequency tracking tasks and align with human motor-learning patterns, highlighting the importance of closed-loop dynamics for biologically plausible learning and offering a framework for interpreting internal representations and spectral shifts during training.

Abstract

Recurrent neural networks (RNNs) trained on neuroscience-inspired tasks offer powerful models of brain computation. However, typical training paradigms rely on open-loop, supervised settings, whereas real-world learning unfolds in closed-loop environments. Here, we develop a mathematical theory describing the learning dynamics of linear RNNs trained in closed-loop contexts. We first demonstrate that two otherwise identical RNNs, trained in either closed- or open-loop modes, follow markedly different learning trajectories. To probe this divergence, we analytically characterize the closed-loop case, revealing distinct stages aligned with the evolution of the training loss. Specifically, we show that the learning dynamics of closed-loop RNNs, in contrast to open-loop ones, are governed by an interplay between two competing objectives: short-term policy improvement and long-term stability of the agent-environment interaction. Finally, we apply our framework to a realistic motor control task, highlighting its broader applicability. Taken together, our results underscore the importance of modeling closed-loop dynamics in a biologically plausible setting.

Figures (19)

• Figure 1: (a) Schematic of learning setups: In the closed-loop setting (top), the RNN output $u$ is fed into the environment, which evolves according to the system dynamics and produces the next input $y$ to the network. In contrast, the open-loop setting (bottom) lacks feedback: a student RNN is trained via supervised learning to imitate a pre-trained teacher mapping i.i.d. inputs $y$ to target outputs ${u}$. (b) Conceptual illustration of our hypothesis: optimization trajectories for closed-loop (blue) and open-loop (red) RNNs explore different regions of the closed-loop loss landscape (darker colors indicate higher loss). (c) Summary of results: the closed-loop RNN (top) progresses through three distinct learning stages, while the open-loop RNN (bottom) shows a sharp test loss peak. Note that the open-loop train loss (dashed black) drops by roughly one order of magnitude while the closed-loop test loss (red solid) drops by four, highlighting its limited sensitivity to the environment. Both panels use log scale; dashed and solid lines show train and test loss on right and left axes, respectively.
• Figure 2: (a) Closed-loop test loss (reproduced from Fig. \ref{['fig:1']}c). Despite identical architectures, the two RNNs exhibit distinct learning dynamics under closed-loop and open-loop training. The closed-loop RNN progresses through three distinct learning stages, whereas the open-loop RNN displays a sharp loss peak. Gray and yellow markers indicate the end of Stages 1 and 2, respectively, for the closed-loop RNN. (b) Numerical estimation of the effective RNN gain $\bm{K}_{\text{eff}}$ during training, projected onto the $(k_1, k_2)$ plane (closed-loop shown in blue; open-loop in red). Although both RNNs begin and end in similar regions, they follow markedly different trajectories across the loss landscape. The background color represents the logarithm of the loss; darker regions correspond to higher loss. The solid black curve separates three distinct stability regimes. (c) Zoomed-in view of panel (b), highlighting the divergence point between the two trajectories (gray marker). Gray and yellow markers match those in panel (a). Notably, the end of Stage 2 for the closed-loop RNN (yellow marker) coincides with a stability phase transition in the $(k_1, k_2)$ plane.
• Figure 3: Stage 1 - Negative position policy. (a) Training loss rapidly decreases during Stage 1. (b) Example trajectory of mass position under RNN control at the end of Stage 1, exhibiting oscillatory divergence; note, the red dashed line indicates the episode length used for training. (c) This oscillatory instability is explained by inspecting the eigenspectrum of $\bm{P}$, which reveals the emergence of a dominant complex-conjugate pair during training with $\rho(\bm{P}) > 1$. Arrow indicates trajectory; darker marker denotes final position. (d) The effective loss (dashed), predicted by the order parameter $\sigma_{\bm{z}\bm{m}}$ alone, shows strong agreement with simulation (solid). (e) As predicted by the asymptotic loss landscape in (d), despite initializing many RNNs with different initial overlaps, all networks converge to a small and negative $\sigma_{\bm{z}\bm{m}}$ by the end of Stage 1.
• Figure 4: Stage 2 - Building a world model. (a) Training loss enters a plateau phase. (b) Example trajectory of mass position under RNN control at the end of Stage 2, showing an underdamped response with sustained oscillations that eventually converge to the target (note the time-step on the x-axis). (c) Empirical evolution of one of the dominant complex-conjugate eigenvalues during Stage 2, exhibiting a zig-zag trajectory inward toward the unit disk (dashed). Arrow indicates trajectory; darker marker denotes final position. (d) Theoretical trajectories of the dominant eigenvalue optimizing the surrogate loss across different $\alpha$ values: optimizing only short-term loss ($\alpha = 0$) drives the eigenvalue upward along the imaginary axis; optimizing only long-term loss ($\alpha = 1$) causes direct descent; intermediate values ($0 < \alpha < 1$) yield trajectories consistent with empirical behavior; red marker indicates the initial condition. (e) Empirical validation: initializing the RNN from the same starting point (red dot) and varying episode length $T$ produces trajectories that qualitatively match the theoretical surrogate predictions.
• Figure 5: Stage 3 - Policy refinement. (a) Final drop in training loss marks the onset of Stage 3. (b) Example trajectory of mass position under RNN control at the end of Stage 3, exhibits fast and non-oscillatory dynamics. (c) Spectrum of the closed-loop matrix $\bm{P}$: the dominant complex-conjugate pair contracts further inward, while $\lambda_3$ grows, indicating the emergence of a second slow mode. Arrow indicates trajectory; darker marker denotes final position. (d) Training loss during Stage 3 for full high-dimensional model (solid) and the low-dimensional effective model (dashed black), across different regularization strengths $\beta$ (colored), initialized from identical Stage 2 endpoints. The effective model shows excellent agreement with the RNN learning dynamics throughout Stage 3.