class_docstr: Simulation manager.
class_docstr_tok: ["Simulation","manager","."]
docstr: Rollout using for loop (no multi-device or ma_training yet).
docstr_tok: ["Rollout","using","for","loop","(","no","multi-device","or","ma_training","yet",")","."]
code: def _for_loop_eval(
      self, params: jnp.ndarray, test: bool
  ) -> Tuple[jnp.ndarray, TaskState]:
      """Rollout using for loop (no multi-device or ma_training yet)."""
      policy_reset_func = self._policy_reset_fn
      policy_act_func = self._policy_act_fn
      if test:
          n_repeats = self._test_n_repeats
          task_reset_func = self._valid_reset_fn
          task_step_func = self._valid_step_fn
          task_max_steps = self._valid_max_steps
          params = duplicate_params(
              params[None, :], self._n_evaluations, False
          )
      else:
          n_repeats = self._n_repeats
          task_reset_func = self._train_reset_fn
          task_step_func = self._train_step_fn
          task_max_steps = self._train_max_steps
  
      params = duplicate_params(params, n_repeats, self._ma_training)
  
      # Start rollout.
      self._key, reset_keys = get_task_reset_keys(
          self._key,
          test,
          self._pop_size,
          self._n_evaluations,
          n_repeats,
          self._ma_training,
      )
      task_state = task_reset_func(reset_keys)
      policy_state = policy_reset_func(task_state)
      scores = jnp.zeros(params.shape[0])
      valid_mask = jnp.ones(params.shape[0])
      start_time = time.perf_counter()
      rollout_steps = 0
      sim_steps = 0
      for i in range(task_max_steps):
          actions, policy_state = policy_act_func(
              task_state, params, policy_state
          )
          task_state, reward, done = task_step_func(task_state, actions)
          scores, valid_mask = update_score_and_mask(
              scores, reward, valid_mask, done
          )
          rollout_steps += 1
          sim_steps = sim_steps + valid_mask
          if all_done(valid_mask):
              break
      time_cost = time.perf_counter() - start_time
      self._logger.debug(
          "{} steps/s, mean.steps={}".format(
              int(rollout_steps * task_state.obs.shape[0] / time_cost),
              sim_steps.sum() / task_state.obs.shape[0],
          )
      )
  
      return report_score(scores, n_repeats), task_state
  
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 3 for calling the SimManager obj's underlying member method code with example usage: obj._for_loop_eval(params, test) and returns: report_score,  task_state