diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py
index 41ca60c..36ab9aa 100644
--- a/torchy_baselines/a2c/a2c.py
+++ b/torchy_baselines/a2c/a2c.py
@@ -92,7 +92,7 @@ class A2C(PPO):
                 action = action.long().flatten()
 
             # TODO: avoid second computation of everything because of the gradient
-            values, log_prob, entropy = self.policy.get_policy_stats(obs, action)
+            values, log_prob, entropy = self.policy.evaluate_actions(obs, action)
             values = values.flatten()
 
             # Normalize advantage (not present in the original implementation)
diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index a8d56ae..36cd0c5 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -107,6 +107,9 @@ class BaseRLModel(object):
         """
         Rescale the action from [low, high] to [-1, 1]
         (no need for symmetric action space)
+
+        :param action: (np.ndarray)
+        :return: (np.ndarray)
         """
         low, high = self.action_space.low, self.action_space.high
         return 2.0 * ((action - low) / (high - low)) - 1.0
@@ -115,6 +118,9 @@ class BaseRLModel(object):
         """
         Rescale the action from [-1, 1] to [low, high]
         (no need for symmetric action space)
+
+        :param scaled_action: (np.ndarray)
+        :return: (np.ndarray)
         """
         low, high = self.action_space.low, self.action_space.high
         return low + (0.5 * (scaled_action + 1.0) * (high - low))
@@ -343,13 +349,11 @@ class BaseRLModel(object):
         if self._vec_normalize_env is not None:
             obs_ = self._vec_normalize_env.get_original_obs()
 
-
         self.rollout_data = None
         if hasattr(self, 'use_sde') and self.use_sde:
             self.actor.reset_noise()
             # Reset rollout data
             self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']}
-            # self.rollout_data = {'observations': [], 'actions': [], 'rewards': [], 'returns': [], 'dones': []}
 
         while total_steps < n_steps or total_episodes < n_episodes:
             done = False
diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py
index 1e3b25c..1cb7ca4 100644
--- a/torchy_baselines/ppo/policies.py
+++ b/torchy_baselines/ppo/policies.py
@@ -5,7 +5,7 @@ import torch as th
 import torch.nn as nn
 import numpy as np
 
-from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp
+from torchy_baselines.common.policies import BasePolicy, register_policy
 from torchy_baselines.common.distributions import make_proba_distribution,\
     DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution
 
@@ -30,7 +30,7 @@ class MlpExtractor(nn.Module):
 
     Adapted from Stable Baselines.
 
-    :param flat_observations: (th.Tensor) The observations to base policy and value function on.
+    :param feature_dim: (int) Dimension of the feature vector (can be the output of a CNN)
     :param net_arch: ([int or dict]) The specification of the policy and value networks.
         See above for details on its formatting.
     :param activation_fn: (nn.Module) The activation function to use for the networks.
@@ -185,7 +185,7 @@ class PPOPolicy(BasePolicy):
         action, _ = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
         return action.detach().cpu().numpy()
 
-    def get_policy_stats(self, obs, action, deterministic=False):
+    def evaluate_actions(self, obs, action, deterministic=False):
         latent_pi, latent_vf = self._get_latent(obs)
         _, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
         log_prob = action_distribution.log_prob(action)
diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py
index f3a738a..771b5d1 100644
--- a/torchy_baselines/ppo/ppo.py
+++ b/torchy_baselines/ppo/ppo.py
@@ -204,7 +204,7 @@ class PPO(BaseRLModel):
                     # Convert discrete action for float to long
                     action = action.long().flatten()
 
-                values, log_prob, entropy = self.policy.get_policy_stats(obs, action)
+                values, log_prob, entropy = self.policy.evaluate_actions(obs, action)
                 values = values.flatten()
                 # Normalize advantage
                 advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8)
@@ -241,7 +241,7 @@ class PPO(BaseRLModel):
                 approx_kl_divs.append(th.mean(old_log_prob - log_prob).detach().cpu().numpy())
 
             if self.target_kl is not None and np.mean(approx_kl_divs) > 1.5 * self.target_kl:
-                print("Early stopping at step {} due to reaching max kl: {:.2f}".format(it, np.mean(approx_kl_divs)))
+                print("Early stopping at step {} due to reaching max kl: {:.2f}".format(gradient_step, np.mean(approx_kl_divs)))
                 break
 
         explained_var = explained_variance(self.rollout_buffer.returns.flatten().cpu().numpy(),