class_docstr: Common superclass for SAC runners, implementing the main training controls.
docstr: See :py:meth:`~maze.train.trainers.common.training_runner.TrainingRunner.setup`.
docstr_tok: ["See",":","py",":","meth",":","`","~maze.train.trainers.common.training_runner.TrainingRunner.setup","`","."]
code: def setup(self, cfg: DictConfig) -> None:
      """
      See :py:meth:`~maze.train.trainers.common.training_runner.TrainingRunner.setup`.
      """
  
      super().setup(cfg)
  
      assert isinstance(
          self._model_composer.critic, TorchStateActionCritic
      ), "Please specify a state action critic for SAC."
  
      worker_policy = self._model_composer.policy
      worker_policy.to("cpu")
  
      model = TorchActorCritic(
          policy=copy.deepcopy(self._model_composer.policy),
          critic=self._model_composer.critic,
          device=cfg.algorithm.device,
      )
  
      worker_env_instance_seeds = [
          self.maze_seeding.generate_env_instance_seed()
          for _ in range(cfg.algorithm.num_actors)
      ]
      replay_buffer_seed = (
          self.maze_seeding.generate_env_instance_seed()
      )
  
      # initialize best model selection
      self._model_selection = BestModelSelection(
          dump_file=self.state_dict_dump_file,
          model=model,
          dump_interval=self.dump_interval,
      )
  
      # initialize the env and enable statistics collection
      evaluator = None
      if cfg.algorithm.rollout_evaluator.n_episodes > 0:
          eval_env = self.create_distributed_eval_env(
              self.env_factory,
              self.eval_concurrency,
              logging_prefix="eval",
          )
          eval_env_instance_seeds = [
              self.maze_seeding.generate_env_instance_seed()
              for _ in range(self.eval_concurrency)
          ]
          eval_env.seed(eval_env_instance_seeds)
  
          # initialize rollout evaluator
          evaluator = Factory(base_type=RolloutEvaluator).instantiate(
              cfg.algorithm.rollout_evaluator,
              eval_env=eval_env,
              model_selection=self._model_selection,
          )
  
      # Init replay buffer
      replay_buffer = UniformReplayBuffer(
          cfg.algorithm.replay_buffer_size, seed=replay_buffer_seed
      )
      if cfg.runner.initial_demonstration_trajectories:
          self.load_replay_buffer(replay_buffer=replay_buffer, cfg=cfg)
      else:
          self.init_replay_buffer(
              replay_buffer=replay_buffer,
              initial_sampling_policy=cfg.algorithm.initial_sampling_policy,
              initial_buffer_size=cfg.algorithm.initial_buffer_size,
              replay_buffer_seed=replay_buffer_seed,
              split_rollouts_into_transitions=cfg.algorithm.split_rollouts_into_transitions,
              n_rollout_steps=cfg.algorithm.n_rollout_steps,
              env_factory=self.env_factory,
          )
  
      distributed_actors = self.create_distributed_rollout_workers(
          env_factory=self.env_factory,
          worker_policy=worker_policy,
          n_rollout_steps=cfg.algorithm.n_rollout_steps,
          n_workers=cfg.algorithm.num_actors,
          batch_size=cfg.algorithm.batch_size,
          rollouts_per_iteration=cfg.algorithm.rollouts_per_iteration,
          split_rollouts_into_transitions=cfg.algorithm.split_rollouts_into_transitions,
          env_instance_seeds=worker_env_instance_seeds,
          replay_buffer=replay_buffer,
      )
  
      # initialize trainer
      self._trainer = SAC(
          algorithm_config=cfg.algorithm,
          learner_model=model,
          distributed_actors=distributed_actors,
          model_selection=self._model_selection,
          evaluator=evaluator,
      )
  
      # initialize model from input_dir
      self._init_trainer_from_input_dir(
          trainer=self._trainer,
          state_dict_dump_file=self.state_dict_dump_file,
          input_dir=cfg.input_dir,
      )
  
imports: import copy
  from abc import abstractmethod
  from dataclasses import dataclass
  from typing import Callable, List, Union
  from omegaconf import DictConfig
  from torch.utils.data import Dataset
  from maze.core.agent.policy import Policy
  from maze.core.agent.torch_actor_critic import TorchActorCritic
  from maze.core.agent.torch_policy import TorchPolicy
  from maze.core.agent.torch_state_action_critic import TorchStateActionCritic
  from maze.core.annotations import override
  from maze.core.env.maze_env import MazeEnv
  from maze.core.log_stats.log_stats import LogStatsAggregator, LogStatsLevel, get_stats_logger
  from maze.core.rollout.rollout_generator import RolloutGenerator
  from maze.core.trajectory_recording.datasets.in_memory_dataset import InMemoryDataset
  from maze.core.trajectory_recording.records.trajectory_record import SpacesTrajectoryRecord
  from maze.core.utils.config_utils import SwitchWorkingDirectoryToInput
  from maze.core.utils.factory import Factory
  from maze.train.parallelization.distributed_actors.base_distributed_workers_with_buffer import BaseDistributedWorkersWithBuffer
  from maze.train.parallelization.distributed_actors.dummy_distributed_workers_with_buffer import DummyDistributedWorkersWithBuffer
  from maze.train.parallelization.vector_env.sequential_vector_env import SequentialVectorEnv
  from maze.train.parallelization.vector_env.structured_vector_env import StructuredVectorEnv
  from maze.train.trainers.common.evaluators.rollout_evaluator import RolloutEvaluator
  from maze.train.trainers.common.model_selection.best_model_selection import BestModelSelection
  from maze.train.trainers.common.replay_buffer.replay_buffer import BaseReplayBuffer
  from maze.train.trainers.common.replay_buffer.uniform_replay_buffer import UniformReplayBuffer
  from maze.train.trainers.common.training_runner import TrainingRunner
  from maze.train.trainers.sac.sac_trainer import SAC
  from maze.utils.process import query_cpu
label_desc: Use image node_id 2 for calling the SACRunner obj's underlying member method code with example usage: obj.setup(cfg) without return types