JackBot is a modular Python framework for controlling a six-legged robot in simulation or on hardware. The current workspace reflects a matured control stack with a PyBullet-based simulator, a GUI-driven manual runtime, and a reinforcement-learning pipeline for PPO training and evaluation.

What JackBot Is

JackBot is a Python-based hexapod 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,
a reinforcement learning pipeline that trains locomotion policies with PPO and behavioral cloning,
and a GUI/gamepad input layer for manual operation.

The project can currently:

run the robot in a PyBullet simulation,
accept commands from a GUI or gamepad,
stream target joint positions to physical hardware (ESP32 / Arduino),
train PPO models and evaluate saved checkpoints,
generate kinematics-based teacher data for behavioral cloning.

Why the Project Exists

A hexapod is hard to control manually because each leg has multiple joints and the robot must maintain balance while moving. The project uses a robot model and physics simulation as a testbed for learning motion strategies, refining kinematic control, and validating behavior before applying commands to real hardware.

In the current codebase, the practical workflow is:

Start the robot in either simulation or hardware mode.
Feed motion commands from GUI/gamepad or a training environment.
Convert target foot positions into joint targets using inverse kinematics.
Apply these targets to the active backend.
Train or evaluate PPO policies against the simulated robot.

Key Components

Robot abstraction (Robot.py): Central handler for the active backend, current joint state, IK/FK helpers, gait handling, and RL action application.
Kinematics (kinematics.py): Converts target foot positions into joint angles using IKPy chains for each leg.
Simulation (simulation.py): Manages PyBullet scene loading, stepping, joint actuation, base pose queries, and physics telemetry.
Inputs (inputs/): Receives commands from GUI sliders, Pygame gamepads, and random command generation.
ML Subsystem (ml/): Contains the Gymnasium environment (env.py), PPO training (run_train.py), evaluation (run_eval.py), BC pretraining (pretrain_bc.py), callbacks, and metrics overlays.
States (states/): Contains additional state-machine classes such as IdleState, WalkingState, and WaveState. These exist in the repository, but the current active runtime in main.py does not currently drive them through Robot.tick().

Project Architecture

JackBot/
├── main.py                  # Manual runtime entry point
├── Robot.py                 # Unified robot wrapper + backend selection
├── config.py                # Central runtime configuration
├── kinematics.py            # IK/FK helpers
├── simulation.py            # PyBullet scene / physics manager
├── robot_init.py            # Initial joint definitions and center points
├── DataTypes.py             # Typed arrays for positions / angles
├── JackBotUrdf.urdf        # Robot URDF
│
├── states/                  # State classes present in the project
│   ├── State.py
│   ├── IdleState.py
│   ├── WalkingState.py
│   ├── WaveState.py
│   ├── ml_walking.py
│   └── __init__.py
│
├── inputs/                  # Command input providers
│   ├── InputProvider.py
│   ├── PygameController.py
│   ├── RandomeInputProvider.py
│   └── RandomInputProvider.py
│
├── gui/                     # GUI/dashboard control layer
│   └── MainWindow.py
│
├── ml/                      # Gymnasium RL training + evaluation stack
│   ├── env.py               # RL environment + reward logic
│   ├── callbacks.py         # SB3 callbacks for logs/curriculum
│   ├── pretrain_bc.py       # Behavior cloning pretraining
│   ├── run_train.py         # PPO training entry point
│   ├── run_eval.py          # Evaluation of saved checkpoints
│   ├── run_eval_training.py # Kinematics-mode benchmark script
│   ├── MetricsOverlay.py    # 3D in-scene metrics HUD
│   └── checkpoints/        # Model checkpoints + saved runs
│
├── EspCommunication.py      # ESP32 UDP communication layer
├── ArduinoCommunication.py  # Arduino serial communication layer
├── Helper Scripts/          # Utility scripts
│   ├── FindCenterPoints.py
│   └── torqueCalc.py
│
├── requirements.txt
├── README.md
└── ...

Important current-state note

The repository contains a few pieces that are still present but not fully connected to the current runtime:

states/WalkingState.py, states/WaveState.py, and states/ml_walking.py exist, but main.py currently drives Robot.tick() directly instead of routing through the active STATE_REGISTRY.
inputs/RandomeInputProvider.py exists, but the active runtime does not currently use it directly. The GUI resolves random walking behavior in gui/MainWindow.py.
WaveEmoteState / LaolaWaveEmoteState are defined in states/WaveState.py, but they are not currently registered in states/__init__.py.

System Requirements

Python: 3.12 recommended
OS: Windows 10/11 or Linux
Dependencies: pybullet, gymnasium, stable-baselines3, torch, numpy, pygame, ikpy, pyserial, matplotlib, dearpygui

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)

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 main operational entry point for driving the robot manually through the GUI or a connected gamepad.

To launch:

python main.py

Configuration (`config.py`)

Before launching main.py, edit config.py to select the backend and connection targets.

Backend Selection (cfg.backend):
- BackendType.SIMULATION: run inside a PyBullet window
- BackendType.ESP32: stream motion over UDP to an ESP32
- BackendType.ARDUINO: stream motion over serial to an Arduino

Supported configuration values in the current code include:

backend
urdf_path
port, baudrate
esp32_ip, esp32_port
tick_rate_hz, step_duration
step_height, step_length

GUI input sources

The current GUI (gui/MainWindow.py) exposes these input sources:

Gamepad
GUI Sliders
Random Walk

The active main.py runtime resolves commands through resolve_active_command(...) and then passes the resulting vector_dirmov directly to Robot.tick().

Machine Learning Pipeline (`ml/`)

The ML subsystem currently uses:

Gymnasium as the environment API
PyBullet as the physics engine
Stable-Baselines3 PPO as the learning algorithm
Curriculum phases to gradually expose more difficult command spaces

Current ML scripts

ml/run_train.py — PPO training entry point
ml/run_eval.py — phase-based model evaluation
ml/run_eval_training.py — reward benchmark for kinematics mode
ml/pretrain_bc.py — collect teacher data from the kinematics solver and pretrain a base policy
ml/env.py — environment definition and reward logic
ml/callbacks.py — reward logging and curriculum callbacks

Training flow

Training is launched with:

python ml/run_train.py --total-timesteps 1500000 --gui

The current training script creates a vectorized PPO environment, optionally loads a pre-trained checkpoint, and saves results into ml/checkpoints/ and ml/logs/.

Evaluation flow

Evaluation is launched with:

python ml/run_eval.py --model ml/checkpoints/jackbot_kinematics_base.zip --episodes-per-phase 5 --gui

This script runs a multi-phase deterministic evaluation suite with fixed command vectors for:

FORWARD
TURN_AND_DIRECTION
OMNI_DIRECTION
FULL_COMMAND

Behavioral cloning pretraining

The repository also contains a behavioral cloning route:

python ml/pretrain_bc.py --num-samples 100000 --epochs 15 --save-path ml/checkpoints/jackbot_kinematics_base.zip

This script gathers (observation, action) data from the kinematics teacher and trains a PPO policy to act as a learned base model.

What the Reward Is Trying to Teach

The reward function in ml/env.py is designed to teach several things at once:

survival and stability,
following the commanded direction,
avoiding lateral drift,
maintaining base height,
reducing abrupt control changes,
staying close to a stable stance when the command is zero.

The environment uses several reward components, including:

height reward,
stability reward,
pose closeness reward,
smoothness reward,
linear velocity tracking,
angular velocity tracking,
jitter penalty,
stand-by penalty when the command is zero.

Detailed reward structure

The reward is built step by step inside JackBotEnv._compute_reward() in ml/env.py. It is not a single binary success signal; it is a dense shaping signal that rewards good behavior continuously throughout the episode.

At each control step, the environment measures:

current body height,
roll and pitch angles,
current joint configuration,
measured linear and angular velocities,
the active command vector [vx, vy, omega],
the difference between the current action and the previous actions.

From these values, it computes a set of sub-rewards:

height reward: a Gaussian-style reward based on how close the robot body is to the target height.
stability reward: a reward for keeping roll/pitch low and the body well balanced.
pose closeness reward: rewards staying near the default standing joint posture.
smoothness reward: rewards actions that change gradually instead of abruptly.
linear velocity reward: encourages the robot to move in the commanded direction and speed.
angular velocity reward: rewards matching the commanded turning rate.

The environment then adds two kinds of correction terms:

jitter penalty: subtracts a small amount when action changes are noisy or jerky.
stand_penalty: when the command is effectively zero, penalizes unintended movement and yaw drift.

This makes the reward function behave as a soft guidance system: the agent gets a steady gradient that says “this is closer to what we want” or “this is worse than the desired behavior.”

Standing mode vs. walking mode

The reward code handles two cases differently:

1. Standing mode

When the command vector is close to zero (cmd_norm < 0.05 and abs(cmd_yaw) < 0.05), the agent is not supposed to move much. In that case the reward focuses on:

keeping the body at the correct height,
staying stable,
maintaining a clean posture,
staying smooth.

A small stand_penalty is then applied to discourage unintended speed and yaw drift while the robot is supposed to hold position.

2. Walking mode

When a non-zero command is active, the robot is rewarded for moving in the intended direction and turning at the requested rate. The reward then emphasizes:

matching the commanded linear velocity,
matching the commanded yaw rate,
continuing to maintain height and stability,
staying smooth in its control inputs.

If the command expects movement but the robot is effectively still, the code now applies a stronger stillness_penalty instead of a neutral reward. This means standing while the command says to move is explicitly discouraged.

Why the reward is shaped this way

The goal is not just to teach the robot to stay alive. The reward is designed so that a PPO agent learns multiple useful habits at once:

do not collapse or tip over,
keep the body at a sensible height,
follow directional commands,
avoid unstable oscillations,
avoid overreactive control jumps,
stay near a normal standing pose when no motion is requested.

This is why the environment is not based on a single sparse reward such as “+1 for success, 0 otherwise.” Instead, it uses dense reward shaping so the policy receives useful feedback on every step.

Are the penalties real penalties?

Yes — in the reward function they are real negative contributions. For example:

jitter_penalty subtracts from the step reward when action changes are too abrupt,
stand_penalty subtracts when the robot moves unnecessarily while standing,
stillness_penalty subtracts when movement is commanded but the robot remains essentially frozen.

In the current implementation, a non-zero command that is not followed by meaningful motion now yields an explicit penalty instead of a neutral reward. This is the main behavior change requested for the training setup: standing while the command says to move is now actively discouraged.

The code does this:

final_reward = step_reward + jitter_penalty + alive_bonus

That means:

negative penalty terms can now reduce the reward below zero,
the reward is no longer clipped to zero in this path,
and the episode is still only ended by the separate failure check in _update_robot_failure().

So the answer is:

the penalties are real reward penalties,
they are now strong enough to discourage command-mismatch behavior,
and the actual terminal condition remains the failure check in _update_robot_failure().

What actually ends an episode?

The episode ends when the robot is considered failed, not when a reward penalty is applied. In ml/env.py, the environment marks the robot as failed if:

it is too tilted (roll or pitch exceed the configured failure threshold), or
it has collapsed below a minimum body-height threshold.

That is a hard termination condition. In other words:

reward penalties discourage bad behavior,
failure conditions stop the episode when the robot is clearly unstable or collapsed.

Practical interpretation

A good mental model is:

the reward function teaches the robot what “good locomotion” looks like,
the failure check prevents the robot from continuing when it is physically broken or unstable,
and the curriculum gradually increases the difficulty of the commands as the robot becomes more capable.

This combination is a common reinforcement-learning setup for locomotion: dense rewards shape the desired behavior, while hard failure conditions protect the training process from degenerate states.

This reward shaping encourages the robot to learn locomotion patterns rather than simply freezing in place.

Curriculum / Phase Progression

The environment currently uses the following curriculum stages:

STAND_ONLY
FORWARD
TURN_AND_DIRECTION
OMNI_DIRECTION
FULL_COMMAND

The training environment gradually advances phase complexity based on survival, stability, and movement metrics. ml/callbacks.py includes custom curriculum-handling logic for the online learning loop.

How the Control Loop Works in Practice

A simplified view of the current runtime is:

main.py starts the joystick controller process and opens the GUI.
The GUI resolves active commands (Gamepad, slider, or random walk).
Robot.tick() receives the current motion vector and applies it through the active backend.
In simulation mode, Simulation.step() advances the PyBullet world.
In training/evaluation mode, JackBotEnv.step() computes reward, updates curriculum, and returns observations.
PPO uses the observation/action loop to improve the policy.

What Makes this Project Useful

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

robot control and kinematics,
physics simulation,
command input sources,
RL environment construction,
PPO training and evaluation,
hardware communication layers.

For a newcomer, the easiest way to think about the project is:

main.py is the manual control entry point,
Robot.py is the core robot wrapper,
simulation.py is the physics layer,
ml/env.py is the environment interface for RL,
ml/run_train.py, ml/run_eval.py, and ml/pretrain_bc.py are the main ML workflows.

Current runtime notes and caveats

State system status

The state classes are present in the repository, but the current runtime path is not using them as the main control loop:

Robot.tick() currently updates the robot directly from vector_dirmov
the STATE_REGISTRY exists, but it is not the path used by the current main.py execution flow
the state subsystem remains partially implemented and should be treated as a legacy or optional extension

Input source status

The repository currently includes both:

inputs/PygameController.py for gamepad input
gui/MainWindow.py for selecting Gamepad, GUI Sliders, and Random Walk

The standalone random input provider file is present, but it is not the path currently used by main.py.

Hardware communication status

The hardware communication classes still exist:

EspCommunication.py for ESP32 UDP
ArduinoCommunication.py for Arduino serial

They are available via cfg.backend, but they are not the primary path in the current example workflows shown here.

Helper Scripts

`Helper Scripts/FindCenterPoints.py`

This script appears to be a utility for analyzing kinematic center point data and related foot-placement experiments. It is not part of the active runtime path.

`Helper Scripts/torqueCalc.py`

This is a stand-alone Tkinter utility for estimating servo torque requirements based on robot mass, leg dimensions, and safety factor. It is a design/support script rather than part of the main robot runtime.

Bottom Line

The current codebase is a working hybrid of:

real-time robot control,
PyBullet simulation,
GUI/manual input,
PPO-based RL training and evaluation,
partial state-machine scaffolding,
hardware communication support.

The information in this README has been updated to match the current files in the workspace, especially around the true ML entry points, active runtime flow, and the fact that some older state-helper modules are present but not currently wired into the main execution path.

README.md

JackBot — Hexapod Control, Simulation & RL Framework