Custom MimicEnvCfg environment for Franka Lift Cube Task

[adityadutt09]

I am trying to work through creating a Custom Mimic environment for the Lift task using the Franka, so that I can collect demonstrations -> annotate them -> then generate a large dataset.

For this I am attempting the following:

• I tried to use the lift_cube_sm.py script by modifying it to add recording functionality and store a small dataset of 10 demos as a replacement to collecting teleop demos.
• Then, I created a FrankaCubeLiftIKAbsMimicEnvCfg and FrankaCubeLiftIKAbsMimicEnv class to allow me to use annotate_demos.py and subsequently generate_dataset.py.

Problem:

• I am not able to run generate_dataset.py on my annotated demo file after using annotate_demos.py with --auto flag on.

I am able to get the example dataset provided here for the Franka Cube Stack working with the annotate_demos.py and generate_dataset.py scripts, so I suspect the problem is in how I am defining my EnvCfg or how I am storing the recorded/annotated datasets.

Modified Code

lift_cube_sm.py (Added recording functionality + disabled observation corruption and resampling of the object goal pose)

"""
Script to run an environment with a pick and lift state machine.

The state machine is implemented in the kernel function `infer_state_machine`.
It uses the `warp` library to run the state machine in parallel on the GPU.

.. code-block:: bash

    ./isaaclab.sh -p scripts/environments/state_machine/lift_cube_sm.py --num_envs 32

"""

"""Launch Omniverse Toolkit first."""

import argparse

from isaaclab.app import AppLauncher

# add argparse arguments
parser = argparse.ArgumentParser(description="Pick and lift state machine for lift environments.")
parser.add_argument(
    "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
)
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
parser.add_argument("--record", action="store_true", default=False, help="Enable recording of demonstrations.")
parser.add_argument(
    "--dataset_file", type=str, default="./datasets/lift_sm_dataset.hdf5", help="File path to export recorded demos."
)
parser.add_argument(
    "--num_demos", type=int, default=0, help="Number of demonstrations to record per environment. Set to 0 for infinite."
)
parser.add_argument(
    "--num_success_steps",
    type=int,
    default=10,
    help="Number of continuous steps with task success for concluding a demo as successful. Default is 10.",
)
# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
# parse the arguments
args_cli = parser.parse_args()

# launch omniverse app
app_launcher = AppLauncher(headless=args_cli.headless)
simulation_app = app_launcher.app

"""Rest everything else."""

import gymnasium as gym
import os
import torch
from collections.abc import Sequence

import warp as wp

from isaaclab.assets.rigid_object.rigid_object_data import RigidObjectData
from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg
from isaaclab.managers import DatasetExportMode, SceneEntityCfg, TerminationTermCfg

import isaaclab_tasks  # noqa: F401
from isaaclab_tasks.manager_based.manipulation.lift.lift_env_cfg import LiftEnvCfg
from isaaclab_tasks.manager_based.manipulation.lift.mdp.terminations import object_reached_goal
from isaaclab_tasks.utils.parse_cfg import parse_env_cfg

# initialize warp
wp.init()


class GripperState:
    """States for the gripper."""

    OPEN = wp.constant(1.0)
    CLOSE = wp.constant(-1.0)


class PickSmState:
    """States for the pick state machine."""

    REST = wp.constant(0)
    APPROACH_ABOVE_OBJECT = wp.constant(1)
    APPROACH_OBJECT = wp.constant(2)
    GRASP_OBJECT = wp.constant(3)
    LIFT_OBJECT = wp.constant(4)


class PickSmWaitTime:
    """Additional wait times (in s) for states for before switching."""

    REST = wp.constant(0.2)
    APPROACH_ABOVE_OBJECT = wp.constant(0.5)
    APPROACH_OBJECT = wp.constant(0.6)
    GRASP_OBJECT = wp.constant(0.3)
    LIFT_OBJECT = wp.constant(1.0)


@wp.func
def distance_below_threshold(current_pos: wp.vec3, desired_pos: wp.vec3, threshold: float) -> bool:
    return wp.length(current_pos - desired_pos) < threshold


@wp.kernel
def infer_state_machine(
    dt: wp.array(dtype=float),
    sm_state: wp.array(dtype=int),
    sm_wait_time: wp.array(dtype=float),
    ee_pose: wp.array(dtype=wp.transform),
    object_pose: wp.array(dtype=wp.transform),
    des_object_pose: wp.array(dtype=wp.transform),
    des_ee_pose: wp.array(dtype=wp.transform),
    gripper_state: wp.array(dtype=float),
    offset: wp.array(dtype=wp.transform),
    position_threshold: float,
):
    # retrieve thread id
    tid = wp.tid()
    # retrieve state machine state
    state = sm_state[tid]
    # decide next state
    if state == PickSmState.REST:
        des_ee_pose[tid] = ee_pose[tid]
        gripper_state[tid] = GripperState.OPEN
        # wait for a while
        if sm_wait_time[tid] >= PickSmWaitTime.REST:
            # move to next state and reset wait time
            sm_state[tid] = PickSmState.APPROACH_ABOVE_OBJECT
            sm_wait_time[tid] = 0.0
    elif state == PickSmState.APPROACH_ABOVE_OBJECT:
        des_ee_pose[tid] = wp.transform_multiply(offset[tid], object_pose[tid])
        gripper_state[tid] = GripperState.OPEN
        if distance_below_threshold(
            wp.transform_get_translation(ee_pose[tid]),
            wp.transform_get_translation(des_ee_pose[tid]),
            position_threshold,
        ):
            # wait for a while
            if sm_wait_time[tid] >= PickSmWaitTime.APPROACH_OBJECT:
                # move to next state and reset wait time
                sm_state[tid] = PickSmState.APPROACH_OBJECT
                sm_wait_time[tid] = 0.0
    elif state == PickSmState.APPROACH_OBJECT:
        des_ee_pose[tid] = object_pose[tid]
        gripper_state[tid] = GripperState.OPEN
        if distance_below_threshold(
            wp.transform_get_translation(ee_pose[tid]),
            wp.transform_get_translation(des_ee_pose[tid]),
            position_threshold,
        ):
            if sm_wait_time[tid] >= PickSmWaitTime.APPROACH_OBJECT:
                # move to next state and reset wait time
                sm_state[tid] = PickSmState.GRASP_OBJECT
                sm_wait_time[tid] = 0.0
    elif state == PickSmState.GRASP_OBJECT:
        des_ee_pose[tid] = object_pose[tid]
        gripper_state[tid] = GripperState.CLOSE
        # wait for a while
        if sm_wait_time[tid] >= PickSmWaitTime.GRASP_OBJECT:
            # move to next state and reset wait time
            sm_state[tid] = PickSmState.LIFT_OBJECT
            sm_wait_time[tid] = 0.0
    elif state == PickSmState.LIFT_OBJECT:
        des_ee_pose[tid] = des_object_pose[tid]
        gripper_state[tid] = GripperState.CLOSE
        if distance_below_threshold(
            wp.transform_get_translation(ee_pose[tid]),
            wp.transform_get_translation(des_ee_pose[tid]),
            position_threshold,
        ):
            # wait for a while
            if sm_wait_time[tid] >= PickSmWaitTime.LIFT_OBJECT:
                # move to next state and reset wait time
                sm_state[tid] = PickSmState.LIFT_OBJECT
                sm_wait_time[tid] = 0.0
    # increment wait time
    sm_wait_time[tid] = sm_wait_time[tid] + dt[tid]


class PickAndLiftSm:
    """A simple state machine in a robot's task space to pick and lift an object.

    The state machine is implemented as a warp kernel. It takes in the current state of
    the robot's end-effector and the object, and outputs the desired state of the robot's
    end-effector and the gripper. The state machine is implemented as a finite state
    machine with the following states:

    1. REST: The robot is at rest.
    2. APPROACH_ABOVE_OBJECT: The robot moves above the object.
    3. APPROACH_OBJECT: The robot moves to the object.
    4. GRASP_OBJECT: The robot grasps the object.
    5. LIFT_OBJECT: The robot lifts the object to the desired pose. This is the final state.
    """

    def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu", position_threshold=0.01):
        """Initialize the state machine.

        Args:
            dt: The environment time step.
            num_envs: The number of environments to simulate.
            device: The device to run the state machine on.
        """
        # save parameters
        self.dt = float(dt)
        self.num_envs = num_envs
        self.device = device
        self.position_threshold = position_threshold
        # initialize state machine
        self.sm_dt = torch.full((self.num_envs,), self.dt, device=self.device)
        self.sm_state = torch.full((self.num_envs,), 0, dtype=torch.int32, device=self.device)
        self.sm_wait_time = torch.zeros((self.num_envs,), device=self.device)

        # desired state
        self.des_ee_pose = torch.zeros((self.num_envs, 7), device=self.device)
        self.des_gripper_state = torch.full((self.num_envs,), 0.0, device=self.device)

        # approach above object offset
        self.offset = torch.zeros((self.num_envs, 7), device=self.device)
        self.offset[:, 2] = 0.1
        self.offset[:, -1] = 1.0  # warp expects quaternion as (x, y, z, w)

        # convert to warp
        self.sm_dt_wp = wp.from_torch(self.sm_dt, wp.float32)
        self.sm_state_wp = wp.from_torch(self.sm_state, wp.int32)
        self.sm_wait_time_wp = wp.from_torch(self.sm_wait_time, wp.float32)
        self.des_ee_pose_wp = wp.from_torch(self.des_ee_pose, wp.transform)
        self.des_gripper_state_wp = wp.from_torch(self.des_gripper_state, wp.float32)
        self.offset_wp = wp.from_torch(self.offset, wp.transform)

    def reset_idx(self, env_ids: Sequence[int] = None):
        """Reset the state machine."""
        if env_ids is None:
            env_ids = slice(None)
        self.sm_state[env_ids] = 0
        self.sm_wait_time[env_ids] = 0.0

    def compute(self, ee_pose: torch.Tensor, object_pose: torch.Tensor, des_object_pose: torch.Tensor) -> torch.Tensor:
        """Compute the desired state of the robot's end-effector and the gripper."""
        # convert all transformations from (w, x, y, z) to (x, y, z, w)
        ee_pose = ee_pose[:, [0, 1, 2, 4, 5, 6, 3]]
        object_pose = object_pose[:, [0, 1, 2, 4, 5, 6, 3]]
        des_object_pose = des_object_pose[:, [0, 1, 2, 4, 5, 6, 3]]

        # convert to warp
        ee_pose_wp = wp.from_torch(ee_pose.contiguous(), wp.transform)
        object_pose_wp = wp.from_torch(object_pose.contiguous(), wp.transform)
        des_object_pose_wp = wp.from_torch(des_object_pose.contiguous(), wp.transform)

        # run state machine
        wp.launch(
            kernel=infer_state_machine,
            dim=self.num_envs,
            inputs=[
                self.sm_dt_wp,
                self.sm_state_wp,
                self.sm_wait_time_wp,
                ee_pose_wp,
                object_pose_wp,
                des_object_pose_wp,
                self.des_ee_pose_wp,
                self.des_gripper_state_wp,
                self.offset_wp,
                self.position_threshold,
            ],
            device=self.device,
        )

        # convert transformations back to (w, x, y, z)
        des_ee_pose = self.des_ee_pose[:, [0, 1, 2, 6, 3, 4, 5]]
        # convert to torch
        return torch.cat([des_ee_pose, self.des_gripper_state.unsqueeze(-1)], dim=-1)


def main():
    # parse configuration
    env_cfg: LiftEnvCfg = parse_env_cfg(
        "Isaac-Lift-Cube-Franka-IK-Abs-v0",
        device=args_cli.device,
        num_envs=args_cli.num_envs,
        use_fabric=not args_cli.disable_fabric,
    )
    
    # Prevent goal pose from changing during episodes
    # This ensures each demonstration has a consistent goal throughout
    env_cfg.commands.object_pose.resampling_time_range = (1.0e9, 1.0e9)
    print(f"Goal resampling disabled: resampling_time_range = {env_cfg.commands.object_pose.resampling_time_range}")
    
    if args_cli.record:
        env_cfg.observations.policy.enable_corruption = False
        env_cfg.observations.policy.concatenate_terms = False
        print("Observation corruption disabled and concatenation disabled for Mimic annotation")
    
    # Configure recording if enabled
    success_term = None
    if args_cli.record:
        # Set up output directories
        output_dir = os.path.dirname(args_cli.dataset_file)
        output_file_name = os.path.splitext(os.path.basename(args_cli.dataset_file))[0]
        
        # Create directory if it doesn't exist
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            print(f"Created output directory: {output_dir}")
        
        # Add success termination term
        success_term = TerminationTermCfg(
            func=object_reached_goal,
            params={
                "command_name": "object_pose",
                "threshold": 0.02,
                "robot_cfg": SceneEntityCfg("robot"),
                "object_cfg": SceneEntityCfg("object"),
            },
        )
        
        # Remove timeout termination to allow environments to run until success
        env_cfg.terminations.time_out = None
        
        # Configure recorder
        env_cfg.recorders = ActionStateRecorderManagerCfg()
        env_cfg.recorders.dataset_export_dir_path = output_dir
        env_cfg.recorders.dataset_filename = output_file_name
        env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY
        
        print(f"Recording enabled: Will save demonstrations to {args_cli.dataset_file}")
        if args_cli.num_demos > 0:
            print(f"Target: {args_cli.num_demos} demonstrations per environment")
        else:
            print("Recording in infinite mode (Ctrl+C to stop)")
    
    # create environment
    env = gym.make("Isaac-Lift-Cube-Franka-IK-Abs-v0", cfg=env_cfg)
    # reset environment at start
    env.reset()

    # create action buffers (position + quaternion)
    actions = torch.zeros(env.unwrapped.action_space.shape, device=env.unwrapped.device)
    actions[:, 3] = 1.0
    # desired object orientation (keep fixed for grasping)
    desired_orientation = torch.zeros((env.unwrapped.num_envs, 4), device=env.unwrapped.device)
    desired_orientation[:, 1] = 1.0
    # create state machine
    pick_sm = PickAndLiftSm(
        env_cfg.sim.dt * env_cfg.decimation, env.unwrapped.num_envs, env.unwrapped.device, position_threshold=0.01
    )
    
    # Initialize recording tracking variables
    if args_cli.record:
        success_step_count = torch.zeros(env.unwrapped.num_envs, dtype=torch.int32, device=env.unwrapped.device)
        recorded_demo_count = torch.zeros(env.unwrapped.num_envs, dtype=torch.int32, device=env.unwrapped.device)
        total_recorded = 0
        print(f"\n{'='*60}")
        print(f"Recording Mode Active")
        print(f"{'='*60}")

    while simulation_app.is_running():
        # run everything in inference mode
        with torch.inference_mode():
            # step environment
            dones = env.step(actions)[-2]

            # observations
            # -- end-effector frame
            ee_frame_sensor = env.unwrapped.scene["ee_frame"]
            tcp_rest_position = ee_frame_sensor.data.target_pos_w[..., 0, :].clone() - env.unwrapped.scene.env_origins
            tcp_rest_orientation = ee_frame_sensor.data.target_quat_w[..., 0, :].clone()
            # -- object frame
            object_data: RigidObjectData = env.unwrapped.scene["object"].data
            object_position = object_data.root_pos_w - env.unwrapped.scene.env_origins
            # -- target object frame (position + orientation from command)
            object_pose_command = env.unwrapped.command_manager.get_command("object_pose")
            desired_position = object_pose_command[..., :3]

            # advance state machine
            actions = pick_sm.compute(
                torch.cat([tcp_rest_position, tcp_rest_orientation], dim=-1),
                torch.cat([object_position, desired_orientation], dim=-1),  # Use fixed orientation for grasping
                torch.cat([desired_position, desired_orientation], dim=-1),  # Use commanded position with fixed orientation
            )
            
            # Handle recording logic
            if args_cli.record:
                # Check success condition
                success_flags = success_term.func(env.unwrapped, **success_term.params)
                
                # Update success step count
                success_step_count = torch.where(success_flags, success_step_count + 1, torch.zeros_like(success_step_count))
                
                # Check if any environment has reached the required number of success steps
                success_threshold_reached = success_step_count >= args_cli.num_success_steps
                
                if success_threshold_reached.any():
                    successful_envs = success_threshold_reached.nonzero(as_tuple=False).squeeze(-1)
                    
                    # Mark episodes as successful and export
                    for env_idx in successful_envs:
                        env_idx_list = [env_idx.item()]
                        env_idx_tensor = torch.tensor([env_idx.item()], device=env.unwrapped.device)
                        
                        # Record pre-reset data and mark as successful
                        env.unwrapped.recorder_manager.record_pre_reset(env_idx_list, force_export_or_skip=False)
                        
                        # Mark as successful
                        success_tensor = torch.tensor([True], dtype=torch.bool, device=env.unwrapped.device)
                        env.unwrapped.recorder_manager.set_success_to_episodes(env_idx_list, success_tensor)
                        
                        # Export the episode
                        env.unwrapped.recorder_manager.export_episodes(env_idx_list)
                        
                        # Update counters
                        recorded_demo_count[env_idx] += 1
                        total_recorded += 1
                        
                        print(f"✓ Env {env_idx.item()}: Recorded demo #{recorded_demo_count[env_idx].item()} | Total: {total_recorded}")
                        
                        # Reset the environment state
                        env.unwrapped.scene.reset(env_idx_tensor)
                        env.unwrapped.recorder_manager.record_post_reset(env_idx_list)
                        
                        # Reset the state machine for this environment
                        pick_sm.reset_idx(env_idx_tensor)
                    
                    # Reset success step count for successful environments
                    success_step_count[successful_envs] = 0
                    
                    # Mark as done to prevent double reset in the main loop
                    dones[successful_envs] = False
                    
                    # Check if we've reached the target number of demos
                    if args_cli.num_demos > 0:
                        envs_reached_target = recorded_demo_count >= args_cli.num_demos
                        if envs_reached_target.all():
                            print(f"\n{'='*60}")
                            print(f"All environments completed {args_cli.num_demos} demonstrations!")
                            print(f"Total recorded: {total_recorded} demonstrations")
                            print(f"Dataset saved to: {args_cli.dataset_file}")
                            print(f"{'='*60}\n")
                            break

            # reset state machine
            if dones.any():
                pick_sm.reset_idx(dones.nonzero(as_tuple=False).squeeze(-1))

    # close the environment
    env.close()
    
    # Print final recording summary
    if args_cli.record:
        print(f"\n{'='*60}")
        print(f"Recording Session Summary")
        print(f"{'='*60}")
        print(f"Total demonstrations recorded: {total_recorded}")
        print(f"Dataset saved to: {args_cli.dataset_file}")
        print(f"{'='*60}\n")


if __name__ == "__main__":
    # run the main function
    main()
    # close sim app
    simulation_app.close()

franka_lift_ik_abs_mimic_env_cfg.py

# Copyright (c) 2024-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

from isaaclab.envs.mimic_env_cfg import MimicEnvCfg, SubTaskConfig
from isaaclab.managers import ObservationGroupCfg as ObsGroup
from isaaclab.managers import ObservationTermCfg as ObsTerm
from isaaclab.managers import SceneEntityCfg
from isaaclab.managers import TerminationTermCfg as DoneTerm
from isaaclab.utils import configclass

from isaaclab_tasks.manager_based.manipulation.lift.config.franka.ik_abs_env_cfg import FrankaCubeLiftEnvCfg
from isaaclab_tasks.manager_based.manipulation.lift.lift_env_cfg import ObservationsCfg as LiftObservationsCfg
from isaaclab_tasks.manager_based.manipulation.lift.lift_env_cfg import TerminationsCfg as LiftTerminationsCfg
import isaaclab_tasks.manager_based.manipulation.lift.mdp as mdp


@configclass
class FrankaCubeLiftIKAbsMimicEnvCfg(FrankaCubeLiftEnvCfg, MimicEnvCfg):
    """
    Isaac Lab Mimic environment config class for Franka Cube Lift IK Abs env.
    """

    def __post_init__(self):
        # post init of parents
        super().__post_init__()

        # Add end-effector observations to the policy group
        self.observations.policy.eef_pos = ObsTerm(func=mdp.ee_frame_pos)
        self.observations.policy.eef_quat = ObsTerm(func=mdp.ee_frame_quat)
        
        # Disable concatenation so we can access individual observation terms
        self.observations.policy.enable_corruption = False
        self.observations.policy.concatenate_terms = False

        # Add subtask observations
        @configclass
        class SubtaskCfg(ObsGroup):
            """Observations for subtask group - termination signals."""

            # Subtask 1: Object is grasped
            object_grasped = ObsTerm(
                func=mdp.object_grasped,
                params={
                    "robot_cfg": SceneEntityCfg("robot"),
                    "ee_frame_cfg": SceneEntityCfg("ee_frame"),
                    "object_cfg": SceneEntityCfg("object"),
                    "diff_threshold": 0.12,  # Relaxed to 12cm for better annotation success rate
                },
            )

            # Subtask 2: Object is lifted to goal height
            object_lifted_to_goal = ObsTerm(
                func=mdp.object_reached_goal,
                params={
                    "threshold": 0.10,  # 10cm tolerance (relaxed from 5cm)
                    "robot_cfg": SceneEntityCfg("robot"),
                    "object_cfg": SceneEntityCfg("object"),
                    "command_name": "object_pose",
                },
            )

            def __post_init__(self):
                self.enable_corruption = False
                self.concatenate_terms = False

        self.observations.subtask_terms = SubtaskCfg()

        # Add success termination based on reaching the goal
        self.terminations.success = DoneTerm(
            func=mdp.object_reached_goal,
            params={
                "threshold": 0.10,  # 10cm tolerance for success (relaxed from 5cm)
                "robot_cfg": SceneEntityCfg("robot"),
                "object_cfg": SceneEntityCfg("object"),
                "command_name": "object_pose",
            },
        )

        # Set resampling time to a very large value so goal stays fixed for entire episode
        self.commands.object_pose.resampling_time_range = (1.0e9, 1.0e9)

        # Add gripper configuration parameters for object_grasped observation
        self.gripper_joint_names = ["panda_finger_.*"]
        self.gripper_open_val = 0.04
        self.gripper_threshold = 0.003  # Reduced from 0.005 to detect grasp earlier (more lenient)

        # Override the existing datagen config values
        self.datagen_config.name = "demo_src_lift_isaac_lab_task_D0"
        self.datagen_config.generation_guarantee = True
        self.datagen_config.generation_keep_failed = True
        self.datagen_config.generation_num_trials = 10
        self.datagen_config.generation_select_src_per_subtask = True
        self.datagen_config.generation_transform_first_robot_pose = False
        self.datagen_config.generation_interpolate_from_last_target_pose = True
        self.datagen_config.max_num_failures = 25
        self.datagen_config.seed = 1

        # Define subtasks for the lift task
        # For a simple lift task, we have:
        # 1. Approach and grasp the object
        # 2. Lift the object to the goal position
        subtask_configs = []
        
        # Subtask 1: Grasp the object
        subtask_configs.append(
            SubTaskConfig(
                object_ref="object",
                subtask_term_signal="object_grasped",
                subtask_term_offset_range=(10, 20),
                selection_strategy="nearest_neighbor_object",
                selection_strategy_kwargs={"nn_k": 3},
                action_noise=0.01,
                num_interpolation_steps=5,
                num_fixed_steps=0,
                apply_noise_during_interpolation=False,
            )
        )
        
        # Subtask 2: Lift object to goal
        # NOTE: Must have a termination signal for annotation to work!
        # The annotation system needs to know when this subtask completes.
        # IMPORTANT: For the LAST subtask, the offset range MUST end with 0
        # This tells the system the episode ends exactly when this subtask completes
        subtask_configs.append(
            SubTaskConfig(
                object_ref="object",
                subtask_term_signal="object_lifted_to_goal",  # Signal when object reaches goal
                subtask_term_offset_range=(0, 0),  # MUST be (0, 0) for the final subtask
                selection_strategy="nearest_neighbor_object",
                selection_strategy_kwargs={"nn_k": 3},
                action_noise=0.01,
                num_interpolation_steps=5,
                num_fixed_steps=0,
                apply_noise_during_interpolation=False,
            )
        )
        
        # Assign subtask configs for the robot
        # The robot name must match the eef_pose key in the annotated dataset
        # For lift task, the dataset uses "robot" as the key
        self.subtask_configs["robot"] = subtask_configs

And franka_lift_ik_abs_mimic_env.py

# Copyright (c) 2024-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

import torch
from collections.abc import Sequence

import isaaclab.utils.math as PoseUtils
from isaaclab.envs import ManagerBasedRLMimicEnv


class FrankaCubeLiftIKAbsMimicEnv(ManagerBasedRLMimicEnv):
    """
    Isaac Lab Mimic environment wrapper class for Franka Cube Lift IK Abs env.
    """

    def get_robot_eef_pose(self, eef_name: str, env_ids: Sequence[int] | None = None) -> torch.Tensor:
        """Get current robot end effector pose.
        
        Args:
            eef_name: Name of the end-effector (not used for single arm robot).
            env_ids: Environment IDs to get poses for. If None, returns poses for all environments.
            
        Returns:
            torch.Tensor: End-effector poses as 4x4 transformation matrices.
        """
        if env_ids is None:
            env_ids = slice(None)

        # Retrieve end effector pose from the observation buffer
        eef_pos = self.obs_buf["policy"]["eef_pos"][env_ids]
        eef_quat = self.obs_buf["policy"]["eef_quat"][env_ids]
        # Quaternion format is w,x,y,z
        return PoseUtils.make_pose(eef_pos, PoseUtils.matrix_from_quat(eef_quat))

    def target_eef_pose_to_action(
        self,
        target_eef_pose_dict: dict,
        gripper_action_dict: dict,
        action_noise_dict: dict | None = None,
        env_id: int = 0,
    ) -> torch.Tensor:
        """Convert target pose to action.

        This method transforms a dictionary of target end-effector poses and gripper actions
        into a single action tensor that can be used by the environment.

        The function:
        1. Extracts target position and rotation from the pose dictionary
        2. Extracts gripper action for the end effector
        3. Concatenates position and quaternion rotation into a pose action
        4. Optionally adds noise to the pose action for exploration
        5. Combines pose action with gripper action into a final action tensor

        Args:
            target_eef_pose_dict: Dictionary containing target end-effector pose(s),
                with keys as eef names and values as pose tensors (4x4 transformation matrices).
            gripper_action_dict: Dictionary containing gripper action(s),
                with keys as eef names and values as action tensors.
            action_noise_dict: Optional dictionary containing noise magnitudes for each component.
                If provided, random noise is generated and added to the pose action.
            env_id: Environment ID for multi-environment setups, defaults to 0.

        Returns:
            torch.Tensor: A single action tensor combining pose and gripper commands.
                Shape: (1, action_dim) where action_dim = pos(3) + quat(4) + gripper(1) = 8
        """
        # Get target position and rotation for single end-effector
        (target_eef_pose,) = target_eef_pose_dict.values()
        target_pos, target_rot = PoseUtils.unmake_pose(target_eef_pose)

        # Get gripper action for single end-effector
        (gripper_action,) = gripper_action_dict.values()

        # Concatenate position and quaternion rotation into pose action
        pose_action = torch.cat([target_pos, PoseUtils.quat_from_matrix(target_rot)], dim=0)
        
        # Add noise to action if specified
        if action_noise_dict is not None:
            # Extract noise values for position and orientation
            eef_key = list(target_eef_pose_dict.keys())[0]
            if eef_key in action_noise_dict:
                noise_config = action_noise_dict[eef_key]
                # Generate random noise scaled by the specified magnitudes
                noise = torch.randn_like(pose_action)
                # Apply different noise scales to position (first 3) and orientation (last 4)
                noise[:3] *= noise_config.get("position", 0.0)
                noise[3:] *= noise_config.get("orientation", 0.0)
                pose_action += noise

        # Combine pose action and gripper action, then add batch dimension
        return torch.cat([pose_action, gripper_action], dim=0).unsqueeze(0)

    def action_to_target_eef_pose(self, action: torch.Tensor) -> dict[str, torch.Tensor]:
        """Convert action to target end-effector pose.
        
        Args:
            action: Action tensor with shape (batch_size, action_dim).
                action_dim = pos(3) + quat(4) + gripper(1) = 8
                
        Returns:
            dict: Dictionary mapping end-effector name to target poses (4x4 transformation matrices).
        """
        # Get the robot/end-effector name from config
        eef_name = list(self.cfg.subtask_configs.keys())[0]

        # Extract position and quaternion from action
        target_pos = action[:, :3]
        target_quat = action[:, 3:7]
        target_rot = PoseUtils.matrix_from_quat(target_quat)

        # Create pose matrices
        target_poses = PoseUtils.make_pose(target_pos, target_rot).clone()

        return {eef_name: target_poses}

    def actions_to_gripper_actions(self, actions: torch.Tensor) -> dict[str, torch.Tensor]:
        """Extract gripper actions from full action tensor.
        
        Args:
            actions: Action tensor with shape (batch_size, action_dim).
                Last dimension contains gripper action.
                
        Returns:
            dict: Dictionary mapping end-effector name to gripper actions.
        """
        eef_name = list(self.cfg.subtask_configs.keys())[0]
        # Last dimension is gripper action
        return {eef_name: actions[:, -1:]}

    def get_object_poses(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]:
        """Get poses of all objects in the scene.
        
        Args:
            env_ids: Environment IDs to get poses for. If None, returns poses for all environments.
            
        Returns:
            dict: Dictionary mapping object names to their poses (4x4 transformation matrices).
        """
        if env_ids is None:
            env_ids = slice(None)

        # Get object pose from scene
        object_data = self.scene["object"].data
        object_pos = object_data.root_pos_w[env_ids]
        object_quat = object_data.root_quat_w[env_ids]  # Format: (w, x, y, z)
        
        object_pose = PoseUtils.make_pose(object_pos, PoseUtils.matrix_from_quat(object_quat))
        
        return {"object": object_pose}

    def get_subtask_term_signals(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]:
        """Get subtask termination signals.
        
        Args:
            env_ids: Environment IDs to get signals for. If None, returns signals for all environments.
            
        Returns:
            dict: Dictionary mapping subtask signal names to boolean tensors indicating completion.
        """
        if env_ids is None:
            env_ids = slice(None)

        signals = dict()
        subtask_terms = self.obs_buf["subtask_terms"]
        
        # Define termination signals for each subtask
        signals["object_grasped"] = subtask_terms["object_grasped"][env_ids]
        signals["object_lifted_to_goal"] = subtask_terms["object_lifted_to_goal"][env_ids]
        
        return signals

The 2 main errors I see are:

1. During annotate_demos.py, most of the recorded demonstrations are skipped as seen here:

2. After taking the few successfully annotated demos from above (~ 10-15 successfully annotated / 100 recorded demos), when running generate_dataset.py the code hangs after loading the demos into the datagen info pool:

For this one, I was able to debug it down to the action_queue being empty. ``` [DEBUG env_loop] Waiting for actions... (qsize=0) [DEBUG env_loop] Waiting for actions... (qsize=0) [DEBUG env_loop] Waiting for actions... (qsize=0) [DEBUG env_loop] Waiting for actions... (qsize=0) [DEBUG env_loop] Waiting for actions... (qsize=0) ```

The main differences between the example approach and my approach are:

1. Start with recording teleop demos instead of recording pre-planned trajectories using the State Machine
2. Commanding the Franka using IK-Rel instead of IK-Abs
3. Consequence of 2 is different Mimic Env Cfgs.

Would really appreciate feedback or some guidance on how to make this work!