scripts/data/convert_agibot.py

import argparse
from copy import deepcopy
from functools import partial
import gc
import json
import logging
from math import ceil
import os
from pathlib import Path
import shutil
from typing import Callable

import einops
import h5py
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
from lerobot.common.datasets.utils import (
    STATS_PATH,
    check_timestamps_sync,
    get_episode_data_index,
    serialize_dict,
    write_json,
)
import numpy as np
import torch
from tqdm import tqdm
from tqdm.contrib.concurrent import process_map


def generate_modality_json(output_dir: str) -> None:
    """Generate modality.json file defining field mappings for the dataset."""
    modality_config = {
        "state": {
            "left_arm_joint_position": {
                "original_key": "observation.state",
                "start": 0,
                "end": 7,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "right_arm_joint_position": {
                "original_key": "observation.state",
                "start": 7,
                "end": 14,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "left_effector_position": {
                "original_key": "observation.state",
                "start": 14,
                "end": 15,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "right_effector_position": {
                "original_key": "observation.state",
                "start": 15,
                "end": 16,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "head_position": {
                "original_key": "observation.state",
                "start": 16,
                "end": 18,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "waist_pitch": {
                "original_key": "observation.state",
                "start": 18,
                "end": 19,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "waist_lift": {
                "original_key": "observation.state",
                "start": 19,
                "end": 20,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
        },
        "action": {
            "left_arm_joint_position": {
                "original_key": "action",
                "start": 0,
                "end": 7,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "right_arm_joint_position": {
                "original_key": "action",
                "start": 7,
                "end": 14,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "left_effector_position": {
                "original_key": "action",
                "start": 14,
                "end": 15,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "right_effector_position": {
                "original_key": "action",
                "start": 15,
                "end": 16,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "head_position": {
                "original_key": "action",
                "start": 16,
                "end": 18,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "waist_pitch": {
                "original_key": "action",
                "start": 18,
                "end": 19,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "waist_lift": {
                "original_key": "action",
                "start": 19,
                "end": 20,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
            "robot_velocity": {
                "original_key": "action",
                "start": 20,
                "end": 22,
                "rotation_type": None,
                "absolute": True,
                "dtype": "float64",
                "range": None
            },
        },
        "video": {
            "top_head": {
                "original_key": "observation.images.top_head"
            },
            "hand_left": {
                "original_key": "observation.images.hand_left"
            },
            "hand_right": {
                "original_key": "observation.images.hand_right"
            },
        },
        "annotation": {
            "language.action_text": {
                "original_key": "task_index"
            },
            "agibot.sub_task": {
                "original_key": "annotation.agibot.sub_task"
            },
            "frame_type": {
                "original_key": "annotation.frame_type"
            },
        },
    }

    modality_path = os.path.join(output_dir, "modality.json")
    with open(modality_path, "w") as f:
        json.dump(modality_config, f, indent=4)
    print(f"Generated modality.json at {modality_path}")


HEAD_COLOR = "head_color.mp4"
HAND_LEFT_COLOR = "hand_left_color.mp4"
HAND_RIGHT_COLOR = "hand_right_color.mp4"

FEATURES = {
    "observation.images.top_head": {
        "dtype": "video",
        "shape": [480, 640, 3],
        "names": ["height", "width", "channel"],
        "video_info": {
            "video.fps": 30.0,
            "video.codec": "av1",
            "video.pix_fmt": "yuv420p",
            "video.is_depth_map": False,
            "has_audio": False,
        },
    },
    "observation.images.hand_left": {
        "dtype": "video",
        "shape": [480, 640, 3],
        "names": ["height", "width", "channel"],
        "video_info": {
            "video.fps": 30.0,
            "video.codec": "av1",
            "video.pix_fmt": "yuv420p",
            "video.is_depth_map": False,
            "has_audio": False,
        },
    },
    "observation.images.hand_right": {
        "dtype": "video",
        "shape": [480, 640, 3],
        "names": ["height", "width", "channel"],
        "video_info": {
            "video.fps": 30.0,
            "video.codec": "av1",
            "video.pix_fmt": "yuv420p",
            "video.is_depth_map": False,
            "has_audio": False,
        },
    },
    "observation.state": {
        "dtype": "float32",
        "shape": [20],
    },
    "action": {
        "dtype": "float32",
        "shape": [22],
    },
    "annotation.language.action_text": {
        "dtype": "int64",
        "shape": [1],
        "names": None,
    },
    "annotation.agibot.tasks": {
        "dtype": "int64",
        "shape": [1],
        "names": None,
    },
    "episode_index": {
        "dtype": "int64",
        "shape": [1],
        "names": None,
    },
    "frame_index": {
        "dtype": "int64",
        "shape": [1],
        "names": None,
    },
    "index": {
        "dtype": "int64",
        "shape": [1],
        "names": None,
    },
    "task_index": {
        "dtype": "int64",
        "shape": [1],
        "names": None,
    },
}


def get_stats_einops_patterns(dataset, num_workers=0):
    """These einops patterns will be used to aggregate batches and compute statistics.

    Note: We assume the images are in channel first format
    """

    dataloader = torch.utils.data.DataLoader(
        dataset,
        num_workers=num_workers,
        batch_size=2,
        shuffle=False,
    )
    batch = next(iter(dataloader))

    stats_patterns = {}

    for key in dataset.features:
        # sanity check that tensors are not float64
        assert batch[key].dtype != torch.float64

        # if isinstance(feats_type, (VideoFrame, Image)):
        if key in dataset.meta.camera_keys:
            # sanity check that images are channel first
            _, c, h, w = batch[key].shape
            assert c < h and c < w, f"expect channel first images, but instead {batch[key].shape}"
            assert (
                batch[key].dtype == torch.float32
            ), f"expect torch.float32, but instead {batch[key].dtype=}"
            # assert batch[key].max() <= 1, f"expect pixels lower than 1, but instead {batch[key].max()=}"
            # assert batch[key].min() >= 0, f"expect pixels greater than 1, but instead {batch[key].min()=}"
            stats_patterns[key] = "b c h w -> c 1 1"
        elif batch[key].ndim == 2:
            stats_patterns[key] = "b c -> c "
        elif batch[key].ndim == 1:
            stats_patterns[key] = "b -> 1"
        else:
            raise ValueError(f"{key}, {batch[key].shape}")

    return stats_patterns


def compute_stats(dataset, batch_size=8, num_workers=4, max_num_samples=None):
    """Compute mean/std and min/max statistics of all data keys in a LeRobotDataset."""
    if max_num_samples is None:
        max_num_samples = len(dataset)

    # for more info on why we need to set the same number of workers, see `load_from_videos`
    stats_patterns = get_stats_einops_patterns(dataset, num_workers)

    # mean and std will be computed incrementally while max and min will track the running value.
    mean, std, max, min = {}, {}, {}, {}
    for key in stats_patterns:
        mean[key] = torch.tensor(0.0).float()
        std[key] = torch.tensor(0.0).float()
        max[key] = torch.tensor(-float("inf")).float()
        min[key] = torch.tensor(float("inf")).float()

    def create_seeded_dataloader(dataset, batch_size, seed):
        generator = torch.Generator()
        generator.manual_seed(seed)
        dataloader = torch.utils.data.DataLoader(
            dataset,
            num_workers=num_workers,
            batch_size=batch_size,
            shuffle=True,
            drop_last=False,
            generator=generator,
        )
        return dataloader

    # Note: Due to be refactored soon. The point of storing `first_batch` is to make sure we don't get
    # surprises when rerunning the sampler.
    first_batch = None
    running_item_count = 0  # for online mean computation
    dataloader = create_seeded_dataloader(dataset, batch_size, seed=1337)
    for i, batch in enumerate(
        tqdm(
            dataloader,
            total=ceil(max_num_samples / batch_size),
            desc="Compute mean, min, max",
        )
    ):
        this_batch_size = len(batch["index"])
        running_item_count += this_batch_size
        if first_batch is None:
            first_batch = deepcopy(batch)
        for key, pattern in stats_patterns.items():
            batch[key] = batch[key].float()
            # Numerically stable update step for mean computation.
            batch_mean = einops.reduce(batch[key], pattern, "mean")
            # Hint: to update the mean we need x̄ₙ = (Nₙ₋₁x̄ₙ₋₁ + Bₙxₙ) / Nₙ, where the subscript represents
            # the update step, N is the running item count, B is this batch size, x̄ is the running mean,
            # and x is the current batch mean. Some rearrangement is then required to avoid risking
            # numerical overflow. Another hint: Nₙ₋₁ = Nₙ - Bₙ. Rearrangement yields
            # x̄ₙ = x̄ₙ₋₁ + Bₙ * (xₙ - x̄ₙ₋₁) / Nₙ
            mean[key] = mean[key] + this_batch_size * (batch_mean - mean[key]) / running_item_count
            max[key] = torch.maximum(max[key], einops.reduce(batch[key], pattern, "max"))
            min[key] = torch.minimum(min[key], einops.reduce(batch[key], pattern, "min"))

        if i == ceil(max_num_samples / batch_size) - 1:
            break

    first_batch_ = None
    running_item_count = 0  # for online std computation
    dataloader = create_seeded_dataloader(dataset, batch_size, seed=1337)
    for i, batch in enumerate(
        tqdm(dataloader, total=ceil(max_num_samples / batch_size), desc="Compute std")
    ):
        this_batch_size = len(batch["index"])
        running_item_count += this_batch_size
        # Sanity check to make sure the batches are still in the same order as before.
        if first_batch_ is None:
            first_batch_ = deepcopy(batch)
            for key in stats_patterns:
                assert torch.equal(first_batch_[key], first_batch[key])
        for key, pattern in stats_patterns.items():
            batch[key] = batch[key].float()
            # Numerically stable update step for mean computation (where the mean is over squared
            # residuals).See notes in the mean computation loop above.
            batch_std = einops.reduce((batch[key] - mean[key]) ** 2, pattern, "mean")
            std[key] = std[key] + this_batch_size * (batch_std - std[key]) / running_item_count

        if i == ceil(max_num_samples / batch_size) - 1:
            break

    for key in stats_patterns:
        std[key] = torch.sqrt(std[key])

    stats = {}
    for key in stats_patterns:
        stats[key] = {
            "mean": mean[key],
            "std": std[key],
            "max": max[key],
            "min": min[key],
        }
    return stats


class AgiBotDataset(LeRobotDataset):
    def __init__(
        self,
        repo_id: str,
        root: str | Path | None = None,
        episodes: list[int] | None = None,
        image_transforms: Callable | None = None,
        delta_timestamps: dict[list[float]] | None = None,
        tolerance_s: float = 1e-4,
        download_videos: bool = True,
        local_files_only: bool = False,
        video_backend: str | None = None,
    ):
        super().__init__(
            repo_id=repo_id,
            root=root,
            episodes=episodes,
            image_transforms=image_transforms,
            delta_timestamps=delta_timestamps,
            tolerance_s=tolerance_s,
            download_videos=download_videos,
            local_files_only=local_files_only,
            video_backend=video_backend,
        )

    def save_episode(
        self, task: str, episode_data: dict | None = None, videos: dict | None = None
    ) -> None:
        """
        We rewrite this method to copy mp4 videos to the target position
        """
        if not episode_data:
            episode_buffer = self.episode_buffer

        episode_length = episode_buffer.pop("size")
        episode_index = episode_buffer["episode_index"]
        if episode_index != self.meta.total_episodes:
            # TODO(aliberts): Add option to use existing episode_index
            raise NotImplementedError(
                "You might have manually provided the episode_buffer with an episode_index that doesn't "
                "match the total number of episodes in the dataset. This is not supported for now."
            )

        if episode_length == 0:
            raise ValueError(
                "You must add one or several frames with `add_frame` before calling `add_episode`."
            )

        # Use our custom task indexing instead of LeRobot's built-in mechanism
        task_index = getattr(self, "_custom_task_to_index", {}).get(task, 0)

        # Remove the 'task' key if it exists (it's passed as a parameter, not needed in buffer)
        episode_buffer.pop("task", None)

        if not set(episode_buffer.keys()) == set(self.features):
            raise ValueError()

        for key, ft in self.features.items():
            if key == "index":
                episode_buffer[key] = np.arange(
                    self.meta.total_frames, self.meta.total_frames + episode_length
                )
            elif key == "episode_index":
                episode_buffer[key] = np.full((episode_length,), episode_index)
            elif key == "task_index":
                episode_buffer[key] = np.full((episode_length,), task_index)
            elif ft["dtype"] in ["image", "video"]:
                continue
            elif ft["dtype"] == "string":
                pass
            elif len(ft["shape"]) == 1 and ft["shape"][0] == 1:
                episode_buffer[key] = np.array(episode_buffer[key], dtype=ft["dtype"])
            elif len(ft["shape"]) == 1 and ft["shape"][0] > 1:
                episode_buffer[key] = np.stack(episode_buffer[key])
            else:
                raise ValueError(key)

        self._wait_image_writer()
        self._save_episode_table(episode_buffer, episode_index)

        # Copy videos first before calling meta.save_episode which might try to read them
        for key in self.meta.video_keys:
            video_path = self.root / self.meta.get_video_file_path(episode_index, key)
            episode_buffer[key] = video_path
            video_path.parent.mkdir(parents=True, exist_ok=True)
            # Copy video files to target location
            shutil.copyfile(str(videos[key]), str(video_path))

        try:
            # FIX: Call meta.save_episode with correct parameters
            # Note: We pass an empty task list to prevent duplicate entries in tasks.jsonl
            # since we create our own tasks.jsonl file with proper indexing
            self.meta.save_episode(episode_index, episode_length, [], {})
        except AttributeError as e:
            if "'NoneType' object has no attribute 'items'" in str(e):
                # Handle the episode stats issue - skip episode stats for now
                print(
                    f"Warning: Episode stats computation failed, proceeding without stats "
                    f"for episode {episode_index}"
                )
                # Just skip the problematic save_episode call - the core data is already saved
                pass
            else:
                raise
        if not episode_data:  # Reset the buffer
            self.episode_buffer = self.create_episode_buffer()
        self.consolidated = False

    def consolidate(self, run_compute_stats: bool = True, keep_image_files: bool = False) -> None:
        self.hf_dataset = self.load_hf_dataset()
        self.episode_data_index = get_episode_data_index(self.meta.episodes, self.episodes)
        check_timestamps_sync(self.hf_dataset, self.episode_data_index, self.fps, self.tolerance_s)
        if len(self.meta.video_keys) > 0:
            self.meta.write_video_info()

        if not keep_image_files:
            img_dir = self.root / "images"
            if img_dir.is_dir():
                shutil.rmtree(self.root / "images")
        video_files = list(self.root.rglob("*.mp4"))
        assert len(video_files) == self.num_episodes * len(self.meta.video_keys)

        parquet_files = list(self.root.rglob("*.parquet"))
        assert len(parquet_files) == self.num_episodes

        if run_compute_stats:
            self.stop_image_writer()
            self.meta.stats = compute_stats(self)
            serialized_stats = serialize_dict(self.meta.stats)
            write_json(serialized_stats, self.root / STATS_PATH)
            self.consolidated = True
        else:
            logging.warning(
                "Skipping computation of the dataset statistics, dataset is not fully consolidated."
            )

    def add_frame(self, frame: dict) -> None:
        """
        This function only adds the frame to the episode_buffer. Apart from images — which are written in a
        temporary directory — nothing is written to disk. To save those frames, the 'save_episode()' method
        then needs to be called.
        """
        # TODO(aliberts, rcadene): Add sanity check for the input, check it's numpy or torch,
        # check the dtype and shape matches, etc.

        if self.episode_buffer is None:
            self.episode_buffer = self.create_episode_buffer()

        frame_index = self.episode_buffer["size"]
        timestamp = frame.pop("timestamp") if "timestamp" in frame else frame_index / self.fps
        self.episode_buffer["frame_index"].append(frame_index)
        self.episode_buffer["timestamp"].append(timestamp)

        for key in frame:
            if key not in self.features:
                raise ValueError(key)
            item = frame[key].numpy() if isinstance(frame[key], torch.Tensor) else frame[key]
            self.episode_buffer[key].append(item)

        self.episode_buffer["size"] += 1


def detect_dataset_format(src_path: str) -> str:
    """Detect whether the dataset follows old or new format structure"""
    src_path = Path(src_path)

    # Check for old format indicators
    if (src_path / "task_info").exists() and (src_path / "proprio_stats").exists():
        return "old"

    # Check for new format indicators
    # Look for pattern: job_id/robot_id/episode_id with aligned_joints.h5
    # The test_data/2810125 directory contains job_id directories
    subdirs = [d for d in src_path.iterdir() if d.is_dir()]
    if subdirs:
        # Check if we have nested structure with aligned_joints.h5
        for job_dir in subdirs:  # These are job_id directories like 3335477
            if not job_dir.is_dir():
                continue
            for (
                robot_dir
            ) in job_dir.iterdir():  # These are robot_id directories like A2D0015AB00061
                if not robot_dir.is_dir():
                    continue
                for (
                    episode_dir
                ) in robot_dir.iterdir():  # These are episode_id directories like 12052353
                    if episode_dir.is_dir() and (episode_dir / "aligned_joints.h5").exists():
                        return "new"

    return "unknown"


def load_local_dataset_old_format(episode_id: int, src_path: str, task_id: int) -> list | None:
    """Load local dataset from old format and return a dict with observations and actions"""

    # --- Load task info for this specific call ---
    task_json_path = Path(src_path) / f"task_info/task_{task_id}.json"
    task_info_list = None
    if task_json_path.exists():
        try:
            with open(task_json_path, "r") as f:
                task_info_list = json.load(f)
        except json.JSONDecodeError:
            print(f"Warning: Failed to decode JSON {task_json_path} for episode {episode_id}")
            task_info_list = []  # Treat as empty if decode fails
    else:
        print(f"Warning: Task info JSON not found at {task_json_path} for episode {episode_id}")
        task_info_list = []  # Treat as empty if not found

    # --- Find action_config for this episode_id ---
    episode_action_config = None
    if isinstance(task_info_list, list):  # Check if loading succeeded and it's a list
        for item in task_info_list:
            # Ensure episode_id exists and compare as int
            if "episode_id" in item and int(item["episode_id"]) == episode_id:
                episode_action_config = item.get("label_info", {}).get("action_config")
                break

    default_action_text = "N/A"

    ob_dir = Path(src_path) / f"observations/{task_id}/{episode_id}"
    proprio_dir = Path(src_path) / f"proprio_stats/{task_id}/{episode_id}"

    with h5py.File(proprio_dir / "proprio_stats.h5") as f:
        state_joint = np.array(f["state/joint/position"])
        state_effector = np.clip((np.array(f["state/effector/position"]) - 35.0) / (120.0 - 35.0), 0.0, 1.0)
        state_head = np.array(f["state/head/position"])
        state_waist = np.array(f["state/waist/position"])
        action_joint = np.array(f["action/joint/position"])
        action_effector = np.clip((np.array(f["action/effector/position"]) - 35.0) / (120.0 - 35.0), 0.0, 1.0)
        action_head = np.array(f["action/head/position"])
        action_waist = np.array(f["action/waist/position"])
        action_velocity = np.array(f["action/robot/velocity"])

    # State (20 DOF): joint(14) + effector(2) + head(2) + waist(2)
    states_value = np.hstack(
        [
            state_joint,
            state_effector,
            state_head,
            state_waist,
        ]
    ).astype(np.float32)
    assert (
        action_joint.shape[0] == action_effector.shape[0]
    ), f"shape of action_joint:{action_joint.shape};shape of action_effector:{action_effector.shape}"
    # Action (22 DOF): joint(14) + effector(2) + head(2) + waist(2) + velocity(2)
    action_value = np.hstack(
        [
            action_joint,
            action_effector,
            action_head,
            action_waist,
            action_velocity,
        ]
    ).astype(np.float32)

    num_frames = len(states_value)

    # --- Create frame -> action_text mapping ---
    frame_action_texts = [default_action_text] * num_frames
    if episode_action_config:  # Only proceed if config was found
        for action in episode_action_config:
            start = action.get("start_frame")
            end = action.get("end_frame")
            text = action.get("action_text", default_action_text)

            if start is None or end is None:
                continue

            clamped_start = max(0, start)
            clamped_end = min(num_frames, end)
            for i in range(clamped_start, clamped_end):
                frame_action_texts[i] = text

    frames = [
        {
            "observation.state": states_value[i],
            "action": action_value[i],
            "annotation.language.action_text": [frame_action_texts[i]],  # Add action_text here
        }
        for i in range(num_frames)
    ]

    v_path = ob_dir / "videos"
    videos = {
        "observation.images.top_head": v_path / HEAD_COLOR,
        "observation.images.hand_left": v_path / HAND_LEFT_COLOR,
        "observation.images.hand_right": v_path / HAND_RIGHT_COLOR,
    }
    return frames, videos


def load_local_dataset_new_format(episode_path: str) -> list | None:
    """Load local dataset from new format and return a dict with observations and actions"""

    episode_dir = Path(episode_path)

    # Load data info JSON
    data_info_path = episode_dir / "data_info.json"
    episode_action_config = None
    default_action_text = "N/A"

    if data_info_path.exists():
        try:
            with open(data_info_path, "r") as f:
                data_info = json.load(f)
                episode_action_config = data_info.get("label_info", {}).get("action_config")
        except json.JSONDecodeError:
            print(f"Warning: Failed to decode JSON {data_info_path}")

    # Load aligned joints data
    joints_path = episode_dir / "aligned_joints.h5"
    if not joints_path.exists():
        print(f"Warning: aligned_joints.h5 not found at {joints_path}")
        return None

    with h5py.File(joints_path) as f:
        # Extract state data - using same mapping as old format for compatibility
        state_joint = np.array(f["state/joint/position"])
        state_head = np.array(f["state/head/position"])
        state_waist = np.array(f["state/waist/position"])

        # For new format, use separate left/right effector position data
        # Normalize from raw [35, 120] to [0, 1]
        state_left_effector = np.array(f["state/left_effector/position"])  # Shape: (N, 1)
        state_right_effector = np.array(f["state/right_effector/position"])  # Shape: (N, 1)
        state_effector = np.clip(
            np.column_stack([state_left_effector.flatten(), state_right_effector.flatten()]) - 35.0,
            0.0, 85.0
        ) / 85.0  # Shape: (N, 2), range [0, 1]

        # Extract action data
        action_joint = np.array(f["action/joint/position"])
        action_head = np.array(f["action/head/position"])
        action_waist = np.array(f["action/waist/position"])

        # For new format, use separate left/right effector position data for actions
        # Normalize from raw [35, 120] to [0, 1]
        action_left_effector = np.array(f["action/left_effector/position"])  # Shape: (N, 1)
        action_right_effector = np.array(f["action/right_effector/position"])  # Shape: (N, 1)
        action_effector = np.clip(
            np.column_stack([action_left_effector.flatten(), action_right_effector.flatten()]) - 35.0,
            0.0, 85.0
        ) / 85.0  # Shape: (N, 2), range [0, 1]

        # Get robot velocity (N,) or (N, 2)
        action_velocity_raw = np.array(f["action/robot/velocity"])
        if action_velocity_raw.ndim == 1:
            # Scalar velocity: pad second component with zeros (linear_x only)
            action_velocity = np.column_stack(
                [action_velocity_raw, np.zeros_like(action_velocity_raw)]
            )
        else:
            action_velocity = action_velocity_raw[:, :2]

    # State (20 DOF): joint(14) + effector(2) + head(2) + waist(2)
    states_value = np.hstack(
        [
            state_joint,
            state_effector,
            state_head,
            state_waist,
        ]
    ).astype(np.float32)
    # Action (22 DOF): joint(14) + effector(2) + head(2) + waist(2) + velocity(2)
    action_value = np.hstack(
        [
            action_joint,
            action_effector,
            action_head,
            action_waist,
            action_velocity,
        ]
    ).astype(np.float32)

    num_frames = len(states_value)

    # --- Create frame -> action_text mapping ---
    frame_action_texts = [default_action_text] * num_frames
    if episode_action_config:  # Only proceed if config was found
        for action in episode_action_config:
            start = action.get("start_frame")
            end = action.get("end_frame")
            # Use English action text if available, otherwise use Chinese
            text = action.get("english_action_text") or action.get(
                "action_text", default_action_text
            )

            if start is None or end is None:
                continue

            clamped_start = max(0, start)
            clamped_end = min(num_frames, end)
            for i in range(clamped_start, clamped_end):
                frame_action_texts[i] = text

    frames = [
        {
            "observation.state": states_value[i],
            "action": action_value[i],
            "annotation.language.action_text": [frame_action_texts[i]],
        }
        for i in range(num_frames)
    ]

    # Videos are at episode level in new format
    videos = {
        "observation.images.top_head": episode_dir / HEAD_COLOR,
        "observation.images.hand_left": episode_dir / HAND_LEFT_COLOR,
        "observation.images.hand_right": episode_dir / HAND_RIGHT_COLOR,
    }
    return frames, videos


def load_local_dataset(
    episode_id: int,
    src_path: str,
    task_id: int = None,
    episode_path: str = None,
    format_type: str = "old",
) -> list | None:
    """Load local dataset and return a dict with observations and actions

    Args:
        episode_id: Episode ID (used for old format)
        src_path: Source path (used for old format)
        task_id: Task ID (used for old format)
        episode_path: Full path to episode directory (used for new format)
        format_type: "old" or "new" format
    """
    if format_type == "old":
        return load_local_dataset_old_format(episode_id, src_path, task_id)
    elif format_type == "new":
        return load_local_dataset_new_format(episode_path)
    else:
        raise ValueError(f"Unknown format type: {format_type}")


def get_task_instruction_old_format(task_json_path: str) -> str:
    """Get task language instruction from old format"""
    with open(task_json_path, "r") as f:
        task_info = json.load(f)
    task_name = task_info[0]["task_name"]
    task_init_scene = task_info[0]["init_scene_text"]
    task_instruction = f"{task_name}.{task_init_scene}"
    print(f"Get Task Instruction <{task_instruction}>")
    return task_instruction


def get_task_instruction_new_format(episode_paths: list) -> str:
    """Get task language instruction from new format - use first episode's data_info.json"""
    if not episode_paths:
        return "Unknown Task"

    first_episode_path = Path(episode_paths[0])
    data_info_path = first_episode_path / "data_info.json"

    if data_info_path.exists():
        try:
            with open(data_info_path, "r") as f:
                data_info = json.load(f)
            # Use English task name if available, otherwise use Chinese
            task_name = data_info.get("english_task_name") or data_info.get(
                "task_name", "Unknown Task"
            )
            task_instruction = task_name
            print(
                f"Get Task Instruction <{task_instruction}> "
                f"(english_task_name: {data_info.get('english_task_name')}, "
                f"task_name: {data_info.get('task_name')})"
            )
            return task_instruction
        except json.JSONDecodeError:
            print(f"Warning: Failed to decode JSON {data_info_path}")

    return "Unknown Task"


def get_task_instruction(
    task_json_path: str = None, episode_paths: list = None, format_type: str = "old"
) -> str:
    """Get task language instruction"""
    if format_type == "old":
        return get_task_instruction_old_format(task_json_path)
    elif format_type == "new":
        return get_task_instruction_new_format(episode_paths)
    else:
        raise ValueError(f"Unknown format type: {format_type}")


def load_new_format_episode(episode_path):
    """Helper function for multiprocessing - load new format episode"""
    return load_local_dataset(
        episode_id=0, src_path="", episode_path=episode_path, format_type="new"
    )


def create_tasks_jsonl(tgt_path: str, repo_id: str, task_name: str, all_action_texts: set) -> dict:
    """Create tasks.jsonl file with unique task names and action texts."""
    meta_path = os.path.join(tgt_path, repo_id, "meta")
    os.makedirs(meta_path, exist_ok=True)

    tasks_jsonl_path = os.path.join(meta_path, "tasks.jsonl")

    # Create a list of unique tasks combining task name and action texts
    # Remove task_name from action texts if it exists to avoid duplicates
    unique_action_texts = all_action_texts - {task_name}

    tasks = [task_name]  # First entry is the main task
    tasks.extend(sorted(unique_action_texts))  # Then all unique action texts (excluding task name)

    # Check if file already exists and has correct content
    should_write = True
    if os.path.exists(tasks_jsonl_path):
        try:
            with open(tasks_jsonl_path, "r") as f:
                existing_content = f.read().strip()

            # Generate expected content
            expected_lines = []
            for i, task in enumerate(tasks):
                task_entry = {"task_index": i, "task": task}
                expected_lines.append(json.dumps(task_entry))
            expected_content = "\n".join(expected_lines)

            # If content matches, don't rewrite
            if existing_content == expected_content:
                should_write = False
        except Exception as e:
            print(f"Warning: Failed to read tasks.jsonl: {e}")
            # If there's any issue reading, we'll rewrite
            pass

    if should_write:
        # Write tasks.jsonl (overwrite to ensure clean content)
        with open(tasks_jsonl_path, "w") as f:
            for i, task in enumerate(tasks):
                task_entry = {"task_index": i, "task": task}
                f.write(json.dumps(task_entry) + "\n")

        print(f"Created tasks.jsonl with {len(tasks)} entries at {tasks_jsonl_path}")
    else:
        print(f"tasks.jsonl already exists with correct content at {tasks_jsonl_path}")

    # Create mapping for lookups
    task_to_index = {task: i for i, task in enumerate(tasks)}

    return task_to_index


def main(
    src_path: str,
    tgt_path: str,
    task_id: int = None,
    repo_id: str = None,
    task_info_json: str = None,
    debug: bool = False,
    chunk_size: int = 10,
):
    # Detect dataset format
    format_type = detect_dataset_format(src_path)
    print(f"Detected dataset format: {format_type}")

    if format_type == "unknown":
        raise ValueError(f"Unable to detect dataset format for path: {src_path}")

    # Collect all unique action texts first
    all_action_texts = set()

    # Initialize dataset
    if not repo_id:
        if format_type == "old":
            repo_id = f"agibotworld/task_{task_id}"
        else:  # new format
            # Use the top-level directory name as task_id
            task_id = Path(src_path).name
            repo_id = f"agibotworld/task_{task_id}"

    dataset = AgiBotDataset.create(
        repo_id=repo_id,
        root=f"{tgt_path}/{repo_id}",
        fps=30,
        robot_type="a2d",
        features=FEATURES,
    )

    if format_type == "old":
        # Old format processing
        task_name = get_task_instruction(task_json_path=task_info_json, format_type="old")

        all_subdir = sorted(
            [f.as_posix() for f in Path(src_path).glob(f"observations/{task_id}/*") if f.is_dir()]
        )

        if debug:
            all_subdir = all_subdir[:2]

        # Get all episode id
        all_subdir_eids = [int(Path(path).name) for path in all_subdir]
        all_subdir_episode_desc = [task_name] * len(all_subdir_eids)

        # First pass: collect all unique action texts
        print("Collecting unique action texts...")
        for episode_id in tqdm(all_subdir_eids, desc="Scanning for action texts"):
            frames_data, _ = load_local_dataset(
                episode_id, src_path=src_path, task_id=task_id, format_type="old"
            )
            if frames_data:
                for frame in frames_data:
                    action_text = frame["annotation.language.action_text"][0]
                    all_action_texts.add(action_text)

        # Create tasks.jsonl with all unique texts
        task_to_index = create_tasks_jsonl(tgt_path, repo_id, task_name, all_action_texts)

        # Pass the task mapping to the dataset for custom indexing
        dataset._custom_task_to_index = task_to_index

        # Process in chunks to reduce memory usage
        for chunk_start in tqdm(
            range(0, len(all_subdir_eids), chunk_size), desc="Processing chunks"
        ):
            chunk_end = min(chunk_start + chunk_size, len(all_subdir_eids))
            chunk_eids = all_subdir_eids[chunk_start:chunk_end]
            chunk_descs = all_subdir_episode_desc[chunk_start:chunk_end]

            # Process only this chunk
            if debug:
                raw_datasets_chunk = [
                    load_local_dataset(
                        subdir, src_path=src_path, task_id=task_id, format_type="old"
                    )
                    for subdir in tqdm(chunk_eids, desc="Loading chunk data")
                ]
            else:
                raw_datasets_chunk = process_map(
                    partial(
                        load_local_dataset, src_path=src_path, task_id=task_id, format_type="old"
                    ),
                    chunk_eids,
                    max_workers=os.cpu_count() // 2,
                    desc=f"Loading chunk {chunk_start//chunk_size + 1}/"
                    f"{(len(all_subdir_eids) + chunk_size - 1)//chunk_size}",
                )

            # Filter out None results
            valid_datasets = [
                (ds, desc) for ds, desc in zip(raw_datasets_chunk, chunk_descs) if ds is not None
            ]

            # Process each dataset in the chunk
            for raw_dataset, episode_desc in tqdm(
                valid_datasets, desc="Processing episodes in chunk"
            ):
                for raw_dataset_sub in tqdm(raw_dataset[0], desc="Processing frames", leave=False):
                    # Convert string annotation to int index
                    action_text = raw_dataset_sub["annotation.language.action_text"][0]
                    raw_dataset_sub["annotation.language.action_text"] = [
                        task_to_index[action_text]
                    ]
                    raw_dataset_sub["annotation.agibot.tasks"] = [task_to_index[episode_desc]]
                    dataset.add_frame(raw_dataset_sub)
                dataset.save_episode(task=episode_desc, videos=raw_dataset[1])

            # Clear memory after each chunk
            raw_datasets_chunk = None
            valid_datasets = None
            gc.collect()

    else:  # new format
        # Find all episode directories
        all_episode_paths = []
        src_path = Path(src_path)

        # Walk through job_id/robot_id/episode_id structure (test_data/2810125 contains job_id directories)
        for job_dir in src_path.iterdir():  # These are job_id directories like 3335477
            if not job_dir.is_dir():
                continue
            for (
                robot_dir
            ) in job_dir.iterdir():  # These are robot_id directories like A2D0015AB00061
                if not robot_dir.is_dir():
                    continue
                for (
                    episode_dir
                ) in robot_dir.iterdir():  # These are episode_id directories like 12052353
                    if episode_dir.is_dir() and (episode_dir / "aligned_joints.h5").exists():
                        all_episode_paths.append(str(episode_dir))

        all_episode_paths = sorted(all_episode_paths)

        if debug:
            all_episode_paths = all_episode_paths[:2]

        # Get task name from first episode
        task_name = get_task_instruction(episode_paths=all_episode_paths, format_type="new")
        all_episode_descs = [task_name] * len(all_episode_paths)

        # First pass: collect all unique action texts
        print("Collecting unique action texts...")
        for episode_path in tqdm(all_episode_paths, desc="Scanning for action texts"):
            frames_data, _ = load_local_dataset(
                episode_id=0, src_path="", episode_path=episode_path, format_type="new"
            )
            if frames_data:
                for frame in frames_data:
                    action_text = frame["annotation.language.action_text"][0]
                    all_action_texts.add(action_text)

        # Create tasks.jsonl with all unique texts
        task_to_index = create_tasks_jsonl(tgt_path, repo_id, task_name, all_action_texts)

        # Pass the task mapping to the dataset for custom indexing
        dataset._custom_task_to_index = task_to_index

        # Process in chunks to reduce memory usage
        for chunk_start in tqdm(
            range(0, len(all_episode_paths), chunk_size), desc="Processing chunks"
        ):
            chunk_end = min(chunk_start + chunk_size, len(all_episode_paths))
            chunk_paths = all_episode_paths[chunk_start:chunk_end]
            chunk_descs = all_episode_descs[chunk_start:chunk_end]

            # Process only this chunk
            if debug:
                raw_datasets_chunk = [
                    load_local_dataset(
                        episode_id=0, src_path="", episode_path=episode_path, format_type="new"
                    )
                    for episode_path in tqdm(chunk_paths, desc="Loading chunk data")
                ]
            else:
                raw_datasets_chunk = process_map(
                    load_new_format_episode,
                    chunk_paths,
                    max_workers=os.cpu_count() // 2,
                    desc=f"Loading chunk {chunk_start//chunk_size + 1}/"
                    f"{(len(all_episode_paths) + chunk_size - 1)//chunk_size}",
                )

            # Filter out None results
            valid_datasets = [
                (ds, desc) for ds, desc in zip(raw_datasets_chunk, chunk_descs) if ds is not None
            ]

            # Process each dataset in the chunk
            for raw_dataset, episode_desc in tqdm(
                valid_datasets, desc="Processing episodes in chunk"
            ):
                for raw_dataset_sub in tqdm(raw_dataset[0], desc="Processing frames", leave=False):
                    # Convert string annotation to int index
                    action_text = raw_dataset_sub["annotation.language.action_text"][0]
                    raw_dataset_sub["annotation.language.action_text"] = [
                        task_to_index[action_text]
                    ]
                    raw_dataset_sub["annotation.agibot.tasks"] = [task_to_index[episode_desc]]
                    dataset.add_frame(raw_dataset_sub)
                dataset.save_episode(task=episode_desc, videos=raw_dataset[1])

            # Clear memory after each chunk
            raw_datasets_chunk = None
            valid_datasets = None
            gc.collect()

    # Only consolidate at the end
    try:
        dataset.consolidate(run_compute_stats=False)
    except Exception as e:
        print(f"Warning: Consolidation failed with error: {e}")
        print("Dataset conversion completed but may not be fully consolidated.")

    # Generate modality.json file in the dataset meta directory
    meta_path = os.path.join(tgt_path, repo_id, "meta")
    os.makedirs(meta_path, exist_ok=True)
    generate_modality_json(meta_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Convert AgiBot dataset to LeRobot format. Supports both old and new format datasets."
    )
    parser.add_argument(
        "--src_path", type=str, required=True, help="Path to source dataset directory"
    )
    parser.add_argument(
        "--task_id",
        type=str,
        required=False,
        help="Task ID (required for old format, optional for new format)",
    )
    parser.add_argument(
        "--tgt_path", type=str, required=True, help="Path to target output directory"
    )
    parser.add_argument(
        "--repo_id",
        type=str,
        required=False,
        help="Repository ID for the dataset (auto-generated if not provided)",
    )
    parser.add_argument(
        "--debug", action="store_true", help="Process only first 2 episodes for debugging"
    )
    parser.add_argument(
        "--chunk_size", type=int, default=10, help="Number of episodes to process at once"
    )
    args = parser.parse_args()

    # Detect format first
    format_type = detect_dataset_format(args.src_path)
    print(f"Auto-detected format: {format_type}")

    # Validate arguments based on format
    if format_type == "old":
        if not args.task_id:
            parser.error("--task_id is required for old format datasets")

        task_id = int(args.task_id)
        json_file = f"{args.src_path}/task_info/task_{args.task_id}.json"

        if not Path(json_file).exists():
            parser.error(f"Cannot find task info file: {json_file}")

        main(
            src_path=args.src_path,
            tgt_path=args.tgt_path,
            task_id=task_id,
            repo_id=args.repo_id,
            task_info_json=json_file,
            debug=args.debug,
            chunk_size=args.chunk_size,
        )

    elif format_type == "new":
        main(
            src_path=args.src_path,
            tgt_path=args.tgt_path,
            task_id=args.task_id,
            repo_id=args.repo_id,
            task_info_json=None,
            debug=args.debug,
            chunk_size=args.chunk_size,
        )

    else:
        parser.error(
            f"Unknown dataset format. Please check the directory structure at: {args.src_path}"
        )