admin/any4lerobot
1
0
Fork 0
mirror of https://github.com/Tavish9/any4lerobot.git synced 2026-05-20 16:39:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix bugs in feature shape

Browse Source
This commit is contained in:
Tavish
2025-02-20 20:54:24 +08:00
parent d2456c1506
commit 6e7bdc3e9e
1 changed files with 3 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
openx_rlds.py
+3 -4
View File
@@ -89,7 +89,7 @@ def generate_features_from_raw(builder: tfds.core.DatasetBuilder, use_videos: bo
if state_encoding == StateEncoding.POS_EULER: if state_encoding == StateEncoding.POS_EULER:
state_names = ["x", "y", "z", "roll", "pitch", "yaw", "pad", "gripper"] state_names = ["x", "y", "z", "roll", "pitch", "yaw", "pad", "gripper"]
if "libero" in dataset_name: if "libero" in dataset_name:
state_names = ["x", "y", "z", "roll", "pitch", "yaw", "gripper", "gripper"] # 2D gripper state state_names = ["x", "y", "z", "roll", "pitch", "yaw", "gripper", "gripper"] # 2D gripper state
elif state_encoding == StateEncoding.POS_QUAT: elif state_encoding == StateEncoding.POS_QUAT:
state_names = ["x", "y", "z", "rx", "ry", "rz", "rw", "gripper"] state_names = ["x", "y", "z", "rx", "ry", "rz", "rw", "gripper"]
elif state_encoding == StateEncoding.JOINT: elif state_encoding == StateEncoding.JOINT:
@@ -107,16 +107,15 @@ def generate_features_from_raw(builder: tfds.core.DatasetBuilder, use_videos: bo
elif action_encoding == ActionEncoding.JOINT_POS: elif action_encoding == ActionEncoding.JOINT_POS:
action_names = [f"motor_{i}" for i in range(7)] + ["gripper"] action_names = [f"motor_{i}" for i in range(7)] + ["gripper"]
DEFAULT_FEATURES = { DEFAULT_FEATURES = {
"observation.state": { "observation.state": {
"dtype": "float32", "dtype": "float32",
"shape": (8,), "shape": (len(state_names),),
"names": {"motors": state_names}, "names": {"motors": state_names},
}, },
"action": { "action": {
"dtype": "float32", "dtype": "float32",
"shape": (7,), "shape": (len(action_names),),
"names": {"motors": action_names}, "names": {"motors": action_names},
}, },
} }