diff --git a/mock/MotionController.py b/mock/MotionController.py new file mode 100644 index 0000000..17b98d9 --- /dev/null +++ b/mock/MotionController.py @@ -0,0 +1,53 @@ +import copy + + +class GaitState: + def __init__(self) -> None: + self.step_length = 0.1 + self.yaw_rate = 0 + self.lateral_fraction = 0 + self.step_velocity = 0.001 + self.swing_period = 0.2 + self.clearance_height = 0.045 + self.penetration_depth = 0.003 + self.contacts = [False] * 4 + + self.target_step_length = 0 + self.target_yaw_rate = 0 + self.target_lateral_fraction = 0 + + def update_gait_state(self, dt): + self.step_length = self.step_length * (1 - dt) + self.target_step_length * dt + self.lateral_fraction = ( + self.lateral_fraction * (1 - dt) + self.target_lateral_fraction * dt + ) + self.yaw_rate = self.yaw_rate * (1 - dt) + self.target_yaw_rate * dt + + +class MotionController: + def __init__( + self, + # env: spotBezierEnv, + # gui: GUI, + # bodyState: BodyState, + # gaitState: GaitState, + spot_model, + gait, + ) -> None: + self.gait = gait + self.gait_state = GaitState() + self.spot_model = spot_model + + self.dt = 0.01 + + def update_gait_state(self, command): + self.gait_state.step_length = abs(command["lx"]) / 255 + + def run(self, model, command): + self.update_gait_state(command) + self.gait_state.contacts = [False] * 4 + self.body_state.worldFeetPositions = copy.deepcopy(self.spot.WorldToFoot) + + model["transformation"]["world_feet_position"] = self.gait.generate_trajectory( + model, self.gait_state, self.dt + ) diff --git a/mock/__init__.py b/mock/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mock/model.py b/mock/model.py new file mode 100644 index 0000000..66fb90f --- /dev/null +++ b/mock/model.py @@ -0,0 +1,34 @@ +import random + + +model = lambda: { + "gait": { + "step_length": 0, + "yaw_rate": 0, + "lateral_fraction": 0, + "step_velocity": 0, + "swing_period": 0, + "clearance_height": 0, + "penetration_depth": 0, + "contacts": 0, + }, + "transformation": { + "world_position": [0, 0, 0], + "position": [0, 0, 0], + "rotation": [0, 0, 0], + "world_feet_positions": {}, + }, + "sensors": { + "mpu": { + "x": 0, + "y": 0, + "z": 0, + }, + "battery": { + "voltage": round(random.uniform(7.6, 8.2), 2), + "ampere": round(random.uniform(0.2, 3), 2), + }, + }, + "logs": ["[2023-02-05 10:00:00] Booting up"], + "settings": {"useMetric": True}, +} diff --git a/mock/server.py b/mock/server.py new file mode 100644 index 0000000..9c96723 --- /dev/null +++ b/mock/server.py @@ -0,0 +1,100 @@ +import asyncio +from enum import Enum +import json +import sys +import websockets +from model import model +from MotionController import GaitState +from simulator.GaitGenerator.Bezier import BezierGait +from simulator.GymEnvs.spot_bezier_env import spotBezierEnv +from simulator.Kinematics.SpotKinematics import SpotModel +from simulator.util.gui import GUI +from simulator.simulator import BodyState, Simulator +import struct + +sys.path.append("./simulator/GymEnvs") + +clients = {} + +env = spotBezierEnv( + render=True, + on_rack=False, + height_field=False, + draw_foot_path=False, + env_randomizer=None, +) +gui = GUI(env.spot.quadruped) +bodyState = BodyState() +gaitState = GaitState() +spot = SpotModel() +bezierGait = BezierGait() +simulator = Simulator() + +class Command(Enum): + ESTOP = 0 + CONTROLLER = 1 + + +def get_controller(buffer): + buffer = struct.unpack("<8b", buffer) + return { + "command": buffer[0], + "estop": buffer[1], + "lx": buffer[2], + "ly": buffer[3], + "rx": buffer[4], + "ry": buffer[5], + "height": buffer[6], + "speed": buffer[7], + } + + +async def handle_binary_message(client, data): + message = get_controller(data) + command = Command(message["command"]) + if command == Command.ESTOP: + client["model"]["running"] = False + 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})) + + +async def handle_json_message(client, message): + data = json.loads(message) + client = client["clientState"] + 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")} + ) + ) + + +async def handle_message(websocket, path): + client_id = id(websocket) + clients[client_id] = { + "clientState": model(), + "websocket": websocket, + } + try: + async for message in websocket: + if isinstance(message, bytes): + await handle_binary_message(clients[client_id], message) + else: + await handle_json_message(clients[client_id], message) + finally: + del clients[client_id] + + +async def main(): + async with websockets.serve(handle_message, "localhost", 2096): + print("Server starting") + await asyncio.Future() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/mock/simulator/GymEnvs/spot_bezier_env.py b/mock/simulator/GymEnvs/spot_bezier_env.py index 157ffa3..95df89e 100644 --- a/mock/simulator/GymEnvs/spot_bezier_env.py +++ b/mock/simulator/GymEnvs/spot_bezier_env.py @@ -10,13 +10,11 @@ import pybullet_data from gym import spaces from gym.utils import seeding from pkg_resources import parse_version -import spot +from .. import spot import pybullet_utils.bullet_client as bullet_client from gym.envs.registration import register -from OpenLoopSM.SpotOL import BezierStepper -from spot_gym_env import spotGymEnv -import Kinematics.LieAlgebra as LA -from spot_env_randomizer import SpotEnvRandomizer +from ..spot_gym_env import spotGymEnv +from ..spot_env_randomizer import SpotEnvRandomizer SENSOR_NOISE_STDDEV = spot.SENSOR_NOISE_STDDEV diff --git a/mock/simulator/Kinematics/SpotKinematics.py b/mock/simulator/Kinematics/SpotKinematics.py index 7e824b4..cc35f30 100644 --- a/mock/simulator/Kinematics/SpotKinematics.py +++ b/mock/simulator/Kinematics/SpotKinematics.py @@ -1,8 +1,8 @@ #!/usr/bin/env python import numpy as np -from Kinematics.LegKinematics import LegIK -from Kinematics.LieAlgebra import RpToTrans, TransToRp, TransInv, RPY, TransformVector +from .LegKinematics import LegIK +from .LieAlgebra import RpToTrans, TransToRp, TransInv, RPY, TransformVector from collections import OrderedDict diff --git a/mock/simulator/simulator.py b/mock/simulator/simulator.py index 53267ce..74e3a1a 100644 --- a/mock/simulator/simulator.py +++ b/mock/simulator/simulator.py @@ -4,12 +4,12 @@ import numpy as np import copy import sys -sys.path.append("../../") +# sys.path.append("../../") -from GymEnvs.spot_bezier_env import spotBezierEnv -from util.gui import GUI -from Kinematics.SpotKinematics import SpotModel -from GaitGenerator.Bezier import BezierGait +from .GymEnvs.spot_bezier_env import spotBezierEnv +from .util.gui import GUI +from .Kinematics.SpotKinematics import SpotModel +from .GaitGenerator.Bezier import BezierGait class GaitState: diff --git a/mock/simulator/spot.py b/mock/simulator/spot.py index b81dfa0..b5c94fe 100644 --- a/mock/simulator/spot.py +++ b/mock/simulator/spot.py @@ -16,12 +16,12 @@ import copy import math import re import numpy as np -import motor +from . import motor from util import pybullet_data print(pybullet_data.getDataPath()) -from Kinematics.SpotKinematics import SpotModel -import Kinematics.LieAlgebra as LA +from .Kinematics.SpotKinematics import SpotModel +from .Kinematics import LieAlgebra as LA INIT_POSITION = [0, 0, 0.25] INIT_RACK_POSITION = [0, 0, 1] diff --git a/mock/simulator/spot_env_randomizer.py b/mock/simulator/spot_env_randomizer.py index 8144488..87be6c0 100644 --- a/mock/simulator/spot_env_randomizer.py +++ b/mock/simulator/spot_env_randomizer.py @@ -12,7 +12,7 @@ https://github.com/bulletphysics/bullet3/blob/master/examples/pybullet/gym/pybul """ import numpy as np -import env_randomizer_base +from . import env_randomizer_base # Relative range. spot_BASE_MASS_ERROR_RANGE = (-0.2, 0.2) # 0.2 means 20% diff --git a/mock/simulator/spot_gym_env.py b/mock/simulator/spot_gym_env.py index c6e5591..d0bca88 100644 --- a/mock/simulator/spot_gym_env.py +++ b/mock/simulator/spot_gym_env.py @@ -20,13 +20,13 @@ import pybullet_data from gym import spaces from gym.utils import seeding from pkg_resources import parse_version -import spot +from . import spot import pybullet_utils.bullet_client as bullet_client from gym.envs.registration import register -from heightfield import HeightField -from OpenLoopSM.SpotOL import BezierStepper -import Kinematics.LieAlgebra as LA -from spot_env_randomizer import SpotEnvRandomizer +from .heightfield import HeightField +from .OpenLoopSM.SpotOL import BezierStepper +from .Kinematics import LieAlgebra as LA +from .spot_env_randomizer import SpotEnvRandomizer NUM_SUBSTEPS = 5 NUM_MOTORS = 12