def __init__(self, xyz_dims=3, xyz_remap=[0, 1, 2], xyz_scale=[1, 1, 1]): """TODO: fill in other params""" from robosuite.devices import SpaceMouse self.xyz_dims = xyz_dims self.xyz_remap = np.array(xyz_remap) self.xyz_scale = np.array(xyz_scale) self.device = SpaceMouse()
class SpaceMouseExpert(Expert): def __init__(self, xyz_dims=3, xyz_remap=[0, 1, 2], xyz_scale=[1, 1, 1]): """TODO: fill in other params""" from robosuite.devices import SpaceMouse self.xyz_dims = xyz_dims self.xyz_remap = np.array(xyz_remap) self.xyz_scale = np.array(xyz_scale) self.device = SpaceMouse() def get_action(self, obs): """Must return (action, valid, reset, accept)""" state = self.device.get_controller_state() dpos, rotation, accept, reset = ( state["dpos"], state["rotation"], state["left_click"], state["right_click"], ) xyz = dpos[self.xyz_remap] * self.xyz_scale a = xyz[:self.xyz_dims] valid = not np.all(np.isclose(a, 0)) return (a, valid, reset, accept)
# Setup printing options for numbers np.set_printoptions(formatter={"float": lambda x: "{0:0.3f}".format(x)}) # initialize device if args.device == "keyboard": from robosuite.devices import Keyboard device = Keyboard(pos_sensitivity=args.pos_sensitivity, rot_sensitivity=args.rot_sensitivity) 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 args.device == "spacemouse": from robosuite.devices import SpaceMouse device = SpaceMouse(pos_sensitivity=args.pos_sensitivity, rot_sensitivity=args.rot_sensitivity) else: 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
tmp_directory = "~/Robotics/{}".format(str(time.time()).replace( ".", "_")) # Change from temp to Robotics folder env = DataCollectionWrapperBaseline(env, tmp_directory) # initialize device if args.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 args.device == "spacemouse": from robosuite.devices import SpaceMouse device = SpaceMouse() else: raise Exception( "Invalid device choice: choose either 'keyboard' or 'spacemouse'.") # make a new timestamped directory t1, t2 = str(time.time()).split(".") new_dir = os.path.join(args.directory, "{}_{}".format(t1, t2)) os.makedirs(new_dir) # Setup network # ---------------------------------------- policy_hidden_size = 100 # matches default def policy_fn(name, ob_space, ac_space, reuse=False): return mlp_policy_sawyer.MlpPolicy(name=name,
from multiworld.envs.mujoco.sawyer_xyz.sawyer_hammer_6dof import SawyerHammer6DOFEnv from multiworld.envs.mujoco.sawyer_xyz.sawyer_box_open_6dof import SawyerBoxOpen6DOFEnv from multiworld.envs.mujoco.sawyer_xyz.sawyer_bin_picking_6dof import SawyerBinPicking6DOFEnv from multiworld.envs.mujoco.sawyer_xyz.sawyer_peg_insertion_side_6dof import SawyerPegInsertionSide6DOFEnv from multiworld.envs.mujoco.sawyer_xyz.sawyer_peg_insertion_topdown_6dof import SawyerPegInsertionTopdown6DOFEnv from multiworld.envs.mujoco.sawyer_xyz.sawyer_peg_unplug_side_6dof import SawyerPegUnplugSide6DOFEnv from robosuite.devices import SpaceMouse from multiworld.envs.mujoco.utils import rotation from robosuite.utils.transform_utils import mat2quat from multiworld.envs.env_util import quat_to_zangle, zangle_to_quat import gym import multiworld space_mouse = SpaceMouse() env = SawyerPegUnplugSide6DOFEnv(random_init=True, obs_type='with_goal') NDIM = env.action_space.low.size lock_action = False obs = env.reset() action = np.zeros(10) closed = False while True: done = False env.render() state = space_mouse.get_controller_state() dpos, rotation, grasp, reset = ( state["dpos"], state["rotation"],
""" Should be run on a machine connected to a spacemouse """ import time import Pyro4 from robosuite.devices import SpaceMouse from rlkit.launchers import config Pyro4.config.SERIALIZERS_ACCEPTED = set( ["pickle", "json", "marshal", "serpent"]) Pyro4.config.SERIALIZER = "pickle" nameserver = Pyro4.locateNS(host=config.SPACEMOUSE_HOSTNAME) uri = nameserver.lookup("example.greeting") device_state = Pyro4.Proxy(uri) device = SpaceMouse() while True: state = device.get_controller_state() print(state) time.sleep(0.1) device_state.set_state(state)
# wrap the environment with data collection wrapper tmp_directory = "/tmp/{}".format(str(time.time()).replace(".", "_")) env = DataCollectionWrapper(env, tmp_directory) # initialize devices if args.device1 == "keyboard": from robosuite.devices import Keyboard device1 = Keyboard() env.viewer.add_keypress_callback("any", device1.on_press) env.viewer.add_keyup_callback("any", device1.on_release) env.viewer.add_keyrepeat_callback("any", device1.on_press) elif args.device1 == "spacemouse": from robosuite.devices import SpaceMouse device1 = SpaceMouse() else: raise Exception( "Invalid device choice: choose either 'keyboard' or 'spacemouse'.") if args.device2 == "keyboard": from robosuite.devices import Keyboard device2 = Keyboard() env.viewer.add_keypress_callback("any", device2.on_press) env.viewer.add_keyup_callback("any", device2.on_release) env.viewer.add_keyrepeat_callback("any", device2.on_press) elif args.device2 == "spacemouse": from robosuite.devices import SpaceMouse device2 = SpaceMouse()
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()