File size: 8,577 Bytes
5b6e415 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 | """
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() |