tangger/lerobot
1
0
Fork 1
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

add support for failed episode

Browse Source
This commit is contained in:
haixuantao
2024-05-28 21:20:43 +02:00
parent 743812e471
commit 1a6602e77c
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
lerobot/common/datasets/push_dataset_to_hub/aloha_dora_format.py
View File

@@ -52,6 +52,8 @@ def load_from_raw(raw_dir: Path, out_dir: Path, fps: int):
key = path.stem # action or observation.state or ... key = path.stem # action or observation.state or ...
if key == reference_key: if key == reference_key:
continue continue
elif "failed_episode_index" in key:
continue
modality_df = pd.read_parquet(path) modality_df = pd.read_parquet(path)
modality_df = modality_df[["timestamp_utc", key]] modality_df = modality_df[["timestamp_utc", key]]
df = pd.merge_asof( df = pd.merge_asof(
@@ -66,7 +68,7 @@ def load_from_raw(raw_dir: Path, out_dir: Path, fps: int):
# because some cameras didnt start recording yet. # because some cameras didnt start recording yet.
df = df.dropna(axis=0) df = df.dropna(axis=0)
# Remove rows with episode_index -1 which indicates a failed episode # Remove rows with episode_index -1 which indicates data that correspond to in-between episodes
df = df[df["episode_index"] != -1] df = df[df["episode_index"] != -1]
# dora only use arrays, so single values are encapsulated into a list # dora only use arrays, so single values are encapsulated into a list