class_docstr: Simulation manager.
class_docstr_tok: ["Simulation","manager","."]
docstr: Initialization function.
  
  Args:
      n_repeats - Number of repeated parameter evaluations.
      pop_size - Population size.
      n_evaluations - Number of evaluations of the best parameter.
      policy_net - Policy network.
      train_vec_task - Vectorized tasks for training.
      valid_vec_task - Vectorized tasks for validation.
      seed - Random seed.
      obs_normalizer - Observation normalization helper.
      use_for_loop - Use for loop for rollout instead of jax.lax.scan.
      logger - Logger.
docstr_tok: ["Initialization","function",".","Args",":","n_repeats","-","Number","of","repeated","parameter","evaluations",".","pop_size","-","Population","size",".","n_evaluations","-","Number","of","evaluations","of","the","best","parameter",".","policy_net","-","Policy","network",".","train_vec_task","-","Vectorized","tasks","for","training",".","valid_vec_task","-","Vectorized","tasks","for","validation",".","seed","-","Random","seed",".","obs_normalizer","-","Observation","normalization","helper",".","use_for_loop","-","Use","for","loop","for","rollout","instead","of","jax.lax.scan",".","logger","-","Logger","."]
code: def __init__(
      self,
      n_repeats: int,
      test_n_repeats: int,
      pop_size: int,
      n_evaluations: int,
      policy_net: PolicyNetwork,
      train_vec_task: VectorizedTask,
      valid_vec_task: VectorizedTask,
      seed: int = 0,
      obs_normalizer: ObsNormalizer = None,
      use_for_loop: bool = False,
      logger: logging.Logger = None,
  ):
      """Initialization function.
  
      Args:
          n_repeats - Number of repeated parameter evaluations.
          pop_size - Population size.
          n_evaluations - Number of evaluations of the best parameter.
          policy_net - Policy network.
          train_vec_task - Vectorized tasks for training.
          valid_vec_task - Vectorized tasks for validation.
          seed - Random seed.
          obs_normalizer - Observation normalization helper.
          use_for_loop - Use for loop for rollout instead of jax.lax.scan.
          logger - Logger.
      """
  
      if logger is None:
          self._logger = create_logger(name="SimManager")
      else:
          self._logger = logger
  
      self._use_for_loop = use_for_loop
      self._logger.info("use_for_loop={}".format(self._use_for_loop))
      self._key = random.PRNGKey(seed=seed)
      self._n_repeats = n_repeats
      self._test_n_repeats = test_n_repeats
      self._pop_size = pop_size
      self._n_evaluations = max(n_evaluations, jax.local_device_count())
      self._ma_training = train_vec_task.multi_agent_training
  
      self.obs_normalizer = obs_normalizer
      if self.obs_normalizer is None:
          self.obs_normalizer = ObsNormalizer(
              obs_shape=train_vec_task.obs_shape,
              dummy=True,
          )
      self.obs_params = self.obs_normalizer.get_init_params()
  
      self._num_device = jax.local_device_count()
      if self._pop_size % self._num_device != 0:
          raise ValueError(
              "pop_size must be multiples of GPU/TPUs: "
              "pop_size={}, #devices={}".format(
                  self._pop_size, self._num_device
              )
          )
      if self._n_evaluations % self._num_device != 0:
          raise ValueError(
              "n_evaluations must be multiples of GPU/TPUs: "
              "n_evaluations={}, #devices={}".format(
                  self._n_evaluations, self._num_device
              )
          )
  
      def step_once(carry, input_data, task):
          (
              task_state,
              policy_state,
              params,
              obs_params,
              accumulated_reward,
              valid_mask,
          ) = carry
          if task.multi_agent_training:
              num_tasks, num_agents = task_state.obs.shape[:2]
              task_state = task_state.replace(
                  obs=task_state.obs.reshape(
                      (-1, *task_state.obs.shape[2:])
                  )
              )
          org_obs = task_state.obs
          normed_obs = self.obs_normalizer.normalize_obs(
              org_obs, obs_params
          )
          task_state = task_state.replace(obs=normed_obs)
          actions, policy_state = policy_net.get_actions(
              task_state, params, policy_state
          )
          if task.multi_agent_training:
              task_state = task_state.replace(
                  obs=task_state.obs.reshape(
                      (num_tasks, num_agents, *task_state.obs.shape[1:])
                  )
              )
              actions = actions.reshape(
                  (num_tasks, num_agents, *actions.shape[1:])
              )
          task_state, reward, done = task.step(task_state, actions)
          if task.multi_agent_training:
              reward = reward.ravel()
              done = jnp.repeat(done, num_agents, axis=0)
          accumulated_reward = accumulated_reward + reward * valid_mask
          valid_mask = valid_mask * (1 - done.ravel())
          return (
              (
                  task_state,
                  policy_state,
                  params,
                  obs_params,
                  accumulated_reward,
                  valid_mask,
              ),
              (org_obs, valid_mask),
          )
  
      def rollout(
          task_states,
          policy_states,
          params,
          obs_params,
          step_once_fn,
          max_steps,
      ):
          accumulated_rewards = jnp.zeros(params.shape[0])
          valid_masks = jnp.ones(params.shape[0])
          (
              (
                  task_states,
                  policy_states,
                  params,
                  obs_params,
                  accumulated_rewards,
                  valid_masks,
              ),
              (obs_set, obs_mask),
          ) = jax.lax.scan(
              step_once_fn,
              (
                  task_states,
                  policy_states,
                  params,
                  obs_params,
                  accumulated_rewards,
                  valid_masks,
              ),
              (),
              max_steps,
          )
          return accumulated_rewards, obs_set, obs_mask, task_states
  
      self._policy_reset_fn = jax.jit(policy_net.reset)
      self._policy_act_fn = jax.jit(policy_net.get_actions)
  
      if (
          hasattr(train_vec_task, "bd_extractor")
          and train_vec_task.bd_extractor is not None
      ):
          self._bd_summarize_fn = jax.jit(
              train_vec_task.bd_extractor.summarize
          )
      else:
          self._bd_summarize_fn = lambda x: x
  
      # Set up training functions.
      self._train_reset_fn = train_vec_task.reset
      self._train_step_fn = train_vec_task.step
      self._train_max_steps = train_vec_task.max_steps
      self._train_rollout_fn = partial(
          rollout,
          step_once_fn=partial(step_once, task=train_vec_task),
          max_steps=train_vec_task.max_steps,
      )
      if self._num_device > 1:
          self._train_rollout_fn = jax.jit(
              jax.pmap(self._train_rollout_fn, in_axes=(0, 0, 0, None))
          )
  
      # Set up validation functions.
      self._valid_reset_fn = valid_vec_task.reset
      self._valid_step_fn = valid_vec_task.step
      self._valid_max_steps = valid_vec_task.max_steps
      self._valid_rollout_fn = partial(
          rollout,
          step_once_fn=partial(step_once, task=valid_vec_task),
          max_steps=valid_vec_task.max_steps,
      )
      if self._num_device > 1:
          self._valid_rollout_fn = jax.jit(
              jax.pmap(self._valid_rollout_fn, in_axes=(0, 0, 0, None))
          )
  
imports: import logging
  from functools import partial
  from typing import Tuple
  from typing import Union
  import time
  import jax
  import jax.numpy
  from jax import random
  from jax.tree_util import tree_map
  from evojax.obs_norm import ObsNormalizer
  from evojax.task.base import TaskState
  from evojax.task.base import VectorizedTask
  from evojax.policy.base import PolicyState
  from evojax.policy.base import PolicyNetwork
  from evojax.util import create_logger
label_desc: Use image node_id 1 to create a new SimManager object from inherited base classes: object with example: obj = SimManager(n_repeats, test_n_repeats, pop_size, n_evaluations, policy_net, train_vec_task, valid_vec_task, seed, obs_normalizer, use_for_loop, logger)