We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Inconsistency in Diagonal Gaussian between log probability and entropy
In the DiagonalGaussian action distribution the logp is calculated like so:

`        return (-0.5 * tf.reduce_sum(
            tf.square((x - self.mean) / self.std), reduction_indices=[1]) -
                0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) -
                tf.reduce_sum(self.log_std, reduction_indices=[1]))`

Here the term  tf.reduce_sum(self.log_std, reduction_indices=[1])) is supposed to compute the part of the logprobability due to the determinant of the covariance matrix. The diagonals of the covariance matrix are the squared standard deviations, so the absence of the 0.5 in front of the term makes sense. But, then, in that case the entropy

`tf.reduce_sum(
            .5 * self.log_std + .5 * np.log(2.0 * np.pi * np.e),
            reduction_indices=[1])`

<img width="310" alt="image" src="https://user-images.githubusercontent.com/7660397/70380361-ed418a80-18ee-11ea-917f-0f266b48ada5.png">

should also not have a 0.5 in front of the log-std term and should be
`tf.reduce_sum(
            self.log_std + .5 * np.log(2.0 * np.pi * np.e),
            reduction_indices=[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/tf/tf_action_dist.py`
Content:
```
1 from __future__ import absolute_import
2 from __future__ import division
3 from __future__ import print_function
4 
5 import numpy as np
6 
7 from ray.rllib.models.action_dist import ActionDistribution
8 from ray.rllib.policy.policy import TupleActions
9 from ray.rllib.utils.annotations import override, DeveloperAPI
10 from ray.rllib.utils import try_import_tf
11 
12 tf = try_import_tf()
13 
14 
15 @DeveloperAPI
16 class TFActionDistribution(ActionDistribution):
17     """TF-specific extensions for building action distributions."""
18 
19     @DeveloperAPI
20     def __init__(self, inputs, model):
21         super(TFActionDistribution, self).__init__(inputs, model)
22         self.sample_op = self._build_sample_op()
23 
24     @DeveloperAPI
25     def _build_sample_op(self):
26         """Implement this instead of sample(), to enable op reuse.
27 
28         This is needed since the sample op is non-deterministic and is shared
29         between sample() and sampled_action_logp().
30         """
31         raise NotImplementedError
32 
33     @override(ActionDistribution)
34     def sample(self):
35         """Draw a sample from the action distribution."""
36         return self.sample_op
37 
38     @override(ActionDistribution)
39     def sampled_action_logp(self):
40         """Returns the log probability of the sampled action."""
41         return self.logp(self.sample_op)
42 
43 
44 class Categorical(TFActionDistribution):
45     """Categorical distribution for discrete action spaces."""
46 
47     @DeveloperAPI
48     def __init__(self, inputs, model=None):
49         super(Categorical, self).__init__(inputs, model)
50 
51     @override(ActionDistribution)
52     def logp(self, x):
53         return -tf.nn.sparse_softmax_cross_entropy_with_logits(
54             logits=self.inputs, labels=tf.cast(x, tf.int32))
55 
56     @override(ActionDistribution)
57     def entropy(self):
58         a0 = self.inputs - tf.reduce_max(
59             self.inputs, reduction_indices=[1], keep_dims=True)
60         ea0 = tf.exp(a0)
61         z0 = tf.reduce_sum(ea0, reduction_indices=[1], keep_dims=True)
62         p0 = ea0 / z0
63         return tf.reduce_sum(p0 * (tf.log(z0) - a0), reduction_indices=[1])
64 
65     @override(ActionDistribution)
66     def kl(self, other):
67         a0 = self.inputs - tf.reduce_max(
68             self.inputs, reduction_indices=[1], keep_dims=True)
69         a1 = other.inputs - tf.reduce_max(
70             other.inputs, reduction_indices=[1], keep_dims=True)
71         ea0 = tf.exp(a0)
72         ea1 = tf.exp(a1)
73         z0 = tf.reduce_sum(ea0, reduction_indices=[1], keep_dims=True)
74         z1 = tf.reduce_sum(ea1, reduction_indices=[1], keep_dims=True)
75         p0 = ea0 / z0
76         return tf.reduce_sum(
77             p0 * (a0 - tf.log(z0) - a1 + tf.log(z1)), reduction_indices=[1])
78 
79     @override(TFActionDistribution)
80     def _build_sample_op(self):
81         return tf.squeeze(tf.multinomial(self.inputs, 1), axis=1)
82 
83     @staticmethod
84     @override(ActionDistribution)
85     def required_model_output_shape(action_space, model_config):
86         return action_space.n
87 
88 
89 class MultiCategorical(TFActionDistribution):
90     """MultiCategorical distribution for MultiDiscrete action spaces."""
91 
92     def __init__(self, inputs, model, input_lens):
93         # skip TFActionDistribution init
94         ActionDistribution.__init__(self, inputs, model)
95         self.cats = [
96             Categorical(input_, model)
97             for input_ in tf.split(inputs, input_lens, axis=1)
98         ]
99         self.sample_op = self._build_sample_op()
100 
101     @override(ActionDistribution)
102     def logp(self, actions):
103         # If tensor is provided, unstack it into list
104         if isinstance(actions, tf.Tensor):
105             actions = tf.unstack(tf.cast(actions, tf.int32), axis=1)
106         logps = tf.stack(
107             [cat.logp(act) for cat, act in zip(self.cats, actions)])
108         return tf.reduce_sum(logps, axis=0)
109 
110     @override(ActionDistribution)
111     def multi_entropy(self):
112         return tf.stack([cat.entropy() for cat in self.cats], axis=1)
113 
114     @override(ActionDistribution)
115     def entropy(self):
116         return tf.reduce_sum(self.multi_entropy(), axis=1)
117 
118     @override(ActionDistribution)
119     def multi_kl(self, other):
120         return [cat.kl(oth_cat) for cat, oth_cat in zip(self.cats, other.cats)]
121 
122     @override(ActionDistribution)
123     def kl(self, other):
124         return tf.reduce_sum(self.multi_kl(other), axis=1)
125 
126     @override(TFActionDistribution)
127     def _build_sample_op(self):
128         return tf.stack([cat.sample() for cat in self.cats], axis=1)
129 
130     @staticmethod
131     @override(ActionDistribution)
132     def required_model_output_shape(action_space, model_config):
133         return np.sum(action_space.nvec)
134 
135 
136 class DiagGaussian(TFActionDistribution):
137     """Action distribution where each vector element is a gaussian.
138 
139     The first half of the input vector defines the gaussian means, and the
140     second half the gaussian standard deviations.
141     """
142 
143     def __init__(self, inputs, model):
144         mean, log_std = tf.split(inputs, 2, axis=1)
145         self.mean = mean
146         self.log_std = log_std
147         self.std = tf.exp(log_std)
148         TFActionDistribution.__init__(self, inputs, model)
149 
150     @override(ActionDistribution)
151     def logp(self, x):
152         return (-0.5 * tf.reduce_sum(
153             tf.square((x - self.mean) / self.std), reduction_indices=[1]) -
154                 0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) -
155                 tf.reduce_sum(self.log_std, reduction_indices=[1]))
156 
157     @override(ActionDistribution)
158     def kl(self, other):
159         assert isinstance(other, DiagGaussian)
160         return tf.reduce_sum(
161             other.log_std - self.log_std +
162             (tf.square(self.std) + tf.square(self.mean - other.mean)) /
163             (2.0 * tf.square(other.std)) - 0.5,
164             reduction_indices=[1])
165 
166     @override(ActionDistribution)
167     def entropy(self):
168         return tf.reduce_sum(
169             .5 * self.log_std + .5 * np.log(2.0 * np.pi * np.e),
170             reduction_indices=[1])
171 
172     @override(TFActionDistribution)
173     def _build_sample_op(self):
174         return self.mean + self.std * tf.random_normal(tf.shape(self.mean))
175 
176     @staticmethod
177     @override(ActionDistribution)
178     def required_model_output_shape(action_space, model_config):
179         return np.prod(action_space.shape) * 2
180 
181 
182 class Deterministic(TFActionDistribution):
183     """Action distribution that returns the input values directly.
184 
185     This is similar to DiagGaussian with standard deviation zero.
186     """
187 
188     @override(TFActionDistribution)
189     def sampled_action_logp(self):
190         return 0.0
191 
192     @override(TFActionDistribution)
193     def _build_sample_op(self):
194         return self.inputs
195 
196     @staticmethod
197     @override(ActionDistribution)
198     def required_model_output_shape(action_space, model_config):
199         return np.prod(action_space.shape)
200 
201 
202 class MultiActionDistribution(TFActionDistribution):
203     """Action distribution that operates for list of actions.
204 
205     Args:
206         inputs (Tensor list): A list of tensors from which to compute samples.
207     """
208 
209     def __init__(self, inputs, model, action_space, child_distributions,
210                  input_lens):
211         # skip TFActionDistribution init
212         ActionDistribution.__init__(self, inputs, model)
213         self.input_lens = input_lens
214         split_inputs = tf.split(inputs, self.input_lens, axis=1)
215         child_list = []
216         for i, distribution in enumerate(child_distributions):
217             child_list.append(distribution(split_inputs[i], model))
218         self.child_distributions = child_list
219 
220     @override(ActionDistribution)
221     def logp(self, x):
222         split_indices = []
223         for dist in self.child_distributions:
224             if isinstance(dist, Categorical):
225                 split_indices.append(1)
226             else:
227                 split_indices.append(tf.shape(dist.sample())[1])
228         split_list = tf.split(x, split_indices, axis=1)
229         for i, distribution in enumerate(self.child_distributions):
230             # Remove extra categorical dimension
231             if isinstance(distribution, Categorical):
232                 split_list[i] = tf.cast(
233                     tf.squeeze(split_list[i], axis=-1), tf.int32)
234         log_list = np.asarray([
235             distribution.logp(split_x) for distribution, split_x in zip(
236                 self.child_distributions, split_list)
237         ])
238         return np.sum(log_list)
239 
240     @override(ActionDistribution)
241     def kl(self, other):
242         kl_list = np.asarray([
243             distribution.kl(other_distribution)
244             for distribution, other_distribution in zip(
245                 self.child_distributions, other.child_distributions)
246         ])
247         return np.sum(kl_list)
248 
249     @override(ActionDistribution)
250     def entropy(self):
251         entropy_list = np.array(
252             [s.entropy() for s in self.child_distributions])
253         return np.sum(entropy_list)
254 
255     @override(ActionDistribution)
256     def sample(self):
257         return TupleActions([s.sample() for s in self.child_distributions])
258 
259     @override(TFActionDistribution)
260     def sampled_action_logp(self):
261         p = self.child_distributions[0].sampled_action_logp()
262         for c in self.child_distributions[1:]:
263             p += c.sampled_action_logp()
264         return p
265 
266 
267 class Dirichlet(TFActionDistribution):
268     """Dirichlet distribution for continuous actions that are between
269     [0,1] and sum to 1.
270 
271     e.g. actions that represent resource allocation."""
272 
273     def __init__(self, inputs, model):
274         """Input is a tensor of logits. The exponential of logits is used to
275         parametrize the Dirichlet distribution as all parameters need to be
276         positive. An arbitrary small epsilon is added to the concentration
277         parameters to be zero due to numerical error.
278 
279         See issue #4440 for more details.
280         """
281         self.epsilon = 1e-7
282         concentration = tf.exp(inputs) + self.epsilon
283         self.dist = tf.distributions.Dirichlet(
284             concentration=concentration,
285             validate_args=True,
286             allow_nan_stats=False,
287         )
288         TFActionDistribution.__init__(self, concentration, model)
289 
290     @override(ActionDistribution)
291     def logp(self, x):
292         # Support of Dirichlet are positive real numbers. x is already be
293         # an array of positive number, but we clip to avoid zeros due to
294         # numerical errors.
295         x = tf.maximum(x, self.epsilon)
296         x = x / tf.reduce_sum(x, axis=-1, keepdims=True)
297         return self.dist.log_prob(x)
298 
299     @override(ActionDistribution)
300     def entropy(self):
301         return self.dist.entropy()
302 
303     @override(ActionDistribution)
304     def kl(self, other):
305         return self.dist.kl_divergence(other.dist)
306 
307     @override(TFActionDistribution)
308     def _build_sample_op(self):
309         return self.dist.sample()
310 
311     @staticmethod
312     @override(ActionDistribution)
313     def required_model_output_shape(action_space, model_config):
314         return np.prod(action_space.shape)
315
```
--- 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/models/tf/tf_action_dist.py b/rllib/models/tf/tf_action_dist.py
--- a/rllib/models/tf/tf_action_dist.py
+++ b/rllib/models/tf/tf_action_dist.py
@@ -166,7 +166,7 @@
     @override(ActionDistribution)
     def entropy(self):
         return tf.reduce_sum(
-            .5 * self.log_std + .5 * np.log(2.0 * np.pi * np.e),
+            self.log_std + .5 * np.log(2.0 * np.pi * np.e),
             reduction_indices=[1])
 
     @override(TFActionDistribution)

{"golden_diff": "diff --git a/rllib/models/tf/tf_action_dist.py b/rllib/models/tf/tf_action_dist.py\n--- a/rllib/models/tf/tf_action_dist.py\n+++ b/rllib/models/tf/tf_action_dist.py\n@@ -166,7 +166,7 @@\n     @override(ActionDistribution)\n     def entropy(self):\n         return tf.reduce_sum(\n-            .5 * self.log_std + .5 * np.log(2.0 * np.pi * np.e),\n+            self.log_std + .5 * np.log(2.0 * np.pi * np.e),\n             reduction_indices=[1])\n \n     @override(TFActionDistribution)\n", "issue": "Inconsistency in Diagonal Gaussian between log probability and entropy\nIn the DiagonalGaussian action distribution the logp is calculated like so:\r\n\r\n`        return (-0.5 * tf.reduce_sum(\r\n            tf.square((x - self.mean) / self.std), reduction_indices=[1]) -\r\n                0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) -\r\n                tf.reduce_sum(self.log_std, reduction_indices=[1]))`\r\n\r\nHere the term  tf.reduce_sum(self.log_std, reduction_indices=[1])) is supposed to compute the part of the logprobability due to the determinant of the covariance matrix. The diagonals of the covariance matrix are the squared standard deviations, so the absence of the 0.5 in front of the term makes sense. But, then, in that case the entropy\r\n\r\n`tf.reduce_sum(\r\n            .5 * self.log_std + .5 * np.log(2.0 * np.pi * np.e),\r\n            reduction_indices=[1])`\r\n\r\n<img width=\"310\" alt=\"image\" src=\"https://user-images.githubusercontent.com/7660397/70380361-ed418a80-18ee-11ea-917f-0f266b48ada5.png\">\r\n\r\nshould also not have a 0.5 in front of the log-std term and should be\r\n`tf.reduce_sum(\r\n            self.log_std + .5 * np.log(2.0 * np.pi * np.e),\r\n            reduction_indices=[1])`\n", "before_files": [{"content": "from __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport numpy as np\n\nfrom ray.rllib.models.action_dist import ActionDistribution\nfrom ray.rllib.policy.policy import TupleActions\nfrom ray.rllib.utils.annotations import override, DeveloperAPI\nfrom ray.rllib.utils import try_import_tf\n\ntf = try_import_tf()\n\n\n@DeveloperAPI\nclass TFActionDistribution(ActionDistribution):\n    \"\"\"TF-specific extensions for building action distributions.\"\"\"\n\n    @DeveloperAPI\n    def __init__(self, inputs, model):\n        super(TFActionDistribution, self).__init__(inputs, model)\n        self.sample_op = self._build_sample_op()\n\n    @DeveloperAPI\n    def _build_sample_op(self):\n        \"\"\"Implement this instead of sample(), to enable op reuse.\n\n        This is needed since the sample op is non-deterministic and is shared\n        between sample() and sampled_action_logp().\n        \"\"\"\n        raise NotImplementedError\n\n    @override(ActionDistribution)\n    def sample(self):\n        \"\"\"Draw a sample from the action distribution.\"\"\"\n        return self.sample_op\n\n    @override(ActionDistribution)\n    def sampled_action_logp(self):\n        \"\"\"Returns the log probability of the sampled action.\"\"\"\n        return self.logp(self.sample_op)\n\n\nclass Categorical(TFActionDistribution):\n    \"\"\"Categorical distribution for discrete action spaces.\"\"\"\n\n    @DeveloperAPI\n    def __init__(self, inputs, model=None):\n        super(Categorical, self).__init__(inputs, model)\n\n    @override(ActionDistribution)\n    def logp(self, x):\n        return -tf.nn.sparse_softmax_cross_entropy_with_logits(\n            logits=self.inputs, labels=tf.cast(x, tf.int32))\n\n    @override(ActionDistribution)\n    def entropy(self):\n        a0 = self.inputs - tf.reduce_max(\n            self.inputs, reduction_indices=[1], keep_dims=True)\n        ea0 = tf.exp(a0)\n        z0 = tf.reduce_sum(ea0, reduction_indices=[1], keep_dims=True)\n        p0 = ea0 / z0\n        return tf.reduce_sum(p0 * (tf.log(z0) - a0), reduction_indices=[1])\n\n    @override(ActionDistribution)\n    def kl(self, other):\n        a0 = self.inputs - tf.reduce_max(\n            self.inputs, reduction_indices=[1], keep_dims=True)\n        a1 = other.inputs - tf.reduce_max(\n            other.inputs, reduction_indices=[1], keep_dims=True)\n        ea0 = tf.exp(a0)\n        ea1 = tf.exp(a1)\n        z0 = tf.reduce_sum(ea0, reduction_indices=[1], keep_dims=True)\n        z1 = tf.reduce_sum(ea1, reduction_indices=[1], keep_dims=True)\n        p0 = ea0 / z0\n        return tf.reduce_sum(\n            p0 * (a0 - tf.log(z0) - a1 + tf.log(z1)), reduction_indices=[1])\n\n    @override(TFActionDistribution)\n    def _build_sample_op(self):\n        return tf.squeeze(tf.multinomial(self.inputs, 1), axis=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 MultiCategorical(TFActionDistribution):\n    \"\"\"MultiCategorical distribution for MultiDiscrete action spaces.\"\"\"\n\n    def __init__(self, inputs, model, input_lens):\n        # skip TFActionDistribution init\n        ActionDistribution.__init__(self, inputs, model)\n        self.cats = [\n            Categorical(input_, model)\n            for input_ in tf.split(inputs, input_lens, axis=1)\n        ]\n        self.sample_op = self._build_sample_op()\n\n    @override(ActionDistribution)\n    def logp(self, actions):\n        # If tensor is provided, unstack it into list\n        if isinstance(actions, tf.Tensor):\n            actions = tf.unstack(tf.cast(actions, tf.int32), axis=1)\n        logps = tf.stack(\n            [cat.logp(act) for cat, act in zip(self.cats, actions)])\n        return tf.reduce_sum(logps, axis=0)\n\n    @override(ActionDistribution)\n    def multi_entropy(self):\n        return tf.stack([cat.entropy() for cat in self.cats], axis=1)\n\n    @override(ActionDistribution)\n    def entropy(self):\n        return tf.reduce_sum(self.multi_entropy(), axis=1)\n\n    @override(ActionDistribution)\n    def multi_kl(self, other):\n        return [cat.kl(oth_cat) for cat, oth_cat in zip(self.cats, other.cats)]\n\n    @override(ActionDistribution)\n    def kl(self, other):\n        return tf.reduce_sum(self.multi_kl(other), axis=1)\n\n    @override(TFActionDistribution)\n    def _build_sample_op(self):\n        return tf.stack([cat.sample() for cat in self.cats], axis=1)\n\n    @staticmethod\n    @override(ActionDistribution)\n    def required_model_output_shape(action_space, model_config):\n        return np.sum(action_space.nvec)\n\n\nclass DiagGaussian(TFActionDistribution):\n    \"\"\"Action distribution where each vector element is a gaussian.\n\n    The first half of the input vector defines the gaussian means, and the\n    second half the gaussian standard deviations.\n    \"\"\"\n\n    def __init__(self, inputs, model):\n        mean, log_std = tf.split(inputs, 2, axis=1)\n        self.mean = mean\n        self.log_std = log_std\n        self.std = tf.exp(log_std)\n        TFActionDistribution.__init__(self, inputs, model)\n\n    @override(ActionDistribution)\n    def logp(self, x):\n        return (-0.5 * tf.reduce_sum(\n            tf.square((x - self.mean) / self.std), reduction_indices=[1]) -\n                0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) -\n                tf.reduce_sum(self.log_std, reduction_indices=[1]))\n\n    @override(ActionDistribution)\n    def kl(self, other):\n        assert isinstance(other, DiagGaussian)\n        return tf.reduce_sum(\n            other.log_std - self.log_std +\n            (tf.square(self.std) + tf.square(self.mean - other.mean)) /\n            (2.0 * tf.square(other.std)) - 0.5,\n            reduction_indices=[1])\n\n    @override(ActionDistribution)\n    def entropy(self):\n        return tf.reduce_sum(\n            .5 * self.log_std + .5 * np.log(2.0 * np.pi * np.e),\n            reduction_indices=[1])\n\n    @override(TFActionDistribution)\n    def _build_sample_op(self):\n        return self.mean + self.std * tf.random_normal(tf.shape(self.mean))\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\n\nclass Deterministic(TFActionDistribution):\n    \"\"\"Action distribution that returns the input values directly.\n\n    This is similar to DiagGaussian with standard deviation zero.\n    \"\"\"\n\n    @override(TFActionDistribution)\n    def sampled_action_logp(self):\n        return 0.0\n\n    @override(TFActionDistribution)\n    def _build_sample_op(self):\n        return self.inputs\n\n    @staticmethod\n    @override(ActionDistribution)\n    def required_model_output_shape(action_space, model_config):\n        return np.prod(action_space.shape)\n\n\nclass MultiActionDistribution(TFActionDistribution):\n    \"\"\"Action distribution that operates for list of actions.\n\n    Args:\n        inputs (Tensor list): A list of tensors from which to compute samples.\n    \"\"\"\n\n    def __init__(self, inputs, model, action_space, child_distributions,\n                 input_lens):\n        # skip TFActionDistribution init\n        ActionDistribution.__init__(self, inputs, model)\n        self.input_lens = input_lens\n        split_inputs = tf.split(inputs, self.input_lens, axis=1)\n        child_list = []\n        for i, distribution in enumerate(child_distributions):\n            child_list.append(distribution(split_inputs[i], model))\n        self.child_distributions = child_list\n\n    @override(ActionDistribution)\n    def logp(self, x):\n        split_indices = []\n        for dist in self.child_distributions:\n            if isinstance(dist, Categorical):\n                split_indices.append(1)\n            else:\n                split_indices.append(tf.shape(dist.sample())[1])\n        split_list = tf.split(x, split_indices, axis=1)\n        for i, distribution in enumerate(self.child_distributions):\n            # Remove extra categorical dimension\n            if isinstance(distribution, Categorical):\n                split_list[i] = tf.cast(\n                    tf.squeeze(split_list[i], axis=-1), tf.int32)\n        log_list = np.asarray([\n            distribution.logp(split_x) for distribution, split_x in zip(\n                self.child_distributions, split_list)\n        ])\n        return np.sum(log_list)\n\n    @override(ActionDistribution)\n    def kl(self, other):\n        kl_list = np.asarray([\n            distribution.kl(other_distribution)\n            for distribution, other_distribution in zip(\n                self.child_distributions, other.child_distributions)\n        ])\n        return np.sum(kl_list)\n\n    @override(ActionDistribution)\n    def entropy(self):\n        entropy_list = np.array(\n            [s.entropy() for s in self.child_distributions])\n        return np.sum(entropy_list)\n\n    @override(ActionDistribution)\n    def sample(self):\n        return TupleActions([s.sample() for s in self.child_distributions])\n\n    @override(TFActionDistribution)\n    def sampled_action_logp(self):\n        p = self.child_distributions[0].sampled_action_logp()\n        for c in self.child_distributions[1:]:\n            p += c.sampled_action_logp()\n        return p\n\n\nclass Dirichlet(TFActionDistribution):\n    \"\"\"Dirichlet distribution for continuous actions that are between\n    [0,1] and sum to 1.\n\n    e.g. actions that represent resource allocation.\"\"\"\n\n    def __init__(self, inputs, model):\n        \"\"\"Input is a tensor of logits. The exponential of logits is used to\n        parametrize the Dirichlet distribution as all parameters need to be\n        positive. An arbitrary small epsilon is added to the concentration\n        parameters to be zero due to numerical error.\n\n        See issue #4440 for more details.\n        \"\"\"\n        self.epsilon = 1e-7\n        concentration = tf.exp(inputs) + self.epsilon\n        self.dist = tf.distributions.Dirichlet(\n            concentration=concentration,\n            validate_args=True,\n            allow_nan_stats=False,\n        )\n        TFActionDistribution.__init__(self, concentration, model)\n\n    @override(ActionDistribution)\n    def logp(self, x):\n        # Support of Dirichlet are positive real numbers. x is already be\n        # an array of positive number, but we clip to avoid zeros due to\n        # numerical errors.\n        x = tf.maximum(x, self.epsilon)\n        x = x / tf.reduce_sum(x, axis=-1, keepdims=True)\n        return self.dist.log_prob(x)\n\n    @override(ActionDistribution)\n    def entropy(self):\n        return self.dist.entropy()\n\n    @override(ActionDistribution)\n    def kl(self, other):\n        return self.dist.kl_divergence(other.dist)\n\n    @override(TFActionDistribution)\n    def _build_sample_op(self):\n        return self.dist.sample()\n\n    @staticmethod\n    @override(ActionDistribution)\n    def required_model_output_shape(action_space, model_config):\n        return np.prod(action_space.shape)\n", "path": "rllib/models/tf/tf_action_dist.py"}], "after_files": [{"content": "from __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport numpy as np\n\nfrom ray.rllib.models.action_dist import ActionDistribution\nfrom ray.rllib.policy.policy import TupleActions\nfrom ray.rllib.utils.annotations import override, DeveloperAPI\nfrom ray.rllib.utils import try_import_tf\n\ntf = try_import_tf()\n\n\n@DeveloperAPI\nclass TFActionDistribution(ActionDistribution):\n    \"\"\"TF-specific extensions for building action distributions.\"\"\"\n\n    @DeveloperAPI\n    def __init__(self, inputs, model):\n        super(TFActionDistribution, self).__init__(inputs, model)\n        self.sample_op = self._build_sample_op()\n\n    @DeveloperAPI\n    def _build_sample_op(self):\n        \"\"\"Implement this instead of sample(), to enable op reuse.\n\n        This is needed since the sample op is non-deterministic and is shared\n        between sample() and sampled_action_logp().\n        \"\"\"\n        raise NotImplementedError\n\n    @override(ActionDistribution)\n    def sample(self):\n        \"\"\"Draw a sample from the action distribution.\"\"\"\n        return self.sample_op\n\n    @override(ActionDistribution)\n    def sampled_action_logp(self):\n        \"\"\"Returns the log probability of the sampled action.\"\"\"\n        return self.logp(self.sample_op)\n\n\nclass Categorical(TFActionDistribution):\n    \"\"\"Categorical distribution for discrete action spaces.\"\"\"\n\n    @DeveloperAPI\n    def __init__(self, inputs, model=None):\n        super(Categorical, self).__init__(inputs, model)\n\n    @override(ActionDistribution)\n    def logp(self, x):\n        return -tf.nn.sparse_softmax_cross_entropy_with_logits(\n            logits=self.inputs, labels=tf.cast(x, tf.int32))\n\n    @override(ActionDistribution)\n    def entropy(self):\n        a0 = self.inputs - tf.reduce_max(\n            self.inputs, reduction_indices=[1], keep_dims=True)\n        ea0 = tf.exp(a0)\n        z0 = tf.reduce_sum(ea0, reduction_indices=[1], keep_dims=True)\n        p0 = ea0 / z0\n        return tf.reduce_sum(p0 * (tf.log(z0) - a0), reduction_indices=[1])\n\n    @override(ActionDistribution)\n    def kl(self, other):\n        a0 = self.inputs - tf.reduce_max(\n            self.inputs, reduction_indices=[1], keep_dims=True)\n        a1 = other.inputs - tf.reduce_max(\n            other.inputs, reduction_indices=[1], keep_dims=True)\n        ea0 = tf.exp(a0)\n        ea1 = tf.exp(a1)\n        z0 = tf.reduce_sum(ea0, reduction_indices=[1], keep_dims=True)\n        z1 = tf.reduce_sum(ea1, reduction_indices=[1], keep_dims=True)\n        p0 = ea0 / z0\n        return tf.reduce_sum(\n            p0 * (a0 - tf.log(z0) - a1 + tf.log(z1)), reduction_indices=[1])\n\n    @override(TFActionDistribution)\n    def _build_sample_op(self):\n        return tf.squeeze(tf.multinomial(self.inputs, 1), axis=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 MultiCategorical(TFActionDistribution):\n    \"\"\"MultiCategorical distribution for MultiDiscrete action spaces.\"\"\"\n\n    def __init__(self, inputs, model, input_lens):\n        # skip TFActionDistribution init\n        ActionDistribution.__init__(self, inputs, model)\n        self.cats = [\n            Categorical(input_, model)\n            for input_ in tf.split(inputs, input_lens, axis=1)\n        ]\n        self.sample_op = self._build_sample_op()\n\n    @override(ActionDistribution)\n    def logp(self, actions):\n        # If tensor is provided, unstack it into list\n        if isinstance(actions, tf.Tensor):\n            actions = tf.unstack(tf.cast(actions, tf.int32), axis=1)\n        logps = tf.stack(\n            [cat.logp(act) for cat, act in zip(self.cats, actions)])\n        return tf.reduce_sum(logps, axis=0)\n\n    @override(ActionDistribution)\n    def multi_entropy(self):\n        return tf.stack([cat.entropy() for cat in self.cats], axis=1)\n\n    @override(ActionDistribution)\n    def entropy(self):\n        return tf.reduce_sum(self.multi_entropy(), axis=1)\n\n    @override(ActionDistribution)\n    def multi_kl(self, other):\n        return [cat.kl(oth_cat) for cat, oth_cat in zip(self.cats, other.cats)]\n\n    @override(ActionDistribution)\n    def kl(self, other):\n        return tf.reduce_sum(self.multi_kl(other), axis=1)\n\n    @override(TFActionDistribution)\n    def _build_sample_op(self):\n        return tf.stack([cat.sample() for cat in self.cats], axis=1)\n\n    @staticmethod\n    @override(ActionDistribution)\n    def required_model_output_shape(action_space, model_config):\n        return np.sum(action_space.nvec)\n\n\nclass DiagGaussian(TFActionDistribution):\n    \"\"\"Action distribution where each vector element is a gaussian.\n\n    The first half of the input vector defines the gaussian means, and the\n    second half the gaussian standard deviations.\n    \"\"\"\n\n    def __init__(self, inputs, model):\n        mean, log_std = tf.split(inputs, 2, axis=1)\n        self.mean = mean\n        self.log_std = log_std\n        self.std = tf.exp(log_std)\n        TFActionDistribution.__init__(self, inputs, model)\n\n    @override(ActionDistribution)\n    def logp(self, x):\n        return (-0.5 * tf.reduce_sum(\n            tf.square((x - self.mean) / self.std), reduction_indices=[1]) -\n                0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) -\n                tf.reduce_sum(self.log_std, reduction_indices=[1]))\n\n    @override(ActionDistribution)\n    def kl(self, other):\n        assert isinstance(other, DiagGaussian)\n        return tf.reduce_sum(\n            other.log_std - self.log_std +\n            (tf.square(self.std) + tf.square(self.mean - other.mean)) /\n            (2.0 * tf.square(other.std)) - 0.5,\n            reduction_indices=[1])\n\n    @override(ActionDistribution)\n    def entropy(self):\n        return tf.reduce_sum(\n            self.log_std + .5 * np.log(2.0 * np.pi * np.e),\n            reduction_indices=[1])\n\n    @override(TFActionDistribution)\n    def _build_sample_op(self):\n        return self.mean + self.std * tf.random_normal(tf.shape(self.mean))\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\n\nclass Deterministic(TFActionDistribution):\n    \"\"\"Action distribution that returns the input values directly.\n\n    This is similar to DiagGaussian with standard deviation zero.\n    \"\"\"\n\n    @override(TFActionDistribution)\n    def sampled_action_logp(self):\n        return 0.0\n\n    @override(TFActionDistribution)\n    def _build_sample_op(self):\n        return self.inputs\n\n    @staticmethod\n    @override(ActionDistribution)\n    def required_model_output_shape(action_space, model_config):\n        return np.prod(action_space.shape)\n\n\nclass MultiActionDistribution(TFActionDistribution):\n    \"\"\"Action distribution that operates for list of actions.\n\n    Args:\n        inputs (Tensor list): A list of tensors from which to compute samples.\n    \"\"\"\n\n    def __init__(self, inputs, model, action_space, child_distributions,\n                 input_lens):\n        # skip TFActionDistribution init\n        ActionDistribution.__init__(self, inputs, model)\n        self.input_lens = input_lens\n        split_inputs = tf.split(inputs, self.input_lens, axis=1)\n        child_list = []\n        for i, distribution in enumerate(child_distributions):\n            child_list.append(distribution(split_inputs[i], model))\n        self.child_distributions = child_list\n\n    @override(ActionDistribution)\n    def logp(self, x):\n        split_indices = []\n        for dist in self.child_distributions:\n            if isinstance(dist, Categorical):\n                split_indices.append(1)\n            else:\n                split_indices.append(tf.shape(dist.sample())[1])\n        split_list = tf.split(x, split_indices, axis=1)\n        for i, distribution in enumerate(self.child_distributions):\n            # Remove extra categorical dimension\n            if isinstance(distribution, Categorical):\n                split_list[i] = tf.cast(\n                    tf.squeeze(split_list[i], axis=-1), tf.int32)\n        log_list = np.asarray([\n            distribution.logp(split_x) for distribution, split_x in zip(\n                self.child_distributions, split_list)\n        ])\n        return np.sum(log_list)\n\n    @override(ActionDistribution)\n    def kl(self, other):\n        kl_list = np.asarray([\n            distribution.kl(other_distribution)\n            for distribution, other_distribution in zip(\n                self.child_distributions, other.child_distributions)\n        ])\n        return np.sum(kl_list)\n\n    @override(ActionDistribution)\n    def entropy(self):\n        entropy_list = np.array(\n            [s.entropy() for s in self.child_distributions])\n        return np.sum(entropy_list)\n\n    @override(ActionDistribution)\n    def sample(self):\n        return TupleActions([s.sample() for s in self.child_distributions])\n\n    @override(TFActionDistribution)\n    def sampled_action_logp(self):\n        p = self.child_distributions[0].sampled_action_logp()\n        for c in self.child_distributions[1:]:\n            p += c.sampled_action_logp()\n        return p\n\n\nclass Dirichlet(TFActionDistribution):\n    \"\"\"Dirichlet distribution for continuous actions that are between\n    [0,1] and sum to 1.\n\n    e.g. actions that represent resource allocation.\"\"\"\n\n    def __init__(self, inputs, model):\n        \"\"\"Input is a tensor of logits. The exponential of logits is used to\n        parametrize the Dirichlet distribution as all parameters need to be\n        positive. An arbitrary small epsilon is added to the concentration\n        parameters to be zero due to numerical error.\n\n        See issue #4440 for more details.\n        \"\"\"\n        self.epsilon = 1e-7\n        concentration = tf.exp(inputs) + self.epsilon\n        self.dist = tf.distributions.Dirichlet(\n            concentration=concentration,\n            validate_args=True,\n            allow_nan_stats=False,\n        )\n        TFActionDistribution.__init__(self, concentration, model)\n\n    @override(ActionDistribution)\n    def logp(self, x):\n        # Support of Dirichlet are positive real numbers. x is already be\n        # an array of positive number, but we clip to avoid zeros due to\n        # numerical errors.\n        x = tf.maximum(x, self.epsilon)\n        x = x / tf.reduce_sum(x, axis=-1, keepdims=True)\n        return self.dist.log_prob(x)\n\n    @override(ActionDistribution)\n    def entropy(self):\n        return self.dist.entropy()\n\n    @override(ActionDistribution)\n    def kl(self, other):\n        return self.dist.kl_divergence(other.dist)\n\n    @override(TFActionDistribution)\n    def _build_sample_op(self):\n        return self.dist.sample()\n\n    @staticmethod\n    @override(ActionDistribution)\n    def required_model_output_shape(action_space, model_config):\n        return np.prod(action_space.shape)\n", "path": "rllib/models/tf/tf_action_dist.py"}]}