wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: Check if robot_type is not None before getting its value

This commit is contained in:
CarolinePascal 2025-04-11 15:46:01 +02:00
parent 5384309e6f
commit 083f72c1d0
No known key found for this signature in database
1 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
lerobot/common/datasets/lerobot_dataset.py
View File

@ -964,8 +964,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
self._save_image(frame[key], img_path)
self.episode_buffer[key].append(str(img_path))
elif self.features[key]["dtype"] == "audio":
if self.meta.robot_type.startswith(
"lekiwi"
if (
self.meta.robot_type is not None and self.meta.robot_type.startswith("lekiwi")
): # Rw data storage should only be triggered for LeKiwi robot, for which audio is stored chunk by chunk in a visual frame-like manner
self.episode_buffer[key].append(frame[key])
else: # Otherwise, only the audio file path is stored in the episode buffer
@ -1034,8 +1034,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
if key in ["index", "episode_index", "task_index"] or ft["dtype"] in ["image", "video"]:
continue
elif ft["dtype"] == "audio":
if self.meta.robot_type.startswith(
"lekiwi"
if (
self.meta.robot_type is not None and self.meta.robot_type.startswith("lekiwi")
): # Raw data storage should only be triggered for LeKiwi robot, for which audio is stored chunk by chunk in a visual frame-like manner
episode_buffer[key] = np.concatenate(episode_buffer[key], axis=0)
continue
@ -1044,8 +1044,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
self._wait_image_writer()
self._save_episode_table(episode_buffer, episode_index)
if self.meta.robot_type.startswith(
"lekiwi"
if (
self.meta.robot_type is not None and self.meta.robot_type.startswith("lekiwi")
): # Raw data storage should only be triggered for LeKiwi robot, for which audio is stored chunk by chunk in a visual frame-like manner
for key in self.meta.audio_keys:
audio_path = self._get_raw_audio_file_path(