hebi-python-examples/kits/arms/ex_AR_kit.py at main · HebiRobotics/hebi-python-examples · 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
142
143
144
145
#!/usr/bin/env python3

import hebi
from os.path import join, dirname, realpath
import numpy as np
from scipy.spatial.transform import Rotation as R
from time import sleep
from hebi_util import create_mobile_io_from_config

# Initialize the interface for network connected modules
lookup = hebi.Lookup()
sleep(2)

# Config file
# Relative to this file directory
example_config_file = join(dirname(__file__), "config/ex_AR_kit.cfg.yaml")
example_config = hebi.config.load_config(example_config_file)

user_data = example_config.user_data
if user_data is None:
    needed_fields = ['homing_duration',
                     'home_position', 'xyz_scale', 'delay_time']
    raise RuntimeError(
        f'This Demo needs the following fields set in the configuration user_data:\n{needed_fields}')

# Set up arm, and mobile_io from config
arm = hebi.arm.create_from_config(example_config, lookup)
mobile_io = create_mobile_io_from_config(example_config, lookup)

# Demo Variables
abort_flag = False
run_mode = "softstart"
goal = hebi.arm.Goal(arm.size)

# Command the softstart to the home position
softstart = hebi.arm.Goal(arm.size)
softstart.add_waypoint(t=user_data['homing_duration'],
                       position=user_data['home_position'])
arm.update()
arm.set_goal(softstart)
arm.send()

# Get the cartesian position and rotation matrix @ home position
xyz_home = np.zeros(3)
rot_home = np.zeros((3, 3))
arm.FK(user_data['home_position'], xyz_out=xyz_home, orientation_out=rot_home)

# Get the states for the mobile device
xyz_phone_init = np.zeros(3)
rot_phone_init = np.zeros((3, 3))

# Print instructions
instructions = """AR KIT EXAMPLE

    🏠 - Home
    📲 - AR Control
    🌍 - Grav Comp
    ❌ - Quit"""

print(instructions)
mobile_io.add_text(instructions)

#######################
## Main Control Loop ##
#######################

home_btn = 1
ar_btn = 3
gravcomp_btn = 6
quit_btn = 8


xyz_scale = np.array(user_data['xyz_scale'])

while not abort_flag:
    arm.update()  # update the arm

    if run_mode == "softstart":
        # End softstart when the arm reaches the home_position
        if arm.at_goal:
            mobile_io.set_led_color("yellow")
            run_mode = "waiting"
            continue
        arm.send()
        continue

    if not mobile_io.update(timeout_ms=0):
        # print("Failed to get feedback from MobileIO")
        arm.send()
        continue

    # Quit
    elif mobile_io.get_button_diff(quit_btn) == 1:
        mobile_io.set_led_color("transparent")
        abort_flag = True
        break

    # Return to home position
    elif mobile_io.get_button_diff(home_btn) == 1:
        mobile_io.set_led_color("yellow")
        run_mode = "waiting"
        arm.set_goal(softstart)

    # Start AR Control
    elif mobile_io.get_button_diff(ar_btn) == 1 and run_mode != "ar_mode":
        mobile_io.set_led_color("green")
        run_mode = "ar_mode"

        # Store initial position and orientation as baseline
        xyz_phone_init = mobile_io.position.copy()
        wxyz = mobile_io.orientation
        xyzw = [*wxyz[1:], wxyz[0]]
        rot_phone_init = R.from_quat(xyzw).as_matrix()

    # Grav Comp Mode
    elif mobile_io.get_button_diff(gravcomp_btn) == 1:  # "ToOn"
        mobile_io.set_led_color("blue")
        run_mode = "grav_comp"
        arm.cancel_goal()

    if run_mode == "ar_mode":
        # Get the latest mobile position and orientation
        xyz_phone = mobile_io.position
        wxyz = mobile_io.orientation
        xyzw = [*wxyz[1:], wxyz[0]]
        rot_phone = R.from_quat(xyzw).as_matrix()

        # Calculate new targets
        xyz_target = xyz_home + \
            rot_phone_init.T @ (xyz_scale * (xyz_phone - xyz_phone_init))
        rot_target = rot_phone_init.T @ rot_phone @ rot_home

        # Calculate new arm joint angles
        target_joints = arm.ik_target_xyz_so3(
            arm.last_feedback.position, xyz_target, rot_target)

        # Set and send new goal to the arm
        goal.clear()
        goal.add_waypoint(position=target_joints,
                          t=float(user_data['delay_time']))
        arm.set_goal(goal)

    arm.send()

mobile_io.set_led_color("red")