code: def impala_loss(
      unroll_fn: types.PolicyValueFn,
      *,
      discount: float,
      max_abs_reward: float = np.inf,
      baseline_cost: float = 1.0,
      entropy_cost: float = 0.0,
  ) -> Callable[[hk.Params, reverb.ReplaySample], jax.Array]:
      """Builds the standard entropy-regularised IMPALA loss function.
  
      Args:
        unroll_fn: A `hk.Transformed` object containing a callable which maps
          (params, observations_sequence, initial_state) -> ((logits, value), state)
        discount: The standard geometric discount rate to apply.
        max_abs_reward: Optional symmetric reward clipping to apply.
        baseline_cost: Weighting of the critic loss relative to the policy loss.
        entropy_cost: Weighting of the entropy regulariser relative to policy loss.
  
      Returns:
        A loss function with signature (params, data) -> (loss_scalar, metrics).
      """
  
      def loss_fn(
          params: hk.Params,
          sample: reverb.ReplaySample,
      ) -> Tuple[jax.Array, Mapping[str, jax.Array]]:
          """Batched, entropy-regularised actor-critic loss with V-trace."""
  
          # Extract the data.
          data = sample.data
          observations, actions, rewards, discounts, extra = (
              data.observation,
              data.action,
              data.reward,
              data.discount,
              data.extras,
          )
          initial_state = tree.map_structure(
              lambda s: s[0], extra["core_state"]
          )
          behaviour_logits = extra["logits"]
  
          # Apply reward clipping.
          rewards = jnp.clip(rewards, -max_abs_reward, max_abs_reward)
  
          # Unroll current policy over observations.
          (logits, values), _ = unroll_fn(
              params, observations, initial_state
          )
  
          # Compute importance sampling weights: current policy / behavior policy.
          rhos = rlax.categorical_importance_sampling_ratios(
              logits[:-1], behaviour_logits[:-1], actions[:-1]
          )
  
          # Critic loss.
          vtrace_returns = rlax.vtrace_td_error_and_advantage(
              v_tm1=values[:-1],
              v_t=values[1:],
              r_t=rewards[:-1],
              discount_t=discounts[:-1] * discount,
              rho_tm1=rhos,
          )
          critic_loss = jnp.square(vtrace_returns.errors)
  
          # Policy gradient loss.
          policy_gradient_loss = rlax.policy_gradient_loss(
              logits_t=logits[:-1],
              a_t=actions[:-1],
              adv_t=vtrace_returns.pg_advantage,
              w_t=jnp.ones_like(rewards[:-1]),
          )
  
          # Entropy regulariser.
          entropy_loss = rlax.entropy_loss(
              logits[:-1], jnp.ones_like(rewards[:-1])
          )
  
          # Combine weighted sum of actor & critic losses, averaged over the sequence.
          mean_loss = jnp.mean(
              policy_gradient_loss
              + baseline_cost * critic_loss
              + entropy_cost * entropy_loss
          )  # []
  
          metrics = {
              "policy_loss": jnp.mean(policy_gradient_loss),
              "critic_loss": jnp.mean(baseline_cost * critic_loss),
              "entropy_loss": jnp.mean(entropy_cost * entropy_loss),
              "entropy": jnp.mean(entropy_loss),
          }
  
          return mean_loss, metrics
  
      return utils.mapreduce(loss_fn, in_axes=(None, 0))
  
file_path: acme/acme/jax/losses/impala.py
imports: from typing import Callable, Mapping, Tuple
  from acme.agents.jax.impala import types
  from acme.jax import utils
  import haiku
  import jax
  import jax.numpy
  import numpy
  import reverb
  import rlax
  import tree
label_desc: Use image node_id 1 for calling a global function with example usage: impala_loss(unroll_fn) and returns: utils, mean_loss,  metrics