[rllib] Better document which methods are abstract and which ones are overrides (#3480)

python/ray/rllib/models/action_dist.py

@@ -4,10 +4,11 @@ from __future__ import print_function
 
 from collections import namedtuple
 import distutils.version
-
 import tensorflow as tf
 import numpy as np
 
+from ray.rllib.utils.annotations import override
+
 use_tf150_api = (distutils.version.LooseVersion(tf.VERSION) >=
                  distutils.version.LooseVersion("1.5.0"))
 
@@ -42,10 +43,12 @@ class ActionDistribution(object):
 class Categorical(ActionDistribution):
     """Categorical distribution for discrete action spaces."""
 
+    @override(ActionDistribution)
     def logp(self, x):
         return -tf.nn.sparse_softmax_cross_entropy_with_logits(
             logits=self.inputs, labels=x)
 
+    @override(ActionDistribution)
     def entropy(self):
         if use_tf150_api:
             a0 = self.inputs - tf.reduce_max(
@@ -61,6 +64,7 @@ class Categorical(ActionDistribution):
         p0 = ea0 / z0
         return tf.reduce_sum(p0 * (tf.log(z0) - a0), reduction_indices=[1])
 
+    @override(ActionDistribution)
     def kl(self, other):
         if use_tf150_api:
             a0 = self.inputs - tf.reduce_max(
@@ -84,6 +88,7 @@ class Categorical(ActionDistribution):
         return tf.reduce_sum(
             p0 * (a0 - tf.log(z0) - a1 + tf.log(z1)), reduction_indices=[1])
 
+    @override(ActionDistribution)
     def sample(self):
         return tf.squeeze(tf.multinomial(self.inputs, 1), axis=1)
 
@@ -102,12 +107,14 @@ class DiagGaussian(ActionDistribution):
         self.log_std = log_std
         self.std = tf.exp(log_std)
 
+    @override(ActionDistribution)
     def logp(self, x):
         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]))
 
+    @override(ActionDistribution)
     def kl(self, other):
         assert isinstance(other, DiagGaussian)
         return tf.reduce_sum(
@@ -116,11 +123,13 @@ class DiagGaussian(ActionDistribution):
             (2.0 * tf.square(other.std)) - 0.5,
             reduction_indices=[1])
 
+    @override(ActionDistribution)
     def entropy(self):
         return tf.reduce_sum(
             .5 * self.log_std + .5 * np.log(2.0 * np.pi * np.e),
             reduction_indices=[1])
 
+    @override(ActionDistribution)
     def sample(self):
         return self.mean + self.std * tf.random_normal(tf.shape(self.mean))
 
@@ -131,6 +140,7 @@ class Deterministic(ActionDistribution):
     This is similar to DiagGaussian with standard deviation zero.
     """
 
+    @override(ActionDistribution)
     def sample(self):
         return self.inputs
 
@@ -150,8 +160,8 @@ class MultiActionDistribution(ActionDistribution):
             child_list.append(distribution(split_inputs[i]))
         self.child_distributions = child_list
 
+    @override(ActionDistribution)
     def logp(self, x):
-        """The log-likelihood of the action distribution."""
         split_indices = []
         for dist in self.child_distributions:
             if isinstance(dist, Categorical):
@@ -170,8 +180,8 @@ class MultiActionDistribution(ActionDistribution):
         ])
         return np.sum(log_list)
 
+    @override(ActionDistribution)
     def kl(self, other):
-        """The KL-divergence between two action distributions."""
         kl_list = np.asarray([
             distribution.kl(other_distribution)
             for distribution, other_distribution in zip(
@@ -179,15 +189,14 @@ class MultiActionDistribution(ActionDistribution):
         ])
         return np.sum(kl_list)
 
+    @override(ActionDistribution)
     def entropy(self):
-        """The entropy of the action distribution."""
         entropy_list = np.array(
             [s.entropy() for s in self.child_distributions])
         return np.sum(entropy_list)
 
+    @override(ActionDistribution)
     def sample(self):
-        """Draw a sample from the action distribution."""
-
         return TupleActions([s.sample() for s in self.child_distributions])