from typing import Optional import pathlib import numpy as np import time import shutil import math from multiprocessing.managers import SharedMemoryManager from umi.real_world.rtde_interpolation_controller import RTDEInterpolationController from umi.real_world.wsg_controller import WSGController from umi.real_world.multi_uvc_camera import MultiUvcCamera from umi.real_world.video_recorder import VideoRecorder from unified_video_action.common.timestamp_accumulator import ( TimestampObsAccumulator, TimestampActionAccumulator, align_timestamps, ) from umi.real_world.multi_camera_visualizer import MultiCameraVisualizer from unified_video_action.common.replay_buffer import ReplayBuffer from unified_video_action.common.cv2_util import get_image_transform, optimal_row_cols from umi.common.usb_util import reset_all_elgato_devices, get_sorted_v4l_paths DEFAULT_OBS_KEY_MAP = { # robot "ActualTCPPose": "robot_eef_pose", "ActualTCPSpeed": "robot_eef_pose_vel", "ActualQ": "robot_joint", "ActualQd": "robot_joint_vel", # gripper "gripper_position": "gripper_position", "gripper_velocity": "gripper_velocity", "gripper_force": "gripper_force", # timestamps "step_idx": "step_idx", "timestamp": "timestamp", } class RealEnv: def __init__( self, # required params output_dir, robot_ip, gripper_ip, gripper_port=1000, # env params frequency=10, n_obs_steps=2, # obs obs_image_resolution=(256, 256), max_obs_buffer_size=30, obs_key_map=DEFAULT_OBS_KEY_MAP, obs_float32=False, # action max_pos_speed=0.25, max_rot_speed=0.6, # robot tcp_offset=0.13, init_joints=False, # video capture params video_capture_fps=60, video_capture_resolution=(1280, 720), # saving params record_raw_video=True, thread_per_video=4, video_crf=21, # vis params enable_multi_cam_vis=True, multi_cam_vis_resolution=(1280, 720), # shared memory shm_manager=None, ): assert frequency <= video_capture_fps output_dir = pathlib.Path(output_dir) assert output_dir.parent.is_dir() video_dir = output_dir.joinpath("videos") video_dir.mkdir(parents=True, exist_ok=True) zarr_path = str(output_dir.joinpath("replay_buffer.zarr").absolute()) replay_buffer = ReplayBuffer.create_from_path(zarr_path=zarr_path, mode="a") if shm_manager is None: shm_manager = SharedMemoryManager() shm_manager.start() # Find and reset all Elgato capture cards. # Required to workaround a firmware bug. reset_all_elgato_devices() # Wait for all v4l cameras to be back online time.sleep(0.1) v4l_paths = get_sorted_v4l_paths() color_tf = get_image_transform( input_res=video_capture_resolution, output_res=obs_image_resolution, # obs output rgb bgr_to_rgb=True, ) color_transform = color_tf if obs_float32: color_transform = lambda x: color_tf(x).astype(np.float32) / 255 def transform(data): data["color"] = color_transform(data["color"]) return data rw, rh, col, row = optimal_row_cols( n_cameras=len(v4l_paths), in_wh_ratio=obs_image_resolution[0] / obs_image_resolution[1], max_resolution=multi_cam_vis_resolution, ) vis_color_transform = get_image_transform( input_res=video_capture_resolution, output_res=(rw, rh), bgr_to_rgb=False ) def vis_transform(data): data["color"] = vis_color_transform(data["color"]) return data recording_transfrom = None recording_fps = video_capture_fps recording_pix_fmt = "bgr24" if not record_raw_video: recording_transfrom = transform recording_fps = frequency recording_pix_fmt = "rgb24" video_recorder = VideoRecorder.create_h264( shm_manager=shm_manager, fps=recording_fps, codec="h264", input_pix_fmt=recording_pix_fmt, crf=video_crf, thread_type="FRAME", thread_count=thread_per_video, ) camera = MultiUvcCamera( dev_video_paths=v4l_paths, shm_manager=shm_manager, resolution=video_capture_resolution, capture_fps=video_capture_fps, put_fps=video_capture_fps, # send every frame immediately after arrival # ignores put_fps put_downsample=False, record_fps=recording_fps, get_max_k=max_obs_buffer_size, transform=transform, vis_transform=vis_transform, recording_transform=recording_transfrom, video_recorder=video_recorder, verbose=False, ) multi_cam_vis = None if enable_multi_cam_vis: multi_cam_vis = MultiCameraVisualizer( camera=camera, row=row, col=col, rgb_to_bgr=False ) cube_diag = np.linalg.norm([1, 1, 1]) j_init = np.array([0, -90, -90, -90, 90, 0]) / 180 * np.pi if not init_joints: j_init = None robot = RTDEInterpolationController( shm_manager=shm_manager, robot_ip=robot_ip, frequency=500, # UR5 CB3 RTDE lookahead_time=0.1, gain=300, max_pos_speed=max_pos_speed * cube_diag, max_rot_speed=max_rot_speed * cube_diag, launch_timeout=3, tcp_offset_pose=[0, 0, tcp_offset, 0, 0, 0], payload_mass=None, payload_cog=None, joints_init=j_init, joints_init_speed=1.05, soft_real_time=False, verbose=False, receive_keys=None, get_max_k=max_obs_buffer_size, ) gripper = WSGController( shm_manager=shm_manager, hostname=gripper_ip, port=gripper_port, ) self.camera = camera self.robot = robot self.gripper = gripper self.multi_cam_vis = multi_cam_vis self.video_capture_fps = video_capture_fps self.frequency = frequency self.n_obs_steps = n_obs_steps self.max_obs_buffer_size = max_obs_buffer_size self.max_pos_speed = max_pos_speed self.max_rot_speed = max_rot_speed self.obs_key_map = obs_key_map # recording self.output_dir = output_dir self.video_dir = video_dir self.replay_buffer = replay_buffer # temp memory buffers self.last_camera_data = None # recording buffers self.robot_obs_accumulator = None self.gripper_obs_accumulator = None self.action_accumulator = None self.stage_accumulator = None self.start_time = None # ======== start-stop API ============= @property def is_ready(self): return self.camera.is_ready and self.robot.is_ready and self.gripper.is_ready def start(self, wait=True): self.camera.start(wait=False) self.gripper.start(wait=False) self.robot.start(wait=False) if self.multi_cam_vis is not None: self.multi_cam_vis.start(wait=False) if wait: self.start_wait() def stop(self, wait=True): self.end_episode() if self.multi_cam_vis is not None: self.multi_cam_vis.stop(wait=False) self.robot.stop(wait=False) self.gripper.stop(wait=False) self.camera.stop(wait=False) if wait: self.stop_wait() def start_wait(self): self.camera.start_wait() self.gripper.start_wait() self.robot.start_wait() if self.multi_cam_vis is not None: self.multi_cam_vis.start_wait() def stop_wait(self): self.robot.stop_wait() self.gripper.stop_wait() self.camera.stop_wait() if self.multi_cam_vis is not None: self.multi_cam_vis.stop_wait() # ========= context manager =========== def __enter__(self): self.start() return self def __exit__(self, exc_type, exc_val, exc_tb): self.stop() # ========= async env API =========== def get_obs(self) -> dict: "observation dict" assert self.is_ready # get data # 30 Hz, camera_receive_timestamp k = math.ceil(self.n_obs_steps * (self.video_capture_fps / self.frequency)) self.last_camera_data = self.camera.get(k=k, out=self.last_camera_data) # 125 hz, robot_receive_timestamp last_robot_data = self.robot.get_all_state() # both have more than n_obs_steps data # 30 hz, gripper_receive_timestamp last_gripper_data = self.gripper.get_all_state() # align camera obs timestamps dt = 1 / self.frequency last_timestamp = np.max( [x["timestamp"][-1] for x in self.last_camera_data.values()] ) obs_align_timestamps = last_timestamp - (np.arange(self.n_obs_steps)[::-1] * dt) camera_obs = dict() for camera_idx, value in self.last_camera_data.items(): this_timestamps = value["timestamp"] this_idxs = list() for t in obs_align_timestamps: is_before_idxs = np.nonzero(this_timestamps < t)[0] this_idx = 0 if len(is_before_idxs) > 0: this_idx = is_before_idxs[-1] this_idxs.append(this_idx) # remap key camera_obs[f"camera_{camera_idx}"] = value["color"][this_idxs] # align robot obs robot_timestamps = last_robot_data["robot_receive_timestamp"] this_timestamps = robot_timestamps this_idxs = list() for t in obs_align_timestamps: is_before_idxs = np.nonzero(this_timestamps < t)[0] this_idx = 0 if len(is_before_idxs) > 0: this_idx = is_before_idxs[-1] this_idxs.append(this_idx) robot_obs_raw = dict() for k, v in last_robot_data.items(): if k in self.obs_key_map: robot_obs_raw[self.obs_key_map[k]] = v robot_obs = dict() for k, v in robot_obs_raw.items(): robot_obs[k] = v[this_idxs] # align gripper obs gripper_timestamps = last_gripper_data["gripper_receive_timestamp"] this_timestamps = gripper_timestamps this_idxs = list() for t in obs_align_timestamps: is_before_idxs = np.nonzero(this_timestamps < t)[0] this_idx = 0 if len(is_before_idxs) > 0: this_idx = is_before_idxs[-1] this_idxs.append(this_idx) gripper_obs_raw = dict() for k, v in last_gripper_data.items(): if k in self.obs_key_map: gripper_obs_raw[self.obs_key_map[k]] = v gripper_obs = dict() for k, v in gripper_obs_raw.items(): gripper_obs[k] = v[this_idxs] # accumulate obs if self.robot_obs_accumulator is not None: self.robot_obs_accumulator.put(robot_obs_raw, robot_timestamps) if self.gripper_obs_accumulator is not None: self.gripper_obs_accumulator.put(gripper_obs_raw, gripper_timestamps) # return obs obs_data = dict(camera_obs) obs_data.update(robot_obs) obs_data.update(gripper_obs) obs_data["timestamp"] = obs_align_timestamps return obs_data def exec_actions( self, actions: np.ndarray, timestamps: np.ndarray, stages: Optional[np.ndarray] = None, ): assert self.is_ready if not isinstance(actions, np.ndarray): actions = np.array(actions) if not isinstance(timestamps, np.ndarray): timestamps = np.array(timestamps) if stages is None: stages = np.zeros_like(timestamps, dtype=np.int64) elif not isinstance(stages, np.ndarray): stages = np.array(stages, dtype=np.int64) # convert action to pose receive_time = time.time() is_new = timestamps > receive_time new_actions = actions[is_new] new_timestamps = timestamps[is_new] new_stages = stages[is_new] # schedule waypoints for i in range(len(new_actions)): r_actions = new_actions[i, :6] g_actions = new_actions[i, 6:] + 1 self.robot.schedule_waypoint(pose=r_actions, target_time=new_timestamps[i]) self.gripper.schedule_waypoint( pos=g_actions, target_time=new_timestamps[i] - 0.02 ) # record actions if self.action_accumulator is not None: self.action_accumulator.put(new_actions, new_timestamps) if self.stage_accumulator is not None: self.stage_accumulator.put(new_stages, new_timestamps) def get_robot_state(self): return self.robot.get_state() # recording API def start_episode(self, start_time=None): "Start recording and return first obs" if start_time is None: start_time = time.time() self.start_time = start_time assert self.is_ready # prepare recording stuff episode_id = self.replay_buffer.n_episodes this_video_dir = self.video_dir.joinpath(str(episode_id)) this_video_dir.mkdir(parents=True, exist_ok=True) n_cameras = self.camera.n_cameras video_paths = list() for i in range(n_cameras): video_paths.append(str(this_video_dir.joinpath(f"{i}.mp4").absolute())) # start recording on camera self.camera.restart_put(start_time=start_time) self.camera.start_recording(video_path=video_paths, start_time=start_time) # create accumulators self.robot_obs_accumulator = TimestampObsAccumulator( start_time=start_time, dt=1 / self.frequency ) self.gripper_obs_accumulator = TimestampObsAccumulator( start_time=start_time, dt=1 / self.frequency ) self.action_accumulator = TimestampActionAccumulator( start_time=start_time, dt=1 / self.frequency ) self.stage_accumulator = TimestampActionAccumulator( start_time=start_time, dt=1 / self.frequency ) print(f"Episode {episode_id} started!") def end_episode(self): "Stop recording" assert self.is_ready # stop video recorder self.camera.stop_recording() if self.robot_obs_accumulator is not None: # recording assert self.gripper_obs_accumulator is not None assert self.action_accumulator is not None assert self.stage_accumulator is not None # Since the only way to accumulate obs and action is by calling # get_obs and exec_actions, which will be in the same thread. # We don't need to worry new data come in here. robot_obs_data = self.robot_obs_accumulator.data robot_obs_timestamps = self.robot_obs_accumulator.timestamps gripper_obs_data = self.gripper_obs_accumulator.data gripper_obs_timestamps = self.gripper_obs_accumulator.timestamps actions = self.action_accumulator.actions action_timestamps = self.action_accumulator.timestamps stages = self.stage_accumulator.actions n_steps = min( len(robot_obs_timestamps), len(gripper_obs_timestamps), len(action_timestamps), ) if n_steps > 0: episode = dict() episode["timestamp"] = robot_obs_timestamps[:n_steps] episode["action"] = actions[:n_steps] episode["stage"] = stages[:n_steps] for key, value in robot_obs_data.items(): episode[key] = value[:n_steps] for key, value in gripper_obs_data.items(): episode[key] = value[:n_steps] self.replay_buffer.add_episode(episode, compressors="disk") episode_id = self.replay_buffer.n_episodes - 1 print(f"Episode {episode_id} saved!") self.robot_obs_accumulator = None self.gripper_obs_accumulator = None self.action_accumulator = None self.stage_accumulator = None def drop_episode(self): self.end_episode() self.replay_buffer.drop_episode() episode_id = self.replay_buffer.n_episodes this_video_dir = self.video_dir.joinpath(str(episode_id)) if this_video_dir.exists(): shutil.rmtree(str(this_video_dir)) print(f"Episode {episode_id} dropped!")