code: def make_networks(
      action_spec: specs.BoundedArray,
      policy_layer_sizes: Sequence[int] = (50, 1024, 1024),
      critic_layer_sizes: Sequence[int] = (50, 1024, 1024),
      vmin: float = -150.0,
      vmax: float = 150.0,
      num_atoms: int = 51,
  ) -> Dict[str, snt.Module]:
      """Creates networks used by the agent."""
  
      num_dimensions = np.prod(action_spec.shape, dtype=int)
  
      policy_network = snt.Sequential(
          [
              networks.LayerNormMLP(
                  policy_layer_sizes,
                  w_init=snt.initializers.Orthogonal(),
                  activation=tf.nn.relu,
                  activate_final=True,
              ),
              networks.MultivariateNormalDiagHead(
                  num_dimensions,
                  tanh_mean=False,
                  init_scale=1.0,
                  fixed_scale=False,
                  use_tfd_independent=True,
                  w_init=snt.initializers.Orthogonal(),
              ),
          ]
      )
  
      # The multiplexer concatenates the (maybe transformed) observations/actions.
      critic_network = networks.CriticMultiplexer(
          observation_network=snt.Sequential(
              [
                  snt.Linear(
                      critic_layer_sizes[0],
                      w_init=snt.initializers.Orthogonal(),
                  ),
                  snt.LayerNorm(
                      axis=slice(1, None),
                      create_scale=True,
                      create_offset=True,
                  ),
                  tf.nn.tanh,
              ]
          ),
          critic_network=snt.nets.MLP(
              critic_layer_sizes[1:],
              w_init=snt.initializers.Orthogonal(),
              activation=tf.nn.relu,
              activate_final=True,
          ),
          action_network=networks.ClipToSpec(action_spec),
      )
      critic_network = snt.Sequential(
          [
              critic_network,
              networks.DiscreteValuedHead(
                  vmin,
                  vmax,
                  num_atoms,
                  w_init=snt.initializers.Orthogonal(),
              ),
          ]
      )
      observation_network = networks.DrQTorso()
  
      return {
          "policy": policy_network,
          "critic": critic_network,
          "observation": observation_network,
      }
  
file_path: acme/examples/tf/control_suite/lp_dmpo_pixels_drqv2.py
imports: import functools
  from typing import Dict, Sequence
  from absl import app
  from absl import flags
  from acme import specs
  from acme.agents.tf import dmpo
  from acme.datasets import image_augmentation
  import helpers
  from acme.tf import networks
  import launchpad
  import numpy
  import sonnet
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: make_networks(action_spec, policy_layer_sizes, critic_layer_sizes, vmin, vmax, num_atoms) and returns: dict