8.8 KiB
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 (RobotBackendprotocol) 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 trainingNparallel 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
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)
python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip setuptools wheel
pip install -r requirements.txt
Windows (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:
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:
python ml/run_train.py --timesteps 100000 --model ml/checkpoints/ppo_joint_command
Train using multi-robot parallel vectorization with visual GUI enabled:
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_degorinit90_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:
python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 5 --gui
Multi-robot evaluation with custom stance pose:
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.
- Severe Tilt: Base roll or pitch orientation exceeds
- 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.