Generalize LeRobot dataset converters to N-dof robots

docs/docs/docs/getting_started/policy_support.md

@@ -53,6 +53,10 @@ pip install numpy==1.26.0
 The available arguments of `isaaclab2lerobotv3.py` are identical to those of `isaaclab2lerobot.py`.
 :::
 
+:::tip
+The conversion scripts are not limited to the SO-101. The joint names, joint count, and value ranges are all taken from the task's environment configuration (`robot_name`, `default_feature_joint_names`, `usd_joint_limits`, `motor_limits`), so a task built around a robot with a different number of joints only needs to override these fields in its environment configuration to export LeRobot datasets.
+:::
+
 ## 2. Policy Training
 
 Taking [GR00T N1.5](https://github.com/NVIDIA/Isaac-GR00T) as an example, which provides a fine-tuning workflow based on the LeRobot Dataset. You can refer to [nvidia/gr00t-n1.5-so101-tuning](https://huggingface.co/blog/nvidia/gr00t-n1-5-so101-tuning) to fine-tune it with your collected lerobot data. We take pick-orange task as an example.

scripts/convert/isaaclab2lerobot.py

@@ -186,6 +186,12 @@ def convert_isaaclab_to_lerobot():
 
         dataset_file_handler.close()
 
+    if now_episode_index == 0:
+        raise RuntimeError(
+            "No episodes were converted: every episode was skipped because it was unsuccessful or"
+            " shorter than 10 frames. Refusing to write an empty LeRobot dataset."
+        )
+
     if args_cli.push_to_hub:
         dataset.push_to_hub()
 

scripts/convert/isaaclab2lerobotv3.py

@@ -187,6 +187,12 @@ def convert_isaaclab_to_lerobot():
 
         dataset_file_handler.close()
 
+    if now_episode_index == 0:
+        raise RuntimeError(
+            "No episodes were converted: every episode was skipped because it was unsuccessful or"
+            " shorter than 10 frames. Refusing to write an empty LeRobot dataset."
+        )
+
     dataset.finalize()
 
     if args_cli.push_to_hub:

source/leisaac/leisaac/tasks/template/bi_arm_env_cfg.py

@@ -18,7 +18,11 @@
 from isaaclab.sensors import TiledCameraCfg
 from isaaclab.utils import configclass
 from isaaclab.utils.datasets.episode_data import EpisodeData
-from leisaac.assets.robots.lerobot import SO101_FOLLOWER_CFG
+from leisaac.assets.robots.lerobot import (
+    SO101_FOLLOWER_CFG,
+    SO101_FOLLOWER_MOTOR_LIMITS,
+    SO101_FOLLOWER_USD_JOINT_LIMLITS,
+)
 from leisaac.devices.action_process import init_action_cfg, preprocess_device_action
 from leisaac.enhance.datasets.lerobot_dataset_handler import LeRobotDatasetCfg
 from leisaac.utils.constant import BI_ARM_JOINT_NAMES
@@ -188,6 +192,10 @@ class BiArmTaskEnvCfg(ManagerBasedRLEnvCfg):
     """Robot name for lerobot dataset export."""
     default_feature_joint_names: list[str] = MISSING
     """Default feature joint names for lerobot dataset export."""
+    usd_joint_limits: dict[str, tuple[float, float]] = SO101_FOLLOWER_USD_JOINT_LIMLITS
+    """Per-joint USD limits (degree) of each arm for lerobot dataset export. Override for non-SO101 robots."""
+    motor_limits: dict[str, tuple[float, float]] = SO101_FOLLOWER_MOTOR_LIMITS
+    """Per-joint motor limits of each arm for lerobot dataset export. Override for non-SO101 robots."""
     task_description: str = MISSING
     """Task description for lerobot dataset export."""
 
@@ -223,17 +231,26 @@ def build_lerobot_frame(self, episode_data: EpisodeData, dataset_cfg: LeRobotDat
         action = episode_data._data["actions"][-1]
         if dataset_cfg.action_align:
             action = action.unsqueeze(0)
-            left_arm_action = convert_leisaac_action_to_lerobot(action[:, :6]).squeeze(0)
-            right_arm_action = convert_leisaac_action_to_lerobot(action[:, 6:]).squeeze(0)
+            arm_dim = len(self.usd_joint_limits)
+            left_arm_action = convert_leisaac_action_to_lerobot(
+                action[:, :arm_dim], self.usd_joint_limits, self.motor_limits
+            ).squeeze(0)
+            right_arm_action = convert_leisaac_action_to_lerobot(
+                action[:, arm_dim:], self.usd_joint_limits, self.motor_limits
+            ).squeeze(0)
             processed_action = np.concatenate([left_arm_action, right_arm_action], axis=0)
         else:
             processed_action = action.cpu().numpy()
         frame = {
             "action": processed_action,
             "observation.state": np.concatenate(
                 [
-                    convert_leisaac_action_to_lerobot(obs_data["left_joint_pos"][-1].unsqueeze(0)).squeeze(0),
-                    convert_leisaac_action_to_lerobot(obs_data["right_joint_pos"][-1].unsqueeze(0)).squeeze(0),
+                    convert_leisaac_action_to_lerobot(
+                        obs_data["left_joint_pos"][-1].unsqueeze(0), self.usd_joint_limits, self.motor_limits
+                    ).squeeze(0),
+                    convert_leisaac_action_to_lerobot(
+                        obs_data["right_joint_pos"][-1].unsqueeze(0), self.usd_joint_limits, self.motor_limits
+                    ).squeeze(0),
                 ],
                 axis=0,
             ),

source/leisaac/leisaac/tasks/template/direct/bi_arm_env.py

@@ -6,6 +6,10 @@
 from isaaclab.managers import SceneEntityCfg
 from isaaclab.utils import configclass
 from isaaclab.utils.datasets.episode_data import EpisodeData
+from leisaac.assets.robots.lerobot import (
+    SO101_FOLLOWER_MOTOR_LIMITS,
+    SO101_FOLLOWER_USD_JOINT_LIMLITS,
+)
 from leisaac.devices.action_process import preprocess_device_action
 from leisaac.enhance.datasets.lerobot_dataset_handler import LeRobotDatasetCfg
 from leisaac.enhance.envs import RecorderEnhanceDirectRLEnv as DirectRLEnv
@@ -62,6 +66,10 @@ class BiArmTaskDirectEnvCfg(DirectRLEnvCfg):
     """Robot name for lerobot dataset export."""
     default_feature_joint_names: list[str] = MISSING
     """Default feature joint names for lerobot dataset export."""
+    usd_joint_limits: dict[str, tuple[float, float]] = SO101_FOLLOWER_USD_JOINT_LIMLITS
+    """Per-joint USD limits (degree) of each arm for lerobot dataset export. Override for non-SO101 robots."""
+    motor_limits: dict[str, tuple[float, float]] = SO101_FOLLOWER_MOTOR_LIMITS
+    """Per-joint motor limits of each arm for lerobot dataset export. Override for non-SO101 robots."""
     task_description: str = MISSING
     """Task description for lerobot dataset export."""
 
@@ -104,17 +112,26 @@ def build_lerobot_frame(self, episode_data: EpisodeData, dataset_cfg: LeRobotDat
         action = episode_data._data["actions"][-1]
         if dataset_cfg.action_align:
             action = action.unsqueeze(0)
-            left_arm_action = convert_leisaac_action_to_lerobot(action[:, :6]).squeeze(0)
-            right_arm_action = convert_leisaac_action_to_lerobot(action[:, 6:]).squeeze(0)
+            arm_dim = len(self.usd_joint_limits)
+            left_arm_action = convert_leisaac_action_to_lerobot(
+                action[:, :arm_dim], self.usd_joint_limits, self.motor_limits
+            ).squeeze(0)
+            right_arm_action = convert_leisaac_action_to_lerobot(
+                action[:, arm_dim:], self.usd_joint_limits, self.motor_limits
+            ).squeeze(0)
             processed_action = np.concatenate([left_arm_action, right_arm_action], axis=0)
         else:
             processed_action = action.cpu().numpy()
         frame = {
             "action": processed_action,
             "observation.state": np.concatenate(
                 [
-                    convert_leisaac_action_to_lerobot(obs_data["left_joint_pos"][-1].unsqueeze(0)).squeeze(0),
-                    convert_leisaac_action_to_lerobot(obs_data["right_joint_pos"][-1].unsqueeze(0)).squeeze(0),
+                    convert_leisaac_action_to_lerobot(
+                        obs_data["left_joint_pos"][-1].unsqueeze(0), self.usd_joint_limits, self.motor_limits
+                    ).squeeze(0),
+                    convert_leisaac_action_to_lerobot(
+                        obs_data["right_joint_pos"][-1].unsqueeze(0), self.usd_joint_limits, self.motor_limits
+                    ).squeeze(0),
                 ],
                 axis=0,
             ),

source/leisaac/leisaac/tasks/template/direct/single_arm_env.py

@@ -5,6 +5,10 @@
 from isaaclab.managers import SceneEntityCfg
 from isaaclab.utils import configclass
 from isaaclab.utils.datasets.episode_data import EpisodeData
+from leisaac.assets.robots.lerobot import (
+    SO101_FOLLOWER_MOTOR_LIMITS,
+    SO101_FOLLOWER_USD_JOINT_LIMLITS,
+)
 from leisaac.devices.action_process import preprocess_device_action
 from leisaac.enhance.datasets.lerobot_dataset_handler import LeRobotDatasetCfg
 from leisaac.enhance.envs import RecorderEnhanceDirectRLEnv as DirectRLEnv
@@ -55,6 +59,10 @@ class SingleArmTaskDirectEnvCfg(DirectRLEnvCfg):
     """Robot name for lerobot dataset export."""
     default_feature_joint_names: list[str] = MISSING
     """Default feature joint names for lerobot dataset export."""
+    usd_joint_limits: dict[str, tuple[float, float]] = SO101_FOLLOWER_USD_JOINT_LIMLITS
+    """Per-joint limits written in USD (degree) for lerobot dataset export. Override for non-SO101 robots."""
+    motor_limits: dict[str, tuple[float, float]] = SO101_FOLLOWER_MOTOR_LIMITS
+    """Per-joint motor limits of the real device for lerobot dataset export. Override for non-SO101 robots."""
     task_description: str = MISSING
     """Task description for lerobot dataset export."""
 
@@ -93,12 +101,16 @@ def build_lerobot_frame(self, episode_data: EpisodeData, dataset_cfg: LeRobotDat
         obs_data = episode_data._data["obs"]
         action = episode_data._data["actions"][-1]
         if dataset_cfg.action_align:
-            processed_action = convert_leisaac_action_to_lerobot(action.unsqueeze(0)).squeeze(0)
+            processed_action = convert_leisaac_action_to_lerobot(
+                action.unsqueeze(0), self.usd_joint_limits, self.motor_limits
+            ).squeeze(0)
         else:
             processed_action = action.cpu().numpy()
         frame = {
             "action": processed_action,
-            "observation.state": convert_leisaac_action_to_lerobot(obs_data["joint_pos"][-1].unsqueeze(0)).squeeze(0),
+            "observation.state": convert_leisaac_action_to_lerobot(
+                obs_data["joint_pos"][-1].unsqueeze(0), self.usd_joint_limits, self.motor_limits
+            ).squeeze(0),
             "task": self.task_description,
         }
         for frame_key in dataset_cfg.features.keys():

source/leisaac/leisaac/tasks/template/lekiwi_env_cfg.py

@@ -18,7 +18,11 @@
 from isaaclab.sensors import FrameTransformerCfg, OffsetCfg, TiledCameraCfg
 from isaaclab.utils import configclass
 from isaaclab.utils.datasets.episode_data import EpisodeData
-from leisaac.assets.robots.lerobot import LEKIWI_CFG
+from leisaac.assets.robots.lerobot import (
+    LEKIWI_CFG,
+    SO101_FOLLOWER_MOTOR_LIMITS,
+    SO101_FOLLOWER_USD_JOINT_LIMLITS,
+)
 from leisaac.devices.action_process import init_action_cfg, preprocess_device_action
 from leisaac.enhance.datasets.lerobot_dataset_handler import LeRobotDatasetCfg
 from leisaac.utils.constant import LEKIWI_JOINT_NAMES
@@ -188,6 +192,10 @@ class LeKiwiTaskEnvCfg(ManagerBasedRLEnvCfg):
     """Robot name for lerobot dataset export."""
     default_feature_joint_names: list[str] = MISSING
     """Default feature joint names for lerobot dataset export."""
+    usd_joint_limits: dict[str, tuple[float, float]] = SO101_FOLLOWER_USD_JOINT_LIMLITS
+    """Per-joint limits of the arm written in USD (degree) for lerobot dataset export. Override for non-SO101 arms."""
+    motor_limits: dict[str, tuple[float, float]] = SO101_FOLLOWER_MOTOR_LIMITS
+    """Per-joint motor limits of the arm on the real device for lerobot dataset export. Override for non-SO101 arms."""
     task_description: str = MISSING
     """Task description for lerobot dataset export."""
 
@@ -221,7 +229,10 @@ def build_lerobot_frame(self, episode_data: EpisodeData, dataset_cfg: LeRobotDat
         action = episode_data._data["actions"][-1]
         if dataset_cfg.action_align:
             action = action.unsqueeze(0)
-            arm_action = convert_leisaac_action_to_lerobot(action[:, :6]).squeeze(0)
+            arm_dim = len(self.usd_joint_limits)
+            arm_action = convert_leisaac_action_to_lerobot(
+                action[:, :arm_dim], self.usd_joint_limits, self.motor_limits
+            ).squeeze(0)
             wheel_action = obs_data["user_vel_action"][-1].cpu().numpy()
             processed_action = np.concatenate([arm_action, wheel_action], axis=0)
         else:
@@ -230,7 +241,9 @@ def build_lerobot_frame(self, episode_data: EpisodeData, dataset_cfg: LeRobotDat
             "action": processed_action,
             "observation.state": np.concatenate(
                 [
-                    convert_leisaac_action_to_lerobot(obs_data["joint_pos"][-1][:-3].unsqueeze(0)).squeeze(0),
+                    convert_leisaac_action_to_lerobot(
+                        obs_data["joint_pos"][-1][:-3].unsqueeze(0), self.usd_joint_limits, self.motor_limits
+                    ).squeeze(0),
                     obs_data["user_vel_state"][-1].cpu().numpy(),
                 ],
                 axis=0,

source/leisaac/leisaac/tasks/template/single_arm_env_cfg.py

@@ -17,7 +17,11 @@
 from isaaclab.sensors import FrameTransformerCfg, OffsetCfg, TiledCameraCfg
 from isaaclab.utils import configclass
 from isaaclab.utils.datasets.episode_data import EpisodeData
-from leisaac.assets.robots.lerobot import SO101_FOLLOWER_CFG
+from leisaac.assets.robots.lerobot import (
+    SO101_FOLLOWER_CFG,
+    SO101_FOLLOWER_MOTOR_LIMITS,
+    SO101_FOLLOWER_USD_JOINT_LIMLITS,
+)
 from leisaac.devices.action_process import init_action_cfg, preprocess_device_action
 from leisaac.enhance.datasets.lerobot_dataset_handler import LeRobotDatasetCfg
 from leisaac.utils.constant import SINGLE_ARM_JOINT_NAMES
@@ -172,6 +176,10 @@ class SingleArmTaskEnvCfg(ManagerBasedRLEnvCfg):
     """Robot name for lerobot dataset export."""
     default_feature_joint_names: list[str] = MISSING
     """Default feature joint names for lerobot dataset export."""
+    usd_joint_limits: dict[str, tuple[float, float]] = SO101_FOLLOWER_USD_JOINT_LIMLITS
+    """Per-joint limits written in USD (degree) for lerobot dataset export. Override for non-SO101 robots."""
+    motor_limits: dict[str, tuple[float, float]] = SO101_FOLLOWER_MOTOR_LIMITS
+    """Per-joint motor limits of the real device for lerobot dataset export. Override for non-SO101 robots."""
     task_description: str = MISSING
     """Task description for lerobot dataset export."""
 
@@ -204,12 +212,16 @@ def build_lerobot_frame(self, episode_data: EpisodeData, dataset_cfg: LeRobotDat
         obs_data = episode_data._data["obs"]
         action = episode_data._data["actions"][-1]
         if dataset_cfg.action_align:
-            processed_action = convert_leisaac_action_to_lerobot(action.unsqueeze(0)).squeeze(0)
+            processed_action = convert_leisaac_action_to_lerobot(
+                action.unsqueeze(0), self.usd_joint_limits, self.motor_limits
+            ).squeeze(0)
         else:
             processed_action = action.cpu().numpy()
         frame = {
             "action": processed_action,
-            "observation.state": convert_leisaac_action_to_lerobot(obs_data["joint_pos"][-1].unsqueeze(0)).squeeze(0),
+            "observation.state": convert_leisaac_action_to_lerobot(
+                obs_data["joint_pos"][-1].unsqueeze(0), self.usd_joint_limits, self.motor_limits
+            ).squeeze(0),
             "task": self.task_description,
         }
         for frame_key in dataset_cfg.features.keys():

source/leisaac/leisaac/utils/robot_utils.py

@@ -14,11 +14,9 @@
 
 @dataclass
 class StateFeatureItem:
+    shape: tuple
+    names: list[str]
     dtype: str = "float32"
-    shape: tuple = (6,)
-    names: list[str] = field(
-        default_factory=lambda: ["joint1.pos", "joint2.pos", "joint3.pos", "joint4.pos", "joint5.pos", "joint6.pos"]
-    )
 
 
 @dataclass
@@ -47,6 +45,14 @@ def build_feature_from_env(env: ManagerBasedEnv | DirectRLEnv, dataset_cfg: LeRo
     features = {}
 
     default_feature_joint_names = env.cfg.default_feature_joint_names
+    if not default_feature_joint_names:
+        raise ValueError(
+            "env.cfg.default_feature_joint_names is empty; it must contain one feature name per robot joint."
+        )
+    if len(set(default_feature_joint_names)) != len(default_feature_joint_names):
+        raise ValueError(
+            f"env.cfg.default_feature_joint_names contains duplicated names: {default_feature_joint_names}"
+        )
     if isinstance(env, ManagerBasedEnv):
         action_dim = env.action_manager.total_action_dim
     else:
@@ -93,16 +99,51 @@ def is_so101_at_rest_pose(joint_pos: torch.Tensor, joint_names: list[str]) -> to
     return is_reset
 
 
-def convert_leisaac_action_to_lerobot(action: torch.Tensor | np.ndarray) -> np.ndarray:
+def resolve_conversion_limits(
+    dim: int,
+    joint_limits: dict[str, tuple[float, float]] | None,
+    motor_limits: dict[str, tuple[float, float]] | None,
+) -> tuple[dict[str, tuple[float, float]], dict[str, tuple[float, float]]]:
+    """
+    Resolve the joint/motor limits used for value conversion, defaulting to the SO-101 follower ones.
+
+    Both dicts must define the joints in order and cover every converted dimension. Raising here
+    prevents dimensions without a configured limit from being silently dropped.
+    """
+    if joint_limits is None:
+        joint_limits = SO101_FOLLOWER_USD_JOINT_LIMLITS
+    if motor_limits is None:
+        motor_limits = SO101_FOLLOWER_MOTOR_LIMITS
+    if len(joint_limits) == 0:
+        raise ValueError("joint_limits must contain at least one joint.")
+    missing_motor_limits = [joint_name for joint_name in joint_limits if joint_name not in motor_limits]
+    if missing_motor_limits:
+        raise ValueError(f"motor_limits is missing entries for joints: {missing_motor_limits}")
+    if dim != len(joint_limits):
+        raise ValueError(
+            f"Value has {dim} dimensions but limits are configured for {len(joint_limits)} joints"
+            f" ({list(joint_limits)}). Configure one joint limit per dimension in the task's"
+            " environment configuration (usd_joint_limits / motor_limits)."
+        )
+    return joint_limits, motor_limits
+
+
+def convert_leisaac_action_to_lerobot(
+    action: torch.Tensor | np.ndarray,
+    joint_limits: dict[str, tuple[float, float]] | None = None,
+    motor_limits: dict[str, tuple[float, float]] | None = None,
+) -> np.ndarray:
     """
     Convert the action from LeIsaac to Lerobot. Just convert value, not include the format.
+
+    ``joint_limits`` and ``motor_limits`` default to the SO-101 follower ones; pass the limits of
+    your own robot to convert values for other robots.
     """
     if isinstance(action, torch.Tensor):
         action = action.cpu().numpy()
 
+    joint_limits, motor_limits = resolve_conversion_limits(action.shape[-1], joint_limits, motor_limits)
     processed_action = np.zeros_like(action)
-    joint_limits = SO101_FOLLOWER_USD_JOINT_LIMLITS
-    motor_limits = SO101_FOLLOWER_MOTOR_LIMITS
     action = action / torch.pi * 180.0  # convert to degree
 
     for idx, joint_name in enumerate(joint_limits):
@@ -116,16 +157,22 @@ def convert_leisaac_action_to_lerobot(action: torch.Tensor | np.ndarray) -> np.n
     return processed_action
 
 
-def convert_lerobot_action_to_leisaac(action: torch.Tensor | np.ndarray) -> np.ndarray:
+def convert_lerobot_action_to_leisaac(
+    action: torch.Tensor | np.ndarray,
+    joint_limits: dict[str, tuple[float, float]] | None = None,
+    motor_limits: dict[str, tuple[float, float]] | None = None,
+) -> np.ndarray:
     """
     Convert the action from Lerobot to LeIsaac. Just convert value, not include the format.
+
+    ``joint_limits`` and ``motor_limits`` default to the SO-101 follower ones; pass the limits of
+    your own robot to convert values for other robots.
     """
     if isinstance(action, torch.Tensor):
         action = action.cpu().numpy()
 
+    joint_limits, motor_limits = resolve_conversion_limits(action.shape[-1], joint_limits, motor_limits)
     processed_action = np.zeros_like(action)
-    joint_limits = SO101_FOLLOWER_USD_JOINT_LIMLITS
-    motor_limits = SO101_FOLLOWER_MOTOR_LIMITS
 
     for idx, joint_name in enumerate(joint_limits):
         motor_limit_range = motor_limits[joint_name]