class_docstr: A trainer that organizes the training logistics.
class_docstr_tok: ["A","trainer","that","organizes","the","training","logistics","."]
docstr: Start the training / test process.
docstr_tok: ["Start","the","training","\/","test","process","."]
code: def run(self, demo_mode: bool = False) -> float:
      """Start the training / test process."""
  
      if self.model_dir is not None:
          params, obs_params = load_model(model_dir=self.model_dir)
          self.sim_mgr.obs_params = obs_params
          self._logger.info(
              "Loaded model parameters from {}.".format(self.model_dir)
          )
      else:
          params = None
  
      if demo_mode:
          if params is None:
              raise ValueError("No policy parameters to evaluate.")
          self._logger.info("Start to test the parameters.")
          scores = np.array(
              self.sim_mgr.eval_params(params=params, test=True)[0]
          )
          self._logger.info(
              "[TEST] #tests={0}, max={1:.4f}, avg={2:.4f}, min={3:.4f}, "
              "std={4:.4f}".format(
                  scores.size,
                  scores.max(),
                  scores.mean(),
                  scores.min(),
                  scores.std(),
              )
          )
          return scores.mean()
      else:
          self._logger.info(
              "Start to train for {} iterations.".format(self._max_iter)
          )
  
          if params is not None:
              # Continue training from the breakpoint.
              self.solver.best_params = params
  
          best_score = -float("Inf")
  
          for i in range(self._max_iter):
              start_time = time.perf_counter()
              params = self.solver.ask()
              self._logger.debug(
                  "solver.ask time: {0:.4f}s".format(
                      time.perf_counter() - start_time
                  )
              )
  
              start_time = time.perf_counter()
              scores, bds = self.sim_mgr.eval_params(
                  params=params, test=False
              )
              self._logger.debug(
                  "sim_mgr.eval_params time: {0:.4f}s".format(
                      time.perf_counter() - start_time
                  )
              )
  
              start_time = time.perf_counter()
              if isinstance(self.solver, QualityDiversityMethod):
                  self.solver.observe_bd(bds)
              self.solver.tell(fitness=scores)
              self._logger.debug(
                  "solver.tell time: {0:.4f}s".format(
                      time.perf_counter() - start_time
                  )
              )
  
              if i > 0 and i % self._log_interval == 0:
                  scores = np.array(scores)
                  self._logger.info(
                      "Iter={0}, size={1}, max={2:.4f}, "
                      "avg={3:.4f}, min={4:.4f}, std={5:.4f}".format(
                          i,
                          scores.size,
                          scores.max(),
                          scores.mean(),
                          scores.min(),
                          scores.std(),
                      )
                  )
                  self._log_scores_fn(i, scores, "train")
  
              if i > 0 and i % self._test_interval == 0:
                  best_params = self.solver.best_params
                  test_scores, _ = self.sim_mgr.eval_params(
                      params=best_params, test=True
                  )
                  self._logger.info(
                      "[TEST] Iter={0}, #tests={1}, max={2:.4f}, avg={3:.4f}, "
                      "min={4:.4f}, std={5:.4f}".format(
                          i,
                          test_scores.size,
                          test_scores.max(),
                          test_scores.mean(),
                          test_scores.min(),
                          test_scores.std(),
                      )
                  )
                  self._log_scores_fn(i, test_scores, "test")
                  mean_test_score = test_scores.mean()
                  save_model(
                      model_dir=self._log_dir,
                      model_name="iter_{}".format(i),
                      params=best_params,
                      obs_params=self.sim_mgr.obs_params,
                      best=mean_test_score > best_score,
                  )
                  best_score = max(best_score, mean_test_score)
  
          # Test and save the final model.
          best_params = self.solver.best_params
          test_scores, _ = self.sim_mgr.eval_params(
              params=best_params, test=True
          )
          self._logger.info(
              "[TEST] Iter={0}, #tests={1}, max={2:.4f}, avg={3:.4f}, "
              "min={4:.4f}, std={5:.4f}".format(
                  self._max_iter,
                  test_scores.size,
                  test_scores.max(),
                  test_scores.mean(),
                  test_scores.min(),
                  test_scores.std(),
              )
          )
          mean_test_score = test_scores.mean()
          save_model(
              model_dir=self._log_dir,
              model_name="final",
              params=best_params,
              obs_params=self.sim_mgr.obs_params,
              best=mean_test_score > best_score,
          )
          best_score = max(best_score, mean_test_score)
          if isinstance(self.solver, QualityDiversityMethod):
              save_lattices(
                  log_dir=self._log_dir,
                  file_name="qd_lattices",
                  fitness_lattice=self.solver.fitness_lattice,
                  params_lattice=self.solver.params_lattice,
                  occupancy_lattice=self.solver.occupancy_lattice,
              )
          self._logger.info(
              "Training done, best_score={0:.4f}".format(best_score)
          )
  
          return best_score
  
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 2 for calling the Trainer obj's underlying member method code with example usage: obj.run(demo_mode) and returns: scores, best_score