作用自动调节策略熵的目标值约束确保温度参数 ≥ 目标熵更新通过拉格朗日乘数法优化temperature_def GeqLagrangeMultiplier( init_valuetemperature_init, # 默认1.0 constraint_shape(), constraint_typegeq, )4.3.2 编码器架构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, # 冻结的预训练权重 ) }4.3.3 损失函数SAC 的熵正则化保证了探索性双Critic的ensemble提供了稳定的价值估计自动温度调节实现了探索-利用的平衡。Critic损失def critic_loss_fn(self, batch, params, rng): # 计算目标Q值 target_next_qs self.forward_target_critic(batch[next_observations], next_actions, rng) target_next_min_q target_next_qs.min(axis0) # 最小Q值保守估计 # TD误差 predicted_qs self.forward_critic(batch[observations], batch[actions], rng, grad_paramsparams) critic_loss jnp.mean((predicted_qs - target_qs) ** 2)Actor损失def policy_loss_fn(self, batch, params, rng): # 最大化Q值-熵 predicted_q predicted_qs.mean(axis0) actor_objective predicted_q - temperature * log_probs actor_loss -jnp.mean(actor_objective)0x05 特色功能5.1 重参数化重参数化Reparameterization Trick直接从分布采样是不可导的。通过 a μ σ · εε 是固定噪声我们把随机性剥离出来让梯度能顺着加法和乘法回传。5.1.1 问题在 SAC 中Actor 输出的是一个概率分布通常是高斯分布。由于我们需要对这个分布进行采样才能得到动作 a但采样这个动作是不可导的这就导致梯度无法直接回传给生成分布的神经网络。为什么采样不能传导梯度这是深度学习中最经典的问题之一。场景神经网络输出 μ10σ2。采样你从这个分布里「随机」抽了一个数 a11。断裂点当你计算 Loss 后你想问「如果我把 μ 从 10 改成 10.1对 a 有什么影响」结论无法回传。因为「采样」这个动作在计算机里是调用了random()。随机数发生器就像一个黑盒梯度传到这里就断了。5.1.2 方案那么在复现 Actor 的更新过程时我们该如何让梯度通过这个采样步骤传回神经网络的参数中重参数化Reparameterization Trick其实就是为了解决 Actor 怎么根据这个带熵的 Q 值更新梯度的问题。目前问题就是我们该怎么把抽样这个动作变成一个加减乘除的公式SERL 不直接采样 a∼N(μ,σ)而是写成aμσ⋅ε,ε∼N(0,1)这里 ε 是一个固定的随机噪声。现在a 就变成了一个关于 μ 和 σ 的确定性函数加法和乘法梯度就可以顺着 a→μ 和 a→σ 传回神经网络了。5.2 输出控制重参数化使用了 aμσε。但在机器人控制中动作通常是有范围的比如 -1 到 1。直接加减可能会超出范围。SAC 论文里用了Tanh 激活函数来把这个 a 限制在 (−1,1)。做法Actor 输出一个原始值 u∼N(μ,σ)然后计算 atanh(u)。用了 Tanh 之后动作就不再是纯粹的高斯分布了。为了计算准确的熵我们需要用到雅可比行列式Jacobian来对概率密度进行修正。在代码里这通常表现为一个修正项loss−logp(u)−log(1−tanh(u)2)。5.2.1 Tanh 挤压: 气球与盒子的数学当你把一个高斯分布的 u 通过 a tanh(u) 映射到 (-1, 1) 时, 概率密度会发生变化。为什么不能直接用高斯公式? 因为 tanh 在靠近 1和-1 的时候非常平。很多个不同的 u 可能会被挤压到极其接近的 a。代码怎么写? 我们需要用到雅可比修正 (Jacobian Correction)。公式如下logπ(a|s)logμ(u|s)−ΣDi1log(1−tanh²(ui))注: μ(u|s) 是原始高斯分布的概率。在代码中, 这通常写成: log_prob dist.log_prob(u) - torch.log(1 - a.pow(2) 1e-6).sum(dim-1)。1e-6 是为了防止数值溢出。5.2.2 Clip 的灾难梯度消失如果我们不使用 Tanh 修正直接强制把超出范围的动作 clip 掉这会给梯度回传带来什么灾难如果用clip(a, -1, 1)Actor 输出 a1.5被clip成了1.0。在反向传播时clip 函数在 1.5 这里的导数是 0。后果梯度传到这里就断了神经网络接收不到任何信号告诉它其实你应该减小输出。Tanh 的好处它是平滑的即便输出很大梯度依然存在虽然很小能指引网络回来。0x06 SAC 的工作流程6.1 工作流程极简版工作流程如下收集数据在环境里跑趟把 (s,a,r,s′,done) 存进经验回放池Replay Buffer。训练 Critic从池子里抓一批数据告诉 Q 网络根据你看到的奖励和下一步的预测修正你对当前状态动作价值的评估。训练 Actor告诉 Actor调整你的参数使得你输出的动作能让 Q 值最大同时熵也要足够大。下面是 SAC 算法的高层结构伪代码。它清晰地展示了 数据流 是如何在 Actor (演员)、Critic (评论家) 和 Buffer (经验池) 之间流动的。class SACAgent: def __init__(self): # 1. 初始化 5 个核心网络 self.actor ActorNetwork() # 策略函数: s - (mu, sigma) self.critic1 CriticNetwork() # Q1函数: (s, a) - q1 self.critic2 CriticNetwork() # Q2函数: (s, a) - q2 self.target_critic1 Target() # Q1的稳定副本 self.target_critic2 Target() # Q2的稳定副本 # 2. 熵自动调节参数 (Temperature) self.log_alpha log(initial_alpha) # 3. 经验回放池 self.replay_buffer ReplayBuffer(capacity1000000) def step(self, state): 与环境交互: 根据当前状态, 喷出一个动作 action self.actor.sample(state) return action def train_step(self): 核心训练逻辑: SAC 的三步走 # 从池子里抓一把数据 batch self.replay_buffer.sample(batch_size256) # --- 第一步: 更新 Critic (练地图) --- self.update_critic(batch) # --- 第二步: 更新 Actor (练向导) --- # 顺着 Critic 指出的梯度方向, 让 Actor 变得更好 self.update_actor(batch) # --- 第三步: 自动调节 Alpha (练灵魂) --- # 如果熵太小, 调大 Alpha 增加探索; 反之调小 self.update_alpha(batch) # --- 最后: 平滑更新 Target 网络 --- self.soft_update_targets() def update_critic(self, batch): 计算带熵的 Bellman 目标 # 核心公式: Target R gamma * (min(Q1_target, Q2_target) - alpha * log_prob) target_q self.calculate_target_q(batch) # 最小化 MSE 误差 loss1 MeanSquaredError(self.critic1(s, a), target_q) loss2 MeanSquaredError(self.critic2(s, a), target_q) # 执行梯度下降...6.2 sac.py SERL我们接下来看看 SERL 开源代码的实现看看其对 SAC 做了什么改变。6.2.1 RLPD 预适配在 SERL 的代码库中sac.py 扮演的是通用底座的角色。原生 SAC 在真机上其实很慢。为了让它起飞SERL 做了若干增强sac.py 其实是一个全能型 SAC。虽然这个文件叫 sac.py但它已经为 RLPD 做好了全部基础准备High UTD 支持update_high_utd 函数把一个大的 Batch 拆成 20 份连续更新 20 次 Critic这是 RLPD 能跑通的前提。LayerNorm 的隐形支持它调用了 MLP 网络。只要在创建时传入 value_layer_normTrue它就会自动在内部插入归一化层。Ensemble Q它支持 critic_ensemble_size10这是 RLPD 抑制 Q 值发散的手段。即在计算 Target 时, 它不是取最小值, 而是计算这 10 个 Q 的均值减去标准差TargetQmean(Q1...10)−std(Q1...10)×ρ这叫悲观备份。在不确定的地方, Q 值会因为标准差大而被拉低。这强迫智能体只信任那些所有 Q 网络都达成共识的高分区域。自动调节 AlphaLagrange它使用了拉格朗日乘子法GeqLagrangeMultiplier来自动调节熵比我们手写的手动更新公式更数学化、更稳定。JAX 异步更新利用 JAXSAC 的 10 个 Critic 可以在不同显卡上并行更新极大地提升了训练吞吐量。缺少的内容如下缺少 50/50 采样逻辑在 sac.py 的 update 函数中它只接收一个 batch。真正的 RLPD 逻辑从两个池子各抽 128 个数据通常是在外部的训练循环中完成的或者是通过更高层的封装实现的。缺少 BC Losssac.py 的 policy_loss_fn中只有 predicted_q - temperature * log_probs。它没有我们之前在 rlpd.py 里看到的那个关键的 bc_alpha * log_prob(batch_actions)。这意味着这个 sac.py 并不具备模仿演示数据的能力。6.2.2 逻辑流程图特色功能 (Special Features)如下Ensemble Support: 通过 jax.vmap 实现的 Q 集成训练速度极快天生支持 REDQ 算法。High UTD Dispatch: 专门的 update_high_utd 逻辑大幅提升采样效率。Modular Encoders: 支持 Shared Encoder (ResNet)节省显存并加速表征学习。Action Chunking: 支持一次输出一串动作适合高频机器人控制场景。6.2.3 四大特色深度解释极致的集成Ensemble与向量化。sac.py 使用了 ensemblize 技巧。黑科技它利用 JAX 的 vmap 将 Q 网络变成了一个并行张量。优势无论你是想要 2 个 Q 还是 10 个 Q在底层计算上几乎一样快。这让算法在保持悲观评估防止高估的同时不会拖累机器人的实时响应。重 Critic、轻 Actor 的高 UTD 架构。SERL 中有一个非常显著的策略在 update_high_utd 里Critic 更新 20 次Actor 才更新 1 次。解释Critic 是 Actor 的导师。如果导师自己都还没把图画清楚Q 值没收敛让 Actor 拼命改参数只会让它学废了。先刷 20 次再更新一次是 SERL 实现 20 分钟学会抓取的硬件级优化。灵活的视觉编码器架构create_pixels。源码中通过 shared_encoder 参数决定了 Actor 和 Critic 是否共用一个视觉大脑。解释在机器人任务中处理像素是最累的活。共用 ResNet 不仅显存省更重要的是能强迫网络去学习那些任务通用的物理特征比如杯子的边缘在哪里、桌子的高度是多少而不是只学习针对自己有用的特征。拉格朗日温度控制GeqLagrangeMultiplier。源码中引入了拉格朗日约束并非简单的梯度下降来更新 α。解释这是一种更稳健的数学方法它能确保熵被强制约束在一个区间内。当熵太低时α 会像踩刹车一样迅速反弹防止智能体陷入死胡同。6.2.4 损失函数在 SAC 中训练目标通常拆成三个部分损失函数更新对象核心目标critic_loss_fncritic / Q 网络学习 Bellman backup让 Q 逼近 TD targetpolicy_loss_fnactor / policy 网络最大化 Q同时最大化熵temperature_loss_fntemperature / α自动调节熵权重使策略熵接近目标熵在这份代码里这三个 loss 会被包装成一个字典def loss_fns(self, batch): return { critic: partial(self.critic_loss_fn, batch), actor: partial(self.policy_loss_fn, batch), temperature: partial(self.temperature_loss_fn, batch),