From 49e743faf4303cf6a26cadda637f5247a0011cd9 Mon Sep 17 00:00:00 2001 From: gursi26 Date: Wed, 15 Jan 2025 00:02:30 -0500 Subject: [PATCH] added robot, cube and container --- .../deformable_cube/deformable_cube_env.py | 194 ++++++++++++++++-- .../deformable_cube_env_cfg.py | 68 ------ .../lab_tasks/direct/deformable_cube/main.py | 18 ++ 3 files changed, 199 insertions(+), 81 deletions(-) delete mode 100644 source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/deformable_cube/deformable_cube_env_cfg.py create mode 100644 source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/deformable_cube/main.py diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/deformable_cube/deformable_cube_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/deformable_cube/deformable_cube_env.py index f04802c0bf..dd81b18ebc 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/deformable_cube/deformable_cube_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/deformable_cube/deformable_cube_env.py @@ -1,32 +1,174 @@ from __future__ import annotations - import math import torch -import torch.nn.functional as F +from dataclasses import MISSING from collections.abc import Sequence - import omni.isaac.lab.sim as sim_utils -from omni.isaac.lab.assets import RigidObjectCfg, RigidObject, Articulation, DeformableObject +import omni.isaac.lab.utils.math as math_utils +from omni.isaac.lab.actuators.actuator_cfg import ImplicitActuatorCfg +from omni.isaac.lab.assets import Articulation, DeformableObject, ArticulationCfg, DeformableObjectCfg, RigidObjectCfg, RigidObject from omni.isaac.lab.envs import DirectRLEnv, DirectRLEnvCfg from omni.isaac.lab.scene import InteractiveSceneCfg from omni.isaac.lab.sim import SimulationCfg -from omni.isaac.lab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane +from omni.isaac.lab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane, UsdFileCfg from omni.isaac.lab.utils import configclass -from omni.isaac.lab.utils.math import quat_from_euler_xyz, euler_xyz_from_quat -from .deformabe_cube_env_cfg import DeformableCubeEnvCfg + +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR +from omni.isaac.lab_assets.franka import FRANKA_PANDA_CFG + +CUBE_SIZE = 0.1 +CUBE_X_MIN, CUBE_X_MAX = 0.1, 1.0 +CUBE_Y_MIN, CUBE_Y_MAX = -0.5, 0.5 + +# TODO: Edit GPU settings for softbody contact buffer size +# TODO: mess with deformable object settings +# TODO: Add container + +def get_robot() -> ArticulationCfg: + return ArticulationCfg( + prim_path="/World/envs/env_.*/Robot", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/Franka/franka_instanceable.usd", + activate_contact_sensors=False, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + max_depenetration_velocity=5.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=False, solver_position_iteration_count=12, solver_velocity_iteration_count=1 + ), + ), + init_state=ArticulationCfg.InitialStateCfg( + joint_pos={ + "panda_joint1": 1.157, + "panda_joint2": -1.066, + "panda_joint3": -0.155, + "panda_joint4": -2.239, + "panda_joint5": -1.841, + "panda_joint6": 1.003, + "panda_joint7": 0.469, + "panda_finger_joint.*": 0.035, + } + ), + actuators={ + "panda_shoulder": ImplicitActuatorCfg( + joint_names_expr=["panda_joint[1-4]"], + effort_limit=87.0, + velocity_limit=2.175, + stiffness=80.0, + damping=4.0, + ), + "panda_forearm": ImplicitActuatorCfg( + joint_names_expr=["panda_joint[5-7]"], + effort_limit=12.0, + velocity_limit=2.61, + stiffness=80.0, + damping=4.0, + ), + "panda_hand": ImplicitActuatorCfg( + joint_names_expr=["panda_finger_joint.*"], + effort_limit=200.0, + velocity_limit=0.2, + stiffness=2e3, + damping=1e2, + ), + }, + ) + +def get_object() -> DeformableObjectCfg: + return DeformableObjectCfg( + prim_path="/World/envs/env_.*/Cube", + spawn=sim_utils.MeshCuboidCfg( + size=(CUBE_SIZE, CUBE_SIZE, CUBE_SIZE), + deformable_props=sim_utils.DeformableBodyPropertiesCfg( + rest_offset=0.0, + contact_offset=0.001 + ), + visual_material=sim_utils.PreviewSurfaceCfg( + diffuse_color=(0.5, 0.1, 0.0) + ), + physics_material=sim_utils.DeformableBodyMaterialCfg( + poissons_ratio=0.4, + youngs_modulus=1e5 + ), + ), + init_state=DeformableObjectCfg.InitialStateCfg(pos=(0.0, 0.0, CUBE_SIZE / 1.9)), + debug_vis=True, + ) + +def get_container() -> RigidObjectCfg: + return RigidObjectCfg( + prim_path="/World/envs/env_.*/Container", + spawn=sim_utils.UsdFileCfg( + usd_path=f"./Container_B04_40x30x12cm_PR_V_NVD_01.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + solver_position_iteration_count=8, + solver_velocity_iteration_count=1, + max_depenetration_velocity=0.01 + ), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg( + collision_enabled=True, + contact_offset=0.02, + rest_offset=0.001 + ), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.0, 1.0), metallic=0.2), + scale=(0.02, 0.02, 0.02), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.5, 0.5, 0.5)), + ) + +@configclass +class DeformableCubeEnvCfg(DirectRLEnvCfg): + num_envs = 5 + env_spacing = 3.0 + dt = 1 / 120 + observation_space = 15 + action_space = 9 + state_space = 0 + action_scale = 7.5 + + # env + decimation = 1 + episode_length_s = 5.0 + + # simulation + sim: SimulationCfg = SimulationCfg(dt=dt, render_interval=decimation) + scene: InteractiveSceneCfg = InteractiveSceneCfg( + num_envs=num_envs, + env_spacing=env_spacing, + replicate_physics=False, + ) + + # entities + robot_cfg: ArticulationCfg = get_robot() + object_cfg: DeformableObjectCfg = get_object() + container_cfg: RigidObjectCfg = get_container() + class DeformableCubeEnv(DirectRLEnv): cfg: DeformableCubeEnvCfg def __init__(self, cfg: DeformableCubeEnvCfg, render_mode: str | None = None, **kwargs): super().__init__(cfg, render_mode, **kwargs) - pass + + self.dt = self.cfg.sim.dt * self.cfg.decimation + + self.robot_dof_lower_limits = self.robot.data.soft_joint_pos_limits[0, :, 0].to(device=self.device) + self.robot_dof_upper_limits = self.robot.data.soft_joint_pos_limits[0, :, 1].to(device=self.device) + + self.robot_dof_speed_scales = torch.ones_like(self.robot_dof_lower_limits) + self.robot_dof_speed_scales[self.robot.find_joints("panda_finger_joint1")[0]] = 0.1 + self.robot_dof_speed_scales[self.robot.find_joints("panda_finger_joint2")[0]] = 0.1 + self.robot_dof_targets = torch.zeros((self.num_envs, self.robot.num_joints), device=self.device) def _setup_scene(self): self.robot = Articulation(self.cfg.robot_cfg) self.object = DeformableObject(self.cfg.object_cfg) - # TODO: spawn container + self.container = RigidObject(self.cfg.container_cfg) + self.scene.articulations["robot"] = self.robot spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) @@ -39,10 +181,12 @@ def _setup_scene(self): light_cfg.func("/World/Light", light_cfg) def _pre_physics_step(self, actions: torch.Tensor) -> None: - pass + self.actions = actions.clone().clamp(-1.0, 1.0) + targets = self.robot_dof_targets + self.robot_dof_speed_scales * self.dt * self.actions * self.cfg.action_scale + self.robot_dof_targets[:] = torch.clamp(targets, self.robot_dof_lower_limits, self.robot_dof_upper_limits) def _apply_action(self) -> None: - pass + self.robot.set_joint_position_target(self.robot_dof_targets) def _get_observations(self) -> dict: pass @@ -51,7 +195,31 @@ def _get_rewards(self) -> torch.Tensor: pass def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: - pass + out_of_bounds = torch.full((self.cfg.num_envs,), fill_value=False, dtype=bool) + time_out = self.episode_length_buf >= self.max_episode_length - 1 + return out_of_bounds, time_out def _reset_idx(self, env_ids: Sequence[int] | None): - pass \ No newline at end of file + print(f"[INFO]: Resetting envs {env_ids}") + self.episode_length_buf[env_ids] = 0.0 + + # reset objects + x_offset, y_offset = torch.rand(2, len(env_ids), 1, device=self.device) + x_offset = x_offset * (CUBE_X_MAX - CUBE_X_MIN) + CUBE_X_MIN + y_offset = y_offset * (CUBE_Y_MAX - CUBE_Y_MIN) + CUBE_Y_MIN + pos_t = torch.cat([x_offset, y_offset, torch.zeros_like(x_offset)], dim=-1) + + quat_t = torch.zeros(len(env_ids), 4, device=self.device) + quat_t[:, 0], quat_t[:, -1] = torch.rand(2, len(env_ids), device=self.device) * 2 - 1.0 + quat_t = quat_t / quat_t.norm(dim=1, keepdim=True) + + nodal_state = self.object.data.default_nodal_state_w[env_ids, :, :] + nodal_state[..., :3] = self.object.transform_nodal_pos(nodal_state[..., :3], pos=pos_t, quat=quat_t) + nodal_state[..., 3:] = 0.0 + self.object.write_nodal_state_to_sim(nodal_state, env_ids) + + # reset robots + joint_pos = self.robot.data.default_joint_pos[env_ids] + joint_vel = torch.zeros_like(joint_pos) + # self.robot.set_joint_position_target(joint_pos, env_ids=env_ids) + self.robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) \ No newline at end of file diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/deformable_cube/deformable_cube_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/deformable_cube/deformable_cube_env_cfg.py deleted file mode 100644 index aedc0d5d4b..0000000000 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/deformable_cube/deformable_cube_env_cfg.py +++ /dev/null @@ -1,68 +0,0 @@ -from __future__ import annotations - -from dataclasses import MISSING - -import omni.isaac.lab.sim as sim_utils -from omni.isaac.lab.assets import AssetBaseCfg, ArticulationCfg, DeformableObjectCfg -from omni.isaac.lab.envs import DirectRLEnv, DirectRLEnvCfg -from omni.isaac.lab.scene import InteractiveSceneCfg -from omni.isaac.lab.sim import SimulationCfg -from omni.isaac.lab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane, UsdFileCfg -from omni.isaac.lab.utils import configclass -from omni.isaac.lab.utils.math import quat_from_euler_xyz, euler_xyz_from_quat -from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR - -from omni.isaac.lab_assets.franka import FRANKA_PANDA_CFG - -@configclass -class DeformableCubeEnvCfg(DirectRLEnvCfg): - # required params - num_envs = 256 - env_spacing = 4.0 - dt = 1 / 120 - observation_space = 15 - action_space = 3 - state_space = 0 - - # env params - decimation = 1 - episode_length_s = 15.0 - - # simulation - sim: SimulationCfg = SimulationCfg( - dt=dt, - render_interval=decimation - ) - - # robot - robot_cfg: ArticulationCfg = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - - # objects - object_cfg: DeformableObjectCfg = DeformableObjectCfg( - prim_path="{ENV_REGEX_NS}/Cube", - spawn=sim_utils.MeshCuboidCfg( - size=(0.2, 0.2, 0.2), - deformable_props=sim_utils.DeformableBodyPropertiesCfg( - rest_offset=0.0, - contact_offset=0.001 - ), - visual_material=sim_utils.PreviewSurfaceCfg( - diffuse_color=(0.5, 0.1, 0.0) - ), - physics_material=sim_utils.DeformableBodyMaterialCfg( - poissons_ratio=0.4, - youngs_modulus=1e5 - ), - ), - init_state=DeformableObjectCfg.InitialStateCfg(), - debug_vis=True, - ) - - container_cfg: AssetBaseCfg = MISSING #TODO: fill this out - - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=num_envs, - env_spacing=env_spacing, - replicate_physics=True, - ) \ No newline at end of file diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/deformable_cube/main.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/deformable_cube/main.py new file mode 100644 index 0000000000..2eea18194e --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/deformable_cube/main.py @@ -0,0 +1,18 @@ +import argparse +import torch +from omni.isaac.lab.app import AppLauncher + +parser = argparse.ArgumentParser() +AppLauncher.add_app_launcher_args(parser) +args_cli = parser.parse_args() +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +from deformable_cube_env import DeformableCubeEnv, DeformableCubeEnvCfg + +if __name__ == "__main__": + env_cfg = DeformableCubeEnvCfg() + env = DeformableCubeEnv(env_cfg) + env.reset() + while simulation_app.is_running(): + env.step(torch.randn(env_cfg.num_envs, env_cfg.action_space, device="cuda"))