# JackBot — Hexapod Control, Simulation & RL Framework

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.

---

## Key Features

* **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).

---

## 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)

```bash
python3.12 -m venv .venv
source .venv/bin/activate
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
python3.12 -m venv .venv
.\.venv\Scripts\Activate.ps1
python -m pip install --upgrade pip setuptools wheel
pip install -r requirements.txt
```

If PowerShell blocks script execution:
```powershell
Set-ExecutionPolicy -Scope Process -ExecutionPolicy RemoteSigned
.\.venv\Scripts\Activate.ps1
```

---

## 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
```

#### Configuration (`config.py`)
Edit `config.py` prior to launching `main.py` to configure execution mode and connections:

* **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
```

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
```

**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.

#### B. Evaluating a Model (`run_eval.py`)

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
```

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
```

---

## Environment Mechanics & Telemetry (`ml/env.py`)

When training with `--gui`, `JackBotEnv` includes dynamic visual feedback mechanisms:

* **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\%$).