[rllib] Add Decentralized DDPPO trainer and documentation (#7088)

This commit is contained in:
Eric Liang
2020-02-10 15:28:27 -08:00
committed by GitHub
parent 6e1c3ea824
commit 026f6884b5
15 changed files with 212 additions and 72 deletions
@@ -13,8 +13,8 @@ class TorchDistributedDataParallelOptimizer(PolicyOptimizer):
def __init__(self,
workers,
expected_batch_size,
num_sgd_iter=1,
train_batch_size=1,
sgd_minibatch_size=0,
standardize_fields=frozenset([]),
keep_local_weights_in_sync=True,
@@ -22,11 +22,12 @@ class TorchDistributedDataParallelOptimizer(PolicyOptimizer):
PolicyOptimizer.__init__(self, workers)
self.learner_stats = {}
self.num_sgd_iter = num_sgd_iter
self.train_batch_size = train_batch_size
self.expected_batch_size = expected_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.sync_down_timer = TimerStat()
self.sync_up_timer = TimerStat()
self.learn_timer = TimerStat()
# Setup the distributed processes.
@@ -52,7 +53,7 @@ class TorchDistributedDataParallelOptimizer(PolicyOptimizer):
# 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:
with self.sync_up_timer:
weights = ray.put(self.workers.local_worker().get_weights())
for e in self.workers.remote_workers():
e.set_weights.remote(weights)
@@ -60,7 +61,7 @@ class TorchDistributedDataParallelOptimizer(PolicyOptimizer):
with self.learn_timer:
results = ray.get([
w.sample_and_learn.remote(
self.train_batch_size, self.num_sgd_iter,
self.expected_batch_size, self.num_sgd_iter,
self.sgd_minibatch_size, self.standardize_fields)
for w in self.workers.remote_workers()
])
@@ -87,8 +88,10 @@ class TorchDistributedDataParallelOptimizer(PolicyOptimizer):
# 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()))
with self.sync_down_timer:
self.workers.local_worker().set_weights(
ray.get(
self.workers.remote_workers()[0].get_weights.remote()))
return self.learner_stats
@@ -96,8 +99,10 @@ class TorchDistributedDataParallelOptimizer(PolicyOptimizer):
def stats(self):
return dict(
PolicyOptimizer.stats(self), **{
"update_weights_time_ms": round(
1000 * self.update_weights_timer.mean, 3),
"sync_weights_up_time": round(1000 * self.sync_up_timer.mean,
3),
"sync_weights_down_time": round(
1000 * self.sync_down_timer.mean, 3),
"learn_time_ms": round(1000 * self.learn_timer.mean, 3),
"learner": self.learner_stats,
})