Source code for pyb_utils.robots

"""This module provides a simple wrapper for a velocity- or torque-controlled
robot."""
import numpy as np
import pybullet as pyb

from .named_tuples import getJointInfo, getJointStates, getLinkState
from .math import quaternion_to_matrix


[docs] class Robot: """Wrapper for a PyBullet robot. Attributes ---------- uid : int The unique ID of the underlying PyBullet robot body. client_id : int The unique ID of the physics server this robot belongs to. tool_idx : int The index of the tool joint/link (in PyBullet each link has the same index as its parent joint). Parameters ---------- uid : int The UID of the body representing the robot. tool_link_name : str If provided, use this link as the end effector. If it is ``None``, then the last link is used. The index of this link is stored in ``self.tool_idx``. actuated_joint_names : iterable An optional list of actuated joint names. These joints will be "actuated", in that they take commands from the ``command_`` methods. If not provided, all moveable joints are considered actuated. To have a free-moving joint (whether actuated or not), see :meth:`set_joint_friction_forces`. client_id : int The physics server UID to connect to, if multiple servers are being used. Raises ------ ValueError If ``tool_link_name`` is provided but the robot has no link with that name. """ def __init__( self, uid, tool_link_name=None, actuated_joint_names=None, client_id=0, ): self.uid = uid self.client_id = client_id n = pyb.getNumJoints(uid, physicsClientId=client_id) # build mappings of joint and link names to their respective indices # for look up later # also record the indices of the moveable joints self._link_index_map = {} self._joint_index_map = {} self._moveable_joint_indices = [] for i in range(n): info = getJointInfo( uid, i, decode="utf-8", physicsClientId=client_id ) self._joint_index_map[info.jointName] = i self._link_index_map[info.linkName] = i if info.jointType != pyb.JOINT_FIXED: self._moveable_joint_indices.append(i) if actuated_joint_names is None: self._actuated_joint_indices = self._moveable_joint_indices else: self._actuated_joint_indices = [ self._joint_index_map[name] for name in actuated_joint_names ] if tool_link_name is None: self.tool_idx = n - 1 else: self.tool_idx = self.get_link_index(tool_link_name) @property def num_total_joints(self): """int: The total number of joints, including fixed and moveable joints.""" return len(self._joint_index_map) @property def num_moveable_joints(self): """int: The number of moveable joints.""" return len(self._moveable_joint_indices) @property def num_actuated_joints(self): """int: The number of actuated (i.e., controlled) joints. This should never be more than `num_joints`.""" return len(self._actuated_joint_indices) @property def link_names(self): """list: A list of all link names.""" return list(self._link_index_map.keys()) @property def all_joint_names(self): """list: A list of all joint names, including fixed and moveable joints.""" return list(self._joint_index_map.keys()) @property def moveable_joint_names(self): """list: A list of the names of the moveable joints.""" return [self.all_joint_names[i] for i in self._moveable_joint_indices] @property def actuated_joint_names(self): """list: A list of the actuated joint names. This is the same as the list of joint names unless ``actuated_joint_names`` is provided to the constructor.""" return [self.all_joint_names[i] for i in self._actuated_joint_indices]
[docs] def get_joint_index(self, name): """Get the index of the joint named ``name``. Parameters ---------- name : str The name of the joint. Returns ------- : The index of the joint. Raises ------ ValueError If the robot has no joint named ``name``. """ try: return self._joint_index_map[name] except KeyError: raise ValueError(f"Robot has no joint named {name}.")
[docs] def get_joint_states(self): """Get position and velocity of the robot's moveable joints. Returns ------- : A tuple ``(q, v)`` where ``q`` is the robot's joint configuration and ``v`` is the joint velocity. """ states = getJointStates( self.uid, self._moveable_joint_indices, physicsClientId=self.client_id, ) q = np.array([state.jointPosition for state in states]) v = np.array([state.jointVelocity for state in states]) return q, v
[docs] def reset_joint_configuration(self, q): """Reset the robot to a particular configuration. It is best not to do this during a simulation, as it overrides all dynamic effects. Parameters ---------- q : iterable The vector of joint values to set. Must be of length ``self.num_moveable_joints``. """ for idx, angle in zip(self._moveable_joint_indices, q): pyb.resetJointState( self.uid, idx, angle, physicsClientId=self.client_id )
[docs] def set_joint_friction_forces(self, forces, joint_indices=None): """Set the friction forces at each joint. Set the force to a small value (or zero) to have free-moving joints. Parameters ---------- forces : iterable The friction forces for the specified joints. joint_indices : iterable An optional list of joint indices to specify forces for. If not provided, then it is assumed forces are being provided for all moveable joints. """ if joint_indices is None: joint_indices = self._moveable_joint_indices assert len(forces) == len( joint_indices ), f"Number of friction forces does not match number of joints." pyb.setJointMotorControlArray( self.uid, joint_indices, controlMode=pyb.VELOCITY_CONTROL, targetVelocities=list(np.zeros_like(forces)), forces=forces, physicsClientId=self.client_id, )
[docs] def command_velocity(self, u): """Send a joint velocity command to the robot. Parameters ---------- u : iterable The joint velocity to command. Must be of length ``self.num_actuated_joints``. """ assert ( len(u) == self.num_actuated_joints ), f"Command length does not match number of actuated joints." pyb.setJointMotorControlArray( self.uid, self._actuated_joint_indices, controlMode=pyb.VELOCITY_CONTROL, targetVelocities=list(u), physicsClientId=self.client_id, )
[docs] def command_effort(self, u): """Send a joint effort (force/torque) command to the robot. Parameters ---------- u : iterable The joint effort to command. Must be of length ``self.num_actuated_joints``. """ assert ( len(u) == self.num_actuated_joints ), f"Command length does not match number of actuated joints." pyb.setJointMotorControlArray( self.uid, self._actuated_joint_indices, controlMode=pyb.TORQUE_CONTROL, forces=list(u), physicsClientId=self.client_id, )