def METHOD_NAME(  # type: ignore
    self,
    step_count: int,
    batch: ObservationType,
    actor_critic_output: ActorCriticOutput[CategoricalDistr],
) -> Dict[str, Tuple[torch.Tensor, Optional[float]]]:
    actions = cast(torch.LongTensor, batch["actions"])
    values = actor_critic_output.values
    action_log_probs = actor_critic_output.distributions.log_prob(actions)
    action_log_probs = action_log_probs.view(
        action_log_probs.shape
        + (1,)
        * (
            len(cast(torch.Tensor, batch["adv_targ"]).shape)
            - len(action_log_probs.shape)
        )
    )
    dist_entropy: torch.FloatTensor = getattr(
        actor_critic_output.distributions, self.entropy_method_name
    )()
    dist_entropy = dist_entropy.view(
        dist_entropy.shape
        + ((1,) * (len(action_log_probs.shape) - len(dist_entropy.shape)))
    )
    value_loss = 0.5 * (cast(torch.FloatTensor, batch["returns"]) - values).pow(2)
    # TODO: Decided not to use normalized advantages here,
    #   is this correct? (it's how it's done in Kostrikov's)
    action_loss = -(
        cast(torch.FloatTensor, batch["adv_targ"]).detach() * action_log_probs
    )
    if self.acktr:
        # TODO: Currently acktr doesn't really work because of this natural gradient stuff
        #   that we should figure out how to integrate properly.
        get_logger().warning("acktr is only partially supported.")
    return {
        "value": (value_loss, self.value_loss_coef),
        "action": (action_loss, None),
        "entropy": (dist_entropy.mul_(-1.0), self.entropy_coef),  # type: ignore
    }