def RobustL2Loss(output, log_std, target) -> float:
    """Robust L2 loss using a Gaussian prior with aleatoric uncertainty estimation."""
    loss = 0.5 * (output - target) ** 2 / (2 * log_std).exp() + log_std
    return loss.mean()