code: def make_discrete_networks(
      environment_spec: specs.EnvironmentSpec,
      hidden_layer_sizes: Sequence[int] = (512,),
      use_conv: bool = True,
  ) -> PPONetworks:
      """Creates networks used by the agent for discrete action environments.
  
      Args:
        environment_spec: Environment spec used to define number of actions.
        hidden_layer_sizes: Network definition.
        use_conv: Whether to use a conv or MLP feature extractor.
      Returns:
        PPONetworks
      """
  
      num_actions = environment_spec.actions.num_values
  
      def forward_fn(inputs):
          layers = []
          if use_conv:
              layers.extend([networks_lib.AtariTorso()])
          layers.extend(
              [hk.nets.MLP(hidden_layer_sizes, activate_final=True)]
          )
          trunk = hk.Sequential(layers)
          h = utils.batch_concat(inputs)
          h = trunk(h)
          logits = hk.Linear(num_actions)(h)
          values = hk.Linear(1)(h)
          values = jnp.squeeze(values, axis=-1)
          return (CategoricalParams(logits=logits), values)
  
      forward_fn = hk.without_apply_rng(hk.transform(forward_fn))
      dummy_obs = utils.zeros_like(environment_spec.observations)
      dummy_obs = utils.add_batch_dim(
          dummy_obs
      )  # Dummy 'sequence' dim.
      network = networks_lib.FeedForwardNetwork(
          lambda rng: forward_fn.init(rng, dummy_obs), forward_fn.apply
      )
      # Create PPONetworks to add functionality required by the agent.
      return make_categorical_ppo_networks(network)
  
imports: import dataclasses
  from typing import Callable, NamedTuple, Optional, Sequence
  from acme import specs
  from acme.agents.jax import actor_core
  from acme.jax import networks
  from acme.jax import utils
  import haiku
  import jax
  import jax.numpy
  import numpy
  import tensorflow_probability
label_desc: Use image node_id 6 for calling a global function with example usage: make_discrete_networks(environment_spec, hidden_layer_sizes, use_conv) and returns: make_categorical_ppo_networks, CategoricalParams,  values