def compute_r2g(rewards, gamma=0.9):
    """Given a rewards array, calculates the rewards-to-go (r2g), where
    r2g[i] = rewards[i] + sum(
        gamma**j * rewards[j] for j in range(i+1, len(rewards))
    )

    Args:
        rewards (List[Int]): The rewards collected in a trajectory
        gamma (float, optional): The discount factor. Defaults to 0.9.

    Returns:
        List[Int]: The Rewards-to-go as calculated by above formula.
    """
    r2g, cumm_reward = [], 0
    for reward in reversed(rewards):
        cumm_reward = reward + gamma*cumm_reward
        r2g.append(cumm_reward)

    return reversed(r2g)