[rllib] Execute PPO using training workflow (#8206)

* wip

* add kl

* kl

* works now

* doc update

* reorg

* add ddppo

* add stats

* fix fetch

* comment

* fix learner stat regression

* test fixes

* fix test

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)

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)

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

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)

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`.

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"

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

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.
 

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

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

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)