"""Dataset recording utility functions for teleoperation data collection.""" import argparse import time import traceback from pathlib import Path from typing import Dict, Optional, Tuple, Any, Union import gymnasium as gym import numpy as np import torch from isaacsim.simulation_app import SimulationApp from isaaclab.envs import DirectRLEnv from isaaclab_tasks.utils import parse_env_cfg from lerobot.datasets.lerobot_dataset import LeRobotDataset from lehome.devices import ( Se3Keyboard, SO101Leader, BiSO101Leader, BiKeyboard, ) from lehome.utils.env_utils import dynamic_reset_gripper_effort_limit_sim from lehome.utils.record import ( get_next_experiment_path_with_gap, append_episode_initial_pose, ) from lehome.utils.logger import get_logger from .common import stabilize_garment_after_reset logger = get_logger(__name__) def validate_task_and_device(args: argparse.Namespace) -> None: """Validate that task name matches the teleop device configuration. Args: args: Command-line arguments containing task and teleop_device. Raises: ValueError: If task is not specified. AssertionError: If task and device configuration mismatch. """ if args.task is None: raise ValueError("Please specify --task.") if "Bi" in args.task: assert ( args.teleop_device == "bi-so101leader" or args.teleop_device == "bi-keyboard" ), "Only support bi-so101leader or bi-keyboard for bi-arm task" else: assert ( args.teleop_device == "so101leader" or args.teleop_device == "keyboard" ), "Only support so101leader or keyboard for single-arm task" def create_teleop_interface( env: DirectRLEnv, args: argparse.Namespace ) -> Union[Se3Keyboard, SO101Leader, BiSO101Leader, BiKeyboard]: """Create teleoperation interface based on device type. Args: env: Environment instance. args: Command-line arguments containing teleop_device and related config. Returns: Teleoperation interface instance. Raises: ValueError: If teleop_device is invalid. """ if args.teleop_device == "keyboard": return Se3Keyboard(env, sensitivity=0.25 * args.sensitivity) if args.teleop_device == "so101leader": return SO101Leader(env, port=args.port, recalibrate=args.recalibrate) if args.teleop_device == "bi-so101leader": return BiSO101Leader( env, left_port=args.left_arm_port, right_port=args.right_arm_port, recalibrate=args.recalibrate, ) if args.teleop_device == "bi-keyboard": return BiKeyboard(env, sensitivity=0.25 * args.sensitivity) raise ValueError( f"Invalid device interface '{args.teleop_device}'. " f"Supported: 'keyboard', 'so101leader', 'bi-so101leader', 'bi-keyboard'." ) def register_teleop_callbacks( teleop_interface: Any, recording_enabled: bool = False ) -> Dict[str, bool]: """Register callback functions for teleoperation control keys. Key bindings: S: Start recording N: Mark current episode as successful (only active during recording) D: Discard current episode and re-record (only active during recording) ESC: Abort entire recording process and clear buffer Args: teleop_interface: Teleoperation interface instance. recording_enabled: Whether recording is enabled. If False, N/D keys are disabled in idle phase. Returns: Dictionary of status flags for recording control. """ flags = { "start": False, # S: Start recording "success": False, # N: Success/early termination of current episode "remove": False, # D: Discard current episode "abort": False, # ESC: Abort entire recording process, clear buffer } def on_start(): flags["start"] = True logger.info("[S] Recording started!") def on_success(): if not recording_enabled or not flags["start"]: # Ignore N key in idle phase (before recording starts) logger.debug("[N] Ignored (recording not started yet)") return flags["success"] = True logger.info("[N] Mark the current episode as successful.") def on_remove(): if not recording_enabled or not flags["start"]: # Ignore D key in idle phase (before recording starts) logger.debug("[D] Ignored (recording not started yet)") return flags["remove"] = True logger.info("[D] Discard the current episode and re-record.") def on_abort(): flags["abort"] = True logger.warning("[ESC] Abort recording, clearing the current episode buffer...") teleop_interface.add_callback("S", on_start) teleop_interface.add_callback("N", on_success) teleop_interface.add_callback("D", on_remove) teleop_interface.add_callback("ESCAPE", on_abort) return flags def create_dataset_if_needed( args: argparse.Namespace, ) -> Tuple[Optional[LeRobotDataset], Optional[Path], Optional[Any], bool]: """Create LeRobotDataset if recording is enabled. Args: args: Command-line arguments containing recording configuration. Returns: Tuple of (dataset, json_path, solver, is_bi_arm): - dataset: LeRobotDataset instance or None if not recording - json_path: Path to object initial pose JSON file or None - solver: RobotKinematics solver instance or None - is_bi_arm: Boolean indicating if dual-arm configuration Raises: ValueError: If record_ee_pose is enabled but ee_urdf_path is not provided. FileNotFoundError: If URDF file is not found. """ if not args.enable_record: return None, None, None, False action_names = [ "shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper", ] is_bi_arm = ("Bi" in (args.task or "")) or ( getattr(args, "teleop_device", "") or "" ).startswith("bi-") if is_bi_arm: left_names = [f"left_{n}" for n in action_names] right_names = [f"right_{n}" for n in action_names] joint_names = left_names + right_names else: joint_names = action_names dim = len(joint_names) features: Dict[str, Dict[str, Any]] = { "observation.state": { "dtype": "float32", "shape": (dim,), "names": joint_names, }, "action": { "dtype": "float32", "shape": (dim,), "names": joint_names, }, } if not getattr(args, "disable_depth", False): features["observation.top_depth"] = { "dtype": "uint16", "shape": (480, 640), "names": ["height", "width"], "info": { "unit": "millimeters", "range_mm": [0, 65535], "range_m": [0.0, 65.535], "precision_mm": 1, "conversion": "depth_meters = uint16_value / 1000.0" } } if is_bi_arm: image_keys = ["top_rgb", "left_rgb", "right_rgb"] else: image_keys = ["top_rgb", "wrist_rgb"] for key in image_keys: features[f"observation.images.{key}"] = { "dtype": "video", "shape": (480, 640, 3), "names": ["height", "width", "channels"], } if getattr(args, "record_ee_pose", False): if is_bi_arm: ee_pose_dim = 16 ee_pose_names = [ "left_x", "left_y", "left_z", "left_qx", "left_qy", "left_qz", "left_qw", "left_gripper", "right_x", "right_y", "right_z", "right_qx", "right_qy", "right_qz", "right_qw", "right_gripper", ] else: ee_pose_dim = 8 ee_pose_names = ["x", "y", "z", "qx", "qy", "qz", "qw", "gripper"] features["observation.ee_pose"] = { "dtype": "float32", "shape": (ee_pose_dim,), "names": ee_pose_names, } features["action.ee_pose"] = { "dtype": "float32", "shape": (ee_pose_dim,), "names": ee_pose_names, } root_path = Path(getattr(args, "dataset_root", "Datasets/record")) dataset = LeRobotDataset.create( repo_id="abc", fps=30, root=get_next_experiment_path_with_gap(root_path), use_videos=True, image_writer_threads=8, image_writer_processes=0, features=features, ) json_path = dataset.root / "meta" / "garment_info.json" solver = None if getattr(args, "record_ee_pose", False): if not args.ee_urdf_path: raise ValueError("--record_ee_pose requires --ee_urdf_path") urdf_path = Path(args.ee_urdf_path) if not urdf_path.exists(): raise FileNotFoundError(f"URDF not found: {urdf_path}") from lehome.utils import RobotKinematics if is_bi_arm: solver_joint_names = [n.replace("left_", "") for n in joint_names[:5]] else: solver_joint_names = joint_names[:5] solver = RobotKinematics( str(urdf_path), target_frame_name="gripper_frame_link", joint_names=solver_joint_names, ) arm_type = "dual-arm" if is_bi_arm else "single-arm" logger.info(f"End-effector pose solver loaded ({arm_type})") return dataset, json_path, solver, is_bi_arm def run_idle_phase( env: DirectRLEnv, teleop_interface: Any, args: argparse.Namespace, count_render: int, ) -> Tuple[Optional[Dict[str, Any]], int]: """Run idle phase before recording starts. Handles environment preparation, stabilization, and waits for user to press S key to start recording. Args: env: Environment instance. teleop_interface: Teleoperation interface. args: Command-line arguments. count_render: Current render count. Returns: Tuple of (object_initial_pose, updated_count_render). """ dynamic_reset_gripper_effort_limit_sim(env, args.teleop_device) actions = teleop_interface.advance() object_initial_pose = None if count_render == 0: logger.info("[Idle Phase] Initializing observations...") env.initialize_obs() count_render += 1 logger.info("[Idle Phase] Stabilizing garment after initialization...") stabilize_garment_after_reset(env, args) logger.info("[Idle Phase] Ready for recording") if actions is None: current_obs = env._get_observations() if "observation.state" in current_obs: current_state = current_obs["observation.state"] if isinstance(current_state, np.ndarray): maintain_action = ( torch.from_numpy(current_state).float().unsqueeze(0).to(env.device) ) else: maintain_action = torch.zeros( 1, len(current_state), dtype=torch.float32, device=env.device ) else: action_dim = 12 if "Bi" in args.task else 6 maintain_action = torch.zeros( 1, action_dim, dtype=torch.float32, device=env.device ) env.step(maintain_action) env.render() else: env.step(actions) object_initial_pose = env.get_all_pose() if object_initial_pose is None: object_initial_pose = env.get_all_pose() return object_initial_pose, count_render def run_recording_phase( env: DirectRLEnv, teleop_interface: Any, args: argparse.Namespace, flags: Dict[str, bool], dataset: LeRobotDataset, json_path: Path, initial_object_pose: Optional[Dict[str, Any]], ee_solver: Optional[Any] = None, is_bi_arm: bool = False, ) -> Dict[str, Any]: """Run recording phase after S key is pressed and recording is enabled. Records episodes until num_episode is reached. Each episode can be marked as successful (N key), discarded (D key), or aborted (ESC key). Args: env: Environment instance. teleop_interface: Teleoperation interface. args: Command-line arguments. flags: Status flags dictionary. dataset: LeRobotDataset instance. json_path: Path to object initial pose JSON file. initial_object_pose: Initial object pose dictionary. ee_solver: Optional kinematic solver for end-effector pose computation. is_bi_arm: Whether using dual-arm configuration. Returns: Final object initial pose dictionary. """ episode_index = 0 object_initial_pose = initial_object_pose # Ensure we have a valid initial pose for the first episode if object_initial_pose is None: object_initial_pose = env.get_all_pose() while episode_index < args.num_episode: # Check if recording should be aborted if flags["abort"]: dataset.clear_episode_buffer() dataset.finalize() logger.warning(f"Recording aborted, completed {episode_index} episodes") return object_initial_pose flags["success"] = False flags["remove"] = False # Loop within a single episode while not flags["success"]: # Check if recording should be aborted if flags["abort"]: dataset.clear_episode_buffer() dataset.finalize() logger.warning(f"Recording aborted, completed {episode_index} episodes") return object_initial_pose try: dynamic_reset_gripper_effort_limit_sim(env, args.teleop_device) actions = teleop_interface.advance() except Exception as e: logger.error(f"[Recording] Error in teleop interface: {e}") traceback.print_exc() actions = None if actions is None: env.render() else: env.step(actions) if args.log_success: success = env._get_success() observations = env._get_observations() if ( getattr(args, "disable_depth", False) and "observation.top_depth" in observations ): observations.pop("observation.top_depth") if getattr(args, "enable_pointcloud", False): # Converting pointcloud online is time-consuming, please convert offline # pointcloud = env._get_workspace_pointcloud( # num_points=4096, use_fps=True # ) print("Converting pointcloud online is time-consuming, please convert offline") _, truncated = env._get_dones() frame = {**observations, "task": args.task_description} if ( ee_solver is not None and "observation.state" in observations and "action" in observations ): from lehome.utils import compute_ee_pose_single_arm obs_state = np.array( observations["observation.state"], dtype=np.float32 ) action_state = np.array(observations["action"], dtype=np.float32) if is_bi_arm: obs_left = compute_ee_pose_single_arm( ee_solver, obs_state[:6], args.ee_state_unit ) obs_right = compute_ee_pose_single_arm( ee_solver, obs_state[6:12], args.ee_state_unit ) frame["observation.ee_pose"] = np.concatenate( [obs_left, obs_right], axis=0 ) act_left = compute_ee_pose_single_arm( ee_solver, action_state[:6], args.ee_state_unit ) act_right = compute_ee_pose_single_arm( ee_solver, action_state[6:12], args.ee_state_unit ) frame["action.ee_pose"] = np.concatenate( [act_left, act_right], axis=0 ) else: frame["observation.ee_pose"] = compute_ee_pose_single_arm( ee_solver, obs_state, args.ee_state_unit ) frame["action.ee_pose"] = compute_ee_pose_single_arm( ee_solver, action_state, args.ee_state_unit ) dataset.add_frame(frame) if truncated or flags["remove"]: dataset.clear_episode_buffer() logger.info(f"Re-recording episode {episode_index}") try: env.reset() stabilize_garment_after_reset(env, args) object_initial_pose = env.get_all_pose() except Exception as e: logger.error( f"[Recording] Failed to reset environment during re-recording: {e}" ) traceback.print_exc() try: object_initial_pose = env.get_all_pose() except Exception: object_initial_pose = None flags["remove"] = False continue save_start_time = time.time() logger.info(f"[Recording] Saving episode {episode_index}...") try: dataset.save_episode() save_duration = time.time() - save_start_time logger.info( f"[Recording] Episode {episode_index} saved (took {save_duration:.1f}s)" ) except Exception as e: logger.error(f"[Recording] Failed to save episode {episode_index}: {e}") traceback.print_exc() garment_name = None if hasattr(env, "cfg") and hasattr(env.cfg, "garment_name"): garment_name = env.cfg.garment_name scale = None if hasattr(env, "object") and hasattr(env.object, "init_scale"): try: scale = env.object.init_scale except Exception: logger.warning("Failed to get scale from garment object") try: append_episode_initial_pose( json_path, episode_index, object_initial_pose, garment_name=garment_name, scale=scale, ) except Exception as e: logger.error( f"[Recording] Failed to save episode metadata for episode {episode_index}: {e}" ) traceback.print_exc() episode_index += 1 logger.info( f"Episode {episode_index - 1} completed, progress: {episode_index}/{args.num_episode}" ) try: env.reset() stabilize_garment_after_reset(env, args) except Exception as e: logger.error(f"[Recording] Failed to reset environment: {e}") traceback.print_exc() try: object_initial_pose = env.get_all_pose() except Exception as e: logger.error(f"[Recording] Failed to get initial pose: {e}") traceback.print_exc() object_initial_pose = None dataset.clear_episode_buffer() dataset.finalize() logger.info(f"All {args.num_episode} episodes recording completed!") return object_initial_pose def run_live_control_without_record( env: DirectRLEnv, teleop_interface: Any, args: argparse.Namespace, ) -> None: """Run live teleoperation control without recording. Handles the case when S key is pressed but recording is not enabled. Performs simple teleoperation control without writing to dataset. Args: env: Environment instance. teleop_interface: Teleoperation interface. args: Command-line arguments. """ dynamic_reset_gripper_effort_limit_sim(env, args.teleop_device) actions = teleop_interface.advance() if actions is None: current_obs = env._get_observations() if "observation.state" in current_obs: current_state = current_obs["observation.state"] if isinstance(current_state, np.ndarray): maintain_action = ( torch.from_numpy(current_state).float().unsqueeze(0).to(env.device) ) else: maintain_action = torch.zeros( 1, len(current_state), dtype=torch.float32, device=env.device ) else: action_dim = 12 if "Bi" in args.task else 6 maintain_action = torch.zeros( 1, action_dim, dtype=torch.float32, device=env.device ) env.step(maintain_action) env.render() else: env.step(actions) if args.log_success: _ = env._get_success() def record_dataset(args: argparse.Namespace, simulation_app: SimulationApp) -> None: """Record dataset.""" # Get device configuration (default to "cpu" for compatibility) device = getattr(args, "device", "cpu") env_cfg = parse_env_cfg( args.task, device=device, ) task_name = args.task env_cfg.garment_name = args.garment_name env_cfg.garment_version = args.garment_version env_cfg.garment_cfg_base_path = args.garment_cfg_base_path env_cfg.particle_cfg_path = args.particle_cfg_path if args.use_random_seed: env_cfg.use_random_seed = True logger.info("Using random seed (no fixed seed)") else: env_cfg.use_random_seed = False env_cfg.random_seed = args.seed logger.info(f"Using fixed random seed: {args.seed}") env: DirectRLEnv = gym.make(task_name, cfg=env_cfg).unwrapped teleop_interface = create_teleop_interface(env, args) flags = register_teleop_callbacks( teleop_interface, recording_enabled=args.enable_record ) teleop_interface.reset() dataset, json_path, ee_solver, is_bi_arm = create_dataset_if_needed(args) count_render = 0 printed_instructions = False idle_frame_counter = 0 object_initial_pose: Optional[Dict[str, Any]] = None try: while simulation_app.is_running(): with torch.inference_mode(): if not flags["start"]: pose, count_render = run_idle_phase( env, teleop_interface, args, count_render, ) if pose is not None: object_initial_pose = pose if count_render > 0: idle_frame_counter += 1 if idle_frame_counter == 100 and not printed_instructions: logger.info("=" * 60) logger.info("🎮 CONTROL INSTRUCTIONS 🎮") logger.info("=" * 60) logger.info(str(teleop_interface)) logger.info("=" * 60 + "\n\n") printed_instructions = True elif args.enable_record and dataset is not None: object_initial_pose = run_recording_phase( env, teleop_interface, args, flags, dataset, json_path, object_initial_pose, ee_solver, is_bi_arm, ) break else: run_live_control_without_record(env, teleop_interface, args) except KeyboardInterrupt: logger.warning("\n[Ctrl+C] Interrupt signal detected") # If Ctrl+C is pressed during recording, clear the current buffer if args.enable_record and dataset is not None and flags["start"]: logger.info("Clearing current episode buffer...") dataset.clear_episode_buffer() logger.info("Buffer cleared, dataset remains intact") dataset.finalize() logger.info("Dataset saved") except Exception as e: logger.error(f"An unexpected error occurred: {e}") finally: env.close()