Complete Restructered Robot Code

Robot into its own Class instead of lose Global Variables that cause circular imports

StateClass usage instead of the old RobotState.py

New Input Class for Controller and randome intputs

states/IdleState.py

@@ -0,0 +1,30 @@
+"""
+states/IdleState.py - Stance & Idle state
+"""
+from __future__ import annotations
+from typing import TYPE_CHECKING, Optional
+from states.State import State
+
+if TYPE_CHECKING:
+    from Robot import Robot
+
+
+class IdleState(State):
+    def enter(self, robot: "Robot") -> None:
+        robot.reset_to_init()
+
+    def execute(self, robot: "Robot") -> Optional[str]:
+        # Check transition conditions based on vector_dirmov or command flags
+        vx, vy, omega = robot.vector_dirmov
+        if abs(vx) > 0.01 or abs(vy) > 0.01 or abs(omega) > 0.01:
+            return "walking"
+
+        if robot.robot_state == "wave":
+            return "wave_emote"
+        if robot.robot_state == "laola":
+            return "laola_emote"
+
+        return None
+
+    def exit(self, robot: "Robot") -> None:
+        pass

states/State.py

@@ -0,0 +1,35 @@
+"""
+states/State.py - Abstract base class for state machine
+"""
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Optional, Dict
+
+if TYPE_CHECKING:
+    from Robot import Robot
+
+
+class State(ABC):
+    """Abstract base class for all robot states."""
+
+    @abstractmethod
+    def enter(self, robot: "Robot") -> None:
+        """Called once when entering the state."""
+        pass
+
+    @abstractmethod
+    def execute(self, robot: "Robot") -> Optional[str]:
+        """
+        Called every control loop tick.
+        Returns Optional[str] containing the name of the next state if transitioning,
+        or None to stay in the current state.
+        """
+        pass
+
+    @abstractmethod
+    def exit(self, robot: "Robot") -> None:
+        """Called once when exiting the state."""
+        pass
+
+
+# Registry mapping state key names to state instances
+STATE_REGISTRY: Dict[str, State] = {}

states/WalkingState.py

@@ -0,0 +1,234 @@
+"""
+states/WalkingState.py - Walking Gait States (Tripod & Four/Wave)
+"""
+from __future__ import annotations
+from typing import TYPE_CHECKING, Optional
+import math
+import numpy as np
+
+import config as cfg
+import DataTypes as dt
+from states.State import State
+
+if TYPE_CHECKING:
+    from Robot import Robot
+
+
+def interpolate_bezier(t: float, p0: float, p1: float, p2: float) -> float:
+    """Quadratic Bezier interpolation."""
+    return (1.0 - t) ** 2 * p0 + 2.0 * (1.0 - t) * t * p1 + t**2 * p2
+
+
+class WalkingState(State):
+    """Tripod Gait Walking State."""
+
+    def __init__(
+        self,
+        duration: float = cfg.standard_duration,
+        tickpersec: float = cfg.standard_tickpersec,
+    ):
+        self.duration = duration
+        self.tickpersec = tickpersec
+        self.ticks = int(duration * tickpersec)
+        self.current_tick = 0
+        self.start_pos: Optional[dt.PosArray] = None
+        self.target_pos: Optional[dt.PosArray] = None
+
+    def enter(self, robot: Robot) -> None:
+        self.current_tick = 0
+        self.start_pos = dt.PosArray(np.copy(robot.current_pos.data))
+        self._calculate_target_positions(robot)
+
+    def _calculate_target_positions(self, robot: Robot) -> None:
+        vx, vy, omega = robot.vector_dirmov
+        vx *= cfg.translation_gain
+        vy *= cfg.translation_gain
+        omega *= cfg.rotation_gain
+
+        target_temp = []
+        for i in range(6):
+            cx, cy, cz = robot.center_points[i]
+
+            # Combine translation and rotation around body origin
+            v_x = vx + (-omega * cy)
+            v_y = vy + (omega * cx)
+
+            length = (v_x**2 + v_y**2) ** 0.5
+            if length > 1.0:
+                v_x /= length
+                v_y /= length
+
+            target_temp.append(
+                [cx + v_x * cfg.step_length, cy + v_y * cfg.step_length, cz]
+            )
+
+        self.target_pos = dt.PosArray(target_temp)
+
+    def execute(self, robot: Robot) -> Optional[str]:
+        # Transition back to idle if velocity is zero and step finished
+        vx, vy, omega = robot.vector_dirmov
+        if (
+            abs(vx) < 0.01
+            and abs(vy) < 0.01
+            and abs(omega) < 0.01
+            and self.current_tick == 0
+        ):
+            return "idle"
+
+        t = self.current_tick / float(self.ticks)
+        tick_pos_temp = []
+
+        for leg_id in range(6):
+            if robot.leg_state[leg_id] == "drag":
+                # Linear sliding toward center reference point
+                pos = self.start_pos[leg_id] + (
+                    robot.center_points[leg_id] - self.start_pos[leg_id]
+                ) * t
+                tick_pos_temp.append(pos)
+
+            elif robot.leg_state[leg_id] == "step":
+                # Bezier curve swing step
+                mid_point = [
+                    (self.start_pos[leg_id][0] + self.target_pos[leg_id][0]) / 2.0,
+                    (self.start_pos[leg_id][1] + self.target_pos[leg_id][1]) / 2.0,
+                    max(self.start_pos[leg_id][2], self.target_pos[leg_id][2])
+                    + cfg.step_height,
+                ]
+
+                x = interpolate_bezier(
+                    t,
+                    self.start_pos[leg_id][0],
+                    mid_point[0],
+                    self.target_pos[leg_id][0],
+                )
+                y = interpolate_bezier(
+                    t,
+                    self.start_pos[leg_id][1],
+                    mid_point[1],
+                    self.target_pos[leg_id][1],
+                )
+                z = interpolate_bezier(
+                    t,
+                    self.start_pos[leg_id][2],
+                    mid_point[2],
+                    self.target_pos[leg_id][2],
+                )
+                tick_pos_temp.append([x, y, z])
+
+        tick_pos = dt.PosArray(tick_pos_temp)
+        target_rad = robot.compute_ik(tick_pos)
+
+        # Update robot instance
+        robot.current_pos = tick_pos
+        robot.set_joint_angles(target_rad)
+
+        self.current_tick += 1
+
+        # Gait phase swap when step finishes
+        if self.current_tick > self.ticks:
+            self.current_tick = 0
+            self.start_pos = dt.PosArray(np.copy(robot.current_pos.data))
+
+            if robot.leg_state[0] == "step":
+                robot.leg_state = np.array(
+                    ["drag", "step", "drag", "step", "drag", "step"]
+                )
+            else:
+                robot.leg_state = np.array(
+                    ["step", "drag", "step", "drag", "step", "drag"]
+                )
+
+            self._calculate_target_positions(robot)
+
+        return None
+
+    def exit(self, robot: Robot) -> None:
+        pass
+
+
+class WalkingFourState(State):
+    """4-Leg/Wave Crawl Gait State."""
+
+    def __init__(
+        self,
+        duration: float = cfg.standard_duration,
+        tickpersec: float = cfg.standard_tickpersec,
+    ):
+        self.duration = duration
+        self.tickpersec = tickpersec
+        self.ticks = int(duration * tickpersec)
+        self.current_tick = 0
+        self.start_pos: Optional[dt.PosArray] = None
+
+    def enter(self, robot: Robot) -> None:
+        self.current_tick = 0
+        self.start_pos = dt.PosArray(np.copy(robot.current_pos.data))
+
+    def execute(self, robot: Robot) -> Optional[str]:
+        t = self.current_tick / float(self.ticks)
+        tick_pos_temp = []
+        dirmov = robot.vector_dirmov
+
+        for leg_id in range(6):
+            target_p = [
+                robot.center_points[leg_id][0] + dirmov[0] * cfg.step_length,
+                robot.center_points[leg_id][1] + dirmov[1] * cfg.step_length,
+                robot.center_points[leg_id][2],
+            ]
+
+            if robot.leg_state[leg_id] == "drag":
+                tick_pos_temp.append(
+                    self.start_pos[leg_id]
+                    + (robot.center_points[leg_id] - self.start_pos[leg_id]) * t
+                )
+            elif robot.leg_state[leg_id] == "step":
+                mid_point = [
+                    (self.start_pos[leg_id][0] + target_p[0]) / 2.0,
+                    (self.start_pos[leg_id][1] + target_p[1]) / 2.0,
+                    max(self.start_pos[leg_id][2], target_p[2]) + cfg.step_height,
+                ]
+                x = interpolate_bezier(
+                    t, self.start_pos[leg_id][0], mid_point[0], target_p[0]
+                )
+                y = interpolate_bezier(
+                    t, self.start_pos[leg_id][1], mid_point[1], target_p[1]
+                )
+                z = interpolate_bezier(
+                    t, self.start_pos[leg_id][2], mid_point[2], target_p[2]
+                )
+                tick_pos_temp.append([x, y, z])
+
+        tick_pos = dt.PosArray(tick_pos_temp)
+        target_rad = robot.compute_ik(tick_pos)
+
+        robot.current_pos = tick_pos
+        robot.set_joint_angles(target_rad)
+
+        self.current_tick += 1
+
+        if self.current_tick > self.ticks:
+            self.current_tick = 0
+            self.start_pos = dt.PosArray(np.copy(robot.current_pos.data))
+
+            # Rotate wave leg sequence
+            if robot.leg_state[0] == "step":
+                robot.leg_state = np.array(
+                    ["drag", "drag", "step", "drag", "drag", "drag"]
+                )
+            elif robot.leg_state[2] == "step":
+                robot.leg_state = np.array(
+                    ["drag", "drag", "drag", "step", "drag", "drag"]
+                )
+            elif robot.leg_state[3] == "step":
+                robot.leg_state = np.array(
+                    ["drag", "drag", "drag", "drag", "drag", "step"]
+                )
+            elif robot.leg_state[5] == "step":
+                robot.leg_state = np.array(
+                    ["step", "drag", "drag", "drag", "drag", "drag"]
+                )
+
+        return None
+
+    def exit(self, robot: Robot) -> None:
+        pass

states/WaveState.py

@@ -0,0 +1,140 @@
+"""
+RobotState/emotes.py - Expressive Emote States (Wave, Laola Wave)
+"""
+from __future__ import annotations
+from typing import TYPE_CHECKING, Optional
+import numpy as np
+import math
+
+import config as cfg
+import DataTypes as dt
+from states.State import State
+if TYPE_CHECKING:
+    from Robot import Robot
+
+
+class WaveEmoteState(State):
+    """Front leg wave greeting emote state."""
+
+    def __init__(
+        self,
+        cycles: int = 3,
+        duration: float = 1.5,
+        tickpersec: float = 20.0,
+        height: float = 40.0,
+        amplitude: float = 25.0,
+    ):
+        self.cycles = cycles
+        self.duration = duration
+        self.tickpersec = tickpersec
+        self.ticks = int(duration * tickpersec)
+        self.height = height
+        self.amplitude = amplitude
+
+        self.current_cycle = 0
+        self.current_tick = 0
+        self.base_pos: Optional[dt.PosArray] = None
+
+    def enter(self, robot: Robot) -> None:
+        self.current_cycle = 0
+        self.current_tick = 0
+        self.base_pos = dt.PosArray(np.copy(robot.current_pos.data))
+
+    def execute(self, robot: Robot) -> Optional[str]:
+        t = self.current_tick / float(self.ticks)
+        tick_pos = np.copy(self.base_pos.data)
+        leg_id = 0  # Front leg
+
+        cx, cy, cz = self.base_pos[leg_id]
+
+        # Wave trajectory calculation
+        y_wave = math.sin(4.0 * math.pi * t)
+        y_offset = self.amplitude * (0.5 * (y_wave + 1.0))
+        z_offset = self.height * math.sin(math.pi * t)
+
+        max_y_dev = 30.0
+        new_y = cx
+        new_y = max(cy - max_y_dev, min(cy + max_y_dev, cy + y_offset))
+
+        tick_pos[leg_id] = [cx + 10.0, new_y, cz + z_offset]
+
+        pos_array = dt.PosArray(tick_pos)
+        target_rad = robot.compute_ik(pos_array)
+
+        robot.current_pos = pos_array
+        robot.set_joint_angles(target_rad)
+
+        self.current_tick += 1
+
+        if self.current_tick >= self.ticks:
+            self.current_tick = 0
+            self.current_cycle += 1
+
+        if self.current_cycle >= self.cycles:
+            return "idle"
+
+        return None
+
+    def exit(self, robot: Robot) -> None:
+        robot.robot_state = "idle"
+
+
+class LaolaWaveEmoteState(State):
+    """Laola Wave side-to-side leg wave emote state."""
+
+    def __init__(
+        self,
+        cycles: int = 3,
+        duration: float = 1.5,
+        tickpersec: float = cfg.standard_tickpersec,
+        height: float = 10.0,
+        amplitude: float = 30.0,
+    ):
+        self.cycles = cycles
+        self.duration = duration
+        self.ticks = int(duration * tickpersec)
+        self.height = height
+        self.amplitude = amplitude
+
+        self.current_cycle = 0
+        self.current_tick = 0
+        self.base_pos: Optional[dt.PosArray] = None
+
+    def enter(self, robot: Robot) -> None:
+        self.current_cycle = 0
+        self.current_tick = 0
+        self.base_pos = dt.PosArray(np.copy(robot.current_pos.data))
+
+    def execute(self, robot: Robot) -> Optional[str]:
+        t = self.current_tick / float(self.ticks)
+        tick_pos = np.copy(self.base_pos.data)
+        wave_legs = [1, 4]
+
+        for leg_id in wave_legs:
+            cx, cy, cz = self.base_pos[leg_id]
+            phase = 0.0 if leg_id == 1 else math.pi
+
+            y_offset = self.amplitude * math.sin(2.0 * math.pi * t + phase)
+            z_offset = self.height * math.sin(2.0 * math.pi * t + phase)
+
+            tick_pos[leg_id] = [cx, cy + y_offset, cz + z_offset]
+
+        pos_array = dt.PosArray(tick_pos)
+        target_rad = robot.compute_ik(pos_array)
+
+        robot.current_pos = pos_array
+        robot.set_joint_angles(target_rad)
+
+        self.current_tick += 1
+
+        if self.current_tick >= self.ticks:
+            self.current_tick = 0
+            self.current_cycle += 1
+
+        if self.current_cycle >= self.cycles:
+            return "idle"
+
+        return None
+
+    def exit(self, robot: Robot) -> None:
+        robot.robot_state = "idle"

states/__init__.py

@@ -0,0 +1,13 @@
+"""
+states/__init__.py - State Machine Initialization
+"""
+from states.State import State, STATE_REGISTRY
+from states.IdleState import IdleState
+from states.WalkingState import WalkingState, WalkingFourState
+
+# Register available state instances
+STATE_REGISTRY["idle"] = IdleState()
+STATE_REGISTRY["walking"] = WalkingState()
+STATE_REGISTRY["walking_four"] = WalkingFourState()
+
+__all__ = ["State", "STATE_REGISTRY", "IdleState", "WalkingState", "WalkingFourState"]

states/ml_walking.py

@@ -0,0 +1,63 @@
+import numpy as np
+import math
+import torch
+import config as cfg
+import GlobalVariables as gv
+import kinematics as kin
+import DataTypes as dt
+
+
+class MLWalkingState:
+    def __init__(self, model_path: str | None = None):
+        self.model_path = model_path or "ml/checkpoints/ppo_joint_command.zip"
+        self.model = None
+        self._load_model()
+        self.step_count = 0
+
+    def _load_model(self):
+        try:
+            from stable_baselines3 import PPO
+        except ImportError:
+            print("stable-baselines3 not installed: ML walking will not be available.")
+            self.model = None
+            return
+
+        try:
+            self.model = PPO.load(self.model_path)
+            print(f"Loaded ML walking model from {self.model_path}")
+        except Exception as exc:
+            print(f"Failed to load ML walking model: {exc}")
+            self.model = None
+
+    def infer_joint_commands(self, current_rad: dt.RadArray, direction: np.ndarray) -> dt.RadArray:
+        if self.model is None:
+            return current_rad
+
+        observation = np.concatenate([current_rad.data.flatten(), direction]).astype(np.float32)
+        action, _ = self.model.predict(observation, deterministic=True)
+        action = np.clip(action, -1.0, 1.0).astype(np.float32)
+
+        new_rad = np.clip(
+            current_rad.data.flatten() + action * math.radians(5.0),
+            -math.pi,
+            math.pi,
+        ).reshape((6, 3))
+        return dt.RadArray(new_rad)
+
+    def update(self, ctx, intent, dt_step):
+        if not intent.walk:
+            return "idle"
+
+        direction = np.array([intent.move_vector.x, intent.move_vector.y, 0.0, intent.turn], dtype=np.float32)
+        target_rad = self.infer_joint_commands(ctx.current_rad, direction)
+
+        if ctx.robotCommunication:
+            ctx.robotCommunication.send_motion(target_rad)
+
+        if ctx.shared_sim:
+            ctx.shared_sim.updatePos(target_rad)
+            ctx.shared_sim.step()
+
+        ctx.current_rad = target_rad
+        self.step_count += 1
+        return None