Flowing Through States: Neural ODE Regularization for Reinforcement Learning

# Flowing Through States: Neural ODE Regularization for Reinforcement Learning
Source: [https://arxiv.org/html/2608.06595](https://arxiv.org/html/2608.06595)
Mohamed Ghanem1Bernd Finkbeiner1,2 1CISPA Helmholtz Center for Information Security 2Technical University of Munich mohamed\.ghanem@cispa\.de,finkbeiner@cispa\.de

###### Abstract

Neural networks applied to sequential decision\-making tasks typically rely on latent representations of environment states\. While environment dynamics dictate how semantic states evolve, the corresponding latent transitions are usually left implicit, creating a potential misalignment between the two\. We propose to model latent dynamics explicitly by drawing an analogy between Markov decision process \(MDP\) trajectories and ordinary differential equation \(ODE\) flows: in both cases, the current state fully determines its successors\. Building on this view, we introduce a neural ODE\-based regularization method that enforces latent embeddings to follow consistent ODE flows, thereby aligning representation learning with environment dynamics\. Although broadly applicable to deep learning agents, we demonstrate its effectiveness in reinforcement learning by integrating it into Actor\-Critic algorithms\. Our approach yields major performance gains across various standard Atari benchmarks for A2C and gridworld environments for PPO\.

## 1Introduction

A central challenge in machine learning is bridging the gap between an object’s semantic meaning and its latent representation\. Because neural networks operate on learned embeddings rather than direct semantics, representation learning has largely focused on designing processes that faithfully encode local object properties\. For instance, convolutional neural networks\(LeCunet al\.,[1989](https://arxiv.org/html/2608.06595#bib.bib40)\)incorporate inductive biases such as translation equivariance, spatial locality, and approximate invariance to scale and rotation\. These architectural choices encode object\-level regularities, ensuring that embeddings reflect structural properties intrinsic to individual objects\.

While such local representations are powerful for perception tasks,*sequential decision\-making*introduces a different challenge: the need for a more*global*understanding of how objects and states relate to one another over time\. In this setting, the relevant inductive biases emerge not from isolated objects but from the dynamics that connect them\. For example, in the context of Markov Decision Processes \(MDPs\), the latent embeddings of a state and its successor should be consistently related by the transition dynamics\. Concretely, if a transition ruleRRconnects states1s\_\{1\}tos2s\_\{2\}, then their embeddings should satisfy a relation of the form:

h​\(s2\)=g​\(h​\(s1\),R\),h\(s\_\{2\}\)=g\(h\(s\_\{1\}\),R\),whereh​\(⋅\)h\(\\cdot\)denotes the embedding function, andggis an arbitrary function\. While the existence of such a mapping is trivial in principle, the structural properties it imposes on the latent space, such as smoothness, consistency, and determinism, are far from trivial and are crucial for reasoning tasks\.

This paper proceeds from the intuition that embeddings of*semantic trajectories*can be understood as discretizations of continuous latent flows\. In other words, each trajectory in the semantic space should correspond to a smooth path in the latent space\. We argue that regularizing latent embeddings to respect this path structure captures an inherent property of transition dynamics, and enhances the model’s ability to learn the task on a more global level\. To operationalize this idea, we define latent flows using neural ordinary differential equations \(neural ODEs\)\(Chenet al\.,[2018](https://arxiv.org/html/2608.06595#bib.bib7)\), which guarantee unique continuous trajectories under mild regularity assumptions such as Lipschitz continuity\(Coddington and Levinson,[1955](https://arxiv.org/html/2608.06595#bib.bib26)\)\. In reasoning contexts, this uniqueness naturally subsumes*the Markov property*: an initial condition \(i\.e\., a state\) completely determines the flow path of subsequent conditions\.

However, directly using neural ODEs for inference is impractical: their reliance on numerical integration makes them significantly slower than standard forward passes, and their application to sequential inference is further complicated by the discontinuities introduced by evolving semantic states\(Duet al\.,[2020](https://arxiv.org/html/2608.06595#bib.bib22); Jia and Benson,[2019](https://arxiv.org/html/2608.06595#bib.bib18); Rubanovaet al\.,[2019](https://arxiv.org/html/2608.06595#bib.bib23)\)\. To overcome these limitations, we propose to train the agent’s semantic embedder tomimicthe flows of a neural ODE through an alignment penalty\. This approach enables the learned embeddings to inherit the topological structure of smooth ODE flows, while avoiding the computational and design burdens of ODE\-based inference\. Our method thus combines the expressivity of continuous\-time dynamics with the efficiency of conventional neural architectures\. Moreover, it adds a layer of global guidance to the agent in the form of a neural ODE that learns to model the latent agent\-environment dynamics in an unsupervised fashion\.

The relevance of this perspective is particularly pronounced in*discrete\-state*MDPs\. In continuous\-state environments, the inherent continuity of the state space naturally induces smoothness in the latent representations: small changes in the input state often correspond to small changes in the embedding\. By contrast, in discrete domains the semantic space consists of isolated states with no*a priori*notion of proximity or smooth transitions\. As a result, continuity must be imposed in the latent space rather than inherited from the state space itself\. Embedding discrete trajectories as smooth latent flows therefore provides a principled way to recover structural regularities that are otherwise absent, enabling latent dynamics to reflect the transition constraints of the underlying MDP\.

Contributions\.In this paper, we introduce flow regularization \(FlowReg\), an unsupervised regularization technique for sequential Markov decision\-making models that aligns the agent’s latent representation field with the underlying semantic environment dynamics\. It does so by learning a neural ODE that acts as a latent surrogate for the environment and aligning its flows with the latent trajectories of the agent’s state embedder\. To showcase our technique, we evaluate FlowReg in the reinforcement learning settings of Advantage Actor\-Critic \(A2C\) on 11 Atari environments\. Our experiments show that FlowReg notably improves the baseline model performance across all environments\. We further examine the resulting latent trajectories and demonstrate their desirable smoothness properties as a result of flow\-regularization\. Lastly, we also show the FlowReg boost to PPO on gridworld environments\.

## 2Related Work

#### Neural ODEs as continuous\-depth networks\.

It has been noted in several existing works that ResNets\(Heet al\.,[2016](https://arxiv.org/html/2608.06595#bib.bib14)\)can be viewed as an Euler discretization of a continuous differential flow\(Balázset al\.,[2021](https://arxiv.org/html/2608.06595#bib.bib8); Luet al\.,[2018](https://arxiv.org/html/2608.06595#bib.bib10); Haber and Ruthotto,[2017](https://arxiv.org/html/2608.06595#bib.bib11)\)\. An implication of this is that an ODE can, in theory, be used to model an infinite\-depth ResNet with a finite number of parameters – making them more parameter efficient\(Chenet al\.,[2018](https://arxiv.org/html/2608.06595#bib.bib7)\)\. In this paper, we take a broader view of sequence transformations modeled by the whole network as an embedder, rather than transformations modeled by the individual layers within the model\. That is, instead of looking at the embedder network as a discretized transformation of an object, we look at the latent trajectories that result from applying the network to a sequence of objects that are sequentially related under well\-defined environment dynamics\.

#### Neural ODEs for continuous control\.

Neural ODEs can model the continuous evolution between discrete events while coupling with event\-triggered mechanisms or classifiers to detect and handle abrupt transitions, e\.g\., collisions or control mode changes\(Jia and Benson,[2019](https://arxiv.org/html/2608.06595#bib.bib18); Auzinaet al\.,[2023](https://arxiv.org/html/2608.06595#bib.bib19)\)\. By integrating traditional neural networks, these models can infer both the continuous flow and the timing or conditions of discrete switches directly from data, bypassing rigid analytical formulations\. The work ofAlvarezet al\.\([2020](https://arxiv.org/html/2608.06595#bib.bib41)\)bears a partial resemblance to ours in that it involves training an ODE to learn entire trajectories of continuous\-space environments\. However, both works fundamentally differ from our approach in that our neural ODE operates on latent trajectories while theirs aim to predict semantic trajectories, which makes them rather cumbersome to apply to discrete\-space tasks since the network’s output is continuous\. Similar toDuet al\.\([2020](https://arxiv.org/html/2608.06595#bib.bib22)\), they use the neural ODE as the main inference model, whereas we only use the neural ODE as a decoupled regularizer\.

#### Shaping representations by predictive coding\.

Enhancing temporal consistency across trajectories requires moving beyond static state discriminators to objectives that model long\-horizon dynamics\. By fusing predictive coding with contrastive learning, representations can be shaped to maximize the mutual information between past history and future outcomes, effectively smoothing the latent space against high\-frequency noise\(Agarwalet al\.,[2021](https://arxiv.org/html/2608.06595#bib.bib49); Schwarzeret al\.,[2020](https://arxiv.org/html/2608.06595#bib.bib47)\)\. Methods like TACO\(Zhenget al\.,[2023](https://arxiv.org/html/2608.06595#bib.bib46)\)enforce a robust temporal structure in the latent space, where state transitions are predictable from their immediate predecessors, preventing the representation from drifting due to task\-irrelevant environmental stochasticity\. Our method enforces a stricter notion of temporal consistency by leveraging the uniqueness of ODE flows at any intermediate point, ensuring that states are predictable givenanyof their predecessors, not only the immediate ones\.

## 3Preliminaries

### 3\.1Markov Decision Processes

We model reinforcement learning \(RL\) problems as*Markov decision processes*\(MDPs\), defined by the tuple

ℳ=\(𝒮,𝒜,P,r,γ\),\\mathcal\{M\}=\(\\mathcal\{S\},\\mathcal\{A\},P,r,\\gamma\),\(1\)where𝒮\\mathcal\{S\}is the state space,𝒜\\mathcal\{A\}the action space,P​\(s′∣s,a\)P\(s^\{\\prime\}\\mid s,a\)the transition kernel,r​\(s,a\)r\(s,a\)the expected immediate reward, andγ∈\[0,1\)\\gamma\\in\[0,1\)a discount factor\. An agent samples actionsat∈𝒜a\_\{t\}\\in\\mathcal\{A\}according to a policyπ​\(a∣s\)\\pi\(a\\mid s\), inducing a trajectoryτ=\(s0,a0,r0,…\)\\tau=\(s\_\{0\},a\_\{0\},r\_\{0\},\\ldots\)The objective is to maximize the expected return

J​\(π\)=𝔼π​\[∑t=0∞γt​r​\(st,at\)\]J\(\\pi\)=\\mathbb\{E\}\_\{\\pi\}\\left\[\\sum\_\{t=0\}^\{\\infty\}\\gamma^\{t\}r\(s\_\{t\},a\_\{t\}\)\\right\]\(2\)We define the following key functions:

- •The state\-value function:Vπ\(s\)=𝔼π\[∑t=0∞γtr\(st,at\)\|s0=s\]V^\{\\pi\}\(s\)=\\mathbb\{E\}\_\{\\pi\}\\left\[\\sum\_\{t=0\}^\{\\infty\}\\gamma^\{t\}r\(s\_\{t\},a\_\{t\}\)\\,\\middle\|\\,s\_\{0\}=s\\right\]
- •The action\-value function:Qπ\(s,a\)=𝔼π\[∑t=0∞γtr\(st,at\)\|s0=s,a0=a\]Q^\{\\pi\}\(s,a\)=\\mathbb\{E\}\_\{\\pi\}\\left\[\\sum\_\{t=0\}^\{\\infty\}\\gamma^\{t\}r\(s\_\{t\},a\_\{t\}\)\\,\\middle\|\\,s\_\{0\}=s,\\,a\_\{0\}=a\\right\]
- •The advantage function:Aπ​\(s,a\)=Qπ​\(s,a\)−Vπ​\(s\)A^\{\\pi\}\(s,a\)=Q^\{\\pi\}\(s,a\)\-V^\{\\pi\}\(s\)

### 3\.2Policy Gradient Methods

Policy gradient algorithms directly optimize a parametric policyπθ​\(a∣s\)\\pi\_\{\\theta\}\(a\\mid s\)\. The policy gradient theorem\(Suttonet al\.,[1999](https://arxiv.org/html/2608.06595#bib.bib35)\)states:

∇θJ​\(πθ\)=𝔼s∼dπθ,a∼πθ​\[∇θlog⁡πθ​\(a∣s\)​Qπθ​\(s,a\)\]\\nabla\_\{\\theta\}J\(\\pi\_\{\\theta\}\)=\\mathbb\{E\}\_\{s\\sim d^\{\\pi\_\{\\theta\}\},\\,a\\sim\\pi\_\{\\theta\}\}\\\!\\left\[\\nabla\_\{\\theta\}\\log\\pi\_\{\\theta\}\(a\\mid s\)\\,Q^\{\\pi\_\{\\theta\}\}\(s,a\)\\right\]\(3\)wheredπθd^\{\\pi\_\{\\theta\}\}denotes the stationary state distribution underπθ\\pi\_\{\\theta\}\. In practice,QπθQ^\{\\pi\_\{\\theta\}\}is approximated and variance is reduced by subtracting a baseline such asVπ​\(s\)V^\{\\pi\}\(s\)\.

### 3\.3Advantage Actor–Critic \(A2C\)

Actor–critic methods\(Mnihet al\.,[2016](https://arxiv.org/html/2608.06595#bib.bib28)\)couple a policy model \(the actor\) with a value function estimator \(the critic\)\. The actor updates its parametersθ\\thetavia the policy gradient, while the critic learns to estimateVπ​\(s\)V^\{\\pi\}\(s\)\(orQπ​\(s,a\)Q^\{\\pi\}\(s,a\)\) using temporal\-difference learning\.

The*Advantage Actor–Critic \(A2C\)*algorithm improves stability by using an advantage estimator\. The policy gradient update is given by

∇θJ​\(πθ\)≈𝔼​\[∇θlog⁡πθ​\(at∣st\)​A^t\]\\nabla\_\{\\theta\}J\(\\pi\_\{\\theta\}\)\\approx\\mathbb\{E\}\\\!\\left\[\\nabla\_\{\\theta\}\\log\\pi\_\{\\theta\}\(a\_\{t\}\\mid s\_\{t\}\)\\,\\hat\{A\}\_\{t\}\\right\]\(4\)with empirical advantage

A^t=rt\+γ​Vθ​\(st\+1\)−Vθ​\(st\)\\hat\{A\}\_\{t\}=r\_\{t\}\+\\gamma V\_\{\\theta\}\(s\_\{t\+1\}\)\-V\_\{\\theta\}\(s\_\{t\}\)\(5\)whereVθV\_\{\\theta\}is the critic parameterized byθ\\theta\. The critic is trained by minimizing the squared error

ℒcritic​\(θ\)=𝔼st∼πθ​\[\(rt\+γ​Vθ​\(st\+1\)−Vθ​\(st\)\)2\]\\mathcal\{L\}\_\{\\text\{critic\}\}\(\\theta\)=\\mathbb\{E\}\_\{s\_\{t\}\\sim\\pi\_\{\\theta\}\}\\Big\[\\big\(r\_\{t\}\+\\gamma V\_\{\\theta\}\(s\_\{t\+1\}\)\-V\_\{\\theta\}\(s\_\{t\}\)\\big\)^\{2\}\\Big\]\(6\)ℒactor​\(θ\)=−𝔼st,at∼πθ​\[log⁡πθ​\(at∣st\)​A^t\]\\mathcal\{L\}\_\{\\text\{actor\}\}\(\\theta\)=\-\\mathbb\{E\}\_\{s\_\{t\},a\_\{t\}\\sim\\pi\_\{\\theta\}\}\\Big\[\\log\\pi\_\{\\theta\}\(a\_\{t\}\\mid s\_\{t\}\)\\,\\hat\{A\}\_\{t\}\\Big\]\(7\)

### 3\.4Neural Ordinary Differential Equations

A Neural Ordinary Differential Equation is defined by the continuous transformation of the hidden stateh​\(t\)h\(t\)given by the differential equation:

d​𝐡​\(t\)d​t=fϕ​\(𝐡​\(t\),t\),𝐡​\(t\)=𝐡​\(t0\)\+∫t0tfϕ​\(𝐡​\(s\),s\)​ds\\displaystyle\\frac\{\\mathrm\{d\}\\mathbf\{h\}\(t\)\}\{\\mathrm\{d\}t\}=f\_\{\\phi\}\(\\mathbf\{h\}\(t\),t\),\\quad\\quad\\mathbf\{h\}\(t\)=\\mathbf\{h\}\(t\_\{0\}\)\+\\int\_\{t\_\{0\}\}^\{t\}f\_\{\{\}\_\{\\phi\}\}\(\\mathbf\{h\}\(s\),s\)\\,\\mathrm\{d\}s\(8\)whereffis a neural network parameterized byϕ\\phi\. As such, neural ODEs differs from classical deep learning in that the neural network is used to model the system dynamics \(through the state derivative\) at a given time instead of modeling the entire system directly\. This framework can be used to model functions that evolve over time\. To seamlessly integrate neural ODEs into traditional deep learning pipeline, a differentiable numeric solver \(e\.g\.,torchdiffeq\(Chenet al\.,[2018](https://arxiv.org/html/2608.06595#bib.bib7)\)orDiffrax\(Kidger,[2021](https://arxiv.org/html/2608.06595#bib.bib21)\)\) is typically used to evaluate the latent state function at given time points\. The continuous\-depth nature of Neural ODEs allows adaptive computation \(e\.g\., varying solver step sizes\), offering memory efficiency and flexible trade\-offs between precision and computational cost compared to fixed\-depth architectures\.

A key mathematical property of Neural ODEs is their invertibility and exact gradient calculation via the adjoint state, which ensures stable training even with long integration intervals\. The framework inherently accommodates irregularly sampled or continuous\-time data, making them suitable for tasks like time\-series modeling and dynamical systems\. However, their performance hinges on numerical solver choices: explicit methods \(e\.g\., Euler\) are computationally light but may struggle with stiff systems, while implicit methods \(e\.g\., backward differentiation\) enhance stability at higher computational cost\. This interplay between numerical precision, stability, and efficiency underscores the importance of solver selection in practice\. Additionally, Neural ODEs enable novel architectures, such as continuous normalizing flows for density estimation, by enforcing invertibility through Lipschitz constraints onff\. By bridging deep learning with differential equations, they provide a principled framework for understanding neural networks as dynamical systems, opening avenues for interpretability and integration with scientific machine learning\.

## 4Approach

![Refer to caption](https://arxiv.org/html/2608.06595v1/images/FlowReg_main-hirez.png)Figure 1:Illustration of the flow regularization landscape\.In this section, we outline the mathematical formulation of our flow regularization technique for a general target model\. As illustrated in Figure[1](https://arxiv.org/html/2608.06595#S4.F1), our setting involves three principal fields: \(1\) the semantic state field defined by the environment, \(2\) the latent observation vector field induced by the semantic state embedder on the environment, and where each point is a vector representation of the corresponding semantic state, and \(3\) the latent flow vector field defined by the neural ODE \(i\.e\., flow model\)\. Field \(2\) is utilized for carrying task information from Field \(1\) into the latent space, while Field \(3\) is utilized for imposing a global latent structure that underpins Field \(1\)\. The essence of our approach is that by aligning \(2\) and \(3\), we get the best of both worlds: a latent field that captures local \(state\-level\) and global \(trajectory\-level\) aspects of the environment\.

### 4\.1Model Setup

Generally, there are two models involved in our framework, namely a target agent modelθ\\thetaand a flow regularizer modelϕ\\phi\. The target model comprises a state embedder network𝐡θ\\mathbf\{h\_\{\\theta\}\}that converts semantic states into their latents, and a downstream headFθF\_\{\\theta\}that produces the final task\-related actions\. For a state trajectory𝐬=s0,s1,…,sN−1\\mathbf\{s\}=s\_\{0\},s\_\{1\},\.\.\.,s\_\{N\-1\}, semantic embeddings are computed as𝐇θ​\(s\)=𝐡θ​\(s0\),𝐡θ​\(s1\),…,𝐡θ​\(sN−1\)\\mathbf\{H\_\{\\theta\}\}\(s\)=\\mathbf\{h\_\{\\theta\}\}\(s\_\{0\}\),\\mathbf\{h\_\{\\theta\}\}\(s\_\{1\}\),\.\.\.,\\mathbf\{h\_\{\\theta\}\}\(s\_\{N\-1\}\), while flow embeddings are obtained by solving the initial value problem on𝐡ϕ​\(0\)=𝐡θ​\(s0\)\\mathbf\{h\_\{\\phi\}\}\(0\)=\\mathbf\{h\_\{\\theta\}\}\(s\_\{0\}\):

𝐇θ​\(s\)=\{𝐡θ​\(si\)\}i=0N−1=𝐡θ​\(\{si\}i=0N−1\)\\mathbf\{H\_\{\\theta\}\}\(s\)=\\\{\\mathbf\{h\_\{\\theta\}\}\(s\_\{i\}\)\\\}\_\{i=0\}^\{N\-1\}=\\mathbf\{h\_\{\\theta\}\}\(\\\{s\_\{i\}\\\}\_\{i=0\}^\{N\-1\}\)\(9\)𝐇ϕ​\(s\)=\{𝐡ϕ​\(si\)\}i=1N−1=ODESolve​\(fϕ,𝐡θ​\(s0\),\{τi\}i=0N−1\)\\mathbf\{H\_\{\\phi\}\}\(s\)=\\\{\\mathbf\{h\_\{\\phi\}\}\(s\_\{i\}\)\\\}\_\{i=1\}^\{N\-1\}=\\text\{ODESolve\}\(f\_\{\\phi\},\\mathbf\{h\_\{\\theta\}\}\(s\_\{0\}\),\\\{\\tau\_\{i\}\\\}\_\{i=0\}^\{N\-1\}\)\(10\)
whereτi\\tau\_\{i\}is the integration time index for statesis\_\{i\}, andfϕf\_\{\\phi\}is a neural network that parameterizes the derivative of the latent state\. MDP states generally do not have timestamps, so we impose a time sampling scheme to associate each state in the trajectory with a time index\. Note that due to the Markov property, the underlying ODE is autonomous \(i\.e\., time\-invariant\)\. However, the choice of the integration times still significantly influences the ODE solver, and our experiments show that it is indeed fairly consequential for performance\. An intuitive option for time sampling would be the step index of the state, i\.e\.,τi=i\\tau\_\{i\}=i\. Another simple approach is using a discounted time horizon with the same discounting factorγ\\gammaused by the agent’s algorithm, i\.e\.,τi=γi\\tau\_\{i\}=\\gamma^\{i\}where0<γ<10<\\gamma<1\. This guarantees that integration times are in\[0,1\]\[0,1\]to avoid arbitrarily large integration times, which might lead to gradient instability\.

### 4\.2Path Alignment

In essence, the flow model defines a smooth latent path that starts at a given semantic state embedding point, whereas the semantic embedder defines a discrete point sequence in the latent space\. Typically, this latent point sequence is topologically unconstrained, which means that the topological structure of the latent space has to be implicitly learned over the course of the training\. The key idea here is that we can speed up this process by imposing a topological structure that we already know to be compatible with the domain\.

Our approach proceeds from the rationale that initially, the flow model carries pure curvature information while the semantic embedder carries task information\. Ideally, we want to fuse both signals into the target model\. To that end, we align the semantic embedding trajectory with the discretized latent flow\. In doing so, each network adapts the information carried by the other\. One straightforward way to incentivize this alignment is by minimizing the MSE between the latent point sequence𝐇θ\\mathbf\{H\_\{\\theta\}\}and the sampled flow path𝐇ϕ\\mathbf\{H\_\{\\phi\}\}\. As such, we can compute the flow regularization loss as follows:

ℒflow​\(s\):=∥𝐇θ​\(s\)−𝐇ϕ​\(s\)∥22N\(FlowReg\)\\mathcal\{L\}\_\{\\text\{flow\}\}\(s\):=\\frac\{\\lVert\\mathbf\{H\_\{\\theta\}\}\(s\)\-\\mathbf\{H\_\{\\phi\}\}\(s\)\\rVert^\{2\}\_\{2\}\}\{N\}\\quad\\quad\(\\text\{FlowReg\}\)\(11\)

### 4\.3Overall Training Objective

Having computed the flow loss on the latent trajectory, this loss is then added to the label\-based task loss:

ℒ​\(s,y\)=ℒtask​\(Fθ​\(𝐇θ​\(s\)\),y\)\+λ​ℒflow​\(s\)\\mathcal\{L\}\(s,y\)=\\mathcal\{L\}\_\{\\text\{task\}\}\(F\_\{\\theta\}\(\\mathbf\{H\_\{\\theta\}\}\(s\)\),y\)\+\\lambda\\mathcal\{L\}\_\{\\text\{flow\}\}\(s\)\(12\)
whereλ\\lambdais the flow\-loss weighting factor\. Note thatℒflow​\(s\)\\mathcal\{L\}\_\{\\text\{flow\}\}\(s\)involves both the semantic embedderθ\\thetaand the neural ODE networkϕ\\phi\. This trainsθ\\thetato follow the continuous ODE flow while optimizingϕ\\phito indirectly adapt to the underlying task modeled byθ\\theta\.

For an Advantage Actor\-Critic agent, the overall training loss would be:

ℒ​\(s,y\)=ℒactor​\(s,y\)\+β​ℒcritic​\(s,y\)\+λ​ℒflow​\(s\)\\mathcal\{L\}\(s,y\)=\\mathcal\{L\}\_\{\\text\{actor\}\}\(s,y\)\+\\beta\\mathcal\{L\}\_\{\\text\{critic\}\}\(s,y\)\+\\lambda\\mathcal\{L\}\_\{\\text\{flow\}\}\(s\)\(13\)
A relevant hyperparameter here is the FlowReg update frequency relative to the agent policy updates\. It is also important to note that the neural ODE is not used for inference, only as a training\-time adaptive regularizer\.

## 5Experiments

We evaluate our method on 11 Atari environments from the Arcade Learning Environment \(ALE\) library\(Bellemareet al\.,[2013](https://arxiv.org/html/2608.06595#bib.bib32)\)\. This is mainly due to A2C being a reasonably simple actor\-critic formulation, which is a cornerstone for many state\-of\-the\-art algorithms like PPO\(Schulmanet al\.,[2017](https://arxiv.org/html/2608.06595#bib.bib30)\)and SAC\(Haarnojaet al\.,[2018](https://arxiv.org/html/2608.06595#bib.bib29)\)\. We build on the Stable\-baselines3 A2C implementation\(Raffinet al\.,[2021](https://arxiv.org/html/2608.06595#bib.bib45)\)to incorporate our regularization loss\. We use the same set of A2C hyperparameters for all environments and agents\. The agent networks for both baseline and flow\-regularized variants are identical for all experiments\. The ultimate goal of our evaluation is to show that flow regularization effectively reduces the training search space by imposing an ODE flow field on the latent space of the agent’s state embedder, hence greatly reducing variance during training, allowing the agent to learn better policies with the same training steps\.

### 5\.1Atari Benchmarks

![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Qbert.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Alien.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Riverraid.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Breakout.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Asterix.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/DemonAttack.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Atlantis.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/BeamRider.png)

Figure 2:Episodic rewards of baseline and flow\-regularized A2C on 8 different Atari environments with a rolling average window of 100 episodes\.#### Hyperparameters\.

We performed 5 independent runs for every RL agent across all environments for 10 million timesteps each\. Our semantic embedder for both baseline and flow\-regularized agents is a commonly used Nature CNN\(Mnihet al\.,[2015](https://arxiv.org/html/2608.06595#bib.bib34)\)feature extractor that embeds game state \(frames\) into a 512\-dimensional vector space\. The ODE flow \(and loss\) is computed on the extracted state feature vectors\. For the FlowReg ODE network, we use a two\-layer MLP with atanh\\tanhactivation on the first layer\. All models are optimized by RMSProp\(Ruder,[2016](https://arxiv.org/html/2608.06595#bib.bib42)\)with an initial learning rate of7×10−47\\times 10^\{\-4\}and a linear decay scheduler\. We apply a global\-norm gradient clipping ratio of 0\.5\(Pascanuet al\.,[2012](https://arxiv.org/html/2608.06595#bib.bib16)\)\. We use thetorchdiffeq\(Chenet al\.,[2018](https://arxiv.org/html/2608.06595#bib.bib7)\)library together withPyTorchfor solving neural ODEs with relative tolerance =10−410^\{\-4\}, and absolute tolerance =10−510^\{\-5\}\. For FlowReg variants, we experiment with both index\-based \(τi=i\\tau\_\{i\}=i\) and exponential decay \(τi=γi\\tau\_\{i\}=\\gamma^\{i\}\) time sampling, along with a regularization frequency \(relative to agent updates\) of\{5,10,20\}\\\{5,10,20\\\}, and take the best configuration averaged over 3 seeds dedicated for hyperparameter search and separate from the 10 seeds of the final comparison runs\. For simplicity, we setλ=1\\lambda=1for all environments\.

#### Flow\-regularized agents consistently outperform the baseline on Atari environments\.

Figure[2](https://arxiv.org/html/2608.06595#S5.F2)highlights the notable performance gap between flow\-regularized A2C and the baseline\. The learning curves on all 11 environments can be found in Figure[5](https://arxiv.org/html/2608.06595#A1.F5)\(Appendix[A](https://arxiv.org/html/2608.06595#A1)\)\. Figure[3](https://arxiv.org/html/2608.06595#S5.F3)shows the overall performance percent gains achieved by applying FlowReg on all 11 environments111The hatched strip in Figure[3](https://arxiv.org/html/2608.06595#S5.F3)indicates values exceeding the y\-axis limit, which was capped for visual clarity to avoid overly downscaling other values\.\. We also find that most FlowReg configurations outperform the baseline across all environments, which means that finding good values for the two FlowReg hyperparameters \(time sampling and update frequency\) is fairly easy\.

![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/stats/tradeoff.png)Figure 3:Trade\-off between performance gain achieved by FlowReg and its runtime overhead\.
#### FlowReg performance gains are robust under time sampling modes\.

As shown in Table[1](https://arxiv.org/html/2608.06595#S5.T1), FlowReg largely improves the baseline performance under bothIndexandExp\-Decaytime sampling modes\. The choice between them, in all likelihood, depends on the granularity of the environment dynamics\. We generally expectExp\-Decayto work better on environments with swifter or more fine\-grained state transitions\. Table[4](https://arxiv.org/html/2608.06595#A2.T4)\(Appendix[B](https://arxiv.org/html/2608.06595#A2)\) shows the specific FlowReg configurations that performed best on each environment along with the corresponding runtimes\.

Table 1:Best mean episode rewards of different time sampling modes\. Each variant was evaluated on 16 episodes averaged across 10 different training seeds\.Indexis whereτi=i\\tau\_\{i\}=iandExp\-Decayis whereτi=γi\\tau\_\{i\}=\\gamma^\{i\}\.A2C AgentQbertRiverraidBeamRiderBase4374\.30±958\.424374\.30\\pm 958\.421862\.27±2399\.581862\.27\\pm 2399\.58960\.66±748\.23960\.66\\pm 748\.23FlowReg \(Index\)8306\.05±1752\.71\\mathbf\{8306\.05\\pm 1752\.71\}2946\.34±2788\.172946\.34\\pm 2788\.171590\.96±1033\.301590\.96\\pm 1033\.30FlowReg \(Exp\-Decay\)6903\.15±2157\.716903\.15\\pm 2157\.712947\.95±2798\.64\\mathbf\{2947\.95\\pm 2798\.64\}1593\.11±961\.77\\mathbf\{1593\.11\\pm 961\.77\}

Table 2:Mean episode rewards of different FlowReg update frequencies relative to agent updates on Atari Qbert\. Each variant was evaluated on 16 episodes averaged across 10 different training seeds\.U\-mmeans the FlowReg loss is applied once everymagent updates\.A2C AgentQbert \(Index\)Qbert \(Exp\-Decay\)Base4374\.30±958\.424374\.30\\pm 958\.424374\.30±958\.424374\.30\\pm 958\.42FlowReg U\-58306\.05±1752\.71\\mathbf\{8306\.05\\pm 1752\.71\}5286\.60±1269\.765286\.60\\pm 1269\.76FlowReg U\-106569\.51±2645\.716569\.51\\pm 2645\.716903\.15±2157\.71\\mathbf\{6903\.15\\pm 2157\.71\}FlowReg U\-205985\.70±2756\.175985\.70\\pm 2756\.176782\.70±1877\.136782\.70\\pm 1877\.13

#### FlowReg loss is still effective under a much lower update frequency compared to the agent loss\.

Table[2](https://arxiv.org/html/2608.06595#S5.T2)points to it being more ideal to apply FlowReg loss once every 10 agent updates under both time sampling modes\. The fourth row \(U\-20\) also shows that FlowReg still results in notable performance gains with half as many updates\. This is good news for runtime as it means the FlowReg loss does not need to be aggressively optimized to improve over the baseline, which allows it to run in a comparable training time\. By contrasting the time\-overhead margins with the performance gains in Figure[3](https://arxiv.org/html/2608.06595#S5.F3), it shows that FlowReg is an overall cost\-effective choice\. Figure[6](https://arxiv.org/html/2608.06595#A2.F6)and Table[4](https://arxiv.org/html/2608.06595#A2.T4)\(Appendix[B](https://arxiv.org/html/2608.06595#A2)\) show the runtime comparison between the baseline and FlowReg in terms of absolute values\.

Table 3:Latent path smoothness measures normalized by trajectory length\.EnvMetricPath LengthNet DisplacementAccel\. EnergyRewardFormula∑t=0N−1∥Δ​𝐡θ​\(𝐬𝐭\)∥\\sum\_\{t=0\}^\{N\-1\}\\lVert\\Delta\\mathbf\{h\_\{\\theta\}\(s\_\{t\}\)\}\\rVert∥𝐡θ​\(𝐬𝐍−𝟏\)−𝐡θ​\(𝐬𝟎\)∥\\lVert\\mathbf\{h\_\{\\theta\}\(s\_\{N\-1\}\)\}\-\\mathbf\{h\_\{\\theta\}\(s\_\{0\}\)\}\\rVert∑t=0N−2∥Δ2​𝐡θ​\(𝐬𝐭\)∥\\sum\_\{t=0\}^\{N\-2\}\\lVert\\Delta^\{2\}\\mathbf\{h\_\{\\theta\}\(s\_\{t\}\)\}\\rVert∑t=0NRt\\sum\_\{t=0\}^\{N\}R\_\{t\}QbertA2C34\.39±2\.1434\.39\\pm 2\.140\.44±0\.170\.44\\pm 0\.174424\.75±521\.764424\.75\\pm 521\.764374\.30±958\.424374\.30\\pm 958\.42A2C\+TACO6\.13±0\.426\.13\\pm 0\.420\.03±0\.01\\mathbf\{0\.03\\pm 0\.01\}106\.38±9\.86106\.38\\pm 9\.862434\.05±2474\.442434\.05\\pm 2474\.44A2C\+FlowReg4\.20±0\.44\\mathbf\{4\.20\\pm 0\.44\}0\.10±0\.020\.10\\pm 0\.0264\.17±7\.05\\mathbf\{64\.17\\pm 7\.05\}8306\.05±1752\.71\\mathbf\{8306\.05\\pm 1752\.71\}BreakoutA2C104\.09±2\.44104\.09\\pm 2\.440\.74±0\.280\.74\\pm 0\.2831432\.59±1698\.8231432\.59\\pm 1698\.8219\.40±1\.8619\.40\\pm 1\.86A2C\+TACO13\.09±1\.0813\.09\\pm 1\.080\.13±0\.050\.13\\pm 0\.05461\.75±125\.72461\.75\\pm 125\.7211\.12±2\.4211\.12\\pm 2\.42A2C\+FlowReg4\.92±0\.23\\mathbf\{4\.92\\pm 0\.23\}0\.06±0\.02\\mathbf\{0\.06\\pm 0\.02\}94\.98±9\.51\\mathbf\{94\.98\\pm 9\.51\}24\.03±0\.84\\mathbf\{24\.03\\pm 0\.84\}RiverraidA2C75\.36±2\.7275\.36\\pm 2\.720\.53±0\.070\.53\\pm 0\.0718298\.55±1487\.2818298\.55\\pm 1487\.281862\.27±2399\.581862\.27\\pm 2399\.58A2C\+TACO50\.35±1\.2650\.35\\pm 1\.260\.36±0\.040\.36\\pm 0\.047599\.32±404\.507599\.32\\pm 404\.502943\.47±1616\.302943\.47\\pm 1616\.30A2C\+FlowReg6\.35±0\.29\\mathbf\{6\.35\\pm 0\.29\}0\.06±0\.02\\mathbf\{0\.06\\pm 0\.02\}137\.25±10\.11\\mathbf\{137\.25\\pm 10\.11\}2947\.95±2798\.64\\mathbf\{2947\.95\\pm 2798\.64\}

### 5\.2Latent Path Smoothness

In addition to the performance results, we set out to investigate some geometric properties of the latent paths \(trajectories\) of flow\-regularized models compared to the baseline\. In particular, we are interested in whether FlowReg induces smoother paths as a result of the ODE alignment\. We measure 3 different smoothness metrics as shown in Table[3](https://arxiv.org/html/2608.06595#S5.T3)\. All 3 metrics are computed on the full dimensionality of the latent space without any reduction, and∥⋅∥\\lVert\\cdot\\rVertis the Euclidean norm\. To control for trajectory length variations, all 3 metrics are normalized by trajectory length, so they correspond to average speed, velocity, and acceleration, respectively\.

Path length measures total segment length along the path, which reflects the jump step size between consecutive states in the latent space\. Ideally, latent representations of consecutive states should be in close proximity, so the smaller the path length, the better the state embedder is from a purely topological standpoint\. Lower net path displacement is desirable for similar reasons, as it indicates that individual trajectories lie in tightly packed regions of the latent space\. Acceleration energy, computed the second\-difference in position:Δ2​𝐡θ​\(si\)=𝐡θ​\(si\+2\)−2​𝐡θ​\(si\+1\)\+𝐡θ​\(si\)\\Delta^\{2\}\\mathbf\{h\_\{\\theta\}\}\(s\_\{i\}\)=\\mathbf\{h\_\{\\theta\}\}\(s\_\{i\+2\}\)\-2\\mathbf\{h\_\{\\theta\}\}\(s\_\{i\+1\}\)\+\\mathbf\{h\_\{\\theta\}\}\(s\_\{i\}\), is a more local measure roughness \(lower is better\)\.

#### FlowReg results in much smoother latent trajectories while improving overall performance\.

Table[3](https://arxiv.org/html/2608.06595#S5.T3)shows that ODE flow alignment notably changes the basic geometric properties of the agent’s latent trajectories, making them much smoother and more tightly wound, consistently across environments\. Naturally, we do not attribute the performance improvement solely to the latent trajectory smoothing effect, since there are many ways to smooth the space while destroying the semantic structure, as evident by the fact that although TACO produces smoother paths than baseline over all 3 environments, it leads to a considerable performance degradation on two of them\. The key distinction in this case is restricting the latent field while respecting the underlying transition dynamics\. In our case, this is achieved by the mutual alignment loss that imposes a diffeomorphic structure on the latent space, resulting in reduced variance as abrupt jumps and crossings are naturally penalized because they violate ODE flows\.

Another takeaway from Table[3](https://arxiv.org/html/2608.06595#S5.T3)is that smoothness and temporal predictability are notably correlated\. Despite the differences in mechanism between TACO and FlowReg, they both aim to instate a notion of predictive temporal structure on the latent representations\. The results of Table[3](https://arxiv.org/html/2608.06595#S5.T3)suggest that this common feature explains the notable reduction in their latent path roughness compared to the baseline\.

### 5\.3Minigrid Environments

We evaluate FlowReg on PPO\(Schulmanet al\.,[2017](https://arxiv.org/html/2608.06595#bib.bib30)\)in Minigrid environments\(Chevalier\-Boisvertet al\.,[2023](https://arxiv.org/html/2608.06595#bib.bib44)\)\. These experiments serve the purposes of showing FlowReg’s efficacy on another major RL algorithm \(PPO\) while also exploring a more radically discrete environment domain than Atari games\. Similar to A2C, we use a modified implementation of the Stable\-Baselines\-3 PPO\(Dhariwalet al\.,[2017](https://arxiv.org/html/2608.06595#bib.bib33)\)\. We use theIndexU\-20 FlowReg configuration for all 3 environments\. We performed 10 runs per agent for 1M timesteps each\.

![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/ppo/FourRooms.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/ppo/Dynamic-Obstacles.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/ppo/DoorKey.png)

Figure 4:Episodic rewards of baseline and flow\-regularized PPO on Minigrid environments with a rolling average window of 100 episodes\.As shown in Figure[4](https://arxiv.org/html/2608.06595#S5.F4), flow\-regularized PPO has a clear advantage onFourRoomsandDynamic\-Obstacleswhile matching the baseline inDoorKey, where both agents practically solve the environment\.

## 6Conclusion

#### Summary\.

In this paper, we presented FlowReg, an unsupervised regularization technique that aligns MDP semantic trajectories with their latent counterparts\. We realized this goal by adding an unsupervised loss term that incentivizes the semantic trajectory embeddings to act like discretizations of a global neural ODE flow\. We chose actor\-critic reinforcement learning on Atari and Minigrid environments to showcase the benefits of applying FlowReg to a target model\. Our results have shown that using FlowReg notably boosts the overall performance of the target agent across almost all attempted environments and results in a more constrained path structure on the learned embedding space\.

#### Limitations\.

Although FlowReg does not require full episodes, it still requires trajectory information to align it with the learned ODE flow\. This means the training pipeline needs to keep track of the episode ID for each state\-action pair\. This is not a significant challenge for the classical RL pipeline structure, where each batch resumes from the environment state after the previous batch\. However, this might impose more implementation demands on more complex pipelines that do not place as much emphasis on episodic structure\. A more fundamental limitation of FlowReg is the fact that ODE flows are unique both forwards and backwards, so flow paths do not intersect themselves or each other\. This can be beneficial for discouraging looping behavior where an agent returns to a previously visited state\. However, this property could present a burden in environments where there are intermediate bottleneck states that need to be passed from different starting states\. An example of that is a maze solver game where the target destination lies in a chamber with only one opening\. Fortunately, this is often not the case for environments with a very large state space \(like Atari\)\.

#### Future Work\.

Since experiments demonstrate the efficacy of FlowReg on a standard on\-policy RL algorithm, it would be of great interest to see how it fares in the off\-policy settings such as DQN\(Mnihet al\.,[2013](https://arxiv.org/html/2608.06595#bib.bib36)\), as well as model\-based algorithms like Dreamer\(Okada and Taniguchi,[2021](https://arxiv.org/html/2608.06595#bib.bib43)\)\. Although the scope of our evaluation pertains to RL, the method itself still lends itself to MDPs in other learning paradigms such as imitation learning or semi\-supervised learning\. As such, these investigations would be a very promising research direction\.

## References

- Contrastive behavioral similarity embeddings for generalization in reinforcement learning\.arXiv preprint arXiv:2101\.05265\.Cited by:[§2](https://arxiv.org/html/2608.06595#S2.SS0.SSS0.Px3.p1.1)\.
- V\. M\. M\. Alvarez, R\. Roşca, and C\. G\. Fălcuţescu \(2020\)Dynode: neural ordinary differential equations for dynamics modeling in continuous control\.arXiv preprint arXiv:2009\.04278\.Cited by:[§2](https://arxiv.org/html/2608.06595#S2.SS0.SSS0.Px2.p1.1)\.
- I\. A\. Auzina, Ç\. Yıldız, S\. Magliacane, M\. Bethge, and E\. Gavves \(2023\)Modulated neural odes\.Advances in Neural Information Processing Systems36,pp\. 44572–44594\.Cited by:[§2](https://arxiv.org/html/2608.06595#S2.SS0.SSS0.Px2.p1.1)\.
- I\. Balázs, P\. Getto, and G\. Röst \(2021\)A continuous semiflow on a space of lipschitz functions for a differential equation with state\-dependent delay from cell biology\.Journal of Differential Equations304,pp\. 73–101\.Cited by:[§2](https://arxiv.org/html/2608.06595#S2.SS0.SSS0.Px1.p1.1)\.
- M\. G\. Bellemare, Y\. Naddaf, J\. Veness, and M\. Bowling \(2013\)The arcade learning environment: an evaluation platform for general agents\.Journal of artificial intelligence research47,pp\. 253–279\.Cited by:[§5](https://arxiv.org/html/2608.06595#S5.p1.1)\.
- R\. T\. Chen, Y\. Rubanova, J\. Bettencourt, and D\. K\. Duvenaud \(2018\)Neural ordinary differential equations\.Advances in neural information processing systems31\.Cited by:[§1](https://arxiv.org/html/2608.06595#S1.p3.1),[§2](https://arxiv.org/html/2608.06595#S2.SS0.SSS0.Px1.p1.1),[§3\.4](https://arxiv.org/html/2608.06595#S3.SS4.p1.3),[§5\.1](https://arxiv.org/html/2608.06595#S5.SS1.SSS0.Px1.p1.8)\.
- M\. Chevalier\-Boisvert, B\. Dai, M\. Towers, R\. de Lazcano, L\. Willems, S\. Lahlou, S\. Pal, P\. S\. Castro, and J\. Terry \(2023\)Minigrid & miniworld: modular & customizable reinforcement learning environments for goal\-oriented tasks\.CoRRabs/2306\.13831\.Cited by:[§5\.3](https://arxiv.org/html/2608.06595#S5.SS3.p1.1)\.
- E\. A\. Coddington and N\. Levinson \(1955\)Theory of ordinary differential equations\.McGraw\-Hill New York\.Cited by:[§1](https://arxiv.org/html/2608.06595#S1.p3.1)\.
- P\. Dhariwal, C\. Hesse, O\. Klimov, A\. Nichol, M\. Plappert, A\. Radford, J\. Schulman, S\. Sidor, Y\. Wu, and P\. Zhokhov \(2017\)OpenAI baselines\.GitHub\.Note:[https://github\.com/openai/baselines](https://github.com/openai/baselines)Cited by:[§5\.3](https://arxiv.org/html/2608.06595#S5.SS3.p1.1)\.
- J\. Du, J\. Futoma, and F\. Doshi\-Velez \(2020\)Model\-based reinforcement learning for semi\-markov decision processes with neural odes\.Advances in Neural Information Processing Systems33,pp\. 19805–19816\.Cited by:[§1](https://arxiv.org/html/2608.06595#S1.p4.1),[§2](https://arxiv.org/html/2608.06595#S2.SS0.SSS0.Px2.p1.1)\.
- T\. Haarnoja, A\. Zhou, P\. Abbeel, and S\. Levine \(2018\)Soft actor\-critic: off\-policy maximum entropy deep reinforcement learning with a stochastic actor\.InInternational conference on machine learning,pp\. 1861–1870\.Cited by:[§5](https://arxiv.org/html/2608.06595#S5.p1.1)\.
- E\. Haber and L\. Ruthotto \(2017\)Stable architectures for deep neural networks\.Inverse problems34\(1\),pp\. 014004\.Cited by:[§2](https://arxiv.org/html/2608.06595#S2.SS0.SSS0.Px1.p1.1)\.
- K\. He, X\. Zhang, S\. Ren, and J\. Sun \(2016\)Deep residual learning for image recognition\.InProceedings of the IEEE conference on computer vision and pattern recognition,pp\. 770–778\.Cited by:[§2](https://arxiv.org/html/2608.06595#S2.SS0.SSS0.Px1.p1.1)\.
- J\. Jia and A\. R\. Benson \(2019\)Neural jump stochastic differential equations\.Advances in Neural Information Processing Systems32\.Cited by:[§1](https://arxiv.org/html/2608.06595#S1.p4.1),[§2](https://arxiv.org/html/2608.06595#S2.SS0.SSS0.Px2.p1.1)\.
- P\. Kidger \(2021\)On Neural Differential Equations\.Ph\.D\. Thesis,University of Oxford\.Cited by:[§3\.4](https://arxiv.org/html/2608.06595#S3.SS4.p1.3)\.
- Y\. LeCun, B\. Boser, J\. S\. Denker, D\. Henderson, R\. E\. Howard, W\. Hubbard, and L\. D\. Jackel \(1989\)Backpropagation applied to handwritten zip code recognition\.Neural computation1\(4\),pp\. 541–551\.Cited by:[§1](https://arxiv.org/html/2608.06595#S1.p1.1)\.
- Y\. Lu, A\. Zhong, Q\. Li, and B\. Dong \(2018\)Beyond finite layer neural networks: bridging deep architectures and numerical differential equations\.InInternational conference on machine learning,pp\. 3276–3285\.Cited by:[§2](https://arxiv.org/html/2608.06595#S2.SS0.SSS0.Px1.p1.1)\.
- V\. Mnih, A\. P\. Badia, M\. Mirza, A\. Graves, T\. Lillicrap, T\. Harley, D\. Silver, and K\. Kavukcuoglu \(2016\)Asynchronous methods for deep reinforcement learning\.InInternational conference on machine learning,pp\. 1928–1937\.Cited by:[§3\.3](https://arxiv.org/html/2608.06595#S3.SS3.p1.3)\.
- V\. Mnih, K\. Kavukcuoglu, D\. Silver, A\. Graves, I\. Antonoglou, D\. Wierstra, and M\. Riedmiller \(2013\)Playing atari with deep reinforcement learning\.arXiv preprint arXiv:1312\.5602\.Cited by:[§6](https://arxiv.org/html/2608.06595#S6.SS0.SSS0.Px3.p1.1)\.
- V\. Mnih, K\. Kavukcuoglu, D\. Silver, A\. A\. Rusu, J\. Veness, M\. G\. Bellemare, A\. Graves, M\. Riedmiller, A\. K\. Fidjeland, G\. Ostrovski,et al\.\(2015\)Human\-level control through deep reinforcement learning\.nature518\(7540\),pp\. 529–533\.Cited by:[§5\.1](https://arxiv.org/html/2608.06595#S5.SS1.SSS0.Px1.p1.8)\.
- M\. Okada and T\. Taniguchi \(2021\)Dreaming: model\-based reinforcement learning by latent imagination without reconstruction\.In2021 ieee international conference on robotics and automation \(icra\),pp\. 4209–4215\.Cited by:[§6](https://arxiv.org/html/2608.06595#S6.SS0.SSS0.Px3.p1.1)\.
- R\. Pascanu, T\. Mikolov, and Y\. Bengio \(2012\)Understanding the exploding gradient problem\.CoRR, abs/1211\.50632\(417\),pp\. 1\.Cited by:[§5\.1](https://arxiv.org/html/2608.06595#S5.SS1.SSS0.Px1.p1.8)\.
- A\. Raffin, A\. Hill, A\. Gleave, A\. Kanervisto, M\. Ernestus, and N\. Dormann \(2021\)Stable\-baselines3: reliable reinforcement learning implementations\.Journal of Machine Learning Research22\(268\),pp\. 1–8\.External Links:[Link](http://jmlr.org/papers/v22/20-1364.html)Cited by:[§5](https://arxiv.org/html/2608.06595#S5.p1.1)\.
- Y\. Rubanova, R\. T\. Chen, and D\. K\. Duvenaud \(2019\)Latent ordinary differential equations for irregularly\-sampled time series\.Advances in neural information processing systems32\.Cited by:[§1](https://arxiv.org/html/2608.06595#S1.p4.1)\.
- S\. Ruder \(2016\)An overview of gradient descent optimization algorithms\.arXiv preprint arXiv:1609\.04747\.Cited by:[§5\.1](https://arxiv.org/html/2608.06595#S5.SS1.SSS0.Px1.p1.8)\.
- J\. Schulman, F\. Wolski, P\. Dhariwal, A\. Radford, and O\. Klimov \(2017\)Proximal policy optimization algorithms\.arXiv preprint arXiv:1707\.06347\.Cited by:[§5\.3](https://arxiv.org/html/2608.06595#S5.SS3.p1.1),[§5](https://arxiv.org/html/2608.06595#S5.p1.1)\.
- M\. Schwarzer, A\. Anand, R\. Goel, R\. D\. Hjelm, A\. Courville, and P\. Bachman \(2020\)Data\-efficient reinforcement learning with self\-predictive representations\.arXiv preprint arXiv:2007\.05929\.Cited by:[§2](https://arxiv.org/html/2608.06595#S2.SS0.SSS0.Px3.p1.1)\.
- R\. S\. Sutton, D\. McAllester, S\. Singh, and Y\. Mansour \(1999\)Policy gradient methods for reinforcement learning with function approximation\.InAdvances in Neural Information Processing Systems,S\. Solla, T\. Leen, and K\. Müller \(Eds\.\),Vol\.12,pp\.\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/1999/file/464d828b85b0bed98e80ade0a5c43b0f-Paper.pdf)Cited by:[§3\.2](https://arxiv.org/html/2608.06595#S3.SS2.p1.1)\.
- R\. Zheng, X\. Wang, Y\. Sun, S\. Ma, J\. Zhao, H\. Xu, H\. Daumé III, and F\. Huang \(2023\)TACO: temporal latent action\-driven contrastive loss for visual reinforcement learning\.Advances in Neural Information Processing Systems36,pp\. 48203–48225\.Cited by:[§2](https://arxiv.org/html/2608.06595#S2.SS0.SSS0.Px3.p1.1)\.

## Appendix AAppendix

### A\.1Learning Curves on all Environments

![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Qbert.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Alien.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Riverraid.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Breakout.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Asterix.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/DemonAttack.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/DoubleDunk.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/MsPacman.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/sep/BeamRider.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Tennis.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/Atlantis.png)

Figure 5:Episodic rewards of baseline and flow\-regularized A2C on all 11 Atari environments with a rolling average window of 100 episodes\.

## Appendix BFlowReg Configurations and Runtime

Table 4:FlowReg configurations used for each environment and their corresponding runtimes\.EnvironmentTimeRel\. UpdateA2CA2C\+FlowRegRuntimeSamplingFrequencyRuntime \(min\.\)Runtime \(min\.\)Overhead \(%\)DemonAttackExp\-Decay10487\.37487\.37554\.00554\.0013\.6713\.67AtlantisExp\-Decay10603\.44603\.44652\.00652\.008\.058\.05BeamRiderExp\-Decay20561\.82561\.82567\.00567\.000\.920\.92TennisExp\-Decay20617\.68617\.68621\.00621\.000\.540\.54RiverraidExp\-Decay5632\.40632\.40699\.00699\.0010\.5310\.53AsterixExp\-Decay5414\.66414\.66478\.02478\.0215\.2815\.28MsPacmanExp\-Decay5538\.13538\.13629\.70629\.7017\.0217\.02QbertIndex5510\.13510\.13565\.70565\.7010\.8910\.89BreakoutIndex5775\.07775\.07791\.00791\.002\.052\.05DoubleDunkIndex51011\.861011\.861202\.001202\.0018\.7918\.79AlienIndex5691\.99691\.99811\.00811\.0017\.2017\.20

![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/stats/runtime.png)Figure 6:Total Training Runtime Comparison \(for 10M timesteps\)\.
## Appendix CHyperparameter Tuning Experiments

![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/lambda/Qbert.png)
![Refer to caption](https://arxiv.org/html/2608.06595v1/images/plots/lcurves/lambda/DemonAttack.png)

Figure 7:Performance of different FlowReg loss weights \(λ\\lambda\)\.