From fae6171f93a2aa72e0a4011e3566fe569fd173ed Mon Sep 17 00:00:00 2001 From: Rune Harlyk Date: Tue, 5 Mar 2024 15:02:31 +0100 Subject: [PATCH] Simplifies spot class and add threaded websocket controller interface --- raspberry_pi/main.py | 10 -- raspberry_pi/run_spot.py | 33 ++++++ raspberry_pi/src/IMU/IMU.py | 2 +- raspberry_pi/src/Kinematics/SpotKinematics.py | 5 +- .../controller_interface/ControllerState.py | 10 ++ .../SharedControllerState.py | 18 +++ .../WebsocketController.py | 108 ++++++++++++++++++ .../hardware_interface/HardwareInterface.py | 10 ++ .../HardwareInterfaceBase.py | 11 ++ .../src/hardware_interface/__init__.py | 0 raspberry_pi/src/spot.py | 35 +++--- 11 files changed, 213 insertions(+), 29 deletions(-) delete mode 100644 raspberry_pi/main.py create mode 100644 raspberry_pi/run_spot.py create mode 100644 raspberry_pi/src/controller_interface/ControllerState.py create mode 100644 raspberry_pi/src/controller_interface/SharedControllerState.py create mode 100644 raspberry_pi/src/controller_interface/WebsocketController.py create mode 100644 raspberry_pi/src/hardware_interface/HardwareInterface.py create mode 100644 raspberry_pi/src/hardware_interface/HardwareInterfaceBase.py create mode 100644 raspberry_pi/src/hardware_interface/__init__.py diff --git a/raspberry_pi/main.py b/raspberry_pi/main.py deleted file mode 100644 index 055a0a9..0000000 --- a/raspberry_pi/main.py +++ /dev/null @@ -1,10 +0,0 @@ -# SCRIPT FOR RUNNING THE RPI ROBOT - -from .src.spot import Spot -# Server - -# Bluetooth controller - -# Motion planing - -# Hardware connections \ No newline at end of file diff --git a/raspberry_pi/run_spot.py b/raspberry_pi/run_spot.py new file mode 100644 index 0000000..0205a1c --- /dev/null +++ b/raspberry_pi/run_spot.py @@ -0,0 +1,33 @@ +# SCRIPT FOR RUNNING THE RPI ROBOT + +import time +from src.spot import Spot +from src.kinematics.SpotKinematics import SpotModel +from src.camera.WebCamera import WebCamera +from src.imu.IMU import IMU +from src.hardware_interface.HardwareInterface import HardwareInterface +from src.controller_interface.WebsocketController import WebsocketController +from src.controller_interface.SharedControllerState import SharedState, shared_state + + +def main(): + kinematics = SpotModel() + camera = WebCamera() + imu = IMU() + hardware_interface = HardwareInterface() + # shared_controller_state = SharedState() + controller_interface = WebsocketController() # shared_controller_state) + + spot = Spot(kinematics, camera, imu, hardware_interface, controller_interface) + + spot.start() + + try: + while True: + spot.run(shared_state.get_latest_state().__dict__) + time.sleep(1) + except KeyboardInterrupt: + spot.stop() + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/raspberry_pi/src/IMU/IMU.py b/raspberry_pi/src/IMU/IMU.py index 4d70e0f..646ecd4 100644 --- a/raspberry_pi/src/IMU/IMU.py +++ b/raspberry_pi/src/IMU/IMU.py @@ -1,5 +1,5 @@ import numpy as np -from .IMUBase import IMUBase +from src.imu.IMUBase import IMUBase class IMU(IMUBase): def __init__(self) -> None: diff --git a/raspberry_pi/src/Kinematics/SpotKinematics.py b/raspberry_pi/src/Kinematics/SpotKinematics.py index cc35f30..c499adc 100644 --- a/raspberry_pi/src/Kinematics/SpotKinematics.py +++ b/raspberry_pi/src/Kinematics/SpotKinematics.py @@ -1,10 +1,11 @@ #!/usr/bin/env python import numpy as np -from .LegKinematics import LegIK -from .LieAlgebra import RpToTrans, TransToRp, TransInv, RPY, TransformVector from collections import OrderedDict +from src.kinematics.LegKinematics import LegIK +from src.kinematics.LieAlgebra import RpToTrans, TransToRp, TransInv, RPY + class SpotModel: def __init__( diff --git a/raspberry_pi/src/controller_interface/ControllerState.py b/raspberry_pi/src/controller_interface/ControllerState.py new file mode 100644 index 0000000..1203016 --- /dev/null +++ b/raspberry_pi/src/controller_interface/ControllerState.py @@ -0,0 +1,10 @@ +class ControllerState: + def __init__(self) -> None: + self.command = 0 + self.estop = 0 + self.lx = 0 + self.ly = 0 + self.rx = 0 + self.ry = 0 + self.height = 0 + self.speed = 0 \ No newline at end of file diff --git a/raspberry_pi/src/controller_interface/SharedControllerState.py b/raspberry_pi/src/controller_interface/SharedControllerState.py new file mode 100644 index 0000000..bedfe61 --- /dev/null +++ b/raspberry_pi/src/controller_interface/SharedControllerState.py @@ -0,0 +1,18 @@ +import threading + +from src.controller_interface.ControllerState import ControllerState + +class SharedState: + def __init__(self): + self.lock = threading.Lock() + self.state = ControllerState() + + def update_state(self, new_state): + with self.lock: + self.state = new_state + + def get_latest_state(self): + with self.lock: + return self.state + +shared_state = SharedState() \ No newline at end of file diff --git a/raspberry_pi/src/controller_interface/WebsocketController.py b/raspberry_pi/src/controller_interface/WebsocketController.py new file mode 100644 index 0000000..930bec6 --- /dev/null +++ b/raspberry_pi/src/controller_interface/WebsocketController.py @@ -0,0 +1,108 @@ +import asyncio +from enum import Enum +import struct +import threading +import websockets +import json + +from src.controller_interface.ControllerState import ControllerState +from src.controller_interface.SharedControllerState import shared_state + +clients = {} + +class Command(Enum): + ESTOP = 0 + CONTROLLER = 1 + +def get_controller(buffer): + buffer = struct.unpack("<8b", buffer) + state = ControllerState() + state.command = buffer[0] + state.estop = buffer[1] + state.lx = buffer[2] + state.ly = buffer[3] + state.rx = buffer[4] + state.ry = buffer[5] + state.height = buffer[6] + state.speed = buffer[7] + + return state + +class WebsocketController: + def __init__(self, shared_controller_state=shared_state, host="localhost", port=2096): + self.shared_controller_state = shared_controller_state + self.host = host + self.port = port + self.loop = None + self.thread = None + + async def handle_message(self, websocket, path): + client_id = id(websocket) + clients[client_id] = { + "websocket": websocket, + } + print("Got a new connection") + try: + async for message in websocket: + if isinstance(message, bytes): + await self.handle_binary_message(clients[client_id], message) + else: + await self.handle_json_message(clients[client_id], message) + finally: + del clients[client_id] + + async def handle_binary_message(self, client, data): + message = get_controller(data) + shared_state.update_state(message) + command = Command(message.command) + if command == Command.ESTOP: + await client["websocket"].send( + json.dumps({"type": "stop", "data": "Servos stopped"}) + ) + + if command == Command.CONTROLLER: + await client["websocket"].send(json.dumps({"type": "echo", "data": message.__dict__})) + + + async def handle_json_message(self, client, message): + data = json.loads(message) + if data["type"] in ("stop", "mode_change"): + client["model"][data["type"]] = data.get("data", False) + await client["websocket"].send( + json.dumps( + {"type": data["type"], "data": data.get("data", "Servos stopped")} + ) + ) + + def start_server(self): + asyncio.set_event_loop(self.loop) + start_server = websockets.serve(self.handle_message, self.host, self.port) + self.loop.run_until_complete(start_server) + self.loop.run_forever() + + def start(self): + # Create a new event loop in a new thread + if not self.thread or not self.thread.is_alive(): + self.loop = asyncio.new_event_loop() + self.thread = threading.Thread(target=self.start_server) + self.thread.start() + + def stop(self): + if self.loop: + self.loop.call_soon_threadsafe(self.loop.stop) + if self.thread and self.thread.is_alive(): + self.thread.join() + +if __name__ == "__main__": + server = WebsocketController("localhost", 2096) + server.run() + + try: + # Your main thread logic here + print("WebSocket server is running...") + # Example: Keep the main thread doing something or waiting + while True: + pass + except KeyboardInterrupt: + print("Stopping WebSocket server...") + server.stop() \ No newline at end of file diff --git a/raspberry_pi/src/hardware_interface/HardwareInterface.py b/raspberry_pi/src/hardware_interface/HardwareInterface.py new file mode 100644 index 0000000..e6a95e4 --- /dev/null +++ b/raspberry_pi/src/hardware_interface/HardwareInterface.py @@ -0,0 +1,10 @@ +from src.hardware_interface.HardwareInterfaceBase import HardwareInterfaceBase + + +class HardwareInterface(HardwareInterfaceBase): + def __init__(self): + super().__init__() + + def set_actuator_positions(self, joint_angles): + pass + \ No newline at end of file diff --git a/raspberry_pi/src/hardware_interface/HardwareInterfaceBase.py b/raspberry_pi/src/hardware_interface/HardwareInterfaceBase.py new file mode 100644 index 0000000..f0953aa --- /dev/null +++ b/raspberry_pi/src/hardware_interface/HardwareInterfaceBase.py @@ -0,0 +1,11 @@ +from abc import abstractmethod + +class HardwareInterfaceBase: + def __init__(self): + pass + + @abstractmethod + def set_actuator_positions(self, joint_angles): + raise NotImplementedError + + diff --git a/raspberry_pi/src/hardware_interface/__init__.py b/raspberry_pi/src/hardware_interface/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/raspberry_pi/src/spot.py b/raspberry_pi/src/spot.py index 7654c65..8252d6b 100644 --- a/raspberry_pi/src/spot.py +++ b/raspberry_pi/src/spot.py @@ -9,15 +9,25 @@ from .State import BodyState, GaitState sys.path.append("../../..") class Spot: - def __init__(self, kinematic, IMU=None): + def __init__(self, kinematic, camera, imu, hardware_interface, controller_interface): + self.kinematic = kinematic + self.camera = camera + self.imu = imu + self.hardware_interface = hardware_interface + self.controller_interface = controller_interface + self.body_state = BodyState() self.gait_state = GaitState() - self.kinematic = kinematic self.gait_controller = BezierGait() - self.IMU = IMU self.body_state.worldFeetPositions = copy.deepcopy(self.kinematic.WorldToFoot) + def start(self): + self.controller_interface.start() + + def stop(self): + self.controller_interface.stop() + def joint_angles(self): return self.kinematic.IK( self.body_state.rotation, @@ -25,27 +35,20 @@ class Spot: self.body_state.worldFeetPositions, ) - def run(self, state, command): - self.gait_state.update_gait_state(self.dt) - self.gui.UserInput(self.body_state, self.gait_state) - self.gait_state.contacts = self.state[-4:] - self.body_state.worldFeetPositions = copy.deepcopy(self.spot.WorldToFoot) + def run(self, dt=.01): + self.gait_state.update_gait_state(dt) + # self.gui.UserInput(self.body_state, self.gait_state) - self.bezier_gait.generate_trajectory(self.body_state, self.gait_state, self.dt) + self.gait_controller.generate_trajectory(self.body_state, self.gait_state, dt) self.update_environment() - self.state, _, done, _ = self.env.step(self.action) - if done: - print("DONE") - return True - def update_environment(self): - joint_angles = self.spot.IK( + joint_angles = self.kinematic.IK( self.body_state.rotation, self.body_state.position, self.body_state.worldFeetPositions, ) - self.env.pass_joint_angles(joint_angles.reshape(-1)) + self.hardware_interface.set_actuator_positions(joint_angles.reshape(-1))