From 807d51e4d5e8ba91ada2a76d91e620e8c7fde507 Mon Sep 17 00:00:00 2001
From: Brian Delhaisse <briandelhaisse@gmail.com>
Date: Thu, 18 Apr 2019 14:05:31 +0200
Subject: [PATCH] fix few errors and update policy, storage, losses, returns,
 algos

---
 pyrobolearn/algos/evaluator.py                | 28 +++++--
 pyrobolearn/algos/explorer.py                 | 34 ++++++---
 pyrobolearn/algos/reinforce.py                |  8 +-
 pyrobolearn/algos/rl_algo.py                  | 26 +++++--
 pyrobolearn/algos/updater.py                  | 74 +++++++++++++------
 .../exploration/actions/action_exploration.py | 50 +++++++++++++
 pyrobolearn/exploration/exploration.py        | 18 +++++
 pyrobolearn/losses/policy_losses.py           | 45 ++++++-----
 pyrobolearn/policies/policy.py                |  6 +-
 pyrobolearn/returns/estimators.py             | 14 +++-
 pyrobolearn/returns/evaluators.py             | 11 ++-
 pyrobolearn/returns/returns.py                | 12 +++
 pyrobolearn/returns/targets.py                | 12 +++
 pyrobolearn/storages/storage.py               | 11 ++-
 14 files changed, 275 insertions(+), 74 deletions(-)

diff --git a/pyrobolearn/algos/evaluator.py b/pyrobolearn/algos/evaluator.py
index 21230b8..25e7a2c 100644
--- a/pyrobolearn/algos/evaluator.py
+++ b/pyrobolearn/algos/evaluator.py
@@ -66,12 +66,23 @@ class Evaluator(object):
     # Methods #
     ###########
 
-    def evaluate(self):
+    def evaluate(self, verbose=False):
         """
-        Evaluate the actions.
+        Evaluate the trajectories performed by the policy.
+
+        Args:
+            verbose (bool): If true, print information on the standard output.
         """
         if self.estimator is not None:
-            self.estimator.evaluate(self.storage)
+            if verbose:
+                print("\n#### Starting the Evaluation phase ####")
+
+            # compute the returns
+            returns = self.estimator.evaluate(self.storage)
+
+            if verbose:
+                print("Returns: {}".format(returns))
+                print("#### End of the Evaluation phase ####")
 
     #############
     # Operators #
@@ -85,6 +96,11 @@ class Evaluator(object):
         """Return the class string."""
         return self.__class__.__name__
 
-    def __call__(self):
-        """Evaluate the estimator on the storage."""
-        self.evaluate()
+    def __call__(self, verbose=False):
+        """
+        Evaluate the trajectories performed by the policy.
+
+        Args:
+            verbose (bool): If true, print information on the standard output.
+        """
+        self.evaluate(verbose=verbose)
diff --git a/pyrobolearn/algos/explorer.py b/pyrobolearn/algos/explorer.py
index a95d433..47e2371 100644
--- a/pyrobolearn/algos/explorer.py
+++ b/pyrobolearn/algos/explorer.py
@@ -122,9 +122,9 @@ class Explorer(object):
     # Methods #
     ###########
 
-    def explore(self, num_steps, rollout_idx=0, deterministic=False, verbose=True):
+    def explore(self, num_steps, rollout_idx=0, deterministic=False, verbose=False):
         """
-        Explore the environment.
+        Explore in the environment.
 
         Args:
             num_steps (int): number of steps
@@ -137,7 +137,9 @@ class Explorer(object):
         """
         # reset environment
         observation = self.env.reset()
-        print("\nExplorer - initial state: {}".format(observation))
+        if verbose:
+            print("\n#### Starting the Exploration phase ####")
+            print("Explorer - initial state: {}".format(observation))
 
         # reset storage
         self.storage.reset(init_states=observation, rollout_idx=rollout_idx)
@@ -174,11 +176,13 @@ class Explorer(object):
                 self.storage.end(rollout_idx)  # fill remaining mask values
                 break
 
-        # print("states: {}".format(self.storage['states']))
-        # print("actions: {}".format(self.storage['actions']))
-        # print("rewards: {}".format(self.storage['rewards']))
-        # print("masks: {}".format(self.storage['masks']))
-        # print("distributions: {}".format(self.storage['distributions']))
+        if verbose:
+            print("#### End of the Exploration phase #####")
+            # print("states: {}".format(self.storage['states']))
+            # print("actions: {}".format(self.storage['actions']))
+            # print("rewards: {}".format(self.storage['rewards']))
+            # print("masks: {}".format(self.storage['masks']))
+            # print("distributions: {}".format(self.storage['distributions']))
 
         # # clear explorer
         # self.explorer.clear()
@@ -198,6 +202,16 @@ class Explorer(object):
         """Return a string describing the class."""
         return self.__class__.__name__
 
-    def __call__(self, num_steps, rollout_idx=0):
-        """Explore in the environment with the specified number of time steps."""
+    def __call__(self, num_steps, rollout_idx=0, deterministic=False, verbose=True):
+        """Explore in the environment.
+
+        Args:
+            num_steps (int): number of steps
+            rollout_idx (int): trajectory/rollout index.
+            deterministic (bool): if deterministic is True, then it does not explore in the environment.
+            verbose (bool): If true, print information on the standard output.
+
+        Returns:
+            DictStorage: updated memory storage
+        """
         self.explore(num_steps, rollout_idx=rollout_idx)
diff --git a/pyrobolearn/algos/reinforce.py b/pyrobolearn/algos/reinforce.py
index 407e13f..d19799c 100755
--- a/pyrobolearn/algos/reinforce.py
+++ b/pyrobolearn/algos/reinforce.py
@@ -16,7 +16,7 @@ from pyrobolearn.exploration import ActionExploration
 
 from pyrobolearn.storages import RolloutStorage
 from pyrobolearn.samplers import StorageSampler
-from pyrobolearn.returns import ActionRewardEstimator
+from pyrobolearn.returns import ActionRewardEstimator, PolicyEvaluator
 from pyrobolearn.losses import PGLoss, ValueL2Loss
 from pyrobolearn.optimizers import Adam
 
@@ -189,12 +189,14 @@ class REINFORCE(GradientRLAlgo):
         # create storage
         states, actions = policy.states, policy.actions
         storage = RolloutStorage(num_steps=1000, state_shapes=states.shape, action_shapes=actions.shape,
-                                 num_trajectories=10)
+                                 num_trajectories=1)
         sampler = StorageSampler(storage)
 
         # create return: R_t = \sum_{t'=t}^{T} \gamma^{t'-t} r_{t'}
         returns = ActionRewardEstimator(storage, gamma=gamma)
 
+        policy_evaluator = PolicyEvaluator(policy=exploration)
+
         # create loss for policy: \mathbb{E}[ \log \pi_{\theta}(a_t | s_t) R_t ]
         loss = PGLoss(returns)
 
@@ -212,7 +214,7 @@ class REINFORCE(GradientRLAlgo):
         # define the 3 main steps in RL: explore, evaluate, and update
         explorer = Explorer(task, exploration, storage, num_workers=num_workers)
         evaluator = Evaluator(returns)
-        updater = Updater(approximators, sampler, loss, optimizer, evaluators=[])
+        updater = Updater(approximators, sampler, loss, optimizer, evaluators=[policy_evaluator])
 
         # initialize RL algorithm
         super(REINFORCE, self).__init__(explorer, evaluator, updater)
diff --git a/pyrobolearn/algos/rl_algo.py b/pyrobolearn/algos/rl_algo.py
index d760ff4..21bdb74 100755
--- a/pyrobolearn/algos/rl_algo.py
+++ b/pyrobolearn/algos/rl_algo.py
@@ -221,6 +221,7 @@ class RLAlgo(object):  # Algo):
 
     @property
     def environment(self):
+        """Return the environment instance."""
         return self.task.environment
 
     @property
@@ -295,6 +296,9 @@ class RLAlgo(object):  # Algo):
         """
         history = {}
 
+        if verbose:
+            print("\n#### Start the RL algo ####")
+
         # init algo with the given parameters
         self.init(num_steps=num_steps, num_rollouts=num_rollouts, num_episodes=num_episodes, seed=seed)
 
@@ -308,18 +312,26 @@ class RLAlgo(object):  # Algo):
             for rollout in range(num_rollouts):
                 # TODO: consider to learn the dynamic model if provided
 
-                # Explore, evaluate, and update
-                self.explorer(num_steps, rollout)
-                if self.evaluator is not None:
-                    self.evaluator()
-                loss = self.updater()
+                if verbose:
+                    print("Episode: {}/{} - Rollout: {}/{}".format(episode+1, num_episodes, rollout+1, num_rollouts))
 
-                # add the loss in the history
-                history.setdefault('loss', []).append(loss)
+                # Explore
+                self.explorer.explore(num_steps, rollout, verbose=verbose)
+
+            # evaluate and update
+            if self.evaluator is not None:
+                self.evaluator.evaluate(verbose=verbose)
+            losses = self.updater()
+
+            # add the loss in the history
+            history.setdefault('losses', []).append(losses)
 
         # set the policy in test mode
         self.policy.eval()
 
+        if verbose:
+            print("\n#### End of the RL algo ####")
+
         return history
 
     def test(self, num_steps, dt=0., use_terminating_condition=False, render=True):  # , storage):
diff --git a/pyrobolearn/algos/updater.py b/pyrobolearn/algos/updater.py
index 1d20e1e..d814fdd 100644
--- a/pyrobolearn/algos/updater.py
+++ b/pyrobolearn/algos/updater.py
@@ -252,13 +252,13 @@ class Updater(object):
 
         # set the tick for each loss
         for loss in self.losses:
-            if loss not in self._ticks:
-                self._ticks[loss] = 1
+            if loss not in ticks:
+                ticks[loss] = 1
 
         # set the tick for each updater
         for updater in self.updaters:
-            if updater not in self._ticks:
-                self._ticks[updater] = 1
+            if updater not in ticks:
+                ticks[updater] = 1
 
         # set the ticks
         self._ticks = ticks
@@ -267,34 +267,42 @@ class Updater(object):
     # Methods #
     ###########
 
-    def update(self, num_batches=10, num_epochs=1):
+    def update(self, num_epochs=1, num_batches=10, verbose=False):
         """
         Update the given approximators (policies, value functions, etc).
 
         Args:
-            num_batches (int): number of batches.
             num_epochs (int): number of epochs.
+            num_batches (int): number of batches.
+            verbose (bool): If true, print information on the standard output.
 
         Returns:
-            list: list of losses
+            dict: dictionary of losses. There is a key for each loss, and the value is a nested list which contains
+                the obtained loss for each epoch and for each batch in the corresponding epoch.
         """
         # set the number of batches
         self.sampler.num_batches = num_batches
 
-        losses = []
+        # keep history of each loss
+        losses = {}
+
+        if verbose:
+            print("\n#### Starting the Update phase ####")
 
         # for each epoch
         for epoch in range(num_epochs):
 
-            # batch losses
-            batch_losses = []
-
             # for each batch
-            for batch in self.sampler:
+            for batch_idx, batch in enumerate(self.sampler):
+
+                if verbose:
+                    print("Epoch: {}/{} - Batch: {}/{}".format(epoch + 1, num_epochs, batch_idx + 1, num_batches))
 
                 # evaluate the evaluators with the current parameters on the given batch and save the results in the
                 # batch's `current` attribute
                 for evaluator in self.evaluators:
+                    if verbose:
+                        print("Subevaluation on the batch using the estimator: {}".format(evaluator))
                     evaluator.evaluate(batch, store=True)
 
                 # update each approximator based on the loss on which it is evaluated and using the specified optimizer
@@ -303,25 +311,36 @@ class Updater(object):
                     # if time to update
                     if self._cnt % self.ticks[loss] == 0:
 
-                        # compute loss on the data (the loss knows what to do with the batch)
-                        loss = loss.compute(batch)
+                        if verbose:
+                            print("\t Compute loss {}".format(loss))
 
-                        # append the loss / batch
-                        batch_losses.append(loss)
+                        # compute loss on the data (the loss knows what to do with the batch)
+                        loss_value = loss.compute(batch)
+
+                        # append the loss value in the history of losses
+                        if loss not in losses:
+                            losses[loss] = [[]] * num_epochs
+                        else:
+                            losses[loss][epoch].append(loss_value)
 
                         # update parameters
-                        optimizer.optimize(approximator.parameters(), loss)
+                        if verbose:
+                            print("\t Optimize the parameters for {} using the loss value: {}".format(approximator,
+                                                                                                      loss_value))
+                        optimizer.optimize(approximator.parameters(), loss_value)
 
                     # call each updater
                     for updater in self.updaters:
                         if self._cnt % self.ticks[updater] == 0:
+                            if verbose:
+                                print("\tRun updater {}".format(updater))
                             updater()
 
                 # increase counter
                 self._cnt += 1
 
-            # append the batch losses into the epoch losses
-            losses.append(batch_losses)
+        if verbose:
+            print("#### End of the Update phase ####")
 
         return losses  # shape=(epochs, batches)
 
@@ -337,6 +356,17 @@ class Updater(object):
         """Return a string describing the class."""
         return self.__class__.__name__
 
-    def __call__(self, num_batches=10):  # , storage, losses):
-        """Update the approximators."""
-        self.update(num_batches=num_batches)
+    def __call__(self, num_epochs=1, num_batches=10, verbose=False):
+        """
+        Update the given approximators (policies, value functions, etc).
+
+        Args:
+            num_epochs (int): number of epochs.
+            num_batches (int): number of batches.
+            verbose (bool): If true, print information on the standard output.
+
+        Returns:
+            dict: dictionary of losses. There is a key for each loss, and the value is a nested list which contains
+                the obtained loss for each epoch and for each batch in the corresponding epoch.
+        """
+        self.update(num_epochs=num_epochs, num_batches=num_batches, verbose=verbose)
diff --git a/pyrobolearn/exploration/actions/action_exploration.py b/pyrobolearn/exploration/actions/action_exploration.py
index 5108fee..fc26bb2 100644
--- a/pyrobolearn/exploration/actions/action_exploration.py
+++ b/pyrobolearn/exploration/actions/action_exploration.py
@@ -150,6 +150,56 @@ class ActionExploration(Exploration):
         """
         raise NotImplementedError
 
+    def predict(self, state=None, deterministic=True, to_numpy=False, return_logits=True):
+        """Predict the action given the state.
+
+        This does not set the action data in the action instances, nor apply the actions in the simulator. Instead,
+        it gets the state data, preprocess it, predict using the actions using the inner model, then post-process
+        the actions, and return the resulting action data.
+
+        Args:
+            state (State): current state
+            deterministic (bool): True by default. It can only be set to False, if the policy is stochastic.
+            to_numpy (bool): If True, it will convert the data (torch.Tensors) to numpy arrays.
+            return_logits (bool): If True, in the case of discrete outputs, it will return the logits.
+
+        Returns:
+            (list of) torch.Tensor: action data
+        """
+        # if deterministic outcome, i.e. we don't explore we just predict the actions using the policy
+        if deterministic:
+            action_data = self.policy.predict(state=state, deterministic=True, to_numpy=to_numpy,
+                                              return_logits=return_logits)
+            action_distribution = None
+
+        # if we should explore
+        else:
+            # get the state data
+            state_data = self.policy.get_state_data(state=state)
+
+            # pre-process the state data
+            state_data = self.policy.preprocess(state_data)
+
+            # predict the actions using the inner model
+            action_data = self.policy.inner_predict(state_data, deterministic=True, to_numpy=False,
+                                                    return_logits=True, set_output_data=False)
+
+            # exploration phase
+
+            # if exploration is a combination of multiple exploration
+            if self._explorations:
+                # explore for each action
+                actions = [explorer.explore(action_data) for explorer in self.explorations]
+                action_data, action_distribution = [a[0] for a in actions], [a[1] for a in actions]
+
+            else:  # there is only one action
+                action_data, action_distribution = self.explore(action_data)
+
+            # post-process the action data
+            action_data = self.policy.postprocess(action_data)
+
+        return action_data, action_distribution
+
     def act(self,  state=None, deterministic=False, to_numpy=False, return_logits=False, apply_action=True):
         r"""
         Act/Explore in the environment given the states.
diff --git a/pyrobolearn/exploration/exploration.py b/pyrobolearn/exploration/exploration.py
index 1ad7d0a..9217de7 100644
--- a/pyrobolearn/exploration/exploration.py
+++ b/pyrobolearn/exploration/exploration.py
@@ -80,6 +80,24 @@ class Exploration(object):  # TODO: inherit from Policy?
     #     """Perform the exploratory action."""
     #     pass
 
+    def predict(self, state=None, deterministic=True, to_numpy=False, return_logits=True):
+        """Predict the action given the state.
+
+        This does not set the action data in the action instances, nor apply the actions in the simulator. Instead,
+        it gets the state data, preprocess it, predict using the actions using the inner model, then post-process
+        the actions, and return the resulting action data.
+
+        Args:
+            state (State): current state
+            deterministic (bool): True by default. It can only be set to False, if the policy is stochastic.
+            to_numpy (bool): If True, it will convert the data (torch.Tensors) to numpy arrays.
+            return_logits (bool): If True, in the case of discrete outputs, it will return the logits.
+
+        Returns:
+            (list of) torch.Tensor: action data
+        """
+        pass
+
     def act(self, state=None, deterministic=False, to_numpy=False, return_logits=False, apply_action=True):
         """Perform the action given the state.
 
diff --git a/pyrobolearn/losses/policy_losses.py b/pyrobolearn/losses/policy_losses.py
index cdd47fa..366dfd8 100644
--- a/pyrobolearn/losses/policy_losses.py
+++ b/pyrobolearn/losses/policy_losses.py
@@ -62,9 +62,10 @@ class PGLoss(BatchLoss):
         Returns:
             torch.Tensor: loss scalar value
         """
-        # evaluate the action
-        log_curr_pi = batch.current['action_distributions']
-        log_curr_pi = log_curr_pi.log_probs(batch.current['actions'])
+        # evaluate the action  # TODO: think when there are multiple actions --> independent joint distribution?
+        log_curr_pi = 0
+        for pi, action in zip(batch.current['action_distributions'], batch.current['actions']):
+            log_curr_pi += pi.log_prob(action)
         estimator = batch[self._estimator]
 
         # compute loss and return it
@@ -115,11 +116,14 @@ class CPILoss(BatchLoss):
         Returns:
             torch.Tensor: loss scalar value
         """
-        # ratio = policy_distribution / old_policy_distribution
-        log_curr_pi = batch.current['action_distributions']
-        log_curr_pi = log_curr_pi.log_probs(batch.current['actions'])
-        log_prev_pi = batch['action_distributions']
-        log_prev_pi = log_prev_pi.log_probs(batch['actions'])
+        # evaluate the actions # TODO: think when there are multiple actions
+        log_curr_pi = 0
+        for pi, action in zip(batch.current['action_distributions'], batch.current['actions']):
+            log_curr_pi += pi.log_prob(action)
+
+        log_prev_pi = 0
+        for pi, action in zip(batch['action_distributions'], batch['actions']):
+            log_prev_pi += pi.log_prob(action)
 
         ratio = torch.exp(log_curr_pi - log_prev_pi)
         estimator = batch[self._estimator]
@@ -172,10 +176,14 @@ class CLIPLoss(BatchLoss):
         Returns:
             torch.Tensor: loss scalar value
         """
-        log_curr_pi = batch.current['action_distributions']
-        log_curr_pi = log_curr_pi.log_probs(batch.current['actions'])
-        log_prev_pi = batch['action_distributions']
-        log_prev_pi = log_prev_pi.log_probs(batch['actions'])
+        # evaluate the actions # TODO: think when there are multiple actions
+        log_curr_pi = 0
+        for pi, action in zip(batch.current['action_distributions'], batch.current['actions']):
+            log_curr_pi += pi.log_prob(action)
+
+        log_prev_pi = 0
+        for pi, action in zip(batch['action_distributions'], batch['actions']):
+            log_prev_pi += pi.log_prob(action)
 
         ratio = torch.exp(log_curr_pi - log_prev_pi)
         estimator = batch[self._estimator]
@@ -216,10 +224,10 @@ class KLPenaltyLoss(BatchLoss):
         Returns:
             torch.Tensor: loss scalar value
         """
-        curr_pi = batch.current['action_distributions']
-        prev_pi = batch['action_distributions']
-
-        return torch.distributions.kl.kl_divergence(prev_pi, curr_pi).mean()
+        kl_div = 0
+        for curr_pi, prev_pi in zip(batch.current['action_distributions'], batch['action_distributions']):
+            kl_div += torch.distributions.kl.kl_divergence(prev_pi, curr_pi).mean()
+        return kl_div
 
     def latex(self):
         """Return a latex formula of the loss."""
@@ -257,6 +265,7 @@ class EntropyLoss(BatchLoss):
         Returns:
             torch.Tensor: loss scalar value
         """
-        distribution = batch.current['action_distributions']
-        entropy = distribution.entropy().mean()
+        entropy = 0
+        for dist in batch.current['action_distributions']:
+            entropy += dist.entropy().mean()
         return entropy
diff --git a/pyrobolearn/policies/policy.py b/pyrobolearn/policies/policy.py
index f71e9bd..c149075 100644
--- a/pyrobolearn/policies/policy.py
+++ b/pyrobolearn/policies/policy.py
@@ -407,8 +407,10 @@ class Policy(object):
         # if the input is an instance of State, get the inner merged data.
         if isinstance(state, State):
             state = state.merged_data
-            if len(state) == 1:
-                state = state[0]
+
+        # if the input state is a list of len(1)
+        if isinstance(state, list) and len(state) == 1:
+            state = state[0]
 
         return state
 
diff --git a/pyrobolearn/returns/estimators.py b/pyrobolearn/returns/estimators.py
index 0a896e2..dd9735c 100644
--- a/pyrobolearn/returns/estimators.py
+++ b/pyrobolearn/returns/estimators.py
@@ -93,7 +93,7 @@ class Estimator(object):
     @property
     def returns(self):
         """Return the returns tensor from the rollout storage."""
-        return self.storage['self']
+        return self.storage[self]
 
     @property
     def states(self):
@@ -166,6 +166,14 @@ class Estimator(object):
     # Operators #
     #############
 
+    def __repr__(self):
+        """Return a representation string of the object."""
+        return self.__class__.__name__
+
+    def __str__(self):
+        """Return a string describing the object."""
+        return self.__class__.__name__ + "(gamma=" + str(self.gamma) + ")"
+
     def __call__(self, storage=None):
         """Evaluate the estimator on the rollout storage.
 
@@ -613,3 +621,7 @@ class GAE(Estimator):
             returns[t] = gae + values[t]
 
         return returns
+
+    def __str__(self):
+        """Return a string describing the object."""
+        return self.__class__.__name__ + "(gamma=" + str(self.gamma) + ", tau=" + str(self.tau) + ")"
diff --git a/pyrobolearn/returns/evaluators.py b/pyrobolearn/returns/evaluators.py
index 79060a4..b8dc620 100644
--- a/pyrobolearn/returns/evaluators.py
+++ b/pyrobolearn/returns/evaluators.py
@@ -73,6 +73,14 @@ class Evaluator(object):
                 batch.current[self] = output
         return output
 
+    def __repr__(self):
+        """Return a representation string of the object."""
+        return self.__class__.__name__
+
+    def __str__(self):
+        """Return a string describing the object."""
+        return self.__class__.__name__
+
     def __call__(self, batch, store=True):
         """Evaluate the evaluator on the given batch."""
         return self.evaluate(batch)
@@ -170,7 +178,8 @@ class PolicyEvaluator(Evaluator):
             torch.Tensor: advantage estimates.
         """
         # evaluate policy
-        actions, action_distributions = self._policy.predict(batch['states'])
+        actions, action_distributions = self._policy.predict(batch['states'], deterministic=False, to_numpy=False,
+                                                             return_logits=False)
 
         # return actions and distribution over actions
         return [('actions', actions), ('action_distributions', action_distributions)]
diff --git a/pyrobolearn/returns/returns.py b/pyrobolearn/returns/returns.py
index 9249b08..5feb543 100644
--- a/pyrobolearn/returns/returns.py
+++ b/pyrobolearn/returns/returns.py
@@ -79,6 +79,14 @@ class Return(object):
             batch.current[self] = output
         return output
 
+    def __repr__(self):
+        """Return a representation string of the object."""
+        return self.__class__.__name__
+
+    def __str__(self):
+        """Return a string describing the object."""
+        return self.__class__.__name__
+
     def __call__(self, batch, store=True):
         """
         Evaluate the return on the given batch.
@@ -126,6 +134,10 @@ class TDReturn(Return):
 
         self._gamma = gamma
 
+    def __str__(self):
+        """Return a string describing the object."""
+        return self.__class__.__name__ + "(gamma=" + str(self.gamma) + ")"
+
 
 class TDValueReturn(TDReturn):
     r"""TD State Value Return
diff --git a/pyrobolearn/returns/targets.py b/pyrobolearn/returns/targets.py
index f97f93e..5f616dd 100644
--- a/pyrobolearn/returns/targets.py
+++ b/pyrobolearn/returns/targets.py
@@ -85,6 +85,14 @@ class Target(object):
             batch.current[self] = output
         return output
 
+    def __repr__(self):
+        """Return a representation string of the object."""
+        return self.__class__.__name__
+
+    def __str__(self):
+        """Return a string describing the object."""
+        return self.__class__.__name__
+
     def __call__(self, batch, store=True):
         """Evaluate the target on the given batch."""
         return self.evaluate(batch)
@@ -119,6 +127,10 @@ class GammaTarget(Target):
 
         self._gamma = gamma
 
+    def __str__(self):
+        """Return a string describing the object."""
+        return self.__class__.__name__ + "(gamma=" + str(self.gamma) + ")"
+
 
 class VTarget(Target):
     r"""Value Target.
diff --git a/pyrobolearn/storages/storage.py b/pyrobolearn/storages/storage.py
index 9e62a1a..67da3a9 100644
--- a/pyrobolearn/storages/storage.py
+++ b/pyrobolearn/storages/storage.py
@@ -861,16 +861,19 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
         # reset the step
         self._step[rollout_idx] = 0
 
+        # reset masks
+        self.masks[:, rollout_idx].copy_(torch.ones_like(self.masks[:, rollout_idx]))
+
         # insert initial states
         if init_states is None:
             for state in self.states:
-                state[0][rollout_idx].copy_(state[-1][rollout_idx])
+                state[0, rollout_idx].copy_(state[-1, rollout_idx])
         else:
             if not isinstance(init_states, list):
                 init_states = [init_states]
             for observation, value in zip(self.states, init_states):
-                observation[0][rollout_idx].copy_(self._convert_to_tensor(value))
-        self.masks[0][rollout_idx].copy_(self.masks[-1][rollout_idx])
+                observation[0, rollout_idx].copy_(self._convert_to_tensor(value))
+        # self.masks[0, rollout_idx].copy_(self.masks[-1, rollout_idx])
         # self.recurrent_hidden_states[0].copy_(self.recurrent_hidden_states[-1])
 
     def update_tensor(self, key, values, step=None, rollout_idx=None, copy=True):
@@ -1013,7 +1016,7 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
 
         # go through each attribute and sample from the tensors
         for key, value in self.iteritems():
-            if isinstance(list, value):  # value = list of tensors
+            if isinstance(value, list):  # value = list of tensors
                 batch[key] = [sample(val, indices) for val in value]
             else:  # value = tensor
                 batch[key] = sample(value, indices)