gello_software/gello/agents/quest_agent.py

178 lines
7.2 KiB
Python

from typing import Dict
import numpy as np
import quaternion
from dm_control import mjcf
from dm_control.utils.inverse_kinematics import qpos_from_site_pose
from oculus_reader.reader import OculusReader
from gello.agents.agent import Agent
from gello.agents.spacemouse_agent import apply_transfer, mj2ur, ur2mj
from gello.dm_control_tasks.arms.ur5e import UR5e
# cartensian space control, controller <> robot relative pose matters. This extrinsics is based on
# our setup, for details please checkout the project page.
quest2ur = np.array([[-1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
ur2quest = np.linalg.inv(quest2ur)
translation_scaling_factor = 2.0
class SingleArmQuestAgent(Agent):
def __init__(self, robot_type: str, which_hand: str, verbose: bool = False) -> None:
"""Interact with the robot using the quest controller.
leftTrig: press to start control (also record the current position as the home position)
leftJS: a tuple of (x,y) for the joystick, only need y to control the gripper
"""
self.which_hand = which_hand
assert self.which_hand in ["l", "r"]
self.oculus_reader = OculusReader()
if robot_type == "ur5":
_robot = UR5e()
else:
raise ValueError(f"Unknown robot type: {robot_type}")
self.physics = mjcf.Physics.from_mjcf_model(_robot.mjcf_model)
self.control_active = False
self.reference_quest_pose = None
self.reference_ee_rot_ur = None
self.reference_ee_pos_ur = None
self.robot_type = robot_type
self._verbose = verbose
def act(self, obs: Dict[str, np.ndarray]) -> np.ndarray:
if self.robot_type == "ur5":
num_dof = 6
current_qpos = obs["joint_positions"][:num_dof] # last one dim is the gripper
current_gripper_angle = obs["joint_positions"][-1]
# run the fk
self.physics.data.qpos[:num_dof] = current_qpos
self.physics.step()
ee_rot_mj = np.array(
self.physics.named.data.site_xmat["attachment_site"]
).reshape(3, 3)
ee_pos_mj = np.array(self.physics.named.data.site_xpos["attachment_site"])
if self.which_hand == "l":
pose_key = "l"
trigger_key = "leftTrig"
# joystick_key = "leftJS"
# left yx
gripper_open_key = "Y"
gripper_close_key = "X"
elif self.which_hand == "r":
pose_key = "r"
trigger_key = "rightTrig"
# joystick_key = "rightJS"
# right ba for the key
gripper_open_key = "B"
gripper_close_key = "A"
else:
raise ValueError(f"Unknown hand: {self.which_hand}")
# check the trigger button state
pose_data, button_data = self.oculus_reader.get_transformations_and_buttons()
if len(pose_data) == 0 or len(button_data) == 0:
print("no data, quest not yet ready")
return np.concatenate([current_qpos, [current_gripper_angle]])
new_gripper_angle = current_gripper_angle
if button_data[gripper_open_key]:
new_gripper_angle = 1
if button_data[gripper_close_key]:
new_gripper_angle = 0
arm_not_move_return = np.concatenate([current_qpos, [new_gripper_angle]])
if len(pose_data) == 0:
print("no data, quest not yet ready")
return arm_not_move_return
trigger_state = button_data[trigger_key][0]
if trigger_state > 0.5:
if self.control_active is True:
if self._verbose:
print("controlling the arm")
current_pose = pose_data[pose_key]
delta_rot = current_pose[:3, :3] @ np.linalg.inv(
self.reference_quest_pose[:3, :3]
)
delta_pos = current_pose[:3, 3] - self.reference_quest_pose[:3, 3]
delta_pos_ur = (
apply_transfer(quest2ur, delta_pos) * translation_scaling_factor
)
# ? is this the case?
delta_rot_ur = quest2ur[:3, :3] @ delta_rot @ ur2quest[:3, :3]
if self._verbose:
print(
f"delta pos and rot in ur space: \n{delta_pos_ur}, {delta_rot_ur}"
)
next_ee_rot_ur = delta_rot_ur @ self.reference_ee_rot_ur
next_ee_pos_ur = delta_pos_ur + self.reference_ee_pos_ur
target_quat = quaternion.as_float_array(
quaternion.from_rotation_matrix(ur2mj[:3, :3] @ next_ee_rot_ur)
)
ik_result = qpos_from_site_pose(
self.physics,
"attachment_site",
target_pos=apply_transfer(ur2mj, next_ee_pos_ur),
target_quat=target_quat,
tol=1e-14,
max_steps=400,
)
self.physics.reset()
if ik_result.success:
new_qpos = ik_result.qpos[:num_dof]
else:
print("ik failed, using the original qpos")
return arm_not_move_return
command = np.concatenate([new_qpos, [new_gripper_angle]])
return command
else: # last state is not in active
self.control_active = True
if self._verbose:
print("control activated!")
self.reference_quest_pose = pose_data[pose_key]
self.reference_ee_rot_ur = mj2ur[:3, :3] @ ee_rot_mj
self.reference_ee_pos_ur = apply_transfer(mj2ur, ee_pos_mj)
return arm_not_move_return
else:
if self._verbose:
print("deactive control")
self.control_active = False
self.reference_quest_pose = None
return arm_not_move_return
class DualArmQuestAgent(Agent):
def __init__(self, robot_type: str) -> None:
self.left_arm = SingleArmQuestAgent(robot_type, "l")
self.right_arm = SingleArmQuestAgent(robot_type, "r")
def act(self, obs: Dict[str, np.ndarray]) -> np.ndarray:
pass
if __name__ == "__main__":
oculus_reader = OculusReader()
while True:
"""
example output:
({'l': array([[-0.828395 , 0.541667 , -0.142682 , 0.219646 ],
[-0.107737 , 0.0958919, 0.989544 , -0.833478 ],
[ 0.549685 , 0.835106 , -0.0210789, -0.892425 ],
[ 0. , 0. , 0. , 1. ]]), 'r': array([[-0.328058, 0.82021 , 0.468652, -1.8288 ],
[ 0.070887, 0.516083, -0.8536 , -0.238691],
[-0.941994, -0.246809, -0.227447, -0.370447],
[ 0. , 0. , 0. , 1. ]])},
{'A': False, 'B': False, 'RThU': True, 'RJ': False, 'RG': False, 'RTr': False, 'X': False, 'Y': False, 'LThU': True, 'LJ': False, 'LG': False, 'LTr': False, 'leftJS': (0.0, 0.0), 'leftTrig': (0.0,), 'leftGrip': (0.0,), 'rightJS': (0.0, 0.0), 'rightTrig': (0.0,), 'rightGrip': (0.0,)})
"""
pose_data, button_data = oculus_reader.get_transformations_and_buttons()
if len(pose_data) == 0:
print("no data")
continue
else:
print(pose_data["l"])