我们在做强化学习或者大模型对齐工作时通常已经形成了某种固定流程在训练环境里用策略梯度把策略跑起来调好超参收敛后保存模型权重部署后就用固定权重做决策。这套流程很成熟但它无意识里藏了一个很大的假设——测试环境和训练环境不会相差太多。一旦部署场景出现分布偏移固定策略就很容易“失灵”。比如推荐系统在某个新流量场景下点击率骤降机器人在某个陌生地形上动作频繁失稳大模型的在线请求分布从代码问答转向长文档摘要。面对这些变化传统的做法是重新收集数据、回炉训练成本高且周期长。那有没有一种思路可以在不改动训练阶段的情况下让策略在测试阶段继续“补课”Test-Time Policy Optimization测试时策略优化简称 TTPO要解决的就是这个问题。它的核心位置不在训练阶段而在部署和推理阶段当策略已经在线上运行时利用当前环境和实际反馈对策略做少而稳的在线更新让模型在真实场景里继续适应。本文会讲清楚 TTPO 解决什么问题、它和常规 RL/PPO 有什么区别然后给出一个可以直接跑通的 PyTorch 最小实现最后讨论工程落地时最容易踩的坑。1. 这篇文章真正要解决的问题先从一个具体场景说起。假设你训练了一个上下文多臂老虎机策略训练时收集到的上下文主要是“用户来自渠道 A、B”模型学到的策略也非常适合这两类用户。上线之后由于业务推广变化请求分布突然变成渠道 C、D 占多数这时候固定策略很可能不是最优的——它可能在渠道 C 上做了一个并不好的动作选择。传统 RL 的处理方式很直观把新数据捞回来重新训练或微调策略。但这个流程有几个现实问题数据回流需要时间问题已经在线发生了很久。重新训练需要大量计算资源临时调度非常麻烦。训练阶段没有见过的新状态即使回炉也不一定能覆盖。TTPO 的做法是把优化窗口从训练期延伸到测试期。部署之后策略每和环境交互一笔就获得一个真实的反馈信号比如点击、转化、奖励、人工修正。基于这些反馈策略在测试阶段做有限步数的策略梯度更新从而更快适应当前环境。更准确地说TTPO 不是一个单独算法而是一类“测试时在线策略优化”的思路。它和普通策略梯度算法的核心区别不是公式而是更新发生的时间点和约束条件。测试阶段的数据量远小于训练阶段更新又不允许让策略彻底崩溃因此 TTPO 几乎总需要加一层额外约束通常是用 KL 散度或者熵正则把新策略拉在旧策略附近。这篇文章最值得读的读者应该是这几类正在做 RL 落地发现固定策略在线上出现分布偏移的开发者和算法工程师。做大模型推理优化关心如何在测试阶段继续利用外部反馈改进输出的工程师。想理解“测试时适应”“推理时优化”这类概念但需要一份可运行代码来辅助理解的算法入门者。如果你只是想把模型训练阶段本身做得更好TTPO 不是最适合的方向如果你手头已经有一个可以接收在线反馈的策略TTPO 能在不改训练流程的情况下给你一个低成本补救手段。2. 核心概念辨析训练时策略优化与测试时策略优化要理解 TTPO先要建立两个基本概念策略和策略优化。策略 (\pi(a|s)) 指的是在状态 (s) 下选择动作 (a) 的概率分布。策略优化则是指调整策略网络的参数使得某些目标函数不断提高。传统流程中策略优化几乎全部发生在训练阶段我们使用静态数据集或者模拟环境通过 PPO、REINFORCE、Actor-Critic 等算法来更新网络。训练时策略优化的关键特征是数据来源是训练集或者仿真器更新过程有完整循环算法可以在多个 epoch 上反复利用样本更新步数通常很大甚至可以跑到收敛。它的目标是得到一个适应训练分布的模型但模型在测试分布上是否依然有效并不在优化目标中。测试时策略优化则完全不同。它把优化目标放宽为“在当前测试环境下逐渐适应当前反馈”。TTPO 的每一条样本都来自真实的在线交互数量不多噪声更大而且不能反复使用太久否则会把策略污染成只适应最近几十条数据。更具体地TTPO 在更新时至少要考虑三个额外问题更新步数必须受到限制防止在线策略漂移。新策略不能离旧策略太远通常需要 KL 约束。奖励信号往往延迟或稀疏需要对优势做基线处理。下面这张表可以帮助快速理解两者的边界维度训练时策略优化测试时策略优化TTPO优化阶段训练期模型保存前部署期模型推理过程中数据来源离线数据集、仿真环境当前真实环境在线交互数据规模大可多轮重复利用小通常只有几十到几百条更新目标最大化训练分布上的回报适应当前测试分布上的短期回报更新约束通常靠 PPO 的 KL 或 clip 控制单步幅度需要更强的 KL、熵正则、步数限制失败风险过拟合或训练不稳定策略在线崩溃、业务指标下降典型算法/工具PPO、REINFORCE、RLHFTTPO、test-time adaptation 等思路这里需要特别区分两个容易混淆的概念Test-Time Training测试时训练和 Test-Time Adaptation测试时适应。前者往往指在测试样本上额外构造无监督任务来更新模型后者通常指调整模型表示以适配新分布例如在图像分类中通过熵最小化来适配测试集。TTPO 更偏策略优化它关心的不是模型表征而是动作分布本身并且必须借助动作产生的反馈信号来作为更新依据。用一个生活化的类比来理解训练时策略优化像是高考前的复习把所有可能考到的知识点都练熟测试时策略优化则像是进了考场做完前几道题发现命题风格和模拟卷差别很大于是当场调整自己的答题节奏和时间分配。目标依然是把这份答卷做好但手段是现场调整而不是重新复习三年。3. TTPO 的核心设计问题与原理要设计一个 TTPO 系统先要回答一个根本问题测试阶段到底优化什么目标在训练阶段目标通常清晰明确最大化累计折扣回报比如 (\max_\theta \mathbb{E}_\tau[R(\tau)])。但在测试阶段我们往往拿不到一个完整轨迹的长度回报只能拿到当前动作的即时反馈甚至是一条稀薄的、延迟到几分钟之后才到达的反馈。因此 TTPO 的目标函数经常被写成当前批量的优势估计[ \mathcal{L}{\text{pg}} -\mathbb{E}{(s,a)\sim\pi_{\text{old}}} \left[ \frac{\pi_\theta(a|s)}{\pi_{\text{old}}(a|s)} A(s,a) \right] ]这和 REINFORCE 非常相似。但真正的区别在于测试阶段的优势 (A(s,a)) 往往来自实时反馈比如当前点击是否发生、当前推荐是否被接受而不是来自环境回报函数的长期折现。于是 TTPO 更接近一种“即时反馈策略梯度”。由于在线样本少直接最大化策略梯度很容易导致策略局部过拟合。比如某条交互样本里动作 3 偶然拿到了高回报更新后策略开始疯狂选择动作 3但下一条样本又发现动作 3 实际并没有那么好策略又得弹回来。这种弹跳会让线上业务指标剧烈波动。解决方案通常有两个方向第一是 KL 散度约束。更新前后策略不能差太远如下式[ \min_\theta \mathcal{L}{\text{pg}} \beta \cdot \mathrm{KL}[\pi{\text{old}} | \pi_\theta] ](\beta) 越大更新越保守策略越不容易崩溃但适应新环境的速度也越慢。(\beta) 太小策略又会过度迎合最近几条反馈。第二是限制更新轮数和批量大小。TTPO 不应该“无限训练”而应该每收到一个批量反馈就更新一次并且用“步数上限 收敛判断”约束总更新量。这也是它和常规在线 RL 的一个区别常规在线 RL 可以持续跑很多个 epochTTPO 更适合短促、受限、可回滚的小批量更新。还有一个常被忽略的问题是方差。单条样本的优势信号方差极大尤其是奖励只有 0/1 且命中概率很低时。实践中至少要做两步降方差一是把某个 batch 中的 reward 减去均值得到优势二是对优势做标准化即除以标准差。这也是下面最小实现里要演示的部分。另外TTPO 并不总是适合所有环境。它比较适合具备以下条件的场景测试阶段能很快拿到动作反馈不需要等待数小时。当前策略已经有一定基础能力只是需要微调而不是从零学习。环境反馈虽然带噪声但方向和真实目标基本一致。如果你的业务动作要数小时后才产生收益那么 TTPO 就需要设计一个短期代理奖励否则在线更新信号会严重滞后。4. 环境准备与最小实验配置下面给出一个可以直接运行的 TTPO 最小实验。实验使用无上下文的 5 臂老虎机环境模拟“测试时最佳动作发生偏移”的场景训练阶段策略对 5 个动作基本均匀测试环境的最佳动作从原本的 0 换成了 4。TTPO 需要利用在线反馈把策略分布逐步拉向动作 4。实验环境建议如下版本以你本机实际环境为准Python 3.9 或更高版本。PyTorch 2.x。NumPy。安装依赖pip install torch numpy文件目录可以这样组织ttpo_demo/ ├── scripts/ │ ├── ttpo_policy.py │ ├── ttpo_update.py │ └── run_ttpo_demo.py需要说明的是TTPO 本身没有标准统一实现不同论文和项目在目标函数细节上会有所差异。下面这段代码演示的是最基础、最容易举一反三的“策略梯度 KL 约束 优势基线”版本。5. 完整示例PyTorch 实现一个 TTPO 最小系统5.1 定义策略网络首先定义一个简单的离散动作策略网络# 文件路径scripts/ttpo_policy.py import torch import torch.nn as nn import torch.nn.functional as F from torch.distributions import Categorical class CategoricalPolicy(nn.Module): 最简单的离散动作策略网络输入状态输出动作分布。 def __init__(self, state_dim: int, action_dim: int, hidden_dim: int 32): super().__init__() self.fc1 nn.Linear(state_dim, hidden_dim) self.fc2 nn.Linear(hidden_dim, hidden_dim) self.fc3 nn.Linear(hidden_dim, action_dim) def forward(self, state: torch.Tensor) - Categorical: x F.relu(self.fc1(state)) x F.relu(self.fc2(x)) logits self.fc3(x) return Categorical(logitslogits)这个网络输入一个状态向量输出一个Categorical分布。为了让代码更通用这里把状态维度保留为参数。在后面的老虎机示例中状态是一个固定的一维向量因为这是一个无上下文问题。5.2 定义 TTPO 更新函数TTPO 的关键逻辑在更新函数里。这里实现了一个批量的策略梯度更新并用 KL 散度把新策略约束在旧策略附近。# 文件路径scripts/ttpo_update.py import torch def ttpo_update(policy, optimizer, states, actions, rewards, kl_coef0.1): TTPO 的核心更新函数 在测试阶段拿到一批 (state, action, reward) 后 通过策略梯度 KL 约束的方式更新策略。 参数说明 - states: 长度为 batch_size 的状态列表 - actions: 长度为 batch_size 的动作列表 - rewards: 长度为 batch_size 的奖励列表 - kl_coef: 与历史策略之间的 KL 散度系数越大越保守 states torch.tensor(states, dtypetorch.float32) actions torch.tensor(actions, dtypetorch.long) rewards torch.tensor(rewards, dtypetorch.float32) # 保留更新前的策略分布用于 KL 约束 with torch.no_grad(): old_dist policy(states) dist policy(states) log_probs dist.log_prob(actions) # 用优势函数降低策略梯度方差 advantages rewards - rewards.mean() if advantages.std() 1e-6: advantages advantages / (advantages.std() 1e-8) pg_loss -(log_probs * advantages).mean() kl torch.distributions.kl_divergence(old_dist, dist).mean() loss pg_loss kl_coef * kl optimizer.zero_grad() loss.backward() optimizer.step() return {pg_loss: pg_loss.item(), kl: kl.item(), total_loss: loss.item()}这个函数做了三件重要的事用当前的 reward 计算优势减去均值并做标准化减少单条样本的更新噪声。计算策略梯度损失让看到正向优势的动作概率增加、负向优势的动作概率降低。加入 KL 散度约束防止新策略和旧策略偏离太远。需要注意的是KL 散度一般是不对称的这里用的是优化后分布相对于旧分布的散度实践中也可以用反向 KL 或者对称版本没有绝对标准。关键是确保测试阶段更新保持保守可控。5.3 主程序测试时分布偏移与在线更新主程序模拟了“策略已部署、测试环境最佳动作偏移”的过程# 文件路径scripts/run_ttpo_demo.py import torch import torch.optim as optim from ttpo_policy import CategoricalPolicy from ttpo_update import ttpo_update torch.manual_seed(42) def collect_rollouts(policy, best_action4, n_rollouts10, reward_prob0.8): 模拟测试环境中的交互。 这里假设环境发生了分布偏移以前最佳动作是 0现在是 4。 TTPO 在测试阶段利用 reward 信号把策略慢慢拉向动作 4。 states, actions, rewards [], [], [] state torch.zeros(1, 1) # 无上下文 bandit固定状态 for _ in range(n_rollouts): dist policy(state) action dist.sample().item() hit (action best_action) reward 1.0 if hit and torch.rand(1).item() reward_prob else -0.1 states.append(state.squeeze(0).tolist()) actions.append(action) rewards.append(reward) return states, actions, rewards def run_demo(): policy CategoricalPolicy(state_dim1, action_dim5, hidden_dim16) optimizer optim.Adam(policy.parameters(), lr0.05) # 模拟“测试过程”持续更新 for round_idx in range(5): states, actions, rewards collect_rollouts(policy, n_rollouts10) info ttpo_update( policy, optimizer, states, actions, rewards, kl_coef0.2 ) avg_reward sum(rewards) / len(rewards) with torch.no_grad(): dist policy(torch.zeros(1, 1)) entropy dist.entropy().item() probs dist.probs.squeeze(0).tolist() print( fround{round_idx} avg_reward{avg_reward:.3f} fentropy{entropy:.3f} pg_loss{info[pg_loss]:.3f} fkl{info[kl]:.3f} ) print(f probs{[round(p, 3) for p in probs]}) if __name__ __main__: run_demo()运行命令python scripts/run_ttpo_demo.py代码里有一个关键参数reward_prob0.8表示选择到最佳动作时有 80% 概率拿到正奖励其他动作统一给一个小的负奖励。由于奖励带噪声单次更新的方向并不总是正确所以需要批量更新并通过优势标准化来平滑噪声。5.4 示例输出的解释实验的输出格式大致如下round0 avg_reward-0.010 entropy1.590 pg_loss... kl... probs[0.201, 0.199, 0.202, 0.198, 0.200]随着更新继续熵会逐渐下降动作 4 对应的概率会逐步上升其他动作的概率下降。具体数值会因为随机种子、KL 系数和学习率而不同。如果发现熵快速降到很低的水平意味着更新可能过于激进需要调大kl_coef或调低学习率。这个实验的重点不是跑到完美收敛而是展示一个 TTPO 可以跑通的完整闭环测试时不断采样、不断收到反馈、不断调整策略分布。6. 运行结果与验证方法判断 TTPO 是否生效不能只看 loss 下降。测试阶段最关心的指标是“策略是否在做正确的事”。在演示环境里验证方法有三个平均奖励是否逐步上升。如果策略逐渐偏向动作 4平均奖励应该整体走高。动作 4 的概率是否在上升。这是最直接的信号。策略熵是否没有过早崩到接近 0。如果熵快速下降到 0说明策略已经停止了有效探索也许只是偶然见到了几次动作 4 的正反馈。更严格的验证方法是设置对照组同一套初始策略一组开启 TTPO 更新另一组冻结策略权重。在同样一组模拟环境反馈下持续比较两组的累计奖励。如果 TTPO 组的累计奖励高于冻结组说明测试时更新确实带来了增量收益。如果实验效果不如预期可以按下面的顺序排查先看avg_reward的趋势如果一直很低可能是奖励信号太少增加n_rollouts。再看kl指标如果一直接近于 0说明策略几乎没动调大学习率或调小kl_coef。最后看entropy如果剧烈下降调大kl_coef或减小更新轮数。这里需要特别提醒TTPO 的效果和奖励设计高度相关。如果奖励本身和真实目标不一致再好的在线更新也只是把策略优化到错误方向。因此在真实项目中设计短期代理奖励时要非常谨慎。7. 常见问题与排查思路实践中 TTPO 容易遇到的问题往往不在公式上而在对更新节奏、噪声和业务指标的权衡上。下面整理了一份排查表问题现象可能原因排查方式解决方案策略完全不更新KL 约等于 0KL 系数太大或学习率太小查看 KL 指标和网络梯度调低kl_coef调大lr策略熵骤降动作几乎固定更新过于激进单批噪声被放大查看 entropy 曲线和最近奖励分布调大kl_coef增加n_rollouts减少轮数在线平均奖励波动很大单条奖励方差高批量太小打印最近 20 条 reward 统计增大批大小做优势标准化引入基线新环境指标上升但旧环境指标下降过度拟合测试样本忘记旧分布同时监控新旧环境的验证指标限制更新次数保留原始模型快照设置回滚阈值测试阶段拿不到即时奖励回报周期太长信号稀疏观察数据流延迟设计短期代理奖励或放弃 TTPO更新后策略推荐内容异常KL 约束没兜住策略漂移检查更新前后输出分布的差异增加 KL 约束权重添加动作级差异限制其中“更新导致旧环境指标下降”是最容易被低估的问题。TTPO 本质上是把一个部署在旧分布上的策略逐渐调整成适应新分布的策略。这个过程中旧分布上的表现往往会有一定程度牺牲。所以生产系统必须有一个决策开关当旧环境指标下降到一定阈值时立即停止更新并回滚到上一版本策略。8. 在真实项目中的最佳实践与工程建议8.1 先判断分布偏移确实存在TTPO 不是万能的增强器。如果测试分布和训练分布基本一致固定策略已经足够好运行 TTPO 会带来额外风险还可能因为在线噪声把好好的策略改坏。建议先做分布检测比较线上请求特征和训练特征的分布差异或者在灰度环境中对比开启与关闭 TTPO 的指标差异。8.2 必须保留模型快照在测试阶段更新策略之前把原始策略权重完整保存下来包括优化器状态如果可能的话。一旦发现更新后指标下降能立刻恢复。快照可以在内存里保存几份分别代表“更新 0 轮”“更新 5 轮”“更新 10 轮”的状态。8.3 更新次数要有限制测试时更新不是训练时训练不应该无限跑。一般做法是设置一个更新轮数上限例如“同一批次反馈最多更新 3 个 gradient step”并且每轮都检查 KL 散度是否超过某个阈值。如果 KL 超过阈值就停止本轮更新。8.4 把 KL 约束和熵正则同时用上KL 约束负责约束新旧策略差异熵正则负责保留最低限度的探索能力。尤其在奖励噪声很大的时候熵正则能避免策略过早陷入“单点确认偏误”。最小实现里没有写熵正则但在真实系统中强烈建议加入loss pg_loss kl_coef * kl - entropy_coef * dist.entropy().mean()这个负号很关键。熵越大策略的探索性越强所以要减去熵的均值让优化器去提高熵。8.5 监控和回滚生产环境里的 TTPO 需要三套监控策略分布监控每一轮更新后输出分布和上一版的分布差异。业务指标监控新场景指标和旧场景指标分开看。更新信号监控奖励有效性是否持续下降是否出现奖励噪声过大。只要发现业务指标异常波动立刻停止更新并回滚到上一版本策略。测试时优化的特点是“快”但也意味着“反应时间短”监控必须跟上。8.6 在大模型场景中的扩展思考TTPO 的思路也可以迁移到大模型推理阶段。比如对话系统在线上收到用户反馈把“这次回答是否被点赞”作为即时奖励用策略梯度方式对生成策略做微调。但大模型推理成本更高在线更新也更危险通常不建议直接对基座模型做全量更新而是在推理链路外层维护一个可插拔的轻量策略层。用 TTPO 更新轻量策略层而不是更新大模型本体。更新前做严格的输出安全评测防止策略漂移影响内容质量。这和最近讨论比较多的“推理时扩展”思想并不冲突推理时扩展是在测试阶段让模型生成更多思考或搜索候选TTPO 则是在测试阶段根据反馈微调策略本身二者可以结合使用。9. 总结与实践建议TTPO 解决的核心问题不是“训练不好”而是“部署之后模型无法继续适应新环境”。它把策略优化的窗口延长到测试阶段让策略可以基于真实反馈做小步快的改进。相比重新训练它的成本更低但相比静态部署它需要更完善的安全机制和监控体系。如果接下来你想自己动手实践建议按这样的顺序先跑通上面的最小示例把策略分布的变化过程打印出来从感性上理解 TTPO 的一次更新到底发生了什么。把示例环境换成一个更贴近你业务的模拟器例如带上下文的推荐环境。加入对照组实验验证 TTPO 带来的净收益。在灰度环境开启 TTPO配合快照和回滚机制小流量试运行。逐步把更新限制条件引入系统做好 KL 约束、熵正则、监控告警后再扩大规模。测试时策略优化的价值不在“多聪明”而在于它给部署系统增加了一个低成本纠偏通道。如果你正面对线上策略性能下滑但回炉训练成本过高的问题TTPO 是值得一试的方向。只要把快照、监控和回滚这三件事做好它带来的收益通常能超过它的实现成本。