From 940120d7ed8e61c74182c35618f231b6a188d8d4 Mon Sep 17 00:00:00 2001
From: Qizhi Chen <tavish9.chen@gmail.com>
Date: Wed, 29 Oct 2025 16:51:15 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=93=9D=20update=20libero=20state=20&=20ac?=
 =?UTF-8?q?tion=20notation=20(#71)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 libero2lerobot/libero_utils/config.py |  6 +++---
 openx2lerobot/openx_rlds.py           | 13 ++++++++++++-
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/libero2lerobot/libero_utils/config.py b/libero2lerobot/libero_utils/config.py
index 35250fa..4df5acb 100644
--- a/libero2lerobot/libero_utils/config.py
+++ b/libero2lerobot/libero_utils/config.py
@@ -12,12 +12,12 @@ LIBERO_FEATURES = {
     "observation.state": {
         "dtype": "float32",
         "shape": (8,),
-        "names": {"motors": ["x", "y", "z", "roll", "pitch", "yaw", "gripper", "gripper"]},
+        "names": {"motors": ["x", "y", "z", "axis_angle1", "axis_angle2", "axis_angle3", "gripper", "gripper"]},
     },
     "observation.states.ee_state": {
         "dtype": "float32",
         "shape": (6,),
-        "names": {"motors": ["x", "y", "z", "roll", "pitch", "yaw"]},
+        "names": {"motors": ["x", "y", "z", "axis_angle1", "axis_angle2", "axis_angle3"]},
     },
     "observation.states.joint_state": {
         "dtype": "float32",
@@ -32,6 +32,6 @@ LIBERO_FEATURES = {
     "action": {
         "dtype": "float32",
         "shape": (7,),
-        "names": {"motors": ["x", "y", "z", "roll", "pitch", "yaw", "gripper"]},
+        "names": {"motors": ["x", "y", "z", "axis_angle1", "axis_angle2", "axis_angle3", "gripper"]},
     },
 }
diff --git a/openx2lerobot/openx_rlds.py b/openx2lerobot/openx_rlds.py
index b8809a1..35fff49 100644
--- a/openx2lerobot/openx_rlds.py
+++ b/openx2lerobot/openx_rlds.py
@@ -89,7 +89,16 @@ def generate_features_from_raw(builder: tfds.core.DatasetBuilder, use_videos: bo
         if state_encoding == StateEncoding.POS_EULER:
             state_names = ["x", "y", "z", "roll", "pitch", "yaw", "pad", "gripper"]
             if "libero" in dataset_name:
-                state_names = ["x", "y", "z", "roll", "pitch", "yaw", "gripper", "gripper"]  # 2D gripper state
+                state_names = [
+                    "x",
+                    "y",
+                    "z",
+                    "axis_angle1",
+                    "axis_angle2",
+                    "axis_angle3",
+                    "gripper",
+                    "gripper",
+                ]  # 2D gripper state
         elif state_encoding == StateEncoding.POS_QUAT:
             state_names = ["x", "y", "z", "rx", "ry", "rz", "rw", "gripper"]
         elif state_encoding == StateEncoding.JOINT:
@@ -104,6 +113,8 @@ def generate_features_from_raw(builder: tfds.core.DatasetBuilder, use_videos: bo
         action_encoding = OXE_DATASET_CONFIGS[dataset_name]["action_encoding"]
         if action_encoding == ActionEncoding.EEF_POS:
             action_names = ["x", "y", "z", "roll", "pitch", "yaw", "gripper"]
+            if "libero" in dataset_name:
+                action_names = ["x", "y", "z", "axis_angle1", "axis_angle2", "axis_angle3", "gripper"]
         elif action_encoding == ActionEncoding.JOINT_POS:
             action_names = [f"motor_{i}" for i in range(7)] + ["gripper"]