code: def stack_and_flatten_spaces(
      input_tensor_dict: Iterable[Dict[str, torch.Tensor]],
      observation_spaces_dict: Dict[StepKeyType, spaces.Dict],
  ) -> Dict[str, torch.Tensor]:
      """Merges an iterable of dictionary spaces (usually observations or actions from subsequent sub-steps)
      into a single dictionary containing all the items.
  
      If one key is present in multiple elements, all its values will be concatenated in the resulting dictionary.
  
      :param input_tensor_dict: Iterable of dictionary spaces (usually observations or actions from subsequent sub-steps).
      :param observation_spaces_dict: The full observation_spaces dict to use in order to infer the concatenation
          dimension
      :return: One flat dictionary, containing all keys and values form the elements of the iterable.
      """
      result = defaultdict(list)
  
      # Collect observations in flat dict
      for input_tensor in input_tensor_dict:
          for key, obs in input_tensor.items():
              result[key].append(obs)
  
      # Infer the concatenation dimension based on raw (unbatched) observation shapes, and then verify it matches for all
      #   keys
      flattened_space = flat_structured_space(observation_spaces_dict)
      key_0 = list(result.keys())[0]
      concat_dim = len(result[key_0][0].shape) - len(
          flattened_space[key_0].shape
      )
      for key in result.keys():
          assert concat_dim == len(result[key][0].shape) - len(
              flattened_space[key].shape
          ), (
              f"Batch dimensions do not seem to align across observations keys.. specifically for "
              f"{(key, result[key][0].shape)} vs {(key_0, result[key_0][0].shape)}"
          )
  
      # Concatenate all at once for efficiency
      for obs_name, observations in result.items():
          if len(observations) == 1:
              result[obs_name] = observations[0]
          else:
              result[obs_name] = torch.stack(
                  observations, dim=concat_dim
              )
  
      # Return an ordinary dict, not default dict
      return dict(result)
  
file_path: maze/maze/perception/perception_utils.py
imports: import collections
  from collections import defaultdict
  from typing import Union, Dict, Any, Callable, Sequence, Iterable
  import gym
  import numpy
  import torch
  from gym import spaces
  from maze.core.env.structured_env import StepKeyType
  from maze.core.utils.structured_env_utils import flat_structured_space
label_desc: Use image node_id 3 for calling a global function with example usage: stack_and_flatten_spaces(input_tensor_dict, observation_spaces_dict) and returns: dict