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 for testing in software before using real hardware, and
• a reinforcement learning pipeline that teaches the robot how to move through trial and error.

The project is not just one script or one model. It 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, and
• train a policy using PPO so the robot can learn locomotion automatically.

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 walking behavior.

In practice, the workflow looks like this:

1. The robot starts from a standing pose.
2. The policy receives sensory information from the simulation.
3. The policy chooses new joint target motions.
4. The simulation updates the physics.
5. The policy receives a reward for surviving and moving in the right direction.
6. Over many iterations, PPO improves the controller.

---

Key Components

• Robot abstraction (Robot.py): wraps the robot body, robot state, and joint control API.
• Kinematics (kinematics.py): turns target leg poses into actual joint angles through inverse kinematics.
• Simulation (simulation.py): creates and updates the PyBullet world where the robot can be tested safely.
• Inputs (inputs/): lets the robot be controlled from either a GUI, a gamepad, or a generated target command.
• ML environment (ml/env.py): provides the Gymnasium environment that the policy interacts with.
• PPO training (ml/train.py): trains a policy using stable-baselines3.
• Evaluation (ml/evaluate.py): loads a saved model and runs policy rollouts to inspect performance.

---

Project Architecture

JackBot/
├── main.py                  # manual control and hardware streaming entry point
├── Robot.py                 # robot wrapper and backend abstraction
├── config.py                # global settings and communication configuration
├── kinematics.py            # inverse kinematics used to map motion commands to joints
├── simulation.py            # PyBullet simulation shell for the robot
├── robot_init.py            # neutral standing pose and initial joint values
├── DataTypes.py             # typed data structures for positions and joint angles
│
├── states/                  # gait/state-machine behavior logic
│   ├── State.py
│   ├── IdleState.py
│   ├── WalkingState.py
│   └── WaveState.py
│
├── inputs/                  # command sources for the robot
│   ├── InputProvider.py
│   ├── PygameController.py
│   └── RandomeInputProvider.py
│
├── gui/                     # visual control frontend
│   └── MainWindow.py
│
├── EspCommunication.py      # ESP32 communication layer
├── ArduinoCommunication.py  # Arduino serial communication layer
├── JackBotUrdf.urdf         # robot mesh and joint definition
│
└── ml/                      # reinforcement learning subsystem
    ├── env.py               # Gymnasium environment used by PPO
    ├── SimManager.py        # PyBullet scene setup and stepping
    ├── MetricsOverlay.py    # HUD and visual overlays
    ├── train.py             # PPO training entry point
    ├── run_train.py         # command-line training launcher
    ├── evaluate.py          # policy evaluation loop
    └── run_eval.py          # CLI wrapper for evaluation

---

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. How the Machine Learning System Works

The machine learning subsystem teaches the robot to move by interacting with a simulation environment instead of relying on a manually written gait controller.

At a high level:

• The policy receives a vector of observations from the simulator.
• The policy outputs a set of continuous joint deltas.
• Those outputs are applied to the robot in the physics simulation.
• The environment computes a reward based on survival, movement, stability, and command-following quality.
• PPO updates the policy over many training iterations to improve the controller.

What the policy sees and does

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:

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:

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.