solving compatbility issues with baselines + update envs, rewards, terminal conditions, states/actions

docs/source/pyrobolearn.rst

@@ -163,8 +163,9 @@ we now provide a brief overview of each submodule and its intended use:
 - ``returns``: this provides the various returns and estimators that are used in RL.
 - ``algos``: this contains the various learning algorithms on how to acquire the data and train the various models
   (policies, values, dynamics, etc).
-- ``metrics``: this contains the various metrics that are used in different learning paradigms. They are not currently 
-  all implemented. You can put different metrics together and plot them by just calling the ``plot`` method.
+- ``metrics``: this contains the various metrics that are used in different learning paradigms (imitation, reinforcement, 
+  transfer, etc). They are not currently all implemented. You can combine different metrics together and plot them by 
+  just calling the ``plot`` method.
 
 Other folders include:
 

docs/source/robots.rst

@@ -236,6 +236,8 @@ To illustrate how to create your own robot, let's assume you want to create a hu
 Sensors and Actuators
 ---------------------
 
+Both sensors and actuators are attached to joints or links, and interact with the simulator interface. They notably both accept a ``noise`` distribution, the number of ``ticks`` (i.e. the number of steps to wait/sleep before the acquisition of the next sensor value), the ``latency`` (currently fixed).
+
 * Sensors
 * Actuators
 

examples/README.rst

@@ -21,4 +21,4 @@ You can check the following folders on:
 - ``rewards``: how to use the reward functions.
 - ``environments``: provide a full example on how to create an environment from scratch in PRL.
 - ``imitation``: how to use imitation learning with the framework.
-- ``gym/cartpole``: policies that are trained with different algorithms on the gym Cartpole environment.
+- ``reinforcement``: how to use reinforcement learning with the framework.

examples/kinematics/inverse/moving_sphere.py

@@ -72,4 +72,5 @@ for t in count():
     robot.set_joint_positions(q, joint_ids=joint_ids)
 
     # step in simulation
+    robot.step()
     world.step(sleep_dt=dt)

examples/reinforcement/README.md

@@ -0,0 +1,9 @@
+## Reinforcement learning task
+
+In this folder, you can run reinforcement learning tasks.
+
+- In the `gym` subfolder, you can run `gym` environments using the models and algorithms available from the PRL 
+  frameworks.
+- In the `baselines` subfolder, you can `PRL` environments using the neural networks models and algorithms provided by 
+  the `stable_baselines` library.
+- Other example files provide PRL environments along with models and algorithms provided by PRL.

examples/reinforcement/baselines/README.rst

@@ -0,0 +1,14 @@
+Baselines
+---------
+
+This folder contains examples when using PRL environments and algorithms defined in the ``stable_baselines`` Python
+library.
+
+Few notes with respect to that:
+
+1. ``stable_baselines`` uses the ``TensorFlow`` backend, and a ``DummyVecEnv`` has to be provided to the algorithms.
+2. Normally, in PRL, the actions can be defined outside the environments and it is the policy that is responsible to
+   apply the action in the world. However, in ``OpenAI gym``, it is the environment that has the ``action_space`` and
+   apply the ``action``. To accommodate with that, the action can also be defined and provided to the PRL environment.
+3. When using PRL with ``stable_baselines``, make sure that each states have the same dimensions; i.e. we can not
+   return a 1D vector state with a 2D matrix state at the same time (at least, not currently).

examples/reinforcement/baselines/acrobot_gym.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+"""Example on how to use the 'Acrobot' OpenAI Gym environments in PRL using the `stable_baselines` library.
+"""
+
+from stable_baselines.common.policies import MlpPolicy
+from stable_baselines.common.vec_env import DummyVecEnv
+from stable_baselines import PPO2
+
+from pyrobolearn.envs import gym  # this is a thin wrapper around the gym library
+
+# create env, state, and action from gym
+env = gym.make('Acrobot-v1')
+state, action = env.state, env.action
+print("State and action space: {} and {}".format(state.space, action.space))
+
+# The algorithms require a vectorized environment to run
+env = DummyVecEnv([lambda: env])
+
+model = PPO2(MlpPolicy, env, verbose=1)
+model.learn(total_timesteps=10000)
+
+obs = env.reset()
+for i in range(1000):
+    action, _states = model.predict(obs)
+    obs, rewards, dones, info = env.step(action)
+    env.render()

examples/reinforcement/baselines/acrobot_prl.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python
+"""Example on how to use the PRL 'Acrobot' environment using the `stable_baselines` library.
+"""
+
+from stable_baselines.common.policies import MlpPolicy
+from stable_baselines.common.vec_env import DummyVecEnv
+from stable_baselines import PPO2
+
+import gym
+
+import pyrobolearn as prl
+from pyrobolearn.envs.control.acrobot import AcrobotEnv
+
+# create env, state, and action from gym
+sim = prl.simulators.Bullet(render=True)
+env = AcrobotEnv(sim)
+print("State and action space: {} and {}".format(env.state.space, env.action.space))
+print("State and action merged space: {} and {}".format(env.state.merged_space, env.action.merged_space))
+
+# The algorithms require a vectorized environment to run
+env = DummyVecEnv([lambda: env])
+
+model = PPO2(MlpPolicy, env, verbose=1)
+model.learn(total_timesteps=10000)
+
+obs = env.reset()
+# env.render()
+for i in range(1000):
+    action, _states = model.predict(obs)
+    obs, rewards, dones, info = env.step(action)
+    # env.render()

examples/reinforcement/baselines/cartpole_gym.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+"""Example on how to use the 'Cartpole' OpenAI Gym environments in PRL using the `stable_baselines` library.
+"""
+
+from stable_baselines.common.policies import MlpPolicy
+from stable_baselines.common.vec_env import DummyVecEnv
+from stable_baselines import PPO2
+
+from pyrobolearn.envs import gym  # this is a thin wrapper around the gym library
+
+# create env, state, and action from gym
+env = gym.make('CartPole-v1')
+state, action = env.state, env.action
+print("State and action space: {} and {}".format(state.space, action.space))
+
+# The algorithms require a vectorized environment to run
+env = DummyVecEnv([lambda: env])
+
+model = PPO2(MlpPolicy, env, verbose=1)
+model.learn(total_timesteps=10000)
+
+obs = env.reset()
+for i in range(1000):
+    action, _states = model.predict(obs)
+    obs, rewards, dones, info = env.step(action)
+    env.render()

examples/reinforcement/baselines/pendulum_gym.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+"""Example on how to use the 'Pendulum' OpenAI Gym environments in PRL using the `stable_baselines` library.
+"""
+
+from stable_baselines.common.policies import MlpPolicy
+from stable_baselines.common.vec_env import DummyVecEnv
+from stable_baselines import PPO2
+
+from pyrobolearn.envs import gym  # this is a thin wrapper around the gym library
+
+# create env, state, and action from gym
+env = gym.make('Pendulum-v0')
+state, action = env.state, env.action
+print("State and action space: {} and {}".format(state.space, action.space))
+
+# The algorithms require a vectorized environment to run
+env = DummyVecEnv([lambda: env])
+
+model = PPO2(MlpPolicy, env, verbose=1)
+model.learn(total_timesteps=10000)
+
+obs = env.reset()
+for i in range(1000):
+    action, _states = model.predict(obs)
+    obs, rewards, dones, info = env.step(action)
+    env.render()

examples/reinforcement/baselines/pendulum_prl.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python
+"""Example on how to use the PRL 'Acrobot' environment using the `stable_baselines` library.
+"""
+
+from stable_baselines.common.policies import MlpPolicy
+from stable_baselines.common.vec_env import DummyVecEnv
+from stable_baselines import PPO2
+
+import gym
+
+import pyrobolearn as prl
+from pyrobolearn.envs.control.pendulum import InvertedPendulumSwingUpEnv
+
+# create env, state, and action from gym
+sim = prl.simulators.Bullet(render=True)
+env = InvertedPendulumSwingUpEnv(sim)
+print("State and action space: {} and {}".format(env.state.space, env.action.space))
+print("State and action merged space: {} and {}".format(env.state.merged_space, env.action.merged_space))
+
+# The algorithms require a vectorized environment to run
+env = DummyVecEnv([lambda: env])
+
+model = PPO2(MlpPolicy, env, verbose=1)
+model.learn(total_timesteps=10000)
+
+obs = env.reset()
+# env.render()
+for i in range(1000):
+    action, _states = model.predict(obs)
+    obs, rewards, dones, info = env.step(action)
+    # env.render()

pyrobolearn/actions/action.py

@@ -161,6 +161,17 @@ class Action(object):
         # one action: change the data
         # if self.has_data():
         else:
+            if self.is_discrete():  # discrete action
+                if isinstance(data, np.ndarray):  # data action is a numpy array
+                    # check if given logits or not
+                    if data.shape[-1] != 1:  # logits
+                        data = np.array([np.argmax(data)])
+                elif isinstance(data, float):
+                    data = int(data)
+                else:
+                    raise TypeError("Expecting the `data` action to be an int, numpy array, instead got: "
+                                    "{}".format(type(data)))
+
             if not isinstance(data, np.ndarray):
                 if isinstance(data, (list, tuple)):
                     data = np.array(data)
@@ -288,14 +299,25 @@ class Action(object):
         """
         return torch.cat([data.reshape(-1) for data in self.merged_torch_data])
 
+    @property
+    def spaces(self):
+        """
+        Get the corresponding spaces as a list of spaces.
+        """
+        if self.has_space():
+            return [self._space]
+        return [action._space for action in self._actions]
+
     @property
     def space(self):
         """
         Get the corresponding space.
         """
         if self.has_space():
-            return [self._space]
-        return [action._space for action in self._actions]
+            # return gym.spaces.Tuple([self._space])
+            return self._space
+        # return [action._space for action in self._actions]
+        return gym.spaces.Tuple([action._space for action in self._actions])
 
     @space.setter
     def space(self, space):
@@ -306,6 +328,40 @@ class Action(object):
                 isinstance(space, (gym.spaces.Box, gym.spaces.Discrete, gym.spaces.MultiDiscrete)):
             self._space = space
 
+    @property
+    def merged_space(self):
+        """
+        Get the corresponding merged space. Note that all the spaces have to be of the same type.
+        """
+        if self.has_space():
+            return self._space
+        spaces = self.spaces
+        result = []
+        dtype, prev_dtype = None, None
+        for space in spaces:
+            if isinstance(space, gym.spaces.Box):
+                dtype = 'box'
+                result.append([space.low, space.high])
+            elif isinstance(space, gym.spaces.Discrete):
+                dtype = 'discrete'
+                result.append(space.n)
+            else:
+                raise NotImplementedError
+
+            if prev_dtype is not None and dtype != prev_dtype:
+                return self.space
+
+            prev_dtype = dtype
+
+        if dtype == 'box':
+            low = np.concatenate([res[0] for res in result])
+            high = np.concatenate([res[1] for res in result])
+            return gym.spaces.Box(low=low, high=high, dtype=np.float32)
+        elif dtype == 'discrete':
+            return gym.spaces.Discrete(n=np.sum(result))
+
+        return self.space
+
     @property
     def name(self):
         """
@@ -456,7 +512,7 @@ class Action(object):
                 if data is None:
                     data = self._data
                 self._write(data)
-            else:  # read each action
+            else:  # write each action
                 if self.actions:
                     if data is None:
                         data = [None] * len(self.actions)

pyrobolearn/actions/robot_actions/joint_actions.py

@@ -548,7 +548,6 @@ class JointTorqueAction(JointAction):
     def _write_continuous(self, data):
         """apply the action data on the robot."""
         data = np.clip(data, self.f_min, self.f_max)
-        print(data)
         self.robot.set_joint_torques(data, self.joints)
 
     def __copy__(self):

pyrobolearn/envs/control/README.rst

@@ -0,0 +1,2 @@
+Control Environments
+--------------------

pyrobolearn/envs/control/acrobot.py

@@ -54,10 +54,15 @@ class AcrobotEnv(ControlEnv):
         Initialize the acrobot environment.
 
         Args:
-            simulator (Simulator): simulator instance.
+            simulator (Simulator): simulator instance. If None, by default, it will instantiate the Bullet
+                simulator.
             use_reward_shaping (bool): if True, it will use a reward that guides how to achieve the goal.
             verbose (bool): if True, it will print information when creating the environment.
         """
+        # simulator
+        if simulator is None:
+            simulator = prl.simulators.Bullet(render=verbose)
+
         # create basic world
         world = prl.worlds.BasicWorld(simulator)
         robot = world.load_robot('acrobot')

pyrobolearn/envs/control/cartpole.py

@@ -64,10 +64,15 @@ class CartpoleEnv(ControlEnv):
         Initialize the Cartpole environment.
 
         Args:
-            simulator (Simulator): simulator instance.
+            simulator (Simulator): simulator instance. If None, by default, it will instantiate the Bullet
+                simulator.
             num_links (int): the number of links that forms the inverted pendulum.
             verbose (bool): if True, it will print information when creating the environment.
         """
+        # simulator
+        if simulator is None:
+            simulator = prl.simulators.Bullet(render=verbose)
+
         # create basic world
         world = prl.worlds.World(simulator)
         robot = prl.robots.CartPole(simulator, position=(0., 0., 0.), num_links=num_links, inverted_pole=False)

pyrobolearn/envs/control/pendulum.py

@@ -51,9 +51,14 @@ class InvertedPendulumSwingUpEnv(ControlEnv):
         Initialize the inverted pendulum swing-up environment.
 
         Args:
-            simulator (Simulator): simulator instance.
+            simulator (Simulator, None): simulator instance. If None, by default, it will instantiate the Bullet
+                simulator.
             verbose (bool): if True, it will print information when creating the environment
         """
+        # simulator
+        if simulator is None:
+            simulator = prl.simulators.Bullet(render=verbose)
+
         # create basic world with the robot
         world = prl.worlds.BasicWorld(simulator)
         robot = world.load_robot('pendulum')
@@ -122,12 +127,15 @@ if __name__ == "__main__":
 
     # create environment
     env = InvertedPendulumSwingUpEnv(sim, verbose=True)
+    action = env.action
+    action.data = 2.
 
     # run simulation
     env.reset()
     for t in prl.count():
-        # if (t % 800) == 0:
-        #     env.reset()  # test reset function
+        if (t % 800) == 0:
+            env.reset()  # test reset function
+        action()
         states, rewards, done, info = env.step(sleep_dt=1./240)
         # print("State: {}".format(states))
-        print("Reward: {}".format(rewards))
+        # print("Reward: {}".format(rewards))

pyrobolearn/envs/env.py

@@ -12,6 +12,7 @@ Dependencies:
 
 import copy
 import pickle
+import numpy as np
 import gym
 
 from pyrobolearn.worlds import World, BasicWorld
@@ -50,9 +51,9 @@ class Env(gym.Env):  # TODO: make it inheriting the gym.Env
     the `gym.Env` class (see `core.py` in `https://github.com/openai/gym/blob/master/gym/core.py`).
 
     References:
-        [1] "Reinforcement Learning: An Introduction", Sutton and Barto, 1998
-        [2] "Wikipedia: Composition over Inheritance", https://en.wikipedia.org/wiki/Composition_over_inheritance
-        [3] "OpenAI gym": https://gym.openai.com/   and    https://github.com/openai/gym
+        - [1] "Reinforcement Learning: An Introduction", Sutton and Barto, 1998
+        - [2] "Wikipedia: Composition over Inheritance", https://en.wikipedia.org/wiki/Composition_over_inheritance
+        - [3] "OpenAI gym": https://gym.openai.com/   and    https://github.com/openai/gym
     """
 
     def __init__(self, world, states, rewards=None, terminal_conditions=None, initial_state_generators=None,
@@ -66,11 +67,11 @@ class Env(gym.Env):  # TODO: make it inheriting the gym.Env
             states ((list of) State): states that are returned by the environment at each time step.
             rewards (None, Reward): The rewards can be None when for instance we are in an imitation learning setting,
                 instead of a reinforcement learning one. If None, only the state is returned by the environment.
-            terminal_conditions (None, callable, TerminalCondition, list of TerminalCondition): A callable function or
+            terminal_conditions (None, callable, TerminalCondition, list[TerminalCondition]): A callable function or
                 object that check if the policy has failed or succeeded the task.
-            initial_state_generators (None, StateGenerator, list of StateGenerator): state generators which are used
+            initial_state_generators (None, StateGenerator, list[StateGenerator]): state generators which are used
                 when resetting the environment to generate the initial states.
-            physics_randomizers (None, PhysicsRandomizer, list of PhysicsRandomizer): physics randomizers. This will be
+            physics_randomizers (None, PhysicsRandomizer, list[PhysicsRandomizer]): physics randomizers. This will be
                 called each time you reset the environment.
             extra_info (None, callable): Extra info returned by the environment at each time step.
             actions ((list of) Action): actions that are given to the environment. Note that this is not used here in
@@ -85,7 +86,7 @@ class Env(gym.Env):  # TODO: make it inheriting the gym.Env
         self.terminal_conditions = terminal_conditions
         self.physics_randomizers = physics_randomizers
         self.state_generators = initial_state_generators
-        self.extra_info = extra_info if extra_info is not None else lambda: False
+        self.extra_info = extra_info if extra_info is not None else lambda: dict()
         self.actions = actions
 
         # state dictionary which contains at least {'policy': State, 'value': State}
@@ -185,6 +186,13 @@ class Env(gym.Env):  # TODO: make it inheriting the gym.Env
         """Return the state space of the first (combined) state."""
         return self.states[0].merged_space
 
+    # alias
+    observations = states
+    observation = state
+    observation_dict = state_dict
+    observation_spaces = state_spaces
+    observation_space = state_space
+
     @property
     def actions(self):
         """Return the actions."""
@@ -351,8 +359,11 @@ class Env(gym.Env):  # TODO: make it inheriting the gym.Env
         for generator in self.state_generators:
             generator(reset_state=False)
 
+        self.world.step()
+
         # reset states and return first states/observations
-        states = [state.reset() for state in self.states]
+        states = [state.reset(merged_data=True) for state in self.states]
+        print("Reset: ", states)
         return self._convert_state_to_data(states)
 
     def step(self, actions=None, sleep_dt=None):
@@ -363,12 +374,14 @@ class Env(gym.Env):  # TODO: make it inheriting the gym.Env
 
         Args:
             actions (None, (list of) Action, (list of) np.array): an action provided by the policy(ies) to the
-                environment. Note that this is not used in this method; calling the actions should be done inside the
-                policy(ies), and not in the environment. The policy decides when to execute an action. Several problems
-                can appear by providing the actions in the environment instead of letting the policy executes them.
-                For instance, think about when there are multiple policies, when using multiprocessing, or when the
-                environment runs in real-time.
-            sleep_dt (float):
+                environment. Note that this is not normally used in this method; calling the actions should be done
+                inside the policy(ies), and not in the environment. The policy decides when to execute an action.
+                Several problems can appear by providing the actions in the environment instead of letting the policy
+                executes them. For instance, think about when there are multiple policies, when using multiprocessing,
+                or when the environment runs in real-time. However, if an action is given as a (list of) np.array,
+                it will be set as the action data, and the action will be executed. If the action is a (list of) Action,
+                it will call each action.
+            sleep_dt (float): time to sleep.
 
         Returns:
             observation (object): agent's observation of the current environment
@@ -391,6 +404,32 @@ class Env(gym.Env):  # TODO: make it inheriting the gym.Env
         # if actions is not None and isinstance(actions, Action):
         #     actions()
 
+        # if the actions are provided, set and apply them in the environment
+        if actions is not None:
+            if isinstance(actions, Action):
+                actions()
+            elif isinstance(actions, np.ndarray) and isinstance(self.actions, list):  # set the data
+                if len(self.actions) == 1:
+                    self.actions[0].data = actions
+                else:
+                    raise ValueError("There are multiple actions defined in the environment, so it is unclear to "
+                                     "which action the data should be set to.")
+            elif isinstance(actions, (list, tuple)):
+                for idx, action in enumerate(actions):
+                    if isinstance(action, Action):
+                        action()
+                    elif isinstance(action, np.ndarray) and self.actions is not None:
+                        if len(actions) != len(self.actions):
+                            raise ValueError("The number of given actions (={}) is different from the number of "
+                                             "actions defined in the environments (={})".format(len(actions),
+                                                                                                len(self.actions)))
+                        self.actions[idx].data = action
+                    else:
+                        raise TypeError("Expecting a list of np.array or `Action` instead got: {}".format(type(action)))
+            else:
+                raise TypeError("Expecting an instance of `Action`, np.array, or a list of the previous ones, but got "
+                                "instead: {}".format(type(actions)))
+
         # perform a step forward in the simulation which computes all the dynamics
         self.world.step(sleep_dt=sleep_dt)
 
@@ -403,7 +442,7 @@ class Env(gym.Env):  # TODO: make it inheriting the gym.Env
 
         # get next state/obs for each policy
         # TODO: this should be before computing the rewards as some rewards need the next state
-        states = [state() for state in self.states]
+        states = [state(merged_data=True) for state in self.states]
         states = self._convert_state_to_data(states, convert=True)
 
         # get extra information
@@ -418,7 +457,7 @@ class Env(gym.Env):  # TODO: make it inheriting the gym.Env
         self.sim.render()
 
     def hide(self):
-        """hide the GUI."""
+        """Hide the GUI."""
         self.is_rendering = False
         self.sim.hide()
 

pyrobolearn/envs/locomotion/README.rst

@@ -0,0 +1,3 @@
+Locomotion Environments
+-----------------------
+

pyrobolearn/envs/locomotion/quadruped.py

@@ -9,8 +9,9 @@ References:
     - [2] RaisimGym: https://github.com/leggedrobotics/raisimGym/blob/master/raisim_gym/env/env/ANYmal/Environment.hpp
 """
 
-import pyrobolearn as prl
+import numpy as np
 
+import pyrobolearn as prl
 from pyrobolearn.envs.locomotion.locomotion import LocomotionEnv
 
 
@@ -24,8 +25,8 @@ __email__ = "briandelhaisse@gmail.com"
 __status__ = "Development"
 
 
-class LocomotionQuadrupedEnv1(LocomotionEnv):
-    r"""Locomotion Quadruped Environment
+class LocomotionQuadrupedBulletEnv(LocomotionEnv):
+    r"""Locomotion Quadruped Bullet Environment
 
     This is based on the locomotion environment provided for the minitaur robot in PyBullet [1] but generalized to
     other quadruped robotic platforms.
@@ -34,52 +35,121 @@ class LocomotionQuadrupedEnv1(LocomotionEnv):
 
     - world: basic world with gravity enabled, a basic floor and the quadruped robot.
     - state:
-    - action:
-    - reward:
+        - joint positions (N)
+        - joint velocities (N)
+        - joint torques (N)
+        - base orientation as quaternion (4)
+    - action: PD joint position targets (or joint torques)
+    - reward: reward = 1.0 * r_f + 0. * c_d + 0. * c_s + 0.005 c_e
+        - forward reward: :math:`r_f = x_t - x_{t-1}` where :math:`x` is the base x-position.
+        - drift cost: :math:`c_d = - |y_t - y_{t-1}|` where :math:`y` is the base y-position.
+        - shake cost: :math:`c_s = - |z_t - z_{t-1}|` where :math:`z` is the base z-position.
+        - energy cost: :math:`c_e = -|\tau * dq| * dt` where :math:`\tau` are the torques, :math:`dq` are the joint
+          velocities, and :math:`dt` is the simulation time step.
     - initial state generator:
+        - reset base position and orientation to initial position / orientation
+        - reset base velocity: [0,0,0,0,0,0]
+        - reset joint positions to initial joint positions
+        - reset joint velocities to 0
     - physics randomizer:
+        - additive base mass noise: U([-0.2, 0.2]) kg
+        - additive leg mass noise: U([-0.2, 0.2]) kg
+        - the coefficient of friction for the feet is sampled from :math:`U([0.8, 1.5])`.
     - terminal condition:
+        - fallen:
+            - orientation: :math:`a_z \cdot [0,0,1] < a` where :math:`a_z` is the z-axis of the base, and :math:`a` is
+              the angle threshold (0.85).
+            - height: :math:`z < h` where :math:`h` is the height threshold.
+        - distance limit: :math:`\sqrt{x^2 + y^2} > \text{threshold}`  where :math:`threshold` is set to inf.
+
+    More information:
+    - inner control_loop = 5.
 
     References:
         - [1] PyBullet:
           https://github.com/bulletphysics/bullet3/blob/master/examples/pybullet/gym/pybullet_envs/bullet/minitaur_gym_env.py
     """
 
-    def __init__(self, simulator=None, robot='minitaur'):
+    def __init__(self, simulator=None, robot='minitaur', verbose=False):
         """
         Initialize the locomotion with quadruped environment.
 
         Args:
-            simulator (Simulator, None): simulator instance.
+            simulator (Simulator, None): simulator instance. If None, by default, it will instantiate the Bullet
+                simulator.
             robot (str): robot name.
+            verbose (bool): if True, it will print information when creating the environment.
         """
+        # create simulator if necessary
+        if simulator is None:
+            simulator = prl.simulators.Bullet(render=verbose)
+
         # create basic world
         world = prl.worlds.BasicWorld(simulator)
-        robot = world.load_robot(robot)
+
+        # load robot in world
+        self.robot = world.load_robot(robot)
+        if not isinstance(self.robot, prl.robots.LeggedRobot):  # prl.robots.QuadrupedRobot
+            raise TypeError("Expecting a legged robot, but got instead {}".format(type(self.robot)))
+        if verbose:
+            self.robot.print_info()
 
         # create state
-        state = None
+        q_state = prl.states.JointPositionState(robot=self.robot)
+        dq_state = prl.states.JointVelocityState(robot=self.robot)
+        tau_state = prl.states.JointForceTorqueState(robot=self.robot)
+        quat_state = prl.states.BaseOrientationState(robot=self.robot)
+        state = q_state + dq_state + tau_state + quat_state
+        if verbose:
+            print(state)
 
         # create action
-        action = None
-
-        # create reward
-        reward = None
+        action = prl.actions.JointPositionAction(robot=self.robot, kp=self.robot.kp, kd=self.robot.kd)
+        if verbose:
+            print(action)
 
         # create terminal condition
-        terminal_condition = None
+        orientation_condition = prl.terminal_conditions.BaseOrientationAxisCondition(self.robot, angle=0.85,
+                                                                                     axis=(0., 0., 1.), dim=2,
+                                                                                     stay=True, out=False)
+        height_condition = prl.terminal_conditions.BaseHeightCondition(self.robot, height=self.robot.base_height/8.,
+                                                                       stay=True, out=True)
+        distance_condition = prl.terminal_conditions.DistanceCondition(self.robot, distance=float("inf"),
+                                                                       dim=[1, 1, 0], stay=True, out=False)
+        terminal_condition = [orientation_condition, height_condition, distance_condition]
+        if verbose:
+            print("Terminal condition: {}".format(terminal_condition))
+
+        # create reward
+        forward_reward = prl.rewards.ForwardProgressReward(self.robot, direction=(1., 0., 0.))
+        base_position_state = prl.states.BasePositionState(self.robot)
+        drift_cost = prl.rewards.DriftCost(base_position_state, update_state=True)  # y component
+        shake_cost = prl.rewards.ShakeCost(base_position_state)  # z component
+        energy_cost = prl.rewards.JointEnergyCost(self.robot, dt=simulator.dt)
+        reward = 1. * forward_reward + 0.005 * energy_cost + 0. * shake_cost + 0. * drift_cost
+        if verbose:
+            print(reward)
 
         # create initial state generator
-        initial_state_generator = None
+        base_pose_gen = prl.states.generators.FixedStateGenerator(state=prl.states.BasePoseState(self.robot))
+        base_vel_gen = prl.states.generators.FixedStateGenerator(state=prl.states.BaseLinearVelocityState(self.robot))
+        q_init = self.robot.get_joint_configurations('home') if self.robot.has_joint_configuration('home') else \
+            np.zeros(len(self.robot.joints))
+        q_gen = prl.states.generators.FixedStateGenerator(state=q_state, data=q_init)
+        dq_gen = prl.states.generators.FixedStateGenerator(state=dq_state, data=np.zeros(len(self.robot.joints)))
+
+        initial_state_generator = [base_pose_gen, base_vel_gen, q_gen, dq_gen]
+        if verbose:
+            print("Initial state generator: {}".format(initial_state_generator))
 
         # create environment using composition
-        super(LocomotionQuadrupedEnv1, self).__init__(world=world, states=state, rewards=reward, actions=action,
-                                                      terminal_conditions=terminal_condition,
-                                                      initial_state_generators=initial_state_generator)
+        super(LocomotionQuadrupedBulletEnv, self).__init__(world=world, states=state, rewards=reward, actions=action,
+                                                           terminal_conditions=terminal_condition,
+                                                           initial_state_generators=initial_state_generator)
 
 
-class LocomotionQuadrupedEnv2(LocomotionEnv):
-    r"""Locomotion Quadruped Environment
+class LocomotionQuadrupedRaisimEnv(LocomotionEnv):
+    r"""Locomotion Quadruped Raisim Environment
 
     This is based on the locomotion environment provided in `raisimGym` for the ANYmal robot in [1]. The Python version
     can be found in `raisimpy` in [2].
@@ -101,19 +171,23 @@ class LocomotionQuadrupedEnv2(LocomotionEnv):
     - initial state generator: fixed state generator for joint positions such that they are set to the home position.
     - terminal condition: if there is contact with a link that is not the foot.
 
+    More information:
+    - inner control_loop = int(control_dt / simulation_dt) where control_dt=0.01 and simulation_dt=0.001.
+
     References:
         - [1] RaisimGym:
           https://github.com/leggedrobotics/raisimGym/blob/master/raisim_gym/env/env/ANYmal/Environment.hpp
         - [2] Raisimpy: https://github.com/robotlearn/raisimpy/blob/master/examples/raisimpy_gym/envs/anymal/env.py
     """
 
-    def __init__(self, simulator=None, robot='anymal'):
+    def __init__(self, simulator=None, robot='anymal', verbose=False):
         """
         Initialize the locomotion quadruped environment.
 
         Args:
             simulator (Simulator, None): simulator instance. If simulator is None, it will use the PyBullet simulator.
             robot (str): robot name.
+            verbose (bool): if True, it will print information when creating the environment.
         """
         # check simulator
         if simulator is None:
@@ -127,25 +201,51 @@ class LocomotionQuadrupedEnv2(LocomotionEnv):
 
         # load robot in world
         self.robot = world.load_robot(robot)
+        if not isinstance(self.robot, prl.robots.LeggedRobot):  # prl.robots.QuadrupedRobot
+            raise TypeError("Expecting a legged robot, but got instead {}".format(type(self.robot)))
+        if verbose:
+            self.robot.print_info()
 
         # create state
-        state = None
+        height_state = prl.states.BaseHeightState(robot=self.robot)
+        z_axis_state = prl.states.BaseAxisState(robot=self.robot, base_axis=2)  # z-axis
+        q_state = prl.states.JointPositionState(robot=self.robot)
+        dq_state = prl.states.JointVelocityState(robot=self.robot)
+        lin_vel_state = prl.states.BaseLinearVelocityState(robot=self.robot)
+        ang_vel_state = prl.states.BaseAngularVelocityState(robot=self.robot)
+        state = height_state + z_axis_state + q_state + dq_state + lin_vel_state + ang_vel_state
+        if verbose:
+            print(state)
 
         # create action
-        action = None
+        action = prl.actions.JointPositionAction(robot=self.robot, kp=self.robot.kp, kd=self.robot.kd)
+        if verbose:
+            print(action)
+
+        # create terminal condition (all links that are not feet must stay out of contact)
+        terminal_condition = prl.terminal_conditions.ContactCondition(robot=self.robot, link_ids=self.robot.feet,
+                                                                      all=True, stay=True, out=True, complement=True)
+        if verbose:
+            print("Terminal condition: {}".format(terminal_condition))
 
         # create reward
-        reward = None
-
-        # create terminal condition
-        terminal_condition = None
+        vel_reward = prl.rewards.BaseLinearVelocityReward(state=lin_vel_state, axis=0)
+        torque_cost = prl.rewards.JointTorqueCost(state=self.robot)
+        terminal_reward = prl.rewards.TerminalReward(terminal_conditions=terminal_condition, final_reward=-10.)
+        reward = 0.3 * vel_reward + 2e-5 * torque_cost + terminal_reward
+        if verbose:
+            print(reward)
 
         # create initial state generator
-        initial_state_generator = None
+        q_init = self.robot.get_joint_configurations('home') if self.robot.has_joint_configuration('home') else \
+            np.zeros(len(self.robot.joints))
+        initial_state_generator = prl.states.generators.FixedStateGenerator(state=q_state, data=q_init)
+        if verbose:
+            print("Initial state generator: {}".format(initial_state_generator))
 
-        super(LocomotionQuadrupedEnv2, self).__init__(world=world, states=state, rewards=reward, actions=action,
-                                                      terminal_conditions=terminal_condition,
-                                                      initial_state_generators=initial_state_generator)
+        super(LocomotionQuadrupedRaisimEnv, self).__init__(world=world, states=state, rewards=reward, actions=action,
+                                                           terminal_conditions=terminal_condition,
+                                                           initial_state_generators=initial_state_generator)
 
 
 # Test
@@ -156,8 +256,16 @@ if __name__ == "__main__":
     sim = prl.simulators.Bullet()
 
     # create environment
-    env = LocomotionQuadrupedEnv1(sim)
+    env = LocomotionQuadrupedRaisimEnv(sim, verbose=True)
+    # env = LocomotionQuadrupedBulletEnv(sim, verbose=True)
 
     # run simulation
     for _ in count():
-        env.step(sleep_dt=1./240)
+        obs, reward, done, info = env.step(sleep_dt=1./240)
+        # print("obs: {}".format(obs))
+        print("reward: {}".format(reward))
+        print("done: {}".format(done))
+        print("info: {}".format(info))
+        if done:
+            print("End")
+            break

pyrobolearn/envs/locomotion/robust_quadruped.py

@@ -56,7 +56,7 @@ class SelfRightingEnv(LocomotionEnv):
         - joint position: :math:`c_{jp} = \sum_{i}^N K(d(q_i, \hat{q}_i), 2.0)` where :math:`\hat{q}_i` is the desired
           target joint position which correspond in this case to the crouching pose, :math:`d(\cdot, \cdot)` is the
           minimum angle difference which maps to :math:`[0,\pi]`, and
-          :math:`K(e, \alpha) = \frac{-1}{e^{\alpha e} + 2 + e^{\alpha e}}` is a kernel function that maps
+          :math:`K(x, \alpha) = \frac{-1}{e^{\alpha x} + 2 + e^{-\alpha x}}` is a kernel function that maps
           :math:`\mathcal{R}` to :math:`[-0.25, 0[`.
         - body impulse: :math:`c_{bi} = \sum_{n \in I_c \backslash I_{c,f}} || i_{c,n} || / (|I_c| - |I_{c,f}|)` where
           :math:`I_c` is the index set of the contact points, :math:`I_{c,f}` is the index set of the foot contact
@@ -98,13 +98,14 @@ class SelfRightingEnv(LocomotionEnv):
         - [2] "Learning agile and dynamic motor skills for legged robots", Hwangbo et al., 2019
     """
 
-    def __init__(self, simulator=None, robot='anymal'):
+    def __init__(self, simulator=None, robot='anymal', verbose=False):
         """
         Initialize the self-righting environment.
 
         Args:
             simulator (Simulator, None): simulator instance. If simulator is None, it will use the PyBullet simulator.
             robot (str): robot name.
+            verbose (bool): if True, it will print information when creating the environment.
         """
         # check simulator
         if simulator is None:
@@ -117,30 +118,58 @@ class SelfRightingEnv(LocomotionEnv):
         world = prl.worlds.BasicWorld(simulator)
 
         # load robot in world
-        self.robot = world.load_robot(robot)
+        robot = world.load_robot(robot)
+        self.robot = robot
 
-        # check if the robot has the crouching pose as joint configuration.
+        # check if the robot is a legged robot
+        if not isinstance(self.robot, prl.robots.LeggedRobot):
+            raise TypeError("Expecting the robot to be a legged robot, but instead got: {}".format(type(self.robot)))
+
+        # check if the robot has the crouching pose as joint configuration
         if not self.robot.has_joint_configuration('crouching'):
             raise TypeError("Expecting the robot to have the 'crouching' joint configuration predefined.")
 
+        # create action
+        action = prl.actions.JointPositionAction(robot, kp=robot.kp, kd=robot.kd)
+
         # create state
+        q_state = prl.states.JointPositionState(robot)
+        dq_state = prl.states.JointVelocityState(robot)
+        action_state = prl.states.PreviousActionState(action)
         state = None
 
-        # create action
-        action = None
-
-        # create reward
-        reward = None
+        # create cost
+        c_tau = prl.rewards.JointTorqueCost(state=robot)
+        c_jslim = prl.rewards.JointSpeedLimitCost(state=robot)
+        c_ad = prl.rewards.ActionDifferentCost(action=action)
+        c_o = prl.rewards.OrientationGravityCost(state=robot)
+        c_jp = prl.rewards.JointAngleDifferenceCost(state=, )
+        c_bi = prl.rewards.BodyImpulseCost(robot)
+        c_bs = prl.rewards.BodySlippageCost(robot)
+        c_cin = prl.rewards.SelfCollisionCost(robot)
+        cost = 0.0005 * c_tau + 0.2 * c_jslim + 0.0025 * c_ad + 6 * c_o + 6 * c_jp + 6 * c_bi + 6 * c_bs + 6 * c_cin
 
         # create terminal condition
-        terminal_condition = None  # prl.terminal_conditions.TimeLimitCondition(time=6)
+        terminal_condition = prl.terminal_conditions.TimeLimitCondition(num_steps=6 * 1./simulator.dt)
 
         # create initial state generator
-        initial_state_generator = None
+        joint_position_generator = prl.states.generators.NormalStateGenerator(state=q_state, )
+        drop_generator = prl.states.generators.DropStateGenerator(robot, height=5, condition='fixed')
+        initial_state_generator = [joint_position_generator, drop_generator]
+
+        # create physics randomizer
+        masses = robot.get_link_masses(link_ids=robot.joints)
+        masses = (masses - masses / 10., masses + masses / 10.)
+        mass_randomizer = prl.physics.LinkPhysicsRandomizer(robot, link_ids=robot.joints, masses=masses)
+        com = (-0.03, 0.03)
+        com_randomizer = prl.physics.LinkPhysicsRandomizer(robot, link_ids=robot.joints, local_inertia_positions=com)
+        friction_randomizer = prl.physics.LinkPhysicsRandomizer(robot, link_ids=robot.feet, lateral_frictions=(0.8, 2.))
+        physics_randomizer = [mass_randomizer, com_randomizer, friction_randomizer]
 
         # create environment using composition
-        super(SelfRightingEnv, self).__init__(world=world, states=state, rewards=reward, actions=action,
+        super(SelfRightingEnv, self).__init__(world=world, states=state, rewards=cost, actions=action,
                                               terminal_conditions=terminal_condition,
+                                              physics_randomizers=physics_randomizer,
                                               initial_state_generators=initial_state_generator)
 
 
@@ -179,7 +208,7 @@ class StandingUpEnv(LocomotionEnv):
         - joint position: :math:`c_{jp} = \sum_{i}^N K(d(q_i, \hat{q}_i), 2.0)` where :math:`\hat{q}_i` is the desired
           target joint position which correspond in this case to the crouching pose, :math:`d(\cdot, \cdot)` is the
           minimum angle difference which maps to :math:`[0,\pi]`, and
-          :math:`K(e, \alpha) = \frac{-1}{e^{\alpha e} + 2 + e^{\alpha e}}` is a kernel function that maps
+          :math:`K(x, \alpha) = \frac{-1}{e^{\alpha x} + 2 + e^{-\alpha x}}` is a kernel function that maps
           :math:`\mathcal{R}` to :math:`[-0.25, 0[`.
     - initial state generator: drop the quadruped from 0.5m about the ground with near-upright pose.
     - physics randomizer:
@@ -214,13 +243,14 @@ class StandingUpEnv(LocomotionEnv):
         - [2] "Learning agile and dynamic motor skills for legged robots", Hwangbo et al., 2019
     """
 
-    def __init__(self, simulator=None, robot='anymal'):
+    def __init__(self, simulator=None, robot='anymal', verbose=False):
         """
         Initialize the standing-up environment.
 
         Args:
             simulator (Simulator, None): simulator instance. If simulator is None, it will use the PyBullet simulator.
             robot (str): robot name.
+            verbose (bool): if True, it will print information when creating the environment.
         """
         # check simulator
         if simulator is None:
@@ -297,7 +327,7 @@ class CommandedLocomotionEnv(LocomotionEnv):
           expressed in the base frame.
         - angular velocity: :math:`c_\omega = K(|\omega^B_B - \hat{\omega}^B_B|, 1.0)`, where :math:`\omega^B_B` is the
           angular velocity of the base expressed in the body frame, :math:`\hat{\omega}` is the desired angular
-          velocity, and :math:`K(e, \alpha) = \frac{-1}{e^{\alpha e} + 2 + e^{\alpha e}}` is a kernel function that
+          velocity, and :math:`K(x, \alpha) = \frac{-1}{e^{\alpha x} + 2 + e^{-\alpha x}}` is a kernel function that
           maps :math:`\mathcal{R}` to :math:`[-0.25, 0[`.
         - linear velocity: :math:`c_v = K(|v^B_B - \hat{v}^B_B|, 4.0)`, where :math:`v^B_B` is the linear velocity of
           the base expressed in the body frame and math:`\hat{v}` is the desired linear velocity.
@@ -344,13 +374,14 @@ class CommandedLocomotionEnv(LocomotionEnv):
         - [2] "Learning agile and dynamic motor skills for legged robots", Hwangbo et al., 2019
     """
 
-    def __init__(self, simulator=None, robot='anymal'):
+    def __init__(self, simulator=None, robot='anymal', verbose=False):
         """
         Initialize the standing-up environment.
 
         Args:
             simulator (Simulator, None): simulator instance. If simulator is None, it will use the PyBullet simulator.
             robot (str): robot name.
+            verbose (bool): if True, it will print information when creating the environment.
         """
         # check simulator
         if simulator is None:
@@ -433,7 +464,7 @@ class BehaviorLocomotionEnv(LocomotionEnv):
           expressed in the base frame.
         - angular velocity: :math:`c_\omega = K(|\omega^B_B - \hat{\omega}^B_B|, 1.0)`, where :math:`\omega^B_B` is the
           angular velocity of the base expressed in the body frame, :math:`\hat{\omega}` is the desired angular
-          velocity, and :math:`K(e, \alpha) = \frac{-1}{e^{\alpha e} + 2 + e^{\alpha e}}` is a kernel function that
+          velocity, and :math:`K(x, \alpha) = \frac{-1}{e^{\alpha x} + 2 + e^{-\alpha x}}` is a kernel function that
           maps :math:`\mathcal{R}` to :math:`[-0.25, 0[`.
         - linear velocity: :math:`c_v = K(|v^B_B - \hat{v}^B_B|, 4.0)`, where :math:`v^B_B` is the linear velocity of
           the base expressed in the body frame and math:`\hat{v}` is the desired linear velocity.
@@ -463,7 +494,7 @@ class BehaviorLocomotionEnv(LocomotionEnv):
         - GAE: discount factor (gamma) = 0.99, lambda = 0.99
         - for value function: Adam optimizer with learning rate = 0.001
 
-    Note that the authors report that they could train the behavior policy in ~5hours on a single desktop
+    Note that the authors report that they could train the behavior policy in ~30min on a single desktop
     machine (32 GB memory, Intel i7-8700K and Geforce GTX 1070) with a fully C++ code.
 
 
@@ -517,11 +548,99 @@ class BehaviorLocomotionEnv(LocomotionEnv):
 class AgileLocomotionEnv(LocomotionEnv):
     r"""Agile locomotion environment.
 
-    This is based on the locomotion environment provided in [1] with the ANYmal robotic platform, where they introduce
-    the actuator net.
+    This is based on the locomotion environment provided in [1, 2] with the ANYmal robotic platform, where they
+    introduce the actuator net.
+
+    - simulator: Raisim
+    - world: basic world with gravity enabled, a basic floor and the quadruped robot.
+    - state:
+        - gravity unit vector (:math:`e_g`) expressed in the base frame (3)
+        - estimated base height (h_e) (1)
+        - base angular velocity in body frame (3)
+        - base linear velocity in body frame (3)
+        - joint position and velocity states (2N)
+        - history of joint position error and velocity: current joint state at t (position error + velocity) and two
+          past states corresponding to t-0.01s and t-0.02s (6N)
+        - previous joint position targets a_{t-1} (N)
+        - desired velocity commands (forward velocity, lateral velocity, yaw rate) (3)
+        - additive noise for observation:
+            - joint velocities U(-0.5, 0.5) rad/s
+            - linear velocity of the base U(-0.08, 0.08) m/s
+            - angular velocity of the base U(-0.16, 0.16) m/s
+    - action: PD joint position targets :math:`q_d = 0.5 o_t + q_n` where :math:`o_t` is the output of the policy and
+      :math:`q_n` is the standing joint configuration.
+    - cost: :math:`0.0005 c_{\tau} + 0.03 c_{jslim} + 0.5c_{ad} + 0.4c_o + 6c_\omega + 10 c_v + 0.1 c_{fc} + 2 c_{fs}`,
+      where:  TODO: the costs are similar but a bit different from the ones reported here
+        - torque: :math:`c_{\tau} = || \tau ||^2` where :math:`\tau` are the joint torques.
+        - joint speed limit: :math:`c_{jslim} = \sum_{i}^{N} \max(\dot{q}_{i,lim} - |q_i|, 0)^2` where :math:`N` is
+          the number of actuated joints, :math:`q_i` is the position of the i-th joint, and :math:`\dot{q}_{i,lim}` is
+          the maximum speed of the i-th joint.
+        - action difference: :math:`c_{ad} = || a_t - a_{t-1} ||^2` where :math:`a_t` is the action vector.
+        - orientation cost: :math:`c_o = || [0,0,-1]^\top - e_g ||` where :math:`e_g` is the unit gravity vector
+          expressed in the base frame.
+        - angular velocity: :math:`c_\omega = K(|\omega^B_B - \hat{\omega}^B_B|, 1.0)`, where :math:`\omega^B_B` is the
+          angular velocity of the base expressed in the body frame, :math:`\hat{\omega}` is the desired angular
+          velocity, and :math:`K(x, \alpha) = \frac{-1}{e^{\alpha x} + 2 + e^{-\alpha x}}` is a kernel function that
+          maps :math:`\mathcal{R}` to :math:`[-0.25, 0[`.
+        - linear velocity: :math:`c_v = K(|v^B_B - \hat{v}^B_B|, 4.0)`, where :math:`v^B_B` is the linear velocity of
+          the base expressed in the body frame and math:`\hat{v}` is the desired linear velocity.
+        - foot clearance: :math:`c_{fc} = \sum (h_{f,i} - 0.07)^2 ||v_{f,i}||,  \forall i s.t. g_i > 0, i \in I_{c,f}`,
+          where :math:`h_{f,i}` is the ze position of the `i`th foot, :math:`v_{f,i}` is the velocity of the `i`th foot,
+          :math:`g_i` is the gap function of the `i`th contact, and :math:`I_{c,f}` is the index set of the foot
+          contact points.
+        - foot slippage: :math:`c_{fs} = \sum ||v_{f,i}||,  \forall i s.t. g_i=0, i \in I_{c,f}`
+    - initial state generator (for ANYmal):
+        - base position: mean = [0,0,0.55], std = 1.5cm
+        - base orientation: mean = [1,0,0,0], std = 0.06 rad about a random axis
+        - joint positions: mean = standing configuration = [0, 0.4, -0.8, 0, 0.4, -0.8, 0, -0.4, 0.8, 0, -0.4, 0.8],
+          std = 0.25 rad
+        - base linear velocity: mean = [0]*3, std = 0.012 m/s
+        - base angular velocity: mean = [0]*3, std = 0.4 rad/s
+        - joint velocities: mean = [0]*12, std = 2 rad/s
+        - sample the desired forward velocity, lateral velocity and yaw rate from U(-1, 1) m/s, U(-0.4, 0.4) m/s, and
+          U(-1.2, 1.2) rad/s respectively for the command-conditioned locomotion motion, or U(-1.6, 1.6) m/s,
+          U(-0.2, 0.2) m/s, and U(-0.3, 0.3) rad/s respectively for the high-speed locomotion motion. Note that this
+          depends on the joystick/game controller that is being used.
+    - physics randomizer:
+        - additive noise for center of mass positions: U(-2, 2) cm
+        - additive noise for the link masses: U(-15, 15)%
+        - additive noise for joint positions: U(-2, 2) cm
+    - terminal condition:
+        - time limit of 12sec
+
+
+    Here are more information about the policy, value function, and algorithm used (with exploration strategy) in the
+    paper [1, 2]:
+
+    - actuator network: 6N input units (=joint position error history and joint velocity history), 3 * [32 (softsign)
+      units], N output units (torques)
+    - policy network: input, 256 (tanh) units, 128 (tanh) units, N output units
+    - value network: input, 256 (tanh) units, 128 (tanh) units, 1 output unit
+    - exploration in the continuous action space.
+    - RL algorithm: TRPO (but also tested PPO)
+        - KL divergence threshold (delta) = 0.01
+        - GAE: discount factor (gamma) = 0.9988, lambda = 0.99
+        - for value function: Adam optimizer with learning rate = 0.001
+    - curriculum learning: constraining cost terms (power, torque, joint speed, action difference and orientation
+      costs) are scaled to 10% of the final value at the first iteration and are scaled up as the training proceeds.
+        - scaling factor: :math:`k_{c,j+1} = (k_{c,j})^{k_d}`, with :math:`k_{c,0} = 0.3` and :math:`k_d = 0.997`.
+    - For ANYmal robot: kp = 50 N·m/rad and kd = 0.1 N·m / (rad·s)
+        - kp = nominal range of torque (30 N·m) / nominal range of motion (0.6 rad)
+
+
+    Notes:
+    - the authors report that they could train the locomotion policy in ~4h on a single desktop machine (32 GB memory,
+      Intel i7-8700K and Geforce GTX 1070) with a fully C++ code.
+    - the choice of the nonlinear activation function has a strong effect on performance on the physical system. The
+      authors advise for bounded soft activation functions such as tanh and softsign instead of ReLU for instance.
+    - with respect to the kernel function for some cost terms: "An Euclidean norm generates a high cost in the
+      beginning of training where the tracking error is high such that termination (i.e. falling) becomes more
+      rewarding strategy. On the other hand, the logistic kernel ensures that the cost is lower-bounded by zero and
+      termination becomes less favorable" [2].
 
     References:
         - [1] "Learning agile and dynamic motor skills for legged robots", Hwangbo et al., 2019
+        - [2] Supp: https://robotics.sciencemag.org/content/robotics/suppl/2019/01/14/4.26.eaau5872.DC1/aau5872_SM.pdf
     """
 
     def __init__(self, simulator=None, robot='anymal'):
@@ -569,10 +688,18 @@ class AgileLocomotionEnv(LocomotionEnv):
 # Test
 if __name__ == "__main__":
     from itertools import count
+    import time
 
     # create simulator
     sim = prl.simulators.Bullet()
 
+    N = int(6 * 1./sim.dt)
+    start = time.time()
+    for t in range(N):
+        sim.step(sleep_time=sim.dt)
+    end = time.time()
+    print("Total time: {}".format(end - start))
+
     # # create environment
     # env = RobustLocomotionQuadrupedEnv(sim)
     #

pyrobolearn/physics/actuators_randomizer.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+"""Define the `ActuatorRandomizer` class which randomizes the physical attributes / properties of an actuator.
+
+Dependencies:
+- `pyrobolearn.physics`
+"""
+
+import collections
+import numpy as np
+
+from pyrobolearn.physics.robot_physics_randomizer import RobotPhysicsRandomizer
+from pyrobolearn.robots.robot import Robot
+
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2019, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "GNU GPLv3"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class ActuatorRandomizer(RobotPhysicsRandomizer):
+    pass

pyrobolearn/physics/sensors_randomizer.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+"""Define the `SensorRandomizer` class which randomizes the physical attributes / properties of a sensor.
+
+Dependencies:
+- `pyrobolearn.physics`
+"""
+
+import collections
+import numpy as np
+
+from pyrobolearn.physics.robot_physics_randomizer import RobotPhysicsRandomizer
+from pyrobolearn.robots.robot import Robot
+
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2019, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "GNU GPLv3"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class SensorRandomizer(RobotPhysicsRandomizer):
+    pass

pyrobolearn/rewards/__init__.py

@@ -4,7 +4,10 @@ from .reward import Reward, ceil, cos, cosh, degrees, exp, expm1, floor, frexp,
     radians, sin, sinh, sqrt, tan, tanh, trunc
 
 # import basic rewards
-from .basic_rewards import *
+from .basic_rewards import FixedReward, DirectiveReward
+
+# import robot rewards
+from .robot_reward import RobotReward, BaseLinearVelocityReward, ForwardProgressReward
 
 # import gym wrapper reward
 from .gym_reward import GymReward
@@ -14,8 +17,9 @@ from .terminal_rewards import TerminalReward
 
 # import costs
 from .cost import *
+from .robot_cost import *
 from .joint_cost import *
 from .link_cost import *
 
-# import processors
+# import reward processors
 from .processors import *

pyrobolearn/rewards/basic_rewards.py

@@ -35,7 +35,7 @@ class FixedReward(Reward):
 
         Args:
             value (int, float): initial value.
-            range (None, tuple of float/int): A tuple corresponding to the min and max possible rewards. By default,
+            range (None, tuple[float/int]): A tuple corresponding to the min and max possible rewards. By default,
                 it is [value, value]. The initial value must be included in the given range.
         """
         super(FixedReward, self).__init__()
@@ -73,67 +73,6 @@ class FixedReward(Reward):
 #         return self.function()
 
 
-class ForwardProgressReward(Reward):
-    r"""Forward progress reward
-
-    Compute the forward progress based on a forward direction, a previous and current positions.
-    """
-
-    def __init__(self, state, direction=(1, 0, 0), normalize=False, update_state=False):
-        """
-        Initialize the Forward Progress Reward.
-
-        Args:
-            state (BasePositionState, PositionState, Robot): robot or base position state.
-            direction (np.float[3], None): forward direction vector. If None, it will take the initial forward vector.
-            normalize (bool): if we should normalize the direction vector.
-            update_state (bool): if we should call the state and update its value.
-        """
-        # check state argument
-        self.update_state = update_state
-        if isinstance(state, Robot):
-            state = states.BasePositionState(state)
-            self.update_state = True
-        elif not isinstance(state, (states.BasePositionState, states.PositionState)):
-            raise TypeError("Expecting the state to be an instance of `BasePositionState`, `PositionState`, or `Robot`"
-                            ", instead got: {}".format(type(state)))
-        super(ForwardProgressReward, self).__init__(state=state)
-
-        # if no direction specified, take the body forward vector
-        if direction is None:
-            self.direction = state.body.forward_vector
-        else:
-            self.direction = np.array(direction)
-
-        # normalize the direction vector if specified
-        if normalize:
-            self.direction = self.normalize(self.direction)
-
-        # remember current position
-        self.prev_pos = np.copy(self.state.data[0])
-        self.value = 0
-
-    @staticmethod
-    def normalize(x):
-        """
-        Normalize the given vector.
-        """
-        if np.allclose(x, 0):
-            return x
-        return x / np.linalg.norm(x)
-
-    def _compute(self):
-        """Compute the difference vector between the current and previous position (i.e. ~ velocity vector), and
-        compute the dot product between this velocity vector and the direction vector."""
-        if self.update_state:
-            self.state()
-        curr_pos = self.state.data[0]
-        velocity = curr_pos - self.prev_pos
-        self.value = self.direction.dot(velocity)
-        self.prev_pos = np.copy(curr_pos)
-        return self.value
-
-
 class DirectiveReward(Reward):
     r"""Directive Reward
 

pyrobolearn/rewards/cost.py

@@ -43,14 +43,22 @@ class Cost(Reward):
     Every classes that defines a cost inherits from this one. A cost is defined as an objective that
     penalizes a certain behavior.
     """
-    def __init__(self):
-        super(Cost, self).__init__()
+    def __init__(self, state=None, action=None, costs=None, range=(-np.infty, np.infty)):
+        super(Cost, self).__init__(state=state, action=action, rewards=costs, range=range)
         # super(Cost, self).__init__(maximize=False)
 
 
 def logistic_kernel_function(error, alpha):
     r"""
-    The logistic kernel function :math:`K(x|\alpha) = \frac{1}{(e^{\alpha x} + 2 + e^{-\alpha x})} \in [-0.25,0)`.
+    The logistic kernel function :math:`K(x|\alpha) = \frac{1}{(e^{\alpha x} + 2 + e^{-\alpha x})} \in [-0.25,0)`,
+    where :math:`x` is an error term, and :math:`\alpha` is a sensitivity factor.
+
+    According to the authors of [1,2]: "We found the logistic kernel function to be more useful than Euclidean norm,
+    which is a more common choice. An Euclidean norm generates a high cost in the beginning of training where the
+    tracking error is high such that termination (i.e. falling) becomes more rewarding strategy. On the other hand,
+    the logistic kernel ensures that the cost is lower-bounded by zero and termination becomes less favorable. Many
+    other bell-shaped kernels (Gaussian, triweight, biweight, etc) have the same functionality and can be used instead
+    of a logistic kernel."
 
     Args:
         error (Cost, float): cost (e.g. error term)
@@ -58,6 +66,10 @@ def logistic_kernel_function(error, alpha):
 
     Return:
         callable, float: logistic kernel function
+
+    References:
+        - [1] "Learning agile and dynamic motor skills for legged robots", Hwangbo et al., 2019
+        - [2] Supp: https://robotics.sciencemag.org/content/robotics/suppl/2019/01/14/4.26.eaau5872.DC1/aau5872_SM.pdf
     """
     if callable(error):
         y = copy.copy(error)  # shallow copy
@@ -140,29 +152,47 @@ class OrientationGravityCost(Cost):
         - [1] "Robust Recovery Controller for a Quadrupedal Robot using Deep Reinforcement Learning", Lee et al., 2019
     """
 
-    def __init__(self, gravity_state, gravity_vector=[0., 0., -1.]):
+    def __init__(self, state, gravity_vector=[0., 0., -1.]):
         super(OrientationGravityCost, self).__init__()
-        self.gravity_state = gravity_state
-        self.gravity = np.array(gravity_vector)
+        self.gravity_state = state
+        self.gravity = np.asarray(gravity_vector)
 
     def _compute(self):
-        return np.linalg.norm(self.gravity_state.data - self.gravity)
+        return np.linalg.norm(self.gravity_state.data[0] - self.gravity)
 
 
 class PowerCost(Cost):
     r"""Power Consumption Cost
 
-    Return the power consumption cost, where the power is computed as the torque times the velocity.
+    Return the power consumption cost, where the power is computed as the torque times the velocity. This is given by
+    [1] as:
+
+    .. math:: \sum
 
     References:
         - [1] "Robust Recovery Controller for a Quadrupedal Robot using Deep Reinforcement Learning", Lee et al., 2019
     """
+
     def __init__(self, torque_state, velocity_state):
         super(PowerCost, self).__init__()
         self.tau = torque_state
         self.vel = velocity_state
 
-    def compute(self):
+    def _compute(self):
+        return - np.sum(np.maximum(self.tau.data[0] * self.vel.data[0], 0))
+
+
+class PowerAbsoluteCost(Cost):
+    r"""Power Consumption Cost
+
+    """
+
+    def __init__(self, torque_state, velocity_state):
+        super(PowerCost, self).__init__()
+        self.tau = torque_state
+        self.vel = velocity_state
+
+    def _compute(self):
         return - np.sum(np.maximum(self.tau.data[0] * self.vel.data[0], 0))
 
 
@@ -288,31 +318,6 @@ class ImpactCost(Cost):
         pass
 
 
-class DriftCost(Cost):
-    """Drift cost.
-
-    Calculates the drift of a moving object wrt a direction.
-    """
-
-    def __init__(self, body, direction):
-        super(DriftCost, self).__init__()
-
-    def _compute(self):
-        pass
-
-
-class ShakeCost(Cost):
-    """Shake cost.
-
-    Calculates the
-    """
-    def __init__(self, body, direction):
-        super(ShakeCost, self).__init__()
-
-    def _compute(self):
-        pass
-
-
 class SpeedCost(Cost):
     """Speed cost.
 

pyrobolearn/rewards/joint_cost.py

@@ -21,6 +21,7 @@ __status__ = "Development"
 
 
 class JointCost(Cost):
+
     r"""(Abstract) Joint Cost."""
     __metaclass__ = ABCMeta
 
@@ -35,88 +36,171 @@ class JointCost(Cost):
         self.update_state = update_state
 
     @staticmethod
-    def _check_state(state, cls, update_state=False):
+    def _check_state(state, cls, update_state=False, **kwargs):
+        """
+        Check that the given state is an instance of the given class. If not, check if it can be constructed.
+
+        Args:
+            state (Robot, State, list/tuple[State]): the state or robot instance that we have to check.
+            cls (State class): the state class that the state should belong to.
+            update_state (bool): if the state should be updated or not by default.
+            **kwargs (dict): dictionary of arguments passed to the `cls` class  if the state is a `Robot` instance.
+
+        Returns:
+            State: an instance of the specified class `cls`.
+            bool: if the state should be updated or not.
+        """
         # check given state
         if isinstance(state, prl.robots.Robot):  # if robot, instantiate state class with robot as param.
-            state = cls(robot=state)
+            state = cls(robot=state, **kwargs)
             update_state = True
-        if not isinstance(state, cls):  # if not an instance of the given state, class, raise error
-            raise TypeError("Expecting the given 'state' to be an instance of `Robot` or `" + cls.__name__ + "`, "
-                            "but instead got: {}".format(type(state)))
+        if not isinstance(state, cls):  # if not an instance of the given state class, look for it (the first instance)
+            if isinstance(state, prl.states.State):
+                state = state.lookfor(cls)
+            elif isinstance(state, (tuple, list)):
+                for s in state:
+                    if isinstance(s, cls):
+                        state = s
+                    elif isinstance(s, prl.states.State):
+                        state = s.lookfor(cls)
+
+                    if state is not None:
+                        break
+            else:
+                raise TypeError("Expecting the given 'state' to be an instance of `Robot`, `{}`, `State` or a list of "
+                                "`State`, but instead got: {}".format(cls.__name__, type(state)))
+
+            if state is None:
+                raise ValueError("Couldn't find the specified state class `{}` in the given "
+                                 "state.".format(cls.__name__))
         return state, update_state
 
     @staticmethod
-    def _check_target_state(state, target_state, cls, update_state=False):
+    def _check_target_state(state, target_state, cls, update_state=False, **kwargs):
+        """
+        Check that the given target state is an instance of the given target state class. If not, check if it can be
+        constructed from it.
+
+        Args:
+            state (Robot, State, list/tuple[State]): the state associated to the given target state. This is used if
+                the target state is an int, float, or np.ndarray.
+            target_state (None, int, float, np.array, Robot, State, list/tuple[State]): the target state or robot
+                instance that we have to check.
+            cls (State class): the state class that the state should belong to.
+            update_state (bool): if the state should be updated or not by default.
+            **kwargs (dict): dictionary of arguments passed to the `cls` class  if the target state is a `Robot`
+                instance.
+
+        Returns:
+            State: an instance of the specified class `cls`.
+            bool: if the state should be updated or not.
+        """
         # check given target state
         if target_state is None:  # if the target is None, initialize it zero
             target_state = np.zeros(state.total_size())
         if isinstance(target_state, (int, float, np.ndarray)):  # if target is a np.array/float/int, create FixedState
+            # TODO: check shape
             target_state = prl.states.FixedState(value=target_state)
             update_state = True
         elif isinstance(target_state, prl.robots.Robot):  # if robot, instantiate state class with robot as param.
-            target_state = cls(robot=target_state)
+            target_state = cls(robot=target_state, **kwargs)
             update_state = True
-        elif not isinstance(target_state, cls):  # if not an instance of the given state class, raise error
-            raise TypeError("Expecting the given 'target_state' to be None, a np.array, or an instance of "
-                            "`Robot` or `" + cls.__name__ + "`, but instead got: {}".format(type(target_state)))
+        elif not isinstance(target_state, cls):  # if not an instance of the given state class, look for it
+            if isinstance(target_state, prl.states.State):
+                target_state = target_state.lookfor(cls)
+            elif isinstance(target_state, (tuple, list)):
+                for s in target_state:
+                    if isinstance(s, cls):
+                        target_state = s
+                    elif isinstance(s, prl.states.State):
+                        target_state = s.lookfor(cls)
+                    if target_state is not None:
+                        break
+            else:
+                raise TypeError("Expecting the given 'target_state' to be None, a np.array, an instance of `Robot`, "
+                                "`{}`, `State`, or a list of `State`, but instead got: "
+                                "{}".format(cls.__name__, type(target_state)))
+
+            if target_state is None:
+                raise ValueError("Couldn't find the specified target state class `{}` in the given "
+                                 "target_state.".format(cls.__name__))
+
         return target_state, update_state
 
 
-# class JointPositionErrorCost(JointCost):
-#     r"""Joint Position Error Cost
-#
-#     Return the joint position error as defined in [1] as :math:`d(\hat{\phi}, \phi) \in [0, \pi]` where :math:`d(.,.)`
-#     is the minimum angle difference between two angles, and :math:`\hat{\phi}` and :math:`\phi` are the target and
-#     current angles.
-#
-#     References:
-#         - [1] "Robust Recovery Controller for a Quadrupedal Robot using Deep Reinforcement Learning", Lee et al., 2019
-#     """
-#
-#     def __init__(self, joint_state, target_joint_state, update_state=False):
-#         """
-#         Initialize the joint position error cost.
-#
-#         Args:
-#             joint_state (JointPositionState):
-#             target_joint_state (JointPositionState):
-#             update_state (bool): if True it will update the given states before computing the cost.
-#         """
-#         super(JointPositionErrorCost, self).__init__(update_state=update_state)
-#         self.state = joint_state
-#         self.target_state = target_joint_state
-#
-#     def _compute(self):
-#         return - min_angle_difference(self.state.data[0], self.target_state.data[0])
-
-
-class JointPositionCost(JointCost):
-    r"""Joint Position Cost
+class JointAngleDifferenceCost(JointCost):
+    r"""Joint Angle Difference Cost
 
     Return the cost such that measures the L2 norm between the current joint positions and the target joint positions:
-    :math:`||d(q_{target},q)||^2` where :math:`d(\cdot, \cdot) \in [-\pi, \pi]` is the minimum distance between two
-    angles.
+
+    .. math:: ||d(q_{target},q)||^2,
+
+    where :math:`d(\cdot, \cdot) \in [-\pi, \pi]` is the minimum distance between two angles as described in [1,2].
 
     References:
         - [1] OpenAI Gym
         - [2] "Robust Recovery Controller for a Quadrupedal Robot using Deep Reinforcement Learning", Lee et al., 2019
     """
 
-    def __init__(self, state, target_state, update_state=False):
+    def __init__(self, state, target_state, joint_ids=None, update_state=False):
         r"""
         Initialize the joint position cost.
 
         Args:
-            state (JointPositionState, Robot): joint position state.
-            target_state (JointPositionState, np.array[N], None): target joint position state. If None, it will be set
-                to 0.
-            update_state (bool): if True it will update the given states before computing the cost.
+            state (JointPositionState, Robot): joint position state, or robot instance.
+            target_state (JointPositionState, np.array[float[N]], None): target joint position state. If None, it will
+                be set to 0.
+            joint_ids (None, int, list[int]): joint ids. This used if `state` is a `Robot` instance.
+            update_state (bool): if True, it will update the given states before computing the cost.
+        """
+        super(JointAngleDifferenceCost, self).__init__(update_state)
+
+        # check given joint position state
+        self.q, self.update_state = self._check_state(state, prl.states.JointPositionState,
+                                                      update_state=self.update_state, joint_ids=joint_ids)
+
+        # check target joint position state
+        self.q_target, self.update_target_state = self._check_target_state(self.q, target_state,
+                                                                           prl.states.JointPositionState,
+                                                                           self.update_state)
+
+        if self.q.total_size() != self.q_target.total_size():
+            raise ValueError("The given state and target_state do not have the same size: "
+                             "{} != {}".format(self.q.total_size(), self.q_target.total_size()))
+
+    def _compute(self):
+        """Compute and return the cost value."""
+        if self.update_state:
+            self.q()
+        if self.update_target_state:
+            self.q_target()
+        return - np.sum(min_angle_difference(self.q.data[0], self.q_target.data[0])**2)
+
+
+class JointPositionCost(JointCost):
+    r"""Joint Position Cost
+
+    Return the cost such that measures the L2 norm between the current joint positions and the target joint positions:
+
+    .. math:: ||q_{target} - q||^2`.
+    """
+
+    def __init__(self, state, target_state, joint_ids=None, update_state=False):
+        r"""
+        Initialize the joint position cost.
+
+        Args:
+            state (JointPositionState, Robot): joint position state, or robot instance.
+            target_state (JointPositionState, np.array[float[N]], None): target joint position state. If None, it will
+                be set to 0.
+            joint_ids (None, int, list[int]): joint ids. This used if `state` is a `Robot` instance.
+            update_state (bool): if True, it will update the given states before computing the cost.
         """
         super(JointPositionCost, self).__init__(update_state)
 
         # check given joint position state
         self.q, self.update_state = self._check_state(state, prl.states.JointPositionState,
-                                                      update_state=self.update_state)
+                                                      update_state=self.update_state, joint_ids=joint_ids)
 
         # check target joint position state
         self.q_target, self.update_target_state = self._check_target_state(self.q, target_state,
@@ -139,25 +223,29 @@ class JointPositionCost(JointCost):
 class JointVelocityCost(JointCost):
     r"""Joint Velocity Cost
 
-    Return the cost due to the joint velocities: :math:`|| \dot{q}_{target} - \dot{q} ||^2`, where
-    :math:`\dot{q}_{target}` can be set to zero if wished.
+    Return the cost due to the joint velocities given by:
+
+    .. math:: c = || \dot{q}_{target} - \dot{q} ||^2
+
+    where :math:`\dot{q}_{target}` can be set to zero if wished.
     """
 
-    def __init__(self, state, target_state=None, update_state=False):
+    def __init__(self, state, target_state=None, joint_ids=None, update_state=False):
         """
         Initialize the joint velocity cost.
 
         Args:
-            state (JointVelocityState, Robot): joint velocity state.
-            target_state (JointVelocityState, np.array[N], Robot, None): target joint velocity state. If None, it
-                will be set to 0.
-            update_state (bool): if True it will update the given states before computing the cost.
+            state (JointVelocityState, Robot): joint velocity state, or robot instance.
+            target_state (JointVelocityState, np.array[float[N]], Robot, None): target joint velocity state. If None,
+                it will be set to 0.
+            joint_ids (None, int, list[int]): joint ids. This used if `state` is a `Robot` instance.
+            update_state (bool): if True, it will update the given states before computing the cost.
         """
         super(JointVelocityCost, self).__init__(update_state)
 
         # check given joint velocity state
         self.dq, self.update_state = self._check_state(state, prl.states.JointVelocityState,
-                                                       update_state=self.update_state)
+                                                       update_state=self.update_state, joint_ids=joint_ids)
 
         # check target joint velocity state
         self.dq_target, self.update_target_state = self._check_target_state(self.dq, target_state,
@@ -176,28 +264,29 @@ class JointVelocityCost(JointCost):
 class JointAccelerationCost(JointCost):
     r"""Joint Acceleration Cost
 
-    Return the joint acceleration cost defined notably in [1] as :math:`cost = || \ddot{q}_{target} - \ddot{q} ||^2`,
-    where :math:`\ddot{q}_{target}` can be set to zero if wished.
+    Return the joint acceleration cost defined as:
 
-    References:
-        - [1] "Robust Recovery Controller for a Quadrupedal Robot using Deep Reinforcement Learning", Lee et al., 2019
+    .. math:: c = || \ddot{q}_{target} - \ddot{q} ||^2
+
+    where :math:`\ddot{q}_{target}` can be set to zero if wished.
     """
 
-    def __init__(self, state, target_state=None, update_state=False):
+    def __init__(self, state, target_state=None, joint_ids=None, update_state=False):
         """
         Initialize the joint acceleration cost.
 
         Args:
             state (JointAccelerationState, Robot): joint acceleration state.
-            target_state (JointAccelerationState, np.array[N], Robot, None): target joint acceleration state. If None,
-                it will be set to 0.
-            update_state (bool): if True it will update the given states before computing the cost.
+            target_state (JointAccelerationState, np.array[float[N]], Robot, None): target joint acceleration state.
+                If None, it will be set to 0.
+            joint_ids (None, int, list[int]): joint ids. This used if `state` is a `Robot` instance.
+            update_state (bool): if True, it will update the given states before computing the cost.
         """
         super(JointAccelerationCost, self).__init__(update_state=update_state)
 
         # check given joint acceleration state
         self.ddq, self.update_state = self._check_state(state, prl.states.JointAccelerationState,
-                                                        update_state=self.update_state)
+                                                        update_state=self.update_state, joint_ids=joint_ids)
 
         # check target joint acceleration state
         self.ddq_target, self.update_target_state = self._check_target_state(self.ddq, target_state,
@@ -216,25 +305,29 @@ class JointAccelerationCost(JointCost):
 class JointTorqueCost(JointCost):
     r"""Torque Cost
 
-    Return the cost due to the joint torques; :math:`|| \tau_{target} - \tau ||^2`, where :math:`\tau_{target}` can
-    be set to zero if wished.
+    Return the cost due to the joint torques given by:
+
+    .. math:: c = || \tau_{target} - \tau ||^2
+
+    where :math:`\tau_{target}` can be set to zero if wished.
     """
 
-    def __init__(self, state, target_state=None, update_state=False):
+    def __init__(self, state, target_state=None, joint_ids=None, update_state=False):
         """
         Initialize the joint torque cost.
 
         Args:
             state (JointForceTorqueState, Robot): joint torque state.
-            target_state (JointForceTorqueState, np.array[N], Robot, None): target joint torque state. If None, it
-                will be set to 0.
-            update_state (bool): if True it will update the given states before computing the cost.
+            target_state (JointForceTorqueState, np.array[float[N]], Robot, None): target joint torque state. If None,
+                it will be set to 0.
+            joint_ids (None, int, list[int]): joint ids. This used if `state` is a `Robot` instance.
+            update_state (bool): if True, it will update the given states before computing the cost.
         """
         super(JointTorqueCost, self).__init__(update_state)
 
         # check given joint torque state
         self.tau, self.update_state = self._check_state(state, prl.states.JointForceTorqueState,
-                                                        update_state=self.update_state)
+                                                        update_state=self.update_state, joint_ids=joint_ids)
 
         # check target joint torque state
         self.tau_target, self.update_target_state = self._check_target_state(self.tau, target_state,
@@ -251,9 +344,14 @@ class JointTorqueCost(JointCost):
 
 
 class JointPowerCost(JointCost):
-    r"""Joint Power Consumption Cost
+    r"""Joint Power Cost
 
-    Return the joint power consumption cost, where the power is computed as the torque times the velocity.
+    Return the joint power cost given by:
+
+    .. math:: c = ||\tau \cdot \dot{q}||^2
+
+    where :math:`\tau \in \mathcal{R}^N` are the torques, and :math:`\dot{q} \in \mathcal{R}^N` are the joint
+    velocities.
     """
 
     def __init__(self, state, joint_ids=None, update_state=False):
@@ -264,95 +362,140 @@ class JointPowerCost(JointCost):
             state (Robot, State): robot instance, or the state. The state must contains the `JointForceTorqueState`
                 and the `JointVelocityState`. Note that if they are multiple torque or velocity states, it will look
                 for the first instance.
-            joint_ids (None, int, list of int): joint ids. This used if `torque` is a `Robot` instance.
-            update_state (bool): if True it will update the given states before computing the cost.
+            joint_ids (None, int, list[int]): joint ids. This used if `state` is a `Robot` instance.
+            update_state (bool): if True, it will update the given states before computing the cost.
         """
-        self.update_state = update_state
+        super(JointPowerCost, self).__init__(update_state)
 
-        # Check the state
-        # if the given state is a robot, create the torque and velocity states
-        if isinstance(state, prl.robots.Robot):
-            torque = prl.states.JointForceTorqueState(state, joint_ids=joint_ids)
-            velocity = prl.states.JointVelocityState(state, joint_ids=joint_ids)
-            self.update_state = True
-            # state = torque + velocity
-
-        # elif the given state is a composite state, look for the torque and velocity states.
-        else:
-
-            if isinstance(state, prl.states.State):
-                state = [state]
-
-            # if the given state is a list of states, check each one of them by looking for the torque/velocity state
-            if isinstance(state, (list, tuple)):
-                # for each state, check if it is a torque, velocity or composite state
-                torque, velocity = None, None
-                for s in state:
-                    if isinstance(s, prl.states.JointForceTorqueState):
-                        if torque is None:
-                            torque = s
-                    elif isinstance(s, prl.states.JointVelocityState):
-                        if velocity is None:
-                            velocity = s
-                    elif isinstance(s, prl.states.State):
-                        if torque is None:
-                            torque = s.lookfor(prl.states.JointForceTorqueState)
-                        if velocity is None:
-                            velocity = s.lookfor(prl.states.JointVelocityState)
-                    else:
-                        raise TypeError("Expecting the state to be an instance of `State` or `Robot`, or a list of "
-                                        "`State`, instead got: {}".format(type(s)))
-                    # if we have found the states, get out of the loop
-                    if torque is not None and velocity is not None:
-                        break
-
-                # check that we have the torque and velocity states
-                if torque is None:
-                    raise ValueError("Didn't find a `JointForceTorqueState` instance in the given states.")
-                if velocity is None:
-                    raise ValueError("Didn't find a `JointVelocityState` instance in the given states.")
-            else:
-                raise TypeError("Expecting the state to be an instance of `State` or `Robot`, instead got: "
-                                "{}".format(type(state)))
-
-        super(JointPowerCost, self).__init__()
-        self.tau = torque
-        self.vel = velocity
-        self.update_state = update_state
+        # check given joint torque and velocity state
+        self.tau, self.update_torque_state = self._check_state(state, prl.states.JointForceTorqueState,
+                                                               update_state=self.update_state, joint_ids=joint_ids)
+        self.vel, self.update_velocity_state = self._check_state(state, prl.states.JointVelocityState,
+                                                                 update_state=self.update_state, joint_ids=joint_ids)
 
     def _compute(self):
         """Compute and return the cost value."""
-        if self.update_state:
+        if self.update_torque_state:
             self.tau()
+        if self.update_velocity_state:
             self.vel()
         return - np.sum((self.tau.data[0] * self.vel.data[0])**2)
 
 
-# class JointSpeedCost(Cost):
-#     r"""Joint Speed Cost
-#
-#     Return the joint speed cost as computed in [1].
-#
-#     .. math:: \text{cost} = || \max(\dot{q}_{max} - |\dot{q}|, 0) ||^2
-#
-#     References:
-#         - [1] "Robust Recovery Controller for a Quadrupedal Robot using Deep Reinforcement Learning", Lee et al., 2019
-#     """
-#
-#     def __init__(self, state, max_joint_speed=None):
-#         """
-#         Initialize the joint speed state.
-#
-#         Args:
-#             state: joint velocity state.
-#             max_joint_speed:
-#         """
-#         super(JointSpeedCost, self).__init__()
-#         self.dq = state
-#         self.dq_max = max_joint_speed
-#         if max_joint_speed is None:
-#             self.dq_max = state.max
-#
-#     def _compute(self):
-#         """Compute and return the cost value."""
-#         return - np.sum(np.maximum(self.dq_max - np.abs(self.dq.data[0]), 0)**2)
+class JointPowerConsumptionCost(JointCost):
+    r"""Joint Power Consumption Cost
+
+    Return the joint power consumption cost given by [1]:
+
+    .. math:: c = \sum_i^N max(\tau_i \dot{q}_i, 0)
+
+    where :math:`\tau \in \mathcal{R}^N` are the torques, and :math:`\dot{q} \in \mathcal{R}^N` are the joint
+    velocities.
+
+    References:
+        - [1] "Robust Recovery Controller for a Quadrupedal Robot using Deep Reinforcement Learning", Lee et al., 2019
+    """
+
+    def __init__(self, state, joint_ids=None, update_state=False):
+        """
+        Initialize the Joint Power Consumption cost.
+
+        Args:
+            state (Robot, State): robot instance, or the state. The state must contains the `JointForceTorqueState`
+                and the `JointVelocityState`. Note that if they are multiple torque or velocity states, it will look
+                for the first instance.
+            joint_ids (None, int, list[int]): joint ids. This used if `state` is a `Robot` instance.
+            update_state (bool): if True, it will update the given states before computing the cost.
+        """
+        super(JointPowerConsumptionCost, self).__init__(update_state)
+
+        # check given joint torque and velocity state
+        self.tau, self.update_torque_state = self._check_state(state, prl.states.JointForceTorqueState,
+                                                               update_state=self.update_state, joint_ids=joint_ids)
+        self.vel, self.update_velocity_state = self._check_state(state, prl.states.JointVelocityState,
+                                                                 update_state=self.update_state, joint_ids=joint_ids)
+
+    def _compute(self):
+        """Compute and return the cost value."""
+        if self.update_torque_state:
+            self.tau()
+        if self.update_velocity_state:
+            self.vel()
+        return - np.sum(np.maximum(self.tau.data[0] * self.vel.data[0], 0))
+
+
+class JointSpeedLimitCost(JointCost):
+    r"""Joint Speed Limit Cost
+
+    Return the joint speed cost as computed in [1].
+
+    .. math:: c = || \max(\dot{q}_{max} - |\dot{q}|, 0) ||^2
+
+    where :math:``
+
+    References:
+        - [1] "Robust Recovery Controller for a Quadrupedal Robot using Deep Reinforcement Learning", Lee et al., 2019
+    """
+
+    def __init__(self, state, max_joint_speed=None, joint_ids=None, update_state=False):
+        """
+        Initialize the joint speed state.
+
+        Args:
+            state (JointVelocityState, Robot): joint velocity state, or robot instance.
+            max_joint_speed (int, float, np.array[float[N]], None):
+            joint_ids (None, int, list[int]): joint ids. This used if `state` is a `Robot` instance.
+            update_state (bool): if True, it will update the given states before computing the cost.
+
+        """
+        super(JointSpeedLimitCost, self).__init__(update_state)
+        self.dq = state
+        self.dq_max = max_joint_speed
+        if max_joint_speed is None:
+            self.dq_max = state.max
+
+    def _compute(self):
+        """Compute and return the cost value."""
+        return - np.sum(np.maximum(self.dq_max - np.abs(self.dq.data[0]), 0)**2)
+
+
+class JointEnergyCost(JointCost):
+    r"""Joint Energy Cost
+
+    Return the joint energy cost given by:
+
+    .. math:: c = | \tau \cdot \dot{q}| * dt
+
+    where :math:`\tau \in \mathcal{R}^N` are the torques, :math:`\dot{q} \in \mathcal{R}^N` are the joint velocities,
+    and :math:`dt` is the simulation time step.
+    """
+
+    def __init__(self, state, dt, joint_ids=None, update_state=False):
+        """
+        Initialize the joint energy cost.
+
+        Args:
+            state (Robot, State): robot instance, or the state. The state must contains the `JointForceTorqueState`
+                and the `JointVelocityState`. Note that if they are multiple torque or velocity states, it will look
+                for the first instance.
+            dt (float): simulation time.
+            joint_ids (None, int, list[int]): joint ids. This used if `state` is a `Robot` instance.
+            update_state (bool): if True, it will update the given states before computing the cost.
+        """
+        super(JointEnergyCost, self).__init__(update_state)
+
+        # check given joint torque and velocity state
+        self.tau, self.update_torque_state = self._check_state(state, prl.states.JointForceTorqueState,
+                                                               update_state=self.update_state, joint_ids=joint_ids)
+        self.vel, self.update_velocity_state = self._check_state(state, prl.states.JointVelocityState,
+                                                                 update_state=self.update_state, joint_ids=joint_ids)
+
+        self.dt = dt
+
+    def _compute(self):
+        """Compute and return the cost value."""
+        if self.update_torque_state:
+            self.tau()
+        if self.update_velocity_state:
+            self.vel()
+        return np.abs(np.dot(self.tau.data[0], self.vel.data[0])) * self.dt

pyrobolearn/rewards/reward.py

@@ -108,6 +108,7 @@ class Reward(object):
         self.action = action
         self.rewards = rewards
         self.range = range
+        self.value = None  # cache the last float value computed by the reward
 
         # # create automatically binary operator methods
         # op_names = ['__add__', '__div__', '__floordiv__', '__iadd__', '__idiv__', '__ifloordiv__', '__imod__',

pyrobolearn/rewards/robot_cost.py

@@ -0,0 +1,193 @@
+#!/usr/bin/env python
+"""Define some basic robot costs used in reinforcement learning and optimization.
+
+Dependencies:
+- `pyrobolearn.states`
+- `pyrobolearn.actions`
+"""
+
+from abc import ABCMeta
+import numpy as np
+
+import pyrobolearn as prl
+from pyrobolearn.robots.robot import Robot
+from pyrobolearn.rewards.cost import Cost
+
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2019, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "GNU GPLv3"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class RobotCost(Cost):
+    r"""Robot reward (abstract).
+
+    Abstract reward class that accepts as input the state and/or action which must depends on a robotic platform.
+    """
+    __metaclass__ = ABCMeta
+
+    def __init__(self, state=None, update_state=False):
+        """
+        Initialize the Robot reward.
+
+        Args:
+            state (State, Robot): robot state.
+            update_state (bool): if we should call the state and update its value.
+        """
+        if isinstance(state, prl.states.State):
+            super(RobotCost, self).__init__(state)
+        else:
+            super(RobotCost, self).__init__()
+        self.update_state = update_state
+
+    @staticmethod
+    def _check_state(state, cls, update_state=False, **kwargs):
+        """
+        Check that the given state is an instance of the given class. If not, check if it can be constructed.
+
+        Args:
+            state (Robot, State, list/tuple[State]): the state or robot instance that we have to check.
+            cls (State class): the state class that the state should belong to.
+            update_state (bool): if the state should be updated or not by default.
+            **kwargs (dict): dictionary of arguments passed to the `cls` class  if the state is a `Robot` instance.
+
+        Returns:
+            State: an instance of the specified class `cls`.
+            bool: if the state should be updated or not.
+        """
+        # check given state
+        if isinstance(state, Robot):  # if robot, instantiate state class with robot as param.
+            state = cls(robot=state, **kwargs)
+            update_state = True
+        if not isinstance(state, cls):  # if not an instance of the given state class, look for it (the first instance)
+            if isinstance(state, prl.states.State):
+                state = state.lookfor(cls)
+            elif isinstance(state, (tuple, list)):
+                for s in state:
+                    if isinstance(s, cls):
+                        state = s
+                    elif isinstance(s, prl.states.State):
+                        state = s.lookfor(cls)
+
+                    if state is not None:
+                        break
+            else:
+                raise TypeError("Expecting the given 'state' to be an instance of `Robot`, `{}`, `State` or a list of "
+                                "`State`, but instead got: {}".format(cls.__name__, type(state)))
+
+            if state is None:
+                raise ValueError("Couldn't find the specified state class `{}` in the given "
+                                 "state.".format(cls.__name__))
+        return state, update_state
+
+    @staticmethod
+    def normalize(x):
+        """
+        Normalize the given vector.
+        """
+        if np.allclose(x, 0):
+            return x
+        return x / np.linalg.norm(x)
+
+
+class DriftCost(RobotCost):
+    """Drift cost.
+
+    Calculates the drift of a moving robot wrt a specified direction.
+    """
+
+    def __init__(self, state, direction=(1, 0, 0), normalize=False, update_state=False):  # TODO: use direction.
+        """
+        Initialize the drift cost.
+
+        Args:
+            state (BasePositionState, Robot): robot or base position state.
+            direction (np.array[float[3]], None): forward direction vector. If None, it will take the initial forward
+              vector.
+            normalize (bool): if we should normalize the direction vector.
+            update_state (bool): if we should call the state and update its value.
+        """
+        super(DriftCost, self).__init__(state=state, update_state=update_state)
+
+        # check given base position state
+        self.state, self.update_state = self._check_state(state, prl.states.BasePositionState,
+                                                          update_state=self.update_state)
+
+        # if no direction specified, take the body forward vector
+        if direction is None:
+            self.direction = self.state.body.forward_vector
+        else:
+            self.direction = np.array(direction)
+
+        # normalize the direction vector if specified
+        if normalize:
+            self.direction = self.normalize(self.direction)
+
+        # remember current position
+        self.prev_pos = np.copy(self.state.data[0])
+        self.value = 0
+
+    def _compute(self):
+        """Compute the difference vector between the current and previous position (i.e. ~ velocity vector), and
+        compute the dot product between this velocity vector and the direction vector."""
+        if self.update_state:
+            self.state()
+        curr_pos = self.state.data[0]
+        velocity = curr_pos - self.prev_pos
+        self.value = -np.abs(velocity[1])
+        self.prev_pos = np.copy(curr_pos)
+        return self.value
+
+
+class ShakeCost(RobotCost):
+    """Shake cost.
+
+    Calculates the shaking cost of a moving robot wrt a specified direction.
+    """
+
+    def __init__(self, state, direction=(1, 0, 0), normalize=False, update_state=False):
+        """
+        Initialize the shake cost.
+
+        Args:
+            state (BasePositionState, Robot): robot or base position state.
+            direction (np.array[float[3]], None): forward direction vector. If None, it will take the initial forward
+              vector.
+            normalize (bool): if we should normalize the direction vector.
+            update_state (bool): if we should call the state and update its value.
+        """
+        super(ShakeCost, self).__init__(state=state, update_state=update_state)
+
+        # check given base position state
+        self.state, self.update_state = self._check_state(state, prl.states.BasePositionState,
+                                                          update_state=self.update_state)
+
+        # if no direction specified, take the body forward vector
+        if direction is None:
+            self.direction = self.state.body.forward_vector
+        else:
+            self.direction = np.array(direction)
+
+        # normalize the direction vector if specified
+        if normalize:
+            self.direction = self.normalize(self.direction)
+
+        # remember current position
+        self.prev_pos = np.copy(self.state.data[0])
+        self.value = 0
+
+    def _compute(self):
+        """Compute the difference vector between the current and previous position (i.e. ~ velocity vector), and
+        compute the dot product between this velocity vector and the direction vector."""
+        if self.update_state:
+            self.state()
+        curr_pos = self.state.data[0]
+        velocity = curr_pos - self.prev_pos
+        self.value = -np.abs(velocity[2])
+        self.prev_pos = np.copy(curr_pos)
+        return self.value

pyrobolearn/rewards/robot_reward.py

@@ -0,0 +1,214 @@
+#!/usr/bin/env python
+"""Define some basic robot rewards used in reinforcement learning and optimization.
+
+Dependencies:
+- `pyrobolearn.states`
+- `pyrobolearn.actions`
+"""
+
+from abc import ABCMeta
+import numpy as np
+
+import pyrobolearn as prl
+from pyrobolearn.robots.robot import Robot
+from pyrobolearn.rewards.reward import Reward
+
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2019, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "GNU GPLv3"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class RobotReward(Reward):
+    r"""Robot reward (abstract).
+
+    Abstract reward class that accepts as input the state and/or action which must depends on a robotic platform.
+    """
+    __metaclass__ = ABCMeta
+
+    def __init__(self, state=None, update_state=False):
+        """
+        Initialize the Robot reward.
+
+        Args:
+            state (State, Robot): robot state.
+            update_state (bool): if we should call the state and update its value.
+        """
+        if isinstance(state, prl.states.State):
+            super(RobotReward, self).__init__(state)
+        else:
+            super(RobotReward, self).__init__()
+        self.update_state = update_state
+
+    @staticmethod
+    def _check_state(state, cls, update_state=False, **kwargs):
+        """
+        Check that the given state is an instance of the given class. If not, check if it can be constructed.
+
+        Args:
+            state (Robot, State, list/tuple[State]): the state or robot instance that we have to check.
+            cls (State class): the state class that the state should belong to.
+            update_state (bool): if the state should be updated or not by default.
+            **kwargs (dict): dictionary of arguments passed to the `cls` class  if the state is a `Robot` instance.
+
+        Returns:
+            State: an instance of the specified class `cls`.
+            bool: if the state should be updated or not.
+        """
+        # check given state
+        if isinstance(state, Robot):  # if robot, instantiate state class with robot as param.
+            state = cls(robot=state, **kwargs)
+            update_state = True
+        if not isinstance(state, cls):  # if not an instance of the given state class, look for it (the first instance)
+            if isinstance(state, prl.states.State):
+                state = state.lookfor(cls)
+            elif isinstance(state, (tuple, list)):
+                for s in state:
+                    if isinstance(s, cls):
+                        state = s
+                    elif isinstance(s, prl.states.State):
+                        state = s.lookfor(cls)
+
+                    if state is not None:
+                        break
+            else:
+                raise TypeError("Expecting the given 'state' to be an instance of `Robot`, `{}`, `State` or a list of "
+                                "`State`, but instead got: {}".format(cls.__name__, type(state)))
+
+            if state is None:
+                raise ValueError("Couldn't find the specified state class `{}` in the given "
+                                 "state.".format(cls.__name__))
+        return state, update_state
+
+    @staticmethod
+    def normalize(x):
+        """
+        Normalize the given vector.
+        """
+        if np.allclose(x, 0):
+            return x
+        return x / np.linalg.norm(x)
+
+
+class BaseLinearVelocityReward(RobotReward):
+    r"""Base Linear velocity reward
+
+    Compute the base linear velocity reward given by:
+
+    .. math:: r = || v ||^2
+
+    where :math:`v = [v_x, v_y, v_z] \in \mathcal{R}^3` is the base linear velocity vector.
+
+    You can also specify to consider only one axis such that `r = v_x` for instance, or specify an axis direction to
+    which we should take the scalar product, that is, instead of computing :math:`r = v \cdot v` it will compute
+    :math:`r = v \cdot a` where :math:`a` is the axis.
+    """
+
+    def __init__(self, state, axis=None, normalize=False, update_state=False):
+        """
+        Initialize the base velocity reward.
+
+        Args:
+            state (BaseLinearVelocityState, Robot): robot or base velocity state.
+            axis (None, int, np.array[float[3]]): axis to consider. If None, it will return the squared norm of the
+              velocity vector. If int, it will take the specified velocity component. If np.array, it will perform the
+              scalar product between the base velocity vector and the axis.
+            normalize (bool): if we should normalize the base velocity vector (and the axis vector if given). This
+              makes sure that the reward is between -1 and 1.
+            update_state (bool): if we should call the state and update its value.
+        """
+        super(BaseLinearVelocityReward, self).__init__(state=state, update_state=update_state)
+
+        # check given base velocity state
+        self.state, self.update_state = self._check_state(state, prl.states.BaseLinearVelocityState,
+                                                          update_state=self.update_state)
+
+        # check axis
+        if isinstance(axis, int):
+            if axis < 0:
+                axis = 0
+            elif axis > 2:
+                axis = 2
+        elif isinstance(axis, (list, tuple, np.ndarray)):
+            if len(axis) != 3:
+                raise ValueError("Expecting the axis to be a list/tuple/np.array of length 3 but got instead: "
+                                 "{}".format(len(axis)))
+            axis = np.asarray(axis)
+        elif axis is not None:
+            raise TypeError("Expecting the given 'axis' to be None, int, or np.array[float[3]], but got instead: "
+                            "{}".format(type(axis)))
+
+        # check normalize
+        self.need_to_normalize = normalize
+        if normalize and isinstance(axis, np.ndarray):
+            axis = self.normalize(axis)
+
+        self._axis = axis
+        self.value = 0
+
+    def _compute(self):
+        """Compute the base linear velocity reward."""
+        if self.update_state:
+            self.state()
+        v = self.state.data[0]
+        if self._axis is None:
+            self.value = self.normalize(v)
+        elif isinstance(self._axis, int):
+            self.value = v[self._axis]
+        else:
+            self.value = v.dot(self._axis)
+        return self.value
+
+
+class ForwardProgressReward(RobotReward):
+    r"""Forward progress reward
+
+    Compute the forward progress based on a forward direction, a previous and current positions.
+    """
+
+    def __init__(self, state, direction=(1, 0, 0), normalize=False, update_state=False):
+        """
+        Initialize the Forward Progress Reward.
+
+        Args:
+            state (BasePositionState, Robot): robot or base position state.
+            direction (np.array[float[3]], None): forward direction vector. If None, it will take the initial forward
+              vector.
+            normalize (bool): if we should normalize the direction vector.
+            update_state (bool): if we should call the state and update its value.
+        """
+        super(ForwardProgressReward, self).__init__(state=state, update_state=update_state)
+
+        # check given base position state
+        self.state, self.update_state = self._check_state(state, prl.states.BasePositionState,
+                                                          update_state=self.update_state)
+
+        # if no direction specified, take the body forward vector
+        if direction is None:
+            self.direction = self.state.body.forward_vector
+        else:
+            self.direction = np.array(direction)
+
+        # normalize the direction vector if specified
+        if normalize:
+            self.direction = self.normalize(self.direction)
+
+        # remember current position
+        self.prev_pos = np.copy(self.state.data[0])
+        self.value = 0
+
+    def _compute(self):
+        """Compute the difference vector between the current and previous position (i.e. ~ velocity vector), and
+        compute the dot product between this velocity vector and the direction vector."""
+        if self.update_state:
+            self.state()
+        curr_pos = self.state.data[0]
+        velocity = curr_pos - self.prev_pos
+        self.value = self.direction.dot(velocity)
+        self.prev_pos = np.copy(curr_pos)
+        return self.value

pyrobolearn/rewards/terminal_rewards.py

@@ -26,12 +26,12 @@ class TerminalReward(Reward):
     value once the goal has been achieved (e.g. games).
     """
 
-    def __init__(self, terminal_conditions, subreward, final_reward):
+    def __init__(self, terminal_conditions, subreward=0., final_reward=0.):
         r"""
         Terminal reward.
 
         Args:
-            terminal_conditions (TerminalCondition, list of TerminalCondition): terminal condition(s).
+            terminal_conditions (TerminalCondition, list[TerminalCondition]): terminal condition(s).
             subreward (Reward, float, int): sub reward that is called until the terminal condition is not fulfilled.
             final_reward (Reward, float, int): final reward that is called when the terminal condition has been reached.
         """

pyrobolearn/robots/anymal.py

@@ -71,6 +71,7 @@ class ANYmal(QuadrupedRobot):
                      if link in self.link_names]
 
         # taken from "Learning agile and dynamic motor skills for legged robots", Hwangbo et al., 2019
+        # nominal range of torque (30 N·m) / nominal range of motion (0.6 rad)
         self.kp = 50. * np.ones(12)
         self.kd = 0.1 * np.ones(12)
 

pyrobolearn/robots/minitaur.py

@@ -110,6 +110,17 @@ class Minitaur(QuadrupedRobot):
         # set feet friction
         self.set_foot_friction(frictions=foot_friction, feet_ids=self.feet)
 
+        h = np.pi / 2  # hip angle from [2]
+        k = 2.1834  # knee angle from [2]
+        right_front_leg_initial_pos = [-h, k, -h, k]  # (outer, inner)
+        right_back_leg_initial_pos = [h, -k, h, -k]  # (outer, inner)
+        left_front_leg_initial_pos = [h, -k, h, -k]  # (outer, inner)
+        left_back_leg_initial_pos = [-h, k, -h, k]  # (outer, inner)
+        self._joint_configuration = {'home': np.array(right_front_leg_initial_pos + right_back_leg_initial_pos +
+                                                      left_front_leg_initial_pos + left_back_leg_initial_pos),
+                                     'standing': 'home',
+                                     'init': 'home'}
+
         # set joint angles to home position
         self.set_home_joint_positions()
 

pyrobolearn/robots/sensors/camera.py

@@ -50,10 +50,10 @@ class CameraSensor(LinkSensor):  # TODO: double-check this class
         plt.show()
 
     References:
-        [1] http://www.codinglabs.net/article_world_view_projection_matrix.aspx
-        [2] https://learnopengl.com/Getting-started/Coordinate-Systems
-        [3] http://www.thecodecrate.com/opengl-es/opengl-transformation-matrices/
-        [4] http://learnwebgl.brown37.net/08_projections/projections_perspective.html
+        - [1] http://www.codinglabs.net/article_world_view_projection_matrix.aspx
+        - [2] https://learnopengl.com/Getting-started/Coordinate-Systems
+        - [3] http://www.thecodecrate.com/opengl-es/opengl-transformation-matrices/
+        - [4] http://learnwebgl.brown37.net/08_projections/projections_perspective.html
     """
 
     def __init__(self, simulator, body_id, link_id, width, height, noise=None, ticks=50, latency=None,
@@ -228,7 +228,7 @@ class DepthCameraSensor(CameraSensor):
             apply_noise (bool): if we should apply the noise or not. Note that the sensor might already have some noise.
 
         Returns:
-            np.array[W,H]: depth image
+            np.array[float[W,H]]: depth image
         """
         data = self.get_depth_image()
         if apply_noise:
@@ -248,7 +248,7 @@ class RGBCameraSensor(CameraSensor):
             apply_noise (bool): if we should apply the noise or not. Note that the sensor might already have some noise.
 
         Returns:
-            np.array[W,H]: depth image
+            np.array[int[W,H,3]]: RGB image
         """
         data = self.get_rgb_image()
         if apply_noise:
@@ -262,13 +262,13 @@ class RGBCameraSensor(CameraSensor):
 #
 #     def _sense(self, apply_noise=True):
 #         """
-#         Sense using the camera RGB sensor.
+#         Sense using the camera segmentation sensor.
 #
 #         Args:
 #             apply_noise (bool): if we should apply the noise or not. Note that the sensor might already have some noise.
 #
 #         Returns:
-#             np.array[W,H]: depth image
+#             np.array[int[W,H]]: segmentation image
 #         """
 #         data = self.get_rgb_image()
 #         if apply_noise:

pyrobolearn/robots/sensors/force_torque.py

@@ -46,7 +46,7 @@ class JointTorqueSensor(JointSensor):
             apply_noise (bool): if we should apply the noise or not. Note that the sensor might already have some noise.
 
         Returns:
-            np.array[N]: torque values
+            np.array[float[N]]: torque values
         """
         # check if the simulator supports that sensor
         if self.sim.supports_sensors("torque"):
@@ -97,7 +97,7 @@ class JointForceTorqueSensor(JointSensor):
             apply_noise (bool): if we should apply the noise or not. Note that the sensor might already have some noise.
 
         Returns:
-            np.array[6*N]: F/T values
+            np.array[float[6*N]]: F/T values
         """
         # check if the simulator supports that sensor
         if self.sim.supports_sensors("force-torque"):

pyrobolearn/robots/sensors/imu.py

@@ -72,7 +72,7 @@ class IMUSensor(LinkSensor):
             apply_noise (bool): if we should apply the noise or not. Note that the sensor might already have some noise.
 
         Returns:
-            np.array[6]: concatenation of linear accelerations and angular velocities
+            np.array[float[6]]: concatenation of linear accelerations and angular velocities
         """
         # if the simulator supports IMU sensors, return the sensed data
         if self.simulator.supports_sensors("imu"):

pyrobolearn/robots/sensors/ray.py

@@ -79,7 +79,7 @@ class RaySensor(LinkSensor):
 
         Args:
             enable (bool): if we should render or not.
-            color (None, tuple/list of 4 float, np.ndarray[float[4]]): RGBA color of all the rays, where each channel
+            color (None, tuple/list[float[4]], np.ndarray[float[4]]): RGBA color of all the rays, where each channel
                 is between 0 and 1.
         """
         pass
@@ -140,7 +140,7 @@ class RayBatchSensor(LinkSensor):
             apply_noise (bool): if we should apply the noise or not. Note that the sensor might already have some noise.
 
         Returns:
-            np.array[N]: hit fractions along each ray in range [0,1] along the ray.
+            np.array[float[N]]: hit fractions along each ray in range [0,1] along the ray.
         """
         if self.simulator.supports_sensors("ray_batch"):
             return self.simulator.get_sensor("ray_batch", self.body_id, self.link_id).sense()
@@ -157,7 +157,7 @@ class RayBatchSensor(LinkSensor):
 
         Args:
             enable (bool): if we should render or not.
-            color (None, tuple/list of 4 float, np.ndarray[float[4]]): RGBA color of all the rays, where each channel
+            color (None, tuple/list[float[4]], np.ndarray[float[4]]): RGBA color of all the rays, where each channel
                 is between 0 and 1.
         """
         pass
@@ -237,8 +237,8 @@ class HeightmapSensor(LinkSensor):
         Return the world positions for the rays to start and end.
 
         Returns:
-            np.array[N,3]: list of starting positions for the rays
-            np.array[N,3]: list of ending positions for the rays
+            np.array[float[N,3]]: list of starting positions for the rays
+            np.array[float[N,3]]: list of ending positions for the rays
         """
         pos = self.position
         w2, h2 = self._width / 2., self._height / 2.
@@ -259,8 +259,8 @@ class HeightmapSensor(LinkSensor):
         Return the heightmap.
 
         Returns:
-            np.array[width, height]: Height map with shape [width, height] where the values are the hit fractions [0,1],
-                you can multiply it by :attr:`max_ray_length` to get the depth in meters.
+            np.array[float[width, height]]: Height map with shape [width, height] where the values are the hit
+                fractions [0,1], you can multiply it by :attr:`max_ray_length` to get the depth in meters.
         """
         from_positions, to_positions = self.get_ray_from_to_positions()
         rays = self.sim.ray_test_batch(from_positions=from_positions, to_positions=to_positions)

pyrobolearn/robots/sensors/sensor.py

@@ -205,7 +205,7 @@ class Sensor(object):  # sensor attached to a link or joint
         """
         if self._enabled:
             self._cnt += 1
-            if (self._cnt % self._ticks) == 0:
+            if (self._cnt % self._ticks) == 0:  # if time to update
                 if self._latency == 0:  # if no latency
                     self._data = self._sense()
                     self._latent_data = self._data

pyrobolearn/simulators/bullet.py

@@ -373,43 +373,44 @@ class Bullet(Simulator):
             if mode == 'rgb':
                 np.array[W,H,D]: RGB image
         """
-        if enable:
-            if mode == 'human':
-                # self.sim.configureDebugVisualizer(self.sim.COV_ENABLE_RENDERING,  1)
-                if self.connection_mode == pybullet.DIRECT:
-                    # save the state of the simulator
-                    filename = 'PYROBOLEARN_RENDERING_STATE.bullet'
-                    self.save(filename=filename)
-                    # change the connection mode
-                    self.connection_mode = pybullet.GUI
-                    self.__init(self.connection_mode)
-                    # load the state of the world in the simulator
-                    self.load(filename)
-                    os.remove(filename)
-                    # reset the camera
-                    self.reset_scene_camera(camera=self._camera)
-            elif mode == 'rgb' or mode == 'rgba':
-                width, height, view_matrix, projection_matrix = self.get_debug_visualizer()[:4]
-                img = np.asarray(self.get_camera_image(width, height, view_matrix, projection_matrix)[2])
-                img = img.reshape(width, height, 4)  # RGBA
-                if mode == 'rgb':
-                    return img[:, :, :3]
-                return img
-        else:
-            if mode == 'human':
-                # self.sim.configureDebugVisualizer(self.sim.COV_ENABLE_RENDERING, 0)
-                if self.connection_mode == pybullet.GUI:
-                    # save the state of the simulator
-                    filename = 'PYROBOLEARN_RENDERING_STATE.bullet'
-                    self.save(filename=filename)
-                    # save main camera configuration (for later)
-                    self._camera = self.get_debug_visualizer()[-4:]
-                    # change the connection mode
-                    self.connection_mode = pybullet.DIRECT
-                    self.__init(self.connection_mode)
-                    # load the state of the world in the simulator
-                    self.load(filename)
-                    os.remove(filename)
+        if not self._render:
+            if enable:
+                if mode == 'human':
+                    # self.sim.configureDebugVisualizer(self.sim.COV_ENABLE_RENDERING,  1)
+                    if self.connection_mode == pybullet.DIRECT:
+                        # save the state of the simulator
+                        filename = 'PYROBOLEARN_RENDERING_STATE.bullet'
+                        self.save(filename=filename)
+                        # change the connection mode
+                        self.connection_mode = pybullet.GUI
+                        self.__init(self.connection_mode)
+                        # load the state of the world in the simulator
+                        self.load(filename)
+                        os.remove(filename)
+                        # reset the camera
+                        self.reset_scene_camera(camera=self._camera)
+                elif mode == 'rgb' or mode == 'rgba':
+                    width, height, view_matrix, projection_matrix = self.get_debug_visualizer()[:4]
+                    img = np.asarray(self.get_camera_image(width, height, view_matrix, projection_matrix)[2])
+                    img = img.reshape(width, height, 4)  # RGBA
+                    if mode == 'rgb':
+                        return img[:, :, :3]
+                    return img
+            else:
+                if mode == 'human':
+                    # self.sim.configureDebugVisualizer(self.sim.COV_ENABLE_RENDERING, 0)
+                    if self.connection_mode == pybullet.GUI:
+                        # save the state of the simulator
+                        filename = 'PYROBOLEARN_RENDERING_STATE.bullet'
+                        self.save(filename=filename)
+                        # save main camera configuration (for later)
+                        self._camera = self.get_debug_visualizer()[-4:]
+                        # change the connection mode
+                        self.connection_mode = pybullet.DIRECT
+                        self.__init(self.connection_mode)
+                        # load the state of the world in the simulator
+                        self.load(filename)
+                        os.remove(filename)
 
         # set the render variable (useful when calling the method `is_rendering`)
         self._render = enable

pyrobolearn/simulators/simulator.py

@@ -2105,20 +2105,20 @@ class Simulator(object):
 
         Returns:
             list:
-                int: contact flag (reserved)
-                int: body unique id of body A
-                int: body unique id of body B
-                int: link index of body A, -1 for base
-                int: link index of body B, -1 for base
-                np.array[float[3]]: contact position on A, in Cartesian world coordinates
-                np.array[float[3]]: contact position on B, in Cartesian world coordinates
-                np.array[float[3]]: contact normal on B, pointing towards A
-                float: contact distance, positive for separation, negative for penetration
-                float: normal force applied during the last `step`
-                float: lateral friction force in the first lateral friction direction (see next returned value)
-                np.array[float[3]]: first lateral friction direction
-                float: lateral friction force in the second lateral friction direction (see next returned value)
-                np.array[float[3]]: second lateral friction direction
+                [0] int: contact flag (reserved)
+                [1] int: body unique id of body A
+                [2] int: body unique id of body B
+                [3] int: link index of body A, -1 for base
+                [4] int: link index of body B, -1 for base
+                [5] np.array[float[3]]: contact position on A, in Cartesian world coordinates
+                [6] np.array[float[3]]: contact position on B, in Cartesian world coordinates
+                [7] np.array[float[3]]: contact normal on B, pointing towards A
+                [8] float: contact distance, positive for separation, negative for penetration
+                [9] float: normal force applied during the last `step`
+                [10] float: lateral friction force in the first lateral friction direction (see next returned value)
+                [11] np.array[float[3]]: first lateral friction direction
+                [12] float: lateral friction force in the second lateral friction direction (see next returned value)
+                [13] np.array[float[3]]: second lateral friction direction
         """
         pass
 

pyrobolearn/states/generators/state_generator.py

@@ -117,15 +117,19 @@ class FixedStateGenerator(StateGenerator):
     This generator returns the same initial state each time it is called.
     """
 
-    def __init__(self, state, fct=None):
+    def __init__(self, state, data=None, fct=None):
         """Initialize the fixed state generator.
 
         Args:
             state (State): state instance.
+            data (int, float, list[float/int] np.array[float/int], None): initial data. If None, it will get the data
+              from the given state, and set this last one as the initial data.
             fct (callable, None): callback function to be called after generating the data.
         """
         super(FixedStateGenerator, self).__init__(state, fct=fct)
-        self.initial_data = self.state.data
+        if data is None:
+            data = self.state.data
+        self.initial_data = data
 
     def _generate(self, set_data=True, reset_state=True):
         """Generate the state.
@@ -478,7 +482,7 @@ class UniformStateGenerator(StateDistributionGenerator):
                 for state, low, high in zip(self.state, self.low, self.high)]
         if set_data:
             self.state.data = data
-            # print("Generate: {}".format(self.state.data))
+            print("Generate: {}".format(self.state.data))
         if reset_state:
             self.state.reset()
         return data
@@ -491,21 +495,87 @@ class NormalStateGenerator(StateDistributionGenerator):
     The states are then truncated / clipped to be inside their corresponding range.
     """
 
-    def __init__(self, state, means=0, scales=1., seed=None, fct=None):
+    def __init__(self, state, mean=0, covariance=1., seed=None, fct=None):
         """
         Initialize the Normal state generator.
 
         Args:
             state (State): state instance.
-            means:
-            scales:
+            mean (int, float, np.array[float[N]]): mean.
+            scale (int, float, np.array[float[N]]): covariance matrix or variance vector.
             seed (None, int): random seed.
             fct (callable, None): callback function to be called after generating the data.
         """
         super(NormalStateGenerator, self).__init__(state, seed=seed, fct=fct)
+        self.mean = mean
+        self.covariance = covariance
+
+    @property
+    def mean(self):
+        """Return the mean vector."""
+        return self._mean
+
+    @mean.setter
+    def mean(self, mean):
+        """Set the mean vector."""
+        if mean is None:
+            mean = [0.] * len(self.state)
+        elif isinstance(mean, (int, float)):
+            mean = [mean] * len(self.state)
+        elif isinstance(mean, (list, tuple, np.ndarray)):
+            if len(mean) != len(self.state):
+                raise ValueError("The mean vector doesn't have the same size as the number of states; len(mean) = {} "
+                                 "and len(state) = {}".format(len(mean), len(self.state)))
+        else:
+            raise TypeError("Expecting the mean vector to be an int, float, or list/tuple/np.array of int/float, "
+                            "instead got: {}".format(type(mean)))
+        self._mean = np.asarray(mean)
+
+    @property
+    def covariance(self):
+        """Return the variance vector and covariance matrix."""
+        return self._covariance
+
+    @covariance.setter
+    def covariance(self, covariance):
+        """Set the variance vector or covariance matrix."""
+        if covariance is None:
+            covariance = [1.] * len(self.state)
+        elif isinstance(covariance, (int, float)):
+            covariance = [covariance] * len(self.state)
+        elif isinstance(covariance, (list, tuple, np.ndarray)):
+            if len(covariance) != len(self.state):
+                raise ValueError("The variance vector or covariance matrix doesn't have the same size as the number of "
+                                 "states; len(covariance) = {} and len(state) = {}".format(len(covariance),
+                                                                                           len(self.state)))
+        else:
+            raise TypeError("Expecting the variance vector or covariance matrix to be an int, float, or list/tuple/"
+                            "np.array of int/float, but instead got: {}".format(type(covariance)))
+        self._covariance = np.asarray(covariance)
 
     def _generate(self, set_data=True, reset_state=True):
-        pass
+        """Generate the state.
+
+        Args:
+            set_data (bool): If True, it will set the generated data to the state.
+            reset_state (bool): If True, it will reset the state with the generated data (if `set_data` has been set
+                to True).
+
+        Returns:
+            (list of) np.array: state data
+        """
+
+        if self.covariance.ndim == 2:
+            data = [np.random.multivariate_normal(self.mean, self.covariance)]
+        else:
+            data = [np.random.normal(self.mean, scale=np.sqrt(self.covariance))]
+
+        if set_data:
+            self.state.data = data
+            # print("Generate: {}".format(self.state.data))
+        if reset_state:
+            self.state.reset()
+        return data
 
 
 class GenerativeStateGenerator(StateGenerator):

pyrobolearn/states/robot_states/__init__.py

@@ -1,7 +1,7 @@
 
 # import the basic robot states
 from .robot_states import RobotState, BasePositionState, BaseHeightState, BaseOrientationState, BasePoseState, \
-    BaseLinearVelocityState, BaseAngularVelocityState, BaseVelocityState
+    BaseLinearVelocityState, BaseAngularVelocityState, BaseVelocityState, BaseAxisState
 
 # import the joint states
 from .joint_states import JointState, JointPositionState, JointTrigonometricPositionState, JointVelocityState, \

pyrobolearn/states/robot_states/joint_states.py

@@ -201,6 +201,13 @@ class JointVelocityState(JointState):
         """
         super(JointVelocityState, self).__init__(robot, joint_ids, window_size=window_size, axis=axis, ticks=ticks)
 
+        # define space
+        max_vel = self.robot.get_joint_max_velocities(joint_ids=self.joints)
+        if np.allclose(max_vel, 0):
+            print("WARNING: joint max velocities are 0, setting low=-10 and high=10.")
+            max_vel = 10 * np.ones(len(self.joints))  # TODO: np.infty instead of 10?
+        self._space = spaces.Box(low=-max_vel, high=max_vel, dtype=np.float32)
+
     def _read(self):
         """Read the next joint velocity state."""
         self.data = self.robot.get_joint_velocities(self.joints)
@@ -212,6 +219,7 @@ class JointVelocityState(JointState):
 
         # reset the robot joint position based on the data
         if len(self.data) > 0:
+            print("reset data: ", self.data[0])
             self.robot.reset_joint_states(dq=self.data[0], joint_ids=self.joints)
 
         # read the next data

pyrobolearn/states/robot_states/robot_states.py

@@ -15,6 +15,8 @@ import numpy as np
 from pyrobolearn.states import State
 from pyrobolearn.robots import Robot
 
+from pyrobolearn.utils.transformation import get_matrix_from_quaternion
+
 
 __author__ = "Brian Delhaisse"
 __copyright__ = "Copyright 2018, PyRoboLearn"
@@ -182,6 +184,50 @@ class BaseOrientationState(RobotState):
         self.data = self.robot.get_base_orientation()
 
 
+class BaseAxisState(RobotState):
+    r"""Base axis state
+
+    This is the state that returns the specified axis computed from the base orientation with respect to the world
+    frame. From the orientation expressed as a quaternion, this is first transformed as a rotation matrix from which a
+    column is then extracted (that column represents one of its axis).
+    """
+
+    def __init__(self, robot, base_axis=2, window_size=1, axis=None, ticks=1):
+        """
+        Initialize the base orientation state.
+
+        Args:
+            robot (Robot): instance of Robot which allows to access to the robot state.
+            base_axis (int): base axis index which is x=0, y=1, or z=2.
+            window_size (int): window size of the state. This is the total number of states we should remember. That
+                is, if the user wants to remember the current state :math:`s_t` and the previous state :math:`s_{t-1}`,
+                the window size is 2. By default, the :attr:`window_size` is one which means we only remember the
+                current state. The window size has to be bigger than 1. If it is below, it will be set automatically
+                to 1. The :attr:`window_size` attribute is only valid when the state is not a combination of states,
+                but is given some :attr:`data`.
+            axis (int, None): axis to concatenate or stack the states in the current window. If you have a state with
+                shape (n,), then if the axis is None (by default), it will just concatenate it such that resulting
+                state has a shape (n*w,) where w is the window size. If the axis is an integer, then it will just stack
+                the states in the specified axis. With the example, for axis=0, the resulting state has a shape of
+                (w,n), and for axis=-1 or 1, it will have a shape of (n,w). The :attr:`axis` attribute is only when the
+                state is not a combination of states, but is given some :attr:`data`.
+            ticks (int): number of ticks to sleep before getting the next state data.
+        """
+        # check basis axis
+        if base_axis < 0:
+            base_axis = 0
+        elif base_axis > 2:
+            base_axis = 2
+        self._base_axis = int(base_axis)
+
+        # call parent class
+        super(BaseAxisState, self).__init__(robot, window_size=window_size, axis=axis, ticks=ticks)
+
+    def _read(self):
+        """Read the orientation state data."""
+        self.data = get_matrix_from_quaternion(self.robot.get_base_orientation())[:, self._base_axis]
+
+
 class BasePoseState(RobotState):
     r"""Base pose state
 
@@ -213,7 +259,7 @@ class BasePoseState(RobotState):
 
     def _read(self):
         """Read the base position state data."""
-        self.data = self.robot.get_base_pose()
+        self.data = self.robot.get_base_pose(concatenate=True)
 
 
 class BaseLinearVelocityState(RobotState):
@@ -222,12 +268,15 @@ class BaseLinearVelocityState(RobotState):
     This is the state that returns the base linear velocity with respect to the world frame.
     """
 
-    def __init__(self, robot, window_size=1, axis=None, ticks=1):
+    def __init__(self, robot, link_frame=None, window_size=1, axis=None, ticks=1):
         """
         Initialize the base linear velocity state.
 
         Args:
-            robot (Robot): instance of Robot which allows to access to the robot state
+            robot (Robot): instance of Robot which allows to access to the robot state.
+            link_frame (None, int): the frame in which the base linear velocity is represented. If None, it will be
+                the world frame, if -1, it will be the base, and if another int, it will be the corresponding link
+                frame.
             window_size (int): window size of the state. This is the total number of states we should remember. That
                 is, if the user wants to remember the current state :math:`s_t` and the previous state :math:`s_{t-1}`,
                 the window size is 2. By default, the :attr:`window_size` is one which means we only remember the
@@ -255,12 +304,15 @@ class BaseAngularVelocityState(RobotState):
     This is the state that returns the base angular velocity with respect to the world frame.
     """
 
-    def __init__(self, robot, window_size=1, axis=None, ticks=1):
+    def __init__(self, robot, link_frame=None, window_size=1, axis=None, ticks=1):
         """
         Initialize the base position state.
 
         Args:
-            robot (Robot): instance of Robot which allows to access to the robot state
+            robot (Robot): instance of Robot which allows to access to the robot state.
+            link_frame (None, int): the frame in which the base angular velocity is represented. If None, it will be
+                the world frame, if -1, it will be the base, and if another int, it will be the corresponding link
+                frame.
             window_size (int): window size of the state. This is the total number of states we should remember. That
                 is, if the user wants to remember the current state :math:`s_t` and the previous state :math:`s_{t-1}`,
                 the window size is 2. By default, the :attr:`window_size` is one which means we only remember the
@@ -313,4 +365,4 @@ class BaseVelocityState(RobotState):
 
     def _read(self):
         """Read the base linear velocity state data."""
-        self.data = self.robot.get_base_velocity()
+        self.data = self.robot.get_base_velocity(concatenate=True)

pyrobolearn/states/state.py

@@ -386,6 +386,9 @@ class State(object):
 
     @property
     def spaces(self):
+        """
+        Get the corresponding spaces as a list of spaces.
+        """
         if self.has_space():
             return [self._space]
         return [state._space for state in self._states]
@@ -397,7 +400,8 @@ class State(object):
         """
         if self.has_space():
             # return [self._space]
-            return gym.spaces.Tuple([self._space])
+            # return gym.spaces.Tuple([self._space])
+            return self._space
         # return [state._space for state in self._states]
         return gym.spaces.Tuple([state._space for state in self._states])
 
@@ -412,9 +416,36 @@ class State(object):
     @property
     def merged_space(self):
         """
-        Get the corresponding merged space.
+        Get the corresponding merged space. Note that all the spaces have to be of the same type.
         """
-        return False
+        if self.has_space():
+            return self._space
+        spaces = self.spaces
+        result = []
+        dtype, prev_dtype = None, None
+        for space in spaces:
+            if isinstance(space, gym.spaces.Box):
+                dtype = 'box'
+                result.append([space.low, space.high])
+            elif isinstance(space, gym.spaces.Discrete):
+                dtype = 'discrete'
+                result.append(space.n)
+            else:
+                raise NotImplementedError
+
+            if prev_dtype is not None and dtype != prev_dtype:
+                return self.space
+
+            prev_dtype = dtype
+
+        if dtype == 'box':
+            low = np.concatenate([res[0] for res in result])
+            high = np.concatenate([res[1] for res in result])
+            return gym.spaces.Box(low=low, high=high, dtype=np.float32)
+        elif dtype == 'discrete':
+            return gym.spaces.Discrete(n=np.sum(result))
+
+        return self.space
 
     @property
     def name(self):
@@ -635,7 +666,7 @@ class State(object):
         """
         pass
 
-    def read(self):
+    def read(self, return_data=True, merged_data=False):
         """
         Read the state values from the simulator for each state, set it and return their values.
         """
@@ -653,7 +684,10 @@ class State(object):
         self._cnt += 1
 
         # return the data
-        return self.data
+        if return_data:
+            if merged_data:
+                return self.merged_data
+            return self.data
 
     def _reset(self):
         """
@@ -662,7 +696,7 @@ class State(object):
         self._cnt = 0
         self._read()
 
-    def reset(self):
+    def reset(self, return_data=True, merged_data=False):
         """
         Some states need to be reset. It returns the initial state.
         """
@@ -675,8 +709,11 @@ class State(object):
         else:  # else, reset this state
             self._reset()
 
-        # return the first state data
-        return self.data  # self.read()
+        # return the first state data if specified
+        if return_data:
+            if merged_data:
+                return self.merged_data
+            return self.data  # self.read()
 
     def max_dimension(self):
         """
@@ -857,7 +894,7 @@ class State(object):
             class_type (type, str): class type or name
 
         Returns:
-            State: the corresponding instance of the State class
+            State, None: the corresponding instance of the State class. None if it was not found.
         """
         # if string, lowercase it
         if isinstance(class_type, str):
@@ -897,11 +934,11 @@ class State(object):
     #         return [str(state) for state in self._states]
     #     return str(self)
 
-    def __call__(self):
+    def __call__(self, return_data=True, merged_data=False):
         """
         Compute/read the state and return it. It is an alias to the `self.read()` method.
         """
-        return self.read()
+        return self.read(return_data=return_data, merged_data=merged_data)
 
     def __len__(self):
         """

pyrobolearn/terminal_conditions/__init__.py

@@ -6,10 +6,11 @@ from .terminal_condition import TerminalCondition, HasFallenCondition, HasReache
 from .basic_conditions import TimeLimitCondition, GymTerminalCondition
 
 # import body terminal condition
-from .body_conditions import BodyCondition, PositionCondition, OrientationCondition
+from .body_conditions import BodyCondition, PositionCondition, OrientationCondition, DistanceCondition, \
+    BaseHeightCondition, BaseOrientationAxisCondition
 
 # import robot terminal condition
-from .robot_condition import RobotCondition
+from .robot_condition import RobotCondition, ContactCondition
 
 # import joint conditions
 from .joint_conditions import JointCondition, JointPositionCondition, JointVelocityCondition, \

pyrobolearn/terminal_conditions/basic_conditions.py

@@ -28,13 +28,13 @@ class TimeLimitCondition(TerminalCondition):
         Initialize the time limit terminal condition.
 
         Args:
-            num_steps (int):
+            num_steps (int): number of steps to perform in the environment.
             btype (bool, str, None): if the terminal condition represents a failure or success condition. If None, it
                 represents a neutral terminal condition (which is neither a failure or success condition, but just
                 means the episode is over). If string, it has to be among {"success", "failure", "neutral"}.
         """
         super(TimeLimitCondition, self).__init__(btype=btype)
-        self.num_steps = num_steps
+        self.num_steps = int(num_steps)
         self.cnt = 0
 
     def reset(self):

pyrobolearn/terminal_conditions/body_conditions.py

@@ -7,7 +7,7 @@ import numpy as np
 
 from pyrobolearn.robots.base import Body
 from pyrobolearn.terminal_conditions.terminal_condition import TerminalCondition
-from pyrobolearn.utils.transformation import get_rpy_from_quaternion
+from pyrobolearn.utils.transformation import get_rpy_from_quaternion, get_matrix_from_quaternion
 
 
 __author__ = "Brian Delhaisse"
@@ -23,7 +23,7 @@ __status__ = "Development"
 class BodyCondition(TerminalCondition):
     r"""Body Terminal Condition
 
-    This terminal condition describes 8 cases:
+    This terminal condition describes 8 cases (4 failure and 4 success cases):
 
     1. all the dimensions of the body state are:
         1. in a certain bounds and must stay between these bounds. Once one gets out, the terminal condition is over,
@@ -48,7 +48,7 @@ class BodyCondition(TerminalCondition):
     """
     __metaclass__ = ABCMeta
 
-    def __init__(self, body, bounds, dim=None, out=False, stay=False, all=False):
+    def __init__(self, body, bounds, dim=None, all=False, stay=False, out=False):
         """
         Initialize the body terminal condition.
 
@@ -58,13 +58,13 @@ class BodyCondition(TerminalCondition):
                 consider all 3 dimensions. If one dimension is provided it will only check along that dimension. If
                 a np.array of 0 and 1 is provided, it will consider the dimensions that are equal to 1. Thus, [1,0,1]
                 means to consider the bounds along the x and z axes.
-            out (bool): if True, we are outside the provided bounds. If False, we are inside the provided bounds.
-            stay (bool): if True, it must stay in the bounds defined by in_bounds or out_bounds; if the state
-                leave the bounds it results in a failure. if :attr:`stay` is False, it must get outside these bounds;
-                if the state leaves the bounds, it results in a success.
             all (bool): this is only used if they are multiple dimensions. if True, all the dimensions of the state
                 are checked if they are inside or outside the bounds depending on the other parameters. if False, any
                 dimensions will be checked.
+            stay (bool): if True, it must stay in the bounds defined by in_bounds or out_bounds; if the state
+                leave the bounds it results in a failure. if :attr:`stay` is False, it must get outside these bounds;
+                if the state leaves the bounds, it results in a success.
+            out (bool): if True, we are outside the provided bounds. If False, we are inside the provided bounds.
         """
         super(BodyCondition, self).__init__()
         self.body = body
@@ -101,8 +101,11 @@ class BodyCondition(TerminalCondition):
         """Set the dimensions."""
         if dim is not None:
             if not isinstance(dim, (int, np.ndarray)):
-                raise TypeError("Expecting the given 'dim' to be an int or an np.array of 3 int, but got instead: "
-                                "{}".format(type(dim)))
+                if isinstance(dim, (list, tuple)):
+                    dim = np.asarray(dim)
+                else:
+                    raise TypeError("Expecting the given 'dim' to be an int or an np.array of 3 int, but got instead: "
+                                    "{}".format(type(dim)))
             if isinstance(dim, np.ndarray):
                 if dim.size != 3:
                     raise ValueError("Expecting the given 'dim' np.array to be of size 3, but got instead a size of: "
@@ -110,6 +113,11 @@ class BodyCondition(TerminalCondition):
                 dim = np.array([bool(d) for d in dim])
         self._dim = dim
 
+    @property
+    def simulator(self):
+        """Return the simulator instance."""
+        return self.body.simulator
+
     ###########
     # Methods #
     ###########
@@ -156,14 +164,14 @@ class BodyCondition(TerminalCondition):
         if self._all:  # all the dimension states
             if self._out:  # are outside a certain bounds
                 if self._stay:  # and must stay outside these ones.
-                    if np.any(self.bounds[0] <= states <= self.bounds[1]):  # one dimension went inside
+                    if np.any((self.bounds[0] <= states) & (states <= self.bounds[1])):  # one dimension went inside
                         self._btype = False     # failure
                         self._over = True       # it is over
                     else:  # they are still all outside
                         self._btype = True      # success
                         self._over = False      # it is not over
                 else:  # and must go inside these ones
-                    if np.all(self.bounds[0] <= states <= self.bounds[1]):  # they all went inside
+                    if np.all((self.bounds[0] <= states) & (states <= self.bounds[1])):  # they all went inside
                         self._btype = True      # success
                         self._over = True       # it is over
                     else:  # they are some still left outside
@@ -171,14 +179,15 @@ class BodyCondition(TerminalCondition):
                         self._over = False      # it is not over
             else:  # are inside a certain bounds
                 if self._stay:  # and must stay inside these ones.
-                    if not np.all(self.bounds[0] <= states <= self.bounds[1]):  # one dimension went outside
+                    if not np.all((self.bounds[0] <= states) &
+                                  (states <= self.bounds[1])):  # one dimension went outside
                         self._btype = False  # failure
                         self._over = True    # it is over
                     else:  # they are still all inside
                         self._btype = True  # success
                         self._over = False  # it is not over
                 else:  # and must go outside these ones.
-                    if np.any(self.bounds[0] <= states <= self.bounds[1]):  # they are still some inside
+                    if np.any((self.bounds[0] <= states) & (states <= self.bounds[1])):  # they are still some inside
                         self._btype = False  # failure
                         self._over = False   # it is not over
                     else:  # they are all outside
@@ -188,14 +197,15 @@ class BodyCondition(TerminalCondition):
         else:  # any of the dimension states
             if self._out:  # is outside a certain bounds
                 if self._stay:  # and still stays outside these ones.
-                    if not np.all(self.bounds[0] <= states <= self.bounds[1]):  # at least one dim. is still outside
+                    if not np.all((self.bounds[0] <= states) &
+                                  (states <= self.bounds[1])):  # at least one dim. is still outside
                         self._btype = True   # success
                         self._over = False   # it is not over
                     else:  # they are all inside
                         self._btype = False   # failure
                         self._over = True     # it is over
                 else:  # and one must at least go inside these ones
-                    if np.any(self.bounds[0] <= states <= self.bounds[1]):  # at least one state is inside
+                    if np.any((self.bounds[0] <= states) & (states <= self.bounds[1])):  # at least one state is inside
                         self._btype = True  # success
                         self._over = True   # it is over
                     else:  # they are still all outside
@@ -203,14 +213,15 @@ class BodyCondition(TerminalCondition):
                         self._over = False   # it is not over
             else:  # is inside a certain bounds
                 if self._stay:  # and must stay inside these ones.
-                    if np.any(self.bounds[0] <= states <= self.bounds[1]):  # at least one state is still inside
+                    if np.any((self.bounds[0] <= states) &
+                              (states <= self.bounds[1])):  # at least one state is still inside
                         self._btype = True   # success
                         self._over = False   # it is not over
                     else:  # they are all outside
                         self._btype = False  # failure
                         self._over = True    # it is over
                 else:  # and must go outside these ones.
-                    if np.all(self.bounds[0] <= states <= self.bounds[1]):  # they are all inside
+                    if np.all((self.bounds[0] <= states) & (states <= self.bounds[1])):  # they are all inside
                         self._btype = False  # failure
                         self._over = False   # it is not over
                     else:  # at least one went outside
@@ -227,7 +238,7 @@ class BodyCondition(TerminalCondition):
 class PositionCondition(BodyCondition):
     r"""World position terminal condition
 
-    This terminal condition describes 8 cases:
+    This terminal condition describes 8 cases (4 failure and 4 success cases):
 
     1. all the dimensions of the body position state are:
         1. in a certain bounds and must stay between these bounds. Once one gets out, the terminal condition is over,
@@ -286,7 +297,7 @@ class PositionCondition(BodyCondition):
 class OrientationCondition(BodyCondition):
     r"""World orientation terminal condition
 
-    This terminal condition describes 8 cases:
+    This terminal condition describes 8 cases (4 failure and 4 success cases):
 
     1. all the dimensions of the body orientation (expressed as roll-pitch-yaw angles) state are:
         1. in a certain bounds and must stay between these bounds. Once one gets out, the terminal condition is over,
@@ -310,7 +321,7 @@ class OrientationCondition(BodyCondition):
 
     def __init__(self, body, bounds=(None, None), dim=None, out=False, stay=False, all=False):
         """
-        Initialize the world position terminal condition.
+        Initialize the world orientation terminal condition.
 
         Args:
             body (Body): body instance.
@@ -338,3 +349,130 @@ class OrientationCondition(BodyCondition):
         if self.dim is None:
             return orientation
         return orientation[self.dim]
+
+
+class BaseOrientationAxisCondition(BodyCondition):
+    r"""Base orientation axis terminal condition
+
+    This uses the cosine similarity function by computing the angle between the given axis and one of the axis
+    of the base orientation (i.e. one of the columns of the rotation matrix).
+
+    This terminal condition describes 4 cases (2 failure and 2 success cases); the angle is in:
+
+    1. in a certain bounds and must stay between these bounds. Once it gets out, the terminal condition is over,
+       and results in a failure. (stay=True, out=False --> must stay in)
+    2. in a certain bounds and must get out of these bounds. Once it gets out, the terminal condition is over,
+       and results in a success. (stay=False, out=False --> must not stay in)
+    3. outside a certain bounds and must get in. Once it gets in, the terminal condition is over, and results
+       in a success. (stay=False, out=True --> must not stay out)
+    4. outside a certain bounds and must stay outside these ones. Once it gets in, the terminal condition is over,
+       and results in a failure. (stay=True, out=True --> must stay out)
+    """
+
+    def __init__(self, body, angle=0.85, axis=(0., 0., 1.), dim=2, stay=False, out=False):
+        """
+        Initialize the base orientation axis terminal condition.
+
+        Args:
+            body (Body): body instance.
+            angle (float): angle bound.
+            axis (tuple/list[float[3]], np.array[float[3]]): axis.
+            dim (int): column that we should consider for the rotation matrix.
+            stay (bool): if True, it must stay in the bounds defined by in_bounds or out_bounds; if the orientation
+                leaves the bounds it results in a failure. if :attr:`stay` is False, it must get outside these bounds;
+                if the orientation leaves the bounds, it results in a success.
+            out (bool): if True, we are outside the provided bounds. If False, we are inside the provided bounds.
+        """
+        bounds = np.array([[angle], [1.1]])  # 1.1 is just to be sure
+        super(BaseOrientationAxisCondition, self).__init__(body, bounds=bounds, dim=dim, stay=stay, out=out)
+        self.axis = np.asarray(axis)
+
+    def _get_states(self):
+        """Return the state."""
+        axis = get_matrix_from_quaternion(self.body.orientation)[self.dim]
+        return np.dot(axis, self.axis)
+
+
+class BaseHeightCondition(BodyCondition):
+    r"""Base Height terminal condition
+
+    This terminal condition describes 4 cases (2 failure and 2 success cases); the base height (i.e. z-position) state
+    is:
+
+    1. in a certain bounds and must stay between these bounds. Once it gets out, the terminal condition is over,
+       and results in a failure. (stay=True, out=False --> must stay in)
+    2. in a certain bounds and must get out of these bounds. Once it gets out, the terminal condition is over,
+       and results in a success. (stay=False, out=False --> must not stay in)
+    3. outside a certain bounds and must get in. Once it gets in, the terminal condition is over, and results
+       in a success. (stay=False, out=True --> must not stay out)
+    4. outside a certain bounds and must stay outside these ones. Once it gets in, the terminal condition is over,
+       and results in a failure. (stay=True, out=True --> must stay out)
+    """
+
+    def __init__(self, body, height, stay=False, out=False):
+        """
+        Initialize the base height terminal condition.
+
+        Args:
+            body (Body): body instance.
+            height (float): max height which defines the bound; the bounds will be defined to be between 0 and height.
+            stay (bool): if True, it must stay in the bounds defined by in_bounds or out_bounds; if the position
+                leaves the bounds it results in a failure. if :attr:`stay` is False, it must get outside these bounds;
+                if the position leaves the bounds, it results in a success.
+            out (bool): if True, we are outside the provided bounds. If False, we are inside the provided bounds.
+        """
+        bounds = np.array([[0.], [height]])
+        super(BaseHeightCondition, self).__init__(body, bounds=bounds, stay=stay, out=out)
+
+    def _get_states(self):
+        """Return the state."""
+        return self.body.position[2]
+
+
+class DistanceCondition(BodyCondition):
+    r"""Distance terminal condition
+
+    This is a bit similar than the ``PositionCondition``. The difference is that this class describes a nd-sphere,
+    while the ``PositionCondition`` describes a nd-rectangle.
+
+    This terminal condition describes 4 cases (2 failure and 2 success cases); the body distance with respect to the
+    provided center must be:
+
+    1. in a certain bounds and must stay between these bounds. Once it gets out, the terminal condition is over,
+       and results in a failure. (stay=True, out=False --> must stay in)
+    2. in a certain bounds and must get out of these bounds. Once it gets out, the terminal condition is over,
+       and results in a success. (stay=False, out=False --> must not stay in)
+    3. outside a certain bounds and must get in. Once it gets in, the terminal condition is over, and results
+       in a success. (stay=False, out=True --> must not stay out)
+    4. outside a certain bounds and must stay outside these ones. Once it gets in, the terminal condition is over,
+       and results in a failure. (stay=True, out=True --> must stay out)
+    """
+
+    def __init__(self, body, distance=float("inf"), center=(0., 0., 0.), dim=None, stay=False, out=False):
+        """
+        Initialize the distance terminal condition.
+
+        Args:
+            body (Body): body instance.
+            distance (float): max distance with respect to the specified :attr:`center`.
+            center (np.array(float[3]), list[float[3]], tuple[float[3]]): center from which take the distance.
+            dim (None, int, int[3]): dimensions that we should consider when looking at the bounds. If None, it will
+                consider all 3 dimensions. If one dimension is provided it will only check along that dimension. If
+                a np.array of 0 and 1 is provided, it will consider the dimensions that are equal to 1. Thus, [1,0,1]
+                means to consider the distance along the x and z axes.
+            stay (bool): if True, it must stay in the bounds defined by in_bounds or out_bounds; if the position
+                leaves the bounds it results in a failure. if :attr:`stay` is False, it must get outside these bounds;
+                if the position leaves the bounds, it results in a success.
+            out (bool): if True, we are outside the provided bounds. If False, we are inside the provided bounds.
+
+        """
+        bounds = np.array([[0.], [distance]])
+        super(DistanceCondition, self).__init__(body, bounds=bounds, dim=dim, stay=stay, out=out)
+        self.center = np.asarray(center)
+
+    def _get_states(self):
+        """Return the state."""
+        position = self.body.position - self.center
+        if self.dim is None:
+            return np.linalg.norm(position)
+        return np.linalg.norm(position[self.dim])

pyrobolearn/terminal_conditions/joint_conditions.py

@@ -22,7 +22,7 @@ __status__ = "Development"
 class JointCondition(RobotCondition):
     r"""Joint Terminal Condition
 
-    This terminal condition describes 8 cases:
+    This terminal condition describes 8 cases (4 failure and 4 success cases):
 
     1. all the joint states are:
         1. in a certain bounds and must stay between these bounds. Once one gets out, the terminal condition is over,
@@ -106,7 +106,7 @@ class JointCondition(RobotCondition):
 class JointPositionCondition(JointCondition):
     r"""Joint position terminal condition
 
-    This terminal condition describes 8 cases:
+    This terminal condition describes 8 cases (4 failure and 4 success cases):
 
     1. all the joint positions are:
         1. in a certain bounds and must stay between these bounds. Once one gets out, the terminal condition is over,
@@ -155,7 +155,7 @@ class JointPositionCondition(JointCondition):
 class JointVelocityCondition(JointCondition):
     r"""Joint velocity terminal condition
 
-    This terminal condition describes 8 cases:
+    This terminal condition describes 8 cases (4 failure and 4 success cases):
 
     1. all the joint velocities are:
         1. in a certain bounds and must stay between these bounds. Once one gets out, the terminal condition is over,
@@ -204,7 +204,7 @@ class JointVelocityCondition(JointCondition):
 class JointAccelerationCondition(JointCondition):
     r"""Joint acceleration terminal condition
 
-    This terminal condition describes 8 cases:
+    This terminal condition describes 8 cases (4 failure and 4 success cases):
 
     1. all the joint accelerations are:
         1. in a certain bounds and must stay between these bounds. Once one gets out, the terminal condition is over,
@@ -253,7 +253,7 @@ class JointAccelerationCondition(JointCondition):
 class JointTorqueCondition(JointCondition):
     r"""Joint torque terminal condition
 
-    This terminal condition describes 8 cases:
+    This terminal condition describes 8 cases (4 failure and 4 success cases):
 
     1. all the joint torques are:
         1. in a certain bounds and must stay between these bounds. Once one gets out, the terminal condition is over,

pyrobolearn/terminal_conditions/link_conditions.py

@@ -23,7 +23,7 @@ __status__ = "Development"
 class LinkCondition(RobotCondition):
     r"""Link Terminal Condition
 
-    This terminal condition describes 8 cases:
+    This terminal condition describes 8 cases (4 failure and 4 success cases):
 
     1. all the dimensions of the link state are:
         1. in a certain bounds and must stay between these bounds. Once one gets out, the terminal condition is over,
@@ -48,7 +48,7 @@ class LinkCondition(RobotCondition):
     """
     __metaclass__ = ABCMeta
 
-    def __init__(self, robot, link_id, wrt_link_id=None, bounds=(None, None), dim=None, out=False, stay=False, all=all):
+    def __init__(self, robot, link_id, wrt_link_id=None, bounds=(None, None), out=False, stay=False, all=False):
         """
         Initialize the link terminal condition.
 
@@ -65,7 +65,7 @@ class LinkCondition(RobotCondition):
                 are checked if they are inside or outside the bounds depending on the other parameters. if False, any
                 dimensions will be checked.
         """
-        super(LinkCondition, self).__init__(robot, bounds=bounds, dim=dim, out=out, stay=stay, all=all)
+        super(LinkCondition, self).__init__(robot, bounds=bounds, out=out, stay=stay, all=all)
 
         # set link
         if not isinstance(link_id, int):
@@ -105,7 +105,7 @@ class LinkCondition(RobotCondition):
 class LinkPositionCondition(LinkCondition):
     r"""Link position terminal condition
 
-    This terminal condition describes 8 cases:
+    This terminal condition describes 8 cases (4 failure and 4 success cases):
 
     1. all the dimensions of the link position state are:
         1. in a certain bounds and must stay between these bounds. Once one gets out, the terminal condition is over,
@@ -127,8 +127,7 @@ class LinkPositionCondition(LinkCondition):
            over, and results in a failure. (all=False, out=True, stay=True)
     """
 
-    def __init__(self, robot, link_id=None, wrt_link_id=None, bounds=(None, None), dim=None, out=False, stay=False,
-                 all=False):
+    def __init__(self, robot, link_id=None, wrt_link_id=None, bounds=(None, None), out=False, stay=False, all=False):
         """
         Initialize the link position terminal condition.
 
@@ -145,7 +144,7 @@ class LinkPositionCondition(LinkCondition):
                 if the link position leaves the bounds, it results in a success.
         """
         super(LinkPositionCondition, self).__init__(robot, link_id=link_id, wrt_link_id=wrt_link_id, bounds=bounds,
-                                                    dim=dim, out=out, stay=stay, all=all)
+                                                    out=out, stay=stay, all=all)
 
     def _get_states(self):
         """Return the link position state."""
@@ -163,7 +162,7 @@ class LinkPositionCondition(LinkCondition):
 class LinkOrientationCondition(LinkCondition):
     r"""Link orientation terminal condition
 
-    This terminal condition describes 8 cases:
+    This terminal condition describes 8 cases (4 failure and 4 success cases):
 
     1. all the dimensions of the link orientation (expressed as roll-pitch-yaw angles) state are:
         1. in a certain bounds and must stay between these bounds. Once one gets out, the terminal condition is over,
@@ -187,8 +186,7 @@ class LinkOrientationCondition(LinkCondition):
     Warnings: the orientation is expressed as roll-pitch-yaw angles.
     """
 
-    def __init__(self, robot, link_id, wrt_link_id=None, bounds=(None, None), dim=None, out=False, stay=False,
-                 all=False):
+    def __init__(self, robot, link_id, wrt_link_id=None, bounds=(None, None), out=False, stay=False, all=False):
         """
         Initialize the link orientation terminal condition.
 
@@ -205,7 +203,7 @@ class LinkOrientationCondition(LinkCondition):
                 these bounds; if the link orientation leaves the bounds, it results in a success.
         """
         super(LinkOrientationCondition, self).__init__(robot, link_id=link_id, wrt_link_id=wrt_link_id, bounds=bounds,
-                                                       dim=dim, out=out, stay=stay, all=all)
+                                                       out=out, stay=stay, all=all)
 
     def _get_states(self):
         """Return the link orientation state."""

pyrobolearn/terminal_conditions/robot_condition.py

@@ -3,7 +3,9 @@
 """
 
 from abc import ABCMeta
+import numpy as np
 
+from pyrobolearn.robots.base import Body
 from pyrobolearn.robots.robot import Robot
 from pyrobolearn.terminal_conditions.body_conditions import BodyCondition
 
@@ -57,3 +59,109 @@ class RobotCondition(BodyCondition):
             raise TypeError("Expecting the given 'robot' to be an instance of `Robot`, instead got: "
                             "{}".format(type(robot)))
         self._robot = robot
+
+
+class ContactCondition(RobotCondition):
+    r"""Contact condition.
+
+    This terminal condition check if the given links are in or not in contact with another body or another link of a
+    body, and describes 8 cases (4 failure and 4 success cases):
+
+    1. all the links are with respect to the specified body/link:
+        1. in contact and must remain in contact. Once one is no more in contact, the terminal condition is over, and
+           results in a failure. (all=True, stay=True, out=False --> all must stay in (contact))
+        2. in contact and must no more be in contact. Once they are all no more in contact, the terminal condition is
+           over, and results in a success. (all=True, stay=False, out=False --> all must not stay in (contact))
+        3. not in contact initially but must all be in contact at the end. Once they all are in contact, the terminal
+           condition is over, and results in a success. (all=True, stay=False, out=True --> all must not stay out)
+        4. not in contact initially and must not be in contact at anytime. Once one link is in contact, the terminal
+           condition is over, and results in a failure. (all=True, stay=True, out=True --> all must stay out)
+    2. any of the links are with respect to the specified body/link:
+        1. in contact and must remain in contact. Once they all are not in contact anymore, the terminal condition is
+           over, and results in a failure. (all=False, stay=True, out=False --> any must stay in (contact))
+        2. in contact and must no more be in contact. Once one link is no more in contact, the terminal condition is
+           over, and results in a success. (all=False, stay=False, out=False --> any must not stay in (contact))
+        3. not in contact but must get in contact. Once one link gets in contact, the terminal condition is over,
+           and results in a success. (all=False, stay=False, out=True --> any must not stay out)
+        4. not in contact and must not be in contact at anytime. Once they all are in contact, the terminal condition
+           is over, and results in a failure. (all=False, stay=True, out=True --> any must stay out)
+    """
+
+    def __init__(self, robot, link_ids, wrt_body=None, wrt_link=-1, out=False, stay=False, all=False, complement=False):
+        """
+        Initialize the robot terminal condition.
+
+        Args:
+            robot (Robot): robot instance.
+            link_ids (int, list[int]): link ids that we must check if they are in contact or not (depending on the
+                next parameter `out`) with the specified body (see :attr:`wrt_body`) or link (see :attr:`wrt_link`).
+            wrt_body (Body, int, None): the body that we have to check if we are colliding with this one.
+            wrt_link (int, None): the link of :attr:`wrt_body` that we have to check if we are colliding with. If None,
+                it will consider all the links of :attr:`wrt_body`.
+            out (bool): if True, we are not in contact. If False, we are in contact.
+            stay (bool): if True, it must stay in contact; if one link or all links (depending on the next parameter
+                :attr:`all`) is/are no more in contact, it results in a failure. if :attr:`stay` is False, it must no
+                more be in contact at a later stage; if one or all (depending on the next parameter :attr:`all`) is/are
+                no more in contact, it results in a success.
+            all (bool): this is only used if they are multiple links. if True, all the links are checked if they are
+                in contact or not depending on the other parameters. if False, any links will be checked if they are
+                in contact or not.
+            complement (bool): if True, it will take the complement set of the specified link ids. For instance, if
+                the feet links are provided as an attribute, it will consider all the other links which are not feet.
+        """
+        super(ContactCondition, self).__init__(robot, out=out, stay=stay, all=all)
+
+        # check body ids
+        self.body_id1 = self.robot.id
+        if isinstance(wrt_body, Body):
+            wrt_body = wrt_body.id
+        self.body_id2 = wrt_body
+
+        # check link ids
+        if not isinstance(link_ids, (list, tuple, np.ndarray)):
+            link_ids = [link_ids]
+        for link_id in link_ids:
+            if not isinstance(link_id, int):
+                raise TypeError("Expecting each given link id must be an int, but we got instead: "
+                                "{}".format(type(link_id)))
+        # if complement
+        if complement:
+            links = [-1] + list(range(self.robot.num_links))
+            links = set(links)
+            link_ids = list(links.difference(link_ids))
+
+        self.link_ids = link_ids
+        self.link_dict = {link_id: idx for idx, link_id in enumerate(self.link_ids)}
+
+        # wrt_link
+        if wrt_body is None:
+            wrt_link = wrt_body
+        if wrt_link is not None and not isinstance(wrt_link, int):
+            raise TypeError("Expecting the given 'wrt_link' to be an int, instead got: {}".format(type(wrt_link)))
+        self.wrt_link = wrt_link
+
+        # define bounds
+        self.bounds = np.ones((2, len(self.link_ids)))  # (2, N)
+
+    def _get_states(self):
+        """Get the contact states."""
+
+        # contact_state = []
+        # for link_id in self.link_ids:
+        #     contacts = self.simulator.get_contact_points(body1=self.body_id1, body2=self.body_id2, link1_id=link_id,
+        #                                                  link2_id=self.wrt_link)
+        #     if contacts is None:
+        #         contact_state.append(0)
+        #     else:
+        #         n = len(contacts)
+        #         n = 1 if n > 0 else 0
+        #         contact_state.append(n)
+        # return np.array(contact_state)
+
+        contact_state = np.zeros(len(self.link_ids))
+        contacts = self.simulator.get_contact_points(body1=self.body_id1, body2=self.body_id2, link2_id=self.wrt_link)
+        for contact in contacts:
+            link_id = contact[3]
+            if link_id in self.link_dict:
+                contact_state[self.link_dict[link_id]] = 1
+        return contact_state