code: def make_control_networks(
      environment_spec: specs.EnvironmentSpec,
      *,
      with_recurrence: bool = False,
      policy_layer_sizes: Sequence[int] = (256, 256, 256),
      critic_layer_sizes: Sequence[int] = (512, 512, 256),
      policy_init_scale: float = 0.7,
      critic_type: types.CriticType = types.CriticType.MIXTURE_OF_GAUSSIANS,
      mog_init_scale: float = 1e-3,  # Used by MoG critic.
      mog_num_components: int = 5,  # Used by MoG critic.
      categorical_num_bins: int = 51,  # Used by CATEGORICAL* critics.
      vmin: float = -150.0,  # Used by CATEGORICAL* critics.
      vmax: float = 150.0,  # Used by CATEGORICAL* critics.
  ) -> MPONetworks:
      """Creates MPONetworks to be used DM Control suite tasks."""
  
      # Unpack the environment spec to get appropriate shapes, dtypes, etc.
      num_dimensions = np.prod(
          environment_spec.actions.shape, dtype=int
      )
  
      # Factory to create the core hk.Module. Must be a factory as the module must
      # be initialized within a hk.transform scope.
      if with_recurrence:
          make_core_module = lambda: GRUWithSkip(16)
      else:
          make_core_module = hk.IdentityCore
  
      def policy_fn(observation: types.NestedArray) -> tfd.Distribution:
          embedding = networks_lib.LayerNormMLP(
              policy_layer_sizes, activate_final=True
          )(observation)
          return networks_lib.MultivariateNormalDiagHead(
              num_dimensions, init_scale=policy_init_scale
          )(embedding)
  
      def critic_fn(
          observation: types.NestedArray, action: types.NestedArray
      ) -> DistributionOrArray:
          # Action is clipped to avoid critic extrapolations outside the spec range.
          clipped_action = networks_lib.ClipToSpec(
              environment_spec.actions
          )(action)
          inputs = jnp.concatenate(
              [observation, clipped_action], axis=-1
          )
          embedding = networks_lib.LayerNormMLP(
              critic_layer_sizes, activate_final=True
          )(inputs)
  
          if critic_type == types.CriticType.MIXTURE_OF_GAUSSIANS:
              return networks_lib.GaussianMixture(
                  num_dimensions=1,
                  num_components=mog_num_components,
                  multivariate=False,
                  init_scale=mog_init_scale,
                  append_singleton_event_dim=False,
                  reinterpreted_batch_ndims=0,
              )(embedding)
          elif critic_type in (
              types.CriticType.CATEGORICAL,
              types.CriticType.CATEGORICAL_2HOT,
          ):
              return networks_lib.CategoricalCriticHead(
                  num_bins=categorical_num_bins, vmin=vmin, vmax=vmax
              )(embedding)
          else:
              return hk.Linear(
                  output_size=1, w_init=hk_init.TruncatedNormal(0.01)
              )(embedding)
  
      # Create unrollable torso.
      torso = make_unrollable_network(make_core_module=make_core_module)
  
      # Create MPONetworks to add functionality required by the agent.
      return MPONetworks(
          policy_head=hk.without_apply_rng(hk.transform(policy_fn)),
          critic_head=hk.without_apply_rng(hk.transform(critic_fn)),
          torso=torso,
      )
  
imports: import dataclasses
  from typing import Callable, NamedTuple, Optional, Sequence, Tuple, Union
  from acme import specs
  from acme.agents.jax.mpo import types
  from acme.jax import networks
  from acme.jax import utils
  import chex
  import haiku
  import haiku.initializers
  import jax
  import jax.numpy
  import numpy
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 3 for calling a global function with example usage: make_control_networks(environment_spec) and returns: MPONetworks, networks_lib, networks_lib, networks_lib, hk