mink/examples/arm_dual_flying_ur5e.py at main · jaikr-dev/mink · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
from pathlib import Path

import mujoco
import mujoco.viewer
import numpy as np
from loop_rate_limiters import RateLimiter

import mink

_HERE = Path(__file__).parent
_XML = _HERE / "universal_robots_ur5e" / "ur5e.xml"


def construct_model() -> mujoco.MjModel:
    root = mujoco.MjSpec()
    root.stat.meansize = 0.08
    root.visual.global_.azimuth = -120
    root.visual.global_.elevation = -20

    root.worldbody.add_light(pos=(0, 0, 1.5), type=mujoco.mjtLightType.mjLIGHT_SPOT)

    base = root.worldbody.add_body(name="base")
    width, height, depth = 0.8, 0.4, 0.2
    base.add_geom(
        type=mujoco.mjtGeom.mjGEOM_BOX,
        size=[width, height, depth],
        density=1e-3,
        rgba=(0.9, 0.8, 0.6, 1),
    )
    base.pos[:] = [-0.0 * width, -0.0 * height, -0.5 * depth]
    base.add_freejoint()
    base.add_site(name="base", pos=(0, 0, depth), group=1)
    left_site = base.add_site(name="l_attachment_site", pos=(0.3, 0, depth), group=5)
    right_site = base.add_site(name="r_attachment_site", pos=(-0.3, 0, depth), group=5)

    left_ur5e = mujoco.MjSpec.from_file(_XML.as_posix())
    left_ur5e.modelname = "l_ur5e"
    left_ur5e.delete(left_ur5e.key("home"))
    root.attach(left_ur5e, prefix="l_ur5e/", site=left_site)

    right_ur5e = mujoco.MjSpec.from_file(_XML.as_posix())
    right_ur5e.modelname = "r_ur5e"
    right_ur5e.delete(right_ur5e.key("home"))
    root.attach(right_ur5e, prefix="r_ur5e/", site=right_site)

    body = root.worldbody.add_body(name="base_target", mocap=True)
    body.add_geom(
        type=mujoco.mjtGeom.mjGEOM_BOX,
        size=(0.05, 0.05, 0.05),
        contype=0,
        conaffinity=0,
        rgba=(0.6, 0.3, 0.3, 0.5),
    )

    body = root.worldbody.add_body(name="l_target", mocap=True)
    body.add_geom(
        type=mujoco.mjtGeom.mjGEOM_BOX,
        size=(0.05, 0.05, 0.05),
        contype=0,
        conaffinity=0,
        rgba=(0.3, 0.6, 0.3, 0.5),
    )

    body = root.worldbody.add_body(name="r_target", mocap=True)
    body.add_geom(
        type=mujoco.mjtGeom.mjGEOM_BOX,
        size=(0.05, 0.05, 0.05),
        contype=0,
        conaffinity=0,
        rgba=(0.3, 0.3, 0.6, 0.5),
    )

    return root.compile()


if __name__ == "__main__":
    model = construct_model()

    configuration = mink.Configuration(model)

    tasks = [
        base_task := mink.FrameTask(
            frame_name="base",
            frame_type="site",
            position_cost=1.0,
            orientation_cost=1.0,
        ),
        left_ee_task := mink.FrameTask(
            frame_name="l_ur5e/attachment_site",
            frame_type="site",
            position_cost=1.0,
            orientation_cost=1.0,
        ),
        right_ee_task := mink.FrameTask(
            frame_name="r_ur5e/attachment_site",
            frame_type="site",
            position_cost=1.0,
            orientation_cost=1.0,
        ),
    ]

    base_mid = model.body("base_target").mocapid[0]
    left_mid = model.body("l_target").mocapid[0]
    right_mid = model.body("r_target").mocapid[0]
    model = configuration.model
    data = configuration.data
    solver = "daqp"

    with mujoco.viewer.launch_passive(
        model=model, data=data, show_left_ui=False, show_right_ui=False
    ) as viewer:
        mujoco.mjv_defaultFreeCamera(model, viewer.cam)
        viewer.opt.frame = mujoco.mjtFrame.mjFRAME_SITE

        for mocap, frame in zip(
            ["base_target", "l_target", "r_target"],
            ["base", "l_ur5e/attachment_site", "r_ur5e/attachment_site"],
        ):
            mink.move_mocap_to_frame(model, data, mocap, frame, "site")

        rate = RateLimiter(frequency=200.0, warn=False)
        t = 0.0
        while viewer.is_running():
            data.mocap_pos[base_mid][2] = 0.3 * np.sin(2.0 * t)
            base_task.set_target(mink.SE3.from_mocap_name(model, data, "base_target"))

            data.mocap_pos[left_mid][1] = 0.5 + 0.2 * np.sin(2.0 * t)
            data.mocap_pos[left_mid][2] = 0.2
            left_ee_task.set_target(mink.SE3.from_mocap_name(model, data, "l_target"))

            data.mocap_pos[right_mid][1] = 0.5 + 0.2 * np.sin(2.0 * t)
            data.mocap_pos[right_mid][2] = 0.2
            right_ee_task.set_target(mink.SE3.from_mocap_name(model, data, "r_target"))

            vel = mink.solve_ik(configuration, tasks, rate.dt, solver, damping=1e-2)
            configuration.integrate_inplace(vel, rate.dt)
            mujoco.mj_camlight(model, data)

            viewer.sync()
            rate.sleep()
            t += rate.dt