[SkillGen] Migrates SkillGen to Lab 3.0 (#4724)

# Description

<!--
Thank you for your interest in sending a pull request. Please make sure
to check the contribution guidelines.

Link:
https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html

💡 Please try to keep PRs small and focused. Large PRs are harder to
review and merge.
-->

Fix cuRobo quaternion handling to use correct convention and update Warp
API calls for compatibility.
Fix Rerun visualization code in cuRobo planner.
Add `--visualizer kit` argument to SkillGen documentation for all
non-headless commands.

## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

- Bug fix (non-breaking change which fixes an issue)
- Documentation update

## Checklist

- [x] I have read and understood the [contribution
guidelines](https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html)
- [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
- [ ] 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

---------

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

Author: njawale42

docs/source/overview/imitation-learning/skillgen.rst

@@ -135,7 +135,7 @@ The dataset contains:
 Download and Setup
 ^^^^^^^^^^^^^^^^^^
 
-1. Download the pre-annotated dataset by clicking `here <https://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/5.0/Isaac/IsaacLab/Mimic/franka_stack_datasets/annotated_dataset_skillgen.hdf5>`__.
+1. Download the pre-annotated dataset by clicking `here <https://omniverse-content-staging.s3-us-west-2.amazonaws.com/Assets/Isaac/6.0/Isaac/IsaacLab/Mimic/franka_stack_datasets/annotated_dataset_skillgen.hdf5>`__.
 
 2. Prepare the datasets directory and move the downloaded file:
 
@@ -173,7 +173,8 @@ Download and Setup
       --task Isaac-Stack-Cube-Franka-IK-Rel-Skillgen-v0 \
       --teleop_device spacemouse \
       --dataset_file ./datasets/dataset_skillgen.hdf5 \
-      --num_demos 10
+      --num_demos 10 \
+      --visualizer kit
 
    **Annotate demonstrations for SkillGen** (writes both term and start boundaries):
 
@@ -184,7 +185,8 @@ Download and Setup
       --task Isaac-Stack-Cube-Franka-IK-Rel-Skillgen-v0 \
       --input_file ./datasets/dataset_skillgen.hdf5 \
       --output_file ./datasets/annotated_dataset_skillgen.hdf5 \
-      --annotate_subtask_start_signals
+      --annotate_subtask_start_signals \
+      --visualizer kit
 
 Understanding Dataset Annotation
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -283,7 +285,7 @@ Start with a small dataset to verify everything works:
    --input_file ./datasets/annotated_dataset_skillgen.hdf5 \
    --output_file ./datasets/generated_dataset_small_skillgen_cube_stack.hdf5 \
    --task Isaac-Stack-Cube-Franka-IK-Rel-Skillgen-v0 \
-   --use_skillgen
+   --use_skillgen --visualizer kit
 
 Full-Scale Generation
 ^^^^^^^^^^^^^^^^^^^^^
@@ -337,7 +339,7 @@ Test the adaptive stacking setup:
    --input_file ./datasets/annotated_dataset_skillgen.hdf5 \
    --output_file ./datasets/generated_dataset_small_skillgen_bin_cube_stack.hdf5 \
    --task Isaac-Stack-Cube-Bin-Franka-IK-Rel-Mimic-v0 \
-   --use_skillgen
+   --use_skillgen --visualizer kit
 
 Full-Scale Generation
 ^^^^^^^^^^^^^^^^^^^^^
@@ -419,7 +421,8 @@ Test your trained policies:
    --device cpu \
    --task Isaac-Stack-Cube-Franka-IK-Rel-Skillgen-v0 \
    --num_rollouts 50 \
-   --checkpoint /path/to/model_checkpoint.pth
+   --checkpoint /path/to/model_checkpoint.pth \
+   --visualizer kit
 
 .. code:: bash
 
@@ -428,7 +431,8 @@ Test your trained policies:
    --device cpu \
    --task Isaac-Stack-Cube-Bin-Franka-IK-Rel-Mimic-v0 \
    --num_rollouts 50 \
-   --checkpoint /path/to/model_checkpoint.pth
+   --checkpoint /path/to/model_checkpoint.pth \
+   --visualizer kit
 
 .. note::
 

source/isaaclab_mimic/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Semantic Versioning is used: https://semver.org/
-version = "1.2.0"
+version = "1.2.1"
 
 # Description
 category = "isaaclab"

source/isaaclab_mimic/docs/CHANGELOG.rst

@@ -1,6 +1,17 @@
 Changelog
 ---------
 
+1.2.1 (2026-02-25)
+~~~~~~~~~~~~~~~~~~~
+
+Fixed
+^^^^^
+
+* Fixed cuRobo planner quaternion handling and Warp API compatibility for Isaac Lab 3.0.
+* Fixed Rerun visualization in cuRobo plan visualizer.
+* Added ``--visualizer kit`` to SkillGen documentation for all non-headless commands.
+
+
 1.2.0 (2026-02-23)
 ~~~~~~~~~~~~~~~~~~~
 

source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 import torch
+import warp as wp
 
 from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModelState
 from curobo.geom.sdf.world import CollisionCheckerType
@@ -193,7 +194,7 @@ def __init__(
 
             # Use env-local base translation for multi-env rendering consistency
             env_origin = self.env.scene.env_origins[env_id, :3]
-            base_translation = (self.robot.data.root_pos_w[env_id, :3] - env_origin).detach().cpu().numpy()
+            base_translation = (wp.to_torch(self.robot.data.root_pos_w)[env_id, :3] - env_origin).detach().cpu().numpy()
             self.plan_visualizer = PlanVisualizer(
                 robot_name=self.config.robot_name,
                 recording_id=f"curobo_plan_{env_id}",
@@ -457,10 +458,12 @@ def get_attached_pose(self, link_name: str, joint_state: JointState | None = Non
         link_poses = {}
         if link_state.links_position is not None and link_state.links_quaternion is not None:
             for i, link in enumerate(link_state.link_names):
+                # cuRobo kinematics returns quaternions in (w, x, y, z) format
                 link_poses[link] = self._make_pose(
                     position=link_state.links_position[..., i, :],
                     quaternion=link_state.links_quaternion[..., i, :],
                     name=link,
+                    quat_is_xyzw=False,
                 )
 
         # For attached object link, use ee_link from robot config as parent
@@ -625,22 +628,23 @@ def _sync_object_poses_with_isaaclab(self) -> None:
             # Get current pose from Lab (may be on CPU or CUDA depending on --device flag)
             obj = rigid_objects[object_name]
             env_origin = self.env.scene.env_origins[self.env_id]
-            current_pos_raw = obj.data.root_pos_w[self.env_id] - env_origin
-            current_quat_raw = obj.data.root_quat_w[self.env_id]  # (w, x, y, z)
+            current_pos_raw = wp.to_torch(obj.data.root_pos_w)[self.env_id] - env_origin
+            current_quat_raw = wp.to_torch(obj.data.root_quat_w)[self.env_id]  # (x, y, z, w)
 
             # Convert to cuRobo device and extract float values for pose list
             current_pos = self._to_curobo_device(current_pos_raw)
             current_quat = self._to_curobo_device(current_quat_raw)
 
-            # Convert to cuRobo pose format [x, y, z, w, x, y, z]
+            # Convert to cuRobo pose format [pos_x, pos_y, pos_z, qw, qx, qy, qz]
+            # Isaac Lab quaternion format: (x, y, z, w) -> cuRobo format: (w, x, y, z)
             pose_list = [
                 float(current_pos[0].item()),
                 float(current_pos[1].item()),
                 float(current_pos[2].item()),
-                float(current_quat[0].item()),
-                float(current_quat[1].item()),
-                float(current_quat[2].item()),
-                float(current_quat[3].item()),
+                float(current_quat[3].item()),  # w
+                float(current_quat[0].item()),  # x
+                float(current_quat[1].item()),  # y
+                float(current_quat[2].item()),  # z
             ]
 
             # Update object pose in cuRobo's world model
@@ -665,8 +669,8 @@ def _sync_object_poses_with_isaaclab(self) -> None:
                 # Get current pose and update in collision checker
                 obj = rigid_objects[object_name]
                 env_origin = self.env.scene.env_origins[self.env_id]
-                current_pos_raw = obj.data.root_pos_w[self.env_id] - env_origin
-                current_quat_raw = obj.data.root_quat_w[self.env_id]
+                current_pos_raw = wp.to_torch(obj.data.root_pos_w)[self.env_id] - env_origin
+                current_quat_raw = wp.to_torch(obj.data.root_quat_w)[self.env_id]
 
                 current_pos = self._to_curobo_device(current_pos_raw)
                 current_quat = self._to_curobo_device(current_quat_raw)
@@ -929,7 +933,7 @@ def _get_current_joint_state_for_curobo(self) -> JointState:
             and zero velocity/acceleration.
         """
         # Fetch joint position (shape: [1, num_joints])
-        joint_pos_raw: torch.Tensor = self.robot.data.joint_pos[self.env_id, :].unsqueeze(0)
+        joint_pos_raw: torch.Tensor = wp.to_torch(self.robot.data.joint_pos)[self.env_id, :].unsqueeze(0)
         joint_vel_raw: torch.Tensor = torch.zeros_like(joint_pos_raw)
         joint_acc_raw: torch.Tensor = torch.zeros_like(joint_pos_raw)
 
@@ -997,41 +1001,48 @@ def _make_pose(
         *,
         name: str | None = None,
         normalize_rotation: bool = False,
+        quat_is_xyzw: bool = True,
     ) -> Pose:
         """Create a cuRobo Pose with sensible defaults and device/dtype alignment.
 
         Auto-populates missing fields with identity values and ensures tensors are
-        on the cuRobo device with the correct dtype.
+        on the cuRobo device with the correct dtype. Handles quaternion format conversion
+        from Isaac Lab's (x, y, z, w) to cuRobo's (w, x, y, z) format when needed.
 
         Args:
             position: Optional position as Tensor/ndarray/list. Defaults to [0, 0, 0].
-            quaternion: Optional quaternion as Tensor/ndarray/list (w, x, y, z). Defaults to [1, 0, 0, 0].
+            quaternion: Optional quaternion as Tensor/ndarray/list. Defaults to identity quaternion.
             name: Optional name of the link that this pose represents.
             normalize_rotation: Whether to normalize the quaternion inside Pose.
+            quat_is_xyzw: If True, quaternion is in Isaac Lab format (x, y, z, w) and will be
+                converted to cuRobo format. If False, quaternion is already in cuRobo (w, x, y, z) format.
 
         Returns:
             Pose: A cuRobo Pose on the configured cuRobo device and dtype.
         """
-        # Defaults
         if position is None:
             position = torch.tensor([0.0, 0.0, 0.0], dtype=self.tensor_args.dtype, device=self.tensor_args.device)
         if quaternion is None:
-            quaternion = torch.tensor(
-                [0.0, 0.0, 0.0, 1.0], dtype=self.tensor_args.dtype, device=self.tensor_args.device
+            quaternion_wxyz = torch.tensor(
+                [1.0, 0.0, 0.0, 0.0], dtype=self.tensor_args.dtype, device=self.tensor_args.device
             )
+        else:
+            if not isinstance(quaternion, torch.Tensor):
+                quaternion = torch.tensor(quaternion, dtype=self.tensor_args.dtype, device=self.tensor_args.device)
+            else:
+                quaternion = self._to_curobo_device(quaternion)
+
+            if quat_is_xyzw:
+                quaternion_wxyz = torch.roll(quaternion, shifts=1, dims=-1)
+            else:
+                quaternion_wxyz = quaternion
 
-        # Convert to tensors if needed
         if not isinstance(position, torch.Tensor):
             position = torch.tensor(position, dtype=self.tensor_args.dtype, device=self.tensor_args.device)
         else:
             position = self._to_curobo_device(position)
 
-        if not isinstance(quaternion, torch.Tensor):
-            quaternion = torch.tensor(quaternion, dtype=self.tensor_args.dtype, device=self.tensor_args.device)
-        else:
-            quaternion = self._to_curobo_device(quaternion)
-
-        return Pose(position=position, quaternion=quaternion, name=name, normalize_rotation=normalize_rotation)
+        return Pose(position=position, quaternion=quaternion_wxyz, name=name, normalize_rotation=normalize_rotation)
 
     def _set_active_links(self, links: list[str], active: bool) -> None:
         """Configure collision checking for specific robot links.
@@ -1563,7 +1574,7 @@ def _update_visualization_at_joint_positions(self, joint_positions: torch.Tensor
         if self.frame_counter % self.sphere_update_freq != 0:
             return
 
-        original_joints: torch.Tensor = self.robot.data.joint_pos[self.env_id].clone()
+        original_joints: torch.Tensor = wp.to_torch(self.robot.data.joint_pos)[self.env_id].clone()
 
         try:
             # Ensure joint positions are on environment device for robot commands
@@ -1690,7 +1701,7 @@ def _create_sphere_config(self, sphere_idx: int, sphere, robot_link_count: int)
         opacity = 0.9 if is_attached else 0.5
 
         # Calculate position in world frame (do not use env_origin)
-        root_translation = (self.robot.data.root_pos_w[self.env_id, :3]).detach().cpu().numpy()
+        root_translation = wp.to_torch(self.robot.data.root_pos_w)[self.env_id, :3].detach().cpu().numpy()
         position = sphere.position.cpu().numpy() if hasattr(sphere.position, "cpu") else sphere.position
         if not is_attached:
             position = position + root_translation
@@ -1777,7 +1788,7 @@ def update_world_and_plan_motion(
         gripper_closed = expected_attached_object is not None
         self._set_gripper_state(gripper_closed)
         current_attached = self.get_attached_objects()
-        gripper_pos = self.robot.data.joint_pos[env_id, -2:]
+        gripper_pos = wp.to_torch(self.robot.data.joint_pos)[env_id, -2:]
 
         self.logger.debug(f"Current attached objects: {current_attached}")
 
@@ -1799,13 +1810,13 @@ def update_world_and_plan_motion(
                 if expected_attached_object in rigid_objects:
                     obj = rigid_objects[expected_attached_object]
                     origin = self.env.scene.env_origins[env_id]
-                    obj_pos = obj.data.root_pos_w[env_id] - origin
+                    obj_pos = wp.to_torch(obj.data.root_pos_w)[env_id] - origin
                     self.logger.debug(f"Isaac Lab object position: {obj_pos}")
 
                     # Debug end-effector position
                     ee_frame_cfg = SceneEntityCfg("ee_frame")
                     ee_frame = self.env.scene[ee_frame_cfg.name]
-                    ee_pos = ee_frame.data.target_pos_w[env_id, 0, :] - origin
+                    ee_pos = wp.to_torch(ee_frame.data.target_pos_w)[env_id, 0, :] - origin
                     self.logger.debug(f"End-effector position: {ee_pos}")
 
                     # Debug distance

source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/plan_visualizer.py

@@ -27,6 +27,8 @@
 except ImportError:
     raise ImportError("Rerun is not installed!")
 
+_RR_HAS_TRANSFORM_AXES = hasattr(rr, "TransformAxes3D")
+
 from curobo.types.state import JointState
 
 import isaaclab.utils.math as PoseUtils
@@ -628,15 +630,17 @@ def _to_rerun_mesh(mesh: "trimesh.Trimesh") -> "rr.Mesh3D":
 
             # Always update transform (objects may move between calls)
             # NOTE: World scene objects are already in correct world coordinates, no offset needed
-            rr.log(
-                rr_path,
-                rr.Transform3D(
-                    translation=tform[:3, 3],
-                    mat3x3=tform[:3, :3],
-                    axis_length=0.25,
-                ),
-                static=False,
-            )
+            if _RR_HAS_TRANSFORM_AXES:
+                # rerun >= 0.28: axis_length moved to TransformAxes3D
+                rr.log(rr_path, rr.Transform3D(translation=tform[:3, 3], mat3x3=tform[:3, :3]), static=False)
+                rr.log(rr_path, rr.TransformAxes3D(axis_length=0.25), static=False)
+            else:
+                # rerun <= 0.27: axis_length on Transform3D directly
+                rr.log(
+                    rr_path,
+                    rr.Transform3D(translation=tform[:3, 3], mat3x3=tform[:3, :3], axis_length=0.25),
+                    static=False,
+                )
 
             # Geometry: send only once per node to avoid duplicates
             if rr_path not in self._logged_geometry:

source/isaaclab_mimic/test/test_curobo_planner_cube_stack.py

@@ -27,6 +27,7 @@
 
 import gymnasium as gym
 import torch
+import warp as wp
 
 import isaaclab.utils.math as math_utils
 from isaaclab.assets import Articulation, RigidObject
@@ -164,7 +165,7 @@ def _pose_from_xy_quat(self, xy: torch.Tensor, z: float, quat: torch.Tensor) ->
     def _get_cube_pos(self, cube_name: str) -> torch.Tensor:
         """Return the current world position of a cube's root (x, y, z)."""
         obj: RigidObject = self.env.scene[cube_name]
-        return obj.data.root_pos_w[0, :3].clone().detach()
+        return wp.to_torch(obj.data.root_pos_w)[0, :3].clone().detach()
 
     def _place_pose_over_cube(self, cube_name: str, height_offset: float) -> torch.Tensor:
         """Compute a goal pose directly above the named cube using the latest pose."""

source/isaaclab_mimic/test/test_curobo_planner_franka.py

@@ -20,6 +20,7 @@
 
 import gymnasium as gym
 import torch
+import warp as wp
 
 import isaaclab.utils.assets as _al_assets
 import isaaclab.utils.math as math_utils
@@ -141,7 +142,7 @@ def _set_arm_positions(self, q: torch.Tensor) -> None:
             q_full = torch.cat([q, fingers], dim=-1)
         else:
             q_full = q
-        self.robot.write_joint_position_to_sim(q_full)
+        self.robot.write_joint_position_to_sim(wp.from_torch(q_full.to(self.env.device)))
 
     @pytest.mark.parametrize("goal_spec, goal_id", predefined_ee_goals_and_ids)
     def test_plan_to_predefined_goal(self, goal_spec, goal_id) -> None: