code: def obtain_normalization_statistics(
      env: Union[MazeEnv, ObservationNormalizationWrapper],
      n_samples: int,
  ) -> Optional[StructuredStatisticsType]:
      """Obtain the normalization statistics of a given environment.
  
      * Returns None, if the ObservationNormalizationWrapper is not implemented
      * Returns the loaded statistics, if available
      * Runs the estimation and returns the newly calculated statistics, if not loaded previously
  
      :param env: Environment with applied ObservationNormalizationWrapper (function returns None immediately if this is
          not the case.
      :param n_samples: Number of samples (=steps) to collect normalization statistics at the beginning of
          the training.
      :return: The normalization statistics or None if the ObservationNormalizationWrapper is not implemented by the env.
      """
      if not isinstance(env, ObservationNormalizationWrapper):
          return None
  
      # first check if loading from file was successful
      normalization_statistics = env.loaded_stats
      if normalization_statistics:
          return normalization_statistics
  
      # no statistics available, run estimation
      estimate_observation_normalization_statistics(
          env, n_samples=n_samples
      )
      normalization_statistics = env.get_statistics()
  
      # deactivate statistics collection
      env.set_observation_collection(False)
  
      return normalization_statistics
  
file_path: maze/maze/core/wrappers/observation_normalization/observation_normalization_utils.py
imports: import os
  from typing import Union, Optional, Callable
  from tqdm import tqdm
  from maze.core.env.maze_env import MazeEnv
  from maze.core.env.structured_env import StructuredEnv
  from maze.core.env.structured_env_spaces_mixin import StructuredEnvSpacesMixin
  from maze.core.rollout.rollout_generator import RolloutGenerator
  from maze.core.wrappers.observation_normalization.normalization_strategies.base import StructuredStatisticsType
  from maze.core.wrappers.observation_normalization.observation_normalization_wrapper import ObservationNormalizationWrapper
label_desc: Use image node_id 2 for calling a global function with example usage: obtain_normalization_statistics(env, n_samples) and returns: normalization_statistics, None, normalization_statistics