convert_lerobot_to_dexdata.py

"""
Example script for converting LeRobot dataset format to DexData format.

This script demonstrates how to convert LeRobot datasets (specifically galaxea_open_world_dataset)
to the DexData format.
"""

import os
import json

import pyarrow.parquet as pq
import pandas as pd
import numpy as np 
from loguru import logger
from tqdm import tqdm
import click


def get_task_list(task_root): 
    """
    Load task list from tasks.jsonl file in the given task root directory.
    Args:
        task_root: Path to the task root directory
    Returns:
        list: List of task strings
    """
    task_file = os.path.join(task_root, "meta", "tasks.jsonl")
    assert os.path.isfile(task_file), f"task file not found: {task_file}"
    tasks = []
    with open(task_file, "r") as f:
        for i, line in enumerate(f): 
            info = json.loads(line)
            task_index = info["task_index"]
            sub_task = info["task"]
            assert task_index == i, f"task index mismatch: {task_index} vs {i}"
            tasks.append(sub_task)
    return tasks


def get_latest_episode_idx(task_dir): 
    """
    Get the latest episode index in the given task directory
    by counting the number of existing episode files.
    Args:
        task_dir: Path to the task directory
    Returns:
        int: Latest episode index
    """
    fnames = os.listdir(task_dir)
    return len(fnames)


def parse_one_episode(
    df: pd.DataFrame,
    task_list,
    head_video_rel_path,
    left_wrist_video_rel_path,
    right_wrist_video_rel_path,
): 
    """
    Parse one episode dataframe into a list of data dictionaries.
    Args:
        df: DataFrame of the episode
        task_list: List of task strings
        head_video_rel_path: Relative path to the head video
        left_wrist_video_rel_path: Relative path to the left wrist video
        right_wrist_video_rel_path: Relative path to the right wrist video
    Returns:
        list: List of data dictionaries for the episode
    """
    data_list = []
    for row_id, row in df.iterrows(): 
        try: 
            # NOTE: parse all data in one row
            left_arm = np.array(row["observation.state.left_arm"])  # (6,) 6dof
            left_arm_vel = np.array(row["observation.state.left_arm.velocities"])  # (6,) 6dof
            right_arm = np.array(row["observation.state.right_arm"])  # (6,) 6dof
            right_arm_vel = np.array(row["observation.state.right_arm.velocities"])  # (6,) 6dof
            chassis = np.array(row["observation.state.chassis"])  # (10,) imu: quat(4) + rot_vel(3) + lin_acc(3)
            torso = np.array(row["observation.state.torso"])  # (4,) 4dof(pad 0 in last)
            torso_vel = np.array(row["observation.state.torso.velocities"])  # (4,) 4dof(pad 0 in last)
            left_gripper = row["observation.state.left_gripper"]  # () mm
            right_gripper = row["observation.state.right_gripper"]  # () mm
            left_ee_pose = np.array(row["observation.state.left_ee_pose"])  # (7,) ee: pos(3) + quat(4)
            right_ee_pose = np.array(row["observation.state.right_ee_pose"])  # (7,) ee: pos(3) + quat(4)

            action_left_gripper = row["action.left_gripper"]  # () 0~100
            action_right_gripper = row["action.right_gripper"]  # () 0~100
            action_chassis_vel = np.array(row["action.chassis.velocities"])  # (3,) lin_x, lin_y, ang_z
            action_torso_vel = np.array(row["action.torso.velocities"])  # (6,) torso velocities: aug(3) + lin(3) 
            action_left_arm = np.array(row["action.left_arm"])  # (6,) 6dof
            action_right_arm = np.array(row["action.right_arm"])  # (6,) 6dof

            timestamp = row["timestamp"]
            frame_index = row["frame_index"]
            episode_index = row["episode_index"]
            index = row["index"]
            coarse_task_index = row["coarse_task_index"]
            task_index = row["task_index"]
            coarse_quality_index = row["coarse_quality_index"]
            quality_index = row["quality_index"]
        
        except: 
            logger.error(f"Error parsing row {row_id}, skip this episode")
            data_list = None
            break

        state = np.concatenate(
            [
                left_arm,
                left_arm_vel,
                right_arm,
                right_arm_vel,
                chassis,
                torso,
                torso_vel,
                np.array([left_gripper]),
                np.array([right_gripper]),
                left_ee_pose,
                right_ee_pose,
            ]
        )
        action = np.concatenate(
            [
                np.array([action_left_gripper]),
                np.array([action_right_gripper]),
                action_chassis_vel,
                action_torso_vel,
                action_left_arm,
                action_right_arm,
            ]
        )

        prompt = task_list[coarse_task_index]
        # NOTE: sometimes, the task is None
        if task_list[task_index] is None: 
            print(f"Invalid task format: {task_list[task_index]}")
            data_list = None
            break
        # NOTE: subtask schema is Chinese@English
        if len(task_list[task_index].split("@")) != 2: 
            print(f"Invalid task format: {task_list[task_index]}")
            data_list = None
            break
        subtask = task_list[task_index].split("@")[1]
        data_list.append(
            {
                "images_1": {
                    "type": "video",
                    "url": head_video_rel_path,
                    "frame_idx": frame_index,
                },
                "images_2": {
                    "type": "video",
                    "url": left_wrist_video_rel_path,
                    "frame_idx": frame_index,
                },
                "images_3": {
                    "type": "video",
                    "url": right_wrist_video_rel_path,
                    "frame_idx": frame_index,
                },
                "prompt": prompt,
                "state": state.tolist(),
                "action": action.tolist(),
                "is_robot": True,
                "extra": {
                    "subtask": subtask,
                    "timestamp": timestamp,
                    "episode_index": episode_index,
                    "index": index,
                    "coarse_quality_index": coarse_quality_index,
                    "quality_index": quality_index,
                }
            }
        )

    return data_list


def save_jsonl(data_list, jsonl_path):
    """
    Save data list to a jsonl file.
    Args:
        data_list: List of data dictionaries
        jsonl_path: Path to the output jsonl file
    Returns:
        None
    """
    with open(jsonl_path, "w") as f:
        for data in data_list:
            line = json.dumps(data, ensure_ascii=False)
            f.write(line + "\n")


@click.command()
@click.option("-i", "--lerobot_dir", type=str, required=True, help="Path to the lerobot dataset root directory.")
@click.option("-o", "--output_dir", type=str, required=True, help="Path to the output dexdata directory.")
def main(lerobot_dir, output_dir): 
    """
    Convert lerobot dataset to dexdata format.
    Args:
        lerobot_dir: Path to the lerobot dataset root directory.
        output_dir: Path to the output dexdata directory.
    Returns:
        None
    """
    for task_name in os.listdir(lerobot_dir): 
        task_root = os.path.join(lerobot_dir, task_name)
        if not os.path.isdir(task_root):
            continue 
        output_jsonl_dir = os.path.join(output_dir, "jsonl", task_name)
        output_video_dir = os.path.join(output_dir, "videos", task_name)
        os.makedirs(output_jsonl_dir, exist_ok=True)
        os.makedirs(output_video_dir, exist_ok=True)
        task_list = get_task_list(task_root)
        
        task_data_dir = os.path.join(task_root, "data")
        for chunk_name in os.listdir(task_data_dir): 
            chunk_path = os.path.join(task_data_dir, chunk_name)
            head_video_dir = os.path.join(task_root, "videos", chunk_name, "observation.images.head_rgb")
            left_wrist_video_dir = os.path.join(task_root, "videos", chunk_name, "observation.images.left_wrist_rgb")
            right_wrist_video_dir = os.path.join(task_root, "videos", chunk_name, "observation.images.right_wrist_rgb")

            for episode_name in tqdm(os.listdir(chunk_path), desc=f"Processing {task_name}/{chunk_name}"): 
                episode_path = os.path.join(chunk_path, episode_name)
                if not episode_name.endswith(".parquet"):
                    continue 
                df = pq.read_table(episode_path).to_pandas()

                episode_id = get_latest_episode_idx(output_jsonl_dir)
                head_video_rel_path = os.path.join(task_name, f"episode_{episode_id:05d}_head.mp4")
                left_wrist_video_rel_path = os.path.join(task_name, f"episode_{episode_id:05d}_left_wrist.mp4")
                right_wrist_video_rel_path = os.path.join(task_name, f"episode_{episode_id:05d}_right_wrist.mp4")
                episode_data_list = parse_one_episode(
                    df,
                    task_list,
                    head_video_rel_path,
                    left_wrist_video_rel_path,
                    right_wrist_video_rel_path,
                )
                if episode_data_list is None: 
                    continue 
                episode_jsonl_path = os.path.join(output_jsonl_dir, f"episode_{episode_id:05d}.jsonl")
                save_jsonl(episode_data_list, episode_jsonl_path)

                # NOTE: Copy video files to output directory
                video_name = episode_name.replace(".parquet", ".mp4")
                head_video_path = os.path.join(head_video_dir, video_name)
                left_wrist_video_path = os.path.join(left_wrist_video_dir, video_name)
                right_wrist_video_path = os.path.join(right_wrist_video_dir, video_name)
                output_head_video_path = os.path.join(output_video_dir, f"episode_{episode_id:05d}_head.mp4")
                cmd = f"cp {head_video_path} {output_head_video_path}"
                os.system(cmd)
                output_left_wrist_video_path = os.path.join(output_video_dir, f"episode_{episode_id:05d}_left_wrist.mp4")
                cmd = f"cp {left_wrist_video_path} {output_left_wrist_video_path}"
                os.system(cmd)
                output_right_wrist_video_path = os.path.join(output_video_dir, f"episode_{episode_id:05d}_right_wrist.mp4")
                cmd = f"cp {right_wrist_video_path} {output_right_wrist_video_path}"
                os.system(cmd)
        

if __name__ == "__main__":
    main()