code: def _compute_advantage(
      networks: CRRNetworks,
      policy_params: networks_lib.Params,
      critic_params: networks_lib.Params,
      transition: types.Transition,
      key: networks_lib.PRNGKey,
      num_action_samples: int = 4,
  ) -> jnp.ndarray:
      """Returns the advantage for the transition."""
      # Sample count actions.
      replicated_observation = jnp.broadcast_to(
          transition.observation,
          (num_action_samples,) + transition.observation.shape,
      )
      dist_params = networks.policy_network.apply(
          policy_params, replicated_observation
      )
      actions = networks.sample(dist_params, key)
      # Compute the state-action values for the sampled actions.
      q_actions = networks.critic_network.apply(
          critic_params, replicated_observation, actions
      )
      # Take the mean as the state-value estimate. It is also possible to take the
      # maximum, aka CRR(max); see table 1 in CRR paper.
      q_estimate = jnp.mean(q_actions, axis=0)
      # Compute the advantage.
      q = networks.critic_network.apply(
          critic_params, transition.observation, transition.action
      )
      return q - q_estimate
  
file_path: acme/acme/agents/jax/crr/losses.py
imports: from typing import Callable
  from acme import types
  from acme.agents.jax.crr.networks import CRRNetworks
  from acme.jax import networks
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: _compute_advantage(networks, policy_params, critic_params, transition, key, num_action_samples) and returns: unknown