def new_epoch(epoch_id):
      """Updates the lists defined above to include the new epoch."""
      all_epochs = list(self._trajectories.keys())
      max_epoch = max(all_epochs) + 1

      # Calculate the numbers of slices for the new epoch.
      epoch_to_ns_slices[epoch_id] = np.array([
          _n_slices(trajectory, max_slice_length, margin)
          for trajectory in self._trajectories[epoch_id]
      ])
      epoch_to_total_n_slices[epoch_id] = np.sum(
          epoch_to_ns_slices[epoch_id]
      )

      # Update the indices of epochs to sample from.
      new_epoch_indices = epochs or all_epochs
      new_epoch_indices = [
          # So -1 means "last".
          ep % max_epoch for ep in new_epoch_indices
      ]
      # Remove duplicates and consider only epochs where some trajectories
      # were recorded and that we have processed in new_epoch.
      new_epoch_indices = [
          epoch_id for epoch_id in set(new_epoch_indices)
          if self._trajectories[epoch_id] and epoch_id in epoch_to_ns_slices
      ]
      epoch_indices[:] = new_epoch_indices

      nonlocal sampling_epoch_weights
      sampling_epoch_weights = np.array(list(
          epoch_to_total_n_slices[ep] for ep in epoch_indices
      ))
Updates the lists defined above to include the new epoch.
trax/rl/task.py