4. 自动微分与vmap/pmap：grad、vmap、pmap、lax.scan等函数式变换的深度解析与RL应用场景

说实话，JAX 最让我着迷的地方，就是它这套函数式变换的思维。你想想看，在 PyTorch 或者 TF 里，我们习惯了“构建图 -> 执行图”这种模式。但 JAX 不一样，它把梯度计算、向量化、并行化、循环展开这些操作，都变成了纯粹的函数变换。说白了，就是“输入一个函数，输出一个新函数”。

这一章，我们就来啃下 JAX 强化学习落地的核心四件套：grad、vmap、pmap 和 lax.scan。我会结合我在实际 RL 项目里踩过的坑，给你讲透它们到底怎么用，以及为什么说它们是 RL 工程的加速器。

4.1 grad：不只是求导，更是策略梯度的引擎

先聊 grad。在 RL 里，我们几乎每天都在跟梯度打交道——策略梯度、Q-learning 的 TD 误差、Actor-Critic 的更新。JAX 的 grad 用起来非常直接：

import jax
import jax.numpy as jnp

def policy_loss(params, obs, actions, advantages):
    logits = policy_network(params, obs)
    log_probs = jax.nn.log_softmax(logits)
    selected_log_probs = log_probs[jnp.arange(len(actions)), actions]
    return -jnp.mean(selected_log_probs * advantages)

# 一行代码搞定梯度计算
loss_and_grad = jax.value_and_grad(policy_loss)
loss, grads = loss_and_grad(params, obs, actions, advantages)

嗯，这里要注意一个细节：grad 默认只对第一个参数求导。如果你的网络参数是嵌套的 dict 或者 tuple，它也能完美处理。我刚开始用的时候，总担心它会不会把学习率或者环境状态也求导了——放心，它只追踪通过 jnp 操作的计算路径。

我的习惯： 在 RL 训练循环里，我几乎不用 grad，而是用 value_and_grad。因为很多时候你既需要 loss 值做日志记录，又需要梯度做更新。一次调用，两个结果，省一次前向传播的开销。

4.2 vmap：隐式批量化，告别手写 for 循环

接下来是 vmap。这个函数我愿称之为“RL 工程师的救星”。

你想想看，在强化学习中，我们经常需要同时处理多个环境实例（并行采样）。传统做法是写一个 for 循环，把每个环境的 obs 塞进网络。但 vmap 让你可以写一个处理单个样本的函数，然后自动把它向量化成处理批量样本的函数。

def single_forward(params, obs):
    # 假设 obs 是 (obs_dim,) 的一维向量
    return policy_network(params, obs)

# 自动变成批量版本
batch_forward = jax.vmap(single_forward, in_axes=(None, 0), out_axes=0)

# 现在可以一次传入 (batch_size, obs_dim) 的 obs
logits = batch_forward(params, obs_batch)

in_axes=(None, 0) 的意思是：第一个参数（params）在所有样本间共享（不拆分），第二个参数（obs）沿着第 0 维拆分。这个参数配置我经常搞混，我的记忆方法是：None 表示“不拆分”，数字表示“沿着哪个轴拆分”。

我曾经踩过的坑： 有一次我在做 PPO 时，用 vmap 批量处理多个环境的 obs，结果发现每个环境的 RNN 隐藏状态没有正确拆分。因为我把 hidden_state 也传了 None，导致所有环境共享了同一个 hidden_state。正确的做法应该是 in_axes=(None, 0, 0)，把 hidden_state 也沿着 batch 维度拆分。

4.3 pmap：多卡并行，从单机到分布式

pmap 是 vmap 的分布式版本。它把计算拆分到多个设备（GPU/TPU）上执行。在 RL 中，当你需要大规模并行采样时，pmap 就是你的利器。

# 假设你有 8 个 GPU
def update_step(params, batch):
    loss, grads = value_and_grad(policy_loss)(params, batch)
    return opt_update(params, grads)

# 每个 GPU 处理一部分 batch
parallel_update = jax.pmap(update_step, axis_name='devices')

# 注意：params 需要在所有设备上复制
replicated_params = jax.device_replicate(params)
replicated_batch = jax.device_put(batch, jax.local_devices())

new_params = parallel_update(replicated_params, replicated_batch)

这里有个关键点：pmap 返回的结果也是分布在多个设备上的。你需要用 jax.device_get 或者跨设备规约操作（比如 jax.lax.psum）来聚合结果。我刚开始用 pmap 时，经常忘记做梯度聚合，导致每个 GPU 各自为政，训练完全乱套。

RL 中的典型用法： 在分布式 PPO 中，我们用 pmap 让每个 GPU 独立采样一批轨迹，计算各自的梯度，然后通过 jax.lax.pmean 对所有设备的梯度求平均，再统一更新参数。这样既保证了数据并行，又保持了参数的一致性。

4.4 lax.scan：优雅地展开时间序列

最后是 lax.scan。这个函数在 RL 中处理时间序列时特别有用。比如 RNN 的 rollout、TD(lambda) 的计算、或者 GAE（广义优势估计）的递推。

scan 本质上是一个函数式的 for 循环。它接受一个初始状态和一个序列，然后依次执行函数，同时传递状态。

def compute_gae(carry, transition):
    # carry: (last_gae, last_value)
    # transition: (reward, value, done)
    last_gae, last_value = carry
    reward, value, done = transition
    
    # GAE 递推公式
    delta = reward + 0.99 * last_value * (1 - done) - value
    gae = delta + 0.99 * 0.95 * last_gae * (1 - done)
    
    return (gae, value), gae

# 从后往前扫描
_, advantages = jax.lax.scan(
    compute_gae,
    (0.0, 0.0),  # 初始 carry
    reversed_transitions,  # 时间序列，从 T 到 0
    reverse=True  # 从后往前
)

你看，用 scan 写 GAE 非常简洁，而且它是完全向量化的，比手写 for 循环快得多。我之前的项目里，用 scan 替换了原来的 Python for 循环后，GAE 的计算速度提升了将近 10 倍。

性能提示： lax.scan 要求所有通过 scan 传递的数组在时间维度上大小固定。如果你的轨迹长度不固定（比如 early termination），需要先 padding 到最大长度，然后用 mask 处理。我一般用 jnp.where 配合 done 标志来屏蔽无效的时间步。

4.5 知识体系总览

为了让你更直观地理解这四者的关系，我画了一张图：

JAX 函数式变换四件套 函数变换核心 grad vmap pmap lax.scan 策略梯度 Q-learning 更新 损失函数求导 批量环境采样 批量网络推理 数据并行处理 多 GPU 训练 分布式采样 模型并行 RNN rollout GAE 计算 TD(lambda) 核心思想：组合变换 grad(vmap(pmap(scan(step_fn)))) 将时间、批量、设备三个维度完全解耦

4.6 实战组合：一个完整的 PPO 更新步骤

最后，我把这四个变换组合起来，给你看一个真实的 PPO 更新流程。这基本上是我在项目中用的模板：

# 1. 定义单步更新
def single_update(carry, transition):
    params, opt_state = carry
    obs, action, reward, done, value = transition
    
    # 计算 loss 和梯度
    loss, grads = value_and_grad(policy_loss)(params, obs, action, reward)
    
    # 更新参数
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    
    return (params, opt_state), loss

# 2. 用 scan 展开时间序列
def trajectory_update(params, opt_state, trajectory):
    # trajectory 是 (T, B, ...) 形状
    (params, opt_state), losses = jax.lax.scan(
        single_update,
        (params, opt_state),
        trajectory
    )
    return params, opt_state, losses

# 3. 用 vmap 处理多个环境
batch_update = jax.vmap(trajectory_update, in_axes=(None, None, 1))

# 4. 用 pmap 分布到多个 GPU
distributed_update = jax.pmap(batch_update, axis_name='devices')

# 实际调用
new_params, new_opt_state, all_losses = distributed_update(
    jax.device_replicate(params),
    jax.device_replicate(opt_state),
    jax.device_put(trajectory_batch, jax.local_devices())
)

你看，通过组合这四个变换，我们把一个复杂的分布式 PPO 更新写成了清晰、可组合的函数式代码。每个变换只负责一件事，组合起来却能处理极其复杂的计算图。

我的感悟： 刚开始接触 JAX 这套函数式变换时，确实有点不适应。但用久了你会发现，这种“小函数组合成大函数”的思维，比面向对象的继承和封装要优雅得多。尤其是在 RL 这种需要频繁调整算法细节的场景下，改一个函数比改一个类要安全得多。

好了，这一章的内容就到这里。记住这四个变换的核心思想：grad 负责微分，vmap 负责批量，pmap 负责分布式，scan 负责时序。把它们用好，你的 RL 代码会既快又清晰。

专注资料整理
← 上一章 📖 返回目录 下一章 →