code: def fairl_reward(
      max_reward_magnitude: Optional[float] = None,
  ) -> ail_networks.ImitationRewardFn:
      """The FAIRL reward function (https://arxiv.org/pdf/1911.02256.pdf).
  
      Args:
        max_reward_magnitude: Clipping value for the reward.
  
      Returns:
        The function from logit to imitation reward.
      """
  
      def imitation_reward(logits: networks_lib.Logits) -> float:
          rewards = jnp.exp(jnp.clip(logits, a_max=20.0)) * -logits
          if max_reward_magnitude is not None:
              # pylint: disable=invalid-unary-operand-type
              rewards = jnp.clip(
                  rewards,
                  a_min=-max_reward_magnitude,
                  a_max=max_reward_magnitude,
              )
          return rewards  # pytype: disable=bad-return-type  # jax-types
  
      return imitation_reward
  
imports: from typing import Optional
  from acme.agents.jax.ail import networks
  from acme.jax import networks
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: fairl_reward(max_reward_magnitude) and returns: imitation_reward, rewards