HiroIshida / snippets

fraction of codes which may be grepped later
6 stars 3 forks source link

skrobot todo #17

Open HiroIshida opened 3 years ago

HiroIshida commented 3 years ago

2.76 sec

HiroIshida commented 3 years ago
#!/usr/bin/env python

import skrobot
import numpy as np
import time
from skrobot.model import Box, MeshLink, Axis
from skrobot.sdf import UnionSDF
from skrobot.planner.utils import get_robot_state, set_robot_state

if __name__ == '__main__':
    robot_model = skrobot.models.urdf.RobotModelFromURDF(
        urdf_file=skrobot.data.pr2_urdfpath())
    robot_model.init_pose()
    viewer = skrobot.viewers.TrimeshSceneViewer(resolution=(640, 480))

    table = Box(extents=[0.7, 1.0, 0.05], with_sdf=True) 
    table.translate([0.8, 0.0, 0.65])

    top = Box(extents=[0.7, 1.0, 0.05], with_sdf=True) 
    top.translate([0.8, 0.0, 1.1])
    viewer.add(top)

    pillar1 = Box(extents=[0.05, 0.05, 2.0], with_sdf=True) 
    pillar1.translate([0.65, 0.0, 1.0])
    viewer.add(pillar1)

    pillar2 = Box(extents=[0.05, 0.05, 2.0], with_sdf=True) 
    pillar2.translate([0.6, -0.4, 1.0])
    viewer.add(pillar2)

    pillar3 = Box(extents=[0.05, 0.05, 2.0], with_sdf=True) 
    pillar3.translate([0.9, -0.2, 1.0])
    viewer.add(pillar3)

    union_sdf = UnionSDF([table.sdf, top.sdf, pillar1.sdf, pillar2.sdf])

    viewer.add(table)

    link_idx_table = {}
    for link_idx in range(len(robot_model.link_list)):
        name = robot_model.link_list[link_idx].name
        link_idx_table[name] = link_idx

    link_names = ["torso_lift_link", "r_shoulder_pan_link", "r_shoulder_lift_link",
                  "r_upper_arm_roll_link", "r_elbow_flex_link",
                  "r_forearm_roll_link", "r_wrist_flex_link",
                  "r_wrist_roll_link"]

    link_list = [robot_model.link_list[link_idx_table[lname]]
                 for lname in link_names]
    joint_list = [link.joint for link in link_list]

    coll_link_names = ["r_upper_arm_link", "r_forearm_link", "r_gripper_palm_link", "r_gripper_r_finger_link", "r_gripper_l_finger_link"]

    coll_link_list = [robot_model.link_list[link_idx_table[lname]]
                 for lname in coll_link_names]

    rarm_end_coords = skrobot.coordinates.CascadedCoords(
        parent=robot_model.r_gripper_tool_frame)
    forarm_coords = skrobot.coordinates.CascadedCoords(
        parent=robot_model.r_forearm_link)

    # set initial angle vector of rarm
    av_init = [0.33, -0.8, 0.35, -0.74, -0.70, -0.17, -0.63, 0.0, -0.25, 0, 0]
    set_robot_state(robot_model, joint_list, av_init, base_also=True)

    target_coords = skrobot.coordinates.Coordinates(
        [0.8, 0.0, 0.9], [0, 0, 0])
    axis = Axis.from_coords(target_coords)
    viewer.add(axis)

    margin = 0.01
    def sdf_with_margin(X):
        return union_sdf(X) - margin

    base_also = True
    av_start = get_robot_state(robot_model, joint_list, base_also=base_also) 

    robot_model.inverse_kinematics_slsqp(
            target_coords, link_list, rarm_end_coords, 
            rot_also=False, base_also=base_also)
    robot_model.inverse_kinematics_slsqp(
            target_coords, link_list, rarm_end_coords, 
            coll_link_list, sdf_with_margin,
            rot_also=False, base_also=base_also)
    av_goal = get_robot_state(robot_model, joint_list, base_also=base_also) 

    traj = robot_model.plan_trajectory(
            av_start, av_goal, link_list, coll_link_list, sdf_with_margin, 15, base_also=base_also)

    set_robot_state(robot_model, joint_list, av_start, base_also=True)
    viewer.add(robot_model)
    viewer.show()

    time.sleep(1.0)
    print("show trajectory")
    for av in traj:
        set_robot_state(robot_model, joint_list, av, base_also=base_also)
        viewer.redraw()
        time.sleep(0.2)