fix bugs in feature shape

2026-05-11 12:09:41 +00:00 · 2025-02-20 20:54:24 +08:00
parent d2456c1506
commit 6e7bdc3e9e
1 changed files with 3 additions and 4 deletions
@@ -107,16 +107,15 @@ def generate_features_from_raw(builder: tfds.core.DatasetBuilder, use_videos: bo
        elif action_encoding == ActionEncoding.JOINT_POS:
            action_names = [f"motor_{i}" for i in range(7)] + ["gripper"]

-
    DEFAULT_FEATURES = {
        "observation.state": {
            "dtype": "float32",
-            "shape": (8,),
+            "shape": (len(state_names),),
            "names": {"motors": state_names},
        },
        "action": {
            "dtype": "float32",
-            "shape": (7,),
+            "shape": (len(action_names),),
            "names": {"motors": action_names},
        },
    }