roboimi/roboimi/demos/diana_air_insert_policy.py

import numpy as np
from pyquaternion import Quaternion

from roboimi.demos.diana_policy import PolicyBase


class TestAirInsertPolicy(PolicyBase):
    ACTION_OBJECT_Z_OFFSET = 0.078
    SOCKET_GRASP_OFFSET = np.array([0.0, 0.0, 0.0], dtype=np.float64)
    PEG_GRASP_OFFSET = np.array([0.0, 0.0, 0.0], dtype=np.float64)
    SOCKET_OUTER_GRASP_STRATEGY = "socket_outer"
    LEGACY_GRASP_STRATEGY = "legacy"
    SOCKET_HOLD_Z = 0.85
    PEG_INSERT_START_OFFSET = np.array([0.105, 0.0, 0.0], dtype=np.float64)
    INSERT_END_T = 580
    LEFT_SOCKET_GRIPPER_CLOSED = -100
    RIGHT_PEG_GRIPPER_CLOSED = -100
    SOCKET_APPROACH_Z = 1.05
    EPISODE_END_T = 600

    def __init__(self, inject_noise=False, grasp_strategy=SOCKET_OUTER_GRASP_STRATEGY):
        super().__init__(inject_noise=inject_noise)
        valid_strategies = {
            self.SOCKET_OUTER_GRASP_STRATEGY,
            self.LEGACY_GRASP_STRATEGY,
        }
        if grasp_strategy not in valid_strategies:
            raise ValueError(
                f"Unsupported air insert grasp_strategy={grasp_strategy!r}; "
                f"expected one of {sorted(valid_strategies)}"
            )
        self.grasp_strategy = grasp_strategy

    def generate_trajectory(self, task_state):
        return self._generate_socket_peg_trajectory(task_state)

    def _generate_socket_peg_trajectory(self, task_state):
        socket_xyz = np.asarray(task_state["socket_pos"], dtype=np.float64)
        peg_xyz = np.asarray(task_state["peg_pos"], dtype=np.float64)

        init_mocap_pose_left = np.array(
            [
                -0.17297014,
                1.00485877,
                1.32773627,
                7.06825181e-01,
                8.20281078e-06,
                -7.07388269e-01,
                -5.20399313e-06,
            ],
            dtype=np.float64,
        )
        init_mocap_pose_right = np.array(
            [
                0.17297014,
                0.9951369,
                1.32773623,
                2.59463975e-06,
                7.07388269e-01,
                5.59551158e-06,
                7.06825181e-01,
            ],
            dtype=np.float64,
        )

        left_init_quat = Quaternion(init_mocap_pose_left[3:])
        right_init_quat = Quaternion(init_mocap_pose_right[3:])

        left_pick_quat = (
            left_init_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=45)
        ).elements
        right_pick_quat = (
            right_init_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=45)
        ).elements

        socket_hold_action = np.array(
            [socket_xyz[0] - 0.078, socket_xyz[1], self.SOCKET_HOLD_Z], dtype=np.float64
        )

        peg_init_xyz = peg_xyz + np.array(
            [0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET + 0.01]
        )
        peg_lift_center = np.array(
            [peg_xyz[0] + 0.078, socket_hold_action[1], self.SOCKET_HOLD_Z - 0.01],
            dtype=np.float64,
        )
        # The front camera looks along +Y, so visual right-to-left insertion is
        # world +X -> -X.  With the socket XML in identity orientation, its
        # tunnel axis is local/world X, so the peg approaches from +X and stops
        # when its leading face reaches the socket's internal pin.
        peg_insert_end_center = np.array(
            [
                socket_hold_action[0] + 0.078 * 2 + 0.04 + 0.06 - 0.01,
                socket_hold_action[1],
                self.SOCKET_HOLD_Z - 0.01,
            ],
            dtype=np.float64,
        )

        self.left_trajectory = [
            {
                "t": 1,
                "xyz": init_mocap_pose_left[:3],
                "quat": init_mocap_pose_left[3:],
                "gripper": 100,
            },
            {
                "t": 130,
                "xyz": socket_xyz
                + np.array([-0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET]),
                "quat": left_pick_quat,
                "gripper": 100,
            },
            {
                "t": 180,
                "xyz": socket_xyz
                + np.array([-0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET]),
                "quat": left_pick_quat,
                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
            },
            {
                "t": 350,
                "xyz": socket_hold_action,
                "quat": left_pick_quat,
                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
            },
            {
                "t": self.EPISODE_END_T,
                "xyz": socket_hold_action,
                "quat": left_pick_quat,
                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
            },
        ]

        self.right_trajectory = [
            {
                "t": 1,
                "xyz": init_mocap_pose_right[:3],
                "quat": init_mocap_pose_right[3:],
                "gripper": 100,
            },
            {
                "t": 80,
                "xyz": peg_init_xyz,
                "quat": right_pick_quat,
                "gripper": 100,
            },
            {
                "t": 150,
                "xyz": peg_init_xyz,
                "quat": right_pick_quat,
                "gripper": 100,
            },
            {
                "t": 180,
                "xyz": peg_init_xyz,
                "quat": right_pick_quat,
                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
            },
            {
                "t": 350,
                "xyz": peg_init_xyz,
                "quat": right_pick_quat,
                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
            },
            {
                "t": 450,
                "xyz": peg_lift_center,
                "quat": right_pick_quat,
                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
            },
            {
                "t": self.INSERT_END_T,
                "xyz": peg_insert_end_center,
                "quat": right_pick_quat,
                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
            },
            {
                "t": self.EPISODE_END_T,
                "xyz": peg_insert_end_center,
                "quat": right_pick_quat,
                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
            },
        ]