diff --git a/pyrobolearn/algos/explorer.py b/pyrobolearn/algos/explorer.py
index f32a528..a95d433 100644
--- a/pyrobolearn/algos/explorer.py
+++ b/pyrobolearn/algos/explorer.py
@@ -45,8 +45,8 @@ class Explorer(object):
         Args:
             task (Task, Env, tuple of Env and Policy): RL task or environment.
             explorer (Exploration): policies.
-            storage (DictStorage): Rollout storage unit (=replay memory). It will save the rollouts in the storage
-                while exploring.
+            storage (DictStorage): Rollout storage unit (=replay memory). It will save the trajectories / rollouts /
+                transitions in the storage while exploring.
             num_workers (int): number of processes / workers to run in parallel.
         """
         self.task = task
@@ -84,7 +84,7 @@ class Explorer(object):
     @property
     def policy(self):
         """Return the policy."""
-        return self.task.policy
+        return self.explorer.policy
 
     @property
     def env(self):
@@ -122,16 +122,18 @@ class Explorer(object):
     # Methods #
     ###########
 
-    def explore(self, num_steps, rollout_idx=0, deterministic=False):
+    def explore(self, num_steps, rollout_idx=0, deterministic=False, verbose=True):
         """
         Explore 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:
-            Rollout: memory storage
+            DictStorage: updated memory storage
         """
         # reset environment
         observation = self.env.reset()
@@ -151,33 +153,37 @@ class Explorer(object):
             # perform one step in the environment
             next_observation, reward, done, info = self.env.step(action)
 
-            # insert in storage
-            print("\nExplorer:")
-            print("1. Observation data: {}".format(observation))
-            print("2. Action data: {}".format(action))
-            print("3. Next observation data: {}".format(next_observation))
-            print("4. Reward: {}".format(reward))
-            print("5. \\pi(.|s): {}".format(distribution))
-            print("6. log \\pi(a|s): {}".format([d.log_prob(action) for d in distribution]))
+            if verbose:
+                print("\nExplorer:")
+                print("1. Observation data: {}".format(observation))
+                print("2. Action data: {}".format(action))
+                print("3. Next observation data: {}".format(next_observation))
+                print("4. Reward: {}".format(reward))
+                print("5. \\pi(.|s): {}".format(distribution))
+                print("6. log \\pi(a|s): {}".format([d.log_prob(action) for d in distribution]))
 
+            # insert in storage
             self.storage.insert(observation, action, next_observation, reward, mask=(1-done),
                                 distributions=distribution, rollout_idx=rollout_idx)
 
+            # set current observation to current one
             observation = next_observation
+
+            # if done, get out of the loop
             if done:
                 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("actions: {}".format(self.storage['actions']))
+        # print("rewards: {}".format(self.storage['rewards']))
         # print("masks: {}".format(self.storage['masks']))
         # print("distributions: {}".format(self.storage['distributions']))
-        raw_input('enter')
 
         # # clear explorer
         # self.explorer.clear()
 
+        # return storage unit
         return self.storage
 
     #############
diff --git a/pyrobolearn/algos/rl_algo.py b/pyrobolearn/algos/rl_algo.py
index e793183..d760ff4 100755
--- a/pyrobolearn/algos/rl_algo.py
+++ b/pyrobolearn/algos/rl_algo.py
@@ -147,9 +147,7 @@ class RLAlgo(object):  # Algo):
         [5] OpenAI - Spinning Up: https://spinningup.openai.com/
     """
 
-    # def __init__(self, rlTask, exploration_strategy, storage, hyperparameters, optimizer=None, dynamic_model=None,
-    #              num_workers=1):
-    def __init__(self, explorer, evaluator, updater, hyperparameters={}, dynamic_model=None):  # , num_workers=1):
+    def __init__(self, explorer, evaluator, updater, dynamic_model=None):  # , hyperparameters={}, num_workers=1):
         """
         Initialize the reinforcement learning algorithm.
 
@@ -157,9 +155,7 @@ class RLAlgo(object):  # Algo):
             explorer (Explorer): explorer that specifies how to explore in the environment
             evaluator (Evaluator): evaluate the actions
             updater (Updater): update the approximators (rl, value-functions,...)
-            hyperparameters (dict): dictionary containing the hyperparameters
             dynamic_model (None): dynamical model
-            num_workers (int): number of workers (useful when parallelizing the code)
         """
 
         super(RLAlgo, self).__init__()
@@ -172,11 +168,6 @@ class RLAlgo(object):  # Algo):
         self.env = self.environment
         self.dynamic_model = dynamic_model
 
-        # self.episodes = hyperparameters.get('episodes', 1) # nb of episodes
-        # self.rollouts = hyperparameters.get('rollouts', 1) # nb of rollouts per episode
-        # self.timesteps = hyperparameters.get('timesteps', 1000) # nb of timesteps per rollout
-        # TODO: define the number of iterations
-
         self.best_reward = -np.infty
         self.best_parameters = None
 
@@ -235,7 +226,7 @@ class RLAlgo(object):  # Algo):
     @property
     def policy(self):
         """Return the policy."""
-        return self.task.policy
+        return self.explorer.policy
 
     @property
     def exploration_strategy(self):
@@ -245,7 +236,9 @@ class RLAlgo(object):  # Algo):
     @property
     def estimator(self):
         """Return the estimator."""
-        return self.evaluator.estimator
+        if self.evaluator is not None:
+            return self.evaluator.estimator
+        return None
 
     @property
     def storage(self):
@@ -266,7 +259,18 @@ class RLAlgo(object):  # Algo):
     # Methods #
     ###########
 
-    def init(self, *args, **kwargs):
+    def init(self, num_steps, num_rollouts, num_episodes, seed=None, *args, **kwargs):
+        """
+        Initialize the reinforcement learning algorithm.
+
+        Args:
+            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)
+            seed (int): random seed
+            *args (list): list of optional arguments.
+            **kwargs (dict): dictionary of optional arguments.
+        """
         pass
 
     # def init(self, explorer, evaluator, updater):
@@ -275,36 +279,6 @@ class RLAlgo(object):  # Algo):
     #     self.evaluator = evaluator
     #     self.updater = updater
 
-    def rollout(self, deterministic=True):
-        """
-        Run the policy in the environment.
-        """
-        # Reset the environment
-        state = self.env.reset()
-
-        if deterministic:
-            self.explorer.disable()
-
-        # Run policy in environment for T time steps
-        total_reward = 0
-        for _ in range(self.timesteps):
-            # run policy given the state
-            prev_state = state
-            action = self.policy.act(state, self.exploration)
-
-            # run one step in the environment
-            state, reward, done, info = self.env.step(action)
-            total_reward += reward
-
-            # save (s,a,s',r) in storage
-            self.storage.add(prev_state, action, state, reward)
-
-            # if episode is done
-            if done:
-                break
-
-        return total_reward
-
     def train(self, num_steps, num_rollouts=1, num_episodes=1, verbose=False, seed=None):
         """
         Train the policy in the provided environment.
@@ -321,6 +295,9 @@ class RLAlgo(object):  # Algo):
         """
         history = {}
 
+        # init algo with the given parameters
+        self.init(num_steps=num_steps, num_rollouts=num_rollouts, num_episodes=num_episodes, seed=seed)
+
         # set the policy in training mode
         self.policy.train()
 
@@ -347,7 +324,7 @@ class RLAlgo(object):  # Algo):
 
     def test(self, num_steps, dt=0., use_terminating_condition=False, render=True):  # , storage):
         """
-        Test the policy in the environment.
+        Test the policy in the environment; perform one rollout.
 
         Args:
             num_steps (int): number of steps
@@ -368,9 +345,11 @@ class RLAlgo(object):  # Algo):
     #############
 
     def __repr__(self):
+        """Return a representation string about the object."""
         return self.__class__.__name__
 
     def __str__(self):
+        """Return a string describing the object."""
         return self.__class__.__name__
 
 
@@ -388,14 +367,13 @@ class GradientRLAlgo(RLAlgo):
         TD residual,...)
     """
 
-    def __init__(self, explorer, evaluator, updater, hyperparameters=None, dynamic_model=None):
-        super(GradientRLAlgo, self).__init__(explorer, evaluator, updater, hyperparameters, dynamic_model)
+    def __init__(self, explorer, evaluator, updater, dynamic_model=None):  # hyperparameters=None)
+        super(GradientRLAlgo, self).__init__(explorer, evaluator, updater, dynamic_model)
 
 
 class EMRLAlgo(RLAlgo):
     r"""Expectation-Maximization reinforcement learning algorithm.
-
     """
 
-    def __init__(self, task, exploration_strategy, storage, hyperparameters):
-        super(EMRLAlgo, self).__init__(task, exploration_strategy, storage, hyperparameters)
+    def __init__(self, task, exploration_strategy, storage, dynamic_model=None):  # hyperparameters=None)
+        super(EMRLAlgo, self).__init__(task, exploration_strategy, storage, dynamic_model)
diff --git a/pyrobolearn/exploration/actions/action_exploration.py b/pyrobolearn/exploration/actions/action_exploration.py
index a26ecf5..5108fee 100644
--- a/pyrobolearn/exploration/actions/action_exploration.py
+++ b/pyrobolearn/exploration/actions/action_exploration.py
@@ -198,7 +198,7 @@ class ActionExploration(Exploration):
                     action_data, action_distribution = self.explore(action_data)
 
                 # post-process the action data
-                self.policy.postprocess(action_data)
+                action_data = self.policy.postprocess(action_data)
 
                 # set the action data
                 self.action_data = self.policy.set_action_data(action_data, to_numpy=to_numpy,
diff --git a/pyrobolearn/policies/policy.py b/pyrobolearn/policies/policy.py
index 1aac9da..714b323 100644
--- a/pyrobolearn/policies/policy.py
+++ b/pyrobolearn/policies/policy.py
@@ -482,17 +482,28 @@ class Policy(object):
         for idx, (action, data) in enumerate(zip(self.actions, action_data)):
             if action.is_discrete():  # discrete action
                 if isinstance(data, np.ndarray):  # data action is a numpy array
-                    discrete_data = np.array([np.argmax(data)])
-                    action.data = discrete_data
-                    if not return_logits:
-                        action_data[idx] = discrete_data
+                    # check if given logits or not
+                    if data.shape[-1] == 1:  # no logits
+                        action.data = data
+                    else:  # given logits
+                        discrete_data = np.array([np.argmax(data)])
+                        action.data = discrete_data
+
+                        # if we do not want the logits in the action data, replace it by the discrete data
+                        if not return_logits:
+                            action_data[idx] = discrete_data
+
                 elif isinstance(data, torch.Tensor):  # data action is a torch.Tensor
-                    discrete_data = torch.argmax(data, dim=0, keepdim=True)
-                    action.torch_data = discrete_data
-                    if not return_logits:
-                        action_data[idx] = self.__convert_to_numpy(discrete_data, to_numpy=to_numpy)
-                    else:
-                        action_data[idx] = self.__convert_to_numpy(data, to_numpy=to_numpy)
+                    # check if given logits or not
+                    if data.shape[-1] == 1:  # no logits
+                        action.torch_data = data
+                    else:  # given logits
+                        discrete_data = torch.argmax(data, dim=-1, keepdim=True)
+                        action.torch_data = discrete_data
+                        if not return_logits:
+                            action_data[idx] = self.__convert_to_numpy(discrete_data, to_numpy=to_numpy)
+                        else:
+                            action_data[idx] = self.__convert_to_numpy(data, to_numpy=to_numpy)
                 # elif isinstance(data, (float, int)):
                 #     discrete_data = np.argmax(data)
                 #     action.data = discrete_data