Example #1
0
        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()