class ik_helper:
def __init__(self, jpos, quat):
self.jpos = jpos
self.quat = quat
self.controller = SawyerIKController(
bullet_data_path=os.path.join(robosuite.models.assets_root,
"bullet_data"),
robot_jpos_getter=self._robot_jpos_getter,
)
def _robot_jpos_getter(self):
return self.jpos
def xyz_to_jvel(self, xyz_action, jpos, quat):
self.jpos = jpos
self.quat = quat
constant_quat = np.array([0, 0, 0, 1])
input_1 = self._make_input(
np.concatenate((xyz_action[:3], constant_quat)), self.quat)
self.controller.sync_ik_robot(self._robot_jpos_getter(), simulate=True)
velocities = self.controller.get_control(**input_1)
gripper_action = np.ones(1)
low_action = np.concatenate([velocities, gripper_action])
return low_action
def _make_input(self, action, old_quat):
"""
Helper function that returns a dictionary with keys dpos, rotation from a raw input
array. The first three elements are taken to be displacement in position, and a
quaternion indicating the change in rotation with respect to @old_quat.
"""
return {
"dpos": action[:3],
# IK controller takes an absolute orientation in robot base frame
"rotation": T.quat2mat(T.quat_multiply(old_quat, action[3:7])),
}
class inverse_kin:
def __init__(self, env):
self.env = env.wrapped_env.env
self.controller = SawyerIKController(
bullet_data_path=os.path.join(robosuite.models.assets_root,
"bullet_data"),
robot_jpos_getter=self._robot_jpos_getter,
)
def _robot_jpos_getter(self):
"""
Helper function to pass to the ik controller for access to the
current robot joint positions.
"""
return np.array(self.env._joint_positions)
def xyz_to_jvel(self, xyz_action):
constant_quat = np.array([0, 0, 0, 1])
input_1 = self._make_input(
np.concatenate((xyz_action[:3], constant_quat)),
self.env._right_hand_quat)
velocities = self.controller.get_control(**input_1)
gripper_action = np.ones(1)
low_action = np.concatenate([velocities, gripper_action])
return low_action
def _make_input(self, action, old_quat):
"""
Helper function that returns a dictionary with keys dpos, rotation from a raw input
array. The first three elements are taken to be displacement in position, and a
quaternion indicating the change in rotation with respect to @old_quat.
"""
return {
"dpos": action[:3],
# IK controller takes an absolute orientation in robot base frame
"rotation": T.quat2mat(T.quat_multiply(old_quat, action[3:7])),
}
class IKWrapper(Wrapper):
env = None
def __init__(self, env, action_repeat=1, gripper=False):
"""
Initializes the inverse kinematics wrapper.
This wrapper allows for controlling the robot through end effector
movements instead of joint velocities.
Args:
env (MujocoEnv instance): The environment to wrap.
action_repeat (int): Determines the number of times low-level joint
control actions will be commanded per high-level end effector
action. Higher values will allow for more precise control of
the end effector to the commanded targets.
"""
super().__init__(env)
if self.env.mujoco_robot.name == "sawyer":
self.controller = SawyerIKController(
bullet_data_path=os.path.join(robosuite.models.assets_root,
"bullet_data"),
robot_jpos_getter=self._robot_jpos_getter,
)
else:
raise Exception("Only Sawyer robot environments supported for IK "
"control currently.")
self.action_repeat = action_repeat
self.gripper = gripper
if self.gripper:
low = -np.array([.1, .1, .1, 1])
high = np.array([.1, .1, .1, 1])
self.action_space = spaces.Box(low=low, high=high)
else:
low = -np.array([.1, .1, .1])
high = np.array([.1, .1, .1])
self.action_space = spaces.Box(low=low, high=high)
self.action_spec = low, high
def set_robot_joint_positions(self, positions):
"""
Overrides the function to set the joint positions directly, since we need to notify
the IK controller of the change.
"""
self.env.set_robot_joint_positions(positions)
self.controller.sync_state()
def _robot_jpos_getter(self):
"""
Helper function to pass to the ik controller for access to the
current robot joint positions.
"""
return np.array(self.env._joint_positions)
def reset(self):
ret = super().reset()
self.controller.sync_state()
return ret
def step(self, action):
"""
Move the end effector(s) according to the input control.
Args:
action (numpy array): The array should have the corresponding elements.
0-2: The desired change in end effector position in x, y, and z.
(now constant orientation and gripper)
previously:
3-6: The desired change in orientation, expressed as a (x, y, z, w) quaternion.
Note that this quaternion encodes a relative rotation with respect to the
current gripper orientation. If the current rotation is r, this corresponds
to a quaternion d such that r * d will be the new rotation.
*: Controls for gripper actuation.
"""
if self.gripper:
gripper_action = action[3:]
else:
gripper_action = np.ones(1)
constant_quat = np.array([0, 0, 0, 1])
input_1 = self._make_input(np.concatenate((action[:3], constant_quat)),
self.env._right_hand_quat)
if self.env.mujoco_robot.name == "sawyer":
velocities = self.controller.get_control(**input_1)
low_action = np.concatenate([velocities, gripper_action])
else:
raise Exception("Only Sawyer robot environments supported for IK "
"control currently.")
# keep trying to reach the target in a closed-loop
for i in range(self.action_repeat):
ret = self.env.step(low_action)
if i + 1 < self.action_repeat:
velocities = self.controller.get_control()
if self.env.mujoco_robot.name == "sawyer":
low_action = np.concatenate([velocities, action[7:]])
else:
raise Exception(
"Only Sawyer robot environments supported for IK "
"control currently.")
return ret
def _make_input(self, action, old_quat):
"""
Helper function that returns a dictionary with keys dpos, rotation from a raw input
array. The first three elements are taken to be displacement in position, and a
quaternion indicating the change in rotation with respect to @old_quat.
"""
return {
"dpos": action[:3],
# IK controller takes an absolute orientation in robot base frame
"rotation": T.quat2mat(T.quat_multiply(old_quat, action[3:7])),
}