300 lines
13 KiB
Markdown
300 lines
13 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.
|
|
|
|
---
|
|
|
|
## 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
|
|
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
|
|
```
|
|
|
|
---
|
|
|
|
### 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 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.
|
|
#### 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. |