[rllib] [experimental] Decentralized Distributed PPO for torch (DD-PPO) (#6918)

rllib/optimizers/torch_distributed_data_parallel_optimizer.py

@@ -0,0 +1,103 @@
+import logging
+
+import ray
+from ray.rllib.optimizers.policy_optimizer import PolicyOptimizer
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.timer import TimerStat
+
+logger = logging.getLogger(__name__)
+
+
+class TorchDistributedDataParallelOptimizer(PolicyOptimizer):
+    """EXPERIMENTAL: torch distributed multi-node SGD."""
+
+    def __init__(self,
+                 workers,
+                 num_sgd_iter=1,
+                 train_batch_size=1,
+                 sgd_minibatch_size=0,
+                 standardize_fields=frozenset([]),
+                 keep_local_weights_in_sync=True,
+                 backend="gloo"):
+        PolicyOptimizer.__init__(self, workers)
+        self.learner_stats = {}
+        self.num_sgd_iter = num_sgd_iter
+        self.train_batch_size = train_batch_size
+        self.sgd_minibatch_size = sgd_minibatch_size
+        self.standardize_fields = standardize_fields
+        self.keep_local_weights_in_sync = keep_local_weights_in_sync
+        self.update_weights_timer = TimerStat()
+        self.learn_timer = TimerStat()
+
+        # Setup the distributed processes.
+        if not self.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)
+            for i, worker in enumerate(workers.remote_workers())
+        ])
+        logger.info("Torch process group init completed")
+
+    @override(PolicyOptimizer)
+    def step(self):
+        # 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 self.keep_local_weights_in_sync:
+            with self.update_weights_timer:
+                weights = ray.put(self.workers.local_worker().get_weights())
+                for e in self.workers.remote_workers():
+                    e.set_weights.remote(weights)
+
+        with self.learn_timer:
+            results = ray.get([
+                w.sample_and_learn.remote(
+                    self.train_batch_size, self.num_sgd_iter,
+                    self.sgd_minibatch_size, self.standardize_fields)
+                for w in self.workers.remote_workers()
+            ])
+        for info, count in results:
+            self.num_steps_sampled += count
+            self.num_steps_trained += count
+        self.learner_stats = results[0][0]
+
+        # In debug mode, check the allreduce successfully synced the weights.
+        if logger.isEnabledFor(logging.DEBUG):
+            weights = ray.get([
+                w.get_weights.remote() for w in self.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
+
+        # Sync down the weights. As with the sync up, this is not really
+        # needed unless the user is reading the local weights.
+        if self.keep_local_weights_in_sync:
+            self.workers.local_worker().set_weights(
+                ray.get(self.workers.remote_workers()[0].get_weights.remote()))
+
+        return self.learner_stats
+
+    @override(PolicyOptimizer)
+    def stats(self):
+        return dict(
+            PolicyOptimizer.stats(self), **{
+                "update_weights_time_ms": round(
+                    1000 * self.update_weights_timer.mean, 3),
+                "learn_time_ms": round(1000 * self.learn_timer.mean, 3),
+                "learner": self.learner_stats,
+            })