def surrogate_loss(policy, all_obs, all_actions, all_adv, old_dist):
    """
    Compute the loss of policy evaluated at current parameter
    given the observation, action, and advantage value

    Parameters
    ----------
    policy (nn.Module):
    all_obs (Variable):
    all_actions (Variable):
    all_adv (Variable):
    old_dist (dict): The dict of means and log_stds Variables of
        collected samples

    Returns
    -------
    surr_loss (Variable): The surrogate loss function wrapped in
        Variable
    """
    new_dist = policy.get_policy_distribution(all_obs)
    old_dist = policy.distribution(old_dist)

    ratio = new_dist.likelihood_ratio(old_dist, all_actions)
    surr_loss = -(ratio * all_adv).mean()

    return surr_loss