lerobot recorder integration in direct env.

source/leisaac/leisaac/enhance/envs/mdp/recorders/recorders.py

@@ -10,6 +10,13 @@ def record_pre_step(self):
         return "actions", self._env.actions
 
 
+class PostStepDirectEnvProcessedActionsRecorder(RecorderTerm):
+    """Direct Env don't have processed actions, use actions to replace processed actions"""
+
+    def record_post_step(self):
+        return "processed_actions", self._env.actions
+
+
 class InitialStateWithParticleObjectsRecorder(RecorderTerm):
     """Recorder term that records the initial state with particle objects."""
 

source/leisaac/leisaac/enhance/envs/mdp/recorders/recorders_cfg.py

@@ -12,6 +12,13 @@ class PreStepDirectEnvActionsRecorderCfg(RecorderTermCfg):
     class_type: type[RecorderTerm] = recorders.PreStepDirectEnvActionsRecorder
 
 
+@configclass
+class PostStepDirectEnvProcessedActionsRecorderCfg(RecorderTermCfg):
+    """Configuration for the post step processed actions recorder term in direct environment."""
+
+    class_type: type[RecorderTerm] = recorders.PostStepDirectEnvProcessedActionsRecorder
+
+
 @configclass
 class InitialStateWithParticleObjectsRecorderCfg(RecorderTermCfg):
     """Configuration for the initial state with particle objects recorder term."""
@@ -24,6 +31,7 @@ class DirectEnvActionStateRecorderManagerCfg(ActionStateRecorderManagerCfg):
     """Recorder configuration for recording actions and states in direct environment."""
 
     record_pre_step_actions = PreStepDirectEnvActionsRecorderCfg()
+    record_post_step_processed_actions = PostStepDirectEnvProcessedActionsRecorderCfg()
 
 
 @configclass

source/leisaac/leisaac/tasks/clean_toy_table/direct/clean_toy_table_bi_arm_env.py

@@ -15,6 +15,8 @@ class CleanToyTableBiArmEnvCfg(BiArmTaskDirectEnvCfg):
 
     scene: CleanToyTableBiArmSceneCfg = CleanToyTableBiArmSceneCfg(env_spacing=8.0)
 
+    task_description: str = "Pick two letter e objects into the box, and reset the arm to rest state."
+
     def __post_init__(self) -> None:
         super().__post_init__()
 

source/leisaac/leisaac/tasks/fold_cloth/direct/fold_cloth_bi_arm_env.py

@@ -25,6 +25,8 @@ class FoldClothBiArmEnvCfg(BiArmTaskDirectEnvCfg):
 
     scene: FoldClothBiArmSceneCfg = FoldClothBiArmSceneCfg(env_spacing=8.0)
 
+    task_description: str = "Fold the cloth, and reset the arm to rest state."
+
     def __post_init__(self) -> None:
         super().__post_init__()
 

source/leisaac/leisaac/tasks/lift_cube/direct/lift_cube_env.py

@@ -20,6 +20,8 @@ class LiftCubeEnvCfg(SingleArmTaskDirectEnvCfg):
 
     scene: LiftCubeSceneCfg = LiftCubeSceneCfg(env_spacing=8.0)
 
+    task_description: str = "Lift the red cube up."
+
     def __post_init__(self) -> None:
         super().__post_init__()
 

source/leisaac/leisaac/tasks/pick_orange/direct/pick_orange_env.py

@@ -20,6 +20,8 @@ class PickOrangeEnvCfg(SingleArmTaskDirectEnvCfg):
 
     scene: PickOrangeSceneCfg = PickOrangeSceneCfg(env_spacing=8.0)
 
+    task_description: str = "Pick three oranges and put them into the plate, then reset the arm to rest state."
+
     def __post_init__(self) -> None:
         super().__post_init__()
 

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

@@ -4,12 +4,14 @@
 import torch
 from isaaclab.managers import SceneEntityCfg
 from isaaclab.utils import configclass
+from isaaclab.utils.datasets.episode_data import EpisodeData
 from leisaac.devices.action_process import preprocess_device_action
 from leisaac.enhance.envs import RecorderEnhanceDirectRLEnv as DirectRLEnv
 from leisaac.enhance.envs import RecorderEnhanceDirectRLEnvCfg as DirectRLEnvCfg
 from leisaac.enhance.envs.mdp.recorders.recorders_cfg import (
     DirectEnvActionStateRecorderManagerCfg as RecordTerm,
 )
+from leisaac.utils.constant import BI_ARM_JOINT_NAMES
 
 from .. import mdp
 from ..bi_arm_env_cfg import BiArmEventCfg, BiArmTaskSceneCfg
@@ -53,7 +55,12 @@ class BiArmTaskDirectEnvCfg(DirectRLEnvCfg):
     dynamic_reset_gripper_effort_limit: bool = True
     """Whether to dynamically reset the gripper effort limit."""
 
-    # TODO: add lerobot dataset export support
+    robot_name: str = "bi_so101_follower"
+    """Robot name for lerobot dataset export."""
+    default_feature_joint_names: list[str] = MISSING
+    """Default feature joint names for lerobot dataset export."""
+    task_description: str = MISSING
+    """Task description for lerobot dataset export."""
 
     def __post_init__(self) -> None:
         super().__post_init__()
@@ -77,13 +84,32 @@ def __post_init__(self) -> None:
                 self.observation_space[cam] = [getattr(self.scene, cam).height, getattr(self.scene, cam).width, 3]
                 self.cameras.append(cam)
 
+        self.default_feature_joint_names = [f"{joint_name}.pos" for joint_name in BI_ARM_JOINT_NAMES]
+
     def use_teleop_device(self, teleop_device) -> None:
         self.task_type = teleop_device
         # self.actions = init_action_cfg(self.actions, device=teleop_device)
 
     def preprocess_device_action(self, action: dict[str, Any], teleop_device) -> torch.Tensor:
         return preprocess_device_action(action, teleop_device)
 
+    def build_lerobot_frame(self, episode_data: EpisodeData, features: dict) -> dict:
+        obs_data = episode_data._data["obs"]
+        frame = {
+            "action": obs_data["actions"][-1].cpu().numpy(),
+            "observation.state": (
+                torch.cat([obs_data["left_joint_pos"][-1], obs_data["right_joint_pos"][-1]]).cpu().numpy()
+            ),
+            "task": self.task_description,
+        }
+        for frame_key in features.keys():
+            if not frame_key.startswith("observation.images"):
+                continue
+            camera_key = frame_key.split(".")[-1]
+            frame[frame_key] = obs_data[camera_key][-1].cpu().numpy()
+
+        return frame
+
 
 class BiArmTaskDirectEnv(DirectRLEnv):
     """Direct RL Environment for bi arm task"""

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

@@ -4,12 +4,14 @@
 import torch
 from isaaclab.managers import SceneEntityCfg
 from isaaclab.utils import configclass
+from isaaclab.utils.datasets.episode_data import EpisodeData
 from leisaac.devices.action_process import preprocess_device_action
 from leisaac.enhance.envs import RecorderEnhanceDirectRLEnv as DirectRLEnv
 from leisaac.enhance.envs import RecorderEnhanceDirectRLEnvCfg as DirectRLEnvCfg
 from leisaac.enhance.envs.mdp.recorders.recorders_cfg import (
     DirectEnvActionStateRecorderManagerCfg as RecordTerm,
 )
+from leisaac.utils.constant import SINGLE_ARM_JOINT_NAMES
 
 from .. import mdp
 from ..single_arm_env_cfg import SingleArmEventCfg, SingleArmTaskSceneCfg
@@ -47,7 +49,12 @@ class SingleArmTaskDirectEnvCfg(DirectRLEnvCfg):
     dynamic_reset_gripper_effort_limit: bool = True
     """Whether to dynamically reset the gripper effort limit."""
 
-    # TODO: add lerobot dataset export support
+    robot_name: str = "so101_follower"
+    """Robot name for lerobot dataset export."""
+    default_feature_joint_names: list[str] = MISSING
+    """Default feature joint names for lerobot dataset export."""
+    task_description: str = MISSING
+    """Task description for lerobot dataset export."""
 
     def __post_init__(self) -> None:
         super().__post_init__()
@@ -70,15 +77,31 @@ def __post_init__(self) -> None:
 
         self.scene.ee_frame.visualizer_cfg.markers["frame"].scale = (0.05, 0.05, 0.05)
 
+        self.default_feature_joint_names = [f"{joint_name}.pos" for joint_name in SINGLE_ARM_JOINT_NAMES]
+
     def use_teleop_device(self, teleop_device) -> None:
         self.task_type = teleop_device
-        # self.actions = init_action_cfg(self.actions, device=teleop_device)
         if teleop_device in ["keyboard", "gamepad"]:
             self.scene.robot.spawn.rigid_props.disable_gravity = True
 
     def preprocess_device_action(self, action: dict[str, Any], teleop_device) -> torch.Tensor:
         return preprocess_device_action(action, teleop_device)
 
+    def build_lerobot_frame(self, episode_data: EpisodeData, features: dict) -> dict:
+        obs_data = episode_data._data["obs"]
+        frame = {
+            "action": obs_data["actions"][-1].cpu().numpy(),
+            "observation.state": obs_data["joint_pos"][-1].cpu().numpy(),
+            "task": self.task_description,
+        }
+        for frame_key in features.keys():
+            if not frame_key.startswith("observation.images"):
+                continue
+            camera_key = frame_key.split(".")[-1]
+            frame[frame_key] = obs_data[camera_key][-1].cpu().numpy()
+
+        return frame
+
 
 class SingleArmTaskDirectEnv(DirectRLEnv):
     """Direct RL Environment for single arm task"""

source/leisaac/leisaac/utils/robot_utils.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import torch
-from isaaclab.envs import ManagerBasedEnv
+from isaaclab.envs import DirectRLEnv, ManagerBasedEnv
 from isaaclab.sensors import Camera
 from leisaac.assets.robots.lerobot import (
     SO101_FOLLOWER_MOTOR_LIMITS,
@@ -40,14 +40,18 @@ class VideoFeatureItem:
     )
 
 
-def build_feature_from_env(env: ManagerBasedEnv, dataset_cfg: LeRobotDatasetCfg) -> dict:
+def build_feature_from_env(env: ManagerBasedEnv | DirectRLEnv, dataset_cfg: LeRobotDatasetCfg) -> dict:
     """
     Build the feature from the environment.
     """
     features = {}
 
     default_feature_joint_names = env.cfg.default_feature_joint_names
-    action_dim = env.action_manager.total_action_dim
+    if isinstance(env, ManagerBasedEnv):
+        action_dim = env.action_manager.total_action_dim
+    else:
+        action_dim = env.actions.shape[-1]
+
     if action_dim != len(default_feature_joint_names):
         # [A bit tricky, currently works because the action dimension matches the joints only when we use leader control]
         action_joint_names = [f"dim_{index}" for index in range(action_dim)]