Improvements

Author: Cadene

lerobot/common/datasets/lerobot_dataset.py

@@ -243,6 +243,10 @@ def write_video_info(self) -> None:
         Warning: this function writes info from first episode videos, implicitly assuming that all videos have
         been encoded the same way. Also, this means it assumes the first episode exists.
         """
+        # TODO(rcadene): What should we do here?
+        if "videos" not in self.info:
+            self.info["videos"] = {}
+
         for key in self.video_keys:
             if key not in self.info["videos"]:
                 video_path = self.root / self.get_video_file_path(ep_index=0, vid_key=key)
@@ -260,6 +264,8 @@ def create(
         robot_type: str | None = None,
         features: dict | None = None,
         use_videos: bool = True,
+        # tags: list[str] | None = None,
+        # license_type: str | None = None,
     ) -> "LeRobotDatasetMetadata":
         """Creates metadata for a LeRobotDataset."""
         obj = cls.__new__(cls)
@@ -283,8 +289,12 @@ def create(
         else:
             features = {**features, **DEFAULT_FEATURES}
 
+        # TODO(rcadene): implement sanity check for features
+
         obj.tasks, obj.stats, obj.episodes = {}, {}, []
         obj.info = create_empty_dataset_info(CODEBASE_VERSION, fps, robot_type, features, use_videos)
+        # obj.tags = tags
+        # obj.license_type = license_type
         if len(obj.video_keys) > 0 and not use_videos:
             raise ValueError()
         write_json(obj.info, obj.root / INFO_PATH)
@@ -646,7 +656,7 @@ def _create_episode_buffer(self, episode_index: int | None = None) -> dict:
         current_ep_idx = self.meta.total_episodes if episode_index is None else episode_index
         return {
             "size": 0,
-            **{key: [] if key != "episode_index" else current_ep_idx for key in self.features},
+            **{key: current_ep_idx if key == "episode_index" else [] for key in self.features},
         }
 
     def _get_image_file_path(self, episode_index: int, image_key: str, frame_index: int) -> Path:
@@ -669,6 +679,9 @@ def add_frame(self, frame: dict) -> None:
         temporary directory — nothing is written to disk. To save those frames, the 'add_episode()' method
         then needs to be called.
         """
+        # TODO(rcadene): Add sanity check for the input, check it's numpy or torch,
+        # check the dtype and shape matches, etc.
+
         frame_index = self.episode_buffer["size"]
         for key, ft in self.features.items():
             if key == "frame_index":
@@ -705,6 +718,11 @@ def add_episode(self, task: str, encode_videos: bool = False) -> None:
             # TODO(aliberts): Add option to use existing episode_index
             raise NotImplementedError()
 
+        if episode_length == 0:
+            raise ValueError(
+                "You must add one or several frames with `add_frame` before calling `add_episode`."
+            )
+
         task_index = self.meta.get_task_index(task)
 
         if not set(self.episode_buffer.keys()) == set(self.features):
@@ -719,11 +737,14 @@ def add_episode(self, task: str, encode_videos: bool = False) -> None:
                 self.episode_buffer[key] = np.full((episode_length,), episode_index)
             elif key == "task_index":
                 self.episode_buffer[key] = np.full((episode_length,), task_index)
-            elif ft["dtype"] in ["image", "video"]:
+            elif ft["dtype"] == "image":
                 continue
+            elif ft["dtype"] == "video":
+                del self.episode_buffer[key]
             elif ft["shape"][0] == 1:
                 self.episode_buffer[key] = torch.tensor(self.episode_buffer[key])
             elif ft["shape"][0] > 1:
+                # TODO(rcadene): why torch over here?
                 self.episode_buffer[key] = torch.stack(self.episode_buffer[key])
             else:
                 raise ValueError()

lerobot/common/datasets/utils.py

@@ -435,7 +435,7 @@ def create_branch(repo_id, *, branch: str, repo_type: str | None = None) -> None
 
 
 def create_lerobot_dataset_card(
-    tags: list | None = None, text: str | None = None, info: dict | None = None
+    tags: list | None = None, text: str | None = None, info: dict | None = None, license: str = "apache-2.0"
 ) -> DatasetCard:
     card = DatasetCard(DATASET_CARD_TEMPLATE)
     card.data.configs = [
@@ -445,6 +445,7 @@ def create_lerobot_dataset_card(
         }
     ]
     card.data.task_categories = ["robotics"]
+    card.data.license = license
     card.data.tags = ["LeRobot"]
     if tags is not None:
         card.data.tags += tags

lerobot/common/robot_devices/control_utils.py

@@ -17,7 +17,6 @@
 
 from lerobot.common.datasets.image_writer import safe_stop_image_writer
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.datasets.utils import _get_info_from_robot
 from lerobot.common.policies.factory import make_policy
 from lerobot.common.robot_devices.robots.utils import Robot
 from lerobot.common.robot_devices.utils import busy_wait
@@ -334,7 +333,8 @@ def sanity_check_dataset_name(repo_id, policy):
 
 
 def sanity_check_dataset_robot_compatibility(dataset, robot, fps, use_videos):
-    robot_type, keys, image_keys, video_keys, shapes, names = _get_info_from_robot(robot, use_videos)
+    # TODO(rcadene): fix that before merging
+    robot_type, keys, image_keys, video_keys, shapes, names = _get_info_from_robot(robot, use_videos)  # noqa
 
     fields = [
         ("robot_type", dataset.meta.info["robot_type"], robot_type),

lerobot/scripts/push_dataset_to_hub.py

@@ -117,10 +117,14 @@ def push_meta_data_to_hub(repo_id: str, meta_data_dir: str | Path, revision: str
 
 
 def push_dataset_card_to_hub(
-    repo_id: str, revision: str | None, tags: list | None = None, text: str | None = None
+    repo_id: str,
+    revision: str | None,
+    tags: list | None = None,
+    text: str | None = None,
+    license: str = "apache-2.0",
 ):
     """Creates and pushes a LeRobotDataset Card with appropriate tags to easily find it on the hub."""
-    card = create_lerobot_dataset_card(tags=tags, text=text)
+    card = create_lerobot_dataset_card(tags=tags, text=text, license=license)
     card.push_to_hub(repo_id=repo_id, repo_type="dataset", revision=revision)