| |
| import argparse |
| import threading |
| import time |
| from collections import deque |
| import sys |
| import termios |
| import tty |
| import select |
| import cv2 |
| import numpy as np |
| import rospy |
| import torch |
| from cv_bridge import CvBridge |
| from geometry_msgs.msg import Twist |
| from nav_msgs.msg import Odometry |
| from openpi_client import image_tools, websocket_client_policy |
| from piper_msgs.msg import PosCmd |
| from sensor_msgs.msg import Image, JointState |
| from std_msgs.msg import Header |
| import os |
| import dm_env |
| import collections |
| from deploy.data_collection.collect_data import save_data |
|
|
| CAMERA_NAMES = ["cam_high", "cam_right_wrist", "cam_left_wrist"] |
|
|
| observation_window = None |
| stream_buffer = None |
|
|
| |
| inference_collection_active = False |
| inference_collection_lock = threading.Lock() |
|
|
| |
| save_data_requested = False |
| save_data_lock = threading.Lock() |
|
|
| lang_embeddings = "Pick and sort bricks on the conveyor." |
| RIGHT_OFFSET = 0.003 |
|
|
|
|
| class InferenceDataCollector: |
| """Inference-mode data collector that saves episodes classified as success or failure.""" |
|
|
| def __init__(self, camera_names, dataset_dir="./data", task_name=None): |
| self.camera_names = camera_names |
| self.dataset_dir = dataset_dir |
| self.task_name = task_name |
| self.is_collecting = False |
| self.timesteps = [] |
| self.actions = [] |
| self.frame_count = 0 |
| self.episode_idx_success = 0 |
| self.episode_idx_fail = 0 |
| self.full_dataset_dir = None |
|
|
| def _find_next_episode_idx(self, subdir_name): |
| """Find next available episode index in the given subdirectory.""" |
| if self.full_dataset_dir is None: |
| return 0 |
| subdir_path = os.path.join(self.full_dataset_dir, subdir_name) |
| if not os.path.exists(subdir_path): |
| os.makedirs(subdir_path, exist_ok=True) |
| print(f"Created directory: {subdir_path}") |
| return 0 |
| existing = [ |
| f for f in os.listdir(subdir_path) |
| if f.startswith('episode_') and f.endswith('.hdf5') |
| ] |
| if existing: |
| indices = [int(f.split('_')[1].split('.')[0]) for f in existing] |
| return max(indices) + 1 |
| return 0 |
|
|
| def update_episode_indices(self): |
| """Update episode indices for success and fail subdirectories.""" |
| self.episode_idx_success = self._find_next_episode_idx('aloha_mobile_success') |
| self.episode_idx_fail = self._find_next_episode_idx('aloha_mobile_fail') |
| print(f"Next success episode: {self.episode_idx_success}") |
| print(f"Next fail episode: {self.episode_idx_fail}") |
|
|
| def start_collection(self): |
| """Start collecting frames and actions.""" |
| self.is_collecting = True |
| self.timesteps = [] |
| self.actions = [] |
| self.frame_count = 0 |
| print(f"\n{'='*70}") |
| print("Data collection started!") |
| print(f" Save dir: {self.full_dataset_dir}") |
| print(f" Success idx: {self.episode_idx_success} | Fail idx: {self.episode_idx_fail}") |
| print(f" Press 's' to stop and save") |
| print(f"{'='*70}\n") |
|
|
| def add_frame(self, observation, action): |
| """Add one frame of observation and action data.""" |
| if not self.is_collecting: |
| return |
| self.frame_count += 1 |
| step_type = dm_env.StepType.FIRST if self.frame_count == 1 else dm_env.StepType.MID |
| self.timesteps.append(dm_env.TimeStep( |
| step_type=step_type, reward=None, discount=None, observation=observation |
| )) |
| self.actions.append(action) |
| if self.frame_count % 50 == 0: |
| print(f"Collected {self.frame_count} frames") |
|
|
| def save_current_episode(self, export_video=True, video_fps=30, |
| video_codec='libx264', video_quality=23): |
| """Stop collection and save episode after asking user for success/fail classification.""" |
| if len(self.actions) == 0: |
| print("\033[31m❌ No data collected, cannot save\033[0m") |
| return False |
|
|
| print(f"\n{'='*70}") |
| print("❓ Was this episode a success or failure?") |
| print(" Press '1' → Success (save to aloha_mobile_success)") |
| print(" Press '0' → Failure (save to aloha_mobile_fail)") |
| print(f"{'='*70}") |
|
|
| local_settings = termios.tcgetattr(sys.stdin) |
| try: |
| tty.setcbreak(sys.stdin.fileno()) |
| is_success = None |
| subdir_name = None |
| episode_idx = None |
| while True: |
| if select.select([sys.stdin], [], [], 0.1)[0]: |
| char = sys.stdin.read(1) |
| if char == '1': |
| is_success = True |
| subdir_name = 'aloha_mobile_success' |
| episode_idx = self.episode_idx_success |
| print("\n✅ Marked as [Success]") |
| break |
| elif char == '0': |
| is_success = False |
| subdir_name = 'aloha_mobile_fail' |
| episode_idx = self.episode_idx_fail |
| print("\n❌ Marked as [Failure]") |
| break |
| else: |
| print(f"⚠️ Invalid input '{char}', press '1' or '0'") |
| finally: |
| termios.tcsetattr(sys.stdin, termios.TCSADRAIN, local_settings) |
|
|
| if is_success is None: |
| print("\n❌ No valid input, cancelling save") |
| return False |
|
|
| save_dir = os.path.join(self.full_dataset_dir, subdir_name) |
| os.makedirs(save_dir, exist_ok=True) |
| dataset_path = os.path.join(save_dir, f"episode_{episode_idx}") |
| print(f"\n💾 Saving Episode {episode_idx} to {subdir_name}...") |
| print(f" Frames: {len(self.actions)}") |
|
|
| class _Args: |
| def __init__(self, camera_names, export_video, video_fps, video_codec, video_quality): |
| self.camera_names = camera_names |
| self.export_video = export_video |
| self.video_fps = video_fps |
| self.video_codec = video_codec |
| self.video_quality = video_quality |
| self.use_robot_base = False |
|
|
| _args = _Args(self.camera_names, export_video, video_fps, video_codec, video_quality) |
| try: |
| save_data(_args, self.timesteps, self.actions, dataset_path) |
| print(f"\n\033[32m✅ Episode {episode_idx} saved!\033[0m") |
| print(f" Path: {dataset_path}.hdf5") |
| if is_success: |
| self.episode_idx_success += 1 |
| else: |
| self.episode_idx_fail += 1 |
| self.timesteps = [] |
| self.actions = [] |
| self.frame_count = 0 |
| self.is_collecting = False |
| print(f"\n{'='*70}") |
| print("✅ Save complete") |
| print(f" Next success episode: {self.episode_idx_success}") |
| print(f" Next fail episode: {self.episode_idx_fail}") |
| print(f"{'='*70}\n") |
| return True |
| except Exception as e: |
| print(f"\033[31m❌ Save failed: {e}\033[0m") |
| import traceback |
| traceback.print_exc() |
| return False |
|
|
| def get_frame_count(self): |
| return self.frame_count |
|
|
| def has_data(self): |
| return len(self.actions) > 0 |
|
|
|
|
| def keyboard_monitor_thread(): |
| """Background thread: press 's' to trigger data save and exit.""" |
| global save_data_requested |
| original_settings = None |
| try: |
| original_settings = termios.tcgetattr(sys.stdin) |
| tty.setcbreak(sys.stdin.fileno()) |
| while not rospy.is_shutdown(): |
| if select.select([sys.stdin], [], [], 0.1)[0]: |
| char = sys.stdin.read(1) |
| if char.lower() == 's': |
| with save_data_lock: |
| save_data_requested = True |
| print("\n" + "💾" * 35) |
| print("💾 Save requested, stopping collection...") |
| print("💾" * 35 + "\n") |
| finally: |
| if original_settings is not None: |
| try: |
| termios.tcsetattr(sys.stdin, termios.TCSADRAIN, original_settings) |
| except Exception: |
| pass |
|
|
|
|
| class StreamActionBuffer: |
| """ |
| Maintains a chunk queue for actions and the current smoothed execution sequence. |
| integrate_new_chunk supports latency trimming and linear overlap blending. |
| """ |
| def __init__(self, max_chunks=10, decay_alpha=0.25, state_dim=14, smooth_method="temporal"): |
| self.chunks = deque() |
| self.max_chunks = max_chunks |
| self.lock = threading.Lock() |
| self.decay_alpha = float(decay_alpha) |
| self.state_dim = state_dim |
| self.smooth_method = smooth_method |
| self.cur_chunk = deque() |
| self.k = 0 |
| self.last_action = None |
|
|
| def integrate_new_chunk(self, actions_chunk: np.ndarray, max_k: int, min_m: int = 8): |
| with self.lock: |
| if actions_chunk is None or len(actions_chunk) == 0: |
| return |
| max_k = max(0, int(max_k)) |
| min_m = max(1, int(min_m)) |
| drop_n = min(self.k, max_k) |
| if drop_n >= len(actions_chunk): |
| return |
| new_chunk = [a.copy() for a in actions_chunk[drop_n:]] |
| if len(self.cur_chunk) == 0 and self.last_action is not None: |
| old_list = [np.asarray(self.last_action, dtype=float).copy() for _ in range(min_m)] |
| self.last_action = None |
| else: |
| old_list = list(self.cur_chunk) |
| if len(old_list) > 0 and len(old_list) < min_m: |
| tail = np.asarray(old_list[-1], dtype=float).copy() |
| old_list.extend([tail.copy() for _ in range(min_m - len(old_list))]) |
| elif len(old_list) == 0: |
| self.cur_chunk = deque(new_chunk, maxlen=None) |
| self.k = 0 |
| return |
| new_list = list(new_chunk) |
| overlap_len = min(len(old_list), len(new_list)) |
| if overlap_len <= 0: |
| self.cur_chunk = deque(new_list, maxlen=None) |
| self.k = 0 |
| return |
| if len(old_list) > len(new_list): |
| old_list = old_list[:len(new_list)] |
| overlap_len = len(new_list) |
| w_old = np.array([1.0], dtype=float) if overlap_len == 1 else np.linspace(1.0, 0.0, overlap_len, dtype=float) |
| w_new = 1.0 - w_old |
| smoothed = [ |
| (w_old[i] * np.asarray(old_list[i], dtype=float) + |
| w_new[i] * np.asarray(new_list[i], dtype=float)) |
| for i in range(overlap_len) |
| ] |
| self.cur_chunk = deque([a.copy() for a in smoothed + new_list[overlap_len:]], maxlen=None) |
| self.k = 0 |
|
|
| def pop_next_action(self) -> np.ndarray | None: |
| with self.lock: |
| if len(self.cur_chunk) == 0: |
| return None |
| if len(self.cur_chunk) == 1: |
| self.last_action = np.asarray(self.cur_chunk[0], dtype=float).copy() |
| act = np.asarray(self.cur_chunk.popleft(), dtype=float) |
| self.k += 1 |
| return act |
|
|
|
|
| def start_inference_thread(args, config, policy, ros_operator): |
| th = threading.Thread(target=inference_fn_non_blocking_fast, args=(args, config, policy, ros_operator)) |
| th.daemon = True |
| th.start() |
|
|
|
|
| def inference_fn_non_blocking_fast(args, config, policy, ros_operator): |
| global stream_buffer |
| rate = rospy.Rate(getattr(args, "inference_rate", 4)) |
| while not rospy.is_shutdown(): |
| try: |
| update_observation_window(args, config, ros_operator) |
| latest_obs = observation_window[-1] |
| imgs = [ |
| latest_obs["images"][config["camera_names"][0]], |
| latest_obs["images"][config["camera_names"][1]], |
| latest_obs["images"][config["camera_names"][2]], |
| ] |
| imgs = [cv2.cvtColor(im, cv2.COLOR_BGR2RGB) for im in imgs] |
| imgs = image_tools.resize_with_pad(np.array(imgs), 224, 224) |
| proprio = latest_obs["qpos"] |
| payload = { |
| "state": proprio, |
| "images": { |
| "top_head": imgs[0].transpose(2, 0, 1), |
| "hand_right": imgs[1].transpose(2, 0, 1), |
| "hand_left": imgs[2].transpose(2, 0, 1), |
| }, |
| "prompt": lang_embeddings, |
| } |
| actions = policy.infer(payload)["actions"] |
| if actions is not None and len(actions) > 0: |
| max_k = int(getattr(args, "latency_k", 0)) |
| min_m = int(getattr(args, "min_smooth_steps", 8)) |
| stream_buffer.integrate_new_chunk(actions, max_k=max_k, min_m=min_m) |
| try: |
| rate.sleep() |
| except Exception: |
| pass |
| except Exception: |
| try: |
| rate.sleep() |
| except Exception: |
| time.sleep(0.005) |
| continue |
|
|
|
|
| def set_seed(seed): |
| torch.manual_seed(seed) |
| np.random.seed(seed) |
|
|
|
|
| class SimpleKalmanFilter: |
| def __init__(self, process_variance=1e-6, measurement_variance=1e-7, initial_value=None): |
| self.process_variance = process_variance |
| self.measurement_variance = measurement_variance |
| self.estimate = initial_value |
| self.error_estimate = 1.0 |
|
|
| def update(self, measurement): |
| if self.estimate is None: |
| self.estimate = measurement.copy() |
| return self.estimate |
| |
| kalman_gain = self.error_estimate / (self.error_estimate + self.measurement_variance) |
| |
| self.estimate = self.estimate + kalman_gain * (measurement - self.estimate) |
| |
| self.error_estimate = (1 - kalman_gain) * self.error_estimate + abs(self.estimate - measurement) * self.process_variance |
| return self.estimate |
|
|
|
|
| def interpolate_action(args, prev_action, cur_action): |
| steps = np.concatenate((np.array(args.arm_steps_length), np.array(args.arm_steps_length)), axis=0) |
| diff = np.abs(cur_action - prev_action) |
| step = np.ceil(diff / steps).astype(int) |
| step = np.max(step) |
| if step <= 1: |
| return cur_action[np.newaxis, :] |
| new_actions = np.linspace(prev_action, cur_action, step + 1) |
| return new_actions[1:] |
|
|
|
|
| def minimum_jerk_interpolation(args, prev_action, cur_action): |
| num_steps = args.jerk_num_steps |
| t_normalized = np.linspace(0, 1, num_steps + 1)[1:] |
| trajectory = [] |
| for tau in t_normalized: |
| factor = 10 * (tau ** 3) - 15 * (tau ** 4) + 6 * (tau ** 5) |
| trajectory.append(prev_action + factor * (cur_action - prev_action)) |
| return np.array(trajectory) |
|
|
|
|
| def get_config(args): |
| return { |
| "episode_len": args.max_publish_step, |
| "state_dim": 14, |
| "chunk_size": args.chunk_size, |
| "camera_names": CAMERA_NAMES, |
| } |
|
|
|
|
| def get_ros_observation(args, ros_operator): |
| rate = rospy.Rate(args.publish_rate) |
| print_flag = True |
| |
| max_wait_s = 0.6 |
| start_t = time.time() |
|
|
| while True and not rospy.is_shutdown(): |
| |
| result = ros_operator.get_frame_peek() |
| if not result: |
| |
| if (time.time() - start_t) > max_wait_s and ros_operator.last_frame_cache is not None: |
| if print_flag: |
| print("syn timeout, using last cached frame in get_ros_observation") |
| print_flag = False |
| (img_front, img_left, img_right, _, _, _, puppet_arm_left, puppet_arm_right, _) = \ |
| ros_operator.last_frame_cache |
| return (img_front, img_left, img_right, puppet_arm_left, puppet_arm_right) |
| if print_flag: |
| print("syn fail when get_ros_observation") |
| print_flag = False |
| rate.sleep() |
| continue |
| print_flag = True |
| (img_front, img_left, img_right, _, _, _, puppet_arm_left, puppet_arm_right, _) = result |
| return (img_front, img_left, img_right, puppet_arm_left, puppet_arm_right) |
|
|
|
|
| def update_observation_window(args, config, ros_operator): |
| |
| def jpeg_mapping(img): |
| img = cv2.imencode(".jpg", img)[1].tobytes() |
| img = cv2.imdecode(np.frombuffer(img, np.uint8), cv2.IMREAD_COLOR) |
| return img |
|
|
| global observation_window |
| if observation_window is None: |
| observation_window = deque(maxlen=2) |
| |
| observation_window.append({ |
| "qpos": None, |
| "images": { |
| config["camera_names"][0]: None, |
| config["camera_names"][1]: None, |
| config["camera_names"][2]: None, |
| }, |
| }) |
|
|
| img_front, img_left, img_right, puppet_arm_left, puppet_arm_right = \ |
| get_ros_observation(args, ros_operator) |
| img_front = jpeg_mapping(img_front) |
| img_left = jpeg_mapping(img_left) |
| img_right = jpeg_mapping(img_right) |
|
|
| qpos = np.concatenate( |
| (np.array(puppet_arm_left.position), np.array(puppet_arm_right.position)), axis=0 |
| ) |
| observation_window.append({ |
| "qpos": qpos, |
| "images": { |
| config["camera_names"][0]: img_front, |
| config["camera_names"][1]: img_right, |
| config["camera_names"][2]: img_left, |
| }, |
| }) |
|
|
|
|
| def inference_fn(args, config, policy): |
| global observation_window |
| global lang_embeddings |
|
|
| while True and not rospy.is_shutdown(): |
| time1 = time.time() |
|
|
| |
| image_arrs = [ |
| observation_window[-1]["images"][config["camera_names"][0]], |
| observation_window[-1]["images"][config["camera_names"][1]], |
| observation_window[-1]["images"][config["camera_names"][2]], |
| ] |
| |
| image_arrs = [cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for img in image_arrs] |
| image_arrs = image_tools.resize_with_pad(np.array(image_arrs), 224, 224) |
|
|
| |
| proprio = observation_window[-1]["qpos"] |
|
|
| payload = { |
| "state": proprio, |
| "images": { |
| "top_head": image_arrs[0].transpose(2, 0, 1), |
| "hand_right": image_arrs[1].transpose(2, 0, 1), |
| "hand_left": image_arrs[2].transpose(2, 0, 1), |
| }, |
| "prompt": lang_embeddings, |
| } |
|
|
| |
| actions = policy.infer(payload)["actions"] |
| print(f"Model inference time: {(time.time() - time1)*1000:.3f} ms") |
| return actions |
|
|
|
|
| def model_inference(args, config, ros_operator): |
| global lang_embeddings |
| global save_data_requested |
| global inference_collection_active |
| global stream_buffer |
|
|
| policy = websocket_client_policy.WebsocketClientPolicy(args.host, args.port) |
| kalman_filters = [SimpleKalmanFilter() for _ in range(config["state_dim"])] |
| print(f"Server metadata: {policy.get_server_metadata()}") |
|
|
| max_publish_step = config["episode_len"] |
| chunk_size = config["chunk_size"] |
|
|
| |
| left0 = [0, 0.32, -0.36, 0, 0.24, 0, 0.07] |
| right0 = [0, 0.32, -0.36, 0, 0.24, 0, 0.07] |
| ros_operator.puppet_arm_publish_continuous(left0, right0) |
|
|
| |
| data_collector = InferenceDataCollector( |
| camera_names=['cam_high', 'cam_left_wrist', 'cam_right_wrist'], |
| dataset_dir=args.dataset_dir, |
| task_name="inference_collect", |
| ) |
| data_collector.full_dataset_dir = os.path.join(args.dataset_dir, "aloha_mobile") |
| os.makedirs(data_collector.full_dataset_dir, exist_ok=True) |
| data_collector.update_episode_indices() |
|
|
| |
| keyboard_thread = threading.Thread(target=keyboard_monitor_thread, daemon=True) |
| keyboard_thread.start() |
| print("\n" + "="*70) |
| print("✅ Keyboard monitor started") |
| print(" Press 's' → Stop and save episode") |
| print("="*70 + "\n") |
| time.sleep(0.3) |
|
|
| print("\n" + "="*70) |
| print("🤖 Robot arms reset to initial position") |
| print("="*70) |
| input("Press [Enter] to start inference and data collection...") |
| print("\n" + "="*70) |
| print("🚀 Starting inference and data collection...") |
| print(" Press 's' to stop and save") |
| print("="*70 + "\n") |
|
|
| |
| data_collector.start_collection() |
| with inference_collection_lock: |
| inference_collection_active = True |
|
|
| pre_action = np.zeros(config["state_dim"]) |
| action = None |
|
|
| with torch.inference_mode(): |
| while True and not rospy.is_shutdown(): |
| t = 0 |
| rate = rospy.Rate(args.publish_rate) |
| action_buffer = np.zeros([chunk_size, config["state_dim"]]) |
| last_stream_act = None |
|
|
| while t < max_publish_step and not rospy.is_shutdown(): |
| |
| with save_data_lock: |
| if save_data_requested: |
| with inference_collection_lock: |
| if inference_collection_active and data_collector.has_data(): |
| print("\n" + "💾" * 35) |
| print("💾 Stopping and saving data...") |
| print("💾" * 35 + "\n") |
| ros_operator.open_grippers() |
| time.sleep(1.0) |
| print("✅ Grippers opened\n") |
| data_collector.save_current_episode( |
| export_video=getattr(args, 'export_video', True), |
| video_fps=getattr(args, 'video_fps', 30), |
| video_codec=getattr(args, 'video_codec', 'libx264'), |
| video_quality=getattr(args, 'video_quality', 23), |
| ) |
| print("\n" + "✅" * 35) |
| print("✅ Data saved, exiting...") |
| print("✅" * 35 + "\n") |
| rospy.signal_shutdown("Data saved, exiting...") |
| return |
| else: |
| print("\n⚠️ No data to save\n") |
| save_data_requested = False |
|
|
| |
| if args.use_temporal_smoothing: |
| if stream_buffer is None: |
| stream_buffer = StreamActionBuffer( |
| max_chunks=args.buffer_max_chunks, |
| decay_alpha=args.exp_decay_alpha, |
| state_dim=config["state_dim"], |
| smooth_method="temporal", |
| ) |
| start_inference_thread(args, config, policy, ros_operator) |
|
|
| act = stream_buffer.pop_next_action() |
| if act is not None: |
| if args.ctrl_type == "joint": |
| left_action = act[:7].copy() |
| right_action = act[7:14].copy() |
| left_action[6] = max(0.0, left_action[6] - RIGHT_OFFSET) |
| right_action[6] = max(0.0, right_action[6] - RIGHT_OFFSET) |
| ros_operator.puppet_arm_publish(left_action, right_action) |
|
|
| |
| with inference_collection_lock: |
| is_collecting = inference_collection_active and data_collector.is_collecting |
| if is_collecting: |
| try: |
| result = ros_operator.get_frame_peek() |
| if result: |
| (img_front, img_left, img_right, _, _, _, |
| puppet_arm_left, puppet_arm_right, _) = result |
| observation = collections.OrderedDict() |
| observation['images'] = { |
| 'cam_high': img_front.copy(), |
| 'cam_left_wrist': img_left.copy(), |
| 'cam_right_wrist': img_right.copy(), |
| } |
| observation['qpos'] = np.concatenate( |
| (puppet_arm_left.position, puppet_arm_right.position) |
| ).copy() |
| observation['qvel'] = np.zeros(14) |
| observation['effort'] = np.zeros(14) |
| observation['base_vel'] = [0.0, 0.0] |
| action_to_save = np.concatenate([left_action, right_action]).copy() |
| data_collector.add_frame(observation, action_to_save) |
| except Exception as e: |
| rospy.logwarn(f"[inference collect] frame error: {e}") |
|
|
| elif args.ctrl_type == "eef": |
| ros_operator.endpose_publish(act[:7], act[7:14]) |
|
|
| last_stream_act = act.copy() |
| else: |
| if last_stream_act is not None: |
| if args.ctrl_type == "joint": |
| ros_operator.puppet_arm_publish(last_stream_act[:7], last_stream_act[7:14]) |
| elif args.ctrl_type == "eef": |
| ros_operator.endpose_publish(last_stream_act[:7], last_stream_act[7:14]) |
|
|
| rate.sleep() |
| t += 1 |
| if t % 50 == 0: |
| print(f"Step {t} | Recording: {data_collector.get_frame_count()} frames | Press 's' to save") |
| continue |
|
|
| |
| update_observation_window(args, config, ros_operator) |
|
|
| if t % chunk_size == 0: |
| action_buffer = inference_fn(args, config, policy).copy() |
|
|
| raw_action = action_buffer[t % chunk_size] |
|
|
| if args.use_kalman_filter: |
| action = np.array([kf.update(raw_action[i]) for i, kf in enumerate(kalman_filters)]) |
| else: |
| action = raw_action |
|
|
| if args.use_actions_interpolation: |
| if args.interpolate_method == "linear": |
| interp_actions = interpolate_action(args, pre_action, action) |
| elif args.interpolate_method == "minimum_jerk": |
| interp_actions = minimum_jerk_interpolation(args, pre_action, action) |
| else: |
| raise NotImplementedError |
| else: |
| interp_actions = action[np.newaxis, :] |
|
|
| for act in interp_actions: |
| if args.ctrl_type == "joint": |
| left_action = act[:7] |
| right_action = act[7:14] |
| if args.gripper_threshold: |
| if left_action[-1] < 0.03: left_action[-1] = 0 |
| if right_action[-1] < 0.03: right_action[-1] = 0 |
| ros_operator.puppet_arm_publish(left_action, right_action) |
| elif args.ctrl_type == "eef": |
| ros_operator.endpose_publish(act[:7], act[7:14]) |
| if args.use_robot_base: |
| ros_operator.robot_base_publish(act[14:16]) |
| rate.sleep() |
|
|
| t += 1 |
| if t % 50 == 0: |
| print(f"Published Step {t}") |
| pre_action = action.copy() |
|
|
|
|
| class RosOperator: |
| def __init__(self, args): |
| self.robot_base_deque = None |
| self.puppet_arm_right_deque = None |
| self.puppet_arm_left_deque = None |
| self.img_front_deque = None |
| self.img_right_deque = None |
| self.img_left_deque = None |
| self.img_front_depth_deque = None |
| self.img_right_depth_deque = None |
| self.img_left_depth_deque = None |
| self.bridge = None |
| self.puppet_arm_left_publisher = None |
| self.puppet_arm_right_publisher = None |
| self.endpose_left_publisher = None |
| self.endpose_right_publisher = None |
| self.robot_base_publisher = None |
| self.puppet_arm_publish_thread = None |
| self.puppet_arm_publish_lock = None |
| self.args = args |
| self.last_frame_cache = None |
| self.init() |
| self.init_ros() |
|
|
| def init(self): |
| self.bridge = CvBridge() |
| self.img_left_deque = deque() |
| self.img_right_deque = deque() |
| self.img_front_deque = deque() |
| self.img_left_depth_deque = deque() |
| self.img_right_depth_deque = deque() |
| self.img_front_depth_deque = deque() |
| self.puppet_arm_left_deque = deque() |
| self.puppet_arm_right_deque = deque() |
| self.robot_base_deque = deque() |
| self.puppet_arm_publish_lock = threading.Lock() |
| self.puppet_arm_publish_lock.acquire() |
|
|
| def open_grippers(self): |
| left_arm = list(self.puppet_arm_left_deque[-1].position) |
| right_arm = list(self.puppet_arm_right_deque[-1].position) |
| left_arm[6] = 1.0 |
| right_arm[6] = 1.0 |
| self.puppet_arm_publish(left_arm, right_arm) |
| print("Grippers opened.") |
|
|
| def puppet_arm_publish(self, left, right): |
| joint_state_msg = JointState() |
| joint_state_msg.header = Header() |
| joint_state_msg.header.stamp = rospy.Time.now() |
| joint_state_msg.name = ["joint0", "joint1", "joint2", "joint3", "joint4", "joint5", "joint6"] |
| joint_state_msg.position = left |
| self.puppet_arm_left_publisher.publish(joint_state_msg) |
| joint_state_msg.position = right |
| self.puppet_arm_right_publisher.publish(joint_state_msg) |
|
|
| def endpose_publish(self, left, right): |
| endpose_msg = PosCmd() |
| endpose_msg.x, endpose_msg.y, endpose_msg.z = left[:3] |
| endpose_msg.roll, endpose_msg.pitch, endpose_msg.yaw = left[3:6] |
| endpose_msg.gripper = left[6] |
| self.endpose_left_publisher.publish(endpose_msg) |
| endpose_msg.x, endpose_msg.y, endpose_msg.z = right[:3] |
| endpose_msg.roll, endpose_msg.pitch, endpose_msg.yaw = right[3:6] |
| endpose_msg.gripper = right[6] |
| self.endpose_right_publisher.publish(endpose_msg) |
|
|
| def robot_base_publish(self, vel): |
| vel_msg = Twist() |
| vel_msg.linear.x = vel[0] |
| vel_msg.linear.y = 0 |
| vel_msg.linear.z = 0 |
| vel_msg.angular.x = 0 |
| vel_msg.angular.y = 0 |
| vel_msg.angular.z = vel[1] |
| self.robot_base_publisher.publish(vel_msg) |
|
|
| def puppet_arm_publish_continuous(self, left, right): |
| rate = rospy.Rate(self.args.publish_rate) |
| left_arm = right_arm = None |
| while True and not rospy.is_shutdown(): |
| if len(self.puppet_arm_left_deque) != 0: left_arm = list(self.puppet_arm_left_deque[-1].position) |
| if len(self.puppet_arm_right_deque) != 0: right_arm = list(self.puppet_arm_right_deque[-1].position) |
| if left_arm is None or right_arm is None: |
| rate.sleep() |
| continue |
| else: |
| break |
| left_symbol = [1 if left[i] - left_arm[i] > 0 else -1 for i in range(len(left))] |
| right_symbol = [1 if right[i] - right_arm[i] > 0 else -1 for i in range(len(right))] |
| flag = True |
| step = 0 |
| while flag and not rospy.is_shutdown(): |
| if self.puppet_arm_publish_lock.acquire(False): |
| return |
| left_diff = [abs(left[i] - left_arm[i]) for i in range(len(left))] |
| right_diff = [abs(right[i] - right_arm[i]) for i in range(len(right))] |
| flag = False |
| for i in range(len(left)): |
| if left_diff[i] < self.args.arm_steps_length[i]: |
| left_arm[i] = left[i] |
| else: |
| left_arm[i] += left_symbol[i] * self.args.arm_steps_length[i] |
| flag = True |
| for i in range(len(right)): |
| if right_diff[i] < self.args.arm_steps_length[i]: |
| right_arm[i] = right[i] |
| else: |
| right_arm[i] += right_symbol[i] * self.args.arm_steps_length[i] |
| flag = True |
| joint_state_msg = JointState() |
| joint_state_msg.header = Header() |
| joint_state_msg.header.stamp = rospy.Time.now() |
| joint_state_msg.name = ["joint0", "joint1", "joint2", "joint3", "joint4", "joint5", "joint6"] |
| joint_state_msg.position = left_arm |
| self.puppet_arm_left_publisher.publish(joint_state_msg) |
| joint_state_msg.position = right_arm |
| self.puppet_arm_right_publisher.publish(joint_state_msg) |
| step += 1 |
| print("puppet_arm_publish_continuous:", step) |
| rate.sleep() |
|
|
| def puppet_arm_publish_linear(self, left, right): |
| num_step = 100 |
| rate = rospy.Rate(200) |
| left_arm = right_arm = None |
| while True and not rospy.is_shutdown(): |
| if len(self.puppet_arm_left_deque) != 0: left_arm = list(self.puppet_arm_left_deque[-1].position) |
| if len(self.puppet_arm_right_deque) != 0: right_arm = list(self.puppet_arm_right_deque[-1].position) |
| if left_arm is None or right_arm is None: |
| rate.sleep() |
| continue |
| else: |
| break |
| traj_left_list = np.linspace(left_arm, left, num_step) |
| traj_right_list = np.linspace(right_arm, right, num_step) |
| for i in range(len(traj_left_list)): |
| traj_left = traj_left_list[i] |
| traj_right = traj_right_list[i] |
| traj_left[-1] = left[-1] |
| traj_right[-1] = right[-1] |
| joint_state_msg = JointState() |
| joint_state_msg.header = Header() |
| joint_state_msg.header.stamp = rospy.Time.now() |
| joint_state_msg.name = ["joint0", "joint1", "joint2", "joint3", "joint4", "joint5", "joint6"] |
| joint_state_msg.position = traj_left |
| self.puppet_arm_left_publisher.publish(joint_state_msg) |
| joint_state_msg.position = traj_right |
| self.puppet_arm_right_publisher.publish(joint_state_msg) |
| rate.sleep() |
|
|
| def puppet_arm_publish_continuous_thread(self, left, right): |
| if self.puppet_arm_publish_thread is not None: |
| self.puppet_arm_publish_lock.release() |
| self.puppet_arm_publish_thread.join() |
| self.puppet_arm_publish_lock.acquire(False) |
| self.puppet_arm_publish_thread = None |
| self.puppet_arm_publish_thread = threading.Thread( |
| target=self.puppet_arm_publish_continuous, args=(left, right) |
| ) |
| self.puppet_arm_publish_thread.start() |
|
|
| def get_frame(self): |
| if ( |
| len(self.img_left_deque) == 0 |
| or len(self.img_right_deque) == 0 |
| or len(self.img_front_deque) == 0 |
| or (self.args.use_depth_image and ( |
| len(self.img_left_depth_deque) == 0 |
| or len(self.img_right_depth_deque) == 0 |
| or len(self.img_front_depth_deque) == 0 |
| )) |
| ): |
| return False |
|
|
| if self.args.use_depth_image: |
| frame_time = min([ |
| self.img_left_deque[-1].header.stamp.to_sec(), |
| self.img_right_deque[-1].header.stamp.to_sec(), |
| self.img_front_deque[-1].header.stamp.to_sec(), |
| self.img_left_depth_deque[-1].header.stamp.to_sec(), |
| self.img_right_depth_deque[-1].header.stamp.to_sec(), |
| self.img_front_depth_deque[-1].header.stamp.to_sec(), |
| ]) |
| else: |
| frame_time = min([ |
| self.img_left_deque[-1].header.stamp.to_sec(), |
| self.img_right_deque[-1].header.stamp.to_sec(), |
| self.img_front_deque[-1].header.stamp.to_sec(), |
| ]) |
|
|
| if len(self.img_left_deque) == 0 or self.img_left_deque[-1].header.stamp.to_sec() < frame_time: return False |
| if len(self.img_right_deque) == 0 or self.img_right_deque[-1].header.stamp.to_sec() < frame_time: return False |
| if len(self.img_front_deque) == 0 or self.img_front_deque[-1].header.stamp.to_sec() < frame_time: return False |
| if len(self.puppet_arm_left_deque) == 0 or self.puppet_arm_left_deque[-1].header.stamp.to_sec() < frame_time: return False |
| if len(self.puppet_arm_right_deque) == 0 or self.puppet_arm_right_deque[-1].header.stamp.to_sec() < frame_time: return False |
| if self.args.use_depth_image and (len(self.img_left_depth_deque) == 0 or self.img_left_depth_deque[-1].header.stamp.to_sec() < frame_time): return False |
| if self.args.use_depth_image and (len(self.img_right_depth_deque) == 0 or self.img_right_depth_deque[-1].header.stamp.to_sec() < frame_time): return False |
| if self.args.use_depth_image and (len(self.img_front_depth_deque) == 0 or self.img_front_depth_deque[-1].header.stamp.to_sec() < frame_time): return False |
| if self.args.use_robot_base and (len(self.robot_base_deque) == 0 or self.robot_base_deque[-1].header.stamp.to_sec() < frame_time): return False |
|
|
| while self.img_left_deque[0].header.stamp.to_sec() < frame_time: self.img_left_deque.popleft() |
| img_left = self.bridge.imgmsg_to_cv2(self.img_left_deque.popleft(), "passthrough") |
| while self.img_right_deque[0].header.stamp.to_sec() < frame_time: self.img_right_deque.popleft() |
| img_right = self.bridge.imgmsg_to_cv2(self.img_right_deque.popleft(), "passthrough") |
| while self.img_front_deque[0].header.stamp.to_sec() < frame_time: self.img_front_deque.popleft() |
| img_front = self.bridge.imgmsg_to_cv2(self.img_front_deque.popleft(), "passthrough") |
| while self.puppet_arm_left_deque[0].header.stamp.to_sec() < frame_time: self.puppet_arm_left_deque.popleft() |
| puppet_arm_left = self.puppet_arm_left_deque.popleft() |
| while self.puppet_arm_right_deque[0].header.stamp.to_sec() < frame_time: self.puppet_arm_right_deque.popleft() |
| puppet_arm_right = self.puppet_arm_right_deque.popleft() |
|
|
| img_left_depth = img_right_depth = img_front_depth = None |
| if self.args.use_depth_image: |
| while self.img_left_depth_deque[0].header.stamp.to_sec() < frame_time: self.img_left_depth_deque.popleft() |
| img_left_depth = self.bridge.imgmsg_to_cv2(self.img_left_depth_deque.popleft(), "passthrough") |
| while self.img_right_depth_deque[0].header.stamp.to_sec() < frame_time: self.img_right_depth_deque.popleft() |
| img_right_depth = self.bridge.imgmsg_to_cv2(self.img_right_depth_deque.popleft(), "passthrough") |
| while self.img_front_depth_deque[0].header.stamp.to_sec() < frame_time: self.img_front_depth_deque.popleft() |
| img_front_depth = self.bridge.imgmsg_to_cv2(self.img_front_depth_deque.popleft(), "passthrough") |
|
|
| robot_base = None |
| if self.args.use_robot_base: |
| while self.robot_base_deque[0].header.stamp.to_sec() < frame_time: self.robot_base_deque.popleft() |
| robot_base = self.robot_base_deque.popleft() |
|
|
| result_tuple = (img_front, img_left, img_right, |
| img_front_depth, img_left_depth, img_right_depth, |
| puppet_arm_left, puppet_arm_right, robot_base) |
| self.last_frame_cache = result_tuple |
| return result_tuple |
|
|
| def get_frame_peek(self): |
| """Thread-safe version that snapshots all deques to avoid concurrent modification.""" |
| if ( |
| len(self.img_left_deque) == 0 |
| or len(self.img_right_deque) == 0 |
| or len(self.img_front_deque) == 0 |
| or (self.args.use_depth_image and ( |
| len(self.img_left_depth_deque) == 0 |
| or len(self.img_right_depth_deque) == 0 |
| or len(self.img_front_depth_deque) == 0 |
| )) |
| ): |
| return False |
|
|
| |
| try: |
| img_left_snapshot = list(self.img_left_deque) |
| img_right_snapshot = list(self.img_right_deque) |
| img_front_snapshot = list(self.img_front_deque) |
| puppet_left_snapshot = list(self.puppet_arm_left_deque) |
| puppet_right_snapshot = list(self.puppet_arm_right_deque) |
| if self.args.use_depth_image: |
| img_left_depth_snapshot = list(self.img_left_depth_deque) |
| img_right_depth_snapshot = list(self.img_right_depth_deque) |
| img_front_depth_snapshot = list(self.img_front_depth_deque) |
| if self.args.use_robot_base: |
| robot_base_snapshot = list(self.robot_base_deque) |
| except Exception: |
| |
| return False |
|
|
| |
| if self.args.use_depth_image: |
| frame_time = min([ |
| img_left_snapshot[-1].header.stamp.to_sec(), |
| img_right_snapshot[-1].header.stamp.to_sec(), |
| img_front_snapshot[-1].header.stamp.to_sec(), |
| img_left_depth_snapshot[-1].header.stamp.to_sec(), |
| img_right_depth_snapshot[-1].header.stamp.to_sec(), |
| img_front_depth_snapshot[-1].header.stamp.to_sec(), |
| ]) |
| else: |
| frame_time = min([ |
| img_left_snapshot[-1].header.stamp.to_sec(), |
| img_right_snapshot[-1].header.stamp.to_sec(), |
| img_front_snapshot[-1].header.stamp.to_sec(), |
| ]) |
|
|
| |
| for snapshot in [img_left_snapshot, img_right_snapshot, img_front_snapshot, |
| puppet_left_snapshot, puppet_right_snapshot]: |
| if len(snapshot) == 0 or snapshot[-1].header.stamp.to_sec() < frame_time: |
| return False |
| if self.args.use_depth_image: |
| for snapshot in [img_left_depth_snapshot, img_right_depth_snapshot, img_front_depth_snapshot]: |
| if len(snapshot) == 0 or snapshot[-1].header.stamp.to_sec() < frame_time: |
| return False |
| if self.args.use_robot_base: |
| if len(robot_base_snapshot) == 0 or robot_base_snapshot[-1].header.stamp.to_sec() < frame_time: |
| return False |
|
|
| def first_ge(snapshot_list, t): |
| """Find first item in snapshot_list with timestamp >= t.""" |
| for item in snapshot_list: |
| if item.header.stamp.to_sec() >= t: |
| return item |
| return snapshot_list[-1] if snapshot_list else None |
|
|
| |
| img_left_msg = first_ge(img_left_snapshot, frame_time) |
| img_right_msg = first_ge(img_right_snapshot, frame_time) |
| img_front_msg = first_ge(img_front_snapshot, frame_time) |
| puppet_arm_left = first_ge(puppet_left_snapshot, frame_time) |
| puppet_arm_right = first_ge(puppet_right_snapshot, frame_time) |
|
|
| |
| img_left = self.bridge.imgmsg_to_cv2(img_left_msg, "passthrough") |
| img_right = self.bridge.imgmsg_to_cv2(img_right_msg, "passthrough") |
| img_front = self.bridge.imgmsg_to_cv2(img_front_msg, "passthrough") |
|
|
| |
| img_left_depth = img_right_depth = img_front_depth = None |
| if self.args.use_depth_image: |
| img_left_depth = self.bridge.imgmsg_to_cv2(first_ge(img_left_depth_snapshot, frame_time), "passthrough") |
| img_right_depth = self.bridge.imgmsg_to_cv2(first_ge(img_right_depth_snapshot, frame_time), "passthrough") |
| img_front_depth = self.bridge.imgmsg_to_cv2(first_ge(img_front_depth_snapshot, frame_time), "passthrough") |
|
|
| |
| robot_base = None |
| if self.args.use_robot_base: |
| robot_base = first_ge(robot_base_snapshot, frame_time) |
|
|
| return (img_front, img_left, img_right, |
| img_front_depth, img_left_depth, img_right_depth, |
| puppet_arm_left, puppet_arm_right, robot_base) |
|
|
| def img_left_callback(self, msg): |
| if len(self.img_left_deque) >= 2000: self.img_left_deque.popleft() |
| self.img_left_deque.append(msg) |
|
|
| def img_right_callback(self, msg): |
| if len(self.img_right_deque) >= 2000: self.img_right_deque.popleft() |
| self.img_right_deque.append(msg) |
|
|
| def img_front_callback(self, msg): |
| if len(self.img_front_deque) >= 2000: self.img_front_deque.popleft() |
| self.img_front_deque.append(msg) |
|
|
| def img_left_depth_callback(self, msg): |
| if len(self.img_left_depth_deque) >= 2000: self.img_left_depth_deque.popleft() |
| self.img_left_depth_deque.append(msg) |
|
|
| def img_right_depth_callback(self, msg): |
| if len(self.img_right_depth_deque) >= 2000: self.img_right_depth_deque.popleft() |
| self.img_right_depth_deque.append(msg) |
|
|
| def img_front_depth_callback(self, msg): |
| if len(self.img_front_depth_deque) >= 2000: self.img_front_depth_deque.popleft() |
| self.img_front_depth_deque.append(msg) |
|
|
| def puppet_arm_left_callback(self, msg): |
| if len(self.puppet_arm_left_deque) >= 2000: self.puppet_arm_left_deque.popleft() |
| self.puppet_arm_left_deque.append(msg) |
|
|
| def puppet_arm_right_callback(self, msg): |
| if len(self.puppet_arm_right_deque) >= 2000: self.puppet_arm_right_deque.popleft() |
| self.puppet_arm_right_deque.append(msg) |
|
|
| def robot_base_callback(self, msg): |
| if len(self.robot_base_deque) >= 2000: self.robot_base_deque.popleft() |
| self.robot_base_deque.append(msg) |
|
|
| def init_ros(self): |
| rospy.init_node("joint_state_publisher", anonymous=True) |
| rospy.Subscriber(self.args.img_left_topic, Image, self.img_left_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.img_right_topic, Image, self.img_right_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.img_front_topic, Image, self.img_front_callback, queue_size=1000, tcp_nodelay=True) |
| if self.args.use_depth_image: |
| rospy.Subscriber(self.args.img_left_depth_topic, Image, self.img_left_depth_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.img_right_depth_topic, Image, self.img_right_depth_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.img_front_depth_topic, Image, self.img_front_depth_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.puppet_arm_left_topic, JointState, self.puppet_arm_left_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.puppet_arm_right_topic, JointState, self.puppet_arm_right_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.robot_base_topic, Odometry, self.robot_base_callback, queue_size=1000, tcp_nodelay=True) |
| self.puppet_arm_left_publisher = rospy.Publisher(self.args.puppet_arm_left_cmd_topic, JointState, queue_size=10) |
| self.puppet_arm_right_publisher = rospy.Publisher(self.args.puppet_arm_right_cmd_topic, JointState, queue_size=10) |
| self.endpose_left_publisher = rospy.Publisher(self.args.endpose_left_cmd_topic, PosCmd, queue_size=10) |
| self.endpose_right_publisher = rospy.Publisher(self.args.endpose_right_cmd_topic, PosCmd, queue_size=10) |
| self.robot_base_publisher = rospy.Publisher(self.args.robot_base_cmd_topic, Twist, queue_size=10) |
|
|
|
|
| def get_arguments(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--max_publish_step", type=int, default=10000) |
| parser.add_argument("--seed", type=int, default=None) |
| parser.add_argument("--img_front_topic", type=str, default="/camera_f/color/image_raw") |
| parser.add_argument("--img_left_topic", type=str, default="/camera_l/color/image_raw") |
| parser.add_argument("--img_right_topic", type=str, default="/camera_r/color/image_raw") |
| parser.add_argument("--img_front_depth_topic", type=str, default="/camera_f/depth/image_raw") |
| parser.add_argument("--img_left_depth_topic", type=str, default="/camera_l/depth/image_raw") |
| parser.add_argument("--img_right_depth_topic", type=str, default="/camera_r/depth/image_raw") |
| parser.add_argument("--puppet_arm_left_cmd_topic", type=str, default="/master/joint_left") |
| parser.add_argument("--puppet_arm_right_cmd_topic", type=str, default="/master/joint_right") |
| parser.add_argument("--puppet_arm_left_topic", type=str, default="/puppet/joint_left") |
| parser.add_argument("--puppet_arm_right_topic", type=str, default="/puppet/joint_right") |
| parser.add_argument("--endpose_left_cmd_topic", type=str, default="/pos_cmd_left") |
| parser.add_argument("--endpose_right_cmd_topic", type=str, default="/pos_cmd_right") |
| parser.add_argument("--robot_base_topic", type=str, default="/odom_raw") |
| parser.add_argument("--robot_base_cmd_topic", type=str, default="/cmd_vel") |
| parser.add_argument("--use_robot_base", action="store_true", default=False) |
| parser.add_argument("--publish_rate", type=int, default=30) |
| |
| parser.add_argument("--use_temporal_smoothing", action="store_true", default=False) |
| parser.add_argument("--latency_k", type=int, default=8) |
| parser.add_argument("--inference_rate", type=float, default=3.0) |
| parser.add_argument("--min_smooth_steps", type=int, default=8) |
| parser.add_argument("--buffer_max_chunks", type=int, default=10) |
| parser.add_argument("--exp_decay_alpha", type=float, default=0.25) |
| parser.add_argument("--gripper_threshold", action="store_true", default=True) |
| parser.add_argument("--chunk_size", type=int, default=50) |
| parser.add_argument("--arm_steps_length", type=float, default=[0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.2]) |
| parser.add_argument("--use_actions_interpolation", action="store_true", default=False) |
| parser.add_argument("--use_kalman_filter", action="store_true", default=True) |
| parser.add_argument("--interpolate_method", type=str, choices=["linear", "minimum_jerk"], default="linear") |
| parser.add_argument("--jerk_num_steps", type=int, default=10) |
| parser.add_argument("--use_depth_image", action="store_true", default=False) |
| |
| parser.add_argument("--export_video", action="store_true", default=True) |
| parser.add_argument("--video_fps", type=int, default=30) |
| parser.add_argument("--video_codec", type=str, choices=["libx264", "libx265", "libsvtav1"], default="libx264") |
| parser.add_argument("--video_quality", type=int, default=23, help="CRF value; lower = higher quality") |
| |
| parser.add_argument("--host", type=str, default="localhost") |
| parser.add_argument("--port", type=int, default=8000) |
| parser.add_argument("--ctrl_type", type=str, choices=["joint", "eef"], default="joint") |
| |
| parser.add_argument("--dataset_dir", type=str, default="/home/agilex/data") |
| parser.add_argument("--task_name", type=str, default="aloha_inference") |
| parser.add_argument("--episode_idx", type=int, default=0) |
| parser.add_argument("--max_timesteps", type=int, default=100000000) |
| parser.add_argument("--frame_rate", type=int, default=30) |
| parser.add_argument("--camera_names", nargs="+", default=["cam_high", "cam_left_wrist", "cam_right_wrist"]) |
| return parser.parse_args() |
|
|
|
|
| def main(): |
| args = get_arguments() |
| try: |
| ros_operator = RosOperator(args) |
| if args.seed is not None: |
| set_seed(args.seed) |
| config = get_config(args) |
| model_inference(args, config, ros_operator) |
| except KeyboardInterrupt: |
| print("\n⚠️ Interrupted, exiting...") |
| except Exception as e: |
| print(f"\n❌ Error: {e}") |
| import traceback |
| traceback.print_exc() |
| finally: |
| print("Program exited") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
