vis_cartesian_dmp.py

"""
====================
Cartesian DMP on UR5
====================

A trajectory is created manually, imitated with a Cartesian DMP, converted
to a joint trajectory by inverse kinematics, and executed with a UR5.
"""
print(__doc__)

import numpy as np
import pytransform3d.visualizer as pv
import pytransform3d.rotations as pr
import pytransform3d.transformations as pt
import pytransform3d.trajectories as ptr
from movement_primitives.kinematics import Kinematics
from movement_primitives.dmp import CartesianDMP


def animation_callback(step, graph, chain, joint_trajectory):
    chain.forward(joint_trajectory[step])
    graph.set_data()
    return graph


rotation_angle = np.deg2rad(45)
n_steps = 201

with open("examples/data/urdf/ur5.urdf", "r") as f:
    kin = Kinematics(f.read(), mesh_path="examples/data/urdf/")
chain = kin.create_chain(
    ["ur5_shoulder_pan_joint", "ur5_shoulder_lift_joint", "ur5_elbow_joint",
     "ur5_wrist_1_joint", "ur5_wrist_2_joint", "ur5_wrist_3_joint"],
    "ur5_base_link", "ur5_tool0")

q0 = np.array([0.0, -0.5, 0.8, -0.5, 0, 0])
chain.forward(q0)

ee2base_start = kin.tm.get_transform("ur5_tool0", "ur5_base_link")

rotation_axis = -pr.unity
start2end = pt.rotate_transform(np.eye(4), pr.matrix_from_compact_axis_angle(
    rotation_axis * rotation_angle))
ee2base_end = pt.concat(ee2base_start, start2end)

start = pt.exponential_coordinates_from_transform(ee2base_start)
end = pt.exponential_coordinates_from_transform(ee2base_end)

T = np.linspace(0, 1, n_steps)
trajectory = start[np.newaxis] + T[:, np.newaxis] * (end[np.newaxis] - start[np.newaxis])

dt = 0.01
execution_time = (n_steps - 1) * dt
T = np.linspace(0, execution_time, n_steps)
Y = ptr.pqs_from_transforms(ptr.transforms_from_exponential_coordinates(trajectory))
dmp = CartesianDMP(execution_time=execution_time, dt=dt, n_weights_per_dim=10)
dmp.imitate(T, Y)
_, Y = dmp.open_loop()

trajectory = ptr.transforms_from_pqs(Y)

random_state = np.random.RandomState(0)
joint_trajectory = chain.inverse_trajectory(
    trajectory, q0, random_state=random_state)

fig = pv.figure()
fig.plot_transform(s=0.3)

graph = fig.plot_graph(
    kin.tm, "ur5_base_link", show_visuals=False, show_collision_objects=True,
    show_frames=True, s=0.1, whitelist=["ur5_base_link", "ur5_tool0"])

fig.plot_transform(ee2base_start, s=0.15)
fig.plot_transform(ee2base_end, s=0.15)
fig.plot_transform(trajectory[-1], s=0.15)
fig.plot_transform(trajectory[-1], s=0.15)

pv.Trajectory(trajectory, s=0.05).add_artist(fig)

fig.view_init()
fig.animate(
    animation_callback, len(trajectory), loop=True,
    fargs=(graph, chain, joint_trajectory))
fig.show()