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

Nits

This commit is contained in:
Simon Alibert 2024-10-23 20:55:54 +02:00
parent 1aba80d93f
commit 0098bd264e
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -389,7 +389,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
@property @property
def features(self) -> datasets.Features: def features(self) -> datasets.Features:
"""Shapes for the different features.""" """Features of the hf_dataset."""
if self.hf_dataset is not None: if self.hf_dataset is not None:
return self.hf_dataset.features return self.hf_dataset.features
elif self.episode_buffer is None: elif self.episode_buffer is None:
@ -664,7 +664,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
self.episode_dicts.append(episode_dict) self.episode_dicts.append(episode_dict)
append_jsonl(episode_dict, self.root / EPISODES_PATH) 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"] episode_index = self.episode_buffer["episode_index"]
if self.image_writer is not None: if self.image_writer is not None:
for cam_key in self.camera_keys: for cam_key in self.camera_keys:

2
lerobot/scripts/control_robot.py
View File

@ -299,7 +299,7 @@ def record(
log_say("Re-record episode", play_sounds) log_say("Re-record episode", play_sounds)
events["rerecord_episode"] = False events["rerecord_episode"] = False
events["exit_early"] = False events["exit_early"] = False
dataset.delete_episode() dataset.clear_episode_buffer()
continue continue
dataset.add_episode(task) dataset.add_episode(task)