diff --git a/docs/source/_static/tasks/manipulation/kuka_allegro_lift.jpg b/docs/source/_static/tasks/manipulation/kuka_allegro_lift.jpg new file mode 100644 index 00000000000..9d19b0e4236 Binary files /dev/null and b/docs/source/_static/tasks/manipulation/kuka_allegro_lift.jpg differ diff --git a/docs/source/_static/tasks/manipulation/kuka_allegro_reorient.jpg b/docs/source/_static/tasks/manipulation/kuka_allegro_reorient.jpg new file mode 100644 index 00000000000..384e7632e44 Binary files /dev/null and b/docs/source/_static/tasks/manipulation/kuka_allegro_reorient.jpg differ diff --git a/docs/source/overview/environments.rst b/docs/source/overview/environments.rst index c4925adfb94..afc1fdc0e99 100644 --- a/docs/source/overview/environments.rst +++ b/docs/source/overview/environments.rst @@ -100,51 +100,55 @@ for the lift-cube environment: .. table:: :widths: 33 37 30 - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | World | Environment ID | Description | - +======================+===========================+=============================================================================+ - | |reach-franka| | |reach-franka-link| | Move the end-effector to a sampled target pose with the Franka robot | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | |reach-ur10| | |reach-ur10-link| | Move the end-effector to a sampled target pose with the UR10 robot | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | |deploy-reach-ur10e| | |deploy-reach-ur10e-link| | Move the end-effector to a sampled target pose with the UR10e robot | - | | | This policy has been deployed to a real robot | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | |lift-cube| | |lift-cube-link| | Pick a cube and bring it to a sampled target position with the Franka robot | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | |stack-cube| | |stack-cube-link| | Stack three cubes (bottom to top: blue, red, green) with the Franka robot. | - | | | Blueprint env used for the NVIDIA Isaac GR00T blueprint for synthetic | - | | |stack-cube-bp-link| | manipulation motion generation | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | |surface-gripper| | |long-suction-link| | Stack three cubes (bottom to top: blue, red, green) | - | | | with the UR10 arm and long surface gripper | - | | |short-suction-link| | or short surface gripper. | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | |cabi-franka| | |cabi-franka-link| | Grasp the handle of a cabinet's drawer and open it with the Franka robot | - | | | | - | | |franka-direct-link| | | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | |cube-allegro| | |cube-allegro-link| | In-hand reorientation of a cube using Allegro hand | - | | | | - | | |allegro-direct-link| | | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | |cube-shadow| | |cube-shadow-link| | In-hand reorientation of a cube using Shadow hand | - | | | | - | | |cube-shadow-ff-link| | | - | | | | - | | |cube-shadow-lstm-link| | | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | |cube-shadow| | |cube-shadow-vis-link| | In-hand reorientation of a cube using Shadow hand using perceptive inputs. | - | | | Requires running with ``--enable_cameras``. | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | |gr1_pick_place| | |gr1_pick_place-link| | Pick up and place an object in a basket with a GR-1 humanoid robot | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | |gr1_pp_waist| | |gr1_pp_waist-link| | Pick up and place an object in a basket with a GR-1 humanoid robot | - | | | with waist degrees-of-freedom enables that provides a wider reach space. | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ - | |galbot_stack| | |galbot_stack-link| | Stack three cubes (bottom to top: blue, red, green) with the left arm of | - | | | a Galbot humanoid robot | - +----------------------+---------------------------+-----------------------------------------------------------------------------+ + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | World | Environment ID | Description | + +=========================+==============================+=============================================================================+ + | |reach-franka| | |reach-franka-link| | Move the end-effector to a sampled target pose with the Franka robot | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |reach-ur10| | |reach-ur10-link| | Move the end-effector to a sampled target pose with the UR10 robot | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |deploy-reach-ur10e| | |deploy-reach-ur10e-link| | Move the end-effector to a sampled target pose with the UR10e robot | + | | | This policy has been deployed to a real robot | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |lift-cube| | |lift-cube-link| | Pick a cube and bring it to a sampled target position with the Franka robot | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |stack-cube| | |stack-cube-link| | Stack three cubes (bottom to top: blue, red, green) with the Franka robot. | + | | | Blueprint env used for the NVIDIA Isaac GR00T blueprint for synthetic | + | | |stack-cube-bp-link| | manipulation motion generation | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |surface-gripper| | |long-suction-link| | Stack three cubes (bottom to top: blue, red, green) | + | | | with the UR10 arm and long surface gripper | + | | |short-suction-link| | or short surface gripper. | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |cabi-franka| | |cabi-franka-link| | Grasp the handle of a cabinet's drawer and open it with the Franka robot | + | | | | + | | |franka-direct-link| | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |cube-allegro| | |cube-allegro-link| | In-hand reorientation of a cube using Allegro hand | + | | | | + | | |allegro-direct-link| | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |cube-shadow| | |cube-shadow-link| | In-hand reorientation of a cube using Shadow hand | + | | | | + | | |cube-shadow-ff-link| | | + | | | | + | | |cube-shadow-lstm-link| | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |cube-shadow| | |cube-shadow-vis-link| | In-hand reorientation of a cube using Shadow hand using perceptive inputs. | + | | | Requires running with ``--enable_cameras``. | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |gr1_pick_place| | |gr1_pick_place-link| | Pick up and place an object in a basket with a GR-1 humanoid robot | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |gr1_pp_waist| | |gr1_pp_waist-link| | Pick up and place an object in a basket with a GR-1 humanoid robot | + | | | with waist degrees-of-freedom enables that provides a wider reach space. | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |kuka-allegro-lift| | |kuka-allegro-lift-link| | Pick up a primitive shape on the table and lift it to target position | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |kuka-allegro-reorient| | |kuka-allegro-reorient-link| | Pick up a primitive shape on the table and orient it to target pose | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + | |galbot_stack| | |galbot_stack-link| | Stack three cubes (bottom to top: blue, red, green) with the left arm of | + | | | a Galbot humanoid robot | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+ .. |reach-franka| image:: ../_static/tasks/manipulation/franka_reach.jpg .. |reach-ur10| image:: ../_static/tasks/manipulation/ur10_reach.jpg @@ -157,8 +161,9 @@ for the lift-cube environment: .. |gr1_pick_place| image:: ../_static/tasks/manipulation/gr-1_pick_place.jpg .. |surface-gripper| image:: ../_static/tasks/manipulation/ur10_stack_surface_gripper.jpg .. |gr1_pp_waist| image:: ../_static/tasks/manipulation/gr-1_pick_place_waist.jpg -.. |surface-gripper| image:: ../_static/tasks/manipulation/ur10_stack_surface_gripper.jpg .. |galbot_stack| image:: ../_static/tasks/manipulation/galbot_stack_cube.jpg +.. |kuka-allegro-lift| image:: ../_static/tasks/manipulation/kuka_allegro_lift.jpg +.. |kuka-allegro-reorient| image:: ../_static/tasks/manipulation/kuka_allegro_reorient.jpg .. |reach-franka-link| replace:: `Isaac-Reach-Franka-v0 `__ .. |reach-ur10-link| replace:: `Isaac-Reach-UR10-v0 `__ @@ -177,8 +182,8 @@ for the lift-cube environment: .. |short-suction-link| replace:: `Isaac-Stack-Cube-UR10-Short-Suction-IK-Rel-v0 `__ .. |gr1_pp_waist-link| replace:: `Isaac-PickPlace-GR1T2-WaistEnabled-Abs-v0 `__ .. |galbot_stack-link| replace:: `Isaac-Stack-Cube-Galbot-Left-Arm-Gripper-RmpFlow-v0 `__ -.. |long-suction-link| replace:: `Isaac-Stack-Cube-UR10-Long-Suction-IK-Rel-v0 `__ -.. |short-suction-link| replace:: `Isaac-Stack-Cube-UR10-Short-Suction-IK-Rel-v0 `__ +.. |kuka-allegro-lift-link| replace:: `Isaac-Dexsuite-Kuka-Allegro-Lift-v0 `__ +.. |kuka-allegro-reorient-link| replace:: `Isaac-Dexsuite-Kuka-Allegro-Reorient-v0 `__ .. |cube-shadow-link| replace:: `Isaac-Repose-Cube-Shadow-Direct-v0 `__ .. |cube-shadow-ff-link| replace:: `Isaac-Repose-Cube-Shadow-OpenAI-FF-Direct-v0 `__ @@ -942,6 +947,14 @@ inferencing, including reading from an already trained checkpoint and disabling - - Manager Based - + * - Isaac-Dexsuite-Kuka-Allegro-Lift-v0 + - Isaac-Dexsuite-Kuka-Allegro-Lift-Play-v0 + - Manager Based + - **rl_games** (PPO), **rsl_rl** (PPO) + * - Isaac-Dexsuite-Kuka-Allegro-Reorient-v0 + - Isaac-Dexsuite-Kuka-Allegro-Reorient-Play-v0 + - Manager Based + - **rl_games** (PPO), **rsl_rl** (PPO) * - Isaac-Stack-Cube-Franka-v0 - - Manager Based diff --git a/source/isaaclab_assets/isaaclab_assets/robots/__init__.py b/source/isaaclab_assets/isaaclab_assets/robots/__init__.py index a5996104680..82a13a05e49 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/__init__.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/__init__.py @@ -18,6 +18,7 @@ from .humanoid import * from .humanoid_28 import * from .kinova import * +from .kuka_allegro import * from .pick_and_place import * from .quadcopter import * from .ridgeback_franka import * diff --git a/source/isaaclab_assets/isaaclab_assets/robots/kuka_allegro.py b/source/isaaclab_assets/isaaclab_assets/robots/kuka_allegro.py new file mode 100644 index 00000000000..d6c86bb3f15 --- /dev/null +++ b/source/isaaclab_assets/isaaclab_assets/robots/kuka_allegro.py @@ -0,0 +1,114 @@ +# 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 Kuka-lbr-iiwa arm robots and Allegro Hand. + +The following configurations are available: + +* :obj:`KUKA_ALLEGRO_CFG`: Kuka Allegro with implicit actuator model. + +Reference: + +* https://www.kuka.com/en-us/products/robotics-systems/industrial-robots/lbr-iiwa +* https://www.wonikrobotics.com/robot-hand + +""" + +import isaaclab.sim as sim_utils +from isaaclab.actuators.actuator_cfg import ImplicitActuatorCfg +from isaaclab.assets.articulation import ArticulationCfg +from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR + +## +# Configuration +## + +KUKA_ALLEGRO_CFG = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAACLAB_NUCLEUS_DIR}/Robots/KukaAllegro/kuka.usd", + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + retain_accelerations=True, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=1000.0, + max_depenetration_velocity=1000.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, + solver_position_iteration_count=32, + solver_velocity_iteration_count=1, + sleep_threshold=0.005, + stabilization_threshold=0.0005, + ), + joint_drive_props=sim_utils.JointDrivePropertiesCfg(drive_type="force"), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.0), + rot=(1.0, 0.0, 0.0, 0.0), + joint_pos={ + "iiwa7_joint_(1|2|7)": 0.0, + "iiwa7_joint_3": 0.7854, + "iiwa7_joint_4": 1.5708, + "iiwa7_joint_(5|6)": -1.5708, + "(index|middle|ring)_joint_0": 0.0, + "(index|middle|ring)_joint_1": 0.3, + "(index|middle|ring)_joint_2": 0.3, + "(index|middle|ring)_joint_3": 0.3, + "thumb_joint_0": 1.5, + "thumb_joint_1": 0.60147215, + "thumb_joint_2": 0.33795027, + "thumb_joint_3": 0.60845138, + }, + ), + actuators={ + "kuka_allegro_actuators": ImplicitActuatorCfg( + joint_names_expr=[ + "iiwa7_joint_(1|2|3|4|5|6|7)", + "index_joint_(0|1|2|3)", + "middle_joint_(0|1|2|3)", + "ring_joint_(0|1|2|3)", + "thumb_joint_(0|1|2|3)", + ], + effort_limit_sim={ + "iiwa7_joint_(1|2|3|4|5|6|7)": 300.0, + "index_joint_(0|1|2|3)": 0.5, + "middle_joint_(0|1|2|3)": 0.5, + "ring_joint_(0|1|2|3)": 0.5, + "thumb_joint_(0|1|2|3)": 0.5, + }, + stiffness={ + "iiwa7_joint_(1|2|3|4)": 300.0, + "iiwa7_joint_5": 100.0, + "iiwa7_joint_6": 50.0, + "iiwa7_joint_7": 25.0, + "index_joint_(0|1|2|3)": 3.0, + "middle_joint_(0|1|2|3)": 3.0, + "ring_joint_(0|1|2|3)": 3.0, + "thumb_joint_(0|1|2|3)": 3.0, + }, + damping={ + "iiwa7_joint_(1|2|3|4)": 45.0, + "iiwa7_joint_5": 20.0, + "iiwa7_joint_6": 15.0, + "iiwa7_joint_7": 15.0, + "index_joint_(0|1|2|3)": 0.1, + "middle_joint_(0|1|2|3)": 0.1, + "ring_joint_(0|1|2|3)": 0.1, + "thumb_joint_(0|1|2|3)": 0.1, + }, + friction={ + "iiwa7_joint_(1|2|3|4|5|6|7)": 1.0, + "index_joint_(0|1|2|3)": 0.01, + "middle_joint_(0|1|2|3)": 0.01, + "ring_joint_(0|1|2|3)": 0.01, + "thumb_joint_(0|1|2|3)": 0.01, + }, + ), + }, + soft_joint_pos_limit_factor=1.0, +) diff --git a/source/isaaclab_tasks/config/extension.toml b/source/isaaclab_tasks/config/extension.toml index f317365d688..1a6d1b88d07 100644 --- a/source/isaaclab_tasks/config/extension.toml +++ b/source/isaaclab_tasks/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.10.51" +version = "0.11.0" # Description title = "Isaac Lab Environments" diff --git a/source/isaaclab_tasks/docs/CHANGELOG.rst b/source/isaaclab_tasks/docs/CHANGELOG.rst index ee84acbafd5..cda01d77b03 100644 --- a/source/isaaclab_tasks/docs/CHANGELOG.rst +++ b/source/isaaclab_tasks/docs/CHANGELOG.rst @@ -1,6 +1,15 @@ Changelog --------- +0.11.0 (2025-09-07) +~~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added dextrous lifting and dextrous reorientation manipulation rl environments. + + 0.10.51 (2025-09-08) ~~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/__init__.py new file mode 100644 index 00000000000..26075d4da25 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/__init__.py @@ -0,0 +1,26 @@ +# 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 + +"""Dexsuite environments. + +Implementation Reference: + +Reorient: +@article{petrenko2023dexpbt, + title={Dexpbt: Scaling up dexterous manipulation for hand-arm systems with population based training}, + author={Petrenko, Aleksei and Allshire, Arthur and State, Gavriel and Handa, Ankur and Makoviychuk, Viktor}, + journal={arXiv preprint arXiv:2305.12127}, + year={2023} +} + +Lift: +@article{singh2024dextrah, + title={Dextrah-rgb: Visuomotor policies to grasp anything with dexterous hands}, + author={Singh, Ritvik and Allshire, Arthur and Handa, Ankur and Ratliff, Nathan and Van Wyk, Karl}, + journal={arXiv preprint arXiv:2412.01791}, + year={2024} +} + +""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/adr_curriculum.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/adr_curriculum.py new file mode 100644 index 00000000000..52fef8b494a --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/adr_curriculum.py @@ -0,0 +1,122 @@ +# 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 + +from isaaclab.managers import CurriculumTermCfg as CurrTerm +from isaaclab.utils import configclass + +from . import mdp + + +@configclass +class CurriculumCfg: + """Curriculum terms for the MDP.""" + + # adr stands for automatic/adaptive domain randomization + adr = CurrTerm( + func=mdp.DifficultyScheduler, params={"init_difficulty": 0, "min_difficulty": 0, "max_difficulty": 10} + ) + + joint_pos_unoise_min_adr = CurrTerm( + func=mdp.modify_term_cfg, + params={ + "address": "observations.proprio.joint_pos.noise.n_min", + "modify_fn": mdp.initial_final_interpolate_fn, + "modify_params": {"initial_value": 0.0, "final_value": -0.1, "difficulty_term_str": "adr"}, + }, + ) + + joint_pos_unoise_max_adr = CurrTerm( + func=mdp.modify_term_cfg, + params={ + "address": "observations.proprio.joint_pos.noise.n_max", + "modify_fn": mdp.initial_final_interpolate_fn, + "modify_params": {"initial_value": 0.0, "final_value": 0.1, "difficulty_term_str": "adr"}, + }, + ) + + joint_vel_unoise_min_adr = CurrTerm( + func=mdp.modify_term_cfg, + params={ + "address": "observations.proprio.joint_vel.noise.n_min", + "modify_fn": mdp.initial_final_interpolate_fn, + "modify_params": {"initial_value": 0.0, "final_value": -0.2, "difficulty_term_str": "adr"}, + }, + ) + + joint_vel_unoise_max_adr = CurrTerm( + func=mdp.modify_term_cfg, + params={ + "address": "observations.proprio.joint_vel.noise.n_max", + "modify_fn": mdp.initial_final_interpolate_fn, + "modify_params": {"initial_value": 0.0, "final_value": 0.2, "difficulty_term_str": "adr"}, + }, + ) + + hand_tips_pos_unoise_min_adr = CurrTerm( + func=mdp.modify_term_cfg, + params={ + "address": "observations.proprio.hand_tips_state_b.noise.n_min", + "modify_fn": mdp.initial_final_interpolate_fn, + "modify_params": {"initial_value": 0.0, "final_value": -0.01, "difficulty_term_str": "adr"}, + }, + ) + + hand_tips_pos_unoise_max_adr = CurrTerm( + func=mdp.modify_term_cfg, + params={ + "address": "observations.proprio.hand_tips_state_b.noise.n_max", + "modify_fn": mdp.initial_final_interpolate_fn, + "modify_params": {"initial_value": 0.0, "final_value": 0.01, "difficulty_term_str": "adr"}, + }, + ) + + object_quat_unoise_min_adr = CurrTerm( + func=mdp.modify_term_cfg, + params={ + "address": "observations.policy.object_quat_b.noise.n_min", + "modify_fn": mdp.initial_final_interpolate_fn, + "modify_params": {"initial_value": 0.0, "final_value": -0.03, "difficulty_term_str": "adr"}, + }, + ) + + object_quat_unoise_max_adr = CurrTerm( + func=mdp.modify_term_cfg, + params={ + "address": "observations.policy.object_quat_b.noise.n_max", + "modify_fn": mdp.initial_final_interpolate_fn, + "modify_params": {"initial_value": 0.0, "final_value": 0.03, "difficulty_term_str": "adr"}, + }, + ) + + object_obs_unoise_min_adr = CurrTerm( + func=mdp.modify_term_cfg, + params={ + "address": "observations.perception.object_point_cloud.noise.n_min", + "modify_fn": mdp.initial_final_interpolate_fn, + "modify_params": {"initial_value": 0.0, "final_value": -0.01, "difficulty_term_str": "adr"}, + }, + ) + + object_obs_unoise_max_adr = CurrTerm( + func=mdp.modify_term_cfg, + params={ + "address": "observations.perception.object_point_cloud.noise.n_max", + "modify_fn": mdp.initial_final_interpolate_fn, + "modify_params": {"initial_value": 0.0, "final_value": -0.01, "difficulty_term_str": "adr"}, + }, + ) + + gravity_adr = CurrTerm( + func=mdp.modify_term_cfg, + params={ + "address": "events.variable_gravity.params.gravity_distribution_params", + "modify_fn": mdp.initial_final_interpolate_fn, + "modify_params": { + "initial_value": ((0.0, 0.0, 0.0), (0.0, 0.0, 0.0)), + "final_value": ((0.0, 0.0, -9.81), (0.0, 0.0, -9.81)), + "difficulty_term_str": "adr", + }, + }, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/__init__.py new file mode 100644 index 00000000000..4240e604428 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/__init__.py @@ -0,0 +1,9 @@ +# 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 + +"""Configurations for the dexsuite 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. diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/__init__.py new file mode 100644 index 00000000000..159ab6727fb --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/__init__.py @@ -0,0 +1,63 @@ +# 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 + +""" +Dextra Kuka Allegro environments. +""" + +import gymnasium as gym + +from . import agents + +## +# Register Gym environments. +## + +# State Observation +gym.register( + id="Isaac-Dexsuite-Kuka-Allegro-Reorient-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.dexsuite_kuka_allegro_env_cfg:DexsuiteKukaAllegroReorientEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DexsuiteKukaAllegroPPORunnerCfg", + }, +) + +gym.register( + id="Isaac-Dexsuite-Kuka-Allegro-Reorient-Play-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.dexsuite_kuka_allegro_env_cfg:DexsuiteKukaAllegroReorientEnvCfg_PLAY", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DexsuiteKukaAllegroPPORunnerCfg", + }, +) + +# Dexsuite Lift Environments +gym.register( + id="Isaac-Dexsuite-Kuka-Allegro-Lift-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.dexsuite_kuka_allegro_env_cfg:DexsuiteKukaAllegroLiftEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DexsuiteKukaAllegroPPORunnerCfg", + }, +) + + +gym.register( + id="Isaac-Dexsuite-Kuka-Allegro-Lift-Play-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.dexsuite_kuka_allegro_env_cfg:DexsuiteKukaAllegroLiftEnvCfg_PLAY", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DexsuiteKukaAllegroPPORunnerCfg", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/__init__.py new file mode 100644 index 00000000000..2e924fbf1b1 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/__init__.py @@ -0,0 +1,4 @@ +# 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 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rl_games_ppo_cfg.yaml b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rl_games_ppo_cfg.yaml new file mode 100644 index 00000000000..3d23b745cc5 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rl_games_ppo_cfg.yaml @@ -0,0 +1,86 @@ +# 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 + +params: + seed: 42 + + # environment wrapper clipping + env: + clip_observations: 100.0 + clip_actions: 100.0 + obs_groups: + obs: ["policy", "proprio", "perception"] + states: ["policy", "proprio", "perception"] + concate_obs_groups: True + + algo: + name: a2c_continuous + + model: + name: continuous_a2c_logstd + + network: + name: actor_critic + separate: True + space: + continuous: + mu_activation: None + sigma_activation: None + mu_init: + name: default + sigma_init: + name: const_initializer + val: 0 + fixed_sigma: True + mlp: + units: [512, 256, 128] + activation: elu + d2rl: False + initializer: + name: default + regularizer: + name: None + + load_checkpoint: False # flag which sets whether to load the checkpoint + load_path: '' # path to the checkpoint to load + + config: + name: reorient + env_name: rlgpu + device: 'cuda:0' + device_name: 'cuda:0' + multi_gpu: False + ppo: True + mixed_precision: False + normalize_input: True + normalize_value: True + value_bootstrap: False + num_actors: -1 + reward_shaper: + scale_value: 0.01 + normalize_advantage: True + gamma: 0.99 + tau: 0.95 + learning_rate: 1e-3 + lr_schedule: adaptive + schedule_type: legacy + kl_threshold: 0.01 + score_to_win: 100000000 + max_epochs: 750000 + save_best_after: 100 + save_frequency: 50 + print_stats: True + grad_norm: 1.0 + entropy_coef: 0.001 + truncate_grads: True + e_clip: 0.2 + horizon_length: 36 + minibatch_size: 36864 + mini_epochs: 5 + critic_coef: 4 + clip_value: True + clip_actions: False + seq_len: 4 + bounds_loss_coef: 0.0001 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rsl_rl_ppo_cfg.py new file mode 100644 index 00000000000..f7965575737 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rsl_rl_ppo_cfg.py @@ -0,0 +1,39 @@ +# 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 + +from isaaclab.utils import configclass + +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg + + +@configclass +class DexsuiteKukaAllegroPPORunnerCfg(RslRlOnPolicyRunnerCfg): + num_steps_per_env = 32 + obs_groups = {"policy": ["policy", "proprio", "perception"], "critic": ["policy", "proprio", "perception"]} + max_iterations = 15000 + save_interval = 250 + experiment_name = "dexsuite_kuka_allegro" + policy = RslRlPpoActorCriticCfg( + init_noise_std=1.0, + actor_obs_normalization=True, + critic_obs_normalization=True, + actor_hidden_dims=[512, 256, 128], + critic_hidden_dims=[512, 256, 128], + activation="elu", + ) + algorithm = RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=0.005, + num_learning_epochs=5, + num_mini_batches=4, + learning_rate=1.0e-3, + schedule="adaptive", + gamma=0.99, + lam=0.95, + desired_kl=0.01, + max_grad_norm=1.0, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/dexsuite_kuka_allegro_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/dexsuite_kuka_allegro_env_cfg.py new file mode 100644 index 00000000000..6c41414f30b --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/dexsuite_kuka_allegro_env_cfg.py @@ -0,0 +1,79 @@ +# 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 + +from isaaclab_assets.robots import KUKA_ALLEGRO_CFG + +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import ContactSensorCfg +from isaaclab.utils import configclass + +from ... import dexsuite_env_cfg as dexsuite +from ... import mdp + + +@configclass +class KukaAllegroRelJointPosActionCfg: + action = mdp.RelativeJointPositionActionCfg(asset_name="robot", joint_names=[".*"], scale=0.1) + + +@configclass +class KukaAllegroReorientRewardCfg(dexsuite.RewardsCfg): + + # bool awarding term if 2 finger tips are in contact with object, one of the contacting fingers has to be thumb. + good_finger_contact = RewTerm( + func=mdp.contacts, + weight=0.5, + params={"threshold": 1.0}, + ) + + +@configclass +class KukaAllegroMixinCfg: + rewards: KukaAllegroReorientRewardCfg = KukaAllegroReorientRewardCfg() + actions: KukaAllegroRelJointPosActionCfg = KukaAllegroRelJointPosActionCfg() + + def __post_init__(self: dexsuite.DexsuiteReorientEnvCfg): + super().__post_init__() + self.commands.object_pose.body_name = "palm_link" + self.scene.robot = KUKA_ALLEGRO_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + finger_tip_body_list = ["index_link_3", "middle_link_3", "ring_link_3", "thumb_link_3"] + for link_name in finger_tip_body_list: + setattr( + self.scene, + f"{link_name}_object_s", + ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Robot/ee_link/" + link_name, + filter_prim_paths_expr=["{ENV_REGEX_NS}/Object"], + ), + ) + self.observations.proprio.contact = ObsTerm( + func=mdp.fingers_contact_force_b, + params={"contact_sensor_names": [f"{link}_object_s" for link in finger_tip_body_list]}, + clip=(-20.0, 20.0), # contact force in finger tips is under 20N normally + ) + self.observations.proprio.hand_tips_state_b.params["body_asset_cfg"].body_names = ["palm_link", ".*_tip"] + self.rewards.fingers_to_object.params["asset_cfg"] = SceneEntityCfg("robot", body_names=["palm_link", ".*_tip"]) + + +@configclass +class DexsuiteKukaAllegroReorientEnvCfg(KukaAllegroMixinCfg, dexsuite.DexsuiteReorientEnvCfg): + pass + + +@configclass +class DexsuiteKukaAllegroReorientEnvCfg_PLAY(KukaAllegroMixinCfg, dexsuite.DexsuiteReorientEnvCfg_PLAY): + pass + + +@configclass +class DexsuiteKukaAllegroLiftEnvCfg(KukaAllegroMixinCfg, dexsuite.DexsuiteLiftEnvCfg): + pass + + +@configclass +class DexsuiteKukaAllegroLiftEnvCfg_PLAY(KukaAllegroMixinCfg, dexsuite.DexsuiteLiftEnvCfg_PLAY): + pass diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/dexsuite_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/dexsuite_env_cfg.py new file mode 100644 index 00000000000..75e40c5c74b --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/dexsuite_env_cfg.py @@ -0,0 +1,466 @@ +# 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 + +from dataclasses import MISSING + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg +from isaaclab.envs import ManagerBasedEnvCfg, ViewerCfg +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim import CapsuleCfg, ConeCfg, CuboidCfg, RigidBodyMaterialCfg, SphereCfg +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise + +from . import mdp +from .adr_curriculum import CurriculumCfg + + +@configclass +class SceneCfg(InteractiveSceneCfg): + """Dexsuite Scene for multi-objects Lifting""" + + # robot + robot: ArticulationCfg = MISSING + + # object + object: RigidObjectCfg = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Object", + spawn=sim_utils.MultiAssetSpawnerCfg( + assets_cfg=[ + CuboidCfg(size=(0.05, 0.1, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CuboidCfg(size=(0.05, 0.05, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CuboidCfg(size=(0.025, 0.1, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CuboidCfg(size=(0.025, 0.05, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CuboidCfg(size=(0.025, 0.025, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CuboidCfg(size=(0.01, 0.1, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + SphereCfg(radius=0.05, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + SphereCfg(radius=0.025, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CapsuleCfg(radius=0.04, height=0.025, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CapsuleCfg(radius=0.04, height=0.01, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CapsuleCfg(radius=0.04, height=0.1, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CapsuleCfg(radius=0.025, height=0.1, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CapsuleCfg(radius=0.025, height=0.2, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CapsuleCfg(radius=0.01, height=0.2, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + ConeCfg(radius=0.05, height=0.1, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + ConeCfg(radius=0.025, height=0.1, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + ], + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=16, + solver_velocity_iteration_count=0, + disable_gravity=False, + ), + collision_props=sim_utils.CollisionPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=0.2), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(-0.55, 0.1, 0.35)), + ) + + # table + table: RigidObjectCfg = RigidObjectCfg( + prim_path="/World/envs/env_.*/table", + spawn=sim_utils.CuboidCfg( + size=(0.8, 1.5, 0.04), + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True), + collision_props=sim_utils.CollisionPropertiesCfg(), + # trick: we let visualizer's color to show the table with success coloring + visible=False, + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(-0.55, 0.0, 0.235), rot=(1.0, 0.0, 0.0, 0.0)), + ) + + # plane + plane = AssetBaseCfg( + prim_path="/World/GroundPlane", + init_state=AssetBaseCfg.InitialStateCfg(), + spawn=sim_utils.GroundPlaneCfg(), + collision_group=-1, + ) + + # lights + sky_light = AssetBaseCfg( + prim_path="/World/skyLight", + spawn=sim_utils.DomeLightCfg( + intensity=750.0, + texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr", + ), + ) + + +@configclass +class CommandsCfg: + """Command terms for the MDP.""" + + object_pose = mdp.ObjectUniformPoseCommandCfg( + asset_name="robot", + object_name="object", + resampling_time_range=(3.0, 5.0), + debug_vis=False, + ranges=mdp.ObjectUniformPoseCommandCfg.Ranges( + pos_x=(-0.7, -0.3), + pos_y=(-0.25, 0.25), + pos_z=(0.55, 0.95), + roll=(-3.14, 3.14), + pitch=(-3.14, 3.14), + yaw=(0.0, 0.0), + ), + success_vis_asset_name="table", + ) + + +@configclass +class ObservationsCfg: + """Observation specifications for the MDP.""" + + @configclass + class PolicyCfg(ObsGroup): + """Observations for policy group.""" + + object_quat_b = ObsTerm(func=mdp.object_quat_b, noise=Unoise(n_min=-0.0, n_max=0.0)) + target_object_pose_b = ObsTerm(func=mdp.generated_commands, params={"command_name": "object_pose"}) + actions = ObsTerm(func=mdp.last_action) + + def __post_init__(self): + self.enable_corruption = True + self.concatenate_terms = True + self.history_length = 5 + + @configclass + class ProprioObsCfg(ObsGroup): + """Observations for proprioception group.""" + + joint_pos = ObsTerm(func=mdp.joint_pos, noise=Unoise(n_min=-0.0, n_max=0.0)) + joint_vel = ObsTerm(func=mdp.joint_vel, noise=Unoise(n_min=-0.0, n_max=0.0)) + hand_tips_state_b = ObsTerm( + func=mdp.body_state_b, + noise=Unoise(n_min=-0.0, n_max=0.0), + # good behaving number for position in m, velocity in m/s, rad/s, + # and quaternion are unlikely to exceed -2 to 2 range + clip=(-2.0, 2.0), + params={ + "body_asset_cfg": SceneEntityCfg("robot"), + "base_asset_cfg": SceneEntityCfg("robot"), + }, + ) + contact: ObsTerm = MISSING + + def __post_init__(self): + self.enable_corruption = True + self.concatenate_terms = True + self.history_length = 5 + + @configclass + class PerceptionObsCfg(ObsGroup): + + object_point_cloud = ObsTerm( + func=mdp.object_point_cloud_b, + noise=Unoise(n_min=-0.0, n_max=0.0), + clip=(-2.0, 2.0), # clamp between -2 m to 2 m + params={"num_points": 64, "flatten": True}, + ) + + def __post_init__(self): + self.enable_corruption = True + self.concatenate_dim = 0 + self.concatenate_terms = True + self.flatten_history_dim = True + self.history_length = 5 + + # observation groups + policy: PolicyCfg = PolicyCfg() + proprio: ProprioObsCfg = ProprioObsCfg() + perception: PerceptionObsCfg = PerceptionObsCfg() + + +@configclass +class EventCfg: + """Configuration for randomization.""" + + # -- pre-startup + randomize_object_scale = EventTerm( + func=mdp.randomize_rigid_body_scale, + mode="prestartup", + params={"scale_range": (0.75, 1.5), "asset_cfg": SceneEntityCfg("object")}, + ) + + robot_physics_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=".*"), + "static_friction_range": [0.5, 1.0], + "dynamic_friction_range": [0.5, 1.0], + "restitution_range": [0.0, 0.0], + "num_buckets": 250, + }, + ) + + object_physics_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("object", body_names=".*"), + "static_friction_range": [0.5, 1.0], + "dynamic_friction_range": [0.5, 1.0], + "restitution_range": [0.0, 0.0], + "num_buckets": 250, + }, + ) + + joint_stiffness_and_damping = EventTerm( + func=mdp.randomize_actuator_gains, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), + "stiffness_distribution_params": [0.5, 2.0], + "damping_distribution_params": [0.5, 2.0], + "operation": "scale", + }, + ) + + joint_friction = EventTerm( + func=mdp.randomize_joint_parameters, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), + "friction_distribution_params": [0.0, 5.0], + "operation": "scale", + }, + ) + + object_scale_mass = EventTerm( + func=mdp.randomize_rigid_body_mass, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("object"), + "mass_distribution_params": [0.2, 2.0], + "operation": "scale", + }, + ) + + reset_table = EventTerm( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "pose_range": {"x": [-0.05, 0.05], "y": [-0.05, 0.05], "z": [0.0, 0.0]}, + "velocity_range": {"x": [-0.0, 0.0], "y": [-0.0, 0.0], "z": [-0.0, 0.0]}, + "asset_cfg": SceneEntityCfg("table"), + }, + ) + + reset_object = EventTerm( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "pose_range": { + "x": [-0.2, 0.2], + "y": [-0.2, 0.2], + "z": [0.0, 0.4], + "roll": [-3.14, 3.14], + "pitch": [-3.14, 3.14], + "yaw": [-3.14, 3.14], + }, + "velocity_range": {"x": [-0.0, 0.0], "y": [-0.0, 0.0], "z": [-0.0, 0.0]}, + "asset_cfg": SceneEntityCfg("object"), + }, + ) + + reset_root = EventTerm( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "pose_range": {"x": [-0.0, 0.0], "y": [-0.0, 0.0], "yaw": [-0.0, 0.0]}, + "velocity_range": {"x": [-0.0, 0.0], "y": [-0.0, 0.0], "z": [-0.0, 0.0]}, + "asset_cfg": SceneEntityCfg("robot"), + }, + ) + + reset_robot_joints = EventTerm( + func=mdp.reset_joints_by_offset, + mode="reset", + params={ + "position_range": [-0.50, 0.50], + "velocity_range": [0.0, 0.0], + }, + ) + + reset_robot_wrist_joint = EventTerm( + func=mdp.reset_joints_by_offset, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names="iiwa7_joint_7"), + "position_range": [-3, 3], + "velocity_range": [0.0, 0.0], + }, + ) + + # Note (Octi): This is a deliberate trick in Remake to accelerate learning. + # By scheduling gravity as a curriculum — starting with no gravity (easy) + # and gradually introducing full gravity (hard) — the agent learns more smoothly. + # This removes the need for a special "Lift" reward (often required to push the + # agent to counter gravity), which has bonus effect of simplifying reward composition overall. + variable_gravity = EventTerm( + func=mdp.randomize_physics_scene_gravity, + mode="reset", + params={ + "gravity_distribution_params": ([0.0, 0.0, 0.0], [0.0, 0.0, 0.0]), + "operation": "abs", + }, + ) + + +@configclass +class ActionsCfg: + pass + + +@configclass +class RewardsCfg: + """Reward terms for the MDP.""" + + action_l2 = RewTerm(func=mdp.action_l2_clamped, weight=-0.005) + + action_rate_l2 = RewTerm(func=mdp.action_rate_l2_clamped, weight=-0.005) + + fingers_to_object = RewTerm(func=mdp.object_ee_distance, params={"std": 0.4}, weight=1.0) + + position_tracking = RewTerm( + func=mdp.position_command_error_tanh, + weight=2.0, + params={ + "asset_cfg": SceneEntityCfg("robot"), + "std": 0.2, + "command_name": "object_pose", + "align_asset_cfg": SceneEntityCfg("object"), + }, + ) + + orientation_tracking = RewTerm( + func=mdp.orientation_command_error_tanh, + weight=4.0, + params={ + "asset_cfg": SceneEntityCfg("robot"), + "std": 1.5, + "command_name": "object_pose", + "align_asset_cfg": SceneEntityCfg("object"), + }, + ) + + success = RewTerm( + func=mdp.success_reward, + weight=10, + params={ + "asset_cfg": SceneEntityCfg("robot"), + "pos_std": 0.1, + "rot_std": 0.5, + "command_name": "object_pose", + "align_asset_cfg": SceneEntityCfg("object"), + }, + ) + + early_termination = RewTerm(func=mdp.is_terminated_term, weight=-1, params={"term_keys": "abnormal_robot"}) + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + time_out = DoneTerm(func=mdp.time_out, time_out=True) + + object_out_of_bound = DoneTerm( + func=mdp.out_of_bound, + params={ + "in_bound_range": {"x": (-1.5, 0.5), "y": (-2.0, 2.0), "z": (0.0, 2.0)}, + "asset_cfg": SceneEntityCfg("object"), + }, + ) + + abnormal_robot = DoneTerm(func=mdp.abnormal_robot_state) + + +@configclass +class DexsuiteReorientEnvCfg(ManagerBasedEnvCfg): + """Dexsuite reorientation task definition, also the base definition for derivative Lift task and evaluation task""" + + # Scene settings + viewer: ViewerCfg = ViewerCfg(eye=(-2.25, 0.0, 0.75), lookat=(0.0, 0.0, 0.45), origin_type="env") + scene: SceneCfg = SceneCfg(num_envs=4096, env_spacing=3, replicate_physics=False) + # Basic settings + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + # MDP settings + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventCfg = EventCfg() + curriculum: CurriculumCfg | None = CurriculumCfg() + + def __post_init__(self): + """Post initialization.""" + # general settings + self.decimation = 2 # 50 Hz + + # *single-goal setup + self.commands.object_pose.resampling_time_range = (10.0, 10.0) + self.commands.object_pose.position_only = False + self.commands.object_pose.success_visualizer_cfg.markers["failure"] = self.scene.table.spawn.replace( + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.25, 0.15, 0.15), roughness=0.25), visible=True + ) + self.commands.object_pose.success_visualizer_cfg.markers["success"] = self.scene.table.spawn.replace( + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.15, 0.25, 0.15), roughness=0.25), visible=True + ) + + self.episode_length_s = 4.0 + self.is_finite_horizon = True + + # simulation settings + self.sim.dt = 1 / 120 + self.sim.render_interval = self.decimation + self.sim.physx.bounce_threshold_velocity = 0.2 + self.sim.physx.bounce_threshold_velocity = 0.01 + self.sim.physx.gpu_max_rigid_patch_count = 4 * 5 * 2**15 + + if self.curriculum is not None: + self.curriculum.adr.params["pos_tol"] = self.rewards.success.params["pos_std"] / 2 + self.curriculum.adr.params["rot_tol"] = self.rewards.success.params["rot_std"] / 2 + + +class DexsuiteLiftEnvCfg(DexsuiteReorientEnvCfg): + """Dexsuite lift task definition""" + + def __post_init__(self): + super().__post_init__() + self.rewards.orientation_tracking = None # no orientation reward + self.commands.object_pose.position_only = True + if self.curriculum is not None: + self.rewards.success.params["rot_std"] = None # make success reward not consider orientation + self.curriculum.adr.params["rot_tol"] = None # make adr not tracking orientation + + +class DexsuiteReorientEnvCfg_PLAY(DexsuiteReorientEnvCfg): + """Dexsuite reorientation task evaluation environment definition""" + + def __post_init__(self): + super().__post_init__() + self.commands.object_pose.resampling_time_range = (2.0, 3.0) + self.commands.object_pose.debug_vis = True + self.curriculum.adr.params["init_difficulty"] = self.curriculum.adr.params["max_difficulty"] + + +class DexsuiteLiftEnvCfg_PLAY(DexsuiteLiftEnvCfg): + """Dexsuite lift task evaluation environment definition""" + + def __post_init__(self): + super().__post_init__() + self.commands.object_pose.resampling_time_range = (2.0, 3.0) + self.commands.object_pose.debug_vis = True + self.commands.object_pose.position_only = True + self.curriculum.adr.params["init_difficulty"] = self.curriculum.adr.params["max_difficulty"] diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.py new file mode 100644 index 00000000000..794113f9253 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.py @@ -0,0 +1,12 @@ +# 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 + +from isaaclab.envs.mdp import * # noqa: F401, F403 + +from .commands import * # noqa: F401, F403 +from .curriculums import * # noqa: F401, F403 +from .observations import * # noqa: F401, F403 +from .rewards import * # noqa: F401, F403 +from .terminations import * # noqa: F401, F403 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/__init__.py new file mode 100644 index 00000000000..a5132558174 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/__init__.py @@ -0,0 +1,6 @@ +# 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 + +from .pose_commands_cfg import * # noqa: F401, F403 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands.py new file mode 100644 index 00000000000..146eee9741d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands.py @@ -0,0 +1,179 @@ +# 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 + + +"""Sub-module containing command generators for pose tracking.""" + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import CommandTerm +from isaaclab.markers import VisualizationMarkers +from isaaclab.utils.math import combine_frame_transforms, compute_pose_error, quat_from_euler_xyz, quat_unique + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from . import pose_commands_cfg as dex_cmd_cfgs + + +class ObjectUniformPoseCommand(CommandTerm): + """Uniform pose command generator for an object (in the robot base frame). + + This command term samples target object poses by: + • Drawing (x, y, z) uniformly within configured Cartesian bounds, and + • Drawing roll-pitch-yaw uniformly within configured ranges, then converting + to a quaternion (w, x, y, z). Optionally makes quaternions unique by enforcing + a positive real part. + + Frames: + Targets are defined in the robot's *base frame*. For metrics/visualization, + targets are transformed into the *world frame* using the robot root pose. + + Outputs: + The command buffer has shape (num_envs, 7): `(x, y, z, qw, qx, qy, qz)`. + + Metrics: + `position_error` and `orientation_error` are computed between the commanded + world-frame pose and the object's current world-frame pose. + + Config: + `cfg` must provide the sampling ranges, whether to enforce quaternion uniqueness, + and optional visualization settings. + """ + + cfg: dex_cmd_cfgs.ObjectUniformPoseCommandCfg + """Configuration for the command generator.""" + + def __init__(self, cfg: dex_cmd_cfgs.ObjectUniformPoseCommandCfg, env: ManagerBasedEnv): + """Initialize the command generator class. + + Args: + cfg: The configuration parameters for the command generator. + env: The environment object. + """ + # initialize the base class + super().__init__(cfg, env) + + # extract the robot and body index for which the command is generated + self.robot: Articulation = env.scene[cfg.asset_name] + self.object: RigidObject = env.scene[cfg.object_name] + self.success_vis_asset: RigidObject = env.scene[cfg.success_vis_asset_name] + + # create buffers + # -- commands: (x, y, z, qw, qx, qy, qz) in root frame + self.pose_command_b = torch.zeros(self.num_envs, 7, device=self.device) + self.pose_command_b[:, 3] = 1.0 + self.pose_command_w = torch.zeros_like(self.pose_command_b) + # -- metrics + self.metrics["position_error"] = torch.zeros(self.num_envs, device=self.device) + self.metrics["orientation_error"] = torch.zeros(self.num_envs, device=self.device) + + self.success_visualizer = VisualizationMarkers(self.cfg.success_visualizer_cfg) + self.success_visualizer.set_visibility(True) + + def __str__(self) -> str: + msg = "UniformPoseCommand:\n" + msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n" + msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n" + return msg + + """ + Properties + """ + + @property + def command(self) -> torch.Tensor: + """The desired pose command. Shape is (num_envs, 7). + + The first three elements correspond to the position, followed by the quaternion orientation in (w, x, y, z). + """ + return self.pose_command_b + + """ + Implementation specific functions. + """ + + def _update_metrics(self): + # transform command from base frame to simulation world frame + self.pose_command_w[:, :3], self.pose_command_w[:, 3:] = combine_frame_transforms( + self.robot.data.root_pos_w, + self.robot.data.root_quat_w, + self.pose_command_b[:, :3], + self.pose_command_b[:, 3:], + ) + # compute the error + pos_error, rot_error = compute_pose_error( + self.pose_command_w[:, :3], + self.pose_command_w[:, 3:], + self.object.data.root_state_w[:, :3], + self.object.data.root_state_w[:, 3:7], + ) + self.metrics["position_error"] = torch.norm(pos_error, dim=-1) + self.metrics["orientation_error"] = torch.norm(rot_error, dim=-1) + + success_id = self.metrics["position_error"] < 0.05 + if not self.cfg.position_only: + success_id &= self.metrics["orientation_error"] < 0.5 + self.success_visualizer.visualize(self.success_vis_asset.data.root_pos_w, marker_indices=success_id.int()) + + def _resample_command(self, env_ids: Sequence[int]): + # sample new pose targets + # -- position + r = torch.empty(len(env_ids), device=self.device) + self.pose_command_b[env_ids, 0] = r.uniform_(*self.cfg.ranges.pos_x) + self.pose_command_b[env_ids, 1] = r.uniform_(*self.cfg.ranges.pos_y) + self.pose_command_b[env_ids, 2] = r.uniform_(*self.cfg.ranges.pos_z) + # -- orientation + euler_angles = torch.zeros_like(self.pose_command_b[env_ids, :3]) + euler_angles[:, 0].uniform_(*self.cfg.ranges.roll) + euler_angles[:, 1].uniform_(*self.cfg.ranges.pitch) + euler_angles[:, 2].uniform_(*self.cfg.ranges.yaw) + quat = quat_from_euler_xyz(euler_angles[:, 0], euler_angles[:, 1], euler_angles[:, 2]) + # make sure the quaternion has real part as positive + self.pose_command_b[env_ids, 3:] = quat_unique(quat) if self.cfg.make_quat_unique else quat + + def _update_command(self): + pass + + def _set_debug_vis_impl(self, debug_vis: bool): + # create markers if necessary for the first tome + if debug_vis: + if not hasattr(self, "goal_visualizer"): + # -- goal pose + self.goal_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg) + # -- current body pose + self.curr_visualizer = VisualizationMarkers(self.cfg.curr_pose_visualizer_cfg) + # set their visibility to true + self.goal_visualizer.set_visibility(True) + self.curr_visualizer.set_visibility(True) + else: + if hasattr(self, "goal_visualizer"): + self.goal_visualizer.set_visibility(False) + self.curr_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + # check if robot is initialized + # note: this is needed in-case the robot is de-initialized. we can't access the data + if not self.robot.is_initialized: + return + # update the markers + if not self.cfg.position_only: + # -- goal pose + self.goal_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:]) + # -- current object pose + self.curr_visualizer.visualize(self.object.data.root_pos_w, self.object.data.root_quat_w) + else: + distance = torch.norm(self.pose_command_w[:, :3] - self.object.data.root_pos_w[:, :3], dim=1) + success_id = (distance < 0.05).int() + # note: since marker indices for position is 1(far) and 2(near), we can simply shift the success_id by 1. + # -- goal position + self.goal_visualizer.visualize(self.pose_command_w[:, :3], marker_indices=success_id + 1) + # -- current object position + self.curr_visualizer.visualize(self.object.data.root_pos_w, marker_indices=success_id + 1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands_cfg.py new file mode 100644 index 00000000000..8501c00116d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands_cfg.py @@ -0,0 +1,92 @@ +# 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 + +from dataclasses import MISSING + +import isaaclab.sim as sim_utils +from isaaclab.managers import CommandTermCfg +from isaaclab.markers import VisualizationMarkersCfg +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + +from . import pose_commands as dex_cmd + +ALIGN_MARKER_CFG = VisualizationMarkersCfg( + markers={ + "frame": sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/UIElements/frame_prim.usd", + scale=(0.1, 0.1, 0.1), + ), + "position_far": sim_utils.SphereCfg( + radius=0.01, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0)), + ), + "position_near": sim_utils.SphereCfg( + radius=0.01, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0)), + ), + } +) + + +@configclass +class ObjectUniformPoseCommandCfg(CommandTermCfg): + """Configuration for uniform pose command generator.""" + + class_type: type = dex_cmd.ObjectUniformPoseCommand + + asset_name: str = MISSING + """Name of the coordinate referencing asset in the environment for which the commands are generated respect to.""" + + object_name: str = MISSING + """Name of the object in the environment for which the commands are generated.""" + + make_quat_unique: bool = False + """Whether to make the quaternion unique or not. Defaults to False. + + If True, the quaternion is made unique by ensuring the real part is positive. + """ + + @configclass + class Ranges: + """Uniform distribution ranges for the pose commands.""" + + pos_x: tuple[float, float] = MISSING + """Range for the x position (in m).""" + + pos_y: tuple[float, float] = MISSING + """Range for the y position (in m).""" + + pos_z: tuple[float, float] = MISSING + """Range for the z position (in m).""" + + roll: tuple[float, float] = MISSING + """Range for the roll angle (in rad).""" + + pitch: tuple[float, float] = MISSING + """Range for the pitch angle (in rad).""" + + yaw: tuple[float, float] = MISSING + """Range for the yaw angle (in rad).""" + + ranges: Ranges = MISSING + """Ranges for the commands.""" + + position_only: bool = True + """Command goal position only. Command includes goal quat if False""" + + # Pose Markers + goal_pose_visualizer_cfg: VisualizationMarkersCfg = ALIGN_MARKER_CFG.replace(prim_path="/Visuals/Command/goal_pose") + """The configuration for the goal pose visualization marker. Defaults to FRAME_MARKER_CFG.""" + + curr_pose_visualizer_cfg: VisualizationMarkersCfg = ALIGN_MARKER_CFG.replace(prim_path="/Visuals/Command/body_pose") + """The configuration for the current pose visualization marker. Defaults to FRAME_MARKER_CFG.""" + + success_vis_asset_name: str = MISSING + """Name of the asset in the environment for which the success color are indicated.""" + + # success markers + success_visualizer_cfg = VisualizationMarkersCfg(prim_path="/Visuals/SuccessMarkers", markers={}) + """The configuration for the success visualization marker. User needs to add the markers""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/curriculums.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/curriculums.py new file mode 100644 index 00000000000..c1a8c0f0d66 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/curriculums.py @@ -0,0 +1,113 @@ +# 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 + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation, RigidObject +from isaaclab.envs import mdp +from isaaclab.managers import ManagerTermBase, SceneEntityCfg +from isaaclab.utils.math import combine_frame_transforms, compute_pose_error + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def initial_final_interpolate_fn(env: ManagerBasedRLEnv, env_id, data, initial_value, final_value, difficulty_term_str): + """ + Interpolate between initial value iv and final value fv, for any arbitrarily + nested structure of lists/tuples in 'data'. Scalars (int/float) are handled + at the leaves. + """ + # get the fraction scalar on the device + difficulty_term: DifficultyScheduler = getattr(env.curriculum_manager.cfg, difficulty_term_str).func + frac = difficulty_term.difficulty_frac + if frac < 0.1: + # no-op during start, since the difficulty fraction near 0 is wasting of resource. + return mdp.modify_env_param.NO_CHANGE + + # convert iv/fv to tensors, but we'll peel them apart in recursion + initial_value_tensor = torch.tensor(initial_value, device=env.device) + final_value_tensor = torch.tensor(final_value, device=env.device) + + return _recurse(initial_value_tensor.tolist(), final_value_tensor.tolist(), data, frac) + + +def _recurse(iv_elem, fv_elem, data_elem, frac): + # If it's a sequence, rebuild the same type with each element recursed + if isinstance(data_elem, Sequence) and not isinstance(data_elem, (str, bytes)): + # Note: we assume initial value element and final value element have the same structure as data + return type(data_elem)(_recurse(iv_e, fv_e, d_e, frac) for iv_e, fv_e, d_e in zip(iv_elem, fv_elem, data_elem)) + # Otherwise it's a leaf scalar: do the interpolation + new_val = frac * (fv_elem - iv_elem) + iv_elem + if isinstance(data_elem, int): + return int(new_val.item()) + else: + # cast floats or any numeric + return new_val.item() + + +class DifficultyScheduler(ManagerTermBase): + """Adaptive difficulty scheduler for curriculum learning. + + Tracks per-environment difficulty levels and adjusts them based on task performance. Difficulty increases when + position/orientation errors fall below given tolerances, and decreases otherwise (unless `promotion_only` is set). + The normalized average difficulty across environments is exposed as `difficulty_frac` for use in curriculum + interpolation. + + Args: + cfg: Configuration object specifying scheduler parameters. + env: The manager-based RL environment. + + """ + + def __init__(self, cfg, env): + super().__init__(cfg, env) + init_difficulty = self.cfg.params.get("init_difficulty", 0) + self.current_adr_difficulties = torch.ones(env.num_envs, device=env.device) * init_difficulty + self.difficulty_frac = 0 + + def get_state(self): + return self.current_adr_difficulties + + def set_state(self, state: torch.Tensor): + self.current_adr_difficulties = state.clone().to(self._env.device) + + def __call__( + self, + env: ManagerBasedRLEnv, + env_ids: Sequence[int], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), + pos_tol: float = 0.1, + rot_tol: float | None = None, + init_difficulty: int = 0, + min_difficulty: int = 0, + max_difficulty: int = 50, + promotion_only: bool = False, + ): + asset: Articulation = env.scene[asset_cfg.name] + object: RigidObject = env.scene[object_cfg.name] + command = env.command_manager.get_command("object_pose") + des_pos_w, des_quat_w = combine_frame_transforms( + asset.data.root_pos_w[env_ids], asset.data.root_quat_w[env_ids], command[env_ids, :3], command[env_ids, 3:7] + ) + pos_err, rot_err = compute_pose_error( + des_pos_w, des_quat_w, object.data.root_pos_w[env_ids], object.data.root_quat_w[env_ids] + ) + pos_dist = torch.norm(pos_err, dim=1) + rot_dist = torch.norm(rot_err, dim=1) + move_up = (pos_dist < pos_tol) & (rot_dist < rot_tol) if rot_tol else pos_dist < pos_tol + demot = self.current_adr_difficulties[env_ids] if promotion_only else self.current_adr_difficulties[env_ids] - 1 + self.current_adr_difficulties[env_ids] = torch.where( + move_up, + self.current_adr_difficulties[env_ids] + 1, + demot, + ).clamp(min=min_difficulty, max=max_difficulty) + self.difficulty_frac = torch.mean(self.current_adr_difficulties) / max(max_difficulty, 1) + return self.difficulty_frac diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/observations.py new file mode 100644 index 00000000000..b48e4bcfb5c --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/observations.py @@ -0,0 +1,197 @@ +# 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 + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import ManagerTermBase, SceneEntityCfg +from isaaclab.utils.math import quat_apply, quat_apply_inverse, quat_inv, quat_mul, subtract_frame_transforms + +from .utils import sample_object_point_cloud + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def object_pos_b( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), +): + """Object position in the robot's root frame. + + Args: + env: The environment. + robot_cfg: Scene entity for the robot (reference frame). Defaults to ``SceneEntityCfg("robot")``. + object_cfg: Scene entity for the object. Defaults to ``SceneEntityCfg("object")``. + + Returns: + Tensor of shape ``(num_envs, 3)``: object position [x, y, z] expressed in the robot root frame. + """ + robot: RigidObject = env.scene[robot_cfg.name] + object: RigidObject = env.scene[object_cfg.name] + return quat_apply_inverse(robot.data.root_quat_w, object.data.root_pos_w - robot.data.root_pos_w) + + +def object_quat_b( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> torch.Tensor: + """Object orientation in the robot's root frame. + + Args: + env: The environment. + robot_cfg: Scene entity for the robot (reference frame). Defaults to ``SceneEntityCfg("robot")``. + object_cfg: Scene entity for the object. Defaults to ``SceneEntityCfg("object")``. + + Returns: + Tensor of shape ``(num_envs, 4)``: object quaternion ``(w, x, y, z)`` in the robot root frame. + """ + robot: RigidObject = env.scene[robot_cfg.name] + object: RigidObject = env.scene[object_cfg.name] + return quat_mul(quat_inv(robot.data.root_quat_w), object.data.root_quat_w) + + +def body_state_b( + env: ManagerBasedRLEnv, + body_asset_cfg: SceneEntityCfg, + base_asset_cfg: SceneEntityCfg, +) -> torch.Tensor: + """Body state (pos, quat, lin vel, ang vel) in the base asset's root frame. + + The state for each body is stacked horizontally as + ``[position(3), quaternion(4)(wxyz), linvel(3), angvel(3)]`` and then concatenated over bodies. + + Args: + env: The environment. + body_asset_cfg: Scene entity for the articulated body whose links are observed. + base_asset_cfg: Scene entity providing the reference (root) frame. + + Returns: + Tensor of shape ``(num_envs, num_bodies * 13)`` with per-body states expressed in the base root frame. + """ + body_asset: Articulation = env.scene[body_asset_cfg.name] + base_asset: Articulation = env.scene[base_asset_cfg.name] + # get world pose of bodies + body_pos_w = body_asset.data.body_pos_w[:, body_asset_cfg.body_ids].view(-1, 3) + body_quat_w = body_asset.data.body_quat_w[:, body_asset_cfg.body_ids].view(-1, 4) + body_lin_vel_w = body_asset.data.body_lin_vel_w[:, body_asset_cfg.body_ids].view(-1, 3) + body_ang_vel_w = body_asset.data.body_ang_vel_w[:, body_asset_cfg.body_ids].view(-1, 3) + num_bodies = int(body_pos_w.shape[0] / env.num_envs) + # get world pose of base frame + root_pos_w = base_asset.data.root_link_pos_w.unsqueeze(1).repeat_interleave(num_bodies, dim=1).view(-1, 3) + root_quat_w = base_asset.data.root_link_quat_w.unsqueeze(1).repeat_interleave(num_bodies, dim=1).view(-1, 4) + # transform from world body pose to local body pose + body_pos_b, body_quat_b = subtract_frame_transforms(root_pos_w, root_quat_w, body_pos_w, body_quat_w) + body_lin_vel_b = quat_apply_inverse(root_quat_w, body_lin_vel_w) + body_ang_vel_b = quat_apply_inverse(root_quat_w, body_ang_vel_w) + # concate and return + out = torch.cat((body_pos_b, body_quat_b, body_lin_vel_b, body_ang_vel_b), dim=1) + return out.view(env.num_envs, -1) + + +class object_point_cloud_b(ManagerTermBase): + """Object surface point cloud expressed in a reference asset's root frame. + + Points are pre-sampled on the object's surface in its local frame and transformed to world, + then into the reference (e.g., robot) root frame. Optionally visualizes the points. + + Args (from ``cfg.params``): + object_cfg: Scene entity for the object to sample. Defaults to ``SceneEntityCfg("object")``. + ref_asset_cfg: Scene entity providing the reference frame. Defaults to ``SceneEntityCfg("robot")``. + num_points: Number of points to sample on the object surface. Defaults to ``10``. + visualize: Whether to draw markers for the points. Defaults to ``True``. + static: If ``True``, cache world-space points on reset and reuse them (no per-step resampling). + + Returns (from ``__call__``): + If ``flatten=False``: tensor of shape ``(num_envs, num_points, 3)``. + If ``flatten=True``: tensor of shape ``(num_envs, 3 * num_points)``. + """ + + def __init__(self, cfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + + self.object_cfg: SceneEntityCfg = cfg.params.get("object_cfg", SceneEntityCfg("object")) + self.ref_asset_cfg: SceneEntityCfg = cfg.params.get("ref_asset_cfg", SceneEntityCfg("robot")) + num_points: int = cfg.params.get("num_points", 10) + self.object: RigidObject = env.scene[self.object_cfg.name] + self.ref_asset: Articulation = env.scene[self.ref_asset_cfg.name] + # lazy initialize visualizer and point cloud + if cfg.params.get("visualize", True): + from isaaclab.markers import VisualizationMarkers + from isaaclab.markers.config import RAY_CASTER_MARKER_CFG + + ray_cfg = RAY_CASTER_MARKER_CFG.replace(prim_path="/Visuals/ObservationPointCloud") + ray_cfg.markers["hit"].radius = 0.0025 + self.visualizer = VisualizationMarkers(ray_cfg) + self.points_local = sample_object_point_cloud( + env.num_envs, num_points, self.object.cfg.prim_path, device=env.device + ) + self.points_w = torch.zeros_like(self.points_local) + + def __call__( + self, + env: ManagerBasedRLEnv, + ref_asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), + num_points: int = 10, + flatten: bool = False, + visualize: bool = True, + ): + """Compute the object point cloud in the reference asset's root frame. + + Note: + Points are pre-sampled at initialization using ``self.num_points``; the ``num_points`` argument is + kept for API symmetry and does not change the sampled set at runtime. + + Args: + env: The environment. + ref_asset_cfg: Reference frame provider (root). Defaults to ``SceneEntityCfg("robot")``. + object_cfg: Object to sample. Defaults to ``SceneEntityCfg("object")``. + num_points: Unused at runtime; see note above. + flatten: If ``True``, return a flattened tensor ``(num_envs, 3 * num_points)``. + visualize: If ``True``, draw markers for the points. + + Returns: + Tensor of shape ``(num_envs, num_points, 3)`` or flattened if requested. + """ + ref_pos_w = self.ref_asset.data.root_pos_w.unsqueeze(1).repeat(1, num_points, 1) + ref_quat_w = self.ref_asset.data.root_quat_w.unsqueeze(1).repeat(1, num_points, 1) + + object_pos_w = self.object.data.root_pos_w.unsqueeze(1).repeat(1, num_points, 1) + object_quat_w = self.object.data.root_quat_w.unsqueeze(1).repeat(1, num_points, 1) + # apply rotation + translation + self.points_w = quat_apply(object_quat_w, self.points_local) + object_pos_w + if visualize: + self.visualizer.visualize(translations=self.points_w.view(-1, 3)) + object_point_cloud_pos_b, _ = subtract_frame_transforms(ref_pos_w, ref_quat_w, self.points_w, None) + + return object_point_cloud_pos_b.view(env.num_envs, -1) if flatten else object_point_cloud_pos_b + + +def fingers_contact_force_b( + env: ManagerBasedRLEnv, + contact_sensor_names: list[str], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """base-frame contact forces from listed sensors, concatenated per env. + + Args: + env: The environment. + contact_sensor_names: Names of contact sensors in ``env.scene.sensors`` to read. + + Returns: + Tensor of shape ``(num_envs, 3 * num_sensors)`` with forces stacked horizontally as + ``[fx, fy, fz]`` per sensor. + """ + force_w = [env.scene.sensors[name].data.force_matrix_w.view(env.num_envs, 3) for name in contact_sensor_names] + force_w = torch.stack(force_w, dim=1) + robot: Articulation = env.scene[asset_cfg.name] + forces_b = quat_apply_inverse(robot.data.root_link_quat_w.unsqueeze(1).repeat(1, force_w.shape[1], 1), force_w) + return forces_b diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py new file mode 100644 index 00000000000..9a6170f1e4f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py @@ -0,0 +1,126 @@ +# 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 + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import RigidObject +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import ContactSensor +from isaaclab.utils import math as math_utils +from isaaclab.utils.math import combine_frame_transforms, compute_pose_error + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def action_rate_l2_clamped(env: ManagerBasedRLEnv) -> torch.Tensor: + """Penalize the rate of change of the actions using L2 squared kernel.""" + return torch.sum(torch.square(env.action_manager.action - env.action_manager.prev_action), dim=1).clamp(-1000, 1000) + + +def action_l2_clamped(env: ManagerBasedRLEnv) -> torch.Tensor: + """Penalize the actions using L2 squared kernel.""" + return torch.sum(torch.square(env.action_manager.action), dim=1).clamp(-1000, 1000) + + +def object_ee_distance( + env: ManagerBasedRLEnv, + std: float, + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """Reward reaching the object using a tanh-kernel on end-effector distance. + + The reward is close to 1 when the maximum distance between the object and any end-effector body is small. + """ + asset: RigidObject = env.scene[asset_cfg.name] + object: RigidObject = env.scene[object_cfg.name] + asset_pos = asset.data.body_pos_w[:, asset_cfg.body_ids] + object_pos = object.data.root_pos_w + object_ee_distance = torch.norm(asset_pos - object_pos[:, None, :], dim=-1).max(dim=-1).values + return 1 - torch.tanh(object_ee_distance / std) + + +def contacts(env: ManagerBasedRLEnv, threshold: float) -> torch.Tensor: + """Penalize undesired contacts as the number of violations that are above a threshold.""" + + thumb_contact_sensor: ContactSensor = env.scene.sensors["thumb_link_3_object_s"] + index_contact_sensor: ContactSensor = env.scene.sensors["index_link_3_object_s"] + middle_contact_sensor: ContactSensor = env.scene.sensors["middle_link_3_object_s"] + ring_contact_sensor: ContactSensor = env.scene.sensors["ring_link_3_object_s"] + # check if contact force is above threshold + thumb_contact = thumb_contact_sensor.data.force_matrix_w.view(env.num_envs, 3) + index_contact = index_contact_sensor.data.force_matrix_w.view(env.num_envs, 3) + middle_contact = middle_contact_sensor.data.force_matrix_w.view(env.num_envs, 3) + ring_contact = ring_contact_sensor.data.force_matrix_w.view(env.num_envs, 3) + + thumb_contact_mag = torch.norm(thumb_contact, dim=-1) + index_contact_mag = torch.norm(index_contact, dim=-1) + middle_contact_mag = torch.norm(middle_contact, dim=-1) + ring_contact_mag = torch.norm(ring_contact, dim=-1) + good_contact_cond1 = (thumb_contact_mag > threshold) & ( + (index_contact_mag > threshold) | (middle_contact_mag > threshold) | (ring_contact_mag > threshold) + ) + + return good_contact_cond1 + + +def success_reward( + env: ManagerBasedRLEnv, + command_name: str, + asset_cfg: SceneEntityCfg, + align_asset_cfg: SceneEntityCfg, + pos_std: float, + rot_std: float | None = None, +) -> torch.Tensor: + """Reward success by comparing commanded pose to the object pose using tanh kernels on error.""" + + asset: RigidObject = env.scene[asset_cfg.name] + object: RigidObject = env.scene[align_asset_cfg.name] + command = env.command_manager.get_command(command_name) + des_pos_w, des_quat_w = combine_frame_transforms( + asset.data.root_pos_w, asset.data.root_quat_w, command[:, :3], command[:, 3:7] + ) + pos_err, rot_err = compute_pose_error(des_pos_w, des_quat_w, object.data.root_pos_w, object.data.root_quat_w) + pos_dist = torch.norm(pos_err, dim=1) + if not rot_std: + # square is not necessary but this help to keep the final value between having rot_std or not roughly the same + return (1 - torch.tanh(pos_dist / pos_std)) ** 2 + rot_dist = torch.norm(rot_err, dim=1) + return (1 - torch.tanh(pos_dist / pos_std)) * (1 - torch.tanh(rot_dist / rot_std)) + + +def position_command_error_tanh( + env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg, align_asset_cfg: SceneEntityCfg +) -> torch.Tensor: + """Reward tracking of commanded position using tanh kernel, gated by contact presence.""" + + asset: RigidObject = env.scene[asset_cfg.name] + object: RigidObject = env.scene[align_asset_cfg.name] + command = env.command_manager.get_command(command_name) + # obtain the desired and current positions + des_pos_b = command[:, :3] + des_pos_w, _ = combine_frame_transforms(asset.data.root_pos_w, asset.data.root_quat_w, des_pos_b) + distance = torch.norm(object.data.root_pos_w - des_pos_w, dim=1) + return (1 - torch.tanh(distance / std)) * contacts(env, 1.0).float() + + +def orientation_command_error_tanh( + env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg, align_asset_cfg: SceneEntityCfg +) -> torch.Tensor: + """Reward tracking of commanded orientation using tanh kernel, gated by contact presence.""" + + asset: RigidObject = env.scene[asset_cfg.name] + object: RigidObject = env.scene[align_asset_cfg.name] + command = env.command_manager.get_command(command_name) + # obtain the desired and current orientations + des_quat_b = command[:, 3:7] + des_quat_w = math_utils.quat_mul(asset.data.root_state_w[:, 3:7], des_quat_b) + quat_distance = math_utils.quat_error_magnitude(object.data.root_quat_w, des_quat_w) + + return (1 - torch.tanh(quat_distance / std)) * contacts(env, 1.0).float() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py new file mode 100644 index 00000000000..3ef9cf14b0a --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py @@ -0,0 +1,49 @@ +# 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 + +"""Common functions that can be used to activate certain terminations for the dexsuite task. + +The functions can be passed to the :class:`isaaclab.managers.TerminationTermCfg` object to enable +the termination introduced by the function. +""" + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def out_of_bound( + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg = SceneEntityCfg("object"), + in_bound_range: dict[str, tuple[float, float]] = {}, +) -> torch.Tensor: + """Termination condition for the object falls out of bound. + + Args: + env: The environment. + asset_cfg: The object configuration. Defaults to SceneEntityCfg("object"). + in_bound_range: The range in x, y, z such that the object is considered in range + """ + object: RigidObject = env.scene[asset_cfg.name] + range_list = [in_bound_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] + ranges = torch.tensor(range_list, device=env.device) + + object_pos_local = object.data.root_pos_w - env.scene.env_origins + outside_bounds = ((object_pos_local < ranges[:, 0]) | (object_pos_local > ranges[:, 1])).any(dim=1) + return outside_bounds + + +def abnormal_robot_state(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Terminating environment when violation of velocity limits detects, this usually indicates unstable physics caused + by very bad, or aggressive action""" + robot: Articulation = env.scene[asset_cfg.name] + return (robot.data.joint_vel.abs() > (robot.data.joint_vel_limits * 2)).any(dim=1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/utils.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/utils.py new file mode 100644 index 00000000000..f7b8e9db59b --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/utils.py @@ -0,0 +1,247 @@ +# 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 + +import hashlib +import logging +import numpy as np +import torch +import trimesh +from trimesh.sample import sample_surface + +import isaacsim.core.utils.prims as prim_utils +from pxr import UsdGeom + +from isaaclab.sim.utils import get_all_matching_child_prims + +# ---- module-scope caches ---- +_PRIM_SAMPLE_CACHE: dict[tuple[str, int], np.ndarray] = {} # (prim_hash, num_points) -> (N,3) in root frame +_FINAL_SAMPLE_CACHE: dict[str, np.ndarray] = {} # env_hash -> (num_points,3) in root frame + + +def clear_pointcloud_caches(): + _PRIM_SAMPLE_CACHE.clear() + _FINAL_SAMPLE_CACHE.clear() + + +def sample_object_point_cloud(num_envs: int, num_points: int, prim_path: str, device: str = "cpu") -> torch.Tensor: + """ + Samples point clouds for each environment instance by collecting points + from all matching USD prims under `prim_path`, then downsamples to + exactly `num_points` per env using farthest-point sampling. + + Caching is in-memory within this module: + - per-prim raw samples: _PRIM_SAMPLE_CACHE[(prim_hash, num_points)] + - final downsampled env: _FINAL_SAMPLE_CACHE[env_hash] + + Returns: + torch.Tensor: Shape (num_envs, num_points, 3) on `device`. + """ + points = torch.zeros((num_envs, num_points, 3), dtype=torch.float32, device=device) + xform_cache = UsdGeom.XformCache() + + for i in range(num_envs): + # Resolve prim path + obj_path = prim_path.replace(".*", str(i)) + + # Gather prims + prims = get_all_matching_child_prims( + obj_path, predicate=lambda p: p.GetTypeName() in ("Mesh", "Cube", "Sphere", "Cylinder", "Capsule", "Cone") + ) + if not prims: + raise KeyError(f"No valid prims under {obj_path}") + + object_prim = prim_utils.get_prim_at_path(obj_path) + world_root = xform_cache.GetLocalToWorldTransform(object_prim) + + # hash each child prim by its rel transform + geometry + prim_hashes = [] + for prim in prims: + prim_type = prim.GetTypeName() + hasher = hashlib.sha256() + + rel = world_root.GetInverse() * xform_cache.GetLocalToWorldTransform(prim) # prim -> root + mat_np = np.array([[rel[r][c] for c in range(4)] for r in range(4)], dtype=np.float32) + hasher.update(mat_np.tobytes()) + + if prim_type == "Mesh": + mesh = UsdGeom.Mesh(prim) + verts = np.asarray(mesh.GetPointsAttr().Get(), dtype=np.float32) + hasher.update(verts.tobytes()) + else: + if prim_type == "Cube": + size = UsdGeom.Cube(prim).GetSizeAttr().Get() + hasher.update(np.float32(size).tobytes()) + elif prim_type == "Sphere": + r = UsdGeom.Sphere(prim).GetRadiusAttr().Get() + hasher.update(np.float32(r).tobytes()) + elif prim_type == "Cylinder": + c = UsdGeom.Cylinder(prim) + hasher.update(np.float32(c.GetRadiusAttr().Get()).tobytes()) + hasher.update(np.float32(c.GetHeightAttr().Get()).tobytes()) + elif prim_type == "Capsule": + c = UsdGeom.Capsule(prim) + hasher.update(np.float32(c.GetRadiusAttr().Get()).tobytes()) + hasher.update(np.float32(c.GetHeightAttr().Get()).tobytes()) + elif prim_type == "Cone": + c = UsdGeom.Cone(prim) + hasher.update(np.float32(c.GetRadiusAttr().Get()).tobytes()) + hasher.update(np.float32(c.GetHeightAttr().Get()).tobytes()) + + prim_hashes.append(hasher.hexdigest()) + + # scale on root (default to 1 if missing) + attr = object_prim.GetAttribute("xformOp:scale") + scale_val = attr.Get() if attr else None + if scale_val is None: + base_scale = torch.ones(3, dtype=torch.float32, device=device) + else: + base_scale = torch.tensor(scale_val, dtype=torch.float32, device=device) + + # env-level cache key (includes num_points) + env_key = "_".join(sorted(prim_hashes)) + f"_{num_points}" + env_hash = hashlib.sha256(env_key.encode()).hexdigest() + + # load from env-level in-memory cache + if env_hash in _FINAL_SAMPLE_CACHE: + arr = _FINAL_SAMPLE_CACHE[env_hash] # (num_points,3) in root frame + points[i] = torch.from_numpy(arr).to(device) * base_scale.unsqueeze(0) + continue + + # otherwise build per-prim samples (with per-prim cache) + all_samples_np: list[np.ndarray] = [] + for prim, ph in zip(prims, prim_hashes): + key = (ph, num_points) + if key in _PRIM_SAMPLE_CACHE: + samples = _PRIM_SAMPLE_CACHE[key] + else: + prim_type = prim.GetTypeName() + if prim_type == "Mesh": + mesh = UsdGeom.Mesh(prim) + verts = np.asarray(mesh.GetPointsAttr().Get(), dtype=np.float32) + faces = _triangulate_faces(prim) + mesh_tm = trimesh.Trimesh(vertices=verts, faces=faces, process=False) + else: + mesh_tm = create_primitive_mesh(prim) + + face_weights = mesh_tm.area_faces + samples_np, _ = sample_surface(mesh_tm, num_points * 2, face_weight=face_weights) + + # FPS to num_points on chosen device + tensor_pts = torch.from_numpy(samples_np.astype(np.float32)).to(device) + prim_idxs = farthest_point_sampling(tensor_pts, num_points) + local_pts = tensor_pts[prim_idxs] + + # prim -> root transform + rel = xform_cache.GetLocalToWorldTransform(prim) * world_root.GetInverse() + mat_np = np.array([[rel[r][c] for c in range(4)] for r in range(4)], dtype=np.float32) + mat_t = torch.from_numpy(mat_np).to(device) + + ones = torch.ones((num_points, 1), device=device) + pts_h = torch.cat([local_pts, ones], dim=1) + root_h = pts_h @ mat_t + samples = root_h[:, :3].detach().cpu().numpy() + + if prim_type == "Cone": + samples[:, 2] -= UsdGeom.Cone(prim).GetHeightAttr().Get() / 2 + + _PRIM_SAMPLE_CACHE[key] = samples # cache in root frame @ num_points + + all_samples_np.append(samples) + + # combine & env-level FPS (if needed) + if len(all_samples_np) == 1: + samples_final = torch.from_numpy(all_samples_np[0]).to(device) + else: + combined = torch.from_numpy(np.concatenate(all_samples_np, axis=0)).to(device) + idxs = farthest_point_sampling(combined, num_points) + samples_final = combined[idxs] + + # store env-level cache in root frame (CPU) + _FINAL_SAMPLE_CACHE[env_hash] = samples_final.detach().cpu().numpy() + + # apply root scale and write out + points[i] = samples_final * base_scale.unsqueeze(0) + + return points + + +def _triangulate_faces(prim) -> np.ndarray: + """Convert a USD Mesh prim into triangulated face indices (N, 3).""" + mesh = UsdGeom.Mesh(prim) + counts = mesh.GetFaceVertexCountsAttr().Get() + indices = mesh.GetFaceVertexIndicesAttr().Get() + faces = [] + it = iter(indices) + for cnt in counts: + poly = [next(it) for _ in range(cnt)] + for k in range(1, cnt - 1): + faces.append([poly[0], poly[k], poly[k + 1]]) + return np.asarray(faces, dtype=np.int64) + + +def create_primitive_mesh(prim) -> trimesh.Trimesh: + """Create a trimesh mesh from a USD primitive (Cube, Sphere, Cylinder, etc.).""" + prim_type = prim.GetTypeName() + if prim_type == "Cube": + size = UsdGeom.Cube(prim).GetSizeAttr().Get() + return trimesh.creation.box(extents=(size, size, size)) + elif prim_type == "Sphere": + r = UsdGeom.Sphere(prim).GetRadiusAttr().Get() + return trimesh.creation.icosphere(subdivisions=3, radius=r) + elif prim_type == "Cylinder": + c = UsdGeom.Cylinder(prim) + return trimesh.creation.cylinder(radius=c.GetRadiusAttr().Get(), height=c.GetHeightAttr().Get()) + elif prim_type == "Capsule": + c = UsdGeom.Capsule(prim) + return trimesh.creation.capsule(radius=c.GetRadiusAttr().Get(), height=c.GetHeightAttr().Get()) + elif prim_type == "Cone": # Cone + c = UsdGeom.Cone(prim) + return trimesh.creation.cone(radius=c.GetRadiusAttr().Get(), height=c.GetHeightAttr().Get()) + else: + raise KeyError(f"{prim_type} is not a valid primitive mesh type") + + +def farthest_point_sampling( + points: torch.Tensor, n_samples: int, memory_threashold=2 * 1024**3 +) -> torch.Tensor: # 2 GiB + """ + Farthest Point Sampling (FPS) for point sets. + + Selects `n_samples` points such that each new point is farthest from the + already chosen ones. Uses a full pairwise distance matrix if memory allows, + otherwise falls back to an iterative version. + + Args: + points (torch.Tensor): Input points of shape (N, D). + n_samples (int): Number of samples to select. + memory_threashold (int): Max allowed bytes for distance matrix. Default 2 GiB. + + Returns: + torch.Tensor: Indices of sampled points (n_samples,). + """ + device = points.device + N = points.shape[0] + elem_size = points.element_size() + bytes_needed = N * N * elem_size + if bytes_needed <= memory_threashold: + dist_mat = torch.cdist(points, points) + sampled_idx = torch.zeros(n_samples, dtype=torch.long, device=device) + min_dists = torch.full((N,), float("inf"), device=device) + farthest = torch.randint(0, N, (1,), device=device) + for j in range(n_samples): + sampled_idx[j] = farthest + min_dists = torch.minimum(min_dists, dist_mat[farthest].view(-1)) + farthest = torch.argmax(min_dists) + return sampled_idx + logging.warning(f"FPS fallback to iterative (needed {bytes_needed} > {memory_threashold})") + sampled_idx = torch.zeros(n_samples, dtype=torch.long, device=device) + distances = torch.full((N,), float("inf"), device=device) + farthest = torch.randint(0, N, (1,), device=device) + for j in range(n_samples): + sampled_idx[j] = farthest + dist = torch.norm(points - points[farthest], dim=1) + distances = torch.minimum(distances, dist) + farthest = torch.argmax(distances) + return sampled_idx