diff --git a/rllib/agents/ppo/ddppo.py b/rllib/agents/ppo/ddppo.py index 1c5f0e01f..330941856 100644 --- a/rllib/agents/ppo/ddppo.py +++ b/rllib/agents/ppo/ddppo.py @@ -1,6 +1,3 @@ -from ray.rllib.agents.ppo import ppo -from ray.rllib.agents.trainer import with_base_config -from ray.rllib.optimizers import TorchDistributedDataParallelOptimizer """Decentralized Distributed PPO implementation. Unlike APPO or PPO, learning is no longer done centralized in the trainer @@ -17,6 +14,23 @@ Paper reference: https://arxiv.org/abs/1911.00357 Note that unlike the paper, we currently do not implement straggler mitigation. """ +import logging +import time + +import ray +from ray.util.iter import LocalIterator +from ray.rllib.agents.ppo import ppo +from ray.rllib.agents.trainer import with_base_config +from ray.rllib.optimizers import TorchDistributedDataParallelOptimizer +from ray.rllib.execution.rollout_ops import ParallelRollouts +from ray.rllib.execution.metric_ops import StandardMetricsReporting +from ray.rllib.execution.common import STEPS_SAMPLED_COUNTER, \ + STEPS_TRAINED_COUNTER, LEARNER_INFO, LEARN_ON_BATCH_TIMER +from ray.rllib.evaluation.rollout_worker import get_global_worker +from ray.rllib.utils.sgd import do_minibatch_sgd + +logger = logging.getLogger(__name__) + # yapf: disable # __sphinx_doc_begin__ DEFAULT_CONFIG = with_base_config(ppo.DEFAULT_CONFIG, { @@ -30,6 +44,9 @@ DEFAULT_CONFIG = with_base_config(ppo.DEFAULT_CONFIG, { "sgd_minibatch_size": 50, # Number of SGD epochs per optimization round. "num_sgd_iter": 10, + # Download weights between each training step. This adds a bit of overhead + # but allows the user to access the weights from the trainer. + "keep_local_weights_in_sync": True, # *** WARNING: configs below are DDPPO overrides over PPO; you # shouldn't need to adjust them. *** @@ -54,10 +71,6 @@ def validate_config(config): raise ValueError( "Set rollout_fragment_length instead of train_batch_size " "for DDPPO.") - ppo.validate_config(config) - - -def make_distributed_allreduce_optimizer(workers, config): if not config["use_pytorch"]: raise ValueError( "Distributed data parallel is only supported for PyTorch") @@ -71,7 +84,10 @@ def make_distributed_allreduce_optimizer(workers, config): raise ValueError( "Distributed data parallel requires truncate_episodes " "batch mode.") + ppo.validate_config(config) + +def make_distributed_allreduce_optimizer(workers, config): return TorchDistributedDataParallelOptimizer( workers, expected_batch_size=config["rollout_fragment_length"] * @@ -81,8 +97,107 @@ def make_distributed_allreduce_optimizer(workers, config): standardize_fields=["advantages"]) +# Experimental distributed execution impl; enable with "use_exec_api": True. +def execution_plan(workers, config): + rollouts = ParallelRollouts(workers, mode="raw") + + # Sync up the weights. In principle we don't need this, but it doesn't + # add too much overhead and handles the case where the user manually + # updates the local weights. + if config["keep_local_weights_in_sync"]: + weights = ray.put(workers.local_worker().get_weights()) + for e in workers.remote_workers(): + e.set_weights.remote(weights) + + # Setup the distributed processes. + if not workers.remote_workers(): + raise ValueError("This optimizer requires >0 remote workers.") + ip = ray.get(workers.remote_workers()[0].get_node_ip.remote()) + port = ray.get(workers.remote_workers()[0].find_free_port.remote()) + address = "tcp://{ip}:{port}".format(ip=ip, port=port) + logger.info("Creating torch process group with leader {}".format(address)) + + # Get setup tasks in order to throw errors on failure. + ray.get([ + worker.setup_torch_data_parallel.remote( + address, i, len(workers.remote_workers()), backend="gloo") + for i, worker in enumerate(workers.remote_workers()) + ]) + logger.info("Torch process group init completed") + + # This function is applied remotely on each rollout worker. + def train_torch_distributed_allreduce(batch): + expected_batch_size = ( + config["rollout_fragment_length"] * config["num_envs_per_worker"]) + this_worker = get_global_worker() + assert batch.count == expected_batch_size, \ + ("Batch size possibly out of sync between workers, expected:", + expected_batch_size, "got:", batch.count) + logger.info("Executing distributed minibatch SGD " + "with epoch size {}, minibatch size {}".format( + batch.count, config["sgd_minibatch_size"])) + info = do_minibatch_sgd(batch, this_worker.policy_map, this_worker, + config["num_sgd_iter"], + config["sgd_minibatch_size"], ["advantages"]) + return info, batch.count + + # Have to manually record stats since we are using "raw" rollouts mode. + class RecordStats: + def _on_fetch_start(self): + self.fetch_start_time = time.perf_counter() + + def __call__(self, items): + for item in items: + info, count = item + metrics = LocalIterator.get_metrics() + metrics.counters[STEPS_SAMPLED_COUNTER] += count + metrics.counters[STEPS_TRAINED_COUNTER] += count + metrics.info[LEARNER_INFO] = info + # Since SGD happens remotely, the time delay between fetch and + # completion is approximately the SGD step time. + metrics.timers[LEARN_ON_BATCH_TIMER].push(time.perf_counter() - + self.fetch_start_time) + + train_op = ( + rollouts.for_each(train_torch_distributed_allreduce) # allreduce + .batch_across_shards() # List[(grad_info, count)] + .for_each(RecordStats())) + + # Sync down the weights. As with the sync up, this is not really + # needed unless the user is reading the local weights. + if config["keep_local_weights_in_sync"]: + + def download_weights(item): + workers.local_worker().set_weights( + ray.get(workers.remote_workers()[0].get_weights.remote())) + return item + + train_op = train_op.for_each(download_weights) + + # In debug mode, check the allreduce successfully synced the weights. + if logger.isEnabledFor(logging.DEBUG): + + def check_sync(item): + weights = ray.get( + [w.get_weights.remote() for w in workers.remote_workers()]) + sums = [] + for w in weights: + acc = 0 + for p in w.values(): + for k, v in p.items(): + acc += v.sum() + sums.append(float(acc)) + logger.debug("The worker weight sums are {}".format(sums)) + assert len(set(sums)) == 1, sums + + train_op = train_op.for_each(check_sync) + + return StandardMetricsReporting(train_op, workers, config) + + DDPPOTrainer = ppo.PPOTrainer.with_updates( name="DDPPO", default_config=DEFAULT_CONFIG, make_policy_optimizer=make_distributed_allreduce_optimizer, + execution_plan=execution_plan, validate_config=validate_config) diff --git a/rllib/agents/ppo/ppo.py b/rllib/agents/ppo/ppo.py index 0a60ee416..52cdafe26 100644 --- a/rllib/agents/ppo/ppo.py +++ b/rllib/agents/ppo/ppo.py @@ -3,6 +3,10 @@ import logging from ray.rllib.agents import with_common_config from ray.rllib.agents.ppo.ppo_tf_policy import PPOTFPolicy from ray.rllib.agents.trainer_template import build_trainer +from ray.rllib.execution.rollout_ops import ParallelRollouts, ConcatBatches, \ + StandardizeFields, SelectExperiences +from ray.rllib.execution.train_ops import TrainOneStep, TrainTFMultiGPU +from ray.rllib.execution.metric_ops import StandardMetricsReporting from ray.rllib.optimizers import SyncSamplesOptimizer, LocalMultiGPUOptimizer from ray.rllib.utils import try_import_tf @@ -71,7 +75,9 @@ DEFAULT_CONFIG = with_common_config({ # Set this to True for debugging on non-GPU machines (set `num_gpus` > 0). "_fake_gpus": False, # Use PyTorch as framework? - "use_pytorch": False + "use_pytorch": False, + # Use the execution plan API instead of policy optimizers. + "use_exec_api": True, }) # __sphinx_doc_end__ # yapf: enable @@ -112,22 +118,26 @@ def update_kl(trainer, fetches): if pi_id in fetches: pi.update_kl(fetches[pi_id]["kl"]) else: - logger.debug("No data for {}, not updating kl".format(pi_id)) + logger.info("No data for {}, not updating kl".format(pi_id)) trainer.workers.local_worker().foreach_trainable_policy(update) def warn_about_bad_reward_scales(trainer, result): + return _warn_about_bad_reward_scales(trainer.config, result) + + +def _warn_about_bad_reward_scales(config, result): if result["policy_reward_mean"]: - return # Punt on handling multiagent case. + return result # Punt on handling multiagent case. # Warn about excessively high VF loss. learner_stats = result["info"]["learner"] if "default_policy" in learner_stats: - scaled_vf_loss = (trainer.config["vf_loss_coeff"] * + scaled_vf_loss = (config["vf_loss_coeff"] * learner_stats["default_policy"]["vf_loss"]) policy_loss = learner_stats["default_policy"]["policy_loss"] - if trainer.config["vf_share_layers"] and scaled_vf_loss > 100: + if config["vf_share_layers"] and scaled_vf_loss > 100: logger.warning( "The magnitude of your value function loss is extremely large " "({}) compared to the policy loss ({}). This can prevent the " @@ -136,12 +146,11 @@ def warn_about_bad_reward_scales(trainer, result): scaled_vf_loss, policy_loss)) # Warn about bad clipping configs - if trainer.config["vf_clip_param"] <= 0: + if config["vf_clip_param"] <= 0: rew_scale = float("inf") else: rew_scale = round( - abs(result["episode_reward_mean"]) / - trainer.config["vf_clip_param"], 0) + abs(result["episode_reward_mean"]) / config["vf_clip_param"], 0) if rew_scale > 200: logger.warning( "The magnitude of your environment rewards are more than " @@ -151,6 +160,8 @@ def warn_about_bad_reward_scales(trainer, result): "function to converge. If this is not intended, consider " "increasing `vf_clip_param`.") + return result + def validate_config(config): if config["entropy_coeff"] < 0: @@ -186,12 +197,62 @@ def get_policy_class(config): return PPOTFPolicy +# Experimental distributed execution impl; enable with "use_exec_api": True. +def execution_plan(workers, config): + rollouts = ParallelRollouts(workers, mode="bulk_sync") + + # Collect large batches of relevant experiences & standardize. + rollouts = rollouts.for_each( + SelectExperiences(workers.trainable_policies())) + rollouts = rollouts.combine( + ConcatBatches(min_batch_size=config["train_batch_size"])) + rollouts = rollouts.for_each(StandardizeFields(["advantages"])) + + if config["simple_optimizer"]: + train_op = rollouts.for_each( + TrainOneStep( + workers, + num_sgd_iter=config["num_sgd_iter"], + sgd_minibatch_size=config["sgd_minibatch_size"])) + else: + train_op = rollouts.for_each( + TrainTFMultiGPU( + workers, + sgd_minibatch_size=config["sgd_minibatch_size"], + num_sgd_iter=config["num_sgd_iter"], + num_gpus=config["num_gpus"], + rollout_fragment_length=config["rollout_fragment_length"], + num_envs_per_worker=config["num_envs_per_worker"], + train_batch_size=config["train_batch_size"], + shuffle_sequences=config["shuffle_sequences"], + _fake_gpus=config["_fake_gpus"])) + + # Callback to update the KL based on optimization info. + def update_kl(item): + _, fetches = item + + def update(pi, pi_id): + if pi_id in fetches: + pi.update_kl(fetches[pi_id]["kl"]) + else: + logger.warning("No data for {}, not updating kl".format(pi_id)) + + workers.local_worker().foreach_trainable_policy(update) + + # Update KL after each round of training. + train_op = train_op.for_each(update_kl) + + return StandardMetricsReporting(train_op, workers, config) \ + .for_each(lambda result: _warn_about_bad_reward_scales(config, result)) + + PPOTrainer = build_trainer( name="PPO", default_config=DEFAULT_CONFIG, default_policy=PPOTFPolicy, get_policy_class=get_policy_class, make_policy_optimizer=choose_policy_optimizer, + execution_plan=execution_plan, validate_config=validate_config, after_optimizer_step=update_kl, after_train_result=warn_about_bad_reward_scales) diff --git a/rllib/agents/trainer.py b/rllib/agents/trainer.py index 715c95b13..e58c739d2 100644 --- a/rllib/agents/trainer.py +++ b/rllib/agents/trainer.py @@ -417,7 +417,7 @@ class Trainer(Trainable): logger.info("Executing eagerly, with eager_tracing={}".format( "True" if config.get("eager_tracing") else "False")) - if tf and not tf.executing_eagerly(): + if tf and not tf.executing_eagerly() and not config.get("use_pytorch"): logger.info("Tip: set 'eager': true or the --eager flag to enable " "TensorFlow eager execution") diff --git a/rllib/agents/trainer_template.py b/rllib/agents/trainer_template.py index 70b791444..a6ff2291b 100644 --- a/rllib/agents/trainer_template.py +++ b/rllib/agents/trainer_template.py @@ -120,10 +120,6 @@ def build_trainer(name, self.execution_plan = execution_plan if use_exec_api: - logger.warning( - "The experimental distributed execution API is enabled " - "for this algorithm. Disable this by setting " - "'use_exec_api': False.") self.train_exec_impl = execution_plan(self.workers, config) elif make_policy_optimizer: self.optimizer = make_policy_optimizer(self.workers, config) diff --git a/rllib/evaluation/worker_set.py b/rllib/evaluation/worker_set.py index 3acabf9aa..1da49b93f 100644 --- a/rllib/evaluation/worker_set.py +++ b/rllib/evaluation/worker_set.py @@ -155,6 +155,11 @@ class WorkerSet: remote_results.extend(res) return local_results + remote_results + @DeveloperAPI + def trainable_policies(self): + """Return the list of trainable policy ids.""" + return self.local_worker().foreach_trainable_policy(lambda _, pid: pid) + @DeveloperAPI def foreach_trainable_policy(self, func): """Apply `func` to all workers' Policies iff in `policies_to_train`. diff --git a/rllib/execution/common.py b/rllib/execution/common.py index 7f865b5be..c0e6ed2cd 100644 --- a/rllib/execution/common.py +++ b/rllib/execution/common.py @@ -18,6 +18,7 @@ WORKER_UPDATE_TIMER = "update" GRAD_WAIT_TIMER = "grad_wait" SAMPLE_TIMER = "sample" LEARN_ON_BATCH_TIMER = "learn" +LOAD_BATCH_TIMER = "load" # Instant metrics (keys for metrics.info). LEARNER_INFO = "learner" diff --git a/rllib/execution/rollout_ops.py b/rllib/execution/rollout_ops.py index 9b9d7f09f..23300d69b 100644 --- a/rllib/execution/rollout_ops.py +++ b/rllib/execution/rollout_ops.py @@ -1,3 +1,4 @@ +import logging from typing import List, Tuple import time @@ -9,7 +10,11 @@ from ray.rllib.evaluation.worker_set import WorkerSet from ray.rllib.execution.common import GradientType, SampleBatchType, \ STEPS_SAMPLED_COUNTER, LEARNER_INFO, SAMPLE_TIMER, \ GRAD_WAIT_TIMER, _check_sample_batch_type -from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.sample_batch import SampleBatch, DEFAULT_POLICY_ID, \ + MultiAgentBatch +from ray.rllib.utils.sgd import standardized + +logger = logging.getLogger(__name__) def ParallelRollouts(workers: WorkerSet, @@ -23,11 +28,14 @@ def ParallelRollouts(workers: WorkerSet, Arguments: workers (WorkerSet): set of rollout workers to use. - mode (str): One of {'async', 'bulk_sync'}. + mode (str): One of {'async', 'bulk_sync', 'raw'}. - In 'async' mode, batches are returned as soon as they are computed by rollout workers with no order guarantees. - In 'bulk_sync' mode, we collect one batch from each worker and concatenate them together into a large batch to return. + - In 'raw' mode, the ParallelIterator object is returned directly + and the caller is responsible for implementing gather and + updating the timesteps counter. async_queue_depth (int): In async mode, the max number of async requests in flight per actor. @@ -76,9 +84,11 @@ def ParallelRollouts(workers: WorkerSet, elif mode == "async": return rollouts.gather_async( async_queue_depth=async_queue_depth).for_each(report_timesteps) + elif mode == "raw": + return rollouts else: - raise ValueError( - "mode must be one of 'bulk_sync', 'async', got '{}'".format(mode)) + raise ValueError("mode must be one of 'bulk_sync', 'async', 'raw', " + "got '{}'".format(mode)) def AsyncGradients( @@ -153,6 +163,12 @@ class ConcatBatches: self.buffer.append(batch) self.count += batch.count if self.count >= self.min_batch_size: + if self.count > self.min_batch_size * 2: + logger.info("Collected more training samples than expected " + "(actual={}, expected={}). ".format( + self.count, self.min_batch_size) + + "This may be because you have many workers or " + "long episodes in 'complete_episodes' batch mode.") out = SampleBatch.concat_samples(self.buffer) timer = LocalIterator.get_metrics().timers[SAMPLE_TIMER] timer.push(time.perf_counter() - self.batch_start_time) @@ -162,3 +178,68 @@ class ConcatBatches: self.count = 0 return [out] return [] + + +class SelectExperiences: + """Callable used to select experiences from a MultiAgentBatch. + + This should be used with the .for_each() operator. + + Examples: + >>> rollouts = ParallelRollouts(...) + >>> rollouts = rollouts.for_each(SelectExperiences(["pol1", "pol2"])) + >>> print(next(rollouts).policy_batches.keys()) + {"pol1", "pol2"} + """ + + def __init__(self, policy_ids: List[str]): + self.policy_ids = policy_ids + + def __call__(self, samples: SampleBatchType) -> SampleBatchType: + _check_sample_batch_type(samples) + + if isinstance(samples, MultiAgentBatch): + samples = MultiAgentBatch({ + k: v + for k, v in samples.policy_batches.items() + if k in self.policy_ids + }, samples.count) + + return samples + + +class StandardizeFields: + """Callable used to standardize fields of batches. + + This should be used with the .for_each() operator. Note that the input + may be mutated by this operator for efficiency. + + Examples: + >>> rollouts = ParallelRollouts(...) + >>> rollouts = rollouts.for_each(StandardizeFields(["advantages"])) + >>> print(np.std(next(rollouts)["advantages"])) + 1.0 + """ + + def __init__(self, fields: List[str]): + self.fields = fields + + def __call__(self, samples: SampleBatchType) -> SampleBatchType: + _check_sample_batch_type(samples) + wrapped = False + + if isinstance(samples, SampleBatch): + samples = MultiAgentBatch({ + DEFAULT_POLICY_ID: samples + }, samples.count) + wrapped = True + + for policy_id in samples.policy_batches: + batch = samples.policy_batches[policy_id] + for field in self.fields: + batch[field] = standardized(batch[field]) + + if wrapped: + samples = samples.policy_batches[DEFAULT_POLICY_ID] + + return samples diff --git a/rllib/execution/train_ops.py b/rllib/execution/train_ops.py index 6efa499c0..ebd76b3ad 100644 --- a/rllib/execution/train_ops.py +++ b/rllib/execution/train_ops.py @@ -1,15 +1,25 @@ +from collections import defaultdict import logging +import numpy as np +import math from typing import List import ray from ray.util.iter import LocalIterator -from ray.rllib.evaluation.metrics import get_learner_stats +from ray.rllib.evaluation.metrics import get_learner_stats, LEARNER_STATS_KEY from ray.rllib.evaluation.worker_set import WorkerSet from ray.rllib.execution.common import SampleBatchType, \ STEPS_SAMPLED_COUNTER, STEPS_TRAINED_COUNTER, LEARNER_INFO, \ APPLY_GRADS_TIMER, COMPUTE_GRADS_TIMER, WORKER_UPDATE_TIMER, \ - LEARN_ON_BATCH_TIMER, LAST_TARGET_UPDATE_TS, NUM_TARGET_UPDATES, \ - _get_global_vars, _check_sample_batch_type + LEARN_ON_BATCH_TIMER, LOAD_BATCH_TIMER, LAST_TARGET_UPDATE_TS, \ + NUM_TARGET_UPDATES, _get_global_vars, _check_sample_batch_type +from ray.rllib.optimizers.multi_gpu_impl import LocalSyncParallelOptimizer +from ray.rllib.policy.sample_batch import SampleBatch, DEFAULT_POLICY_ID, \ + MultiAgentBatch +from ray.rllib.utils import try_import_tf +from ray.rllib.utils.sgd import do_minibatch_sgd, averaged + +tf = try_import_tf() logger = logging.getLogger(__name__) @@ -30,8 +40,15 @@ class TrainOneStep: local iterator context. """ - def __init__(self, workers: WorkerSet): + def __init__(self, + workers: WorkerSet, + num_sgd_iter: int = 1, + sgd_minibatch_size: int = 0): self.workers = workers + self.policies = dict(self.workers.local_worker() + .foreach_trainable_policy(lambda p, i: (i, p))) + self.num_sgd_iter = num_sgd_iter + self.sgd_minibatch_size = sgd_minibatch_size def __call__(self, batch: SampleBatchType) -> (SampleBatchType, List[dict]): @@ -39,10 +56,18 @@ class TrainOneStep: metrics = LocalIterator.get_metrics() learn_timer = metrics.timers[LEARN_ON_BATCH_TIMER] with learn_timer: - info = self.workers.local_worker().learn_on_batch(batch) + if self.num_sgd_iter > 1 or self.sgd_minibatch_size > 0: + info = do_minibatch_sgd(batch, self.policies, + self.workers.local_worker(), + self.num_sgd_iter, + self.sgd_minibatch_size, []) + # TODO(ekl) shouldn't be returning learner stats directly here + metrics.info[LEARNER_INFO] = info + else: + info = self.workers.local_worker().learn_on_batch(batch) + metrics.info[LEARNER_INFO] = get_learner_stats(info) learn_timer.push_units_processed(batch.count) metrics.counters[STEPS_TRAINED_COUNTER] += batch.count - metrics.info[LEARNER_INFO] = get_learner_stats(info) if self.workers.remote_workers(): with metrics.timers[WORKER_UPDATE_TIMER]: weights = ray.put(self.workers.local_worker().get_weights()) @@ -53,6 +78,149 @@ class TrainOneStep: return batch, info +class TrainTFMultiGPU: + """TF Multi-GPU version of TrainOneStep. + + This should be used with the .for_each() operator. A tuple of the input + and learner stats will be returned. + + Examples: + >>> rollouts = ParallelRollouts(...) + >>> train_op = rollouts.for_each(TrainMultiGPU(workers, ...)) + >>> print(next(train_op)) # This trains the policy on one batch. + SampleBatch(...), {"learner_stats": ...} + + Updates the STEPS_TRAINED_COUNTER counter and LEARNER_INFO field in the + local iterator context. + """ + + def __init__(self, + workers: WorkerSet, + sgd_minibatch_size: int, + num_sgd_iter: int, + num_gpus: int, + rollout_fragment_length: int, + num_envs_per_worker: int, + train_batch_size: int, + shuffle_sequences: bool, + _fake_gpus: bool = False): + self.workers = workers + self.policies = dict(self.workers.local_worker() + .foreach_trainable_policy(lambda p, i: (i, p))) + self.num_sgd_iter = num_sgd_iter + self.sgd_minibatch_size = sgd_minibatch_size + self.shuffle_sequences = shuffle_sequences + + # Collect actual devices to use. + if not num_gpus: + _fake_gpus = True + num_gpus = 1 + type_ = "cpu" if _fake_gpus else "gpu" + self.devices = [ + "/{}:{}".format(type_, i) for i in range(int(math.ceil(num_gpus))) + ] + + self.batch_size = int(sgd_minibatch_size / len(self.devices)) * len( + self.devices) + assert self.batch_size % len(self.devices) == 0 + assert self.batch_size >= len(self.devices), "batch size too small" + self.per_device_batch_size = int(self.batch_size / len(self.devices)) + + # per-GPU graph copies created below must share vars with the policy + # reuse is set to AUTO_REUSE because Adam nodes are created after + # all of the device copies are created. + self.optimizers = {} + with self.workers.local_worker().tf_sess.graph.as_default(): + with self.workers.local_worker().tf_sess.as_default(): + for policy_id, policy in self.policies.items(): + with tf.variable_scope(policy_id, reuse=tf.AUTO_REUSE): + if policy._state_inputs: + rnn_inputs = policy._state_inputs + [ + policy._seq_lens + ] + else: + rnn_inputs = [] + self.optimizers[policy_id] = ( + LocalSyncParallelOptimizer( + policy._optimizer, self.devices, + [v + for _, v in policy._loss_inputs], rnn_inputs, + self.per_device_batch_size, policy.copy)) + + self.sess = self.workers.local_worker().tf_sess + self.sess.run(tf.global_variables_initializer()) + + def __call__(self, + samples: SampleBatchType) -> (SampleBatchType, List[dict]): + _check_sample_batch_type(samples) + + # Handle everything as if multiagent + if isinstance(samples, SampleBatch): + samples = MultiAgentBatch({ + DEFAULT_POLICY_ID: samples + }, samples.count) + + metrics = LocalIterator.get_metrics() + load_timer = metrics.timers[LOAD_BATCH_TIMER] + learn_timer = metrics.timers[LEARN_ON_BATCH_TIMER] + with load_timer: + # (1) Load data into GPUs. + num_loaded_tuples = {} + for policy_id, batch in samples.policy_batches.items(): + if policy_id not in self.policies: + continue + + policy = self.policies[policy_id] + policy._debug_vars() + tuples = policy._get_loss_inputs_dict( + batch, shuffle=self.shuffle_sequences) + data_keys = [ph for _, ph in policy._loss_inputs] + if policy._state_inputs: + state_keys = policy._state_inputs + [policy._seq_lens] + else: + state_keys = [] + num_loaded_tuples[policy_id] = ( + self.optimizers[policy_id].load_data( + self.sess, [tuples[k] for k in data_keys], + [tuples[k] for k in state_keys])) + + with learn_timer: + # (2) Execute minibatch SGD on loaded data. + fetches = {} + for policy_id, tuples_per_device in num_loaded_tuples.items(): + optimizer = self.optimizers[policy_id] + num_batches = max( + 1, + int(tuples_per_device) // int(self.per_device_batch_size)) + logger.debug("== sgd epochs for {} ==".format(policy_id)) + for i in range(self.num_sgd_iter): + iter_extra_fetches = defaultdict(list) + permutation = np.random.permutation(num_batches) + for batch_index in range(num_batches): + batch_fetches = optimizer.optimize( + self.sess, permutation[batch_index] * + self.per_device_batch_size) + for k, v in batch_fetches[LEARNER_STATS_KEY].items(): + iter_extra_fetches[k].append(v) + logger.debug("{} {}".format(i, + averaged(iter_extra_fetches))) + fetches[policy_id] = averaged(iter_extra_fetches) + + load_timer.push_units_processed(samples.count) + learn_timer.push_units_processed(samples.count) + + metrics.counters[STEPS_TRAINED_COUNTER] += samples.count + metrics.info[LEARNER_INFO] = fetches + if self.workers.remote_workers(): + with metrics.timers[WORKER_UPDATE_TIMER]: + weights = ray.put(self.workers.local_worker().get_weights()) + for e in self.workers.remote_workers(): + e.set_weights.remote(weights, _get_global_vars()) + # Also update global vars of the local worker. + self.workers.local_worker().set_global_vars(_get_global_vars()) + return samples, fetches + + class ComputeGradients: """Callable that computes gradients with respect to the policy loss. diff --git a/rllib/models/modelv2.py b/rllib/models/modelv2.py index 0a8374498..ea77986b9 100644 --- a/rllib/models/modelv2.py +++ b/rllib/models/modelv2.py @@ -187,7 +187,7 @@ class ModelV2: except AttributeError: raise ValueError("Output is not a tensor: {}".format(outputs)) else: - if len(shape) != 2 or shape[1] != self.num_outputs: + if len(shape) != 2 or int(shape[1]) != self.num_outputs: raise ValueError( "Expected output shape of [None, {}], got {}".format( self.num_outputs, shape)) diff --git a/rllib/optimizers/tests/test_optimizers.py b/rllib/optimizers/tests/test_optimizers.py index e274009ed..ebf431e45 100644 --- a/rllib/optimizers/tests/test_optimizers.py +++ b/rllib/optimizers/tests/test_optimizers.py @@ -72,8 +72,8 @@ class PPOCollectTest(unittest.TestCase): "train_batch_size": 128, "num_workers": 3, }) - ppo.train() - self.assertEqual(ppo.optimizer.num_steps_sampled, 600) + result = ppo.train() + self.assertEqual(result["info"]["num_steps_sampled"], 600) ppo.stop() # Check we collect at least the specified amount of samples @@ -84,8 +84,8 @@ class PPOCollectTest(unittest.TestCase): "train_batch_size": 900, "num_workers": 3, }) - ppo.train() - self.assertEqual(ppo.optimizer.num_steps_sampled, 1000) + result = ppo.train() + self.assertEqual(result["info"]["num_steps_sampled"], 1200) ppo.stop() # Check in vectorized mode @@ -97,8 +97,8 @@ class PPOCollectTest(unittest.TestCase): "train_batch_size": 900, "num_workers": 3, }) - ppo.train() - self.assertEqual(ppo.optimizer.num_steps_sampled, 1200) + result = ppo.train() + self.assertEqual(result["info"]["num_steps_sampled"], 1200) ppo.stop() diff --git a/rllib/tests/test_execution.py b/rllib/tests/test_execution.py index 9e6a1dfeb..efca91954 100644 --- a/rllib/tests/test_execution.py +++ b/rllib/tests/test_execution.py @@ -1,3 +1,4 @@ +import numpy as np import pytest import time import gym @@ -11,7 +12,7 @@ from ray.rllib.execution.concurrency_ops import Concurrently, Enqueue, Dequeue from ray.rllib.execution.metric_ops import StandardMetricsReporting from ray.rllib.execution.replay_ops import StoreToReplayBuffer, Replay from ray.rllib.execution.rollout_ops import ParallelRollouts, AsyncGradients, \ - ConcatBatches + ConcatBatches, StandardizeFields from ray.rllib.execution.train_ops import TrainOneStep, ComputeGradients, \ AverageGradients from ray.rllib.optimizers.async_replay_optimizer import LocalReplayBuffer, \ @@ -132,6 +133,15 @@ def test_concat_batches(ray_start_regular_shared): assert "sample" in timers +def test_standardize(ray_start_regular_shared): + workers = make_workers(0) + a = ParallelRollouts(workers, mode="async") + b = a.for_each(StandardizeFields(["t"])) + batch = next(b) + assert abs(np.mean(batch["t"])) < 0.001, batch + assert abs(np.std(batch["t"]) - 1.0) < 0.001, batch + + def test_async_grads(ray_start_regular_shared): workers = make_workers(2) a = AsyncGradients(workers)