
.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/simulator/example_curi_cube_ik_osc.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_simulator_example_curi_cube_ik_osc.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_simulator_example_curi_cube_ik_osc.py:


CURI cube pick
==============

Use Jacobian matrix and inverse kinematics control of curi robot to pick up a box.
Damped Least Squares method from: https://www.math.ucsd.edu/~sbuss/ResearchWeb/ikmethods/iksurvey.pdf

.. GENERATED FROM PYTHON SOURCE LINES 8-423

.. code-block:: default


    from isaacgym import gymapi
    from isaacgym import gymutil
    from isaacgym import gymtorch
    from isaacgym.torch_utils import *

    import math
    import numpy as np
    import torch
    import random
    import time

    DOF = 25


    def quat_axis(q, axis=0):
        basis_vec = torch.zeros(q.shape[0], 3, device=q.device)
        basis_vec[:, axis] = 1
        return quat_rotate(q, basis_vec)


    def orientation_error(desired, current):
        cc = quat_conjugate(current)
        q_r = quat_mul(desired, cc)
        return q_r[:, 0:3] * torch.sign(q_r[:, 3]).unsqueeze(-1)


    def cube_grasping_yaw(q, corners):
        """ returns horizontal rotation required to grasp cube """
        rc = quat_rotate(q, corners)
        yaw = (torch.atan2(rc[:, 1], rc[:, 0]) - 0.25 * math.pi) % (0.5 * math.pi)
        theta = 0.5 * yaw
        w = theta.cos()
        x = torch.zeros_like(w)
        y = torch.zeros_like(w)
        z = theta.sin()
        yaw_quats = torch.stack([x, y, z, w], dim=-1)
        return yaw_quats


    def control_ik(dpose):
        global damping, j_eef, num_envs
        # solve damped least squares
        j_eef_T = torch.transpose(j_eef, 1, 2)
        lmbda = torch.eye(6, device=device) * (damping ** 2)
        u = (j_eef_T @ torch.inverse(j_eef @ j_eef_T + lmbda) @ dpose).view(num_envs, 7)
        return u


    def control_osc(dpose):
        global kp, kd, kp_null, kd_null, default_dof_pos_tensor, mm, j_eef, num_envs, dof_pos, dof_vel, hand_vel
        mm_inv = torch.inverse(mm)
        m_eef_inv = j_eef @ mm_inv @ torch.transpose(j_eef, 1, 2)
        m_eef = torch.inverse(m_eef_inv)
        u = torch.transpose(j_eef, 1, 2) @ m_eef @ (
                kp * dpose - kd * hand_vel.unsqueeze(-1))

        # Nullspace control torques `u_null` prevents large changes in joint configuration
        # They are added into the nullspace of OSC so that the end effector orientation remains constant
        # roboticsproceedings.org/rss07/p31.pdf
        j_eef_inv = m_eef @ j_eef @ mm_inv
        u_null = kd_null * -dof_vel + kp_null * (
                (default_dof_pos_tensor.view(1, -1, 1) - dof_pos + np.pi) % (2 * np.pi) - np.pi)
        u_null = u_null[:, 7:14]
        u_null = mm @ u_null
        u += (torch.eye(7, device=device).unsqueeze(0) - torch.transpose(j_eef, 1, 2) @ j_eef_inv) @ u_null
        return u.squeeze(-1)


    # set random seed
    np.random.seed(42)

    torch.set_printoptions(precision=4, sci_mode=False)

    # acquire gym interface
    gym = gymapi.acquire_gym()

    # parse arguments

    # Add custom arguments
    custom_parameters = [
        {"name": "--controller", "type": str, "default": "osc",
         "help": "Controller to use for curi. Options are {ik, osc}"},
        {"name": "--num_envs", "type": int, "default": 1, "help": "Number of environments to create"},
    ]
    args = gymutil.parse_arguments(
        description="curi Jacobian Inverse Kinematics (IK) + Operational Space Control (OSC) Example",
        custom_parameters=custom_parameters,
    )

    # Grab controller
    controller = args.controller
    assert controller in {"ik", "osc"}, f"Invalid controller specified -- options are (ik, osc). Got: {controller}"

    # set torch device
    device = args.sim_device if args.use_gpu_pipeline else 'cpu'

    # configure sim
    sim_params = gymapi.SimParams()
    sim_params.up_axis = gymapi.UP_AXIS_Z
    sim_params.gravity = gymapi.Vec3(0.0, 0.0, -9.8)
    sim_params.dt = 1.0 / 60.0
    sim_params.substeps = 2
    sim_params.use_gpu_pipeline = args.use_gpu_pipeline
    if args.physics_engine == gymapi.SIM_PHYSX:
        sim_params.physx.solver_type = 1
        sim_params.physx.num_position_iterations = 8
        sim_params.physx.num_velocity_iterations = 1
        sim_params.physx.rest_offset = 0.0
        sim_params.physx.contact_offset = 0.001
        sim_params.physx.friction_offset_threshold = 0.001
        sim_params.physx.friction_correlation_distance = 0.0005
        sim_params.physx.num_threads = args.num_threads
        sim_params.physx.use_gpu = args.use_gpu
    else:
        raise Exception("This example can only be used with PhysX")

    # Set controller parameters
    # IK params
    damping = 0.05

    # OSC params
    kp = 150.
    kd = 2.0 * np.sqrt(kp)
    kp_null = 10.
    kd_null = 2.0 * np.sqrt(kp_null)

    # create sim
    sim = gym.create_sim(args.compute_device_id, args.graphics_device_id, args.physics_engine, sim_params)
    if sim is None:
        raise Exception("Failed to create sim")

    # create viewer
    viewer = gym.create_viewer(sim, gymapi.CameraProperties())
    if viewer is None:
        raise Exception("Failed to create viewer")

    asset_root = "../assets"

    # create table asset
    table_dims = gymapi.Vec3(0.6, 2.5, 0.5)
    asset_options = gymapi.AssetOptions()
    asset_options.fix_base_link = True
    table_asset = gym.create_box(sim, table_dims.x, table_dims.y, table_dims.z, asset_options)

    # create box asset
    box_size = 0.045
    asset_options = gymapi.AssetOptions()
    box_asset = gym.create_box(sim, box_size, box_size, box_size, asset_options)

    # load curi asset
    curi_asset_file = "urdf/curi/urdf/curi_isaacgym.urdf"
    asset_options = gymapi.AssetOptions()
    asset_options.armature = 0.01
    asset_options.fix_base_link = True
    asset_options.disable_gravity = True
    asset_options.flip_visual_attachments = True
    curi_asset = gym.load_asset(sim, asset_root, curi_asset_file, asset_options)

    # configure curi dofs
    curi_dof_props = gym.get_asset_dof_properties(curi_asset)
    curi_lower_limits = curi_dof_props["lower"]
    curi_upper_limits = curi_dof_props["upper"]
    curi_ranges = curi_upper_limits - curi_lower_limits
    curi_mids = 0.3 * (curi_upper_limits + curi_lower_limits)

    # use position drive for all dofs
    if controller == "ik":
        curi_dof_props["driveMode"][7:].fill(gymapi.DOF_MODE_POS)
        curi_dof_props["stiffness"][7:].fill(300.0)
        curi_dof_props["damping"][7:].fill(80.0)
    else:  # osc
        curi_dof_props["driveMode"][7:].fill(gymapi.DOF_MODE_EFFORT)
        curi_dof_props["stiffness"][7:].fill(0.0)
        curi_dof_props["damping"][7:].fill(0.0)
    # grippers
    curi_dof_props["driveMode"][14:16].fill(gymapi.DOF_MODE_POS)
    curi_dof_props["stiffness"][14:16].fill(800.0)
    curi_dof_props["damping"][14:16].fill(40.0)
    curi_dof_props["driveMode"][23:25].fill(gymapi.DOF_MODE_POS)
    curi_dof_props["stiffness"][23:25].fill(800.0)
    curi_dof_props["damping"][23:25].fill(40.0)

    # default dof states and position targets
    curi_num_dofs = gym.get_asset_dof_count(curi_asset)
    default_dof_pos = np.zeros(curi_num_dofs, dtype=np.float32)
    default_dof_pos[7:] = curi_mids[7:]
    # grippers open
    default_dof_pos[14:16] = curi_upper_limits[14:16]
    default_dof_pos[23:25] = curi_upper_limits[23:25]

    default_dof_state = np.zeros(curi_num_dofs, gymapi.DofState.dtype)
    default_dof_state["pos"] = default_dof_pos

    # send to torch
    default_dof_pos_tensor = to_torch(default_dof_pos, device=device)

    # get link index of panda hand, which we will use as end effector
    curi_link_dict = gym.get_asset_rigid_body_dict(curi_asset)
    curi_hand_index = curi_link_dict["panda_left_hand"]

    # configure env grid
    num_envs = args.num_envs
    num_per_row = int(math.sqrt(num_envs))
    spacing = 1.0
    env_lower = gymapi.Vec3(-spacing, -spacing, 0.0)
    env_upper = gymapi.Vec3(spacing, spacing, spacing)
    print("Creating %d environments" % num_envs)

    curi_pose = gymapi.Transform()
    curi_pose.p = gymapi.Vec3(0, 0, 0)

    table_pose = gymapi.Transform()
    table_pose.p = gymapi.Vec3(1, 0.0, 0.5 * table_dims.z)

    box_pose = gymapi.Transform()

    envs = []
    box_idxs = []
    hand_idxs = []
    init_pos_list = []
    init_rot_list = []

    # add ground plane
    plane_params = gymapi.PlaneParams()
    plane_params.normal = gymapi.Vec3(0, 0, 1)
    gym.add_ground(sim, plane_params)

    for i in range(num_envs):
        # create env
        env = gym.create_env(sim, env_lower, env_upper, num_per_row)
        envs.append(env)

        # add table
        table_handle = gym.create_actor(env, table_asset, table_pose, "table", i, 0)

        # add box
        box_pose.p.x = table_pose.p.x + np.random.uniform(-0.2, 0.1)
        box_pose.p.y = table_pose.p.y + np.random.uniform(-0.3, 0.3)
        box_pose.p.z = table_dims.z + 0.5 * box_size
        box_pose.r = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 0, 1), np.random.uniform(-math.pi, math.pi))
        box_handle = gym.create_actor(env, box_asset, box_pose, "box", i, 0)
        color = gymapi.Vec3(np.random.uniform(0, 1), np.random.uniform(0, 1), np.random.uniform(0, 1))
        gym.set_rigid_body_color(env, box_handle, 0, gymapi.MESH_VISUAL_AND_COLLISION, color)

        # get global index of box in rigid body state tensor
        box_idx = gym.get_actor_rigid_body_index(env, box_handle, 0, gymapi.DOMAIN_SIM)
        box_idxs.append(box_idx)

        # add curi
        curi_handle = gym.create_actor(env, curi_asset, curi_pose, "curi", i, 2)

        # set dof properties
        gym.set_actor_dof_properties(env, curi_handle, curi_dof_props)

        # set initial dof states
        gym.set_actor_dof_states(env, curi_handle, default_dof_state, gymapi.STATE_ALL)

        # set initial position targets
        gym.set_actor_dof_position_targets(env, curi_handle, default_dof_pos)

        # get inital hand pose
        hand_handle = gym.find_actor_rigid_body_handle(env, curi_handle, "panda_left_hand")
        hand_pose = gym.get_rigid_transform(env, hand_handle)
        init_pos_list.append([hand_pose.p.x, hand_pose.p.y, hand_pose.p.z])
        init_rot_list.append([hand_pose.r.x, hand_pose.r.y, hand_pose.r.z, hand_pose.r.w])

        # get global index of hand in rigid body state tensor
        hand_idx = gym.find_actor_rigid_body_index(env, curi_handle, "panda_left_hand", gymapi.DOMAIN_SIM)
        hand_idxs.append(hand_idx)

    # point camera at middle env
    cam_pos = gymapi.Vec3(4, 3, 2)
    cam_target = gymapi.Vec3(-4, -3, 0)
    middle_env = envs[num_envs // 2 + num_per_row // 2]
    gym.viewer_camera_look_at(viewer, middle_env, cam_pos, cam_target)

    # ==== prepare tensors =====
    # from now on, we will use the tensor API that can run on CPU or GPU
    gym.prepare_sim(sim)

    # initial hand position and orientation tensors
    init_pos = torch.Tensor(init_pos_list).view(num_envs, 3).to(device)
    init_rot = torch.Tensor(init_rot_list).view(num_envs, 4).to(device)

    # hand orientation for grasping
    down_q = torch.stack(num_envs * [torch.tensor([1.0, 0.0, 0.0, 0.0])]).to(device).view((num_envs, 4))

    # box corner coords, used to determine grasping yaw
    box_half_size = 0.5 * box_size
    corner_coord = torch.Tensor([box_half_size, box_half_size, box_half_size])
    corners = torch.stack(num_envs * [corner_coord]).to(device)

    # downard axis
    down_dir = torch.Tensor([0, 0, -1]).to(device).view(1, 3)

    # get jacobian tensor
    # for fixed-base curi, tensor has shape (num envs, 10, 6, 9)
    _jacobian = gym.acquire_jacobian_tensor(sim, "curi")
    jacobian = gymtorch.wrap_tensor(_jacobian)

    # jacobian entries corresponding to curi hand
    j_eef = jacobian[:, curi_hand_index - 1, :, 7:14]

    # get mass matrix tensor
    _massmatrix = gym.acquire_mass_matrix_tensor(sim, "curi")
    mm = gymtorch.wrap_tensor(_massmatrix)
    mm = mm[:, 7:14, 7:14]  # only need elements corresponding to the curi arm

    # get rigid body state tensor
    _rb_states = gym.acquire_rigid_body_state_tensor(sim)
    rb_states = gymtorch.wrap_tensor(_rb_states)

    # get dof state tensor
    _dof_states = gym.acquire_dof_state_tensor(sim)
    dof_states = gymtorch.wrap_tensor(_dof_states)
    dof_pos = dof_states[:, 0].view(num_envs, DOF, 1)
    dof_vel = dof_states[:, 1].view(num_envs, DOF, 1)

    # Create a tensor noting whether the hand should return to the initial position
    hand_restart = torch.full([num_envs], False, dtype=torch.bool).to(device)

    # Set action tensors
    pos_action = torch.zeros_like(dof_pos).squeeze(-1)
    effort_action = torch.zeros_like(pos_action)

    # simulation loop
    while not gym.query_viewer_has_closed(viewer):

        # step the physics
        gym.simulate(sim)
        gym.fetch_results(sim, True)

        # refresh tensors
        gym.refresh_rigid_body_state_tensor(sim)
        gym.refresh_dof_state_tensor(sim)
        gym.refresh_jacobian_tensors(sim)
        gym.refresh_mass_matrix_tensors(sim)

        box_pos = rb_states[box_idxs, :3]
        box_rot = rb_states[box_idxs, 3:7]

        hand_pos = rb_states[hand_idxs, :3]
        hand_rot = rb_states[hand_idxs, 3:7]
        hand_vel = rb_states[hand_idxs, 7:]

        to_box = box_pos - hand_pos
        box_dist = torch.norm(to_box, dim=-1).unsqueeze(-1)
        box_dir = to_box / box_dist
        box_dot = box_dir @ down_dir.view(3, 1)

        # how far the hand should be from box for grasping
        grasp_offset = 0.11 if controller == "ik" else 0.10

        # determine if we're holding the box (grippers are closed and box is near)
        gripper_sep = dof_pos[:, 14] + dof_pos[:, 15]
        gripped = (gripper_sep < 0.045) & (box_dist < grasp_offset + 0.5 * box_size)

        yaw_q = cube_grasping_yaw(box_rot, corners)
        box_yaw_dir = quat_axis(yaw_q, 0)
        hand_yaw_dir = quat_axis(hand_rot, 0)
        yaw_dot = torch.bmm(box_yaw_dir.view(num_envs, 1, 3), hand_yaw_dir.view(num_envs, 3, 1)).squeeze(-1)

        # determine if we have reached the initial position; if so allow the hand to start moving to the box
        to_init = init_pos - hand_pos
        init_dist = torch.norm(to_init, dim=-1)
        hand_restart = (hand_restart & (init_dist > 0.02)).squeeze(-1)
        return_to_start = (hand_restart | gripped.squeeze(-1)).unsqueeze(-1)

        # if hand is above box, descend to grasp offset
        # otherwise, seek a position above the box
        above_box = ((box_dot >= 0.99) & (yaw_dot >= 0.95) & (box_dist < grasp_offset * 3)).squeeze(-1)
        grasp_pos = box_pos.clone()
        grasp_pos[:, 2] = torch.where(above_box, box_pos[:, 2] + grasp_offset, box_pos[:, 2] + grasp_offset * 2.5)

        # compute goal position and orientation
        goal_pos = torch.where(return_to_start, init_pos, grasp_pos)
        goal_rot = torch.where(return_to_start, init_rot, quat_mul(down_q, quat_conjugate(yaw_q)))

        # compute position and orientation error
        pos_err = goal_pos - hand_pos
        orn_err = orientation_error(goal_rot, hand_rot)
        dpose = torch.cat([pos_err, orn_err], -1).unsqueeze(-1)
        import rofunc as rf
        rf.logger.beauty_print("pos_err: {}".format(pos_err), type="info")
        rf.logger.beauty_print("orn_err: {}".format(orn_err), type="info")

        # Deploy control based on type
        if controller == "ik":
            pos_action[:, 7:14] = dof_pos.squeeze(-1)[:, 7:14] + control_ik(dpose)
        else:  # osc
            effort_action[:, 7:14] = control_osc(dpose)

        # gripper actions depend on distance between hand and box
        close_gripper = (box_dist < grasp_offset + 0.02) | gripped
        # always open the gripper above a certain height, dropping the box and restarting from the beginning
        hand_restart = hand_restart | (box_pos[:, 2] > 0.6)
        keep_going = torch.logical_not(hand_restart)
        close_gripper = close_gripper & keep_going.unsqueeze(-1)
        grip_acts = torch.where(close_gripper, torch.Tensor([[0., 0.]] * num_envs).to(device),
                                torch.Tensor([[0.04, 0.04]] * num_envs).to(device))
        pos_action[:, 14:16] = grip_acts

        # Deploy actions
        gym.set_dof_position_target_tensor(sim, gymtorch.unwrap_tensor(pos_action))
        gym.set_dof_actuation_force_tensor(sim, gymtorch.unwrap_tensor(effort_action))

        # update viewer
        gym.step_graphics(sim)
        gym.draw_viewer(viewer, sim, False)
        gym.sync_frame_time(sim)

    # cleanup
    gym.destroy_viewer(viewer)
    gym.destroy_sim(sim)


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 0.000 seconds)


.. _sphx_glr_download_examples_simulator_example_curi_cube_ik_osc.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example




    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: example_curi_cube_ik_osc.py <example_curi_cube_ik_osc.py>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: example_curi_cube_ik_osc.ipynb <example_curi_cube_ik_osc.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_
