text: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import ray
  from ray.rllib.evaluation.sample_batch import SampleBatch
  
  
  def collect_samples(agents, train_batch_size):
      num_timesteps_so_far = 0
      trajectories = []
      # This variable maps the object IDs of trajectories that are currently
      # computed to the agent that they are computed on; we start some initial
      # tasks here.
  
      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(fut_sample)
          num_timesteps_so_far += next_sample.count
          trajectories.append(next_sample)
      return SampleBatch.concat_samples(trajectories)
  