import torch


def calc_returns(rewards, values, dones, reward_discount):
    """
    Calculate temporal difference returns
    Args:
        rewards: Rewards from environment
        values: State-values
        dones: Episode ended flags
        reward_discount: Discount factor for state-values

    Returns: Temporal difference returns

    """
    assert len(rewards) == len(values) - 1 == len(dones)

    R = values[-1]
    returns = torch.zeros_like(rewards)
    for t in reversed(range(len(rewards))):
        nonterminal = 1 - dones[t]
        R = rewards[t] + nonterminal * reward_discount * R
        returns[t] = R

    return returns