def save_to_file(self, file_name='rl.pkl',
                   task_file_name='trajectories.pkl'):
    """Save current epoch number and average returns to file."""
    assert self._output_dir is not None
    task_path = os.path.join(self._output_dir, task_file_name)
    self._task.save_to_file(task_path)
    file_path = os.path.join(self._output_dir, file_name)
    dictionary = {'epoch': self._epoch, 'avg_returns': self._avg_returns}
    for eval_t in self._eval_temperatures:
      dictionary['avg_returns_temperature_{}'.format(
          eval_t)] = self._avg_returns_temperatures[eval_t]
    with tf.io.gfile.GFile(file_path, 'wb') as f:
      pickle.dump(dictionary, f)
Save current epoch number and average returns to file.
trax/rl/training.py