class_docstr: Simulation manager.
class_docstr_tok: ["Simulation","manager","."]
docstr: Rollout using jax.lax.scan.
docstr_tok: ["Rollout","using","jax.lax.scan","."]
code: def _scan_loop_eval(
      self, params: jnp.ndarray, test: bool
  ) -> Tuple[jnp.ndarray, TaskState]:
      """Rollout using jax.lax.scan."""
      policy_reset_func = self._policy_reset_fn
      if test:
          n_repeats = self._test_n_repeats
          task_reset_func = self._valid_reset_fn
          rollout_func = self._valid_rollout_fn
          params = duplicate_params(
              params[None, :], self._n_evaluations, False
          )
      else:
          n_repeats = self._n_repeats
          task_reset_func = self._train_reset_fn
          rollout_func = self._train_rollout_fn
  
      # Suppose pop_size=2 and n_repeats=3.
      # For multi-agents training, params become
      #   a1, a2, ..., an  (individual 1 params)
      #   b1, b2, ..., bn  (individual 2 params)
      #   a1, a2, ..., an  (individual 1 params)
      #   b1, b2, ..., bn  (individual 2 params)
      #   a1, a2, ..., an  (individual 1 params)
      #   b1, b2, ..., bn  (individual 2 params)
      # For non-ma training, params become
      #   a1, a2, ..., an  (individual 1 params)
      #   a1, a2, ..., an  (individual 1 params)
      #   a1, a2, ..., an  (individual 1 params)
      #   b1, b2, ..., bn  (individual 2 params)
      #   b1, b2, ..., bn  (individual 2 params)
      #   b1, b2, ..., bn  (individual 2 params)
      params = duplicate_params(params, n_repeats, self._ma_training)
  
      self._key, reset_keys = get_task_reset_keys(
          self._key,
          test,
          self._pop_size,
          self._n_evaluations,
          n_repeats,
          self._ma_training,
      )
  
      # Reset the tasks and the policy.
      task_state = task_reset_func(reset_keys)
      policy_state = policy_reset_func(task_state)
      if self._num_device > 1:
          params = split_params_for_pmap(params)
          task_state = split_states_for_pmap(task_state)
          policy_state = split_states_for_pmap(policy_state)
  
      # Do the rollouts.
      scores, all_obs, masks, final_states = rollout_func(
          task_state, policy_state, params, self.obs_params
      )
      if self._num_device > 1:
          all_obs = reshape_data_from_pmap(all_obs)
          masks = reshape_data_from_pmap(masks)
          final_states = merge_state_from_pmap(final_states)
  
      if not test and not self.obs_normalizer.is_dummy:
          self.obs_params = (
              self.obs_normalizer.update_normalization_params(
                  obs_buffer=all_obs,
                  obs_mask=masks,
                  obs_params=self.obs_params,
              )
          )
  
      if self._ma_training:
          if not test:
              # In training, each agent has different parameters.
              scores = jnp.mean(
                  scores.ravel().reshape((n_repeats, -1)), axis=0
              )
          else:
              # In tests, they share the same parameters.
              scores = jnp.mean(
                  scores.ravel().reshape((n_repeats, -1)), axis=1
              )
      else:
          scores = jnp.mean(
              scores.ravel().reshape((-1, n_repeats)), axis=-1
          )
  
      # Note: QD methods do not support ma_training for now.
      if not self._ma_training:
          final_states = tree_map(
              lambda x: x.reshape(
                  (scores.shape[0], n_repeats, *x.shape[1:])
              ),
              final_states,
          )
  
      return scores, self._bd_summarize_fn(final_states)
  
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 4 for calling the SimManager obj's underlying member method code with example usage: obj._scan_loop_eval(params, test) and returns: scores,  self