class_docstr: Base class for training runner implementations.
docstr: Sets up prerequisites to training.
  Includes wrapping the environment for observation normalization, instantiating the model composer etc.
  :param cfg: DictConfig defining components to initialize.
docstr_tok: ["Sets","up","prerequisites","to","training",".","Includes","wrapping","the","environment","for","observation","normalization",",","instantiating","the","model","composer","etc",".",":","param","cfg",":","DictConfig","defining","components","to","initialize","."]
code: def setup(self, cfg: DictConfig) -> None:
      """
      Sets up prerequisites to training.
      Includes wrapping the environment for observation normalization, instantiating the model composer etc.
      :param cfg: DictConfig defining components to initialize.
      """
  
      self._cfg = cfg
  
      # Generate a random state used for sampling random seeds for the envs and agents
      self.maze_seeding = MazeSeeding(
          env_seed=cfg.seeding.env_base_seed,
          agent_seed=cfg.seeding.agent_base_seed,
          cudnn_determinism_flag=cfg.seeding.cudnn_determinism_flag,
          explicit_agent_seeds=None,
          explicit_env_seeds=None,
          shuffle_seeds=False,
      )
  
      with SwitchWorkingDirectoryToInput(cfg.input_dir):
          assert isinstance(cfg.env, DictConfig) or isinstance(
              cfg.env, Callable
          )
          wrapper_cfg = (
              omegaconf.OmegaConf.to_object(cfg["wrappers"])
              if "wrappers" in cfg
              else {}
          )
  
          # if the observation normalization is already available, read it from the input directory
          if isinstance(cfg.env, DictConfig):
              self.env_factory = EnvFactory(
                  omegaconf.OmegaConf.to_object(cfg["env"]), wrapper_cfg
              )
          elif isinstance(cfg.env, Callable):
              env_fn = omegaconf.OmegaConf.to_container(cfg)["env"]
              self.env_factory = (
                  lambda: WrapperFactory.wrap_from_config(
                      env_fn(), wrapper_cfg
                  )
              )
  
          normalization_env = self.env_factory()
          normalization_env.seed(
              self.maze_seeding.generate_env_instance_seed()
          )
  
      # Observation normalization
      self._normalization_statistics = obtain_normalization_statistics(
          normalization_env, n_samples=self.normalization_samples
      )
      if self._normalization_statistics:
          self.env_factory = make_normalized_env_factory(
              self.env_factory, self._normalization_statistics
          )
          # dump statistics to current working directory
          assert isinstance(
              normalization_env, ObservationNormalizationWrapper
          )
          normalization_env.dump_statistics()
  
      # Generate an agent seed and set the seed globally for the model initialization
      set_seeds_globally(
          self.maze_seeding.global_seed,
          self.maze_seeding.cudnn_determinism_flag,
          info_txt=f"training runner (Pid:{os.getpid()})",
      )
  
      # init model composer
      composer_type = Factory(
          base_type=BaseModelComposer
      ).type_from_name(cfg.model["_target_"])
      composer_type.check_model_config(cfg.model)
  
      # todo Factory.instantiate returns specified dicts as DictConfig, i.e. many specified types are wrong. How do we
      #  go about this? DictConfig behaves similarly to Dict for all intents and purposes, but typing is still off/
      #  misleading. This is independent from our Python training API and can apparently not be changed, i.e. kwargs
      #  seems to be always converted to DictConfig/ListConfig.
      self._model_composer = Factory(
          base_type=BaseModelComposer
      ).instantiate(
          cfg.model,
          action_spaces_dict=normalization_env.action_spaces_dict,
          observation_spaces_dict=normalization_env.observation_spaces_dict,
          agent_counts_dict=normalization_env.agent_counts_dict,
      )
  
      SpacesConfig(
          self._model_composer.action_spaces_dict,
          self._model_composer.observation_spaces_dict,
          self._model_composer.agent_counts_dict,
      ).save(self.spaces_config_dump_file)
  
      # Should be done after the normalization runs, otherwise stats from those will get logged as well.
      setup_logging(job_config=cfg)
  
      # close normalization env
      normalization_env.close()
  
imports: import logging
  import os
  import dataclasses
  from typing import Optional, Callable, Union
  import omegaconf
  from omegaconf import DictConfig
  from maze.core.env.structured_env import StructuredEnv
  from maze.core.env.structured_env_spaces_mixin import StructuredEnvSpacesMixin
  from maze.core.utils.config_utils import EnvFactory, SwitchWorkingDirectoryToInput
  from maze.core.utils.factory import Factory
  from maze.core.wrappers.observation_normalization.normalization_strategies.base import StructuredStatisticsType
  from maze.core.utils.seeding import set_seeds_globally, MazeSeeding
  from maze.core.wrappers.observation_normalization.observation_normalization_utils import obtain_normalization_statistics, make_normalized_env_factory
  from maze.core.wrappers.observation_normalization.observation_normalization_wrapper import ObservationNormalizationWrapper
  from maze.core.wrappers.wrapper_factory import WrapperFactory
  from maze.perception.models.model_composer import BaseModelComposer
  from maze.perception.models.space_config import SpacesConfig
  from maze.runner import Runner
  from maze.train.trainers.common.model_selection.best_model_selection import BestModelSelection
  from maze.train.trainers.common.trainer import Trainer
  from maze.utils.bcolors import BColors
  from maze.utils.log_stats_utils import setup_logging
label_desc: Use image node_id 1 for calling the TrainingRunner obj's underlying member method code with example usage: obj.setup(cfg) without return types