Flowing Through States: Neural ODE Regularization for Reinforcement Learning
Mohamed Ghanem, Bernd Finkbeiner
cs.LG, cs.AI
Submitted: 2026-08-06
Updated: 2026-08-10
Code: https://github.com/openai/baselines
License: http://creativecommons.org/licenses/by/4.0/
The gist: Neural networks applied to sequential decision-making tasks typically rely on latent representations of environment states.
Terminology
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.
Sources
- Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning
- DyNODE: Neural Ordinary Differential Equations for Dynamics Modeling in Continuous Control
- Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Tasks
- Playing Atari with Deep Reinforcement Learning
- On the difficulty of training Recurrent Neural Networks
- An overview of gradient descent optimization algorithms
- Proximal Policy Optimization Algorithms
- Data-Efficient Reinforcement Learning with Self-Predictive Representations
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks