diff --git a/python/ray/rllib/agents/ppo/ppo.py b/python/ray/rllib/agents/ppo/ppo.py index 620ef5314..59eba8ace 100644 --- a/python/ray/rllib/agents/ppo/ppo.py +++ b/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( diff --git a/python/ray/rllib/agents/ppo/rollout.py b/python/ray/rllib/agents/ppo/rollout.py deleted file mode 100644 index 4084e9ba0..000000000 --- a/python/ray/rllib/agents/ppo/rollout.py +++ /dev/null @@ -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) diff --git a/python/ray/rllib/optimizers/multi_gpu_optimizer.py b/python/ray/rllib/optimizers/multi_gpu_optimizer.py index e6f2a427c..73c2416b1 100644 --- a/python/ray/rllib/optimizers/multi_gpu_optimizer.py +++ b/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 diff --git a/python/ray/rllib/optimizers/rollout.py b/python/ray/rllib/optimizers/rollout.py new file mode 100644 index 000000000..d13f7c5ef --- /dev/null +++ b/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) diff --git a/python/ray/rllib/test/test_optimizers.py b/python/ray/rllib/test/test_optimizers.py index ee08dcfdb..074e0c081 100644 --- a/python/ray/rllib/test/test_optimizers.py +++ b/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])})