import torch


def compute_r2g(rewards, done, gamma):
    """
        Computes discounted rewards-to-go.
    """
    rewards2go = []
    running_sum = 0
    for r, is_terminal in zip(reversed(rewards), reversed(done)):
        running_sum = r + gamma * running_sum * (1 - is_terminal)
        rewards2go.insert(0, running_sum)

    rewards2go = torch.as_tensor(rewards2go, dtype=torch.float32)
    return rewards2go