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}".
