kai0_code/scripts/split_lerobot.py

"""
Split a LeRobot-format dataset into multiple disjoint subsets by episode.

Each subset is written as a full LeRobot dataset under dst_path/split_0, split_1, ...
Episodes are shuffled (with a fixed seed) then divided so you can train separate
models on each subset and later mix them with model_arithmetic/arithmetic.py.
"""
from pathlib import Path
import json
import argparse
import shutil
import random
import pandas as pd
import numpy as np
import os
import multiprocessing as mp
from tqdm import tqdm

# Video keys expected in LeRobot dataset layout
VIDEO_KEYS = [
    "observation.images.hand_left",
    "observation.images.hand_right",
    "observation.images.top_head",
]


def split_lerobot_data(
    source_path: Path,
    dst_path: Path,
    episode_index: list[int],
    num: int,
) -> None:
    """
    Write one subset of the source LeRobot dataset to dst_path.

    episode_index: list of source episode indices to include (will be renumbered 0..n-1).
    num: subset index (used only for progress messages).
    """
    # Remap episode indices to 0, 1, 2, ...
    old_episode_index = sorted(episode_index)
    new_episode_index = list(range(len(old_episode_index)))
    old2new_episode_index = dict(zip(old_episode_index, new_episode_index))
    new2old_episode_index = dict(zip(new_episode_index, old_episode_index))

    dst_path.mkdir(parents=True, exist_ok=True)

    # Load episode stats for the selected episode indices
    with open(source_path / "meta" / "episodes_stats.jsonl", "r") as f:
        episodes_stats = [json.loads(line) for line in f if line.strip()]
    episodes_stats = [
        ep for ep in episodes_stats if ep["episode_index"] in episode_index
    ]
    assert len(episodes_stats) == len(episode_index), (
        f"episode_index count mismatch: {len(episodes_stats)} vs {len(episode_index)}"
    )
    episodes_stats.sort(key=lambda x: x["episode_index"])

    with open(source_path / "meta" / "info.json", "r") as f:
        info = json.load(f)
    chunks_size = info["chunks_size"]

    # Build new episodes_stats with renumbered episode and frame indices
    new_episodes_stats = []
    new_frame_index = 0
    for episode_stat in tqdm(
        episodes_stats,
        desc=f"Split {num}: episodes_stats",
        total=len(episodes_stats),
    ):
        index_count = episode_stat["stats"]["index"]["count"][0]
        new_idx = old2new_episode_index[episode_stat["episode_index"]]
        episode_stat["episode_index"] = new_idx
        episode_stat["stats"]["index"]["min"] = [new_frame_index]
        episode_stat["stats"]["index"]["max"] = [new_frame_index + index_count - 1]
        episode_stat["stats"]["index"]["mean"] = [(new_frame_index + new_frame_index + index_count - 1) / 2]
        episode_stat["stats"]["index"]["std"] = [np.std(range(new_frame_index, new_frame_index + index_count))]
        episode_stat["stats"]["index"]["count"] = [index_count]
        new_frame_index += index_count
        new_episodes_stats.append(episode_stat)

    (dst_path / "meta").mkdir(parents=True, exist_ok=True)
    with open(dst_path / "meta" / "episodes_stats.jsonl", "w") as f:
        for episode_stat in new_episodes_stats:
            f.write(json.dumps(episode_stat) + "\n")

    # Write info.json with updated totals
    with open(dst_path / "meta" / "info.json", "w") as f:
        info["total_episodes"] = len(old_episode_index)
        info["total_frames"] = new_frame_index
        info["total_videos"] = len(old_episode_index) * 3
        info["total_chunks"] = len(old_episode_index) // chunks_size + 1
        info["splits"] = {"train": f"0:{len(old_episode_index)}"}
        json.dump(info, f, indent=4)
    # Copy and reindex parquet files
    for new_stat in new_episodes_stats:
        new_index = new_stat["episode_index"]
        old_index = new2old_episode_index[new_index]
        old_episode_path = source_path / "data" / f"chunk-{old_index // chunks_size:03d}" / f"episode_{old_index:06d}.parquet"
        new_episode_path = dst_path / "data" / f"chunk-{new_index // chunks_size:03d}" / f"episode_{new_index:06d}.parquet"
        if not new_episode_path.parent.exists():
            new_episode_path.parent.mkdir(parents=True, exist_ok=True)
        parquet = pd.read_parquet(old_episode_path)
        parquet["index"] = parquet["index"] - parquet["index"].min() + new_stat["stats"]["index"]["min"][0]
        parquet["episode_index"] = new_index
        parquet.to_parquet(new_episode_path, index=False)

    # Update record CSV files (e.g. v1/record.csv)
    record_files = [
        p for p in os.listdir(source_path / "meta")
        if "record.csv" in p and p.lower().startswith("v")
    ]
    for record_file in record_files:
        record_data = pd.read_csv(source_path / "meta" / record_file)
        record_data["episode_index"] = record_data["video"].str.split("_").str[-1].str.split(".").str[0].astype(int)
        record_data = record_data.loc[record_data["episode_index"].isin(episode_index)]
        record_data["episode_index"] = record_data["episode_index"].map(old2new_episode_index)
        record_data["video"] = record_data["video"].map(
            lambda x: f"episode_{old2new_episode_index[int(x.split('.')[0].split('_')[-1])]:06d}.mp4"
        )
        record_data.drop(columns=["episode_index"], inplace=True)
        record_data.to_csv(dst_path / "meta" / record_file, index=False)

    # Write episodes.jsonl with renumbered episode_index
    with open(source_path / "meta" / "episodes.jsonl", "r") as f:
        episodes = [json.loads(line) for line in f if line.strip()]
    episodes = [ep for ep in episodes if ep["episode_index"] in old_episode_index]
    assert len(episodes) == len(episode_index), "episode count mismatch in episodes.jsonl"
    episodes.sort(key=lambda x: x["episode_index"])
    for episode in episodes:
        episode["episode_index"] = old2new_episode_index[episode["episode_index"]]
    with open(dst_path / "meta" / "episodes.jsonl", "w") as f:
        for episode in episodes:
            f.write(json.dumps(episode) + "\n")

    shutil.copy(source_path / "meta" / "tasks.jsonl", dst_path / "meta" / "tasks.jsonl")

    # Copy video files
    for new_stat in tqdm(
        new_episodes_stats,
        desc=f"Split {num}: videos",
        total=len(new_episodes_stats),
    ):
        new_index = new_stat["episode_index"]
        old_index = new2old_episode_index[new_index]
        for video_key in VIDEO_KEYS:
            old_episode_path = source_path / "videos" / f"chunk-{old_index // chunks_size:03d}" / video_key / f"episode_{old_index:06d}.mp4"
            new_episode_path = dst_path / "videos" / f"chunk-{new_index // chunks_size:03d}" / video_key / f"episode_{new_index:06d}.mp4"
            if not new_episode_path.parent.exists():
                new_episode_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy(old_episode_path, new_episode_path)

    

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Split a LeRobot dataset into disjoint subsets by episode for training separate models."
    )
    parser.add_argument(
        "--source_path",
        type=str,
        required=True,
        help="Path to the source LeRobot dataset (must contain meta/, data/, videos/).",
    )
    parser.add_argument(
        "--dst_path",
        type=str,
        required=True,
        help="Output directory; subsets will be written as dst_path/split_0, split_1, ...",
    )
    parser.add_argument(
        "--split_num",
        type=int,
        default=4,
        help="Number of disjoint subsets to create (default: 4).",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for shuffling episodes before splitting (default: 42).",
    )
    args = parser.parse_args()
    source_path = Path(args.source_path)
    dst_path = Path(args.dst_path)
    dst_path.mkdir(parents=True, exist_ok=True)

    with open(source_path / "meta" / "info.json", "r") as f:
        info = json.load(f)
    total_episodes = info["total_episodes"]
    episode_indices = list(range(total_episodes))
    random.seed(args.seed)
    random.shuffle(episode_indices)
    splits = np.array_split(episode_indices, args.split_num)
    max_workers = min(args.split_num, os.cpu_count() or 1)
    with mp.Pool(processes=max_workers) as pool:
        pool.starmap(
            split_lerobot_data,
            [(source_path, dst_path / f"split_{i}", splits[i].tolist(), i) for i in range(args.split_num)],
        )


if __name__ == "__main__":
    main()