def _filter_trajectories(self, trajectory_stream):
    """Filter trajectories based on the quantile."""
    def trajectory_return(trajectory):
      return trajectory.timesteps[0].return_

    trajectory_buffer = []
    for trajectory in trajectory_stream:
      trajectory_buffer.append(trajectory)
      if len(trajectory_buffer) == self._filter_buffer_size:
        n_best = int((1 - self._quantile) * self._filter_buffer_size) or 1
        trajectory_buffer.sort(key=trajectory_return, reverse=True)
        yield from trajectory_buffer[:n_best]
        trajectory_buffer.clear()
Filter trajectories based on the quantile.
trax/rl/training.py