def METHOD_NAME(self):
    values = torch.tensor([[2.] * 5], dtype=torch.float32)
    step_types = torch.tensor([[StepType.MID] * 5], dtype=torch.int64)
    rewards = torch.tensor([[3.] * 5], dtype=torch.float32)
    discounts = torch.tensor([[0.9] * 5], dtype=torch.float32)
    td_lambda = 0.6 / 0.9
    d = 2 * 0.9 + 1
    expected = torch.tensor([[((d * 0.6 + d) * 0.6 + d) * 0.6 + d,
                              (d * 0.6 + d) * 0.6 + d, d * 0.6 + d, d]],
                            dtype=torch.float32)
    self._check(
        rewards=rewards,
        values=values,
        step_types=step_types,
        discounts=discounts,
        td_lambda=td_lambda,
        expected=expected)
    # two episodes, and exceed by time limit (discount=1)
    step_types = torch.tensor([[
        StepType.MID, StepType.MID, StepType.LAST, StepType.MID,
        StepType.MID
    ]],
                              dtype=torch.int32)
    expected = torch.tensor([[d * 0.6 + d, d, 0, d]], dtype=torch.float32)
    self._check(
        rewards=rewards,
        values=values,
        step_types=step_types,
        discounts=discounts,
        td_lambda=td_lambda,
        expected=expected)
    # tow episodes, and end normal (discount=0)
    step_types = torch.tensor([[
        StepType.MID, StepType.MID, StepType.LAST, StepType.MID,
        StepType.MID
    ]],
                              dtype=torch.int32)
    discounts = torch.tensor([[0.9, 0.9, 0.0, 0.9, 0.9]],
                             dtype=torch.float32)
    expected = torch.tensor([[1 * 0.6 + d, 1, 0, d]], dtype=torch.float32)
    self._check(
        rewards=rewards,
        values=values,
        step_types=step_types,
        discounts=discounts,
        td_lambda=td_lambda,
        expected=expected)