code: def make_discriminator(
      environment_spec: specs.EnvironmentSpec,
      discriminator_transformed: hk.TransformedWithState,
      logpi_fn: Optional[
          Callable[
              [
                  networks_lib.Params,
                  networks_lib.Observation,
                  networks_lib.Action,
              ],
              jnp.ndarray,
          ]
      ] = None,
  ) -> networks_lib.FeedForwardNetwork:
      """Creates the discriminator network.
  
      Args:
        environment_spec: Environment spec
        discriminator_transformed: Haiku transformed of the discriminator.
        logpi_fn: If the policy logpi function is provided, its output will be
          removed from the discriminator logit.
  
      Returns:
        The network.
      """
  
      def apply_fn(
          params: hk.Params,
          policy_params: networks_lib.Params,
          state: hk.State,
          transitions: types.Transition,
          is_training: bool,
          rng: networks_lib.PRNGKey,
      ) -> networks_lib.Logits:
          output, state = discriminator_transformed.apply(
              params,
              state,
              transitions.observation,
              transitions.action,
              transitions.next_observation,
              is_training,
              rng,
          )
          if logpi_fn is not None:
              logpi = logpi_fn(
                  policy_params,
                  transitions.observation,
                  transitions.action,
              )
  
              # Quick Maths:
              # D = exp(output)/(exp(output) + pi(a|s))
              # logit(D) = log(D/(1-D)) = log(exp(output)/pi(a|s))
              # logit(D) = output - logpi
              return (
                  output - logpi,
                  state,
              )  # pytype: disable=bad-return-type  # jax-ndarray
          return (
              output,
              state,
          )  # pytype: disable=bad-return-type  # jax-ndarray
  
      dummy_obs = utils.zeros_like(environment_spec.observations)
      dummy_obs = utils.add_batch_dim(dummy_obs)
      dummy_actions = utils.zeros_like(environment_spec.actions)
      dummy_actions = utils.add_batch_dim(dummy_actions)
  
      return networks_lib.FeedForwardNetwork(
          # pylint: disable=g-long-lambda
          init=lambda rng: discriminator_transformed.init(
              rng, dummy_obs, dummy_actions, dummy_obs, False, rng
          ),
          apply=apply_fn,
      )
  
imports: import dataclasses
  import functools
  from typing import Any, Callable, Generic, Iterable, Optional
  from acme import specs
  from acme import types
  from acme.jax import networks
  from acme.jax import utils
  from acme.jax.imitation_learning_types import DirectRLNetworks
  import haiku
  import jax
  from jax import numpy
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: make_discriminator(environment_spec, discriminator_transformed, logpi_fn) and returns: networks_lib, output,  state, unknown,  state