176 lines
8.8 KiB
Markdown
176 lines
8.8 KiB
Markdown
# 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
|
|
```
|
|
|
|
### 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
|
|
|
|
### 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
|
|
```
|
|
|
|
---
|
|
|
|
## Reward Function & Termination Mechanics (`ml/env.py`)
|
|
|
|
### 1. Reward Function Formulation
|
|
|
|
The per-robot step reward ($R_{\text{step}}$) incentivizes tracking directional velocity commands while maintaining body stability and smooth joint actuation. The total environment step reward is the sum of all individual active robot rewards:
|
|
|
|
$$R_{\text{step}} = R_{\text{alive}} + R_{\text{tracking}} + R_{\text{rotation}} - P_{\text{stability}} - P_{\text{action}}$$
|
|
|
|
* **Alive Bonus ($R_{\text{alive}} = +0.1$):** A constant positive baseline awarded every timestep the robot remains upright.
|
|
* **Velocity Tracking ($R_{\text{tracking}} = v_{x,\text{cmd}} \cdot v_x + v_{y,\text{cmd}} \cdot v_y$):** Rewards linear movement in the target command direction ($v_x, v_y$).
|
|
* **Yaw Rotation Tracking ($R_{\text{rotation}} = \omega_{\text{cmd}} \cdot \omega_z$):** Rewards turning along the vertical yaw axis according to angular command $\omega_{\text{cmd}}$.
|
|
* **Body Stability Penalty ($P_{\text{stability}} = 0.2 \cdot (|\text{roll}| + |\text{pitch}|)$):** Penalizes tilting away from a level horizontal posture.
|
|
* **Action Energy Penalty ($P_{\text{action}} = 0.01 \cdot \sum a_i^2$):** Penalizes excessive joint delta actions to encourage smooth, energy-efficient leg movements and reduce jitter.
|
|
|
|
---
|
|
|
|
### 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.7\text{ radians}$ ($\approx 40^\circ$).
|
|
* **Base Collapse:** Base height drops below $0.05\text{ meters}$ above the ground plane.
|
|
* **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 entire environment step resets when any Robot fails.
|
|
* **Environment Truncation (`truncated=True`):** Occurs when the episode reaches the maximum allowable step budget (`max_episode_steps = 3000`).
|
|
|
|
---
|
|
|
|
## 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. |