curi_cabinet_image.py

# Copyright 2023, Junjia LIU, jjliu@mae.cuhk.edu.hk
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os

import tqdm
from isaacgym import gymapi, gymtorch
from isaacgym.torch_utils import *

import rofunc as rf
from rofunc.learning.RofuncRL.tasks.isaacgymenv.base.curi_base_task import CURIBaseTask
from rofunc.learning.RofuncRL.tasks.isaacgymenv.curi_cabinet import CURICabinetTask
from rofunc.utils.oslab.path import get_rofunc_path


class CURICabinetImageTask(CURICabinetTask, CURIBaseTask):

    def __init__(self, cfg, rl_device, sim_device, graphics_device_id, headless, virtual_screen_capture, force_render):
        CURIBaseTask.__init__(self, cfg, rl_device, sim_device, graphics_device_id, headless, virtual_screen_capture,
                              force_render)
        self.obs_buf = torch.zeros((self.num_envs, 128, 128, 4), device=self.device, dtype=torch.float)

        # get gym GPU state tensors
        actor_root_state_tensor = self.gym.acquire_actor_root_state_tensor(self.sim)
        dof_state_tensor = self.gym.acquire_dof_state_tensor(self.sim)
        rigid_body_tensor = self.gym.acquire_rigid_body_state_tensor(self.sim)

        self.gym.refresh_actor_root_state_tensor(self.sim)
        self.gym.refresh_dof_state_tensor(self.sim)
        self.gym.refresh_rigid_body_state_tensor(self.sim)

        # create some wrapper tensors for different slices
        # set default pos (seven links and left and right gripper, [:9] for the left arm, [9:] for the right arm)
        self.curi_default_dof_pos = to_torch(
            [0., 0., 0.3863, 0.5062, -0.1184, -0.6105, 0.023, 1.6737, 0.9197, 0.04, 0.04, -0.5349, 0, 0.1401, -1.7951,
             0.0334, 3.2965, 0.6484, 0.04, 0.04], device=self.device)
        self.dof_state = gymtorch.wrap_tensor(dof_state_tensor)
        self.curi_dof_state = self.dof_state.view(self.num_envs, -1, 2)[:, :self.num_curi_dofs]
        self.curi_dof_pos = self.curi_dof_state[..., 0]
        self.curi_dof_vel = self.curi_dof_state[..., 1]
        self.cabinet_dof_state = self.dof_state.view(self.num_envs, -1, 2)[:, self.num_curi_dofs:]
        self.cabinet_dof_pos = self.cabinet_dof_state[..., 0]
        self.cabinet_dof_vel = self.cabinet_dof_state[..., 1]

        self.rigid_body_states = gymtorch.wrap_tensor(rigid_body_tensor).view(self.num_envs, -1, 13)
        self.num_bodies = self.rigid_body_states.shape[1]

        self.root_state_tensor = gymtorch.wrap_tensor(actor_root_state_tensor).view(self.num_envs, -1, 13)

        if self.num_props > 0:
            self.prop_states = self.root_state_tensor[:, 2:]

        self.num_dofs = self.gym.get_sim_dof_count(self.sim) // self.num_envs
        self.curi_dof_targets = torch.zeros((self.num_envs, self.num_dofs), dtype=torch.float, device=self.device)

        self.global_indices = torch.arange(self.num_envs * (2 + self.num_props), dtype=torch.int32,
                                           device=self.device).view(self.num_envs, -1)

        self.reset_idx(torch.arange(self.num_envs, device=self.device))

    def _create_envs(self, num_envs, spacing, num_per_row):
        lower = gymapi.Vec3(-spacing, -spacing, 0.0)
        upper = gymapi.Vec3(spacing, spacing, spacing)

        # get rofunc path from rofunc package metadata
        rofunc_path = get_rofunc_path()
        asset_root = os.path.join(rofunc_path, "simulator/assets")

        curi_asset_file = "urdf/curi/urdf/curi_isaacgym_dual_arm_w_head.urdf"
        cabinet_asset_file = "urdf/sektion_cabinet_model/urdf/sektion_cabinet_2.urdf"

        # load curi asset
        asset_options = gymapi.AssetOptions()
        asset_options.flip_visual_attachments = True
        asset_options.fix_base_link = True
        asset_options.collapse_fixed_joints = True
        asset_options.disable_gravity = True
        # asset_options.thickness = 0.001
        asset_options.default_dof_drive_mode = gymapi.DOF_MODE_POS
        # asset_options.use_mesh_materials = True
        curi_asset = self.gym.load_asset(self.sim, asset_root, curi_asset_file, asset_options)

        """
        Rigid body dict for curi_isaacgym_dual_arm_w_head.urdf (since we use `collapse_fixed_joints` here)
        {'head_link1': 2,
         'head_link2': 3,
         'panda_left_leftfinger': 11,
         'panda_left_link1': 4,
         'panda_left_link2': 5,
         'panda_left_link3': 6,
         'panda_left_link4': 7,
         'panda_left_link5': 8,
         'panda_left_link6': 9,
         'panda_left_link7': 10,
         'panda_left_rightfinger': 12,
         'panda_right_leftfinger': 20,
         'panda_right_link1': 13,
         'panda_right_link2': 14,
         'panda_right_link3': 15,
         'panda_right_link4': 16,
         'panda_right_link5': 17,
         'panda_right_link6': 18,
         'panda_right_link7': 19,
         'panda_right_rightfinger': 21,
         'summit_xls_base_footprint': 1,
         'world': 0}
         """

        # load cabinet asset
        asset_options.flip_visual_attachments = False
        asset_options.collapse_fixed_joints = True
        asset_options.disable_gravity = False
        asset_options.default_dof_drive_mode = gymapi.DOF_MODE_NONE
        asset_options.armature = 0.005
        cabinet_asset = self.gym.load_asset(self.sim, asset_root, cabinet_asset_file, asset_options)

        curi_dof_stiffness = to_torch(
            [400, 400, 400, 400, 400, 400, 400, 400, 400, 1.0e6, 1.0e6, 400, 400, 400, 400, 400, 400, 400, 1.0e6,
             1.0e6], dtype=torch.float, device=self.device)
        curi_dof_damping = to_torch(
            [80, 80, 80, 80, 80, 80, 80, 80, 80, 1.0e2, 1.0e2, 80, 80, 80, 80, 80, 80, 80, 1.0e2, 1.0e2],
            dtype=torch.float,
            device=self.device)

        self.num_curi_bodies = self.gym.get_asset_rigid_body_count(curi_asset)
        self.num_curi_dofs = self.gym.get_asset_dof_count(curi_asset)
        self.num_cabinet_bodies = self.gym.get_asset_rigid_body_count(cabinet_asset)
        self.num_cabinet_dofs = self.gym.get_asset_dof_count(cabinet_asset)

        print("num env: ", num_envs)
        print("num curi bodies: ", self.num_curi_bodies)
        print("num curi dofs: ", self.num_curi_dofs)
        print("num cabinet bodies: ", self.num_cabinet_bodies)
        print("num cabinet dofs: ", self.num_cabinet_dofs)

        # set curi dof properties
        curi_dof_props = self.gym.get_asset_dof_properties(curi_asset)
        self.curi_dof_lower_limits = []
        self.curi_dof_upper_limits = []
        for i in range(self.num_curi_dofs):
            curi_dof_props['driveMode'][i] = gymapi.DOF_MODE_POS
            if self.physics_engine == gymapi.SIM_PHYSX:
                curi_dof_props['stiffness'][i] = curi_dof_stiffness[i]
                curi_dof_props['damping'][i] = curi_dof_damping[i]
            else:
                curi_dof_props['stiffness'][i] = 7000.0
                curi_dof_props['damping'][i] = 50.0

            self.curi_dof_lower_limits.append(curi_dof_props['lower'][i])
            self.curi_dof_upper_limits.append(curi_dof_props['upper'][i])

        self.curi_dof_lower_limits = to_torch(self.curi_dof_lower_limits, device=self.device)
        self.curi_dof_upper_limits = to_torch(self.curi_dof_upper_limits, device=self.device)
        self.curi_dof_speed_scales = torch.ones_like(self.curi_dof_lower_limits)
        self.curi_dof_speed_scales[[7, 8]] = 0.1
        self.curi_dof_speed_scales[[16, 17]] = 0.1
        curi_dof_props['effort'][7] = 200
        curi_dof_props['effort'][8] = 200
        curi_dof_props['effort'][16] = 200
        curi_dof_props['effort'][17] = 200

        # set cabinet dof properties
        cabinet_dof_props = self.gym.get_asset_dof_properties(cabinet_asset)
        for i in range(self.num_cabinet_dofs):
            cabinet_dof_props['damping'][i] = 10.0

        # create prop assets
        box_opts = gymapi.AssetOptions()
        box_opts.density = 400
        prop_asset = self.gym.create_box(self.sim, self.prop_width, self.prop_height, self.prop_width, box_opts)

        curi_start_pose = gymapi.Transform()
        curi_start_pose.p = gymapi.Vec3(1.5, 0.0, 0.0)
        curi_start_pose.r = gymapi.Quat(0.0, 0.0, 1.0, 0.0)

        cabinet_start_pose = gymapi.Transform()
        cabinet_start_pose.p = gymapi.Vec3(*get_axis_params(0.4, self.up_axis_idx))

        # compute aggregate size
        num_curi_bodies = self.gym.get_asset_rigid_body_count(curi_asset)
        num_curi_shapes = self.gym.get_asset_rigid_shape_count(curi_asset)
        num_cabinet_bodies = self.gym.get_asset_rigid_body_count(cabinet_asset)
        num_cabinet_shapes = self.gym.get_asset_rigid_shape_count(cabinet_asset)
        num_prop_bodies = self.gym.get_asset_rigid_body_count(prop_asset)
        num_prop_shapes = self.gym.get_asset_rigid_shape_count(prop_asset)
        max_agg_bodies = num_curi_bodies + num_cabinet_bodies + self.num_props * num_prop_bodies
        max_agg_shapes = num_curi_shapes + num_cabinet_shapes + self.num_props * num_prop_shapes

        camera_props = gymapi.CameraProperties()
        camera_props.enable_tensors = True
        camera_props.width = self.cfg['env']['image_size']
        camera_props.height = self.cfg['env']['image_size']
        attached_body = "head_link2"
        local_transform = gymapi.Transform()
        local_transform.p = gymapi.Vec3(0.12, 0, 0.18)
        local_transform.r = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), np.radians(90.0)) * \
                            gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), np.radians(-90.0))

        self.curis = []
        self.cabinets = []
        self.default_prop_states = []
        self.prop_start = []
        self.envs = []
        self.camera_handles = []
        self.camera_tensors = []

        rf.logger.beauty_print("Creating {} environments".format(num_envs), type="info")
        for i in tqdm.trange(self.num_envs):
            # create env instance
            env_ptr = self.gym.create_env(self.sim, lower, upper, num_per_row)

            if self.aggregate_mode >= 3:
                self.gym.begin_aggregate(env_ptr, max_agg_bodies, max_agg_shapes, True)

            curi_actor = self.gym.create_actor(env_ptr, curi_asset, curi_start_pose, "curi", i, 0)
            self.gym.set_actor_dof_properties(env_ptr, curi_actor, curi_dof_props)

            if self.aggregate_mode == 2:
                self.gym.begin_aggregate(env_ptr, max_agg_bodies, max_agg_shapes, True)

            # Camera Sensor
            camera_handle = self.gym.create_camera_sensor(env_ptr, camera_props)
            body_handle = self.gym.find_actor_rigid_body_handle(env_ptr, curi_actor, attached_body)
            self.gym.attach_camera_to_body(camera_handle, env_ptr, body_handle, local_transform,
                                           gymapi.FOLLOW_TRANSFORM)
            # obtain camera tensor
            cam_tensor = self.gym.get_camera_image_gpu_tensor(self.sim, env_ptr, camera_handle, gymapi.IMAGE_COLOR)
            # wrap camera tensor in a pytorch tensor
            torch_camera_tensor = gymtorch.wrap_tensor(cam_tensor)

            cabinet_pose = cabinet_start_pose
            cabinet_pose.p.x += self.start_position_noise * (np.random.rand() - 0.5)
            dz = 0.5 * np.random.rand()
            dy = np.random.rand() - 0.5
            cabinet_pose.p.y += self.start_position_noise * dy
            cabinet_pose.p.z += self.start_position_noise * dz
            cabinet_actor = self.gym.create_actor(env_ptr, cabinet_asset, cabinet_pose, "cabinet", i, 2, 0)
            self.gym.set_actor_dof_properties(env_ptr, cabinet_actor, cabinet_dof_props)

            if self.aggregate_mode == 1:
                self.gym.begin_aggregate(env_ptr, max_agg_bodies, max_agg_shapes, True)

            if self.num_props > 0:
                self.prop_start.append(self.gym.get_sim_actor_count(self.sim))
                drawer_handle = self.gym.find_actor_rigid_body_handle(env_ptr, cabinet_actor, "drawer_top")
                drawer_pose = self.gym.get_rigid_transform(env_ptr, drawer_handle)

                props_per_row = int(np.ceil(np.sqrt(self.num_props)))
                xmin = -0.5 * self.prop_spacing * (props_per_row - 1)
                yzmin = -0.5 * self.prop_spacing * (props_per_row - 1)

                prop_count = 0
                for j in range(props_per_row):
                    prop_up = yzmin + j * self.prop_spacing
                    for k in range(props_per_row):
                        if prop_count >= self.num_props:
                            break
                        propx = xmin + k * self.prop_spacing
                        prop_state_pose = gymapi.Transform()
                        prop_state_pose.p.x = drawer_pose.p.x + propx
                        propz, propy = 0, prop_up
                        prop_state_pose.p.y = drawer_pose.p.y + propy
                        prop_state_pose.p.z = drawer_pose.p.z + propz
                        prop_state_pose.r = gymapi.Quat(0, 0, 0, 1)
                        prop_handle = self.gym.create_actor(env_ptr, prop_asset, prop_state_pose,
                                                            "prop{}".format(prop_count), i, 0, 0)
                        prop_count += 1

                        prop_idx = j * props_per_row + k
                        self.default_prop_states.append([prop_state_pose.p.x, prop_state_pose.p.y, prop_state_pose.p.z,
                                                         prop_state_pose.r.x, prop_state_pose.r.y, prop_state_pose.r.z,
                                                         prop_state_pose.r.w,
                                                         0, 0, 0, 0, 0, 0])
            if self.aggregate_mode > 0:
                self.gym.end_aggregate(env_ptr)

            self.envs.append(env_ptr)
            self.curis.append(curi_actor)
            self.cabinets.append(cabinet_actor)
            self.camera_handles.append(camera_handle)
            self.camera_tensors.append(torch_camera_tensor)

        self.hand_handle = self.gym.find_actor_rigid_body_handle(env_ptr, curi_actor, "panda_left_link7")
        self.drawer_handle = self.gym.find_actor_rigid_body_handle(env_ptr, cabinet_actor, "drawer_top")
        self.lfinger_handle = self.gym.find_actor_rigid_body_handle(env_ptr, curi_actor, "panda_left_leftfinger")
        self.rfinger_handle = self.gym.find_actor_rigid_body_handle(env_ptr, curi_actor, "panda_left_rightfinger")
        self.default_prop_states = to_torch(self.default_prop_states, device=self.device, dtype=torch.float).view(
            self.num_envs, self.num_props, 13)

        # self.create_cameras()
        self.init_data()

    def create_cameras(self):
        # Camera Sensor
        camera_props = gymapi.CameraProperties()
        camera_props.enable_tensors = True
        camera_props.width = 128
        camera_props.height = 128

        attached_body = "head_link2"
        local_transform = gymapi.Transform()
        local_transform.p = gymapi.Vec3(0.12, 0, 0.18)
        local_transform.r = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), np.radians(90.0)) * \
                            gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), np.radians(-90.0))

        self.camera_handles = []
        self.camera_tensors = []
        rf.logger.beauty_print("Attaching {} cameras".format(self.num_envs), type="info")

        # def parallel(i):
        #     camera_handle = self.gym.create_camera_sensor(self.envs[i], camera_props)
        #     body_handle = self.gym.find_actor_rigid_body_handle(self.envs[i], self.curis[i], attached_body)
        #     self.gym.attach_camera_to_body(camera_handle, self.envs[i], body_handle, local_transform,
        #                                    gymapi.FOLLOW_TRANSFORM)
        #
        #     # obtain camera tensor
        #     cam_tensor = self.gym.get_camera_image_gpu_tensor(self.sim, self.envs[i], camera_handle, gymapi.IMAGE_COLOR)
        #     # wrap camera tensor in a pytorch tensor
        #     torch_camera_tensor = gymtorch.wrap_tensor(cam_tensor)
        #
        #     self.camera_handles.append(camera_handle)
        #     self.camera_tensors.append(torch_camera_tensor)
        #
        # from multiprocessing.pool import ThreadPool
        # # multiprocessing.set_start_method('fork')
        # pool = ThreadPool()
        # env_list = np.arange(self.num_envs)
        # pool.map(parallel, env_list)
        # pool.close()

        for i in tqdm.trange(self.num_envs):
            camera_handle = self.gym.create_camera_sensor(self.envs[i], camera_props)
            body_handle = self.gym.find_actor_rigid_body_handle(self.envs[i], self.curis[i], attached_body)
            self.gym.attach_camera_to_body(camera_handle, self.envs[i], body_handle, local_transform,
                                           gymapi.FOLLOW_TRANSFORM)

            # obtain camera tensor
            cam_tensor = self.gym.get_camera_image_gpu_tensor(self.sim, self.envs[i], camera_handle, gymapi.IMAGE_COLOR)
            # wrap camera tensor in a pytorch tensor
            # torch_camera_tensor = gymtorch.wrap_tensor(cam_tensor)

            self.camera_handles.append(camera_handle)
            # self.camera_tensors.append(torch_camera_tensor)

    def compute_observations(self):
        self.gym.refresh_actor_root_state_tensor(self.sim)
        self.gym.refresh_dof_state_tensor(self.sim)
        self.gym.refresh_rigid_body_state_tensor(self.sim)

        # Obtain image observations from camera sensors
        self.gym.render_all_camera_sensors(self.sim)
        self.gym.start_access_image_tensors(self.sim)

        self.obs_buf = torch.permute(torch.stack(self.camera_tensors).cuda()[:, :, :, :3], (0, 3, 1, 2))

        self.gym.end_access_image_tensors(self.sim)

        # hand_pos = self.rigid_body_states[:, self.hand_handle][:, 0:3]
        # hand_rot = self.rigid_body_states[:, self.hand_handle][:, 3:7]
        # drawer_pos = self.rigid_body_states[:, self.drawer_handle][:, 0:3]
        # drawer_rot = self.rigid_body_states[:, self.drawer_handle][:, 3:7]

        # self.curi_grasp_rot[:], self.curi_grasp_pos[:], self.drawer_grasp_rot[:], self.drawer_grasp_pos[:] = \
        #     compute_grasp_transforms(hand_rot, hand_pos, self.curi_local_grasp_rot, self.curi_local_grasp_pos,
        #                              drawer_rot, drawer_pos, self.drawer_local_grasp_rot, self.drawer_local_grasp_pos
        #                              )
        #
        # self.curi_lfinger_pos = self.rigid_body_states[:, self.lfinger_handle][:, 0:3]
        # self.curi_rfinger_pos = self.rigid_body_states[:, self.rfinger_handle][:, 0:3]
        # self.curi_lfinger_rot = self.rigid_body_states[:, self.lfinger_handle][:, 3:7]
        # self.curi_rfinger_rot = self.rigid_body_states[:, self.rfinger_handle][:, 3:7]
        #
        # dof_pos_scaled = (2.0 * (self.curi_dof_pos - self.curi_dof_lower_limits)
        #                   / (self.curi_dof_upper_limits - self.curi_dof_lower_limits) - 1.0)
        # to_target = self.drawer_grasp_pos - self.curi_grasp_pos
        # self.obs_buf = torch.cat((dof_pos_scaled, self.curi_dof_vel * self.dof_vel_scale, to_target,
        #                           self.cabinet_dof_pos[:, 3].unsqueeze(-1), self.cabinet_dof_vel[:, 3].unsqueeze(-1)),
        #                          dim=-1)
        return self.obs_buf