import torch


def get_return(reward, args):
    """Compute episodic return given trajectory

    Args:
        reward (list): Contains rewards across trajectories for specific agent
        args (argparse): Python argparse that contains arguments

    Returns:
        return_ (torch.Tensor): Episodic return with shape: (batch, ep_horizon)
    """
    reward = torch.stack(reward, dim=1)
    assert reward.shape == (args.traj_batch_size, args.ep_horizon), \
        "Shape must be: (batch, ep_horizon)"

    R, return_ = 0., []
    for timestep in reversed(range(args.ep_horizon)):
        R = reward[:, timestep] + args.discount * R
        return_.insert(0, R)
    return_ = torch.stack(return_, dim=1)

    return return_