import torch


def get_entropy_loss(memory, args, i_agent):
    """Compute entropy loss for exploration

    Args:
        memory (ReplayMemory): Class that includes trajectories
        args (argparse): Python argparse that contains arguments
        i_agent (int): Index of agent to compute entropy loss

    Returns:
        entropy_loss (torch.Tensor): Entropy loss for encouraging exploration
    """
    _, _, entropies, _, _ = memory.sample()
    entropy = torch.stack(entropies[i_agent], dim=1)
    assert entropy.shape == (args.traj_batch_size, args.ep_horizon), \
        "Shape must be: (batch, ep_horizon)"

    entropy_loss = -args.entropy_weight * torch.mean(torch.sum(entropy, dim=1))
    return entropy_loss