From 0098bd264ec073cb97bab3916755a97143698451 Mon Sep 17 00:00:00 2001 From: Simon Alibert Date: Wed, 23 Oct 2024 20:55:54 +0200 Subject: [PATCH] Nits --- lerobot/common/datasets/lerobot_dataset.py | 4 ++-- lerobot/scripts/control_robot.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/lerobot/common/datasets/lerobot_dataset.py b/lerobot/common/datasets/lerobot_dataset.py index 6caec09c..b32e1008 100644 --- a/lerobot/common/datasets/lerobot_dataset.py +++ b/lerobot/common/datasets/lerobot_dataset.py @@ -389,7 +389,7 @@ class LeRobotDataset(torch.utils.data.Dataset): @property def features(self) -> datasets.Features: - """Shapes for the different features.""" + """Features of the hf_dataset.""" if self.hf_dataset is not None: return self.hf_dataset.features elif self.episode_buffer is None: @@ -664,7 +664,7 @@ class LeRobotDataset(torch.utils.data.Dataset): self.episode_dicts.append(episode_dict) append_jsonl(episode_dict, self.root / EPISODES_PATH) - def delete_episode(self) -> None: + def clear_episode_buffer(self) -> None: episode_index = self.episode_buffer["episode_index"] if self.image_writer is not None: for cam_key in self.camera_keys: diff --git a/lerobot/scripts/control_robot.py b/lerobot/scripts/control_robot.py index 9ef50ced..1185db20 100644 --- a/lerobot/scripts/control_robot.py +++ b/lerobot/scripts/control_robot.py @@ -299,7 +299,7 @@ def record( log_say("Re-record episode", play_sounds) events["rerecord_episode"] = False events["exit_early"] = False - dataset.delete_episode() + dataset.clear_episode_buffer() continue dataset.add_episode(task)