wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-06 05:16:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

[RLlib] Fix KL method of MultiCategorial tf distribution (issue #7009). (#7119)

Browse Source 
* Fix KL method of MultiCategorial tf distribution.

* Fix KL method of MultiCategorial tf distribution.

* Merge AsyncReplayOptimizer fixes into this branch.
This commit is contained in:
Sven Mika
2020-02-12 21:46:15 +01:00
committed by GitHub
parent 275fd343fb
commit f41a9b9813
1 changed files with 8 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/models/tf/tf_action_dist.py
+8 -9
View File
@@ -61,17 +61,14 @@ class Categorical(TFActionDistribution):
@override(ActionDistribution)
def kl(self, other):
a0 = self.inputs - tf.reduce_max(
self.inputs, reduction_indices=[1], keep_dims=True)
a1 = other.inputs - tf.reduce_max(
other.inputs, reduction_indices=[1], keep_dims=True)
a0 = self.inputs - tf.reduce_max(self.inputs, axis=1, keep_dims=True)
a1 = other.inputs - tf.reduce_max(other.inputs, axis=1, keep_dims=True)
ea0 = tf.exp(a0)
ea1 = tf.exp(a1)
z0 = tf.reduce_sum(ea0, reduction_indices=[1], keep_dims=True)
z1 = tf.reduce_sum(ea1, reduction_indices=[1], keep_dims=True)
z0 = tf.reduce_sum(ea0, axis=1, keep_dims=True)
z1 = tf.reduce_sum(ea1, axis=1, keep_dims=True)
p0 = ea0 / z0
return tf.reduce_sum(
p0 * (a0 - tf.log(z0) - a1 + tf.log(z1)), reduction_indices=[1])
return tf.reduce_sum(p0 * (a0 - tf.log(z0) - a1 + tf.log(z1)), axis=1)
@override(TFActionDistribution)
def _build_sample_op(self):
@@ -114,7 +111,9 @@ class MultiCategorical(TFActionDistribution):
@override(ActionDistribution)
def multi_kl(self, other):
return [cat.kl(oth_cat) for cat, oth_cat in zip(self.cats, other.cats)]
return tf.stack(
[cat.kl(oth_cat) for cat, oth_cat in zip(self.cats, other.cats)],
axis=1)
@override(ActionDistribution)
def kl(self, other):