强化学习KL散度：约束、正则还是度量？策略优化中的正确用法与调参经验

强化学习里的KL散度是个特别容易被“用错但看起来没错”的东西。我见过不少做策略优化的朋友代码里加了一行KL惩罚项训练日志上KL值也确实在动但策略该崩还是崩该不收敛还是不收敛。问题往往不在KL本身而在于没搞清楚自己加的到底是“约束”“正则”还是“度量”以及它作用在策略空间的哪个位置上。这篇就围绕强化学习中的KL散度展开把它在策略优化里的几种典型用法、背后的数学直觉、实现时的坑以及我实际调参中攒下来的经验尽量讲透。不管你是刚接触PPO、TRPO这类算法的新手还是已经在做策略约束但总觉得效果不稳的从业者应该都能从里面找到能直接抄作业的东西。1. 先搞清楚KL散度在策略优化里到底扮演什么角色很多人第一次接触KL散度是在信息论课上公式是那个熟悉的求和形式KL(p||q) Σ p(x) log(p(x)/q(x))。它衡量的是用一个分布q去近似另一个分布p时多付出的信息代价。放到强化学习里p和q通常就是新旧两个策略在某个状态下的动作分布。但真正决定你怎么用它、怎么调它的是它在优化目标里出现的位置和方向。1.1 三种常见用法别混为一谈在策略优化里KL散度大致有三种出场方式每一种的意图完全不同。第一种是信赖域约束。TRPO是典型代表它把新旧策略的KL散度当作一个硬约束要求更新后的策略不能离旧策略太远超过阈值就缩放步长。这里的KL是一个“安全边界”目的是防止一次更新把策略推到不可恢复的坏区域。它的方向通常是KL(旧||新)因为旧策略是采样来源我们关心的是在新策略下旧策略会多“惊讶”。第二种是惩罚项正则。PPO的KL-penalty变体就是这类把KL散度乘上一个系数β加到损失里作为软约束。这里KL更像一个“拉力”β越大策略越不敢偏离。它的好处是实现简单坏处是β很难调因为KL的量级会随训练阶段剧烈变化。第三种是分布度量或诊断指标。有时候我们不算它进损失只是监控KL值来判断训练是否稳定。比如PPO的KL-adaptive机制会根据实测KL动态调整裁剪范围。这里的KL是“仪表盘”不直接参与梯度但影响超参。提示动手前先问自己一句——我加的这个KL是要它产生梯度、限制步长还是只用来观察三种用法的实现细节和调参逻辑差别很大混着用是很多训练不稳的根源。1.2 方向问题KL(旧||新)和KL(新||旧)不是一回事KL散度不对称这一点在策略优化里非常关键但经常被忽略。KL(旧||新)在旧策略有支撑但新策略概率极低的地方会给出很大的惩罚也就是它更在意“新策略不要把旧策略认为可能的行为彻底否定”。而KL(新||旧)则相反它惩罚的是新策略去探索旧策略几乎不会做的动作。在TRPO和PPO的默认实现里约束项一般用KL(旧||新)因为采样来自旧策略我们希望在旧策略覆盖到的动作空间内做保守更新。如果你在实现时把方向搞反约束的语义就变了可能出现“明明KL很小但策略已经跑偏”的诡异现象。我早期就踩过这个坑监控指标用的是KL(新||旧)约束项用的是KL(旧||新)两个数对不上排查了半天才发现是方向不一致。1.3 采样估计为什么你算出来的KL和论文对不上理论上KL是对整个状态空间和动作空间求期望但实际只能用一个batch的样本去估。常见做法是对每个状态单独算动作分布间的KL然后对batch取平均。这里有两个细节会显著影响数值。一是是否按状态分布加权。严格来说应该按当前策略访问状态的分布加权但实现里通常就是batch内均匀平均。如果batch里状态分布偏移很大估计值会有偏。二是动作分布是离散还是连续。离散策略直接对概率向量求和连续策略比如高斯有解析形式均值差和方差比都能算进去。高斯策略下KL的解析式是KL log(σ_新/σ_旧) (σ_旧² (μ_旧 - μ_新)²)/(2σ_新²) - 0.5方向不同公式里新旧位置要换。手写这个公式时符号错一位KL就可能变成负数训练直接发散。2. PPO里KL惩罚与裁剪的取舍我为什么最后都选了裁剪PPO原论文其实给了两个变体KL-penalty和clip。现在主流实现几乎都用clip但理解KL-penalty为什么被冷落对搞懂KL在优化里的脾气很有帮助。2.1 KL-penalty的β为什么这么难调KL-penalty的损失长这样L -E[ratio * A] β * KL。β固定时训练早期策略变化快KL容易冲高惩罚不够训练后期策略趋于稳定KL很小惩罚又显得多余。更麻烦的是KL的量级和奖励尺度、动作维度都相关换个环境β就得重调。原论文提过一个自适应方案设定目标KL值d_target如果实测KL 1.5 * d_target就把β乘2如果 d_target/1.5就把β除2。这个思路本身没问题但实现时“实测KL”用哪个batch、多久调一次β都会影响稳定性。我试过每步都调结果β震荡得很厉害改成每若干步调一次又会出现滞后。相比之下clip机制用一个固定的ε就把“信任区域”表达清楚了超参少一个心智负担小很多。2.2 裁剪本质上是KL约束的一种粗粒度近似clip的目标函数是min(ratio * A, clip(ratio, 1-ε, 1ε) * A)。它没有显式计算KL但通过限制概率比ratio的范围间接限制了新旧策略的差异。可以理解为clip给了一个“每步最多走多远”的硬边界而KL-penalty给的是一个“越远越贵”的软拉力。从实践角度看clip的ε常见0.1到0.3比KL-penalty的β直观得多因为它直接对应概率比的范围。但要注意clip只约束了ratio不等于严格约束了KL。当动作维度很高时单个动作的ratio都在范围内累积起来的KL仍可能不小。这就是为什么有些实现会在clip之外再监控KL作为辅助诊断。2.3 一个容易忽略的细节early stopping该不该用有些PPO实现会在一次迭代内做多轮minibatch更新如果某一轮实测KL超过阈值就提前停止。这个early stopping用的是KL但它的作用和clip是互补的clip管的是单样本的ratioearly stopping管的是整批更新的整体偏移。我一般会设一个比较宽松的阈值比如目标KL的1.5到2倍只在明显异常时触发避免正常更新被误杀。阈值设太紧训练会变得畏手畏脚样本效率反而下降。3. 连续动作空间下KL的实现细节与数值陷阱离散动作的KL好算连续动作尤其是高斯策略坑就多了。这一节把实现层面容易出问题的地方摊开讲。3.1 高斯策略KL的解析式与手写风险假设策略输出动作的均值μ和标准差σ动作分布是N(μ, σ²)。两个高斯分布之间的KL以旧为参考、新为待评估解析式前面提过。手写时最常见的错误有三个把log里的σ比值写反、方差和标准差混用、常数项0.5漏掉或符号错。任何一个错了KL都可能在数值上表现为负或者异常大。我的建议是如果框架里有现成的分布类比如某些概率编程库的kl_divergence优先用库函数别手写。如果必须手写写完拿两个相同分布测一下KL应该严格为0再拿两个已知解析解的小例子对一遍确认方向和量级都对。3.2 标准差参数化方式对KL的影响很多实现不直接输出σ而是输出log σ或者用softplus保证正性。这本身没问题但要注意KL计算时用的是真实的σ还是变换后的量。如果策略网络输出的是log σ算KL前必须先exp回来。我见过有代码直接把log σ代进高斯KL公式结果KL值完全失真训练看起来在跑实际约束根本没生效。另外σ的下界裁剪也会影响KL。如果对σ做了clamp防止它趋近0那么当新旧策略的σ都撞到下界时KL里方差相关的项会退化成只由均值差决定这时候KL的语义就变了。调试时如果发现KL异常小但策略明显在变先检查σ是不是被裁到下界了。3.3 数值稳定性log_prob的减法比直接算KL更稳一个实用技巧与其显式计算KL不如用log_prob的差值来估。对同一批动作KL(旧||新) ≈ E_旧[log π_旧(a|s) - log π_新(a|s)]。这个估计在实现上更稳因为它复用了采样时已经算过的log_prob避免了重复计算分布参数带来的数值误差。当然它是有偏的因为是在旧策略样本上估但在PPO这类on-policy场景里偏差可以接受而且计算开销小。注意用log_prob差值估KL时样本必须来自旧策略。如果你混入了新策略采样的数据这个估计就失去意义了。on-policy的纪律在这里体现得很直接。4. 把KL用进损失函数时的梯度与缩放问题KL进入损失后它的梯度怎么流、量级怎么和策略梯度平衡是另一个容易翻车的地方。4.1 KL项的梯度会不会“污染”策略梯度当损失是 -E[ratioA] βKL 时KL项对策略参数的梯度是存在的它会推动策略往减小KL的方向走。问题在于KL的梯度量级和策略梯度项往往不在一个数量级。如果β设得偏大KL梯度会主导更新策略几乎不动表现为奖励曲线平、KL很小但学不到东西β偏小则约束形同虚设。我的做法是先把β设为0跑一段观察纯策略梯度下KL的自然增长速率再据此选β让KL惩罚在初期大约占损失总量的一个可控比例比如10%到30%。这个比例不是金科玉律但比拍脑袋定β靠谱。4.2 用detach切断不需要的梯度路径有些场景下我们只想用KL做监控或自适应调参不希望它产生梯度。这时候必须对KL项做detach否则它会悄悄影响优化方向。反过来如果KL是约束项就要确保它对新策略参数是可导的别不小心把新策略的分布参数也detach了那样约束就完全失效。这两种情况在代码里往往只差一个.detach()但行为天差地别review代码时要特别留意。4.3 多步更新下KL的累积效应PPO一次采样会做多轮更新每轮更新后策略都在变但用来算ratio和KL的旧策略log_prob是固定的。这意味着随着更新轮数增加实测KL会单调上升。如果不做early stopping后面几轮的更新可能已经超出信任区域但因为clip的存在损失看起来还算正常。我的经验是把更新轮数控制在合理范围常见3到10轮并监控每轮的KL一旦发现某一轮KL跳变明显就该停。5. 自适应KL与目标KL的设定经验KL-adaptive机制在不少实现里都有但目标KL设多少、怎么调资料里往往一笔带过。这里说说我的实际取值习惯。5.1 目标KL的经验区间目标KL的设定和动作维度、任务难度都有关。动作维度低比如几个离散动作时目标KL可以设小一些0.01到0.02量级连续高维动作时同样的KL数值对应的策略偏移可能更大目标可以适当放宽到0.02到0.05。这些数字不是绝对的关键是保持一致性一旦定了目标监控和调参都围绕它来别今天0.01明天0.1。5.2 调整倍率的节奏比倍率本身更重要常见的自适应规则是KL超过目标1.5倍就增大惩罚、低于目标除以1.5就减小惩罚。倍率1.5到2都行真正影响稳定性的是调整频率。每步都调会让β高频震荡建议每隔固定步数比如一个采样周期调一次并且给β设上下界防止它跑到极端值。β一旦趋近0约束就没了β过大策略冻结。设个比如[1e-4, 10]的范围能兜住大部分异常。5.3 监控KL的分布而不只是均值只看batch平均KL会掩盖问题。同一个batch里不同状态的KL可能差异很大少数状态KL极高但被平均掉了。我的习惯是同时记录KL的均值和最大值或者95分位如果最大值远高于均值说明策略在某些状态上偏移剧烈这时候即使均值达标也该警惕。这个习惯帮我提前发现过好几次局部策略崩塌。6. 几个真实调试场景里的KL异常与排查思路理论讲完落到实际。下面几个场景是我或身边同行真实遇到过的排查链路尽量还原。6.1 KL突然变成NaN从分布参数查起训练中途KL变NaN第一反应应该是分布参数出了非有限值。常见原因σ经过exp后溢出、log_prob里出现log(0)、均值网络输出爆炸。排查顺序是先在算KL前打印μ和σ的min/max确认没有inf/nan再检查σ的裁剪是否生效最后看奖励尺度是否过大导致梯度爆炸间接污染了策略参数。我遇到过一次是奖励没做归一化梯度把均值推到了极大值σ的exp直接溢出。6.2 KL很小但策略明显退化方向或detach的锅如果监控的KL一直很小但评估奖励在掉先确认KL的方向和计算对象对不对。有可能是约束项算的是KL(新||旧)而监控算的是KL(旧||新)两者不一致也有可能是KL项被意外detach约束根本没生效。还有一种隐蔽情况σ被裁到下界KL里方差项退化数值上很小但均值其实已经跑很远。这时候要单独看均值差。6.3 训练后期KL持续爬升不回落更新轮数或学习率的问题后期KL单调上升通常意味着单次迭代内更新过头或者学习率相对当前策略尺度偏大。先减更新轮数再考虑降学习率。如果用了自适应KL检查β是不是已经撞到上界还在往上顶那说明目标KL设得太小策略根本没法在约束内有效更新需要放宽目标。7. 一些不写在文档里的实操习惯最后分享几个我长期用下来觉得值当的习惯都是些小动作但能省不少排查时间。第一把KL和奖励、熵、ratio的均值画在同一张监控图上。单独看KL没感觉放在一起才能看出“KL涨的时候奖励是不是在涨”“ratio是不是已经贴边”。这种联动观察比任何单一指标都直观。第二给KL设一个报警阈值而不只是记录。训练脚本里加一行判断KL超过某个明显异常的值就打印警告甚至中断比事后翻日志高效得多。阈值可以设成目标KL的若干倍具体看任务容忍度。第三换环境或换动作维度时重新校准KL的量级。别把上一个任务调好的β或目标KL直接搬过来动作维度一变同样的KL数值含义就变了。花十分钟重新观察一下KL的自然范围比盲目迁移超参省事。第四离散和连续策略的KL代码分开写、分开测。混在一个函数里用if分支容易在某个分支上漏掉方向或参数变换。分开写虽然啰嗦但测试用例清晰出问题定位快。KL散度这东西公式简单但在强化学习里怎么用、用在哪、用多大全是经验活。把它当成一个需要理解脾气的工具而不是一个填进去就完事的数字训练稳定性会好很多。