code: def multigrid_obs_preproc(
      obs: Dict[str, Any],
      conv_filters: int = 8,
      conv_kernel: int = 3,
      scalar_fc: int = 5,
      scalar_name: str = "direction",
      scalar_dim: int = 4,
  ) -> jnp.ndarray:
      """Conducts preprocessing on 'multigrid' environment dict observations.
  
      The preprocessing applied here is similar to those in:
      https://github.com/google-research/google-research/blob/master/social_rl/multiagent_tfagents/multigrid_networks.py
  
      Args:
        obs: multigrid observation dict, which can include observation inputs such
          as 'image', 'position', and a custom additional observation (defined by
          scalar_name).
        conv_filters: Number of convolution filters.
        conv_kernel: Size of the convolution kernel.
        scalar_fc: Number of neurons in the fully connected layer processing the
          scalar input.
        scalar_name: a special observation key, which is set to
          `direction` in most multigrid environments (and can be overridden here if
          otherwise).
        scalar_dim: Highest possible value for the scalar input. Used to convert to
          one-hot representation.
  
      Returns:
        out: output observation.
      """
  
      def _cast_and_scale(x, scale_by=10.0):
          if isinstance(x, jnp.ndarray):
              x = x.astype(jnp.float32)
          return x / scale_by
  
      outputs = []
  
      if "image" in obs.keys():
          image_preproc = hk.Sequential(
              [
                  _cast_and_scale,
                  hk.Conv2D(
                      output_channels=conv_filters,
                      kernel_shape=conv_kernel,
                  ),
                  jax.nn.relu,
                  hk.Flatten(),
              ]
          )
          outputs.append(image_preproc(obs["image"]))
  
      if "position" in obs.keys():
          position_preproc = hk.Sequential(
              [_cast_and_scale, hk.Linear(scalar_fc)]
          )
          outputs.append(position_preproc(obs["position"]))
  
      if scalar_name in obs.keys():
          direction_preproc = hk.Sequential(
              [
                  functools.partial(
                      jax.nn.one_hot, num_classes=scalar_dim
                  ),
                  hk.Flatten(),
                  hk.Linear(scalar_fc),
              ]
          )
          outputs.append(direction_preproc(obs[scalar_name]))
  
      out = jnp.concatenate(outputs, axis=-1)
      return out
  
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 1 for calling a global function with example usage: multigrid_obs_preproc(obs, conv_filters, conv_kernel, scalar_fc, scalar_name, scalar_dim) and returns: out, unknown