from typing import Tuple

import jax

from jax_rl.agents.common import InfoDict, Model
from jax_rl.datasets import Batch


def update(actor: Model, batch: Batch) -> Tuple[Model, InfoDict]:
    def loss_fn(actor_params):
        dist = actor.apply({'params': actor_params}, batch.observations)
        log_probs = dist.log_prob(batch.actions)
        actor_loss = -log_probs.mean()
        return actor_loss, {'actor_loss': actor_loss}

    info, new_actor = actor.apply_gradient(loss_fn)

    return new_actor, info
