""" Usage: (umi): python scripts_real/eval_real_umi.py -i data/outputs/2023.10.26/02.25.30_train_diffusion_unet_timm_umi/checkpoints/latest.ckpt -o data_local/cup_test_data ================ Human in control ============== Robot movement: Move your SpaceMouse to move the robot EEF (locked in xy plane). Press SpaceMouse right button to unlock z axis. Press SpaceMouse left button to enable rotation axes. Recording control: Click the opencv window (make sure it's in focus). Press "C" to start evaluation (hand control over to policy). Press "Q" to exit program. ================ Policy in control ============== Make sure you can hit the robot hardware emergency-stop button quickly! Recording control: Press "S" to stop evaluation and gain control back. """ # %% import sys import os ROOT_DIR = os.path.dirname(os.path.dirname(__file__)) sys.path.append(ROOT_DIR) os.chdir(ROOT_DIR) # %% import os import pathlib import time from multiprocessing.managers import SharedMemoryManager from moviepy.editor import VideoFileClip import av import click import cv2 import dill import hydra import numpy as np import scipy.spatial.transform as st import torch from omegaconf import OmegaConf import json from diffusion_policy.common.replay_buffer import ReplayBuffer from diffusion_policy.common.cv2_util import ( get_image_transform ) from umi.common.cv_util import ( parse_fisheye_intrinsics, FisheyeRectConverter ) from diffusion_policy.common.pytorch_util import dict_apply from diffusion_policy.policy.base_image_policy import BaseImagePolicy from diffusion_policy.workspace.base_workspace import BaseWorkspace from umi.common.precise_sleep import precise_wait from umi.real_world.umi_env import UmiEnv from umi.real_world.keystroke_counter import ( KeystrokeCounter, Key, KeyCode ) from umi.real_world.real_inference_util import (get_real_obs_dict, get_real_obs_resolution, get_real_umi_obs_dict, get_real_umi_action) from umi.real_world.spacemouse_shared_memory import Spacemouse OmegaConf.register_new_resolver("eval", eval, replace=True) @click.command() @click.option('--input', '-i', required=True, help='Path to checkpoint') @click.option('--output', '-o', required=True, help='Directory to save recording') # @click.option('--robot_ip', '-ri', default='172.24.95.9') # @click.option('--gripper_ip', '-gi', default='172.24.95.17') @click.option('--robot_ip', default='192.168.0.9') @click.option('--gripper_ip', default='192.168.0.27') # @click.option('--robot_ip', default='172.16.0.3') # @click.option('--gripper_ip', default='172.24.95.27') @click.option('--match_dataset', '-m', default=None, help='Dataset used to overlay and adjust initial condition') @click.option('--match_episode', '-me', default=None, type=int, help='Match specific episode from the match dataset') @click.option('--match_camera', '-mc', default=0, type=int) @click.option('--camera_reorder', '-cr', default='021') @click.option('--vis_camera_idx', default=0, type=int, help="Which RealSense camera to visualize.") @click.option('--init_joints', '-j', is_flag=True, default=False, help="Whether to initialize robot joint configuration in the beginning.") @click.option('--steps_per_inference', '-si', default=6, type=int, help="Action horizon for inference.") @click.option('--max_duration', '-md', default=60, help='Max duration for each epoch in seconds.') @click.option('--frequency', '-f', default=10, type=float, help="Control frequency in Hz.") @click.option('--command_latency', '-cl', default=0.01, type=float, help="Latency between receiving SapceMouse command to executing on Robot in Sec.") @click.option('-nm', '--no_mirror', is_flag=True, default=False) @click.option('-sf', '--sim_fov', type=float, default=None) @click.option('-ci', '--camera_intrinsics', type=str, default=None) @click.option('-rt', '--robot_type', default='ur5') @click.option('--mirror_crop', is_flag=True, default=False) @click.option('--mirror_swap', is_flag=True, default=False) def main(input, output, robot_ip, gripper_ip, match_dataset, match_episode, match_camera, camera_reorder, vis_camera_idx, init_joints, steps_per_inference, max_duration, frequency, command_latency, no_mirror, sim_fov, camera_intrinsics, robot_type, mirror_crop, mirror_swap): max_gripper_width = 0.09 gripper_speed = 0.2 # load checkpoint ckpt_path = input if not ckpt_path.endswith('.ckpt'): ckpt_path = os.path.join(ckpt_path, 'checkpoints', 'latest.ckpt') payload = torch.load(open(ckpt_path, 'rb'), map_location='cpu', pickle_module=dill) cfg = payload['cfg'] print("model_name:", cfg.policy.obs_encoder.model_name) print("dataset_path:", cfg.task.dataset.dataset_path) # setup experiment dt = 1/frequency obs_res = get_real_obs_resolution(cfg.task.shape_meta) # load fisheye converter fisheye_converter = None if sim_fov is not None: assert camera_intrinsics is not None opencv_intr_dict = parse_fisheye_intrinsics( json.load(open(camera_intrinsics, 'r'))) fisheye_converter = FisheyeRectConverter( **opencv_intr_dict, out_size=obs_res, out_fov=sim_fov ) print("steps_per_inference:", steps_per_inference) with SharedMemoryManager() as shm_manager: with Spacemouse(shm_manager=shm_manager) as sm, \ KeystrokeCounter() as key_counter, \ UmiEnv( output_dir=output, robot_ip=robot_ip, gripper_ip=gripper_ip, frequency=frequency, obs_image_resolution=obs_res, obs_float32=True, camera_reorder=[int(x) for x in camera_reorder], init_joints=init_joints, enable_multi_cam_vis=True, # latency camera_obs_latency=0.17, robot_obs_latency=0.0001, gripper_obs_latency=0.01, robot_action_latency=0.18, gripper_action_latency=0.1, # camera_obs_latency=0.0, # robot_obs_latency=0.0, # gripper_obs_latency=0.0, # robot_action_latency=0.0, # gripper_action_latency=0.0, # obs camera_obs_horizon=cfg.task.shape_meta.obs.camera0_rgb.horizon, robot_obs_horizon=cfg.task.shape_meta.obs.robot0_eef_pos.horizon, gripper_obs_horizon=cfg.task.shape_meta.obs.robot0_gripper_width.horizon, no_mirror=no_mirror, fisheye_converter=fisheye_converter, mirror_crop=mirror_crop, mirror_swap=mirror_swap, # action max_pos_speed=2.0, max_rot_speed=6.0, robot_type=robot_type, shm_manager=shm_manager) as env: cv2.setNumThreads(2) print("Waiting for camera") time.sleep(1.0) # load match_dataset episode_first_frame_map = dict() match_replay_buffer = None if match_dataset is not None: match_dir = pathlib.Path(match_dataset) match_zarr_path = match_dir.joinpath('replay_buffer.zarr') match_replay_buffer = ReplayBuffer.create_from_path(str(match_zarr_path), mode='r') match_video_dir = match_dir.joinpath('videos') for vid_dir in match_video_dir.glob("*/"): episode_idx = int(vid_dir.stem) match_video_path = vid_dir.joinpath(f'{match_camera}.mp4') if match_video_path.exists(): img = None with av.open(str(match_video_path)) as container: stream = container.streams.video[0] for frame in container.decode(stream): img = frame.to_ndarray(format='rgb24') break # img = VideoFileClip(str(match_video_path)).get_frame(0) episode_first_frame_map[episode_idx] = img print(f"Loaded initial frame for {len(episode_first_frame_map)} episodes") # creating model # have to be done after fork to prevent # duplicating CUDA context with ffmpeg nvenc cls = hydra.utils.get_class(cfg._target_) workspace = cls(cfg) workspace: BaseWorkspace workspace.load_payload(payload, exclude_keys=None, include_keys=None) policy = workspace.model if cfg.training.use_ema: policy = workspace.ema_model policy.num_inference_steps = 16 # DDIM inference iterations obs_pose_rep = cfg.task.pose_repr.obs_pose_repr action_pose_repr = cfg.task.pose_repr.action_pose_repr print('obs_pose_rep', obs_pose_rep) print('action_pose_repr', action_pose_repr) device = torch.device('cuda') policy.eval().to(device) print("Warming up policy inference") obs = env.get_obs() with torch.no_grad(): policy.reset() obs_dict_np = get_real_umi_obs_dict( env_obs=obs, shape_meta=cfg.task.shape_meta, obs_pose_repr=obs_pose_rep) obs_dict = dict_apply(obs_dict_np, lambda x: torch.from_numpy(x).unsqueeze(0).to(device)) result = policy.predict_action(obs_dict) action = result['action_pred'][0].detach().to('cpu').numpy() assert action.shape[-1] == 10 action = get_real_umi_action(action, obs, action_pose_repr) assert action.shape[-1] == 7 del result print('Ready!') while True: # ========= human control loop ========== print("Human in control!") state = env.get_robot_state() target_pose = state['ActualTCPPose'] gripper_state = env.gripper.get_state() gripper_target_pos = gripper_state['gripper_position'] t_start = time.monotonic() iter_idx = 0 while True: # calculate timing t_cycle_end = t_start + (iter_idx + 1) * dt t_sample = t_cycle_end - command_latency t_command_target = t_cycle_end + dt # pump obs obs = env.get_obs() # visualize episode_id = env.replay_buffer.n_episodes vis_img = obs[f'camera{match_camera}_rgb'][-1] match_episode_id = episode_id if match_episode is not None: match_episode_id = match_episode if match_episode_id in episode_first_frame_map: match_img = episode_first_frame_map[match_episode_id] ih, iw, _ = match_img.shape oh, ow, _ = vis_img.shape tf = get_image_transform( input_res=(iw, ih), output_res=(ow, oh), bgr_to_rgb=False) match_img = tf(match_img).astype(np.float32) / 255 vis_img = (vis_img + match_img) / 2 obs_img = obs['camera0_rgb'][-1] if mirror_crop: crop_img = obs['camera0_rgb_mirror_crop'][-1] vis_img = np.concatenate([obs_img, crop_img, vis_img], axis=1) else: vis_img = np.concatenate([obs_img, vis_img], axis=1) text = f'Episode: {episode_id}' cv2.putText( vis_img, text, (10,20), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, lineType=cv2.LINE_AA, thickness=3, color=(0,0,0) ) cv2.putText( vis_img, text, (10,20), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, thickness=1, color=(255,255,255) ) cv2.imshow('default', vis_img[...,::-1]) _ = cv2.pollKey() press_events = key_counter.get_press_events() start_policy = False for key_stroke in press_events: if key_stroke == KeyCode(char='q'): # Exit program env.end_episode() exit(0) elif key_stroke == KeyCode(char='c'): # Exit human control loop # hand control over to the policy start_policy = True elif key_stroke == KeyCode(char='e'): # Next episode if match_episode is not None: match_episode = min(match_episode + 1, env.replay_buffer.n_episodes-1) elif key_stroke == KeyCode(char='w'): # Prev episode if match_episode is not None: match_episode = max(match_episode - 1, 0) elif key_stroke == KeyCode(char='m'): # move the robot duration = 3.0 ep = match_replay_buffer.get_episode(match_episode_id) pos = ep['robot0_eef_pos'][0] rot = ep['robot0_eef_rot_axis_angle'][0] grip = ep['robot0_gripper_width'][0] start_pose = np.concatenate([pos, rot]) start_grip = grip[0] env.robot.servoL(start_pose, duration=duration) env.gripper.schedule_waypoint(start_grip, target_time=time.time() + duration) time.sleep(duration) target_pose = start_pose gripper_target_pos = start_grip elif key_stroke == Key.backspace: if click.confirm('Are you sure to drop an episode?'): env.drop_episode() key_counter.clear() if start_policy: break precise_wait(t_sample) # get teleop command sm_state = sm.get_motion_state_transformed() # print(sm_state) dpos = sm_state[:3] * (0.5 / frequency) drot_xyz = sm_state[3:] * (1.5 / frequency) drot = st.Rotation.from_euler('xyz', drot_xyz) target_pose[:3] += dpos target_pose[3:] = (drot * st.Rotation.from_rotvec( target_pose[3:])).as_rotvec() target_pose[2] = np.maximum(target_pose[2], 0.055) dpos = 0 if sm.is_button_pressed(0): # close gripper dpos = -gripper_speed / frequency if sm.is_button_pressed(1): dpos = gripper_speed / frequency gripper_target_pos = np.clip(gripper_target_pos + dpos, 0, max_gripper_width) action = np.zeros((7,)) action[:6] = target_pose action[-1] = gripper_target_pos # execute teleop command env.exec_actions( actions=[action], timestamps=[t_command_target-time.monotonic()+time.time()], compensate_latency=False) precise_wait(t_cycle_end) iter_idx += 1 # ========== policy control loop ============== try: # start episode policy.reset() start_delay = 1.0 eval_t_start = time.time() + start_delay t_start = time.monotonic() + start_delay env.start_episode(eval_t_start) # wait for 1/30 sec to get the closest frame actually # reduces overall latency frame_latency = 1/60 precise_wait(eval_t_start - frame_latency, time_func=time.time) print("Started!") iter_idx = 0 perv_target_pose = None while True: # calculate timing t_cycle_end = t_start + (iter_idx + steps_per_inference) * dt # get obs obs = env.get_obs() obs_timestamps = obs['timestamp'] print(f'Obs latency {time.time() - obs_timestamps[-1]}') # run inference with torch.no_grad(): s = time.time() obs_dict_np = get_real_umi_obs_dict( env_obs=obs, shape_meta=cfg.task.shape_meta, obs_pose_repr=obs_pose_rep) obs_dict = dict_apply(obs_dict_np, lambda x: torch.from_numpy(x).unsqueeze(0).to(device)) result = policy.predict_action(obs_dict) raw_action = result['action_pred'][0].detach().to('cpu').numpy() action = get_real_umi_action(raw_action, obs, action_pose_repr) print('Inference latency:', time.time() - s) # convert policy action to env actions this_target_poses = action # this_target_poses[:,2] = np.maximum(this_target_poses[:,2], 0.055) # deal with timing # the same step actions are always the target for action_timestamps = (np.arange(len(action), dtype=np.float64) ) * dt + obs_timestamps[-1] action_exec_latency = 0.01 curr_time = time.time() is_new = action_timestamps > (curr_time + action_exec_latency) if np.sum(is_new) == 0: # exceeded time budget, still do something this_target_poses = this_target_poses[[-1]] # schedule on next available step next_step_idx = int(np.ceil((curr_time - eval_t_start) / dt)) action_timestamp = eval_t_start + (next_step_idx) * dt print('Over budget', action_timestamp - curr_time) action_timestamps = np.array([action_timestamp]) else: this_target_poses = this_target_poses[is_new] action_timestamps = action_timestamps[is_new] # execute actions env.exec_actions( actions=this_target_poses, timestamps=action_timestamps, compensate_latency=True ) print(f"Submitted {len(this_target_poses)} steps of actions.") # visualize episode_id = env.replay_buffer.n_episodes if mirror_crop: vis_img = obs[f'camera{vis_camera_idx}_rgb'][-1] crop_img = obs['camera0_rgb_mirror_crop'][-1] vis_img = np.concatenate([vis_img, crop_img], axis=1) else: vis_img = obs[f'camera{vis_camera_idx}_rgb'][-1] text = 'Episode: {}, Time: {:.1f}'.format( episode_id, time.monotonic() - t_start ) cv2.putText( vis_img, text, (10,20), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, thickness=1, color=(255,255,255) ) cv2.imshow('default', vis_img[...,::-1]) _ = cv2.pollKey() press_events = key_counter.get_press_events() stop_episode = False for key_stroke in press_events: if key_stroke == KeyCode(char='s'): # Stop episode # Hand control back to human print('Stopped.') stop_episode = True t_since_start = time.time() - eval_t_start if t_since_start > max_duration: print("Max Duration reached.") stop_episode = True if stop_episode: env.end_episode() break # wait for execution precise_wait(t_cycle_end - frame_latency) iter_idx += steps_per_inference except KeyboardInterrupt: print("Interrupted!") # stop robot. env.end_episode() print("Stopped.") # %% if __name__ == '__main__': main()