We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
"lr_schedule" option ignored using torch framework and PPO algorithm
<!--Please include [tune], [rllib], [autoscaler] etc. in the issue title if relevant-->

*Ray version and other system information (Python version, TensorFlow version, OS):*
- Ray: [0.9.0.dev0 (2c599dbf05e41e338920ee2fbe692658bcbec4dd)](https://s3-us-west-2.amazonaws.com/ray-wheels/releases/0.8.5/02c1ab0ec6d615ad54ebf33bd93c51c04000534e/ray-0.8.5-cp36-cp36m-manylinux1_x86_64.whl)
- CUDA: 10.1
- Pytorch: 1.4.0 with GPU support
- Ubuntu 18.04
- Python 3.6

### What is the problem?

Setting the hyperparameter "lr_schedule" as no effect when using PyTorch as backend framework and PPO learning algorithm.

### Reproduction (REQUIRED)

```
import ray 
from ray.rllib.agents.ppo import PPOTrainer, DEFAULT_CONFIG 
 
config = DEFAULT_CONFIG.copy() 
for key, val in { 
    "env": "CartPole-v0", 
    "num_workers": 0, 
    "use_pytorch": False, 
    "lr": 1.0e-5, 
    "lr_schedule": [ 
        [0, 1.0e-6], 
        [1, 1.0e-7], 
    ] 
}.items(): config[key] = val 
 
ray.init() 

for use_pytorch in [False, True]: 
    config["use_pytorch"] = use_pytorch 
    agent = PPOTrainer(config, "CartPole-v0") 
    for _ in range(2): 
        result = agent.train() 
        print(f"use_pytorch: {use_pytorch} - Current learning rate: "\
              f"{result['info']['learner']['default_policy']['cur_lr']}")
```

- [x] I have verified my script runs in a clean environment and reproduces the issue.
- [x] I have verified the issue also occurs with the [latest wheels](https://docs.ray.io/en/latest/installation.html).


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `rllib/agents/ppo/ppo_torch_policy.py`
Content:
```
1 import logging
2 
3 import ray
4 from ray.rllib.agents.a3c.a3c_torch_policy import apply_grad_clipping
5 from ray.rllib.agents.ppo.ppo_tf_policy import postprocess_ppo_gae, \
6     setup_config
7 from ray.rllib.evaluation.postprocessing import Postprocessing
8 from ray.rllib.policy.sample_batch import SampleBatch
9 from ray.rllib.policy.torch_policy import EntropyCoeffSchedule, \
10     LearningRateSchedule
11 from ray.rllib.policy.torch_policy_template import build_torch_policy
12 from ray.rllib.utils.explained_variance import explained_variance
13 from ray.rllib.utils.torch_ops import sequence_mask
14 from ray.rllib.utils import try_import_torch
15 
16 torch, nn = try_import_torch()
17 
18 logger = logging.getLogger(__name__)
19 
20 
21 class PPOLoss:
22     def __init__(self,
23                  dist_class,
24                  model,
25                  value_targets,
26                  advantages,
27                  actions,
28                  prev_logits,
29                  prev_actions_logp,
30                  vf_preds,
31                  curr_action_dist,
32                  value_fn,
33                  cur_kl_coeff,
34                  valid_mask,
35                  entropy_coeff=0,
36                  clip_param=0.1,
37                  vf_clip_param=0.1,
38                  vf_loss_coeff=1.0,
39                  use_gae=True):
40         """Constructs the loss for Proximal Policy Objective.
41 
42         Arguments:
43             dist_class: action distribution class for logits.
44             value_targets (Placeholder): Placeholder for target values; used
45                 for GAE.
46             actions (Placeholder): Placeholder for actions taken
47                 from previous model evaluation.
48             advantages (Placeholder): Placeholder for calculated advantages
49                 from previous model evaluation.
50             prev_logits (Placeholder): Placeholder for logits output from
51                 previous model evaluation.
52             prev_actions_logp (Placeholder): Placeholder for prob output from
53                 previous model evaluation.
54             vf_preds (Placeholder): Placeholder for value function output
55                 from previous model evaluation.
56             curr_action_dist (ActionDistribution): ActionDistribution
57                 of the current model.
58             value_fn (Tensor): Current value function output Tensor.
59             cur_kl_coeff (Variable): Variable holding the current PPO KL
60                 coefficient.
61             valid_mask (Tensor): A bool mask of valid input elements (#2992).
62             entropy_coeff (float): Coefficient of the entropy regularizer.
63             clip_param (float): Clip parameter
64             vf_clip_param (float): Clip parameter for the value function
65             vf_loss_coeff (float): Coefficient of the value function loss
66             use_gae (bool): If true, use the Generalized Advantage Estimator.
67         """
68         if valid_mask is not None:
69             num_valid = torch.sum(valid_mask)
70 
71             def reduce_mean_valid(t):
72                 return torch.sum(t * valid_mask) / num_valid
73 
74         else:
75 
76             def reduce_mean_valid(t):
77                 return torch.mean(t)
78 
79         prev_dist = dist_class(prev_logits, model)
80         # Make loss functions.
81         logp_ratio = torch.exp(
82             curr_action_dist.logp(actions) - prev_actions_logp)
83         action_kl = prev_dist.kl(curr_action_dist)
84         self.mean_kl = reduce_mean_valid(action_kl)
85 
86         curr_entropy = curr_action_dist.entropy()
87         self.mean_entropy = reduce_mean_valid(curr_entropy)
88 
89         surrogate_loss = torch.min(
90             advantages * logp_ratio,
91             advantages * torch.clamp(logp_ratio, 1 - clip_param,
92                                      1 + clip_param))
93         self.mean_policy_loss = reduce_mean_valid(-surrogate_loss)
94 
95         if use_gae:
96             vf_loss1 = torch.pow(value_fn - value_targets, 2.0)
97             vf_clipped = vf_preds + torch.clamp(value_fn - vf_preds,
98                                                 -vf_clip_param, vf_clip_param)
99             vf_loss2 = torch.pow(vf_clipped - value_targets, 2.0)
100             vf_loss = torch.max(vf_loss1, vf_loss2)
101             self.mean_vf_loss = reduce_mean_valid(vf_loss)
102             loss = reduce_mean_valid(
103                 -surrogate_loss + cur_kl_coeff * action_kl +
104                 vf_loss_coeff * vf_loss - entropy_coeff * curr_entropy)
105         else:
106             self.mean_vf_loss = 0.0
107             loss = reduce_mean_valid(-surrogate_loss +
108                                      cur_kl_coeff * action_kl -
109                                      entropy_coeff * curr_entropy)
110         self.loss = loss
111 
112 
113 def ppo_surrogate_loss(policy, model, dist_class, train_batch):
114     logits, state = model.from_batch(train_batch)
115     action_dist = dist_class(logits, model)
116 
117     mask = None
118     if state:
119         max_seq_len = torch.max(train_batch["seq_lens"])
120         mask = sequence_mask(train_batch["seq_lens"], max_seq_len)
121         mask = torch.reshape(mask, [-1])
122 
123     policy.loss_obj = PPOLoss(
124         dist_class,
125         model,
126         train_batch[Postprocessing.VALUE_TARGETS],
127         train_batch[Postprocessing.ADVANTAGES],
128         train_batch[SampleBatch.ACTIONS],
129         train_batch[SampleBatch.ACTION_DIST_INPUTS],
130         train_batch[SampleBatch.ACTION_LOGP],
131         train_batch[SampleBatch.VF_PREDS],
132         action_dist,
133         model.value_function(),
134         policy.kl_coeff,
135         mask,
136         entropy_coeff=policy.entropy_coeff,
137         clip_param=policy.config["clip_param"],
138         vf_clip_param=policy.config["vf_clip_param"],
139         vf_loss_coeff=policy.config["vf_loss_coeff"],
140         use_gae=policy.config["use_gae"],
141     )
142 
143     return policy.loss_obj.loss
144 
145 
146 def kl_and_loss_stats(policy, train_batch):
147     return {
148         "cur_kl_coeff": policy.kl_coeff,
149         "cur_lr": policy.cur_lr,
150         "total_loss": policy.loss_obj.loss,
151         "policy_loss": policy.loss_obj.mean_policy_loss,
152         "vf_loss": policy.loss_obj.mean_vf_loss,
153         "vf_explained_var": explained_variance(
154             train_batch[Postprocessing.VALUE_TARGETS],
155             policy.model.value_function(),
156             framework="torch"),
157         "kl": policy.loss_obj.mean_kl,
158         "entropy": policy.loss_obj.mean_entropy,
159         "entropy_coeff": policy.entropy_coeff,
160     }
161 
162 
163 def vf_preds_fetches(policy, input_dict, state_batches, model, action_dist):
164     """Adds value function outputs to experience train_batches."""
165     return {
166         SampleBatch.VF_PREDS: policy.model.value_function(),
167     }
168 
169 
170 class KLCoeffMixin:
171     def __init__(self, config):
172         # KL Coefficient.
173         self.kl_coeff = config["kl_coeff"]
174         self.kl_target = config["kl_target"]
175 
176     def update_kl(self, sampled_kl):
177         if sampled_kl > 2.0 * self.kl_target:
178             self.kl_coeff *= 1.5
179         elif sampled_kl < 0.5 * self.kl_target:
180             self.kl_coeff *= 0.5
181         return self.kl_coeff
182 
183 
184 class ValueNetworkMixin:
185     def __init__(self, obs_space, action_space, config):
186         if config["use_gae"]:
187 
188             def value(ob, prev_action, prev_reward, *state):
189                 model_out, _ = self.model({
190                     SampleBatch.CUR_OBS: self._convert_to_tensor([ob]),
191                     SampleBatch.PREV_ACTIONS: self._convert_to_tensor(
192                         [prev_action]),
193                     SampleBatch.PREV_REWARDS: self._convert_to_tensor(
194                         [prev_reward]),
195                     "is_training": False,
196                 }, [self._convert_to_tensor(s) for s in state],
197                                           self._convert_to_tensor([1]))
198                 return self.model.value_function()[0]
199 
200         else:
201 
202             def value(ob, prev_action, prev_reward, *state):
203                 return 0.0
204 
205         self._value = value
206 
207 
208 def setup_mixins(policy, obs_space, action_space, config):
209     ValueNetworkMixin.__init__(policy, obs_space, action_space, config)
210     KLCoeffMixin.__init__(policy, config)
211     EntropyCoeffSchedule.__init__(policy, config["entropy_coeff"],
212                                   config["entropy_coeff_schedule"])
213     LearningRateSchedule.__init__(policy, config["lr"], config["lr_schedule"])
214 
215 
216 PPOTorchPolicy = build_torch_policy(
217     name="PPOTorchPolicy",
218     get_default_config=lambda: ray.rllib.agents.ppo.ppo.DEFAULT_CONFIG,
219     loss_fn=ppo_surrogate_loss,
220     stats_fn=kl_and_loss_stats,
221     extra_action_out_fn=vf_preds_fetches,
222     postprocess_fn=postprocess_ppo_gae,
223     extra_grad_process_fn=apply_grad_clipping,
224     before_init=setup_config,
225     after_init=setup_mixins,
226     mixins=[KLCoeffMixin, ValueNetworkMixin])
227
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/rllib/agents/ppo/ppo_torch_policy.py b/rllib/agents/ppo/ppo_torch_policy.py
--- a/rllib/agents/ppo/ppo_torch_policy.py
+++ b/rllib/agents/ppo/ppo_torch_policy.py
@@ -223,4 +223,7 @@
     extra_grad_process_fn=apply_grad_clipping,
     before_init=setup_config,
     after_init=setup_mixins,
-    mixins=[KLCoeffMixin, ValueNetworkMixin])
+    mixins=[
+        LearningRateSchedule, EntropyCoeffSchedule, KLCoeffMixin,
+        ValueNetworkMixin
+    ])

{"golden_diff": "diff --git a/rllib/agents/ppo/ppo_torch_policy.py b/rllib/agents/ppo/ppo_torch_policy.py\n--- a/rllib/agents/ppo/ppo_torch_policy.py\n+++ b/rllib/agents/ppo/ppo_torch_policy.py\n@@ -223,4 +223,7 @@\n     extra_grad_process_fn=apply_grad_clipping,\n     before_init=setup_config,\n     after_init=setup_mixins,\n-    mixins=[KLCoeffMixin, ValueNetworkMixin])\n+    mixins=[\n+        LearningRateSchedule, EntropyCoeffSchedule, KLCoeffMixin,\n+        ValueNetworkMixin\n+    ])\n", "issue": "\"lr_schedule\" option ignored using torch framework and PPO algorithm\n<!--Please include [tune], [rllib], [autoscaler] etc. in the issue title if relevant-->\r\n\r\n*Ray version and other system information (Python version, TensorFlow version, OS):*\r\n- Ray: [0.9.0.dev0 (2c599dbf05e41e338920ee2fbe692658bcbec4dd)](https://s3-us-west-2.amazonaws.com/ray-wheels/releases/0.8.5/02c1ab0ec6d615ad54ebf33bd93c51c04000534e/ray-0.8.5-cp36-cp36m-manylinux1_x86_64.whl)\r\n- CUDA: 10.1\r\n- Pytorch: 1.4.0 with GPU support\r\n- Ubuntu 18.04\r\n- Python 3.6\r\n\r\n### What is the problem?\r\n\r\nSetting the hyperparameter \"lr_schedule\" as no effect when using PyTorch as backend framework and PPO learning algorithm.\r\n\r\n### Reproduction (REQUIRED)\r\n\r\n```\r\nimport ray \r\nfrom ray.rllib.agents.ppo import PPOTrainer, DEFAULT_CONFIG \r\n \r\nconfig = DEFAULT_CONFIG.copy() \r\nfor key, val in { \r\n    \"env\": \"CartPole-v0\", \r\n    \"num_workers\": 0, \r\n    \"use_pytorch\": False, \r\n    \"lr\": 1.0e-5, \r\n    \"lr_schedule\": [ \r\n        [0, 1.0e-6], \r\n        [1, 1.0e-7], \r\n    ] \r\n}.items(): config[key] = val \r\n \r\nray.init() \r\n\r\nfor use_pytorch in [False, True]: \r\n    config[\"use_pytorch\"] = use_pytorch \r\n    agent = PPOTrainer(config, \"CartPole-v0\") \r\n    for _ in range(2): \r\n        result = agent.train() \r\n        print(f\"use_pytorch: {use_pytorch} - Current learning rate: \"\\\r\n              f\"{result['info']['learner']['default_policy']['cur_lr']}\")\r\n```\r\n\r\n- [x] I have verified my script runs in a clean environment and reproduces the issue.\r\n- [x] I have verified the issue also occurs with the [latest wheels](https://docs.ray.io/en/latest/installation.html).\r\n\n", "before_files": [{"content": "import logging\n\nimport ray\nfrom ray.rllib.agents.a3c.a3c_torch_policy import apply_grad_clipping\nfrom ray.rllib.agents.ppo.ppo_tf_policy import postprocess_ppo_gae, \\\n    setup_config\nfrom ray.rllib.evaluation.postprocessing import Postprocessing\nfrom ray.rllib.policy.sample_batch import SampleBatch\nfrom ray.rllib.policy.torch_policy import EntropyCoeffSchedule, \\\n    LearningRateSchedule\nfrom ray.rllib.policy.torch_policy_template import build_torch_policy\nfrom ray.rllib.utils.explained_variance import explained_variance\nfrom ray.rllib.utils.torch_ops import sequence_mask\nfrom ray.rllib.utils import try_import_torch\n\ntorch, nn = try_import_torch()\n\nlogger = logging.getLogger(__name__)\n\n\nclass PPOLoss:\n    def __init__(self,\n                 dist_class,\n                 model,\n                 value_targets,\n                 advantages,\n                 actions,\n                 prev_logits,\n                 prev_actions_logp,\n                 vf_preds,\n                 curr_action_dist,\n                 value_fn,\n                 cur_kl_coeff,\n                 valid_mask,\n                 entropy_coeff=0,\n                 clip_param=0.1,\n                 vf_clip_param=0.1,\n                 vf_loss_coeff=1.0,\n                 use_gae=True):\n        \"\"\"Constructs the loss for Proximal Policy Objective.\n\n        Arguments:\n            dist_class: action distribution class for logits.\n            value_targets (Placeholder): Placeholder for target values; used\n                for GAE.\n            actions (Placeholder): Placeholder for actions taken\n                from previous model evaluation.\n            advantages (Placeholder): Placeholder for calculated advantages\n                from previous model evaluation.\n            prev_logits (Placeholder): Placeholder for logits output from\n                previous model evaluation.\n            prev_actions_logp (Placeholder): Placeholder for prob output from\n                previous model evaluation.\n            vf_preds (Placeholder): Placeholder for value function output\n                from previous model evaluation.\n            curr_action_dist (ActionDistribution): ActionDistribution\n                of the current model.\n            value_fn (Tensor): Current value function output Tensor.\n            cur_kl_coeff (Variable): Variable holding the current PPO KL\n                coefficient.\n            valid_mask (Tensor): A bool mask of valid input elements (#2992).\n            entropy_coeff (float): Coefficient of the entropy regularizer.\n            clip_param (float): Clip parameter\n            vf_clip_param (float): Clip parameter for the value function\n            vf_loss_coeff (float): Coefficient of the value function loss\n            use_gae (bool): If true, use the Generalized Advantage Estimator.\n        \"\"\"\n        if valid_mask is not None:\n            num_valid = torch.sum(valid_mask)\n\n            def reduce_mean_valid(t):\n                return torch.sum(t * valid_mask) / num_valid\n\n        else:\n\n            def reduce_mean_valid(t):\n                return torch.mean(t)\n\n        prev_dist = dist_class(prev_logits, model)\n        # Make loss functions.\n        logp_ratio = torch.exp(\n            curr_action_dist.logp(actions) - prev_actions_logp)\n        action_kl = prev_dist.kl(curr_action_dist)\n        self.mean_kl = reduce_mean_valid(action_kl)\n\n        curr_entropy = curr_action_dist.entropy()\n        self.mean_entropy = reduce_mean_valid(curr_entropy)\n\n        surrogate_loss = torch.min(\n            advantages * logp_ratio,\n            advantages * torch.clamp(logp_ratio, 1 - clip_param,\n                                     1 + clip_param))\n        self.mean_policy_loss = reduce_mean_valid(-surrogate_loss)\n\n        if use_gae:\n            vf_loss1 = torch.pow(value_fn - value_targets, 2.0)\n            vf_clipped = vf_preds + torch.clamp(value_fn - vf_preds,\n                                                -vf_clip_param, vf_clip_param)\n            vf_loss2 = torch.pow(vf_clipped - value_targets, 2.0)\n            vf_loss = torch.max(vf_loss1, vf_loss2)\n            self.mean_vf_loss = reduce_mean_valid(vf_loss)\n            loss = reduce_mean_valid(\n                -surrogate_loss + cur_kl_coeff * action_kl +\n                vf_loss_coeff * vf_loss - entropy_coeff * curr_entropy)\n        else:\n            self.mean_vf_loss = 0.0\n            loss = reduce_mean_valid(-surrogate_loss +\n                                     cur_kl_coeff * action_kl -\n                                     entropy_coeff * curr_entropy)\n        self.loss = loss\n\n\ndef ppo_surrogate_loss(policy, model, dist_class, train_batch):\n    logits, state = model.from_batch(train_batch)\n    action_dist = dist_class(logits, model)\n\n    mask = None\n    if state:\n        max_seq_len = torch.max(train_batch[\"seq_lens\"])\n        mask = sequence_mask(train_batch[\"seq_lens\"], max_seq_len)\n        mask = torch.reshape(mask, [-1])\n\n    policy.loss_obj = PPOLoss(\n        dist_class,\n        model,\n        train_batch[Postprocessing.VALUE_TARGETS],\n        train_batch[Postprocessing.ADVANTAGES],\n        train_batch[SampleBatch.ACTIONS],\n        train_batch[SampleBatch.ACTION_DIST_INPUTS],\n        train_batch[SampleBatch.ACTION_LOGP],\n        train_batch[SampleBatch.VF_PREDS],\n        action_dist,\n        model.value_function(),\n        policy.kl_coeff,\n        mask,\n        entropy_coeff=policy.entropy_coeff,\n        clip_param=policy.config[\"clip_param\"],\n        vf_clip_param=policy.config[\"vf_clip_param\"],\n        vf_loss_coeff=policy.config[\"vf_loss_coeff\"],\n        use_gae=policy.config[\"use_gae\"],\n    )\n\n    return policy.loss_obj.loss\n\n\ndef kl_and_loss_stats(policy, train_batch):\n    return {\n        \"cur_kl_coeff\": policy.kl_coeff,\n        \"cur_lr\": policy.cur_lr,\n        \"total_loss\": policy.loss_obj.loss,\n        \"policy_loss\": policy.loss_obj.mean_policy_loss,\n        \"vf_loss\": policy.loss_obj.mean_vf_loss,\n        \"vf_explained_var\": explained_variance(\n            train_batch[Postprocessing.VALUE_TARGETS],\n            policy.model.value_function(),\n            framework=\"torch\"),\n        \"kl\": policy.loss_obj.mean_kl,\n        \"entropy\": policy.loss_obj.mean_entropy,\n        \"entropy_coeff\": policy.entropy_coeff,\n    }\n\n\ndef vf_preds_fetches(policy, input_dict, state_batches, model, action_dist):\n    \"\"\"Adds value function outputs to experience train_batches.\"\"\"\n    return {\n        SampleBatch.VF_PREDS: policy.model.value_function(),\n    }\n\n\nclass KLCoeffMixin:\n    def __init__(self, config):\n        # KL Coefficient.\n        self.kl_coeff = config[\"kl_coeff\"]\n        self.kl_target = config[\"kl_target\"]\n\n    def update_kl(self, sampled_kl):\n        if sampled_kl > 2.0 * self.kl_target:\n            self.kl_coeff *= 1.5\n        elif sampled_kl < 0.5 * self.kl_target:\n            self.kl_coeff *= 0.5\n        return self.kl_coeff\n\n\nclass ValueNetworkMixin:\n    def __init__(self, obs_space, action_space, config):\n        if config[\"use_gae\"]:\n\n            def value(ob, prev_action, prev_reward, *state):\n                model_out, _ = self.model({\n                    SampleBatch.CUR_OBS: self._convert_to_tensor([ob]),\n                    SampleBatch.PREV_ACTIONS: self._convert_to_tensor(\n                        [prev_action]),\n                    SampleBatch.PREV_REWARDS: self._convert_to_tensor(\n                        [prev_reward]),\n                    \"is_training\": False,\n                }, [self._convert_to_tensor(s) for s in state],\n                                          self._convert_to_tensor([1]))\n                return self.model.value_function()[0]\n\n        else:\n\n            def value(ob, prev_action, prev_reward, *state):\n                return 0.0\n\n        self._value = value\n\n\ndef setup_mixins(policy, obs_space, action_space, config):\n    ValueNetworkMixin.__init__(policy, obs_space, action_space, config)\n    KLCoeffMixin.__init__(policy, config)\n    EntropyCoeffSchedule.__init__(policy, config[\"entropy_coeff\"],\n                                  config[\"entropy_coeff_schedule\"])\n    LearningRateSchedule.__init__(policy, config[\"lr\"], config[\"lr_schedule\"])\n\n\nPPOTorchPolicy = build_torch_policy(\n    name=\"PPOTorchPolicy\",\n    get_default_config=lambda: ray.rllib.agents.ppo.ppo.DEFAULT_CONFIG,\n    loss_fn=ppo_surrogate_loss,\n    stats_fn=kl_and_loss_stats,\n    extra_action_out_fn=vf_preds_fetches,\n    postprocess_fn=postprocess_ppo_gae,\n    extra_grad_process_fn=apply_grad_clipping,\n    before_init=setup_config,\n    after_init=setup_mixins,\n    mixins=[KLCoeffMixin, ValueNetworkMixin])\n", "path": "rllib/agents/ppo/ppo_torch_policy.py"}], "after_files": [{"content": "import logging\n\nimport ray\nfrom ray.rllib.agents.a3c.a3c_torch_policy import apply_grad_clipping\nfrom ray.rllib.agents.ppo.ppo_tf_policy import postprocess_ppo_gae, \\\n    setup_config\nfrom ray.rllib.evaluation.postprocessing import Postprocessing\nfrom ray.rllib.policy.sample_batch import SampleBatch\nfrom ray.rllib.policy.torch_policy import EntropyCoeffSchedule, \\\n    LearningRateSchedule\nfrom ray.rllib.policy.torch_policy_template import build_torch_policy\nfrom ray.rllib.utils.explained_variance import explained_variance\nfrom ray.rllib.utils.torch_ops import sequence_mask\nfrom ray.rllib.utils import try_import_torch\n\ntorch, nn = try_import_torch()\n\nlogger = logging.getLogger(__name__)\n\n\nclass PPOLoss:\n    def __init__(self,\n                 dist_class,\n                 model,\n                 value_targets,\n                 advantages,\n                 actions,\n                 prev_logits,\n                 prev_actions_logp,\n                 vf_preds,\n                 curr_action_dist,\n                 value_fn,\n                 cur_kl_coeff,\n                 valid_mask,\n                 entropy_coeff=0,\n                 clip_param=0.1,\n                 vf_clip_param=0.1,\n                 vf_loss_coeff=1.0,\n                 use_gae=True):\n        \"\"\"Constructs the loss for Proximal Policy Objective.\n\n        Arguments:\n            dist_class: action distribution class for logits.\n            value_targets (Placeholder): Placeholder for target values; used\n                for GAE.\n            actions (Placeholder): Placeholder for actions taken\n                from previous model evaluation.\n            advantages (Placeholder): Placeholder for calculated advantages\n                from previous model evaluation.\n            prev_logits (Placeholder): Placeholder for logits output from\n                previous model evaluation.\n            prev_actions_logp (Placeholder): Placeholder for prob output from\n                previous model evaluation.\n            vf_preds (Placeholder): Placeholder for value function output\n                from previous model evaluation.\n            curr_action_dist (ActionDistribution): ActionDistribution\n                of the current model.\n            value_fn (Tensor): Current value function output Tensor.\n            cur_kl_coeff (Variable): Variable holding the current PPO KL\n                coefficient.\n            valid_mask (Tensor): A bool mask of valid input elements (#2992).\n            entropy_coeff (float): Coefficient of the entropy regularizer.\n            clip_param (float): Clip parameter\n            vf_clip_param (float): Clip parameter for the value function\n            vf_loss_coeff (float): Coefficient of the value function loss\n            use_gae (bool): If true, use the Generalized Advantage Estimator.\n        \"\"\"\n        if valid_mask is not None:\n            num_valid = torch.sum(valid_mask)\n\n            def reduce_mean_valid(t):\n                return torch.sum(t * valid_mask) / num_valid\n\n        else:\n\n            def reduce_mean_valid(t):\n                return torch.mean(t)\n\n        prev_dist = dist_class(prev_logits, model)\n        # Make loss functions.\n        logp_ratio = torch.exp(\n            curr_action_dist.logp(actions) - prev_actions_logp)\n        action_kl = prev_dist.kl(curr_action_dist)\n        self.mean_kl = reduce_mean_valid(action_kl)\n\n        curr_entropy = curr_action_dist.entropy()\n        self.mean_entropy = reduce_mean_valid(curr_entropy)\n\n        surrogate_loss = torch.min(\n            advantages * logp_ratio,\n            advantages * torch.clamp(logp_ratio, 1 - clip_param,\n                                     1 + clip_param))\n        self.mean_policy_loss = reduce_mean_valid(-surrogate_loss)\n\n        if use_gae:\n            vf_loss1 = torch.pow(value_fn - value_targets, 2.0)\n            vf_clipped = vf_preds + torch.clamp(value_fn - vf_preds,\n                                                -vf_clip_param, vf_clip_param)\n            vf_loss2 = torch.pow(vf_clipped - value_targets, 2.0)\n            vf_loss = torch.max(vf_loss1, vf_loss2)\n            self.mean_vf_loss = reduce_mean_valid(vf_loss)\n            loss = reduce_mean_valid(\n                -surrogate_loss + cur_kl_coeff * action_kl +\n                vf_loss_coeff * vf_loss - entropy_coeff * curr_entropy)\n        else:\n            self.mean_vf_loss = 0.0\n            loss = reduce_mean_valid(-surrogate_loss +\n                                     cur_kl_coeff * action_kl -\n                                     entropy_coeff * curr_entropy)\n        self.loss = loss\n\n\ndef ppo_surrogate_loss(policy, model, dist_class, train_batch):\n    logits, state = model.from_batch(train_batch)\n    action_dist = dist_class(logits, model)\n\n    mask = None\n    if state:\n        max_seq_len = torch.max(train_batch[\"seq_lens\"])\n        mask = sequence_mask(train_batch[\"seq_lens\"], max_seq_len)\n        mask = torch.reshape(mask, [-1])\n\n    policy.loss_obj = PPOLoss(\n        dist_class,\n        model,\n        train_batch[Postprocessing.VALUE_TARGETS],\n        train_batch[Postprocessing.ADVANTAGES],\n        train_batch[SampleBatch.ACTIONS],\n        train_batch[SampleBatch.ACTION_DIST_INPUTS],\n        train_batch[SampleBatch.ACTION_LOGP],\n        train_batch[SampleBatch.VF_PREDS],\n        action_dist,\n        model.value_function(),\n        policy.kl_coeff,\n        mask,\n        entropy_coeff=policy.entropy_coeff,\n        clip_param=policy.config[\"clip_param\"],\n        vf_clip_param=policy.config[\"vf_clip_param\"],\n        vf_loss_coeff=policy.config[\"vf_loss_coeff\"],\n        use_gae=policy.config[\"use_gae\"],\n    )\n\n    return policy.loss_obj.loss\n\n\ndef kl_and_loss_stats(policy, train_batch):\n    return {\n        \"cur_kl_coeff\": policy.kl_coeff,\n        \"cur_lr\": policy.cur_lr,\n        \"total_loss\": policy.loss_obj.loss,\n        \"policy_loss\": policy.loss_obj.mean_policy_loss,\n        \"vf_loss\": policy.loss_obj.mean_vf_loss,\n        \"vf_explained_var\": explained_variance(\n            train_batch[Postprocessing.VALUE_TARGETS],\n            policy.model.value_function(),\n            framework=\"torch\"),\n        \"kl\": policy.loss_obj.mean_kl,\n        \"entropy\": policy.loss_obj.mean_entropy,\n        \"entropy_coeff\": policy.entropy_coeff,\n    }\n\n\ndef vf_preds_fetches(policy, input_dict, state_batches, model, action_dist):\n    \"\"\"Adds value function outputs to experience train_batches.\"\"\"\n    return {\n        SampleBatch.VF_PREDS: policy.model.value_function(),\n    }\n\n\nclass KLCoeffMixin:\n    def __init__(self, config):\n        # KL Coefficient.\n        self.kl_coeff = config[\"kl_coeff\"]\n        self.kl_target = config[\"kl_target\"]\n\n    def update_kl(self, sampled_kl):\n        if sampled_kl > 2.0 * self.kl_target:\n            self.kl_coeff *= 1.5\n        elif sampled_kl < 0.5 * self.kl_target:\n            self.kl_coeff *= 0.5\n        return self.kl_coeff\n\n\nclass ValueNetworkMixin:\n    def __init__(self, obs_space, action_space, config):\n        if config[\"use_gae\"]:\n\n            def value(ob, prev_action, prev_reward, *state):\n                model_out, _ = self.model({\n                    SampleBatch.CUR_OBS: self._convert_to_tensor([ob]),\n                    SampleBatch.PREV_ACTIONS: self._convert_to_tensor(\n                        [prev_action]),\n                    SampleBatch.PREV_REWARDS: self._convert_to_tensor(\n                        [prev_reward]),\n                    \"is_training\": False,\n                }, [self._convert_to_tensor(s) for s in state],\n                                          self._convert_to_tensor([1]))\n                return self.model.value_function()[0]\n\n        else:\n\n            def value(ob, prev_action, prev_reward, *state):\n                return 0.0\n\n        self._value = value\n\n\ndef setup_mixins(policy, obs_space, action_space, config):\n    ValueNetworkMixin.__init__(policy, obs_space, action_space, config)\n    KLCoeffMixin.__init__(policy, config)\n    EntropyCoeffSchedule.__init__(policy, config[\"entropy_coeff\"],\n                                  config[\"entropy_coeff_schedule\"])\n    LearningRateSchedule.__init__(policy, config[\"lr\"], config[\"lr_schedule\"])\n\n\nPPOTorchPolicy = build_torch_policy(\n    name=\"PPOTorchPolicy\",\n    get_default_config=lambda: ray.rllib.agents.ppo.ppo.DEFAULT_CONFIG,\n    loss_fn=ppo_surrogate_loss,\n    stats_fn=kl_and_loss_stats,\n    extra_action_out_fn=vf_preds_fetches,\n    postprocess_fn=postprocess_ppo_gae,\n    extra_grad_process_fn=apply_grad_clipping,\n    before_init=setup_config,\n    after_init=setup_mixins,\n    mixins=[\n        LearningRateSchedule, EntropyCoeffSchedule, KLCoeffMixin,\n        ValueNetworkMixin\n    ])\n", "path": "rllib/agents/ppo/ppo_torch_policy.py"}]}