How to remove a specfic episode in collected data

[milong26]

When I collect data, I may fail at grasping, I want to remove these failed episodes from whole dataset, how can I do?

Thanks!

[Tavish9]

Hi, there are two ways you can do.

1. Re-generate a new dataset.

basic steps:

old_dataset = LeRobotDataset(
    repo_id="",
    root="your_old_local_dir",
)

new_dataset = LeRobotDataset(
    repo_id="",
    root="your_new_local_dir",
)

current_episode_index = 0
for data_item in old_dataset:
    episode_index = data_item["episode_index"]

    if episode_index != current_episode_index:
        new_dataset.save_episode()
        current_episode_index = episode_index

    # your filter logic here, if failed, continue to next item

    new_dataset.add_frame(data_item)  # check key and dtype

2. filter failed episodes on the fly (recommended)

basic steps:

dataset = LeRobotDataset(
    repo_id="",
    root="your_old_local_dir",
)

failed_episodes = set()

for data_item in dataset:
    episode_index = data_item["episode_index"]

    if failed:
        failed_episodes.add(episode_index)

successful_episodes = list(dataset.meta.episodes.keys() - failed_episodes)

# save successful_episodes on the disk for future use

dataset = LeRobotDataset(
    repo_id="",
    root="your_old_local_dir",
    episodes=list(successful_episodes),
)