Created
March 5, 2024 10:41
-
-
Save justagist/a00a58fb960ff84eaaf44ca6f510fc71 to your computer and use it in GitHub Desktop.
Pybullet interface for any robot urdfs. Primary purpose is for testing controllers and planners. Not implemented for handling vision/camera/sensor information.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| """This is a standalone file for interfacing with any robot in pybullet. | |
| FEATURES: | |
| - Main use-case is for controlling the robot. Not implemented for testing vision/camera stuff. | |
| - Handles prismatic, revolute and continuous joints. | |
| - Control in torque (Position, velocity, torque targets) or in direct position mode. | |
| - Retrieve kinematic and dynamics data for links and robot. | |
| - Get contact states for specified end-effectors (or links) | |
| TODO: | |
| - place_on_ground: does not check if the collision is with the ground object only. It can be with any object in the world. | |
| - optimise further using joint indices instead of names wherever possible. | |
| - retrieve contact forces for links as well | |
| """ | |
| import pybullet as pb | |
| from threading import Lock | |
| import numpy as np | |
| from typing import List, Dict, Any, Tuple, TypeAlias | |
| from pprint import pprint | |
| QuatType: TypeAlias = np.ndarray | |
| """Numpy array representating quaternion in format [x,y,z,w]""" | |
| Vector3D: TypeAlias = np.ndarray | |
| """Numpy array representating 3D cartesian vector in format [x,y,z]""" | |
| def wrap_angle(angle: float | np.ndarray) -> float | np.ndarray: | |
| """Wrap the provided angle(s) (radians) to be within -pi and pi. | |
| Args: | |
| angle (float | np.ndarray): Input angle (or array of angles) in radians. | |
| Returns: | |
| float | np.ndarray: Output after wrapping. | |
| """ | |
| return (angle + np.pi) % (2 * np.pi) - np.pi | |
| def get_pybullet_object_joint_info( | |
| objectId: int, cid: int = 0, verbose: bool = False | |
| ) -> Dict[str, Any]: | |
| """Get dictionarised version of joint information from pybullet's getJointInfo for all joints. | |
| Args: | |
| objectId (int): Pybullet object ID | |
| cid (int, optional): Pybullet physics client ID. Defaults to 0. | |
| verbose (bool, optional): Verbosity flag. Defaults to False. | |
| Returns: | |
| Dict[str, Any]: Key-value pairs of joint info. | |
| Keys: | |
| "jointIndex", | |
| "jointName", | |
| "jointType", | |
| "qIndex", | |
| "uIndex", | |
| "flags", | |
| "jointDamping", | |
| "jointFriction", | |
| "jointLowerLimit", | |
| "jointUpperLimit", | |
| "jointMaxForce", | |
| "jointMaxVelocity", | |
| "linkName", | |
| "jointAxis", | |
| "parentFramePos", | |
| "parentFrameOrn", | |
| "parentIndex", | |
| "numActuatedJoint", | |
| "actuatedJointName", | |
| "actuatedJointIndex", | |
| "actuatedJointLowerLimit", | |
| "actuatedJointUpperLimit" | |
| """ | |
| PYBULLET_JOINT_INFO_KEYS = [ | |
| "jointIndex", | |
| "jointName", | |
| "jointType", | |
| "qIndex", | |
| "uIndex", | |
| "flags", | |
| "jointDamping", | |
| "jointFriction", | |
| "jointLowerLimit", | |
| "jointUpperLimit", | |
| "jointMaxForce", | |
| "jointMaxVelocity", | |
| "linkName", | |
| "jointAxis", | |
| "parentFramePos", | |
| "parentFrameOrn", | |
| "parentIndex", | |
| ] | |
| numJoints = pb.getNumJoints(objectId, physicsClientId=cid) | |
| jointInfoValues = [[] for _ in range(len(PYBULLET_JOINT_INFO_KEYS))] | |
| for i in range(numJoints): | |
| jointInfoTuple = pb.getJointInfo(objectId, i, physicsClientId=cid) | |
| for jointInfoValue, jointInfoItem in zip(jointInfoValues, jointInfoTuple): | |
| if isinstance(jointInfoItem, bytes): | |
| jointInfoItem = jointInfoItem.decode() | |
| jointInfoValue.append(jointInfoItem) | |
| jointInfo = dict(zip(PYBULLET_JOINT_INFO_KEYS, jointInfoValues)) | |
| jointInfo["numJoint"] = numJoints | |
| actuated_joint_name = [] | |
| actuated_joint_index = [] | |
| actuated_joint_lower_limits = [] | |
| actuated_joint_upper_limits = [] | |
| continuous_joint_name = [] | |
| continuous_joint_index = [] | |
| for n, name in enumerate(jointInfo["jointName"]): | |
| jtype = jointInfo["jointType"][n] | |
| if jtype == pb.JOINT_FIXED: | |
| continue | |
| j_upper_lim = jointInfo["jointUpperLimit"][n] | |
| j_lower_lim = jointInfo["jointLowerLimit"][n] | |
| actuated_joint_index.append(n) | |
| actuated_joint_name.append(name) | |
| actuated_joint_lower_limits.append(j_lower_lim) | |
| actuated_joint_upper_limits.append(j_upper_lim) | |
| if jtype == pb.JOINT_REVOLUTE and j_upper_lim == 0 and j_lower_lim == 0: | |
| continuous_joint_name.append(name) | |
| continuous_joint_index.append(n) | |
| jointInfo["numActuatedJoint"] = len(actuated_joint_name) | |
| jointInfo["actuatedJointName"] = actuated_joint_name | |
| jointInfo["actuatedJointIndex"] = actuated_joint_index | |
| jointInfo["continuousJointName"] = continuous_joint_name | |
| jointInfo["continuousJointIndex"] = continuous_joint_index | |
| jointInfo["actuatedJointLowerLimit"] = actuated_joint_lower_limits | |
| jointInfo["actuatedJointUpperLimit"] = actuated_joint_upper_limits | |
| if verbose: | |
| print("jointInfo:") | |
| pprint(jointInfo) | |
| return jointInfo | |
| def get_pybullet_object_link_info( | |
| objectId: int, cid: int = 0, verbose: bool = False | |
| ) -> Dict[str, Any]: | |
| """Get dictionarised version of link information from links of a pybullet object. | |
| Args: | |
| objectId (int): Pybullet object ID | |
| cid (int, optional): Pybullet physics client ID. Defaults to 0. | |
| verbose (bool, optional): Verbosity flag. Defaults to False. | |
| Returns: | |
| Dict[str, Any]: Key-value pairs of joint info. | |
| Keys: | |
| "mass", | |
| "lateralFriction", | |
| "localInertiaDiagonal", | |
| "localInertialPos", | |
| "localInertialOrn", | |
| "restitution", | |
| "rollingFriction", | |
| "spinningFriction", | |
| "contactDamping", | |
| "contactStiffness", | |
| "bodyType", | |
| "collisionMargin", | |
| "index", | |
| "name" | |
| """ | |
| PYBULLET_LINK_INFO_KEYS = [ | |
| "mass", | |
| "lateralFriction", | |
| "localInertiaDiagonal", | |
| "localInertialPos", | |
| "localInertialOrn", | |
| "restitution", | |
| "rollingFriction", | |
| "spinningFriction", | |
| "contactDamping", | |
| "contactStiffness", | |
| "bodyType", | |
| "collisionMargin", | |
| ] | |
| numJoints = pb.getNumJoints(objectId, physicsClientId=cid) | |
| baseName = pb.getBodyInfo(objectId, physicsClientId=cid)[0].decode() | |
| linkIndexes = [-1] + list(range(numJoints)) | |
| linkNames = [baseName] | |
| for i in range(numJoints): | |
| jointInfoTuple = pb.getJointInfo(objectId, i, physicsClientId=cid) | |
| linkNames.append(jointInfoTuple[12].decode()) | |
| linkInfoValues = [[] for _ in range(len(PYBULLET_LINK_INFO_KEYS))] | |
| for linkIndex in linkIndexes: | |
| linkInfoTuple = pb.getDynamicsInfo(objectId, linkIndex, physicsClientId=cid) | |
| for linkInfoValue, linkInfoIterm in zip(linkInfoValues, linkInfoTuple): | |
| linkInfoValue.append(linkInfoIterm) | |
| linkInfo = dict(zip(PYBULLET_LINK_INFO_KEYS, linkInfoValues)) | |
| linkInfo["index"] = linkIndexes | |
| linkInfo["name"] = linkNames | |
| if verbose: | |
| print("linkInfo:") | |
| pprint(linkInfo) | |
| return linkInfo | |
| def get_pybullet_object_info( | |
| objectId: int, cid: int = 0, verbose: bool = True | |
| ) -> Dict[str, Any]: | |
| """Get dictionarised version of information about a pybullet object. | |
| Args: | |
| objectId (int): Pybullet object ID | |
| cid (int, optional): Pybullet physics client ID. Defaults to 0. | |
| verbose (bool, optional): Verbosity flag. Defaults to False. | |
| Returns: | |
| Dict[str, Any]: Key-value pairs of joint info. | |
| Keys: | |
| "objectId", | |
| "objectName", | |
| "objectMass", | |
| "baseName", | |
| "jointInfo", | |
| "linkInfo", | |
| """ | |
| objectId = objectId | |
| objectName = pb.getBodyInfo(objectId, physicsClientId=cid)[1].decode() | |
| baseName = pb.getBodyInfo(objectId, physicsClientId=cid)[0].decode() | |
| jointInfo = get_pybullet_object_joint_info(objectId, cid=cid) | |
| linkInfo = get_pybullet_object_link_info(objectId, cid=cid) | |
| objectMass = np.sum(linkInfo["mass"]) | |
| objectInfo = { | |
| "objectId": objectId, | |
| "objectName": objectName, | |
| "objectMass": objectMass, | |
| "baseName": baseName, | |
| "jointInfo": jointInfo, | |
| "linkInfo": linkInfo, | |
| } | |
| if verbose: | |
| print("\n" + "*" * 100 + "\nPybullet Info " + "\u2193 " * 20 + "\n" + "*" * 100) | |
| print("objectInfo:") | |
| pprint(objectInfo) | |
| print("*" * 100 + "\nPybullet Info " + "\u2191 " * 20 + "\n" + "*" * 100) | |
| return objectInfo | |
| def fix_pb_link_mass_inertia( | |
| body_id: int, | |
| cid: int = 0, | |
| mass: float = 1e-9, | |
| local_inertia_diagonal: float = (1e-9, 1e-9, 1e-9), | |
| ): | |
| """Fix pybullet virtual links. | |
| When loading urdf, pybullet assigned default value for link with no inertial (default value mass=1, | |
| localinertiadiagonal = 1,1,1; identity local inertial frame). This function fixes the link inertial | |
| values to the given mass and local_inertia_diagonal; small values by default. | |
| Args: | |
| body_id (int): Pybullet object ID | |
| cid (int, optional): Pybullet physics client ID. Defaults to 0. | |
| mass (float, optional): Virtual mass to assign to virtual link. Defaults to 1e-9. | |
| local_inertia_diagonal (float, optional): Values to assign to inertia matrix diagonals for virtual links. | |
| Defaults to (1e-9, 1e-9, 1e-9). | |
| """ | |
| link_ids = [-1] + [-1] + list(range(pb.getNumJoints(body_id, physicsClientId=cid))) | |
| for link_index in link_ids: | |
| link_info = pb.getDynamicsInfo(body_id, link_index, physicsClientId=cid) | |
| if link_info[0] == 1.0 and link_info[2] == (1.0, 1.0, 1.0): | |
| pb.changeDynamics( | |
| bodyUniqueId=body_id, | |
| linkIndex=link_index, | |
| mass=mass, | |
| localInertiaDiagonal=local_inertia_diagonal, | |
| physicsClientId=cid, | |
| ) | |
| class PybulletSimulatorInterface: | |
| def __init__( | |
| self, | |
| urdf_path: str, | |
| cid: int = None, | |
| run_async: bool = True, | |
| ee_names: List[str] = None, | |
| default_joint_positions: List[float] = None, | |
| place_on_ground: bool = False, | |
| default_base_position: Vector3D = np.zeros(3), | |
| default_base_orientation: QuatType = np.array([0, 0, 0, 1]), | |
| enable_torque_mode: bool = True, | |
| verbose: bool = True, | |
| load_ground_plane: bool = True, | |
| ): | |
| """Create a RobotInterface object for a pybullet simulation of a robot. | |
| Args: | |
| urdf_path (str): Path to urdf file of robot. | |
| cid (int, optional): Physics client id of pybullet physics engine (if already running). | |
| If None provided, will create a new physics GUI instance. Defaults to None. | |
| run_async (bool, optional): Whether to run physics in a separate thread. If set to False, | |
| step() method has to be called in the main thread. Defaults to True. | |
| ee_names (List[str], optional): List of end-effectors for the robot. Defaults to None. | |
| default_joint_positions (List[float], optional): Optional starting values for joints. | |
| Defaults to None. NOTE: These values should be in the order of joints in pybullet. | |
| place_on_ground (bool): If true, the base position height will automatically be adjusted | |
| such that the robot is on the ground with the given joint positions. Defaults to False. | |
| NOTE: This is a naive implementation. It does not actually check if the collision is | |
| with the ground object only. It can be with any object in the world. | |
| default_base_position (Vector3D, optional): Default position of the base of the robot | |
| in the world frame during start-up. Note that the height value (z) is only used if | |
| 'place_on_ground' is set to False. Defaults to np.zeros(3). | |
| default_base_orientation (QuatType, optional): Default orientation quaternion in the | |
| world frame for the base of the robot during start-up. Defaults to np.array([0, 0, 0, 1]). | |
| enable_torque_mode (bool, optional): Flag to enable torque controlling of the robot (control commands | |
| have to be sent continuously to keep the robot from falling). Defaults to True. | |
| verbose (bool, optional): Verbosity flag for debugging robot info during construction. Defaults to True. | |
| """ | |
| self.cid = cid | |
| self.ground_id = None | |
| if self.cid is None: | |
| self.cid = pb.connect(pb.GUI_SERVER) | |
| if load_ground_plane: | |
| import pybullet_data | |
| pb.setAdditionalSearchPath(pybullet_data.getDataPath()) | |
| self.ground_id = pb.loadURDF("plane.urdf") | |
| else: | |
| if place_on_ground: | |
| print( | |
| f"{__class__.__name__}: Cannot place robot on ground without creating a plane via this interface." | |
| " Set `load_ground_plane=True` to create a ground plane." | |
| ) | |
| place_on_ground = False | |
| self.sync_mode = not run_async | |
| self.default_start_pose = [default_base_position, default_base_orientation] | |
| self.ee_names = ee_names | |
| if self.ee_names is None: | |
| self.ee_names = [] | |
| self.robot_id = pb.loadURDF( | |
| urdf_path, | |
| basePosition=self.default_start_pose[0], | |
| physicsClientId=self.cid, | |
| ) | |
| pb.setGravity(0, 0, -9.81, physicsClientId=self.cid) | |
| self._state_mutex = Lock() | |
| fix_pb_link_mass_inertia(body_id=self.robot_id, cid=self.cid) | |
| self._retrieve_robot_info() | |
| if verbose: | |
| self._print_robot_info() | |
| self.default_joint_positions = ( | |
| default_joint_positions | |
| if default_joint_positions is not None | |
| else np.zeros_like(self.actuated_joint_ids) | |
| ) | |
| self.reset_actuated_joint_positions( | |
| joint_positions=self.default_joint_positions | |
| ) | |
| self.urdf_path = urdf_path | |
| if place_on_ground: | |
| self._place_robot_on_ground() | |
| self._in_torque_mode = False | |
| if enable_torque_mode: | |
| self.set_torque_control_mode() | |
| pb.setRealTimeSimulation(0 if self.sync_mode else 1, physicsClientId=self.cid) | |
| def _retrieve_robot_info(self): | |
| self.objectInfo = get_pybullet_object_info( | |
| self.robot_id, cid=self.cid, verbose=False | |
| ) | |
| self.name = self.objectInfo["objectName"] | |
| self.mass = self.objectInfo["objectMass"] | |
| self.base_name = self.objectInfo["baseName"] | |
| self.link_names = self.objectInfo["linkInfo"]["name"] | |
| self.link_ids = self.objectInfo["linkInfo"]["index"] | |
| self.link_masses = self.objectInfo["linkInfo"]["mass"] | |
| self.num_joints = self.objectInfo["jointInfo"]["numJoint"] | |
| self.joint_names = self.objectInfo["jointInfo"]["jointName"] | |
| self.joint_ids = self.objectInfo["jointInfo"]["jointIndex"] | |
| self.joint_dampings = self.objectInfo["jointInfo"]["jointDamping"] | |
| self.joint_frictions = self.objectInfo["jointInfo"]["jointFriction"] | |
| self.joint_lower_limits = self.objectInfo["jointInfo"]["jointLowerLimit"] | |
| self.joint_upper_limits = self.objectInfo["jointInfo"]["jointUpperLimit"] | |
| self.joint_neutral_positions = ( | |
| np.array(self.joint_lower_limits) + np.array(self.joint_upper_limits) | |
| ) / 2 | |
| self.num_of_actuated_joints = self.objectInfo["jointInfo"]["numActuatedJoint"] | |
| self.actuated_joint_names = self.objectInfo["jointInfo"]["actuatedJointName"] | |
| self.actuated_joint_ids = self.objectInfo["jointInfo"]["actuatedJointIndex"] | |
| self.actuated_joint_lower_limits = self.objectInfo["jointInfo"][ | |
| "actuatedJointLowerLimit" | |
| ] | |
| self.actuated_joint_upper_limits = self.objectInfo["jointInfo"][ | |
| "actuatedJointUpperLimit" | |
| ] | |
| self.continuous_joint_ids = self.objectInfo["jointInfo"]["continuousJointIndex"] | |
| self.continuous_joint_names = self.objectInfo["jointInfo"][ | |
| "continuousJointName" | |
| ] | |
| self.has_continuous_joints = len(self.continuous_joint_ids) > 0 | |
| self.actuated_joint_neutral_positions = ( | |
| np.array(self.actuated_joint_lower_limits) | |
| + np.array(self.actuated_joint_upper_limits) | |
| ) / 2 | |
| self.link_name_to_index = dict(zip(self.link_names, self.link_ids)) | |
| self.joint_name_to_index = dict(zip(self.joint_names, self.joint_ids)) | |
| self.ee_ids = [self.link_name_to_index[name] for name in self.ee_names] | |
| def _print_robot_info(self): | |
| print("\n") | |
| print("*" * 100 + "\nSimRobot Info " + "\u2193 " * 20 + "\n" + "*" * 100) | |
| print("robot ID: ", self.robot_id) | |
| print("robot name: ", self.name) | |
| print("robot mass: ", self.mass) | |
| print("base link name: ", self.base_name) | |
| print("link names: ", len(self.link_names), self.link_names) | |
| print("link indexes: ", len(self.link_ids), self.link_ids) | |
| print("link masses: ", len(self.link_masses), self.link_masses) | |
| print("end-effectors ", self.ee_names) | |
| print("num of joints: ", self.num_joints) | |
| print("num of actuated joints:", self.num_of_actuated_joints) | |
| print("joint names: ", len(self.joint_names), self.joint_names) | |
| print("joint indexes: ", len(self.joint_ids), self.joint_ids) | |
| print("joint dampings: ", len(self.joint_dampings), self.joint_dampings) | |
| print( | |
| "joint frictions: ", len(self.joint_frictions), self.joint_frictions | |
| ) | |
| print( | |
| "joint lower limits: ", | |
| len(self.joint_lower_limits), | |
| self.joint_lower_limits, | |
| ) | |
| print( | |
| "joint higher limits: ", | |
| len(self.joint_upper_limits), | |
| self.joint_upper_limits, | |
| ) | |
| print( | |
| "actuated joint names: ", | |
| len(self.actuated_joint_names), | |
| self.actuated_joint_names, | |
| ) | |
| print( | |
| "actuated joint indexes:", | |
| len(self.actuated_joint_ids), | |
| self.actuated_joint_ids, | |
| ) | |
| print( | |
| "continuous joint names: ", | |
| len(self.continuous_joint_names), | |
| self.continuous_joint_names, | |
| ) | |
| print( | |
| "continuous joint indexes:", | |
| len(self.continuous_joint_ids), | |
| self.continuous_joint_ids, | |
| ) | |
| print("*" * 100 + "\nSimRobot Info " + "\u2191 " * 20 + "\n" + "*" * 100) | |
| print("\n") | |
| def step(self): | |
| """Step simulation (if this object was created with `async=False` argument).""" | |
| if self.sync_mode: | |
| pb.stepSimulation(physicsClientId=self.cid) | |
| def _place_robot_on_ground( | |
| self, move_resolution: float = 0.01, default_position=None | |
| ): | |
| pb.setRealTimeSimulation(0) | |
| in_collision = True | |
| if default_position is None: | |
| default_position = [ | |
| self.default_start_pose[0][0], | |
| self.default_start_pose[0][1], | |
| -move_resolution, | |
| ] | |
| joint_pos = self.get_actuated_joint_positions( | |
| actuated_joint_names=self.actuated_joint_names | |
| ).copy() | |
| # print(f"{self.__class__.__name__}: Trying to place robot on the ground.") | |
| while in_collision: | |
| default_position[2] += move_resolution | |
| self.reset_base_pose( | |
| position=default_position, orientation=self.default_start_pose[1] | |
| ) | |
| self.reset_actuated_joint_positions( | |
| joint_positions=joint_pos, joint_names=self.actuated_joint_names | |
| ) | |
| pb.stepSimulation(physicsClientId=self.cid) | |
| in_collision = np.any(self.get_contact_states_of_links(self.link_ids)) | |
| # print(f"{self.__class__.__name__}: Robot should be just above the ground now.") | |
| self.default_start_pose[0] = np.array(default_position) | |
| def set_position_control_mode(self): | |
| """Set control mode to position-based (robot will stay in place if no commands sent).""" | |
| pb.setJointMotorControlArray( | |
| self.robot_id, | |
| self.joint_ids, | |
| pb.POSITION_CONTROL, | |
| targetPositions=self.joint_neutral_positions, | |
| physicsClientId=self.cid, | |
| ) | |
| self._in_torque_mode = False | |
| print("SimRobot position control mode enabled!") | |
| def get_physics_parameters(self) -> Dict[str, Any]: | |
| return pb.getPhysicsEngineParameters(physicsClientId=self.cid) | |
| def set_torque_control_mode(self): | |
| """Set control mode to torque-based (robot will NOT stay in place if no commands sent).""" | |
| pb.setJointMotorControlArray( | |
| self.robot_id, | |
| self.actuated_joint_ids, | |
| pb.VELOCITY_CONTROL, | |
| forces=[0.0] * self.num_of_actuated_joints, | |
| physicsClientId=self.cid, | |
| ) | |
| self._in_torque_mode = True | |
| print("SimRobot torque control mode enabled!") | |
| def reset_base_pose(self, position: Vector3D = None, orientation: QuatType = None): | |
| """Reset the robot's base pose to default pose. | |
| Args: | |
| position (Vector3D, optional): If specified, sets robot to this position. Defaults to None. | |
| orientation (QuatType, optional): If specified, sets robot to this orientation. | |
| Quaternion order: x,y,z,w. Defaults to None. | |
| """ | |
| if position is None: | |
| position = self.default_start_pose[0] | |
| if orientation is None: | |
| orientation = self.default_start_pose[1] | |
| pb.resetBasePositionAndOrientation( | |
| self.robot_id, position, orientation, physicsClientId=self.cid | |
| ) | |
| def reset_actuated_joint_positions( | |
| self, joint_positions: np.ndarray = None, joint_names: List[str] = None | |
| ): | |
| """Reset the joint positions of the robot. | |
| Args: | |
| joint_positions (np.ndarray, optional): If specified, sets joints to these values, otherwise | |
| uses default. Defaults to None. | |
| joint_names (List[str], optional): Joint names in the order the position values were given. | |
| If None provided, uses default order of self.actuated_joint_names. | |
| """ | |
| if joint_positions is None: | |
| joint_positions = self.default_joint_positions | |
| if joint_names is None: | |
| joint_names = self.actuated_joint_names | |
| for jointName, jointPosition in zip(joint_names, joint_positions): | |
| pb.resetJointState( | |
| self.robot_id, | |
| self.get_joint_id(jointName), | |
| jointPosition, | |
| physicsClientId=self.cid, | |
| ) | |
| def get_object_id(self) -> int: | |
| return self.robot_id | |
| def get_base_link_name(self) -> str: | |
| return pb.getBodyInfo(self.robot_id, physicsClientId=self.cid)[0].decode() | |
| def get_joint_name(self, joint_id: int) -> str: | |
| return pb.getJointInfo(self.robot_id, joint_id, physicsClientId=self.cid)[ | |
| 1 | |
| ].decode() | |
| def get_joint_id(self, joint_name: str) -> int: | |
| return self.joint_name_to_index[joint_name] | |
| def get_joint_ids(self, joint_names: List[str]) -> List[int]: | |
| return [self.joint_name_to_index[joint_name] for joint_name in joint_names] | |
| def get_link_name(self, link_id: int) -> str: | |
| if link_id == -1: | |
| return self.get_base_link_name() | |
| else: | |
| return pb.getJointInfo(self.robot_id, link_id, physicsClientId=self.cid)[ | |
| 12 | |
| ].decode() | |
| def get_link_index(self, link_name: str) -> int: | |
| return self.link_name_to_index[link_name] | |
| def get_joint_type(self, joint_id: int) -> str: | |
| return pb.getJointInfo(self.robot_id, joint_id, physicsClientId=self.cid)[2] | |
| def get_joint_damping(self, joint_id: int) -> float: | |
| return pb.getJointInfo(self.robot_id, joint_id, physicsClientId=self.cid)[6] | |
| def get_joint_friction(self, joint_id: int) -> float: | |
| return pb.getJointInfo(self.robot_id, joint_id, physicsClientId=self.cid)[7] | |
| def get_joint_lower_limit_position(self, joint_id: int) -> float: | |
| return pb.getJointInfo(self.robot_id, joint_id, physicsClientId=self.cid)[8] | |
| def get_joint_upper_limit_position(self, joint_id: int) -> float: | |
| return pb.getJointInfo(self.robot_id, joint_id, physicsClientId=self.cid)[9] | |
| def get_joint_max_force(self, joint_id: int) -> float: | |
| return pb.getJointInfo(self.robot_id, joint_id, physicsClientId=self.cid)[10] | |
| def get_joint_max_velocity(self, joint_id: int) -> float: | |
| return pb.getJointInfo(self.robot_id, joint_id, physicsClientId=self.cid)[11] | |
| def get_link_mass(self, link_id: int) -> float: | |
| return pb.getDynamicsInfo(self.robot_id, link_id, physicsClientId=self.cid)[0] | |
| def get_link_local_inertial_transform( | |
| self, link_id: int | |
| ) -> Tuple[np.ndarray, QuatType]: | |
| return pb.getDynamicsInfo(self.robot_id, link_id, physicsClientId=self.cid)[3:5] | |
| def get_link_local_inertial_position(self, link_id: int) -> np.ndarray: | |
| return np.array( | |
| pb.getDynamicsInfo(self.robot_id, link_id, physicsClientId=self.cid)[3] | |
| ) | |
| def get_link_local_inertial_quaternion(self, link_id: int) -> QuatType: | |
| return np.array( | |
| pb.getDynamicsInfo(self.robot_id, link_id, physicsClientId=self.cid)[4] | |
| ) | |
| # Get Base states | |
| def get_base_local_inertia_transform(self) -> Tuple[Vector3D, QuatType]: | |
| return self.get_link_local_inertial_transform(-1) | |
| def get_base_local_inertia_position(self) -> Vector3D: | |
| return self.get_link_local_inertial_position(-1) | |
| def get_base_local_inertia_quaternion(self) -> QuatType: | |
| return self.get_link_local_inertial_quaternion(-1) | |
| def get_base_com_position(self) -> Vector3D: | |
| worldTransCom = pb.getBasePositionAndOrientation( | |
| self.robot_id, physicsClientId=self.cid | |
| ) | |
| localTransCom = self.get_base_local_inertia_transform() | |
| comTransLocal = pb.invertTransform( | |
| position=localTransCom[0], | |
| orientation=localTransCom[1], | |
| physicsClientId=self.cid, | |
| ) | |
| worldTransLocal = pb.multiplyTransforms( | |
| positionA=worldTransCom[0], | |
| orientationA=worldTransCom[1], | |
| positionB=comTransLocal[0], | |
| orientationB=comTransLocal[1], | |
| physicsClientId=self.cid, | |
| ) | |
| return np.array(worldTransLocal[0]) | |
| def get_base_pose(self) -> Tuple[Vector3D, QuatType]: | |
| p, q = pb.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.cid) | |
| return np.array(p), np.array(q) | |
| def get_base_velocity(self) -> Tuple[Vector3D, Vector3D]: | |
| lin, ang = pb.getBaseVelocity(self.robot_id, physicsClientId=self.cid) | |
| return np.array(lin), np.array(ang) | |
| # Get actuated joint states | |
| def get_actuated_joint_positions( | |
| self, actuated_joint_names: List[str] | |
| ) -> np.ndarray: | |
| return np.array( | |
| [ | |
| joint_state[0] | |
| for joint_state in pb.getJointStates( | |
| self.robot_id, | |
| self.get_joint_ids(actuated_joint_names), | |
| physicsClientId=self.cid, | |
| ) | |
| ] | |
| ) | |
| def get_actuated_joint_velocities( | |
| self, actuated_joint_names: List[str] | |
| ) -> np.ndarray: | |
| return np.array( | |
| [ | |
| joint_state[1] | |
| for joint_state in pb.getJointStates( | |
| self.robot_id, | |
| self.get_joint_ids(actuated_joint_names), | |
| physicsClientId=self.cid, | |
| ) | |
| ] | |
| ) | |
| def get_actuated_joint_torques(self, actuated_joint_names: List[str]) -> np.ndarray: | |
| return np.array( | |
| [ | |
| joint_state[3] | |
| for joint_state in pb.getJointStates( | |
| self.robot_id, | |
| self.get_joint_ids(actuated_joint_names), | |
| physicsClientId=self.cid, | |
| ) | |
| ] | |
| ) | |
| def get_actuated_joint_name_to_position(self) -> Dict[str, float]: | |
| return dict( | |
| zip( | |
| self.actuated_joint_names, | |
| [ | |
| joint_state[0] | |
| for joint_state in pb.getJointStates( | |
| self.robot_id, self.actuated_joint_ids, physicsClientId=self.cid | |
| ) | |
| ], | |
| ) | |
| ) | |
| def get_actuated_joint_name_to_velocity(self) -> Dict[str, float]: | |
| return dict( | |
| zip( | |
| self.actuated_joint_names, | |
| [ | |
| joint_state[1] | |
| for joint_state in pb.getJointStates( | |
| self.robot_id, self.actuated_joint_ids, physicsClientId=self.cid | |
| ) | |
| ], | |
| ) | |
| ) | |
| def get_actuated_joint_name_to_torque(self) -> Dict[str, float]: | |
| return dict( | |
| zip( | |
| self.actuated_joint_names, | |
| [ | |
| joint_state[3] | |
| for joint_state in pb.getJointStates( | |
| self.robot_id, self.actuated_joint_ids, physicsClientId=self.cid | |
| ) | |
| ], | |
| ) | |
| ) | |
| def get_link_pose(self, link_id: int) -> Tuple[Vector3D, QuatType]: | |
| if link_id == -1: | |
| return self.get_base_pose() | |
| p, o = pb.getLinkState(self.robot_id, link_id, physicsClientId=self.cid)[4:6] | |
| return np.array(p), np.array(o) | |
| def get_ee_contact_states(self, ee_names: List[str] = None) -> np.ndarray: | |
| if ee_names is None: | |
| ee_names = self.ee_names | |
| return self.get_contact_states_of_links( | |
| link_ids=[self.get_link_index(link_name=name) for name in ee_names] | |
| ) | |
| def get_contact_states_of_links(self, link_ids: List[int]) -> np.ndarray: | |
| return np.array( | |
| [ | |
| int( | |
| len( | |
| pb.getContactPoints( | |
| bodyA=self.robot_id, | |
| linkIndexA=idx, | |
| physicsClientId=self.cid, | |
| ) | |
| ) | |
| > 0 | |
| ) | |
| for idx in link_ids | |
| ] | |
| ) | |
| # Get robot states | |
| def get_robot_states( | |
| self, actuated_joint_names: List[str] = None, ee_names: List[str] = None | |
| ) -> Dict[str, Any]: | |
| if actuated_joint_names is None: | |
| actuated_joint_names = self.actuated_joint_names | |
| joint_ids = self.actuated_joint_ids | |
| else: | |
| joint_ids = self.get_joint_ids(joint_names=actuated_joint_names) | |
| if ee_names is None: | |
| ee_ids = self.ee_ids | |
| ee_names = self.ee_names | |
| else: | |
| ee_ids = [self.get_link_index(link_name=name) for name in ee_names] | |
| joint_states = self.get_joint_states(joint_ids=joint_ids) | |
| base_pose = self.get_base_pose() | |
| base_vel = self.get_base_velocity() | |
| return { | |
| "base_position": base_pose[0], | |
| "base_com_position": self.get_base_com_position(), | |
| "base_quaternion": base_pose[1], | |
| "base_velocity_linear": base_vel[0], | |
| "base_velocity_angular": base_vel[1], | |
| "actuated_joint_positions": joint_states[0], | |
| "actuated_joint_velocities": joint_states[1], | |
| "actuated_joint_torques": joint_states[2], | |
| "joint_order": actuated_joint_names, | |
| "ee_order": ee_names, | |
| "ee_contact_states": self.get_contact_states_of_links(ee_ids), | |
| } | |
| def get_joint_states( | |
| self, joint_ids: List[int] = None | |
| ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: | |
| """Get joint positions, velocities and torques. | |
| Args: | |
| joint_ids (List[int], Optional): List of joint ids (optional) | |
| Returns: | |
| Tuple[np.ndarray, np.ndarray, np.ndarray]: joint positions, joint velocities, joint torques | |
| """ | |
| if joint_ids is None: | |
| joint_ids = self.actuated_joint_ids | |
| dof = len(joint_ids) | |
| q, v, tau = np.zeros(dof), np.zeros(dof), np.zeros(dof) | |
| for n, idx in enumerate(joint_ids): | |
| q[n], v[n], _, tau[n] = pb.getJointState( | |
| self.robot_id, idx, physicsClientId=self.cid | |
| ) | |
| return q, v, tau | |
| def compute_joint_error( | |
| self, target_q: np.ndarray, joint_names_order: List[str] | |
| ) -> np.ndarray: | |
| if not self.has_continuous_joints: | |
| return target_q - self.get_actuated_joint_positions(joint_names_order) | |
| curr_q = self.get_actuated_joint_positions(joint_names_order) | |
| err = target_q - curr_q | |
| for n, joint_name in enumerate(joint_names_order): | |
| if joint_name in self.continuous_joint_names: | |
| err[n] = wrap_angle(target_q[n] - curr_q[n]) | |
| return err | |
| def set_actuated_joint_commands( | |
| self, | |
| actuated_joint_names: List[str] = None, | |
| q: float | np.ndarray = 0, | |
| Kp: float | np.ndarray = 0, | |
| dq: float | np.ndarray = 0, | |
| Kd: float | np.ndarray = 0, | |
| tau: float | np.ndarray = 0, | |
| ): | |
| # can be further optimised for speed by storing last used actuated joint names, e.g. | |
| # if using in a control loop | |
| if actuated_joint_names is None: | |
| actuated_joint_names = self.actuated_joint_names | |
| if self._in_torque_mode: | |
| tau_cmd = ( | |
| Kp | |
| * self.compute_joint_error( | |
| target_q=q, joint_names_order=actuated_joint_names | |
| ) | |
| + Kd * (dq - self.get_actuated_joint_velocities(actuated_joint_names)) | |
| + tau | |
| ) | |
| pb.setJointMotorControlArray( | |
| self.robot_id, | |
| self.get_joint_ids(actuated_joint_names), | |
| pb.TORQUE_CONTROL, | |
| forces=tau_cmd, | |
| physicsClientId=self.cid, | |
| ) | |
| else: | |
| tau = [50] * len(actuated_joint_names) | |
| pb.setJointMotorControlArray( | |
| self.robot_id, | |
| self.get_joint_ids(actuated_joint_names), | |
| pb.POSITION_CONTROL, | |
| targetPositions=q, | |
| targetVelocities=dq, | |
| forces=tau, | |
| positionGains=Kp, | |
| velocityGains=Kd, | |
| physicsClientId=self.cid, | |
| ) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment