def run(self, n_epochs=1, n_epochs_is_total_epochs=False):
    """Runs this loop for n epochs.

    Args:
      n_epochs: Stop training after completing n steps.
      n_epochs_is_total_epochs: if True, consider n_epochs as the total
        number of epochs to train, including previously trained ones
    """
    with self._open_summary_writer() as sw:
      n_epochs_to_run = n_epochs
      if n_epochs_is_total_epochs:
        n_epochs_to_run -= self._epoch
      cur_n_interactions = 0
      for _ in range(n_epochs_to_run):
        self._epoch += 1
        cur_time = time.time()
        avg_return = self._collect_trajectories()
        self._avg_returns.append(avg_return)
        if self._n_trajectories_per_epoch:
          supervised.trainer_lib.log(
              'Collecting %d episodes took %.2f seconds.'
              % (self._n_trajectories_per_epoch, time.time() - cur_time))
        else:
          supervised.trainer_lib.log(
              'Collecting %d interactions took %.2f seconds.'
              % (self._n_interactions_per_epoch, time.time() - cur_time))
        supervised.trainer_lib.log(
            'Average return in epoch %d was %.2f.' % (self._epoch, avg_return))
        if self._n_eval_episodes > 0:
          for steps in self._eval_steps:
            for eval_t in self._eval_temperatures:
              avg_return_temperature = self.task.collect_trajectories(
                  functools.partial(self.policy, temperature=eval_t),
                  n_trajectories=self._n_eval_episodes,
                  max_steps=steps,
                  only_eval=True)
              supervised.trainer_lib.log(
                  'Eval return in epoch %d with temperature %.2f was %.2f.'
                  % (self._epoch, eval_t, avg_return_temperature))
              self._avg_returns_temperatures[eval_t][steps].append(
                  avg_return_temperature)

        if sw is not None:
          sw.scalar('timing/collect', time.time() - cur_time,
                    step=self._epoch)
          sw.scalar('rl/avg_return', avg_return, step=self._epoch)
          if self._n_eval_episodes > 0:
            for steps in self._eval_steps:
              for eval_t in self._eval_temperatures:
                sw.scalar(
                    'rl/avg_return_temperature%.2f_steps%d' % (eval_t, steps),
                    self._avg_returns_temperatures[eval_t][steps][-1],
                    step=self._epoch)
          sw.scalar('rl/n_interactions', self.task.n_interactions(),
                    step=self._epoch)
          sw.scalar('rl/n_interactions_per_second',
                    (self.task.n_interactions() - cur_n_interactions)/ \
                    (time.time() - cur_time),
                    step=self._epoch)
          cur_n_interactions = self.task.n_interactions()
          sw.scalar('rl/n_trajectories', self.task.n_trajectories(),
                    step=self._epoch)
          sw.flush()

        cur_time = time.time()
        self.train_epoch()
        supervised.trainer_lib.log(
            'RL training took %.2f seconds.' % (time.time() - cur_time))

        if self._output_dir is not None and self._epoch == 1:
          self.save_gin(sw)
        if self._output_dir is not None:
          self.save_to_file()
Runs this loop for n epochs.

    Args:
      n_epochs: Stop training after completing n steps.
      n_epochs_is_total_epochs: if True, consider n_epochs as the total
        number of epochs to train, including previously trained ones
    
trax/rl/training.py