From 083f72c1d0d3bd4a6b46425d4925e8b897e14fea Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Fri, 11 Apr 2025 15:46:01 +0200 Subject: [PATCH] fix: Check if robot_type is not None before getting its value --- lerobot/common/datasets/lerobot_dataset.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/lerobot/common/datasets/lerobot_dataset.py b/lerobot/common/datasets/lerobot_dataset.py index e51a163d..5e13f9e5 100644 --- a/lerobot/common/datasets/lerobot_dataset.py +++ b/lerobot/common/datasets/lerobot_dataset.py @@ -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(