RoboRally/remote_control/mpc_controller.py

import numpy as np
import time

from casadi_opt import OpenLoopSolver


class MPCController:
    def __init__(self, estimator):
        self.t = time.time()

        self.estimator = estimator
        self.controlling = False

        self.mstep = 2
        self.ols = OpenLoopSolver(N=20, T=1.0)
        self.ols.setup()
        self.dt = self.ols.T / self.ols.N

        # integrator
        self.omega_max = 5.0
        self.control_scaling = 0.4

    def move_to_pos(self, target_pos, robot, near_target_counter=5):
        near_target = 0
        while near_target < near_target_counter:
            while not self.estimator.event_queue.empty():
                event = self.estimator.event_queue.get()
                print("event: ", event)
                if event[0] == 'click':
                    pass
                elif event[0] == 'key':
                    key = event[1]

                    if key == 84:  # arrow up
                        self.controlling = True
                        self.t = time.time()
                    elif key == 82:  # arrow down
                        self.controlling = False
                        robot.send_cmd()
                    elif key == 48:  # 0
                        target_pos = np.array([0.0, 0.0, 0.0])
                    elif key == 43:  # +
                        self.control_scaling += 0.1
                        self.control_scaling = min(self.control_scaling, 1.0)
                        print("control scaling = ", self.control_scaling)
                    elif key == 45:  # -
                        self.control_scaling -= 0.1
                        self.control_scaling = max(self.control_scaling, 0.1)
                        print("control scaling = ", self.control_scaling)
                    elif key == 113:
                        print("quit!")
                        self.controlling = False
                        robot.send_cmd()
                        return
                    elif key == 27:  # escape
                        print("quit!")
                        self.controlling = False
                        robot.send_cmd()
                        return

            x_pred = self.get_measurement(robot.id)

            if x_pred is not None:
                error_pos = np.linalg.norm(x_pred[0:2] - target_pos[0:2])
                angles_unwrapped = np.unwrap([x_pred[2], target_pos[2]])  # unwrap angle to avoid jump in data
                error_ang = np.abs(angles_unwrapped[0] - angles_unwrapped[1])
                # print("error pos = ", error_pos)
                # print("error_pos = {}, error_ang = {}".format(error_pos, error_ang))

                # if error_pos > 0.075 or error_ang > 0.35:
                if error_pos > 0.05 or error_ang > 0.1:
                    # solve mpc open loop problem
                    res = self.ols.solve(x_pred, target_pos)

                    # us1 = res[0]
                    # us2 = res[1]
                    us1 = res[0] * self.control_scaling
                    us2 = res[1] * self.control_scaling
                    # print("u = {}", (us1, us2))

                    # print("---------------- mpc solution took {} seconds".format(tmpc_end - tmpc_start))
                    dt_mpc = time.time() - self.t
                    if dt_mpc < self.dt:  # wait until next control can be applied
                        # print("sleeping for {} seconds...".format(self.dt - dt_mpc))
                        time.sleep(self.dt - dt_mpc)
                else:
                    us1 = [0] * self.mstep
                    us2 = [0] * self.mstep
                    near_target += 1

                # send controls to the robot
                for i in range(0, self.mstep):  # option to use multistep mpc if len(range) > 1
                    u1 = us1[i]
                    u2 = us2[i]
                    robot.send_cmd(u1, u2)
                    if i < self.mstep:
                        time.sleep(self.dt)
                    self.t = time.time()  # save time the most recent control was applied
            else:
                print("robot not detected yet!")

    def interactive_control(self, robots):
        controlled_robot_number = 0
        robot = robots[controlled_robot_number]

        target_pos = np.array([0.0, 0.0, 0.0])

        running = True
        while running:
            # handle events from opencv window
            while not self.estimator.event_queue.empty():
                event = self.estimator.event_queue.get()
                print("event: ", event)
                if event[0] == 'click':
                    target_pos = event[1]
                elif event[0] == 'key':
                    key = event[1]

                    if key == 32:  # arrow up
                        self.controlling = not self.controlling
                        if not self.controlling:
                            print("disable control")
                            robot.send_cmd()  # stop robot
                        else:
                            print("enable control")
                            self.t = time.time()
                    elif key == 48:  # 0
                        target_pos = np.array([0.0, 0.0, 0.0])  # TODO: use center of board for target pos
                    elif key == 43:  # +
                        self.control_scaling += 0.1
                        self.control_scaling = min(self.control_scaling, 1.0)
                        print("control scaling = ", self.control_scaling)
                    elif key == 45:  # -
                        self.control_scaling -= 0.1
                        self.control_scaling = max(self.control_scaling, 0.1)
                        print("control scaling = ", self.control_scaling)
                    elif key == 9:  # TAB
                        # switch controlled robot
                        robot.send_cmd()  # stop current robot
                        controlled_robot_number = (controlled_robot_number + 1) % len(robots)
                        robot = robots[controlled_robot_number]
                        print(f"controlled robot: {robot.id}")
                    elif key == 113 or key == 27:  # q or ESCAPE
                        print("quit!")
                        self.controlling = False
                        robot.send_cmd()
                        return

            if self.controlling:
                # measure state
                x_pred = self.get_measurement(robot.id)

                # print(np.linalg.norm(x_pred-target_pos))

                # solve mpc open loop problem
                res = self.ols.solve(x_pred, target_pos)

                us1 = res[0] * self.control_scaling
                us2 = res[1] * self.control_scaling

                dt_mpc = time.time() - self.t
                if dt_mpc < self.dt:  # wait until next control can be applied
                    time.sleep(self.dt - dt_mpc)

                # send controls to the robot
                for i in range(0, self.mstep):  # option to use multistep mpc if len(range) > 1
                    u1 = us1[i]
                    u2 = us2[i]
                    robot.send_cmd(u1, u2)
                    if i < self.mstep:
                        time.sleep(self.dt)
                    self.t = time.time()  # save time the most recent control was applied

    def get_measurement(self, robot_id):
        return self.estimator.get_robot_state_estimate(robot_id)