class_docstr: A trainer that organizes the training logistics.
class_docstr_tok: ["A","trainer","that","organizes","the","training","logistics","."]
docstr: Initialization.
  
  Args:
      policy - The policy network to use.
      solver - The ES algorithm for optimization.
      train_task - The task for training.
      test_task - The task for evaluation.
      max_iter - Maximum number of training iterations.
      log_interval - Interval for logging.
      test_interval - Interval for tests.
      n_repeats - Number of rollout repetitions.
      n_evaluations - Number of tests to conduct.
      seed - Random seed to use.
      debug - Whether to turn on the debug flag.
      use_for_loop - Use for loop for rollouts.
      normalize_obs - Whether to use an observation normalizer.
      model_dir - Directory to save/load model.
      log_dir - Directory to dump logs.
      logger - Logger.
      log_scores_fn - custom function to log the scores array. Expects input:
          `current_iter`: int, `scores`: jnp.ndarray, 'stage': str = "train" | "test"
docstr_tok: ["Initialization",".","Args",":","policy","-","The","policy","network","to","use",".","solver","-","The","ES","algorithm","for","optimization",".","train_task","-","The","task","for","training",".","test_task","-","The","task","for","evaluation",".","max_iter","-","Maximum","number","of","training","iterations",".","log_interval","-","Interval","for","logging",".","test_interval","-","Interval","for","tests",".","n_repeats","-","Number","of","rollout","repetitions",".","n_evaluations","-","Number","of","tests","to","conduct",".","seed","-","Random","seed","to","use",".","debug","-","Whether","to","turn","on","the","debug","flag",".","use_for_loop","-","Use","for","loop","for","rollouts",".","normalize_obs","-","Whether","to","use","an","observation","normalizer",".","model_dir","-","Directory","to","save\/load","model",".","log_dir","-","Directory","to","dump","logs",".","logger","-","Logger",".","log_scores_fn","-","custom","function","to","log","the","scores","array",".","Expects","input",":","`","current_iter","`",":","int",",","`","scores","`",":","jnp.ndarray",",","'stage","'",":","str","=","``","train","''","|","``","test","''"]
code: def __init__(
      self,
      policy: PolicyNetwork,
      solver: NEAlgorithm,
      train_task: VectorizedTask,
      test_task: VectorizedTask,
      max_iter: int = 1000,
      log_interval: int = 20,
      test_interval: int = 100,
      n_repeats: int = 1,
      test_n_repeats: int = 1,
      n_evaluations: int = 100,
      seed: int = 42,
      debug: bool = False,
      use_for_loop: bool = False,
      normalize_obs: bool = False,
      model_dir: str = None,
      log_dir: str = None,
      logger: logging.Logger = None,
      log_scores_fn: Optional[
          Callable[[int, jnp.ndarray, str], None]
      ] = None,
  ):
      """Initialization.
  
      Args:
          policy - The policy network to use.
          solver - The ES algorithm for optimization.
          train_task - The task for training.
          test_task - The task for evaluation.
          max_iter - Maximum number of training iterations.
          log_interval - Interval for logging.
          test_interval - Interval for tests.
          n_repeats - Number of rollout repetitions.
          n_evaluations - Number of tests to conduct.
          seed - Random seed to use.
          debug - Whether to turn on the debug flag.
          use_for_loop - Use for loop for rollouts.
          normalize_obs - Whether to use an observation normalizer.
          model_dir - Directory to save/load model.
          log_dir - Directory to dump logs.
          logger - Logger.
          log_scores_fn - custom function to log the scores array. Expects input:
              `current_iter`: int, `scores`: jnp.ndarray, 'stage': str = "train" | "test"
      """
  
      if logger is None:
          self._logger = create_logger(
              name="Trainer", log_dir=log_dir, debug=debug
          )
      else:
          self._logger = logger
  
      self._log_interval = log_interval
      self._test_interval = test_interval
      self._max_iter = max_iter
      self.model_dir = model_dir
      self._log_dir = log_dir
  
      self._log_scores_fn = log_scores_fn or (lambda x, y, z: None)
  
      self._obs_normalizer = ObsNormalizer(
          obs_shape=train_task.obs_shape,
          dummy=not normalize_obs,
      )
  
      self.solver = solver
      self.sim_mgr = SimManager(
          n_repeats=n_repeats,
          test_n_repeats=test_n_repeats,
          pop_size=solver.pop_size,
          n_evaluations=n_evaluations,
          policy_net=policy,
          train_vec_task=train_task,
          valid_vec_task=test_task,
          seed=seed,
          obs_normalizer=self._obs_normalizer,
          use_for_loop=use_for_loop,
          logger=self._logger,
      )
  
file_path: evojax/evojax/trainer.py
imports: import logging
  import time
  from typing import Optional, Callable
  import jax.numpy
  import numpy
  from evojax.task.base import VectorizedTask
  from evojax.policy import PolicyNetwork
  from evojax.algo import NEAlgorithm
  from evojax.algo import QualityDiversityMethod
  from evojax.sim_mgr import SimManager
  from evojax.obs_norm import ObsNormalizer
  from evojax.util import create_logger
  from evojax.util import load_model
  from evojax.util import save_model
  from evojax.util import save_lattices
label_desc: Use image node_id 1 to create a new Trainer object from inherited base classes: object with example: obj = Trainer(policy, solver, train_task, test_task, max_iter, log_interval, test_interval, n_repeats, test_n_repeats, n_evaluations, seed, debug, use_for_loop, normalize_obs, model_dir, log_dir, logger, log_scores_fn)