Introducing manipulation modules: factory extensions and track goals

source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) 2024-2025, The UW Lab Project Developers.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Manipulation environments for fixed-arm robots."""

source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/factory_extension/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) 2024-2025, The UW Lab Project Developers.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Factory Task Environments"""

source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/factory_extension/assembly_keypoints.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2024-2025, The UW Lab Project Developers.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+from typing import TYPE_CHECKING
+
+import isaaclab.utils.math as math_utils
+from isaaclab.utils import configclass
+
+if TYPE_CHECKING:
+    from isaaclab.assets import Articulation, RigidObject
+
+
+@configclass
+class Offset:
+    pos: tuple[float, float, float] = (0.0, 0.0, 0.0)
+    quat: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0)
+
+    @property
+    def pose(self) -> tuple[float, float, float, float, float, float, float]:
+        return self.pos + self.quat
+
+    def apply(self, root: RigidObject | Articulation) -> tuple[torch.Tensor, torch.Tensor]:
+        data = root.data.root_pos_w
+        pos_w, quat_w = math_utils.combine_frame_transforms(
+            root.data.root_pos_w,
+            root.data.root_quat_w,
+            torch.tensor(self.pos).to(data.device).repeat(data.shape[0], 1),
+            torch.tensor(self.quat).to(data.device).repeat(data.shape[0], 1),
+        )
+        return pos_w, quat_w
+
+    def combine(self, pos: torch.Tensor, quat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        num_data = pos.shape[0]
+        device = pos.device
+        return math_utils.combine_frame_transforms(
+            pos,
+            quat,
+            torch.tensor(self.pos).to(device).repeat(num_data, 1),
+            torch.tensor(self.quat).to(device).repeat(num_data, 1),
+        )
+
+
+@configclass
+class KeyPointsNistBoard:
+    bolt_m16: Offset = Offset(pos=(0.145, -0.1495, -0.01))
+    hole_8mm: Offset = Offset(pos=(-0.0895, 0.0, 0.01))
+    gear_base: Offset = Offset(pos=(0.145, 0.021, 0.0092), quat=(0.70711, 0.0, 0.0, -0.70711))
+    small_gear: Offset = Offset(pos=(0.145, 0.021, 0.0092), quat=(0.70711, 0.0, 0.0, -0.70711))
+    medium_gear: Offset = Offset(pos=(0.145, 0.021, 0.0092), quat=(0.70711, 0.0, 0.0, -0.70711))
+    large_gear: Offset = Offset(pos=(0.145, 0.021, 0.0092), quat=(0.70711, 0.0, 0.0, -0.70711))
+
+
+@configclass
+class KeyPointsBoltM16:
+    bolt_tip_offset: Offset = Offset(pos=(0, 0, 0.035))
+    bolt_base_offset: Offset = Offset(pos=(0, 0, 0.01))
+    screwed_nut_offset: Offset = Offset(pos=(0, 0, 0.032))
+
+
+@configclass
+class KeyPointsNutM16:
+    center_axis_bottom: Offset = Offset(pos=(0.0, 0.0, 0.01))
+    center_axis_middle: Offset = Offset(pos=(0.0, 0.0, 0.0165))
+    center_axis_top: Offset = Offset(pos=(0.0, 0.0, 0.023))
+    grasp_point: Offset = Offset(pos=(0.0, 0.0, 0.01), quat=(0.70711, 0.0, 0.0, -0.70711))
+    grasp_diameter: float = 0.024
+
+
+@configclass
+class KeyPointsGearBase:
+    small_gear_tip_offset: Offset = Offset(pos=(0.05075, 0.0, 0.025))
+    small_gear_assembled_bottom_offset = Offset(pos=(0.05075, 0.0, 0.005))
+    medium_gear_tip_offset: Offset = Offset(pos=(0.02025, 0.0, 0.025))
+    medium_gear_assembled_bottom_offset = Offset(pos=(0.02025, 0.0, 0.005))
+    large_gear_tip_offset: Offset = Offset(pos=(-0.03025, 0.0, 0.025))
+    large_gear_assembled_bottom_offset = Offset(pos=(-0.03025, 0.0, 0.005))
+
+
+@configclass
+class KeyPointsSmallGear:
+    center_axis_bottom: Offset = Offset(pos=(0.05075, 0.0, 0.005))
+    center_axis_top: Offset = Offset(pos=(0.05075, 0.0, 0.03))
+    grasp_point: Offset = Offset(pos=(0.05075, 0.0, 0.022))
+    grasp_diameter: float = 0.03
+
+
+@configclass
+class KeyPointsMediumGear:
+    center_axis_bottom: Offset = Offset(pos=(0.02025, 0.0, 0.005))
+    center_axis_top: Offset = Offset(pos=(0.02025, 0.0, 0.03))
+    grasp_point: Offset = Offset(pos=(0.02025, 0.0, 0.022), quat=(0.70711, 0.0, 0.0, -0.70711))
+    grasp_diameter: float = 0.03
+
+
+@configclass
+class KeyPointsLargeGear:
+    center_axis_bottom: Offset = Offset(pos=(-3.025e-2, 0.0, 0.005))
+    center_axis_top: Offset = Offset(pos=(-3.025e-2, 0.0, 0.03))
+    grasp_point: Offset = Offset(pos=(-3.025e-2, 0.0, 0.022))
+    grasp_diameter: float = 0.03
+
+
+@configclass
+class KeyPointsHole8MM:
+    hole_tip_offset: Offset = Offset(pos=(0, 0, 0.025))
+    inserted_peg_base_offset: Offset = Offset(pos=(0, 0, 0.0))
+
+
+@configclass
+class KeyPointsPeg8MM:
+    center_axis_bottom: Offset = Offset(pos=(0.0, 0.0, 0.0))
+    center_axis_middle: Offset = Offset(pos=(0.0, 0.0, 0.025))
+    center_axis_top: Offset = Offset(pos=(0.0, 0.0, 0.05))
+    grasp_point: Offset = Offset(pos=(0.0, 0.0, 0.035))
+    grasp_diameter: float = 0.007986
+
+
+@configclass
+class KeyPointPandaHand:
+    object_grasped_point: Offset = Offset(pos=(0.0, 0.0, 0.107), quat=(0.0, 0.0, 1.0, 0.0))
+
+
+KEYPOINTS_NISTBOARD = KeyPointsNistBoard()
+KEYPOINTS_BOLTM16 = KeyPointsBoltM16()
+KEYPOINTS_NUTM16 = KeyPointsNutM16()
+KEYPOINTS_GEARBASE = KeyPointsGearBase()
+KEYPOINTS_SMALLGEAR = KeyPointsSmallGear()
+KEYPOINTS_MEDIUMGEAR = KeyPointsMediumGear()
+KEYPOINTS_LARGEGEAR = KeyPointsLargeGear()
+KEYPOINTS_HOLE8MM = KeyPointsHole8MM()
+KEYPOINTS_PEG8MM = KeyPointsPeg8MM()
+KEYPOINTS_PANDAHAND = KeyPointPandaHand()

source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/factory_extension/config/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2024-2025, The UW Lab Project Developers.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Configurations for the factory environments."""
+
+# We leave this file empty since we don't want to expose any configs in this package directly.
+# We still need this file to import the "config" module in the parent package.

source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/factory_extension/config/franka/__init__.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2024-2025, The UW Lab Project Developers.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+import gymnasium as gym
+
+from . import agents
+
+gym.register(
+    id="UW-NutThread-Franka-JointPos-v0",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaNutThreadEnvCfg",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FactoryPPORunnerCfg",
+    },
+)
+
+
+gym.register(
+    id="UW-PegInsert-Franka-JointPos-v0",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaPegInsertEnvCfg",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FactoryPPORunnerCfg",
+    },
+)
+
+gym.register(
+    id="UW-GearMesh-Franka-JointPos-v0",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaGearMeshEnvCfg",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FactoryPPORunnerCfg",
+    },
+)

source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/factory_extension/config/franka/agents/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) 2024-2025, The UW Lab Project Developers.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause

source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/factory_extension/config/franka/agents/rsl_rl_ppo_cfg.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2024-2025, The UW Lab Project Developers.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from isaaclab.utils import configclass
+
+from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg
+
+
+@configclass
+class FactoryPPORunnerCfg(RslRlOnPolicyRunnerCfg):
+    num_steps_per_env = 128
+    max_iterations = 400
+    save_interval = 50
+    experiment_name = "factory"
+    empirical_normalization = False
+    policy = RslRlPpoActorCriticCfg(
+        init_noise_std=1.0,
+        actor_hidden_dims=[512, 128, 64],
+        critic_hidden_dims=[512, 128, 64],
+        activation="elu",
+    )
+    algorithm = RslRlPpoAlgorithmCfg(
+        value_loss_coef=1.0,
+        use_clipped_value_loss=True,
+        clip_param=0.2,
+        entropy_coef=1e-3,
+        num_learning_epochs=5,
+        num_mini_batches=4,
+        learning_rate=5.0e-4,
+        schedule="adaptive",
+        gamma=0.99,
+        lam=0.95,
+        desired_kl=0.02,
+        max_grad_norm=1.0,
+    )

source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/factory_extension/config/franka/ik_del_env_cfg.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2024-2025, The UW 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 import BinaryJointPositionActionCfg, DifferentialInverseKinematicsActionCfg
+from isaaclab.utils import configclass
+
+from ...factory_tasks import FRANKA_PANDA_CFG
+from ...gearmesh_env_cfg import GearMeshEnvCfg
+from ...nutthread_env_cfg import NutThreadEnvCfg
+from ...peginsert_env_cfg import PegInsertEnvCfg
+from . import joint_pos_env_cfg
+
+
+@configclass
+class FrankaFactoryIkDelEnvMixIn:
+    def __post_init__(self: joint_pos_env_cfg.FrankaNutThreadEnvCfg):
+        super().__post_init__()
+        self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
+        self.scene.robot.actuators["panda_arm1"].stiffness = 800
+        self.scene.robot.actuators["panda_arm2"].stiffness = 700
+        self.scene.robot.actuators["panda_arm1"].damping = 30
+        self.scene.robot.actuators["panda_arm2"].damping = 30
+        self.actions.arm_action = DifferentialInverseKinematicsActionCfg(
+            asset_name="robot",
+            joint_names=["panda_joint.*"],
+            body_name="panda_hand",
+            controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=True, ik_method="dls"),
+            scale=1.0,
+            body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=(0.0, 0.0, 0.107)),
+        )
+        self.actions.gripper_action = BinaryJointPositionActionCfg(
+            asset_name="robot",
+            joint_names=["panda_finger.*"],
+            open_command_expr={"panda_finger_.*": 0.04},
+            close_command_expr={"panda_finger_.*": 0.0},
+        )
+
+
+@configclass
+class FrankaNutThreadIkDeltaEnvCfg(FrankaFactoryIkDelEnvMixIn, NutThreadEnvCfg):
+    pass
+
+
+@configclass
+class FrankaGearMeshIkDeltaEnvCfg(FrankaFactoryIkDelEnvMixIn, GearMeshEnvCfg):
+    pass
+
+
+@configclass
+class FrankaPegInsertIkDeltaEnvCfg(FrankaFactoryIkDelEnvMixIn, PegInsertEnvCfg):
+    pass

source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/factory_extension/config/franka/joint_pos_env_cfg.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2024-2025, The UW Lab Project Developers.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from isaaclab.utils import configclass
+
+import uwlab_tasks.manager_based.manipulation.factory_extension.mdp as mdp
+
+from ...factory_assets_cfg import FRANKA_PANDA_CFG
+from ...factory_env_base import FactoryBaseEnvCfg
+from ...gearmesh_env_cfg import GearMeshEnvCfg
+from ...nutthread_env_cfg import NutThreadEnvCfg
+from ...peginsert_env_cfg import PegInsertEnvCfg
+
+
+@configclass
+class ActionCfg:
+    arm_action = mdp.RelativeJointPositionActionCfg(
+        asset_name="robot",
+        joint_names=["panda_joint.*"],
+        scale=0.02,
+        use_zero_offset=True,
+    )
+
+    gripper_action = mdp.BinaryJointPositionActionCfg(
+        asset_name="robot",
+        joint_names=["panda_finger.*"],
+        open_command_expr={"panda_finger_.*": 0.0},
+        close_command_expr={"panda_finger_.*": 0.0},
+    )
+
+
+@configclass
+class FrankaFactoryEnvMixIn:
+    actions: ActionCfg = ActionCfg()
+
+    def __post_init__(self: FactoryBaseEnvCfg):
+        super().__post_init__()
+        self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
+        self.scene.robot.actuators["panda_arm1"].stiffness = 80.0
+        self.scene.robot.actuators["panda_arm1"].damping = 4.0
+        self.scene.robot.actuators["panda_arm2"].stiffness = 80.0
+        self.scene.robot.actuators["panda_arm2"].damping = 4.0
+
+
+@configclass
+class FrankaNutThreadEnvCfg(FrankaFactoryEnvMixIn, NutThreadEnvCfg):
+    pass
+
+
+@configclass
+class FrankaGearMeshEnvCfg(FrankaFactoryEnvMixIn, GearMeshEnvCfg):
+    pass
+
+
+@configclass
+class FrankaPegInsertEnvCfg(FrankaFactoryEnvMixIn, PegInsertEnvCfg):
+    pass