Hardware Connection fix

readme was outdated
env rewards got a penalty for standing still while it should move instead of 0 reward

.gitignore

@@ -1,10 +1,14 @@
 # Ignore dependency folders
 node_modules/
 .venv/
-.vs/
+.vscode/
 __pycache__/
-ml/checkpoints/
-ml/tensorboard/
+ml/checkpoints/*
+!ml/checkpoints/jackbot_kinematics_base.zip
+!ml/checkpoints/jackbot_ppo6_149994_steps.zip
+!ml/checkpoints/jackbot_ppo6_199992_steps.zip
+!ml/checkpoints/jackbot_ppo6_499980_steps.zip
+#ml/logs/
 
 # Ignore environment files with private passwords/keys
 .env

ArduinoCommunication.py

@@ -4,14 +4,14 @@ import numpy as np
 import serial
 import time
 
-import config as cfg
+from config import cfg
 import DataTypes as dt
 
 class ArduinoCommunication(Thread):
-    def __init__(self, port=cfg.port, baudrate=cfg.baudrate, timeout=cfg.timeout):
+    def __init__(self, port=cfg.port, baudrate=cfg.baudrate, timeout=cfg.comm_timeout):
         super().__init__()
-        self.daemon = True  # Thread schließt sich beim Programmende
-        self.serial_conn  = serial.Serial(port, baudrate, timeout=timeout)
+        self.daemon = True  # Thread schlie�t sich beim Programmende
+        self.serial_conn = serial.Serial(port, baudrate, timeout=timeout)
         time.sleep(2)  # Warten bis Arduino ready
 
         self.command_queue = Queue()
@@ -19,6 +19,9 @@ class ArduinoCommunication(Thread):
         self.running = Event()
         self.running.set()
 
+    def send_motion(self, radial_array):
+        self.write(radial_array)
+
     def run(self):
         while self.running.is_set():
             # 1. Befehle senden
@@ -70,5 +73,12 @@ class ArduinoCommunication(Thread):
 
     def stop(self):
         self.running.clear()
-        self.join()
-        self.serial_conn.close()
+        try:
+            self.join(timeout=0.5)
+        except RuntimeError:
+            pass
+        if self.serial_conn and self.serial_conn.is_open:
+            self.serial_conn.close()
+
+    def close(self):
+        self.stop()

Controller.py

@@ -1,200 +0,0 @@
-import pygame
-import math
-import GlobalVariables as gv
-
-def controller():
-    pygame.init()
-    
-    # This is a simple class that will help us print to the screen.
-    # It has nothing to do with the joysticks, just outputting the
-    # information.
-    class TextPrint:
-        def __init__(self):
-            self.reset()
-            self.font = pygame.font.Font(None, 25)
-
-        def tprint(self, screen, text):
-            text_bitmap = self.font.render(text, True, (0, 0, 0))
-            screen.blit(text_bitmap, (self.x, self.y))
-            self.y += self.line_height
-
-        def reset(self):
-            self.x = 10
-            self.y = 10
-            self.line_height = 15
-
-        def indent(self):
-            self.x += 10
-
-        def unindent(self):
-            self.x -= 10
-
-
-    def main():
-        
-        def normalize_controller_input(x, y):
-            # Berechne die Laenge des Vektors
-            magnitude = math.sqrt(x**2 + y**2)
-    
-            # Wenn die Laenge groesser als 1 ist, normalisiere sie
-            if magnitude > 1.0:
-                x /= magnitude
-                y /= magnitude
-    
-            return x, y
-    
-        # Set the width and height of the screen (width, height), and name the window.
-        screen = pygame.display.set_mode((500, 700))
-        pygame.display.set_caption("Controller Inputs")
-
-        # Used to manage how fast the screen updates.
-        clock = pygame.time.Clock()
-
-        # Get ready to print.
-        text_print = TextPrint()
-
-        # This dict can be left as-is, since pygame will generate a
-        # pygame.JOYDEVICEADDED event for every joystick connected
-        # at the start of the program.
-        joysticks = {}
-
-        done = False
-        while not done:
-            # Event processing step.
-            # Possible joystick events: JOYAXISMOTION, JOYBALLMOTION, JOYBUTTONDOWN,
-            # JOYBUTTONUP, JOYHATMOTION, JOYDEVICEADDED, JOYDEVICEREMOVED
-            for event in pygame.event.get():
-                if event.type == pygame.QUIT:
-                    done = True  # Flag that we are done so we exit this loop.
-
-                if event.type == pygame.JOYBUTTONDOWN:
-                    print("Joystick button pressed.")
-                    if event.button == 3 and gv.emote is None:
-                        gv.emote = "wave"
-                    if event.button == 2 and gv.emote is None:
-                        gv.robotCommunication.send_text("Huhrensohn")
-                    if event.button == 0:
-                        joystick = joysticks[event.instance_id]
-                        gv
-                        if joystick.rumble(0, 0.7, 500):
-                            print(f"Rumble effect played on joystick {event.instance_id}")
-
-                if event.type == pygame.JOYBUTTONUP:
-                    print("Joystick button released.")
-
-                # Handle hotplugging
-                if event.type == pygame.JOYDEVICEADDED:
-                    # This event will be generated when the program starts for every
-                    # joystick, filling up the list without needing to create them manually.
-                    joy = pygame.joystick.Joystick(event.device_index)
-                    joysticks[joy.get_instance_id()] = joy
-                    print(f"Joystick {joy.get_instance_id()} connencted")
-
-                if event.type == pygame.JOYDEVICEREMOVED:
-                    del joysticks[event.instance_id]
-                    print(f"Joystick {event.instance_id} disconnected")
-
-            # Drawing step
-            # First, clear the screen to white. Don't put other drawing commands
-            # above this, or they will be erased with this command.
-            screen.fill((255, 255, 255))
-            text_print.reset()
-
-            # Get count of joysticks.
-            joystick_count = pygame.joystick.get_count()
-
-            text_print.tprint(screen, f"Number of joysticks: {joystick_count}")
-            text_print.indent()
-
-            # For each joystick:
-            for joystick in joysticks.values():
-                jid = joystick.get_instance_id()
-
-                text_print.tprint(screen, f"Joystick {jid}")
-                text_print.indent()
-
-                # Get the name from the OS for the controller/joystick.
-                name = joystick.get_name()
-                text_print.tprint(screen, f"Joystick name: {name}")
-
-                guid = joystick.get_guid()
-                text_print.tprint(screen, f"GUID: {guid}")
-
-                power_level = joystick.get_power_level()
-                text_print.tprint(screen, f"Joystick's power level: {power_level}")
-
-                # Usually axis run in pairs, up/down for one, and left/right for
-                # the other. Triggers count as axes.
-                axes = joystick.get_numaxes()
-                text_print.tprint(screen, f"Number of axes: {axes}")
-                text_print.indent()
-
-                for i in range(axes):
-                    axis = joystick.get_axis(i)
-                    gv.vector_dirmov = [axis]
-                    text_print.tprint(screen, f"Axis {i} value: {axis:>6.3f}")
-                text_print.unindent()
-                
-                # Get Movement Direction Vector
-                # Left stick (translation)
-                vx = joystick.get_axis(1)   # forward/back
-                vy = joystick.get_axis(0)   # strafe
-
-                vx, vy = normalize_controller_input(vx, vy)
-
-                # Right stick X (rotation) - adjust axis index if needed
-                omega = joystick.get_axis(2)
-
-                # Deadzone for rotation
-                if abs(omega) < 0.2:
-                    omega = 0.0
-
-                # Apply axis snapping (only for translation)
-                if 0.8 < vx and -0.2 < vy < 0.2:
-                    vx, vy = 1.0, 0.0
-                elif vx < -0.8 and -0.2 < vy < 0.2:
-                    vx, vy = -1.0, 0.0
-                elif -0.2 < vx < 0.2 and 0.8 < vy:
-                    vx, vy = 0.0, 1.0
-                elif -0.2 < vx < 0.2 and vy < -0.8:
-                    vx, vy = 0.0, -1.0
-
-                # Final movement vector
-                gv.vector_dirmov = [vx, vy, omega]
-
-                # Robot state
-                if abs(vx) < 0.2 and abs(vy) < 0.2 and abs(omega) < 0.2:
-                    gv.robot_state = "idle"
-                else:
-                    gv.robot_state = "walking"
-
-                buttons = joystick.get_numbuttons()
-                text_print.tprint(screen, f"Number of buttons: {buttons}")
-                text_print.indent()
-
-                for i in range(buttons):
-                    button = joystick.get_button(i)
-                    text_print.tprint(screen, f"Button {i:>2} value: {button}")
-                text_print.unindent()
-
-                hats = joystick.get_numhats()
-                text_print.tprint(screen, f"Number of hats: {hats}")
-                text_print.indent()
-
-                # Hat position. All or nothing for direction, not a float like
-                # get_axis(). Position is a tuple of int values (x, y).
-                for i in range(hats):
-                    hat = joystick.get_hat(i)
-                    text_print.tprint(screen, f"Hat {i} value: {str(hat)}")
-                text_print.unindent()
-
-                text_print.unindent()
-
-            # Go ahead and update the screen with what we've drawn.
-            pygame.display.flip()
-
-            # Limit to 30 frames per second.
-            clock.tick(30)
-            
-    main()
-    pygame.quit()

EspCommunication.py

@@ -5,7 +5,7 @@ import numpy as np
 from threading import Thread, Event
 from queue import Queue
 
-import config as cfg
+from config import cfg
 import DataTypes as dt
 
 # ============================================================
@@ -25,10 +25,14 @@ HEADER_FMT  = "<B B H I" # ID, Ver, Len, Timestamp
 TLV_FMT = "<B H"
 
 class ESP32Communication(Thread):
-    def __init__(self, host=cfg.esp32_ip, port=cfg.esp32_port):
+    def __init__(self, host=None, port=None, ip=None):
         super().__init__(daemon=True)
+        if ip is not None:
+            host = ip
+        host = cfg.esp32_ip if host is None else host
+        port = cfg.esp32_port if port is None else port
         self.addr = (host, port)
-        
+
         # UDP Socket - Zero Lag
         self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
         
@@ -126,4 +130,10 @@ class ESP32Communication(Thread):
 
     def stop(self):
         self.running.clear()
-        self.sock.close()
+        try:
+            self.sock.close()
+        except Exception:
+            pass
+
+    def close(self):
+        self.stop()

GlobalVariables.py

@@ -1,39 +0,0 @@
-import numpy as np
-from ArduinoCommunication import ArduinoCommunication
-from EspCommunication import ESP32Communication
-from simulation import Simulation
-import DataTypes as dt
-import config as cfg
-from robot_init import init_deg, init90_deg, init_pos, center_points
-
-# Globale Variablen
-robotCommunication = None
-shared_sim = None
-emote = None
-display_text = ""
-
-if cfg.sim:
-    shared_sim = Simulation()
-else:
-    if cfg.arduinoConnection:
-        robotCommunication = ArduinoCommunication()
-    else:
-        robotCommunication = ESP32Communication()
-
-vector_dirmov = [0, 0, 0]  # Direction Movement [vx, vy, omega]
-current_rad: dt.RadArray  # Current Rad
-current_pos: dt.PosArray  # Current Position
-
-# current_deg: dt.DegArray  # Current Degrees
-# legarray: dt.DegArray  # Working Leg Array
-# target_pos: dt.PosArray  # Target Position
-
-robot_state = "idle"  # Current State
-# Init for 6 Legs
-#leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
-# Init for 4 Legs
-leg_state = np.array(["step", "drag", "drag", "drag", "drag", "drag"])
-control_pause: bool = False
-
-# Initilize mit Start Position
-## Shared init pose definitions are imported from robot_init.py

README.md

@@ -1,21 +1,123 @@
+# JackBot — Hexapod Control, Simulation & RL Framework
 
-# JackBot — Hexapod Control & Simulation
+JackBot is a modular Python framework for controlling a six-legged robot in simulation or on hardware. The current workspace reflects a matured control stack with a PyBullet-based simulator, a GUI-driven manual runtime, and a reinforcement-learning pipeline for PPO training and evaluation.
 
-JackBot is a Python project for controlling and simulating a six-legged hexapod robot.
-It uses IKPy for inverse kinematics, PyBullet for optional simulation, and can send joint commands to ESP32 or Arduino hardware.
+---
 
-## Requirements
+## What JackBot Is
 
-- Python 3.12 is safest for `pygame` compatibility
-- Required Python packages:
-	- `numpy`
-	- `pygame`
-	- `ikpy`
-	- `pybullet`
-	- `pyserial`
-	- `matplotlib`
+JackBot is a Python-based hexapod project that combines:
 
-## Setup
+* a robot control stack for a six-legged walking robot,
+* a physics simulator (`simulation.py`) for testing in software before using real hardware,
+* a reinforcement learning pipeline that trains locomotion policies with PPO and behavioral cloning,
+* and a GUI/gamepad input layer for manual operation.
+
+The project can currently:
+
+* run the robot in a PyBullet simulation,
+* accept commands from a GUI or gamepad,
+* stream target joint positions to physical hardware (ESP32 / Arduino),
+* train PPO models and evaluate saved checkpoints,
+* generate kinematics-based teacher data for behavioral cloning.
+
+---
+
+## Why the Project Exists
+
+A hexapod is hard to control manually because each leg has multiple joints and the robot must maintain balance while moving. The project uses a robot model and physics simulation as a testbed for learning motion strategies, refining kinematic control, and validating behavior before applying commands to real hardware.
+
+In the current codebase, the practical workflow is:
+
+1. Start the robot in either simulation or hardware mode.
+2. Feed motion commands from GUI/gamepad or a training environment.
+3. Convert target foot positions into joint targets using inverse kinematics.
+4. Apply these targets to the active backend.
+5. Train or evaluate PPO policies against the simulated robot.
+
+---
+
+## Key Components
+
+* **Robot abstraction (`Robot.py`)**: Central handler for the active backend, current joint state, IK/FK helpers, gait handling, and RL action application.
+* **Kinematics (`kinematics.py`)**: Converts target foot positions into joint angles using IKPy chains for each leg.
+* **Simulation (`simulation.py`)**: Manages PyBullet scene loading, stepping, joint actuation, base pose queries, and physics telemetry.
+* **Inputs (`inputs/`)**: Receives commands from GUI sliders, Pygame gamepads, and random command generation.
+* **ML Subsystem (`ml/`)**: Contains the Gymnasium environment (`env.py`), PPO training (`run_train.py`), evaluation (`run_eval.py`), BC pretraining (`pretrain_bc.py`), callbacks, and metrics overlays.
+* **States (`states/`)**: Contains additional state-machine classes such as `IdleState`, `WalkingState`, and `WaveState`. These exist in the repository, but the current active runtime in `main.py` does not currently drive them through `Robot.tick()`.
+
+---
+
+## Project Architecture
+
+```text
+JackBot/
+├── main.py                  # Manual runtime entry point
+├── Robot.py                 # Unified robot wrapper + backend selection
+├── config.py                # Central runtime configuration
+├── kinematics.py            # IK/FK helpers
+├── simulation.py            # PyBullet scene / physics manager
+├── robot_init.py            # Initial joint definitions and center points
+├── DataTypes.py             # Typed arrays for positions / angles
+├── JackBotUrdf.urdf        # Robot URDF
+│
+├── states/                  # State classes present in the project
+│   ├── State.py
+│   ├── IdleState.py
+│   ├── WalkingState.py
+│   ├── WaveState.py
+│   ├── ml_walking.py
+│   └── __init__.py
+│
+├── inputs/                  # Command input providers
+│   ├── InputProvider.py
+│   ├── PygameController.py
+│   ├── RandomeInputProvider.py
+│   └── RandomInputProvider.py
+│
+├── gui/                     # GUI/dashboard control layer
+│   └── MainWindow.py
+│
+├── ml/                      # Gymnasium RL training + evaluation stack
+│   ├── env.py               # RL environment + reward logic
+│   ├── callbacks.py         # SB3 callbacks for logs/curriculum
+│   ├── pretrain_bc.py       # Behavior cloning pretraining
+│   ├── run_train.py         # PPO training entry point
+│   ├── run_eval.py          # Evaluation of saved checkpoints
+│   ├── run_eval_training.py # Kinematics-mode benchmark script
+│   ├── MetricsOverlay.py    # 3D in-scene metrics HUD
+│   └── checkpoints/        # Model checkpoints + saved runs
+│
+├── EspCommunication.py      # ESP32 UDP communication layer
+├── ArduinoCommunication.py  # Arduino serial communication layer
+├── Helper Scripts/          # Utility scripts
+│   ├── FindCenterPoints.py
+│   └── torqueCalc.py
+│
+├── requirements.txt
+├── README.md
+└── ...
+```
+
+### Important current-state note
+
+The repository contains a few pieces that are still present but not fully connected to the current runtime:
+
+* `states/WalkingState.py`, `states/WaveState.py`, and `states/ml_walking.py` exist, but `main.py` currently drives `Robot.tick()` directly instead of routing through the active `STATE_REGISTRY`.
+* `inputs/RandomeInputProvider.py` exists, but the active runtime does not currently use it directly. The GUI resolves random walking behavior in `gui/MainWindow.py`.
+* `WaveEmoteState` / `LaolaWaveEmoteState` are defined in `states/WaveState.py`, but they are not currently registered in `states/__init__.py`.
+
+---
+
+## System Requirements
+
+* **Python:** 3.12 recommended
+* **OS:** Windows 10/11 or Linux
+* **Dependencies:** `pybullet`, `gymnasium`, `stable-baselines3`, `torch`, `numpy`, `pygame`, `ikpy`, `pyserial`, `matplotlib`, `dearpygui`
+
+---
+
+## Installation & Setup
 
 ### Linux (Bash)
 
@@ -29,81 +131,353 @@ pip install -r requirements.txt
 ### Windows (PowerShell)
 
 ```powershell
-python3.12 -m venv .venv
+Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope Process
+python -m venv .venv
 .\.venv\Scripts\Activate.ps1
 python -m pip install --upgrade pip setuptools wheel
 pip install -r requirements.txt
 ```
 
-If PowerShell blocks activation:
+---
 
-```powershell
-Set-ExecutionPolicy -Scope Process -ExecutionPolicy RemoteSigned
-.\.venv\Scripts\Activate.ps1
-```
+## Usage Guide
 
-## Run
+## Manual Control & Hardware Streaming (`main.py`)
 
-From the repository root:
+`main.py` is the main operational entry point for driving the robot manually through the GUI or a connected gamepad.
 
-```
+To launch:
+
+```bash
 python main.py
 ```
 
-If Linux breaks with 
-```
-ExampleBrowserThreadFunc started
-X11 functions dynamically loaded using dlopen/dlsym OK!
+### Configuration (`config.py`)
 
-        cannot connect to X server
-```
-run:
-```
-export DISPLAY=:0
-python main.py
+Before launching `main.py`, edit `config.py` to select the backend and connection targets.
+
+* **Backend Selection (`cfg.backend`)**:
+  * `BackendType.SIMULATION`: run inside a PyBullet window
+  * `BackendType.ESP32`: stream motion over UDP to an ESP32
+  * `BackendType.ARDUINO`: stream motion over serial to an Arduino
+
+Supported configuration values in the current code include:
+
+* `backend`
+* `urdf_path`
+* `port`, `baudrate`
+* `esp32_ip`, `esp32_port`
+* `tick_rate_hz`, `step_duration`
+* `step_height`, `step_length`
+
+### GUI input sources
+
+The current GUI (`gui/MainWindow.py`) exposes these input sources:
+
+* `Gamepad`
+* `GUI Sliders`
+* `Random Walk`
+
+The active `main.py` runtime resolves commands through `resolve_active_command(...)` and then passes the resulting `vector_dirmov` directly to `Robot.tick()`.
+
+---
+
+## Machine Learning Pipeline (`ml/`)
+
+The ML subsystem currently uses:
+
+* **Gymnasium** as the environment API
+* **PyBullet** as the physics engine
+* **Stable-Baselines3 PPO** as the learning algorithm
+* **Curriculum phases** to gradually expose more difficult command spaces
+
+### Current ML scripts
+
+* `ml/run_train.py` — PPO training entry point
+* `ml/run_eval.py` — phase-based model evaluation
+* `ml/run_eval_training.py` — reward benchmark for kinematics mode
+* `ml/pretrain_bc.py` — collect teacher data from the kinematics solver and pretrain a base policy
+* `ml/env.py` — environment definition and reward logic
+* `ml/callbacks.py` — reward logging and curriculum callbacks
+
+### Training flow
+
+Training is launched with:
+
+```bash
+python ml/run_train.py --total-timesteps 1500000 --gui
 ```
 
-## Configuration
+The current training script creates a vectorized PPO environment, optionally loads a pre-trained checkpoint, and saves results into `ml/checkpoints/` and `ml/logs/`.
 
-Edit `config.py` before running:
+### Evaluation flow
 
-- `sim = True` to enable PyBullet simulation
-- `sim = False` to use hardware control
-- `arduinoConnection = True` to use Arduino
-- `arduinoConnection = False` to use ESP32
-- `port` and `baudrate` for Arduino
-- `esp32_ip` and `esp32_port` for ESP32
-- `urdf_path = "JackBotUrdf.urdf"`
+Evaluation is launched with:
 
-## Project structure
+```bash
+python ml/run_eval.py --model ml/checkpoints/jackbot_kinematics_base.zip --episodes-per-phase 5 --gui
+```
 
-- `main.py` — main application entry point
-- `Controller.py` — Pygame-based controller and input display
-- `kinematics.py` — IKPy forward/inverse kinematics
-- `simulation.py` — PyBullet simulation wrapper
-- `GlobalVariables.py` — shared runtime state and comms
-- `DataTypes.py` — typed arrays and control intent
-- `RobotState/idle.py` — idle robot state
-- `RobotState/walking.py` — walking robot state
-- `EspCommunication.py` — ESP32 communication
-- `ArduinoCommunication.py` — Arduino communication
-- `JackBotUrdf.urdf` — robot model file
+This script runs a multi-phase deterministic evaluation suite with fixed command vectors for:
 
-## Notes
+* `FORWARD`
+* `TURN_AND_DIRECTION`
+* `OMNI_DIRECTION`
+* `FULL_COMMAND`
 
-- `GlobalVariables.py` initializes either `Simulation()` or the selected hardware comm class.
-- `DataTypes.py` declares `PosArray`, `DegArray`, `RadArray`, `RobotCommand`, and `ControlIntent`.
-- `Controller.py` updates `gv.vector_dirmov` and `gv.robot_state` from joystick input.
-- `kinematics.py` loads leg chains from `JackBotUrdf.urdf` and computes IK.
+### Behavioral cloning pretraining
 
-## Troubleshooting
+The repository also contains a behavioral cloning route:
 
-- If `pygame` installation fails on Windows, use Python 3.12 and upgrade `pip setuptools wheel` first.
-- If `python` is not found on Windows, install Python and enable "Add Python to PATH".
-- If the program crashes on startup, verify `JackBotUrdf.urdf` path and `config.py` settings.
+```bash
+python ml/pretrain_bc.py --num-samples 100000 --epochs 15 --save-path ml/checkpoints/jackbot_kinematics_base.zip
+```
 
-## Suggested improvements
+This script gathers `(observation, action)` data from the kinematics teacher and trains a PPO policy to act as a learned base model.
 
-- Add a `requirements.txt` or `pyproject.toml`.
-- Add a `LICENSE` file before sharing the project.
-- Document hardware wiring and packet formats for ESP32/Arduino.
+---
+
+## What the Reward Is Trying to Teach
+
+The reward function in `ml/env.py` is designed to teach several things at once:
+
+* survival and stability,
+* following the commanded direction,
+* avoiding lateral drift,
+* maintaining base height,
+* reducing abrupt control changes,
+* staying close to a stable stance when the command is zero.
+
+The environment uses several reward components, including:
+
+* height reward,
+* stability reward,
+* pose closeness reward,
+* smoothness reward,
+* linear velocity tracking,
+* angular velocity tracking,
+* jitter penalty,
+* stand-by penalty when the command is zero.
+
+### Detailed reward structure
+
+The reward is built step by step inside `JackBotEnv._compute_reward()` in `ml/env.py`. It is not a single binary success signal; it is a dense shaping signal that rewards good behavior continuously throughout the episode.
+
+At each control step, the environment measures:
+
+* current body height,
+* roll and pitch angles,
+* current joint configuration,
+* measured linear and angular velocities,
+* the active command vector `[vx, vy, omega]`,
+* the difference between the current action and the previous actions.
+
+From these values, it computes a set of sub-rewards:
+
+* `height reward`: a Gaussian-style reward based on how close the robot body is to the target height.
+* `stability reward`: a reward for keeping roll/pitch low and the body well balanced.
+* `pose closeness reward`: rewards staying near the default standing joint posture.
+* `smoothness reward`: rewards actions that change gradually instead of abruptly.
+* `linear velocity reward`: encourages the robot to move in the commanded direction and speed.
+* `angular velocity reward`: rewards matching the commanded turning rate.
+
+The environment then adds two kinds of correction terms:
+
+* `jitter penalty`: subtracts a small amount when action changes are noisy or jerky.
+* `stand_penalty`: when the command is effectively zero, penalizes unintended movement and yaw drift.
+
+This makes the reward function behave as a soft guidance system: the agent gets a steady gradient that says “this is closer to what we want” or “this is worse than the desired behavior.”
+
+### Standing mode vs. walking mode
+
+The reward code handles two cases differently:
+
+#### 1. Standing mode
+When the command vector is close to zero (`cmd_norm < 0.05` and `abs(cmd_yaw) < 0.05`), the agent is not supposed to move much. In that case the reward focuses on:
+
+* keeping the body at the correct height,
+* staying stable,
+* maintaining a clean posture,
+* staying smooth.
+
+A small `stand_penalty` is then applied to discourage unintended speed and yaw drift while the robot is supposed to hold position.
+
+#### 2. Walking mode
+When a non-zero command is active, the robot is rewarded for moving in the intended direction and turning at the requested rate. The reward then emphasizes:
+
+* matching the commanded linear velocity,
+* matching the commanded yaw rate,
+* continuing to maintain height and stability,
+* staying smooth in its control inputs.
+
+If the command expects movement but the robot is effectively still, the code now applies a stronger `stillness_penalty` instead of a neutral reward. This means standing while the command says to move is explicitly discouraged.
+
+### Why the reward is shaped this way
+
+The goal is not just to teach the robot to stay alive. The reward is designed so that a PPO agent learns multiple useful habits at once:
+
+* do not collapse or tip over,
+* keep the body at a sensible height,
+* follow directional commands,
+* avoid unstable oscillations,
+* avoid overreactive control jumps,
+* stay near a normal standing pose when no motion is requested.
+
+This is why the environment is not based on a single sparse reward such as “+1 for success, 0 otherwise.” Instead, it uses dense reward shaping so the policy receives useful feedback on every step.
+
+### Are the penalties real penalties?
+
+Yes — in the reward function they are real negative contributions. For example:
+
+* `jitter_penalty` subtracts from the step reward when action changes are too abrupt,
+* `stand_penalty` subtracts when the robot moves unnecessarily while standing,
+* `stillness_penalty` subtracts when movement is commanded but the robot remains essentially frozen.
+
+In the current implementation, a non-zero command that is not followed by meaningful motion now yields an explicit penalty instead of a neutral reward. This is the main behavior change requested for the training setup: standing while the command says to move is now actively discouraged.
+
+The code does this:
+
+```python
+final_reward = step_reward + jitter_penalty + alive_bonus
+```
+
+That means:
+
+* negative penalty terms can now reduce the reward below zero,
+* the reward is no longer clipped to zero in this path,
+* and the episode is still only ended by the separate failure check in `_update_robot_failure()`.
+
+So the answer is:
+
+* the penalties are real reward penalties,
+* they are now strong enough to discourage command-mismatch behavior,
+* and the actual terminal condition remains the failure check in `_update_robot_failure()`.
+
+### What actually ends an episode?
+
+The episode ends when the robot is considered failed, not when a reward penalty is applied. In `ml/env.py`, the environment marks the robot as failed if:
+
+* it is too tilted (`roll` or `pitch` exceed the configured failure threshold), or
+* it has collapsed below a minimum body-height threshold.
+
+That is a hard termination condition. In other words:
+
+* reward penalties discourage bad behavior,
+* failure conditions stop the episode when the robot is clearly unstable or collapsed.
+
+### Practical interpretation
+
+A good mental model is:
+
+* the reward function teaches the robot what “good locomotion” looks like,
+* the failure check prevents the robot from continuing when it is physically broken or unstable,
+* and the curriculum gradually increases the difficulty of the commands as the robot becomes more capable.
+
+This combination is a common reinforcement-learning setup for locomotion: dense rewards shape the desired behavior, while hard failure conditions protect the training process from degenerate states.
+
+This reward shaping encourages the robot to learn locomotion patterns rather than simply freezing in place.
+
+---
+
+## Curriculum / Phase Progression
+
+The environment currently uses the following curriculum stages:
+
+* **STAND_ONLY**
+* **FORWARD**
+* **TURN_AND_DIRECTION**
+* **OMNI_DIRECTION**
+* **FULL_COMMAND**
+
+The training environment gradually advances phase complexity based on survival, stability, and movement metrics. `ml/callbacks.py` includes custom curriculum-handling logic for the online learning loop.
+
+---
+
+## How the Control Loop Works in Practice
+
+A simplified view of the current runtime is:
+
+1. `main.py` starts the joystick controller process and opens the GUI.
+2. The GUI resolves active commands (`Gamepad`, slider, or random walk).
+3. `Robot.tick()` receives the current motion vector and applies it through the active backend.
+4. In simulation mode, `Simulation.step()` advances the PyBullet world.
+5. In training/evaluation mode, `JackBotEnv.step()` computes reward, updates curriculum, and returns observations.
+6. PPO uses the observation/action loop to improve the policy.
+
+---
+
+## What Makes this Project Useful
+
+This repository is useful because it combines several layers that are often separate:
+
+* robot control and kinematics,
+* physics simulation,
+* command input sources,
+* RL environment construction,
+* PPO training and evaluation,
+* hardware communication layers.
+
+For a newcomer, the easiest way to think about the project is:
+
+- `main.py` is the manual control entry point,
+- `Robot.py` is the core robot wrapper,
+- `simulation.py` is the physics layer,
+- `ml/env.py` is the environment interface for RL,
+- `ml/run_train.py`, `ml/run_eval.py`, and `ml/pretrain_bc.py` are the main ML workflows.
+
+---
+
+## Current runtime notes and caveats
+
+### State system status
+
+The state classes are present in the repository, but the current runtime path is not using them as the main control loop:
+
+* `Robot.tick()` currently updates the robot directly from `vector_dirmov`
+* the `STATE_REGISTRY` exists, but it is not the path used by the current `main.py` execution flow
+* the state subsystem remains partially implemented and should be treated as a legacy or optional extension
+
+### Input source status
+
+The repository currently includes both:
+
+* `inputs/PygameController.py` for gamepad input
+* `gui/MainWindow.py` for selecting `Gamepad`, `GUI Sliders`, and `Random Walk`
+
+The standalone random input provider file is present, but it is not the path currently used by `main.py`.
+
+### Hardware communication status
+
+The hardware communication classes still exist:
+
+* `EspCommunication.py` for ESP32 UDP
+* `ArduinoCommunication.py` for Arduino serial
+
+They are available via `cfg.backend`, but they are not the primary path in the current example workflows shown here.
+
+---
+
+## Helper Scripts
+
+### `Helper Scripts/FindCenterPoints.py`
+
+This script appears to be a utility for analyzing kinematic center point data and related foot-placement experiments. It is not part of the active runtime path.
+
+### `Helper Scripts/torqueCalc.py`
+
+This is a stand-alone Tkinter utility for estimating servo torque requirements based on robot mass, leg dimensions, and safety factor. It is a design/support script rather than part of the main robot runtime.
+
+---
+
+## Bottom Line
+
+The current codebase is a working hybrid of:
+
+* real-time robot control,
+* PyBullet simulation,
+* GUI/manual input,
+* PPO-based RL training and evaluation,
+* partial state-machine scaffolding,
+* hardware communication support.
+
+The information in this README has been updated to match the current files in the workspace, especially around the true ML entry points, active runtime flow, and the fact that some older state-helper modules are present but not currently wired into the main execution path.

Robot.py

@@ -0,0 +1,296 @@
+"""
+Robot.py - Central Robot Control, Kinematics, State Machine & Hardware Abstraction
+"""
+from typing import Protocol, Optional, Union, Tuple, List
+import numpy as np
+import math
+
+from states import STATE_REGISTRY
+from states.State import State
+import DataTypes as dt
+import kinematics as kin
+import robot_init as ri
+from config import cfg, BackendType
+
+from simulation import Simulation
+from EspCommunication import ESP32Communication
+from ArduinoCommunication import ArduinoCommunication
+
+
+class RobotBackend(Protocol):
+    """Protocol defining hardware abstraction for both Simulation and Hardware backends."""
+    def send_angles(self, rad_array: dt.RadArray) -> None: ...
+    def step_simulation(self) -> None: ...
+    def hard_reset_joints(self, target_angles: np.ndarray) -> None: ...
+    def reset_base(self, position: Optional[List[float]] = None, orientation: Optional[List[float]] = None) -> None: ...
+    def get_joint_angles(self) -> np.ndarray: ...
+    def get_base_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]: ...
+    def get_base_velocity(self) -> Tuple[List[float], List[float]]: ...
+    def cleanup(self) -> None: ...
+
+
+class HardwareBackend:
+    """Backend for physical ESP32 or Arduino microcontrollers."""
+    def __init__(self, comm_channel):
+        self.comm_channel = comm_channel
+        self.internal_angles = np.zeros(18, dtype=np.float32)
+
+    def send_angles(self, rad_array: dt.RadArray) -> None:
+        self.internal_angles = rad_array.data.flatten().copy()
+        if self.comm_channel:
+            self.comm_channel.send_motion(rad_array)
+
+    def step_simulation(self) -> None:
+        pass
+
+    def hard_reset_joints(self, target_angles: np.ndarray) -> None:
+        self.internal_angles = target_angles.flatten().copy()
+
+    def reset_base(self, position: Optional[List[float]] = None, orientation: Optional[List[float]] = None) -> None:
+        pass
+
+    def get_joint_angles(self) -> np.ndarray:
+        return self.internal_angles.copy()
+
+    def get_base_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]:
+        return [0.0, 0.0, 0.122], (0.0, 0.0, 0.0)
+
+    def get_base_velocity(self) -> Tuple[List[float], List[float]]:
+        return [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]
+
+    def cleanup(self) -> None:
+        if self.comm_channel:
+            if hasattr(self.comm_channel, 'stop'):
+                self.comm_channel.stop()
+            elif hasattr(self.comm_channel, 'close'):
+                self.comm_channel.close()
+
+
+class PyBulletBackend:
+    """Backend mapping Robot operations directly to PyBullet simulation engine."""
+    def __init__(self, sim_instance: Simulation):
+        self.sim = sim_instance
+
+    def send_angles(self, rad_array: dt.RadArray) -> None:
+        if self.sim:
+            self.sim.set_robot_joint_angles(rad_array)
+
+    def step_simulation(self) -> None:
+        if self.sim:
+            self.sim.step()
+
+    def hard_reset_joints(self, target_angles: np.ndarray) -> None:
+        if self.sim:
+            self.sim.hard_reset_joint_angles(target_angles)
+
+    def reset_base(self, position: Optional[List[float]] = None, orientation: Optional[List[float]] = None) -> None:
+        if self.sim:
+            self.sim.reset_robot_base(pos=position, orn=orientation)
+
+    def get_joint_angles(self) -> np.ndarray:
+        return self.sim.get_robot_joint_angles() if self.sim else np.zeros(18, dtype=np.float32)
+
+    def get_base_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]:
+        return self.sim.get_robot_pose_and_rpy() if self.sim else ([0, 0, 0], (0, 0, 0))
+
+    def get_base_velocity(self) -> Tuple[List[float], List[float]]:
+        return self.sim.get_robot_velocity() if self.sim else ([0, 0, 0], [0, 0, 0])
+
+    def cleanup(self) -> None:
+        if self.sim:
+            self.sim.disconnect()
+
+
+class Robot:
+    """
+    Unified JackBot Class.
+    Coordinates joint memory, IK solvers, procedural tripods, and backend communication.
+    """
+
+    def __init__(
+        self,
+        backend_type: Union[BackendType, RobotBackend] = BackendType.SIMULATION,
+        start_pose: str = "init_deg",
+        urdf_path: str = cfg.urdf_path,
+        mode: str = "kinematics"  # "kinematics", "residual", or "direct"
+    ):
+        self.urdf_path = urdf_path
+        self.start_pose = start_pose
+        self.mode = mode
+
+        # --- BACKEND INSTANTIATION ---
+        if isinstance(backend_type, BackendType):
+            if backend_type == BackendType.SIMULATION:
+                sim_instance = Simulation(urdf_path=self.urdf_path, use_gui=True)
+                sim_instance.load_scene()
+                self.backend: RobotBackend = PyBulletBackend(sim_instance)
+            elif backend_type == BackendType.ESP32:
+                comm = ESP32Communication(ip=cfg.esp32_ip, port=cfg.esp32_port)
+                comm.start()
+                self.backend = HardwareBackend(comm)
+            elif backend_type == BackendType.ARDUINO:
+                comm = ArduinoCommunication(port=cfg.port, baudrate=cfg.baudrate)
+                comm.start()
+                self.backend = HardwareBackend(comm)
+        else:
+            self.backend = backend_type
+
+        # Kinematics initialization
+        pose_deg = ri.init_deg if start_pose == "init_deg" else ri.init90_deg
+        self.current_rad: dt.RadArray = pose_deg.to_rad()
+        self.current_pos: dt.PosArray = kin.ikpyForward(self.current_rad)
+        self.center_points: dt.PosArray = (
+            ri.get_center_points() if hasattr(ri, "get_center_points") else ri.center_points
+        )
+
+        # RL configuration & Gait state variables
+        self.action_scale = 0.1  # Radian step scale for RL deltas
+        self.gait_phase = 0.0
+        self.vector_dirmov = [0.0, 0.0, 0.0]  # [vx, vy, omega]
+        self.robot_state = "idle"
+        self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
+
+        # State Machine Initialization
+        self.current_state_key: str = "idle"
+        self.current_state: State = STATE_REGISTRY["idle"]
+        self.current_state.enter(self)
+
+    def set_joint_angles(self, target_rad: dt.RadArray) -> None:
+        """Updates internal Python memory state and sends angles to active backend."""
+        self.current_rad = target_rad
+        if self.backend:
+            self.backend.send_angles(target_rad)
+
+    def step_sim(self) -> None:
+        if self.backend:
+            self.backend.step_simulation()
+
+    def reset_to_init(self) -> None:
+        """Resets kinematics state and forces instant joint alignment in backend."""
+        pose_deg = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg
+        self.current_rad = pose_deg.to_rad()
+        self.current_pos = kin.ikpyForward(self.current_rad)
+        self.gait_phase = 0.0
+        self.robot_state = "idle"
+        self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
+
+        init_flat = self.current_rad.data.flatten()
+        if self.backend:
+            self.backend.hard_reset_joints(init_flat)
+            self.backend.send_angles(self.current_rad)
+
+    def compute_ik(self, target_pos: dt.PosArray) -> dt.RadArray:
+        return kin.ikpyInverse(target_pos, initial_rad=self.current_rad)
+
+    def compute_fk(self, target_rad: Optional[dt.RadArray] = None) -> dt.PosArray:
+        rads = target_rad if target_rad is not None else self.current_rad
+        return kin.ikpyForward(rads)
+
+    def transition_to(self, next_state_key: str) -> None:
+        if next_state_key in STATE_REGISTRY and next_state_key != self.current_state_key:
+            self.current_state.exit(self)
+            self.current_state_key = next_state_key
+            self.current_state = STATE_REGISTRY[next_state_key]
+            self.current_state.enter(self)
+
+    def tick(self, action: Optional[np.ndarray] = None) -> None:
+        """
+        Unified control loop tick.
+        Processes commands through direct RL, residual RL, or State Machine kinematics.
+        """
+        vx, vy, omega = self.vector_dirmov
+
+        if self.mode == "direct":
+            if action is not None:
+                self.apply_rl_action(action)
+
+        elif self.mode == "residual":
+            self.step_kinematic_gait(vx, vy, omega)
+            if action is not None:
+                self.apply_rl_action_delta(action)
+
+        else:  # "kinematics" / standard State Machine execution
+            next_state_key = self.current_state.execute(self)
+            if next_state_key:
+                self.transition_to(next_state_key)
+
+        self.step_sim()
+
+    def step_kinematic_gait(self, vx: float, vy: float, omega: float) -> None:
+        """Procedural Tripod Gait solver."""
+        cmd_mag = math.hypot(vx, vy) + abs(omega)
+        if cmd_mag < 0.03:
+            target_rad = self.compute_ik(self.center_points)
+            self.set_joint_angles(target_rad)
+            return
+
+        self.gait_phase = (self.gait_phase + 0.22) % (2.0 * math.pi)
+        stride_len = cfg.step_length
+        step_height = cfg.step_height
+
+        center_data = (
+            self.center_points.data if hasattr(self.center_points, 'data') else np.array(self.center_points)
+        )
+
+        target_positions = []
+        for leg_id in range(6):
+            base_pos = np.array(center_data[leg_id], dtype=np.float32)
+            phase_offset = 0.0 if (leg_id % 2 == 0) else math.pi
+            leg_phase = (self.gait_phase + phase_offset) % (2.0 * math.pi)
+
+            lx, ly = base_pos[0], base_pos[1]
+            rot_dx = -omega * ly
+            rot_dy = omega * lx
+
+            dx_dir = vx + rot_dx
+            dy_dir = vy + rot_dy
+            dir_norm = math.hypot(dx_dir, dy_dir) + 1e-6
+
+            dx_unit = dx_dir / dir_norm
+            dy_unit = dy_dir / dir_norm
+
+            if leg_phase < math.pi:
+                # Swing phase (leg lifted in air, moving forward)
+                progress = math.cos(leg_phase)
+                lift = math.sin(leg_phase) * step_height
+                dx = progress * stride_len * dx_unit
+                dy = progress * stride_len * dy_unit
+                dz = lift
+            else:
+                # Stance phase (leg on ground, pushing body forward)
+                progress = math.cos(leg_phase - math.pi)
+                dx = -progress * stride_len * dx_unit
+                dy = -progress * stride_len * dy_unit
+                dz = 0.0
+
+            target_positions.append(base_pos + np.array([dx, dy, dz], dtype=np.float32))
+
+        target_pos_array = dt.PosArray(np.array(target_positions))
+        target_rad = self.compute_ik(target_pos_array)
+        self.set_joint_angles(target_rad)
+
+    def apply_rl_action(self, action: np.ndarray) -> None:
+        action = np.asarray(action, dtype=np.float32)
+        new_rad = dt.RadArray(data=action.reshape(self.current_rad.data.shape))
+        self.set_joint_angles(new_rad)
+
+    def apply_rl_action_delta(self, action: np.ndarray) -> None:
+        """Applies action deltas on top of joint state for Residual RL."""
+        action = np.asarray(action, dtype=np.float32)
+        scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale
+        current_flat = self.backend.get_joint_angles().flatten()
+        updated_flat = np.clip(current_flat + scaled_action, -np.pi / 2, np.pi / 2)
+        new_rad = dt.RadArray(data=updated_flat.reshape(self.current_rad.data.shape))
+        self.set_joint_angles(new_rad)
+
+    def get_observation(self, command: Optional[np.ndarray] = None) -> np.ndarray:
+        """Extracts joint positions directly from active backend."""
+        joint_angles = self.backend.get_joint_angles().flatten().astype(np.float32)
+        if command is not None:
+            cmd = np.asarray(command, dtype=np.float32).flatten()
+            return np.concatenate([joint_angles, cmd]).astype(np.float32)
+        return joint_angles
+
+    def cleanup(self) -> None:
+        if self.backend:
+            self.backend.cleanup()

RobotState.py

@@ -1,319 +0,0 @@
-import GlobalVariables as gv
-import config as cfg
-import DataTypes as dt
-import kinematics as kin
-import numpy as np
-import math
-import time
-
-# Leg_Pair 1 {leg[num] = 0,2,4}
-# Leg Pair 2 {leg[num] = 1,3,5}
-
-def emitCalculation(target_rad: dt.RadArray):
-    if gv.shared_sim != None:
-        gv.shared_sim.updatePos(target_rad)
-        gv.shared_sim.step()
-    if gv.robotCommunication != None:
-        gv.robotCommunication.send_motion(target_rad)
-    gv.current_rad = target_rad
-
-def initPos():
-    # Init Position and gv.current_rad
-    gv.current_rad = gv.init_deg.to_rad()
-    gv.current_pos = gv.center_points
-    target_rad: dt.RadArray = kin.ikpyInverse(gv.center_points)
-    emitCalculation(target_rad)
-
-
-def walking(duration=cfg.standard_duration, tickpersec=cfg.standard_tickpersec, curve_height=cfg.step_height):
-    # init Values
-    ticks = int(duration * tickpersec)
-    tick_duration = 1 / tickpersec
-    tick_pos: dt.PosArray  # aktuelle XYZ-Positionen
-    current_pos_copy: dt.PosArray = gv.current_pos
-    vx, vy, omega = gv.vector_dirmov
-
-    # Apply config scaling
-    vx *= cfg.translation_gain
-    vy *= cfg.translation_gain
-    omega *= cfg.rotation_gain
-
-    target_pos_temp = []
-
-    for i in range(6):
-        cx, cy, cz = gv.center_points[i]
-
-        # relative position (assuming body center = 0,0)
-        rx = cx
-        ry = cy
-
-        # rotation component
-        v_rot_x = -omega * ry
-        v_rot_y =  omega * rx
-
-        # combine translation + rotation
-        v_x = vx + v_rot_x
-        v_y = vy + v_rot_y
-
-        # normalize combined vector (important!)
-        length = (v_x**2 + v_y**2) ** 0.5
-        if length > 1.0:
-            v_x /= length
-            v_y /= length
-
-        target_pos_temp.append([
-            cx + v_x * cfg.step_length,
-            cy + v_y * cfg.step_length,
-            cz
-        ])
-    target_pos: dt.PosArray = target_pos_temp
-
-    def interpolate(t, p0, p1, p2):
-        return (1 - t) ** 2 * p0 + 2 * (1 - t) * t * p1 + t**2 * p2
-
-    # Smooth Step
-    for tick in range(ticks + 1):
-        loop_start = time.perf_counter()
-        t = tick / ticks
-        tick_pos_temp = []
-
-        for leg_id in range(6):
-            if gv.leg_state[leg_id] == "drag":
-                # lineares Gleiten in die Center-Position
-                tick_pos_temp.append(current_pos_copy[leg_id] + (gv.center_points[leg_id] - current_pos_copy[leg_id]) * t)
-
-            elif gv.leg_state[leg_id] == "step":
-                # Mittlerer Kontrollpunkt für Bezier-Kurve
-                mid_point = [
-                    (current_pos_copy[leg_id][0] + target_pos[leg_id][0]) / 2,
-                    (current_pos_copy[leg_id][1] + target_pos[leg_id][1]) / 2,
-                    max(current_pos_copy[leg_id][2], target_pos[leg_id][2])
-                    + cfg.step_height,
-                ]
-                x = interpolate(
-                    t, current_pos_copy[leg_id][0], mid_point[0], target_pos[leg_id][0]
-                )
-                y = interpolate(
-                    t, current_pos_copy[leg_id][1], mid_point[1], target_pos[leg_id][1]
-                )
-                z = interpolate(
-                    t, current_pos_copy[leg_id][2], mid_point[2], target_pos[leg_id][2]
-                )
-                tick_pos_temp.append([x, y, z])
-
-        tick_pos = dt.PosArray(tick_pos_temp)
-        # IK mit letzter Winkelstellung als Startpunkt
-        emitCalculation(kin.ikpyInverse(tick_pos))
-        print("Tick Position: \n", tick_pos)
-        if (cfg.sim == True):
-            gv.shared_sim.step()
-
-        # Timing anpassen, damit Loop gleichmäßig bleibt
-        elapsed = time.perf_counter() - loop_start
-        sleep_time = tick_duration - elapsed
-        if sleep_time > 0:
-            time.sleep(sleep_time)
-
-    # Endposition sichern
-    #emitCalculation(kin.ikpyInverse(target_pos))
-    #print("END TargetPos:\n", target_pos, "\n\n\n")
-
-    # Update current pos
-    gv.current_pos = tick_pos
-    gv.current_rad = kin.ikpyInverse(tick_pos)
-    #emitCalculation(kin.ikpyInverse(tick_pos))
-
-    # Gait-Zustand wechseln
-    if gv.leg_state[0] == "step":
-        gv.leg_state = np.array(["drag", "step", "drag", "step", "drag", "step"])
-    else:
-        gv.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
-
-def walking_four(duration=cfg.standard_duration, tickpersec=cfg.standard_tickpersec, curve_height=cfg.step_height):
-    # init Values
-    ticks = duration * tickpersec
-    tick_duration = 1 / tickpersec
-    tick_pos: dt.PosArray  # aktuelle XYZ-Positionen
-    current_rad_copy: dt.RadArray = gv.current_rad
-    current_pos_copy: dt.PosArray = gv.current_pos
-    dirmov_copy = gv.vector_dirmov
-
-    # Zielpunkte für jede Beinspitze berechnen
-    target_pos_temp = []
-    for i in range(6):
-        target_pos_temp.append(
-            [
-                gv.center_points[i][0] + dirmov_copy[0] * cfg.step_length,
-                gv.center_points[i][1] + dirmov_copy[1] * cfg.step_length,
-                gv.center_points[i][2]
-            ]
-        )
-    target_pos: dt.PosArray = target_pos_temp
-
-    def interpolate(t, p0, p1, p2):
-        return (1 - t) ** 2 * p0 + 2 * (1 - t) * t * p1 + t**2 * p2
-
-    # Smooth Step
-    for tick in range(int(ticks) + 1):
-        loop_start = time.perf_counter()
-        t = tick / ticks
-        tick_pos_temp = []
-
-        for leg_id in range(6):
-            if gv.leg_state[leg_id] == "drag":
-                # lineares Gleiten in die Center-Position
-                tick_pos_temp.append(current_pos_copy[leg_id] + (gv.center_points[leg_id] - current_pos_copy[leg_id]) * t)
-
-            elif gv.leg_state[leg_id] == "step":
-                # Mittlerer Kontrollpunkt für Bezier-Kurve
-                mid_point = [
-                    (current_pos_copy[leg_id][0] + target_pos[leg_id][0]) / 2,
-                    (current_pos_copy[leg_id][1] + target_pos[leg_id][1]) / 2,
-                    max(current_pos_copy[leg_id][2], target_pos[leg_id][2])
-                    + cfg.step_height,
-                ]
-                x = interpolate(
-                    t, current_pos_copy[leg_id][0], mid_point[0], target_pos[leg_id][0]
-                )
-                y = interpolate(
-                    t, current_pos_copy[leg_id][1], mid_point[1], target_pos[leg_id][1]
-                )
-                z = interpolate(
-                    t, current_pos_copy[leg_id][2], mid_point[2], target_pos[leg_id][2]
-                )
-                tick_pos_temp.append([x, y, z])
-
-        tick_pos = dt.PosArray(tick_pos_temp)
-        # IK mit letzter Winkelstellung als Startpunkt
-        emitCalculation(kin.ikpyInverse(tick_pos))
-        print("Tick Position: \n", tick_pos)
-
-        # Timing anpassen, damit Loop gleichmäßig bleibt
-        elapsed = time.perf_counter() - loop_start
-        sleep_time = tick_duration - elapsed
-        if sleep_time > 0:
-            time.sleep(sleep_time)
-
-    # Endposition sichern
-    #emitCalculation(kin.ikpyInverse(target_pos))
-    print("END TargetPos:\n", target_pos, "\n\n\n")
-
-    # Gait-Zustand wechseln
-    if gv.leg_state[0] == "step":
-        gv.leg_state = np.array(["drag", "drag", "step", "drag", "drag", "drag"])
-    elif gv.leg_state[2] == "step":
-        gv.leg_state = np.array(["drag", "drag", "drag", "step", "drag", "drag"])
-    elif gv.leg_state[3] == "step":
-        gv.leg_state = np.array(["drag", "drag", "drag", "drag", "drag", "step"])
-    elif gv.leg_state[5] == "step":
-        gv.leg_state = np.array(["step", "drag", "drag", "drag", "drag", "drag"])
-
-def wave_emote(cycles=3, duration=1.5, tickpersec=20,
-                   height=40, amplitude=25, inward_offset=10):
-    """
-    Greeting wave using front leg (leg 0)
-
-    Motion:
-    - Z: lifts leg up
-    - Y: waves left/right
-    - X: slightly pulled inward to avoid IK limits
-    """
-
-    ticks = int(duration * tickpersec)
-    tick_duration = 1 / tickpersec
-
-    # Safe copy
-    base_pos = dt.PosArray(np.copy(gv.current_pos.data))
-
-    leg_id = 0  # front leg
-
-    for cycle in range(cycles):
-        for tick in range(ticks):
-            loop_start = time.perf_counter()
-
-            t = tick / ticks
-            tick_pos = np.copy(base_pos.data)
-
-            cx, cy, cz = base_pos[leg_id]
-
-            # smooth outward-only wave
-            y_wave = math.sin(4 * math.pi * t)
-            y_offset = amplitude * (0.5 * (y_wave + 1))
-
-            # vertical lift
-            z_offset = height * math.sin(math.pi * t)
-
-            # clamp sideways motion
-            max_y_dev = 30
-            new_y = cy + y_offset
-            new_y = max(cy - max_y_dev, min(cy + max_y_dev, new_y))
-
-            tick_pos[leg_id] = [
-                cx + 10,     # small forward bias (IMPORTANT)
-                new_y,
-                cz + z_offset
-            ]
-
-            tick_pos = dt.PosArray(tick_pos)
-
-            emitCalculation(kin.ikpyInverse(tick_pos))
-
-            # Timing
-            elapsed = time.perf_counter() - loop_start
-            sleep_time = tick_duration - elapsed
-            if sleep_time > 0:
-                time.sleep(sleep_time)
-
-    # Return to base pose
-    emitCalculation(kin.ikpyInverse(base_pos))
-    gv.current_pos = base_pos
-
-def laola_wave_emote(cycles=3, duration=1.5, tickpersec=cfg.standard_tickpersec, height=10, amplitude=30):
-    ticks = int(duration * tickpersec)
-    tick_duration = 1 / tickpersec
-
-    # Safe copy of base position
-    base_pos = dt.PosArray(np.copy(gv.current_pos.data))
-
-    # Legs that perform the wave
-    wave_legs = [1, 4]
-
-    for cycle in range(cycles):
-        for tick in range(ticks):
-            loop_start = time.perf_counter()
-            t = tick / ticks  # normalized 0 → 1
-            tick_pos = np.copy(base_pos.data)
-
-            for leg_id in wave_legs:
-                cx, cy, cz = base_pos[leg_id]
-
-                # Phase shift for Laola wave
-                phase = 0 if leg_id == 1 else math.pi
-
-                # Sideways motion (Y) — reduced amplitude to avoid overextension
-                y_offset = amplitude * math.sin(2 * math.pi * t + phase)
-
-                # Vertical motion (Z) — full up/down oscillation
-                z_offset = height * math.sin(2 * math.pi * t + phase)
-
-                # Apply movement in YZ plane, X stays fixed
-                tick_pos[leg_id] = [
-                    cx,
-                    cy + y_offset,
-                    cz + z_offset
-                ]
-
-            tick_pos = dt.PosArray(tick_pos)
-
-            # IK + send command
-            emitCalculation(kin.ikpyInverse(tick_pos))
-
-            # Timing control
-            elapsed = time.perf_counter() - loop_start
-            sleep_time = tick_duration - elapsed
-            if sleep_time > 0:
-                time.sleep(sleep_time)
-
-    # Return to original stance
-    emitCalculation(kin.ikpyInverse(base_pos))
-    gv.current_pos = base_pos

RobotState/RobotState.py

@@ -1,20 +0,0 @@
-import numpy as np
-
-class RobotState:
-    def on_enter(self, ctx):
-        pass
-
-    def on_exit(self, ctx):
-        pass
-
-    def update(self, ctx, intent, dt):
-        pass
-
-class RobotContext:
-    def __init__(self):
-        self.current_rad = None
-        self.current_pos = None
-        self.leg_state = None
-
-        self.robotCommunication = None
-        self.shared_sim = None

RobotState/idle.py

@@ -1,7 +0,0 @@
-from RobotState.RobotState import RobotState
-
-class IdleState(RobotState):
-    def update(self, ctx, intent, dt):
-        if intent.walk:
-            return "walking"
-        return None

RobotState/walking.py

@@ -1,70 +0,0 @@
-# RobotState/walking.py
-import time
-import numpy as np
-import config as cfg
-import kinematics as kin
-import DataTypes as dt
-from RobotState.RobotState import RobotState
-
-class WalkingState(RobotState):
-
-    def on_enter(self, ctx):
-        self.t = 0.0
-        self.tick_duration = 1 / cfg.standard_tickpersec
-        self.ticks = cfg.standard_duration * cfg.standard_tickpersec
-
-        self.current_pos_copy = ctx.current_pos.copy()
-        self.dirmov = [0.0, 0.0]
-
-    def update(self, ctx, intent, dt):
-        # 🛑 Transition check
-        if not intent.walk:
-            return "idle"
-
-        self.dirmov = [intent.move_x, intent.move_y]
-
-        t = self.t / self.ticks
-        tick_pos_temp = []
-
-        for leg_id in range(6):
-            if ctx.leg_state[leg_id] == "drag":
-                tick_pos_temp.append(
-                    self.current_pos_copy[leg_id]
-                    + (ctx.center_points[leg_id] - self.current_pos_copy[leg_id]) * t
-                )
-
-            elif ctx.leg_state[leg_id] == "step":
-                mid = [
-                    (self.current_pos_copy[leg_id][0] + ctx.center_points[leg_id][0]) / 2,
-                    (self.current_pos_copy[leg_id][1] + ctx.center_points[leg_id][1]) / 2,
-                    ctx.center_points[leg_id][2] + cfg.step_height,
-                ]
-
-                def bez(a, b, c):
-                    return (1 - t)**2 * a + 2*(1 - t)*t*b + t*t*c
-
-                x = bez(self.current_pos_copy[leg_id][0], mid[0], ctx.center_points[leg_id][0])
-                y = bez(self.current_pos_copy[leg_id][1], mid[1], ctx.center_points[leg_id][1])
-                z = bez(self.current_pos_copy[leg_id][2], mid[2], ctx.center_points[leg_id][2])
-
-                tick_pos_temp.append([x, y, z])
-
-        tick_pos = dt.PosArray(tick_pos_temp)
-        target_rad = kin.ikpyInverse(tick_pos)
-
-        ctx.current_pos = tick_pos
-        ctx.current_rad = target_rad
-
-        if ctx.robotCommunication:
-            ctx.robotCommunication.send_command(target_rad)
-
-        self.t += 1
-        if self.t >= self.ticks:
-            self.t = 0
-            ctx.leg_state = (
-                np.array(["drag","step","drag","step","drag","step"])
-                if ctx.leg_state[0] == "step"
-                else np.array(["step","drag","step","drag","step","drag"])
-            )
-
-        return None

config.py

@@ -1,31 +1,55 @@
-# Connection
-port = "COM4"
-baudrate = 38400
+"""
+config.py - Central Configuration & Parameters for JackBot
+"""
+from dataclasses import dataclass, field
+from enum import Enum
 
-esp32_ip = "192.168.188.32"
-esp32_port = 3323
 
-timeout = 2.0
+class BackendType(Enum):
+    SIMULATION = "sim"
+    ESP32 = "esp32"
+    ARDUINO = "arduino"
 
-sim: bool = True
-arduinoConnection: bool = False
-urdf_path = "JackBotUrdf.urdf"
 
-# row = legnumber (left 0,1,2 right 3,4,5)
-# column = servo position from torso(0) to feet(2)
+@dataclass
+class RobotConfig:
+    # -------------------------------------------------------------------------
+    # 1. Execution & Backend Target
+    # -------------------------------------------------------------------------
+    backend: BackendType = BackendType.SIMULATION
+    urdf_path: str = "JackBotUrdf.urdf"
 
-# Dimensions in m
-#robot_height = -0.1
-step_height = 0.03
-step_length = 0.04
-translation_gain = 1.0
-rotation_gain = 2.0
-# Leglength in m
-# L1 = 0.068
-# L2 = 0.0602
-# L3 = 0.070
+    # Hardware Connection Settings
+    port: str = "COM4"                  # Serial Port for Arduino
+    baudrate: int = 38400               # Baudrate for Serial
+    esp32_ip: str = "192.168.188.32"    # Wi-Fi IP for ESP32 UDP communication
+    esp32_port: int = 3323              # UDP Target Port
+    comm_timeout: float = 2.0           # Connection timeout threshold [s]
 
-# standard_tickpersec = # 20 Fluessige Bewegung
-standard_tickpersec:float = 25 # [ticks/s]
-standard_duration:float = 0.8  # [s]
-standard_tickduration:float = 1 / standard_tickpersec
+    # -------------------------------------------------------------------------
+    # 2. Gait & Motion Kinematics Parameters
+    # -------------------------------------------------------------------------
+    step_height: float = 0.05          # Clearance height of foot swing [m]
+    step_length: float = 0.08          # Maximum stride distance [m]
+    translation_gain: float = 1.0      # Speed multiplier for directional translation
+    rotation_gain: float = 4.0         # Speed multiplier for yaw rotation
+
+    # Timing & Execution Frequency
+    tick_rate_hz: float = 25.0         # Motion loop execution rate [ticks/sec]
+    step_duration: float = 0.8         # Time to complete full gait stride [s]
+
+    @property
+    def standard_tickpersec(self) -> float:
+        return self.tick_rate_hz
+
+    @standard_tickpersec.setter
+    def standard_tickpersec(self, value: float) -> None:
+        self.tick_rate_hz = value
+
+    @property
+    def tick_duration(self) -> float:
+        return 1.0 / self.tick_rate_hz
+
+
+# Global Active Instance
+cfg = RobotConfig()

gui/MainWindow.py

@@ -0,0 +1,184 @@
+"""
+MainWindow.py - Master Dashboard with Dynamic Input Selection
+"""
+import random
+import time
+import dearpygui.dearpygui as dpg
+from config import cfg, BackendType
+
+# Input State Tracking
+input_config = {
+    "source": "Gamepad",  # Options: "Gamepad", "GUI Sliders", "Random Walk"
+    "manual_vx": 0.0,
+    "manual_vy": 0.0,
+    "manual_omega": 0.0,
+    "random_interval": 2.0,  # Change random vector every N seconds
+    "last_random_time": 0.0,
+    "random_vx": 0.0,
+    "random_vy": 0.0,
+    "random_omega": 0.0,
+}
+
+
+def create_parameter_gui(on_start_callback=None, on_stop_callback=None):
+    dpg.create_context()
+    dpg.create_viewport(title="JackBot Master Control Deck", width=460, height=800)
+
+    is_running = {"value": False}
+
+    def toggle_start(sender, app_data, user_data):
+        is_running["value"] = not is_running["value"]
+        if is_running["value"]:
+            dpg.configure_item("start_btn", label="STOP ROBOT")
+            dpg.configure_item("status_text", default_value="Status: RUNNING", color=(0, 255, 0))
+            dpg.configure_item("backend_combo", enabled=False)
+            if on_start_callback:
+                on_start_callback()
+        else:
+            dpg.configure_item("start_btn", label="START ROBOT")
+            dpg.configure_item("status_text", default_value="Status: STOPPED", color=(255, 100, 100))
+            dpg.configure_item("backend_combo", enabled=True)
+            if on_stop_callback:
+                on_stop_callback()
+
+    with dpg.window(label="Robot Control & Telemetry Deck", width=440, height=770, pos=(5, 5)):
+        
+        # --- Execution Control ---
+        dpg.add_text("Control State", color=(0, 200, 255))
+        dpg.add_text("Status: STOPPED", tag="status_text", color=(255, 100, 100))
+        dpg.add_button(
+            label="START ROBOT", 
+            tag="start_btn", 
+            width=-1, 
+            height=40, 
+            callback=toggle_start
+        )
+        dpg.add_separator()
+
+        # --- Input Source Switcher ---
+        dpg.add_text("Input Controller Source", color=(0, 200, 255))
+        dpg.add_combo(
+            tag="input_source_combo",
+            label="Input Source",
+            items=["Gamepad", "GUI Sliders", "Random Walk"],
+            default_value=input_config["source"],
+            callback=lambda s, d: input_config.update({"source": d})
+        )
+
+        # Manual Sliders (Active when GUI Sliders selected)
+        dpg.add_slider_float(
+            label="Vx (Fwd/Bwd)",
+            default_value=0.0, min_value=-1.0, max_value=1.0,
+            callback=lambda s, d: input_config.update({"manual_vx": d})
+        )
+        dpg.add_slider_float(
+            label="Vy (Strafe)",
+            default_value=0.0, min_value=-1.0, max_value=1.0,
+            callback=lambda s, d: input_config.update({"manual_vy": d})
+        )
+        dpg.add_slider_float(
+            label="Omega (Turn)",
+            default_value=0.0, min_value=-1.0, max_value=1.0,
+            callback=lambda s, d: input_config.update({"manual_omega": d})
+        )
+        dpg.add_separator()
+
+        # --- Active Command Readouts ---
+        dpg.add_text("Active Command Output", color=(0, 200, 255))
+        dpg.add_text("Joysticks Connected: 0", tag="txt_joy_count")
+        dpg.add_text("Motion State: IDLE", tag="txt_motion_state", color=(200, 200, 200))
+        dpg.add_text("Active Vx   :  0.00", tag="txt_vx")
+        dpg.add_text("Active Vy   :  0.00", tag="txt_vy")
+        dpg.add_text("Active Omega:  0.00", tag="txt_omega")
+        dpg.add_separator()
+
+        # --- Backend Selection ---
+        dpg.add_text("Backend Selection", color=(0, 200, 255))
+        dpg.add_combo(
+            tag="backend_combo",
+            label="Target Backend",
+            items=[e.value for e in BackendType],
+            default_value=cfg.backend.value,
+            callback=lambda sender, data: setattr(cfg, 'backend', BackendType(data))
+        )
+        dpg.add_separator()
+
+        # --- Kinematics Adjustments ---
+        dpg.add_text("Gait & Kinematics", color=(0, 200, 255))
+        dpg.add_slider_float(
+            label="Step Height (m)",
+            default_value=cfg.step_height,
+            min_value=0.01, max_value=0.10, format="%.3f m",
+            callback=lambda sender, data: setattr(cfg, 'step_height', data)
+        )
+        dpg.add_slider_float(
+            label="Step Length (m)",
+            default_value=cfg.step_length,
+            min_value=0.01, max_value=0.10, format="%.3f m",
+            callback=lambda sender, data: setattr(cfg, 'step_length', data)
+        )
+        dpg.add_separator()
+
+        # --- Hardware Comms ---
+        dpg.add_text("Hardware Communications", color=(0, 200, 255))
+        dpg.add_input_text(label="ESP32 IP", default_value=cfg.esp32_ip, callback=lambda s, d: setattr(cfg, 'esp32_ip', d))
+        dpg.add_input_text(label="Serial Port", default_value=cfg.port, callback=lambda s, d: setattr(cfg, 'port', d))
+
+    dpg.setup_dearpygui()
+    dpg.show_viewport()
+
+
+def resolve_active_command(gamepad_cmd: dict | None) -> dict:
+    """Selects motion vector based on current active source in GUI."""
+    source = input_config["source"]
+    now = time.time()
+
+    if source == "GUI Sliders":
+        vx = input_config["manual_vx"]
+        vy = input_config["manual_vy"]
+        omega = input_config["manual_omega"]
+        joy_count = gamepad_cmd.get("joysticks_count", 0) if gamepad_cmd else 0
+
+    elif source == "Random Walk":
+        # Generate new random direction vector periodically
+        if now - input_config["last_random_time"] > input_config["random_interval"]:
+            input_config["random_vx"] = round(random.uniform(-1.0, 1.0), 2)
+            input_config["random_vy"] = round(random.uniform(-1.0, 1.0), 2)
+            input_config["random_omega"] = round(random.uniform(-1.0, 1.0), 2)
+            input_config["last_random_time"] = now
+
+        vx = input_config["random_vx"]
+        vy = input_config["random_vy"]
+        omega = input_config["random_omega"]
+        joy_count = gamepad_cmd.get("joysticks_count", 0) if gamepad_cmd else 0
+
+    else:  # "Gamepad"
+        if gamepad_cmd:
+            vx, vy, omega = gamepad_cmd.get("dirmov", [0.0, 0.0, 0.0])
+            joy_count = gamepad_cmd.get("joysticks_count", 0)
+        else:
+            vx, vy, omega, joy_count = 0.0, 0.0, 0.0, 0
+
+    state = "walking" if (vx != 0.0 or vy != 0.0 or omega != 0.0) else "idle"
+
+    # Update GUI Labels
+    if dpg.is_dearpygui_running():
+        dpg.set_value("txt_joy_count", f"Joysticks Connected: {joy_count}")
+        dpg.set_value("txt_motion_state", f"Motion State: {state.upper()} ({source})")
+        dpg.set_value("txt_vx", f"Active Vx   : {vx:>6.2f}")
+        dpg.set_value("txt_vy", f"Active Vy   : {vy:>6.2f}")
+        dpg.set_value("txt_omega", f"Active Omega: {omega:>6.2f}")
+
+    return {
+        "state": state,
+        "dirmov": [vx, vy, omega]
+    }
+
+
+def render_gui_frame():
+    if dpg.is_dearpygui_running():
+        dpg.render_dearpygui_frame()
+
+
+def stop_gui():
+    dpg.destroy_context()

inputs/InputProvider.py

@@ -0,0 +1,27 @@
+"""
+inputs/InputProvider.py - Abstraction layer for robot control inputs
+"""
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import List
+
+
+@dataclass
+class CommandFrame:
+    """Represents a snapshot of control inputs at a single tick."""
+    state: str = "idle"
+    vector_dirmov: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])  # [vx, vy, omega]
+    emote: str = ""
+
+
+class InputProvider(ABC):
+    """Abstract Base Class for input providers (Joystick, Random/RL, Scripted, etc.)."""
+
+    @abstractmethod
+    def get_command(self) -> CommandFrame:
+        """Polls or generates the latest input command frame."""
+        pass
+
+    def stop(self) -> None:
+        """Optional cleanup when shutting down input loop."""
+        pass

inputs/PygameController.py

@@ -0,0 +1,101 @@
+"""
+inputs/PygameController.py - Headless Pygame Joystick Provider
+"""
+import math
+from multiprocessing import Queue
+import pygame
+
+from inputs.InputProvider import InputProvider, CommandFrame
+
+
+class PygameController(InputProvider):
+    def __init__(self, deadzone: float = 0.2):
+        pygame.init()
+        pygame.joystick.init()
+        
+        self.deadzone = deadzone
+        self.joysticks = {}
+        self.clock = pygame.time.Clock()
+        self.current_command = CommandFrame(state="idle", vector_dirmov=[0.0, 0.0, 0.0])
+
+    def _normalize_input(self, x: float, y: float) -> tuple[float, float]:
+        magnitude = math.sqrt(x**2 + y**2)
+        if magnitude > 1.0:
+            x /= magnitude
+            y /= magnitude
+        return x, y
+
+    def update(self) -> CommandFrame:
+        vx, vy, omega = 0.0, 0.0, 0.0
+        emote = ""
+
+        # Pump Pygame events in headless mode
+        for event in pygame.event.get():
+            if event.type == pygame.JOYDEVICEADDED:
+                joy = pygame.joystick.Joystick(event.device_index)
+                self.joysticks[joy.get_instance_id()] = joy
+                print(f"[Controller] Joystick {joy.get_instance_id()} connected: {joy.get_name()}")
+            elif event.type == pygame.JOYDEVICEREMOVED:
+                if event.instance_id in self.joysticks:
+                    print(f"[Controller] Joystick {event.instance_id} disconnected")
+                    del self.joysticks[event.instance_id]
+            elif event.type == pygame.JOYBUTTONDOWN:
+                if event.button == 3:
+                    emote = "wave"
+
+        # Read stick values if joystick is connected
+        if self.joysticks:
+            for joystick in self.joysticks.values():
+                raw_vx = -joystick.get_axis(1)  # Forward (+) / Backward (-)
+                raw_vy = joystick.get_axis(0)   # Right (+) / Left (-)
+                
+                num_axes = joystick.get_numaxes()
+                raw_omega = joystick.get_axis(3) if num_axes > 3 else joystick.get_axis(2)
+
+                norm_vx, norm_vy = self._normalize_input(raw_vx, raw_vy)
+
+                if abs(raw_omega) >= self.deadzone:
+                    omega = raw_omega
+                if abs(norm_vx) >= self.deadzone or abs(norm_vy) >= self.deadzone:
+                    vx, vy = norm_vx, norm_vy
+
+                # Cardinal snapping
+                if 0.8 < vx and -0.2 < vy < 0.2:
+                    vx, vy = 1.0, 0.0
+                elif vx < -0.8 and -0.2 < vy < 0.2:
+                    vx, vy = -1.0, 0.0
+                elif -0.2 < vx < 0.2 and 0.8 < vy:
+                    vx, vy = 0.0, 1.0
+                elif -0.2 < vx < 0.2 and vy < -0.8:
+                    vx, vy = 0.0, -1.0
+
+        state = "walking" if (vx != 0.0 or vy != 0.0 or omega != 0.0) else "idle"
+
+        self.current_command = CommandFrame(
+            state=state,
+            vector_dirmov=[vx, vy, omega],
+            emote=emote
+        )
+        return self.current_command
+
+    def get_command(self) -> CommandFrame:
+        return self.update()
+
+    def stop(self) -> None:
+        pygame.quit()
+
+
+def controller_loop(control_queue: Queue):
+    provider = PygameController()
+    try:
+        while True:
+            cmd = provider.get_command()
+            control_queue.put({
+                "state": cmd.state,
+                "dirmov": cmd.vector_dirmov,
+                "emote": cmd.emote,
+                "joysticks_count": len(provider.joysticks)
+            })
+            provider.clock.tick(30)
+    finally:
+        provider.stop()

inputs/RandomeInputProvider.py

@@ -0,0 +1,35 @@
+"""
+inputs/RandomInputProvider.py - Random Command Generator for Simulation / Training
+"""
+import random
+import time
+from inputs.InputProvider import InputProvider, CommandFrame
+
+
+class RandomInputProvider(InputProvider):
+    def __init__(self, change_interval: float = 2.0):
+        self.change_interval = change_interval
+        self.last_change = time.time()
+        self.current_command = CommandFrame()
+
+    def get_command(self) -> CommandFrame:
+        now = time.time()
+        if now - self.last_change > self.change_interval:
+            self.last_change = now
+            
+            # 20% chance to idle, 80% to walk randomly
+            if random.random() < 0.2:
+                self.current_command = CommandFrame(
+                    state="idle",
+                    vector_dirmov=[0.0, 0.0, 0.0]
+                )
+            else:
+                vx = random.uniform(-1.0, 1.0)
+                vy = random.uniform(-1.0, 1.0)
+                omega = random.uniform(-1.0, 1.0)
+                self.current_command = CommandFrame(
+                    state="walking",
+                    vector_dirmov=[vx, vy, omega]
+                )
+
+        return self.current_command

kinematics.py

@@ -1,11 +1,14 @@
 from ikpy.chain import Chain
 from ikpy.link import OriginLink, URDFLink
+from typing import Optional
 import numpy as np
 import math
 import time
 
-import GlobalVariables as gv
-import config as cfg
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning, module="ikpy")
+
+from config import cfg
 import DataTypes as dt
 
 leg_chains = {
@@ -110,33 +113,38 @@ def ikpyForward(target_rad: dt.RadArray) -> dt.PosArray:
         target_pos_temp.append([x, y, z])
     return dt.RadArray(np.array(target_pos_temp))
 
+def ikpyInverse(target_pos: dt.PosArray, initial_rad: Optional[dt.RadArray] = None) -> dt.RadArray:
+    calculated_rads = []
 
-def ikpyInverse(target_pos: dt.PosArray) -> dt.RadArray:
-    target_rad_temp = []
-    print("target_pos:\n",target_pos)
-    for leg_id, chain in leg_chains.items():
-        # prepare initial guess
-        guess = np.array([0] + list(gv.current_rad[leg_id]) + [0], dtype=float)
+    for leg_id in range(6):
+        chain = leg_chains[leg_id]  # Your IKPy Chain object
+        target_xyz = target_pos[leg_id]
 
-        try:
-            ik_result = chain.inverse_kinematics(
-                target_pos[leg_id],
-                initial_position=guess,
-                max_iter=100
+        full_initial_position = None
+        if initial_rad is not None:
+            # Create a zero array matching the total number of links in the chain
+            full_initial_position = [0.0] * len(chain.links)
+            
+            # Map the 3 active joint angles into active link indices
+            active_indices = [i for i, active in enumerate(chain.active_links_mask) if active]
+            
+            # Match 3 active joints to the 3 active link positions in the chain
+            for idx, angle in zip(active_indices, initial_rad[leg_id]):
+                full_initial_position[idx] = angle
+
+        if full_initial_position is not None:
+            angles = chain.inverse_kinematics(
+                target_position=target_xyz, 
+                initial_position=full_initial_position
             )
-        except ValueError:
-            # fallback → keep the clipped guess
-            ik_result = guess
+        else:
+            angles = chain.inverse_kinematics(target_position=target_xyz)
 
-        # keep only the 3 actuated joint angles (indices 1,2,3)
-        target_rad_temp.append(ik_result[1:4])
+        # Extract only the active joint angles (3 revolute joints) from IKPy result
+        active_angles = chain.active_from_full(angles)
+        calculated_rads.append(active_angles)
 
-    return dt.RadArray(np.array(target_rad_temp))
-    # for leg_id, chain in leg_chains.items():
-    #     ik_result = chain.inverse_kinematics(target_pos[leg_id])
-    #     # Remove first element (IKPy adds a "dummy" fixed base joint)
-    #     target_rad_temp.append(ik_result[1:4])
-    return dt.RadArray(np.array(target_rad_temp))
+    return dt.RadArray(calculated_rads)
 
 
 def ikpytest():
@@ -161,132 +169,4 @@ def ikpytest():
         ]
     )
     current_rad: dt.RadArray = joint_angles_deg.to_rad()
-    ikpyInverse(targets, current_rad)
-
-
-# walk old
-"""
-def walk(duration=standard_duration, ticks=standard_tickrate, curve_height=gv.step_height):
-    tick_duration = duration / ticks
-    tick_positions = np.copy(gv.current_pos)
-    dirmov_copy = gv.vector_dirmov 
-    target_pos =[]
-    for i in range(6): # Zielposition berechnen 
-        target_pos.append([gv.center_points[i][0] + dirmov_copy[0] * gv.step_length, gv.center_points[i][1] + dirmov_copy[1] * gv.step_length, gv.robot_height ])
-    def interpolate(t, p0, p1, p2):
-        return (1 - t)**2 * p0 + 2 * (1 - t) * t * p1 + t**2 * p2
-
-    for tick in range(int(ticks) + 1):
-        loop_start = time.perf_counter()
-        t = tick / ticks
-
-        for leg_id in range(6):
-            if gv.leg_state[leg_id] == "drag":
-                # Linear interpolation
-                tick_positions[leg_id] = gv.current_pos[leg_id] + (gv.center_points[leg_id] - gv.current_pos[leg_id]) * t
-
-            elif gv.leg_state[leg_id] == "step":
-                # Curve movement (Bezier path)
-                help_pos = [
-                    target_pos[leg_id][0] - gv.current_pos[leg_id][0],
-                    target_pos[leg_id][1] - gv.current_pos[leg_id][1],
-                    target_pos[leg_id][2] + curve_height
-                ]
-                x = interpolate(t, gv.current_pos[leg_id][0], help_pos[0], target_pos[leg_id][0])
-                y = interpolate(t, gv.current_pos[leg_id][1], help_pos[1], target_pos[leg_id][1])
-                z = interpolate(t, gv.current_pos[leg_id][2], help_pos[2], target_pos[leg_id][2])
-                tick_positions[leg_id] = [x, y, z]
-
-        # Send all legs at once through IK
-        ikpyInverse(tick_positions, gv.current_pos)
-
-        elapsed = time.perf_counter() - loop_start
-        sleep_time = tick_duration - elapsed
-        if sleep_time > 0:
-            time.sleep(sleep_time)
-
-    # Final position correction
-    ikpyInverse(target_pos, gv.current_pos)
-
-    if (gv.leg_state[0] == "step"):
-        gv.leg_state = np.array(["drag", "step", "drag", "step", "drag", "step"])
-    if (gv.leg_state[0] == "drag"):
-        gv.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
-    time.sleep(0.05)
-
-
-
-# Berechnung der Bein Bewegung - Grade
-def drag_leg(
-    leg_num,
-    current_pos,
-    target_pos,
-    duration=cfg.standard_duration,
-    ticks=cfg.standard_tickrate,
-):
-    tick_duration = duration / ticks
-    dragtickvec = (target_pos - current_pos[leg_num]) / ticks
-
-    tick_positions = np.copy(current_pos)
-
-    for tick in range(int(ticks)):
-        start_time = time.perf_counter()
-
-        tick_positions[leg_num] += dragtickvec
-        ikpyInverse(tick_positions, current_pos)
-
-        elapsed = time.perf_counter() - start_time
-        sleep_time = tick_duration - elapsed
-        if sleep_time > 0:
-            time.sleep(sleep_time)
-
-    tick_positions[leg_num] = target_pos
-    ikpyInverse(tick_positions, current_pos)
-    time.sleep(0.05)
-
-
-# Berechnung der Bein Bewegung - Kurve
-def curve_leg(
-    leg_num,
-    current_pos,
-    target_pos,
-    curve_height=cfg.step_height,
-    duration=cfg.standard_duration,
-    ticks=cfg.standard_tickrate,
-):
-    tick_duration = duration / ticks
-
-    help_pos = [
-        target_pos[0] - current_pos[leg_num][0],
-        target_pos[1] - current_pos[leg_num][1],
-        target_pos[2] + curve_height,
-    ]
-
-    def interpolate(t, p0, p1, p2):
-        return (1 - t) ** 2 * p0 + 2 * (1 - t) * t * p1 + t**2 * p2
-
-    tick_positions = np.copy(current_pos)
-
-    for tick in range(int(ticks) + 1):
-        start_time = time.perf_counter()
-
-        t = tick / ticks
-        x = interpolate(t, current_pos[leg_num][0], help_pos[0], target_pos[0])
-        y = interpolate(t, current_pos[leg_num][1], help_pos[1], target_pos[1])
-        z = interpolate(t, current_pos[leg_num][2], help_pos[2], target_pos[2])
-
-        tick_positions[leg_num] = [x, y, z]
-        ikpyInverse(tick_positions, current_pos)
-
-        elapsed = time.perf_counter() - start_time
-        sleep_time = tick_duration - elapsed
-        if sleep_time > 0:
-            time.sleep(sleep_time)
-
-    tick_positions[leg_num] = target_pos
-    ikpyInverse(tick_positions, current_pos)
-    time.sleep(0.05)
-
-"""
-if __name__ == "__main__":
-    print(ikpyForward(gv.test_deg.to_rad()))
+    ikpyInverse(targets, current_rad)

main.py

@@ -1,60 +1,85 @@
-from threading import Thread
-from multiprocessing import Process, Manager, Queue
+"""
+main.py - Entry point handling dynamic input source selection
+"""
+from multiprocessing import Queue, Process
 import time
 
-# Selfmade Libraries
-import Controller as ctr
-import RobotState as rs
+from config import cfg
+from Robot import Robot
+from gui.MainWindow import (
+    create_parameter_gui, 
+    render_gui_frame, 
+    stop_gui, 
+    resolve_active_command
+)
+from inputs.PygameController import controller_loop
 
-# Global Variables
-import GlobalVariables as gv
+robot_instance: Robot | None = None
+is_robot_active = False
 
 
-def robot_control():
-    # Connection
-    if gv.robotCommunication != None:
-        gv.robotCommunication.start()
+def start_robot():
+    global robot_instance, is_robot_active
+    print(f"[Main] Starting Robot with Backend: {cfg.backend.value.upper()}")
+    robot_instance = Robot(backend_type=cfg.backend)
+    robot_instance.reset_to_init()
+    is_robot_active = True
 
-    # Start Position
-    rs.initPos()
-    time.sleep(1)
 
-    ############################## Main Loop ##############################
-    tick = 0.05
+def stop_robot():
+    global robot_instance, is_robot_active
+    print("[Main] Stopping Robot...")
+    is_robot_active = False
+    if robot_instance is not None:
+        robot_instance.cleanup()
+        robot_instance = None
+
+
+def main_event_loop(control_queue: Queue):
+    global robot_instance, is_robot_active
+
+    create_parameter_gui(
+        on_start_callback=start_robot,
+        on_stop_callback=stop_robot
+    )
+
     next_time = time.time()
 
-    while gv.control_pause == False:
-        next_time += tick
+    try:
+        while True:
+            next_time += cfg.tick_duration
 
-        if gv.emote == "wave":
-            rs.initPos()
-            time.sleep(0.3)
-            rs.wave_emote()
-            gv.emote = None
-            gv.robot_state = "idle"
-            continue
+            # Drain latest gamepad frame if available
+            gamepad_cmd = None
+            while not control_queue.empty():
+                gamepad_cmd = control_queue.get_nowait()
 
-        # Robot
-        if gv.robot_state == "idle":
-            rs.initPos()
-            time.sleep(0.2)
-        elif gv.robot_state == "walking":
-            rs.walking()
+            # Resolve motion vector based on current active dropdown mode
+            active_cmd = resolve_active_command(gamepad_cmd)
 
-        # Simulation
-        if gv.shared_sim != None:
-            gv.shared_sim.step()
+            # Pass inputs to robot and tick state machine
+            if is_robot_active and robot_instance is not None:
+                robot_instance.robot_state = active_cmd["state"]
+                robot_instance.vector_dirmov = active_cmd["dirmov"]
+                robot_instance.tick()
 
-        sleep_time = next_time - time.time()
-        if sleep_time > 0:
-            time.sleep(sleep_time)
+            render_gui_frame()
+
+            sleep_time = next_time - time.time()
+            if sleep_time > 0:
+                time.sleep(sleep_time)
+
+    finally:
+        stop_gui()
 
 
 if __name__ == "__main__":
-    controller_queue = Queue()
-    robot_queue = Queue()
+    control_queue = Queue()
+    
+    controller_process = Process(target=controller_loop, args=(control_queue,))
+    controller_process.start()
 
-    controller_thread = Thread(target=ctr.controller)
-    controller_thread.start()
-    robot_control_thread = Thread(target=robot_control)
-    robot_control_thread.start()
+    try:
+        main_event_loop(control_queue)
+    finally:
+        controller_process.terminate()

ml/MetricsOverlay.py

@@ -0,0 +1,121 @@
+"""
+ml/MetricsOverlay.py - 3D HUD overlay for live simulation telemetry.
+
+This file provides a small PyBullet HUD renderer that draws live robot metrics in the
+simulation scene. It is used during GUI runs to show the current phase, command vector,
+reward information, and basic motion statistics without leaving the 3D view.
+"""
+from typing import List, Tuple, Optional, Dict
+import numpy as np
+import pybullet as p
+
+
+class MetricsHUD:
+    """Renders real-time telemetry as black text floating in 3D, always facing the active camera."""
+
+    def __init__(self, physics_client_id: int = 0):
+        self.client_id = physics_client_id
+        self._text_id: Optional[int] = None
+
+    def _get_camera_facing_orientation(self) -> List[float]:
+        """Calculates a quaternion that points the text towards the current GUI camera."""
+        try:
+            cam_info = p.getDebugVisualizerCamera(physicsClientId=self.client_id)
+            # cam_info index 8: yaw, index 9: pitch
+            yaw = cam_info[8]
+            pitch = cam_info[9]
+            
+            pitch_rad = np.radians(pitch + 90.0)
+            yaw_rad = np.radians(yaw)
+            
+            text_orientation = p.getQuaternionFromEuler(
+                [pitch_rad, 0.0, yaw_rad], 
+                physicsClientId=self.client_id
+            )
+            return text_orientation
+        except Exception:
+            return [0.0, 0.0, 0.0, 1.0]
+
+    def update(
+        self,
+        episode: int,
+        step: int,
+        robot_rewards: List[float],
+        cmd_vel: np.ndarray,
+        fps: float = 0.0,
+        height: float = 0.0,
+        roll_pitch: Tuple[float, float] = (0.0, 0.0),
+        mode: str = "direct",
+        phase: str = "STAND_ONLY",
+        distance: float = 0.0,
+        status: str = "ALIVE",
+        reward_components: Optional[Dict[str, float]] = None,
+        ep_step: int = 0,
+    ) -> None:
+        """Updates floating text block in 3D space with expanded telemetry."""
+        sorted_rewards = sorted(robot_rewards, reverse=True)
+        top1 = f"{sorted_rewards[0]:+.2f}" if len(sorted_rewards) > 0 else "0.00"
+
+        vx = cmd_vel[0] if len(cmd_vel) > 0 else 0.0
+        vy = cmd_vel[1] if len(cmd_vel) > 1 else 0.0
+        omega = cmd_vel[2] if len(cmd_vel) > 2 else 0.0
+
+        lines = [
+            "=== JACKBOT TELEMETRY ===",
+            f"Mode:        {mode.upper()}",
+            f"Curriculum:  {phase}",
+            f"Status:      {status}",
+            f"Episode:     {episode} (Step {ep_step})",
+            f"Global Step: {step}",
+            f"FPS:         {fps:.1f}",
+            "-------------------------",
+            f"Episode Rew: {top1}",
+            f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]",
+            f"Height:      {height:.3f} m",
+            f"Roll/Pitch:  {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°",
+            f"Max Dist:    {distance:.2f} m",
+        ]
+
+        if reward_components:
+            lin_v = reward_components.get("lin_vel", 0.0)
+            stab = reward_components.get("stability", 0.0)
+            h_rew = reward_components.get("height", 0.0)
+            jit = reward_components.get("jitter_penalty", 0.0)
+            lines.append("--- Reward Components ---")
+            lines.append(f"LinVel: {lin_v:.2f} | Stab: {stab:.2f}")
+            lines.append(f"Height: {h_rew:.2f} | Jitter: {jit:+.3f}")
+
+        hud_text = "\n".join(lines)
+
+        # Position above origin in simulation world
+        text_position = [-0.8, -0.8, 1.2]
+        text_color = [0, 0, 0]  # Pure black
+
+        # Calculate dynamic orientation to align text flat against camera plane
+        text_orientation = self._get_camera_facing_orientation()
+
+        # Safely remove old text to prevent PyBullet ghosting/overlapping
+        if self._text_id is not None:
+            try:
+                p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
+            except Exception:
+                pass
+                
+        # Draw fresh text
+        self._text_id = p.addUserDebugText(
+            text=hud_text,
+            textPosition=text_position,
+            textColorRGB=text_color,
+            textSize=0.085,
+            textOrientation=text_orientation,
+            physicsClientId=self.client_id
+        )
+
+    def reset(self) -> None:
+        """Removes the active debug text item so a new episode starts with a clean overlay."""
+        if self._text_id is not None:
+            try:
+                p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
+            except Exception:
+                pass
+        self._text_id = None

ml/README.md

@@ -1,59 +0,0 @@
-JackBot ML training and evaluation
-
-Quick-start
-
-1. Create a Python virtualenv and activate it:
-
-```bash
-python -m venv .venv
-source .venv/bin/activate
-```
-
-2. Install dependencies (GPU users should install `torch` appropriate for their CUDA):
-
-```bash
-pip install -r requirements.txt
-```
-
-3. Quick training (short, for smoke test):
-
-```bash
-python ml/run_train.py --timesteps 50000 --model ml/checkpoints/ppo_joint_command
-```
-
-Use multi-robot training with `--num-robots` and `--robot-spacing`:
-
-```bash
-python ml/run_train.py --timesteps 50000 --model ml/checkpoints/ppo_joint_command --num-robots 3 --robot-spacing 0.75
-```
-
-Choose the start pose at reset with `--start-pose`:
-
-```bash
-python ml/run_train.py --timesteps 50000 --model ml/checkpoints/ppo_joint_command --start-pose init_deg
-```
-
-If you want to watch the agent train in the PyBullet window, add `--gui`:
-
-```bash
-python ml/run_train.py --timesteps 50000 --model ml/checkpoints/ppo_joint_command --gui
-```
-
-4. Evaluate the trained model in GUI mode:
-
-```bash
-python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3 --gui
-```
-
-For evaluation with multiple robots and start pose:
-
-```bash
-python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3 --gui --num-robots 2 --robot-spacing 0.6 --start-pose init_deg
-```
-
-Notes
-
-- `ml/env.py` exposes a `JackBotEnv` gym environment that uses your `JackBotUrdf.urdf`.
-- The observation is `[joint_angles..., vx, vy, vz, omega]` and the action is per-joint delta in [-1,1].
-- Start with small timesteps and in `use_gui=False` to speed up iteration. Increase timesteps and tune the reward for better walking quality.
-- Keep the trained models off hardware until they behave well in simulation.

ml/__init__.py

@@ -1,6 +1,3 @@
-from .env import JackBotEnv
-from .model import ActorCritic
-from .train import train
-from .evaluate import evaluate
+from .env import JackBotEnv, CurriculumPhase
 
-__all__ = ["JackBotEnv", "ActorCritic", "train", "evaluate"]
+__all__ = ["JackBotEnv", "CurriculumPhase"]

ml/callbacks.py

@@ -0,0 +1,102 @@
+"""
+ml/callbacks.py - Stable-Baselines3 callbacks for training diagnostics.
+
+These callbacks extend SB3 training with two responsibilities: logging reward-component
+statistics for TensorBoard/console output, and checking whether the curriculum should
+advance to a harder set of commands based on recent training performance.
+"""
+
+import numpy as np
+from stable_baselines3.common.callbacks import BaseCallback
+
+
+class RewardLoggerCallback(BaseCallback):
+    """
+    Logs individual reward component averages to TensorBoard and prints 
+    the best worker's performance breakdown to the console per iteration.
+    """
+
+    def __init__(self, verbose: int = 1):
+        super().__init__(verbose)
+        self.iteration = 0
+
+    def _on_step(self) -> bool:
+        return True
+
+    def _on_rollout_end(self) -> None:
+        self.iteration += 1
+        if self.training_env is None:
+            return
+
+        try:
+            # Safely query method across all parallel worker processes
+            all_worker_averages = self.training_env.env_method("get_reward_component_averages")
+        except Exception:
+            return
+
+        if not all_worker_averages or len(all_worker_averages) == 0:
+            return
+
+        # 1. Log mean component values across ALL workers to TensorBoard
+        component_keys = all_worker_averages[0].keys()
+        for key in component_keys:
+            mean_val = float(np.mean([w.get(key, 0.0) for w in all_worker_averages]))
+            self.logger.record(f"reward_components/{key}", mean_val)
+
+        # 2. Identify the best performing worker of this iteration
+        worker_totals = [sum(w.values()) for w in all_worker_averages]
+        best_worker_idx = int(np.argmax(worker_totals))
+        best_averages = all_worker_averages[best_worker_idx]
+        best_total = worker_totals[best_worker_idx]
+
+        # 3. Print best worker breakdown to console
+        if self.verbose > 0:
+            print(f"\n" + "=" * 65)
+            print(f" ITERATION {self.iteration} | BEST WORKER (#{best_worker_idx}) REWARD BREAKDOWN")
+            print(f" Total Avg Reward / Step: {best_total:+.4f}")
+            print("-" * 65)
+            for key, val in best_averages.items():
+                print(f"   • {key:<26}: {val:+.5f}")
+            print("=" * 65 + "\n")
+
+
+class CurriculumCallback(BaseCallback):
+    """
+    Monitors training metrics using SB3's native ep_info_buffer and 
+    dynamically advances curriculum phases across worker processes.
+    """
+
+    def __init__(self, reward_threshold: float = 100.0, verbose: int = 1):
+        super().__init__(verbose)
+        self.reward_threshold = reward_threshold
+
+    def _on_step(self) -> bool:
+        return True
+
+    def _on_rollout_end(self) -> None:
+        if self.training_env is None:
+            return
+
+        # SB3 natively records finished episode stats in self.model.ep_info_buffer
+        if hasattr(self.model, "ep_info_buffer") and len(self.model.ep_info_buffer) > 0:
+            recent_rewards = [ep_info["r"] for ep_info in self.model.ep_info_buffer]
+            mean_reward = float(np.mean(recent_rewards[-50:]))
+
+            try:
+                # Query current phase from worker 0
+                phases = self.training_env.get_attr("curriculum_phase")
+                current_phase = phases[0]
+
+                # Advance curriculum if mean reward exceeds threshold
+                if mean_reward >= self.reward_threshold:
+                    if hasattr(current_phase, "next"):
+                        next_phase = current_phase.next()
+                        if next_phase != current_phase:
+                            self.training_env.set_attr("curriculum_phase", next_phase)
+                            if self.verbose > 0:
+                                print(
+                                    f"\n[Curriculum] 🚀 Promoted workers to phase: {next_phase.name} "
+                                    f"(Mean Reward: {mean_reward:.2f})"
+                                )
+            except Exception:
+                pass  # Keep rollout loop running safely if phase check fails

ml/env.py

@@ -1,317 +1,456 @@
+"""
+ml/env.py - Gymnasium environment for JackBot RL training and evaluation.
+
+This file defines JackBotEnv, the main training/evaluation environment used by PPO.
+It wraps the PyBullet simulation and Robot interfaces into a Gymnasium-compatible
+step/reset loop, manages command sampling, curriculum progression, and reward
+calculation, and exposes metrics that the training callbacks can log.
+"""
+import time
 import math
-import os
-import random
-import numpy as np
-import pybullet as p
-import pybullet_data
+from enum import IntEnum
+from typing import Optional, Tuple, Dict, Any, List
+from collections import defaultdict
+
 import gymnasium as gym
 from gymnasium import spaces
-import config as cfg
-import robot_init as ri
+import numpy as np
+
+from config import cfg
+from simulation import Simulation
+from Robot import Robot, PyBulletBackend
+from ml.MetricsOverlay import MetricsHUD
+
+
+class CurriculumPhase(IntEnum):
+    STAND_ONLY = 0
+    FORWARD = 1
+    TURN_AND_DIRECTION = 2
+    OMNI_DIRECTION = 3
+    FULL_COMMAND = 4
 
 
 class JackBotEnv(gym.Env):
-    """Gymnasium environment for joint-command learning with the JackBot URDF."""
-
-    metadata = {"render_modes": ["human", "rgb_array"]}
+    """Gymnasium environment wrapping JackBot hexapod simulation."""
 
     def __init__(
         self,
-        urdf_path: str | None = None,
         use_gui: bool = True,
-        frame_skip: int = 4,
         random_command: bool = True,
-        max_episode_steps: int = 2000,
-        num_robots: int = 1,
-        robot_spacing: float = 0.5,
-        start_pose: str = "init_deg",
+        max_episode_steps: int = 3000,
+        urdf_path: str = cfg.urdf_path,
+        robot_mode: str = "direct",  # "direct", "residual", or "kinematics"
     ):
-        """
-        Initialize the JackBot environment.
-
-        Args:
-            urdf_path: Path to the robot's URDF file. Defaults to config value.
-            use_gui: If True, uses p.GUI (shows window), else p.DIRECT (headless).
-            frame_skip: Number of physics simulation steps per agent step. 
-                Higher values mean the agent takes actions less frequently (e.g., 4 steps at 240Hz = 60Hz control).
-            random_command: If True, assigns a random velocity command (vx, vy, omega) to each robot every reset.
-            max_episode_steps: Maximum number of steps (interactions) before the episode is truncated.
-            num_robots: Number of robots to spawn for parallel training in the same scene.
-            robot_spacing: Distance in meters between robots when spawning multiple.
-            start_pose: The initial joint configuration name (e.g., 'init_deg' or 'init90_deg').
-        """
-        self.urdf_path = urdf_path or cfg.urdf_path
+        super().__init__()
+        self.robot_mode = robot_mode
         self.use_gui = use_gui
-        self.frame_skip = frame_skip
         self.random_command = random_command
         self.max_episode_steps = max_episode_steps
-        self.num_robots = max(1, num_robots)
-        self.robot_spacing = robot_spacing
-        self.start_pose = start_pose
-        self.action_scale = math.radians(8.0)
-        self.physics_client = None
-        self.robots = []
-        self.plane = None
-        self.robot_joint_indices = []
-        self.joint_lower = None
-        self.joint_upper = None
-        self.initial_angles = None
-        self.commands = None
+        self.urdf_path = urdf_path
+        self.max_robot_speed = 0.6
+
         self.step_count = 0
-        self.last_action = None
+        self.total_steps = 0
+        self.cumulative_reward = 0.0
+        self.robot_reward = 0.0
+        self.is_failed = False
+
+        self.episode_count = 0
+        self.episode_height_sum = 0.0
+        self.episode_roll_sum = 0.0
+        self.episode_pitch_sum = 0.0
+        self._curriculum_advanced = False
         self._first_reset = True
-        self._min_steps_before_done = 150
 
-        self._connect_sim()
+        self.last_reward_components: Dict[str, float] = {}
+        self.episode_reward_components_sum: Dict[str, float] = defaultdict(float)
 
-        total_joints = self.num_robots * len(self.robot_joint_indices[0])
-        observation_dim = total_joints + self.num_robots * 4
-        action_dim = total_joints
+        self.control_freq = 60
+        self.min_cmd_hold_steps = int(2.0 * self.control_freq)
+        self.max_cmd_hold_steps = int(6.0 * self.control_freq)
+        self.next_cmd_resample_step = 0
+        self.initial_stand_steps = 120
 
-        self.observation_space = spaces.Box(
-            low=-np.inf,
-            high=np.inf,
-            shape=(observation_dim,),
-            dtype=np.float32,
-        )
-        self.action_space = spaces.Box(
-            low=-1.0,
-            high=1.0,
-            shape=(action_dim,),
-            dtype=np.float32,
-        )
+        # --- INITIALIZE PYBULLET SIMULATION ENGINE ---
+        self.sim = Simulation(urdf_path=self.urdf_path, use_gui=self.use_gui)
+        self.plane_id, self.pb_robot, self.revolute_joints = self.sim.load_scene()
 
-    def _connect_sim(self):
-        if self.physics_client is not None:
-            p.disconnect(self.physics_client)
+        # --- INITIALIZE ROBOT WITH BACKEND ---
+        self.backend = PyBulletBackend(self.sim)
+        self.robot = Robot(backend_type=self.backend, urdf_path=self.urdf_path, mode=self.robot_mode)
 
-        flags = p.GUI if self.use_gui else p.DIRECT
-        self.physics_client = p.connect(flags)
-        p.setAdditionalSearchPath(pybullet_data.getDataPath())
-        p.setGravity(0, 0, -9.81)
+        action_dim = 18
+        obs_dim = 18 + 3  # 18 Joint Angles + 3 Command Inputs [vx, vy, omega]
 
-        self.plane = p.loadURDF("plane.urdf")
-        self.robots = []
-        self.robot_joint_indices = []
+        self.action_space = spaces.Box(-1.0, 1.0, shape=(action_dim,), dtype=np.float32)
+        self.observation_space = spaces.Box(-np.inf, np.inf, shape=(obs_dim,), dtype=np.float32)
+        self.min_joint_limits, self.max_joint_limits = self.sim._get_urdf_joint_limits()
 
-        for robot_id in range(self.num_robots):
-            base_pos = self._robot_base_position(robot_id)
-            robot = p.loadURDF(self.urdf_path, basePosition=base_pos, useFixedBase=False)
-            self.robots.append(robot)
-            joint_indices = [
-                i
-                for i in range(p.getNumJoints(robot))
-                if p.getJointInfo(robot, i)[2] == p.JOINT_REVOLUTE
-            ]
-            self.robot_joint_indices.append(joint_indices)
+        self.command = np.zeros(3, dtype=np.float32)  # [vx, vy, omega]
+        self.last_action = np.zeros(action_dim, dtype=np.float32)
+        self.last_last_action = np.zeros(action_dim, dtype=np.float32)
+        self.target_height = 0.122
+        self.collapse_height_fraction = 0.55
+        self.tilt_failure_rad = 0.9
 
-        joint_count = len(self.robot_joint_indices[0])
-        if any(len(indices) != joint_count for indices in self.robot_joint_indices):
-            raise ValueError("All robots must have the same number of revolute joints.")
+        self.start_position = [0.0, 0.0, 0.0]
+        self.max_distance_from_start = 0.0
+        self.max_survival_steps = 0
+        self.default_joint_angles = np.zeros(18, dtype=np.float32)
 
-        lower_limits = []
-        upper_limits = []
-        for joint_index in self.robot_joint_indices[0]:
-            info = p.getJointInfo(self.robots[0], joint_index)
-            lower = info[8]
-            upper = info[9]
-            if lower >= upper:
-                lower = -math.pi
-                upper = math.pi
-            lower_limits.append(lower)
-            upper_limits.append(upper)
+        self.curriculum_phase = CurriculumPhase.STAND_ONLY
+        self.curriculum_stage_requirements = {
+            CurriculumPhase.FORWARD: {"survival_steps": 300, "min_avg_height_ratio": 0.88, "max_avg_roll_pitch": 0.18},
+            CurriculumPhase.TURN_AND_DIRECTION: {"survival_steps": 500, "min_forward_distance": 2.5, "max_lateral_drift": 0.8, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25},
+            CurriculumPhase.OMNI_DIRECTION: {"survival_steps": 600, "min_distance": 5.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25},
+            CurriculumPhase.FULL_COMMAND: {"survival_steps": 750, "min_distance": 8.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.20},
+        }
 
-        limits = np.array([lower_limits, upper_limits], dtype=np.float32)
-        self.joint_lower = np.tile(limits[0], (self.num_robots, 1))
-        self.joint_upper = np.tile(limits[1], (self.num_robots, 1))
-
-        pose = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg
-        initial_pose = pose.to_rad().data.flatten()
-        self.initial_angles = np.tile(initial_pose, (self.num_robots, 1)).astype(np.float32)
-        self.initial_angles = np.clip(
-            self.initial_angles,
-            self.joint_lower,
-            self.joint_upper,
-        )
-
-        self.last_action = np.zeros(self.num_robots * joint_count, dtype=np.float32)
-        self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
-
-    def reset(self, command: np.ndarray | None = None, seed: int | None = None, options: dict | None = None):
-        if self._first_reset:
-            self.step_count = 0
-            self.last_action = np.zeros(self.num_robots * len(self.robot_joint_indices[0]), dtype=np.float32)
-
-            if seed is not None:
-                self.seed(seed)
-
-            if command is not None:
-                self.commands = np.array(command, dtype=np.float32).reshape(self.num_robots, 4)
-            elif self.random_command:
-                self.commands = np.stack([self.sample_command() for _ in range(self.num_robots)])
-            else:
-                self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
-
-            self._reset_pose()
-            self._first_reset = False
-            return self._get_obs(), {}
-
-        p.resetSimulation()
-        p.setGravity(0, 0, -9.81)
-        p.setAdditionalSearchPath(pybullet_data.getDataPath())
-        self.plane = p.loadURDF("plane.urdf")
-        self.robots = []
-        self.robot_joint_indices = []
-
-        for robot_id in range(self.num_robots):
-            base_pos = self._robot_base_position(robot_id)
-            robot = p.loadURDF(self.urdf_path, basePosition=base_pos, useFixedBase=False)
-            self.robots.append(robot)
-            joint_indices = [
-                i
-                for i in range(p.getNumJoints(robot))
-                if p.getJointInfo(robot, i)[2] == p.JOINT_REVOLUTE
-            ]
-            self.robot_joint_indices.append(joint_indices)
-
-        self.step_count = 0
-        self.last_action = np.zeros(self.num_robots * len(self.robot_joint_indices[0]), dtype=np.float32)
-
-        if seed is not None:
-            self.seed(seed)
-
-        if command is not None:
-            self.commands = np.array(command, dtype=np.float32).reshape(self.num_robots, 4)
-        elif self.random_command:
-            self.commands = np.stack([self.sample_command() for _ in range(self.num_robots)])
-        else:
-            self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
-
-        self._reset_pose()
-        return self._get_obs(), {}
-
-    def seed(self, seed: int | None = None):
-        """Set RNG seed for reproducibility (Gym compatibility)."""
-        if seed is None:
-            seed = int(np.random.randint(0, 2**31 - 1))
-        random.seed(seed)
-        np.random.seed(seed)
-        self._seed = seed
-        return [seed]
-
-    def _reset_pose(self):
-        for robot, joint_indices, init_angles in zip(self.robots, self.robot_joint_indices, self.initial_angles):
-            for joint_index, target_angle in zip(joint_indices, init_angles):
-                p.resetJointState(robot, joint_index, target_angle)
-                p.setJointMotorControl2(
-                    bodyIndex=robot,
-                    jointIndex=joint_index,
-                    controlMode=p.POSITION_CONTROL,
-                    targetPosition=target_angle,
-                    force=250,
-                )
-
-        for _ in range(400):
-            p.stepSimulation()
-
-        self.joint_angles = self.initial_angles.copy()
+        self.hud = MetricsHUD(physics_client_id=self.sim.physics_client)
+        self.last_time = time.time()
 
     def sample_command(self) -> np.ndarray:
-        vx = np.random.uniform(-1.0, 1.0)
-        vy = np.random.uniform(-0.5, 0.5)
-        vz = 0.0
-        omega = np.random.uniform(-1.0, 1.0)
-        return np.array([vx, vy, vz, omega], dtype=np.float32)
+        """Samples a command vector [vx, vy, omega] based on active curriculum phase."""
+        phase = self.curriculum_phase
+        stand_probs = {
+            CurriculumPhase.STAND_ONLY: 1.0,
+            CurriculumPhase.FORWARD: 0.25,
+            CurriculumPhase.TURN_AND_DIRECTION: 0.20,
+            CurriculumPhase.OMNI_DIRECTION: 0.15,
+            CurriculumPhase.FULL_COMMAND: 0.15,
+        }
+
+        if np.random.random() < stand_probs.get(phase, 0.15):
+            return np.zeros(3, dtype=np.float32)
+
+        if phase == CurriculumPhase.FORWARD:
+            vx, vy, omega = np.random.uniform(0.15, 0.50), 0.0, 0.0
+        elif phase == CurriculumPhase.TURN_AND_DIRECTION:
+            vx, vy, omega = np.random.uniform(-0.8, 0.8), 0.0, np.random.uniform(-0.8, 0.8)
+        elif phase == CurriculumPhase.OMNI_DIRECTION:
+            vx, vy, omega = np.random.uniform(-0.8, 0.8), np.random.uniform(-0.5, 0.5), np.random.uniform(-0.8, 0.8)
+        else:
+            vx, vy, omega = np.random.uniform(-1.0, 1.0), np.random.uniform(-1.0, 1.0), np.random.uniform(-1.0, 1.0)
+
+        return np.array([vx, vy, omega], dtype=np.float32)
+
+    def reset(self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None):
+        super().reset(seed=seed)
+        self.episode_count += 1
+        self.step_count = 0
+        self.cumulative_reward = 0.0
+        self.robot_reward = 0.0
+        self.is_failed = False
+
+        self.episode_height_sum = 0.0
+        self.episode_roll_sum = 0.0
+        self.episode_pitch_sum = 0.0
+        self.last_reward_components = {}
+        self.episode_reward_components_sum = defaultdict(float)
+
+        spawn_pos = [0.0, 0.0, 0.15]
+        spawn_orn = [0.0, 0.0, 0.0, 1.0]
+
+        # 1. Reset base pose and velocities
+        self.sim.reset_robot_base(spawn_pos, spawn_orn)
+
+        # 2. Reset internal kinematics & hard reset joints in PyBullet
+        self.robot.reset_to_init()
+
+        action_dim = self.action_space.shape[0]
+        self.last_action = np.zeros(action_dim, dtype=np.float32)
+        self.last_last_action = np.zeros(action_dim, dtype=np.float32)
+        self.max_distance_from_start = 0.0
+        self.max_survival_steps = 0
+
+        if self._first_reset:
+            self.curriculum_phase = CurriculumPhase.STAND_ONLY
+            self._first_reset = False
+
+        self.command = np.zeros(3, dtype=np.float32)
+        self.next_cmd_resample_step = self.initial_stand_steps
+
+        pos, _ = self.sim.get_robot_pose()
+        self.start_position = list(pos)
+
+        # Drop settlement
+        self.target_height = self.sim.settle_and_measure_height(
+            target_angles=self.robot.current_rad, 
+            steps=300, 
+            fallback_height=0.122
+        )
+        self.default_joint_angles = self.sim.get_robot_joint_angles()
+        self.default_action = np.clip(self.default_joint_angles / (np.pi / 2.0), -1.0, 1.0)
+
+        if self.use_gui:
+            self.hud.reset()
+            self._update_hud()
+
+        return self._get_obs(), {}
 
     def _get_obs(self) -> np.ndarray:
-        return np.concatenate([self.joint_angles.flatten(), self.commands.flatten()]).astype(np.float32)
+        # Read raw joint angles from backend
+        raw_angles = np.asarray(self.robot.backend.get_joint_angles(), dtype=np.float32).flatten()
 
-    def step(self, action: np.ndarray):
-        action = np.clip(action, self.action_space.low, self.action_space.high).astype(np.float32)
-        self.last_action = action
-        action_matrix = action.reshape(self.num_robots, -1)
-        self.joint_angles = np.clip(
-            self.joint_angles + action_matrix * self.action_scale,
-            self.joint_lower,
-            self.joint_upper,
-        )
+        min_lim = self.min_joint_limits.flatten()
+        max_lim = self.max_joint_limits.flatten()
 
-        for robot, joint_indices, angles in zip(self.robots, self.robot_joint_indices, self.joint_angles):
-            for joint_index, target_angle in zip(joint_indices, angles):
-                p.setJointMotorControl2(
-                    bodyIndex=robot,
-                    jointIndex=joint_index,
-                    controlMode=p.POSITION_CONTROL,
-                    targetPosition=target_angle,
-                    force=250,
+        # Map raw joint radians [min, max] -> normalized [-1, 1]
+        normalized_joints = 2.0 * (raw_angles - min_lim) / (max_lim - min_lim) - 1.0
+        normalized_joints = np.clip(normalized_joints, -1.0, 1.0)
+
+        # Concatenate normalized joints with active command vector
+        obs = np.concatenate([normalized_joints, self.command]).astype(np.float32)
+        return obs
+
+    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
+        previous_action = self.last_action.copy()
+        self.step_count += 1
+        self.total_steps += 1
+
+        self.last_last_action = self.last_action.copy()
+        self.last_action = action.copy()
+
+        # Command resampling ONLY if random_command is True
+        if self.random_command and (
+            self.step_count >= self.next_cmd_resample_step or self._curriculum_advanced):
+            self.command = self.sample_command()
+            random_interval = np.random.randint(self.min_cmd_hold_steps, self.max_cmd_hold_steps + 1)
+            self.next_cmd_resample_step = self.step_count + random_interval
+
+        # Extract active [vx, vy, omega]
+        cmd_vx, cmd_vy, cmd_omega = self.command
+
+        # Mirror input resolution logic to keep robot state synchronized
+        self.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle"
+        self.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
+
+        # Direct Mode: Target joint scaling
+        action_flat = np.asarray(action, dtype=np.float32).flatten()
+        action_clipped = np.clip(action_flat, -1.0, 1.0)
+
+        if self.robot_mode == "direct":
+            # Map [-1, 1] linearly to physical joint limits [min, max]
+            min_lim = self.min_joint_limits.flatten()
+            max_lim = self.max_joint_limits.flatten()
+            
+            target_angles = min_lim + (action_clipped + 1.0) * 0.5 * (max_lim - min_lim)
+        else:
+            # Residual mode mapping logic
+            target_angles = self.default_joint_angles.flatten() + action_clipped * 0.20
+
+        # Apply target joint angles to physics engine
+        self.robot.tick(action=target_angles, physics_substeps=4)
+
+        if self.robot_mode != "kinematics" and self.step_count % 60 == 0:
+            random_force = np.random.uniform(-2.0, 2.0, size=2)
+            self.sim.apply_external_force(force=[random_force[0], random_force[1], 0.0])
+
+        self._update_robot_failure()
+        self._update_distance_metrics()
+        self._update_curriculum()
+
+        # Build next observation preserving active command
+        obs = self._get_obs()
+        reward = self._compute_reward(action_flat, previous_action)
+
+        self.cumulative_reward += reward
+        self.robot_reward += reward
+
+        terminated = self.is_failed
+        truncated = self.step_count >= self.max_episode_steps
+        info = {"reward_components": self.last_reward_components.copy()}
+
+        if self.step_count % 120 == 0 and self.use_gui:
+            self._update_hud()
+
+        if self.use_gui:
+            time.sleep(1.0 / self.control_freq)
+
+        return obs, reward, terminated, truncated, info
+
+    def _update_distance_metrics(self):
+        pos, _ = self.sim.get_robot_pose()
+        self.episode_height_sum += float(pos[2])
+        start_x, start_y, _ = self.start_position
+        dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
+        self.max_distance_from_start = max(self.max_distance_from_start, dist)
+        self.max_survival_steps = max(self.max_survival_steps, self.step_count)
+
+    def _phase_progress_ready(self, next_phase: CurriculumPhase) -> bool:
+        if next_phase not in self.curriculum_stage_requirements:
+            return False
+
+        req = self.curriculum_stage_requirements[next_phase]
+        survival_ok = self.max_survival_steps >= req["survival_steps"]
+        avg_roll = self.episode_roll_sum / max(1, self.step_count)
+        avg_pitch = self.episode_pitch_sum / max(1, self.step_count)
+        max_allowed_angle = req.get("max_avg_roll_pitch", 0.20)
+        stability_ok = (avg_roll <= max_allowed_angle) and (avg_pitch <= max_allowed_angle)
+
+        avg_height = self.episode_height_sum / max(1, self.step_count)
+        required_min_avg_height = self.target_height * req.get("min_avg_height_ratio", 0.85)
+        height_ok = avg_height >= required_min_avg_height
+
+        pos, _ = self.sim.get_robot_pose()
+        dx, dy = pos[0] - self.start_position[0], pos[1] - self.start_position[1]
+        dist_2d = math.hypot(dx, dy)
+
+        distance_ok = True
+        if "min_forward_distance" in req:
+            distance_ok = dx >= req["min_forward_distance"]
+        elif "min_distance" in req:
+            distance_ok = dist_2d >= req["min_distance"]
+
+        drift_ok = abs(dy) <= req["max_lateral_drift"] if "max_lateral_drift" in req else True
+        return survival_ok and height_ok and stability_ok and distance_ok and drift_ok
+
+    def _update_curriculum(self):
+        self._curriculum_advanced = False
+
+        if self.curriculum_phase < CurriculumPhase.FORWARD and (self._phase_progress_ready(CurriculumPhase.FORWARD)):
+            self.curriculum_phase = CurriculumPhase.FORWARD
+            self._curriculum_advanced = True
+        elif self.curriculum_phase < CurriculumPhase.TURN_AND_DIRECTION and self._phase_progress_ready(CurriculumPhase.TURN_AND_DIRECTION):
+            self.curriculum_phase = CurriculumPhase.TURN_AND_DIRECTION
+            self._curriculum_advanced = True
+        elif self.curriculum_phase < CurriculumPhase.OMNI_DIRECTION and self._phase_progress_ready(CurriculumPhase.OMNI_DIRECTION):
+            self.curriculum_phase = CurriculumPhase.OMNI_DIRECTION
+            self._curriculum_advanced = True
+        elif self.curriculum_phase < CurriculumPhase.FULL_COMMAND and self._phase_progress_ready(CurriculumPhase.FULL_COMMAND):
+            self.curriculum_phase = CurriculumPhase.FULL_COMMAND
+            self._curriculum_advanced = True
+
+    def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> float:
+        pos, (roll, pitch, yaw) = self.sim.get_robot_pose_and_rpy()
+        linear_vel, angular_vel = self.sim.get_robot_velocity()
+        current_joints = self.sim.get_robot_joint_angles()
+
+        cmd_vx, cmd_vy, cmd_yaw = self.command
+        cmd_norm = math.hypot(cmd_vx, cmd_vy)
+
+        raw_speed = math.hypot(linear_vel[0], linear_vel[1])
+        # Forgiving zones: ignore noise below 0.04 m/s and 0.05 rad/s
+        filtered_vx, filtered_vy = (linear_vel[0], linear_vel[1]) if raw_speed >= 0.04 else (0.0, 0.0)
+        filtered_speed = raw_speed if raw_speed >= 0.04 else 0.0
+        raw_yaw_rate = abs(angular_vel[2])
+        filtered_yaw_rate = raw_yaw_rate if raw_yaw_rate >= 0.05 else 0.0
+
+        # --- 1. FIXED JITTER PENALTY ---
+        # Scaled way down (0.005) and capped so it can never dominate the reward
+        action_accel = action - 2.0 * self.last_action + self.last_last_action
+        raw_jitter = float(np.mean(np.square(action_accel)))
+        jitter_penalty = -0.005 * min(raw_jitter, 10.0)
+
+        # --- Base Components ---
+        height_error = pos[2] - self.target_height
+        r_height = math.exp(-150.0 * (height_error ** 2))
+        r_stability = math.exp(-25.0 * (roll**2 + pitch**2))
+        r_pose = math.exp(-2.0 * np.mean(np.square(current_joints - self.default_joint_angles)))
+        r_smoothness = math.exp(-0.1 * np.mean(np.square(action - previous_action)))
+
+        r_lin_vel, r_ang_vel, stillness_penalty = 0.0, 0.0, 0.0
+        stand_penalty = 0.0
+
+        # --- 2. COMMAND IS ZERO: STANDING MODE ---
+        if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05:
+            w_height, w_stability, w_pose, w_smoothness = 0.35, 0.35, 0.20, 0.10
+            base_reward = (w_height * r_height) + (w_stability * r_stability) + (w_pose * r_pose) + (w_smoothness * r_smoothness)
+
+            # Soft quadratic penalty on filtered speed (gives a forgiving gap near 0)
+            stand_penalty = -0.5 * filtered_speed - 0.1 * filtered_yaw_rate
+            total_reward = base_reward + stand_penalty
+
+        # --- 3. COMMAND IS NON-ZERO: WALKING MODE ---
+        else:
+            is_moving = (filtered_speed > 0.0) or (filtered_yaw_rate > 0.0)
+            target_vx, target_vy = cmd_vx * self.max_robot_speed, cmd_vy * self.max_robot_speed
+            target_speed = math.hypot(target_vx, target_vy)
+
+            if not is_moving:
+                # If the command says move but the robot stays effectively still,
+                # give a real penalty instead of a neutral reward.
+                stillness_penalty = -0.10
+                total_reward = stillness_penalty
+            else:
+                lin_vel_error = (filtered_vx - target_vx)**2 + (filtered_vy - target_vy)**2
+                r_lin_vel = math.exp(-25.0 * lin_vel_error)
+                r_ang_vel = math.exp(-15.0 * ((filtered_yaw_rate - cmd_yaw)**2))
+
+                if target_speed > 0.08 and raw_speed < 0.03:
+                    r_lin_vel = 0.0
+                    stillness_penalty = -0.05
+
+                w_lin_vel, w_ang_vel, w_height, w_stability, w_smoothness = 0.55, 0.15, 0.10, 0.12, 0.08
+                total_reward = (
+                    (w_lin_vel * r_lin_vel) + (w_ang_vel * r_ang_vel) + (w_height * r_height)
+                    + (w_stability * r_stability) + (w_smoothness * r_smoothness) + stillness_penalty
                 )
 
-        for _ in range(self.frame_skip):
-            p.stepSimulation()
+        step_reward = float(total_reward / 10.0)
+        alive_bonus = 0.01
+        final_reward = step_reward + jitter_penalty + alive_bonus
 
-        self.step_count += 1
-        observation = self._get_obs()
-        reward = self._compute_reward()
-        done = self._is_done()
-        if self.step_count <= self._min_steps_before_done:
-            done = False
-        info = {"step": self.step_count}
+        self.last_reward_components = {
+            "height": float(r_height),
+            "stability": float(r_stability),
+            "pose": float(r_pose),
+            "smoothness": float(r_smoothness),
+            "lin_vel": float(r_lin_vel),
+            "ang_vel": float(r_ang_vel),
+            "jitter_penalty": float(jitter_penalty),
+            "stand_penalty": float(stand_penalty),
+            "stillness_penalty": float(stillness_penalty),
+            "total": final_reward,
+        }
+        for k, v in self.last_reward_components.items():
+            self.episode_reward_components_sum[k] += v
 
-        # Gymnasium-style return: (obs, reward, terminated, truncated, info)
-        if done:
-            terminated = True
-            truncated = False
-        else:
-            terminated = False
-            truncated = False
+        return final_reward
+    
+    def get_reward_component_averages(self) -> Dict[str, float]:
+        steps = max(1, self.step_count)
+        return {k: v / steps for k, v in self.episode_reward_components_sum.items()}
 
-        return observation, float(reward), terminated, truncated, info
+    def get_current_robot_metrics(self) -> List[Dict[str, Any]]:
+        linear_vel, angular_vel = self.sim.get_robot_velocity()
+        speed = float(math.hypot(linear_vel[0], linear_vel[1]))
+        yaw_rate = float(abs(angular_vel[2]))
 
-    def _compute_reward(self) -> float:
-        rewards = []
-        for robot_id, robot in enumerate(self.robots):
-            linear_vel, angular_vel = p.getBaseVelocity(robot)
-            _, orientation = p.getBasePositionAndOrientation(robot)
-            roll, pitch, _ = p.getEulerFromQuaternion(orientation)
-            command = self.commands[robot_id]
+        metrics = {
+            "alive": not self.is_failed,
+            "phase_name": self.curriculum_phase.name,
+            "reward": float(self.cumulative_reward),
+            "speed": speed,
+            "yaw_rate": yaw_rate,
+            "distance_from_start": float(self.max_distance_from_start),
+            "survival_steps": int(self.max_survival_steps),
+        }
+        return [metrics]
 
-            forward_reward = command[0] * linear_vel[0] + command[1] * linear_vel[1]
-            rotation_reward = command[3] * angular_vel[2]
-            stability_penalty = abs(roll) + abs(pitch)
-            action_penalty = float(np.sum(np.square(self.last_action.reshape(self.num_robots, -1)[robot_id]))) * 0.01
+    def _update_hud(self):
+        if not self.use_gui:
+            return
+        now = time.time()
+        fps = 1.0 / max(now - self.last_time, 1e-5)
+        self.last_time = now
+        pos, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy()
+        self.hud.update(
+            episode=self.episode_count, step=self.total_steps, robot_rewards=[self.robot_reward],
+            cmd_vel=self.command, fps=fps, height=pos[2], roll_pitch=(math.degrees(roll), math.degrees(pitch))
+        )
 
-            rewards.append(0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty)
+    def _update_robot_failure(self):
+        if self.is_failed:
+            return
 
-        return float(np.sum(rewards))
+        position, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy()
+        collapse_threshold = max(0.04, self.collapse_height_fraction * self.target_height)
+        is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad
+        is_collapsed = position[2] < collapse_threshold
 
-    def _is_done(self) -> bool:
-        for robot in self.robots:
-            _, orientation = p.getBasePositionAndOrientation(robot)
-            roll, pitch, _ = p.getEulerFromQuaternion(orientation)
-            if abs(roll) > 0.7 or abs(pitch) > 0.7:
-                return True
-            position, _ = p.getBasePositionAndOrientation(robot)
-            if position[2] < 0.02:
-                return True
-
-        if self.step_count >= self.max_episode_steps:
-            return True
-        return False
-
-    def render(self, mode="human"):
-        if self.use_gui:
-            return None
-        raise NotImplementedError("Render is only supported with use_gui=True")
-
-    def _robot_base_position(self, robot_id: int) -> list[float]:
-        cols = int(math.sqrt(self.num_robots - 1)) + 1
-        row = robot_id // cols
-        col = robot_id % cols
-        x = (col - (cols - 1) / 2.0) * self.robot_spacing
-        y = (row - (cols - 1) / 2.0) * self.robot_spacing
-        return [x, y, 0.1]
+        if is_tilted or is_collapsed:
+            self.is_failed = True
 
     def close(self):
-        if self.physics_client is not None:
-            p.disconnect(self.physics_client)
-            self.physics_client = None
+        self.sim.disconnect()

ml/evaluate.py

@@ -1,78 +0,0 @@
-import argparse
-import numpy as np
-
-from .env import JackBotEnv
-
-
-def evaluate(
-    model_path: str,
-    episodes: int = 5,
-    use_gui: bool = False,
-    num_robots: int = 1,
-    robot_spacing: float = 0.5,
-    start_pose: str = "init_deg",
-):
-    try:
-        from stable_baselines3 import PPO
-    except ImportError as exc:
-        raise ImportError(
-            "stable-baselines3 is required for evaluation. Install with: pip install stable-baselines3"
-        ) from exc
-
-    env = JackBotEnv(
-        use_gui=use_gui,
-        random_command=False,
-        num_robots=num_robots,
-        robot_spacing=robot_spacing,
-        start_pose=start_pose,
-    )
-    model = PPO.load(model_path)
-
-    for episode in range(episodes):
-        reset_res = env.reset()
-        # handle Gym / Gymnasium compatibility: reset may return (obs, info)
-        if isinstance(reset_res, tuple) and len(reset_res) == 2:
-            obs, _ = reset_res
-        else:
-            obs = reset_res
-
-        done = False
-        episode_reward = 0.0
-
-        while not done:
-            # pass only the observation to the policy
-            action, _ = model.predict(obs, deterministic=True)
-
-            step_res = env.step(action)
-            # Gymnasium-style: (obs, reward, terminated, truncated, info)
-            if isinstance(step_res, tuple) and len(step_res) == 5:
-                obs, reward, terminated, truncated, info = step_res
-                done = bool(terminated or truncated)
-            else:
-                # legacy Gym: (obs, reward, done, info)
-                obs, reward, done, info = step_res
-
-            episode_reward += float(reward)
-
-        print(f"Episode {episode + 1}: reward={episode_reward:.2f}")
-
-    env.close()
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Evaluate a trained JackBot policy.")
-    parser.add_argument("--model-path", type=str, required=True)
-    parser.add_argument("--episodes", type=int, default=5)
-    parser.add_argument("--gui", action="store_true")
-    parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the environment")
-    parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
-    parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
-    args = parser.parse_args()
-    evaluate(
-        args.model_path,
-        episodes=args.episodes,
-        use_gui=args.gui,
-        num_robots=args.num_robots,
-        robot_spacing=args.robot_spacing,
-        start_pose=args.start_pose,
-    )

ml/model.py

@@ -1,49 +0,0 @@
-import torch
-import torch.nn as nn
-from torch.distributions import Normal
-
-
-class ActorCritic(nn.Module):
-    def __init__(self, obs_dim: int, action_dim: int, hidden_sizes=(256, 256)):
-        super().__init__()
-        self.backbone = nn.Sequential(
-            nn.Linear(obs_dim, hidden_sizes[0]),
-            nn.ReLU(),
-            nn.Linear(hidden_sizes[0], hidden_sizes[1]),
-            nn.ReLU(),
-        )
-
-        self.mean_head = nn.Linear(hidden_sizes[1], action_dim)
-        self.value_head = nn.Linear(hidden_sizes[1], 1)
-        self.log_std = nn.Parameter(torch.zeros(action_dim, dtype=torch.float32))
-
-    def forward(self, obs: torch.Tensor):
-        x = self.backbone(obs)
-        mean = self.mean_head(x)
-        std = self.log_std.exp()
-        value = self.value_head(x).squeeze(-1)
-        return mean, std, value
-
-    def get_action(self, obs: torch.Tensor):
-        mean, std, value = self.forward(obs)
-        dist = Normal(mean, std)
-        action = dist.sample()
-        log_prob = dist.log_prob(action).sum(-1)
-        return action, log_prob, value
-
-    def evaluate_actions(self, obs: torch.Tensor, actions: torch.Tensor):
-        mean, std, value = self.forward(obs)
-        dist = Normal(mean, std)
-        log_prob = dist.log_prob(actions).sum(-1)
-        entropy = dist.entropy().sum(-1)
-        return value, log_prob, entropy
-
-    def save(self, path: str):
-        torch.save(self.state_dict(), path)
-
-    @classmethod
-    def load(cls, path: str, obs_dim: int, action_dim: int, hidden_sizes=(256, 256)):
-        model = cls(obs_dim, action_dim, hidden_sizes)
-        model.load_state_dict(torch.load(path, map_location=torch.device("cpu")))
-        model.eval()
-        return model

ml/pretrain_bc.py

@@ -0,0 +1,155 @@
+"""
+ml/pretrain_bc.py - Behavioral Cloning teacher-data pipeline.
+
+This script collects observation/action pairs from the kinematics solver, then uses
+those pairs as training data for a PPO policy. In practice it serves as a teacher-
+student pretraining step: the kinematics controller generates sample trajectories, and
+this file teaches the policy to imitate those behavior patterns before PPO training.
+"""
+import sys
+import time
+import argparse
+from pathlib import Path
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, TensorDataset
+from stable_baselines3 import PPO
+from tqdm import tqdm  # <--- Progress Bar Support
+
+# Ensure project root is in sys.path
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+from ml.env import JackBotEnv, CurriculumPhase
+
+
+def collect_kinematics_dataset(num_samples: int = 100_000, use_gui: bool = False):
+    """Collects (Observation, Action) pairs directly from Kinematics Teacher."""
+    print(f"\n[Pretrain] Collecting {num_samples} samples from Kinematics Teacher (GUI={use_gui})...")
+    
+    # Disable random command resampling inside env.step so manual command locks persist
+    env = JackBotEnv(use_gui=use_gui, robot_mode="kinematics", random_command=False)
+    env.curriculum_phase = CurriculumPhase.FULL_COMMAND
+
+    observations = []
+    actions = []
+
+    obs, _ = env.reset()
+    
+    # --- PROGRESS BAR: Data Collection ---
+    pbar = tqdm(range(num_samples), desc=" Collecting Data", unit="step")
+    for i in pbar:
+        # 1. Update command and vector targets every 120 steps
+        if i % 120 == 0:
+            env.command = env.sample_command()
+            cmd_vx, cmd_vy, cmd_omega = env.command
+            env.robot.robot_state = (
+                "walking"
+                if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01)
+                else "idle"
+            )
+            env.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
+
+        # 2. Capture observation BEFORE stepping environment
+        current_obs = env._get_obs()
+
+        # 3. Step environment ONCE (updates kinematics solver, PyBullet physics, and computes IK)
+        obs, _, terminated, truncated, _ = env.step(np.zeros(18, dtype=np.float32))
+
+        # 4. Extract procedural IK joint targets computed during this step
+        target_ik_rad = env.robot.current_rad.data.flatten().copy()
+
+        # Step 5: Convert target radians directly to [-1, 1] relative to joint limits
+        min_lim = env.min_joint_limits.flatten()
+        max_lim = env.max_joint_limits.flatten()
+
+        normalized_action = 2.0 * (target_ik_rad - min_lim) / (max_lim - min_lim) - 1.0
+        normalized_action = np.clip(normalized_action, -1.0, 1.0)
+
+        # Step 6: Store matching input (obs) and target ground truth (normalized_action)
+        observations.append(current_obs.copy())
+        actions.append(normalized_action.copy())
+
+        if use_gui:
+            time.sleep(1.0 / 60.0)
+
+        # 7. Handle episode boundaries using terminated and truncated
+        if terminated or truncated:
+            obs, _ = env.reset()
+
+    env.close()
+    print("[Pretrain] Data collection complete!\n")
+    return np.array(observations, dtype=np.float32), np.array(actions, dtype=np.float32)
+
+
+def pretrain_policy(
+    save_path: str = "ml/checkpoints/jackbot_kinematics_base.zip",
+    epochs: int = 15,
+    batch_size: int = 256,
+    num_samples: int = 100_000,
+    use_gui: bool = False
+):
+    # Collect dataset from Kinematics teacher
+    obs_data, action_data = collect_kinematics_dataset(num_samples=num_samples, use_gui=use_gui)
+
+    # Initialize Dummy Env & Fresh SB3 PPO Model
+    dummy_env = JackBotEnv(use_gui=False, robot_mode="direct")
+    model = PPO("MlpPolicy", dummy_env, learning_rate=5e-4, verbose=0, device="cpu")
+
+    # Extract PyTorch Policy Network & Optimizer
+    policy = model.policy
+    optimizer = torch.optim.Adam(policy.parameters(), lr=5e-4)
+    loss_fn = nn.MSELoss()
+
+    # Convert to PyTorch Dataloader
+    dataset = TensorDataset(torch.tensor(obs_data), torch.tensor(action_data))
+    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+    print(f"[Pretrain] Pre-training Policy Network ({epochs} Epochs)...")
+    policy.train()
+
+    # --- PROGRESS BAR: Epoch Training ---
+    for epoch in range(epochs):
+        epoch_loss = 0.0
+        
+        batch_pbar = tqdm(loader, desc=f" Epoch {epoch + 1:02d}/{epochs:02d}", leave=True, unit="batch")
+        for batch_obs, batch_actions in batch_pbar:
+            optimizer.zero_grad()
+            
+            distribution = policy.get_distribution(batch_obs)
+            predicted_actions = distribution.distribution.mean
+
+            loss = loss_fn(predicted_actions, batch_actions)
+            loss.backward()
+            optimizer.step()
+
+            current_loss = loss.item()
+            epoch_loss += current_loss * len(batch_obs)
+            
+            # Dynamic loss update in progress bar tail
+            batch_pbar.set_postfix({"loss": f"{current_loss:.6f}"})
+
+        avg_loss = epoch_loss / len(dataset)
+        tqdm.write(f"  └─ Epoch {epoch + 1:02d}/{epochs:02d} Complete | Mean MSE Loss: {avg_loss:.6f}")
+
+    # Save SB3 Model Checkpoint
+    out_file = Path(save_path)
+    out_file.parent.mkdir(parents=True, exist_ok=True)
+    model.save(out_file)
+    dummy_env.close()
+    print(f"\n[Pretrain] Successfully saved pre-trained base model to: {out_file.resolve()}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="JackBot Behavioral Cloning Pre-trainer")
+    parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering during collection")
+    parser.add_argument("--num-samples", type=int, default=100_000, help="Number of dataset samples to collect")
+    parser.add_argument("--epochs", type=int, default=15, help="Number of BC training epochs")
+    parser.add_argument("--save-path", type=str, default="ml/checkpoints/jackbot_kinematics_base.zip", help="Output path for pre-trained model .zip")
+    args = parser.parse_args()
+
+    pretrain_policy(
+        save_path=args.save_path,
+        epochs=args.epochs,
+        num_samples=args.num_samples,
+        use_gui=args.gui
+    )

ml/run_eval.py

@@ -1,39 +1,153 @@
-"""Run a trained policy in the PyBullet sim for quick inspection.
+"""
+ml/run_eval.py - Evaluate a saved JackBot policy across fixed command phases.
 
-Usage:
-    python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3 --gui
+This script loads a trained PPO checkpoint, instantiates the Gymnasium environment in
+non-random mode, and runs deterministic evaluation episodes for several command
+regimes. It is used to measure whether a policy can survive, move, and maintain
+stability under forward, turning, lateral, and omni-direction commands.
 """
 
 import argparse
+import json
+import time
 import sys
 from pathlib import Path
+import numpy as np
+from stable_baselines3 import PPO
 
-ROOT_DIR = Path(__file__).resolve().parent.parent
-if str(ROOT_DIR) not in sys.path:
-    sys.path.insert(0, str(ROOT_DIR))
-
-from ml.evaluate import evaluate
+# Ensure project root is in sys.path
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+from ml.env import JackBotEnv, CurriculumPhase
 
 
 def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model", type=str, required=True, help="Path to the trained model file (.zip)")
-    parser.add_argument("--episodes", type=int, default=3, help="Number of 'rounds' to run. One episode lasts from reset until the robot falls over or the time limit is reached.")
-    parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during evaluation")
-    parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the evaluation environment")
-    parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
-    parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
+    parser = argparse.ArgumentParser(description="JackBot Phase-by-Phase Policy Evaluator")
+    parser.add_argument("--model", type=str, required=True, help="Path to trained model checkpoint (.zip)")
+    parser.add_argument("--episodes-per-phase", type=int, default=1, help="Number of test episodes per phase")
+    parser.add_argument("--max-steps-per-episode", type=int, default=600, help="Max simulation steps per episode")
+    parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
+    parser.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to save evaluation summary")
     args = parser.parse_args()
 
-    evaluate(
-        model_path=args.model,
-        episodes=args.episodes,
+    print(f"[Eval] Loading policy model from: {args.model}")
+    model = PPO.load(args.model, device="cpu")
+
+    # Instantiate environment with random_command disabled so our test command stays locked
+    env = JackBotEnv(
         use_gui=args.gui,
-        num_robots=args.num_robots,
-        robot_spacing=args.robot_spacing,
-        start_pose=args.start_pose,
+        random_command=False,
+        max_episode_steps=args.max_steps_per_episode,
+        robot_mode="direct"
     )
 
+    # Multi-Phase Configurations Suite
+    phase_configs = [
+        #(CurriculumPhase.STAND_ONLY,          "STAND",                 np.array([0.0, 0.0, 0.0], dtype=np.float32)),
+        (CurriculumPhase.FORWARD,             "FORWARD GAIT",           np.array([1.0, 0.0, 0.0], dtype=np.float32)),
+        (CurriculumPhase.TURN_AND_DIRECTION,  "FORWARD + YAW TURN",     np.array([0.5, 0.0, 0.4], dtype=np.float32)),
+        (CurriculumPhase.OMNI_DIRECTION,      "STRIDE LATERAL",         np.array([0.5, 0.5, 0.0], dtype=np.float32)),
+        (CurriculumPhase.FULL_COMMAND,        "FULL OMNI COMBINATION",  np.array([0.3, 0.2, 0.3], dtype=np.float32)),
+    ]
+
+    phase_summary = []
+
+    try:
+        print("\n" + "=" * 75)
+        print(" STARTING MULTI-PHASE EVALUATION SUITE")
+        print("=" * 75)
+
+        for phase_enum, phase_name, test_cmd in phase_configs:
+            print(f"\n▶ Testing Phase [{phase_enum.value}]: {phase_enum.name} ({phase_name})")
+            print(f"  Target Command Vector [vx, vy, omega]: {test_cmd.tolist()}")
+
+            ep_rewards = []
+            ep_steps = []
+            ep_distances = []
+
+            for ep in range(args.episodes_per_phase):
+                obs, _ = env.reset()
+
+                # Force environment into active curriculum phase and lock command BEFORE getting obs
+                env.curriculum_phase = phase_enum
+                env.command = test_cmd.copy()
+                cmd_vx, cmd_vy, cmd_omega = test_cmd
+                env.robot.robot_state = (
+                    "walking"
+                    if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01)
+                    else "idle"
+                )
+                env.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
+
+                # Get correct observation with test_cmd attached
+                obs = env._get_obs()
+
+                done = False
+                total_reward = 0.0
+                steps = 0
+
+                while not done:
+                    # Predict deterministic action from policy
+                    action, _ = model.predict(obs, deterministic=True)
+                    
+                    # Step environment
+                    obs, reward, terminated, truncated, _ = env.step(action)
+                    done = terminated or truncated
+
+                    total_reward += float(reward)
+                    steps += 1
+
+                    if args.gui:
+                        time.sleep(1.0 / 60.0)
+
+                status_str = "FAILED (Collapsed)" if terminated else "SUCCESS (Completed)"
+                dist = float(env.max_distance_from_start)
+                print(
+                    f"   └─ Ep {ep + 1}/{args.episodes_per_phase}: {status_str:<19} | "
+                    f"Steps: {steps:<4} | Reward: {total_reward:+.2f} | Max Dist: {dist:.2f}m"
+                )
+
+                ep_rewards.append(total_reward)
+                ep_steps.append(steps)
+                ep_distances.append(dist)
+
+            phase_summary.append({
+                "phase_id": phase_enum.value,
+                "phase_name": phase_enum.name,
+                "label": phase_name,
+                "command": test_cmd.tolist(),
+                "mean_reward": float(np.mean(ep_rewards)),
+                "mean_steps": float(np.mean(ep_steps)),
+                "mean_distance": float(np.mean(ep_distances)),
+            })
+
+    finally:
+        env.close()
+
+    # Print Summary Table
+    print("\n" + "=" * 80)
+    print(" EVALUATION SUMMARY REPORT")
+    print("=" * 80)
+    print(f"{'Phase ID & Name':<25} | {'Label':<22} | {'Reward':<8} | {'Steps':<6} | {'Max Dist':<8}")
+    print("-" * 80)
+    for res in phase_summary:
+        phase_str = f"[{res['phase_id']}] {res['phase_name']}"
+        print(
+            f"{phase_str:<25} | "
+            f"{res['label']:<22} | "
+            f"{res['mean_reward']:<+8.2f} | "
+            f"{res['mean_steps']:<6.0f} | "
+            f"{res['mean_distance']:<8.2f}m"
+        )
+    print("=" * 80)
+
+    # Save JSON Report if requested
+    if args.save_metrics:
+        out_path = Path(args.save_metrics)
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(out_path, "w") as f:
+            json.dump({"model_path": str(args.model), "summary": phase_summary}, f, indent=4)
+        print(f"\n[Eval] Saved report to: {out_path.resolve()}")
+
 
 if __name__ == "__main__":
-    main()
+    main()

ml/run_eval_training.py

@@ -0,0 +1,78 @@
+"""
+ml/run_eval_training.py - Run a kinematics-mode reward benchmark.
+
+This script repeatedly resets the JackBot environment in kinematics mode and applies
+fixed command vectors for each curriculum phase. It is intended as a lightweight
+benchmark to inspect reward components, movement quality, and survival behavior without
+requiring an already-trained PPO model.
+"""
+import time
+import numpy as np
+import sys
+from pathlib import Path
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+from ml.env import JackBotEnv, CurriculumPhase
+
+def evaluate_kinematics(episode_length: int = 1000):
+    env = JackBotEnv(
+        use_gui=True, 
+        random_command=False, 
+        max_episode_steps=episode_length,
+        robot_mode="kinematics"
+    )
+
+    print("\n" + "=" * 70)
+    print(" RUNNING MULTI-PHASE REWARD BENCHMARK (KINEMATICS MODE)")
+    print("=" * 70 + "\n")
+
+    phase_configs = [
+        (CurriculumPhase.STAND_ONLY,          "STAND",               np.array([1.0, 0.0, 0.0], dtype=np.float32)),
+        (CurriculumPhase.FORWARD,             "FORWARD GAIT",        np.array([1.0, 0.0, 0.0], dtype=np.float32)),
+        (CurriculumPhase.TURN_AND_DIRECTION,  "FORWARD + YAW TURN",  np.array([0.5, 0.0, 0.4], dtype=np.float32)),
+        (CurriculumPhase.OMNI_DIRECTION,      "STRIDE LATERAL",      np.array([0.5, 0.5, 0.0], dtype=np.float32)),
+        (CurriculumPhase.FULL_COMMAND,        "FULL OMNI COMBINATION",np.array([0.3, 0.2, 0.3], dtype=np.float32)),
+    ]
+
+    for phase_enum, label, cmd in phase_configs:
+        obs, _ = env.reset()
+        
+        env.curriculum_phase = phase_enum
+        env.command = cmd.copy()
+        
+        done = False
+        total_reward = 0.0
+        step_count = 0
+
+        while not done:
+            dummy_action = np.zeros(18, dtype=np.float32)
+            obs, reward, terminated, truncated, _ = env.step(dummy_action)
+            total_reward += reward
+            step_count += 1
+            done = terminated or truncated
+
+            time.sleep(1.0 / 60.0)
+
+        comp_averages = env.get_reward_component_averages()
+
+        print(f"\n--- Episode Stage: [{phase_enum.name}] ({label}) ---")
+        print(f"Command Applied: vx={cmd[0]:.2f}, vy={cmd[1]:.2f}, yaw={cmd[2]:.2f}")
+        print(f"Total Episode Reward: {total_reward:.4f}")
+        print("Component Step Averages:")
+        for name, value in comp_averages.items():
+            print(f"  • {name:<20}: {value:+.5f}")
+
+        metrics = env.get_current_robot_metrics()
+        dist = metrics[0]["distance_from_start"] if metrics else 0.0
+        speed = metrics[0]["speed"] if metrics else 0.0
+        avg_reward = total_reward / max(1, step_count)
+
+        print(f"  ├─ Average Reward / Step: {avg_reward:.4f}")
+        print(f"  ├─ Distance Travelled: {dist:.2f} m")
+        print(f"  ├─ Actual Avg Speed: {speed:.2f} m/s")
+        print(f"  └─ Steps Survived: {step_count} / {episode_length}")
+        print("-" * 70)
+
+    env.close()
+
+if __name__ == "__main__":
+    evaluate_kinematics()

ml/run_train.py

@@ -1,47 +1,161 @@
-"""Minimal training launcher for quick experiments.
+"""PPO training launcher for JackBot.
+
+This script is the main entry point for training a policy in the JackBot Gymnasium
+environment. It creates a vectorized environment, optionally loads a pretrained base
+model, runs Stable-Baselines3 PPO for a configured number of timesteps, and saves
+checkpoints plus evaluation artifacts during training.
 
 Usage:
-    python ml/run_train.py --timesteps 50000 --model ml/checkpoints/ppo_joint_command
-
-This is a convenience wrapper around `ml.train.train` with friendly defaults
-for interactive experimentation.
+    python ml/run_train.py --total-timesteps 1500000 --gui
 """
 
 import argparse
+import os
+import re
 import sys
+import warnings
 from pathlib import Path
+from stable_baselines3 import PPO
+from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv
+from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
 
-ROOT_DIR = Path(__file__).resolve().parent.parent
-if str(ROOT_DIR) not in sys.path:
-    sys.path.insert(0, str(ROOT_DIR))
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+from ml.env import JackBotEnv
+from ml.callbacks import CurriculumCallback, RewardLoggerCallback
 
-from ml.train import train
+# Silence SB3's UserWarning about SubprocVecEnv vs DummyVecEnv
+warnings.filterwarnings("ignore", category=UserWarning, module="stable_baselines3")
+
+
+def get_next_run_number(save_dir: str) -> int:
+    """Scans the save directory for existing ppo<number> patterns and returns the next integer."""
+    if not os.path.exists(save_dir):
+        return 1
+    
+    existing_numbers = []
+    for item in os.listdir(save_dir):
+        # Match pattern ppo followed by numbers (e.g., ppo1, ppo_1, jackbot_ppo12)
+        matches = re.findall(r"ppo_?(\d+)", item, re.IGNORECASE)
+        for m in matches:
+            existing_numbers.append(int(m))
+            
+    return max(existing_numbers, default=0) + 1
+
+
+def make_env(rank: int, use_gui: bool = False, seed: int = 0):
+    """Utility helper to instantiate parallel JackBot environments."""
+    def _init():
+        env = JackBotEnv(
+            use_gui=use_gui,
+            random_command=True,
+        )
+        env.reset(seed=seed + rank)
+        return env
+    return _init
 
 
 def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--timesteps", type=int, default=50000, help="Total number of 'practice steps'.")
-    parser.add_argument("--model", type=str, default="ml/checkpoints/ppo_joint_command", help="Path to save the trained model")
-    parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducibility")
-    parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto'")
-    parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during training")
-    parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the training environment")
-    parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
-    parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
+    parser = argparse.ArgumentParser(description="JackBot PPO Curriculum Trainer")
+    parser.add_argument("--num-workers", type=int, default=16, help="Number of parallel sub-process environments")
+    parser.add_argument("--total-timesteps", type=int, default=1_500_000, help="Total training timesteps")
+    parser.add_argument("--log-dir", type=str, default="ml/logs", help="Directory for TensorBoard logs")
+    parser.add_argument("--save-dir", type=str, default="ml/checkpoints", help="Directory for model checkpoints")
+    parser.add_argument("--save-freq", type=int, default=50_000, help="Checkpoint save frequency (steps)")
+    parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
+    parser.add_argument(
+        "--pretrained-model",
+        type=str,
+        default=None,
+        help="Path to pre-trained base model checkpoint (.zip) to start PPO training from"
+    )
     args = parser.parse_args()
 
-    Path(args.model).parent.mkdir(parents=True, exist_ok=True)
-    train(
-        total_timesteps=args.timesteps,
-        model_path=args.model,
-        seed=args.seed,
-        device=args.device,
-        use_gui=args.gui,
-        num_robots=args.num_robots,
-        robot_spacing=args.robot_spacing,
-        start_pose=args.start_pose,
+    os.makedirs(args.log_dir, exist_ok=True)
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    # Automatically determine next run number (e.g. ppo1, ppo2, ppo3...)
+    run_num = get_next_run_number(args.save_dir)
+    ppo_name = f"ppo{run_num}"
+
+    print(f"[Train] Initializing Run #{run_num} ('{ppo_name}') with {args.num_workers} parallel workers...")
+    
+    env_fns = [
+        make_env(rank=i, use_gui=(args.gui if i == 0 else False)) 
+        for i in range(args.num_workers)
+    ]
+    vec_env = SubprocVecEnv(env_fns)
+
+    # Initialize or Load PPO Policy Model
+    if args.pretrained_model and os.path.exists(args.pretrained_model):
+        print(f"[Train] Loading pre-trained base knowledge from: {args.pretrained_model}")
+        model = PPO.load(
+            args.pretrained_model,
+            env=vec_env,
+            learning_rate=5e-5,  # Lower learning rate so RL fine-tunes without destroying base gait
+            ent_coef=0.0,
+            target_kl=0.05,
+            vf_coef=0.5,
+            max_grad_norm=0.5,
+            verbose=2,
+            tensorboard_log=args.log_dir,
+            device="cpu",
+        )
+    else:
+        print("[Train] No base model provided. Starting training from scratch...")
+        model = PPO(
+            policy="MlpPolicy",
+            env=vec_env,
+            learning_rate=5e-5,
+            n_steps=256,
+            batch_size=256,
+            n_epochs=10,
+            gamma=0.99,
+            gae_lambda=0.95,
+            clip_range=0.2,
+            ent_coef=0.0,
+            target_kl=0.05,
+            vf_coef=0.5,
+            max_grad_norm=0.5,
+            verbose=2,
+            tensorboard_log=args.log_dir,
+            device="cpu",
+        )
+    model.policy.log_std.data.fill_(-2.0)
+    # Setup Callbacks with ppo<number> naming
+    checkpoint_callback = CheckpointCallback(
+        save_freq=max(1, args.save_freq // args.num_workers),
+        save_path=args.save_dir,
+        name_prefix=f"jackbot_{ppo_name}",
     )
+    reward_logger_callback = RewardLoggerCallback(verbose=1)
+    curriculum_callback = CurriculumCallback()
+
+    eval_env = DummyVecEnv([lambda: JackBotEnv(use_gui=False, random_command=True)])
+    best_model_path = os.path.join(args.save_dir, f"best_model_{ppo_name}")
+
+    eval_callback = EvalCallback(
+        eval_env,
+        best_model_save_path=best_model_path,
+        log_path="ml/logs/results",
+        eval_freq=max(1, 50_000 // args.num_workers),
+        deterministic=True,
+        render=False,
+    )
+
+    print(f"[Train] Starting training for {args.total_timesteps} timesteps...")
+    try:
+        model.learn(
+            total_timesteps=args.total_timesteps,
+            callback=[checkpoint_callback, reward_logger_callback, curriculum_callback, eval_callback],
+            progress_bar=True,
+        )
+        final_model_path = os.path.join(args.save_dir, f"jackbot_{ppo_name}_final.zip")
+        model.save(final_model_path)
+        print(f"[Train] Training complete! Saved final model to {final_model_path}")
+    finally:
+        vec_env.close()
+        eval_env.close()
 
 
 if __name__ == "__main__":
-    main()
+    main()

ml/sim_manager.py

@@ -1,105 +0,0 @@
-from typing import Dict, Any, List, Tuple
-import pybullet as p
-import pybullet_data
-import numpy as np
-
-
-class SimManager:
-    """
-    Manages the PyBullet simulation lifecycle and live HUD overlays.
-    """
-
-    def __init__(self, use_gui: bool = True):
-        self.use_gui = use_gui
-        self.physics_client = None
-        self.debug_text_ids: Dict[str, int] = {}
-
-    def connect(self):
-        """Connects to PyBullet and sets up the basic physics world."""
-        if self.physics_client is not None:
-            p.disconnect(self.physics_client)
-
-        flags = p.GUI if self.use_gui else p.DIRECT
-        self.physics_client = p.connect(flags)
-
-        p.setAdditionalSearchPath(pybullet_data.getDataPath())
-        p.setGravity(0, 0, -9.81)
-
-        if self.use_gui:
-            # Disable unnecessary side panels for a clean UI
-            p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)
-            p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0)
-            p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0)
-            p.configureDebugVisualizer(p.COV_ENABLE_RGB_OUTPUT_PREVIEW, 0)
-
-    def load_scene(
-        self, urdf_path: str, num_robots: int, robot_spacing: float, base_pos_fn
-    ) -> Tuple[int, List[int], List[List[int]]]:
-        """Loads plane and robots into the active PyBullet simulation."""
-        plane_id = p.loadURDF("plane.urdf")
-        robots = []
-        robot_joint_indices = []
-
-        for robot_id in range(num_robots):
-            base_pos = base_pos_fn(robot_id, num_robots, robot_spacing)
-            robot = p.loadURDF(urdf_path, basePosition=base_pos, useFixedBase=False)
-            robots.append(robot)
-
-            joint_indices = [
-                i
-                for i in range(p.getNumJoints(robot))
-                if p.getJointInfo(robot, i)[2] == p.JOINT_REVOLUTE
-            ]
-            robot_joint_indices.append(joint_indices)
-
-        return plane_id, robots, robot_joint_indices
-
-    def disconnect(self):
-        """Safely disconnects from the simulation."""
-        if self.physics_client is not None:
-            p.disconnect(self.physics_client)
-            self.physics_client = None
-            self.debug_text_ids.clear()
-
-    def update_hud(self, stats: Dict[str, Any]):
-        """
-        Renders live training metrics onto the PyBullet 3D viewport.
-        Uses replaceItemUniqueId to prevent flicker.
-        """
-        if not self.use_gui or self.physics_client is None:
-            return
-
-        # Fixed position near the top-left of the origin in 3D world coordinates
-        x_pos, y_pos, z_start = -1.2, -1.2, 1.6
-        line_height = 0.10
-
-        for i, (label, value) in enumerate(stats.items()):
-            if isinstance(value, float):
-                display_text = f"{label}: {value:.3f}"
-            elif isinstance(value, (list, np.ndarray)):
-                formatted_vals = ", ".join(f"{v:.2f}" for v in np.atleast_1d(value))
-                display_text = f"{label}: [{formatted_vals}]"
-            else:
-                display_text = f"{label}: {value}"
-
-            color = [0, 0, 0]  # Black text for clear visibility against the light plane
-
-            if label in self.debug_text_ids:
-                p.addUserDebugText(
-                    display_text,
-                    [x_pos, y_pos, z_start - i * line_height],
-                    textColorRGB=color,
-                    textSize=1.1,
-                    replaceItemUniqueId=self.debug_text_ids[label],
-                )
-            else:
-                self.debug_text_ids[label] = p.addUserDebugText(
-                    display_text,
-                    [x_pos, y_pos, z_start - i * line_height],
-                    textColorRGB=color,
-                    textSize=1.1,
-                )
-
-    def reset_hud(self):
-        """Clears debug text tracking."""
-        self.debug_text_ids.clear()

ml/train.py

@@ -1,113 +0,0 @@
-import argparse
-import os
-from pathlib import Path
-
-from .env import JackBotEnv
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Train a joint-command policy for JackBot.")
-    parser.add_argument("--timesteps", type=int, default=500_000, help="Total training timesteps")
-    parser.add_argument("--model-path", type=str, default="ml/checkpoints/ppo_joint_command", help="Where to save the trained model")
-    parser.add_argument("--seed", type=int, default=0, help="Random seed")
-    parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto' to autodetect")
-    parser.add_argument("--use-gui", action="store_true", help="Enable PyBullet GUI during training")
-    parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the training environment")
-    parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
-    parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
-    return parser.parse_args()
-
-
-def train(
-    total_timesteps: int,
-    model_path: str,
-    seed: int = 0,
-    device: str = "auto",
-    use_gui: bool = False,
-    num_robots: int = 1,
-    robot_spacing: float = 0.5,
-    start_pose: str = "init_deg",
-):
-    try:
-        from stable_baselines3 import PPO
-        from stable_baselines3.common.vec_env import DummyVecEnv
-    except ImportError as exc:
-        raise ImportError(
-            "stable-baselines3 and gym are required for training. "
-            "Install them with: pip install stable-baselines3 gym"
-        ) from exc
-
-    env = DummyVecEnv([
-        lambda: JackBotEnv(
-            use_gui=use_gui,
-            random_command=True,
-            num_robots=num_robots,
-            robot_spacing=robot_spacing,
-            start_pose=start_pose,
-        )
-    ])
-
-    def resolve_device(requested_device: str) -> str:
-        try:
-            import torch
-        except ImportError:
-            if requested_device != "cpu":
-                raise RuntimeError(
-                    "PyTorch is not installed in the active environment. "
-                    "Install torch with a GPU-enabled build before using --device cuda."
-                )
-            return "cpu"
-
-        hip_supported = getattr(torch.version, "hip", None) is not None
-        cuda_available = torch.cuda.is_available()
-        hip_available = hip_supported and getattr(torch.backends, "hip", None) is not None and torch.backends.hip.is_available()
-
-        if requested_device == "auto":
-            if hip_available or cuda_available:
-                return "cuda"
-            return "cpu"
-
-        if requested_device in {"cuda", "gpu", "hip"}:
-            if hip_available or cuda_available:
-                return "cuda"
-            raise RuntimeError(
-                f"GPU device requested ({requested_device}) but no CUDA/ROCm-capable PyTorch is available. "
-                f"Installed torch build: {torch.__version__} (hip={getattr(torch.version, 'hip', None)}, cuda={cuda_available})"
-            )
-
-        if requested_device == "cpu":
-            return "cpu"
-
-        raise ValueError(
-            f"Unsupported device '{requested_device}'. Use 'cpu', 'cuda', or 'auto'."
-        )
-
-    device = resolve_device(device)
-
-    model = PPO(
-        "MlpPolicy",
-        env,
-        verbose=1,
-        seed=seed,
-        device=device,
-        tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"),
-    )
-    model.learn(total_timesteps=total_timesteps)
-
-    Path(model_path).parent.mkdir(parents=True, exist_ok=True)
-    model.save(model_path)
-    env.close()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    train(
-        args.timesteps,
-        args.model_path,
-        seed=args.seed,
-        device=args.device,
-        use_gui=args.use_gui,
-        num_robots=args.num_robots,
-        robot_spacing=args.robot_spacing,
-        start_pose=args.start_pose,
-    )

requirements.txt

@@ -4,11 +4,12 @@ ikpy
 pybullet
 pyserial
 matplotlib
+dearpygui
 
 # Deep learning / RL
 torch
 stable-baselines3
-gym
-gymnasium[box2d]
+stable-baselines3[extra]
+gymnasium
 shimmy
 tensorboard

simulation.py

@@ -1,68 +1,217 @@
-import pybullet as p
-import numpy as np
-import RobotState as rs
-import GlobalVariables as gv
-import config as cfg
-import DataTypes as dt
+"""
+simulation.py - PyBullet Simulation Interface & Physics Engine
+Consolidates scene management, physics queries, motor control, and rendering.
+"""
 import time
-import os
 import math
+from typing import List, Tuple, Optional, Union
+import numpy as np
+import pybullet as p
+import pybullet_data
+
+from config import cfg
+import DataTypes as dt
 
 
 class Simulation:
-    def __init__(self):
-        self.physics_client = p.connect(p.GUI)
-        self.robot = p.loadURDF(cfg.urdf_path, useFixedBase=True)
+    """Single PyBullet simulation manager providing getters/setters for JackBot."""
 
-        self.revolute_joints = [
-            i
-            for i in range(p.getNumJoints(self.robot))
-            if p.getJointInfo(self.robot, i)[2] == p.JOINT_REVOLUTE
-        ]
+    def __init__(self, urdf_path: str = cfg.urdf_path, use_gui: bool = True):
+        self.urdf_path = urdf_path
+        self.use_gui = use_gui
+        self.physics_client: Optional[int] = None
+        self.plane_id: Optional[int] = None
+        self.robot_id: Optional[int] = None
+        self.revolute_joints: List[int] = []
 
-        self.set_all_joints_to_90()
-        p.resetDebugVisualizerCamera(
-            cameraDistance=1.0,
-            cameraYaw=50,
-            cameraPitch=-35,
-            cameraTargetPosition=[0, 0, 0],
-        )
+        self.connect()
 
-    def set_all_joints_to_90(self):
-        for joint_index in range(p.getNumJoints(self.robot)):
-            joint_info = p.getJointInfo(self.robot, joint_index)
-            joint_type = joint_info[2]
-            if joint_type == p.JOINT_REVOLUTE:
-                p.resetJointState(self.robot, joint_index, math.radians(90))
+    def connect(self) -> None:
+        """Establishes connection to PyBullet GUI or DIRECT mode."""
+        if self.physics_client is not None and p.isConnected(self.physics_client):
+            return
 
-    def updatePos(self, current_rad: dt.RadArray):
-        radflat = current_rad.data.flatten()
-        for joint_index, target_angle in zip(self.revolute_joints, radflat):
-            p.setJointMotorControl2(
-                bodyIndex=self.robot,
-                jointIndex=joint_index,
-                controlMode=p.POSITION_CONTROL,
-                targetPosition=target_angle,
-                force=500,
+        flags = p.GUI if self.use_gui else p.DIRECT
+        self.physics_client = p.connect(flags)
+
+        p.setAdditionalSearchPath(pybullet_data.getDataPath())
+        p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client)
+
+        if self.use_gui:
+            p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.physics_client)
+            p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self.physics_client)
+            p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
+            p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
+            p.resetDebugVisualizerCamera(
+                cameraDistance=1.0, cameraYaw=50, cameraPitch=-35, cameraTargetPosition=[0, 0, 0],
+                physicsClientId=self.physics_client
             )
 
-    def step(self):
-        p.stepSimulation()
+    def load_scene(self, spawn_pos: Optional[List[float]] = None) -> Tuple[int, int, List[int]]:
+        """Loads plane and robot URDF, discovering revolute joint indices dynamically."""
+        if spawn_pos is None:
+            spawn_pos = [0.0, 0.0, 0.20]
 
-    def disconnect(self):
-        p.disconnect(self.physics_client)
+        self.plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
+        self.robot_id = p.loadURDF(self.urdf_path, spawn_pos, physicsClientId=self.physics_client)
+
+        self.revolute_joints = []
+        for j in range(p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)):
+            info = p.getJointInfo(self.robot_id, j, physicsClientId=self.physics_client)
+            if info[2] == p.JOINT_REVOLUTE:
+                self.revolute_joints.append(j)
+
+        return self.plane_id, self.robot_id, self.revolute_joints
+
+    # --- ACTUATION (SETTERS) ---
+
+    def set_robot_joint_angles(
+        self, target_angles: Union[np.ndarray, List[float], dt.RadArray]
+    ) -> None:
+        """Applies motor torque to pull joints toward target position angles."""
+        if isinstance(target_angles, dt.RadArray):
+            radflat = target_angles.data.flatten()
+        elif isinstance(target_angles, np.ndarray):
+            radflat = target_angles.flatten()
+        else:
+            radflat = target_angles
+
+        for joint_index, target_angle in zip(self.revolute_joints, radflat):
+            p.setJointMotorControl2(
+                bodyIndex=self.robot_id,
+                jointIndex=joint_index,
+                controlMode=p.POSITION_CONTROL,
+                targetPosition=float(target_angle),
+                force=30,
+                physicsClientId=self.physics_client
+            )
+            
+    def hard_reset_joint_angles(self, target_angles: Union[np.ndarray, dt.RadArray]) -> None:
+        """Instantly teleports joint angles to target positions, clearing velocity state."""
+        radflat = target_angles.data.flatten() if isinstance(target_angles, dt.RadArray) else target_angles.flatten()
+        for joint_index, target_angle in zip(self.revolute_joints, radflat):
+            p.resetJointState(
+                bodyUniqueId=self.robot_id,
+                jointIndex=joint_index,
+                targetValue=float(target_angle),
+                targetVelocity=0.0,
+                physicsClientId=self.physics_client
+            )
+
+    def reset_robot_base(
+        self,
+        pos: Optional[List[float]] = None,
+        orn: Optional[List[float]] = None,
+        linear_velocity: Optional[List[float]] = None,
+        angular_velocity: Optional[List[float]] = None
+    ) -> None:
+        """Resets root torso position, orientation quaternion, and clears base velocities."""
+        if pos is None:
+            pos = [0.0, 0.0, 0.20]
+        if orn is None:
+            orn = [0.0, 0.0, 0.0, 1.0]
+
+        lin_v = linear_velocity if linear_velocity is not None else [0.0, 0.0, 0.0]
+        ang_v = angular_velocity if angular_velocity is not None else [0.0, 0.0, 0.0]
+
+        p.resetBasePositionAndOrientation(self.robot_id, pos, orn, physicsClientId=self.physics_client)
+        p.resetBaseVelocity(self.robot_id, linearVelocity=lin_v, angularVelocity=ang_v, physicsClientId=self.physics_client)
+    
+    # --- TELEMETRY (GETTERS) ---
+
+    def get_robot_pose(self) -> Tuple[List[float], List[float]]:
+        pos, orn = p.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.physics_client)
+        return list(pos), list(orn)
+
+    def get_robot_rpy(self) -> Tuple[float, float, float]:
+        _, orn = p.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.physics_client)
+        roll, pitch, yaw = p.getEulerFromQuaternion(orn)
+        return float(roll), float(pitch), float(yaw)
+
+    def get_robot_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]:
+        pos, orn = p.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.physics_client)
+        roll, pitch, yaw = p.getEulerFromQuaternion(orn)
+        return list(pos), (float(roll), float(pitch), float(yaw))
+
+    def get_robot_velocity(self) -> Tuple[List[float], List[float]]:
+        lin_v, ang_v = p.getBaseVelocity(self.robot_id, physicsClientId=self.physics_client)
+        return list(lin_v), list(ang_v)
+
+    def get_robot_joint_angles(self) -> np.ndarray:
+        joint_states = p.getJointStates(self.robot_id, self.revolute_joints, physicsClientId=self.physics_client)
+        return np.array([state[0] for state in joint_states], dtype=np.float32)
+
+    def _get_urdf_joint_limits(self) -> Tuple[np.ndarray, np.ndarray]:
+        """Dynamically reads lower and upper limits for all revolute joints from PyBullet."""
+        lower_limits = []
+        upper_limits = []
+        
+        # Iterate through joints in PyBullet
+        for j_idx in range(p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)):
+            info = p.getJointInfo(self.robot_id, j_idx, physicsClientId=self.physics_client)
+            joint_type = info[2]
+            
+            # Only collect limits for revolute joints
+            if joint_type == p.JOINT_REVOLUTE:
+                lower_limits.append(info[8])  # Index 8 = jointLowerLimit
+                upper_limits.append(info[9])  # Index 9 = jointUpperLimit
+                
+        return np.array(lower_limits, dtype=np.float32), np.array(upper_limits, dtype=np.float32)
+
+    # --- SIMULATION LIFECYCLE CONTROLS ---
+
+    def step(self) -> None:
+        """Advances physics simulation by 1 time step."""
+        p.stepSimulation(physicsClientId=self.physics_client)
+
+    def settle_and_measure_height(
+        self, target_angles: Optional[np.ndarray] = None, steps: int = 100, fallback_height: float = 0.122
+    ) -> float:
+        """Settles the robot into the ground while actively holding target joint angles."""
+        for _ in range(steps):
+            if target_angles is not None:
+                self.set_robot_joint_angles(target_angles)
+            self.step()
+        pos, _ = self.get_robot_pose()
+        return pos[2] if pos[2] > 0.0 else fallback_height
+
+    def apply_external_force(
+        self, force: Union[List[float], np.ndarray], link_index: int = -1, position: Tuple[float, float, float] = (0.0, 0.0, 0.0)
+    ) -> None:
+        p.applyExternalForce(
+            objectUniqueId=self.robot_id,
+            linkIndex=link_index,
+            forceObj=list(force),
+            posObj=list(position),
+            flags=p.WORLD_FRAME,
+            physicsClientId=self.physics_client,
+        )
+
+    def disconnect(self) -> None:
+        if self.physics_client is not None and p.isConnected(self.physics_client):
+            p.disconnect(self.physics_client)
+            self.physics_client = None
+
+    def close(self) -> None:
+        self.disconnect()
 
 
 if __name__ == "__main__":
-    count = 0
-    current_rad: dt.RadArray = gv.init_deg.to_rad()
-    gv.shared_sim.updatePos(current_rad)
+    from Robot import Robot, PyBulletBackend
 
-    rs.walking()
-    while True:
-        count = +1
-        gv.shared_sim.step()
-        time.sleep(1 / 240)
-        if count > 50:
-            rs.walking()
-            count = 0
+    sim_instance = Simulation(use_gui=True)
+    sim_instance.load_scene()
+
+    backend = PyBulletBackend(sim_instance)
+    robot = Robot(backend_type=backend, mode="kinematics")
+    robot.reset_to_init()
+
+    # Forward velocity command
+    robot.vector_dirmov = [0.3, 0.0, 0.0]
+
+    try:
+        while True:
+            robot.tick()
+            time.sleep(1.0 / 60.0)
+    except KeyboardInterrupt:
+        sim_instance.disconnect()

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
+
+from config import 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.step_duration,
+        tickpersec: float = cfg.tick_rate_hz,
+    ):
+        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.step_duration,
+        tickpersec: float = cfg.tick_rate_hz,
+    ):
+        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
+
+from config import 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"]