From 32abb04ab2858726b15869f87f44a1e88af80434 Mon Sep 17 00:00:00 2001 From: Brian Delhaisse Date: Wed, 11 Sep 2019 11:43:54 +0200 Subject: [PATCH] update simulators --- pyrobolearn/simulators/bullet.py | 9 +- pyrobolearn/simulators/mujoco.py | 1650 +++++++++++++++-- pyrobolearn/simulators/simulator.py | 34 +- .../utils/parsers/robots/data_structures.py | 58 +- 4 files changed, 1606 insertions(+), 145 deletions(-) diff --git a/pyrobolearn/simulators/bullet.py b/pyrobolearn/simulators/bullet.py index 7c5aa0c..7aae413 100644 --- a/pyrobolearn/simulators/bullet.py +++ b/pyrobolearn/simulators/bullet.py @@ -91,8 +91,8 @@ class Bullet(Simulator): Args: render (bool): if True, it will open the GUI, otherwise, it will just run the server. num_instances (int): number of simulator instances. - **kwargs (dict): optional arguments (this is not used here). middleware (MiddleWare, None): middleware instance. + **kwargs (dict): optional arguments (this is not used here). """ # try to import the pybullet library # normally that should be done outside the class but because it might have some conflicts with other libraries @@ -1616,7 +1616,7 @@ class Bullet(Simulator): states[idx][2] = np.asarray(state[2]) return states - def reset_joint_state(self, body_id, joint_id, position, velocity=0.): + def reset_joint_state(self, body_id, joint_id, position, velocity=None): """ Reset the state of the joint. It is best only to do this at the start, while not running the simulation: `reset_joint_state` overrides all physics simulation. Note that we only support 1-DOF motorized joints at @@ -1628,7 +1628,10 @@ class Bullet(Simulator): position (float): the joint position (angle in radians [rad] or position [m]) velocity (float): the joint velocity (angular [rad/s] or linear velocity [m/s]) """ - self.sim.resetJointState(body_id, joint_id, position, velocity) + if velocity is None: + self.sim.resetJointState(body_id, joint_id, position) + else: + self.sim.resetJointState(body_id, joint_id, position, velocity) def enable_joint_force_torque_sensor(self, body_id, joint_ids, enable=True): """ diff --git a/pyrobolearn/simulators/mujoco.py b/pyrobolearn/simulators/mujoco.py index 0b917da..79d192e 100644 --- a/pyrobolearn/simulators/mujoco.py +++ b/pyrobolearn/simulators/mujoco.py @@ -76,7 +76,7 @@ import glfw from pyrobolearn.simulators.simulator import Simulator from pyrobolearn.utils.parsers.robots import URDFParser, MuJoCoParser, SDFParser import pyrobolearn.utils.parsers.robots.data_structures as struct -from pyrobolearn.utils.transformation import get_homogeneous_matrix +from pyrobolearn.utils.transformation import get_homogeneous_matrix, get_quaternion_from_matrix # check Python version @@ -110,42 +110,10 @@ class Texture(object): self.material = material -# class Body(object): -# """Body.""" -# -# def __init__(self, body_id, body_tag): -# """ -# Initialize the Body. -# -# Args: -# body_id (int): unique body id. -# body_tag (xml.etree.ElementTree.Element): body tag element in the xml file. -# """ -# self.id = body_id -# if not isinstance(body_tag, ET.Element): -# raise TypeError("Expecting the given 'body_tag' to be an instance of `ET.Element`, but got instead: " -# "{}".format(type(body_tag))) -# self.body = body_tag -# -# self.q_start = 0 # starting index in the whole state -# self.q_end = 0 # end index in the whole state -# self.fixed_base = False -# -# # list of inner bodies (=links) -# self.bodies = [] -# self.joints = [] -# self.joint = None -# -# @property -# def name(self): -# """Return the body name.""" -# return self.body.attrib.get("name") - - class Body(object): """Body.""" - def __init__(self, body_id, body_tag, body=None, fixed_base=False): + def __init__(self, body_id, body_tag, body): """ Initialize the MultiBody. @@ -153,24 +121,180 @@ class Body(object): body_id (int): unique body id in the Mujoco model. body_tag (ET.Element): body tag element in the XML file. body (MultiBody, None): multibody data structure. - fixed_base (bool): if True, the body is fixed in the world. """ self.id = body_id self.tag = body_tag - self.q_start = 0 # starting index in the whole state - self.q_end = 0 # end index in the whole state - self.fixed_base = fixed_base + # check given body + if body is None: + body = struct.MultiBody() + if not isinstance(body, struct.MultiBody): + raise TypeError("Expecting the given 'body' to be an instance of `MultiBody` but got instead: " + "{}".format(type(body))) - if body is not None: - self.num_bodies = body.num_bodies # number of links - self.num_joints = body.num_joints # number of joints (including fixed joints) - self.num_actuated_joints = body.num_actuated_joints # number of actuated joints + # get information from body + self.num_bodies = body.num_bodies # number of links/bodies + self.num_joints = body.num_joints # nb of joints (include fixed joints but exclude free joints) = num_bodies + self.num_actuated_joints = body.num_actuated_joints # nb of actuated joints (exclude fixed and free joints) + self.num_free_joints = body.num_free_joints # nb of free joints (including free but excluding fixed joints) + self.num_dofs = body.num_dofs # nb of DoFs + self.q_length = self.num_dofs # length of q + self.fixed = body.fixed_base if body.fixed_base is not None else True + if not self.fixed: # if free joint, add 1 because in Mujoco the pose is represented as position vector + self.q_length += 1 # (3) + quaternion (4) = 7, so one more than 6 DoFs + + # define variables for indices that appears in the various vectors and matrices returned by mjModel and mjData + self._q_idx0, self._q_idxf = 0, 0 # initial and final q indices + self._b_idx0, self._b_idxf = 0, 0 # initial and final body indices + self._j_idx0, self._j_idxf = 0, 0 # initial and final free joint indices + self._v_idx0, self._v_idxf = 0, 0 # initial and final dq (velocity) indices + + # keep in memory the body + # self.body = body + self.joints = list(body.joints.values()) + self.links = list(body.bodies.values()) + + # compute mapping from joint ids to q indices + idx, jnt_to_q = 0, [] + for joint in body.joints.values(): + if joint.dtype == 'fixed': + jnt_to_q.append(-1) + else: + jnt_to_q.append(idx) + idx += 1 + self.jnt_to_q = np.array(jnt_to_q) @property - def num_dofs(self): - """Return the number of DoFs.""" - return self.q_end - self.q_start + def num_links(self): + """Alias to `num_bodies`.""" + return self.num_bodies + + @property + def q_idx0(self): + """Return the initial q index.""" + return self._q_idx0 + + @q_idx0.setter + def q_idx0(self, q): + """Set the the initial q index.""" + q = int(q) + if q < 0: + raise ValueError("Error while setting the initial q index, this index has to be bigger than 0!") + self._q_idx0 = q + self._q_idxf = q + self.q_length # set the final q index + + @property + def q_idxf(self): + """Return the final q index.""" + return self._q_idxf + + @q_idxf.setter + def q_idxf(self, q): + """Set the final q index.""" + q = int(q) + if q < 0: + raise ValueError("Error while setting the final q index, this index has to be bigger than 0!") + self._q_idxf = q + self._q_idx0 = q - self.q_length # set initial q index + if self._q_idx0 < 0: + raise ValueError("Error while setting the final q index, by computing automatically the initial q index " + "from it, it appears it is smaller than 0. The initial q index has to be bigger than 0!") + + @property + def b_idx0(self): + """Return the initial body (link) index.""" + return self._b_idx0 + + @b_idx0.setter + def b_idx0(self, b): + """Set the initial body (link) index.""" + b = int(b) + if b < 0: + raise ValueError("Error while setting the initial body index, this index has to be bigger than 0!") + self._b_idx0 = b + self._b_idxf = b + self.num_bodies # set the final body index + + @property + def b_idxf(self): + """Return the final body (link) index.""" + return self._b_idxf + + @b_idxf.setter + def b_idxf(self, b): + """Set the final body (link) index.""" + b = int(b) + if b < 0: + raise ValueError("Error while setting the final body index, this index has to be bigger than 0!") + self._b_idxf = b + self._b_idx0 = b - self.num_bodies # set initial body index + if self._b_idx0 < 0: + raise ValueError("Error while setting the final body index, by computing automatically the initial body " + "index from it, it appears it is smaller than 0. The initial body index has to be bigger " + "than 0!") + + @property + def j_idx0(self): + """Return the initial free joint index.""" + return self._j_idx0 + + @j_idx0.setter + def j_idx0(self, j): + """Set the initial free joint index.""" + j = int(j) + if j < 0: + raise ValueError("Error while setting the initial free joint index, this index has to be bigger than 0!") + self._j_idx0 = j + self._j_idxf = j + self.num_free_joints # set the final free joint index + + @property + def j_idxf(self): + """Return the initial free joint index.""" + return self._j_idxf + + @j_idxf.setter + def j_idxf(self, j): + """Set the final free joint index.""" + j = int(j) + if j < 0: + raise ValueError("Error while setting the final free joint index, this index has to be bigger than 0!") + self._j_idxf = j + self._j_idx0 = j - self.num_free_joints # set initial free joint index + if self._j_idx0 < 0: + raise ValueError("Error while setting the final joint index, by computing automatically the initial joint " + "index from it, it appears it is smaller than 0. The initial joint index has to be bigger " + "than 0!") + + @property + def v_idx0(self): + """Return the initial velocity index.""" + return self._v_idx0 + + @v_idx0.setter + def v_idx0(self, v): + """Set the initial velocity index.""" + v = int(v) + if v < 0: + raise ValueError("Error while setting the initial velocity index, this index has to be bigger than 0!") + self._v_idx0 = v + self._v_idxf = v + self.num_dofs # set the final velocity index + + @property + def v_idxf(self): + """Return the initial velocity index.""" + return self._v_idxf + + @v_idxf.setter + def v_idxf(self, v): + """Set the final velocity index.""" + v = int(v) + if v < 0: + raise ValueError("Error while setting the final velocity index, this index has to be bigger than 0!") + self._v_idxf = v + self._v_idx0 = v - self.num_dofs # set initial velocity index + if self._v_idx0 < 0: + raise ValueError("Error while setting the final velocity index, by computing automatically the initial " + "velocity index from it, it appears it is smaller than 0. The initial velocity index has " + "to be bigger than 0!") @property def name(self): @@ -183,6 +307,30 @@ class Body(object): """Return the body tag name.""" return self.tag.attrib.get("name") + def get_q_idx(self, joint_id, keep=False): + q = self.jnt_to_q[joint_id] + if keep: # keep fixed joints (-1) + return q + if isinstance(q, float): + if q!=-1: + return q + return [] + return q[q!=-1] # remove fixed joints + + def get_dq_idx(self, joint_id, keep=False): + return self.get_q_idx(joint_id, keep) + + def check_joint_id(self, joint_id): + if joint_id < 0 or joint_id > (self.num_joints - 1): + raise ValueError("joint_id should belong to [0, `num_joints-1`].") + + def check_link_id(self, link_id): + if link_id < -1 or link_id > (self.num_bodies - 2): # -1 is for the base + raise ValueError("link_id should belong to [-1, `num_links-2`].") + + def get_joint_type(self, joint_id): + return self.joints[joint_id].dtype + class Mujoco(Simulator): r"""Mujoco Simulator interface. @@ -260,7 +408,9 @@ class Mujoco(Simulator): # counters for Mujoco self._q_cnt = 0 - self._link_cnt = 0 + self._dq_cnt = 0 + self._joint_cnt = 0 + self._link_cnt = 1 # this is the number of bodies (=links) in Mujoco, 0 is for the worldbody. self._mjc_body_id = 0 self.default_timestep = 0.002 @@ -592,7 +742,7 @@ class Mujoco(Simulator): body.add_parent_joint(joint) tree.add_joint(joint, idx=0) - return self._create_body(tree, fixed_base=use_fixed_base, verbose=1) + return self._create_body(tree, verbose=2) def load_sdf(self, filename, scaling=1., *args, **kwargs): # TODO """Load a SDF file in the simulator. @@ -705,7 +855,7 @@ class Mujoco(Simulator): floor = self._parser.add_element(name="geom", parent_element=self._worldbody, attributes={"type": "plane", "size": str(dim) + " " + str(dim) + " 1."}) - body = Body(body_id=0, body_tag=floor, fixed_base=True) # 0 is only for the floor + body = Body(body_id=0, body_tag=floor, body=None) # 0 is only for the floor self._bodies[0] = body # the floor is always the first body in the ordered dict self._floor_id = 0 @@ -722,7 +872,7 @@ class Mujoco(Simulator): # notify that the Mujoco model has changed (this will be checked in the `step` method) self._model_changed = True - def _create_body(self, tree, body_id=None, fixed_base=False, verbose=1): + def _create_body(self, tree, body_id=None, verbose=2): """ Create inner body in the simulator; given the tree data structure, it generates all the XML tags using the inner parser/generator, wraps the returned tree XML tag to a Body structure, sets its attributes (such as @@ -731,7 +881,6 @@ class Mujoco(Simulator): Args: tree (struct.MultiBody): the tree / multi-body data structure. body_id (int, None): unique body id to save the wrapped body. - fixed_base (bool): if the base of the multibody/tree is fixed or not. verbose (bool, int): if True, it will print the current Mujoco XML file that we have which is useful for debug. If int, it represents the level of verbosity. @@ -748,21 +897,25 @@ class Mujoco(Simulator): # save body self._mjc_body_id += 1 - body = Body(body_id=self._mjc_body_id, body_tag=tree_tag, body=tree, fixed_base=fixed_base) + body = Body(body_id=self._mjc_body_id, body_tag=tree_tag, body=tree) self._bodies[body_id] = body - num_dofs = tree.num_dofs if verbose > 0: print("\nNum DoFs: {}".format(tree.num_dofs)) + print("Length of q: {}".format(body.q_length)) print("Num links/bodies: {}".format(tree.num_bodies)) print("Num joints: {}".format(tree.num_joints)) print("Num actuated joints: {}\n".format(tree.num_actuated_joints)) - # set generalized coordinates counter - body.q_start = self._q_cnt - num_dofs = num_dofs if fixed_base else num_dofs + 7 - self._q_cnt += num_dofs - body.q_end = self._q_cnt + # update indices + body.q_idx0 = self._q_cnt + self._q_cnt += body.q_length + body.b_idx0 = self._link_cnt + self._link_cnt += body.num_bodies + body.j_idx0 = self._joint_cnt + self._joint_cnt += body.num_free_joints + body.v_idx0 = self._dq_cnt + self._dq_cnt += body.num_dofs # update mujoco model if necessary self._update_sim() @@ -819,7 +972,7 @@ class Mujoco(Simulator): if not static: tree.add_joint(joint) - return self._create_body(tree, body_id=self._body_cnt, fixed_base=static, verbose=1) + return self._create_body(tree, body_id=self._body_cnt, verbose=2) def remove_body(self, body_id): # DONE """Remove a particular body in the simulator. @@ -828,19 +981,24 @@ class Mujoco(Simulator): body_id (int): unique body id. """ # remove body from the bodies - # This is an O(N) operation because I have to modify the id, q_start, and q_end of the bodies that appears + # This is an O(N) operation because I have to modify the id, q_idx0, and q_idxf of the bodies that appears # after the given body body_to_remove = self._bodies[body_id] - found, num_dofs = False, 0 + found, q_length, num_links, num_free_joints = False, 0, 0, 0 for body in self._bodies: if body_to_remove == body: found = True - num_dofs = body.num_dofs - if found: # for the bodies after body_to_remove, shift their id and q_start/q_end to the left - body.id -= 1 # shift id to the left - body.q_start -= num_dofs - body.q_end -= num_dofs - self._q_cnt = body.q_end + q_length = body.q_length + num_links = body.num_bodies + num_free_joints = body.num_free_joints + self._q_cnt -= q_length + self._joint_cnt -= num_free_joints + self._link_cnt -= num_links + if found: # for the bodies after body_to_remove, shift their indices to the left + body.id -= 1 + body.q_idx0 -= q_length + body.b_idx0 -= num_links + body.j_idx0 -= num_free_joints body = self._bodies.pop(body_id) self._mjc_body_id -= 1 @@ -1022,10 +1180,7 @@ class Mujoco(Simulator): float: total mass of the robot [kg] """ body = self._bodies[body_id] - mass = self.get_base_mass(body_id) - for b in body.bodies: - mass += sim.get_base_mass(b.id) - return mass + return sum(self.sim.model.body_mass[body.b_idx0:body.b_idxf]) def get_base_mass(self, body_id): """Return the base mass of the robot. @@ -1033,7 +1188,8 @@ class Mujoco(Simulator): Args: body_id (int): unique object id. """ - return self.sim.model.body_mass[body_id] + body = self._bodies[body_id] + return self.sim.model.body_mass[body.b_idx0] def get_base_name(self, body_id): """ @@ -1045,8 +1201,8 @@ class Mujoco(Simulator): Returns: str: base name """ - body_id = self._bodies[body_id].id - return self.sim.model.body_id2name(body_id) + body = self._bodies[body_id] + return self.sim.model.body_id2name(body.b_idx0) def get_center_of_mass_position(self, body_id, link_ids=None): # TODO """ @@ -1060,7 +1216,8 @@ class Mujoco(Simulator): Returns: np.array[float[3]]: center of mass position in the Cartesian world coordinates """ - return self.sim.data.subtree_com[body_id] + body = self._bodies[body_id] + return self.sim.data.subtree_com[body.b_idx0] # sim.data.body_xipos def get_center_of_mass_velocity(self, body_id, link_ids=None): # TODO """ @@ -1074,7 +1231,23 @@ class Mujoco(Simulator): Returns: np.array[float[3]]: center of mass linear velocity. """ - return self.sim.data.subtree_linvel[body_id] + body = self._bodies[body_id] + return self.sim.data.subtree_linvel[body.b_idx0] # sim.data.cvel[3:] + + def get_center_of_mass_angular_momentum(self, body_id, link_ids=None): # TODO + """ + Return the link angular momentum around its CoM. + + Args: + body_id (int): unique body id. + link_ids (list[int]): link ids associated with the given body id. If None, it will take all the links + of the specified body. + + Returns: + np.array[float[3]]: angular momentum. + """ + body = self._bodies[body_id] + return self.sim.data.subtree_angmom[body.b_idx0] # sim.data.cvel[:3] def get_base_pose(self, body_id): """ @@ -1088,14 +1261,10 @@ class Mujoco(Simulator): np.array[float[4]]: base orientation (quaternion [x,y,z,w]) """ # WARNING: body_xpos is one step late compared to qpos - # return self.sim.data.body_xpos[body_id], self._convert_wxyz_to_xyzw(self.sim.data.body_xquat[body_id]) - - # print(self.sim.data.body_xpos, q[body.q_start:body.q_start+3]) - # print(self.sim.data.body_xquat, q[body.q_start+3:body.q_end]) - body = self._bodies[body_id] - q = self.sim.data.qpos - return q[body.q_start:body.q_start+3], self._convert_wxyz_to_xyzw(q[body.q_start+3:body.q_start+7]) + position = self.sim.data.body_xpos[body.b_idx0] + orientation = self._convert_wxyz_to_xyzw(self.sim.data.body_xquat[body.b_idx0]) + return position, orientation def get_base_position(self, body_id): """ @@ -1107,10 +1276,8 @@ class Mujoco(Simulator): Returns: np.array[float[3]]: base position. """ - # return self.sim.data.body_xpos[body_id] body = self._bodies[body_id] - q = self.sim.data.qpos - return q[body.q_start:body.q_start + 3] + return self.sim.data.body_xpos[body.b_idx0] def get_base_orientation(self, body_id): """ @@ -1122,10 +1289,8 @@ class Mujoco(Simulator): Returns: np.array[float[4]]: base orientation in the form of a quaternion (x,y,z,w) """ - # return self._convert_wxyz_to_xyzw(self.sim.data.body_xquat[body_id]) body = self._bodies[body_id] - q = self.sim.data.qpos - return self._convert_wxyz_to_xyzw(q[body.q_start+3:body.q_start+7]) + return self._convert_wxyz_to_xyzw(self.sim.data.body_xquat[body.b_idx0]) def reset_base_pose(self, body_id, position, orientation): """ @@ -1141,9 +1306,12 @@ class Mujoco(Simulator): orientation (np.array[float[4]]): new base orientation (expressed as a quaternion [x,y,z,w]) """ body = self._bodies[body_id] - q = self.sim.data.qpos - q[body.q_start:body.q_start + 3] = position - q[body.q_start + 3:body.q_start + 7] = self._convert_xyzw_to_wxyz(orientation) + if body.fixed: + self.model.body_pos[body.b_idx0] = position + self.model.body_quat[body.b_idx0] = self._convert_xyzw_to_wxyz(orientation) + else: + self.sim.data.qpos[body.q_idx0:body.q_idx0 + 3] = position + self.sim.data.qpos[body.q_idx0 + 3:body.q_idx0 + 7] = self._convert_xyzw_to_wxyz(orientation) def reset_base_position(self, body_id, position): """ @@ -1153,11 +1321,11 @@ class Mujoco(Simulator): body_id (int): unique object id. position (np.array[float[3]]): new base position. """ - # self.sim.data.body_xpos[body_id] = position - # self.sim.forward() body = self._bodies[body_id] - q = self.sim.data.qpos - q[body.q_start:body.q_start + 3] = position + if body.fixed: # fixed base + self.model.body_pos[body.b_idx0] = position + else: # free joint + self.sim.data.qpos[body.q_idx0:body.q_idx0 + 3] = position def reset_base_orientation(self, body_id, orientation): """ @@ -1168,8 +1336,10 @@ class Mujoco(Simulator): orientation (np.array[float[4]]): new base orientation (expressed as a quaternion [x,y,z,w]) """ body = self._bodies[body_id] - q = self.sim.data.qpos - q[body.q_start+3:body.q_start+7] = self._convert_xyzw_to_wxyz(orientation) + if body.fixed: + self.model.body_quat[body.b_idx0] = self._convert_xyzw_to_wxyz(orientation) + else: + self.sim.data.qpos[body.q_idx0+3:body.q_idx0+7] = self._convert_xyzw_to_wxyz(orientation) def get_base_velocity(self, body_id): """ @@ -1183,8 +1353,10 @@ class Mujoco(Simulator): np.array[float[3]]: angular velocity of the base in Cartesian world space coordinates """ body = self._bodies[body_id] + if body.fixed: + return np.zeros(3), np.zeros(3) dq = self.sim.data.qvel - return dq[body.q_start:body.q_start+3], dq[body.q_start+3:body.q_start+6] + return dq[body.q_idx0:body.q_idx0+3], dq[body.q_idx0+3:body.q_idx0+6] def get_base_linear_velocity(self, body_id): """ @@ -1197,8 +1369,9 @@ class Mujoco(Simulator): np.array[float[3]]: linear velocity of the base in Cartesian world space coordinates """ body = self._bodies[body_id] - dq = self.sim.data.qvel - return dq[body.q_start:body.q_start + 3] + if body.fixed: + return np.zeros(3) + return self.sim.data.qvel[body.q_idx0:body.q_idx0 + 3] def get_base_angular_velocity(self, body_id): """ @@ -1211,8 +1384,9 @@ class Mujoco(Simulator): np.array[float[3]]: angular velocity of the base in Cartesian world space coordinates """ body = self._bodies[body_id] - dq = self.sim.data.qvel - return dq[body.q_start+3:body.q_start+6] + if body.fixed: + return np.zeros(3) + return self.sim.data.qvel[body.q_idx0+3:body.q_idx0+6] def reset_base_velocity(self, body_id, linear_velocity=None, angular_velocity=None): """ @@ -1224,9 +1398,10 @@ class Mujoco(Simulator): angular_velocity (np.array[float[3]]): new angular velocity of the base. """ body = self._bodies[body_id] - dq = self.sim.data.qvel - dq[body.q_start:body.q_start + 3] = linear_velocity - dq[body.q_start + 3:body.q_start + 6] = angular_velocity + if not body.fixed: + dq = self.sim.data.qvel + dq[body.q_idx0:body.q_idx0 + 3] = linear_velocity + dq[body.q_idx0 + 3:body.q_idx0 + 6] = angular_velocity def reset_base_linear_velocity(self, body_id, linear_velocity): """ @@ -1237,8 +1412,8 @@ class Mujoco(Simulator): linear_velocity (np.array[float[3]]): new linear velocity of the base """ body = self._bodies[body_id] - dq = self.sim.data.qvel - dq[body.q_start:body.q_start + 3] = linear_velocity + if not body.fixed: + self.sim.data.qvel[body.q_idx0:body.q_idx0 + 3] = linear_velocity def reset_base_angular_velocity(self, body_id, angular_velocity): """ @@ -1249,8 +1424,8 @@ class Mujoco(Simulator): angular_velocity (np.array[float[3]]): new angular velocity of the base """ body = self._bodies[body_id] - dq = self.sim.data.qvel - dq[body.q_start + 3:body.q_start + 6] = angular_velocity + if not body.fixed: + self.sim.data.qvel[body.q_idx0 + 3:body.q_idx0 + 6] = angular_velocity def get_base_acceleration(self, body_id): """ @@ -1264,10 +1439,13 @@ class Mujoco(Simulator): np.array[float[3]]: angular acceleration [rad/s^2] """ body = self._bodies[body_id] + if body.fixed: + return np.zeros(3), np.zeros(3) ddq = self.sim.data.qacc - return ddq[body.q_start:body.q_start+3], ddq[body.q_start+3:body.q_start+6] + return ddq[body.q_idx0:body.q_idx0+3], ddq[body.q_idx0+3:body.q_idx0+6] - def apply_external_force(self, body_id, link_id=-1, force=(0., 0., 0.), position=(0., 0., 0.), frame=1): + def apply_external_force(self, body_id, link_id=-1, force=(0., 0., 0.), position=(0., 0., 0.), + frame=Simulator.LINK_FRAME): """ Apply the specified external force on the specified position on the body / link. @@ -1280,9 +1458,27 @@ class Mujoco(Simulator): frame (int): if frame = 1, then the force / position is described in the link frame. If frame = 2, they are described in the world frame. """ - pass + body = self._bodies[body_id] + if link_id < -1 or link_id > (body.num_bodies - 2): + raise ValueError("link_id should belong to [-1, `num_links-2`].") + idx = body.b_idx0 + link_id + 1 + # self.sim.data.xfrc_applied[idx, :3] = force - def apply_external_torque(self, body_id, link_id=-1, torque=(0., 0., 0.), frame=1): + # if position expressed in link frame, map it to world frame + if frame == Simulator.LINK_FRAME: + xpos = np.zeros(3) + xmat = np.zeros(9) + quat = np.array([1., 0., 0., 0.]) # TODO: get orientation of the link + sameframe = 1 + mujoco.functions.mj_local2Global(self.sim.data, xpos, xmat, position, quat, idx, sameframe) + position = xpos + + # apply force + qfrc_target = np.zeros(self.model.nv) + mujoco.functions.mj_applyFT(self.model, self.sim.data, force, np.zeros(3), position, idx, qfrc_target) + return qfrc_target[body.v_idx0:body.v_idxf] + + def apply_external_torque(self, body_id, link_id=-1, torque=(0., 0., 0.), frame=Simulator.LINK_FRAME): """ Apply an external torque on a body, or a link of the body. Note that after each simulation step, the external torques are cleared to 0. @@ -1291,10 +1487,18 @@ class Mujoco(Simulator): body_id (int): unique body id. link_id (int): link id to apply the torque, if -1 it will apply the torque on the base torque (float[3]): Cartesian torques to be applied on the body - frame (int): Specify the coordinate system of force/position: either `pybullet.WORLD_FRAME` (=2) for - Cartesian world coordinates or `pybullet.LINK_FRAME` (=1) for local link coordinates. + frame (int): Specify the coordinate system of force/position: either `Simulator.WORLD_FRAME` (=2) for + Cartesian world coordinates or `Simulator.LINK_FRAME` (=1) for local link coordinates. """ - pass + body = self._bodies[body_id] + if link_id < -1 or link_id > (body.num_bodies - 2): + raise ValueError("link_id should belong to [-1, `num_links-2`].") + idx = body.b_idx0 + link_id + 1 + # self.sim.data.xfrc_applied[idx, 3:] = torque + qfrc_target = np.zeros(self.model.nv) + position = np.zeros(3) + mujoco.functions.mj_applyFT(self.model, self.sim.data, np.zeros(3), torque, position, idx, qfrc_target) + return qfrc_target[body.v_idx0:body.v_idxf] ############################# # Robots (joints and links) # @@ -1310,8 +1514,7 @@ class Mujoco(Simulator): Returns: int: number of joints with the associated body id. """ - body = self._bodies[body_id] - return len(body.links) + return self._bodies[body_id].num_joints def num_actuated_joints(self, body_id): """ @@ -1323,8 +1526,837 @@ class Mujoco(Simulator): Returns: int: number of actuated joints of the specified body. """ + return self._bodies[body_id].num_actuated_joints + + def num_links(self, body_id): + """ + Return the total number of links of the specified body. This is the same as calling `num_joints`. + + Args: + body_id (int): unique body id. + + Returns: + int: number of links with the associated body id. + """ + return self._bodies[body_id].num_links + + def get_joint_info(self, body_id, joint_id): + """ + Return information about the given joint about the specified body. + + Note that this method returns a lot of information, so specific methods have been implemented that return + only the desired information. Also, note that we do not convert the data here. + + Args: + body_id (int): unique body id. + joint_id (int): joint id is included in [0..`num_joints(body_id)`]. + + Returns: + [0] int: the same joint id as the input parameter + [1] str: name of the joint (as specified in the URDF/SDF/etc file) + [2] int: type of the joint which implies the number of position and velocity variables. + The types include JOINT_REVOLUTE (=0), JOINT_PRISMATIC (=1), JOINT_SPHERICAL (=2), + JOINT_PLANAR (=3), and JOINT_FIXED (=4). + [3] int: q index - the first position index in the positional state variables for this body + [4] int: dq index - the first velocity index in the velocity state variables for this body + [5] int: flags (reserved) + [6] float: the joint damping value (as specified in the URDF file) + [7] float: the joint friction value (as specified in the URDF file) + [8] float: the positional lower limit for slider and revolute joints + [9] float: the positional upper limit for slider and revolute joints + [10] float: maximum force specified in URDF. Note that this value is not automatically used. + You can use maxForce in 'setJointMotorControl2'. + [11] float: maximum velocity specified in URDF. Note that this value is not used in actual + motor control commands at the moment. + [12] str: name of the link (as specified in the URDF/SDF/etc file) + [13] np.array[float[3]]: joint axis in local frame (ignored for JOINT_FIXED) + [14] np.array[float[3]]: joint position in parent frame + [15] np.array[float[4]]: joint orientation in parent frame + [16] int: parent link index, -1 for base + """ body = self._bodies[body_id] - return len(body.joints) + + if joint_id < 0 or joint_id > (body.num_joints - 1): + raise ValueError("joint_id should belong to [0, `num_joints-1`].") + + dtype = body.get_joint_type(joint_id) + if dtype == 'fixed': # special care for fixed joints (as they are usually not specified in Mujoco models) + pass + idx = body.j_idx0 + joint_id + + name = self.model.joint_id2name() + dtype = self.model.jnt_type # ['free', 'ball', 'slide', 'hinge'] + q_idx = body.get_q_idx(joint_id) + dq_idx = body.get_dq_idx(joint_id) + flag = -1 + damping = self.model.dof_damping + friction = self.model.dof_frictionloss + limited = self.model.limited + limits = self.model.jnt_range + axis = self.model.jnt_axis + pos = self.model.jnt_pos + orientation = 0 + + axis_pos = self.sim.data.xaxis + # xanchor + + # stiffness + + return joint_id + + def get_joint_state(self, body_id, joint_id): + """ + Get the joint state. + + Args: + body_id (int): unique body id. + joint_id (int): joint index in range [0..num_joints(body_id)] + + Returns: + float: The position value of this joint. + float: The velocity value of this joint. + np.array[float[6]]: These are the joint reaction forces, if a torque sensor is enabled for this joint it is + [Fx, Fy, Fz, Mx, My, Mz]. Without torque sensor, it is [0, 0, 0, 0, 0, 0]. + float: This is the motor torque applied during the last stepSimulation. Note that this only applies in + VELOCITY_CONTROL and POSITION_CONTROL. If you use TORQUE_CONTROL then the applied joint motor torque + is exactly what you provide, so there is no need to report it separately. + """ + body = self._bodies[body_id] + if not isinstance(joint_id, int): + raise TypeError("Expecting the given joint id to be an int, but got instead: {}".format(type(joint_id))) + if joint_id < 0 or joint_id > (body.num_joints - 1): + raise ValueError("joint_id should belong to [0, `num_joints-1`].") + q = body.get_q_idx(joint_id, keep=True) + if q == -1: + return 0, 0, np.zeros(6), 0 + qpos_idx = body.q_idx0 + q + qvel_idx = body.v_idx0 + q + if not body.fixed: + qpos_idx += 7 + qvel_idx += 6 + + pos = self.sim.data.qpos[qpos_idx] + vel = self.sim.data.qvel[qvel_idx] + reaction_forces = np.zeros(6) + torque = self.sim.data.qfrc_applied[qvel_idx] + return pos, vel, reaction_forces, torque + + def get_joint_states(self, body_id, joint_ids): + """ + Get the joint state of the specified joints. + + Args: + body_id (int): unique body id. + joint_ids (list[int]): list of joint ids. + + Returns: + list: + float: The position value of this joint. + float: The velocity value of this joint. + np.array[float[6]]: These are the joint reaction forces, if a torque sensor is enabled for this joint + it is [Fx, Fy, Fz, Mx, My, Mz]. Without torque sensor, it is [0, 0, 0, 0, 0, 0]. + float: This is the motor torque applied during the last `step`. Note that this only applies in + VELOCITY_CONTROL and POSITION_CONTROL. If you use TORQUE_CONTROL then the applied joint motor + torque is exactly what you provide, so there is no need to report it separately. + """ + return [self.get_joint_state(body_id, joint_id) for joint_id in joint_ids] + + def reset_joint_state(self, body_id, joint_id, position, velocity=0.): + """ + Reset the state of the joint. It is best only to do this at the start, while not running the simulation: + `reset_joint_state` overrides all physics simulation. + + Args: + body_id (int): unique body id. + joint_id (int): joint index in range [0..num_joints(body_id)] + position (float): the joint position (angle in radians [rad] or position [m]) + velocity (float): the joint velocity (angular [rad/s] or linear velocity [m/s]) + """ + body = self._bodies[body_id] + if not isinstance(joint_id, int): + raise TypeError("Expecting the given joint id to be an int, but got instead: {}".format(type(joint_id))) + if joint_id < 0 or joint_id > (body.num_joints - 1): + raise ValueError("joint_id should belong to [0, `num_joints-1`].") + q = body.get_q_idx(joint_id, keep=True) + if q == -1: + return + qpos_idx = body.q_idx0 + q + qvel_idx = body.v_idx0 + q + if not body.fixed: + qpos_idx += 7 + qvel_idx += 6 + + self.sim.data.qpos[qpos_idx] = position + self.sim.data.qvel[qvel_idx] = velocity + + def enable_joint_force_torque_sensor(self, body_id, joint_ids, enable=True): + """ + You can enable or disable a joint force/torque sensor in each joint. + + Args: + body_id (int): body unique id. + joint_ids (int, int[N]): joint index in range [0..num_joints(body_id)], or list of joint ids. + enable (bool): True to enable, False to disable the force/torque sensor + """ + pass # TODO + + def set_joint_motor_control(self, body_id, joint_ids, control_mode=2, positions=None, + velocities=None, forces=None, kp=None, kd=None, max_velocity=None): + r""" + Set the joint motor control. + + In position control: + .. math:: error = Kp (x_{des} - x) + Kd (\dot{x}_{des} - \dot{x}) + + In velocity control: + .. math:: error = \dot{x}_{des} - \dot{x} + + Note that the maximum forces and velocities are not automatically used for the different control schemes. + + Args: + body_id (int): body unique id. + joint_ids (int): joint/link id, or list of joint ids. + control_mode (int): POSITION_CONTROL (=2) (which is in fact CONTROL_MODE_POSITION_VELOCITY_PD), + VELOCITY_CONTROL (=0), TORQUE_CONTROL (=1) and PD_CONTROL (=3). + positions (float, np.array[float[N]]): target joint position(s) (used in POSITION_CONTROL). + velocities (float, np.array[float[N]]): target joint velocity(ies). In VELOCITY_CONTROL and + POSITION_CONTROL, the target velocity(ies) is(are) the desired velocity of the joint. Note that the + target velocity(ies) is(are) not the maximum joint velocity(ies). In PD_CONTROL and + POSITION_CONTROL/CONTROL_MODE_POSITION_VELOCITY_PD, the final target velocities are computed using: + `kp*(erp*(desiredPosition-currentPosition)/dt)+currentVelocity+kd*(m_desiredVelocity - currentVelocity)` + forces (float, list[float]): in POSITION_CONTROL and VELOCITY_CONTROL, these are the maximum motor + forces used to reach the target values. In TORQUE_CONTROL these are the forces / torques to be applied + each simulation step. + kp (float, list[float]): position (stiffness) gain(s) (used in POSITION_CONTROL). + kd (float, list[float]): velocity (damping) gain(s) (used in POSITION_CONTROL). + max_velocity (float): in POSITION_CONTROL this limits the velocity to a maximum. + """ + pass + + def get_link_state(self, body_id, link_id, compute_velocity=False, compute_forward_kinematics=False): + """ + Get the state of the associated link. + + Args: + body_id (int): body unique id. + link_id (int): link index. + compute_velocity (bool): If True, the Cartesian world velocity will be computed and returned. + compute_forward_kinematics (bool): if True, the Cartesian world position/orientation will be recomputed + using forward kinematics. + + Returns: + [0] np.array[float[3]]: Cartesian world position of CoM + [1] np.array[float[4]]: Cartesian world orientation of CoM, in quaternion [x,y,z,w] + [2] np.array[float[3]]: local position offset of inertial frame (center of mass) expressed in the URDF + link frame + [3] np.array[float[4]]: local orientation (quaternion [x,y,z,w]) offset of the inertial frame expressed + in URDF link frame + [4] np.array[float[3]]: Cartesian world position of the URDF link frame + [5] np.array[float[4]]: Cartesian world orientation of the URDF link frame + [6] np.array[float[3]]: Cartesian world linear velocity. Only returned if `compute_velocity` is True. + [7] np.array[float[3]]: Cartesian world angular velocity. Only returned if `compute_velocity` is True. + """ + body = self._bodies[body_id] + if link_id < -1 or link_id > (body.num_bodies - 2): # -1 is for the base + raise ValueError("link_id should belong to [-1, `num_links-2`].") + idx = body.b_idx0 + 1 + link_id + + pos = self.sim.data.body_xpos[idx] # Cartesian position of body frame (same as xipos) + quat = self._convert_wxyz_to_xyzw(self.sim.data.body_xquat[idx]) # Cartesian orientation of body frame + # Note that in Mujoco the body frame is defined at the CoM + + # link frame in URDF is the joint frame in Mujoco + # TODO: read from the Tree data structure and perform the correct transformations + + if compute_velocity: + vel = self.sim.data.cvel[idx] # com-based velocity [3D rot; 3D tran] + return pos, quat, vel[3:], vel[:3] + + return pos, quat + + def get_link_states(self, body_id, link_ids, compute_velocity=False, compute_forward_kinematics=False): + """ + Get the state of the associated links. + + Args: + body_id (int): body unique id. + link_ids (list[int]): list of link index. + compute_velocity (bool): If True, the Cartesian world velocity will be computed and returned. + compute_forward_kinematics (bool): if True, the Cartesian world position/orientation will be recomputed + using forward kinematics. + + Returns: + list: + np.array[float[3]]: Cartesian position of CoM + np.array[float[4]]: Cartesian orientation of CoM, in quaternion [x,y,z,w] + np.array[float[3]]: local position offset of inertial frame (center of mass) expressed in the URDF + link frame + np.array[float[4]]: local orientation (quaternion [x,y,z,w]) offset of the inertial frame expressed + in URDF link frame + np.array[float[3]]: world position of the URDF link frame + np.array[float[4]]: world orientation of the URDF link frame + np.array[float[3]]: Cartesian world linear velocity. Only returned if `compute_velocity` is True. + np.array[float[3]]: Cartesian world angular velocity. Only returned if `compute_velocity` is True. + """ + return [self.get_link_state(body_id, link_id, compute_velocity, compute_forward_kinematics) + for link_id in link_ids] + + def get_link_names(self, body_id, link_ids): + """ + Return the name of the given link(s). + + Args: + body_id (int): unique body id. + link_ids (int, list[int]): link id, or list of link ids. + + Returns: + if 1 link: + str: link name + if multiple links: + str[N]: link names + """ + pass + + def _get_link_result(self, body_id, link_ids, mujoco_data, fct=None, slice=None): + body = self._bodies[body_id] # TODO: maybe use the tree data structure instead... + one_link = isinstance(link_ids, int) + if one_link: + link_ids = [link_ids] + results = [] + for link_id in link_ids: + if link_id < -1 or link_id > (body.num_bodies - 2): # -1 is for the base + raise ValueError("link_id should belong to [-1, `num_links-2`].") + idx = body.b_idx0 + 1 + link_id + data = mujoco_data[idx] + if slice is not None: + data = data[slice] + if fct is not None: + data = fct(data) + results.append(data) + if one_link and len(results) > 0: + return results[0] + return results + + def get_link_masses(self, body_id, link_ids): + """ + Return the mass of the given link(s). + + Args: + body_id (int): unique body id. + link_ids (int, list[int]): link id, or list of link ids. + + Returns: + if 1 link: + float: mass of the given link + else: + float[N]: mass of each link + """ + # TODO: maybe use the tree data structure instead... + return self._get_link_result(body_id, link_ids, self.sim.model.body_mass) + + def get_link_frames(self, body_id, link_ids): + r""" + Return the link world frame position(s) and orientation(s). + + Args: + body_id (int): body id. + link_ids (int, int[N]): link id, or list of desired link ids. + + Returns: + if 1 link: + np.array[float[3]]: the link frame position in the world space + np.array[float[4]]: Cartesian orientation of the link frame [x,y,z,w] + if multiple links: + np.array[float[N,3]]: link frame position of each link in world space + np.array[float[N,4]]: orientation of each link frame [x,y,z,w] + """ + pass + + def get_link_world_positions(self, body_id, link_ids): + """ + Return the CoM position (in the Cartesian world space coordinates) of the given link(s). + + Args: + body_id (int): unique body id. + link_ids (int, list[int]): link index, or list of link indices. + + Returns: + if 1 link: + np.array[float[3]]: the link CoM position in the world space + if multiple links: + np.array[float[N,3]]: CoM position of each link in world space + """ + # Cartesian position of body frame (same as xipos) + return np.asarray(self._get_link_result(body_id, link_ids, self.sim.data.body_xpos)) + + def get_link_positions(self, body_id, link_ids): + pass + + def get_link_world_orientations(self, body_id, link_ids): + """ + Return the CoM orientation (in the Cartesian world space) of the given link(s). + + Args: + body_id (int): unique body id. + link_ids (int, list[int]): link index, or list of link indices. + + Returns: + if 1 link: + np.array[float[4]]: Cartesian orientation of the link CoM (x,y,z,w) + if multiple links: + np.array[float[N,4]]: CoM orientation of each link (x,y,z,w) + """ + # Cartesian orientation of body frame + return np.asarray(self._get_link_result(body_id, link_ids, self.sim.data.body_xquat, + fct=self._convert_wxyz_to_xyzw)) + + def get_link_orientations(self, body_id, link_ids): + pass + + def get_link_world_linear_velocities(self, body_id, link_ids): + """ + Return the linear velocity of the link(s) expressed in the Cartesian world space coordinates. + + Args: + body_id (int): unique body id. + link_ids (int, list[int]): link index, or list of link indices. + + Returns: + if 1 link: + np.array[float[3]]: linear velocity of the link in the Cartesian world space + if multiple links: + np.array[float[N,3]]: linear velocity of each link + """ + return np.asarray(self._get_link_result(body_id, link_ids, self.sim.data.cvel, slice=slice(3, 6))) + + def get_link_world_angular_velocities(self, body_id, link_ids): + """ + Return the angular velocity of the link(s) in the Cartesian world space coordinates. + + Args: + body_id (int): unique body id. + link_ids (int, list[int]): link index, or list of link indices. + + Returns: + if 1 link: + np.array[float[3]]: angular velocity of the link in the Cartesian world space + if multiple links: + np.array[float[N,3]]: angular velocity of each link + """ + return np.asarray(self._get_link_result(body_id, link_ids, self.sim.data.cvel, slice=slice(3))) + + def get_link_world_velocities(self, body_id, link_ids): + """ + Return the linear and angular velocities (expressed in the Cartesian world space coordinates) for the given + link(s). + + Args: + body_id (int): unique body id. + link_ids (int, list[int]): link index, or list of link indices. + + Returns: + if 1 link: + np.array[float[6]]: linear and angular velocity of the link in the Cartesian world space + if multiple links: + np.array[float[N,6]]: linear and angular velocity of each link + """ + velocities = np.asarray(self._get_link_result(body_id, link_ids, self.sim.data.cvel)) + if velocities.ndim == 1: + return np.roll(velocities, shift=3) + return np.roll(velocities, shift=3, axis=1) + + def get_link_velocities(self, body_id, link_ids): + pass + + def get_link_world_linear_accelerations(self, body_id, link_ids): + """ + Return the linear acceleration of the link(s) expressed in the Cartesian world space coordinates. + + Args: + body_id (int): unique body id. + link_ids (int, list[int]): link index, or list of link indices. + + Returns: + if 1 link: + np.array[float[3]]: linear acceleration of the link in the Cartesian world space + if multiple links: + np.array[float[N,3]]: linear acceleration of each link + """ + # com-based acceleration + return np.asarray(self._get_link_result(body_id, link_ids, self.sim.data.cacc, slice=slice(3, 6))) + + def get_link_world_angular_accelerations(self, body_id, link_ids): + """ + Return the angular acceleration of the link(s) in the Cartesian world space coordinates. + + Args: + body_id (int): unique body id. + link_ids (int, list[int]): link index, or list of link indices. + + Returns: + if 1 link: + np.array[float[3]]: angular acceleration of the link in the Cartesian world space + if multiple links: + np.array[float[N,3]]: angular acceleration of each link + """ + # com-based acceleration + return np.asarray(self._get_link_result(body_id, link_ids, self.sim.data.cacc, slice=slice(3))) + + def get_link_world_accelerations(self, body_id, link_ids): + """ + Return the linear and angular accelerations (expressed in the Cartesian world space coordinates) for the given + link(s). This is only valid if the simulator `supports_acceleration`. + + Args: + body_id (int): unique body id. + link_ids (int, list[int]): link index, or list of link indices. + + Returns: + if 1 link: + np.array[float[6]]: linear and angular acceleration of the link in the Cartesian world space + if multiple links: + np.array[float[N,6]]: linear and angular acceleration of each link + """ + accelerations = np.asarray(self._get_link_result(body_id, link_ids, self.sim.data.cacc)) + if accelerations.ndim == 1: + return np.roll(accelerations, shift=3) + return np.roll(accelerations, shift=3, axis=1) + + def get_q_indices(self, body_id, joint_ids): + """ + Get the corresponding q index of the given joint(s). + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): a joint id, or list of joint ids. + + Returns: + if 1 joint: + int: q index + if multiple joints: + list[int]: q indices + """ + body = self._bodies[body_id] + q_idx = body.get_q_idx(joint_ids) + return q_idx + + def get_actuated_joint_ids(self, body_id): + """ + Get the actuated joint ids associated with the given body id. + + Args: + body_id (int): unique body id. + + Returns: + list[int]: actuated joint ids. + """ + body = self._bodies[body_id] + return [i for i, joint in enumerate(body.joints) if joint.dtype != 'fixed' and joint.dtype != 'floating'] + + def get_joint_names(self, body_id, joint_ids): + """ + Return the name of the given joint(s). + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): a joint id, or list of joint ids. + + Returns: + if 1 joint: + str: name of the joint + if multiple joints: + str[N]: name of each joint + """ + body = self._bodies[body_id] + one_joint = isinstance(joint_ids, int) + if one_joint: + joint_ids = [joint_ids] + names = [] + for joint_id in joint_ids: + name = body.joints[joint_id].name + if name.startswith('prl_'): + name = '_'.join(name.split('_')[1:-1]) + names.append(name) + if one_joint and len(names) > 1: + return names[0] + return names + + def get_joint_type_ids(self, body_id, joint_ids): + """ + Get the joint type ids. + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): a joint id, or list of joint ids. + + Returns: + if 1 joint: + int: joint type id. + if multiple joints: list of above + """ + pass + + def get_joint_type_names(self, body_id, joint_ids): + """ + Get joint type names. + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): a joint id, or list of joint ids. + + Returns: + if 1 joint: + str: joint type name. + if multiple joints: list of above + """ + pass + + def get_joint_dampings(self, body_id, joint_ids): + """ + Get the damping coefficient of the given joint(s). + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): a joint id, or list of joint ids. + + Returns: + if 1 joint: + float: damping coefficient of the given joint + if multiple joints: + np.array[float[N]]: damping coefficient for each specified joint + """ + pass + + def get_joint_frictions(self, body_id, joint_ids): + """ + Get the friction coefficient of the given joint(s). + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): a joint id, or list of joint ids. + + Returns: + if 1 joint: + float: friction coefficient of the given joint + if multiple joints: + np.array[float[N]]: friction coefficient for each specified joint + """ + pass + + def get_joint_limits(self, body_id, joint_ids): + """ + Get the joint limits of the given joint(s). + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): a joint id, or list of joint ids. + + Returns: + if 1 joint: + np.array[float[2]]]: lower and upper limit + if multiple joints: + np.array[N,2]: lower and upper limit for each specified joint + """ + pass + + def get_joint_max_forces(self, body_id, joint_ids): + """ + Get the maximum force that can be applied on the given joint(s). + + Warning: Note that this is not automatically used in position, velocity, or torque control. + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): a joint id, or list of joint ids. + + Returns: + if 1 joint: + float: maximum force [N] + if multiple joints: + np.array[float[N]]: maximum force for each specified joint [N] + """ + pass + + def get_joint_max_velocities(self, body_id, joint_ids): + """ + Get the maximum velocity that can be applied on the given joint(s). + + Warning: Note that this is not automatically used in position, velocity, or torque control. + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): a joint id, or list of joint ids. + + Returns: + if 1 joint: + float: maximum velocity [rad/s] + if multiple joints: + np.array[float[N]]: maximum velocities for each specified joint [rad/s] + """ + pass + + def get_joint_axes(self, body_id, joint_ids): + """ + Get the joint axis about the given joint(s). + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): a joint id, or list of joint ids. + + Returns: + if 1 joint: + np.array[float[3]]: joint axis + if multiple joint: + np.array[float[N,3]]: list of joint axis + """ + pass + + def set_joint_positions(self, body_id, joint_ids, positions, velocities=None, kps=None, kds=None, forces=None): + """ + Set the position of the given joint(s) (using position control). + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): joint id, or list of joint ids. + positions (float, np.array[float[N]]): desired position, or list of desired positions [rad] + velocities (None, float, np.array[float[N]]): desired velocity, or list of desired velocities [rad/s] + kps (None, float, np.array[float[N]]): position gain(s) + kds (None, float, np.array[float[N]]): velocity gain(s) + forces (None, float, np.array[float[N]]): maximum motor force(s)/torque(s) used to reach the target values. + """ + pass + + def get_joint_positions(self, body_id, joint_ids): + """ + Get the position of the given joint(s). + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): joint id, or list of joint ids. + + Returns: + if 1 joint: + float: joint position [rad] + if multiple joints: + np.array[float[N]]: joint positions [rad] + """ + body = self._bodies[body_id] + pass + + def set_joint_velocities(self, body_id, joint_ids, velocities, max_force=None): + """ + Set the velocity of the given joint(s) (using velocity control). + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): joint id, or list of joint ids. + velocities (float, np.array[float[N]]): desired velocity, or list of desired velocities [rad/s] + max_force (None, float, np.array[float[N]]): maximum motor forces/torques + """ + pass + + def get_joint_velocities(self, body_id, joint_ids): + """ + Get the velocity of the given joint(s). + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): joint id, or list of joint ids. + + Returns: + if 1 joint: + float: joint velocity [rad/s] + if multiple joints: + np.array[float[N]]: joint velocities [rad/s] + """ + pass + + def set_joint_accelerations(self, body_id, joint_ids, accelerations, q=None, dq=None): + """ + Set the acceleration of the given joint(s) (using force control). This is achieved by performing inverse + dynamic which given the joint accelerations compute the joint torques to be applied. + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): joint id, or list of joint ids. + accelerations (float, np.array[float[N]]): desired joint acceleration, or list of desired joint + accelerations [rad/s^2] + """ + pass + + def get_joint_accelerations(self, body_id, joint_ids): # , q=None, dq=None): + """ + Get the acceleration of the specified joint(s). This is only valid if the simulator `supports_acceleration`. + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): joint id, or list of joint ids. + + Returns: + if 1 joint: + float: joint acceleration [rad/s^2] + if multiple joints: + np.array[float[N]]: joint accelerations [rad/s^2] + """ + pass + + def set_joint_torques(self, body_id, joint_ids, torques): + """ + Set the torque/force to the given joint(s) (using force/torque control). + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): joint id, or list of joint ids. + torques (float, list[float], np.array[float]): desired torque(s) to apply to the joint(s) [N]. + """ + pass + + def get_joint_torques(self, body_id, joint_ids): + """ + Get the applied torque(s) on the given joint(s). + + Args: + body_id (int): unique body id. + joint_ids (int, list[int]): a joint id, or list of joint ids. + + Returns: + if 1 joint: + float: torque [Nm] + if multiple joints: + np.array[float[N]]: torques associated to the given joints [Nm] + """ + pass + + def get_joint_reaction_forces(self, body_id, joint_ids): + """Return the joint reaction forces at the given joint. Note that the torque sensor must be enabled, otherwise + it will always return [0,0,0,0,0,0]. + + Args: + body_id (int): unique body id. + joint_ids (int, int[N]): joint id, or list of joint ids + + Returns: + if 1 joint: + np.array[float[6]]: joint reaction force (fx,fy,fz,mx,my,mz) [N,Nm] + if multiple joints: + np.array[float[N,6]]: joint reaction forces [N, Nm] + """ + pass + + def get_joint_powers(self, body_id, joint_ids): + """Return the applied power at the given joint(s). Power = torque * velocity. + + Args: + body_id (int): unique body id. + joint_ids (int, int[N]): joint id, or list of joint ids + + Returns: + if 1 joint: + float: joint power [W] + if multiple joints: + np.array[float[N]]: power at each joint [W] + """ + pass ################# # Visualization # @@ -1680,6 +2712,386 @@ class Mujoco(Simulator): vec = to_position - from_position return self.sim.ray(pnt=from_position, vec=vec) # this return the distance and id of the geom + ########################### + # Kinematics and Dynamics # + ########################### + + def get_dynamics_info(self, body_id, link_id=-1): + """ + Get dynamic information about the mass, center of mass, friction and other properties of the base and links. + + Args: + body_id (int): body/object unique id. + link_id (int): link/joint index or -1 for the base. + + Returns: + [0] float: mass in kg + [1] float: lateral friction coefficient + [2] np.array[float[3]]: local inertia diagonal. Note that links and base are centered around the center of + mass and aligned with the principal axes of inertia. + [3] np.array[float[3]]: position of inertial frame in local coordinates of the joint frame + [4] np.array[float[4]]: orientation of inertial frame in local coordinates of joint frame + [5] float: coefficient of restitution + [6] float: rolling friction coefficient orthogonal to contact normal + [7] float: spinning friction coefficient around contact normal + [8] float: damping of contact constraints. -1 if not available. + [9] float: stiffness of contact constraints. -1 if not available. + """ + body = self._bodies[body_id] + if link_id < -1 or link_id > (body.num_bodies - 2): # -1 is for the base + raise ValueError("link_id should belong to [-1, `num_links-2`].") + + idx = body.b_idx0 + 1 + link_id + mass = self.sim.model.body_mass[idx] + + inertia = self.model.body_inertia[idx] # diagonal inertia in ipos/iquat frame + position = self.model.body_ipos[idx] + orientation = self._convert_wxyz_to_xyzw(self.model.body_iquat[idx]) + + slide, spin, roll = self.model.geom_friction[idx] # TODO: one body can have multiple geoms + + # TODO: see solref and solimp + # From Todorov: "The solref and solimp parameters can be adjusted to obtain different restitution effects + # but they are not in one-to-one correspondence with the notion of restitution, and it is generally not + # possible to guarantee a fixed coefficient of restitution. You can make contacts more bouncy by increasing + # the second parameter in solref; it corresponds to a damping coefficient: 1 = critical damping, less than + # 1 = under-damped, more than 1 = over-damped. + # http://www.mujoco.org/forum/index.php?threads/coefficient-of-restitution.3426/ + + # Check: self.model.geom_solmix[idx], self.model.geom_solref[idx], self.model.geom_solimp[idx] + # Check section 'Restitution' in http://mujoco.org/book/modeling.html + restitution = 0 + + # 1. http://www.mujoco.org/book/modeling.html#CContact + # 2. http://www.mujoco.org/book/modeling.html#CSolver + damping = -1 + stiffness = -1 + + return mass, slide, inertia, position, orientation, restitution, roll, spin, damping, stiffness + + def change_dynamics(self, body_id, link_id=-1, mass=None, lateral_friction=None, spinning_friction=None, + rolling_friction=None, restitution=None, linear_damping=None, angular_damping=None, + contact_stiffness=None, contact_damping=None, friction_anchor=None, + local_inertia_diagonal=None, joint_damping=None): + """ + Change dynamic properties of the given body (or link) such as mass, friction and restitution coefficients, etc. + + Args: + body_id (int): object unique id, as returned by `load_urdf`, etc. + link_id (int): link index or -1 for the base. + mass (float): change the mass of the link (or base for link index -1) + lateral_friction (float): lateral (linear) contact friction + spinning_friction (float): torsional friction around the contact normal + rolling_friction (float): torsional friction orthogonal to contact normal + restitution (float): bouncyness of contact. Keep it a bit less than 1. + linear_damping (float): linear damping of the link (0.04 by default) + angular_damping (float): angular damping of the link (0.04 by default) + contact_stiffness (float): stiffness of the contact constraints, used together with `contact_damping` + contact_damping (float): damping of the contact constraints for this body/link. Used together with + `contact_stiffness`. This overrides the value if it was specified in the URDF file in the contact + section. + friction_anchor (int): enable or disable a friction anchor: positional friction correction (disabled by + default, unless set in the URDF contact section) + local_inertia_diagonal (np.array[float[3]]): diagonal elements of the inertia tensor. Note that the base + and links are centered around the center of mass and aligned with the principal axes of inertia so + there are no off-diagonal elements in the inertia tensor. + joint_damping (float): joint damping coefficient applied at each joint. This coefficient is read from URDF + joint damping field. Keep the value close to 0. + `joint_damping_force = -damping_coefficient * joint_velocity`. + """ + body = self._bodies[body_id] + if link_id < -1 or link_id > (body.num_bodies - 2): # -1 is for the base + raise ValueError("link_id should belong to [-1, `num_links-2`].") + + idx = body.b_idx0 + 1 + link_id + if mass is not None: + self.sim.model.body_mass[idx] = mass + + if lateral_friction is not None: # TODO: one body can have multiple geoms + self.model.geom_friction[idx][0] = lateral_friction + + if spinning_friction is not None: + self.model.geom_friction[idx][1] = spinning_friction + + if rolling_friction is not None: + self.model.geom_friction[idx][2] = rolling_friction + + if restitution is not None: # TODO + # play with solref and solimp + # self.model.geom_solmix[idx] = + # self.model.geom_solref[idx] = + # self.model.geom_solimp[idx] = + pass + + if linear_damping is not None: + pass + + if angular_damping is not None: + pass + + if contact_stiffness is not None: + pass + + if contact_damping is not None: + pass + + if friction_anchor is not None: + pass + + # inertia + if local_inertia_diagonal is not None: + self.model.body_inertia[idx] = local_inertia_diagonal + if inertia_position is not None: + self.model.body_ipos[idx] = inertia_position + if inertia_orientation is not None: + self.model.body_iquat[idx] = inertia_orientation + + if joint_damping is not None: + if link_id < 0 or link_id > (body.num_joints - 1): + raise ValueError("link_id should belong to [0, `num_joints-1`] when setting the joint damping.") + idx = body.v_idx0 + link_id + if not body.fixed: + idx += 6 # + self.model.dof_damping[idx] = joint_damping + + def calculate_jacobian(self, body_id, link_id, local_position, q=None, dq=None, des_ddq=None): + r""" + Return the full geometric Jacobian matrix :math:`J(q) = [J_{lin}(q), J_{ang}(q)]^T`, such that: + + .. math:: v = [\dot{p}, \omega]^T = J(q) \dot{q} + + where :math:`\dot{p}` is the Cartesian linear velocity of the link, and :math:`\omega` is its angular velocity. + + Warnings: if we have a floating base then the Jacobian will also include columns corresponding to the root + link DoFs (at the beginning). If it is a fixed base, it will only have columns associated with the joints. + + Args: + body_id (int): unique body id. + link_id (int): link id. + local_position (np.array[float[3]]): the point on the specified link to compute the Jacobian (in link local + coordinates around its center of mass). If None, it will use the CoM position (in the link frame). + q (np.array[float[N]]): joint positions of size N, where N is the number of DoFs. + dq (np.array[float[N]]): joint velocities of size N, where N is the number of DoFs. + des_ddq (np.array[float[N]]): desired joint accelerations of size N. + + Returns: + np.array[float[6,N]], np.array[float[6,6+N]]: full geometric (linear and angular) Jacobian matrix. The + number of columns depends if the base is fixed or floating. + """ + body = self._bodies[body_id] + if link_id < -1 or link_id > (body.num_bodies - 2): # -1 is for the base + raise ValueError("link_id should belong to [-1, `num_links-2`].") + + # TODO: use q, dq, des_ddq by setting it in the data and then restoring the data + + idx = body.b_idx0 + 1 + link_id + jacp, jacr = np.zeros(3 * self.model.nv), np.zeros(3 * self.model.nv) + if local_position is None: + local_position = np.zeros(3) + mujoco.functions.mj_jac(self.model, self.sim.data, jacp, jacr, local_position, idx) + jacp = jacp.reshape(3, self.model.nv)[:, body.v_idx0:body.v_idxf] + jacr = jacr.reshape(3, self.model.nv)[:, body.v_idx0:body.v_idxf] + return np.vstack((jacp, jacr)) + + def calculate_mass_matrix(self, body_id, q): + r""" + Return the mass/inertia matrix :math:`H(q)`, which is used in the rigid-body equation of motion (EoM) in joint + space given by (see [1]): + + .. math:: \tau = H(q)\ddot{q} + C(q,\dot{q}) + + where :math:`\tau` is the vector of applied torques, :math:`H(q)` is the inertia matrix, and + :math:`C(q,\dot{q}) \dot{q}` is the vector accounting for Coriolis, centrifugal forces, gravity, and any + other forces acting on the system except the applied torques :math:`\tau`. + + Warnings: If the base is floating, it will return a [6+N,6+N] inertia matrix, where N is the number of actuated + joints. If the base is fixed, it will return a [N,N] inertia matrix + + Args: + body_id (int): body unique id. + q (np.array[float[N]]): joint positions of size N, where N is the total number of DoFs. + + Returns: + np.array[float[N,N]], np.array[float[6+N,6+N]]: inertia matrix + """ + body = self._bodies[body_id] + + # TODO: use q + + # get sparse matrix + sparse_inertia = self.sim.data.qM + + # Convert sparse inertia matrix M into full (i.e. dense) matrix + inertia = np.zeros(self.model.nv * self.model.nv) + mujoco.functions.mj_fullM(self.model, inertia, sparse_inertia) + inertia = inertia.reshape(self.model.nv, self.model.nv)[body.v_idx0:body.v_idxf, body.v_idx0:body.v_idxf] + return inertia + + def calculate_inverse_kinematics(self, body_id, link_id, position, orientation=None, lower_limits=None, + upper_limits=None, joint_ranges=None, rest_poses=None, joint_dampings=None, + solver=None, q_curr=None, max_iters=None, threshold=None): + r""" + Compute the FULL Inverse kinematics; it will return a position for all the actuated joints. + + "You can compute the joint angles that makes the end-effector reach a given target position in Cartesian world + space. Internally, Bullet uses an improved version of Samuel Buss Inverse Kinematics library. At the moment + only the Damped Least Squares method with or without Null Space control is exposed, with a single end-effector + target. Optionally you can also specify the target orientation of the end effector. In addition, there is an + option to use the null-space to specify joint limits and rest poses. This optional null-space support requires + all 4 lists (lower_limits, upper_limits, joint_ranges, rest_poses), otherwise regular IK will be used." [1] + + Args: + body_id (int): body unique id, as returned by `load_urdf`, etc. + link_id (int): end effector link index. + position (np.array[float[3]]): target position of the end effector (its link coordinate, not center of mass + coordinate!). By default this is in Cartesian world space, unless you provide `q_curr` joint angles. + orientation (np.array[float[4]]): target orientation in Cartesian world space, quaternion [x,y,w,z]. If not + specified, pure position IK will be used. + lower_limits (np.array[float[N]], list of N floats): lower joint limits. Optional null-space IK. + upper_limits (np.array[float[N]], list of N floats): upper joint limits. Optional null-space IK. + joint_ranges (np.array[float[N]], list of N floats): range of value of each joint. + rest_poses (np.array[float[N]], list of N floats): joint rest poses. Favor an IK solution closer to a + given rest pose. + joint_dampings (np.array[float[N]], list of N floats): joint damping factors. Allow to tune the IK solution + using joint damping factors. + solver (int): p.IK_DLS (=0) or p.IK_SDLS (=1), Damped Least Squares or Selective Damped Least Squares, as + described in the paper by Samuel Buss "Selectively Damped Least Squares for Inverse Kinematics". + q_curr (np.array[float[N]]): list of joint positions. By default PyBullet uses the joint positions of the + body. If provided, the target_position and targetOrientation is in local space! + max_iters (int): maximum number of iterations. Refine the IK solution until the distance between target + and actual end effector position is below this threshold, or the `max_iters` is reached. + threshold (float): residual threshold. Refine the IK solution until the distance between target and actual + end effector position is below this threshold, or the `max_iters` is reached. + + Returns: + np.array[float[N]]: joint positions (for each actuated joint). + """ + pass + + def calculate_inverse_dynamics(self, body_id, q, dq, des_ddq): + r""" + Starting from the specified joint positions :math:`q` and velocities :math:`\dot{q}`, it computes the joint + torques :math:`\tau` required to reach the desired joint accelerations :math:`\ddot{q}_{des}`. That is, + :math:`\tau = ID(model, q, \dot{q}, \ddot{q}_{des})`. + + Specifically, it uses the rigid-body equation of motion in joint space given by (see [1]): + + .. math:: \tau = H(q)\ddot{q} + C(q,\dot{q}) + + where :math:`\tau` is the vector of applied torques, :math:`H(q)` is the inertia matrix, and + :math:`C(q,\dot{q}) \dot{q}` is the vector accounting for Coriolis, centrifugal forces, gravity, and any + other forces acting on the system except the applied torques :math:`\tau`. + + Normally, a more popular form of this equation of motion (in joint space) is given by: + + .. math:: H(q) \ddot{q} + S(q,\dot{q}) \dot{q} + g(q) = \tau + J^T(q) F + + which is the same as the first one with :math:`C = S\dot{q} + g(q) - J^T(q) F`. However, this last formulation + is useful to understand what happens when we set some variables to 0. + Assuming that there are no forces acting on the system, and giving desired joint accelerations of 0, this + method will return :math:`\tau = S(q,\dot{q}) \dot{q} + g(q)`. If in addition joint velocities are also 0, + it will return :math:`\tau = g(q)` which can for instance be useful for gravity compensation. + + For forward dynamics, which computes the joint accelerations given the joint positions, velocities, and + torques (that is, :math:`\ddot{q} = FD(model, q, \dot{q}, \tau)`, this can be computed using + :math:`\ddot{q} = H^{-1} (\tau - C)` (see also `computeFullFD`). For more information about different + control schemes (position, force, impedance control and others), or about the formulation of the equation + of motion in task/operational space (instead of joint space), check the references [1-4]. + + Args: + body_id (int): body unique id. + q (np.array[float[N]]): joint positions + dq (np.array[float[N]]): joint velocities + des_ddq (np.array[float[N]]): desired joint accelerations + + Returns: + np.array[float[N]]: joint torques computed using the rigid-body equation of motion + + References: + - [1] "Rigid Body Dynamics Algorithms", Featherstone, 2008, chap1.1 + - [2] "Robotics: Modelling, Planning and Control", Siciliano et al., 2010 + - [3] "Springer Handbook of Robotics", Siciliano et al., 2008 + - [4] Lecture on "Impedance Control" by Prof. De Luca, Universita di Roma, + http://www.diag.uniroma1.it/~deluca/rob2_en/15_ImpedanceControl.pdf + """ + body = self._bodies[body_id] + + # copy data + dest = mujoco.cymj.PyMjData() + mujoco.functions.mj_copyData(dest, self.model, self.sim.data) + dest.qpos[body.q_idx0:body.q_idxf] = q + dest.qvel[body.v_idx0:body.v_idxf] = dq + dest.qacc[body.v_idx0:body.v_idxf] = des_ddq + + # inverse dynamics + mujoco.functions.mj_inverse(self.model, dest) + + # get resulting torques and return it + torques = dest.qfrc_applied[body.v_idx0:body.v_idxf] + return torques + + def calculate_forward_dynamics(self, body_id, q, dq, torques): + r""" + Given the specified joint positions :math:`q` and velocities :math:`\dot{q}`, and joint torques :math:`\tau`, + it computes the joint accelerations :math:`\ddot{q}`. That is, :math:`\ddot{q} = FD(model, q, \dot{q}, \tau)`. + + Specifically, it uses the rigid-body equation of motion in joint space given by (see [1]): + + .. math:: \ddot{q} = H(q)^{-1} (\tau - C(q,\dot{q})) + + where :math:`\tau` is the vector of applied torques, :math:`H(q)` is the inertia matrix, and + :math:`C(q,\dot{q}) \dot{q}` is the vector accounting for Coriolis, centrifugal forces, gravity, and any + other forces acting on the system except the applied torques :math:`\tau`. + + Normally, a more popular form of this equation of motion (in joint space) is given by: + + .. math:: H(q) \ddot{q} + S(q,\dot{q}) \dot{q} + g(q) = \tau + J^T(q) F + + which is the same as the first one with :math:`C = S\dot{q} + g(q) - J^T(q) F`. However, this last formulation + is useful to understand what happens when we set some variables to 0. + Assuming that there are no forces acting on the system, and giving desired joint torques of 0, this + method will return :math:`\ddot{q} = - H(q)^{-1} (S(q,\dot{q}) \dot{q} + g(q))`. If in addition + the joint velocities are also 0, it will return :math:`\ddot{q} = - H(q)^{-1} g(q)` which are + the accelerations due to gravity. + + For inverse dynamics, which computes the joint torques given the joint positions, velocities, and + accelerations (that is, :math:`\tau = ID(model, q, \dot{q}, \ddot{q})`, this can be computed using + :math:`\tau = H(q)\ddot{q} + C(q,\dot{q})`. For more information about different + control schemes (position, force, impedance control and others), or about the formulation of the equation + of motion in task/operational space (instead of joint space), check the references [1-4]. + + Args: + body_id (int): unique body id. + q (np.array[float[N]]): joint positions + dq (np.array[float[N]]): joint velocities + torques (np.array[float[N]]): desired joint torques + + Returns: + np.array[float[N]]: joint accelerations computed using the rigid-body equation of motion + + References: + - [1] "Rigid Body Dynamics Algorithms", Featherstone, 2008, chap1.1 + - [2] "Robotics: Modelling, Planning and Control", Siciliano et al., 2010 + - [3] "Springer Handbook of Robotics", Siciliano et al., 2008 + - [4] Lecture on "Impedance Control" by Prof. De Luca, Universita di Roma, + http://www.diag.uniroma1.it/~deluca/rob2_en/15_ImpedanceControl.pdf + """ + body = self._bodies[body_id] + + # copy data and set q, dq, tau + dest = mujoco.cymj.PyMjData() + mujoco.functions.mj_copyData(dest, self.model, self.sim.data) + dest.qpos[body.q_idx0:body.q_idxf] = q + dest.qvel[body.v_idx0:body.v_idxf] = dq + dest.qfrc_applied[body.v_idx0:body.v_idxf] = torques + + # forward dynamics + mujoco.functions.mj_forward(self.model, dest) + + # get ddq and return it + qacc = dest.qacc[body.v_idx0:body.v_idxf] + return qacc + # Test if __name__ == '__main__': diff --git a/pyrobolearn/simulators/simulator.py b/pyrobolearn/simulators/simulator.py index dc2d837..ed06381 100644 --- a/pyrobolearn/simulators/simulator.py +++ b/pyrobolearn/simulators/simulator.py @@ -1219,7 +1219,7 @@ class Simulator(object): """ pass - def reset_joint_state(self, body_id, joint_id, position, velocity=0.): + def reset_joint_state(self, body_id, joint_id, position, velocity=None): """ Reset the state of the joint. It is best only to do this at the start, while not running the simulation: `reset_joint_state` overrides all physics simulation. @@ -1467,6 +1467,38 @@ class Simulator(object): def get_link_velocities(self, body_id, link_ids): pass + def get_link_world_linear_accelerations(self, body_id, link_ids): + """ + Return the linear acceleration of the link(s) expressed in the Cartesian world space coordinates. + + Args: + body_id (int): unique body id. + link_ids (int, list[int]): link index, or list of link indices. + + Returns: + if 1 link: + np.array[float[3]]: linear acceleration of the link in the Cartesian world space + if multiple links: + np.array[float[N,3]]: linear acceleration of each link + """ + pass + + def get_link_world_angular_accelerations(self, body_id, link_ids): + """ + Return the angular acceleration of the link(s) in the Cartesian world space coordinates. + + Args: + body_id (int): unique body id. + link_ids (int, list[int]): link index, or list of link indices. + + Returns: + if 1 link: + np.array[float[3]]: angular acceleration of the link in the Cartesian world space + if multiple links: + np.array[float[N,3]]: angular acceleration of each link + """ + pass + def get_link_world_accelerations(self, body_id, link_ids): """ Return the linear and angular accelerations (expressed in the Cartesian world space coordinates) for the given diff --git a/pyrobolearn/utils/parsers/robots/data_structures.py b/pyrobolearn/utils/parsers/robots/data_structures.py index 0a6fb0c..c86b92f 100644 --- a/pyrobolearn/utils/parsers/robots/data_structures.py +++ b/pyrobolearn/utils/parsers/robots/data_structures.py @@ -822,10 +822,7 @@ class MultiBody(object): @property def num_dofs(self): """Return the total number of degrees of freedom.""" - num_dofs = 0 - for joint in self.joints.values(): - num_dofs += joint.num_dofs - return num_dofs + return sum([joint.num_dofs for joint in self.joints.values()]) @property def num_bodies(self): @@ -834,23 +831,21 @@ class MultiBody(object): @property def num_joints(self): - """Return the total number of joints in this multi-body (this accounts for fixed joints as well, but not - free joints).""" + """Return the total number of joints which are not free joints (so this accounts for fixed joints as well, but + not free joints).""" # return len(self.joints) - num_joints = 0 - for joint in self.joints.values(): - if joint.dtype != 'floating': - num_joints += 1 - return num_joints + return sum([1 for joint in self.joints.values() if joint.dtype != 'floating']) + + @property + def num_free_joints(self): + """Return the total number of free joints (this does not include the fixed joints). Basically it is the joints + that have at least 1 DoF.""" + return sum([1 for joint in self.joints.values() if joint.dtype != 'fixed']) @property def num_actuated_joints(self): """Return the total number of joints which are not fixed nor free.""" - num_actuated_joints = 0 - for joint in self.joints.values(): - if joint.dtype != 'fixed' and joint.dtype != 'floating': - num_actuated_joints += 1 - return num_actuated_joints + return sum([1 for joint in self.joints.values() if joint.dtype != 'fixed' and joint.dtype != 'floating']) @property def root(self): @@ -873,6 +868,11 @@ class MultiBody(object): """Return if the root element in the tree is static or not.""" if self.root is not None: return self.root.static + if self.joints: + joint = self.joints[next(iter(self.joints))] + if joint.dtype == 'free' or joint.dtype == 'floating': + return False + return True @static.setter def static(self, static): @@ -880,6 +880,10 @@ class MultiBody(object): if self.root is not None: self.root.static = static + # aliases + fixed = static + fixed_base = static + @property def position(self): """Return the tree frame position.""" @@ -963,15 +967,25 @@ class MultiBody(object): if not isinstance(joint, Joint): raise TypeError("Expecting the given 'joint' to be an instance of `Joint`, but got instead: " "{}".format(type(joint))) - if idx is not None: + + if idx is None: self.joints[joint.name] = joint else: - # this insert - joints = OrderedDict() - for i, (joint_name, joint_instance) in enumerate(self.joints.items()): - if i == idx: + # inserts a joint at the specified index + if idx == 0 and len(self.joints) == 0: # first joint ever to insert + self.joints[joint.name] = joint + else: + joints = OrderedDict() + for i, (joint_name, joint_instance) in enumerate(self.joints.items()): + if i == idx: + joints[joint.name] = joint + joints[joint_name] = joint_instance + + if idx == len(self.joints): # last joint joints[joint.name] = joint - joint[joint_name] = joint_instance + + # replace old joint dictionary + self.joints = joints # alias