import numpy as np import matplotlib.pyplot as plt from pyquaternion import Quaternion from constants import SIM_TASK_CONFIGS from ee_sim_env import make_ee_sim_env import IPython e = IPython.embed class BasePolicy: def __init__(self, inject_noise=False): self.inject_noise = inject_noise self.step_count = 0 self.left_trajectory = None self.right_trajectory = None def generate_trajectory(self, ts_first): raise NotImplementedError @staticmethod def interpolate(curr_waypoint, next_waypoint, t): t_frac = (t - curr_waypoint["t"]) / (next_waypoint["t"] - curr_waypoint["t"]) curr_xyz = curr_waypoint['xyz'] curr_quat = curr_waypoint['quat'] curr_grip = curr_waypoint['gripper'] next_xyz = next_waypoint['xyz'] next_quat = next_waypoint['quat'] next_grip = next_waypoint['gripper'] xyz = curr_xyz + (next_xyz - curr_xyz) * t_frac quat = curr_quat + (next_quat - curr_quat) * t_frac gripper = curr_grip + (next_grip - curr_grip) * t_frac return xyz, quat, gripper def __call__(self, ts): # generate trajectory at first timestep, then open-loop execution if self.step_count == 0: self.generate_trajectory(ts) # obtain left and right waypoints if self.left_trajectory[0]['t'] == self.step_count: self.curr_left_waypoint = self.left_trajectory.pop(0) next_left_waypoint = self.left_trajectory[0] if self.right_trajectory[0]['t'] == self.step_count: self.curr_right_waypoint = self.right_trajectory.pop(0) next_right_waypoint = self.right_trajectory[0] # interpolate between waypoints to obtain current pose and gripper command left_xyz, left_quat, left_gripper = self.interpolate(self.curr_left_waypoint, next_left_waypoint, self.step_count) right_xyz, right_quat, right_gripper = self.interpolate(self.curr_right_waypoint, next_right_waypoint, self.step_count) # Inject noise if self.inject_noise: scale = 0.01 left_xyz = left_xyz + np.random.uniform(-scale, scale, left_xyz.shape) right_xyz = right_xyz + np.random.uniform(-scale, scale, right_xyz.shape) action_left = np.concatenate([left_xyz, left_quat, [left_gripper]]) action_right = np.concatenate([right_xyz, right_quat, [right_gripper]]) self.step_count += 1 return np.concatenate([action_left, action_right]) class PickAndTransferPolicy(BasePolicy): def generate_trajectory(self, ts_first): init_mocap_pose_right = ts_first.observation['mocap_pose_right'] init_mocap_pose_left = ts_first.observation['mocap_pose_left'] box_info = np.array(ts_first.observation['env_state']) box_xyz = box_info[:3] box_quat = box_info[3:] # print(f"Generate trajectory for {box_xyz=}") gripper_pick_quat = Quaternion(init_mocap_pose_right[3:]) gripper_pick_quat = gripper_pick_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=-60) meet_left_quat = Quaternion(axis=[1.0, 0.0, 0.0], degrees=90) meet_xyz = np.array([0, 0.5, 0.25]) self.left_trajectory = [ {"t": 0, "xyz": init_mocap_pose_left[:3], "quat": init_mocap_pose_left[3:], "gripper": 0}, # sleep {"t": 100, "xyz": meet_xyz + np.array([-0.1, 0, -0.02]), "quat": meet_left_quat.elements, "gripper": 1}, # approach meet position {"t": 260, "xyz": meet_xyz + np.array([0.02, 0, -0.02]), "quat": meet_left_quat.elements, "gripper": 1}, # move to meet position {"t": 310, "xyz": meet_xyz + np.array([0.02, 0, -0.02]), "quat": meet_left_quat.elements, "gripper": 0}, # close gripper {"t": 360, "xyz": meet_xyz + np.array([-0.1, 0, -0.02]), "quat": np.array([1, 0, 0, 0]), "gripper": 0}, # move left {"t": 400, "xyz": meet_xyz + np.array([-0.1, 0, -0.02]), "quat": np.array([1, 0, 0, 0]), "gripper": 0}, # stay ] self.right_trajectory = [ {"t": 0, "xyz": init_mocap_pose_right[:3], "quat": init_mocap_pose_right[3:], "gripper": 0}, # sleep {"t": 90, "xyz": box_xyz + np.array([0, 0, 0.08]), "quat": gripper_pick_quat.elements, "gripper": 1}, # approach the cube {"t": 130, "xyz": box_xyz + np.array([0, 0, -0.015]), "quat": gripper_pick_quat.elements, "gripper": 1}, # go down {"t": 170, "xyz": box_xyz + np.array([0, 0, -0.015]), "quat": gripper_pick_quat.elements, "gripper": 0}, # close gripper {"t": 200, "xyz": meet_xyz + np.array([0.05, 0, 0]), "quat": gripper_pick_quat.elements, "gripper": 0}, # approach meet position {"t": 220, "xyz": meet_xyz, "quat": gripper_pick_quat.elements, "gripper": 0}, # move to meet position {"t": 310, "xyz": meet_xyz, "quat": gripper_pick_quat.elements, "gripper": 1}, # open gripper {"t": 360, "xyz": meet_xyz + np.array([0.1, 0, 0]), "quat": gripper_pick_quat.elements, "gripper": 1}, # move to right {"t": 400, "xyz": meet_xyz + np.array([0.1, 0, 0]), "quat": gripper_pick_quat.elements, "gripper": 1}, # stay ] class InsertionPolicy(BasePolicy): def generate_trajectory(self, ts_first): init_mocap_pose_right = ts_first.observation['mocap_pose_right'] init_mocap_pose_left = ts_first.observation['mocap_pose_left'] peg_info = np.array(ts_first.observation['env_state'])[:7] peg_xyz = peg_info[:3] peg_quat = peg_info[3:] socket_info = np.array(ts_first.observation['env_state'])[7:] socket_xyz = socket_info[:3] socket_quat = socket_info[3:] gripper_pick_quat_right = Quaternion(init_mocap_pose_right[3:]) gripper_pick_quat_right = gripper_pick_quat_right * Quaternion(axis=[0.0, 1.0, 0.0], degrees=-60) gripper_pick_quat_left = Quaternion(init_mocap_pose_right[3:]) gripper_pick_quat_left = gripper_pick_quat_left * Quaternion(axis=[0.0, 1.0, 0.0], degrees=60) meet_xyz = np.array([0, 0.5, 0.15]) lift_right = 0.00715 self.left_trajectory = [ {"t": 0, "xyz": init_mocap_pose_left[:3], "quat": init_mocap_pose_left[3:], "gripper": 0}, # sleep {"t": 120, "xyz": socket_xyz + np.array([0, 0, 0.08]), "quat": gripper_pick_quat_left.elements, "gripper": 1}, # approach the cube {"t": 170, "xyz": socket_xyz + np.array([0, 0, -0.03]), "quat": gripper_pick_quat_left.elements, "gripper": 1}, # go down {"t": 220, "xyz": socket_xyz + np.array([0, 0, -0.03]), "quat": gripper_pick_quat_left.elements, "gripper": 0}, # close gripper {"t": 285, "xyz": meet_xyz + np.array([-0.1, 0, 0]), "quat": gripper_pick_quat_left.elements, "gripper": 0}, # approach meet position {"t": 340, "xyz": meet_xyz + np.array([-0.05, 0, 0]), "quat": gripper_pick_quat_left.elements,"gripper": 0}, # insertion {"t": 400, "xyz": meet_xyz + np.array([-0.05, 0, 0]), "quat": gripper_pick_quat_left.elements, "gripper": 0}, # insertion ] self.right_trajectory = [ {"t": 0, "xyz": init_mocap_pose_right[:3], "quat": init_mocap_pose_right[3:], "gripper": 0}, # sleep {"t": 120, "xyz": peg_xyz + np.array([0, 0, 0.08]), "quat": gripper_pick_quat_right.elements, "gripper": 1}, # approach the cube {"t": 170, "xyz": peg_xyz + np.array([0, 0, -0.03]), "quat": gripper_pick_quat_right.elements, "gripper": 1}, # go down {"t": 220, "xyz": peg_xyz + np.array([0, 0, -0.03]), "quat": gripper_pick_quat_right.elements, "gripper": 0}, # close gripper {"t": 285, "xyz": meet_xyz + np.array([0.1, 0, lift_right]), "quat": gripper_pick_quat_right.elements, "gripper": 0}, # approach meet position {"t": 340, "xyz": meet_xyz + np.array([0.05, 0, lift_right]), "quat": gripper_pick_quat_right.elements, "gripper": 0}, # insertion {"t": 400, "xyz": meet_xyz + np.array([0.05, 0, lift_right]), "quat": gripper_pick_quat_right.elements, "gripper": 0}, # insertion ] def test_policy(task_name): # example rolling out pick_and_transfer policy onscreen_render = True inject_noise = False # setup the environment episode_len = SIM_TASK_CONFIGS[task_name]['episode_len'] if 'sim_transfer_cube' in task_name: env = make_ee_sim_env('sim_transfer_cube') elif 'sim_insertion' in task_name: env = make_ee_sim_env('sim_insertion') else: raise NotImplementedError for episode_idx in range(2): ts = env.reset() episode = [ts] if onscreen_render: ax = plt.subplot() plt_img = ax.imshow(ts.observation['images']['angle']) plt.ion() policy = PickAndTransferPolicy(inject_noise) for step in range(episode_len): action = policy(ts) ts = env.step(action) episode.append(ts) if onscreen_render: plt_img.set_data(ts.observation['images']['angle']) plt.pause(0.02) plt.close() episode_return = np.sum([ts.reward for ts in episode[1:]]) if episode_return > 0: print(f"{episode_idx=} Successful, {episode_return=}") else: print(f"{episode_idx=} Failed") if __name__ == '__main__': test_task_name = 'sim_transfer_cube_scripted' test_policy(test_task_name)