Enhancement: create urdf

[iory]

cc: @hiroya1224

import numpy as np

from skrobot.coordinates.math import random_quaternion
from skrobot.coordinates.math import random_translation
from skrobot.coordinates.math import quaternion2rpy
from skrobot.model import RobotModel
from skrobot.model import RotationalJoint
from skrobot.model import Cylinder
from skrobot.coordinates import Coordinates
from skrobot.viewers import TrimeshSceneViewer


class JointRelative(object):
    def __init__(self,
                 reference_jointname, target_jointname,
                 relative_position, relative_rotation):
        # frame_id of reference_joint and target_joint
        self.reference_jointname = reference_jointname
        self.target_jointname = target_jointname
        # target_joint's position and rotation w.r.t. reference_joint
        # 3D vector
        self.relative_position = relative_position
        # quaternion (wxyz order) or rotation
        self.relative_rotation = relative_rotation


joint_relatives = [
    JointRelative("joint{}".format(0), "joint{}".format(1),
                      [1, 0, 0],
                      [1, 0, 0, 0],
                      ),
    JointRelative("joint{}".format(1), "joint{}".format(2),
                      [1, 0, 0],
                      [1, 0, 0, 0],
                      )
]


def generate_urdf(joint_relatives):
    urdf_template = """
<robot name="simple_robot">
  {links}
  {joints}
</robot>
"""

    link_template = """
  <link name="{name}">
    <visual>
      <geometry>
        <cylinder length="{length}" radius="0.1"/>
      </geometry>
      <origin rpy="{rpy}" xyz="{xyz}"/>
    </visual>
  </link>
"""

    joint_template = """
  <joint name="{name}" type="revolute">
    <parent link="{parent}"/>
    <child link="{child}"/>
    <origin xyz="{xyz}" rpy="{rpy}"/>
    <axis xyz="0 1 0"/>
    <limit effort="10" lower="-3.141592653589793" upper="3.141592653589793" velocity="10" />
  </joint>
"""

    links = []
    joints = []

    # Add root link
    cylinder_rpy = '0.0 1.5707963267948966 0.0'
    links.append(link_template.format(name='base_link',
                                      length=0.01,
                                      xyz='0.01 0 0',
                                      rpy=cylinder_rpy))
    joints.append(joint_template.format(
        name=f"joint0",
        parent='base_link',
        child='link_joint0',
        xyz="0 0 0",
        rpy="0 0 0",
    ))

    for i, joint_relative in enumerate(joint_relatives):
        # Add link
        length = np.linalg.norm(joint_relative.relative_position)
        links.append(link_template.format(
            name='link_' + joint_relative.reference_jointname,
            length=length,
            xyz=" ".join(map(str, [length / 2.0, 0, 0])),
            # xyz=" ".join(map(str, [0, 0, 0])),
            rpy=cylinder_rpy))

        # Add joint
        parent_link = 'link_' + joint_relative.reference_jointname
        joints.append(joint_template.format(
            name=f"joint{i + 1}",
            parent=parent_link,
            child='link_' + joint_relative.target_jointname,
            xyz=" ".join(map(str, joint_relative.relative_position)),
            rpy=" ".join(map(str, quaternion2rpy(joint_relative.relative_rotation)[0])),
        ))

    links.append(link_template.format(
        name='link_' + joint_relative.target_jointname,
        length=0.01,
        xyz="0.005 0 0",
        rpy=cylinder_rpy))

    urdf_str = urdf_template.format(links="\n".join(links), joints="\n".join(joints))
    return urdf_str


urdf_str = generate_urdf(joint_relatives)
out_urdfpath = '/tmp/robot.urdf'
with open(out_urdfpath, 'w') as f:
    f.write(urdf_str)

robot_model = RobotModel()
robot_model.load_urdf_file(open(out_urdfpath))
viewer = TrimeshSceneViewer()
viewer.add(robot_model)
viewer.show()
robot_model.joint0.joint_angle(0.4)
robot_model.joint1.joint_angle(-0.5)

robot_model.joint0.joint_angle(0.4)
robot_model.joint1.joint_angle(-0.5)
robot_model.inverse_kinematics(
    target_coords=Coordinates(pos=[1.9, 0, 0]),
    rotation_axis=False,
    move_target=robot_model.link_joint2)
# array([-0.3183119 ,  0.63661397,  0.        ], dtype=float32)
viewer.redraw()
print(robot_model.link_joint2.copy_worldcoords())
#<Coordinates 0x176d61c50 1.900 0.000 0.000 / 0.0 0.3 0.0>