[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)

python/ray/rllib/optimizers/multi_gpu_optimizer.py

@@ -12,6 +12,8 @@ import ray
 from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph
 from ray.rllib.optimizers.policy_optimizer import PolicyOptimizer
 from ray.rllib.optimizers.multi_gpu_impl import LocalSyncParallelOptimizer
+from ray.rllib.optimizers.rollout import collect_samples, \
+    collect_samples_straggler_mitigation
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.timer import TimerStat
 from ray.rllib.evaluation.sample_batch import SampleBatch, DEFAULT_POLICY_ID, \
@@ -40,12 +42,18 @@ class LocalMultiGPUOptimizer(PolicyOptimizer):
     def _init(self,
               sgd_batch_size=128,
               num_sgd_iter=10,
+              sample_batch_size=200,
+              num_envs_per_worker=1,
               train_batch_size=1024,
               num_gpus=0,
-              standardize_fields=[]):
+              standardize_fields=[],
+              straggler_mitigation=False):
         self.batch_size = sgd_batch_size
         self.num_sgd_iter = num_sgd_iter
+        self.num_envs_per_worker = num_envs_per_worker
+        self.sample_batch_size = sample_batch_size
         self.train_batch_size = train_batch_size
+        self.straggler_mitigation = straggler_mitigation
         if not num_gpus:
             self.devices = ["/cpu:0"]
         else:
@@ -108,10 +116,20 @@ class LocalMultiGPUOptimizer(PolicyOptimizer):
 
         with self.sample_timer:
             if self.remote_evaluators:
-                # TODO(rliaw): remove when refactoring
-                from ray.rllib.agents.ppo.rollout import collect_samples
-                samples = collect_samples(self.remote_evaluators,
-                                          self.train_batch_size)
+                if self.straggler_mitigation:
+                    samples = collect_samples_straggler_mitigation(
+                        self.remote_evaluators, self.train_batch_size)
+                else:
+                    samples = collect_samples(
+                        self.remote_evaluators, self.sample_batch_size,
+                        self.num_envs_per_worker, self.train_batch_size)
+                if samples.count > self.train_batch_size * 2:
+                    logger.info(
+                        "Collected more training samples than expected "
+                        "(actual={}, train_batch_size={}). ".format(
+                            samples.count, self.train_batch_size) +
+                        "This may be because you have many workers or "
+                        "long episodes in 'complete_episodes' batch mode.")
             else:
                 samples = self.local_evaluator.sample()
             # Handle everything as if multiagent

python/ray/rllib/optimizers/rollout.py

@@ -0,0 +1,71 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+
+import ray
+from ray.rllib.evaluation.sample_batch import SampleBatch
+
+logger = logging.getLogger(__name__)
+
+
+def collect_samples(agents, sample_batch_size, num_envs_per_worker,
+                    train_batch_size):
+    """Collects at least train_batch_size samples, never discarding any."""
+
+    num_timesteps_so_far = 0
+    trajectories = []
+    agent_dict = {}
+
+    for agent in agents:
+        fut_sample = agent.sample.remote()
+        agent_dict[fut_sample] = agent
+
+    while agent_dict:
+        [fut_sample], _ = ray.wait(list(agent_dict))
+        agent = agent_dict.pop(fut_sample)
+        next_sample = ray.get(fut_sample)
+        assert next_sample.count >= sample_batch_size * num_envs_per_worker
+        num_timesteps_so_far += next_sample.count
+        trajectories.append(next_sample)
+
+        # Only launch more tasks if we don't already have enough pending
+        pending = len(agent_dict) * sample_batch_size * num_envs_per_worker
+        if num_timesteps_so_far + pending < train_batch_size:
+            fut_sample2 = agent.sample.remote()
+            agent_dict[fut_sample2] = agent
+
+    return SampleBatch.concat_samples(trajectories)
+
+
+def collect_samples_straggler_mitigation(agents, train_batch_size):
+    """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.
+    """
+
+    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(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)

python/ray/rllib/test/test_optimizers.py

@@ -9,6 +9,7 @@ import time
 import unittest
 
 import ray
+from ray.rllib.agents.ppo import PPOAgent
 from ray.rllib.agents.ppo.ppo_policy_graph import PPOPolicyGraph
 from ray.rllib.evaluation import SampleBatch
 from ray.rllib.evaluation.policy_evaluator import PolicyEvaluator
@@ -31,6 +32,64 @@ class AsyncOptimizerTest(unittest.TestCase):
         self.assertTrue(all(local.get_weights() == 0))
 
 
+class PPOCollectTest(unittest.TestCase):
+    def tearDown(self):
+        ray.shutdown()
+
+    def testPPOSampleWaste(self):
+        ray.init(num_cpus=4)
+
+        # Check we at least collect the initial wave of samples
+        ppo = PPOAgent(
+            env="CartPole-v0",
+            config={
+                "sample_batch_size": 200,
+                "train_batch_size": 128,
+                "num_workers": 3,
+            })
+        ppo.train()
+        self.assertEqual(ppo.optimizer.num_steps_sampled, 600)
+        ppo.stop()
+
+        # Check we collect at least the specified amount of samples
+        ppo = PPOAgent(
+            env="CartPole-v0",
+            config={
+                "sample_batch_size": 200,
+                "train_batch_size": 900,
+                "num_workers": 3,
+            })
+        ppo.train()
+        self.assertEqual(ppo.optimizer.num_steps_sampled, 1000)
+        ppo.stop()
+
+        # Check in vectorized mode
+        ppo = PPOAgent(
+            env="CartPole-v0",
+            config={
+                "sample_batch_size": 200,
+                "num_envs_per_worker": 2,
+                "train_batch_size": 900,
+                "num_workers": 3,
+            })
+        ppo.train()
+        self.assertEqual(ppo.optimizer.num_steps_sampled, 1200)
+        ppo.stop()
+
+        # Check legacy mode
+        ppo = PPOAgent(
+            env="CartPole-v0",
+            config={
+                "sample_batch_size": 200,
+                "train_batch_size": 128,
+                "num_workers": 3,
+                "straggler_mitigation": True,
+            })
+        ppo.train()
+        self.assertEqual(ppo.optimizer.num_steps_sampled, 200)
+        ppo.stop()
+
+
 class SampleBatchTest(unittest.TestCase):
     def testConcat(self):
         b1 = SampleBatch({"a": np.array([1, 2, 3]), "b": np.array([4, 5, 6])})