reward2go

class torchrl.objectives.value.functional.reward2go(reward, done, gamma, *, time_dim: int = - 2)

根据多个轨迹和剧集结束情况，计算奖励的折扣累积和。

参数：

• reward (torch.Tensor) – 一个张量，包含多个轨迹在每个时间步收到的奖励。

• done (Tensor) – 表示剧集结束的布尔标志。与截断（剧集未结束但中断）不同。

• gamma (float, optional) – 用于计算奖励折扣累积和的折扣因子。默认为 1.0。

• time_dim (int) – 时间维度展开所在的维度。默认为 -2。

返回值：

一个形状为 [B, T] 的张量，包含折扣累积

每个时间步的奖励总和 (reward-to-go)。

返回值类型：

torch.Tensor

示例

>>> reward = torch.ones(1, 10)
>>> done = torch.zeros(1, 10, dtype=torch.bool)
>>> done[:, [3, 7]] = True
>>> reward2go(reward, done, 0.99, time_dim=-1)
tensor([[3.9404],
        [2.9701],
        [1.9900],
        [1.0000],
        [3.9404],
        [2.9701],
        [1.9900],
        [1.0000],
        [1.9900],
        [1.0000]])