code: def stacked_shapes(
      shapes: Dict[StepKeyType, Dict[str, Sequence[int]]],
      agent_counts_dict: Dict[StepKeyType, int],
  ) -> Dict[StepKeyType, Dict[str, Sequence[int]]]:
      """Adopt shapes dict for stacked multi-agent scenario. (Mainly for observation stacking.)
  
      Takes a dict of shapes for a structured step. Augments shapes for sub-steps that have multiple agents
      by the agent dimension, so the observations from multiple agents can be concatenated.
      """
      shapes = shapes.copy()
  
      for sub_step_key, agent_count in agent_counts_dict.items():
          # Dynamic number of agents is not supported
          assert agent_count != -1, (
              "for observation concatenation, the number of agents "
              "needs to be known upfront"
          )
  
          # Single-agent spaces are left as-is
          if agent_count == 1:
              continue
  
          # Multi-agent spaces are stacked
          for obs_name, original_shape in shapes[sub_step_key].items():
              shapes[sub_step_key][obs_name] = (
                  agent_count,
                  *original_shape,
              )
  
      return shapes
  
imports: from typing import Dict, Union, Sequence
  from gym import spaces
  from maze.core.env.structured_env import StepKeyType
label_desc: Use image node_id 3 for calling a global function with example usage: stacked_shapes(shapes, agent_counts_dict) and returns: shapes