SERL：让真机强化学习从“难用”走向“可复现”的强化学习框架 ----

0x00 概要SERL 的核心使命是在真实世界中让机器人在 20-40 分钟内学会高精度的机械操作。它通过集成 SAC、RLPD、DrQ 和 VICE将原本需要数百万次尝试的 RL压缩到了人类演示水平的量级。RLPDReinforcement Learning with Prior Data是一种基于 off-policy 的 actor-critic强化学习算法借鉴了 soft-actor-critic 等时序差分算法的成功经验但为满足上述需求做出了一些关键修改它其实就是 SAC 之前的数据Prior Data 极高的更新频率High UTD。注本系列的最终目标是“通过一系列相关项目/算法的解读来深入学习/分析/反推 LWDLearning while Deploying这篇论文的机理和可能实现”。之所以从SERL入手是因为 SERLHIL-SERLSOP没有开源都是罗剑岚博士的一系列论文可以从中管窥作者的思路脉络。本文依然是从工程/论文进行反推还请读者不吝指出问题多谢。0x01 基础 背景1.1 总体流程图SERL 的总体流程图如下其中SACSoft Actor-Critic是算法底座 是整个系统的引擎。RLPD (RL with Prior Data) 是性能加速器是 SERL 实现20 分钟学会的关键它通过暴力更新和不忘初心来榨取每一条数据的价值。High UTD迫使模型对有限的样本进行 深度研读从而在极短的物理时间内捕捉到成功的信号。1.2 面临的问题在许多场景中强化学习的优异表现依赖于与环境进行大量的在线交互这通常通过使用模拟器来实现。然而在实际问题中常常面临样本获取成本高昂的情况。此外奖励信号稀疏且高维的状态和动作空间往往使这一问题更加严重。一些先前的研究致力于通过预训练利用这些数据而其他方法则在在线训练时引入约束以应对分布转移问题。然而每种方法都有其缺点例如需要额外的训练时间和超参数或者在行为策略之外的提升有限。0x02 RLPD 基础RLPD关注的是是否可以在 在线学习时直接应用现有的离策略方法以充分利用离线数据。在每一步训练中RLPD 在先验(离线)数据和on-policy数据之间等概率采样以形成一个训练批次即“对称采样”即每个批次有50%的数据来自在线回放缓冲区另外50%来自离线数据缓冲区先验数据。我们对比如下原生纯靠自己试SERL在 SAC 的 Batch 里塞进人类演示数据相当于给 SAC 考试时递了一张带有参考答案的纸让它不用从头瞎猜那RLPD这四个字母贵在哪里秘密不在怎么算Loss Function而在喂什么Batch Composition。RLPD 50/50 抽样 强行保证每个 Batch 都有 128 条 Demo。这相当于给机器人装了一个强制记忆模块让它每一秒钟都在看正确答案。2.1 RLPD 的三大支柱High UTDUpdate-to-DataSERL 每跑一步环境就进行 20 次网络更新。这种暴力刷题让每一条物理数据都被反复压榨是 20 分钟收敛的硬件级保证。Layer NormalizationLN问题在普通的 SAC 中, 网络通常就是纯 MLP。但在 RLPD 中, 由于我们要做 High UTD (高频更新), Q 值的估算非常容易发散。解决方案因此RLPD 会在 Critic 网络 (甚至 Actor 网络) 的每一个隐藏层之后, 都加上 LayerNorm。直观理解高 UTD 会产生巨大的梯度冲击。LN 就给高速赛车装上强力悬挂确保每一层神经元的输出在 20 倍更新频率下依然稳定不至于飞出去。50/50 混合采样做法在 update 过程中, 我们不再只从 replay_buffer 抽数据, 而是将演示数据Demo与在线数据Online强行等比例混合确保智能体在探索新坑时始终能看到正确路标。抽 128 个在线数据 (自己跑出来的)。抽 128 个先验数据 (人类演示的 Demo)。把它们拼成一个 256 的 Batch 进行训练。为什么有效?演示数据告诉智能体正确的路长啥样, 在线数据告诉智能体这里的坑别踩。两者结合, 收敛速度能提升数倍。2.2 Prior Data 策略Prior Data 策略决定了演示数据与在线数据的黄金采样配比。真机强化学习最怕冷启动— 机器人像没头苍蝇一样乱撞。SERL 引入了 RLPDReinforcement Learning with Prior Data机制数据混洗在训练的每一个 Batch 中系统会强制性地混合50% 的演示数据这是由人类老师录制好的标准答案50% 的在线数据这是机器人自己折腾出来的实战记录算法价值这种混合采样确保了模型在进化的每一秒都在不断对比正确答案与自己的尝试。它解决了强化学习初期的探索困境让机器人即使在完全没拿高分的情况下也能通过模仿专家数据来迅速建立起任务的初步认知。我们可以这样理解prior data 像老师给的标准答案online replay 像学生自己的练习记录每次训练都同时看标准答案和自己的错题策略就不会偏离任务太远这也是 SERL 样本效率高的关键之一。它不是让机器人从零开始乱试而是在 demonstrations 的引导下进行强化学习微调。论文中明确写到每次更新使用 sample-based approximation其中 half of the samples drawn from prior datahalf drawn from replay buffer。2.3 算法流程图下面是根据 rlpd.py 的代码逻辑整理的 RLPD 训练流程逻辑图。这个图展示了从数据采样到网络更新的完整路径, 特别标注了 RLPD 相对于普通 SAC 的核心改进点 (如 BC Loss 和 Pessimistic Backup)。2.4 代码细节RLPD 关键组件说明Ensemble Q (10 Qs): 使用 10 个 Q 网络而非 2 个, 增加评估的多样性。rho (Pessimism): 通过 均值 - rho * 标准差 来实现对不确定区域的惩罚。BC Loss: 在 Actor 更新中加入行为克隆, 强迫智能体初始阶段不要偏离演示数据。High UTD (Scan): JAX 的lax.scan允许在一个硬件循环内执行 20 次上述流程。LayerNorm (代码内部) 在所有隐藏层后强制执行归一化, 支撑高频更新。结合 rlpd.py 代码的深度解读:关于 rho 的计算:next_qs self.network.select(target_critic)(batch[next_observations][..., -1, :], next_actions) next_q next_qs.mean(axis0) - next_qs.std(axis0) * self.config[rho]这是 RLPD 的精髓。普通的 SAC 是 min(Q1, Q2), 而这里是用标准差来量化不确定性。如果 10 个 Q 网络对某个状态动作意见不统一 (std 大), next_q 就会被压得很低。关于 BC Lossbc_loss -(dist.log_prob(jnp.clip(batch_actions, -1 1e-5, 1 - 1e-5)) * batch[valid][..., -1]).mean() * self.config[bc_alpha]这行代码在告诉 Actor: 不管 Q 值怎么说, 你输出的动作最好和 Buffer 里的真实动作 (演示数据) 接近一些。这对机器人任务极其关键, 因为它防止了机器人在训练初期因为乱甩而撞坏硬件。关于 High UTD:jax.jit def batch_update(self, batch): agent, infos jax.lax.scan(self._update, self, batch)这里使用了 JAX 的 scan 原语。这比 Python 的 for 循环快得多, 它能把 20 次更新编译成一个高效的 GPU 算子。演示数据Prior DataSERL 所基于的 RLPD 算法中作者发现最简单、最有效的办法就是一视同仁比如从 Replay Buffer自己跑的数据取 128 个。从 Demo Buffer演示数据取 128 个。凑成一个 256 的 Batch直接丢进 Loss 函数。意义不需要复杂的权重计算这种简单的对半开采样就能极大地提升效率。看未来与看现在的逻辑rlpd.py 中计算 next_q 用的是 target_critic, 而计算 actor_loss 时用的是 critic (当前网络)。Target Critic (看未来): 用于计算r γ Q_{target}。由于 Q_{target} 更新得很慢 (Soft Update), 它提供了一个稳定的地基, 防止 Q 值计算产生正反馈螺旋 (即自己把自己估高)。Critic (看现在): 用于 Actor 的更新。Actor 问: 我现在的动作好吗? 。由于 Critic 正在被最快地训练, 它能给 Actor 提供最及时的反馈。矛盾解决: 这就是评估要稳 (Target), 改进要快 (Current)的权衡。0x03 BC在 SERL 的复现中, 通常的步骤是:收集 20 个 Demo。跑 bc.py 进行预训练, 让机器人学会手往哪放。跑 rlpd.py 进行正式训练, 利用 BC 训练好的模型作为起点, 通过 50/50 采样快速进化。在 SERL 的完整流程中bc.py 的代码非常关键, 它揭示了 SERL 系统中 Behavioral Cloning (BC) 环节是如何运作的。3.1 BCAgentBCAgent 负责预训练冷启动是纯监督学习的行为克隆实现架构最为简洁。SERL先用 BC 模仿 Demo让机器人学会手往哪放再开启 RL 寻找怎么抓取。3.1.1 BCAgent 在SERL中的作用冷启动为RL算法提供初始策略演示数据利用从专家演示中学习安全基线在RL训练初期提供安全策略核心优势BCAgent 的简洁性使其成为从演示到强化学习的理想桥梁通过监督学习快速获得可用的策略然后可以在此基础上进行RL微调。BCAgent 的特性如下特性BCAgent说明网络数量仅1个Policy网络无Critic无Temperaturetanh_squashFalse不使用tanh压缩输出分布MultivariateNormalDiag标准高斯分布训练目标最小化MSE 负对数似然监督学习3.1.2 BCAgent 核心组件唯一的 Policy 网络network_kwargs[activate_final] True networks { actor: Policy( encoder_def, # 视觉编码器 MLP(**network_kwargs), # 默认 [256, 256] action_dimactions.shape[-1], tanh_squash_distributionFalse, # 关键差异 ) }组件输入网络结构输出特点Policy图像观测编码器MLP[256,256]动作分布(μ,σ)纯监督学习3.1.3 BCAgent 编码器架构small 编码器encoders { image_key: SmallEncoder( features(32, 64, 128, 256), kernel_sizes(3, 3, 3, 3), strides(2, 2, 2, 2), paddingVALID, pool_methodavg, bottleneck_dim256, spatial_block_size8, ) }resnet 编码器encoders { image_key: resnetv1_configs[resnetv1-10]( pooling_methodspatial_learned_embeddings, num_spatial_blocks8, bottleneck_dim256, ) }resnet-pretrained 编码器pretrained_encoder resnetv1_configs[resnetv1-10-frozen]( pre_poolingTrue, ) encoders { image_key: PreTrainedResNetEncoder( pooling_methodspatial_learned_embeddings, num_spatial_blocks8, bottleneck_dim256, pretrained_encoderpretrained_encoder, ) }3.1.4 BCAgent 损失函数def loss_fn(params, rng): # 前向传播 dist self.state.apply_fn( {params: params}, batch[observations], temperature1.0, trainTrue, rngs{dropout: key}, nameactor, ) pi_actions dist.mode() # 预测动作 log_probs dist.log_prob(batch[actions]) # 对数概率 # 多重损失 mse ((pi_actions - batch[actions]) ** 2).sum(-1) # MSE损失 actor_loss -(log_probs).mean() # 负对数似然 return actor_loss, { actor_loss: actor_loss, mse: mse.mean(), }3.1.5 BCAgent 动作采样def sample_actions(self, observations, seedNone, temperature1.0, argmaxFalse): dist self.state.apply_fn( {params: self.state.params}, observations, temperaturetemperature, nameactor, ) if argmax: actions dist.mode() # 确定性采样 else: actions dist.sample(seedseed) # 随机采样 return actions3.2 bc.py 流程图BC Agent (模仿学习) 核心流程图如下。BC 关键组件说明ResNet-10: SERL 的默认视觉骨架用于从原始像素中提取物理特征。Random Crop: 极其重要的 Trick通过对画面进行 ±4 像素的裁剪来模拟环境扰动。TanhNormal: 动作分布模型确保输出的动作符合机械臂的物理范围。Pre-training: BC 在 SERL 中扮演“冷启动”的角色将 RL 的搜索空间缩减到目标附近。3.3 bc.py 技术细节核心逻辑: update 函数不仅仅是回归: 虽然它计算了 mse, 但实际更新用的是 actor_loss -log_probs.mean()。这是一种概率视角下的模仿: 让策略在演示数据给出的状态下, 输出演示动作的概率尽可能大。没有 Critic这里完全没有 Q 网络。BC 纯粹是看着答案抄答案, 不需要奖励信号。视觉处理: 数据增强 (Data Augmentation) data_augmentation_fn这是 DrQ-v2 风格的随机裁剪。为什么重要?: 在机器人任务中, 摄像头画面可能会有轻微抖动。通过对演示图片进行随机裁剪, 可以让策略学会忽略这种位移, 提高鲁棒性。这也是 SERL 能在 20 分钟内学会任务的秘诀之一。网络架构: ResNet-10在 create 方法中, 它定义了三种编码器。SERL 默认推荐的是 resnet:elif encoder_type resnet: encoders { image_key: resnetv1_configs[resnetv1-10](...) }ResNet-10 比传统的 CNN 效果好得多, 因为它能提取更深层的特征, 同时又不像 ResNet-50 那样运算缓慢。为什么会有 mse 却不用它更新?在 update 函数的 loss_fn 中, 作者计算了 mse ((pi_actions - batch[actions]) ** 2).sum(-1), 但返回值的第一项 (真正的 Loss) 是 actor_loss -(log_probs).mean()。原因: MSE 只关心均值对不对, 而 log_prob 关心的是整个概率分布。如果人类演示同一个动作时有微小的偏差, log_prob 能更好地捕获这种容错性。EncodingWrapper 的作用这个包装器能把视觉图像和机械臂自身的状态 (关节角度、末端坐标) 揉在一起。这意味着机器人不仅知道自己看到了什么, 还知道自己现在手在哪。encoder_def EncodingWrapper(..., use_propriouse_proprio, enable_stackingTrue, ...)冷启动与热切换如果复现SERL一般会把 bc.py 练出来的模型会作为 RLPDAgent.create 时的初始参数 (或权重)。这相当于把原本需要几百万次尝试才能学会的动作, 压缩成了几千步的模仿。bc_loss 的数据生效为paddingVALID。这是因为不能对在线数据做 BC Loss在线数据 (智能体自己乱跑出来的) 非常乱。物理后果: 如果强迫智能体去模仿这些乱七八糟的动作, 就像是让一个正在学走路的孩子去模仿自己摔跤的动作。这会导致策略陷入低水平的循环。0x04 High UTD 稳定性机制High UTD 的意义它强迫神经网络在极短的时间内吃透每一张图片。4.1 High UTD把每条真机样本反复研读4.1.1 理解 UTDUpdate-to-Data RatioUTDUpdate-to-Data Ratio表示每采集一条环境数据算法进行多少次梯度更新。传统 RL 常用 UTD1采一步训一步。SERL / RLPD 使用更高 UTD通常为 20 甚至更高采集一条昂贵的真机数据后learner 会多次从 buffer 中采样并更新网络。4.1.2 为什么 SERL 需要高 UTD我们可以把 High UTD 理解成真机数据太贵所以每一帧都要反复研读不能看一遍就扔。没有 UTD 的后果普通的 SAC 每采样一个数据才更新一次。对于机器人这种高维度ResNet 图像且数据量极小只有 2.5 小时数据的任务收敛太慢你可能需要练上 10 天半个月。不稳定性由于视觉特征ResNet需要海量更新才能稳定如果更新频率太低视觉头会一直处于模糊状态无法提取有效的位姿信息。4.1.3 极致的采样效率High UTD 将数据的价值榨取到了极致20x 的复习强度 利用 High UTDUpdate-to-Data策略机器人每在现实中走一步后台 Learner 就会对现有数据进行 20 到 40 次的高频更新。REDQ 保驾护航 为了防止这种高强度学习产生幻觉系统利用 10–20 个 Critic 组成 陪审团Ensemble通过取最小值的方式压制过估计偏差。成果数据压榨逻辑真机采集的数据中隐藏着极其细微的物理交互特征如手爪与工件的摩擦。通过 High UTD模型被迫对有限的样本进行深度研读从而在极短的物理时间内捕捉到成功的信号。这让原本需要几周的训练过程被压缩到了喝杯咖啡的时间。工程代价High UTD 对 Learner 的算力提出了严苛要求。这要求 JAX 必须在毫秒级时间内完成多轮反向传播以确保学习速度始终领先于采样速度。4.1.4 实现partial(jax.jit, static_argnames(utd_ratio, pmap_axis)) def update_high_utd( self, batch: Batch, *, utd_ratio: int, pmap_axis: Optional[str] None, ) - Tuple[SACAgent, dict]: Fast JITted high-UTD version of .update. Splits the batch into minibatches, performs utd_ratio critic (and target) updates, and then one actor/temperature update. Batch dimension must be divisible by utd_ratio. batch_size batch[rewards].shape[0] assert ( batch_size % utd_ratio 0 ), fBatch size {batch_size} must be divisible by UTD ratio {utd_ratio} minibatch_size batch_size // utd_ratio chex.assert_tree_shape_prefix(batch, (batch_size,)) def scan_body(carry: Tuple[SACAgent], data: Tuple[Batch]): (agent,) carry (minibatch,) data agent, info agent.update( minibatch, pmap_axispmap_axis, networks_to_updatefrozenset({critic}) ) return (agent,), info def make_minibatch(data: jnp.ndarray):