raise Exception("Invalid device choice: choose either 'keyboard' or 'spacemouse'.")
while True:
# Reset the environment
obs = env.reset()
# Setup rendering
cam_id = 0
num_cam = len(env.sim.model.camera_names)
env.render()
# Initialize variables that should the maintained between resets
last_grasp = 0
# Initialize device control
device.start_control()
while True:
# Set active robot
active_robot = env.robots[0] if args.config == "bimanual" else env.robots[args.arm == "left"]
# Get the newest action
action, grasp = input2action(
device=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),
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()