dp_rob.py — RobDataParallelPPOActor 是机器人 VLA 模型的 PPO Actor 实...

文件路径: verl/experimental/vla/dp_rob.py 模块路径: verl.experimental.vla.dp_rob

文件概述

RobDataParallelPPOActor 是机器人 VLA 模型的 PPO Actor 实现。它处理动作 token 的特殊 logit 切片、轨迹 mask 生成、对数概率计算和策略更新。核心难点在于 VLA 模型的动作预测不是标准的文本生成，而是从词表的特定区间采样。

核心类：RobDataParallelPPOActor

动作 Token 的特殊处理

OpenVLA 模型将连续动作离散化为 256 个 bin，这些 bin 对应词表末尾的 256 个 token。因此 logit 计算需要特殊的切片：

def process_tensor(self, logits):
    """从完整 logits 中切出动作 token 的 logits

    vocab_size 末尾的 256+64 个 token 中的 256 个是动作 token。
    切片范围: logits[..., vocab_size-256-64 : vocab_size-64]
    """
    start_index = self.vocab_size - 256 - 64
    end_index = self.vocab_size - 64
    action_logits = logits[..., start_index:end_index]
    return action_logits

轨迹 Mask 生成

def generate_traj_mask(self, data):
    """生成轨迹的有效 token mask

    在机器人轨迹中，需要标识出哪些 token 是动作预测，
    哪些是 padding。这个 mask 用于 PPO 损失计算。
    """
    response_mask = data.batch['response_mask']  # (B, seq_len)
    # 只对有效的动作 token 计算损失
    traj_mask = response_mask.bool()
    return traj_mask

对数概率计算

def compute_log_prob(self, data):
    """计算动作的对数概率

    PPO 需要 \(\log \pi(a|s)\) 来计算重要性采样比率。
    """
    # 前向传播获取 logits
    logits = self.model.forward(data)

    # 切出动作 logits
    action_logits = self.process_tensor(logits)

    # 计算 log_softmax
    log_probs = F.log_softmax(action_logits, dim=-1)

    # 取实际采样动作的 log_prob
    action_ids = data.batch['response_ids']
    selected_log_probs = log_probs.gather(-1, action_ids.unsqueeze(-1))

    return selected_log_probs.squeeze(-1)

策略更新

def update_policy(self, data):
    """PPO 策略梯度更新

    使用 clipped PPO 目标函数:
    \(L = \min\left(r \cdot A,\; \text{clip}(r,\; 1-\epsilon,\; 1+\epsilon) \cdot A\right)\)
    其中 \(r = \frac{\pi_{\text{new}}(a|s)}{\pi_{\text{old}}(a|s)}\)
    """
    # 计算新策略下的 log_prob
    new_log_probs = self.compute_log_prob(data)
    old_log_probs = data.batch['old_log_probs']

    # 重要性采样比率
    ratio = torch.exp(new_log_probs - old_log_probs)

    # Clipped PPO 损失
    advantages = data.batch['advantages']
    surr1 = ratio * advantages
    surr2 = torch.clamp(ratio, 1-self.clip_eps, 1+self.clip_eps) * advantages
    loss = -torch.min(surr1, surr2).mean()

    # 反向传播和优化
    loss.backward()
    self.optimizer.step()

核心类/函数列表

名称	类型	说明
RobDataParallelPPOActor	类	机器人 PPO Actor
process_tensor	方法	动作 logit 切片（vocab 末尾 256 个）
generate_traj_mask	方法	生成轨迹有效 token mask
compute_log_prob	方法	计算动作对数概率
update_policy	方法	PPO 策略梯度更新

与其他模块的关系

• 被 RobActorRolloutRefWorker（fsdp_workers.py）使用
• 使用 OpenVLA 模型（models/openvla_oft/）做前向传播
• 训练数据来自 EnvLoop（env_loop.py）收集的轨迹

小结

RobDataParallelPPOActor 的核心贡献是将标准 PPO 算法适配到 VLA 模型。关键适配点是动作 token 的特殊 logit 切片（vocab_size-256-64:-64），这反映了 OpenVLA 将连续动作离散化到词表末尾 token 的设计。