Hardware Connection fix

readme was outdated
env rewards got a penalty for standing still while it should move instead of 0 reward

.gitignore

@@ -3,8 +3,12 @@ node_modules/
 .venv/
 .vscode/
 __pycache__/
-ml/checkpoints/
+ml/checkpoints/*
-ml/logs/
+!ml/checkpoints/jackbot_kinematics_base.zip
+!ml/checkpoints/jackbot_ppo6_149994_steps.zip
+!ml/checkpoints/jackbot_ppo6_199992_steps.zip
+!ml/checkpoints/jackbot_ppo6_499980_steps.zip
+#ml/logs/
 
 # Ignore environment files with private passwords/keys
 .env

ArduinoCommunication.py

@@ -11,7 +11,7 @@ class ArduinoCommunication(Thread):
     def __init__(self, port=cfg.port, baudrate=cfg.baudrate, timeout=cfg.comm_timeout):
         super().__init__()
         self.daemon = True  # Thread schlie�t sich beim Programmende
-        self.serial_conn  = serial.Serial(port, baudrate, timeout=timeout)
+        self.serial_conn = serial.Serial(port, baudrate, timeout=timeout)
         time.sleep(2)  # Warten bis Arduino ready
 
         self.command_queue = Queue()
@@ -19,6 +19,9 @@ class ArduinoCommunication(Thread):
         self.running = Event()
         self.running.set()
 
+    def send_motion(self, radial_array):
+        self.write(radial_array)
+
     def run(self):
         while self.running.is_set():
             # 1. Befehle senden
@@ -70,5 +73,12 @@ class ArduinoCommunication(Thread):
 
     def stop(self):
         self.running.clear()
-        self.join()
+        try:
-        self.serial_conn.close()
+            self.join(timeout=0.5)
+        except RuntimeError:
+            pass
+        if self.serial_conn and self.serial_conn.is_open:
+            self.serial_conn.close()
+
+    def close(self):
+        self.stop()

EspCommunication.py

@@ -25,10 +25,14 @@ HEADER_FMT  = "<B B H I" # ID, Ver, Len, Timestamp
 TLV_FMT = "<B H"
 
 class ESP32Communication(Thread):
-    def __init__(self, host=cfg.esp32_ip, port=cfg.esp32_port):
+    def __init__(self, host=None, port=None, ip=None):
         super().__init__(daemon=True)
+        if ip is not None:
+            host = ip
+        host = cfg.esp32_ip if host is None else host
+        port = cfg.esp32_port if port is None else port
         self.addr = (host, port)
-        
+
         # UDP Socket - Zero Lag
         self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
         
@@ -126,4 +130,10 @@ class ESP32Communication(Thread):
 
     def stop(self):
         self.running.clear()
-        self.sock.close()
+        try:
+            self.sock.close()
+        except Exception:
+            pass
+
+    def close(self):
+        self.stop()

README.md

@@ -1,46 +1,50 @@
 # 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.
+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 robot project that combines:
+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, and
+* a physics simulator (`simulation.py`) for testing in software before using real hardware,
-* a reinforcement learning pipeline that teaches the robot how to move through trial and error or pretraining from kinematics.
+* a reinforcement learning pipeline that trains locomotion policies with PPO and behavioral cloning,
+* and a GUI/gamepad input layer for manual operation.
 
-The project is a full control system that can:
+The project can currently:
 
 * run the robot in a PyBullet simulation,
-* accept human commands from a GUI or gamepad,
+* accept commands from a GUI or gamepad,
-* stream target joint positions to physical hardware (ESP32 / Arduino), and
+* stream target joint positions to physical hardware (ESP32 / Arduino),
-* train and evaluate policy checkpoints using PPO and Behavioral Cloning (BC).
+* 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 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.
+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 practice, the RL workflow looks like this:
+In the current codebase, the practical workflow is:
 
-1. The robot starts from a standing pose.
+1. Start the robot in either simulation or hardware mode.
-2. The policy receives sensory information (joint angles + command vector) from the simulation.
+2. Feed motion commands from GUI/gamepad or a training environment.
-3. The policy chooses continuous target joint positions.
+3. Convert target foot positions into joint targets using inverse kinematics.
-4. The simulation steps forward and computes physics.
+4. Apply these targets to the active backend.
-5. The policy receives shaped rewards for surviving, following commands, and remaining stable.
+5. Train or evaluate PPO policies against the simulated robot.
-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.
+* **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`)**: Turns target leg positions into actual joint angles through inverse kinematics.
+* **Kinematics (`kinematics.py`)**: Converts target foot positions into joint angles using IKPy chains for each leg.
-* **Simulation (`simulation.py`)**: Consolidates PyBullet scene management, motor positioning, physics queries, and environment stepping into a single manager.
+* **Simulation (`simulation.py`)**: Manages PyBullet scene loading, stepping, joint actuation, base pose queries, and physics telemetry.
-* **Inputs (`inputs/`)**: Interface for receiving control commands from a GUI, gamepad, or procedural random generators.
+* **Inputs (`inputs/`)**: Receives commands from GUI sliders, Pygame gamepads, and random command generation.
-* **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`).
+* **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()`.
 
 ---
 
@@ -48,49 +52,68 @@ In practice, the RL workflow looks like this:
 
 ```text
 JackBot/
-├── main.py                  # Manual control, GUI, and hardware streaming entry point
+├── main.py                  # Manual runtime entry point
-├── Robot.py                 # Robot wrapper, gait execution, and backend abstraction
+├── Robot.py                 # Unified robot wrapper + backend selection
-├── config.py                # Global settings and hardware/backend configuration
+├── config.py                # Central runtime configuration
-├── kinematics.py            # Inverse kinematics mapping cartesian targets to joint angles
+├── kinematics.py            # IK/FK helpers
-├── simulation.py            # PyBullet simulation manager (handles engine, URDF loading, & physics)
+├── simulation.py            # PyBullet scene / physics manager
-├── robot_init.py            # Neutral standing pose and initial joint definitions
+├── robot_init.py            # Initial joint definitions and center points
-├── DataTypes.py             # Typed numpy structures for positions and joint angles
+├── DataTypes.py             # Typed arrays for positions / angles
+├── JackBotUrdf.urdf        # Robot URDF
 │
-├── states/                  # Gait and state-machine behavior logic
+├── states/                  # State classes present in the project
 │   ├── State.py
 │   ├── IdleState.py
 │   ├── WalkingState.py
-│   └── WaveState.py
+│   ├── WaveState.py
+│   ├── ml_walking.py
+│   └── __init__.py
 │
-├── inputs/                  # Command input sources
+├── inputs/                  # Command input providers
 │   ├── InputProvider.py
 │   ├── PygameController.py
-│   └── RandomeInputProvider.py
+│   ├── RandomeInputProvider.py
+│   └── RandomInputProvider.py
 │
-├── gui/                     # Visual control frontend (Tkinter/CustomTkinter)
+├── gui/                     # GUI/dashboard control layer
 │   └── MainWindow.py
 │
-├── EspCommunication.py      # ESP32 socket communication layer
+├── ml/                      # Gymnasium RL training + evaluation stack
-├── ArduinoCommunication.py  # Arduino serial communication layer
+│   ├── env.py               # RL environment + reward logic
-├── JackBotUrdf.urdf         # Robot URDF mesh and joint axis definition
+│   ├── 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
 │
-└── ml/                      # Reinforcement Learning Subsystem
+├── EspCommunication.py      # ESP32 UDP communication layer
-    ├── env.py               # Gymnasium environment wrapping PyBullet & Robot control
+├── ArduinoCommunication.py  # Arduino serial communication layer
-    ├── callbacks.py         # Custom SB3 callbacks for TensorBoard logging & Curriculum progression
+├── Helper Scripts/          # Utility scripts
-    ├── pretrain_bc.py       # Behavioral Cloning (BC) script to pre-train policy from IK teacher
+│   ├── FindCenterPoints.py
-    ├── run_train.py         # Main training entry point (runs parallelized PPO via SubprocVecEnv)
+│   └── torqueCalc.py
-    ├── run_eval.py          # Phase-by-phase policy evaluator with CLI reports
+│
-    ├── run_eval_training.py # Benchmark reward system across phases using pure kinematics
+├── requirements.txt
-    └── MetricsOverlay.py    # PyBullet HUD debug text and visual overlays
+├── 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)
+* **Python:** 3.12 recommended
 * **OS:** Windows 10/11 or Linux
-* **Dependencies:** `pybullet`, `gymnasium`, `stable-baselines3`, `torch`, `numpy`, `pygame`, `ikpy`, `pyserial`, `matplotlib`
+* **Dependencies:** `pybullet`, `gymnasium`, `stable-baselines3`, `torch`, `numpy`, `pygame`, `ikpy`, `pyserial`, `matplotlib`, `dearpygui`
 
 ---
 
@@ -117,138 +140,269 @@ pip install -r requirements.txt
 
 ---
 
-### Usage Guide
+## 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.
+`main.py` is the main operational entry point for driving the robot manually through the GUI or a connected gamepad.
 
 To launch:
+
 ```bash
 python main.py
 ```
 
-#### Configuration (`config.py`)
+### Configuration (`config.py`)
-Edit `config.py` prior to launching `main.py` to configure execution mode and connections:
 
-* **Backend Selection (`cfg.backend`):**
+Before launching `main.py`, edit `config.py` to select the backend and connection targets.
-  * `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`).
+* **Backend Selection (`cfg.backend`)**:
-  * `BackendType.ARDUINO`: Streams target joint angles over Serial to an Arduino (`cfg.port`, `cfg.baudrate`).
+  * `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/`)
+## 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.
+The ML subsystem currently uses:
 
-### ML Workflow Steps
+* **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
 
-The policy input contains:
+### Current ML scripts
 
-* 18 joint angles from the robot, and
+* `ml/run_train.py` — PPO training entry point
-* 4 command dimensions describing the desired motion direction and yaw rate.
+* `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
 
-The policy output is:
+### Training flow
 
-* 18 continuous values, one per joint, describing how much each joint should change.
+Training is launched with:
 
-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:
 ```bash
-python ml/run_train.py --timesteps 100000 --model ml/checkpoints/ppo_joint_command
+python ml/run_train.py --total-timesteps 1500000 --gui
 ```
 
-Training uses the simulator as the environment, runs PPO updates, and periodically saves checkpoint models.
+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 flow
 
-Evaluation uses `ml/run_eval.py` and `ml/evaluate.py`.
+Evaluation is launched with:
 
-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
+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:
+
+```bash
+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 is not a single number with one purpose. It combines several terms so the policy learns to:
+The reward function in `ml/env.py` is designed to teach several things at once:
 
-* survive in the environment,
+* survival and stability,
-* move in the commanded direction,
+* following the commanded direction,
-* avoid unwanted sideways drift,
+* avoiding lateral drift,
-* stay upright without excessive tilt,
+* maintaining base height,
-* avoid giant control jumps,
+* reducing abrupt control changes,
-* keep the robot away from a low collapsed posture.
+* staying close to a stable stance when the command is zero.
 
-The main ideas are:
+The environment uses several reward components, including:
 
-* a small alive bonus for staying upright,
+* height reward,
-* directional movement rewards,
+* stability reward,
-* penalties for drifting or standing still when a command is active,
+* pose closeness reward,
-* penalties for excessive tilt or collapse,
+* smoothness reward,
-* a penalty for large abrupt control changes.
+* linear velocity tracking,
+* angular velocity tracking,
+* jitter penalty,
+* stand-by penalty when the command is zero.
 
-This style of reward shaping encourages the policy to learn real locomotion rather than just freezing in place.
+### Detailed reward structure
 
-### Curriculum / phase progression
+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.
 
-The environment uses a staged curriculum to expose command complexity gradually.
+At each control step, the environment measures:
 
-* **Phase 0**: basic standing / forward-focused behavior
+* current body height,
-* **Phase 1**: turning and directional regularization
+* roll and pitch angles,
-* **Phase 2**: omni-directional movement
+* current joint configuration,
-* **Phase 3**: full-command challenge
+* measured linear and angular velocities,
+* the active command vector `[vx, vy, omega]`,
+* the difference between the current action and the previous actions.
 
-The goal is to prevent the learner from being asked to master all difficult motion goals at once.
+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:
+
+```python
+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 simple, human-readable pipeline is:
+A simplified view of the current runtime is:
 
-1. `main.py` or the training/evaluation wrapper creates the robot environment.
+1. `main.py` starts the joystick controller process and opens the GUI.
-2. The simulation runs the robot in PyBullet.
+2. The GUI resolves active commands (`Gamepad`, slider, or random walk).
-3. The policy observes the latest joint states and command vectors.
+3. `Robot.tick()` receives the current motion vector and applies it through the active backend.
-4. PPO predicts a new action.
+4. In simulation mode, `Simulation.step()` advances the PyBullet world.
-5. The action is applied to the robot joints.
+5. In training/evaluation mode, `JackBotEnv.step()` computes reward, updates curriculum, and returns observations.
-6. The simulator steps forward one frame.
+6. PPO uses the observation/action loop to improve the policy.
-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.
 
 ---
 
@@ -257,55 +411,73 @@ That loop is the essence of the machine learning part of the project.
 This repository is useful because it combines several layers that are often separate:
 
 * robot control and kinematics,
-* physical simulation,
+* physics simulation,
 * command input sources,
 * RL environment construction,
-* PPO training and evaluation.
+* PPO training and evaluation,
+* hardware communication layers.
 
-For a beginner, the easiest way to think about the repo is:
+For a newcomer, the easiest way to think about the project is:
 
-- `main.py` is for direct manual control,
+- `main.py` is the manual control entry point,
-- `ml/env.py` is the simulator-to-policy interface,
+- `Robot.py` is the core robot wrapper,
-- `ml/train.py` is where training happens,
+- `simulation.py` is the physics layer,
-- `ml/evaluate.py` is where you check whether the learned policy is actually good.
+- `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.
 
 ---
 
-## Reward Function & Termination Mechanics (`ml/env.py`)
+## Current runtime notes and caveats
 
-### 1. Reward Function Formulation
+### State system status
 
-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:
+The state classes are present in the repository, but the current runtime path is not using them as the main control loop:
 
-$$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}}$$
+* `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
 
-* **Alive Bonus ($R_{\text{alive}} \approx +0.02$):** A small positive baseline awarded every timestep the robot remains upright.
+### Input source status
-* **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.
+The repository currently includes both:
-* **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.
+* `inputs/PygameController.py` for gamepad input
-* **Height Penalty ($P_{\text{height}}$):** Penalizes falling below the measured post-settle standing height.
+* `gui/MainWindow.py` for selecting `Gamepad`, `GUI Sliders`, and `Random Walk`
-* **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.
+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.
 
 ---
 
-### 2. Failure Detection & Termination Logic
+## Helper Scripts
 
-In multi-robot vectorized training (`JackBotEnv`), individual robot failures are handled independently to allow maximum simulation efficiency:
+### `Helper Scripts/FindCenterPoints.py`
 
-* **Individual Failure Masking:** A robot is flagged as failed (`failed_robots_mask[idx] = True`) if either condition is met:
+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.
-  * **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.
+### `Helper Scripts/torqueCalc.py`
-* **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.
+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.
-* **Environment Truncation (`truncated=True`):** Occurs when the episode reaches the maximum allowable step budget (`max_episode_steps`).
 
 ---
 
-## Environment Mechanics & Telemetry (`ml/env.py`)
+## Bottom Line
 
-When training with `--gui`, `JackBotEnv` includes dynamic visual feedback mechanisms:
+The current codebase is a working hybrid of:
 
-* **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.
+* real-time robot control,
-* **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.
+* 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.

Robot.py

@@ -59,8 +59,11 @@ class HardwareBackend:
         return [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]
 
     def cleanup(self) -> None:
-        if self.comm_channel and hasattr(self.comm_channel, 'close'):
+        if self.comm_channel:
-            self.comm_channel.close()
+            if hasattr(self.comm_channel, 'stop'):
+                self.comm_channel.stop()
+            elif hasattr(self.comm_channel, 'close'):
+                self.comm_channel.close()
 
 
 class PyBulletBackend:
@@ -123,9 +126,11 @@ class Robot:
                 self.backend: RobotBackend = PyBulletBackend(sim_instance)
             elif backend_type == BackendType.ESP32:
                 comm = ESP32Communication(ip=cfg.esp32_ip, port=cfg.esp32_port)
+                comm.start()
                 self.backend = HardwareBackend(comm)
             elif backend_type == BackendType.ARDUINO:
                 comm = ArduinoCommunication(port=cfg.port, baudrate=cfg.baudrate)
+                comm.start()
                 self.backend = HardwareBackend(comm)
         else:
             self.backend = backend_type

config.py

@@ -38,6 +38,14 @@ class RobotConfig:
     tick_rate_hz: float = 25.0         # Motion loop execution rate [ticks/sec]
     step_duration: float = 0.8         # Time to complete full gait stride [s]
 
+    @property
+    def standard_tickpersec(self) -> float:
+        return self.tick_rate_hz
+
+    @standard_tickpersec.setter
+    def standard_tickpersec(self, value: float) -> None:
+        self.tick_rate_hz = value
+
     @property
     def tick_duration(self) -> float:
         return 1.0 / self.tick_rate_hz

ml/MetricsOverlay.py

@@ -1,5 +1,9 @@
 """
-ml/MetricsOverlay.py - Camera-Facing (Billboard) 3D Floating Text Overlay
+ml/MetricsOverlay.py - 3D HUD overlay for live simulation telemetry.
+
+This file provides a small PyBullet HUD renderer that draws live robot metrics in the
+simulation scene. It is used during GUI runs to show the current phase, command vector,
+reward information, and basic motion statistics without leaving the 3D view.
 """
 from typing import List, Tuple, Optional, Dict
 import numpy as np

ml/callbacks.py

@@ -1,6 +1,9 @@
 """
-ml/callbacks.py - Stable-Baselines3 Custom Callbacks for Logging & Curriculum Advancement
+ml/callbacks.py - Stable-Baselines3 callbacks for training diagnostics.
-Fully compatible with SubprocVecEnv and DummyVecEnv.
+
+These callbacks extend SB3 training with two responsibilities: logging reward-component
+statistics for TensorBoard/console output, and checking whether the curriculum should
+advance to a harder set of commands based on recent training performance.
 """
 
 import numpy as np

ml/env.py

@@ -1,5 +1,10 @@
 """
-ml/env.py - Gymnasium Environment for JackBot Hexapod RL Training
+ml/env.py - Gymnasium environment for JackBot RL training and evaluation.
+
+This file defines JackBotEnv, the main training/evaluation environment used by PPO.
+It wraps the PyBullet simulation and Robot interfaces into a Gymnasium-compatible
+step/reset loop, manages command sampling, curriculum progression, and reward
+calculation, and exposes metrics that the training callbacks can log.
 """
 import time
 import math
@@ -184,7 +189,19 @@ class JackBotEnv(gym.Env):
         return self._get_obs(), {}
 
     def _get_obs(self) -> np.ndarray:
-        return self.robot.get_observation(command=self.command)
+        # Read raw joint angles from backend
+        raw_angles = np.asarray(self.robot.backend.get_joint_angles(), dtype=np.float32).flatten()
+
+        min_lim = self.min_joint_limits.flatten()
+        max_lim = self.max_joint_limits.flatten()
+
+        # Map raw joint radians [min, max] -> normalized [-1, 1]
+        normalized_joints = 2.0 * (raw_angles - min_lim) / (max_lim - min_lim) - 1.0
+        normalized_joints = np.clip(normalized_joints, -1.0, 1.0)
+
+        # Concatenate normalized joints with active command vector
+        obs = np.concatenate([normalized_joints, self.command]).astype(np.float32)
+        return obs
 
     def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
         previous_action = self.last_action.copy()
@@ -194,25 +211,36 @@ class JackBotEnv(gym.Env):
         self.last_last_action = self.last_action.copy()
         self.last_action = action.copy()
 
-        # Command resampling
+        # Command resampling ONLY if random_command is True
-        if self.random_command and (self.step_count >= self.next_cmd_resample_step or self._curriculum_advanced):
+        if self.random_command and (
+            self.step_count >= self.next_cmd_resample_step or self._curriculum_advanced):
             self.command = self.sample_command()
             random_interval = np.random.randint(self.min_cmd_hold_steps, self.max_cmd_hold_steps + 1)
             self.next_cmd_resample_step = self.step_count + random_interval
 
-        # Extract [vx, vy, omega]
+        # Extract active [vx, vy, omega]
         cmd_vx, cmd_vy, cmd_omega = self.command
 
-        # Mirror main.py input resolution logic: update robot_state and vector_dirmov directly
+        # Mirror input resolution logic to keep robot state synchronized
         self.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle"
         self.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
-        
+
-        joint_range = np.minimum(
+        # Direct Mode: Target joint scaling
-            self.default_joint_angles - self.min_joint_limits,
+        action_flat = np.asarray(action, dtype=np.float32).flatten()
-            self.max_joint_limits - self.default_joint_angles
+        action_clipped = np.clip(action_flat, -1.0, 1.0)
-        )
+
-        target_angles = self.default_joint_angles + action * joint_range
+        if self.robot_mode == "direct":
-        self.robot.tick(action=target_angles)
+            # Map [-1, 1] linearly to physical joint limits [min, max]
+            min_lim = self.min_joint_limits.flatten()
+            max_lim = self.max_joint_limits.flatten()
+            
+            target_angles = min_lim + (action_clipped + 1.0) * 0.5 * (max_lim - min_lim)
+        else:
+            # Residual mode mapping logic
+            target_angles = self.default_joint_angles.flatten() + action_clipped * 0.20
+
+        # Apply target joint angles to physics engine
+        self.robot.tick(action=target_angles, physics_substeps=4)
 
         if self.robot_mode != "kinematics" and self.step_count % 60 == 0:
             random_force = np.random.uniform(-2.0, 2.0, size=2)
@@ -222,8 +250,9 @@ class JackBotEnv(gym.Env):
         self._update_distance_metrics()
         self._update_curriculum()
 
+        # Build next observation preserving active command
         obs = self._get_obs()
-        reward = self._compute_reward(action, previous_action)
+        reward = self._compute_reward(action_flat, previous_action)
 
         self.cumulative_reward += reward
         self.robot_reward += reward
@@ -235,6 +264,9 @@ class JackBotEnv(gym.Env):
         if self.step_count % 120 == 0 and self.use_gui:
             self._update_hud()
 
+        if self.use_gui:
+            time.sleep(1.0 / self.control_freq)
+
         return obs, reward, terminated, truncated, info
 
     def _update_distance_metrics(self):
@@ -336,7 +368,10 @@ class JackBotEnv(gym.Env):
             target_speed = math.hypot(target_vx, target_vy)
 
             if not is_moving:
-                total_reward = 0.0
+                # If the command says move but the robot stays effectively still,
+                # give a real penalty instead of a neutral reward.
+                stillness_penalty = -0.10
+                total_reward = stillness_penalty
             else:
                 lin_vel_error = (filtered_vx - target_vx)**2 + (filtered_vy - target_vy)**2
                 r_lin_vel = math.exp(-25.0 * lin_vel_error)
@@ -344,7 +379,7 @@ class JackBotEnv(gym.Env):
 
                 if target_speed > 0.08 and raw_speed < 0.03:
                     r_lin_vel = 0.0
-                    stillness_penalty = -0.1  # Softened from -0.25
+                    stillness_penalty = -0.05
 
                 w_lin_vel, w_ang_vel, w_height, w_stability, w_smoothness = 0.55, 0.15, 0.10, 0.12, 0.08
                 total_reward = (
@@ -354,7 +389,7 @@ class JackBotEnv(gym.Env):
 
         step_reward = float(total_reward / 10.0)
         alive_bonus = 0.01
-        final_reward = max(0.0, step_reward + jitter_penalty + alive_bonus)
+        final_reward = step_reward + jitter_penalty + alive_bonus
 
         self.last_reward_components = {
             "height": float(r_height),
@@ -365,6 +400,7 @@ class JackBotEnv(gym.Env):
             "ang_vel": float(r_ang_vel),
             "jitter_penalty": float(jitter_penalty),
             "stand_penalty": float(stand_penalty),
+            "stillness_penalty": float(stillness_penalty),
             "total": final_reward,
         }
         for k, v in self.last_reward_components.items():

ml/pretrain_bc.py

@@ -1,5 +1,10 @@
 """
-ml/pretrain_bc.py - Behavioral Cloning from Kinematics Teacher
+ml/pretrain_bc.py - Behavioral Cloning teacher-data pipeline.
+
+This script collects observation/action pairs from the kinematics solver, then uses
+those pairs as training data for a PPO policy. In practice it serves as a teacher-
+student pretraining step: the kinematics controller generates sample trajectories, and
+this file teaches the policy to imitate those behavior patterns before PPO training.
 """
 import sys
 import time
@@ -21,7 +26,8 @@ def collect_kinematics_dataset(num_samples: int = 100_000, use_gui: bool = False
     """Collects (Observation, Action) pairs directly from Kinematics Teacher."""
     print(f"\n[Pretrain] Collecting {num_samples} samples from Kinematics Teacher (GUI={use_gui})...")
     
-    env = JackBotEnv(use_gui=use_gui, robot_mode="kinematics")
+    # Disable random command resampling inside env.step so manual command locks persist
+    env = JackBotEnv(use_gui=use_gui, robot_mode="kinematics", random_command=False)
     env.curriculum_phase = CurriculumPhase.FULL_COMMAND
 
     observations = []
@@ -32,39 +38,41 @@ def collect_kinematics_dataset(num_samples: int = 100_000, use_gui: bool = False
     # --- PROGRESS BAR: Data Collection ---
     pbar = tqdm(range(num_samples), desc=" Collecting Data", unit="step")
     for i in pbar:
-        # 1. Sample random movement command
+        # 1. Update command and vector targets every 120 steps
         if i % 120 == 0:
             env.command = env.sample_command()
+            cmd_vx, cmd_vy, cmd_omega = env.command
+            env.robot.robot_state = (
+                "walking"
+                if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01)
+                else "idle"
+            )
+            env.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
 
-        # 2. Let Robot compute Kinematics target angles
+        # 2. Capture observation BEFORE stepping environment
-        cmd_vx, cmd_vy, cmd_omega = env.command
+        current_obs = env._get_obs()
-        env.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle"
+
-        env.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
+        # 3. Step environment ONCE (updates kinematics solver, PyBullet physics, and computes IK)
-        
+        obs, _, terminated, truncated, _ = env.step(np.zeros(18, dtype=np.float32))
-        # Advance internal IK tick to calculate joint angles
+
-        env.robot.tick()
+        # 4. Extract procedural IK joint targets computed during this step
-        
-        # Extract .data and flatten (6, 3) matrix to 18-dim 1D array
         target_ik_rad = env.robot.current_rad.data.flatten().copy()
 
-        # 3. Convert target angles back to normalized [-1, 1] action space
+        # Step 5: Convert target radians directly to [-1, 1] relative to joint limits
-        normalized_action = np.where(
+        min_lim = env.min_joint_limits.flatten()
-            target_ik_rad >= env.default_joint_angles,
+        max_lim = env.max_joint_limits.flatten()
-            (target_ik_rad - env.default_joint_angles) / np.maximum(1e-5, env.max_joint_limits - env.default_joint_angles),
+
-            (target_ik_rad - env.default_joint_angles) / np.maximum(1e-5, env.default_joint_angles - env.min_joint_limits)
+        normalized_action = 2.0 * (target_ik_rad - min_lim) / (max_lim - min_lim) - 1.0
-        )
         normalized_action = np.clip(normalized_action, -1.0, 1.0)
 
-        # 4. Save sample
+        # Step 6: Store matching input (obs) and target ground truth (normalized_action)
-        observations.append(obs.copy())
+        observations.append(current_obs.copy())
         actions.append(normalized_action.copy())
 
-        # Step simulation environment
-        obs, _, terminated, truncated, _ = env.step(normalized_action)
-
         if use_gui:
             time.sleep(1.0 / 60.0)
 
+        # 7. Handle episode boundaries using terminated and truncated
         if terminated or truncated:
             obs, _ = env.reset()
 

ml/run_eval.py

@@ -1,6 +1,10 @@
 """
-ml/run_eval.py - Phase-by-Phase Policy Evaluator for JackBot
+ml/run_eval.py - Evaluate a saved JackBot policy across fixed command phases.
-Evaluates a trained model across all curriculum phases with fixed command vectors.
+
+This script loads a trained PPO checkpoint, instantiates the Gymnasium environment in
+non-random mode, and runs deterministic evaluation episodes for several command
+regimes. It is used to measure whether a policy can survive, move, and maintain
+stability under forward, turning, lateral, and omni-direction commands.
 """
 
 import argparse
@@ -33,11 +37,12 @@ def main():
         use_gui=args.gui,
         random_command=False,
         max_episode_steps=args.max_steps_per_episode,
+        robot_mode="direct"
     )
 
     # Multi-Phase Configurations Suite
     phase_configs = [
-        (CurriculumPhase.STAND_ONLY,          "STAND",                 np.array([0.0, 0.0, 0.0], dtype=np.float32)),
+        #(CurriculumPhase.STAND_ONLY,          "STAND",                 np.array([0.0, 0.0, 0.0], dtype=np.float32)),
         (CurriculumPhase.FORWARD,             "FORWARD GAIT",           np.array([1.0, 0.0, 0.0], dtype=np.float32)),
         (CurriculumPhase.TURN_AND_DIRECTION,  "FORWARD + YAW TURN",     np.array([0.5, 0.0, 0.4], dtype=np.float32)),
         (CurriculumPhase.OMNI_DIRECTION,      "STRIDE LATERAL",         np.array([0.5, 0.5, 0.0], dtype=np.float32)),
@@ -62,9 +67,18 @@ def main():
             for ep in range(args.episodes_per_phase):
                 obs, _ = env.reset()
 
-                # Force environment into active curriculum phase and lock command
+                # Force environment into active curriculum phase and lock command BEFORE getting obs
                 env.curriculum_phase = phase_enum
                 env.command = test_cmd.copy()
+                cmd_vx, cmd_vy, cmd_omega = test_cmd
+                env.robot.robot_state = (
+                    "walking"
+                    if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01)
+                    else "idle"
+                )
+                env.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
+
+                # Get correct observation with test_cmd attached
                 obs = env._get_obs()
 
                 done = False
@@ -72,12 +86,10 @@ def main():
                 steps = 0
 
                 while not done:
-                    # Enforce locked command each step
-                    env.command = test_cmd.copy()
-
                     # Predict deterministic action from policy
                     action, _ = model.predict(obs, deterministic=True)
                     
+                    # Step environment
                     obs, reward, terminated, truncated, _ = env.step(action)
                     done = terminated or truncated
 

ml/run_eval_training.py

@@ -1,5 +1,10 @@
 """
-ml/run_eval_training.py - Benchmark reward system across all curriculum phases.
+ml/run_eval_training.py - Run a kinematics-mode reward benchmark.
+
+This script repeatedly resets the JackBot environment in kinematics mode and applies
+fixed command vectors for each curriculum phase. It is intended as a lightweight
+benchmark to inspect reward components, movement quality, and survival behavior without
+requiring an already-trained PPO model.
 """
 import time
 import numpy as np
@@ -21,7 +26,7 @@ def evaluate_kinematics(episode_length: int = 1000):
     print("=" * 70 + "\n")
 
     phase_configs = [
-        (CurriculumPhase.STAND_ONLY,           "STAND",              np.array([1.0, 0.0, 0.0], dtype=np.float32)),
+        (CurriculumPhase.STAND_ONLY,          "STAND",               np.array([1.0, 0.0, 0.0], dtype=np.float32)),
         (CurriculumPhase.FORWARD,             "FORWARD GAIT",        np.array([1.0, 0.0, 0.0], dtype=np.float32)),
         (CurriculumPhase.TURN_AND_DIRECTION,  "FORWARD + YAW TURN",  np.array([0.5, 0.0, 0.4], dtype=np.float32)),
         (CurriculumPhase.OMNI_DIRECTION,      "STRIDE LATERAL",      np.array([0.5, 0.5, 0.0], dtype=np.float32)),

ml/run_train.py

@@ -1,4 +1,9 @@
-"""Minimal training launcher for quick experiments.
+"""PPO training launcher for JackBot.
+
+This script is the main entry point for training a policy in the JackBot Gymnasium
+environment. It creates a vectorized environment, optionally loads a pretrained base
+model, runs Stable-Baselines3 PPO for a configured number of timesteps, and saves
+checkpoints plus evaluation artifacts during training.
 
 Usage:
     python ml/run_train.py --total-timesteps 1500000 --gui
@@ -86,7 +91,12 @@ def main():
         model = PPO.load(
             args.pretrained_model,
             env=vec_env,
-            learning_rate=1e-4,  # Lower learning rate so RL fine-tunes without destroying base gait
+            learning_rate=5e-5,  # Lower learning rate so RL fine-tunes without destroying base gait
+            ent_coef=0.0,
+            target_kl=0.05,
+            vf_coef=0.5,
+            max_grad_norm=0.5,
+            verbose=2,
             tensorboard_log=args.log_dir,
             device="cpu",
         )
@@ -95,22 +105,22 @@ def main():
         model = PPO(
             policy="MlpPolicy",
             env=vec_env,
-            learning_rate=1e-4,
+            learning_rate=5e-5,
             n_steps=256,
             batch_size=256,
             n_epochs=10,
             gamma=0.99,
             gae_lambda=0.95,
             clip_range=0.2,
-            ent_coef=0.01,
+            ent_coef=0.0,
             target_kl=0.05,
             vf_coef=0.5,
             max_grad_norm=0.5,
-            verbose=1,
+            verbose=2,
             tensorboard_log=args.log_dir,
             device="cpu",
         )
-
+    model.policy.log_std.data.fill_(-2.0)
     # Setup Callbacks with ppo<number> naming
     checkpoint_callback = CheckpointCallback(
         save_freq=max(1, args.save_freq // args.num_workers),
@@ -127,7 +137,7 @@ def main():
         eval_env,
         best_model_save_path=best_model_path,
         log_path="ml/logs/results",
-        eval_freq=max(1, 20_000 // args.num_workers),
+        eval_freq=max(1, 50_000 // args.num_workers),
         deterministic=True,
         render=False,
     )

simulation.py

@@ -85,7 +85,7 @@ class Simulation:
                 force=30,
                 physicsClientId=self.physics_client
             )
-
+            
     def hard_reset_joint_angles(self, target_angles: Union[np.ndarray, dt.RadArray]) -> None:
         """Instantly teleports joint angles to target positions, clearing velocity state."""
         radflat = target_angles.data.flatten() if isinstance(target_angles, dt.RadArray) else target_angles.flatten()