StanfordQuadruped/run_robot.py.save at main · simonchamorro/StanfordQuadruped · 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
import pickle

import numpy as np
import time
import os

from pupper.HardwareInterfaceArm import HardwareInterfaceArm
from src.IMU import IMU
from src.Controller import Controller
from src.JoystickInterface import JoystickInterface
from src.State import State, ArmState, GripperState
from pupper.HardwareInterface import HardwareInterface
from pupper.Config import Configuration
from pupper.Kinematics import four_legs_inverse_kinematics


def main(use_imu=False):
    """Main program
    """

    # Create config
    config = Configuration()
    hardware_interface = HardwareInterface()
    #hardware_interface_arm = HardwareInterfaceArm()
    #hardware_interface_arm.start()

    # Create imu handle
    if use_imu:
        imu = IMU(port="/dev/ttyACM0")
        imu.flush_buffer()

    # Create controller and user input handles
    controller = Controller(config, four_legs_inverse_kinematics,)
    state = State()
    print("Creating joystick listener...")
    joystick_interface = JoystickInterface(config)
    print("Done.")

    last_loop = time.time()

    print("Summary of gait parameters:")
    print("overlap time: ", config.overlap_time)
    print("swing time: ", config.swing_time)
    print("z clearance: ", config.z_clearance)
    print("x shift: ", config.x_shift)

    currently_recording = False
    currently_running_arm = False
    recording_pause_timer = time.time()
    states = {}
    state.arm_joint_angles[:] = hardware_interface_arm.pos_rest
    hardware_interface_arm.move(state.arm_joint_angles)

    # Wait until the activate button has been pressed
    while True:
        print("Waiting for L1 to activate robot.")
        while True:
            command = joystick_interface.get_command(state)
            joystick_interface.set_color(config.ps4_deactivated_color)
            if command.activate_event == 1:
                break
            time.sleep(0.1)
        print("Robot activated.")
        joystick_interface.set_color(config.ps4_color)

        while True:
            now = time.time()
            if now - last_loop < config.dt:
                continue
            last_loop = time.time()

            # Parse the udp joystick commands and then update the robot controller's parameters
            command = joystick_interface.get_command(state)
            if command.activate_event == 1:
                print("Deactivating Robot")
                break

            if state.arm_state is not ArmState.DEACTIVATED:
                if not currently_running_arm:
                    joystick_interface.set_color(config.ps4_arm_color)
                    if (state.arm_joint_angles == hardware_interface_arm.pos_rest).all():
                        state.arm_joint_angles = hardware_interface_arm.pos_erect
                currently_running_arm = True
                hardware_interface_arm.move(state.arm_joint_angles)
            else:
                if currently_running_arm:
                    joystick_interface.set_color(config.ps4_color)
                    currently_running_arm = False

            if command.record_event == 1 and not currently_recording and (time.time() - recording_pause_timer) > 1:
                print("starting recording")
                joystick_interface.set_color(config.ps4_record_color)
                currently_recording = True
                recording_pause_timer = time.time()
                states = {}
            elif command.record_event == 1 and currently_recording and (time.time() - recording_pause_timer) > 1:
                print("finishing recording")
                joystick_interface.set_color(config.ps4_color)
                currently_recording = False
                with open(f"/home/pi/recordings/{time.time()}.pkl", "wb") as handle:
                    pickle.dump(states, handle, protocol=pickle.HIGHEST_PROTOCOL)
                del states
                recording_pause_timer = time.time()

            if command.gripper_event == 1:

                if state.gripper_state is GripperState.NEUTRAL:
                    hardware_interface_arm.gripper_open()
                elif state.gripper_state is GripperState.OPEN:
                    hardware_interface_arm.gripper_close()
                elif state.gripper_state is GripperState.CLOSED:
                    hardware_interface_arm.gripper_neutral()

            # Read imu data. Orientation will be None if no data was available
            quat_orientation = imu.read_orientation() if use_imu else np.array([1, 0, 0, 0])
            state.quat_orientation = quat_orientation

            # Step the controller forward by dt
            controller.run(state, command)

            if currently_recording:
                states[str(time.time())] = {
                    "feet": np.copy(state.foot_locations),
                    "joints": np.copy(state.joint_angles),
                    "height": np.copy(state.height),
                    "yaw": np.copy(state.yaw_rate),
                    "pitch": np.copy(state.pitch),
                    "roll": np.copy(state.roll),
                }

            # Update the pwm widths going to the servos
            hardware_interface.set_actuator_postions(state.joint_angles)


main()