Genesis/examples/rigid/diffik_controller.py at 9139dfbcfcf3b584fec75812930118bb86a8bd7a · Genesis-Embodied-AI/Genesis · 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
import argparse

import numpy as np

import genesis as gs

config = {
    "ur5e": {"mjcf_file": "xml/universal_robots_ur5e/ur5e.xml", "end_effector_link": "ee_virtual_link"},
    "panda": {"mjcf_file": "xml/franka_emika_panda/panda.xml", "end_effector_link": "left_finger"},
}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-v", "--vis", action="store_true", default=False)
    parser.add_argument("-c", "--cpu", action="store_true", default=False)
    parser.add_argument(
        "-r", "--robot", choices=["panda", "ur5e"], default="ur5e", help="Select robot model (panda or ur5e)"
    )
    args = parser.parse_args()

    ########################## init ##########################
    backend = gs.cpu if args.cpu else gs.gpu
    gs.init(backend=backend)

    ########################## create a scene ##########################
    scene = gs.Scene(
        viewer_options=gs.options.ViewerOptions(
            camera_pos=(0.0, -2, 1.5),
            camera_lookat=(0.0, 0.0, 0.5),
            camera_fov=40,
            max_FPS=200,
        ),
        rigid_options=gs.options.RigidOptions(
            enable_joint_limit=False,
            enable_collision=False,
            gravity=(0, 0, -0),
        ),
        show_viewer=args.vis,
        show_FPS=False,
    )

    ########################## entities ##########################

    plane = scene.add_entity(
        gs.morphs.Plane(),
    )

    robot_config = config[args.robot]
    mjcf_file = robot_config["mjcf_file"]
    robot = scene.add_entity(
        gs.morphs.MJCF(file=mjcf_file),
    )

    print("links=", robot.links)

    target_entity = scene.add_entity(
        gs.morphs.Mesh(
            file="meshes/axis.obj",
            scale=0.10,
        ),
        surface=gs.surfaces.Default(color=(1, 0.5, 0.5, 1)),
    )
    ########################## build ##########################
    scene.build()

    target_quat = np.array([0, 1, 0, 0])
    center = np.array([0.5, 0, 0.5])
    r = 0.1
    damping = 1e-4
    diag = damping * np.eye(6)

    end_effector_link = robot_config["end_effector_link"]
    ee_link = robot.get_link(end_effector_link)

    for i in range(0, 2000):
        target_pos = center + np.array([np.cos(i / 360 * np.pi), np.sin(i / 360 * np.pi), 0]) * r

        target_entity.set_qpos(np.concatenate([target_pos, target_quat]))

        # Position error.
        error_pos = target_pos - ee_link.get_pos().cpu().numpy()

        # Orientation error.
        ee_quat = ee_link.get_quat().cpu().numpy()
        error_quat = gs.transform_quat_by_quat(gs.inv_quat(ee_quat), target_quat)
        error_rotvec = gs.quat_to_rotvec(error_quat)

        error = np.concatenate([error_pos, error_rotvec])

        # jacobian
        jac = robot.get_jacobian(link=ee_link).cpu().numpy()
        dq = jac.T @ np.linalg.solve(jac @ jac.T + diag, error)
        q = robot.get_qpos().cpu().numpy() + dq

        # control
        robot.control_dofs_position(q)
        scene.step()


if __name__ == "__main__":
    main()