diff --git a/pyrobolearn/algos/evaluator.py b/pyrobolearn/algos/evaluator.py
index 288f232..7b79bff 100644
--- a/pyrobolearn/algos/evaluator.py
+++ b/pyrobolearn/algos/evaluator.py
@@ -77,7 +77,10 @@ class Evaluator(object):
         Evaluate the trajectories performed by the policy.
 
         Args:
-            verbose (bool): If true, print information on the standard output.
+            verbose (int, bool): verbose level, select between {0=False, 1=True, 2}. If 1 or 2, it will print
+                information about the evaluation process. The level 2 will print more detailed information. Do not use
+                it when the states / actions are big or high dimensional, as it could be very hard to make sense of
+                the data.
         """
         if self.estimator is not None:
             if verbose:
@@ -87,30 +90,34 @@ class Evaluator(object):
             # compute the returns
             returns = self.estimator.evaluate(self.storage)
 
-            if verbose:
+            if verbose > 1:
                 # print("Returns: {}".format(returns))
 
                 print("\nFinal storage status: ")
                 states = self.storage['states'][0]
+                num_step, num_traj = states.shape[:2]
                 states = states.view(-1, *states.size()[2:])
-                print("states: {}".format(torch.cat((torch.arange(len(states), dtype=torch.float).view(-1, 1),
-                                                       states), dim=1)))
+                print("states: {}".format(torch.cat((torch.Tensor(list(range(num_step)) * num_traj).view(-1, 1),
+                                                     states), dim=1)))
                 actions = self.storage['actions'][0]
                 actions = actions.view(-1, *actions.size()[2:])
-                print("actions: {}".format(torch.cat((torch.arange(len(actions), dtype=torch.float).view(-1, 1),
-                                                        actions), dim=1)))
+                print("actions: {}".format(torch.cat((torch.Tensor(list(range(num_step - 1)) * num_traj).view(-1, 1),
+                                                      actions), dim=1)))
                 rewards = self.storage['rewards'][:, :, 0]
                 print("rewards: {}".format(torch.cat((torch.arange(len(rewards), dtype=torch.float).view(-1, 1),
-                                                        rewards), dim=1)))
+                                                      rewards), dim=1)))
                 masks = self.storage['masks'][:, :, 0]
                 print("masks: {}".format(torch.cat((torch.arange(len(masks), dtype=torch.float).view(-1, 1),
-                                                        masks), dim=1)))
+                                                    masks), dim=1)))
                 returns = self.storage[self.estimator][:, :, 0]
                 print("returns: {}".format(torch.cat((torch.arange(len(returns), dtype=torch.float).view(-1, 1),
-                                                        returns), dim=1)))
+                                                      returns), dim=1)))
 
                 print("\n#### End of the Evaluation phase ####")
 
+            elif verbose:
+                print("#### End of the Evaluation phase ####")
+
     #############
     # Operators #
     #############
diff --git a/pyrobolearn/algos/explorer.py b/pyrobolearn/algos/explorer.py
index ba63cbf..c4cff91 100644
--- a/pyrobolearn/algos/explorer.py
+++ b/pyrobolearn/algos/explorer.py
@@ -20,6 +20,7 @@ from pyrobolearn.exploration import Exploration
 from pyrobolearn.storages import RolloutStorage, ExperienceReplay
 
 from pyrobolearn import logger
+from pyrobolearn.metrics import Metric
 
 
 __author__ = "Brian Delhaisse"
@@ -46,7 +47,7 @@ class Explorer(object):
     storage unit.
     """
 
-    def __init__(self, task, explorer, storage, num_workers=1, backend='multiprocessing'):
+    def __init__(self, task, explorer, storage, num_workers=1, backend='multiprocessing', metrics=None):
         """
         Initialize the exploration phase.
 
@@ -64,6 +65,7 @@ class Explorer(object):
                 PyTorch has been built from source with MPI support). For more information, we refer the reader to
                 references [2,4]. If the backend is 'mpi', you have to run the code using the following command:
                 `mpirun -n 4 python <code>.py`.
+            metrics ((list of) Metric, None): metrics that are used to evaluate the algorithm.
 
         References:
             [1] Multiprocessing best practices: https://pytorch.org/docs/stable/notes/multiprocessing.html
@@ -74,6 +76,7 @@ class Explorer(object):
         self.task = task
         self.explorer = explorer
         self.storage = storage
+        self.metrics = metrics
 
         # check the number of workers
         if not isinstance(num_workers, (int, long)):
@@ -214,6 +217,31 @@ class Explorer(object):
                             "instead got: {}".format(type(storage)))
         self._storage = storage
 
+    @property
+    def metrics(self):
+        """Return the metric instances."""
+        return self._metrics
+
+    @metrics.setter
+    def metrics(self, metrics):
+        """Set the metrics."""
+        # check metrics type
+        if metrics is None:
+            metrics = []
+        elif isinstance(metrics, Metric):
+            metrics = [metrics]
+        elif not isinstance(metrics, list):
+            raise TypeError("Expecting the given 'metrics' to be an instance of `Metric` or a list of `Metric`, but "
+                            "got instead: {}".format(type(metrics)))
+
+        # check each metric type
+        for i, metric in enumerate(metrics):
+            if not isinstance(metric, Metric):
+                raise TypeError("The {}th metric is not an instance of `Metric`, but: {}".format(i, type(metric)))
+
+        # set metrics
+        self._metrics = metrics
+
     ###########
     # Methods #
     ###########
@@ -523,7 +551,10 @@ class Explorer(object):
             num_rollouts (int): number of trajectories/rollouts (only valid in the on-policy case).
             deterministic (bool): if deterministic is True, then it does not explore in the environment.
             render (bool): if we should render the environment.
-            verbose (bool): If true, print information on the standard output.
+            verbose (int, bool): verbose level, select between {0=False, 1=True, 2}. If 1 or 2, it will print
+                information about the exploration process. The level 2 will print more detailed information. Do not use
+                it when the states / actions are big or high dimensional, as it could be very hard to make sense of
+                the data.
 
         Returns:
             DictStorage: updated memory storage
diff --git a/pyrobolearn/algos/rl_algo.py b/pyrobolearn/algos/rl_algo.py
index 46f6e69..4b249eb 100755
--- a/pyrobolearn/algos/rl_algo.py
+++ b/pyrobolearn/algos/rl_algo.py
@@ -15,6 +15,8 @@ from pyrobolearn.algos.explorer import Explorer
 from pyrobolearn.algos.evaluator import Evaluator
 from pyrobolearn.algos.updater import Updater
 
+from pyrobolearn.metrics import Metric
+
 
 __author__ = "Brian Delhaisse"
 __copyright__ = "Copyright 2018, PyRoboLearn"
@@ -148,7 +150,7 @@ class RLAlgo(object):  # Algo):
         [5] OpenAI - Spinning Up: https://spinningup.openai.com/
     """
 
-    def __init__(self, explorer, evaluator, updater, dynamic_model=None):
+    def __init__(self, explorer, evaluator, updater, dynamic_model=None, metrics=None):
         """
         Initialize the reinforcement learning algorithm.
 
@@ -157,6 +159,7 @@ class RLAlgo(object):  # Algo):
             evaluator (Evaluator): evaluate the actions
             updater (Updater): update the approximators (rl, value-functions,...)
             dynamic_model (None): dynamical model
+            metrics ((list of) Metric, None): metrics that are used to evaluate the algorithm.
         """
 
         super(RLAlgo, self).__init__()
@@ -172,6 +175,8 @@ class RLAlgo(object):  # Algo):
         self.best_reward = -np.infty
         self.best_parameters = None
 
+        self.metrics = metrics
+
     ##############
     # Properties #
     ##############
@@ -265,6 +270,31 @@ class RLAlgo(object):  # Algo):
         """Return the losses."""
         return self.updater.losses
 
+    @property
+    def metrics(self):
+        """Return the metric instances."""
+        return self._metrics
+
+    @metrics.setter
+    def metrics(self, metrics):
+        """Set the metrics."""
+        # check metrics type
+        if metrics is None:
+            metrics = []
+        elif isinstance(metrics, Metric):
+            metrics = [metrics]
+        elif not isinstance(metrics, list):
+            raise TypeError("Expecting the given 'metrics' to be an instance of `Metric` or a list of `Metric`, but "
+                            "got instead: {}".format(type(metrics)))
+
+        # check each metric type
+        for i, metric in enumerate(metrics):
+            if not isinstance(metric, Metric):
+                raise TypeError("The {}th metric is not an instance of `Metric`, but: {}".format(i, type(metric)))
+
+        # set metrics
+        self._metrics = metrics
+
     ###########
     # Methods #
     ###########
@@ -298,7 +328,7 @@ class RLAlgo(object):  # Algo):
     #     self.evaluator = evaluator
     #     self.updater = updater
 
-    def train(self, num_steps, num_rollouts=1, num_episodes=1, verbose=False, seed=None):
+    def train(self, num_steps, num_rollouts=1, num_episodes=1, verbose=0, seed=None):
         """
         Train the policy in the provided environment.
 
@@ -306,11 +336,14 @@ class RLAlgo(object):  # Algo):
             num_steps (int): number of step per rollout/trajectory
             num_rollouts (int): number of rollouts/trajectories per episode (default: 1)
             num_episodes (int): number of episodes (default: 1)
-            verbose (bool): if True, print details about the optimization process
+            verbose (int, bool): verbose level, select between {0=False, 1=True, 2}. If 1 or 2, it will print
+                information about the training process. The level 2 will print more detailed information. Do not use
+                it when the states / actions are big or high dimensional, as it could be very hard to make sense of
+                the data.
             seed (int): random seed
 
         Returns:
-            dict: history
+           Metric, list of Metric: metric instance(s).
         """
         history = {}
 
@@ -351,6 +384,10 @@ class RLAlgo(object):  # Algo):
             # 3. update
             losses = self.updater.update(verbose=verbose)
 
+            # compute metrics
+            for metric in self.metrics:
+                metric.end_episode_update(episode_idx=episode, num_episodes=num_episodes)
+
             # add the loss in the history
             history.setdefault('losses', []).append(losses)
 
@@ -362,7 +399,10 @@ class RLAlgo(object):  # Algo):
         if verbose:
             print("\n#### End of the RL algo ####")
 
-        return history
+        # return history
+        if len(self.metrics) == 1:
+            return self.metrics[0]
+        return self.metrics
 
     def test(self, num_steps, dt=0., use_terminating_condition=False, render=True):  # , storage):
         """
@@ -409,13 +449,24 @@ class GradientRLAlgo(RLAlgo):
         TD residual,...)
     """
 
-    def __init__(self, explorer, evaluator, updater, dynamic_model=None):  # hyperparameters=None)
-        super(GradientRLAlgo, self).__init__(explorer, evaluator, updater, dynamic_model)
+    def __init__(self, explorer, evaluator, updater, dynamic_model=None, metrics=None):  # hyperparameters=None)
+        """
+        Initialize the gradient reinforcement learning algorithm.
+
+        Args:
+            explorer (Explorer): explorer that specifies how to explore in the environment
+            evaluator (Evaluator): evaluate the actions
+            updater (Updater): update the approximators (rl, value-functions,...)
+            dynamic_model (None): dynamical model
+            metrics ((list of) Metric, None): metrics that are used to evaluate the algorithm.
+        """
+        super(GradientRLAlgo, self).__init__(explorer, evaluator, updater, dynamic_model=dynamic_model, metrics=metrics)
 
 
 class EMRLAlgo(RLAlgo):
     r"""Expectation-Maximization reinforcement learning algorithm.
     """
 
-    def __init__(self, task, exploration_strategy, storage, dynamic_model=None):  # hyperparameters=None)
-        super(EMRLAlgo, self).__init__(task, exploration_strategy, storage, dynamic_model)
+    def __init__(self, task, exploration_strategy, storage, dynamic_model=None, metrics=None):  # hyperparameters=None)
+        super(EMRLAlgo, self).__init__(task, exploration_strategy, storage, dynamic_model=dynamic_model,
+                                       metrics=metrics)
diff --git a/pyrobolearn/algos/updater.py b/pyrobolearn/algos/updater.py
index 8ab97ec..b4e1d06 100644
--- a/pyrobolearn/algos/updater.py
+++ b/pyrobolearn/algos/updater.py
@@ -2,13 +2,13 @@
 """Provide the Updater class used in the third and final step of RL algorithms
 
 The updater update the approximator (such as the policy and/or value function) parameters based on the loss, and
-using the specified optmizer.
+using the specified optimizer.
 
 Dependencies:
 - `pyrobolearn/approximators`: models (which contain parameters to update)
 - `pyrobolearn/losses`: to compute the loss
 - `pyrobolearn/optimizers`: the optimizers used to update the model parameters
-- `pyrobolearn/samplers`:
+- `pyrobolearn/samplers`: to sample from batches
 """
 
 import collections
@@ -27,6 +27,8 @@ from pyrobolearn.samplers import StorageSampler
 from pyrobolearn.returns import Return, Target, Evaluator
 from pyrobolearn.parameters.updater import ParameterUpdater
 
+from pyrobolearn.metrics import Metric
+
 
 __author__ = "Brian Delhaisse"
 __copyright__ = "Copyright 2018, PyRoboLearn"
@@ -50,7 +52,8 @@ class Updater(object):
     This class focuses on the third step of RL algorithms.
     """
 
-    def __init__(self, approximators, sampler, losses, optimizers, evaluators=None, updaters=None, ticks=None):
+    def __init__(self, approximators, sampler, losses, optimizers, evaluators=None, updaters=None, ticks=None,
+                 metrics=None):
         """
         Initialize the update phase.
 
@@ -67,6 +70,7 @@ class Updater(object):
             ticks (None, dictionary): dictionary containing as the key (updater or loss) and the value are the number
                 of time steps to wait before updating the corresponding key. By default, it will evaluate the given
                 losses and updaters at each time step.
+            metrics ((list of) Metric, None): metrics that are used to evaluate the algorithm.
         """
         self.approximators = approximators
         self.sampler = sampler
@@ -75,6 +79,7 @@ class Updater(object):
         self.evaluators = evaluators
         self.updaters = updaters
         self.ticks = ticks
+        self.metrics = metrics
 
         # counter
         self._cnt = 0
@@ -259,6 +264,31 @@ class Updater(object):
         # set the ticks
         self._ticks = ticks
 
+    @property
+    def metrics(self):
+        """Return the metric instances."""
+        return self._metrics
+
+    @metrics.setter
+    def metrics(self, metrics):
+        """Set the metrics."""
+        # check metrics type
+        if metrics is None:
+            metrics = []
+        elif isinstance(metrics, Metric):
+            metrics = [metrics]
+        elif not isinstance(metrics, list):
+            raise TypeError("Expecting the given 'metrics' to be an instance of `Metric` or a list of `Metric`, but "
+                            "got instead: {}".format(type(metrics)))
+
+        # check each metric type
+        for i, metric in enumerate(metrics):
+            if not isinstance(metric, Metric):
+                raise TypeError("The {}th metric is not an instance of `Metric`, but: {}".format(i, type(metric)))
+
+        # set metrics
+        self._metrics = metrics
+
     ###########
     # Methods #
     ###########
@@ -270,7 +300,10 @@ class Updater(object):
         Args:
             num_epochs (int): number of epochs.
             num_batches (int): number of batches.
-            verbose (bool): If true, print information on the standard output.
+            verbose (int, bool): verbose level, select between {0=False, 1=True, 2}. If 1 or 2, it will print
+                information about the update process. The level 2 will print more detailed information. Do not use
+                it when the states / actions are big or high dimensional, as it could be very hard to make sense of
+                the data.
 
         Returns:
             dict: dictionary of losses. There is a key for each loss, and the value is a nested list which contains
@@ -292,8 +325,8 @@ class Updater(object):
             for batch_idx, batch in enumerate(self.sampler):
 
                 if verbose:
-                    print("Epoch: {}/{} - Batch: {}/{} with size {}".format(epoch + 1, num_epochs, batch_idx + 1,
-                                                                            num_batches, batch.size))
+                    print("\nEpoch: {}/{} - Batch: {}/{} with size {}".format(epoch + 1, num_epochs, batch_idx + 1,
+                                                                              num_batches, batch.size))
 
                 # evaluate the evaluators with the current parameters on the given batch and save the results in the
                 # batch's `current` attribute
@@ -332,11 +365,21 @@ class Updater(object):
                                 print("\tRun updater {}".format(updater))
                             updater()
 
+                # compute metrics
+                for metric in self.metrics:
+                    metric.end_batch_update(batch_idx=batch_idx, num_batches=num_batches)
+
                 # increase counter
                 self._cnt += 1
 
-        if verbose:
+            # compute metrics
+            for metric in self.metrics:
+                metric.end_epoch_update(epoch_idx=epoch, num_epochs=num_epochs)
+
+        if verbose > 1:
             print("Losses: {}".format(losses))
+
+        if verbose:
             print("#### End of the Update phase ####")
 
         return losses  # shape=(epochs, batches)
diff --git a/pyrobolearn/losses/losses.py b/pyrobolearn/losses/losses.py
index 013d4d4..241e050 100644
--- a/pyrobolearn/losses/losses.py
+++ b/pyrobolearn/losses/losses.py
@@ -24,6 +24,12 @@ class BatchLoss(Loss):
     r"""Loss evaluated on a batch.
     """
 
+    def __init__(self):
+        super(BatchLoss, self).__init__()
+
+        # cache the last value that was computed by the loss
+        self.value = None
+
     def _compute(self, batch):
         """Compute the loss on the given batch. This method has to be implemented in the child classes."""
         raise NotImplementedError
@@ -42,7 +48,8 @@ class BatchLoss(Loss):
         if not isinstance(batch, Batch):
             raise TypeError("Expecting the given 'batch' to be an instance of `Batch`, instead got: "
                             "{}".format(type(batch)))
-        return self._compute(batch)
+        self.value = self._compute(batch)
+        return self.value
 
 
 class FixedLoss(BatchLoss):
@@ -142,9 +149,9 @@ class HuberLoss(BatchLoss):
     :math:`a = Q(s,a) - (r + \gamma \max_a Q(s',a))` where :math:`(r + \gamma \max_a Q(s',a))` is the target function.
 
     References:
-        [1] Huber Loss (on Wikipedia): https://en.wikipedia.org/wiki/Huber_loss
-        [2] "Reinforcement Learning (DQN) Tutorial":
-            https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html
+        - [1] Huber Loss (on Wikipedia): https://en.wikipedia.org/wiki/Huber_loss
+        - [2] "Reinforcement Learning (DQN) Tutorial":
+          https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html
     """
 
     def __init__(self, loss, delta=1.):
@@ -191,7 +198,7 @@ class PseudoHuberLoss(BatchLoss):
     While the above is the most common form, other smooth approximations of the Huber loss function also exist." [1]
 
     References:
-        [1] Huber Loss (on Wikipedia): https://en.wikipedia.org/wiki/Huber_loss#Pseudo-Huber_loss_function
+        - [1] Huber Loss (on Wikipedia): https://en.wikipedia.org/wiki/Huber_loss#Pseudo-Huber_loss_function
     """
 
     def __init__(self, loss, delta=1.):
diff --git a/pyrobolearn/losses/policy_losses.py b/pyrobolearn/losses/policy_losses.py
index ccf2389..7a3524f 100644
--- a/pyrobolearn/losses/policy_losses.py
+++ b/pyrobolearn/losses/policy_losses.py
@@ -35,8 +35,8 @@ class PGLoss(BatchLoss):
     .. math:: g = \mathbb{E}[ \nabla_{\theta} \log \pi_{\theta}(a_t | s_t) \psi_t ]
 
     References:
-        [1] "Proximal Policy Optimization Algorithms", Schulman et al., 2017
-        [2] "High-Dimensional Continuous Control using Generalized Advantage Estimation", Schulman et al., 2016
+        - [1] "Proximal Policy Optimization Algorithms", Schulman et al., 2017
+        - [2] "High-Dimensional Continuous Control using Generalized Advantage Estimation", Schulman et al., 2016
     """
 
     def __init__(self, estimator):
@@ -89,8 +89,8 @@ class CPILoss(BatchLoss):
     :math:`r_t(\theta) = \frac{ \pi_{\theta}(a_t|s_t) }{ \pi_{\theta_{old}}(a_t|s_t) }`.
 
     References:
-        [1] "Approximately optimal approximate reinforcement learning", Kakade et al., 2002
-        [2] "Proximal Policy Optimization Algorithms", Schulman et al., 2017
+        - [1] "Approximately optimal approximate reinforcement learning", Kakade et al., 2002
+        - [2] "Proximal Policy Optimization Algorithms", Schulman et al., 2017
     """
 
     def __init__(self, estimator):
@@ -154,7 +154,7 @@ class CLIPLoss(BatchLoss):
     :math:`r_t(\theta) = \frac{ \pi_{\theta}(a_t|s_t) }{ \pi_{\theta_{old}}(a_t|s_t) }`.
 
     References:
-        [1] "Proximal Policy Optimization Algorithms", Schulman et al., 2017
+        - [1] "Proximal Policy Optimization Algorithms", Schulman et al., 2017
     """
 
     def __init__(self, estimator, clip=0.2):
@@ -256,9 +256,9 @@ class EntropyLoss(BatchLoss):
     where :math:`H[.]` is the Shannon entropy of the given probability distribution.
 
     References:
-        [1] "Simple Statistical Gradient-following Algorithms for Connectionist Reinforcement Learning", Williams, 1992
-        [2] "Asynchronous Methods for Deep Reinforcement Learning", Mnih et al., 2016
-        [3] "Proximal Policy Optimization Algorithms", Schulman et al., 2017
+        - [1] "Simple Statistical Gradient-following Algorithms for Connectionist Reinforcement Learning", Williams, 1992
+        - [2] "Asynchronous Methods for Deep Reinforcement Learning", Mnih et al., 2016
+        - [3] "Proximal Policy Optimization Algorithms", Schulman et al., 2017
     """
 
     def __init__(self):
diff --git a/pyrobolearn/losses/value_losses.py b/pyrobolearn/losses/value_losses.py
index 7e6c021..a23c29d 100644
--- a/pyrobolearn/losses/value_losses.py
+++ b/pyrobolearn/losses/value_losses.py
@@ -137,8 +137,8 @@ class MSBELoss(BatchLoss):
     minimize them, we try to enforce the Bellman equations.
 
     References:
-        [1] https://spinningup.openai.com/en/latest/algorithms/ddpg.html
-        [2] "Reinforcement Learning: An Introduction", Sutton and Barto, 2018
+        - [1] https://spinningup.openai.com/en/latest/algorithms/ddpg.html
+        - [2] "Reinforcement Learning: An Introduction", Sutton and Barto, 2018
     """
 
     def __init__(self, td_return):
diff --git a/pyrobolearn/metrics/__init__.py b/pyrobolearn/metrics/__init__.py
index d5a3b8e..695c0c0 100644
--- a/pyrobolearn/metrics/__init__.py
+++ b/pyrobolearn/metrics/__init__.py
@@ -1,12 +1,13 @@
 
 # import the metrics
-from .metric import *
+from .metric import Metric
 
 # import reinforcement learning metrics
-from .rl_metrics import *
+from .rl_metrics import RLMetric, AverageReturnMetric, LossMetric
 
 # import imitation learning metrics
-from .il_metrics import *
+from .il_metrics import ILMetric
 
 # import transfer learning metrics
-from .tl_metrics import *
+from .tl_metrics import TLMetric, JumpstartMetric, AsymptoticPerformanceMetric, TotalRewardMetric, \
+    TransferRatioMetric, TimeToThresholdMetric
diff --git a/pyrobolearn/metrics/metric.py b/pyrobolearn/metrics/metric.py
index 5b2059b..83ba898 100644
--- a/pyrobolearn/metrics/metric.py
+++ b/pyrobolearn/metrics/metric.py
@@ -33,7 +33,8 @@ class Metric(object):
         Args:
             metrics (None, Metric, list of Metric): inner metric objects. Each metric will be plot in a subplot.
         """
-        self._metrics = metrics
+        self.metrics = metrics
+        self._data = None
 
     ##############
     # Properties #
@@ -51,16 +52,39 @@ class Metric(object):
             metrics = []
         if not isinstance(metrics, collections.Iterable):
             metrics = [metrics]
-        for i, metric in enumerate(metrics):
-            if not isinstance(metric, Metric):
-                raise TypeError("Expecting the given {}th metric to be an instance of `Metric`, instead got: "
-                                "{}".format(i, type(metric)))
+        self._check_recursively_metric_type(metrics)
         self._metrics = metrics
 
+    @property
+    def data(self):
+        """Return the data."""
+        if self.metrics:
+            return [metric.data for metric in self.metrics]
+        return self._get_data()
+
     ###########
     # Methods #
     ###########
 
+    def _check_recursively_metric_type(self, metrics):
+        """Check recursively the metric types."""
+        if isinstance(metrics, collections.Iterable):
+            for metric in metrics:
+                self._check_recursively_metric_type(metric)
+        elif not isinstance(metrics, Metric):
+            raise TypeError("Expecting the given 'metric' to be an instance of `Metric`, instead got: "
+                            "{}".format(type(metrics)))
+
+    def _get_data(self):
+        """Return the inner data."""
+        return self._data
+
+    def has_inner_metrics(self):
+        """Check if the metric has inner metrics."""
+        if self.metrics:
+            return True
+        return False
+
     def append(self, metric):
         """Append the given metric to the list of metrics."""
         if not isinstance(metric, Metric):
@@ -68,102 +92,220 @@ class Metric(object):
                             "{}".format(type(metric)))
         self.metrics.append(metric)
 
-    def update(self, *args, **kwargs):
-        """Update the metrics."""
-        pass
-
-    def step_update(self, step_idx=None):
-        """Update at each time step."""
+    def __call(self, method_name, *args, **kwargs):  # TODO do it in a recursive way
+        """Call the method for each metric with the provided arguments."""
         if self.metrics:
             for metric in self.metrics:
-                metric._step_update(step_idx=step_idx)
+                fct = getattr(metric, method_name)
+                fct(*args, **kwargs)
         else:
-            self._step_update(step_idx=step_idx)
+            fct = getattr(self, method_name)
+            fct(*args, **kwargs)
 
-    def _step_update(self, step_idx=None):
-        """Update at each time step; this has to be implemented in the child class."""
+    def update(self, *args, **kwargs):
+        """Update the metrics."""
+        self.__call('_update', *args, **kwargs)
+
+    def _update(self, *args, **kwargs):
+        """Update the metric; this has to be implemented in the child class."""
+        pass
+
+    def start_algo_update(self, algo):
+        """Update each time an algorithm starts."""
+        self.__call('_start_algo_update', algo=algo)
+
+    def _start_algo_update(self, algo):
+        """Update each time an algorithm starts; this has to be implemented in the child class."""
+        pass
+
+    def end_algo_update(self, algo):
+        """Update each time an algorithm ends."""
+        self.__call('_end_algo_update', algo=algo)
+
+    def _end_algo_update(self, algo):
+        """Update each time an algorithm ends; this has to be implemented in the child class."""
+        pass
+
+    def start_episode_update(self, episode_idx=None, num_episodes=None):
+        """Update each time an episode starts."""
+        self.__call('_start_episode_update', episode_idx=episode_idx, num_episodes=num_episodes)
+
+    def _start_episode_update(self, episode_idx=None, num_episodes=None):
+        """Update each time an episode starts; this has to be implemented in the child class."""
         pass
 
     def episode_update(self, episode_idx=None):
         """Update at each episode."""
-        if self.metrics:
-            for metric in self.metrics:
-                metric._episode_step(episode_idx=episode_idx)
-        else:
-            self._episode_update(episode_idx=episode_idx)
+        self.__call('_episode_update', episode_idx=episode_idx)
 
     def _episode_update(self, episode_idx=None):
         """Update at each episode; this has to be implemented in the child class."""
         pass
 
-    def _plot(self, ax):
-        """
-        Plot the metric in the given axis. This has to be implemented in the child classes.
-        """
+    def end_episode_update(self, episode_idx=None, num_episodes=None):
+        """Update each time an episode ends."""
+        self.__call('_end_episode_update', episode_idx=episode_idx, num_episodes=num_episodes)
+
+    def _end_episode_update(self, episode_idx=None, num_episodes=None):
+        """Update each time an episode ends."""
         pass
 
-    def plot(self, nrows=-1, ncols=-1, block=False, filename=None):
+    def start_rollout_update(self, rollout_idx=None, num_rollouts=None):
+        """Update each time a rollout starts."""
+        self.__call('_start_rollout_update', rollout_idx=rollout_idx, num_rollouts=num_rollouts)
+
+    def _start_rollout_update(self, rollout_idx=None, num_rollouts=None):
+        """Update each time a rollout starts; this has to be implemented in the child class."""
+        pass
+
+    def rollout_update(self, rollout_idx=None):
+        """Update at each rollout."""
+        self.__call('_rollout_update', rollout_idx=rollout_idx)
+
+    def _rollout_update(self, rollout_idx=None):
+        """Update at each rollout; this has to be implemented in the child class."""
+        pass
+
+    def end_rollout_update(self, rollout_idx=None, num_rollouts=None):
+        """Update each time a rollout ends."""
+        self.__call('_end_rollout_update', rollout_idx=rollout_idx, num_rollouts=num_rollouts)
+
+    def _end_rollout_update(self, rollout_idx=None, num_rollouts=None):
+        """Update each time a rollout ends; this has to be implemented in the child class."""
+        pass
+
+    def start_step_update(self, step_idx=None, num_steps=None):
+        """Update each time a step starts."""
+        self.__call('_start_step_update', step_idx=step_idx, num_steps=num_steps)
+
+    def _start_step_update(self, step_idx=None, num_steps=None):
+        """Update each time a step starts; this has to be implemented in the child class."""
+        pass
+
+    def step_update(self, step_idx=None):
+        """Update at each time step."""
+        self.__call('_step_update', step_idx=step_idx)
+
+    def _step_update(self, step_idx=None):
+        """Update at each time step; this has to be implemented in the child class."""
+        pass
+
+    def end_step_update(self, step_idx=None, num_steps=None):
+        """Update each time a step ends."""
+        self.__call('_end_step_update', step_idx=step_idx, num_steps=num_steps)
+
+    def _end_step_update(self, step_idx=None, num_steps=None):
+        """Update each time a step ends; this has to be implemented in the child class."""
+        pass
+
+    def start_epoch_update(self, epoch_idx=None, num_epochs=None):
+        """Update each time an epoch starts."""
+        self.__call('_start_epoch_update', epoch_idx=epoch_idx, num_epochs=num_epochs)
+
+    def _start_epoch_update(self, epoch_idx=None, num_epochs=None):
+        """Update each time a epoch starts; this has to be implemented in the child class."""
+        pass
+
+    def end_epoch_update(self, epoch_idx=None, num_epochs=None):
+        """Update each time a epoch ends."""
+        self.__call('_end_epoch_update', epoch_idx=epoch_idx, num_epochs=num_epochs)
+
+    def _end_epoch_update(self, epoch_idx=None, num_epochs=None):
+        """Update each time a epoch ends; this has to be implemented in the child class."""
+        pass
+
+    def start_batch_update(self, batch_idx=None, num_batches=None):
+        """Update each time a batch starts."""
+        self.__call('_start_batch_update', batch_idx=batch_idx, num_batches=num_batches)
+
+    def _start_batch_update(self, batch_idx=None, num_batches=None):
+        """Update each time a batch starts; this has to be implemented in the child class."""
+        pass
+
+    def end_batch_update(self, batch_idx=None, num_batches=None):
+        """Update each time a batch ends."""
+        self.__call('_end_batch_update', batch_idx=batch_idx, num_batches=num_batches)
+
+    def _end_batch_update(self, batch_idx=None, num_batches=None):
+        """Update each time a batch ends; this has to be implemented in the child class."""
+        pass
+
+    def plot(self, nrows=-1, ncols=-1, block=True, filename=None, _ax=None):
         """
         Plot the metric(s).
 
         Args:
             nrows (int): number of rows in the subplot.
             ncols (int): number of columns in the subplot.
-            ax (plt.Axes): axis to plot the figure.
             block (bool): if True, it will block when showing the graphs.
             filename (str, None): if a string is given, it will save the plot in the given filename.
+            _ax (plt.Axes, None): axis to plot the figure. Do not use this parameter, this is used internally in a
+                recursive way.
 
         Returns:
             matplotlib.figure.Figure: figure
             np.array of matplotlib.axes._subplots.AxesSubplot: axes
         """
-        metrics = self.metrics if self.metrics else [self]
+        if not self.metrics and _ax is not None:
+            self._plot(ax=_ax)
 
-        # get nrows and ncols
-        if nrows < 1 or ncols < 1:
-            if nrows < 1 and ncols < 1:
-                if len(metrics) <= 4:
-                    if len(metrics) <= 2:
-                        nrows = 1
-                        ncols = len(metrics)
+        else:
+            metrics = self.metrics if self.metrics else [self]
+
+            # get nrows and ncols
+            if nrows < 1 or ncols < 1:
+                if nrows < 1 and ncols < 1:
+                    if len(metrics) <= 4:
+                        if len(metrics) <= 2:
+                            nrows = 1
+                            ncols = len(metrics)
+                        else:
+                            nrows = 2
+                            ncols = int(len(metrics) / 2)
                     else:
-                        nrows = 2
-                        ncols = int(len(metrics) / 2)
-                else:
-                    ncols = 4
+                        ncols = 4
 
-            if nrows < 1:  # ncols is given
-                nrows = int(len(metrics) / ncols)
-                if len(metrics) % ncols != 0:
-                    nrows += 1
+                if nrows < 1:  # ncols is given
+                    nrows = int(len(metrics) / ncols)
+                    if len(metrics) % ncols != 0:
+                        nrows += 1
 
-            elif ncols < 1:  # nrows is given
-                ncols = int(len(metrics) / nrows)
-                if len(metrics) % nrows != 0:
-                    ncols += 1
+                elif ncols < 1:  # nrows is given
+                    ncols = int(len(metrics) / nrows)
+                    if len(metrics) % nrows != 0:
+                        ncols += 1
 
-        # get number of subplots
-        nplots = nrows * ncols
+            # get number of subplots
+            nplots = nrows * ncols
 
-        # create figure and axes
-        fig, axes = plt.subplots(nrows=nrows, ncols=ncols)
-        if not isinstance(axes, np.ndarray):
-            axes = np.array(axes)
-        axes = axes.reshape(-1)
+            # create figure and axes
+            fig, axes = plt.subplots(nrows=nrows, ncols=ncols)
+            if not isinstance(axes, np.ndarray):
+                axes = np.array(axes)
+            axes = axes.reshape(-1)
 
-        # plot each metric
-        for i, metric in enumerate(self.metrics):
-            metric._plot(ax=axes[i])
+            # plot each metric
+            for i, metric in enumerate(metrics):
+                metric.plot(_ax=axes[i])
 
-        # save figure if specified
-        if filename is not None:
-            fig.savefig(filename)
+            # tight the layout
+            fig.tight_layout()
 
-        # show plot
-        plt.show(block=block)
+            # save figure if specified
+            if filename is not None:
+                fig.savefig(filename)
 
-        # return figure and axes
-        return fig, axes
+            # show plot
+            plt.show(block=block)
+
+            # return figure and axes
+            return fig, axes
+
+    def _plot(self, ax):
+        """
+        Plot the metric in the given axis. This has to be implemented in the child classes.
+        """
+        pass
 
     #############
     # Operators #
@@ -178,15 +320,21 @@ class Metric(object):
     def __str__(self):
         """Return a string describing the object."""
         if self.metrics:
-            return ' + '.join(self.metrics)
+            return ' + '.join([str(metric) for metric in self.metrics])
         return self.__class__.__name__
 
     def __add__(self, other):
-        """Add two sets of metrics together."""
+        """Add two sets of metrics together; they will be in the same figure but in different subplots."""
         if not isinstance(other, Metric):
             raise TypeError("Expecting the given other metric to be an instance of `Metric`, but got instead: "
                             "{}".format(type(other)))
-        return Metric(metrics=self.metrics + other.metrics)
+        if self.metrics:
+            if other.metrics:
+                return Metric(metrics=self.metrics + other.metrics)
+            return Metric(metrics=self.metrics + [other])
+        if other.metrics:
+            return Metric(metrics=[self] + other.metrics)
+        return Metric(metrics=[self, other])
 
     def __radd__(self, other):
         """Add two sets of metrics together."""
diff --git a/pyrobolearn/metrics/rl_metrics.py b/pyrobolearn/metrics/rl_metrics.py
index 8d43739..4785974 100644
--- a/pyrobolearn/metrics/rl_metrics.py
+++ b/pyrobolearn/metrics/rl_metrics.py
@@ -6,12 +6,12 @@ import numpy as np
 import matplotlib.pyplot as plt
 
 from pyrobolearn.tasks import RLTask
-from pyrobolearn.algos import RLAlgo
 from pyrobolearn.metrics import Metric
 from pyrobolearn.losses import BatchLoss
 
+
 __author__ = "Brian Delhaisse"
-__copyright__ = "Copyright 2018, PyRoboLearn"
+__copyright__ = "Copyright 2019, PyRoboLearn"
 __credits__ = ["Brian Delhaisse"]
 __license__ = "GNU GPLv3"
 __version__ = "1.0.0"
@@ -26,7 +26,7 @@ class RLMetric(Metric):
     Metrics used in reinforcement learning.
 
     References:
-        [1] "Deep Reinforcement Learning that Matters", Henderson et al., 2018
+        - [1] "Deep Reinforcement Learning that Matters", Henderson et al., 2018
     """
 
     def __init__(self):
@@ -44,9 +44,12 @@ class AverageReturnMetric(RLMetric):
     .. math:: J(\pi_{\theta}) = \mathcal{E}_{\tau \sim \pi_{\theta}}[ R(\tau) ]
 
     where :math:`R(\tau) = \sum_{t=0}^T \gamma^t r_t` is the discounted return.
+
+    References:
+        - [1] "Deep Reinforcement Learning that Matters", Henderson et al., 2018
     """
 
-    def __init__(self, task, gamma=1., num_episodes=10, num_steps=100):
+    def __init__(self, task, gamma=1., num_rollouts=10, num_steps=100):
         """
         Initialize the average return metric.
 
@@ -58,7 +61,7 @@ class AverageReturnMetric(RLMetric):
         self.task = task
         self.returns = []
         self._num_steps = num_steps
-        self._num_episodes = num_episodes
+        self._num_rollouts = num_rollouts
 
     ##############
     # Properties #
@@ -96,40 +99,84 @@ class AverageReturnMetric(RLMetric):
     # Methods #
     ###########
 
-    def _episode_update(self, episode_idx=None):
-        """Update the metric."""
+    def _get_data(self):
+        """Return the average rewards."""
+        return self.returns
+
+    def _end_episode_update(self, episode_idx=None, num_episodes=None):
+        """Update the metric at the end of an episode."""
         rewards = []
-        for ep in range(self._num_episodes):
+        for _ in range(self._num_rollouts):
             reward = self.task.run(num_steps=self._num_steps)
             rewards.append(reward)
 
         rewards = np.asarray(rewards).mean()
+
+        print("\nAverage return metric: {}".format(rewards))
+
         self.returns.append(rewards)
 
     def _plot(self, ax):
         """
         Plot the average return metric in the given axis.
         """
-        ax.set_title('Average Return per iteration')  # per epoch, per iteration=epoch*batch
-        ax.set_xlabel('iterations')
+        ax.set_title('Average Return per episode')  # per epoch, per iteration=epoch*batch
+        ax.set_xlabel('episodes')
         ax.set_ylabel('Average return')
         ax.plot(self.returns)
+        ax.set_ylim(bottom=0)
 
 
 class LossMetric(RLMetric):
     r"""Loss Metric
+
+    This provides the loss value with respect to the number of episodes/iterations.
+
+    Warnings: As described in [1] and copied-pasted here for completeness:
+
+    This loss metric should not be confused with "a loss function in the typical sense from supervised learning. There
+    are two main differences from standard loss functions.
+
+    1. The data distribution depends on the parameters. A loss function is usually defined on a fixed data distribution
+      which is independent of the parameters we aim to optimize. Not so here, where the data must be sampled on the
+      most recent policy.
+
+    2. It doesn't measure performance. A loss function usually evaluates the performance metric that we care about.
+      Here, we care about expected return, :math:`J(\pi_{\theta})`, but our 'loss' function does not approximate this
+      at all, even in expectation. This 'loss' function is only useful to us because, when evaluated at the current
+      parameters, with data generated by the current parameters, it has the negative gradient of performance.
+
+    But after that first step of gradient descent, there is no more connection to performance. This means that
+    minimizing this 'loss' function, for a given batch of data, has no guarantee whatsoever of improving expected
+    return. You can send this loss to :math:`-\infty` and policy performance could crater; in fact, it usually will.
+    Sometimes a deep RL researcher might describe this outcome as the policy 'overfitting' to a batch of data.
+    This is descriptive, but should not be taken literally because it does not refer to generalization error.
+
+    We raise this point because it is common for ML practitioners to interpret a loss function as a useful signal
+    during training - 'if the loss goes down, all is well.' In policy gradients, this intuition is wrong, and you
+    should only care about average return. The loss function means nothing." [1]
+
+    References:
+        - [1] https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html
     """
 
-    def __init__(self, loss):
+    def __init__(self, loss, wrt='iteration'):
         """
         Initialize the loss metric.
 
         Args:
             loss (BatchLoss): batch loss.
+            wrt (str): string that specify with respect to what we want to plot; the number of episodes, the number of
+                epochs (episodes * epochs), or the total number of iterations (episodes * epochs * batches). Select
+                between {'episode', 'epoch', 'batch'/'iteration'}. If set to something else, it will be set to
+                'iteration' by default.
         """
         super(LossMetric, self).__init__()
         self.loss = loss
+        self.batch_losses = []
+        self.epoch_losses = []
         self.losses = []
+        self.wrt = wrt
 
     ##############
     # Properties #
@@ -148,17 +195,132 @@ class LossMetric(RLMetric):
                             "{}".format(type(loss)))
         self._loss = loss
 
+    @property
+    def wrt(self):
+        """Return with respect to what we plot the loss."""
+        return self._wrt
+
+    @wrt.setter
+    def wrt(self, wrt):
+        """Set with respect to what we plot the loss."""
+        if wrt is None:
+            wrt = 'iteration'
+        else:
+            wrt = wrt.lower()
+            if wrt[:7] == 'episode':
+                wrt = 'episode'
+            elif wrt[:5] == 'epoch':
+                wrt = 'epoch'
+            else:
+                wrt = 'iteration'
+        self._wrt = wrt
+
     ###########
     # Methods #
     ###########
 
-    def update(self):
-        pass
+    def _get_data(self):
+        """Return the losses."""
+        return self.losses
+
+    def _end_batch_update(self, batch_idx=None, num_batches=None):
+        """Update the metric at the end of a batch."""
+        # print("Adding loss value: {}, {}, {}".format(self.loss.value, self.loss.value.detach(),
+        #                                              self.loss.value.item()))
+        self.batch_losses.append(self.loss.value.item())
+
+    def _end_epoch_update(self, epoch_idx=None, num_epochs=None):
+        """Update the metric at the end of an epoch."""
+        self.epoch_losses.append(self.batch_losses)
+        self.batch_losses = []
+
+    def _end_episode_update(self, episode_idx=None, num_episodes=None):
+        """Update the metric at the end of an episode."""
+        self.losses.append(self.epoch_losses)
+        self.epoch_losses = []
 
     def _plot(self, ax):
         """
         Plot the loss in the given axis.
         """
-        ax.set_title(self.loss.__class__.__name__ + ' per iteration')
-        ax.set_xlabel('iterations')
+        ax.set_title(self.loss.__class__.__name__ + ' per ' + self.wrt)
+        ax.set_xlabel(self.wrt + 's')
         ax.set_ylabel('Loss')
+
+        losses = np.asarray(self.losses)  # (num_episode, num_epoch, num_batch)
+
+        if self.wrt == 'iteration':  # iteration / batch
+            losses = losses.reshape(-1)
+        elif self.wrt == 'epoch':  # epoch
+            losses = losses.mean(axis=2).reshape(-1)
+        elif self.wrt == 'episode':  # episode
+            losses = losses.mean(axis=2).mean(axis=1)
+
+        ax.plot(losses)
+
+
+class NumberOfIterationsMetric(RLMetric):
+    r"""Number of iterations metric.
+
+    This metric checked the number of iterations it took to solve the RL task, or achieved a certain performance
+    threshold level.
+
+    References:
+        - [1] "Deep Reinforcement Learning that Matters", Henderson et al., 2018
+    """
+
+    def __init__(self, threshold=None):
+        """
+        Initialize the number of iterations metric.
+
+        Args:
+            (float, None): desired performance threshold level.
+        """
+        super(NumberOfIterationsMetric, self).__init__()
+
+
+class MaxAverageReturnMetric(RLMetric):
+    r"""Max average return metric.
+
+    This metric checked the max average return per episode / iteration.
+
+    References:
+        - [1] "Deep Reinforcement Learning that Matters", Henderson et al., 2018
+    """
+
+    def __init__(self):
+        """
+        Initialize the max average return metric.
+        """
+        super(MaxAverageReturnMetric, self).__init__()
+
+
+class ExplorationMetric(RLMetric):
+    r"""Exploration metric.
+
+    This measures how much an agent explores in an environment as the number of episodes / iterations increases. This
+    is useful for instance if the policy (given the state) predicts as well the variance / covariance of its actions.
+    We can use the entropy to measure the uncertainty on each action.
+    """
+
+    def __init__(self):
+        """
+        Initialize the exploration metric.
+        """
+        super(ExplorationMetric, self).__init__()
+
+
+class KLMetric(RLMetric):
+    r"""KL divergence metric.
+
+    This measures how much the distance between two distributions decreases as the number of episodes / iterations
+    increases. For instance, it can be used to check how distant is the learned policy :math:`\pi_\theta(\cdot | s)`
+    from the optimal one :math:`\pi^*(\cdot | s)`, or how distant is the learned dynamic transition probability
+    function :math:`p(\cdot | s, a)` from the true one (or an approximation of it) :math:`p^*(\cdot | s, a)`.
+    """
+
+    def __init__(self):
+        """
+        Initialize the uncertainty metric.
+        """
+        super(KLMetric, self).__init__()
diff --git a/pyrobolearn/metrics/tl_metrics.py b/pyrobolearn/metrics/tl_metrics.py
index 6458bb7..2bd8075 100644
--- a/pyrobolearn/metrics/tl_metrics.py
+++ b/pyrobolearn/metrics/tl_metrics.py
@@ -55,7 +55,7 @@ class JumpstartMetric(TLMetric):
         super(JumpstartMetric, self).__init__()
 
 
-class AsymptoticPerformance(TLMetric):
+class AsymptoticPerformanceMetric(TLMetric):
     r"""Asymptotic performance metric
 
     The asymptotic performance metric measures how much the final learned performance of an agent in the target task
@@ -69,7 +69,7 @@ class AsymptoticPerformance(TLMetric):
         """
         Initialize the asymptotic performance metric.
         """
-        super(AsymptoticPerformance, self).__init__()
+        super(AsymptoticPerformanceMetric, self).__init__()
 
 
 class TotalRewardMetric(TLMetric):
diff --git a/pyrobolearn/optimizers/torch_optimizer.py b/pyrobolearn/optimizers/torch_optimizer.py
index 73ab4a1..bc2c123 100644
--- a/pyrobolearn/optimizers/torch_optimizer.py
+++ b/pyrobolearn/optimizers/torch_optimizer.py
@@ -44,12 +44,31 @@ __status__ = "Development"
 class Adam(Optimizer):
     r"""Adam Optimizer
 
+    This provides a wrapper around the Adam optimizer [1, 2] where the parameters can be passed at a later stage.
+    The documentation is taken from [2].
+
     References:
-        [1] "Adam: A Method for Stochastic Optimization", Kingma et al., 2014
+        - [1] "Adam: A Method for Stochastic Optimization", Kingma et al., 2014
+        - [2] torch.optim: https://pytorch.org/docs/stable/optim.html
     """
 
     def __init__(self, learning_rate=1e-3, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False,
-                 max_grad_norm=None, *args, **kwargs):  # 0.5
+                 max_grad_norm=None, verbose=0, *args, **kwargs):  # 0.5
+        """
+        Initialize the Adam optimizer.
+
+        Args:
+            learning_rate (float): learning rate.
+            betas (tuple[float, float]): coefficients used for computing running averages of gradient and its square.
+            eps (float): term added to the denominator to improve numerical stability.
+            weight_decay (float): weight decay (L2 penalty).
+            amsgrad (bool): whether to use the AMSGrad variant of this algorithm from the paper "On the Convergence of
+                Adam and Beyond".
+            max_grad_norm (float, None): the allowed maximum gradient norm. If provided, it will clip the gradients.
+            verbose (bool, int): if verbose=2, it will print the gradient norm for each parameter.
+            *args (list): list of arguments that are given to the parent class `Optimizer`.
+            **kwargs (dict): dictionary of arguments that are given to the parent class `Optimizer`.
+        """
         super(Adam, self).__init__(*args, **kwargs)
         self.optimizer = None
         self.learning_rate = learning_rate
@@ -58,11 +77,20 @@ class Adam(Optimizer):
         self.weight_decay = weight_decay
         self.amsgrad = amsgrad
         self.max_grad_norm = max_grad_norm
+        self.verbose = verbose
 
     def reset(self):
+        """Reset the optimizer."""
         self.optimizer = None
 
     def optimize(self, params, loss):
+        """
+        Optimize the given parameters on the specified loss.
+
+        Args:
+            params (list of torch.Tensor): model parameters.
+            loss (torch.Tensor): loss value that were computed using the model parameters.
+        """
         # create optimizer if necessary
         if self.optimizer is None:
             self.optimizer = optim.Adam(params, lr=self.learning_rate, betas=self.betas, eps=self.eps,
@@ -73,17 +101,44 @@ class Adam(Optimizer):
         loss.backward(retain_graph=True)
         if self.max_grad_norm is not None:
             nn.utils.clip_grad_norm_(params, self.max_grad_norm)
+
+        if self.verbose > 1:
+            total_norm = 0
+            for p in params:
+                param_norm = p.grad.data.norm(2)
+                total_norm += param_norm.item() ** 2  # item() convert a one element tensor to a scalar
+                print("{} - grad norm = {}".format(p, 1./2 * param_norm.item()))
+            total_norm = total_norm ** (1. / 2)
+            print("Total gradient norm: {}".format(total_norm))
+
         self.optimizer.step()
 
 
 class Adadelta(Optimizer):
     r"""Adadelta Optimizer
 
+    This provides a wrapper around the Adadelta optimizer [1, 2] where the parameters can be passed at a later stage.
+    The documentation is taken from [2].
+
     References:
-        [1] "ADADELTA: An Adaptive Learning Rate Method", Zeiler, 2012
+        - [1] "ADADELTA: An Adaptive Learning Rate Method", Zeiler, 2012
+        - [2] torch.optim: https://pytorch.org/docs/stable/optim.html
     """
 
     def __init__(self, learning_rate=1., rho=0.9, eps=1e-6, weight_decay=0, max_grad_norm=None, *args, **kwargs):  # 0.5
+        """
+        Initialize the Adadelta optimizer.
+
+        Args:
+            learning_rate (float): learning rate.
+            rho (float): coefficient used for computing a running average of squared gradients.
+            eps (float): term added to the denominator to improve numerical stability.
+            weight_decay (float): weight decay (L2 penalty).
+            max_grad_norm (float, None): the allowed maximum gradient norm. If provided, it will clip the gradients.
+            verbose (bool, int): if verbose=2, it will print the gradient norm for each parameter.
+            *args (list): list of arguments that are given to the parent class `Optimizer`.
+            **kwargs (dict): dictionary of arguments that are given to the parent class `Optimizer`.
+        """
         super(Adadelta, self).__init__(*args, **kwargs)
         self.optimizer = None
         self.learning_rate = learning_rate
@@ -93,6 +148,13 @@ class Adadelta(Optimizer):
         self.max_grad_norm = max_grad_norm
 
     def optimize(self, params, loss):
+        """
+        Optimize the given parameters on the specified loss.
+
+        Args:
+            params (list of torch.Tensor): model parameters.
+            loss (torch.Tensor): loss value that were computed using the model parameters.
+        """
         if self.optimizer is None:
             self.optimizer = optim.Adadelta(params, lr=self.learning_rate, rho=self.rho, eps=self.eps,
                                             weight_decay=self.weight_decay)
@@ -108,12 +170,29 @@ class Adadelta(Optimizer):
 class Adagrad(Optimizer):
     r"""Adagrad Optimizer
 
+    This provides a wrapper around the Adagrad optimizer [1, 2] where the parameters can be passed at a later stage.
+    The documentation is taken from [2].
+
     References:
-        [1] "Adaptive Subgradient Methods for Online Learning and Stochastic Optimization", Duchi et al., 2011
+        - [1] "Adaptive Subgradient Methods for Online Learning and Stochastic Optimization", Duchi et al., 2011
+        - [2] torch.optim: https://pytorch.org/docs/stable/optim.html
     """
 
     def __init__(self, learning_rate=0.01, learning_rate_decay=0, weight_decay=0, initial_accumumaltor_value=0,
                  max_grad_norm=None, *args, **kwargs):  # 0.5
+        """
+        Initialize the Adagrad optimizer.
+
+        Args:
+            learning_rate (float): learning rate.
+            learning_rate_decay (float): learning rate decay.
+            weight_decay (float): weight decay (L2 penalty).
+            initial_accumulator_value (float): the initial accumulator value.
+            max_grad_norm (float, None): the allowed maximum gradient norm. If provided, it will clip the gradients.
+            verbose (bool, int): if verbose=2, it will print the gradient norm for each parameter.
+            *args (list): list of arguments that are given to the parent class `Optimizer`.
+            **kwargs (dict): dictionary of arguments that are given to the parent class `Optimizer`.
+        """
         super(Adagrad, self).__init__(*args, **kwargs)
         self.optimizer = None
         self.learning_rate = learning_rate
@@ -123,6 +202,13 @@ class Adagrad(Optimizer):
         self.max_grad_norm = max_grad_norm
 
     def optimize(self, params, loss):
+        """
+        Optimize the given parameters on the specified loss.
+
+        Args:
+            params (list of torch.Tensor): model parameters.
+            loss (torch.Tensor): loss value that were computed using the model parameters.
+        """
         if self.optimizer is None:
             self.optimizer = optim.Adagrad(params, lr=self.learning_rate, lr_decay=self.learning_rate_decay,
                                            weight_decay=self.weight_decay,
@@ -139,14 +225,34 @@ class Adagrad(Optimizer):
 class RMSprop(Optimizer):
     r"""RMSprop
 
+    This provides a wrapper around the RMSprop optimizer [1, 2, 3] where the parameters can be passed at a later stage.
+    The documentation is taken from [3].
+
     References:
-        [1] "RMSprop:  Divide the gradient by a running average of its recent magnitude" (lecture 6.5), Tieleman and
+        - [1] "RMSprop:  Divide the gradient by a running average of its recent magnitude" (lecture 6.5), Tieleman and
             Hinton, 2012
-        [2] "Generating Sequences With Recurrent Neural Networks", Graves, 2014
+        - [2] "Generating Sequences With Recurrent Neural Networks", Graves, 2014
+        - [3] torch.optim: https://pytorch.org/docs/stable/optim.html
     """
 
     def __init__(self, learning_rate=1e-2, alpha=0.99, eps=1e-8, weight_decay=0, momentum=0, centered=False,
                  max_grad_norm=None, *args, **kwargs):  # 0.5
+        """
+        Initialize the RMSprop optimizer.
+
+        Args:
+            learning_rate (float): learning rate.
+            alpha (float): smoothing constant.
+            eps (float): term added to the denominator to improve numerical stability.
+            weight_decay (float): weight decay (L2 penalty).
+            momentum (float): momentum factor.
+            centered (bool): if True, compute the centered RMSProp, the gradient is normalized by an estimation of
+                its variance.
+            max_grad_norm (float, None): the allowed maximum gradient norm. If provided, it will clip the gradients.
+            verbose (bool, int): if verbose=2, it will print the gradient norm for each parameter.
+            *args (list): list of arguments that are given to the parent class `Optimizer`.
+            **kwargs (dict): dictionary of arguments that are given to the parent class `Optimizer`.
+        """
         super(RMSprop, self).__init__(*args, **kwargs)
         self.optimizer = None
         self.learning_rate = learning_rate
@@ -158,6 +264,13 @@ class RMSprop(Optimizer):
         self.max_grad_norm = max_grad_norm
 
     def optimize(self, params, loss):
+        """
+        Optimize the given parameters on the specified loss.
+
+        Args:
+            params (list of torch.Tensor): model parameters.
+            loss (torch.Tensor): loss value that were computed using the model parameters.
+        """
         if self.optimizer is None:
             self.optimizer = optim.RMSprop(params, lr=self.learning_rate, alpha=self.alpha, eps=self.eps,
                                            weight_decay=self.weight_decay, momentum=self.momentum,
@@ -174,13 +287,31 @@ class RMSprop(Optimizer):
 class SGD(Optimizer):
     r"""Stochastic Gradient Descent
 
+    This provides a wrapper around the SGD optimizer [1, 2, 3] where the parameters can be passed at a later stage.
+    The documentation is taken from [3].
+
     References:
-        [1] "A Stochastic Approximation Method", Robbins and Monro, 1951
-        [2] "On the importance of initialization and momentum in deep learning", Sutskever et al., 2013
+        - [1] "A Stochastic Approximation Method", Robbins and Monro, 1951
+        - [2] "On the importance of initialization and momentum in deep learning", Sutskever et al., 2013
+        - [3] torch.optim: https://pytorch.org/docs/stable/optim.html
     """
 
     def __init__(self, learning_rate=1e-3, momentum=0, dampening=0, weight_decay=0, nesterov=False, max_grad_norm=None,
                  *args, **kwargs):  # 0.5
+        """
+        Initialize the SGD optimizer.
+
+        Args:
+            learning_rate (float): learning rate.
+            momentum (float): momentum factor.
+            dampening (float): dampening for momentum.
+            weight_decay (float): weight decay (L2 penalty).
+            nesterov (bool): enables Nesterov momentum.
+            max_grad_norm (float, None): the allowed maximum gradient norm. If provided, it will clip the gradients.
+            verbose (bool, int): if verbose=2, it will print the gradient norm for each parameter.
+            *args (list): list of arguments that are given to the parent class `Optimizer`.
+            **kwargs (dict): dictionary of arguments that are given to the parent class `Optimizer`.
+        """
         super(SGD, self).__init__(*args, **kwargs)
         self.optimizer = None
         self.learning_rate = learning_rate
@@ -191,6 +322,13 @@ class SGD(Optimizer):
         self.max_grad_norm = max_grad_norm
 
     def optimize(self, params, loss):
+        """
+        Optimize the given parameters on the specified loss.
+
+        Args:
+            params (list of torch.Tensor): model parameters.
+            loss (torch.Tensor): loss value that were computed using the model parameters.
+        """
         # create optimizer if necessary
         if self.optimizer is None:
             self.optimizer = optim.SGD(params, lr=self.learning_rate, momentum=self.momentum, dampening=self.dampening,
@@ -223,13 +361,22 @@ class CG(Optimizer):
     matrix, and :math:`g` is the gradient.
 
     References:
-        [1] Conjugate Gradient (Wikipedia): https://en.wikipedia.org/wiki/Conjugate_gradient_method
-        [2] Hessian matrix (Wikipedia): https://en.wikipedia.org/wiki/Hessian_matrix#Use_in_optimization
-        [3] Hessian-Vector products: https://justindomke.wordpress.com/2009/01/17/hessian-vector-products/
-        [4] Torch CG: https://github.com/sbarratt/torch_cg
+        - [1] Conjugate Gradient (Wikipedia): https://en.wikipedia.org/wiki/Conjugate_gradient_method
+        - [2] Hessian matrix (Wikipedia): https://en.wikipedia.org/wiki/Hessian_matrix#Use_in_optimization
+        - [3] Hessian-Vector products: https://justindomke.wordpress.com/2009/01/17/hessian-vector-products/
+        - [4] Torch CG: https://github.com/sbarratt/torch_cg
     """
 
     def __init__(self, threshold=1.e-8, max_iters=10, *args, **kwargs):
+        """
+        Initialize the conjugate gradient.
+
+        Args:
+            threshold (float): threshold level.
+            max_iters (int): maximum number of iterations.
+            *args (list): list of arguments that are given to the parent class `Optimizer`.
+            **kwargs (dict): dictionary of arguments that are given to the parent class `Optimizer`.
+        """
         super(CG, self).__init__(*args, **kwargs)
         self._threshold = threshold
         self._max_iters = int(max_iters)
diff --git a/pyrobolearn/returns/estimators.py b/pyrobolearn/returns/estimators.py
index c6cf7e0..62ed55d 100644
--- a/pyrobolearn/returns/estimators.py
+++ b/pyrobolearn/returns/estimators.py
@@ -44,7 +44,7 @@ class Estimator(object):
     than 1, then it will reduces the variance but at the cost of introducing a bias.
 
     Reference:
-        [1] "High-Dimensional Continuous Control using Generalized Advantage Estimation", Schulman et al., 2016
+        - [1] "High-Dimensional Continuous Control using Generalized Advantage Estimation", Schulman et al., 2016
     """
 
     def __init__(self, storage, gamma=1.):
@@ -160,6 +160,9 @@ class Estimator(object):
         if self not in self.storage:
             self.storage.create_new_entry(key=self, shapes=1, num_steps=self.num_steps+1)
 
+            # add alias reference (note that modifying storage['returns'] will modify storage[self])
+            self.storage['returns'] = self.storage[self]
+
         return self._evaluate()
 
     #############
@@ -575,7 +578,7 @@ class GAE(Estimator):
     Good values for GAE are obtained when :math:`\gamma` and :math:`\tau` are in :math:`[0.9,0.99]`.
 
     References:
-        [1] "High-Dimensional Continuous Control using Generalized Advantage Estimation", Schulman et al., 2016
+        - [1] "High-Dimensional Continuous Control using Generalized Advantage Estimation", Schulman et al., 2016
     """
 
     def __init__(self, storage, value, gamma=0.98, tau=0.99):
diff --git a/pyrobolearn/samplers/sampler.py b/pyrobolearn/samplers/sampler.py
index 193668b..a80591e 100644
--- a/pyrobolearn/samplers/sampler.py
+++ b/pyrobolearn/samplers/sampler.py
@@ -52,7 +52,7 @@ class StorageSampler(Sampler):
     """
 
     def __init__(self, storage, sampler=None, num_batches=10, batch_size=None, batch_size_bounds=None,
-                 replacement=True):
+                 replacement=True, verbose=0):
         """
         Initialize the storage sampler.
 
@@ -71,6 +71,9 @@ class StorageSampler(Sampler):
                 one is too small (<16), it will be set to 16, and if this one is too big (>128), it will be set to 128.
             replacement (bool): if we should sample each element only one time, or we can sample the same ones multiple
                 times.
+            verbose (int, bool): verbose level, select between {0, 1, 2}. If 0=False, it won't print anything. If
+                1=True, it will print basic information about the sampler. If verbose=2, it will print detailed
+                information.
         """
         # set the storage
         self.storage = storage
@@ -80,6 +83,7 @@ class StorageSampler(Sampler):
         self._replacement = bool(replacement)
         self._batch_size_bounds = batch_size_bounds
         self._batch_size_given = batch_size is not None
+        self._verbose = verbose
 
         # set the sampler
         if sampler is None:
@@ -113,7 +117,9 @@ class StorageSampler(Sampler):
 
         self.sampler = sampler
 
-        print("Sampler: size: {} - num batches: {} - batch size: {}".format(self.size, num_batches, self.batch_size))
+        if verbose:
+            print("\nCreating sampler with size: {} - num batches: {} - batch size: {}".format(self.size, num_batches,
+                                                                                               self.batch_size))
 
     ##############
     # Properties #
@@ -244,8 +250,8 @@ class StorageSampler(Sampler):
         # get the filled size
         size = self.filled_size
 
-        print("Storage size: {}".format(self.size))
-        print("Storage filled size: {}".format(size))
+        if self._verbose:
+            print("Storage filled size: {} - size: {}".format(size, self.size))
 
         # modify the sampler (by changing the size)
         # check if there is a sub-sampler
@@ -268,6 +274,10 @@ class StorageSampler(Sampler):
             elif hasattr(self.sampler, 'indices'):
                 self.sampler.indices = range(size)
 
+        if self._verbose:
+            print("\nCreating sampler with size: {} - num batches: {} - batch size: {}".format(size, self.num_batches,
+                                                                                               self.batch_size))
+
         # provide the batches
         batch_idx = 0
         while True:  # this is to account for replacement = True
@@ -286,7 +296,7 @@ class BatchRandomSampler(StorageSampler):
 
     """
 
-    def __init__(self, storage, num_batches=10, batch_size=None, batch_size_bounds=None, replacement=True):
+    def __init__(self, storage, num_batches=10, batch_size=None, batch_size_bounds=None, replacement=True, verbose=0):
         """
         Initialize the storage sampler.
 
@@ -303,6 +313,10 @@ class BatchRandomSampler(StorageSampler):
                 one is too small (<16), it will be set to 16, and if this one is too big (>128), it will be set to 128.
             replacement (bool): if we should sample each element only one time, or we can sample the same ones multiple
                 times.
+            verbose (int, bool): verbose level, select between {0, 1, 2}. If 0=False, it won't print anything. If
+                1=True, it will print basic information about the sampler. If verbose=2, it will print detailed
+                information.
         """
         super(BatchRandomSampler, self).__init__(storage=storage, num_batches=num_batches, batch_size=batch_size,
-                                                 batch_size_bounds=batch_size_bounds, replacement=replacement)
+                                                 batch_size_bounds=batch_size_bounds, replacement=replacement,
+                                                 verbose=verbose)
diff --git a/pyrobolearn/storages/storage.py b/pyrobolearn/storages/storage.py
index 64b001d..f7ffbed 100644
--- a/pyrobolearn/storages/storage.py
+++ b/pyrobolearn/storages/storage.py
@@ -629,7 +629,7 @@ class Batch(DictStorage):
     are filled by the exploration phase in RL algorithms (see `pyrobolearn/algos/explorer`).
     """
 
-    def __init__(self, kwargs=None, device=None, dtype=None, size=None):
+    def __init__(self, kwargs=None, device=None, dtype=None, size=None, verbose=0):
         """
         Initialize the Batch storage.
 
@@ -641,7 +641,12 @@ class Batch(DictStorage):
             dtype (torch.dtype, None): convert the `torch.Tensor` to the specified data type. If None, it will keep
                 the original dtype
             size (int, None): size of the batch.
+            verbose (int, bool): if True, it will print information about the batch status, such as what is being
+                inserted, removed, etc. Don't use it on states or actions that are big.
         """
+        if verbose == 1:
+            print("\nCreating batch...")
+
         super(Batch, self).__init__(kwargs=kwargs, device=device, dtype=dtype, update=False)
         # contains the current values evaluated during the update phase of RL algorithms.
         self.current = DictStorage(kwargs={}, device=device, dtype=dtype, update=False)
@@ -649,6 +654,45 @@ class Batch(DictStorage):
         self.indices = None
         self._size = size if size is not None else len(kwargs['masks'])   # TODO: need to generalize this
 
+        if verbose > 1:
+            print("\nCreating batch with the following variables: ")
+
+            if 'states' in kwargs:
+                states = kwargs['states'][0]  # only take the first state for now
+                print("states: {}".format(torch.cat((torch.arange(len(states),
+                                                                  dtype=torch.float).view(-1, 1),
+                                                     states), dim=1)))
+
+            if 'actions' in kwargs:
+                actions = kwargs['actions'][0]  # only take the first action for now
+                print("actions: {}".format(torch.cat((torch.arange(len(actions),
+                                                                   dtype=torch.float).view(-1, 1),
+                                                      actions), dim=1)))
+
+            if 'rewards' in kwargs:
+                rewards = kwargs['rewards']
+                print("rewards: {}".format(torch.cat((torch.arange(len(rewards), dtype=torch.float).view(-1, 1),
+                                                      rewards), dim=1)))
+
+            if 'masks' in kwargs:
+                masks = kwargs['masks']
+                print("masks: {}".format(torch.cat((torch.arange(len(masks), dtype=torch.float).view(-1, 1),
+                                                    masks), dim=1)))
+
+            if 'returns' in kwargs:
+                returns =kwargs['returns']
+                print("returns: {}".format(torch.cat((torch.arange(len(returns), dtype=torch.float).view(-1, 1),
+                                                      returns), dim=1)))
+
+            # print other variables
+            tmp = {'states', 'actions', 'rewards', 'masks', 'returns'}
+            for key, value in kwargs.iteritems():
+                if key not in tmp:
+                    print("{}: {}".format(key, value))
+
+        if verbose:
+            print("Batch created.")
+
     @property
     def size(self):
         """Return the size of the batch."""
@@ -692,15 +736,15 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
     has its advantages, as you can for example store multiple value scalars from multiple value function approximators.
 
     Also, in contrast to [1], we do not compute the returns / estimators here. This is done by the `Estimator` class
-    which takes as input a `RolloutStorage`.
+    which takes as input a `RolloutStorage`, and will insert them inside the storage.
 
     In PRL, this storage is notably used by `RLAlgo` (`Explorator`, `Evaluator`, `Updater`), `Loss`, `Estimators`, etc.
 
     References:
-        [1] https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail/blob/master/a2c_ppo_acktr/storage.py
+        - [1] https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail/blob/master/a2c_ppo_acktr/storage.py
     """
 
-    def __init__(self, num_steps, state_shapes, action_shapes, num_trajectories=1):
+    def __init__(self, num_steps, state_shapes, action_shapes, num_trajectories=1, verbose=0):
         # , recurrent_hidden_state_size=0):
         """
         Initialize the rollout storage.
@@ -710,6 +754,9 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
             state_shapes (list of tuple of int, tuple of int): each tuple represents the shape of an observation/state.
             action_shapes (list of tuple of int, tuple of int): each tuple represents the shape of an action.
             num_trajectories (int): number of trajectories.
+            verbose (int, bool): verbose level, if False (=0) it won't print anything. If True (=1) it will print basic
+                information. If verbose=2, it will print information about what is being inserted and removed in the
+                storage, and the batch status.
         """
         # recurrent_hidden_state_size (int): size of the internal state
         super(RolloutStorage, self).__init__()
@@ -719,8 +766,17 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
         self._state_shapes = state_shapes
         self._action_shapes = action_shapes
         self._shifts = {}  # dictionary that maps the key to the time shift; this is add to the current time step
+        self.verbose = verbose
+
+        if self.verbose:
+            print("\nCreating RolloutStorage with num_steps={} and num_rollouts={}".format(self._num_steps,
+                                                                                           self._num_trajectories))
+
         self.init(self.num_steps, state_shapes, action_shapes, self.num_trajectories)
 
+        if self.verbose:
+            print("RolloutStorage created.")
+
     ##############
     # Properties #
     ##############
@@ -838,6 +894,9 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
             raise TypeError("Expecting the given shapes {} to be a list of tuple of int, a tuple of int, or an int, "
                             "instead got: {}".format({}, type(shapes)))
 
+        if self.verbose:
+            print("Storage: creating new entry for {} with shapes {}".format(key, shapes))
+
         # add shift
         self._shifts[key] = num_steps - self.num_steps
 
@@ -993,12 +1052,13 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
         """
         t = self._step[rollout_idx]
 
-        print("\nStorage: rollout = {}, step = {}".format(rollout_idx, t))
-        print("Storage: insert state: {}".format(states))
-        print("Storage: insert action: {}".format(actions))
-        print("Storage: insert next state: {}".format(next_states))
-        print("Storage: insert reward: {}".format(reward))
-        print("Storage: insert mask: {}".format(mask))
+        if self.verbose > 1:
+            print("\nStorage: rollout = {}, step = {}".format(rollout_idx, t))
+            print("Storage: insert state: {}".format(states))
+            print("Storage: insert action: {}".format(actions))
+            print("Storage: insert next state: {}".format(next_states))
+            print("Storage: insert reward: {}".format(reward))
+            print("Storage: insert mask: {}".format(mask))
 
         # check given observations/states and actions
         if not isinstance(next_states, list):
@@ -1082,37 +1142,45 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
         #             item = item[:-1]  # take only the T steps
         #         return item.reshape(-1, *item.shape[2:])[indices]  # reshape to (T*P, *shape) and from T*P takes I
 
-        print("Indices: {} - length: {}".format(indices, len(indices)))
+        if self.verbose:
+            print("\nStorage: get batch with size: {} and indices: {}".format(len(indices), indices))
 
         def sample(item, indices):
             """Given indices where each index is between 0 and `self.filled_size` (=number of masks that are equal
             to 1), it returns the corresponding entries in the item.
             """
-            if isinstance(item, torch.Tensor):  # TODO: improve performance of this
-                masks = (self.masks[:, :, 0] == 1).nonzero()  # [F, 2] --> indices for (steps, trajs)
-                masks[:, 0] -= 1  # remove 1 because indices for masks are in [1, t+1] --> [0, t]
-                indices = masks[indices]  # [I,2] --> allowed indices for (steps, trajs)
-                print("Torch mask length: {}".format(len(masks)))
-                print("Indices: {}".format(indices))
-                return item[indices[:, 0], indices[:, 1]]  # [I, *shape]
-            elif isinstance(item, np.ndarray):  # TODO: improve performance of this
-                masks = np.vstack(((self.masks[:, :, 0] == 1).nonzero()))  # [F,2] --> indices for (steps, trajs)
-                masks[:, 0] -= 1  # remove 1 because indices for masks are in [1, t+1] --> [0, t]
-                indices = masks[indices]  # [I,2] --> allowed indices for (steps, trajs)
-                print("Numpy Mask length: {}".format(len(masks)))
-                print("Indices: {}".format(indices))
+            if isinstance(item, (torch.Tensor, np.ndarray)):
                 return item[indices[:, 0], indices[:, 1]]  # [I, *shape]
+            else:
+                raise TypeError("Expecting the given 'item' to be a torch.Tensor or np.ndarray, but got: "
+                                "{}".format(type(item)))
+
+        # compute the step and traj indices
+        original_indices = indices
+        masks = (self.masks[:, :, 0] == 1).nonzero()  # [F, 2] --> indices for (step_idx, traj_idx)
+        masks[:, 0] -= 1  # remove 1 because indices for masks are in [1, t+1] --> [0, t]
+        indices = masks[indices]  # [I,2] --> allowed indices for (step_idx, traj_idx)
+
+        # if self.verbose:
+        #     print("Mask length: {}".format(len(masks)))
+        #     print("Indices: {}".format(indices))
 
         # go through each attribute and sample from the tensors
         for key, value in self.iteritems():
-            print("batch - add key: {}".format(key))
+            # print("batch - add key: {}".format(key))
             if isinstance(value, list):  # value = list of tensors
-                batch[key] = [sample(val, indices) for val in value]
+                batch[key] = [sample(val, indices) for val in value]  # [[I, *shape] for each shape]
             else:  # value = tensor
-                batch[key] = sample(value, indices)
+                # if key == 'masks':  # TODO: need to shift masks?
+                #     m = torch.clone(masks)
+                #     m[:, 0] += 1  # indices [0, t] --> [1, t+1]
+                #     idx = m[original_indices]
+                #     batch[key] = value[idx[:, 0], idx[:, 1]]
+                # else:
+                batch[key] = sample(value, indices)  # [I, *shape]
 
         # TODO: add the following lines in the `end` method? Need to be called only one time as long we don't clear
-        # replace the zeros in the action_distributions field by dummy distributions
+        #  replace the zeros in the action_distributions field by dummy distributions
         # take the first distribution for each action distribution
         dists = [dist[dist != 0][0] for dist in self['action_distributions']]
         # replace the zeros by that first distribution
@@ -1128,8 +1196,9 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
             distributions[i] = dist.__class__.from_list(distribution)
 
         # create Batch object
-        batch = Batch(batch, device=self.device, dtype=self.dtype, size=size)
+        batch = Batch(batch, device=self.device, dtype=self.dtype, size=size, verbose=self.verbose)
         batch.indices = torch.tensor(range(len(indices)))[batch['masks'][:, 0] != 0].tolist()
+        # print("Batch indices: {}".format(batch.indices))  # TODO: need to shift masks?
 
         # return batch (which is given to the updater (and loss))
         return batch