import torch


def n_step_returns(q_values, rewards, kls, discount=0.99):
    """
    Calculates all n-step returns.

    Args:
        q_values (torch.Tensor): the Q-value estimates at each time step [time_steps+1, batch_size, 1]
        rewards (torch.Tensor): the rewards at each time step [time_steps, batch_size, 1]
        kls (torch.Tensor): the scaled kl divergences at each time step [time_steps, batch_size, 1]
        discount (float): the temporal discount factor
    """
    discounts = torch.cat([(discount*torch.ones_like(q_values[:1]))**i for i in range(rewards.shape[0])], 0)
    rewards[1:] = rewards[1:] - kls[:-1]
    discounted_returns = torch.cumsum(discounts * rewards, dim=0)
    terminal_values = discount * discounts * (q_values[1:] - kls)
    # return torch.cat([q_values[:1], discounted_returns], dim=0)
    return torch.cat([q_values[:1], discounted_returns + terminal_values], dim=0)