We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[rllib] TorchDiagGaussian doesn’t handle multiple actions correctly.
This is not a contribution.

Ray version: 0.8.2
Python version: 3.6.8
Pytorch version: 1.4
OS: Ubuntu 18.04 Docker

TorchDiagGaussian doesn’t handle multiple actions correctly. As a result, training PPO with Pytorch will crash when the action space has more than 1 action. Here’s minimal reproduction script:

```python
import gym
from gym.spaces import Box
from ray import tune

class ContinuousEnv(gym.Env):
   def __init__(self, config):
       self.action_space = Box(0.0, 1.0, shape=(2,))
       self.observation_space = Box(0.0, 1.0, shape=(1, ))

   def reset(self):
       return [0.0]

   def step(self, action):
       return [0.0], 1.0, False, {}

tune.run(
   "PPO",
   config={"env": ContinuousEnv, "use_pytorch": True, "num_workers": 1})
```


--- 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/models/torch/torch_action_dist.py`
Content:
```
1 import numpy as np
2 
3 from ray.rllib.models.action_dist import ActionDistribution
4 from ray.rllib.utils.annotations import override
5 from ray.rllib.utils import try_import_torch
6 
7 torch, nn = try_import_torch()
8 
9 
10 class TorchDistributionWrapper(ActionDistribution):
11     """Wrapper class for torch.distributions."""
12 
13     def __init_(self, inputs):
14         super().__init__(inputs)
15         # Store the last sample here.
16         self.last_sample = None
17 
18     @override(ActionDistribution)
19     def logp(self, actions):
20         return self.dist.log_prob(actions)
21 
22     @override(ActionDistribution)
23     def entropy(self):
24         return self.dist.entropy()
25 
26     @override(ActionDistribution)
27     def kl(self, other):
28         return torch.distributions.kl.kl_divergence(self.dist, other.dist)
29 
30     @override(ActionDistribution)
31     def sample(self):
32         self.last_sample = self.dist.sample()
33         return self.last_sample
34 
35     @override(ActionDistribution)
36     def sampled_action_logp(self):
37         assert self.last_sample is not None
38         return self.logp(self.last_sample)
39 
40 
41 class TorchCategorical(TorchDistributionWrapper):
42     """Wrapper class for PyTorch Categorical distribution."""
43 
44     @override(ActionDistribution)
45     def __init__(self, inputs, model):
46         super().__init__(inputs, model)
47         self.dist = torch.distributions.categorical.Categorical(logits=inputs)
48 
49     @override(ActionDistribution)
50     def deterministic_sample(self):
51         return self.dist.probs.argmax(dim=1)
52 
53     @staticmethod
54     @override(ActionDistribution)
55     def required_model_output_shape(action_space, model_config):
56         return action_space.n
57 
58 
59 class TorchDiagGaussian(TorchDistributionWrapper):
60     """Wrapper class for PyTorch Normal distribution."""
61 
62     @override(ActionDistribution)
63     def __init__(self, inputs, model):
64         super().__init__(inputs, model)
65         mean, log_std = torch.chunk(inputs, 2, dim=1)
66         self.dist = torch.distributions.normal.Normal(mean, torch.exp(log_std))
67 
68     @override(ActionDistribution)
69     def deterministic_sample(self):
70         return self.dist.mean
71 
72     @override(TorchDistributionWrapper)
73     def logp(self, actions):
74         return TorchDistributionWrapper.logp(self, actions).sum(-1)
75 
76     @staticmethod
77     @override(ActionDistribution)
78     def required_model_output_shape(action_space, model_config):
79         return np.prod(action_space.shape) * 2
80
```

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

{"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@@ -67,9 +67,15 @@\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 \n-        def reduce_mean_valid(t):\n-            return torch.mean(t * valid_mask)\n+            def reduce_mean_valid(t):\n+                return torch.mean(t * valid_mask)\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@@ -109,13 +115,11 @@\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-    else:\n-        mask = torch.ones_like(\n-            train_batch[Postprocessing.ADVANTAGES], dtype=torch.bool)\n \n     policy.loss_obj = PPOLoss(\n         dist_class,\ndiff --git a/rllib/models/torch/torch_action_dist.py b/rllib/models/torch/torch_action_dist.py\n--- a/rllib/models/torch/torch_action_dist.py\n+++ b/rllib/models/torch/torch_action_dist.py\n@@ -71,7 +71,15 @@\n \n     @override(TorchDistributionWrapper)\n     def logp(self, actions):\n-        return TorchDistributionWrapper.logp(self, actions).sum(-1)\n+        return super().logp(actions).sum(-1)\n+\n+    @override(TorchDistributionWrapper)\n+    def entropy(self):\n+        return super().entropy().sum(-1)\n+\n+    @override(TorchDistributionWrapper)\n+    def kl(self, other):\n+        return super().kl(other).sum(-1)\n \n     @staticmethod\n     @override(ActionDistribution)\n", "issue": "[rllib] TorchDiagGaussian doesn\u2019t handle multiple actions correctly.\nThis is not a contribution.\r\n\r\nRay version: 0.8.2\r\nPython version: 3.6.8\r\nPytorch version: 1.4\r\nOS: Ubuntu 18.04 Docker\r\n\r\nTorchDiagGaussian doesn\u2019t handle multiple actions correctly. As a result, training PPO with Pytorch will crash when the action space has more than 1 action. Here\u2019s minimal reproduction script:\r\n\r\n```python\r\nimport gym\r\nfrom gym.spaces import Box\r\nfrom ray import tune\r\n\r\nclass ContinuousEnv(gym.Env):\r\n   def __init__(self, config):\r\n       self.action_space = Box(0.0, 1.0, shape=(2,))\r\n       self.observation_space = Box(0.0, 1.0, shape=(1, ))\r\n\r\n   def reset(self):\r\n       return [0.0]\r\n\r\n   def step(self, action):\r\n       return [0.0], 1.0, False, {}\r\n\r\ntune.run(\r\n   \"PPO\",\r\n   config={\"env\": ContinuousEnv, \"use_pytorch\": True, \"num_workers\": 1})\r\n```\r\n\n", "before_files": [{"content": "import numpy as np\n\nfrom ray.rllib.models.action_dist import ActionDistribution\nfrom ray.rllib.utils.annotations import override\nfrom ray.rllib.utils import try_import_torch\n\ntorch, nn = try_import_torch()\n\n\nclass TorchDistributionWrapper(ActionDistribution):\n    \"\"\"Wrapper class for torch.distributions.\"\"\"\n\n    def __init_(self, inputs):\n        super().__init__(inputs)\n        # Store the last sample here.\n        self.last_sample = None\n\n    @override(ActionDistribution)\n    def logp(self, actions):\n        return self.dist.log_prob(actions)\n\n    @override(ActionDistribution)\n    def entropy(self):\n        return self.dist.entropy()\n\n    @override(ActionDistribution)\n    def kl(self, other):\n        return torch.distributions.kl.kl_divergence(self.dist, other.dist)\n\n    @override(ActionDistribution)\n    def sample(self):\n        self.last_sample = self.dist.sample()\n        return self.last_sample\n\n    @override(ActionDistribution)\n    def sampled_action_logp(self):\n        assert self.last_sample is not None\n        return self.logp(self.last_sample)\n\n\nclass TorchCategorical(TorchDistributionWrapper):\n    \"\"\"Wrapper class for PyTorch Categorical distribution.\"\"\"\n\n    @override(ActionDistribution)\n    def __init__(self, inputs, model):\n        super().__init__(inputs, model)\n        self.dist = torch.distributions.categorical.Categorical(logits=inputs)\n\n    @override(ActionDistribution)\n    def deterministic_sample(self):\n        return self.dist.probs.argmax(dim=1)\n\n    @staticmethod\n    @override(ActionDistribution)\n    def required_model_output_shape(action_space, model_config):\n        return action_space.n\n\n\nclass TorchDiagGaussian(TorchDistributionWrapper):\n    \"\"\"Wrapper class for PyTorch Normal distribution.\"\"\"\n\n    @override(ActionDistribution)\n    def __init__(self, inputs, model):\n        super().__init__(inputs, model)\n        mean, log_std = torch.chunk(inputs, 2, dim=1)\n        self.dist = torch.distributions.normal.Normal(mean, torch.exp(log_std))\n\n    @override(ActionDistribution)\n    def deterministic_sample(self):\n        return self.dist.mean\n\n    @override(TorchDistributionWrapper)\n    def logp(self, actions):\n        return TorchDistributionWrapper.logp(self, actions).sum(-1)\n\n    @staticmethod\n    @override(ActionDistribution)\n    def required_model_output_shape(action_space, model_config):\n        return np.prod(action_space.shape) * 2\n", "path": "rllib/models/torch/torch_action_dist.py"}, {"content": "import logging\n\nimport ray\nfrom ray.rllib.agents.impala.vtrace_policy import BEHAVIOUR_LOGITS\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.policy import ACTION_LOGP\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\n        def reduce_mean_valid(t):\n            return torch.mean(t * valid_mask)\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    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    else:\n        mask = torch.ones_like(\n            train_batch[Postprocessing.ADVANTAGES], dtype=torch.bool)\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[BEHAVIOUR_LOGITS],\n        train_batch[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_and_logits_fetches(policy, input_dict, state_batches, model,\n                                action_dist):\n    \"\"\"Adds value function and logits outputs to experience train_batches.\"\"\"\n    return {\n        SampleBatch.VF_PREDS: policy.model.value_function(),\n        BEHAVIOUR_LOGITS: policy.model.last_output(),\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: torch.Tensor([ob]).to(self.device),\n                    SampleBatch.PREV_ACTIONS: torch.Tensor([prev_action]).to(\n                        self.device),\n                    SampleBatch.PREV_REWARDS: torch.Tensor([prev_reward]).to(\n                        self.device),\n                    \"is_training\": False,\n                }, [torch.Tensor([s]).to(self.device) for s in state],\n                                          torch.Tensor([1]).to(self.device))\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_and_logits_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 numpy as np\n\nfrom ray.rllib.models.action_dist import ActionDistribution\nfrom ray.rllib.utils.annotations import override\nfrom ray.rllib.utils import try_import_torch\n\ntorch, nn = try_import_torch()\n\n\nclass TorchDistributionWrapper(ActionDistribution):\n    \"\"\"Wrapper class for torch.distributions.\"\"\"\n\n    def __init_(self, inputs):\n        super().__init__(inputs)\n        # Store the last sample here.\n        self.last_sample = None\n\n    @override(ActionDistribution)\n    def logp(self, actions):\n        return self.dist.log_prob(actions)\n\n    @override(ActionDistribution)\n    def entropy(self):\n        return self.dist.entropy()\n\n    @override(ActionDistribution)\n    def kl(self, other):\n        return torch.distributions.kl.kl_divergence(self.dist, other.dist)\n\n    @override(ActionDistribution)\n    def sample(self):\n        self.last_sample = self.dist.sample()\n        return self.last_sample\n\n    @override(ActionDistribution)\n    def sampled_action_logp(self):\n        assert self.last_sample is not None\n        return self.logp(self.last_sample)\n\n\nclass TorchCategorical(TorchDistributionWrapper):\n    \"\"\"Wrapper class for PyTorch Categorical distribution.\"\"\"\n\n    @override(ActionDistribution)\n    def __init__(self, inputs, model):\n        super().__init__(inputs, model)\n        self.dist = torch.distributions.categorical.Categorical(logits=inputs)\n\n    @override(ActionDistribution)\n    def deterministic_sample(self):\n        return self.dist.probs.argmax(dim=1)\n\n    @staticmethod\n    @override(ActionDistribution)\n    def required_model_output_shape(action_space, model_config):\n        return action_space.n\n\n\nclass TorchDiagGaussian(TorchDistributionWrapper):\n    \"\"\"Wrapper class for PyTorch Normal distribution.\"\"\"\n\n    @override(ActionDistribution)\n    def __init__(self, inputs, model):\n        super().__init__(inputs, model)\n        mean, log_std = torch.chunk(inputs, 2, dim=1)\n        self.dist = torch.distributions.normal.Normal(mean, torch.exp(log_std))\n\n    @override(ActionDistribution)\n    def deterministic_sample(self):\n        return self.dist.mean\n\n    @override(TorchDistributionWrapper)\n    def logp(self, actions):\n        return super().logp(actions).sum(-1)\n\n    @override(TorchDistributionWrapper)\n    def entropy(self):\n        return super().entropy().sum(-1)\n\n    @override(TorchDistributionWrapper)\n    def kl(self, other):\n        return super().kl(other).sum(-1)\n\n    @staticmethod\n    @override(ActionDistribution)\n    def required_model_output_shape(action_space, model_config):\n        return np.prod(action_space.shape) * 2\n", "path": "rllib/models/torch/torch_action_dist.py"}, {"content": "import logging\n\nimport ray\nfrom ray.rllib.agents.impala.vtrace_policy import BEHAVIOUR_LOGITS\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.policy import ACTION_LOGP\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\n            def reduce_mean_valid(t):\n                return torch.mean(t * valid_mask)\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[BEHAVIOUR_LOGITS],\n        train_batch[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_and_logits_fetches(policy, input_dict, state_batches, model,\n                                action_dist):\n    \"\"\"Adds value function and logits outputs to experience train_batches.\"\"\"\n    return {\n        SampleBatch.VF_PREDS: policy.model.value_function(),\n        BEHAVIOUR_LOGITS: policy.model.last_output(),\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: torch.Tensor([ob]).to(self.device),\n                    SampleBatch.PREV_ACTIONS: torch.Tensor([prev_action]).to(\n                        self.device),\n                    SampleBatch.PREV_REWARDS: torch.Tensor([prev_reward]).to(\n                        self.device),\n                    \"is_training\": False,\n                }, [torch.Tensor([s]).to(self.device) for s in state],\n                                          torch.Tensor([1]).to(self.device))\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_and_logits_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"}]}