Readme Updated

README.md

@@ -1,21 +1,65 @@
+# JackBot — Hexapod Control, Simulation & RL Framework
 
-# JackBot — Hexapod Control & Simulation
+JackBot is a modular 3D hexapod robot control and machine learning framework built in Python. It supports real-time kinematics, multi-input options (GUI, gamepads), hardware streaming (ESP32 / Arduino), and vectorized Reinforcement Learning (PPO) using PyBullet and Gymnasium.
 
-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
+## Key Features
 
-- Python 3.12 is safest for `pygame` compatibility
-- Required Python packages:
-	- `numpy`
-	- `pygame`
-	- `ikpy`
-	- `pybullet`
-	- `pyserial`
-	- `matplotlib`
+* **Unified Robot Abstraction (`Robot.py`):** Virtual backends (`RobotBackend` protocol) allow seamless switching between 3D PyBullet simulation and physical hardware (ESP32 / Arduino) without changing high-level logic.
+* **Flexible Input Pipeline:** Pluggable input handlers supporting Pygame gamepad controllers, manual GUI sliders, or randomized direction vectors.
+* **Parallel Multi-Robot Training:** Vectorized Gymnasium environment (`JackBotEnv`) capable of simulating and training $N$ parallel hexapods simultaneously in PyBullet for PPO reinforcement learning.
+* **Live Telemetry & Visual Tracking:** Built-in PyBullet overlay features including real-time performance HUDs, floating leader crown tracking ($\text{👑}$) for top-reward robots, and visual failure feedback (failed robots turn semi-transparent dark gray).
 
-## Setup
+---
+
+## Project Architecture
+
+```text
+JackBot/
+├── main.py                  # Primary application entry point for manual & hardware control
+├── Robot.py                 # Core Robot class, kinematics wrapper, and Backend protocols
+├── config.py                # Global settings (backend selection, URDF path, communication specs)
+├── kinematics.py            # Forward and Inverse Kinematics (IKPy)
+├── simulation.py            # Base PyBullet GUI wrapper for single-robot interactive simulation
+├── robot_init.py            # Default stance angles and neutral leg positions
+├── DataTypes.py             # Strongly typed arrays (PosArray, RadArray, DegArray) & structs
+│
+├── states/                  # Finite State Machine (FSM) gait states
+│   ├── State.py             # Base State class
+│   ├── idle.py              # Neutral stance state
+│   └── walking.py           # Inverse-kinematics tripod gait state
+│
+├── inputs/                  # Input providers
+│   ├── InputProvider.py     # Base input abstraction
+│   └── PygameController.py  # Asynchronous gamepad loop (process-isolated)
+│
+├── gui/                     # Control interface
+│   └── MainWindow.py        # Pygame / parameter GUI layout and command resolver
+│
+├── EspCommunication.py      # WiFi socket sender for ESP32 hardware
+├── ArduinoCommunication.py  # Serial communication wrapper for Arduino hardware
+├── JackBotUrdf.urdf         # Kinematic 3D model definition (18 active joints)
+│
+└── ml/                      # Machine Learning Subsystem
+    ├── env.py               # JackBotEnv (Gymnasium multi-robot vector environment)
+    ├── SimManager.py        # Physics server initialization and scene loading
+    ├── MetricsOverlay.py    # PyBullet HUD (MetricsHUD) & leader crown tracking (LeaderCrown)
+    ├── run_train.py         # PPO training execution script
+    └── run_eval.py          # Model evaluation script
+```
+
+---
+
+## System Requirements
+
+* **Python 3.12** (Recommended)
+* **OS:** Linux (Ubuntu/Debian) or Windows 10/11
+* **Dependencies:** `pybullet`, `gymnasium`, `stable-baselines3`, `torch`, `numpy`, `pygame`, `ikpy`, `pyserial`, `matplotlib`
+
+---
+
+## Installation & Setup
 
 ### Linux (Bash)
 
@@ -26,6 +70,12 @@ python -m pip install --upgrade pip setuptools wheel
 pip install -r requirements.txt
 ```
 
+> **Linux X11 Headless Note:** If running PyBullet GUI on Linux gives an X11 server connection error (`cannot connect to X server`), ensure your display environment variable is set:
+> ```bash
+> export DISPLAY=:0
+> python main.py
+> ```
+
 ### Windows (PowerShell)
 
 ```powershell
@@ -35,75 +85,76 @@ python -m pip install --upgrade pip setuptools wheel
 pip install -r requirements.txt
 ```
 
-If PowerShell blocks activation:
-
+If PowerShell blocks script execution:
 ```powershell
 Set-ExecutionPolicy -Scope Process -ExecutionPolicy RemoteSigned
 .\.venv\Scripts\Activate.ps1
 ```
 
-## Run
+---
 
-From the repository root:
+## Usage Guide
 
-```
+### 1. Manual Control & Hardware Streaming (`main.py`)
+
+`main.py` is the operational entry point for driving the robot manually via GUI sliders or a gamepad, running either in 3D PyBullet simulation or connected to physical hardware.
+
+To launch:
+```bash
 python main.py
 ```
 
-If Linux breaks with 
-```
-ExampleBrowserThreadFunc started
-X11 functions dynamically loaded using dlopen/dlsym OK!
+#### Configuration (`config.py`)
+Edit `config.py` prior to launching `main.py` to configure execution mode and connections:
 
-        cannot connect to X server
-```
-run:
-```
-export DISPLAY=:0
-python main.py
+* **Backend Selection (`cfg.backend`):**
+  * `BackendType.SIMULATION`: Executes motion inside a 3D PyBullet window.
+  * `BackendType.ESP32`: Streams target joint angles over WiFi sockets to an ESP32 micro-controller (`cfg.esp32_ip`, `cfg.esp32_port`).
+  * `BackendType.ARDUINO`: Streams target joint angles over Serial to an Arduino (`cfg.port`, `cfg.baudrate`).
+
+---
+
+### 2. Machine Learning: PPO Training & Evaluation (`ml/`)
+
+The `ml/` directory contains tools to train RL policies using **Proximal Policy Optimization (PPO)**. The agent receives observations ($18\text{ joint angles} + 4\text{ velocity/turning commands}$) and outputs continuous joint delta actions in $[-1, 1]$.
+
+#### A. Training a Model (`run_train.py`)
+
+Train a single robot policy:
+```bash
+python ml/run_train.py --timesteps 100000 --model ml/checkpoints/ppo_joint_command
 ```
 
-## Configuration
+Train using multi-robot parallel vectorization with visual GUI enabled:
+```bash
+python ml/run_train.py --timesteps 500000 --model ml/checkpoints/ppo_joint_command --num-robots 16 --robot-spacing 0.75 --gui
+```
 
-Edit `config.py` before running:
+**Training CLI Arguments:**
+* `--timesteps`: Total training timesteps.
+* `--num-robots`: Number of parallel robot instances spawned in a grid layout (e.g., 16 to 64).
+* `--robot-spacing`: Distance in meters between robot spawn origins.
+* `--start-pose`: Stance pose at environment reset (`init_deg` or `init90_deg`).
+* `--gui`: Renders the live PyBullet GUI with metrics HUD, leader crown, and failure graying.
 
-- `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"`
+#### B. Evaluating a Model (`run_eval.py`)
 
-## Project structure
+Run an evaluation loop using a saved model checkpoint:
+```bash
+python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 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
+Multi-robot evaluation with custom stance pose:
+```bash
+python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3 --gui --num-robots 4 --robot-spacing 0.8 --start-pose init_deg
+```
 
-## Notes
+---
 
-- `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.
+## Environment Mechanics & Telemetry (`ml/env.py`)
 
-## Troubleshooting
+When training with `--gui`, `JackBotEnv` includes dynamic visual feedback mechanisms:
 
-- 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.
-
-## Suggested improvements
-
-- Add a `requirements.txt` or `pyproject.toml`.
-- Add a `LICENSE` file before sharing the project.
-- Document hardware wiring and packet formats for ESP32/Arduino.
+* **Failure Detection & Graying:** Robots are continuously evaluated for roll/pitch tilt ($> 0.7\text{ rad}$) or base collapse ($< 0.05\text{ m}$ height). When a robot fails, its state mask is flagged and its 3D mesh automatically turns **semi-transparent dark gray**.
+* **Leader Crown ($\text{👑}$):** A floating crown indicator tracks and sits directly above the robot currently achieving the highest cumulative reward in the multi-robot grid.
+* **Termination Threshold:** The environment episode terminates automatically when the percentage of failed robots exceeds the configured threshold (default: $30\%$).

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/train.py

@@ -13,7 +13,7 @@ def parse_args():
     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("--robot-spacing", type=float, default=3.0, 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()