code: def make_multigrid_ppo_networks(
      environment_spec: specs.EnvironmentSpec,
      hidden_layer_sizes: Sequence[int] = (64, 64),
  ) -> ppo.PPONetworks:
      """Returns PPO networks used by the agent in the multigrid environments."""
  
      # Check that multigrid environment is defined with discrete actions, 0-indexed
      assert np.issubdtype(
          environment_spec.actions.dtype, np.integer
      ), (
          "Expected multigrid environment to have discrete actions with int dtype"
          f" but environment_spec.actions.dtype == {environment_spec.actions.dtype}"
      )
      assert environment_spec.actions.minimum == 0, (
          "Expected multigrid environment to have 0-indexed action indices, but"
          f" environment_spec.actions.minimum == {environment_spec.actions.minimum}"
      )
      num_actions = environment_spec.actions.maximum + 1
  
      def forward_fn(inputs):
          processed_inputs = multigrid_obs_preproc(inputs)
          trunk = hk.nets.MLP(hidden_layer_sizes, activation=jnp.tanh)
          h = trunk(processed_inputs)
          logits = hk.Linear(num_actions)(h)
          values = hk.Linear(1)(h)
          values = jnp.squeeze(values, axis=-1)
          return (CategoricalParams(logits=logits), values)
  
      # Transform into pure functions.
      forward_fn = hk.without_apply_rng(hk.transform(forward_fn))
  
      dummy_obs = acme_jax_utils.zeros_like(
          environment_spec.observations
      )
      dummy_obs = acme_jax_utils.add_batch_dim(
          dummy_obs
      )  # Dummy 'sequence' dim.
      network = networks_lib.FeedForwardNetwork(
          lambda rng: forward_fn.init(rng, dummy_obs), forward_fn.apply
      )
      return make_categorical_ppo_networks(network)
  
file_path: acme/examples/multiagent/multigrid/helpers.py
imports: import functools
  from typing import Any, Dict, NamedTuple, Sequence
  from acme import specs
  from acme.agents.jax import ppo
  from acme.agents.jax.multiagent.decentralized import factories
  from acme.jax import networks
  from acme.jax import utils
  from acme.multiagent import types
  import haiku
  import jax
  import jax.numpy
  import numpy
  import tensorflow_probability
label_desc: Use image node_id 3 for calling a global function with example usage: make_multigrid_ppo_networks(environment_spec, hidden_layer_sizes) and returns: make_categorical_ppo_networks, CategoricalParams,  values