def teleoperate(self, demons_config):
env = self.get_env()
# Need to use inverse-kinematics controller to set position using device
env = IKWrapper(env)
if demons_config.device == "keyboard":
from robosuite.devices import Keyboard
device = Keyboard()
env.viewer.add_keypress_callback("any", device.on_press)
env.viewer.add_keyup_callback("any", device.on_release)
env.viewer.add_keyrepeat_callback("any", device.on_press)
elif demons_config.device == "spacemouse":
from robosuite.devices import SpaceMouse
device = SpaceMouse()
for run in range(demons_config.n_runs):
obs = env.reset()
env.set_robot_joint_positions([0, -1.18, 0.00, 2.18, 0.00, 0.57, 1.5708])
# rotate the gripper so we can see it easily - NOTE : REMOVE MAYBE
env.viewer.set_camera(camera_id=2)
env.render()
device.start_control()
reset = False
task_completion_hold_count = -1
step = 0
tr_vobvs, tr_dof, tr_actions = [], [], []
while not reset:
if int(step % demons_config.collect_freq) == 0:
tr_vobvs.append(np.array(obs[self.vis_obv_key]))
tr_dof.append(np.array(obs[self.dof_obv_key].flatten()))
device_state = device.get_controller_state()
dpos, rotation, grasp, reset = (
device_state["dpos"],
device_state["rotation"],
device_state["grasp"],
device_state["reset"],
)
current = env._right_hand_orn
drotation = current.T.dot(rotation)
dquat = T.mat2quat(drotation)
grasp = grasp - 1.
ik_action = np.concatenate([dpos, dquat, [grasp]])
obs, _, done, _ = env.step(ik_action)
env.render()
joint_velocities = np.array(env.controller.commanded_joint_velocities)
if env.env.mujoco_robot.name == "sawyer":
gripper_actuation = np.array(ik_action[7:])
elif env.env.mujoco_robot.name == "baxter":
gripper_actuation = np.array(ik_action[14:])
# NOTE: Action for the normal environment (not inverse kinematic)
action = np.concatenate([joint_velocities, gripper_actuation], axis=0)
if int(step % demons_config.collect_freq) == 0:
tr_actions.append(action)
if (int(step % demons_config.flush_freq) == 0) or (demons_config.break_traj_success and task_completion_hold_count == 0):
print("Storing Trajectory")
trajectory = {self.vis_obv_key : np.array(tr_vobvs), self.dof_obv_key : np.array(tr_dof), 'action' : np.array(tr_actions)}
store_trajectoy(trajectory, 'teleop')
trajectory, tr_vobvs, tr_dof, tr_actions = {}, [], [], []
if demons_config.break_traj_success and env._check_success():
if task_completion_hold_count > 0:
task_completion_hold_count -= 1 # latched state, decrement count
else:
task_completion_hold_count = 10 # reset count on first success timestep
else:
task_completion_hold_count = -1
step += 1
env.close()
else:
raise Exception(
"Invalid device choice: choose either 'keyboard' or 'spacemouse'.")
while True:
obs = env.reset()
env.viewer.set_camera(camera_id=2)
env.render()
# rotate the gripper so we can see it easily
env.set_robot_joint_positions(
[0, -1.18, 0.00, 2.18, 0.00, 0.57, 1.5708])
device.start_control()
while True:
state = device.get_controller_state()
dpos, rotation, grasp, reset = (
state["dpos"],
state["rotation"],
state["grasp"],
state["reset"],
)
if reset:
break
# convert into a suitable end effector action for the environment
current = env._right_hand_orn
# relative rotation of desired from current
drotation = current.T.dot(rotation)
dquat = T.mat2quat(drotation)
env.render()
# Initialize variables that should the maintained between resets
last_grasp = 0
# Initialize device control
device.start_control()
# Set active robot
active_robot = env.robots[
0] if args.config == "bimanual" else env.robots[args.arm == "left"]
# If using the Vive Tracker, adjust origin of the position tracking system such that
# the absolute coordinates agree with the initial pose of the end-effector
if args.device == "vivetracker":
tracker_pos = device.get_controller_state()['dpos']
device.set_origin(tracker_pos)
while True:
# Get the newest action
action, grasp = input2action(device=device,
grasp_device=grasp_device,
robot=active_robot,
active_arm=args.arm,
env_configuration=args.config)
# If action is none, then this a reset so we should break
if action is None:
break
# If the current grasp is active (1) and last grasp is not (-1) (i.e.: grasping input just pressed),