mirror of
https://github.com/wassname/ray.git
synced 2026-07-07 21:34:11 +08:00
Eric Liang
995ac24a2c
It's possible to configure PPO in a way that ends up discarding most of the samples (they are treated as "stragglers"). Add a warning when this happens, and raise an exception if the waste is particularly egregious.
35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
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)
|