Implementation of the RA-L2023 paper: Part-Guided 3D RL for Sim2Real Articulated Object Manipulation
__version__ = "0.1.0"
from pathlib import Path

ROOT_DIR = Path(__file__).parent.resolve()
REPO_DIR = Path(__file__).parent.parent.resolve()
ASSET_DIR = ROOT_DIR / "assets"
AGENT_CONFIG_DIR = ASSET_DIR / "configs"
DESCRIPTION_DIR = ASSET_DIR / "descriptions"
ALG_DIR = ROOT_DIR / "algorithms"
# with DR: arti_data_seg, arti_data_seg_randombg, artidata384_seg_DR_randombg_mixfaucet
# w/o DR: arti_data_seg_noDR, arti_data_seg_noDR_randombg
SEGDATA_DIR = ROOT_DIR / "algorithms/data_process/artidata384_seg_DR_randombg"
KPTDATA_DIR = ROOT_DIR / "algorithms/data_process/artidata384_keypoints_noDR_norandombg"
REAL_DIR = ROOT_DIR / "algorithms/data_process/real_capdata_doordrawerfaucet_360_640"
REALEXP_DIR = ROOT_DIR / "test/real_exps"
VISUALMODEL_DIR = ROOT_DIR / "algorithms/visual_net/scripts/log"
RLMODEL_DIR = ROOT_DIR / "algorithms/rl_iam/sac/logs"

arti_mani/__init__.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from collections import OrderedDict
from typing import Dict, Tuple

import numpy as np
import sapien.core as sapien
import transforms3d as t3d
from arti_mani.agents.camera import get_camera_seg, get_texture
from sapien.core import ActorBase, Pose, Scene
from sapien.sensor.depth_processor import calc_main_depth_from_left_right_ir


class ActiveLightSensor:
    def __init__(
        self,
        name: str,
        scene: Scene,
        mount: ActorBase,
        rgb_resolution: Tuple[int, int],
        ir_resolution: Tuple[int, int],
        rgb_intrinsic: np.ndarray,
        ir_intrinsic: np.ndarray,
        trans_pose_l: Pose,
        trans_pose_r: Pose,
        light_pattern: str,
        max_depth: float = 8.0,
        min_depth: float = 0.3,
        ir_ambient_strength: float = 0.002,
        ir_light_dim_factor: float = 0.05,
        ir_light_fov: float = 1.57,
        ir_intensity: float = 10.0,
    ):
        self.name = name
        self._scene = scene
        self._cam_mount = mount
        self._rgb_w, self._rgb_h = rgb_resolution
        self._ir_w, self._ir_h = ir_resolution
        self._rgb_intrinsic = rgb_intrinsic
        self._ir_intrinsic = ir_intrinsic
        self._trans_pose_l = trans_pose_l
        self._trans_pose_r = trans_pose_r
        self._light_pattern = light_pattern
        self._max_depth = max_depth
        self._min_depth = min_depth
        self._ir_ambient_strength = ir_ambient_strength
        self._ir_light_dim_factor = ir_light_dim_factor
        self._ir_light_fov = ir_light_fov
        self._ir_intensity = ir_intensity

        self._create_cameras()
        self.alight = self._scene.add_active_light(
            pose=Pose([0, 0, 0]),
            color=[0, 0, 0],
            fov=self._ir_light_fov,
            tex_path=self._light_pattern,
        )

    def _create_cameras(self):
        tran_pose0 = sapien.Pose([0, 0, 0])
        camera0 = self._scene.add_mounted_camera(
            f"{self.name}",
            self._cam_mount,
            tran_pose0,
            self._rgb_w,
            self._rgb_h,
            fovy=0.0,
            near=0.001,
            far=100,
        )
        camera0.set_perspective_parameters(
            0.1,
            100.0,
            self._rgb_intrinsic[0, 0],
            self._rgb_intrinsic[1, 1],
            self._rgb_intrinsic[0, 2],
            self._rgb_intrinsic[1, 2],
            self._rgb_intrinsic[0, 1],
        )

        camera1 = self._scene.add_mounted_camera(
            f"{self.name}_ir_left",
            self._cam_mount,
            self._trans_pose_l,
            self._ir_w,
            self._ir_h,
            fovy=0.0,
            near=0.001,
            far=100,
        )
        camera1.set_perspective_parameters(
            0.1,
            100.0,
            self._ir_intrinsic[0, 0],
            self._ir_intrinsic[1, 1],
            self._ir_intrinsic[0, 2],
            self._ir_intrinsic[1, 2],
            self._ir_intrinsic[0, 1],
        )
        camera2 = self._scene.add_mounted_camera(
            f"{self.name}_ir_right",
            self._cam_mount,
            self._trans_pose_r,
            self._ir_w,
            self._ir_h,
            fovy=0.0,
            near=0.001,
            far=100,
        )
        camera2.set_perspective_parameters(
            0.1,
            100.0,
            self._ir_intrinsic[0, 0],
            self._ir_intrinsic[1, 1],
            self._ir_intrinsic[0, 2],
            self._ir_intrinsic[1, 2],
            self._ir_intrinsic[0, 1],
        )

        self._cam_rgb, self._cam_ir_l, self._cam_ir_r = camera0, camera1, camera2

    def get_image_dict(self, rgb=True, visual_seg=True, actor_seg=True) -> OrderedDict:
        apos = t3d.quaternions.mat2quat(
            self._cam_mount.get_pose().to_transformation_matrix()[:3, :3]
            @ t3d.quaternions.quat2mat((-0.5, 0.5, 0.5, -0.5))
        )
        self.alight.set_pose(Pose(self._cam_mount.get_pose().p, apos))
        self.alight.set_color([0, 0, 0])

        self._scene.update_render()
        self._cam_rgb.take_picture()

        self._ir_mode()
        self._scene.update_render()
        self._cam_ir_l.take_picture()
        self._cam_ir_r.take_picture()

        self._normal_mode()
        self._scene.update_render()

        ir_l = get_texture(self._cam_ir_l, "Color")[:, :, 0]
        ir_r = get_texture(self._cam_ir_r, "Color")[:, :, 0]
        ir_l = self._float2uint8(ir_l)
        ir_r = self._float2uint8(ir_r)

        clean_depth = -get_texture(self._cam_rgb, "Position")[:, :, [2]]  # unit: meter

        # stereo
        ex_l = self._pose2cv2ex(self._trans_pose_l)
        ex_r = self._pose2cv2ex(self._trans_pose_r)
        ex_main = self._pose2cv2ex(Pose())

        depth = calc_main_depth_from_left_right_ir(
            ir_l,
            ir_r,
            ex_l,
            ex_r,
            ex_main,
            self._ir_intrinsic,
            self._ir_intrinsic,
            self._rgb_intrinsic,
            lr_consistency=False,
            main_cam_size=(self._rgb_w, self._rgb_h),
            ndisp=128,
            use_census=True,
            register_depth=True,
            census_wsize=7,
            use_noise=True,
        )
        depth[depth > self._max_depth] = 0
        depth[depth < self._min_depth] = 0

        sensor_dict = OrderedDict()
        sensor_dict["ir_l"] = ir_l
        sensor_dict["ir_r"] = ir_r
        sensor_dict["clean_depth"] = clean_depth
        sensor_dict["stereo_depth"] = depth
        if rgb:
            rgb = get_texture(self._cam_rgb, "Color")[:, :, :3]
            rgb = self._float2uint8(rgb)
            sensor_dict["rgb"] = rgb
        if visual_seg or actor_seg:
            seg = get_camera_seg(self._cam_rgb)
            if visual_seg:
                sensor_dict["visual_seg"] = seg[..., 0]
            if actor_seg:
                sensor_dict["actor_seg"] = seg[..., 1]

        return sensor_dict

    def _ir_mode(self):
        self._light_d = {}
        for l in self._scene.get_all_lights():
            self._light_d[l] = l.color
            l.set_color(l.color * self._ir_light_dim_factor)

        self._light_a = self._scene.ambient_light
        self._scene.set_ambient_light([self._ir_ambient_strength, 0, 0])
        self.alight.set_color([self._ir_intensity, 0, 0])

    def _normal_mode(self):
        for l in self._scene.get_all_lights():
            l.set_color(self._light_d[l])
        self._scene.set_ambient_light(self._light_a)
        self.alight.set_color([0, 0, 0])

    @staticmethod
    def _float2uint8(x):
        return (x * 255).clip(0, 255).astype(np.uint8)

    @staticmethod
    def _pose2cv2ex(pose):
        ros2opencv = np.array(
            [
                [0.0, -1.0, 0.0, 0.0],
                [0.0, 0.0, -1.0, 0.0],
                [1.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 1.0],
            ],
            dtype=np.float32,
        )
        return ros2opencv @ np.linalg.inv(pose.to_transformation_matrix())

arti_mani/agents/active_light_sensor.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import copy
from collections import OrderedDict, deque
from typing import Dict, List, Tuple

import attr
import numpy as np
import sapien.core as sapien
import yaml
from arti_mani import ASSET_DIR, DESCRIPTION_DIR
from arti_mani.agents.camera import get_camera_images, get_camera_pcd
from arti_mani.agents.controllers.combined_controller import CombinedController
from arti_mani.utils.cv_utils import depth2pts_np
from arti_mani.utils.sapien_utils import get_entity_by_name, make_actor_visible
from gym import spaces
from sapien.sensor import StereoDepthSensor, StereoDepthSensorConfig


@attr.s(auto_attribs=True, kw_only=True)
class MountedCameraConfig:
    name: str
    # extrinsic parameters
    mount: str  # name of link to mount
    mount_p: List[float]  # position relative to link
    mount_q: List[float]  # quaternion relative to link
    hide_mount_link: bool  # whether to hide the mount link
    # intrinsic parameters
    width: int
    height: int
    near: float
    far: float
    fx: float
    fy: float
    cx: float
    cy: float
    skew: float


@attr.s(auto_attribs=True, kw_only=True)
class MountedActiveLightSensorConfig:
    name: str
    # extrinsic parameters
    mount: str  # name of link to mount
    # intrinsic parameters
    rgb_resolution: Tuple[int, int]
    ir_resolution: Tuple[int, int]
    rgb_intrinsic: List[float]
    ir_intrinsic: List[float]
    trans_pose_l: List[float]
    trans_pose_r: List[float]
    light_pattern_path: str
    max_depth: float
    min_depth: float
    max_disp: int
    ir_ambient_strength: float
    ir_light_dim_factor: float
    ir_light_fov: float
    ir_intensity: float


def create_mounted_camera(
    config: MountedCameraConfig, robot: sapien.Articulation, scene: sapien.Scene
) -> sapien.CameraEntity:
    camera_mount = get_entity_by_name(robot.get_links(), config.mount)
    mount_pose = sapien.Pose(config.mount_p, config.mount_q)
    camera = scene.add_mounted_camera(
        config.name,
        camera_mount,
        mount_pose,
        width=config.width,
        height=config.height,
        fovy=0,  # focal will be set later.
        near=config.near,
        far=config.far,
    )
    camera.set_perspective_parameters(
        near=config.near,
        far=config.far,
        fx=config.fx,
        fy=config.fy,
        cx=config.cx,
        cy=config.cy,
        skew=config.skew,
    )
    if config.hide_mount_link:
        make_actor_visible(camera_mount, False)
    return camera


def create_mounted_sensor(
    config: MountedActiveLightSensorConfig,
    robot: sapien.Articulation,
    scene: sapien.Scene,
):
    sensor_mount = get_entity_by_name(robot.get_links(), config.mount)
    rgb_intrinsic = np.array(config.rgb_intrinsic).astype(np.float)
    ir_intrinsic = np.array(config.ir_intrinsic).astype(np.float)
    trans_pose_l = sapien.Pose(config.trans_pose_l[:3], config.trans_pose_l[3:])
    trans_pose_r = sapien.Pose(config.trans_pose_r[:3], config.trans_pose_r[3:])

    sensor_config = StereoDepthSensorConfig()
    sensor_config.rgb_resolution = config.rgb_resolution
    sensor_config.rgb_intrinsic = rgb_intrinsic
    sensor_config.ir_resolution = config.ir_resolution
    sensor_config.ir_intrinsic = ir_intrinsic
    sensor_config.trans_pose_l = trans_pose_l
    sensor_config.trans_pose_r = trans_pose_r
    sensor_config.light_pattern = str(ASSET_DIR / config.light_pattern_path)
    sensor_config.min_depth = config.min_depth
    sensor_config.max_depth = config.max_depth
    sensor_config.max_disp = config.max_disp
    sensor = StereoDepthSensor(config.name, scene, sensor_config, mount=sensor_mount)
    return sensor


@attr.s(auto_attribs=True, kw_only=True)
class AgentConfig:
    agent_class: str
    name: str
    urdf_file: str
    fix_root_link: bool = True
    urdf_config: Dict = attr.ib(factory=dict)

    default_init_qpos: List[float]
    default_control_mode: str

    controllers: Dict[str, List[Dict]]
    cameras: List[Dict] = attr.ib(factory=list)
    sensors: List[Dict] = attr.ib(factory=list)

    torque_freq: int


def parse_urdf_config(config_dict: Dict, scene: sapien.Scene) -> Dict:
    """Parse config from dict for SAPIEN URDF loader.

    Args:
        config_dict (Dict): dict containing link physical properties.
        scene (sapien.Scene): simualtion scene

    Returns:
        Dict: urdf config passed to `sapien.URDFLoader.load`.
    """
    config_dict = copy.deepcopy(config_dict)
    urdf_config = {}

    phy_mtls = {}
    for k, v in config_dict.get("materials", {}).items():
        phy_mtls[k] = scene.create_physical_material(**v)

    link_configs = config_dict.get("links", {})
    for link_name, link_config in link_configs.items():
        # replace with actual material
        link_config["material"] = phy_mtls[link_config["material"]]

    urdf_config["link"] = link_configs
    return urdf_config


class BaseAgent:
    """Base class for agents.

    Agent is an interface of the robot (sapien.Articulation).

    Args:
        config (AgentConfig):  agent configuration.
        scene (sapien.Scene): simulation scene instance.
        control_freq (int): control frequency (Hz).

    """

    _config: AgentConfig
    _scene: sapien.Scene
    _robot: sapien.Articulation
    _combined_controllers: Dict[str, CombinedController]
    _control_mode: str
    _cameras: Dict[str, sapien.CameraEntity]
    _sensors: Dict[str, StereoDepthSensor]
    _q_vel_ex: np.ndarray  # cached q velocity of previous step for external torque
    _sim_time_step: float

    def __init__(self, config: AgentConfig, scene: sapien.Scene, control_freq: int):
        self._config = copy.deepcopy(config)

        self._scene = scene
        self._sim_time_step = scene.timestep
        self._control_freq = control_freq
        assert control_freq % config.torque_freq == 0
        torque_queue_len = control_freq // config.torque_freq
        self._torque_queue = deque(maxlen=torque_queue_len)

        self._initialize_robot()
        self._initialize_controllers()
        self._initialize_cameras()
        self._initialize_sensors()

    def _initialize_robot(self):
        loader = self._scene.create_urdf_loader()
        loader.fix_root_link = self._config.fix_root_link

        urdf_file = DESCRIPTION_DIR / self._config.urdf_file
        urdf_config = parse_urdf_config(self._config.urdf_config, self._scene)
        self._robot = loader.load(str(urdf_file), urdf_config)
        self._robot.set_name(self._config.name)

    def _initialize_controllers(self):
        self._combined_controllers = OrderedDict()
        for control_mode, controller_configs in self._config.controllers.items():
            self._combined_controllers[control_mode] = CombinedController(
                controller_configs, self._robot, self._control_freq
            )
        self._control_mode = self._config.default_control_mode
        self._combined_controllers[self._control_mode].set_joint_drive_property()

    def _initialize_cameras(self):
        self._cameras = OrderedDict()
        for config in self._config.cameras:
            config = MountedCameraConfig(**config)
            if config.name in self._cameras:
                raise KeyError("Non-unique camera name: {}".format(config.name))
            cam = create_mounted_camera(config, self._robot, self._scene)
            self._cameras[config.name] = cam

    def _initialize_sensors(self):
        self._sensors = OrderedDict()
        for config in self._config.sensors:
            config = MountedActiveLightSensorConfig(**config)
            if config.name in self._sensors:
                raise KeyError(
                    "Non-unique active light sensor name: {}".format(config.name)
                )
            sensor = create_mounted_sensor(config, self._robot, self._scene)
            self._sensors[config.name] = sensor

    @classmethod
    def from_config_file(cls, config_path: str, scene: sapien.Scene, control_freq: int):
        with open(config_path, "r") as f:
            config_dict = yaml.safe_load(f)
        config = AgentConfig(**config_dict)
        return cls(config, scene, control_freq)

    @property
    def supported_control_modes(self) -> Tuple[str]:
        return tuple(self._combined_controllers.keys())

    @property
    def control_mode(self) -> str:
        return self._control_mode

    def set_control_mode(self, control_mode: str):
        assert (
            control_mode in self.supported_control_modes
        ), f"Unknown control mode: [{control_mode}]. All supported control modes: [{self.supported_control_modes}]"
        self._control_mode = control_mode
        self._combined_controllers[control_mode].set_joint_drive_property()
        self._combined_controllers[control_mode].reset()

    @property
    def action_space(self):
        return spaces.Dict(
            {
                mode: spaces.Box(
                    controller.action_range[:, 0], controller.action_range[:, 1]
                )
                for mode, controller in self._combined_controllers.items()
            }
        )

    @property
    def action_range(self) -> spaces.Box:
        return spaces.Box(
            self._combined_controllers[self._control_mode].action_range[:, 0],
            self._combined_controllers[self._control_mode].action_range[:, 1],
        )

    def set_action(self, action: np.ndarray):
        self._combined_controllers[self._control_mode].set_action(action)

    def simulation_step(self):
        self._combined_controllers[self._control_mode].simulation_step()
        self._q_vel_ex = self._robot.get_qvel()

    def reset(self, init_qpos=None):
        if init_qpos is None:
            init_qpos = self._config.default_init_qpos
        self._robot.set_qpos(init_qpos)
        self._robot.set_qvel(np.zeros(self._robot.dof))
        self._robot.set_qacc(np.zeros(self._robot.dof))
        self._robot.set_qf(np.zeros(self._robot.dof))
        self._combined_controllers[self._control_mode].reset()
        self._q_vel_ex = self._robot.get_qvel()
        self._torque_queue.clear()

    def get_proprioception(self) -> Dict:
        raise NotImplementedError

    def get_state(self) -> Dict:
        """Get current state for MPC, including robot state and controller state"""
        state = OrderedDict()

        # robot state
        root_link = self._robot.get_links()[0]
        state["robot_root_pose"] = root_link.get_pose()
        state["robot_root_vel"] = root_link.get_velocity()
        state["robot_root_qvel"] = root_link.get_angular_velocity()
        state["robot_qpos"] = self._robot.get_qpos()
        state["robot_qvel"] = self._robot.get_qvel()
        state["robot_qacc"] = self._robot.get_qacc()
        return state

    def set_state(self, state: Dict):
        # robot state
        self._robot.set_root_pose(state["robot_root_pose"])
        self._robot.set_root_velocity(state["robot_root_vel"])
        self._robot.set_root_angular_velocity(state["robot_root_qvel"])
        self._robot.set_qpos(state["robot_qpos"])
        self._robot.set_qvel(state["robot_qvel"])
        self._robot.set_qacc(state["robot_qacc"])

    def get_fused_pointcloud(
        self, rgb=True, visual_seg=False, actor_seg=False
    ) -> Dict[str, np.ndarray]:
        """get pointcloud from each camera, transform them to the *world* frame, and fuse together"""
        assert self._cameras
        pcds = []
        # self._scene.update_render()
        for cam in self._cameras.values():
            cam.take_picture()
            pcd = get_camera_pcd(cam, rgb, visual_seg, actor_seg)  # dict
            # Model matrix is the transformation from OpenGL camera space to SAPIEN world space
            # camera.get_model_matrix() must be called after scene.update_render()!
            T = cam.get_model_matrix()
            R = T[:3, :3]
            t = T[:3, 3]
            pcd["xyz"] = pcd["xyz"] @ R.transpose() + t
            pcds.append(pcd)
        fused_pcd = {}
        for key in pcds[0].keys():
            fused_pcd[key] = np.concatenate([pcd[key] for pcd in pcds], axis=0)
        return fused_pcd

    def get_sensor_fused_point_cloud(self, rgb=True, visual_seg=True, actor_seg=True):
        assert self._sensors
        pcds = []
        for sensor_name, sensor in self._sensors.items():
            sensor_dict = sensor.get_image_dict(rgb, visual_seg, actor_seg)
            camera_intrinsic = sensor._rgb_intrinsic
            cam_extrinsic_world_frame = sensor._cam_rgb.get_extrinsic_matrix()
            pcd = {}
            xyz = depth2pts_np(
                sensor_dict["stereo_depth"], camera_intrinsic, cam_extrinsic_world_frame
            )
            num_pts = xyz.shape[0]
            mask = sensor_dict["stereo_depth"].reshape(-1) > 0
            xyz = xyz[mask]
            pcd["xyz"] = xyz

            for k in ("rgb", "visual_seg", "actor_seg"):
                if k in sensor_dict:
                    pcd[k] = sensor_dict[k].reshape(num_pts, -1)[mask]
            pcds.append(pcd)

        fused_pcd = {}
        for k in pcds[0].keys():
            fused_pcd[k] = np.concatenate([pcd[k] for pcd in pcds], axis=0)

        return fused_pcd

    def get_images(
        self, rgb=True, depth=False, visual_seg=False, actor_seg=False
    ) -> Dict[str, Dict[str, np.ndarray]]:
        assert self._cameras
        all_images = OrderedDict()
        # self._scene.update_render()
        for cam_name, cam in self._cameras.items():
            cam.take_picture()
            cam_images = get_camera_images(
                cam, rgb=rgb, depth=depth, visual_seg=visual_seg, actor_seg=actor_seg
            )
            cam_images["camera_intrinsic"] = cam.get_intrinsic_matrix()
            cam_extrinsic_world_frame = cam.get_extrinsic_matrix()
            robot_base_frame = (
                self._robot.get_root_pose().to_transformation_matrix()
            )  # robot base -> world
            cam_images["camera_extrinsic_base_frame"] = (
                cam_extrinsic_world_frame @ robot_base_frame
            )  # robot base -> camera
            all_images[cam_name] = cam_images

        if depth:
            for sensor_name, sensor in self._sensors.items():
                sensor_dict = OrderedDict()
                sensor.take_picture()
                sensor.compute_depth()
                sensor_dict["rgb"] = sensor.get_rgb()
                sensor_dict["depth"] = sensor.get_depth()[..., None]
                # sensor_dict = sensor.get_image_dict()
                # sensor_seg_dict = get_camera_images(
                #     sensor._cam_rgb,
                #     rgb=False,
                #     depth=False,
                #     visual_seg=visual_seg,
                #     actor_seg=actor_seg,
                # )
                # sensor_dict.update(sensor_seg_dict)
                # sensor_dict["camera_intrinsic"] = sensor._cam_rgb.get_intrinsic_matrix()
                # cam_extrinsic_world_frame = sensor._cam_rgb.get_extrinsic_matrix()
                # robot_base_frame = (
                #     self._robot.get_root_pose().to_transformation_matrix()
                # )  # robot base -> world
                # sensor_dict["camera_extrinsic_base_frame"] = (
                #     cam_extrinsic_world_frame @ robot_base_frame
                # )  # robot base -> camera
                all_images[sensor_name] = sensor_dict

        return all_images

    def get_camera_poses(self) -> Dict[str, np.ndarray]:
        assert self._cameras
        pose_list = OrderedDict()
        for cam_name, cam in self._cameras.items():
            pose_list[cam_name] = cam.get_pose().to_transformation_matrix()
        return pose_list

    def update_generalized_external_forces(self):
        """Note: only called in step_action()"""
        q_pos = self._robot.get_qpos()
        q_vel = self._robot.get_qvel()
        actual_acc = (q_vel - self._q_vel_ex) / self._sim_time_step
        actual_torques = self._robot.compute_inverse_dynamics(actual_acc)
        command_torques = []
        for j_idx, j in enumerate(self._robot.get_active_joints()):
            command_torque = j.stiffness * (
                j.get_drive_target() - q_pos[j_idx]
            ) + j.damping * (j.get_drive_velocity_target() - q_vel[j_idx])
            command_torques.append(command_torque)

        command_torques = np.array(command_torques)
        F = (
            actual_torques
            - command_torques
            - self._robot.get_qf()
            + self._robot.compute_passive_force()
        )
        self._torque_queue.append(F)

    def get_generalized_external_forces(self):
        if self._torque_queue:
            return np.mean(np.stack(self._torque_queue), axis=0)
        else:
            return np.zeros_like(self._robot.get_qf())

arti_mani/agents/base_agent.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
from collections import OrderedDict
from typing import Dict, List, Optional

import numpy as np
import sapien.core as sapien

try:
    import cupy
except ImportError as e:
    print("cupy is not installed.")


def get_texture(camera: sapien.CameraEntity, name, dtype="float"):
    """Get texture from camera.
    Faster data communication are supported by cupy.

    Args:
        camera (sapien.CameraEntity): SAPIEN camera
        name (str): texture name
        dtype (str, optional): texture dtype, [float/uint32].
            Defaults to "float".

    Returns:
        np.ndarray: texture
    """
    if os.environ.get("WITH_CUPY", 0) == 1:
        # dtype info is included already in the dl_tensor
        dlpack = camera.get_dl_tensor(name)
        return cupy.asnumpy(cupy.from_dlpack(dlpack))
    else:
        if dtype == "float":
            return camera.get_float_texture(name)
        elif dtype == "uint32":
            return camera.get_uint32_texture(name)
        else:
            raise NotImplementedError(f"Unsupported texture type: {dtype}")


def get_camera_rgb(camera: sapien.CameraEntity, uint8=True):
    image = get_texture(camera, "Color")[..., :3]  # [H, W, 3]
    if uint8:
        image = np.clip(image * 255, 0, 255).astype(np.uint8)
    return image


def get_camera_depth(camera: sapien.CameraEntity):
    # position texture is in OpenGL frame, and thus depth is negative.
    # The unit is meter
    depth = -get_texture(camera, "Position")[..., [2]]  # [H, W, 1]
    # mask = get_texture(camera, "Position")[..., [3]]
    # depth[mask == 1] = camera.far
    return depth


def get_camera_seg(camera: sapien.CameraEntity):
    seg = get_texture(camera, "Segmentation", "uint32")  # [H, W, 4]
    # channel 0 is visual id (mesh-level)
    # channel 1 is actor id (actor-level)
    return seg[..., :2]


def get_camera_images(
    camera: sapien.CameraEntity,
    rgb=True,
    depth=False,
    visual_seg=False,
    actor_seg=False,
) -> Dict[str, np.ndarray]:
    # Assume camera.take_picture() is called
    images = OrderedDict()
    if rgb:
        images["rgb"] = get_camera_rgb(camera)
    if depth:
        images["depth"] = get_camera_depth(camera)
    if visual_seg or actor_seg:
        seg = get_camera_seg(camera)
        if visual_seg:
            images["visual_seg"] = seg[..., 0]
        if actor_seg:
            images["actor_seg"] = seg[..., 1]
    return images


def get_camera_pcd(
    camera: sapien.CameraEntity,
    rgb=True,
    visual_seg=False,
    actor_seg=False,
) -> Dict[str, np.ndarray]:
    pcd = OrderedDict()
    # Each pixel is (x, y, z, z_buffer_depth) in OpenGL camera space
    position = get_texture(camera, "Position")  # [H, W, 4]
    # Remove invalid points
    mask = position[..., -1] < 1
    pcd["xyz"] = position[..., :3][mask]
    if rgb:
        pcd["rgb"] = get_camera_rgb(camera)[mask]
    if visual_seg or actor_seg:
        seg = get_camera_seg(camera)
        if visual_seg:
            pcd["visual_seg"] = seg[..., 0][mask]
        if actor_seg:
            pcd["actor_seg"] = seg[..., 1][mask]
    return pcd

arti_mani/agents/camera.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
# general controllers
from .floatinggripper_ee_pose import FloatGripperPoseController
from .general_pd_joint_pos import GeneralPDJointPosController

# gripper controllers
from .gripper_pd_joint_pos_mimic import GripperPDJointPosMimicController

arti_mani/agents/controllers/__init__.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
# Controller takes in action and transforms it to sapien joint control
from collections.abc import Iterable
from typing import List

import numpy as np
import sapien.core as sapien


class BaseController:
    """Base controller class"""

    action_dimension: int
    control_type: str

    def __init__(
        self, controller_config: dict, robot: sapien.Articulation, control_freq: int
    ):
        self.robot = robot
        self.control_freq = control_freq
        self.sim_timestep = self.robot.get_builder().get_scene().get_timestep()
        self.control_step = (
            round(1 / self.sim_timestep) // control_freq
        )  # used for action interpolation
        self.curr_step = 0
        self.interpolate = controller_config[
            "interpolate"
        ]  # whether to linear interpolate between simulation steps

        # list of strings
        self.control_joint_names = controller_config["control_joints"]
        self.control_joints: List[sapien.Joint] = []
        self.control_joint_index = []  # the index of controlled joint in `robot`
        for joint_name in self.control_joint_names:
            found = False
            for j_idx, j in enumerate(self.robot.get_active_joints()):
                if j.get_name() == joint_name:
                    self.control_joints.append(j)
                    self.control_joint_index.append(j_idx)
                    found = True
                    break
            if not found:
                raise AssertionError(
                    f"Can NOT find joint [{joint_name}] in robot."
                    f" All joint names: {[j.get_name() for j in self.robot.get_active_joints()]}"
                )
        self.control_joint_index = np.array(self.control_joint_index)
        # check the control joint index is continuous
        assert (
            len(self.control_joint_index)
            == self.control_joint_index.max() - self.control_joint_index.min() + 1
        )
        self.num_control_joints = len(self.control_joints)
        self.start_link = self.control_joints[0].get_parent_link()
        self.end_link = self.control_joints[-1].get_child_link()
        self.start_link_idx = 0
        self.end_link_idx = 0
        for link_idx, link in enumerate(robot.get_links()):
            if link == self.start_link:
                self.start_link_idx = link_idx
            if link == self.end_link:
                self.end_link_idx = link_idx
        self.start_joint_idx = 0
        self.end_joint_idx = 0

        for joint_idx, j in enumerate(
            robot.get_joints()
        ):  # find the joint idx in all the robot joints (including fix)
            if joint_idx == 0:  # first joint is the joint from None to base link
                continue
            if j.get_name() == self.control_joint_names[0]:
                self.start_joint_idx = joint_idx
                continue
            if j.get_name() == self.control_joint_names[-1]:
                self.end_joint_idx = joint_idx
                continue

        assert self.end_joint_idx > self.start_joint_idx > 0
        # create a new articulation to compute mass matrix and Jacobian
        builder: sapien.ArticulationBuilder = (
            self.robot.get_builder().get_scene().create_articulation_builder()
        )
        root: sapien.LinkBuilder = builder.create_link_builder()
        root.set_mass_and_inertia(
            self.start_link.get_mass(),
            self.start_link.cmass_local_pose,
            self.start_link.get_inertia(),
        )
        links = [root]
        all_joints = robot.get_joints()[self.start_joint_idx : self.end_joint_idx + 1]
        for j_idx, j in enumerate(all_joints):
            link = builder.create_link_builder(links[-1])
            link.set_mass_and_inertia(
                j.get_child_link().get_mass(),
                j.get_child_link().cmass_local_pose,
                j.get_child_link().get_inertia(),
            )
            link.set_joint_properties(
                j.type, j.get_limits(), j.get_pose_in_parent(), j.get_pose_in_child()
            )
            links.append(link)

        self.controller_articulation = builder.build(
            fix_root_link=True
        )  # fix=False will make inverse dynamic wrong
        self.controller_articulation.set_name(controller_config["controller_name"])

    @staticmethod
    def nums2array(nums, dim):
        """
        Convert input @nums into numpy array of length @dim. If @nums is a single number, broadcasts it to the
        corresponding dimension size @dim before converting into a numpy array

        Args:
            nums (numeric or Iterable): Either single value or array of numbers
            dim (int): Size of array to broadcast input to env.sim.data.actuator_force

        Returns:
            np.array: Array filled with values specified in @nums
        """
        # First run sanity check to make sure no strings are being inputted
        if isinstance(nums, str):
            raise TypeError(
                "Error: Only numeric inputs are supported for this function, nums2array!"
            )
        if isinstance(nums, Iterable):
            if len(nums) != dim:
                raise TypeError(
                    f"Error: the dimension of nums [{len(nums)}]is not consistent with dim[{dim}]."
                )
        # Check if input is an Iterable, if so, we simply convert the input to np.array and return
        # Else, input is a single value, so we map to a numpy array of correct size and return
        return np.array(nums) if isinstance(nums, Iterable) else np.ones(dim) * nums

    def set_action(self, action: np.ndarray):
        pass

    def simulation_step(self):
        pass

    @property
    def action_range(self) -> np.ndarray:
        return np.arange([])

    def _get_curr_joint_pos(self):
        return self.robot.get_qpos()[self.control_joint_index].copy()

    def _get_curr_joint_vel(self):
        return self.robot.get_qvel()[self.control_joint_index].copy()

    def _sync_articulation(self):
        # sync the robot qpos to controller articulation to compute mass matrix and Jacobian
        self.controller_articulation.set_qpos(self._get_curr_joint_pos())
        self.controller_articulation.set_qvel(self._get_curr_joint_vel())
        self.controller_articulation.set_root_pose(self.start_link.get_pose())
        self.controller_articulation.set_root_velocity(self.start_link.get_velocity())
        self.controller_articulation.set_root_angular_velocity(
            self.start_link.get_angular_velocity()
        )

    def set_joint_drive_property(self):
        """Set the joint drive propertry according to the controller type"""
        raise NotImplementedError()

    def reset(self):
        raise NotImplementedError()

arti_mani/agents/controllers/base_controller.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from copy import deepcopy
from typing import Dict, List, Tuple

import arti_mani.agents.controllers as controller_zoo
import numpy as np
import sapien.core as sapien
from arti_mani.agents.controllers.base_controller import BaseController


class CombinedController:
    """A class to combine a list of controllers"""

    _controllers: List[BaseController]
    _controller_action_idx: List[Tuple[int]]

    def __init__(
        self,
        controller_configs: List[Dict],
        robot: sapien.Articulation,
        control_freq: int,
    ):
        self._robot = robot
        self._control_freq = control_freq
        self._controllers = []
        self._controller_action_idx = []
        self.total_torque = None
        self.total_torque_imp = None

        idx_before = 0
        controlled_joint_idx = []
        combined_action_range = []

        for controller_config in controller_configs:
            ControllerClass = getattr(
                controller_zoo, controller_config["controller_type"]
            )
            controller: BaseController = ControllerClass(
                controller_config, robot, control_freq
            )
            self._controllers.append(controller)
            controlled_joint_idx.extend(list(controller.control_joint_index))
            combined_action_range.append(controller.action_range)
            self._controller_action_idx.append(
                (idx_before, idx_before + controller.action_dimension)
            )
            idx_before += controller.action_dimension

        controlled_joint_idx = np.array(controlled_joint_idx)
        assert np.allclose(
            controlled_joint_idx, np.arange(len(self._robot.get_active_joints()))
        ), "All active joints should be controlled by a controller"

        self._action_range = np.concatenate(combined_action_range)

    @property
    def action_range(self) -> np.ndarray:
        return self._action_range

    def simulation_step(self):
        qf_total = self._robot.compute_passive_force(external=False)
        self.total_torque_imp = deepcopy(qf_total)
        for i, controller in enumerate(self._controllers):
            if controller.control_type == "torque":
                qf_total += controller.simulation_step()
            else:
                controller.simulation_step()
        self.total_torque = deepcopy(qf_total)
        # torque_max = np.max(np.abs(qf_total))
        # torque_min = np.min(np.abs(qf_total))
        # print("qf_total max, min:", torque_max, torque_min)
        # self._robot.torque_max_min[0] = torque_max if self.torque_max_min[0] < torque_max else self.torque_max_min[0]
        # self._robot.torque_max_min[1] = torque_min if self.torque_max_min[1] > torque_min else self.torque_max_min[1]
        # print("qf_total: ", qf_total - self._robot.compute_passive_force(external=False))
        qf_total_addnoise = qf_total.copy()
        mean, sigma = 0, 0
        noise = np.clip(
            np.random.normal(mean, sigma, qf_total[:-2].size), -2 * sigma, 2 * sigma
        )
        qf_total_addnoise[:-2] += noise
        # add (m=0, s=100, 10, 1, 0.1, 0.01) norm noise
        self._robot.set_qf(qf_total_addnoise)

    def set_action(self, action: np.ndarray):
        assert action.shape == self._action_range.shape[:1]
        for i, (idx_start, idx_end) in enumerate(self._controller_action_idx):
            self._controllers[i].set_action(action[idx_start:idx_end])

    def set_joint_drive_property(self):
        """set the joint drive property of all controllers"""
        for controller in self._controllers:
            controller.set_joint_drive_property()

    def reset(self):
        for controller in self._controllers:
            controller.reset()

arti_mani/agents/controllers/combined_controller.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import numpy as np
import sapien.core as sapien
from arti_mani.agents.controllers.base_controller import BaseController
from transforms3d.quaternions import axangle2quat, qmult


class FloatGripperPoseController(BaseController):
    def __init__(
        self, controller_config: dict, robot: sapien.Articulation, control_freq: int
    ):
        super().__init__(controller_config, robot, control_freq)
        self.action_dimension = 6
        self.control_type = "pos"
        assert not self.interpolate, "Do NOT support interpolation"

        self.ee_pos_min = self.nums2array(controller_config["ee_pos_min"], 3)
        self.ee_pos_max = self.nums2array(controller_config["ee_pos_max"], 3)
        self.ee_rot_min = self.nums2array(controller_config["ee_rot_min"], 3)
        self.ee_rot_max = self.nums2array(controller_config["ee_rot_max"], 3)
        self.joint_stiffness = self.nums2array(
            controller_config["joint_stiffness"], self.num_control_joints
        )
        self.joint_damping = self.nums2array(
            controller_config["joint_damping"], self.num_control_joints
        )
        self.joint_friction = self.nums2array(
            controller_config["joint_friction"], self.num_control_joints
        )

        # self.pmodel = self.robot.create_pinocchio_model()
        self.qmask = np.zeros(self.robot.dof)
        self.qmask[self.control_joint_index] = 1
        self.target_joint_pose = self._get_curr_joint_pos()

    def set_joint_drive_property(self):
        for j_idx, j in enumerate(self.control_joints):
            j.set_drive_property(self.joint_stiffness[j_idx], self.joint_damping[j_idx])
            j.set_friction(self.joint_friction[j_idx])

    def reset(self):
        self.target_joint_pose = self._get_curr_joint_pos()

    @property
    def action_range(self) -> np.ndarray:
        lower_bound = np.hstack([self.ee_pos_min, self.ee_rot_min])
        upper_bound = np.hstack([self.ee_pos_max, self.ee_rot_max])
        return np.stack([lower_bound, upper_bound], axis=1)

    def set_action(self, action: np.ndarray):
        assert action.shape[0] == self.action_dimension
        target_ee_pos = self.end_link.pose.p + action[0:3]

        angle = np.linalg.norm(action[3:6])
        if angle < 1e-6:
            axis = (0, 0, 1)
            angle = 0
        else:
            axis = action[3:6] / angle
        delta_quat = axangle2quat(axis, angle)
        target_ee_quat = qmult(self.end_link.pose.q, delta_quat)
        target_ee_pose = sapien.Pose(target_ee_pos, target_ee_quat)
        # print(self.target_ee_pos, self.end_link.pose.p)
        if target_ee_pose is not None:
            self.target_joint_pos = target_ee_pose

    def simulation_step(self):
        for j_idx, j in enumerate(self.control_joints):
            j.set_drive_target(self.target_joint_pos[j_idx])

arti_mani/agents/controllers/floatinggripper_ee_pose.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import numpy as np
import sapien.core as sapien
from arti_mani.agents.controllers.base_controller import BaseController


class GeneralPDJointPosController(BaseController):
    """General PD joint position controller"""

    def __init__(
        self, controller_config: dict, robot: sapien.Articulation, control_freq: int
    ):
        super(GeneralPDJointPosController, self).__init__(
            controller_config, robot, control_freq
        )
        self.action_dimension = self.num_control_joints
        self.control_type = "pos"
        self.use_delta = controller_config[
            "use_delta"
        ]  # interpret action as delta or absolute pose

        if self.use_delta:
            self.joint_delta_pos_min = self.nums2array(
                controller_config["joint_delta_pos_min"], self.num_control_joints
            )
            self.joint_delta_pos_max = self.nums2array(
                controller_config["joint_delta_pos_max"], self.num_control_joints
            )
        else:
            self.joint_pos_min = self.nums2array(
                controller_config["joint_pos_min"], self.num_control_joints
            )
            self.joint_pos_max = self.nums2array(
                controller_config["joint_pos_max"], self.num_control_joints
            )

        self.joint_stiffness = self.nums2array(
            controller_config["joint_stiffness"], self.num_control_joints
        )
        self.joint_damping = self.nums2array(
            controller_config["joint_damping"], self.num_control_joints
        )
        self.joint_friction = self.nums2array(
            controller_config["joint_friction"], self.num_control_joints
        )

        self.curr_joint_pos: np.ndarray = self._get_curr_joint_pos()
        self.start_joint_pos: np.ndarray = self.curr_joint_pos.copy()
        self.final_target_joint_pos: np.ndarray = self.curr_joint_pos.copy()

    @property
    def action_range(self) -> np.ndarray:
        if self.use_delta:
            return np.stack(
                [self.joint_delta_pos_min, self.joint_delta_pos_max], axis=1
            )
        else:
            return np.stack([self.joint_pos_min, self.joint_pos_max], axis=1)

    def set_action(self, action: np.ndarray):
        assert action.shape[0] == self.action_dimension
        self.curr_step = 0
        self.start_joint_pos = self._get_curr_joint_pos()
        if self.use_delta:
            self.final_target_joint_pos = action + self.start_joint_pos
        else:
            self.final_target_joint_pos = action

    def simulation_step(self):
        self.curr_step += 1
        if self.interpolate:
            curr_target_joint_pos = (
                self.final_target_joint_pos - self.start_joint_pos
            ) / self.control_step * self.curr_step + self.start_joint_pos
        else:
            curr_target_joint_pos = self.final_target_joint_pos

        for j_idx, j in enumerate(self.control_joints):
            j.set_drive_target(curr_target_joint_pos[j_idx])

    def set_joint_drive_property(self):
        for j_idx, j in enumerate(self.control_joints):
            j.set_drive_property(self.joint_stiffness[j_idx], self.joint_damping[j_idx])
            j.set_friction(self.joint_friction[j_idx])

    def reset(self):
        self.curr_joint_pos: np.ndarray = self._get_curr_joint_pos()
        self.start_joint_pos: np.ndarray = self.curr_joint_pos.copy()
        self.final_target_joint_pos: np.ndarray = self.curr_joint_pos.copy()

arti_mani/agents/controllers/general_pd_joint_pos.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import numpy as np
import sapien.core as sapien
from arti_mani.agents.controllers.base_controller import BaseController


class GripperPDJointPosMimicController(BaseController):
    def __init__(
        self, controller_config: dict, robot: sapien.Articulation, control_freq: int
    ):
        super(GripperPDJointPosMimicController, self).__init__(
            controller_config, robot, control_freq
        )
        self.action_dimension = 1
        self.control_type = "pos"
        self.joint_stiffness = controller_config["joint_stiffness"]
        self.joint_damping = controller_config["joint_damping"]
        self.joint_friction = controller_config["joint_friction"]
        self.joint_force_limit = controller_config["joint_force_limit"]

        self.joint_pos_min = controller_config["joint_pos_min"]
        self.joint_pos_max = controller_config["joint_pos_max"]

        self.curr_joint_pos: np.ndarray = self._get_curr_joint_pos()
        self.start_joint_pos: np.ndarray = self.curr_joint_pos.copy()
        self.final_target_joint_pos: np.ndarray = self.curr_joint_pos.copy()

    @property
    def action_range(self) -> np.ndarray:
        return np.array([self.joint_pos_min, self.joint_pos_max]).reshape(1, 2)

    def set_action(self, action: np.ndarray):
        assert action.shape[0] == self.action_dimension
        self.curr_step = 0
        self.start_joint_pos = self._get_curr_joint_pos()
        self.final_target_joint_pos = np.ones(self.num_control_joints) * action

    def simulation_step(self):
        self.curr_step += 1
        if self.interpolate:
            curr_target_joint_pos = (
                self.final_target_joint_pos - self.start_joint_pos
            ) / self.control_step * self.curr_step + self.start_joint_pos
        else:
            curr_target_joint_pos = self.final_target_joint_pos

        for j_idx, j in enumerate(self.control_joints):
            j.set_drive_target(curr_target_joint_pos[j_idx])

    def set_joint_drive_property(self):
        for j_idx, j in enumerate(self.control_joints):
            j.set_drive_property(
                self.joint_stiffness, self.joint_damping, self.joint_force_limit
            )
            j.set_friction(self.joint_friction)

    def reset(self):
        self.curr_joint_pos: np.ndarray = self._get_curr_joint_pos()
        self.start_joint_pos: np.ndarray = self.curr_joint_pos.copy()
        self.final_target_joint_pos: np.ndarray = self.curr_joint_pos.copy()

arti_mani/agents/controllers/gripper_pd_joint_pos_mimic.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from collections import OrderedDict

import numpy as np
import sapien.core as sapien
from arti_mani.agents.base_agent import BaseAgent
from arti_mani.utils.common import compute_angle_between
from arti_mani.utils.geometry import transform_points
from arti_mani.utils.sapien_utils import (
    check_joint_stuck,
    get_actor_by_name,
    get_entity_by_name,
    get_pairwise_contact_impulse,
)
from arti_mani.utils.trimesh_utils import get_actor_mesh
from sapien.core import Pose


class FixedXmate3Robotiq(BaseAgent):
    def __init__(self, *args, **kwargs):
        super(FixedXmate3Robotiq, self).__init__(*args, **kwargs)
        self.finger1_link: sapien.LinkBase = get_actor_by_name(
            self._robot.get_links(), "left_inner_finger_pad"
        )
        self.finger2_link: sapien.LinkBase = get_actor_by_name(
            self._robot.get_links(), "right_inner_finger_pad"
        )
        self.finger1_mesh = get_actor_mesh(self.finger1_link, False)
        self.finger2_mesh = get_actor_mesh(self.finger2_link, False)
        self.fingers_visid = []
        for visual_body in self.finger1_link.get_visual_bodies():
            self.fingers_visid.append(visual_body.get_visual_id())
        for visual_body in self.finger2_link.get_visual_bodies():
            self.fingers_visid.append(visual_body.get_visual_id())
        self.finger_size = (0.03, 0.07, 0.0075)  # values from URDF
        self.grasp_site: sapien.Link = get_entity_by_name(
            self._robot.get_links(), "grasp_convenient_link"
        )

    def get_proprioception(self):
        state_dict = OrderedDict()
        qpos = self._robot.get_qpos()
        qvel = self._robot.get_qvel()

        state_dict["qpos"] = qpos
        state_dict["qvel"] = qvel
        state_dict["tcp_wrench"] = self.get_tcp_wrench()
        state_dict["joint_external_torque"] = self.get_generalized_external_forces()
        state_dict["gripper_grasp"] = np.array(
            [self.check_gripper_grasp_real()]
        ).astype(np.float)

        return state_dict

    def check_grasp(self, actor: sapien.ActorBase, min_impulse=1e-6, max_angle=85):
        assert isinstance(actor, sapien.ActorBase), type(actor)
        contacts = self._scene.get_contacts()

        limpulse = get_pairwise_contact_impulse(contacts, self.finger1_link, actor)
        rimpulse = get_pairwise_contact_impulse(contacts, self.finger2_link, actor)

        # direction to open the gripper
        ldirection = self.finger1_link.pose.to_transformation_matrix()[:3, 2]
        rdirection = self.finger2_link.pose.to_transformation_matrix()[:3, 2]

        # angle between impulse and open direction
        langle = compute_angle_between(ldirection, limpulse)
        rangle = compute_angle_between(rdirection, rimpulse)

        lflag = (
            np.linalg.norm(limpulse) >= min_impulse and np.rad2deg(langle) <= max_angle
        )
        rflag = (
            np.linalg.norm(rimpulse) >= min_impulse and np.rad2deg(rangle) <= max_angle
        )

        return all([lflag, rflag])

    def sample_ee_coords(self, num_sample=10) -> np.ndarray:
        """Uniformly sample points on the two finger meshes. Used for dense reward computation
        return: ee_coords (2, num_sample, 3)"""
        finger_points = (
            np.arange(num_sample) / (num_sample - 1) - 0.5
        ) * self.finger_size[1]
        finger_points = np.stack(
            [np.zeros(num_sample), finger_points, np.zeros(num_sample)], axis=1
        )  # (num_sample, 3)

        finger1_points = transform_points(
            self.finger1_link.get_pose().to_transformation_matrix(), finger_points
        )
        finger2_points = transform_points(
            self.finger2_link.get_pose().to_transformation_matrix(), finger_points
        )

        ee_coords = np.stack((finger1_points, finger2_points))

        return ee_coords

    def get_tcp_wrench(self):
        joint_tau = self.get_generalized_external_forces()[:7]
        controller = self._combined_controllers[self._control_mode]._controllers[0]
        assert controller.control_joint_names == [
            "joint1",
            "joint2",
            "joint3",
            "joint4",
            "joint5",
            "joint6",
            "joint7",
        ]
        controller._sync_articulation()
        J_full = controller.controller_articulation.compute_world_cartesian_jacobian()[
            -6:
        ]
        J = np.linalg.pinv(J_full.T)

        TCP_wrench = J @ joint_tau
        return TCP_wrench

    @staticmethod
    def build_grasp_pose(forward, flat, center):
        extra = np.cross(flat, forward)
        ans = np.eye(4)
        ans[:3, :3] = np.array([forward, flat, -extra]).T
        ans[:3, 3] = center
        return Pose.from_transformation_matrix(ans)

    def check_gripper_grasp_real(self) -> bool:
        """check whether the gripper is grasping something by checking the joint position and velocity"""
        from arti_mani.agents.controllers import GripperPDJointPosMimicController

        assert isinstance(
            self._combined_controllers[self._control_mode]._controllers[1],
            GripperPDJointPosMimicController,
        )
        for joint_idx, joint in enumerate(self._robot.get_active_joints()):
            if joint.name == "robotiq_2f_140_left_driver_joint":
                active_joint1_idx = joint_idx
            if joint.name == "robotiq_2f_140_right_driver_joint":
                active_joint2_idx = joint_idx

        joint1_stuck = check_joint_stuck(
            self._robot, active_joint1_idx, pos_diff_threshold=1e-3, vel_threshold=1
        )
        joint2_stuck = check_joint_stuck(
            self._robot, active_joint2_idx, pos_diff_threshold=1e-3, vel_threshold=1
        )

        return joint1_stuck and joint2_stuck

arti_mani/agents/fixed_xmate3_robotiq.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import copy
from collections import OrderedDict
from typing import Dict

import numpy as np
import sapien.core as sapien
import yaml
from arti_mani import DESCRIPTION_DIR
from arti_mani.agents.base_agent import (
    ActiveLightSensor,
    AgentConfig,
    MountedActiveLightSensorConfig,
    MountedCameraConfig,
    create_mounted_camera,
    create_mounted_sensor,
    parse_urdf_config,
)
from arti_mani.agents.camera import get_camera_images
from arti_mani.agents.controllers.combined_controller import CombinedController
from arti_mani.utils.common import compute_angle_between
from arti_mani.utils.geometry import transform_points
from arti_mani.utils.sapien_utils import (
    check_joint_stuck,
    get_actor_by_name,
    get_entity_by_name,
    get_pairwise_contact_impulse,
)
from gym import spaces
from sapien.core import Pose


class FloatingRobotiq:
    _config: AgentConfig
    _scene: sapien.Scene
    _robot: sapien.Articulation
    _cameras: Dict[str, sapien.CameraEntity]
    _sensors: Dict[str, ActiveLightSensor]

    def __init__(self, config: AgentConfig, scene: sapien.Scene, control_freq: int):
        self._config = copy.deepcopy(config)
        self._scene = scene
        self._control_freq = control_freq
        self._initialize_robot()
        self._initialize_controllers()
        self._initialize_cameras()
        self._initialize_sensors()

        self.finger1_link: sapien.LinkBase = get_actor_by_name(
            self._robot.get_links(), "left_inner_finger_pad"
        )
        self.finger2_link: sapien.LinkBase = get_actor_by_name(
            self._robot.get_links(), "right_inner_finger_pad"
        )
        self.fingers_visid = []
        for visual_body in self.finger1_link.get_visual_bodies():
            self.fingers_visid.append(visual_body.get_visual_id())
        for visual_body in self.finger2_link.get_visual_bodies():
            self.fingers_visid.append(visual_body.get_visual_id())
        self.finger_size = (0.03, 0.07, 0.0075)  # values from URDF
        self.grasp_site: sapien.Link = get_entity_by_name(
            self._robot.get_links(), "grasp_convenient_link"
        )

    def _initialize_robot(self):
        loader = self._scene.create_urdf_loader()
        loader.fix_root_link = self._config.fix_root_link

        urdf_file = DESCRIPTION_DIR / self._config.urdf_file
        urdf_config = parse_urdf_config(self._config.urdf_config, self._scene)
        self._robot = loader.load(str(urdf_file), urdf_config)
        self._robot.set_name(self._config.name)
        self._control_mode = self._config.default_control_mode

    def _initialize_cameras(self):
        self._cameras = OrderedDict()
        for config in self._config.cameras:
            config = MountedCameraConfig(**config)
            if config.name in self._cameras:
                raise KeyError("Non-unique camera name: {}".format(config.name))
            cam = create_mounted_camera(config, self._robot, self._scene)
            self._cameras[config.name] = cam

    def _initialize_sensors(self):
        self._sensors = OrderedDict()
        for config in self._config.sensors:
            config = MountedActiveLightSensorConfig(**config)
            if config.name in self._sensors:
                raise KeyError(
                    "Non-unique active light sensor name: {}".format(config.name)
                )
            sensor = create_mounted_sensor(config, self._robot, self._scene)
            self._sensors[config.name] = sensor

    def _initialize_controllers(self):
        self._combined_controllers = OrderedDict()
        for control_mode, controller_configs in self._config.controllers.items():
            self._combined_controllers[control_mode] = CombinedController(
                controller_configs, self._robot, self._control_freq
            )
        self._control_mode = self._config.default_control_mode
        self._combined_controllers[self._control_mode].set_joint_drive_property()

    def sample_ee_coords(self, num_sample=10) -> np.ndarray:
        """Uniformly sample points on the two finger meshes. Used for dense reward computation
        return: ee_coords (2, num_sample, 3)"""
        finger_points = (
            np.arange(num_sample) / (num_sample - 1) - 0.5
        ) * self.finger_size[1]
        finger_points = np.stack(
            [np.zeros(num_sample), finger_points, np.zeros(num_sample)], axis=1
        )  # (num_sample, 3)

        finger1_points = transform_points(
            self.finger1_link.get_pose().to_transformation_matrix(), finger_points
        )
        finger2_points = transform_points(
            self.finger2_link.get_pose().to_transformation_matrix(), finger_points
        )

        ee_coords = np.stack((finger1_points, finger2_points))

        return ee_coords

    def get_images(
        self,
        rgb=True,
        depth=False,
        visual_seg=False,
        actor_seg=False,
        irsensor_mode=False,
    ) -> Dict[str, Dict[str, np.ndarray]]:
        assert self._cameras
        all_images = OrderedDict()
        # self._scene.update_render()
        if not irsensor_mode:
            for cam_name, cam in self._cameras.items():
                cam.take_picture()
                cam_images = get_camera_images(
                    cam,
                    rgb=rgb,
                    depth=depth,
                    visual_seg=visual_seg,
                    actor_seg=actor_seg,
                )
                cam_images["camera_intrinsic"] = cam.get_intrinsic_matrix()
                cam_extrinsic_world_frame = cam.get_extrinsic_matrix()
                robot_base_frame = (
                    self._robot.get_root_pose().to_transformation_matrix()
                )  # robot base -> world
                cam_images["camera_extrinsic_base_frame"] = (
                    cam_extrinsic_world_frame @ robot_base_frame
                )  # robot base -> camera
                all_images[cam_name] = cam_images
        else:
            for sensor_name, sensor in self._sensors.items():
                sensor_dict = sensor.get_image_dict()
                # sensor_seg_dict = get_camera_images(
                #     sensor._cam_rgb,
                #     rgb=False,
                #     depth=False,
                #     visual_seg=visual_seg,
                #     actor_seg=actor_seg,
                # )
                # sensor_dict.update(sensor_seg_dict)
                sensor_dict["camera_intrinsic"] = sensor._cam_rgb.get_intrinsic_matrix()
                cam_extrinsic_world_frame = sensor._cam_rgb.get_extrinsic_matrix()
                robot_base_frame = (
                    self._robot.get_root_pose().to_transformation_matrix()
                )  # robot base -> world
                sensor_dict["camera_extrinsic_base_frame"] = (
                    cam_extrinsic_world_frame @ robot_base_frame
                )  # robot base -> camera
                all_images[sensor_name] = sensor_dict

        return all_images

    def get_proprioception(self):
        state_dict = OrderedDict()
        qpos = self._robot.get_qpos()
        qvel = self._robot.get_qvel()

        state_dict["qpos"] = qpos
        state_dict["qvel"] = qvel
        return state_dict

    @classmethod
    def from_config_file(cls, config_path: str, scene: sapien.Scene, control_freq: int):
        with open(config_path, "r") as f:
            config_dict = yaml.safe_load(f)
        config = AgentConfig(**config_dict)
        return cls(config, scene, control_freq)

    @property
    def control_mode(self) -> str:
        return self._control_mode

    def set_control_mode(self, control_mode: str):
        self._control_mode = control_mode

    @property
    def action_space(self):
        return spaces.Dict(
            {
                mode: spaces.Box(
                    controller.action_range[:, 0], controller.action_range[:, 1]
                )
                for mode, controller in self._combined_controllers.items()
            }
        )

    @property
    def action_range(self) -> spaces.Box:
        return spaces.Box(
            self._combined_controllers[self._control_mode].action_range[:, 0],
            self._combined_controllers[self._control_mode].action_range[:, 1],
        )

arti_mani/agents/floating_robotiq.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from collections import OrderedDict

import numpy as np
import sapien.core as sapien
from arti_mani.agents.base_agent import BaseAgent
from arti_mani.utils.common import compute_angle_between
from arti_mani.utils.sapien_utils import (
    check_joint_stuck,
    get_actor_by_name,
    get_entity_by_name,
    get_pairwise_contact_impulse,
)
from sapien.core import Pose


class Panda(BaseAgent):
    def __init__(self, *args, **kwargs):
        super(Panda, self).__init__(*args, **kwargs)
        self.finger1_link, self.finger2_link = get_actor_by_name(
            self._robot.get_links(), ["panda_leftfinger", "panda_rightfinger"]
        )
        self.finger1_joint, self.finger2_joint = get_actor_by_name(
            self._robot.get_joints(), ["panda_finger_joint1", "panda_finger_joint2"]
        )
        self.grasp_site: sapien.Link = get_entity_by_name(
            self._robot.get_links(), "grasp_site"
        )

    def get_proprioception(self):
        state_dict = OrderedDict()
        qpos = self._robot.get_qpos()
        qvel = self._robot.get_qvel()

        state_dict["qpos"] = qpos
        state_dict["qvel"] = qvel
        state_dict["tcp_wrench"] = self.get_tcp_wrench().astype(np.float32)
        state_dict["joint_external_torque"] = self.get_generalized_external_forces()
        state_dict["gripper_grasp"] = np.array(
            [self.check_gripper_grasp_real()], dtype=np.float32
        )

        return state_dict

    def check_grasp(self, actor: sapien.ActorBase, min_impulse=1e-6, max_angle=85):
        assert isinstance(actor, sapien.ActorBase), type(actor)
        contacts = self._scene.get_contacts()

        limpulse = get_pairwise_contact_impulse(contacts, self.finger1_link, actor)
        rimpulse = get_pairwise_contact_impulse(contacts, self.finger2_link, actor)

        # direction to open the gripper
        ldirection = self.finger1_link.pose.to_transformation_matrix()[:3, 1]
        rdirection = -self.finger2_link.pose.to_transformation_matrix()[:3, 1]

        # angle between impulse and open direction
        langle = compute_angle_between(ldirection, limpulse)
        rangle = compute_angle_between(rdirection, rimpulse)

        lflag = (
            np.linalg.norm(limpulse) >= min_impulse and np.rad2deg(langle) <= max_angle
        )
        rflag = (
            np.linalg.norm(rimpulse) >= min_impulse and np.rad2deg(rangle) <= max_angle
        )

        return all([lflag, rflag])

    def check_contact_fingers(self, actor: sapien.ActorBase, min_impulse=1e-6):
        assert isinstance(actor, sapien.ActorBase), type(actor)
        contacts = self._scene.get_contacts()

        limpulse = get_pairwise_contact_impulse(contacts, self.finger1_link, actor)
        rimpulse = get_pairwise_contact_impulse(contacts, self.finger2_link, actor)

        return (
            np.linalg.norm(limpulse) >= min_impulse,
            np.linalg.norm(rimpulse) >= min_impulse,
        )

    def sample_ee_coords(self, num_sample=10) -> np.ndarray:
        """Uniformly sample points on the two finger meshes. Used for dense reward computation
        return: ee_coords (2, num_sample, 3)"""
        l = 0.0355
        r = 0.052
        ret = []
        for i in range(num_sample):
            x = (l * i + (4 - i) * r) / 4
            finger_tips = [
                self.finger2_joint.get_global_pose().transform(Pose([0, x, 0])).p,
                self.finger1_joint.get_global_pose().transform(Pose([0, -x, 0])).p,
            ]
            ret.append(finger_tips)
        return np.array(ret).transpose((1, 0, 2))

    def get_tcp_wrench(self):
        joint_tau = self.get_generalized_external_forces()[:7]
        controller = self._combined_controllers[self._control_mode]._controllers[0]
        assert controller.control_joint_names == [
            "panda_joint1",
            "panda_joint2",
            "panda_joint3",
            "panda_joint4",
            "panda_joint5",
            "panda_joint6",
            "panda_joint7",
        ]
        controller._sync_articulation()
        J_full = controller.controller_articulation.compute_world_cartesian_jacobian()[
            -6:
        ]
        J = np.linalg.pinv(J_full.T)

        TCP_wrench = J @ joint_tau
        return TCP_wrench

    @staticmethod
    def build_grasp_pose(forward, flat, center):
        extra = np.cross(flat, forward)
        ans = np.eye(4)
        ans[:3, :3] = np.array([extra, flat, forward]).T
        ans[:3, 3] = center
        return Pose.from_transformation_matrix(ans)

    def check_gripper_grasp_real(self) -> bool:
        """check whether the gripper is grasping something by checking the joint position and velocity"""
        from arti_mani.agents.controllers import GripperPDJointPosMimicController

        assert isinstance(
            self._combined_controllers[self._control_mode]._controllers[1],
            GripperPDJointPosMimicController,
        )
        for joint_idx, joint in enumerate(self._robot.get_active_joints()):
            if joint.name == "panda_finger_joint1":
                active_joint1_idx = joint_idx
            if joint.name == "panda_finger_joint2":
                active_joint2_idx = joint_idx

        joint1_stuck = check_joint_stuck(self._robot, active_joint1_idx)
        joint2_stuck = check_joint_stuck(self._robot, active_joint2_idx)

        return joint1_stuck or joint2_stuck

arti_mani/agents/panda.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import logging
import os
import random
import time
from typing import Optional

import numpy as np
import torch
import yaml
from arti_mani import ALG_DIR
from path import Path


def set_seed(seed=None):
    # type: (Optional[int]) -> int
    """
    set the random seed using the required value (`seed`)
    or a random value if `seed` is `None`
    :return: the newly set seed
    """
    if seed is None:
        seed = random.randint(1, 10000)
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    return seed


class Logger:
    def __init__(self, path, clevel=logging.DEBUG, Flevel=logging.INFO):
        self.logger = logging.getLogger(path)
        self.logger.setLevel(logging.DEBUG)
        fmt = logging.Formatter(
            "[%(asctime)s] [%(levelname)s] %(message)s", "%Y%m%d_%H:%M:%S"
        )
        # 设置CMD日志
        sh = logging.StreamHandler()
        sh.setFormatter(fmt)
        sh.setLevel(clevel)
        # 设置文件日志
        fh = logging.FileHandler(path)
        fh.setFormatter(fmt)
        fh.setLevel(Flevel)
        self.logger.addHandler(sh)
        self.logger.addHandler(fh)

    def debug(self, message):
        self.logger.debug(message)

    def info(self, message):
        self.logger.info(message)

    def war(self, message):
        self.logger.warning(message)

    def error(self, message):
        self.logger.error(message)

    def cri(self, message):
        self.logger.critical(message)


class Config(object):
    def __init__(
        self,
        conf_file_path=None,
        seed=None,
        exp_name="pretrain_unet",
        mode="train",
        log_name=None,
        log=True,
    ):
        assert mode in ["train", "test"]
        self.log_each_step = log

        # if the configuration file is not specified
        # try to load a configuration file based on the experiment name
        tmp = Path(__file__).parent / (exp_name + ".yaml")
        if conf_file_path is None and tmp.exists():
            conf_file_path = tmp

        # read the YAML configuation file
        if conf_file_path is None:
            cfg = {}
        else:
            conf_file = open(conf_file_path, "r")
            cfg = yaml.safe_load(conf_file)

        # read configuration parameters from YAML file
        # or set their default value
        self.seed = cfg.get("seed", set_seed(seed))  # type: int
        self.hmap_h = cfg.get("img_h", 72)  # type: int
        self.hmap_w = cfg.get("img_w", 128)  # type: int
        self.num_classes = cfg.get("num_classes", 2)  # type: int
        self.lr = cfg.get("lr", 0.0001)  # type: float
        self.epochs = cfg.get("epochs", 1000)  # type: int
        self.n_workers = cfg.get("n_workers", 8)  # type: int
        self.batch_size = cfg.get("batch_size", 1)  # type: int
        self.train_ratio = cfg.get("train_ratio", 0.7)  # type: float
        self.val_ratio = cfg.get("val_ratio", 0.1)  # type: float
        self.test_ratio = cfg.get("test_ratio", 0.2)  # type: float
        self.loss = cfg.get("loss", None)  # type: str
        self.exp_suffix = cfg.get("exp_suffix", None)  # type: str
        self.data_path = cfg.get("data_path", None)  # type: str
        self.segnet_config = cfg.get("segnet_config", {})  # type: dict
        self.smp_config = cfg.get("smp_config", {})  # type: dict

        if cfg.get("device", None) is not None and cfg["device"] != "cpu":
            os.environ["CUDA_VISIBLE_DEVICES"] = str(
                cfg.get("device").split(":")[1]
            )  # type: str
            self.device = "cuda:0"
        elif cfg.get("device", None) is not None and cfg["device"] == "cpu":
            self.device = "cpu"
        else:
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        assert self.exp_suffix is not None, "set exp_suffix in config yaml! "

        assert os.path.exists(
            ALG_DIR / self.data_path
        ), "the specified directory for the Dataset does not exists"

        # define and build output path
        if log_name is None:
            self.exp_log_path = f"log/{exp_name}/{time.strftime('%Y%m%d_%H%M%S', time.localtime())}_{self.exp_suffix}"
        else:
            self.exp_log_path = f"log/{exp_name}/{log_name}"
        if not os.path.exists(self.exp_log_path):
            os.makedirs(self.exp_log_path)

        # set logger
        self.logger = Logger(f"{self.exp_log_path}/{mode}_log.txt")

        self.logger.info(f"Config:\n{cfg}")

arti_mani/algorithms/config/config.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import random

import cv2
import gym
import matplotlib.pyplot as plt
import numpy as np
from arti_mani import ASSET_DIR
from arti_mani.utils.cv_utils import visualize_depth, visualize_seg
from arti_mani.utils.sapien_utils import look_at
from PIL import Image


def gen_grid_campos(dist_xy_xz_num, length, width, height, mode):
    dist_limit = np.sqrt(width**2 + length**2) / 2
    dist_num, rotxy_num, rotxz_num = dist_xy_xz_num

    if mode == "faucet":
        # get more data for faucet
        dist_range = np.linspace(
            dist_limit + 0.125, dist_limit + 0.6, int(1.5 * dist_num)
        )  # 0.475
        rot_xy_range = np.linspace(
            45 / 180 * np.pi, 135 / 180 * np.pi, int(1.5 * rotxy_num)
        )  # 90
        rot_xz_range = np.linspace(
            -60 / 180 * np.pi, 60 / 180 * np.pi, int(1.5 * rotxz_num)
        )  # 120
    else:
        dist_range = np.linspace(dist_limit + 0.2, dist_limit + 0.6, dist_num)  # 0.4
        rot_xy_range = np.linspace(-10 / 180 * np.pi, 30 / 180 * np.pi, rotxy_num)  # 40
        rot_xz_range = np.linspace(-30 / 180 * np.pi, 30 / 180 * np.pi, rotxz_num)  # 60

    grid_sample = np.array(np.meshgrid(dist_range, rot_xy_range, rot_xz_range)).reshape(
        3, -1
    )
    ori_sample = grid_sample.copy()

    # add some noise to lookat angle
    grid_sample[-2:] += (
        np.random.normal(0, 2, size=grid_sample[-2:].shape) / 180 * np.pi
    )
    print(np.abs(grid_sample[-2:] - ori_sample[-2:]).mean() * 180 / np.pi)

    sample_pos = np.stack(
        (
            -grid_sample[0] * np.cos(grid_sample[1]) * np.cos(grid_sample[2]),
            grid_sample[0] * np.cos(grid_sample[1]) * np.sin(grid_sample[2]),
            grid_sample[0] * np.sin(grid_sample[1]) + height / 2,
        ),
        axis=1,
    )
    sample_campos = sample_pos[sample_pos[:, 2] > 0]
    return sample_campos


def random_crop(image, crop_shape):
    nh = random.randint(0, image.shape[0] - crop_shape[0])
    nw = random.randint(0, image.shape[1] - crop_shape[1])
    image_crop = image[nh : nh + crop_shape[0], nw : nw + crop_shape[1]]
    return image_crop


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    data_mode = "seg"  # "seg" "keypoints"
    assert data_mode in ["seg", "keypoints"]
    save_mode = True
    ir_mode = True
    domain_random = True
    random_bg_mode = False
    render_mode = False
    H, W = 144, 256
    dist_xy_xz_num = (8, 8, 12)
    num_classes = 6
    data_dir_name = f"artidata_rebalance_{dist_xy_xz_num}_{data_mode}"
    other_handle_visible = True if data_mode == "seg" else False
    print("data_dir_name:", data_dir_name)

    if domain_random:
        data_dir_name = data_dir_name + "_DR"
    else:
        data_dir_name = data_dir_name + "_noDR"

    if random_bg_mode:
        data_dir_name = data_dir_name + "_randombg"
        random_bg_num = 180
        bg_rgbs = []
        bg_ds = []
        for ind in range(random_bg_num):
            rgb_data = np.asarray(Image.open(ASSET_DIR / f"random_bg/{ind:03}.jpg"))
            d_data = np.asarray(Image.open(ASSET_DIR / f"random_bg/{ind:03}.png"))
            # rgb_data = rgb_data / 255.0
            d_data = (d_data.astype(np.float32) / 1000.0).clip(0.2, 2)
            bg_rgbs.append(rgb_data)
            bg_ds.append(d_data)
    else:
        data_dir_name = data_dir_name + "_norandombg"

    mode_list = ["door", "drawer", "faucet"]
    arti_items = {
        "door": [
            0,
            1000,
            1001,
            1002,
            1006,
            1007,
            1014,
            1017,
            1018,
            1025,
            1026,
            1027,
            1030,
            1031,
            1034,
            1036,
            1038,
            1039,
            1041,
            1042,
            1044,
            1045,
            1047,
            1049,
            1051,
            1052,
            1054,
            1057,
            1060,
            1061,
            1062,
            1063,
            1064,
            1065,
            1067,
            1073,
            1075,
            1077,
            1078,
            1081,
            # (1028, 1046, 1068)
        ],
        # 40 + 3
        "drawer": [
            1,
            1004,
            1005,
            1013,
            1016,
            1021,
            1024,
            1032,
            1033,
            1035,
            1040,
            1056,
            1066,
            1076,
            1079,
            1082,
        ],
        # 16,
        "faucet": [
            5002,
            5004,
            5005,
            5007,
            5018,
            5023,
            5024,
            5028,
            5034,
            5049,
            5052,
            5053,
            5063,
            5069,
        ],
        # 14 + 1 (60)
    }

    train_list = [
        0,
        1,
        1000,
        1001,
        1004,
        1006,
        1013,
        1014,
        1016,
        1018,
        1024,
        1025,
        1026,
        1031,
        1032,
        1034,
        1035,
        1038,
        1039,
        1040,
        1041,
        1042,
        1044,
        1047,
        1049,
        1051,
        1052,
        1057,
        1060,
        1061,
        1062,
        1064,
        1065,
        1066,
        1073,
        1076,
        1077,
        1078,
        1082,
        5002,
        5004,
        5005,
        5018,
        5023,
        5024,
        5034,
        5049,
        5052,
        5053,
        5063,
    ]  # 39+11

    val_list = [
        1002,
        1005,
        1007,
        1017,
        1021,
        1027,
        1030,
        1033,
        1036,
        1045,
        1054,
        1056,
        1063,
        1067,
        1075,
        1079,
        1081,
        5007,
        5028,
        5069,
    ]  # 17+3

    if data_mode == "keypoints":
        train_list.remove(1035)
        val_list.remove(1056)
        arti_items["drawer"].remove(1035)
        arti_items["drawer"].remove(1056)
        plt_x, plt_y = 2, 2
    else:
        plt_x, plt_y = 1, 3

    obs_mode = "state_egorgbd"
    control_mode = "pd_joint_delta_pos"
    for mode in mode_list:
        if mode == "door":
            env_id = "CabinetDoor-v0"
        elif mode == "drawer":
            env_id = "CabinetDrawer-v0"
        elif mode == "faucet":
            env_id = "Faucet-v0"
        else:
            raise NotImplementedError(mode)
        env = gym.make(
            env_id,
            articulation_ids=[arti_items[mode][0]],
            articulation_config_path=None,
            obs_mode=obs_mode,
            control_mode=control_mode,
            reward_mode="dense",
            irsensor_mode=ir_mode,
            other_handle_visible=other_handle_visible,
            domain_random=domain_random,
        )
        # random.shuffle(arti_items[mode])
        for arti_ind in arti_items[mode]:
            render_rgbd_path = f"./{data_dir_name}/{mode}/{arti_ind}/"
            print(render_rgbd_path)
            if not os.path.exists(render_rgbd_path):
                os.makedirs(render_rgbd_path)

            obs_global = env.reset(articulation_id=arti_ind)
            if arti_ind == 0:
                width, length, height = 0.304, 0.751, 0.54
                obj_center_pos = np.array([0, 0, 0.27])
            elif arti_ind == 1:
                width, length, height = 0.474, 0.28, 0.453
                obj_center_pos = np.array([0, 0, 0.23])
            else:
                obj_center_pos = env.unwrapped._articulation.get_root_pose().p
                bbox_max, bbox_min = (
                    np.array(env.unwrapped._articulation_config["bbox_max"]),
                    np.array(env.unwrapped._articulation_config["bbox_min"]),
                )
                scale = env.unwrapped._articulation_config["scale"]
                width, length, height = (bbox_max - bbox_min) * scale

            # sample camera position
            sample_campos = gen_grid_campos(dist_xy_xz_num, length, width, height, mode)

            rgbd_seg_files = {"rgb": [], "depth": [], "seg": []}
            part_keypoints = []
            for ind, pos_it in enumerate(sample_campos):
                gripper_qpos = np.random.uniform(0, 0.068)
                env.unwrapped._agent._robot.set_qpos([gripper_qpos, gripper_qpos])
                env.unwrapped._initialize_articulations()

                # door and draw have a 0.6 offset on x-axis
                if mode != "faucet":
                    # random lookat center
                    sample_poscenter = obj_center_pos + np.hstack(
                        (
                            np.random.normal(loc=0, scale=width / 12, size=1),
                            np.random.normal(loc=0, scale=length / 10, size=1),
                            np.random.normal(loc=0, scale=height / 10, size=1),
                        )
                    )
                    pos_it[0] += 0.6
                    sample_poscenter[0] += 0.6
                else:
                    # random lookat center
                    sample_poscenter = obj_center_pos + np.hstack(
                        (
                            np.random.normal(loc=0, scale=0.05, size=1),
                            np.random.normal(loc=0, scale=0.05, size=1),
                            np.random.normal(loc=0, scale=0.01, size=1),
                        )
                    )

                # random sample camera in-plane rotation
                sample_gripper_roty = (np.random.rand() - 0.5) * 2  # (-1, 1)
                sample_gripper_rotz = np.sqrt(
                    1 - sample_gripper_roty**2
                ) * np.random.choice([-1, 1])
                env.unwrapped._agent._robot.set_root_pose(
                    look_at(
                        eye=pos_it,
                        target=sample_poscenter,
                        up=[0, sample_gripper_roty, sample_gripper_rotz],
                    )
                )
                obs_global = env.get_obs()
                rgb = obs_global["rgb"]
                clean_depth = obs_global["clean_depth"]
                depth = obs_global["depth"]
                seg = obs_global["seg"]
                bg_seg = obs_global["bg_seg"]
                if random_bg_mode:
                    random_id = np.random.choice(range(random_bg_num))
                    bg_rgb_sample, bg_d_sample = bg_rgbs[random_id], bg_ds[random_id]
                    bg_rgb_sample_crop = random_crop(bg_rgb_sample, (H, W))
                    bg_d_sample_crop = random_crop(bg_d_sample, (H, W))
                    bg_seg_mask = bg_seg.astype(np.bool_)
                    rgb = (
                        rgb * ~bg_seg_mask
                        + bg_rgb_sample_crop.transpose(2, 0, 1) * bg_seg_mask
                    )
                    clean_depth = (
                        clean_depth * ~bg_seg_mask + bg_d_sample_crop * bg_seg_mask
                    )
                    depth = depth * ~bg_seg_mask + bg_d_sample_crop * bg_seg_mask

                if data_mode == "seg":
                    rgbd_seg_files["rgb"].append(rgb)
                    rgbd_seg_files["depth"].append(depth)
                    rgbd_seg_files["seg"].append(seg)
                elif data_mode == "keypoints":
                    kpts = obs_global["kpts"]
                    uvz = obs_global["uvz"]
                    uvz_visable = obs_global["uvz_visable"]  # (3, 3)
                    print(
                        "+++ tip camxyz, uvz, visable: ",
                        kpts[0],
                        uvz[0],
                        uvz_visable[0],
                    )
                    print(
                        "+++ center camxyz, uvz, visable: ",
                        kpts[1],
                        uvz[1],
                        uvz_visable[1],
                    )
                    print(
                        "+++ bottom camxyz, uvz, visable: ",
                        kpts[2],
                        uvz[2],
                        uvz_visable[2],
                    )
                    if not np.any(uvz_visable[:, 2]):
                        continue
                    part_keypoints.append(
                        np.stack((kpts, uvz, uvz_visable))
                    )  # (3, 3, 3)

                    rgbd_seg_files["rgb"].append(rgb)
                    rgbd_seg_files["depth"].append(depth)

                    plot_rgb = rgb.transpose((1, 2, 0)).copy()
                    for idn in range(uvz.shape[0]):
                        if uvz_visable[idn, 0] and uvz_visable[idn, 1]:
                            if uvz_visable[idn, 2]:
                                cv2.circle(
                                    plot_rgb,
                                    (int(uvz[idn, 0]), int(uvz[idn, 1])),
                                    radius=2,
                                    color=(255, 0, 0),
                                    thickness=-1,
                                )
                            else:
                                cv2.circle(
                                    plot_rgb,
                                    (int(uvz[idn, 0]), int(uvz[idn, 1])),
                                    radius=2,
                                    color=(0, 255, 0),
                                    thickness=-1,
                                )
                else:
                    raise NotImplementedError(data_mode)

                plt.figure()
                plt.subplot(plt_x, plt_y, 1)
                plt.imshow(rgb.transpose((1, 2, 0)) / 255.0)
                plt.title("rgb")
                plt.axis("off")
                plt.subplot(plt_x, plt_y, 2)
                plt.imshow(visualize_depth(depth.squeeze())[..., ::-1])
                plt.title("depth")
                plt.axis("off")
                plt.subplot(plt_x, plt_y, 3)
                plt.imshow(visualize_seg(seg, num_classes))
                plt.title("seg")
                plt.axis("off")
                if data_mode == "keypoints":
                    plt.subplot(plt_x, plt_y, 4)
                    plt.imshow(plot_rgb)
                    plt.title("keypoints")
                    plt.axis("off")
                # plt.tight_layout()
                if save_mode:
                    plt.savefig(
                        f"./{render_rgbd_path}/{arti_ind}_{ind}.jpg",
                        bbox_inches="tight",
                        pad_inches=0.0,
                    )
                else:
                    plt.show()
                plt.close()

                if render_mode:
                    viewer = env.render()
                    print("Press [e] to start")
                    while True:
                        if viewer.window.key_down("e"):
                            break
                        env.render()
                print(mode, arti_ind, ind)
            if save_mode:
                if arti_ind in train_list:
                    data_rgbd_path = f"./{data_dir_name}/train/"
                    data_kpts_path = f"./{data_dir_name}/train_kpts/"
                elif arti_ind in val_list:
                    data_rgbd_path = f"./{data_dir_name}/val/"
                    data_kpts_path = f"./{data_dir_name}/val_kpts/"
                else:
                    raise NotImplementedError
                if not os.path.exists(data_rgbd_path):
                    os.makedirs(data_rgbd_path)
                if not os.path.exists(data_kpts_path):
                    os.makedirs(data_kpts_path)

                np.savez(data_rgbd_path + f"{arti_ind}.npz", **rgbd_seg_files)
                if data_mode == "keypoints":
                    kpts_data = np.array(part_keypoints)
                    np.save(data_kpts_path + f"{arti_ind}_keypoints.npy", kpts_data)

arti_mani/algorithms/data_process/gen_rgbd_data.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from collections import deque

import numpy as np
import torch
from arti_mani.utils.o3d_utils import pcd_uni_down_sample_with_crop
from scipy.spatial.distance import cdist


class UniDownsampler:
    def __init__(self, sample_num: int):
        self.sample_num = sample_num
        self.sample_indices = None

    def sampling(
        self,
        raw_xyz: np.ndarray,  # (N, 3), H*W=N
        episode_rng,  # RandomState from env
    ):
        ## crop the pts & uniform downsample to num pts
        sample_indices = pcd_uni_down_sample_with_crop(
            raw_xyz,
            self.sample_num,
            min_bound=np.array([-0.5, -1, 1e-3]),
            max_bound=np.array([2, 2, 2]),
        )
        indices_len = len(sample_indices)
        if indices_len < 1000:
            sample_indices.extend(
                episode_rng.choice(raw_xyz.shape[0], self.sample_num - indices_len)
            )
        self.sample_indices = sample_indices
        return sample_indices


class ScoreSampler:
    def __init__(self, sample_num: int):
        self.sample_num = sample_num
        self.sample_indices = None

    def sampling(
        self,
        seg_map: torch.Tensor,  # (C, H, W)
    ):
        num_classes = seg_map.shape[0]
        score_map = seg_map.view(num_classes, -1)  # (C, N)
        _, indices = torch.topk(score_map, self.sample_num, dim=1)  # (C, sample_num)
        pts_index = indices.detach().cpu().numpy().reshape(-1)  # (C*sample_num)
        self.sample_indices = pts_index
        return pts_index


def farthest_point_sample(xyz, npoint, episode_rng):
    N, C = xyz.shape
    centroids = np.zeros((npoint,), dtype=np.int_)
    distance = np.ones((N,)) * 1e10
    farthest = episode_rng.randint(0, N)

    for i in range(npoint):
        centroids[i] = farthest
        centroid = xyz[farthest, :]
        dist = np.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = np.argmax(distance, -1)
    return centroids


class FPSSampler:
    def __init__(self, sample_num: int):
        self.sample_num = sample_num
        self.sample_indices = None

    def sampling(
        self,
        raw_xyz: np.ndarray,
        num_classes: int,
        seg_mc_map: np.ndarray,  # (H*W)
        episode_rng,  # RandomState from env
    ):
        part_inds = []
        for ind in range(num_classes):
            one_indices = np.where(seg_mc_map == ind)[0]
            ones_num = one_indices.shape[0]
            if ones_num == 0:
                cur_indices = np.zeros(self.sample_num, dtype=np.int_)
            elif ones_num < self.sample_num:
                repeat_times = self.sample_num // ones_num
                remain_num = self.sample_num - repeat_times * ones_num
                cur_indices = np.concatenate(
                    (one_indices.repeat(repeat_times), one_indices[:remain_num])
                )
            else:
                part_xyz = raw_xyz[one_indices]
                rel_indices = farthest_point_sample(
                    part_xyz, self.sample_num, episode_rng
                )
                cur_indices = one_indices[rel_indices]
            part_inds.append(cur_indices)
        pts_index = np.concatenate(part_inds)  # (C*sample_num)
        self.sample_indices = pts_index
        return pts_index


class RandomSampler:
    def __init__(self, sample_num: int):
        self.sample_num = sample_num
        self.sample_indices = None

    def sampling(
        self,
        seg_map_mc: np.ndarray,  # (H*W)
        num_classes: int,
        episode_rng,  # RandomState from env
    ):
        part_inds = []
        for ind in range(num_classes):
            one_indices = np.where(seg_map_mc == ind)[0]
            ones_num = one_indices.shape[0]
            if ones_num == 0:
                cur_indices = np.zeros(self.sample_num, dtype=np.int_)
            elif ones_num < self.sample_num:
                repeat_times = self.sample_num // ones_num
                remain_num = self.sample_num - repeat_times * ones_num
                cur_indices = np.concatenate(
                    (one_indices.repeat(repeat_times), one_indices[:remain_num])
                )
            else:
                cur_indices = episode_rng.choice(
                    one_indices, self.sample_num, replace=False
                )
            part_inds.append(cur_indices)
        pts_index = np.concatenate(part_inds)  # (C*sample_num)
        self.sample_indices = pts_index
        return pts_index


class FCUWSampler:
    """Frame-Consistent Uncertainty-aware Weighted Sampler"""

    def __init__(self, frame_num: int, sample_num: int):
        self.partpts_q = deque(maxlen=frame_num)
        self.frame_num = frame_num
        self.sample_num = sample_num
        self.sample_indices = None

    def sampling(
        self,
        raw_xyz: np.ndarray,
        num_classes: int,
        seg_map_mc: np.ndarray,  # (H*W)
        uncertainty_map: np.ndarray,
        episode_rng,  # RandomState from env
    ):
        # (H*W)
        uncertainty_indices = uncertainty_map.reshape((-1,))
        weight_map = np.zeros((num_classes, seg_map_mc.shape[0]))
        pts_index = []
        for ind in range(num_classes):
            part_indices = np.where(seg_map_mc == ind)[0]
            part_pts_num = part_indices.shape[0]
            # (H*W, ) => (PN, ), (0, 1.7)
            if part_pts_num == 0:
                pts_index += np.zeros(self.sample_num, dtype=np.int_).tolist()
            elif part_pts_num < self.sample_num:
                repeat_times = self.sample_num // part_pts_num
                remain_num = self.sample_num - repeat_times * part_pts_num
                pts_index += np.concatenate(
                    (part_indices.repeat(repeat_times), part_indices[:remain_num])
                ).tolist()
            else:
                part_xyz = raw_xyz[part_indices]
                # (PN, 3), (L*30, 3) => (PN, L*30)
                part_confidence = np.exp(
                    -uncertainty_indices[part_indices]
                )  # w_uncertainty
                if len(self.partpts_q) != 0:
                    cache_partpts = (
                        np.array(self.partpts_q)
                        .transpose((1, 0, 2, 3))
                        .reshape((num_classes, -1, 3))
                    )
                    part_dist = cdist(part_xyz, cache_partpts[ind], metric="euclidean")
                    w_frameconsist = 2 ** (-40 * (part_dist.min(axis=1)))
                    part_confidence = w_frameconsist * part_confidence
                norm_w = part_confidence / np.sum(part_confidence)
                weight_map[ind, part_indices] = norm_w
                pts_index += episode_rng.choice(
                    part_indices, self.sample_num, replace=True, p=norm_w
                ).tolist()
        # (C*sample_num, 3)
        sample_xyz = raw_xyz[pts_index]
        self.partpts_q.append(
            sample_xyz.reshape(num_classes, self.sample_num, 3)
        )  # (C, sample_num, 3)
        self.sample_indices = pts_index
        return pts_index, weight_map

arti_mani/algorithms/rl_iam/Sampling.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import random
import time

import gym
import segmentation_models_pytorch as smp
import torch
import torch.nn.functional as F
import yaml
from arti_mani import VISUALMODEL_DIR
from arti_mani.algorithms.visual_net.Networks.conv2dnet import (
    Attention2DUNet,
    AttentionDepthUNet,
    SiameseNet,
    UNetEncoder,
)
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import CustomUnet
from arti_mani.algorithms.visual_net.Networks.img_encoder import ImgEncoder
from arti_mani.algorithms.visual_net.Networks.pointnet import (
    PointNetDenseCls,
    PointNetfeat,
)
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor


class CustomStateExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space: gym.Space):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomStateExtractor, self).__init__(observation_space, features_dim=1)

        self._features_dim = 222

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]  # (N, 9)
        ee_pos_base = observations["ee_pos_base"]  # (N, 3)
        bs = qpos.shape[0]
        ee_xyz = observations["eepad_pts"].view(bs, -1)  # (N, 40*3)
        target_xyz = observations["target_pcd_ee"].view(bs, -1)  # (N, 30*3)

        state = torch.cat([qpos, ee_pos_base, ee_xyz, target_xyz], dim=1)

        return state


class CustomProcessStateExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space: gym.Space, mode="robot_xyz"):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomProcessStateExtractor, self).__init__(
            observation_space, features_dim=1
        )
        support_modes = ["robot_xyz", "cam_xyz", "cam_zvu"]
        assert mode in support_modes, f"only support modes: {support_modes}"
        print("training mode： ", mode)

        self.mode = mode
        self.base_cam_intrin = torch.tensor(
            [[45.5882, 0.0, 31.7436], [0.0, 45.547, 17.6573], [0.0, 0.0, 1.0]],
            requires_grad=True,
            device="cuda:0",
        )
        self.base_cam_extrin = torch.tensor(
            [
                [0.4034, -0.915, -0.0018, -0.6329],
                [-0.4556, -0.1991, -0.8676, 0.3805],
                [0.7936, 0.3508, -0.4972, 0.6297],
                [0.0, 0.0, 0.0, 1.0],
            ],
            requires_grad=True,
            device="cuda:0",
        )

        # Update the features dim manually
        self._features_dim = 15

    def forward(self, observations) -> torch.Tensor:
        if self.mode == "robot_xyz":
            return observations
        elif self.mode == "cam_xyz":
            point_target_robot = observations[:, -3:]
            point_target_cam = (
                torch.mm(
                    self.base_cam_extrin[:3, :3],
                    point_target_robot.transpose(0, 1).contiguous(),
                )
                .transpose(0, 1)
                .contiguous()
            )
            point_target_cam += self.base_cam_extrin[:3, 3]
            return torch.cat([observations[:, :-3], point_target_cam], dim=1)
        elif self.mode == "cam_zvu":
            point_target_robot = observations[:, -3:]
            point_target_cam = (
                torch.mm(
                    self.base_cam_extrin[:3, :3],
                    point_target_robot.transpose(0, 1).contiguous(),
                )
                .transpose(0, 1)
                .contiguous()
            )
            point_target_cam += self.base_cam_extrin[:3, 3]
            uvz = (
                torch.mm(
                    self.base_cam_intrin, point_target_cam.transpose(0, 1).contiguous()
                )
                .transpose(0, 1)
                .contiguous()
            )

            base_target_uvz = torch.zeros_like(uvz)
            base_target_uvz[:, 0] = torch.round(uvz[:, 0] / uvz[:, 2]) / 64.0
            base_target_uvz[:, 1] = torch.round(uvz[:, 1] / uvz[:, 2]) / 36.0
            base_target_uvz[:, 2] = (uvz[:, 2] - 0.5) / (2 - 0.5)
            base_target_uvz = base_target_uvz[:, [2, 1, 0]]

            # Return a (B, self._features_dim) PyTorch tensor, where B is batch dimension.
            return torch.cat([observations[:, :-3], base_target_uvz], dim=1)


class CustomDepthSegExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space: gym.spaces.Dict, state_repeat_times=2):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomDepthSegExtractor, self).__init__(observation_space, features_dim=1)
        self.state_repeat_times = state_repeat_times
        self.ego_cam_intrin_inv = torch.tensor(
            [
                [0.01090477, 0.0, -0.69031547],
                [0.0, 0.01089966, -0.40202298],
                [0.0, 0.0, 1.0],
            ],
            device="cuda:0",
        )
        self.ego_cam_extrin = torch.tensor(
            [
                [-0.0127, -0.0079, -0.9999, -0.2265],
                [-0.9999, 0.0038, 0.0127, 0.0253],
                [0.0037, 1.0, -0.0079, 0.0774],
                [0.0, 0.0, 0.0, 1.0],
            ],
            device="cuda:0",
        )
        # load the Unet model
        self.attn = AttentionDepthUNet(
            depth_feat_channels=8,
            filters=[16, 16],
            kernel_sizes=[5, 5],
            strides=[2, 2],
            output_channels=2,
            norm=None,
            activation="lrelu",
            skip_connections=True,
        ).to(torch.device("cuda:0"))
        self.attn.train()

        # Update the features dim manually
        self._features_dim = 15 * self.state_repeat_times

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]
        ee_pos_base = observations["ee_pos_base"]

        depth = observations["depth"]  # (N, 1, H, W)
        N, ch, h, w = depth.shape
        depth_norm = torch.clip(depth - torch.mean(depth), -1, 1)
        attnmap = F.softmax(self.attn(depth_norm), dim=1)[
            :, 0
        ]  # (N, 2, H, W)=>(N, H, W)

        ## argmax to get maxpoint ind
        max_indice = torch.argmax(attnmap.view(N, -1), dim=1)

        # zc = torch.stack([-depth.view(N, -1)[i, max_indice[i]] for i in range(N)])  # (N, sample_num), too much time
        zc = ((-depth.view(N, -1)).gather(1, max_indice.unsqueeze(1))).squeeze(1)
        v = torch.div(max_indice, w, rounding_mode="floor")  # (N, sample_num)
        u = max_indice - v * w  # (N, sample_num)
        uvz_sample = torch.stack((u * zc, v * zc, zc), dim=1).unsqueeze(
            2
        )  # (N, 3, sample_num)
        camxyz_sample = torch.matmul(
            self.ego_cam_intrin_inv, uvz_sample
        )  # (N, 3, pts_num)
        camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, pts_num)
        eexyz_sample = (
            torch.matmul(
                camxyz_sample.transpose(1, 2).contiguous(),
                (self.ego_cam_extrin[:3, :3]).t(),
            )
            + self.ego_cam_extrin[:3, 3]
        )  # (N,sample_num,3)
        eepts_pred = eexyz_sample[:, 0]

        state = torch.cat([qpos, ee_pos_base, eepts_pred], dim=1).repeat(
            1, self.state_repeat_times
        )

        return state


class CustomRGBDSegExtractor(BaseFeaturesExtractor):
    def __init__(
        self, observation_space: gym.spaces.Dict, state_repeat_times=2, sample_num=50
    ):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomRGBDSegExtractor, self).__init__(observation_space, features_dim=1)
        self.state_repeat_times = state_repeat_times
        self.sample_num = sample_num
        # self.ego_cam_intrin = torch.tensor(
        #     [[91.703, 0., 63.304],
        #      [0., 91.746, 36.884],
        #      [0., 0., 1.]],
        #     device="cuda:0",
        # )
        self.ego_cam_intrin_inv = torch.tensor(
            [
                [0.01090477, 0.0, -0.69031547],
                [0.0, 0.01089966, -0.40202298],
                [0.0, 0.0, 1.0],
            ],
            device="cuda:0",
        )
        self.ego_cam_extrin = torch.tensor(
            [
                [-0.0127, -0.0079, -0.9999, -0.2265],
                [-0.9999, 0.0038, 0.0127, 0.0253],
                [0.0037, 1.0, -0.0079, 0.0774],
                [0.0, 0.0, 0.0, 1.0],
            ],
            device="cuda:0",
        )
        # self.ego_cam_extrin_inv = torch.tensor(
        #     [[-0.01272954, -0.99992464, 0.00369915, 0.02212854],
        #      [-0.00785192, 0.00379982, 0.99992353, -0.07926868],
        #      [-0.99987629, 0.01267029, -0.00794717, -0.22617743],
        #      [0., 0., 0., 1.]],
        #     device="cuda:0",
        # )

        # load the Attention model
        siamese_net = SiameseNet(
            input_channels=[3, 1],
            filters=[8],
            kernel_sizes=[5],
            strides=[1],
            activation="lrelu",
            norm=None,
        )
        self.attn = Attention2DUNet(
            siamese_net=siamese_net,
            filters=[16, 16],
            kernel_sizes=[5, 5],
            strides=[2, 2],
            output_channels=5,
            norm=None,
            activation="lrelu",
            skip_connections=True,
        ).to(torch.device("cuda:0"))
        self.attn.train()

        # self.weight_gen = WeightGen(
        #     feat_chan=16,
        #     weight_num=5,
        #     norm="layer",
        #     activation="lrelu",
        # ).to(torch.device("cuda:0"))
        # self.weight_gen.train()
        # self.weight = None

        # self.feat_agg = FeatAggNet(
        #     in_chan=5,
        #     out_chan=1,
        #     norm=None,
        #     activation="relu",
        # ).to(torch.device("cuda:0"))
        # self.feat_agg.train()

        # self.pn = PointNetfeat(in_ch=3, mlp_specs=[64, 128, 1024], xyz_transform=False, feature_transform=False)
        # self.pn.train()
        # self.weight_centers = torch.nn.Parameter(torch.ones(5), requires_grad=True)

        # Update the features dim manually
        self._features_dim = 12 * self.state_repeat_times + self.sample_num * 3 * 5

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]
        ee_pos_base = observations["ee_pos_base"]
        # check_grasp = observations["check_grasp"]
        # cam_xyz = observations["cam_xyz"]
        # ee_xyz = observations["ee_xyz"]
        rgb = observations["rgb"]
        rgb = rgb / 255.0  # Normalization
        N, ch, h, w = rgb.shape
        depth = observations["depth"]  # (N, 1, H, W)
        depth_norm = torch.clip(depth - torch.mean(depth), -1, 1)

        seg_feat, latent_feat = self.attn(
            [rgb, depth_norm]
        )  # (N, 5, H, W), (N, 16, H//4, W//4)
        seg_map = F.softmax(seg_feat, dim=1)  # (N, 5, H, W)
        # seg_map = F.softmax(self.attn([rgb, depth_norm]), dim=1)  # (N, 5, H, W)
        # seg_map = self.feat_agg(seg_map)  # (N, 1, H, W)

        ## argmax to get maxpoint ind
        # max_indice = torch.argmax(seg_map.view(N, 5, -1), dim=2)  # (N, 5)
        ## soft_argmax to get maxpoint ind
        # max_indice = spatial_soft_argmax2d(attnmap, normalized_coordinates=True)  # (N, 1, H, W)->(N,1,2)
        ## multinomial sampling
        # weights = torch.clone(attnmap)
        # weights[weights.ge(0.5)] = 200  # handle_ratio: 72*128/42 ~ 220
        # weights[weights.lt(0.5)] = 1
        # indices = torch.multinomial(weights, self.pts_num, replacement=False)  # (N, pts_num)
        ## gumbel_softmax to sample sample_num pixel coords
        # indices = torch.nn.functional.gumbel_softmax(
        #     seg_map.view(N, -1).unsqueeze(1).repeat(1, self.sample_num, 1),
        #     tau=0.1,
        #     hard=True
        # )  # (N, sample_num, h*w)

        ## gumber_softmax to get sample_num mask indexes
        # mask_ind_onehot = torch.nn.functional.gumbel_softmax(
        #     self.weight_centers.unsqueeze(1).expand(-1, self.sample_num),
        #     tau=0.1,
        #     hard=True
        # )  # (5, sample_num)
        # print("===", mask_ind_onehot.requires_grad)
        # sample_masks = torch.einsum("ijkl,jm->imkl", seg_map, mask_ind_onehot)  # (N, sample_num, H, W)
        # indices = spatial_soft_argmax2d(sample_masks, normalized_coordinates=True)  # (N, sample_num, 2)
        # zc = F.grid_sample(-depth, indices.unsqueeze(1), align_corners=False).squeeze(1).squeeze(1)  # (N, sample_num)
        # v = (indices[:, :, 1] + 1) / 2 * (h - 1)  # (N, sample_num)
        # u = (indices[:, :, 0] + 1) / 2 * (w - 1)  # (N, sample_num)
        # uvz_sample = torch.stack((u * zc, v * zc, zc), dim=1)  # (N, 3, sample_num)
        # camxyz_sample = torch.matmul(self.ego_cam_intrin_inv, uvz_sample)  # (N, 3, sample_num)
        # camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, sample_num)
        # eexyz_sample = torch.matmul(camxyz_sample.transpose(1, 2).contiguous(),
        #                       (self.ego_cam_extrin[:3, :3]).t()) \
        #         + self.ego_cam_extrin[:3, 3]  # (N, sample_num, 3)
        # eepts_pred = eexyz_sample.view(N, -1)  # (N, sample_num*3)

        ### topk to get topk ind
        _, indices = torch.topk(
            seg_map.view(N, 5, -1), self.sample_num, dim=2
        )  # (N, 5, sample_num)
        zc = (
            -depth.view(N, -1).unsqueeze(1).repeat(1, 5, 1).gather(2, indices)
        )  # (N, 5, sample_num)
        v = torch.div(indices, w, rounding_mode="floor")  # (N, 5, sample_num)
        u = indices - v * w  # (N, 5, sample_num)
        ## mean point from topk
        # uvz_sample = torch.stack((u * zc, v * zc, zc), dim=1).mean(3)  # (N, 3, 5)
        ## all k points
        uvz_sample = torch.stack((u * zc, v * zc, zc), dim=1).view(
            N, 3, -1
        )  # (N, 3, 5 * sample_num)
        camxyz_sample = torch.matmul(
            self.ego_cam_intrin_inv, uvz_sample
        )  # (N, 3, 5 * sample_num)
        camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, 5 * sample_num)
        eexyz_sample = (
            torch.matmul(
                camxyz_sample.transpose(1, 2).contiguous(),
                (self.ego_cam_extrin[:3, :3]).t(),
            )
            + self.ego_cam_extrin[:3, 3]
        )  # (N, 5 * sample_num, 3)
        # weight = F.softmax(self.weight_centers, dim=0)  # (5,)
        ## mean point from topk
        # eepts_pred = torch.einsum("ijl,j->il", eexyz_sample, weight)  # (N, 3)
        ## weighted k points
        # eexyz_part_sample = eexyz_sample.view(N, 5, self.sample_num, 3)  # (N, 5, sample_num, 3)
        # eepts_pred = torch.einsum("ijkl,j->ikl", eexyz_part_sample, weight).view(N, -1)  # (N, sample_num * 3)
        ## all 5*k points
        eepts_pred = eexyz_sample.view(N, -1)  # (N, 5*sample_num*3)

        # zc = torch.stack([-depth.view(N, -1)[i, indices[i]] for i in range(N)])  # (N, sample_num)
        # zc = (-depth.view(N, -1)).gather(1, indices)  # (N, sample_num)
        # v = torch.div(indices, w, rounding_mode="floor")  # (N, sample_num)
        # u = (indices - v * w)  # (N, sample_num)
        # uvz_sample = torch.stack((u * zc, v * zc, zc), dim=1)  # (N, 3, sample_num)
        # camxyz_sample = torch.matmul(self.ego_cam_intrin_inv, uvz_sample)  # (N, 3, pts_num)
        # camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, pts_num)
        # ee_xyz_sample = torch.matmul(camxyz_sample.transpose(1, 2).contiguous(),
        #                       (self.ego_cam_extrin[:3, :3]).t()) \
        #         + self.ego_cam_extrin[:3, 3]  # (N,sample_num,3)
        ## vote one handle_center point from the sample_num points
        # eepts_pred = ee_xyz_sample[:, 0]  # (N,3)

        # zc = torch.stack([-depth.view(N, -1)[i, max_indice[i]] for i in range(N)])  # (N, sample_num), too much time
        # zc = ((-depth.view(N, -1)).gather(1, max_indice))  # (N, 5)
        # v = torch.div(max_indice, w, rounding_mode="floor")  # (N, 5)
        # u = max_indice - v * w  # (N, 5)
        # uvz_sample = torch.stack((u * zc, v * zc, zc), dim=1)  # (N, 3, 5)
        # camxyz_sample = torch.matmul(self.ego_cam_intrin_inv, uvz_sample)  # (N, 3, 5)
        # camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, 5)
        # eexyz_sample = torch.matmul(camxyz_sample.transpose(1, 2).contiguous(),
        #                             (self.ego_cam_extrin[:3, :3]).t()) \
        #                + self.ego_cam_extrin[:3, 3]  # (N, 5, 3)
        # eepts_pred = eexyz_sample.view(N, -1)  # (N, 5*3)

        # zc = F.grid_sample(-depth, max_indice.unsqueeze(1), align_corners=False).squeeze(1).squeeze(1)  # (N, sample_num)
        # v = (max_indice[:, :, 1] + 1) / 2 * (h - 1)  # (N, sample_num)
        # u = (max_indice[:, :, 0] + 1) / 2 * (w - 1)  # (N, sample_num)
        # uvz_sample = torch.stack((u * zc, v * zc, zc), dim=1)  # (N, 3, sample_num)

        # zc = torch.einsum("ij,ikj->ik", -depth.view(N, -1), indices)  # (N, sample_num)
        # u, v = torch.arange(0, w, dtype=depth.dtype, device=depth.device), \
        #        torch.arange(0, h, dtype=depth.dtype, device=depth.device)
        # grid_v, grid_u = torch.meshgrid(v, u, indexing="ij")
        # idx_v = torch.einsum("ijk,k->ij", indices, grid_v.reshape(-1))  # (N, sample_num)
        # idx_u = torch.einsum("ijk,k->ij", indices, grid_u.reshape(-1))  # (N, sample_num)
        # uvz_sample = torch.stack((idx_u * zc, idx_v * zc, zc), dim=1)  # (N, 3, sample_num)
        # camxyz_sample = torch.matmul(self.ego_cam_intrin_inv, uvz_sample)  # (N, 3, sample_num)
        # camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, sample_num)
        # eexyz_sample = torch.matmul(camxyz_sample.transpose(1, 2).contiguous(),
        #                             (self.ego_cam_extrin[:3, :3]).t()) \
        #                + self.ego_cam_extrin[:3, 3]  # (N, sample_num, 3)
        # eepts_pred = eexyz_sample.view(N, -1)

        ## for weight_center
        # zc = ((-depth.view(N, -1)).gather(1, max_indice)).squeeze(1)  # (N, 5)
        # v = torch.div(max_indice, w, rounding_mode="floor")  # (N, 5)
        # u = max_indice - v * w  # (N, 5)
        # uvz_sample = torch.stack((u * zc, v * zc, zc), dim=1)  # (N, 3, 5)
        # camxyz_sample = torch.matmul(self.ego_cam_intrin_inv, uvz_sample)  # (N, 3, 5)
        # camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, 5)
        # eexyz_sample = torch.matmul(camxyz_sample.transpose(1, 2).contiguous(),
        #                             (self.ego_cam_extrin[:3, :3]).t()) \
        #                + self.ego_cam_extrin[:3, 3]  # (N, 5, 3)
        # weight = F.softmax(self.weight_centers, dim=0)  # (5,)
        # eepts_pred = torch.einsum("ijk,j->ik", eexyz_sample, weight)  # (N, 3)
        # self.weight = F.softmax(self.weight_gen(latent_feat.detach()), dim=1)  # (N, 5)
        # eepts_pred = torch.einsum("ijk,ij->ik", eexyz_sample, self.weight)  # (N, 3)

        # ee_xyz = observations["ee_xyz"]
        # eexyz_sample = torch.stack([ee_xyz.view(N, 3, -1)[i, :, indices[i]] for i in range(N)])  # (N, 3, pts_num)
        # eergb_pred = torch.stack([rgb.view(N, 3, -1)[i, :, indices[i]] for i in range(N)])  # (N, 3, pts_num)
        # pn_feature, _, _, _ = self.pn(torch.cat((eexyz_pred, eergb_pred), dim=1))  # (N, 6, pts_num) -> (N, 128)
        # pn_feature, _, _, _ = self.pn(ee_xyz)  # (N, 3, pts_num) -> (N, 128)
        # eepts_pred = torch.stack([ee_xyz.view(N, 3, -1)[i, :, max_indice[i]] for i in range(N)])  # (N, 3)

        state = torch.cat([qpos, ee_pos_base], dim=1).repeat(1, self.state_repeat_times)

        return torch.cat([state, eepts_pred], dim=1)


class CustomRGBDPretrainSegExtractor(BaseFeaturesExtractor):
    def __init__(
        self, observation_space: gym.spaces.Dict, state_repeat_times=2, sample_num=50
    ):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomRGBDPretrainSegExtractor, self).__init__(
            observation_space, features_dim=1
        )

        self.state_repeat_times = state_repeat_times
        self.sample_num = sample_num

        self.ego_cam_intrin_inv = torch.tensor(
            [
                [0.01090477, 0.0, -0.69031547],
                [0.0, 0.01089966, -0.40202298],
                [0.0, 0.0, 1.0],
            ],
            device="cuda:0",
        )
        self.ego_cam_extrin = torch.tensor(
            [
                [-0.0127, -0.0079, -0.9999, -0.2265],
                [-0.9999, 0.0038, 0.0127, 0.0253],
                [0.0037, 1.0, -0.0079, 0.0774],
                [0.0, 0.0, 0.0, 1.0],
            ],
            device="cuda:0",
        )

        # log_name = "pretrain_unet/20221212_152355_cellb005_allcabfaucetfilter_bs512_lr1e-3_bnrelu_2siam4updown163264128"
        # log_name = "pretrain_unet/20221214_210516_cellbs0.1_allcabfaucetfilter_bs512_0.9step20lr1e-3_bnrelu_2siam2updown1616"
        log_name = (
            "smp_model/20221227_000530_allcabfaucet_bs256_0.9step50lr1e-3_unet++3_effb1"
        )
        model_path = VISUALMODEL_DIR / f"{log_name}/best.pth"
        config_path = VISUALMODEL_DIR / f"{log_name}/config.yaml"
        with open(config_path, "r") as f:
            cfg = yaml.safe_load(f)
        smp_cfg = cfg["smp_config"]
        self.attn = smp.UnetPlusPlus(
            encoder_name=smp_cfg["encoder"],
            encoder_depth=smp_cfg["encoder_depth"],
            decoder_channels=smp_cfg["decoder_channels"],
            encoder_weights=smp_cfg["encoder_weights"],
            in_channels=smp_cfg["in_channels"],
            classes=cfg["num_classes"],
            activation=smp_cfg["activation"],
        )
        self.attn.load_state_dict(torch.load(model_path))
        self.attn.to(torch.device("cuda:0"))
        # siamese_cfg = cfg["segnet_config"]["siamese_config"]
        # unet_cfg = cfg["segnet_config"]["unet_config"]
        # siamese_net = SiameseNet(
        #     input_channels=siamese_cfg["input_channels"],
        #     filters=siamese_cfg["filters"],
        #     kernel_sizes=siamese_cfg["kernel_sizes"],
        #     strides=siamese_cfg["strides"],
        #     norm=siamese_cfg["norm"],
        #     activation=siamese_cfg["activation"]
        # )
        # self.attn = Attention2DUNet(
        #     siamese_net=siamese_net,
        #     filters=unet_cfg["filters"],
        #     kernel_sizes=unet_cfg["kernel_sizes"],
        #     strides=unet_cfg["strides"],
        #     output_channels=cfg["num_classes"],
        #     norm=unet_cfg["norm"],
        #     activation=unet_cfg["activation"],
        #     skip_connections=unet_cfg["skip_connections"]
        # ).to(torch.device("cuda:0"))
        # self.attn.load_w(model_path)
        self.attn.eval()
        # self.attn.requires_grad(False)
        self.num_classes = cfg["num_classes"]

        # Update the features dim manually
        self._features_dim = (
            12 * self.state_repeat_times + self.sample_num * 3 * self.num_classes
        )

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]
        ee_pos_base = observations["ee_pos_base"]
        rgb = observations["rgb"]
        rgb = rgb / 255.0  # (N, 3, H, W)
        N, ch, h, w = rgb.shape
        depth = observations["depth"]  # (N, 1, H, W)
        # depth_norm = torch.clip(depth - torch.mean(depth), -1, 1)
        depth_norm = depth / torch.max(depth)

        # seg_feat, latent_feat = self.attn([rgb, depth_norm])  # (N, C, H, W), (N, 16, H//4, W//4)
        # seg_map = F.softmax(seg_feat, dim=1)  # (N, C, H, W)
        seg_feat = self.attn.predict(
            torch.cat([rgb, depth_norm], dim=1)
        )  # (N, C, H, W)
        seg_map = F.softmax(seg_feat, dim=1)  # (N, C, H, W)

        # handle masks delete other handles
        # other_handle_seg = observations["other_handleseg"]
        # seg_map[:, 1][other_handle_seg == 1] = 0

        ## argmax to get maxpoint ind
        # max_indice = torch.argmax(seg_map.view(N, self.num_classes, -1), dim=2)  # (N, C)
        # zc = ((-depth.view(N, -1)).gather(1, max_indice))  # (N, C)
        # v = torch.div(max_indice, w, rounding_mode="floor")  # (N, C)
        # u = max_indice - v * w  # (N, C)
        # uvz_sample = torch.stack((u * zc, v * zc, zc), dim=1)  # (N, 3, C)
        # camxyz_sample = torch.matmul(self.ego_cam_intrin_inv, uvz_sample)  # (N, 3, C)
        # camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, C)
        # eexyz_sample = torch.matmul(camxyz_sample.transpose(1, 2).contiguous(),
        #                             (self.ego_cam_extrin[:3, :3]).t()) \
        #                + self.ego_cam_extrin[:3, 3]  # (N, C, 3)
        # eepts_pred = eexyz_sample.view(N, -1)  # (N, C*3)

        ### topk to get topk ind
        _, indices = torch.topk(
            seg_map.view(N, self.num_classes, -1), self.sample_num, dim=2
        )  # (N, C, sample_num)
        zc = (
            -depth.view(N, -1)
            .unsqueeze(1)
            .repeat(1, self.num_classes, 1)
            .gather(2, indices)
        )  # (N, C, sample_num)
        v = torch.div(indices, w, rounding_mode="floor")  # (N, C, sample_num)
        u = indices - v * w  # (N, C, sample_num)
        ## all k points
        uvz_sample = torch.stack((u * zc, v * zc, zc), dim=1).view(
            N, 3, -1
        )  # (N, 3, C * sample_num)
        camxyz_sample = torch.matmul(
            self.ego_cam_intrin_inv, uvz_sample
        )  # (N, 3, C * sample_num)
        camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, C * sample_num)
        eexyz_sample = (
            torch.matmul(
                camxyz_sample.transpose(1, 2).contiguous(),
                (self.ego_cam_extrin[:3, :3]).t(),
            )
            + self.ego_cam_extrin[:3, 3]
        )  # (N, C * sample_num, 3)
        ## all 5*k points
        eepts_pred = eexyz_sample.view(N, -1)  # (N, C*sample_num*3)

        state = torch.cat([qpos, ee_pos_base], dim=1).repeat(1, self.state_repeat_times)

        return torch.cat([state, eepts_pred], dim=1)


class CustomRGBDPredSegResNetExtractor(BaseFeaturesExtractor):
    def __init__(
        self,
        observation_space: gym.spaces.Dict,
        state_repeat_times=2,
        vismodel_path=None,
        device="cuda",
    ):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomRGBDPredSegResNetExtractor, self).__init__(
            observation_space, features_dim=1
        )

        self.state_repeat_times = state_repeat_times
        self.device = device

        # vismodel_path = VISUALMODEL_DIR / f"smp_model/20230228_220446_384_DR_randombg_mixfaucet_aug_dropout0.5_stereo_bs16_focalloss_0.5step50lr0.001_RGBDunet-163264128_mobilenet_v2"
        model_path = vismodel_path / "best.pth"
        config_path = vismodel_path / "config.yaml"
        with open(config_path, "r") as f:
            cfg = yaml.safe_load(f)
        smp_cfg = cfg["smp_config"]
        if smp_cfg["mode"] == "RGBD":
            in_ch = 4
        elif smp_cfg["mode"] == "RGB":
            in_ch = 3
        elif smp_cfg["mode"] == "D":
            in_ch = 1
        else:
            raise NotImplementedError
        self.segmodel = CustomUnet(
            has_dropout=False,
            encoder_name=smp_cfg["encoder"],
            encoder_depth=smp_cfg["encoder_depth"],
            decoder_channels=smp_cfg["decoder_channels"],
            encoder_weights=smp_cfg["encoder_weights"],
            in_channels=in_ch,
            classes=cfg["num_classes"],
            activation=smp_cfg["activation"],
        )
        self.segmodel.load_state_dict(torch.load(model_path))
        self.segmodel.to(torch.device(self.device))
        self.segmodel.eval()

        self.num_classes = cfg["num_classes"]

        # self.resnet = ResNet(
        #     block=BasicBlock,
        #     num_blocks=[1, 1, 1, 1],
        #     in_dim=self.num_classes + 1,
        #     planes=16
        # )
        # self.resnet.to(torch.device(self.device))
        # self.resnet.train()
        self.encoder = ImgEncoder(
            in_ch=self.num_classes + 1, feature_dim=256, num_layers=4, num_filters=32
        )

        # Update the features dim manually
        self._features_dim = 12 * self.state_repeat_times + 256

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]
        ee_pos_base = observations["ee_pos_base"]
        rgb = observations["rgb"]
        rgb = rgb / 255.0  # (N, 3, H, W)
        depth = observations["depth"]  # (N, 1, H, W)

        with torch.no_grad():
            seg_map = self.segmodel.predict(
                torch.cat([rgb, depth], dim=1)
            )  # (N, C, H, W)

        # img_feat_norm = self.resnet(torch.cat([seg_map, depth], dim=1))  # (N, 256)
        # img_feat_norm = F.layer_norm(img_feat, img_feat.size()[1:])
        img_feat_norm = self.encoder.forward(
            torch.cat([seg_map, depth], dim=1)
        )  # (N, 256)

        state = torch.cat([qpos, ee_pos_base], dim=1).repeat(1, self.state_repeat_times)

        return torch.cat([state, img_feat_norm], dim=1)


class CustomDSegResNetExtractor(BaseFeaturesExtractor):
    def __init__(
        self,
        observation_space: gym.spaces.Dict,
        state_repeat_times=2,
        num_classes=6,
        device="cuda",
    ):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomDSegResNetExtractor, self).__init__(
            observation_space, features_dim=1
        )

        self.state_repeat_times = state_repeat_times
        self.num_classes = num_classes
        self.device = device

        # self.resnet = models.resnet18(pretrained=False)
        # self.resnet.conv1 = torch.nn.Conv2d(
        #     2,
        #     64,
        #     kernel_size=(7, 7),
        #     stride=(2, 2),
        #     padding=(3, 3),
        #     bias=False,
        # )
        # self.resnet.fc = torch.nn.Linear(in_features=512, out_features=256, bias=True)
        # self.resnet.to(torch.device(self.device))
        # self.resnet.train()
        self.encoder = ImgEncoder(
            in_ch=self.num_classes + 1, feature_dim=256, num_layers=4, num_filters=32
        )

        # Update the features dim manually
        self._features_dim = 12 * self.state_repeat_times + 256

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]
        ee_pos_base = observations["ee_pos_base"]
        depth = observations["depth"]  # (N, H, W)
        seg = observations["seg"]  # (N, H, W)
        seg_map = (
            F.one_hot(seg.long(), self.num_classes)
            .permute(0, 3, 1, 2)
            .to(self.device, dtype=torch.float32)
        )  # (N, C, H, W)

        # img_feat = self.resnet(torch.cat([depth.unsqueeze(1), seg.unsqueeze(1)], dim=1))  # (N, 256)
        # img_feat_norm = F.layer_norm(img_feat, img_feat.size()[1:])
        img_feat_norm = self.encoder.forward(
            torch.cat([seg_map, depth], dim=1)
        )  # (N, 256)

        state = torch.cat([qpos, ee_pos_base], dim=1).repeat(1, self.state_repeat_times)

        return torch.cat([state, img_feat_norm], dim=1)


class CustomRGBDPretrainSegPointNetExtractor(BaseFeaturesExtractor):
    def __init__(
        self, observation_space: gym.spaces.Dict, state_repeat_times=2, sample_num=50
    ):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomRGBDPretrainSegPointNetExtractor, self).__init__(
            observation_space, features_dim=1
        )

        self.state_repeat_times = state_repeat_times
        self.sample_num = sample_num

        self.ego_cam_intrin_inv = torch.tensor(
            [
                [0.01090477, 0.0, -0.69031547],
                [0.0, 0.01089966, -0.40202298],
                [0.0, 0.0, 1.0],
            ],
            device="cuda:0",
        )
        self.ego_cam_extrin = torch.tensor(
            [
                [-0.0127, -0.0079, -0.9999, -0.2265],
                [-0.9999, 0.0038, 0.0127, 0.0253],
                [0.0037, 1.0, -0.0079, 0.0774],
                [0.0, 0.0, 0.0, 1.0],
            ],
            device="cuda:0",
        )

        # log_name = "smp_model/20230110_103746_cab39-17_faucet11-3_newaug_bs256_0.5step50lr0.001_RGBDunet3-3264128_mobilenet_v2"
        log_name = "smp_model/20230228_014637_384_noDR_norandombg_aug_dropout0.2_stereo_bs16_focalloss_0.5step50lr0.001_RGBDunet-163264128_mobilenet_v2"
        model_path = VISUALMODEL_DIR / f"{log_name}/best.pth"
        config_path = VISUALMODEL_DIR / f"{log_name}/config.yaml"
        with open(config_path, "r") as f:
            cfg = yaml.safe_load(f)
        smp_cfg = cfg["smp_config"]
        if smp_cfg["mode"] == "RGBD":
            in_ch = 4
        elif smp_cfg["mode"] == "RGB":
            in_ch = 3
        elif smp_cfg["mode"] == "D":
            in_ch = 1
        else:
            raise NotImplementedError
        self.segmodel = CustomUnet(
            has_dropout=False,
            encoder_name=smp_cfg["encoder"],
            encoder_depth=smp_cfg["encoder_depth"],
            decoder_channels=smp_cfg["decoder_channels"],
            encoder_weights=smp_cfg["encoder_weights"],
            in_channels=in_ch,
            classes=cfg["num_classes"],
            activation=smp_cfg["activation"],
        )
        self.segmodel.load_state_dict(torch.load(model_path))
        self.segmodel.to("cuda:0")
        self.segmodel.eval()
        self.num_classes = cfg["num_classes"]

        self.pn = PointNetfeat(
            in_ch=3 + self.num_classes + 1,
            global_feat=True,
            mlp_specs=[64, 128, 256],
            xyz_transform=False,
            feature_transform=False,
        )
        self.pn.train()
        seg_info = torch.zeros(
            self.num_classes * self.sample_num + 40,
            self.num_classes + 1,
            device=torch.device("cuda:0"),
        )
        for nc in range(self.num_classes):
            seg_info[nc * self.sample_num : (nc + 1) * self.sample_num, nc] = 1
        seg_info[-40:, self.num_classes] = 1

        # shuffle random 5 points label
        # random_label = torch.eye(
        #     self.num_classes, device=torch.device("cuda:0")
        # )[torch.randint(0, self.num_classes, (5 * self.num_classes,))]
        # random_label = torch.zeros(
        #     self.num_classes * 5, self.num_classes + 1,
        #     device=torch.device("cuda:0")
        # )
        # indices = torch.randint(0, self.num_classes, (self.num_classes * 5, ), device=torch.device("cuda:0"))
        # random_label.scatter_(1, indices.unsqueeze(1), 1)
        # seg_info[
        #     np.random.choice(range(self.sample_num*self.num_classes), 5 * self.num_classes), :
        # ] = random_label

        self.seg_info = seg_info

        # Update the features dim manually
        self._features_dim = 12 * self.state_repeat_times + 256

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]
        ee_pos_base = observations["ee_pos_base"]
        eepad_pts = observations["eepad_pts"]  # (N, 40, 3)
        rgb = observations["rgb"]
        # rgb = rgb / 255.0  # (N, 3, H, W)
        N, ch, h, w = rgb.shape
        # depth = observations["depth"]  # (N, 1, H, W)
        depth = observations["depth"]  # (N, 1, H, W)

        with torch.no_grad():
            # t0 = time.time()
            seg_map = self.segmodel.predict(
                torch.cat([rgb, depth], dim=1)
            )  # (N, C, H, W)
            # t1 = time.time()
            # print("+", t1-t0)

            ### random sampling
            indices = []
            num_classes = seg_map.shape[1]
            seg_map_mc = torch.argmax(seg_map, dim=1).view(N, -1)  # (N, H, W)
            t2 = time.time()
            # print("++", t2 - t1)
            for num in range(N):
                part_inds = []
                for ind in range(num_classes):
                    one_indices = torch.where(seg_map_mc[num] == ind)[0]
                    ones_num = one_indices.shape[0]
                    if ones_num == 0:
                        cur_indices = torch.zeros(
                            self.sample_num, dtype=torch.int64, device="cuda:0"
                        )
                    elif ones_num < self.sample_num:
                        repeat_times = self.sample_num // ones_num
                        remain_num = self.sample_num - repeat_times * ones_num
                        cur_indices = torch.cat(
                            (one_indices.repeat(repeat_times), one_indices[:remain_num])
                        )
                    else:
                        cur_indices = one_indices[
                            random.sample(range(len(one_indices)), self.sample_num)
                        ]
                    part_inds.append(cur_indices[None])
                pts_index = torch.cat(part_inds, dim=0)  # (C, sample_num)
                indices.append(pts_index[None])
            indices = torch.cat(indices, dim=0)  # (N, C, sample_num)
            # t3 = time.time()
            # print("+++", t3-t2)

        ### topk to get topk ind
        # values, indices = torch.topk(
        #     seg_map.view(N, self.num_classes, -1), self.sample_num, dim=2
        # )  # (N, C, sample_num)
        zc = (
            -depth.view(N, -1)
            .unsqueeze(1)
            .repeat(1, self.num_classes, 1)
            .gather(2, indices)
        )  # (N, C, sample_num)
        v = torch.div(indices, w, rounding_mode="floor")  # (N, C, sample_num)
        u = indices - v * w  # (N, C, sample_num)
        ## all k points
        uvz_sample = torch.stack((u * zc, v * zc, zc), dim=1).view(
            N, 3, -1
        )  # (N, 3, C * sample_num)
        camxyz_sample = torch.matmul(
            self.ego_cam_intrin_inv, uvz_sample
        )  # (N, 3, C * sample_num)
        camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, C * sample_num)
        eexyz_sample = (
            torch.matmul(
                camxyz_sample.transpose(1, 2).contiguous(),
                (self.ego_cam_extrin[:3, :3]).t(),
            )
            + self.ego_cam_extrin[:3, 3]
        )  # (N, C * sample_num, 3)

        eewpad_xyz_sample = torch.cat(
            (eexyz_sample, eepad_pts), dim=1
        )  # (N, C*sample_num+40, 3)

        ## concat class info
        eexyzC_sample = torch.cat(
            (eewpad_xyz_sample, self.seg_info.repeat(N, 1, 1)), dim=2
        )  # (N, C*sample_num+40, 3+C+1)
        ## concat class & confidence info
        # conf_scores = torch.cat(
        #     (values.view(N, -1), torch.ones((N, 40), device=torch.device("cuda:0"))), dim=1
        # ).unsqueeze(2)  # (N, C*sample_num+40, 1)
        # eexyzS_sample = torch.cat(
        #     (eewpad_xyz_sample, self.seg_info.repeat(N, 1, 1), conf_scores), dim=2
        # )

        ## PointNet to get 128 feat
        pn_feature, _, _ = self.pn(
            eexyzC_sample.transpose(2, 1).contiguous()
        )  # (N, 3+C+1, C*sample_num+40)
        # pn_feature, _, _ = self.pn(eexyzS_sample.transpose(2, 1).contiguous())  # (N, 3+C+1+1, C*sample_num+40)

        state = torch.cat([qpos, ee_pos_base], dim=1).repeat(1, self.state_repeat_times)

        return torch.cat([state, pn_feature], dim=1)


class CustomPretrainSegPNExtractor(BaseFeaturesExtractor):
    def __init__(
        self,
        observation_space: gym.spaces.Dict,
        num_classes=6,
        state_repeat_times=2,
        sample_num=50,
    ):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomPretrainSegPNExtractor, self).__init__(
            observation_space, features_dim=1
        )

        self.num_classes = num_classes
        self.state_repeat_times = state_repeat_times
        self.sample_num = sample_num

        self.pn = PointNetfeat(
            in_ch=3 + self.num_classes + 1,
            global_feat=True,
            mlp_specs=[64, 128, 256],
            xyz_transform=False,
            feature_transform=False,
        )
        self.pn.train()
        seg_info = torch.zeros(
            self.num_classes * self.sample_num + 40,
            self.num_classes + 1,
            device=torch.device("cuda:0"),
        )
        for nc in range(self.num_classes):
            seg_info[nc * self.sample_num : (nc + 1) * self.sample_num, nc] = 1
        seg_info[-40:, self.num_classes] = 1

        # shuffle random 5 points label
        # random_label = torch.eye(
        #     self.num_classes, device=torch.device("cuda:0")
        # )[torch.randint(0, self.num_classes, (5 * self.num_classes,))]
        # random_label = torch.zeros(
        #     self.num_classes * 5, self.num_classes + 1,
        #     device=torch.device("cuda:0")
        # )
        # indices = torch.randint(0, self.num_classes, (self.num_classes * 5, ), device=torch.device("cuda:0"))
        # random_label.scatter_(1, indices.unsqueeze(1), 1)
        # seg_info[
        #     np.random.choice(range(self.sample_num*self.num_classes), 5 * self.num_classes), :
        # ] = random_label

        self.seg_info = seg_info

        # Update the features dim manually
        self._features_dim = 12 * self.state_repeat_times + 256

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]  # (N, 9)
        ee_pos_base = observations["ee_pos_base"]  # (N, 3)
        eepad_pts = observations["eepad_pts"]  # (N, 40, 3)
        eexyz_sample = observations["segsampled_pts"]  # (N, C*sample_num, 3)
        N = qpos.shape[0]

        eewpad_xyz_sample = torch.cat(
            (eexyz_sample, eepad_pts), dim=1
        )  # (N, C*sample_num+40, 3)

        ## concat class info
        eexyzC_sample = torch.cat(
            (eewpad_xyz_sample, self.seg_info.repeat(N, 1, 1)), dim=2
        )  # (N, C*sample_num+40, 3+C+1)
        ## concat class & confidence info
        # conf_scores = torch.cat(
        #     (values.view(N, -1), torch.ones((N, 40), device=torch.device("cuda:0"))), dim=1
        # ).unsqueeze(2)  # (N, C*sample_num+40, 1)
        # eexyzS_sample = torch.cat(
        #     (eewpad_xyz_sample, self.seg_info.repeat(N, 1, 1), conf_scores), dim=2
        # )

        ## PointNet to get 128 feat
        pn_feature, _, _ = self.pn(
            eexyzC_sample.transpose(2, 1).contiguous()
        )  # (N, 3+C+1, C*sample_num+40)
        # pn_feature, _, _ = self.pn(eexyzS_sample.transpose(2, 1).contiguous())  # (N, 3+C+1+1, C*sample_num+40)

        state = torch.cat([qpos, ee_pos_base], dim=1).repeat(1, self.state_repeat_times)

        return torch.cat([state, pn_feature], dim=1)


class CustomSegPNExtractor(BaseFeaturesExtractor):
    def __init__(
        self,
        observation_space: gym.spaces.Dict,
        num_classes=6,
        state_repeat_times=2,
        sample_num=50,
    ):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomSegPNExtractor, self).__init__(observation_space, features_dim=1)

        self.num_classes = num_classes
        self.state_repeat_times = state_repeat_times
        self.sample_num = sample_num

        self.pn = PointNetfeat(
            in_ch=3 + self.num_classes,
            # in_ch=3 + self.num_classes+1,
            global_feat=True,
            mlp_specs=[64, 128, 256],
            xyz_transform=False,
            feature_transform=False,
        )
        self.pn.train()

        # Update the features dim manually
        self._features_dim = 12 * self.state_repeat_times + 256

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]  # (N, 9)
        ee_pos_base = observations["ee_pos_base"]  # (N, 3)
        eexyzc_sample = observations["segsampled_ptsC"]  # (N, SN+40, 3+C+1)

        ## PointNet to get 128 feat
        pn_feature, _, _ = self.pn(
            eexyzc_sample.transpose(2, 1).contiguous()
        )  # (N, 3+C+1, SN+40)
        # pn_feature, _, _ = self.pn(eexyzS_sample.transpose(2, 1).contiguous())  # (N, 3+C+1+1, SN+40)

        state = torch.cat([qpos, ee_pos_base], dim=1).repeat(1, self.state_repeat_times)

        return torch.cat([state, pn_feature], dim=1)


class CustomKptExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space: gym.spaces.Dict, num_kpts=3):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomKptExtractor, self).__init__(observation_space, features_dim=1)

        # self.ego_cam_intrin_inv = torch.tensor(
        #     [
        #         [0.01090477, 0.0, -0.69031547],
        #         [0.0, 0.01089966, -0.40202298],
        #         [0.0, 0.0, 1.0],
        #     ],
        #     device="cuda:0",
        # )
        # self.ego_cam_extrin = torch.tensor(
        #     [
        #         [-0.0127, -0.0079, -0.9999, -0.2265],
        #         [-0.9999, 0.0038, 0.0127, 0.0253],
        #         [0.0037, 1.0, -0.0079, 0.0774],
        #         [0.0, 0.0, 0.0, 1.0],
        #     ],
        #     device="cuda:0",
        # )

        self.num_kpts = num_kpts

        # Update the features dim manually
        self._features_dim = 12 + num_kpts * 3

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]  # (N, 9)
        ee_pos_base = observations["ee_pos_base"]  # (N, 3)
        N = qpos.shape[0]
        # uvz_pred = observations["uvz_pred"]  # (N, K, 3)
        eexyz_sample = observations["xyz_ee_pred"]  # (N, K, 3)

        ## all k points
        # camxyz_sample = torch.matmul(
        #     self.ego_cam_intrin_inv, uvz_pred.transpose(1, 2)
        # )  # (N, 3, K)
        # camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, K)
        # eexyz_sample = (
        #         torch.matmul(
        #             camxyz_sample.transpose(1, 2).contiguous(),
        #             (self.ego_cam_extrin[:3, :3]).t(),
        #         )
        #         + self.ego_cam_extrin[:3, 3]
        # )  # (N, K, 3)
        ## all 5*k points
        eepts_pred = eexyz_sample.view(N, -1)  # (N, K*3)

        state = torch.cat([qpos, ee_pos_base, eepts_pred], dim=1)

        return state


class CustomKptGTExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space: gym.spaces.Dict, num_kpts=3):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomKptGTExtractor, self).__init__(observation_space, features_dim=1)

        # self.ego_cam_intrin_inv = torch.tensor(
        #     [
        #         [0.01090477, 0.0, -0.69031547],
        #         [0.0, 0.01089966, -0.40202298],
        #         [0.0, 0.0, 1.0],
        #     ],
        #     device="cuda:0",
        # )
        # self.ego_cam_extrin = torch.tensor(
        #     [
        #         [-0.0127, -0.0079, -0.9999, -0.2265],
        #         [-0.9999, 0.0038, 0.0127, 0.0253],
        #         [0.0037, 1.0, -0.0079, 0.0774],
        #         [0.0, 0.0, 0.0, 1.0],
        #     ],
        #     device="cuda:0",
        # )

        self.num_kpts = num_kpts

        # Update the features dim manually
        self._features_dim = 12 + num_kpts * 3

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]  # (N, 9)
        ee_pos_base = observations["ee_pos_base"]  # (N, 3)
        N = qpos.shape[0]
        # uvz_gt = observations["uvz"]  # (N, K, 3)

        ## all k points
        # camxyz_sample = torch.matmul(
        #     self.ego_cam_intrin_inv, uvz_gt.transpose(1, 2)
        # )  # (N, 3, K)
        # camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, K)
        # eexyz_sample = (
        #     torch.matmul(
        #         camxyz_sample.transpose(1, 2).contiguous(),
        #         (self.ego_cam_extrin[:3, :3]).t(),
        #     )
        #     + self.ego_cam_extrin[:3, 3]
        # )  # (N, K, 3)
        # ## all 5*k points
        # eepts_pred = eexyz_sample.view(N, -1)  # (N, K*3)

        kpts_gt_noise = observations["kpts"]  # (N, K, 3)
        eepts_pred = kpts_gt_noise.view(N, -1)  # (N, K*3)

        state = torch.cat([qpos, ee_pos_base, eepts_pred], dim=1)

        return state


class CustomRGBDSegGTPointNetExtractor(BaseFeaturesExtractor):
    def __init__(
        self, observation_space: gym.spaces.Dict, state_repeat_times=2, sample_num=50
    ):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomRGBDSegGTPointNetExtractor, self).__init__(
            observation_space, features_dim=1
        )

        self.state_repeat_times = state_repeat_times
        self.sample_num = sample_num

        self.ego_cam_intrin_inv = torch.tensor(
            [
                [0.01090477, 0.0, -0.69031547],
                [0.0, 0.01089966, -0.40202298],
                [0.0, 0.0, 1.0],
            ],
            device="cuda:0",
        )
        self.ego_cam_extrin = torch.tensor(
            [
                [-0.0127, -0.0079, -0.9999, -0.2265],
                [-0.9999, 0.0038, 0.0127, 0.0253],
                [0.0037, 1.0, -0.0079, 0.0774],
                [0.0, 0.0, 0.0, 1.0],
            ],
            device="cuda:0",
        )
        self.num_classes = 6

        self.pn = PointNetfeat(
            in_ch=3 + self.num_classes + 1,
            mlp_specs=[64, 128, 256],
            global_feat=True,
            xyz_transform=False,
            feature_transform=False,
        )
        self.pn.train()
        seg_info = torch.zeros(
            self.num_classes * self.sample_num + 40,
            self.num_classes + 1,
            device=torch.device("cuda:0"),
        )
        for nc in range(self.num_classes):
            seg_info[nc * self.sample_num : (nc + 1) * self.sample_num, nc] = 1
        seg_info[-40:, self.num_classes] = 1
        self.seg_info = seg_info

        # Update the features dim manually
        self._features_dim = 12 * self.state_repeat_times + 256

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]
        ee_pos_base = observations["ee_pos_base"]
        eepad_pts = observations["eepad_pts"]  # (N, 40, 3)
        rgb = observations["rgb"]
        rgb = rgb / 255.0  # (N, 3, H, W)
        N, ch, h, w = rgb.shape
        depth = observations["depth"]  # (N, 1, H, W)
        seg = observations["seg"]  # (N, H, W)

        seg_map = (
            F.one_hot(seg.long(), num_classes=self.num_classes)
            .permute(0, 3, 1, 2)
            .contiguous()
        )  # (N, C, H, W)
        seg_map_noise = (
            seg_map + torch.randn(seg_map.shape, device=seg_map.device) * 0.1
        )

        ### topk to get topk ind
        values, indices = torch.topk(
            seg_map_noise.view(N, self.num_classes, -1), self.sample_num, dim=2
        )  # (N, C, sample_num)
        zc = (
            -depth.view(N, -1)
            .unsqueeze(1)
            .repeat(1, self.num_classes, 1)
            .gather(2, indices)
        )  # (N, C, sample_num)
        v = torch.div(indices, w, rounding_mode="floor")  # (N, C, sample_num)
        u = indices - v * w  # (N, C, sample_num)
        ## all k points
        uvz_sample = torch.stack((u * zc, v * zc, zc), dim=1).view(
            N, 3, -1
        )  # (N, 3, C * sample_num)
        camxyz_sample = torch.matmul(
            self.ego_cam_intrin_inv, uvz_sample
        )  # (N, 3, C * sample_num)
        camxyz_sample[:, 0] = -camxyz_sample[:, 0]  # (N, 3, C * sample_num)
        eexyz_sample = (
            torch.matmul(
                camxyz_sample.transpose(1, 2).contiguous(),
                (self.ego_cam_extrin[:3, :3]).t(),
            )
            + self.ego_cam_extrin[:3, 3]
        )  # (N, C * sample_num, 3)

        eewpad_xyz_sample = torch.cat(
            (eexyz_sample, eepad_pts), dim=1
        )  # (N, C*sample_num+40, 3)

        ## concat class info
        eexyzC_sample = torch.cat(
            (eewpad_xyz_sample, self.seg_info.repeat(N, 1, 1)), dim=2
        )  # (N, C*sample_num+40, 3+C+1)

        ## PointNet to get 128 feat
        pn_feature, _, _ = self.pn(
            eexyzC_sample.transpose(2, 1)
        )  # (N, 3+C+1, C*sample_num+40)

        state = torch.cat([qpos, ee_pos_base], dim=1).repeat(1, self.state_repeat_times)

        return torch.cat([state, pn_feature], dim=1)


class CustomRGBDEncodExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space: gym.spaces.Dict, state_repeat_times=2):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomRGBDEncodExtractor, self).__init__(
            observation_space, features_dim=1
        )
        self.state_repeat_times = state_repeat_times
        # load the Attention model
        # siamese_net = SiameseNet(
        #     input_channels=[3, 1],
        #     filters=[8],
        #     kernel_sizes=[5],
        #     strides=[1],
        #     activation="lrelu",
        #     norm=None,
        # )
        # self.gene_cnn = GeneralCNN(
        #     siamese_net=siamese_net,
        #     filters=[16, 16],
        #     kernel_sizes=[5, 5],
        #     strides=[2, 2],
        #     norm=None,
        #     activation="lrelu",
        # )
        # self.gene_cnn.train()
        siamese_net = SiameseNet(
            input_channels=[3, 1],
            filters=[8],
            kernel_sizes=[5],
            strides=[1],
            activation="lrelu",
            norm=None,
        )
        self.unet_encod = UNetEncoder(
            siamese_net=siamese_net,
            filters=[16, 16],
            kernel_sizes=[5, 5],
            strides=[2, 2],
            norm=None,
            activation="lrelu",
        ).to(torch.device("cuda:0"))
        self.unet_encod.train()

        # Update the features dim manually
        self._features_dim = 12 * self.state_repeat_times + 16

    def forward(self, observations) -> torch.Tensor:
        qpos = observations["qpos"]
        ee_pos_base = observations["ee_pos_base"]
        rgb = observations["rgb"]
        rgb = rgb / 255.0  # Normalization

        depth = observations["depth"]  # (N, 1, H, W)
        depth_norm = torch.clip(depth - torch.mean(depth), -1, 1)
        encod_latent = self.unet_encod([rgb, depth_norm])  # (N, 16)
        # gene_latent = self.gene_cnn([rgb, depth_norm])  # (N, 128)

        # state = torch.cat([qpos, ee_pos_base, encod_latent], dim=1).repeat(1, self.state_repeat_times)
        state = torch.cat([qpos, ee_pos_base], dim=1).repeat(1, self.state_repeat_times)

        return torch.cat([state, encod_latent], dim=1)


class CustomPCSegExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space: gym.spaces.Dict, state_repeat_times=1):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomPCSegExtractor, self).__init__(observation_space, features_dim=1)
        self.state_repeat_times = state_repeat_times

        self.attn = PointNetDenseCls(in_ch=3, mlp_specs=[64, 128, 512], k=2).to(
            torch.device("cuda:0")
        )
        self.attn.train()

        # Update the features dim manually
        self._features_dim = 15 * self.state_repeat_times

    def forward(self, observations) -> torch.Tensor:
        world_xyz = observations["xyz"]
        bs, N, _ = world_xyz.shape
        # handle_seg = observations["handle_seg"]
        qpos = observations["qpos"]
        ee_pos_base = observations["ee_pos_base"]
        ee_posemat = observations["ee_posemat_world"]  # (bs,4,4)

        ## normalization
        # compute the lenght of max(X,Y,Z), scale to [-1,1]
        # m = torch.max(torch.sqrt(torch.sum(xyz ** 2, dim=1)))
        # xyz_norm = xyz / m
        # pn2_sigmap = self.pn2score_net(torch.cat((xyz_norm, rgb), dim=2))  # (bs, N, 6) => (bs, N, 2)
        # attnmap = torch.max(pn2_sigmap, dim=2)[0]  # (bs, N)

        # attnmap, _, _ = self.pn_seg(torch.cat((xyz_norm, rgb), dim=2).transpose(2, 1).contiguous())   # (bs, N, 1)
        # # max_indice = torch.argmax(attnmap[:, :, 0], dim=1)  # (bs, N) => (bs,1)
        # max_indice = torch.argmax(attnmap[:, :, 0], dim=1)  # (bs, N) => (bs,1)
        # max_ee_xyz = torch.stack([xyz[i, max_indice[i], :] for i in range(bs)])  # (N, 3)
        # pn_feature, _, _ = self.pn(torch.cat((xyz_norm, rgb, attnmap), dim=2).transpose(2, 1).contiguous())  # (bs, N, 7) => (bs, 128)

        # (bs, 128), (bs, N, 2)
        _, attnmap, _, _ = self.attn(world_xyz.transpose(2, 1).contiguous())
        max_indice = torch.argmax(attnmap[:, :, 0], dim=1)  # (bs, N) => (bs,1)
        max_world_xyz = torch.stack(
            [world_xyz[i, max_indice[i], :] for i in range(bs)]
        )  # (bs, 3)

        max_ee_xyz = (
            torch.einsum("ij,ikj->ik", max_world_xyz, ee_posemat[:, :3, :3])
            + ee_posemat[:, :3, 3]
        )  # (bs, 3)

        state = torch.cat([qpos, ee_pos_base, max_ee_xyz], dim=1).repeat(
            1, self.state_repeat_times
        )  # (bs, 15)

        return state


class CustomPCRGBSegGTExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space: gym.spaces.Dict):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomPCRGBSegGTExtractor, self).__init__(
            observation_space, features_dim=1
        )

        self.pn = PointNetfeat(
            in_ch=5,
            mlp_specs=[64, 128, 1024],
            xyz_transform=False,
            feature_transform=False,
        )
        self.pn.train()

        # Update the features dim manually
        self._features_dim = 128 + 12 * 10

    def forward(self, observations) -> torch.Tensor:
        xyz = observations["xyz"]
        # rgb = observations["rgb"]
        handle_seg = observations["handle_seg"]
        qpos = observations["qpos"]
        ee_pos_base = observations["ee_pos_base"]
        # check_grasp = observations["check_grasp"]

        ## normalization
        # rgb = rgb * 2 - 1
        # compute the lenght of max(X,Y,Z)
        # m = torch.max(torch.sqrt(torch.sum(xyz ** 2, dim=1)))
        # scale to [-1,1]
        # xyz_norm = xyz / m

        bs, N, _ = xyz.shape
        # (bs, 128), (bs, N, 1)
        pn_feature, _, _, _ = self.pn(
            torch.cat((xyz, handle_seg), dim=2).transpose(2, 1).contiguous()
        )
        # max_indice = torch.argmax(handle_seg[:, :, 0], dim=1)  # (bs, N) => (bs,1)
        # max_ee_xyz = torch.stack([xyz[i, max_indice[i], :] for i in range(bs)])  # (bs, 3)

        state = torch.cat([qpos, ee_pos_base], dim=1).repeat(1, 10)  # (bs, 120)

        return torch.cat([state, pn_feature], dim=1)


class CustomPCSegGTExtractor(BaseFeaturesExtractor):
    def __init__(
        self, observation_space: gym.spaces.Dict, state_repeat_times=2, pts_num=300
    ):
        # We do not know features-dim here before going over all the items,
        # so put something dummy for now. PyTorch requires calling
        # nn.Module.__init__ before adding modules
        super(CustomPCSegGTExtractor, self).__init__(observation_space, features_dim=1)
        self.state_repeat_times = state_repeat_times
        # self.pts_num = pts_num
        # self.pn = PointNetfeat(in_ch=3, mlp_specs=[64, 128, 1024], xyz_transform=False, feature_transform=False)
        # self.pn.train()
        self.weight_centers = torch.nn.Parameter(
            torch.ones(2) * 0.5, requires_grad=True
        )

        # Update the features dim manually
        self._features_dim = 15 * self.state_repeat_times

    def forward(self, observations) -> torch.Tensor:
        ee_xyz = observations["ee_xyz"]  ## (N, 3, H, W)
        # rgb = observations["rgb"]
        handle_seg = observations["handle_seg"]  ## (N, 1, H, W)
        qpos = observations["qpos"]
        ee_pos_base = observations["ee_pos_base"]
        # check_grasp = observations["check_grasp"]

        N, ch, h, w = ee_xyz.shape
        eepts = ee_xyz.view(N, 3, -1)
        handel_mask = handle_seg.view(N, -1).bool()
        background_mask = ~handel_mask
        ee_ptcenter = torch.zeros((N, 3), device="cuda:0")
        background_ptcenter = torch.zeros((N, 3), device="cuda:0")
        for i in range(N):
            if torch.any(handel_mask[i]):
                ee_ptcenter[i] = (eepts[i, :, handel_mask[i]]).mean(1)  # (3)
                # zc = torch.stack([-depth.view(N, -1)[i, max_indice[i]] for i in range(N)])  # (N, sample_num), too much time
                # eepts[i] = ((-depth.view(N, -1)).gather(1, max_indice.unsqueeze(1))).squeeze(1)
            # else:
            #     eepts.append((ee_xyz.view(N, 3, -1))[i, :, np.random.choice(np.arange(h * w))])
            background_ptcenter[i] = eepts[i, :, background_mask[i]].mean(1)
        weight = F.softmax(self.weight_centers, dim=0)
        eepts_pred = weight[0] * ee_ptcenter + weight[1] * background_ptcenter
        # weights = handle_seg.view(N, -1)
        ## argmax to get maxpoint ind
        # max_indice = torch.argmax(weights, dim=1)

        ## multinomial sampling
        # weights[weights.ge(0.5)] = 200  # handle_ratio: 72*128/42 ~ 220
        # weights[weights.lt(0.5)] = 1
        # indices = torch.multinomial(weights, self.pts_num, replacement=False)  # (N, pts_num)
        # eexyz_pred = torch.stack([ee_xyz.view(N, 3, -1)[i, :, indices[i]] for i in range(N)])  # (N, 3, pts_num)
        # eepts_pred = torch.stack([ee_xyz.view(N, 3, -1)[i, :, max_indice[i]] for i in range(N)])  # (N, 3)

        # pn_feature, _, _, _ = self.pn(eexyz_pred)
        # max_indice = torch.argmax(handle_seg[:, :, 0], dim=1)  # (bs, N) => (bs,1)
        # max_ee_xyz = torch.stack([xyz[i, max_indice[i], :] for i in range(bs)])  # (bs, 3)

        state = torch.cat([qpos, ee_pos_base, eepts_pred], dim=1).repeat(
            1, self.state_repeat_times
        )  # (bs, 120)

        return state

arti_mani/algorithms/rl_iam/feature_extract.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import time
from typing import Optional, Tuple

import gym
import numpy as np
import torch
from arti_mani import ASSET_DIR
from arti_mani.utils.wrappers import NormalizeActionWrapper, RenderInfoWrapper
from stable_baselines3.common.vec_env import VecMonitor
from stable_baselines3.common.vec_env.base_vec_env import (
    VecEnv,
    VecEnvObs,
    VecEnvStepReturn,
    VecEnvWrapper,
)


class Custom_VecMonitor(VecMonitor):
    def __init__(
        self,
        venv: VecEnv,
        filename: Optional[str] = None,
        info_keywords: Tuple[str, ...] = (),
    ):
        super().__init__(
            venv,
            filename,
            info_keywords,
        )
        self.success_flag = np.zeros(self.num_envs, dtype=np.bool_)

    def reset(self) -> VecEnvObs:
        obs = self.venv.reset()
        self.episode_returns = np.zeros(self.num_envs, dtype=np.float32)
        self.episode_lengths = np.zeros(self.num_envs, dtype=np.int32)
        self.success_flag = np.zeros(self.num_envs, dtype=np.bool_)
        return obs

    def step_wait(self) -> VecEnvStepReturn:
        obs, rewards, dones, infos = self.venv.step_wait()
        for i in range(self.num_envs):
            if infos[i]["is_success"] == 1:
                arti_fail_rate = self.venv.unwrapped.get_attr("_arti_fail_rate")

        self.episode_returns += rewards
        self.episode_lengths += 1
        new_infos = list(infos[:])
        for i in range(len(dones)):
            if dones[i]:
                info = infos[i].copy()
                episode_return = self.episode_returns[i]
                episode_length = self.episode_lengths[i]
                episode_info = {
                    "r": episode_return,
                    "l": episode_length,
                    "t": round(time.time() - self.t_start, 6),
                }
                for key in self.info_keywords:
                    episode_info[key] = info[key]
                info["episode"] = episode_info
                self.episode_count += 1
                self.episode_returns[i] = 0
                self.episode_lengths[i] = 0
                if self.results_writer:
                    self.results_writer.write_row(episode_info)
                new_infos[i] = info
        return obs, rewards, dones, new_infos


def sb3_make_env(env_id, seed=0, obs_mode="state_depth", control_mode="pd_joint_vel"):
    """
    Utility function for custom env.

    :param env_id: (str) the environment ID
    :param seed: (int) the inital seed for RNG
    :param rank: (int) index of the subprocess
    """
    env = gym.make(
        env_id,
        articulation_config_path=ASSET_DIR
        / "partnet_mobility_configs/1999-cabinet-door-real-1.yml",
        obs_mode=obs_mode,  # "state", "state_dict", "rgbd", "pointcloud"
        control_mode=control_mode,  # "pd_joint_delta_pos", "pd_ee_delta_pose", "pd_ee_delta_pos"
        process_pc_mode="ms2_cam",  # "ms2_cam", "ms2_gt", "ms2_real", "ms2_realfull"
        gen_pc_mode="cam_depth",  ## "cam_depth", "sensor_depth"
        reward_mode="dense",  # "sparse", "dense"
    )
    env = NormalizeActionWrapper(env)
    env.seed(seed)
    return env


def sb3_make_evalenv(
    env_id, seed=0, obs_mode="state_depth", control_mode="pd_joint_vel"
):
    """
    Utility function for custom env.

    :param env_id: (str) the environment ID
    :param seed: (int) the inital seed for RNG
    :param rank: (int) index of the subprocess
    """
    env = gym.make(
        env_id,
        articulation_config_path=ASSET_DIR
        / "partnet_mobility_configs/1999-cabinet-door-real-1.yml",
        obs_mode=obs_mode,  # "state", "state_dict", "rgbd", "pointcloud"
        control_mode=control_mode,  # "pd_joint_delta_pos", "pd_ee_delta_pose", "pd_ee_delta_pos"
        process_pc_mode="ms2_cam",  # "ms2_cam", "ms2_gt", "ms2_real", "ms2_realfull"
        gen_pc_mode="cam_depth",  ## "cam_depth", "sensor_depth"
        reward_mode="dense",  # "sparse", "dense"
    )
    env = NormalizeActionWrapper(env)
    env = RenderInfoWrapper(env)
    env.seed(seed)
    return env


def sb3_make_env_multiinput(
    env_id, rank, seed=0, obs_mode="state_depth", control_mode="pd_joint_vel"
):
    """
    Utility function for multiprocessed env.

    :param env_id: (str) the environment ID
    :param rank: (int) index of the subprocess
    :param seed: (int) the inital seed for RNG
    :param obs_mode: (str) observation mode
    :param control_mode: (str) control mode
    """

    def _init():
        env = gym.make(
            env_id,
            articulation_ids=np.array(["1999-cabinet-door-real-1"]),
            articulation_config_path=ASSET_DIR / "partnet_mobility_configs/",
            obs_mode=obs_mode,  # "state", "state_dict", "rgbd", "pointcloud", "state_depth"
            control_mode=control_mode,  # "pd_joint_delta_pos", "pd_joint_vel", "pd_ee_delta_pose", "pd_ee_delta_pos"
            reward_mode="dense",  # "sparse", "dense"
        )
        env = NormalizeActionWrapper(env)
        # Important: use a different seed for each environment
        env.seed(seed + rank)
        return env

    return _init


def sb3_make_env_multicab(
    env_id,
    arti_ids,
    arti_config_path,
    obs_mode="state_depth",
    control_mode="pd_joint_vel",
    rank=0,
    seed=0,
):
    """
    Utility function for multiprocessed env.

    :param env_id: (str) the environment ID
    :param rank: (int) index of the subprocess
    :param seed: (int) the inital seed for RNG
    :param obs_mode: (str) observation mode
    :param control_mode: (str) control mode
    """

    def _init():
        env = gym.make(
            env_id,
            articulation_ids=arti_ids,
            articulation_config_path=arti_config_path,
            obs_mode=obs_mode,  # "state", "state_dict", "rgbd", "pointcloud", "state_depth"
            control_mode=control_mode,  # "pd_joint_delta_pos", "pd_joint_vel", "pd_ee_delta_pose", "pd_ee_delta_pos"
            reward_mode="dense",  # "sparse", "dense"
        )
        env = NormalizeActionWrapper(env)
        # Important: use a different seed for each environment
        env.seed(seed + rank)
        return env

    return _init


def sb3_make_env_multiarti(
    env_id,
    arti_ids,
    segmodel_path,
    sample_mode,
    frame_num,
    sample_num,
    other_handle_visible,
    obs_mode,
    control_mode,
    num_classes=6,
    device="",
    rank=0,
    seed=0,
):
    """
    Utility function for multiprocessed env.

    :param env_id: (str) the environment ID
    :param rank: (int) index of the subprocess
    :param seed: (int) the inital seed for RNG
    :param obs_mode: (str) observation mode
    :param control_mode: (str) control mode
    """

    def _init():
        env = gym.make(
            env_id,
            articulation_ids=arti_ids,
            segmodel_path=segmodel_path,
            sample_mode=sample_mode,
            frame_num=frame_num,
            sample_num=sample_num,
            other_handle_visible=other_handle_visible,
            num_classes=num_classes,
            obs_mode=obs_mode,  # "state", "state_dict", "rgbd", "pointcloud", "state_depth"
            control_mode=control_mode,  # "pd_joint_delta_pos", "pd_joint_vel", "pd_ee_delta_pose", "pd_ee_delta_pos"
            reward_mode="dense",  # "sparse", "dense"
            device=device,
        )
        env = NormalizeActionWrapper(env)
        # Important: use a different seed for each environment
        env.seed(seed + rank)
        return env

    return _init


def compute_zvupred(pred_vhp):
    N, Z, H, W = pred_vhp.shape
    indexes = torch.argmax(pred_vhp.view(N, -1), dim=1)
    zvu_pred = torch.zeros((pred_vhp.shape[0], 3), device=indexes.device)
    zvu_pred[:, 0] = torch.div(indexes, (H * W), rounding_mode="floor")
    zvu_pred[:, 1] = torch.div(
        indexes - (H * W) * zvu_pred[:, 0], W, rounding_mode="floor"
    )
    zvu_pred[:, 2] = indexes - (H * W) * zvu_pred[:, 0] - W * zvu_pred[:, 1]
    # zvu_pred = get_local_maxima_3d(vhmaps=pred_vhp, threshold=0.1, device=pred_vhp.device)
    # zvu_pred[:, 0] = zvu_pred[:, 0] / 39.0
    # zvu_pred[:, 1] = zvu_pred[:, 1] / 36.0
    # zvu_pred[:, 2] = zvu_pred[:, 2] / 64.0
    # zvu_pred = zvu_pred / torch.tensor((39.0, 36.0, 64.0), device=zvu_pred.device)
    # zvu_pred = zvu_pred / torch.tensor(hmap_size, device=zvu_pred.device)
    hmap_size = torch.tensor([32, 32, 32], device=indexes.device)
    xyz_min = torch.tensor([0.4, -0.7, 0.1], device=indexes.device)
    xyz_max = torch.tensor([1.2, 0.1, 0.9], device=indexes.device)
    zvu_pred = (zvu_pred + 0.5) / hmap_size * (
        xyz_max - xyz_min
    ) + xyz_min  # pred zvu in unit (m)

    return zvu_pred

arti_mani/algorithms/rl_iam/rl_utils.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
from collections import defaultdict

import gym
from arti_mani import RLMODEL_DIR
from arti_mani.envs import *
from arti_mani.utils.common import convert_np_bool_to_float
from arti_mani.utils.visualization import images_to_video
from arti_mani.utils.wrappers import NormalizeActionWrapper, RenderInfoWrapper
from stable_baselines3.sac import SAC
from tqdm import tqdm

if __name__ == "__main__":
    device = "cuda:0"
    mode = "door"  # "door", "drawer", "faucet", "arti"
    test_num = 50
    save_num = 10
    MAX_STEPS = 100
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069]
    if mode == "door":
        arti_ids = [0]  # , 1006, 1030, 1047, 1081
    elif mode == "drawer":
        arti_ids = [1]
    elif mode == "faucet":
        arti_ids = [5052]
    elif mode == "arti":
        arti_ids = [0, 1, 5052]
    else:
        raise NotImplementedError(mode)

    other_handle_visible = False
    obs_mode = "state_egostereo_dexpoints"
    control_mode = "pd_joint_delta_pos"
    segmodel_path = None
    sample_mode = None
    frame_num = 0
    num_classes = 2
    state_repeat_times = 10
    pts_sample_num = 0
    env = gym.make(
        "ArtiMani-v0",
        articulation_ids=arti_ids,
        segmodel_path=segmodel_path,
        sample_mode=sample_mode,
        frame_num=frame_num,
        sample_num=pts_sample_num,
        other_handle_visible=other_handle_visible,
        obs_mode=obs_mode,
        control_mode=control_mode,
        device=device,
        reward_mode="dense",  # "sparse", "dense"
    )
    env = NormalizeActionWrapper(env)
    env = RenderInfoWrapper(env)
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    env.seed(eval_seed)

    if mode == "arti":
        rl_exp_names = sorted(
            [
                f"hybrid/Dexpoints-based RL/{name}"
                for name in os.listdir(RLMODEL_DIR / "hybrid/Dexpoints-based RL")
                if mode in name
            ]
        )
    else:
        rl_exp_names = sorted(
            [
                f"Dexpoints-based RL/{name}"
                for name in os.listdir(RLMODEL_DIR / "Dexpoints-based RL")
                if mode in name
            ]
        )

    success_rates = defaultdict(list)
    success_steps = defaultdict(list)
    for rl_exp_name in rl_exp_names:
        model_path = RLMODEL_DIR / f"{rl_exp_name}/best_model"
        RL_model = SAC.load(
            model_path,
            env=env,
            print_system_info=True,
        )
        for arti_id in arti_ids:
            eval_videos_path = (
                RLMODEL_DIR / f"{rl_exp_name}/eval_videos/{arti_id}_best_model/"
            )
            if not os.path.exists(eval_videos_path):
                os.makedirs(eval_videos_path)

            is_success = []
            exp_steps = []
            with torch.no_grad():
                for num in tqdm(
                    range(test_num),
                    desc=f"Processing best_model",
                    colour="green",
                    leave=True,
                ):
                    obs = env.reset(articulation_id=arti_id)
                    if num < save_num:
                        frames = []
                        base_rgbd = env.render("cameras")
                        frames.append(base_rgbd)
                    success_flag = False
                    total_steps = 0
                    for step in tqdm(range(MAX_STEPS), colour="red", leave=False):
                        action, _states = RL_model.predict(obs, deterministic=True)
                        obs, rewards, dones, info = env.step(action)
                        # print(f"{step} step: {info['is_success']}")
                        if num < save_num:
                            base_rgbd = env.render("cameras")
                            frames.append(base_rgbd)
                        if not success_flag and info["is_success"] == 1.0:
                            success_flag = True
                            total_steps = step
                            break
                    if num < save_num:
                        images_to_video(
                            images=frames,
                            output_dir=eval_videos_path,
                            video_name=f"{num}",
                            verbose=False,
                        )
                    is_success.append(convert_np_bool_to_float(success_flag))
                    exp_steps.append(total_steps)
            success_rate = np.sum(is_success) / test_num
            exp_steps = np.array(exp_steps)
            success_mask = exp_steps > 0
            if np.any(success_mask):
                mean_steps = np.mean(exp_steps[success_mask])
            else:
                mean_steps = MAX_STEPS
            print(f"{arti_id}, success_rate, mean_steps: {success_rate}, {mean_steps}")
            success_rates[arti_id].append(success_rate)
            success_steps[arti_id].append(mean_steps)

    for key in success_rates.keys():
        print("success rates: ", key, success_rates[key])
    for key in success_steps.keys():
        print("success steps: ", key, success_steps[key])

arti_mani/algorithms/rl_iam/sac/sac_eval_dexpoints_PNfeat.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import gym
from arti_mani import RLMODEL_DIR, VISUALMODEL_DIR
from arti_mani.envs import *
from arti_mani.utils.common import convert_np_bool_to_float
from arti_mani.utils.wrappers import NormalizeActionWrapper, RenderInfoWrapper
from stable_baselines3.sac import SAC
from tqdm import tqdm


def load_vismodel(vismodel_path, load_device):
    model_path = VISUALMODEL_DIR / f"smp_model/{vismodel_path}/best.pth"
    config_path = VISUALMODEL_DIR / f"smp_model/{vismodel_path}/config.yaml"
    with open(config_path, "r") as f:
        cfg = yaml.safe_load(f)
    smp_cfg = cfg["smp_config"]
    if smp_cfg["mode"] == "RGBD":
        in_ch = 4
    elif smp_cfg["mode"] == "RGB":
        in_ch = 3
    elif smp_cfg["mode"] == "D":
        in_ch = 1
    else:
        raise NotImplementedError
    segmodel = CustomUnet(
        has_dropout=False,
        encoder_name=smp_cfg["encoder"],
        encoder_depth=smp_cfg["encoder_depth"],
        decoder_channels=smp_cfg["decoder_channels"],
        encoder_weights=smp_cfg["encoder_weights"],
        in_channels=in_ch,
        classes=cfg["num_classes"],
        activation=smp_cfg["activation"],
    )
    segmodel.load_state_dict(torch.load(model_path))
    segmodel.to(torch.device(load_device))
    segmodel.eval()
    return segmodel


if __name__ == "__main__":
    device = "cuda:0"
    # Ours, Oracle-based RL, Image-based RL, Keypoints-based RL, Dexpoints-based RL
    baseline_mode = "hybrid/Dexpoints-based RL"
    mode = "arti"  # "door", "drawer", "faucet", "arti"
    test_num = 50
    MAX_STEPS = 100
    if mode == "door":
        arti_ids = [0]
    elif mode == "drawer":
        arti_ids = [1]
    elif mode == "faucet":
        arti_ids = [5052]
    elif mode == "arti":
        arti_ids = [0, 1, 5052]
    else:
        raise NotImplementedError(mode)

    if "Image" in baseline_mode or "Ours" == baseline_mode:
        smp_exp_name = "20230219_000940_train52-val18_384_noDR_randombg_aug_stereo_bs16_focalloss_0.5step50lr0.001_RGBDunet-163264128_mobilenet_v2"
        segmodel_path = VISUALMODEL_DIR / f"smp_model/{smp_exp_name}"
    elif "Keypoints" in baseline_mode:
        smp_exp_name = "20230307_182821_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2_dropout0.2_newdatafilter2drawer"
        segmodel_path = VISUALMODEL_DIR / f"kpt_model/{smp_exp_name}"
    else:
        segmodel_path = None

    other_handle_visible = False
    # state_egostereo_segpoints, state_egostereo_dseg, state_egostereo_keypoints, state_egostereo_segpoints_gt, state_egostereo_dexpoints
    obs_mode = "state_egostereo_keypoints"
    control_mode = "pd_joint_delta_pos"
    sample_mode = "random_sample"  # full_downsample, score_sample, random_sample, frameweight_sample, fps_sample
    frame_num = 1
    num_classes = 6
    state_repeat_times = 10
    pts_sample_num = 32
    rl_exp_name = f"{baseline_mode}/sac4_tr12eval3arti_state_egostereo_keypoints3_bs256_tr16up16_5"

    env = gym.make(
        "ArtiMani-v0",
        articulation_ids=arti_ids,
        segmodel_path=segmodel_path,
        sample_mode=sample_mode,
        frame_num=frame_num,
        sample_num=pts_sample_num,
        other_handle_visible=other_handle_visible,
        obs_mode=obs_mode,
        control_mode=control_mode,
        device=device,
        reward_mode="dense",  # "sparse", "dense"
    )
    env = NormalizeActionWrapper(env)
    env = RenderInfoWrapper(env)
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    env.seed(eval_seed)

    model_path = RLMODEL_DIR / f"{rl_exp_name}/best_model"
    RL_model = SAC.load(
        model_path,
        env=env,
        print_system_info=True,
    )
    for arti_id in arti_ids:
        is_success = []
        exp_steps = []
        with torch.no_grad():
            for num in tqdm(
                range(test_num),
                desc=f"Processing best model",
                colour="green",
                leave=True,
            ):
                obs = env.reset(articulation_id=arti_id)
                success_flag = False
                total_steps = 0
                # for step in range(MAX_STEPS):
                for step in tqdm(range(MAX_STEPS), colour="red", leave=False):
                    action, _states = RL_model.predict(obs, deterministic=True)
                    obs, rewards, dones, info = env.step(action)
                    # print(f"{step} step: {info['is_success']}")
                    if not success_flag and info["is_success"] == 1.0:
                        success_flag = True
                        total_steps = step
                        break
                is_success.append(convert_np_bool_to_float(success_flag))
                exp_steps.append(total_steps)
        success_rate = np.sum(is_success) / test_num
        exp_steps = np.array(exp_steps)
        success_mask = exp_steps > 0
        if np.any(success_mask):
            mean_steps = np.mean(exp_steps[success_mask])
        else:
            mean_steps = MAX_STEPS
        print(f"{arti_id}, success_rate, mean_steps: {success_rate}, {mean_steps}")

arti_mani/algorithms/rl_iam/sac/sac_eval_hybrid_multiarti.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import glob
import os
from collections import defaultdict

import cv2
import gym
import matplotlib.pyplot as plt
from arti_mani import RLMODEL_DIR, VISUALMODEL_DIR
from arti_mani.envs import *
from arti_mani.utils.common import convert_np_bool_to_float
from arti_mani.utils.visualization import images_to_video
from arti_mani.utils.wrappers import NormalizeActionWrapper, RenderInfoWrapper
from stable_baselines3.sac import SAC
from tqdm import tqdm


def vis_keypoints(rgb, uvz, uvz_pred, uvz_visable, save_filename):
    rgb_kptgt = rgb.copy()
    for idn in range(uvz.shape[0]):
        if uvz_visable[idn, 0] and uvz_visable[idn, 1]:
            if uvz_visable[idn, 2]:
                cv2.circle(
                    rgb_kptgt,
                    (round(uvz[idn, 0]), round(uvz[idn, 1])),
                    radius=2,
                    color=(1, 0, 0),
                    thickness=-1,
                )
            else:
                cv2.circle(
                    rgb_kptgt,
                    (round(uvz[idn, 0]), round(uvz[idn, 1])),
                    radius=2,
                    color=(0, 1, 0),
                    thickness=-1,
                )
    rgb_kptpred = rgb.copy()
    for idn in range(uvz_pred.shape[0]):
        if uvz_visable[idn, 0] and uvz_visable[idn, 1]:
            if uvz_visable[idn, 2]:
                cv2.circle(
                    rgb_kptpred,
                    (round(uvz_pred[idn, 0]), round(uvz_pred[idn, 1])),
                    radius=2,
                    color=(1, 0, 0),
                    thickness=-1,
                )
            else:
                cv2.circle(
                    rgb_kptpred,
                    (round(uvz_pred[idn, 0]), round(uvz_pred[idn, 1])),
                    radius=2,
                    color=(0, 1, 0),
                    thickness=-1,
                )
    plt.subplot(121)
    plt.imshow(rgb_kptgt)
    plt.title("GT")
    plt.axis("off")
    plt.subplot(122)
    plt.imshow(rgb_kptpred)
    plt.title("Pred")
    plt.axis("off")
    plt.savefig(save_filename, bbox_inches="tight", pad_inches=0.0)


if __name__ == "__main__":
    device = "cuda:0"
    mode = "door"  # "door", "drawer", "faucet"
    test_num = 10
    save_num = 10
    MAX_STEPS = 100
    rl_exp_name = "sac4_tr4eval1door_state_egostereo_keypoints3_bs128_tr16up16"
    smp_exp_name = "20230307_182821_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2_dropout0.2_newdatafilter2drawer"
    eval_steps = []
    for file in glob.glob(str(RLMODEL_DIR / f"{rl_exp_name}/rl_model_*_steps.zip")):
        eval_steps.append(int(os.path.basename(file).split("_")[2]))
    eval_steps.sort()
    eval_steps = np.array(eval_steps)

    env_id = "ArtiMani-v0"
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069]
    if mode == "door":
        arti_ids = [0, 1006, 1030, 1047, 1081]
    elif mode == "drawer":
        arti_ids = [1, 1004, 1005, 1016, 1024]
    elif mode == "faucet":
        arti_ids = [5052, 5004, 5007, 5023, 5069]
    else:
        raise NotImplementedError(mode)
    kptmodel_path = VISUALMODEL_DIR / f"kpt_model/{smp_exp_name}"
    other_handle_visible = False
    obs_mode = "state_egostereo_keypoints"
    control_mode = "pd_joint_delta_pos"

    env = gym.make(
        env_id,
        articulation_ids=arti_ids,
        segmodel_path=kptmodel_path,
        sample_mode=None,
        frame_num=0,
        sample_num=0,
        other_handle_visible=other_handle_visible,
        obs_mode=obs_mode,
        control_mode=control_mode,
        device=device,
        reward_mode="dense",  # "sparse", "dense"
    )
    env = NormalizeActionWrapper(env)
    env = RenderInfoWrapper(env)
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    env.seed(eval_seed)

    success_rates = defaultdict(list)
    success_steps = defaultdict(list)
    uvz_errors = defaultdict(list)
    for eval_step in eval_steps:
        model_path = RLMODEL_DIR / f"{rl_exp_name}/rl_model_{eval_step}_steps"
        RL_model = SAC.load(
            model_path,
            env=env,
            print_system_info=True,
        )
        for arti_id in arti_ids:
            eval_videos_path = (
                RLMODEL_DIR
                / f"{rl_exp_name}/eval_videos/{arti_id}/{eval_step / 1000000}M/"
            )
            if not os.path.exists(eval_videos_path):
                os.makedirs(eval_videos_path)

            is_success = []
            exp_steps = []
            uvz_error = []
            with torch.no_grad():
                for num in tqdm(
                    range(test_num),
                    desc=f"Processing {eval_step / 1000000}M",
                    colour="green",
                    leave=True,
                ):
                    obs = env.reset(articulation_id=arti_id)
                    if num < save_num:
                        frames = []
                        base_rgbd = env.render("cameras")
                        frames.append(base_rgbd)
                    success_flag = False
                    total_steps = MAX_STEPS
                    uvz_err = []
                    # for step in range(MAX_STEPS):
                    for step in tqdm(range(MAX_STEPS), colour="red", leave=False):
                        action, _states = RL_model.predict(obs, deterministic=True)
                        obs, rewards, dones, info = env.step(action)
                        # print(f"{step} step: {info['is_success']}")
                        if num < save_num:
                            base_rgbd = env.render("cameras")
                            frames.append(base_rgbd)
                        uvz = obs["uvz"]
                        uvz_pred = obs["uvz_pred"]
                        uvz_visable = obs["uvz_visable"]
                        uvz_err.append(np.abs(uvz_pred - uvz))
                        if num == 0:
                            rgb = env.get_handsensor_rgb()
                            vis_filename = os.path.join(
                                eval_videos_path, f"{num}_{step}.png"
                            )
                            vis_keypoints(rgb, uvz, uvz_pred, uvz_visable, vis_filename)
                        if not success_flag and info["is_success"] == 1.0:
                            success_flag = True
                            total_steps = step
                            break
                    if num < save_num:
                        images_to_video(
                            images=frames,
                            output_dir=eval_videos_path,
                            video_name=f"{num}",
                            verbose=False,
                        )
                    is_success.append(convert_np_bool_to_float(success_flag))
                    exp_steps.append(total_steps)
                    uvz_error.append(
                        np.array(uvz_err).mean(0)
                    )  # (steps, J, 3) -> (J, 3)
            success_rate = np.sum(is_success) / test_num
            mean_steps = np.sum(exp_steps) / test_num
            mean_kpt_err = np.array(uvz_error).mean(0)  # (test_num, J, 3) -> (J, 3)
            print(f"{arti_id}, success_rate, mean_steps: {success_rate}, {mean_steps}")
            print(
                f"{arti_id}, mean tip, center, bottom error: {mean_kpt_err[0]}, {mean_kpt_err[1]}, {mean_kpt_err[2]}"
            )
            success_rates[arti_id].append(success_rate)
            success_steps[arti_id].append(mean_steps)
            uvz_errors[arti_id].append(mean_kpt_err)

    fig, ax = plt.subplots()
    for arti_id in arti_ids:
        ax.plot(
            eval_steps / 1000000, success_rates[arti_id], label=arti_id, color="red"
        )
    ax.legend(loc="lower left")
    ax.set_xlabel("eval_model_steps(M)")
    ax.set_ylabel("success rate")
    ax2 = ax.twinx()
    for arti_id in arti_ids:
        ax2.plot(
            eval_steps / 1000000, success_steps[arti_id], label=arti_id, color="blue"
        )
    ax2.legend(loc="upper right")
    ax2.set_ylabel("mean steps")
    plt.savefig(RLMODEL_DIR / f"{rl_exp_name}/eval_videos/sr-evalsteps.png")

arti_mani/algorithms/rl_iam/sac/sac_eval_keypts.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
from collections import defaultdict

import gym
from arti_mani import RLMODEL_DIR
from arti_mani.envs import *
from arti_mani.utils.common import convert_np_bool_to_float
from arti_mani.utils.visualization import images_to_video
from arti_mani.utils.wrappers import NormalizeActionWrapper, RenderInfoWrapper
from stable_baselines3.sac import SAC

if __name__ == "__main__":
    device = "cuda:0"
    mode = "drawer"  # "door", "drawer", "faucet", "arti"
    test_num = 10
    save_num = 10
    success_nums = 0
    MAX_STEPS = 100
    if mode == "arti":
        rl_exp_names = sorted(
            [
                name
                for name in os.listdir(RLMODEL_DIR / "hybrid/Oracle-based RL")
                if mode in name
            ]
        )
    else:
        rl_exp_names = sorted(
            [
                name
                for name in os.listdir(RLMODEL_DIR / "Oracle-based RL")
                if mode in name
            ]
        )

    env_id = "ArtiMani-v0"
    if mode == "door":
        arti_ids = [0]
    elif mode == "drawer":
        arti_ids = [1]
    elif mode == "faucet":
        arti_ids = [5052]
    elif mode == "arti":
        arti_ids = [0, 1, 5052]
    else:
        raise NotImplementedError(mode)
    other_handle_visible = False
    obs_mode = "state_egostereo_segpoints_gt"
    control_mode = "pd_joint_delta_pos"
    sample_mode = "random_sample"  # full_downsample, score_sample, random_sample, frameweight_sample
    frame_num = 1
    num_classes = 6
    state_repeat_times = 10
    pts_sample_num = 32

    env = gym.make(
        env_id,
        articulation_ids=arti_ids,
        segmodel_path=None,
        sample_mode=sample_mode,
        frame_num=frame_num,
        sample_num=pts_sample_num,
        other_handle_visible=other_handle_visible,
        obs_mode=obs_mode,
        control_mode=control_mode,
        device=device,
        reward_mode="dense",  # "sparse", "dense"
    )
    env = NormalizeActionWrapper(env)
    env = RenderInfoWrapper(env)
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    env.seed(eval_seed)

    success_rates = defaultdict(list)
    success_steps = defaultdict(list)
    for rl_exp_name in rl_exp_names:
        model_path = RLMODEL_DIR / f"{rl_exp_name}/best_model"
        RL_model = SAC.load(
            model_path,
            env=env,
            print_system_info=True,
        )
        for arti_id in arti_ids:
            eval_videos_path = (
                RLMODEL_DIR / f"{rl_exp_name}/eval_videos/{arti_id}_best_model/"
            )
            if not os.path.exists(eval_videos_path):
                os.makedirs(eval_videos_path)

            is_success = []
            exp_steps = []
            with torch.no_grad():
                for num in range(test_num):
                    obs = env.reset(articulation_id=arti_id)
                    if num < save_num:
                        frames = []
                        base_rgbd = env.render("cameras")
                        frames.append(base_rgbd)
                    success_flag = False
                    total_steps = 0
                    for step in range(MAX_STEPS):
                        action, _states = RL_model.predict(obs, deterministic=True)
                        obs, rewards, dones, info = env.step(action)
                        # print(f"{i} step: {info['is_success']}")
                        if not success_flag and info["is_success"] == 1.0:
                            success_nums += 1
                            total_steps = step
                            success_flag = True
                        if num < save_num:
                            base_rgbd = env.render("cameras")
                            frames.append(base_rgbd)
                    if num < save_num:
                        images_to_video(
                            images=frames,
                            output_dir=eval_videos_path,
                            video_name=f"{num}",
                        )
                    is_success.append(convert_np_bool_to_float(success_flag))
                    exp_steps.append(total_steps)
            success_rate = np.sum(is_success) / test_num
            exp_steps = np.array(exp_steps)
            success_mask = exp_steps > 0
            if np.any(success_mask):
                mean_steps = np.mean(exp_steps[success_mask])
            else:
                mean_steps = MAX_STEPS
            print(f"{arti_id}, success_rate, mean_steps: {success_rate}, {mean_steps}")
            success_rates[arti_id].append(success_rate)
            success_steps[arti_id].append(mean_steps)

    for key in success_rates.keys():
        print("success rates: ", key, success_rates[key])
    for key in success_steps.keys():
        print("success steps: ", key, success_steps[key])

arti_mani/algorithms/rl_iam/sac/sac_eval_segGTpts_PNfeat.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import faulthandler
import os
from collections import defaultdict

import gym
from arti_mani import RLMODEL_DIR
from arti_mani.envs import *
from arti_mani.utils.common import convert_np_bool_to_float
from arti_mani.utils.visualization import images_to_video
from arti_mani.utils.wrappers import NormalizeActionWrapper, RenderInfoWrapper
from stable_baselines3.sac import SAC

faulthandler.enable()


def draw_gripper_pts(env, max_pos=0.068, radius=0.005, color=(1, 0, 0), name=""):
    builder = env.unwrapped._scene.create_actor_builder()
    cur_gripper_pos = np.mean(env.unwrapped.agent._robot.get_qpos()[-2:])
    # gripper_pts_ee = sample_grasp_points_ee(max_pos - cur_gripper_pos, z_offset=0.03)
    gripper_pts_ee = sample_grasp_multipoints_ee(
        max_pos - cur_gripper_pos, num_points_perlink=10, x_offset=0.02
    )
    trans_ee2world = env.unwrapped.agent.grasp_site.pose.to_transformation_matrix()
    gripper_pts_world = transform_points(trans_ee2world, gripper_pts_ee)
    actors = []
    for pts in gripper_pts_world:
        builder.add_sphere_visual(radius=radius, color=color)
        gripper_pts = builder.build_static(name=name)
        gripper_pts.set_pose(sapien.Pose(pts))
        actors.append(gripper_pts)
    return actors


def draw_target_pts(env, gripper_pts_world, radius=0.005, color=(0, 0, 1), name=""):
    builder = env.unwrapped._scene.create_actor_builder()
    actors = []
    for pts in gripper_pts_world:
        builder.add_sphere_visual(radius=radius, color=color)
        gripper_pts = builder.build_static(name=name)
        gripper_pts.set_pose(sapien.Pose(pts))
        actors.append(gripper_pts)
    return actors


if __name__ == "__main__":
    device = "cuda:0"
    mode = "arti"  # "door", "drawer", "faucet", "laptop", "kitchen_pot"
    test_num = 50
    save_num = 5
    success_nums = 0
    MAX_STEPS = 100
    # rl_exp_name = "sac4_tr4eval1laptop_state_egostereo_segpoints_gt_random_samplef1s32_gt_0"
    # rl_exp_name = "hybrid/5arti/sac4_tr20eval5arti_state_ego_segpoints_gt_random_samplef1s32_multiarti_5class_graspP2turnP1_potaligndrawer"
    rl_exp_names = [
        "hybrid/5arti/sac4_tr20eval5arti_state_ego_segpoints_gt_random_samplef1s32_multiarti_5class_graspP2turnP1_3",
        "hybrid/5arti/sac4_tr20eval5arti_state_ego_segpoints_gt_random_samplef1s32_multiarti_5class_graspP2turnP1_4",
        "hybrid/5arti/sac4_tr20eval5arti_state_ego_segpoints_gt_random_samplef1s32_multiarti_5class_graspP2turnP1_5",
        "hybrid/5arti/sac4_tr20eval5arti_state_ego_segpoints_gt_random_samplef1s32_multiarti_5class_graspP2turnP1_6",
    ]
    env_id = "ArtiMani-v0"
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069]
    arti_ids = {
        "door": [1006, 1030, 1047, 1081, 0],
        "drawer": [1005, 1016, 1024, 1076, 1],
        "faucet": [5004, 5007, 5023, 5069, 5052],
        "laptop": [9960, 9968, 9992, 9996, 9748],
        "kitchen_pot": [100051, 100054, 100055, 100060, 100015],
    }
    if mode == "door":
        arti_ids = [1006, 1030, 1047, 1081, 0]
    elif mode == "drawer":
        arti_ids = [1005, 1016, 1024, 1076, 1]
    elif mode == "faucet":
        arti_ids = [5004, 5007, 5023, 5069, 5052]
    elif mode == "laptop":
        arti_ids = [9960, 9968, 9992, 9996, 9748]
    elif mode == "kitchen_pot":
        arti_ids = [100051, 100054, 100055, 100060, 100015]
    elif mode == "arti":
        # arti_ids = [1006, 1030, 1047, 1081, 1005, 1016, 1024, 1076, 5004, 5007, 5023, 5069,
        #             9960, 9968, 9992, 9996, 100051, 100054, 100055, 100060, 0, 1, 5052, 9748, 100015]
        arti_ids = [0, 1, 5052, 9748, 100015]
    else:
        raise NotImplementedError(mode)
    other_handle_visible = False
    obs_mode = "state_ego_segpoints_gt"
    control_mode = "pd_joint_delta_pos"
    sample_mode = "random_sample"  # full_downsample, score_sample, random_sample, frameweight_sample
    frame_num = 1
    num_classes = 5
    state_repeat_times = 10
    pts_sample_num = 32
    # gripper_actors = None
    # target_grasp_actors = None

    env = gym.make(
        env_id,
        articulation_ids=arti_ids,
        segmodel_path=None,
        sample_mode=sample_mode,
        frame_num=frame_num,
        sample_num=pts_sample_num,
        other_handle_visible=other_handle_visible,
        num_classes=num_classes,
        obs_mode=obs_mode,
        control_mode=control_mode,
        device=device,
        reward_mode="dense",  # "sparse", "dense"
    )
    env = NormalizeActionWrapper(env)
    env = RenderInfoWrapper(env)
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    env.seed(eval_seed)

    success_rates = defaultdict(list)
    for rl_exp_name in rl_exp_names:
        model_path = RLMODEL_DIR / f"{rl_exp_name}/best_model"
        RL_model = SAC.load(
            model_path,
            env=env,
            print_system_info=True,
        )
        for arti_id in arti_ids:
            eval_videos_path = (
                RLMODEL_DIR / f"{rl_exp_name}/eval_videos/{arti_id}/best_model/"
            )
            if not os.path.exists(eval_videos_path):
                os.makedirs(eval_videos_path)

            is_success = []
            # gripper_actors, target_grasp_actors = None, None
            with torch.no_grad():
                for num in range(test_num):
                    obs = env.reset(articulation_id=arti_id)
                    if num < save_num:
                        frames = []
                        base_rgbd = env.render("cameras")
                        frames.append(base_rgbd)
                    success_flag = False
                    for i in range(MAX_STEPS):
                        action, _states = RL_model.predict(obs, deterministic=True)
                        obs, rewards, dones, info = env.step(action)
                        if not success_flag and info["is_success"] == 1.0:
                            success_nums += 1
                            success_flag = True
                        if num < save_num:
                            base_rgbd = env.render("cameras")
                            frames.append(base_rgbd)
                    if num < save_num:
                        images_to_video(
                            images=frames,
                            output_dir=eval_videos_path,
                            video_name=f"{num}",
                        )
                    is_success.append(convert_np_bool_to_float(success_flag))
            sum_num = np.sum(is_success)
            print(f"{arti_id}, success rate: {sum_num}/{test_num}")
            success_rates[arti_id].append(sum_num / test_num)
    print(success_rates)

arti_mani/algorithms/rl_iam/sac/sac_eval_segGTpts_PNfeat_multiarti.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
from collections import defaultdict

import cv2
import gym
import matplotlib.pyplot as plt
from arti_mani import RLMODEL_DIR, VISUALMODEL_DIR
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import CustomUnet
from arti_mani.envs import *
from arti_mani.utils.common import convert_np_bool_to_float
from arti_mani.utils.cv_utils import visualize_depth, visualize_seg
from arti_mani.utils.visualization import images_to_video
from arti_mani.utils.wrappers import NormalizeActionWrapper, RenderInfoWrapper
from stable_baselines3.sac import SAC
from tqdm import tqdm

CMAPS = [
    (1, 0, 0),  # "red"
    (0, 0, 1),  # "blue"
    (1, 1, 0),  # "yellow"
    (0, 1, 0),  # "green"
    (0.627, 0.125, 0.941),  # "purple"
    (0.753, 0.753, 0.753),  # "grey"
]


def load_vismodel(vismodel_path, load_device):
    model_path = VISUALMODEL_DIR / f"smp_model/{vismodel_path}/best.pth"
    config_path = VISUALMODEL_DIR / f"smp_model/{vismodel_path}/config.yaml"
    with open(config_path, "r") as f:
        cfg = yaml.safe_load(f)
    smp_cfg = cfg["smp_config"]
    if smp_cfg["mode"] == "RGBD":
        in_ch = 4
    elif smp_cfg["mode"] == "RGB":
        in_ch = 3
    elif smp_cfg["mode"] == "D":
        in_ch = 1
    else:
        raise NotImplementedError
    segmodel = CustomUnet(
        has_dropout=False,
        encoder_name=smp_cfg["encoder"],
        encoder_depth=smp_cfg["encoder_depth"],
        decoder_channels=smp_cfg["decoder_channels"],
        encoder_weights=smp_cfg["encoder_weights"],
        in_channels=in_ch,
        classes=cfg["num_classes"],
        activation=smp_cfg["activation"],
    )
    segmodel.load_state_dict(torch.load(model_path))
    segmodel.to(torch.device(load_device))
    segmodel.eval()
    return segmodel


def vis_sample_pts(
    rgbdseg, indices, num_classes, pts_sample_num, vis_model, device, save_filename
):
    assert indices.shape[0] == num_classes * pts_sample_num
    rgb, depth, seg = rgbdseg["rgb"] / 255.0, rgbdseg["depth"], rgbdseg["seg"]
    h, w = 144, 256
    v = indices // w  # (C * sample_num)
    u = indices - v * w  # (C * sample_num)
    rgb_plt = rgb.transpose(1, 2, 0)
    rgb_sample = rgb_plt.copy()
    for idn in range(num_classes):
        for pt_num in range(pts_sample_num):
            cv2.circle(
                rgb_sample,
                (
                    round(u[idn * pts_sample_num + pt_num]),
                    round(v[idn * pts_sample_num + pt_num]),
                ),
                radius=2,
                color=CMAPS[idn],
                thickness=-1,
            )

    with torch.no_grad():
        rgb_tensor = torch.from_numpy(rgb[None]).float().to(device)
        depth_tensor = torch.from_numpy(depth[None, None]).float().to(device)
        seg_pred = vis_model.predict(torch.cat((rgb_tensor, depth_tensor), dim=1))[0]
        seg_mc_label = torch.argmax(seg_pred, dim=0).cpu().numpy()

    plt.subplot(231)
    plt.imshow(rgb_plt)
    plt.title("rgb")
    plt.axis("off")
    plt.subplot(232)
    plt.imshow(visualize_depth(depth))
    plt.title("depth")
    plt.axis("off")
    plt.subplot(233)
    plt.imshow(visualize_seg(seg))
    plt.title("seg")
    plt.axis("off")
    plt.subplot(234)
    plt.imshow(visualize_seg(seg_mc_label, 6))
    plt.title("segpred")
    plt.axis("off")
    plt.subplot(235)
    plt.imshow(rgb_sample)
    plt.title("sample_pts")
    plt.axis("off")
    plt.savefig(save_filename, bbox_inches="tight", pad_inches=0.0)


if __name__ == "__main__":
    device = "cuda:0"
    mode = "arti"  # "door", "drawer", "faucet", "arti"
    test_num = 50
    save_num = 10
    MAX_STEPS = 100
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069]
    if mode == "door":
        arti_ids = [0]  # , 1006, 1030, 1047, 1081
    elif mode == "drawer":
        arti_ids = [1]
    elif mode == "faucet":
        arti_ids = [5004, 5007, 5023, 5069, 5052]
    elif mode == "arti":
        arti_ids = [0, 1, 5052]
    else:
        raise NotImplementedError(mode)

    smp_exp_name = "20230219_000940_train52-val18_384_noDR_randombg_aug_stereo_bs16_focalloss_0.5step50lr0.001_RGBDunet-163264128_mobilenet_v2"
    vis_model = load_vismodel(smp_exp_name, device)
    segmodel_path = VISUALMODEL_DIR / f"smp_model/{smp_exp_name}"
    other_handle_visible = False
    obs_mode = "state_egostereo_segpoints"
    control_mode = "pd_joint_delta_pos"
    sample_mode = "random_sample"  # full_downsample, score_sample, random_sample, frameweight_sample, fps_sample
    frame_num = 1
    num_classes = 6
    state_repeat_times = 10
    pts_sample_num = 32

    rl_exp_names = [
        "sac4_tr12eval3arti_state_egostereo_segpoints_random_samplef1s32_100steps_smp0219_largerange_cabinetrew_ptsaddnoise0.01"
    ]

    for rl_exp_name in rl_exp_names:
        eval_steps = [750000, 1000000, 1500000, 1750000, 2250000, 3000000]
        # for file in glob.glob(str(RLMODEL_DIR / f"{rl_exp_name}/rl_model_*_steps.zip")):
        #     eval_steps.append(int(os.path.basename(file).split("_")[2]))
        # eval_steps.sort()
        eval_steps = np.array(eval_steps)

        env = gym.make(
            "ArtiMani-v0",
            articulation_ids=arti_ids,
            segmodel_path=segmodel_path,
            sample_mode=sample_mode,
            frame_num=frame_num,
            sample_num=pts_sample_num,
            other_handle_visible=other_handle_visible,
            obs_mode=obs_mode,
            control_mode=control_mode,
            device=device,
            reward_mode="dense",  # "sparse", "dense"
        )
        env = NormalizeActionWrapper(env)
        env = RenderInfoWrapper(env)
        eval_seed = np.random.RandomState().randint(2**32)
        print("experiment eval random seed: ", eval_seed)
        env.seed(eval_seed)

        success_rates = defaultdict(list)
        success_steps = defaultdict(list)
        for eval_step in eval_steps:
            model_path = RLMODEL_DIR / f"{rl_exp_name}/rl_model_{eval_step}_steps"
            RL_model = SAC.load(
                model_path,
                env=env,
                print_system_info=True,
            )
            for arti_id in arti_ids:
                eval_videos_path = (
                    RLMODEL_DIR
                    / f"{rl_exp_name}/eval_videos/{arti_id}/{eval_step / 1000000}M/"
                )
                if not os.path.exists(eval_videos_path):
                    os.makedirs(eval_videos_path)

                is_success = []
                exp_steps = []
                with torch.no_grad():
                    # for num in range(test_num):
                    for num in tqdm(
                        range(test_num),
                        desc=f"Processing {eval_step / 1000000}M",
                        colour="green",
                        leave=True,
                    ):
                        obs = env.reset(articulation_id=arti_id)
                        if num < save_num:
                            frames = []
                            base_rgbd = env.render("cameras")
                            frames.append(base_rgbd)
                        success_flag = False
                        total_steps = 0
                        # for step in range(MAX_STEPS):
                        for step in tqdm(range(MAX_STEPS), colour="red", leave=False):
                            action, _states = RL_model.predict(obs, deterministic=True)
                            obs, rewards, dones, info = env.step(action)
                            # print(f"{step} step: {info['is_success']}")
                            if num < save_num:
                                base_rgbd = env.render("cameras")
                                frames.append(base_rgbd)
                            if num == 0:
                                rgbdseg = env.get_handsensor_rgbdseg()
                                indices = env.sampler.sample_indices  # C*sample_num
                                vis_filename = os.path.join(
                                    eval_videos_path, f"{num}_{step}.png"
                                )
                                vis_sample_pts(
                                    rgbdseg,
                                    indices,
                                    num_classes,
                                    pts_sample_num,
                                    vis_model,
                                    device,
                                    vis_filename,
                                )
                            if not success_flag and info["is_success"] == 1.0:
                                success_flag = True
                                total_steps = step
                                break
                        if num < save_num:
                            images_to_video(
                                images=frames,
                                output_dir=eval_videos_path,
                                video_name=f"{num}",
                                verbose=False,
                            )
                        is_success.append(convert_np_bool_to_float(success_flag))
                        exp_steps.append(total_steps)
                success_rate = np.sum(is_success) / test_num
                exp_steps = np.array(exp_steps)
                success_mask = exp_steps > 0
                if np.any(success_mask):
                    mean_steps = np.mean(exp_steps[success_mask])
                else:
                    mean_steps = MAX_STEPS
                print(
                    f"{arti_id}, success_rate, mean_steps: {success_rate}, {mean_steps}"
                )
                success_rates[arti_id].append(success_rate)
                success_steps[arti_id].append(mean_steps)

        fig, ax = plt.subplots()
        for arti_id in arti_ids:
            ax.plot(eval_steps / 1000000, success_rates[arti_id], "-", label=arti_id)
        ax.legend(loc="lower left")
        ax.set_xlabel("eval_model_steps(M)")
        ax.set_ylabel("success rate")
        ax2 = ax.twinx()
        for arti_id in arti_ids:
            ax2.plot(eval_steps / 1000000, success_steps[arti_id], "--", label=arti_id)
        ax2.legend(loc="upper right")
        ax2.set_ylabel("mean steps")
        plt.savefig(RLMODEL_DIR / f"{rl_exp_name}/eval_videos/sr-evalsteps.png")

arti_mani/algorithms/rl_iam/sac/sac_eval_segpts_PNfeat.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
from collections import defaultdict

import cv2
import gym
import matplotlib.pyplot as plt
from arti_mani import RLMODEL_DIR, VISUALMODEL_DIR
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import CustomUnet
from arti_mani.envs import *
from arti_mani.utils.common import convert_np_bool_to_float
from arti_mani.utils.cv_utils import visualize_depth, visualize_seg
from arti_mani.utils.visualization import images_to_video
from arti_mani.utils.wrappers import NormalizeActionWrapper, RenderInfoWrapper
from stable_baselines3.sac import SAC
from tqdm import tqdm

CMAPS = [
    (1, 0, 0),  # "red"
    (0, 0, 1),  # "blue"
    (1, 1, 0),  # "yellow"
    (0, 1, 0),  # "green"
    (0.627, 0.125, 0.941),  # "purple"
    (0.753, 0.753, 0.753),  # "grey"
]


def load_vismodel(vismodel_path, load_device):
    model_path = VISUALMODEL_DIR / f"smp_model/{vismodel_path}/best.pth"
    config_path = VISUALMODEL_DIR / f"smp_model/{vismodel_path}/config.yaml"
    with open(config_path, "r") as f:
        cfg = yaml.safe_load(f)
    smp_cfg = cfg["smp_config"]
    if smp_cfg["mode"] == "RGBD":
        in_ch = 4
    elif smp_cfg["mode"] == "RGB":
        in_ch = 3
    elif smp_cfg["mode"] == "D":
        in_ch = 1
    else:
        raise NotImplementedError
    segmodel = CustomUnet(
        has_dropout=False,
        encoder_name=smp_cfg["encoder"],
        encoder_depth=smp_cfg["encoder_depth"],
        decoder_channels=smp_cfg["decoder_channels"],
        encoder_weights=smp_cfg["encoder_weights"],
        in_channels=in_ch,
        classes=cfg["num_classes"],
        activation=smp_cfg["activation"],
    )
    segmodel.load_state_dict(torch.load(model_path))
    segmodel.to(torch.device(load_device))
    segmodel.eval()
    return segmodel


def vis_sample_pts(
    rgbdseg, indices, num_classes, pts_sample_num, vis_model, device, save_filename
):
    assert indices.shape[0] == num_classes * pts_sample_num
    rgb, depth, seg = (
        rgbdseg["rgb"].transpose((1, 2, 0)) / 255.0,
        rgbdseg["depth"],
        rgbdseg["seg"],
    )
    h, w = 144, 256
    v = indices // w  # (C * sample_num)
    u = indices - v * w  # (C * sample_num)
    rgb_plt = rgb.copy()
    for idn in range(num_classes):
        for pt_num in range(pts_sample_num):
            cv2.circle(
                rgb_plt,
                (
                    round(u[idn * pts_sample_num + pt_num]),
                    round(v[idn * pts_sample_num + pt_num]),
                ),
                radius=2,
                color=CMAPS[idn],
                thickness=-1,
            )

    with torch.no_grad():
        rgb_tensor = torch.from_numpy(rgb.transpose(2, 0, 1)[None]).to(
            device, dtype=torch.float32
        )
        depth_tensor = torch.from_numpy(depth[None, None]).to(
            device, dtype=torch.float32
        )
        seg_pred = vis_model.predict(torch.cat((rgb_tensor, depth_tensor), dim=1))[0]
        seg_mc_label = torch.argmax(seg_pred, dim=0).cpu().numpy()

    plt.subplot(231)
    plt.imshow(rgb)
    plt.title("rgb")
    plt.axis("off")
    plt.subplot(232)
    plt.imshow(visualize_depth(depth))
    plt.title("depth")
    plt.axis("off")
    plt.subplot(233)
    plt.imshow(visualize_seg(seg))
    plt.title("seg")
    plt.axis("off")
    plt.subplot(234)
    plt.imshow(visualize_seg(seg_mc_label, 6))
    plt.title("segpred")
    plt.axis("off")
    plt.subplot(235)
    plt.imshow(rgb_plt)
    plt.title("sample_pts")
    plt.axis("off")
    plt.savefig(save_filename, bbox_inches="tight", pad_inches=0.0)


if __name__ == "__main__":
    device = "cuda:0"
    mode = "door"  # "door", "drawer", "faucet", "arti"
    test_trials = 3
    test_num = 50
    save_num = 0
    MAX_STEPS = 100
    arti_id = 0

    smp_exp_name = "20230219_000940_train52-val18_384_noDR_randombg_aug_stereo_bs16_focalloss_0.5step50lr0.001_RGBDunet-163264128_mobilenet_v2"
    vis_model = load_vismodel(smp_exp_name, device)
    segmodel_path = VISUALMODEL_DIR / f"smp_model/{smp_exp_name}"
    other_handle_visible = False
    obs_mode = "state_egostereo_segpoints"
    control_mode = "pd_joint_delta_pos"
    sample_mode = "random_sample"  # full_downsample, score_sample, random_sample, frameweight_sample, fps_sample
    frame_num = 1
    num_classes = 6
    state_repeat_times = 10
    pts_sample_num = 8

    rl_exp_name = "ablations/sac4_tr4eval1door_state_egostereo_segpoints_random_samplef1s8_smp0219_2"
    model_path = (
        RLMODEL_DIR / f"{rl_exp_name}/best_model"
    )  # rl_model_2000000_steps, best_model

    env = gym.make(
        "ArtiMani-v0",
        articulation_ids=[arti_id],
        segmodel_path=segmodel_path,
        sample_mode=sample_mode,
        frame_num=frame_num,
        sample_num=pts_sample_num,
        other_handle_visible=other_handle_visible,
        obs_mode=obs_mode,
        control_mode=control_mode,
        device=device,
        reward_mode="dense",  # "sparse", "dense"
    )
    env = NormalizeActionWrapper(env)
    env = RenderInfoWrapper(env)
    RL_model = SAC.load(
        model_path,
        env=env,
        print_system_info=True,
    )

    success_rates = defaultdict(list)
    success_steps = defaultdict(list)
    for test_id in range(test_trials):
        eval_videos_path = (
            RLMODEL_DIR / f"{rl_exp_name}/eval_videos/0_best_model_{test_id}/"
        )
        if not os.path.exists(eval_videos_path):
            os.makedirs(eval_videos_path)

        eval_seed = np.random.RandomState().randint(2**32)
        print("experiment eval random seed: ", eval_seed)
        env.seed(eval_seed)

        is_success = []
        exp_steps = []
        with torch.no_grad():
            # for num in range(test_num):
            for num in tqdm(
                range(test_num),
                desc=f"Processing best RL model",
                colour="green",
                leave=True,
            ):
                obs = env.reset(articulation_id=arti_id)
                if num < save_num:
                    frames = []
                    base_rgbd = env.render("cameras")
                    frames.append(base_rgbd)
                success_flag = False
                total_steps = 0
                for step in tqdm(range(MAX_STEPS), colour="red", leave=False):
                    action, _states = RL_model.predict(obs, deterministic=True)
                    obs, rewards, dones, info = env.step(action)
                    # print(f"{step} step: {info['is_success']}")
                    if num < save_num:
                        base_rgbd = env.render("cameras")
                        frames.append(base_rgbd)
                    if num == 0:
                        rgbdseg = env.get_handsensor_rgbdseg()
                        indices = env.sampler.sample_indices  # C*sample_num
                        vis_filename = os.path.join(
                            eval_videos_path, f"{num}_{step}.png"
                        )
                        vis_sample_pts(
                            rgbdseg,
                            indices,
                            num_classes,
                            pts_sample_num,
                            vis_model,
                            device,
                            vis_filename,
                        )
                    if not success_flag and info["is_success"] == 1.0:
                        success_flag = True
                        total_steps = step
                        break
                if num < save_num:
                    images_to_video(
                        images=frames,
                        output_dir=eval_videos_path,
                        video_name=f"{num}",
                        verbose=False,
                    )
                is_success.append(convert_np_bool_to_float(success_flag))
                exp_steps.append(total_steps)
        success_rate = np.sum(is_success) / test_num
        exp_steps = np.array(exp_steps)
        success_mask = exp_steps > 0
        if np.any(success_mask):
            mean_steps = np.mean(exp_steps[success_mask])
        else:
            mean_steps = MAX_STEPS
        print(f"{arti_id}, success_rate, mean_steps: {success_rate}, {mean_steps}")
        success_rates[arti_id].append(success_rate)
        success_steps[arti_id].append(mean_steps)
    print(f"{arti_id}, {test_trials} trials, success_rates: {success_rates[arti_id]}")
    print(f"{arti_id}, {test_trials} trials, success_steps: {success_steps[arti_id]}")
    sr_mean, sr_std = np.mean(success_rates[arti_id]), np.std(success_rates[arti_id])
    ss_mean, ss_std = np.mean(success_steps[arti_id]), np.std(success_steps[arti_id])
    print(f"{test_trials} trials, success_rates mean, std: {sr_mean}, {sr_std}")
    print(f"{test_trials} trials, success_steps mean, std: {ss_mean}, {ss_std}")

arti_mani/algorithms/rl_iam/sac/sac_eval_segpts_PNfeat_ablations.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import glob
import logging
import os
import time
from collections import defaultdict

import cv2
import gym
import matplotlib.pyplot as plt
import segmentation_models_pytorch as smp
from arti_mani import RLMODEL_DIR, VISUALMODEL_DIR
from arti_mani.envs import *
from arti_mani.utils.cv_utils import visualize_depth
from arti_mani.utils.wrappers import NormalizeActionWrapper
from PIL import Image
from stable_baselines3.sac import SAC


def make_gif(frame_path, result_prefix):
    file_list = sorted(
        glob.glob(f"{frame_path}/*.jpg"), key=lambda name: int(name[-7:-4])
    )
    frames = [Image.open(image) for image in file_list]
    frame_one = frames[0]
    frame_one.save(
        f"{frame_path}/{result_prefix}.gif",
        format="GIF",
        append_images=frames,
        save_all=True,
        duration=500,
        loop=0,
    )


if __name__ == "__main__":
    device = torch.device("cuda:0")
    # visual_models = ["51200_epo", "102400_epo", "716800_epo", "1024000_epo"]
    eval_steps = np.array([5000000])
    color_maps = [
        (0, 255, 0),
        (0, 0, 255),
        (255, 255, 0),
        (0, 255, 255),
        (255, 0, 0),
        (255, 255, 255),
    ]
    subplot_num = [3, 3]
    num_classes = 6
    obs_mode = "state_egostereo_rgbd"
    control_mode = "pd_joint_delta_pos"
    rl_exp_name = "sac4_4doors_rgbdunet-mobilev2_6segmap180_pn256ln_state12pad10"

    smp_exps = [
        "20230111_005201_cab39-17_faucet11-3_stereo_bs256_0.5step50lr0.001_RGBDunet3-163264_mobilenet_v2",
    ]

    logfile_path = VISUALMODEL_DIR / f"smp_model"
    logfile_name = f"{time.strftime('%Y%m%d_%H%M%S', time.localtime())}_eval_smp_result"
    logger = logging.getLogger(logfile_name)
    logger.setLevel(logging.DEBUG)
    # set two handlers
    log_file = "{}.log".format(logfile_name)
    # rm_file(log_file)
    fileHandler = logging.FileHandler(os.path.join(logfile_path, log_file), mode="w")
    fileHandler.setLevel(logging.DEBUG)
    # set formatter
    formatter = logging.Formatter(
        "[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    fileHandler.setFormatter(formatter)
    # add
    logger.addHandler(fileHandler)

    for smp_exp_name in smp_exps:
        smp_model_path = VISUALMODEL_DIR / f"smp_model/{smp_exp_name}/best.pth"
        config_path = VISUALMODEL_DIR / f"smp_model/{smp_exp_name}/config.yaml"
        with open(config_path, "r") as f:
            cfg = yaml.safe_load(f)
        smp_cfg = cfg["smp_config"]
        if smp_cfg["mode"] == "RGBD":
            in_ch = 4
        elif smp_cfg["mode"] == "RGB":
            in_ch = 3
        elif smp_cfg["mode"] == "D":
            in_ch = 1
        else:
            raise NotImplementedError
        unet_model = smp.Unet(
            encoder_name=smp_cfg["encoder"],
            encoder_depth=smp_cfg["encoder_depth"],
            decoder_channels=smp_cfg["decoder_channels"],
            encoder_weights=smp_cfg["encoder_weights"],
            in_channels=in_ch,
            classes=cfg["num_classes"],
            activation=smp_cfg["activation"],
        )
        unet_model.load_state_dict(torch.load(smp_model_path))
        unet_model.to(device)
        unet_model.eval()
        model_params = sum(p.numel() for p in unet_model.parameters())

        # door
        env_id = "OpenCabinetDoor-v0"
        arti_ids = [0, 1006, 1030, 1047, 1081]
        arti_config_path = (
            ASSET_DIR / "partnet_mobility_configs/fixed_cabinet_doors_new.yml"
        )

        # create env
        env = gym.make(
            env_id,
            articulation_ids=arti_ids,
            articulation_config_path=arti_config_path,
            obs_mode=obs_mode,
            control_mode=control_mode,
            reward_mode="dense",
        )
        env = NormalizeActionWrapper(env)

        for eval_step in eval_steps:
            model_path = RLMODEL_DIR / f"{rl_exp_name}/rl_model_{eval_step}_steps.zip"

            # load models
            model = SAC.load(
                model_path,
                env=env,
                print_system_info=True,
            )
            sample_pts = model.actor.features_extractor.sample_num

            for arti_id in arti_ids:
                eval_result_path = (
                    VISUALMODEL_DIR
                    / f"smp_model/{smp_exp_name}/vis_results/{arti_id}_{eval_step / 1000000}M_new"
                )
                # print("visual results dir: ", eval_result_path)
                logger.info(
                    f"{smp_exp_name}: {arti_id}, {rl_exp_name}-{eval_step / 1000000}M"
                )
                if not os.path.exists(eval_result_path):
                    os.makedirs(eval_result_path)

                logger.info(f"model params: {model_params/1e6:.2f} M")

                eval_seed = np.random.RandomState().randint(2**32)
                # print("experiment eval random seed: ", eval_seed)
                logger.info(f"experiment eval random seed: {eval_seed}")
                env.seed(eval_seed)

                with torch.no_grad():
                    obs = env.reset(articulation_id=arti_id)
                    infer_times = []
                    eval_items = defaultdict(list)
                    for ind in range(30):
                        rgb = obs["rgb"]  # (3,72,128)
                        rgb_norm = torch.from_numpy((rgb / 255.0)[None]).to(
                            device
                        )  # Normalization
                        ch, h, w = rgb.shape

                        depth = obs["depth"]  # (1,72,128)
                        # depth_norm = torch.from_numpy((depth / np.max(depth))[None]).to(device=device)
                        depth_norm = torch.from_numpy(depth[None]).to(device=device)
                        seg = obs["seg"].astype(np.uint8)  # (72, 128)
                        seg_torch = torch.from_numpy(seg[None]).long().to(device)
                        seg_gt_onehot = F.one_hot(
                            seg_torch, num_classes=cfg["num_classes"]
                        ).permute(
                            0, 3, 1, 2
                        )  # (1,C,72,128)
                        # other_seg = obs["other_handleseg"]
                        # print(rgb_norm.shape, depth_norm.shape, seg.shape)
                        if smp_cfg["mode"] == "RGBD":
                            img = torch.cat([rgb_norm, depth_norm], dim=1)
                        elif smp_cfg["mode"] == "RGB":
                            img = rgb_norm
                        elif smp_cfg["mode"] == "D":
                            img = depth_norm
                        else:
                            raise NotImplementedError
                        time0 = time.time()
                        # seg_feat = unet_model.predict(img)  # (1, C, H, W)
                        # attnmap = F.softmax(seg_feat, dim=1)  # (1, C, H, W)
                        attnmap = unet_model.predict(img)  # (1, C, H, W)
                        time1 = time.time()
                        infer_times.append(time1 - time0)

                        tp, fp, fn, tn = smp.metrics.get_stats(
                            attnmap, seg_gt_onehot, mode="multilabel", threshold=0.5
                        )

                        iou_metric = smp.metrics.iou_score(
                            tp, fp, fn, tn, reduction="none"
                        )  # (bs, C)
                        f1_metric = smp.metrics.f1_score(
                            tp, fp, fn, tn, reduction="none"
                        )  # (bs, C)
                        for cls_ind, key in enumerate(smp_cfg["classes"]):
                            eval_items[f"iou_{key}"].append(
                                torch.mean(iou_metric, dim=0)[cls_ind].item()
                            )
                            eval_items[f"f1_{key}"].append(
                                torch.mean(f1_metric, dim=0)[cls_ind].item()
                            )
                        # then compute metrics with required reduction (see metric docs)
                        eval_items["iou"].append(
                            smp.metrics.iou_score(
                                tp, fp, fn, tn, reduction="micro"
                            ).item()
                        )
                        eval_items["f1"].append(
                            smp.metrics.f1_score(
                                tp, fp, fn, tn, reduction="micro"
                            ).item()
                        )
                        eval_items["acc"].append(
                            smp.metrics.accuracy(
                                tp, fp, fn, tn, reduction="macro"
                            ).item()
                        )
                        eval_items["prec"].append(
                            smp.metrics.precision(
                                tp, fp, fn, tn, reduction="micro-imagewise"
                            ).item()
                        )
                        eval_items["recall"].append(
                            smp.metrics.recall(
                                tp, fp, fn, tn, reduction="micro-imagewise"
                            ).item()
                        )

                        plot_rgb = (rgb.transpose((1, 2, 0))).astype(np.uint8)
                        attnmap = attnmap.squeeze().cpu().numpy()

                        plt.subplot(subplot_num[0], subplot_num[1], 1)
                        plt.imshow(plot_rgb / 255.0)
                        plt.subplot(subplot_num[0], subplot_num[1], 2)
                        plt.imshow(depth.squeeze())
                        plt.subplot(subplot_num[0], subplot_num[1], 3)
                        plt.imshow(seg)
                        for pid in range(num_classes):
                            plt.subplot(subplot_num[0], subplot_num[1], pid + 4)
                            plt.imshow(attnmap[pid])

                        plt.savefig(f"{eval_result_path}/{ind:03}.jpg", dpi=200)
                        plt.close()

                        depth_colormap = visualize_depth(depth.squeeze())
                        cv2.imwrite(
                            f"{eval_result_path}/{ind:03}_depth_colormap.png",
                            depth_colormap,
                        )

                        action, _states = model.predict(obs, deterministic=True)
                        obs, rewards, dones, info = env.step(action)
                        # print(f"{ind} step: {info['is_success']}, {info['check_grasp']}")
                        if info["is_success"] == 1.0:
                            break
                # print(f"inference time: {np.mean(infer_times)}")
                logger.info(f"inference time: {np.mean(infer_times)*1000:.2f} ms")
                for key in eval_items.keys():
                    logger.info(f"metrics {key}: {np.mean(eval_items[key]):.4f}")
                make_gif(eval_result_path, f"{eval_step / 1000000}M")

arti_mani/algorithms/rl_iam/sac/sac_eval_smpmodel.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import glob
import os
from collections import defaultdict

import gym
import matplotlib.pyplot as plt
import torch
from arti_mani import RLMODEL_DIR
from arti_mani.envs import *
from arti_mani.utils.common import convert_np_bool_to_float
from arti_mani.utils.visualization import images_to_video
from arti_mani.utils.wrappers import NormalizeActionWrapper, RenderInfoWrapper
from stable_baselines3.sac import SAC

if __name__ == "__main__":
    device = "cuda:0"
    mode = "door"  # "door", "drawer", "faucet"
    test_num = 10
    save_num = 10
    success_nums = 0
    MAX_STEPS = 50
    rl_exp_name = "sac4_4faucet_state_30tip40finger_gripperclose_leftrew"
    eval_steps = []
    for file in glob.glob(str(RLMODEL_DIR / f"{rl_exp_name}/rl_model_*_steps.zip")):
        eval_steps.append(int(os.path.basename(file).split("_")[2]))
    eval_steps.sort()
    eval_steps = np.array(eval_steps)

    env_id = "ArtiMani-v0"
    other_handle_visible = False
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069]
    if mode == "door":
        arti_ids = [0, 1006, 1030, 1047, 1081]
    elif mode == "drawer":
        arti_ids = [1, 1004, 1005, 1016, 1024]
    elif mode == "faucet":
        arti_ids = [5052, 5004, 5007, 5023, 5069]
    else:
        raise NotImplementedError(mode)
    obs_mode = "state_dict"
    control_mode = "pd_joint_delta_pos"
    sample_mode = "random_sample"  # full_downsample, score_sample, random_sample, frameweight_sample
    frame_num = 1
    num_classes = 6
    state_repeat_times = 10
    pts_sample_num = 30

    env = gym.make(
        env_id,
        articulation_ids=arti_ids,
        segmodel_path=None,
        load_device="cuda:0",
        sample_mode=sample_mode,
        frame_num=frame_num,
        sample_num=pts_sample_num,
        other_handle_visible=False,
        obs_mode=obs_mode,
        control_mode=control_mode,
        reward_mode="dense",  # "sparse", "dense"
    )
    env = NormalizeActionWrapper(env)
    env = RenderInfoWrapper(env)
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    env.seed(eval_seed)

    success_rates = defaultdict(list)
    for eval_step in eval_steps:
        model_path = RLMODEL_DIR / f"{rl_exp_name}/rl_model_{eval_step}_steps"
        model = SAC.load(
            model_path,
            env=env,
            print_system_info=True,
        )
        for arti_id in arti_ids:
            eval_videos_path = (
                RLMODEL_DIR
                / f"{rl_exp_name}/eval_videos/{arti_id}/{eval_step / 1000000}M/"
            )
            if not os.path.exists(eval_videos_path):
                os.makedirs(eval_videos_path)

            is_success = []
            with torch.no_grad():
                for num in range(test_num):
                    obs = env.reset(articulation_id=arti_id)
                    if num < save_num:
                        frames = []
                        base_rgbd = env.render("cameras")
                        frames.append(base_rgbd)
                    success_flag = False
                    for i in range(MAX_STEPS):
                        action, _states = model.predict(obs, deterministic=True)
                        obs, rewards, dones, info = env.step(action)
                        # print(f"{i} step: {info['is_success']}")
                        if not success_flag and info["is_success"] == 1.0:
                            success_nums += 1
                            success_flag = True
                        if num < save_num:
                            base_rgbd = env.render("cameras")
                            frames.append(base_rgbd)
                    if num < save_num:
                        images_to_video(
                            images=frames,
                            output_dir=eval_videos_path,
                            video_name=f"{num}",
                        )
                    is_success.append(convert_np_bool_to_float(success_flag))
            sum_num = np.sum(is_success)
            print(f"success rate: {sum_num}/{test_num}")
            success_rates[arti_id].append(sum_num / test_num)

    plt.figure()
    for arti_id in arti_ids:
        plt.plot(eval_steps / 1000000, success_rates[arti_id], label=arti_id)
    plt.legend(loc="lower right")
    plt.xlabel("eval_model_steps(M)")
    plt.ylabel("success_rate")
    plt.title(f"Eval on {len(arti_ids)} {mode}(s)")
    plt.savefig(RLMODEL_DIR / f"{rl_exp_name}/eval_videos/sr-evalsteps.png")

arti_mani/algorithms/rl_iam/sac/sac_eval_state.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import random

from arti_mani.algorithms.rl_iam.feature_extract import CustomSegPNExtractor
from arti_mani.algorithms.rl_iam.rl_utils import sb3_make_env_multiarti
from arti_mani.envs import *
from stable_baselines3 import SAC
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
from stable_baselines3.common.logger import configure
from stable_baselines3.common.vec_env import SubprocVecEnv, VecMonitor


def setup_seed(seed=1029):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "door"  # "door", "drawer", "faucet", "arti"
    train_procs = 4
    eval_procs = 1
    save_eval_freq = 250000 // train_procs
    seed = np.random.RandomState().randint(2**32)
    print("experiment random seed: ", seed)
    setup_seed(seed)
    env_id = "ArtiMani-v0"
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069], [5004, 5007, 5023, 5069]
    if mode == "door":
        arti_ids = [1006, 1030, 1047, 1081]
        eval_ids = [0]
    elif mode == "drawer":
        arti_ids = [1005, 1016, 1024, 1076]
        eval_ids = [1]
    elif mode == "faucet":
        arti_ids = [5004, 5007, 5023, 5069]
        eval_ids = [5052]
    elif mode == "arti":
        arti_ids = [
            1006,
            1030,
            1047,
            1081,
            1005,
            1016,
            1024,
            1076,
            5004,
            5007,
            5023,
            5069,
        ]
        eval_ids = [0, 1, 5052]
    else:
        raise NotImplementedError(mode)
    other_handle_visible = False
    obs_mode = "state_egostereo_dexpoints"
    control_mode = "pd_joint_delta_pos"
    num_classes = 2
    state_repeat_times = 10

    # Create log dir
    exp_suffix = "0"
    log_dir = f"./logs/sac{train_procs}_tr{len(arti_ids)}eval{len(eval_ids)}{mode}_{obs_mode}_{exp_suffix}/"
    os.makedirs(log_dir, exist_ok=True)

    vec_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=arti_ids,
                segmodel_path=None,
                sample_mode=None,
                frame_num=0,
                sample_num=0,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(train_procs)
        ],
        start_method="spawn",
    )
    vec_env = VecMonitor(vec_env, log_dir)

    print("Observation Space: ", vec_env.observation_space)
    print("Action Space: ", vec_env.action_space)

    # setup callbacks
    checkpoint_callback = CheckpointCallback(
        save_freq=save_eval_freq, save_path=log_dir
    )
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    vec_eval_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=eval_ids,
                segmodel_path=None,
                sample_mode=None,
                frame_num=0,
                sample_num=0,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(eval_procs)
        ],
        start_method="spawn",
    )
    vec_eval_env = VecMonitor(vec_eval_env, log_dir)
    eval_callback = EvalCallback(
        vec_eval_env,
        n_eval_episodes=50,
        best_model_save_path=log_dir,
        log_path=log_dir,
        eval_freq=save_eval_freq,
        deterministic=True,
        render=False,
    )

    # set up logger
    new_logger = configure(log_dir, ["stdout", "csv", "tensorboard"])

    model = SAC(
        "MultiInputPolicy",
        vec_env,
        batch_size=256,  # 1024, 400
        ent_coef="auto_0.2",
        gamma=0.99,
        train_freq=16,  # 4, 64
        gradient_steps=16,  # 2, 4
        buffer_size=100000,
        learning_starts=800,
        use_sde=True,
        policy_kwargs=dict(
            log_std_init=-3.67,
            net_arch=[256, 256],
            features_extractor_class=CustomSegPNExtractor,
            features_extractor_kwargs=dict(
                num_classes=num_classes,
                state_repeat_times=state_repeat_times,
                sample_num=1000,
            ),
            normalize_images=False,
            share_features_extractor=True,
        ),
        tensorboard_log=log_dir + "sac_opendoor_tb/",
        device=device,
        verbose=1,
    )
    # Set new logger
    model.set_logger(new_logger)
    model.learn(
        total_timesteps=2_000_000,
        callback=[checkpoint_callback, eval_callback],
    )

arti_mani/algorithms/rl_iam/sac/sac_train_dexpoints_PNfeat.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import random

from arti_mani.algorithms.rl_iam.feature_extract import CustomDSegResNetExtractor
from arti_mani.algorithms.rl_iam.rl_utils import sb3_make_env_multiarti
from arti_mani.envs import *
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
from stable_baselines3.common.logger import configure
from stable_baselines3.common.vec_env import SubprocVecEnv, VecMonitor
from stable_baselines3.sac import SAC


def setup_seed(seed=1029):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "door"  # "door", "drawer", "faucet"
    train_procs = 4
    eval_procs = 1
    save_eval_freq = 250000 // train_procs
    seed = np.random.RandomState().randint(2**32)
    print("experiment random seed: ", seed)
    setup_seed(seed)
    env_id = "ArtiMani-v0"
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069], [5004, 5007, 5023, 5069]
    if mode == "door":
        arti_ids = [1006, 1030, 1047, 1081]
        eval_ids = [0]
    elif mode == "drawer":
        arti_ids = [1005, 1016, 1024, 1076]
        eval_ids = [1]
    elif mode == "faucet":
        arti_ids = [5004, 5007, 5023, 5069]
        eval_ids = [5052]
    elif mode == "arti":
        arti_ids = [
            1006,
            1030,
            1047,
            1081,
            1005,
            1016,
            1024,
            1076,
            5004,
            5007,
            5023,
            5069,
        ]
        eval_ids = [0, 1, 5052]
    else:
        raise NotImplementedError(mode)
    obs_mode = "state_egostereo_dseg_gt"
    control_mode = "pd_joint_delta_pos"
    other_handle_visible = False
    state_repeat_times = 10
    num_classes = 6

    # Create log dir
    exp_suffix = "100steps_padfriction1_p3000d100_tr16up16_2"
    log_dir = f"./logs/sac{train_procs}_tr{len(arti_ids)}eval{len(eval_ids)}{mode}_{obs_mode}_{exp_suffix}/"
    os.makedirs(log_dir, exist_ok=True)

    vec_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=arti_ids,
                segmodel_path=None,
                sample_mode=None,
                frame_num=0,
                sample_num=0,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(train_procs)
        ],
        start_method="spawn",
    )
    vec_env = VecMonitor(vec_env, log_dir)

    print("Observation Space: ", vec_env.observation_space)
    print("Action Space: ", vec_env.action_space)

    # setup callbacks
    checkpoint_callback = CheckpointCallback(
        save_freq=save_eval_freq, save_path=log_dir
    )
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    vec_eval_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=eval_ids,
                segmodel_path=None,
                sample_mode=None,
                frame_num=0,
                sample_num=0,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(eval_procs)
        ],
        start_method="spawn",
    )
    vec_eval_env = VecMonitor(vec_eval_env, log_dir)
    eval_callback = EvalCallback(
        vec_eval_env,
        n_eval_episodes=50,
        best_model_save_path=log_dir,
        log_path=log_dir,
        eval_freq=save_eval_freq,
        deterministic=True,
        render=False,
    )

    # set up logger
    new_logger = configure(log_dir, ["stdout", "csv", "tensorboard"])

    model = SAC(
        "MultiInputPolicy",
        vec_env,
        batch_size=256,  # 1024, 400
        ent_coef="auto_0.2",
        gamma=0.99,
        train_freq=16,  # 4, 64
        gradient_steps=16,  # 2, 4
        buffer_size=100000,
        learning_starts=800,
        use_sde=True,
        policy_kwargs=dict(
            log_std_init=-3.67,
            net_arch=[256, 256],
            features_extractor_class=CustomDSegResNetExtractor,
            features_extractor_kwargs=dict(
                state_repeat_times=state_repeat_times,
                num_classes=num_classes,
                device=device,
            ),
            normalize_images=False,
            share_features_extractor=True,
        ),
        tensorboard_log=log_dir + "sac_opendoor_tb/",
        device=device,
        verbose=1,
    )
    # Set new logger
    model.set_logger(new_logger)
    model.learn(
        total_timesteps=2_000_000,
        callback=[checkpoint_callback, eval_callback],
    )

arti_mani/algorithms/rl_iam/sac/sac_train_dsegGT_RN18feat.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import random

from arti_mani import VISUALMODEL_DIR
from arti_mani.algorithms.rl_iam.feature_extract import CustomRGBDPredSegResNetExtractor
from arti_mani.algorithms.rl_iam.rl_utils import sb3_make_env_multiarti
from arti_mani.envs import *
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
from stable_baselines3.common.logger import configure
from stable_baselines3.common.vec_env import SubprocVecEnv, VecMonitor
from stable_baselines3.sac import SAC


def setup_seed(seed=1029):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "arti"  # "door", "drawer", "faucet", "arti"
    train_procs = 4
    eval_procs = 1
    save_eval_freq = 250000 // train_procs
    seed = np.random.RandomState().randint(2**32)
    print("experiment random seed: ", seed)
    setup_seed(seed)
    env_id = "ArtiMani-v0"
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069], [5004, 5007, 5023, 5069]
    if mode == "door":
        arti_ids = [1006, 1030, 1047, 1081]
        eval_ids = [0]
    elif mode == "drawer":
        arti_ids = [1005, 1016, 1024, 1076]
        eval_ids = [1]
    elif mode == "faucet":
        arti_ids = [5004, 5007, 5023, 5069]
        eval_ids = [5052]
    elif mode == "arti":
        arti_ids = [
            1006,
            1030,
            1047,
            1081,
            1005,
            1016,
            1024,
            1076,
            5004,
            5007,
            5023,
            5069,
        ]
        eval_ids = [0, 1, 5052]
    else:
        raise NotImplementedError(mode)
    segmodel_path = (
        VISUALMODEL_DIR
        / "smp_model/20230310_210426_noDR_norandombg_aug_dropout0.2_bs16_focalloss_0.5step40lr0.001_RGBDunet-163264128_mobilenet_v2"
    )
    obs_mode = "state_egostereo_dseg"
    control_mode = "pd_joint_delta_pos"
    other_handle_visible = False
    state_repeat_times = 10

    # Create log dir
    exp_suffix = "100steps_padfriction1_p3000d100_bs128tr16up16"
    log_dir = f"./logs/sac{train_procs}_tr{len(arti_ids)}eval{len(eval_ids)}{mode}_{obs_mode}_{exp_suffix}/"
    os.makedirs(log_dir, exist_ok=True)

    vec_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=arti_ids,
                segmodel_path=None,
                sample_mode=None,
                frame_num=0,
                sample_num=0,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(train_procs)
        ],
        start_method="spawn",
    )
    vec_env = VecMonitor(vec_env, log_dir)

    print("Observation Space: ", vec_env.observation_space)
    print("Action Space: ", vec_env.action_space)

    # setup callbacks
    checkpoint_callback = CheckpointCallback(
        save_freq=save_eval_freq, save_path=log_dir
    )
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    vec_eval_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=eval_ids,
                segmodel_path=None,
                sample_mode=None,
                frame_num=0,
                sample_num=0,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(eval_procs)
        ],
        start_method="spawn",
    )
    vec_eval_env = VecMonitor(vec_eval_env, log_dir)
    eval_callback = EvalCallback(
        vec_eval_env,
        n_eval_episodes=50,
        best_model_save_path=log_dir,
        log_path=log_dir,
        eval_freq=save_eval_freq,
        deterministic=True,
        render=False,
    )

    # set up logger
    new_logger = configure(log_dir, ["stdout", "csv", "tensorboard"])

    model = SAC(
        "MultiInputPolicy",
        vec_env,
        batch_size=128,  # 1024, 400
        ent_coef="auto_0.2",
        gamma=0.99,
        train_freq=16,  # 4, 64
        gradient_steps=16,  # 2, 4
        buffer_size=100000,
        learning_starts=800,
        use_sde=True,
        policy_kwargs=dict(
            log_std_init=-3.67,
            net_arch=[256, 256],
            features_extractor_class=CustomRGBDPredSegResNetExtractor,
            features_extractor_kwargs=dict(
                state_repeat_times=state_repeat_times,
                vismodel_path=segmodel_path,
                device=device,
            ),
            normalize_images=False,
            share_features_extractor=True,
        ),
        tensorboard_log=log_dir + "sac_opendoor_tb/",
        device=device,
        verbose=1,
    )
    # Set new logger
    model.set_logger(new_logger)
    model.learn(
        total_timesteps=2_000_000,
        callback=[checkpoint_callback, eval_callback],
    )

arti_mani/algorithms/rl_iam/sac/sac_train_dseg_RN18feat.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import random

from arti_mani import VISUALMODEL_DIR
from arti_mani.algorithms.rl_iam.feature_extract import CustomKptExtractor
from arti_mani.algorithms.rl_iam.rl_utils import sb3_make_env_multiarti
from arti_mani.envs import *
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
from stable_baselines3.common.logger import configure
from stable_baselines3.common.vec_env import SubprocVecEnv, VecMonitor
from stable_baselines3.sac import SAC


def setup_seed(seed=1029):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "arti"  # "door", "drawer", "faucet", "arti"
    train_procs = 4
    eval_procs = 1
    save_eval_freq = 250000 // train_procs
    seed = np.random.RandomState().randint(2**32)
    print("experiment random seed: ", seed)
    setup_seed(seed)
    env_id = "ArtiMani-v0"
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069], [5004, 5007, 5023, 5069]
    if mode == "door":
        arti_ids = [1006, 1030, 1047, 1081]
        eval_ids = [0]
    elif mode == "drawer":
        arti_ids = [1005, 1016, 1024, 1076]
        eval_ids = [1]
    elif mode == "faucet":
        arti_ids = [5004, 5007, 5023, 5069]
        eval_ids = [5052]
    elif mode == "arti":
        arti_ids = [
            1006,
            1030,
            1047,
            1081,
            1005,
            1016,
            1024,
            1076,
            5004,
            5007,
            5023,
            5069,
        ]
        eval_ids = [0, 1, 5052]
    else:
        raise NotImplementedError(mode)
    segmodel_path = (
        VISUALMODEL_DIR
        / "kpt_model/20230307_182821_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2_dropout0.2_newdatafilter2drawer"
    )
    other_handle_visible = False
    num_kpts = 3
    obs_mode = "state_egostereo_keypoints"
    control_mode = "pd_joint_delta_pos"

    # Create log dir
    exp_suffix = "bs256_tr16up16_3"
    log_dir = f"./logs/sac{train_procs}_tr{len(arti_ids)}eval{len(eval_ids)}{mode}_{obs_mode}{num_kpts}_{exp_suffix}/"
    os.makedirs(log_dir, exist_ok=True)

    vec_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=arti_ids,
                segmodel_path=segmodel_path,
                sample_mode=None,
                frame_num=0,
                sample_num=0,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(train_procs)
        ],
        start_method="spawn",
    )
    vec_env = VecMonitor(vec_env, log_dir)

    print("Observation Space: ", vec_env.observation_space)
    print("Action Space: ", vec_env.action_space)

    # setup callbacks
    checkpoint_callback = CheckpointCallback(
        save_freq=save_eval_freq, save_path=log_dir
    )
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    vec_eval_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=eval_ids,
                segmodel_path=segmodel_path,
                sample_mode=None,
                frame_num=0,
                sample_num=0,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(eval_procs)
        ],
        start_method="spawn",
    )
    vec_eval_env = VecMonitor(vec_eval_env, log_dir)
    eval_callback = EvalCallback(
        vec_eval_env,
        n_eval_episodes=50,
        best_model_save_path=log_dir,
        log_path=log_dir,
        eval_freq=save_eval_freq,
        deterministic=True,
        render=False,
    )

    # set up logger
    new_logger = configure(log_dir, ["stdout", "csv", "tensorboard"])

    model = SAC(
        "MultiInputPolicy",
        vec_env,
        batch_size=256,  # 1024, 400
        ent_coef="auto_0.2",
        gamma=0.99,
        train_freq=16,  # 4, 64
        gradient_steps=16,  # 2, 4
        buffer_size=100000,
        learning_starts=800,
        use_sde=True,
        policy_kwargs=dict(
            log_std_init=-3.67,
            net_arch=[256, 256],
            features_extractor_class=CustomKptExtractor,
            features_extractor_kwargs=dict(num_kpts=num_kpts),
            normalize_images=False,
            share_features_extractor=True,
        ),
        tensorboard_log=log_dir + "sac_opendoor_tb/",
        device=device,
        verbose=1,
    )
    # Set new logger
    model.set_logger(new_logger)
    model.learn(
        total_timesteps=2_000_000,
        callback=[checkpoint_callback, eval_callback],
    )

arti_mani/algorithms/rl_iam/sac/sac_train_keypts.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import random

from arti_mani import VISUALMODEL_DIR
from arti_mani.algorithms.rl_iam.feature_extract import CustomKptGTExtractor
from arti_mani.algorithms.rl_iam.rl_utils import sb3_make_env_multiarti
from arti_mani.envs import *
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
from stable_baselines3.common.logger import configure
from stable_baselines3.common.vec_env import SubprocVecEnv, VecMonitor
from stable_baselines3.sac import SAC


def setup_seed(seed=1029):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "door"  # "door", "drawer", "faucet"
    train_procs = 4
    eval_procs = 1
    save_eval_freq = 250000 // train_procs
    seed = np.random.RandomState().randint(2**32)
    print("experiment random seed: ", seed)
    setup_seed(seed)
    env_id = "ArtiMani-v0"
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069], [5004, 5007, 5023, 5069]
    if mode == "door":
        arti_ids = [1006, 1030, 1047, 1081]
        eval_ids = [0]
    elif mode == "drawer":
        arti_ids = [1005, 1016, 1024, 1076]
        eval_ids = [1]
    elif mode == "faucet":
        arti_ids = [5004, 5007, 5023, 5069]
        eval_ids = [5052]
    elif mode == "arti":
        arti_ids = [
            1006,
            1030,
            1047,
            1081,
            1005,
            1016,
            1024,
            1076,
            5004,
            5007,
            5023,
            5069,
        ]
        eval_ids = [0, 1, 5052]
    else:
        raise NotImplementedError(mode)
    segmodel_path = (
        VISUALMODEL_DIR
        / "kpt_model/20230307_182821_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2_dropout0.2_newdatafilter2drawer"
    )
    other_handle_visible = False
    num_kpts = 3
    obs_mode = "state_egostereo_keypoints_gt"
    control_mode = "pd_joint_delta_pos"

    # Create log dir
    exp_suffix = "bs128_tr16up16"
    log_dir = f"./logs/sac{train_procs}_tr{len(arti_ids)}eval{len(eval_ids)}{mode}_{obs_mode}{num_kpts}_{exp_suffix}/"
    os.makedirs(log_dir, exist_ok=True)

    vec_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=arti_ids,
                segmodel_path=segmodel_path,
                sample_mode=None,
                frame_num=0,
                sample_num=0,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(train_procs)
        ],
        start_method="spawn",
    )
    vec_env = VecMonitor(vec_env, log_dir)

    print("Observation Space: ", vec_env.observation_space)
    print("Action Space: ", vec_env.action_space)

    # setup callbacks
    checkpoint_callback = CheckpointCallback(
        save_freq=save_eval_freq, save_path=log_dir
    )
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    vec_eval_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=eval_ids,
                segmodel_path=segmodel_path,
                sample_mode=None,
                frame_num=0,
                sample_num=0,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(eval_procs)
        ],
        start_method="spawn",
    )
    vec_eval_env = VecMonitor(vec_eval_env, log_dir)
    eval_callback = EvalCallback(
        vec_eval_env,
        n_eval_episodes=50,
        best_model_save_path=log_dir,
        log_path=log_dir,
        eval_freq=save_eval_freq,
        deterministic=True,
        render=False,
    )

    # set up logger
    new_logger = configure(log_dir, ["stdout", "csv", "tensorboard"])

    # check_env(env)
    model = SAC(
        "MultiInputPolicy",
        vec_env,
        batch_size=128,  # 1024, 400
        ent_coef="auto_0.2",
        gamma=0.99,
        train_freq=16,  # 4, 64
        gradient_steps=16,  # 2, 4
        buffer_size=100000,
        learning_starts=800,
        use_sde=True,
        policy_kwargs=dict(
            log_std_init=-3.67,
            net_arch=[256, 256],
            features_extractor_class=CustomKptGTExtractor,
            features_extractor_kwargs=dict(num_kpts=num_kpts),
            normalize_images=False,
            share_features_extractor=True,
        ),
        tensorboard_log=log_dir + "sac_opendoor_tb/",
        device=device,
        verbose=1,
    )
    # Set new logger
    model.set_logger(new_logger)
    model.learn(
        total_timesteps=2_000_000,
        callback=[checkpoint_callback, eval_callback],
    )

arti_mani/algorithms/rl_iam/sac/sac_train_keyptsGT.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import random

from arti_mani.algorithms.rl_iam.feature_extract import CustomSegPNExtractor
from arti_mani.algorithms.rl_iam.rl_utils import sb3_make_env_multiarti
from arti_mani.envs import *
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
from stable_baselines3.common.logger import configure
from stable_baselines3.common.vec_env import SubprocVecEnv, VecMonitor
from stable_baselines3.sac import SAC


def setup_seed(seed=1029):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "arti"  # "door", "drawer", "faucet", "arti"
    train_procs = 4
    eval_procs = 1
    save_eval_freq = 250000 // train_procs
    seed = np.random.RandomState().randint(2**32)
    print("experiment random seed: ", seed)
    setup_seed(seed)
    env_id = "ArtiMani-v0"
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069] [5004, 5007, 5023, 5069]
    if mode == "door":
        arti_ids = [1006, 1030, 1047, 1081]
        eval_ids = [0]
    elif mode == "drawer":
        arti_ids = [1005, 1016, 1024, 1076]
        eval_ids = [1]
    elif mode == "faucet":
        arti_ids = [5004, 5007, 5023, 5069]
        eval_ids = [5052]
    elif mode == "arti":
        arti_ids = [
            1006,
            1030,
            1047,
            1081,
            1005,
            1016,
            1024,
            1076,
            5004,
            5007,
            5023,
            5069,
        ]
        eval_ids = [0, 1, 5052]
    else:
        raise NotImplementedError(mode)
    other_handle_visible = False
    obs_mode = "state_egostereo_segpoints_gt"
    control_mode = "pd_joint_delta_pos"
    num_classes = 6
    state_repeat_times = 10
    sample_mode = "random_sample"  # full_downsample, score_sample, random_sample, frameweight_sample
    frame_num = 1
    pts_sample_num = 32

    # Create log dir
    exp_suffix = "gt_4"  # "adjustreward_2"
    log_dir = f"./logs/sac{train_procs}_tr{len(arti_ids)}eval{len(eval_ids)}{mode}_{obs_mode}_{sample_mode}f{frame_num}s{pts_sample_num}_{exp_suffix}/"
    os.makedirs(log_dir, exist_ok=True)

    vec_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=arti_ids,
                segmodel_path=None,
                sample_mode=sample_mode,
                frame_num=frame_num,
                sample_num=pts_sample_num,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(train_procs)
        ],
        start_method="spawn",
    )
    vec_env = VecMonitor(vec_env, log_dir)

    print("Observation Space: ", vec_env.observation_space)
    print("Action Space: ", vec_env.action_space)

    # setup callbacks
    checkpoint_callback = CheckpointCallback(
        save_freq=save_eval_freq, save_path=log_dir
    )
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    vec_eval_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=eval_ids,
                segmodel_path=None,
                sample_mode=sample_mode,
                frame_num=frame_num,
                sample_num=pts_sample_num,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(eval_procs)
        ],
        start_method="spawn",
    )
    vec_eval_env = VecMonitor(vec_eval_env, log_dir)
    eval_callback = EvalCallback(
        vec_eval_env,
        n_eval_episodes=50,
        best_model_save_path=log_dir,
        log_path=log_dir,
        eval_freq=save_eval_freq,
        deterministic=True,
        render=False,
    )

    # set up logger
    new_logger = configure(log_dir, ["stdout", "csv", "tensorboard"])

    # check_env(env)
    model = SAC(
        "MultiInputPolicy",
        vec_env,
        batch_size=256,  # 1024, 400
        ent_coef="auto_0.2",
        gamma=0.99,
        train_freq=16,  # 4, 64
        gradient_steps=16,  # 2, 4
        buffer_size=100000,
        learning_starts=800,
        use_sde=True,
        policy_kwargs=dict(
            log_std_init=-3.67,
            net_arch=[256, 256],
            features_extractor_class=CustomSegPNExtractor,
            features_extractor_kwargs=dict(
                num_classes=num_classes,
                state_repeat_times=state_repeat_times,
                sample_num=pts_sample_num,
            ),
            normalize_images=False,
            share_features_extractor=True,
        ),
        tensorboard_log=log_dir + "sac_opendoor_tb/",
        device=device,
        verbose=1,
    )
    # Set new logger
    model.set_logger(new_logger)
    model.learn(
        total_timesteps=2_000_000,
        callback=[checkpoint_callback, eval_callback],
    )

arti_mani/algorithms/rl_iam/sac/sac_train_segGTpts_PNfeat.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import random

from arti_mani.algorithms.rl_iam.feature_extract import CustomSegPNExtractor
from arti_mani.algorithms.rl_iam.rl_utils import sb3_make_env_multiarti
from arti_mani.envs import *
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
from stable_baselines3.common.logger import configure
from stable_baselines3.common.vec_env import SubprocVecEnv, VecMonitor
from stable_baselines3.sac import SAC


def setup_seed(seed=1029):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "arti"  # "door", "drawer", "faucet", "laptop", "kitchen_pot", "arti"
    train_procs = 4
    eval_procs = 1
    save_eval_freq = 200000 // train_procs
    seed = np.random.RandomState().randint(2**32)
    print("experiment random seed: ", seed)
    setup_seed(seed)
    env_id = "ArtiMani-v0"
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069] [5004, 5007, 5023, 5069]
    if mode == "door":
        arti_ids = [1006, 1030, 1047, 1081]
        eval_ids = [0]
    elif mode == "drawer":
        arti_ids = [1005, 1016, 1024, 1076]
        eval_ids = [1]
    elif mode == "faucet":
        arti_ids = [5004, 5007, 5023, 5069]
        eval_ids = [5052]
    elif mode == "laptop":
        arti_ids = [9960, 9968, 9992, 9996]
        eval_ids = [9748]
    elif mode == "kitchen_pot":
        arti_ids = [100051, 100054, 100055, 100060]
        eval_ids = [100015]
    elif mode == "arti":
        arti_ids = [
            1006,
            1030,
            1047,
            1081,
            1005,
            1016,
            1024,
            1076,
            5004,
            5007,
            5023,
            5069,
            9960,
            9968,
            9992,
            9996,
            100051,
            100054,
            100055,
            100060,
        ]
        eval_ids = [0, 1, 5052, 9748, 100015]
    else:
        raise NotImplementedError(mode)
    other_handle_visible = False
    obs_mode = "state_ego_segpoints_gt"
    control_mode = "pd_joint_delta_pos"
    num_classes = 5
    state_repeat_times = 10
    sample_mode = "random_sample"  # full_downsample, score_sample, random_sample, frameweight_sample
    frame_num = 1
    pts_sample_num = 32

    # Create log dir
    exp_suffix = "multiarti_5class_test"  # "adjustreward_2"
    log_dir = f"./logs/sac{train_procs}_tr{len(arti_ids)}eval{len(eval_ids)}{mode}_{obs_mode}_{sample_mode}f{frame_num}s{pts_sample_num}_{exp_suffix}/"
    os.makedirs(log_dir, exist_ok=True)

    vec_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=arti_ids,
                segmodel_path=None,
                sample_mode=sample_mode,
                frame_num=frame_num,
                sample_num=pts_sample_num,
                other_handle_visible=other_handle_visible,
                num_classes=num_classes,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(train_procs)
        ],
        start_method="spawn",
    )
    vec_env = VecMonitor(vec_env, log_dir)

    print("Observation Space: ", vec_env.observation_space)
    print("Action Space: ", vec_env.action_space)

    # setup callbacks
    checkpoint_callback = CheckpointCallback(
        save_freq=save_eval_freq, save_path=log_dir
    )
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    vec_eval_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=eval_ids,
                segmodel_path=None,
                sample_mode=sample_mode,
                frame_num=frame_num,
                sample_num=pts_sample_num,
                other_handle_visible=other_handle_visible,
                num_classes=num_classes,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(eval_procs)
        ],
        start_method="spawn",
    )
    vec_eval_env = VecMonitor(vec_eval_env, log_dir)
    eval_callback = EvalCallback(
        vec_eval_env,
        n_eval_episodes=50,
        best_model_save_path=log_dir,
        log_path=log_dir,
        eval_freq=save_eval_freq,
        deterministic=True,
        render=False,
    )

    # set up logger
    new_logger = configure(log_dir, ["stdout", "csv", "tensorboard"])

    model = SAC(
        "MultiInputPolicy",
        vec_env,
        batch_size=256,  # 1024, 400
        ent_coef="auto_0.2",
        gamma=0.99,
        train_freq=16,  # 4, 64
        gradient_steps=16,  # 2, 4
        buffer_size=100000,
        learning_starts=800,
        use_sde=True,
        policy_kwargs=dict(
            log_std_init=-3.67,
            net_arch=[256, 256],
            features_extractor_class=CustomSegPNExtractor,
            features_extractor_kwargs=dict(
                num_classes=num_classes,
                state_repeat_times=state_repeat_times,
                sample_num=pts_sample_num,
            ),
            normalize_images=False,
            share_features_extractor=True,
        ),
        tensorboard_log=log_dir + "sac_opendoor_tb/",
        device=device,
        verbose=1,
    )
    # Set new logger
    model.set_logger(new_logger)
    model.learn(
        total_timesteps=3_000_000,
        callback=[checkpoint_callback, eval_callback],
    )

arti_mani/algorithms/rl_iam/sac/sac_train_segGTpts_PNfeat_multiarti.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import random

from arti_mani import VISUALMODEL_DIR
from arti_mani.algorithms.rl_iam.feature_extract import CustomSegPNExtractor
from arti_mani.algorithms.rl_iam.rl_utils import sb3_make_env_multiarti
from arti_mani.envs import *
from stable_baselines3 import SAC
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
from stable_baselines3.common.logger import configure
from stable_baselines3.common.vec_env import SubprocVecEnv, VecMonitor


def setup_seed(seed=1029):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "arti"  # "door", "drawer", "faucet", "arti"
    train_procs = 4
    eval_procs = 1
    save_eval_freq = 250000 // train_procs
    seed = np.random.RandomState().randint(2**32)
    print("experiment random seed: ", seed)
    setup_seed(seed)
    env_id = "ArtiMani-v0"
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069], [5004, 5007, 5023, 5069]
    if mode == "door":
        arti_ids = [1006, 1030, 1047, 1081]
        eval_ids = [0]
    elif mode == "drawer":
        arti_ids = [1005, 1016, 1024, 1076]
        eval_ids = [1]
    elif mode == "faucet":
        arti_ids = [5004, 5007, 5023, 5069]
        eval_ids = [5052]
    elif mode == "arti":
        arti_ids = [
            1006,
            1030,
            1047,
            1081,
            1005,
            1016,
            1024,
            1076,
            5004,
            5007,
            5023,
            5069,
        ]
        eval_ids = [0, 1, 5052]
    else:
        raise NotImplementedError(mode)
    segmodel_path = (
        VISUALMODEL_DIR
        / "smp_model/20230219_000940_train52-val18_384_noDR_randombg_aug_stereo_bs16_focalloss_0.5step50lr0.001_RGBDunet-163264128_mobilenet_v2"
    )
    other_handle_visible = False
    obs_mode = "state_egostereo_segpoints"
    control_mode = "pd_joint_delta_pos"
    num_classes = 6
    state_repeat_times = 10
    sample_mode = "random_sample"  # full_downsample, score_sample, random_sample, fps_sample, frameweight_sample
    frame_num = 1
    pts_sample_num = 32
    # Create log dir
    exp_suffix = "smp0219"  # _largerange_adjustgripper_ptsaddnoise0.01
    log_dir = f"./logs/sac{train_procs}_tr{len(arti_ids)}eval{len(eval_ids)}{mode}_{obs_mode}_{sample_mode}f{frame_num}s{pts_sample_num}_{exp_suffix}/"
    os.makedirs(log_dir, exist_ok=True)

    vec_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=arti_ids,
                segmodel_path=segmodel_path,
                sample_mode=sample_mode,
                frame_num=frame_num,
                sample_num=pts_sample_num,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(train_procs)
        ],
        start_method="spawn",
    )
    vec_env = VecMonitor(vec_env, log_dir)

    print("Observation Space: ", vec_env.observation_space)
    print("Action Space: ", vec_env.action_space)

    # setup callbacks
    checkpoint_callback = CheckpointCallback(
        save_freq=save_eval_freq, save_path=log_dir
    )
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    vec_eval_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=eval_ids,
                segmodel_path=segmodel_path,
                sample_mode=sample_mode,
                frame_num=frame_num,
                sample_num=pts_sample_num,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(eval_procs)
        ],
        start_method="spawn",
    )
    vec_eval_env = VecMonitor(vec_eval_env, log_dir)
    eval_callback = EvalCallback(
        vec_eval_env,
        n_eval_episodes=50,
        best_model_save_path=log_dir,
        log_path=log_dir,
        eval_freq=save_eval_freq,
        deterministic=True,
        render=False,
    )

    # set up logger
    new_logger = configure(log_dir, ["stdout", "csv", "tensorboard"])

    # check_env(env)
    model = SAC(
        "MultiInputPolicy",
        vec_env,
        batch_size=256,  # 1024, 400
        ent_coef="auto_0.2",
        gamma=0.99,
        train_freq=16,  # 4, 64
        gradient_steps=16,  # 2, 4
        buffer_size=100000,
        learning_starts=800,
        use_sde=True,
        policy_kwargs=dict(
            log_std_init=-3.67,
            net_arch=[256, 256],
            features_extractor_class=CustomSegPNExtractor,
            features_extractor_kwargs=dict(
                num_classes=num_classes,
                state_repeat_times=state_repeat_times,
                sample_num=pts_sample_num,
            ),
            normalize_images=False,
            share_features_extractor=True,
        ),
        tensorboard_log=log_dir + "sac_opendoor_tb/",
        device=device,
        verbose=1,
    )
    # Set new logger
    model.set_logger(new_logger)
    model.learn(
        total_timesteps=2_000_000,
        callback=[checkpoint_callback, eval_callback],
    )

arti_mani/algorithms/rl_iam/sac/sac_train_segpts_PNfeat.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import random

from arti_mani.algorithms.rl_iam.feature_extract import CustomStateExtractor
from arti_mani.algorithms.rl_iam.rl_utils import sb3_make_env_multiarti
from arti_mani.envs import *
from stable_baselines3 import SAC
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
from stable_baselines3.common.logger import configure
from stable_baselines3.common.vec_env import SubprocVecEnv, VecMonitor


def setup_seed(seed=1029):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "door"  # "door", "drawer", "faucet"
    train_procs = 4
    eval_procs = 1
    save_eval_freq = 100000 // train_procs
    seed = np.random.RandomState().randint(2**32)
    print("experiment random seed: ", seed)
    setup_seed(seed)
    env_id = "ArtiMani-v0"
    other_handle_visible = False
    # door: [0, 1006, 1030, 1047, 1081]
    # drawer: 1, [1045, 1054, 1063, 1067], [1004, 1005, 1016, 1024]
    # faucet: [5002, 5023, 5052, 5069] [5004, 5007, 5023, 5069]
    if mode == "door":
        arti_ids = [1006, 1030, 1047, 1081]
        eval_ids = [0]
    elif mode == "drawer":
        arti_ids = [1004, 1005, 1016, 1024]
        eval_ids = [1]
    elif mode == "faucet":
        arti_ids = [5004, 5007, 5023, 5069]
        eval_ids = [5052]
    else:
        raise NotImplementedError(mode)
    obs_mode = "state_dict"
    control_mode = "pd_joint_delta_pos"
    sample_mode = "full_downsample"  # full_downsample, score_sample, random_sample, frameweight_sample
    frame_num = 1
    num_classes = 6
    state_repeat_times = 10
    pts_sample_num = 30

    # Create log dir
    exp_suffix = "gripperclose_leftrew"
    log_dir = f"./logs/sac{train_procs}_tr{len(arti_ids)}eval{len(eval_ids)}{mode}_{obs_mode}_{pts_sample_num}tip40finger_{exp_suffix}/"
    os.makedirs(log_dir, exist_ok=True)

    vec_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=arti_ids,
                segmodel_path=None,
                sample_mode=sample_mode,
                frame_num=frame_num,
                sample_num=pts_sample_num,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(train_procs)
        ],
        start_method="spawn",
    )
    vec_env = VecMonitor(vec_env, log_dir)

    print("Observation Space: ", vec_env.observation_space)
    print("Action Space: ", vec_env.action_space)

    # setup callbacks
    checkpoint_callback = CheckpointCallback(
        save_freq=save_eval_freq, save_path=log_dir
    )
    eval_seed = np.random.RandomState().randint(2**32)
    print("experiment eval random seed: ", eval_seed)
    vec_eval_env = SubprocVecEnv(
        [
            sb3_make_env_multiarti(
                env_id=env_id,
                arti_ids=eval_ids,
                segmodel_path=None,
                sample_mode=sample_mode,
                frame_num=frame_num,
                sample_num=pts_sample_num,
                other_handle_visible=other_handle_visible,
                obs_mode=obs_mode,
                control_mode=control_mode,
                device=device,
                rank=i,
                seed=seed,
            )
            for i in range(eval_procs)
        ],
        start_method="spawn",
    )
    vec_eval_env = VecMonitor(vec_eval_env, log_dir)
    eval_callback = EvalCallback(
        vec_eval_env,
        n_eval_episodes=50,
        best_model_save_path=log_dir,
        log_path=log_dir,
        eval_freq=save_eval_freq,
        deterministic=True,
        render=False,
    )

    # set up logger
    new_logger = configure(log_dir, ["stdout", "csv", "tensorboard"])

    # check_env(env)
    model = SAC(
        "MultiInputPolicy",
        vec_env,
        batch_size=256,
        ent_coef="auto_0.2",
        gamma=0.99,
        train_freq=4,
        gradient_steps=2,
        buffer_size=100000,
        learning_starts=800,
        use_sde=True,
        policy_kwargs=dict(
            log_std_init=-3.67,
            net_arch=[256, 256],
            features_extractor_class=CustomStateExtractor,
        ),
        tensorboard_log=log_dir + "sac_opendoor_tb/",
        device=device,
        verbose=1,
    )
    # Set new logger
    model.set_logger(new_logger)
    model.learn(
        total_timesteps=1_000_000,
        callback=[checkpoint_callback, eval_callback],
    )

arti_mani/algorithms/rl_iam/sac/sac_train_state.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import random
from copy import deepcopy
from typing import List, Optional, Union

import albumentations as albu
import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from arti_mani.algorithms.visual_net.scripts.dataset import AddSaltPepperNoise
from segmentation_models_pytorch.base import (
    ClassificationHead,
    SegmentationHead,
    SegmentationModel,
)
from segmentation_models_pytorch.decoders.unet.decoder import UnetDecoder
from segmentation_models_pytorch.decoders.unet.model import Unet


class CustomUnet(Unet):
    def __init__(
        self,
        has_dropout: bool = True,
        encoder_name: str = "resnet34",
        encoder_depth: int = 5,
        encoder_weights: Optional[str] = "imagenet",
        decoder_use_batchnorm: bool = True,
        decoder_channels: List[int] = (256, 128, 64, 32, 16),
        decoder_attention_type: Optional[str] = None,
        in_channels: int = 3,
        classes: int = 1,
        activation: Optional[Union[str, callable]] = None,
        aux_params: Optional[dict] = None,
    ):
        super().__init__(
            encoder_name,
            encoder_depth,
            encoder_weights,
            decoder_use_batchnorm,
            decoder_channels,
            decoder_attention_type,
            in_channels,
            classes,
            activation,
            aux_params,
        )

        self.has_dropput = has_dropout
        self.dropout = nn.Dropout2d(p=0.5, inplace=False)

    def forward(self, x, activate_dropout=False):
        """Sequentially pass `x` trough model`s encoder, decoder and heads"""
        if not activate_dropout:
            has_dropout = self.has_dropput
            self.has_dropput = False

        self.check_input_shape(x)

        features = self.encoder(x)
        if self.has_dropput:
            features[-1] = self.dropout(features[-1])

        decoder_output = self.decoder(*features)
        if self.has_dropput:
            decoder_output = self.dropout(decoder_output)

        masks = self.segmentation_head(decoder_output)

        if self.classification_head is not None:
            labels = self.classification_head(features[-1])
            return masks, labels

        if not activate_dropout:
            self.has_dropput = has_dropout

        return masks

    @torch.no_grad()
    def predict(self, x):
        """Inference method. Switch model to `eval` mode, call `.forward(x)` with `torch.no_grad()`

        Args:
            x: 4D torch tensor with shape (batch_size, channels, height, width)

        Return:
            prediction: 4D torch tensor with shape (batch_size, classes, height, width)

        """
        if self.training:
            self.eval()

        x = self.forward(x, activate_dropout=False)

        return x


rgb_aug = albu.Compose(
    [
        # albu.GaussNoise(var_limit=0.001, p=1),
        albu.ColorJitter(brightness=0.01, contrast=0.01, saturation=0.01, hue=0.01, p=1)
    ]
)
depth_aug = albu.Compose(
    [
        # albu.GaussNoise(var_limit=0.001, p=1),
        AddSaltPepperNoise(density=0.001, p=1)
    ]
)


def test_flip(x: np.ndarray, flip_mode):
    # flip_mode in [0, 1, 2, 3]
    if flip_mode == 3:
        return x
    # 0: both 1: hori 2: vert
    if flip_mode == 0:
        return x[::-1, ::-1]
    if flip_mode == 1:
        return x[:, ::-1]
    if flip_mode == 2:
        return x[::-1]


def feature_exchange(f_rgb, f_d, p=0.2):
    exchange_f_rgb = f_rgb.clone()
    exchange_f_d = f_d.clone()
    C = f_rgb.shape[1]
    sC = int(C * p)
    for i in range(f_rgb.shape[0]):
        exchange_f_rgb[i, -sC:] = f_d[i, -sC:]
        exchange_f_d[i, -sC:] = f_rgb[i, -sC:]
    return exchange_f_rgb, exchange_f_d


class CustomMobileNetEncoder(nn.Module):
    def __init__(self, in_channels=4, depth=5, dropout_p=0.2, **kwargs):
        super().__init__(**kwargs)
        self._depth = depth
        self.in_channels = in_channels
        self.out_channels = (3, 16, 24, 32, 96)

        self.dropout = nn.Dropout2d(dropout_p, inplace=True)
        # copy encoders for rgb and depth
        self.features_rgb = torchvision.models.MobileNetV2()  # weight init inside
        self.features_depth = torchvision.models.MobileNetV2()
        del self.features_rgb.classifier, self.features_depth.classifier
        self.stages_rgb = [
            nn.Identity(),
            self.features_rgb.features[:2],
            self.features_rgb.features[2:4],
            self.features_rgb.features[4:7],
            self.features_rgb.features[7:14],
            self.features_rgb.features[14:],
        ]
        self.stages_d = [
            nn.Identity(),
            self.features_depth.features[:2],
            self.features_depth.features[2:4],
            self.features_depth.features[4:7],
            self.features_depth.features[7:14],
            self.features_depth.features[14:],
        ]
        # conv skip connection
        self.connections = nn.ModuleList([nn.Identity()])
        for i in range(1, depth + 1):
            self.connections.append(
                nn.Sequential(
                    nn.Conv2d(
                        self.out_channels[i],
                        self.out_channels[i],
                        kernel_size=3,
                        padding=1,
                    ),
                    nn.BatchNorm2d(self.out_channels[i]),
                    nn.ReLU(),
                )
            )

    def forward(self, x):
        features = []
        f_rgb, f_d = x[:, :3], x[:, 3:4]
        f_d = f_d.repeat(1, 3, 1, 1)
        for i in range(self._depth + 1):
            # print(f_rgb.shape, f_d.shape)
            if i == 0:
                f_rgb = self.stages_rgb[i](f_rgb)
                f_d = self.stages_d[i](f_d)
            else:
                f_rgb, f_d = feature_exchange(f_rgb, f_d)
                f_rgb = self.dropout(self.stages_rgb[i](f_rgb))
                f_d = self.dropout(self.stages_d[i](f_d))
            features.append(self.connections[i](f_rgb + f_d))

        return features


class SplitUnet(SegmentationModel):
    def __init__(
        self,
        in_channels: int = 4,
        dropout_p: float = 0.2,
        encoder_name: str = "resnet34",
        encoder_depth: int = 5,
        encoder_weights: Optional[str] = "imagenet",
        decoder_use_batchnorm: bool = True,
        decoder_channels: List[int] = (256, 128, 64, 32, 16),
        decoder_attention_type: Optional[str] = None,
        classes: int = 1,
        activation=None,
    ):
        super().__init__()

        self.encoder = CustomMobileNetEncoder(
            in_channels=in_channels, depth=encoder_depth
        )

        self.decoder = UnetDecoder(
            encoder_channels=self.encoder.out_channels,
            decoder_channels=decoder_channels,
            n_blocks=encoder_depth,
            use_batchnorm=decoder_use_batchnorm,
            center=True if encoder_name.startswith("vgg") else False,
            attention_type=decoder_attention_type,
        )

        self.segmentation_head = SegmentationHead(
            in_channels=decoder_channels[-1],
            out_channels=classes,
            activation=activation,
            kernel_size=3,
        )

        self.in_channels = in_channels

        self.dropout_p = dropout_p
        self.dropout = nn.Dropout2d(p=dropout_p)

        self.classification_head = None

        self.name = "u-{}".format(encoder_name)
        self.initialize()

    def forward(self, x, activate_dropout=False):
        """Sequentially pass `x` trough model`s encoder, decoder and heads."""
        # self.check_input_shape(x)
        if activate_dropout:
            cur_dropout_status = self.dropout.training
            self.dropout.train()
            self.encoder.dropout.train()

        # use split backbone to extract features
        features = self.encoder(x)

        decoder_output = self.decoder(*features)
        decoder_output = self.dropout(decoder_output)

        masks = self.segmentation_head(decoder_output)

        if activate_dropout:
            self.dropout.train(cur_dropout_status)
            self.encoder.dropout.train(cur_dropout_status)

        return masks

    def predict(self, x, T=10, dropout=False):
        # no aug
        if T == 0:
            return self.forward(x, activate_dropout=dropout)

        # predict with test-time augmentation
        device = x.device
        x = x.cpu().numpy().squeeze()
        x = x.transpose(1, 2, 0)

        masks = []
        for _ in range(T):
            # aug
            flip_mode = np.random.randint(0, 4)
            aug_x = test_flip(x, flip_mode)
            aug_x[..., :3] = rgb_aug(image=aug_x[..., :3])["image"]
            aug_x[..., 3:4] = depth_aug(image=aug_x[..., 3:4])["image"]
            # aug as np.array
            aug_x = aug_x.transpose(2, 0, 1)[None].copy()
            aug_x = torch.from_numpy(aug_x).to(device=device, dtype=torch.float32)
            # mean
            cur_seg = self.forward(aug_x, activate_dropout=dropout)
            cur_seg = cur_seg.detach().cpu().numpy().squeeze().transpose(1, 2, 0)
            cur_seg = test_flip(cur_seg, flip_mode)
            cur_seg = torch.from_numpy(cur_seg.copy().transpose(2, 0, 1))
            masks.append(cur_seg.to(device=device, dtype=torch.float32))
        masks = F.softmax(torch.stack(masks), 1)
        masks = masks.mean(0, keepdim=True)
        return masks


class CustomUnetNew(Unet):
    def __init__(
        self,
        dropout_p: float = 0.2,
        encoder_name: str = "resnet34",
        encoder_depth: int = 5,
        encoder_weights: Optional[str] = "imagenet",
        decoder_use_batchnorm: bool = True,
        decoder_channels: List[int] = (256, 128, 64, 32, 16),
        decoder_attention_type: Optional[str] = None,
        in_channels: int = 3,
        classes: int = 1,
    ):
        super().__init__(
            encoder_name,
            encoder_depth,
            encoder_weights,
            decoder_use_batchnorm,
            decoder_channels,
            decoder_attention_type,
            in_channels,
            classes,
        )

        self.dropout_p = dropout_p
        self.dropout = nn.Dropout2d(p=self.dropout_p)

    def forward(self, x, activate_dropout=False):
        """Sequentially pass `x` trough model`s encoder, decoder and heads."""
        # self.check_input_shape(x)
        if activate_dropout:
            cur_dropout_status = self.dropout.training
            self.dropout.train()

        features = self.encoder(x)
        features[-1] = self.dropout(features[-1])

        decoder_output = self.decoder(*features)
        decoder_output = self.dropout(decoder_output)

        masks = self.segmentation_head(decoder_output)

        if activate_dropout:
            self.dropout.train(cur_dropout_status)

        if self.classification_head is not None:
            labels = self.classification_head(features[-1])
            return masks, labels
        return masks

    def predict(self, x, T=10, dropout=False):
        # no aug
        if T == 0:
            return self.forward(x, activate_dropout=dropout)

        # predict with test-time augmentation
        device = x.device
        x = x.cpu().numpy().squeeze()
        x = x.transpose(1, 2, 0)

        masks = []
        for _ in range(T):
            # aug
            flip_mode = np.random.randint(0, 4)
            aug_x = test_flip(x, flip_mode)
            aug_x[..., :3] = rgb_aug(image=aug_x[..., :3])["image"]
            aug_x[..., 3:4] = depth_aug(image=aug_x[..., 3:4])["image"]
            # aug as np.array
            aug_x = aug_x.transpose(2, 0, 1)[None].copy()
            aug_x = torch.from_numpy(aug_x).to(device=device, dtype=torch.float32)
            # mean
            cur_seg = self.forward(aug_x, activate_dropout=dropout)
            cur_seg = cur_seg.detach().cpu().numpy().squeeze().transpose(1, 2, 0)
            cur_seg = test_flip(cur_seg, flip_mode)
            cur_seg = torch.from_numpy(cur_seg.copy().transpose(2, 0, 1))
            masks.append(cur_seg.to(device=device, dtype=torch.float32))
        masks = F.softmax(torch.stack(masks), 1)
        masks = masks.mean(0, keepdim=True)
        return masks


def weighted_label_smoothed_nll_loss(
    lprobs: torch.Tensor,
    target: torch.Tensor,
    weight: torch.Tensor,
    epsilon: float,
    ignore_index=None,
    reduction="mean",
    dim=-1,
) -> torch.Tensor:
    """NLL loss with label smoothing.

    References:
        https://github.com/pytorch/fairseq/blob/master/fairseq/criterions/label_smoothed_cross_entropy.py

    Args:
        lprobs (torch.Tensor): Log-probabilities of predictions (e.g after log_softmax)
    """
    if target.dim() == lprobs.dim() - 1:
        target = target.unsqueeze(dim)

    if ignore_index is not None:
        pad_mask = target.eq(ignore_index)
        target = target.masked_fill(pad_mask, 0)
        nll_loss = -lprobs.gather(dim=dim, index=target)
        smooth_loss = -lprobs.sum(dim=dim, keepdim=True)

        # nll_loss.masked_fill_(pad_mask, 0.0)
        # smooth_loss.masked_fill_(pad_mask, 0.0)
        nll_loss = nll_loss.masked_fill(pad_mask, 0.0)
        smooth_loss = smooth_loss.masked_fill(pad_mask, 0.0)
    else:
        nll_loss = -lprobs.gather(dim=dim, index=target)
        smooth_loss = -lprobs.sum(dim=dim, keepdim=True)

        nll_loss = nll_loss.squeeze(dim)
        smooth_loss = smooth_loss.squeeze(dim)

    # multiply weight
    # print(nll_loss.shape, weight.shape)
    nll_loss *= weight
    smooth_loss *= weight

    if reduction == "sum":
        nll_loss = nll_loss.sum()
        smooth_loss = smooth_loss.sum()
    if reduction == "mean":
        nll_loss = nll_loss.mean()
        smooth_loss = smooth_loss.mean()

    eps_i = epsilon / lprobs.size(dim)
    loss = (1.0 - epsilon) * nll_loss + eps_i * smooth_loss
    return loss


class WCELoss(nn.Module):
    __constants__ = ["reduction", "ignore_index", "smooth_factor"]

    def __init__(
        self,
        reduction: str = "mean",
        smooth_factor: Optional[float] = None,
        ignore_index: Optional[int] = -100,
        dim: int = 1,
    ):
        """Drop-in replacement for torch.nn.CrossEntropyLoss with
        label_smoothing.

        Args:
            smooth_factor: Factor to smooth target (e.g. if smooth_factor=0.1 then [1, 0, 0] -> [0.9, 0.05, 0.05])

        Shape
             - **y_pred** - torch.Tensor of shape (N, C, H, W)
             - **y_true** - torch.Tensor of shape (N, H, W)

        Reference
            https://github.com/BloodAxe/pytorch-toolbelt
        """
        super().__init__()
        self.smooth_factor = smooth_factor
        self.ignore_index = ignore_index
        self.reduction = reduction
        self.dim = dim

    def forward(
        self, y_pred: torch.Tensor, y_true: torch.Tensor, weight: torch.Tensor
    ) -> torch.Tensor:
        log_prob = F.log_softmax(y_pred, dim=self.dim)
        return weighted_label_smoothed_nll_loss(
            log_prob,
            y_true,
            weight,
            epsilon=self.smooth_factor,
            ignore_index=self.ignore_index,
            reduction=self.reduction,
            dim=self.dim,
        )

arti_mani/algorithms/visual_net/Networks/Custom_Unet.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from pathlib import Path
from typing import Union

import torch
import torch.nn as nn


class BaseModel(nn.Module):
    def __init__(self):
        super().__init__()

    @property
    def n_param(self):
        # type: (BaseModel) -> int
        """
        :return: number of parameters
        """
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    @property
    def is_cuda(self):
        # type: () -> bool
        """
        :return: `True` if the model is on Cuda; `False` otherwise
        """
        return next(self.parameters()).is_cuda

    def save_w(self, path):
        # type: (Union[str, Path]) -> None
        """
        save model weights in the specified path
        """
        torch.save(self.state_dict(), path)

    def load_w(self, path):
        # type: (Union[str, Path]) -> None
        """
        load model weights from the specified path
        """
        self.load_state_dict(torch.load(path))

    def requires_grad(self, flag):
        # type: (bool) -> None
        """
        param flag: True if the model requires gradient, False otherwise
        """

        for p in self.parameters():
            p.requires_grad = flag

arti_mani/algorithms/visual_net/Networks/base_model.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import copy
from typing import List

import torch
import torch.nn as nn
from arti_mani.algorithms.visual_net.Networks.base_model import BaseModel
from arti_mani.algorithms.visual_net.Networks.utils import (
    act_layer,
    norm_layer1d,
    norm_layer2d,
)

LRELU_SLOPE = 0.02


class DenseBlock(nn.Module):
    def __init__(self, in_features, out_features, norm=None, activation=None):
        super(DenseBlock, self).__init__()
        self.linear = nn.Linear(in_features, out_features)

        if activation is None:
            nn.init.xavier_uniform_(
                self.linear.weight, gain=nn.init.calculate_gain("linear")
            )
            nn.init.zeros_(self.linear.bias)
        elif activation == "tanh":
            nn.init.xavier_uniform_(
                self.linear.weight, gain=nn.init.calculate_gain("tanh")
            )
            nn.init.zeros_(self.linear.bias)
        elif activation == "lrelu":
            nn.init.kaiming_uniform_(
                self.linear.weight, a=LRELU_SLOPE, nonlinearity="leaky_relu"
            )
            nn.init.zeros_(self.linear.bias)
        elif activation == "relu":
            nn.init.kaiming_uniform_(self.linear.weight, nonlinearity="relu")
            nn.init.zeros_(self.linear.bias)
        else:
            raise ValueError()

        self.activation = None
        self.norm = None
        if norm is not None:
            self.norm = norm_layer1d(norm, out_features)
        if activation is not None:
            self.activation = act_layer(activation)

    def forward(self, x):
        x = self.linear(x)
        x = self.norm(x) if self.norm is not None else x
        x = self.activation(x) if self.activation is not None else x
        return x


class Conv2DBlock(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_sizes,
        strides,
        norm=None,
        activation=None,
        padding_mode="replicate",
    ):
        super(Conv2DBlock, self).__init__()
        padding = (
            kernel_sizes // 2
            if isinstance(kernel_sizes, int)
            else (kernel_sizes[0] // 2, kernel_sizes[1] // 2)
        )
        self.conv2d = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_sizes,
            strides,
            padding=padding,
            padding_mode=padding_mode,
        )

        if activation is None:
            nn.init.xavier_uniform_(
                self.conv2d.weight, gain=nn.init.calculate_gain("linear")
            )
            nn.init.zeros_(self.conv2d.bias)
        elif activation == "tanh":
            nn.init.xavier_uniform_(
                self.conv2d.weight, gain=nn.init.calculate_gain("tanh")
            )
            nn.init.zeros_(self.conv2d.bias)
        elif activation == "lrelu":
            nn.init.kaiming_uniform_(
                self.conv2d.weight, a=LRELU_SLOPE, nonlinearity="leaky_relu"
            )
            nn.init.zeros_(self.conv2d.bias)
        elif activation == "relu":
            nn.init.kaiming_uniform_(self.conv2d.weight, nonlinearity="relu")
            nn.init.zeros_(self.conv2d.bias)
        else:
            raise ValueError()

        self.activation = None
        self.norm = None
        if norm is not None:
            self.norm = norm_layer2d(norm, out_channels)
        if activation is not None:
            self.activation = act_layer(activation)

    def forward(self, x):
        x = self.conv2d(x)
        x = self.norm(x) if self.norm is not None else x
        x = self.activation(x) if self.activation is not None else x
        return x


class Conv2DUpsampleBlock(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_sizes,
        strides,
        norm=None,
        activation=None,
    ):
        super(Conv2DUpsampleBlock, self).__init__()
        layer = [
            Conv2DBlock(in_channels, out_channels, kernel_sizes, 1, norm, activation)
        ]
        if strides > 1:
            layer.append(
                nn.Upsample(scale_factor=strides, mode="bilinear", align_corners=False)
            )
        convt_block = Conv2DBlock(
            out_channels, out_channels, kernel_sizes, 1, norm, activation
        )
        layer.append(convt_block)
        self.conv_up = nn.Sequential(*layer)

    def forward(self, x):
        return self.conv_up(x)


class SiameseNet(BaseModel):
    def __init__(
        self,
        input_channels: List[int],
        filters: List[int],
        kernel_sizes: List[int],
        strides: List[int],
        norm: str = None,
        activation: str = "relu",
    ):
        super(SiameseNet, self).__init__()
        self._input_channels = input_channels
        self._filters = filters
        self._kernel_sizes = kernel_sizes
        self._strides = strides
        self._norm = norm
        self._activation = activation
        self.output_channels = filters[-1]  # * len(input_channels)

        self._siamese_blocks = nn.ModuleList()
        for i, ch in enumerate(self._input_channels):
            blocks = []
            for i, (filt, ksize, stride) in enumerate(
                zip(self._filters, self._kernel_sizes, self._strides)
            ):
                conv_block = Conv2DBlock(
                    ch, filt, ksize, stride, self._norm, self._activation
                )
                blocks.append(conv_block)
                ch = filt
            self._siamese_blocks.append(nn.Sequential(*blocks))
        self._fuse = Conv2DBlock(
            self._filters[-1] * len(self._siamese_blocks),
            self._filters[-1],
            1,
            1,
            self._norm,
            self._activation,
        )

    def forward(self, x):
        if len(x) != len(self._siamese_blocks):
            raise ValueError(
                "Expected a list of tensors of size %d." % len(self._siamese_blocks)
            )
        self.streams = [stream(y) for y, stream in zip(x, self._siamese_blocks)]
        y = self._fuse(torch.cat(self.streams, 1))
        return y


class Attention2DUNet(BaseModel):
    def __init__(
        self,
        siamese_net: SiameseNet,
        filters: List[int],
        kernel_sizes: List[int],
        strides: List[int],
        norm: str = None,
        activation: str = "relu",
        output_channels: int = 1,
        skip_connections: bool = True,
    ):
        super(Attention2DUNet, self).__init__()
        self._siamese_net = copy.deepcopy(siamese_net)
        self._input_channels = self._siamese_net.output_channels
        self._filters = filters
        self._kernel_sizes = kernel_sizes
        self._strides = strides
        self._norm = norm
        self._activation = activation
        self._output_channels = output_channels
        self._skip_connections = skip_connections

        self._down = []
        ch = self._input_channels
        for filt, ksize, stride in zip(
            self._filters, self._kernel_sizes, self._strides
        ):
            conv_block = Conv2DBlock(
                ch,
                filt,
                ksize,
                stride,
                self._norm,
                self._activation,
                padding_mode="replicate",
            )
            ch = filt
            self._down.append(conv_block)
        self._down = nn.ModuleList(self._down)

        reverse_conv_data = list(zip(self._filters, self._kernel_sizes, self._strides))
        reverse_conv_data.reverse()

        self._up = []
        for i, (filt, ksize, stride) in enumerate(reverse_conv_data):
            if i > 0 and self._skip_connections:
                # ch += reverse_conv_data[-i-1][0]
                ch += reverse_conv_data[i][0]
            convt_block = Conv2DUpsampleBlock(
                ch, filt, ksize, stride, self._norm, self._activation
            )
            ch = filt
            self._up.append(convt_block)
        self._up = nn.ModuleList(self._up)

        self._final_conv = Conv2DBlock(
            ch, self._output_channels, 3, 1, padding_mode="replicate"
        )

    def forward(self, observations):
        x = self._siamese_net(observations)
        _, _, h, w = x.shape
        layers_for_skip = []
        for l in self._down:
            x = l(x)
            layers_for_skip.append(x)
        self.latent = x  # (N, 16, H/4, W/4)
        layers_for_skip.reverse()
        for i, l in enumerate(self._up):
            if i > 0 and self._skip_connections:
                # Skip connections. Skip the first up layer.
                x = torch.cat([layers_for_skip[i], x], 1)
            x = l(x)
        x = self._final_conv(x)  # (N, 2, H, W)
        return x, self.latent


class WeightGen(BaseModel):
    """
    Unet seg feature encoding maps to generate weights
    """

    def __init__(
        self,
        feat_chan: int = 16,
        weight_num: int = 5,
        norm: str = None,
        activation: str = None,
    ):
        super(WeightGen, self).__init__()

        self.maxpool = nn.AdaptiveMaxPool2d(1)
        self.flatten = nn.Flatten()
        self.fc1 = DenseBlock(
            in_features=feat_chan,
            out_features=feat_chan * 4,
            norm=norm,
            activation=activation,
        )
        self.fc2 = DenseBlock(
            in_features=feat_chan * 4,
            out_features=feat_chan * 4,
            norm=norm,
            activation=activation,
        )
        self.fc3 = DenseBlock(
            in_features=feat_chan * 4,
            out_features=weight_num,
            norm=norm,
            activation=activation,
        )

    def forward(self, x):
        feat_max = self.flatten(self.maxpool(x))
        feat1 = self.fc1(feat_max)
        feat2 = self.fc2(feat1)
        weight = self.fc3(feat2)
        return weight


class FeatAggNet(BaseModel):
    """
    Unet seg feature maps aggregation net
    """

    def __init__(
        self,
        in_chan: int = 5,
        out_chan: int = 1,
        norm: str = None,
        activation: str = None,
    ):
        super(FeatAggNet, self).__init__()

        self.conv1 = Conv2DBlock(
            in_channels=in_chan,
            out_channels=out_chan,
            kernel_sizes=1,
            strides=1,
            norm=norm,
            activation=activation,
        )

    def forward(self, x):
        feat_norm = self.conv1(x)
        attn_map = torch.sigmoid(feat_norm)
        return attn_map


class AttentionDepthUNet(BaseModel):
    def __init__(
        self,
        depth_feat_channels: int,
        filters: List[int],
        kernel_sizes: List[int],
        strides: List[int],
        norm: str = None,
        activation: str = "relu",
        output_channels: int = 1,
        skip_connections: bool = True,
    ):
        super(AttentionDepthUNet, self).__init__()
        self._depth_feat_channels = depth_feat_channels
        self._filters = filters
        self._kernel_sizes = kernel_sizes
        self._strides = strides
        self._norm = norm
        self._activation = activation
        self._output_channels = output_channels
        self._skip_connections = skip_connections

        self.depth_conv_block = Conv2DBlock(
            in_channels=1,
            out_channels=self._depth_feat_channels,
            kernel_sizes=5,
            strides=1,
            norm=None,
            activation=self._norm,
        )

        self._down = []
        ch = self._depth_feat_channels
        for filt, ksize, stride in zip(
            self._filters, self._kernel_sizes, self._strides
        ):
            conv_block = Conv2DBlock(
                ch,
                filt,
                ksize,
                stride,
                self._norm,
                self._activation,
                padding_mode="replicate",
            )
            ch = filt
            self._down.append(conv_block)
        self._down = nn.ModuleList(self._down)

        reverse_conv_data = list(zip(self._filters, self._kernel_sizes, self._strides))
        reverse_conv_data.reverse()

        self._up = []
        for i, (filt, ksize, stride) in enumerate(reverse_conv_data):
            if i > 0 and self._skip_connections:
                # ch += reverse_conv_data[-i-1][0]
                ch += reverse_conv_data[i][0]
            convt_block = Conv2DUpsampleBlock(
                ch, filt, ksize, stride, self._norm, self._activation
            )
            ch = filt
            self._up.append(convt_block)
        self._up = nn.ModuleList(self._up)

        self._final_conv = Conv2DBlock(
            ch, self._output_channels, 3, 1, padding_mode="replicate"
        )
        # self.weight = torch.nn.Parameter(
        #                     torch.ones(ch) / ch, requires_grad=True
        #                 )

    def forward(self, observations):
        x = self.depth_conv_block(observations)
        _, _, h, w = x.shape
        self.ups = []
        self.downs = []
        layers_for_skip = []
        for l in self._down:
            x = l(x)
            layers_for_skip.append(x)
            self.downs.append(x)
        self.latent = x  # (N, 16, H/4, W/4)
        layers_for_skip.reverse()
        for i, l in enumerate(self._up):
            if i > 0 and self._skip_connections:
                # Skip connections. Skip the first up layer.
                x = torch.cat([layers_for_skip[i], x], 1)
            x = l(x)
            self.ups.append(x)
        x = self._final_conv(x)  # (N, 2, H, W)
        # x = F.softmax(x, dim=1)  # (N, 16, H, W)
        # weight = F.softmax(self.weight, dim=0)
        # soft_weight_map = torch.einsum("ijkl,j->ikl", [x, weight]) # (N, H, W)
        return x


class UNetEncoder(BaseModel):
    def __init__(
        self,
        siamese_net: SiameseNet,
        filters: List[int],
        kernel_sizes: List[int],
        strides: List[int],
        norm: str = None,
        activation: str = "relu",
    ):
        super(UNetEncoder, self).__init__()
        self._siamese_net = copy.deepcopy(siamese_net)
        self._input_channels = self._siamese_net.output_channels
        self._filters = filters
        self._kernel_sizes = kernel_sizes
        self._strides = strides
        self._norm = norm
        self._activation = activation
        self._down = []
        ch = self._input_channels
        for filt, ksize, stride in zip(
            self._filters, self._kernel_sizes, self._strides
        ):
            conv_block = Conv2DBlock(
                ch,
                filt,
                ksize,
                stride,
                self._norm,
                self._activation,
                padding_mode="replicate",
            )
            ch = filt
            self._down.append(conv_block)
        self._down = nn.ModuleList(self._down)
        self.maxpool = nn.AdaptiveMaxPool2d(1)  # (N, C, 1, 1)
        self.flatten = nn.Flatten()

    def forward(self, observations):
        x = self._siamese_net(observations)
        _, _, h, w = x.shape
        self.downs = []
        for l in self._down:
            x = l(x)
            self.downs.append(x)
        x = self.maxpool(x)  # (N, 16, H/4, W/4) => (N, 16)
        max_feat = self.flatten(x)  # (N, 16)
        return max_feat


class GeneralCNN(BaseModel):
    def __init__(
        self,
        siamese_net: SiameseNet,
        filters: List[int],
        kernel_sizes: List[int],
        strides: List[int],
        norm: str = None,
        activation: str = "relu",
    ):
        super(GeneralCNN, self).__init__()
        self._siamese_net = copy.deepcopy(siamese_net)
        self._input_channels = self._siamese_net.output_channels
        self._filters = filters
        self._kernel_sizes = kernel_sizes
        self._strides = strides
        self._norm = norm
        self._activation = activation
        self._down = []
        blocks_num = len(self._filters)
        ch = self._input_channels
        for filt, ksize, stride in zip(
            self._filters, self._kernel_sizes, self._strides
        ):
            conv_block = Conv2DBlock(
                ch,
                filt,
                ksize,
                stride,
                self._norm,
                self._activation,
                padding_mode="replicate",
            )
            ch = filt
            self._down.append(conv_block)
        self._down = nn.ModuleList(self._down)
        self.flatten = nn.Flatten()
        cnnfeat_size = ch * 72 * 128 // 4**blocks_num
        self.final = nn.Sequential(
            DenseBlock(cnnfeat_size, 1024, activation=self._activation),
            DenseBlock(1024, 128, activation=self._activation),
        )

    def forward(self, observations):
        x = self._siamese_net(observations)
        _, _, h, w = x.shape
        self.downs = []
        for l in self._down:
            x = l(x)
            self.downs.append(x)
        x = self.flatten(x)  # (N, 16, H/4, W/4) => (N, -1)
        final_feat = self.final(x)  # (N, 128)
        return final_feat


if __name__ == "__main__":
    data = [torch.zeros([4, 3, 72, 128]), torch.zeros([4, 1, 72, 128])]
    siamese_net = SiameseNet(
        input_channels=[3, 1],
        filters=[8],
        kernel_sizes=[5],
        strides=[1],
        activation="lrelu",
        norm=None,
    )
    attention_net = Attention2DUNet(
        siamese_net=siamese_net,
        filters=[16, 32, 64],
        kernel_sizes=[5, 5, 5],
        strides=[2, 2, 2],
        output_channels=2,
        norm=None,
        activation="lrelu",
        skip_connections=True,
    )
    for param in attention_net.parameters():
        param.requires_grad = True
    y, latent = attention_net(data)
    print(f"Attention2DUNet output size: {y.shape}, {latent.shape}")

    seg_map = torch.randn(4, 5, 72, 128)
    featagg = FeatAggNet(
        in_chan=5,
        out_chan=1,
        norm=None,
        activation="relu",
    )
    y = featagg(seg_map)
    print(f"FeatAggNet output size: {y.shape}")

    unet_encoder = UNetEncoder(
        siamese_net=siamese_net,
        filters=[16, 16],
        kernel_sizes=[5, 5],
        strides=[2, 2],
        norm=None,
        activation="lrelu",
    )
    y = unet_encoder(data)
    print(f"UNetEncoder output size: {y.shape}")

    gene_cnn = GeneralCNN(
        siamese_net=siamese_net,
        filters=[16, 16],
        kernel_sizes=[5, 5],
        strides=[2, 2],
        norm=None,
        activation="lrelu",
    )
    y = gene_cnn(data)
    print(f"GeneralCNN output size: {y.shape}")

    attn_depth = AttentionDepthUNet(
        depth_feat_channels=8,
        filters=[16, 16],
        kernel_sizes=[5, 5],
        strides=[2, 2],
        output_channels=2,
        norm=None,
        activation="lrelu",
        skip_connections=True,
    )
    y = attn_depth(data[1])
    print(f"AttentionDepthUNet output size: {y.shape}")

arti_mani/algorithms/visual_net/Networks/conv2dnet.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import torch
import torch.nn as nn


class ImgEncoder(nn.Module):
    """Convolutional encoder for image-based observations."""

    def __init__(self, in_ch, feature_dim, num_layers=2, num_filters=32):
        super().__init__()

        # assert len(obs_shape) == 3
        self.num_layers = num_layers
        self.num_filters = num_filters
        self.output_dim = 8 * 15
        self.output_logits = False
        self.feature_dim = feature_dim

        self.convs = nn.ModuleList(
            [
                nn.Conv2d(in_ch, self.num_filters, 3, stride=2),
            ]
        )

        for i in range(num_layers - 1):
            self.convs.append(
                nn.Conv2d(self.num_filters, self.num_filters, 3, stride=2)
            )

        self.head = nn.Sequential(
            nn.Linear(self.num_filters * self.output_dim, self.feature_dim),
            nn.LayerNorm(self.feature_dim),
        )

        self.outputs = dict()

    def forward_conv(self, obs):
        obs = obs / 255.0

        conv = torch.relu(self.convs[0](obs))

        for i in range(1, self.num_layers):
            conv = torch.relu(self.convs[i](conv))

        h = conv.view(conv.size(0), -1)
        return h

    def forward(self, obs, detach=False):
        h = self.forward_conv(obs)

        if detach:
            h = h.detach()

        out = self.head(h)
        if not self.output_logits:
            out = torch.tanh(out)

        return out


if __name__ == "__main__":
    data = torch.randn(4, 7, 144, 256)
    net_encoder = ImgEncoder(in_ch=7, feature_dim=256, num_layers=4, num_filters=32)
    print(net_encoder)
    out = net_encoder(data)
    print(out.shape)

arti_mani/algorithms/visual_net/Networks/img_encoder.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import math

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models


def _assert_no_grad(tensor):
    assert not tensor.requires_grad, (
        "nn criterions don't compute the gradient w.r.t. targets - please "
        "mark these tensors as not requiring gradients"
    )


def soft_argmax(heatmaps, joint_num, output_shape):
    """
    Args:
        heatmaps: (N, J*D, H, W)
        joint_num: J, num of keypoints
        output_shape: (D, H, W)

    Returns:

    """
    assert isinstance(heatmaps, torch.Tensor)
    d, h, w = output_shape
    heatmaps = heatmaps.reshape((-1, joint_num, d * h * w))
    heatmaps = F.softmax(heatmaps, 2)
    heatmaps = heatmaps.reshape((-1, joint_num, d, h, w))

    accu_x = heatmaps.sum(dim=(2, 3))  # (N, J, W)
    accu_y = heatmaps.sum(dim=(2, 4))  # (N, J, H)
    accu_z = heatmaps.sum(dim=(3, 4))  # (N, J, D)

    # accu_x = accu_x * torch.arange(1, w + 1).float().to(accu_x.device)
    # accu_y = accu_y * torch.arange(1, h + 1).float().to(accu_y.device)
    # accu_z = accu_z * torch.arange(1, d + 1).float().to(accu_z.device)
    accu_x = accu_x * (torch.arange(0, w) / w).float().to(
        accu_x.device
    )  # normalized the coord [0, 1)
    accu_y = accu_y * (torch.arange(0, h) / h).float().to(accu_y.device)
    accu_z = accu_z * (torch.arange(0, d) / d).float().to(accu_z.device)

    accu_x = accu_x.sum(dim=2, keepdim=True)  # (N, J, 1)
    accu_y = accu_y.sum(dim=2, keepdim=True)  # (N, J, 1)
    accu_z = accu_z.sum(dim=2, keepdim=True)  # (N, J, 1)

    coord_out = torch.cat((accu_x, accu_y, accu_z), dim=2)  # (N, J, 3)

    return coord_out


class JointLocationLoss(nn.Module):
    def __init__(self):
        super(JointLocationLoss, self).__init__()

    def forward(self, heatmap_out, gt_coord, out_shape, gt_vis):
        joint_num = gt_coord.shape[1]
        coord_out = soft_argmax(heatmap_out, joint_num, out_shape)  # (N, J, 3)

        ## set not visable coord_out == gt_coord to make it not useful
        not_vis_coord = gt_vis == 0  # (N, J)
        gt_coord[not_vis_coord] = coord_out[not_vis_coord].clone().detach()

        _assert_no_grad(gt_coord)  # (N, J, 3)
        _assert_no_grad(gt_vis)  # (N, J)

        ## l1 loss
        # loss = F.l1_loss(coord_out, gt_coord, reduction="mean")
        # loss = torch.abs(coord_out - gt_coord)  # (N, J, 3)
        # loss = torch.einsum('ijk,ij->ijk', torch.abs(coord_out - gt_coord), gt_vis)  # (N, J, 3)
        ## smooth l1 loss
        # loss = F.smooth_l1_loss(coord_out, gt_coord) * 10
        ## l2 loss
        loss = F.mse_loss(coord_out, gt_coord) * 100

        # loss = (loss[:, :, 0] + loss[:, :, 1] + loss[:, :, 2]) / 3.
        # loss = loss.mean()

        return loss


class IntegralHumanPoseModel(nn.Module):
    def __init__(
        self,
        num_keypoints=3,
        num_deconv_layers=3,
        depth_dim=64,
        has_dropout=True,
        mode="RGBD",
    ):
        super(IntegralHumanPoseModel, self).__init__()

        self.mode = mode
        init_in_ch = 32  # 32, 64
        # init RGBD input layers
        if mode == "RGBD":
            self.init_layer = nn.Sequential(
                nn.Conv2d(
                    4,
                    init_in_ch,
                    kernel_size=(3, 3),
                    stride=(2, 2),
                    padding=(1, 1),
                    bias=False,
                ),
                nn.BatchNorm2d(init_in_ch),  # 32, 64
                nn.ReLU(inplace=True),
                # nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
            )
        elif self.mode == "RGB":
            pass
        elif self.mode == "D":
            self.init_layer = nn.Sequential(
                nn.Conv2d(
                    1,
                    init_in_ch,
                    kernel_size=(3, 3),
                    stride=(2, 2),
                    padding=(1, 1),
                    bias=False,
                ),
                nn.BatchNorm2d(init_in_ch),  # 32, 64
                nn.ReLU(inplace=True),
                nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
            )
        else:
            raise NotImplementedError(self.mode)

        # Feature extraction network
        # resnet = models.resnet34(pretrained=True)
        # resnet = models.resnet18(pretrained=True)
        # self.feature_extraction = nn.Sequential(*list(resnet.children())[4:-2])

        mobilenet = models.mobilenet_v2(pretrained=True)
        # self.feature_extraction = nn.Sequential(*list(mobilenet.features[1:-1]))
        # self.feature_extraction = nn.Sequential(*list(mobilenet.features[:-1]))
        if self.mode == "RGBD" or self.mode == "D":
            self.feature_extraction = nn.Sequential(*list(mobilenet.features[1:-1]))
            # self.feature_extraction = nn.Sequential(*list(mobilenet.features[1:8]))
        elif self.mode == "RGB":
            # self.feature_extraction = nn.Sequential(*list(mobilenet.features[:-1]))
            self.feature_extraction = nn.Sequential(*list(mobilenet.features[:8]))

        # self.middle_net = nn.Sequential(
        #     nn.Conv2d(64, 512, kernel_size=(1, 1), stride=(1, 1), bias=False),
        #     nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
        #     nn.ReLU(inplace=True),
        #     nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False),
        #     nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
        #     nn.ReLU(inplace=True),
        #     nn.Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1), bias=False),
        #     nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
        # )

        self.has_dropout = has_dropout
        self.dropout = nn.Dropout2d(p=0.2, inplace=False)

        # Intermediate supervision layers
        # in_ch = 64
        in_ch = 320
        out_ch = 256
        deconv_layers = []
        for i in range(num_deconv_layers):
            deconv_layers.append(
                nn.ConvTranspose2d(
                    in_channels=in_ch,
                    out_channels=out_ch,
                    kernel_size=4,
                    stride=2,
                    padding=1,
                    output_padding=0,
                    bias=False,
                )
            )
            deconv_layers.append(nn.BatchNorm2d(out_ch))
            deconv_layers.append(nn.ReLU(inplace=True))
            in_ch = out_ch
        self.deconv_layers = nn.Sequential(*deconv_layers)

        # Output layers
        self.final_layer = nn.Conv2d(
            in_channels=in_ch,
            out_channels=num_keypoints * depth_dim,
            kernel_size=1,
            stride=1,
            padding=0,
        )

    def forward(self, x):
        # input layer
        if self.mode == "RGBD" or self.mode == "D":
            x = self.init_layer(x)

        # Feature extraction
        features = self.feature_extraction(x)
        # features = self.middle_net(features)
        if self.has_dropout:
            features = self.dropout(features)

        # Intermediate supervision
        deconv_result = self.deconv_layers(features)
        if self.has_dropout:
            deconv_result = self.dropout(deconv_result)

        # Output
        out = self.final_layer(deconv_result)

        return out

    def init_weights(self):
        for name, m in self.named_modules():
            if "feature_extraction" not in name:
                if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                    n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                    m.weight.data.normal_(0, math.sqrt(2.0 / n))
                elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d, nn.BatchNorm2d)):
                    nn.init.constant_(m.weight, 1)
                    nn.init.constant_(m.bias, 0)
                elif isinstance(m, nn.Linear):
                    nn.init.kaiming_normal_(m.weight)
                    if m.bias is not None:
                        nn.init.constant_(m.bias, 0)


if __name__ == "__main__":
    estmodel = IntegralHumanPoseModel()
    estmodel.init_weights()
    print(estmodel)

    data = torch.randn((10, 4, 144, 256))

    out = estmodel(data)
    print(data.shape, out.shape)
    model_params = sum(p.numel() for p in estmodel.parameters() if p.requires_grad)
    print(f"* number of parameters: {model_params}")

arti_mani/algorithms/visual_net/Networks/keypoint_detection.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from __future__ import print_function

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.parallel
import torch.utils.data
from torch.autograd import Variable


class STN3d(nn.Module):
    def __init__(self):
        super(STN3d, self).__init__()
        self.conv1 = torch.nn.Conv1d(3, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 9)
        self.relu = nn.ReLU()

        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.bn4 = nn.BatchNorm1d(512)
        self.bn5 = nn.BatchNorm1d(256)

    def forward(self, x):
        batchsize = x.size()[0]
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(-1, 1024)

        x = F.relu(self.bn4(self.fc1(x)))
        x = F.relu(self.bn5(self.fc2(x)))
        x = self.fc3(x)

        iden = (
            Variable(
                torch.from_numpy(
                    np.array([1, 0, 0, 0, 1, 0, 0, 0, 1]).astype(np.float32)
                )
            )
            .view(1, 9)
            .repeat(batchsize, 1)
        )
        if x.is_cuda:
            iden = iden.cuda()
        x = x + iden
        x = x.view(-1, 3, 3)
        return x


class STNkd(nn.Module):
    def __init__(self, k=64):
        super(STNkd, self).__init__()
        self.conv1 = torch.nn.Conv1d(k, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, k * k)
        self.relu = nn.ReLU()

        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.bn4 = nn.BatchNorm1d(512)
        self.bn5 = nn.BatchNorm1d(256)

        self.k = k

    def forward(self, x):
        batchsize = x.size()[0]
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(-1, 1024)

        x = F.relu(self.bn4(self.fc1(x)))
        x = F.relu(self.bn5(self.fc2(x)))
        x = self.fc3(x)

        iden = (
            Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32)))
            .view(1, self.k * self.k)
            .repeat(batchsize, 1)
        )
        if x.is_cuda:
            iden = iden.cuda()
        x = x + iden
        x = x.view(-1, self.k, self.k)
        return x


class PointNetfeat(nn.Module):
    def __init__(
        self,
        in_ch=3,
        mlp_specs=None,
        global_feat=True,
        xyz_transform=False,
        feature_transform=False,
    ):
        super(PointNetfeat, self).__init__()
        if mlp_specs is None:
            mlp_specs = [64, 128, 512]
        assert in_ch >= 3
        self.in_ch = in_ch
        self.mlp_specs = mlp_specs
        self.global_feat = global_feat
        self.xyz_transform = xyz_transform
        if self.xyz_transform:
            self.stn = STN3d()
        # self.fc = nn.Sequential(
        #     nn.Linear(mlp_specs[2], 256),
        #     nn.Linear(256, 128)
        # )
        self.conv1 = torch.nn.Conv1d(in_ch, mlp_specs[0], 1)
        self.conv2 = torch.nn.Conv1d(mlp_specs[0], mlp_specs[1], 1)
        self.conv3 = torch.nn.Conv1d(mlp_specs[1], mlp_specs[2], 1)
        # self.bn1 = nn.BatchNorm1d(mlp_specs[0])
        # self.bn2 = nn.BatchNorm1d(mlp_specs[1])
        # self.bn3 = nn.BatchNorm1d(mlp_specs[2])
        self.norm1 = nn.LayerNorm(mlp_specs[0], eps=1e-6)
        self.norm2 = nn.LayerNorm(mlp_specs[1], eps=1e-6)
        self.norm3 = nn.LayerNorm(mlp_specs[2], eps=1e-6)
        self.feature_transform = feature_transform
        if self.feature_transform:
            self.fstn = STNkd(k=64)

    def forward(self, x):
        n_pts = x.size()[2]
        if self.xyz_transform:
            if self.in_ch > 3:
                xyz, feat = x[:, :3], x[:, 3:]
                trans = self.stn(xyz)
                xyz = xyz.transpose(2, 1)
                xyz = torch.bmm(xyz, trans)
                xyz = xyz.transpose(2, 1)
                x = torch.cat((xyz, feat), dim=1)
            else:
                trans = self.stn(x)
                x = x.transpose(2, 1)
                x = torch.bmm(x, trans)
                x = x.transpose(2, 1)
        else:
            trans = None
            x = F.relu(self.norm1(self.conv1(x).transpose(2, 1)))
            x = x.transpose(2, 1)

        if self.feature_transform:
            trans_feat = self.fstn(x)
            x = x.transpose(2, 1)
            x = torch.bmm(x, trans_feat)
            x = x.transpose(2, 1)
        else:
            trans_feat = None

        pointfeat = x
        x = F.relu(self.norm2(self.conv2(x).transpose(2, 1)))
        x = x.transpose(2, 1)
        x = self.norm3(self.conv3(x).transpose(2, 1))
        x = x.transpose(2, 1)
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(-1, self.mlp_specs[2])
        if self.global_feat:
            # x = self.fc(x)
            return x, trans, trans_feat
        else:
            x = x.view(-1, self.mlp_specs[2], 1).repeat(1, 1, n_pts)
            return torch.cat([x, pointfeat], 1), trans, trans_feat
        # global_x = self.fc(x)
        # cat_x = torch.cat([x.view(-1, self.mlp_specs[2], 1).repeat(1, 1, n_pts), pointfeat], 1)
        # return global_x, cat_x, trans, trans_feat


class PointNetCls(nn.Module):
    def __init__(self, k=2, feature_transform=False):
        super(PointNetCls, self).__init__()
        self.feature_transform = feature_transform
        self.feat = PointNetfeat(feature_transform=feature_transform)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, k)
        self.dropout = nn.Dropout(p=0.3)
        self.bn1 = nn.BatchNorm1d(512)
        self.bn2 = nn.BatchNorm1d(256)
        self.relu = nn.ReLU()

    def forward(self, x):
        x, trans, trans_feat = self.feat(x)
        x = F.relu(self.bn1(self.fc1(x)))
        x = F.relu(self.bn2(self.dropout(self.fc2(x))))
        x = self.fc3(x)
        return F.log_softmax(x, dim=1), trans, trans_feat


class PointNetDenseCls(nn.Module):
    def __init__(
        self, in_ch=3, mlp_specs=None, k=2, xyz_transform=False, feature_transform=False
    ):
        super(PointNetDenseCls, self).__init__()
        if mlp_specs is None:
            mlp_specs = [64, 128, 512]
        self.k = k
        self.feature_transform = feature_transform
        self.feat = PointNetfeat(
            in_ch=in_ch,
            mlp_specs=mlp_specs,
            xyz_transform=xyz_transform,
            feature_transform=feature_transform,
        )
        self.conv1 = torch.nn.Conv1d(mlp_specs[0] + mlp_specs[2], 512, 1)
        self.conv2 = torch.nn.Conv1d(512, 256, 1)
        self.conv3 = torch.nn.Conv1d(256, 128, 1)
        self.conv4 = torch.nn.Conv1d(128, self.k, 1)
        self.bn1 = nn.BatchNorm1d(512)
        self.bn2 = nn.BatchNorm1d(256)
        self.bn3 = nn.BatchNorm1d(128)

    def forward(self, x):
        # batchsize = x.size()[0]
        # n_pts = x.size()[2]
        gloabal_x, cat_x, trans, trans_feat = self.feat(x)  # (N, 128), (1088, N)
        cat_x = F.relu(self.bn1(self.conv1(cat_x)))
        cat_x = F.relu(self.bn2(self.conv2(cat_x)))
        cat_x = F.relu(self.bn3(self.conv3(cat_x)))
        cat_x = self.conv4(cat_x)
        cat_x = cat_x.transpose(2, 1).contiguous()  # (N, 2)
        # cat_x = F.log_softmax(cat_x.view(-1,self.k), dim=-1)
        # cat_x = cat_x.view(batchsize, n_pts, self.k)
        return gloabal_x, cat_x, trans, trans_feat


def feature_transform_regularizer(trans):
    d = trans.size()[1]
    batchsize = trans.size()[0]
    I = torch.eye(d)[None, :, :]
    if trans.is_cuda:
        I = I.cuda()
    loss = torch.mean(
        torch.norm(torch.bmm(trans, trans.transpose(2, 1)) - I, dim=(1, 2))
    )
    return loss


if __name__ == "__main__":
    sim_data = Variable(torch.rand(32, 3, 2500))
    trans = STN3d()
    out = trans(sim_data)
    print("stn", out.size())
    print("loss", feature_transform_regularizer(out))

    sim_data_64d = Variable(torch.rand(32, 64, 2500))
    trans = STNkd(k=64)
    out = trans(sim_data_64d)
    print("stn64d", out.size())
    print("loss", feature_transform_regularizer(out))

    pointfeat = PointNetfeat()
    global_x, cat_x, _, _ = pointfeat(sim_data)
    print("global feat, point feat: ", global_x.size(), cat_x.size())

    # cls = PointNetCls(k = 5)
    # out, _, _ = cls(sim_data)
    # print('class', out.size())

    seg = PointNetDenseCls(k=2)
    global_x, cat_x, _, _ = seg(sim_data)
    print("global feat, seg map: ", global_x.size(), cat_x.size())

arti_mani/algorithms/visual_net/Networks/pointnet.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from typing import Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

LRELU_SLOPE = 0.02


def act_layer(act):
    if act == "relu":
        return nn.ReLU()
    elif act == "lrelu":
        return nn.LeakyReLU(LRELU_SLOPE)
    elif act == "elu":
        return nn.ELU()
    elif act == "tanh":
        return nn.Tanh()
    elif act == "prelu":
        return nn.PReLU()
    else:
        raise ValueError("%s not recognized." % act)


def norm_layer2d(norm, channels):
    if norm == "batch":
        return nn.BatchNorm2d(channels)
    elif norm == "instance":
        return nn.InstanceNorm2d(channels, affine=True)
    elif norm == "layer":
        return nn.GroupNorm(1, channels, affine=True)
    elif norm == "group":
        return nn.GroupNorm(4, channels, affine=True)
    else:
        raise ValueError("%s not recognized." % norm)


def norm_layer1d(norm, num_channels):
    if norm == "batch":
        return nn.BatchNorm1d(num_channels)
    elif norm == "instance":
        return nn.InstanceNorm1d(num_channels, affine=True)
    elif norm == "layer":
        return nn.LayerNorm(num_channels)
    else:
        raise ValueError("%s not recognized." % norm)


def create_meshgrid(
    x: torch.Tensor, normalized_coordinates: Optional[bool]
) -> (torch.Tensor, torch.Tensor):
    assert len(x.shape) == 4, x.shape
    _, _, height, width = x.shape
    _device, _dtype = x.device, x.dtype
    if normalized_coordinates:
        xs = torch.linspace(-1.0, 1.0, width, device=_device, dtype=_dtype)
        ys = torch.linspace(-1.0, 1.0, height, device=_device, dtype=_dtype)
    else:
        xs = torch.linspace(0, width - 1, width, device=_device, dtype=_dtype)
        ys = torch.linspace(0, height - 1, height, device=_device, dtype=_dtype)
    return torch.meshgrid(ys, xs)  # pos_y, pos_x


class SpatialSoftArgmax2d(nn.Module):
    r"""Creates a module that computes the Spatial Soft-Argmax 2D
    of a given input heatmap.

    Returns the index of the maximum 2d coordinates of the give map.
    The output order is x-coord and y-coord.

    Arguments:
        normalized_coordinates (Optional[bool]): wether to return the
          coordinates normalized in the range of [-1, 1]. Otherwise,
          it will return the coordinates in the range of the input shape.
          Default is True.

    Shape:
        - Input: :math:`(B, N, H, W)`
        - Output: :math:`(B, N, 2)`

    Examples::
        >>> input = torch.rand(1, 4, 2, 3)
        >>> m = tgm.losses.SpatialSoftArgmax2d()
        >>> coords = m(input)  # 1x4x2
        >>> x_coord, y_coord = torch.chunk(coords, dim=-1, chunks=2)
    """

    def __init__(self, normalized_coordinates: Optional[bool] = True) -> None:
        super(SpatialSoftArgmax2d, self).__init__()
        self.normalized_coordinates: Optional[bool] = normalized_coordinates
        self.eps: float = 1e-6

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if not torch.is_tensor(input):
            raise TypeError(
                "Input input type is not a torch.Tensor. Got {}".format(type(input))
            )
        if not len(input.shape) == 4:
            raise ValueError(
                "Invalid input shape, we expect BxCxHxW. Got: {}".format(input.shape)
            )
        # unpack shapes and create view from input tensor
        batch_size, channels, height, width = input.shape
        x: torch.Tensor = input.view(batch_size, channels, -1)

        # compute softmax with max substraction trick
        exp_x = torch.exp(x - torch.max(x, dim=-1, keepdim=True)[0])
        exp_x_sum = 1.0 / (exp_x.sum(dim=-1, keepdim=True) + self.eps)

        # create coordinates grid
        pos_y, pos_x = create_meshgrid(input, self.normalized_coordinates)
        pos_x = pos_x.reshape(-1)
        pos_y = pos_y.reshape(-1)

        # compute the expected coordinates
        expected_y: torch.Tensor = torch.sum(
            (pos_y * exp_x) * exp_x_sum, dim=-1, keepdim=True
        )
        expected_x: torch.Tensor = torch.sum(
            (pos_x * exp_x) * exp_x_sum, dim=-1, keepdim=True
        )
        output: torch.Tensor = torch.cat([expected_x, expected_y], dim=-1)
        return output.view(batch_size, channels, 2)  # BxNx2


######################
# functional interface
######################


def spatial_soft_argmax2d(
    input: torch.Tensor, normalized_coordinates: Optional[bool] = True
) -> torch.Tensor:
    r"""Function that computes the Spatial Soft-Argmax 2D of a given heatmap.

    See :class:`torchgeometry.contrib.SpatialSoftArgmax2d` for details.
    """
    return SpatialSoftArgmax2d(normalized_coordinates)(input)


class SpatialSoftmax(torch.nn.Module):
    def __init__(self, height, width, channel, temperature=None, data_format="NCHW"):
        super(SpatialSoftmax, self).__init__()
        self.data_format = data_format
        self.height = height
        self.width = width
        self.channel = channel

        if temperature:
            self.temperature = nn.Parameter(torch.ones(1) * temperature)
        else:
            self.temperature = 1.0

        pos_x, pos_y = np.meshgrid(
            np.linspace(-1.0, 1.0, self.height), np.linspace(-1.0, 1.0, self.width)
        )
        pos_x = torch.from_numpy(pos_x.reshape(self.height * self.width)).float()
        pos_y = torch.from_numpy(pos_y.reshape(self.height * self.width)).float()
        self.register_buffer("pos_x", pos_x)
        self.register_buffer("pos_y", pos_y)

    def forward(self, feature):
        # Output:
        #   (N, C*2) x_0 y_0 ...
        if self.data_format == "NHWC":
            feature = (
                feature.transpose(1, 3)
                .tranpose(2, 3)
                .view(-1, self.height * self.width)
            )
        else:
            feature = feature.view(-1, self.height * self.width)

        softmax_attention = F.softmax(feature / self.temperature, dim=-1)
        expected_x = torch.sum(self.pos_x * softmax_attention, dim=1, keepdim=True)
        expected_y = torch.sum(self.pos_y * softmax_attention, dim=1, keepdim=True)
        expected_xy = torch.cat([expected_x, expected_y], 1)
        feature_keypoints = expected_xy.view(-1, self.channel * 2)

        return feature_keypoints

arti_mani/algorithms/visual_net/Networks/utils.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import random
import time

import albumentations as albu
import cv2
import matplotlib.pyplot as plt
import numpy as np
import segmentation_models_pytorch as smp
import torch
from albumentations.core.transforms_interface import ImageOnlyTransform
from arti_mani import KPTDATA_DIR, SEGDATA_DIR

# from copy_paste import CopyPaste
from segmentation_models_pytorch import utils as smp_utils
from torch.utils.data import DataLoader
from torch.utils.data import Dataset as BaseDataset

IMAGE_HEIGHT = 144
IMAGE_WIDTH = 256


def seed_everything(seed):
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.set_num_threads(4)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = True


# helper function for data visualization
def visualize(**images):
    """PLot images in one row."""
    plt.figure(figsize=(16, 16))
    for i, (name, image) in enumerate(images.items()):
        plt.subplot(4, 4, i + 1)
        plt.xticks([])
        plt.yticks([])
        plt.title(" ".join(name.split("_")).title())
        plt.imshow(image)
    plt.show()


def vis_keypoints(rgb, depth, keypoints, color=(255, 0, 0), diameter=2):
    plot_rgb = rgb.copy()

    for x, y in keypoints[:, :2]:
        cv2.circle(plot_rgb, (round(x), round(y)), diameter, color, -1)

    plt.figure(figsize=(8, 8))
    plt.axis("off")
    plt.subplot(1, 3, 1)
    plt.imshow(rgb)
    plt.subplot(1, 3, 2)
    plt.imshow(depth)
    plt.subplot(1, 3, 3)
    plt.imshow(plot_rgb)
    plt.show()


def spatial_augmentation(copy_paste=False):
    transforms = [
        albu.HorizontalFlip(p=0.5),
        albu.VerticalFlip(p=0.5),
        albu.RandomSizedCrop(
            (int(IMAGE_HEIGHT * 0.85), IMAGE_HEIGHT),
            height=IMAGE_HEIGHT,
            width=IMAGE_WIDTH,
            w2h_ratio=IMAGE_WIDTH / IMAGE_HEIGHT,
            p=1,
        ),
    ]
    # if copy_paste:
    #     transforms = [CopyPaste(blend=False, p=0.2)] + transforms
    return albu.Compose(transforms)


def rgb_pixel_augmentation():
    transforms = [
        albu.GaussNoise(var_limit=(10, 50), p=0.5),
        albu.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2, p=0.5),
    ]
    return albu.Compose(transforms)


def depth_pixel_augmention():
    train_transform = [
        albu.GaussNoise(var_limit=0.03, p=0.5),
        AddSaltPepperNoise(density=0.05, p=0.5),
    ]
    return albu.Compose(train_transform)


class AddSaltPepperNoise(ImageOnlyTransform):
    def __init__(self, density=0.0, always_apply=False, p=0.5):
        super(AddSaltPepperNoise, self).__init__(always_apply, p)
        self.density = density
        self.p = p

    def apply(self, img, mask=None, **params):
        if random.uniform(0, 1) < self.p:
            img[mask == 0] = 0  # 椒
            img[mask == 1] = 2  # 盐
            return img
        else:
            return img

    def get_params_dependent_on_targets(self, params):
        image = params["image"]
        Nd = self.density
        Sd = 1 - Nd
        h, w, c = image.shape
        mask = np.random.choice(
            (0, 1, 2), size=(h, w, 1), p=[Nd / 2.0, Nd / 2.0, Sd]
        )  # h,w,q
        mask = np.repeat(mask, c, axis=2)  # h,w,c
        return {"mask": mask}

    @property
    def targets_as_params(self):
        return ["image"]

    def get_transform_init_args_names(self):
        return "density"


def get_validation_augmentation():
    """Add paddings to make image shape divisible by 32."""
    test_transform = [
        albu.PadIfNeeded(IMAGE_HEIGHT, IMAGE_WIDTH)
        # albu.PadIfNeeded(80, 128)
    ]
    return albu.Compose(test_transform)


def to_tensor(x, **kwargs):
    return x.transpose(2, 0, 1).astype("float32")


def get_preprocessing(preprocessing_fn):
    """Construct preprocessing transform.

    Args:
        preprocessing_fn (callbale): data normalization function
            (can be specific for each pretrained neural network)
    Return:
        transform: albumentations.Compose
    """

    _transform = [
        albu.Lambda(image=preprocessing_fn),
        albu.Lambda(image=to_tensor, mask=to_tensor),
    ]
    return albu.Compose(_transform)


class SegDataset(BaseDataset):
    """Read images, apply augmentation and preprocessing transformations.

    Args:
        data_dir (str): path to RGBD data and seg masks
        class_values (list): values of classes to extract from segmentation mask
        augmentation (bool): data transfromation pipeline
            (e.g. flip, scale, etc.)
        normalization (bool): data preprocessing
            (e.g. normalization, shape manipulation, etc.)
    """

    def __init__(
        self,
        data_mode,
        classes=None,
        augmentation=False,
        copy_paste=False,
        normalization=True,
        mode="RGBD",
    ):
        self.ids = os.listdir(SEGDATA_DIR / data_mode)
        self.rgbd_seg = {"rgb": [], "depth": [], "seg": []}
        for obj_id in sorted(self.ids):
            with np.load(SEGDATA_DIR / data_mode / obj_id) as obj:
                print(obj_id)
                for key in self.rgbd_seg.keys():
                    self.rgbd_seg[key].append(obj[key])
        for k in self.rgbd_seg:
            self.rgbd_seg[k] = np.vstack(self.rgbd_seg[k])
        # for key in self.rgbd_seg.keys():
        #     self.rgbd_seg[key] = np.array(self.rgbd_seg[key])

        # rgb: (N, 3, H, W), depth: (N, H, W), seg: (N, H, W)

        # convert str names to class values on masks
        self.class_values = list(range(len(classes)))

        self.augmentation = augmentation
        self.normalization = normalization
        self.data_mode = data_mode
        self.mode = mode
        self.copy_paste = copy_paste

        self.sp_aug = spatial_augmentation(copy_paste=self.copy_paste)
        self.rgb_px_aug = rgb_pixel_augmentation()
        self.d_px_aug = depth_pixel_augmention()

    def __getitem__(self, i):
        """
        self.data: rgb, stereo_depth, clean_depth, seg, bg_seg
        """
        # read data
        ori_image_rgbd = np.concatenate(
            (self.rgbd_seg["rgb"][i], self.rgbd_seg["depth"][i][None]), axis=0
        ).transpose(
            1, 2, 0
        )  # (H, W, 4)
        mask = self.rgbd_seg["seg"][i]  # (H, W)

        # extract certain classes from mask (e.g. cars)
        ori_mask = [(mask == v).astype(np.float32) for v in self.class_values]

        aug_image, aug_mask = None, None

        # apply augmentations
        if self.augmentation and self.data_mode == "train":
            if self.copy_paste is True:
                # paste faucet on door/drawer, vice versa
                is_faucet = (ori_mask[3].sum() + ori_mask[4].sum()) > 0
                paste_is_faucet = is_faucet
                while is_faucet == paste_is_faucet:
                    paste_i = random.randint(0, len(self.rgbd_seg["rgb"]) - 1)
                    paste_masks = [
                        (self.rgbd_seg["seg"][paste_i] == v).astype(np.float32)
                        for v in self.class_values
                    ]  # (H, W, 6)
                    paste_is_faucet = (paste_masks[3].sum() + paste_masks[4].sum()) > 0
                paste_rgbd = np.concatenate(
                    (
                        self.rgbd_seg["rgb"][paste_i],
                        self.rgbd_seg["depth"][paste_i][None],
                    ),
                    axis=0,
                ).transpose(
                    1, 2, 0
                )  # (H, W, 4)
                sample = self.sp_aug(
                    image=ori_image_rgbd,
                    masks=ori_mask,
                    paste_image=paste_rgbd,
                    paste_masks=paste_masks,
                    paste_bboxes=None,
                )  # (H, W, C)
            else:
                sample = self.sp_aug(image=ori_image_rgbd, masks=ori_mask)  # (H, W, C)
            aug_mask = sample["masks"]  # [(H,W)]*6
            if self.mode == "RGBD":
                sample_rgb = self.rgb_px_aug(
                    image=sample["image"][..., :3].astype(np.uint8)
                )  # (H, W, 3)
                sample_d = self.d_px_aug(
                    image=sample["image"][..., 3][..., None]
                )  # (H, W, 1)
                aug_image = np.concatenate(
                    [
                        sample_rgb["image"].transpose(2, 0, 1),
                        sample_d["image"].transpose(2, 0, 1),
                    ],
                    axis=0,
                )  # (4, H, W)
            elif self.mode == "RGD":
                sample_rgb = self.rgb_px_aug(
                    image=sample["image"][..., :3].astype(np.uint8)
                )  # (H, W, 3)
                sample_d = self.d_px_aug(
                    image=sample["image"][..., 3][..., None]
                )  # (H, W, 1)
                aug_image = np.concatenate(
                    [
                        sample_rgb["image"].transpose(2, 0, 1)[:2],
                        sample_d["image"].transpose(2, 0, 1),
                    ],
                    axis=0,
                )  # (3, H, W)
            elif self.mode == "RGB":
                sample_rgb = self.rgb_px_aug(
                    image=sample["image"][..., :3].astype(np.uint8)
                )  # (H, W, 3)
                aug_image = sample_rgb["image"].transpose(2, 0, 1)  # (3, H, W)
            elif self.mode == "D":
                sample_d = self.d_px_aug(
                    image=sample["image"][..., 3][..., None]
                )  # (H, W, 1)
                aug_image = sample_d["image"].transpose(2, 0, 1)  # (1, H, W)
            else:
                raise NotImplementedError
        else:
            if self.mode == "RGBD":
                aug_image, aug_mask = ori_image_rgbd.transpose((2, 0, 1)), ori_mask
            elif self.mode == "RGD":
                aug_image = np.concatenate(
                    (self.rgbd_seg["rgb"][i][:2], self.rgbd_seg["depth"][i][None]),
                    axis=0,
                )  # (3, H, W)
                aug_mask = ori_mask
            elif self.mode == "RGB":
                aug_image, aug_mask = ori_image_rgbd.transpose((2, 0, 1))[:3], ori_mask
            elif self.mode == "D":
                aug_image, aug_mask = ori_image_rgbd.transpose((2, 0, 1))[3:4], ori_mask
            else:
                raise NotImplementedError
        # apply normalization
        if self.normalization:
            if self.mode == "RGBD":
                aug_image[:3] = aug_image[:3] / 255.0
                # aug_image[3] = aug_image[3]/np.max(aug_image[3])
            elif self.mode == "RGD":
                aug_image[:2] = aug_image[:2] / 255.0
            elif self.mode == "RGB":
                aug_image = aug_image / 255.0
            elif self.mode == "D":
                pass
                # aug_image[3] = aug_image[3]/np.max(aug_image[3])
            else:
                raise NotImplementedError

        return ori_image_rgbd, ori_mask, aug_image, aug_mask

    def __len__(self):
        return len(self.rgbd_seg["seg"])


class KeypointDataset(BaseDataset):
    def __init__(
        self,
        data_mode,
        augmentation=False,
        normalization=True,
        mode="RGBD",
    ):
        self.ids = sorted(
            [file[:-4] for file in os.listdir(KPTDATA_DIR / f"{data_mode}/")]
        )

        rgb_data, depth_data, kpts_data = [], [], []
        for obj_id in self.ids:
            rgbd = np.load(KPTDATA_DIR / f"{data_mode}/{obj_id}.npz")
            rgb_data.extend(rgbd["rgb"])
            depth_data.extend(rgbd["depth"][:, None])
            kpts = np.load(
                KPTDATA_DIR / f"{data_mode}_kpts/{obj_id}_keypoints.npy"
            )  # (N, 3, 3, 3)
            kpts_data.extend(kpts)
        self.rgbd_data = np.concatenate((rgb_data, depth_data), axis=1)  # (N, 4, H, W)
        self.kpts_data = np.array(kpts_data)  # (N, J, 3, 3)

        self.augmentation = augmentation
        self.normalization = normalization
        self.mode = mode

    def __getitem__(self, i):
        """
        self.data: rgb, stereo_depth, clean_depth, seg, bg_seg
        """
        # read data
        ori_image_rgbd = self.rgbd_data[i].transpose(1, 2, 0)  # (H, W, 4)
        ori_kpts = self.kpts_data[i, 1]  # (J, 3), uvz
        # ori_kpts = self.kpts_data[i, 0]  # (J, 3), cam_xyz
        ori_kpt_visable = self.kpts_data[i, 2, :, 2]  # (J), uvz visable

        aug_image, aug_kpts, aug_kpt_visable = None, ori_kpts, ori_kpt_visable

        # apply augmentations
        if self.augmentation:
            sample = spatial_augmentation()(image=ori_image_rgbd)  # (H, W, C)
            if self.mode == "RGBD":
                sample_rgb = rgb_pixel_augmentation()(
                    image=sample["image"][..., :3].astype(np.uint8)
                )  # (H, W, 3)
                sample_d = depth_pixel_augmention()(
                    image=sample["image"][..., 3][..., None]
                )  # (H, W, 1)
                aug_image = np.concatenate(
                    [
                        sample_rgb["image"].transpose(2, 0, 1),
                        sample_d["image"].transpose(2, 0, 1),
                    ],
                    axis=0,
                )  # (4, H, W)
            elif self.mode == "RGB":
                sample_rgb = rgb_pixel_augmentation()(
                    image=sample["image"][..., :3].astype(np.uint8)
                )  # (H, W, 3)
                aug_image = sample_rgb["image"].transpose(2, 0, 1)  # (3, H, W)
            elif self.mode == "D":
                sample_d = depth_pixel_augmention()(
                    image=sample["image"][..., 3][..., None]
                )  # (H, W, 1)
                aug_image = sample_d["image"].transpose(2, 0, 1)  # (1, H, W)
            else:
                raise NotImplementedError
        else:
            if self.mode == "RGBD":
                aug_image = ori_image_rgbd.transpose(2, 0, 1)
            elif self.mode == "RGB":
                aug_image = ori_image_rgbd.transpose(2, 0, 1)[:3]
            elif self.mode == "D":
                aug_image = ori_image_rgbd.transpose(2, 0, 1)[3]
            else:
                raise NotImplementedError
        # apply normalization
        if self.normalization:
            ## uvz min, max: (11, 0, 0.185), (242, 143, 0.88)
            kpts_min = np.array([0, 0, 0.18])
            kpts_max = np.array([255, 143, 1.0])
            ## camxyz min, max: (-0.25, -0.19, 0.18), (0.29, 0.30, 0.88)
            # kpts_min = np.array([-0.3, -0.2, 0.18])
            # kpts_max = np.array([0.3, 0.3, 1.0])
            aug_kpts = (aug_kpts - kpts_min) / (kpts_max - kpts_min)
            if self.mode == "RGBD":
                aug_image[:3] = aug_image[:3] / 255.0
                # aug_image[3] = aug_image[3]/np.max(aug_image[3])
            elif self.mode == "RGB":
                aug_image = aug_image / 255.0
            elif self.mode == "D":
                pass
                # aug_image[3] = aug_image[3]/np.max(aug_image[3])
            else:
                raise NotImplementedError

        return ori_image_rgbd, ori_kpts, aug_image, aug_kpts, aug_kpt_visable

    def __len__(self):
        return self.kpts_data.shape[0]


if __name__ == "__main__":
    MODE = "RGBD"  # RGBD, RGB, D
    CLASSES = ["handle", "door", "cabinet", "switchlink", "fixlink", "other"]

    # preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)

    seg_dataset = SegDataset(
        data_mode="train",
        augmentation=True,
        copy_paste=True,
        classes=CLASSES,
        mode=MODE,
    )

    cnt = 0
    while cnt < 5:
        ind = np.random.choice(range(len(seg_dataset)))
        image_rgbd, mask, aug_image_rgbd, aug_mask = seg_dataset[ind]  # get some sample
        ori_bin_masks = {"o-" + CLASSES[i]: mask[i] for i in range(len(CLASSES))}
        aug_bin_masks = {"a-" + CLASSES[i]: aug_mask[i] for i in range(len(CLASSES))}
        visualize(
            ori_rgb=image_rgbd[:, :, :3] / 255.0,
            ori_depth=image_rgbd[:, :, 3],
            aug_rgb=aug_image_rgbd[:3].transpose(1, 2, 0),
            aug_depth=aug_image_rgbd[3],
            **ori_bin_masks,
            **aug_bin_masks,
        )
        time.sleep(0.5)
        cnt += 1

    # kpts_dataset = KeypointDataset(
    #     data_mode="train",
    #     augmentation=True,
    #     mode=MODE
    # )
    #
    # cnt = 0
    # while cnt < 5:
    #     ind = np.random.choice(range(len(kpts_dataset)))
    #     image_rgbd, kpts, aug_image_rgbd, aug_kpts, aug_kpt_visable = kpts_dataset[ind]  # get some sample
    #     vis_keypoints(aug_image_rgbd[:3].transpose((1, 2, 0)), aug_image_rgbd[3], aug_kpts)
    #     cnt += 1

arti_mani/algorithms/visual_net/scripts/dataset.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os

import cv2
import gym
import matplotlib.pyplot as plt
from arti_mani import VISUALMODEL_DIR
from arti_mani.envs import *


def main():
    np.set_printoptions(suppress=True, precision=4)
    log_name = "20230307_182821_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2_dropout0.2_newdatafilter2drawer"
    vis_result_path = (
        VISUALMODEL_DIR / f"kpt_model/{log_name}/visual_results/test_kpts_new/"
    )
    if not os.path.exists(vis_result_path):
        os.makedirs(vis_result_path)

    visualization = False

    uvz_errors = []
    for mode in ["door", "drawer", "faucet"]:
        if mode == "door":
            arti_ids = [0, 1006, 1030, 1047, 1081]
        elif mode == "drawer":
            arti_ids = [1, 1004, 1005, 1016, 1024]
        elif mode == "faucet":
            arti_ids = [5052, 5004, 5007, 5023, 5069]
        else:
            raise NotImplementedError(mode)
        kpts_min = np.array([0, 0, 0.18])
        kpts_max = np.array([255, 143, 1.0])
        plt_x, plt_y = 2, 2

        # init env
        env = gym.make(
            "ArtiMani-v0",
            articulation_ids=arti_ids,
            segmodel_path=VISUALMODEL_DIR / f"kpt_model/{log_name}",
            device="cuda:0",
            obs_mode="state_egostereo_keypoints",
            control_mode="pd_joint_delta_pos",
        )
        uvz_error = []

        for arti_ind in arti_ids:
            _ = env.reset(articulation_id=arti_ind)
            obs_global = env.get_obs()
            rgbd = env.get_handsensor_rgbdseg()
            print(f"visualize result of {mode}-{arti_ind}")
            # viewer = env.render()
            # print("Press [e] to start")
            # while True:
            #     if viewer.window.key_down("e"):
            #         break
            #     env.render()

            uvz = obs_global["uvz"]
            uvz_pred = obs_global["uvz_pred"]
            uvz_visable = obs_global["uvz_visable"]
            uvz_norm_tensor = torch.from_numpy((uvz - kpts_min) / (kpts_max - kpts_min))
            uvz_pred_norm_tensor = torch.from_numpy(
                (uvz_pred - kpts_min) / (kpts_max - kpts_min)
            )
            loss = F.l1_loss(uvz_pred_norm_tensor, uvz_norm_tensor, reduction="mean")

            print("loss: ", loss.item())
            uvz_error.append(np.abs(uvz_pred - uvz))

            rgb = rgbd["rgb"]  # (H, W, 3)
            depth = rgbd["depth"].squeeze()  # (H, W)

            rgb_kptgt = rgb.copy()
            for idn in range(uvz.shape[0]):
                if uvz_visable[idn, 0] and uvz_visable[idn, 1]:
                    if uvz_visable[idn, 2]:
                        cv2.circle(
                            rgb_kptgt,
                            (round(uvz[idn, 0]), round(uvz[idn, 1])),
                            radius=2,
                            color=(1, 0, 0),
                            thickness=-1,
                        )
                    else:
                        cv2.circle(
                            rgb_kptgt,
                            (round(uvz[idn, 0]), round(uvz[idn, 1])),
                            radius=2,
                            color=(0, 1, 0),
                            thickness=-1,
                        )

            rgb_kptpred = rgb.copy()
            for idn in range(uvz_pred.shape[0]):
                if uvz_visable[idn, 0] and uvz_visable[idn, 1]:
                    if uvz_visable[idn, 2]:
                        cv2.circle(
                            rgb_kptpred,
                            (round(uvz_pred[idn, 0]), round(uvz_pred[idn, 1])),
                            radius=2,
                            color=(1, 0, 0),
                            thickness=-1,
                        )
                    else:
                        cv2.circle(
                            rgb_kptpred,
                            (round(uvz_pred[idn, 0]), round(uvz_pred[idn, 1])),
                            radius=2,
                            color=(0, 1, 0),
                            thickness=-1,
                        )

            plt.subplot(plt_x, plt_y, 1)
            plt.imshow(rgb)
            plt.subplot(plt_x, plt_y, 2)
            plt.imshow(depth)
            plt.subplot(plt_x, plt_y, 3)
            plt.imshow(rgb_kptgt)
            plt.subplot(plt_x, plt_y, 4)
            plt.imshow(rgb_kptpred)
            plt.tight_layout()
            plt.savefig(vis_result_path / f"{mode}_{arti_ind}.jpg")
            if visualization:
                plt.show()
            plt.close()
        mean_arti_uvzerr = np.array(uvz_error).mean(0)
        print(f"{mode} mean uvz error: ", mean_arti_uvzerr)
        uvz_errors.append(mean_arti_uvzerr)
    print("mean uvz erros: ", np.array(uvz_errors).mean(0))


if __name__ == "__main__":
    main()

arti_mani/algorithms/visual_net/scripts/test_kpts_new.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
from arti_mani import KPTDATA_DIR, VISUALMODEL_DIR
from arti_mani.algorithms.visual_net.Networks.keypoint_detection import (
    IntegralHumanPoseModel,
    soft_argmax,
)
from torch.nn import functional as F


def load_kptmodel(log_name, device="cuda:0"):
    vis_result_path = (
        VISUALMODEL_DIR / f"kpt_model/{log_name}/visual_results/test_kpts_val/"
    )
    if not os.path.exists(vis_result_path):
        os.makedirs(vis_result_path)
    # load model
    model_path = VISUALMODEL_DIR / f"kpt_model/{log_name}/best.pth"
    kpt_model = IntegralHumanPoseModel(
        num_keypoints=3, num_deconv_layers=3, depth_dim=64, has_dropout=False
    )
    kpt_model.load_state_dict(torch.load(model_path))
    kpt_model.to(torch.device(device))
    kpt_model.eval()
    model_params = sum(p.numel() for p in kpt_model.parameters())
    print(f"model params: {model_params / 1e6:.2f} M")
    return kpt_model


def load_data(arti_ids):
    rgbd_data = []
    kpts_data = []
    train_ids = [int(name[:-14]) for name in os.listdir(KPTDATA_DIR / "train_kpts")]
    val_ids = [int(name[:-14]) for name in os.listdir(KPTDATA_DIR / "val_kpts")]
    for val_id in arti_ids:
        if val_id in val_ids:
            rgbd = np.load(KPTDATA_DIR / f"val/{val_id}.npz")
            kpts = np.load(KPTDATA_DIR / f"val_kpts/{val_id}_keypoints.npy")
        elif val_id in train_ids:
            rgbd = np.load(KPTDATA_DIR / f"train/{val_id}.npz")
            kpts = np.load(KPTDATA_DIR / f"train_kpts/{val_id}_keypoints.npy")
        rgbd_data.extend(
            np.concatenate((rgbd["rgb"] / 255.0, rgbd["depth"][:, None]), axis=1)
        )  # SN, 4, H, W
        kpts_data.extend(kpts)  # SN, 3, 3, 3
    rgbd_data = np.array(rgbd_data)  # SN, 4, H, W
    kpts_data = np.array(kpts_data)  # SN, 3, 3, 3
    return rgbd_data, kpts_data


def vis_kpt(rgb, depth, ind, save_dir):
    # save qualified results
    rgb_kptgt = rgb.copy()
    for idn in range(joint_num):
        if uvz_visable[idn, 0] and uvz_visable[idn, 1]:
            if uvz_visable[idn, 2]:
                cv2.circle(
                    rgb_kptgt,
                    (round(uvz[idn, 0]), round(uvz[idn, 1])),
                    radius=2,
                    color=(1, 0, 0),
                    thickness=-1,
                )
            else:
                cv2.circle(
                    rgb_kptgt,
                    (round(uvz[idn, 0]), round(uvz[idn, 1])),
                    radius=2,
                    color=(0, 1, 0),
                    thickness=-1,
                )

    rgb_kptpred = rgb.copy()
    uvz_pred_norm = uvz_pred_norm_tensor.detach().cpu().numpy()
    uvz_pred = uvz_pred_norm * (kpts_max - kpts_min) + kpts_min
    for idn in range(joint_num):
        if uvz_visable[idn, 0] and uvz_visable[idn, 1]:
            if uvz_visable[idn, 2]:
                cv2.circle(
                    rgb_kptpred,
                    (round(uvz_pred[idn, 0]), round(uvz_pred[idn, 1])),
                    radius=2,
                    color=(1, 0, 0),
                    thickness=-1,
                )
            else:
                cv2.circle(
                    rgb_kptpred,
                    (round(uvz_pred[idn, 0]), round(uvz_pred[idn, 1])),
                    radius=2,
                    color=(0, 1, 0),
                    thickness=-1,
                )
    plt.subplot(plt_x, plt_y, 1)
    plt.imshow(rgb)
    plt.subplot(plt_x, plt_y, 2)
    plt.imshow(depth)
    plt.subplot(plt_x, plt_y, 3)
    plt.imshow(rgb_kptgt)
    plt.subplot(plt_x, plt_y, 4)
    plt.imshow(rgb_kptpred)
    plt.tight_layout()
    plt.savefig(save_dir / f"val_{ind}.jpg")
    plt.close()


def compute_err(uvz_err, xyz, fx, fy):
    uvz_mean_err = uvz_err.mean(0)
    print("mean uvz error: ", uvz_mean_err)
    xy_mean_err = uvz_mean_err[:2] * xyz[:, :, 2].mean() / np.array([fx, fy])
    xyz_mean_err = np.append(xy_mean_err, uvz_mean_err[2])
    print("mean xyz error: ", xyz_mean_err)
    print("mean kpt error: ", xyz_mean_err.mean())


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = torch.device("cuda:0")
    joint_num = 3
    batch_size = 256
    IM_H, IM_W = 144, 256
    out_shape = (64, 40, 64)
    kpts_min = np.array([0, 0, 0.18])
    kpts_max = np.array([255, 143, 1.0])
    fx, fy = 183.406, 183.491
    cx, cy = 126.608, 73.768
    plt_x, plt_y = 2, 2

    # load model
    log_name = "20230307_182821_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2_dropout0.2_newdatafilter2drawer"
    kpt_model = load_kptmodel(log_name)

    # load sim data and preprocess
    val_list = [
        0,
        1006,
        1030,
        1047,
        1081,
        1,
        1005,
        1016,
        1024,
        1076,
        5004,
        5007,
        5023,
        5069,
        5052,
    ]
    # rgbd: (N, 4, H, W), kpts: (N, 3, 3, 3)
    rgbd_data, kpts_data = load_data(val_list)
    data_num = kpts_data.shape[0]
    xyz = kpts_data[:, 0]  # N, 3, 3
    uvz = kpts_data[:, 1]  # N, 3, 3
    uvz_norm = (uvz - kpts_min) / (kpts_max - kpts_min)  # N, 3, 3
    uvz_visable = kpts_data[:, 2]  # N, 3, 3

    ids_slice = np.append(np.arange(data_num)[::batch_size], data_num)
    loss, uvz_error_all, uvz_error_occlu, uvz_error_noocclu = [], [], [], []
    with torch.no_grad():
        for id_start, id_end in zip(ids_slice[:-1], ids_slice[1:]):
            uvz_gt, xyz_gt = uvz[id_start:id_end], xyz[id_start:id_end]
            uvz_visable_gt = uvz_visable[id_start:id_end]  # (bs, 3, 3)
            sim_rgbd = (
                torch.from_numpy(rgbd_data[id_start:id_end]).float().to(device)
            )  # (bs, ch, H, W)
            uvz_norm_tensor = torch.from_numpy(uvz_norm[id_start:id_end]).to(device)
            # pred
            pred_heatmap = kpt_model.forward(sim_rgbd)  # (bs, 3*D, H', W')
            uvz_pred_norm_tensor = soft_argmax(
                pred_heatmap, joint_num, out_shape
            )  # (bs, 3, 3)
            loss.append(
                F.l1_loss(
                    uvz_pred_norm_tensor, uvz_norm_tensor, reduction="mean"
                ).item()
            )

            uvz_pred_norm = uvz_pred_norm_tensor.detach().cpu().numpy()  # (bs, 3, 3)
            uvz_pred = uvz_pred_norm * (kpts_max - kpts_min) + kpts_min  # (bs, 3, 3)
            uvz_error_all.append(np.abs(uvz_pred - uvz_gt))  # (bs, 3, 3) total
            uvz_error_noocclu.append(
                np.abs(uvz_pred - uvz_gt) * uvz_visable_gt
            )  # (bs, 3, 3) w/o occlusion
            uvz_error_occlu.append(
                np.abs(uvz_pred - uvz_gt) * (1 - uvz_visable_gt)
            )  # (bs, 3, 3) only occlusion

    print("loss: ", np.mean(loss))
    total_nums = uvz_visable.shape[0]
    visable_nums = uvz_visable.sum(0)  # (3,3)
    invisable_nums = total_nums - uvz_visable.sum(0)
    print(invisable_nums)
    mean_uvz_err = np.concatenate(np.array(uvz_error_all), axis=0).sum(0) / total_nums
    mean_uvz_noocclu_err = (
        np.concatenate(np.array(uvz_error_noocclu), axis=0).sum(0) / visable_nums
    )
    mean_uvz_occlu_err = (
        np.concatenate(np.array(uvz_error_occlu), axis=0).sum(0) / invisable_nums
    )
    compute_err(mean_uvz_err, xyz, fx, fy)
    compute_err(mean_uvz_noocclu_err, xyz, fx, fy)
    compute_err(mean_uvz_occlu_err, xyz, fx, fy)

arti_mani/algorithms/visual_net/scripts/test_kpts_simenv.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import argparse
import logging
import os
import time

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
from arti_mani import KPTDATA_DIR, VISUALMODEL_DIR
from arti_mani.algorithms.visual_net.Networks.keypoint_detection import (
    IntegralHumanPoseModel,
    soft_argmax,
)
from torch.nn import functional as F


def main(sample_num, logger, log_name, device):
    vis_result_path = (
        VISUALMODEL_DIR / f"kpt_model/{log_name}/visual_results/test_kpts_val/"
    )
    if not os.path.exists(vis_result_path):
        os.makedirs(vis_result_path)
    # load model
    model_path = VISUALMODEL_DIR / f"kpt_model/{log_name}/best.pth"
    kpt_model = IntegralHumanPoseModel(
        num_keypoints=3, num_deconv_layers=3, depth_dim=64, has_dropout=False
    )
    kpt_model.load_state_dict(torch.load(model_path))
    kpt_model.to(torch.device(device))
    kpt_model.eval()
    model_params = sum(p.numel() for p in kpt_model.parameters())
    logger.info(f"model params: {model_params / 1e6:.2f} M")
    IM_H, IM_W = 144, 256
    out_shape = (64, 40, 64)
    kpts_min = np.array([0, 0, 0.18])
    kpts_max = np.array([255, 143, 1.0])
    plt_x, plt_y = 2, 2

    # load sim data and preprocess
    rgbd_data = []
    kpts_data = []
    val_list = [
        1002,
        1005,
        1007,
        1017,
        1021,
        1027,
        1030,
        1033,
        1036,
        1045,
        1054,
        1063,
        1067,
        1075,
        1079,
        1081,
        5007,
        5028,
        5069,
    ]  # 1056,
    num_perinst = sample_num // len(val_list)
    for val_id in val_list:
        rgbd = np.load(KPTDATA_DIR / f"val/{val_id}.npz")
        kpts = np.load(KPTDATA_DIR / f"val_kpts/{val_id}_keypoints.npy")
        indices = np.random.choice(kpts.shape[0], num_perinst)
        rgbd_data.extend(
            np.concatenate(
                (rgbd["rgb"][indices] / 255.0, rgbd["depth"][indices][:, None]), axis=1
            )
        )  # SN, 4, H, W
        kpts_data.extend(kpts[indices])  # SN, 3, 3, 3
    rgbd_data = np.array(rgbd_data)  # SN, 4, H, W
    kpts_data = np.array(kpts_data)  # SN, 3, 3, 3
    uvz = kpts_data[:, 1]  # SN, 3, 3
    uvz_norm = (uvz - kpts_min) / (kpts_max - kpts_min)
    uvz_visable = kpts_data[:, 2]  # SN, 3, 3

    # transfer to torch data
    sim_rgbd = torch.from_numpy(rgbd_data).float().to(device)  # (SN, ch, H, W)
    uvz_norm_tensor = torch.from_numpy(uvz_norm).to(device)

    pred_heatmap = kpt_model.forward(sim_rgbd)  # (SN, J*D, H', W')
    joint_num = uvz.shape[1]

    uvz_pred_norm_tensor = soft_argmax(pred_heatmap, joint_num, out_shape)  # (SN, J, 3)

    loss = F.l1_loss(uvz_pred_norm_tensor, uvz_norm_tensor, reduction="mean")
    print("loss: ", loss.item())

    uvz_error = []
    for ind in range(kpts_data.shape[0]):
        rgb = rgbd_data[ind, :3].transpose(1, 2, 0)
        depth = rgbd_data[ind, 3]
        # save qualified results
        rgb_kptgt = rgb.copy()
        for idn in range(joint_num):
            if uvz_visable[ind, idn, 0] and uvz_visable[ind, idn, 1]:
                if uvz_visable[ind, idn, 2]:
                    cv2.circle(
                        rgb_kptgt,
                        (round(uvz[ind, idn, 0]), round(uvz[ind, idn, 1])),
                        radius=2,
                        color=(1, 0, 0),
                        thickness=-1,
                    )
                else:
                    cv2.circle(
                        rgb_kptgt,
                        (round(uvz[ind, idn, 0]), round(uvz[ind, idn, 1])),
                        radius=2,
                        color=(0, 1, 0),
                        thickness=-1,
                    )

        rgb_kptpred = rgb.copy()
        uvz_pred_norm = uvz_pred_norm_tensor.detach().cpu().numpy()
        uvz_pred = uvz_pred_norm * (kpts_max - kpts_min) + kpts_min
        for idn in range(joint_num):
            if uvz_visable[ind, idn, 0] and uvz_visable[ind, idn, 1]:
                if uvz_visable[ind, idn, 2]:
                    cv2.circle(
                        rgb_kptpred,
                        (round(uvz_pred[ind, idn, 0]), round(uvz_pred[ind, idn, 1])),
                        radius=2,
                        color=(1, 0, 0),
                        thickness=-1,
                    )
                else:
                    cv2.circle(
                        rgb_kptpred,
                        (round(uvz_pred[ind, idn, 0]), round(uvz_pred[ind, idn, 1])),
                        radius=2,
                        color=(0, 1, 0),
                        thickness=-1,
                    )
        uvz_error.append(np.abs(uvz_pred[ind] - uvz[ind]))
        plt.subplot(plt_x, plt_y, 1)
        plt.imshow(rgb)
        plt.subplot(plt_x, plt_y, 2)
        plt.imshow(depth)
        plt.subplot(plt_x, plt_y, 3)
        plt.imshow(rgb_kptgt)
        plt.subplot(plt_x, plt_y, 4)
        plt.imshow(rgb_kptpred)
        plt.tight_layout()
        plt.savefig(vis_result_path / f"val_{ind}.jpg")
        plt.close()
    mean_kpt_err = np.array(uvz_error).mean(0)
    print("mean uvz error: ", mean_kpt_err)


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = torch.device("cuda:0")
    test_num = 80
    parser = argparse.ArgumentParser(description=f"test for {test_num} sim images")
    # log_name = "20230305_173407_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3"
    # log_name = "20230305_181802_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_resnet18"
    # log_name = "20230305_185213_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2"
    # log_name = "20230305_193703_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2_dropout0.2"
    # log_name = "20230305_201857_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2_block8_dropout0.2"
    # log_name = "20230307_175848_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2_block8_dropout0.2_newdatafilter2drawer"
    log_name = "20230307_182821_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2_dropout0.2_newdatafilter2drawer"
    # log_name = "20230308_135416_kpts3_uvz_modeRGBDepoch60bs32lr0.001_mobilenetv2_dropout0.2_newdatafilter2drawer"
    # log_name = "20230308_145353_kpts3_uvz_modeRGBDepoch60bs32lr0.001_mobilenetv2_dropout0.2_newdatafilter2drawer_smoothl1"
    parser.add_argument("--log-name", "-ln", type=str, default=log_name)
    parser.add_argument("--sample-num", type=int, default=test_num)
    args = parser.parse_args()

    # generate log
    logfile_path = VISUALMODEL_DIR / f"kpt_model/{args.log_name}/simval_results/"
    if not os.path.exists(logfile_path):
        os.makedirs(logfile_path)
    logfile_name = f"{time.strftime('%Y%m%d_%H%M%S', time.localtime())}"
    logger = logging.getLogger(logfile_name)
    logger.setLevel(logging.DEBUG)
    log_file = "{}.log".format(logfile_name)
    fileHandler = logging.FileHandler(os.path.join(logfile_path, log_file), mode="w")
    fileHandler.setLevel(logging.DEBUG)
    formatter = logging.Formatter(
        "[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    fileHandler.setFormatter(formatter)
    logger.addHandler(fileHandler)

    main(
        sample_num=args.sample_num,
        logger=logger,
        log_name=args.log_name,
        device=device,
    )

arti_mani/algorithms/visual_net/scripts/test_kpts_val.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import argparse
import glob
import logging
import os
import time

import cv2
import matplotlib.pyplot as plt
import numpy as np
import segmentation_models_pytorch as smp
import torch
import yaml
from arti_mani import REAL_DIR, VISUALMODEL_DIR
from arti_mani.utils.cv_utils import visualize_depth, visualize_seg


def main(logger, log_name, logfile_path, device, classes, only_vis):
    # load model
    model_path = VISUALMODEL_DIR / f"{log_name}/best.pth"
    config_path = VISUALMODEL_DIR / f"{log_name}/config.yaml"
    with open(config_path, "r") as f:
        cfg = yaml.safe_load(f)
    smp_cfg = cfg["smp_config"]
    if smp_cfg["mode"] == "RGBD":
        in_ch = 4
    elif smp_cfg["mode"] == "RGB":
        in_ch = 3
    elif smp_cfg["mode"] == "D":
        in_ch = 1
    else:
        raise NotImplementedError
    seg_model = smp.Unet(
        encoder_name=smp_cfg["encoder"],
        encoder_depth=smp_cfg["encoder_depth"],
        decoder_channels=smp_cfg["decoder_channels"],
        encoder_weights=smp_cfg["encoder_weights"],
        in_channels=in_ch,
        classes=cfg["num_classes"],
        activation=smp_cfg["activation"],
    )
    seg_model.load_state_dict(torch.load(model_path))
    seg_model.to(device)
    seg_model.eval()
    model_params = sum(p.numel() for p in seg_model.parameters())
    logger.info(f"model params: {model_params / 1e6:.2f} M")
    num_classes = cfg["num_classes"]
    IM_SIZE = (256, 144)
    if smp_cfg["encoder_depth"] == 5:
        PAD_SIZE = (256, 160)
    else:
        PAD_SIZE = (256, 144)

    # load real data and preprocess
    leftrightpad = int((PAD_SIZE[0] - IM_SIZE[0]) / 2)
    topbottompad = int((PAD_SIZE[1] - IM_SIZE[1]) / 2)
    real_img = {"rgb": [], "depth": [], "masks": []}
    for mode in ["door", "drawer", "faucet"]:
        for im_mode in ["rgb", "depth"]:
            files = glob.glob(str(REAL_DIR / f"{mode}/*{im_mode}.npy"))
            files.sort()
            for file in files:
                data = np.load(file)
                data = cv2.resize(
                    src=data, dsize=IM_SIZE, interpolation=cv2.INTER_NEAREST
                )
                if leftrightpad > 0 or topbottompad > 0:
                    data = cv2.copyMakeBorder(
                        data,
                        topbottompad,
                        topbottompad,
                        leftrightpad,
                        leftrightpad,
                        cv2.BORDER_DEFAULT,
                    )
                if im_mode == "rgb":
                    data /= 255.0
                    data = data[..., ::-1]
                real_img[im_mode].append(data)
    # load real masks
    if not only_vis:
        mask_files = []
        for mode in ["door", "drawer", "faucet"]:
            files = glob.glob(str(REAL_DIR / f"masks/*{mode}.npy"))
            files.sort()
            mask_files.extend(files)
        for mf in mask_files:
            data = np.load(mf)
            data = cv2.resize(src=data, dsize=IM_SIZE, interpolation=cv2.INTER_NEAREST)
            if leftrightpad > 0 or topbottompad > 0:
                data = cv2.copyMakeBorder(
                    data,
                    topbottompad,
                    topbottompad,
                    leftrightpad,
                    leftrightpad,
                    cv2.BORDER_DEFAULT,
                )
            real_img["masks"].append(data)
    for key in real_img.keys():
        real_img[key] = np.array(real_img[key])

    if in_ch == 1:
        real_data = np.expand_dims(real_img["depth"], 3)
    elif in_ch == 3:
        real_data = np.array(real_img["rgb"])
    elif in_ch == 4:
        real_data = np.concatenate(
            [real_img["rgb"], np.expand_dims(real_img["depth"], 3)], axis=3
        )
    else:
        raise NotImplementedError
    # transfer to torch data
    real_rgbd = (
        torch.from_numpy(real_data).float().permute(0, 3, 1, 2).to(device)
    )  # (N, ch, H, W)
    pred_seg = seg_model.forward(real_rgbd)  # (N, 6, H, W)
    pred_mc = torch.argmax(pred_seg, dim=1).to(device)  # (N, H, W)

    # save qualified results
    subplot_num = [3, 3]
    seg_pred_img = pred_seg.detach().cpu().numpy()
    for ind in range(real_rgbd.shape[0]):
        plt.subplot(subplot_num[0], subplot_num[1], 1)
        plt.imshow(real_img["rgb"][ind])
        plt.subplot(subplot_num[0], subplot_num[1], 2)
        plt.imshow(visualize_depth(real_img["depth"][ind])[..., ::-1])
        plt.subplot(subplot_num[0], subplot_num[1], 3)
        plt.imshow(visualize_seg(real_img["masks"][ind], num_classes))
        for pid in range(num_classes):
            plt.subplot(subplot_num[0], subplot_num[1], pid + 4)
            plt.imshow(seg_pred_img[ind][pid])

        plt.savefig(f"{logfile_path}/{ind:03}.jpg", dpi=200)

    # save quantified results
    if not only_vis:
        real_seg = torch.from_numpy(real_img["masks"]).to(device)  # (N, H, W)
        tp, fp, fn, tn = smp.metrics.get_stats(
            pred_mc.long(), real_seg.long(), mode="multiclass", num_classes=num_classes
        )
        iou_metric = smp.metrics.iou_score(tp, fp, fn, tn, reduction="none")  # (bs, C)
        f1_metric = smp.metrics.f1_score(tp, fp, fn, tn, reduction="none")  # (bs, C)
        logger.info(
            f"\n miou: {smp.metrics.iou_score(tp, fp, fn, tn, reduction='micro').item()}"
        )
        for ind, key in enumerate(classes):
            logger.info(f"iou_{key}: {torch.mean(iou_metric, dim=0)[ind].item()}")
        logger.info(
            f"\n mf1: {smp.metrics.f1_score(tp, fp, fn, tn, reduction='micro').item()}"
        )
        for ind, key in enumerate(classes):
            logger.info(f"f1_{key}: {torch.mean(f1_metric, dim=0)[ind].item()}")


if __name__ == "__main__":
    device = torch.device("cuda:0")
    classes_name = ["handle", "door", "cabinet", "switchlink", "fixlink", "other"]

    parser = argparse.ArgumentParser(description="test for 90 real images")
    parser.add_argument(
        "--log-name",
        "-ln",
        type=str,
        default="smp_model/20230216_151449_train50-val20_noDR_norandombg_aug_stereo_bs64_cel_0.5step50lr0.001_RGBDunet4-163264128_mobilenet_v2",
    )
    parser.add_argument(
        "--only-vis", action="store_true", help="only visualize results or not"
    )
    args = parser.parse_args()

    # generate log
    logfile_path = VISUALMODEL_DIR / f"{args.log_name}/real_results/"
    if not os.path.exists(logfile_path):
        os.makedirs(logfile_path)
    logfile_name = f"{time.strftime('%Y%m%d_%H%M%S', time.localtime())}_47arti"
    logger = logging.getLogger(logfile_name)
    logger.setLevel(logging.DEBUG)
    log_file = "{}.log".format(logfile_name)
    fileHandler = logging.FileHandler(os.path.join(logfile_path, log_file), mode="w")
    fileHandler.setLevel(logging.DEBUG)
    consoleHandler = logging.StreamHandler()
    consoleHandler.setLevel(logging.DEBUG)
    formatter = logging.Formatter(
        "[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    fileHandler.setFormatter(formatter)
    consoleHandler.setFormatter(formatter)
    logger.addHandler(fileHandler)
    logger.addHandler(consoleHandler)

    main(
        logger=logger,
        log_name=args.log_name,
        logfile_path=logfile_path,
        device=device,
        classes=classes_name,
        only_vis=args.only_vis,
    )

arti_mani/algorithms/visual_net/scripts/test_real.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import argparse
import glob
import json
import logging
import os
import time

import cv2
import matplotlib.pyplot as plt
import numpy as np
import segmentation_models_pytorch as smp
import torch
import yaml
from arti_mani import REAL_DIR, VISUALMODEL_DIR
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import CustomUnet, SplitUnet
from arti_mani.utils.cv_utils import visualize_depth, visualize_seg
from path import Path
from PIL import Image
from torch.nn import functional as F
from tqdm import tqdm


def load_real_data_old(
    REAL_DIR="/ssd/pwral_real_seg_old",
    PAD_SIZE=(256, 144),
    IM_SIZE=(256, 144),
    in_ch=4,
    only_vis=False,
):
    # load real data and preprocess
    REAL_DIR = Path(REAL_DIR)
    leftrightpad = int((PAD_SIZE[0] - IM_SIZE[0]) / 2)
    topbottompad = int((PAD_SIZE[1] - IM_SIZE[1]) / 2)
    real_img = {"rgb": [], "depth": [], "masks": []}
    for mode in ["door", "drawer", "faucet"]:
        for im_mode in ["rgb", "depth"]:
            files = glob.glob(str(REAL_DIR / f"{mode}/*{im_mode}.npy"))
            files.sort()
            for file in files:
                data = np.load(file)
                data = cv2.resize(
                    src=data, dsize=IM_SIZE, interpolation=cv2.INTER_NEAREST
                )
                if leftrightpad > 0 or topbottompad > 0:
                    data = cv2.copyMakeBorder(
                        data,
                        topbottompad,
                        topbottompad,
                        leftrightpad,
                        leftrightpad,
                        cv2.BORDER_DEFAULT,
                    )
                if im_mode == "rgb":
                    data = data[..., ::-1] / 255.0
                real_img[im_mode].append(data)
    # load real masks
    if not only_vis:
        mask_files = []
        for mode in ["door", "drawer", "faucet"]:
            files = glob.glob(str(REAL_DIR / f"masks/*{mode}.npy"))
            files.sort()
            mask_files.extend(files)
        for mf in mask_files:
            data = np.load(mf)
            data = cv2.resize(src=data, dsize=IM_SIZE, interpolation=cv2.INTER_NEAREST)
            if leftrightpad > 0 or topbottompad > 0:
                data = cv2.copyMakeBorder(
                    data,
                    topbottompad,
                    topbottompad,
                    leftrightpad,
                    leftrightpad,
                    cv2.BORDER_DEFAULT,
                )
            real_img["masks"].append(data)

    for key in real_img.keys():
        real_img[key] = np.array(real_img[key])

    if in_ch == 1:
        real_data = np.expand_dims(real_img["depth"], 3)
    elif in_ch == 3:
        real_data = np.array(real_img["rgb"])
    elif in_ch == 4:
        real_data = np.concatenate(
            [real_img["rgb"], np.expand_dims(real_img["depth"], 3)], axis=3
        )
    else:
        raise NotImplementedError
    # transfer to torch data
    real_rgbd = torch.from_numpy(real_data).permute(0, 3, 1, 2)  # (N, ch, H, W)
    return real_rgbd, real_img


def load_real_data(REAL_DIR="/ssd/pwral_real", in_ch=4):
    # load real data and preprocess
    REAL_DIR = Path(REAL_DIR)
    real_img = {"rgbd": [], "mask": []}
    corr_map = {"handle": 0, "door": 1, "cabinet": 2, "switch_link": 4, "fix_link": 3}
    for mode in ["doors", "drawers", "faucets"]:
        cate_dir = str(REAL_DIR / f"{mode}_labels/")
        dirs = os.listdir(cate_dir)
        for dir in sorted(dirs):
            # load data
            data = np.load(cate_dir + f"{dir}/{dir}_rgbd.npy")
            label = np.array(Image.open(cate_dir + f"{dir}/label.png"))
            # load color config
            with open(str(REAL_DIR / f"{mode}/") + f"{dir}_rgb.json") as f:
                conf = json.load(f)
            label_names = []
            for shape in conf["shapes"]:
                if shape["label"] not in label_names:
                    label_names.append(shape["label"])

            real_img["rgbd"].append(data)
            # fix mask index
            fix_label = label.copy()
            fix_label[label == 0] = 5
            for i in range(1, len(label_names) + 1):
                fix_label[label == i] = corr_map[label_names[i - 1]]
            real_img["mask"].append(fix_label)

    for key in real_img.keys():
        real_img[key] = np.array(real_img[key])
    print(real_img["rgbd"].shape, real_img["mask"].shape)

    # transfer to torch data
    real_rgbd = torch.from_numpy(real_img["rgbd"][..., :in_ch]).permute(0, 3, 1, 2)
    return real_rgbd, real_img


def main(logger, log_name, logfile_path, device, classes, only_vis):
    # load model
    model_path = VISUALMODEL_DIR / f"{log_name}/best.pth"
    config_path = VISUALMODEL_DIR / f"{log_name}/config.yaml"
    with open(config_path, "r") as f:
        cfg = yaml.safe_load(f)
    smp_cfg = cfg["smp_config"]
    if smp_cfg["mode"] == "RGBD":
        in_ch = 4
    elif smp_cfg["mode"] == "RGB":
        in_ch = 3
    elif smp_cfg["mode"] == "D":
        in_ch = 1
    else:
        raise NotImplementedError
    seg_model = SplitUnet(
        dropout_p=smp_cfg["dropout_p"],
        encoder_name=smp_cfg["encoder"],
        encoder_depth=smp_cfg["encoder_depth"],
        decoder_channels=smp_cfg["decoder_channels"],
        encoder_weights=smp_cfg["encoder_weights"],
        in_channels=in_ch,
        classes=cfg["num_classes"],
    )
    seg_model.load_state_dict(torch.load(model_path))
    seg_model.to(device)
    seg_model.eval()
    model_params = sum(p.numel() for p in seg_model.parameters())
    logger.info(f"model params: {model_params / 1e6:.2f} M")
    num_classes = cfg["num_classes"]

    real_data, real_img = load_real_data(in_ch=in_ch)

    print("real data:", real_data.shape)

    pred_mc_img = []
    uncertain_map_img = []

    with torch.no_grad():
        for i in tqdm(range(len(real_data)), desc="Evaluating"):
            d = real_data[i : i + 1].to(device=device, dtype=torch.float32)
            # get uncertainty
            uncertain_pred = []
            for _ in range(4):
                pred_seg_mtcsamp = seg_model(d, activate_dropout=True)  # (1, 6, H, W)
                uncertain_pred.append(F.softmax(pred_seg_mtcsamp, 1))
            uncertain_mean = torch.mean(
                torch.stack(uncertain_pred), dim=0
            )  # (1, 6, H, W)
            uncertain_map = -1.0 * torch.sum(
                uncertain_mean * torch.log(uncertain_mean + 1e-6), dim=1
            )  # (1, H, W)

            # predict segmentation
            seg_model.eval()
            pred_seg = seg_model.predict(
                d, T=10, dropout=True
            )  # (1, 6, H, W) # test-time aug
            # pred_seg = torch.mean(torch.stack(uncertain_pred), dim=0) # test-time dropout
            pred_mc = torch.argmax(pred_seg, dim=1)  # (1, H, W)
            pred_mc_img.append(pred_mc.cpu().numpy())

            uncertain_map -= uncertain_map.mean()
            uncertain_map = torch.maximum(
                uncertain_map, torch.zeros(uncertain_map.shape, device="cuda")
            )
            uncertain_map_img.append(uncertain_map.cpu().numpy())

        # save qualified results
        subplot_num = [2, 3]
        # seg_pred_img = pred_seg.detach().cpu().numpy()
        pred_mc_img = np.concatenate(pred_mc_img, 0)  # (N, H, W)
        uncertain_map_img = np.concatenate(uncertain_map_img, 0)  # (N, H, W)
        uncertain_map_img = (uncertain_map_img - uncertain_map_img.min()) / (
            uncertain_map_img.max() - uncertain_map_img.min()
        )
        kernel1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))
        kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
        for ind in tqdm(range(real_data.shape[0]), desc="Saving"):
            # print(f'saving image {ind}')
            uncertain_fix = cv2.dilate(uncertain_map_img[ind], kernel, iterations=1)
            uncertain_fix = cv2.morphologyEx(
                uncertain_fix, cv2.MORPH_OPEN, kernel1, iterations=1
            )
            plt.subplot(subplot_num[0], subplot_num[1], 1)
            plt.imshow(real_img["rgbd"][ind][..., :3])
            plt.subplot(subplot_num[0], subplot_num[1], 2)
            plt.imshow(visualize_depth(real_img["rgbd"][ind][..., 3])[..., ::-1])
            plt.subplot(subplot_num[0], subplot_num[1], 3)
            plt.imshow(visualize_seg(real_img["mask"][ind], num_classes))
            plt.subplot(subplot_num[0], subplot_num[1], 4)
            plt.imshow(visualize_seg(pred_mc_img[ind], num_classes=num_classes))
            plt.subplot(subplot_num[0], subplot_num[1], 5)
            plt.imshow(uncertain_map_img[ind], cmap="jet")
            plt.subplot(subplot_num[0], subplot_num[1], 6)
            plt.imshow(uncertain_fix, cmap="jet")

            plt.savefig(f"{logfile_path}/images/{ind:03}.jpg", dpi=200)
            plt.close()

        # save quantified results
        real_seg = torch.from_numpy(real_img["mask"]).to(device)  # (N, H, W)
        tp, fp, fn, tn = smp.metrics.get_stats(
            torch.from_numpy(pred_mc_img).long().to(device),
            real_seg.long(),
            mode="multiclass",
            num_classes=num_classes,
        )
        iou_metric = smp.metrics.iou_score(tp, fp, fn, tn, reduction="none")  # (bs, C)
        # f1_metric = smp.metrics.f1_score(tp, fp, fn, tn, reduction='none')  # (bs, C)
        logger.info(
            f"\n miou: {smp.metrics.iou_score(tp, fp, fn, tn, reduction='micro').item()}"
        )
        for ind, key in enumerate(classes):
            logger.info(f"iou_{key}: {torch.mean(iou_metric, dim=0)[ind].item()}")


if __name__ == "__main__":
    device = torch.device("cuda:0")
    classes_name = ["handle", "door", "cabinet", "switchlink", "fixlink", "other"]

    parser = argparse.ArgumentParser(description="test for 90 real images")
    parser.add_argument(
        "--log-name",
        "-ln",
        type=str,
        default="smp_model/20230327_231728_full_splitnet_new_data_aug_ce0.0_dropout0.2",
    )
    parser.add_argument(
        "--only-vis", action="store_true", help="only visualize results or not"
    )
    args = parser.parse_args()

    # generate log
    logfile_path = VISUALMODEL_DIR / f"{args.log_name}/real_results"
    if not os.path.exists(logfile_path):
        os.makedirs(logfile_path / "images")
    logfile_name = f"{time.strftime('%Y%m%d_%H%M%S', time.localtime())}_47arti"
    logger = logging.getLogger(logfile_name)
    logger.setLevel(logging.DEBUG)
    log_file = "{}.log".format(logfile_name)
    fileHandler = logging.FileHandler(os.path.join(logfile_path, log_file), mode="w")
    fileHandler.setLevel(logging.DEBUG)
    consoleHandler = logging.StreamHandler()
    consoleHandler.setLevel(logging.DEBUG)
    formatter = logging.Formatter(
        "[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    fileHandler.setFormatter(formatter)
    consoleHandler.setFormatter(formatter)
    logger.addHandler(fileHandler)
    logger.addHandler(consoleHandler)

    main(
        logger=logger,
        log_name=args.log_name,
        logfile_path=logfile_path,
        device=device,
        classes=classes_name,
        only_vis=args.only_vis,
    )

arti_mani/algorithms/visual_net/scripts/test_real_uncertainty.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import math
import os
import time
from datetime import datetime

import numpy as np
import torch
import torch.backends.cudnn as cudnn
from arti_mani.algorithms.config.config import Logger, set_seed
from arti_mani.algorithms.visual_net.Networks.keypoint_detection import (
    IntegralHumanPoseModel,
    JointLocationLoss,
)
from dataset import KeypointDataset
from path import Path
from torch import optim
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

cudnn.benchmark = True


class KPTSTrainer(object):
    def __init__(self):
        self.batch_size = 32
        self.lr = 1e-3
        self.epochs = 60
        self.mode = "RGBD"
        self.out_shape = (64, 40, 64)
        self.num_keypoints = 3
        self.num_deconv_layers = 3
        self.has_dropout = True
        self.device = torch.device("cuda:0")
        self.seed = set_seed()
        self.log_each_step = False
        exp_name = "kpt_model"
        exp_suffix = (
            f"kpts{self.num_keypoints}"
            f"_uvz_mode{self.mode}epoch{self.epochs}bs{self.batch_size}lr{self.lr}"
            f"_mobilenetv2_dropout0.2_newdatafilter2drawer_L2visable"
        )

        self.exp_log_path = f"log/{exp_name}/{time.strftime('%Y%m%d_%H%M%S', time.localtime())}_{exp_suffix}"
        if not os.path.exists(self.exp_log_path):
            os.makedirs(self.exp_log_path)
        self.logger = Logger(f"{self.exp_log_path}/log.txt")
        self.logger.info(
            f"\n▶ Starting Experiment '{exp_name}_{exp_suffix}' [seed: {self.seed}]"
        )

        self.model = IntegralHumanPoseModel(
            num_keypoints=self.num_keypoints,
            num_deconv_layers=self.num_deconv_layers,
            depth_dim=self.out_shape[0],
            has_dropout=self.has_dropout,
        ).to(self.device)

        model_params = sum(
            p.numel() for p in self.model.parameters() if p.requires_grad
        )
        self.logger.info(f"* number of parameters: {model_params}")

        # init optimizer
        self.optimizer = optim.Adam(params=self.model.parameters(), lr=self.lr)
        self.loss = JointLocationLoss()
        ### StepLR
        self.scheduler = torch.optim.lr_scheduler.MultiStepLR(
            self.optimizer,
            milestones=[30, 40, 50],
            gamma=0.5,
        )

        training_set = KeypointDataset(
            data_mode="train", augmentation=True, normalization=True, mode=self.mode
        )

        val_set = KeypointDataset(
            data_mode="val", augmentation=True, normalization=True, mode=self.mode
        )

        self.train_loader = DataLoader(
            training_set, batch_size=self.batch_size, shuffle=True, num_workers=8
        )
        self.val_loader = DataLoader(
            val_set, batch_size=self.batch_size, shuffle=True, num_workers=8
        )

        self.logger.info(
            f"training_set size: {len(training_set)}, val_set size: {len(val_set)}"
        )
        # init logging stuff
        self.log_path = Path(self.exp_log_path)
        tb_logdir = self.log_path.abspath()
        # print(f'tensorboard --logdir={tb_logdir}\n')
        self.logger.info(f"tensorboard --logdir={tb_logdir}\n")
        self.sw = SummaryWriter(self.log_path)
        self.log_freq = len(self.train_loader)

        # starting values values
        self.epoch = 0
        self.train_losses = None
        self.val_losses = None
        self.best_val_epoch = 0
        self.best_val_loss = None

    def load_ck(self):
        """
        load training checkpoint
        """
        ck_path = self.log_path / "training.ck"
        if ck_path.exists():
            ck = torch.load(ck_path, map_location=torch.device("cpu"))
            # print('[loading checkpoint \'{}\']'.format(ck_path))
            self.logger.info("[loading checkpoint '{}']".format(ck_path))
            self.epoch = ck["epoch"]
            self.model.load_state_dict(ck["model"], strict=True)
            self.model.to(self.device)
            self.best_val_loss = ck["best_val_loss"]
            if ck.get("optimizer", None) is not None:
                self.optimizer.load_state_dict(ck["optimizer"])

    def save_ck(self):
        """
        save training checkpoint
        """
        ck = {
            "epoch": self.epoch,
            "model": self.model.state_dict(),
            "optimizer": self.optimizer.state_dict(),
            "best_val_loss": self.best_val_loss,
        }
        torch.save(ck, self.log_path / "training.ck")

    def train(self):
        """
        train model for one epoch on the Training-Set.
        """
        self.model.train()

        start_time = time.time()
        times = []
        self.train_losses = []

        t = time.time()
        for step, sample in enumerate(self.train_loader):
            self.optimizer.zero_grad()
            img, pts = sample[2].float().to(self.device), sample[3].float().to(
                self.device
            )  # (N, 4, H, W)  (N, J, 3)
            pts_visable = sample[4].float().to(self.device)  # (N, J)

            pred = self.model.forward(img)  # (N, J*D, H', W')

            loss = self.loss(pred, pts, self.out_shape, pts_visable)
            loss.backward()
            self.train_losses.append(loss.item())

            self.optimizer.step(None)

            # print an incredible progress bar
            progress = (step + 1) / len(self.train_loader)
            progress_bar = ("█" * int(50 * progress)) + (
                "┈" * (50 - int(50 * progress))
            )
            times.append(time.time() - t)
            t = time.time()
            if self.log_each_step or (not self.log_each_step and progress == 1):
                self.logger.info(
                    "\r[{}] Epoch {:0{e}d}.{:0{s}d}: │{}│ {:6.2f}% │ Loss: {:.8f} │ ↯: {:5.2f} step/s".format(
                        datetime.now().strftime("%m-%d@%H:%M"),
                        self.epoch,
                        step + 1,
                        progress_bar,
                        100 * progress,
                        np.mean(self.train_losses),
                        1 / np.mean(times),
                        e=math.ceil(math.log10(self.epochs)),
                        s=math.ceil(math.log10(self.log_freq)),
                    )
                )

            # if step >= self.epochs - 1:
            #     break

        # log average loss of this epoch
        mean_epoch_loss = np.mean(self.train_losses)
        self.sw.add_scalar(
            tag="train_loss", scalar_value=mean_epoch_loss, global_step=self.epoch
        )

        cur_lr = self.optimizer.state_dict()["param_groups"][0]["lr"]
        self.sw.add_scalar(tag="train_lr", scalar_value=cur_lr, global_step=self.epoch)

        # log epoch duration
        self.logger.info(f" │ T: {time.time() - start_time:.2f} s")

    def test(self):
        """
        test model on the Test-Set
        """
        self.val_losses = []

        self.model.eval()

        t = time.time()
        for step, sample in enumerate(self.val_loader):
            img, pts = sample[2].float().to(self.device), sample[3].float().to(
                self.device
            )  # (N, 4, H, W)  (N, J, 3)
            pts_visable = sample[4].float().to(self.device)  # (N, J)

            pred = self.model.forward(img)  # (N, J*D, H', W')

            loss = self.loss(pred, pts, self.out_shape, pts_visable)
            self.val_losses.append(loss.item())

        # log average loss on test set
        mean_val_loss = np.mean(self.val_losses)
        # print(f'\t● AVG Loss on VAL-set: {mean_val_loss:.6f} │ T: {time() - t:.2f} s')
        self.logger.info(
            f"\t● AVG Loss on VAL-set: {mean_val_loss:.8f} │ T: {time.time() - t:.2f} s"
        )

        self.sw.add_scalar(
            tag="val_loss", scalar_value=mean_val_loss, global_step=self.epoch
        )

        # save best model
        if self.best_val_loss is None or mean_val_loss < self.best_val_loss:
            self.best_val_loss = mean_val_loss
            self.best_val_epoch = self.epoch
            torch.save(self.model.state_dict(), self.log_path / "best.pth")
            self.logger.info(f"\t● Save best model on epoch: {self.epoch}")

    def run(self):
        """
        start model training procedure (train > test > checkpoint > repeat)
        """
        for e in range(self.epoch, self.epochs):
            # if self.epoch == 13:
            #     pass
            self.train()
            if e % 1 == 0 and e != 0:
                self.test()
            self.epoch += 1
            self.scheduler.step()
            self.save_ck()
            if e % 50 == 0 and e != 0:
                torch.save(self.model.state_dict(), self.log_path / f"{e}.pth")


if __name__ == "__main__":
    trainer = KPTSTrainer()
    trainer.run()
    trainer.logger.info(f"\n▶ best val model: ")
    trainer.logger.info(
        f"\n  epoch, loss: {trainer.best_val_epoch}, {trainer.best_val_loss}"
    )

arti_mani/algorithms/visual_net/scripts/train_kpts.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import glob
import math
from collections import defaultdict
from copy import deepcopy
from datetime import datetime
from time import time

import cv2
import numpy as np
import segmentation_models_pytorch as smp
import torch
import yaml
from arti_mani.algorithms.config.config import Config
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import (
    CustomUnet,
    SplitUnet,
    WCELoss,
)
from arti_mani.utils.cv_utils import visualize_depth, visualize_seg
from dataset import SegDataset, seed_everything
from matplotlib import pyplot as plt
from path import Path
from torch import optim
from torch.nn import functional as F
from torch.nn.utils.clip_grad import clip_grad_norm_
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter


def load_real_data(
    REAL_DIR="/ssd2/pwral_real_seg",
    PAD_SIZE=(256, 144),
    IM_SIZE=(256, 144),
    in_ch=4,
    only_vis=False,
):
    # load real data and preprocess
    REAL_DIR = Path(REAL_DIR)
    leftrightpad = int((PAD_SIZE[0] - IM_SIZE[0]) / 2)
    topbottompad = int((PAD_SIZE[1] - IM_SIZE[1]) / 2)
    real_img = {"rgb": [], "depth": [], "masks": []}
    for mode in ["door", "drawer", "faucet"]:
        for im_mode in ["rgb", "depth"]:
            files = glob.glob(str(REAL_DIR / f"{mode}/*{im_mode}.npy"))
            files.sort()
            for file in files:
                data = np.load(file)
                data = cv2.resize(
                    src=data, dsize=IM_SIZE, interpolation=cv2.INTER_NEAREST
                )
                if leftrightpad > 0 or topbottompad > 0:
                    data = cv2.copyMakeBorder(
                        data,
                        topbottompad,
                        topbottompad,
                        leftrightpad,
                        leftrightpad,
                        cv2.BORDER_DEFAULT,
                    )
                if im_mode == "rgb":
                    data = data[..., ::-1] / 255.0
                real_img[im_mode].append(data)
    # load real masks
    if not only_vis:
        mask_files = []
        for mode in ["door", "drawer", "faucet"]:
            files = glob.glob(str(REAL_DIR / f"masks/*{mode}.npy"))
            files.sort()
            mask_files.extend(files)
        for mf in mask_files:
            data = np.load(mf)
            data = cv2.resize(src=data, dsize=IM_SIZE, interpolation=cv2.INTER_NEAREST)
            if leftrightpad > 0 or topbottompad > 0:
                data = cv2.copyMakeBorder(
                    data,
                    topbottompad,
                    topbottompad,
                    leftrightpad,
                    leftrightpad,
                    cv2.BORDER_DEFAULT,
                )
            real_img["masks"].append(data)

    for key in real_img.keys():
        real_img[key] = np.array(real_img[key])

    if in_ch == 1:
        real_data = np.expand_dims(real_img["depth"], 3)
    elif in_ch == 3:
        real_data = np.array(real_img["rgb"])
    elif in_ch == 4:
        real_data = np.concatenate(
            [real_img["rgb"], np.expand_dims(real_img["depth"], 3)], axis=3
        )
    else:
        raise NotImplementedError
    # transfer to torch data
    real_rgbd = torch.from_numpy(real_data).permute(0, 3, 1, 2)  # (N, ch, H, W)
    return real_rgbd, real_img


class SMPTrainer(object):
    def __init__(self, cfg):
        # super().__init__(cfg)
        self.cfg = cfg
        smp_cfg = cfg.smp_config

        if smp_cfg["mode"] == "RGBD":
            in_ch = 4
        elif smp_cfg["mode"] == "RGD":
            in_ch = 3
        elif smp_cfg["mode"] == "RGB":
            in_ch = 3
        elif smp_cfg["mode"] == "D":
            in_ch = 1
        else:
            raise NotImplementedError

        if smp_cfg["encoder"] == "splitnet":
            self.model = SplitUnet(
                dropout_p=smp_cfg["dropout_p"],
                encoder_name=smp_cfg["encoder"],
                encoder_depth=smp_cfg["encoder_depth"],
                decoder_channels=smp_cfg["decoder_channels"],
                encoder_weights=smp_cfg["encoder_weights"],
                in_channels=in_ch,
                classes=cfg.num_classes,
            ).to(cfg.device)
        else:
            self.model = CustomUnet(
                dropout_p=smp_cfg["dropout_p"],
                encoder_name=smp_cfg["encoder"],
                encoder_depth=smp_cfg["encoder_depth"],
                decoder_channels=smp_cfg["decoder_channels"],
                encoder_weights=smp_cfg["encoder_weights"],
                in_channels=in_ch,
                classes=cfg.num_classes,
            ).to(cfg.device)

        self.classes = smp_cfg["classes"]
        self.num_classes = cfg.num_classes

        model_params = sum(
            p.numel() for p in self.model.parameters() if p.requires_grad
        )
        self.cfg.logger.info(f"* number of parameters: {model_params}")

        # init optimizer
        self.optimizer = optim.Adam(params=self.model.parameters(), lr=cfg.lr)
        self.val_loss = smp.losses.SoftCrossEntropyLoss(smooth_factor=0)
        if self.cfg.loss == "focal":
            self.loss = smp.losses.FocalLoss(mode="multiclass")
        elif self.cfg.loss == "wce":
            self.loss = WCELoss(smooth_factor=0.1)
        elif self.cfg.loss == "ce":
            self.loss = smp.losses.SoftCrossEntropyLoss(smooth_factor=0.1)
        elif self.cfg.loss == "dice":
            self.loss = smp.losses.DiceLoss(mode="multiclass")
        elif self.cfg.loss == "tver":
            self.loss = smp.losses.TverskyLoss(mode="multiclass")
        else:
            raise NotImplementedError
        ### StepLR
        self.scheduler = torch.optim.lr_scheduler.StepLR(
            self.optimizer, step_size=40, gamma=0.3
        )

        training_set = SegDataset(
            data_mode="train",  # "train", "train_lowres"
            augmentation=smp_cfg["augmentation"],
            copy_paste=smp_cfg["copy_paste"],
            classes=smp_cfg["classes"],
            mode=smp_cfg["mode"],
        )

        val_set = SegDataset(
            data_mode="val",  # "val", "val_lowres"
            augmentation=False,
            copy_paste=False,
            classes=smp_cfg["classes"],
            mode=smp_cfg["mode"],
        )

        self.real_data, self.real_img = load_real_data(in_ch=in_ch)
        self.real_data = self.real_data.to(device=cfg.device, dtype=torch.float32)

        print("real data:", self.real_data.shape)

        self.train_loader = DataLoader(
            training_set,
            batch_size=cfg.batch_size,
            shuffle=True,
            pin_memory=True,
            num_workers=cfg.n_workers,
        )
        self.val_loader = DataLoader(
            val_set,
            batch_size=cfg.batch_size,
            shuffle=True,
            pin_memory=True,
            num_workers=cfg.n_workers,
        )

        self.cfg.logger.info(
            f"training_set size: {len(training_set)}, val_set size: {len(val_set)}"
        )
        # init logging stuff
        self.log_path = Path(cfg.exp_log_path)
        tb_logdir = self.log_path.abspath()
        # print(f'tensorboard --logdir={tb_logdir}\n')
        self.cfg.logger.info(f"tensorboard --logdir={tb_logdir}\n")
        self.sw = SummaryWriter(self.log_path)
        self.log_freq = len(self.train_loader)

        # starting values values
        self.epoch = 0
        self.train_losses = None
        self.val_losses = None
        self.best_val_epoch = 0
        self.best_val_loss = None
        self.best_val_items = None

    def load_ck(self):
        """load training checkpoint."""
        ck_path = self.log_path / "training.ck"
        if ck_path.exists():
            ck = torch.load(ck_path, map_location=torch.device("cpu"))
            # print('[loading checkpoint \'{}\']'.format(ck_path))
            self.cfg.logger.info("[loading checkpoint '{}']".format(ck_path))
            self.epoch = ck["epoch"]
            self.model.load_state_dict(ck["model"], strict=True)
            self.model.to(self.cfg.device)
            self.best_val_loss = ck["best_val_loss"]
            if ck.get("optimizer", None) is not None:
                self.optimizer.load_state_dict(ck["optimizer"])

    def save_ck(self):
        """save training checkpoint."""
        ck = {
            "epoch": self.epoch,
            "model": self.model.state_dict(),
            "optimizer": self.optimizer.state_dict(),
            "best_val_loss": self.best_val_loss,
        }
        torch.save(ck, self.log_path / "training.ck")

    def train(self):
        """train model for one epoch on the Training-Set."""
        self.model.train()

        start_time = time()
        times = []
        self.train_losses = []
        self.uncertaintys = []
        # train_items = dict.fromkeys(["iou", "f1", "acc", "prec", "recall"], [])
        train_items = defaultdict(list)

        t = time()
        for step, sample in enumerate(self.train_loader):
            self.optimizer.zero_grad()
            img = sample[2].to(device=self.cfg.device, dtype=torch.float32)
            seg_gt_onehot = (
                torch.stack(sample[3]).transpose(1, 0).to(self.cfg.device)
            )  # (N, C, H, W)
            seg_gt = torch.argmax(seg_gt_onehot, dim=1)  # (N, H, W)

            pred = self.model.forward(img)  # (N, C, H, W)

            # pred_mc = torch.argmax(pred, dim=1)  # (N, H, W)
            # first compute statistics for true positives, false positives, false negative and
            # true negative "pixels"
            # tp, fp, fn, tn = smp.metrics.get_stats(pred, seg_gt_onehot.long(), mode='multilabel', threshold=0.5)
            # tp, fp, fn, tn = smp.metrics.get_stats(pred_mc.long(),
            #                                        seg_gt.long(),
            #                                        mode='multiclass',
            #                                        num_classes=self.num_classes)

            # iou_metric = smp.metrics.iou_score(tp, fp, fn, tn, reduction='none')  # (bs, C)
            # f1_metric = smp.metrics.f1_score(tp, fp, fn, tn, reduction='none')  # (bs, C)
            # for ind, key in enumerate(self.classes):
            #     train_items[f'iou_{key}'].append(torch.mean(iou_metric, dim=0)[ind].item())
            #     train_items[f'f1_{key}'].append(torch.mean(f1_metric, dim=0)[ind].item())

            # then compute metrics with required reduction (see metric docs)
            # train_items['miou'].append(
            #     smp.metrics.iou_score(tp, fp, fn, tn, reduction='micro').item())
            # train_items['f1'].append(smp.metrics.f1_score(tp, fp, fn, tn, reduction='micro').item())
            # train_items["acc"].append(smp.metrics.accuracy(tp, fp, fn, tn, reduction="macro").item())
            # train_items["prec"].append(smp.metrics.precision(tp, fp, fn, tn, reduction="micro-imagewise").item())
            # train_items["recall"].append(smp.metrics.recall(tp, fp, fn, tn, reduction="micro-imagewise").item())

            # loss = softmax_focal_loss(seg_feat, seg_gt_onehot, alpha=0.75, reduction='mean')
            # loss = soft_dice_loss(seg_feat, seg_gt_onehot)
            # loss = log_cosh_dice_loss(seg_feat, seg_gt_onehot)
            # loss = self.wfocal_loss(seg_feat, seg_gt)
            # loss = F.cross_entropy(pred, seg_gt_onehot)
            if self.cfg.loss == "wce":
                # get uncertainty
                # get uncertainty map
                uncertain_pred = []
                with torch.no_grad():
                    for _ in range(self.cfg.smp_config["sample"]):
                        pred_seg_mtcsamp = self.model(
                            img, activate_dropout=True
                        )  # (N, 6, H, W)
                        uncertain_pred.append(F.softmax(pred_seg_mtcsamp, dim=1))
                    uncertain_mean = torch.mean(
                        torch.stack(uncertain_pred), dim=0
                    )  # (N, 6, H, W)
                    # print(torch.stack(uncertain_pred).shape, uncertain_mean.shape)
                    uncertain_map = -1.0 * torch.sum(
                        uncertain_mean * torch.log(uncertain_mean + 1e-6),
                        dim=1,
                        keepdim=True,
                    )  # (N, H, W)
                # print(uncertain_map.min(), uncertain_map.max())
                min_w = self.cfg.smp_config["min_weight"]
                norm_uncertainty = (uncertain_map - uncertain_map.min()) / (
                    uncertain_map.max() - uncertain_map.min()
                )
                weight = min_w + (1 - min_w) * norm_uncertainty**2
                loss = self.loss(pred, seg_gt, weight)
                self.uncertaintys.append(uncertain_map.mean().item())
            else:
                loss = self.loss(pred, seg_gt)
            loss.backward()
            self.train_losses.append(loss.item())
            clip_grad_norm_(self.model.parameters(), 10)
            self.optimizer.step()

            # print an incredible progress bar
            progress = (step + 1) / len(self.train_loader)
            bar_len = 40
            progress_bar = ("■" * int(bar_len * progress)) + (
                "┈" * (bar_len - int(bar_len * progress))
            )
            times.append(time() - t)
            t = time()
            if self.cfg.log_step or (not self.cfg.log_step and progress == 1):
                if step > 0 and step % self.cfg.log_step == 0:
                    self.cfg.logger.info(
                        "\r[{}] Epoch {:0{e}d}.{:0{s}d}: │{}│ {:6.2f}% │ Loss: {:.6f} │ ↯: {:5.2f} step/s".format(
                            datetime.now().strftime("%m-%d@%H:%M"),
                            self.epoch,
                            step,
                            progress_bar,
                            100 * progress,
                            np.mean(self.train_losses[-self.cfg.log_step :]),
                            1 / np.mean(times[-self.cfg.log_step :]),
                            e=math.ceil(math.log10(self.cfg.epochs)),
                            s=math.ceil(math.log10(self.log_freq)),
                        )
                    )

            # if step >= self.cfg.epoch_len - 1:
            #     break

        # log average loss of this epoch
        mean_uncertainty = np.mean(self.uncertaintys)
        self.sw.add_scalar(
            tag="uncertainty", scalar_value=mean_uncertainty, global_step=self.epoch
        )
        mean_epoch_loss = np.mean(self.train_losses)
        self.sw.add_scalar(
            tag="train/loss", scalar_value=mean_epoch_loss, global_step=self.epoch
        )
        for key in train_items.keys():
            mean_epoch_item = np.mean(train_items[key])
            self.sw.add_scalar(
                tag=f"train/iou/{key}",
                scalar_value=mean_epoch_item,
                global_step=self.epoch,
            )

        cur_lr = self.optimizer.state_dict()["param_groups"][0]["lr"]
        self.sw.add_scalar(tag="train_lr", scalar_value=cur_lr, global_step=self.epoch)

        # log epoch duration
        self.cfg.logger.info(f" │ T: {time() - start_time:.2f} s")

    def test(self):
        """test model on the Test-Set."""
        self.val_losses = []
        # val_items = dict.fromkeys(["iou", "f1", "acc", "prec", "recall"], [])
        val_items = defaultdict(list)

        self.model.eval()

        t = time()
        with torch.no_grad():
            for step, sample in enumerate(self.val_loader):
                img, masks = (
                    sample[2].to(device=self.cfg.device, dtype=torch.float32),
                    sample[3],
                )
                seg_gt_onehot = (
                    torch.stack(masks).transpose(1, 0).to(self.cfg.device)
                )  # (N, C, H, W)
                seg_gt = torch.argmax(seg_gt_onehot, dim=1)  # (N, H, W)

                pred = self.model(img)  # (N, C, H, W)
                pred_mc = torch.argmax(pred, dim=1)  # (N, H, W)

                # loss = softmax_focal_loss(seg_feat, seg_gt_onehot, alpha=0.75, reduction='mean')
                # loss = soft_dice_loss(seg_feat, seg_gt_onehot)
                # loss = log_cosh_dice_loss(seg_feat, seg_gt_onehot)
                # loss = self.wfocal_loss(seg_feat, seg_gt)
                # loss = F.cross_entropy(pred, seg_gt_onehot)

                loss = self.val_loss(pred, seg_gt)
                self.val_losses.append(loss.item())

                # first compute statistics for true positives, false positives, false negative and
                # true negative "pixels"
                # tp, fp, fn, tn = smp.metrics.get_stats(pred, seg_gt_onehot.long(), mode='multilabel', threshold=0.5)
                tp, fp, fn, tn = smp.metrics.get_stats(
                    pred_mc.long(),
                    seg_gt.long(),
                    mode="multiclass",
                    num_classes=self.num_classes,
                )

                iou_metric = smp.metrics.iou_score(
                    tp, fp, fn, tn, reduction="none"
                )  # (bs, C)
                # f1_metric = smp.metrics.f1_score(tp, fp, fn, tn, reduction='none')  # (bs, C)
                for ind, key in enumerate(self.classes):
                    val_items[f"iou_{key}"].append(
                        torch.mean(iou_metric, dim=0)[ind].item()
                    )
                    # val_items[f'f1_{key}'].append(torch.mean(f1_metric, dim=0)[ind].item())

                # then compute metrics with required reduction (see metric docs)
                val_items["miou"].append(
                    smp.metrics.iou_score(tp, fp, fn, tn, reduction="micro").item()
                )
                # val_items['f1'].append(smp.metrics.f1_score(tp, fp, fn, tn, reduction='micro').item())
                # val_items["acc"].append(smp.metrics.accuracy(tp, fp, fn, tn, reduction="macro").item())
                # val_items["prec"].append(smp.metrics.precision(tp, fp, fn, tn, reduction="micro-imagewise").item())
                # val_items["recall"].append(smp.metrics.recall(tp, fp, fn, tn, reduction="micro-imagewise").item())

        # log average loss on test set
        mean_val_loss = np.mean(self.val_losses)
        # print(f'\t● AVG Loss on VAL-set: {mean_val_loss:.6f} │ T: {time() - t:.2f} s')
        self.cfg.logger.info(
            f"\t● AVG Loss on VAL-set: {mean_val_loss:.8f} │ T: {time() - t:.2f} s"
        )

        self.sw.add_scalar(
            tag="val/loss", scalar_value=mean_val_loss, global_step=self.epoch
        )

        cur_mean_epoch_item = defaultdict(float)
        for key in val_items.keys():
            cur_mean_epoch_item[key] = float(np.mean(val_items[key]))
            self.sw.add_scalar(
                tag=f"val/iou/{key}",
                scalar_value=cur_mean_epoch_item[key],
                global_step=self.epoch,
            )

        # save best model
        if self.best_val_loss is None or mean_val_loss < self.best_val_loss:
            self.best_val_loss = mean_val_loss
            self.best_val_epoch = self.epoch
            torch.save(self.model.state_dict(), self.log_path / "best.pth")
            self.cfg.logger.info(f"\t● Save best model on epoch: {self.epoch}")
            self.cfg.logger.info(f"\t● best model metrics:")
            iou_str, f1_str = "", ""
            for k in cur_mean_epoch_item:
                if k.startswith("iou"):
                    iou_str += f"{k[4:]}: {cur_mean_epoch_item[k]:.3f} "
                # if 'f1' in k:
                #     obj = k if k == 'f1' else k[3:]
                #     f1_str += f'{obj}: {cur_mean_epoch_item[k]:.3f} '
            self.cfg.logger.info(f"test on val dataset:")
            self.cfg.logger.info(f'miou: {cur_mean_epoch_item["miou"]:.3f}')
            self.cfg.logger.info(f"{iou_str}")
            # self.cfg.logger.info('f1:')
            # self.cfg.logger.info(f'{f1_str}')
            self.best_val_items = cur_mean_epoch_item

    def test_real(self, vis=False):
        """test model on the Real-Test-Set."""
        self.model.eval()

        with torch.no_grad():
            # get uncertainty map
            uncertain_pred = []
            for ind in range(self.cfg.smp_config["sample"]):
                pred_seg_mtcsamp = self.model(
                    self.real_data, activate_dropout=True
                )  # (N, 6, H, W)
                uncertain_pred.append(pred_seg_mtcsamp)
            uncertain_mean = torch.mean(
                torch.stack(uncertain_pred), dim=0
            )  # (N, 6, H, W)
            uncertain_map = -1.0 * torch.sum(
                uncertain_mean * torch.log(uncertain_mean + 1e-6), dim=1
            )  # (N, H, W)
            # get segmentation prediction
            pred_seg = self.model(self.real_data)  # (N, 6, H, W)
            pred_mc = torch.argmax(pred_seg, dim=1).to(self.cfg.device)  # (N, H, W)
            real_seg = torch.from_numpy(self.real_img["masks"]).to(
                self.cfg.device
            )  # (N, H, W)

            # save qualified results
            if vis:
                subplot_num = [2, 3]
                # seg_pred_img = pred_seg.detach().cpu().numpy()
                pred_mc_img = pred_mc.cpu().numpy()  # (N, H, W)
                uncertain_map_img = uncertain_map.cpu().numpy()  # (N, H, W)
                uncertain_map_img = (uncertain_map_img - uncertain_map_img.min()) / (
                    uncertain_map_img.max() - uncertain_map_img.min()
                )
                kernel1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))
                kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
                for ind in range(self.real_data.shape[0]):
                    uncertain_fix = cv2.dilate(
                        uncertain_map_img[ind], kernel, iterations=1
                    )
                    # uncertain_fix = cv2.dilate(uncertain_fix, kernel2, iterations=1)
                    # uncertain_fix = cv2.erode(uncertain_fix, kernel, iterations=1)
                    # uncertain_fix = cv2.morphologyEx(uncertain_fix, cv2.MORPH_CLOSE, kernel1, iterations=1)
                    uncertain_fix = cv2.morphologyEx(
                        uncertain_fix, cv2.MORPH_OPEN, kernel1, iterations=1
                    )
                    plt.subplot(subplot_num[0], subplot_num[1], 1)
                    plt.imshow(self.real_img["rgb"][ind])
                    plt.subplot(subplot_num[0], subplot_num[1], 2)
                    plt.imshow(visualize_depth(self.real_img["depth"][ind])[..., ::-1])
                    plt.subplot(subplot_num[0], subplot_num[1], 3)
                    plt.imshow(
                        visualize_seg(self.real_img["masks"][ind], self.cfg.num_classes)
                    )
                    plt.subplot(subplot_num[0], subplot_num[1], 4)
                    plt.imshow(
                        visualize_seg(
                            pred_mc_img[ind], num_classes=self.cfg.num_classes
                        )
                    )
                    plt.subplot(subplot_num[0], subplot_num[1], 5)
                    plt.imshow(uncertain_map_img[ind], cmap="jet")
                    plt.subplot(subplot_num[0], subplot_num[1], 6)
                    plt.imshow(uncertain_fix, cmap="jet")
                    # for pid in range(num_classes):
                    #     plt.subplot(subplot_num[0], subplot_num[1], pid + 4)
                    #     plt.imshow(seg_pred_img[ind][pid])

                    plt.savefig(f"./vis_seg/{ind:03}.jpg", dpi=200)

            # real_seg_onehot = F.one_hot(
            #     real_seg.long(), num_classes=num_classes
            # ).long().permute(0, 3, 1, 2).to(device)  # (N, 6, H, W)
            tp, fp, fn, tn = smp.metrics.get_stats(
                pred_mc.long(),
                real_seg.long(),
                mode="multiclass",
                num_classes=self.num_classes,
            )
            iou_metric = smp.metrics.iou_score(
                tp, fp, fn, tn, reduction="none"
            )  # (bs, C)
            # f1_metric = smp.metrics.f1_score(tp, fp, fn, tn, reduction="none")  # (bs, C)
            miou = smp.metrics.iou_score(tp, fp, fn, tn, reduction="micro").item()
            self.cfg.logger.info(f"test on real dataset:")
            self.cfg.logger.info(f"miou: {miou:.3f}")
            self.sw.add_scalar(
                tag=f"test_real/iou/miou", scalar_value=miou, global_step=self.epoch
            )
            logstr = ""
            for ind, key in enumerate(self.classes):
                cur_iou = torch.mean(iou_metric, dim=0)[ind].item()
                logstr += f"{key}: {cur_iou:.3f} "
                self.sw.add_scalar(
                    tag=f"test_real/iou/{key}",
                    scalar_value=cur_iou,
                    global_step=self.epoch,
                )
            self.cfg.logger.info(logstr)

    def run(self):
        """start model training procedure (train > test > checkpoint >
        repeat)"""
        for e in range(self.epoch, self.cfg.epochs):
            # if self.epoch == 13:
            #     pass
            self.train()
            if e % 1 == 0:
                self.test()
                self.test_real()
            self.epoch += 1
            self.scheduler.step()
            self.save_ck()
            if e % 10 == 0 and e != 0:
                torch.save(self.model.state_dict(), self.log_path / f"{e}.pth")


def main(exp_name=None, seed=None):
    cfg = Config(seed=seed, exp_name=exp_name, mode="train", log=True)

    seed_everything(cfg.seed)

    ## save config file
    with open(f"{cfg.exp_log_path}/config.yaml", "w") as file:
        cfg_text = deepcopy(cfg)
        del cfg_text.logger
        yaml.emitter.Emitter.process_tag = lambda self, *args, **kw: None
        yaml.dump(cfg_text, file)

    cfg.logger.info(f"\n▶ Starting Experiment '{exp_name}' [seed: {cfg.seed}]")

    # trainer = Trainer(cfg=cfg)
    trainer = SMPTrainer(cfg=cfg)
    trainer.run()
    cfg.logger.info(f"\n▶ best val model: ")
    cfg.logger.info(
        f"\n  epoch, loss: {trainer.best_val_epoch}, {trainer.best_val_loss}"
    )
    cfg.logger.info(f"\n  metrics: {trainer.best_val_items}")


if __name__ == "__main__":
    main(exp_name="smp_model", seed=2)  # pretrain_unet, smp_model

arti_mani/algorithms/visual_net/scripts/train_seg.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from .arti_mani import *
from .fixed_robotenv import *
from .open_cabinet_door import *
from .open_cabinet_drawer import *
from .turn_faucet import *

arti_mani/envs/__init__.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from collections import OrderedDict
from typing import List

import numpy as np
import open3d as o3d
import sapien.core as sapien
import torch
import torch.nn.functional as F
import trimesh
import yaml
from arti_mani import ASSET_DIR
from arti_mani.agents.camera import get_camera_rgb, get_camera_seg, get_texture
from arti_mani.algorithms.rl_iam.Sampling import (
    FPSSampler,
    RandomSampler,
    ScoreSampler,
    UniDownsampler,
)
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import (
    CustomUnet,
    CustomUnetNew,
    SplitUnet,
)
from arti_mani.algorithms.visual_net.Networks.keypoint_detection import (
    IntegralHumanPoseModel,
    soft_argmax,
)
from arti_mani.envs.fixed_robotenv import FixedXmate3RobotiqSensorLowResEnv
from arti_mani.utils.articulation_utils import ignore_collision
from arti_mani.utils.common import (
    convert_np_bool_to_float,
    flatten_state_dict,
    np_random,
    random_choice,
    register_gym_env,
)
from arti_mani.utils.contrib import (
    angle_distance_ms,
    apply_pose_to_points,
    norm,
    normalize_and_clip_in_interval,
    o3d_to_trimesh,
    trimesh_to_o3d,
    uv2xyz,
)
from arti_mani.utils.geometry import sample_grasp_multipoints_ee, transform_points
from arti_mani.utils.o3d_utils import (
    check_coplanar,
    get_visual_body_meshes,
    merge_mesh,
    merge_meshes,
    np2mesh,
    pcd_uni_down_sample_with_crop,
)
from arti_mani.utils.sapien_utils import (
    get_entity_by_name,
    hex2rgba,
    set_articulation_render_material,
)
from arti_mani.utils.trimesh_utils import get_actor_mesh
from sapien.core import Pose
from transforms3d.euler import euler2quat

CMAPS = [
    (1, 0, 0),  # "red"
    (0, 0, 1),  # "blue"
    (1, 1, 0),  # "yellow"
    (0, 1, 0),  # "green"
    (0.627, 0.125, 0.941),  # "purple"
    (0.753, 0.753, 0.753),  # "grey"
]


@register_gym_env("ArtiMani-v0", max_episode_steps=100)
class ArtiMani(FixedXmate3RobotiqSensorLowResEnv):
    SUPPORTED_OBS_MODES = (
        "state",
        "state_dict",
        "state_egorgbd",
        "state_ego_segpoints_gt",
        "state_egosegpoints",
        "state_egostereo_rgbd",
        "state_egostereo_segpoints",
        "state_egostereo_segpoints_gt",
        "state_egostereo_keypoints",
        "state_egostereo_keypoints_gt",
        "state_egostereo_dseg",
        "state_egostereo_dseg_gt",
        "state_egostereo_dexpoints",
    )
    SUPPORTED_REWARD_MODES = ("dense", "sparse")
    SUPPORTED_CONTROL_MODES = (
        "pd_joint_delta_pos",
        "pd_ee_delta_pose",
        "pd_ee_delta_pos",
    )

    def __init__(
        self,
        articulation_ids: List[int] = (),
        segmodel_path: str = None,
        sample_mode: str = None,
        frame_num: int = 1,
        sample_num: int = 30,
        other_handle_visible=False,
        num_classes: int = 6,
        obs_mode=None,
        reward_mode=None,
        sim_freq=500,
        control_freq=20,
        control_mode=None,
        add_eepadpts=False,
        device: str = "cuda:0",
    ):
        self._arti_info = OrderedDict()
        with open(
            ASSET_DIR / f"partnet_mobility_configs/fixed_artis_new.yml", "r"
        ) as f:
            self._arti_info = yaml.safe_load(f)
        # self._arti_fail_rate = dict.fromkeys(self._arti_info.keys(), 1)
        self._arti_mode = None
        self._articulation_info = None
        if isinstance(articulation_ids, int):
            articulation_ids = [articulation_ids]

        assert len(articulation_ids) > 0

        self._arti_fail_total = OrderedDict()
        for arti_id in articulation_ids:
            self._arti_fail_total[arti_id] = [0, 0]

        self.articulation_ids = articulation_ids
        self.articulation_id = None
        self.articulation_scale = None
        self.other_handle_visible = other_handle_visible
        self.num_classes = num_classes
        self.device = device
        self.sample_mode = sample_mode
        self.sample_num = sample_num
        if sample_mode == "full_downsample":
            self.sampler = UniDownsampler(sample_num=sample_num)
        elif sample_mode == "score_sample":
            self.sampler = ScoreSampler(sample_num=sample_num)
        elif sample_mode == "fps_sample":
            self.sampler = FPSSampler(sample_num=sample_num)
        elif sample_mode == "random_sample":
            self.sampler = RandomSampler(sample_num=sample_num)
        elif sample_mode == None:
            self.sampler = None
        else:
            raise NotImplementedError(sample_mode)

        if obs_mode in [
            "state_egosegpoints",
            "state_egostereo_segpoints",
            "state_egostereo_keypoints",
        ]:
            self._load_segmodel(segmodel_path, obs_mode)
        self.out_shape = (64, 40, 64)
        self.kpts_max = np.array([255, 143, 1.0])
        self.kpts_min = np.array([0, 0, 0.18])
        self.cam_para = OrderedDict()
        self.add_eepadpts = add_eepadpts

        super().__init__(
            obs_mode, reward_mode, control_mode, sim_freq, control_freq, device
        )

    def _load_segmodel(self, segmodel_path, obs_mode):
        if obs_mode == "state_egostereo_keypoints":
            model_path = f"{segmodel_path}/best.pth"
            self.kpt_model = IntegralHumanPoseModel(
                num_keypoints=3, num_deconv_layers=3, depth_dim=64, has_dropout=False
            )
            self.kpt_model.load_state_dict(torch.load(model_path))
            self.kpt_model.to(torch.device(self.device))
            self.kpt_model.eval()
        else:
            model_path = f"{segmodel_path}/best.pth"
            config_path = f"{segmodel_path}/config.yaml"
            with open(config_path, "r") as f:
                cfg = yaml.safe_load(f)
            assert self.num_classes == cfg["num_classes"]
            smp_cfg = cfg["smp_config"]
            if smp_cfg["mode"] == "RGBD":
                in_ch = 4
            elif smp_cfg["mode"] == "RGB":
                in_ch = 3
            elif smp_cfg["mode"] == "D":
                in_ch = 1
            else:
                raise NotImplementedError
            if "has_dropout" in smp_cfg.keys():
                self.segmodel = CustomUnet(
                    has_dropout=False,
                    encoder_name=smp_cfg["encoder"],
                    encoder_depth=smp_cfg["encoder_depth"],
                    decoder_channels=smp_cfg["decoder_channels"],
                    encoder_weights=smp_cfg["encoder_weights"],
                    in_channels=in_ch,
                    classes=cfg["num_classes"],
                    activation=smp_cfg["activation"],
                )
            elif smp_cfg["encoder"] == "splitnet":
                self.segmodel = SplitUnet(
                    dropout_p=smp_cfg["dropout_p"],
                    encoder_name=smp_cfg["encoder"],
                    encoder_depth=smp_cfg["encoder_depth"],
                    decoder_channels=smp_cfg["decoder_channels"],
                    encoder_weights=smp_cfg["encoder_weights"],
                    in_channels=in_ch,
                    classes=cfg["num_classes"],
                )
            else:
                self.segmodel = CustomUnetNew(
                    dropout_p=smp_cfg["dropout_p"],
                    encoder_name=smp_cfg["encoder"],
                    encoder_depth=smp_cfg["encoder_depth"],
                    decoder_channels=smp_cfg["decoder_channels"],
                    encoder_weights=smp_cfg["encoder_weights"],
                    in_channels=in_ch,
                    classes=cfg["num_classes"],
                )
            self.segmodel.load_state_dict(torch.load(model_path))
            self.segmodel.to(torch.device(self.device))
            self.segmodel.eval()

    def _setup_camera(self):
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            self.render_camera = self._scene.add_camera(
                "topview", 512, 512, 1, 0.01, 10
            )
            self.render_camera.set_local_pose(
                Pose([-0.1, 0.0, 1.5], euler2quat(0, 1.57, 0))
            )
        elif self._arti_mode in ["faucet", "laptop", "kitchen_pot"]:
            self.render_camera = self._scene.add_camera(
                "front_view", 512, 512, 1, 0.01, 10
            )
            self.render_camera.set_local_pose(
                Pose([1, 0, 1.2], euler2quat(0, 0.6, 3.14))
            )
        else:
            raise NotImplementedError

    def _setup_viewer(self):
        super()._setup_viewer()
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            ### top view
            self._viewer.set_camera_xyz(0.0, 0.0, 1.2)
            self._viewer.set_camera_rpy(0, -3.14, 0)
        elif self._arti_mode in ["faucet", "laptop", "kitchen_pot"]:
            ### front view
            self._viewer.set_camera_xyz(1.0, 0.0, 1.2)
            self._viewer.set_camera_rpy(0, -0.5, 3.14)
            ### top view
            # self._viewer.set_camera_xyz(0.0, 0.0, 1.2)
            # self._viewer.set_camera_rpy(0, -3.14, 0)
        else:
            raise NotImplementedError

    def reset(
        self,
        seed=None,
        reconfigure=False,
        articulation_id=None,
        articulation_scale=None,
    ):
        self.set_episode_rng(seed)
        _reconfigure = self._set_model(articulation_id, articulation_scale)
        reconfigure = _reconfigure or reconfigure
        ret = super().reset(seed=self._episode_seed, reconfigure=reconfigure)
        return ret

    def _set_model(self, articulation_id, articulation_scale):
        """Set the model id and scale. If not provided, choose one randomly."""
        reconfigure = False

        # Model ID
        if articulation_id is None:
            articulation_id = self.articulation_ids[
                self._episode_rng.choice(len(self.articulation_ids), p=None)
            ]
        if articulation_id != self.articulation_id:
            reconfigure = True
        self.articulation_id = articulation_id
        for mode in self._arti_info.keys():
            if self.articulation_id in self._arti_info[mode].keys():
                self._articulation_info = self._arti_info[mode]
                self._arti_mode = mode
                break
        assert self._articulation_info is not None
        self._articulation_config = self._articulation_info[self.articulation_id]

        if self._arti_mode in ["faucet", "laptop", "kitchen_pot"]:
            # Scale
            if articulation_scale is None:
                articulation_scale = self._articulation_config.get("scale")
            if articulation_scale is None:
                bbox_max, bbox_min = (
                    self._articulation_config["bbox_max"],
                    self._articulation_config["bbox_min"],
                )
                bbox_size = np.float32(bbox_max) - np.float32(bbox_min)
                articulation_scale = 0.3 / max(bbox_size)  # hardcode
            if articulation_scale != self.articulation_scale:
                reconfigure = True
            self.articulation_scale = articulation_scale

            if "offset" in self._articulation_config:
                self.articulation_offset = np.float32(
                    self._articulation_config["offset"]
                )
            else:
                bbox_min = self._articulation_config["bbox_min"]
                self.articulation_offset = -np.float32(bbox_min) * articulation_scale
            # Add a small clearance
            self.articulation_offset[2] += 0.01

        return reconfigure

    def _initialize_articulations(self):
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            posxy_rotz_min = self._articulation_config["posxy_rotz_min"]
            posxy_rotz_max = self._articulation_config["posxy_rotz_max"]
            posxy_rotz = self._episode_rng.uniform(posxy_rotz_min, posxy_rotz_max)
            posz = (
                -self._articulation_config["scale"]
                * self._articulation_config["bbox_min"][2]
            )
            arti_root_pose = Pose(
                [posxy_rotz[0], posxy_rotz[1], posz],
                [np.sqrt(1 - posxy_rotz[2] ** 2), 0, 0, posxy_rotz[2]],
            )
            if self._articulation_config["flip"]:
                arti_root_pose = arti_root_pose * Pose(q=np.array([0, 1, 0, 0]))
            self._articulation.set_root_pose(arti_root_pose)

            # [[lmin, lmax]] = self._target_joint.get_limits()
            lmin, lmax = self._articulation_config["target_joint_range"]
            init_open_extent = self._episode_rng.uniform(
                0, self._articulation_config["init_open_extent_range"]
            )
            qpos = np.zeros(self._articulation.dof)
            for i in range(self._articulation.dof):
                qpos[i] = self._articulation.get_active_joints()[i].get_limits()[0][0]
            qpos[self._target_joint_idx] = lmin + (lmax - lmin) * init_open_extent
            self._articulation.set_qpos(qpos)

            self.target_qpos = (
                lmin + (lmax - lmin) * self._articulation_config["open_extent"]
            )
            self._get_handle_info_in_target_link()
        elif self._arti_mode == "faucet":
            p = np.array([0.1, 0, 0])
            p[:2] += self._episode_rng.uniform(-0.05, 0.05, [2])
            p[2] = self.articulation_offset[2]
            ori = self._episode_rng.uniform(0, np.pi / 3)
            q = euler2quat(0, 0, ori)
            self._articulation.set_pose(Pose(p, q))
            self._initialize_task()
        elif self._arti_mode == "laptop":
            p = np.array([0.15, 0, 0])
            p[:2] += self._episode_rng.uniform(-0.05, 0.05, [2])
            p[2] = self.articulation_offset[2]
            ori = self._episode_rng.uniform(-np.pi / 6, np.pi / 6)
            q = euler2quat(0, 0, ori)
            self._articulation.set_pose(Pose(p, q))

            # initialize task
            qmin, qmax, q_range = self.qmin_max_range
            self.init_angle = self._episode_rng.uniform(
                qmin + q_range / 6, qmin + q_range / 3
            )
            self.target_angle = qmin + q_range * 0.9
            # initialize qpos
            qpos = self._articulation_init_qpos.copy()
            qpos[self.target_joint_idx] = self.init_angle
            self._articulation.set_qpos(qpos)
            self.last_angle_diff = self.target_angle - self.current_angle
        elif self._arti_mode == "kitchen_pot":
            p = np.array([0.025, 0.05, 0.0])
            p[:2] += self._episode_rng.uniform(0.0, 0.05, [2])
            p[2] = self.articulation_offset[2]
            ori = np.pi / 2 + self._episode_rng.uniform(-np.pi / 12, np.pi / 12)
            q = euler2quat(0, 0, ori)
            self._articulation.set_pose(Pose(p, q))

            # initialize task
            qmin, qmax, q_range = self.qmin_max_range
            # self.init_angle = self._episode_rng.uniform(qmin + q_range / 6, qmin + q_range / 3)
            self.init_angle = 0
            self.target_angle = qmin + q_range * 0.9
            # initialize qpos
            qpos = self._articulation_init_qpos.copy()
            qpos[self.target_joint_idx] = self.init_angle
            self._articulation.set_qpos(qpos)
            self.last_angle_diff = self.target_angle - self.current_angle
        else:
            raise NotImplementedError

    def _initialize_task(self):
        self._set_target_link()
        self._set_init_and_target_angle()

        qpos = self._articulation_init_qpos.copy()
        qpos[self.target_joint_idx] = self.init_angle
        self._articulation.set_qpos(qpos)

        self.last_angle_diff = self.target_angle - self.current_angle

    def _set_target_link(self):
        n_switch_links = len(self.switch_link_names)
        idx = random_choice(np.arange(n_switch_links), self._episode_rng)

        self.target_link_name = self.switch_link_names[idx]
        self._target_link: sapien.Link = self.switch_links[idx]
        self.target_joint: sapien.Joint = self.switch_joints[idx]
        self.target_joint_idx = self._articulation.get_active_joints().index(
            self.target_joint
        )

        # x-axis is the revolute joint direction
        assert self.target_joint.type == "revolute", self.target_joint.type
        joint_pose = self.target_joint.get_global_pose().to_transformation_matrix()
        self.target_joint_axis = joint_pose[:3, 0]

        self.target_link_mesh = self.switch_links_mesh[idx]
        with np_random(self._episode_seed):
            self.target_pcd = self.target_link_mesh.sample(1000)
        # trimesh.PointCloud(self.target_pcd).show()

        # get points center
        target_link_frame_pos = self._target_link.get_pose().p
        # get points far, near from joint axis
        pointing_dir = (
            transform_points(
                self.target_link.pose.to_transformation_matrix(), self.target_pcd
            )
            - target_link_frame_pos
        )  # N X 3
        distance = np.sqrt(
            np.linalg.norm(pointing_dir, axis=1) ** 2
            - np.dot(pointing_dir, self.target_joint_axis) ** 2
        )
        ind_sort = np.argsort(distance)
        points_num = int(0.5 * self.target_pcd.shape[0])
        tip_index = ind_sort[-points_num:]
        self.tip_points = self.target_pcd[tip_index]

        # get tip, med, bottom keypoints
        if self.articulation_id in [5024, 5034]:
            norm_sw_pcd = apply_pose_to_points(
                self.target_pcd, self._target_joint.get_pose_in_parent().inv()
            )
            x_max, x_min = norm_sw_pcd[:, 0].max(), norm_sw_pcd[:, 0].min()
            center_pts, bottom_pts = (
                norm_sw_pcd[norm_sw_pcd[:, 0] - x_max > -0.01],
                norm_sw_pcd[norm_sw_pcd[:, 0] - x_min < 0.01],
            )
            center = np.mean(center_pts, axis=0)
            center[0] = x_max
            bottom = np.mean(bottom_pts, axis=0)
            bottom[0] = x_min
            tip = norm_sw_pcd[norm_sw_pcd[:, 1].argmin()]
            self._handle_keypoints = apply_pose_to_points(
                np.stack([tip, center, bottom]), self._target_joint.get_pose_in_parent()
            )
        else:
            xyzmax_val, xyzmin_val = self.target_pcd.max(0), self.target_pcd.min(0)
            xyzmax_id, xyzmin_id = np.argmax(self.target_pcd, 0), np.argmin(
                self.target_pcd, 0
            )
            bbox_size = xyzmax_val - xyzmin_val
            long_side = np.argmax(bbox_size)
            tip = self.target_pcd[xyzmax_id[long_side]]

            joint_pts_ids = np.where(
                np.linalg.norm(self.target_pcd[:, [0, 2]], axis=1) < 0.02
            )  # dist to joint axis < 0.01
            joint_pts = self.target_pcd[joint_pts_ids][:, 1]
            joint_pts_max = joint_pts.max()
            center = np.array([0, joint_pts_max, 0])
            bottom = np.array([0, self.target_pcd[:, 1].min(), 0])
            self._handle_keypoints = np.stack([tip, center, bottom])

        cmass_pose = self.target_link.pose * self.target_link.cmass_local_pose
        self.target_link_pos = cmass_pose.p

    def _set_init_and_target_angle(self):
        qmin, qmax = self.target_joint.get_limits()[0]
        qmin = 0 if np.isinf(qmin) else qmin
        qmax = np.pi / 2 if np.isinf(qmax) else qmax
        q_range = qmax - qmin
        self.init_angle = self._episode_rng.uniform(
            qmin + q_range / 6, qmin + q_range / 3
        )
        self.target_angle = qmin + q_range * 0.9

        # The angle to go
        self.target_angle_diff = self.target_angle - self.init_angle

    def _initialize_agent(self):
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            ### close to handle 1, half open gripper
            qpos = np.array([1.4, -1.053, -2.394, 1.662, 1.217, 1.05, -0.8, 0.04, 0.04])
        elif self._arti_mode in ["faucet", "laptop"]:
            qpos = np.array([-0.5, -0.143, 0, np.pi / 3, 0, 1.57, 1.57, 0.068, 0.068])
            # qpos = np.array([-0.488665, 0.340129, -1.14341, 1.18789, 0.346452, 1.73497, 0.7321, 0.068, 0.068])
        elif self._arti_mode == "kitchen_pot":
            qpos = np.array([-0.5, -0.143, 0, np.pi / 3, 0, 1.57, 1.57, 0.04, 0.04])
        else:
            raise NotImplementedError
        qpos[:-2] += self._episode_rng.normal(0, 0.02, len(qpos) - 2)
        # qpos = np.array([0., 0., 0., 0., 0., 0., 0., 0.04, 0.04])
        self._agent.reset(qpos)
        self._agent._robot.set_pose(Pose([-0.6, 0.4, 0]))
        ### get agent camera extrinsic, intrinsic para
        self.update_render()
        for cam_name, cam in self.agent._cameras.items():
            self.cam_para[cam_name + "_intrinsic"] = cam.get_intrinsic_matrix()
            cam_extrinsic_world_frame = cam.get_model_matrix()  # cam -> world
            if cam_name == "base_camera":
                base_frame = (
                    self.agent._robot.get_root_pose().inv().to_transformation_matrix()
                )  # world -> robot_base
            elif cam_name == "hand_camera":
                base_frame = (
                    self.agent.grasp_site.get_pose().inv().to_transformation_matrix()
                )  # world -> ee
            self.cam_para[cam_name + "_extrinsic_base_frame"] = (
                base_frame @ cam_extrinsic_world_frame
            )  # cam -> robot_base/ee

    def _load_articulations(self):
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            self._articulation = self._load_cabinet()
        elif self._arti_mode == "faucet":
            self._articulation = self._load_faucet()
            # Cache qpos to restore
            self._articulation_init_qpos = self._articulation.get_qpos()
        elif self._arti_mode == "laptop":
            self._articulation = self._load_laptop()
            # Cache qpos to restore
            self._articulation_init_qpos = self._articulation.get_qpos()
        elif self._arti_mode == "kitchen_pot":
            self._articulation = self._load_kitchenpot()
            # Cache qpos to restore
            self._articulation_init_qpos = self._articulation.get_qpos()
        else:
            raise NotImplementedError

        # set physical properties for all the joints
        self._joint_stiffness_range = (0.0, 0.0)
        if self._arti_mode in ["cabinet_door", "laptop"]:
            self._joint_friction_range = (0.8, 1.0)
            self._joint_damping_range = (20.0, 30.0)
        elif self._arti_mode in ["cabinet_drawer", "kitchen_pot"]:
            self._joint_friction_range = (0.2, 0.8)
            self._joint_damping_range = (70.0, 80.0)
        elif self._arti_mode == "faucet":
            self._joint_friction_range = (0.2, 0.8)
            self._joint_damping_range = (4.0, 5.0)  # faucet
        else:
            raise NotImplementedError

        joint_friction = self._episode_rng.uniform(
            self._joint_friction_range[0], self._joint_friction_range[1]
        )
        joint_stiffness = self._episode_rng.uniform(
            self._joint_stiffness_range[0], self._joint_stiffness_range[1]
        )
        joint_damping = self._episode_rng.uniform(
            self._joint_damping_range[0], self._joint_damping_range[1]
        )
        for joint in self._articulation.get_active_joints():
            joint.set_friction(joint_friction)
            joint.set_drive_property(joint_stiffness, joint_damping)
        ignore_collision(self._articulation)

        self._set_switch_links()

    def _load_cabinet(self):
        loader = self._scene.create_urdf_loader()
        loader.load_multiple_collisions_from_file = self._articulation_config[
            "multiple_collisions"
        ]
        loader.scale = self._articulation_config["scale"]
        loader.fix_root_link = True

        articulation_path = (
            ASSET_DIR
            / f"partnet_mobility_dataset/{self._arti_mode}/{self.articulation_id}"
        )
        urdf_path = articulation_path / "mobility_cvx.urdf"
        assert (
            urdf_path.exists()
        ), f"{urdf_path} is not found. Please download Partnet-Mobility (ManiSkill2022) first."
        config = {
            "material": self._physical_materials["default"],
            # "density": 8e3
        }
        articulation = loader.load(str(urdf_path), config=config)
        articulation.set_name("cabinet")
        return articulation

    def _load_faucet(self):
        loader = self._scene.create_urdf_loader()
        loader.scale = self.articulation_scale
        loader.fix_root_link = True

        model_dir = (
            ASSET_DIR
            / f"partnet_mobility_dataset/{self._arti_mode}/{self.articulation_id}"
        )
        urdf_path = model_dir / "mobility_cvx.urdf"
        loader.load_multiple_collisions_from_file = True
        assert (
            urdf_path.exists()
        ), f"{urdf_path} is not found. Please download Partnet-Mobility (ManiSkill2022) first."

        density = self._articulation_config.get("density", 8e3)
        articulation = loader.load(str(urdf_path), config={"density": density})
        articulation.set_name("faucet")

        cam = self._scene.add_camera("dummy_camera", 1, 1, 1, 0.01, 10)
        self._scene.remove_camera(cam)
        set_articulation_render_material(
            articulation, color=hex2rgba("#AAAAAA"), metallic=1, roughness=0.4
        )

        return articulation

    def _load_laptop(self):
        loader = self._scene.create_urdf_loader()
        loader.scale = self.articulation_scale
        loader.fix_root_link = True

        model_dir = (
            ASSET_DIR
            / f"partnet_mobility_dataset/{self._arti_mode}/{self.articulation_id}"
        )
        urdf_path = model_dir / "mobility.urdf"
        loader.load_multiple_collisions_from_file = True
        assert (
            urdf_path.exists()
        ), f"{urdf_path} is not found. Please download Partnet-Mobility Dataset first."

        density = self._articulation_config.get("density", 8e3)
        articulation = loader.load(str(urdf_path), config={"density": density})
        articulation.set_name("laptop_kitchenpot")

        cam = self._scene.add_camera("dummy_camera", 1, 1, 1, 0.01, 10)
        self._scene.remove_camera(cam)
        set_articulation_render_material(
            articulation, color=hex2rgba("#AAAAAA"), metallic=1, roughness=0.4
        )

        return articulation

    def _load_kitchenpot(self):
        loader = self._scene.create_urdf_loader()
        loader.scale = self.articulation_scale
        loader.fix_root_link = True

        model_dir = (
            ASSET_DIR
            / f"partnet_mobility_dataset/{self._arti_mode}/{self.articulation_id}"
        )
        urdf_path = model_dir / "mobility.urdf"
        loader.load_multiple_collisions_from_file = True
        assert (
            urdf_path.exists()
        ), f"{urdf_path} is not found. Please download Partnet-Mobility Dataset first."

        density = self._articulation_config.get("density", 1e3)
        articulation = loader.load(str(urdf_path), config={"density": density})
        articulation.set_name("laptop_kitchenpot")

        cam = self._scene.add_camera("dummy_camera", 1, 1, 1, 0.01, 10)
        self._scene.remove_camera(cam)
        set_articulation_render_material(
            articulation, color=hex2rgba("#AAAAAA"), metallic=1, roughness=0.4
        )

        return articulation

    def _set_switch_links(self):
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            all_links = self._articulation.get_links()
            all_active_joints = self._articulation.get_active_joints()

            # set target link & joint, handle & door of target_link
            self._target_joint_idx = self._articulation_config["target_joint_idx"]
            self._target_joint = all_active_joints.pop(self._target_joint_idx)
            self._target_link = self._target_joint.get_child_link()
            all_links.remove(self._target_link)

            handle_visid = []
            door_visid = []
            for visual_body in self._target_link.get_visual_bodies():
                if "handle" in visual_body.get_name():
                    handle_visid.append(visual_body.get_visual_id())
                else:
                    door_visid.append(visual_body.get_visual_id())
            self.handle_visid = handle_visid
            self.door_visid = door_visid

            other_handle_visid = []
            other_door_visid = []
            for ajoint in all_active_joints:
                door_link = ajoint.get_child_link()
                all_links.remove(door_link)
                for visual_body in door_link.get_visual_bodies():
                    if "handle" in visual_body.get_name():
                        other_handle_visid.append(visual_body.get_visual_id())
                        visual_body.set_visibility(
                            1.0 if self.other_handle_visible else 0.0
                        )
                    else:
                        other_door_visid.append(visual_body.get_visual_id())
            self.other_handle_visid = other_handle_visid
            self.other_door_visid = other_door_visid

            cabinet_visid = []
            for link in all_links:
                for visual_body in link.get_visual_bodies():
                    cabinet_visid.append(visual_body.get_visual_id())
            self.cabinet_visid = cabinet_visid
        elif self._arti_mode == "faucet":
            all_links = self._articulation.get_links()
            all_joints = self._articulation.get_joints()

            switch_link_names = []
            fix_link_visid = []
            # self.fix_links = []
            # self.fix_links_mesh = []
            for semantic in self._articulation_config["semantics"]:
                if semantic[2] == "switch":
                    switch_link_names.append(semantic[0])
                else:
                    fix_link = get_entity_by_name(all_links, semantic[0])
                    # self.fix_links.append(fix_link)
                    # link_mesh = get_actor_mesh(fix_link, False)
                    # self.fix_links_mesh.append(link_mesh)
                    for visual_body in fix_link.get_visual_bodies():
                        fix_link_visid.append(visual_body.get_visual_id())
            self.fix_link_visid = fix_link_visid

            if len(switch_link_names) == 0:
                raise RuntimeError(self.articulation_id)
            self.switch_link_names = switch_link_names

            self.switch_links = []
            self.switch_links_visid = []
            self.switch_links_mesh = []
            self.switch_joints = []
            for name in self.switch_link_names:
                link = get_entity_by_name(all_links, name)
                self.switch_links.append(link)
                for visual_body in link.get_visual_bodies():
                    self.switch_links_visid.append(visual_body.get_visual_id())

                # cache mesh
                link_mesh = get_actor_mesh(link, False)
                self.switch_links_mesh.append(link_mesh)

                # hardcode
                joint = all_joints[link.get_index()]
                # joint.set_friction(0.1)
                # joint.set_drive_property(0.0, 2.0)
                self.switch_joints.append(joint)
        elif self._arti_mode == "laptop":
            all_links = self._articulation.get_links()
            all_joints = self._articulation.get_joints()
            switch_link_names = []
            fix_link_visid = []
            for semantic in self._articulation_config["semantics"]:
                if semantic[2] == "switch":
                    switch_link_names.append(semantic[0])
                else:
                    fix_link = get_entity_by_name(all_links, semantic[0])
                    for visual_body in fix_link.get_visual_bodies():
                        fix_link_visid.append(visual_body.get_visual_id())
            self.fix_link_visid = fix_link_visid
            self.fix_link = fix_link
            if len(switch_link_names) == 0:
                raise RuntimeError(self.articulation_id)
            self.switch_link_names = switch_link_names
            self.switch_links = []
            self.switch_links_visid = []
            switch_links_mesh = []
            switch_joints = []
            for name in self.switch_link_names:
                link = get_entity_by_name(all_links, name)
                self.switch_links.append(link)
                for visual_body in link.get_visual_bodies():
                    self.switch_links_visid.append(visual_body.get_visual_id())
                # cache mesh
                link_mesh = get_actor_mesh(link, False)
                switch_links_mesh.append(link_mesh)
                # hardcode
                joint = all_joints[link.get_index()]

                switch_joints.append(joint)

            n_switch_links = len(self.switch_link_names)
            idx = random_choice(np.arange(n_switch_links), self._episode_rng)
            self.target_link_name = self.switch_link_names[idx]
            self._target_link: sapien.Link = self.switch_links[idx]
            self.target_link_mesh: trimesh.Trimesh = switch_links_mesh[idx]
            self.target_joint: sapien.Joint = switch_joints[idx]
            self.target_joint_idx = self._articulation.get_active_joints().index(
                self.target_joint
            )

            joint_pose = self.target_joint.get_global_pose().to_transformation_matrix()
            self.target_joint_axis = joint_pose[:3, 0]
            with np_random(self._episode_seed):
                self.target_pcd = self.target_link_mesh.sample(1000)

            distance = self.target_pcd[:, 1]
            ind_sort = np.argsort(distance)
            points_num = int(0.2 * self.target_pcd.shape[0])
            tip_index = ind_sort[-points_num:]
            self.tip_points = self.target_pcd[tip_index]

            qmin, qmax = self.target_joint.get_limits()[0]
            qmin = 0 if np.isinf(qmin) else qmin
            qmax = np.pi / 2 if np.isinf(qmax) else qmax
            q_range = qmax - qmin
            self.qmin_max_range = [qmin, qmax, q_range]
        elif self._arti_mode == "kitchen_pot":
            all_links = self._articulation.get_links()
            all_joints = self._articulation.get_joints()
            switch_link_names = []
            fix_link_visid = []
            for semantic in self._articulation_config["semantics"]:
                if semantic[2] == "switch":
                    switch_link_names.append(semantic[0])
                else:
                    fix_link = get_entity_by_name(all_links, semantic[0])
                    for visual_body in fix_link.get_visual_bodies():
                        fix_link_visid.append(visual_body.get_visual_id())
            self.fix_link_visid = fix_link_visid
            self.fix_link = fix_link
            if len(switch_link_names) == 0:
                raise RuntimeError(self.articulation_id)
            self.switch_link_names = switch_link_names

            # set target link, joint
            n_switch_links = len(self.switch_link_names)
            idx = random_choice(np.arange(n_switch_links), self._episode_rng)
            self.target_link_name = self.switch_link_names[idx]
            self._target_link: sapien.Link = get_entity_by_name(
                all_links, self.target_link_name
            )
            self.target_joint: sapien.Joint = all_joints[self._target_link.get_index()]
            self.target_joint_idx = self._articulation.get_active_joints().index(
                self.target_joint
            )

            self.lid_visid, lid_meshes = [], []
            self.handle_visid, handle_meshes = [], []

            for visual_body in self._target_link.get_visual_bodies():
                vis_id = visual_body.get_visual_id()
                if "handle" in visual_body.get_name():
                    self.handle_visid.append(vis_id)
                    # cache mesh
                    handle_meshes.extend(get_visual_body_meshes(visual_body))
                else:
                    self.lid_visid.append(vis_id)
                    # cache mesh
                    lid_meshes.extend(get_visual_body_meshes(visual_body))
            self.handle_mesh = merge_meshes(handle_meshes)
            self.lid_mesh = merge_meshes(lid_meshes)

            ### get handle grasp
            handle_convex_mesh = trimesh.convex.convex_hull(self.handle_mesh)
            handle_pcd = handle_convex_mesh.sample(200)
            self.handle_pcd = handle_pcd
            handle_pcd_world = apply_pose_to_points(
                handle_pcd, self._target_link.get_pose()
            )
            bbox_size = (handle_pcd_world.max(0) - handle_pcd_world.min(0)) / 2
            center = (handle_pcd_world.max(0) + handle_pcd_world.min(0)) / 2 + np.array(
                [0, 0.04, 0]
            )
            forward = np.array([0, -1, 0])
            if bbox_size[0] > bbox_size[2]:
                flat = np.array([0, 0, 1])
            else:
                flat = np.array([-1, 0, 0])
            grasp_pose = (
                self.agent.build_grasp_pose(forward, flat, center),
                self.agent.build_grasp_pose(forward, -flat, center),
            )
            self.handle_grasp = grasp_pose

            qmin, qmax = self.target_joint.get_limits()[0]
            qmin = 0 if np.isinf(qmin) else qmin
            qmax = 0.2 if np.isinf(qmax) else qmax
            q_range = qmax - qmin
            self.qmin_max_range = [qmin, qmax, q_range]
        else:
            raise NotImplementedError

    def _get_handle_info_in_target_link(self):
        """
        build a mesh and a point cloud of handle in target link, compute grasp poses of handle
        """
        handle_meshes = []
        others_meshes = []
        for visual_body in self._target_link.get_visual_bodies():
            for render_shape in visual_body.get_render_shapes():
                vertices = apply_pose_to_points(
                    render_shape.mesh.vertices * visual_body.scale,
                    visual_body.local_pose,
                )
                shape_mesh = np2mesh(vertices, render_shape.mesh.indices.reshape(-1, 3))
                if "handle" in visual_body.get_name():
                    handle_meshes.append(shape_mesh)
                else:
                    others_meshes.append(shape_mesh)
        handle_mesh = merge_mesh(handle_meshes)

        ### get handle grasp
        handle_convex_mesh = trimesh.convex.convex_hull(o3d_to_trimesh(handle_mesh))
        handle_pcd = handle_convex_mesh.sample(200)
        handle_pcd_world = apply_pose_to_points(
            handle_pcd, self._target_link.get_pose()
        )
        bbox_size = (handle_pcd_world.max(0) - handle_pcd_world.min(0)) / 2
        center = (handle_pcd_world.max(0) + handle_pcd_world.min(0)) / 2
        R = self._target_link.get_pose().to_transformation_matrix()[:3, :3]
        # choose the axis closest to X
        idx = np.argmax(np.abs(R[0]))
        forward = R[:3, idx]
        if forward[0] < 0:
            forward *= -1
        if bbox_size[1] > bbox_size[2]:
            flat = np.array([0, 0, 1])
        else:
            flat = np.cross(forward, np.array([0, 0, 1]))
        grasp_pose = (
            self._target_link.get_pose().inv()
            * self._agent.build_grasp_pose(forward, flat, center),
            self._target_link.get_pose().inv()
            * self._agent.build_grasp_pose(forward, -flat, center),
        )

        # trans_handle_mesh = handle_mesh.transform(self._target_link.get_pose().to_transformation_matrix())
        self.target_link_mesh = o3d_to_trimesh(handle_mesh)
        self.target_pcd = self.target_link_mesh.sample(300)  ## target_link frame

        # get tip, med, bottom keypoints
        xyzmax_val, xyzmin_val = self.target_pcd.max(0), self.target_pcd.min(0)
        xyzmax_id, xyzmin_id = np.argmax(self.target_pcd, 0), np.argmin(
            self.target_pcd, 0
        )
        bbox_size = xyzmax_val - xyzmin_val
        long_side = np.argmax(bbox_size)
        tip, bottom = (
            self.target_pcd[xyzmax_id[long_side]],
            self.target_pcd[xyzmin_id[long_side]],
        )
        center = self.target_pcd.mean(0)
        if self.articulation_id in [0, 1]:
            center[1] = xyzmax_val[1]
        else:
            center[2] = xyzmax_val[2]
        self._handle_keypoints = np.stack([tip, center, bottom])

        self._handle_info = OrderedDict()
        self._handle_info["mesh"] = trimesh_to_o3d(handle_convex_mesh)
        self._handle_info["pcd"] = handle_pcd
        self._handle_info["grasp"] = grasp_pose
        self._handle_info["center"] = handle_pcd.mean(0)

    def _get_keypart_vid(self, part_name):
        """
        get the visual id from each link in (cabinet, door, robot, gripper)
        """
        assert part_name in ["cabinet", "robot", "gripper"]

        vis_ids = []
        for link in self._part_info[part_name]["links"]:
            for visual_body in link.get_visual_bodies():
                vis_ids.append(visual_body.get_visual_id())
        self._part_info[part_name]["vis_id"] = vis_ids
        if part_name == "cabinet":
            hanele_ids = []
            others_ids = []
            for visual_body in self._target_link.get_visual_bodies():
                if "handle" in visual_body.get_name():
                    hanele_ids.append(visual_body.get_visual_id())
                else:
                    others_ids.append(visual_body.get_visual_id())
            self._part_info["handle"]["vis_id"] = hanele_ids
            self._part_info["door"]["vis_id"] = others_ids

    def _get_keypart_info(self, part_name="cabinet"):
        """
        get the visual id, sampled pcd from each link in (cabinet, door, robot, gripper)
        """
        assert part_name in ["cabinet", "robot", "gripper"]

        idx = 0
        for link in self._part_info[part_name]["links"]:
            meshes = []
            for visual_body in link.get_visual_bodies():
                if visual_body.type == "mesh":  ### for gripper finger pad
                    scale = visual_body.scale
                elif visual_body.type == "box":
                    scale = visual_body.half_lengths
                elif visual_body.type == "sphere":  ### from maniskill actor_to_o3d
                    scale = visual_body.radius
                else:
                    raise TypeError
                for render_shape in visual_body.get_render_shapes():
                    if render_shape.mesh.indices.reshape(-1, 3).shape[
                        0
                    ] < 4 or check_coplanar(render_shape.mesh.vertices * scale):
                        continue
                    vertices = apply_pose_to_points(
                        render_shape.mesh.vertices * scale, visual_body.local_pose
                    )
                    meshes.append(
                        np2mesh(vertices, render_shape.mesh.indices.reshape(-1, 3))
                    )
            if len(meshes) > 0:
                link_mesh = merge_mesh(meshes)  # link frame
                num_faces = 1024
                if np.asarray(link_mesh.triangles).shape[0] > num_faces:
                    link_mesh = link_mesh.simplify_vertex_clustering(
                        voxel_size=0.02,
                        contraction=o3d.geometry.SimplificationContraction.Average,
                    )
                link_trimesh = o3d_to_trimesh(link_mesh)
                link_pcd_gt = link_trimesh.sample(
                    self._part_info[part_name]["pts_num"][idx]
                )  ## link frame
                self._part_info[part_name]["pcd"][link.get_name()] = link_pcd_gt
            idx += 1

    def _get_obs_state_dict(self) -> OrderedDict:
        state_dict = OrderedDict()
        ## propriception
        arm_qpos = self.agent._robot.get_qpos()
        state_dict["qpos"] = arm_qpos.astype(np.float32)
        ee_pos_robot = (
            self._agent._robot.get_root_pose().inv() * self._agent.grasp_site.get_pose()
        ).p
        state_dict["ee_pos_base"] = ee_pos_robot.astype(np.float32)

        return state_dict

    def get_handcam_seg(self):
        self.update_render()
        ego_cam = self._agent._cameras["hand_camera"]
        ego_cam.take_picture()
        visual_id_seg = get_camera_seg(ego_cam)[..., 0]
        cam_shape = visual_id_seg.shape
        cam_seg = (
            np.ones(cam_shape, dtype=np.uint8) * 5
        )  # gripper, handle, door, cabient, ..., ..., other
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            seg_visids = [
                self.handle_visid,
                self.door_visid + self.other_door_visid,
                self.cabinet_visid,
            ]
            seg_id = [0, 1, 2]
            for id, seg_visid in enumerate(seg_visids):
                mask = np.zeros(cam_shape, dtype=np.bool)
                for visual_id in seg_visid:
                    mask = mask | (visual_id_seg == visual_id)
                cam_seg[mask] = seg_id[id]
        elif self._arti_mode == "faucet":
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [3, 4]
            for id, seg_visid in enumerate(seg_visids):
                mask = np.zeros(cam_shape, dtype=np.bool)
                for visual_id in seg_visid:
                    mask = mask | (visual_id_seg == visual_id)
                cam_seg[mask] = seg_id[id]
        else:
            raise NotImplementedError
        return cam_seg

    def get_handcam_rgbdseg(self):
        rgbdseg = OrderedDict()
        ### get visual info
        self.update_render()
        ego_cam = self._agent._cameras["hand_camera"]
        ego_cam.take_picture()
        ## process rgb
        cam_rgb = get_camera_rgb(ego_cam)  # (H, W, 3)
        cam_H, cam_W = cam_rgb.shape[:2]
        rgbdseg["rgb"] = cam_rgb.transpose((2, 0, 1)).astype(np.uint8)  # (3, H, W)
        ## process depth & point clouds
        position = get_texture(ego_cam, "Position")
        cam_depth = -position[..., [2]]  # [H, W, 1]
        rgbdseg["depth"] = (cam_depth.transpose((2, 0, 1))).astype(
            np.float32
        )  # (1, H, W)

        visual_id_seg = get_camera_seg(ego_cam)[..., 0]
        cam_seg = (
            np.ones((cam_H, cam_W), dtype=np.bool) * 5
        )  # gripper, handle, door, cabient, ..., ..., other
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            seg_visids = [
                self.handle_visid,
                self.door_visid + self.other_door_visid,
                self.cabinet_visid,
            ]
            seg_id = [0, 1, 2]
            for id, seg_visid in enumerate(seg_visids):
                mask = np.zeros((cam_H, cam_W), dtype=np.bool)
                for visual_id in seg_visid:
                    mask = mask | (visual_id_seg == visual_id)
                cam_seg[mask] = seg_id[id]
            rgbdseg["seg"] = convert_np_bool_to_float(cam_seg).astype(
                np.uint8
            )  # (H, W)
        elif self._arti_mode == "faucet":
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [3, 4]
            for id, seg_visid in enumerate(seg_visids):
                mask = np.zeros((cam_H, cam_W), dtype=np.bool)
                for visual_id in seg_visid:
                    mask = mask | (visual_id_seg == visual_id)
                cam_seg[mask] = seg_id[id]
            rgbdseg["seg"] = convert_np_bool_to_float(cam_seg).astype(
                np.uint8
            )  # (H, W)
        else:
            raise NotImplementedError
        return rgbdseg

    def get_handsensor_rgb(self):
        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        rgb = (hand_sensor.get_rgb().transpose(2, 0, 1) * 255).astype(
            np.uint8
        )  # (3, H, W)
        return rgb

    def get_handsensor_rgbdseg(self):
        state_egopts = OrderedDict()
        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        hand_sensor.compute_depth()

        state_egopts["rgb"] = (hand_sensor.get_rgb().transpose(2, 0, 1) * 255).astype(
            np.uint8
        )  # (3, H, W)
        state_egopts["depth"] = hand_sensor.get_depth().astype(np.float32)

        # process seg groundtruth
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        cam_H, cam_W = visual_id_seg.shape
        # gripper, handle, door, cabient, switch_link, fix_link, other
        cam_seg = np.ones((cam_H, cam_W), dtype=np.bool) * 5
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            seg_visids = [
                self.handle_visid,
                self.door_visid + self.other_door_visid,
                self.cabinet_visid,
            ]
            seg_id = [0, 1, 2]
        elif self._arti_mode == "faucet":
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [3, 4]
        elif self._arti_mode == "laptop":
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [5, 6]
        elif self._arti_mode == "kitchen_pot":
            seg_visids = [self.handle_visid, self.lid_visid, self.fix_link_visid]
            seg_id = [7, 8, 9]
        else:
            raise NotImplementedError
        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros((cam_H, cam_W), dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
        state_egopts["seg"] = convert_np_bool_to_float(cam_seg).astype(
            np.uint8
        )  # (H, W)
        return state_egopts

    def get_handsensor_pc(self):
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        hand_sensor.compute_depth()
        # process point cloud
        cam_xyz = hand_sensor.get_pointcloud(False)  # [H*W, 3]
        ## cam to world
        T_cam2world = sensor_rgb.get_model_matrix()
        world_xyz = transform_points(
            T_cam2world, cam_xyz * np.array([1, -1, -1])
        )  # [H*W, 3]
        return world_xyz

    def get_egostereo_segpts(self):
        # get state
        state_egopts = self._get_obs_state_dict()

        ## ee points in ee frame
        if self.add_eepadpts:
            ee_cords = self._agent.sample_ee_coords(num_sample=20).reshape(
                -1, 3
            )  # [40, 3], world frame
            eepad_pts = apply_pose_to_points(
                ee_cords, self.agent.grasp_site.get_pose().inv()
            )  # [40, 3], ee frame
            eepad_pts_t = eepad_pts
            eepat_pts_label = np.zeros(
                (eepad_pts.shape[0], self.num_classes + 1)
            )  # (40, C+1)
            eepat_pts_label[:, -1] = 1.0
            eepad_ptsC = np.concatenate(
                (eepad_pts_t, eepat_pts_label), axis=1
            )  # (40, 3+C+1)

        # get visual info
        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        hand_sensor.compute_depth()

        # process seg groundtruth
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        cam_H, cam_W = visual_id_seg.shape
        cam_seg = (
            np.ones((cam_H, cam_W), dtype=np.bool) * 5
        )  # gripper, handle, door, cabient, switch_link, fix_link, other
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            seg_visids = [
                self.handle_visid,
                self.door_visid + self.other_door_visid,
                self.cabinet_visid,
            ]
            seg_id = [0, 1, 2]
        elif self._arti_mode == "faucet":
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [3, 4]
        else:
            raise NotImplementedError

        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros((cam_H, cam_W), dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
        if self.add_eepadpts:
            state_egopts["seg"] = convert_np_bool_to_float(cam_seg).astype(
                np.float32
            )  # (H, W)
        else:
            state_egopts["seg"] = convert_np_bool_to_float(cam_seg).astype(
                np.uint8
            )  # (H, W)

        # process seg model prediction
        cam_rgb_tensor = (
            torch.utils.dlpack.from_dlpack(hand_sensor.get_rgba_dl_tensor())
            .clone()[None, ..., :3]
            .clip(0, 1)
            .permute((0, 3, 1, 2))
        )  # [1,3,H,W]
        cam_depth_tensor = torch.utils.dlpack.from_dlpack(
            hand_sensor.get_depth_dl_tensor()
        ).clone()[
            None, None
        ]  # [1,1,H,W]

        # predict seg & sample K points
        seg_map = self.segmodel.predict(
            torch.cat([cam_rgb_tensor, cam_depth_tensor], dim=1)
        )[
            0
        ]  # (C, H, W)
        # seg_map_noise = (
        #     seg_map + torch.randn(seg_map.shape, device=seg_map.device) * 0.1
        # )
        seg_mc_label = torch.argmax(seg_map, dim=0).reshape(-1)  # (H*W)
        if self.add_eepadpts:
            seg_onehot_label = (
                F.one_hot(seg_mc_label, num_classes=self.num_classes + 1)
                .cpu()
                .detach()
                .numpy()
            )  # (H*W, C+1)
        else:
            seg_onehot_label = (
                F.one_hot(seg_mc_label, num_classes=self.num_classes)
                .cpu()
                .detach()
                .numpy()
            )  # (H*W, C)
        # if self._arti_mode == "cabinet_drawer":
        #     seg_onehot_label *= np.array([-1, -1, -1, 1, 1, 1])

        # process point cloud
        cam_xyz = hand_sensor.get_pointcloud(False)  # [H*W, 3]
        ## cam to world
        T_cam2world = sensor_rgb.get_model_matrix()
        world_xyz = transform_points(
            T_cam2world, cam_xyz * np.array([1, -1, -1])
        )  # [H*W, 3]
        T_world2ee = (
            self.agent.grasp_site.get_pose().inv().to_transformation_matrix()
        )  ## world => ee
        if self.sample_mode == "full_downsample":
            # ## crop the pts & uniform downsample to num pts
            sample_indices = self.sampler.sampling(world_xyz, self._episode_rng)
            ds_world_xyz = world_xyz[sample_indices]
            ds_onehot_label = seg_onehot_label[sample_indices]
            ## world to ee
            ds_ee_xyz = transform_points(T_world2ee, ds_world_xyz)
            # add noise on downsampled ee_xyz
            # with np_random(self._episode_seed):
            #     ds_ee_xyz = ds_ee_xyz + np.random.normal(loc=0, scale=0.01, size=ds_ee_xyz.shape)
            segptsC = np.concatenate((ds_ee_xyz, ds_onehot_label), axis=1)  # (SN, 3+C)
            if self.add_eepadpts:
                state_egopts["segsampled_ptsC"] = np.concatenate(
                    (segptsC, eepad_ptsC), axis=0
                ).astype(
                    np.float32
                )  # (SN+40, 3+C+1)
            else:
                state_egopts["segsampled_ptsC"] = segptsC.astype(
                    np.float32
                )  # (SN, 3+C)
        else:
            # pts_index = self.sampler.sampling(world_xyz, seg_map, self._episode_rng)
            if self.sample_mode == "score_sample":
                pts_index = self.sampler.sampling(seg_map)
            elif self.sample_mode == "random_sample":
                seg_mc_map = seg_mc_label.cpu().numpy()
                pts_index = self.sampler.sampling(
                    seg_mc_map, self.num_classes, self._episode_rng
                )
            elif self.sample_mode == "fps_sample":
                seg_mc_map = seg_mc_label.cpu().numpy()
                pts_index = self.sampler.sampling(
                    world_xyz, self.num_classes, seg_mc_map, self._episode_rng
                )
            else:
                raise NotImplementedError
            world_segsampled_xyz = world_xyz[pts_index]  # (C*sample_num, 3)
            sampled_onehot_label = seg_onehot_label[pts_index]  # (C*sample_num, C)
            ee_segsampled_xyz = transform_points(
                T_world2ee, world_segsampled_xyz
            )  # (C*sample_num, 3)
            # add noise on downsampled ee_xyz
            # with np_random(self._episode_seed):
            #     ee_segsampled_xyz = ee_segsampled_xyz + np.random.normal(loc=0, scale=0.01, size=ee_segsampled_xyz.shape)
            segptsC = np.concatenate(
                (ee_segsampled_xyz, sampled_onehot_label), axis=1
            )  # (C*sample_num, 3+C)
            if self.add_eepadpts:
                state_egopts["segsampled_ptsC"] = np.concatenate(
                    (segptsC, eepad_ptsC), axis=0
                ).astype(
                    np.float32
                )  # (C*sample_num+40, 3+C+1)
            else:
                state_egopts["segsampled_ptsC"] = segptsC.astype(
                    np.float32
                )  # (C*sample_num, 3+C)
        return state_egopts

    def get_egostereo_segpts_gt(self):
        # get state
        state_egopts = self._get_obs_state_dict()

        ## ee points in ee frame
        # ee_cords = self._agent.sample_ee_coords(num_sample=20).reshape(
        #     -1, 3
        # )  # [40, 3], world frame
        # eepad_pts = apply_pose_to_points(
        #     ee_cords, self.agent.grasp_site.get_pose().inv()
        # )  # [40, 3], ee frame
        # eepad_pts_t = eepad_pts
        # eepat_pts_label = np.zeros((eepad_pts.shape[0], self.num_classes + 1))  # (40, C+1)
        # eepat_pts_label[:, -1] = 1.0
        # eepad_ptsC = np.concatenate((eepad_pts_t, eepat_pts_label), axis=1)  # (40, 3+C+1)

        # get visual info
        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture(infrared_only=False)
        hand_sensor.compute_depth()

        # get set gt
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        cam_H, cam_W = visual_id_seg.shape
        cam_seg = (
            np.ones((cam_H, cam_W), dtype=np.bool) * 5
        )  # gripper, handle, door, cabient, switch_link, fix_link, other
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            seg_visids = [
                self.handle_visid,
                self.door_visid + self.other_door_visid,
                self.cabinet_visid,
            ]
            seg_id = [0, 1, 2]
        elif self._arti_mode == "faucet":
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [3, 4]
        elif self._arti_mode == "laptop":
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [5, 6]
        elif self._arti_mode == "kitchen_pot":
            seg_visids = [self.handle_visid, self.lid_visid, self.fix_link_visid]
            seg_id = [7, 8, 9]
        else:
            raise NotImplementedError
        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros((cam_H, cam_W), dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
        state_egopts["seg"] = convert_np_bool_to_float(cam_seg).astype(
            np.uint8
        )  # (H, W)

        # get C * sample_num indices from cam_seg
        seg_label = np.eye(self.num_classes)
        seg_onehot_label = seg_label[cam_seg.reshape(-1)].astype(np.float_)  # (H*W, C)

        indices = []
        for ind in range(seg_onehot_label.shape[1]):
            one_indices = np.nonzero(seg_onehot_label[:, ind])[0]  # S
            ones_num = one_indices.shape[0]
            if ones_num == 0:
                cur_indices = np.zeros(self.sample_num, dtype=np.int_)
            elif ones_num < self.sample_num:
                repeat_times = self.sample_num // ones_num
                remain_num = self.sample_num - repeat_times * ones_num
                cur_indices = np.concatenate(
                    (one_indices.repeat(repeat_times), one_indices[:remain_num])
                )
            else:
                cur_indices = np.random.choice(one_indices, self.sample_num)
            indices.append(cur_indices)
        pts_index = np.concatenate(indices)  # (C*sample_num)

        # get sampled point cloud (ee frame)
        cam_xyz = hand_sensor.get_pointcloud(with_rgb=False)
        cam_xyz = cam_xyz * np.array([1, -1, -1])
        ds_cam_xyz = cam_xyz[pts_index]  # (C*sample_num, 3)
        T_cam2ee = self.cam_para["hand_camera_extrinsic_base_frame"]
        ds_ee_xyz = transform_points(T_cam2ee, ds_cam_xyz)  # (C*sample_num, 3)
        # get sampled one-hot label
        sampled_onehot_label = seg_onehot_label[pts_index]  # (C*sample_num, C)
        segptsC = np.concatenate(
            (ds_ee_xyz, sampled_onehot_label), axis=1
        )  # (C*sample_num, 3+C)
        state_egopts["segsampled_ptsC"] = segptsC.astype(
            np.float_
        )  # (C*sample_num+40, 3+C)
        # state_egopts["segsampled_ptsC"] = np.concatenate((segptsC, eepad_ptsC), axis=0).astype(
        #     np.float_)  # (C*sample_num+40, 3+C+1)
        return state_egopts

    def get_egostereo_rgbd_seg_kpts(self):
        # get state
        state_egopts = self._get_obs_state_dict()

        # get visual info
        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        hand_sensor.compute_depth()
        cam_intrin = self.cam_para["hand_camera_intrinsic"]
        T_cam2ee = self.cam_para["hand_camera_extrinsic_base_frame"]

        # process seg groundtruth
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        cam_H, cam_W = visual_id_seg.shape
        cam_seg = (
            np.ones((cam_H, cam_W), dtype=np.bool) * 5
        )  # gripper, handle, door, cabient, switch_link, fix_link, other
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            seg_visids = [
                self.handle_visid,
                self.door_visid + self.other_door_visid,
                self.cabinet_visid,
            ]
            seg_id = [0, 1, 2]
        elif self._arti_mode == "faucet":
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [3, 4]
        else:
            raise NotImplementedError

        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros((cam_H, cam_W), dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
        state_egopts["seg"] = convert_np_bool_to_float(cam_seg).astype(
            np.uint8
        )  # (H, W)

        # process seg model prediction
        cam_rgb_tensor = (
            torch.utils.dlpack.from_dlpack(hand_sensor.get_rgba_dl_tensor())
            .clone()[None, ..., :3]
            .clip(0, 1)
            .permute((0, 3, 1, 2))
        )  # [1,3,H,W]
        cam_depth_tensor = torch.utils.dlpack.from_dlpack(
            hand_sensor.get_depth_dl_tensor()
        ).clone()[
            None, None
        ]  # [1,1,H,W]

        # predict seg & sample K points
        heatmap_out = self.kpt_model.forward(
            torch.cat([cam_rgb_tensor, cam_depth_tensor], dim=1)
        )  # (1, J*D', H', W')
        uvz_pred_norm = (
            soft_argmax(heatmap_out, 3, self.out_shape)[0].detach().cpu().numpy()
        )  # (J, 3)
        uvz_pred = uvz_pred_norm * (self.kpts_max - self.kpts_min) + self.kpts_min
        # state_egopts["uvz_pred"] = uvz_pred.astype(np.float32)  # (J, 3)
        xyz_cam_pred = uv2xyz(uvz_pred, cam_intrin)
        xyz_ee_pred = transform_points(T_cam2ee, xyz_cam_pred)
        state_egopts["xyz_ee_pred"] = xyz_ee_pred.astype(np.float32)  # [J, 3]
        return state_egopts

    def get_egostereo_rgbd_seg_kpts_gt(self):
        # get state
        state_egopts = self._get_obs_state_dict()

        # get visual info
        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()

        # process seg groundtruth
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        cam_H, cam_W = visual_id_seg.shape
        cam_seg = (
            np.ones((cam_H, cam_W), dtype=np.bool) * 5
        )  # gripper, handle, door, cabient, switch_link, fix_link, other
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            seg_visids = [
                self.handle_visid,
                self.door_visid + self.other_door_visid,
                self.cabinet_visid,
            ]
            seg_id = [0, 1, 2]
        elif self._arti_mode == "faucet":
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [3, 4]
        else:
            raise NotImplementedError

        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros((cam_H, cam_W), dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
        state_egopts["seg"] = convert_np_bool_to_float(cam_seg).astype(
            np.uint8
        )  # (H, W)

        kpts_world = apply_pose_to_points(
            self._handle_keypoints, self._target_link.get_pose()
        )  ## world frame
        kpts_ee = apply_pose_to_points(
            kpts_world, self.agent.grasp_site.get_pose().inv()
        )  # ee frame
        kpts = kpts_ee + np.random.normal(loc=0, scale=0.015, size=kpts_ee.shape)
        state_egopts["kpts"] = kpts.astype(np.float32)  # (J, 3)
        return state_egopts

    def get_egostereo_rgbdseg(self):
        # get state
        state_egorgbdseg = self._get_obs_state_dict()

        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        hand_sensor.compute_depth()

        state_egorgbdseg["rgb"] = (
            hand_sensor.get_rgb().transpose(2, 0, 1) * 255
        ).astype(
            np.uint8
        )  # [3, H, W]
        state_egorgbdseg["depth"] = hand_sensor.get_depth()[None].astype(
            np.float32
        )  # [1, H, W]

        # process seg groundtruth
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        cam_H, cam_W = visual_id_seg.shape
        # gripper, handle, door, cabient, switch_link, fix_link, other
        cam_seg = np.ones((cam_H, cam_W), dtype=np.bool) * 5
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            seg_visids = [
                self.handle_visid,
                self.door_visid + self.other_door_visid,
                self.cabinet_visid,
            ]
            seg_id = [0, 1, 2]
        elif self._arti_mode == "faucet":
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [3, 4]
        else:
            raise NotImplementedError
        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros((cam_H, cam_W), dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
        state_egorgbdseg["seg"] = convert_np_bool_to_float(cam_seg).astype(
            np.uint8
        )  # (H, W)
        return state_egorgbdseg

    def get_egostereo_rgbdseg_gt(self):
        # get state
        state_egorgbdseg = self._get_obs_state_dict()

        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        hand_sensor.compute_depth()

        # state_egorgbdseg["rgb"] = (hand_sensor.get_rgb().transpose(2, 0, 1) * 255).astype(np.uint8)  # [3, H, W]
        state_egorgbdseg["depth"] = hand_sensor.get_depth()[None].astype(
            np.float32
        )  # [1, H, W]

        # process seg groundtruth
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        cam_H, cam_W = visual_id_seg.shape
        # gripper, handle, door, cabient, switch_link, fix_link, other
        cam_seg = np.ones((cam_H, cam_W), dtype=np.bool) * 5
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            seg_visids = [self.handle_visid, self.door_visid, self.cabinet_visid]
            seg_id = [0, 1, 2]
        elif self._arti_mode == "faucet":
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [3, 4]
        else:
            raise NotImplementedError
        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros((cam_H, cam_W), dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
        state_egorgbdseg["seg"] = convert_np_bool_to_float(cam_seg).astype(
            np.uint8
        )  # (H, W)
        return state_egorgbdseg

    def get_egostereo_state_rgbdseg(self):
        # get state
        state_egorgbdseg = self._get_obs_state_dict()

        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        hand_sensor.compute_depth()

        state_egorgbdseg["rgb"] = (
            hand_sensor.get_rgb().transpose(2, 0, 1) * 255
        ).astype(
            np.uint8
        )  # [3, H, W]
        state_egorgbdseg["depth"] = hand_sensor.get_depth()[None]  # [1, H, W]

        if self.sample_mode in ["full_downsample", "score_sample", "random_sample"]:
            # process point cloud
            cam_xyz = hand_sensor.get_pointcloud(False)  # [H*W, 3]
            if self.sample_mode in ["score_sample", "random_sample"]:
                # cam to ee frame directly
                T_cam2ee = self.cam_para["hand_camera_extrinsic_base_frame"]
                ee_xyz = (
                    transform_points(T_cam2ee, cam_xyz)
                    .reshape(-1, 3)
                    .astype(np.float32)
                )  # [H*W, 3]
                state_egorgbdseg["ee_xyz"] = ee_xyz
            else:
                ## cam to world
                T_cam2world = sensor_rgb.get_model_matrix()
                world_xyz = transform_points(
                    T_cam2world, cam_xyz * np.array([1, -1, -1])
                )  # [H*W, 3]
                T_world2ee = (
                    self.agent.grasp_site.get_pose().inv().to_transformation_matrix()
                )  ## world => ee
                state_egorgbdseg["world_xyz"] = world_xyz.astype(np.float32)  # [H*W, 3]
                state_egorgbdseg["world_ee"] = T_world2ee

        # process seg groundtruth
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        cam_H, cam_W = visual_id_seg.shape
        # gripper, handle, door, cabient, switch_link, fix_link, other
        cam_seg = np.ones((cam_H, cam_W), dtype=np.bool) * 5
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            seg_visids = [
                self.handle_visid,
                self.door_visid + self.other_door_visid,
                self.cabinet_visid,
            ]
            seg_id = [0, 1, 2]
        elif self._arti_mode == "faucet":
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [3, 4]
        else:
            raise NotImplementedError
        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros((cam_H, cam_W), dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
        state_egorgbdseg["seg"] = convert_np_bool_to_float(cam_seg).astype(
            np.uint8
        )  # (H, W)
        return state_egorgbdseg

    def get_egostereo_dexpoints(self):
        # get state
        state_egodexpts = self._get_obs_state_dict()

        # get ee pad sampled points
        eepad_pts_world = self._agent.sample_ee_coords(num_sample=20).reshape(
            -1, 3
        )  # [40, 3], world frame
        eepad_pts_robot = apply_pose_to_points(
            eepad_pts_world, self._agent._robot.get_root_pose().inv()
        )
        # eepad_pts = apply_pose_to_points(
        #     eepad_pts_world, self.agent.grasp_site.get_pose().inv()
        # )  # [40, 3], ee frame
        eepat_pts_label = np.zeros((eepad_pts_robot.shape[0], 2))  # (40, 2)
        eepat_pts_label[:, -1] = 1.0
        eepad_ptsC_robot = np.concatenate(
            (eepad_pts_robot, eepat_pts_label), axis=1
        )  # (40, 3+2)

        # get visual info
        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        hand_sensor.compute_depth()

        # process seg groundtruth
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        cam_H, cam_W = visual_id_seg.shape
        cam_seg = (
            np.ones((cam_H, cam_W), dtype=np.bool) * 5
        )  # gripper, handle, door, cabient, switch_link, fix_link, other
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            seg_visids = [
                self.handle_visid,
                self.door_visid + self.other_door_visid,
                self.cabinet_visid,
            ]
            seg_id = [0, 1, 2]
        elif self._arti_mode in ["faucet", "laptop", "kitchen_pot"]:
            seg_visids = [self.switch_links_visid, self.fix_link_visid]
            seg_id = [3, 4]
        else:
            raise NotImplementedError

        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros((cam_H, cam_W), dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
            state_egodexpts["seg"] = convert_np_bool_to_float(cam_seg).astype(
                np.uint8
            )  # (H, W)

        # process point cloud
        cam_xyz = hand_sensor.get_pointcloud(False)  # [H*W, 3]
        ## cam to world
        T_cam2world = sensor_rgb.get_model_matrix()
        world_xyz = transform_points(
            T_cam2world, cam_xyz * np.array([1, -1, -1])
        )  # [H*W, 3]
        robot_xyz = apply_pose_to_points(
            world_xyz, self._agent._robot.get_root_pose().inv()
        )
        sample_indices = pcd_uni_down_sample_with_crop(
            world_xyz,
            1000,
            min_bound=np.array([-0.5, -1, 1e-3]),
            max_bound=np.array([1, 1, 2]),
        )
        indices_len = len(sample_indices)
        if indices_len < 1000:
            sample_indices.extend(
                self._episode_rng.choice(world_xyz.shape[0], 1000 - indices_len)
            )
        ds_robot_xyz = robot_xyz[sample_indices]
        ds_robot_label = np.zeros((ds_robot_xyz.shape[0], 2))  # (1000, 2)
        ds_robot_label[:, 0] = 1.0
        ds_ptsC_robot = np.concatenate(
            (ds_robot_xyz, ds_robot_label), axis=1
        )  # (1000, 3+2)

        state_egodexpts["segsampled_ptsC"] = np.concatenate(
            (ds_ptsC_robot, eepad_ptsC_robot), axis=0
        ).astype(
            np.float32
        )  # (1000+40, 3+2)
        return state_egodexpts

    def get_obs(self):
        if self._obs_mode == "state":
            state_dict = self._get_obs_state_dict()
            return flatten_state_dict(state_dict)
        elif self._obs_mode == "state_dict":
            state_dict = self._get_obs_state_dict()
            return state_dict
        elif self._obs_mode == "state_egorgbd":
            ### get state
            state_rgbd = self._get_obs_state_dict()

            ### get visual info
            self.update_render()
            ego_cam = self._agent._cameras["hand_camera"]
            ego_cam.take_picture()
            ## process rgb
            cam_rgb = get_camera_rgb(ego_cam)  # (H, W, 3)
            cam_H, cam_W = cam_rgb.shape[:2]
            state_rgbd["rgb"] = (cam_rgb.transpose((2, 0, 1))).astype(
                np.uint8
            )  # (3, H, W)
            ## process depth & point clouds
            position = get_texture(ego_cam, "Position")
            cam_depth = -position[..., [2]]  # [H, W, 1]
            state_rgbd["depth"] = (cam_depth.transpose((2, 0, 1))).astype(
                np.float32
            )  # (1, H, W)

            visual_id_seg = get_camera_seg(ego_cam)[..., 0]
            cam_seg = (
                np.ones((cam_H, cam_W), dtype=np.bool) * 5
            )  # gripper, handle, door, cabient, ..., ..., other
            if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
                seg_visids = [self.handle_visid, self.door_visid, self.cabinet_visid]
                seg_id = [0, 1, 2]
                for id, seg_visid in enumerate(seg_visids):
                    mask = np.zeros((cam_H, cam_W), dtype=np.bool)
                    for visual_id in seg_visid:
                        mask = mask | (visual_id_seg == visual_id)
                    cam_seg[mask] = seg_id[id]
                state_rgbd["seg"] = convert_np_bool_to_float(cam_seg).astype(
                    np.float32
                )  # (H, W)
            elif self._arti_mode in ["faucet", "laptop", "kitchen_pot"]:
                seg_visids = [self.switch_links_visid, self.fix_link_visid]
                seg_id = [3, 4]
                for id, seg_visid in enumerate(seg_visids):
                    mask = np.zeros((cam_H, cam_W), dtype=np.bool)
                    for visual_id in seg_visid:
                        mask = mask | (visual_id_seg == visual_id)
                    cam_seg[mask] = seg_id[id]
                state_rgbd["seg"] = convert_np_bool_to_float(cam_seg).astype(
                    np.float32
                )  # (H, W)
            else:
                raise NotImplementedError

            return state_rgbd
        elif self._obs_mode == "state_egosegpoints":
            # get state
            state_egopts = self._get_obs_state_dict()

            # get visual info
            self.update_render()
            ego_cam = self._agent._cameras["hand_camera"]
            ego_cam.take_picture()
            # process rgb
            cam_rgb = get_camera_rgb(ego_cam)  # (H, W, 3)
            cam_H, cam_W = cam_rgb.shape[:2]
            rgb_proc = torch.from_numpy(
                (cam_rgb.transpose((2, 0, 1)) / 255.0)[None].astype(np.float32)
            ).to(
                torch.device(self.device)
            )  # (1, 3, H, W)

            # process depth
            position = get_texture(ego_cam, "Position")  # [H, W, 4]
            cam_depth = -position[..., [2]]  # [H, W, 1]
            depth_proc = torch.from_numpy(
                (cam_depth.transpose((2, 0, 1)))[None].astype(np.float32)
            ).to(
                torch.device(self.device)
            )  # (1, 1, H, W)
            # process points
            cam_xyz = position[..., :3]
            with np_random(self._episode_seed):
                cam_xyz_noise = cam_xyz + np.random.normal(
                    loc=0, scale=0.01, size=cam_xyz.shape
                )
            # invalid_mask = position[..., -1] < 1  # Remove invalid points
            # filter_cam_xyz = cam_xyz[invalid_mask]  # [H, W, 3]
            # cam to ee frame directly
            T_cam2ee = self.cam_para["hand_camera_extrinsic_base_frame"]
            ee_xyz = (
                (cam_xyz_noise @ T_cam2ee[:3, :3].transpose() + T_cam2ee[:3, 3])
                .reshape(-1, 3)
                .astype(np.float32)
            )  # [H*W, 3]

            # predict seg & sample K points
            seg_map = self.segmodel.predict(torch.cat([rgb_proc, depth_proc], dim=1))[
                0
            ]  # (C, H, W)
            seg_map_noise = (
                seg_map + torch.randn(seg_map.shape, device=seg_map.device) * 0.1
            )
            _, indices = torch.topk(
                seg_map_noise.view(self.num_classes, -1), self.sample_num, dim=1
            )  # (C, sample_num)
            pts_index = indices.detach().cpu().numpy().reshape(-1)  # (C*sample_num)
            segsampled_pts = ee_xyz[pts_index]  # (C*sample_num, 3)
            state_egopts["segsampled_pts"] = segsampled_pts

            visual_id_seg = get_camera_seg(ego_cam)[..., 0]
            cam_seg = (
                np.ones((cam_H, cam_W), dtype=np.bool) * 5
            )  # gripper, handle, door, cabient, switch_link, fix_link, other
            if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
                seg_visids = [self.handle_visid, self.door_visid, self.cabinet_visid]
                seg_id = [0, 1, 2]
            elif self._arti_mode == "faucet":
                seg_visids = [self.switch_links_visid, self.fix_link_visid]
                seg_id = [3, 4]
            else:
                raise NotImplementedError

            for id, seg_visid in enumerate(seg_visids):
                mask = np.zeros((cam_H, cam_W), dtype=np.bool)
                for visual_id in seg_visid:
                    mask = mask | (visual_id_seg == visual_id)
                cam_seg[mask] = seg_id[id]
            state_egopts["seg"] = convert_np_bool_to_float(cam_seg).astype(
                np.float32
            )  # (H, W)

            return state_egopts
        elif self._obs_mode == "state_ego_segpoints_gt":
            ### get state
            state_ego_segpoints_gt = self._get_obs_state_dict()

            ### get visual info
            self.update_render()
            ego_cam = self._agent._cameras["hand_camera"]
            ego_cam.take_picture()
            ## process rgb
            # cam_rgb = get_camera_rgb(ego_cam)  # (H, W, 3)
            # cam_H, cam_W = cam_rgb.shape[:2]
            # state_ego_segpoints_gt["rgb"] = (cam_rgb.transpose((2, 0, 1))).astype(np.uint8)  # (3, H, W)
            ## process depth & point clouds
            position = get_texture(ego_cam, "Position")
            # cam_depth = -position[..., [2]]  # [H, W, 1]
            # state_rgbd["depth"] = (cam_depth.transpose((2, 0, 1))).astype(np.float32)  # (1, H, W)
            cam_xyz = position[..., :3].reshape(-1, 3)  # (H*W, 3)

            # get set gt
            visual_id_seg = get_camera_seg(ego_cam)[..., 0]  # [H, W]
            cam_H, cam_W = visual_id_seg.shape
            cam_seg = np.ones((cam_H, cam_W), dtype=np.bool) * (
                self.num_classes - 1
            )  # gripper, handle, door, cabient, switch_link, fix_link, other
            if self._arti_mode == "cabinet_door":
                seg_visids = [
                    self.handle_visid,
                    self.door_visid + self.other_door_visid,
                    self.cabinet_visid,
                ]
                seg_id = [0, 2, 3]
            elif self._arti_mode == "cabinet_drawer":
                seg_visids = [
                    self.handle_visid,
                    self.door_visid + self.other_door_visid,
                    self.cabinet_visid,
                ]
                seg_id = [0, 2, 3]
            elif self._arti_mode == "faucet":
                seg_visids = [self.switch_links_visid, self.fix_link_visid]
                seg_id = [1, 3]
            elif self._arti_mode == "laptop":
                seg_visids = [self.switch_links_visid, self.fix_link_visid]
                seg_id = [1, 3]
            elif self._arti_mode == "kitchen_pot":
                seg_visids = [self.handle_visid, self.lid_visid, self.fix_link_visid]
                seg_id = [0, 2, 3]
            else:
                raise NotImplementedError
            for id, seg_visid in enumerate(seg_visids):
                mask = np.zeros((cam_H, cam_W), dtype=np.bool)
                for visual_id in seg_visid:
                    mask = mask | (visual_id_seg == visual_id)
                cam_seg[mask] = seg_id[id]
            state_ego_segpoints_gt["seg"] = convert_np_bool_to_float(cam_seg).astype(
                np.uint8
            )  # (H, W)
            # get C * sample_num indices from cam_seg
            seg_label = np.eye(self.num_classes)
            seg_onehot_label = seg_label[cam_seg.reshape(-1)].astype(
                np.float_
            )  # (H*W, C)
            indices = []
            for ind in range(seg_onehot_label.shape[1]):
                one_indices = np.nonzero(seg_onehot_label[:, ind])[0]  # S
                ones_num = one_indices.shape[0]
                if ones_num == 0:
                    cur_indices = np.zeros(self.sample_num, dtype=np.int_)
                elif ones_num < self.sample_num:
                    repeat_times = self.sample_num // ones_num
                    remain_num = self.sample_num - repeat_times * ones_num
                    cur_indices = np.concatenate(
                        (one_indices.repeat(repeat_times), one_indices[:remain_num])
                    )
                else:
                    cur_indices = np.random.choice(one_indices, self.sample_num)
                indices.append(cur_indices)
            pts_index = np.concatenate(indices)  # (C*sample_num)

            # get sampled point cloud (ee frame)
            # cam_xyz = cam_xyz * np.array([1, -1, -1])
            ds_cam_xyz = cam_xyz[pts_index]  # (C*sample_num, 3)
            T_cam2ee = self.cam_para["hand_camera_extrinsic_base_frame"]
            ds_ee_xyz = transform_points(T_cam2ee, ds_cam_xyz)  # (C*sample_num, 3)
            # get sampled one-hot label
            sampled_onehot_label = seg_onehot_label[pts_index]  # (C*sample_num, C)
            segptsC = np.concatenate(
                (ds_ee_xyz, sampled_onehot_label), axis=1
            )  # (C*sample_num, 3+C)
            state_ego_segpoints_gt["segsampled_ptsC"] = segptsC.astype(
                np.float_
            )  # (C*sample_num+40, 3+C)

            return state_ego_segpoints_gt
        elif self._obs_mode == "state_egostereo_rgbd":
            state_egorgbdseg = self.get_egostereo_state_rgbdseg()
            return state_egorgbdseg
        elif self._obs_mode == "state_egostereo_segpoints":
            state_egosegpts = self.get_egostereo_segpts()
            return state_egosegpts
        elif self._obs_mode == "state_egostereo_segpoints_gt":
            state_egosegptsgt = self.get_egostereo_segpts_gt()
            return state_egosegptsgt
        elif self._obs_mode == "state_egostereo_keypoints":
            state_egokpts = self.get_egostereo_rgbd_seg_kpts()
            return state_egokpts
        elif self._obs_mode == "state_egostereo_keypoints_gt":
            state_egokptsgt = self.get_egostereo_rgbd_seg_kpts_gt()
            return state_egokptsgt
        elif self._obs_mode == "state_egostereo_dseg":
            state_egodseg = self.get_egostereo_rgbdseg()
            return state_egodseg
        elif self._obs_mode == "state_egostereo_dseg_gt":
            state_egodseggt = self.get_egostereo_rgbdseg_gt()
            return state_egodseggt
        elif self._obs_mode == "state_egostereo_dexpoints":
            state_egodexpoints = self.get_egostereo_dexpoints()
            return state_egodexpoints
        else:
            raise NotImplementedError(self._obs_mode)

    @property
    def current_angle(self):
        return self._articulation.get_qpos()[self.target_joint_idx]

    def _compute_distance(self):
        ## transfer to target frame
        pcdee_wf = self._agent.sample_ee_coords(num_sample=20).reshape(-1, 3)  # (40, 3)
        T = self.target_link.get_pose().inv().to_transformation_matrix()
        pcdee_tf = transform_points(T, pcdee_wf)  # (40, 3), target frame

        pcdee_center_tf = pcdee_tf.mean(0)
        # pcdtip_center_tf = self.tip_points.mean(0)
        # l1_distance = pcdee_center_tf - pcdtip_center_tf
        ee_tip_dist = pcdee_center_tf - self.tip_points
        l1x_distance = np.max(ee_tip_dist[:, 0])
        l2_distance = np.min(np.linalg.norm(ee_tip_dist[:, 1:], axis=1))
        return l1x_distance, l2_distance

    def _get_visual_state(self) -> OrderedDict:  # from ManiSkill1
        joint_pos = (
            self._articulation.get_qpos()[self._target_joint_idx] / self.target_qpos
        )
        current_handle = apply_pose_to_points(
            self._handle_info["pcd"], self.target_link.get_pose()
        ).mean(0)

        visual_state = OrderedDict()
        visual_state["joint_pos"] = np.array([joint_pos])  # 1
        visual_state["open_enough"] = convert_np_bool_to_float(joint_pos > 1.0)  # 1
        visual_state["current_handle"] = current_handle  # 3

        return visual_state

    def check_handle_exist(self, obs):
        assert "seg" in obs.keys()
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            partseg_id = 0
        elif self._arti_mode == "faucet":
            partseg_id = 3
        elif self._arti_mode == "laptop":
            partseg_id = 5
        elif self._arti_mode == "kitchen_pot":
            partseg_id = 7
        else:
            raise NotImplementedError
        handle_exist_rew = 1 if np.any(obs["seg"] == partseg_id) else 0
        return handle_exist_rew

    def compute_other_flag_dict(self):
        ee_cords = self._agent.sample_ee_coords()  # [2, 10, 3]
        current_handle = apply_pose_to_points(
            self._handle_info["pcd"], self._target_link.get_pose()
        )  # [200, 3]
        ee_to_handle = (
            ee_cords[..., None, :] - current_handle
        )  # [2, 10, 200, 3] = [2, 10, 1, 3] - [200, 3]
        dist_ee_to_handle = (
            np.linalg.norm(ee_to_handle, axis=-1).min(-1).min(-1)
        )  # [2, 10, 200, 3] -> [2, 10, 200](dist) -> [2, 10] -> [2] min_dist between (left_finger, right_finger) <-> handle

        handle_center_world = apply_pose_to_points(
            self._handle_info["center"][None], self._target_link.get_pose()
        )[0]

        handle_mesh = trimesh.Trimesh(
            vertices=apply_pose_to_points(
                np.asarray(self._handle_info["mesh"].vertices),
                self._target_link.get_pose(),
            ),
            faces=np.asarray(np.asarray(self._handle_info["mesh"].triangles)),
        )
        handle_obj = trimesh.proximity.ProximityQuery(handle_mesh)
        # sd between ee_left, ee_right <-> handle, >0 means inside the mesh; <0 means outside the mesh
        sd_ee_to_handle = (
            handle_obj.signed_distance(ee_cords.reshape(-1, 3)).reshape(2, -1).max(1)
        )  # [2,10,3]->[20,3]->[20]->[2,10]->2
        # sd between mid_ee <-> handle, >0 means inside the mesh; <0 means outside the mesh
        sd_ee_mid_to_handle = handle_obj.signed_distance(
            ee_cords.mean(0)
        ).max()  # [2,10,3]->[10,3]->[10]->1

        ee_close_to_handle = convert_np_bool_to_float(
            dist_ee_to_handle.max() <= 0.025
            and sd_ee_mid_to_handle > -0.01  # 0 or -0.01
        )

        # Grasp = mid close almost in cvx and both sides has similar sign distance.
        close_to_grasp = (
            sd_ee_to_handle.min() > -1e-2
        )  ## grasp happen: one side close to handle
        ee_in_grasp_pose = (
            sd_ee_mid_to_handle > 0
        )  ## one mean gripper point in handle mesh, 0/-1e-2
        grasp_happen = ee_in_grasp_pose and close_to_grasp

        other_info = {
            "dist_ee_to_handle": np.array(dist_ee_to_handle, dtype=np.float_),  ## 2
            "handle_center_world": np.array(handle_center_world, dtype=np.float_),  ## 1
            "ee_close_to_handle": convert_np_bool_to_float(ee_close_to_handle),  ## 1
            "sd_ee_to_handle": np.array(sd_ee_to_handle, dtype=np.float_),  ## 2
            "sd_ee_mid_to_handle": np.array(sd_ee_mid_to_handle, dtype=np.float_),  ## 1
            "grasp_happen": convert_np_bool_to_float(grasp_happen),  ## 1
        }
        return other_info

    def compute_eval_flag_dict(self):
        flag_dict = OrderedDict()
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            flag_dict["cabinet_static"] = convert_np_bool_to_float(
                self.check_actor_static(self._target_link, max_v=0.1, max_ang_v=1.0)
            )
            flag_dict["open_enough"] = convert_np_bool_to_float(
                self._articulation.get_qpos()[
                    self._articulation_config["target_joint_idx"]
                ]
                >= 0.9 * self.target_qpos
            )
        elif self._arti_mode in ["faucet", "laptop", "kitchen_pot"]:
            cur_qpos = self._articulation.get_qpos()[self.target_joint_idx]
            flag_dict["target_achieved"] = (
                (cur_qpos - self.init_angle) / (self.target_angle - self.init_angle)
            ) > 0.5
        else:
            raise NotImplementedError

        flag_dict["is_success"] = convert_np_bool_to_float(all(flag_dict.values()))
        return flag_dict

    def check_success(self):
        flag_dict = self.compute_eval_flag_dict()
        return flag_dict["is_success"]

    def compute_dense_reward_cabinet(self, **kwargs):
        np.set_printoptions(suppress=True)
        flag_dict = self.compute_eval_flag_dict()
        other_info = self.compute_other_flag_dict()
        dist_ee_to_handle = other_info["dist_ee_to_handle"]
        sd_ee_mid_to_handle = other_info["sd_ee_mid_to_handle"]  # > 0 means good
        sd_ee_to_handle = other_info[
            "sd_ee_to_handle"
        ]  # [2,], min > -1e-2 means grasp happen: one side close to handle

        agent_proprio = self.agent.get_proprioception()
        gripper_qpos = agent_proprio["qpos"][-2:]

        grasp_site_pose = self._agent.grasp_site.get_pose()
        target_pose_2 = self._target_link.get_pose() * self._handle_info["grasp"][1]
        gripper_angle_err = (
            np.abs(angle_distance_ms(grasp_site_pose, target_pose_2)) / np.pi
        )

        handle_center_world = other_info["handle_center_world"]
        ## handle center xyz in EE frame, abs: 0 ~ (0.4, 0.1, 0.1)
        handle_center_ee = apply_pose_to_points(
            handle_center_world[None], grasp_site_pose.inv()
        )[0]

        cabinet_qpos = self._articulation.get_qpos()[self._target_joint_idx]
        cabinet_qvel = self._articulation.get_qvel()[self._target_joint_idx]

        door_vel_norm = min(norm(self._target_link.get_velocity()), 1)  ## [0, 1]
        door_ang_vel_norm = min(
            norm(self._target_link.get_angular_velocity()), 1
        )  ## [0, 1]

        scale = 0.5
        stage_index = 0

        handle_exist_rew = self.check_handle_exist(kwargs["obs"])

        # [-1.5, 0]: (0.01, 0.5) -> [-0.5, 0] + (-0.35, 0.01) => [0, 1]
        close_to_handle_rew = (
            -1
            - dist_ee_to_handle.mean()
            + normalize_and_clip_in_interval(sd_ee_mid_to_handle, -0.35, 0.01)
        )
        dir_rew = (
            -np.clip(gripper_angle_err, 0, 0.5)
            * 2
            # - np.clip(np.sum(np.abs(handle_center_ee[:2])), 0, 0.1) * 5
        )

        cabinet_qvel_rew = 0
        cabinet_qpos_rew = 0
        cabinet_static_rew = 0
        ## open half, [-1, 0]
        gripper_rew = -np.abs(gripper_qpos.mean() - 0.034) / 0.034

        if dist_ee_to_handle.max() < 0.03 and sd_ee_mid_to_handle > -0.01:
            stage_index = 1
            ## close, [-2, 2]
            gripper_rew = (gripper_qpos.mean() - 0.034) / 0.034 * 2
            cabinet_qvel_rew = (
                np.clip(cabinet_qvel, -0.5, 0.5) * 2
            )  ## Push vel to positive, [-1, 1], init 0
            cabinet_qpos_rew = (
                normalize_and_clip_in_interval(cabinet_qpos, 0, self.target_qpos * 1.1)
                * 2
            )  ## [0, 2]
            if flag_dict["open_enough"]:
                ## cabinet_qpos >= self.target_qpos
                stage_index = 2
                gripper_rew = 2
                cabinet_qvel_rew = 1
                cabinet_qvel_rew = 2
                cabinet_static_rew = -(door_vel_norm + door_ang_vel_norm)  ## [-2, 0]

        reward = (
            handle_exist_rew
            + close_to_handle_rew
            + dir_rew
            + gripper_rew
            + cabinet_qvel_rew
            + cabinet_qpos_rew
            + cabinet_static_rew
        ) * scale

        action = kwargs["action"]
        if isinstance(action, dict):
            cur_action = action["action"]
        elif isinstance(action, np.ndarray):
            cur_action = action
        else:
            raise NotImplementedError
        info_dict = {
            "stage_index": np.array(stage_index),
            ### stage 0
            "dist_ee_to_handle": np.array(dist_ee_to_handle, dtype=np.float_),
            "sd_ee_to_handle": np.array(sd_ee_to_handle, dtype=np.float_),
            "sd_ee_mid_to_handle": np.array(sd_ee_mid_to_handle, dtype=np.float_),
            "gripper_angle_err": np.array(gripper_angle_err * 180, dtype=np.float_),
            "handle_center_ee": np.array(handle_center_ee, dtype=np.float_),
            "close_to_handle_rew": np.array(close_to_handle_rew, dtype=np.float_),
            "dir_rew": np.array(dir_rew, dtype=np.float_),
            "gripper_rew": np.array(gripper_rew, dtype=np.float_),
            ### stage 1
            "cabinet_qpos": np.array(cabinet_qpos, dtype=np.float_),
            "cabinet_qvel": np.array(cabinet_qvel, dtype=np.float_),
            ##  np.clip(cabinet_qvel, -0.5, 0.5) * 2
            "cabinet_qvel_rew": np.array(cabinet_qvel_rew, dtype=np.float_),
            ## normalize_and_clip_in_interval(cabinet_qpos, 0, self.target_qpos * 1.1) * 2
            "cabinet_qpos_rew": np.array(cabinet_qpos_rew, dtype=np.float_),
            ### stage 2
            "door_vel_norm": np.array(door_vel_norm, dtype=np.float_),
            "door_ang_vel_norm": np.array(door_ang_vel_norm, dtype=np.float_),
            ## -(door_vel_norm + door_ang_vel_norm * 0.5)
            "cabinet_static_rew": np.array(cabinet_static_rew, dtype=np.float_),
            "handle_exist_rew": np.array(handle_exist_rew, dtype=np.float_),
            "reward_raw": np.array(reward, dtype=np.float_),
            "check_grasp": np.array(agent_proprio["gripper_grasp"], dtype=np.float_),
            "action": np.array(cur_action, dtype=np.float_),
        }
        self._cache_info = info_dict
        return reward

    def compute_dense_reward_door(self, **kwargs):
        np.set_printoptions(suppress=True)
        flag_dict = self.compute_eval_flag_dict()
        other_info = self.compute_other_flag_dict()
        dist_ee_to_handle = other_info["dist_ee_to_handle"]
        sd_ee_mid_to_handle = other_info["sd_ee_mid_to_handle"]  # > 0 means good
        sd_ee_to_handle = other_info[
            "sd_ee_to_handle"
        ]  # [2,], min > -1e-2 means grasp happen: one side close to handle

        agent_proprio = self.agent.get_proprioception()
        gripper_qpos = agent_proprio["qpos"][-2:]

        grasp_site_pose = self._agent.grasp_site.get_pose()
        target_pose = self._target_link.get_pose() * self._handle_info["grasp"][0]
        target_pose_2 = self._target_link.get_pose() * self._handle_info["grasp"][1]
        angle1 = np.abs(angle_distance_ms(grasp_site_pose, target_pose))
        angle2 = np.abs(angle_distance_ms(grasp_site_pose, target_pose_2))
        gripper_angle_err = min(angle1, angle2) / np.pi

        handle_center_world = other_info["handle_center_world"]
        ## handle center xyz in EE frame, abs: 0 ~ (0.4, 0.1, 0.1)
        handle_center_ee = apply_pose_to_points(
            handle_center_world[None], grasp_site_pose.inv()
        )[0]

        cabinet_qpos = self._articulation.get_qpos()[self._target_joint_idx]
        cabinet_qvel = self._articulation.get_qvel()[self._target_joint_idx]

        door_vel_norm = min(norm(self._target_link.get_velocity()), 1)  ## [0, 1]
        door_ang_vel_norm = min(
            norm(self._target_link.get_angular_velocity()), 1
        )  ## [0, 1]

        scale = 0.5
        stage_index = 0

        handle_exist_rew = self.check_handle_exist(kwargs["obs"])

        # (0.07, 0.35) -> [-oo, 0], (-0.35, 0.01) => [-0.7, 0]
        close_to_handle_rew = (
            -dist_ee_to_handle.mean()
            + normalize_and_clip_in_interval(sd_ee_mid_to_handle, -0.5, 0.01)
        )  # direction reward, [0, 0.25] -> [-1, 0]
        dir_rew = (
            -np.clip(gripper_angle_err, 0, 0.5)
            - np.clip(np.sum(np.abs(handle_center_ee[:2])), 0, 0.1) * 10
        )

        ## open half, [-0.5, 0.5]
        open_gripper_rew = 0.5 - np.abs(gripper_qpos.mean() - 0.034) / 0.034
        ## close, [-1, 1]
        close_gripper_rew = (gripper_qpos.mean() - 0.034) / 0.034

        cabinet_qvel_rew = 0
        cabinet_qpos_rew = 0
        cabinet_static_rew = 0
        gripper_rew = open_gripper_rew

        if dist_ee_to_handle.max() < 0.03 and sd_ee_mid_to_handle > 0:
            ## dist_ee_to_handle.max() <= 0.025 and sd_ee_mid_to_handle > -0.01 & sd_ee_mid_to_handle > 0
            stage_index = 1
            dir_rew = 0.5
            # close_to_handle_rew = 0.5
            gripper_rew = close_gripper_rew
            cabinet_qvel_rew = (
                np.clip(cabinet_qvel, -0.5, 0.5) * 2
            )  ## Push vel to positive, [-1, 1], init 0
            cabinet_qpos_rew = (
                normalize_and_clip_in_interval(cabinet_qpos, 0, self.target_qpos * 1.1)
                * 2
            )  ## [0, 2]
            if flag_dict["open_enough"]:
                ## cabinet_qpos >= self.target_qpos
                stage_index = 2
                gripper_rew = 1
                cabinet_qvel_rew = 1
                cabinet_qvel_rew = 2
                cabinet_static_rew = -(door_vel_norm + door_ang_vel_norm)  ## [-2, 0]

        reward = (
            +handle_exist_rew
            + close_to_handle_rew
            + dir_rew
            + gripper_rew
            + cabinet_qvel_rew
            + cabinet_qpos_rew
            + cabinet_static_rew
        ) * scale

        action = kwargs["action"]
        if isinstance(action, dict):
            cur_action = action["action"]
        elif isinstance(action, np.ndarray):
            cur_action = action
        else:
            raise NotImplementedError
        info_dict = {
            "stage_index": np.array(stage_index),
            ### stage 0
            "dist_ee_to_handle": np.array(dist_ee_to_handle, dtype=np.float_),
            "sd_ee_to_handle": np.array(sd_ee_to_handle, dtype=np.float_),
            "sd_ee_mid_to_handle": np.array(sd_ee_mid_to_handle, dtype=np.float_),
            "gripper_angle_err": np.array(gripper_angle_err * 180, dtype=np.float_),
            "close_to_handle_rew": np.array(close_to_handle_rew, dtype=np.float_),
            "dir_rew": np.array(dir_rew, dtype=np.float_),
            "gripper_rew": np.array(gripper_rew, dtype=np.float_),
            ### stage 1
            "cabinet_qpos": np.array(cabinet_qpos, dtype=np.float_),
            "cabinet_qvel": np.array(cabinet_qvel, dtype=np.float_),
            ##  np.clip(cabinet_qvel, -0.5, 0.5) * 2
            "cabinet_qvel_rew": np.array(cabinet_qvel_rew, dtype=np.float_),
            ## normalize_and_clip_in_interval(cabinet_qpos, 0, self.target_qpos * 1.1) * 2
            "cabinet_qpos_rew": np.array(cabinet_qpos_rew, dtype=np.float_),
            ### stage 2
            "door_vel_norm": np.array(door_vel_norm, dtype=np.float_),
            "door_ang_vel_norm": np.array(door_ang_vel_norm, dtype=np.float_),
            ## -(door_vel_norm + door_ang_vel_norm * 0.5)
            "cabinet_static_rew": np.array(cabinet_static_rew, dtype=np.float_),
            "handle_exist_rew": np.array(handle_exist_rew, dtype=np.float_),
            "reward_raw": np.array(reward, dtype=np.float_),
            "check_grasp": np.array(agent_proprio["gripper_grasp"], dtype=np.float_),
            "action": np.array(cur_action, dtype=np.float_),
        }
        self._cache_info = info_dict
        return reward

    def compute_dense_reward_drawer(self, **kwargs):
        np.set_printoptions(suppress=True)
        flag_dict = self.compute_eval_flag_dict()
        other_info = self.compute_other_flag_dict()
        dist_ee_to_handle = other_info["dist_ee_to_handle"]
        sd_ee_mid_to_handle = other_info["sd_ee_mid_to_handle"]  # > 0 means good
        sd_ee_to_handle = other_info[
            "sd_ee_to_handle"
        ]  # [2,], min > -1e-2 means grasp happen: one side close to handle

        agent_proprio = self.agent.get_proprioception()
        gripper_qpos = agent_proprio["qpos"][-2:]

        grasp_site_pose = self._agent.grasp_site.get_pose()
        target_pose_2 = self._target_link.get_pose() * self._handle_info["grasp"][1]
        gripper_angle_err = (
            np.abs(angle_distance_ms(grasp_site_pose, target_pose_2)) / np.pi
        )

        handle_center_world = other_info["handle_center_world"]
        ## handle center xyz in EE frame, abs: 0 ~ (0.4, 0.1, 0.1)
        handle_center_ee = apply_pose_to_points(
            handle_center_world[None], grasp_site_pose.inv()
        )[0]

        cabinet_qpos = self._articulation.get_qpos()[self._target_joint_idx]
        cabinet_qvel = self._articulation.get_qvel()[self._target_joint_idx]

        door_vel_norm = min(norm(self._target_link.get_velocity()), 1)  ## [0, 1]
        door_ang_vel_norm = min(
            norm(self._target_link.get_angular_velocity()), 1
        )  ## [0, 1]

        scale = 0.5
        stage_index = 0

        handle_exist_rew = self.check_handle_exist(kwargs["obs"])

        # [-1.5, 0]: (0.01, 0.5) -> [-0.5, 0] + (-0.35, 0.01) => [0, 1]
        close_to_handle_rew = (
            -1
            - dist_ee_to_handle.mean()
            + normalize_and_clip_in_interval(sd_ee_mid_to_handle, -0.35, 0.01)
        )
        dir_rew = -np.clip(gripper_angle_err, 0, 0.5) * 2

        cabinet_qvel_rew = 0
        cabinet_qpos_rew = 0
        cabinet_static_rew = 0
        ## open half, [-1, 0]
        gripper_rew = -np.abs(gripper_qpos.mean() - 0.034) / 0.034

        if gripper_angle_err < 0.15:  # 27degree
            dir_rew = 0.5
            if dist_ee_to_handle.max() < 0.03 and sd_ee_mid_to_handle > -0.01:
                stage_index = 1
                # dir_rew = 0.5
                close_to_handle_rew = 0.5
                ## close, [-2, 2]
                gripper_rew = (gripper_qpos.mean() - 0.034) / 0.034 * 2
                cabinet_qvel_rew = (
                    np.clip(cabinet_qvel, -0.5, 0.5) * 2
                )  ## Push vel to positive, [-1, 1], init 0
                cabinet_qpos_rew = (
                    normalize_and_clip_in_interval(
                        cabinet_qpos, 0, self.target_qpos * 1.1
                    )
                    * 2
                )  ## [0, 2]
                if flag_dict["open_enough"]:
                    ## cabinet_qpos >= self.target_qpos
                    stage_index = 2
                    gripper_rew = 2
                    cabinet_qvel_rew = 1
                    cabinet_qvel_rew = 2
                    cabinet_static_rew = -(
                        door_vel_norm + door_ang_vel_norm
                    )  ## [-2, 0]

        reward = (
            handle_exist_rew
            + close_to_handle_rew
            + dir_rew
            + gripper_rew
            + cabinet_qvel_rew
            + cabinet_qpos_rew
            + cabinet_static_rew
        ) * scale

        action = kwargs["action"]
        if isinstance(action, dict):
            cur_action = action["action"]
        elif isinstance(action, np.ndarray):
            cur_action = action
        else:
            raise NotImplementedError
        info_dict = {
            "stage_index": np.array(stage_index),
            ### stage 0
            "dist_ee_to_handle": np.array(dist_ee_to_handle, dtype=np.float_),
            "sd_ee_to_handle": np.array(sd_ee_to_handle, dtype=np.float_),
            "sd_ee_mid_to_handle": np.array(sd_ee_mid_to_handle, dtype=np.float_),
            "gripper_angle_err": np.array(gripper_angle_err * 180, dtype=np.float_),
            "handle_center_ee": np.array(handle_center_ee, dtype=np.float_),
            "close_to_handle_rew": np.array(close_to_handle_rew, dtype=np.float_),
            "dir_rew": np.array(dir_rew, dtype=np.float_),
            "gripper_rew": np.array(gripper_rew, dtype=np.float_),
            ### stage 1
            "cabinet_qpos": np.array(cabinet_qpos, dtype=np.float_),
            "cabinet_qvel": np.array(cabinet_qvel, dtype=np.float_),
            ##  np.clip(cabinet_qvel, -0.5, 0.5) * 2
            "cabinet_qvel_rew": np.array(cabinet_qvel_rew, dtype=np.float_),
            ## normalize_and_clip_in_interval(cabinet_qpos, 0, self.target_qpos * 1.1) * 2
            "cabinet_qpos_rew": np.array(cabinet_qpos_rew, dtype=np.float_),
            ### stage 2
            "door_vel_norm": np.array(door_vel_norm, dtype=np.float_),
            "door_ang_vel_norm": np.array(door_ang_vel_norm, dtype=np.float_),
            ## -(door_vel_norm + door_ang_vel_norm * 0.5)
            "cabinet_static_rew": np.array(cabinet_static_rew, dtype=np.float_),
            "handle_exist_rew": np.array(handle_exist_rew, dtype=np.float_),
            "reward_raw": np.array(reward, dtype=np.float_),
            "check_grasp": np.array(agent_proprio["gripper_grasp"], dtype=np.float_),
            "action": np.array(cur_action, dtype=np.float_),
        }
        self._cache_info = info_dict
        return reward

    def compute_dense_reward_faucet(self, **kwargs):
        reward = 0.0

        flag_dict = self.compute_eval_flag_dict()
        if flag_dict["is_success"]:
            return 4.0

        handle_exist_rew = self.check_handle_exist(kwargs["obs"])
        reward += handle_exist_rew
        l1_distance, l2_distance = self._compute_distance()
        close_reward = 1 - 10 * np.clip(l2_distance, 0, 0.2)  # [-1, 1]
        reward += close_reward
        left_reward = -20.0 * np.clip(l1_distance, 0, 0.1)  # [-2, 0]
        reward += left_reward

        angle_diff = self.target_angle - self.current_angle
        turn_reward_1 = 5 * (1 - np.tanh(max(angle_diff, 0) * 2.0))  # [0, 1] => [0, 5]
        reward += turn_reward_1

        delta_angle = angle_diff - self.last_angle_diff  # <=> last_angle - cur_angle
        if angle_diff > 0:
            turn_reward_2 = -np.tanh(delta_angle * 5)  # [0, 1]
        else:
            turn_reward_2 = np.tanh(delta_angle * 5)
        turn_reward_2 *= 5
        reward += turn_reward_2
        self.last_angle_diff = angle_diff

        # scale
        reward = reward * 0.5

        action = kwargs["action"]
        if isinstance(action, dict):
            cur_action = action["action"]
        elif isinstance(action, np.ndarray):
            cur_action = action
        else:
            raise NotImplementedError
        info_dict = {
            "handle_exist_rew": np.array(handle_exist_rew, dtype=np.float_),
            ### stage 0
            "l1_distance": np.array(l1_distance, dtype=np.float_),
            "left_reward": np.array(left_reward, dtype=np.float_),
            "distance": np.array(l2_distance, dtype=np.float_),
            "close_reward": np.array(close_reward, dtype=np.float_),
            ### stage 1
            "angle_diff": np.array(angle_diff, dtype=np.float_),
            "turn_reward_1": np.array(turn_reward_1, dtype=np.float_),
            ### stage 2
            "delta_angle": np.array(delta_angle, dtype=np.float_),
            "turn_reward_2": np.array(turn_reward_2, dtype=np.float_),
            "reward": np.array(reward, dtype=np.float_),
            # other info
            "cur_angle": np.array(self.current_angle, dtype=np.float_),
            "init_angle": np.array(self.init_angle, dtype=np.float_),
            "target_angle": np.array(self.target_angle, dtype=np.float_),
            "action": np.array(cur_action, dtype=np.float_),
        }
        self._cache_info = info_dict

        return reward

    def compute_dense_reward_laptop(self, **kwargs):
        reward = 0.0
        flag_dict = self.compute_eval_flag_dict()
        if flag_dict["is_success"]:
            return 4.0

        handle_exist_rew = self.check_handle_exist(kwargs["obs"])
        reward += handle_exist_rew

        ## transfer to target frame
        pcdee_wf = self._agent.sample_ee_coords(num_sample=20).reshape(-1, 3)  # (40, 3)
        T = self.target_link.get_pose().inv().to_transformation_matrix()
        pcdee_tf = transform_points(T, pcdee_wf)  # (40, 3), target frame
        pcdee_center_tf = pcdee_tf.mean(0)
        ee_tip_dist = self.tip_points - pcdee_center_tf
        l1z_distance = np.max(ee_tip_dist[:, 2])
        l2_distance = np.min(np.linalg.norm(ee_tip_dist[:, :2], axis=1))

        close_reward = 1 - 10 * np.clip(l2_distance, 0, 0.2)  # [-1, 1]
        down_reward = -20.0 * np.clip(l1z_distance, 0, 0.1)  # [-2, 0]
        reward += close_reward + down_reward

        angle_diff = self.target_angle - self.current_angle
        turn_reward_1 = 5 * (1 - np.tanh(max(angle_diff, 0) * 2.0))  # [0, 1] => [0, 5]
        reward += turn_reward_1

        delta_angle = angle_diff - self.last_angle_diff  # <=> last_angle - cur_angle
        if angle_diff > 0:
            turn_reward_2 = -np.tanh(delta_angle * 5)  # [0, 1]
        else:
            turn_reward_2 = np.tanh(delta_angle * 5)
        turn_reward_2 *= 5
        reward += turn_reward_2
        self.last_angle_diff = angle_diff

        # scale
        reward = reward * 0.5

        action = kwargs["action"]
        if isinstance(action, dict):
            cur_action = action["action"]
        elif isinstance(action, np.ndarray):
            cur_action = action
        else:
            raise NotImplementedError
        info_dict = {
            "handle_exist_rew": np.array(handle_exist_rew, dtype=np.float_),
            ### stage 0
            "l1z_distance": np.array(l1z_distance, dtype=np.float_),
            "down_reward": np.array(down_reward, dtype=np.float_),
            "l2_distance": np.array(l2_distance, dtype=np.float_),
            "close_reward": np.array(close_reward, dtype=np.float_),
            ### stage 1
            "angle_diff": np.array(angle_diff, dtype=np.float_),
            "turn_reward_1": np.array(turn_reward_1, dtype=np.float_),
            ### stage 2
            "delta_angle": np.array(delta_angle, dtype=np.float_),
            "turn_reward_2": np.array(turn_reward_2, dtype=np.float_),
            "reward": np.array(reward, dtype=np.float_),
            # other info
            "cur_angle": np.array(self.current_angle, dtype=np.float_),
            "init_angle": np.array(self.init_angle, dtype=np.float_),
            "target_angle": np.array(self.target_angle, dtype=np.float_),
            "action": np.array(cur_action, dtype=np.float_),
        }
        self._cache_info = info_dict

        return reward

    def compute_dense_reward_kitchenpot(self, **kwargs):
        flag_dict = self.compute_eval_flag_dict()
        if flag_dict["is_success"]:
            return 4.0

        handle_exist_rew = self.check_handle_exist(kwargs["obs"])

        ## transfer to target frame
        gripper_pts_ee = sample_grasp_multipoints_ee(
            0.03, num_points_perlink=10, x_offset=0.02
        )
        trans_tarlink2ee = (
            self._agent.grasp_site.get_pose().inv().to_transformation_matrix()
            @ self.target_link.get_pose().to_transformation_matrix()
        )
        grasp_mat_ee = (
            trans_tarlink2ee @ self.handle_grasp[0].to_transformation_matrix()
        )
        gripper_pts_targetpose_ee = transform_points(grasp_mat_ee, gripper_pts_ee)
        gripper_pts_dist = np.linalg.norm(
            gripper_pts_targetpose_ee - gripper_pts_ee, axis=-1
        )  # (k, )
        grasp_dist = np.mean(gripper_pts_dist, axis=-1)
        approach_reward = 1 - 3.0 * np.tanh(10.0 * grasp_dist)  # [-2, 1]

        grasp_reward, turn_reward_1, turn_reward_2 = 0, 0, 0

        angle_diff = self.target_angle - self.current_angle
        delta_angle = angle_diff - self.last_angle_diff  # <=> last_angle - cur_angle
        if grasp_dist < 0.02:
            # grasp reward
            is_grasped = self.agent.check_grasp(self.target_link, max_angle=30)
            grasp_reward = 1.0 if is_grasped else 0.0

            # reach-goal reward
            if is_grasped:
                ## close, [-2, 2]
                turn_reward_1 = 1 * (1 - np.tanh(max(angle_diff, 0) * 2.0))  # [0, 1]

                if angle_diff > 0:
                    turn_reward_2 = -np.tanh(delta_angle * 5)  # [0, 1]
                else:
                    turn_reward_2 = np.tanh(delta_angle * 5)
                turn_reward_2 *= 1  # [0, 1]

        self.last_angle_diff = angle_diff

        # scale
        reward = 0.5 * (
            approach_reward
            + grasp_reward
            + turn_reward_1
            + turn_reward_2
            + handle_exist_rew
        )

        action = kwargs["action"]
        if isinstance(action, dict):
            cur_action = action["action"]
        elif isinstance(action, np.ndarray):
            cur_action = action
        else:
            raise NotImplementedError
        info_dict = {
            "handle_exist_rew": np.array(handle_exist_rew, dtype=np.float_),
            ### stage 0
            "grasp_dist": np.array(grasp_dist, dtype=np.float_),
            "approach_reward": np.array(approach_reward, dtype=np.float_),
            "grasp_reward": np.array(grasp_reward, dtype=np.float_),
            ### stage 2
            "angle_diff": np.array(angle_diff, dtype=np.float_),
            "turn_reward_1": np.array(turn_reward_1, dtype=np.float_),
            "delta_angle": np.array(delta_angle, dtype=np.float_),
            "turn_reward_2": np.array(turn_reward_2, dtype=np.float_),
            "reward": np.array(reward, dtype=np.float_),
            # other info
            "cur_angle": np.array(self.current_angle, dtype=np.float_),
            "init_angle": np.array(self.init_angle, dtype=np.float_),
            "target_angle": np.array(self.target_angle, dtype=np.float_),
            "action": np.array(cur_action, dtype=np.float_),
        }
        self._cache_info = info_dict

        return reward

    def compute_dense_reward_kitchenpot_aligndrawer(self, **kwargs):
        np.set_printoptions(suppress=True)

        ee_cords = self._agent.sample_ee_coords()  # [2, 10, 3]
        current_handle = apply_pose_to_points(
            self.handle_pcd, self._target_link.get_pose()
        )  # [200, 3]
        ee_to_handle = (
            ee_cords[..., None, :] - current_handle
        )  # [2, 10, 200, 3] = [2, 10, 1, 3] - [200, 3]
        # [2, 10, 200, 3] -> [2, 10, 200](dist) -> [2, 10] -> [2] min_dist between (left_finger, right_finger) <-> handle
        dist_ee_to_handle = np.linalg.norm(ee_to_handle, axis=-1).min(-1).min(-1)
        dist_ee_mid_to_handle = -np.linalg.norm(ee_to_handle.mean(0), axis=-1).min()
        # [-1.5, 0]: (0.01, 0.5) -> [-0.5, 0] + (-0.35, 0.01) => [0, 1]
        close_to_handle_rew = (
            -1
            - dist_ee_to_handle.mean()
            + normalize_and_clip_in_interval(dist_ee_mid_to_handle, -0.35, 0.01)
        )

        grasp_site_pose = self._agent.grasp_site.get_pose()
        target_pose_2 = self._target_link.get_pose() * self.handle_grasp[1]
        grasp_site_pose = self._agent.grasp_site.get_pose()
        gripper_angle_err = (
            np.abs(angle_distance_ms(grasp_site_pose, target_pose_2)) / np.pi
        )
        # direction reward, [-1.5, 0]
        dir_rew = -np.clip(gripper_angle_err, 0, 0.5) * 2

        cabinet_qpos = self._articulation.get_qpos()[self.target_joint_idx]
        cabinet_qvel = self._articulation.get_qvel()[self.target_joint_idx]
        door_vel_norm = min(norm(self.target_link.get_velocity()), 1)  ## [0, 1]
        door_ang_vel_norm = min(
            norm(self.target_link.get_angular_velocity()), 1
        )  ## [0, 1]

        scale = 0.5
        stage_index = 0
        cabinet_qvel_rew = 0
        cabinet_qpos_rew = 0
        cabinet_static_rew = 0
        handle_exist_rew = self.check_handle_exist(kwargs["obs"])
        agent_proprio = self.agent.get_proprioception()
        gripper_qpos = agent_proprio["qpos"][-2:]
        ## open half, [-1, 0]
        gripper_rew = -np.abs(gripper_qpos.mean() - 0.034) / 0.034

        if gripper_angle_err < 0.15:  # 27degree
            dir_rew = 0.5
            if dist_ee_to_handle.max() < 0.03 and dist_ee_mid_to_handle > -0.01:
                stage_index = 1
                # dir_rew = 0.5
                close_to_handle_rew = 0.5
                ## close, [-2, 2]
                gripper_rew = (gripper_qpos.mean() - 0.034) / 0.034 * 2
                cabinet_qvel_rew = (
                    np.clip(cabinet_qvel, -0.5, 0.5) * 2
                )  ## Push vel to positive, [-1, 1], init 0
                cabinet_qpos_rew = (
                    normalize_and_clip_in_interval(
                        cabinet_qpos, 0, self.target_qpos * 1.1
                    )
                    * 2
                )  ## [0, 2]

                cur_qpos = self._articulation.get_qpos()[self.target_joint_idx]
                target_achieved = (
                    (cur_qpos - self.init_angle) / (self.target_angle - self.init_angle)
                ) > 0.5
                if target_achieved:
                    stage_index = 2
                    gripper_rew = 2
                    cabinet_qvel_rew = 2
                    cabinet_static_rew = -(
                        door_vel_norm + door_ang_vel_norm
                    )  ## [-2, 0]

        reward = (
            handle_exist_rew
            + close_to_handle_rew
            + dir_rew
            + gripper_rew
            + cabinet_qvel_rew
            + cabinet_qpos_rew
            + cabinet_static_rew
        ) * scale

        action = kwargs["action"]
        if isinstance(action, dict):
            cur_action = action["action"]
        elif isinstance(action, np.ndarray):
            cur_action = action
        else:
            raise NotImplementedError
        info_dict = {
            "stage_index": np.array(stage_index),
            ### stage 0
            "dist_ee_to_handle": np.array(dist_ee_to_handle, dtype=np.float_),
            "dist_ee_mid_to_handle": np.array(dist_ee_mid_to_handle, dtype=np.float_),
            "gripper_angle_err": np.array(gripper_angle_err * 180, dtype=np.float_),
            "close_to_handle_rew": np.array(close_to_handle_rew, dtype=np.float_),
            "dir_rew": np.array(dir_rew, dtype=np.float_),
            "gripper_rew": np.array(gripper_rew, dtype=np.float_),
            ### stage 1
            "cabinet_qpos": np.array(cabinet_qpos, dtype=np.float_),
            "cabinet_qvel": np.array(cabinet_qvel, dtype=np.float_),
            ##  np.clip(cabinet_qvel, -0.5, 0.5) * 2
            "cabinet_qvel_rew": np.array(cabinet_qvel_rew, dtype=np.float_),
            ## normalize_and_clip_in_interval(cabinet_qpos, 0, self.target_qpos * 1.1) * 2
            "cabinet_qpos_rew": np.array(cabinet_qpos_rew, dtype=np.float_),
            ### stage 2
            "door_vel_norm": np.array(door_vel_norm, dtype=np.float_),
            "door_ang_vel_norm": np.array(door_ang_vel_norm, dtype=np.float_),
            ## -(door_vel_norm + door_ang_vel_norm * 0.5)
            "cabinet_static_rew": np.array(cabinet_static_rew, dtype=np.float_),
            "handle_exist_rew": np.array(handle_exist_rew, dtype=np.float_),
            "reward_raw": np.array(reward, dtype=np.float_),
            "action": np.array(cur_action, dtype=np.float_),
        }
        self._cache_info = info_dict
        return reward

    def step_action(self, action):
        if action is None:  # simulation without action
            pass
        elif isinstance(action, np.ndarray):
            if self._arti_mode in ["faucet", "laptop"]:
                action[-1] = 1  # set gripper close forcibly
            self._agent.set_action(action)
        elif isinstance(action, dict):
            if action["control_mode"] != self._agent.control_mode:
                self._agent.set_control_mode(action["control_mode"])
            if self._arti_mode in ["faucet", "laptop"]:
                action["action"][-1] = 1  # set gripper close forcibly
            self._agent.set_action(action["action"])
        else:
            raise TypeError(type(action))

        for _ in range(self._sim_step_per_control):
            self._agent.simulation_step()
            self._scene.step()

        self._agent.update_generalized_external_forces()

    def get_reward(self, **kwargs):
        if self._reward_mode == "sparse":
            return float(self.check_success())
        elif self._reward_mode == "dense":
            if self._arti_mode == "cabinet_door":
                return self.compute_dense_reward_door(**kwargs)
                # return self.compute_dense_reward_cabinet(**kwargs)
            elif self._arti_mode == "cabinet_drawer":
                return self.compute_dense_reward_drawer(**kwargs)
                # return self.compute_dense_reward_cabinet(**kwargs)
            elif self._arti_mode == "faucet":
                return self.compute_dense_reward_faucet(**kwargs)
            elif self._arti_mode == "laptop":
                return self.compute_dense_reward_laptop(**kwargs)
            elif self._arti_mode == "kitchen_pot":
                return self.compute_dense_reward_kitchenpot(**kwargs)
            else:
                raise NotImplementedError
        else:
            raise NotImplementedError(self._reward_mode)

    def get_info(self, **kwargs):
        info = super().get_info()
        info.update(self._cache_info)
        return info

    def get_state(self) -> np.ndarray:
        state = super().get_state()
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            return state
        elif self._arti_mode == "faucet":
            return np.hstack([state, self.target_angle])
        else:
            raise NotImplementedError

    def set_state(self, state):
        if self._arti_mode in ["cabinet_door", "cabinet_drawer"]:
            super().set_state(state)
        elif self._arti_mode == "faucet":
            self.target_angle = state[-1]
            super().set_state(state[:-1])
            self.last_angle_diff = self.target_angle - self.current_angle
        else:
            raise NotImplementedError

    @property
    def articulation(self):
        return self._articulation

    @property
    def target_link(self):
        return self._target_link

    @property
    def handle_info(self):
        return self._handle_info

    @property
    def table(self):
        return self._table

arti_mani/envs/arti_mani.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from collections import OrderedDict
from typing import Dict, Optional, Tuple, Union

import gym
import numpy as np
import sapien.core as sapien
from arti_mani.agents.base_agent import BaseAgent
from arti_mani.agents.camera import get_camera_images, get_camera_pcd
from arti_mani.utils.articulation_utils import (
    get_actor_state,
    get_articulation_state,
    set_actor_state,
    set_articulation_state,
)
from arti_mani.utils.common import (
    convert_observation_to_space,
    flatten_state_dict,
    merge_dicts,
    warn,
)
from arti_mani.utils.geometry import transform_points
from arti_mani.utils.trimesh_utils import (
    get_actor_meshes,
    get_articulation_meshes,
    merge_meshes,
)
from arti_mani.utils.visualization import observations_to_images, tile_images
from sapien.utils import Viewer


class BaseEnv(gym.Env):
    """Superclass for ManiSkill environments.

    Args:
        obs_mode: observation mode registered in @SUPPORTED_OBS_MODES.
        reward_mode: reward mode registered in @SUPPORTED_REWARD_MODES.
        sim_freq: simulation frequency (Hz).
        control_freq: control frequency (Hz).
    """

    SUPPORTED_OBS_MODES: Tuple[str] = ()
    SUPPORTED_REWARD_MODES: Tuple[str] = ()
    SUPPORTED_CONTROL_MODES: Tuple[str] = ()
    _scene: sapien.Scene  # simulation scene instance
    _agent: BaseAgent
    _cameras: Dict[str, sapien.CameraEntity]

    def __init__(
        self,
        obs_mode=None,
        reward_mode=None,
        control_mode=None,
        sim_freq=500,
        control_freq=20,
        device="",
    ):
        # super().__init__(sim_freq, control_freq, device=device)
        self._engine = sapien.Engine()
        self._engine.set_log_level("off")
        sapien.VulkanRenderer.set_log_level("off")
        # self._renderer = sapien.VulkanRenderer(default_mipmap_levels=1, device=device)
        self._renderer = sapien.SapienRenderer(device=device)
        self._engine.set_renderer(self._renderer)
        # sapien.render_config.camera_shader_dir = "rt"
        sapien.render_config.rt_samples_per_pixel = 8
        sapien.render_config.rt_max_path_depth = 8
        sapien.render_config.rt_use_denoiser = True
        self._viewer = None

        self._sim_freq = sim_freq
        self._control_freq = control_freq
        if sim_freq % control_freq != 0:
            warn(
                f"sim_freq({sim_freq}) is not divisible by control_freq({control_freq}).",
            )
        self._sim_step_per_control = sim_freq // control_freq

        if obs_mode is None:
            obs_mode = self.SUPPORTED_OBS_MODES[0]
        if obs_mode not in self.SUPPORTED_OBS_MODES:
            raise NotImplementedError("Unsupported obs mode: {}".format(obs_mode))
        self._obs_mode = obs_mode

        if reward_mode is None:
            reward_mode = self.SUPPORTED_REWARD_MODES[0]
        if reward_mode not in self.SUPPORTED_REWARD_MODES:
            raise NotImplementedError("Unsupported reward mode: {}".format(reward_mode))
        self._reward_mode = reward_mode

        if control_mode is None:
            control_mode = self.SUPPORTED_CONTROL_MODES[0]
        if control_mode not in self.SUPPORTED_CONTROL_MODES:
            raise NotImplementedError(
                "Unsupported control mode: {}".format(control_mode)
            )
        self._control_mode = control_mode

        self.seed()

        obs = self.reset(reconfigure=True)
        self.observation_space = convert_observation_to_space(obs)
        self.action_space = self._agent.action_space[self._control_mode]

    def seed(self, seed=None):
        # For each episode, seed can be passed through `reset(seed=...)`,
        # or generated by `_main_rng`
        if seed is None:
            # Explicitly generate a seed for reproducibility
            seed = np.random.RandomState().randint(2**32)
        self._main_seed = seed
        self._main_rng = np.random.RandomState(self._main_seed)
        return [self._main_seed]

    @property
    def agent(self):
        return self._agent

    @property
    def control_mode(self):
        return self._control_mode

    @property
    def sim_freq(self):
        return self._sim_freq

    @property
    def control_freq(self):
        return self._control_freq

    @property
    def sim_time_step(self):
        return 1.0 / self._sim_freq

    @property
    def control_time_step(self):
        return 1.0 / self._control_freq

    # ---------------------------------------------------------------------------- #
    # Observation
    # ---------------------------------------------------------------------------- #
    @property
    def obs_mode(self):
        return self._obs_mode

    def get_obs(self):
        if self._obs_mode == "state":
            state_dict = self._get_obs_state_dict()
            return flatten_state_dict(state_dict)
        elif self._obs_mode == "state_dict":
            return self._get_obs_state_dict()
        elif self._obs_mode == "rgbd":
            return self._get_obs_rgbd()
        elif self._obs_mode == "depth":
            return self._get_obs_depth()
        elif self._obs_mode == "pointcloud":
            return self._get_obs_pointcloud()
        else:
            raise NotImplementedError(self._obs_mode)

    def _get_obs_state_dict(self) -> OrderedDict:
        """Get (GT) state-based observations."""
        return OrderedDict(
            agent=self._get_proprioception(),
            extra=self._get_obs_extra(),
        )

    def _get_proprioception(self) -> OrderedDict:
        """Get observations from proprioceptive sensors."""
        return OrderedDict(self.agent.get_proprioception())

    def _get_obs_extra(self) -> OrderedDict:
        """Get task-relevant extra observations."""
        return OrderedDict()

    def _get_obs_rgbd(self) -> OrderedDict:
        return OrderedDict(
            image=self._get_images(),
            agent=self._get_proprioception(),
            extra=self._get_obs_extra(),
        )

    def _get_obs_depth(self) -> OrderedDict:
        return OrderedDict(
            image=self._get_images(rgb=False),
            agent=self._get_proprioception(),
            extra=self._get_obs_extra(),
        )

    def _get_images(self, rgb=True, depth=True, **kwargs) -> OrderedDict:
        """Get observations from cameras.
        The key is the camera name, and the value is an OrderedDict
        containing camera observations (rgb, depth, ...).
        """
        self.update_render()
        obs_dict = self.agent.get_images(rgb=rgb, depth=depth, **kwargs)
        return OrderedDict(obs_dict)

    def _get_camera_images(self, camera: sapien.CameraEntity, **kwargs) -> OrderedDict:
        camera.take_picture()
        obs_dict = get_camera_images(camera, **kwargs)
        obs_dict.update(
            camera_intrinsic=camera.get_intrinsic_matrix(),
            camera_extrinsic_world_frame=camera.get_extrinsic_matrix(),
        )
        return OrderedDict(obs_dict)

    def _get_obs_pointcloud(self):
        """get pointcloud from each camera, transform them to the *world* frame, and fuse together"""
        self.update_render()

        fused_pcd = self._agent.get_fused_pointcloud()

        pcds = []
        for _, camera in self._cameras.items():
            camera.take_picture()
            pcd = get_camera_pcd(camera)
            T = camera.get_model_matrix()
            pcd["xyz"] = transform_points(T, pcd["xyz"])
            pcds.append(pcd)

        if len(pcds) > 0:
            fused_pcd = merge_dicts([fused_pcd, merge_dicts(pcds, True)], True)

        return OrderedDict(
            pointcloud=fused_pcd,
            agent=self._get_proprioception(),
            extra=self._get_obs_extra(),
        )

    # -------------------------------------------------------------------------- #
    # Reward mode
    # -------------------------------------------------------------------------- #
    @property
    def reward_mode(self):
        return self._reward_mode

    @reward_mode.setter
    def reward_mode(self, mode: str):
        if mode not in self.SUPPORTED_REWARD_MODES:
            raise NotImplementedError("Unsupported reward mode: {}".format(mode))
        self._reward_mode = mode

    def get_reward(self, obs, action):
        raise NotImplementedError

    def compute_dense_reward(self, **kwargs):
        raise NotImplementedError

    # -------------------------------------------------------------------------- #
    # Reconfigure
    # -------------------------------------------------------------------------- #
    def reconfigure(self):
        """Reconfigure the simulation scene instance.
        This function should clear the previous scene, and create a new one.
        """
        self._clear()

        self._setup_scene()
        self._setup_physical_materials()
        self._setup_render_materials()
        self._load_actors()
        self._load_articulations()
        self._load_agent()
        self._setup_cameras()
        self._setup_lighting()

        if self._viewer is not None:
            self._setup_viewer()

        # Cache actors and articulations
        self._actors = self.get_actors()
        self._articulations = self.get_articulations()
        # Cache initial simulation state
        # self._initial_sim_state = self.get_sim_state()

    def _clear(self):
        """Clear the simulation scene instance and other buffers.
        The function can be called in reset() before a new scene is created.
        """
        self._close_viewer()
        self._agent = None
        self._cameras = OrderedDict()
        self._scene = None
        self._physical_materials = {}
        self._render_materials = {}

    def _close_viewer(self):
        if self._viewer is not None:
            self._viewer.close()
            self._viewer = None

    def _get_default_scene_config(self):
        scene_config = sapien.SceneConfig()
        scene_config.default_dynamic_friction = 1.0
        scene_config.default_static_friction = 1.0
        scene_config.default_restitution = 0.0
        scene_config.contact_offset = 0.02
        scene_config.enable_pcm = False
        scene_config.solver_iterations = 25
        scene_config.solver_velocity_iterations = 0
        return scene_config

    def _setup_scene(self, scene_config: Optional[sapien.SceneConfig] = None):
        """Setup the simulation scene instance.
        The function should be called in reset(), and might be overrided.
        """
        if scene_config is None:
            scene_config = self._get_default_scene_config()
        self._scene = self._engine.create_scene(scene_config)
        self._scene.set_timestep(1.0 / self._sim_freq)

    # -------------------------------------------------------------------------- #
    # Utilities
    # -------------------------------------------------------------------------- #
    """Methods below are provided for convenience."""

    _physical_materials: Dict[str, sapien.PhysicalMaterial]
    _render_materials: Dict[str, sapien.RenderMaterial]

    def add_physical_material(self, name: str, *args, **kwargs):
        if name in self._physical_materials:
            warn(f"Overwrite physical material {name}")
        self._physical_materials[name] = self._scene.create_physical_material(
            *args, **kwargs
        )

    def add_render_material(self, name: str, **kwargs):
        if name in self._render_materials:
            warn(f"Overwrite render material {name}")
        m = self._renderer.create_material()
        for k, v in kwargs.items():
            if k == "color":
                m.set_base_color(v)
            else:
                getattr(m, f"set_{k}")(v)
        self._render_materials[name] = m

    def _setup_physical_materials(self):
        pass

    def _setup_render_materials(self):
        self.add_render_material(
            "ground",
            color=[0.5, 0.5, 0.5, 1],
            metallic=1.0,
            roughness=0.7,
            specular=0.04,
        )
        self.add_render_material(
            "default",
            color=[0.8, 0.8, 0.8, 1],
            metallic=0,
            roughness=0.9,
            specular=0.0,
        )

    def _add_ground(self, altitude=0.0):
        return self._scene.add_ground(
            altitude=altitude, render_material=self._render_materials["ground"]
        )

    def _create_box(
        self, pose: sapien.Pose, half_size, color=None, name=""
    ) -> sapien.Actor:
        half_size = np.array(half_size)
        builder: sapien.ActorBuilder = self._scene.create_actor_builder()
        builder.add_box_visual(half_size=half_size, color=color)
        box: sapien.Actor = builder.build_static(name=name)
        box.set_pose(pose)
        return box

    def _load_actors(self):
        color = np.array([0.1, 0.1, 0.1, 1.0])
        self._bg_front = self._create_box(
            pose=sapien.Pose([1, 0, 1], [1, 0, 0, 0]),
            half_size=np.float32([0.01, 1, 2]),
            color=color,
            name="background_front",
        )
        self._bg_left = self._create_box(
            pose=sapien.Pose([0, 1, 1], [1, 0, 0, 0]),
            half_size=np.float32([1, 0.01, 2]),
            color=color,
            name="background_left",
        )
        self._bg_right = self._create_box(
            pose=sapien.Pose([0, -1, 1], [1, 0, 0, 0]),
            half_size=np.float32([1, 0.01, 2]),
            color=color,
            name="background_right",
        )

    def _load_articulations(self):
        pass

    def _load_agent(self):
        pass

    def _setup_cameras(self):
        self._cameras = OrderedDict()

    def _setup_lighting(self):
        pass

    # -------------------------------------------------------------------------- #
    # Reset
    # -------------------------------------------------------------------------- #
    def reset(self, seed=None, reconfigure=False):
        self.set_episode_rng(seed)

        if reconfigure:  # Reconfigure the scene if assets change
            self.reconfigure()
        else:
            # self.set_sim_state(self._initial_sim_state)
            self._clear_sim_state()

        self.initialize_episode()
        self._agent.set_control_mode(self._control_mode)
        return self.get_obs()

    def set_episode_rng(self, seed):
        """Set the random generator for current episode."""
        if seed is None:
            self._episode_seed = self._main_rng.randint(2**32)
        else:
            self._episode_seed = seed
        self._episode_rng = np.random.RandomState(self._episode_seed)

    def initialize_episode(self):
        """Initialize the episode, e.g., poses of actors and articulations,
        and robot configuration. No new assets are created."""
        self._initialize_actors()
        self._initialize_articulations()
        self._initialize_agent()

    def _initialize_actors(self):
        pass

    def _initialize_articulations(self):
        pass

    def _initialize_agent(self):
        pass

    def _clear_sim_state(self):
        """Clear simulation state (velocities)"""
        for actor in self._scene.get_all_actors():
            if actor.type != "static":
                actor.set_velocity([0, 0, 0])
                actor.set_angular_velocity([0, 0, 0])
        for articulation in self._scene.get_all_articulations():
            articulation.set_qvel(np.zeros(articulation.dof))
            articulation.set_root_velocity([0, 0, 0])
            articulation.set_root_angular_velocity([0, 0, 0])

    # -------------------------------------------------------------------------- #
    # Step
    # -------------------------------------------------------------------------- #
    def step(self, action: Union[None, np.ndarray, Dict]):
        self.step_action(action)
        obs = self.get_obs()
        info = self.get_info()
        reward = self.get_reward(obs=obs, action=action, info=info)
        done = self.get_done()
        return obs, reward, done, info

    def step_action(self, action):
        if action is None:  # simulation without action
            pass
        elif isinstance(action, np.ndarray):
            self._agent.set_action(action)
        elif isinstance(action, dict):
            if action["control_mode"] != self._agent.control_mode:
                self._agent.set_control_mode(action["control_mode"])
            self._agent.set_action(action["action"])
        else:
            raise TypeError(type(action))

        for _ in range(self._sim_step_per_control):
            self._agent.simulation_step()
            self._scene.step()

        self._agent.update_generalized_external_forces()

    def check_success(self) -> bool:
        raise NotImplementedError

    def get_done(self):
        # return self.check_success()
        return False

    def get_info(self):
        return dict(is_success=self.check_success())

    # -------------------------------------------------------------------------- #
    # Simulation state (required for MPC)
    # -------------------------------------------------------------------------- #
    def get_actors(self):
        return self._scene.get_all_actors()

    def get_articulations(self):
        return [self._agent._robot]

    def get_sim_state(self) -> np.ndarray:
        """Get simulation state."""
        state = []
        for actor in self._actors:
            state.append(get_actor_state(actor))
        for articulation in self._articulations:
            state.append(get_articulation_state(articulation))
        return np.hstack(state)

    def set_sim_state(self, state: np.ndarray):
        """Set simulation state."""
        KINEMANTIC_DIM = 13  # [pos, quat, lin_vel, ang_vel]
        start = 0
        for actor in self._actors:
            set_actor_state(actor, state[start : start + KINEMANTIC_DIM])
            start += KINEMANTIC_DIM
        for articulation in self._articulations:
            ndim = KINEMANTIC_DIM + 2 * articulation.dof
            set_articulation_state(articulation, state[start : start + ndim])
            start += ndim

    def get_state(self):
        """Get environment state. Override to include task information (e.g., goal)"""
        return self.get_sim_state()

    def set_state(self, state: np.ndarray):
        """Set environment state. Override to include task information (e.g., goal)"""
        return self.set_sim_state(state)

    # -------------------------------------------------------------------------- #
    # Visualization
    # -------------------------------------------------------------------------- #
    """Due to some known bugs of sapien.utils.Viewer, do not close it after created."""
    _viewer: Viewer
    render_camera: sapien.CameraEntity  # should be created in `_setup_cameras`

    def _setup_viewer(self):
        """Setup the interactive viewer.
        The function should be called in reset(), and overrided to adjust camera.
        """
        # CAUTION: call `set_scene` after assets are loaded.
        self._viewer.set_scene(self._scene)
        self._viewer.toggle_axes(False)
        self._viewer.toggle_camera_lines(False)

    def update_render(self):
        """Update renderer(s). This function should be called before any rendering,
        to sync simulator and renderer."""
        self._scene.update_render()

    def render(self, mode="human"):
        if mode == "human":
            self.update_render()
            if self._viewer is None:
                self._viewer = Viewer(self._renderer)
                self._setup_viewer()
            self._viewer.render()
            return self._viewer
        if mode == "rgb_array":
            return super().render(mode=mode)
        elif mode == "cameras":
            images = [self.render("rgb_array")]
            self.update_render()
            for obs_dict in self._get_images().values():
                images.extend(observations_to_images(obs_dict))
            return tile_images(images)
        else:
            raise NotImplementedError(f"Unsupported render mode {mode}.")

    def gen_scene_pcd(self, num_points: int = int(1e5)) -> np.ndarray:
        """generate scene point cloud for motion planning, excluding the robot"""
        meshes = []
        articulations = self._scene.get_all_articulations()
        if self._agent is not None:
            articulations.pop(articulations.index(self._agent._robot))
        for articulation in articulations:
            articulation_mesh = merge_meshes(get_articulation_meshes(articulation))
            if articulation_mesh:
                meshes.append(articulation_mesh)

        for actor in self._scene.get_all_actors():
            actor_mesh = merge_meshes(get_actor_meshes(actor))
            if actor_mesh:
                meshes.append(
                    actor_mesh.apply_transform(
                        actor.get_pose().to_transformation_matrix()
                    )
                )

        scene_mesh = merge_meshes(meshes)
        scene_pcd = scene_mesh.sample(num_points)
        return scene_pcd

arti_mani/envs/base.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from collections import OrderedDict
from typing import Dict, Tuple, Union

import numpy as np
import sapien.core as sapien
from arti_mani import AGENT_CONFIG_DIR, ASSET_DIR
from arti_mani.agents.camera import get_camera_rgb
from arti_mani.agents.fixed_xmate3_robotiq import FixedXmate3Robotiq
from arti_mani.agents.panda import Panda
from arti_mani.envs.base import BaseEnv
from arti_mani.utils.articulation_utils import (
    get_actor_state,
    get_articulation_state,
    set_actor_state,
    set_articulation_state,
)
from arti_mani.utils.geometry import angle_distance
from transforms3d.euler import euler2quat


class FixedXmate3RobotiqEnv(BaseEnv):
    SUPPORTED_OBS_MODES: Tuple[str] = ()
    SUPPORTED_REWARD_MODES: Tuple[str] = ()
    SUPPORTED_CONTROL_MODES: Tuple[str] = ()
    _agent: FixedXmate3Robotiq

    def __init__(
        self,
        obs_mode=None,
        reward_mode=None,
        control_mode=None,
        sim_freq=500,
        control_freq=20,
        device="",
    ):
        self._articulation = None
        self.step_in_ep = 0
        self._prev_actor_pose = sapien.Pose()
        self._cache_obs_state_dict: OrderedDict = (
            OrderedDict()
        )  # save obs state dict to save reward computation time
        self._cache_info = {}
        super().__init__(
            obs_mode, reward_mode, control_mode, sim_freq, control_freq, device
        )

    # -------------------------------------------------------------------------- #
    # Reset
    # -------------------------------------------------------------------------- #

    def _setup_physical_materials(self):
        self.add_physical_material("default", 1.0, 1.0, 0.0)

    def _setup_render_materials(self):
        self.add_render_material(
            "ground",
            color=[0.5, 0.5, 0.5, 1],
            metallic=1.0,
            roughness=0.7,
            specular=0.04,
        )
        self.add_render_material(
            "default",
            color=[0.8, 0.8, 0.8, 1],
            metallic=0,
            roughness=0.9,
            specular=0.0,
        )

    def _load_articulations(self):
        pass

    def _load_agent(self):
        self._agent = FixedXmate3Robotiq.from_config_file(
            AGENT_CONFIG_DIR / "fixed_xmate3_robotiq.yml",
            self._scene,
            self._control_freq,
        )

    def _initialize_articulations(self):
        pass

    def _initialize_agent(self):
        raise NotImplementedError

    def _load_table(self):
        loader = self._scene.create_actor_builder()
        loader.add_visual_from_file(
            str(ASSET_DIR / "descriptions/optical_table/visual/optical_table.dae"),
            # scale=[3, 3, 1]
        )
        loader.add_collision_from_file(
            str(ASSET_DIR / "descriptions/optical_table/visual/optical_table.dae"),
            # scale=[3, 3, 1]
        )
        self._table = loader.build_static(name="table")
        self._table.set_pose(sapien.Pose([0.0, 0.0, -0.04]))

    def reconfigure(self):
        self._prev_actor_pose = sapien.Pose()
        self._clear()

        self._setup_scene()
        self._setup_physical_materials()
        self._setup_render_materials()
        self._load_table()
        self._load_actors()
        self._load_agent()
        self._load_articulations()
        self._setup_camera()
        self._setup_lighting()

        if self._viewer is not None:
            self._setup_viewer()

        # cache actors and articulations for sim state
        self._actors = self.get_actors()
        self._articulations = self.get_articulations()
        # Cache initial simulation state
        # self._initial_sim_state = self.get_sim_state()

    def reset(self, seed=None, reconfigure=False):
        super().reset(seed, reconfigure)
        self.step_in_ep = 0
        # for actor in self._agent._robot.get_links():
        #     make_actor_visible(actor, False)
        self._cache_obs_state_dict.clear()
        self._cache_info.clear()
        return self.get_obs()

    # -------------------------------------------------------------------------- #
    # Step
    # -------------------------------------------------------------------------- #
    def check_success(self) -> bool:
        raise NotImplementedError

    def get_done(self):
        # return self.check_success()
        return False

    def get_info(self):
        return dict(is_success=self.check_success())

    def step(self, action: Union[None, np.ndarray, Dict]):
        self.step_in_ep += 1
        self.step_action(action)
        obs = self.get_obs()
        reward = self.get_reward(obs=obs, action=action)
        info = self.get_info()
        done = self.get_done()
        return obs, reward, done, info

    # ---------------------------------------------------------------------------- #
    # Visualization
    # ---------------------------------------------------------------------------- #
    def _setup_lighting(self):
        self._scene.set_ambient_light([0.3, 0.3, 0.3])
        self._scene.add_point_light([2, 2, 2], [1, 1, 1])
        self._scene.add_point_light([2, -2, 2], [1, 1, 1])
        self._scene.add_point_light([-2, 0, 2], [1, 1, 1])
        self._scene.add_directional_light([1, -1, -1], [0.3, 0.3, 0.3])
        self._scene.add_directional_light([0, 0, -1], [1, 1, 1])

    def _setup_viewer(self):
        super()._setup_viewer()
        ### front view
        # self._viewer.set_camera_xyz(1.0, 0.0, 1.2)
        # self._viewer.set_camera_rpy(0, -0.5, 3.14)
        ### right side view
        # self._viewer.set_camera_xyz(0, -1.0, 1.2)
        # self._viewer.set_camera_rpy(0, -0.5, -1.57)
        ### top view
        self._viewer.set_camera_xyz(0.0, 0.0, 1.2)
        self._viewer.set_camera_rpy(0, -3.14, 0)

    def _setup_camera(self):
        # self.render_camera = self._scene.add_camera("sideview", 512, 512, 1, 0.01, 10)
        # self.render_camera.set_local_pose(
        #     sapien.Pose([0, -1.0, 1.2], euler2quat(0, 0.5, 1.57))
        # )
        self.render_camera = self._scene.add_camera("topview", 512, 512, 1, 0.01, 10)
        self.render_camera.set_local_pose(
            sapien.Pose([-0.1, 0.0, 1.5], euler2quat(0, 1.57, 0))
        )

    def render(self, mode="human"):
        if mode == "rgb_array":
            self._scene.update_render()
            self.render_camera.take_picture()
            rgb = get_camera_rgb(self.render_camera)
            # rgb = self.render_camera.get_color_rgba()[..., :3]
            # rgb = np.clip(rgb * 255, 0, 255).astype(np.uint8)
            return rgb
        else:
            return super().render(mode=mode)

    # utilization
    def check_actor_static(self, actor, max_v=None, max_ang_v=None):
        if self.step_in_ep <= 1:
            flag_v = (max_v is None) or (np.linalg.norm(actor.get_velocity()) <= max_v)
            flag_ang_v = (max_ang_v is None) or (
                np.linalg.norm(actor.get_angular_velocity()) <= max_ang_v
            )
        else:
            pose = actor.get_pose()
            t = 1.0 / self._control_freq
            flag_v = (max_v is None) or (
                np.linalg.norm(pose.p - self._prev_actor_pose.p) <= max_v * t
            )
            flag_ang_v = (max_ang_v is None) or (
                angle_distance(self._prev_actor_pose, pose) <= max_ang_v * t
            )
        self._prev_actor_pose = actor.get_pose()
        return flag_v and flag_ang_v

    # -------------------------------------------------------------------------- #
    # Simulation state (required for MPC)
    # -------------------------------------------------------------------------- #
    def get_actors(self):
        return [x for x in self._scene.get_all_actors()]

    def get_articulations(self):
        return [self._agent._robot, self._articulation]

    def get_sim_state(self) -> np.ndarray:
        state = []
        # for actor in self._scene.get_all_actors():
        for actor in self._actors:
            state.append(get_actor_state(actor))
        # for articulation in self._scene.get_all_articulations():
        for articulation in self._articulations:
            state.append(get_articulation_state(articulation))
        return np.hstack(state)

    def set_sim_state(self, state: np.ndarray):
        KINEMANTIC_DIM = 13  # [pos, quat, lin_vel, ang_vel]
        start = 0
        for actor in self._actors:
            set_actor_state(actor, state[start : start + KINEMANTIC_DIM])
            start += KINEMANTIC_DIM
        for articulation in self._articulations:
            ndim = KINEMANTIC_DIM + 2 * articulation.dof
            set_articulation_state(articulation, state[start : start + ndim])
            start += ndim

    def get_state(self):
        return self.get_sim_state()

    def set_state(self, state: np.ndarray):
        return self.set_sim_state(state)


class PandaEnv(FixedXmate3RobotiqEnv):
    _agent: Panda

    def _load_agent(self):
        self._agent = Panda.from_config_file(
            AGENT_CONFIG_DIR / "panda.yml", self._scene, self._control_freq
        )


class FixedXmate3RobotiqSensorEnv(FixedXmate3RobotiqEnv):
    def _load_agent(self):
        self._agent = FixedXmate3Robotiq.from_config_file(
            AGENT_CONFIG_DIR / "fixed_xmate3_robotiq_sensors.yml",
            self._scene,
            self._control_freq,
        )


class FixedXmate3RobotiqSensorLowResEnv(FixedXmate3RobotiqEnv):
    def _load_agent(self):
        self._agent = FixedXmate3Robotiq.from_config_file(
            AGENT_CONFIG_DIR / "fixed_xmate3_robotiq_sensors_low_res.yml",
            self._scene,
            self._control_freq,
        )


class FixedXmate3RobotiqLowResEnv(FixedXmate3RobotiqEnv):
    def _load_agent(self):
        self._agent = FixedXmate3Robotiq.from_config_file(
            AGENT_CONFIG_DIR / "fixed_xmate3_robotiq_low_res.yml",
            self._scene,
            self._control_freq,
        )

arti_mani/envs/fixed_robotenv.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
from collections import OrderedDict
from typing import List, Tuple

import numpy as np
import sapien.core as sapien
import yaml
from arti_mani import AGENT_CONFIG_DIR, ASSET_DIR
from arti_mani.agents.camera import get_camera_rgb, get_camera_seg, get_texture
from arti_mani.agents.floating_robotiq import FloatingRobotiq
from arti_mani.envs.base import BaseEnv
from arti_mani.utils.articulation_utils import (
    get_actor_state,
    get_articulation_state,
    get_entity_by_name,
    ignore_collision,
    set_actor_state,
    set_articulation_state,
)
from arti_mani.utils.common import (
    convert_np_bool_to_float,
    np_random,
    random_choice,
    register_gym_env,
)
from arti_mani.utils.contrib import apply_pose_to_points, o3d_to_trimesh, xyz2uv
from arti_mani.utils.geometry import angle_distance, transform_points
from arti_mani.utils.o3d_utils import merge_mesh, np2mesh
from arti_mani.utils.sapien_utils import (
    get_entity_by_name,
    hex2rgba,
    look_at,
    set_actor_render_material,
    set_articulation_render_material,
    set_cabinet_render_material,
)
from arti_mani.utils.trimesh_utils import get_actor_mesh
from transforms3d.euler import euler2quat


class FloatingRobotiqEnv(BaseEnv):
    SUPPORTED_OBS_MODES: Tuple[str] = ()
    SUPPORTED_REWARD_MODES: Tuple[str] = ()
    SUPPORTED_CONTROL_MODES: Tuple[str] = ()
    _agent: FloatingRobotiq

    def __init__(
        self,
        obs_mode=None,
        reward_mode=None,
        control_mode=None,
        irsensor_mode=False,
        sim_freq=500,
        control_freq=20,
    ):
        self._articulation = None
        self.step_in_ep = 0
        self._prev_actor_pose = sapien.Pose()
        self._cache_obs_state_dict: OrderedDict = (
            OrderedDict()
        )  # save obs state dict to save reward computation time
        self._cache_info = {}
        self.irsensor_mode = irsensor_mode
        self.height, self.width = 144, 256
        self.texture_filelist = os.listdir(ASSET_DIR / "textures/random_textures")
        self.colortexture_filelist = os.listdir(ASSET_DIR / "textures/color_textures")
        super().__init__(obs_mode, reward_mode, control_mode, sim_freq, control_freq)

    # -------------------------------------------------------------------------- #
    # Reset
    # -------------------------------------------------------------------------- #

    def _setup_physical_materials(self):
        self.add_physical_material("default", 1.0, 1.0, 0.0)

    def _setup_render_materials(self):
        self.add_render_material(
            "ground",
            color=[0.5, 0.5, 0.5, 1],
            metallic=1.0,
            roughness=0.7,
            specular=0.04,
        )
        self.add_render_material(
            "default",
            color=[0.8, 0.8, 0.8, 1],
            metallic=0,
            roughness=0.9,
            specular=0.0,
        )

    def _random_lighting(self):
        for light in self._scene.get_all_lights():
            if type(light) in [sapien.PointLightEntity, sapien.DirectionalLightEntity]:
                self._scene.remove_light(light)
        pointlight_num = self._episode_rng.choice(np.arange(1, 4))
        direclight_num = self._episode_rng.choice(np.arange(1, 4))
        for num in range(pointlight_num):
            self._scene.add_point_light(
                [
                    self._episode_rng.uniform(-1, 1),
                    self._episode_rng.uniform(-0.5, 0.5),
                    self._episode_rng.uniform(1, 3),
                ],
                [self._episode_rng.uniform(0.5, 1)] * 3,
                shadow=True,
            )
        for num in range(direclight_num):
            self._scene.add_directional_light(
                [
                    self._episode_rng.uniform(0.1, 1),
                    self._episode_rng.uniform(-1, 1),
                    self._episode_rng.uniform(-1, -0.25),
                ],
                [self._episode_rng.uniform(0.5, 1)] * 3,
                shadow=True,
            )

    def _domain_random(self):
        # self._random_lighting()
        if self._episode_rng.rand() > 0.2:
            rand_img_file = self._episode_rng.choice(self.texture_filelist, 1)[0]
            set_actor_render_material(
                self._table,
                metallic=self._episode_rng.uniform(0.4, 0.9),  # 1, 0
                roughness=self._episode_rng.uniform(0.4, 0.9),  # 0.4, 0.9
                specular=self._episode_rng.uniform(0.01, 0.8),  # /, 0.0
                diffuse_texture_filename=str(
                    ASSET_DIR / f"textures/random_textures/{rand_img_file}"
                ),
            )

    def _load_articulations(self):
        pass

    def _load_agent(self):
        self._agent = FloatingRobotiq.from_config_file(
            AGENT_CONFIG_DIR / "floating_robotiq_low_res.yml",
            self._scene,
            self._control_freq,
        )

    def _initialize_articulations(self):
        pass

    def _initialize_agent(self):
        qpos = self._episode_rng.uniform(0, 0.068, [2])
        self._agent._robot.set_qpos(qpos)
        self._agent._robot.set_pose(sapien.Pose([-0.6, 0.4, 0], [1, 0, 0, 0]))

    def _load_table(self):
        loader = self._scene.create_actor_builder()
        loader.add_visual_from_file(
            str(ASSET_DIR / "descriptions/optical_table/visual/optical_table.dae"),
            # scale=[3, 3, 1]
        )
        loader.add_collision_from_file(
            str(ASSET_DIR / "descriptions/optical_table/visual/optical_table.dae"),
            # scale=[3, 3, 1]
        )
        self._table = loader.build_static(name="table")
        self._table.set_pose(sapien.Pose([0.0, 0.0, -0.04]))

    def reconfigure(self):
        self._prev_actor_pose = sapien.Pose()
        self._clear()

        self._setup_scene()
        self._setup_physical_materials()
        self._setup_render_materials()
        self._load_table()
        self._load_actors()
        self._load_articulations()
        self._load_agent()
        self._setup_camera()
        self._setup_lighting()

        if self._viewer is not None:
            self._setup_viewer()

        # cache actors and articulations for sim state
        self._actors = self.get_actors()
        self._articulations = self.get_articulations()
        # Cache initial simulation state
        # self._initial_sim_state = self.get_sim_state()

    def reset(self, seed=None, reconfigure=False):
        super().reset(seed, reconfigure)
        self.step_in_ep = 0
        # for actor in self._agent._robot.get_links():
        #     make_actor_visible(actor, False)
        self._cache_obs_state_dict.clear()
        self._cache_info.clear()
        return self.get_obs()

    # ---------------------------------------------------------------------------- #
    # Visualization
    # ---------------------------------------------------------------------------- #
    def _setup_lighting(self):
        self._scene.set_ambient_light([0.3, 0.3, 0.3])
        self._scene.add_point_light([2, 2, 2], [1, 1, 1])
        self._scene.add_point_light([2, -2, 2], [1, 1, 1])
        self._scene.add_point_light([-2, 0, 2], [1, 1, 1])
        self._scene.add_point_light([1, -1, 1], [1, 1, 1])
        self._scene.add_directional_light([1, -1, -1], [0.3, 0.3, 0.3])
        self._scene.add_directional_light([0, 0, -1], [1, 1, 1])

    def _setup_viewer(self):
        super()._setup_viewer()
        ### top view
        self._viewer.set_camera_xyz(0.0, 0.0, 1.2)
        self._viewer.set_camera_rpy(0, -3.14, 0)

    def _setup_camera(self):
        # self.render_camera = self._scene.add_camera("sideview", 512, 512, 1, 0.01, 10)
        # self.render_camera.set_local_pose(
        #     sapien.Pose([0, -1.0, 1.2], euler2quat(0, 0.5, 1.57))
        # )
        self.render_camera = self._scene.add_camera("topview", 512, 512, 1, 0.01, 10)
        self.render_camera.set_local_pose(
            sapien.Pose([-0.1, 0.0, 1.5], euler2quat(0, 1.57, 0))
        )

    def render(self, mode="human"):
        if mode == "rgb_array":
            self._scene.update_render()
            self.render_camera.take_picture()
            rgb = get_camera_rgb(self.render_camera)
            # rgb = self.render_camera.get_color_rgba()[..., :3]
            # rgb = np.clip(rgb * 255, 0, 255).astype(np.uint8)
            return rgb
        else:
            return super().render(mode=mode)

    # utilization
    def check_actor_static(self, actor, max_v=None, max_ang_v=None):
        if self.step_in_ep <= 1:
            flag_v = (max_v is None) or (np.linalg.norm(actor.get_velocity()) <= max_v)
            flag_ang_v = (max_ang_v is None) or (
                np.linalg.norm(actor.get_angular_velocity()) <= max_ang_v
            )
        else:
            pose = actor.get_pose()
            t = 1.0 / self._control_freq
            flag_v = (max_v is None) or (
                np.linalg.norm(pose.p - self._prev_actor_pose.p) <= max_v * t
            )
            flag_ang_v = (max_ang_v is None) or (
                angle_distance(self._prev_actor_pose, pose) <= max_ang_v * t
            )
        self._prev_actor_pose = actor.get_pose()
        return flag_v and flag_ang_v

    # -------------------------------------------------------------------------- #
    # Simulation state (required for MPC)
    # -------------------------------------------------------------------------- #
    def get_actors(self):
        return [x for x in self._scene.get_all_actors()]

    def get_articulations(self):
        return [self._agent._robot, self._articulation]

    def get_sim_state(self) -> np.ndarray:
        state = []
        # for actor in self._scene.get_all_actors():
        for actor in self._actors:
            state.append(get_actor_state(actor))
        # for articulation in self._scene.get_all_articulations():
        for articulation in self._articulations:
            state.append(get_articulation_state(articulation))
        return np.hstack(state)

    def set_sim_state(self, state: np.ndarray):
        KINEMANTIC_DIM = 13  # [pos, quat, lin_vel, ang_vel]
        start = 0
        for actor in self._actors:
            set_actor_state(actor, state[start : start + KINEMANTIC_DIM])
            start += KINEMANTIC_DIM
        for articulation in self._articulations:
            ndim = KINEMANTIC_DIM + 2 * articulation.dof
            set_articulation_state(articulation, state[start : start + ndim])
            start += ndim

    def get_state(self):
        return self.get_sim_state()

    def set_state(self, state: np.ndarray):
        return self.set_sim_state(state)


@register_gym_env("CabinetDoor-v0", max_episode_steps=200)
class CabinetDoor(FloatingRobotiqEnv):
    SUPPORTED_OBS_MODES = ("state", "state_dict", "state_egorgbd", "egorgbd_pc")
    SUPPORTED_REWARD_MODES = ("dense", "sparse")
    SUPPORTED_CONTROL_MODES = ("pd_joint_delta_pos", "pd_ee_delta_pose")
    DEFAULT_MODEL_PATH = (
        ASSET_DIR / "partnet_mobility_configs/fixed_cabinet_doors_new.yml"
    )

    def __init__(
        self,
        articulation_ids: List[int] = (),
        articulation_config_path="",
        obs_mode=None,
        reward_mode=None,
        irsensor_mode=False,
        other_handle_visible=False,
        domain_random=False,
        sim_freq=500,
        control_freq=20,
        control_mode=None,
    ):
        if articulation_config_path is None:
            articulation_config_path = self.DEFAULT_MODEL_PATH
        with open(articulation_config_path, "r") as f:
            self._articulation_info = yaml.safe_load(f)
        if isinstance(articulation_ids, int):
            articulation_ids = [articulation_ids]
        if len(articulation_ids) == 0:
            articulation_ids = sorted(self._articulation_info.keys())
        self.articulation_ids = articulation_ids
        self.articulation_id = None
        self.other_handle_visible = other_handle_visible
        self.domain_random = domain_random

        super().__init__(
            obs_mode,
            reward_mode,
            control_mode,
            irsensor_mode,
            sim_freq,
            control_freq,
        )

    def _domain_random(self):
        self._scene.remove_articulation(self._articulation)
        self._load_articulations()
        self._initialize_articulations()
        if self._episode_rng.rand() > 0.2:
            color_rand_val = self._episode_rng.uniform(0, 1)
            a_rand_val = self._episode_rng.uniform(0.9, 1)
            rand_img_file = self._episode_rng.choice(self.texture_filelist, 1)[0]
            rand_color_file = self._episode_rng.choice(self.colortexture_filelist, 1)[0]
            set_cabinet_render_material(
                self._articulation,
                # color=[color_rand_val, color_rand_val, color_rand_val, a_rand_val],
                color=[
                    self._episode_rng.uniform(0, 1),
                    self._episode_rng.uniform(0, 1),
                    self._episode_rng.uniform(0, 1),
                    1,
                ],  # hex2rgba("#AAAAAA")<=>[0.40982574, 0.40982574, 0.40982574, 1.], [0.8, 0.8, 0.8, 1]
                color_texture_filename=str(
                    ASSET_DIR / f"textures/color_textures/{rand_color_file}"
                ),
                metallic=self._episode_rng.uniform(0, 0.5),  # 1, 0
                roughness=self._episode_rng.uniform(0, 1),  # 0.4, 0.9
                specular=self._episode_rng.uniform(0, 0.8),  # /, 0.0
                diffuse_texture_filename=str(
                    ASSET_DIR / f"textures/random_textures/{rand_img_file}"
                ),
            )
        super()._domain_random()

    def _initialize_agent(self):
        qpos = self._episode_rng.uniform(0, 0.068, [2])
        self._agent._robot.set_qpos(qpos)
        self._agent._robot.set_pose(
            look_at(
                eye=[-0.3, 0, 0.5],
                target=self._articulation.get_root_pose().p,
                up=[0, 0, 1],
            )
        )

    def _load_cabinet(self):
        # with open(self._articulation_config_path.joinpath(self.articulation_id+".yml"), "r") as f:
        #     config_dict = yaml.safe_load(f)
        #     self._articulation_config = ArticulationConfig(**config_dict)
        loader = self._scene.create_urdf_loader()
        loader.load_multiple_collisions_from_file = self._articulation_config[
            "multiple_collisions"
        ]
        loader.scale = self._articulation_config["scale"]
        loader.fix_root_link = True

        articulation_path = (
            ASSET_DIR / f"partnet_mobility_dataset/{self.articulation_id}"
        )
        urdf_path = articulation_path / "mobility_cvx.urdf"
        assert (
            urdf_path.exists()
        ), f"{urdf_path} is not found. Please download Partnet-Mobility (ManiSkill2022) first."
        config = {
            "material": self._physical_materials["default"],
            # "density": 8e3
        }
        articulation = loader.load(str(urdf_path), config=config)
        articulation.set_name("cabinet")
        return articulation

    def _load_articulations(self):
        self._articulation = self._load_cabinet()

        # set physical properties for all the joints
        joint_friction = self._episode_rng.uniform(
            self._articulation_config["joint_friction_range"][0],
            self._articulation_config["joint_friction_range"][1],
        )
        joint_stiffness = self._episode_rng.uniform(
            self._articulation_config["joint_stiffness_range"][0],
            self._articulation_config["joint_stiffness_range"][1],
        )
        joint_damping = self._episode_rng.uniform(
            self._articulation_config["joint_damping_range"][0],
            self._articulation_config["joint_damping_range"][1],
        )
        for joint in self._articulation.get_active_joints():
            joint.set_friction(joint_friction)
            joint.set_drive_property(joint_stiffness, joint_damping)
        ignore_collision(self._articulation)

        self._set_semantic_visid()

    def _set_semantic_visid(self):
        all_links = self._articulation.get_links()
        all_active_joints = self._articulation.get_active_joints()

        # set target link & joint, handle & door of target_link
        self._target_joint_idx = self._articulation_config["target_joint_idx"]
        self._target_joint = all_active_joints[self._target_joint_idx]
        self._target_link = self._target_joint.get_child_link()
        # all_links.remove(self._target_link)

        door_links = []
        for joint in all_active_joints:
            link = joint.get_child_link()
            door_links.append(link)

        self.handles_visid = []
        self.doors_visid = []
        self.cabinet_visid = []
        for link in all_links:
            if link in door_links:
                for visual_body in link.get_visual_bodies():
                    if "handle" in visual_body.get_name():
                        self.handles_visid.append(visual_body.get_visual_id())
                        if link != self._target_link:
                            visual_body.set_visibility(
                                1.0 if self.other_handle_visible else 0.0
                            )
                    else:
                        self.doors_visid.append(visual_body.get_visual_id())
            else:
                for visual_body in link.get_visual_bodies():
                    self.cabinet_visid.append(visual_body.get_visual_id())

    def _initialize_articulations(self):
        # sim init
        posz = (
            -self._articulation_config["scale"]
            * self._articulation_config["bbox_min"][2]
        )
        if self.articulation_id == 0:
            posx, posy, rotz = 0.152, -0.188, 0.7
        elif self.articulation_id == 1:
            posx, posy, rotz = 0.237, -0.14, 0
        else:
            posx, posy, rotz = 0, 0, 0
        arti_root_pose = sapien.Pose(
            [posx + 0.6, posy, posz], [np.sqrt(1 - rotz**2), 0, 0, rotz]
        )
        if self._articulation_config["flip"]:
            arti_root_pose = arti_root_pose * sapien.Pose(q=np.array([0, 1, 0, 0]))
        self._articulation.set_root_pose(arti_root_pose)

        # init_open_extent = self._episode_rng.uniform(0, self._articulation_config["init_open_extent_range"])
        init_open_extent = self._episode_rng.uniform(0, 0.1, self._articulation.dof)
        qpos = np.zeros(self._articulation.dof)
        for i in range(self._articulation.dof):
            [[lmin, lmax]] = self._articulation.get_active_joints()[i].get_limits()
            qpos[i] = lmin + (lmax - lmin) * init_open_extent[i]
            # print(f"++++ {self.articulation_id}, activejoint_{i}", lmin, lmax)
        ## sim init
        self._articulation.set_qpos(qpos)

        # [[lmin, lmax]] = self._target_joint.get_limits()
        lmin, lmax = self._articulation_config["target_joint_range"]
        # print("++++", lmin, lmax)
        self.target_qpos = (
            lmin + (lmax - lmin) * self._articulation_config["open_extent"]
        )

        self._get_handle_info_in_target_link()

    def _get_handle_info_in_target_link(self):
        handle_meshes = []
        for visual_body in self._target_link.get_visual_bodies():
            for render_shape in visual_body.get_render_shapes():
                vertices = apply_pose_to_points(
                    render_shape.mesh.vertices * visual_body.scale,
                    visual_body.local_pose,
                )
                shape_mesh = np2mesh(vertices, render_shape.mesh.indices.reshape(-1, 3))
                if "handle" in visual_body.get_name():
                    handle_meshes.append(shape_mesh)
        handle_mesh = merge_mesh(handle_meshes)
        handle_trimesh = o3d_to_trimesh(handle_mesh)
        handle_pcd = handle_trimesh.sample(300)  ## target_link frame, (300, 3)
        # get tip, med, bottom keypoints
        xyzmax_val, xyzmin_val = handle_pcd.max(0), handle_pcd.min(0)
        xyzmax_id, xyzmin_id = np.argmax(handle_pcd, 0), np.argmin(handle_pcd, 0)
        bbox_size = xyzmax_val - xyzmin_val
        long_side = np.argmax(bbox_size)
        tip, bottom = handle_pcd[xyzmax_id[long_side]], handle_pcd[xyzmin_id[long_side]]
        center = handle_pcd.mean(0)
        if self.articulation_id in [0, 1]:
            center[1] = xyzmax_val[1]
        else:
            center[2] = xyzmax_val[2]
        self._handle_keypoints = np.stack([tip, center, bottom])
        self._handle_pcd = handle_pcd

    def reset(self, seed=None, reconfigure=False, articulation_id=None):
        self.set_episode_rng(seed)
        _reconfigure = self._set_model(articulation_id)
        reconfigure = _reconfigure or reconfigure
        ret = super().reset(seed=self._episode_seed, reconfigure=reconfigure)
        return ret

    def _set_model(self, articulation_id):
        """Set the model id and scale. If not provided, choose one randomly."""
        reconfigure = False

        # Model ID
        if articulation_id is None:
            articulation_id = random_choice(self.articulation_ids, self._episode_rng)
        if articulation_id != self.articulation_id:
            reconfigure = True
        self.articulation_id = articulation_id
        self._articulation_config = self._articulation_info[self.articulation_id]

        return reconfigure

    def get_obs_state_egorgbd(self):
        ### get state
        state_rgbd = self._get_obs_state_dict()

        ### get visual info
        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        hand_sensor.compute_depth()

        rgb = (hand_sensor.get_rgb() * 255).astype(np.uint8)  # [H, W, 3]
        ir_l, ir_r = hand_sensor.get_ir()  # [H, W]
        depth = hand_sensor.get_depth()  # [H, W]
        clean_depth = -get_texture(sensor_rgb, "Position")[..., 2]  # [H, W]
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        imh, imw = visual_id_seg.shape

        state_rgbd["rgb"] = rgb.transpose((2, 0, 1))
        if self.irsensor_mode:
            state_rgbd["ir_l"] = ir_l.astype(np.float32)
            state_rgbd["ir_r"] = ir_r.astype(np.float32)
            state_rgbd["depth"] = depth.astype(np.float32)
            state_rgbd["clean_depth"] = clean_depth.astype(np.float32)
        else:
            state_rgbd["clean_depth"] = depth.astype(np.float32)

        seg_visids = [self.handles_visid, self.doors_visid, self.cabinet_visid]
        cam_seg = np.ones(visual_id_seg.shape, dtype=np.float32) * 5
        seg_id = [0, 1, 2]
        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros(visual_id_seg.shape, dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
        state_rgbd["seg"] = convert_np_bool_to_float(cam_seg).astype(np.uint8)  # (H, W)
        background_seg = np.zeros(visual_id_seg.shape, dtype=np.bool)
        bg_visids = [
            0,
            self._bg_front.get_id(),
            self._bg_left.get_id(),
            self._bg_right.get_id(),
        ]
        for bg_visid in bg_visids:
            background_seg[visual_id_seg == bg_visid] = True
        state_rgbd["bg_seg"] = convert_np_bool_to_float(background_seg).astype(
            np.uint8
        )  # (H, W)
        # got pcd target_link => world => ee
        # handle_pcd_world = apply_pose_to_points(self._handle_pcd, self._target_link.get_pose())  ## world frame
        kpts_world = apply_pose_to_points(
            self._handle_keypoints, self._target_link.get_pose()
        )  ## world frame
        # handle_pcd_ee = apply_pose_to_points(handle_pcd_world, self.agent.grasp_site.get_pose().inv())  ## ee frame
        # kpts_ee = apply_pose_to_points(handle_keypoints_world, self.agent.grasp_site.get_pose().inv())  ## ee frame
        # kpts_cam = apply_pose_to_points(kpts_ee, sensor_rgb.get_pose().inv())  ## cam frame
        kpts_cam = transform_points(sensor_rgb.get_extrinsic_matrix(), kpts_world)
        uvz = xyz2uv(kpts_cam, sensor_rgb.get_camera_matrix())
        u_visable = (uvz[:, 0] >= 0) & (uvz[:, 0] <= imw - 1)
        v_visable = (uvz[:, 1] >= 0) & (uvz[:, 1] <= imh - 1)
        uvz_visable = np.stack(
            [u_visable, v_visable, u_visable & v_visable]
        ).T  # (3, 3)
        for ind in range(uvz_visable.shape[0]):
            if uvz_visable[ind, 2]:  # not out of img range
                u, v = round(uvz[ind, 0]), round(uvz[ind, 1])
                if cam_seg[v, u] not in seg_id:  # handle keypoint is occluded
                    uvz_visable[ind, 2] = False
        state_rgbd["kpts"] = kpts_cam.astype(np.float32)
        state_rgbd["uvz"] = uvz.astype(np.float32)
        state_rgbd["uvz_visable"] = uvz_visable
        return state_rgbd

    def get_obs_egorgbd_pc(self):
        ### get state
        rgbd_pc = OrderedDict()

        ### get visual info
        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        hand_sensor.compute_depth()

        rgb = (hand_sensor.get_rgb() * 255).astype(np.uint8)  # [H, W, 3]
        ir_l, ir_r = hand_sensor.get_ir()  # [H, W]
        depth = hand_sensor.get_depth()  # [H, W]
        clean_depth = -get_texture(sensor_rgb, "Position")[..., 2]  # [H, W]
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        imh, imw = visual_id_seg.shape

        # process point cloud
        cam_xyz = hand_sensor.get_pointcloud(False)  # [H*W, 3]
        ## cam to world
        T_cam2world = sensor_rgb.get_model_matrix()
        world_xyz = transform_points(
            T_cam2world, cam_xyz * np.array([1, -1, -1])
        )  # [H*W, 3]
        T_world2ee = (
            self.agent.grasp_site.get_pose().inv().to_transformation_matrix()
        )  ## world => ee
        rgbd_pc["world_xyz"] = world_xyz

        rgbd_pc["rgb"] = rgb.transpose((2, 0, 1))
        if self.irsensor_mode:
            rgbd_pc["ir_l"] = ir_l.astype(np.float32)
            rgbd_pc["ir_r"] = ir_r.astype(np.float32)
            rgbd_pc["depth"] = depth.astype(np.float32)
            rgbd_pc["clean_depth"] = clean_depth.astype(np.float32)
        else:
            rgbd_pc["clean_depth"] = depth.astype(np.float32)

        seg_visids = [self.handles_visid, self.doors_visid, self.cabinet_visid]
        cam_seg = np.ones(visual_id_seg.shape, dtype=np.float32) * 5
        seg_id = [0, 1, 2]
        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros(visual_id_seg.shape, dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
        rgbd_pc["seg"] = convert_np_bool_to_float(cam_seg).astype(np.uint8)  # (H, W)
        background_seg = np.zeros(visual_id_seg.shape, dtype=np.bool)
        bg_visids = [
            0,
            self._bg_front.get_id(),
            self._bg_left.get_id(),
            self._bg_right.get_id(),
        ]
        for bg_visid in bg_visids:
            background_seg[visual_id_seg == bg_visid] = True
        rgbd_pc["bg_seg"] = convert_np_bool_to_float(background_seg).astype(
            np.uint8
        )  # (H, W)
        return rgbd_pc

    def get_obs(self):
        if self.domain_random:
            self._domain_random()
        if self._obs_mode == "state_dict":
            state_dict = self._get_obs_state_dict()
            return state_dict
        elif self._obs_mode == "state_egorgbd":
            state_rgbd = self.get_obs_state_egorgbd()
            return state_rgbd
        elif self._obs_mode == "egorgbd_pc":
            rgbd_pc = self.get_obs_egorgbd_pc()
            return rgbd_pc
        else:
            raise NotImplementedError(self._obs_mode)


@register_gym_env("CabinetDrawer-v0", max_episode_steps=200)
class CabinetDrawer(CabinetDoor):
    DEFAULT_MODEL_PATH = (
        ASSET_DIR / "partnet_mobility_configs/fixed_cabinet_drawers_new.yml"
    )

    def __init__(self, *args, **kwargs):
        super(CabinetDrawer, self).__init__(*args, **kwargs)


@register_gym_env("Faucet-v0", max_episode_steps=200)
class Faucet(FloatingRobotiqEnv):
    SUPPORTED_OBS_MODES = ("state", "state_dict", "state_egorgbd", "egorgbd_pc")
    SUPPORTED_REWARD_MODES = ("dense", "sparse")
    SUPPORTED_CONTROL_MODES = ("pd_joint_delta_pos", "pd_ee_delta_pose")
    DEFAULT_MODEL_PATH = ASSET_DIR / "partnet_mobility_configs/fixed_faucets_new.yml"

    def __init__(
        self,
        articulation_ids: List[int] = (),
        articulation_config_path="",
        obs_mode=None,
        reward_mode=None,
        irsensor_mode=False,
        other_handle_visible=False,
        domain_random=False,
        sim_freq=500,
        control_freq=20,
        control_mode=None,
    ):
        if articulation_config_path is None:
            articulation_config_path = self.DEFAULT_MODEL_PATH
        with open(articulation_config_path, "r") as f:
            self._articulation_info = yaml.safe_load(f)
        if isinstance(articulation_ids, int):
            articulation_ids = [articulation_ids]
        if len(articulation_ids) == 0:
            articulation_ids = sorted(self._articulation_info.keys())
        self.articulation_ids = articulation_ids
        self.articulation_id = None
        self.articulation_scale = None
        self.domain_random = domain_random

        super().__init__(
            obs_mode,
            reward_mode,
            control_mode,
            irsensor_mode,
            sim_freq,
            control_freq,
        )

    def _domain_random(self):
        self._scene.remove_articulation(self._articulation)
        self._load_articulations()
        self._initialize_articulations()
        if self._episode_rng.rand() > 0.2:
            rand_img_file = self._episode_rng.choice(self.texture_filelist, 1)[0]
            set_articulation_render_material(
                self._articulation,
                metallic=self._episode_rng.uniform(0.4, 1),  # 1, 0
                roughness=self._episode_rng.uniform(0.4, 0.9),  # 0.4, 0.9
                specular=self._episode_rng.uniform(0.0, 0.8),  # /, 0.0
                diffuse_texture_filename=str(
                    ASSET_DIR / f"textures/random_textures/{rand_img_file}"
                ),
            )
        super()._domain_random()

    def _initialize_agent(self):
        qpos = self._episode_rng.uniform(0, 0.068, [2])
        self._agent._robot.set_qpos(qpos)
        self._agent._robot.set_root_pose(
            look_at(eye=[-0.2, 0, 0.8], target=[-0.1, 0, 0], up=[0, 0, 1])
        )

    def reset(
        self,
        seed=None,
        reconfigure=False,
        articulation_id=None,
        articulation_scale=None,
    ):
        self.set_episode_rng(seed)
        _reconfigure = self._set_model(articulation_id, articulation_scale)
        reconfigure = _reconfigure or reconfigure
        ret = super().reset(seed=self._episode_seed, reconfigure=reconfigure)
        return ret

    def _set_model(self, articulation_id, articulation_scale):
        """Set the articulation id and scale. If not provided, choose one randomly."""
        reconfigure = False

        # Model ID
        if articulation_id is None:
            articulation_id = random_choice(self.articulation_ids, self._episode_rng)
        if articulation_id != self.articulation_id:
            reconfigure = True
        self.articulation_id = articulation_id
        self._articulation_config = self._articulation_info[self.articulation_id]

        # Scale
        if articulation_scale is None:
            articulation_scale = self._articulation_config.get("scale")
        if articulation_scale is None:
            bbox_max, bbox_min = (
                self._articulation_config["bbox_max"],
                self._articulation_config["bbox_min"],
            )
            bbox_size = np.float32(bbox_max) - np.float32(bbox_min)
            articulation_scale = 0.3 / max(bbox_size)  # hardcode
        if articulation_scale != self.articulation_scale:
            reconfigure = True
        self.articulation_scale = articulation_scale

        if "offset" in self._articulation_config:
            self.articulation_offset = np.float32(self._articulation_config["offset"])
        else:
            bbox_min = self._articulation_config["bbox_min"]
            self.articulation_offset = -np.float32(bbox_min) * articulation_scale
        # Add a small clearance
        self.articulation_offset[2] += 0.01

        return reconfigure

    def _load_articulations(self):
        self._articulation = self._load_faucet()
        # Cache qpos to restore
        self._articulation_init_qpos = self._articulation.get_qpos()

        # Set friction and damping for all joints
        for joint in self._articulation.get_active_joints():
            joint.set_friction(1.0)
            joint.set_drive_property(0.0, 10.0)

        self._set_switch_links()

    def _load_faucet(self):
        loader = self._scene.create_urdf_loader()
        loader.scale = self.articulation_scale
        loader.fix_root_link = True

        model_dir = ASSET_DIR / f"partnet_mobility_dataset/{self.articulation_id}"
        urdf_path = model_dir / "mobility_cvx.urdf"
        loader.load_multiple_collisions_from_file = True
        assert (
            urdf_path.exists()
        ), f"{urdf_path} is not found. Please download Partnet-Mobility (ManiSkill2022) first."

        density = self._articulation_config.get("density", 8e3)
        articulation = loader.load(str(urdf_path), config={"density": density})
        articulation.set_name("faucet")

        cam = self._scene.add_camera("dummy_camera", 1, 1, 1, 0.01, 10)
        self._scene.remove_camera(cam)
        set_articulation_render_material(
            articulation, color=hex2rgba("#AAAAAA"), metallic=1, roughness=0.4
        )

        return articulation

    def _set_switch_links(self):
        all_links = self._articulation.get_links()
        all_joints = self._articulation.get_joints()

        switch_link_names = []
        fix_link_visid = []
        for semantic in self._articulation_config["semantics"]:
            if semantic[2] == "switch":
                switch_link_names.append(semantic[0])
            else:
                fix_link = get_entity_by_name(all_links, semantic[0])
                for visual_body in fix_link.get_visual_bodies():
                    fix_link_visid.append(visual_body.get_visual_id())
        self.fix_link_visid = fix_link_visid

        if len(switch_link_names) == 0:
            raise RuntimeError(self.articulation_id)
        self.switch_link_names = switch_link_names

        self.switch_links = []
        self.switch_links_visid = []
        self.switch_links_mesh = []
        self.switch_joints = []
        for name in self.switch_link_names:
            link = get_entity_by_name(all_links, name)
            self.switch_links.append(link)
            for visual_body in link.get_visual_bodies():
                self.switch_links_visid.append(visual_body.get_visual_id())

            # cache mesh
            link_mesh = get_actor_mesh(link, False)
            self.switch_links_mesh.append(link_mesh)

            # hardcode
            joint = all_joints[link.get_index()]
            joint.set_friction(0.1)
            joint.set_drive_property(0.0, 2.0)
            self.switch_joints.append(joint)

    def _initialize_articulations(self):
        p = np.zeros(3)
        p[2] = self.articulation_offset[2]
        ori = self._episode_rng.uniform(-np.pi / 12, np.pi / 12)
        q = euler2quat(0, 0, ori)
        self._articulation.set_pose(sapien.Pose(p, q))

        # init_open_extent = 1
        init_open_extent = self._episode_rng.uniform(0, 1)
        # lmin, lmax = 0, 3.14
        qpos = np.zeros(self._articulation.dof)
        for i, joint in enumerate(self.switch_joints):
            [[lmin, lmax]] = joint.get_limits()
            lmin = 0 if np.isinf(lmin) else lmin
            lmax = 3.14 if np.isinf(lmax) else lmax
            # print(f"{self.articulation_id} joint limits: {lmin, lmax}")
            qpos[i] = lmin + (lmax - lmin) * init_open_extent
        self._articulation.set_qpos(qpos)
        self._set_target_link()

    def _set_target_link(self):
        n_switch_links = len(self.switch_link_names)
        idx = random_choice(np.arange(n_switch_links), self._episode_rng)

        self._target_link_name = self.switch_link_names[idx]
        self._target_link: sapien.Link = self.switch_links[idx]
        self._target_joint: sapien.Joint = self.switch_joints[idx]
        self._target_joint_idx = self._articulation.get_active_joints().index(
            self._target_joint
        )

        # x-axis is the revolute joint direction
        assert self._target_joint.type == "revolute", self._target_joint.type
        joint_pose = self._target_joint.get_global_pose().to_transformation_matrix()
        self.target_joint_axis = joint_pose[:3, 0]

        self.target_link_mesh = self.switch_links_mesh[idx]
        with np_random(self._episode_seed):
            self._sw_pcd = self.target_link_mesh.sample(1000)

        # get tip, med, bottom keypoints
        if self.articulation_id in [5024, 5034]:
            norm_sw_pcd = apply_pose_to_points(
                self._sw_pcd, self._target_joint.get_pose_in_parent().inv()
            )
            x_max, x_min = norm_sw_pcd[:, 0].max(), norm_sw_pcd[:, 0].min()
            center_pts, bottom_pts = (
                norm_sw_pcd[norm_sw_pcd[:, 0] - x_max > -0.01],
                norm_sw_pcd[norm_sw_pcd[:, 0] - x_min < 0.01],
            )
            center = np.mean(center_pts, axis=0)
            center[0] = x_max
            bottom = np.mean(bottom_pts, axis=0)
            bottom[0] = x_min
            tip = norm_sw_pcd[norm_sw_pcd[:, 1].argmin()]
            self._sw_keypoints = apply_pose_to_points(
                np.stack([tip, center, bottom]), self._target_joint.get_pose_in_parent()
            )
        else:
            xyzmax_val, xyzmin_val = self._sw_pcd.max(0), self._sw_pcd.min(0)
            xyzmax_id, xyzmin_id = np.argmax(self._sw_pcd, 0), np.argmin(
                self._sw_pcd, 0
            )
            bbox_size = xyzmax_val - xyzmin_val
            long_side = np.argmax(bbox_size)
            tip = self._sw_pcd[xyzmax_id[long_side]]
            # bottom = self._sw_pcd[xyzmin_id[long_side]]
            # center = self._sw_pcd.mean(0)
            # center[1] = xyzmax_val[1]
            joint_pts_ids = np.where(
                np.linalg.norm(self._sw_pcd[:, [0, 2]], axis=1) < 0.02
            )  # dist to joint axis < 0.01
            joint_pts = self._sw_pcd[joint_pts_ids][:, 1]
            joint_pts_max = joint_pts.max()
            center = np.array([0, joint_pts_max, 0])
            bottom = np.array([0, self._sw_pcd[:, 1].min(), 0])
            self._sw_keypoints = np.stack([tip, center, bottom])

    def get_obs_state_egorgbd(self):
        ### get state
        state_rgbd = self._get_obs_state_dict()

        ### get visual info
        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        hand_sensor.compute_depth()

        rgb = (hand_sensor.get_rgb() * 255).astype(np.uint8)  # [H, W, 3]
        ir_l, ir_r = hand_sensor.get_ir()  # [H, W]
        depth = hand_sensor.get_depth()  # [H, W]
        clean_depth = -get_texture(sensor_rgb, "Position")[..., 2]  # [H, W]
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        imh, imw = visual_id_seg.shape

        state_rgbd["rgb"] = rgb.transpose((2, 0, 1))
        if self.irsensor_mode:
            state_rgbd["ir_l"] = ir_l.astype(np.float32)
            state_rgbd["ir_r"] = ir_r.astype(np.float32)
            state_rgbd["depth"] = depth.astype(np.float32)
            state_rgbd["clean_depth"] = clean_depth.astype(np.float32)
        else:
            state_rgbd["clean_depth"] = depth.astype(np.float32)

        cam_seg = np.ones(visual_id_seg.shape, dtype=np.float32) * 5
        seg_visids = [self.switch_links_visid, self.fix_link_visid]
        seg_id = [3, 4]
        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros(visual_id_seg.shape, dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
        state_rgbd["seg"] = convert_np_bool_to_float(cam_seg).astype(np.uint8)  # (H, W)
        background_seg = np.zeros(visual_id_seg.shape, dtype=np.bool)
        bg_visids = [
            0,
            self._bg_front.get_id(),
            self._bg_left.get_id(),
            self._bg_right.get_id(),
        ]
        for bg_visid in bg_visids:
            background_seg[visual_id_seg == bg_visid] = True
        state_rgbd["bg_seg"] = convert_np_bool_to_float(background_seg).astype(
            np.uint8
        )  # (H, W)
        # got pcd target_link => world => ee
        # sw_pcd_world = apply_pose_to_points(self._sw_pcd, self._target_link.get_pose())  ## world frame
        kpts_world = apply_pose_to_points(
            self._sw_keypoints, self._target_link.get_pose()
        )  ## world frame
        # sw_pcd_ee = apply_pose_to_points(sw_pcd_world, self.agent.grasp_site.get_pose().inv())  ## world frame
        # sw_kpts_ee = apply_pose_to_points(sw_kpts_world, self.agent.grasp_site.get_pose().inv())  ## ee frame
        kpts_cam = transform_points(sensor_rgb.get_extrinsic_matrix(), kpts_world)
        uvz = xyz2uv(kpts_cam, sensor_rgb.get_camera_matrix())
        u_visable = (uvz[:, 0] >= 0) & (uvz[:, 0] <= imw - 1)
        v_visable = (uvz[:, 1] >= 0) & (uvz[:, 1] <= imh - 1)
        uvz_visable = np.stack(
            [u_visable, v_visable, u_visable & v_visable]
        ).T  # (3, 3)
        for ind in range(uvz_visable.shape[0]):
            if uvz_visable[ind, 2]:  # not out of img range
                u, v = round(uvz[ind, 0]), round(uvz[ind, 1])
                if cam_seg[v, u] not in seg_id:  # handle keypoint is occluded
                    uvz_visable[ind, 2] = False
        state_rgbd["kpts"] = kpts_cam.astype(np.float32)
        state_rgbd["uvz"] = uvz.astype(np.float32)
        state_rgbd["uvz_visable"] = uvz_visable
        return state_rgbd

    def get_obs_egorgbd_pc(self):
        rgbd_pc = OrderedDict()

        ### get visual info
        self.update_render()
        hand_sensor = self.agent._sensors["hand"]
        sensor_rgb = hand_sensor._cam_rgb
        hand_sensor.take_picture()
        hand_sensor.compute_depth()

        rgb = (hand_sensor.get_rgb() * 255).astype(np.uint8)  # [H, W, 3]
        ir_l, ir_r = hand_sensor.get_ir()  # [H, W]
        depth = hand_sensor.get_depth()  # [H, W]
        clean_depth = -get_texture(sensor_rgb, "Position")[..., 2]  # [H, W]
        visual_id_seg = get_camera_seg(sensor_rgb)[..., 0]  # [H, W]
        imh, imw = visual_id_seg.shape

        # process point cloud
        cam_xyz = hand_sensor.get_pointcloud(False)  # [H*W, 3]
        ## cam to world
        T_cam2world = sensor_rgb.get_model_matrix()
        world_xyz = transform_points(
            T_cam2world, cam_xyz * np.array([1, -1, -1])
        )  # [H*W, 3]
        T_world2ee = (
            self.agent.grasp_site.get_pose().inv().to_transformation_matrix()
        )  ## world => ee
        rgbd_pc["world_xyz"] = world_xyz

        rgbd_pc["rgb"] = rgb.transpose((2, 0, 1))
        if self.irsensor_mode:
            rgbd_pc["ir_l"] = ir_l.astype(np.float32)
            rgbd_pc["ir_r"] = ir_r.astype(np.float32)
            rgbd_pc["depth"] = depth.astype(np.float32)
            rgbd_pc["clean_depth"] = clean_depth.astype(np.float32)
        else:
            rgbd_pc["clean_depth"] = depth.astype(np.float32)

        background_seg = np.zeros(visual_id_seg.shape, dtype=np.bool)
        bg_visids = [
            0,
            self._bg_front.get_id(),
            self._bg_left.get_id(),
            self._bg_right.get_id(),
        ]
        for bg_visid in bg_visids:
            background_seg[visual_id_seg == bg_visid] = True
        rgbd_pc["bg_seg"] = convert_np_bool_to_float(background_seg).astype(
            np.uint8
        )  # (H, W)

        cam_seg = np.ones(visual_id_seg.shape, dtype=np.float32) * 5
        seg_visids = [self.switch_links_visid, self.fix_link_visid]
        seg_id = [3, 4]
        for id, seg_visid in enumerate(seg_visids):
            mask = np.zeros(visual_id_seg.shape, dtype=np.bool)
            for visual_id in seg_visid:
                mask = mask | (visual_id_seg == visual_id)
            cam_seg[mask] = seg_id[id]
        rgbd_pc["seg"] = convert_np_bool_to_float(cam_seg).astype(np.uint8)  # (H, W)
        return rgbd_pc

    def get_obs(self):
        if self.domain_random:
            self._domain_random()
        if self._obs_mode == "state_dict":
            return self._get_obs_state_dict()
        elif self._obs_mode == "state_egorgbd":
            return self.get_obs_state_egorgbd()
        elif self._obs_mode == "egorgbd_pc":
            return self.get_obs_egorgbd_pc()
        else:
            raise NotImplementedError(self._obs_mode)

arti_mani/envs/floating_robotiqenv.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import time
import warnings
from collections import OrderedDict, defaultdict

import cv2
import gym
import numpy as np
import open3d as o3d
import sapien.core as sapien
import torch
import torch.nn.functional as F
import yaml
from arti_mani import RLMODEL_DIR, VISUALMODEL_DIR
from arti_mani.algorithms.rl_iam.Sampling import (
    FCUWSampler,
    FPSSampler,
    RandomSampler,
    ScoreSampler,
    UniDownsampler,
)
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import (
    CustomUnet,
    CustomUnetNew,
    SplitUnet,
)
from arti_mani.envs.arti_mani import ArtiMani
from arti_mani.test.real_exps.utils import mat2pose
from arti_mani.utils.geometry import transform_points
from arti_mani.utils.sapien_utils import get_entity_by_name
from arti_mani.utils.wrappers import NormalizeActionWrapper
from observer import Observer
from stable_baselines3 import SAC
from xmate3Robotiq import Robotiq2FGripper_Listener, ROS_ImpOpendoor

warnings.filterwarnings("ignore")


def setup_seed(seed=1029):
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    rng = np.random.RandomState(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    return rng


def scale_para(para, output_low, output_high):
    """
    para: norm to [-1, 1]
    output_low, output_high: reproject to (output_low, output_high)
    """
    bias = 0.5 * (output_high + output_low)
    weight = 0.5 * (output_high - output_low)
    output = weight * para + bias

    return output


def load_segmodel(path):
    ### load visual model
    smp_model_path = VISUALMODEL_DIR / f"smp_model/{path}/best.pth"
    config_path = VISUALMODEL_DIR / f"smp_model/{path}/config.yaml"
    with open(config_path, "r") as f:
        cfg = yaml.safe_load(f)
    smp_cfg = cfg["smp_config"]
    data_mode = smp_cfg["mode"]
    if data_mode == "RGBD":
        in_ch = 4
    elif data_mode == "RGB":
        in_ch = 3
    elif data_mode == "D":
        in_ch = 1
    else:
        raise NotImplementedError

    if "has_dropout" in smp_cfg.keys():
        unet_model = CustomUnet(
            has_dropout=True,
            encoder_name=smp_cfg["encoder"],
            encoder_depth=smp_cfg["encoder_depth"],
            decoder_channels=smp_cfg["decoder_channels"],
            encoder_weights=smp_cfg["encoder_weights"],
            in_channels=in_ch,
            classes=cfg["num_classes"],
            activation=None,
            # activation=smp_cfg["activation"],
        )
    elif smp_cfg["encoder"] == "splitnet":
        unet_model = SplitUnet(
            dropout_p=smp_cfg["dropout_p"],
            encoder_name=smp_cfg["encoder"],
            encoder_depth=smp_cfg["encoder_depth"],
            decoder_channels=smp_cfg["decoder_channels"],
            encoder_weights=smp_cfg["encoder_weights"],
            in_channels=in_ch,
            classes=cfg["num_classes"],
        )
    else:
        unet_model = CustomUnetNew(
            dropout_p=smp_cfg["dropout_p"],
            encoder_name=smp_cfg["encoder"],
            encoder_depth=smp_cfg["encoder_depth"],
            decoder_channels=smp_cfg["decoder_channels"],
            encoder_weights=smp_cfg["encoder_weights"],
            in_channels=in_ch,
            classes=cfg["num_classes"],
        )
    unet_model.load_state_dict(torch.load(smp_model_path))
    unet_model.to(device)
    return unet_model, data_mode


def get_uncertainty(seg_model, rgbd_tensor):
    uncertain_pred = []
    for ind in range(4):
        pred_seg_mtcsamp = seg_model(rgbd_tensor, True)  # (N, 6, H, W)
        uncertain_pred.append(F.softmax(pred_seg_mtcsamp))
    uncertain_mean = torch.mean(torch.stack(uncertain_pred), dim=0)  # (N, 6, H, W)
    print("mean uncertainty ==", (uncertain_pred[0] - uncertain_mean).mean())
    uncertain_map = -1.0 * torch.sum(
        uncertain_mean * torch.log(uncertain_mean + 1e-6), dim=1
    )  # (N, H, W)
    uncertain_map = uncertain_map.cpu().numpy().squeeze()  # (N, H, W)
    norm_uncertain_map = (uncertain_map - uncertain_map.min()) / (
        uncertain_map.max() - uncertain_map.min()
    )
    norm_uncertain_map -= norm_uncertain_map[norm_uncertain_map > 1e-3].mean()
    norm_uncertain_map = np.clip(norm_uncertain_map, 0, 1)
    return norm_uncertain_map


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "door"  # door, drawer, faucet
    rs_mode = "rgbd"  # "rgbd", "depth", "rgb"

    eval_seed = np.random.RandomState().randint(2**32)
    rng = setup_seed(eval_seed)
    print("experiment eval random seed: ", eval_seed)

    save_mode = False
    record_id = 0
    plot_mode = False
    points_vis_mode = False
    add_eepadpts = False
    gripper_wait = 0.3
    MAX_STEPS = 50
    subplot_num = [2, 3]
    camH, camW = 144, 256
    sample_mode = "frameweight_sample"  # full_downsample, score_sample, fps_sample, random_sample, frameweight_sample
    frame_num = 3
    sample_num = 32
    num_classes = 6
    test_name = (
        f"{mode}0_demo{record_id}_{sample_mode}_{frame_num}_{sample_num}_smp0327"
    )
    fig_modes = [
        "rgb",
        "depth",
        "segpred",
        "samplepts",
        "uncertainty_norm",
        "weight_norm",
    ]
    cmaps = [
        (1, 0, 0),  # "red"
        (0, 0, 1),  # "blue"
        (1, 1, 0),  # "yellow"
        (0, 1, 0),  # "green"
        (0.627, 0.125, 0.941),  # "purple"
        (0.753, 0.753, 0.753),  # "grey"
    ]
    if mode == "door":
        arti_ids = [0]  # 0
    elif mode == "drawer":
        arti_ids = [1]  # 1
    elif mode == "faucet":
        arti_ids = [5052]
    elif mode == "arti":
        arti_ids = [0, 1, 5052]
    else:
        raise NotImplementedError
    obs_mode = "state_egostereo_segpoints"
    control_mode = "pd_joint_delta_pos"
    smp_exps = ["share/20230327_231728_full_splitnet_new_data_aug_ce0.0_dropout0.2"]
    smp_exp_name = smp_exps[0]
    if mode == "door":
        rl_exp_name = "sac4_tr4eval1door_state_egostereo_segpoints_random_samplef1s32_100steps_padfriction1_p3000d100_tr16up16_noeepadpts_newextric_ptsaddnoise0.01_smp0219"
    elif mode == "drawer":
        rl_exp_name = "sac2_tr4eval1drawer_state_egostereo_segpoints_random_samplef1s32_100steps_padfriction1_p3000d100_tr16up16_noeepadpts_newextric_smp0219_adjustreward_2close_initrot"
    elif mode == "faucet":
        rl_exp_name = "sac4_tr12eval3arti_state_egostereo_segpoints_random_samplef1s32_100steps_smp0219_largerange"
    else:
        raise NotImplementedError
    model_path = RLMODEL_DIR / f"{rl_exp_name}/best_model"

    real_results_path = RLMODEL_DIR / f"{rl_exp_name}/real_result/{test_name}/"
    if save_mode:
        if not os.path.exists(real_results_path):
            os.makedirs(real_results_path)

    if sample_mode == "full_downsample":
        sampler = UniDownsampler(sample_num=sample_num)
    elif sample_mode == "score_sample":
        sampler = ScoreSampler(sample_num=sample_num)
    elif sample_mode == "fps_sample":
        sampler = FPSSampler(sample_num=sample_num)
    elif sample_mode == "random_sample":
        sampler = RandomSampler(sample_num=sample_num)
    elif sample_mode == "frameweight_sample":
        sampler = FCUWSampler(frame_num=frame_num, sample_num=sample_num)
    else:
        raise NotImplementedError(sample_mode)

    print("log name: ", rl_exp_name)

    print("+++++ Build Real2Sim Env +++++++")
    env: ArtiMani = gym.make(
        "ArtiMani-v0",
        articulation_ids=arti_ids,
        segmodel_path=VISUALMODEL_DIR / f"smp_model/{smp_exp_name}",
        sample_mode="random_sample",
        frame_num=frame_num,
        sample_num=sample_num,
        other_handle_visible=False,
        obs_mode=obs_mode,
        control_mode=control_mode,
        reward_mode="dense",  # "sparse", "dense"
        add_eepadpts=add_eepadpts,
        device=device,
    )
    env = NormalizeActionWrapper(env)
    env.seed(eval_seed)

    obs = env.reset()
    robot_ee: sapien.Link = get_entity_by_name(
        env.agent._robot.get_links(), "xmate3_link7"
    )
    ### Set Cabinet Pose
    env.unwrapped._articulation.set_root_pose(sapien.Pose([2, 0, 0], [1, 0, 0, 0]))

    ### Set Robot Pose
    ## real robot init
    if mode == "door" or mode == "drawer":
        qpos = np.array([1.4, -1.053, -2.394, 1.662, 1.217, 1.05, -0.8, 0.0, 0.0])
    elif mode == "faucet":
        qpos = np.array([-0.5, -0.143, 0, np.pi / 3, 0, 1.57, 1.57, 0.068, 0.068])
    else:
        raise NotImplementedError(mode)
    env.unwrapped._agent._robot.set_qpos(qpos)

    # viewer = env.unwrapped.render()
    # print("Press [b] to start")
    # while True:
    #     if viewer.window.key_down("b"):
    #         break
    #     env.render()

    print(f"Action space {env.action_space}")
    print(f"Control mode {env.control_mode}")
    action_range_low = env.unwrapped.action_space.low
    action_range_high = env.unwrapped.action_space.high
    print("Input Action space range (low):", action_range_low)
    print("Input Action space range (high):", action_range_high)

    print("+++++ Setup robotiq controller +++++++")
    gripper_listener = Robotiq2FGripper_Listener()
    print("+++++ Setup robot arm controller +++++++")
    imp_opendoor = ROS_ImpOpendoor(
        gripper_listener=gripper_listener, control_mode=control_mode
    )
    if control_mode == "pd_joint_delta_pos":
        joint_kp = [200.0, 200.0, 200.0, 200.0, 50.0, 50.0, 50.0]
        joint_kd = [20.0, 20.0, 20.0, 20.0, 10.0, 10.0, 10.0]
        # joint_kp = [300.0, 300.0, 300.0, 300.0, 100.0, 100.0, 100.0]
        # joint_kd = [10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0]
    else:
        raise NotImplementedError
    print("current joint kp: ", joint_kp)
    print("current joint kd: ", joint_kd)

    ### load RL model
    RL_model = SAC.load(model_path, env=env, print_system_info=True)

    unet_model, data_mode = load_segmodel(smp_exp_name)
    unet_model.eval()

    step = 0
    sim2real_step_num = 5
    prev_target_qpos = None
    prev_target_gripper_pos = None
    figs_dict = defaultdict(list)

    with torch.no_grad():
        rs_observer = Observer(
            base_camera_SN="918512071887",
            mode=rs_mode,
        )  # d435

        while step < MAX_STEPS:
            ### get agent qpos in real
            imp_opendoor.get_realstate()
            robot_arm_state = imp_opendoor.real_state.q_m
            gripper_state = np.repeat(
                gripper_listener.gPO / 255 * 0.068, 2
            )  # 0-255 => 0-0.068
            real_robot_qpos = np.concatenate((robot_arm_state, gripper_state))
            print("=== real robot current qpos: ", real_robot_qpos)
            ### reset agent qpos and get grasp_site pos in sim
            env.agent._robot.set_qpos(real_robot_qpos)
            sim_ee_pose_base = (
                env.agent._robot.get_root_pose().inv() * env.agent.grasp_site.get_pose()
            )  # in robot root coord
            sim_robotee_pose_base = (
                env.agent._robot.get_root_pose().inv() * robot_ee.get_pose()
            )  # in robot root coord
            real_ee_pose_mat = imp_opendoor.real_state.toolTobase_pos_m
            real_ee_pose_base = mat2pose(np.reshape(real_ee_pose_mat, (4, 4)))
            ee_pos_base = sim_ee_pose_base.p
            print("sim robot current qpos: ", env.agent._robot.get_qpos())
            print("sim ee pose (robot root frame): ", sim_ee_pose_base)
            print("sim robot link7 pose (robot root frame): ", sim_robotee_pose_base)
            print("real ee pose (robot root frame): ", real_ee_pose_base)

            # # get ee pad pts
            ee_cords = env.agent.sample_ee_coords(num_sample=20).reshape(
                -1, 3
            )  # [40, 3]
            T_world2ee = (
                env.agent.grasp_site.get_pose().inv().to_transformation_matrix()
            )
            eepad_pts = transform_points(T_world2ee, ee_cords)  # [40, 3]

            eepat_pts_label = np.zeros(
                (eepad_pts.shape[0], num_classes + 1)
            )  # (40, C+1)
            eepat_pts_label[:, -1] = 1.0
            eepad_ptsC = np.concatenate(
                (eepad_pts, eepat_pts_label), axis=1
            )  # (40, 10)

            ### concate to total states
            frame, _, _, cam_xyz, rs_depth_intrinsics = rs_observer.get_observation()
            cur_rgb, cur_depth = frame[:, :, :3] / 255.0, frame[:, :, 3]
            H, W = cur_depth.shape

            # get pred masks
            rgb_tensor = (
                torch.from_numpy(cur_rgb[None].transpose(0, 3, 1, 2)).float().to(device)
            )  # 1,3,H,W
            depth_tensor = (
                torch.from_numpy(cur_depth[None, None]).float().to(device)
            )  # 1,1,H,W
            if data_mode == "RGBD":
                img = torch.cat([rgb_tensor, depth_tensor], dim=1)
            elif data_mode == "RGB":
                img = rgb_tensor
            elif data_mode == "D":
                img = depth_tensor
            else:
                raise NotImplementedError
            seg_map = unet_model.predict(img, T=8, dropout=True)[0]  # (C, H, W)
            seg_mc_label = torch.argmax(seg_map, dim=0).reshape(-1)  # (H*W)
            if add_eepadpts:
                seg_onehot_label = (
                    F.one_hot(seg_mc_label, num_classes=num_classes + 1)
                    .cpu()
                    .detach()
                    .numpy()
                )  # (H*W, C+1)
            else:
                seg_onehot_label = (
                    F.one_hot(seg_mc_label, num_classes=num_classes)
                    .cpu()
                    .detach()
                    .numpy()
                )  # (H*W, C)
            seg_mc_map = seg_mc_label.cpu().detach().numpy()  # (H*W)

            # uncertainty map
            uncertainty_norm = get_uncertainty(unet_model, img)

            ## cam to world
            hand_sensor = env.agent._sensors["hand"]
            sensor_rgb = hand_sensor._cam_rgb
            T_cam2world_sim = sensor_rgb.get_model_matrix()
            print("=== cam to world (wrong): ", T_cam2world_sim)
            T_cam2ee = env.cam_para["hand_camera_extrinsic_base_frame"]
            print("=== cam to ee: ", T_cam2ee)
            T_ee2world = (
                env.agent.grasp_site.get_pose().to_transformation_matrix()
            )  ## world => ee
            print("=== ee to world: ", T_ee2world)
            T_cam2world = T_ee2world @ T_cam2ee
            print("=== cam to world (correct): ", T_cam2world)
            world_xyz = transform_points(
                T_cam2world, cam_xyz * np.array([1, -1, -1])
            )  # [H*W, 3]

            weight_map = np.zeros((1, camH, camW))
            if sample_mode == "full_downsample":
                pts_index = sampler.sampling(world_xyz, rng)
                ds_world_xyz = world_xyz[pts_index]
                ds_onehot_label = seg_onehot_label[pts_index]
                # ## world to ee
                T_world2ee = (
                    env.unwrapped.agent.grasp_site.get_pose()
                    .inv()
                    .to_transformation_matrix()
                )  ## world => ee
                ds_ee_xyz = transform_points(T_world2ee, ds_world_xyz)
                segptsC = np.concatenate(
                    (ds_ee_xyz, ds_onehot_label), axis=1
                )  # (SN, 3+C+1)
                if points_vis_mode:
                    # pcd0 = o3d.geometry.PointCloud()
                    # pcd0.points = o3d.utility.Vector3dVector(transform_points(T_world2ee, world_xyz))
                    # pcd0.paint_uniform_color([0,0,0])
                    pcd1 = o3d.geometry.PointCloud()
                    pcd1.points = o3d.utility.Vector3dVector(eepad_pts)
                    pcd1.paint_uniform_color([0, 1, 0])
                    ds_mc_label = seg_mc_map[pts_index]  # (1000, )
                    pcd_parts = [
                        o3d.geometry.PointCloud(),
                        o3d.geometry.PointCloud(),
                        o3d.geometry.PointCloud(),
                        o3d.geometry.PointCloud(),
                        o3d.geometry.PointCloud(),
                        o3d.geometry.PointCloud(),
                    ]
                    for ind in range(6):
                        pcd_parts[ind].points = o3d.utility.Vector3dVector(
                            ds_ee_xyz[ds_mc_label == ind]
                        )
                        pcd_parts[ind].paint_uniform_color(cmaps[ind])
                    XYZ = o3d.geometry.TriangleMesh.create_coordinate_frame(
                        size=0.1, origin=[0.0, 0.0, 0.0]
                    )
                    objs = [
                        pcd1,
                        XYZ,
                        pcd_parts[0],
                        pcd_parts[1],
                        pcd_parts[2],
                        pcd_parts[3],
                        pcd_parts[4],
                        pcd_parts[5],
                    ]
                    o3d.visualization.draw_geometries(objs)
            else:
                # get seg points
                if sample_mode == "score_sample":
                    pts_index = sampler.sampling(seg_map)
                elif sample_mode == "random_sample":
                    pts_index = sampler.sampling(seg_mc_map, num_classes, rng)
                elif sample_mode == "fps_sample":
                    pts_index = sampler.sampling(
                        world_xyz, num_classes, seg_mc_map, rng
                    )
                elif sample_mode == "frameweight_sample":
                    pts_index, weight_map = sampler.sampling(
                        world_xyz, num_classes, seg_mc_map, uncertainty_norm, rng
                    )
                world_segsampled_xyz = world_xyz[pts_index]  # (C*sample_num, 3)
                ee_segsampled_xyz = transform_points(
                    T_world2ee, world_segsampled_xyz
                )  # (C*sample_num, 3)
                sampled_onehot_label = seg_onehot_label[
                    pts_index
                ]  # (C*sample_num, C+1)
                segptsC = np.concatenate(
                    (ee_segsampled_xyz, sampled_onehot_label), axis=1
                )  # (C*sample_num, 3+C+1)
                # debug point clouds
                if points_vis_mode:
                    ee_xyz = transform_points(T_world2ee, world_xyz)
                    # pcd0 = o3d.geometry.PointCloud()
                    # pcd0.points = o3d.utility.Vector3dVector(ee_xyz)
                    # pcd0.paint_uniform_color([0, 0, 0])
                    pcd1 = o3d.geometry.PointCloud()
                    pcd1.points = o3d.utility.Vector3dVector(eepad_pts)
                    pcd1.paint_uniform_color([0, 1, 0])
                    XYZ = o3d.geometry.TriangleMesh.create_coordinate_frame(
                        size=0.1, origin=[0.0, 0.0, 0.0]
                    )
                    # objs = [pcd0, pcd1, XYZ]
                    pcd_parts = [
                        o3d.geometry.PointCloud(),
                        o3d.geometry.PointCloud(),
                        o3d.geometry.PointCloud(),
                        o3d.geometry.PointCloud(),
                        o3d.geometry.PointCloud(),
                        o3d.geometry.PointCloud(),
                    ]
                    for ind in range(6):
                        pcd_parts[ind].points = o3d.utility.Vector3dVector(
                            segptsC[ind * 30 : (ind + 1) * 30, :3]
                        )
                        pcd_parts[ind].paint_uniform_color(cmaps[ind])
                    objs = [
                        pcd1,
                        XYZ,
                        pcd_parts[0],
                        pcd_parts[1],
                        pcd_parts[2],
                        pcd_parts[3],
                        pcd_parts[4],
                        pcd_parts[5],
                    ]
                    o3d.visualization.draw_geometries(objs)

            state_egopts = OrderedDict()
            state_egopts["qpos"] = torch.from_numpy(
                real_robot_qpos
            ).float()  # .to(device)  # (9)
            state_egopts["ee_pos_base"] = torch.from_numpy(
                ee_pos_base
            ).float()  # .to(device)  # (3)
            if add_eepadpts:
                state_egopts["segsampled_ptsC"] = np.concatenate(
                    (segptsC, eepad_ptsC), axis=0
                ).astype(
                    np.float32
                )  # (SN+40, 3+C+1)
            else:
                state_egopts["segsampled_ptsC"] = segptsC.astype(
                    np.float32
                )  # (C*sample_num+40, 3+C)

            rgb_plot = cur_rgb.copy()
            if sample_mode == "full_downsample":
                pts_inds = np.array(pts_index)
                u = pts_inds // camW
                v = pts_inds - u * camW
                for idn in range(u.shape[0]):
                    cls = np.argmax(seg_onehot_label, axis=1)[pts_inds[idn]]
                    cv2.circle(
                        rgb_plot,
                        (int(v[idn]), int(u[idn])),
                        radius=2,
                        color=cmaps[cls],
                        thickness=-1,
                    )
            else:
                for cls in range(num_classes):
                    cls_inds = np.array(
                        pts_index[cls * sample_num : (cls + 1) * sample_num]
                    )
                    u = cls_inds // camW
                    v = cls_inds - u * camW
                    for idn in range(u.shape[0]):
                        cv2.circle(
                            rgb_plot,
                            (int(v[idn]), int(u[idn])),
                            radius=2,
                            color=cmaps[cls],
                            thickness=-1,
                        )
            print(f"+++++++++ step {step} +++++++++++")

            action, _states = RL_model.predict(state_egopts, deterministic=True)
            if save_mode:
                real_robot_qpos_eepos_action = np.concatenate(
                    (real_robot_qpos, ee_pos_base, action), axis=0
                )  # 9 + 3 + 8
                np.save(
                    real_results_path / f"{step:02}_qpos_eepos_action.npy",
                    real_robot_qpos_eepos_action,
                )

            real_delta_qpos = scale_para(
                action[:7], action_range_low[:7], action_range_high[:7]
            )
            target_qpos = real_robot_qpos[:-2] + real_delta_qpos
            print("real qpos, action: ", real_robot_qpos[:-2], action)
            print("prev target robot qpos: ", prev_target_qpos)
            print("target robot qpos: ", np.asarray(target_qpos))
            for k in range(sim2real_step_num):
                imp_opendoor.exec_trajs(
                    target_qpos, stiffness=joint_kp, damping=joint_kd
                )

            ## control gripper
            if mode == "door" or mode == "drawer":
                target_gripper_pos = int(
                    (action[-1] + 1) / 2 * 255
                )  # Gripper Open->Close: -1->1 => 0->255
            elif mode == "faucet":
                target_gripper_pos = 255
            else:
                raise NotImplementedError(mode)
            print("target gripper qpos: ", target_gripper_pos)
            imp_opendoor.exec_gripper(target_gripper_pos)
            if mode != "faucet":
                time.sleep(gripper_wait)
            real_gripper_pos = gripper_listener.gPO
            print("real gripper qpos: ", real_gripper_pos)

            step += 1
    print("Total steps: ", step)

arti_mani/test/real_exps/SB3_opencabinet.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import time
import warnings
from collections import OrderedDict, defaultdict
from copy import deepcopy

import cv2
import gym
import matplotlib.pyplot as plt
import numpy as np
import open3d as o3d
import sapien.core as sapien
import segmentation_models_pytorch as smp
import torch
import torch.nn.functional as F
import yaml
from arti_mani import ASSET_DIR, RLMODEL_DIR, VISUALMODEL_DIR
from arti_mani.algorithms.rl_iam.Sampling import (
    FCUWSampler,
    FPSSampler,
    RandomSampler,
    ScoreSampler,
    UniDownsampler,
)
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import (
    CustomUnet,
    CustomUnetNew,
    SplitUnet,
)
from arti_mani.envs.arti_mani import ArtiMani
from arti_mani.test.real_exps.utils import mat2pose
from arti_mani.utils.contrib import apply_pose_to_points
from arti_mani.utils.cv_utils import visualize_depth, visualize_seg
from arti_mani.utils.geometry import transform_points
from arti_mani.utils.o3d_utils import pcd_uni_down_sample_with_crop
from arti_mani.utils.sapien_utils import get_entity_by_name
from arti_mani.utils.wrappers import NormalizeActionWrapper
from observer import Observer
from stable_baselines3 import SAC
from xmate3Robotiq import Robotiq2FGripper_Listener, ROS_ImpOpendoor

# from xmate3Robotiq_new import Robotiq2FGripper_Listener, ROS_ImpOpendoor

warnings.filterwarnings("ignore")


def setup_seed(seed=1029):
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    rng = np.random.RandomState(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    return rng


def load_segmodel(path):
    ### load visual model
    smp_model_path = VISUALMODEL_DIR / f"smp_model/{path}/best.pth"
    config_path = VISUALMODEL_DIR / f"smp_model/{path}/config.yaml"
    with open(config_path, "r") as f:
        cfg = yaml.safe_load(f)
    smp_cfg = cfg["smp_config"]
    data_mode = smp_cfg["mode"]
    if data_mode == "RGBD":
        in_ch = 4
    elif data_mode == "RGB":
        in_ch = 3
    elif data_mode == "D":
        in_ch = 1
    else:
        raise NotImplementedError

    if "has_dropout" in smp_cfg.keys():
        unet_model = CustomUnet(
            has_dropout=True,
            encoder_name=smp_cfg["encoder"],
            encoder_depth=smp_cfg["encoder_depth"],
            decoder_channels=smp_cfg["decoder_channels"],
            encoder_weights=smp_cfg["encoder_weights"],
            in_channels=in_ch,
            classes=cfg["num_classes"],
            activation=None,
            # activation=smp_cfg["activation"],
        )
    elif smp_cfg["encoder"] == "splitnet":
        unet_model = SplitUnet(
            dropout_p=smp_cfg["dropout_p"],
            encoder_name=smp_cfg["encoder"],
            encoder_depth=smp_cfg["encoder_depth"],
            decoder_channels=smp_cfg["decoder_channels"],
            encoder_weights=smp_cfg["encoder_weights"],
            in_channels=in_ch,
            classes=cfg["num_classes"],
        )
    else:
        unet_model = CustomUnetNew(
            dropout_p=smp_cfg["dropout_p"],
            encoder_name=smp_cfg["encoder"],
            encoder_depth=smp_cfg["encoder_depth"],
            decoder_channels=smp_cfg["decoder_channels"],
            encoder_weights=smp_cfg["encoder_weights"],
            in_channels=in_ch,
            classes=cfg["num_classes"],
        )
    unet_model.load_state_dict(torch.load(smp_model_path))
    unet_model.to(torch.device("cuda:0"))
    return unet_model, data_mode


def scale_para(para, output_low, output_high):
    """
    para: norm to [-1, 1]
    output_low, output_high: reproject to (output_low, output_high)
    """
    bias = 0.5 * (output_high + output_low)
    weight = 0.5 * (output_high - output_low)
    output = weight * para + bias

    return output


def save_subplots(real_results_path, step, subplot_num, rgb, depth, segmap, mode=None):
    plt.subplot(subplot_num[0], subplot_num[1], 1)
    plt.imshow(rgb)
    plt.title("rgb")
    plt.axis("off")
    plt.subplot(subplot_num[0], subplot_num[1], 2)
    plt.imshow(visualize_depth(depth)[..., ::-1])
    plt.title("depth")
    plt.axis("off")
    plt.subplot(subplot_num[0], subplot_num[1], 3)
    plt.imshow(visualize_seg(segmap))
    plt.title("seg")
    plt.axis("off")
    if mode == "plot":
        plt.show()
    elif mode == "save":
        plt.savefig(
            real_results_path / f"test_{step:02}.png",
            bbox_inches="tight",
            pad_inches=0.1,
        )
    else:
        raise NotImplementedError


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "faucet"  # door, drawer, faucet
    rs_mode = "rgbd"  # "rgbd", "depth", "rgb"

    eval_seed = np.random.RandomState().randint(2**32)
    rng = setup_seed(eval_seed)
    print("experiment eval random seed: ", eval_seed)

    save_plot_mode = "plot"  # "plot", "save"
    record_id = 0
    gripper_wait = 0.3
    MAX_STEPS = 50
    subplot_num = [1, 3]
    test_name = f"{mode}0_demo{record_id}_2dseg"
    fig_modes = [
        "rgb",
        "depth",
        "segpred",
        "samplepts",
        "uncertainty_norm",
        "weight_norm",
    ]
    cmaps = [
        (1, 0, 0),  # "red"
        (0, 0, 1),  # "blue"
        (1, 1, 0),  # "yellow"
        (0, 1, 0),  # "green"
        (0.627, 0.125, 0.941),  # "purple"
        (0.753, 0.753, 0.753),  # "grey"
    ]
    if mode == "door":
        arti_ids = [0]  # 0
    elif mode == "drawer":
        arti_ids = [1]  # 1
    elif mode == "faucet":
        arti_ids = [5052]
    elif mode == "arti":
        arti_ids = [0, 1, 5052]
    else:
        raise NotImplementedError
    obs_mode = "state_egostereo_dseg"
    control_mode = "pd_joint_delta_pos"

    if mode == "door":
        rl_exp_name = "sac4_tr4eval1door_state_egostereo_dexpoints_0"  # 0, 1, 2
    elif mode == "drawer":
        rl_exp_name = "sac4_tr4eval1drawer_state_egostereo_dexpoints_2"  # 0, 1, 2
    elif mode == "faucet":
        rl_exp_name = "baselines/sac4_tr4eval1faucet_state_egostereo_dseg_100steps_padfriction1_p3000d100_bs128tr16up4_0"  # 0, 1, 5
    else:
        raise NotImplementedError
    model_path = (
        RLMODEL_DIR / f"{rl_exp_name}/best_model"
    )  # best_model, rl_model_1750000_steps, 2000000, 2250000, 2500000

    real_results_path = RLMODEL_DIR / f"{rl_exp_name}/real_result/{test_name}/"
    if save_plot_mode == "save":
        if not os.path.exists(real_results_path):
            os.makedirs(real_results_path)

    print("log name: ", rl_exp_name)

    print("+++++ Build Real2Sim Env +++++++")
    env: ArtiMani = gym.make(
        "ArtiMani-v0",
        articulation_ids=arti_ids,
        segmodel_path=None,
        sample_mode=None,
        frame_num=0,
        sample_num=0,
        other_handle_visible=False,
        obs_mode=obs_mode,
        control_mode=control_mode,
        reward_mode="dense",  # "sparse", "dense"
        add_eepadpts=False,
        device=device,
    )
    env = NormalizeActionWrapper(env)
    env.seed(eval_seed)

    obs = env.reset()
    robot_ee: sapien.Link = get_entity_by_name(
        env.agent._robot.get_links(), "xmate3_link7"
    )
    ### Set Cabinet Pose
    env.unwrapped._articulation.set_root_pose(sapien.Pose([2, 0, 0], [1, 0, 0, 0]))

    ### Set Robot Pose
    ## real robot init
    if mode == "door" or mode == "drawer":
        qpos = np.array([1.4, -1.053, -2.394, 1.662, 1.217, 1.05, -0.8, 0.0, 0.0])
    elif mode == "faucet":
        qpos = np.array([-0.5, -0.143, 0, np.pi / 3, 0, 1.57, 1.57, 0.068, 0.068])
    else:
        raise NotImplementedError(mode)
    env.unwrapped._agent._robot.set_qpos(qpos)

    print(f"Action space {env.action_space}")
    print(f"Control mode {env.control_mode}")
    action_range_low = env.unwrapped.action_space.low
    action_range_high = env.unwrapped.action_space.high
    print("Input Action space range (low):", action_range_low)
    print("Input Action space range (high):", action_range_high)

    print("+++++ Setup robotiq controller +++++++")
    gripper_listener = Robotiq2FGripper_Listener()
    print("+++++ Setup robot arm controller +++++++")
    imp_opendoor = ROS_ImpOpendoor(
        gripper_listener=gripper_listener, control_mode=control_mode
    )
    if control_mode == "pd_joint_delta_pos":
        joint_kp = [200.0, 200.0, 200.0, 200.0, 50.0, 50.0, 50.0]
        joint_kd = [20.0, 20.0, 20.0, 20.0, 10.0, 10.0, 10.0]
        # joint_kp = [300.0, 300.0, 300.0, 300.0, 100.0, 100.0, 100.0]
        # joint_kd = [10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0]
    else:
        raise NotImplementedError
    print("current joint kp: ", joint_kp)
    print("current joint kd: ", joint_kd)

    ### load seg model
    smp_exp_name = "20230228_220446_384_DR_randombg_mixfaucet_aug_dropout0.5_stereo_bs16_focalloss_0.5step50lr0.001_RGBDunet-163264128_mobilenet_v2"
    unet_model, data_mode = load_segmodel(smp_exp_name)
    unet_model.eval()

    ### load RL model
    print("RL model path: ", model_path)
    RL_model = SAC.load(model_path, env=env, print_system_info=True)

    step = 0
    sim2real_step_num = 5
    prev_target_qpos = None
    prev_target_gripper_pos = None

    figs_dict = defaultdict(list)

    with torch.no_grad():
        rs_observer = Observer(
            base_camera_SN="918512071887",
            mode=rs_mode,
        )  # d435

        while step < MAX_STEPS:
            ### get agent qpos in real
            imp_opendoor.get_realstate()
            robot_arm_state = imp_opendoor.real_state.q_m
            gripper_state = np.repeat(
                gripper_listener.gPO / 255 * 0.068, 2
            )  # 0-255 => 0-0.068
            real_robot_qpos = np.concatenate((robot_arm_state, gripper_state))
            print("=== real robot current qpos: ", real_robot_qpos)
            ### reset agent qpos and get grasp_site pos in sim
            env.agent._robot.set_qpos(real_robot_qpos)
            sim_ee_pose_base = (
                env.agent._robot.get_root_pose().inv() * env.agent.grasp_site.get_pose()
            )  # in robot root coord
            sim_robotee_pose_base = (
                env.agent._robot.get_root_pose().inv() * robot_ee.get_pose()
            )  # in robot root coord
            real_ee_pose_mat = imp_opendoor.real_state.toolTobase_pos_m
            real_ee_pose_base = mat2pose(np.reshape(real_ee_pose_mat, (4, 4)))
            ee_pos_base = sim_ee_pose_base.p
            print("sim robot current qpos: ", env.agent._robot.get_qpos())
            print("sim ee pose (robot root frame): ", sim_ee_pose_base)
            print("sim robot link7 pose (robot root frame): ", sim_robotee_pose_base)
            print("real ee pose (robot root frame): ", real_ee_pose_base)

            ### concate to total states
            frame, _, _, cam_xyz, rs_depth_intrinsics = rs_observer.get_observation()
            cur_rgb, cur_depth = frame[:, :, :3], frame[:, :, 3]
            H, W = cur_depth.shape

            state_egorgbd = OrderedDict()
            state_egorgbd["qpos"] = torch.from_numpy(
                real_robot_qpos
            ).float()  # .to(device)  # (9)
            state_egorgbd["ee_pos_base"] = torch.from_numpy(
                ee_pos_base
            ).float()  # .to(device)  # (3)
            state_egorgbd["rgb"] = torch.from_numpy(
                cur_rgb.transpose((2, 0, 1))
            ).float()  # (3, 144, 256)
            state_egorgbd["depth"] = torch.from_numpy(
                cur_depth[None]
            ).float()  # (1, 144, 256)

            rgb_tensor = (
                torch.from_numpy((cur_rgb[None] / 255.0).transpose(0, 3, 1, 2))
                .float()
                .to(device)
            )  # 1,3,H,W
            depth_tensor = (
                torch.from_numpy(cur_depth[None, None]).float().to(device)
            )  # 1,1,H,W
            seg_map = unet_model.predict(torch.cat([rgb_tensor, depth_tensor], dim=1))[
                0
            ]
            seg_mc_map = torch.argmax(seg_map, dim=0).cpu().detach().numpy()  # (H, W)
            print(seg_mc_map.shape)

            print(f"+++++++++ step {step} +++++++++++")
            if save_plot_mode:
                save_subplots(
                    real_results_path,
                    step,
                    subplot_num,
                    cur_rgb / 255.0,
                    cur_depth,
                    seg_mc_map,
                    save_plot_mode,
                )

            action, _states = RL_model.predict(state_egorgbd, deterministic=True)
            if save_plot_mode == "save":
                real_robot_qpos_eepos_action = np.concatenate(
                    (real_robot_qpos, ee_pos_base, action), axis=0
                )  # 9 + 3 + 8
                np.save(
                    real_results_path / f"{step:02}_qpos_eepos_action.npy",
                    real_robot_qpos_eepos_action,
                )

            real_delta_qpos = scale_para(
                action[:7], action_range_low[:7], action_range_high[:7]
            )
            target_qpos = real_robot_qpos[:-2] + real_delta_qpos
            print("real qpos, action: ", real_robot_qpos[:-2], action)
            print("prev target robot qpos: ", prev_target_qpos)
            print("target robot qpos: ", np.asarray(target_qpos))
            for k in range(sim2real_step_num):
                imp_opendoor.exec_trajs(
                    target_qpos, stiffness=joint_kp, damping=joint_kd
                )

            ## control gripper
            if mode == "door" or mode == "drawer":
                target_gripper_pos = int(
                    (action[-1] + 1) / 2 * 255
                )  # Gripper Open->Close: -1->1 => 0->255
            elif mode == "faucet":
                target_gripper_pos = 255
            else:
                raise NotImplementedError(mode)
            print("target gripper qpos: ", target_gripper_pos)
            imp_opendoor.exec_gripper(target_gripper_pos)
            if mode != "faucet":
                time.sleep(gripper_wait)
            real_gripper_pos = gripper_listener.gPO
            print("real gripper qpos: ", real_gripper_pos)

            step += 1
    print("Total steps: ", step)

arti_mani/test/real_exps/SB3_opencabinet_2dseg.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import time
import warnings
from collections import OrderedDict, defaultdict
from copy import deepcopy

import cv2
import gym
import matplotlib.pyplot as plt
import numpy as np
import open3d as o3d
import sapien.core as sapien
import segmentation_models_pytorch as smp
import torch
import torch.nn.functional as F
import yaml
from arti_mani import ASSET_DIR, RLMODEL_DIR, VISUALMODEL_DIR
from arti_mani.algorithms.rl_iam.Sampling import (
    FCUWSampler,
    FPSSampler,
    RandomSampler,
    ScoreSampler,
    UniDownsampler,
)
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import (
    CustomUnet,
    CustomUnetNew,
    SplitUnet,
)
from arti_mani.algorithms.visual_net.Networks.keypoint_detection import (
    IntegralHumanPoseModel,
    soft_argmax,
)
from arti_mani.envs.arti_mani import ArtiMani
from arti_mani.test.real_exps.utils import mat2pose
from arti_mani.utils.contrib import uv2xyz
from arti_mani.utils.cv_utils import visualize_depth, visualize_seg
from arti_mani.utils.geometry import transform_points
from arti_mani.utils.o3d_utils import pcd_uni_down_sample_with_crop
from arti_mani.utils.sapien_utils import get_entity_by_name
from arti_mani.utils.wrappers import NormalizeActionWrapper
from observer import Observer
from stable_baselines3 import SAC
from xmate3Robotiq import Robotiq2FGripper_Listener, ROS_ImpOpendoor

# from xmate3Robotiq_new import Robotiq2FGripper_Listener, ROS_ImpOpendoor

warnings.filterwarnings("ignore")


def setup_seed(seed=1029):
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    rng = np.random.RandomState(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    return rng


def scale_para(para, output_low, output_high):
    """
    para: norm to [-1, 1]
    output_low, output_high: reproject to (output_low, output_high)
    """
    bias = 0.5 * (output_high + output_low)
    weight = 0.5 * (output_high - output_low)
    output = weight * para + bias

    return output


def load_kpt_model(path):
    model_path = VISUALMODEL_DIR / f"kpt_model/{path}/best.pth"
    kpt_model = IntegralHumanPoseModel(
        num_keypoints=3, num_deconv_layers=3, depth_dim=64, has_dropout=False
    )
    kpt_model.load_state_dict(torch.load(model_path))
    kpt_model.to(torch.device("cuda:0"))
    kpt_model.eval()
    return kpt_model


def save_subplots(real_results_path, step, subplot_num, rgb, depth, rgbplot, mode=None):
    plt.subplot(subplot_num[0], subplot_num[1], 1)
    plt.imshow(rgb)
    plt.title("rgb")
    plt.axis("off")
    plt.subplot(subplot_num[0], subplot_num[1], 2)
    plt.imshow(visualize_depth(depth)[..., ::-1])
    plt.title("depth")
    plt.axis("off")
    plt.subplot(subplot_num[0], subplot_num[1], 3)
    plt.imshow(rgbplot)
    plt.title("keypoints")
    plt.axis("off")
    if mode == "plot":
        plt.show()
    elif mode == "save":
        plt.savefig(
            real_results_path / f"test_{step:02}.png",
            bbox_inches="tight",
            pad_inches=0.1,
        )
    else:
        raise NotImplementedError


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "drawer"  # door, drawer, faucet
    rs_mode = "rgbd"  # "rgbd", "depth", "rgb"

    eval_seed = np.random.RandomState().randint(2**32)
    rng = setup_seed(eval_seed)
    print("experiment eval random seed: ", eval_seed)

    save_plot_mode = "plot"  # None, "plot", "save"
    record_id = 0
    gripper_wait = 0.3
    MAX_STEPS = 50
    subplot_num = [1, 3]
    out_shape = (64, 40, 64)
    kpts_max = np.array([255, 143, 1.0])
    kpts_min = np.array([0, 0, 0.18])
    test_name = f"{mode}0_demo{record_id}_3dkpt"
    if mode == "door":
        arti_ids = [0]  # 0
    elif mode == "drawer":
        arti_ids = [1]  # 1
    elif mode == "faucet":
        arti_ids = [5052]
    elif mode == "arti":
        arti_ids = [0, 1, 5052]
    else:
        raise NotImplementedError
    obs_mode = "state_egostereo_keypoints"
    control_mode = "pd_joint_delta_pos"

    if mode == "door":
        rl_exp_name = "baselines/sac4_tr4eval1door_state_egostereo_keypoints3_bs256_tr16up16_0"  # 0, 1, 2
    elif mode == "drawer":
        rl_exp_name = "baselines/sac4_tr4eval1drawer_state_egostereo_keypoints3_bs256_tr16up16_0"  # 0, 1, 2
    elif mode == "faucet":
        rl_exp_name = "baselines/sac4_tr4eval1faucet_state_egostereo_keypoints3_bs256_tr16up16_0"  # 0, 1, 2
    else:
        raise NotImplementedError
    model_path = (
        RLMODEL_DIR / f"{rl_exp_name}/best_model"
    )  # best_model, rl_model_1750000_steps, 2000000, 2250000, 2500000

    real_results_path = RLMODEL_DIR / f"{rl_exp_name}/real_result/{test_name}/"
    if save_plot_mode == "save":
        if not os.path.exists(real_results_path):
            os.makedirs(real_results_path)

    ### load kpt model
    kpt_exp_name = "20230307_182821_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2_dropout0.2_newdatafilter2drawer"
    kpt_model = load_kpt_model(kpt_exp_name)
    kpt_model.eval()

    print("log name: ", rl_exp_name)

    print("+++++ Build Real2Sim Env +++++++")
    env: ArtiMani = gym.make(
        "ArtiMani-v0",
        articulation_ids=arti_ids,
        segmodel_path=VISUALMODEL_DIR / f"kpt_model/{kpt_exp_name}",
        sample_mode=None,
        frame_num=0,
        sample_num=0,
        other_handle_visible=False,
        obs_mode=obs_mode,
        control_mode=control_mode,
        reward_mode="dense",  # "sparse", "dense"
        add_eepadpts=False,
        device=device,
    )
    env = NormalizeActionWrapper(env)
    env.seed(eval_seed)

    obs = env.reset()
    robot_ee: sapien.Link = get_entity_by_name(
        env.agent._robot.get_links(), "xmate3_link7"
    )
    ### Set Cabinet Pose
    env.unwrapped._articulation.set_root_pose(sapien.Pose([2, 0, 0], [1, 0, 0, 0]))

    ### Set Robot Pose
    ## real robot init
    if mode == "door" or mode == "drawer":
        qpos = np.array([1.4, -1.053, -2.394, 1.662, 1.217, 1.05, -0.8, 0.0, 0.0])
    elif mode == "faucet":
        qpos = np.array([-0.5, -0.143, 0, np.pi / 3, 0, 1.57, 1.57, 0.068, 0.068])
    else:
        raise NotImplementedError(mode)
    env.unwrapped._agent._robot.set_qpos(qpos)

    print(f"Action space {env.action_space}")
    print(f"Control mode {env.control_mode}")
    action_range_low = env.unwrapped.action_space.low
    action_range_high = env.unwrapped.action_space.high
    print("Input Action space range (low):", action_range_low)
    print("Input Action space range (high):", action_range_high)

    print("+++++ Setup robotiq controller +++++++")
    gripper_listener = Robotiq2FGripper_Listener()
    print("+++++ Setup robot arm controller +++++++")
    imp_opendoor = ROS_ImpOpendoor(
        gripper_listener=gripper_listener, control_mode=control_mode
    )
    if control_mode == "pd_joint_delta_pos":
        joint_kp = [200.0, 200.0, 200.0, 200.0, 50.0, 50.0, 50.0]
        joint_kd = [20.0, 20.0, 20.0, 20.0, 10.0, 10.0, 10.0]
        # joint_kp = [300.0, 300.0, 300.0, 300.0, 100.0, 100.0, 100.0]
        # joint_kd = [10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0]
    else:
        raise NotImplementedError
    print("current joint kp: ", joint_kp)
    print("current joint kd: ", joint_kd)

    ### load seg model
    kpt_exp_name = "20230307_182821_D64H40W64_deconv3_kpts3norm01addvis_uvz_lr1e-3_mobilenetv2_dropout0.2_newdatafilter2drawer"
    kpt_model = load_kpt_model(kpt_exp_name)
    kpt_model.eval()

    ### load RL model
    print("RL model path: ", model_path)
    RL_model = SAC.load(model_path, env=env, print_system_info=True)

    step = 0
    sim2real_step_num = 5
    prev_target_qpos = None
    prev_target_gripper_pos = None

    figs_dict = defaultdict(list)

    with torch.no_grad():
        rs_observer = Observer(
            base_camera_SN="918512071887",
            mode=rs_mode,
        )  # d435

        while step < MAX_STEPS:
            ### get agent qpos in real
            imp_opendoor.get_realstate()
            robot_arm_state = imp_opendoor.real_state.q_m
            gripper_state = np.repeat(
                gripper_listener.gPO / 255 * 0.068, 2
            )  # 0-255 => 0-0.068
            real_robot_qpos = np.concatenate((robot_arm_state, gripper_state))
            print("=== real robot current qpos: ", real_robot_qpos)
            ### reset agent qpos and get grasp_site pos in sim
            env.agent._robot.set_qpos(real_robot_qpos)
            sim_ee_pose_base = (
                env.agent._robot.get_root_pose().inv() * env.agent.grasp_site.get_pose()
            )  # in robot root coord
            sim_robotee_pose_base = (
                env.agent._robot.get_root_pose().inv() * robot_ee.get_pose()
            )  # in robot root coord
            real_ee_pose_mat = imp_opendoor.real_state.toolTobase_pos_m
            real_ee_pose_base = mat2pose(np.reshape(real_ee_pose_mat, (4, 4)))
            ee_pos_base = sim_ee_pose_base.p
            print("sim robot current qpos: ", env.agent._robot.get_qpos())
            print("sim ee pose (robot root frame): ", sim_ee_pose_base)
            print("sim robot link7 pose (robot root frame): ", sim_robotee_pose_base)
            print("real ee pose (robot root frame): ", real_ee_pose_base)

            ### concate to total states
            frame, _, _, cam_xyz, rs_depth_intrinsics = rs_observer.get_observation()
            cur_rgb, cur_depth = frame[:, :, :3] / 255.0, frame[:, :, 3]
            camH, camW = cur_depth.shape

            state_egokpts = OrderedDict()
            state_egokpts["qpos"] = torch.from_numpy(
                real_robot_qpos
            ).float()  # .to(device)  # (9)
            state_egokpts["ee_pos_base"] = torch.from_numpy(
                ee_pos_base
            ).float()  # .to(device)  # (3)

            rgb_tensor = (
                torch.from_numpy(cur_rgb[None].transpose(0, 3, 1, 2)).float().to(device)
            )  # 1,3,H,W
            depth_tensor = (
                torch.from_numpy(cur_depth[None, None]).float().to(device)
            )  # 1,1,H,W
            heatmap_out = kpt_model.forward(
                torch.cat([rgb_tensor, depth_tensor], dim=1)
            )  # (1, 3*D', H', W')
            uvz_pred_norm = (
                soft_argmax(heatmap_out, 3, out_shape)[0].detach().cpu().numpy()
            )  # (3, 3)
            uvz_pred = uvz_pred_norm * (kpts_max - kpts_min) + kpts_min
            state_egokpts["uvz_pred"] = uvz_pred.astype(np.float32)  # (J, 3)
            # xyz_cam_pred = uv2xyz(uvz_pred, rs_depth_intrinsics)
            # T_cam2ee = env.cam_para["hand_camera_extrinsic_base_frame"]
            # print("=== cam to ee: ", T_cam2ee)
            # xyz_ee_pred = transform_points(T_cam2ee, xyz_cam_pred)
            # state_egokpts["xyz_ee_pred"] = xyz_ee_pred.astype(np.float32)  # [J, 3]

            print(f"+++++++++ step {step} +++++++++++")
            if save_plot_mode:
                rgb_plot = cur_rgb.copy()
                for kpt_ind in range(uvz_pred.shape[0]):
                    cv2.circle(
                        rgb_plot,
                        (int(uvz_pred[kpt_ind, 1]), int(uvz_pred[kpt_ind, 0])),
                        radius=2,
                        color=(1, 0, 0),
                        thickness=-1,
                    )
                save_subplots(
                    real_results_path,
                    step,
                    subplot_num,
                    cur_rgb,
                    cur_depth,
                    rgb_plot,
                    save_plot_mode,
                )

            action, _states = RL_model.predict(state_egokpts, deterministic=True)
            if save_plot_mode == "save":
                real_robot_qpos_eepos_action = np.concatenate(
                    (real_robot_qpos, ee_pos_base, action), axis=0
                )  # 9 + 3 + 8
                np.save(
                    real_results_path / f"{step:02}_qpos_eepos_action.npy",
                    real_robot_qpos_eepos_action,
                )

            real_delta_qpos = scale_para(
                action[:7], action_range_low[:7], action_range_high[:7]
            )
            target_qpos = real_robot_qpos[:-2] + real_delta_qpos
            print("real qpos, action: ", real_robot_qpos[:-2], action)
            print("prev target robot qpos: ", prev_target_qpos)
            print("target robot qpos: ", np.asarray(target_qpos))
            for k in range(sim2real_step_num):
                imp_opendoor.exec_trajs(
                    target_qpos, stiffness=joint_kp, damping=joint_kd
                )

            ## control gripper
            if mode == "door" or mode == "drawer":
                target_gripper_pos = int(
                    (action[-1] + 1) / 2 * 255
                )  # Gripper Open->Close: -1->1 => 0->255
            elif mode == "faucet":
                target_gripper_pos = 255
            else:
                raise NotImplementedError(mode)
            print("target gripper qpos: ", target_gripper_pos)
            imp_opendoor.exec_gripper(target_gripper_pos)
            if mode != "faucet":
                time.sleep(gripper_wait)
            real_gripper_pos = gripper_listener.gPO
            print("real gripper qpos: ", real_gripper_pos)

            step += 1
    print("Total steps: ", step)

arti_mani/test/real_exps/SB3_opencabinet_3dkpt.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import time
import warnings
from collections import OrderedDict, defaultdict
from copy import deepcopy

import cv2
import gym
import matplotlib.pyplot as plt
import numpy as np
import open3d as o3d
import sapien.core as sapien
import segmentation_models_pytorch as smp
import torch
import torch.nn.functional as F
import yaml
from arti_mani import ASSET_DIR, RLMODEL_DIR, VISUALMODEL_DIR
from arti_mani.algorithms.rl_iam.Sampling import (
    FCUWSampler,
    FPSSampler,
    RandomSampler,
    ScoreSampler,
    UniDownsampler,
)
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import (
    CustomUnet,
    CustomUnetNew,
    SplitUnet,
)
from arti_mani.envs.arti_mani import ArtiMani
from arti_mani.test.real_exps.utils import mat2pose
from arti_mani.utils.contrib import apply_pose_to_points
from arti_mani.utils.cv_utils import visualize_depth, visualize_seg
from arti_mani.utils.geometry import transform_points
from arti_mani.utils.o3d_utils import pcd_uni_down_sample_with_crop
from arti_mani.utils.sapien_utils import get_entity_by_name
from arti_mani.utils.wrappers import NormalizeActionWrapper
from observer import Observer
from stable_baselines3 import SAC
from xmate3Robotiq import Robotiq2FGripper_Listener, ROS_ImpOpendoor

# from xmate3Robotiq_new import Robotiq2FGripper_Listener, ROS_ImpOpendoor

warnings.filterwarnings("ignore")


def setup_seed(seed=1029):
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    rng = np.random.RandomState(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    return rng


def scale_para(para, output_low, output_high):
    """
    para: norm to [-1, 1]
    output_low, output_high: reproject to (output_low, output_high)
    """
    bias = 0.5 * (output_high + output_low)
    weight = 0.5 * (output_high - output_low)
    output = weight * para + bias

    return output


def save_subplots(real_results_path, step, subplot_num, rgb, depth, mode=None):
    plt.subplot(subplot_num[0], subplot_num[1], 1)
    plt.imshow(rgb)
    plt.title("rgb")
    plt.axis("off")
    plt.subplot(subplot_num[0], subplot_num[1], 2)
    plt.imshow(visualize_depth(depth)[..., ::-1])
    plt.title("depth")
    plt.axis("off")
    if mode == "plot":
        plt.show()
    elif mode == "save":
        plt.savefig(
            real_results_path / f"test_{step:02}.png",
            bbox_inches="tight",
            pad_inches=0.1,
        )
    else:
        raise NotImplementedError


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "faucet"  # door, drawer, faucet
    rs_mode = "rgbd"  # "rgbd", "depth", "rgb"

    eval_seed = np.random.RandomState().randint(2**32)
    rng = setup_seed(eval_seed)
    print("experiment eval random seed: ", eval_seed)

    save_plot_mode = None  # "plot", "save"
    record_id = 0
    points_vis_mode = False
    gripper_wait = 0.1
    MAX_STEPS = 50
    subplot_num = [2, 3]
    test_name = f"{mode}1_demo{record_id}_dexpoints"
    fig_modes = [
        "rgb",
        "depth",
        "segpred",
        "samplepts",
        "uncertainty_norm",
        "weight_norm",
    ]
    cmaps = [
        (1, 0, 0),  # "red"
        (0, 0, 1),  # "blue"
        (1, 1, 0),  # "yellow"
        (0, 1, 0),  # "green"
        (0.627, 0.125, 0.941),  # "purple"
        (0.753, 0.753, 0.753),  # "grey"
    ]
    if mode == "door":
        arti_ids = [0]  # 0
    elif mode == "drawer":
        arti_ids = [1]  # 1
    elif mode == "faucet":
        arti_ids = [5052]
    elif mode == "arti":
        arti_ids = [0, 1, 5052]
    else:
        raise NotImplementedError
    obs_mode = "state_egostereo_dexpoints"  # state_egostereo_segpoints, state_egostereo_dexpoints
    control_mode = "pd_joint_delta_pos"

    if mode == "door":
        rl_exp_name = (
            "baselines/sac4_tr4eval1door_state_egostereo_dexpoints_0"  # 0, 1, 2
        )
    elif mode == "drawer":
        rl_exp_name = (
            "baselines/sac4_tr4eval1drawer_state_egostereo_dexpoints_2"  # 0, 1, 2
        )
    elif mode == "faucet":
        rl_exp_name = (
            "baselines/sac4_tr4eval1faucet_state_egostereo_dexpoints_3"  # 0, 2, 3
        )
    else:
        raise NotImplementedError
    # model_path = RLMODEL_DIR / f"{rl_exp_name}/rl_model_2000000_steps"  # best_model, rl_model_1750000_steps, 2000000, 2250000, 2500000
    model_path = (
        RLMODEL_DIR / f"{rl_exp_name}/best_model"
    )  # best_model, rl_model_1750000_steps, 2000000, 2250000, 2500000

    real_results_path = RLMODEL_DIR / f"{rl_exp_name}/real_result/{test_name}/"
    if save_plot_mode:
        if not os.path.exists(real_results_path):
            os.makedirs(real_results_path)

    print("log name: ", rl_exp_name)

    print("+++++ Build Real2Sim Env +++++++")
    env: ArtiMani = gym.make(
        "ArtiMani-v0",
        articulation_ids=arti_ids,
        segmodel_path=None,
        sample_mode=None,
        frame_num=0,
        sample_num=0,
        other_handle_visible=False,
        obs_mode=obs_mode,
        control_mode=control_mode,
        reward_mode="dense",  # "sparse", "dense"
        add_eepadpts=False,
        device=device,
    )
    env = NormalizeActionWrapper(env)
    env.seed(eval_seed)

    obs = env.reset()
    robot_ee: sapien.Link = get_entity_by_name(
        env.agent._robot.get_links(), "xmate3_link7"
    )
    ### Set Cabinet Pose
    env.unwrapped._articulation.set_root_pose(sapien.Pose([2, 0, 0], [1, 0, 0, 0]))

    ### Set Robot Pose
    ## real robot init
    if mode == "door" or mode == "drawer":
        qpos = np.array([1.4, -1.053, -2.394, 1.662, 1.217, 1.05, -0.8, 0.0, 0.0])
    elif mode == "faucet":
        qpos = np.array([-0.5, -0.143, 0, np.pi / 3, 0, 1.57, 1.57, 0.068, 0.068])
    else:
        raise NotImplementedError(mode)
    env.unwrapped._agent._robot.set_qpos(qpos)

    print(f"Action space {env.action_space}")
    print(f"Control mode {env.control_mode}")
    action_range_low = env.unwrapped.action_space.low
    action_range_high = env.unwrapped.action_space.high
    print("Input Action space range (low):", action_range_low)
    print("Input Action space range (high):", action_range_high)

    print("+++++ Setup robotiq controller +++++++")
    gripper_listener = Robotiq2FGripper_Listener()
    print("+++++ Setup robot arm controller +++++++")
    imp_opendoor = ROS_ImpOpendoor(
        gripper_listener=gripper_listener, control_mode=control_mode
    )
    if control_mode == "pd_joint_delta_pos":
        joint_kp = [200.0, 200.0, 200.0, 200.0, 50.0, 50.0, 50.0]
        joint_kd = [20.0, 20.0, 20.0, 20.0, 10.0, 10.0, 10.0]
        # joint_kp = [300.0, 300.0, 300.0, 300.0, 100.0, 100.0, 100.0]
        # joint_kd = [10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0]
    else:
        raise NotImplementedError
    print("current joint kp: ", joint_kp)
    print("current joint kd: ", joint_kd)

    ### load RL model
    print("RL model path: ", model_path)
    RL_model = SAC.load(model_path, env=env, print_system_info=True)

    step = 0
    sim2real_step_num = 5
    prev_target_qpos = None
    prev_target_gripper_pos = None

    figs_dict = defaultdict(list)

    with torch.no_grad():
        rs_observer = Observer(
            base_camera_SN="918512071887",
            mode=rs_mode,
        )  # d435

        while step < MAX_STEPS:
            ### get agent qpos in real
            imp_opendoor.get_realstate()
            robot_arm_state = imp_opendoor.real_state.q_m
            gripper_state = np.repeat(
                gripper_listener.gPO / 255 * 0.068, 2
            )  # 0-255 => 0-0.068
            real_robot_qpos = np.concatenate((robot_arm_state, gripper_state))
            print("=== real robot current qpos: ", real_robot_qpos)
            ### reset agent qpos and get grasp_site pos in sim
            env.agent._robot.set_qpos(real_robot_qpos)
            sim_ee_pose_base = (
                env.agent._robot.get_root_pose().inv() * env.agent.grasp_site.get_pose()
            )  # in robot root coord
            sim_robotee_pose_base = (
                env.agent._robot.get_root_pose().inv() * robot_ee.get_pose()
            )  # in robot root coord
            real_ee_pose_mat = imp_opendoor.real_state.toolTobase_pos_m
            real_ee_pose_base = mat2pose(np.reshape(real_ee_pose_mat, (4, 4)))
            ee_pos_base = sim_ee_pose_base.p
            print("sim robot current qpos: ", env.agent._robot.get_qpos())
            print("sim ee pose (robot root frame): ", sim_ee_pose_base)
            print("sim robot link7 pose (robot root frame): ", sim_robotee_pose_base)
            print("real ee pose (robot root frame): ", real_ee_pose_base)

            # # get ee pad pts
            # if step == 0:
            #     target_gripper_pos = gripper_listener.gPO  # 0 step, target pos == cur pos
            # gripper_qpos = target_gripper_pos / 255 * 0.068
            # print("real target gripper pos: ", target_gripper_pos)
            # cur_qpos = env.agent._robot.get_qpos().copy()
            # print("cur_qpos, : ", cur_qpos)
            # cur_qpos[-2:] = np.array([gripper_qpos] * 2)
            # env.agent._robot.set_qpos(cur_qpos)
            # print("real calibrated qpos, : ", cur_qpos)
            ee_cords = env.agent.sample_ee_coords(num_sample=20).reshape(
                -1, 3
            )  # [40, 3]
            # ee_points_eef = apply_pose_to_points(ee_cords, env.agent.grasp_site.get_pose().inv())
            T_world2robot = (
                env.agent._robot.get_root_pose().inv().to_transformation_matrix()
            )
            print("=== world to robot (correct): ", T_world2robot)
            eepad_pts_robot = transform_points(T_world2robot, ee_cords)  # [40, 3]
            ## compensate the calibration error
            # eepad_pts += np.array([-0.01, 0, 0])

            eepat_pts_label = np.zeros((eepad_pts_robot.shape[0], 2))  # (40, 2)
            eepat_pts_label[:, -1] = 1.0
            eepad_ptsC_robot = np.concatenate(
                (eepad_pts_robot, eepat_pts_label), axis=1
            )  # (40, 5)

            ### concate to total states
            frame, _, _, cam_xyz, rs_depth_intrinsics = rs_observer.get_observation()
            cur_rgb, cur_depth = frame[:, :, :3] / 255.0, frame[:, :, 3]
            H, W = cur_depth.shape

            ## cam to world
            hand_sensor = env.agent._sensors["hand"]
            sensor_rgb = hand_sensor._cam_rgb
            T_cam2world_sim = sensor_rgb.get_model_matrix()
            print("=== cam to world (wrong): ", T_cam2world_sim)
            T_cam2ee = env.cam_para["hand_camera_extrinsic_base_frame"]
            print("=== cam to ee: ", T_cam2ee)
            T_ee2world = env.agent.grasp_site.get_pose().to_transformation_matrix()
            print("=== ee to world: ", T_ee2world)
            T_cam2world = T_ee2world @ T_cam2ee
            print("=== cam to world (correct): ", T_cam2world)
            world_xyz = transform_points(
                T_cam2world, cam_xyz * np.array([1, -1, -1])
            )  # [H*W, 3]
            robot_xyz = transform_points(T_world2robot, world_xyz)  # [H*W, 3]
            ## compensate the calibration error
            # world_xyz += np.array([0.01, 0.01, 0])
            # world_xyz += np.array([0., 0., -0.01])

            sample_indices = pcd_uni_down_sample_with_crop(
                world_xyz,
                1000,
                min_bound=np.array([-0.5, -1, 1e-3]),
                max_bound=np.array([1, 1, 2]),
            )
            indices_len = len(sample_indices)
            if indices_len < 1000:
                sample_indices.extend(
                    np.random.choice(world_xyz.shape[0], 1000 - indices_len)
                )
            ds_robot_xyz = robot_xyz[sample_indices]
            ds_robot_label = np.zeros((ds_robot_xyz.shape[0], 2))  # (1000, 2)
            ds_robot_label[:, 0] = 1.0
            ds_ptsC_robot = np.concatenate((ds_robot_xyz, ds_robot_label), axis=1)

            if points_vis_mode:
                pcd0 = o3d.geometry.PointCloud()
                pcd0.points = o3d.utility.Vector3dVector(ds_robot_xyz)
                pcd0.paint_uniform_color([0, 0, 0])
                pcd1 = o3d.geometry.PointCloud()
                pcd1.points = o3d.utility.Vector3dVector(eepad_pts_robot)
                pcd1.paint_uniform_color([1, 0, 0])
                XYZ = o3d.geometry.TriangleMesh.create_coordinate_frame(
                    size=0.1, origin=[0.0, 0.0, 0.0]
                )
                objs = [pcd0, pcd1, XYZ]
                o3d.visualization.draw_geometries(objs)

            state_egodexpts = OrderedDict()
            state_egodexpts["qpos"] = torch.from_numpy(
                real_robot_qpos
            ).float()  # .to(device)  # (9)
            state_egodexpts["ee_pos_base"] = torch.from_numpy(
                ee_pos_base
            ).float()  # .to(device)  # (3)
            state_egodexpts["segsampled_ptsC"] = np.concatenate(
                (ds_ptsC_robot, eepad_ptsC_robot), axis=0
            ).astype(
                np.float32
            )  # (1000+40, 3+2)

            print(f"+++++++++ step {step} +++++++++++")
            if save_plot_mode:
                save_subplots(
                    real_results_path,
                    step,
                    subplot_num,
                    cur_rgb,
                    cur_depth,
                    save_plot_mode,
                )

            action, _states = RL_model.predict(state_egodexpts, deterministic=True)
            if save_plot_mode:
                real_robot_qpos_eepos_action = np.concatenate(
                    (real_robot_qpos, ee_pos_base, action), axis=0
                )  # 9 + 3 + 8
                np.save(
                    real_results_path / f"{step:02}_qpos_eepos_action.npy",
                    real_robot_qpos_eepos_action,
                )

            real_delta_qpos = scale_para(
                action[:7], action_range_low[:7], action_range_high[:7]
            )
            target_qpos = real_robot_qpos[:-2] + real_delta_qpos
            print("real qpos, action: ", real_robot_qpos[:-2], action)
            print("prev target robot qpos: ", prev_target_qpos)
            print("target robot qpos: ", np.asarray(target_qpos))
            for k in range(sim2real_step_num):
                imp_opendoor.exec_trajs(
                    target_qpos, stiffness=joint_kp, damping=joint_kd
                )

            ## control gripper
            if mode == "door" or mode == "drawer":
                target_gripper_pos = int(
                    (action[-1] + 1) / 2 * 255
                )  # Gripper Open->Close: -1->1 => 0->255
            elif mode == "faucet":
                target_gripper_pos = 255
            else:
                raise NotImplementedError(mode)
            print("target gripper qpos: ", target_gripper_pos)
            imp_opendoor.exec_gripper(target_gripper_pos)
            if mode != "faucet":
                time.sleep(gripper_wait)
            real_gripper_pos = gripper_listener.gPO
            print("real gripper qpos: ", real_gripper_pos)

            step += 1
    print("Total steps: ", step)

arti_mani/test/real_exps/SB3_opencabinet_dexpoints.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import time
from collections import defaultdict

import cv2
import numpy as np
import pyrealsense2 as rs
from arti_mani.utils.cv_utils import visualize_depth

color_cache = defaultdict(lambda: {})
WIDTH, HEIGHT = 640, 360
RS_WIDTH, RS_HEIGHT = 256, 144
RS_RATIO = 0.4
save_path = "./"


def get_depth_real():
    # ...from Camera 1
    pipeline_1 = rs.pipeline()
    config_1 = rs.config()
    config_1.enable_device("918512071887")  # hand_camera d435
    config_1.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
    config_1.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
    # ...from Camera 2
    pipeline_2 = rs.pipeline()
    config_2 = rs.config()
    config_2.enable_device("104122061117")  # base_camera d415
    config_2.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
    config_2.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)

    # Start streaming from both cameras
    pipeline_1.start(config_1)
    pipeline_2.start(config_2)

    # Camera 1
    # Wait for a coherent pair of frames: depth and color
    frames_1 = pipeline_1.wait_for_frames()
    depth_frame_1 = frames_1.get_depth_frame()

    # Convert images to numpy arrays
    depth_image_1 = np.asanyarray(depth_frame_1.get_data())
    depth_image_1_small = cv2.resize(depth_image_1, (320, 180))

    # Camera 2
    # Wait for a coherent pair of frames: depth and color
    frames_2 = pipeline_2.wait_for_frames()
    depth_frame_2 = frames_2.get_depth_frame()
    depth_image_2 = np.asanyarray(depth_frame_2.get_data())
    depth_image_2_small = cv2.resize(depth_image_2, (320, 180))
    print(depth_image_1_small.shape, depth_image_2_small.shape)
    depth_images = np.stack((depth_image_1_small, depth_image_2_small), axis=2)
    return depth_images


def depth2xyz(depth_map, depth_cam_matrix, flatten=True, depth_scale=0.0010):
    """
    depth_map = np.random.randint(0, 10000,(720, 1280))
    depth_cam_matrix = np.array([[540, 0,  640],
                                 [0,   540,360],
                                 [0,   0,    1]])
    pc = depth2xyz(depth_map, depth_cam_matrix)
    """
    fx, fy = depth_cam_matrix[0, 0], depth_cam_matrix[1, 1]
    cx, cy = depth_cam_matrix[0, 2], depth_cam_matrix[1, 2]
    h, w = np.mgrid[0 : depth_map.shape[0], 0 : depth_map.shape[1]]
    z = depth_map * depth_scale
    x = (w - cx) * z / fx
    y = (h - cy) * z / fy

    xyz = np.dstack((x, y, z)).reshape(-1, 3) if flatten else np.dstack((x, y, z))
    return xyz  # [N,3]


class Observer:
    def __init__(self, base_camera_SN="104122061117", mode="depth") -> None:
        self.pipeline = rs.pipeline()
        config = rs.config()
        config.enable_device(base_camera_SN)
        config.enable_stream(rs.stream.depth, WIDTH, HEIGHT, rs.format.z16, 30)
        config.enable_stream(rs.stream.color, WIDTH, HEIGHT, rs.format.rgb8, 30)

        # set align mode
        align_to = rs.stream.color  # depth align to color
        # align_to = rs.stream.depth  # color align to depth
        self.alignedFs = rs.align(align_to)

        self.cfg = self.pipeline.start(config)
        time.sleep(2)

        # Getting the depth sensor's depth scale (see rs-align example for explanation)
        depth_sensor = self.cfg.get_device().first_depth_sensor()
        depth_sensor.set_option(rs.option.exposure, 5000.000)
        print("depth sensor exposure: ", depth_sensor.get_option(rs.option.exposure))

        color_sensor = self.cfg.get_device().first_color_sensor()
        color_sensor.set_option(rs.option.exposure, 300.000)
        print("color sensor exposure: ", color_sensor.get_option(rs.option.exposure))

        self.depth_scale = depth_sensor.get_depth_scale()
        print("Depth Scale is: ", self.depth_scale)

        self.mode = mode

        ###################################################################################
        # realsense-viewer, orders of the postprocess:
        # decimation_filter --> HDR Merge --> threshold_filter --> Depth to Disparity --> spatial_filter
        # --> temporal_filter --> Disparity to Depth
        g_rs_downsample_filter = rs.decimation_filter(magnitude=1)  # 2 ** 1,  # 下采样率
        g_rs_thres_filter = rs.threshold_filter(min_dist=0.2, max_dist=2.0)
        g_rs_depth2disparity_trans = rs.disparity_transform(True)
        g_rs_spatical_filter = rs.spatial_filter(
            magnitude=2,
            smooth_alpha=0.5,
            smooth_delta=20,
            hole_fill=0,
        )
        g_rs_templ_filter = rs.temporal_filter(
            smooth_alpha=0.1, smooth_delta=40.0, persistence_control=3
        )
        g_rs_disparity2depth_trans = rs.disparity_transform(False)
        self.g_rs_depth_postprocess_list = [
            g_rs_downsample_filter,
            g_rs_thres_filter,
            g_rs_depth2disparity_trans,
            g_rs_spatical_filter,
            g_rs_templ_filter,
            g_rs_disparity2depth_trans,
        ]
        ###################################################################################

        self.colorizer = rs.colorizer()

    def get_observation(self):
        try:
            while True:
                # Wait for a coherent pair of frames: depth and color
                frames = self.pipeline.wait_for_frames()
                aligned_frames = self.alignedFs.process(frames)
                frame_num = aligned_frames.frame_number
                timestamp = aligned_frames.timestamp
                print("timestramp: ", timestamp)
                print("frame_num: ", frame_num)

                depth_frame = aligned_frames.get_depth_frame()
                color_frame = aligned_frames.get_color_frame()

                if not depth_frame or not color_frame:
                    continue

                depth_frame_filter = depth_frame
                for trans in self.g_rs_depth_postprocess_list:
                    depth_frame_filter = trans.process(depth_frame_filter)
                depth_frame = depth_frame_filter

                # Convert images to numpy arrays
                color_image = np.asanyarray(color_frame.get_data())
                depth_image = np.asanyarray(depth_frame.get_data())
                colorized_depth = np.asanyarray(
                    self.colorizer.colorize(depth_frame).get_data()
                )

                print(color_image.shape, depth_image.shape, colorized_depth.shape)
                print(color_image.dtype, depth_image.dtype, colorized_depth.dtype)
                print(color_image.max(), depth_image.max(), colorized_depth.max())
                print(color_image.min(), depth_image.min(), colorized_depth.min())

                # resize color and depth
                rs_color_image = cv2.resize(
                    color_image, (RS_WIDTH, RS_HEIGHT), interpolation=cv2.INTER_NEAREST
                )
                rs_depth_image = cv2.resize(
                    depth_image, (RS_WIDTH, RS_HEIGHT), interpolation=cv2.INTER_NEAREST
                )

                rs_depth_colormap = visualize_depth(rs_depth_image)

                print(
                    rs_color_image.shape,
                    rs_depth_image.shape,
                    colorized_depth.shape,
                    rs_depth_colormap.shape,
                )
                print(
                    rs_color_image.dtype,
                    rs_depth_image.dtype,
                    colorized_depth.dtype,
                    rs_depth_colormap.dtype,
                )
                print(
                    rs_color_image.max(),
                    rs_depth_image.max(),
                    colorized_depth.max(),
                    rs_depth_colormap.max(),
                )
                print(
                    rs_color_image.min(),
                    rs_depth_image.min(),
                    colorized_depth.min(),
                    rs_depth_colormap.min(),
                )

                if self.mode == "rgbd":
                    frame = np.concatenate(
                        [rs_color_image, rs_depth_image[..., None] * self.depth_scale],
                        axis=2,
                    ).astype(np.float32)
                elif self.mode == "rgb":
                    frame = rs_color_image
                elif self.mode == "depth":
                    frame = rs_depth_image[..., None] * self.depth_scale
                else:
                    raise NotImplemented

                depth_intrin = (
                    depth_frame.get_profile().as_video_stream_profile().get_intrinsics()
                )

                extrinsics = depth_frame.get_profile().get_extrinsics_to(
                    color_frame.get_profile()
                )
                rotation, translation = extrinsics.rotation, extrinsics.translation
                print("extrinsics: ")
                print(rotation, translation)

                rs_g_depth_intrinsics_matrix = np.array(
                    [
                        [depth_intrin.fx * RS_RATIO, 0.0, depth_intrin.ppx * RS_RATIO],
                        [0.0, depth_intrin.fy * RS_RATIO, depth_intrin.ppy * RS_RATIO],
                        [0, 0, 1.0],
                    ]
                )
                # print(rs_g_depth_intrinsics_matrix)
                points_xyz = depth2xyz(
                    rs_depth_image, rs_g_depth_intrinsics_matrix, self.depth_scale
                )
                # map to color frame (use extrinsics)
                points_xyz = points_xyz @ np.array(rotation).reshape(
                    (3, 3)
                ).transpose() + np.array(translation)

                print(points_xyz.shape, points_xyz.min(), points_xyz.max())

                return (
                    frame,
                    colorized_depth,
                    rs_depth_colormap,
                    points_xyz,
                    rs_g_depth_intrinsics_matrix,
                )

        finally:
            # self.pipeline.stop()
            pass

arti_mani/test/real_exps/observer.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import glob
import os

import gym
import matplotlib.pyplot as plt
import numpy as np
import sapien.core as sapien
import torch
from arti_mani import REALEXP_DIR
from arti_mani.envs.arti_mani import ArtiMani
from arti_mani.utils.sapien_utils import get_entity_by_name
from arti_mani.utils.visualization import images_to_video
from arti_mani.utils.wrappers import NormalizeActionWrapper


def setup_seed(seed=1029):
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    rng = np.random.RandomState(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    return rng


def scale_para(para, output_low, output_high):
    """
    para: norm to [-1, 1]
    output_low, output_high: reproject to (output_low, output_high)
    """
    bias = 0.5 * (output_high + output_low)
    weight = 0.5 * (output_high - output_low)
    output = weight * para + bias

    return output


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    device = "cuda:0"
    mode = "faucet"  # door, drawer, faucet
    record_id = 3
    save_mode = True
    eval_seed = np.random.RandomState().randint(2**32)
    rng = setup_seed(eval_seed)
    print("experiment eval random seed: ", eval_seed)

    ## load real trajs
    real_traj_path = REALEXP_DIR / f"real_results/{mode}s/real2sim"
    real_traj_list = sorted(os.listdir(real_traj_path))
    real_traj_file = real_traj_list[record_id]
    print("real2sim record name: ", real_traj_file)
    traj_len = len(glob.glob(str(real_traj_path / real_traj_file) + "/test_*.png"))
    print("real2sim record length: ", traj_len)
    traj_qpos = []
    traj_eepos = []
    traj_action = []
    for ind in range(traj_len):
        qpos_eepos_action = np.load(
            real_traj_path / f"{real_traj_file}/{ind:02}_qpos_eepos_action.npy"
        )
        traj_qpos.append(qpos_eepos_action[:9])
        traj_eepos.append(qpos_eepos_action[9:12])
        traj_action.append(qpos_eepos_action[12:])

    # real2sim results
    real2sim_results_path = REALEXP_DIR / f"real_results/real2sim_videos"

    if mode == "door":
        arti_ids = [0]  # 0
    elif mode == "drawer":
        arti_ids = [1]  # 1
    elif mode == "faucet":
        arti_ids = [5052]
    elif mode == "arti":
        arti_ids = [0, 1, 5052]
    else:
        raise NotImplementedError
    obs_mode = "state_egostereo_rgbd"
    control_mode = "pd_joint_delta_pos"

    print("+++++ Build Real2Sim Env +++++++")
    env: ArtiMani = gym.make(
        "ArtiMani-v0",
        articulation_ids=arti_ids,
        segmodel_path=None,
        sample_mode="random_sample",
        frame_num=0,
        sample_num=32,
        other_handle_visible=False,
        obs_mode=obs_mode,
        control_mode=control_mode,
        reward_mode="dense",  # "sparse", "dense"
        device=device,
    )
    env = NormalizeActionWrapper(env)
    env.seed(eval_seed)
    print(f"Action space {env.action_space}")
    print(f"Control mode {env.control_mode}")
    action_range_low = env.unwrapped.action_space.low
    action_range_high = env.unwrapped.action_space.high
    print("Input Action space range (low):", action_range_low)
    print("Input Action space range (high):", action_range_high)

    obs = env.reset()
    robot_ee: sapien.Link = get_entity_by_name(
        env.agent._robot.get_links(), "xmate3_link7"
    )
    ### Set Cabinet0 Pose
    # z = env.unwrapped._articulation.get_root_pose().p[2]
    # # real init
    # x = -0.6 + 100 * 0.025  # 52
    # y = 0.4 - 16 * 0.025  # 16
    # rz = 106  # arctan(2/7) + 90
    # q4 = np.sin(0.5 * (rz / 180 * 3.1415926))
    # env.unwrapped._articulation.set_root_pose(
    #     sapien.Pose([x, y, z], [np.sqrt(1 - q4 ** 2), 0, 0, q4])
    # )
    # env.unwrapped._articulation.set_qpos([0, 0])
    env.unwrapped._articulation.set_root_pose(sapien.Pose([2, 0, 0], [1, 0, 0, 0]))

    ### Set Robot Pose
    ## real robot init
    if mode == "door" or mode == "drawer":
        qpos = np.array([1.4, -1.053, -2.394, 1.662, 1.217, 1.05, -0.8, 0.0, 0.0])
    elif mode == "faucet":
        qpos = np.array([-0.5, -0.143, 0, np.pi / 3, 0, 1.57, 1.57, 0.068, 0.068])
    else:
        raise NotImplementedError(mode)
    env.unwrapped._agent._robot.set_qpos(qpos)

    viewer = env.unwrapped.render()
    print("Press [b] to start")
    while True:
        if viewer.window.key_down("b"):
            break
        env.render()

    if save_mode:
        frames = []
        base_rgbd = env.render("cameras")
        frames.append(base_rgbd)
        qpos_err = []
    for step in range(traj_len):
        cur_qpos = traj_qpos[step]
        cur_eepos = traj_eepos[step]
        # env.unwrapped._agent._robot.set_qpos(cur_qpos)

        # task action
        cur_action = traj_action[step]
        print(step, cur_qpos, cur_eepos, cur_action)
        obs, rew, done, info = env.step(cur_action)
        # print(f"STEP: {step:5d}, rew: {rew}, done: {done}, info: {info}")
        base_rgbd = env.render(mode="cameras")
        if save_mode:
            frames.append(base_rgbd)
            qpos_err.append(env.unwrapped._agent._robot.get_qpos() - cur_qpos)
        viewer = env.unwrapped.render()
        # print("Press [c] to start")
        # while True:
        #     if viewer.window.key_down("c"):
        #         break
        #     env.render()

    if save_mode:
        images_to_video(
            images=frames,
            output_dir=str(real2sim_results_path),
            video_name=f"{record_id}_trajectory_{real_traj_file[:15]}",
        )
        qpos_err = np.array(qpos_err)
        labels = np.array([f"j{i}" for i in range(7)])
        plt.subplot(2, 1, 1)
        plt.plot(np.arange(traj_len), qpos_err[:, -2], label=labels)
        plt.title("qpos err (sim - real)")
        plt.legend()
        plt.subplot(2, 1, 2)
        plt.plot(np.arange(traj_len), qpos_err[:, -2:])
        plt.title("gripper pos err (sim - real)")
        plt.savefig(
            real2sim_results_path / f"{record_id}_{mode}_real2sim_err.png",
            dpi=300,
            pad_inches=0.2,
        )

arti_mani/test/real_exps/real2sim_replay.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import time
import warnings
from collections import OrderedDict
from copy import deepcopy

import cv2
import gym
import matplotlib.pyplot as plt
import numpy as np
import open3d as o3d
import sapien.core as sapien
import segmentation_models_pytorch as smp
import torch
import torch.nn.functional as F
import yaml
from arti_mani import ASSET_DIR, RLMODEL_DIR, VISUALMODEL_DIR
from arti_mani.algorithms.rl_iam.Sampling import (
    FCUWSampler,
    FPSSampler,
    RandomSampler,
    ScoreSampler,
    UniDownsampler,
)
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import (
    CustomUnet,
    CustomUnetNew,
    SplitUnet,
)
from arti_mani.envs.arti_mani import ArtiMani
from arti_mani.test.real_exps.utils import mat2pose
from arti_mani.utils.contrib import apply_pose_to_points
from arti_mani.utils.cv_utils import visualize_depth, visualize_seg
from arti_mani.utils.geometry import transform_points
from arti_mani.utils.o3d_utils import pcd_uni_down_sample_with_crop
from arti_mani.utils.sapien_utils import get_entity_by_name
from arti_mani.utils.visualization import images_to_video
from arti_mani.utils.wrappers import NormalizeActionWrapper
from observer import Observer
from stable_baselines3 import SAC
from tqdm import tqdm
from xmate3Robotiq import Robotiq2FGripper_Listener, ROS_ImpOpendoor


def interp_steps(traj_qpos_init, traj_qpos_end, step_num):
    delta_qpos = traj_qpos_end - traj_qpos_init
    interp_qpos = delta_qpos / step_num
    traj_qpos_seq = np.tile(traj_qpos_init, (step_num, 1))
    for step in range(step_num):
        traj_qpos_seq[step] += interp_qpos * (step + 1)
    return traj_qpos_seq


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)

    test_name = "drawer0_demo2_frameweight_sample_3_32_smp0327"
    rl_exp_name = "sac2_tr4eval1drawer_state_egostereo_segpoints_random_samplef1s32_100steps_padfriction1_p3000d100_tr16up16_noeepadpts_newextric_smp0219_adjustreward_2close_initrot"
    real_results_path = RLMODEL_DIR / f"{rl_exp_name}/real_result/{test_name}/"
    real_qpos = []
    for ind in range(43):
        real_qpos_ind = np.load(real_results_path / f"{ind:02}_qpos_eepos_action.npy")
        real_qpos.append(real_qpos_ind[:9])

    control_mode = "pd_joint_delta_pos"

    # logger.info(f"Reward mode {env.reward_mode}")
    print("+++++ Setup robotiq controller +++++++")
    gripper_listener = Robotiq2FGripper_Listener()
    print("+++++ Setup robot arm controller +++++++")
    imp_opendoor = ROS_ImpOpendoor(
        gripper_listener=gripper_listener, control_mode=control_mode
    )
    if control_mode == "pd_joint_delta_pos":
        # joint_kp = [200.0, 200.0, 200.0, 200.0, 50.0, 50.0, 50.0]
        # joint_kd = [20.0, 20.0, 20.0, 20.0, 10.0, 10.0, 10.0]
        joint_kp = [300, 300, 300, 300, 100, 100, 100]
        joint_kd = [20, 20, 20, 20, 5, 5, 5]
    else:
        raise NotImplementedError
    print("current joint kp: ", joint_kp)
    print("current joint kd: ", joint_kd)
    qpos_interp = interp_steps(real_qpos[16], real_qpos[29], 5)
    real_interp_qpos = np.concatenate(
        (real_qpos[:17], qpos_interp, real_qpos[29:]), axis=0
    )

    step = 0
    for real_robot_qpos in real_interp_qpos:
        # if step > 15 and step < 29:
        #     step += 1
        #     continue
        print(f"+++++++++ step {step} +++++++++++")
        print("=== real robot current qpos: ", real_robot_qpos)
        target_qpos = real_robot_qpos[:-2]

        print("target robot qpos: ", np.asarray(target_qpos))
        for k in range(5):
            imp_opendoor.exec_trajs(target_qpos, stiffness=joint_kp, damping=joint_kd)
        target_gripper_pos = int(real_robot_qpos[-1] / 0.068 * 255)
        print("target gripper qpos: ", target_gripper_pos)
        imp_opendoor.exec_gripper(target_gripper_pos)
        time.sleep(0.4)
        real_gripper_pos = gripper_listener.gPO
        print("real gripper qpos: ", real_gripper_pos)

        step += 1

arti_mani/test/real_exps/realTRAJ_replay.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import time
import warnings
from collections import OrderedDict
from copy import deepcopy

import cv2
import gym
import matplotlib.pyplot as plt
import numpy as np
import open3d as o3d
import sapien.core as sapien
import segmentation_models_pytorch as smp
import torch
import torch.nn.functional as F
import yaml
from arti_mani import ASSET_DIR, RLMODEL_DIR, VISUALMODEL_DIR
from arti_mani.algorithms.rl_iam.Sampling import (
    FCUWSampler,
    FPSSampler,
    RandomSampler,
    ScoreSampler,
    UniDownsampler,
)
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import (
    CustomUnet,
    CustomUnetNew,
    SplitUnet,
)
from arti_mani.envs.arti_mani import ArtiMani
from arti_mani.test.real_exps.utils import mat2pose
from arti_mani.utils.contrib import apply_pose_to_points
from arti_mani.utils.cv_utils import visualize_depth, visualize_seg
from arti_mani.utils.geometry import transform_points
from arti_mani.utils.o3d_utils import pcd_uni_down_sample_with_crop
from arti_mani.utils.sapien_utils import get_entity_by_name
from arti_mani.utils.visualization import images_to_video
from arti_mani.utils.wrappers import NormalizeActionWrapper
from observer import Observer
from stable_baselines3 import SAC
from tqdm import tqdm
from xmate3Robotiq import Robotiq2FGripper_Listener, ROS_ImpOpendoor


def setup_seed(seed=1029):
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    rng = np.random.RandomState(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    return rng


def scale_para(para, output_low, output_high):
    """
    para: norm to [-1, 1]
    output_low, output_high: reproject to (output_low, output_high)
    """
    bias = 0.5 * (output_high + output_low)
    weight = 0.5 * (output_high - output_low)
    output = weight * para + bias

    return output


def load_vismodel(vismodel_path, load_device):
    model_path = VISUALMODEL_DIR / f"smp_model/{vismodel_path}/best.pth"
    config_path = VISUALMODEL_DIR / f"smp_model/{vismodel_path}/config.yaml"
    with open(config_path, "r") as f:
        cfg = yaml.safe_load(f)
    smp_cfg = cfg["smp_config"]
    if smp_cfg["mode"] == "RGBD":
        in_ch = 4
    elif smp_cfg["mode"] == "RGB":
        in_ch = 3
    elif smp_cfg["mode"] == "D":
        in_ch = 1
    else:
        raise NotImplementedError
    segmodel = CustomUnet(
        has_dropout=False,
        encoder_name=smp_cfg["encoder"],
        encoder_depth=smp_cfg["encoder_depth"],
        decoder_channels=smp_cfg["decoder_channels"],
        encoder_weights=smp_cfg["encoder_weights"],
        in_channels=in_ch,
        classes=cfg["num_classes"],
        activation=smp_cfg["activation"],
    )
    segmodel.load_state_dict(torch.load(model_path))
    segmodel.to(torch.device(load_device))
    segmodel.eval()
    return segmodel


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    eval_seed = np.random.RandomState().randint(2**32)
    rng = setup_seed(eval_seed)
    print("experiment eval random seed: ", eval_seed)
    device = "cuda:0"

    record_id = 3
    sim2real_step_num = 5
    gripper_wait = 0.5
    name_suffix = (
        f"_repeatstep{sim2real_step_num}_gripperwait{gripper_wait}_realreset_kp500kd45"
    )
    mode = "door"  # door, drawer, faucet
    save_mode = True
    render_mode = False
    MAX_STEPS = 5

    smp_exp_name = "20230219_000940_train52-val18_384_noDR_randombg_aug_stereo_bs16_focalloss_0.5step50lr0.001_RGBDunet-163264128_mobilenet_v2"
    vis_model = load_vismodel(smp_exp_name, device)
    segmodel_path = VISUALMODEL_DIR / f"smp_model/{smp_exp_name}"

    if mode == "door":
        # rl_exp_name = "sac4_tr4eval1door_state_egostereo_segpoints_random_samplef1s32_100steps_padfriction1_p3000d100_tr16up16_noeepadpts_newextric_ptsaddnoise0.01_smp0219"
        rl_exp_name = "sac2_tr12eval3arti_state_egostereo_segpoints_random_samplef1s32_100steps_smp0219_largerange_cabinetrew"
    elif mode == "drawer":
        rl_exp_name = "sac2_tr4eval1drawer_state_egostereo_segpoints_random_samplef1s32_100steps_padfriction1_p3000d100_tr16up16_noeepadpts_newextric_smp0219_adjustreward_2close_initrot"
    elif mode == "faucet":
        rl_exp_name = "sac4_tr4eval1faucet_state_egostereo_segpoints_random_samplef1s32_100steps_padfriction1_p3000d100_tr16up16_noeepadpts_newextric_ptsaddnoise0.01_smp0219"
    else:
        raise NotImplementedError
    model_name = "best_model"  # "rl_model_1000000_steps"
    model_path = RLMODEL_DIR / f"{rl_exp_name}/{model_name}"
    eval_videos_path = RLMODEL_DIR / f"{rl_exp_name}/sim2real_videos/{model_name}/"
    if not os.path.exists(eval_videos_path):
        os.makedirs(eval_videos_path)

    print("+++++ Build Real2Sim Env +++++++")
    if mode == "door":
        arti_ids = [0]  # 0
    elif mode == "drawer":
        arti_ids = [1]  # 1
    elif mode == "faucet":
        arti_ids = [5052]
    elif mode == "arti":
        arti_ids = [0, 1, 5052]
    else:
        raise NotImplementedError
    obs_mode = "state_egostereo_segpoints"
    control_mode = "pd_joint_delta_pos"
    sample_mode = "random_sample"  # full_downsample, score_sample, random_sample, frameweight_sample, fps_sample
    frame_num = 1
    pts_sample_num = 32
    env: ArtiMani = gym.make(
        "ArtiMani-v0",
        articulation_ids=arti_ids,
        segmodel_path=segmodel_path,
        sample_mode=sample_mode,
        frame_num=frame_num,
        sample_num=pts_sample_num,
        other_handle_visible=False,
        obs_mode=obs_mode,
        control_mode=control_mode,
        reward_mode="dense",  # "sparse", "dense"
        device=device,
    )
    env = NormalizeActionWrapper(env)
    env.seed(eval_seed)

    obs = env.reset()
    robot_ee: sapien.Link = get_entity_by_name(
        env.agent._robot.get_links(), "xmate3_link7"
    )
    ### Set Cabinet Pose
    # env.unwrapped._articulation.set_root_pose(
    #     sapien.Pose([2, 0, 0], [1, 0, 0, 0])
    # )

    ### load RL model
    RL_model = SAC.load(
        model_path,
        env=env,
        print_system_info=True,
    )

    ### Set Robot Pose
    ## real robot init
    if mode == "door" or mode == "drawer":
        qpos = np.array([1.4, -1.053, -2.394, 1.662, 1.217, 1.05, -0.8, 0.0, 0.0])
    elif mode == "faucet":
        qpos = np.array([-0.5, -0.143, 0, np.pi / 3, 0, 1.57, 1.57, 0.068, 0.068])
    else:
        raise NotImplementedError(mode)
    env.unwrapped._agent._robot.set_qpos(qpos)

    if render_mode:
        viewer = env.unwrapped.render()
        print("Press [b] to start")
        while True:
            if viewer.window.key_down("b"):
                break
            env.render()

    # logger.info(f"Observation space {env.observation_space}")
    print(f"Action space {env.action_space}")
    print(f"Control mode {env.control_mode}")
    action_range_low = env.unwrapped.action_space.low
    action_range_high = env.unwrapped.action_space.high
    print("Input Action space range (low):", action_range_low)
    print("Input Action space range (high):", action_range_high)

    # logger.info(f"Reward mode {env.reward_mode}")
    print("+++++ Setup robotiq controller +++++++")
    gripper_listener = Robotiq2FGripper_Listener()
    print("+++++ Setup robot arm controller +++++++")
    imp_opendoor = ROS_ImpOpendoor(
        gripper_listener=gripper_listener, control_mode=control_mode
    )
    if control_mode == "pd_joint_delta_pos":
        # joint_kp = [100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0]
        # joint_kd = [5.5, 5.5, 5.5, 5.5, 5.5, 5.5, 5.5]
        joint_kp = [200.0, 200.0, 200.0, 200.0, 50.0, 50.0, 50.0]
        joint_kd = [20.0, 20.0, 20.0, 20.0, 10.0, 10.0, 10.0]
        # joint_kp = [300.0, 300.0, 300.0, 300.0, 100.0, 100.0, 100.0]
        # joint_kd = [20.0, 20.0, 20.0, 20.0, 10.0, 10.0, 10.0]
        # joint_kp = [1000, 1000, 1000, 1000, 100, 100, 100]
        # joint_kd = [60, 60, 60, 60, 20, 20, 20]
        # joint_kp = [500, 500, 500, 500, 50, 50, 50]
        # joint_kd = [45, 45, 45, 45, 14, 14, 14]
        # joint_kp = [200, 200, 200, 200, 20, 20, 20]
        # joint_kd = [20, 20, 20, 20, 5, 5, 5]
    else:
        raise NotImplementedError
    print("current joint kp: ", joint_kp)
    print("current joint kd: ", joint_kd)

    ### load RL model
    RL_model = SAC.load(model_path, env=env, print_system_info=True)

    obs = env.reset()
    success_flag = False
    if save_mode:
        frames = []
        base_rgbd = env.render("cameras")
        frames.append(base_rgbd)
        sim_qpos_traj = []
        real_qpos_traj = []
    for step in tqdm(range(MAX_STEPS), colour="green", leave=False):
        ## get current real state
        imp_opendoor.get_realstate()
        robot_arm_state = imp_opendoor.real_state.q_m
        gripper_state = np.repeat(
            gripper_listener.gPO / 255 * 0.068, 2
        )  # 0-255 => 0-0.068
        real_robot_qpos = np.concatenate((robot_arm_state, gripper_state))
        ### reset agent qpos and get grasp_site pos in sim
        # env.agent._robot.set_qpos(real_robot_qpos)

        action, _states = RL_model.predict(obs, deterministic=True)
        ## sim step
        obs, rewards, dones, info = env.step(action)

        if render_mode:
            viewer = env.unwrapped.render()
            print("Press [c] to start")
            while True:
                if viewer.window.key_down("c"):
                    break
                env.render()
        if save_mode:
            base_rgbd = env.render("cameras")
            frames.append(base_rgbd)
            sim_qpos_traj.append(env.unwrapped._agent._robot.get_qpos())

        ## get current real state
        imp_opendoor.get_realstate()
        robot_arm_state = imp_opendoor.real_state.q_m
        gripper_state = np.repeat(
            gripper_listener.gPO / 255 * 0.068, 2
        )  # 0-255 => 0-0.068
        real_robot_qpos = np.concatenate((robot_arm_state, gripper_state))
        print("=== real robot current qpos: ", real_robot_qpos)

        ## real step
        real_delta_qpos = scale_para(
            action[:7], action_range_low[:7], action_range_high[:7]
        )
        target_qpos = real_robot_qpos[:-2] + real_delta_qpos
        print("target robot qpos: ", np.asarray(target_qpos))
        for k in range(sim2real_step_num):
            imp_opendoor.exec_trajs(target_qpos, stiffness=joint_kp, damping=joint_kd)
        if mode == "door" or mode == "drawer":
            target_gripper_pos = int(
                (action[-1] + 1) / 2 * 255
            )  # Gripper Open->Close: -1->1 => 0->255
        elif mode == "faucet":
            target_gripper_pos = 255
        else:
            raise NotImplementedError(mode)
        print("target gripper qpos: ", target_gripper_pos)
        imp_opendoor.exec_gripper(target_gripper_pos)
        if mode != "faucet":
            time.sleep(gripper_wait)
        real_gripper_pos = gripper_listener.gPO
        print("real gripper qpos: ", real_gripper_pos)

        ## get current real state
        imp_opendoor.get_realstate()
        robot_arm_state = imp_opendoor.real_state.q_m
        gripper_state = np.repeat(
            gripper_listener.gPO / 255 * 0.068, 2
        )  # 0-255 => 0-0.068
        real_robot_qpos = np.concatenate((robot_arm_state, gripper_state))
        if save_mode:
            real_qpos_traj.append(real_robot_qpos)

        if not success_flag and info["is_success"] == 1.0:
            success_flag = True
            total_steps = step
            break

    if save_mode:
        sim_qpos_traj = np.array(sim_qpos_traj)
        real_qpos_traj = np.array(real_qpos_traj)
        qpos_err = real_qpos_traj - sim_qpos_traj
        labels = np.array([f"j{i}" for i in range(7)])
        plt.subplot(2, 1, 1)
        plt.plot(np.arange(sim_qpos_traj.shape[0]), qpos_err[:, :-2], label=labels)
        plt.title("qpos err (real -sim)")
        plt.legend()
        plt.subplot(2, 1, 2)
        plt.plot(np.arange(sim_qpos_traj.shape[0]), qpos_err[:, -2:])
        plt.title("gripper pos err (real -sim)")
        plt.savefig(
            eval_videos_path / f"{record_id}_{mode}_sim2real_err_{name_suffix}.png",
            dpi=300,
            pad_inches=0.2,
        )
        images_to_video(
            images=frames,
            output_dir=eval_videos_path,
            video_name=f"{record_id}_{mode}_sim2real_{name_suffix}",
            verbose=False,
        )
        np.save(
            eval_videos_path / f"{record_id}_{mode}_sim2real_qpos_{name_suffix}.npy",
            np.concatenate((sim_qpos_traj, real_qpos_traj), axis=0),
        )

arti_mani/test/real_exps/sim2real_replay.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
import time
import warnings
from collections import OrderedDict
from copy import deepcopy

import cv2
import gym
import matplotlib.pyplot as plt
import numpy as np
import open3d as o3d
import sapien.core as sapien
import segmentation_models_pytorch as smp
import torch
import torch.nn.functional as F
import yaml
from arti_mani import ASSET_DIR, RLMODEL_DIR, VISUALMODEL_DIR
from arti_mani.algorithms.rl_iam.Sampling import (
    FCUWSampler,
    FPSSampler,
    RandomSampler,
    ScoreSampler,
    UniDownsampler,
)
from arti_mani.algorithms.visual_net.Networks.Custom_Unet import (
    CustomUnet,
    CustomUnetNew,
    SplitUnet,
)
from arti_mani.envs.arti_mani import ArtiMani
from arti_mani.utils.visualization import images_to_video
from arti_mani.utils.wrappers import NormalizeActionWrapper
from stable_baselines3 import SAC
from tqdm import tqdm


def setup_seed(seed=1029):
    os.environ["PYTHONHASHSEED"] = str(seed)  # 为了禁止hash随机化，使得实验可复现
    np.random.seed(seed)
    rng = np.random.RandomState(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    return rng


def load_vismodel(vismodel_path, load_device):
    model_path = VISUALMODEL_DIR / f"smp_model/{vismodel_path}/best.pth"
    config_path = VISUALMODEL_DIR / f"smp_model/{vismodel_path}/config.yaml"
    with open(config_path, "r") as f:
        cfg = yaml.safe_load(f)
    smp_cfg = cfg["smp_config"]
    if smp_cfg["mode"] == "RGBD":
        in_ch = 4
    elif smp_cfg["mode"] == "RGB":
        in_ch = 3
    elif smp_cfg["mode"] == "D":
        in_ch = 1
    else:
        raise NotImplementedError
    segmodel = CustomUnet(
        has_dropout=False,
        encoder_name=smp_cfg["encoder"],
        encoder_depth=smp_cfg["encoder_depth"],
        decoder_channels=smp_cfg["decoder_channels"],
        encoder_weights=smp_cfg["encoder_weights"],
        in_channels=in_ch,
        classes=cfg["num_classes"],
        activation=smp_cfg["activation"],
    )
    segmodel.load_state_dict(torch.load(model_path))
    segmodel.to(torch.device(load_device))
    segmodel.eval()
    return segmodel


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    eval_seed = np.random.RandomState().randint(2**32)
    rng = setup_seed(eval_seed)
    print("experiment eval random seed: ", eval_seed)
    device = "cuda:0"

    record_id = 5
    mode = "door"  # door, drawer, faucet
    save_mode = True
    render_mode = True
    MAX_STEPS = 30

    smp_exp_name = "20230219_000940_train52-val18_384_noDR_randombg_aug_stereo_bs16_focalloss_0.5step50lr0.001_RGBDunet-163264128_mobilenet_v2"
    vis_model = load_vismodel(smp_exp_name, device)
    segmodel_path = VISUALMODEL_DIR / f"smp_model/{smp_exp_name}"

    if mode == "door":
        # rl_exp_name = "sac4_tr4eval1door_state_egostereo_segpoints_random_samplef1s32_100steps_padfriction1_p3000d100_tr16up16_noeepadpts_newextric_ptsaddnoise0.01_smp0219"
        rl_exp_name = "sac2_tr12eval3arti_state_egostereo_segpoints_random_samplef1s32_100steps_smp0219_largerange_cabinetrew"
    elif mode == "drawer":
        rl_exp_name = "sac2_tr4eval1drawer_state_egostereo_segpoints_random_samplef1s32_100steps_padfriction1_p3000d100_tr16up16_noeepadpts_newextric_smp0219_adjustreward_2close_initrot"
    elif mode == "faucet":
        rl_exp_name = "sac4_tr4eval1faucet_state_egostereo_segpoints_random_samplef1s32_100steps_padfriction1_p3000d100_tr16up16_noeepadpts_newextric_ptsaddnoise0.01_smp0219"
    else:
        raise NotImplementedError
    model_name = "best_model"  # "rl_model_1000000_steps"
    model_path = RLMODEL_DIR / f"{rl_exp_name}/{model_name}"
    eval_videos_path = RLMODEL_DIR / f"{rl_exp_name}/sim_trajs/{model_name}/"
    if not os.path.exists(eval_videos_path):
        os.makedirs(eval_videos_path)

    print("+++++ Build Real2Sim Env +++++++")
    if mode == "door":
        arti_ids = [0]  # 0
    elif mode == "drawer":
        arti_ids = [1]  # 1
    elif mode == "faucet":
        arti_ids = [5052]
    elif mode == "arti":
        arti_ids = [0, 1, 5052]
    else:
        raise NotImplementedError
    obs_mode = "state_egostereo_segpoints"
    control_mode = "pd_joint_delta_pos"
    sample_mode = "random_sample"  # full_downsample, score_sample, random_sample, frameweight_sample, fps_sample
    frame_num = 1
    pts_sample_num = 32
    env: ArtiMani = gym.make(
        "ArtiMani-v0",
        articulation_ids=arti_ids,
        segmodel_path=segmodel_path,
        sample_mode=sample_mode,
        frame_num=frame_num,
        sample_num=pts_sample_num,
        other_handle_visible=False,
        obs_mode=obs_mode,
        control_mode=control_mode,
        reward_mode="dense",  # "sparse", "dense"
        device=device,
    )
    env = NormalizeActionWrapper(env)
    env.seed(eval_seed)
    obs = env.reset()

    ### Set Cabinet Pose
    # env.unwrapped._articulation.set_root_pose(
    #     sapien.Pose([2, 0, 0], [1, 0, 0, 0])
    # )

    ### load RL model
    RL_model = SAC.load(
        model_path,
        env=env,
        print_system_info=True,
    )

    ### Set Robot Pose
    ## real robot init
    if mode == "door" or mode == "drawer":
        qpos = np.array([1.4, -1.053, -2.394, 1.662, 1.217, 1.05, -0.8, 0.0, 0.0])
    elif mode == "faucet":
        qpos = np.array([-0.5, -0.143, 0, np.pi / 3, 0, 1.57, 1.57, 0.068, 0.068])
    else:
        raise NotImplementedError(mode)
    env.unwrapped._agent._robot.set_qpos(qpos)

    if render_mode:
        viewer = env.unwrapped.render()
        print("Press [b] to start")
        while True:
            if viewer.window.key_down("b"):
                break
            env.render()

    # logger.info(f"Observation space {env.observation_space}")
    print(f"Action space {env.action_space}")
    print(f"Control mode {env.control_mode}")
    action_range_low = env.unwrapped.action_space.low
    action_range_high = env.unwrapped.action_space.high
    print("Input Action space range (low):", action_range_low)
    print("Input Action space range (high):", action_range_high)

    success_flag = False
    if save_mode:
        frames = []
        base_rgbd = env.render("cameras")
        frames.append(base_rgbd)
        sim_qpos_action = []
    for step in tqdm(range(MAX_STEPS), colour="green", leave=False):
        action, _states = RL_model.predict(obs, deterministic=True)
        # action[-1] = 1
        if step % 2 == 0:
            action[-1] = -1
        else:
            action[-1] = 1
        print("+++", action)
        ## sim step
        obs, rewards, dones, info = env.step(action)

        if render_mode:
            viewer = env.unwrapped.render()
            print("Press [c] to start")
            while True:
                if viewer.window.key_down("c"):
                    break
                env.render()
        if save_mode:
            base_rgbd = env.render("cameras")
            frames.append(base_rgbd)
            sim_qpos_action.append(
                np.concatenate((env.unwrapped._agent._robot.get_qpos(), action), axis=0)
            )

        if not success_flag and info["is_success"] == 1.0:
            success_flag = True
            total_steps = step
            break

    if save_mode:
        sim_qpos_traj = np.array(sim_qpos_action)
        images_to_video(
            images=frames,
            output_dir=eval_videos_path,
            video_name=f"{record_id}_{mode}_sim2real",
            verbose=False,
        )
        np.save(
            eval_videos_path / f"{record_id}_{mode}_sim2real_qpos_action.npy",
            sim_qpos_traj,
        )

arti_mani/test/real_exps/sim_eval.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import glob
import logging
import os
import time

import cv2
import gym
import matplotlib.pyplot as plt
import segmentation_models_pytorch as smp
from arti_mani import ROOT_DIR, VISUALMODEL_DIR
from arti_mani.envs import *
from arti_mani.utils.cv_utils import visualize_depth
from arti_mani.utils.wrappers import NormalizeActionWrapper
from PIL import Image


def make_gif(frame_path, result_prefix):
    file_list = sorted(
        glob.glob(f"{frame_path}/*.jpg"), key=lambda name: int(name[-7:-4])
    )
    frames = [Image.open(image) for image in file_list]
    frame_one = frames[0]
    frame_one.save(
        f"{frame_path}/{result_prefix}.gif",
        format="GIF",
        append_images=frames,
        save_all=True,
        duration=500,
        loop=0,
    )


if __name__ == "__main__":
    device = torch.device("cuda:0")
    inds = [92]  # list(range(0, 16))
    color_maps = [
        (0, 255, 0),
        (0, 0, 255),
        (255, 255, 0),
        (0, 255, 255),
        (255, 0, 0),
        (255, 255, 255),
    ]
    subplot_num = [3, 3]
    num_classes = 6

    smp_exps = [
        # '20230106_222106_cab41-15_faucet15-3_newaug_bs256_0.8step20lr0.001_unet3-mobilenet_v2_good',
        # '20230108_150002_cab41-15_faucet15-3_newaug_bs256_0.5step50lr0.001_unet3_3264128-mobilenet_v2_good',
        # '20230108_161133_cab41-15_faucet15-3_newaug_bs256_0.5step50lr0.001_unet3_163264-mobilenet_v2_good',
        # '20230109_133653_cab41-15_faucet11-3_newaug_bs256_0.5step50lr0.001_unet3-163264_mobilenet_v2_good',
        #
        # '20230110_021737_cab39-17_faucet11-3_newaug_bs256_0.5step50lr0.001_RGBDunet3-163264_mobilenet_v2',
        # '20230110_103746_cab39-17_faucet11-3_newaug_bs256_0.5step50lr0.001_RGBDunet3-3264128_mobilenet_v2',
        # '20230110_030002_cab39-17_faucet11-3_newaug_bs256_0.5step50lr0.001_RGBunet3-163264_mobilenet_v2',
        # '20230110_110619_cab39-17_faucet11-3_newaug_bs256_0.5step50lr0.001_RGBunet3-3264128_mobilenet_v2',
        # '20230110_021934_cab39-17_faucet11-3_newaug_bs256_0.5step50lr0.001_Dunet3-163264_mobilenet_v2',
        # '20230110_133947_cab39-17_faucet11-3_newaug_bs256_0.5step50lr0.001_Dunet3-3264128_mobilenet_v2',
        "20230111_005201_cab39-17_faucet11-3_stereo_bs256_0.5step50lr0.001_RGBDunet3-163264_mobilenet_v2",
        # '20230111_014201_cab39-17_faucet11-3_stereo_bs256_0.5step50lr0.001_RGBDunet3-3264128_mobilenet_v2',
        # '20230111_130024_cab39-17_faucet11-3_stereo_bs256_0.5step50lr0.001_RGBunet3-163264_mobilenet_v2',
        # '20230111_100917_cab39-17_faucet11-3_stereo_bs256_0.5step50lr0.001_RGBunet3-3264128_mobilenet_v2',
        # '20230111_005327_cab39-17_faucet11-3_stereo_bs256_0.5step50lr0.001_Dunet3-163264_mobilenet_v2',
        # '20230111_012400_cab39-17_faucet11-3_stereo_bs256_0.5step50lr0.001_Dunet3-3264128_mobilenet_v2',
    ]

    logfile_path = VISUALMODEL_DIR / f"smp_model/"
    logfile_name = (
        f"{time.strftime('%Y%m%d_%H%M%S', time.localtime())}_real_smp_evalresults"
    )
    logger = logging.getLogger(logfile_name)
    logger.setLevel(logging.DEBUG)
    # set two handlers
    log_file = "{}.log".format(logfile_name)
    # rm_file(log_file)
    fileHandler = logging.FileHandler(os.path.join(logfile_path, log_file), mode="w")
    fileHandler.setLevel(logging.DEBUG)
    # consoleHandler = logging.StreamHandler()
    # consoleHandler.setLevel(logging.DEBUG)
    # set formatter
    formatter = logging.Formatter(
        "[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    # consoleHandler.setFormatter(formatter)
    fileHandler.setFormatter(formatter)
    # add
    logger.addHandler(fileHandler)
    # logger.addHandler(consoleHandler)

    # "state", "state_dict", "rgbd", "pointcloud", "state_depth", "state_egorgbd", "state_egopcrgb", "state_crossview"
    obs_mode = "state_egorgbd"
    # "pd_joint_delta_pos", "pd_ee_delta_pose", "pd_ee_delta_pos"
    control_mode = "pd_joint_delta_pos"
    env_id = "OpenCabinetDoor-v0"
    arti_ids = [0]
    arti_config_path = (
        ASSET_DIR / "partnet_mobility_configs/fixed_cabinet_doors_new.yml"
    )

    env = gym.make(
        env_id,
        articulation_ids=arti_ids,
        articulation_config_path=arti_config_path,
        obs_mode=obs_mode,
        control_mode=control_mode,
        reward_mode="dense",
    )
    env = NormalizeActionWrapper(env)
    eval_seed = np.random.RandomState().randint(2**32)
    logger.info(f"eval env random seed: {eval_seed}")
    env.seed(eval_seed)
    imgs = env.unwrapped._agent.get_images()
    hand_camera_extrinsic_base_frame = np.asarray(
        imgs["hand_camera"]["camera_extrinsic_base_frame"]
    )
    logger.info(f"hand_camera_extrinsic_base_frame: {hand_camera_extrinsic_base_frame}")

    for smp_exp_name in smp_exps:
        logger.info(f"{smp_exp_name} real visual results.")
        smp_model_path = VISUALMODEL_DIR / f"smp_model/{smp_exp_name}/best.pth"
        config_path = VISUALMODEL_DIR / f"smp_model/{smp_exp_name}/config.yaml"
        with open(config_path, "r") as f:
            cfg = yaml.safe_load(f)
        smp_cfg = cfg["smp_config"]
        if smp_cfg["mode"] == "RGBD":
            in_ch = 4
        elif smp_cfg["mode"] == "RGB":
            in_ch = 3
        elif smp_cfg["mode"] == "D":
            in_ch = 1
        else:
            raise NotImplementedError
        unet_model = smp.Unet(
            encoder_name=smp_cfg["encoder"],
            encoder_depth=smp_cfg["encoder_depth"],
            decoder_channels=smp_cfg["decoder_channels"],
            encoder_weights=smp_cfg["encoder_weights"],
            in_channels=in_ch,
            classes=cfg["num_classes"],
            activation=smp_cfg["activation"],
        )
        unet_model.load_state_dict(torch.load(smp_model_path))
        unet_model.to(device)
        unet_model.eval()
        model_params = sum(p.numel() for p in unet_model.parameters())

        eval_result_path = (
            VISUALMODEL_DIR / f"smp_model/{smp_exp_name}/real_results_final/"
        )
        if not os.path.exists(eval_result_path):
            os.makedirs(eval_result_path)

        logger.info(f"model params: {model_params/1e6:.2f} M")

        for ind in inds:
            real_data_path = ROOT_DIR / f"test/real_exps/real_images/"
            # cur_rgb = np.array(Image.open(real_data_path / f"{ind:02}_rgb.png"))
            # cur_depth = np.array(Image.open(real_data_path / f"{ind:02}_depth.png"))
            cur_rgb = np.load(real_data_path / f"{ind:02}_rgb.npy")[..., ::-1]
            cur_depth = np.load(real_data_path / f"{ind:02}_depth.npy")
            logger.info(f"real rgb, depth shape: {cur_rgb.shape}, {cur_depth.shape}")
            logger.info(f"real rgb min, max: {np.min(cur_rgb)}, {np.max(cur_rgb)}")
            logger.info(
                f"real depth min, max: {np.min(cur_depth)}, {np.max(cur_depth)}"
            )
            with torch.no_grad():
                # proc_rgb = cv2.resize(cur_rgb, (128, 72)) / 255.0
                # proc_depth = cv2.resize(cur_depth, (128, 72))
                proc_rgb = cur_rgb / 255.0
                proc_depth = cur_depth

                real_rgb = (
                    torch.from_numpy(proc_rgb).float().to(device).permute((2, 0, 1))
                )  # (3, H, W)
                real_depth = (
                    torch.from_numpy(proc_depth[None]).float().to(device)
                )  # (1, H, W)

                if smp_cfg["mode"] == "RGBD":
                    img_input = torch.cat([real_rgb, real_depth], dim=0)[None]
                elif smp_cfg["mode"] == "RGB":
                    img_input = real_rgb[None]
                elif smp_cfg["mode"] == "D":
                    img_input = real_depth[None]
                else:
                    raise NotImplementedError

                time0 = time.time()
                attnmap = unet_model.predict(img_input)  # (1, 6, H, W)
                time1 = time.time()
                infer_times = (time1 - time0) * 1000

                plot_rgb = real_rgb.detach().cpu().numpy().transpose((1, 2, 0))
                plot_depth = real_depth.detach().squeeze().cpu().numpy()
                attnmap = attnmap.squeeze().cpu().numpy()

                plt.subplot(subplot_num[0], subplot_num[1], 1)
                plt.imshow(plot_rgb)
                plt.subplot(subplot_num[0], subplot_num[1], 2)
                plt.imshow(plot_depth)
                for pid in range(num_classes):
                    plt.subplot(subplot_num[0], subplot_num[1], pid + 3)
                    plt.imshow(attnmap[pid])

                plt.savefig(f"{eval_result_path}/test_{ind:02}.jpg", dpi=200)
                plt.close()

                depth_colormap = visualize_depth(cur_depth)
                cv2.imwrite(
                    f"{eval_result_path}/{ind:03}_depth_colormap.png", depth_colormap
                )

                logger.info(f"inference time: {infer_times:.2f} ms")

arti_mani/test/real_exps/smp_eval.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import numpy as np
import sapien.core as sapien
import transforms3d


def normalize_image(image):
    v_max = np.max(image)
    v_min = np.min(image)
    return (image - v_min) / (v_max - v_min)


def mat2pose(mat: np.ndarray):
    quat = transforms3d.quaternions.mat2quat(mat[:3, :3])
    pos = mat[:3, 3]
    # pose = np.concatenate((pos, quat))
    pose = sapien.Pose(pos, quat)
    return pose


def trans_axangle2pose(trans_axangle):
    RT = np.eye(4)
    RT[:3, 3] = trans_axangle[:3]
    angle = np.linalg.norm(trans_axangle[3:6])
    if angle < 1e-6:
        axis = (0, 0, 1)
    else:
        axis = trans_axangle[3:6] / angle
    RT[:3, :3] = transforms3d.axangles.axangle2mat(axis, angle)
    return mat2pose(RT)

arti_mani/test/real_exps/utils.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import numpy as np
import rospy
import sapien.core as sapien
import transforms3d
from impedance_control.msg import (
    CartesianControlCommand,
    ImpedanceRobotState,
    JointControlCommand,
)
from robotiq_2f_gripper_control.msg import (
    Robotiq2FGripper_robot_input,
    Robotiq2FGripper_robot_output,
)


def mat2pose(mat: np.ndarray):
    quat = transforms3d.quaternions.mat2quat(mat[:3, :3])
    pos = mat[:3, 3]
    # pose = np.concatenate((pos, quat))
    pose = sapien.Pose(pos, quat)
    return pose


def trans_axangle2pose(trans_axangle):
    RT = np.eye(4)
    RT[:3, 3] = trans_axangle[:3]
    angle = np.linalg.norm(trans_axangle[3:6])
    if angle < 1e-6:
        axis = (0, 0, 1)
    else:
        axis = trans_axangle[3:6] / angle
    RT[:3, :3] = transforms3d.axangles.axangle2mat(axis, angle)
    return mat2pose(RT)


class Robotiq2FGripper_Listener:
    def __init__(self) -> None:
        self.gripper_is_ready = True
        self.gPO = 0

    def listen(self, input):
        # print("input.gACT: ", input.gACT)
        # print("input.gOBJ: ", input.gOBJ)
        if input.gACT == 1 and input.gOBJ > 0:
            self.gripper_is_ready = True
        else:
            self.gripper_is_ready = False
        self.gPO = input.gPO


class ROS_ImpOpendoor:
    def __init__(self, gripper_listener, control_mode="pd_joint_delta_pos") -> None:
        rospy.init_node("open_cabinet_node")
        self.gripper_listener = gripper_listener
        self.rate = rospy.Rate(10)  # 10hz
        self.control_mode = control_mode
        self._initialize_robotcontrol()
        self._initialize_gripper()

    def _initialize_robotcontrol(self):
        if self.control_mode == "pd_joint_delta_pos":
            self.command_pub = rospy.Publisher(
                "joint_control_command", JointControlCommand, queue_size=1
            )
        elif self.control_mode == "pd_ee_pose":
            self.command_pub = rospy.Publisher(
                "cartesian_control_command", CartesianControlCommand, queue_size=1
            )
        else:
            raise NotImplementedError

    def _initialize_gripper(self):
        rospy.Subscriber(
            "Robotiq2FGripperRobotInput",
            Robotiq2FGripper_robot_input,
            self.gripper_listener.listen,
        )

        self.gripper_input = Robotiq2FGripper_robot_input()
        self.gripper_output = Robotiq2FGripper_robot_output()
        self.gripper_pub = rospy.Publisher(
            "Robotiq2FGripperRobotOutput", Robotiq2FGripper_robot_output
        )
        self.gripper_output.rACT = 1  # 1: Active, 0: Not active
        self.gripper_output.rGTO = 1
        self.gripper_output.rATR = 0
        self.gripper_output.rPR = 150  # 0~255: Placement
        self.gripper_output.rFR = 50  # 0~255: Force
        self.gripper_output.rSP = 150  # 0~255: Speed

        while self.gripper_listener.gripper_is_ready:
            self.gripper_pub.publish(self.gripper_output)
        self.gripper_output.rPR = 0  # Open Gripper
        while not self.gripper_listener.gripper_is_ready:
            pass
        while self.gripper_listener.gripper_is_ready:
            self.gripper_pub.publish(self.gripper_output)

    def get_realstate(self):
        self.real_state: ImpedanceRobotState = rospy.wait_for_message(
            "impedance_robot_state", ImpedanceRobotState, timeout=10000
        )
        if self.real_state == None:
            print("no real robot state received in 10000s")
            return

    def exec_gripper(self, gripper_norm_rPR):
        self.gripper_output.rPR = gripper_norm_rPR
        while not self.gripper_listener.gripper_is_ready:
            pass
        # while self.gripper_listener.gripper_is_ready: self.gripper_pub.publish(self.gripper_output)
        self.gripper_pub.publish(self.gripper_output)
        while not self.gripper_listener.gripper_is_ready:
            pass

    def exec_trajs(self, target, stiffness=None, damping=None):
        """
        target: (joint_pos): qpos[7], (ee_pose): mat(16)
        stiffness: (joint_): 7, (ee_): 6
        damping: (joint_): 7, (ee_): 6
        """
        self.get_realstate()
        assert (
            stiffness is not None and damping is not None
        ), "joint/ee stiffness and damping must be set!"
        if self.control_mode == "pd_joint_delta_pos":
            msg = JointControlCommand(target, stiffness, damping, False)
            # print("++++ get robot arm control msg ++++")
        elif self.control_mode == "pd_ee_pose":
            msg = CartesianControlCommand(target, stiffness, damping, False)
        else:
            raise NotImplementedError
        self.command_pub.publish(msg)
        # print("++++ robot arm executation finished ++++")
        self.rate.sleep()


if __name__ == "__main__":
    np.set_printoptions(suppress=True, precision=4)
    control_mode = "pd_ee_pose"  # "pd_joint_delta_pos",

    if control_mode == "pd_joint_delta_pos":
        kp = [200.0, 200.0, 200.0, 200.0, 50.0, 50.0, 50.0]
        kd = [20.0, 20.0, 20.0, 20.0, 10.0, 10.0, 10.0]
    elif control_mode == "pd_ee_pose":
        kp = [300.0, 300.0, 300.0, 50.0, 50.0, 50.0]
        kd = [20.0, 20.0, 20.0, 10.0, 10.0, 10.0]
    else:
        raise NotImplementedError
    print("current kp: ", kp)
    print("current kd: ", kd)

    print("+++++ Setup robotiq controller +++++++")
    gripper_listener = Robotiq2FGripper_Listener()
    print("+++++ Setup robot arm controller +++++++")
    imp_opendoor = ROS_ImpOpendoor(
        gripper_listener=gripper_listener, control_mode=control_mode
    )

    imp_opendoor.get_realstate()
    ## robot qpos
    robot_arm_state = imp_opendoor.real_state.q_m
    gripper_state = np.repeat(gripper_listener.gPO / 255 * 0.068, 2)  # 0-255 => 0-0.068
    real_robot_qpos = np.concatenate((robot_arm_state, gripper_state))
    print("=== real robot current qpos: ", real_robot_qpos)
    ## robot ee pose (robot frame)
    real_ee_pose_mat = imp_opendoor.real_state.toolTobase_pos_m
    real_ee_pose_base = mat2pose(np.reshape(real_ee_pose_mat, (4, 4)))
    print("=== real robot ee pose: ", real_ee_pose_base)
    ## target ee pose
    target_link7_pose = sapien.Pose(
        np.array([0.6, 0, 0.4]), np.array([0, 0.707, 0.707, 0])
    )
    ### link7 pose: Pose([0, 0, 0.431991], [0, -0.707107, 0.707107, 0])
    ### handcolorframe pose: Pose([-0.0773266, -0.0402823, 0.419322], [-0.00315852, 0.704288, 0.00583035, -0.709884])
    target_ee_mat = np.reshape(target_link7_pose.to_transformation_matrix(), 16)
    print("=== target robot ee pose: ", target_link7_pose)
    print("=== target robot ee mat: ", target_ee_mat)
    for k in range(10):
        imp_opendoor.exec_trajs(target_ee_mat, stiffness=kp, damping=kd)
    print("Done!")

    ## camera pos (robot frame): (0.685, 0.025, 0.393~0.4 - 0.03 = 0.363~0.37)
    ## xmate7_link pos (robot frame): (0.6, 0, 0.4)
    ## diff: (0.085, 0.025, -0.037~-0.03)
    # <joint name="realsense_hand_joint" type="fixed">
    #     <origin rpy="2.12574149 -1.55582192  2.58286879" xyz="0.04028231  0.07736792  0.01266915"/>
    #     <!--  <origin rpy="3.9290746 ,   -1.6206431 , -124.82640542" xyz="-0.0188 0.0531 -0.0298"/>  11 samples from d435 -->
    #     <parent link="xmate3_link7"/>
    #     <child link="camera_hand_color_frame"/>
    # </joint>

arti_mani/test/real_exps/xmate3Robotiq.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import numpy as np
import open3d as o3d
import sapien.core as sapien
import trimesh
from arti_mani.utils.geometry import all_close, apply_pose, check_coplanar, convex_hull
from arti_mani.utils.o3d_utils import merge_mesh, np2mesh
from arti_mani.utils.string_utils import regex_match
from sapien.core import ActorBase, Articulation, CollisionShape, Pose
from transforms3d.euler import euler2mat


def get_actor_names(actors):
    single = False
    if isinstance(actors, ActorBase):
        actors = [
            actors,
        ]
        single = True
    elif isinstance(actors, Articulation):
        actors = actors.get_links()
    ret = [actor.get_name() for actor in actors]
    if single:
        ret = ret[0]
    return ret


def get_actors_by_names(actors, names):
    assert isinstance(actors, (list, tuple))
    # Actors can be joint and link
    if isinstance(names, str):
        names = [names]
        sign = True
    else:
        sign = False
    ret = [None for _ in names]
    for actor in actors:
        if actor.get_name() in names:
            ret[names.index(actor.get_name())] = actor
    return ret[0] if sign else ret


def collision_shape_to_o3d(
    collision_shape: CollisionShape,
) -> o3d.geometry.TriangleMesh:
    if collision_shape.type == "capsule":
        capsule = collision_shape.geometry
        mesh = trimesh.creation.capsule(
            capsule.half_length * 2, capsule.radius, count=[4, 4]
        )
        vertices = mesh.vertices
        indices = mesh.faces
        to_x = euler2mat(0, np.pi / 2, 0)
        vertices = apply_pose(to_x, vertices)
    elif collision_shape.type == "convex_mesh":
        convex = collision_shape.geometry
        vertices = convex.vertices * convex.scale
        vertices = apply_pose(Pose(p=[0, 0, 0], q=convex.rotation), vertices)
        indices = convex.indices.reshape(-1, 3)
    elif collision_shape.type == "sphere":
        sphere = collision_shape.geometry
        mesh = trimesh.creation.icosphere(radius=sphere.radius, subdivisions=1)
        vertices = mesh.vertices
        indices = mesh.faces
    elif collision_shape.type == "box":
        box = collision_shape.geometry
        mesh = trimesh.creation.box(extents=box.half_lengths * 2)
        vertices = mesh.vertices
        indices = mesh.faces
    else:
        print(collision_shape.type, "????")
        raise NotImplementedError("")
    vertices = apply_pose(collision_shape.get_local_pose(), vertices)
    vertices = o3d.utility.Vector3dVector(vertices)
    indices = o3d.utility.Vector3iVector(indices)
    return o3d.geometry.TriangleMesh(vertices=vertices, triangles=indices)


cached_o3d = {}


def actor_to_o3d(
    actor,
    convex=False,
    collision=False,
    pattern_name="(.*?)",
    actor_exclude=[],
    part_exclude=[],
    actor_name="",
    use_cache=True,
):
    if actor.get_name() in actor_exclude:
        return None
    ret = []
    is_actor = False
    actor_name_curr = actor_name + f"/{actor.get_name()}"
    if isinstance(actor, Articulation):
        for link in actor.get_links():
            tmp_o3d_mesh = actor_to_o3d(
                link,
                convex,
                collision,
                pattern_name,
                actor_exclude,
                part_exclude,
                actor_name=actor_name_curr,
                use_cache=use_cache,
            )
            if tmp_o3d_mesh:
                ret.append(tmp_o3d_mesh)
    elif isinstance(actor, ActorBase):
        is_actor = True
        if not collision:
            actor_name_curr = actor_name_curr + "_visual"
            print(list(cached_o3d.keys()))
            if actor_name_curr in cached_o3d and use_cache:
                ori_mesh, pose, trans_mesh = cached_o3d[actor_name_curr]
                print("Hit cache", use_cache, actor_name, actor.get_pose(), pose)
                if trans_mesh is None or all_close(pose, actor.get_pose()):
                    return trans_mesh
                ret = ori_mesh
            else:
                for visual_body in actor.get_visual_bodies():
                    if not regex_match(visual_body.get_name(), pattern_name):
                        continue
                    if visual_body.get_name() in part_exclude:
                        continue
                    if visual_body.type == "mesh":
                        scale = visual_body.scale
                    elif visual_body.type == "box":
                        scale = visual_body.half_lengths
                    elif visual_body.type == "sphere":
                        scale = visual_body.radius
                    else:
                        scale = 1.0
                    for i in visual_body.get_render_shapes():
                        if i.mesh.indices.reshape(-1, 3).shape[0] < 4 or check_coplanar(
                            i.mesh.vertices * scale
                        ):
                            continue
                        vertices = apply_pose(
                            visual_body.local_pose, i.mesh.vertices * scale
                        )
                        mesh = np2mesh(vertices, i.mesh.indices.reshape(-1, 3))
                        if convex:
                            mesh = convex_hull(mesh)
                        ret.append(mesh)
        else:
            actor_name_curr = actor_name_curr + "_collision"
            for collision_shape in actor.get_collision_shapes():
                mesh = collision_shape_to_o3d(collision_shape)
                ret.append(mesh)
    else:
        print(type(actor))
        raise NotImplementedError()
    # print(ret)
    if not ret:
        # raise ValueError(f"No entity in {actor}")
        return None
    if isinstance(ret, list):
        ret = merge_mesh(ret)
    if is_actor:
        num_faces = 1024
        if ret is not None and np.asarray(ret.triangles).shape[0] > num_faces:
            ret = ret.simplify_vertex_clustering(
                voxel_size=0.02,
                contraction=o3d.geometry.SimplificationContraction.Average,
            )
        trans_mesh = apply_pose(actor.get_pose(), ret)
        if actor_name_curr not in cached_o3d and use_cache:
            cached_o3d[actor_name_curr] = [ret, actor.get_pose(), trans_mesh]
        ret = trans_mesh
    return ret


def actor_to_o3d_mesh(
    actor: ActorBase,
    collision: bool = False,
    convex: bool = False,
    max_num_faces: int = 1024,
) -> o3d.geometry.TriangleMesh:
    assert isinstance(actor, ActorBase), type(actor)
    shape_list = []
    if collision:
        for collision_shape in actor.get_collision_shapes():
            mesh = collision_shape_to_o3d(collision_shape)
            shape_list.append(mesh)
    else:
        for visual_body in actor.get_visual_bodies():
            if visual_body.type == "mesh":
                scale = visual_body.scale
            elif visual_body.type == "box":
                scale = visual_body.half_lengths
            elif visual_body.type == "sphere":
                scale = visual_body.radius
            else:
                scale = 1.0
            for render_shape in visual_body.get_render_shapes():
                if render_shape.mesh.indices.reshape(-1, 3).shape[
                    0
                ] < 4 or check_coplanar(render_shape.mesh.vertices * scale):
                    continue
                vertices = apply_pose(
                    visual_body.local_pose, render_shape.mesh.vertices * scale
                )
                mesh = np2mesh(vertices, render_shape.mesh.indices.reshape(-1, 3))
                if convex:
                    mesh = convex_hull(mesh)
                shape_list.append(mesh)
    if not shape_list:
        return None
    mesh = merge_mesh(shape_list)
    if np.asarray(mesh.triangles).shape[0] > max_num_faces:
        mesh = mesh.simplify_vertex_clustering(
            voxel_size=0.02, contraction=o3d.geometry.SimplificationContraction.Average
        )
    return mesh


def project_qpos_to_valid_range(robot, qpos):
    for i, joint in enumerate(robot.get_active_joints()):
        l, r = joint.get_limits()[0]
        if l <= qpos[i] <= r:
            continue
        if joint.type == "revolute":
            print(qpos[i], l, r, qpos[i] - l)

        elif joint.type == "prismatic":
            return None

    exit(0)


def get_entity_by_name(entities, name: str, is_unique=True):
    """Get a Sapien.Entity given the name.

    Args:
        entities (List[sapien.Entity]): entities (link, joint, ...) to query.
        name (str): name for query.
        is_unique (bool, optional):
            whether the name should be unique. Defaults to True.

    Raises:
        RuntimeError: The name is not unique when @is_unique is True.

    Returns:
        sapien.Entity or List[sapien.Entity]:
            matched entity or entities. None if no matches.
    """
    matched_entities = [x for x in entities if x.get_name() == name]
    if len(matched_entities) > 1:
        if not is_unique:
            return matched_entities
        else:
            raise RuntimeError(f"Multiple entities with the same name {name}.")
    elif len(matched_entities) == 1:
        return matched_entities[0]
    else:
        return None


# -------------------------------------------------------------------------- #
# Entity state
# -------------------------------------------------------------------------- #
def get_actor_state(actor: sapien.Actor):
    pose = actor.get_pose()
    if actor.type == "static":
        vel = np.zeros(3)
        ang_vel = np.zeros(3)
    else:
        vel = actor.get_velocity()  # [3]
        ang_vel = actor.get_angular_velocity()  # [3]
    return np.hstack([pose.p, pose.q, vel, ang_vel])  # 3+4+3+3=13


def set_actor_state(actor: sapien.Actor, state: np.ndarray):
    assert len(state) == 13, len(state)
    actor.set_pose(Pose(state[0:3], state[3:7]))
    if actor.type != "static":
        actor.set_velocity(state[7:10])
        actor.set_angular_velocity(state[10:13])


def get_articulation_state(articulation: sapien.Articulation):
    root_link = articulation.get_links()[0]
    pose = root_link.get_pose()
    vel = root_link.get_velocity()  # [3]
    ang_vel = root_link.get_angular_velocity()  # [3]
    qpos = articulation.get_qpos()
    qvel = articulation.get_qvel()
    return np.hstack(
        [pose.p, pose.q, vel, ang_vel, qpos, qvel]
    )  # 7+2 robot: 3+4+3+3+9+9=31


def set_articulation_state(articulation: sapien.Articulation, state: np.ndarray):
    articulation.set_root_pose(Pose(state[0:3], state[3:7]))
    articulation.set_root_velocity(state[7:10])
    articulation.set_root_angular_velocity(state[10:13])
    qpos, qvel = np.split(state[13:], 2)
    articulation.set_qpos(qpos)
    articulation.set_qvel(qvel)


def get_pad_articulation_state(
    articulation: sapien.Articulation, max_dof: int
) -> np.ndarray:
    root_link = articulation.get_links()[0]
    pose = root_link.get_pose()
    base_pos, base_quat = pose.p, pose.q
    base_vel = root_link.get_velocity()  # [3]
    base_ang_vel = root_link.get_angular_velocity()  # [3]
    qpos = articulation.get_qpos()
    qvel = articulation.get_qvel()
    k = len(qpos)
    # pad_obj_internal_state = np.zeros(max_dof)
    # pad_obj_internal_state[:k] = qpos
    pad_obj_internal_state = np.zeros(2 * max_dof)
    pad_obj_internal_state[:k] = qpos
    pad_obj_internal_state[max_dof : max_dof + k] = qvel
    return np.concatenate(
        [base_pos, base_quat, base_vel, base_ang_vel, pad_obj_internal_state]
    )


def get_actor_AABB(actor):
    mins = np.array([np.inf, np.inf, np.inf])
    maxs = -mins
    if isinstance(actor, Articulation):
        for link in actor.get_links():
            mins_i, maxs_i = get_actor_AABB(link)
            mins = np.minimum(mins, mins_i)
            maxs = np.maximum(maxs, maxs_i)
    else:
        assert isinstance(actor, ActorBase)
        actor_pose = actor.get_pose()
        for s in actor.get_collision_shapes():
            p = actor_pose * s.get_local_pose()
            T = p.to_transformation_matrix()
            vertices = s.geometry.vertices * s.geometry.scale
            vertices = vertices @ T[:3, :3].T + T[:3, 3]
            mins = np.minimum(mins, vertices.min(0))
            maxs = np.maximum(maxs, vertices.max(0))
    return mins, maxs


def ignore_collision(articulation: sapien.Articulation):
    """ignore collision among all movable links for acceleration"""
    for joint, link in zip(articulation.get_joints(), articulation.get_links()):
        if joint.type in ["revolute", "prismatic"]:
            shapes = link.get_collision_shapes()
            for s in shapes:
                g0, g1, g2, g3 = s.get_collision_groups()
                s.set_collision_groups(g0, g1, g2 | 1 << 31, g3)

arti_mani/utils/articulation_utils.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import gzip
import json
import warnings
from collections import OrderedDict
from pathlib import Path
from typing import Dict, List, Sequence, Union

import gym
import numpy as np
from arti_mani import ASSET_DIR
from gym import spaces
from gym.utils.colorize import colorize


# -------------------------------------------------------------------------- #
# Basic
# -------------------------------------------------------------------------- #
def merge_dicts(ds: List[Dict], asarray=False):
    ret = {k: [d[k] for d in ds] for k in ds[0].keys()}
    if asarray:
        ret = {k: np.concatenate(v) for k, v in ret.items()}
    return ret


def warn(message, stacklevel=2):
    warnings.warn(colorize(message, "yellow"), stacklevel=stacklevel)


# -------------------------------------------------------------------------- #
# Numpy
# -------------------------------------------------------------------------- #
def normalize_vector(x, eps=1e-6):
    x = np.asarray(x)
    assert x.ndim == 1, x.ndim
    norm = np.linalg.norm(x)
    return np.zeros_like(x) if norm < eps else (x / norm)


def compute_angle_between(x1, x2):
    """Compute angle (radian) between two vectors."""
    x1, x2 = normalize_vector(x1), normalize_vector(x2)
    dot_prod = np.clip(np.dot(x1, x2), -1, 1)
    return np.arccos(dot_prod).item()


class np_random:
    """Context manager for numpy random state"""

    def __init__(self, seed):
        self.seed = seed
        self.state = None

    def __enter__(self):
        self.state = np.random.get_state()
        np.random.seed(self.seed)
        return self.state

    def __exit__(self, exc_type, exc_val, exc_tb):
        np.random.set_state(self.state)


def random_choice(x: Sequence, rng: np.random.RandomState = np.random):
    assert len(x) > 0
    if len(x) == 1:
        return x[0]
    else:
        return x[rng.choice(len(x))]


# ---------------------------------------------------------------------------- #
# OpenAI gym
# ---------------------------------------------------------------------------- #
def register_gym_env(name, **kwargs):
    """A decorator to register gym environments.

    Args:
        name (str): a unique id to register in gym.
    """

    def _register(cls):
        entry_point = "{}:{}".format(cls.__module__, cls.__name__)
        gym.register(name, entry_point=entry_point, **kwargs)
        return cls

    return _register


def register_gym_env_for_xpw(name, **kwargs):
    def _register(cls):
        entry_point = "{}:{}".format(cls.__module__, cls.__name__)
        for obs_mode in ["state", "rgbd", "depth"]:
            tag = obs_mode
            gym.register(
                id=name.replace("@", tag),
                entry_point=entry_point,
                kwargs={
                    "obs_mode": obs_mode,
                    "reward_mode": "dense",
                    "articulation_config_path": ASSET_DIR
                    / "partnet_mobility_configs/1999-cabinet-door-real-1.yml",
                },
                **kwargs,
            )
        return cls

    return _register


def convert_observation_to_space(observation):
    """Convert observation to OpenAI gym observation space (recursively).
    Modified from gym.envs.mujoco_env
    """
    if isinstance(observation, (dict)):
        if not isinstance(observation, OrderedDict):
            warn(
                "observation is not an OrderedDict. Keys are {}".format(
                    observation.keys()
                )
            )
        space = spaces.Dict(
            OrderedDict(
                [
                    (key, convert_observation_to_space(value))
                    for key, value in observation.items()
                ]
            )
        )
    elif isinstance(observation, np.ndarray):
        if observation.dtype == np.uint8:
            low = np.zeros(observation.shape, dtype=np.uint8)
            high = np.ones(observation.shape, dtype=np.uint8) * 255
        else:
            low = np.full(observation.shape, -float("inf"))
            high = np.full(observation.shape, float("inf"))
        space = spaces.Box(low, high, dtype=observation.dtype)
    else:
        raise NotImplementedError(type(observation), observation)

    return space


def normalize_action_space(action_space: spaces.Box):
    assert isinstance(action_space, spaces.Box), type(action_space)
    return spaces.Box(-1, 1, shape=action_space.shape, dtype=action_space.dtype)


def clip_and_scale_action(action, low, high, reverse=False):
    """Clip action to [-1, 1] and scale according to a range [low, high]."""
    low, high = np.asarray(low), np.asarray(high)
    action = np.clip(action, -1, 1)
    return 0.5 * (high + low) + 0.5 * (high - low) * action


def normalize_action(action, low, high):
    """Normalize and clip action to [-1, 1] according to a range [low, high]."""
    low, high = np.asarray(low), np.asarray(high)
    action = 2.0 * (action - low) / (high - low) - 1.0
    return np.clip(action, -1.0, 1.0)


def flatten_state_dict(state_dict: OrderedDict):
    """Flatten an ordered dict containing states recursively.

    Args:
        state_dict (OrderedDict): an ordered dict containing states.

    Raises:
        TypeError: If @state_dict is not an OrderedDict.
        TypeError: If a value of @state_dict is a dict instead of OrderedDict.
        AssertionError: If a value of @state_dict is an ndarray with ndim > 1.

    Returns:
        np.ndarray: flattened states.
    """
    if not isinstance(state_dict, OrderedDict):
        raise TypeError(
            "Must be an OrderedDict, but received {}".format(type(state_dict))
        )
    if len(state_dict) == 0:
        return np.empty(0)
    states = []
    for key, value in state_dict.items():
        if isinstance(value, OrderedDict):
            states.append(flatten_state_dict(value))
        elif isinstance(value, dict):
            raise TypeError(
                "Must be an OrderedDict, but received dict for {}".format(key)
            )
        elif isinstance(value, (int, float, tuple, list)):
            states.append(value)
        elif isinstance(value, np.ndarray):
            assert value.ndim <= 1, "Too many dimensions({}) for {}".format(
                value.ndim, key
            )
            states.append(value)
        elif isinstance(value, np.bool_):
            # x = np.array(1) > 0 is np.bool_ instead of ndarray
            states.append(value.astype(int))
        else:
            raise TypeError("Unsupported type: {}".format(type(value)))
    return np.hstack(states)


def convert_np_bool_to_float(x) -> np.ndarray:
    return np.array(x).astype(float)


def load_json(filename: Union[str, Path]):
    filename = str(filename)
    if filename.endswith(".gz"):
        f = gzip.open(filename, "rt")
    elif filename.endswith(".json"):
        f = open(filename, "rt")
    else:
        raise RuntimeError(f"Unsupported extension: {filename}")
    ret = json.loads(f.read())
    f.close()
    return ret

arti_mani/utils/common.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from copy import deepcopy
from math import atan2, pi

import numpy as np
import open3d as o3d
import trimesh
from sapien.core import Pose
from shapely.geometry import Polygon
from transforms3d.quaternions import quat2axangle


def skew(x):
    return np.array([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]])


def norm(x, keepdims=False):
    return np.sqrt((x**2).sum(axis=-1, keepdims=keepdims))


def angle_distance(q0, q1):
    qd = (q0.inv() * q1).q
    return 2 * atan2(norm(qd[1:]), qd[0]) / pi


def angle_distance_ms(q0, q1):
    theta = 2 * np.arccos(np.clip(np.abs(q0.q @ q1.q), a_min=0, a_max=1))
    return theta


def compute_relative_vel(frame_pose, frame_vel, frame_ang_vel, p_world, p_world_vel):
    p_frame = frame_pose.inv().transform(Pose(p_world)).p
    H = frame_pose.to_transformation_matrix()
    R = H[:3, :3]
    o = H[:3, 3]
    S = skew(frame_ang_vel)
    return S @ (R @ p_frame) + frame_vel - p_world_vel


def get_unit_box_corners():
    corners = np.array(
        [
            [0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [0.0, 0.0, 1.0],
            [1.0, 1.0, 1.0],
            [0.0, 1.0, 1.0],
            [1.0, 0.0, 1.0],
            [1.0, 1.0, 0.0],
        ]
    )
    return corners - [0.5, 0.5, 0.5]


def to_generalized(x):
    if x.shape[-1] == 4:
        return x
    assert x.shape[-1] == 3
    output_shape = list(x.shape)
    output_shape[-1] = 4
    ret = np.ones(output_shape)
    ret[..., :3] = x
    return ret


def to_normal(x):
    if x.shape[-1] == 3:
        return x
    assert x.shape[-1] == 4
    return x[..., :3] / x[..., 3:]


def o3d_to_trimesh(x):
    vertices = np.asarray(x.vertices)
    faces = np.asarray(x.triangles)
    return trimesh.Trimesh(vertices=vertices, faces=faces)


def trimesh_to_o3d(x):
    return o3d.geometry.TriangleMesh(
        o3d.utility.Vector3dVector(x.vertices), o3d.utility.Vector3iVector(x.faces)
    )


def depth2xyz(depth_cam_matrix, depth_map, depth_scale=1, flatten=False):
    fx, fy = depth_cam_matrix[0, 0], depth_cam_matrix[1, 1]
    cx, cy = depth_cam_matrix[0, 2], depth_cam_matrix[1, 2]
    h, w = np.mgrid[0 : depth_map.shape[0], 0 : depth_map.shape[1]]
    z = depth_map / depth_scale
    x = -(w - cx) * z / fx
    y = (h - cy) * z / fy
    xyz = (
        np.dstack((x, y, z))
        if flatten == False
        else np.dstack((x, y, z)).reshape(-1, 3)
    )
    return xyz


def xyz2uv(pc_cam, depth_cam_matrix):
    intrinsic = depth_cam_matrix[:3, :3]
    pc_norm = np.dot(intrinsic, pc_cam.T).T
    pc_norm[:, 0] /= pc_norm[:, 2]
    pc_norm[:, 1] /= pc_norm[:, 2]
    # u = fx * raw_xyz[0] / raw_xyz[2] + cx
    # v = fy * raw_xyz[1] / raw_xyz[2] + cy
    return pc_norm


def uv2xyz(uvz, depth_cam_matrix):
    intrinsic_inv = np.linalg.inv(depth_cam_matrix[:3, :3])
    zuzvz = deepcopy(uvz)
    zuzvz[:, 0] *= zuzvz[:, 2]
    zuzvz[:, 1] *= zuzvz[:, 2]
    xyz = np.dot(intrinsic_inv, zuzvz.T).T
    return xyz


def apply_pose_to_points(x, pose):
    assert x.ndim == 2 and x.shape[1] == 3
    return to_normal(to_generalized(x) @ pose.to_transformation_matrix().T)


def apply_pose(pose, x):
    return apply_pose_to_points(x, pose)


def apply_pose_to_point(x, pose, with_translation=True):
    mat = pose.to_transformation_matrix()
    if with_translation:
        return mat[:3, :3] @ x + mat[:3, 3]
    else:
        return mat[:3, :3] @ x


def get_articulations_obs(articulations, with_root_info=True):
    if not isinstance(articulations, list):
        articulations = [articulations]
    ret = []
    for articulation in articulations:
        ret = [articulation.get_qpos(), articulation.get_qvel()]
        if with_root_info:
            root_pose = articulation.get_root_pose()
            ret += [root_pose.p, root_pose.q]
    return np.concatenate(ret)


def get_actors_obs(actors, base_pose=Pose()):
    if not isinstance(actors, list):
        actors = [actors]
    ret = []
    inv_base = base_pose.inv()
    inv_rot = inv_base.to_transformation_matrix()[:3, :3]
    for actor in actors:
        pose_in_base = actor.get_pose() * inv_base
        vel_in_base = inv_rot @ actor.get_velocity()
        ang_vel_in_base = inv_rot @ actor.get_angular_velocity()
        ret += [pose_in_base.p, pose_in_base.q, vel_in_base, ang_vel_in_base]
    return np.concatenate(ret)


def normalize_vec(x):
    return x / np.clip(np.linalg.norm(x, keepdims=True), a_min=1e-6, a_max=1e6)


def angle_between_vec(a, b):
    return np.arccos(np.dot(a, b)) / pi


def rew_close_to_target(
    pos,
    vel,
    target_pos,
    max_dist=1,
    max_vel=1,
    dist_coeff=1,
    vel_coeff=2,
    angle_coeff=1,
):
    ret = 0
    info = {}
    target_vel = normalize_vec(target_pos - pos)
    info["target_vel"] = target_vel
    if dist_coeff > 0:
        dist = np.linalg.norm(pos - target_pos)
        info["rew_dist"] = (
            1 - normalize_and_clip_in_interval(dist, 0, max_dist)
        ) * dist_coeff
        ret += info[
            "rew_dist"
        ]  # - normalize_and_clip_in_interval(dist, 0, max_dist) * dist_coeff + dist_coeff
    if vel_coeff > 0:
        vel_norm = np.dot(vel, target_vel)
        info["rew_vel"] = normalize_and_clip_in_interval(vel_norm, max_vel) * vel_coeff
        ret += info["rew_vel"]
    if angle_coeff > 0:
        vel_angle = angle_between_vec(normalize_vec(vel), target_vel)
        info["rew_vel_angle"] = (1 - vel_angle) * angle_coeff
        ret += info["rew_vel_angle"]
    return ret, info


def normalize_and_clip_in_interval(x, min_x, max_x=None):
    if max_x is None:
        min_x = -abs(min_x)
        max_x = abs(min_x)
    len_x = max_x - min_x
    return (min(max(x, min_x), max_x) - min_x) / len_x


def normalize_reward(x, norm_x):
    return x / norm_x


def quaternion_distance(q1, q2):
    q1, q2 = Pose(q=q1), Pose(q=q2)
    q = q1.inv() * q2
    # print(q.q)
    return 1 - np.abs(q.q[0])  # 1-dot(q1, q2)


def build_pose(forward, flat):
    extra = np.cross(forward, flat)
    ans = np.eye(4)
    ans[:3, :3] = np.array([forward, flat, extra])
    return Pose.from_transformation_matrix(ans)


def pose_vec_distance(pose1, pose2):
    dist_p = np.linalg.norm(pose1.p - pose2.p)
    dist_q = quaternion_distance(pose1.q, pose2.q)
    return dist_p + 0.01 * dist_q


def pose_corner_distance(pose1, pose2):
    unit_box = get_unit_box_corners()
    t1 = pose1.to_transformation_matrix()
    t2 = pose2.to_transformation_matrix()

    corner1 = to_generalized(unit_box) @ t1.T
    corner2 = to_generalized(unit_box) @ t2.T
    # print(corner1.shape, corner2.shape)
    return np.mean(np.linalg.norm(corner1 - corner2, axis=-1))


def generate_ducttape_mesh(
    inner_radius_range, width_range, height_range, n_polygons, num
):
    # Range: [low, high], low>=0, high>=low
    duct_tapes = []
    for _ in range(num):
        for i in range(n_polygons):
            r1 = np.random.uniform(inner_radius_range[0], inner_radius_range[0])
            r2 = r1 + np.random.uniform(width_range[0], width_range[1])
            height = np.random.uniform(height_range[0], height_range[1])
            scene = trimesh.Scene()
            theta1 = 2 * np.pi * i / n_polygons
            theta2 = 2 * np.pi * (i + 1) / n_polygons
            coord1 = np.array([np.cos(theta1), np.sin(theta1)])
            coord2 = np.array([np.cos(theta2), np.sin(theta2)])
            p = [coord1 * r2, coord1 * r1, coord2 * r1, coord2 * r2]
            g = trimesh.creation.extrude_polygon(Polygon(p), height)
            scene.add_geometry(g)
            duct_tapes.append(scene)
    return duct_tapes


def compute_dist2pcd(triangle_vertices, point):
    return np.min(np.linalg.norm(triangle_vertices - point, axis=-1))


def compute_dist2surface(triangle_vertices, triangle_indices, point):
    triangles = triangle_vertices[triangle_indices.reshape(-1, 3)]
    p = trimesh.triangles.closest_point(
        triangles, np.tile(point, (triangles.shape[0], 1))
    )
    return np.min(np.linalg.norm(p - point, axis=1))


def compute_dist2object(obj, point):
    point = obj.get_pose().inv().transform(Pose(point)).p
    ds = [
        compute_dist2surface(g.geometry.vertices, g.geometry.indices, point)
        for g in obj.get_collision_shapes()
    ]
    return np.min(ds)

arti_mani/utils/contrib.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import cv2
import matplotlib.cm as cm
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
import numpy as np


def show_cmap(cmap, norm=None, extend=None):
    """show the colormap."""
    if norm is None:
        norm = mcolors.Normalize(vmin=0, vmax=cmap.N)
    im = cm.ScalarMappable(norm=norm, cmap=cmap)

    fig, ax = plt.subplots(figsize=(6, 1))
    fig.subplots_adjust(bottom=0.5)
    fig.colorbar(im, cax=ax, orientation="horizontal", extend=extend)
    plt.show()


def depth2pts_np(
    depth_map: np.ndarray,
    cam_intrinsic: np.ndarray,
    cam_extrinsic: np.ndarray = np.eye(4),
) -> np.ndarray:
    assert (len(depth_map.shape) == 2) or (
        len(depth_map.shape) == 3 and depth_map.shape[2] == 1
    )
    assert cam_intrinsic.shape == (3, 3)
    assert cam_extrinsic.shape == (4, 4)
    feature_grid = get_pixel_grids_np(depth_map.shape[0], depth_map.shape[1])

    uv = np.matmul(np.linalg.inv(cam_intrinsic), feature_grid)
    cam_points = uv * np.reshape(depth_map, (1, -1))  # (3, N)

    R = cam_extrinsic[:3, :3]
    t = cam_extrinsic[:3, 3:4]
    R_inv = np.linalg.inv(R)

    world_points = np.matmul(R_inv, cam_points - t).transpose()  # (N, 3)
    return world_points


def get_pixel_grids_np(height: int, width: int):
    x_linspace = np.linspace(0.5, width - 0.5, width)
    y_linspace = np.linspace(0.5, height - 0.5, height)
    x_coordinates, y_coordinates = np.meshgrid(x_linspace, y_linspace)
    x_coordinates = np.reshape(x_coordinates, (1, -1))
    y_coordinates = np.reshape(y_coordinates, (1, -1))
    ones = np.ones_like(x_coordinates).astype(np.float)
    grid = np.concatenate([x_coordinates, y_coordinates, ones], axis=0)

    return grid


def visualize_depth(depth):
    MAX_DEPTH = 2.0
    cmap = plt.get_cmap("rainbow")
    if depth.dtype == np.uint16:
        depth = depth.astype(np.float32) / 1000.0
    if len(depth.shape) == 3:
        depth = depth[..., 0]
    depth = np.clip(depth / MAX_DEPTH, 0.0, 1.0)
    vis_depth = cmap(depth)
    vis_depth = (vis_depth[:, :, :3] * 255.0).astype(np.uint8)
    vis_depth = cv2.cvtColor(vis_depth, cv2.COLOR_RGB2BGR)
    return vis_depth


def get_cmaps(num_colors):
    if num_colors == 6:
        cmaps = [
            (1, 0, 0),  # "red"
            (0, 0, 1),  # "blue"
            (1, 1, 0),  # "yellow"
            (0, 1, 0),  # "green"
            (0.627, 0.125, 0.941),  # "purple"
            (0.753, 0.753, 0.753),  # "grey"
        ]
    elif num_colors == 5:
        cmaps = [
            (1, 0, 0),  # "red"
            (0, 0, 1),  # "blue"
            (1, 1, 0),  # "yellow"
            (0, 1, 0),  # "green"
            (0.753, 0.753, 0.753),  # "grey"
        ]
    else:
        cmaps = cm.get_cmap("jet", num_colors)(range(num_colors))[:, :3]
    return cmaps


def visualize_seg(seg, num_classes=6):
    H, W = seg.shape
    cmaps = get_cmaps(num_classes)
    segvis_mask = np.zeros((H, W, 3))
    for ind in range(num_classes):
        segvis_mask[seg == ind] = cmaps[ind]
    return segvis_mask


class OpenCVViewer:
    def __init__(self, name="OpenCVViewer", is_rgb=True, exit_on_esc=True):
        self.name = name
        cv2.namedWindow(name, cv2.WINDOW_AUTOSIZE)
        self.is_rgb = is_rgb
        self.exit_on_esc = exit_on_esc

    def imshow(self, image: np.ndarray, is_rgb=None, non_blocking=False, delay=0):
        if image.ndim == 2:
            image = np.tile(image[..., np.newaxis], (1, 1, 3))
        elif image.ndim == 3 and image.shape[-1] == 1:
            image = np.tile(image, (1, 1, 3))
        assert image.ndim == 3, image.shape

        if self.is_rgb or is_rgb:
            image = image[..., ::-1]
        cv2.imshow(self.name, image)

        if non_blocking:
            return
        else:
            key = cv2.waitKey(delay)
            if key == 27:  # escape
                if self.exit_on_esc:
                    exit(0)
                else:
                    return None
            elif key == -1:  # timeout
                pass
            else:
                return chr(key)

    def close(self):
        cv2.destroyWindow(self.name)

    def __del__(self):
        self.close()

arti_mani/utils/cv_utils.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from typing import Dict

import numpy as np
import sapien.core as sapien
import sklearn
import transforms3d
import trimesh
from arti_mani.utils.o3d_utils import is_o3d, np2mesh, to_o3d
from arti_mani.utils.trimesh_utils import is_trimesh, to_trimesh
from sapien.core import Articulation, Pose
from scipy.spatial.transform import Rotation
from transforms3d.quaternions import nearly_equivalent


def sample_on_unit_sphere(rng):
    """
    Algo from http://corysimon.github.io/articles/uniformdistn-on-sphere/
    """
    v = np.zeros(3)
    while np.linalg.norm(v) < 1e-4:
        v[0] = rng.normal()  # random standard normal
        v[1] = rng.normal()
        v[2] = rng.normal()

    v = v / np.linalg.norm(v)
    return v


def sample_on_unit_circle(rng):
    v = np.zeros(2)
    while np.linalg.norm(v) < 1e-4:
        v[0] = rng.normal()  # random standard normal
        v[1] = rng.normal()

    v = v / np.linalg.norm(v)
    return v


def sample_grasp_points_ee(gripper_pos, x_offset=0.03):
    """
    sample 6 points representing gripper, in EE frame
    """
    x_trans = gripper_pos + 0.002
    finger_points = np.array(
        [
            [-0.14, 0, 0],
            [-0.07, 0, 0],
            [-0.07, x_trans, 0],
            [0, x_trans, 0],
            [-0.07, -x_trans, 0],
            [0, -x_trans, 0],
        ]
    )
    adjust_finger_points = finger_points + np.array([x_offset, 0, 0])
    return adjust_finger_points


def sample_grasp_multipoints_ee(gripper_pos, num_points_perlink=10, x_offset=0.03):
    """
    sample 6 points representing gripper, in EE frame
    """
    x_trans = gripper_pos + 0.002
    left_finger = np.linspace(
        [0, -x_trans, 0], [-0.07, -x_trans, 0], num_points_perlink
    )[:-1]
    left_knuckle = np.linspace([-0.07, -x_trans, 0], [-0.07, 0, 0], num_points_perlink)[
        :-1
    ]
    right_finger = np.linspace(
        [0, x_trans, 0], [-0.07, x_trans, 0], num_points_perlink
    )[:-1]
    right_knuckle = np.linspace([-0.07, x_trans, 0], [-0.07, 0, 0], num_points_perlink)[
        :-1
    ]
    base = np.linspace([-0.14, 0, 0], [-0.07, 0, 0], num_points_perlink)
    finger_points = np.concatenate(
        (base, left_knuckle, left_finger, right_knuckle, right_finger)
    )
    adjust_finger_points = finger_points + np.array([x_offset, 0, 0])
    return adjust_finger_points


def rotation_between_vec(a, b):  # from a to b
    a = a / np.linalg.norm(a)
    b = b / np.linalg.norm(b)
    axis = np.cross(a, b)
    axis = axis / np.linalg.norm(axis)  # norm might be 0
    angle = np.arccos(a @ b)
    R = Rotation.from_rotvec(axis * angle)
    return R


def angle_between_vec(a, b):  # from a to b
    a = a / np.linalg.norm(a)
    b = b / np.linalg.norm(b)
    angle = np.arccos(a @ b)
    return angle


def wxyz_to_xyzw(q):
    return np.concatenate([q[1:4], q[0:1]])


def xyzw_to_wxyz(q):
    return np.concatenate([q[3:4], q[0:3]])


def rotate_2d_vec_by_angle(vec, theta):
    rot_mat = np.array(
        [[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]
    )
    return rot_mat @ vec


def angle_distance(q0: sapien.Pose, q1: sapien.Pose):
    qd = (q0.inv() * q1).q
    return 2 * np.arctan2(np.linalg.norm(qd[1:]), qd[0])


def get_axis_aligned_bbox_for_articulation(art: Articulation):
    mins = np.array([np.inf, np.inf, np.inf])
    maxs = -mins
    for link in art.get_links():
        lp = link.pose
        for s in link.get_collision_shapes():
            p = lp * s.get_local_pose()
            T = p.to_transformation_matrix()
            vertices = s.geometry.vertices * s.geometry.scale
            vertices = vertices @ T[:3, :3].T + T[:3, 3]
            mins = np.minimum(mins, vertices.min(0))
            maxs = np.maximum(maxs, vertices.max(0))
    return mins, maxs


def get_axis_aligned_bbox_for_actor(actor):
    mins = np.ones(3) * np.inf
    maxs = -mins

    for shape in actor.get_collision_shapes():  # this is CollisionShape
        scaled_vertices = shape.geometry.vertices * shape.geometry.scale
        local_pose = shape.get_local_pose()
        mat = (actor.get_pose() * local_pose).to_transformation_matrix()
        world_vertices = scaled_vertices @ (mat[:3, :3].T) + mat[:3, 3]
        mins = np.minimum(mins, world_vertices.min(0))
        maxs = np.maximum(maxs, world_vertices.max(0))

    return mins, maxs


def get_local_axis_aligned_bbox_for_link(link):
    mins = np.array([np.inf, np.inf, np.inf])
    maxs = -mins
    for s in link.get_collision_shapes():
        p = s.get_local_pose()
        T = p.to_transformation_matrix()
        vertices = s.geometry.vertices * s.geometry.scale
        vertices = vertices @ T[:3, :3].T + T[:3, 3]
        mins = np.minimum(mins, vertices.min(0))
        maxs = np.maximum(maxs, vertices.max(0))
    return mins, maxs


def transform_points(H: np.ndarray, pts: np.ndarray) -> np.ndarray:
    assert H.shape == (4, 4), H.shape
    assert pts.ndim == 2 and pts.shape[1] == 3, pts.shape
    return pts @ H[:3, :3].T + H[:3, 3]


def get_oriented_bounding_box_for_2d_points(
    points_2d: np.ndarray, resolution=0.0
) -> Dict:
    assert len(points_2d.shape) == 2 and points_2d.shape[1] == 2
    if resolution > 0.0:
        points_2d = np.round(points_2d / resolution) * resolution
        points_2d = np.unique(points_2d, axis=0)
    ca = np.cov(points_2d, y=None, rowvar=0, bias=1)

    v, vect = np.linalg.eig(ca)
    tvect = np.transpose(vect)

    # use the inverse of the eigenvectors as a rotation matrix and
    # rotate the points so they align with the x and y axes
    ar = np.dot(points_2d, np.linalg.inv(tvect))

    # get the minimum and maximum x and y
    mina = np.min(ar, axis=0)
    maxa = np.max(ar, axis=0)
    half_size = (maxa - mina) * 0.5

    # the center is just half way between the min and max xy
    center = mina + half_size
    # get the 4 corners by subtracting and adding half the bounding boxes height and width to the center
    corners = np.array(
        [
            center + [-half_size[0], -half_size[1]],
            center + [half_size[0], -half_size[1]],
            center + [half_size[0], half_size[1]],
            center + [-half_size[0], half_size[1]],
        ]
    )

    # use the the eigenvectors as a rotation matrix and
    # rotate the corners and the centerback
    corners = np.dot(corners, tvect)
    center = np.dot(center, tvect)

    return {"center": center, "half_size": half_size, "axes": vect, "corners": corners}


def pose2mat(pose):
    mat = np.eye(4)
    mat[:3, 3] = pose[:3]
    mat[:3, :3] = transforms3d.quaternions.quat2mat(pose[3:])
    return mat


def mat2pose(mat: np.ndarray):
    quat = transforms3d.quaternions.mat2quat(mat[:3, :3])
    pos = mat[:3, 3]
    pose = np.concatenate((pos, quat))
    return pose


def pose2trans_axangle(pose):
    rot_mat = pose2mat(pose)
    axis, angle = transforms3d.axangles.mat2axangle(
        rot_mat[:3, :3]
    )  # (normalized pivot axis:3, angle:1)
    trans_axangle = np.concatenate((rot_mat[:3, 3], angle * axis))
    return trans_axangle


def trans_axangle2pose(trans_axangle):
    RT = np.eye(4)
    RT[:3, 3] = trans_axangle[:3]
    angle = np.linalg.norm(trans_axangle[3:6])
    if angle < 1e-6:
        axis = (0, 0, 1)
    else:
        axis = trans_axangle[3:6] / angle
    RT[:3, :3] = transforms3d.axangles.axangle2mat(axis, angle)
    return mat2pose(RT)


def euler2quat(rot_xyz, axes="sxyz", unit="deg"):  ## rx, ry, rz, axes='sxyz'
    if unit == "deg":
        rx, ry, rz = (
            rot_xyz[0] * np.pi / 180,
            rot_xyz[1] * np.pi / 180,
            rot_xyz[2] * np.pi / 180,
        )
    return transforms3d.euler.euler2quat(rx, ry, rz, axes)


def quat2euler(quat, axes="sxyz", unit="deg"):  ## rx, ry, rz, axes='sxyz'
    rot_xyz = np.array(transforms3d.euler.quat2euler(quat, axes))
    if unit == "deg":
        rot_xyz = rot_xyz * 180 / np.pi
    return rot_xyz


def all_close(p1, p2, t_eps=1e-4, r_eps=1e-5):
    assert type(p1) == type(p2)
    return np.allclose(p1.p, p2.p, atol=t_eps) and nearly_equivalent(
        p1.q, p2.q, atol=r_eps
    )


def convex_hull(x, o3d=True):
    x = to_trimesh(x)
    x = trimesh.convex.convex_hull(x)
    if o3d:
        x = to_o3d(x)
    return x


def apply_pose(pose, x):
    if x is None:
        return x
    if isinstance(pose, np.ndarray):
        pose = Pose.from_transformation_matrix(pose)
    assert isinstance(pose, Pose)
    if isinstance(x, Pose):
        return pose * x
    elif isinstance(x, np.ndarray):
        assert x.ndim == 2 and x.shape[1] == 3
        x = np.concatenate((x, np.ones((x.shape[0], 1), dtype=x.dtype)), axis=1)
        return (x @ pose.to_transformation_matrix().T)[:, :3]
    elif is_trimesh(x) or is_o3d(x):
        sign = is_o3d(x)
        x = to_trimesh(x)
        if isinstance(x, trimesh.Trimesh):
            vertices = x.vertices
            faces = x.faces
            vertices = apply_pose(pose, vertices)
            x = trimesh.Trimesh(vertices=vertices, faces=faces)
        elif isinstance(x, trimesh.points.PointCloud):
            vertices = x.vertices
            vertices = apply_pose(pose, vertices)
            x = trimesh.points.PointCloud(vertices=vertices)
        if sign:
            x = to_o3d(x)
        return x
    else:
        print(x, type(x))
        raise NotImplementedError("")


def check_coplanar(vertices):
    pca = sklearn.decomposition.PCA(n_components=3)
    pca.fit(vertices)
    return pca.singular_values_[-1] < 1e-3

arti_mani/utils/geometry.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from typing import List

import numpy as np
import open3d as o3d
import sapien.core as sapien
import sklearn
import trimesh


def check_coplanar(vertices):
    pca = sklearn.decomposition.PCA(n_components=3)
    pca.fit(vertices)
    return pca.singular_values_[-1] < 1e-3


# ---------------------------------------------------------------------------- #
# Convert in opne3d
# ---------------------------------------------------------------------------- #


def get_visual_body_meshes(visual_body: sapien.RenderBody):
    meshes = []
    for render_shape in visual_body.get_render_shapes():
        vertices = render_shape.mesh.vertices * visual_body.scale  # [n, 3]
        faces = render_shape.mesh.indices.reshape(-1, 3)  # [m * 3]
        mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
        mesh.apply_transform(visual_body.local_pose.to_transformation_matrix())
        meshes.append(mesh)
    return meshes


def get_actor_visual_meshes(actor: sapien.ActorBase):
    """Get actor (visual) meshes in the actor frame."""
    meshes = []
    for vb in actor.get_visual_bodies():
        meshes.extend(get_visual_body_meshes(vb))
    return meshes


def merge_meshes(meshes: List[trimesh.Trimesh]):
    n, vs, fs = 0, [], []
    for mesh in meshes:
        v, f = mesh.vertices, mesh.faces
        vs.append(v)
        fs.append(f + n)
        n = n + v.shape[0]
    if n:
        return trimesh.Trimesh(np.vstack(vs), np.vstack(fs))
    else:
        return None


def merge_mesh(meshes: List[o3d.geometry.TriangleMesh]) -> o3d.geometry.TriangleMesh:
    if not meshes:
        return None
    # Merge without color and normal
    vertices = np.zeros((0, 3))
    triangles = np.zeros((0, 3))

    for mesh in meshes:
        vertices_i = np.asarray(mesh.vertices)
        triangles_i = np.asarray(mesh.triangles)
        triangles_i += vertices.shape[0]
        vertices = np.append(vertices, vertices_i, axis=0)
        triangles = np.append(triangles, triangles_i, axis=0)

    vertices = o3d.utility.Vector3dVector(vertices)
    triangles = o3d.utility.Vector3iVector(triangles)
    # print(vertices, triangles)
    # exit(0)
    mesh = o3d.geometry.TriangleMesh(vertices, triangles)
    mesh.compute_vertex_normals(normalized=True)
    mesh.compute_triangle_normals(normalized=True)
    return mesh


def get_actor_meshes(actor: sapien.ActorBase):
    """Get actor (collision) meshes in the actor frame."""
    meshes = []
    for col_shape in actor.get_collision_shapes():
        geom = col_shape.geometry
        if isinstance(geom, sapien.BoxGeometry):
            mesh = trimesh.creation.box(extents=2 * geom.half_lengths)
        elif isinstance(geom, sapien.CapsuleGeometry):
            mesh = trimesh.creation.capsule(
                height=2 * geom.half_length, radius=geom.radius
            )
        elif isinstance(geom, sapien.SphereGeometry):
            mesh = trimesh.creation.icosphere(radius=geom.radius)
        elif isinstance(geom, sapien.PlaneGeometry):
            continue
        elif isinstance(
            geom, (sapien.ConvexMeshGeometry, sapien.NonconvexMeshGeometry)
        ):
            vertices = geom.vertices  # [n, 3]
            faces = geom.indices.reshape(-1, 3)  # [m * 3]
            vertices = vertices * geom.scale
            mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
        else:
            raise TypeError(type(geom))
        mesh.apply_transform(col_shape.get_local_pose().to_transformation_matrix())
        meshes.append(mesh)
    return meshes


def get_actor_mesh(actor: sapien.ActorBase, to_world_frame=True):
    mesh = merge_meshes(get_actor_meshes(actor))
    if mesh is None:
        return None
    if to_world_frame:
        T = actor.pose.to_transformation_matrix()
        mesh.apply_transform(T)
    return mesh


def get_actor_visual_mesh(actor: sapien.ActorBase):
    mesh = merge_meshes(get_actor_visual_meshes(actor))
    if mesh is None:
        return None
    return mesh


def get_articulation_meshes(
    articulation: sapien.ArticulationBase, exclude_link_names=()
):
    """Get link meshes in the world frame."""
    meshes = []
    for link in articulation.get_links():
        if link.name in exclude_link_names:
            continue
        mesh = get_actor_mesh(link, True)
        if mesh is None:
            continue
        meshes.append(mesh)
    return meshes


def mesh2pcd(mesh, sample_density, num_points=None) -> o3d.geometry.PointCloud:
    pcd_tmp = mesh.sample_points_uniformly(number_of_points=sample_density)
    points = np.asarray(pcd_tmp.points)
    normals = np.asarray(pcd_tmp.normals)

    pcd = o3d.geometry.PointCloud()
    if num_points:
        idx = np.arange(points.shape[0])
        np.random.shuffle(idx)
        idx = idx[:num_points]
        points = points[idx]
        normals = normals[idx]
    # print(vertices.shape, normals.shape)
    pcd.points = o3d.utility.Vector3dVector(points)
    pcd.normals = o3d.utility.Vector3dVector(normals)
    return pcd


# ---------------------------------------------------------------------------- #
# Build from numpy
# ---------------------------------------------------------------------------- #


def np2mesh(
    vertices, triangles, colors=None, vertex_normals=None, triangle_normals=None
) -> o3d.geometry.TriangleMesh:
    """Convert numpy array to open3d TriangleMesh."""
    vertices = o3d.utility.Vector3dVector(vertices)
    triangles = o3d.utility.Vector3iVector(triangles)
    mesh = o3d.geometry.TriangleMesh(vertices, triangles)
    if colors is not None:
        colors = np.array(colors)
        if colors.ndim == 2:
            assert len(colors) == len(vertices)
        elif colors.ndim == 1:
            colors = np.tile(colors, (len(vertices), 1))
        else:
            raise RuntimeError(colors.shape)
        mesh.vertex_colors = o3d.utility.Vector3dVector(colors)
    if vertex_normals is not None:
        assert len(triangles) == len(vertex_normals)
        mesh.vertex_normals = o3d.utility.Vector3dVector(vertex_normals)
    else:
        mesh.compute_vertex_normals(normalized=True)

    if triangle_normals is not None:
        assert len(triangles) == len(triangle_normals)
        mesh.triangle_normals = o3d.utility.Vector3dVector(triangle_normals)
    else:
        mesh.compute_triangle_normals(normalized=True)
    return mesh


def np2pcd(points, colors=None, normals=None) -> o3d.geometry.PointCloud:
    """Convert numpy array to open3d PointCloud."""
    pc = o3d.geometry.PointCloud()
    pc.points = o3d.utility.Vector3dVector(points)
    if colors is not None:
        colors = np.array(colors)
        if colors.ndim == 2:
            assert len(colors) == len(points)
        elif colors.ndim == 1:
            colors = np.tile(colors, (len(points), 1))
        else:
            raise RuntimeError(colors.shape)
        pc.colors = o3d.utility.Vector3dVector(colors)
    if normals is not None:
        assert len(points) == len(normals)
        pc.normals = o3d.utility.Vector3dVector(normals)
    return pc


def create_aabb(bbox, color=(0, 1, 0)) -> o3d.geometry.AxisAlignedBoundingBox:
    """Draw an axis-aligned bounding box."""
    assert (
        len(bbox) == 6
    ), f"The format of bbox should be xyzwlh, but received {len(bbox)}."
    bbox = np.asarray(bbox)
    abb = o3d.geometry.AxisAlignedBoundingBox(
        bbox[0:3] - bbox[3:6] * 0.5, bbox[0:3] + bbox[3:6] * 0.5
    )
    abb.color = color
    return abb


def create_aabb_from_pcd(
    pcd: np.ndarray, color=(0, 1, 0)
) -> o3d.geometry.AxisAlignedBoundingBox:
    """Draw an axis-aligned bounding box."""
    assert (
        pcd.shape[-1] == 3
    ), f"The format of bbox should be xyzwlh, but received {pcd.shape}."
    abb = o3d.geometry.AxisAlignedBoundingBox.create_from_points(
        o3d.utility.Vector3dVector(pcd)
    )
    abb.color = color
    return abb


def create_obb(bbox, R, color=(0, 1, 0)):
    """Draw an oriented bounding box."""
    assert (
        len(bbox) == 6
    ), f"The format of bbox should be xyzwlh, but received {len(bbox)}."
    obb = o3d.geometry.OrientedBoundingBox(bbox[0:3], R, bbox[3:6])
    obb.color = color
    return obb


def create_obb_from_pcd(pcd, color=(0, 1, 0)):
    """Draw an axis-aligned bounding box."""
    assert (
        pcd.shape[-1] == 3
    ), f"The format of bbox should be xyzwlh, but received {pcd.shape}."
    obb = o3d.geometry.OrientedBoundingBox.create_from_points(
        o3d.utility.Vector3dVector(pcd)
    )
    obb.color = color
    return obb


# ---------------------------------------------------------------------------- #
# Computation
# ---------------------------------------------------------------------------- #
def compute_pcd_normals(points, search_param=None, camera_location=(0.0, 0.0, 0.0)):
    """Compute normals."""
    pcd = np2pcd(points)
    if search_param is None:
        pcd.estimate_normals()
    else:
        pcd.estimate_normals(search_param=search_param)
    pcd.orient_normals_towards_camera_location(camera_location)
    normals = np.array(pcd.normals)
    return normals


def pcd_voxel_down_sample_with_crop(
    points,
    voxel_size,
    min_bound: np.ndarray,
    max_bound: np.ndarray,
) -> List[int]:
    """Crop and downsample the point cloud and return sample indices."""
    crop_mask = np.logical_and(
        np.logical_and.reduce(points > min_bound[None, :], axis=1),
        np.logical_and.reduce(points < max_bound[None, :], axis=1),
    )
    if not crop_mask.any():
        return []
    else:
        crop_indices = np.where(crop_mask)[0]
        points = points[crop_mask]
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points)

        downsample_pcd, mapping, index_buckets = pcd.voxel_down_sample_and_trace(
            voxel_size, min_bound[:, None], max_bound[:, None]
        )
        sample_indices = [crop_indices[int(x[0])] for x in index_buckets]
        return sample_indices


def pcd_uni_down_sample_with_crop(
    points,
    num,
    min_bound: np.ndarray,
    max_bound: np.ndarray,
) -> List[int]:
    """Crop and downsample the point cloud and return sample indices."""
    crop_mask = np.logical_and(
        np.logical_and.reduce(points > min_bound[None, :], axis=1),
        np.logical_and.reduce(points < max_bound[None, :], axis=1),
    )
    if not crop_mask.any():
        return []
    else:
        crop_indices = np.where(crop_mask)[0]
        sample_points = points[crop_mask]

        N = sample_points.shape[0]
        index = np.arange(N)
        if N > num:
            np.random.shuffle(index)
            index = index[:num]
        else:
            num_repeat = num // N
            index = np.concatenate([index for i in range(num_repeat)])
            index = np.concatenate([index, index[: num - N * num_repeat]])
        sample_indices = crop_indices[index]
        return sample_indices.tolist()


def is_o3d(x):
    return isinstance(x, (o3d.geometry.TriangleMesh, o3d.geometry.PointCloud))


def to_o3d(x):
    """
    Numpy support is for pcd!
    """
    if is_o3d(x):
        return x
    elif isinstance(x, np.ndarray):
        assert x.ndim == 2 and x.shape[-1] == 3
        return o3d.geometry.PointCloud(o3d.utility.Vector3dVector(x))
    elif isinstance(x, trimesh.Trimesh):
        return o3d.geometry.TriangleMesh(
            o3d.utility.Vector3dVector(x.vertices), o3d.utility.Vector3iVector(x.faces)
        )
    elif isinstance(x, trimesh.points.PointCloud):
        return o3d.geometry.PointCloud(x.vertices)
    else:
        print(type(x))
        raise NotImplementedError()


def draw_3d_example(env, obs):
    pcd_points = obs["pointcloud"]["xyz"]
    pcd_points = env.unwrapped._get_obs_pointcloud_gt()["pointcloud"]["xyz"]
    pt0 = pcd_points[:300]
    pt1 = pcd_points[300:600]
    pt2 = pcd_points[600:900]
    pt3 = pcd_points[900:1200]
    pt4 = env.unwrapped.get_obs()["extra"]["point_target"][None]
    pcd0 = o3d.geometry.PointCloud()
    pcd0.points = o3d.utility.Vector3dVector(pt0)
    pcd0.paint_uniform_color([1, 0, 0])
    pcd1 = o3d.geometry.PointCloud()
    pcd1.points = o3d.utility.Vector3dVector(pt1)
    pcd1.paint_uniform_color([0, 1, 0])
    pcd2 = o3d.geometry.PointCloud()
    pcd2.points = o3d.utility.Vector3dVector(pt2)
    pcd2.paint_uniform_color([0, 0, 1])
    pcd3 = o3d.geometry.PointCloud()
    pcd3.points = o3d.utility.Vector3dVector(pt3)
    pcd3.paint_uniform_color([0, 1, 1])
    pcd4 = o3d.geometry.PointCloud()
    pcd4.points = o3d.utility.Vector3dVector(pt4)
    pcd4.paint_uniform_color([1, 0, 1])
    XYZ = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
    o3d.visualization.draw_geometries([pcd0, pcd1, pcd2, pcd3, pcd4, XYZ])

    pcd_points = obs["pointcloud"]["xyz"]
    pt0 = pcd_points[:300]
    pt1 = pcd_points[300:600]
    pcd0 = o3d.geometry.PointCloud()
    pcd0.points = o3d.utility.Vector3dVector(pt0)
    pcd0.paint_uniform_color([1, 0, 0])
    pcd1 = o3d.geometry.PointCloud()
    pcd1.points = o3d.utility.Vector3dVector(pt1)
    pcd1.paint_uniform_color([0, 1, 0])
    XYZ = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
    o3d.visualization.draw_geometries([pcd0, pcd1, XYZ])

    import open3d as o3d
    from arti_mani.agents.camera import get_texture
    from arti_mani.utils.contrib import apply_pose_to_points
    from sapien.core import Pose

    pcgt = obs_global["pointcloud"]["xyz"][:300]
    pcd0 = o3d.geometry.PointCloud()
    pcd0.points = o3d.utility.Vector3dVector(pcgt)
    pcd0.paint_uniform_color([1, 0, 0])
    ego_cam = env.unwrapped._agent._cameras["hand_camera"]
    ego_cam.take_picture()
    position = get_texture(ego_cam, "Position")
    cam_xyz = position[..., :3]
    H, W, _ = cam_xyz.shape
    pc_cam = cam_xyz.reshape(-1, 3)
    ego_cam_extrin = env.unwrapped.cam_para["hand_camera_extrinsic_base_frame"]
    pc_ee = pc_cam @ ego_cam_extrin[:3, :3].transpose() + ego_cam_extrin[:3, 3]
    pcd1 = o3d.geometry.PointCloud()
    pcd1.points = o3d.utility.Vector3dVector(pc_ee)
    pcd1.paint_uniform_color([0, 1, 0])
    XYZ = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
    o3d.visualization.draw_geometries([pcd0, pcd1, XYZ])

    tipm = o3d.geometry.TriangleMesh.create_sphere(0.003)
    tipm.translate(point_pos)
    tipm.paint_uniform_color([0, 0, 1])
    XYZ = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
    o3d.visualization.draw_geometries([tipm, XYZ])

arti_mani/utils/o3d_utils.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from contextlib import contextmanager
from typing import List, Optional, Tuple, Union

import numpy as np
import sapien.core as sapien
from sapien.core import Pose
from transforms3d.quaternions import mat2quat


def normalize_vector(x, eps=1e-6):
    x = np.asarray(x)
    assert x.ndim == 1, x.ndim
    norm = np.linalg.norm(x)
    if norm < eps:
        return np.zeros_like(x)
    else:
        return x / norm


def vectorize_pose(pose: sapien.Pose):
    return np.hstack([pose.p, pose.q])


def get_actor_by_name(actors, names: Union[List[str], str]):
    assert isinstance(actors, (list, tuple))
    # Actors can be joint and link
    if isinstance(names, str):
        names = [names]
        sign = True
    else:
        sign = False
    ret = [None for _ in names]
    for actor in actors:
        if actor.get_name() in names:
            ret[names.index(actor.get_name())] = actor
    return ret[0] if sign else ret


def make_actor_visible(actor: sapien.Actor, visible=True):
    for v in actor.get_visual_bodies():
        v.set_visibility(1.0 if visible else 0.0)


@contextmanager
def set_default_physical_material(
    material: sapien.PhysicalMaterial, scene: sapien.Scene
):
    """Set default physical material within the context.

    Args:
        material (sapien.PhysicalMaterial): physical material to use as default.
        scene (sapien.Scene): scene instance.

    Yields:
        sapien.PhysicalMaterial: original default physical material.

    Example:
        with set_default_physical_material(material, scene):
            ...
    """
    old_material = scene.default_physical_material
    scene.default_physical_material = material
    try:
        yield old_material
    finally:
        scene.default_physical_material = old_material


def get_entity_by_name(entities, name: str, is_unique=True):
    """Get a Sapien.Entity given the name.

    Args:
        entities (List[sapien.Entity]): entities (link, joint, ...) to query.
        name (str): name for query.
        is_unique (bool, optional):
            whether the name should be unique. Defaults to True.

    Raises:
        RuntimeError: The name is not unique when @is_unique is True.

    Returns:
        sapien.Entity or List[sapien.Entity]:
            matched entity or entities. None if no matches.
    """
    matched_entities = [x for x in entities if x.get_name() == name]
    if len(matched_entities) > 1:
        if not is_unique:
            return matched_entities
        else:
            raise RuntimeError(f"Multiple entities with the same name {name}.")
    elif len(matched_entities) == 1:
        return matched_entities[0]
    else:
        return None


# -------------------------------------------------------------------------- #
# Entity state
# -------------------------------------------------------------------------- #
def get_actor_state(actor: sapien.Actor):
    pose = actor.get_pose()
    if actor.type == "static":
        vel = np.zeros(3)
        ang_vel = np.zeros(3)
    else:
        vel = actor.get_velocity()  # [3]
        ang_vel = actor.get_angular_velocity()  # [3]
    return np.hstack([pose.p, pose.q, vel, ang_vel])  # 3+4+3+3=13


def set_actor_state(actor: sapien.Actor, state: np.ndarray):
    assert len(state) == 13, len(state)
    actor.set_pose(Pose(state[0:3], state[3:7]))
    if actor.type != "static":
        actor.set_velocity(state[7:10])
        actor.set_angular_velocity(state[10:13])


def get_articulation_state(articulation: sapien.Articulation):
    root_link = articulation.get_links()[0]
    pose = root_link.get_pose()
    vel = root_link.get_velocity()  # [3]
    ang_vel = root_link.get_angular_velocity()  # [3]
    qpos = articulation.get_qpos()
    qvel = articulation.get_qvel()
    return np.hstack(
        [pose.p, pose.q, vel, ang_vel, qpos, qvel]
    )  # 7+2 robot: 3+4+3+3+9+9=31


def set_articulation_state(articulation: sapien.Articulation, state: np.ndarray):
    articulation.set_root_pose(Pose(state[0:3], state[3:7]))
    articulation.set_root_velocity(state[7:10])
    articulation.set_root_angular_velocity(state[10:13])
    qpos, qvel = np.split(state[13:], 2)
    articulation.set_qpos(qpos)
    articulation.set_qvel(qvel)


def get_pad_articulation_state(
    articulation: sapien.Articulation, max_dof: int
) -> np.ndarray:
    root_link = articulation.get_links()[0]
    pose = root_link.get_pose()
    base_pos, base_quat = pose.p, pose.q
    base_vel = root_link.get_velocity()  # [3]
    base_ang_vel = root_link.get_angular_velocity()  # [3]
    qpos = articulation.get_qpos()
    qvel = articulation.get_qvel()
    k = len(qpos)
    # pad_obj_internal_state = np.zeros(max_dof)
    # pad_obj_internal_state[:k] = qpos
    pad_obj_internal_state = np.zeros(2 * max_dof)
    pad_obj_internal_state[:k] = qpos
    pad_obj_internal_state[max_dof : max_dof + k] = qvel
    return np.concatenate(
        [base_pos, base_quat, base_vel, base_ang_vel, pad_obj_internal_state]
    )


# -------------------------------------------------------------------------- #
# Contact
# -------------------------------------------------------------------------- #
def get_pairwise_contacts(
    contacts: List[sapien.Contact], actor0: sapien.ActorBase, actor1: sapien.ActorBase
) -> List[Tuple[sapien.Contact, bool]]:
    pairwise_contacts = []
    for contact in contacts:
        if contact.actor0 == actor0 and contact.actor1 == actor1:
            pairwise_contacts.append((contact, True))
        elif contact.actor0 == actor1 and contact.actor1 == actor0:
            pairwise_contacts.append((contact, False))
    return pairwise_contacts


def compute_total_impulse(contact_infos: List[Tuple[sapien.Contact, bool]]):
    total_impulse = np.zeros(3)
    for contact, flag in contact_infos:
        contact_impulse = np.sum([point.impulse for point in contact.points], axis=0)
        # Impulse is applied on the first actor
        total_impulse += contact_impulse * (1 if flag else -1)
    return total_impulse


def get_pairwise_contact_impulse(
    contacts: List[sapien.Contact], actor0: sapien.ActorBase, actor1: sapien.ActorBase
):
    pairwise_contacts = get_pairwise_contacts(contacts, actor0, actor1)
    total_impulse = compute_total_impulse(pairwise_contacts)
    return total_impulse


def get_actor_contacts(
    contacts: List[sapien.Contact], actor: sapien.ActorBase
) -> List[Tuple[sapien.Contact, bool]]:
    actor_contacts = []
    for contact in contacts:
        if contact.actor0 == actor:
            actor_contacts.append((contact, True))
        elif contact.actor1 == actor:
            actor_contacts.append((contact, False))
    return actor_contacts


def get_articulation_contacts(
    contacts: List[sapien.Contact],
    articulation: sapien.Articulation,
    excluded_actors: Optional[List[sapien.Actor]] = None,
    included_links: Optional[List[sapien.Link]] = None,
) -> List[Tuple[sapien.Contact, bool]]:
    articulation_contacts = []
    links = articulation.get_links()
    if excluded_actors is None:
        excluded_actors = []
    if included_links is None:
        included_links = links
    for contact in contacts:
        if contact.actor0 in included_links:
            if contact.actor1 in links:
                continue
            if contact.actor1 in excluded_actors:
                continue
            articulation_contacts.append((contact, True))
            # print(contact.actor0, contact.actor1)
        elif contact.actor1 in included_links:
            if contact.actor0 in links:
                continue
            if contact.actor0 in excluded_actors:
                continue
            articulation_contacts.append((contact, False))
            # print(contact.actor0, contact.actor1)
    return articulation_contacts


def compute_max_impulse_norm(contact_infos: List[Tuple[sapien.Contact, bool]]):
    max_impulse_norms = [0]
    for contact, flag in contact_infos:
        max_impulse_norm = max(
            [np.linalg.norm(point.impulse) for point in contact.points]
        )
        max_impulse_norms.append(max_impulse_norm)
    return max(max_impulse_norms)


def get_articulation_max_impulse_norm(
    contacts: List[sapien.Contact],
    articulation: sapien.Articulation,
    excluded_actors: Optional[List[sapien.Actor]] = None,
):
    articulation_contacts = get_articulation_contacts(
        contacts, articulation, excluded_actors
    )
    max_impulse_norm = compute_max_impulse_norm(articulation_contacts)
    return max_impulse_norm


# -------------------------------------------------------------------------- #
# Camera
# -------------------------------------------------------------------------- #
def sapien_pose_to_opencv_extrinsic(sapien_pose_matrix: np.ndarray) -> np.ndarray:
    sapien2opencv = np.array(
        [
            [0.0, -1.0, 0.0, 0.0],
            [0.0, 0.0, -1.0, 0.0],
            [1.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 1.0],
        ],
        dtype=np.float32,
    )
    ex = sapien2opencv @ np.linalg.inv(sapien_pose_matrix)  # world -> camera

    return ex


def look_at(eye, target, up=(0, 0, 1)) -> sapien.Pose:
    """Get the camera pose in SAPIEN by the Look-At method.

    Note:
        https://www.scratchapixel.com/lessons/mathematics-physics-for-computer-graphics/lookat-function
        The SAPIEN camera follows the convention: (forward, right, up) = (x, -y, z)
        while the OpenGL camera follows (forward, right, up) = (-z, x, y)
        Note that the camera coordinate system (OpenGL) is left-hand.

    Args:
        eye: camera location
        target: looking-at location
        up: a general direction of "up" from the camera.

    Raises:
        RuntimeError: @up and @right are both None.

    Returns:
        sapien.Pose: camera pose
    """
    forward = normalize_vector(np.array(target) - np.array(eye))
    up = normalize_vector(up)
    left = np.cross(up, forward)
    up = np.cross(forward, left)
    rotation = np.stack([forward, left, up], axis=1)
    return sapien.Pose(p=eye, q=mat2quat(rotation))


def hex2rgba(h, correction=True):
    # https://stackoverflow.com/a/29643643
    h = h.lstrip("#")
    r, g, b = tuple(int(h[i : i + 2], 16) / 255 for i in (0, 2, 4))
    rgba = np.array([r, g, b, 1])
    if correction:  # reverse gamma correction in sapien
        rgba = rgba**2.2
    return rgba


def set_render_material(material: sapien.RenderMaterial, **kwargs):
    for k, v in kwargs.items():
        if k == "color":
            material.set_base_color(v)
        else:
            setattr(material, k, v)
    return material


def set_articulation_render_material(articulation: sapien.Articulation, **kwargs):
    for link in articulation.get_links():
        for b in link.get_visual_bodies():
            for s in b.get_render_shapes():
                mat = s.material
                set_render_material(mat, **kwargs)
                s.set_material(mat)


def set_cabinet_render_material(articulation: sapien.Articulation, **kwargs):
    handle_kwargs = {
        "color": kwargs.pop("color"),
        "diffuse_texture_filename": kwargs.pop("color_texture_filename"),
    }
    for link in articulation.get_links():
        for b in link.get_visual_bodies():
            if "handle" in b.get_name():
                for s in b.get_render_shapes():
                    mat = s.material
                    set_render_material(mat, **handle_kwargs)
                    s.set_material(mat)
            else:
                for s in b.get_render_shapes():
                    mat = s.material
                    set_render_material(mat, **kwargs)
                    s.set_material(mat)


def set_actor_render_material(actor: sapien.Actor, **kwargs):
    for b in actor.get_visual_bodies():
        for s in b.get_render_shapes():
            mat = s.material
            set_render_material(mat, **kwargs)
            s.set_material(mat)


def check_joint_stuck(
    articulation: sapien.Articulation,
    active_joint_idx: int,
    pos_diff_threshold: float = 1e-3,
    vel_threshold: float = 1e-4,
):
    actual_pos = articulation.get_qpos()[active_joint_idx]
    target_pos = articulation.get_drive_target()[active_joint_idx]
    actual_vel = articulation.get_qvel()[active_joint_idx]

    return (
        abs(actual_pos - target_pos) > pos_diff_threshold
        and abs(actual_vel) < vel_threshold
    )


def ignore_collision(articulation: sapien.Articulation):
    """ignore collision among all movable links for acceleration"""
    for joint, link in zip(articulation.get_joints(), articulation.get_links()):
        if joint.type in ["revolute", "prismatic"]:
            shapes = link.get_collision_shapes()
            for s in shapes:
                g0, g1, g2, g3 = s.get_collision_groups()
                s.set_collision_groups(g0, g1, g2 | 1 << 31, g3)

arti_mani/utils/sapien_utils.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import re


def custom_format(template_string, **kwargs):
    template_string = template_string.replace("{", "{{")
    template_string = template_string.replace("}", "}}")
    template_string = template_string.replace("&lformat ", "{")
    template_string = template_string.replace(" &rformat", "}")
    return template_string.format_map(kwargs)


def regex_match(string, pattern):
    return re.match(pattern, string) is not None


def prefix_match(string, prefix=None):
    """Check if the string matches the given prefix"""
    if prefix is None or len(prefix) == 0:
        return True
    return re.match(f"({prefix})+(.*?)", string) is not None


def float_str(num, precision):
    format_str = "%.{0}f".format(precision)
    return format_str % num


def num_to_str(num, unit=None, precision=2, number_only=False, auto_select_unit=False):
    unit_list = ["K", "M", "G", "T", "P"]
    if auto_select_unit and unit is None:
        for i, tmp in enumerate(unit_list):
            unit_num = 1024 ** (i + 1)
            if num < unit_num:
                break
            unit = tmp
    if unit is not None:
        unit_num = 1024 ** (unit_list.index(unit) + 1)
        num = num * 1.0 / unit_num
    else:
        unit = ""
    if number_only:
        return num
    else:
        return float_str(num, precision) + unit

arti_mani/utils/string_utils.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from typing import List

import numpy as np
import open3d as o3d
import sapien.core as sapien
import trimesh


def get_actor_meshes(actor: sapien.ActorBase):
    """Get actor meshes in the actor frame."""
    meshes = []
    for col_shape in actor.get_collision_shapes():
        geom = col_shape.geometry
        if isinstance(geom, sapien.BoxGeometry):
            mesh = trimesh.creation.box(extents=2 * geom.half_lengths)
        elif isinstance(geom, sapien.CapsuleGeometry):
            mesh = trimesh.creation.capsule(
                height=2 * geom.half_length, radius=geom.radius
            )
        elif isinstance(geom, sapien.SphereGeometry):
            mesh = trimesh.creation.icosphere(radius=geom.radius)
        elif isinstance(geom, sapien.PlaneGeometry):
            continue
        elif isinstance(
            geom, (sapien.ConvexMeshGeometry, sapien.NonconvexMeshGeometry)
        ):
            vertices = geom.vertices  # [n, 3]
            faces = geom.indices  # [m * 3]
            faces = [faces[i : i + 3] for i in range(0, len(faces), 3)]
            vertices = vertices * geom.scale
            mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
        else:
            raise TypeError(type(geom))
        mesh.apply_transform(col_shape.get_local_pose().to_transformation_matrix())
        meshes.append(mesh)
    return meshes


def merge_meshes(meshes: List[trimesh.Trimesh]):
    n, vs, fs = 0, [], []
    for mesh in meshes:
        v, f = mesh.vertices, mesh.faces
        vs.append(v)
        fs.append(f + n)
        n = n + v.shape[0]
    if n:
        return trimesh.Trimesh(np.vstack(vs), np.vstack(fs))
    else:
        return None


def get_actor_mesh(actor: sapien.ActorBase, to_world_frame=True):
    mesh = merge_meshes(get_actor_meshes(actor))
    if mesh is None:
        return None
    if to_world_frame:
        T = actor.pose.to_transformation_matrix()
        mesh.apply_transform(T)
    return mesh


def get_articulation_meshes(
    articulation: sapien.ArticulationBase, exclude_link_names=()
):
    """Get link meshes in the world frame."""
    meshes = []
    for link in articulation.get_links():
        if link.name in exclude_link_names:
            continue
        link_mesh = get_actor_mesh(link, True)
        if link_mesh is None:
            continue
        meshes.append(link_mesh)
    return meshes


def to_trimesh(x):
    if is_trimesh(x):
        return x
    elif isinstance(x, np.ndarray):
        assert x.ndim == 2 and x.shape[-1] == 3
        return trimesh.points.PointCloud(x)
    elif isinstance(x, o3d.geometry.TriangleMesh):
        vertices = np.asarray(x.vertices)
        faces = np.asarray(x.triangles)
        return trimesh.Trimesh(vertices=vertices, faces=faces)
    elif isinstance(x, o3d.geometry.PointCloud):
        points = np.asarray(x.points)
        return trimesh.points.PointCloud(vertices=points)
    else:
        print(type(x))
        raise NotImplementedError()


def is_trimesh(x):
    return isinstance(x, (trimesh.Trimesh, trimesh.points.PointCloud))

arti_mani/utils/trimesh_utils.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
import os
from typing import Dict, List, Optional

import cv2
import imageio
import numpy as np
import tqdm
from arti_mani.utils.cv_utils import visualize_depth


def images_to_video(
    images: List[np.ndarray],
    output_dir: str,
    video_name: str,
    fps: int = 10,
    quality: Optional[float] = 5,
    verbose: bool = True,
    **kwargs,
):
    r"""Calls imageio to run FFMPEG on a list of images. For more info on
    parameters, see https://imageio.readthedocs.io/en/stable/format_ffmpeg.html

    Args:
        images: The list of images. Images should be HxWx3 in RGB order.
        output_dir: The folder to put the video in.
        video_name: The name for the video.
        fps: Frames per second for the video. Not all values work with FFMPEG,
            use at your own risk.
        quality: Default is 5. Uses variable bit rate. Highest quality is 10,
            lowest is 0.  Set to None to prevent variable bitrate flags to
            FFMPEG so you can manually specify them using output_params
            instead. Specifying a fixed bitrate using ‘bitrate’ disables
            this parameter.

    References:
        https://github.com/facebookresearch/habitat-lab/blob/main/habitat/utils/visualizations/utils.py
    """
    assert 0 <= quality <= 10
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    video_name = video_name.replace(" ", "_").replace("\n", "_") + ".mp4"
    output_path = os.path.join(output_dir, video_name)
    writer = imageio.get_writer(output_path, fps=fps, quality=quality, **kwargs)
    if verbose:
        print(f"Video created: {output_path}")
        images_iter = tqdm.tqdm(images)
    else:
        images_iter = images
    for im in images_iter:
        writer.append_data(im)
    writer.close()


def normalize_depth(depth, min_depth=0, max_depth=None):
    if min_depth is None:
        min_depth = np.min(depth)
    if max_depth is None:
        max_depth = np.max(depth)
    depth = (depth - min_depth) / (max_depth - min_depth)
    depth = np.clip(depth, 0, 1)
    return depth


def observations_to_images(observations, max_depth=None) -> List[np.ndarray]:
    """Parse images from camera observations."""
    images = []
    for sensor_name in observations:
        if "rgb" in sensor_name:
            rgb = observations[sensor_name]
            if rgb.dtype == np.float32:
                rgb = np.clip(rgb * 255, 0, 255).astype(np.uint8)
            images.append(rgb)
        elif "depth" in sensor_name:
            depth = observations[sensor_name]  # [H, W, 1]
            depth = visualize_depth(depth)[..., ::-1]
            # depth = normalize_depth(depth, max_depth=max_depth)
            # depth = np.clip(depth * 255, 0, 255).astype(np.uint8)
            # depth = np.repeat(depth, 3, axis=-1)
            images.append(depth)
        elif "seg" in sensor_name:
            seg = observations[sensor_name]  # [H, W]
            seg = (seg * 255).astype(np.uint8)
            seg = np.repeat(seg[..., np.newaxis], 3, axis=-1)
            images.append(seg)
    return images


def tile_images(images: List[np.ndarray]) -> np.ndarray:
    """Tile multiple images to a single image. Support non-equal size."""
    # Sort images in descending order of vertical height
    images = sorted(images, key=lambda x: x.shape[0], reverse=True)

    columns = []
    max_h = images[0].shape[0]
    cur_h = 0
    cur_w = images[0].shape[1]

    # Arrange images in columns from left to right
    column = []
    for im in images:
        if cur_h + im.shape[0] <= max_h and cur_w == im.shape[1]:
            column.append(im)
            cur_h += im.shape[0]
        else:
            columns.append(column)
            column = [im]
            cur_h, cur_w = im.shape[0:2]
    columns.append(column)

    # Tile columns
    total_width = sum(x[0].shape[1] for x in columns)
    output_image = np.zeros((max_h, total_width, 3), dtype=images[0].dtype)
    cur_x = 0
    for column in columns:
        cur_w = column[0].shape[1]
        next_x = cur_x + cur_w
        column_image = np.concatenate(column, axis=0)
        cur_h = column_image.shape[0]
        output_image[:cur_h, cur_x:next_x] = column_image
        cur_x = next_x
    return output_image


def put_text_on_image(image: np.ndarray, lines: List[str]):
    assert image.dtype == np.uint8, image.dtype
    image = image.copy()

    font_size = 0.5
    font_thickness = 1
    font = cv2.FONT_HERSHEY_SIMPLEX

    y = 0
    for line in lines:
        textsize = cv2.getTextSize(line, font, font_size, font_thickness)[0]
        y += textsize[1] + 10
        x = 10
        cv2.putText(
            image,
            line,
            (x, y),
            font,
            font_size,
            (0, 255, 0),
            font_thickness,
            lineType=cv2.LINE_AA,
        )
    return image


def append_text_to_image(image: np.ndarray, lines: List[str]):
    r"""Appends text left to an image of size (height, width, channels).
    The returned image has white text on a black background.

    Args:
        image: the image to put text
        text: a string to display

    Returns:
        A new image with text inserted left to the input image

    See also:
        habitat.utils.visualization.utils
    """
    # h, w, c = image.shape
    font_size = 0.5
    font_thickness = 1
    font = cv2.FONT_HERSHEY_SIMPLEX
    blank_image = np.zeros(image.shape, dtype=np.uint8)

    y = 0
    for line in lines:
        textsize = cv2.getTextSize(line, font, font_size, font_thickness)[0]
        y += textsize[1] + 10
        x = 10
        cv2.putText(
            blank_image,
            line,
            (x, y),
            font,
            font_size,
            (255, 255, 255),
            font_thickness,
            lineType=cv2.LINE_AA,
        )
    # text_image = blank_image[0 : y + 10, 0:w]
    # final = np.concatenate((image, text_image), axis=0)
    final = np.concatenate((blank_image, image), axis=1)
    return final


def put_info_on_image(image, info: Dict[str, float], extras=None, overlay=True):
    lines = [f"{k}: {v:.3f}" for k, v in info.items()]
    if extras is not None:
        lines.extend(extras)
    if overlay:
        return put_text_on_image(image, lines)
    else:
        return append_text_to_image(image, lines)

arti_mani/utils/visualization.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from collections import OrderedDict
from copy import deepcopy
from typing import *

import cv2
import numpy as np
from arti_mani.utils.common import (
    clip_and_scale_action,
    convert_np_bool_to_float,
    normalize_action,
    normalize_action_space,
)
from arti_mani.utils.o3d_utils import pcd_voxel_down_sample_with_crop
from gym import ActionWrapper, ObservationWrapper, Wrapper, spaces


class NormalizeActionWrapper(ActionWrapper):
    def __init__(self, env):
        super().__init__(env)
        self.action_space = normalize_action_space(self.env.action_space)

    def action(self, action):
        return clip_and_scale_action(
            action, self.env.action_space.low, self.env.action_space.high
        )

    def reverse_action(self, action):
        return normalize_action(
            action, self.env.action_space.low, self.env.action_space.high
        )


def put_texts_on_image(image: np.ndarray, lines: List[str]):
    assert image.dtype == np.uint8, image.dtype
    image = image.copy()

    font_size = 0.4
    font_thickness = 1
    font = cv2.FONT_HERSHEY_SIMPLEX

    y = 0
    for line in lines:
        textsize = cv2.getTextSize(line, font, font_size, font_thickness)[0]
        y += textsize[1] + 10
        x = 10
        cv2.putText(
            image,
            line,
            (x, y),
            font,
            font_size,
            (255, 0, 0),
            font_thickness,
            lineType=cv2.LINE_AA,
        )
    return image


def put_infos_on_image(image, info: Dict[str, np.ndarray], overlay=True):
    lines = [f"{k}: {v.round(3)}" for k, v in info.items()]
    return put_texts_on_image(image, lines)


class RenderInfoWrapper(Wrapper):
    def step(self, action):
        obs, rew, done, info = super().step(action)
        if "TimeLimit.truncated" in info.keys():
            info["TimeLimit.truncated"] = convert_np_bool_to_float(
                info["TimeLimit.truncated"]
            )
        self._info_for_render = info
        return obs, rew, done, info

    def reset(self, **kwargs):
        obs = super().reset(**kwargs)
        # self._info_for_render = self.env.get_info()
        self._info_for_render = {}
        return obs

    def render(self, mode="rgb_array", **kwargs):
        if mode == "rgb_array":
            img = super().render(mode=mode, **kwargs)
            img = (img * 255).astype(np.uint8)
            return put_infos_on_image(img, self._info_for_render, overlay=True)
        elif mode == "cameras":
            img = super().render(mode=mode, **kwargs)
            return put_infos_on_image(img, self._info_for_render, overlay=True)
        else:
            return super().render(mode=mode, **kwargs)


class ActionControlWrapper(ActionWrapper):
    def __init__(self, env, control_mode=None):
        super().__init__(env)
        if control_mode is None:
            control_mode = self.env.unwrapped.agent.control_mode
        else:
            self.env.unwrapped.agent.set_control_mode(control_mode)
        self._control_mode = control_mode
        self.action_space = self.env.action_space[self._control_mode]

    def reset(self, *args, **kwargs):
        ret = super().reset(*args, **kwargs)
        self.env.unwrapped.agent.set_control_mode(self._control_mode)
        return ret

    def action(self, action):
        return {"control_mode": self._control_mode, "action": action}

    @property
    def control_mode(self):
        return self._control_mode


class ObsWrapper(ObservationWrapper):
    def __init__(self, env):
        super().__init__(env)
        if self.obs_mode in ["state", "state_dict", "state_depth"]:
            pass
        else:
            raise NotImplementedError()

    def observation(self, observation):
        from arti_mani.utils.common import flatten_state_dict

        if self.obs_mode == "state":
            return observation
        elif self.obs_mode == "state_dict":
            obs = observation
            obs.pop("task", None)
            obs.pop("articulation", None)
            return obs
        elif self.obs_mode == "state_depth":
            obs = observation
            ### process uvz
            base_target_uvz = obs.pop("base_target_uvz", None)
            base_target_uvz[:2] = base_target_uvz[:2] / 4
            base_target_uvz[2] = 39 * (base_target_uvz[2] - 0.5) / (2 - 0.5)
            base_target_uvz = np.round(base_target_uvz).astype(np.int_)
            obs["zvu"] = deepcopy(base_target_uvz[:, [2, 1, 0]])  # (3, )
            ### process depth
            depth = obs.pop("base_depth", None)
            # depth = depth.transpose(2,0,1)[:, ::4, ::4]  # (1, H, W)
            depth = depth.transpose(2, 0, 1).squeeze()[::4, ::4]
            depth = (depth - np.mean(depth)) / np.std(depth)
            obs["depth"] = deepcopy(depth)  # (H//4, W//4)
            return obs
        else:
            raise NotImplementedError()


class ManiObsWrapper(ObservationWrapper):
    def __init__(self, env):
        super().__init__(env)
        if self.obs_mode in ["state", "state_dict"]:
            pass
        else:
            raise NotImplementedError()

    def observation(self, observation):
        from arti_mani.utils.common import flatten_state_dict

        if self.obs_mode == "state":
            return observation
        elif self.obs_mode == "state_dict":
            obs = observation
            obs.pop("task", None)
            obs.pop("articulation", None)
            return obs
        else:
            raise NotImplementedError()

    @property
    def _max_episode_steps(self):
        return self.env.unwrapped._max_episode_steps

    def render(self, mode="human", *args, **kwargs):
        if mode == "human":
            self.env.render(mode, *args, **kwargs)
            return

        if mode in ["rgb_array", "color_image"]:
            img = self.env.render(mode="rgb_array", *args, **kwargs)
        else:
            img = self.env.render(mode=mode, *args, **kwargs)
        if isinstance(img, dict):
            if "world" in img:
                img = img["world"]
            elif "main" in img:
                img = img["main"]
            else:
                print(img.keys())
                exit(0)
        if isinstance(img, dict):
            img = img["rgb"]
        if img.ndim == 4:
            assert img.shape[0] == 1
            img = img[0]
        if img.dtype in [np.float32, np.float64]:
            img = np.clip(img, a_min=0, a_max=1) * 255
        img = img[..., :3]
        img = img.astype(np.uint8)
        return img


class EgoDepthCutoff(ObservationWrapper):
    """For Egocentric camera realsense D415, depth should be cutoff due to real sensor range."""

    def __init__(self, env, depth_cut=False, cutlen=0.225):
        ### default cutlen: sensor range(0.45)-grasp_site_z(0.225)
        super().__init__(env)
        self.depth_cut = depth_cut
        self.cutlen = cutlen
        # self.grasp_z = self.env.grasp_site.pose.p[2]
        # self.gripper_pose = self.env.grasp_site.pose

    def observation(self, observation):
        obs = observation
        if (self.obs_mode == "depth" or self.obs_mode == "rgbd") and self.depth_cut:
            if self.env.grasp_site.pose.p[2] < self.cutlen:
                obs["depth"][..., 0] = np.zeros_like(obs["depth"][..., 0])
        return obs


class PointCloudPreprocessObsWrapper(ObservationWrapper):
    """Preprocess point cloud, crop and voxel downsample"""

    _num_pts: int
    _vox_size: float

    def __init__(
        self,
        env,
        num_pts=2048,
        vox_size=0.003,
        min_bound=np.array([-1, -1, 1e-3]),
        max_bound=np.array([1, 1, 1]),
    ):
        super().__init__(env)
        self._num_pts = num_pts
        self._vox_size = vox_size
        self._min_bound = min_bound
        self._max_bound = max_bound

    def observation(self, observation):
        if not self.obs_mode == "pointcloud":
            return observation

        obs = observation
        pointcloud = obs["pointcloud"]
        xyz = pointcloud["xyz"]
        sample_indices = pcd_voxel_down_sample_with_crop(
            xyz, self._vox_size, self._min_bound, self._max_bound
        )
        if len(sample_indices) >= self._num_pts:
            sample_indices = np.random.choice(
                sample_indices, self._num_pts, replace=False
            )
        else:
            ext_sample_indices = np.random.choice(
                sample_indices, (self._num_pts - len(sample_indices))
            )
            sample_indices.extend(ext_sample_indices)
        for k, v in pointcloud.items():
            pointcloud[k] = v[sample_indices]
        obs["pointcloud"] = pointcloud
        return obs

arti_mani/utils/wrappers.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation
from arti_mani import __version__
from setuptools import setup

setup(
    name="arti_mani",
    version=__version__,
    author="Xpw",
    zip_safe=False,
    packages=["arti_mani"],
)

setup.py
THU-VCLab/Part-Guided-3D-RL-for-Sim2Real-Articulated-Object-Manipulation