Temperature (自动熵调节) 网络：

• 作用：自动调节策略熵的目标值
• 约束：确保温度参数 ≥ 目标熵
• 更新：通过拉格朗日乘数法优化

temperature_def = GeqLagrangeMultiplier( init_value=temperature_init, # 默认1.0 constraint_shape=(), constraint_type="geq", )

4.3.2 编码器架构

"small" 编码器：

encoders = { image_key: SmallEncoder( features=(32, 64, 128, 256), kernel_sizes=(3, 3, 3, 3), strides=(2, 2, 2, 2), padding="VALID", pool_method="avg", bottleneck_dim=256, spatial_block_size=8, ) }

"resnet" 编码器：

encoders = { image_key: resnetv1_configs["resnetv1-10"]( pooling_method="spatial_learned_embeddings", num_spatial_blocks=8, bottleneck_dim=256, ) }

"resnet-pretrained" 编码器：

pretrained_encoder = resnetv1_configs["resnetv1-10-frozen"]( pre_pooling=True, ) encoders = { image_key: PreTrainedResNetEncoder( pooling_method="spatial_learned_embeddings", num_spatial_blocks=8, bottleneck_dim=256, pretrained_encoder=pretrained_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(axis=0) # 最小Q值（保守估计） # TD误差 predicted_qs = self.forward_critic(batch["observations"], batch["actions"], rng, grad_params=params) critic_loss = jnp.mean((predicted_qs - target_qs) ** 2)

Actor损失：

def policy_loss_fn(self, batch, params, rng): # 最大化Q值-熵 predicted_q = predicted_qs.mean(axis=0) 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。
• 采样：你从这个分布里「随机」抽了一个数 a=11。
• 断裂点：当你计算 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(μ,σ)，然后计算 a=tanh(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)−ΣDi=1log(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 输出 a=1.5，被clip成了1.0。
• 在反向传播时，clip 函数在 1.5 这里的导数是 0。
• 后果：梯度传到这里就断了！神经网络接收不到任何信号告诉它"其实你应该减小输出"。
• Tanh 的好处：它是平滑的，即便输出很大，梯度依然存在（虽然很小），能指引网络回来。

0x06 SAC 的工作流程

6.1 工作流程

极简版工作流程如下：

1. 收集数据：在环境里跑趟，把 (s,a,r,s′,done) 存进"经验回放池"（Replay Buffer）。
2. 训练 Critic：从池子里抓一批数据，告诉 Q 网络："根据你看到的奖励和下一步的预测，修正你对当前状态动作价值的评估"。
3. 训练 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(capacity=1000000) def step(self, state): """与环境交互: 根据当前状态, 喷出一个动作""" action = self.actor.sample(state) return action def train_step(self): """核心训练逻辑: SAC 的三步走""" # 从池子里抓一把数据 batch = self.replay_buffer.sample(batch_size=256) # --- 第一步: 更新 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_norm=True，它就会自动在内部插入归一化层。

• Ensemble Q：它支持 critic_ensemble_size=10，这是 RLPD 抑制 Q 值发散的手段。即，在计算 Target 时, 它不是取最小值, 而是计算这 10 个 Q 的均值减去标准差：TargetQ=mean(Q1...10)−std(Q1...10)×ρ

  这叫"悲观备份"。在不确定的地方, Q 值会因为标准差大而被拉低。这强迫智能体只信任那些所有 Q 网络都达成共识的高分区域。

• 自动调节 Alpha（Lagrange）：它使用了拉格朗日乘子法（GeqLagrangeMultiplier）来自动调节熵，比我们手写的手动更新公式更数学化、更稳定。

• JAX 异步更新：利用 JAX，SAC 的 10 个 Critic 可以在不同显卡上并行更新，极大地提升了训练吞吐量。

缺少的内容如下：

1. 缺少 50/50 采样逻辑：在 sac.py 的 update 函数中，它只接收一个 batch。真正的 RLPD 逻辑（从两个池子各抽 128 个数据）通常是在外部的训练循环中完成的，或者是通过更高层的封装实现的。
2. 缺少 BC Loss：sac.py 的 policy_loss_fn中，只有 predicted_q - temperature * log_probs。它没有我们之前在 rlpd.py 里看到的那个关键的 bc_alpha * log_prob(batch_actions)。这意味着这个 sac.py 并不具备"模仿演示数据"的能力。

6.2.2 逻辑流程图

特色功能 (Special Features)如下：

1. Ensemble Support: 通过 jax.vmap 实现的 Q 集成，训练速度极快，天生支持 REDQ 算法。
2. High UTD Dispatch: 专门的 update_high_utd 逻辑，大幅提升采样效率。
3. Modular Encoders: 支持 Shared Encoder (ResNet)，节省显存并加速表征学习。
4. 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 中，训练目标通常拆成三个部分：

在这份代码里，这三个 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),