Add explained variance

torchy_baselines/common/utils.py

@@ -52,3 +52,22 @@ def discount_cumsum(x, discount):
          x2]
     """
     return scipy.signal.lfilter([1], [1, float(-discount)], x[::-1], axis=0)[::-1]
+
+# From stable baselines
+def explained_variance(y_pred, y_true):
+    """
+    Computes fraction of variance that ypred explains about y.
+    Returns 1 - Var[y-ypred] / Var[y]
+
+    interpretation:
+        ev=0  =>  might as well have predicted zero
+        ev=1  =>  perfect prediction
+        ev<0  =>  worse than just predicting zero
+
+    :param y_pred: (np.ndarray) the prediction
+    :param y_true: (np.ndarray) the expected value
+    :return: (float) explained variance of ypred and y
+    """
+    assert y_true.ndim == 1 and y_pred.ndim == 1
+    var_y = np.var(y_true)
+    return np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y

torchy_baselines/ppo/ppo.py

@@ -8,6 +8,7 @@ from torchy_baselines.common.base_class import BaseRLModel
 from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.ppo.policies import PPOPolicy
 from torchy_baselines.common.buffers import RolloutBuffer
+from torchy_baselines.common.utils import explained_variance
 
 
 class PPO(BaseRLModel):
@@ -137,6 +138,7 @@ class PPO(BaseRLModel):
                 # TODO: clip grad norm?
                 # nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
                 self.policy.optimizer.step()
+        print(explained_variance(return_batch.numpy()[:, 0], values[:, 0].detach().cpu().numpy()))
 
     def learn(self, total_timesteps, callback=None, log_interval=100,
               eval_freq=-1, n_eval_episodes=5, tb_log_name="PPO", reset_num_timesteps=True):