From 0db1a67eaf416ed3775752dbe0d83ce2e4e3f5d2 Mon Sep 17 00:00:00 2001
From: Steven Palma <imstevenpmwork@ieee.org>
Date: Wed, 28 May 2025 15:08:10 +0200
Subject: [PATCH] fix(dataset): key is an action if it starts with such prefix
 in dataset_to_policy_features (#1156)

---
 lerobot/common/datasets/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lerobot/common/datasets/utils.py b/lerobot/common/datasets/utils.py
index 581b3c1d..a16fe8ac 100644
--- a/lerobot/common/datasets/utils.py
+++ b/lerobot/common/datasets/utils.py
@@ -461,7 +461,7 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea
             type = FeatureType.ENV
         elif key.startswith("observation"):
             type = FeatureType.STATE
-        elif key == "action":
+        elif key.startswith("action"):
             type = FeatureType.ACTION
         else:
             continue