[rllib] Avoid sample wastage with bad PPO configurations (#3552)

## What do these changes do?

Previously we logged a warning if the PPO configuration would waste many samples. However, this didn't apply in the case of long episodes in `complete_episodes` batch mode, and also the amount of waste is up to 2x in common cases.

This pr:
- Estimates the number of sampling tasks needed to avoid over-sampling.
- Collects all sample results and never discards any. In principle this can degrade performance at large scale if certain machines are slower. Add a config flag to enable this legacy behavior.

## Related issue number

Closes: https://github.com/ray-project/ray/issues/3549

python/ray/rllib/agents/ppo/ppo.py

@@ -53,6 +53,9 @@ DEFAULT_CONFIG = with_common_config({
     # Uses the sync samples optimizer instead of the multi-gpu one. This does
     # not support minibatches.
     "simple_optimizer": False,
+    # (Deprecated) Use the sampling behavior as of 0.6, which launches extra
+    # sampling tasks for performance but can waste a large portion of samples.
+    "straggler_mitigation": False,
 })
 # __sphinx_doc_end__
 # yapf: enable
@@ -84,8 +87,12 @@ class PPOAgent(Agent):
                     "sgd_batch_size": self.config["sgd_minibatch_size"],
                     "num_sgd_iter": self.config["num_sgd_iter"],
                     "num_gpus": self.config["num_gpus"],
+                    "sample_batch_size": self.config["sample_batch_size"],
+                    "num_envs_per_worker": self.config["num_envs_per_worker"],
                     "train_batch_size": self.config["train_batch_size"],
                     "standardize_fields": ["advantages"],
+                    "straggler_mitigation": (
+                        self.config["straggler_mitigation"]),
                 })
 
     @override(Agent)
@@ -108,17 +115,6 @@ class PPOAgent(Agent):
         return res
 
     def _validate_config(self):
-        waste_ratio = (
-            self.config["sample_batch_size"] * self.config["num_workers"] /
-            self.config["train_batch_size"])
-        if waste_ratio > 1:
-            msg = ("sample_batch_size * num_workers >> train_batch_size. "
-                   "This means that many steps will be discarded. Consider "
-                   "reducing sample_batch_size, or increase train_batch_size.")
-            if waste_ratio > 1.5:
-                raise ValueError(msg)
-            else:
-                logger.warning(msg)
         if self.config["sgd_minibatch_size"] > self.config["train_batch_size"]:
             raise ValueError(
                 "Minibatch size {} must be <= train batch size {}.".format(

python/ray/rllib/agents/ppo/rollout.py

@@ -1,34 +0,0 @@
-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)