import torch


def J_E(Q_t, action_tensor, is_demo, margin=0.8):
    """
    Large Margin Classification Loss
    --------------------------------
    Parameters
      Q_t:                 (b,|A|) tensor of q-values from model
      action_tensor:       (b,|A|) 1-hot vector of actions taken
      is_demo:             (b,1)   binary value indicating if expert action
      margin:              int of supervised loss margin

    Returns
      loss:                (1,) tensor of loss magnitude
    """
    action_tensor = action_tensor.clone().detach().requires_grad_(False)
    # creates a (b, 1) tensor
    loss, _ = torch.max(Q_t + margin*action_tensor,dim=1)
    loss    = loss - torch.sum(Q_t * action_tensor,dim=1)

    # masks out losses for non-expert examples
    return is_demo * loss