get_rewards_ppo

torchtune.rlhf.get_rewards_ppo(scores: Tensor, logprobs: Tensor, ref_logprobs: Tensor, kl_coeff: float, valid_score_idxs: Optional[Tensor] = None) → Tuple[Tensor, Tensor, Tensor]

计算给定分数、logprobs 和参考 logprobs 的 PPO 奖励。

参数:

• scores (torch.Tensor) – 奖励模型分数，形状 (b,)。

• logprobs (torch.Tensor) – 策略 logprobs，形状 (b, response_len)。

• ref_logprobs (torch.Tensor) – 参考基础模型 logprobs，形状 (b, response_len)。

• kl_coeff (float) – KL 奖励贡献系数。

• valid_score_idxs (Optional[torch.Tensor]) – 有效（非填充）token 预测的索引张量。这在计算填充序列的奖励时很有用，因为分数和价值估计是为最后一个有效的预测 token 定义的。形状： (b,)。默认为 None。

返回:

形状为 (b, response_len) 的张量元组，每个张量对应：

• total_reward：总奖励，结合了每个 token 的 kl 奖励和奖励模型分数。

• kl：策略和参考策略 logprobs 之间的 kl 散度。

• kl_reward：kl 散度，按 kl_coeff 缩放。

返回类型:

Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

用于张量形状的符号

• b：批大小

• response_len：模型响应长度