From a21f53e91afb34c0cf4f221715501b90b736af72 Mon Sep 17 00:00:00 2001 From: chuchen0210 <167320956+chuchen0210@users.noreply.github.com> Date: Wed, 7 Jan 2026 11:58:50 +0800 Subject: [PATCH 1/9] Add LeIsaac + Cosmos + IDM tutorial --- .../cosmos_tutorials/generate_batch_config.py | 125 ++++ .../cosmos2.5_to_step2_format.py | 279 +++++++++ .../preprocess_leisaac/dump_idm_actions.py | 481 +++++++++++++++ .../lerobot_to_step2_format.py | 161 +++++ .../preprocess_leisaac/preprocess_video.py | 376 ++++++++++++ .../preprocess_leisaac/raw_to_lerobot.py | 568 ++++++++++++++++++ .../split_video_instruction.py | 101 ++++ scripts/tutorials/cosmos_tutorials/replay.py | 495 +++++++++++++++ scripts/tutorials/cosmos_tutorials/so101.sh | 149 +++++ 9 files changed, 2735 insertions(+) create mode 100644 scripts/tutorials/cosmos_tutorials/generate_batch_config.py create mode 100644 scripts/tutorials/cosmos_tutorials/preprocess_leisaac/cosmos2.5_to_step2_format.py create mode 100644 scripts/tutorials/cosmos_tutorials/preprocess_leisaac/dump_idm_actions.py create mode 100644 scripts/tutorials/cosmos_tutorials/preprocess_leisaac/lerobot_to_step2_format.py create mode 100644 scripts/tutorials/cosmos_tutorials/preprocess_leisaac/preprocess_video.py create mode 100644 scripts/tutorials/cosmos_tutorials/preprocess_leisaac/raw_to_lerobot.py create mode 100644 scripts/tutorials/cosmos_tutorials/preprocess_leisaac/split_video_instruction.py create mode 100644 scripts/tutorials/cosmos_tutorials/replay.py create mode 100644 scripts/tutorials/cosmos_tutorials/so101.sh diff --git a/scripts/tutorials/cosmos_tutorials/generate_batch_config.py b/scripts/tutorials/cosmos_tutorials/generate_batch_config.py new file mode 100644 index 00000000..ec17ba87 --- /dev/null +++ b/scripts/tutorials/cosmos_tutorials/generate_batch_config.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +""" +Generate batch inference configuration files (JSONL). + +This script scans a directory of input videos and generates a JSONL file +for batch video2world inference. Each line corresponds to one video. +Optionally, task prompts can be loaded from episodes.jsonl. +""" + +import json +import os +from pathlib import Path + +# --------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------- + +# Input video directory used to generate inference configs +VIDEO_DIR = "/videos/chunk-000/observation.images.front" + +# Metadata directory (used when loading text prompts) +META_DIR = "/meta" + +# Output JSONL file +OUTPUT_JSONL = "batch_inference_config.jsonl" + +# Task name prefix +TASK_NAME = "liftcube" + +# --------------------------------------------------------------------- +# Base configuration template shared by all videos +# --------------------------------------------------------------------- + +BASE_CONFIG = { + "inference_type": "video2world", + "seed": 21, + "guidance": 7, + "resolution": "480,640", + "enable_autoregressive": True, + # Number of output frames (e.g. 110 ≈ 6s, 140 ≈ 8s, 210 ≈ 16s) + "num_output_frames": 210, + "chunk_size": 77, + "chunk_overlap": 1, + # Default prompt (will be overridden if episode-specific prompt is found) + "prompt": "The robot arm is performing a task", + "negative_prompt": ( + "The video captures a series of frames showing ugly scenes, static with no motion, " + "motion blur, over-saturation, shaky footage, low resolution, grainy texture, " + "pixelated images, poorly lit areas, underexposed and overexposed scenes, " + "poor color balance, washed out colors, choppy sequences, jerky movements, " + "low frame rate, artifacting, color banding, unnatural transitions, " + "outdated special effects, fake elements, unconvincing visuals, " + "poorly edited content, jump cuts, visual noise, and flickering. " + "Overall, the video is of poor quality." + ), +} + +# --------------------------------------------------------------------- +# Main logic +# --------------------------------------------------------------------- + +def main(use_prompt=True): + video_dir = Path(VIDEO_DIR) + meta_dir = Path(META_DIR) + + # Collect all mp4 video files + video_files = sorted(video_dir.glob("*.mp4")) + print(f"Found {len(video_files)} video files") + + # Load prompts from episodes.jsonl if enabled + episode_prompts = {} + if use_prompt: + episodes_file = meta_dir / "episodes.jsonl" + if episodes_file.exists(): + with open(episodes_file, 'r') as f: + for line in f: + episode_data = json.loads(line) + episode_index = episode_data.get("episode_index") + tasks = episode_data.get("tasks", []) + if episode_index is not None and tasks: + # Use the first task as the episode prompt + prompt = f"The robot arm is performing a task. {tasks[0]}" + episode_prompts[episode_index] = prompt + print(f"✓ Loaded {len(episode_prompts)} prompts from episodes.jsonl") + + # Generate JSONL configuration file + with open(OUTPUT_JSONL, 'w') as f: + for video_file in video_files: + # Video filename without extension (e.g. episode_000001) + video_name = video_file.stem + + # Create per-video config based on the base template + config = BASE_CONFIG.copy() + config["name"] = f"{TASK_NAME}_{video_name}" + config["input_path"] = str(video_file) + + # Assign episode-specific prompt if available + if use_prompt: + # Extract episode index from video filename + episode_index = int(video_name.split('_')[-1]) + if episode_index in episode_prompts: + config["prompt"] = episode_prompts[episode_index] + + # Write one JSON object per line + f.write(json.dumps(config) + '\n') + + # Summary information + print(f"✓ Config file generated: {OUTPUT_JSONL}") + print(f" - Total videos: {len(video_files)}") + print(f" - Use prompts: {'Yes' if use_prompt else 'No'}") + if use_prompt: + print( + f" - Videos with matched prompts: " + f"{sum(1 for vf in video_files if int(vf.stem.split('_')[-1]) in episode_prompts)}" + ) + +# --------------------------------------------------------------------- +# Script entry point +# --------------------------------------------------------------------- + +if __name__ == "__main__": + import sys + # Text prompts are enabled only when '--use-prompt' is provided + use_prompt = "--use-prompt" in sys.argv + main(use_prompt=use_prompt) diff --git a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/cosmos2.5_to_step2_format.py b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/cosmos2.5_to_step2_format.py new file mode 100644 index 00000000..8d5ffbce --- /dev/null +++ b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/cosmos2.5_to_step2_format.py @@ -0,0 +1,279 @@ +#!/usr/bin/env python3 +""" +Convert cosmos2.5 outputs (+ optional LeRobot meta) to convert_directory output format. + +Input: + cosmos_dir/ + ├── *.mp4 + ├── *.json(same stem as mp4) + └── ... + +Optional: + lerobot_dir/ + └── meta/ + ├── tasks.jsonl + └── episodes.jsonl + +Output: + output_dir/ + ├── / + │ ├── 0.mp4 + │ ├── 1.mp4 + │ └── ... + └── ... + +Rule: + - If json has "prompt": use it as task name (sanitize to dir) + - Else: fallback to LeRobot meta (requires --lerobot_dir) + - If lerobot_dir not provided: print warning and skip that sample +""" + +import argparse +import json +import shutil +from pathlib import Path +from collections import defaultdict +from typing import Dict, List, Optional, Tuple + + +def sanitize_task_name(task: str) -> str: + """ + Convert task description to valid directory name. + Example: "Lift the red cube up." -> "Lift_the_red_cube_up" + """ + task = task.strip() + for ch in [".", ",", "!", "?", ":", ";", "\"", "'"]: + task = task.replace(ch, "") + task = task.replace(" ", "_") + task = "_".join(filter(None, task.split("_"))) + return task + + +def load_tasks(tasks_file: Path) -> Dict[int, str]: + """Load task definitions. Returns {task_index: task_string}""" + tasks: Dict[int, str] = {} + with open(tasks_file, "r") as f: + for line in f: + if not line.strip(): + continue + data = json.loads(line) + tasks[int(data["task_index"])] = data["task"] + return tasks + + +def load_episodes(episodes_file: Path) -> List[dict]: + """Load episode information. Returns list of dict (episode_index, tasks, length, ...)""" + episodes: List[dict] = [] + with open(episodes_file, "r") as f: + for line in f: + if not line.strip(): + continue + episodes.append(json.loads(line)) + return episodes + + +def build_episode_to_task_map(lerobot_dir: Path) -> Dict[int, str]: + """ + Build mapping: episode_index -> task_string (first task in episode['tasks']) + """ + meta_dir = lerobot_dir / "meta" + tasks_path = meta_dir / "tasks.jsonl" + episodes_path = meta_dir / "episodes.jsonl" + + if not tasks_path.exists() or not episodes_path.exists(): + raise FileNotFoundError(f"LeRobot meta files not found under: {meta_dir}") + + _tasks = load_tasks(tasks_path) # not strictly needed, but useful if episodes store indices in some setups + episodes = load_episodes(episodes_path) + + ep2task: Dict[int, str] = {} + for ep in episodes: + ep_idx = int(ep["episode_index"]) + ep_tasks = ep.get("tasks", []) + if not ep_tasks: + continue + + # In many LeRobot datasets, ep["tasks"] stores task strings already. + # If it stores indices, you can extend here. We'll support both: + t0 = ep_tasks[0] + if isinstance(t0, int): + task_str = _tasks.get(int(t0), str(t0)) + else: + task_str = str(t0) + + ep2task[ep_idx] = task_str + + return ep2task + + +def read_json(path: Path) -> Optional[dict]: + try: + with open(path, "r") as f: + return json.load(f) + except Exception as e: + print(f"[WARN] Failed to read json: {path} ({e})") + return None + + +def parse_episode_index_from_stem(stem: str) -> Optional[int]: + """ + Try to extract episode index from filename stem. + Examples it can handle: + - episode_000123 + - ..._000123 + - 000123 + If cannot parse, return None. + """ + # Most common: "episode_000123" + if stem.startswith("episode_"): + tail = stem[len("episode_") :] + if tail.isdigit(): + return int(tail) + + # Try last underscore chunk + parts = stem.split("_") + for candidate in reversed(parts): + if candidate.isdigit(): + return int(candidate) + + # Entire stem digits? + if stem.isdigit(): + return int(stem) + + return None + + +def convert_cosmos_to_step2( + cosmos_dir: Path, + output_dir: Path, + lerobot_dir: Optional[Path] = None, + chunk_missing_prompt_policy: str = "skip", +): + """ + Convert cosmos2.5 outputs to step2 format. + + chunk_missing_prompt_policy: + - "skip": if no prompt and no lerobot_dir mapping, skip the sample + """ + cosmos_dir = cosmos_dir.resolve() + output_dir = output_dir.resolve() + output_dir.mkdir(parents=True, exist_ok=True) + + ep2task: Dict[int, str] = {} + if lerobot_dir is not None: + ep2task = build_episode_to_task_map(lerobot_dir.resolve()) + + # Collect pairs (json, mp4) by stem + json_files = sorted(cosmos_dir.glob("*.json")) + if not json_files: + raise FileNotFoundError(f"No .json files found in cosmos_dir: {cosmos_dir}") + + # task_name -> list of source mp4 paths (ordered) + task_to_videos: Dict[str, List[Path]] = defaultdict(list) + skipped: List[Tuple[Path, str]] = [] + + for jpath in json_files: + stem = jpath.stem + mpath = cosmos_dir / f"{stem}.mp4" + if not mpath.exists(): + print(f"[WARN] Missing mp4 for json: {jpath.name} -> expected {mpath.name}, skip.") + skipped.append((jpath, "missing_mp4")) + continue + + data = read_json(jpath) + if data is None: + skipped.append((jpath, "bad_json")) + continue + + prompt = data.get("prompt", None) + if isinstance(prompt, str) and prompt.strip(): + task_str = prompt.strip() + task_dir_name = sanitize_task_name(task_str) + task_to_videos[task_dir_name].append(mpath) + continue + + # No prompt -> fallback to lerobot meta mapping + if not ep2task: + print( + f"[WARN] {jpath.name} has no 'prompt'. " + f"You didn't provide --lerobot_dir (or mapping is empty), cannot infer task. Skipping." + ) + skipped.append((jpath, "no_prompt_no_lerobot")) + continue + + ep_idx = parse_episode_index_from_stem(stem) + if ep_idx is None: + print( + f"[WARN] {jpath.name} has no 'prompt' and episode index cannot be parsed from name '{stem}'. Skipping." + ) + skipped.append((jpath, "no_prompt_cannot_parse_episode")) + continue + + task_str = ep2task.get(ep_idx) + if not task_str: + print( + f"[WARN] {jpath.name} has no 'prompt'. Parsed episode_index={ep_idx}, " + f"but it's not found in lerobot meta. Skipping." + ) + skipped.append((jpath, "episode_not_in_meta")) + continue + + task_dir_name = sanitize_task_name(task_str) + task_to_videos[task_dir_name].append(mpath) + + # Copy into output folders with sequential numbering per task + total_copied = 0 + for task_dir_name, vids in sorted(task_to_videos.items(), key=lambda x: x[0]): + dst_task_dir = output_dir / task_dir_name + dst_task_dir.mkdir(parents=True, exist_ok=True) + + # keep deterministic order + vids_sorted = sorted(vids, key=lambda p: p.name) + for i, src in enumerate(vids_sorted): + dst = dst_task_dir / f"{i}.mp4" + shutil.copy2(src, dst) + total_copied += 1 + + print(f"[OK] Task '{task_dir_name}': {len(vids_sorted)} videos") + + print("\nConversion complete!") + print(f" Total videos copied: {total_copied}") + print(f" Output directory: {output_dir}") + + if skipped: + print(f"\n[SUMMARY] Skipped {len(skipped)} samples:") + # print a few for readability + for p, reason in skipped[:20]: + print(f" - {p.name}: {reason}") + if len(skipped) > 20: + print(f" ... and {len(skipped) - 20} more") + + +def main(): + parser = argparse.ArgumentParser( + description="Convert cosmos2.5 outputs (+ optional LeRobot meta) to convert_directory output format" + ) + parser.add_argument("--cosmos_dir", type=str, required=True, help="Cosmos output directory containing *.mp4 and *.json") + parser.add_argument("--output_dir", type=str, required=True, help="Output directory") + parser.add_argument( + "--lerobot_dir", + type=str, + default=None, + help="Optional LeRobot dataset directory (only needed when json has no 'prompt')", + ) + + args = parser.parse_args() + + cosmos_dir = Path(args.cosmos_dir) + output_dir = Path(args.output_dir) + lerobot_dir = Path(args.lerobot_dir) if args.lerobot_dir else None + + convert_cosmos_to_step2( + cosmos_dir=cosmos_dir, + output_dir=output_dir, + lerobot_dir=lerobot_dir, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/dump_idm_actions.py b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/dump_idm_actions.py new file mode 100644 index 00000000..003c27b0 --- /dev/null +++ b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/dump_idm_actions.py @@ -0,0 +1,481 @@ +import argparse +import json +import multiprocessing as mp +import os + +import numpy as np +import torch +import yaml +from hydra.utils import instantiate +from omegaconf import OmegaConf +from tqdm import tqdm +from concurrent.futures import ProcessPoolExecutor, as_completed +from tianshou.data import Batch +from huggingface_hub import hf_hub_download +from gr00t.data.dataset import LeRobotSingleDataset +from gr00t.model.idm import IDM +from gr00t.utils.video import get_all_frames_and_timestamps +from gr00t.data.embodiment_tags import EmbodimentTag + + + +def load_dataset_and_config(checkpoint_path, validation_dataset_path, video_indices): + # Check if checkpoint_path is a HuggingFace model repo + is_hf_repo = not os.path.exists(checkpoint_path) and '/' in checkpoint_path + + if is_hf_repo: + # For HuggingFace repos, we need to download config differently + try: + # Download the config file from the repo + config_file = hf_hub_download( + repo_id=checkpoint_path, + filename="experiment_cfg/conf.yaml", + repo_type="model" + ) + with open(config_file, "r") as f: + config = yaml.safe_load(f) + except Exception as e: + print(f"Error loading config from HuggingFace repo: {e}") + raise + else: + # Original local path handling + config_path = os.path.join(checkpoint_path, "experiment_cfg", "conf.yaml") + with open(config_path, "r") as f: + config = yaml.safe_load(f) + + cfg = OmegaConf.create(config) + + dataset_name = os.path.basename(validation_dataset_path).split(".")[0] + embodiment = dataset_name + print(f"Dataset name: {dataset_name}") + + modality_configs = cfg["modality_configs"][embodiment] + if video_indices is not None: + video_delta_indices = video_indices.split(" ") + print(f"Using provided video_delta_indices: {video_delta_indices}") + modality_configs["video"]["delta_indices"] = video_delta_indices + modality_configs = instantiate(modality_configs) + + + if "all_transforms" in cfg: + transform = cfg["all_transforms"][embodiment] + else: + transform = cfg["train_dataset"]["all_transforms"][embodiment] + + # Filter out VideoColorJitter transform for inference + if "transforms" in transform: + filtered_transforms = [] + for t in transform["transforms"]: + if t.get("_target_") != "groot.data.transform.VideoColorJitter": + filtered_transforms.append(t) + + transform["transforms"] = filtered_transforms + transform_inst = instantiate(transform) + + metadata_versions = cfg["metadata_versions"] + metadata_version = metadata_versions[embodiment] + + if "gr1" in embodiment: + embodiment_tag = EmbodimentTag.GR1_unified + elif "franka" in embodiment: + embodiment_tag = EmbodimentTag.FRANKA + elif "so101" in embodiment: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + elif "so100" in embodiment: + embodiment_tag = EmbodimentTag.SO100 + elif "robocasa" in embodiment: + embodiment_tag = EmbodimentTag.ROBOCASA + elif "piper" in embodiment: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + elif "piper double" in embodiment: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + else: + raise ValueError(f"Unknown embodiment: {embodiment}") + + dataset = LeRobotSingleDataset( + dataset_path=validation_dataset_path, + modality_configs=modality_configs, + # metadata_version=metadata_version, + transforms=transform_inst, + embodiment_tag=embodiment_tag, + ) + + return cfg, dataset, modality_configs + + +def collate_fn(features_list, device): + batch_dict = {} + keys = features_list[0].keys() + for key in keys: + if key in ["images", "view_ids"]: + vals = [f[key] for f in features_list] + batch_dict[key] = torch.as_tensor(np.concatenate(vals), device=device) + else: + vals = [f[key] for f in features_list] + batch_dict[key] = torch.as_tensor(np.stack(vals), device=device) + + return batch_dict + +def get_step_data(dataset, trajectory_id, base_index): + data = {} + dataset.curr_traj_data = dataset.get_trajectory_data(trajectory_id) + # Get the data for all modalities + for modality in dataset.modality_keys: + # Get the data corresponding to each key in the modality + for key in dataset.modality_keys[modality]: + if modality == "video": + pass + elif modality == "state" or modality == "action": + data[key] = dataset.get_state_or_action(trajectory_id, modality, key, base_index) + elif modality == "language": + data[key] = dataset.get_language(trajectory_id, key, base_index) + + return data + +def save_trajectory_data(trajectory_data, dataset, trajectory_id, output_dir): + chunk_index = dataset.get_episode_chunk(trajectory_id) + chunk_dir = f"chunk-{chunk_index:03d}" + os.makedirs(os.path.join(output_dir, "data", chunk_dir), exist_ok=True) + + episode_id = f"episode_{int(trajectory_id):06d}" + output_file_path = os.path.join(output_dir, "data", chunk_dir, f"{episode_id}.parquet") + trajectory_data.to_parquet(output_file_path) + +## +# Worker function, now handles either validation or writing based on a boolean flag. +## +def worker_func( + gpu_id: int, + traj_id_list: list, + checkpoint_path: str, + validation_dataset_path: str, + output_dir: str, + batch_size: int, + num_workers: int = 1, + video_indices=None, + dataset=None, + modality_configs=None, +): + """This function runs in a separate process on GPU `gpu_id` and processes all `traj_id_list`.""" + # Load model on this GPU + # from_pretrained should handle both local paths and HF model IDs + print(f"Loading model from: {checkpoint_path}") + model = IDM.from_pretrained(checkpoint_path) + + model.requires_grad_(False) + model.eval() + device = torch.device(f"cuda:{gpu_id}") + model.to(device) + + + for tid in tqdm(traj_id_list, desc=f"GPU {gpu_id}", position=gpu_id): + action_dict = {} + traj_data = dataset.get_trajectory_data(tid) + length = len(traj_data) + + + all_features = [] + + # Create a simple prefetching mechanism with a thread pool + from concurrent.futures import ThreadPoolExecutor + + def load_and_transform_step(step_idx, video_data): + import torch + from einops import rearrange + + step_data = get_step_data(dataset, tid, step_idx) + timestamp = dataset.curr_traj_data["timestamp"].to_numpy() + for key in video_data: + frames, whole_indices = video_data[key] + step_indices = dataset.delta_indices[key] + step_idx + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, dataset.trajectory_lengths[tid] - 1) + indices = np.array([np.where(np.isclose(whole_indices, val))[0][0] for val in timestamp[step_indices]]) + step_data[key] = frames[indices] + + output = dataset.transforms(step_data) + return output + + + video_data = {} + for key in dataset.modality_keys["video"]: + video_path = dataset.get_video_path(tid, key.replace("video.", "")) + video_backend = dataset.video_backend + video_backend_kwargs = dataset.video_backend_kwargs + frames, whole_indices = get_all_frames_and_timestamps(video_path.as_posix(), video_backend, video_backend_kwargs) + video_data[key] = (frames, whole_indices) + + + with ThreadPoolExecutor(max_workers=num_workers) as executor: + futures = [] + for step_idx in range(length): + future = executor.submit(load_and_transform_step, step_idx, video_data) + futures.append(future) + + all_features = [] + for future in as_completed(futures): + all_features.append(future.result()) + + + for start_idx in range(0, length, batch_size): + end_idx = min(start_idx + batch_size, length) + step_ids = list(range(start_idx, end_idx)) + + batch_features = all_features[start_idx:end_idx] + batch_dict = collate_fn(batch_features, device) + + with torch.no_grad(): + out = model.get_action(batch_dict) + + pred_actions = out["action_pred"].cpu() + pred_actions = dataset.transforms.unapply(Batch(action=pred_actions)) + + # Load modality.json to get the proper structure + modality_json_path = os.path.join(validation_dataset_path, 'meta', 'modality.json') + + with open(modality_json_path, 'r') as f: + modality_config = json.load(f) + + # Get action part configurations + action_parts = modality_config.get('action', {}) + + # Calculate total action dimension + total_dim = 0 + for part, indices in action_parts.items(): + total_dim = max(total_dim, indices.get('end', 0)) + + # Check if we have an action horizon dimension + has_horizon = False + action_horizon = 1 + sample_key = list(pred_actions.keys())[0] + if len(pred_actions[sample_key].shape) == 3: + has_horizon = True + batch_size, action_horizon, _ = pred_actions[sample_key].shape + else: + batch_size = pred_actions[sample_key].shape[0] + + if has_horizon: + final_actions = np.zeros((batch_size, action_horizon, total_dim)) + + # Fill in the actions for parts that exist in pred_actions + for part, indices in action_parts.items(): + action_key = f'action.{part}' + start_idx = indices.get('start', 0) + end_idx = indices.get('end', 0) + + if action_key in pred_actions: + final_actions[:, :, start_idx:end_idx] = pred_actions[action_key] + else: + final_actions = np.zeros((batch_size, total_dim)) + + # Fill in the actions for parts that exist in pred_actions + for part, indices in action_parts.items(): + action_key = f'action.{part}' + start_idx = indices.get('start', 0) + end_idx = indices.get('end', 0) + + if action_key in pred_actions: + final_actions[:, start_idx:end_idx] = pred_actions[action_key] + + pred_actions = final_actions + + # Not validating => we do the usual writing + for i, s in enumerate(step_ids): + for j in range(action_horizon): + if s+j >= length: + break + if s+j not in action_dict: + action_dict[s+j] = [] + if has_horizon: + action_dict[s+j].append(pred_actions[i, j].flatten()) + else: + action_dict[s+j].append(pred_actions[i].flatten()) + + for s in action_dict: + mean_action = np.mean(action_dict[s], axis=0) + traj_data.at[s, "action"] = mean_action + + save_trajectory_data(traj_data, dataset, tid, output_dir) + + del model + if torch.cuda.is_available(): + torch.cuda.empty_cache() + +def validate_checkpoint( + checkpoint_path, + validation_dataset_path, + output_dir=None, + num_gpus=8, + batch_size=16, + max_episodes=None, + num_workers=1, + video_indices=None, +): + device_count = torch.cuda.device_count() + print(f"Found {device_count} GPUs available.") + if device_count < num_gpus: + print( + f"WARNING: You requested num_gpus={num_gpus} but only {device_count} GPUs are visible." + ) + num_gpus = device_count + + # If not validation, create the output directory for saving predicted actions + if output_dir: + os.makedirs(output_dir, exist_ok=True) + print(f"Will save predicted actions to: {output_dir}") + + _, dataset, modality_configs = load_dataset_and_config(checkpoint_path, validation_dataset_path, video_indices) + + dataset.transforms.eval() + + trajectory_ids = dataset.trajectory_ids + if max_episodes is not None and max_episodes > 0: + trajectory_ids = trajectory_ids[:max_episodes] + + print(f"Processing {len(trajectory_ids)} trajectories total.") + + # Split trajectories among GPUs + chunk_size = (len(trajectory_ids) + num_gpus - 1) // num_gpus + + tasks_by_gpu = {} + for i in range(num_gpus): + start = i * chunk_size + end = min(start + chunk_size, len(trajectory_ids)) + if start >= end: + break + gpu_traj_list = trajectory_ids[start:end] + tasks_by_gpu[i] = gpu_traj_list + + try: + with ProcessPoolExecutor(max_workers=num_gpus) as executor: + futures = [] + for gpu_id, gpu_traj_list in tasks_by_gpu.items(): + future = executor.submit( + worker_func, + gpu_id, + gpu_traj_list, + checkpoint_path, + validation_dataset_path, + output_dir, + batch_size, + num_workers, + video_indices, + dataset, + modality_configs, + ) + futures.append(future) + + # Wait for all futures to complete or handle interruption + for future in as_completed(futures): + try: + # Get the result to catch any exceptions + future.result() + except Exception as e: + print(f"Error in worker process: {e}") + import traceback + traceback.print_exc() + except KeyboardInterrupt: + print("\nProcess interrupted by user. Cleaning up...") + # The context manager will handle cancellation of pending futures + except Exception as e: + print(f"Error in main process: {e}") + import traceback + traceback.print_exc() + finally: + # Force cleanup of CUDA memory + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # If not validating, optionally copy the meta & videos folders + if output_dir: + # Copy the metadata directory + meta_src = os.path.join(validation_dataset_path, "meta") + meta_dst = os.path.join(output_dir, "meta") + if os.path.exists(meta_src): + import shutil + + if not os.path.exists(meta_dst): + shutil.copytree(meta_src, meta_dst) + print(f"Copied metadata to: {meta_dst}") + + tasks_path = os.path.join(meta_dst, "tasks.jsonl") + if os.path.exists(tasks_path): + # Read the existing tasks + tasks = [] + with open(tasks_path, "r") as f: + for line in f: + tasks.append(json.loads(line)) + + # Update tasks with prefix if not already present + updated_tasks = [] + for task in tasks: + if "task" in task and not task["task"].startswith(""): + task["task"] = f"{task['task']}" + updated_tasks.append(task) + + # Write the updated tasks back to the file + with open(tasks_path, "w") as f: + for task in updated_tasks: + f.write(json.dumps(task) + "\n") + + print("Updated tasks.jsonl with prefix") + + + # Copy the videos directory if it exists + videos_src = os.path.join(validation_dataset_path, "videos") + videos_dst = os.path.join(output_dir, "videos") + if os.path.exists(videos_src): + import shutil + + if not os.path.exists(videos_dst): + shutil.copytree(videos_src, videos_dst) + print(f"Copied videos to: {videos_dst}") + + return {"message": "Dataset with predicted actions written to disk"} + + +# +# Main +# +if __name__ == "__main__": + # If you're doing spawn/forkserver across multiple processes: + mp.set_start_method("spawn", force=True) + + parser = argparse.ArgumentParser( + description="Validate a checkpoint on a dataset with multi-GPU parallelism" + ) + parser.add_argument("--checkpoint", type=str, required=True, help="Path to the checkpoint") + parser.add_argument("--dataset", type=str, required=True, help="Path to the validation dataset") + parser.add_argument( + "--output_dir", type=str, help="Path to save the output dataset with predicted actions" + ) + parser.add_argument("--num_gpus", type=int, default=8, help="Number of GPUs to use in parallel") + parser.add_argument("--batch_size", type=int, default=32, help="Batch size for inference steps") + parser.add_argument( + "--max_episodes", type=int, default=None, help="Maximum number of episodes to validate" + ) + parser.add_argument( + "--num_workers", + type=int, + default=16, + help="Number of worker threads for data loading per GPU", + ) + parser.add_argument( + "--video_indices", + type=str, + default=None, + help="Video frame indices to use for inference", + ) + + args = parser.parse_args() + + result = validate_checkpoint( + checkpoint_path=args.checkpoint, + validation_dataset_path=args.dataset, + output_dir=args.output_dir, + num_gpus=args.num_gpus, + batch_size=args.batch_size, + max_episodes=args.max_episodes, + num_workers=args.num_workers, + video_indices=args.video_indices, + ) diff --git a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/lerobot_to_step2_format.py b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/lerobot_to_step2_format.py new file mode 100644 index 00000000..33ffe696 --- /dev/null +++ b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/lerobot_to_step2_format.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +""" +Convert LeRobot format dataset to convert_directory output format. + +Input format (LeRobot): + dataset_dir/ + ├── videos/chunk-000/observation.images.front/ + │ ├── episode_000000.mp4 + │ ├── episode_000001.mp4 + │ └── ... + └── meta/ + ├── tasks.jsonl (task definitions) + └── episodes.jsonl (episode -> task mapping) + +Output format (after convert_directory): + output_dir/ + ├── Lift_the_red_cube_up/ + │ ├── 0.mp4 + │ ├── 1.mp4 + │ └── ... + └── Another_task/ + ├── 0.mp4 + └── ... + +Usage: + python lerobot_to_step2_format.py \ + --input_dir /path/to/lerobot/dataset \ + --output_dir /path/to/output \ + --video_key observation.images.front +""" + +import argparse +import json +import shutil +from pathlib import Path +from collections import defaultdict + + +def load_tasks(tasks_file: Path) -> dict: + """Load task definitions. Returns {task_index: task_description}""" + tasks = {} + with open(tasks_file, 'r') as f: + for line in f: + data = json.loads(line.strip()) + tasks[data['task_index']] = data['task'] + return tasks + + +def load_episodes(episodes_file: Path) -> list: + """Load episode information. Returns list of {episode_index, tasks, length}""" + episodes = [] + with open(episodes_file, 'r') as f: + for line in f: + data = json.loads(line.strip()) + episodes.append(data) + return episodes + + +def sanitize_task_name(task: str) -> str: + """ + Convert task description to valid directory name. + Example: "Lift the red cube up." -> "Lift_the_red_cube_up" + """ + task = task.replace('.', '').replace(',', '').replace('!', '').replace('?', '') + task = task.replace(' ', '_') + task = '_'.join(filter(None, task.split('_'))) + return task + + +def convert_lerobot_to_step2( + input_dir: Path, + output_dir: Path, + video_key: str = "observation.images.front", + chunk_name: str = "chunk-000" +): + """Convert LeRobot format to convert_directory output format.""" + + # Load tasks and episodes + tasks = load_tasks(input_dir / "meta" / "tasks.jsonl") + episodes = load_episodes(input_dir / "meta" / "episodes.jsonl") + + print(f"Loaded {len(tasks)} tasks") + print(f"Loaded {len(episodes)} episodes") + + video_dir = input_dir / "videos" / chunk_name / video_key + + # Group episodes by task + task_episodes = defaultdict(list) + for episode in episodes: + episode_idx = episode['episode_index'] + episode_tasks = episode.get('tasks', []) + if episode_tasks: + task = episode_tasks[0] + task_episodes[task].append(episode_idx) + + print(f"\nStarting conversion...") + + total_copied = 0 + for task, episode_indices in task_episodes.items(): + # Create task directory + task_dir_name = sanitize_task_name(task) + task_dir = output_dir / task_dir_name + task_dir.mkdir(parents=True, exist_ok=True) + + print(f"\nTask: {task}") + print(f" Directory: {task_dir_name}") + print(f" Episodes: {len(episode_indices)}") + + # Copy video files + for idx, episode_idx in enumerate(sorted(episode_indices)): + src_video = video_dir / f"episode_{episode_idx:06d}.mp4" + dst_video = task_dir / f"{idx}.mp4" + shutil.copy2(src_video, dst_video) + total_copied += 1 + + print(f"\nConversion complete!") + print(f" Total videos copied: {total_copied}") + print(f" Output directory: {output_dir}") + + +def main(): + parser = argparse.ArgumentParser( + description="Convert LeRobot format to convert_directory output format" + ) + parser.add_argument( + "--input_dir", + type=str, + required=True, + help="LeRobot dataset directory" + ) + parser.add_argument( + "--output_dir", + type=str, + required=True, + help="Output directory" + ) + parser.add_argument( + "--video_key", + type=str, + default="observation.images.front", + help="Video key name (default: observation.images.front)" + ) + parser.add_argument( + "--chunk_name", + type=str, + default="chunk-000", + help="Chunk name (default: chunk-000)" + ) + + args = parser.parse_args() + + convert_lerobot_to_step2( + Path(args.input_dir), + Path(args.output_dir), + args.video_key, + args.chunk_name + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/preprocess_video.py b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/preprocess_video.py new file mode 100644 index 00000000..6ddd8755 --- /dev/null +++ b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/preprocess_video.py @@ -0,0 +1,376 @@ +import os +import cv2 +import numpy as np +from pathlib import Path +import multiprocessing as mp +from tqdm import tqdm +import shutil +import argparse +import concurrent.futures +from functools import partial +import decord +import imageio + + +def custom_crop_pad_resize_gr1(img, target_size=(256, 256)): + """ + Custom crop, pad, and resize operation that maintains the aspect ratio. + + Args: + img: Input frame + original_width: Original video width + original_height: Original video height + target_size: Target resolution (width, height) + + Returns: + Processed frame at target_size resolution + """ + # For 832x480 videos, adjust the crop parameters proportionally from the 1280x800 example + # Original crop for 1280x800: (310, 770, 110, 1130) - (top, bottom, left, right) + original_height, original_width = img.shape[:2] + + + # Calculate proportional crop values + crop_top_ratio = 310 / 800 + crop_bottom_ratio = 770 / 800 + crop_left_ratio = 110 / 1280 + crop_right_ratio = 1130 / 1280 + + # Apply ratios to the current dimensions + crop_top = int(original_height * crop_top_ratio) + crop_bottom = int(original_height * crop_bottom_ratio) + crop_left = int(original_width * crop_left_ratio) + crop_right = int(original_width * crop_right_ratio) + + # Ensure crop boundaries are within image dimensions + crop_top = max(0, min(crop_top, original_height - 1)) + crop_bottom = max(crop_top + 1, min(crop_bottom, original_height)) + crop_left = max(0, min(crop_left, original_width - 1)) + crop_right = max(crop_left + 1, min(crop_right, original_width)) + + # Crop the image + img_cropped = img[crop_top:crop_bottom, crop_left:crop_right] + + # Calculate intermediate size while maintaining aspect ratio + cropped_height, cropped_width = img_cropped.shape[:2] + aspect_ratio = cropped_width / cropped_height + + # Resize to intermediate size (similar to 720x480 in the example) + intermediate_height = 480 + intermediate_width = 720 + img_resized = cv2.resize(img_cropped, (intermediate_width, intermediate_height), cv2.INTER_AREA) + + # Pad to make it square + if intermediate_width > intermediate_height: + # Width is larger, pad height + height_pad = (intermediate_width - intermediate_height) // 2 + img_pad = np.pad( + img_resized, ((height_pad, height_pad), (0, 0), (0, 0)), mode="constant", constant_values=0 + ) + else: + # Height is larger, pad width + width_pad = (intermediate_height - intermediate_width) // 2 + img_pad = np.pad( + img_resized, ((0, 0), (width_pad, width_pad), (0, 0)), mode="constant", constant_values=0 + ) + + # Final resize to target size + final_img = cv2.resize(img_pad, target_size, cv2.INTER_AREA) + + return final_img + +def resize_with_padding(img, ratio=1.0, target_size=(256, 256)): + # Original aspect ratio is 1280:800 (or 16:10) + target_ratio = ratio + + h, w = img.shape[:2] + current_ratio = w / h + if target_ratio >= 1: + # Width is the limiting factor + new_w = target_size[0] + new_h = int(new_w / target_ratio) + resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR) + + # Add padding to height + pad_top = (target_size[1] - new_h) // 2 + pad_bottom = target_size[1] - new_h - pad_top + padded = cv2.copyMakeBorder(resized, pad_top, pad_bottom, 0, 0, cv2.BORDER_CONSTANT, value=[0, 0, 0]) + else: + # Height is the limiting factor + new_h = target_size[1] + new_w = int(new_h * target_ratio) + resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR) + + # Add padding to width + pad_left = (target_size[0] - new_w) // 2 + pad_right = target_size[0] - new_w - pad_left + padded = cv2.copyMakeBorder(resized, 0, 0, pad_left, pad_right, cv2.BORDER_CONSTANT, value=[0, 0, 0]) + + return padded + +def extract_subimages_franka(frame, original_width, original_height): + h, w = frame.shape[:2] # h=480, w=832 + + # Calculate dimensions for even division + half_width = w // 2 + half_height = h // 2 + + # Extract subimages + image_side_0 = frame[:half_height, :half_width] # Top-left (240x416) + image_side_1 = frame[:half_height, half_width:] # Top-right (240x416) + wrist_image = frame[half_height:, :half_width] # Bottom-left (240x416) + + + image_side_0 = cv2.resize(image_side_0, (original_width, original_height), interpolation=cv2.INTER_LINEAR) + image_side_1 = cv2.resize(image_side_1, (original_width, original_height), interpolation=cv2.INTER_LINEAR) + wrist_image = cv2.resize(wrist_image, (original_width, original_height), interpolation=cv2.INTER_LINEAR) + + return image_side_0, image_side_1, wrist_image + + +def extract_subimages(frame, ratio): + """Extract subimages from a frame and resize to 256x256 while preserving aspect ratio with padding.""" + h, w = frame.shape[:2] # h=480, w=832 + + # Calculate dimensions for even division + half_width = w // 2 # 416 + half_height = h // 2 # 240 + + # Extract subimages + image_side_0 = frame[:half_height, :half_width] # Top-left (240x416) + image_side_1 = frame[:half_height, half_width:] # Top-right (240x416) + wrist_image = frame[half_height:, :half_width] # Bottom-left (240x416) + + image_side_0 = resize_with_padding(image_side_0, ratio) + image_side_1 = resize_with_padding(image_side_1, ratio) + wrist_image = resize_with_padding(wrist_image, ratio) + + return image_side_0, image_side_1, wrist_image + +def process_batch_frames(frames, output_videos, src_path, dataset, original_width, original_height): + """Process a batch of frames.""" + ratio = original_width / original_height + for frame in frames: + # Extract subimages + if dataset == 'robocasa': + image_side_0, image_side_1, wrist_image = extract_subimages(frame, ratio) + output_videos['observation.images.left_view'].append_data(image_side_0) + output_videos['observation.images.right_view'].append_data(image_side_1) + output_videos['observation.images.wrist_view'].append_data(wrist_image) + elif dataset == 'gr1': + image = custom_crop_pad_resize_gr1(frame) + output_videos['observation.images.ego_view'].append_data(image) + elif dataset == 'franka': + image_side_0, image_side_1, wrist_image = extract_subimages(frame, ratio) + output_videos['observation.images.exterior_image_1_left_pad_res256_freq15'].append_data(image_side_0) + output_videos['observation.images.exterior_image_2_left_pad_res256_freq15'].append_data(image_side_1) + output_videos['observation.images.wrist_image_left_pad_res256_freq15'].append_data(wrist_image) + elif dataset == 'so100': + image = resize_with_padding(frame, ratio) + output_videos['observation.images.front'].append_data(image) + elif dataset == 'so101': + image = resize_with_padding(frame, ratio) + output_videos['observation.images.front'].append_data(image) + elif dataset == 'piper': + # Piper: 720x720 -> resize to 256x256 (square, no padding needed) + image = cv2.resize(frame, (256, 256), interpolation=cv2.INTER_LINEAR) + output_videos['observation.images.agentview_all'].append_data(image) + elif dataset == 'piper double': + image = cv2.resize(frame, (256, 256), interpolation=cv2.INTER_LINEAR) + output_videos['observation.images.agentview_all'].append_data(image) + else: + raise ValueError(f"Unknown task: {src_path}") + +def process_video(args): + """Process a single video file.""" + src_path, dst_dir, video_name, dataset, original_width, original_height = args + + # Create output directories if they don't exist + if dataset == 'robocasa': + output_dirs = { + 'observation.images.left_view': os.path.join(dst_dir, 'videos', 'observation.images.left_view'), + 'observation.images.right_view': os.path.join(dst_dir, 'videos', 'observation.images.right_view'), + 'observation.images.wrist_view': os.path.join(dst_dir, 'videos', 'observation.images.wrist_view'), + } + elif dataset == 'gr1': + output_dirs = { + 'observation.images.ego_view': os.path.join(dst_dir, 'videos', 'observation.images.ego_view') + } + elif dataset == 'franka': + output_dirs = { + 'observation.images.exterior_image_1_left_pad_res256_freq15': os.path.join(dst_dir, 'videos', 'observation.images.exterior_image_1_left_pad_res256_freq15'), + 'observation.images.exterior_image_2_left_pad_res256_freq15': os.path.join(dst_dir, 'videos', 'observation.images.exterior_image_2_left_pad_res256_freq15'), + 'observation.images.wrist_image_left_pad_res256_freq15': os.path.join(dst_dir, 'videos', 'observation.images.wrist_image_left_pad_res256_freq15'), + } + elif dataset == 'so100': + output_dirs = { + 'observation.images.front': os.path.join(dst_dir, 'videos', 'observation.images.front'), + } + elif dataset == 'so101': + output_dirs = { + 'observation.images.front': os.path.join(dst_dir, 'videos', 'observation.images.front'), + } + elif dataset == 'piper': + output_dirs = { + 'observation.images.agentview_all': os.path.join(dst_dir, 'videos', 'observation.images.agentview_all'), + } + elif dataset == 'piper double': + output_dirs = { + 'observation.images.agentview_all': os.path.join(dst_dir, 'videos', 'observation.images.agentview_all'), + } + + for dir_path in output_dirs.values(): + os.makedirs(dir_path, exist_ok=True) + + # Open the source video + vr = decord.VideoReader(src_path) + + # Get video properties + fps = vr.get_avg_fps() + frame_count = len(vr) + + # Create VideoWriter objects for each subimage + output_videos = {} + for name, dir_path in output_dirs.items(): + output_path = os.path.join(dir_path, f"{video_name}.mp4") + output_videos[name] = imageio.get_writer(output_path, fps=fps) + + + # Process frames in batches + batch_size = 32 # Adjust based on available memory + frames_batch = [] + + # Process each frame with progress bar + pbar = tqdm(total=frame_count, desc=f"Processing {video_name}", leave=False) + + for frame in vr: + frames_batch.append(frame.asnumpy()) + + # Process batch when it reaches the desired size + if len(frames_batch) >= batch_size: + process_batch_frames(frames_batch, output_videos, src_path, dataset, original_width, original_height) + frames_batch = [] + pbar.update(batch_size) + + # Process remaining frames + if frames_batch: + process_batch_frames(frames_batch, output_videos, src_path, dataset, original_width, original_height) + pbar.update(len(frames_batch)) + + # Close progress bar + pbar.close() + + # Release resources + for writer in output_videos.values(): + writer.close() + + +def copy_labels(src_dir, dst_dir): + """Copy label files from source to destination.""" + src_labels_dir = os.path.join(src_dir, 'labels') + dst_labels_dir = os.path.join(dst_dir, 'labels') + + if os.path.exists(src_labels_dir): + os.makedirs(dst_labels_dir, exist_ok=True) + for label_file in os.listdir(src_labels_dir): + if label_file.endswith('.txt'): + shutil.copy2( + os.path.join(src_labels_dir, label_file), + os.path.join(dst_labels_dir, label_file) + ) + +def process_subdirectory(subdir, src_dir, dst_dir, num_workers, max_videos=None, dataset=None, original_width=None, original_height=None): + """Process a single subdirectory.""" + print(f"Processing subdirectory: {src_dir}, {subdir}") + src_subdir = os.path.join(src_dir, subdir) + dst_subdir = os.path.join(dst_dir, subdir) + + # Copy label files + copy_labels(src_subdir, dst_subdir) + + # Process videos + src_videos_dir = os.path.join(src_subdir, 'videos') + if os.path.exists(src_videos_dir): + video_files = [f for f in os.listdir(src_videos_dir) if f.endswith('.mp4')] + + # Sort files numerically to ensure consistent ordering (1.mp4, 2.mp4, ..., 10.mp4, not 1.mp4, 10.mp4, 2.mp4) + def sort_key(filename): + stem = os.path.splitext(filename)[0] + try: + return int(stem) + except ValueError: + return float('inf') + video_files.sort(key=sort_key) + + # Limit number of videos if max_videos is specified + if max_videos is not None: + video_files = video_files[:max_videos] + print(f"Processing limited set of {len(video_files)} videos in {subdir}") + + # Prepare arguments for multiprocessing + args_list = [ + (os.path.join(src_videos_dir, video_file), dst_subdir, os.path.splitext(video_file)[0], dataset, original_width, original_height) + for video_file in video_files + ] + + # Process videos in parallel + with mp.Pool(num_workers) as pool: + list(tqdm(pool.imap(process_video, args_list), total=len(args_list), desc=f"Processing {subdir}")) + +def process_directory(src_dir, dst_dir, num_workers=None, num_subdirs_parallel=1, max_videos=None, dataset=None, original_width=None, original_height=None, recursive=False): + """Process all videos in the source directory.""" + if num_workers is None: + num_workers = max(1, mp.cpu_count() // 2) # Leave some cores for system + + # Create destination directory structure + os.makedirs(dst_dir, exist_ok=True) + + + if recursive: + # Get all subdirectories in the source directory + subdirs = [d for d in os.listdir(src_dir) if os.path.isdir(os.path.join(src_dir, d))] + else: + # Get all subdirectories in the source directory + subdirs = [''] + + # Process subdirectories in parallel + with concurrent.futures.ThreadPoolExecutor(max_workers=num_subdirs_parallel) as executor: + process_subdir_fn = partial(process_subdirectory, + src_dir=src_dir, + dst_dir=dst_dir, + num_workers=num_workers // num_subdirs_parallel, + max_videos=max_videos, + dataset=dataset, + original_width=original_width, + original_height=original_height) + list(tqdm(executor.map(process_subdir_fn, subdirs), total=len(subdirs), desc="Processing subdirectories")) + +def main(): + parser = argparse.ArgumentParser(description='Split videos into subimages and save them.') + parser.add_argument('--src_dir', type=str, default='/mnt/amlfs-01/home/seonghyeony/data/dreams_robocasa_70K_0228', + help='Source directory containing videos') + parser.add_argument('--dst_dir', type=str, default='/mnt/amlfs-01/home/seonghyeony/data/dreams_robocasa_70K_0228_split', + help='Destination directory for processed videos') + parser.add_argument('--workers', type=int, default=8, + help='Number of worker processes (default: half of CPU cores)') + parser.add_argument('--parallel_subdirs', type=int, default=1, + help='Number of subdirectories to process in parallel') + parser.add_argument('--batch_size', type=int, default=32, + help='Number of frames to process in a batch') + parser.add_argument('--max_videos', type=int, default=None, + help='Maximum number of videos to process per subdirectory (for debugging)') + parser.add_argument('--dataset', type=str, default='robocasa', + help='Dataset name', choices=['robocasa', 'gr1', 'franka', 'so100', 'so101', 'piper', 'piper double']) + parser.add_argument("--recursive", action="store_true", help="Process subdirectories recursively, maintaining directory structure") + parser.add_argument("--original_width", type=int, default=1280, help="Original width of the video") + parser.add_argument("--original_height", type=int, default=800, help="Original height of the video") + + args = parser.parse_args() + # Set OpenCV thread optimization + cv2.setNumThreads(1) # Disable OpenCV's internal multithreading to avoid conflicts + + process_directory(args.src_dir, args.dst_dir, args.workers, args.parallel_subdirs, + max_videos=args.max_videos, dataset=args.dataset, original_width=args.original_width, original_height=args.original_height, recursive=args.recursive) + +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/raw_to_lerobot.py b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/raw_to_lerobot.py new file mode 100644 index 00000000..07d3d1ca --- /dev/null +++ b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/raw_to_lerobot.py @@ -0,0 +1,568 @@ +import json +import shutil +from pathlib import Path +from typing import Dict, List, Any, Tuple +import concurrent.futures +import os + +import numpy as np +import pandas as pd +from tqdm import tqdm +import argparse + +import subprocess + +import multiprocessing +from multiprocessing import Pool +import math + + +CHUNKS_SIZE = 1000 +DATA_PATH = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet" +VIDEO_PATH = "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4" + + +def get_video_metadata(video_path): + """Get video metadata using ffprobe.""" + cmd = [ + "ffprobe", "-v", "error", + "-select_streams", "v:0", + "-show_entries", "stream=height,width,codec_name,pix_fmt,r_frame_rate", + "-of", "json", video_path, + ] + + try: + output = subprocess.check_output(cmd).decode("utf-8") + probe_data = json.loads(output) + stream = probe_data["streams"][0] + + # Parse frame rate + num, den = map(int, stream["r_frame_rate"].split("/")) + fps = num / den + + return { + "dtype": "video", + "shape": [stream["height"], stream["width"], 3], + "names": ["height", "width", "channel"], + "video_info": { + "video.fps": fps, + "video.codec": stream["codec_name"], + "video.pix_fmt": stream["pix_fmt"], + "video.is_depth_map": False, + }, + } + except Exception as e: + print(f"Error getting video metadata: {e}") + return None + +def dump_jsonl(data: List[Dict[str, Any]], path: Path) -> None: + """Dump list of dictionaries as JSONL file.""" + with open(path, 'w') as f: + for item in data: + json_str = json.dumps(item) + f.write(json_str + '\n') + +def json_dump(data: Dict[str, Any], path: Path, indent: int = 4) -> None: + """Dump dictionary as JSON file.""" + with open(path, 'w') as f: + json.dump(data, f, indent=indent) + + +def process_video_chunk(args): + """Process a chunk of videos in parallel.""" + video_files, labels_dir, output_dir, cosmos_predict2, data_type, videos_dir, video_key = args + results = [] + + for video_file in video_files: + video_id = video_file.stem + label_file = labels_dir / f"{video_id}.txt" + + # Read annotation + with open(label_file, "r") as f: + annotation = f.read().strip() + + # Get video frame count (if not in cosmos_predict2 mode) + if cosmos_predict2: + frame_count = 93 # Fixed frame count for cosmos_predict2 + else: + try: + cmd = ["ffprobe", "-v", "error", "-select_streams", "v:0", + "-show_entries", "stream=duration,r_frame_rate", "-of", "json", str(videos_dir / video_key / f"{video_file}.mp4")] + output = subprocess.check_output(cmd).decode() + probe_data = json.loads(output) + + if "streams" in probe_data and len(probe_data["streams"]) > 0: + stream = probe_data["streams"][0] + + # Check if duration is available + if "duration" in stream: + duration = float(stream["duration"]) + fps_str = stream.get("r_frame_rate", "30/1") + try: + num, den = map(int, fps_str.split("/")) + fps = num / den + frame_count = int(duration * fps) + except Exception: + # Default to 30fps if fps parsing fails + frame_count = int(duration * 30) + else: + # If no duration, use a default frame count + print(f"Warning: Could not determine duration for {video_file}, using default frame count") + frame_count = 93 + else: + # No streams found + print(f"Warning: No valid video streams found in {video_file}, using default frame count") + frame_count = 93 + + except Exception as e: + print(f"Error getting frame count for {video_file}: {e}") + # Default frame count if all methods fail + frame_count = 93 + + # Ensure frame count is at least 1 + frame_count = max(1, frame_count) + + results.append((video_id, annotation, frame_count)) + + return results + + +def copy_videos_parallel(video_copy_tasks, max_workers=16): + """Copy multiple videos in parallel using ThreadPoolExecutor.""" + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [] + for source, dest in video_copy_tasks: + futures.append(executor.submit(shutil.copy2, source, dest)) + + # Wait for all copy operations to complete + for future in concurrent.futures.as_completed(futures): + try: + future.result() + except Exception as e: + print(f"Error copying file: {e}") + +def process_folder(args): + """Process a single input folder.""" + folder_path, output_base_dir, annotation_source, fps, max_videos, num_workers_per_folder, cosmos_predict2, data_type, embodiment, video_key = args + + + + # Create folder-specific output directory + folder_name = folder_path.name + + output_dir = output_base_dir / f"{embodiment}.{folder_name}" + + # Process this folder + result = convert_raw_to_lerobot( + raw_dir=folder_path, + output_dir=output_dir, + annotation_source=annotation_source, + fps=fps, + max_videos=max_videos, + num_workers=num_workers_per_folder, + cosmos_predict2=cosmos_predict2, + data_type=data_type, + video_key=video_key, + ) + + return folder_name, result + +def convert_raw_to_lerobot( + raw_dir: Path, + output_dir: Path, + annotation_source: str = "human", + fps: int = 8, + max_videos: int | None = None, + num_workers: int = None, + cosmos_predict2: bool = False, + data_type: str = "lapa", + video_key: str = None, +): + """Convert raw dataset to LeRobot format.""" + + + + # Setup directories + videos_dir = raw_dir / "videos" + labels_dir = raw_dir / "labels" + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + # Create metadata directory + meta_dir = output_dir / "meta" + meta_dir.mkdir(parents=True, exist_ok=True) + + # Check for the new folder structure + video_subfolders = [f for f in videos_dir.iterdir() if f.is_dir()] + if video_subfolders: + # New folder structure detected + print(f"Detected new folder structure with subfolders: {[f.name for f in video_subfolders]}") + + # Get all unique video IDs across all subfolders + all_video_ids = set() + for folder in video_subfolders: + video_files = folder.glob("*.mp4") + all_video_ids.update(video_file.stem for video_file in video_files) + + # Convert to sorted list with numeric sorting + # Try to sort as integers if all IDs are numeric, otherwise use string sort + try: + video_ids = sorted(list(all_video_ids), key=lambda x: int(x) if x.isdigit() else float('inf')) + except ValueError: + # Fallback to string sort if conversion fails + video_ids = sorted(list(all_video_ids)) + + # Create dummy video_files list with just the IDs + video_files = [Path(video_id) for video_id in video_ids] + else: + raise ValueError("No video subfolders found in the input directory") + + if max_videos is not None: + video_files = video_files[:max_videos] + print(f"Processing only first {max_videos} videos for debugging") + + print(f"Processing {len(video_files)} videos in {videos_dir}") + + # Setup multiprocessing + if num_workers is None: + num_workers = max(1, multiprocessing.cpu_count() // 2) # Use half of available cores by default + + # Split videos into chunks for parallel processing + chunk_size = math.ceil(len(video_files) / num_workers) + video_chunks = [video_files[i:i + chunk_size] for i in range(0, len(video_files), chunk_size)] + + # Prepare arguments for parallel processing + args_list = [ + (chunk, labels_dir, output_dir, cosmos_predict2, data_type, videos_dir, video_key) + for chunk in video_chunks + ] + + # Process videos in parallel + total_frames = 0 + annotation_to_index = {} + episodes_info = [] + + print(f"Processing {len(video_files)} videos in {raw_dir} using {num_workers} workers") + print(f"cosmos_predict2 mode: {cosmos_predict2} (fixed FPS=16, frames=93)" if cosmos_predict2 else "") + with Pool(num_workers) as pool: + all_results = list(tqdm( + pool.imap(process_video_chunk, args_list), + total=len(args_list), + desc=f"Processing {raw_dir.name}" + )) + + # Process results and create videos + print(f"Processing {len(all_results)} chunks of videos...") + + # Flatten all results and sort by video_id to ensure episode_index matches video order + all_video_results = [] + for chunk_results in all_results: + all_video_results.extend(chunk_results) + + # Sort by video_id numerically to ensure episode 0 -> video 1, episode 1 -> video 2, etc. + try: + all_video_results.sort(key=lambda x: int(x[0]) if x[0].isdigit() else float('inf')) + except (ValueError, IndexError): + # Fallback to string sort if conversion fails + all_video_results.sort(key=lambda x: x[0]) + + # We'll collect all video copy tasks here + all_video_copy_tasks = [] + + + for video_id, annotation, frame_count in tqdm( + all_video_results, + desc="Processing videos in order" + ): + if annotation not in annotation_to_index: + annotation_to_index[annotation] = len(annotation_to_index) + + episode_index = len(episodes_info) + + # If cosmos_predict2 mode is enabled, use fixed frame count + actual_frame_count = 93 if cosmos_predict2 else frame_count + actual_fps = fps # Use the actual fps parameter + + # Determine state/action dimensions based on data_type + if data_type == "piper" or "piper" in str(output_dir).lower(): + state_action_dim = 8 # Piper: 6 joints + 2 grippers + elif data_type == "so101" or "so101" in str(output_dir).lower(): + state_action_dim = 6 # SO101: 5 joints + 1 gripper + else: + state_action_dim = 44 # Default + + # Create episode data + episode_data = { + "observation.state": [np.zeros(state_action_dim, dtype=np.float32)] * actual_frame_count, + "action": [np.zeros(state_action_dim, dtype=np.float32)] * actual_frame_count, + "timestamp": [i/actual_fps for i in range(actual_frame_count)], + "episode_index": [episode_index] * actual_frame_count, + "index": np.arange(total_frames, total_frames + actual_frame_count), + "task_index": [annotation_to_index[annotation]] * actual_frame_count, + f"annotation.{annotation_source}": [[annotation_to_index[annotation]]] * actual_frame_count + } + + # Save episode data + episode_chunk = episode_index // CHUNKS_SIZE + data_path = DATA_PATH.format(episode_chunk=episode_chunk, episode_index=episode_index) + save_path = output_dir / data_path + save_path.parent.mkdir(parents=True, exist_ok=True) + + df = pd.DataFrame(episode_data) + df.to_parquet(save_path) + + + view_list = [] + for folder in video_subfolders: + source_video_path = folder / f"{video_id}.mp4" + if source_video_path.exists(): + view_list.append((folder.name, source_video_path)) + + for video_key, source_video_path in view_list: + # Get the LeRobot format key + + # Destination path in LeRobot format + video_save_path = output_dir / VIDEO_PATH.format( + episode_chunk=episode_chunk, + video_key=video_key, + episode_index=episode_index + ) + video_save_path.parent.mkdir(parents=True, exist_ok=True) + + # Add to copy tasks instead of copying immediately + all_video_copy_tasks.append((source_video_path, video_save_path)) + + # Update episodes info + episodes_info.append({ + "episode_index": episode_index, + "tasks": [annotation], + "length": actual_frame_count, + "video_id": video_id + }) + + total_frames += actual_frame_count + + # Now copy all videos in parallel + print(f"Copying {len(all_video_copy_tasks)} videos in parallel...") + copy_videos_parallel(all_video_copy_tasks, max_workers=min(32, num_workers)) + + # Generate metadata files + # 1. tasks.jsonl + tasks_path = meta_dir / "tasks.jsonl" + tasks = [{"task_index": idx, "task": task} for task, idx in annotation_to_index.items()] + dump_jsonl(tasks, tasks_path) + + # 2. episodes.jsonl + episodes_path = meta_dir / "episodes.jsonl" + dump_jsonl(episodes_info, episodes_path) + + # 3. info.json + # Find a sample video for metadata from each view + + # Determine robot-specific configuration + if data_type == "piper" or "piper" in str(output_dir).lower(): + robot_type = "Piper-Abs" + state_action_shape = (8,) + state_action_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", + "finger_joint_left", "finger_joint_right"] + elif data_type == "so101" or "so101" in str(output_dir).lower(): + robot_type = "SO101" + state_action_shape = (6,) + state_action_names = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", + "wrist_roll", "gripper"] + else: + robot_type = data_type + state_action_shape = (44,) + state_action_names = [f"motor_{i}" for i in range(44)] + + info = { + "robot_type": robot_type, + "total_episodes": len(video_files), + "total_frames": total_frames, + "total_tasks": len(annotation_to_index), + "total_videos": len(video_subfolders), + "chunks_size": CHUNKS_SIZE, + "total_chunks": (len(video_files) + CHUNKS_SIZE - 1) // CHUNKS_SIZE, + "fps": fps, # Use the actual fps parameter + "data_path": DATA_PATH, + "video_path": VIDEO_PATH, + "features": { + "observation.state": { + "dtype": "float32", + "shape": state_action_shape, + "names": state_action_names + }, + "action": { + "dtype": "float32", + "shape": state_action_shape, + "names": state_action_names + }, + f"annotation.{annotation_source}": { + "dtype": "int64", + "shape": (1,) + } + } + } + + + for folder in video_subfolders: + # Find the first video in this folder + sample_videos = list(folder.glob("*.mp4")) + if sample_videos: + view_key = folder.name + sample_path = sample_videos[0] + info["features"][view_key] = get_video_metadata(sample_path) + + + # Remove None values from features + info["features"] = {k: v for k, v in info["features"].items() if v is not None} + + info_path = meta_dir / "info.json" + json_dump(info, info_path, indent=4) + + return output_dir + +def process_multiple_folders( + input_base_dir: Path, + output_base_dir: Path, + annotation_source: str = "human", + fps: int = 8, + max_videos: int | None = None, + num_workers: int = None, + cosmos_predict2: bool = False, + data_type: str = "lapa", + embodiment: str = None +): + """Process multiple input folders in parallel.""" + # Get all subdirectories that contain videos/ and labels/ folders + input_folders = [] + for folder in sorted(input_base_dir.iterdir()): + if folder.is_dir() and (folder / "videos").exists() and (folder / "labels").exists(): + input_folders.append(folder) + + print(f"Found {len(input_folders)} folders to process") + + # Determine number of workers per folder + total_workers = multiprocessing.cpu_count() if num_workers is None else num_workers + workers_per_folder = max(1, total_workers) + + # Prepare arguments for parallel folder processing + args_list = [ + (folder, output_base_dir, annotation_source, fps, max_videos, workers_per_folder, cosmos_predict2, data_type, embodiment) + for folder in input_folders + ] + + # Process folders in parallel using ThreadPoolExecutor + # (ProcessPoolExecutor would create too many processes) + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + futures = {executor.submit(process_folder, args): args[0] for args in args_list} + + # Process results as they complete + for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Processing folders"): + folder_path = futures[future] + try: + folder_name, result = future.result() + print(f"Completed processing folder: {folder_name}") + except Exception as e: + print(f"Error processing folder {folder_path.name}: {e}") + import traceback + traceback.print_exc() + +def main(): + parser = argparse.ArgumentParser(description="Convert raw dataset to LeRobot format") + parser.add_argument("--input_dir", type=str, required=True, help="Input directory containing multiple folders with videos/ and labels/") + parser.add_argument("--output_dir", type=str, required=True, help="Output directory for LeRobot dataset") + parser.add_argument("--fps", type=int, default=16, help="Video FPS") + parser.add_argument("--max_videos", type=int, default=None, help="Maximum number of videos to process per folder (for debugging)") + parser.add_argument("--num_workers", type=int, default=16, help="Total number of worker processes") + parser.add_argument("--cosmos_predict2", action="store_true", help="Process videos for cosmos_predict2 video models (fixed FPS=8, frames=81)") + parser.add_argument("--recursive", action="store_true", help="Process a single folder instead of multiple folders") + parser.add_argument("--data_type", type=str, default="dream", choices=["lapa", "dream"]) + parser.add_argument("--embodiment", type=str, default=None, help="Embodiment") + parser.add_argument("--video_key", type=str, default=None, help="Video key if cosmos_predict2 is false") + + args = parser.parse_args() + + if args.embodiment is None: + if 'robocasa' in args.output_dir: + args.embodiment = "robocasa_panda_omron" + elif 'gr1' in args.output_dir: + args.embodiment = "gr1_unified" + elif 'franka' in args.output_dir: + args.embodiment = "franka" + elif 'so101' in args.output_dir: + args.embodiment = "so101" + elif 'so100' in args.output_dir: + args.embodiment = "so100" + elif 'piper' in args.output_dir: + args.embodiment = "piper" + elif 'piper double' in args.output_dir: + args.embodiment = "piper double" + else: + raise ValueError(f"Unknown embodiment for {args.output_dir}")\ + + if args.embodiment == "robocasa_panda_omron": + args.annotation_source = "human.action.task_description" + elif args.embodiment == "gr1_unified": + args.annotation_source = "human.coarse_action" + elif args.embodiment == "franka": + args.annotation_source = "language.language_instruction" + elif args.embodiment == "so100": + args.annotation_source = "human.task_description" + elif args.embodiment == "so101": + args.annotation_source = "human.task_description" + elif args.embodiment == "piper": + args.annotation_source = "human.task_description" + elif args.embodiment == "piper double": + args.annotation_source = "human.task_description" + + if args.recursive: + # Process a single folder (original behavior) + process_multiple_folders( + input_base_dir=Path(args.input_dir), + output_base_dir=Path(args.output_dir), + annotation_source=args.annotation_source, + fps=args.fps, + max_videos=args.max_videos, + num_workers=args.num_workers, + cosmos_predict2=args.cosmos_predict2, + data_type=args.data_type, + embodiment=args.embodiment + ) + else: + convert_raw_to_lerobot( + raw_dir=Path(args.input_dir), + output_dir=Path(args.output_dir), + annotation_source=args.annotation_source, + fps=args.fps, + max_videos=args.max_videos, + num_workers=args.num_workers, + cosmos_predict2=args.cosmos_predict2, + data_type=args.data_type, + video_key=args.video_key + ) + + if args.embodiment == "gr1_unified": + source_dir = "IDM_dump/global_metadata/gr1" + elif args.embodiment == "robocasa_panda_omron": + source_dir = "IDM_dump/global_metadata/robocasa" + elif args.embodiment == "franka": + source_dir = "IDM_dump/global_metadata/franka" + elif args.embodiment == "so100": + source_dir = "IDM_dump/global_metadata/so100" + elif args.embodiment == "so101": + source_dir = "IDM_dump/global_metadata/so101" + elif args.embodiment == "piper": + source_dir = "IDM_dump/global_metadata/piper" + elif args.embodiment == "piper double": + source_dir = "IDM_dump/global_metadata/piper double" + + # copy modality.json + shutil.copy(source_dir + "/modality.json", args.output_dir + "/meta/modality.json") + + # # copy stats.json + shutil.copy(source_dir + "/stats.json", args.output_dir + "/meta/stats.json") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/split_video_instruction.py b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/split_video_instruction.py new file mode 100644 index 00000000..3620670e --- /dev/null +++ b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/split_video_instruction.py @@ -0,0 +1,101 @@ +import os +import shutil +import re +import argparse + +def process_mp4_files(source_dir, output_dir, recursive=False, task_name=None): + """ + Process MP4 files from source_dir and its subdirectories, extract instructions from filenames, + and save them to output_dir with appropriate structure. + + Args: + source_dir: Directory containing MP4 files or subdirectories with MP4 files + output_dir: Directory to save processed data + recursive: If True, maintain the directory structure from source_dir in output_dir + task_name: Optional task name to use as instruction when filename is just a number + """ + if not recursive: + # Original behavior: process all MP4 files into a single output directory + labels_dir = os.path.join(output_dir, "labels") + videos_dir = os.path.join(output_dir, "videos") + + os.makedirs(labels_dir, exist_ok=True) + os.makedirs(videos_dir, exist_ok=True) + + # Find all MP4 files in source_dir + mp4_files = [] + for file in os.listdir(source_dir): + if file.endswith('.mp4'): + mp4_files.append(os.path.join(source_dir, file)) + + # Sort files numerically to ensure consistent ordering (1.mp4, 2.mp4, ..., 10.mp4, not 1.mp4, 10.mp4, 2.mp4) + def sort_key(filepath): + filename = os.path.basename(filepath) + stem = os.path.splitext(filename)[0] + try: + return int(stem) + except ValueError: + return float('inf') + mp4_files.sort(key=sort_key) + + # Process each MP4 file + for idx, mp4_path in enumerate(mp4_files, 1): + mp4_file = os.path.basename(mp4_path) + + # Extract instruction from filename (remove number prefix and .mp4 extension) + instruction = re.sub(r'^\d+_', '', mp4_file).replace('.mp4', '').replace('_', ' ') + + # If instruction is just a number and task_name is provided, use task_name instead + if task_name and instruction.isdigit(): + # Convert task_name from directory format (e.g., "Lift_the_red_cube_up") to readable format + instruction = task_name.replace('_', ' ') + + label_file = os.path.join(labels_dir, f"{idx}.txt") + with open(label_file, 'w') as f: + f.write(instruction) + + # Copy video with new name + target_video = os.path.join(videos_dir, f"{idx}.mp4") + shutil.copy2(mp4_path, target_video) + + print(f"Processed {mp4_path} -> {idx}.mp4, instruction: {instruction}") + + print(f"Processed {len(mp4_files)} files. Results saved to {output_dir}") + else: + # Recursive behavior: maintain directory structure + print(f"Processing in recursive mode, maintaining directory structure...") + + # Get all subdirectories in the source directory + subdirs = [] + for item in os.listdir(source_dir): + item_path = os.path.join(source_dir, item) + if os.path.isdir(item_path): + subdirs.append(item) + + if not subdirs: + # If no subdirectories, process the source directory directly + print(f"No subdirectories found in {source_dir}, processing directly...") + process_mp4_files(source_dir, output_dir, recursive=False) + return + + # Process each subdirectory + for subdir in subdirs: + source_subdir = os.path.join(source_dir, subdir) + output_subdir = os.path.join(output_dir, subdir) + + # Create output subdirectory + os.makedirs(output_subdir, exist_ok=True) + + # Process files in this subdirectory, passing the directory name as task_name + process_mp4_files(source_subdir, output_subdir, recursive=False, task_name=subdir) + + print(f"Recursive processing complete. Results saved to {output_dir}") + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Process MP4 files from source directory and save to output directory") + parser.add_argument("--source_dir", type=str, required=True, help="Path to source directory containing MP4 files") + parser.add_argument("--output_dir", type=str, required=True, help="Path to output directory") + parser.add_argument("--recursive", action="store_true", help="Process subdirectories recursively, maintaining directory structure") + args = parser.parse_args() + + process_mp4_files(args.source_dir, args.output_dir, args.recursive) \ No newline at end of file diff --git a/scripts/tutorials/cosmos_tutorials/replay.py b/scripts/tutorials/cosmos_tutorials/replay.py new file mode 100644 index 00000000..de61e13e --- /dev/null +++ b/scripts/tutorials/cosmos_tutorials/replay.py @@ -0,0 +1,495 @@ +""" +Hybrid Replay: Parquet actions/states + HDF5 initial states + +This script replays recorded robot episodes by: +- Loading initial states from HDF5 datasets +- Loading actions or states from Parquet files +- Running replay inside IsaacLab environments +- Supporting parallel environments and success statistics +""" + + +import multiprocessing +if multiprocessing.get_start_method() != "spawn": + multiprocessing.set_start_method("spawn", force=True) + +# --------------------------------------------------------------------- +# CLI arguments and IsaacLab app launcher +# --------------------------------------------------------------------- +import argparse +from isaaclab.app import AppLauncher + +parser = argparse.ArgumentParser(description="Hybrid Replay for Leisaac") +parser.add_argument("--task", type=str, required=True, help="LeIsaac Task name,e.g. LeIsaac-SO101-LiftCube-v0") +parser.add_argument("--num_envs", type=int, default=1, help="Number of parallel environments") +parser.add_argument("--step_hz", type=int, default=60, help="Simulation frequency") +parser.add_argument("--parquet_dir", type=str, default="", + help="Directory containing Parquet episode files (expects episode_*.parquet)." +parser.add_argument("--hdf5_file", type=str, default="", + help="HDF5 dataset file containing per-episode initial states and metadata.") +parser.add_argument("--select_episodes", type=int, nargs="+", default=[], + help="Episode indices to replay. If empty, replays all available parquet episodes.") +parser.add_argument("--task_type",type=str,default=None,help="Optional task type override. If not set, inferred from --task.") +parser.add_argument("--replay_mode",type=str,default="action",choices=["action", "state"], + help="Replay mode: 'action' replays action vectors; 'state' replays joint states.") + +# Add IsaacLab standard launcher args (e.g., device/headless/experience settings). +AppLauncher.add_app_launcher_args(parser) +args_cli = parser.parse_args() + +# Launch IsaacLab app (must happen before importing certain sim-related modules). +app_launcher = AppLauncher(vars(args_cli)) +simulation_app = app_launcher.app + +# --------------------------------------------------------------------- +# Imports (after app launch) +# --------------------------------------------------------------------- +import contextlib +import os +import time +import re +import gymnasium as gym +import torch +import h5py +import numpy as np +import pandas as pd +from pathlib import Path +from typing import List, Optional, Dict + +from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.utils.datasets import HDF5DatasetFileHandler +from isaaclab_tasks.utils import parse_env_cfg +from leisaac.utils.env_utils import dynamic_reset_gripper_effort_limit_sim, get_task_type + + +# --------------------------------------------------------------------- +# Episode container: parquet actions/states + HDF5 initial states +# --------------------------------------------------------------------- +class HybridEpisodeData: + """ + Holds a single episode's replay payload: + - initial_state: from HDF5 + - actions: from parquet column "action" (optional) + - states: from parquet column "observation.state" (optional) + + Data is denormalized from LeRobot format into IsaacLab radians. + """ + + def __init__(self, actions_df: pd.DataFrame, initial_state: dict, device: str): + self.device = device + self.step_idx = 0 + self.initial_state = initial_state + self.actions = None + self.states = None + + # Parquet data uses normalized angles in [-100, 100]. + # We denormalize into IsaacLab joint limits and convert degrees -> radians. + if "action" in actions_df.columns: + actions_np = np.stack(actions_df["action"].to_numpy()) + actions_np = self._denormalize_joint_pos(actions_np) + self.actions = torch.tensor(actions_np, device=device, dtype=torch.float32) + + if "observation.state" in actions_df.columns: + states_np = np.stack(actions_df["observation.state"].to_numpy()) + states_np = self._denormalize_joint_pos(states_np) + self.states = torch.tensor(states_np, device=device, dtype=torch.float32) + + @staticmethod + def _denormalize_joint_pos(joint_pos: np.ndarray) -> np.ndarray: + """ + Denormalize: LeRobot normalized joint positions [-100, 100] -> IsaacLab joint limits (degrees), + then convert degrees to radians. + + NOTE: The limits here must match the embodiment/task joint conventions. + """ + ISAACLAB_LIMITS = [ + (-110.0, 110.0), + (-100.0, 100.0), + (-100.0, 90.0), + (-95.0, 95.0), + (-160.0, 160.0), + (-10, 100.0), + ] + LEROBOT_LIMITS = [ + (-100, 100), + (-100, 100), + (-100, 100), + (-100, 100), + (-100, 100), + (0, 100), + ] + + result = joint_pos.copy() + for i in range(6): + isaaclab_min, isaaclab_max = ISAACLAB_LIMITS[i] + lerobot_min, lerobot_max = LEROBOT_LIMITS[i] + isaac_range = isaaclab_max - isaaclab_min + lerobot_range = lerobot_max - lerobot_min + result[:, i] = (result[:, i] - lerobot_min) / lerobot_range * isaac_range + isaaclab_min + + # Convert degrees to radians + result = result * np.pi / 180.0 + return result + + def get_next_action(self): + """Return the next action for replay_mode='action', or None if episode finished/unavailable.""" + if self.actions is not None and self.step_idx < len(self.actions): + action = self.actions[self.step_idx] + self.step_idx += 1 + return action + return None + + def get_next_state(self): + """Return the next state for replay_mode='state', or None if episode finished/unavailable.""" + if self.states is not None and self.step_idx < len(self.states): + state = self.states[self.step_idx] + self.step_idx += 1 + return state + return None + + def get_initial_state(self): + """Return the initial state dict loaded from HDF5.""" + return self.initial_state + + +# --------------------------------------------------------------------- +# Dataset loader: align parquet episodes with HDF5 episodes and filter invalid ones +# --------------------------------------------------------------------- +class HybridDatasetLoader: + """ + Loads: + - Parquet episode files: actions/states + - HDF5 episode initial states + + Performs a pre-scan over HDF5 to filter out failed episodes (based on attrs and required keys). + """ + + def __init__(self, hdf5_path: str, parquet_dir: str): + self.hdf5_path = hdf5_path + self.parquet_dir = Path(parquet_dir) + + self.hdf5_handler = HDF5DatasetFileHandler() + self.hdf5_handler.open(self.hdf5_path) + + # Collect parquet episodes and sort by the integer index in their filename (episode_XXXX.parquet). + self.all_parquet_files = sorted( + self.parquet_dir.glob("**/episode_*.parquet"), + key=lambda p: int(re.search(r"episode_(\d+)", p.name).group(1)) + if re.search(r"episode_(\d+)", p.name) + else 0, + ) + + # HDF5 episode names as stored in the dataset (e.g. "demo_0", "demo_1", ...). + self.all_hdf5_names = list(self.hdf5_handler.get_episode_names()) + self.valid_indices = [] + + # Pre-scan HDF5 episodes to filter out failures/missing keys. + print("Pre-scanning HDF5 episode attributes to filter out failed data...") + with h5py.File(self.hdf5_path, "r") as f: + # Use the max range to iterate; assumes parquet and HDF5 correspond broadly. + max_range = max(len(self.all_hdf5_names), len(self.all_parquet_files)) + + for i in range(max_range): + demo_name = self.all_hdf5_names[i] + demo_path = f"data/{demo_name}" + + if demo_path in f: + demo_group = f[demo_path] + # If 'success' attr does not exist, default to True. + is_success = demo_group.attrs.get("success", True) + + if is_success: + try: + # Validate required datasets exist. + _ = demo_group["actions"] + _ = demo_group["obs/joint_pos"] + _ = demo_group["obs/front"] + self.valid_indices.append(i) + except KeyError: + print(f"Pre-scan: Filtering out failed episode {demo_name}") + else: + print(f"Pre-scan: Filtering out failed episode {demo_name}") + + self.num_episodes = len(self.valid_indices) + print(f"Pre-scan complete: {self.num_episodes} valid episodes (original total: {max_range}).") + + def load_episode(self, logical_index: int, device: str): + """ + Load an episode: + - Parquet is addressed directly by logical_index (episode ordering from parquet list) + - HDF5 episode is chosen via valid_indices mapping + + Returns: + HybridEpisodeData or None if out-of-range. + """ + if logical_index >= len(self.all_parquet_files): + return None + + parquet_path = self.all_parquet_files[logical_index] + + # Map logical index (parquet ordering) into an actual HDF5 index among valid ones. + if logical_index < len(self.valid_indices): + actual_hdf5_idx = self.valid_indices[logical_index] + else: + # If parquet list is longer than valid HDF5 list, fall back to last valid index. + actual_hdf5_idx = self.valid_indices[-1] if self.valid_indices else 0 + + demo_name = self.all_hdf5_names[actual_hdf5_idx] + hdf5_ep_data = self.hdf5_handler.load_episode(demo_name, device) + initial_state = hdf5_ep_data.get_initial_state() + + df = pd.read_parquet(parquet_path) + return HybridEpisodeData(df, initial_state, device) + + def close(self): + """Close underlying HDF5 handler.""" + self.hdf5_handler.close() + + +# --------------------------------------------------------------------- +# Rate limiter: attempt to replay at a fixed step rate while rendering +# --------------------------------------------------------------------- +class RateLimiter: + """ + Simple real-time rate limiter: + - sleeps to maintain a target hz + - calls env.sim.render() during waiting to keep UI responsive + """ + + def __init__(self, hz): + self.hz = hz + self.last_time = time.time() + self.sleep_duration = 1.0 / hz + self.render_period = min(0.0166, self.sleep_duration) + + def sleep(self, env): + next_wakeup_time = self.last_time + self.sleep_duration + while time.time() < next_wakeup_time: + time.sleep(self.render_period) + env.sim.render() + + # Advance the timeline by exactly one tick. + self.last_time = self.last_time + self.sleep_duration + + # Catch up if we are behind. + if self.last_time < time.time(): + while self.last_time < time.time(): + self.last_time += self.sleep_duration + + +# --------------------------------------------------------------------- +# Replay helpers +# --------------------------------------------------------------------- +def get_next_action_or_state(episode_data, replay_mode: str): + """ + Depending on replay_mode: + - "state": replay joint states (episode_data.states) + - "action": replay actions (episode_data.actions) + """ + if replay_mode == "state": + next_state = episode_data.get_next_state() + if next_state is None: + return None + return next_state + else: + return episode_data.get_next_action() + + +# --------------------------------------------------------------------- +# Replay manager: multi-env episode scheduling, stepping, success tracking +# --------------------------------------------------------------------- +class ReplayManager: + """ + Runs replay in N parallel envs: + - schedules episodes to env slots + - steps the env with action/state vectors + - tracks success based on termination signal + - prints per-episode success/failure summary + """ + + def __init__(self, env, dataset_loader, config): + self.env = env + self.loader = dataset_loader + self.cfg = config + self.num_envs = env.num_envs + self.device = env.device + + # If env provides a configured idle_action, use it; otherwise default to zeros. + if hasattr(env.cfg, "idle_action"): + self.idle_action = env.cfg.idle_action.repeat(self.num_envs, 1) + else: + self.idle_action = torch.zeros(env.action_space.shape, device=self.device) + + self.rate_limiter = RateLimiter(config.step_hz) + self.task_type = get_task_type(config.task, config.task_type) + self.success_count = 0 + self.total_episodes = 0 + + def run(self, episode_indices: List[int]): + """ + Replay the given list of episode indices. + + Episode lifecycle per env: + - if env has no active episode: load next from queue, reset_to initial state + - each step: pull next action/state and env.step() + - if episode ends: record success based on termination flag observed during replay + """ + episode_queue = list(episode_indices) + + # Per-env bookkeeping: + # env_data_map: active HybridEpisodeData for each env_id + # env_idx_map: which episode index is running in each env_id + # env_success_map: whether that env has triggered termination (treated as success) + env_data_map: Dict[int, Optional[HybridEpisodeData]] = {i: None for i in range(self.num_envs)} + env_idx_map: Dict[int, Optional[int]] = {i: None for i in range(self.num_envs)} + env_success_map: Dict[int, bool] = {i: False for i in range(self.num_envs)} + recorded = set() + + print(f"Replaying {len(episode_indices)} episodes...") + + # Use inference_mode to avoid autograd overhead; suppress Ctrl+C traceback noise. + with contextlib.suppress(KeyboardInterrupt), torch.inference_mode(): + while simulation_app.is_running() and not simulation_app.is_exiting(): + actions = self.idle_action.clone() + has_next = False + + for env_id in range(self.num_envs): + curr_data = env_data_map[env_id] + next_act = get_next_action_or_state(curr_data, self.cfg.replay_mode) if curr_data else None + + if next_act is None: + # Episode finished: record success/failure once per episode index. + curr_idx = env_idx_map[env_id] + if curr_idx is not None and curr_idx not in recorded: + recorded.add(curr_idx) + self.total_episodes += 1 + if env_success_map[env_id]: + self.success_count += 1 + print(f"[Episode #{curr_idx}] ✓ SUCCESS | {self.success_count}/{self.total_episodes}") + else: + print(f"[Episode #{curr_idx}] ✗ FAILED | {self.success_count}/{self.total_episodes}") + + # Load next episode from queue for this env slot. + next_idx = None + while episode_queue: + next_idx = episode_queue.pop(0) + if next_idx < len(self.loader.all_parquet_files): + break + next_idx = None + + if next_idx is not None: + new_data = self.loader.load_episode(next_idx, self.device) + if new_data: + env_data_map[env_id] = new_data + env_idx_map[env_id] = next_idx + env_success_map[env_id] = False + + # Reset only the selected env_id using provided initial state. + self.env.reset_to( + new_data.get_initial_state(), + torch.tensor([env_id], device=self.device), + is_relative=True, + ) + + # After reset, immediately fetch the first action/state for this episode. + next_act = get_next_action_or_state(new_data, self.cfg.replay_mode) + has_next = True + else: + # No more episodes available. + env_data_map[env_id] = None + env_idx_map[env_id] = None + continue + else: + has_next = True + + # Apply action/state for this env slot. + if next_act is not None: + actions[env_id] = next_act + + # If no env has any next step, exit. + if not has_next: + break + + # Optionally update gripper effort limits dynamically between steps. + if self.cfg.replay_mode == "action" and self.env.cfg.dynamic_reset_gripper_effort_limit: + dynamic_reset_gripper_effort_limit_sim(self.env, self.task_type) + + # Step simulation with batched actions. + _, _, reset_terminated, _, _ = self.env.step(actions) + + # Mark success if termination triggered for that env slot. + for env_id in range(self.num_envs): + if reset_terminated[env_id] and env_idx_map[env_id] is not None: + env_success_map[env_id] = True + + # Maintain wall-clock rate and render while waiting. + self.rate_limiter.sleep(self.env) + + # Print final success rate summary. + success_rate = (self.success_count / max(1, self.total_episodes)) * 100 + print(f"\n{'='*50}") + print(f"Success Rate: {self.success_count}/{self.total_episodes} ({success_rate:.1f}%)") + print(f"{'='*50}") + + +# --------------------------------------------------------------------- +# Environment creation helper +# --------------------------------------------------------------------- +def create_env(args): + """ + Create and initialize IsaacLab environment with appropriate task config. + Also configures teleop device and disables recorders/timeouts where appropriate. + """ + env_cfg = parse_env_cfg(args.task, device=args.device, num_envs=args.num_envs) + task_type = get_task_type(args.task, args.task_type) + env_cfg.use_teleop_device(task_type) + + # Some tasks (Direct) use manual termination, so disable timeouts/recorders. + if "Direct" in args.task: + env_cfg.never_time_out = True + env_cfg.manual_terminate = True + env_cfg.recorders = {} + else: + env_cfg.recorders = {} + # Disable time_out termination if present. + if hasattr(env_cfg, "terminations") and hasattr(env_cfg.terminations, "time_out"): + env_cfg.terminations.time_out = None + + env = gym.make(args.task, cfg=env_cfg).unwrapped + if hasattr(env, "initialize"): + env.initialize() + env.reset() + return env + + +# --------------------------------------------------------------------- +# Main entry +# --------------------------------------------------------------------- +def main(): + # Load datasets (parquet + hdf5) and pre-filter invalid episodes. + dataset_loader = HybridDatasetLoader(args_cli.hdf5_file, args_cli.parquet_dir) + print(f"Datasets loaded. Available episodes: {dataset_loader.num_episodes}") + + # If no episodes are explicitly selected, replay all parquet episodes found. + episode_indices = args_cli.select_episodes if args_cli.select_episodes else list(range(len(dataset_loader.all_parquet_files))) + episode_indices = [idx for idx in episode_indices if idx < len(dataset_loader.all_parquet_files)] + + if not episode_indices: + print("No valid episodes to replay.") + return + + print(f"Selected {len(episode_indices)} episodes for replay.") + + # Create env and run replay. + env = create_env(args_cli) + manager = ReplayManager(env, dataset_loader, args_cli) + manager.run(episode_indices) + + # Cleanup resources. + env.close() + dataset_loader.close() + simulation_app.close() + + +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/cosmos_tutorials/so101.sh b/scripts/tutorials/cosmos_tutorials/so101.sh new file mode 100644 index 00000000..e72003d4 --- /dev/null +++ b/scripts/tutorials/cosmos_tutorials/so101.sh @@ -0,0 +1,149 @@ +#!/bin/bash +set -e + +# ============================================================================= +# Configuration +# ============================================================================= + +# Directory containing Cosmos Predict 2.5 inference outputs. +# This directory should include generated videos produced by the Cosmos model. +LEROBOT_INPUT="" + +# Working directory used to store all intermediate and final IDM outputs. +# It is recommended to place this on a fast local disk. +WORK_DIR="" + +# Robot embodiment type (used by IDM and LeRobot). +ROBOT_TYPE="so101" + +# Key name used to store videos in the LeRobot dataset structure. +# This should match the observation key expected by downstream IDM scripts. +VIDEO_KEY="observation.images.front" + +# ============================================================================= +# Intermediate directories (auto-generated) +# ============================================================================= + +# Step 1 output: task-named directories converted from Cosmos outputs +STEP1_DIR="${WORK_DIR}/step1" + +# Step 2 output: split videos and text instructions +STEP2_DIR="${WORK_DIR}/step2" + +# Step 3 output: preprocessed videos (e.g., resized, normalized) +STEP3_DIR="${WORK_DIR}/step3" + +# Step 4 output: final LeRobot-format dataset +STEP4_DIR="${WORK_DIR}/${ROBOT_TYPE}.data" + +# ============================================================================= +# Step 1: Convert Cosmos outputs to task-based directory structure +# ============================================================================= +# - Reads Cosmos Predict outputs +# - Groups videos by task name +# - Prepares data for downstream preprocessing +python3 IDM_dump/scripts/preprocess_leisaac/cosmos2.5_to_step2_format.py \ + --cosmos_dir "${LEROBOT_INPUT}" \ + --output_dir "${STEP1_DIR}" + +# ============================================================================= +# Step 2: Split videos and instructions +# ============================================================================= +# - Separates raw videos and text instructions into: +# - videos/ +# - labels/ +# - The --recursive flag allows processing nested task directories +python3 IDM_dump/scripts/preprocess_leisaac/split_video_instruction.py \ + --source_dir "${STEP1_DIR}" \ + --output_dir "${STEP2_DIR}" \ + --recursive + +# ============================================================================= +# Step 3: Preprocess videos +# ============================================================================= +# - Resizes videos to the resolution expected by IDM +# - Converts video format if necessary +# - Preserves directory structure across tasks +python3 IDM_dump/scripts/preprocess_leisaac/preprocess_video.py \ + --src_dir "${STEP2_DIR}" \ + --dst_dir "${STEP3_DIR}" \ + --dataset "${ROBOT_TYPE}" \ + --original_width 640 \ + --original_height 480 \ + --recursive + +# ============================================================================= +# IMPORTANT USAGE NOTE +# ============================================================================= +# It is STRONGLY RECOMMENDED to: +# +# 1. Run Step 1–3 first +# 2. Inspect the contents of ${STEP3_DIR} +# 3. Identify the generated task directory name(s) +# 4. Manually copy the desired task directory name into Step 4 +# +# This avoids hard-coding task names before they are known and +# allows flexible reuse of this script across different tasks. +# +# Example: +# ls ${STEP3_DIR} +# → Lift_the_red_cube_up +# +# Then use: +# --input_dir "${STEP3_DIR}/Lift_the_red_cube_up" +# +# ============================================================================= + +# ============================================================================= +# Step 4: Convert preprocessed data to LeRobot dataset reminder +# ============================================================================= +# --input_dir: +# Path to a task-specific directory under STEP3_DIR. +# The directory name MUST match the task name generated in Step 3. +# +# --output_dir: +# Target directory for the LeRobot-format dataset. +# +# --fps: +# Target frames per second for the output dataset. +# +# --embodiment: +# Robot embodiment identifier used by LeRobot/IDM. +# +# --video_key: +# Observation key used to store video data. +python3 IDM_dump/scripts/preprocess_leisaac/raw_to_lerobot.py \ + --input_dir "${STEP3_DIR}/Lift_the_red_cube_up" \ + --output_dir "${STEP4_DIR}" \ + --fps 16 \ + --embodiment "${ROBOT_TYPE}" \ + --video_key "${VIDEO_KEY}" \ + --embodiment "so101" + +# ============================================================================= +# Step 5: Dump IDM actions from the LeRobot dataset +# ============================================================================= +# - Loads a pretrained IDM checkpoint +# - Runs IDM inference on the LeRobot dataset +# - Exports predicted action trajectories +# +# --checkpoint: +# Path to a trained IDM checkpoint. +# +# --dataset: +# Path to the LeRobot dataset generated in Step 4. +# +# --output_dir: +# Output directory where IDM predictions will be stored. +# +# --num_gpus: +# Number of GPUs used for IDM inference. +# +# --video_indices: +# Indices of videos to process (e.g., "0 16" processes videos 0–16). +python3 IDM_dump/scripts/preprocess_leisaac/dump_idm_actions.py \ + --checkpoint "" \ + --dataset "${STEP4_DIR}" \ + --output_dir "${STEP4_DIR}_idm_cosmos" \ + --num_gpus 1 \ + --video_indices "0 16" From c6c0492d905d5f30d573f03c6aeb80580b6c7cd4 Mon Sep 17 00:00:00 2001 From: EverNorif <1320346985@qq.com> Date: Wed, 7 Jan 2026 16:19:42 +0800 Subject: [PATCH 2/9] add tutorial doc. --- docs/docs/docs/tutorials/cosmos_tutorial.md | 517 ++++++++++++++++++++ docs/docusaurus.config.js | 2 +- docs/sidebars.js | 6 + 3 files changed, 524 insertions(+), 1 deletion(-) create mode 100644 docs/docs/docs/tutorials/cosmos_tutorial.md diff --git a/docs/docs/docs/tutorials/cosmos_tutorial.md b/docs/docs/docs/tutorials/cosmos_tutorial.md new file mode 100644 index 00000000..7ea00885 --- /dev/null +++ b/docs/docs/docs/tutorials/cosmos_tutorial.md @@ -0,0 +1,517 @@ +# LeIsaac × Cosmos: Video-to-Action Data Generation Pipeline + +This tutorial extends **LeIsaac** by integrating **Cosmos-Predict2.5** and **GR00T-Dreams IDM** into a LeIsaac-native data generation loop. **LeIsaac** is used to collect teleoperated demonstrations (HDF5) and convert them into LeRobot datasets. **Cosmos-Predict 2.5** is post-trained on these videos to synthesize additional rollout videos at scale, and **IDM** is fine-tuned on the same dataset to infer robot actions from the generated videos. Together, this produces **a scalable pipeline for constructing synthetic, complete LeRobot datasets** , which can be replayed and evaluated directly in **LeIsaac**. + +## 🎥 From a Single Demonstration to Large-Scale Synthetic Rollouts via Cosmos +
+ +
+ +## Overview + +1. Use **LeIsaac** to collect an **HDF5 dataset** and convert it into a **LeRobot dataset** +2. Post-train **Cosmos-Predict2.5**, then run inference to **generate synthetic videos** +3. Fine-tune **IDM**, then run inference to **generate synthetic LeRobot trajectories(parquet)** +4. Use **LeIsaac** to **replay and evaluate** using the HDF5 dataset + inferred action trajectories + +## Step 1: Data Collection with LeIsaac + +::::tip[What you get from this step] +- A real demonstration dataset in **HDF5** format (recorded in Isaac Sim through LeIsaac teleoperation) +- A converted dataset in **LeRobot** format (videos + metadata + parquet) +:::: + +### 1.1 Collect the HDF5 Dataset via LeIsaac Teleoperation + +First, collect demonstrations using **LeIsaac** teleoperation. + +- Reference: **[Teleoperation | LeIsaac Document](https://lightwheelai.github.io/leisaac/docs/getting_started/teleoperation)** +- Output: an HDF5 file (for example, `dataset.hdf5`) + +The HDF5 dataset serves as the primary source of truth for replay and evaluation. + +It is also the source used to construct the initial **LeRobot dataset** for both Cosmos and IDM. + +### 1.2 Convert HDF5 to a LeRobot Dataset + +Convert the recorded HDF5 dataset into **LeRobot format**. + +- Reference: **[Data Convention | LeIsaac Document](https://lightwheelai.github.io/leisaac/docs/getting_started/policy_support#1-data-convention)** +- Output: a LeRobot dataset used for post-training **Cosmos-Predict2.5** and fine-tuning **IDM** + +::::info +**💡Important (Video Encoding)** + +Ensure that all output videos are encoded using **H.264 (h264)**. + +Avoid **AV1** encoding, as Cosmos and IDM may fail during decoding or processing. + +You can directly modify the `leisaac/scripts/convert/isaaclab2lerobot.py` to ensure correct video encoding: + +```diff +- "video.codec": "av1" ++ "video.codec": "h264" +``` + +:::: + +## Step 2: Video Generation with Cosmos-predict2.5 + +::::tip[What you get from this step] +- A **Cosmos-Predict2.5 checkpoint** post-trained on your task-specific LeRobot videos +- A set of **synthetic rollout videos** generated by Cosmos + +In this pipeline, **Cosmos-Predict 2.5** is used purely as a **video generator**. + +It learns the visual rollout distribution from LeRobot videos collected via **LeIsaac**, and generates new rollout videos conditioned on: +- a text prompt (task description), +- and the first few frames of an example videos. + +These generated videos will later be converted into executable robot actions using **IDM**. +:::: + +### 2.1 Install Cosmos-Predict2.5 + +Set up the Cosmos-Predict2.5 environment by following the official installation guide:**[Set up the Cosmos-Predict2.5](https://github.com/nvidia-cosmos/cosmos-predict2.5/blob/main/docs/setup.md)** + + +### 2.2 Prepare the Cosmos-Predict2.5 Post-Training Dataset (Video + Prompt) + +Dataset folder format should be: +```text +cosmos-predict2.5/datasets/benchmark_train// +├── metas/ +│ ├── *.txt +├── videos/ +│ ├── *.mp4 +``` +Construct it from your LeRobot dataset: + +1. Copy MP4 videos from your LeRobot dataset from `/videos/` to `cosmos-predict2.5/datasets/benchmark_train//videos/` + +2. Rename the copied videos as a clean numeric sequence:`1.mp4`, `2.mp4`, `3.mp4`, ... + +3. Create the same number of prompt files under `cosmos-predict2.5/datasets/benchmark_train//metas/` + +4. Fill each prompt file using the task text in `/meta/tasks.jsonl`.For example, if one line in tasks.jsonl is `{"task_index": 0, "task": "Lift the red cube up."}`,then 1.txt should contain `Lift_the_red_cube_up`.In many single-task cases, you will write the same prompt into all *.txt files, but the format supports per-video prompts if needed. + +### 2.3 Post-train Cosmos-Predict2.5 + +Post-train Cosmos-Predict2.5 on the prepared dataset by following **[the official post-training instructions](https://github.com/nvidia-cosmos/cosmos-predict2.5/blob/main/docs/post-training_video2world_gr00t.md#21-post-training-cosmos-predict25-2b-model)**. + +This step produces a Cosmos-Predict2.5 checkpoint specialized for your robot embodiment and task distribution. + +### 2.4 Run Inference to Generate Videos + +After post-training, the Cosmos-Predict2.5 model checkpoints are typically saved in **Distributed Checkpoint (DCP)** format.Before running inference, these checkpoints need to be converted into a **consolidated PyTorch format** that can be loaded by the inference scripts. + +**Follow the [converting-dcp-checkpoint-to-consolidated-pytorch-format](https://github.com/nvidia-cosmos/cosmos-predict2.5/blob/main/docs/post-training_video2world_gr00t.md#31-converting-dcp-checkpoint-to-consolidated-pytorch-format) to convert DCP checkpoints**. + +Once the checkpoint has been converted, run video generation inference following the [official running-inference instructions ](https://github.com/nvidia-cosmos/cosmos-predict2.5/blob/main/docs/post-training_video2world_gr00t.md#32-running-inference). + +#### 🔬**Batch Inference (Used for fine-tuning IDM)** + +For large-scale video generation, batch inference is supported. +Copy the batch inference helper script `generate_batch_config.py` into the folder:`cosmos-predict2.5/scripts/`. + +You can generate a batch inference configuration `jsonl` file using: +```bash +python scripts/generate_batch_config.py --use-prompt +``` + +## Step 3: Action Inference with IDM + +::::tip[What you get from this step] +- An **IDM checkpoint** fine-tuned on your LeRobot dataset and robot embodiment +- A set of **LeRobot-format trajectories** (e.g. parquet files) inferred from Cosmos-generated videos + +In this pipeline, **IDM (Inverse Dynamics Model)** from **GR00T-Dreams** is used to convert **Cosmos-generated rollout videos** into executable robot actions. + +IDM is first fine-tuned on the original LeRobot dataset collected via **LeIsaac**, and then applied to infer actions from synthetic videos, producing a new, fully compatible **LeRobot dataset**. +:::: + +### 3.1 Install IDM Environment + +IDM requires the **Cosmos-Predict2** environment (**not 2.5**). + +Follow the official prerequisites guide:**[Cosmos-Predict2 – Prerequisites](https://github.com/nvidia-cosmos/cosmos-predict2/blob/main/documentations/post-training_video2world_gr00t.md#prerequisites)** + +💡 **Notes on Dependency Installation** +- When installing dependencies such as `openai` and `tyro`, use: + ```bash + uv pip install openai tyro numpydantic albumentations tianshou + ``` +- For `pytorch3d`, install with `no build isolation`: + ```bash + uv pip install --no-build-isolation git+https://github.com/facebookresearch/pytorch3d.git + ``` +- If you encounter issues related to `APEX`, it can be safely removed from the uv environment.APEX is not required for IDM training or inference.: + +### 3.2 Fine-tune IDM +For IDM fine-tuning, please refer to **[training-custom-idm-model](https://github.com/NVIDIA/GR00T-Dreams?tab=readme-ov-file#optional-33-training-custom-idm-model)** + +#### 3.2.1 Preparation: Modality Metadata and DataConfig + +##### Step 1: Add modality.json + +Create `modality.json` under `GR00T-Dreams/IDM_dump/global_metadata/{embodiment_name}/` and copy the same file to `/meta/` + +
+Example: SO101 modality.json +```json +{ + "state": { + "shoulder_pan": { "start": 0, "end": 1 }, + "shoulder_lift": { "start": 1, "end": 2 }, + "elbow_flex": { "start": 2, "end": 3 }, + "wrist_flex": { "start": 3, "end": 4 }, + "wrist_roll": { "start": 4, "end": 5 }, + "gripper": { "start": 5, "end": 6 } + }, + "action": { + "shoulder_pan": { "start": 0, "end": 1, "absolute": false }, + "shoulder_lift": { "start": 1, "end": 2, "absolute": false }, + "elbow_flex": { "start": 2, "end": 3, "absolute": false }, + "wrist_flex": { "start": 3, "end": 4, "absolute": false }, + "wrist_roll": { "start": 4, "end": 5, "absolute": false }, + "gripper": { "start": 5, "end": 6, "absolute": false } + }, + "video": { + "front": { "original_key": "observation.images.front" } + }, + "annotation": { + "human.task_description": { "original_key": "task_index" } + } +} +``` +
+ +##### Step 2: Add a New DataConfig (So101DataConfig) + +Add a new `So101DataConfig` class defining in `GR00T-Dreams/gr00t/experiment/data_config_idm.py`: + +
+GR00T-Dreams/gr00t/experiment/data_config_idm.py +```python +class So101DataConfig(BaseDataConfig): + video_keys = ["video.front"] + state_keys = ["state.shoulder_pan", "state.shoulder_lift", "state.elbow_flex", "state.wrist_flex", "state.wrist_roll", "state.gripper"] + action_keys = ["action.shoulder_pan", "action.shoulder_lift", "action.elbow_flex", "action.wrist_flex", "action.wrist_roll", "action.gripper"] + language_keys = ["annotation.human.task_description"] + observation_indices = [0, 16] + action_indices = list(range(16)) + + def modality_config(self) -> dict[str, ModalityConfig]: + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + + return modality_configs + + def transform(self) -> ModalityTransform: + transforms = [ + # video transforms + VideoToTensor(apply_to=self.video_keys), + VideoCrop(apply_to=self.video_keys, scale=0.95), + VideoResize(apply_to=self.video_keys, height=224, width=224, interpolation="linear"), + VideoColorJitter( + apply_to=self.video_keys, + brightness=0.3, + contrast=0.4, + saturation=0.5, + hue=0.08, + ), + VideoToNumpy(apply_to=self.video_keys), + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={key: "min_max" for key in self.state_keys}, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={key: "min_max" for key in self.action_keys}, + ), + # concat transforms + ConcatTransform( + video_concat_order=self.video_keys, + state_concat_order=self.state_keys, + action_concat_order=self.action_keys, + ), + # model-specific transform + GR00TIDMTransform( + state_horizon=len(self.observation_indices), + action_horizon=len(self.action_indices), + max_state_dim=64, + max_action_dim=32, + ), + ] + return ComposedModalityTransform(transforms=transforms) +``` + +Register the new config in `DATA_CONFIG_MAP`: +```python +DATA_CONFIG_MAP = { + "gr1_arms_waist": Gr1ArmsWaistDataConfig(), + "gr1_arms_only": Gr1ArmsOnlyDataConfig(), + "gr1_full_upper_body": Gr1FullUpperBodyDataConfig(), + "bimanual_panda_gripper": BimanualPandaGripperDataConfig(), + "bimanual_panda_hand": BimanualPandaHandDataConfig(), + "single_panda_gripper": SinglePandaGripperDataConfig(), + "so100": So100DataConfig(), + "franka": FrankaDataConfig(), + "so101": So101DataConfig(),#add +} +``` +
+ +#### 3.2.2 Run IDM Post-training + +Post-train **IDM (GR00T-Dreams)** using the LeRobot dataset and the newly registered data configuration. + +```bash +PYTHONPATH=. torchrun scripts/idm_training.py \ + --dataset-path \ + --data-config \ + --output_dir +``` + +After loading the dataset, `stats.json` will be automatically generated under `/meta/`. + +**Copy this file to `GR00T-Dreams/IDM_dump/global_metadata/{embodiment_name}/`.** + +### 3.3 Extract Robot Actions to LeRobot Format + +After IDM post-training, use the trained model to infer actions from Cosmos-generated videos. + +#### Step 1: Prepare Inference Configuration + +Select a checkpoint directory (e.g. checkpoint-10000/) and create `checkpoint-10000/experiment_cfg/conf.yaml`: + +
+checkpoint-10000/experiment_cfg/conf.yaml + +```yaml +# Configuration for so101 IDM (LiftCube dataset) +# SO101 robot with 6 DOF: shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper + +modality_configs: + so101: + video: + _target_: gr00t.data.dataset.ModalityConfig + delta_indices: + - 0 + - 16 + modality_keys: + - video.front + state: + _target_: gr00t.data.dataset.ModalityConfig + delta_indices: + - 0 + - 16 + modality_keys: + - state.shoulder_pan + - state.shoulder_lift + - state.elbow_flex + - state.wrist_flex + - state.wrist_roll + - state.gripper + action: + _target_: gr00t.data.dataset.ModalityConfig + delta_indices: + - 0 + - 1 + - 2 + - 3 + - 4 + - 5 + - 6 + - 7 + - 8 + - 9 + - 10 + - 11 + - 12 + - 13 + - 14 + - 15 + modality_keys: + - action.shoulder_pan + - action.shoulder_lift + - action.elbow_flex + - action.wrist_flex + - action.wrist_roll + - action.gripper + language: + _target_: gr00t.data.dataset.ModalityConfig + delta_indices: + - 0 + modality_keys: + - annotation.human.task_description + +all_transforms: + so101: + _target_: gr00t.data.transform.base.ComposedModalityTransform + transforms: + - _target_: gr00t.data.transform.video.VideoToTensor + apply_to: + - video.front + - _target_: gr00t.data.transform.video.VideoCrop + apply_to: + - video.front + scale: 0.95 + - _target_: gr00t.data.transform.video.VideoResize + apply_to: + - video.front + height: 224 + width: 224 + interpolation: linear + - _target_: gr00t.data.transform.video.VideoColorJitter + apply_to: + - video.front + brightness: 0.3 + contrast: 0.4 + saturation: 0.5 + hue: 0.08 + - _target_: gr00t.data.transform.video.VideoToNumpy + apply_to: + - video.front + - _target_: gr00t.data.transform.state_action.StateActionToTensor + apply_to: + - state.shoulder_pan + - state.shoulder_lift + - state.elbow_flex + - state.wrist_flex + - state.wrist_roll + - state.gripper + - _target_: gr00t.data.transform.state_action.StateActionTransform + apply_to: + - state.shoulder_pan + - state.shoulder_lift + - state.elbow_flex + - state.wrist_flex + - state.wrist_roll + - state.gripper + normalization_modes: + state.shoulder_pan: min_max + state.shoulder_lift: min_max + state.elbow_flex: min_max + state.wrist_flex: min_max + state.wrist_roll: min_max + state.gripper: min_max + - _target_: gr00t.data.transform.state_action.StateActionToTensor + apply_to: + - action.shoulder_pan + - action.shoulder_lift + - action.elbow_flex + - action.wrist_flex + - action.wrist_roll + - action.gripper + - _target_: gr00t.data.transform.state_action.StateActionTransform + apply_to: + - action.shoulder_pan + - action.shoulder_lift + - action.elbow_flex + - action.wrist_flex + - action.wrist_roll + - action.gripper + normalization_modes: + action.shoulder_pan: min_max + action.shoulder_lift: min_max + action.elbow_flex: min_max + action.wrist_flex: min_max + action.wrist_roll: min_max + action.gripper: min_max + - _target_: gr00t.data.transform.concat.ConcatTransform + video_concat_order: + - video.front + state_concat_order: + - state.shoulder_pan + - state.shoulder_lift + - state.elbow_flex + - state.wrist_flex + - state.wrist_roll + - state.gripper + action_concat_order: + - action.shoulder_pan + - action.shoulder_lift + - action.elbow_flex + - action.wrist_flex + - action.wrist_roll + - action.gripper + - _target_: gr00t.model.transforms_idm.GR00TIDMTransform + state_horizon: 2 + action_horizon: 16 + max_state_dim: 64 + max_action_dim: 32 + +metadata_versions: + so101: v2.1 +``` + +
+ +Use a configuration matching the SO101 modality definition and transforms (as shown above). + +#### Step 2: Run IDM Inference + +Copy the preprocessing scripts folder `preprocess_leisaac` to `GR00T-Dreams/dump/scripts/` and `so101.sh` to `GR00T-Dreams/IDM_dump/scripts/preprocess/` + +Run the inference script: +```bash +PYTHONPATH=. bash IDM_dump/scripts/preprocess/so101.sh +``` + +This step produces the complete LeRobot-format outputs based on the Cosmos-generated videos. + + + +## Step 4: Replay and Evaluation in LeIsaac + +Switch to the **LeIsaac** environment and project directory. + +Copy the replay helper script `replay.py` into the LeIsaac scripts folder:`leisaac/scripts/`. + +Run the replay script with the **target task** enabled: + +```bash +python scripts/replay.py \ + --task LeIsaac-SO101-LiftCube-v0 \ + --enable_cameras +``` +In this step, the **original HDF5 dataset** and the **IDM-generated LeRobot trajectories (parquet)** are replayed in **Isaac Sim** using **LeIsaac**.This replay process is used to validate the quality and physical plausibility of the inferred action trajectories. + +
+ +
\ No newline at end of file diff --git a/docs/docusaurus.config.js b/docs/docusaurus.config.js index 04041e43..ed65749a 100644 --- a/docs/docusaurus.config.js +++ b/docs/docusaurus.config.js @@ -130,7 +130,7 @@ const config = { prism: { theme: prismThemes.github, darkTheme: prismThemes.dracula, - additionalLanguages: ['bash'] + additionalLanguages: ['bash', 'diff'] }, }), }; diff --git a/docs/sidebars.js b/docs/sidebars.js index 6bd6fc4e..e782e7f9 100644 --- a/docs/sidebars.js +++ b/docs/sidebars.js @@ -75,6 +75,12 @@ const sidebars = { link: { type: 'doc', id: 'docs/tutorials/marble_tutorial' }, items: [], }, + { + type: 'category', + label: 'LeIsaac × Cosmos', + link: { type: 'doc', id: 'docs/tutorials/cosmos_tutorial' }, + items: [], + }, ], }, { From 3593ad2dfecc839a00baeafb315f7f92c30009c7 Mon Sep 17 00:00:00 2001 From: EverNorif <1320346985@qq.com> Date: Wed, 7 Jan 2026 16:22:05 +0800 Subject: [PATCH 3/9] fix docs. --- docs/docs/docs/tutorials/cosmos_tutorial.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/docs/docs/tutorials/cosmos_tutorial.md b/docs/docs/docs/tutorials/cosmos_tutorial.md index 7ea00885..614f6862 100644 --- a/docs/docs/docs/tutorials/cosmos_tutorial.md +++ b/docs/docs/docs/tutorials/cosmos_tutorial.md @@ -112,7 +112,7 @@ After post-training, the Cosmos-Predict2.5 model checkpoints are typically saved Once the checkpoint has been converted, run video generation inference following the [official running-inference instructions ](https://github.com/nvidia-cosmos/cosmos-predict2.5/blob/main/docs/post-training_video2world_gr00t.md#32-running-inference). -#### 🔬**Batch Inference (Used for fine-tuning IDM)** +🔬 **Batch Inference (Used for fine-tuning IDM)** For large-scale video generation, batch inference is supported. Copy the batch inference helper script `generate_batch_config.py` into the folder:`cosmos-predict2.5/scripts/`. @@ -155,7 +155,7 @@ For IDM fine-tuning, please refer to **[training-custom-idm-model](https://githu #### 3.2.1 Preparation: Modality Metadata and DataConfig -##### Step 1: Add modality.json +**Step 1: Add modality.json** Create `modality.json` under `GR00T-Dreams/IDM_dump/global_metadata/{embodiment_name}/` and copy the same file to `/meta/` @@ -189,7 +189,7 @@ Create `modality.json` under `GR00T-Dreams/IDM_dump/global_metadata/{embodiment_ ``` -##### Step 2: Add a New DataConfig (So101DataConfig) +**Step 2: Add a New DataConfig (So101DataConfig)** Add a new `So101DataConfig` class defining in `GR00T-Dreams/gr00t/experiment/data_config_idm.py`: From 43540d9926732f511ce0fd524e03df3458664671 Mon Sep 17 00:00:00 2001 From: chuchen0210 <167320956+chuchen0210@users.noreply.github.com> Date: Thu, 8 Jan 2026 15:35:49 +0800 Subject: [PATCH 4/9] Add lerobot2isaaclab converter(feature:cosmos+IDM) --- scripts/convert/lerobot2isaaclab.py | 416 +++++++++++++ scripts/environments/list_envs.py | 3 +- scripts/environments/teleoperation/replay.py | 3 +- scripts/evaluation/policy_inference.py | 3 +- scripts/mimic/annotate_demos.py | 3 +- scripts/mimic/generate_dataset.py | 3 +- .../cosmos_tutorials/generate_batch_config.py | 125 ---- .../cosmos2.5_to_step2_format.py | 279 --------- .../preprocess_leisaac/dump_idm_actions.py | 481 --------------- .../lerobot_to_step2_format.py | 161 ----- .../preprocess_leisaac/preprocess_video.py | 376 ------------ .../preprocess_leisaac/raw_to_lerobot.py | 568 ------------------ .../split_video_instruction.py | 101 ---- scripts/tutorials/cosmos_tutorials/replay.py | 495 --------------- scripts/tutorials/cosmos_tutorials/so101.sh | 149 ----- 15 files changed, 421 insertions(+), 2745 deletions(-) create mode 100755 scripts/convert/lerobot2isaaclab.py delete mode 100644 scripts/tutorials/cosmos_tutorials/generate_batch_config.py delete mode 100644 scripts/tutorials/cosmos_tutorials/preprocess_leisaac/cosmos2.5_to_step2_format.py delete mode 100644 scripts/tutorials/cosmos_tutorials/preprocess_leisaac/dump_idm_actions.py delete mode 100644 scripts/tutorials/cosmos_tutorials/preprocess_leisaac/lerobot_to_step2_format.py delete mode 100644 scripts/tutorials/cosmos_tutorials/preprocess_leisaac/preprocess_video.py delete mode 100644 scripts/tutorials/cosmos_tutorials/preprocess_leisaac/raw_to_lerobot.py delete mode 100644 scripts/tutorials/cosmos_tutorials/preprocess_leisaac/split_video_instruction.py delete mode 100644 scripts/tutorials/cosmos_tutorials/replay.py delete mode 100644 scripts/tutorials/cosmos_tutorials/so101.sh diff --git a/scripts/convert/lerobot2isaaclab.py b/scripts/convert/lerobot2isaaclab.py new file mode 100755 index 00000000..ffc619c6 --- /dev/null +++ b/scripts/convert/lerobot2isaaclab.py @@ -0,0 +1,416 @@ +#!/usr/bin/env python3 + +""" +Convert IDM parquet (LeRobot-style rows) + source LeIsaac HDF5 (initial_state/template) +to a new replayable LeIsaac HDF5. + +Input parquet row example: +{ + "observation.state":[...6...], + "action":[...6...], + "timestamp":0, + "episode_index":0, + "index":0 +} + +Source HDF5 episode structure (example): +/actions (T,6) +/initial_state/... +/obs/actions (T,6) +/processed_actions (T,6) +/states/articulation/robot/joint_position (T,6) +... + +We will: +- Copy root metadata groups (everything except /data) from source HDF5. +- For each episode: + - Copy initial_state + attrs from a template episode in source HDF5 + - Write /actions, /obs/actions, /processed_actions from parquet "action" + - Optionally write /states/.../joint_position from parquet "observation.state" + - Optionally write /timestamps from parquet "timestamp" + - Optionally copy template obs/states arrays (including images) from source HDF5(can be huge) +""" + +import argparse +import os +import re +from pathlib import Path + +import h5py +import numpy as np +import pandas as pd +from tqdm import tqdm + +# ------------------------ +# Denormalize function (same as isaaclab2lerobot) +# ------------------------ + +ISAACLAB_LIMITS = [(-110.0, 110.0), (-100.0, 100.0), (-100.0, 90.0), (-95.0, 95.0), (-160.0, 160.0), (-10.0, 100.0)] +LEROBOT_LIMITS = [(-100.0, 100.0), (-100.0, 100.0), (-100.0, 100.0), (-100.0, 100.0), (-100.0, 100.0), (0.0, 100.0)] + + +def denormalize_lerobot_to_isaaclab_radians(joint_pos: np.ndarray) -> np.ndarray: + """ + LeRobot normalized degrees in roughly [-100,100] -> IsaacLab joint limits (degrees) -> radians. + joint_pos: (T,6) + """ + result = joint_pos.astype(np.float32).copy() + for i in range(6): + isa_min, isa_max = ISAACLAB_LIMITS[i] + le_min, le_max = LEROBOT_LIMITS[i] + isa_range = isa_max - isa_min + le_range = le_max - le_min + result[:, i] = (result[:, i] - le_min) / le_range * isa_range + isa_min + # degrees -> radians + result = result * (np.pi / 180.0) + return result + + +# ------------------------ +# HDF5 copy helpers +# ------------------------ + + +def copy_attrs(src_obj, dst_obj): + for k, v in src_obj.attrs.items(): + dst_obj.attrs[k] = v + + +def copy_group_recursive(src: h5py.Group, dst: h5py.Group, skip: list[str] | None = None): + """ + Recursively copy a group, skipping relative paths in `skip`. + """ + skip = set(skip or []) + + def _rec(s: h5py.Group | h5py.File, d: h5py.Group, rel: str): + copy_attrs(s, d) + for name, item in s.items(): + child_rel = f"{rel}/{name}" if rel else name + if child_rel in skip: + continue + if isinstance(item, h5py.Dataset): + ds = d.create_dataset(name, data=item[()], compression=item.compression) + for ak, av in item.attrs.items(): + ds.attrs[ak] = av + else: + g = d.create_group(name) + _rec(item, g, child_rel) + + _rec(src, dst, "") + + +def ensure_group(root: h5py.Group, path: str) -> h5py.Group: + """ + Ensure group path exists under root. Return the group. + """ + path = path.strip("/") + cur = root + if not path: + return cur + for part in path.split("/"): + if part not in cur: + cur = cur.create_group(part) + else: + cur = cur[part] + if not isinstance(cur, h5py.Group): + raise ValueError(f"Path conflict: {path} has non-group at {part}") + return cur + + +def write_dataset(root: h5py.Group, path: str, data: np.ndarray, compression: str | None = None): + """ + Overwrite dataset at path with given data. + """ + path = path.strip("/") + parent_path, name = path.rsplit("/", 1) if "/" in path else ("", path) + parent = ensure_group(root, parent_path) + if name in parent: + del parent[name] + parent.create_dataset(name, data=data, compression=compression) + + +# ------------------------ +# Parquet loading +# ------------------------ + +_EP_RE = re.compile(r"episode_(\d+)\.parquet$") + + +def list_episode_parquet_files(parquet_dir: str) -> list[tuple[int, Path]]: + """Return [(episode_id_from_filename, path), ...] sorted by episode_id.""" + p = Path(parquet_dir) + if not p.is_dir(): + raise ValueError(f"--parquet must be a directory when using filename mapping: {parquet_dir}") + files = sorted(p.glob("**/episode_*.parquet")) + items: list[tuple[int, Path]] = [] + for f in files: + m = _EP_RE.search(f.name) + if not m: + continue + items.append((int(m.group(1)), f)) + if not items: + raise FileNotFoundError(f"No episode_*.parquet found under: {parquet_dir}") + items.sort(key=lambda x: x[0]) + return items + + +def load_parquet_input(parquet_path_or_dir: str) -> pd.DataFrame: + p = Path(parquet_path_or_dir) + if p.is_dir(): + files = sorted( + p.glob("**/*.parquet"), + key=lambda x: ( + int(re.search(r"episode_(\d+)", x.name).group(1)) if re.search(r"episode_(\d+)", x.name) else 0 + ), + ) + if not files: + raise FileNotFoundError(f"No parquet files found in dir: {parquet_path_or_dir}") + dfs = [pd.read_parquet(f) for f in files] + df = pd.concat(dfs, ignore_index=True) + else: + if not p.exists(): + raise FileNotFoundError(parquet_path_or_dir) + df = pd.read_parquet(p) + return df + + +def stack_vector_column(df: pd.DataFrame, col: str) -> np.ndarray: + # column entries are list/np.ndarray per row + return np.stack(df[col].to_numpy()) + + +# ------------------------ +# Main conversion +# ------------------------ + + +def convert( # noqa: C901 + parquet_path_or_dir: str, + src_hdf5: str, + out_hdf5: str, + overwrite: bool, + episode_col: str, + sort_col: str | None, + action_col: str, + state_col: str, + timestamp_col: str | None, + write_states: bool, + keep_template_obs: bool, +): + if os.path.exists(out_hdf5): + if overwrite: + os.remove(out_hdf5) + else: + raise FileExistsError(f"Output exists: {out_hdf5} (use --overwrite)") + + # Use filename-based mapping when input is a directory containing episode_*.parquet + parquet_items = list_episode_parquet_files(parquet_path_or_dir) if Path(parquet_path_or_dir).is_dir() else None + + if parquet_items is None: + # fallback: single parquet file mode (still uses episode_index) + df = load_parquet_input(parquet_path_or_dir) + + for need in [episode_col, action_col]: + if need not in df.columns: + raise ValueError(f"Missing required column '{need}' in parquet. columns={list(df.columns)}") + + if write_states and state_col not in df.columns: + raise ValueError(f"--write_states enabled but missing column '{state_col}'") + + if sort_col and sort_col not in df.columns: + raise ValueError(f"--sort_col '{sort_col}' not in columns") + + if timestamp_col and timestamp_col not in df.columns: + raise ValueError(f"--timestamp_col '{timestamp_col}' not in columns") + + ep_ids = sorted(df[episode_col].unique().tolist()) + print(f"[INFO] parquet episodes (from column {episode_col}): {len(ep_ids)}") + else: + print(f"[INFO] parquet files (from filename): {len(parquet_items)}") + + with h5py.File(src_hdf5, "r") as f_src, h5py.File(out_hdf5, "w") as f_out: + if "data" not in f_src: + raise ValueError("Source HDF5 missing /data") + + # copy root metadata except /data + for k, item in f_src.items(): + if k == "data": + continue + if isinstance(item, h5py.Dataset): + ds = f_out.create_dataset(k, data=item[()], compression=item.compression) + for ak, av in item.attrs.items(): + ds.attrs[ak] = av + else: + g = f_out.create_group(k) + copy_group_recursive(item, g) + + g_data_out = f_out.create_group("data") + g_data_src = f_src["data"] + copy_attrs(g_data_src, g_data_out) + + src_ep_names = list(g_data_src.keys()) + if not src_ep_names: + raise ValueError("Source HDF5 has zero episodes under /data") + template_name = src_ep_names[0] + + # We will copy: + # - attrs + # - initial_state group + # optionally: obs/states skeleton (but not the big image arrays unless user wants to keep) + # Note: copying images can explode file size; default keep_template_obs=False. + base_skip = [] + if not keep_template_obs: + # skip heavy video/images and other per-step arrays from template + # we'll only write the minimum required datasets. + base_skip += [ + "obs/front", + "obs/ee_frame_state", + "obs/joint_pos", + "obs/joint_pos_rel", + "obs/joint_pos_target", + "obs/joint_vel", + "obs/joint_vel_rel", + "states", + "actions", + "obs/actions", + "processed_actions", + ] + + if parquet_items is None: + # old behavior: split by episode_col inside one big df + iterator = [(int(ep_id), None) for ep_id in ep_ids] + else: + # new behavior: each file defines one episode, use filename episode id + iterator = parquet_items # [(ep_id, path), ...] + + for ep_id, pq_path in tqdm(iterator, desc="Writing episodes"): + if pq_path is None: + dfe = df[df[episode_col] == ep_id].copy() + else: + dfe = pd.read_parquet(pq_path) + + # basic column checks (per-file for filename mode) + if action_col not in dfe.columns: + raise ValueError(f"Missing required column '{action_col}' in parquet for ep={ep_id}: {pq_path}") + if write_states and state_col not in dfe.columns: + raise ValueError(f"--write_states enabled but missing '{state_col}' for ep={ep_id}: {pq_path}") + if sort_col and sort_col not in dfe.columns: + raise ValueError(f"--sort_col '{sort_col}' not in parquet columns for ep={ep_id}: {pq_path}") + if timestamp_col and timestamp_col not in dfe.columns: + raise ValueError(f"--timestamp_col '{timestamp_col}' not in parquet columns for ep={ep_id}: {pq_path}") + + if sort_col: + dfe = dfe.sort_values(sort_col, ascending=True) + + # actions (T,6) in lerobot normalized degrees + act_raw = stack_vector_column(dfe, action_col).astype(np.float32) + if act_raw.ndim != 2 or act_raw.shape[1] != 6: + raise ValueError(f"Action shape expected (T,6), got {act_raw.shape} for ep={ep_id}") + act = denormalize_lerobot_to_isaaclab_radians(act_raw) + + # states (T,6) optional + st = None + if write_states: + st_raw = stack_vector_column(dfe, state_col).astype(np.float32) + if st_raw.ndim != 2 or st_raw.shape[1] != 6: + raise ValueError(f"State shape expected (T,6), got {st_raw.shape} for ep={ep_id}") + st = denormalize_lerobot_to_isaaclab_radians(st_raw) + + ts = None + if timestamp_col: + ts = dfe[timestamp_col].to_numpy(dtype=np.int64) + + # choose a source episode to copy initial_state from (now: use ep_id from filename) + src_name = f"demo_{int(ep_id)}" + if src_name not in g_data_src: + # fallback map by index + if int(ep_id) < len(src_ep_names): + src_name = src_ep_names[int(ep_id)] + else: + src_name = template_name + + g_src_ep = g_data_src[src_name] + dst_name = f"demo_{int(ep_id)}" + g_dst_ep = g_data_out.create_group(dst_name) + + # copy template stuff (attrs + initial_state + maybe skeleton) + copy_attrs(g_src_ep, g_dst_ep) + if "initial_state" in g_src_ep: + g_init = g_dst_ep.create_group("initial_state") + copy_group_recursive(g_src_ep["initial_state"], g_init) + else: + raise ValueError(f"Source episode {src_name} missing /initial_state") + + for grp_name in ["obs", "states"]: + if grp_name in g_src_ep and keep_template_obs: + g = g_dst_ep.create_group(grp_name) + copy_group_recursive(g_src_ep[grp_name], g) + + # Write actions into required places + write_dataset(g_dst_ep, "actions", act, compression=None) + write_dataset(g_dst_ep, "obs/actions", act, compression=None) + write_dataset(g_dst_ep, "processed_actions", act, compression=None) + + if write_states and st is not None: + write_dataset(g_dst_ep, "states/articulation/robot/joint_position", st, compression=None) + + if ts is not None: + write_dataset(g_dst_ep, "timestamps", ts, compression=None) + + g_dst_ep.attrs["num_steps"] = int(act.shape[0]) + g_dst_ep.attrs["action_dim"] = int(act.shape[1]) + g_dst_ep.attrs["converted_from_src_episode"] = src_name + g_dst_ep.attrs["idm_source"] = str(Path(parquet_path_or_dir).resolve()) + if pq_path is not None: + g_dst_ep.attrs["idm_parquet_file"] = pq_path.name + + print(f"[DONE] Output written: {out_hdf5}") + print("Run replay with:") + print(f" --dataset_file {out_hdf5}") + + +def build_parser(): + p = argparse.ArgumentParser("IDM parquet -> LeIsaac HDF5 converter") + p.add_argument("--parquet", type=str, required=True, help="Parquet file OR directory containing parquet(s).") + p.add_argument("--src_hdf5", type=str, required=True, help="Source teleop HDF5 for initial_state/template.") + p.add_argument("--out_hdf5", type=str, required=True, help="Output HDF5 for replay.") + p.add_argument("--overwrite", action="store_true") + + p.add_argument("--episode_col", type=str, default="episode_index") + p.add_argument("--sort_col", type=str, default="index", help="Order within episode. default: index") + p.add_argument("--action_col", type=str, default="action") + p.add_argument("--state_col", type=str, default="observation.state") + p.add_argument("--timestamp_col", type=str, default="timestamp") + + p.add_argument( + "--write_states", + action="store_true", + help="Also write /states/articulation/robot/joint_position from observation.state", + ) + p.add_argument( + "--keep_template_obs", + action="store_true", + help="Copy template obs/states arrays (including images) from source (can be huge). Default off.", + ) + return p + + +def main(): + args = build_parser().parse_args() + convert( + parquet_path_or_dir=args.parquet, + src_hdf5=args.src_hdf5, + out_hdf5=args.out_hdf5, + overwrite=args.overwrite, + episode_col=args.episode_col, + sort_col=args.sort_col, + action_col=args.action_col, + state_col=args.state_col, + timestamp_col=args.timestamp_col, + write_states=args.write_states, + keep_template_obs=args.keep_template_obs, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/environments/list_envs.py b/scripts/environments/list_envs.py index e2b9e813..a43e953d 100644 --- a/scripts/environments/list_envs.py +++ b/scripts/environments/list_envs.py @@ -24,9 +24,8 @@ """Rest everything follows.""" import gymnasium as gym -from prettytable import PrettyTable - import leisaac # noqa: F401 +from prettytable import PrettyTable def main(): diff --git a/scripts/environments/teleoperation/replay.py b/scripts/environments/teleoperation/replay.py index c653a13e..50a4fc77 100644 --- a/scripts/environments/teleoperation/replay.py +++ b/scripts/environments/teleoperation/replay.py @@ -57,6 +57,7 @@ import time import gymnasium as gym +import leisaac # noqa: F401 import torch from isaaclab.envs import DirectRLEnv, ManagerBasedRLEnv from isaaclab.utils.datasets import EpisodeData, HDF5DatasetFileHandler @@ -66,8 +67,6 @@ get_task_type, ) -import leisaac # noqa: F401 - class RateLimiter: """Convenience class for enforcing rates in loops.""" diff --git a/scripts/evaluation/policy_inference.py b/scripts/evaluation/policy_inference.py index 7ee7a3dc..9dfd3dd5 100644 --- a/scripts/evaluation/policy_inference.py +++ b/scripts/evaluation/policy_inference.py @@ -53,6 +53,7 @@ import carb import gymnasium as gym +import leisaac # noqa: F401 import omni import torch from isaaclab.envs import ManagerBasedRLEnv @@ -62,8 +63,6 @@ get_task_type, ) -import leisaac # noqa: F401 - class RateLimiter: """Convenience class for enforcing rates in loops.""" diff --git a/scripts/mimic/annotate_demos.py b/scripts/mimic/annotate_demos.py index f99ef769..5ec29da6 100644 --- a/scripts/mimic/annotate_demos.py +++ b/scripts/mimic/annotate_demos.py @@ -74,6 +74,7 @@ from isaaclab.devices import Se3Keyboard import isaaclab_tasks # noqa: F401 +import leisaac # noqa: F401 from isaaclab.envs import ManagerBasedRLMimicEnv from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg from isaaclab.managers import RecorderTerm, RecorderTermCfg, TerminationTermCfg @@ -85,8 +86,6 @@ get_task_type, ) -import leisaac # noqa: F401 - is_paused = False current_action_index = 0 marked_subtask_action_indices = [] diff --git a/scripts/mimic/generate_dataset.py b/scripts/mimic/generate_dataset.py index c2212aaf..8928f182 100644 --- a/scripts/mimic/generate_dataset.py +++ b/scripts/mimic/generate_dataset.py @@ -79,6 +79,7 @@ import isaaclab_mimic.envs.pinocchio_envs # noqa: F401 import isaaclab_tasks # noqa: F401 +import leisaac # noqa: F401 from isaaclab_mimic.datagen.generation import ( env_loop, setup_async_generation, @@ -87,8 +88,6 @@ from isaaclab_mimic.datagen.utils import get_env_name_from_dataset, setup_output_paths from leisaac.utils.env_utils import get_task_type -import leisaac # noqa: F401 - def main(): num_envs = args_cli.num_envs diff --git a/scripts/tutorials/cosmos_tutorials/generate_batch_config.py b/scripts/tutorials/cosmos_tutorials/generate_batch_config.py deleted file mode 100644 index ec17ba87..00000000 --- a/scripts/tutorials/cosmos_tutorials/generate_batch_config.py +++ /dev/null @@ -1,125 +0,0 @@ -#!/usr/bin/env python3 -""" -Generate batch inference configuration files (JSONL). - -This script scans a directory of input videos and generates a JSONL file -for batch video2world inference. Each line corresponds to one video. -Optionally, task prompts can be loaded from episodes.jsonl. -""" - -import json -import os -from pathlib import Path - -# --------------------------------------------------------------------- -# Configuration -# --------------------------------------------------------------------- - -# Input video directory used to generate inference configs -VIDEO_DIR = "/videos/chunk-000/observation.images.front" - -# Metadata directory (used when loading text prompts) -META_DIR = "/meta" - -# Output JSONL file -OUTPUT_JSONL = "batch_inference_config.jsonl" - -# Task name prefix -TASK_NAME = "liftcube" - -# --------------------------------------------------------------------- -# Base configuration template shared by all videos -# --------------------------------------------------------------------- - -BASE_CONFIG = { - "inference_type": "video2world", - "seed": 21, - "guidance": 7, - "resolution": "480,640", - "enable_autoregressive": True, - # Number of output frames (e.g. 110 ≈ 6s, 140 ≈ 8s, 210 ≈ 16s) - "num_output_frames": 210, - "chunk_size": 77, - "chunk_overlap": 1, - # Default prompt (will be overridden if episode-specific prompt is found) - "prompt": "The robot arm is performing a task", - "negative_prompt": ( - "The video captures a series of frames showing ugly scenes, static with no motion, " - "motion blur, over-saturation, shaky footage, low resolution, grainy texture, " - "pixelated images, poorly lit areas, underexposed and overexposed scenes, " - "poor color balance, washed out colors, choppy sequences, jerky movements, " - "low frame rate, artifacting, color banding, unnatural transitions, " - "outdated special effects, fake elements, unconvincing visuals, " - "poorly edited content, jump cuts, visual noise, and flickering. " - "Overall, the video is of poor quality." - ), -} - -# --------------------------------------------------------------------- -# Main logic -# --------------------------------------------------------------------- - -def main(use_prompt=True): - video_dir = Path(VIDEO_DIR) - meta_dir = Path(META_DIR) - - # Collect all mp4 video files - video_files = sorted(video_dir.glob("*.mp4")) - print(f"Found {len(video_files)} video files") - - # Load prompts from episodes.jsonl if enabled - episode_prompts = {} - if use_prompt: - episodes_file = meta_dir / "episodes.jsonl" - if episodes_file.exists(): - with open(episodes_file, 'r') as f: - for line in f: - episode_data = json.loads(line) - episode_index = episode_data.get("episode_index") - tasks = episode_data.get("tasks", []) - if episode_index is not None and tasks: - # Use the first task as the episode prompt - prompt = f"The robot arm is performing a task. {tasks[0]}" - episode_prompts[episode_index] = prompt - print(f"✓ Loaded {len(episode_prompts)} prompts from episodes.jsonl") - - # Generate JSONL configuration file - with open(OUTPUT_JSONL, 'w') as f: - for video_file in video_files: - # Video filename without extension (e.g. episode_000001) - video_name = video_file.stem - - # Create per-video config based on the base template - config = BASE_CONFIG.copy() - config["name"] = f"{TASK_NAME}_{video_name}" - config["input_path"] = str(video_file) - - # Assign episode-specific prompt if available - if use_prompt: - # Extract episode index from video filename - episode_index = int(video_name.split('_')[-1]) - if episode_index in episode_prompts: - config["prompt"] = episode_prompts[episode_index] - - # Write one JSON object per line - f.write(json.dumps(config) + '\n') - - # Summary information - print(f"✓ Config file generated: {OUTPUT_JSONL}") - print(f" - Total videos: {len(video_files)}") - print(f" - Use prompts: {'Yes' if use_prompt else 'No'}") - if use_prompt: - print( - f" - Videos with matched prompts: " - f"{sum(1 for vf in video_files if int(vf.stem.split('_')[-1]) in episode_prompts)}" - ) - -# --------------------------------------------------------------------- -# Script entry point -# --------------------------------------------------------------------- - -if __name__ == "__main__": - import sys - # Text prompts are enabled only when '--use-prompt' is provided - use_prompt = "--use-prompt" in sys.argv - main(use_prompt=use_prompt) diff --git a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/cosmos2.5_to_step2_format.py b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/cosmos2.5_to_step2_format.py deleted file mode 100644 index 8d5ffbce..00000000 --- a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/cosmos2.5_to_step2_format.py +++ /dev/null @@ -1,279 +0,0 @@ -#!/usr/bin/env python3 -""" -Convert cosmos2.5 outputs (+ optional LeRobot meta) to convert_directory output format. - -Input: - cosmos_dir/ - ├── *.mp4 - ├── *.json(same stem as mp4) - └── ... - -Optional: - lerobot_dir/ - └── meta/ - ├── tasks.jsonl - └── episodes.jsonl - -Output: - output_dir/ - ├── / - │ ├── 0.mp4 - │ ├── 1.mp4 - │ └── ... - └── ... - -Rule: - - If json has "prompt": use it as task name (sanitize to dir) - - Else: fallback to LeRobot meta (requires --lerobot_dir) - - If lerobot_dir not provided: print warning and skip that sample -""" - -import argparse -import json -import shutil -from pathlib import Path -from collections import defaultdict -from typing import Dict, List, Optional, Tuple - - -def sanitize_task_name(task: str) -> str: - """ - Convert task description to valid directory name. - Example: "Lift the red cube up." -> "Lift_the_red_cube_up" - """ - task = task.strip() - for ch in [".", ",", "!", "?", ":", ";", "\"", "'"]: - task = task.replace(ch, "") - task = task.replace(" ", "_") - task = "_".join(filter(None, task.split("_"))) - return task - - -def load_tasks(tasks_file: Path) -> Dict[int, str]: - """Load task definitions. Returns {task_index: task_string}""" - tasks: Dict[int, str] = {} - with open(tasks_file, "r") as f: - for line in f: - if not line.strip(): - continue - data = json.loads(line) - tasks[int(data["task_index"])] = data["task"] - return tasks - - -def load_episodes(episodes_file: Path) -> List[dict]: - """Load episode information. Returns list of dict (episode_index, tasks, length, ...)""" - episodes: List[dict] = [] - with open(episodes_file, "r") as f: - for line in f: - if not line.strip(): - continue - episodes.append(json.loads(line)) - return episodes - - -def build_episode_to_task_map(lerobot_dir: Path) -> Dict[int, str]: - """ - Build mapping: episode_index -> task_string (first task in episode['tasks']) - """ - meta_dir = lerobot_dir / "meta" - tasks_path = meta_dir / "tasks.jsonl" - episodes_path = meta_dir / "episodes.jsonl" - - if not tasks_path.exists() or not episodes_path.exists(): - raise FileNotFoundError(f"LeRobot meta files not found under: {meta_dir}") - - _tasks = load_tasks(tasks_path) # not strictly needed, but useful if episodes store indices in some setups - episodes = load_episodes(episodes_path) - - ep2task: Dict[int, str] = {} - for ep in episodes: - ep_idx = int(ep["episode_index"]) - ep_tasks = ep.get("tasks", []) - if not ep_tasks: - continue - - # In many LeRobot datasets, ep["tasks"] stores task strings already. - # If it stores indices, you can extend here. We'll support both: - t0 = ep_tasks[0] - if isinstance(t0, int): - task_str = _tasks.get(int(t0), str(t0)) - else: - task_str = str(t0) - - ep2task[ep_idx] = task_str - - return ep2task - - -def read_json(path: Path) -> Optional[dict]: - try: - with open(path, "r") as f: - return json.load(f) - except Exception as e: - print(f"[WARN] Failed to read json: {path} ({e})") - return None - - -def parse_episode_index_from_stem(stem: str) -> Optional[int]: - """ - Try to extract episode index from filename stem. - Examples it can handle: - - episode_000123 - - ..._000123 - - 000123 - If cannot parse, return None. - """ - # Most common: "episode_000123" - if stem.startswith("episode_"): - tail = stem[len("episode_") :] - if tail.isdigit(): - return int(tail) - - # Try last underscore chunk - parts = stem.split("_") - for candidate in reversed(parts): - if candidate.isdigit(): - return int(candidate) - - # Entire stem digits? - if stem.isdigit(): - return int(stem) - - return None - - -def convert_cosmos_to_step2( - cosmos_dir: Path, - output_dir: Path, - lerobot_dir: Optional[Path] = None, - chunk_missing_prompt_policy: str = "skip", -): - """ - Convert cosmos2.5 outputs to step2 format. - - chunk_missing_prompt_policy: - - "skip": if no prompt and no lerobot_dir mapping, skip the sample - """ - cosmos_dir = cosmos_dir.resolve() - output_dir = output_dir.resolve() - output_dir.mkdir(parents=True, exist_ok=True) - - ep2task: Dict[int, str] = {} - if lerobot_dir is not None: - ep2task = build_episode_to_task_map(lerobot_dir.resolve()) - - # Collect pairs (json, mp4) by stem - json_files = sorted(cosmos_dir.glob("*.json")) - if not json_files: - raise FileNotFoundError(f"No .json files found in cosmos_dir: {cosmos_dir}") - - # task_name -> list of source mp4 paths (ordered) - task_to_videos: Dict[str, List[Path]] = defaultdict(list) - skipped: List[Tuple[Path, str]] = [] - - for jpath in json_files: - stem = jpath.stem - mpath = cosmos_dir / f"{stem}.mp4" - if not mpath.exists(): - print(f"[WARN] Missing mp4 for json: {jpath.name} -> expected {mpath.name}, skip.") - skipped.append((jpath, "missing_mp4")) - continue - - data = read_json(jpath) - if data is None: - skipped.append((jpath, "bad_json")) - continue - - prompt = data.get("prompt", None) - if isinstance(prompt, str) and prompt.strip(): - task_str = prompt.strip() - task_dir_name = sanitize_task_name(task_str) - task_to_videos[task_dir_name].append(mpath) - continue - - # No prompt -> fallback to lerobot meta mapping - if not ep2task: - print( - f"[WARN] {jpath.name} has no 'prompt'. " - f"You didn't provide --lerobot_dir (or mapping is empty), cannot infer task. Skipping." - ) - skipped.append((jpath, "no_prompt_no_lerobot")) - continue - - ep_idx = parse_episode_index_from_stem(stem) - if ep_idx is None: - print( - f"[WARN] {jpath.name} has no 'prompt' and episode index cannot be parsed from name '{stem}'. Skipping." - ) - skipped.append((jpath, "no_prompt_cannot_parse_episode")) - continue - - task_str = ep2task.get(ep_idx) - if not task_str: - print( - f"[WARN] {jpath.name} has no 'prompt'. Parsed episode_index={ep_idx}, " - f"but it's not found in lerobot meta. Skipping." - ) - skipped.append((jpath, "episode_not_in_meta")) - continue - - task_dir_name = sanitize_task_name(task_str) - task_to_videos[task_dir_name].append(mpath) - - # Copy into output folders with sequential numbering per task - total_copied = 0 - for task_dir_name, vids in sorted(task_to_videos.items(), key=lambda x: x[0]): - dst_task_dir = output_dir / task_dir_name - dst_task_dir.mkdir(parents=True, exist_ok=True) - - # keep deterministic order - vids_sorted = sorted(vids, key=lambda p: p.name) - for i, src in enumerate(vids_sorted): - dst = dst_task_dir / f"{i}.mp4" - shutil.copy2(src, dst) - total_copied += 1 - - print(f"[OK] Task '{task_dir_name}': {len(vids_sorted)} videos") - - print("\nConversion complete!") - print(f" Total videos copied: {total_copied}") - print(f" Output directory: {output_dir}") - - if skipped: - print(f"\n[SUMMARY] Skipped {len(skipped)} samples:") - # print a few for readability - for p, reason in skipped[:20]: - print(f" - {p.name}: {reason}") - if len(skipped) > 20: - print(f" ... and {len(skipped) - 20} more") - - -def main(): - parser = argparse.ArgumentParser( - description="Convert cosmos2.5 outputs (+ optional LeRobot meta) to convert_directory output format" - ) - parser.add_argument("--cosmos_dir", type=str, required=True, help="Cosmos output directory containing *.mp4 and *.json") - parser.add_argument("--output_dir", type=str, required=True, help="Output directory") - parser.add_argument( - "--lerobot_dir", - type=str, - default=None, - help="Optional LeRobot dataset directory (only needed when json has no 'prompt')", - ) - - args = parser.parse_args() - - cosmos_dir = Path(args.cosmos_dir) - output_dir = Path(args.output_dir) - lerobot_dir = Path(args.lerobot_dir) if args.lerobot_dir else None - - convert_cosmos_to_step2( - cosmos_dir=cosmos_dir, - output_dir=output_dir, - lerobot_dir=lerobot_dir, - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/dump_idm_actions.py b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/dump_idm_actions.py deleted file mode 100644 index 003c27b0..00000000 --- a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/dump_idm_actions.py +++ /dev/null @@ -1,481 +0,0 @@ -import argparse -import json -import multiprocessing as mp -import os - -import numpy as np -import torch -import yaml -from hydra.utils import instantiate -from omegaconf import OmegaConf -from tqdm import tqdm -from concurrent.futures import ProcessPoolExecutor, as_completed -from tianshou.data import Batch -from huggingface_hub import hf_hub_download -from gr00t.data.dataset import LeRobotSingleDataset -from gr00t.model.idm import IDM -from gr00t.utils.video import get_all_frames_and_timestamps -from gr00t.data.embodiment_tags import EmbodimentTag - - - -def load_dataset_and_config(checkpoint_path, validation_dataset_path, video_indices): - # Check if checkpoint_path is a HuggingFace model repo - is_hf_repo = not os.path.exists(checkpoint_path) and '/' in checkpoint_path - - if is_hf_repo: - # For HuggingFace repos, we need to download config differently - try: - # Download the config file from the repo - config_file = hf_hub_download( - repo_id=checkpoint_path, - filename="experiment_cfg/conf.yaml", - repo_type="model" - ) - with open(config_file, "r") as f: - config = yaml.safe_load(f) - except Exception as e: - print(f"Error loading config from HuggingFace repo: {e}") - raise - else: - # Original local path handling - config_path = os.path.join(checkpoint_path, "experiment_cfg", "conf.yaml") - with open(config_path, "r") as f: - config = yaml.safe_load(f) - - cfg = OmegaConf.create(config) - - dataset_name = os.path.basename(validation_dataset_path).split(".")[0] - embodiment = dataset_name - print(f"Dataset name: {dataset_name}") - - modality_configs = cfg["modality_configs"][embodiment] - if video_indices is not None: - video_delta_indices = video_indices.split(" ") - print(f"Using provided video_delta_indices: {video_delta_indices}") - modality_configs["video"]["delta_indices"] = video_delta_indices - modality_configs = instantiate(modality_configs) - - - if "all_transforms" in cfg: - transform = cfg["all_transforms"][embodiment] - else: - transform = cfg["train_dataset"]["all_transforms"][embodiment] - - # Filter out VideoColorJitter transform for inference - if "transforms" in transform: - filtered_transforms = [] - for t in transform["transforms"]: - if t.get("_target_") != "groot.data.transform.VideoColorJitter": - filtered_transforms.append(t) - - transform["transforms"] = filtered_transforms - transform_inst = instantiate(transform) - - metadata_versions = cfg["metadata_versions"] - metadata_version = metadata_versions[embodiment] - - if "gr1" in embodiment: - embodiment_tag = EmbodimentTag.GR1_unified - elif "franka" in embodiment: - embodiment_tag = EmbodimentTag.FRANKA - elif "so101" in embodiment: - embodiment_tag = EmbodimentTag.NEW_EMBODIMENT - elif "so100" in embodiment: - embodiment_tag = EmbodimentTag.SO100 - elif "robocasa" in embodiment: - embodiment_tag = EmbodimentTag.ROBOCASA - elif "piper" in embodiment: - embodiment_tag = EmbodimentTag.NEW_EMBODIMENT - elif "piper double" in embodiment: - embodiment_tag = EmbodimentTag.NEW_EMBODIMENT - else: - raise ValueError(f"Unknown embodiment: {embodiment}") - - dataset = LeRobotSingleDataset( - dataset_path=validation_dataset_path, - modality_configs=modality_configs, - # metadata_version=metadata_version, - transforms=transform_inst, - embodiment_tag=embodiment_tag, - ) - - return cfg, dataset, modality_configs - - -def collate_fn(features_list, device): - batch_dict = {} - keys = features_list[0].keys() - for key in keys: - if key in ["images", "view_ids"]: - vals = [f[key] for f in features_list] - batch_dict[key] = torch.as_tensor(np.concatenate(vals), device=device) - else: - vals = [f[key] for f in features_list] - batch_dict[key] = torch.as_tensor(np.stack(vals), device=device) - - return batch_dict - -def get_step_data(dataset, trajectory_id, base_index): - data = {} - dataset.curr_traj_data = dataset.get_trajectory_data(trajectory_id) - # Get the data for all modalities - for modality in dataset.modality_keys: - # Get the data corresponding to each key in the modality - for key in dataset.modality_keys[modality]: - if modality == "video": - pass - elif modality == "state" or modality == "action": - data[key] = dataset.get_state_or_action(trajectory_id, modality, key, base_index) - elif modality == "language": - data[key] = dataset.get_language(trajectory_id, key, base_index) - - return data - -def save_trajectory_data(trajectory_data, dataset, trajectory_id, output_dir): - chunk_index = dataset.get_episode_chunk(trajectory_id) - chunk_dir = f"chunk-{chunk_index:03d}" - os.makedirs(os.path.join(output_dir, "data", chunk_dir), exist_ok=True) - - episode_id = f"episode_{int(trajectory_id):06d}" - output_file_path = os.path.join(output_dir, "data", chunk_dir, f"{episode_id}.parquet") - trajectory_data.to_parquet(output_file_path) - -## -# Worker function, now handles either validation or writing based on a boolean flag. -## -def worker_func( - gpu_id: int, - traj_id_list: list, - checkpoint_path: str, - validation_dataset_path: str, - output_dir: str, - batch_size: int, - num_workers: int = 1, - video_indices=None, - dataset=None, - modality_configs=None, -): - """This function runs in a separate process on GPU `gpu_id` and processes all `traj_id_list`.""" - # Load model on this GPU - # from_pretrained should handle both local paths and HF model IDs - print(f"Loading model from: {checkpoint_path}") - model = IDM.from_pretrained(checkpoint_path) - - model.requires_grad_(False) - model.eval() - device = torch.device(f"cuda:{gpu_id}") - model.to(device) - - - for tid in tqdm(traj_id_list, desc=f"GPU {gpu_id}", position=gpu_id): - action_dict = {} - traj_data = dataset.get_trajectory_data(tid) - length = len(traj_data) - - - all_features = [] - - # Create a simple prefetching mechanism with a thread pool - from concurrent.futures import ThreadPoolExecutor - - def load_and_transform_step(step_idx, video_data): - import torch - from einops import rearrange - - step_data = get_step_data(dataset, tid, step_idx) - timestamp = dataset.curr_traj_data["timestamp"].to_numpy() - for key in video_data: - frames, whole_indices = video_data[key] - step_indices = dataset.delta_indices[key] + step_idx - step_indices = np.maximum(step_indices, 0) - step_indices = np.minimum(step_indices, dataset.trajectory_lengths[tid] - 1) - indices = np.array([np.where(np.isclose(whole_indices, val))[0][0] for val in timestamp[step_indices]]) - step_data[key] = frames[indices] - - output = dataset.transforms(step_data) - return output - - - video_data = {} - for key in dataset.modality_keys["video"]: - video_path = dataset.get_video_path(tid, key.replace("video.", "")) - video_backend = dataset.video_backend - video_backend_kwargs = dataset.video_backend_kwargs - frames, whole_indices = get_all_frames_and_timestamps(video_path.as_posix(), video_backend, video_backend_kwargs) - video_data[key] = (frames, whole_indices) - - - with ThreadPoolExecutor(max_workers=num_workers) as executor: - futures = [] - for step_idx in range(length): - future = executor.submit(load_and_transform_step, step_idx, video_data) - futures.append(future) - - all_features = [] - for future in as_completed(futures): - all_features.append(future.result()) - - - for start_idx in range(0, length, batch_size): - end_idx = min(start_idx + batch_size, length) - step_ids = list(range(start_idx, end_idx)) - - batch_features = all_features[start_idx:end_idx] - batch_dict = collate_fn(batch_features, device) - - with torch.no_grad(): - out = model.get_action(batch_dict) - - pred_actions = out["action_pred"].cpu() - pred_actions = dataset.transforms.unapply(Batch(action=pred_actions)) - - # Load modality.json to get the proper structure - modality_json_path = os.path.join(validation_dataset_path, 'meta', 'modality.json') - - with open(modality_json_path, 'r') as f: - modality_config = json.load(f) - - # Get action part configurations - action_parts = modality_config.get('action', {}) - - # Calculate total action dimension - total_dim = 0 - for part, indices in action_parts.items(): - total_dim = max(total_dim, indices.get('end', 0)) - - # Check if we have an action horizon dimension - has_horizon = False - action_horizon = 1 - sample_key = list(pred_actions.keys())[0] - if len(pred_actions[sample_key].shape) == 3: - has_horizon = True - batch_size, action_horizon, _ = pred_actions[sample_key].shape - else: - batch_size = pred_actions[sample_key].shape[0] - - if has_horizon: - final_actions = np.zeros((batch_size, action_horizon, total_dim)) - - # Fill in the actions for parts that exist in pred_actions - for part, indices in action_parts.items(): - action_key = f'action.{part}' - start_idx = indices.get('start', 0) - end_idx = indices.get('end', 0) - - if action_key in pred_actions: - final_actions[:, :, start_idx:end_idx] = pred_actions[action_key] - else: - final_actions = np.zeros((batch_size, total_dim)) - - # Fill in the actions for parts that exist in pred_actions - for part, indices in action_parts.items(): - action_key = f'action.{part}' - start_idx = indices.get('start', 0) - end_idx = indices.get('end', 0) - - if action_key in pred_actions: - final_actions[:, start_idx:end_idx] = pred_actions[action_key] - - pred_actions = final_actions - - # Not validating => we do the usual writing - for i, s in enumerate(step_ids): - for j in range(action_horizon): - if s+j >= length: - break - if s+j not in action_dict: - action_dict[s+j] = [] - if has_horizon: - action_dict[s+j].append(pred_actions[i, j].flatten()) - else: - action_dict[s+j].append(pred_actions[i].flatten()) - - for s in action_dict: - mean_action = np.mean(action_dict[s], axis=0) - traj_data.at[s, "action"] = mean_action - - save_trajectory_data(traj_data, dataset, tid, output_dir) - - del model - if torch.cuda.is_available(): - torch.cuda.empty_cache() - -def validate_checkpoint( - checkpoint_path, - validation_dataset_path, - output_dir=None, - num_gpus=8, - batch_size=16, - max_episodes=None, - num_workers=1, - video_indices=None, -): - device_count = torch.cuda.device_count() - print(f"Found {device_count} GPUs available.") - if device_count < num_gpus: - print( - f"WARNING: You requested num_gpus={num_gpus} but only {device_count} GPUs are visible." - ) - num_gpus = device_count - - # If not validation, create the output directory for saving predicted actions - if output_dir: - os.makedirs(output_dir, exist_ok=True) - print(f"Will save predicted actions to: {output_dir}") - - _, dataset, modality_configs = load_dataset_and_config(checkpoint_path, validation_dataset_path, video_indices) - - dataset.transforms.eval() - - trajectory_ids = dataset.trajectory_ids - if max_episodes is not None and max_episodes > 0: - trajectory_ids = trajectory_ids[:max_episodes] - - print(f"Processing {len(trajectory_ids)} trajectories total.") - - # Split trajectories among GPUs - chunk_size = (len(trajectory_ids) + num_gpus - 1) // num_gpus - - tasks_by_gpu = {} - for i in range(num_gpus): - start = i * chunk_size - end = min(start + chunk_size, len(trajectory_ids)) - if start >= end: - break - gpu_traj_list = trajectory_ids[start:end] - tasks_by_gpu[i] = gpu_traj_list - - try: - with ProcessPoolExecutor(max_workers=num_gpus) as executor: - futures = [] - for gpu_id, gpu_traj_list in tasks_by_gpu.items(): - future = executor.submit( - worker_func, - gpu_id, - gpu_traj_list, - checkpoint_path, - validation_dataset_path, - output_dir, - batch_size, - num_workers, - video_indices, - dataset, - modality_configs, - ) - futures.append(future) - - # Wait for all futures to complete or handle interruption - for future in as_completed(futures): - try: - # Get the result to catch any exceptions - future.result() - except Exception as e: - print(f"Error in worker process: {e}") - import traceback - traceback.print_exc() - except KeyboardInterrupt: - print("\nProcess interrupted by user. Cleaning up...") - # The context manager will handle cancellation of pending futures - except Exception as e: - print(f"Error in main process: {e}") - import traceback - traceback.print_exc() - finally: - # Force cleanup of CUDA memory - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - # If not validating, optionally copy the meta & videos folders - if output_dir: - # Copy the metadata directory - meta_src = os.path.join(validation_dataset_path, "meta") - meta_dst = os.path.join(output_dir, "meta") - if os.path.exists(meta_src): - import shutil - - if not os.path.exists(meta_dst): - shutil.copytree(meta_src, meta_dst) - print(f"Copied metadata to: {meta_dst}") - - tasks_path = os.path.join(meta_dst, "tasks.jsonl") - if os.path.exists(tasks_path): - # Read the existing tasks - tasks = [] - with open(tasks_path, "r") as f: - for line in f: - tasks.append(json.loads(line)) - - # Update tasks with prefix if not already present - updated_tasks = [] - for task in tasks: - if "task" in task and not task["task"].startswith(""): - task["task"] = f"{task['task']}" - updated_tasks.append(task) - - # Write the updated tasks back to the file - with open(tasks_path, "w") as f: - for task in updated_tasks: - f.write(json.dumps(task) + "\n") - - print("Updated tasks.jsonl with prefix") - - - # Copy the videos directory if it exists - videos_src = os.path.join(validation_dataset_path, "videos") - videos_dst = os.path.join(output_dir, "videos") - if os.path.exists(videos_src): - import shutil - - if not os.path.exists(videos_dst): - shutil.copytree(videos_src, videos_dst) - print(f"Copied videos to: {videos_dst}") - - return {"message": "Dataset with predicted actions written to disk"} - - -# -# Main -# -if __name__ == "__main__": - # If you're doing spawn/forkserver across multiple processes: - mp.set_start_method("spawn", force=True) - - parser = argparse.ArgumentParser( - description="Validate a checkpoint on a dataset with multi-GPU parallelism" - ) - parser.add_argument("--checkpoint", type=str, required=True, help="Path to the checkpoint") - parser.add_argument("--dataset", type=str, required=True, help="Path to the validation dataset") - parser.add_argument( - "--output_dir", type=str, help="Path to save the output dataset with predicted actions" - ) - parser.add_argument("--num_gpus", type=int, default=8, help="Number of GPUs to use in parallel") - parser.add_argument("--batch_size", type=int, default=32, help="Batch size for inference steps") - parser.add_argument( - "--max_episodes", type=int, default=None, help="Maximum number of episodes to validate" - ) - parser.add_argument( - "--num_workers", - type=int, - default=16, - help="Number of worker threads for data loading per GPU", - ) - parser.add_argument( - "--video_indices", - type=str, - default=None, - help="Video frame indices to use for inference", - ) - - args = parser.parse_args() - - result = validate_checkpoint( - checkpoint_path=args.checkpoint, - validation_dataset_path=args.dataset, - output_dir=args.output_dir, - num_gpus=args.num_gpus, - batch_size=args.batch_size, - max_episodes=args.max_episodes, - num_workers=args.num_workers, - video_indices=args.video_indices, - ) diff --git a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/lerobot_to_step2_format.py b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/lerobot_to_step2_format.py deleted file mode 100644 index 33ffe696..00000000 --- a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/lerobot_to_step2_format.py +++ /dev/null @@ -1,161 +0,0 @@ -#!/usr/bin/env python3 -""" -Convert LeRobot format dataset to convert_directory output format. - -Input format (LeRobot): - dataset_dir/ - ├── videos/chunk-000/observation.images.front/ - │ ├── episode_000000.mp4 - │ ├── episode_000001.mp4 - │ └── ... - └── meta/ - ├── tasks.jsonl (task definitions) - └── episodes.jsonl (episode -> task mapping) - -Output format (after convert_directory): - output_dir/ - ├── Lift_the_red_cube_up/ - │ ├── 0.mp4 - │ ├── 1.mp4 - │ └── ... - └── Another_task/ - ├── 0.mp4 - └── ... - -Usage: - python lerobot_to_step2_format.py \ - --input_dir /path/to/lerobot/dataset \ - --output_dir /path/to/output \ - --video_key observation.images.front -""" - -import argparse -import json -import shutil -from pathlib import Path -from collections import defaultdict - - -def load_tasks(tasks_file: Path) -> dict: - """Load task definitions. Returns {task_index: task_description}""" - tasks = {} - with open(tasks_file, 'r') as f: - for line in f: - data = json.loads(line.strip()) - tasks[data['task_index']] = data['task'] - return tasks - - -def load_episodes(episodes_file: Path) -> list: - """Load episode information. Returns list of {episode_index, tasks, length}""" - episodes = [] - with open(episodes_file, 'r') as f: - for line in f: - data = json.loads(line.strip()) - episodes.append(data) - return episodes - - -def sanitize_task_name(task: str) -> str: - """ - Convert task description to valid directory name. - Example: "Lift the red cube up." -> "Lift_the_red_cube_up" - """ - task = task.replace('.', '').replace(',', '').replace('!', '').replace('?', '') - task = task.replace(' ', '_') - task = '_'.join(filter(None, task.split('_'))) - return task - - -def convert_lerobot_to_step2( - input_dir: Path, - output_dir: Path, - video_key: str = "observation.images.front", - chunk_name: str = "chunk-000" -): - """Convert LeRobot format to convert_directory output format.""" - - # Load tasks and episodes - tasks = load_tasks(input_dir / "meta" / "tasks.jsonl") - episodes = load_episodes(input_dir / "meta" / "episodes.jsonl") - - print(f"Loaded {len(tasks)} tasks") - print(f"Loaded {len(episodes)} episodes") - - video_dir = input_dir / "videos" / chunk_name / video_key - - # Group episodes by task - task_episodes = defaultdict(list) - for episode in episodes: - episode_idx = episode['episode_index'] - episode_tasks = episode.get('tasks', []) - if episode_tasks: - task = episode_tasks[0] - task_episodes[task].append(episode_idx) - - print(f"\nStarting conversion...") - - total_copied = 0 - for task, episode_indices in task_episodes.items(): - # Create task directory - task_dir_name = sanitize_task_name(task) - task_dir = output_dir / task_dir_name - task_dir.mkdir(parents=True, exist_ok=True) - - print(f"\nTask: {task}") - print(f" Directory: {task_dir_name}") - print(f" Episodes: {len(episode_indices)}") - - # Copy video files - for idx, episode_idx in enumerate(sorted(episode_indices)): - src_video = video_dir / f"episode_{episode_idx:06d}.mp4" - dst_video = task_dir / f"{idx}.mp4" - shutil.copy2(src_video, dst_video) - total_copied += 1 - - print(f"\nConversion complete!") - print(f" Total videos copied: {total_copied}") - print(f" Output directory: {output_dir}") - - -def main(): - parser = argparse.ArgumentParser( - description="Convert LeRobot format to convert_directory output format" - ) - parser.add_argument( - "--input_dir", - type=str, - required=True, - help="LeRobot dataset directory" - ) - parser.add_argument( - "--output_dir", - type=str, - required=True, - help="Output directory" - ) - parser.add_argument( - "--video_key", - type=str, - default="observation.images.front", - help="Video key name (default: observation.images.front)" - ) - parser.add_argument( - "--chunk_name", - type=str, - default="chunk-000", - help="Chunk name (default: chunk-000)" - ) - - args = parser.parse_args() - - convert_lerobot_to_step2( - Path(args.input_dir), - Path(args.output_dir), - args.video_key, - args.chunk_name - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/preprocess_video.py b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/preprocess_video.py deleted file mode 100644 index 6ddd8755..00000000 --- a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/preprocess_video.py +++ /dev/null @@ -1,376 +0,0 @@ -import os -import cv2 -import numpy as np -from pathlib import Path -import multiprocessing as mp -from tqdm import tqdm -import shutil -import argparse -import concurrent.futures -from functools import partial -import decord -import imageio - - -def custom_crop_pad_resize_gr1(img, target_size=(256, 256)): - """ - Custom crop, pad, and resize operation that maintains the aspect ratio. - - Args: - img: Input frame - original_width: Original video width - original_height: Original video height - target_size: Target resolution (width, height) - - Returns: - Processed frame at target_size resolution - """ - # For 832x480 videos, adjust the crop parameters proportionally from the 1280x800 example - # Original crop for 1280x800: (310, 770, 110, 1130) - (top, bottom, left, right) - original_height, original_width = img.shape[:2] - - - # Calculate proportional crop values - crop_top_ratio = 310 / 800 - crop_bottom_ratio = 770 / 800 - crop_left_ratio = 110 / 1280 - crop_right_ratio = 1130 / 1280 - - # Apply ratios to the current dimensions - crop_top = int(original_height * crop_top_ratio) - crop_bottom = int(original_height * crop_bottom_ratio) - crop_left = int(original_width * crop_left_ratio) - crop_right = int(original_width * crop_right_ratio) - - # Ensure crop boundaries are within image dimensions - crop_top = max(0, min(crop_top, original_height - 1)) - crop_bottom = max(crop_top + 1, min(crop_bottom, original_height)) - crop_left = max(0, min(crop_left, original_width - 1)) - crop_right = max(crop_left + 1, min(crop_right, original_width)) - - # Crop the image - img_cropped = img[crop_top:crop_bottom, crop_left:crop_right] - - # Calculate intermediate size while maintaining aspect ratio - cropped_height, cropped_width = img_cropped.shape[:2] - aspect_ratio = cropped_width / cropped_height - - # Resize to intermediate size (similar to 720x480 in the example) - intermediate_height = 480 - intermediate_width = 720 - img_resized = cv2.resize(img_cropped, (intermediate_width, intermediate_height), cv2.INTER_AREA) - - # Pad to make it square - if intermediate_width > intermediate_height: - # Width is larger, pad height - height_pad = (intermediate_width - intermediate_height) // 2 - img_pad = np.pad( - img_resized, ((height_pad, height_pad), (0, 0), (0, 0)), mode="constant", constant_values=0 - ) - else: - # Height is larger, pad width - width_pad = (intermediate_height - intermediate_width) // 2 - img_pad = np.pad( - img_resized, ((0, 0), (width_pad, width_pad), (0, 0)), mode="constant", constant_values=0 - ) - - # Final resize to target size - final_img = cv2.resize(img_pad, target_size, cv2.INTER_AREA) - - return final_img - -def resize_with_padding(img, ratio=1.0, target_size=(256, 256)): - # Original aspect ratio is 1280:800 (or 16:10) - target_ratio = ratio - - h, w = img.shape[:2] - current_ratio = w / h - if target_ratio >= 1: - # Width is the limiting factor - new_w = target_size[0] - new_h = int(new_w / target_ratio) - resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR) - - # Add padding to height - pad_top = (target_size[1] - new_h) // 2 - pad_bottom = target_size[1] - new_h - pad_top - padded = cv2.copyMakeBorder(resized, pad_top, pad_bottom, 0, 0, cv2.BORDER_CONSTANT, value=[0, 0, 0]) - else: - # Height is the limiting factor - new_h = target_size[1] - new_w = int(new_h * target_ratio) - resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR) - - # Add padding to width - pad_left = (target_size[0] - new_w) // 2 - pad_right = target_size[0] - new_w - pad_left - padded = cv2.copyMakeBorder(resized, 0, 0, pad_left, pad_right, cv2.BORDER_CONSTANT, value=[0, 0, 0]) - - return padded - -def extract_subimages_franka(frame, original_width, original_height): - h, w = frame.shape[:2] # h=480, w=832 - - # Calculate dimensions for even division - half_width = w // 2 - half_height = h // 2 - - # Extract subimages - image_side_0 = frame[:half_height, :half_width] # Top-left (240x416) - image_side_1 = frame[:half_height, half_width:] # Top-right (240x416) - wrist_image = frame[half_height:, :half_width] # Bottom-left (240x416) - - - image_side_0 = cv2.resize(image_side_0, (original_width, original_height), interpolation=cv2.INTER_LINEAR) - image_side_1 = cv2.resize(image_side_1, (original_width, original_height), interpolation=cv2.INTER_LINEAR) - wrist_image = cv2.resize(wrist_image, (original_width, original_height), interpolation=cv2.INTER_LINEAR) - - return image_side_0, image_side_1, wrist_image - - -def extract_subimages(frame, ratio): - """Extract subimages from a frame and resize to 256x256 while preserving aspect ratio with padding.""" - h, w = frame.shape[:2] # h=480, w=832 - - # Calculate dimensions for even division - half_width = w // 2 # 416 - half_height = h // 2 # 240 - - # Extract subimages - image_side_0 = frame[:half_height, :half_width] # Top-left (240x416) - image_side_1 = frame[:half_height, half_width:] # Top-right (240x416) - wrist_image = frame[half_height:, :half_width] # Bottom-left (240x416) - - image_side_0 = resize_with_padding(image_side_0, ratio) - image_side_1 = resize_with_padding(image_side_1, ratio) - wrist_image = resize_with_padding(wrist_image, ratio) - - return image_side_0, image_side_1, wrist_image - -def process_batch_frames(frames, output_videos, src_path, dataset, original_width, original_height): - """Process a batch of frames.""" - ratio = original_width / original_height - for frame in frames: - # Extract subimages - if dataset == 'robocasa': - image_side_0, image_side_1, wrist_image = extract_subimages(frame, ratio) - output_videos['observation.images.left_view'].append_data(image_side_0) - output_videos['observation.images.right_view'].append_data(image_side_1) - output_videos['observation.images.wrist_view'].append_data(wrist_image) - elif dataset == 'gr1': - image = custom_crop_pad_resize_gr1(frame) - output_videos['observation.images.ego_view'].append_data(image) - elif dataset == 'franka': - image_side_0, image_side_1, wrist_image = extract_subimages(frame, ratio) - output_videos['observation.images.exterior_image_1_left_pad_res256_freq15'].append_data(image_side_0) - output_videos['observation.images.exterior_image_2_left_pad_res256_freq15'].append_data(image_side_1) - output_videos['observation.images.wrist_image_left_pad_res256_freq15'].append_data(wrist_image) - elif dataset == 'so100': - image = resize_with_padding(frame, ratio) - output_videos['observation.images.front'].append_data(image) - elif dataset == 'so101': - image = resize_with_padding(frame, ratio) - output_videos['observation.images.front'].append_data(image) - elif dataset == 'piper': - # Piper: 720x720 -> resize to 256x256 (square, no padding needed) - image = cv2.resize(frame, (256, 256), interpolation=cv2.INTER_LINEAR) - output_videos['observation.images.agentview_all'].append_data(image) - elif dataset == 'piper double': - image = cv2.resize(frame, (256, 256), interpolation=cv2.INTER_LINEAR) - output_videos['observation.images.agentview_all'].append_data(image) - else: - raise ValueError(f"Unknown task: {src_path}") - -def process_video(args): - """Process a single video file.""" - src_path, dst_dir, video_name, dataset, original_width, original_height = args - - # Create output directories if they don't exist - if dataset == 'robocasa': - output_dirs = { - 'observation.images.left_view': os.path.join(dst_dir, 'videos', 'observation.images.left_view'), - 'observation.images.right_view': os.path.join(dst_dir, 'videos', 'observation.images.right_view'), - 'observation.images.wrist_view': os.path.join(dst_dir, 'videos', 'observation.images.wrist_view'), - } - elif dataset == 'gr1': - output_dirs = { - 'observation.images.ego_view': os.path.join(dst_dir, 'videos', 'observation.images.ego_view') - } - elif dataset == 'franka': - output_dirs = { - 'observation.images.exterior_image_1_left_pad_res256_freq15': os.path.join(dst_dir, 'videos', 'observation.images.exterior_image_1_left_pad_res256_freq15'), - 'observation.images.exterior_image_2_left_pad_res256_freq15': os.path.join(dst_dir, 'videos', 'observation.images.exterior_image_2_left_pad_res256_freq15'), - 'observation.images.wrist_image_left_pad_res256_freq15': os.path.join(dst_dir, 'videos', 'observation.images.wrist_image_left_pad_res256_freq15'), - } - elif dataset == 'so100': - output_dirs = { - 'observation.images.front': os.path.join(dst_dir, 'videos', 'observation.images.front'), - } - elif dataset == 'so101': - output_dirs = { - 'observation.images.front': os.path.join(dst_dir, 'videos', 'observation.images.front'), - } - elif dataset == 'piper': - output_dirs = { - 'observation.images.agentview_all': os.path.join(dst_dir, 'videos', 'observation.images.agentview_all'), - } - elif dataset == 'piper double': - output_dirs = { - 'observation.images.agentview_all': os.path.join(dst_dir, 'videos', 'observation.images.agentview_all'), - } - - for dir_path in output_dirs.values(): - os.makedirs(dir_path, exist_ok=True) - - # Open the source video - vr = decord.VideoReader(src_path) - - # Get video properties - fps = vr.get_avg_fps() - frame_count = len(vr) - - # Create VideoWriter objects for each subimage - output_videos = {} - for name, dir_path in output_dirs.items(): - output_path = os.path.join(dir_path, f"{video_name}.mp4") - output_videos[name] = imageio.get_writer(output_path, fps=fps) - - - # Process frames in batches - batch_size = 32 # Adjust based on available memory - frames_batch = [] - - # Process each frame with progress bar - pbar = tqdm(total=frame_count, desc=f"Processing {video_name}", leave=False) - - for frame in vr: - frames_batch.append(frame.asnumpy()) - - # Process batch when it reaches the desired size - if len(frames_batch) >= batch_size: - process_batch_frames(frames_batch, output_videos, src_path, dataset, original_width, original_height) - frames_batch = [] - pbar.update(batch_size) - - # Process remaining frames - if frames_batch: - process_batch_frames(frames_batch, output_videos, src_path, dataset, original_width, original_height) - pbar.update(len(frames_batch)) - - # Close progress bar - pbar.close() - - # Release resources - for writer in output_videos.values(): - writer.close() - - -def copy_labels(src_dir, dst_dir): - """Copy label files from source to destination.""" - src_labels_dir = os.path.join(src_dir, 'labels') - dst_labels_dir = os.path.join(dst_dir, 'labels') - - if os.path.exists(src_labels_dir): - os.makedirs(dst_labels_dir, exist_ok=True) - for label_file in os.listdir(src_labels_dir): - if label_file.endswith('.txt'): - shutil.copy2( - os.path.join(src_labels_dir, label_file), - os.path.join(dst_labels_dir, label_file) - ) - -def process_subdirectory(subdir, src_dir, dst_dir, num_workers, max_videos=None, dataset=None, original_width=None, original_height=None): - """Process a single subdirectory.""" - print(f"Processing subdirectory: {src_dir}, {subdir}") - src_subdir = os.path.join(src_dir, subdir) - dst_subdir = os.path.join(dst_dir, subdir) - - # Copy label files - copy_labels(src_subdir, dst_subdir) - - # Process videos - src_videos_dir = os.path.join(src_subdir, 'videos') - if os.path.exists(src_videos_dir): - video_files = [f for f in os.listdir(src_videos_dir) if f.endswith('.mp4')] - - # Sort files numerically to ensure consistent ordering (1.mp4, 2.mp4, ..., 10.mp4, not 1.mp4, 10.mp4, 2.mp4) - def sort_key(filename): - stem = os.path.splitext(filename)[0] - try: - return int(stem) - except ValueError: - return float('inf') - video_files.sort(key=sort_key) - - # Limit number of videos if max_videos is specified - if max_videos is not None: - video_files = video_files[:max_videos] - print(f"Processing limited set of {len(video_files)} videos in {subdir}") - - # Prepare arguments for multiprocessing - args_list = [ - (os.path.join(src_videos_dir, video_file), dst_subdir, os.path.splitext(video_file)[0], dataset, original_width, original_height) - for video_file in video_files - ] - - # Process videos in parallel - with mp.Pool(num_workers) as pool: - list(tqdm(pool.imap(process_video, args_list), total=len(args_list), desc=f"Processing {subdir}")) - -def process_directory(src_dir, dst_dir, num_workers=None, num_subdirs_parallel=1, max_videos=None, dataset=None, original_width=None, original_height=None, recursive=False): - """Process all videos in the source directory.""" - if num_workers is None: - num_workers = max(1, mp.cpu_count() // 2) # Leave some cores for system - - # Create destination directory structure - os.makedirs(dst_dir, exist_ok=True) - - - if recursive: - # Get all subdirectories in the source directory - subdirs = [d for d in os.listdir(src_dir) if os.path.isdir(os.path.join(src_dir, d))] - else: - # Get all subdirectories in the source directory - subdirs = [''] - - # Process subdirectories in parallel - with concurrent.futures.ThreadPoolExecutor(max_workers=num_subdirs_parallel) as executor: - process_subdir_fn = partial(process_subdirectory, - src_dir=src_dir, - dst_dir=dst_dir, - num_workers=num_workers // num_subdirs_parallel, - max_videos=max_videos, - dataset=dataset, - original_width=original_width, - original_height=original_height) - list(tqdm(executor.map(process_subdir_fn, subdirs), total=len(subdirs), desc="Processing subdirectories")) - -def main(): - parser = argparse.ArgumentParser(description='Split videos into subimages and save them.') - parser.add_argument('--src_dir', type=str, default='/mnt/amlfs-01/home/seonghyeony/data/dreams_robocasa_70K_0228', - help='Source directory containing videos') - parser.add_argument('--dst_dir', type=str, default='/mnt/amlfs-01/home/seonghyeony/data/dreams_robocasa_70K_0228_split', - help='Destination directory for processed videos') - parser.add_argument('--workers', type=int, default=8, - help='Number of worker processes (default: half of CPU cores)') - parser.add_argument('--parallel_subdirs', type=int, default=1, - help='Number of subdirectories to process in parallel') - parser.add_argument('--batch_size', type=int, default=32, - help='Number of frames to process in a batch') - parser.add_argument('--max_videos', type=int, default=None, - help='Maximum number of videos to process per subdirectory (for debugging)') - parser.add_argument('--dataset', type=str, default='robocasa', - help='Dataset name', choices=['robocasa', 'gr1', 'franka', 'so100', 'so101', 'piper', 'piper double']) - parser.add_argument("--recursive", action="store_true", help="Process subdirectories recursively, maintaining directory structure") - parser.add_argument("--original_width", type=int, default=1280, help="Original width of the video") - parser.add_argument("--original_height", type=int, default=800, help="Original height of the video") - - args = parser.parse_args() - # Set OpenCV thread optimization - cv2.setNumThreads(1) # Disable OpenCV's internal multithreading to avoid conflicts - - process_directory(args.src_dir, args.dst_dir, args.workers, args.parallel_subdirs, - max_videos=args.max_videos, dataset=args.dataset, original_width=args.original_width, original_height=args.original_height, recursive=args.recursive) - -if __name__ == "__main__": - main() diff --git a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/raw_to_lerobot.py b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/raw_to_lerobot.py deleted file mode 100644 index 07d3d1ca..00000000 --- a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/raw_to_lerobot.py +++ /dev/null @@ -1,568 +0,0 @@ -import json -import shutil -from pathlib import Path -from typing import Dict, List, Any, Tuple -import concurrent.futures -import os - -import numpy as np -import pandas as pd -from tqdm import tqdm -import argparse - -import subprocess - -import multiprocessing -from multiprocessing import Pool -import math - - -CHUNKS_SIZE = 1000 -DATA_PATH = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet" -VIDEO_PATH = "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4" - - -def get_video_metadata(video_path): - """Get video metadata using ffprobe.""" - cmd = [ - "ffprobe", "-v", "error", - "-select_streams", "v:0", - "-show_entries", "stream=height,width,codec_name,pix_fmt,r_frame_rate", - "-of", "json", video_path, - ] - - try: - output = subprocess.check_output(cmd).decode("utf-8") - probe_data = json.loads(output) - stream = probe_data["streams"][0] - - # Parse frame rate - num, den = map(int, stream["r_frame_rate"].split("/")) - fps = num / den - - return { - "dtype": "video", - "shape": [stream["height"], stream["width"], 3], - "names": ["height", "width", "channel"], - "video_info": { - "video.fps": fps, - "video.codec": stream["codec_name"], - "video.pix_fmt": stream["pix_fmt"], - "video.is_depth_map": False, - }, - } - except Exception as e: - print(f"Error getting video metadata: {e}") - return None - -def dump_jsonl(data: List[Dict[str, Any]], path: Path) -> None: - """Dump list of dictionaries as JSONL file.""" - with open(path, 'w') as f: - for item in data: - json_str = json.dumps(item) - f.write(json_str + '\n') - -def json_dump(data: Dict[str, Any], path: Path, indent: int = 4) -> None: - """Dump dictionary as JSON file.""" - with open(path, 'w') as f: - json.dump(data, f, indent=indent) - - -def process_video_chunk(args): - """Process a chunk of videos in parallel.""" - video_files, labels_dir, output_dir, cosmos_predict2, data_type, videos_dir, video_key = args - results = [] - - for video_file in video_files: - video_id = video_file.stem - label_file = labels_dir / f"{video_id}.txt" - - # Read annotation - with open(label_file, "r") as f: - annotation = f.read().strip() - - # Get video frame count (if not in cosmos_predict2 mode) - if cosmos_predict2: - frame_count = 93 # Fixed frame count for cosmos_predict2 - else: - try: - cmd = ["ffprobe", "-v", "error", "-select_streams", "v:0", - "-show_entries", "stream=duration,r_frame_rate", "-of", "json", str(videos_dir / video_key / f"{video_file}.mp4")] - output = subprocess.check_output(cmd).decode() - probe_data = json.loads(output) - - if "streams" in probe_data and len(probe_data["streams"]) > 0: - stream = probe_data["streams"][0] - - # Check if duration is available - if "duration" in stream: - duration = float(stream["duration"]) - fps_str = stream.get("r_frame_rate", "30/1") - try: - num, den = map(int, fps_str.split("/")) - fps = num / den - frame_count = int(duration * fps) - except Exception: - # Default to 30fps if fps parsing fails - frame_count = int(duration * 30) - else: - # If no duration, use a default frame count - print(f"Warning: Could not determine duration for {video_file}, using default frame count") - frame_count = 93 - else: - # No streams found - print(f"Warning: No valid video streams found in {video_file}, using default frame count") - frame_count = 93 - - except Exception as e: - print(f"Error getting frame count for {video_file}: {e}") - # Default frame count if all methods fail - frame_count = 93 - - # Ensure frame count is at least 1 - frame_count = max(1, frame_count) - - results.append((video_id, annotation, frame_count)) - - return results - - -def copy_videos_parallel(video_copy_tasks, max_workers=16): - """Copy multiple videos in parallel using ThreadPoolExecutor.""" - with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = [] - for source, dest in video_copy_tasks: - futures.append(executor.submit(shutil.copy2, source, dest)) - - # Wait for all copy operations to complete - for future in concurrent.futures.as_completed(futures): - try: - future.result() - except Exception as e: - print(f"Error copying file: {e}") - -def process_folder(args): - """Process a single input folder.""" - folder_path, output_base_dir, annotation_source, fps, max_videos, num_workers_per_folder, cosmos_predict2, data_type, embodiment, video_key = args - - - - # Create folder-specific output directory - folder_name = folder_path.name - - output_dir = output_base_dir / f"{embodiment}.{folder_name}" - - # Process this folder - result = convert_raw_to_lerobot( - raw_dir=folder_path, - output_dir=output_dir, - annotation_source=annotation_source, - fps=fps, - max_videos=max_videos, - num_workers=num_workers_per_folder, - cosmos_predict2=cosmos_predict2, - data_type=data_type, - video_key=video_key, - ) - - return folder_name, result - -def convert_raw_to_lerobot( - raw_dir: Path, - output_dir: Path, - annotation_source: str = "human", - fps: int = 8, - max_videos: int | None = None, - num_workers: int = None, - cosmos_predict2: bool = False, - data_type: str = "lapa", - video_key: str = None, -): - """Convert raw dataset to LeRobot format.""" - - - - # Setup directories - videos_dir = raw_dir / "videos" - labels_dir = raw_dir / "labels" - output_dir = Path(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - # Create metadata directory - meta_dir = output_dir / "meta" - meta_dir.mkdir(parents=True, exist_ok=True) - - # Check for the new folder structure - video_subfolders = [f for f in videos_dir.iterdir() if f.is_dir()] - if video_subfolders: - # New folder structure detected - print(f"Detected new folder structure with subfolders: {[f.name for f in video_subfolders]}") - - # Get all unique video IDs across all subfolders - all_video_ids = set() - for folder in video_subfolders: - video_files = folder.glob("*.mp4") - all_video_ids.update(video_file.stem for video_file in video_files) - - # Convert to sorted list with numeric sorting - # Try to sort as integers if all IDs are numeric, otherwise use string sort - try: - video_ids = sorted(list(all_video_ids), key=lambda x: int(x) if x.isdigit() else float('inf')) - except ValueError: - # Fallback to string sort if conversion fails - video_ids = sorted(list(all_video_ids)) - - # Create dummy video_files list with just the IDs - video_files = [Path(video_id) for video_id in video_ids] - else: - raise ValueError("No video subfolders found in the input directory") - - if max_videos is not None: - video_files = video_files[:max_videos] - print(f"Processing only first {max_videos} videos for debugging") - - print(f"Processing {len(video_files)} videos in {videos_dir}") - - # Setup multiprocessing - if num_workers is None: - num_workers = max(1, multiprocessing.cpu_count() // 2) # Use half of available cores by default - - # Split videos into chunks for parallel processing - chunk_size = math.ceil(len(video_files) / num_workers) - video_chunks = [video_files[i:i + chunk_size] for i in range(0, len(video_files), chunk_size)] - - # Prepare arguments for parallel processing - args_list = [ - (chunk, labels_dir, output_dir, cosmos_predict2, data_type, videos_dir, video_key) - for chunk in video_chunks - ] - - # Process videos in parallel - total_frames = 0 - annotation_to_index = {} - episodes_info = [] - - print(f"Processing {len(video_files)} videos in {raw_dir} using {num_workers} workers") - print(f"cosmos_predict2 mode: {cosmos_predict2} (fixed FPS=16, frames=93)" if cosmos_predict2 else "") - with Pool(num_workers) as pool: - all_results = list(tqdm( - pool.imap(process_video_chunk, args_list), - total=len(args_list), - desc=f"Processing {raw_dir.name}" - )) - - # Process results and create videos - print(f"Processing {len(all_results)} chunks of videos...") - - # Flatten all results and sort by video_id to ensure episode_index matches video order - all_video_results = [] - for chunk_results in all_results: - all_video_results.extend(chunk_results) - - # Sort by video_id numerically to ensure episode 0 -> video 1, episode 1 -> video 2, etc. - try: - all_video_results.sort(key=lambda x: int(x[0]) if x[0].isdigit() else float('inf')) - except (ValueError, IndexError): - # Fallback to string sort if conversion fails - all_video_results.sort(key=lambda x: x[0]) - - # We'll collect all video copy tasks here - all_video_copy_tasks = [] - - - for video_id, annotation, frame_count in tqdm( - all_video_results, - desc="Processing videos in order" - ): - if annotation not in annotation_to_index: - annotation_to_index[annotation] = len(annotation_to_index) - - episode_index = len(episodes_info) - - # If cosmos_predict2 mode is enabled, use fixed frame count - actual_frame_count = 93 if cosmos_predict2 else frame_count - actual_fps = fps # Use the actual fps parameter - - # Determine state/action dimensions based on data_type - if data_type == "piper" or "piper" in str(output_dir).lower(): - state_action_dim = 8 # Piper: 6 joints + 2 grippers - elif data_type == "so101" or "so101" in str(output_dir).lower(): - state_action_dim = 6 # SO101: 5 joints + 1 gripper - else: - state_action_dim = 44 # Default - - # Create episode data - episode_data = { - "observation.state": [np.zeros(state_action_dim, dtype=np.float32)] * actual_frame_count, - "action": [np.zeros(state_action_dim, dtype=np.float32)] * actual_frame_count, - "timestamp": [i/actual_fps for i in range(actual_frame_count)], - "episode_index": [episode_index] * actual_frame_count, - "index": np.arange(total_frames, total_frames + actual_frame_count), - "task_index": [annotation_to_index[annotation]] * actual_frame_count, - f"annotation.{annotation_source}": [[annotation_to_index[annotation]]] * actual_frame_count - } - - # Save episode data - episode_chunk = episode_index // CHUNKS_SIZE - data_path = DATA_PATH.format(episode_chunk=episode_chunk, episode_index=episode_index) - save_path = output_dir / data_path - save_path.parent.mkdir(parents=True, exist_ok=True) - - df = pd.DataFrame(episode_data) - df.to_parquet(save_path) - - - view_list = [] - for folder in video_subfolders: - source_video_path = folder / f"{video_id}.mp4" - if source_video_path.exists(): - view_list.append((folder.name, source_video_path)) - - for video_key, source_video_path in view_list: - # Get the LeRobot format key - - # Destination path in LeRobot format - video_save_path = output_dir / VIDEO_PATH.format( - episode_chunk=episode_chunk, - video_key=video_key, - episode_index=episode_index - ) - video_save_path.parent.mkdir(parents=True, exist_ok=True) - - # Add to copy tasks instead of copying immediately - all_video_copy_tasks.append((source_video_path, video_save_path)) - - # Update episodes info - episodes_info.append({ - "episode_index": episode_index, - "tasks": [annotation], - "length": actual_frame_count, - "video_id": video_id - }) - - total_frames += actual_frame_count - - # Now copy all videos in parallel - print(f"Copying {len(all_video_copy_tasks)} videos in parallel...") - copy_videos_parallel(all_video_copy_tasks, max_workers=min(32, num_workers)) - - # Generate metadata files - # 1. tasks.jsonl - tasks_path = meta_dir / "tasks.jsonl" - tasks = [{"task_index": idx, "task": task} for task, idx in annotation_to_index.items()] - dump_jsonl(tasks, tasks_path) - - # 2. episodes.jsonl - episodes_path = meta_dir / "episodes.jsonl" - dump_jsonl(episodes_info, episodes_path) - - # 3. info.json - # Find a sample video for metadata from each view - - # Determine robot-specific configuration - if data_type == "piper" or "piper" in str(output_dir).lower(): - robot_type = "Piper-Abs" - state_action_shape = (8,) - state_action_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", - "finger_joint_left", "finger_joint_right"] - elif data_type == "so101" or "so101" in str(output_dir).lower(): - robot_type = "SO101" - state_action_shape = (6,) - state_action_names = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", - "wrist_roll", "gripper"] - else: - robot_type = data_type - state_action_shape = (44,) - state_action_names = [f"motor_{i}" for i in range(44)] - - info = { - "robot_type": robot_type, - "total_episodes": len(video_files), - "total_frames": total_frames, - "total_tasks": len(annotation_to_index), - "total_videos": len(video_subfolders), - "chunks_size": CHUNKS_SIZE, - "total_chunks": (len(video_files) + CHUNKS_SIZE - 1) // CHUNKS_SIZE, - "fps": fps, # Use the actual fps parameter - "data_path": DATA_PATH, - "video_path": VIDEO_PATH, - "features": { - "observation.state": { - "dtype": "float32", - "shape": state_action_shape, - "names": state_action_names - }, - "action": { - "dtype": "float32", - "shape": state_action_shape, - "names": state_action_names - }, - f"annotation.{annotation_source}": { - "dtype": "int64", - "shape": (1,) - } - } - } - - - for folder in video_subfolders: - # Find the first video in this folder - sample_videos = list(folder.glob("*.mp4")) - if sample_videos: - view_key = folder.name - sample_path = sample_videos[0] - info["features"][view_key] = get_video_metadata(sample_path) - - - # Remove None values from features - info["features"] = {k: v for k, v in info["features"].items() if v is not None} - - info_path = meta_dir / "info.json" - json_dump(info, info_path, indent=4) - - return output_dir - -def process_multiple_folders( - input_base_dir: Path, - output_base_dir: Path, - annotation_source: str = "human", - fps: int = 8, - max_videos: int | None = None, - num_workers: int = None, - cosmos_predict2: bool = False, - data_type: str = "lapa", - embodiment: str = None -): - """Process multiple input folders in parallel.""" - # Get all subdirectories that contain videos/ and labels/ folders - input_folders = [] - for folder in sorted(input_base_dir.iterdir()): - if folder.is_dir() and (folder / "videos").exists() and (folder / "labels").exists(): - input_folders.append(folder) - - print(f"Found {len(input_folders)} folders to process") - - # Determine number of workers per folder - total_workers = multiprocessing.cpu_count() if num_workers is None else num_workers - workers_per_folder = max(1, total_workers) - - # Prepare arguments for parallel folder processing - args_list = [ - (folder, output_base_dir, annotation_source, fps, max_videos, workers_per_folder, cosmos_predict2, data_type, embodiment) - for folder in input_folders - ] - - # Process folders in parallel using ThreadPoolExecutor - # (ProcessPoolExecutor would create too many processes) - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: - futures = {executor.submit(process_folder, args): args[0] for args in args_list} - - # Process results as they complete - for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Processing folders"): - folder_path = futures[future] - try: - folder_name, result = future.result() - print(f"Completed processing folder: {folder_name}") - except Exception as e: - print(f"Error processing folder {folder_path.name}: {e}") - import traceback - traceback.print_exc() - -def main(): - parser = argparse.ArgumentParser(description="Convert raw dataset to LeRobot format") - parser.add_argument("--input_dir", type=str, required=True, help="Input directory containing multiple folders with videos/ and labels/") - parser.add_argument("--output_dir", type=str, required=True, help="Output directory for LeRobot dataset") - parser.add_argument("--fps", type=int, default=16, help="Video FPS") - parser.add_argument("--max_videos", type=int, default=None, help="Maximum number of videos to process per folder (for debugging)") - parser.add_argument("--num_workers", type=int, default=16, help="Total number of worker processes") - parser.add_argument("--cosmos_predict2", action="store_true", help="Process videos for cosmos_predict2 video models (fixed FPS=8, frames=81)") - parser.add_argument("--recursive", action="store_true", help="Process a single folder instead of multiple folders") - parser.add_argument("--data_type", type=str, default="dream", choices=["lapa", "dream"]) - parser.add_argument("--embodiment", type=str, default=None, help="Embodiment") - parser.add_argument("--video_key", type=str, default=None, help="Video key if cosmos_predict2 is false") - - args = parser.parse_args() - - if args.embodiment is None: - if 'robocasa' in args.output_dir: - args.embodiment = "robocasa_panda_omron" - elif 'gr1' in args.output_dir: - args.embodiment = "gr1_unified" - elif 'franka' in args.output_dir: - args.embodiment = "franka" - elif 'so101' in args.output_dir: - args.embodiment = "so101" - elif 'so100' in args.output_dir: - args.embodiment = "so100" - elif 'piper' in args.output_dir: - args.embodiment = "piper" - elif 'piper double' in args.output_dir: - args.embodiment = "piper double" - else: - raise ValueError(f"Unknown embodiment for {args.output_dir}")\ - - if args.embodiment == "robocasa_panda_omron": - args.annotation_source = "human.action.task_description" - elif args.embodiment == "gr1_unified": - args.annotation_source = "human.coarse_action" - elif args.embodiment == "franka": - args.annotation_source = "language.language_instruction" - elif args.embodiment == "so100": - args.annotation_source = "human.task_description" - elif args.embodiment == "so101": - args.annotation_source = "human.task_description" - elif args.embodiment == "piper": - args.annotation_source = "human.task_description" - elif args.embodiment == "piper double": - args.annotation_source = "human.task_description" - - if args.recursive: - # Process a single folder (original behavior) - process_multiple_folders( - input_base_dir=Path(args.input_dir), - output_base_dir=Path(args.output_dir), - annotation_source=args.annotation_source, - fps=args.fps, - max_videos=args.max_videos, - num_workers=args.num_workers, - cosmos_predict2=args.cosmos_predict2, - data_type=args.data_type, - embodiment=args.embodiment - ) - else: - convert_raw_to_lerobot( - raw_dir=Path(args.input_dir), - output_dir=Path(args.output_dir), - annotation_source=args.annotation_source, - fps=args.fps, - max_videos=args.max_videos, - num_workers=args.num_workers, - cosmos_predict2=args.cosmos_predict2, - data_type=args.data_type, - video_key=args.video_key - ) - - if args.embodiment == "gr1_unified": - source_dir = "IDM_dump/global_metadata/gr1" - elif args.embodiment == "robocasa_panda_omron": - source_dir = "IDM_dump/global_metadata/robocasa" - elif args.embodiment == "franka": - source_dir = "IDM_dump/global_metadata/franka" - elif args.embodiment == "so100": - source_dir = "IDM_dump/global_metadata/so100" - elif args.embodiment == "so101": - source_dir = "IDM_dump/global_metadata/so101" - elif args.embodiment == "piper": - source_dir = "IDM_dump/global_metadata/piper" - elif args.embodiment == "piper double": - source_dir = "IDM_dump/global_metadata/piper double" - - # copy modality.json - shutil.copy(source_dir + "/modality.json", args.output_dir + "/meta/modality.json") - - # # copy stats.json - shutil.copy(source_dir + "/stats.json", args.output_dir + "/meta/stats.json") - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/split_video_instruction.py b/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/split_video_instruction.py deleted file mode 100644 index 3620670e..00000000 --- a/scripts/tutorials/cosmos_tutorials/preprocess_leisaac/split_video_instruction.py +++ /dev/null @@ -1,101 +0,0 @@ -import os -import shutil -import re -import argparse - -def process_mp4_files(source_dir, output_dir, recursive=False, task_name=None): - """ - Process MP4 files from source_dir and its subdirectories, extract instructions from filenames, - and save them to output_dir with appropriate structure. - - Args: - source_dir: Directory containing MP4 files or subdirectories with MP4 files - output_dir: Directory to save processed data - recursive: If True, maintain the directory structure from source_dir in output_dir - task_name: Optional task name to use as instruction when filename is just a number - """ - if not recursive: - # Original behavior: process all MP4 files into a single output directory - labels_dir = os.path.join(output_dir, "labels") - videos_dir = os.path.join(output_dir, "videos") - - os.makedirs(labels_dir, exist_ok=True) - os.makedirs(videos_dir, exist_ok=True) - - # Find all MP4 files in source_dir - mp4_files = [] - for file in os.listdir(source_dir): - if file.endswith('.mp4'): - mp4_files.append(os.path.join(source_dir, file)) - - # Sort files numerically to ensure consistent ordering (1.mp4, 2.mp4, ..., 10.mp4, not 1.mp4, 10.mp4, 2.mp4) - def sort_key(filepath): - filename = os.path.basename(filepath) - stem = os.path.splitext(filename)[0] - try: - return int(stem) - except ValueError: - return float('inf') - mp4_files.sort(key=sort_key) - - # Process each MP4 file - for idx, mp4_path in enumerate(mp4_files, 1): - mp4_file = os.path.basename(mp4_path) - - # Extract instruction from filename (remove number prefix and .mp4 extension) - instruction = re.sub(r'^\d+_', '', mp4_file).replace('.mp4', '').replace('_', ' ') - - # If instruction is just a number and task_name is provided, use task_name instead - if task_name and instruction.isdigit(): - # Convert task_name from directory format (e.g., "Lift_the_red_cube_up") to readable format - instruction = task_name.replace('_', ' ') - - label_file = os.path.join(labels_dir, f"{idx}.txt") - with open(label_file, 'w') as f: - f.write(instruction) - - # Copy video with new name - target_video = os.path.join(videos_dir, f"{idx}.mp4") - shutil.copy2(mp4_path, target_video) - - print(f"Processed {mp4_path} -> {idx}.mp4, instruction: {instruction}") - - print(f"Processed {len(mp4_files)} files. Results saved to {output_dir}") - else: - # Recursive behavior: maintain directory structure - print(f"Processing in recursive mode, maintaining directory structure...") - - # Get all subdirectories in the source directory - subdirs = [] - for item in os.listdir(source_dir): - item_path = os.path.join(source_dir, item) - if os.path.isdir(item_path): - subdirs.append(item) - - if not subdirs: - # If no subdirectories, process the source directory directly - print(f"No subdirectories found in {source_dir}, processing directly...") - process_mp4_files(source_dir, output_dir, recursive=False) - return - - # Process each subdirectory - for subdir in subdirs: - source_subdir = os.path.join(source_dir, subdir) - output_subdir = os.path.join(output_dir, subdir) - - # Create output subdirectory - os.makedirs(output_subdir, exist_ok=True) - - # Process files in this subdirectory, passing the directory name as task_name - process_mp4_files(source_subdir, output_subdir, recursive=False, task_name=subdir) - - print(f"Recursive processing complete. Results saved to {output_dir}") - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Process MP4 files from source directory and save to output directory") - parser.add_argument("--source_dir", type=str, required=True, help="Path to source directory containing MP4 files") - parser.add_argument("--output_dir", type=str, required=True, help="Path to output directory") - parser.add_argument("--recursive", action="store_true", help="Process subdirectories recursively, maintaining directory structure") - args = parser.parse_args() - - process_mp4_files(args.source_dir, args.output_dir, args.recursive) \ No newline at end of file diff --git a/scripts/tutorials/cosmos_tutorials/replay.py b/scripts/tutorials/cosmos_tutorials/replay.py deleted file mode 100644 index de61e13e..00000000 --- a/scripts/tutorials/cosmos_tutorials/replay.py +++ /dev/null @@ -1,495 +0,0 @@ -""" -Hybrid Replay: Parquet actions/states + HDF5 initial states - -This script replays recorded robot episodes by: -- Loading initial states from HDF5 datasets -- Loading actions or states from Parquet files -- Running replay inside IsaacLab environments -- Supporting parallel environments and success statistics -""" - - -import multiprocessing -if multiprocessing.get_start_method() != "spawn": - multiprocessing.set_start_method("spawn", force=True) - -# --------------------------------------------------------------------- -# CLI arguments and IsaacLab app launcher -# --------------------------------------------------------------------- -import argparse -from isaaclab.app import AppLauncher - -parser = argparse.ArgumentParser(description="Hybrid Replay for Leisaac") -parser.add_argument("--task", type=str, required=True, help="LeIsaac Task name,e.g. LeIsaac-SO101-LiftCube-v0") -parser.add_argument("--num_envs", type=int, default=1, help="Number of parallel environments") -parser.add_argument("--step_hz", type=int, default=60, help="Simulation frequency") -parser.add_argument("--parquet_dir", type=str, default="", - help="Directory containing Parquet episode files (expects episode_*.parquet)." -parser.add_argument("--hdf5_file", type=str, default="", - help="HDF5 dataset file containing per-episode initial states and metadata.") -parser.add_argument("--select_episodes", type=int, nargs="+", default=[], - help="Episode indices to replay. If empty, replays all available parquet episodes.") -parser.add_argument("--task_type",type=str,default=None,help="Optional task type override. If not set, inferred from --task.") -parser.add_argument("--replay_mode",type=str,default="action",choices=["action", "state"], - help="Replay mode: 'action' replays action vectors; 'state' replays joint states.") - -# Add IsaacLab standard launcher args (e.g., device/headless/experience settings). -AppLauncher.add_app_launcher_args(parser) -args_cli = parser.parse_args() - -# Launch IsaacLab app (must happen before importing certain sim-related modules). -app_launcher = AppLauncher(vars(args_cli)) -simulation_app = app_launcher.app - -# --------------------------------------------------------------------- -# Imports (after app launch) -# --------------------------------------------------------------------- -import contextlib -import os -import time -import re -import gymnasium as gym -import torch -import h5py -import numpy as np -import pandas as pd -from pathlib import Path -from typing import List, Optional, Dict - -from isaaclab.envs import ManagerBasedRLEnv -from isaaclab.utils.datasets import HDF5DatasetFileHandler -from isaaclab_tasks.utils import parse_env_cfg -from leisaac.utils.env_utils import dynamic_reset_gripper_effort_limit_sim, get_task_type - - -# --------------------------------------------------------------------- -# Episode container: parquet actions/states + HDF5 initial states -# --------------------------------------------------------------------- -class HybridEpisodeData: - """ - Holds a single episode's replay payload: - - initial_state: from HDF5 - - actions: from parquet column "action" (optional) - - states: from parquet column "observation.state" (optional) - - Data is denormalized from LeRobot format into IsaacLab radians. - """ - - def __init__(self, actions_df: pd.DataFrame, initial_state: dict, device: str): - self.device = device - self.step_idx = 0 - self.initial_state = initial_state - self.actions = None - self.states = None - - # Parquet data uses normalized angles in [-100, 100]. - # We denormalize into IsaacLab joint limits and convert degrees -> radians. - if "action" in actions_df.columns: - actions_np = np.stack(actions_df["action"].to_numpy()) - actions_np = self._denormalize_joint_pos(actions_np) - self.actions = torch.tensor(actions_np, device=device, dtype=torch.float32) - - if "observation.state" in actions_df.columns: - states_np = np.stack(actions_df["observation.state"].to_numpy()) - states_np = self._denormalize_joint_pos(states_np) - self.states = torch.tensor(states_np, device=device, dtype=torch.float32) - - @staticmethod - def _denormalize_joint_pos(joint_pos: np.ndarray) -> np.ndarray: - """ - Denormalize: LeRobot normalized joint positions [-100, 100] -> IsaacLab joint limits (degrees), - then convert degrees to radians. - - NOTE: The limits here must match the embodiment/task joint conventions. - """ - ISAACLAB_LIMITS = [ - (-110.0, 110.0), - (-100.0, 100.0), - (-100.0, 90.0), - (-95.0, 95.0), - (-160.0, 160.0), - (-10, 100.0), - ] - LEROBOT_LIMITS = [ - (-100, 100), - (-100, 100), - (-100, 100), - (-100, 100), - (-100, 100), - (0, 100), - ] - - result = joint_pos.copy() - for i in range(6): - isaaclab_min, isaaclab_max = ISAACLAB_LIMITS[i] - lerobot_min, lerobot_max = LEROBOT_LIMITS[i] - isaac_range = isaaclab_max - isaaclab_min - lerobot_range = lerobot_max - lerobot_min - result[:, i] = (result[:, i] - lerobot_min) / lerobot_range * isaac_range + isaaclab_min - - # Convert degrees to radians - result = result * np.pi / 180.0 - return result - - def get_next_action(self): - """Return the next action for replay_mode='action', or None if episode finished/unavailable.""" - if self.actions is not None and self.step_idx < len(self.actions): - action = self.actions[self.step_idx] - self.step_idx += 1 - return action - return None - - def get_next_state(self): - """Return the next state for replay_mode='state', or None if episode finished/unavailable.""" - if self.states is not None and self.step_idx < len(self.states): - state = self.states[self.step_idx] - self.step_idx += 1 - return state - return None - - def get_initial_state(self): - """Return the initial state dict loaded from HDF5.""" - return self.initial_state - - -# --------------------------------------------------------------------- -# Dataset loader: align parquet episodes with HDF5 episodes and filter invalid ones -# --------------------------------------------------------------------- -class HybridDatasetLoader: - """ - Loads: - - Parquet episode files: actions/states - - HDF5 episode initial states - - Performs a pre-scan over HDF5 to filter out failed episodes (based on attrs and required keys). - """ - - def __init__(self, hdf5_path: str, parquet_dir: str): - self.hdf5_path = hdf5_path - self.parquet_dir = Path(parquet_dir) - - self.hdf5_handler = HDF5DatasetFileHandler() - self.hdf5_handler.open(self.hdf5_path) - - # Collect parquet episodes and sort by the integer index in their filename (episode_XXXX.parquet). - self.all_parquet_files = sorted( - self.parquet_dir.glob("**/episode_*.parquet"), - key=lambda p: int(re.search(r"episode_(\d+)", p.name).group(1)) - if re.search(r"episode_(\d+)", p.name) - else 0, - ) - - # HDF5 episode names as stored in the dataset (e.g. "demo_0", "demo_1", ...). - self.all_hdf5_names = list(self.hdf5_handler.get_episode_names()) - self.valid_indices = [] - - # Pre-scan HDF5 episodes to filter out failures/missing keys. - print("Pre-scanning HDF5 episode attributes to filter out failed data...") - with h5py.File(self.hdf5_path, "r") as f: - # Use the max range to iterate; assumes parquet and HDF5 correspond broadly. - max_range = max(len(self.all_hdf5_names), len(self.all_parquet_files)) - - for i in range(max_range): - demo_name = self.all_hdf5_names[i] - demo_path = f"data/{demo_name}" - - if demo_path in f: - demo_group = f[demo_path] - # If 'success' attr does not exist, default to True. - is_success = demo_group.attrs.get("success", True) - - if is_success: - try: - # Validate required datasets exist. - _ = demo_group["actions"] - _ = demo_group["obs/joint_pos"] - _ = demo_group["obs/front"] - self.valid_indices.append(i) - except KeyError: - print(f"Pre-scan: Filtering out failed episode {demo_name}") - else: - print(f"Pre-scan: Filtering out failed episode {demo_name}") - - self.num_episodes = len(self.valid_indices) - print(f"Pre-scan complete: {self.num_episodes} valid episodes (original total: {max_range}).") - - def load_episode(self, logical_index: int, device: str): - """ - Load an episode: - - Parquet is addressed directly by logical_index (episode ordering from parquet list) - - HDF5 episode is chosen via valid_indices mapping - - Returns: - HybridEpisodeData or None if out-of-range. - """ - if logical_index >= len(self.all_parquet_files): - return None - - parquet_path = self.all_parquet_files[logical_index] - - # Map logical index (parquet ordering) into an actual HDF5 index among valid ones. - if logical_index < len(self.valid_indices): - actual_hdf5_idx = self.valid_indices[logical_index] - else: - # If parquet list is longer than valid HDF5 list, fall back to last valid index. - actual_hdf5_idx = self.valid_indices[-1] if self.valid_indices else 0 - - demo_name = self.all_hdf5_names[actual_hdf5_idx] - hdf5_ep_data = self.hdf5_handler.load_episode(demo_name, device) - initial_state = hdf5_ep_data.get_initial_state() - - df = pd.read_parquet(parquet_path) - return HybridEpisodeData(df, initial_state, device) - - def close(self): - """Close underlying HDF5 handler.""" - self.hdf5_handler.close() - - -# --------------------------------------------------------------------- -# Rate limiter: attempt to replay at a fixed step rate while rendering -# --------------------------------------------------------------------- -class RateLimiter: - """ - Simple real-time rate limiter: - - sleeps to maintain a target hz - - calls env.sim.render() during waiting to keep UI responsive - """ - - def __init__(self, hz): - self.hz = hz - self.last_time = time.time() - self.sleep_duration = 1.0 / hz - self.render_period = min(0.0166, self.sleep_duration) - - def sleep(self, env): - next_wakeup_time = self.last_time + self.sleep_duration - while time.time() < next_wakeup_time: - time.sleep(self.render_period) - env.sim.render() - - # Advance the timeline by exactly one tick. - self.last_time = self.last_time + self.sleep_duration - - # Catch up if we are behind. - if self.last_time < time.time(): - while self.last_time < time.time(): - self.last_time += self.sleep_duration - - -# --------------------------------------------------------------------- -# Replay helpers -# --------------------------------------------------------------------- -def get_next_action_or_state(episode_data, replay_mode: str): - """ - Depending on replay_mode: - - "state": replay joint states (episode_data.states) - - "action": replay actions (episode_data.actions) - """ - if replay_mode == "state": - next_state = episode_data.get_next_state() - if next_state is None: - return None - return next_state - else: - return episode_data.get_next_action() - - -# --------------------------------------------------------------------- -# Replay manager: multi-env episode scheduling, stepping, success tracking -# --------------------------------------------------------------------- -class ReplayManager: - """ - Runs replay in N parallel envs: - - schedules episodes to env slots - - steps the env with action/state vectors - - tracks success based on termination signal - - prints per-episode success/failure summary - """ - - def __init__(self, env, dataset_loader, config): - self.env = env - self.loader = dataset_loader - self.cfg = config - self.num_envs = env.num_envs - self.device = env.device - - # If env provides a configured idle_action, use it; otherwise default to zeros. - if hasattr(env.cfg, "idle_action"): - self.idle_action = env.cfg.idle_action.repeat(self.num_envs, 1) - else: - self.idle_action = torch.zeros(env.action_space.shape, device=self.device) - - self.rate_limiter = RateLimiter(config.step_hz) - self.task_type = get_task_type(config.task, config.task_type) - self.success_count = 0 - self.total_episodes = 0 - - def run(self, episode_indices: List[int]): - """ - Replay the given list of episode indices. - - Episode lifecycle per env: - - if env has no active episode: load next from queue, reset_to initial state - - each step: pull next action/state and env.step() - - if episode ends: record success based on termination flag observed during replay - """ - episode_queue = list(episode_indices) - - # Per-env bookkeeping: - # env_data_map: active HybridEpisodeData for each env_id - # env_idx_map: which episode index is running in each env_id - # env_success_map: whether that env has triggered termination (treated as success) - env_data_map: Dict[int, Optional[HybridEpisodeData]] = {i: None for i in range(self.num_envs)} - env_idx_map: Dict[int, Optional[int]] = {i: None for i in range(self.num_envs)} - env_success_map: Dict[int, bool] = {i: False for i in range(self.num_envs)} - recorded = set() - - print(f"Replaying {len(episode_indices)} episodes...") - - # Use inference_mode to avoid autograd overhead; suppress Ctrl+C traceback noise. - with contextlib.suppress(KeyboardInterrupt), torch.inference_mode(): - while simulation_app.is_running() and not simulation_app.is_exiting(): - actions = self.idle_action.clone() - has_next = False - - for env_id in range(self.num_envs): - curr_data = env_data_map[env_id] - next_act = get_next_action_or_state(curr_data, self.cfg.replay_mode) if curr_data else None - - if next_act is None: - # Episode finished: record success/failure once per episode index. - curr_idx = env_idx_map[env_id] - if curr_idx is not None and curr_idx not in recorded: - recorded.add(curr_idx) - self.total_episodes += 1 - if env_success_map[env_id]: - self.success_count += 1 - print(f"[Episode #{curr_idx}] ✓ SUCCESS | {self.success_count}/{self.total_episodes}") - else: - print(f"[Episode #{curr_idx}] ✗ FAILED | {self.success_count}/{self.total_episodes}") - - # Load next episode from queue for this env slot. - next_idx = None - while episode_queue: - next_idx = episode_queue.pop(0) - if next_idx < len(self.loader.all_parquet_files): - break - next_idx = None - - if next_idx is not None: - new_data = self.loader.load_episode(next_idx, self.device) - if new_data: - env_data_map[env_id] = new_data - env_idx_map[env_id] = next_idx - env_success_map[env_id] = False - - # Reset only the selected env_id using provided initial state. - self.env.reset_to( - new_data.get_initial_state(), - torch.tensor([env_id], device=self.device), - is_relative=True, - ) - - # After reset, immediately fetch the first action/state for this episode. - next_act = get_next_action_or_state(new_data, self.cfg.replay_mode) - has_next = True - else: - # No more episodes available. - env_data_map[env_id] = None - env_idx_map[env_id] = None - continue - else: - has_next = True - - # Apply action/state for this env slot. - if next_act is not None: - actions[env_id] = next_act - - # If no env has any next step, exit. - if not has_next: - break - - # Optionally update gripper effort limits dynamically between steps. - if self.cfg.replay_mode == "action" and self.env.cfg.dynamic_reset_gripper_effort_limit: - dynamic_reset_gripper_effort_limit_sim(self.env, self.task_type) - - # Step simulation with batched actions. - _, _, reset_terminated, _, _ = self.env.step(actions) - - # Mark success if termination triggered for that env slot. - for env_id in range(self.num_envs): - if reset_terminated[env_id] and env_idx_map[env_id] is not None: - env_success_map[env_id] = True - - # Maintain wall-clock rate and render while waiting. - self.rate_limiter.sleep(self.env) - - # Print final success rate summary. - success_rate = (self.success_count / max(1, self.total_episodes)) * 100 - print(f"\n{'='*50}") - print(f"Success Rate: {self.success_count}/{self.total_episodes} ({success_rate:.1f}%)") - print(f"{'='*50}") - - -# --------------------------------------------------------------------- -# Environment creation helper -# --------------------------------------------------------------------- -def create_env(args): - """ - Create and initialize IsaacLab environment with appropriate task config. - Also configures teleop device and disables recorders/timeouts where appropriate. - """ - env_cfg = parse_env_cfg(args.task, device=args.device, num_envs=args.num_envs) - task_type = get_task_type(args.task, args.task_type) - env_cfg.use_teleop_device(task_type) - - # Some tasks (Direct) use manual termination, so disable timeouts/recorders. - if "Direct" in args.task: - env_cfg.never_time_out = True - env_cfg.manual_terminate = True - env_cfg.recorders = {} - else: - env_cfg.recorders = {} - # Disable time_out termination if present. - if hasattr(env_cfg, "terminations") and hasattr(env_cfg.terminations, "time_out"): - env_cfg.terminations.time_out = None - - env = gym.make(args.task, cfg=env_cfg).unwrapped - if hasattr(env, "initialize"): - env.initialize() - env.reset() - return env - - -# --------------------------------------------------------------------- -# Main entry -# --------------------------------------------------------------------- -def main(): - # Load datasets (parquet + hdf5) and pre-filter invalid episodes. - dataset_loader = HybridDatasetLoader(args_cli.hdf5_file, args_cli.parquet_dir) - print(f"Datasets loaded. Available episodes: {dataset_loader.num_episodes}") - - # If no episodes are explicitly selected, replay all parquet episodes found. - episode_indices = args_cli.select_episodes if args_cli.select_episodes else list(range(len(dataset_loader.all_parquet_files))) - episode_indices = [idx for idx in episode_indices if idx < len(dataset_loader.all_parquet_files)] - - if not episode_indices: - print("No valid episodes to replay.") - return - - print(f"Selected {len(episode_indices)} episodes for replay.") - - # Create env and run replay. - env = create_env(args_cli) - manager = ReplayManager(env, dataset_loader, args_cli) - manager.run(episode_indices) - - # Cleanup resources. - env.close() - dataset_loader.close() - simulation_app.close() - - -if __name__ == "__main__": - main() diff --git a/scripts/tutorials/cosmos_tutorials/so101.sh b/scripts/tutorials/cosmos_tutorials/so101.sh deleted file mode 100644 index e72003d4..00000000 --- a/scripts/tutorials/cosmos_tutorials/so101.sh +++ /dev/null @@ -1,149 +0,0 @@ -#!/bin/bash -set -e - -# ============================================================================= -# Configuration -# ============================================================================= - -# Directory containing Cosmos Predict 2.5 inference outputs. -# This directory should include generated videos produced by the Cosmos model. -LEROBOT_INPUT="" - -# Working directory used to store all intermediate and final IDM outputs. -# It is recommended to place this on a fast local disk. -WORK_DIR="" - -# Robot embodiment type (used by IDM and LeRobot). -ROBOT_TYPE="so101" - -# Key name used to store videos in the LeRobot dataset structure. -# This should match the observation key expected by downstream IDM scripts. -VIDEO_KEY="observation.images.front" - -# ============================================================================= -# Intermediate directories (auto-generated) -# ============================================================================= - -# Step 1 output: task-named directories converted from Cosmos outputs -STEP1_DIR="${WORK_DIR}/step1" - -# Step 2 output: split videos and text instructions -STEP2_DIR="${WORK_DIR}/step2" - -# Step 3 output: preprocessed videos (e.g., resized, normalized) -STEP3_DIR="${WORK_DIR}/step3" - -# Step 4 output: final LeRobot-format dataset -STEP4_DIR="${WORK_DIR}/${ROBOT_TYPE}.data" - -# ============================================================================= -# Step 1: Convert Cosmos outputs to task-based directory structure -# ============================================================================= -# - Reads Cosmos Predict outputs -# - Groups videos by task name -# - Prepares data for downstream preprocessing -python3 IDM_dump/scripts/preprocess_leisaac/cosmos2.5_to_step2_format.py \ - --cosmos_dir "${LEROBOT_INPUT}" \ - --output_dir "${STEP1_DIR}" - -# ============================================================================= -# Step 2: Split videos and instructions -# ============================================================================= -# - Separates raw videos and text instructions into: -# - videos/ -# - labels/ -# - The --recursive flag allows processing nested task directories -python3 IDM_dump/scripts/preprocess_leisaac/split_video_instruction.py \ - --source_dir "${STEP1_DIR}" \ - --output_dir "${STEP2_DIR}" \ - --recursive - -# ============================================================================= -# Step 3: Preprocess videos -# ============================================================================= -# - Resizes videos to the resolution expected by IDM -# - Converts video format if necessary -# - Preserves directory structure across tasks -python3 IDM_dump/scripts/preprocess_leisaac/preprocess_video.py \ - --src_dir "${STEP2_DIR}" \ - --dst_dir "${STEP3_DIR}" \ - --dataset "${ROBOT_TYPE}" \ - --original_width 640 \ - --original_height 480 \ - --recursive - -# ============================================================================= -# IMPORTANT USAGE NOTE -# ============================================================================= -# It is STRONGLY RECOMMENDED to: -# -# 1. Run Step 1–3 first -# 2. Inspect the contents of ${STEP3_DIR} -# 3. Identify the generated task directory name(s) -# 4. Manually copy the desired task directory name into Step 4 -# -# This avoids hard-coding task names before they are known and -# allows flexible reuse of this script across different tasks. -# -# Example: -# ls ${STEP3_DIR} -# → Lift_the_red_cube_up -# -# Then use: -# --input_dir "${STEP3_DIR}/Lift_the_red_cube_up" -# -# ============================================================================= - -# ============================================================================= -# Step 4: Convert preprocessed data to LeRobot dataset reminder -# ============================================================================= -# --input_dir: -# Path to a task-specific directory under STEP3_DIR. -# The directory name MUST match the task name generated in Step 3. -# -# --output_dir: -# Target directory for the LeRobot-format dataset. -# -# --fps: -# Target frames per second for the output dataset. -# -# --embodiment: -# Robot embodiment identifier used by LeRobot/IDM. -# -# --video_key: -# Observation key used to store video data. -python3 IDM_dump/scripts/preprocess_leisaac/raw_to_lerobot.py \ - --input_dir "${STEP3_DIR}/Lift_the_red_cube_up" \ - --output_dir "${STEP4_DIR}" \ - --fps 16 \ - --embodiment "${ROBOT_TYPE}" \ - --video_key "${VIDEO_KEY}" \ - --embodiment "so101" - -# ============================================================================= -# Step 5: Dump IDM actions from the LeRobot dataset -# ============================================================================= -# - Loads a pretrained IDM checkpoint -# - Runs IDM inference on the LeRobot dataset -# - Exports predicted action trajectories -# -# --checkpoint: -# Path to a trained IDM checkpoint. -# -# --dataset: -# Path to the LeRobot dataset generated in Step 4. -# -# --output_dir: -# Output directory where IDM predictions will be stored. -# -# --num_gpus: -# Number of GPUs used for IDM inference. -# -# --video_indices: -# Indices of videos to process (e.g., "0 16" processes videos 0–16). -python3 IDM_dump/scripts/preprocess_leisaac/dump_idm_actions.py \ - --checkpoint "" \ - --dataset "${STEP4_DIR}" \ - --output_dir "${STEP4_DIR}_idm_cosmos" \ - --num_gpus 1 \ - --video_indices "0 16" From 238202d191ef16d5527377fe0cb1ea0557114886 Mon Sep 17 00:00:00 2001 From: chuchen0210 <167320956+chuchen0210@users.noreply.github.com> Date: Thu, 8 Jan 2026 16:05:32 +0800 Subject: [PATCH 5/9] Update cosmos_tutorial.md --- docs/docs/docs/tutorials/cosmos_tutorial.md | 774 +++++++++++++++++++- 1 file changed, 741 insertions(+), 33 deletions(-) diff --git a/docs/docs/docs/tutorials/cosmos_tutorial.md b/docs/docs/docs/tutorials/cosmos_tutorial.md index 614f6862..9a134a96 100644 --- a/docs/docs/docs/tutorials/cosmos_tutorial.md +++ b/docs/docs/docs/tutorials/cosmos_tutorial.md @@ -12,10 +12,11 @@ This tutorial extends **LeIsaac** by integrating **Cosmos-Predict2.5** and **GR0 ## Overview -1. Use **LeIsaac** to collect an **HDF5 dataset** and convert it into a **LeRobot dataset** -2. Post-train **Cosmos-Predict2.5**, then run inference to **generate synthetic videos** -3. Fine-tune **IDM**, then run inference to **generate synthetic LeRobot trajectories(parquet)** -4. Use **LeIsaac** to **replay and evaluate** using the HDF5 dataset + inferred action trajectories +1. Use **LeIsaac** to collect an **HDF5 dataset** and convert it into a **LeRobot dataset** +2. Post-train **Cosmos-Predict2.5**, then run inference to **generate synthetic videos** +3. Fine-tune **IDM**, then run inference to **generate synthetic LeRobot trajectories(parquet)** +4. Convert the **original HDF5 dataset** and the **IDM-generated LeRobot trajectories (parquet)** to a **replayable LeIsaac HDF5 dataset** and use **LeIsaac** to **replay and evaluate** + ## Step 1: Data Collection with LeIsaac @@ -31,7 +32,7 @@ First, collect demonstrations using **LeIsaac** teleoperation. - Reference: **[Teleoperation | LeIsaac Document](https://lightwheelai.github.io/leisaac/docs/getting_started/teleoperation)** - Output: an HDF5 file (for example, `dataset.hdf5`) -The HDF5 dataset serves as the primary source of truth for replay and evaluation. +The HDF5 dataset serves as the primary source of truth for replay and evaluation. It is also the source used to construct the initial **LeRobot dataset** for both Cosmos and IDM. @@ -44,7 +45,7 @@ Convert the recorded HDF5 dataset into **LeRobot format**. ::::info **💡Important (Video Encoding)** - + Ensure that all output videos are encoded using **H.264 (h264)**. Avoid **AV1** encoding, as Cosmos and IDM may fail during decoding or processing. @@ -61,8 +62,8 @@ You can directly modify the `leisaac/scripts/convert/isaaclab2lerobot.py` to ens ## Step 2: Video Generation with Cosmos-predict2.5 ::::tip[What you get from this step] -- A **Cosmos-Predict2.5 checkpoint** post-trained on your task-specific LeRobot videos -- A set of **synthetic rollout videos** generated by Cosmos +- A **Cosmos-Predict2.5 checkpoint** post-trained on your task-specific LeRobot videos +- A set of **synthetic rollout videos** generated by Cosmos In this pipeline, **Cosmos-Predict 2.5** is used purely as a **video generator**. @@ -110,23 +111,159 @@ After post-training, the Cosmos-Predict2.5 model checkpoints are typically saved **Follow the [converting-dcp-checkpoint-to-consolidated-pytorch-format](https://github.com/nvidia-cosmos/cosmos-predict2.5/blob/main/docs/post-training_video2world_gr00t.md#31-converting-dcp-checkpoint-to-consolidated-pytorch-format) to convert DCP checkpoints**. -Once the checkpoint has been converted, run video generation inference following the [official running-inference instructions ](https://github.com/nvidia-cosmos/cosmos-predict2.5/blob/main/docs/post-training_video2world_gr00t.md#32-running-inference). +Once the checkpoint has been converted, run video generation inference following the [official running-inference instructions ](https://github.com/nvidia-cosmos/cosmos-predict2.5/blob/main/docs/post-training_video2world_gr00t.md#32-running-inference). 🔬 **Batch Inference (Used for fine-tuning IDM)** -For large-scale video generation, batch inference is supported. -Copy the batch inference helper script `generate_batch_config.py` into the folder:`cosmos-predict2.5/scripts/`. +For large-scale video generation, batch inference is supported. +Create the batch inference helper script `generate_batch_config.py` under `cosmos-predict2.5/scripts/`. + +
+cosmos-predict2.5/scripts/generate_batch_config.py + +```python +#!/usr/bin/env python3 +""" +Generate batch inference configuration files (JSONL). + +This script scans a directory of input videos and generates a JSONL file +for batch video2world inference. Each line corresponds to one video. +Optionally, task prompts can be loaded from episodes.jsonl. +""" + +import json +import os +from pathlib import Path + +# --------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------- + +# Input video directory used to generate inference configs +VIDEO_DIR = "/videos/chunk-000/observation.images.front" + +# Metadata directory (used when loading text prompts) +META_DIR = "/meta" + +# Output JSONL file +OUTPUT_JSONL = "batch_inference_config.jsonl" + +# Task name prefix +TASK_NAME = "liftcube" + +# --------------------------------------------------------------------- +# Base configuration template shared by all videos +# --------------------------------------------------------------------- + +BASE_CONFIG = { + "inference_type": "video2world", + "seed": 21, + "guidance": 7, + "resolution": "480,640", + "enable_autoregressive": True, + # Number of output frames (e.g. 110 ≈ 6s, 140 ≈ 8s, 210 ≈ 16s) + "num_output_frames": 210, + "chunk_size": 77, + "chunk_overlap": 1, + # Default prompt (will be overridden if episode-specific prompt is found) + "prompt": "The robot arm is performing a task", + "negative_prompt": ( + "The video captures a series of frames showing ugly scenes, static with no motion, " + "motion blur, over-saturation, shaky footage, low resolution, grainy texture, " + "pixelated images, poorly lit areas, underexposed and overexposed scenes, " + "poor color balance, washed out colors, choppy sequences, jerky movements, " + "low frame rate, artifacting, color banding, unnatural transitions, " + "outdated special effects, fake elements, unconvincing visuals, " + "poorly edited content, jump cuts, visual noise, and flickering. " + "Overall, the video is of poor quality." + ), +} + +# --------------------------------------------------------------------- +# Main logic +# --------------------------------------------------------------------- + +def main(use_prompt=True): + video_dir = Path(VIDEO_DIR) + meta_dir = Path(META_DIR) + + # Collect all mp4 video files + video_files = sorted(video_dir.glob("*.mp4")) + print(f"Found {len(video_files)} video files") + + # Load prompts from episodes.jsonl if enabled + episode_prompts = {} + if use_prompt: + episodes_file = meta_dir / "episodes.jsonl" + if episodes_file.exists(): + with open(episodes_file, 'r') as f: + for line in f: + episode_data = json.loads(line) + episode_index = episode_data.get("episode_index") + tasks = episode_data.get("tasks", []) + if episode_index is not None and tasks: + # Use the first task as the episode prompt + prompt = f"The robot arm is performing a task. {tasks[0]}" + episode_prompts[episode_index] = prompt + print(f"✓ Loaded {len(episode_prompts)} prompts from episodes.jsonl") + + # Generate JSONL configuration file + with open(OUTPUT_JSONL, 'w') as f: + for video_file in video_files: + # Video filename without extension (e.g. episode_000001) + video_name = video_file.stem + + # Create per-video config based on the base template + config = BASE_CONFIG.copy() + config["name"] = f"{TASK_NAME}_{video_name}" + config["input_path"] = str(video_file) + + # Assign episode-specific prompt if available + if use_prompt: + # Extract episode index from video filename + episode_index = int(video_name.split('_')[-1]) + if episode_index in episode_prompts: + config["prompt"] = episode_prompts[episode_index] + + # Write one JSON object per line + f.write(json.dumps(config) + '\n') + + # Summary information + print(f"✓ Config file generated: {OUTPUT_JSONL}") + print(f" - Total videos: {len(video_files)}") + print(f" - Use prompts: {'Yes' if use_prompt else 'No'}") + if use_prompt: + print( + f" - Videos with matched prompts: " + f"{sum(1 for vf in video_files if int(vf.stem.split('_')[-1]) in episode_prompts)}" + ) + +# --------------------------------------------------------------------- +# Script entry point +# --------------------------------------------------------------------- + +if __name__ == "__main__": + import sys + # Text prompts are enabled only when '--use-prompt' is provided + use_prompt = "--use-prompt" in sys.argv + main(use_prompt=use_prompt) + +``` + +
You can generate a batch inference configuration `jsonl` file using: ```bash +# Run this command from the cosmos-predict2.5 project root +# cd python scripts/generate_batch_config.py --use-prompt ``` ## Step 3: Action Inference with IDM ::::tip[What you get from this step] -- An **IDM checkpoint** fine-tuned on your LeRobot dataset and robot embodiment -- A set of **LeRobot-format trajectories** (e.g. parquet files) inferred from Cosmos-generated videos +- An **IDM checkpoint** fine-tuned on your LeRobot dataset and robot embodiment +- A set of **LeRobot-format trajectories** (e.g. parquet files) inferred from Cosmos-generated videos In this pipeline, **IDM (Inverse Dynamics Model)** from **GR00T-Dreams** is used to convert **Cosmos-generated rollout videos** into executable robot actions. @@ -135,7 +272,7 @@ IDM is first fine-tuned on the original LeRobot dataset collected via **LeIsaac* ### 3.1 Install IDM Environment -IDM requires the **Cosmos-Predict2** environment (**not 2.5**). +IDM requires the **Cosmos-Predict2** environment (**not 2.5**). Follow the official prerequisites guide:**[Cosmos-Predict2 – Prerequisites](https://github.com/nvidia-cosmos/cosmos-predict2/blob/main/documentations/post-training_video2world_gr00t.md#prerequisites)** @@ -150,17 +287,17 @@ Follow the official prerequisites guide:**[Cosmos-Predict2 – Prerequisites](ht ``` - If you encounter issues related to `APEX`, it can be safely removed from the uv environment.APEX is not required for IDM training or inference.: -### 3.2 Fine-tune IDM +### 3.2 Fine-tune IDM For IDM fine-tuning, please refer to **[training-custom-idm-model](https://github.com/NVIDIA/GR00T-Dreams?tab=readme-ov-file#optional-33-training-custom-idm-model)** #### 3.2.1 Preparation: Modality Metadata and DataConfig -**Step 1: Add modality.json** - +**1.Add modality.json** Create `modality.json` under `GR00T-Dreams/IDM_dump/global_metadata/{embodiment_name}/` and copy the same file to `/meta/`
Example: SO101 modality.json + ```json { "state": { @@ -189,12 +326,12 @@ Create `modality.json` under `GR00T-Dreams/IDM_dump/global_metadata/{embodiment_ ```
-**Step 2: Add a New DataConfig (So101DataConfig)** - +**2.Add a New DataConfig (So101DataConfig)** Add a new `So101DataConfig` class defining in `GR00T-Dreams/gr00t/experiment/data_config_idm.py`:
GR00T-Dreams/gr00t/experiment/data_config_idm.py + ```python class So101DataConfig(BaseDataConfig): video_keys = ["video.front"] @@ -304,7 +441,7 @@ PYTHONPATH=. torchrun scripts/idm_training.py \ --output_dir ``` -After loading the dataset, `stats.json` will be automatically generated under `/meta/`. +After loading the dataset, `stats.json` will be automatically generated under `/meta/`. **Copy this file to `GR00T-Dreams/IDM_dump/global_metadata/{embodiment_name}/`.** @@ -312,7 +449,7 @@ After loading the dataset, `stats.json` will be automatically generated under `< After IDM post-training, use the trained model to infer actions from Cosmos-generated videos. -#### Step 1: Prepare Inference Configuration +#### 3.3.1 Prepare Inference Configuration Select a checkpoint directory (e.g. checkpoint-10000/) and create `checkpoint-10000/experiment_cfg/conf.yaml`: @@ -481,37 +618,608 @@ metadata_versions: Use a configuration matching the SO101 modality definition and transforms (as shown above). -#### Step 2: Run IDM Inference +#### 3.3.2 Run IDM Inference + +**1.Modify conversion scripts to add a specific embodiment (SO101)** +To support the SO101 embodiment, we extend the conversion utilities under +`GR00T-Dreams/IDM_dump/` to handle **SO101-specific** video streams and metadata. + +
+GR00T-DreamsIDM_dump/preprocess_video.py -Copy the preprocessing scripts folder `preprocess_leisaac` to `GR00T-Dreams/dump/scripts/` and `so101.sh` to `GR00T-Dreams/IDM_dump/scripts/preprocess/` +```diff ++ # === [1/3] Video frame writing: add SO101 output key (front view) === + elif dataset == 'so100': + image = resize_with_padding(frame, ratio) + output_videos['observation.images.webcam'].append_data(image) ++ elif dataset == 'so101': ++ image = resize_with_padding(frame, ratio) ++ output_videos['observation.images.front'].append_data(image) + else: + raise ValueError(f"Unknown task: {src_path}") + + ---------------------------------------------------------------------- + ++ # === [2/3] Output directory mapping: add videos/observation.images.front for SO101 === + elif dataset == 'so100': + output_dirs = { + 'observation.images.webcam': os.path.join(dst_dir, 'videos', 'observation.images.webcam'), + } ++ elif dataset == 'so101': ++ output_dirs = { ++ 'observation.images.front': os.path.join(dst_dir, 'videos', 'observation.images.front'), ++ } + + for dir_path in output_dirs.values(): + os.makedirs(dir_path, exist_ok=True) + + ---------------------------------------------------------------------- + ++ # === [3/3] CLI argument: add so101 to dataset choices === + parser.add_argument('--max_videos', type=int, default=None, + help='Maximum number of videos to process per subdirectory (for debugging)') +- parser.add_argument('--dataset', type=str, default='robocasa', +- help='Dataset name', choices=['robocasa', 'gr1', 'franka', 'so100']) ++ parser.add_argument('--dataset', type=str, default='robocasa', ++ help='Dataset name', choices=['robocasa', 'gr1', 'franka', 'so100', 'so101']) + + parser.add_argument("--recursive", action="store_true", help="Process subdirectories recursively, maintaining directory structure") +``` +
+
+GR00T-DreamsIDM_dump/raw_to_lerobot.py + +```diff ++ # === [1/2] Embodiment inference & annotation source: add SO101 === + if args.embodiment is None: + if 'robocasa' in args.output_dir: + args.embodiment = "robocasa_panda_omron" + elif 'gr1' in args.output_dir: + args.embodiment = "gr1_unified" + elif 'franka' in args.output_dir: + args.embodiment = "franka" + elif 'so100' in args.output_dir: + args.embodiment = "so100" ++ elif 'so101' in args.output_dir: ++ args.embodiment = "so101" + else: + raise ValueError(f"Unknown embodiment for {args.output_dir}") + + if args.embodiment == "robocasa_panda_omron": + args.annotation_source = "human.action.task_description" + elif args.embodiment == "gr1_unified": + args.annotation_source = "human.coarse_action" + elif args.embodiment == "franka": + args.annotation_source = "language.language_instruction" + elif args.embodiment == "so100": + args.annotation_source = "human.task_description" ++ elif args.embodiment == "so101": ++ args.annotation_source = "human.task_description" + + ---------------------------------------------------------------------- + ++ # === [2/2] Global metadata source: add SO101 modality definition === + elif args.embodiment == "so100": + source_dir = "IDM_dump/global_metadata/so100" ++ elif args.embodiment == "so101": ++ source_dir = "IDM_dump/global_metadata/so101" + + # copy modality.json + shutil.copy( + source_dir + "/modality.json", + args.output_dir + "/meta/modality.json" + ) +``` +
+ +
+GR00T-DreamsIDM_dump/dump_idm_actions.py + +```diff ++ # === [1/1] Embodiment tag mapping: add SO101 === + if "gr1" in embodiment: + embodiment_tag = EmbodimentTag.GR1_unified + elif "franka" in embodiment: + embodiment_tag = EmbodimentTag.FRANKA + elif "so100" in embodiment: + embodiment_tag = EmbodimentTag.SO100 ++ elif "so101" in embodiment: ++ embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + elif "robocasa" in embodiment: + embodiment_tag = EmbodimentTag.ROBOCASA + else: + raise ValueError(f"Unknown embodiment: {embodiment}") +``` +
+ + +**2.Create a format conversion interface from Cosmos to IDM** +Create the convert helper script `cosmos2.5_to_step2_format.py` under `GR00T-Dreams/IDM_dump/`. +
+GR00T-Dreams/IDM_dump/cosmos2.5_to_step2_format.py + +```python +#!/usr/bin/env python3 +""" +Convert cosmos2.5 outputs (+ optional LeRobot meta) to convert_directory output format. + +Input: + cosmos_dir/ + ├── *.mp4 + ├── *.json(same stem as mp4) + └── ... + +Optional: + lerobot_dir/ + └── meta/ + ├── tasks.jsonl + └── episodes.jsonl + +Output: + output_dir/ + ├── / + │ ├── 0.mp4 + │ ├── 1.mp4 + │ └── ... + └── ... + +Rule: + - If json has "prompt": use it as task name (sanitize to dir) + - Else: fallback to LeRobot meta (requires --lerobot_dir) + - If lerobot_dir not provided: print warning and skip that sample +""" + +import argparse +import json +import shutil +from pathlib import Path +from collections import defaultdict +from typing import Dict, List, Optional, Tuple + + +def sanitize_task_name(task: str) -> str: + """ + Convert task description to valid directory name. + Example: "Lift the red cube up." -> "Lift_the_red_cube_up" + """ + task = task.strip() + for ch in [".", ",", "!", "?", ":", ";", "\"", "'"]: + task = task.replace(ch, "") + task = task.replace(" ", "_") + task = "_".join(filter(None, task.split("_"))) + return task + + +def load_tasks(tasks_file: Path) -> Dict[int, str]: + """Load task definitions. Returns {task_index: task_string}""" + tasks: Dict[int, str] = {} + with open(tasks_file, "r") as f: + for line in f: + if not line.strip(): + continue + data = json.loads(line) + tasks[int(data["task_index"])] = data["task"] + return tasks + + +def load_episodes(episodes_file: Path) -> List[dict]: + """Load episode information. Returns list of dict (episode_index, tasks, length, ...)""" + episodes: List[dict] = [] + with open(episodes_file, "r") as f: + for line in f: + if not line.strip(): + continue + episodes.append(json.loads(line)) + return episodes + + +def build_episode_to_task_map(lerobot_dir: Path) -> Dict[int, str]: + """ + Build mapping: episode_index -> task_string (first task in episode['tasks']) + """ + meta_dir = lerobot_dir / "meta" + tasks_path = meta_dir / "tasks.jsonl" + episodes_path = meta_dir / "episodes.jsonl" + + if not tasks_path.exists() or not episodes_path.exists(): + raise FileNotFoundError(f"LeRobot meta files not found under: {meta_dir}") + + _tasks = load_tasks(tasks_path) # not strictly needed, but useful if episodes store indices in some setups + episodes = load_episodes(episodes_path) + + ep2task: Dict[int, str] = {} + for ep in episodes: + ep_idx = int(ep["episode_index"]) + ep_tasks = ep.get("tasks", []) + if not ep_tasks: + continue + + # In many LeRobot datasets, ep["tasks"] stores task strings already. + # If it stores indices, you can extend here. We'll support both: + t0 = ep_tasks[0] + if isinstance(t0, int): + task_str = _tasks.get(int(t0), str(t0)) + else: + task_str = str(t0) + + ep2task[ep_idx] = task_str + + return ep2task + + +def read_json(path: Path) -> Optional[dict]: + try: + with open(path, "r") as f: + return json.load(f) + except Exception as e: + print(f"[WARN] Failed to read json: {path} ({e})") + return None + + +def parse_episode_index_from_stem(stem: str) -> Optional[int]: + """ + Try to extract episode index from filename stem. + Examples it can handle: + - episode_000123 + - ..._000123 + - 000123 + If cannot parse, return None. + """ + # Most common: "episode_000123" + if stem.startswith("episode_"): + tail = stem[len("episode_") :] + if tail.isdigit(): + return int(tail) + + # Try last underscore chunk + parts = stem.split("_") + for candidate in reversed(parts): + if candidate.isdigit(): + return int(candidate) + + # Entire stem digits? + if stem.isdigit(): + return int(stem) + + return None + + +def convert_cosmos_to_step2( + cosmos_dir: Path, + output_dir: Path, + lerobot_dir: Optional[Path] = None, + chunk_missing_prompt_policy: str = "skip", +): + """ + Convert cosmos2.5 outputs to step2 format. + + chunk_missing_prompt_policy: + - "skip": if no prompt and no lerobot_dir mapping, skip the sample + """ + cosmos_dir = cosmos_dir.resolve() + output_dir = output_dir.resolve() + output_dir.mkdir(parents=True, exist_ok=True) + + ep2task: Dict[int, str] = {} + if lerobot_dir is not None: + ep2task = build_episode_to_task_map(lerobot_dir.resolve()) + + # Collect pairs (json, mp4) by stem + json_files = sorted(cosmos_dir.glob("*.json")) + if not json_files: + raise FileNotFoundError(f"No .json files found in cosmos_dir: {cosmos_dir}") + + # task_name -> list of source mp4 paths (ordered) + task_to_videos: Dict[str, List[Path]] = defaultdict(list) + skipped: List[Tuple[Path, str]] = [] + + for jpath in json_files: + stem = jpath.stem + mpath = cosmos_dir / f"{stem}.mp4" + if not mpath.exists(): + print(f"[WARN] Missing mp4 for json: {jpath.name} -> expected {mpath.name}, skip.") + skipped.append((jpath, "missing_mp4")) + continue + + data = read_json(jpath) + if data is None: + skipped.append((jpath, "bad_json")) + continue + + prompt = data.get("prompt", None) + if isinstance(prompt, str) and prompt.strip(): + task_str = prompt.strip() + task_dir_name = sanitize_task_name(task_str) + task_to_videos[task_dir_name].append(mpath) + continue + + # No prompt -> fallback to lerobot meta mapping + if not ep2task: + print( + f"[WARN] {jpath.name} has no 'prompt'. " + f"You didn't provide --lerobot_dir (or mapping is empty), cannot infer task. Skipping." + ) + skipped.append((jpath, "no_prompt_no_lerobot")) + continue + + ep_idx = parse_episode_index_from_stem(stem) + if ep_idx is None: + print( + f"[WARN] {jpath.name} has no 'prompt' and episode index cannot be parsed from name '{stem}'. Skipping." + ) + skipped.append((jpath, "no_prompt_cannot_parse_episode")) + continue + + task_str = ep2task.get(ep_idx) + if not task_str: + print( + f"[WARN] {jpath.name} has no 'prompt'. Parsed episode_index={ep_idx}, " + f"but it's not found in lerobot meta. Skipping." + ) + skipped.append((jpath, "episode_not_in_meta")) + continue + + task_dir_name = sanitize_task_name(task_str) + task_to_videos[task_dir_name].append(mpath) + + # Copy into output folders with sequential numbering per task + total_copied = 0 + for task_dir_name, vids in sorted(task_to_videos.items(), key=lambda x: x[0]): + dst_task_dir = output_dir / task_dir_name + dst_task_dir.mkdir(parents=True, exist_ok=True) + + # keep deterministic order + vids_sorted = sorted(vids, key=lambda p: p.name) + for i, src in enumerate(vids_sorted): + dst = dst_task_dir / f"{i}.mp4" + shutil.copy2(src, dst) + total_copied += 1 + + print(f"[OK] Task '{task_dir_name}': {len(vids_sorted)} videos") + + print("\nConversion complete!") + print(f" Total videos copied: {total_copied}") + print(f" Output directory: {output_dir}") + + if skipped: + print(f"\n[SUMMARY] Skipped {len(skipped)} samples:") + # print a few for readability + for p, reason in skipped[:20]: + print(f" - {p.name}: {reason}") + if len(skipped) > 20: + print(f" ... and {len(skipped) - 20} more") + + +def main(): + parser = argparse.ArgumentParser( + description="Convert cosmos2.5 outputs (+ optional LeRobot meta) to convert_directory output format" + ) + parser.add_argument("--cosmos_dir", type=str, required=True, help="Cosmos output directory containing *.mp4 and *.json") + parser.add_argument("--output_dir", type=str, required=True, help="Output directory") + parser.add_argument( + "--lerobot_dir", + type=str, + default=None, + help="Optional LeRobot dataset directory (only needed when json has no 'prompt')", + ) + + args = parser.parse_args() + + cosmos_dir = Path(args.cosmos_dir) + output_dir = Path(args.output_dir) + lerobot_dir = Path(args.lerobot_dir) if args.lerobot_dir else None + + convert_cosmos_to_step2( + cosmos_dir=cosmos_dir, + output_dir=output_dir, + lerobot_dir=lerobot_dir, + ) + + +if __name__ == "__main__": + main() +``` +
+ +**3.Create a preprocessing script** +Create the preprocessing helper script `so101.sh` under `GR00T-Dreams/IDM_dump/scripts/preprocess/`. +
+GR00T-Dreams/IDM_dump/scripts/preprocess/so101.sh + +```sh +#!/bin/bash +set -e + +# ============================================================================= +# Configuration +# ============================================================================= + +# Directory containing Cosmos Predict 2.5 inference outputs. +# This directory should include generated videos produced by the Cosmos model. +LEROBOT_INPUT="" + +# Working directory used to store all intermediate and final IDM outputs. +# It is recommended to place this on a fast local disk. +WORK_DIR="" + +# Robot embodiment type (used by IDM and LeRobot). +ROBOT_TYPE="so101" + +# Key name used to store videos in the LeRobot dataset structure. +# This should match the observation key expected by downstream IDM scripts. +VIDEO_KEY="observation.images.front" + +# ============================================================================= +# Intermediate directories (auto-generated) +# ============================================================================= + +# Step 1 output: task-named directories converted from Cosmos outputs +STEP1_DIR="${WORK_DIR}/step1" + +# Step 2 output: split videos and text instructions +STEP2_DIR="${WORK_DIR}/step2" + +# Step 3 output: preprocessed videos (e.g., resized, normalized) +STEP3_DIR="${WORK_DIR}/step3" + +# Step 4 output: final LeRobot-format dataset +STEP4_DIR="${WORK_DIR}/${ROBOT_TYPE}.data" + +# ============================================================================= +# Step 1: Convert Cosmos outputs to task-based directory structure +# ============================================================================= +# - Reads Cosmos Predict outputs +# - Groups videos by task name +# - Prepares data for downstream preprocessing +python3 IDM_dump/scripts/preprocess_leisaac/cosmos2.5_to_step2_format.py \ + --cosmos_dir "${LEROBOT_INPUT}" \ + --output_dir "${STEP1_DIR}" + +# ============================================================================= +# Step 2: Split videos and instructions +# ============================================================================= +# - Separates raw videos and text instructions into: +# - videos/ +# - labels/ +# - The --recursive flag allows processing nested task directories +python3 IDM_dump/scripts/preprocess_leisaac/split_video_instruction.py \ + --source_dir "${STEP1_DIR}" \ + --output_dir "${STEP2_DIR}" \ + --recursive + +# ============================================================================= +# Step 3: Preprocess videos +# ============================================================================= +# - Resizes videos to the resolution expected by IDM +# - Converts video format if necessary +# - Preserves directory structure across tasks +python3 IDM_dump/scripts/preprocess_leisaac/preprocess_video.py \ + --src_dir "${STEP2_DIR}" \ + --dst_dir "${STEP3_DIR}" \ + --dataset "${ROBOT_TYPE}" \ + --original_width 640 \ + --original_height 480 \ + --recursive + +# ============================================================================= +# IMPORTANT USAGE NOTE +# ============================================================================= +# It is STRONGLY RECOMMENDED to: +# +# 1. Run Step 1–3 first +# 2. Inspect the contents of ${STEP3_DIR} +# 3. Identify the generated task directory name(s) +# 4. Manually copy the desired task directory name into Step 4 +# +# This avoids hard-coding task names before they are known and +# allows flexible reuse of this script across different tasks. +# +# Example: +# ls ${STEP3_DIR} +# → Lift_the_red_cube_up +# +# Then use: +# --input_dir "${STEP3_DIR}/Lift_the_red_cube_up" +# +# ============================================================================= + +# ============================================================================= +# Step 4: Convert preprocessed data to LeRobot dataset reminder +# ============================================================================= +# --input_dir: +# Path to a task-specific directory under STEP3_DIR. +# The directory name MUST match the task name generated in Step 3. +# +# --output_dir: +# Target directory for the LeRobot-format dataset. +# +# --fps: +# Target frames per second for the output dataset. +# +# --embodiment: +# Robot embodiment identifier used by LeRobot/IDM. +# +# --video_key: +# Observation key used to store video data. +python3 IDM_dump/scripts/preprocess_leisaac/raw_to_lerobot.py \ + --input_dir "${STEP3_DIR}/Lift_the_red_cube_up" \ + --output_dir "${STEP4_DIR}" \ + --fps 16 \ + --embodiment "${ROBOT_TYPE}" \ + --video_key "${VIDEO_KEY}" + +# ============================================================================= +# Step 5: Dump IDM actions from the LeRobot dataset +# ============================================================================= +# - Loads a pretrained IDM checkpoint +# - Runs IDM inference on the LeRobot dataset +# - Exports predicted action trajectories +# +# --checkpoint: +# Path to a trained IDM checkpoint. +# +# --dataset: +# Path to the LeRobot dataset generated in Step 4. +# +# --output_dir: +# Output directory where IDM predictions will be stored. +# +# --num_gpus: +# Number of GPUs used for IDM inference. +# +# --video_indices: +# Indices of videos to process (e.g., "0 16" processes videos 0–16). +python3 IDM_dump/scripts/preprocess_leisaac/dump_idm_actions.py \ + --checkpoint "" \ + --dataset "${STEP4_DIR}" \ + --output_dir "${STEP4_DIR}_idm_cosmos" \ + --num_gpus 1 \ + --video_indices "0 16" +``` + +
+ +**4.Run IDM inference** Run the inference script: ```bash +# Run this command from the GR00T-Dreams project root +# cd PYTHONPATH=. bash IDM_dump/scripts/preprocess/so101.sh ``` - This step produces the complete LeRobot-format outputs based on the Cosmos-generated videos. +## Step 4: Replay and Evaluate in LeIsaac -## Step 4: Replay and Evaluation in LeIsaac +In this step, the **original HDF5 dataset** and the **IDM-generated LeRobot trajectories (parquet)** are first converted into a **replayable LeIsaac HDF5 dataset**,and then replayed in **Isaac Sim** using **LeIsaac**.This replay process is used to validate the quality and physical plausibility of the inferred action trajectories. -Switch to the **LeIsaac** environment and project directory. -Copy the replay helper script `replay.py` into the LeIsaac scripts folder:`leisaac/scripts/`. +### 4.1 Convert IDM Outputs to LeIsaac HDF5 -Run the replay script with the **target task** enabled: +IDM inference produces action trajectories in **LeRobot parquet format**. +To replay these trajectories in LeIsaac, they must be merged with a source +LeIsaac HDF5 file (used as an initial-state and metadata template). +Switch to the **LeIsaac** environment and run the conversion script from the **LeIsaac project directory**: ```bash -python scripts/replay.py \ - --task LeIsaac-SO101-LiftCube-v0 \ - --enable_cameras +python scripts/convert/lerobot2isaaclab.py \ + --parquet \ + --src_hdf5 \ + --out_hdf5 \ + --overwrite ``` -In this step, the **original HDF5 dataset** and the **IDM-generated LeRobot trajectories (parquet)** are replayed in **Isaac Sim** using **LeIsaac**.This replay process is used to validate the quality and physical plausibility of the inferred action trajectories. +This conversion creates a new, fully replayable LeIsaac HDF5 dataset. + +### 4.2 Replay the Generated Dataset in LeIsaac + +After conversion, the newly generated **HDF5 dataset** can be replayed using **LeIsaac’s action replay mode**.You can refer to **[dataset_replay](https://lightwheelai.github.io/leisaac/docs/getting_started/dataset_replay)** + +
-
\ No newline at end of file + From b6ea1fd0d3e66118a3a7ec38d5bb5168d9e7b182 Mon Sep 17 00:00:00 2001 From: chuchen0210 <167320956+chuchen0210@users.noreply.github.com> Date: Tue, 13 Jan 2026 16:24:33 +0800 Subject: [PATCH 6/9] Update Cosmos tutorial and merge utilities --- docs/docs/docs/tutorials/cosmos_tutorial.md | 33 +- scripts/convert/lerobot2isaaclab.py | 738 ++++++++++---------- scripts/tutorials/cosmos_merge.py | 481 +++++++++++++ 3 files changed, 863 insertions(+), 389 deletions(-) mode change 100755 => 100644 scripts/convert/lerobot2isaaclab.py create mode 100644 scripts/tutorials/cosmos_merge.py diff --git a/docs/docs/docs/tutorials/cosmos_tutorial.md b/docs/docs/docs/tutorials/cosmos_tutorial.md index 9a134a96..dacc8d76 100644 --- a/docs/docs/docs/tutorials/cosmos_tutorial.md +++ b/docs/docs/docs/tutorials/cosmos_tutorial.md @@ -1,6 +1,6 @@ # LeIsaac × Cosmos: Video-to-Action Data Generation Pipeline -This tutorial extends **LeIsaac** by integrating **Cosmos-Predict2.5** and **GR00T-Dreams IDM** into a LeIsaac-native data generation loop. **LeIsaac** is used to collect teleoperated demonstrations (HDF5) and convert them into LeRobot datasets. **Cosmos-Predict 2.5** is post-trained on these videos to synthesize additional rollout videos at scale, and **IDM** is fine-tuned on the same dataset to infer robot actions from the generated videos. Together, this produces **a scalable pipeline for constructing synthetic, complete LeRobot datasets** , which can be replayed and evaluated directly in **LeIsaac**. +This tutorial extends **[LeIsaac](https://github.com/LightwheelAI/leisaac)** by integrating **[Cosmos-Predict2.5](https://github.com/nvidia-cosmos/cosmos-predict2.5)** and **[GR00T-Dreams IDM](https://github.com/nvidia/GR00T-dreams)** into a LeIsaac-native data generation loop. **LeIsaac** is used to collect teleoperated demonstrations (HDF5) and convert them into LeRobot datasets. **Cosmos-Predict 2.5** is post-trained on these videos to synthesize additional rollout videos at scale, and **IDM** is fine-tuned on the same dataset to infer robot actions from the generated videos. Together, this produces **a scalable pipeline for constructing synthetic, complete LeRobot datasets** , which can be replayed and evaluated directly in **LeIsaac**. ## 🎥 From a Single Demonstration to Large-Scale Synthetic Rollouts via Cosmos
@@ -1192,28 +1192,35 @@ This step produces the complete LeRobot-format outputs based on the Cosmos-gener ## Step 4: Replay and Evaluate in LeIsaac -In this step, the **original HDF5 dataset** and the **IDM-generated LeRobot trajectories (parquet)** are first converted into a **replayable LeIsaac HDF5 dataset**,and then replayed in **Isaac Sim** using **LeIsaac**.This replay process is used to validate the quality and physical plausibility of the inferred action trajectories. - +In this step, the **original HDF5 dataset** and the **IDM-generated LeRobot trajectories (parquet)** are first converted and merged into a **replayable LeIsaac HDF5 dataset**, and then replayed in **Isaac Sim** using **LeIsaac**. This replay process is used to validate the quality and physical plausibility of the inferred action trajectories. ### 4.1 Convert IDM Outputs to LeIsaac HDF5 -IDM inference produces action trajectories in **LeRobot parquet format**. -To replay these trajectories in LeIsaac, they must be merged with a source -LeIsaac HDF5 file (used as an initial-state and metadata template). +IDM inference produces action trajectories in **LeRobot parquet format**. To process these in LeIsaac, they must first be converted into a LeIsaac-compatible HDF5 format. Switch to the **LeIsaac** environment and run the conversion script from the **LeIsaac project directory**: + ```bash python scripts/convert/lerobot2isaaclab.py \ - --parquet \ - --src_hdf5 \ - --out_hdf5 \ - --overwrite + --lerobot_dir \ + --output_hdf5 \ + --include_videos +``` + +### 4.2 Merge with Source HDF5 Dataset + +The HDF5 file generated in the step4.1 needs to be merged with the source LeIsaac HDF5 file (from chapter 1.1) to restore initial states required for replay. + +```bash +python scripts/tutorials/cosmos_merge.py \ + --lerobot_hdf5 \ + --source_hdf5 \ + --output_hdf5 ``` -This conversion creates a new, fully replayable LeIsaac HDF5 dataset. -### 4.2 Replay the Generated Dataset in LeIsaac +### 4.3 Replay the Generated Dataset in LeIsaac -After conversion, the newly generated **HDF5 dataset** can be replayed using **LeIsaac’s action replay mode**.You can refer to **[dataset_replay](https://lightwheelai.github.io/leisaac/docs/getting_started/dataset_replay)** +After conversion and merging, the final **HDF5 dataset** can be replayed using **LeIsaac’s action replay mode**. You can refer to **[dataset_replay](https://lightwheelai.github.io/leisaac/docs/getting_started/dataset_replay)** for detailed instructions. diff --git a/scripts/convert/lerobot2isaaclab.py b/scripts/convert/lerobot2isaaclab.py old mode 100755 new mode 100644 index ffc619c6..8ae7088c --- a/scripts/convert/lerobot2isaaclab.py +++ b/scripts/convert/lerobot2isaaclab.py @@ -1,414 +1,400 @@ -#!/usr/bin/env python3 - """ -Convert IDM parquet (LeRobot-style rows) + source LeIsaac HDF5 (initial_state/template) -to a new replayable LeIsaac HDF5. - -Input parquet row example: -{ - "observation.state":[...6...], - "action":[...6...], - "timestamp":0, - "episode_index":0, - "index":0 -} - -Source HDF5 episode structure (example): -/actions (T,6) -/initial_state/... -/obs/actions (T,6) -/processed_actions (T,6) -/states/articulation/robot/joint_position (T,6) -... - -We will: -- Copy root metadata groups (everything except /data) from source HDF5. -- For each episode: - - Copy initial_state + attrs from a template episode in source HDF5 - - Write /actions, /obs/actions, /processed_actions from parquet "action" - - Optionally write /states/.../joint_position from parquet "observation.state" - - Optionally write /timestamps from parquet "timestamp" - - Optionally copy template obs/states arrays (including images) from source HDF5(can be huge) +Convert a local LeRobot dataset folder (v2.1 episode-based layout) to a single HDF5 file. + +Compared to the original version, this script can optionally: + 1) Sort frames inside each episode by a chosen column (default: "index" if present), + 2) Convert LeRobot normalized joint values (roughly degrees in [-100, 100]) into + IsaacLab joint angles in RADIANS, using the same limits mapping as your replay-correct script. + +Expected input layout (v2.1): +dataset_root/ +├── data/chunk-000/episode_000000.parquet +├── data/chunk-000/episode_000001.parquet +├── videos/chunk-000//episode_000000.mp4 +└── meta/episodes.jsonl (optional) + other meta files + +Output HDF5 layout (mirrors folder-like structure): +/ + attrs: source_dir, layout_version, created_at + /meta/ (text dataset, utf-8) + /data///... (datasets from parquet columns) + /videos///.mp4 (uint8 bytes dataset, optional) + """ import argparse -import os -import re +import datetime as dt +import json from pathlib import Path +from typing import Any import h5py import numpy as np -import pandas as pd from tqdm import tqdm -# ------------------------ -# Denormalize function (same as isaaclab2lerobot) -# ------------------------ +ISAACLAB_LIMITS_DEG = [ + (-110.0, 110.0), + (-100.0, 100.0), + (-100.0, 90.0), + (-95.0, 95.0), + (-160.0, 160.0), + (-10.0, 100.0), +] +LEROBOT_LIMITS_DEG = [ + (-100.0, 100.0), + (-100.0, 100.0), + (-100.0, 100.0), + (-100.0, 100.0), + (-100.0, 100.0), + (0.0, 100.0), +] + + +def denormalize_lerobot_to_isaaclab_radians(joint_values_lerobot: np.ndarray) -> np.ndarray: + """ + LeRobot normalized degrees (roughly within LEROBOT_LIMITS_DEG) -> + IsaacLab joint limits (degrees) -> radians. -ISAACLAB_LIMITS = [(-110.0, 110.0), (-100.0, 100.0), (-100.0, 90.0), (-95.0, 95.0), (-160.0, 160.0), (-10.0, 100.0)] -LEROBOT_LIMITS = [(-100.0, 100.0), (-100.0, 100.0), (-100.0, 100.0), (-100.0, 100.0), (-100.0, 100.0), (0.0, 100.0)] + joint_values_lerobot: + shape (T, D) where D >= 6. + If D == 12, treated as bimanual (6+6). + Otherwise, treated as single arm (first 6 processed, others converted deg->rad only). + """ + joint_values_lerobot = np.asarray(joint_values_lerobot, dtype=np.float32) + if joint_values_lerobot.ndim != 2: + raise ValueError(f"Expected 2D array (T,D), got shape {joint_values_lerobot.shape}") + + dimension = joint_values_lerobot.shape[1] + if dimension < 6: + raise ValueError(f"Expected D>=6, got D={dimension}") + + result_deg = joint_values_lerobot.copy() + + if dimension == 12: + # bimanual: apply same 6-dim mapping to left [0:6] and right [6:12] + for arm_offset in (0, 6): + for joint_index in range(6): + isa_min, isa_max = ISAACLAB_LIMITS_DEG[joint_index] + le_min, le_max = LEROBOT_LIMITS_DEG[joint_index] + col = arm_offset + joint_index + result_deg[:, col] = (result_deg[:, col] - le_min) / (le_max - le_min) * (isa_max - isa_min) + isa_min + else: + # single arm (or single arm + extra features): only map the first 6 + for joint_index in range(6): + isa_min, isa_max = ISAACLAB_LIMITS_DEG[joint_index] + le_min, le_max = LEROBOT_LIMITS_DEG[joint_index] + result_deg[:, joint_index] = (result_deg[:, joint_index] - le_min) / (le_max - le_min) * ( + isa_max - isa_min + ) + isa_min -def denormalize_lerobot_to_isaaclab_radians(joint_pos: np.ndarray) -> np.ndarray: - """ - LeRobot normalized degrees in roughly [-100,100] -> IsaacLab joint limits (degrees) -> radians. - joint_pos: (T,6) - """ - result = joint_pos.astype(np.float32).copy() - for i in range(6): - isa_min, isa_max = ISAACLAB_LIMITS[i] - le_min, le_max = LEROBOT_LIMITS[i] - isa_range = isa_max - isa_min - le_range = le_max - le_min - result[:, i] = (result[:, i] - le_min) / le_range * isa_range + isa_min # degrees -> radians - result = result * (np.pi / 180.0) - return result + result_rad = result_deg * (np.pi / 180.0) + return result_rad -# ------------------------ -# HDF5 copy helpers -# ------------------------ +def write_text_dataset(h5_group: h5py.Group, dataset_name: str, text: str) -> None: + """Store UTF-8 text as variable-length string dataset.""" + string_dtype = h5py.string_dtype(encoding="utf-8") + if dataset_name in h5_group: + del h5_group[dataset_name] + h5_group.create_dataset(dataset_name, data=text, dtype=string_dtype) -def copy_attrs(src_obj, dst_obj): - for k, v in src_obj.attrs.items(): - dst_obj.attrs[k] = v +def read_file_as_uint8_bytes(file_path: Path) -> np.ndarray: + data = file_path.read_bytes() + return np.frombuffer(data, dtype=np.uint8) -def copy_group_recursive(src: h5py.Group, dst: h5py.Group, skip: list[str] | None = None): - """ - Recursively copy a group, skipping relative paths in `skip`. - """ - skip = set(skip or []) - - def _rec(s: h5py.Group | h5py.File, d: h5py.Group, rel: str): - copy_attrs(s, d) - for name, item in s.items(): - child_rel = f"{rel}/{name}" if rel else name - if child_rel in skip: - continue - if isinstance(item, h5py.Dataset): - ds = d.create_dataset(name, data=item[()], compression=item.compression) - for ak, av in item.attrs.items(): - ds.attrs[ak] = av - else: - g = d.create_group(name) - _rec(item, g, child_rel) - - _rec(src, dst, "") - - -def ensure_group(root: h5py.Group, path: str) -> h5py.Group: - """ - Ensure group path exists under root. Return the group. - """ - path = path.strip("/") - cur = root - if not path: - return cur - for part in path.split("/"): - if part not in cur: - cur = cur.create_group(part) - else: - cur = cur[part] - if not isinstance(cur, h5py.Group): - raise ValueError(f"Path conflict: {path} has non-group at {part}") - return cur - - -def write_dataset(root: h5py.Group, path: str, data: np.ndarray, compression: str | None = None): +def stack_object_column_to_numpy(column_values: Any) -> tuple[np.ndarray, str | None]: """ - Overwrite dataset at path with given data. + Convert a parquet column to a numpy array. + - If it's already a numpy array -> return directly + - If it's a 1D array/list of per-row arrays/lists -> stack to (T, D) + - If it's scalar-like -> return (T,) or scalar + Returns: (array, note) """ - path = path.strip("/") - parent_path, name = path.rsplit("/", 1) if "/" in path else ("", path) - parent = ensure_group(root, parent_path) - if name in parent: - del parent[name] - parent.create_dataset(name, data=data, compression=compression) - - -# ------------------------ -# Parquet loading -# ------------------------ - -_EP_RE = re.compile(r"episode_(\d+)\.parquet$") - - -def list_episode_parquet_files(parquet_dir: str) -> list[tuple[int, Path]]: - """Return [(episode_id_from_filename, path), ...] sorted by episode_id.""" - p = Path(parquet_dir) - if not p.is_dir(): - raise ValueError(f"--parquet must be a directory when using filename mapping: {parquet_dir}") - files = sorted(p.glob("**/episode_*.parquet")) - items: list[tuple[int, Path]] = [] - for f in files: - m = _EP_RE.search(f.name) - if not m: - continue - items.append((int(m.group(1)), f)) - if not items: - raise FileNotFoundError(f"No episode_*.parquet found under: {parquet_dir}") - items.sort(key=lambda x: x[0]) - return items - - -def load_parquet_input(parquet_path_or_dir: str) -> pd.DataFrame: - p = Path(parquet_path_or_dir) - if p.is_dir(): - files = sorted( - p.glob("**/*.parquet"), - key=lambda x: ( - int(re.search(r"episode_(\d+)", x.name).group(1)) if re.search(r"episode_(\d+)", x.name) else 0 - ), - ) - if not files: - raise FileNotFoundError(f"No parquet files found in dir: {parquet_path_or_dir}") - dfs = [pd.read_parquet(f) for f in files] - df = pd.concat(dfs, ignore_index=True) - else: - if not p.exists(): - raise FileNotFoundError(parquet_path_or_dir) - df = pd.read_parquet(p) - return df - - -def stack_vector_column(df: pd.DataFrame, col: str) -> np.ndarray: - # column entries are list/np.ndarray per row - return np.stack(df[col].to_numpy()) - - -# ------------------------ -# Main conversion -# ------------------------ - - -def convert( # noqa: C901 - parquet_path_or_dir: str, - src_hdf5: str, - out_hdf5: str, - overwrite: bool, - episode_col: str, - sort_col: str | None, - action_col: str, - state_col: str, - timestamp_col: str | None, - write_states: bool, - keep_template_obs: bool, -): - if os.path.exists(out_hdf5): - if overwrite: - os.remove(out_hdf5) - else: - raise FileExistsError(f"Output exists: {out_hdf5} (use --overwrite)") + note = None + + if isinstance(column_values, np.ndarray): + if ( + column_values.dtype == object + and len(column_values) > 0 + and isinstance(column_values[0], (list, tuple, np.ndarray)) + ): + try: + stacked = np.stack([np.asarray(x) for x in column_values], axis=0) + return stacked, note + except Exception as exc: + note = f"object-stack-failed: {exc}" + return np.asarray(column_values), note + return column_values, note + + if isinstance(column_values, (list, tuple)): + if len(column_values) > 0 and isinstance(column_values[0], (list, tuple, np.ndarray)): + try: + stacked = np.stack([np.asarray(x) for x in column_values], axis=0) + return stacked, note + except Exception as exc: + note = f"list-stack-failed: {exc}" + return np.asarray(column_values, dtype=object), note + return np.asarray(column_values), note + + return np.asarray(column_values), note + + +def write_numpy_dataset( + h5_group: h5py.Group, + dataset_name: str, + array: np.ndarray, + compression: str | None, + compression_level: int | None, +) -> None: + """Write/overwrite a dataset. If dtype=object, store JSON strings.""" + if dataset_name in h5_group: + del h5_group[dataset_name] + + # Object dtype: store each row as json string + if array.dtype == object: + string_dtype = h5py.string_dtype(encoding="utf-8") + string_list = [] + for value in array.tolist(): + try: + string_list.append(json.dumps(value, ensure_ascii=False)) + except Exception: + string_list.append(str(value)) + h5_group.create_dataset(dataset_name, data=np.asarray(string_list, dtype=object), dtype=string_dtype) + h5_group[dataset_name].attrs["stored_as"] = "json_string_list" + return + + h5_group.create_dataset( + dataset_name, + data=array, + compression=compression, + compression_opts=compression_level, + shuffle=True if compression else None, + ) - # Use filename-based mapping when input is a directory containing episode_*.parquet - parquet_items = list_episode_parquet_files(parquet_path_or_dir) if Path(parquet_path_or_dir).is_dir() else None - if parquet_items is None: - # fallback: single parquet file mode (still uses episode_index) - df = load_parquet_input(parquet_path_or_dir) +def load_episode_parquet(parquet_path: Path): + import pandas as pd - for need in [episode_col, action_col]: - if need not in df.columns: - raise ValueError(f"Missing required column '{need}' in parquet. columns={list(df.columns)}") + return pd.read_parquet(parquet_path) - if write_states and state_col not in df.columns: - raise ValueError(f"--write_states enabled but missing column '{state_col}'") - if sort_col and sort_col not in df.columns: - raise ValueError(f"--sort_col '{sort_col}' not in columns") +def sort_episode_dataframe_inplace(dataframe, sort_column: str | None) -> str: + """ + Sort dataframe by a chosen column if present. + If sort_column is None, try common ordering keys in priority: + index > frame_index > timestamp + Returns the column used, or "" if not sorted. + """ + candidate_columns = [] + if sort_column: + candidate_columns.append(sort_column) + candidate_columns += ["index", "frame_index", "timestamp"] + + for col in candidate_columns: + if col in dataframe.columns: + dataframe.sort_values(col, ascending=True, inplace=True) + dataframe.reset_index(drop=True, inplace=True) + return col + + return "" + + +def convert_lerobot_folder_to_hdf5( + dataset_directory: str, + output_hdf5_path: str, + include_videos: bool, + compression: str | None, + compression_level: int, + sort_column: str | None, + convert_to_isaaclab_radians: bool, + lerobot_action_column_name: str, + lerobot_state_column_name: str, +): + dataset_root = Path(dataset_directory).expanduser().resolve() + output_file_path = Path(output_hdf5_path).expanduser().resolve() + output_file_path.parent.mkdir(parents=True, exist_ok=True) + + data_directory = dataset_root / "data" + videos_directory = dataset_root / "videos" + meta_directory = dataset_root / "meta" + + episode_parquet_files = sorted(data_directory.glob("chunk-*/episode_*.parquet")) + file_parquet_files = sorted(data_directory.glob("chunk-*/file-*.parquet")) + + if not episode_parquet_files: + if file_parquet_files: + raise RuntimeError( + "Detected LeRobot Dataset v3.0 style (file-xxx.parquet). This script targets v2.1 (episode_*.parquet)." + ) + raise RuntimeError( + "No parquet files found under data/chunk-*/episode_*.parquet. " + f"Please check dataset_directory: {dataset_root}" + ) - if timestamp_col and timestamp_col not in df.columns: - raise ValueError(f"--timestamp_col '{timestamp_col}' not in columns") + with h5py.File(output_file_path, "w") as output_h5: + output_h5.attrs["source_dir"] = str(dataset_root) + output_h5.attrs["layout_version"] = "lerobot_v2.1_episode_based__optionally_converted_to_isaaclab_radians" + output_h5.attrs["created_at"] = dt.datetime.now().isoformat() + output_h5.attrs["convert_to_isaaclab_radians"] = bool(convert_to_isaaclab_radians) + output_h5.attrs["sort_column_requested"] = "" if sort_column is None else str(sort_column) + output_h5.attrs["lerobot_action_column_name"] = lerobot_action_column_name + output_h5.attrs["lerobot_state_column_name"] = lerobot_state_column_name + + # Write meta files + if meta_directory.exists(): + meta_group = output_h5.require_group("meta") + for meta_file_path in sorted(meta_directory.glob("*")): + if not meta_file_path.is_file(): + continue + try: + text = meta_file_path.read_text(encoding="utf-8", errors="ignore") + except Exception: + text = meta_file_path.read_bytes().decode("utf-8", errors="ignore") + write_text_dataset(meta_group, meta_file_path.name, text) + + # Convert parquet episodes + for parquet_path in tqdm(episode_parquet_files, desc="Converting episodes (parquet)"): + chunk_name = parquet_path.parent.name + episode_name = parquet_path.stem + + episode_dataframe = load_episode_parquet(parquet_path) + + used_sort_column = sort_episode_dataframe_inplace(episode_dataframe, sort_column) + + episode_group = output_h5.require_group(f"data/{chunk_name}/{episode_name}") + episode_group.attrs["parquet_relpath"] = str(parquet_path.relative_to(dataset_root)) + episode_group.attrs["num_frames"] = int(len(episode_dataframe)) + episode_group.attrs["sorted_by"] = used_sort_column + + # Write each column + for column_name in episode_dataframe.columns: + column_values = episode_dataframe[column_name].to_numpy() + numpy_array, note = stack_object_column_to_numpy(column_values) + + # Normalize float dtype if possible + if numpy_array.dtype == np.float64: + numpy_array = numpy_array.astype(np.float32) + + # If enabled: convert action/state columns from LeRobot normalized degrees -> IsaacLab radians + if convert_to_isaaclab_radians and numpy_array.ndim == 2: + if column_name in (lerobot_action_column_name, lerobot_state_column_name): + numpy_array = denormalize_lerobot_to_isaaclab_radians(numpy_array) + + safe_dataset_name = column_name.replace("/", "_") + write_numpy_dataset( + h5_group=episode_group, + dataset_name=safe_dataset_name, + array=numpy_array, + compression=compression, + compression_level=compression_level if compression else None, + ) + + if note: + episode_group[safe_dataset_name].attrs["note"] = note + + # Optionally store mp4 bytes in HDF5 + if include_videos and videos_directory.exists(): + video_files = sorted(videos_directory.glob("chunk-*/**/episode_*.mp4")) + videos_group = output_h5.require_group("videos") + + for video_path in tqdm(video_files, desc="Packing videos into HDF5"): + relative_video_path = video_path.relative_to(dataset_root) + parts = relative_video_path.parts + if len(parts) < 4: + continue + + _, chunk_name, camera_name = parts[0], parts[1], parts[2] + episode_filename = parts[-1] + episode_stem = Path(episode_filename).stem + + camera_group = videos_group.require_group(f"{chunk_name}/{camera_name}") + dataset_name = f"{episode_stem}.mp4" + + if dataset_name in camera_group: + del camera_group[dataset_name] + + video_bytes = read_file_as_uint8_bytes(video_path) + camera_group.create_dataset( + dataset_name, + data=video_bytes, + dtype=np.uint8, + compression=compression, + compression_opts=compression_level if compression else None, + shuffle=True if compression else None, + ) + camera_group[dataset_name].attrs["video_relpath"] = str(relative_video_path) + + print(f"[OK] Wrote HDF5 to: {output_file_path}") - ep_ids = sorted(df[episode_col].unique().tolist()) - print(f"[INFO] parquet episodes (from column {episode_col}): {len(ep_ids)}") - else: - print(f"[INFO] parquet files (from filename): {len(parquet_items)}") - - with h5py.File(src_hdf5, "r") as f_src, h5py.File(out_hdf5, "w") as f_out: - if "data" not in f_src: - raise ValueError("Source HDF5 missing /data") - - # copy root metadata except /data - for k, item in f_src.items(): - if k == "data": - continue - if isinstance(item, h5py.Dataset): - ds = f_out.create_dataset(k, data=item[()], compression=item.compression) - for ak, av in item.attrs.items(): - ds.attrs[ak] = av - else: - g = f_out.create_group(k) - copy_group_recursive(item, g) - - g_data_out = f_out.create_group("data") - g_data_src = f_src["data"] - copy_attrs(g_data_src, g_data_out) - - src_ep_names = list(g_data_src.keys()) - if not src_ep_names: - raise ValueError("Source HDF5 has zero episodes under /data") - template_name = src_ep_names[0] - - # We will copy: - # - attrs - # - initial_state group - # optionally: obs/states skeleton (but not the big image arrays unless user wants to keep) - # Note: copying images can explode file size; default keep_template_obs=False. - base_skip = [] - if not keep_template_obs: - # skip heavy video/images and other per-step arrays from template - # we'll only write the minimum required datasets. - base_skip += [ - "obs/front", - "obs/ee_frame_state", - "obs/joint_pos", - "obs/joint_pos_rel", - "obs/joint_pos_target", - "obs/joint_vel", - "obs/joint_vel_rel", - "states", - "actions", - "obs/actions", - "processed_actions", - ] - - if parquet_items is None: - # old behavior: split by episode_col inside one big df - iterator = [(int(ep_id), None) for ep_id in ep_ids] - else: - # new behavior: each file defines one episode, use filename episode id - iterator = parquet_items # [(ep_id, path), ...] - - for ep_id, pq_path in tqdm(iterator, desc="Writing episodes"): - if pq_path is None: - dfe = df[df[episode_col] == ep_id].copy() - else: - dfe = pd.read_parquet(pq_path) - - # basic column checks (per-file for filename mode) - if action_col not in dfe.columns: - raise ValueError(f"Missing required column '{action_col}' in parquet for ep={ep_id}: {pq_path}") - if write_states and state_col not in dfe.columns: - raise ValueError(f"--write_states enabled but missing '{state_col}' for ep={ep_id}: {pq_path}") - if sort_col and sort_col not in dfe.columns: - raise ValueError(f"--sort_col '{sort_col}' not in parquet columns for ep={ep_id}: {pq_path}") - if timestamp_col and timestamp_col not in dfe.columns: - raise ValueError(f"--timestamp_col '{timestamp_col}' not in parquet columns for ep={ep_id}: {pq_path}") - - if sort_col: - dfe = dfe.sort_values(sort_col, ascending=True) - - # actions (T,6) in lerobot normalized degrees - act_raw = stack_vector_column(dfe, action_col).astype(np.float32) - if act_raw.ndim != 2 or act_raw.shape[1] != 6: - raise ValueError(f"Action shape expected (T,6), got {act_raw.shape} for ep={ep_id}") - act = denormalize_lerobot_to_isaaclab_radians(act_raw) - - # states (T,6) optional - st = None - if write_states: - st_raw = stack_vector_column(dfe, state_col).astype(np.float32) - if st_raw.ndim != 2 or st_raw.shape[1] != 6: - raise ValueError(f"State shape expected (T,6), got {st_raw.shape} for ep={ep_id}") - st = denormalize_lerobot_to_isaaclab_radians(st_raw) - - ts = None - if timestamp_col: - ts = dfe[timestamp_col].to_numpy(dtype=np.int64) - - # choose a source episode to copy initial_state from (now: use ep_id from filename) - src_name = f"demo_{int(ep_id)}" - if src_name not in g_data_src: - # fallback map by index - if int(ep_id) < len(src_ep_names): - src_name = src_ep_names[int(ep_id)] - else: - src_name = template_name - - g_src_ep = g_data_src[src_name] - dst_name = f"demo_{int(ep_id)}" - g_dst_ep = g_data_out.create_group(dst_name) - - # copy template stuff (attrs + initial_state + maybe skeleton) - copy_attrs(g_src_ep, g_dst_ep) - if "initial_state" in g_src_ep: - g_init = g_dst_ep.create_group("initial_state") - copy_group_recursive(g_src_ep["initial_state"], g_init) - else: - raise ValueError(f"Source episode {src_name} missing /initial_state") - - for grp_name in ["obs", "states"]: - if grp_name in g_src_ep and keep_template_obs: - g = g_dst_ep.create_group(grp_name) - copy_group_recursive(g_src_ep[grp_name], g) - - # Write actions into required places - write_dataset(g_dst_ep, "actions", act, compression=None) - write_dataset(g_dst_ep, "obs/actions", act, compression=None) - write_dataset(g_dst_ep, "processed_actions", act, compression=None) - - if write_states and st is not None: - write_dataset(g_dst_ep, "states/articulation/robot/joint_position", st, compression=None) - - if ts is not None: - write_dataset(g_dst_ep, "timestamps", ts, compression=None) - - g_dst_ep.attrs["num_steps"] = int(act.shape[0]) - g_dst_ep.attrs["action_dim"] = int(act.shape[1]) - g_dst_ep.attrs["converted_from_src_episode"] = src_name - g_dst_ep.attrs["idm_source"] = str(Path(parquet_path_or_dir).resolve()) - if pq_path is not None: - g_dst_ep.attrs["idm_parquet_file"] = pq_path.name - - print(f"[DONE] Output written: {out_hdf5}") - print("Run replay with:") - print(f" --dataset_file {out_hdf5}") - - -def build_parser(): - p = argparse.ArgumentParser("IDM parquet -> LeIsaac HDF5 converter") - p.add_argument("--parquet", type=str, required=True, help="Parquet file OR directory containing parquet(s).") - p.add_argument("--src_hdf5", type=str, required=True, help="Source teleop HDF5 for initial_state/template.") - p.add_argument("--out_hdf5", type=str, required=True, help="Output HDF5 for replay.") - p.add_argument("--overwrite", action="store_true") - - p.add_argument("--episode_col", type=str, default="episode_index") - p.add_argument("--sort_col", type=str, default="index", help="Order within episode. default: index") - p.add_argument("--action_col", type=str, default="action") - p.add_argument("--state_col", type=str, default="observation.state") - p.add_argument("--timestamp_col", type=str, default="timestamp") - - p.add_argument( - "--write_states", + +def main(): + argument_parser = argparse.ArgumentParser( + description="Convert local LeRobot dataset folder (v2.1) to a single HDF5 file." + ) + argument_parser.add_argument("--lerobot_dir", type=str, required=True, help="Local LeRobot dataset folder path") + argument_parser.add_argument("--output_hdf5", type=str, required=True, help="Output HDF5 file path") + + argument_parser.add_argument( + "--include_videos", action="store_true", - help="Also write /states/articulation/robot/joint_position from observation.state", + help="Store mp4 bytes into HDF5 (may be large)", ) - p.add_argument( - "--keep_template_obs", + argument_parser.add_argument( + "--no_compression", action="store_true", - help="Copy template obs/states arrays (including images) from source (can be huge). Default off.", + help="Disable HDF5 compression", + ) + argument_parser.add_argument( + "--compression_level", + type=int, + default=4, + help="gzip compression level (1-9)", + ) + argument_parser.add_argument( + "--sort_column", + type=str, + default=None, + help=( + "Sort frames inside each episode by this column if present. " + "If not provided, tries: index > frame_index > timestamp." + ), + ) + argument_parser.add_argument( + "--no_convert_to_isaaclab_radians", + action="store_true", + help="Disable conversion of LeRobot normalized degrees to IsaacLab radians (default: conversion is ENABLED).", + ) + argument_parser.add_argument( + "--lerobot_action_column_name", + type=str, + default="action", + help="Column name in parquet for action vectors (default: action)", + ) + argument_parser.add_argument( + "--lerobot_state_column_name", + type=str, + default="observation.state", + help="Column name in parquet for state vectors (default: observation.state)", ) - return p + args = argument_parser.parse_args() -def main(): - args = build_parser().parse_args() - convert( - parquet_path_or_dir=args.parquet, - src_hdf5=args.src_hdf5, - out_hdf5=args.out_hdf5, - overwrite=args.overwrite, - episode_col=args.episode_col, - sort_col=args.sort_col, - action_col=args.action_col, - state_col=args.state_col, - timestamp_col=args.timestamp_col, - write_states=args.write_states, - keep_template_obs=args.keep_template_obs, + compression = None if args.no_compression else "gzip" + + convert_lerobot_folder_to_hdf5( + dataset_directory=args.lerobot_dir, + output_hdf5_path=args.output_hdf5, + include_videos=args.include_videos, + compression=compression, + compression_level=args.compression_level, + sort_column=args.sort_column, + convert_to_isaaclab_radians=not args.no_convert_to_isaaclab_radians, + lerobot_action_column_name=args.lerobot_action_column_name, + lerobot_state_column_name=args.lerobot_state_column_name, ) diff --git a/scripts/tutorials/cosmos_merge.py b/scripts/tutorials/cosmos_merge.py new file mode 100644 index 00000000..c51d6b69 --- /dev/null +++ b/scripts/tutorials/cosmos_merge.py @@ -0,0 +1,481 @@ +""" +Merge a LeRobot-converted HDF5 with a source/template LeIsaac HDF5. + +Inputs: + - lerobot_converted_hdf5_path: + HDF5 produced from local LeRobot dataset conversion (episode parquet columns stored as datasets) + - source_template_hdf5_path: + LeIsaac HDF5 providing initial_state + episode structure template + metadata +Output: + - output_merged_hdf5_path: + New replayable HDF5 that follows source structure, but actions/states come from the LeRobot-converted HDF5. + +Default behavior: + - Copy everything from source root EXCEPT '/data' into output root. + - Choose a template episode from source '/data' (first one, or specified). + - For each lerobot episode under /data/chunk-*/episode_*: + - Copy template episode group into output '/data/' + - Overwrite: + /actions + /obs/actions (if exists) + /processed_actions + - Optionally overwrite a user-specified joint/state dataset path with LeRobot observation.state +""" + +import argparse +from contextlib import suppress +from pathlib import Path + +import h5py +import numpy as np +from tqdm import tqdm + + +def list_template_episode_names(source_hdf5_file: h5py.File) -> list[str]: + """List episode group names under source '/data'.""" + if "data" not in source_hdf5_file: + return [] + return sorted(list(source_hdf5_file["data"].keys())) + + +def ensure_parent_groups_exist(episode_root_group: h5py.Group, dataset_path: str) -> h5py.Group: + """ + Ensure all parent groups exist for a dataset path like + 'states/articulation/robot/joint_position'. Returns the parent group. + """ + path_parts = [part for part in dataset_path.split("/") if part] + if len(path_parts) <= 1: + return episode_root_group + + parent_parts = path_parts[:-1] + current_group = episode_root_group + for group_name in parent_parts: + current_group = current_group.require_group(group_name) + return current_group + + +def copy_hdf5_attributes(source_object, destination_object) -> None: + """Copy all HDF5 attributes from source object to destination object.""" + for attribute_key, attribute_value in source_object.attrs.items(): + destination_object.attrs[attribute_key] = attribute_value + + +def copy_source_root_except_data_group(source_hdf5_file: h5py.File, output_hdf5_file: h5py.File) -> None: + """ + Copy everything at the root of source HDF5 into output HDF5 except '/data'. + This preserves metadata, configs, and other root-level structures. + """ + copy_hdf5_attributes(source_hdf5_file, output_hdf5_file) + for root_item_name in source_hdf5_file.keys(): + if root_item_name == "data": + continue + source_hdf5_file.copy(root_item_name, output_hdf5_file, name=root_item_name) + + +def copy_template_episode_group( + source_hdf5_file: h5py.File, + output_hdf5_file: h5py.File, + template_episode_name: str, + output_episode_name: str, +) -> None: + """ + Copy '/data/' from source into output at '/data/'. + """ + output_data_group = output_hdf5_file.require_group("data") + + if output_episode_name in output_data_group: + del output_data_group[output_episode_name] + + source_hdf5_file["data"].copy(template_episode_name, output_data_group, name=output_episode_name) + + # Preserve attrs on /data itself if present + copy_hdf5_attributes(source_hdf5_file["data"], output_data_group) + + +def write_or_overwrite_hdf5_dataset( + parent_group: h5py.Group, + dataset_name: str, + dataset_data: np.ndarray, + template_dataset: h5py.Dataset | None = None, +) -> None: + """ + Overwrite dataset 'dataset_name' inside 'parent_group' with numpy 'dataset_data'. + If 'template_dataset' is provided, mirror its dtype/chunking/compression settings when possible. + """ + if dataset_name in parent_group: + del parent_group[dataset_name] + + create_dataset_kwargs = {} + + if template_dataset is not None: + # Mirror chunking/compression options + if template_dataset.chunks is not None: + create_dataset_kwargs["chunks"] = template_dataset.chunks + if template_dataset.compression is not None: + create_dataset_kwargs["compression"] = template_dataset.compression + if template_dataset.compression_opts is not None: + create_dataset_kwargs["compression_opts"] = template_dataset.compression_opts + if template_dataset.shuffle: + create_dataset_kwargs["shuffle"] = True + + # Try to match dtype (if possible) + with suppress(Exception): + dataset_data = dataset_data.astype(template_dataset.dtype, copy=False) + + parent_group.create_dataset(dataset_name, data=dataset_data, **create_dataset_kwargs) + + +def iterate_lerobot_episode_groups( + lerobot_hdf5_file: h5py.File, +) -> list[tuple[str, h5py.Group]]: + """ + Iterate episodes from a lerobot-converted HDF5: + + Expected layout: + /data/chunk-000/episode_000000 + /data/chunk-000/episode_000001 + ... + + Returns list of: + (episode_identifier_string, episode_group) + where episode_identifier_string is like 'chunk-000/episode_000000'. + """ + episodes: list[tuple[str, h5py.Group]] = [] + + if "data" not in lerobot_hdf5_file: + return episodes + + for chunk_group_name in sorted(lerobot_hdf5_file["data"].keys()): + chunk_group = lerobot_hdf5_file["data"][chunk_group_name] + if not isinstance(chunk_group, h5py.Group): + continue + + for episode_group_name in sorted(chunk_group.keys()): + if not episode_group_name.startswith("episode_"): + continue + + episode_group = chunk_group[episode_group_name] + if isinstance(episode_group, h5py.Group): + episode_identifier = f"{chunk_group_name}/{episode_group_name}" + episodes.append((episode_identifier, episode_group)) + + return episodes + + +def read_action_and_state_arrays_from_lerobot_group( + lerobot_episode_group: h5py.Group, + lerobot_action_dataset_key: str, + lerobot_state_dataset_key: str, +) -> tuple[np.ndarray, np.ndarray | None]: + """ + Read action and state arrays from a lerobot episode group. + Supports common naming variants for 'observation.state' due to prior sanitization. + + Returns: + actions_array: np.ndarray + state_array: Optional[np.ndarray] (None if not found) + """ + if lerobot_action_dataset_key not in lerobot_episode_group: + raise KeyError( + f"Missing action dataset '{lerobot_action_dataset_key}' " + f"in lerobot episode group: {lerobot_episode_group.name}" + ) + + actions_array = np.array(lerobot_episode_group[lerobot_action_dataset_key]) + + state_candidate_keys = [lerobot_state_dataset_key] + + # Accept common sanitized variants between 'observation.state' and 'observation_state' + if lerobot_state_dataset_key == "observation.state": + state_candidate_keys.append("observation_state") + elif lerobot_state_dataset_key == "observation_state": + state_candidate_keys.append("observation.state") + + state_array = None + for state_key in state_candidate_keys: + if state_key in lerobot_episode_group: + state_array = np.array(lerobot_episode_group[state_key]) + break + + return actions_array, state_array + + +def merge_lerobot_converted_hdf5_with_source_template( + lerobot_converted_hdf5_path: str, + source_template_hdf5_path: str, + output_merged_hdf5_path: str, + source_template_episode_name: str | None, + output_episode_name_prefix: str, + lerobot_action_dataset_key: str, + lerobot_state_dataset_key: str, + lerobot_timestamp_dataset_key: str | None, + write_state_to_output_hdf5: bool, + output_state_target_dataset_path: str | None, + skip_first_n_lerobot_episodes: int, + max_number_of_episodes_to_merge: int | None, +) -> None: + lerobot_converted_hdf5_path = str(Path(lerobot_converted_hdf5_path).expanduser().resolve()) + source_template_hdf5_path = str(Path(source_template_hdf5_path).expanduser().resolve()) + output_merged_hdf5_path = str(Path(output_merged_hdf5_path).expanduser().resolve()) + Path(output_merged_hdf5_path).parent.mkdir(parents=True, exist_ok=True) + + with h5py.File(source_template_hdf5_path, "r") as source_hdf5_file, h5py.File( + lerobot_converted_hdf5_path, "r" + ) as lerobot_hdf5_file, h5py.File(output_merged_hdf5_path, "w") as output_hdf5_file: + + # 1) Copy root-level template/metadata except '/data' + copy_source_root_except_data_group(source_hdf5_file, output_hdf5_file) + + # 2) Choose a template episode in source '/data' + available_template_episode_names = list_template_episode_names(source_hdf5_file) + if not available_template_episode_names: + raise RuntimeError("source_template_hdf5 has no '/data/' groups to use as template.") + + if source_template_episode_name is None: + source_template_episode_name = available_template_episode_names[0] + else: + if source_template_episode_name not in source_hdf5_file["data"]: + preview = available_template_episode_names[:10] + raise RuntimeError( + f"template_episode '{source_template_episode_name}' not found in source '/data'. " + f"Available examples: {preview}{'...' if len(available_template_episode_names) > 10 else ''}" + ) + + # 3) Iterate lerobot episodes + lerobot_episode_list = iterate_lerobot_episode_groups(lerobot_hdf5_file) + if not lerobot_episode_list: + raise RuntimeError("lerobot_converted_hdf5 has no episodes under '/data/chunk-*/episode_*'.") + + # Apply slicing + lerobot_episode_list = lerobot_episode_list[skip_first_n_lerobot_episodes:] + if max_number_of_episodes_to_merge is not None: + lerobot_episode_list = lerobot_episode_list[:max_number_of_episodes_to_merge] + + # Ensure /data exists in output + output_hdf5_file.require_group("data") + + for output_episode_index, (lerobot_episode_identifier, lerobot_episode_group) in enumerate( + tqdm(lerobot_episode_list, desc="Merging episodes") + ): + output_episode_name = f"{output_episode_name_prefix}{output_episode_index}" + + # 3.1 Copy template episode group to output + copy_template_episode_group( + source_hdf5_file=source_hdf5_file, + output_hdf5_file=output_hdf5_file, + template_episode_name=source_template_episode_name, + output_episode_name=output_episode_name, + ) + output_episode_group = output_hdf5_file["data"][output_episode_name] + + # 3.2 Read lerobot action/state arrays + actions_array, state_array = read_action_and_state_arrays_from_lerobot_group( + lerobot_episode_group=lerobot_episode_group, + lerobot_action_dataset_key=lerobot_action_dataset_key, + lerobot_state_dataset_key=lerobot_state_dataset_key, + ) + + # 3.3 Overwrite /actions + template_actions_dataset = output_episode_group["actions"] if "actions" in output_episode_group else None + if isinstance(template_actions_dataset, h5py.Group): + template_actions_dataset = None + + write_or_overwrite_hdf5_dataset( + parent_group=output_episode_group, + dataset_name="actions", + dataset_data=actions_array, + template_dataset=( + template_actions_dataset if isinstance(template_actions_dataset, h5py.Dataset) else None + ), + ) + + # 3.4 Overwrite /obs/actions if exists (or create) + if "obs" in output_episode_group and isinstance(output_episode_group["obs"], h5py.Group): + output_obs_group = output_episode_group["obs"] + else: + output_obs_group = output_episode_group.require_group("obs") + + template_obs_actions_dataset = output_obs_group["actions"] if "actions" in output_obs_group else None + write_or_overwrite_hdf5_dataset( + parent_group=output_obs_group, + dataset_name="actions", + dataset_data=actions_array, + template_dataset=( + template_obs_actions_dataset if isinstance(template_obs_actions_dataset, h5py.Dataset) else None + ), + ) + + # 3.5 Overwrite /processed_actions (create if missing) + template_processed_actions_dataset = ( + output_episode_group["processed_actions"] + if "processed_actions" in output_episode_group + and isinstance(output_episode_group["processed_actions"], h5py.Dataset) + else None + ) + write_or_overwrite_hdf5_dataset( + parent_group=output_episode_group, + dataset_name="processed_actions", + dataset_data=actions_array, + template_dataset=template_processed_actions_dataset, + ) + + # 3.6 Optional: overwrite timestamps + if lerobot_timestamp_dataset_key is not None and lerobot_timestamp_dataset_key in lerobot_episode_group: + timestamps_array = np.array(lerobot_episode_group[lerobot_timestamp_dataset_key]) + template_timestamps_dataset = ( + output_episode_group["timestamps"] + if "timestamps" in output_episode_group + and isinstance(output_episode_group["timestamps"], h5py.Dataset) + else None + ) + write_or_overwrite_hdf5_dataset( + parent_group=output_episode_group, + dataset_name="timestamps", + dataset_data=timestamps_array, + template_dataset=template_timestamps_dataset, + ) + + # 3.7 Optional: write state/joint positions into a target dataset path + if write_state_to_output_hdf5 and state_array is not None and output_state_target_dataset_path is not None: + parent_group = ensure_parent_groups_exist(output_episode_group, output_state_target_dataset_path) + target_dataset_name = output_state_target_dataset_path.split("/")[-1] + + # Mirror template dataset properties if it exists + template_state_dataset = None + if output_state_target_dataset_path in output_episode_group and isinstance( + output_episode_group[output_state_target_dataset_path], h5py.Dataset + ): + template_state_dataset = output_episode_group[output_state_target_dataset_path] + + write_or_overwrite_hdf5_dataset( + parent_group=parent_group, + dataset_name=target_dataset_name, + dataset_data=state_array, + template_dataset=template_state_dataset, + ) + + # Provenance + output_episode_group.attrs["lerobot_episode_identifier"] = lerobot_episode_identifier + output_episode_group.attrs["source_template_episode_name"] = source_template_episode_name + + # Root provenance + output_hdf5_file.attrs["merged_from_lerobot_converted_hdf5"] = lerobot_converted_hdf5_path + output_hdf5_file.attrs["merged_from_source_template_hdf5"] = source_template_hdf5_path + + print(f"[OK] Wrote merged replayable HDF5: {output_merged_hdf5_path}") + + +def main(): + argument_parser = argparse.ArgumentParser( + description="Merge lerobot-converted HDF5 with source IsaacLab/LeIsaac template HDF5." + ) + + argument_parser.add_argument( + "--lerobot_hdf5", + type=str, + required=True, + help="Path to HDF5 converted from local LeRobot dataset", + ) + argument_parser.add_argument( + "--source_hdf5", + type=str, + required=True, + help="Path to source/template IsaacLab/LeIsaac HDF5 (provides initial_state + structure template)", + ) + argument_parser.add_argument( + "--output_hdf5", + type=str, + required=True, + help="Path to output merged HDF5", + ) + + argument_parser.add_argument( + "--template_episode", + type=str, + default=None, + help="Episode name under source '/data' to use as template (default: first episode found)", + ) + + argument_parser.add_argument( + "--output_episode_prefix", + type=str, + default="demo_", + help="Output episode name prefix under '/data' (default: demo)", + ) + + # Keys in lerobot episode group + argument_parser.add_argument( + "--lerobot_action_key", + type=str, + default="action", + help="Dataset key inside lerobot episode group for actions (default: action)", + ) + argument_parser.add_argument( + "--lerobot_state_key", + type=str, + default="observation.state", + help=( + "Dataset key inside lerobot episode group for state/joint pos (default: observation.state; also tries" + " observation_state)" + ), + ) + argument_parser.add_argument( + "--lerobot_timestamp_key", + type=str, + default=None, + help="Optional dataset key inside lerobot episode group for timestamps (e.g. timestamp)", + ) + + # Where to write state in output episode + argument_parser.add_argument( + "--write_state", + action="store_true", + help="If set, write lerobot state/joint positions into output at --output_state_target_path", + ) + argument_parser.add_argument( + "--output_state_target_path", + type=str, + default=None, + help=( + "Target dataset path inside each output episode to store state " + "(e.g. 'obs/joint_pos' or 'states/articulation/robot/joint_position')" + ), + ) + + # Episode slicing + argument_parser.add_argument( + "--skip_first_n_episodes", + type=int, + default=0, + help="Skip first N lerobot episodes (default: 0)", + ) + argument_parser.add_argument( + "--max_episodes", + type=int, + default=None, + help="Merge at most N episodes (default: all)", + ) + + parsed_arguments = argument_parser.parse_args() + + if parsed_arguments.write_state and not parsed_arguments.output_state_target_path: + raise SystemExit("--write_state requires --output_state_target_path to be set") + + merge_lerobot_converted_hdf5_with_source_template( + lerobot_converted_hdf5_path=parsed_arguments.lerobot_hdf5, + source_template_hdf5_path=parsed_arguments.source_hdf5, + output_merged_hdf5_path=parsed_arguments.output_hdf5, + source_template_episode_name=parsed_arguments.template_episode, + output_episode_name_prefix=parsed_arguments.output_episode_prefix, + lerobot_action_dataset_key=parsed_arguments.lerobot_action_key, + lerobot_state_dataset_key=parsed_arguments.lerobot_state_key, + lerobot_timestamp_dataset_key=parsed_arguments.lerobot_timestamp_key, + write_state_to_output_hdf5=parsed_arguments.write_state, + output_state_target_dataset_path=parsed_arguments.output_state_target_path, + skip_first_n_lerobot_episodes=parsed_arguments.skip_first_n_episodes, + max_number_of_episodes_to_merge=parsed_arguments.max_episodes, + ) + + +if __name__ == "__main__": + main() From ac5fc61af4c4c8a95da2ac275bc2181d085d5a7a Mon Sep 17 00:00:00 2001 From: chuchen0210 <167320956+chuchen0210@users.noreply.github.com> Date: Tue, 13 Jan 2026 18:45:08 +0800 Subject: [PATCH 7/9] Refine lerobot2isaaclab converter and cosmos merge utility --- scripts/convert/lerobot2isaaclab.py | 207 ++++++++++++++++++++++++---- scripts/tutorials/cosmos_merge.py | 38 ++--- 2 files changed, 199 insertions(+), 46 deletions(-) diff --git a/scripts/convert/lerobot2isaaclab.py b/scripts/convert/lerobot2isaaclab.py index 8ae7088c..ac2be622 100644 --- a/scripts/convert/lerobot2isaaclab.py +++ b/scripts/convert/lerobot2isaaclab.py @@ -25,11 +25,20 @@ import argparse import datetime as dt import json +import os +from contextlib import suppress from pathlib import Path from typing import Any +# Disable HDF5 file locking to avoid BlockingIOError on some filesystems +os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" + import h5py import numpy as np +import pandas as pd +import torch +from lerobot.datasets.compute_stats import aggregate_stats, compute_episode_stats +from lerobot.datasets.lerobot_dataset import LeRobotDataset from tqdm import tqdm ISAACLAB_LIMITS_DEG = [ @@ -167,6 +176,19 @@ def write_numpy_dataset( h5_group[dataset_name].attrs["stored_as"] = "json_string_list" return + # Numpy unicode string dtype (e.g. str: return "" +def load_episode_data(dataset: LeRobotDataset, episode_idx: int) -> dict[str, Any]: + """ + Manually load all frames for a specific episode from LeRobotDataset. + Reconstructs the dictionary of arrays (T, D) for the episode. + """ + start_frame = dataset.episode_data_index["from"][episode_idx].item() + end_frame = dataset.episode_data_index["to"][episode_idx].item() + + # Pre-allocate lists for each key + # We peek at the first frame to know the keys + first_frame = dataset[start_frame] + episode_dict = {k: [] for k in first_frame.keys()} + + for i in range(start_frame, end_frame): + frame = dataset[i] + for k, v in frame.items(): + episode_dict[k].append(v) + + # Stack into arrays/tensors + # dataset[i] returns torch tensors or numpy arrays depending on backend, usually tensors + collapsed_dict = {} + for k, v_list in episode_dict.items(): + if len(v_list) > 0: + if isinstance(v_list[0], torch.Tensor): + collapsed_dict[k] = torch.stack(v_list) + else: + collapsed_dict[k] = np.stack(v_list) + else: + collapsed_dict[k] = np.array([]) + + return collapsed_dict + + +def generate_stats_if_missing(dataset_root: Path): + """ + Check if meta/episodes_stats.jsonl exists. If not, generate it and meta/stats.json + using LeRobot compute_stats utilities. + """ + episodes_stats_path = dataset_root / "meta/episodes_stats.jsonl" + stats_path = dataset_root / "meta/stats.json" + + if episodes_stats_path.exists() and stats_path.exists(): + print(f"Stats files found in {dataset_root / 'meta'}, skipping generation.") + return + + print("Missing stats files! Generating them now (this may take a while)...") + + info_path = dataset_root / "meta/info.json" + if not info_path.exists(): + raise FileNotFoundError(f"meta/info.json not found in {dataset_root}") + + with open(info_path) as f: + info = json.load(f) + features = info.get("features", {}) + + if not features: + raise ValueError("No 'features' key found in info.json") + + data_dir = dataset_root / "data" + parquet_files = sorted(data_dir.glob("chunk-*/episode_*.parquet")) + + if not parquet_files: + raise FileNotFoundError(f"No parquet files found in {data_dir}") + + print(f"Found {len(parquet_files)} episodes to process for stats generation.") + + all_stats = [] + episode_stats_list = [] + + for pq_path in tqdm(parquet_files, desc="Computing stats"): + # Load episode + df = pd.read_parquet(pq_path) + + # Convert to dictionary of arrays + episode_data = {} + for col in df.columns: + if col not in features: + continue + + val = df[col].values + if val.dtype == object and len(val) > 0 and isinstance(val[0], (list, tuple, np.ndarray)): + with suppress(Exception): + val = np.stack(val) + episode_data[col] = val + + try: + stats = compute_episode_stats(episode_data, features) + all_stats.append(stats) + + # Prepare for per-episode jsonl + stats_with_index = stats.copy() + episode_index = int(pq_path.stem.split("_")[-1]) + stats_with_index["episode_index"] = episode_index + episode_stats_list.append(stats_with_index) + + except Exception as e: + print(f"Error computing stats for {pq_path}: {e}") + raise e + + print("Aggregating statistics...") + aggregated = aggregate_stats(all_stats) + + def numpy_converter(obj): + if isinstance(obj, np.integer): + return int(obj) + elif isinstance(obj, np.floating): + return float(obj) + elif isinstance(obj, np.ndarray): + return obj.tolist() + return obj + + # Save aggregated stats + with open(stats_path, "w") as f: + json.dump(aggregated, f, indent=4, default=numpy_converter) + print(f"Saved aggregated stats to {stats_path}") + + # Save per-episode stats + with open(episodes_stats_path, "w") as f: + for stat in episode_stats_list: + clean_stat = json.loads(json.dumps(stat, default=numpy_converter)) + output_item = { + "episode_index": clean_stat["episode_index"], + "stats": {k: v for k, v in clean_stat.items() if k != "episode_index"}, + } + f.write(json.dumps(output_item) + "\n") + + print(f"Saved per-episode stats to {episodes_stats_path}") + + def convert_lerobot_folder_to_hdf5( dataset_directory: str, output_hdf5_path: str, @@ -215,31 +366,22 @@ def convert_lerobot_folder_to_hdf5( lerobot_state_column_name: str, ): dataset_root = Path(dataset_directory).expanduser().resolve() + + # Check and generate stats if missing + generate_stats_if_missing(dataset_root) + output_file_path = Path(output_hdf5_path).expanduser().resolve() output_file_path.parent.mkdir(parents=True, exist_ok=True) - data_directory = dataset_root / "data" videos_directory = dataset_root / "videos" meta_directory = dataset_root / "meta" - episode_parquet_files = sorted(data_directory.glob("chunk-*/episode_*.parquet")) - file_parquet_files = sorted(data_directory.glob("chunk-*/file-*.parquet")) - - if not episode_parquet_files: - if file_parquet_files: - raise RuntimeError( - "Detected LeRobot Dataset v3.0 style (file-xxx.parquet). This script targets v2.1 (episode_*.parquet)." - ) - raise RuntimeError( - "No parquet files found under data/chunk-*/episode_*.parquet. " - f"Please check dataset_directory: {dataset_root}" - ) - with h5py.File(output_file_path, "w") as output_h5: output_h5.attrs["source_dir"] = str(dataset_root) output_h5.attrs["layout_version"] = "lerobot_v2.1_episode_based__optionally_converted_to_isaaclab_radians" output_h5.attrs["created_at"] = dt.datetime.now().isoformat() output_h5.attrs["convert_to_isaaclab_radians"] = bool(convert_to_isaaclab_radians) + # Note: sort_column is not used when using LeRobotDataset as it handles loading output_h5.attrs["sort_column_requested"] = "" if sort_column is None else str(sort_column) output_h5.attrs["lerobot_action_column_name"] = lerobot_action_column_name output_h5.attrs["lerobot_state_column_name"] = lerobot_state_column_name @@ -256,24 +398,32 @@ def convert_lerobot_folder_to_hdf5( text = meta_file_path.read_bytes().decode("utf-8", errors="ignore") write_text_dataset(meta_group, meta_file_path.name, text) - # Convert parquet episodes - for parquet_path in tqdm(episode_parquet_files, desc="Converting episodes (parquet)"): - chunk_name = parquet_path.parent.name - episode_name = parquet_path.stem + # Initialize LeRobotDataset + # Note: repo_id is required but for local usage we can just pass the name. + # root must be the direct path to the dataset folder. + dataset = LeRobotDataset(repo_id=dataset_root.name, root=dataset_root) - episode_dataframe = load_episode_parquet(parquet_path) + # Convert episodes using LeRobotDataset + for episode_idx in tqdm(range(dataset.num_episodes), desc="Converting episodes"): + # Load episode data manually + episode_data = load_episode_data(dataset, episode_idx) + episode_name = f"episode_{episode_idx:06d}" - used_sort_column = sort_episode_dataframe_inplace(episode_dataframe, sort_column) + # NOTE: Flattened structure to match isaaclab2lerobot expectations: /data/ + episode_group = output_h5.require_group(f"data/{episode_name}") + # episode_group.attrs["parquet_relpath"] = str(parquet_path.relative_to(dataset_root)) # Not applicable - episode_group = output_h5.require_group(f"data/{chunk_name}/{episode_name}") - episode_group.attrs["parquet_relpath"] = str(parquet_path.relative_to(dataset_root)) - episode_group.attrs["num_frames"] = int(len(episode_dataframe)) - episode_group.attrs["sorted_by"] = used_sort_column + # Use 'index' column for length if available, else use first column length + num_frames = len(episode_data[next(iter(episode_data))]) + episode_group.attrs["num_frames"] = int(num_frames) + # episode_group.attrs["sorted_by"] = used_sort_column # LeRobotDataset returns sorted frames # Write each column - for column_name in episode_dataframe.columns: - column_values = episode_dataframe[column_name].to_numpy() - numpy_array, note = stack_object_column_to_numpy(column_values) + for column_name, column_tensor in episode_data.items(): + if isinstance(column_tensor, torch.Tensor): + numpy_array = column_tensor.numpy() + else: + numpy_array = np.array(column_tensor) # Normalize float dtype if possible if numpy_array.dtype == np.float64: @@ -293,8 +443,7 @@ def convert_lerobot_folder_to_hdf5( compression_level=compression_level if compression else None, ) - if note: - episode_group[safe_dataset_name].attrs["note"] = note + # Note: `note` logic from stack_object_column_to_numpy is skipped as LeRobotDataset returns tensors # Optionally store mp4 bytes in HDF5 if include_videos and videos_directory.exists(): diff --git a/scripts/tutorials/cosmos_merge.py b/scripts/tutorials/cosmos_merge.py index c51d6b69..79dc002a 100644 --- a/scripts/tutorials/cosmos_merge.py +++ b/scripts/tutorials/cosmos_merge.py @@ -131,33 +131,35 @@ def iterate_lerobot_episode_groups( """ Iterate episodes from a lerobot-converted HDF5: - Expected layout: - /data/chunk-000/episode_000000 - /data/chunk-000/episode_000001 - ... + Supports layouts: + 1) /data/chunk-000/episode_000000 (nested) + 2) /data/episode_000000 (flat) Returns list of: (episode_identifier_string, episode_group) - where episode_identifier_string is like 'chunk-000/episode_000000'. + where episode_identifier_string is like 'chunk-000/episode_000000' or 'episode_000000'. """ episodes: list[tuple[str, h5py.Group]] = [] if "data" not in lerobot_hdf5_file: return episodes - for chunk_group_name in sorted(lerobot_hdf5_file["data"].keys()): - chunk_group = lerobot_hdf5_file["data"][chunk_group_name] - if not isinstance(chunk_group, h5py.Group): - continue + for key in sorted(lerobot_hdf5_file["data"].keys()): + item = lerobot_hdf5_file["data"][key] - for episode_group_name in sorted(chunk_group.keys()): - if not episode_group_name.startswith("episode_"): - continue + # Case 1: Flat structure /data/episode_XXXXXX + if key.startswith("episode_") and isinstance(item, h5py.Group): + episodes.append((key, item)) + continue - episode_group = chunk_group[episode_group_name] - if isinstance(episode_group, h5py.Group): - episode_identifier = f"{chunk_group_name}/{episode_group_name}" - episodes.append((episode_identifier, episode_group)) + # Case 2: Nested structure /data/chunk-XXX/episode_XXXXXX + # We assume any other group might contain episodes (e.g. chunk-000) + if isinstance(item, h5py.Group): + for sub_key in sorted(item.keys()): + if sub_key.startswith("episode_"): + sub_item = item[sub_key] + if isinstance(sub_item, h5py.Group): + episodes.append((f"{key}/{sub_key}", sub_item)) return episodes @@ -244,7 +246,9 @@ def merge_lerobot_converted_hdf5_with_source_template( # 3) Iterate lerobot episodes lerobot_episode_list = iterate_lerobot_episode_groups(lerobot_hdf5_file) if not lerobot_episode_list: - raise RuntimeError("lerobot_converted_hdf5 has no episodes under '/data/chunk-*/episode_*'.") + raise RuntimeError( + "lerobot_converted_hdf5 has no episodes under '/data/chunk-*/episode_*' or '/data/episode_*'." + ) # Apply slicing lerobot_episode_list = lerobot_episode_list[skip_first_n_lerobot_episodes:] From 45e31c519d339bdb454f303aabf054614fbd68bd Mon Sep 17 00:00:00 2001 From: chuchen0210 <167320956+chuchen0210@users.noreply.github.com> Date: Tue, 13 Jan 2026 18:59:02 +0800 Subject: [PATCH 8/9] Apply isort import ordering --- scripts/environments/list_envs.py | 3 ++- scripts/environments/teleoperation/replay.py | 3 ++- scripts/evaluation/policy_inference.py | 3 ++- scripts/mimic/annotate_demos.py | 3 ++- scripts/mimic/generate_dataset.py | 3 ++- 5 files changed, 10 insertions(+), 5 deletions(-) diff --git a/scripts/environments/list_envs.py b/scripts/environments/list_envs.py index a43e953d..e2b9e813 100644 --- a/scripts/environments/list_envs.py +++ b/scripts/environments/list_envs.py @@ -24,9 +24,10 @@ """Rest everything follows.""" import gymnasium as gym -import leisaac # noqa: F401 from prettytable import PrettyTable +import leisaac # noqa: F401 + def main(): """Print all environments registered in `leisaac` package.""" diff --git a/scripts/environments/teleoperation/replay.py b/scripts/environments/teleoperation/replay.py index 50a4fc77..c653a13e 100644 --- a/scripts/environments/teleoperation/replay.py +++ b/scripts/environments/teleoperation/replay.py @@ -57,7 +57,6 @@ import time import gymnasium as gym -import leisaac # noqa: F401 import torch from isaaclab.envs import DirectRLEnv, ManagerBasedRLEnv from isaaclab.utils.datasets import EpisodeData, HDF5DatasetFileHandler @@ -67,6 +66,8 @@ get_task_type, ) +import leisaac # noqa: F401 + class RateLimiter: """Convenience class for enforcing rates in loops.""" diff --git a/scripts/evaluation/policy_inference.py b/scripts/evaluation/policy_inference.py index 9dfd3dd5..7ee7a3dc 100644 --- a/scripts/evaluation/policy_inference.py +++ b/scripts/evaluation/policy_inference.py @@ -53,7 +53,6 @@ import carb import gymnasium as gym -import leisaac # noqa: F401 import omni import torch from isaaclab.envs import ManagerBasedRLEnv @@ -63,6 +62,8 @@ get_task_type, ) +import leisaac # noqa: F401 + class RateLimiter: """Convenience class for enforcing rates in loops.""" diff --git a/scripts/mimic/annotate_demos.py b/scripts/mimic/annotate_demos.py index 5ec29da6..f99ef769 100644 --- a/scripts/mimic/annotate_demos.py +++ b/scripts/mimic/annotate_demos.py @@ -74,7 +74,6 @@ from isaaclab.devices import Se3Keyboard import isaaclab_tasks # noqa: F401 -import leisaac # noqa: F401 from isaaclab.envs import ManagerBasedRLMimicEnv from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg from isaaclab.managers import RecorderTerm, RecorderTermCfg, TerminationTermCfg @@ -86,6 +85,8 @@ get_task_type, ) +import leisaac # noqa: F401 + is_paused = False current_action_index = 0 marked_subtask_action_indices = [] diff --git a/scripts/mimic/generate_dataset.py b/scripts/mimic/generate_dataset.py index 8928f182..c2212aaf 100644 --- a/scripts/mimic/generate_dataset.py +++ b/scripts/mimic/generate_dataset.py @@ -79,7 +79,6 @@ import isaaclab_mimic.envs.pinocchio_envs # noqa: F401 import isaaclab_tasks # noqa: F401 -import leisaac # noqa: F401 from isaaclab_mimic.datagen.generation import ( env_loop, setup_async_generation, @@ -88,6 +87,8 @@ from isaaclab_mimic.datagen.utils import get_env_name_from_dataset, setup_output_paths from leisaac.utils.env_utils import get_task_type +import leisaac # noqa: F401 + def main(): num_envs = args_cli.num_envs From 27b67b2bd8d8c02c1e7a712b0c88a8812e8da5ff Mon Sep 17 00:00:00 2001 From: chuchen0210 <167320956+chuchen0210@users.noreply.github.com> Date: Tue, 13 Jan 2026 19:09:38 +0800 Subject: [PATCH 9/9] Add LeIsaacXCosmos docs --- scripts/convert/lerobot2isaaclab.py | 551 ---------------------------- scripts/tutorials/cosmos_merge.py | 485 ------------------------ 2 files changed, 1036 deletions(-) delete mode 100644 scripts/convert/lerobot2isaaclab.py delete mode 100644 scripts/tutorials/cosmos_merge.py diff --git a/scripts/convert/lerobot2isaaclab.py b/scripts/convert/lerobot2isaaclab.py deleted file mode 100644 index ac2be622..00000000 --- a/scripts/convert/lerobot2isaaclab.py +++ /dev/null @@ -1,551 +0,0 @@ -""" -Convert a local LeRobot dataset folder (v2.1 episode-based layout) to a single HDF5 file. - -Compared to the original version, this script can optionally: - 1) Sort frames inside each episode by a chosen column (default: "index" if present), - 2) Convert LeRobot normalized joint values (roughly degrees in [-100, 100]) into - IsaacLab joint angles in RADIANS, using the same limits mapping as your replay-correct script. - -Expected input layout (v2.1): -dataset_root/ -├── data/chunk-000/episode_000000.parquet -├── data/chunk-000/episode_000001.parquet -├── videos/chunk-000//episode_000000.mp4 -└── meta/episodes.jsonl (optional) + other meta files - -Output HDF5 layout (mirrors folder-like structure): -/ - attrs: source_dir, layout_version, created_at - /meta/ (text dataset, utf-8) - /data///... (datasets from parquet columns) - /videos///.mp4 (uint8 bytes dataset, optional) - -""" - -import argparse -import datetime as dt -import json -import os -from contextlib import suppress -from pathlib import Path -from typing import Any - -# Disable HDF5 file locking to avoid BlockingIOError on some filesystems -os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" - -import h5py -import numpy as np -import pandas as pd -import torch -from lerobot.datasets.compute_stats import aggregate_stats, compute_episode_stats -from lerobot.datasets.lerobot_dataset import LeRobotDataset -from tqdm import tqdm - -ISAACLAB_LIMITS_DEG = [ - (-110.0, 110.0), - (-100.0, 100.0), - (-100.0, 90.0), - (-95.0, 95.0), - (-160.0, 160.0), - (-10.0, 100.0), -] -LEROBOT_LIMITS_DEG = [ - (-100.0, 100.0), - (-100.0, 100.0), - (-100.0, 100.0), - (-100.0, 100.0), - (-100.0, 100.0), - (0.0, 100.0), -] - - -def denormalize_lerobot_to_isaaclab_radians(joint_values_lerobot: np.ndarray) -> np.ndarray: - """ - LeRobot normalized degrees (roughly within LEROBOT_LIMITS_DEG) -> - IsaacLab joint limits (degrees) -> radians. - - joint_values_lerobot: - shape (T, D) where D >= 6. - If D == 12, treated as bimanual (6+6). - Otherwise, treated as single arm (first 6 processed, others converted deg->rad only). - """ - joint_values_lerobot = np.asarray(joint_values_lerobot, dtype=np.float32) - - if joint_values_lerobot.ndim != 2: - raise ValueError(f"Expected 2D array (T,D), got shape {joint_values_lerobot.shape}") - - dimension = joint_values_lerobot.shape[1] - if dimension < 6: - raise ValueError(f"Expected D>=6, got D={dimension}") - - result_deg = joint_values_lerobot.copy() - - if dimension == 12: - # bimanual: apply same 6-dim mapping to left [0:6] and right [6:12] - for arm_offset in (0, 6): - for joint_index in range(6): - isa_min, isa_max = ISAACLAB_LIMITS_DEG[joint_index] - le_min, le_max = LEROBOT_LIMITS_DEG[joint_index] - col = arm_offset + joint_index - result_deg[:, col] = (result_deg[:, col] - le_min) / (le_max - le_min) * (isa_max - isa_min) + isa_min - else: - # single arm (or single arm + extra features): only map the first 6 - for joint_index in range(6): - isa_min, isa_max = ISAACLAB_LIMITS_DEG[joint_index] - le_min, le_max = LEROBOT_LIMITS_DEG[joint_index] - result_deg[:, joint_index] = (result_deg[:, joint_index] - le_min) / (le_max - le_min) * ( - isa_max - isa_min - ) + isa_min - - # degrees -> radians - result_rad = result_deg * (np.pi / 180.0) - return result_rad - - -def write_text_dataset(h5_group: h5py.Group, dataset_name: str, text: str) -> None: - """Store UTF-8 text as variable-length string dataset.""" - string_dtype = h5py.string_dtype(encoding="utf-8") - if dataset_name in h5_group: - del h5_group[dataset_name] - h5_group.create_dataset(dataset_name, data=text, dtype=string_dtype) - - -def read_file_as_uint8_bytes(file_path: Path) -> np.ndarray: - data = file_path.read_bytes() - return np.frombuffer(data, dtype=np.uint8) - - -def stack_object_column_to_numpy(column_values: Any) -> tuple[np.ndarray, str | None]: - """ - Convert a parquet column to a numpy array. - - If it's already a numpy array -> return directly - - If it's a 1D array/list of per-row arrays/lists -> stack to (T, D) - - If it's scalar-like -> return (T,) or scalar - Returns: (array, note) - """ - note = None - - if isinstance(column_values, np.ndarray): - if ( - column_values.dtype == object - and len(column_values) > 0 - and isinstance(column_values[0], (list, tuple, np.ndarray)) - ): - try: - stacked = np.stack([np.asarray(x) for x in column_values], axis=0) - return stacked, note - except Exception as exc: - note = f"object-stack-failed: {exc}" - return np.asarray(column_values), note - return column_values, note - - if isinstance(column_values, (list, tuple)): - if len(column_values) > 0 and isinstance(column_values[0], (list, tuple, np.ndarray)): - try: - stacked = np.stack([np.asarray(x) for x in column_values], axis=0) - return stacked, note - except Exception as exc: - note = f"list-stack-failed: {exc}" - return np.asarray(column_values, dtype=object), note - return np.asarray(column_values), note - - return np.asarray(column_values), note - - -def write_numpy_dataset( - h5_group: h5py.Group, - dataset_name: str, - array: np.ndarray, - compression: str | None, - compression_level: int | None, -) -> None: - """Write/overwrite a dataset. If dtype=object, store JSON strings.""" - if dataset_name in h5_group: - del h5_group[dataset_name] - - # Object dtype: store each row as json string - if array.dtype == object: - string_dtype = h5py.string_dtype(encoding="utf-8") - string_list = [] - for value in array.tolist(): - try: - string_list.append(json.dumps(value, ensure_ascii=False)) - except Exception: - string_list.append(str(value)) - h5_group.create_dataset(dataset_name, data=np.asarray(string_list, dtype=object), dtype=string_dtype) - h5_group[dataset_name].attrs["stored_as"] = "json_string_list" - return - - # Numpy unicode string dtype (e.g. str: - """ - Sort dataframe by a chosen column if present. - If sort_column is None, try common ordering keys in priority: - index > frame_index > timestamp - Returns the column used, or "" if not sorted. - """ - candidate_columns = [] - if sort_column: - candidate_columns.append(sort_column) - candidate_columns += ["index", "frame_index", "timestamp"] - - for col in candidate_columns: - if col in dataframe.columns: - dataframe.sort_values(col, ascending=True, inplace=True) - dataframe.reset_index(drop=True, inplace=True) - return col - - return "" - - -def load_episode_data(dataset: LeRobotDataset, episode_idx: int) -> dict[str, Any]: - """ - Manually load all frames for a specific episode from LeRobotDataset. - Reconstructs the dictionary of arrays (T, D) for the episode. - """ - start_frame = dataset.episode_data_index["from"][episode_idx].item() - end_frame = dataset.episode_data_index["to"][episode_idx].item() - - # Pre-allocate lists for each key - # We peek at the first frame to know the keys - first_frame = dataset[start_frame] - episode_dict = {k: [] for k in first_frame.keys()} - - for i in range(start_frame, end_frame): - frame = dataset[i] - for k, v in frame.items(): - episode_dict[k].append(v) - - # Stack into arrays/tensors - # dataset[i] returns torch tensors or numpy arrays depending on backend, usually tensors - collapsed_dict = {} - for k, v_list in episode_dict.items(): - if len(v_list) > 0: - if isinstance(v_list[0], torch.Tensor): - collapsed_dict[k] = torch.stack(v_list) - else: - collapsed_dict[k] = np.stack(v_list) - else: - collapsed_dict[k] = np.array([]) - - return collapsed_dict - - -def generate_stats_if_missing(dataset_root: Path): - """ - Check if meta/episodes_stats.jsonl exists. If not, generate it and meta/stats.json - using LeRobot compute_stats utilities. - """ - episodes_stats_path = dataset_root / "meta/episodes_stats.jsonl" - stats_path = dataset_root / "meta/stats.json" - - if episodes_stats_path.exists() and stats_path.exists(): - print(f"Stats files found in {dataset_root / 'meta'}, skipping generation.") - return - - print("Missing stats files! Generating them now (this may take a while)...") - - info_path = dataset_root / "meta/info.json" - if not info_path.exists(): - raise FileNotFoundError(f"meta/info.json not found in {dataset_root}") - - with open(info_path) as f: - info = json.load(f) - features = info.get("features", {}) - - if not features: - raise ValueError("No 'features' key found in info.json") - - data_dir = dataset_root / "data" - parquet_files = sorted(data_dir.glob("chunk-*/episode_*.parquet")) - - if not parquet_files: - raise FileNotFoundError(f"No parquet files found in {data_dir}") - - print(f"Found {len(parquet_files)} episodes to process for stats generation.") - - all_stats = [] - episode_stats_list = [] - - for pq_path in tqdm(parquet_files, desc="Computing stats"): - # Load episode - df = pd.read_parquet(pq_path) - - # Convert to dictionary of arrays - episode_data = {} - for col in df.columns: - if col not in features: - continue - - val = df[col].values - if val.dtype == object and len(val) > 0 and isinstance(val[0], (list, tuple, np.ndarray)): - with suppress(Exception): - val = np.stack(val) - episode_data[col] = val - - try: - stats = compute_episode_stats(episode_data, features) - all_stats.append(stats) - - # Prepare for per-episode jsonl - stats_with_index = stats.copy() - episode_index = int(pq_path.stem.split("_")[-1]) - stats_with_index["episode_index"] = episode_index - episode_stats_list.append(stats_with_index) - - except Exception as e: - print(f"Error computing stats for {pq_path}: {e}") - raise e - - print("Aggregating statistics...") - aggregated = aggregate_stats(all_stats) - - def numpy_converter(obj): - if isinstance(obj, np.integer): - return int(obj) - elif isinstance(obj, np.floating): - return float(obj) - elif isinstance(obj, np.ndarray): - return obj.tolist() - return obj - - # Save aggregated stats - with open(stats_path, "w") as f: - json.dump(aggregated, f, indent=4, default=numpy_converter) - print(f"Saved aggregated stats to {stats_path}") - - # Save per-episode stats - with open(episodes_stats_path, "w") as f: - for stat in episode_stats_list: - clean_stat = json.loads(json.dumps(stat, default=numpy_converter)) - output_item = { - "episode_index": clean_stat["episode_index"], - "stats": {k: v for k, v in clean_stat.items() if k != "episode_index"}, - } - f.write(json.dumps(output_item) + "\n") - - print(f"Saved per-episode stats to {episodes_stats_path}") - - -def convert_lerobot_folder_to_hdf5( - dataset_directory: str, - output_hdf5_path: str, - include_videos: bool, - compression: str | None, - compression_level: int, - sort_column: str | None, - convert_to_isaaclab_radians: bool, - lerobot_action_column_name: str, - lerobot_state_column_name: str, -): - dataset_root = Path(dataset_directory).expanduser().resolve() - - # Check and generate stats if missing - generate_stats_if_missing(dataset_root) - - output_file_path = Path(output_hdf5_path).expanduser().resolve() - output_file_path.parent.mkdir(parents=True, exist_ok=True) - - videos_directory = dataset_root / "videos" - meta_directory = dataset_root / "meta" - - with h5py.File(output_file_path, "w") as output_h5: - output_h5.attrs["source_dir"] = str(dataset_root) - output_h5.attrs["layout_version"] = "lerobot_v2.1_episode_based__optionally_converted_to_isaaclab_radians" - output_h5.attrs["created_at"] = dt.datetime.now().isoformat() - output_h5.attrs["convert_to_isaaclab_radians"] = bool(convert_to_isaaclab_radians) - # Note: sort_column is not used when using LeRobotDataset as it handles loading - output_h5.attrs["sort_column_requested"] = "" if sort_column is None else str(sort_column) - output_h5.attrs["lerobot_action_column_name"] = lerobot_action_column_name - output_h5.attrs["lerobot_state_column_name"] = lerobot_state_column_name - - # Write meta files - if meta_directory.exists(): - meta_group = output_h5.require_group("meta") - for meta_file_path in sorted(meta_directory.glob("*")): - if not meta_file_path.is_file(): - continue - try: - text = meta_file_path.read_text(encoding="utf-8", errors="ignore") - except Exception: - text = meta_file_path.read_bytes().decode("utf-8", errors="ignore") - write_text_dataset(meta_group, meta_file_path.name, text) - - # Initialize LeRobotDataset - # Note: repo_id is required but for local usage we can just pass the name. - # root must be the direct path to the dataset folder. - dataset = LeRobotDataset(repo_id=dataset_root.name, root=dataset_root) - - # Convert episodes using LeRobotDataset - for episode_idx in tqdm(range(dataset.num_episodes), desc="Converting episodes"): - # Load episode data manually - episode_data = load_episode_data(dataset, episode_idx) - episode_name = f"episode_{episode_idx:06d}" - - # NOTE: Flattened structure to match isaaclab2lerobot expectations: /data/ - episode_group = output_h5.require_group(f"data/{episode_name}") - # episode_group.attrs["parquet_relpath"] = str(parquet_path.relative_to(dataset_root)) # Not applicable - - # Use 'index' column for length if available, else use first column length - num_frames = len(episode_data[next(iter(episode_data))]) - episode_group.attrs["num_frames"] = int(num_frames) - # episode_group.attrs["sorted_by"] = used_sort_column # LeRobotDataset returns sorted frames - - # Write each column - for column_name, column_tensor in episode_data.items(): - if isinstance(column_tensor, torch.Tensor): - numpy_array = column_tensor.numpy() - else: - numpy_array = np.array(column_tensor) - - # Normalize float dtype if possible - if numpy_array.dtype == np.float64: - numpy_array = numpy_array.astype(np.float32) - - # If enabled: convert action/state columns from LeRobot normalized degrees -> IsaacLab radians - if convert_to_isaaclab_radians and numpy_array.ndim == 2: - if column_name in (lerobot_action_column_name, lerobot_state_column_name): - numpy_array = denormalize_lerobot_to_isaaclab_radians(numpy_array) - - safe_dataset_name = column_name.replace("/", "_") - write_numpy_dataset( - h5_group=episode_group, - dataset_name=safe_dataset_name, - array=numpy_array, - compression=compression, - compression_level=compression_level if compression else None, - ) - - # Note: `note` logic from stack_object_column_to_numpy is skipped as LeRobotDataset returns tensors - - # Optionally store mp4 bytes in HDF5 - if include_videos and videos_directory.exists(): - video_files = sorted(videos_directory.glob("chunk-*/**/episode_*.mp4")) - videos_group = output_h5.require_group("videos") - - for video_path in tqdm(video_files, desc="Packing videos into HDF5"): - relative_video_path = video_path.relative_to(dataset_root) - parts = relative_video_path.parts - if len(parts) < 4: - continue - - _, chunk_name, camera_name = parts[0], parts[1], parts[2] - episode_filename = parts[-1] - episode_stem = Path(episode_filename).stem - - camera_group = videos_group.require_group(f"{chunk_name}/{camera_name}") - dataset_name = f"{episode_stem}.mp4" - - if dataset_name in camera_group: - del camera_group[dataset_name] - - video_bytes = read_file_as_uint8_bytes(video_path) - camera_group.create_dataset( - dataset_name, - data=video_bytes, - dtype=np.uint8, - compression=compression, - compression_opts=compression_level if compression else None, - shuffle=True if compression else None, - ) - camera_group[dataset_name].attrs["video_relpath"] = str(relative_video_path) - - print(f"[OK] Wrote HDF5 to: {output_file_path}") - - -def main(): - argument_parser = argparse.ArgumentParser( - description="Convert local LeRobot dataset folder (v2.1) to a single HDF5 file." - ) - argument_parser.add_argument("--lerobot_dir", type=str, required=True, help="Local LeRobot dataset folder path") - argument_parser.add_argument("--output_hdf5", type=str, required=True, help="Output HDF5 file path") - - argument_parser.add_argument( - "--include_videos", - action="store_true", - help="Store mp4 bytes into HDF5 (may be large)", - ) - argument_parser.add_argument( - "--no_compression", - action="store_true", - help="Disable HDF5 compression", - ) - argument_parser.add_argument( - "--compression_level", - type=int, - default=4, - help="gzip compression level (1-9)", - ) - argument_parser.add_argument( - "--sort_column", - type=str, - default=None, - help=( - "Sort frames inside each episode by this column if present. " - "If not provided, tries: index > frame_index > timestamp." - ), - ) - argument_parser.add_argument( - "--no_convert_to_isaaclab_radians", - action="store_true", - help="Disable conversion of LeRobot normalized degrees to IsaacLab radians (default: conversion is ENABLED).", - ) - argument_parser.add_argument( - "--lerobot_action_column_name", - type=str, - default="action", - help="Column name in parquet for action vectors (default: action)", - ) - argument_parser.add_argument( - "--lerobot_state_column_name", - type=str, - default="observation.state", - help="Column name in parquet for state vectors (default: observation.state)", - ) - - args = argument_parser.parse_args() - - compression = None if args.no_compression else "gzip" - - convert_lerobot_folder_to_hdf5( - dataset_directory=args.lerobot_dir, - output_hdf5_path=args.output_hdf5, - include_videos=args.include_videos, - compression=compression, - compression_level=args.compression_level, - sort_column=args.sort_column, - convert_to_isaaclab_radians=not args.no_convert_to_isaaclab_radians, - lerobot_action_column_name=args.lerobot_action_column_name, - lerobot_state_column_name=args.lerobot_state_column_name, - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/tutorials/cosmos_merge.py b/scripts/tutorials/cosmos_merge.py deleted file mode 100644 index 79dc002a..00000000 --- a/scripts/tutorials/cosmos_merge.py +++ /dev/null @@ -1,485 +0,0 @@ -""" -Merge a LeRobot-converted HDF5 with a source/template LeIsaac HDF5. - -Inputs: - - lerobot_converted_hdf5_path: - HDF5 produced from local LeRobot dataset conversion (episode parquet columns stored as datasets) - - source_template_hdf5_path: - LeIsaac HDF5 providing initial_state + episode structure template + metadata -Output: - - output_merged_hdf5_path: - New replayable HDF5 that follows source structure, but actions/states come from the LeRobot-converted HDF5. - -Default behavior: - - Copy everything from source root EXCEPT '/data' into output root. - - Choose a template episode from source '/data' (first one, or specified). - - For each lerobot episode under /data/chunk-*/episode_*: - - Copy template episode group into output '/data/' - - Overwrite: - /actions - /obs/actions (if exists) - /processed_actions - - Optionally overwrite a user-specified joint/state dataset path with LeRobot observation.state -""" - -import argparse -from contextlib import suppress -from pathlib import Path - -import h5py -import numpy as np -from tqdm import tqdm - - -def list_template_episode_names(source_hdf5_file: h5py.File) -> list[str]: - """List episode group names under source '/data'.""" - if "data" not in source_hdf5_file: - return [] - return sorted(list(source_hdf5_file["data"].keys())) - - -def ensure_parent_groups_exist(episode_root_group: h5py.Group, dataset_path: str) -> h5py.Group: - """ - Ensure all parent groups exist for a dataset path like - 'states/articulation/robot/joint_position'. Returns the parent group. - """ - path_parts = [part for part in dataset_path.split("/") if part] - if len(path_parts) <= 1: - return episode_root_group - - parent_parts = path_parts[:-1] - current_group = episode_root_group - for group_name in parent_parts: - current_group = current_group.require_group(group_name) - return current_group - - -def copy_hdf5_attributes(source_object, destination_object) -> None: - """Copy all HDF5 attributes from source object to destination object.""" - for attribute_key, attribute_value in source_object.attrs.items(): - destination_object.attrs[attribute_key] = attribute_value - - -def copy_source_root_except_data_group(source_hdf5_file: h5py.File, output_hdf5_file: h5py.File) -> None: - """ - Copy everything at the root of source HDF5 into output HDF5 except '/data'. - This preserves metadata, configs, and other root-level structures. - """ - copy_hdf5_attributes(source_hdf5_file, output_hdf5_file) - for root_item_name in source_hdf5_file.keys(): - if root_item_name == "data": - continue - source_hdf5_file.copy(root_item_name, output_hdf5_file, name=root_item_name) - - -def copy_template_episode_group( - source_hdf5_file: h5py.File, - output_hdf5_file: h5py.File, - template_episode_name: str, - output_episode_name: str, -) -> None: - """ - Copy '/data/' from source into output at '/data/'. - """ - output_data_group = output_hdf5_file.require_group("data") - - if output_episode_name in output_data_group: - del output_data_group[output_episode_name] - - source_hdf5_file["data"].copy(template_episode_name, output_data_group, name=output_episode_name) - - # Preserve attrs on /data itself if present - copy_hdf5_attributes(source_hdf5_file["data"], output_data_group) - - -def write_or_overwrite_hdf5_dataset( - parent_group: h5py.Group, - dataset_name: str, - dataset_data: np.ndarray, - template_dataset: h5py.Dataset | None = None, -) -> None: - """ - Overwrite dataset 'dataset_name' inside 'parent_group' with numpy 'dataset_data'. - If 'template_dataset' is provided, mirror its dtype/chunking/compression settings when possible. - """ - if dataset_name in parent_group: - del parent_group[dataset_name] - - create_dataset_kwargs = {} - - if template_dataset is not None: - # Mirror chunking/compression options - if template_dataset.chunks is not None: - create_dataset_kwargs["chunks"] = template_dataset.chunks - if template_dataset.compression is not None: - create_dataset_kwargs["compression"] = template_dataset.compression - if template_dataset.compression_opts is not None: - create_dataset_kwargs["compression_opts"] = template_dataset.compression_opts - if template_dataset.shuffle: - create_dataset_kwargs["shuffle"] = True - - # Try to match dtype (if possible) - with suppress(Exception): - dataset_data = dataset_data.astype(template_dataset.dtype, copy=False) - - parent_group.create_dataset(dataset_name, data=dataset_data, **create_dataset_kwargs) - - -def iterate_lerobot_episode_groups( - lerobot_hdf5_file: h5py.File, -) -> list[tuple[str, h5py.Group]]: - """ - Iterate episodes from a lerobot-converted HDF5: - - Supports layouts: - 1) /data/chunk-000/episode_000000 (nested) - 2) /data/episode_000000 (flat) - - Returns list of: - (episode_identifier_string, episode_group) - where episode_identifier_string is like 'chunk-000/episode_000000' or 'episode_000000'. - """ - episodes: list[tuple[str, h5py.Group]] = [] - - if "data" not in lerobot_hdf5_file: - return episodes - - for key in sorted(lerobot_hdf5_file["data"].keys()): - item = lerobot_hdf5_file["data"][key] - - # Case 1: Flat structure /data/episode_XXXXXX - if key.startswith("episode_") and isinstance(item, h5py.Group): - episodes.append((key, item)) - continue - - # Case 2: Nested structure /data/chunk-XXX/episode_XXXXXX - # We assume any other group might contain episodes (e.g. chunk-000) - if isinstance(item, h5py.Group): - for sub_key in sorted(item.keys()): - if sub_key.startswith("episode_"): - sub_item = item[sub_key] - if isinstance(sub_item, h5py.Group): - episodes.append((f"{key}/{sub_key}", sub_item)) - - return episodes - - -def read_action_and_state_arrays_from_lerobot_group( - lerobot_episode_group: h5py.Group, - lerobot_action_dataset_key: str, - lerobot_state_dataset_key: str, -) -> tuple[np.ndarray, np.ndarray | None]: - """ - Read action and state arrays from a lerobot episode group. - Supports common naming variants for 'observation.state' due to prior sanitization. - - Returns: - actions_array: np.ndarray - state_array: Optional[np.ndarray] (None if not found) - """ - if lerobot_action_dataset_key not in lerobot_episode_group: - raise KeyError( - f"Missing action dataset '{lerobot_action_dataset_key}' " - f"in lerobot episode group: {lerobot_episode_group.name}" - ) - - actions_array = np.array(lerobot_episode_group[lerobot_action_dataset_key]) - - state_candidate_keys = [lerobot_state_dataset_key] - - # Accept common sanitized variants between 'observation.state' and 'observation_state' - if lerobot_state_dataset_key == "observation.state": - state_candidate_keys.append("observation_state") - elif lerobot_state_dataset_key == "observation_state": - state_candidate_keys.append("observation.state") - - state_array = None - for state_key in state_candidate_keys: - if state_key in lerobot_episode_group: - state_array = np.array(lerobot_episode_group[state_key]) - break - - return actions_array, state_array - - -def merge_lerobot_converted_hdf5_with_source_template( - lerobot_converted_hdf5_path: str, - source_template_hdf5_path: str, - output_merged_hdf5_path: str, - source_template_episode_name: str | None, - output_episode_name_prefix: str, - lerobot_action_dataset_key: str, - lerobot_state_dataset_key: str, - lerobot_timestamp_dataset_key: str | None, - write_state_to_output_hdf5: bool, - output_state_target_dataset_path: str | None, - skip_first_n_lerobot_episodes: int, - max_number_of_episodes_to_merge: int | None, -) -> None: - lerobot_converted_hdf5_path = str(Path(lerobot_converted_hdf5_path).expanduser().resolve()) - source_template_hdf5_path = str(Path(source_template_hdf5_path).expanduser().resolve()) - output_merged_hdf5_path = str(Path(output_merged_hdf5_path).expanduser().resolve()) - Path(output_merged_hdf5_path).parent.mkdir(parents=True, exist_ok=True) - - with h5py.File(source_template_hdf5_path, "r") as source_hdf5_file, h5py.File( - lerobot_converted_hdf5_path, "r" - ) as lerobot_hdf5_file, h5py.File(output_merged_hdf5_path, "w") as output_hdf5_file: - - # 1) Copy root-level template/metadata except '/data' - copy_source_root_except_data_group(source_hdf5_file, output_hdf5_file) - - # 2) Choose a template episode in source '/data' - available_template_episode_names = list_template_episode_names(source_hdf5_file) - if not available_template_episode_names: - raise RuntimeError("source_template_hdf5 has no '/data/' groups to use as template.") - - if source_template_episode_name is None: - source_template_episode_name = available_template_episode_names[0] - else: - if source_template_episode_name not in source_hdf5_file["data"]: - preview = available_template_episode_names[:10] - raise RuntimeError( - f"template_episode '{source_template_episode_name}' not found in source '/data'. " - f"Available examples: {preview}{'...' if len(available_template_episode_names) > 10 else ''}" - ) - - # 3) Iterate lerobot episodes - lerobot_episode_list = iterate_lerobot_episode_groups(lerobot_hdf5_file) - if not lerobot_episode_list: - raise RuntimeError( - "lerobot_converted_hdf5 has no episodes under '/data/chunk-*/episode_*' or '/data/episode_*'." - ) - - # Apply slicing - lerobot_episode_list = lerobot_episode_list[skip_first_n_lerobot_episodes:] - if max_number_of_episodes_to_merge is not None: - lerobot_episode_list = lerobot_episode_list[:max_number_of_episodes_to_merge] - - # Ensure /data exists in output - output_hdf5_file.require_group("data") - - for output_episode_index, (lerobot_episode_identifier, lerobot_episode_group) in enumerate( - tqdm(lerobot_episode_list, desc="Merging episodes") - ): - output_episode_name = f"{output_episode_name_prefix}{output_episode_index}" - - # 3.1 Copy template episode group to output - copy_template_episode_group( - source_hdf5_file=source_hdf5_file, - output_hdf5_file=output_hdf5_file, - template_episode_name=source_template_episode_name, - output_episode_name=output_episode_name, - ) - output_episode_group = output_hdf5_file["data"][output_episode_name] - - # 3.2 Read lerobot action/state arrays - actions_array, state_array = read_action_and_state_arrays_from_lerobot_group( - lerobot_episode_group=lerobot_episode_group, - lerobot_action_dataset_key=lerobot_action_dataset_key, - lerobot_state_dataset_key=lerobot_state_dataset_key, - ) - - # 3.3 Overwrite /actions - template_actions_dataset = output_episode_group["actions"] if "actions" in output_episode_group else None - if isinstance(template_actions_dataset, h5py.Group): - template_actions_dataset = None - - write_or_overwrite_hdf5_dataset( - parent_group=output_episode_group, - dataset_name="actions", - dataset_data=actions_array, - template_dataset=( - template_actions_dataset if isinstance(template_actions_dataset, h5py.Dataset) else None - ), - ) - - # 3.4 Overwrite /obs/actions if exists (or create) - if "obs" in output_episode_group and isinstance(output_episode_group["obs"], h5py.Group): - output_obs_group = output_episode_group["obs"] - else: - output_obs_group = output_episode_group.require_group("obs") - - template_obs_actions_dataset = output_obs_group["actions"] if "actions" in output_obs_group else None - write_or_overwrite_hdf5_dataset( - parent_group=output_obs_group, - dataset_name="actions", - dataset_data=actions_array, - template_dataset=( - template_obs_actions_dataset if isinstance(template_obs_actions_dataset, h5py.Dataset) else None - ), - ) - - # 3.5 Overwrite /processed_actions (create if missing) - template_processed_actions_dataset = ( - output_episode_group["processed_actions"] - if "processed_actions" in output_episode_group - and isinstance(output_episode_group["processed_actions"], h5py.Dataset) - else None - ) - write_or_overwrite_hdf5_dataset( - parent_group=output_episode_group, - dataset_name="processed_actions", - dataset_data=actions_array, - template_dataset=template_processed_actions_dataset, - ) - - # 3.6 Optional: overwrite timestamps - if lerobot_timestamp_dataset_key is not None and lerobot_timestamp_dataset_key in lerobot_episode_group: - timestamps_array = np.array(lerobot_episode_group[lerobot_timestamp_dataset_key]) - template_timestamps_dataset = ( - output_episode_group["timestamps"] - if "timestamps" in output_episode_group - and isinstance(output_episode_group["timestamps"], h5py.Dataset) - else None - ) - write_or_overwrite_hdf5_dataset( - parent_group=output_episode_group, - dataset_name="timestamps", - dataset_data=timestamps_array, - template_dataset=template_timestamps_dataset, - ) - - # 3.7 Optional: write state/joint positions into a target dataset path - if write_state_to_output_hdf5 and state_array is not None and output_state_target_dataset_path is not None: - parent_group = ensure_parent_groups_exist(output_episode_group, output_state_target_dataset_path) - target_dataset_name = output_state_target_dataset_path.split("/")[-1] - - # Mirror template dataset properties if it exists - template_state_dataset = None - if output_state_target_dataset_path in output_episode_group and isinstance( - output_episode_group[output_state_target_dataset_path], h5py.Dataset - ): - template_state_dataset = output_episode_group[output_state_target_dataset_path] - - write_or_overwrite_hdf5_dataset( - parent_group=parent_group, - dataset_name=target_dataset_name, - dataset_data=state_array, - template_dataset=template_state_dataset, - ) - - # Provenance - output_episode_group.attrs["lerobot_episode_identifier"] = lerobot_episode_identifier - output_episode_group.attrs["source_template_episode_name"] = source_template_episode_name - - # Root provenance - output_hdf5_file.attrs["merged_from_lerobot_converted_hdf5"] = lerobot_converted_hdf5_path - output_hdf5_file.attrs["merged_from_source_template_hdf5"] = source_template_hdf5_path - - print(f"[OK] Wrote merged replayable HDF5: {output_merged_hdf5_path}") - - -def main(): - argument_parser = argparse.ArgumentParser( - description="Merge lerobot-converted HDF5 with source IsaacLab/LeIsaac template HDF5." - ) - - argument_parser.add_argument( - "--lerobot_hdf5", - type=str, - required=True, - help="Path to HDF5 converted from local LeRobot dataset", - ) - argument_parser.add_argument( - "--source_hdf5", - type=str, - required=True, - help="Path to source/template IsaacLab/LeIsaac HDF5 (provides initial_state + structure template)", - ) - argument_parser.add_argument( - "--output_hdf5", - type=str, - required=True, - help="Path to output merged HDF5", - ) - - argument_parser.add_argument( - "--template_episode", - type=str, - default=None, - help="Episode name under source '/data' to use as template (default: first episode found)", - ) - - argument_parser.add_argument( - "--output_episode_prefix", - type=str, - default="demo_", - help="Output episode name prefix under '/data' (default: demo)", - ) - - # Keys in lerobot episode group - argument_parser.add_argument( - "--lerobot_action_key", - type=str, - default="action", - help="Dataset key inside lerobot episode group for actions (default: action)", - ) - argument_parser.add_argument( - "--lerobot_state_key", - type=str, - default="observation.state", - help=( - "Dataset key inside lerobot episode group for state/joint pos (default: observation.state; also tries" - " observation_state)" - ), - ) - argument_parser.add_argument( - "--lerobot_timestamp_key", - type=str, - default=None, - help="Optional dataset key inside lerobot episode group for timestamps (e.g. timestamp)", - ) - - # Where to write state in output episode - argument_parser.add_argument( - "--write_state", - action="store_true", - help="If set, write lerobot state/joint positions into output at --output_state_target_path", - ) - argument_parser.add_argument( - "--output_state_target_path", - type=str, - default=None, - help=( - "Target dataset path inside each output episode to store state " - "(e.g. 'obs/joint_pos' or 'states/articulation/robot/joint_position')" - ), - ) - - # Episode slicing - argument_parser.add_argument( - "--skip_first_n_episodes", - type=int, - default=0, - help="Skip first N lerobot episodes (default: 0)", - ) - argument_parser.add_argument( - "--max_episodes", - type=int, - default=None, - help="Merge at most N episodes (default: all)", - ) - - parsed_arguments = argument_parser.parse_args() - - if parsed_arguments.write_state and not parsed_arguments.output_state_target_path: - raise SystemExit("--write_state requires --output_state_target_path to be set") - - merge_lerobot_converted_hdf5_with_source_template( - lerobot_converted_hdf5_path=parsed_arguments.lerobot_hdf5, - source_template_hdf5_path=parsed_arguments.source_hdf5, - output_merged_hdf5_path=parsed_arguments.output_hdf5, - source_template_episode_name=parsed_arguments.template_episode, - output_episode_name_prefix=parsed_arguments.output_episode_prefix, - lerobot_action_dataset_key=parsed_arguments.lerobot_action_key, - lerobot_state_dataset_key=parsed_arguments.lerobot_state_key, - lerobot_timestamp_dataset_key=parsed_arguments.lerobot_timestamp_key, - write_state_to_output_hdf5=parsed_arguments.write_state, - output_state_target_dataset_path=parsed_arguments.output_state_target_path, - skip_first_n_lerobot_episodes=parsed_arguments.skip_first_n_episodes, - max_number_of_episodes_to_merge=parsed_arguments.max_episodes, - ) - - -if __name__ == "__main__": - main()