Add so-100 envs

scripts/tools/record_demos.py

@@ -19,6 +19,8 @@
     --step_hz                 Environment stepping rate in Hz. (default: 30)
     --num_demos               Number of demonstrations to record. (default: 0)
     --num_success_steps       Number of continuous steps with task success for concluding a demo as successful. (default: 10)
+    --pos_sensitivity         Position sensitivity for teleoperation devices. (default: 0.2)
+    --rot_sensitivity         Rotation sensitivity for teleoperation devices. (default: 0.5)
 """
 
 """Launch Isaac Sim Simulator first."""
@@ -53,6 +55,18 @@
     default=False,
     help="Enable Pinocchio.",
 )
+parser.add_argument(
+    "--pos_sensitivity",
+    type=float,
+    default=0.2,
+    help="Position sensitivity for teleoperation devices. Default is 0.2.",
+)
+parser.add_argument(
+    "--rot_sensitivity",
+    type=float,
+    default=0.5,
+    help="Rotation sensitivity for teleoperation devices. Default is 0.5.",
+)
 
 # append AppLauncher cli args
 AppLauncher.add_app_launcher_args(parser)
@@ -269,9 +283,13 @@ def setup_teleop_device(callbacks: dict[str, Callable]) -> object:
             omni.log.warn(f"No teleop device '{args_cli.teleop_device}' found in environment config. Creating default.")
             # Create fallback teleop device
             if args_cli.teleop_device.lower() == "keyboard":
-                teleop_interface = Se3Keyboard(Se3KeyboardCfg(pos_sensitivity=0.2, rot_sensitivity=0.5))
+                teleop_interface = Se3Keyboard(
+                    Se3KeyboardCfg(pos_sensitivity=args_cli.pos_sensitivity, rot_sensitivity=args_cli.rot_sensitivity)
+                )
             elif args_cli.teleop_device.lower() == "spacemouse":
-                teleop_interface = Se3SpaceMouse(Se3SpaceMouseCfg(pos_sensitivity=0.2, rot_sensitivity=0.5))
+                teleop_interface = Se3SpaceMouse(
+                    Se3SpaceMouseCfg(pos_sensitivity=args_cli.pos_sensitivity, rot_sensitivity=args_cli.rot_sensitivity)
+                )
             else:
                 omni.log.error(f"Unsupported teleop device: {args_cli.teleop_device}")
                 omni.log.error("Supported devices: keyboard, spacemouse, handtracking")

source/isaaclab_assets/isaaclab_assets/robots/so_100.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Configuration for the Franka Emika robots.
+
+The following configurations are available:
+
+* :obj:`SO_100_CFG`: SO-100 robot
+
+Reference: https://github.com/TheRobotStudio/SO-ARM100
+"""
+
+import isaaclab.sim as sim_utils
+from isaaclab.actuators import ImplicitActuatorCfg
+from isaaclab.assets.articulation import ArticulationCfg
+from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR
+
+##
+# Configuration
+##
+
+SO_100_CFG = ArticulationCfg(
+    spawn=sim_utils.UsdFileCfg(
+        usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/RobotStudio/so100/so100.usd",
+        activate_contact_sensors=False,
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+            disable_gravity=True,
+            retain_accelerations=False,
+            linear_damping=0.0,
+            angular_damping=0.0,
+            max_linear_velocity=1000.0,
+            max_angular_velocity=1000.0,
+            max_depenetration_velocity=1.0,
+        ),
+        articulation_props=sim_utils.ArticulationRootPropertiesCfg(
+            enabled_self_collisions=True, solver_position_iteration_count=8, solver_velocity_iteration_count=4
+        ),
+    ),
+    init_state=ArticulationCfg.InitialStateCfg(
+        joint_pos={
+            "shoulder_pan": 0.0,
+            "shoulder_lift": 0.0,
+            "elbow_flex": 0.0,
+            "wrist_flex": 0.0,
+            "wrist_roll": 0.0,
+            "gripper": 0.0,
+        },
+    ),
+    actuators={
+        "all_joints": ImplicitActuatorCfg(
+            joint_names_expr=["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll"],
+            effort_limit_sim=None,
+            velocity_limit_sim=None,
+            stiffness=None,
+            damping=None,
+            armature=0.0,
+        ),
+        "gripper": ImplicitActuatorCfg(
+            joint_names_expr=["gripper"],
+            effort_limit_sim=0.1,
+            velocity_limit_sim=2.175,
+            stiffness=8.0,
+            damping=0.4,
+            armature=0.0,
+        ),
+    },
+    soft_joint_pos_limit_factor=1.0,
+)
+"""Configuration of SO‑100 robot arm."""

source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/__init__.py

@@ -11,6 +11,8 @@
 from .nutpour_gr1t2_mimic_env_cfg import NutPourGR1T2MimicEnvCfg
 from .pickplace_gr1t2_mimic_env import PickPlaceGR1T2MimicEnv
 from .pickplace_gr1t2_mimic_env_cfg import PickPlaceGR1T2MimicEnvCfg
+from .stack_so100_mimic_env import StackSO100MimicEnv
+from .stack_so100_mimic_env_cfg import StackSO100MimicEnvCfg
 
 gym.register(
     id="Isaac-PickPlace-GR1T2-Abs-Mimic-v0",
@@ -34,3 +36,12 @@
     kwargs={"env_cfg_entry_point": exhaustpipe_gr1t2_mimic_env_cfg.ExhaustPipeGR1T2MimicEnvCfg},
     disable_env_checker=True,
 )
+
+gym.register(
+    id="Isaac-Stack-SO100-Abs-Mimic-v0",
+    entry_point="isaaclab_mimic.envs.pinocchio_envs:StackSO100MimicEnv",
+    kwargs={
+        "env_cfg_entry_point": stack_so100_mimic_env_cfg.StackSO100MimicEnvCfg,
+    },
+    disable_env_checker=True,
+)

source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/stack_so100_mimic_env.py

@@ -0,0 +1,142 @@
+# Copyright (c) 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
+
+# Copyright (c) 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 StackSO100MimicEnv(ManagerBasedRLMimicEnv):
+
+    def get_robot_eef_pose(self, eef_name: str, env_ids: Sequence[int] | None = None) -> torch.Tensor:
+        """
+        Get current robot end effector pose. Should be the same frame as used by the robot end-effector controller.
+
+        Args:
+            eef_name: Name of the end effector (for SO100, this is typically "gripper").
+            env_ids: Environment indices to get the pose for. If None, all envs are considered.
+
+        Returns:
+            A torch.Tensor eef pose matrix. Shape is (len(env_ids), 4, 4)
+        """
+        if env_ids is None:
+            env_ids = slice(None)
+
+        # For SO100, the end effector pose is stored in policy observations
+        eef_pos_name = "eef_pos"
+        eef_quat_name = "eef_quat"
+
+        target_eef_position = self.obs_buf["policy"][eef_pos_name][env_ids]
+        target_rot_mat = PoseUtils.matrix_from_quat(self.obs_buf["policy"][eef_quat_name][env_ids])
+
+        return PoseUtils.make_pose(target_eef_position, target_rot_mat)
+
+    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:
+        """
+        Takes a target pose and gripper action for the end effector controller and returns an action
+        (usually a normalized delta pose action) to try and achieve that target pose.
+        Noise is added to the target pose action if specified.
+
+        Args:
+            target_eef_pose_dict: Dictionary containing target eef pose. For SO100, expects key "gripper".
+            gripper_action_dict: Dictionary containing gripper action. For SO100, expects key "gripper".
+            action_noise_dict: Noise to add to the action. If None, no noise is added.
+            env_id: Environment index to get the action for.
+
+        Returns:
+            An action torch.Tensor that's compatible with env.step().
+        """
+
+        # target position and rotation for single arm
+        target_eef_pos, target_rot = PoseUtils.unmake_pose(target_eef_pose_dict["gripper"])
+        target_eef_rot_quat = PoseUtils.quat_from_matrix(target_rot)
+
+        # gripper action
+        gripper_action = gripper_action_dict["gripper"]
+
+        if action_noise_dict is not None:
+            pos_noise = action_noise_dict["gripper"] * torch.randn_like(target_eef_pos)
+            quat_noise = action_noise_dict["gripper"] * torch.randn_like(target_eef_rot_quat)
+
+            target_eef_pos += pos_noise
+            target_eef_rot_quat += quat_noise
+
+        return torch.cat(
+            (
+                target_eef_pos,
+                target_eef_rot_quat,
+                gripper_action,
+            ),
+            dim=0,
+        )
+
+    def action_to_target_eef_pose(self, action: torch.Tensor) -> dict[str, torch.Tensor]:
+        """
+        Converts action (compatible with env.step) to a target pose for the end effector controller.
+        Inverse of @target_eef_pose_to_action. Usually used to infer a sequence of target controller poses
+        from a demonstration trajectory using the recorded actions.
+
+        Args:
+            action: Environment action. Shape is (num_envs, action_dim).
+
+        Returns:
+            A dictionary of eef pose torch.Tensor that @action corresponds to.
+        """
+        target_poses = {}
+
+        # For SO100: action dimensions are [pos(3), quat(4), gripper(1)]
+        target_eef_position = action[:, 0:3]
+        target_rot_mat = PoseUtils.matrix_from_quat(action[:, 3:7])
+        target_pose = PoseUtils.make_pose(target_eef_position, target_rot_mat)
+        target_poses["gripper"] = target_pose
+
+        return target_poses
+
+    def actions_to_gripper_actions(self, actions: torch.Tensor) -> dict[str, torch.Tensor]:
+        """
+        Extracts the gripper actuation part from a sequence of env actions (compatible with env.step).
+
+        Args:
+            actions: environment actions. The shape is (num_envs, num steps in a demo, action_dim).
+
+        Returns:
+            A dictionary of torch.Tensor gripper actions. Key is "gripper".
+        """
+        print("actions", actions)
+        print("actions", actions.shape)
+        return {"gripper": actions[:, -1:]}  # Gripper is at index 7 (single dimension)
+
+    def get_subtask_term_signals(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]:
+        """
+        Gets a dictionary of termination signal flags for each subtask in a task. The flag is 1
+        when the subtask has been completed and 0 otherwise.
+
+        Args:
+            env_ids: Environment indices to get the termination signals for. If None, all envs are considered.
+
+        Returns:
+            A dictionary termination signal flags (False or True) for each subtask.
+        """
+        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]
+        return signals

source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/stack_so100_mimic_env_cfg.py

@@ -0,0 +1,98 @@
+# Copyright (c) 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
+
+# Copyright (c) 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.so_100.stack_pink_ik_abs_visuomotor_env_cfg import (
+    SO100CubeStackPinkIKAbsVisuomotorEnvCfg,
+)
+
+
+@configclass
+class StackSO100MimicEnvCfg(SO100CubeStackPinkIKAbsVisuomotorEnvCfg, MimicEnvCfg):
+
+    def __post_init__(self):
+        # Calling post init of parents
+        super().__post_init__()
+
+        # Override the existing values
+        self.datagen_config.name = "demo_src_so100_stack_task_D0"
+        self.datagen_config.generation_guarantee = True
+        self.datagen_config.generation_keep_failed = False
+        self.datagen_config.generation_num_trials = 1000
+        self.datagen_config.generation_select_src_per_subtask = False
+        self.datagen_config.generation_select_src_per_arm = False
+        self.datagen_config.generation_relative = False
+        self.datagen_config.generation_joint_pos = False
+        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.num_demo_to_render = 10
+        self.datagen_config.num_fail_demo_to_render = 25
+        self.datagen_config.seed = 1
+
+        # The following are the subtask configurations for the stack task.
+        # For SO100, we only have one "gripper" arm, so all subtasks are for this single arm
+        subtask_configs = []
+
+        # Subtask 1: Pick up cube_2
+        subtask_configs.append(
+            SubTaskConfig(
+                # Each subtask involves manipulation with respect to a single object frame.
+                # For picking up cube_2, the object of interest is cube_2
+                object_ref="cube_2",
+                # This key corresponds to the binary indicator in "datagen_info" that signals
+                # when this subtask is finished (e.g., on a 0 to 1 edge).
+                subtask_term_signal="grasp_1",
+                first_subtask_start_offset_range=(0, 0),
+                # Randomization range for starting index of the first subtask
+                subtask_term_offset_range=(0, 0),
+                # Selection strategy for the source subtask segment during data generation
+                selection_strategy="nearest_neighbor_object",
+                # Optional parameters for the selection strategy function
+                selection_strategy_kwargs={"nn_k": 3},
+                # Amount of action noise to apply during this subtask
+                action_noise=0.0,
+                # Number of interpolation steps to bridge to this subtask segment
+                num_interpolation_steps=0,
+                # Additional fixed steps for the robot to reach the necessary pose
+                num_fixed_steps=0,
+                # If True, apply action noise during the interpolation phase and execution
+                apply_noise_during_interpolation=False,
+            )
+        )
+
+        # Subtask 2: Place cube_2 on cube_1
+        subtask_configs.append(
+            SubTaskConfig(
+                # For placing cube_2 on cube_1, the object of interest is cube_1 (target location)
+                object_ref="cube_1",
+                # Corresponding key for the binary indicator in "datagen_info" for completion
+                subtask_term_signal=None,  # This is the final subtask, so no termination signal
+                # Time offsets for data generation when splitting a trajectory
+                subtask_term_offset_range=(0, 0),
+                # Selection strategy for source subtask segment
+                selection_strategy="nearest_neighbor_object",
+                # Optional parameters for the selection strategy function
+                selection_strategy_kwargs={"nn_k": 3},
+                # Amount of action noise to apply during this subtask
+                action_noise=0.0,
+                # Number of interpolation steps to bridge to this subtask segment
+                num_interpolation_steps=3,
+                # Additional fixed steps for the robot to reach the necessary pose
+                num_fixed_steps=0,
+                # If True, apply action noise during the interpolation phase and execution
+                apply_noise_during_interpolation=False,
+            )
+        )
+
+        # For SO100, we only have one arm "gripper", so we assign all subtasks to it
+        self.subtask_configs["gripper"] = subtask_configs