diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index d7b4042e53..1b540f9f42 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -64,6 +64,7 @@ Guidelines for modifications: * Nuralem Abizov * Oyindamola Omotuyi * Özhan Özen +* Peter Du * Qian Wan * Qinxi Yu * René Zurbrügg diff --git a/docs/source/_static/tasks/manipulation/franka_stack.jpg b/docs/source/_static/tasks/manipulation/franka_stack.jpg new file mode 100644 index 0000000000..1da36113ec Binary files /dev/null and b/docs/source/_static/tasks/manipulation/franka_stack.jpg differ diff --git a/docs/source/overview/environments.rst b/docs/source/overview/environments.rst index 531588a8de..5d08c22b6e 100644 --- a/docs/source/overview/environments.rst +++ b/docs/source/overview/environments.rst @@ -109,6 +109,8 @@ for the reach environment: +--------------------+-------------------------+-----------------------------------------------------------------------------+ | |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 | + +--------------------+-------------------------+-----------------------------------------------------------------------------+ | |cabi-franka| | |cabi-franka-link| | Grasp the handle of a cabinet's drawer and open it with the Franka robot | | | | | | | |franka-direct-link| | | @@ -132,6 +134,7 @@ for the reach environment: .. |cabi-franka| image:: ../_static/tasks/manipulation/franka_open_drawer.jpg .. |cube-allegro| image:: ../_static/tasks/manipulation/allegro_cube.jpg .. |cube-shadow| image:: ../_static/tasks/manipulation/shadow_cube.jpg +.. |stack-franka| image:: ../_static/tasks/manipulation/franka_stack.jpg .. |reach-franka-link| replace:: `Isaac-Reach-Franka-v0 `__ .. |reach-ur10-link| replace:: `Isaac-Reach-UR10-v0 `__ @@ -142,6 +145,7 @@ for the reach environment: .. |franka-direct-link| replace:: `Isaac-Franka-Cabinet-Direct-v0 `__ .. |cube-allegro-link| replace:: `Isaac-Repose-Cube-Allegro-v0 `__ .. |allegro-direct-link| replace:: `Isaac-Repose-Cube-Allegro-Direct-v0 `__ +.. |stack-cube-link| replace:: `Isaac-Stack-Cube-Franka-v0 `__ .. |cube-shadow-link| replace:: `Isaac-Repose-Cube-Shadow-Direct-v0 `__ .. |cube-shadow-ff-link| replace:: `Isaac-Repose-Cube-Shadow-OpenAI-FF-Direct-v0 `__ diff --git a/source/extensions/omni.isaac.lab_tasks/config/extension.toml b/source/extensions/omni.isaac.lab_tasks/config/extension.toml index c4beba4ecd..ba976faf92 100644 --- a/source/extensions/omni.isaac.lab_tasks/config/extension.toml +++ b/source/extensions/omni.isaac.lab_tasks/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.10.13" +version = "0.10.14" # Description title = "Isaac Lab Environments" diff --git a/source/extensions/omni.isaac.lab_tasks/docs/CHANGELOG.rst b/source/extensions/omni.isaac.lab_tasks/docs/CHANGELOG.rst index c7912dcc6d..068db9b64a 100644 --- a/source/extensions/omni.isaac.lab_tasks/docs/CHANGELOG.rst +++ b/source/extensions/omni.isaac.lab_tasks/docs/CHANGELOG.rst @@ -1,6 +1,15 @@ Changelog --------- +0.10.14 (2024-12-03) +~~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added ``Isaac-Stack-Cube-Franka-IK-Rel-v0`` and ``Isaac-Stack-Cube-Instance-Randomize-Franka-IK-Rel-v0`` environments + as manager-based RL envs that implement a three cube stacking task. + 0.10.13 (2024-10-30) ~~~~~~~~~~~~~~~~~~~~ diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/__init__.py new file mode 100644 index 0000000000..a6be077a39 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configurations for the object stack 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/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/__init__.py new file mode 100644 index 0000000000..a6be077a39 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configurations for the object stack 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/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/__init__.py new file mode 100644 index 0000000000..5a564fc996 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/__init__.py @@ -0,0 +1,64 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +import gymnasium as gym +import os + +from . import ( + agents, + stack_ik_rel_env_cfg, + stack_ik_rel_instance_randomize_env_cfg, + stack_joint_pos_env_cfg, + stack_joint_pos_instance_randomize_env_cfg, +) + +## +# Register Gym environments. +## + +## +# Joint Position Control +## + +gym.register( + id="Isaac-Stack-Cube-Franka-v0", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": stack_joint_pos_env_cfg.FrankaCubeStackEnvCfg, + }, + disable_env_checker=True, +) + +gym.register( + id="Isaac-Stack-Cube-Instance-Randomize-Franka-v0", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": stack_joint_pos_instance_randomize_env_cfg.FrankaCubeStackInstanceRandomizeEnvCfg, + }, + disable_env_checker=True, +) + + +## +# Inverse Kinematics - Relative Pose Control +## + +gym.register( + id="Isaac-Stack-Cube-Franka-IK-Rel-v0", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": stack_ik_rel_env_cfg.FrankaCubeStackEnvCfg, + "robomimic_bc_cfg_entry_point": os.path.join(agents.__path__[0], "robomimic/bc_rnn_low_dim.json"), + }, + disable_env_checker=True, +) + +gym.register( + id="Isaac-Stack-Cube-Instance-Randomize-Franka-IK-Rel-v0", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": stack_ik_rel_instance_randomize_env_cfg.FrankaCubeStackInstanceRandomizeEnvCfg, + }, + disable_env_checker=True, +) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/agents/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/agents/__init__.py new file mode 100644 index 0000000000..c3ee657052 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/agents/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/agents/robomimic/bc_rnn_low_dim.json b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/agents/robomimic/bc_rnn_low_dim.json new file mode 100644 index 0000000000..24d807ac34 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/agents/robomimic/bc_rnn_low_dim.json @@ -0,0 +1,101 @@ +{ + "algo_name": "bc", + "experiment": { + "name": "bc", + "validate": false, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 100, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "env": null, + "additional_envs": null, + "render": false, + "render_video": false, + "rollout": { + "enabled": false + } + }, + "train": { + "data": null, + "num_data_workers": 4, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": null, + "hdf5_validation_filter_key": null, + "seq_length": 10, + "dataset_keys": [ + "actions" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 101 + }, + "algo": { + "optim_params": { + "policy": { + "optimizer_type": "adam", + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [], + "scheduler_type": "multistep" + }, + "regularization": { + "L2": 0.0 + } + } + }, + "loss": { + "l2_weight": 1.0, + "l1_weight": 0.0, + "cos_weight": 0.0 + }, + "actor_layer_dims": [], + "gmm": { + "enabled": true, + "num_modes": 5, + "min_std": 0.0001, + "std_activation": "softplus", + "low_noise_eval": true + }, + "rnn": { + "enabled": true, + "horizon": 10, + "hidden_dim": 400, + "rnn_type": "LSTM", + "num_layers": 2, + "open_loop": false, + "kwargs": { + "bidirectional": false + } + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "eef_pos", + "eef_quat", + "gripper_pos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + } + } + } +} diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_env_cfg.py new file mode 100644 index 0000000000..f95f1c61dc --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_env_cfg.py @@ -0,0 +1,36 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from omni.isaac.lab.controllers.differential_ik_cfg import DifferentialIKControllerCfg +from omni.isaac.lab.envs.mdp.actions.actions_cfg import DifferentialInverseKinematicsActionCfg +from omni.isaac.lab.utils import configclass + +from . import stack_joint_pos_env_cfg + +## +# Pre-defined configs +## +from omni.isaac.lab_assets.franka import FRANKA_PANDA_HIGH_PD_CFG # isort: skip + + +@configclass +class FrankaCubeStackEnvCfg(stack_joint_pos_env_cfg.FrankaCubeStackEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Set Franka as robot + # We switch here to a stiffer PD controller for IK tracking to be better. + self.scene.robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Set actions for the specific robot type (franka) + self.actions.arm_action = DifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["panda_joint.*"], + body_name="panda_hand", + controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=True, ik_method="dls"), + scale=0.5, + body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.107]), + ) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_instance_randomize_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_instance_randomize_env_cfg.py new file mode 100644 index 0000000000..91d1d02730 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_instance_randomize_env_cfg.py @@ -0,0 +1,38 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from omni.isaac.lab.controllers.differential_ik_cfg import DifferentialIKControllerCfg +from omni.isaac.lab.envs.mdp.actions.actions_cfg import DifferentialInverseKinematicsActionCfg +from omni.isaac.lab.utils import configclass + +from . import stack_joint_pos_instance_randomize_env_cfg + +## +# Pre-defined configs +## +from omni.isaac.lab_assets.franka import FRANKA_PANDA_HIGH_PD_CFG # isort: skip + + +@configclass +class FrankaCubeStackInstanceRandomizeEnvCfg( + stack_joint_pos_instance_randomize_env_cfg.FrankaCubeStackInstanceRandomizeEnvCfg +): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Set Franka as robot + # We switch here to a stiffer PD controller for IK tracking to be better. + self.scene.robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Set actions for the specific robot type (franka) + self.actions.arm_action = DifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["panda_joint.*"], + body_name="panda_hand", + controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=True, ik_method="dls"), + scale=0.5, + body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.107]), + ) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_env_cfg.py new file mode 100644 index 0000000000..3b7dc1dc2f --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_env_cfg.py @@ -0,0 +1,180 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import RigidObjectCfg +from omni.isaac.lab.managers import EventTermCfg as EventTerm +from omni.isaac.lab.managers import SceneEntityCfg +from omni.isaac.lab.sensors import CameraCfg, FrameTransformerCfg +from omni.isaac.lab.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg +from omni.isaac.lab.sim.schemas.schemas_cfg import RigidBodyPropertiesCfg +from omni.isaac.lab.sim.spawners.from_files.from_files_cfg import UsdFileCfg +from omni.isaac.lab.utils import configclass +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR + +from omni.isaac.lab_tasks.manager_based.manipulation.stack import mdp +from omni.isaac.lab_tasks.manager_based.manipulation.stack.mdp import franka_stack_events +from omni.isaac.lab_tasks.manager_based.manipulation.stack.stack_env_cfg import StackEnvCfg + +## +# Pre-defined configs +## +from omni.isaac.lab.markers.config import FRAME_MARKER_CFG # isort: skip +from omni.isaac.lab_assets.franka import FRANKA_PANDA_CFG # isort: skip + + +@configclass +class EventCfg: + """Configuration for events.""" + + init_franka_arm_pose = EventTerm( + func=franka_stack_events.set_default_joint_pose, + mode="startup", + params={ + "default_pose": [0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400, 0.0400], + }, + ) + + randomize_franka_joint_state = EventTerm( + func=franka_stack_events.randomize_joint_by_gaussian_offset, + mode="reset", + params={ + "mean": 0.0, + "std": 0.02, + "asset_cfg": SceneEntityCfg("robot"), + }, + ) + + randomize_cube_positions = EventTerm( + func=franka_stack_events.randomize_object_pose, + mode="reset", + params={ + "pose_range": {"x": (0.4, 0.6), "y": (-0.10, 0.10), "z": (0.0203, 0.0203), "yaw": (-1.0, 1, 0)}, + "min_separation": 0.1, + "asset_cfgs": [SceneEntityCfg("cube_1"), SceneEntityCfg("cube_2"), SceneEntityCfg("cube_3")], + }, + ) + + +@configclass +class FrankaCubeStackEnvCfg(StackEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Set events + self.events = EventCfg() + + # Set Franka as robot + self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Set actions for the specific robot type (franka) + self.actions.arm_action = mdp.JointPositionActionCfg( + asset_name="robot", joint_names=["panda_joint.*"], scale=0.5, use_default_offset=True + ) + self.actions.gripper_action = mdp.BinaryJointPositionActionCfg( + asset_name="robot", + joint_names=["panda_finger.*"], + open_command_expr={"panda_finger_.*": 0.04}, + close_command_expr={"panda_finger_.*": 0.0}, + ) + + # Rigid body properties of each cube + cube_properties = RigidBodyPropertiesCfg( + solver_position_iteration_count=16, + solver_velocity_iteration_count=1, + max_angular_velocity=1000.0, + max_linear_velocity=1000.0, + max_depenetration_velocity=5.0, + disable_gravity=False, + ) + + # Set each stacking cube deterministically + self.scene.cube_1 = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_1", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.4, 0.0, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/blue_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ) + self.scene.cube_2 = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_2", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.55, 0.05, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/red_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ) + self.scene.cube_3 = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_3", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.60, -0.1, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/green_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ) + + # Set wrist camera + self.scene.wrist_cam = CameraCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_hand/wrist_cam", + update_period=0.0333, + height=84, + width=84, + data_types=["rgb", "distance_to_image_plane"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 1.0e5) + ), + offset=CameraCfg.OffsetCfg(pos=(0.025, 0.0, 0.0), rot=(0.707, 0.0, 0.0, 0.707), convention="ros"), + ) + + # Set table view camera + self.scene.table_cam = CameraCfg( + prim_path="{ENV_REGEX_NS}/table_cam", + update_period=0.0333, + height=84, + width=84, + data_types=["rgb", "distance_to_image_plane"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 1.0e5) + ), + offset=CameraCfg.OffsetCfg(pos=(1.0, 0.0, 0.33), rot=(-0.3799, 0.5963, 0.5963, -0.3799), convention="ros"), + ) + + # Listens to the required transforms + marker_cfg = FRAME_MARKER_CFG.copy() + marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1) + marker_cfg.prim_path = "/Visuals/FrameTransformer" + self.scene.ee_frame = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_link0", + debug_vis=False, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_hand", + name="end_effector", + offset=OffsetCfg( + pos=[0.0, 0.0, 0.1034], + ), + ), + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_rightfinger", + name="tool_rightfinger", + offset=OffsetCfg( + pos=(0.0, 0.0, 0.046), + ), + ), + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_leftfinger", + name="tool_leftfinger", + offset=OffsetCfg( + pos=(0.0, 0.0, 0.046), + ), + ), + ], + ) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_instance_randomize_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_instance_randomize_env_cfg.py new file mode 100644 index 0000000000..f9038c8ae9 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_instance_randomize_env_cfg.py @@ -0,0 +1,224 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import RigidObjectCfg, RigidObjectCollectionCfg +from omni.isaac.lab.managers import EventTermCfg as EventTerm +from omni.isaac.lab.managers import SceneEntityCfg +from omni.isaac.lab.sensors import CameraCfg, FrameTransformerCfg +from omni.isaac.lab.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg +from omni.isaac.lab.sim.schemas.schemas_cfg import RigidBodyPropertiesCfg +from omni.isaac.lab.sim.spawners.from_files.from_files_cfg import UsdFileCfg +from omni.isaac.lab.utils import configclass +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR + +from omni.isaac.lab_tasks.manager_based.manipulation.stack import mdp +from omni.isaac.lab_tasks.manager_based.manipulation.stack.mdp import franka_stack_events +from omni.isaac.lab_tasks.manager_based.manipulation.stack.stack_instance_randomize_env_cfg import ( + StackInstanceRandomizeEnvCfg, +) + +## +# Pre-defined configs +## +from omni.isaac.lab.markers.config import FRAME_MARKER_CFG # isort: skip +from omni.isaac.lab_assets.franka import FRANKA_PANDA_CFG # isort: skip + + +@configclass +class EventCfg: + """Configuration for events.""" + + init_franka_arm_pose = EventTerm( + func=franka_stack_events.set_default_joint_pose, + mode="startup", + params={ + "default_pose": [0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400, 0.0400], + }, + ) + + randomize_franka_joint_state = EventTerm( + func=franka_stack_events.randomize_joint_by_gaussian_offset, + mode="reset", + params={ + "mean": 0.0, + "std": 0.02, + "asset_cfg": SceneEntityCfg("robot"), + }, + ) + + randomize_cubes_in_focus = EventTerm( + func=franka_stack_events.randomize_rigid_objects_in_focus, + mode="reset", + params={ + "asset_cfgs": [SceneEntityCfg("cube_1"), SceneEntityCfg("cube_2"), SceneEntityCfg("cube_3")], + "out_focus_state": torch.tensor([10.0, 10.0, 10.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + "pose_range": {"x": (0.4, 0.6), "y": (-0.10, 0.10), "z": (0.0203, 0.0203), "yaw": (-1.0, 1, 0)}, + "min_separation": 0.1, + }, + ) + + +@configclass +class FrankaCubeStackInstanceRandomizeEnvCfg(StackInstanceRandomizeEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Set events + self.events = EventCfg() + + # Set Franka as robot + self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Set actions for the specific robot type (franka) + self.actions.arm_action = mdp.JointPositionActionCfg( + asset_name="robot", joint_names=["panda_joint.*"], scale=0.5, use_default_offset=True + ) + self.actions.gripper_action = mdp.BinaryJointPositionActionCfg( + asset_name="robot", + joint_names=["panda_finger.*"], + open_command_expr={"panda_finger_.*": 0.04}, + close_command_expr={"panda_finger_.*": 0.0}, + ) + + # Rigid body properties of each cube + cube_properties = RigidBodyPropertiesCfg( + solver_position_iteration_count=16, + solver_velocity_iteration_count=1, + max_angular_velocity=1000.0, + max_linear_velocity=1000.0, + max_depenetration_velocity=5.0, + disable_gravity=False, + ) + + # Set each stacking cube to be a collection of rigid objects + cube_1_config_dict = { + "blue_cube": RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_1_Blue", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.4, 0.0, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/blue_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ), + "red_cube": RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_1_Red", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.4, 0.0, 0.0403], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/red_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ), + } + + cube_2_config_dict = { + "red_cube": RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_2_Red", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.55, 0.05, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/red_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ), + "yellow_cube": RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_2_Yellow", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.55, 0.05, 0.0403], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/yellow_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ), + } + + cube_3_config_dict = { + "yellow_cube": RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_3_Yellow", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.60, -0.1, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/yellow_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ), + "green_cube": RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_2_Green", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.60, -0.1, 0.0403], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/green_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ), + } + + self.scene.cube_1 = RigidObjectCollectionCfg(rigid_objects=cube_1_config_dict) + self.scene.cube_2 = RigidObjectCollectionCfg(rigid_objects=cube_2_config_dict) + self.scene.cube_3 = RigidObjectCollectionCfg(rigid_objects=cube_3_config_dict) + + # Set wrist camera + self.scene.wrist_cam = CameraCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_hand/wrist_cam", + update_period=0.0333, + height=84, + width=84, + data_types=["rgb", "distance_to_image_plane"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 1.0e5) + ), + offset=CameraCfg.OffsetCfg(pos=(0.025, 0.0, 0.0), rot=(0.707, 0.0, 0.0, 0.707), convention="ros"), + ) + + # Set table view camera + self.scene.table_cam = CameraCfg( + prim_path="{ENV_REGEX_NS}/table_cam", + update_period=0.0333, + height=84, + width=84, + data_types=["rgb", "distance_to_image_plane"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 1.0e5) + ), + offset=CameraCfg.OffsetCfg(pos=(1.0, 0.0, 0.33), rot=(-0.3799, 0.5963, 0.5963, -0.3799), convention="ros"), + ) + + # Listens to the required transforms + marker_cfg = FRAME_MARKER_CFG.copy() + marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1) + marker_cfg.prim_path = "/Visuals/FrameTransformer" + self.scene.ee_frame = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_link0", + debug_vis=False, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_hand", + name="end_effector", + offset=OffsetCfg( + pos=[0.0, 0.0, 0.1034], + ), + ), + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_rightfinger", + name="tool_rightfinger", + offset=OffsetCfg( + pos=(0.0, 0.0, 0.046), + ), + ), + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_leftfinger", + name="tool_leftfinger", + offset=OffsetCfg( + pos=(0.0, 0.0, 0.046), + ), + ), + ], + ) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/__init__.py new file mode 100644 index 0000000000..8a6d3a1a71 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""This sub-module contains the functions that are specific to the lift environments.""" + +from omni.isaac.lab.envs.mdp import * # noqa: F401, F403 + +from .observations import * # noqa: F401, F403 +from .terminations import * # noqa: F401, F403 diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py new file mode 100644 index 0000000000..f741470a34 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py @@ -0,0 +1,186 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +from __future__ import annotations + +import math +import random +import torch +from typing import TYPE_CHECKING + +import omni.isaac.lab.utils.math as math_utils +from omni.isaac.lab.assets import Articulation, AssetBase +from omni.isaac.lab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from omni.isaac.lab.envs import ManagerBasedEnv + + +def set_default_joint_pose( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + default_pose: torch.Tensor, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +): + # Set the default pose for robots in all envs + asset = env.scene[asset_cfg.name] + asset.data.default_joint_pos = torch.tensor(default_pose, device=env.device).repeat(env.num_envs, 1) + + +def randomize_joint_by_gaussian_offset( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + mean: float, + std: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +): + asset: Articulation = env.scene[asset_cfg.name] + + # Add gaussian noise to joint states + joint_pos = asset.data.default_joint_pos[env_ids].clone() + joint_vel = asset.data.default_joint_vel[env_ids].clone() + joint_pos += math_utils.sample_gaussian(mean, std, joint_pos.shape, joint_pos.device) + + # Clamp joint pos to limits + joint_pos_limits = asset.data.soft_joint_pos_limits[env_ids] + joint_pos = joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) + + # Don't noise the gripper poses + joint_pos[:, -2:] = asset.data.default_joint_pos[env_ids, -2:] + + # Set into the physics simulation + asset.set_joint_position_target(joint_pos, env_ids=env_ids) + asset.set_joint_velocity_target(joint_vel, env_ids=env_ids) + asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids) + + +def randomize_scene_lighting_domelight( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + intensity_range: tuple[float, float], + asset_cfg: SceneEntityCfg = SceneEntityCfg("light"), +): + asset: AssetBase = env.scene[asset_cfg.name] + light_prim = asset.prims[0] + + # Sample new light intensity + new_intensity = random.uniform(intensity_range[0], intensity_range[1]) + + # Set light intensity to light prim + intensity_attr = light_prim.GetAttribute("inputs:intensity") + intensity_attr.Set(new_intensity) + + +def sample_object_poses( + num_objects: int, + min_separation: float = 0.0, + pose_range: dict[str, tuple[float, float]] = {}, + max_sample_tries: int = 5000, +): + range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] + pose_list = [] + + for i in range(num_objects): + for j in range(max_sample_tries): + sample = [random.uniform(range[0], range[1]) for range in range_list] + + # Accept pose if it is the first one, or if reached max num tries + if len(pose_list) == 0 or j == max_sample_tries - 1: + pose_list.append(sample) + break + + # Check if pose of object is sufficiently far away from all other objects + separation_check = [math.dist(sample[:3], pose[:3]) > min_separation for pose in pose_list] + if False not in separation_check: + pose_list.append(sample) + break + + return pose_list + + +def randomize_object_pose( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + asset_cfgs: list[SceneEntityCfg], + min_separation: float = 0.0, + pose_range: dict[str, tuple[float, float]] = {}, + max_sample_tries: int = 5000, +): + if env_ids is None: + return + + # Randomize poses in each environment independently + for cur_env in env_ids.tolist(): + pose_list = sample_object_poses( + num_objects=len(asset_cfgs), + min_separation=min_separation, + pose_range=pose_range, + max_sample_tries=max_sample_tries, + ) + + # Randomize pose for each object + for i in range(len(asset_cfgs)): + asset_cfg = asset_cfgs[i] + asset = env.scene[asset_cfg.name] + + # Write pose to simulation + pose_tensor = torch.tensor([pose_list[i]], device=env.device) + positions = pose_tensor[:, 0:3] + env.scene.env_origins[cur_env, 0:3] + orientations = math_utils.quat_from_euler_xyz(pose_tensor[:, 3], pose_tensor[:, 4], pose_tensor[:, 5]) + asset.write_root_pose_to_sim( + torch.cat([positions, orientations], dim=-1), env_ids=torch.tensor([cur_env], device=env.device) + ) + asset.write_root_velocity_to_sim( + torch.zeros(1, 6, device=env.device), env_ids=torch.tensor([cur_env], device=env.device) + ) + + +def randomize_rigid_objects_in_focus( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + asset_cfgs: list[SceneEntityCfg], + out_focus_state: torch.Tensor, + min_separation: float = 0.0, + pose_range: dict[str, tuple[float, float]] = {}, + max_sample_tries: int = 5000, +): + if env_ids is None: + return + + # List of rigid objects in focus for each env (dim = [num_envs, num_rigid_objects]) + env.rigid_objects_in_focus = [] + + for cur_env in env_ids.tolist(): + # Sample in focus object poses + pose_list = sample_object_poses( + num_objects=len(asset_cfgs), + min_separation=min_separation, + pose_range=pose_range, + max_sample_tries=max_sample_tries, + ) + + selected_ids = [] + for asset_idx in range(len(asset_cfgs)): + asset_cfg = asset_cfgs[asset_idx] + asset = env.scene[asset_cfg.name] + + # Randomly select an object to bring into focus + object_id = random.randint(0, asset.num_objects - 1) + selected_ids.append(object_id) + + # Create object state tensor + object_states = torch.stack([out_focus_state] * asset.num_objects).to(device=env.device) + pose_tensor = torch.tensor([pose_list[asset_idx]], device=env.device) + positions = pose_tensor[:, 0:3] + env.scene.env_origins[cur_env, 0:3] + orientations = math_utils.quat_from_euler_xyz(pose_tensor[:, 3], pose_tensor[:, 4], pose_tensor[:, 5]) + object_states[object_id, 0:3] = positions + object_states[object_id, 3:7] = orientations + + asset.write_object_state_to_sim( + object_state=object_states, env_ids=torch.tensor([cur_env], device=env.device) + ) + + env.rigid_objects_in_focus.append(selected_ids) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/observations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/observations.py new file mode 100644 index 0000000000..54831e9d75 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/observations.py @@ -0,0 +1,326 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from omni.isaac.lab.assets import Articulation, RigidObject, RigidObjectCollection +from omni.isaac.lab.managers import SceneEntityCfg +from omni.isaac.lab.sensors import FrameTransformer + +if TYPE_CHECKING: + from omni.isaac.lab.envs import ManagerBasedRLEnv + + +def cube_positions_in_world_frame( + env: ManagerBasedRLEnv, + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), +) -> torch.Tensor: + """The position of the cubes in the world frame.""" + cube_1: RigidObject = env.scene[cube_1_cfg.name] + cube_2: RigidObject = env.scene[cube_2_cfg.name] + cube_3: RigidObject = env.scene[cube_3_cfg.name] + + return torch.cat((cube_1.data.root_pos_w, cube_2.data.root_pos_w, cube_3.data.root_pos_w), dim=1) + + +def instance_randomize_cube_positions_in_world_frame( + env: ManagerBasedRLEnv, + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), +) -> torch.Tensor: + """The position of the cubes in the world frame.""" + if not hasattr(env, "rigid_objects_in_focus"): + return torch.full((env.num_envs, 9), fill_value=-1) + + cube_1: RigidObjectCollection = env.scene[cube_1_cfg.name] + cube_2: RigidObjectCollection = env.scene[cube_2_cfg.name] + cube_3: RigidObjectCollection = env.scene[cube_3_cfg.name] + + cube_1_pos_w = [] + cube_2_pos_w = [] + cube_3_pos_w = [] + for env_id in range(env.num_envs): + cube_1_pos_w.append(cube_1.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][0], :3]) + cube_2_pos_w.append(cube_2.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][1], :3]) + cube_3_pos_w.append(cube_3.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][2], :3]) + cube_1_pos_w = torch.stack(cube_1_pos_w) + cube_2_pos_w = torch.stack(cube_2_pos_w) + cube_3_pos_w = torch.stack(cube_3_pos_w) + + return torch.cat((cube_1_pos_w, cube_2_pos_w, cube_3_pos_w), dim=1) + + +def cube_orientations_in_world_frame( + env: ManagerBasedRLEnv, + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), +): + """The orientation of the cubes in the world frame.""" + cube_1: RigidObject = env.scene[cube_1_cfg.name] + cube_2: RigidObject = env.scene[cube_2_cfg.name] + cube_3: RigidObject = env.scene[cube_3_cfg.name] + + return torch.cat((cube_1.data.root_quat_w, cube_2.data.root_quat_w, cube_3.data.root_quat_w), dim=1) + + +def instance_randomize_cube_orientations_in_world_frame( + env: ManagerBasedRLEnv, + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), +) -> torch.Tensor: + """The orientation of the cubes in the world frame.""" + if not hasattr(env, "rigid_objects_in_focus"): + return torch.full((env.num_envs, 9), fill_value=-1) + + cube_1: RigidObjectCollection = env.scene[cube_1_cfg.name] + cube_2: RigidObjectCollection = env.scene[cube_2_cfg.name] + cube_3: RigidObjectCollection = env.scene[cube_3_cfg.name] + + cube_1_quat_w = [] + cube_2_quat_w = [] + cube_3_quat_w = [] + for env_id in range(env.num_envs): + cube_1_quat_w.append(cube_1.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][0], :4]) + cube_2_quat_w.append(cube_2.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][1], :4]) + cube_3_quat_w.append(cube_3.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][2], :4]) + cube_1_quat_w = torch.stack(cube_1_quat_w) + cube_2_quat_w = torch.stack(cube_2_quat_w) + cube_3_quat_w = torch.stack(cube_3_quat_w) + + return torch.cat((cube_1_quat_w, cube_2_quat_w, cube_3_quat_w), dim=1) + + +def object_obs( + env: ManagerBasedRLEnv, + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), + ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"), +): + """ + Object observations (in world frame): + cube_1 pos, + cube_1 quat, + cube_2 pos, + cube_2 quat, + cube_3 pos, + cube_3 quat, + gripper to cube_1, + gripper to cube_2, + gripper to cube_3, + cube_1 to cube_2, + cube_2 to cube_3, + cube_1 to cube_3, + """ + cube_1: RigidObject = env.scene[cube_1_cfg.name] + cube_2: RigidObject = env.scene[cube_2_cfg.name] + cube_3: RigidObject = env.scene[cube_3_cfg.name] + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + + cube_1_pos_w = cube_1.data.root_pos_w + cube_1_quat_w = cube_1.data.root_quat_w + + cube_2_pos_w = cube_2.data.root_pos_w + cube_2_quat_w = cube_2.data.root_quat_w + + cube_3_pos_w = cube_3.data.root_pos_w + cube_3_quat_w = cube_3.data.root_quat_w + + ee_pos_w = ee_frame.data.target_pos_w[:, 0, :] + gripper_to_cube_1 = cube_1_pos_w - ee_pos_w + gripper_to_cube_2 = cube_2_pos_w - ee_pos_w + gripper_to_cube_3 = cube_3_pos_w - ee_pos_w + + cube_1_to_2 = cube_1_pos_w - cube_2_pos_w + cube_2_to_3 = cube_2_pos_w - cube_3_pos_w + cube_1_to_3 = cube_1_pos_w - cube_3_pos_w + + return torch.cat( + ( + cube_1_pos_w, + cube_1_quat_w, + cube_2_pos_w, + cube_2_quat_w, + cube_3_pos_w, + cube_3_quat_w, + gripper_to_cube_1, + gripper_to_cube_2, + gripper_to_cube_3, + cube_1_to_2, + cube_2_to_3, + cube_1_to_3, + ), + dim=1, + ) + + +def instance_randomize_object_obs( + env: ManagerBasedRLEnv, + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), + ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"), +): + """ + Object observations (in world frame): + cube_1 pos, + cube_1 quat, + cube_2 pos, + cube_2 quat, + cube_3 pos, + cube_3 quat, + gripper to cube_1, + gripper to cube_2, + gripper to cube_3, + cube_1 to cube_2, + cube_2 to cube_3, + cube_1 to cube_3, + """ + if not hasattr(env, "rigid_objects_in_focus"): + return torch.full((env.num_envs, 9), fill_value=-1) + + cube_1: RigidObjectCollection = env.scene[cube_1_cfg.name] + cube_2: RigidObjectCollection = env.scene[cube_2_cfg.name] + cube_3: RigidObjectCollection = env.scene[cube_3_cfg.name] + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + + cube_1_pos_w = [] + cube_2_pos_w = [] + cube_3_pos_w = [] + cube_1_quat_w = [] + cube_2_quat_w = [] + cube_3_quat_w = [] + for env_id in range(env.num_envs): + cube_1_pos_w.append(cube_1.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][0], :3]) + cube_2_pos_w.append(cube_2.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][1], :3]) + cube_3_pos_w.append(cube_3.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][2], :3]) + cube_1_quat_w.append(cube_1.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][0], :4]) + cube_2_quat_w.append(cube_2.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][1], :4]) + cube_3_quat_w.append(cube_3.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][2], :4]) + cube_1_pos_w = torch.stack(cube_1_pos_w) + cube_2_pos_w = torch.stack(cube_2_pos_w) + cube_3_pos_w = torch.stack(cube_3_pos_w) + cube_1_quat_w = torch.stack(cube_1_quat_w) + cube_2_quat_w = torch.stack(cube_2_quat_w) + cube_3_quat_w = torch.stack(cube_3_quat_w) + + ee_pos_w = ee_frame.data.target_pos_w[:, 0, :] + gripper_to_cube_1 = cube_1_pos_w - ee_pos_w + gripper_to_cube_2 = cube_2_pos_w - ee_pos_w + gripper_to_cube_3 = cube_3_pos_w - ee_pos_w + + cube_1_to_2 = cube_1_pos_w - cube_2_pos_w + cube_2_to_3 = cube_2_pos_w - cube_3_pos_w + cube_1_to_3 = cube_1_pos_w - cube_3_pos_w + + return torch.cat( + ( + cube_1_pos_w, + cube_1_quat_w, + cube_2_pos_w, + cube_2_quat_w, + cube_3_pos_w, + cube_3_quat_w, + gripper_to_cube_1, + gripper_to_cube_2, + gripper_to_cube_3, + cube_1_to_2, + cube_2_to_3, + cube_1_to_3, + ), + dim=1, + ) + + +def ee_frame_pos(env: ManagerBasedRLEnv, ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame")) -> torch.Tensor: + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + ee_frame_pos = ee_frame.data.target_pos_w[:, 0, :] + + return ee_frame_pos + + +def ee_frame_quat(env: ManagerBasedRLEnv, ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame")) -> torch.Tensor: + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + ee_frame_quat = ee_frame.data.target_quat_w[:, 0, :] + + return ee_frame_quat + + +def gripper_pos(env: ManagerBasedRLEnv, robot_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + robot: Articulation = env.scene[robot_cfg.name] + finger_joint_1 = robot.data.joint_pos[:, -1].clone().unsqueeze(1) + finger_joint_2 = -1 * robot.data.joint_pos[:, -2].clone().unsqueeze(1) + + return torch.cat((finger_joint_1, finger_joint_2), dim=1) + + +def object_grasped( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg, + ee_frame_cfg: SceneEntityCfg, + object_cfg: SceneEntityCfg, + diff_threshold: float = 0.06, + gripper_open_val: torch.tensor = torch.tensor([0.04]), + gripper_threshold: float = 0.005, +) -> torch.Tensor: + """Check if an object is grasped by the specified robot.""" + + robot: Articulation = env.scene[robot_cfg.name] + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + object: RigidObject = env.scene[object_cfg.name] + + object_pos = object.data.root_pos_w + end_effector_pos = ee_frame.data.target_pos_w[:, 0, :] + pose_diff = torch.linalg.vector_norm(object_pos - end_effector_pos, dim=1) + + grasped = torch.logical_and( + pose_diff < diff_threshold, + torch.abs(robot.data.joint_pos[:, -1] - gripper_open_val.to(env.device)) > gripper_threshold, + ) + grasped = torch.logical_and( + grasped, torch.abs(robot.data.joint_pos[:, -2] - gripper_open_val.to(env.device)) > gripper_threshold + ) + + return grasped + + +def object_stacked( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg, + upper_object_cfg: SceneEntityCfg, + lower_object_cfg: SceneEntityCfg, + xy_threshold: float = 0.05, + height_threshold: float = 0.005, + height_diff: float = 0.0468, + gripper_open_val: torch.tensor = torch.tensor([0.04]), +) -> torch.Tensor: + """Check if an object is stacked by the specified robot.""" + + robot: Articulation = env.scene[robot_cfg.name] + upper_object: RigidObject = env.scene[upper_object_cfg.name] + lower_object: RigidObject = env.scene[lower_object_cfg.name] + + pos_diff = upper_object.data.root_pos_w - lower_object.data.root_pos_w + height_dist = torch.linalg.vector_norm(pos_diff[:, 2:], dim=1) + xy_dist = torch.linalg.vector_norm(pos_diff[:, :2], dim=1) + + stacked = torch.logical_and(xy_dist < xy_threshold, (height_dist - height_diff) < height_threshold) + + stacked = torch.logical_and( + torch.isclose(robot.data.joint_pos[:, -1], gripper_open_val.to(env.device), atol=1e-4, rtol=1e-4), stacked + ) + stacked = torch.logical_and( + torch.isclose(robot.data.joint_pos[:, -2], gripper_open_val.to(env.device), atol=1e-4, rtol=1e-4), stacked + ) + + return stacked diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/terminations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/terminations.py new file mode 100644 index 0000000000..9e5e8e161a --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/terminations.py @@ -0,0 +1,66 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Common functions that can be used to activate certain terminations for the lift task. + +The functions can be passed to the :class:`omni.isaac.lab.managers.TerminationTermCfg` object to enable +the termination introduced by the function. +""" + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from omni.isaac.lab.assets import Articulation, RigidObject +from omni.isaac.lab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from omni.isaac.lab.envs import ManagerBasedRLEnv + + +def cubes_stacked( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), + xy_threshold: float = 0.05, + height_threshold: float = 0.005, + height_diff: float = 0.0468, + gripper_open_val: torch.tensor = torch.tensor([0.04]), + atol=0.0001, + rtol=0.0001, +): + robot: Articulation = env.scene[robot_cfg.name] + cube_1: RigidObject = env.scene[cube_1_cfg.name] + cube_2: RigidObject = env.scene[cube_2_cfg.name] + cube_3: RigidObject = env.scene[cube_3_cfg.name] + + pos_diff_c12 = cube_1.data.root_pos_w - cube_2.data.root_pos_w + pos_diff_c23 = cube_2.data.root_pos_w - cube_3.data.root_pos_w + + # Compute cube position difference in x-y plane + xy_dist_c12 = torch.norm(pos_diff_c12[:, :2], dim=1) + xy_dist_c23 = torch.norm(pos_diff_c23[:, :2], dim=1) + + # Compute cube height difference + h_dist_c12 = torch.norm(pos_diff_c12[:, 2:], dim=1) + h_dist_c23 = torch.norm(pos_diff_c23[:, 2:], dim=1) + + # Check cube positions + stacked = torch.logical_and(xy_dist_c12 < xy_threshold, xy_dist_c23 < xy_threshold) + stacked = torch.logical_and(h_dist_c12 - height_diff < height_threshold, stacked) + stacked = torch.logical_and(h_dist_c23 - height_diff < height_threshold, stacked) + + # Check gripper positions + stacked = torch.logical_and( + torch.isclose(robot.data.joint_pos[:, -1], gripper_open_val.to(env.device), atol=atol, rtol=rtol), stacked + ) + stacked = torch.logical_and( + torch.isclose(robot.data.joint_pos[:, -2], gripper_open_val.to(env.device), atol=atol, rtol=rtol), stacked + ) + + return stacked diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/stack_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/stack_env_cfg.py new file mode 100644 index 0000000000..6d07861a59 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/stack_env_cfg.py @@ -0,0 +1,205 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import ArticulationCfg, AssetBaseCfg +from omni.isaac.lab.envs import ManagerBasedRLEnvCfg +from omni.isaac.lab.managers import ObservationGroupCfg as ObsGroup +from omni.isaac.lab.managers import ObservationTermCfg as ObsTerm +from omni.isaac.lab.managers import SceneEntityCfg +from omni.isaac.lab.managers import TerminationTermCfg as DoneTerm +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.sensors import CameraCfg +from omni.isaac.lab.sensors.frame_transformer.frame_transformer_cfg import FrameTransformerCfg +from omni.isaac.lab.sim.spawners.from_files.from_files_cfg import GroundPlaneCfg, UsdFileCfg +from omni.isaac.lab.utils import configclass +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR + +from . import mdp + + +## +# Scene definition +## +@configclass +class ObjectTableSceneCfg(InteractiveSceneCfg): + """Configuration for the lift scene with a robot and a object. + This is the abstract base implementation, the exact scene is defined in the derived classes + which need to set the target object, robot and end-effector frames + """ + + # robots: will be populated by agent env cfg + robot: ArticulationCfg = MISSING + # end-effector sensor: will be populated by agent env cfg + ee_frame: FrameTransformerCfg = MISSING + + # Cameras + wrist_cam: CameraCfg = MISSING + table_cam: CameraCfg = MISSING + + # Table + table = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/Table", + init_state=AssetBaseCfg.InitialStateCfg(pos=[0.5, 0, 0], rot=[0.707, 0, 0, 0.707]), + spawn=UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd"), + ) + + # plane + plane = AssetBaseCfg( + prim_path="/World/GroundPlane", + init_state=AssetBaseCfg.InitialStateCfg(pos=[0, 0, -1.05]), + spawn=GroundPlaneCfg(), + ) + + # lights + light = AssetBaseCfg( + prim_path="/World/light", + spawn=sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0), + ) + + +## +# MDP settings +## +@configclass +class ActionsCfg: + """Action specifications for the MDP.""" + + # will be set by agent env cfg + arm_action: mdp.JointPositionActionCfg = MISSING + gripper_action: mdp.BinaryJointPositionActionCfg = MISSING + + +@configclass +class ObservationsCfg: + """Observation specifications for the MDP.""" + + @configclass + class PolicyCfg(ObsGroup): + """Observations for policy group with state values.""" + + actions = ObsTerm(func=mdp.last_action) + joint_pos = ObsTerm(func=mdp.joint_pos_rel) + joint_vel = ObsTerm(func=mdp.joint_vel_rel) + object = ObsTerm(func=mdp.object_obs) + cube_positions = ObsTerm(func=mdp.cube_positions_in_world_frame) + cube_orientations = ObsTerm(func=mdp.cube_orientations_in_world_frame) + eef_pos = ObsTerm(func=mdp.ee_frame_pos) + eef_quat = ObsTerm(func=mdp.ee_frame_quat) + gripper_pos = ObsTerm(func=mdp.gripper_pos) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + @configclass + class RGBCameraPolicyCfg(ObsGroup): + """Observations for policy group with RGB images.""" + + table_cam = ObsTerm( + func=mdp.image, params={"sensor_cfg": SceneEntityCfg("table_cam"), "data_type": "rgb", "normalize": False} + ) + wrist_cam = ObsTerm( + func=mdp.image, params={"sensor_cfg": SceneEntityCfg("wrist_cam"), "data_type": "rgb", "normalize": False} + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + @configclass + class SubtaskCfg(ObsGroup): + """Observations for subtask group.""" + + grasp_1 = ObsTerm( + func=mdp.object_grasped, + params={ + "robot_cfg": SceneEntityCfg("robot"), + "ee_frame_cfg": SceneEntityCfg("ee_frame"), + "object_cfg": SceneEntityCfg("cube_2"), + }, + ) + stack_1 = ObsTerm( + func=mdp.object_stacked, + params={ + "robot_cfg": SceneEntityCfg("robot"), + "upper_object_cfg": SceneEntityCfg("cube_2"), + "lower_object_cfg": SceneEntityCfg("cube_1"), + }, + ) + grasp_2 = ObsTerm( + func=mdp.object_grasped, + params={ + "robot_cfg": SceneEntityCfg("robot"), + "ee_frame_cfg": SceneEntityCfg("ee_frame"), + "object_cfg": SceneEntityCfg("cube_3"), + }, + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + # observation groups + policy: PolicyCfg = PolicyCfg() + rgb_camera: RGBCameraPolicyCfg = RGBCameraPolicyCfg() + subtask_terms: SubtaskCfg = SubtaskCfg() + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + time_out = DoneTerm(func=mdp.time_out, time_out=True) + + cube_1_dropping = DoneTerm( + func=mdp.root_height_below_minimum, params={"minimum_height": -0.05, "asset_cfg": SceneEntityCfg("cube_1")} + ) + + cube_2_dropping = DoneTerm( + func=mdp.root_height_below_minimum, params={"minimum_height": -0.05, "asset_cfg": SceneEntityCfg("cube_2")} + ) + + cube_3_dropping = DoneTerm( + func=mdp.root_height_below_minimum, params={"minimum_height": -0.05, "asset_cfg": SceneEntityCfg("cube_3")} + ) + + success = DoneTerm(func=mdp.cubes_stacked) + + +@configclass +class StackEnvCfg(ManagerBasedRLEnvCfg): + """Configuration for the stacking environment.""" + + # Scene settings + scene: ObjectTableSceneCfg = ObjectTableSceneCfg(num_envs=4096, env_spacing=2.5, replicate_physics=False) + # Basic settings + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + # MDP settings + terminations: TerminationsCfg = TerminationsCfg() + + # Unused managers + commands = None + rewards = None + events = None + curriculum = None + + def __post_init__(self): + """Post initialization.""" + # general settings + self.decimation = 5 + self.episode_length_s = 30.0 + # simulation settings + self.sim.dt = 0.01 # 100Hz + 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_found_lost_aggregate_pairs_capacity = 1024 * 1024 * 4 + self.sim.physx.gpu_total_aggregate_pairs_capacity = 16 * 1024 + self.sim.physx.friction_correlation_distance = 0.00625 diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/stack_instance_randomize_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/stack_instance_randomize_env_cfg.py new file mode 100644 index 0000000000..1a5766b918 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/stack_instance_randomize_env_cfg.py @@ -0,0 +1,157 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import ArticulationCfg, AssetBaseCfg +from omni.isaac.lab.envs import ManagerBasedRLEnvCfg +from omni.isaac.lab.managers import ObservationGroupCfg as ObsGroup +from omni.isaac.lab.managers import ObservationTermCfg as ObsTerm +from omni.isaac.lab.managers import SceneEntityCfg +from omni.isaac.lab.managers import TerminationTermCfg as DoneTerm +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.sensors import CameraCfg +from omni.isaac.lab.sensors.frame_transformer.frame_transformer_cfg import FrameTransformerCfg +from omni.isaac.lab.sim.spawners.from_files.from_files_cfg import GroundPlaneCfg, UsdFileCfg +from omni.isaac.lab.utils import configclass +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR + +from . import mdp + + +## +# Scene definition +## +@configclass +class ObjectTableSceneCfg(InteractiveSceneCfg): + """Configuration for the lift scene with a robot and a object. + This is the abstract base implementation, the exact scene is defined in the derived classes + which need to set the target object, robot and end-effector frames + """ + + # robots: will be populated by agent env cfg + robot: ArticulationCfg = MISSING + # end-effector sensor: will be populated by agent env cfg + ee_frame: FrameTransformerCfg = MISSING + + # Cameras + wrist_cam: CameraCfg = MISSING + table_cam: CameraCfg = MISSING + + # Table + table = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/Table", + init_state=AssetBaseCfg.InitialStateCfg(pos=[0.5, 0, 0], rot=[0.707, 0, 0, 0.707]), + spawn=UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd"), + ) + + # plane + plane = AssetBaseCfg( + prim_path="/World/GroundPlane", + init_state=AssetBaseCfg.InitialStateCfg(pos=[0, 0, -1.05]), + spawn=GroundPlaneCfg(), + ) + + # lights + light = AssetBaseCfg( + prim_path="/World/light", + spawn=sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0), + ) + + +## +# MDP settings +## +@configclass +class ActionsCfg: + """Action specifications for the MDP.""" + + # will be set by agent env cfg + arm_action: mdp.JointPositionActionCfg = MISSING + gripper_action: mdp.BinaryJointPositionActionCfg = MISSING + + +@configclass +class ObservationsCfg: + """Observation specifications for the MDP.""" + + @configclass + class PolicyCfg(ObsGroup): + """Observations for policy group with state values.""" + + actions = ObsTerm(func=mdp.last_action) + joint_pos = ObsTerm(func=mdp.joint_pos_rel) + joint_vel = ObsTerm(func=mdp.joint_vel_rel) + object = ObsTerm(func=mdp.instance_randomize_object_obs) + cube_positions = ObsTerm(func=mdp.instance_randomize_cube_positions_in_world_frame) + cube_orientations = ObsTerm(func=mdp.instance_randomize_cube_orientations_in_world_frame) + eef_pos = ObsTerm(func=mdp.ee_frame_pos) + eef_quat = ObsTerm(func=mdp.ee_frame_quat) + gripper_pos = ObsTerm(func=mdp.gripper_pos) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + @configclass + class RGBCameraPolicyCfg(ObsGroup): + """Observations for policy group with RGB images.""" + + table_cam = ObsTerm( + func=mdp.image, params={"sensor_cfg": SceneEntityCfg("table_cam"), "data_type": "rgb", "normalize": False} + ) + wrist_cam = ObsTerm( + func=mdp.image, params={"sensor_cfg": SceneEntityCfg("wrist_cam"), "data_type": "rgb", "normalize": False} + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + # observation groups + policy: PolicyCfg = PolicyCfg() + rgb_camera: RGBCameraPolicyCfg = RGBCameraPolicyCfg() + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + time_out = DoneTerm(func=mdp.time_out, time_out=True) + + +@configclass +class StackInstanceRandomizeEnvCfg(ManagerBasedRLEnvCfg): + """Configuration for the stacking environment.""" + + # Scene settings + scene: ObjectTableSceneCfg = ObjectTableSceneCfg(num_envs=4096, env_spacing=2.5, replicate_physics=False) + # Basic settings + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + # MDP settings + terminations: TerminationsCfg = TerminationsCfg() + + # Unused managers + commands = None + rewards = None + events = None + curriculum = None + + def __post_init__(self): + """Post initialization.""" + # general settings + self.decimation = 5 + self.episode_length_s = 30.0 + # simulation settings + self.sim.dt = 0.01 # 100Hz + 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_found_lost_aggregate_pairs_capacity = 1024 * 1024 * 4 + self.sim.physx.gpu_total_aggregate_pairs_capacity = 16 * 1024 + self.sim.physx.friction_correlation_distance = 0.00625