From ea51e94729724d66964d1766da7c5b633d189a6f Mon Sep 17 00:00:00 2001
From: Chua Cheow Huan <17569306+ChuaCheowHuan@users.noreply.github.com>
Date: Sat, 15 Aug 2020 15:51:45 +0800
Subject: [PATCH] [rllib] Learning rate schedule for DDPPO. (#10006)

* Get shared metrics, increment counter & set global vars for remote workers.

* Add unit test to test lr_schedule for DDPPO.

* Broadcast the local set of global vars to remote workers instead of independently setting the global vars on each rollout worker.
---
 rllib/agents/ppo/ddppo.py            | 11 ++++++++++-
 rllib/agents/ppo/tests/test_ddppo.py | 15 +++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/rllib/agents/ppo/ddppo.py b/rllib/agents/ppo/ddppo.py
index eed80b6d2..cb16b9b3b 100644
--- a/rllib/agents/ppo/ddppo.py
+++ b/rllib/agents/ppo/ddppo.py
@@ -24,7 +24,7 @@ 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, \
-    _get_shared_metrics
+    _get_shared_metrics, _get_global_vars
 from ray.rllib.evaluation.rollout_worker import get_global_worker
 from ray.rllib.utils.sgd import do_minibatch_sgd
 
@@ -121,6 +121,13 @@ def execution_plan(workers, config):
                                 config["sgd_minibatch_size"], ["advantages"])
         return info, batch.count
 
+    # Broadcast the local set of global vars.
+    def update_worker_global_vars(item):
+        global_vars = _get_global_vars()
+        for w in workers.remote_workers():
+            w.set_global_vars.remote(global_vars)
+        return item
+
     # Have to manually record stats since we are using "raw" rollouts mode.
     class RecordStats:
         def _on_fetch_start(self):
@@ -143,6 +150,8 @@ def execution_plan(workers, config):
         .batch_across_shards()  # List[(grad_info, count)]
         .for_each(RecordStats()))
 
+    train_op = train_op.for_each(update_worker_global_vars)
+
     # 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"]:
diff --git a/rllib/agents/ppo/tests/test_ddppo.py b/rllib/agents/ppo/tests/test_ddppo.py
index 4c56a22ee..5b3d659e4 100644
--- a/rllib/agents/ppo/tests/test_ddppo.py
+++ b/rllib/agents/ppo/tests/test_ddppo.py
@@ -28,6 +28,21 @@ class TestDDPPO(unittest.TestCase):
             check_compute_single_action(trainer)
             trainer.stop()
 
+    def test_ddppo_schedule(self):
+        """Test whether lr_schedule will anneal lr to 0"""
+        config = ppo.ddppo.DEFAULT_CONFIG.copy()
+        config["num_gpus_per_worker"] = 0
+        config["lr_schedule"] = [[0, config["lr"]], [1000, 0.0]]
+        num_iterations = 3
+
+        for _ in framework_iterator(config, "torch"):
+            trainer = ppo.ddppo.DDPPOTrainer(config=config, env="CartPole-v0")
+            for _ in range(num_iterations):
+                result = trainer.train()
+                lr = result["info"]["learner"]["default_policy"]["cur_lr"]
+            trainer.stop()
+            assert lr == 0.0, "lr should anneal to 0.0"
+
 
 if __name__ == "__main__":
     import pytest