From 478c32f7e4e6ed99ed87323c7dbd298e6ddc0796 Mon Sep 17 00:00:00 2001
From: Brian Delhaisse <briandelhaisse@gmail.com>
Date: Wed, 10 Apr 2019 04:37:51 +0200
Subject: [PATCH] update explorations and policies

---
 pyrobolearn/distributions/modules.py          |  38 ++--
 .../exploration/actions/action_exploration.py | 100 ++++++---
 pyrobolearn/exploration/actions/boltzmann.py  |  16 ++
 pyrobolearn/exploration/actions/continuous.py |   4 +
 pyrobolearn/exploration/actions/discrete.py   |   4 +
 pyrobolearn/exploration/actions/eps_greedy.py |  24 +-
 pyrobolearn/exploration/actions/gaussian.py   |  16 ++
 pyrobolearn/exploration/exploration.py        |  14 +-
 .../exploration/parameters/gaussian.py        |  16 +-
 .../parameters/parameter_exploration.py       |  53 +++--
 pyrobolearn/policies/basic_policy.py          |  27 ++-
 pyrobolearn/policies/cpg_policy.py            |  17 +-
 pyrobolearn/policies/dmp_policy.py            |  15 +-
 pyrobolearn/policies/nn_policy.py             |  14 +-
 pyrobolearn/policies/policy.py                | 211 ++++++++++++------
 15 files changed, 395 insertions(+), 174 deletions(-)

diff --git a/pyrobolearn/distributions/modules.py b/pyrobolearn/distributions/modules.py
index afe9a79..a8e045e 100644
--- a/pyrobolearn/distributions/modules.py
+++ b/pyrobolearn/distributions/modules.py
@@ -635,37 +635,35 @@ class DiscreteModule(torch.nn.Module):
             logits (torch.nn.Module): event logits module.
         """
         super(DiscreteModule, self).__init__()
-        self.logits = logits
-        self.probs = probs
+        if probs is None and logits is None:
+            raise ValueError("Expectingt the given 'probs' xor 'logits' to be different than None.")
+        if probs is not None and logits is not None:
+            raise ValueError("Expecting the given 'probs' xor 'logits' to be None.")
+
+        if probs is not None:
+            if not isinstance(probs, torch.nn.Module):
+                raise TypeError("Expecting the probs to be an instance of `torch.nn.Module`, instead got: "
+                                "{}".format(type(probs)))
+            self._probs = probs
+            self._logits = lambda x: None
+
+        if logits is not None:
+            if not isinstance(logits, torch.nn.Module):
+                raise TypeError("Expecting the logits to be an instance of `torch.nn.Module`, instead got: "
+                                "{}".format(type(logits)))
+            self._logits = logits
+            self._probs = lambda x: None
 
     @property
     def logits(self):
         """Return the logits module."""
         return self._logits
 
-    @logits.setter
-    def logits(self, logits):
-        """Set the logits module."""
-        if logits is not None and not isinstance(logits, torch.nn.Module):
-            raise TypeError("Expecting the logits to be an instance of `torch.nn.Module`, instead got: "
-                            "{}".format(type(logits)))
-        self._logits = logits
-        self._probs = lambda x: None
-
     @property
     def probs(self):
         """Return the probabilities module."""
         return self._probs
 
-    @probs.setter
-    def probs(self, probs):
-        """Set the probabilities module."""
-        if probs is not None and not isinstance(probs, torch.nn.Module):
-            raise TypeError("Expecting the probs to be an instance of `torch.nn.Module`, instead got: "
-                            "{}".format(type(probs)))
-        self._probs = probs
-        self._logits = lambda x: None
-
 
 class CategoricalModule(DiscreteModule):
     r"""Categorical Module
diff --git a/pyrobolearn/exploration/actions/action_exploration.py b/pyrobolearn/exploration/actions/action_exploration.py
index 5cba87d..a26ecf5 100644
--- a/pyrobolearn/exploration/actions/action_exploration.py
+++ b/pyrobolearn/exploration/actions/action_exploration.py
@@ -115,6 +115,8 @@ class ActionExploration(Exploration):
 
         # set the exploration strategies
         self._explorations = explorations
+        self.action_data = None
+        self.action_distribution = None
 
     ##############
     # Properties #
@@ -134,7 +136,21 @@ class ActionExploration(Exploration):
     # Methods #
     ###########
 
-    def act(self,  state=None, deterministic=True, to_numpy=True, return_logits=False, apply_action=True):
+    def explore(self, outputs):
+        r"""
+        Explore in the action space. Note that this does not run the policy; it is assumed that it has been called
+        outside.
+
+        Args:
+            outputs (torch.Tensor): action outputs (=logits) returned by the model.
+
+        Returns:
+            torch.Tensor: action
+            torch.distributions.Distribution: distribution on the action :math:`\pi_{\theta}(.|s)`
+        """
+        raise NotImplementedError
+
+    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.
 
@@ -149,47 +165,79 @@ class ActionExploration(Exploration):
             (list of) torch.Tensor: action(s)
             (list of) torch.distributions.Distribution: policy distribution(s) :math:`\pi_{\theta}(\cdot | s)`
         """
-        # TODO: finish to clean
-        print(state)
-        actions = self.policy.act(state, to_numpy=False, return_logits=True)
-
+        # if deterministic outcome, i.e. we don't explore just run the policy
         if deterministic:
-            return actions, None
+            self.action_data = self.policy.act(state, deterministic=True, to_numpy=to_numpy,
+                                               return_logits=return_logits, apply_action=apply_action)
 
-        # From deterministic output into stochastic outputs
-        # print("Actions before dist: {}").format(actions.train_data)
-        print("Exploration strategy - actions: {}".format(actions))
-        self.dist = self.distribution(actions)
-        actions = self.dist.sample()
-        print("Exploration strategy - sampled action: {}".format(actions))
-        if isinstance(actions, torch.Tensor):
-            if actions.requires_grad:
-                self.policy.actions.data = actions.detach().numpy()
-            else:
-                self.policy.actions.data = actions.numpy()
+            self.action_distribution = None
+
+        # if we should explore
         else:
-            self.policy.actions.data = actions
+            # if we should predict
+            if (self.policy.cnt % self.policy.rate) == 0:
+                # get the state data
+                state_data = self.policy.get_state_data(state=state)
 
-        return actions, self.dist
+                # 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
+                self.policy.postprocess(action_data)
+
+                # set the action data
+                self.action_data = self.policy.set_action_data(action_data, to_numpy=to_numpy,
+                                                               return_logits=return_logits)
+                self.action_distribution = action_distribution
+
+            # apply action
+            if apply_action:
+                self.policy.actions()
+
+            # increment policy's tick counter
+            self.policy.cnt = self.policy.cnt + 1
+
+        return self.action_data, self.action_distribution
 
     def mode(self):
         """Return the mode of the distributions."""
-        actions = self.dist.mode()
-        return actions
+        if self.action_distribution is not None:
+            if isinstance(self.action_distribution, list):
+                return [dist.mode() for dist in self.action_distribution]
+            return self.action_distribution.mode()
+        return self.action_data
 
     def sample(self):
         """Sample an action from the distribution."""
-        actions = self.dist.sample()
-        return actions
+        if self.action_distribution is None:
+            raise NotImplementedError("The action distribution has not been set.")
+        if isinstance(self.action_distribution, list):
+            return [dist.sample() for dist in self.action_distribution]
+        return self.action_distribution.sample()
 
     def action_log_prob(self, actions):
         """Return the log probability evaluated at the given actions."""
-        return self.dist.log_probs(actions)
+        return self.action_distribution.log_probs(actions)
 
     def action_prob(self, actions):
         """Return the probability evaluated at the given actions."""
-        return torch.exp(self.dist.log_probs(actions))
+        return torch.exp(self.action_distribution.log_probs(actions))
 
     def entropy(self):
         """Return the entropy of the distribution."""
-        return self.dist.entropy().mean()
+        return self.action_distribution.entropy().mean()
diff --git a/pyrobolearn/exploration/actions/boltzmann.py b/pyrobolearn/exploration/actions/boltzmann.py
index eb98a1a..26290fe 100644
--- a/pyrobolearn/exploration/actions/boltzmann.py
+++ b/pyrobolearn/exploration/actions/boltzmann.py
@@ -48,3 +48,19 @@ class BoltzmannActionExploration(DiscreteActionExploration):
         # create Categorical module
         logits = IdentityModule()
         self._module = CategoricalModule(logits=logits)
+
+    def explore(self, outputs):
+        r"""
+        Explore in the action space. Note that this does not run the policy; it is assumed that it has been called
+        outside.
+
+        Args:
+            outputs (torch.Tensor): action outputs (=logits) returned by the model.
+
+        Returns:
+            torch.Tensor: action
+            torch.distributions.Distribution: distribution on the action :math:`\pi_{\theta}(.|s)`
+        """
+        distribution = self._module(outputs)
+        action = distribution.sample((1,))
+        return action, distribution
diff --git a/pyrobolearn/exploration/actions/continuous.py b/pyrobolearn/exploration/actions/continuous.py
index 43109b8..456c85c 100644
--- a/pyrobolearn/exploration/actions/continuous.py
+++ b/pyrobolearn/exploration/actions/continuous.py
@@ -17,6 +17,8 @@ References:
     [1] "Reinforcement Learning: An Introduction", Sutton and Barto, 2018
 """
 
+from abc import ABCMeta
+
 from pyrobolearn.exploration.actions.action_exploration import ActionExploration
 
 __author__ = "Brian Delhaisse"
@@ -35,6 +37,8 @@ class ContinuousActionExploration(ActionExploration):
     Continuous action exploration strategies use continuous probability distributions on the (continuous) actions.
     """
 
+    __metaclass__ = ABCMeta
+
     def __init__(self, policy, action):
         """
         Initialize the continuous action exploration strategy.
diff --git a/pyrobolearn/exploration/actions/discrete.py b/pyrobolearn/exploration/actions/discrete.py
index 8e23b27..8dad5b9 100644
--- a/pyrobolearn/exploration/actions/discrete.py
+++ b/pyrobolearn/exploration/actions/discrete.py
@@ -17,6 +17,8 @@ References:
     [1] "Reinforcement Learning: An Introduction", Sutton and Barto, 2018
 """
 
+from abc import ABCMeta
+
 from pyrobolearn.exploration.actions.action_exploration import ActionExploration
 
 __author__ = "Brian Delhaisse"
@@ -35,6 +37,8 @@ class DiscreteActionExploration(ActionExploration):
     Discrete action exploration strategies use discrete probability distributions on the (discrete) actions.
     """
 
+    __metaclass__ = ABCMeta
+
     def __init__(self, policy, action):
         """
         Initialize the discrete action exploration strategy.
diff --git a/pyrobolearn/exploration/actions/eps_greedy.py b/pyrobolearn/exploration/actions/eps_greedy.py
index ab52693..e086e2e 100644
--- a/pyrobolearn/exploration/actions/eps_greedy.py
+++ b/pyrobolearn/exploration/actions/eps_greedy.py
@@ -8,6 +8,7 @@ it selects the best action :math:`a*` with probability :math:`p = (1 - \epsilon)
 
 import torch
 
+from pyrobolearn.distributions.categorical import Categorical
 from pyrobolearn.exploration.actions.discrete import DiscreteActionExploration
 
 
@@ -29,12 +30,33 @@ class EpsilonGreedyActionExploration(DiscreteActionExploration):
     :math:`a \in A\{a*}` randomly (based on uniform distribution) with probability :math:`p = \frac{\epsilon}{|A|-1}`.
     """
 
-    def __init__(self, policy, action):
+    def __init__(self, policy, action, epsilon=0.1):
         """
         Initialize the epsilon-greedy action exploration strategy.
 
         Args:
             policy (Policy): policy to wrap.
             action (Action): discrete actions.
+            epsilon (float): epsilon probability.
         """
         super(EpsilonGreedyActionExploration, self).__init__(policy, action=action)
+        self.epsilon = epsilon
+
+    def explore(self, outputs):
+        r"""
+        Explore in the action space. Note that this does not run the policy; it is assumed that it has been called
+        outside.
+
+        Args:
+            outputs (torch.Tensor): action outputs (=logits) returned by the model.
+
+        Returns:
+            torch.Tensor: action
+            torch.distributions.Distribution: distribution on the action :math:`\pi_{\theta}(.|s)`
+        """
+        idx = torch.argmax(outputs)
+        probs = self.epsilon/(outputs.size()[-1] - 1) * torch.ones_like(outputs)
+        probs[idx] = (1. - self.epsilon)
+        distribution = Categorical(probs=probs)
+        action = distribution.sample((1,))
+        return action, distribution
diff --git a/pyrobolearn/exploration/actions/gaussian.py b/pyrobolearn/exploration/actions/gaussian.py
index efafcff..668eb07 100644
--- a/pyrobolearn/exploration/actions/gaussian.py
+++ b/pyrobolearn/exploration/actions/gaussian.py
@@ -52,3 +52,19 @@ class GaussianActionExploration(ContinuousActionExploration):
                             "{}".format(type(module)))
 
         self._module = module
+
+    def explore(self, outputs):
+        r"""
+        Explore in the action space. Note that this does not run the policy; it is assumed that it has been called
+        outside.
+
+        Args:
+            outputs (torch.Tensor): action outputs (=logits) returned by the model.
+
+        Returns:
+            torch.Tensor: action
+            torch.distributions.Distribution: distribution on the action :math:`\pi_{\theta}(.|s)`
+        """
+        distribution = self._module(outputs)
+        action = distribution.rsample((1,))
+        return action, distribution
diff --git a/pyrobolearn/exploration/exploration.py b/pyrobolearn/exploration/exploration.py
index 6ca3e75..51333d2 100644
--- a/pyrobolearn/exploration/exploration.py
+++ b/pyrobolearn/exploration/exploration.py
@@ -76,11 +76,11 @@ class Exploration(object):
         """Reset the exploration strategy, which can be useful at the beginning of an episode."""
         self.policy.reset()
 
-    def _act(self, state=None, to_numpy=True, return_logits=False, apply_action=True):
-        """Perform the exploratory action."""
-        raise NotImplementedError
+    # def explore(self, *args, **kwargs):
+    #     """Perform the exploratory action."""
+    #     pass
 
-    def act(self, state=None, deterministic=False, to_numpy=True, return_logits=False, apply_action=True):
+    def act(self, state=None, deterministic=False, to_numpy=False, return_logits=False, apply_action=True):
         """Perform the action given the state.
 
         Args:
@@ -93,11 +93,7 @@ class Exploration(object):
         Returns:
             (list of) np.array / torch.Tensor: action data
         """
-        if deterministic:
-            return self.policy.act(state, deterministic=True, to_numpy=to_numpy, return_logits=return_logits,
-                                   apply_action=apply_action)
-        else:  # explore using the distribution
-            return self._act(state, to_numpy=to_numpy, return_logits=return_logits, apply_action=apply_action)
+        pass
 
     # def step(self, states):
     #     """Perform one step using the policy with the corresponding exploration strategy."""
diff --git a/pyrobolearn/exploration/parameters/gaussian.py b/pyrobolearn/exploration/parameters/gaussian.py
index 137e569..6f7bf6c 100644
--- a/pyrobolearn/exploration/parameters/gaussian.py
+++ b/pyrobolearn/exploration/parameters/gaussian.py
@@ -50,19 +50,27 @@ class GaussianParameterExploration(ParameterExploration):
             module = Gaussian(mean=mean, covariance=covariance)
 
         # check that the module is a Gaussian Module
-        if not isinstance(module, GaussianModule):
-            raise TypeError("Expecting the given 'module' to be an instance of `GaussianModule`, instead got: "
-                            "{}".format(type(module)))
+        if not isinstance(module, (torch.distributions.Normal, torch.distributions.MultivariateNormal)):
+            raise TypeError("Expecting the given 'module' to be an instance of `torch.distributions.Normal` or "
+                            "`torch.distributions.MultivariateNormal`, instead got: {}".format(type(module)))
 
         self._module = module
 
+    ##############
+    # Properties #
+    ##############
+
     @property
     def module(self):
         """Return the module instance."""
         return self._module
 
+    ###########
+    # Methods #
+    ###########
+
     def sample(self):
         """Sample the parameters from the """
         parameters = self.module.rsample((1,))  # rsample allows to get the gradients
-        self.policy.set_vectorized_parameters(vector=parameters)
+        return parameters
 
diff --git a/pyrobolearn/exploration/parameters/parameter_exploration.py b/pyrobolearn/exploration/parameters/parameter_exploration.py
index c131a2d..5da9429 100644
--- a/pyrobolearn/exploration/parameters/parameter_exploration.py
+++ b/pyrobolearn/exploration/parameters/parameter_exploration.py
@@ -58,43 +58,54 @@ class ParameterExploration(Exploration):
 
     def __init__(self, policy):
         super(ParameterExploration, self).__init__(policy)
+        # initial parameters
         self._parameters = policy.get_vectorized_parameters(to_numpy=False)
 
     ##############
     # Properties #
     ##############
 
-    @property
-    def parameters(self):
-        """Returns the parameters."""
-        return self._parameters
+    # @property
+    # def parameters(self):
+    #     """Returns the parameters."""
+    #     return self._parameters
 
     @property
     def size(self):
         """Returns the dimension of the parameters."""
-        return self.parameters.size(-1)
+        return self._parameters.size(-1)
 
     ###########
     # Methods #
     ###########
 
-    def reset(self):
-        # sample new set of parameters for policy
+    def sample(self):
+        """Sample a new set of parameters. To be overridden in the child class."""
         pass
 
-    def act(self, state=None, deterministic=True, to_numpy=True, return_logits=False, apply_action=True):
-        """Perform the action given the state."""
-        actions = self.policy.act(state)
+    def reset(self):
+        """Reset the parameter explorer: it samples a new set of parameters for the policy."""
+        # sample new set of parameters for policy
+        parameters = self.sample()
+        # set the parameters
+        self.policy.set_vectorized_parameters(vector=parameters)
 
+    def act(self, state=None, deterministic=True, to_numpy=False, return_logits=False, apply_action=True):
+        r"""
+        Act/Explore in the environment given the states.
 
-# class ModelUncertaintyExploration(Exploration):
-#     r"""Model Uncertainty Exploration
-#
-#     Exploration based on the uncertainty of a learned dynamic model.
-#
-#     References:
-#     [1]
-#     """
-#
-#     def __init__(self, policy):
-#         super(ModelUncertaintyExploration, self).__init__(policy)
+        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.
+            apply_action (bool): If True, it will call and execute the action.
+
+        Returns:
+            (list of) torch.Tensor: action(s)
+            (list of) torch.distributions.Distribution: policy distribution(s) :math:`\pi_{\theta}(\cdot | s)`
+        """
+        actions = self.policy.act(state=state, deterministic=deterministic, to_numpy=to_numpy,
+                                  return_logits=return_logits, apply_action=apply_action)
+        # return distribution
+        return actions, None
diff --git a/pyrobolearn/policies/basic_policy.py b/pyrobolearn/policies/basic_policy.py
index a131ef1..59bc737 100644
--- a/pyrobolearn/policies/basic_policy.py
+++ b/pyrobolearn/policies/basic_policy.py
@@ -107,8 +107,8 @@ class LinearPolicy(Policy):
         super(LinearPolicy, self).__init__(states, actions, model, rate=rate, *args, **kwargs)
 
 
-class PolicyFromValue(Policy):
-    r"""Policy From state-action value function approximator
+class PolicyFromQValue(Policy):
+    r"""Policy from state-action value function approximator
 
     This computes the optimal discrete action :math:`a` using the underlying value function approximator
     :math:`Q(s,a)` which given the state as input computes the Q-value for each discrete action. The policy select
@@ -134,9 +134,9 @@ class PolicyFromValue(Policy):
             **kwargs (dict): dictionary of arguments
         """
         self.value = value
-        super(PolicyFromValue, self).__init__(value.state, value.action, model=value, rate=rate,
-                                              preprocessors=preprocessors, postprocessors=postprocessors,
-                                              *args, **kwargs)
+        super(PolicyFromQValue, self).__init__(value.state, value.action, model=value, rate=rate,
+                                               preprocessors=preprocessors, postprocessors=postprocessors,
+                                               *args, **kwargs)
 
     ##############
     # Properties #
@@ -161,9 +161,20 @@ class PolicyFromValue(Policy):
     # Methods #
     ###########
 
-    def _predict(self, state, to_numpy=False, return_logits=True, set_output_data=False):
-        """Inner prediction step."""
-        action = self.model.compute(state, to_numpy=to_numpy)
+    def inner_predict(self, state, to_numpy=False, return_logits=True, set_output_data=False):
+        """Inner prediction step.
+
+        Args:
+            state ((list of) torch.Tensor, (list of) np.array): state data.
+            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.
+            set_output_data (bool): If True, it will set the predicted output data to the outputs given to the
+                approximator.
+
+        Returns:
+            (list of) torch.Tensor, (list of) np.array: predicted action data.
+        """
+        action = self.model.evaluate(state, to_numpy=to_numpy)
         if to_numpy:
             return np.argmax(action)
         return torch.argmax(action, dim=0, keepdim=True)
diff --git a/pyrobolearn/policies/cpg_policy.py b/pyrobolearn/policies/cpg_policy.py
index 8bb7605..d79a026 100644
--- a/pyrobolearn/policies/cpg_policy.py
+++ b/pyrobolearn/policies/cpg_policy.py
@@ -139,7 +139,7 @@ class CPGPolicy(Policy):
             # create the CPG network based on the robot kinematic structures
             cpg_network = {}
             for leg_idx, leg in enumerate(legs):
-                # compute parent/child leg coupling weight
+                # evaluate parent/child leg coupling weight
                 if len(leg) != 0:
                     if parent_coupling and init_parent_coupling_weight is None:
                         parent_coupling_weight = 1. / len(leg)
@@ -196,8 +196,19 @@ class CPGPolicy(Policy):
                                 frequency_bounds=frequency_bounds, weight_bounds=weight_bounds,
                                 bias_bounds=bias_bounds, *args, **kwargs)
 
-    def _predict(self, state, to_numpy=False, return_logits=True, set_output_data=False):
-        """Inner prediction step."""
+    def inner_predict(self, state, to_numpy=False, return_logits=True, set_output_data=False):
+        """Inner prediction step.
+
+        Args:
+            state ((list of) torch.Tensor, (list of) np.array): state data.
+            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.
+            set_output_data (bool): If True, it will set the predicted output data to the outputs given to the
+                approximator.
+
+        Returns:
+            (list of) torch.Tensor, (list of) np.array: predicted action data.
+        """
         action_data = self.model.step()
         return action_data
 
diff --git a/pyrobolearn/policies/dmp_policy.py b/pyrobolearn/policies/dmp_policy.py
index 80b37f9..a7ac9e3 100644
--- a/pyrobolearn/policies/dmp_policy.py
+++ b/pyrobolearn/policies/dmp_policy.py
@@ -39,8 +39,19 @@ class DMPPolicy(Policy):
         if not (self.is_joint_position_action or self.is_joint_velocity_action or self.is_joint_acceleration_action):
             raise ValueError("The actions do not have a joint position, velocity, or acceleration action.")
 
-    def _predict(self, state, to_numpy=False, return_logits=True, set_output_data=False):
-        """Inner prediction step."""
+    def inner_predict(self, state, to_numpy=False, return_logits=True, set_output_data=False):
+        """Inner prediction step.
+
+        Args:
+            state ((list of) torch.Tensor, (list of) np.array): state data.
+            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.
+            set_output_data (bool): If True, it will set the predicted output data to the outputs given to the
+                approximator.
+
+        Returns:
+            (list of) torch.Tensor, (list of) np.array: predicted action data.
+        """
         if isinstance(state, (np.ndarray, list, tuple)):
             state = state[0]
         y, dy, ddy = self.model.step(state)
diff --git a/pyrobolearn/policies/nn_policy.py b/pyrobolearn/policies/nn_policy.py
index 6106e44..c0ea81b 100644
--- a/pyrobolearn/policies/nn_policy.py
+++ b/pyrobolearn/policies/nn_policy.py
@@ -71,8 +71,8 @@ class MLPPolicy(NNPolicy):
     activation functions.
     """
 
-    def __init__(self, states, actions, hidden_units=(), activation_fct='linear', last_activation_fct=None,
-                 dropout_prob=None, rate=1, preprocessors=None, postprocessors=None):
+    def __init__(self, states, actions, hidden_units=(), activation='linear', last_activation=None,
+                 dropout=None, rate=1, preprocessors=None, postprocessors=None):
         """Initialize MLP policy.
 
         Args:
@@ -81,12 +81,12 @@ class MLPPolicy(NNPolicy):
             actions (Action): 1D-actions outputted by the policy and will be applied in the simulator (the output
                               dimensions will be inferred from the actions)
             hidden_units (list/tuple of int): number of hidden units in the corresponding layer
-            activation_fct (None, str, or list/tuple of str/None): activation function to be applied after each layer.
+            activation (None, str, or list/tuple of str/None): activation function to be applied after each layer.
                                                                    If list/tuple, then it has to match the
-            last_activation_fct (None or str): last activation function to be applied. If not specified, it will check
+            last_activation (None or str): last activation function to be applied. If not specified, it will check
                                                if it is in the list/tuple of activation functions provided for the
                                                previous argument.
-            dropout_prob (None, float, or list/tuple of float/None): dropout probability.
+            dropout (None, float, or list/tuple of float/None): dropout probability.
             rate (int, float): rate (float) at which the policy operates if we are operating in real-time. If we are
                 stepping deterministically in the simulator, it represents the number of ticks (int) to sleep before
                 executing the model.
@@ -94,8 +94,8 @@ class MLPPolicy(NNPolicy):
             postprocessors (Processor, list of Processor, None): post-processors to be applied to the output
         """
         model = MLPApproximator(states, actions, hidden_units=hidden_units,
-                                activation_fct=activation_fct, last_activation_fct=last_activation_fct,
-                                dropout_prob=dropout_prob, preprocessors=preprocessors, postprocessors=postprocessors)
+                                activation=activation, last_activation=last_activation,
+                                dropout=dropout, preprocessors=preprocessors, postprocessors=postprocessors)
         super(MLPPolicy, self).__init__(states, actions, model, rate=rate)
 
     # def act(self, state, deterministic=True):
diff --git a/pyrobolearn/policies/policy.py b/pyrobolearn/policies/policy.py
index 7bc8183..2d8e205 100644
--- a/pyrobolearn/policies/policy.py
+++ b/pyrobolearn/policies/policy.py
@@ -144,7 +144,6 @@ class Policy(object):
         self.states = states
         self.actions = actions
         self.model = model
-        self.train_mode = False
         self.rate = rate
         self.cnt = 0
         self.action_data = None
@@ -391,13 +390,136 @@ class Policy(object):
             return x.numpy()
         return x
 
-    def _predict(self, state, to_numpy=False, return_logits=True, set_output_data=False):
-        """Inner prediction step."""
+    def get_state_data(self, state=None):
+        """Get the state data to be feed to the preprocessors and to the inner model / approximator.
+
+        Args:
+            state (State, None, (list of) torch.Tensor, (list of) np.array): state
+
+        Returns:
+            (list of) torch.Tensor, (list of) np.array: state data
+        """
+        # if no input is given, take the provided inputs at the beginning
+        if state is None:
+            state = self.states
+
+        # 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]
+
+        return state
+
+    def preprocess(self, state_data):
+        """Pre-process the given state data.
+
+        Args:
+            state_data ((list of) torch.Tensor, (list of) np.array): state data.
+
+        Returns:
+            (list of) torch.Tensor, (list of) np.array: preprocessed state data
+        """
+        # go through each preprocessor
+        for processor in self.preprocessors:
+            state_data = processor(state_data)
+        return state_data
+
+    def inner_predict(self, state_data, deterministic=True, to_numpy=False, return_logits=True, set_output_data=False):
+        """Inner prediction step.
+
+        Args:
+            state_data ((list of) torch.Tensor, (list of) np.array): state data.
+            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.
+            set_output_data (bool): If True, it will set the predicted output data to the outputs given to the
+                approximator.
+
+        Returns:
+            (list of) torch.Tensor, (list of) np.array: predicted action data.
+        """
         if isinstance(self.model, Approximator):  # inner model is an approximator
-            return self.model.predict(state, to_numpy=to_numpy, return_logits=return_logits,
+            return self.model.predict(state_data, to_numpy=to_numpy, return_logits=return_logits,
                                       set_output_data=set_output_data)
         # inner model is a learning model
-        return self.model.predict(state, to_numpy=to_numpy)
+        return self.model.predict(state_data, to_numpy=to_numpy)
+
+    def postprocess(self, action_data):
+        """Post-process the given action data.
+
+        Args:
+            action_data ((list of) torch.Tensor, (list of) np.array): action data.
+
+        Returns:
+            (list of) torch.Tensor, (list of) np.array: post-processed action data.
+        """
+        # go through each postprocessor
+        for processor in self.postprocessors:
+            action_data = processor(action_data)
+        return action_data
+
+    def set_action_data(self, action_data, to_numpy=True, return_logits=False):
+        """Set the given action data.
+
+        Args:
+            action_data ((list of) torch.Tensor, (list of) np.array): action data.
+            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, (list of) np.array: action data
+        """
+        # set the action data
+        if action_data is None:
+            action_data = self.action_data
+
+        # if action data is not a list, make it a list as we will iterate through it
+        if not isinstance(action_data, list):
+            action_data = [action_data]
+
+        # go through each action and data
+        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
+                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)
+                # elif isinstance(data, (float, int)):
+                #     discrete_data = np.argmax(data)
+                #     action.data = discrete_data
+                #     if not return_logits:
+                #         action_data[idx] = discrete_data
+                else:
+                    raise TypeError(
+                        "Expecting the `data` action to be a numpy array or torch.Tensor, instead got: "
+                        "{}".format(type(data)))
+            else:  # continuous action
+                if isinstance(data, np.ndarray):
+                    action.data = data
+                elif isinstance(data, torch.Tensor):
+                    action.torch_data = data
+                    action_data[idx] = self.__convert_to_numpy(data, to_numpy=to_numpy)
+                # elif isinstance(data, (float, int)):
+                #     pass
+                else:
+                    raise TypeError(
+                        "Expecting `data` to be a numpy array or torch.Tensor, instead got: "
+                        "{}".format(type(data)))
+
+        # if action_data is a list and has one element, return just that element
+        if isinstance(action_data, list) and len(action_data) == 1:
+            action_data = action_data[0]
+
+        return action_data
 
     def act(self, state=None, deterministic=True, to_numpy=True, return_logits=False, apply_action=True):
         """Perform the action given the state.
@@ -414,79 +536,25 @@ class Policy(object):
         """
         # if we should predict
         if (self.cnt % self.rate) == 0:
+            # get the state data
+            state_data = self.get_state_data(state=state)
 
-            # if no input is given, take the provided inputs at the beginning
-            if state is None:
-                state = self.states
-
-            # 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]
-
-            # go through each preprocessor
-            for processor in self.preprocessors:
-                state = processor(state)
+            # pre-process the state data
+            state_data = self.preprocess(state_data)
 
             # predict the output using the inner model
-            self.action_data = self._predict(state, to_numpy=False, return_logits=True, set_output_data=False)
+            action_data = self.inner_predict(state_data, to_numpy=False, return_logits=True, set_output_data=False)
             # if isinstance(self.model, Approximator):  # inner model is an approximator
             #     self.action_data = self.model.predict(state, to_numpy=False, return_logits=True,
             #                                           set_output_data=False)
             # else:  # inner model is a learning model
             #     self.action_data = self.model.predict(state, to_numpy=to_numpy)
 
-            # go through each postprocessor
-            for processor in self.postprocessors:
-                self.action_data = processor(self.action_data)
+            # post-process the action data
+            action_data = self.postprocess(action_data)
 
             # set the action data
-
-            # if action data is not a list, make it a list as we will iterate through it
-            if not isinstance(self.action_data, list):
-                self.action_data = [self.action_data]
-
-            # go through each action and data
-            for idx, (action, data) in enumerate(zip(self.actions, self.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:
-                            self.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:
-                            self.action_data[idx] = self.__convert_to_numpy(discrete_data, to_numpy=to_numpy)
-                        else:
-                            self.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
-                    #     if not return_logits:
-                    #         self.action_data[idx] = discrete_data
-                    else:
-                        raise TypeError(
-                            "Expecting the `data` action to be a numpy array or torch.Tensor, instead got: "
-                            "{}".format(type(data)))
-                else:  # continuous action
-                    if isinstance(data, np.ndarray):
-                        action.data = data
-                    elif isinstance(data, torch.Tensor):
-                        action.torch_data = data
-                        self.action_data[idx] = self.__convert_to_numpy(data, to_numpy=to_numpy)
-                    # elif isinstance(data, (float, int)):
-                    #     pass
-                    else:
-                        raise TypeError(
-                            "Expecting `data` to be a numpy array or torch.Tensor, instead got: "
-                            "{}".format(type(data)))
-
-            # if action_data is a list and has one element, return just that element
-            if isinstance(self.action_data, list) and len(self.action_data) == 1:
-                self.action_data = self.action_data[0]
+            self.action_data = self.set_action_data(action_data, to_numpy=to_numpy, return_logits=return_logits)
 
         # apply action
         if apply_action:
@@ -510,20 +578,17 @@ class Policy(object):
         """
         pass
 
-    def train(self, mode=True):
+    def train(self):
         """
         Set the policy in training mode.
-
-        Args:
-            mode (bool): if True, set the policy in train mode.
         """
-        self.train_mode = mode
+        self.model.train()
 
     def eval(self):
         """
         Set the policy in evaluation mode.
         """
-        self.train(mode=False)
+        self.model.eval()
 
     def reset(self, reset_processors=False, *args, **kwargs):
         """