[论文解读] On Wasserstein Reinforcement Learning and the Fokker-Planck equation

Pierre H. Richemond, Brendan Maginnis|arXiv (Cornell University)|Dec 19, 2017

Fault Detection and Control Systems参考文献 18被引用 13

一句话总结

本文提出了一种基于Wasserstein正则化的策略梯度框架，表明在$W_2$距离下进行小步长策略更新时，其动力学受Fokker-Planck（热）方程支配。这揭示了最优策略通过扩散和向高奖励动作的输运过程自然涌现，从而通过最优传输理论解释了诸如高斯噪声和熵正则化等经验实践的有效性。

ABSTRACT

Policy gradients methods often achieve better performance when the change in policy is limited to a small Kullback-Leibler divergence. We derive policy gradients where the change in policy is limited to a small Wasserstein distance (or trust region). This is done in the discrete and continuous multi-armed bandit settings with entropy regularisation. We show that in the small steps limit with respect to the Wasserstein distance $W_2$, policy dynamics are governed by the Fokker-Planck (heat) equation, following the Jordan-Kinderlehrer-Otto result. This means that policies undergo diffusion and advection, concentrating near actions with high reward. This helps elucidate the nature of convergence in the probability matching setup, and provides justification for empirical practices such as Gaussian policy priors and additive gradient noise.

研究动机与目标

• 建立最优传输与基于$W_2$距离的熵正则化强化学习之间的理论联系。
• 证明在连续时间极限下，基于Wasserstein信任区域的策略梯度更新收敛于Fokker-Planck方程的解。
• 通过策略空间中的扩散与输运过程，解释强化学习中经验成功做法（如梯度噪声和高斯先验）的机理。
• 将信任区域方法推广至KL散度之外，采用Wasserstein度量（特别是$W_2$）实现更平滑的策略更新。
• 为无限维策略空间中的策略传输提供数学基础，基于Jordan-Kinderlehrer-Otto（JKO）框架。

提出的方法

• 将策略迭代形式化为在概率测度空间中使用$W_2$ Wasserstein距离作为度量的梯度流。
• 推导在$W_2$信任区域下策略更新的连续时间极限，表明其收敛于Fokker-Planck方程。
• 利用Jordan-Kinderlehrer-Otto（JKO）结果，将$W_2$空间中的梯度流与热方程关联，以建模策略的扩散与输运。
• 对奖励泛函施加熵正则化，导出类似于统计力学中的自由能泛函。
• 证明Fokker-Planck方程的稳态解即为Gibbs测度$\pi^*(a|s) \propto e^{r(s,a)/\beta}$，对应最优策略。
• 利用Sinkhorn算法和正则化最优传输近似$W_2$近端映射，实现实际应用中的可行计算。

实验结果

研究问题

• RQ1 在策略梯度中使用$W_2$ Wasserstein距离作为信任区域，如何影响策略优化的动力学？
• RQ2 $W_2$正则化策略梯度更新的连续时间极限由何种偏微分方程支配？
• RQ3 Fokker-Planck方程如何通过扩散与输运过程解释最优策略的涌现？
• RQ4 Fokker-Planck动力学与添加梯度高斯噪声等经验实践之间有何关联？
• RQ5 基于$W_2$的策略更新能否与已知的熵正则化策略优化框架（如Trust Region Policy Optimization（TRPO）或soft Q-learning）建立联系？

主要发现

• 在$W_2$信任区域下的策略更新，在连续时间极限下收敛于Fokker-Planck方程的解，可建模策略的扩散与输运。
• Fokker-Planck方程的稳态解为Gibbs测度$\pi^*(a|s) \propto e^{r(s,a)/\beta}$，与熵正则化最优策略一致。
• Fokker-Planck动力学解释了为何策略会集中于高奖励动作附近，为概率匹配中的收敛提供了机制。
• 基于$W_2$的梯度流解释了为何梯度噪声和高斯策略先验在实践中有效，其根源在于随机扩散过程。
• $W_2$近端映射生成一系列收敛于最优策略的Gibbs测度，其动力学由连续的Kantorovich势函数所支配。
• 通过Sinkhorn算法实现的正则化最优传输，为$W_2$基础策略更新提供了数值稳定的近似，从而支持实际实现。