跳转至

Primal-Dual Representation Learning for Low-Rank Constrained MDPs

会议: NeurIPS 2026
arXiv: 2610.30756
论文: https://arxiv.org/abs/2610.30756
领域: 强化学习
关键词: 约束马尔可夫决策过程, 低秩MDP, 表示学习, 原对偶优化, 不确定性奖励

一句话总结

针对未知转移特征的连续状态低秩约束MDP,本文提出首个具备理论保证的原对偶表示学习算法体系,通过三阶段辅助探测与合并协方差奖励消除估计偏差,分别实现了软约束下的最优 \(\widetilde{O}(\sqrt{K})\) 遗憾与软违例界,以及硬约束下基于置信下界证书的严格累积违例控制。

研究背景与动机

约束马尔可夫决策过程(Constrained Markov Decision Processes, CMDPs)为兼顾奖励最大化与安全性约束的强化学习提供了坚实的数学理论框架。在机器人控制、自动驾驶与工业调度等安全攸关领域,智能体不仅要在未知环境中最大化累积收益,更必须严格满足安全性或预算约束。早期的可证明安全强化学习研究主要集中在表格型(tabular)CMDP,其样本复杂度与状态空间基数 \(|\mathcal{S}|\) 呈多项式依赖,无法推广至高维或连续状态空间;后续线性 CMDP 方法虽然通过低维特征假设消除了状态维度的显式依赖,但完全依赖于先验已知的环境潜在转移特征 \(\phi^\star(s,a)\),当环境特征未知且需要从交互数据中在线归纳时,现有理论保证完全失效。

在无约束连续状态强化学习中,以 FLAMBE、REP-UCB 为代表的低秩 MDP 表示学习方法证明了利用最大似然估计(MLE)能够高效发现潜在转移结构并驱动乐观探索。然而,将表示学习推广至 CMDP 存在深层次的理论瓶颈:表示学习的误差会同时污染奖励函数的价值估计与约束效用函数的合规性评估。为了学习未知特征,算法需要具备动作覆盖度的探测策略,但盲目探索会导致严重的安全性违规;反之,若严格拘泥于当前策略的占用测度,又会陷入动作覆盖不足、无法准确识别环境动态的死锁。更为关键的是,理论推导表明朴素的双层探测无法提供覆盖度转移所需的下界支撑,造成协方差矩阵与真实转移的不匹配。

面对这一表征覆盖与安全约束之间的根本冲突,本文旨在解决连续状态未知特征低秩 CMDP 的在线探索与约束控制难题。核心 idea:设计带有衰减 \(\epsilon\)-贪心动作探测的三阶段辅助转移采样流,通过最大似然估计学习潜在转移因子并构建合并经验协方差奖励,在软约束下结合乐观原对偶规划实现平衡探索,在硬约束下引入基于单边置信下界的效用证书与稳定化对偶搜索,在完全未知特征的连续环境下首次实现可证明的次线性遗憾与约束违例界。

方法详解

整体框架

算法围绕连续状态低秩 CMDP 的在线交互展开,核心目标是在未知转移函数 \(P^\star(s'|s,a) = \langle \phi^\star(s,a), \mu^\star(s') \rangle\) 的情况下,既能学得精确的转移特征表示,又能控制累积约束违例。针对不同的违规度量标准,文章提出了两个核心算法:面向软累积违例(soft violation,允许不同回合间赤字与盈余相互抵消)的 REP-PD-TEG,以及面向更严格的逐回合硬累积违例(hard violation,仅统计各回合正向赤字)的 REP-PD-H-TEG。

整体流程在每轮迭代中依次执行:(1) 结合上一轮返回策略与均匀动作先验构建衰减 \(\epsilon_k\)-贪心执行策略;(2) 通过随机预言机采集包含连续三步状态-动作转移的三阶段辅助元组(tertiary tuple),并分流存入三个经验重放池;(3) 跨三个数据流联合执行 MLE 拟合低秩转移特征与状态嵌入;(4) 基于主、次两级数据流构造合并经验协方差矩阵及其 Mahalanobis 不确定性奖励;(5) 在软约束下直接执行带有乐观奖励与效用加权拉格朗日目标函数的策略规划与投影对偶变量更新,或在硬约束下基于单边效用置信下界证书开展截断对偶网格搜索以选定安全返回策略。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["起始状态采样<br/>s ~ d(π_k-1)"] --> B["三阶段探测执行策略<br/>q_k = (1-ε_k)π_k-1 + ε_k U(A)"]
    B --> C["三级辅助转移收集<br/>生成 (s,a,s'), (s',a',s''), (s'',a'',s_bar) 分入 D_k, D'_k, D''_k"]
    C --> D["跨流联合最大似然估计<br/>拟合低秩转移特征 (μ_hat_k, φ_hat_k)"]
    D --> E["合并协方差不确定性奖励<br/>基于 D_k 与 D'_k 构建 Σ_hat_k 并计算 b_hat_k"]
    E -->|软约束分支| F["乐观拉格朗日原对偶更新<br/>REP-PD-TEG: 奖励/效用双重乐观 + 投影对偶步"]
    E -->|硬约束分支| G["下置信效用证书与截断对偶搜索<br/>REP-PD-H-TEG: U_bar_k(π) ≥ b 验证可行性"]
    F --> H["输出返回策略 π_k 并进入下一轮"]
    G --> H

关键设计

1. 三阶段辅助探测机制与跨流联合 MLE:破解两级探测协方差失配的理论障碍

在低秩表示学习分析中,若仅使用主探测样本 \(D_k\)(从当前策略占用分布采样)和次级一步转移样本 \(D'_k\),两级探测的特征覆盖率无法提供闭合的集中性不等式支撑。理论证明(Proposition D.1)指出,存在两状态对抗 MDP 使得主探测协方差矩阵在特定方向上的范数与合并协方差矩阵之比发散,导致标准的一步回退(one-step-back)悲观界失效。为此,本文设计了三阶段辅助探测序列(tertiary tuple):智能体在每一回合不仅执行两步转移,而是推进三步连续交互: $\(s \sim d_{P^\star}^{\pi_{k-1}}, \quad a \sim q_k(\cdot|s), \quad s' \sim P^\star(\cdot|s,a), \quad a' \sim q_k(\cdot|s'), \quad s'' \sim P^\star(\cdot|s',a'), \quad a'' \sim q_k(\cdot|s''), \quad \bar{s} \sim P^\star(\cdot|s'',a''),\)$ 其中探测策略为衰减 \(\epsilon\)-贪心混合策略 \(q_k = (1-\epsilon_k)\pi_{k-1} + \epsilon_k U(\mathcal{A})\),参数满足 \(\epsilon_k = k^{-p}\)。所得三组转移分别存入 \(D_k, D'_k, D''_k\)。随后,算法通过联合最大似然估计(MLE)求解低秩分解模型: $\(\widehat{P}_k = (\widehat{\mu}_k, \widehat{\phi}_k) \in \arg\max_{(\mu,\phi)\in\mathcal{M}} \mathbb{E}_{D_k + D'_k + D''_k} \left[\log \mu(x')^\top \phi(x,u)\right].\)$ 这一设计的精妙之处在于,第三级数据流 \(D''_k\) 为次级转移的推前(pushforward)分布提供了受到 MLE 控制的母体分母,从而严密控制了分布漂移引起的误差放大,彻底消除了两级探测的结构性障碍。

2. 合并经验协方差探索奖励:统一量化表示不确定性

在获得拟合特征 \(\widehat{\phi}_k\) 后,算法将主探测样本 \(D_k\) 与次级探测样本 \(D'_k\) 中的特征外积合并,构建合并经验协方差矩阵: $\(\widehat{\Sigma}_k^{\mathrm{mix}} = \sum_{(x,u,x')\in D_k} \widehat{\phi}_k(x,u)\widehat{\phi}_k(x,u)^\top + \sum_{(x,u,x')\in D'_k} \widehat{\phi}_k(x,u)\widehat{\phi}_k(x,u)^\top + \lambda_k I.\)$ 注意第三级数据流 \(D''_k\) 仅用于约束 MLE 误差界,并不直接加入协方差矩阵,以防止引起无限级联的递推传播。基于该矩阵,算法定义了截断 Mahalanobis 探索奖励: $\(\widehat{b}_k(x,u) = \min\left\{\alpha_k \sqrt{\widehat{\phi}_k(x,u)^\top (\widehat{\Sigma}_k^{\mathrm{mix}})^{-1} \widehat{\phi}_k(x,u)}, \; \bar{b}\right\},\)$ 其中正则化参数 \(\lambda_k\) 与奖励尺度 \(\alpha_k\) 经过精心配置,与动作空间大小 \(|\mathcal{A}|\)、探测衰减指数 \(k^p\) 以及模型复杂度 \(\ln(|\mathcal{M}|k/\delta)\) 紧密挂钩。该奖励项不仅衡量了模型对特定状态-动作对转移动态的认识不确定性,而且通过合理的参数缩放,确保了后续无论是对于奖励函数还是约束效用函数,都能在悲观与乐观分析中吸收由于潜在特征估计引起的预测偏差。

3. 乐观拉格朗日原对偶规划(REP-PD-TEG):面向软累积违例的最优权衡

在软约束场景下,由于不同回合之间的效用盈余允许补偿亏损,因此可以直接在拉格朗日函数中对效用估计施加乐观奖励。在第 \(k\) 轮,REP-PD-TEG 将奖励函数 \(r\) 与效用函数 \(g\) 分别通过探索奖励 \(\widehat{b}_k\) 进行乐观增广,并求解乐观复合动作价值函数的最优策略: $\(\pi_k \in \arg\max_{\pi} \left( Q^{\pi}_{\widehat{P}_k, r+\widehat{b}_k} + Y_k Q^{\pi}_{\widehat{P}_k, g+\widehat{b}_k} \right).\)$ 随后,拉格朗日乘子(对偶变量)按照次梯度下降方式进行投影更新: $\(Y_{k+1} = \Pi_{[0,\xi]} \left( Y_k + \eta \left( b - V^{\pi_k}_{\widehat{P}_k, g+\widehat{b}_k}(s_1) \right) \right),\)$ 其中截断阈值 \(\xi = 2/\theta\),\(\theta\) 为 Slater 条件给出的可行裕度。通过椭圆势能引理与原对偶性能差异分解,REP-PD-TEG 将返回策略(return policy)的遗憾界与软违例界收敛至 \(\widetilde{O}(K^{1/2+p})\)。特别地,当 \(p=0\)(即全程采用均匀辅助探测)时,返回策略直接达到理论最优的 \(\widetilde{O}(\sqrt{K})\) 速率;而若同时考核执行探测策略自身的代价,选取 \(p=1/4\) 则可达成兼顾探测损耗与策略性能的整体 \(\widetilde{O}(K^{3/4})\) 样本复杂度。

4. 单边置信下界效用证书与截断对偶搜索(REP-PD-H-TEG):实现逐回合严格硬约束控制

在硬约束场景中,每个回合的效用赤字必须逐个进行惩罚(\([b - V_g^{\pi_k}]_+\)),因此乐观效用估计会致命地高估不安全策略的可行性。REP-PD-H-TEG 沿用相同的特征学习与奖励构造机制,但彻底重构了策略选择与可行性验证逻辑。算法引入带有单边统计修正项 \(\ell_k = C_{\mathrm{cert}}\sqrt{\frac{|\mathcal{A}|k^p \ln(|\mathcal{M}|k/\delta)}{k(1-\gamma)}}\) 的效用置信下界证书: $\(\underline{\mathcal{U}}_k(\pi) := \left[ V^{\pi}_{\widehat{P}_k, g-\widehat{b}_k}(s_1) - \ell_k \right]_+.\)$ 在事件概率成立时,该证书严格构成真实效用函数的一致下界 \(\underline{\mathcal{U}}_k(\pi) \leq V^{\pi}_{P^\star, g}(s_1)\)。在每个回合内部,算法在网格 \(Y \in [0, Y_{\max}]\)(其中 \(Y_{\max} = K^{1/4}\))上执行截断对偶搜索:对每个网格点计算基于乐观拉格朗日目标的候选策略 \(\pi_{k,Y}\)(采用温度参数为 \(\tau_h\) 的状态级 Softmax 近似以保证对偶敏感度平滑),并逐一检验证书 \(\underline{\mathcal{U}}_k(\pi_{k,Y}) \geq b\)。算法部署首个通过证书测试的候选策略;若遍历至 \(Y_{\max}\) 仍未通过,则部署截断策略。基于 Slater 条件所诱导的对偶惩罚分析,未通过测试的回合对硬违例的累积贡献被严格限制在低阶项,最终保证了返回策略具备 \(\widetilde{O}(K^{1/2+p})\) 的硬违例界与 \(\widetilde{O}(K^{3/4+p})\) 的遗憾界。

损失函数 / 训练策略

在理论算法层面,MLE 阶段以经验对数似然最大化为损失函数;超参数 \(\alpha_k\) 与 \(\lambda_k\) 随回合数按 \(\Theta(\sqrt{k^p \ln(k)})\) 与 \(\Theta(k^p \ln(k))\) 动态调度。在面向高维连续控制的实用神经网络版本(REP-PD-NN)中,算法将全量 MLE 替换为小批量基于谱分解嵌入的 SPEDER 表现学习目标 \(\mathcal{J}_{\mathrm{rep}}(\eta)\),并结合双 Critic 架构与带有拉格朗日乘子自适应调整的软动作-评价(SAC)框架进行联合端到端优化: $\(\mathcal{J}_\pi(\theta) = \mathbb{E}\left[\alpha_{\mathrm{ent}} \log\pi_\theta(a|s) - \min_{i=1,2} Q_{\psi_i}^r(\phi_\eta(s,a)) + \lambda \min_{i=1,2} Q_{\omega_i}^c(\phi_\eta(s,a))\right].\)$

实验关键数据

主实验

论文在理论层面给出了软约束与硬约束下与各基准算法的样本复杂度全面对比,并在连续控制安全基准环境 Safety Gym(PointGoal1, PointPush1, PointPush2)上针对神经网络实用版本 REP-PD-NN 与无表示学习的原始状态输入基线(Raw)进行了 11M 步交互、33 个随机种子的对比实验。

表 1 展示了软约束设置下理论界对比(包含经典表格型、已知特征线性 CMDP 以及本文低秩 CMDP 算法):

MDP 类型 算法 特征已知 遗憾界(Regret) 软约束违例界(Violation)
Tabular Efroni et al. (2020) 是 \(\widetilde{O}(\sqrt{\vert\mathcal{S}\vert^3 \vert\mathcal{A}\vert H^4 K})\) \(\widetilde{O}(\sqrt{\vert\mathcal{S}\vert^3 \vert\mathcal{A}\vert H^4 K})\)
Tabular Liu et al. (2021) 是 \(\widetilde{O}(\sqrt{\vert\mathcal{S}\vert^3 \vert\mathcal{A}\vert H^6 K})\) \(0\)
Linear Ding et al. (2021) 是 \(\widetilde{O}(\sqrt{d^2 H^6 K})\) \(\widetilde{O}(\sqrt{d^2 H^6 K})\)
Linear Ghosh et al. (2022) 是 \(\widetilde{O}(\sqrt{d^3 H^3 K})\) \(\widetilde{O}(\sqrt{d^3 H^3 K})\)
Linear Ghosh et al. (2022) 是 \(\widetilde{O}(\sqrt{d^3 H^5 K})\) \(0\)
Low-rank REP-PD-TEG (返回策略, \(p=0\)) 否 \(\widetilde{O}\left(\frac{\vert\mathcal{A}\vert d^2 \sqrt{K}}{\theta (1-\gamma)^2}\right)\) \(\widetilde{O}\left(\frac{\vert\mathcal{A}\vert d^2 \sqrt{K}}{(1-\gamma)^2}\right)\)
Low-rank REP-PD-TEG (执行策略, \(p=1/4\)) 否 \(\widetilde{O}\left(\frac{\vert\mathcal{A}\vert d^2 K^{3/4}}{(1-\gamma)^2}\right)\) \(\widetilde{O}\left(\frac{\vert\mathcal{A}\vert d^2 K^{3/4}}{(1-\gamma)^2}\right)\)

表 2 展示了硬约束设置下理论界的对比结果:

MDP 类型 算法 特征已知 安全动作已知 遗憾界(Regret) 硬约束违例界(Violation)
Tabular Efroni et al. (2020) 是 是 \(\widetilde{O}(\sqrt{\vert\mathcal{S}\vert^3 \vert\mathcal{A}\vert H^4 K})\) \(\widetilde{O}(\sqrt{\vert\mathcal{S}\vert^3 \vert\mathcal{A}\vert H^4 K})\)
Tabular Bura et al. (2022) 是 是 \(\widetilde{O}(\sqrt{\vert\mathcal{S}\vert^2 \vert\mathcal{A}\vert H^6 K})\) \(0\)
Linear Amani et al. (2021) 是 是 \(\widetilde{O}(\sqrt{d^3 H^4 K})\) \(0\)
Linear Kitamura et al. (2025) 是 是 \(\widetilde{O}(\sqrt{d^5 H^8 K})\) \(0\)
Linear Wei et al. (2024) 是 否 \(\widetilde{O}(\sqrt{d^3 H^4 K})\) \(\widetilde{O}(\sqrt{d H^2 K})\)
Linear Ghosh et al. (2024) 是 否 \(\widetilde{O}(\sqrt{d^3 H^4 K})\) \(\widetilde{O}(\sqrt{d^3 H^4 K})\)
Low-rank REP-PD-H-TEG (返回策略, \(p=0\)) 否 否 \(\widetilde{O}\left(K^{3/4}\right)\) \(\widetilde{O}\left(\sqrt{K}\right)\)
Low-rank REP-PD-H-TEG (执行策略, \(p=1/8\)) 否 否 \(\widetilde{O}\left(K^{7/8}\right)\) \(\widetilde{O}\left(K^{7/8}\right)\)

表 3 汇总了实用神经网络实现 REP-PD-NN 在 Safety Gym 连续控制基准上的评估表现(评测指标为 33 个种子在 11M 步时的最终确定性测试累积回报与累积安全代价):

实验环境 原始状态基线 回报 (Raw Ret.) 原始状态基线 代价 (Raw Cost) REP-PD-NN 回报 (Ours Ret.) REP-PD-NN 代价 (Ours Cost)
PointGoal1 \(-1.38 \pm 0.58\) \(27.85 \pm 18.99\) \(\mathbf{-0.14 \pm 0.03}\) \(\mathbf{16.45 \pm 12.69}\)
PointPush1 \(-1.40 \pm 0.60\) \(32.17 \pm 14.66\) \(\mathbf{-0.17 \pm 0.12}\) \(\mathbf{12.13 \pm 6.95}\)
PointPush2 \(-1.40 \pm 0.71\) \(13.00 \pm 8.17\) \(\mathbf{-0.03 \pm 0.01}\) \(\mathbf{10.90 \pm 3.85}\)

消融实验

为了验证表示学习和约束对偶更新各自的作用机制,论文在任务调度(Job Scheduling)离散玩具环境上进行了对比分析(图 1 结果): 1. 与无约束低秩表示学习(REP-UCB)对比:REP-UCB 虽然具有更低的回报遗憾,但因为忽略了效用安全边界,产生了极其巨大的累积约束违例;而 REP-PD-TEG 成功将约束违例大幅抑制在极低水平,直观展示了原对偶机制对安全违例的严格约束效果。 2. 与已知特征线性 CMDP 方法对比:在无需任何先验特征的前提下,REP-PD-TEG 在遗憾与约束违例两个维度上的收敛尺度均与完美掌握真实转移特征的线性基线(Ghosh et al., 2022)旗鼓相当,验证了在线表示学习机制从零捕捉环境流形的优越性。

关键发现

  • 探测强度参数 \(p\) 的理论平衡点:在多项式探测调度 \(\epsilon_k = k^{-p}\) 下,存在明显的“统计学习误差”与“实际部署损失”之间的权衡。在软约束下,\(p=0\) 为返回策略带来理论最优的 \(\widetilde{O}(\sqrt{K})\) 误差,而考虑探测执行代价时,\(p=1/4\) 是两者的最佳平衡点,总代价为 \(\widetilde{O}(K^{3/4})\);在硬约束下,平衡点调整为 \(p=1/8\),总代价为 \(\widetilde{O}(K^{7/8})\)。
  • 表示学习对高维连续安全控制的提效作用:在 Safety Gym 实验中,直接利用原始连续状态训练的 SAC-Lagrangian 基线因探索效率低、维度灾难而陷入极低的回报和偏高的违规代价;引入 SPEDER 风格的紧凑潜空间表示学习后,智能体在 PointGoal1 和两个 PointPush 任务中的回报均提升近一个数量级,同时安全违规代价降低了 \(40\%\sim 60\%\)。

亮点与洞察

  • 理论结构性创举:首次填补了低秩 MDP 从无约束到约束环境的技术空白,针对朴素双层探测导致协方差失效的反例(Proposition D.1),构思了三阶段辅助探测机制,为高维低秩 CMDP 建立了首个完整的集中性与椭圆势能证明闭环。
  • 软硬违例的双轨解决方案:深刻剖析了“软累积违例允许盈余互补”与“硬累积违例必须逐回合鉴权”的本质差异,分别给出了“乐观拉格朗日双重奖励规划”与“置信下界单边证书+截断对偶搜索”的两套针对性算法,兼顾了理论优美度与针对性场景的实用性。
  • 理论与实践结合紧密:在完成完备的 PAC/Regret 理论分析之余,给出了基于自监督谱嵌入与 SAC-Lagrangian 的神经近似版本(REP-PD-NN),并在标准安全强化学习基准 Safety Gym 上验证了潜表示在缓解维度灾难与强化安全合规方面的双重收益。

局限与展望

  • 理论假设的严格性:理论推导依赖于转移模型类的有限可实现性假设(\(|\mathcal{M}| < \infty\) 且 \(P^\star \in \mathcal{M}\))以及严格的最优 MLE 解算子;在实际大模型或复杂神经网络拟合中,模型误设(misspecification)和优化近似误差尚未在理论证明中闭合。
  • 硬约束搜索的计算开销:REP-PD-H-TEG 需要在每个回合内部执行覆盖 \([0, K^{1/4}]\) 的细粒度对偶网格搜索,且需要通过平滑 Softmax 近似策略来控制对偶敏感度,当动作空间增大或精度要求极高时,回合内的规划耗时将显著增加。
  • 未来方向:探索无模型(model-free)的低秩 CMDP 表示学习框架以降低计算负担,并将分析拓展至部分可观测 CMDP(POCMDP)以及多智能体安全马尔可夫博弈场景。