流经状态：用于强化学习的神经ODE正则化

# 在状态间流动：用于强化学习的神经 ODE 正则化

来源：https://arxiv.org/html/2608.06595
Mohamed Ghanem1 Bernd Finkbeiner1,2 1CISPA 亥姆霍兹信息安全中心 2慕尼黑工业大学 mohamed\.ghanem@cispa\.de, finkbeiner@cispa\.de

###### 摘要

应用于序列决策任务的神经网络通常依赖于环境状态的潜在表示。虽然环境动态决定了语义状态如何演化，但相应的潜在转换通常被隐式保留，从而在这两者之间产生潜在的错位。我们通过将马尔可夫决策过程（MDP）轨迹与常微分方程（ODE）流进行类比，提出显式建模潜在动态：在这两种情况下，当前状态完全决定其后继状态。基于这一观点，我们引入了一种基于神经 ODE 的正则化方法，迫使潜在嵌入遵循一致的 ODE 流，从而将表示学习与环境动态对齐。尽管该方法广泛适用于深度智能体，我们通过将其集成到 Actor-Critic 算法中，在强化学习中展示了其有效性。我们的方法在 11 个 Atari 基准环境（用于 A2C）和网格世界环境（用于 PPO）上均带来了显著的性能提升。

## 1 引言

机器学习中的一个核心挑战是弥合对象的语义含义与其潜在表示之间的差距。由于神经网络作用于学习到的嵌入而非直接语义，表示学习主要集中在设计能够忠实编码局部对象属性的过程上。例如，卷积神经网络（LeCun 等，1989 (https://arxiv.org/html/2608.06595#bib.bib40)）融入了诸如平移等变性、空间局部性以及对缩放和旋转的近似不变性等归纳偏置。这些架构选择编码了对象级的规律性，确保嵌入反映单个对象内在的结构属性。

虽然这种局部表示对于感知任务非常强大，*序列决策*引入了不同的挑战：需要更*全局*地理解对象和状态如何随时间相互关联。在这种情况下，相关的归纳偏置并非来自孤立对象，而是来自连接它们的动态。例如，在马尔可夫决策过程（MDP）的背景下，一个状态及其后继状态的潜在嵌入应当由转换动态一致地关联。具体来说，如果转换规则 R 将状态 s1 连接到 s2，那么它们的嵌入应满足如下形式的关系：

h(s2)=g(h(s1),R)，

其中 h(·) 表示嵌入函数，g 是任意函数。虽然原则上这种映射的存在是平凡的，但它在潜在空间上施加的结构属性（如平滑性、一致性和确定性）远非平凡，并且对于推理任务至关重要。

本文源于这样的直觉：*语义轨迹*的嵌入可以被理解为连续潜在流的离散化。换言之，语义空间中的每条轨迹都应对应于潜在空间中的一条平滑路径。我们主张，对潜在嵌入进行正则化以尊重这种路径结构，能够捕捉转换动态的固有属性，并增强模型在更全局层面上学习任务的能力。为了实现这一想法，我们使用神经常微分方程（神经 ODE）（Chen 等，2018 (https://arxiv.org/html/2608.06595#bib.bib7)）来定义潜在流，该方程在 Lipschitz 连续性等温和正则性假设下保证唯一的连续轨迹（Coddington and Levinson, 1955 (https://arxiv.org/html/2608.06595#bib.bib26)）。在推理背景下，这种唯一性自然蕴含了*马尔可夫性质*：初始条件（即状态）完全决定了后续条件的流动路径。

然而，直接使用神经 ODE 进行推理并不实用：它们对数值积分的依赖使其比标准前向传播慢得多，并且它们应用于序列推理时，还会因语义状态演化所引入的不连续性而进一步复杂化（Du 等，2020 (https://arxiv.org/html/2608.06595#bib.bib22)；Jia and Benson, 2019 (https://arxiv.org/html/2608.06595#bib.bib18)；Rubanova 等，2019 (https://arxiv.org/html/2608.06595#bib.bib23)）。为克服这些限制，我们提出训练智能体的语义嵌入器，通过对齐惩罚来*模仿*神经 ODE 的流。这种方法使学习到的嵌入能够继承平滑 ODE 流的拓扑结构，同时避免了基于 ODE 推理的计算和设计负担。因此，我们的方法结合了连续时间动态的表达能力与传统神经架构的高效性。此外，它以一种神经 ODE 的形式为智能体添加了一层全局指导，该神经 ODE 以无监督方式学习建模潜在的智能体-环境动态。

这一视角的相关性在*离散状态* MDP 中尤为突出。在连续状态环境中，状态空间固有的连续性自然会在潜在表示中引发平滑性：输入状态的微小变化通常对应嵌入的微小变化。相比之下，在离散领域中，语义空间由孤立状态组成，没有*先验*的邻近性或平滑转换概念。因此，必须在潜在空间中施加连续性，而不是从状态空间本身继承。将离散轨迹嵌入为平滑潜在流，因此提供了一种有原则的方法来恢复原本缺失的结构规律性，从而使潜在动态能够反映底层 MDP 的转换约束。

贡献。在本文中，我们引入了流正则化（FlowReg），一种用于序列马尔可夫决策模型的无监督正则化技术，它将智能体的潜在表示场与底层语义环境动态对齐。为此，它学习一个神经 ODE，作为环境的潜在代理，并将其流与智能体状态嵌入器的潜在轨迹对齐。为了展示我们的技术，我们在强化学习环境中评估了 FlowReg，即在 11 个 Atari 环境中使用 Advantage Actor-Critic（A2C）。我们的实验表明，FlowReg 在所有环境中显著提高了基线模型的性能。我们进一步检查了由此产生的潜在轨迹，并展示了流正则化所赋予的理想平滑性。最后，我们还展示了 FlowReg 在网格世界环境中对 PPO 的促进效果。

## 2 相关工作

#### 作为连续深度网络的神经 ODE。

已有若干工作指出，ResNet（He 等，2016 (https://arxiv.org/html/2608.06595#bib.bib14)）可以被视为连续微分流的欧拉离散化（Balázs 等，2021 (https://arxiv.org/html/2608.06595#bib.bib8)；Lu 等，2018 (https://arxiv.org/html/2608.06595#bib.bib10)；Haber and Ruthotto, 2017 (https://arxiv.org/html/2608.06595#bib.bib11)）。这意味着，理论上可以用 ODE 以有限数量的参数建模无限深度的 ResNet——使其更具参数效率（Chen 等，2018 (https://arxiv.org/html/2608.06595#bib.bib7)）。在本文中，我们采取更广阔的视角，将整个网络建模的序列变换视为嵌入器，而不是模型中单个层所建模的变换。也就是说，我们并不是将嵌入器网络视为对象的离散化变换，而是考察将网络应用于一系列在明确定义的环境动态下顺序相关的对象所产生的潜在轨迹。

#### 用于连续控制的神经 ODE。

神经 ODE 可以建模离散事件之间的连续演化，同时与事件触发机制或分类器耦合，以检测和处理突变转换，例如碰撞或控制模式变化（Jia and Benson, 2019 (https://arxiv.org/html/2608.06595#bib.bib18)；Auzina 等，2023 (https://arxiv.org/html/2608.06595#bib.bib19)）。通过整合传统神经网络，这些模型可以直接从数据中推断连续流以及离散切换的时机或条件，从而绕开僵化的解析公式。Alvarez 等人（2020 (https://arxiv.org/html/2608.06595#bib.bib41)）的工作与我们的部分相似，因为它涉及训练一个 ODE 来学习连续空间环境的完整轨迹。然而，这两项工作与我们的方法有根本区别：我们的神经 ODE 作用于潜在轨迹，而它们旨在预测语义轨迹，这使得它们很难应用于离散空间任务，因为网络的输出是连续的。与 Du 等人（2020 (https://arxiv.org/html/2608.06595#bib.bib22)）类似，它们将神经 ODE 作为主要推理模型，而我们仅将神经 ODE 用作解耦的正则化器。

#### 通过预测编码塑造表示。

增强跨轨迹的时间一致性，需要超越静态状态判别器，转向对长时程动态进行建模的目标函数。通过将预测编码与对比学习相结合，可以塑造表示，以最大化过去历史与未来结果之间的互信息，从而有效地使潜在空间对高频噪声平滑化（Agarwal 等，2021 (https://arxiv.org/html/2608.06595#bib.bib49)；Schwarzer 等，2020 (https://arxiv.org/html/2608.06595#bib.bib47)）。诸如 TACO（Zheng 等，2023 (https://arxiv.org/html/2608.06595#bib.bib46)）之类的方法在潜在空间中强制实现稳健的时间结构，其中状态转换可由其直接前驱预测，从而防止表示因与任务无关的环境随机性而漂移。我们的方法通过利用 ODE 流在任意中间点的唯一性，强制实现更严格的时间一致性概念，确保状态可由*任意*前驱（而不仅仅是直接前驱）预测。

## 3 预备知识

### 3.1 马尔可夫决策过程

我们将强化学习（RL）问题建模为*马尔可夫决策过程*（MDP），由元组定义

M=(S,A,P,r,γ)，(1)

其中 S 是状态空间，A 是动作空间，P(s′∣s,a) 是转移核，r(s,a) 是期望即时奖励，γ∈[0,1) 是折扣因子。智能体根据策略 π(a∣s) 采样动作 at∈A，从而产生轨迹 τ=(s0,a0,r0,...)。目标是最大化期望回报

J(π)=Eπ[∑t=0∞γtr(st,at)] (2)

我们定义以下关键函数：

- 状态值函数：Vπ(s)=Eπ[∑t=0∞γtr(st,at) | s0=s]
- 动作值函数：Qπ(s,a)=Eπ[∑t=0∞γtr(st,at) | s0=s, a0=a]
- 优势函数：Aπ(s,a)=Qπ(s,a)−Vπ(s)

### 3.2 策略梯度方法

策略梯度算法直接优化参数化策略 πθ(a∣s)。策略梯度定理（Sutton 等，1999 (https://arxiv.org/html/2608.06595#bib.bib35)）指出：

∇θJ(πθ)=Es∼dπθ,a∼πθ[∇θlogπθ(a∣s) Qπθ(s,a)] (3)

其中 dπθ 表示 πθ 下的平稳状态分布。在实践中，Qπθ 被近似，并通过减去诸如 Vπ(s) 之类的基线来降低方差。

### 3.3 优势 Actor-Critic（A2C）

Actor-Critic 方法（Mnih 等，2016 (https://arxiv.org/html/2608.06595#bib.bib28)）将策略模型（Actor）与值函数估计器（Critic）耦合。Actor 通过策略梯度更新参数 θ，而 Critic 则使用时序差分学习来估计 Vπ(s)（或 Qπ(s,a)）。

*优势 Actor-Critic（A2C）*算法通过使用优势估计器来提高稳定性。策略梯度更新为

∇θJ(πθ)≈E[∇θlogπθ(at∣st) Ât] (4)

其中经验优势为

Ât=rt+γVθ(st+1)−Vθ(st) (5)

其中 Vθ 是由 θ 参数化的 Critic。Critic 通过最小化平方误差来训练

Lcritic(θ)=Est∼πθ[(rt+γVθ(st+1)−Vθ(st))2] (6)

Lactor(θ)=−Est,at∼πθ[logπθ(at∣st) Ât] (7)

### 3.4 神经常微分方程

神经常微分方程由隐藏状态 h(t) 的连续变换定义，由微分方程给出：

dh(t)dt=fφ(h(t),t)， h(t)=h(t0)+∫t0t fφ(h(s),s) ds (8)

其中 f 是由 φ 参数化的神经网络。因此，神经 ODE 与传统深度学习的不同之处在于，神经网络用于建模系统在给定时间的动态（通过状态导数），而不是直接建模整个系统。该框架可用于建模随时间演化的函数。为了将神经 ODE 无缝集成到传统深度学习流程中，通常使用可微分的数值求解器（例如 torchdiffeq（Chen 等，2018 (https://arxiv.org/html/2608.06595#bib.bib7)）或 Diffrax（Kidger, 2021 (https://arxiv.org/html/2608.06595#bib.bib21)））在给定时间点评估潜在状态函数。神经 ODE 的连续深度特性允许自适应计算（例如，变化的求解器步长），与固定深度架构相比，提供了内存效率以及在精度和计算成本之间的灵活权衡。

神经 ODE 的一个关键数学性质是它们的可逆性以及通过伴随状态进行的精确梯度计算，这确保了即使在长积分区间内也能稳定训练。该框架天然地适应不规则采样的数据。