Refactor dataset features

lerobot/common/robot_devices/robots/manipulator.py

@@ -226,13 +226,42 @@ class ManipulatorRobot:
         self.is_connected = False
         self.logs = {}
 
-        action_names = [f"{arm}_{motor}" for arm, bus in self.leader_arms.items() for motor in bus.motors]
-        state_names = [f"{arm}_{motor}" for arm, bus in self.follower_arms.items() for motor in bus.motors]
-        self.names = {
-            "action": action_names,
-            "observation.state": state_names,
+    def get_motor_names(self, arm: dict[str, MotorsBus]) -> list:
+        return [f"{arm}_{motor}" for arm, bus in arm.items() for motor in bus.motors]
+
+    @property
+    def camera_features(self) -> dict:
+        cam_ft = {}
+        for cam_key, cam in self.cameras.items():
+            key = f"observation.images.{cam_key}"
+            cam_ft[key] = {
+                "shape": (cam.width, cam.height, cam.channels),
+                "names": ["width", "height", "channels"],
+                "info": None,
+            }
+        return cam_ft
+
+    @property
+    def motor_features(self) -> dict:
+        action_names = self.get_motor_names(self.leader_arms)
+        state_names = self.get_motor_names(self.leader_arms)
+        return {
+            "action": {
+                "dtype": "float32",
+                "shape": (len(action_names),),
+                "names": action_names,
+            },
+            "observation.state": {
+                "dtype": "float32",
+                "shape": (len(state_names),),
+                "names": state_names,
+            },
         }
 
+    @property
+    def features(self):
+        return {**self.motor_features, **self.camera_features}
+
     @property
     def has_camera(self):
         return len(self.cameras) > 0