code: def check_timestep_batch(batch, array_type, ignored_fields=()):
      """Check a TimeStepBatch of any array type that has .shape.
  
      Args:
          batch (TimeStepBatch): Batch of timesteps.
          array_type (type): Array type.
          ignored_fields (set[str]): Set of fields to ignore checking on.
  
      Raises:
          ValueError: If an invariant of TimeStepBatch is broken.
      """
      # pylint:disable=too-many-branches
      fields = {
          field: getattr(batch, field)
          for field in [
              "env_spec",
              "rewards",
              "rewards",
              "observations",
              "actions",
              "next_observations",
              "step_types",
              "agent_infos",
              "episode_infos",
              "env_infos",
          ]
          if field not in ignored_fields
      }
      env_spec = fields.get("env_spec", None)
      inferred_batch_size = None
      inferred_batch_size_field = None
      for field, value in fields.items():
          if field in [
              "observations",
              "actions",
              "rewards",
              "next_observations",
              "step_types",
          ]:
              if not isinstance(value, array_type):
                  raise ValueError(
                      f"{field} is not of type {array_type!r}"
                  )
          if hasattr(value, "shape"):
              if inferred_batch_size is None:
                  inferred_batch_size = value.shape[0]
                  inferred_batch_size_field = field
              elif value.shape[0] != inferred_batch_size:
                  raise ValueError(
                      f"{field} has batch size {value.shape[0]}, but "
                      f"must have batch size {inferred_batch_size} "
                      f"to match {inferred_batch_size_field}"
                  )
              if env_spec and field in [
                  "observations",
                  "next_observations",
              ]:
                  if not _space_soft_contains(
                      env_spec.observation_space, value[0]
                  ):
                      raise ValueError(
                          f"Each {field[:-1]} has shape {value[0].shape} "
                          f"but must match the observation_space "
                          f"{env_spec.observation_space}"
                      )
                  if isinstance(
                      value[0], np.ndarray
                  ) and not env_spec.observation_space.contains(
                      value[0]
                  ):
                      warnings.warn(
                          f"Observation {value[0]!r} is outside "
                          f"observation_space {env_spec.observation_space}"
                      )
              if env_spec and field == "actions":
                  if not _space_soft_contains(
                      env_spec.action_space, value[0]
                  ):
                      raise ValueError(
                          f"Each {field[:-1]} has shape {value[0].shape} "
                          f"but must match the action_space "
                          f"{env_spec.action_space}"
                      )
              if field in ["rewards", "step_types"]:
                  if value.shape != (inferred_batch_size,):
                      raise ValueError(
                          f"{field} has shape {value.shape} "
                          f"but must have batch size "
                          f"{inferred_batch_size} to match "
                          f"{inferred_batch_size_field}"
                      )
          if field in ["agent_infos", "env_infos", "episode_infos"]:
              for key, val in value.items():
                  if not isinstance(val, (array_type, dict)):
                      raise ValueError(
                          f"Entry {key!r} in {field} is of type {type(val)}"
                          f"but must be {array_type!r} or dict"
                      )
                  if hasattr(val, "shape"):
                      if val.shape[0] != inferred_batch_size:
                          raise ValueError(
                              f"Entry {key!r} in {field} has batch size "
                              f"{val.shape[0]} but must have batch size "
                              f"{inferred_batch_size} to match "
                              f"{inferred_batch_size_field}"
                          )
  
          if (
              field == "step_types"
              and isinstance(value, np.ndarray)
              and value.dtype  # Only numpy arrays support custom dtypes.
              != StepType
          ):
              raise ValueError(
                  f"step_types has dtype {value.dtype} but must have "
                  f"dtype StepType"
              )
  
imports: from dataclasses import dataclass
  import enum
  from typing import Dict, List
  import warnings
  import numpy
  from garage.np import concat_tensor_dict_list, pad_batch_array, slice_nested_dict, stack_tensor_dict_list
label_desc: Use image node_id 2 for calling a global function with example usage: check_timestep_batch(batch, array_type, ignored_fields) without return types