import torch


def prune_duplicate_goal(z_batch, z_goal):
    """
    Ensure goal is in z_batch only once
    """
    l = []
    for idx, z in enumerate(z_batch):
        s = 0.5 * (z - z_goal).pow(2).sum(dim=1)
        if s > 1e-8:
            app_z = z.unsqueeze(0)
            l.append(app_z)
    l.append(z_goal)
    return torch.cat(l, dim=0).unsqueeze(0)