Flowing Through States: Neural ODE Regularization for Reinforcement Learning

2026-08-06

AI Summary

This work addresses the mismatch between latent state representations and the true environmental dynamics in reinforcement learning by introducing Neural Ordinary Differential Equations (Neural ODEs) into representation learning. By modeling trajectories of Markov Decision Processes as continuous ODE flows, the proposed approach explicitly constrains latent embeddings to evolve according to consistent, continuous dynamics. A novel, general-purpose regularization mechanism is developed to effectively align the dynamics in latent space with the underlying semantic dynamics of the environment. Empirical evaluations on Atari benchmarks using A2C and on grid-world environments with PPO demonstrate significant performance improvements, validating both the efficacy and generalization capability of the method.

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.

Problem

• latent dynamics
• environment dynamics
• representation alignment
• reinforcement learning
• Markov decision process

Innovation

• Neural ODE
• Reinforcement Learning
• Representation Learning
• Dynamics Alignment
• Regularization