From 5ed2ddd3c3af6284139f8e950fb663d6a013e7e8 Mon Sep 17 00:00:00 2001
From: Brian Delhaisse <briandelhaisse@gmail.com>
Date: Wed, 15 May 2019 09:45:10 +0200
Subject: [PATCH] fix few bugs in algos and others

---
 pyrobolearn/actions/gym_actions.py            |  2 +-
 .../actions/robot_actions/actuator_actions.py | 62 +++++++++++++++++++
 .../actions/robot_actions/joint_actions.py    | 10 +--
 .../actions/robot_actions/link_actions.py     |  2 +-
 .../actions/robot_actions/robot_actions.py    |  2 +-
 pyrobolearn/actorcritics/actorcritic.py       | 54 +++++++++++++++-
 pyrobolearn/algos/ddpg.py                     | 16 +++--
 pyrobolearn/algos/dqn.py                      | 10 ++-
 pyrobolearn/algos/evaluator.py                |  2 +-
 pyrobolearn/algos/reinforce.py                |  2 +-
 pyrobolearn/algos/sac.py                      |  6 +-
 pyrobolearn/algos/td3.py                      | 14 +++--
 pyrobolearn/algos/updater.py                  |  2 +-
 pyrobolearn/dynamics/dynamic.py               | 14 ++++-
 pyrobolearn/robots/actuators/actuator.py      |  7 +--
 pyrobolearn/robots/sensors/joints.py          |  3 +-
 pyrobolearn/robots/sensors/links.py           |  2 +-
 pyrobolearn/robots/sensors/sensor.py          |  4 +-
 pyrobolearn/states/gym_states.py              | 20 ++++++
 .../states/robot_states/joint_states.py       | 23 +++++++
 .../states/robot_states/link_states.py        | 23 +++++++
 .../states/robot_states/robot_states.py       | 20 ++++++
 .../states/robot_states/sensor_states.py      | 20 ++++++
 pyrobolearn/storages/er.py                    |  9 +--
 pyrobolearn/storages/storage.py               |  5 ++
 .../terminal_conditions/terminal_condition.py | 50 +++++++++++++--
 pyrobolearn/values/basic_value.py             | 60 +++++++++++++++++-
 pyrobolearn/values/value.py                   | 36 ++++++++++-
 28 files changed, 429 insertions(+), 51 deletions(-)
 create mode 100644 pyrobolearn/actions/robot_actions/actuator_actions.py

diff --git a/pyrobolearn/actions/gym_actions.py b/pyrobolearn/actions/gym_actions.py
index c7d53bf..b2c082e 100644
--- a/pyrobolearn/actions/gym_actions.py
+++ b/pyrobolearn/actions/gym_actions.py
@@ -61,7 +61,7 @@ class GymAction(Action):
         """
         if self in memo:
             return memo[self]
-        env = copy.deepcopy(self.env)
+        env = memo.get(self.env, self.env)  # copy.deepcopy(self.env, memo)
         action = self.__class__(gym_env=env)
         memo[self] = action
         return action
diff --git a/pyrobolearn/actions/robot_actions/actuator_actions.py b/pyrobolearn/actions/robot_actions/actuator_actions.py
new file mode 100644
index 0000000..699a118
--- /dev/null
+++ b/pyrobolearn/actions/robot_actions/actuator_actions.py
@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+"""Define the various actuator actions
+"""
+
+from abc import ABCMeta
+import collections
+import numpy as np
+
+from pyrobolearn.actions.action import Action
+from pyrobolearn.robots.actuators.actuator import Actuator
+
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class ActuatorAction(Action):
+    r"""Actuator action (abstract class)
+    """
+    __metaclass__ = ABCMeta
+
+    def __init__(self, actuators, ticks=1):
+        """
+        Initialize the sensor state.
+
+        Args:
+            actuators (A, list of Actuator): actuator(s).
+            ticks (int): number of ticks to sleep before setting the next action data.
+        """
+        if not isinstance(actuators, collections.Iterable):
+            actuators = [actuators]
+        for actuator in actuators:
+            if not isinstance(actuator, Actuator):
+                raise TypeError("Expecting the given 'actuator' to be an instance of `Actuator`, instead got: "
+                                "{}".format(type(actuator)))
+        self.actuators = actuators
+        super(ActuatorAction, self).__init__(ticks=ticks)
+
+    def __copy__(self):
+        """Return a shallow copy of the action. This can be overridden in the child class."""
+        return self.__class__(actuators=self.actuators, ticks=self.ticks)
+
+    def __deepcopy__(self, memo={}):
+        """Return a deep copy of the action. This can be overridden in the child class.
+
+        Args:
+            memo (dict): memo dictionary of objects already copied during the current copying pass
+        """
+        if self in memo:
+            return memo[self]
+
+        actuators = copy.deepcopy(self.actuators, memo)
+        action = self.__class__(actuators=actuators, ticks=self.ticks)
+
+        memo[self] = action
+        return action
diff --git a/pyrobolearn/actions/robot_actions/joint_actions.py b/pyrobolearn/actions/robot_actions/joint_actions.py
index d7060f4..cf4ca26 100644
--- a/pyrobolearn/actions/robot_actions/joint_actions.py
+++ b/pyrobolearn/actions/robot_actions/joint_actions.py
@@ -62,7 +62,7 @@ class JointAction(RobotAction):
         """
         if self in memo:
             return memo[self]
-        robot = copy.deepcopy(self.robot, memo)
+        robot = memo.get(self.robot, self.robot)  # copy.deepcopy(self.robot, memo)
         joints = copy.deepcopy(self.joints)
         action = self.__class__(robot=robot, joint_ids=joints)
         memo[self] = action
@@ -98,7 +98,7 @@ class JointPositionAction(JointAction):
         """
         if self in memo:
             return memo[self]
-        robot = copy.deepcopy(self.robot, memo)
+        robot = memo.get(self.robot, self.robot)  # copy.deepcopy(self.robot, memo)
         joints = copy.deepcopy(self.joints)
         kp = copy.deepcopy(self.kp)
         kd = copy.deepcopy(self.kd)
@@ -159,7 +159,7 @@ class JointPositionAndVelocityAction(JointAction):
         """
         if self in memo:
             return memo[self]
-        robot = copy.deepcopy(self.robot, memo)
+        robot = memo.get(self.robot, self.robot)  # copy.deepcopy(self.robot, memo)
         joints = copy.deepcopy(self.joints)
         kp = copy.deepcopy(self.kp)
         kd = copy.deepcopy(self.kd)
@@ -209,7 +209,7 @@ class JointForceAction(JointAction):
         """
         if self in memo:
             return memo[self]
-        robot = copy.deepcopy(self.robot, memo)
+        robot = memo.get(self.robot, self.robot)  # copy.deepcopy(self.robot, memo)
         joints = copy.deepcopy(self.joints)
         f_min = copy.deepcopy(self.f_min)
         f_max = copy.deepcopy(self.f_max)
@@ -251,7 +251,7 @@ class JointAccelerationAction(JointAction):
         """
         if self in memo:
             return memo[self]
-        robot = copy.deepcopy(self.robot, memo)
+        robot = memo.get(self.robot, self.robot)  # copy.deepcopy(self.robot, memo)
         joints = copy.deepcopy(self.joints)
         a_min = copy.deepcopy(self.a_min)
         a_max = copy.deepcopy(self.a_max)
diff --git a/pyrobolearn/actions/robot_actions/link_actions.py b/pyrobolearn/actions/robot_actions/link_actions.py
index cbf832e..eca4d6f 100644
--- a/pyrobolearn/actions/robot_actions/link_actions.py
+++ b/pyrobolearn/actions/robot_actions/link_actions.py
@@ -52,7 +52,7 @@ class LinkAction(RobotAction):
         """
         if self in memo:
             return memo[self]
-        robot = copy.deepcopy(self.robot, memo)
+        robot = memo.get(self.robot, self.robot)  # copy.deepcopy(self.robot, memo)
         links = copy.deepcopy(self.links)
         action = self.__class__(robot, links)
         memo[self] = action
diff --git a/pyrobolearn/actions/robot_actions/robot_actions.py b/pyrobolearn/actions/robot_actions/robot_actions.py
index dd3fb57..d9d6508 100644
--- a/pyrobolearn/actions/robot_actions/robot_actions.py
+++ b/pyrobolearn/actions/robot_actions/robot_actions.py
@@ -61,7 +61,7 @@ class RobotAction(Action):
         """
         if self in memo:
             return memo[self]
-        robot = copy.deepcopy(self.robot, memo)
+        robot = memo.get(self.robot, self.robot)  # copy.deepcopy(self.robot, memo)
         action = self.__class__(robot=robot)
         memo[self] = action
         return action
diff --git a/pyrobolearn/actorcritics/actorcritic.py b/pyrobolearn/actorcritics/actorcritic.py
index aeaf138..6df4abf 100644
--- a/pyrobolearn/actorcritics/actorcritic.py
+++ b/pyrobolearn/actorcritics/actorcritic.py
@@ -13,7 +13,7 @@ import itertools
 import torch
 
 from pyrobolearn.policies import Policy
-from pyrobolearn.values import ValueApproximator
+from pyrobolearn.values import Value
 
 __author__ = "Brian Delhaisse"
 __copyright__ = "Copyright 2018, PyRoboLearn"
@@ -68,8 +68,8 @@ class ActorCritic(object):
     @critic.setter
     def critic(self, critic):
         """Set the critic."""
-        if not isinstance(critic, (ValueApproximator, torch.nn.Module)):
-            raise TypeError("Expecting the critic to be an instance of 'ValueApproximator' or 'torch.nn.Module', "
+        if not isinstance(critic, (Value, torch.nn.Module)):
+            raise TypeError("Expecting the critic to be an instance of 'Value' or 'torch.nn.Module', "
                             "instead got {}".format(type(critic)))
         self._critic = critic
 
@@ -87,11 +87,59 @@ class ActorCritic(object):
     # Methods #
     ###########
 
+    def train(self):
+        """
+        Set the actor and critic in training mode.
+        """
+        self.actor.train()
+        self.critic.train()
+
+    def eval(self):
+        """
+        Set the actor and critic in evaluation mode.
+        """
+        self.actor.eval()
+        self.critic.eval()
+
     def parameters(self):
         """Return the parameters of first the actor then the critic."""
         generator = itertools.chain(self.actor.parameters(), self.critic.parameters())
         return generator
 
+    def named_parameters(self):
+        """
+        Return an iterator over the learning model parameters; yielding both the name and the parameter itself.
+        """
+        generator = itertools.chain(self.actor.named_parameters(), self.critic.named_parameters())
+        return generator
+
+    def list_parameters(self):
+        """
+        Return the learning model parameters.
+        """
+        return self.actor.list_parameters() + self.critic.list_parameters()
+
+    def get_vectorized_parameters(self, to_numpy=True):
+        """
+        Get the parameters in a vectorized form.
+
+        Args:
+            to_numpy (bool): if True, it will convert the 1D parameter vector into a numpy array.
+        """
+        return self.actor.get_vectorized_parameters(to_numpy=to_numpy), \
+               self.critic.get_vectorized_parameters(to_numpy=to_numpy)
+
+    def set_vectorized_parameters(self, actor_parameters, critic_parameters):
+        """
+        Set the vectorized parameters.
+
+        Args:
+            actor_parameters (np.array, torch.Tensor): 1D parameter vector for the actor.
+            critic_parameters (np.array, torch.Tensor): 1D parameter vector for the critic.
+        """
+        self.actor.set_vectorized_parameters(vector=actor_parameters)
+        self.critic.set_vectorized_parameters(vector=critic_parameters)
+
     def value(self, x):
         """Compute the value function."""
         return self.evaluate(x)
diff --git a/pyrobolearn/algos/ddpg.py b/pyrobolearn/algos/ddpg.py
index d908d04..0a0ebcc 100755
--- a/pyrobolearn/algos/ddpg.py
+++ b/pyrobolearn/algos/ddpg.py
@@ -231,7 +231,7 @@ class DDPG(GradientRLAlgo):
         Args:
             task (RLTask, Env): RL task/env to run
             approximators ([Policy, QValue]): policy and Q-value function approximator to optimize.
-            gamma (float): discount factor (which is a bias-variance tradeoff). This parameter describes how much
+            gamma (float): discount factor (which is a bias-variance trade-off). This parameter describes how much
                 importance has the future rewards we get.
             lr (float): learning rate
             polyak (float): coefficient (between 0 and 1) used in the polyak averaging when updating the target
@@ -261,20 +261,24 @@ class DDPG(GradientRLAlgo):
         else:
             raise TypeError("Expecting a list/tuple of a policy and a Q-value function.")
 
+        # get states and actions from policy
+        states, actions = policy.states, policy.actions
+
         # check that the actions are continuous
-        actions = policy.actions
         if not actions.is_continuous():
             raise ValueError("The DDPG assumes that the actions are continuous, however got an action which is not.")
 
         # Set target parameters equal to main parameters
-        q_target = copy.deepcopy(q_value)
-        policy_target = copy.deepcopy(policy)
+        memo = {}
+        q_target = copy.deepcopy(q_value, memo=memo)
+        policy_target = copy.deepcopy(policy, memo=memo)
 
         # create action exploration strategy
-        exploration = ActionExploration(policy=policy, action=policy.actions)
+        exploration = ActionExploration(policy=policy, action=actions)
 
         # create experience replay
-        storage = ExperienceReplay(state_shapes=policy.states, action_shapes=policy.actions, capacity=capacity)
+        storage = ExperienceReplay(state_shapes=states.merged_shape, action_shapes=actions.merged_shape,
+                                   capacity=capacity)
         sampler = BatchRandomSampler(storage)
 
         # create target return estimator
diff --git a/pyrobolearn/algos/dqn.py b/pyrobolearn/algos/dqn.py
index cbb317b..e71da8b 100755
--- a/pyrobolearn/algos/dqn.py
+++ b/pyrobolearn/algos/dqn.py
@@ -138,13 +138,17 @@ class DQN(GradientRLAlgo):
                             "`ParametrizedQValueOutput`, instead got: {}".format(type(approximator)))
 
         # evaluate target Q-value fct by copying Q-value function approximator
-        q_target = copy.deepcopy(q_value)
+        q_target = copy.deepcopy(q_value, memo={})
+
+        # get states and actions from policy
+        states, actions = policy.states, policy.actions
 
         # create action exploration strategy
-        exploration = EpsilonGreedyActionExploration(policy=policy, action=policy.actions)
+        exploration = EpsilonGreedyActionExploration(policy=policy, action=actions)
 
         # create experience replay and sampler
-        storage = ExperienceReplay(capacity=capacity)
+        storage = ExperienceReplay(state_shapes=states.merged_shape, action_shapes=actions.merged_shape,
+                                   capacity=capacity)
         sampler = BatchRandomSampler(storage)
 
         # create target return estimator
diff --git a/pyrobolearn/algos/evaluator.py b/pyrobolearn/algos/evaluator.py
index edfd5c3..e781aaf 100644
--- a/pyrobolearn/algos/evaluator.py
+++ b/pyrobolearn/algos/evaluator.py
@@ -52,7 +52,7 @@ class Evaluator(object):
     @estimator.setter
     def estimator(self, estimator):
         """Set the estimator."""
-        if not None and not isinstance(estimator, Estimator):
+        if estimator is not None and not isinstance(estimator, Estimator):
             raise TypeError("Expecting estimator to be an instance of `Estimator` or None, instead got: "
                             "{}".format(type(estimator)))
         self._estimator = estimator
diff --git a/pyrobolearn/algos/reinforce.py b/pyrobolearn/algos/reinforce.py
index 52c796f..bc2a355 100755
--- a/pyrobolearn/algos/reinforce.py
+++ b/pyrobolearn/algos/reinforce.py
@@ -188,7 +188,7 @@ class REINFORCE(GradientRLAlgo):
 
         # create storage
         states, actions = policy.states, policy.actions
-        storage = RolloutStorage(num_steps=1000, state_shapes=states.shape, action_shapes=actions.shape,
+        storage = RolloutStorage(num_steps=1000, state_shapes=states.merged_shape, action_shapes=actions.merged_shape,
                                  num_trajectories=1)
         sampler = StorageSampler(storage)
 
diff --git a/pyrobolearn/algos/sac.py b/pyrobolearn/algos/sac.py
index 4ec2c51..c9153db 100755
--- a/pyrobolearn/algos/sac.py
+++ b/pyrobolearn/algos/sac.py
@@ -305,10 +305,12 @@ class SAC(GradientRLAlgo):
             raise TypeError("No Q-value function approximators were given to the algorithm.")
 
         # set target parameters equal to main parameters for the value function
-        value_target = copy.deepcopy(value)
+        value_target = copy.deepcopy(value, memo={})
 
         # create experience replay
-        storage = ExperienceReplay(state_shapes=policy.states, action_shapes=policy.actions, capacity=capacity)
+        states, actions = policy.states, policy.actions
+        storage = ExperienceReplay(state_shapes=states.merged_shape, action_shapes=actions.merged_shape,
+                                   capacity=capacity)
         sampler = BatchRandomSampler(storage)
 
         # create action exploration
diff --git a/pyrobolearn/algos/td3.py b/pyrobolearn/algos/td3.py
index 4fab72f..da7aea0 100755
--- a/pyrobolearn/algos/td3.py
+++ b/pyrobolearn/algos/td3.py
@@ -202,20 +202,24 @@ class TD3(GradientRLAlgo):
         if len(q_values) < 2:
             raise ValueError("Expecting at least 2 Q-value function approximators for the TD3 algorithm.")
 
+        # get states and actions from policy
+        states, actions = policy.states, policy.actions
+
         # check that the actions are continuous
-        actions = policy.actions
         if not actions.is_continuous():
             raise ValueError("The TD3 assumes that the actions are continuous, however got an action which is not.")
 
         # evaluate target Q-value fct by copying Q-value function approximator
-        q_targets = [copy.deepcopy(q_value) for q_value in q_values]
-        policy_target = copy.deepcopy(policy)
+        memo = {}
+        q_targets = [copy.deepcopy(q_value, memo=memo) for q_value in q_values]
+        policy_target = copy.deepcopy(policy, memo=memo)
 
         # create action exploration strategy
-        exploration = ActionExploration(policy=policy, action=policy.actions)
+        exploration = ActionExploration(policy=policy, action=actions)
 
         # create experience replay
-        storage = ExperienceReplay(state_shapes=policy.states, action_shapes=policy.actions, capacity=capacity)
+        storage = ExperienceReplay(state_shapes=states.merged_shape, action_shapes=actions.merged_shape,
+                                   capacity=capacity)
         sampler = BatchRandomSampler(storage)
 
         # create target return estimator
diff --git a/pyrobolearn/algos/updater.py b/pyrobolearn/algos/updater.py
index 3515bbf..6a08166 100644
--- a/pyrobolearn/algos/updater.py
+++ b/pyrobolearn/algos/updater.py
@@ -308,7 +308,7 @@ class Updater(object):
                     if self._cnt % self.ticks[loss] == 0:
 
                         if verbose:
-                            print("\t Compute loss {}".format(loss))
+                            print("\t Compute loss: {}".format(loss))
 
                         # compute loss on the data (the loss knows what to do with the batch)
                         loss_value = loss.compute(batch)
diff --git a/pyrobolearn/dynamics/dynamic.py b/pyrobolearn/dynamics/dynamic.py
index 56e310e..74be6ef 100644
--- a/pyrobolearn/dynamics/dynamic.py
+++ b/pyrobolearn/dynamics/dynamic.py
@@ -442,6 +442,18 @@ class ParametrizedDynamicModel(DynamicModel):
     # Methods #
     ###########
 
+    def train(self):
+        """
+        Set the dynamic transition model in training mode.
+        """
+        self.model.train()
+
+    def eval(self):
+        """
+        Set the dynamic transition model in evaluation mode.
+        """
+        self.model.eval()
+
     def parameters(self):
         """
         Return an iterator over the learning model parameters.
@@ -493,7 +505,7 @@ class ParametrizedDynamicModel(DynamicModel):
         Set the vectorized parameters.
 
         Args:
-            np.array, torch.Tensor: 1D parameter vector.
+            vector (np.array, torch.Tensor): 1D parameter vector.
         """
         self.model.set_vectorized_parameters(vector=vector)
 
diff --git a/pyrobolearn/robots/actuators/actuator.py b/pyrobolearn/robots/actuators/actuator.py
index 9943218..d840f90 100644
--- a/pyrobolearn/robots/actuators/actuator.py
+++ b/pyrobolearn/robots/actuators/actuator.py
@@ -69,11 +69,8 @@ class Actuator(object):
         Args:
             memo (dict): memo dictionary of objects already copied during the current copying pass
         """
+        if self in memo:
+            return memo[self]
         actuator = self.__class__()
-
-        # update the memodict (note that `copy.deepcopy` will automatically check this dictionary and return the
-        # reference if already present)
         memo[self] = actuator
-
-        # return the copy
         return actuator
diff --git a/pyrobolearn/robots/sensors/joints.py b/pyrobolearn/robots/sensors/joints.py
index 143409c..869f5e8 100644
--- a/pyrobolearn/robots/sensors/joints.py
+++ b/pyrobolearn/robots/sensors/joints.py
@@ -62,7 +62,7 @@ class JointSensor(Sensor):
         Args:
             memo (dict): memo dictionary of objects already copied during the current copying pass
         """
-        simulator = copy.deepcopy(self.simulator, memo)
+        simulator = memo.get(self.simulator, self.simulator)  # copy.deepcopy(self.simulator, memo)
         body_id = copy.deepcopy(self.body_id)
         joint_id = copy.deepcopy(self.joint_id)
         position = copy.deepcopy(self.local_position)
@@ -73,7 +73,6 @@ class JointSensor(Sensor):
         return sensor
 
 
-
 class Encoder(JointSensor):
     r"""Encoder joint sensor
 
diff --git a/pyrobolearn/robots/sensors/links.py b/pyrobolearn/robots/sensors/links.py
index ad462dc..fbd0a92 100644
--- a/pyrobolearn/robots/sensors/links.py
+++ b/pyrobolearn/robots/sensors/links.py
@@ -76,7 +76,7 @@ class LinkSensor(Sensor):
         Args:
             memo (dict): memo dictionary of objects already copied during the current copying pass
         """
-        simulator = copy.deepcopy(self.simulator, memo)
+        simulator = memo.get(self.simulator, self.simulator)  # copy.deepcopy(self.simulator, memo)
         body_id = copy.deepcopy(self.body_id)
         link_id = copy.deepcopy(self.link_id)
         position = copy.deepcopy(self.local_position)
diff --git a/pyrobolearn/robots/sensors/sensor.py b/pyrobolearn/robots/sensors/sensor.py
index 2fad3fc..d0cf4a9 100644
--- a/pyrobolearn/robots/sensors/sensor.py
+++ b/pyrobolearn/robots/sensors/sensor.py
@@ -130,7 +130,9 @@ class Sensor(object):  # sensor attached to a link or joint
         Args:
             memo (dict): memo dictionary of objects already copied during the current copying pass
         """
-        simulator = copy.deepcopy(self.simulator, memo)
+        if self in memo:
+            return memo[self]
+        simulator = memo.get(self.simulator, self.simulator)  # copy.deepcopy(self.simulator, memo)
         body_id = copy.deepcopy(self.body_id)
         position = copy.deepcopy(self.local_position)
         orientation = copy.deepcopy(self.local_orientation)
diff --git a/pyrobolearn/states/gym_states.py b/pyrobolearn/states/gym_states.py
index ee2cc99..eacb8cd 100644
--- a/pyrobolearn/states/gym_states.py
+++ b/pyrobolearn/states/gym_states.py
@@ -8,6 +8,7 @@ the various policies defined in the pyrobolearn framework.
 """
 
 import gym
+import copy
 
 from pyrobolearn.states.state import State
 
@@ -71,6 +72,25 @@ class GymState(State):
     def _read(self):
         pass
 
+    def __copy__(self):
+        """Return a shallow copy of the state. This can be overridden in the child class."""
+        return self.__class__(gym_env=self.env, window_size=self.window_size, axis=self.axis, ticks=self.ticks)
+
+    def __deepcopy__(self, memo={}):
+        """Return a deep copy of the state. This can be overridden in the child class.
+
+        Args:
+            memo (dict): memo dictionary of objects already copied during the current copying pass
+        """
+        if self in memo:
+            return memo[self]
+
+        env = memo.get(self.env, self.env)  # copy.deepcopy(self.env, memo)
+        state = self.__class__(gym_env=env, window_size=self.window_size, axis=self.axis, ticks=self.ticks)
+
+        memo[self] = state
+        return state
+
 
 # Tests
 if __name__ == '__main__':
diff --git a/pyrobolearn/states/robot_states/joint_states.py b/pyrobolearn/states/robot_states/joint_states.py
index 27edcb9..06b5f94 100644
--- a/pyrobolearn/states/robot_states/joint_states.py
+++ b/pyrobolearn/states/robot_states/joint_states.py
@@ -4,6 +4,7 @@
 This includes notably the joint positions, velocities, and force/torque states.
 """
 
+import copy
 from abc import ABCMeta
 
 from pyrobolearn.states.robot_states.robot_states import RobotState, Robot
@@ -58,6 +59,28 @@ class JointState(RobotState):
 
         super(JointState, self).__init__(robot, window_size=window_size, axis=axis, ticks=ticks)
 
+    def __copy__(self):
+        """Return a shallow copy of the state. This can be overridden in the child class."""
+        return self.__class__(robot=self.robot, joint_ids=self.joints, window_size=self.window_size, axis=self.axis,
+                              ticks=self.ticks)
+
+    def __deepcopy__(self, memo={}):
+        """Return a deep copy of the state. This can be overridden in the child class.
+
+        Args:
+            memo (dict): memo dictionary of objects already copied during the current copying pass
+        """
+        if self in memo:
+            return memo[self]
+
+        robot = memo.get(self.robot, self.robot)  # copy.deepcopy(self.robot, memo)
+        joint_ids = copy.deepcopy(self.joints)
+        state = self.__class__(robot=robot, joint_ids=joint_ids, window_size=self.window_size, axis=self.axis,
+                               ticks=self.ticks)
+
+        memo[self] = state
+        return state
+
 
 class JointPositionState(JointState):
     r"""Joint Position State
diff --git a/pyrobolearn/states/robot_states/link_states.py b/pyrobolearn/states/robot_states/link_states.py
index e123234..f7c4be0 100644
--- a/pyrobolearn/states/robot_states/link_states.py
+++ b/pyrobolearn/states/robot_states/link_states.py
@@ -4,6 +4,7 @@
 This includes notably the link positions and velocities.
 """
 
+import copy
 from abc import ABCMeta
 
 from pyrobolearn.states.robot_states.robot_states import RobotState, Robot
@@ -57,6 +58,28 @@ class LinkState(RobotState):
         # call parent constructor
         super(LinkState, self).__init__(robot, window_size=window_size, axis=axis, ticks=ticks)
 
+    def __copy__(self):
+        """Return a shallow copy of the state. This can be overridden in the child class."""
+        return self.__class__(robot=self.robot, link_ids=self.links, window_size=self.window_size, axis=self.axis,
+                              ticks=self.ticks)
+
+    def __deepcopy__(self, memo={}):
+        """Return a deep copy of the state. This can be overridden in the child class.
+
+        Args:
+            memo (dict): memo dictionary of objects already copied during the current copying pass
+        """
+        if self in memo:
+            return memo[self]
+
+        robot = memo.get(self.robot, self.robot)  # copy.deepcopy(self.robot, memo)
+        link_ids = copy.deepcopy(self.links)
+        state = self.__class__(robot=robot, link_ids=link_ids, window_size=self.window_size, axis=self.axis,
+                               ticks=self.ticks)
+
+        memo[self] = state
+        return state
+
 
 class LinkPositionState(LinkState):
     r"""Link Position state
diff --git a/pyrobolearn/states/robot_states/robot_states.py b/pyrobolearn/states/robot_states/robot_states.py
index 49d3135..b098177 100644
--- a/pyrobolearn/states/robot_states/robot_states.py
+++ b/pyrobolearn/states/robot_states/robot_states.py
@@ -8,6 +8,7 @@ Dependencies:
 - `pyrobolearn.robots`
 """
 
+# import copy
 from abc import ABCMeta
 import numpy as np
 
@@ -62,6 +63,25 @@ class RobotState(State):
         """Return the robot instance"""
         return self._robot
 
+    def __copy__(self):
+        """Return a shallow copy of the state. This can be overridden in the child class."""
+        return self.__class__(robot=self.robot, window_size=self.window_size, axis=self.axis, ticks=self.ticks)
+
+    def __deepcopy__(self, memo={}):
+        """Return a deep copy of the state. This can be overridden in the child class.
+
+        Args:
+            memo (dict): memo dictionary of objects already copied during the current copying pass
+        """
+        if self in memo:
+            return memo[self]
+
+        robot = memo.get(self.robot, self.robot)  # copy.deepcopy(self.robot, memo)
+        state = self.__class__(robot=robot, window_size=self.window_size, axis=self.axis, ticks=self.ticks)
+
+        memo[self] = state
+        return state
+
 
 class BasePositionState(RobotState):
     r"""Base position state
diff --git a/pyrobolearn/states/robot_states/sensor_states.py b/pyrobolearn/states/robot_states/sensor_states.py
index 288b351..6c69fca 100644
--- a/pyrobolearn/states/robot_states/sensor_states.py
+++ b/pyrobolearn/states/robot_states/sensor_states.py
@@ -4,6 +4,7 @@
 This includes notably the camera, contact, IMU, force/torque sensors and others.
 """
 
+import copy
 from abc import ABCMeta
 import collections
 import numpy as np
@@ -59,6 +60,25 @@ class SensorState(State):  # RobotState # TODO: define refresh_rate & frequency
         self.sensors = sensors
         super(SensorState, self).__init__(window_size=window_size, axis=axis, ticks=ticks)
 
+    def __copy__(self):
+        """Return a shallow copy of the state. This can be overridden in the child class."""
+        return self.__class__(sensors=self.sensors, window_size=self.window_size, axis=self.axis, ticks=self.ticks)
+
+    def __deepcopy__(self, memo={}):
+        """Return a deep copy of the state. This can be overridden in the child class.
+
+        Args:
+            memo (dict): memo dictionary of objects already copied during the current copying pass
+        """
+        if self in memo:
+            return memo[self]
+
+        sensors = copy.deepcopy(self.sensors, memo)
+        state = self.__class__(sensors=sensors, window_size=self.window_size, axis=self.axis, ticks=self.ticks)
+
+        memo[self] = state
+        return state
+
 
 class CameraState(SensorState):
     r"""Camera state
diff --git a/pyrobolearn/storages/er.py b/pyrobolearn/storages/er.py
index 8bbc402..a329291 100644
--- a/pyrobolearn/storages/er.py
+++ b/pyrobolearn/storages/er.py
@@ -203,9 +203,10 @@ class ExperienceReplay(DictStorage):  # ExperienceReplayStorage(DictStorage):
     @property
     def size(self):
         """Return the size of the experience replay storage."""
-        if self.full:
-            return self.capacity
-        return self.position
+        # if self.full:
+        #     return self.capacity
+        # return self.position
+        return self.capacity
 
     ###########
     # Methods #
@@ -373,7 +374,7 @@ class ExperienceReplay(DictStorage):  # ExperienceReplayStorage(DictStorage):
 
         # go through each attribute in the  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] = [val[indices] for val in value]
             else:  # value = tensor
                 batch[key] = value[indices]
diff --git a/pyrobolearn/storages/storage.py b/pyrobolearn/storages/storage.py
index 139da19..658225f 100644
--- a/pyrobolearn/storages/storage.py
+++ b/pyrobolearn/storages/storage.py
@@ -728,6 +728,11 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
         """Return the size (=number of steps * number of processes) of the rollout storage."""
         return self._num_steps * self._num_trajectories
 
+    @property
+    def capacity(self):
+        """Return the capacity of the rollout storage (=number of steps * number of processes)."""
+        return self._num_steps * self._num_trajectories
+
     @property
     def curr_step(self):
         """Return the current time step."""
diff --git a/pyrobolearn/terminal_conditions/terminal_condition.py b/pyrobolearn/terminal_conditions/terminal_condition.py
index fb8a7a8..d9b7070 100644
--- a/pyrobolearn/terminal_conditions/terminal_condition.py
+++ b/pyrobolearn/terminal_conditions/terminal_condition.py
@@ -2,6 +2,7 @@
 """Define some common terminal conditions for the environment.
 """
 
+import copy
 import numpy as np
 
 from pyrobolearn.robots import Robot
@@ -40,7 +41,7 @@ class TerminalCondition(object):
         """
         return False
 
-    def __repr__(self):
+    def __str__(self):
         return self.__class__.__name__
 
     def __call__(self, *args, **kwargs):
@@ -123,12 +124,31 @@ class HasFallen(FailedCondition):
         angle_condition = angle > self.angle_threshold
         return height_condition or angle_condition
 
-    def __repr__(self):
+    def __str__(self):
         description = '{} (\n\tbase_height={} ?<? height_threshold={}, \n\tangle_up_vector={} ?>? angle_threshold={}' \
                       '\n)'.format(self.__class__.__name__, self._compute_height(), self.height_threshold,
                                    self._compute_angle(), self.angle_threshold)
         return description
 
+    def __copy__(self):
+        """Return a shallow copy of the terminal condition. This can be overridden in the child class."""
+        return self.__class__(robot=self.robot, height_threshold=self.height_threshold,
+                              angle_threshold=self.angle_threshold)
+
+    def __deepcopy__(self, memo={}):
+        """Return a deep copy of the terminal condition. This can be overridden in the child class.
+
+        Args:
+            memo (dict): memo dictionary of objects already copied during the current copying pass
+        """
+        if self in memo:
+            return memo[self]
+        robot = memo.get(self.robot, self.robot)
+        terminal = self.__class__(robot=robot, height_threshold=self.height_threshold,
+                                  angle_threshold=self.angle_threshold)
+        memo[self] = terminal
+        return terminal
+
 
 class HasReached(SucceededCondition):
     r"""Has Reached Condition
@@ -175,7 +195,7 @@ class LinkInSpecifiedDirection(HasReached):
         Returns:
             bool: True if the condition is satisfied.
         """
-        pos = self.state._data
+        pos = self.state.data[0]
         pos = self.normalize(pos)
         value = np.dot(pos, self.direction)
         # check if the direction belongs to the cone domain
@@ -188,5 +208,25 @@ class LinkInSpecifiedDirection(HasReached):
             self.cnt = 0
         return False
 
-    def __repr__(self):
-        return self.__class__.__name__ + '(direction=' + str(self.direction) + ')'
\ No newline at end of file
+    def __str__(self):
+        return self.__class__.__name__ + '(direction=' + str(self.direction) + ')'
+
+    def __copy__(self):
+        """Return a shallow copy of the terminal condition. This can be overridden in the child class."""
+        return self.__class__(state=self.state, direction=self.direction, domain=self.domain,
+                              total_steps=self.total_steps)
+
+    def __deepcopy__(self, memo={}):
+        """Return a deep copy of the terminal condition. This can be overridden in the child class.
+
+        Args:
+            memo (dict): memo dictionary of objects already copied during the current copying pass
+        """
+        if self in memo:
+            return memo[self]
+        state = copy.deepcopy(self.state, memo)
+        direction = copy.deepcopy(self.direction)
+        domain = copy.deepcopy(self.domain)
+        terminal = self.__class__(state=state, direction=direction, domain=domain, total_steps=self.total_steps)
+        memo[self] = terminal
+        return terminal
diff --git a/pyrobolearn/values/basic_value.py b/pyrobolearn/values/basic_value.py
index 266c132..7424bfb 100644
--- a/pyrobolearn/values/basic_value.py
+++ b/pyrobolearn/values/basic_value.py
@@ -2,7 +2,7 @@
 """Provides the various basic value function approximators (e.g. table and linear value approximators)
 """
 
-from abc import ABCMeta
+import copy
 import torch
 
 # from pyrobolearn.models import Linear
@@ -52,6 +52,24 @@ class LinearValue(ParametrizedValue):
         model = LinearApproximator(inputs=state, outputs=torch.Tensor([1]), preprocessors=preprocessors)
         super(LinearValue, self).__init__(state, model=model)
 
+    def __copy__(self):
+        """Return a shallow copy of the value approximator. This can be overridden in the child class."""
+        return self.__class__(state=self.state, preprocessors=self.model.preprocessors)
+
+    def __deepcopy__(self, memo={}):
+        """Return a deep copy of the value approximator. This can be overridden in the child class.
+
+        Args:
+            memo (dict): memo dictionary of objects already copied during the current copying pass
+        """
+        if self in memo:
+            return memo[self]
+        state = copy.deepcopy(self.state, memo)
+        preprocessors = [copy.deepcopy(preprocessor, memo) for preprocessor in self.model.preprocessors]
+        value = self.__class__(state=state, preprocessors=preprocessors)
+        memo[self] = value
+        return value
+
 
 class LinearQValue(ParametrizedQValue):
     r"""Linear Q-value function approximator (which accepts as inputs the states and actions)
@@ -74,6 +92,26 @@ class LinearQValue(ParametrizedQValue):
         model = LinearApproximator(inputs=[state, action], outputs=torch.Tensor([1]), preprocessors=preprocessors)
         super(LinearQValue, self).__init__(state, action, model=model)
 
+    def __copy__(self):
+        """Return a shallow copy of the value approximator. This can be overridden in the child class."""
+        return self.__class__(state=self.state, action=self.action, preprocessors=self.model.preprocessors)
+
+    def __deepcopy__(self, memo={}):
+        """Return a deep copy of the value approximator. This can be overridden in the child class.
+
+        Args:
+            memo (dict): memo dictionary of objects already copied during the current copying pass
+        """
+        if self in memo:
+            return memo[self]
+
+        state = copy.deepcopy(self.state, memo)
+        action = copy.deepcopy(self.action, memo)
+        preprocessors = [copy.deepcopy(preprocessor, memo) for preprocessor in self.model.preprocessors]
+        value = self.__class__(state=state, action=action, preprocessors=preprocessors)
+        memo[self] = value
+        return value
+
 
 class LinearQValueOutput(ParametrizedQValueOutput):
     r"""Linear Q-value function approximator (which accepts as inputs the states and outputs a Q-value for each
@@ -96,3 +134,23 @@ class LinearQValueOutput(ParametrizedQValueOutput):
         """
         model = LinearApproximator(inputs=state, outputs=action, preprocessors=preprocessors)
         super(LinearQValueOutput, self).__init__(state, action, model=model)
+
+    def __copy__(self):
+        """Return a shallow copy of the value approximator. This can be overridden in the child class."""
+        return self.__class__(state=self.state, action=self.action, preprocessors=self.model.preprocessors)
+
+    def __deepcopy__(self, memo={}):
+        """Return a deep copy of the value approximator. This can be overridden in the child class.
+
+        Args:
+            memo (dict): memo dictionary of objects already copied during the current copying pass
+        """
+        if self in memo:
+            return memo[self]
+
+        state = copy.deepcopy(self.state, memo)
+        action = copy.deepcopy(self.action, memo)
+        preprocessors = [copy.deepcopy(preprocessor, memo) for preprocessor in self.model.preprocessors]
+        value = self.__class__(state=state, action=action, preprocessors=preprocessors)
+        memo[self] = value
+        return value
diff --git a/pyrobolearn/values/value.py b/pyrobolearn/values/value.py
index 7a59888..0876465 100644
--- a/pyrobolearn/values/value.py
+++ b/pyrobolearn/values/value.py
@@ -90,6 +90,18 @@ class ValueApproximator(object):
         """Predict the value."""
         pass
 
+    def train(self):
+        """
+        Set the value approximator in training mode.
+        """
+        pass
+
+    def eval(self):
+        """
+        Set the value approximator in evaluation mode.
+        """
+        pass
+
     #############
     # Operators #
     #############
@@ -116,6 +128,8 @@ class ValueApproximator(object):
         Args:
             memo (dict): memo dictionary of objects already copied during the current copying pass
         """
+        if self in memo:
+            return memo[self]
         state = copy.deepcopy(self.state, memo) if isinstance(self.state, State) else copy.deepcopy(self.state)
         value = self.__class__(state=state)
         memo[self] = value
@@ -174,6 +188,8 @@ class QValueApproximator(ValueApproximator):
         Args:
             memo (dict): memo dictionary of objects already copied during the current copying pass
         """
+        if self in memo:
+            return memo[self]
         state = copy.deepcopy(self.state, memo) if isinstance(self.state, State) else copy.deepcopy(self.state)
         action = copy.deepcopy(self.action, memo) if isinstance(self.action, Action) else copy.deepcopy(self.action)
         value = self.__class__(state=state, action=action)
@@ -255,6 +271,18 @@ class ParametrizedValue(ValueApproximator):
     # Methods #
     ###########
 
+    def train(self):
+        """
+        Set the value approximator in training mode.
+        """
+        self.model.train()
+
+    def eval(self):
+        """
+        Set the value approximator in evaluation mode.
+        """
+        self.model.eval()
+
     def parameters(self):
         """
         Return an iterator over the learning model parameters.
@@ -306,7 +334,7 @@ class ParametrizedValue(ValueApproximator):
         Set the vectorized parameters.
 
         Args:
-            np.array, torch.Tensor: 1D parameter vector.
+            vector (np.array, torch.Tensor): 1D parameter vector.
         """
         self.model.set_vectorized_parameters(vector=vector)
 
@@ -351,6 +379,8 @@ class ParametrizedValue(ValueApproximator):
         Args:
             memo (dict): memo dictionary of objects already copied during the current copying pass
         """
+        if self in memo:
+            return memo[self]
         state = copy.deepcopy(self.state, memo) if isinstance(self.state, State) else copy.deepcopy(self.state)
         model = copy.deepcopy(self.model, memo)
         value = self.__class__(state=state, model=model)
@@ -551,6 +581,8 @@ class ParametrizedQValue(QValueApproximator):  # ParametrizedValue, QValueApprox
         Args:
             memo (dict): memo dictionary of objects already copied during the current copying pass
         """
+        if self in memo:
+            return memo[self]
         state = copy.deepcopy(self.state, memo) if isinstance(self.state, State) else copy.deepcopy(self.state)
         action = copy.deepcopy(self.action, memo) if isinstance(self.action, Action) else copy.deepcopy(self.action)
         model = copy.deepcopy(self.model, memo)
@@ -646,6 +678,8 @@ class ParametrizedQValueOutput(ParametrizedQValue):
         Args:
             memo (dict): memo dictionary of objects already copied during the current copying pass
         """
+        if self in memo:
+            return memo[self]
         state = copy.deepcopy(self.state, memo) if isinstance(self.state, State) else copy.deepcopy(self.state)
         action = copy.deepcopy(self.action, memo) if isinstance(self.action, Action) else copy.deepcopy(self.action)
         model = copy.deepcopy(self.model, memo)