code: def _get_single_agent_spec(spec):
      """Returns a single-agent spec from multiagent multigrid spec.
  
      Primarily used for converting multigrid specs to multiagent Acme specs,
      wherein actions and observations specs are expected to be lists (each entry
      corresponding to the spec of that particular agent). Note that this function
      assumes homogeneous observation / action specs across all agents, which is the
      case in multigrid.
  
      Args:
        spec: multigrid environment spec.
      """
  
      def make_single_agent_spec(spec):
          if not spec.shape:  # Rewards & discounts
              shape = ()
          elif len(spec.shape) == 1:  # Actions
              shape = ()
          else:  # Observations
              shape = spec.shape[1:]
  
          if isinstance(spec, specs.BoundedArray):
              # Bounded rewards and discounts often have no dimensions as they are
              # amongst the agents, whereas observations are of shape [num_agents, ...].
              # The following pair of if statements handle both cases accordingly.
              minimum = (
                  spec.minimum
                  if spec.minimum.ndim == 0
                  else spec.minimum[0]
              )
              maximum = (
                  spec.maximum
                  if spec.maximum.ndim == 0
                  else spec.maximum[0]
              )
              return specs.BoundedArray(
                  shape=shape,
                  name=spec.name,
                  minimum=minimum,
                  maximum=maximum,
                  dtype=spec.dtype,
              )
          elif isinstance(spec, specs.DiscreteArray):
              return specs.DiscreteArray(
                  num_values=spec.num_values,
                  dtype=spec.dtype,
                  name=spec.name,
              )
          elif isinstance(spec, specs.Array):
              return specs.Array(
                  shape=shape, dtype=spec.dtype, name=spec.name
              )
          else:
              raise ValueError(f"Unexpected spec type {type(spec)}.")
  
      single_agent_spec = jax.tree_map(make_single_agent_spec, spec)
      return single_agent_spec
  
imports: from typing import Any, Dict, List, Optional
  import warnings
  from acme import specs
  from acme import types
  from acme import wrappers
  from acme.multiagent import types
  from acme.wrappers import multiagent_dict_key_wrapper
  import dm_env
  import gym
  from gym import spaces
  import jax
  import numpy
  import tree
label_desc: Use image node_id 1 for calling a global function with example usage: _get_single_agent_spec(spec) and returns: single_agent_spec, specs, specs, specs