def collect_human_trajectory(env, device, arm, env_configuration): """ Use the device (keyboard or SpaceNav 3D mouse) to collect a demonstration. The rollout trajectory is saved to files in npz format. Modify the DataCollectionWrapper wrapper to add new fields or change data formats. Args: env (MujocoEnv): environment to control device (Device): to receive controls from the device arms (str): which arm to control (eg bimanual) 'right' or 'left' env_configuration (str): specified environment configuration """ env.reset() # ID = 2 always corresponds to agentview env.render() is_first = True task_completion_hold_count = -1 # counter to collect 10 timesteps after reaching goal device.start_control() # Loop until we get a reset from the input or the task completes while True: # Set active robot active_robot = env.robots[0] if env_configuration == "bimanual" else env.robots[arm == "left"] # Get the newest action action, grasp = input2action( device=device, robot=active_robot, active_arm=arm, env_configuration=env_configuration ) # If action is none, then this a reset so we should break if action is None: break # Run environment step env.step(action) env.render() # Also break if we complete the task if task_completion_hold_count == 0: break # state machine to check for having a success for 10 consecutive timesteps if 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 # null the counter if there's no success # cleanup for end of data collection episodes env.close()
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), # toggle arm control and / or camera viewing angle if requested if last_grasp < 0 < grasp: if args.switch_on_grasp: args.arm = "left" if args.arm == "right" else "right" if args.toggle_camera_on_grasp: cam_id = (cam_id + 1) % num_cam env.viewer.set_camera(camera_id=cam_id) # Update last grasp
def collect_human_trajectory(env, device, arm, env_configuration): """ Use the device (keyboard or SpaceNav 3D mouse) to collect a demonstration. The rollout trajectory is saved to files in npz format. Modify the DataCollectionWrapper wrapper to add new fields or change data formats. Args: env (MujocoEnv): environment to control device (Device): to receive controls from the device arms (str): which arm to control (eg bimanual) 'right' or 'left' env_configuration (str): specified environment configuration """ env.reset() # ID = 2 always corresponds to agentview env.render() is_first = True task_completion_hold_count = -1 # counter to collect 10 timesteps after reaching goal device.start_control() # Loop until we get a reset from the input or the task completes while True: # Set active robot active_robot = env.robots[ 0] if env_configuration == "bimanual" else env.robots[arm == "left"] # Get the newest action action, grasp = input2action(device=device, robot=active_robot, active_arm=arm, env_configuration=env_configuration) # If action is none, then this a reset so we should break if action is None: break # Run environment step env.step(action) if is_first: is_first = False # We grab the initial model xml and state and reload from those so that # we can support deterministic playback of actions from our demonstrations. # This is necessary due to rounding issues with the model xml and with # env.sim.forward(). We also have to do this after the first action is # applied because the data collector wrapper only starts recording # after the first action has been played. initial_mjstate = env.sim.get_state().flatten() xml_str = env.sim.model.get_xml() env.reset_from_xml_string(xml_str) env.sim.reset() env.sim.set_state_from_flattened(initial_mjstate) env.sim.forward() env.render() # Also break if we complete the task if task_completion_hold_count == 0: break # state machine to check for having a success for 10 consecutive timesteps if 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 # null the counter if there's no success # cleanup for end of data collection episodes env.close()
def main_osc(controller_kp=[1500, 1500, 50, 150, 150, 150], controller_zeta=[1.5, 1.5, 1, 10, 10, 10], perception_error=0.0, offscreen=False): # Task configuration # option: # board : Hole 12mm, Hole 9mm # peg : 16mm. 12mm, 9mm # USB : USB-C task_config = {'board': 'Square_hole_16mm', 'peg': 'cylinder_16mm'} # IMP_OSC is a custom controller written by us. # find out the source code at https://github.com/garlicbutter/robosuite # Theory based on the paper by Valency: https://doi.org/10.1115/1.1590685. controller_config = { "type": "OSC_POSE", "input_max": 1, "input_min": -1, "output_max": [0.05, 0.05, 0.05, 0.5, 0.5, 0.5], "output_min": [-0.05, -0.05, -0.05, -0.5, -0.5, -0.5], "kp": controller_kp, "damping_ratio": controller_zeta, "impedance_mode": "fixed", "kp_limits": [[0, 300], [0, 300], [0, 300], [0, 300], [0, 300], [0, 300]], "damping_ratio_limits": [0, 10], "position_limits": None, "orientation_limits": None, "uncouple_pos_ori": True, "control_delta": True, "interpolation": None, "ramp_ratio": 0.2 } if offscreen: # create on-screen environment instance env = MyEnvOffScreen(robots="UR5e", task_configs=task_config, controller_configs=controller_config, gripper_types='Robotiq85Gripper', ignore_done=True) else: # create on-screen environment instance env = MyEnv(robots="UR5e", task_configs=task_config, controller_configs=controller_config, gripper_types='Robotiq85Gripper', has_renderer=True, has_offscreen_renderer=False, use_object_obs=True, use_camera_obs=False, ignore_done=True, render_camera=None) # create motion planning class motion_planning = motion_planning_osc.Policy_action(env.control_timestep, P=3, I=0.1) # manual control via keyboard manual_control = False if manual_control: from robosuite.devices import Keyboard from robosuite.utils.input_utils import input2action device = Keyboard(pos_sensitivity=0.05, rot_sensitivity=1.0) 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) # Initial recorder eeff_record = np.empty([1, 3]) # end effector force record robot_torque_record = np.empty([1, 6]) # Actuation torque record eeft_record = np.empty([1, 3]) # end effector torque record eefd_record = np.empty([1, 3]) # end effector displacement record eefd_desire_record = np.empty( [1, 3]) # desired end effector displacement record eefd_desire = np.empty( [1, 3]) # to calculate desired displacement at each moment t_record = np.empty(1) # time record # Initial action action = np.zeros(env.robots[0].dof) action_status = None # simulate termination condition done = False # simulation # for perception_error import random random_radiant = 2 * np.pi * random.random() perception_err_x = np.cos(random_radiant) * perception_error perception_err_y = np.sin(random_radiant) * perception_error while not done: obs, reward, done, _ = env.step( action) # take action in the environment if manual_control: action, grasp = input2action( device=device, robot=env.robots[0], ) else: obs['plate_pos'][0] += perception_err_x obs['plate_pos'][1] += perception_err_y # update observation to motion planning11 motion_planning.update_obs(obs) # decide which action to take for next simulation action, action_status = motion_planning.get_policy_action() if not offscreen: env.render() eeff_record = np.append(eeff_record, [env.robots[0].ee_force], axis=0) eeft_record = np.append(eeft_record, [env.robots[0].ee_torque], axis=0) eefd_record = np.append(eefd_record, [obs['robot0_eef_pos']], axis=0) eefd_desire = np.add(eefd_desire, action[0:3]) eefd_desire_record = np.append(eefd_desire_record, eefd_desire, axis=0) robot_torque_record = np.append(robot_torque_record, [env.robots[0].torques], axis=0) t_record = np.append(t_record, np.array([env.cur_time]), axis=0) if action_status['done']: break if t_record[-1] > 30: # failure case eeff_record = 0 eeft_record = 0 eefd_record = 0 t_record = 0 break return eeff_record, eeft_record, eefd_record, t_record, robot_torque_record
def collect_human_trajectory(env, device, arm, env_configuration): """ Use the device (keyboard or SpaceNav 3D mouse) to collect a demonstration. The rollout trajectory is saved to files in npz format. Modify the DataCollectionWrapper wrapper to add new fields or change data formats. Args: env (MujocoEnv): environment to control device (Device): to receive controls from the device arms (str): which arm to control (eg bimanual) 'right' or 'left' env_configuration (str): specified environment configuration """ env.reset() # ID = 2 always corresponds to agentview env.render() is_first = True task_completion_hold_count = -1 # counter to collect 10 timesteps after reaching goal device.start_control() # create a video writer with imageio #writer = imageio.get_writer(args.video_path, fps=20) #frames = [] # Loop until we get a reset from the input or the task completes while True: # Set active robot active_robot = env.robots[0] if env_configuration == "bimanual" else env.robots[arm == "left"] # Get the newest action action, grasp = input2action( device=device, robot=active_robot, active_arm=arm, env_configuration=env_configuration ) # If action is none, then this a reset so we should break if action is None: break # Run environment step obs, reward, done, info = env.step(action) env.render() #print(obs, action) # read camera observation #im = np.flip(obs[args.camera + "_image"].transpose((1, 0, 2)), 1) #pygame.pixelcopy.array_to_surface(screen, im) #pygame.display.update() #frame = obs[args.camera + "_image"][::-1] #writer.append_data(frame) #print("Saving frame #{}".format(i)) ''' if is_first: is_first = False # We grab the initial model xml and state and reload from those so that # we can support deterministic playback of actions from our demonstrations. # This is necessary due to rounding issues with the model xml and with # env.sim.forward(). We also have to do this after the first action is # applied because the data collector wrapper only starts recording # after the first action has been played. initial_mjstate = env.sim.get_state().flatten() xml_str = env.model.get_xml() env.reset_from_xml_string(xml_str) env.sim.reset() env.sim.set_state_from_flattened(initial_mjstate) env.sim.forward() env.sim.render(512, 512) ''' # Also break if we complete the task if task_completion_hold_count == 0: break # state machine to check for having a success for 10 consecutive timesteps if 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 # null the counter if there's no success
device = Keyboard(pos_sensitivity=0.05, rot_sensitivity=1.0) 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) # Initial action action = np.zeros(env.robots[0].dof) # simulate termination condition done = False # simulation while not done: obs, reward, done, _ = env.step( action) # take action in the environment F_int = env.robots[0].ee_force if manual_control: action, grasp = input2action( device=device, robot=env.robots[0], ) else: # update observation to motion planning motion_planning.update_obs(obs, F_int) # decide which action to take for next simulation action = motion_planning.get_policy_action() os.system('clear') print("Robot: {}, Gripper:{}".format(env.robots[0].name, env.robots[0].gripper_type)) print("Control Frequency:{}".format(env.robots[0].control_freq)) print("eef_force:\n \t x: {a[0]:2.4f}, y: {a[1]:2.4f}, z: {a[2]:2.4f}". format(a=env.robots[0].ee_force)) print( "eef_torque:\n \t x: {a[0]:2.4f}, y: {a[1]:2.4f}, z: {a[2]:2.4f}".