code: def logp() -> BCLossWithoutAux:
      """Log probability loss."""
  
      def loss(
          networks: bc_networks.BCNetworks,
          params: networks_lib.Params,
          key: jax_types.PRNGKey,
          transitions: types.Transition,
      ) -> jnp.ndarray:
          logits = networks.policy_network.apply(
              params, transitions.observation, is_training=True, key=key
          )
          logp_action = networks.log_prob(logits, transitions.action)
          return -jnp.mean(logp_action)
  
      return loss
  
file_path: acme/acme/agents/jax/bc/losses.py
imports: from typing import Callable, Optional, Tuple, Union
  from acme import types
  from acme.agents.jax.bc import networks
  from acme.jax import networks
  from acme.jax import types
  from acme.utils import loggers
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: logp() and returns: loss, unknown