Flowing Through States: Neural ODE Regularization for Reinforcement Learning
Authors: Mohamed Ghanem
Bernd Finkbeiner
Venue: ICLR 2026 Poster (2026)
Keywords: Neural ODE
Reinforcement Learning
MDP
Regularization
Actor-Critic
A2C
PPO
Atari
Minigrid
Primary area: reinforcement learning
TL;DR: We regularize reinforcement learning by modeling latent state transitions in MDPs as neural ODE flows, leading to improved stability and performance across standard benchmarks.

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.

FLOWING THROUGH STATES: NEURAL ODE REGU-
LARIZATION FOR REINFORCEMENT LEARNING
Mohamed Ghanem' Bernd Finkbeiner
CCIPPA Helmholtz Center for Information Security
Technical University of Munich
mohamed, ghanem@cispa de, Finkbeinerdcispa.co
ABSTRACT
Neural networks applied to sequential decision-making tasks typically rely on la- tent representations of enviromment states. While environment dynamics dictate how semantic states evolve, the corresponding latent transitions are usually left implicit, creating potential misalignment between the two. We propose to model latent dynamics explicitly by drawing analogy between Markov decision pro- cess (MDP) trajectories and ordinary differential equation (ODE) flows: in both cases, the current state fully determines its successors. Building on this view, we introduce neural ODE-based regularization method that enforces latent em- beddings follow consistent ODE flows, thereby aligning representation learning with environment dynamics. Although broadly applicable to deep learning agents, we demonstrate its effectiveness reinforcement learning by integrating it into Actor-Criie algorithms. Our approach yields major performance gains across various standard Atari benchmarks for A2C and gridworld environments for PPO.
INTRODUCTION
A central challenge in machine learning is bridging the gap between an object's '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 (LeCun al. 1989) incorporate inductive biases such as translation equivariance, spatial locality, and approximate in- variance to scale and rotation. These architectural choices encode object-level regularities, ensuring that embeddings reflect structural properties intrinsic individual objects.
While such local representations are powerful for perception tasks, sequential decision-making in- troduces a different challenge: the need for 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 transition rule R connects state 81 82, then their embeddings should satisfy a relation of the form:
h(s2) = g(h(s)),
where h(-) denotes the embedding function, and is an arbitrary function. While the existence of such a mapping 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 smooth path the latent space. We argue that regularizing latent embeddings to respect this path structure captures inherent property of transition dynamics, and enhances the model's 's ability to learn the task on a more global level. To operationalize this idea, we define la- tent flows using neural ordinary differential equations (neural ODEs) (Chen al.) 2018), which guarantee unique continuous trajectories under mild regularity assumptions such as Lipschitz con- tinuity (Coddington & Levinson 1955). In reasoning contexts, this uniqueness naturally subsumes

---

the Markov property: initial condition (i.e., state) completely determines the flow path of sub- sequent conditions.
However, directly using neural ODEs for inference is impractical: their reliance numerical in- tegration 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 (Du et al.) 2020 Jia & Benson 2019 Rubanova et al.) 2019). To overcome these limitations, propose train the agent's semantic embedder mimic the flows of 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 infer- ence. 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 neural ODE that learns to model the latent agent-envvrrnment dynamics in an unsupervised fashion.
The relevance of this perspective is particularly pronounced in discrete-stale 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 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 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 reg- ularization technique for sequential Markov decision-mail 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 a latent surrogate for the environment and aligning its flows with the la- tent trajectories of the agent's '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 result of tow-regulatization. Lastly, we also show the FlowReg boost to PPO on gridworld environments.
2 RELATED WORK
Neural ODEs continuous-depth networks. It has been noted in several existing works that ResNets (He et al. 2016) can be viewed as an Euler discretization of continuous differential flow
(Balázs al. 2021 Lu et al.] 2018 Haber & Ruthotto, (017). An implication of this is that an ODE can, in theory, be used to model an infinite-depth ResNet with finite number of parameters making them more parameter efficient (Chen al., 2018). In this paper, we take a broader view of sequence transformations modeled by the whole network as an embedder, rather than transforma- tions modeled by the individual layers within the model. That is. instead of looking at the embedder network 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 & Benson] 2019: Auzina et al., 2023). 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 of Alvarez et al. (2020) bears partial resemblance to ours in that it involves training an ODE to learn entire trajectories of continuous- space environments. However, both works fundamen- tally 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
2

---

tasks since the network's 's output is continuous. Similar to Du et al. (2020), they 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 trajecto- ries requires moving beyond static state discriminators to objectives that model long-horizon dy- namics. 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 segh-freuuency noise (Agarwal et al. 2021; Schwarzer et al., 2020). Meth- ods like TACO (Zheng al. 2023) enforce 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 stricter no- tion of temporal consistency by leveraging the uniqueness of ODE flows at any intermediate point, ensuring that states are predictable given any of their predecessors, not only the immediate ones.
PRELIMINARIES
3.1 MARKOV DECISION PROCESSES
We model reinforcement learning (RI.) problems as Markov decision processes (MDPs), defined by the tuple
= (,r,, where S is the state space, AA the action space, P(s' a) the transition kernel, r(s, a) the expected immediate reward, and E [0,1] discount factor. An agent samples actions a, E A according to
(1)
policy π (a 8), inducing trajectory T (8o, ao, ... The objective is to maximize the expected return
((T) = E
Σ y (r(s),)))
(2)
We define the following key functions:
The state-value function: V*(s) = E, (
The action-value function: (s, E, EE ('r(se, = =, 00 a
The advantage function: (s (s, (*(
3.2 POLICY GRADIENT METHODS
Policy gradient algorithms directly optimize parametric policy Ta(a | s). The policy gradient theorem (Sutton et al. 1999) states:
(o)( = Eg amom Vo ma
(3)
where de denotes the stationary state distribution under no- In practice, QT is approximated and variance is reduced by subtracting a baseline such as 1 (s).
3.3 ADVANTAGE ACTOR-CRITE (A2C)
Actor-critic methods (Mnih et al. 2016) couple policy model (the actor) with a value function estimator (the critic). The actor updates its parameters 0 via the policy gradient, while the critic learns to estimate V (T(s) (or QT a)) using temporal-difference learning.
The Advantage Acior-Critic (A2C) algorithm improves stability by using an advantage estimator. The policy gradient update is given by
VØJ(mT) Rs V log o (8)) Ag
(4)
with empirical advantage
A, rt + Vo( Vo(s.)
(5)
3

---

where Vo is the critic parameterized by 8. The critic is trained by minimizing the squared enror
Leritie (0) = -Es (re + +er++%Vlle (Vo(8++1) - vo(s)))
(6)
Lactor (0) = - Ea log Ta(aT |s) Ag
(7)
3.4 NEURAL ORDINARY DIFFERENTIAL EQUATIONS
Neural Ordinary Differential Equation is defined by the continuous transformation of the hidden state h(t) given by the differential equation:
dh(t) (b(t),,, dt
h(t) h(to) [(
(8)
where f is a neural network parameterized by 6. As such, neural ODEs differs from classical deep learning in that the neural network is used model the system dynamics (through the state derivative) at a given time instead of modeling the entire system directly. This framework can be used model functions that evolve over time. To seamlessly integrate neural ODEs into traditional deep learning pipeline, a differentiable numeric solver (e.g.. TORCHDIFFEQ (Chen et al. 2018) or DIFFRAX (Kidger 202I)) is typically used to evaluate the latent state function at given time points. The continuous-depth nature of Neural ODEs allows adaptive computation (c.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 frame- work 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 strug- gle with stiff systems, while implicit methods (e.g.. backward differentiation) enhance stability at higher computational cost. This interplay between numerical precision, stability, and efficiency un- derscores the importance of solver selection in practice. Additionally, Neural ODEs enable novel architectures, such as continuous normalizing flows for density estimation, by enforcing invertibil- ity through Lipschitz constraints on f. By bridging deep learning with differential equations, they provide principled framework for understanding neural networks dynamical systems, opening avenues for interpretability and integration with scientific machine learning.
APPROACH
In this section, we outline the mathematical formulation of our flow regularization technique for general target model. As illustrated in Figure | our setting involves three principal fields: (1) the semantic state field defined by the environment, (2) the latent observation vector field induced by the environment, and where each point is vector representation
the semantic state embedder 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), get the best of both worlds: latent field that captures local (state-level) and global (lajecttory-eveel aspects of the environment.
4.1 MODEL SETUP
Generally, there are two models involved in our framework, namely a target agent model 0 and flow regularizer model é. The target model comprises state embedder network he that con- verts semantic states into their latents, and a downstream head Fg that produces the final task- related actions. For state trajectory Hs(s) = he (s0), (s))), ..., he(sN-1), while flow embeddings are obtained by solving the initial value problem on (0) (s())):
So, Si,..., SN-1. semantic embeddings are computed as
Hs(s) (ho (st) ho((())
4

---

11
FR
R3
Semantic Observations
Fi2
$2
ho(s;)
(as)))
Latent Observations
ha (5o)
ho(ss)
Lnow
ODESolve($,
(s1)
h,(#s)
Latent Flow
h, (5o)
(52)
Figure 1: Illustration of the flow regularization landscape.
H (hp(s:)) ODESolve(So: (o())), (Tr)
(10)
where is the integration time index for state and fg is neural network that parameterizes the derivative 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 time index. Note that due to the Markov property, the underlying ODE is autonomous (i.c., time-invariant). However, the choice of the integration times still significantly influences the ODE solver, and experiments show that indeed fairly consequential for performance. An intuitive option for time sampling would be the i. Another simple approach is using discounted time horizon
step index of the state, i.e., with the same discounting factor y used by the agent's 's algorithm, i.e., т; = 7 where 0 < 7 < 1. This guarantees that integration times are in [0, 1] avoid arbitrarily large integration times, which might lead to gradient instability.
4.2 PATH ALIGNMENT
In essence, the flow model defines smooth latent path that starts a given semantic state embed- ding 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 topologi- cal structure of the latent space has to be implicitly learned over the course of the training. The key idea here that we can speed up this process by imposing a topological structure that we already know be compatible with the domain.
Our approach proceeds from the rationale that initially, the flow model carries pure curvature infor- mation while the semantic embedder carries task information. Ideally, 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 straightfor- ward way to incentivize this alignment is by minimizing the MSE between the latent point sequence Ha and the sampled flow path H As such,
compute the flow regularization loss follows:
Luow (s) H)()) H (s)13
(FlowReg)
(11)
N
5

---

4.3 OVERALL TRAINING OBJECTIVE
Having computed the flow loss the latent trajectory, this loss is then added to the label-based task loss:
L(s, v) = Liask (Fo (Ho(()))) + AL
(12)
where À is the flow-loss weighting factor. Note that Lnow (s) involves both the semantic embedder and the neural ODE network 6. This trains 0 to follow the continuous ODE flow while optimizing & to indirectly adapt to the underlying task modeled by B.
For an Advantage Actor-Critic agent, the overall training loss would be:
[(s, y) = Lactor (s, y) + B (s, v) + AC
(13)
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 training- waining-timme adaptive regularizer.
EXPERIMENTS
We evaluate our method on 11 Atari environments from the Arcade Learning Environment (ALE) library (Bellemare et al. 2013). This is mainly due to A2C being a reasonably simple actor-critic formulation, which is comerstone for many state-of-the-ar algorithms like PPO (Schulman al. 2017) and SAC (Haarnoja et al. 2018). We build on the Stable-baselines3 A2C implementation (Raffin al. 2021) incorporate eters for all environments and agents. The agent networks for both baseline and flow-regularized variants 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.
regularization loss. We the same set of A2C hyperparam-
5.1
ATARI BENCHMARKS
Obart
Allen
Riverald
Breakoot
300
www. pauale
500
Anturia
Damooddttack
Allordec
Beardilder
Ma www.
Figure 2: Episodic rewards of baseline and flow-regularized A2C on 8 different Atari environments with rolling average window of 100 episodes.

---

Hyperparameters. We performed 5 independent rums for every RL agent across all environments for 10 million timesteps each. Our semantic embedder for both baseline and feow-regularized agents is commonly used Nature CNN (Mnih al., 2015) feature extractor that embeds game state (frames) into 512-dimensional vector space. The ODE flow (and loss) is computed on the extracted state feature vectors. For the FlowReg ODE network, tion on the first layer. All models are optimized by RMSProp (Ruder) 2016) with an initial learning rate of 7 10 and linear decay scheduler. We apply a global-nonn -nom gradient clipping ratio of 0.5 (Pascanu et al. 2012). We use the TORCHDIFFEQ (Chen et al. 2018) library together with PY- TORCH for solving neural ODEs with relative tolerance =10 4 and absolute tolerance =10- For FlowReg variants, we experiment with both index-based (7: i) and exponential decay (T; = 7²) time sampling, along with a regularization frequency (relative to agent updates) of (5, 10, 20), and take the best configuration averaged over 3 seeds dedicated for hyperparameter search and separate from the 10 seeds the final comparison runs. For simplicity, we set
two-layer MLP with tanh activa-
for all environments.
Flow-regularized agents consistently outperform the baseline on Atari environments. Figure 2 highlights the notable performance gap between flow-regularized A2C and the baseline. The learning curves on all 11 environments can found in Figure 5 (Appendix A Figure [ shows the overall performance percent gains achieved by applying FlowReg all 11 environment 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.
Performance
Runtime Overhoad (%)
11:99
15.80
4X0%
200L
Figure 3: Trade-off 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 FlowReg largely improves the baseline performance under both Index and Exp-Decay time sampling modes. The choice between them, in all likelihood, depends on the granularity of the environment dynamics. We generally expect Exp-Decay to work better on environments with swifter or more fine-grained state transitions. Table (Appendix BB 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. Index is where T; = i and Exp-Decay is where T; 7
A2C AGENT
QBERT
RIVERRAID
BEAMRIDER
BASE
4374.30 + 958.95884 42
1862.27 +: 2399.58
960.66 + 748.23
FLOWREG (INDEX)
8306.00 05 1752.71
2946.34 1: 2788.17
1590.96 1 1033.30
FLOWREG (EXP-DECAY)
6903. 15 7157.71
2947. .95 1 2798 64
1593 11 1 961.77
'The hatched strip in Figure 3 indicates values exceeding the y-axis limit, which was capped for visual clarity avoid overly downscaling other values.
7

---

Table 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-m means the FlowReg loss applied once every m agent updates.
A2C AGENT
QBERT (INDEX)
QBERT SEXP-DECAY
4374.30
4374.30 + 958.42
FLOWREG U-5
8306.05 +: 1752.71 71
3286.60 1269.76
FLOWREG U-10
6569.51 # 2645.71
6903.15 15 + 2157. 71
FLOWREG U-20
5985.70 70 + 2756.17
6782.70 1877.13
FlowReg loss is still effective under much lower update frequency compared to the agent loss. Table [ 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 be aggressively optimized improve over the baseline, which allows it to run in comparable training time. By contrasting the time-overhead margins with the performance gains in Figure 3 shows that FlowReg an overall cost-effective choice. Figure and Table T (Appendix B) show the runtime comparison between the baseline and FlowReg in terms of absolute values.
Table 3: Latent path smoothness measures normalized by trajectory length.
METRIC FORMULA
NET
EXV
XX-1 [[ho(s))]
[ha(se-1) 1) ho(so)]
Σ
A2C
31.39 + 2.14
0.41 0.17
112475 521.76
1371.0 958.
QBERT
AZC+TACO
6.13 0.42
0.01
106.38 +
2434.05 2474.44
AZC+FLOWREG
4.20 0.44
0.10 + 0.02
64.17 + 7.05
8306.05 +: 1752.71
A2C
104.09 + 2.44
0.74 + 0.28
31432.50 + 1608 .82
10.40 = 1.86
BREAKOUT
A2C+TACO
13.09 +: 1.0S
0.13 1: 0.05
461.75 725.72
2.442
FLOWFLOWREG
4.92 0.23
0.06 + 0.02
94.98 9.51
24.03 0.84
A2C
75.36 + 2.72
0.53 - 0.07
18298. 55 1487 228
1862.27 2399.38
RIVERRAID
A2C+TACO
50.35 + 1.26
0.36 + 0.04
7599.32 + 404.50
2943.47 + 1616.30
AZC+FLOWREG
8.35 0.29
0.06 1 0.02
137.25 10.11
2947.90 95 2798.66 64
5.2 LATENT 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 result of the ODE alignment. We measure 3 different smoothness metrics shown in Table All 3 metrics are computed on the full dimensionality of the latent space without any reduction, and ||-I| is the Euclidean norm. To control for trajectory length variations, all 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: A Aha(s:) (s(42) 2ho(s:+1) + ho(si,, is more local measure roughness (lower is better).
FlowReg results in much smoother latent trajectories while improving overall performance. 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 envi- ronments. Naturally, we do not attribute the performance improvement solely the latent trajectory
8

---

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 dis- tinction 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 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 is that smoothness and temporal predictability are notably corre- lated. Despite the differences in mechanism between TACO and FlowReg, they both aim to instate notion of predictive temporal structure on the latent representations. The results of Table suggesst that this common feature explains the notable reduction in their latent path roughness compared to the baseline.
5.3 MINIGRID ENVIRONMENTS
We evaluate FlowReg on PPO (chhulman et 2017) in Minigrid environments (Chevalier-Boinz)
the purposes of showing FlowReg's efficacy another major RL allalooro algorithm (PPO) while also exploring a more radically discrete environment domain than Atari modified implementation of the ctt al. 2017). We use the Index U-20 FlowReg contiguration for all 3 environments. We performed 10 runs per agent for IM timesteps each.
al. 2023). These experiments
games. Similar to A2C,
PPO (Dhariwal
FourBoora
Dysanii-Gbabbcle
DonrKay
Figure 4: Episodic rewards of baseline and flow-regularized PPO on Minigrid environments with rolling average window of 100 episodes.
As shown in Figure flow-regularized PPO has clear advantage FourRooms and Dynamic- Obstacles while matching the baseline in DoorKey, where both agents practically solve the environ- ment.
CONCLUSION
Summary. In this paper, presented FlowReg, 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 discretiza- tions of a global neural ODE flow. We chose actor-critic reinforcement learning on Atari and Min- igrid environments have shown that using FlowReg notably boosts the overall performance of the target agent across almost all attempted environments and results in more constrained path structure on the learned embedding space.
showcase the benefits of applying FlowReg target model. Our results
Limitations. Although FlowReg does not require full episodes, still requires trajectory infor- mation 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 much emphasis episodic structure. A more fundamental limitation of FlowReg is the fact

---

that ODE flows 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 burden in environments where there are intermediate bottleneck states that need to be passed from different starting states, An example of that is 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 very large state space (like Atari).
Future Work. Since experiments demonstrate the efficacy of FlowReg on a standard on on-policy RL algorithm, it would be of great interest to see how it fares in the off off-policy settings such DQN (Mnih et al. 2013), as well as model-based algorithms like Dreamer (Okada & Taniguchi 202T). Although the scope of our evaluation pertains 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 very promising research direction.
REFERENCES
Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Mare G Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. arXiv preprint arXiv:2101.05265 2021.
Victor M Martinez Alvarez, Rares Rosca, and Cristian G Falcujescu. Dynode: Neural or- dinary differential equations for dynamics modeling in continuous control. arXiv preprint 0wX12 2020.
Ilze Amanda Auzina, Çagatay Yıldız, Sara Magliacane, Matthias Bethge, and Efstratios Gavves. Modulated neural odes. Advances in Neural Information Processing Systems, 36:44572-44594, 2023.
István Balázs, Philipp Getto, and Gergely Röst. A continuous semiflow on a space of lipschitz functions for differential equation with state-dependent delay from cell biology. Journal of Differential Equations, 304:73-101, 2021.
Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of artificial intelligence research, 47: 253-279, 2013.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Maxime Choiveier-Boisvertt, Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. CoRR, abs/2306. 13831, 2023.
Earl A Coddington and Norman Levinson. Theory of ordinary differential equations. McGraw-Hill New York, 1955.
Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, Yuhuai Wu, and Peter Zhokhov. Openai baselines. https: /basel nnss 2017.
//github. watoo//peeaa///aaeee
Jianzhun Du, Joseph Futoma, and Finale Doshi-Vedes Model-based reinforcement learning for semi-markov decision processes with neural odes. Advances in Neural Information Processing Systems, 33:19805-19886, 2020.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning, pp. 1861-1870. Pmlr, 2018.
Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse problems, 34(1):014004, 2017.
10

---

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Junteng Jia and Austin R Benson. Neural jump stochastic differential equations. Advances Neural Information Processing Systems, 32, 2019.
Patrick Kidger. On Neural Differential Equations. PhD thesis, University of Oxford, 2021.
Yann LeCun, Berhhard Boser. John S Denker, Donnie Henderson, Richard E Howard. Wayne Hub- bard, and Lawrence Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 1(4):541-551 1989.
Yiping Lu, Aoxiao Zhong, Quanzheng Li, and Bin Dong. Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations. In International conference on machine learning. pp. 3276-3285, PMLR, 2018.
Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv: 1312.560, 2013.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 877544)) 2015.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pp.
PmLR, 2016.
Masashi Okada and Tadahiro Taniguchi. Dreaming: Model-based reinforcement learning by la- tent imagination without reconstruction. In 2021 ieee international conference on robotics and automation (icra), pp.
IEEE, 2021.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. Understanding the exploding gradient prob- lem. CoRR, abs/1211.5063, 2,417):1, 2012.
Antonin Raflin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-8, 2021. URL httt:::/jjllr.rr//appp. 20 1364. html
Yulia Rubanova, Ricky TQ Chen, and David K Duvenaud. Latent ordinary differential equations for irreuulrly-samppled time series. Advances in neural information processing systems, 32, 2019.
Sebastian Ruder. An overview of gradient descent optimization algorithms. arXiv preprint arXiv:1609.04747, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06347, 2017.
Max Schwarzer, Ankesh Anand, Rishab Goel, Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self self-predictive representations, arXiv preprint arXiv:2007.05929 2020.
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. S. Solla, T. Leen, and K. Müller (eds.). Advances in Neural Information Processing Systems, volume 12. MIT Press, 1999. URL httpps://procceddggg.... www.:ys...uu.pp... wtt.ppps.ccppppp...ppp httpppaarrrreeepppppaaapeeeeeeee files//aper/1999/ Ei 4 E11e/336622055000ceeeeccccaaaa00
Ruijie Zheng, Xiyao Wang, Yanchao Sun, Shuang Ma, Jieyu Zhao, Huazhe Xu, Hal Daumé III, and Furong Huang. Taco: Temporal latent action-driven contrastive loss for visual reinforcement learning. Advances in Neural Information Processing Systems, 36:48203 48225, 2023.

---

A APPENDIX
A.1 LEARNING CURVES ON ALL ENVIRONMENTS
Demorstttack
DouthBurk
McPacinn
100
BaaeeRider
Figure 5: Episodic rewards of baseline and flow-regularized A2C on all 11 Atari environments with a rolling average window of 100 episodes.
12

---

B FLOWREG CONFIGURATIONS AND RUNTIME
Table 4: FlowReg configurations used for each environment and their corresponding runtimes.
TIME
REL. UPDATE
A2C
AZC+FLOWREG
RUNTIME
ENVIRONMENT
SAMPLING
FREQUENCY
RUNTIME (MIN.)
RUNTIME (MIN.)
OVERIEAD (%
DEMONAFTACK
EXP-DECAY
10
487.37
554.00
13.67
ATLANTIS
EXP-DECAY
10
603.44
632,00
8.05
BEAMRIDER
EXP-DECAY
561.82
567.00
0.92
TENNIS
EXP-DECAY
617.68
621.00
0.34
RIVERRAID
EXP-DECAY
632.40
609.00
10.53
ASTERIX
EXP-DECAY
414.66
478.02
15.28
MsPACMAN
EXP-DECAY
538,13
6229.70
17.02
QBERT
INDEX
510. 13
566.70
10.89
BREAKOUT
INDEX
775.07
791.00
2.05
1.86
1202.00
18.79
ALIEN
INDEX
691.99
811.00
17.20
AUC
000
Runtim
600
400
200
Almm
Almms
Reammaater
Breskott
Aromaed
Figure 6: Total Training Runtime Comparison (for 10M timesteps).
13

---

C HYPERPARAMETER TUNING EXPERIMENTS
Qbert
DemonAttack
A2C+FlowRag (A=1.0)
A2C-FlowReg (A-0.5)
8000
A2C-FlowReg (A=0.1)
A2C
0000
4000
8000
2000
2000
0.0
0.6
0.8
0.8
1.0
0.0
044 Timesteps
888
Timestape
Figure 7: Performance of different FlowReg loss weights (A).
14


Reviews:
[{"review_id": "wJQDyKA5UG", "reviewer": "Reviewer_RVQ7", "summary": "This paper proposes a regularization strategy for deep reinforcement learning that incentivizes latent state trajectories to reflect the semantic path of observations in the environment. This idea is implemented by introducing a neural ordinary differential equation (ODE) that is trained to reflect the behavior of the environment, and then comparing the ODE flows to the latent flows and minimizing the difference between their embeddings. The paper includes experiments on 10 Atari games, providing evidence that the flow-regularized model improves over the baseline non-regularized model.", "questions": "1. Can you expand on Limitations a bit more? I was having trouble following the second half of that paragraph, but it feels important. I would love more detail. Felt similar confusion about the second half of the \"Neural ODEs as continuous-depth networks\" paragraph. These parts were confusing.\n\n2. I also got a bit lost in lines 246--252.\n\n3. What's going on with MsPacman and Tennis? And why aren't those learning curves in the main text? They feel a bit buried in the appendix.\n\n4. \"We find that _Exp-Decay_ outperforms _Index_ more often than otherwise.\" What time sampling strategy do you use for Figs 2, 3, 4? Is it an environment-dependent mix or the single setting that performs best well across all environments?\n\n5. $\\lambda$ is an additional hyperparameter not mentioned in line 360; how is it selected?\n\n6. Atari environments are discrete, but they aren't _that_ discrete. What happens if you do this on something like a visual gridworld?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "This paper was a joy to read. What a cool idea. I was genuinely excited to learn something new. Extremely clear, and very interesting idea.\n\nThe problem that it points to in the introduction---states that are near each other in time ought to be near each other in latent space---is one that I have experienced firsthand. Until now, all the solutions to this problem that I had encountered (or tried myself) felt unsatisfying. This solution feels absolutely beautiful.", "weaknesses": "My main concerns come from the experimental validation.\n\nThe paper notes, \"We performed 5 independent runs for every RL agent\", and that \"we experiment with [different hyperparameter settings] and take the best configuration.\"\n\nRunning only 5 seeds is rather low. The results are likely significant despite having few seeds, because there are 8-10 environments and FlowReg performed better across the board. However, 10-20 seeds per environment would be better.\n\nMeanwhile, I'm wondering how the seeds fit into the experimentation process here. By \"and take the best configuration\"... does this mean you then _re-run_ the best configuration with _entirely new seeds_? (I hope so!) Or is this effectively taking the max over hyperparameter configurations without re-running on new seeds? If it's the latter, FlowReg could be winning solely due to selection bias and having more attempts (due to having more hyperparameter configurations) rather than because that particular hyperparameter configuration is actually better.\n\nI am prepared to increase my score if I can gain more confidence in the experimental validation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iKkBJdtClk", "reviewer": "Reviewer_Z2UD", "summary": "The paper proposes FlowReg, a neural ODE–based regularization technique for reinforcement learning that enforces smooth and dynamically consistent latent representations of environment states. The key idea is to treat Markov decision process (MDP) trajectories as discretized samples from continuous ODE flows and to align the latent embeddings of a policy’s state encoder with trajectories generated by a learned neural ODE. This regularizer introduces global structural constraints in the latent space without using the ODE for inference. Implemented on top of A2C, FlowReg improves performance across ten Atari environments, producing smoother latent trajectories and more stable training while adding only moderate computational overhead.", "questions": "- Can you provide runtime or FLOP comparisons to the A2C baseline to clarify FlowReg’s efficiency?\n\n- Could the alignment loss collapse diversity in latent representations, harming exploration?\n\n- Does FlowReg interact with representation learning methods like contrastive or predictive coding?\n\n- What happens if the ODE model is underparameterized or unstable—does it bias the policy?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Simple and general mechanism. FlowReg can be applied to most latent-state architectures with minimal modification—only adding an auxiliary ODE and an alignment loss—making it broadly usable beyond RL.\n\n- Empirical consistency. The method improves A2C on all tested Atari benchmarks and yields interpretable geometric effects (e.g., reduced latent path length and curvature).\n\n- Strong experimental hygiene. Results are averaged across runs, include multiple time-sampling and update-frequency settings, and discuss trade-offs in stability and runtime.", "weaknesses": "- No theoretical justification. The link between ODE smoothness and policy improvement is argued intuitively but not proven; there is no analysis of convergence, variance reduction, or representational bias.\n\n- Hyperparameter sensitivity. FlowReg introduces new parameters (update frequency, λ, time sampling) yet the paper gives limited guidance on tuning or robustness.\n\n- Computational overhead unquantified. The paper notes runtime remains “comparable” but lacks actual wall-clock comparisons; ODE solvers can be nontrivial in cost.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "beoFepLUij", "reviewer": "Reviewer_4aQj", "summary": "To align the representations in neural networks and environment dynamics in reinforcement learning, this paper proposes a new framework, FlowReg. It works 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.  Experiments show the effectiveness of the proposed method.", "questions": "1. Why should global (trajectory-level) aspects be considered in MDPs?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. It is an interesting idea to align the representations in neural networks and environmental dynamics.  I'm unaware of other work in this direction in reinforcement learning.\n\n2. In addition, the experimental results of A2C-FlowReg are indeed better than the experimental results of A2C.", "weaknesses": "1. This paper can be connected to Context Markov Decision Processes [1].\n2. There are many repeated subgraphs in Figures 2 and 4.\n3. The experiment is insufficient. The proposed FlowReg is only performed on one reinforcement learning algorithm, A2C.\n4. Lack of theoretical analysis. The author should theoretically prove that by the FlowReg framework, the latent transitions will better align with real state transitions.\n5. A2C does not need to be introduced in detail as it is not directly related to the new method in Section 4.\n6. The author should provide intuitive examples, theoretical analysis, and more metrics to explain why the return of existing algorithms can be improved by FlowReg.\n\n[1] Hallak, Assaf, Dotan Di Castro, and Shie Mannor. \"Contextual Markov decision processes.\" arXiv preprint arXiv:1502.02259 (2015).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Bf169qAOi3", "reviewer": "Reviewer_KaXh", "summary": "The work proposes a new regularization technique to facilitate learning of useful representations for reinforcement learning based on neural ODEs. Motivated by trajectories within MDPs to follow the dynamics of the transition function, the authors argue that latent representations learned by the agent should try to capture such dynamics. In addition to a encoder network that encodes states into a latent space, but might not capture the structure of the MDP transitions, they propose to learn a neural ODE model that captures the flow of latent embeddings throughout time. To encourage the latent space to follow the flow of the ODE model, they introduce an additional regularization loss term that minimizes the difference between the embeddings obtained by both models. The proposed FlowReg approach is evaluated on top of A2C in 10 Atari game environments and shown to improve sample efficiency. Additional analyzes shows the effect of different components of the ODE model, including time sampling approaches, update frequency, and shows that the latent space is smoother compared to without the regularization.", "questions": "1. What do you refer to with the term \"semantic space\" and how is it different from the latent space obtained from any encoder network? (Weakness 12)\n2. Function $g$ in the unnumbered equation within the introduction appears undefined. What does $g$ correspond to? (Weakness 15)\n3. How does FlowReg differ from existing auxiliary objectives to shape representations for sequential decision making (discussion of related literature above)? Does FlowReg have any unique advantages? (Weaknesses 1/2)\n4. The neural ODE process appears to be just conditioned on states and time. Is my understanding correct that the learning objective of this process is non-stationary since the flow of states throughout time depends on the policy -- that is learned and continually changes -- used to collect these trajectories? If so, do you find this to be a problem for the learning process of the ODE? (Weakness 4)\n5. Would you be able to show some qualitative analyzes and visualizations of the latent space learned by the ODE model and the latent space resulting from FlowReg? This might help to better understand what is being learned, and how it might be beneficial. (Weaknesses 5/6)\n6. How are the \"overall relative performance gains\" values reported in Figure 3 computed? (Weakness 8)\n7. What do lines and shading in Figure 2 correspond to? (Weakness 9)\n8. Do you have any insight into what makes FlowReg particularly effective in BeamRider and Qbert but less effective in Tennis? Are there any particular properties within the tasks with more significant gains that allow for the ODE model to be more precise? (Weakness 10)\n9. The authors state that they use the same hyperparameters for all environments and agents (Section 5). How were these hyperparameters determined? (Weakness 11)", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "The discussed problem of learning representations for sequential decision making that smoothly capture the dynamics of the environment is highly relevant and important. To the best of my knowledge, the application of neural ODEs to learn and regularize the flow of latent states is novel and original. The method is largely clearly defined in Section 4 and is conceptually fairly simple which I appreciate.\n\nExperiments in 10 Atari environments show that FlowReg consistently leads to improved performance with gains varying between tasks. There appears to be no task in which FlowReg harms performance which is encouraging. Additional ablations/ experiments demonstrate the robustness of the approach to different configurations and quantitatively analyzes the learned latent space in Table 3.\n\nOverall, I would consider this paper a largely well executed work, albeit with potentially limited significance and a lack of contextualization within related literature, as per weaknesses below.", "weaknesses": "In its current form, I am afraid that the work is not sufficient to justify acceptance at a venue like ICLR. Below, I try to give concrete weaknesses and highlight any weaknesses that I see as critical / major with (**Major**). I'd expect these to be addressed for this work to be considered for acceptance. \n\n## Originality & Significance\nAs mentioned above, I consider the combination of neural ODEs to model flow of latent states in MDPs as a regularization technique original and interesting. However, I am not convinced that it is a significant contribution to the field.\n\nIn its current form, this work does not at all acknowledge, discuss, or compare to alternative approaches of shaping representations for reinforcement learning. This is a rich literature space that I would highly suggest the authors to review and discuss in detail within their work in order to establish any potential unique features that their approach might offer. Almost all of these approaches consider the structure of a MDP in some form (temporal proximity, transition dynamics, ...) to shape representations that facilitate more efficient learning. The following might provide some starting points for such a review but is really just the tip of the iceberg:\n- Approaches that learn representations via constrastive learning\n\t- Uses contrastive learning to learn representations in which observations and noise-injected / augmented observations are close in latent space: Laskin, Michael, Aravind Srinivas, and Pieter Abbeel. \"Curl: Contrastive unsupervised representations for reinforcement learning.\" In _International conference on machine learning_, pp. 5639-5650. PMLR, 2020.\n\t- Use contrastive learning to learn similar representations for states and actions that are close in time: Zheng, Ruijie, Xiyao Wang, Yanchao Sun, Shuang Ma, Jieyu Zhao, Huazhe Xu, Hal Daumé III, and Furong Huang. \"TACO: Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning.\" _Advances in Neural Information Processing Systems_ 36 (2023): 48203-48225.\n- Approaches that learn representations leveraging temporal structure of transitions via predicting the future.\n\t- Schwarzer, Max, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron Courville, and Philip Bachman. \"Data-efficient reinforcement learning with self-predictive representations.\" In _International conference on learning representations_, 2021.\n\t- McInroe, Trevor, Lukas Schäfer, and Stefano V. Albrecht. \"Multi-horizon representations with hierarchical forward models for reinforcement learning.\" In _Transactions on Machine Learning Research_, 2024.\n- Approaches that learn similar representations for states that share certain properties according to bisimulation metrics\n\t- Agarwal, Rishabh, Marlos C. Machado, Pablo Samuel Castro, and Marc G. Bellemare. \"Contrastive behavioral similarity embeddings for generalization in reinforcement learning.\" In _International conference on learning representations_, 2021.\n\t- Castro, Pablo Samuel, Tyler Kastner, Prakash Panangaden, and Mark Rowland. \"MICo: Improved representations via sampling-based state similarity for Markov decision processes.\" _Advances in Neural Information Processing Systems_ 34 (2021): 30113-30126.\n- Approaches that leverage inverse dynamics models to learn representations for decision making\n\t- Pathak, Deepak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell. \"Curiosity-driven exploration by self-supervised prediction.\" In _International conference on machine learning_, pp. 2778-2787. PMLR, 2017.\n- ...\n\nThere are three main questions/ expectations within this that I'd all consider **major** weaknesses that need addressing:\n1. Given all these works in some form are motivated by the same question as this work -- how can we leverage the structure of the sequential problem of a MDP to learn better representations for decision making -- I would expect a detailed discussion that makes any differences in learned representations or assumptions clear between this work and prior work.\n2. Does FlowReg have any unique benefits over these alternatives? \n3. To establish any significance of this work in the context of the prior literature, I would expect comparisons of the effects and benefits of FlowReg with alternative recent representation learning approaches. And just to be clear, I am not stating that this method needs to be state-of-the-art on some benchmark to warrant publication, but given the similarity in objectives, there should be an empirical comparison to better understand potential differences.\n\n## Methodology\n4. **Major:** The proposed ODE model defined in Eq. (9, 10) appears to be purely state and time conditioned. That is somewhat confusing to me as transitions to future states depend on the actions being taken by the agent. In that sense, the function that the neural ODE is trying to learn changes whenever the policy changes which is constantly happening throughout training. Is my understanding correct that the objective of the neural ODE is non-stationary as it changes with the current policy generating trajectories? If so, do you find this to be a problem for the learning process of the ODE? What does the ODE tend to learn?\n\n## Experiments\n5. **Major**: To further substantiate weakness 3. above, I believe that comparisons to existing auxiliary representation learning objectives would be valuable. The main benefits and claims of FlowReg appears to be (1) improved performance/ learning efficiency and (2) smoothness of the latent space. However, approaches such as temporal contrastive auxiliary losses or future-predictive auxiliary objectives tend to also learn a latent space in which embeddings that are close in time are also close in embedding space, resulting in a latent states that evolve smoothly in time. Is FlowReg different in its learned latent space from these alternatives? To answer this question, I would expect to see some results for recent representatives of such alternative approaches that compares their latent space to FlowReg either qualitatively using visualisations or quantitatively using metrics as presented in Table 3.\n6. **Major**: The metrics presented in Table 3 to evaluate the smoothness of the latent space appear to be potentially uninformative. While a smaller path length and displacement might indicate a smoothly evolving latent space throughout the trajectory, it is also possible that the learned latent space merely compressed embeddings into a smaller space while maintaining identical smoothness and structure. For example, one could scale a latent space with an arbitrary small positive factor and obtain a new latent space that exhibits significantly lower path length and displacement. However, I would argue this latent space would be no smoother than the unscaled previous latent space. Such shortcuts can easily be learned through the use of regularization techniques as Eq. (12) merely encourages to minimize the distance of embeddings. That could likely be achieved by scaling down the latent space. To convincingly show the smoothness of the latent space, it would be helpful to visualize PCA/ t-SNE/ UMAP projections of the trajectories within latent space and see whether they are indeed more smooth compared to baselines.\n7. The novel contribution of this work is its ODE regularization process. Connected to questions in Weakness 4. Is there any qualitative analysis that you could provide that sheds insight into what the ODE tends to learn?\n8. Figure 3 states to show \"overall relative performance gains\" but it is unclear what exactly is being reported. Would the authors be able to clarify what they compute for this Figure?\n9. What do lines and shading in Figure 2 correspond to? I would have typically expected them to correspond to mean and standard deviation or confidence intervals but given the noisy lines with small shading for example in DemonAttack and Qbert, I would doubt that is the case.\n10. As per Figure 3, relative performance gains of FlowReg differ quite substantially across environments. Do you have any insight into what makes FlowReg particularly effective in BeamRider and Qbert but less effective in Tennis? Are there any particular properties within the tasks with more significant gains that allow for the ODE model to be more precise?\n11. The authors state that they use the same hyperparameters for all environments and agents (Section 5). How were these hyperparameters determined?\n\n## Clarity\n12. The work frequently uses the terms of \"semantic space\", \"semantic trajectories\", and \"semantic observations\" in contrast to latent space and information. However, it is not clear to me what exactly these terms are meant to express. It would be helpful if these terms are properly introduced and/ or defined.\n13. Within the introduction, the authors state that their method \"combines the expressivity of continuous-time dynamics with the efficiency of conventional neural architectures\" but it is unclear to me why modelling continuous-time dynamics within MDPs is valuable given the MDP itself operates on discrete steps.\n14. Within the introduction, the authors states \"the structure properties it [embedding function] imposes on the latent space -- such as smoothness, consistency, and determinism -- are far from trivial and are crucial for reasoning tasks.\" It is unclear to me why it is crucial for the embedding function to impose determinism. Do the authors refer to the encoding from states to latent embeddings to be deterministic, or the transitions between states to be deterministic within the latent space? Any clarification and substantiation of this claim would be appreciated.\n15. Function $g$ in the unnumbered equation within the introduction appears undefined. What does $g$ correspond to?\n16. Most parameterised functions are written with parameters as a subscript but in Eq. (8), you write $f(..., \\phi)$ instead of $f_\\phi(...)$. Also, the right hand equation of Eq. (8) also write $f$ but without any notation to indicate its parameters.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)