[RLlib] MB-MPO cleanup (comments, docstrings, type annotations). (#11033)

rllib/evaluation/postprocessing.py

@@ -4,8 +4,19 @@ from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.annotations import DeveloperAPI
 
 
-def discount(x: np.ndarray, gamma: float):
-    return scipy.signal.lfilter([1], [1, -gamma], x[::-1], axis=0)[::-1]
+def discount_cumsum(x: np.ndarray, gamma: float) -> float:
+    """Calculates the discounted cumulative sum over a reward sequence `x`.
+
+    y[t] - discount*y[t+1] = x[t]
+    reversed(y)[t] - discount*reversed(y)[t-1] = reversed(x)[t]
+
+    Args:
+        gamma (float): The discount factor gamma.
+
+    Returns:
+        float: The discounted cumulative sum over the reward sequence `x`.
+    """
+    return scipy.signal.lfilter([1], [1, float(-gamma)], x[::-1], axis=0)[::-1]
 
 
 class Postprocessing:
@@ -54,7 +65,8 @@ def compute_advantages(rollout: SampleBatch,
             rollout[SampleBatch.REWARDS] + gamma * vpred_t[1:] - vpred_t[:-1])
         # This formula for the advantage comes from:
         # "Generalized Advantage Estimation": https://arxiv.org/abs/1506.02438
-        rollout[Postprocessing.ADVANTAGES] = discount(delta_t, gamma * lambda_)
+        rollout[Postprocessing.ADVANTAGES] = discount_cumsum(
+            delta_t, gamma * lambda_)
         rollout[Postprocessing.VALUE_TARGETS] = (
             rollout[Postprocessing.ADVANTAGES] +
             rollout[SampleBatch.VF_PREDS]).astype(np.float32)
@@ -62,8 +74,8 @@ def compute_advantages(rollout: SampleBatch,
         rewards_plus_v = np.concatenate(
             [rollout[SampleBatch.REWARDS],
              np.array([last_r])])
-        discounted_returns = discount(rewards_plus_v,
-                                      gamma)[:-1].astype(np.float32)
+        discounted_returns = discount_cumsum(rewards_plus_v,
+                                             gamma)[:-1].astype(np.float32)
 
         if use_critic:
             rollout[Postprocessing.