improve robots, CPG model/policy, and others

pyrobolearn/actions/robot_actions/joint_actions.py

@@ -55,16 +55,16 @@ class JointPositionAction(JointAction):
     Set the joint positions using position control.
     """
 
-    def __init__(self, robot, joint_ids=None, kp=None, kd=None):
-        self.kp, self.kd = kp, kd
+    def __init__(self, robot, joint_ids=None, kp=None, kd=None, max_force=None):
+        self.kp, self.kd, self.max_force = kp, kd, max_force
         super(JointPositionAction, self).__init__(robot, joint_ids)
-        self._data = robot.getJointPositions(self.joints)
+        self.data = robot.getJointPositions(self.joints)
 
     def _write(self, data=None):
         if data is None:
-            self.robot.setJointPositions(self._data, self.joints, kp=self.kp, kd=self.kd)
+            self.robot.setJointPositions(self._data, self.joints, kp=self.kp, kd=self.kd, maxTorque=self.max_force)
         else:
-            self.robot.setJointPositions(data, self.joints, kp=self.kp, kd=self.kd)
+            self.robot.setJointPositions(data, self.joints, kp=self.kp, kd=self.kd, maxTorque=self.max_force)
 
 
 class JointVelocityAction(JointAction):
@@ -75,7 +75,7 @@ class JointVelocityAction(JointAction):
 
     def __init__(self, robot, joint_ids=None):
         super(JointVelocityAction, self).__init__(robot, joint_ids)
-        self._data = robot.getJointVelocities(self.joints)
+        self.data = robot.getJointVelocities(self.joints)
 
     def _write(self, data=None):
         if data is None:
@@ -92,7 +92,7 @@ class JointForceAction(JointAction):
 
     def __init__(self, robot, joint_ids=None, f_min=-np.infty, f_max=np.infty):
         super(JointForceAction, self).__init__(robot, joint_ids)
-        self._data = robot.getJointTorques(self.joints)
+        self.data = robot.getJointTorques(self.joints)
         self.f_min = f_min
         self.f_max = f_max
 
@@ -114,7 +114,7 @@ class JointAccelerationAction(JointAction):
 
     def __init__(self, robot, joint_ids=None, a_min=-np.infty, a_max=np.infty):
         super(JointAccelerationAction, self).__init__(robot, joint_ids)
-        self._data = robot.getJointAccelerations(self.joints)
+        self.data = robot.getJointAccelerations(self.joints)
         self.a_min = a_min
         self.a_max = a_max
 

pyrobolearn/models/cpg.py

@@ -35,7 +35,64 @@ class CPGNode(object):
         [1] "Central pattern generators for locomotion control in animals and robots: a review", Ijspeert, 2008
     """
 
-    def __init__(self, id, phi=0, offset=0, amplitude=1., timesteps=100, freq=None):
+    def __init__(self, id, phi=0., offset=0., amplitude=1., timesteps=100, freq=1., update_amplitude=True,
+                 update_offset=True, update_frequency=True, update_init_phase=True, update_weights=True,
+                 update_biases=True, amplitude_bounds=np.pi, offset_bounds=np.pi, phase_bounds=np.pi,
+                 frequency_bounds=5., weight_bounds=2., bias_bounds=np.pi,  *args, **kwargs):
+        """
+        Initialize the Central Pattern generator
+
+        Args:
+            id (int): unique node id.
+            phi (float): initial phase value.
+            offset (float): desired offset.
+            amplitude (float): desired amplitude.
+            timesteps (int): total number of timesteps for the phase to do a complete cycle.
+            freq (float): desired frequency. From this value, the desired angular velocity (omega) is computed.
+            update_amplitude (bool): If True, it will allow to train the desired amplitudes.
+            update_offset (bool): If True, it will allow to train the desired offsets.
+            update_frequency (bool): If True, it will allow to train the desired frequencies.
+            update_phase (bool): If True, it will allow to optimize the initial phases.
+            update_weights (bool): If True, it will allow to optimize the coupling weights.
+            update_biases (bool): If True, it will allow to optimize the coupling biases.
+            amplitude_bounds (float, tuple of float): bounds / limits to the amplitude parameter (useful when training).
+            offset_bounds (float, tuple of float): bounds / limits to the offset parameter (useful when training).
+            phase_bounds (float, tuple of float): bounds / limits to the phase parameter (useful when training).
+            frequency_bounds (float, tuple of float): bounds / limits to the frequency parameter (useful when training).
+            weight_bounds (float, tuple of float): bounds / limits to the weight parameters (useful when training).
+            bias_bounds (float, tuple of float): bounds / limits to the bias parameters (useful when training).
+            *args (list): list of other arguments given to the CPG node. NOT CURRENTLY USED.
+            **kwargs (dict): dictionary of other arguments given to the CPG node. NOT CURRENTLY USED.
+        """
+        # set which parameters we can update
+        self.update_amplitude = update_amplitude
+        self.update_offset = update_offset
+        self.update_frequency = update_frequency
+        self.update_init_phase = update_init_phase
+        self.update_weights = update_weights
+        self.update_biases = update_biases
+
+        # set bounds
+        def set_bounds(bounds):
+            if isinstance(bounds, (list, tuple, np.ndarray)):
+                if len(bounds) != 2:
+                    raise ValueError("Got more than 2 bounds for one of the parameters. Expecting an upper and lower "
+                                     "bound, instead got: {}".format(bounds))
+            elif isinstance(bounds, (float, int)):
+                bounds = (-bounds, bounds)
+            else:
+                raise TypeError("Expecting bounds to a float, list or tuple of an upper bound and lower bound, "
+                                "instead got {}".format(type(bounds)))
+            return bounds
+
+        self.amplitude_bounds = set_bounds(amplitude_bounds)
+        self.offset_bounds = set_bounds(offset_bounds)
+        self.phase_bounds = set_bounds(phase_bounds)
+        self.frequency_bounds = set_bounds(frequency_bounds)
+        self.weight_bounds = set_bounds(weight_bounds)
+        self.bias_bounds = set_bounds(bias_bounds)
+
+        # usual variables
         self.id = id
         self.timesteps = timesteps
         self.dt = 1./self.timesteps
@@ -146,26 +203,49 @@ class CPGNode(object):
     #     self.curr_phi, self.phi = phi, phi
     #     self.theta = self.offset + self.amp * np.cos(self.phi)
 
-    def set_phi(self, phi):
-        self.curr_phi, self.phi = phi, phi
-        self.theta = self.offset + self.amp * np.cos(self.phi)
-
     @property
     def num_parameters(self):
-        return 3 + 2 * len(self.nodes)
+        """Return the total number of parameters."""
+        return len(self.list_parameters())
 
     ###########
     # Methods #
     ###########
 
+    def set_phi(self, phi):
+        """
+        Set/overwrite the phase value by the new provided one.
+
+        Args:
+            phi (float): new phase value.
+        """
+        self.curr_phi, self.phi = phi, phi
+        self.theta = self.offset + self.amp * np.cos(self.phi)
+
     def add_node(self, cpg_node, weight=0., bias=0.):
-        """Add a coupling node."""
+        """
+        Add a new coupling node.
+
+        Args:
+            cpg_node (CPGNode): node that is coupled with the current one.
+            weight (float): coupling weight.
+            bias (float): coupling bias.
+        """
         self.nodes[cpg_node] = {'weight': weight, 'bias': bias}
 
     def remove_node(self, cpg_node):
-        """Remove a coupling node"""
+        """
+        Remove a coupling node, and return it if it exists.
+
+        Args:
+            cpg_node (CPGNode): coupling node to be removed.
+
+        Returns:
+            CPGNode, None: coupling node that has been removed. None if it is not a node that is coupled with the
+                current one.
+        """
         if cpg_node in self.nodes:
-            self.nodes.pop(cpg_node)
+            return self.nodes.pop(cpg_node)
 
     def reset(self):
         """Reset the phase of the CPG node; this can be useful for phase resetting."""
@@ -176,40 +256,86 @@ class CPGNode(object):
     def parameters(self):
         """Returns an iterator over the model parameters."""
         # proper node parameters
-        yield self.des_amp
-        yield self.des_offset
-        yield self.des_freq
+        if self.update_amplitude:
+            yield self.des_amp
+        if self.update_offset:
+            yield self.des_offset
+        if self.update_frequency:
+            yield self.des_freq
+        if self.update_init_phase:
+            yield self.init_phi
 
         # coupling parameters
         for node in self.nodes:
-            yield self.nodes[node]['weight']
-            yield self.nodes[node]['bias']
+            if self.update_weights:
+                yield self.nodes[node]['weight']
+            if self.update_biases:
+                yield self.nodes[node]['bias']
+
+    def bounds(self):
+        """Return an iterator over the upper and lower bounds of the parameters."""
+        # bounds for the node parameters
+        if self.update_amplitude:
+            yield self.amplitude_bounds
+        if self.update_offset:
+            yield self.offset_bounds
+        if self.update_frequency:
+            yield self.frequency_bounds
+        if self.update_init_phase:
+            yield self.phase_bounds
+
+        # bounds for coupling parameters
+        for node in self.nodes:
+            if self.update_weights:
+                yield self.weight_bounds
+            if self.update_biases:
+                yield self.bias_bounds
 
     def named_parameters(self):
-        """Returns an iterator over the model parameters, yielding both the name and the parameter itself"""
+        """Returns an iterator over the model parameters, yielding both the name and the parameter itself."""
         # proper node parameters
-        yield str(self) + ':amplitude', self.des_amp
-        yield str(self) + ':offset', self.des_offset
-        yield str(self) + ':frequency', self.des_freq
+        if self.update_amplitude:
+            yield 'amplitude ' + str(self.id), self.des_amp
+        if self.update_offset:
+            yield 'offset ' + str(self.id), self.des_offset
+        if self.update_frequency:
+            yield 'frequency ' + str(self.id), self.des_freq
+        if self.update_init_phase:
+            yield 'init_phase ' + str(self.id), self.init_phi
 
         # coupling parameters
         for node in self.nodes:
-            yield str(self) + ':weight:' + str(node), self.nodes[node]['weight']
-            yield str(self) + ':bias:' + str(node), self.nodes[node]['bias']
+            if self.update_weights:
+                yield str(self.id) + ':weight:' + str(node.id), self.nodes[node]['weight']
+            if self.update_biases:
+                yield str(self.id) + ':bias:' + str(node.id), self.nodes[node]['bias']
 
     def list_parameters(self):
-        """Return a list of parameters"""
+        """Return a list of parameters."""
         return list(self.parameters())
 
     def get_vectorized_parameters(self, to_numpy=True):
-        """Return a vectorized form of the parameters (weights, biases, offsets)."""
+        """
+        Return a vectorized form of the parameters (weights, biases, offsets).
+
+        Args:
+            to_numpy (bool): If True, it will return a np.array for the vector, otherwise it will return a
+                torch.Tensor.
+
+        Returns:
+            np.array, torch.Tensor: parameter vector
+        """
         if to_numpy:
             return np.array(self.list_parameters())
-        else:
-            return torch.from_numpy(np.array(self.list_parameters()))
+        return torch.from_numpy(np.array(self.list_parameters()))
 
     def set_vectorized_parameters(self, vector):
-        """Set the vector parameters."""
+        """
+        Set the vector parameters.
+
+        Args:
+            vector (list of float, np.array): vector containing the parameter values.
+        """
         # set the parameters from the vectorized one
         if len(vector) != self.num_parameters:
             raise ValueError("Expecting the size of the vectorized parameters to match the number of parameters "
@@ -217,13 +343,14 @@ class CPGNode(object):
         self.des_amp = vector[0]
         self.des_offset = vector[1]
         self.des_freq = vector[2]
+        self.init_phi = vector[3]
 
         # coupling parameters
         for idx, node in enumerate(self.nodes):
             self.nodes[node]['weight'] = vector[2*idx + 3]
             self.nodes[node]['bias'] = vector[2*idx + 4]
 
-    def step(self):
+    def step(self, *args, **kwargs):
         """
         Perform a step by integrating (i.e. Euler integration) the differential equations governing the CPG.
 
@@ -483,22 +610,78 @@ class CPGNode(object):
 
     def __str__(self):
         """Return a string describing the class."""
-        return self.__class__.__name__ + "(" + str(self.id) + ")"
+        coupling_description = '\n\t\t'.join(['id: ' + str(node.id) + ' - weight: ' + str(values['weight']) +
+                                              ' - bias: ' + str(values['bias']) for node, values in self.nodes.items()])
+        if len(coupling_description) == 0:
+            coupling_description = str(None)
+
+        description = [self.__class__.__name__ + "(",
+                       '\n\tnode id: ' + str(self.id),
+                       '\n\tdesired offset: ' + str(self.des_offset),
+                       '\n\tdesired amplitude: ' + str(self.des_amp),
+                       '\n\tdesired frequency: ' + str(self.des_freq),
+                       '\n\tdesired angular velocity: ' + str(self.des_omega),
+                       '\n\tinitial phase: ' + str(self.init_phi),
+                       '\n\tcoupled nodes: ' + coupling_description + ')']
+
+        return ''.join(description)
 
 
 class CPGNetwork(object):
     r"""Central Pattern Generator Network
 
+    This creates a CPG network that couples different CPG node together. Their phase becomes dependent on other phase
+    in the network/graph.
     """
 
-    def __init__(self, nodes, timesteps=100):
+    def __init__(self, nodes, timesteps=100, update_amplitudes=True, update_offsets=True, update_init_phases=True,
+                 update_frequencies=True, update_weights=True, update_biases=True, amplitude_bounds=np.pi,
+                 offset_bounds=np.pi, phase_bounds=np.pi, frequency_bounds=5., weight_bounds=2., bias_bounds=np.pi,
+                 *args, **kwargs):
+        """
+        Initialize the CPG network.
+
+        Args:
+            nodes (dict, int, None): dictionary describing the CPG network. The syntax is the following:
+                nodes = {<node_id>: {'phi': <phi>, 'offset': <offset>, 'amplitude': <amplitude>, 'freq': <freq>,
+                                     'nodes': [{'id': <coupling_node_id>, 'bias': <coupling_bias>,
+                                                'weight': <coupling_weight>}, ...]}
+                where <node_id> is the id of the current node, <phi> is the initial phase, <offset> is the initial and
+                desired offset, <amplitude> is the amplitude of the signal sent by the CPG node, <freq> is the
+                frequency at which operates the node, then we add inside the list associated with the 'nodes' key in
+                the dictionary, each node which is coupled to the current node by specifying the id of the coupled
+                node, the coupling weight and bias.
+                If nodes = integer, it will be assumed it is the total number of nodes, and a fully connected CPG
+                network will be built. That is, it will connect all the nodes to each other.
+            timesteps (int): total number of timesteps for the phase to do a complete cycle.
+            update_amplitudes (bool): If True, it will allow to train the desired amplitudes.
+            update_offsets (bool): If True, it will allow to train the desired offsets.
+            update_init_phases (bool): If True, it will allow to optimize the initial phases.
+            update_frequencies (bool): If True, it will allow to train the desired frequencies.
+            update_weights (bool): If True, it will allow to optimize the coupling weights.
+            update_biases (bool): If True, it will allow to optimize the coupling biases.
+            amplitude_bounds (float, tuple of float): bounds / limits to the amplitude parameter (useful when training).
+            offset_bounds (float, tuple of float): bounds / limits to the offset parameter (useful when training).
+            phase_bounds (float, tuple of float): bounds / limits to the phase parameter (useful when training).
+            frequency_bounds (float, tuple of float): bounds / limits to the frequency parameter (useful when training).
+            weight_bounds (float, tuple of float): bounds / limits to the weight parameters (useful when training).
+            bias_bounds (float, tuple of float): bounds / limits to the bias parameters (useful when training).
+            *args (list): list of other arguments given to the CPG network. NOT CURRENTLY USED.
+            **kwargs (dict): dictionary of other arguments given to the CPG network. NOT CURRENTLY USED.
+        """
         nodes = nodes if isinstance(nodes, dict) else self.fully_connected_network(nodes)
 
         # create each node
         self.nodes, init_params = {}, {'phi', 'offset', 'amplitude', 'freq'}
         for node_id in nodes.keys():
             d = {key: val for key, val in nodes[node_id].items() if key in init_params}
-            self.nodes[node_id] = CPGNode(node_id, timesteps=timesteps, **d)
+            self.nodes[node_id] = CPGNode(node_id, timesteps=timesteps, update_amplitude=update_amplitudes,
+                                          update_offset=update_offsets, update_init_phase=update_init_phases,
+                                          update_frequency=update_frequencies, update_weights=update_weights,
+                                          update_biases=update_biases, amplitude_bounds=amplitude_bounds,
+                                          offset_bounds=offset_bounds, phase_bounds=phase_bounds,
+                                          frequency_bounds=frequency_bounds, weight_bounds=weight_bounds,
+                                          bias_bounds=bias_bounds, **d)
 
         # couple the nodes
         for node_id, node in self.nodes.items():
@@ -512,11 +695,19 @@ class CPGNetwork(object):
         self.nodes_id = self.nodes.keys()
         self.nodes_id.sort()
 
+    ##############
+    # Properties #
+    ##############
+
     @property
     def num_parameters(self):
         """Return the total number of parameters in this CPG network."""
         return sum([node.num_parameters for node in self.nodes.values()])
 
+    ###########
+    # Methods #
+    ###########
+
     def parameters(self):
         """Returns an iterator over the model parameters."""
         for node in self.nodes.values():
@@ -566,11 +757,11 @@ class CPGNetwork(object):
         for node in self.nodes.values():
             node.reset()
 
-    def step(self):
+    def step(self, *args, **kwargs):
         """Perform a step with the CPG network."""
         # perform one step
         for node in self.nodes.values():
-            node.step()
+            node.step(*args, **kwargs)
         # perform update
         for node in self.nodes.values():
             node.update()
@@ -606,6 +797,17 @@ class CPGNetwork(object):
              for i in node_ids}
         return d
 
+    #############
+    # Operators #
+    #############
+
+    def __str__(self):
+        """Return a string that describes the CPG network model."""
+        description = [self.__class__.__name__ + '(\n\t',
+                       '\n\n\t'.join([str(node) for node in self.nodes.values()]),
+                       '\n)']
+        return ''.join(description)
+
 
 # Tests
 if __name__ == "__main__":
@@ -703,10 +905,13 @@ if __name__ == "__main__":
         for _ in range(10*T):
             # Get angles from CPG network and set them to the minitaur
             act = network.step()
-            for i in range(len(act)//2): # Left
+            print("kp: ", minitaur._kp)
+            print("kd: ", minitaur._kd)
+            print("max force: ", minitaur._max_force)
+            for i in range(len(act)//2):  # Left
                 minitaur._SetDesiredMotorAngleById(minitaur._motor_id_list[2*i], act[i])
                 minitaur._SetDesiredMotorAngleById(minitaur._motor_id_list[2*i+1], -np.pi-act[i])
-            for i in range(len(act)//2, len(act)): # Right
+            for i in range(len(act)//2, len(act)):  # Right
                 minitaur._SetDesiredMotorAngleById(minitaur._motor_id_list[2*i], act[i])
                 minitaur._SetDesiredMotorAngleById(minitaur._motor_id_list[2*i+1], np.pi - act[i])
 

pyrobolearn/policies/cpg_policy.py

@@ -9,8 +9,8 @@ import torch
 
 from pyrobolearn.models import CPGNetwork
 from policy import Policy
-from pyrobolearn.states import State
-from pyrobolearn.actions import Action, JointPositionAction
+from pyrobolearn.states import State, PhaseState
+from pyrobolearn.actions import JointAction, Action
 
 
 __author__ = "Brian Delhaisse"
@@ -27,84 +27,177 @@ class CPGPolicy(Policy):
     r"""Central Pattern Generator (CPG) Network policy
     """
 
-    def __init__(self, states, actions, timesteps=100, rate=1, *args, **kwargs):
+    def __init__(self, states, actions, rate=1, cpg_network=None, amplitude=np.pi/4, offset=0., init_phase=0.,
+                 freq=1., couple_hips=False, parent_coupling=True, child_coupling=True,
+                 hip_coupling_weight=None, hip_coupling_bias=0., parent_coupling_weight=None, parent_coupling_bias=0.,
+                 child_coupling_weight=None, child_coupling_bias=0., update_amplitudes=True, update_offsets=True,
+                 update_init_phases=True, update_frequencies=True, update_weights=True, update_biases=True,
+                 amplitude_bounds=np.pi, offset_bounds=np.pi, phase_bounds=np.pi, frequency_bounds=5.,
+                 weight_bounds=2., bias_bounds=np.pi, *args, **kwargs):
+        """
+        Initialize the CPG Network Policy.
+
+        Args:
+            states (PhaseState): phase state.
+            actions (JointAction): joint action. Normally, it will be JointPositionAction.
+            rate (int): number of steps to wait before going to the next step with the policy.
+            cpg_network (dict, None): dictionary describing the CPG network. The syntax is the following:
+                cpg_network = {<node_id>: {'phi': <phi>, 'offset': <offset>, 'amplitude': <amplitude>, 'freq': <freq>,
+                                           'nodes': [{'id': <coupling_node_id>, 'bias': <coupling_bias>,
+                                                      'weight': <coupling_weight>}, ...]}
+                where <node_id> is the id of the current node, <phi> is the initial phase, <offset> is the initial and
+                desired offset, <amplitude> is the amplitude of the signal sent by the CPG node, <freq> is the
+                frequency at which operates the node, then we add inside the list associated with the 'nodes' key in
+                the dictionary, each node which is coupled to the current node by specifying the id of the coupled
+                node, the coupling weight and bias.
+                If None, it will create the CPG network automatically by using the other specified arguments.
+            amplitude (float, list of float): desired amplitude of each CPG node. If it is a list, it has to match the
+                number and order of joints returned by the given `JointAction`.
+            offset (float, list of float): desired offset of each CPG node. If it is a list, it has to match the number
+                and order of joints returned by the given `JointAction`.
+            init_phase (float, list of float): initial phase of each CPG node. If it is a list, it has to match the
+                number of joints returned by the given `JointAction`.
+            freq (float, list of float): initial desired frequency of each CPG node. If it is a list, it has to match
+                the number of joints returned by the given `JointAction`.
+            couple_hips (bool): If True it will couple in a bidirectional way the hips together. That is, the phase
+                of one hip joint influences the phase of another hip joint.
+            parent_coupling (bool): if enabled, then it will couple each node in the leg with its parent. That is,
+                the phase of the parent node influences / has an effect on the phase of its child. Assume a leg with
+                3 joints [hip, knee, ankle], the phase of the hip will influence the phase of the knee, and the phase
+                of the knee will have an effect on the phase of the ankle.
+            child_coupling (bool): if enabled, then it will couple each node with its child. That is, the phase of
+                the child node influences / has an effect on the phase of its child. Assume a leg with 3 joints [hip,
+                knee, ankle], the phase of the ankle influences the phase of the knee, and the phase of the knee has
+                an effect on the phase of the hip.
+            hip_coupling_weight (float, None): coupling weight between the hips. If None, it will be 1 divided by
+                the number of hips.
+            hip_coupling_bias (float): hip coupling bias.
+            parent_coupling_weight (float, None): coupling weight between the parent and the current node. It is used
+                when `parent_coupling` is enabled. If None, it will be 1 divided by the number of joints in the leg.
+            parent_coupling_bias (float): parent coupling bias; this bias is used when `parent_coupling` is enabled.
+            child_coupling_weight (float, None): coupling weight between the child and the current node. It is used
+                when `child_coupling` is enabled. If None, it will be 1 divided by the number of joints in the leg.
+            child_coupling_bias (float): child coupling bias; this bias is used when `child_coupling` is enabled.
+            update_amplitudes (bool): If True, it will allow to train the desired amplitudes.
+            update_offsets (bool): If True, it will allow to train the desired offsets.
+            update_frequencies (bool): If True, it will allow to train the desired frequencies.
+            update_init_phases (bool): If True, it will allow to optimize the initial phases.
+            update_weights (bool): If True, it will allow to optimize the coupling weights.
+            update_biases (bool): If True, it will allow to optimize the coupling biases.
+            amplitude_bounds (float, tuple of float): bounds / limits to the amplitude parameter (useful when training).
+            offset_bounds (float, tuple of float): bounds / limits to the offset parameter (useful when training).
+            phase_bounds (float, tuple of float): bounds / limits to the phase parameter (useful when training).
+            frequency_bounds (float, tuple of float): bounds / limits to the frequency parameter (useful when training).
+            weight_bounds (float, tuple of float): bounds / limits to the weight parameters (useful when training).
+            bias_bounds (float, tuple of float): bounds / limits to the bias parameters (useful when training).
+            *args (list): other arguments given to the CPG network learning model.
+            **kwargs (dict): other key + value arguments given to the CPG network learning model.
+        """
         super(CPGPolicy, self).__init__(states, actions, rate=rate, *args, **kwargs)
 
         # check actions
-        if not isinstance(actions, JointPositionAction):
-            raise TypeError("Expecting the actions to be an instance of JointPositionAction, instead got: "
+        if not isinstance(actions, JointAction):
+            raise TypeError("Expecting the actions to be an instance of JointAction, instead got: "
                             "{}".format(type(actions)))
 
-        # create CPG network based on the robot kinematic structures
+        # check states
+        if not isinstance(states, PhaseState):
+            raise TypeError("Expecting the states to be an instance of PhaseState, instead got: "
+                            "{}".format(type(states)))
 
-        # get specified legs
+        # get useful information from the state/action
+        timesteps = states.num_steps
         robot = actions.robot
         joints = set(actions.joints)
-        legs = []
-        for robot_leg in robot.legs:
-            leg = []
-            for joint in robot_leg:
-                if joint in joints:
-                    leg.append(joint)
-            legs.append(leg)
 
-        num_legs = len(legs)
+        # create CPG network based on the robot kinematic structures if not provided
+        if cpg_network is None:
+            # get specified legs
+            legs = []
+            for robot_leg in robot.legs:
+                leg = []
+                for joint in robot_leg:
+                    if joint in joints:
+                        leg.append(joint)
+                legs.append(leg)
 
-        # define few variables to initialize the CPG nodes
-        # variables for the node
-        init_phi = 0.
-        offset = 0.
-        amplitude = 1.
-        freq = 1.
-        # variables for coupling the nodes
-        weight_legs = 1. / len(legs)
-        if len(legs) > 0 and len(legs[0]) > 0:
-            weight_leg = 1. / len(legs[0])
+            num_legs = len(legs)
+
+            # variables for coupling the nodes
+            if couple_hips:
+                if hip_coupling_weight is None:
+                    hip_coupling_weight = 1. / num_legs
+
+            init_parent_coupling_weight = parent_coupling_weight
+            init_child_coupling_weight = child_coupling_weight
+
+            # create the CPG network based on the robot kinematic structures
+            cpg_network = {}
+            for leg_idx, leg in enumerate(legs):
+                # compute parent/child leg coupling weight
+                if len(leg) != 0:
+                    if parent_coupling and init_parent_coupling_weight is None:
+                        parent_coupling_weight = 1. / len(leg)
+                    if child_coupling and init_child_coupling_weight is None:
+                        child_coupling_weight = 1. / len(leg)
+
+                for idx, joint in enumerate(leg):
+                    # proper node parameters
+                    node = {'phi': init_phase, 'offset': offset, 'amplitude': amplitude, 'freq': freq}
+
+                    # coupling parameters
+
+                    # if first upper joint in the leg, connect it with the other upper joints (in the other legs)
+                    if couple_hips and idx == 0:
+                        for l in legs:
+                            if len(l) > 0 and joint != l[0]:  # i.e. not the same node
+                                coupling_params = {'id': l[0], 'weight': hip_coupling_weight, 'bias': hip_coupling_bias}
+                                node.setdefault('nodes', []).append(coupling_params)
+
+                    # add coupling to current joint (except the last one) with the next joint in the leg
+                    if parent_coupling and idx < len(leg)-1:
+                        # add next node
+                        coupling_params = {'id': leg[idx + 1], 'weight': parent_coupling_weight,
+                                           'bias': parent_coupling_bias}
+                        node.setdefault('nodes', []).append(coupling_params)
+
+                    # add coupling to current joint (except the first one) with the previous joint in the leg
+                    if child_coupling and idx > 0:
+                        # add next node
+                        coupling_params = {'id': leg[idx - 1], 'weight': child_coupling_weight,
+                                           'bias': child_coupling_bias}
+                        node.setdefault('nodes', []).append(coupling_params)
+
+                    # add node in the CPG network dictionary
+                    cpg_network[joint] = node
         else:
-            weight_leg = 0.
-        bias = 0.
+            # Quick check
+            if isinstance(joints, int):
+                joints = [joints]
+            if len(cpg_network) != len(joints):
+                cpg_network_ids = [idx for idx in cpg_network]
+                raise ValueError("The number of joints doesn't match up the number of CPG id in the CPG network. "
+                                 "The joints specified in the joint actions are {}, and the joint ids in the CPG "
+                                 "network are {}".format(joints, cpg_network_ids))
 
-        # create the CPG network based on the robot kinematic structures
-        nodes = {}
-        for leg_idx, leg in enumerate(legs):
-            for idx, joint in enumerate(leg):
-                # proper node parameters
-                node = {'phi': init_phi, 'offset': offset, 'amplitude': amplitude, 'freq': freq}
-
-                # coupling parameters
-
-                # if first upper joint in the leg, connect it with the other upper joints (in the other legs)
-                if idx == 0:
-                    for l in legs:
-                        if len(l) > 0 and joint != l[0]:  # i.e. not the same node
-                            coupling_params = {'id': l[0], 'weight': weight_legs, 'bias': bias}
-                            node.setdefault('nodes', []).append(coupling_params)
-
-                # add coupling to current joint (except the last one) with the next joint in the leg
-                if idx < len(leg)-1:
-                    # add next node
-                    coupling_params = {'id': leg[idx + 1], 'weight': weight_leg, 'bias': bias}
-                    node.setdefault('nodes', []).append(coupling_params)
-
-                # add coupling to current joint (except the first one) with the previous joint in the leg
-                if idx > 0:
-                    # add next node
-                    coupling_params = {'id': leg[idx - 1], 'weight': weight_leg, 'bias': bias}
-                    node.setdefault('nodes', []).append(coupling_params)
-
-                # add node in the CPG network dictionary
-                nodes[joint] = node
+        self.cpg_network = cpg_network
 
         # create learning model
-        self.model = CPGNetwork(nodes=nodes, timesteps=timesteps)
+        self.model = CPGNetwork(nodes=cpg_network, timesteps=timesteps, update_amplitudes=update_amplitudes,
+                                update_offsets=update_offsets, update_init_phases=update_init_phases,
+                                update_frequencies=update_frequencies, update_weights=update_weights,
+                                update_biases=update_biases, amplitude_bounds=amplitude_bounds,
+                                offset_bounds=offset_bounds, phase_bounds=phase_bounds,
+                                frequency_bounds=frequency_bounds, weight_bounds=weight_bounds,
+                                bias_bounds=bias_bounds, *args, **kwargs)
 
     def _size(self, x):
         size = 0
         if isinstance(x, (State, Action)):
-            if x.isDiscrete():
+            if x.is_discrete():
                 size = x.space[0].n
             else:
-                size = x.totalSize()
+                size = x.total_size()
         elif isinstance(x, np.ndarray):
             size = x.size
         elif isinstance(x, torch.Tensor):
@@ -113,7 +206,7 @@ class CPGPolicy(Policy):
             size = x
         return size
 
-    def act(self, state=None, deterministic=True, to_numpy=True):
+    def act(self, state=None, deterministic=True, to_numpy=True, return_logits=False):
         if (self.cnt % self.rate) == 0:
             self.last_action = self.model.step()
         self.cnt += 1
@@ -127,4 +220,12 @@ class CPGPolicy(Policy):
 
     def phase_resetting(self):
         self.model.reset()
-        pass
+
+    #############
+    # Operators #
+    #############
+
+    def __str__(self):
+        """Return a string that describes the CPG policy."""
+        description = self.__class__.__name__ + '(' + str(self.model) + ')'
+        return description

pyrobolearn/robots/legged_robot.py

@@ -2,6 +2,8 @@
 """Provide the Legged robot abstract classes.
 """
 
+import collections
+import numpy as np
 from robot import Robot
 
 
@@ -133,8 +135,37 @@ class LeggedRobot(Robot):
     def turnRight(self):
         raise NotImplementedError
 
+    def setFootFriction(self, friction, feet_id=None):
+        """
+        Set the foot friction in the simulator.
+
+        Warnings: only available in the simulator.
+
+        Args:
+            friction (float, list of float): friction value(s).
+            feet_id (int, list of int): list of foot/feet id(s).
+        """
+        if feet_id is None:
+            foot_id = self.feet
+        if isinstance(feet_id, int):
+            feet_id = [feet_id]
+        if isinstance(friction, (float, int)):
+            friction = friction * np.ones(len(feet_id))
+        for foot_id, frict in zip(feet_id, friction):
+            if isinstance(foot_id, int):
+                self.sim.changeDynamics(self.id, foot_id, lateralFriction=frict)
+            elif isinstance(foot_id, collections.Iterable):
+                for idx in foot_id:
+                    self.sim.changeDynamics(self.id, idx, lateralFriction=frict)
+            else:
+                raise TypeError("Expecting foot_id to be a list of int, or an int. Instead got: "
+                                "{}".format(type(foot_id)))
+
 
 class BipedRobot(LeggedRobot):
+    r"""Biped Robot
+
+    """
 
     def __init__(self, simulator, urdf_path, init_pos=(0,0,1.5), init_orient=(0,0,0,1), useFixedBase=False, scaling=1.):
         super(BipedRobot, self).__init__(simulator, urdf_path, init_pos, init_orient, useFixedBase, scaling)
@@ -184,6 +215,9 @@ class BipedRobot(LeggedRobot):
 
 
 class QuadrupedRobot(LeggedRobot):
+    r"""Quadruped robot
+
+    """
 
     def __init__(self, simulator, urdf_path, init_pos=(0,0,1.), init_orient=(0,0,0,1), useFixedBase=False, scaling=1.):
         super(QuadrupedRobot, self).__init__(simulator, urdf_path, init_pos, init_orient, useFixedBase, scaling)

pyrobolearn/robots/minitaur.py

@@ -3,6 +3,8 @@
 """
 
 import os
+import collections
+import numpy as np
 from legged_robot import QuadrupedRobot
 
 
@@ -13,20 +15,23 @@ class Minitaur(QuadrupedRobot):
 
     References:
         [1] pybullet_envs/bullet/minitaur.py
+        [2] https://github.com/bulletphysics/bullet3/blob/master/examples/pybullet/gym/pybullet_envs/bullet/minitaur.py
     """
 
     def __init__(self,
                  simulator,
-                 init_pos=(0, 0, .2),
+                 init_pos=(0, 0, .3),
                  init_orient=(0, 0, 0, 1),
                  useFixedBase=False,
                  scaling=1.,
+                 couple_legs=True,
+                 foot_friction=1.,
                  urdf_path=os.path.dirname(__file__) + '/urdfs/minitaur/minitaur.urdf'):
         # check parameters
         if init_pos is None:
-            init_pos = (0., 0., 0.2)
+            init_pos = (0., 0., 0.3)
         if len(init_pos) == 2:  # assume x, y are given
-            init_pos = tuple(init_pos) + (0.2,)
+            init_pos = tuple(init_pos) + (0.3,)
         if init_orient is None:
             init_orient = (0, 0, 0, 1)
         if useFixedBase is None:
@@ -36,14 +41,14 @@ class Minitaur(QuadrupedRobot):
         self.name = 'minitaur'
 
         self.legs = [[self.getLinkIds(link) for link in links if link in self.link_names]
-                     for links in [['motor_front_leftL_link', 'motor_front_leftR_link',
-                                    'lower_leg_front_leftL_link', 'lower_leg_front_leftR_link'],
-                                   ['motor_front_rightL_link', 'motor_front_rightR_link',
-                                    'lower_leg_front_rightL_link', 'lower_leg_front_rightR_link'],
-                                   ['motor_back_leftL_link', 'motor_back_leftR_link',
-                                    'lower_leg_back_leftL_link', 'lower_leg_back_leftR_link'],
-                                   ['motor_back_rightL_link', 'motor_back_rightR_link',
-                                    'lower_leg_back_rightL_link', 'lower_leg_back_rightR_link']]]
+                     for links in [['motor_front_leftL_link', 'lower_leg_front_leftL_link',
+                                    'motor_front_leftR_link', 'lower_leg_front_leftR_link'],
+                                   ['motor_front_rightL_link', 'lower_leg_front_rightL_link',
+                                    'motor_front_rightR_link', 'lower_leg_front_rightR_link'],
+                                   ['motor_back_leftL_link', 'lower_leg_back_leftL_link',
+                                    'motor_back_leftR_link', 'lower_leg_back_leftR_link'],
+                                   ['motor_back_rightL_link', 'lower_leg_back_rightL_link',
+                                    'motor_back_rightR_link', 'lower_leg_back_rightR_link']]]
 
         self.feet = [[self.getLinkIds(link) for link in links if link in self.link_names]
                      for links in [['lower_leg_front_leftL_link', 'lower_leg_front_leftR_link'],
@@ -51,6 +56,156 @@ class Minitaur(QuadrupedRobot):
                                    ['lower_leg_back_leftL_link', 'lower_leg_back_leftR_link'],
                                    ['lower_leg_back_rightL_link', 'lower_leg_back_rightR_link']]]
 
+        self.outer_legs = [self.left_front_outer_leg, self.right_front_outer_leg,
+                           self.left_back_outer_leg, self.right_back_outer_leg]
+        self.inner_legs = [self.left_front_inner_leg, self.right_front_inner_leg,
+                           self.left_back_inner_leg, self.right_back_inner_leg]
+
+        self.outer_hips = [leg[0] for leg in self.outer_legs]
+        self.inner_hips = [leg[0] for leg in self.inner_legs]
+        self.outer_knees = [leg[1] for leg in self.outer_legs]
+        self.inner_knees = [leg[1] for leg in self.inner_legs]
+
+        self.hip_directions = [1, -1] * 4  # outer/inner
+        self.knee_directions = [-1, 1] * 4  # outer/inner
+
+        self.outer_joints = set([joint for leg in self.outer_legs for joint in leg])
+
+        self.init_joint_positions = self.getHomeJointPositions()
+
+        # constraints
+        self.couple_legs = couple_legs
+        if self.couple_legs:
+            for leg in self.legs:
+                self.sim.createConstraint(parentBodyUniqueId=self.id, parentLinkIndex=leg[3],
+                                          childBodyUniqueId=self.id, childLinkIndex=leg[1],
+                                          jointType=self.sim.JOINT_POINT2POINT,
+                                          jointAxis=[0, 0, 0], parentFramePosition=[0, 0.005, 0.2],
+                                          childFramePosition=[0, 0.01, 0.2])
+
+        # disable motors
+        self.disableMotor(self.knees)
+
+        # kp, kd gains
+        self.kp = 1.
+        self.kd = 1.
+        self.max_force = 3.5
+
+        # set feet friction
+        self.setFootFriction(friction=foot_friction, feet_id=self.feet)
+
+        # set joint angles to home position
+        self.setJointHomePositions()
+
+    ##############
+    # Properties #
+    ##############
+
+    @property
+    def hips(self):
+        hip_ids = []
+        for o, i in zip(self.outer_hips, self.inner_hips):
+            hip_ids.append(o)
+            hip_ids.append(i)
+        return hip_ids
+
+    @property
+    def knees(self):
+        knee_ids = []
+        for o, i in zip(self.outer_knees, self.inner_knees):
+            knee_ids.append(o)
+            knee_ids.append(i)
+        return knee_ids
+
+    @property
+    def left_front_outer_leg(self):
+        return self.left_front_leg[:2]
+
+    @property
+    def left_front_inner_leg(self):
+        return self.left_front_leg[2:]
+
+    @property
+    def right_front_outer_leg(self):
+        return self.right_front_leg[2:]
+
+    @property
+    def right_front_inner_leg(self):
+        return self.right_front_leg[:2]
+
+    @property
+    def left_back_outer_leg(self):
+        return self.left_back_leg[:2]
+
+    @property
+    def left_back_inner_leg(self):
+        return self.left_back_leg[2:]
+
+    @property
+    def right_back_outer_leg(self):
+        return self.right_back_leg[2:]
+
+    @property
+    def right_back_inner_leg(self):
+        return self.right_back_leg[:2]
+
+    ###########
+    # Methods #
+    ###########
+
+    def getHomeJointPositions(self):
+        """Return the joint positions for the home position"""
+        h = np.pi/2  # hip angle from [2]
+        k = 2.1834  # knee angle from [2]
+        # joint positions
+        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)
+        return np.array(right_front_leg_initial_pos + right_back_leg_initial_pos + left_front_leg_initial_pos +
+                        left_back_leg_initial_pos)
+
+    def setJointPositions(self, position, jointId=None, kp=None, kd=None, velocity=None, maxTorque=None):
+        if self.couple_legs:  # assume the joint ids are for the outer legs
+
+            if jointId is None:
+                pass
+
+            # if the given joint is just one id
+            elif isinstance(jointId, int):
+                if jointId not in self.outer_joints:
+                    raise ValueError("Expecting the jointId to be an outer joint as the legs of the minitaur are "
+                                     "coupled")
+                jointId = [jointId, jointId + 3]
+                if isinstance(position, collections.Iterable):
+                    position = position[0]
+
+                position = np.array([position, -position])
+                position += self.init_joint_positions[self.getQIndex(jointId)]
+
+            # if multiple joint ids
+            elif isinstance(jointId, collections.Iterable):
+                # for each outer joint id, get the corresponding inner joint id
+                joints = []
+                for joint in jointId:
+                    if joint not in self.outer_joints:
+                        raise ValueError("One of the jointId is not an outer joint which is a problem as the legs of "
+                                         "the minitaur are coupled")
+                    joints.append(joint+3)
+
+                # increase the list of joint ids to take into account inner joint ids
+                jointId = list(jointId) + joints
+
+                # compute the positions (offset original position, and compute positions for inner joints)
+                position = list(position) + list(-position)
+                position = np.array(position) + self.init_joint_positions[self.getQIndex(jointId)]
+
+            else:
+                raise TypeError("Unknown type of for jointId; expecting a list of int, or an int, got instead :"
+                                "{}".format(type(jointId)))
+        super(Minitaur, self).setJointPositions(position, jointId=jointId, kp=kp, kd=kd, velocity=velocity,
+                                                maxTorque=maxTorque)
+
 
 # Test
 if __name__ == "__main__":
@@ -65,17 +220,23 @@ if __name__ == "__main__":
     world = BasicWorld(sim)
 
     # create robot
-    robot = Minitaur(sim)
+    robot = Minitaur(sim, couple_legs=True)
 
     # print information about the robot
     robot.printRobotInfo()
+    print("Robot leg ids: {}".format(robot.legs))
+    print("Robot feet ids: {}".format(robot.feet))
 
     # Position control using sliders
-    robot.addJointSlider(robot.getLeftFrontLegIds())
+    # robot.addJointSlider(robot.getLeftFrontLegIds())
 
+    t = 0
     # run simulator
     for _ in count():
-        robot.updateJointSlider()
+        t += 0.01
+        position = np.pi/4 * np.sin(2 * np.pi * t) * np.ones(len(robot.outer_hips))
+        robot.setJointPositions(position, robot.outer_hips)
+        # robot.updateJointSlider()
         # robot.computeAndDrawCoMPosition()
         # robot.computeAndDrawProjectedCoMPosition()
         world.step(sleep_dt=1./240)

pyrobolearn/robots/robot.py

@@ -37,7 +37,7 @@ class Robot(object):
     """
 
     def __init__(self, simulator, urdf_path, init_pos=(0, 0, 1.5), init_orient=(0, 0, 0, 1),
-                 useFixedBase=False, scaling=1.):
+                 useFixedBase=False, scaling=1., *args, **kwargs):
         """
         Initialize the robot.
 
@@ -57,6 +57,9 @@ class Robot(object):
         if useFixedBase is None:
             useFixedBase = False
 
+        self.init_position = init_pos
+        self.init_orientation = init_orient
+
         # set the simulator
         self.sim = simulator
         # self.name = urdf_path.split('/')[-1].split('.urdf')[0]
@@ -993,40 +996,62 @@ class Robot(object):
         return torque * velocity
 
     # TODO: desVel, maxVel, and maxTorque
-    def setJointPositions(self, position, jointId=None, kp=None, kd=None, maxVelocity=True, maxTorque=True):
+    def setJointPositions(self, position, jointId=None, kp=None, kd=None, velocity=None, maxTorque=None):
         """
         Set the position of the given joint(s) (using position control).
 
         Args:
             jointId (int, int[N], None): joint id, or list of joint ids. If None, get all the actuated joints.
             position (float, float[N]): desired position, or list of desired positions [rad]
-            velocity (None, float, float[N]): desired velocity, or list of desired velocities [rad/s]
-            kp (None, float, float[N]): position gain(s)
-            kd (None, float, float[N]): velocity gain(s)
-            maxVelocity (bool): NOT CURRENTLY USED
-            maxTorque (bool, float, float[N]): maximum motor torques
+            velocity (float, float[N], None): desired velocity, or list of desired velocities [rad/s]
+            kp (float, float[N], None): position gain(s)
+            kd (float, float[N], None): velocity gain(s)
+            maxTorque (float, float[N], None): maximum motor torques
         """
         if isinstance(jointId, int):
-            if kp is not None and kd is not None:
-                self.sim.setJointMotorControl2(self.id, jointId, self.sim.POSITION_CONTROL, targetPosition=position,
-                                               positionGain=kp, velocityGain=kd)
-            else:
-                self.sim.setJointMotorControl2(self.id, jointId, self.sim.POSITION_CONTROL, targetPosition=position)
+            kwargs = {}
+            if kp is not None:
+                kwargs['positionGain'] = kp
+            if kd is not None:
+                kwargs['velocityGain'] = kd
+            if velocity is not None:
+                kwargs['targetVelocity'] = velocity
+            if maxTorque is not None:
+                kwargs['force'] = maxTorque
+            self.sim.setJointMotorControl2(self.id, jointId, self.sim.POSITION_CONTROL, targetPosition=position,
+                                           **kwargs)
         else:
             if jointId is None:
                 jointId = self.joints
-            if kp is not None and kd is not None:
+            kwargs = {}
+            if kp is not None:
                 if isinstance(kp, (float, int)):
                     kp = kp * np.ones(len(jointId))
+                kwargs['positionGains'] = kp
+            if kd is not None:
                 if isinstance(kd, (float, int)):
                     kd = kd * np.ones(len(jointId))
-                qIdx = self.getQIndex(jointId)
-                position = np.clip(position, self.joint_limits[qIdx, 0], self.joint_limits[qIdx, 1])
-                self.sim.setJointMotorControlArray(self.id, jointId, self.sim.POSITION_CONTROL,
-                                                   targetPositions=position, positionGains=kp, velocityGains=kd)
-            else:
-                self.sim.setJointMotorControlArray(self.id, jointId, self.sim.POSITION_CONTROL,
-                                                   targetPositions=position)
+                kwargs['velocityGains'] = kd
+            # qIdx = self.getQIndex(jointId)
+            # print("pos: ", position)
+            # print(self.joint_limits[qIdx, 0], self.joint_limits[qIdx, 1])
+            # TODO: the following clip causes an error... Check Minitaur...
+            # position = np.clip(position, self.joint_limits[qIdx, 0], self.joint_limits[qIdx, 1])
+            # kp = kp.tolist()
+            # kd = kd.tolist()
+            # print("pos: ", position)
+            # print("kp: ", kp)
+            # print("kd: ", kd)
+            if velocity is not None:
+                if isinstance(velocity, (float, int)):
+                    velocity = velocity * np.ones(len(jointId))
+                kwargs['targetVelocities'] = velocity
+            if maxTorque is not None:
+                if isinstance(maxTorque, (float, int)):
+                    maxTorque = maxTorque * np.ones(len(jointId))
+                kwargs['forces'] = maxTorque
+            self.sim.setJointMotorControlArray(self.id, jointId, self.sim.POSITION_CONTROL, targetPositions=position,
+                                               **kwargs)
 
     # TODO: maxVel and maxTorque
     def setJointVelocities(self, velocity, jointId=None, maxVelocity=True, maxTorque=True):
@@ -1173,10 +1198,14 @@ class Robot(object):
         # check jointIds
         if not jointIds:
             jointIds = self.joints
+        if isinstance(jointIds, int):
+            jointIds = [jointIds]
 
         # check q
         if q is None:
             q = np.zeros(len(jointIds))
+        elif isinstance(q, (int, float)):
+            q = [q]
         else:
             if len(q) != len(jointIds):
                 raise ValueError("The number of joint ids does not match up with the number of q's")
@@ -1184,6 +1213,8 @@ class Robot(object):
         # check dq
         if dq is None:
             dq = np.zeros(len(jointIds))
+        elif isinstance(dq, (int, float)):
+            dq = [dq]
         else:
             if len(dq) != len(jointIds):
                 raise ValueError("The number of joint ids does not match with the number of dq's")

pyrobolearn/states/time_states.py

@@ -90,6 +90,7 @@ class PhaseState(TimeState):
 
     def __init__(self, num_steps=100, start=0, end=1., rate=1):
         self.cnt = 0
+        self.num_steps = num_steps
         self.rate = rate
         self.end_value = end
         self.start_value = start
@@ -169,6 +170,27 @@ class ExponentialPhaseState(TimeState):
 DecayPhaseState = ExponentialPhaseState
 
 
+class RhythmicPhase(PhaseState):
+    r"""Rhythmic Phase state
+
+    The PhaseState starts from `start` and ends with `end`, calling after `num_steps` will just return `end`.
+    In this class, we cycle through the phase `[start, end[`; that is, once the end is reached it restarts
+    automatically from `start`
+    """
+
+    def __init__(self, num_steps=100, start=0, end=1., rate=1):
+        super(RhythmicPhase, self).__init__(num_steps=num_steps, start=start, end=end, rate=rate)
+
+    def _read(self):
+        if (self.cnt % self.rate) == 0:
+            self.data = self._data + self.dphase
+            if self.sign > 0 and self._data[0] >= self.end_value:
+                self.data = np.array([self.start_value])
+            if self.sign < 0 and self._data[0] <= self.end_value:
+                self.data = np.array([self.start_value])
+        self.cnt += 1
+
+
 # Tests the different time states
 if __name__ == '__main__':
     s = AbsoluteTimeState()
@@ -202,6 +224,12 @@ if __name__ == '__main__':
     for i in range(200):
         print(s())
 
+    s = RhythmicPhase(num_steps=10, start=1, end=2)
+    print("\nPhase Time State:")
+    print(s.reset())
+    for i in range(100):
+        print(s())
+
     combined = AbsoluteTimeState() + RelativeTimeState() + CumulativeTimeState()
     fused = AbsoluteTimeState() & RelativeTimeState() & CumulativeTimeState()
 

pyrobolearn/worlds/world.py

@@ -409,12 +409,11 @@ class World(object):
         [2] Open3D: http://www.open3d.org/
     """
 
-    def __init__(self, simulator, set_gravity=True):
+    def __init__(self, simulator, gravity=(0., 0., -9.81)):
         self.sim = simulator
         # self.sim.setAdditionalSearchPath(pybullet_data.getDataPath())
 
         self.robots = {}
-        self.robot_init_states = {}
         self.movable_objects = {} # set()
         self.immovable_objects = {} # set()
         self.visual_objects = {} # set()
@@ -426,8 +425,7 @@ class World(object):
         self.quaternion_converter = QuaternionListConverter(convention=1)
 
         # By default, set the gravity
-        if set_gravity:
-            self.setGravity()
+        self.setGravity(gravity)
 
         self.worldState = None
 
@@ -564,6 +562,13 @@ class World(object):
             # reset simulation: remove all objects from the world and reset the world to initial conditions
             self.sim.resetSimulation()
 
+        # print("\nRobot reset state:")
+        # for robot_id, robot in self.robots.items():
+        #     print("Robot base position and orientation: {}".format(robot.getBasePositionAndOrientation()))
+        #     print("Robot base velocities: {}".format(robot.getBaseVelocity()))
+        #     print("Robot joint positions: {}".format(robot.getJointPositions()))
+        #     print("Robot joint velocities: {}".format(robot.getJointVelocities()))
+
     def resetSimulator(self):
         """
         Reset the simulator.
@@ -582,14 +587,14 @@ class World(object):
         if sleep_dt is not None:
             time.sleep(sleep_dt)
 
-    def setGravity(self, xyz=(0.,0.,-9.81)):
+    def setGravity(self, xyz=(0., 0., -9.81)):
         """
         Set the given gravity.
 
         Args:
             xyz (float[3]): gravity (acceleration) along the 3 axis.
         """
-        x,y,z = xyz
+        x, y, z = xyz
         self.sim.setGravity(x, y, z)
 
     def getMainCamera(self):
@@ -598,7 +603,7 @@ class World(object):
         """
         return WorldCamera(self.sim)
 
-    def loadRobot(self, robot, position=None, orientation=None, useFixedBase=None):
+    def loadRobot(self, robot, position=None, orientation=None, useFixedBase=None, *args, **kwargs):
         """
         Load the robot into the world. If the robot parameter is a known robot name or the path to the urdf file,
         it will create a `Robot` instance and return it. If the robot is already an instance of `Robot` it will
@@ -622,34 +627,35 @@ class World(object):
 
             if robot in robot_names_to_classes:
                 robot_class = robot_names_to_classes[robot]
-                robot = robot_class(self.sim, init_pos=position, init_orient=orientation, useFixedBase=useFixedBase)
+                robot = robot_class(self.sim, init_pos=position, init_orient=orientation, useFixedBase=useFixedBase,
+                                    *args, **kwargs)
 
             else:  # robot is the path to the urdf
                 robot = Robot(self.sim, urdf_path=robot, init_pos=position, init_orient=orientation,
-                              useFixedBase=useFixedBase)
+                              useFixedBase=useFixedBase, *args, **kwargs)
 
         elif isClass(robot):  # robot class
-            robot = robot(self.sim, init_pos=position, init_orient=orientation, useFixedBase=useFixedBase)
+            robot = robot(self.sim, init_pos=position, init_orient=orientation, useFixedBase=useFixedBase,
+                          *args, **kwargs)
 
-        else: # unknown type
+        else:  # unknown type
             raise TypeError('Unknown type for robot: {}. It must be a string or '
                             'an instance of Robot'.format(type(robot)))
 
         self.robots[robot.id] = robot
-        self.robot_init_states[robot.id] = (robot.getJointPositions(), robot.getJointVelocities())
         return robot
 
-    def isRobotId(self, robotId):
+    def isRobotId(self, robot_id):
         """
         Check if the given id is a robot id.
 
         Args:
-            robotId (int): the possible robot id
+            robot_id (int): the possible robot id
 
         Returns:
             bool: True if the id is a robot id, False otherwise
         """
-        return robotId in self.robots
+        return robot_id in self.robots
 
     def getRobot(self, robotId):
         """
@@ -668,12 +674,18 @@ class World(object):
 
     def resetRobots(self):
         """
-        Reset the joint states of each robot
+        Reset the base and joint states of each robot
         """
-        for robotId, robot in self.robots.items():
-            pos, vel = self.robot_init_states[robotId]
-            for jointId, p, v in zip(robot.joints, pos, vel):
-                self.sim.resetJointState(robotId, jointId, p, v)
+        for robot_id, robot in self.robots.items():
+            # reset base
+            self.sim.resetBasePositionAndOrientation(robot_id, robot.init_position, robot.init_orientation)
+            self.sim.resetBaseVelocity(robot_id, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 0])
+
+            # reset joint positions
+            positions = robot.init_joint_positions
+            velocities = np.zeros(len(positions))
+            for joint_id, position, velocity in zip(robot.joints, positions, velocities):
+                self.sim.resetJointState(robot_id, joint_id, position, velocity)
 
     def loadURDF(self, filename, position, orientation, useFixedBase=False, scaling=1., objectName=None):
         """
@@ -1742,13 +1754,13 @@ class World(object):
     def distribute_objects(self, distributor, objects):
         pass
 
-    def getDynamicsInfo(self, bodyId, linkId):
+    def getDynamicsInfo(self, bodyId, linkId=-1):
         """
         Return the dynamics information about objects that are in the world.
 
         Args:
-            bodyId: object unique id
-            linkId: link index (or -1 for the base)
+            bodyId (int): object unique id.
+            linkId (int): link index (or -1 for the base).
 
         Returns:
             float: mass in kg
@@ -1769,11 +1781,32 @@ class World(object):
         local_inertial_orn = np.array(local_inertial_orn)
         return info[:2] + [local_inertia_diag, local_inertial_pos, local_inertial_orn] + info[5:]
 
-    def changeDynamics(self, lateralFriction, spinningFriction, rollingFriction, linearDamping, angularDamping,
-                       contactStiffness=-1, contactDamping=-1):
-        self.sim.changeDynamics(bodyUniqueId=self.floor_id, linkIndex=-1, lateralFriction=lateralFriction,
-                                spinningFriction=spinningFriction, rollingFriction=rollingFriction,
-                                linearDamping=linearDamping, angularDamping=angularDamping)
+    def print_dynamics_info(self, body_id, link_id=-1):
+        """
+        Print the dynamics information related to the given body id and link id.
+
+        Args:
+            bodyId (int): object unique id.
+            linkId (int): link index (or -1 for the base).
+        """
+        info = self.sim.getDynamicsInfo(body_id, link_id)
+        print("Mass: {}".format(info[0]))
+        print("Lateral friction coefficient: {}".format(info[1]))
+        print("Local inertia diagonal: {}".format(info[2]))
+        print("Local inertial position: {}".format(info[3]))
+        print("Local inertial orientation (quat=[x,y,z,w]): {}".format(info[4]))
+        print("Restitution coefficient (bouncyness): {}".format(info[5]))
+        print("Rolling friction coefficient: {}".format(info[6]))
+        print("Spinning friction coefficient: {}".format(info[7]))
+        print("Contact damping coefficient (-1 if not available): {}".format(info[8]))
+        print("Contact stiffness coefficient (-1 if not available): {}".format(info[9]))
+
+    def changeDynamics(self, lateral_friction, spinning_friction, rolling_friction, linear_damping, angular_damping,
+                       contact_stiffness=-1, contact_damping=-1):
+        self.sim.changeDynamics(bodyUniqueId=self.floor_id, linkIndex=-1, lateralFriction=lateral_friction,
+                                spinningFriction=spinning_friction, rollingFriction=rolling_friction,
+                                linearDamping=linear_damping, angularDamping=angular_damping,
+                                contactStiffness=contact_stiffness, contactDamping=contact_damping)
 
 
 class BasicWorld(World):
@@ -1782,20 +1815,20 @@ class BasicWorld(World):
     It creates a basic world with a floor and set the gravity.
     """
 
-    def __init__(self, simulator, floor_path=None, set_gravity=True, scaling=1., lateralFriction=.9,
-                 spinningFriction=0., rollingFriction=0., contactStiffness=-1, contactDamping=-1):
-        super(BasicWorld, self).__init__(simulator, set_gravity=set_gravity)
+    def __init__(self, simulator, floor_path=None, gravity=(0., 0., -9.81), scaling=1., lateral_friction=.9,
+                 spinning_friction=0., rolling_friction=0., contact_stiffness=-1, contact_damping=-1):
+        super(BasicWorld, self).__init__(simulator, gravity=gravity)
 
         if floor_path is None:
             self.loadFloor(scaling=scaling)
-            self.changeDynamics(lateralFriction=lateralFriction, spinningFriction=spinningFriction,
-                                rollingFriction=rollingFriction, linearDamping=0, angularDamping=0,
-                                contactStiffness=30000)
+            self.changeDynamics(lateral_friction=lateral_friction, spinning_friction=spinning_friction,
+                                rolling_friction=rolling_friction, linear_damping=0, angular_damping=0,
+                                contact_stiffness=contact_stiffness, contact_damping=contact_damping)
             self.simulator.setDefaultContactERP(0.9)
         else:
             self.loadTerrain(floor_path, replace_floor=True)
 
-        print(self.sim.getDynamicsInfo(self.floor_id, -1))
+        self.print_dynamics_info(self.floor_id)
 
 
 class RobotPartyWorld(BasicWorld):