Adds absolute pose franka cube stacking environment for mimic (#267)

This PR adds an absolute pose variation of the franka cube stacking
environment for mimic.

- New feature (non-breaking change which adds functionality)
- This change requires a documentation update

Please attach before and after screenshots of the change if applicable.

<!--
Example:

| Before | After |
| ------ | ----- |
| _gif/png before_ | _gif/png after_ |

To upload images to a PR -- simply drag and drop an image while in edit
mode and it should upload the image directly. You can then paste that
source into the above before/after sections.
-->

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

---------
Co-authored-by: Jiakai Zhang <jaczhang@nvidia.com>
Co-authored-by: Kelly Guo <kellyg@nvidia.com>

source/isaaclab_mimic/config/extension.toml

 [package]
 
 # Semantic Versioning is used: https://semver.org/
-version = "1.0.3"
+version = "1.0.4"
 
 # Description
 category = "isaaclab"

source/isaaclab_mimic/docs/CHANGELOG.rst

 Changelog
 ---------
 
-1.0.3 (2025-03-10)
+1.0.4 (2025-03-10)
 ~~~~~~~~~~~~~~~~~~
 
 Changed
@@ -15,6 +15,15 @@ Added
 * Added ``Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-Mimic-v0`` environment for blueprint vision stacking.
 
 
+1.0.3 (2025-03-06)
+~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^^
+
+* Added absolute pose mimic environment for Franka cube stacking task (:class:`FrankaCubeStackIKAbsMimicEnv`)
+
+
 1.0.2 (2025-01-10)
 ~~~~~~~~~~~~~~~~~~
 

source/isaaclab_mimic/isaaclab_mimic/envs/__init__.py

@@ -7,6 +7,8 @@
 
 import gymnasium as gym
 
+from .franka_stack_ik_abs_mimic_env import FrankaCubeStackIKAbsMimicEnv
+from .franka_stack_ik_abs_mimic_env_cfg import FrankaCubeStackIKAbsMimicEnvCfg
 from .franka_stack_ik_rel_blueprint_mimic_env_cfg import FrankaCubeStackIKRelBlueprintMimicEnvCfg
 from .franka_stack_ik_rel_mimic_env import FrankaCubeStackIKRelMimicEnv
 from .franka_stack_ik_rel_mimic_env_cfg import FrankaCubeStackIKRelMimicEnvCfg
@@ -32,3 +34,12 @@ gym.register(
     },
     disable_env_checker=True,
 )
+
+gym.register(
+    id="Isaac-Stack-Cube-Franka-IK-Abs-Mimic-v0",
+    entry_point="isaaclab_mimic.envs:FrankaCubeStackIKAbsMimicEnv",
+    kwargs={
+        "env_cfg_entry_point": franka_stack_ik_abs_mimic_env_cfg.FrankaCubeStackIKAbsMimicEnvCfg,
+    },
+    disable_env_checker=True,
+)

source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_abs_mimic_env.py

+# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
+# 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 FrankaCubeStackIKAbsMimicEnv(ManagerBasedRLMimicEnv):
+    """
+    Isaac Lab Mimic environment wrapper class for Franka Cube Stack 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."""
+        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, noise: float | 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.
+            gripper_action_dict: Dictionary containing gripper action(s),
+                with keys as eef names and values as action tensors.
+            noise: Optional noise magnitude to apply to the pose action for exploration.
+                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.
+        """
+        # target position and rotation
+        (target_eef_pose,) = target_eef_pose_dict.values()
+        target_pos, target_rot = PoseUtils.unmake_pose(target_eef_pose)
+
+        # get gripper action for single eef
+        (gripper_action,) = gripper_action_dict.values()
+
+        # add noise to action
+        pose_action = torch.cat([target_pos, PoseUtils.quat_from_matrix(target_rot)], dim=0)
+        if noise is not None:
+            noise = noise * torch.randn_like(pose_action)
+            pose_action += noise
+
+        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 pose."""
+        eef_name = list(self.cfg.subtask_configs.keys())[0]
+
+        target_pos = action[:, :3]
+        target_quat = action[:, 3:7]
+        target_rot = PoseUtils.matrix_from_quat(target_quat)
+
+        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."""
+        # last dimension is gripper action
+        return {list(self.cfg.subtask_configs.keys())[0]: actions[:, -1:]}
+
+    def get_subtask_term_signals(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]:
+        """Get subtask termination signals."""
+        if env_ids is None:
+            env_ids = slice(None)
+
+        signals = dict()
+        subtask_terms = self.obs_buf["subtask_terms"]
+        signals["grasp_1"] = subtask_terms["grasp_1"][env_ids]
+        signals["grasp_2"] = subtask_terms["grasp_2"][env_ids]
+        signals["stack_1"] = subtask_terms["stack_1"][env_ids]
+        return signals

source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_abs_mimic_env_cfg.py

+# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from isaaclab.envs.mimic_env_cfg import MimicEnvCfg, SubTaskConfig
+from isaaclab.utils import configclass
+
+from isaaclab_tasks.manager_based.manipulation.stack.config.franka.stack_ik_abs_env_cfg import FrankaCubeStackEnvCfg
+
+
+@configclass
+class FrankaCubeStackIKAbsMimicEnvCfg(FrankaCubeStackEnvCfg, MimicEnvCfg):
+    """
+    Isaac Lab Mimic environment config class for Franka Cube Stack IK Abs env.
+    """
+
+    def __post_init__(self):
+        # post init of parents
+        super().__post_init__()
+
+        # Override the existing values
+        self.datagen_config.name = "demo_src_stack_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
+
+        # The following are the subtask configurations for the stack task.
+        subtask_configs = []
+        subtask_configs.append(
+            SubTaskConfig(
+                object_ref="cube_2",
+                subtask_term_signal="grasp_1",
+                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_configs.append(
+            SubTaskConfig(
+                object_ref="cube_1",
+                subtask_term_signal="stack_1",
+                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_configs.append(
+            SubTaskConfig(
+                object_ref="cube_3",
+                subtask_term_signal="grasp_2",
+                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_configs.append(
+            SubTaskConfig(
+                object_ref="cube_2",
+                subtask_term_signal=None,
+                subtask_term_offset_range=(0, 0),
+                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,
+            )
+        )
+        self.subtask_configs["franka"] = subtask_configs

source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/franka/__init__.py

@@ -7,6 +7,7 @@ import os
 
 from . import (
     agents,
+    stack_ik_abs_env_cfg,
     stack_ik_rel_blueprint_env_cfg,
     stack_ik_rel_env_cfg,
     stack_ik_rel_instance_randomize_env_cfg,
@@ -55,6 +56,16 @@ gym.register(
     disable_env_checker=True,
 )
 
+gym.register(
+    id="Isaac-Stack-Cube-Franka-IK-Abs-v0",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    kwargs={
+        "env_cfg_entry_point": stack_ik_abs_env_cfg.FrankaCubeStackEnvCfg,
+        "robomimic_bc_cfg_entry_point": os.path.join(agents.__path__[0], "robomimic/bc_rnn_low_dim.json"),
+    },
+    disable_env_checker=True,
+)
+
 gym.register(
     id="Isaac-Stack-Cube-Instance-Randomize-Franka-IK-Rel-v0",
     entry_point="isaaclab.envs:ManagerBasedRLEnv",

source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_abs_env_cfg.py

+# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from isaaclab.controllers.differential_ik_cfg import DifferentialIKControllerCfg
+from isaaclab.envs.mdp.actions.actions_cfg import DifferentialInverseKinematicsActionCfg
+from isaaclab.utils import configclass
+
+from . import stack_joint_pos_env_cfg
+
+##
+# Pre-defined configs
+##
+from isaaclab_assets.robots.franka import FRANKA_PANDA_HIGH_PD_CFG  # isort: skip
+
+
+@configclass
+class FrankaCubeStackEnvCfg(stack_joint_pos_env_cfg.FrankaCubeStackEnvCfg):
+    def __post_init__(self):
+        # post init of parent
+        super().__post_init__()
+
+        # Set Franka as robot
+        # We switch here to a stiffer PD controller for IK tracking to be better.
+        self.scene.robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
+
+        # Set actions for the specific robot type (franka)
+        self.actions.arm_action = DifferentialInverseKinematicsActionCfg(
+            asset_name="robot",
+            joint_names=["panda_joint.*"],
+            body_name="panda_hand",
+            controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls"),
+        )