首页 / 资讯中心 / 文章详情

SLAC强化学习实战:随机潜在变量模型提升样本效率

SLAC强化学习实战:随机潜在变量模型提升样本效率 ★ FEATURED ARTICLE
强化学习在连续控制任务里有个老毛病样本效率低得让人抓狂。尤其是那些高维观测输入比如图像、多模态传感器读数的场景模型动辄要几百万步交互才能学到像样的策略训练成本高到很多团队根本玩不起。我最早接触随机潜在变量模型这条路线就是因为在一个机械臂抓取项目里被样本效率卡了脖子——环境交互一次要几秒钟一天跑下来也就攒个几万条数据用常规的无模型算法根本喂不饱网络。后来把目光转向“在潜空间里做强化学习”这个思路才算是找到了突破口。Stochastic Latent Actor-Critic简称 SLAC就是这条路线里一个相当有代表性的工作它把随机潜在变量模型和 actor-critic 框架揉在一起专门解决高维观测下样本效率低的问题。这篇博文我会从设计思路、核心原理、实操复现、踩坑排查几个维度把它拆开讲透适合有一定强化学习基础、想在自己的项目里落地潜空间策略学习的同学参考也适合刚入门想搞清楚“为什么要在潜空间做 RL”这个问题的朋友。1. 为什么要在潜空间里做强化学习1.1 高维观测带来的样本效率困境先说清楚问题背景。在标准的无模型强化学习设定里智能体直接从原始观测比如一张 84x84 的图片展开就是七千多维映射到动作。这个映射关系由神经网络来拟合听起来很直接但实际训练时问题一大堆。最典型的就是维度灾难观测维度越高策略网络和价值网络需要拟合的函数就越复杂需要的样本量呈指数级上升。你在 Atari 上跑 DQN 能work是因为那些任务的状态空间相对规整、奖励信号密集一旦换到真实机器人控制或者复杂物理仿真观测维度上去了、奖励又稀疏直接端到端学习基本就是灾难。另一个隐性问题是观测里混杂了大量与决策无关的信息。一张机械臂视角的图像里背景纹理、光照变化、相机噪声占了绝大部分像素真正决定动作的只有机械臂末端和目标的相对位置那一小块区域。端到端网络被迫把这些无关信息也编码进去等于让网络在噪声里找信号样本效率自然上不去。我实测过一个对比同样的任务用原始图像输入和用人工提取的低维状态输入达到相同性能所需的交互步数差了将近一个数量级。这就是为什么大家要想办法先把观测“压缩”成紧凑的表示再在这个表示上做决策。潜空间方法的核心动机就在这儿用一个学习到的低维潜在状态来概括原始观测把决策问题从高维观测空间搬到低维潜空间。这样一来策略网络面对的输入维度可能从几千维降到几十维函数拟合难度大幅下降样本效率自然提升。而且潜空间表示是通过无监督或自监督的方式从数据里学出来的不需要人工设计特征泛化性也更好。1.2 随机潜在变量模型相比确定性编码的优势把观测压缩成潜表示最朴素的做法是用一个确定性编码器比如普通的自编码器把观测映射成一个固定向量。但 SLAC 选择的是随机潜在变量模型也就是潜状态本身是一个概率分布而不是一个点。这个选择不是拍脑袋的背后有明确的理由。确定性编码有个致命弱点它无法表达观测的不确定性。真实环境里传感器有噪声、部分观测会遮挡、同一时刻的观测可能对应多种合理的潜在状态。确定性编码器会强行把这些不确定性压成一个点估计导致学到的表示对噪声敏感、对未见过的观测泛化差。更麻烦的是在部分可观测的设定下POMDP当前观测不足以确定真实状态必须结合历史信息推断这时候一个点估计根本没法表达“我对当前状态有多确定”。随机潜在变量模型把潜状态建模成一个分布 $q(z_t \mid \text{history})$用采样的方式得到具体的潜状态。这样做的好处是模型能显式地表达不确定性采样过程相当于一种正则化让策略对潜状态的扰动更鲁棒同时随机潜变量配合序列模型比如 RNN可以自然地处理部分可观测性把历史观测的信息累积到潜状态的分布里。SLAC 正是利用了这一点用一个序列潜变量模型来同时完成表示学习和状态推断。提示随机潜变量模型不是万能药。如果你的任务是完全可观测的、观测维度本来就不高那上随机潜变量模型属于杀鸡用牛刀反而增加训练难度。它真正发挥价值的地方是高维观测加部分可观测的组合场景。1.3 SLAC 的整体定位与它解决的问题把 SLAC 放到整个方法谱系里看它属于基于模型的表示学习 无模型策略优化的混合路线。注意这里有个容易混淆的点SLAC 虽然用了潜变量模型但它并不是严格意义上的“基于模型的强化学习”MBRL因为它不直接用学到的模型去做规划或者生成虚拟轨迹来训练策略。它的做法是先用潜变量模型学一个好的表示然后在这个表示空间里跑标准的 actor-critic 算法。所以它继承了无模型方法的稳定性和通用性又拿到了表示学习带来的样本效率红利。具体来说SLAC 要解决三个层面的问题。第一是表示层面如何从高维观测序列里学到一个紧凑、信息充分、能表达不确定性的潜状态。第二是推断层面在部分可观测条件下如何根据历史观测在线推断当前潜状态的后验分布。第三是控制层面如何在潜空间里高效地做策略评估和策略改进同时保证训练稳定。这三个层面环环相扣任何一个环节出问题都会拖垮整体性能后面的章节我会逐一拆解。2. 核心架构拆解潜变量模型与 actor-critic 如何咬合2.1 序列潜变量模型的结构设计SLAC 的潜变量模型是一个序列化的随机模型核心结构可以理解成三层编码器、转移模型、解码器。编码器负责把当前观测 $o_t$ 编码成潜变量的近似后验参数转移模型描述潜状态随时间的演化 $p(z_t \mid z_{t-1}, a_{t-1})$解码器负责从潜状态重建观测 $p(o_t \mid z_t)$。这三者构成一个变分自编码器VAE的序列版本训练目标是最大化观测的对数似然下界ELBO。这里的关键设计在于潜状态的时序依赖。普通的 VAE 对每一帧独立编码丢掉了时间信息SLAC 的潜状态 $z_t$ 同时依赖当前观测和历史潜状态通过一个循环结构实践中常用 GRU 或 LSTM 来承载历史信息把过去的信息压缩进隐状态。这样即使当前观测有遮挡或者噪声模型也能借助历史推断出合理的潜状态。我在复现时对比过带时序和不带时序两个版本在部分可观测的抓取任务上带时序的版本收敛速度快了大约 40%最终成功率也高出十几个百分点差距非常明显。从实现角度看编码器输出的不是潜状态本身而是高斯分布的均值和对数方差然后通过重参数化技巧采样得到 $z_t$。这个重参数化是让梯度能穿过采样操作回传的关键具体做法是 $z \mu \sigma \odot \epsilon$其中 $\epsilon$ 从标准正态分布采样。这个技巧看似简单但它是整个模型能端到端训练的基础少了它梯度就断了。2.2 潜空间中的 actor-critic 设计表示学好了接下来是在潜空间里做控制。SLAC 的 actor 和 critic 都以潜状态 $z_t$ 为输入而不是原始观测。Actor 网络输出动作分布的参数连续动作通常是高斯分布的均值和方差critic 网络输出状态价值或动作价值。这里有个设计细节值得说critic 的输入是潜状态加动作而不是像某些方法那样直接用观测。这样做的好处是输入维度低、函数简单但也带来一个隐患——如果潜状态丢失了决策相关的信息critic 的估计就会有偏。所以潜变量模型的质量直接决定了控制性能的上限这也是为什么 SLAC 在表示学习上下了很大功夫。另一个关键点是潜空间的奖励信号处理。奖励是从环境里拿到的原始标量它不依赖潜状态所以可以直接用在潜空间的时序差分学习里。SLAC 用的是类似 SACSoft Actor-Critic的最大熵框架在 TD 目标里加入了熵项鼓励策略保持一定的随机性。这个选择对训练稳定性帮助很大我在早期版本里试过不带熵项的普通 actor-critic策略很快就塌缩到确定性动作探索能力丧失在稀疏奖励任务上直接卡死。2.3 两个模块的联合训练与梯度流动SLAC 最精妙也最容易出问题的地方是潜变量模型和 actor-critic 的联合训练方式。理论上两者可以分开训先训好 VAE冻结编码器再训策略。但这样有个问题——VAE 的表示是为重建观测优化的不一定对控制任务友好。SLAC 采用的是联合优化让策略学习的梯度能够回传到编码器促使表示朝着“有利于决策”的方向调整。具体实现上actor 的损失函数里包含了对潜状态的依赖梯度会经过采样操作回传到编码器参数。这里重参数化技巧再次发挥作用保证了梯度路径的畅通。不过联合训练也带来了稳定性挑战表示学习和策略学习的目标不一致训练早期表示还很粗糙策略梯度噪声大容易把编码器带偏。SLAC 的应对策略是给不同损失项配不同的权重并且用目标网络target network来稳定 critic 的更新。这些超参数的调节我在第 4 章会详细讲这里先记住一个原则表示学习的损失权重要足够大保证潜空间不会在策略优化的过程中退化。模块输入输出训练目标关键技巧编码器观测 $o_t$ 历史隐状态潜变量后验参数ELBO 重建项 KL 项重参数化采样转移模型$z_{t-1}$ $a_{t-1}$$z_t$ 的先验分布KL 散度对齐时序循环结构解码器$z_t$重建观测重建误差与编码器共享部分结构Actor$z_t$动作分布参数最大化 Q 值 熵潜空间输入Critic$z_t$ $a_t$Q 值TD 误差双 Q 网络 目标网络3. 从零复现 SLAC 的关键步骤3.1 环境与依赖准备复现 SLAC 之前先把环境搭好。我用的技术栈是 Python 3.8 以上、PyTorch 1.10 以上仿真环境选的是 MuJoCo 系列的连续控制任务比如 HalfCheetah、Walker2d 这类因为 SLAC 原文就是在这些任务上验证的。如果你手头没有 MuJoCo 授权也可以用 PyBullet 的对应任务替代接口基本兼容只是物理参数略有差异不影响方法本身的验证。依赖清单大致如下torch、numpy、gym、mujoco-py或pybullet、tensorboard用于训练监控。我建议额外装一个tqdm来看训练进度长任务没有进度条会很难受。环境配置这块有个坑要提醒MuJoCo 的版本和 gym 的版本耦合比较紧装之前先确认版本兼容性不然会卡在 import 阶段。我踩过一次gym 0.21 配 mujoco-py 2.1 死活跑不起来换成 gym 0.18 加 mujoco-py 2.0 才通。pip install torch numpy gym tqdm tensorboard pip install mujoco-py2.0.2.13装完之后跑一个简单的环境交互测试确认env.reset()和env.step()能正常返回观测维度和动作维度符合预期。这一步别省很多后续的维度不匹配错误都是环境没配好导致的。3.2 潜变量模型的实现要点潜变量模型是整个复现的核心我把它拆成编码器、转移模型、解码器三块分别实现。编码器用一个 MLP 加 GRU 的组合MLP 先把观测压成一个特征向量GRU 维护历史隐状态最后输出潜变量的均值和对数方差。这里有个维度设计的经验潜变量维度不要设太大原文在 MuJoCo 任务上用的是 8 到 16 维我实测 16 维在大部分任务上够用设到 32 维以上反而容易过拟合、训练变慢。转移模型是一个以 $z_{t-1}$ 和 $a_{t-1}$ 为输入的小 MLP输出下一时刻潜变量的先验分布参数。解码器则是从 $z_t$ 重建观测的 MLP输出维度等于观测维度。训练损失由三部分组成重建误差高斯观测用 MSE二值观测用交叉熵、KL 散度后验和先验之间、以及一个可选的时序一致性项。KL 项的权重需要仔细调太小了潜空间不规范太大了重建质量下降我一般从 1.0 开始试根据重建损失的量级微调。class LatentModel(nn.Module): def __init__(self, obs_dim, action_dim, latent_dim16, hidden_dim256): super().__init__() self.encoder nn.Sequential( nn.Linear(obs_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU() ) self.rnn nn.GRUCell(hidden_dim, hidden_dim) self.fc_mu nn.Linear(hidden_dim, latent_dim) self.fc_logvar nn.Linear(hidden_dim, latent_dim) self.transition nn.Sequential( nn.Linear(latent_dim action_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, latent_dim * 2) ) self.decoder nn.Sequential( nn.Linear(latent_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, obs_dim) )上面这段是骨架代码实际实现时要注意 GRU 的隐状态要在序列维度上正确传递别把 batch 维和时间维搞混了。我第一版就是这里写错了导致时序信息完全没利用上训练出来的模型和单帧 VAE 没区别白跑了两天。3.3 actor-critic 网络的搭建与超参数选择Actor 和 critic 都是简单的 MLP输入是潜状态critic 额外拼接动作。Actor 输出高斯动作分布的均值和标准差标准差要经过一个 softplus 或者 exp 变换保证为正并且做上下界裁剪防止数值爆炸。Critic 用双 Q 网络结构取两个 Q 值的最小值作为目标这是 SAC 系列的标准做法能有效缓解 Q 值高估问题。超参数方面我整理了一份实测比较稳的配置学习率 actor 和 critic 用 3e-4潜变量模型用 1e-3表示学习通常可以快一点batch size 用 256目标网络软更新系数 tau 取 0.005熵温度系数用自动调节SAC 的自适应 alpha初始值 0.1。经验回放池大小设 1e6但实际在样本效率敏感的任务上1e5 到 2e5 就够启动训练了。这些数值不是金科玉律不同任务需要微调但作为起点能省不少试错时间。超参数推荐值调节范围说明actor/critic 学习率3e-41e-4 ~ 1e-3太大不稳定太小收敛慢潜变量模型学习率1e-35e-4 ~ 3e-3表示学习可略快batch size256128 ~ 512影响梯度噪声软更新系数 tau0.0050.001 ~ 0.01控制目标网络更新速度潜变量维度168 ~ 32过高过拟合过低欠拟合回放池大小1e61e5 ~ 1e6视内存而定3.4 训练循环的组织与关键细节训练循环的组织方式直接影响复现成败。我的做法是每个环境步做一次模型更新但潜变量模型的更新频率可以比 actor-critic 高一些因为表示学习需要更多数据。具体来说每收集一个 batch 的数据潜变量模型更新两次actor-critic 更新一次。这个比例是我试出来的表示学习欠拟合时策略学不动所以给它更多更新预算是合理的。序列采样这块有个细节从回放池里采样时要采连续的序列片段而不是独立帧否则 RNN 的历史信息就断了。我一般采长度为 8 到 16 的子序列太短了时序信息不足太长了显存吃不消。采样的时候要注意处理序列边界episode 结束的地方要重置 RNN 隐状态不然会把上一个 episode 的信息错误地串到下一个。注意训练初期不要急着看策略性能。潜变量模型需要几千步才能学出一个像样的表示这之前 actor-critic 基本是在噪声上学习。我一般等潜变量模型的重建损失降到稳定区间后才开始认真评估策略否则容易被早期的糟糕曲线误导。4. 实操中的典型问题与排查手册4.1 潜空间坍塌表示退化的识别与修复潜空间坍塌是复现 SLAC 时最常见的坑表现是潜变量模型的重建损失看着还行但潜变量本身几乎不携带信息——不同观测编码出来的潜状态几乎一样策略网络拿到的输入等于常数自然学不出东西。这个问题在 KL 项权重过大时特别容易出现因为模型为了最小化 KL 散度干脆让后验等于先验潜变量退化成纯噪声。识别方法很简单训练过程中定期统计潜变量在不同样本间的方差如果方差趋近于零基本就是坍塌了。修复手段有几个降低 KL 项权重、用 KL 退火训练初期 KL 权重设小逐渐增大、或者改用 free bits 技巧给每个潜变量维度设一个最小的 KL 下界低于这个下界就不惩罚。我实测下来 KL 退火最稳从 0.1 的权重开始线性增加到 1.0坍塌问题基本消失。4.2 训练不稳定的三大诱因除了坍塌训练不稳定还有几个常见诱因。第一个是梯度爆炸RNN 结构在序列较长时容易出现解决办法是加梯度裁剪阈值设 1.0 到 5.0 之间我一般用 5.0。第二个是Q 值发散critic 的估计值越来越大最后变成 NaN这通常是目标网络更新太快或者学习率太大导致的把 tau 调小、学习率降一档通常能缓解。第三个是策略熵塌缩策略过早变得确定探索不足这时候检查熵温度系数是不是太小或者自动调节的 alpha 是不是卡在下界了。问题现象可能原因排查方法解决手段潜变量方差趋零KL 权重过大统计潜变量跨样本方差KL 退火 / free bits损失变 NaN梯度爆炸监控梯度范数梯度裁剪Q 值持续增大目标网络更新过快打印 Q 值均值调小 tau / 降学习率策略过早确定熵系数过小监控策略熵调大 alpha 初值重建模糊解码器容量不足看重建可视化加宽解码器 / 降 KL 权重4.3 样本效率不达预期的诊断路径有时候训练能跑通、也不崩但样本效率就是不如预期跑了几十万步性能还上不去。这时候要按顺序排查几个环节。先看潜变量模型的重建质量如果重建出来的观测和原始观测差太多说明表示本身就没学好后面全是白搭。再看潜状态和真实状态的对应关系可以在简单任务上做个可视化看看潜空间是否保留了决策相关的维度。最后看 actor-critic 的更新是否正常Q 值有没有在合理范围内波动、策略熵有没有维持在健康水平。我遇到过一次样本效率异常低的情况排查了半天发现是回放池采样时序列片段的重叠处理有问题导致同一段数据被反复采样有效样本多样性不足。改成按 episode 为单位采样、保证片段不重叠之后样本效率立刻恢复正常。这种问题不会有任何报错只能靠对训练曲线的敏感度去发现所以训练监控一定要做细。4.4 部分可观测场景下的额外注意事项如果你的任务本身是部分可观测的比如观测有遮挡、只有局部视野SLAC 的时序结构会发挥更大作用但也有额外要注意的地方。RNN 隐状态的初始化要合理episode 开始时重置为全零序列长度要足够覆盖任务的关键历史窗口太短了推断不出真实状态。另外部分可观测场景下潜变量的不确定性会更大KL 项的调节要更保守否则模型会倾向于忽略观测、只靠先验瞎猜。我在一个带遮挡的导航任务上试过把序列长度从 8 加到 16成功率提升了将近 20%这个投入产出比很划算。5. 我踩过的坑与几条实操心得5.1 关于表示学习和策略学习的平衡联合训练最大的难点是平衡。我早期犯过一个错误为了让策略尽快出效果把 actor-critic 的学习率调得很高结果策略梯度把编码器带偏了潜空间朝着“有利于当前策略”而不是“有利于表示”的方向退化最后策略和表示一起崩。后来我调整策略让表示学习先行——训练前一万步只更新潜变量模型等表示稳定了再引入策略梯度整体稳定性好了很多。这个“预热”阶段看起来浪费了样本实际上省下了后面反复调参的时间。5.2 监控指标比损失曲线更重要光看总损失是不够的SLAC 的训练过程里有很多隐性问题不会直接反映在总损失上。我建议至少监控这几个量潜变量的跨样本方差防坍塌、重建误差的绝对值防表示退化、Q 值的均值和方差防发散、策略熵防探索不足、以及梯度范数防爆炸。这些指标用 tensorboard 画出来训练异常时一眼就能定位到是哪个环节出了问题。我现在的习惯是训练脚本里把这些监控全部埋好跑起来之后先看半小时曲线再决定要不要继续。5.3 小规模验证再上大任务SLAC 的超参数对任务比较敏感直接上复杂任务调参成本太高。我的做法是先在简化版任务上把流程跑通、把超参数调到一个大致合理的区间再迁移到目标任务上微调。简化版任务可以是降维的观测、缩短的 episode、简化的奖励只要能验证表示学习和策略学习的基本流程就行。这个习惯帮我省下了大量算力也避免了在复杂任务上被各种耦合问题搞晕。最后分享一个我觉得挺有用的小技巧潜变量模型训练好之后可以单独拿出来做可视化分析把潜状态用降维方法比如 PCA 或 t-SNE画出来看看不同状态在潜空间里的分布。如果同类状态聚在一起、不同状态分得开说明表示质量不错如果混成一团那策略再调也白搭。这个分析花不了多少时间但对判断问题出在表示层还是控制层特别有帮助。
阅读完成 · 觉得有帮助?
咨询建站