def truncate(self, indices=None, num_to_keep=1):
    """Truncates trajectories at the specified indices."""

    if indices is None:
      indices = np.arange(self.batch_size)

    self.trajectories.truncate_trajectories(indices, num_to_keep=num_to_keep)
Truncates trajectories at the specified indices.
tensor2tensor/envs/env_problem.py