code: def gail_reward(
      reward_balance: float = 0.5,
      max_reward_magnitude: Optional[float] = None,
  ) -> ail_networks.ImitationRewardFn:
      """GAIL reward function (https://arxiv.org/pdf/1606.03476.pdf).
  
      Args:
        reward_balance: 1 means log(D) reward, 0 means -log(1-D) and other values
          mean an average of the two.
        max_reward_magnitude: Clipping value for the reward.
  
      Returns:
        The function from logit to imitation reward.
      """
  
      def imitation_reward(logits: networks_lib.Logits) -> float:
          # Quick Maths:
          # logits = ln(D) - ln(1-D)
          # -softplus(-logits) = ln(D)
          # softplus(logits) = -ln(1-D)
          rewards = reward_balance * -jax.nn.softplus(-logits) + (
              1 - reward_balance
          ) * jax.nn.softplus(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  # jnp-type
  
      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 2 for calling a global function with example usage: gail_reward(reward_balance, max_reward_magnitude) and returns: imitation_reward, rewards