跳转至

Decision-Focused Learning in MDPs: An Occupancy Measure Approach

会议: NeurIPS 2026
arXiv: 2610.08384
代码: https://github.com/zzhao628/DFL-MDP
领域: 强化学习
关键词: 决策聚焦学习, 马尔可夫决策过程, 占用度量, 线性规划层, 状态聚合

一句话总结

将未知转移与奖励条件下的马尔可夫决策过程重新表述为占用度量线性规划层,通过可行基矩阵求逆导出闭式梯度,并结合增广拉格朗日代理损失、随机行草图平滑与可学习软状态聚合,在大幅降低反向传播计算复杂度的同时显著优化下游决策遗憾。

研究背景与动机

在很多现实序列决策场景中,底层环境的转移概率与奖励参数往往未知,需要利用上下文特征借助预测模型先行估计,例如不确定需求下的库存控制与动态天气下的机器人导航。传统的两阶段方法先使用监督学习最小化环境参数的预测误差(如均方误差),再将预测参数送入下游优化器求解策略。然而,预测误差的降低并不能等价转化为下游决策质量的提升,因为微小的关键参数偏差可能引发优化解拓扑结构的剧烈跳变。决策聚焦学习(Decision-Focused Learning, DFL)试图通过将策略后悔度梯度端到端反向传播至预测模型来打破这种错配。

然而,现有面向马尔可夫决策过程(MDP)的 DFL 方法大多针对贝尔曼最优性方程的 KKT 条件进行隐式微分。这一路线在反向传播时需要求解一个维度高达 \(|\mathcal{S}||\mathcal{A}| \times |\mathcal{S}||\mathcal{A}|\) 的大型线性系统,不仅其计算复杂度随状态和动作空间急剧膨胀,而且贝尔曼算子的非线性极易在折扣因子趋近 1 时引发严重的矩阵病态问题。与此同时,在约束条件中包含预测参数的 DFL 场景下,预测动态诱导出的占用度量通常会违反真实环境的流平衡约束,导致传统的后悔度目标产生系统性偏移。

针对这一困境,本文的核心切入点是回归 MDP 的对偶空间,利用占用度量(Occupancy Measure)在线性规划(LP)下的几何与代数特性构建可微层。核心 idea:将 MDP 重构为基于占用度量的线性规划,通过在基可行解处求逆仅为 \(|\mathcal{S}| \times |\mathcal{S}|\) 维度的基矩阵导出闭式良态梯度,并配合增广拉格朗日代理损失纠正可行域偏差、随机行草图平滑边界跳跃,以及可微软状态聚合扩展至大规模乃至连续状态空间。

方法详解

整体框架

本文构建的端到端 DFL 流程以环境上下文特征为输入,通过预测网络生成参数化的转移概率与即时奖励。为了实现决策感知梯度的低成本传递,框架首先通过软状态聚合层将原始状态空间压缩为低维抽象状态,进而构建低维占用度量线性规划。在求解线性规划时,引入随机行草图机制对流平衡约束进行扰动平滑并平均解向量,消除基变更引起的梯度断裂;接着利用增广拉格朗日函数计算兼顾真实回报、对偶修正与流守恒惩罚的代理损失;最后根据基矩阵导出的闭式雅可比矩阵完成预测器权重的高效梯度更新。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入上下文特征 x"] --> B["预测器生成环境参数<br/>预测转移 P_theta 与奖励 r_theta"]
    B --> C["可微软状态聚合<br/>温度缩放映射与对角化投影压缩维度"]
    C --> D["随机行草图平滑<br/>高斯随机约束投影消除基跳变"]
    D --> E["占用度量线性规划层<br/>S*S 基矩阵逆计算闭式梯度"]
    E --> F["增广拉格朗日代理损失<br/>对偶校正与可行性惩罚保证全局一致性"]
    F --> G["端到端梯度回传更新预测器 w 与聚合参数 xi"]

关键设计

1. 占用度量线性规划层与闭式梯度:摆脱 KKT 联立求解的高效反向传播 传统基于贝尔曼方程的 KKT 微分将值函数与策略变量深度耦合,导致反向传播时必须对整个状态-动作联合空间 \(|\mathcal{S}||\mathcal{A}| \times |\mathcal{S}||\mathcal{A}|\) 的雅可比矩阵求逆。本文将 MDP 等价重塑为以折现占用度量 \(y(s,a) = \sum_{t=0}^\infty \beta^t \Pr(s_t=s, a_t=a)\) 为决策变量的线性规划: $\(\max_{y \ge 0} r_\theta^\top y \quad \text{s.t.} \quad H_\theta y = \gamma\)$ 其中流约束矩阵定义为 \([H_\theta]_{s, (s',a')} = \delta_{s,s'} - \beta P_\theta(s \mid s',a')\),\(\gamma\) 为严格正的初始状态分布。在严格正初始分布下,每个基可行解均非退化,且基矩阵 \(H_\theta^B \in \mathbb{R}^{|\mathcal{S}| \times |\mathcal{S}|}\) 严格可逆。通过在流约束等式 \(H_\theta y_\theta^* = \gamma\) 上对预测参数 \(\theta\) 进行全微分,并利用非基变量分量 \(y_\theta^N = 0\) 的互补松弛性质,推导出基变量关于模型参数的精确解析梯度: $\(\frac{dy_\theta^B}{d\theta} = -(H_\theta^B)^{-1} \left( \frac{\partial H_\theta^B}{\partial \theta} y_\theta^B \right)\)$ 这一设计使矩阵求逆维度从 \(|\mathcal{S}||\mathcal{A}|\) 降至 \(|\mathcal{S}|\),不仅在动作空间较大时带来多数量级的计算加速,而且避免了传统目标扰动 LP 中解梯度恒等于零的奇异性问题。

2. 增广拉格朗日代理损失:纠正动态不匹配诱导的目标偏差 当预测模型输出的转移概率偏离真实动态时,最优解 \(y_\theta^*\) 仅在预测流矩阵 \(H_\theta\) 下可行,而在真实流约束 \(H_{\mathrm{true}}\) 下产生违例。如果直接以真实线性回报 \(-r_{\mathrm{true}}^\top y_\theta^*\) 作为损失函数,预测器可能会倾向于生成能虚拟制造出“虚假高流”的虚假动态,使损失函数的全局极小点偏离真实物理参数。为此,本文设计了融合对偶校准项与二次违例惩罚的增广拉格朗日代理损失函数: $\(L(y_\theta^*) = -r_{\mathrm{true}}^\top y_\theta^* + (H_{\mathrm{true}} y_\theta^* - \gamma)^\top \nu^* + \rho \|H_{\mathrm{true}} y_\theta^* - \gamma\|^2\)$ 其中 \(\nu^*\) 是真实 MDP 对应的最优对偶向量(由离线真实 LP 求解阶段直接得出,无需额外计算)。由于强对偶性 \(r_{\mathrm{true}}^\top y_{\mathrm{true}}^* = \gamma^\top \nu^*\),代理损失与理论最优界满足: $\(L(y_\theta^*) - L(y_{\mathrm{true}}^*) = (\eta^*)^\top y_\theta^* + \rho \|H_{\mathrm{true}} y_\theta^* - \gamma\|^2 \ge 0\)$ 式中 \(\eta^* \ge 0\) 为非负乘子。当且仅当预测策略在真实环境下完全满足流平衡且达到理论最大收益时差值取零。该损失在理论上严格证明了真实参数 \(\theta^*\) 是增广拉格朗日代理目标的全局极小点。

3. 约束随机行草图平滑:缓解多面体边界基变更的梯度突变 由于线性规划的最优解始终位于可行多面体的极点上,微小的环境参数扰动可能导致最优基发生离散切换,进而使得梯度在多面体边界处产生突变甚至不连续,阻碍一阶优化算法的收敛。与传统仅扰动目标函数的做法不同,本文提出约束随机行草图方法。在每次前向传播中独立抽取 \(K\) 个高斯草图矩阵 \(S^{(i)} \in \mathbb{R}^{k \times |\mathcal{S}|}\)(保持行保留比 \(\alpha = k/|\mathcal{S}|\)),同时显式引入质量守恒等式 \(\mathbf{1}^\top y = 1/(1-\beta)\) 保证草图后可行域的有界性。通过分别求解 \(K\) 个草图子 LP: $\(\max_{y \ge 0} r_\theta^\top y \quad \text{s.t.} \quad S^{(i)} H_\theta y = S^{(i)} \gamma, \quad \mathbf{1}^\top y = \frac{1}{1-\beta}\)$ 并将得到的多个解向量在极点之间取经验平均 \(\bar{y}_\theta^* = K^{-1} \sum_{i=1}^K y_\theta^{*(i)}\)。这种多极点随机凸组合从经验上有效平滑了有效基边界的跃变,为预测网络反向传播提供连续稳定的指导梯度。

4. 可微软状态聚合与连续函数近似:打破维度瓶颈的端到端表征压缩 当状态空间极大或处于连续域时,\(|\mathcal{S}| \times |\mathcal{S}|\) 的矩阵求逆依然沉重。本文设计了端到端可训练的软状态聚合层,引入参数化归属度矩阵 \(W \in \mathbb{R}^{|\mathcal{S}| \times M}\)(满足行单形约束 \(\sum_{m} W_{sm} = 1\)),将状态归约到 \(M \ll |\mathcal{S}|\) 个隐式聚合簇。为保证投影后的流约束能够精确对应一个合法的抽象 MDP,本文提出对角化投影替代方案,利用 \(\mathrm{diag}(\mu)\)(其中 \(\mu_m = \sum_s W_{sm}\))逼近投影交叉项 \(W^\top W\),构造出维度仅为 \(M \times M\) 的聚合流平衡矩阵 \(\hat{H}\): $\(\hat{r}_{m,a} = \frac{\sum_s W_{sm} r_{s,a}}{\mu_m}, \quad \hat{P}(n' \mid m,a) = \frac{\sum_{s,s'} W_{sm} P(s' \mid s,a) W_{s'n'}}{\mu_m}, \quad \hat{\gamma}_m = \sum_s W_{sm} \gamma_s\)$ 对于连续状态空间,进一步泛化为函数近似范式:采用全连接网络表示占用特征 \(\Phi(s,a)\) 与测试检验基 \(\Psi(s)\),并在离线转移经验回放上通过蒙特卡洛期望无偏估计流约束矩阵项 \((\Psi^\top P \Phi)_{ij} \approx \frac{1}{N} \sum_{t=1}^N \psi_i(s_{t+1}) \phi_j(s_t, a_t)\),使计算复杂度从 \(\mathcal{O}(|\mathcal{S}|^\omega)\) 彻底收敛至 \(\mathcal{O}(M^\omega)\) 或 \(\mathcal{O}(d^\omega)\)。

损失函数 / 训练策略

训练过程采用端到端交替更新。离线阶段首先通过真实环境参数预计算真实对偶向量 \(\nu^*\)。在线迭代时,采样微批次环境特征 \(x\),预测器输出 \(\theta = m_w(x)\) 生成预测转移核与奖励;接着执行带温度缩放的加权聚合,生成降维参数 \((\hat{r}, \hat{P}, \hat{\gamma})\);随后并行执行 \(K\) 次独立的随机行草图求解,通过平均解计算解嵌入并提升还原至全状态空间;最后依据增广拉格朗日公式计算标量损失 \(\ell\),通过基矩阵闭式求逆和自动微分工具,联合反向更新预测器权重 \(w\) 与聚合结构参数 \(\xi\)。在测试推理阶段,直接使用预测参数求解未做草图投影的标准 LP 并根据占用度量比值恢复确定性策略。

实验关键数据

主实验

论文在离散库存管理(Inventory)、带潜在天气转换的随机悬崖行走(Cliff Walking)以及连续控制 CartPole 任务上进行了全面评测,对比了标准二阶段法(Two-stage)、基于 KKT 系统的二次规划松弛 DFL(DFL-QP)、可行性感知方法(DFL-Feas)以及本文的 DFL-LP、DFL-Sketch 与结合状态聚合的 DFL-SA。

在随机悬崖行走任务(Cliff Walking, 状态数 48,动作数 4,测试集存在严重天气分布漂移)中,各方法在 10 个随机种子下的成对测试遗憾对比见下表:

方法 平均测试遗憾 (Test Regret ↓) 标准误 (SE) 相对两阶段胜场 (Wins vs. TS) 计算耗时/复杂度特性
Two-stage (MSE基线) 14.75 0.29 — 仅回归拟合,决策感知缺失
DFL-LP (本文闭式基逆) 14.14 0.18 7/10 耗时相对 KKT 下降数倍
DFL-Sketch (本文草图平滑) 11.34 1.44 8/10 边界平滑改善断裂梯度
DFL-QP (KKT 基线 [Wang et al., 2021]) 9.53 0.84 10/10 系统庞大,计算极耗时
DFL-SA (\(M=6\), 本文聚合) 7.21 0.15 10/10 兼顾低维基求解与决策感知
DFL-SA (\(M=12\), 本文聚合) 7.76 0.15 10/10 鲁棒适应分布漂移
DFL-SA (\(M=24\), 本文聚合) 7.87 0.64 10/10 全面超越传统 DFL 基线

在更大状态空间的扩展库存管理任务(Large Inventory MDP,状态数 101,动作数 31,优化变量 \(101 \times 31 = 3131\))中,各方法的测试遗憾与训练时间对比如下:

方法 测试遗憾 (Test Regret ↓) 总训练耗时 (秒) 相对 DFL-QP 耗时加速比
Two-stage \(53.788 \pm 7.710\) \(227.3 \pm 6.4\) —
DFL-QP \(44.545 \pm 5.823\) \(34791.0 \pm 756.4\) 基准 (\(1.0\times\))
DFL-LP (Ours) \(40.690 \pm 4.946\) \(19915.3 \pm 527.1\) 约 \(1.75\times\)
DFL-Feas \(62.904 \pm 8.230\) \(19832.6 \pm 618.7\) 约 \(1.75\times\)
DFL-Sketch (Ours) \(48.591 \pm 5.718\) \(13790.6 \pm 513.4\) 约 \(2.52\times\)
DFL-SA (\(M=10\)) (Ours) \(42.568 \pm 5.814\) \(540.6 \pm 3.8\) 约 \(64.4\times\)
DFL-SA (\(M=20\)) (Ours) \(43.459 \pm 5.432\) \(605.2 \pm 6.0\) 约 \(57.5\times\)
DFL-SA (\(M=30\)) (Ours) \(\mathbf{39.187 \pm 5.198}\) \(862.8 \pm 12.2\) 约 \(\mathbf{40.3\times}\)

消融实验

论文针对增广拉格朗日违例惩罚权重 \(\rho\) 的敏感性以及软状态聚合中的对角化流平衡投影设计进行了消融分析。

在库存控制场景下对软状态聚合投影机制的消融结果如下表所示(聚合大小 \(M=10\)):

投影配置 测试遗憾 (Test Regret ↓) 训练求解不可行比率 (%) 机制分析
Hard (不可导硬聚类) \(63.024 \pm 4.911\) 0.0% 无法进行端到端表征梯度回传
Proj-LP (直接投影,未做对角化) \(24.147 \pm 5.758\) 61.8% 破坏流守恒约束,造成大量非法求解
Proj-MDP (DFL-SA 对角化投影) \(\mathbf{14.966 \pm 2.319}\) \(\mathbf{0.0\%}\) 严格保持有效 MDP 结构,实现最佳决策表现

关键发现

  • 状态聚合层显著释放了计算瓶颈:在 \(|S|=101\) 的大状态库存环境中,DFL-SA (\(M=30\)) 在取得全场最低测试遗憾(39.187)的同时,将训练时间从传统 DFL-QP 的近 10 小时骤降至不到 15 分钟(862.8 秒),取得了超过 40 倍的计算加速。
  • 随机行草图平滑在跨基跳跃严重的场景表现卓越:在 CartPole 与悬崖行走任务中,草图平均机制显著缓解了多面体边界处的梯度震荡,在 CartPole 连续任务中将平均测试遗憾从未平滑的约 24 降至 9。
  • 增广拉格朗日对偶校正项不可或缺:直接在未校正的占用度量上执行 DFL 会因动态违规产生严重的目标偏移,而增大惩罚权重 \(\rho\) 能单调降低流违例均方误差,并在 \(\rho \in [10^3, 10^5]\) 区间内形成稳健的低遗憾决策边界。

亮点与洞察

  • 矩阵维度大幅解耦:将反向传播计算核心从 \(|\mathcal{S}||\mathcal{A}| \times |\mathcal{S}||\mathcal{A}|\) 压缩至 \(|\mathcal{S}| \times |\mathcal{S}|\) 甚至 \(M \times M\)。不仅在理论上绕开了贝尔曼非线性算子的数值不稳定性,而且巧妙利用流平衡的基选择导出了显式闭式梯度。
  • 约束空间随机行草图平滑:打破了传统 DFL 仅对目标函数系数施加高斯扰动的思维定势,针对动态不确定性下的约束多面体采用高斯行随机投影,通过多解集成实现了梯度流的经验平滑。
  • 统一离散聚合与连续函数近似:通过对偶对角化投影技术巧妙保持了软聚类后的有效 MDP 结构,并与蒙特卡洛离线转移流匹配无缝融合,使 DFL 首次能够顺畅扩展至连续状态强化学习环境。

局限与展望

  • 基切换处的理论平滑性缺乏严格闭式上界:虽然随机行草图平均从经验上大幅平滑了优化边界,但有限采样平均无法在理论上严格消除测度为零处的导数间断点。
  • 对超参数选择存在一定敏感度:增广拉格朗日惩罚权重 \(\rho\) 与草图保留比 \(\alpha\) 在不同环境尺度下最佳量级差异较大,仍需借助离线交叉验证进行挑选。
  • 依赖离线转移样本质量:连续状态下的函数近似极度依赖行为策略采集的经验缓冲池覆盖度,在未知探索区域的外推泛化性仍有待进一步验证。

相关工作与启发

  • vs DFL-QP (Wang et al., NeurIPS 2021): DFL-QP 直接针对贝尔曼最优性方程的 KKT 条件展开隐式微分,需要求解高达 \(|\mathcal{S}||\mathcal{A}| \times |\mathcal{S}||\mathcal{A}|\) 的庞大代数系统且面临病态条件数。本文通过占用度量 LP 重新表述,反向传播核心缩减为仅需对 \(|\mathcal{S}| \times |\mathcal{S}|\) 的局部基矩阵求逆,计算开销呈指数级下降。
  • vs DFL-Feas (Mandi et al., NeurIPS 2025): DFL-Feas 主要针对静态规划中参数进入约束的问题引入可行性罚项,但忽略了 MDP 中动态失配导致的占用度量系统性偏差。本文引入带真实对偶向量修正的增广拉格朗日损失,严格保证真实参数是代理目标的全局驻点。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从占用度量线性规划视角构建可微 MDP 决策聚焦学习框架,设计高度原创。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖离散控制、随机网格与连续控制,并包含细致的复杂度缩放与投影机制消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨流畅,问题定位精准,方法脉络清晰。
  • 价值: ⭐⭐⭐⭐⭐ 成功解决了 DFL 在强化学习领域难以扩展至大状态动作空间的核心瓶颈,工程与理论意义重大。