code: def collect_samples_straggler_mitigation(agents, train_batch_size):
      
  
      num_timesteps_so_far = 0
      trajectories = []
      agent_dict = {}
  
      for agent in agents:
          fut_sample = agent.sample.remote()
          agent_dict[fut_sample] = agent
  
      while num_timesteps_so_far < train_batch_size:
          # TODO(pcm): Make wait support arbitrary iterators and remove the
          # conversion to list here.
          [fut_sample], _ = ray.wait(list(agent_dict))
          agent = agent_dict.pop(fut_sample)
          # Start task with next trajectory and record it in the dictionary.
          fut_sample2 = agent.sample.remote()
          agent_dict[fut_sample2] = agent
  
          next_sample = ray_get_and_free(fut_sample)
          num_timesteps_so_far += next_sample.count
          trajectories.append(next_sample)
  
      logger.info("Discarding {} sample tasks".format(len(agent_dict)))
      return SampleBatch.concat_samples(trajectories)
docs: Collects at least train_batch_size samples.
  
      This is the legacy behavior as of 0.6, and launches extra sample tasks to
      potentially improve performance but can result in many wasted samples.