JackBot/README.md

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

---

## What JackBot Is

JackBot is a Python-based hexapod robot project that combines:

* a robot control stack for a six-legged walking robot,
* a physics simulator (`simulation.py`) for testing in software before using real hardware, and
* a reinforcement learning pipeline that teaches the robot how to move through trial and error or pretraining from kinematics.

The project is a full control system that can:

* run the robot in a PyBullet simulation,
* accept human commands from a GUI or gamepad,
* stream target joint positions to physical hardware (ESP32 / Arduino), and
* train and evaluate policy checkpoints using PPO and Behavioral Cloning (BC).

## Why the Project Exists

A hexapod is hard to control manually because each leg has several joints and the robot has to maintain balance while moving. Instead of hard-coding every motion rule, this project uses the robot model and physics simulation as a testbed to learn or refine walking behavior.

In practice, the RL workflow looks like this:

1. The robot starts from a standing pose.
2. The policy receives sensory information (joint angles + command vector) from the simulation.
3. The policy chooses continuous target joint positions.
4. The simulation steps forward and computes physics.
5. The policy receives shaped rewards for surviving, following commands, and remaining stable.
6. Over many iterations, PPO improves the locomotion policy.

---

## Key Components

* **Robot abstraction (`Robot.py`)**: Wraps the robot body, state machine execution, joint control API, and backend selection.
* **Kinematics (`kinematics.py`)**: Turns target leg positions into actual joint angles through inverse kinematics.
* **Simulation (`simulation.py`)**: Consolidates PyBullet scene management, motor positioning, physics queries, and environment stepping into a single manager.
* **Inputs (`inputs/`)**: Interface for receiving control commands from a GUI, gamepad, or procedural random generators.
* **ML Subsystem (`ml/`)**: Contains Gymnasium environments (`env.py`), custom callbacks (`callbacks.py`), behavioral cloning pretraining (`pretrain_bc.py`), PPO training (`run_train.py`), and multi-phase evaluation scripts (`run_eval.py`, `run_eval_training.py`).

---

## Project Architecture

```text
JackBot/
├── main.py                  # Manual control, GUI, and hardware streaming entry point
├── Robot.py                 # Robot wrapper, gait execution, and backend abstraction
├── config.py                # Global settings and hardware/backend configuration
├── kinematics.py            # Inverse kinematics mapping cartesian targets to joint angles
├── simulation.py            # PyBullet simulation manager (handles engine, URDF loading, & physics)
├── robot_init.py            # Neutral standing pose and initial joint definitions
├── DataTypes.py             # Typed numpy structures for positions and joint angles
│
├── states/                  # Gait and state-machine behavior logic
│   ├── State.py
│   ├── IdleState.py
│   ├── WalkingState.py
│   └── WaveState.py
│
├── inputs/                  # Command input sources
│   ├── InputProvider.py
│   ├── PygameController.py
│   └── RandomeInputProvider.py
│
├── gui/                     # Visual control frontend (Tkinter/CustomTkinter)
│   └── MainWindow.py
│
├── EspCommunication.py      # ESP32 socket communication layer
├── ArduinoCommunication.py  # Arduino serial communication layer
├── JackBotUrdf.urdf         # Robot URDF mesh and joint axis definition
│
└── ml/                      # Reinforcement Learning Subsystem
    ├── env.py               # Gymnasium environment wrapping PyBullet & Robot control
    ├── callbacks.py         # Custom SB3 callbacks for TensorBoard logging & Curriculum progression
    ├── pretrain_bc.py       # Behavioral Cloning (BC) script to pre-train policy from IK teacher
    ├── run_train.py         # Main training entry point (runs parallelized PPO via SubprocVecEnv)
    ├── run_eval.py          # Phase-by-phase policy evaluator with CLI reports
    ├── run_eval_training.py # Benchmark reward system across phases using pure kinematics
    └── MetricsOverlay.py    # PyBullet HUD debug text and visual overlays
```

---

## System Requirements

* **Python 3.12** (Recommended)
* **OS:** Windows 10/11 or Linux
* **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
```

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

---

### Usage Guide

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

---

### Machine Learning Pipeline (`ml/`)

The machine learning system trains the hexapod to move using a combination of Behavioral Cloning (BC) (optional) and Proximal Policy Optimization (PPO) driven by a multi-phase curriculum.

### ML Workflow Steps

The policy input contains:

* 18 joint angles from the robot, and
* 4 command dimensions describing the desired motion direction and yaw rate.

The policy output is:

* 18 continuous values, one per joint, describing how much each joint should change.

This means the learning method is not controlling a discrete gait state directly. Instead, it learns a continuous control policy for the whole body.

#### Methods implemented

The current ML stack uses:

* **Gymnasium** as the environment API.
* **PyBullet** as the physics engine.
* **Stable-Baselines3 PPO** as the learning algorithm.
* **Curriculum learning** to gradually unlock harder command regimes.

The key ideas are:

1. **Standing first**: the robot must stay upright and stable.
2. **Forward motion second**: the robot must learn to move in a commanded direction.
3. **Turning next**: after forward command-following is stable, the policy gets a yaw/turning challenge.
4. **Omni-direction later**: the full command space is gradually exposed only once the simpler skills are reliable.

This staged approach is important because learning to balance and walk all at once is difficult for a real hexapod.

#### Training flow

Training is launched through `ml/run_train.py` and hands off to `ml/train.py`.

Typical training starts as:
```bash
python ml/run_train.py --timesteps 100000 --model ml/checkpoints/ppo_joint_command
```

Training uses the simulator as the environment, runs PPO updates, and periodically saves checkpoint models.

#### Evaluation flow

Evaluation uses `ml/run_eval.py` and `ml/evaluate.py`.

A saved PPO model is loaded and run in deterministic evaluation mode. This is used to answer questions like:

* Does the policy actually move?
* Does it survive longer than before?
* Does it remain stable under the commanded motion?

Example:
```bash
python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 5 --gui
```

---

## What the Reward Is Trying to Teach

The reward function is not a single number with one purpose. It combines several terms so the policy learns to:

* survive in the environment,
* move in the commanded direction,
* avoid unwanted sideways drift,
* stay upright without excessive tilt,
* avoid giant control jumps,
* keep the robot away from a low collapsed posture.

The main ideas are:

* a small alive bonus for staying upright,
* directional movement rewards,
* penalties for drifting or standing still when a command is active,
* penalties for excessive tilt or collapse,
* a penalty for large abrupt control changes.

This style of reward shaping encourages the policy to learn real locomotion rather than just freezing in place.

### Curriculum / phase progression

The environment uses a staged curriculum to expose command complexity gradually.

* **Phase 0**: basic standing / forward-focused behavior
* **Phase 1**: turning and directional regularization
* **Phase 2**: omni-directional movement
* **Phase 3**: full-command challenge

The goal is to prevent the learner from being asked to master all difficult motion goals at once.

---

## How the Control Loop Works in Practice

A simple, human-readable pipeline is:

1. `main.py` or the training/evaluation wrapper creates the robot environment.
2. The simulation runs the robot in PyBullet.
3. The policy observes the latest joint states and command vectors.
4. PPO predicts a new action.
5. The action is applied to the robot joints.
6. The simulator steps forward one frame.
7. The reward is computed from the new state.
8. The policy improves based on that reward.

That loop is the essence of the machine learning part of the project.

---

## What Makes this Project Useful

This repository is useful because it combines several layers that are often separate:

* robot control and kinematics,
* physical simulation,
* command input sources,
* RL environment construction,
* PPO training and evaluation.

For a beginner, the easiest way to think about the repo is:

- `main.py` is for direct manual control,
- `ml/env.py` is the simulator-to-policy interface,
- `ml/train.py` is where training happens,
- `ml/evaluate.py` is where you check whether the learned policy is actually good.

---

## Reward Function & Termination Mechanics (`ml/env.py`)

### 1. Reward Function Formulation

The per-robot step reward ($R_{\text{step}}$) is shaped to reward tracked motion and survival while strongly discouraging stillness once a command is present. The current implementation uses the following structure:

$$R_{\text{step}} = R_{\text{alive}} + R_{\text{tracking}} + R_{\text{rotation}} - P_{\text{drift}} - P_{\text{still}}(t) - P_{\text{height}} - P_{\text{stability}} - P_{\text{large-delta}}$$

* **Alive Bonus ($R_{\text{alive}} \approx +0.02$):** A small positive baseline awarded every timestep the robot remains upright.
* **Velocity Tracking ($R_{\text{tracking}}$):** Rewards linear movement aligned with the commanded direction vector.
* **Yaw Rotation Tracking ($R_{\text{rotation}}$):** Rewards turning in the commanded yaw direction.
* **Drift Penalty ($P_{\text{drift}}$):** Penalizes lateral motion that does not align with the commanded direction.
* **Age-Ramped Stillness Penalty ($P_{\text{still}}(t)$):** Penalizes remaining stationary when a command is active; the penalty ramps up as the episode gets older so the policy cannot settle into a frozen local optimum.
* **Height Penalty ($P_{\text{height}}$):** Penalizes falling below the measured post-settle standing height.
* **Body Stability Penalty ($P_{\text{stability}}$):** Penalizes roll and pitch tilt.
* **Large-Delta Penalty ($P_{\text{large-delta}}$):** Penalizes only large joint-control jumps, not normal continuous command usage, so the robot is free to use its joints continuously.

---

### 2. Failure Detection & Termination Logic

In multi-robot vectorized training (`JackBotEnv`), individual robot failures are handled independently to allow maximum simulation efficiency:

* **Individual Failure Masking:** A robot is flagged as failed (`failed_robots_mask[idx] = True`) if either condition is met:
  * **Severe Tilt:** Base roll or pitch orientation exceeds $0.9\text{ radians}$.
  * **Base Collapse:** Base height drops below a relative threshold derived from the measured settled standing height.
* **Visual Failure Feedback:** When a robot fails during GUI execution (`--gui`), its 3D URDF mesh immediately updates to a **semi-transparent dark gray** visual state (`COLOR_FAILED = [0.3, 0.3, 0.3, 0.6]`) to distinguish it from active learners.
* **Environment Termination (`terminated=True`):** The environment as a whole is marked as terminated when the failure ratio reaches the configured threshold.
* **Environment Truncation (`truncated=True`):** Occurs when the episode reaches the maximum allowable step budget (`max_episode_steps`).

---

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

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

* **Metrics HUD:** A live on-screen text overlay tracking active episode count, total step rate (FPS), average base height, roll/pitch angles, and cumulative per-robot rewards.
* **Leader Crown ($\text{👑}$):** A floating crown debug indicator tracks and positions itself directly above the base of whichever robot is achieving the highest cumulative reward in the multi-robot grid.