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/tensorboard/
+ml/checkpoints/*
+!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()
-        self.serial_conn.close()
+        try:
+            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,48 +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 for testing in software before using real hardware, and
-* a reinforcement learning pipeline that teaches the robot how to move through trial and error.
+* 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 is not just one script or one model. It 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,
-* stream target joint positions to physical hardware, and
-* train a policy using PPO so the robot can learn locomotion automatically.
+* 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 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.
+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 workflow looks like this:
+In the current codebase, the practical workflow is:
 
-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.
+1. Start the robot in either simulation or hardware mode.
+2. Feed motion commands from GUI/gamepad or a training environment.
+3. Convert target foot positions into joint targets using inverse kinematics.
+4. Apply these targets to the active backend.
+5. Train or evaluate PPO policies against the simulated robot.
 
 ---
 
 ## 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.
+* **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()`.
 
 ---
 
@@ -50,49 +52,68 @@ In practice, the workflow looks like this:
 
 ```text
 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
+├── 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/                  # gait/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 sources for the robot
+├── inputs/                  # Command input providers
 │   ├── InputProvider.py
 │   ├── PygameController.py
-│   └── RandomeInputProvider.py
+│   ├── RandomeInputProvider.py
+│   └── RandomInputProvider.py
 │
-├── gui/                     # visual control frontend
+├── gui/                     # GUI/dashboard control layer
 │   └── MainWindow.py
 │
-├── EspCommunication.py      # ESP32 communication layer
-├── ArduinoCommunication.py  # Arduino serial communication layer
-├── JackBotUrdf.urdf         # robot mesh and joint definition
+├── 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
 │
-└── 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
+├── 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:** Linux (Ubuntu/Debian) or Windows 10/11
-* **Dependencies:** `pybullet`, `gymnasium`, `stable-baselines3`, `torch`, `numpy`, `pygame`, `ikpy`, `pyserial`, `matplotlib`
+* **Python:** 3.12 recommended
+* **OS:** Windows 10/11 or Linux
+* **Dependencies:** `pybullet`, `gymnasium`, `stable-baselines3`, `torch`, `numpy`, `pygame`, `ikpy`, `pyserial`, `matplotlib`, `dearpygui`
 
 ---
 
@@ -110,7 +131,8 @@ pip install -r requirements.txt
 ### Windows (PowerShell)
 
 ```powershell
-python3.12 -m venv .venv
+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
@@ -120,144 +142,267 @@ pip install -r requirements.txt
 
 ## Usage Guide
 
-### 1. Manual Control & Hardware Streaming (`main.py`)
+## 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`)
-Edit `config.py` prior to launching `main.py` to configure execution mode and connections:
+### Configuration (`config.py`)
 
-* **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`).
+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()`.
 
 ---
 
-### 2. How the Machine Learning System Works
+## Machine Learning Pipeline (`ml/`)
 
-The machine learning subsystem teaches the robot to move by interacting with a simulation environment instead of relying on a manually written gait controller.
+The ML subsystem currently uses:
 
-At a high level:
+* **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 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.
+### Current ML scripts
 
-#### What the policy sees and does
+* `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
 
-The policy input contains:
+### Training flow
 
-* 18 joint angles from the robot, and
-* 4 command dimensions describing the desired motion direction and yaw rate.
+Training is launched with:
 
-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:
 ```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,
-* 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.
+* 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 main ideas are:
+The environment uses several reward components, including:
 
-* 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.
+* 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.
 
-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
-* **Phase 1**: turning and directional regularization
-* **Phase 2**: omni-directional movement
-* **Phase 3**: full-command challenge
+* 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.
 
-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.
-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.
+1. `main.py` starts the joystick controller process and opens the GUI.
+2. The GUI resolves active commands (`Gamepad`, slider, or random walk).
+3. `Robot.tick()` receives the current motion vector and applies it through the active backend.
+4. In simulation mode, `Simulation.step()` advances the PyBullet world.
+5. In training/evaluation mode, `JackBotEnv.step()` computes reward, updates curriculum, and returns observations.
+6. PPO uses the observation/action loop to improve the policy.
 
 ---
 
@@ -266,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,
-- `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.
+- `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.
 
 ---
 
-## 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.
-* **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.
+### 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.
 
 ---
 
-### 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:
-  * **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`).
+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.
 
 ---
 
-## 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.
-* **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.
+* 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.

Robot.py

@@ -1,12 +1,9 @@
 """
-Robot.py - Unified Robot Class for JackBot
-Handles state, kinematics, backends (Hardware/Simulation), and motion execution.
+Robot.py - Central Robot Control, Kinematics, State Machine & Hardware Abstraction
 """
-
-from typing import Protocol, Optional, Union
+from typing import Protocol, Optional, Union, Tuple, List
 import numpy as np
 import math
-import pybullet as p
 
 from states import STATE_REGISTRY
 from states.State import State
@@ -15,93 +12,130 @@ import kinematics as kin
 import robot_init as ri
 from config import cfg, BackendType
 
-# Import communications and simulation modules
 from simulation import Simulation
 from EspCommunication import ESP32Communication
 from ArduinoCommunication import ArduinoCommunication
 
 
 class RobotBackend(Protocol):
-    """Abstraction layer for hardware vs simulation output."""
-    def send_angles(self, rad_array: dt.RadArray) -> None:
-        ...
-    def step_simulation(self) -> None:
-        ...
-    def cleanup(self) -> None:
-        ...
+    """Protocol defining hardware abstraction for both Simulation and Hardware backends."""
+    def send_angles(self, rad_array: dt.RadArray) -> None: ...
+    def step_simulation(self) -> None: ...
+    def hard_reset_joints(self, target_angles: np.ndarray) -> None: ...
+    def reset_base(self, position: Optional[List[float]] = None, orientation: Optional[List[float]] = None) -> None: ...
+    def get_joint_angles(self) -> np.ndarray: ...
+    def get_base_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]: ...
+    def get_base_velocity(self) -> Tuple[List[float], List[float]]: ...
+    def cleanup(self) -> None: ...
 
 
 class HardwareBackend:
-    """Backend for physical ESP32 or Arduino robot."""
+    """Backend for physical ESP32 or Arduino microcontrollers."""
     def __init__(self, comm_channel):
         self.comm_channel = comm_channel
+        self.internal_angles = np.zeros(18, dtype=np.float32)
 
     def send_angles(self, rad_array: dt.RadArray) -> None:
+        self.internal_angles = rad_array.data.flatten().copy()
         if self.comm_channel:
             self.comm_channel.send_motion(rad_array)
 
     def step_simulation(self) -> None:
-        pass  # Physical hardware steps in real-time
+        pass
+
+    def hard_reset_joints(self, target_angles: np.ndarray) -> None:
+        self.internal_angles = target_angles.flatten().copy()
+
+    def reset_base(self, position: Optional[List[float]] = None, orientation: Optional[List[float]] = None) -> None:
+        pass
+
+    def get_joint_angles(self) -> np.ndarray:
+        return self.internal_angles.copy()
+
+    def get_base_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]:
+        return [0.0, 0.0, 0.122], (0.0, 0.0, 0.0)
+
+    def get_base_velocity(self) -> Tuple[List[float], List[float]]:
+        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'):
-            self.comm_channel.close()
+        if self.comm_channel:
+            if hasattr(self.comm_channel, 'stop'):
+                self.comm_channel.stop()
+            elif hasattr(self.comm_channel, 'close'):
+                self.comm_channel.close()
 
 
 class PyBulletBackend:
-    """Backend for PyBullet simulation execution."""
-    def __init__(self, sim_instance, body_id: Optional[int] = None):
+    """Backend mapping Robot operations directly to PyBullet simulation engine."""
+    def __init__(self, sim_instance: Simulation):
         self.sim = sim_instance
-        self.body_id = body_id
 
     def send_angles(self, rad_array: dt.RadArray) -> None:
         if self.sim:
-            # If body_id is set, target that specific robot body
-            if self.body_id is not None and hasattr(self.sim, 'updatePosForBody'):
-                self.sim.updatePosForBody(self.body_id, rad_array)
-            else:
-                self.sim.updatePos(rad_array)
+            self.sim.set_robot_joint_angles(rad_array)
 
     def step_simulation(self) -> None:
         if self.sim:
             self.sim.step()
 
+    def hard_reset_joints(self, target_angles: np.ndarray) -> None:
+        if self.sim:
+            self.sim.hard_reset_joint_angles(target_angles)
+
+    def reset_base(self, position: Optional[List[float]] = None, orientation: Optional[List[float]] = None) -> None:
+        if self.sim:
+            self.sim.reset_robot_base(pos=position, orn=orientation)
+
+    def get_joint_angles(self) -> np.ndarray:
+        return self.sim.get_robot_joint_angles() if self.sim else np.zeros(18, dtype=np.float32)
+
+    def get_base_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]:
+        return self.sim.get_robot_pose_and_rpy() if self.sim else ([0, 0, 0], (0, 0, 0))
+
+    def get_base_velocity(self) -> Tuple[List[float], List[float]]:
+        return self.sim.get_robot_velocity() if self.sim else ([0, 0, 0], [0, 0, 0])
+
     def cleanup(self) -> None:
-        if self.sim and hasattr(self.sim, 'disconnect'):
+        if self.sim:
             self.sim.disconnect()
 
 
 class Robot:
     """
-    Encapsulates a single JackBot hexapod instance.
-    Maintains joint states, leg positions, kinematics, and backend control.
+    Unified JackBot Class.
+    Coordinates joint memory, IK solvers, procedural tripods, and backend communication.
     """
 
     def __init__(
         self,
         backend_type: Union[BackendType, RobotBackend] = BackendType.SIMULATION,
         start_pose: str = "init_deg",
-        urdf_path: str = cfg.urdf_path
+        urdf_path: str = cfg.urdf_path,
+        mode: str = "kinematics"  # "kinematics", "residual", or "direct"
     ):
         self.urdf_path = urdf_path
         self.start_pose = start_pose
+        self.mode = mode
 
-        # --- BACKEND FACTORY CREATION ---
+        # --- BACKEND INSTANTIATION ---
         if isinstance(backend_type, BackendType):
             if backend_type == BackendType.SIMULATION:
-                # Launch PyBullet 3D Simulation GUI
-                sim_instance = Simulation(urdf_path=self.urdf_path)
+                sim_instance = Simulation(urdf_path=self.urdf_path, use_gui=True)
+                sim_instance.load_scene()
                 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
 
-        # Kinematics and position initialization
+        # Kinematics initialization
         pose_deg = ri.init_deg if start_pose == "init_deg" else ri.init90_deg
         self.current_rad: dt.RadArray = pose_deg.to_rad()
         self.current_pos: dt.PosArray = kin.ikpyForward(self.current_rad)
@@ -109,31 +143,20 @@ class Robot:
             ri.get_center_points() if hasattr(ri, "get_center_points") else ri.center_points
         )
 
-        # RL configuration
-        self.action_scale = 0.1  # Joint delta step size (radians)
-
-        # Gait / motion variables
-        self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
-        self.robot_state = "idle"
+        # RL configuration & Gait state variables
+        self.action_scale = 0.1  # Radian step scale for RL deltas
+        self.gait_phase = 0.0
         self.vector_dirmov = [0.0, 0.0, 0.0]  # [vx, vy, omega]
+        self.robot_state = "idle"
+        self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
 
-        # State machine initialization
+        # State Machine Initialization
         self.current_state_key: str = "idle"
         self.current_state: State = STATE_REGISTRY["idle"]
         self.current_state.enter(self)
 
-    def change_state(self, new_state: State) -> None:
-        if self.current_state:
-            self.current_state.exit(self)
-        self.current_state = new_state
-        self.current_state.enter(self)
-
-    def update(self) -> None:
-        if self.current_state:
-            self.current_state.execute(self)
-        self.step_sim()
-
     def set_joint_angles(self, target_rad: dt.RadArray) -> None:
+        """Updates internal Python memory state and sends angles to active backend."""
         self.current_rad = target_rad
         if self.backend:
             self.backend.send_angles(target_rad)
@@ -143,11 +166,18 @@ class Robot:
             self.backend.step_simulation()
 
     def reset_to_init(self) -> None:
+        """Resets kinematics state and forces instant joint alignment in backend."""
         pose_deg = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg
         self.current_rad = pose_deg.to_rad()
         self.current_pos = kin.ikpyForward(self.current_rad)
-        self.set_joint_angles(self.current_rad)
-        self.step_sim()
+        self.gait_phase = 0.0
+        self.robot_state = "idle"
+        self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
+
+        init_flat = self.current_rad.data.flatten()
+        if self.backend:
+            self.backend.hard_reset_joints(init_flat)
+            self.backend.send_angles(self.current_rad)
 
     def compute_ik(self, target_pos: dt.PosArray) -> dt.RadArray:
         return kin.ikpyInverse(target_pos, initial_rad=self.current_rad)
@@ -163,57 +193,104 @@ class Robot:
             self.current_state = STATE_REGISTRY[next_state_key]
             self.current_state.enter(self)
 
-    def tick(self) -> None:
-        next_state_key = self.current_state.execute(self)
-        if next_state_key:
-            self.transition_to(next_state_key)
+    def tick(self, action: Optional[np.ndarray] = None) -> None:
+        """
+        Unified control loop tick.
+        Processes commands through direct RL, residual RL, or State Machine kinematics.
+        """
+        vx, vy, omega = self.vector_dirmov
+
+        if self.mode == "direct":
+            if action is not None:
+                self.apply_rl_action(action)
+
+        elif self.mode == "residual":
+            self.step_kinematic_gait(vx, vy, omega)
+            if action is not None:
+                self.apply_rl_action_delta(action)
+
+        else:  # "kinematics" / standard State Machine execution
+            next_state_key = self.current_state.execute(self)
+            if next_state_key:
+                self.transition_to(next_state_key)
+
         self.step_sim()
 
-    # --- RL METHODS ---
+    def step_kinematic_gait(self, vx: float, vy: float, omega: float) -> None:
+        """Procedural Tripod Gait solver."""
+        cmd_mag = math.hypot(vx, vy) + abs(omega)
+        if cmd_mag < 0.03:
+            target_rad = self.compute_ik(self.center_points)
+            self.set_joint_angles(target_rad)
+            return
 
-    def apply_rl_action(self, action: np.ndarray) -> None:
-        """
-        Applies continuous RL action deltas [-1, 1] to current joint angles.
-        """
-        action = np.asarray(action, dtype=np.float32)
-        scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale
+        self.gait_phase = (self.gait_phase + 0.22) % (2.0 * math.pi)
+        stride_len = cfg.step_length
+        step_height = cfg.step_height
 
-        current_flat = self.current_rad.data.flatten()
-        updated_flat = np.clip(
-            current_flat + scaled_action,
-            -np.pi / 2,
-            np.pi / 2
+        center_data = (
+            self.center_points.data if hasattr(self.center_points, 'data') else np.array(self.center_points)
         )
 
+        target_positions = []
+        for leg_id in range(6):
+            base_pos = np.array(center_data[leg_id], dtype=np.float32)
+            phase_offset = 0.0 if (leg_id % 2 == 0) else math.pi
+            leg_phase = (self.gait_phase + phase_offset) % (2.0 * math.pi)
+
+            lx, ly = base_pos[0], base_pos[1]
+            rot_dx = -omega * ly
+            rot_dy = omega * lx
+
+            dx_dir = vx + rot_dx
+            dy_dir = vy + rot_dy
+            dir_norm = math.hypot(dx_dir, dy_dir) + 1e-6
+
+            dx_unit = dx_dir / dir_norm
+            dy_unit = dy_dir / dir_norm
+
+            if leg_phase < math.pi:
+                # Swing phase (leg lifted in air, moving forward)
+                progress = math.cos(leg_phase)
+                lift = math.sin(leg_phase) * step_height
+                dx = progress * stride_len * dx_unit
+                dy = progress * stride_len * dy_unit
+                dz = lift
+            else:
+                # Stance phase (leg on ground, pushing body forward)
+                progress = math.cos(leg_phase - math.pi)
+                dx = -progress * stride_len * dx_unit
+                dy = -progress * stride_len * dy_unit
+                dz = 0.0
+
+            target_positions.append(base_pos + np.array([dx, dy, dz], dtype=np.float32))
+
+        target_pos_array = dt.PosArray(np.array(target_positions))
+        target_rad = self.compute_ik(target_pos_array)
+        self.set_joint_angles(target_rad)
+
+    def apply_rl_action(self, action: np.ndarray) -> None:
+        action = np.asarray(action, dtype=np.float32)
+        new_rad = dt.RadArray(data=action.reshape(self.current_rad.data.shape))
+        self.set_joint_angles(new_rad)
+
+    def apply_rl_action_delta(self, action: np.ndarray) -> None:
+        """Applies action deltas on top of joint state for Residual RL."""
+        action = np.asarray(action, dtype=np.float32)
+        scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale
+        current_flat = self.backend.get_joint_angles().flatten()
+        updated_flat = np.clip(current_flat + scaled_action, -np.pi / 2, np.pi / 2)
         new_rad = dt.RadArray(data=updated_flat.reshape(self.current_rad.data.shape))
         self.set_joint_angles(new_rad)
 
     def get_observation(self, command: Optional[np.ndarray] = None) -> np.ndarray:
-        """
-        Returns observation vector [18 joint angles] + [optional 4 command dimensions].
-        Queries PyBullet if backend is PyBulletBackend; otherwise falls back to internal state.
-        """
-        if isinstance(self.backend, PyBulletBackend) and self.backend.sim and hasattr(self.backend.sim, 'physics_client'):
-            physics_client = self.backend.sim.physics_client
-            body_id = self.backend.body_id if self.backend.body_id is not None else 0
-            
-            # Retrieve joint mapping from SimManager/Simulation if available
-            if hasattr(self.backend.sim, 'robot_joints') and body_id in self.backend.sim.robot_joints:
-                joint_indices = self.backend.sim.robot_joints[body_id]
-            else:
-                joint_indices = list(range(18))
-
-            joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=physics_client)
-            joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32)
-        else:
-            joint_angles = self.current_rad.data.flatten().astype(np.float32)
-
+        """Extracts joint positions directly from active backend."""
+        joint_angles = self.backend.get_joint_angles().flatten().astype(np.float32)
         if command is not None:
             cmd = np.asarray(command, dtype=np.float32).flatten()
             return np.concatenate([joint_angles, cmd]).astype(np.float32)
-
-        return joint_angles.astype(np.float32)
+        return joint_angles
 
     def cleanup(self) -> None:
-        if self.backend and hasattr(self.backend, 'cleanup'):
+        if self.backend:
             self.backend.cleanup()

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,7 +1,11 @@
 """
-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
+from typing import List, Tuple, Optional, Dict
 import numpy as np
 import pybullet as p
 
@@ -21,20 +25,15 @@ class MetricsHUD:
             yaw = cam_info[8]
             pitch = cam_info[9]
             
-            # Orient the text normal toward the camera view direction
-            # PyBullet text default faces local +Z/-Y depending on roll,
-            # converting visualizer yaw/pitch to Euler angles (roll, pitch, yaw in radians)
-            roll_rad = 0.0
             pitch_rad = np.radians(pitch + 90.0)
             yaw_rad = np.radians(yaw)
             
             text_orientation = p.getQuaternionFromEuler(
-                [pitch_rad, roll_rad, yaw_rad], 
+                [pitch_rad, 0.0, yaw_rad], 
                 physicsClientId=self.client_id
             )
             return text_orientation
         except Exception:
-            # Fallback default orientation if camera info call fails
             return [0.0, 0.0, 0.0, 1.0]
 
     def update(
@@ -44,30 +43,49 @@ class MetricsHUD:
         robot_rewards: List[float],
         cmd_vel: np.ndarray,
         fps: float = 0.0,
-        avg_height: float = 0.0,
-        roll_pitch: Tuple[float, float] = (0.0, 0.0)
+        height: float = 0.0,
+        roll_pitch: Tuple[float, float] = (0.0, 0.0),
+        mode: str = "direct",
+        phase: str = "STAND_ONLY",
+        distance: float = 0.0,
+        status: str = "ALIVE",
+        reward_components: Optional[Dict[str, float]] = None,
+        ep_step: int = 0,
     ) -> None:
-        """Updates floating black text block in 3D space with billboarding."""
+        """Updates floating text block in 3D space with expanded telemetry."""
         sorted_rewards = sorted(robot_rewards, reverse=True)
         top1 = f"{sorted_rewards[0]:+.2f}" if len(sorted_rewards) > 0 else "0.00"
-        top2 = f"{sorted_rewards[1]:+.2f}" if len(sorted_rewards) > 1 else "0.00"
-        top3 = f"{sorted_rewards[2]:+.2f}" if len(sorted_rewards) > 2 else "0.00"
 
         vx = cmd_vel[0] if len(cmd_vel) > 0 else 0.0
         vy = cmd_vel[1] if len(cmd_vel) > 1 else 0.0
-        omega = cmd_vel[3] if len(cmd_vel) > 3 else 0.0
+        omega = cmd_vel[2] if len(cmd_vel) > 2 else 0.0
 
-        hud_text = (
-            f"=== JACKBOT METRICS ===\n"
-            f"Episode:     {episode}\n"
-            f"Global Step: {step}\n"
-            f"FPS:         {fps:.1f}\n"
-            f"----------------------\n"
-            f"Top Rewards: [{top1}, {top2}, {top3}]\n"
-            f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]\n"
-            f"Height:      {avg_height:.3f} m\n"
-            f"Roll/Pitch:  {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°"
-        )
+        lines = [
+            "=== JACKBOT TELEMETRY ===",
+            f"Mode:        {mode.upper()}",
+            f"Curriculum:  {phase}",
+            f"Status:      {status}",
+            f"Episode:     {episode} (Step {ep_step})",
+            f"Global Step: {step}",
+            f"FPS:         {fps:.1f}",
+            "-------------------------",
+            f"Episode Rew: {top1}",
+            f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]",
+            f"Height:      {height:.3f} m",
+            f"Roll/Pitch:  {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°",
+            f"Max Dist:    {distance:.2f} m",
+        ]
+
+        if reward_components:
+            lin_v = reward_components.get("lin_vel", 0.0)
+            stab = reward_components.get("stability", 0.0)
+            h_rew = reward_components.get("height", 0.0)
+            jit = reward_components.get("jitter_penalty", 0.0)
+            lines.append("--- Reward Components ---")
+            lines.append(f"LinVel: {lin_v:.2f} | Stab: {stab:.2f}")
+            lines.append(f"Height: {h_rew:.2f} | Jitter: {jit:+.3f}")
+
+        hud_text = "\n".join(lines)
 
         # Position above origin in simulation world
         text_position = [-0.8, -0.8, 1.2]
@@ -76,69 +94,25 @@ class MetricsHUD:
         # Calculate dynamic orientation to align text flat against camera plane
         text_orientation = self._get_camera_facing_orientation()
 
-        if self._text_id is None:
-            self._text_id = p.addUserDebugText(
-                text=hud_text,
-                textPosition=text_position,
-                textColorRGB=text_color,
-                textSize=0.1,
-                textOrientation=text_orientation,
-                physicsClientId=self.client_id
-            )
-        else:
-            self._text_id = p.addUserDebugText(
-                text=hud_text,
-                textPosition=text_position,
-                textColorRGB=text_color,
-                textSize=0.1,
-                textOrientation=text_orientation,
-                replaceItemUniqueId=self._text_id,
-                physicsClientId=self.client_id
-            )
+        # Safely remove old text to prevent PyBullet ghosting/overlapping
+        if self._text_id is not None:
+            try:
+                p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
+            except Exception:
+                pass
+                
+        # Draw fresh text
+        self._text_id = p.addUserDebugText(
+            text=hud_text,
+            textPosition=text_position,
+            textColorRGB=text_color,
+            textSize=0.085,
+            textOrientation=text_orientation,
+            physicsClientId=self.client_id
+        )
 
     def reset(self) -> None:
         """Removes the active debug text item so a new episode starts with a clean overlay."""
-        if self._text_id is not None:
-            try:
-                p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
-            except Exception:
-                pass
-        self._text_id = None
-
-
-class LeaderCrown:
-    """Renders a floating crown or star emoji above the leading robot in PyBullet."""
-
-    def __init__(self, physics_client_id: int = 0):
-        self.client_id = physics_client_id
-        self._text_id = None
-
-    def update(self, leader_pos: list[float]):
-        """Positions a floating crown ~0.35m directly above the lead robot's base."""
-        crown_pos = [leader_pos[0], leader_pos[1], leader_pos[2] + 0.35]
-        
-        # You can use "👑 CROWN", "⭐ LEADER", or "★ TOP1"
-        crown_text = "👑"
-
-        if self._text_id is None:
-            self._text_id = p.addUserDebugText(
-                text=crown_text,
-                textPosition=crown_pos,
-                textColorRGB=[1.0, 0.84, 0.0],
-                textSize=2.0,
-                physicsClientId=self.client_id
-            )
-        else:
-            self._text_id = p.addUserDebugText(
-                text=crown_text,
-                textPosition=crown_pos,
-                textColorRGB=[1.0, 0.84, 0.0],
-                textSize=2.0,
-                replaceItemUniqueId=self._text_id,
-                physicsClientId=self.client_id
-            )
-
-    def reset(self):
         if self._text_id is not None:
             try:
                 p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)

ml/SimManager.py

@@ -1,87 +0,0 @@
-"""
-ml/SimManager.py - PyBullet Simulation & Multi-Body Manager
-"""
-from typing import Dict, List, Tuple
-import pybullet as p
-import pybullet_data
-import DataTypes as dt
-
-
-class SimManager:
-    """Manages PyBullet simulation lifecycle and multi-robot physics."""
-
-    def __init__(self, use_gui: bool = True):
-        self.use_gui = use_gui
-        self.physics_client = None
-        self.robot_joints: Dict[int, List[int]] = {}
-
-    def connect(self):
-        """Connects to PyBullet and hides side GUI panels."""
-        if self.physics_client is not None and p.isConnected(self.physics_client):
-            return
-
-        flags = p.GUI if self.use_gui else p.DIRECT
-        self.physics_client = p.connect(flags)
-
-        p.setAdditionalSearchPath(pybullet_data.getDataPath())
-        p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client)
-
-        if self.use_gui:
-            # Disable PyBullet side panel and preview windows
-            p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.physics_client)
-            p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self.physics_client)
-            p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
-            p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
-
-    def load_scene(
-        self, urdf_path: str, robot_spacing: float, base_pos_fn
-    ) -> Tuple[int, List[int], List[List[int]]]:
-        """Loads the plane and a single hexapod body into the simulation scene."""
-        plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
-        robots = []
-        robot_joint_indices = []
-        self.robot_joints.clear()
-
-        base_pos = base_pos_fn(0, robot_spacing)
-        robot = p.loadURDF(
-            urdf_path,
-            basePosition=base_pos,
-            useFixedBase=False,
-            physicsClientId=self.physics_client
-        )
-        robots.append(robot)
-
-        joint_indices = [
-            i for i in range(p.getNumJoints(robot, physicsClientId=self.physics_client))
-            if p.getJointInfo(robot, i, physicsClientId=self.physics_client)[2] == p.JOINT_REVOLUTE
-        ]
-        robot_joint_indices.append(joint_indices)
-        self.robot_joints[robot] = joint_indices
-
-        return plane_id, robots, robot_joint_indices
-
-    def updatePosForBody(self, body_id: int, current_rad: dt.RadArray):
-        """Sets joint motor position targets on individual robot bodies."""
-        if body_id not in self.robot_joints:
-            return
-            
-        joint_indices = self.robot_joints[body_id]
-        radflat = current_rad.data.flatten()
-        
-        for joint_index, target_angle in zip(joint_indices, radflat):
-            p.setJointMotorControl2(
-                bodyIndex=body_id,
-                jointIndex=joint_index,
-                controlMode=p.POSITION_CONTROL,
-                targetPosition=float(target_angle),
-                force=250,
-                physicsClientId=self.physics_client
-            )
-
-    def step(self):
-        p.stepSimulation(physicsClientId=self.physics_client)
-
-    def disconnect(self):
-        if self.physics_client is not None and p.isConnected(self.physics_client):
-            p.disconnect(self.physics_client)
-            self.physics_client = None

ml/__init__.py

@@ -1,6 +1,3 @@
-from .env import JackBotEnv
-from .model import ActorCritic
-from .train import train
-from .evaluate import evaluate
+from .env import JackBotEnv, CurriculumPhase
 
-__all__ = ["JackBotEnv", "ActorCritic", "train", "evaluate"]
+__all__ = ["JackBotEnv", "CurriculumPhase"]

ml/callbacks.py

@@ -0,0 +1,102 @@
+"""
+ml/callbacks.py - Stable-Baselines3 callbacks for training diagnostics.
+
+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
+from stable_baselines3.common.callbacks import BaseCallback
+
+
+class RewardLoggerCallback(BaseCallback):
+    """
+    Logs individual reward component averages to TensorBoard and prints 
+    the best worker's performance breakdown to the console per iteration.
+    """
+
+    def __init__(self, verbose: int = 1):
+        super().__init__(verbose)
+        self.iteration = 0
+
+    def _on_step(self) -> bool:
+        return True
+
+    def _on_rollout_end(self) -> None:
+        self.iteration += 1
+        if self.training_env is None:
+            return
+
+        try:
+            # Safely query method across all parallel worker processes
+            all_worker_averages = self.training_env.env_method("get_reward_component_averages")
+        except Exception:
+            return
+
+        if not all_worker_averages or len(all_worker_averages) == 0:
+            return
+
+        # 1. Log mean component values across ALL workers to TensorBoard
+        component_keys = all_worker_averages[0].keys()
+        for key in component_keys:
+            mean_val = float(np.mean([w.get(key, 0.0) for w in all_worker_averages]))
+            self.logger.record(f"reward_components/{key}", mean_val)
+
+        # 2. Identify the best performing worker of this iteration
+        worker_totals = [sum(w.values()) for w in all_worker_averages]
+        best_worker_idx = int(np.argmax(worker_totals))
+        best_averages = all_worker_averages[best_worker_idx]
+        best_total = worker_totals[best_worker_idx]
+
+        # 3. Print best worker breakdown to console
+        if self.verbose > 0:
+            print(f"\n" + "=" * 65)
+            print(f" ITERATION {self.iteration} | BEST WORKER (#{best_worker_idx}) REWARD BREAKDOWN")
+            print(f" Total Avg Reward / Step: {best_total:+.4f}")
+            print("-" * 65)
+            for key, val in best_averages.items():
+                print(f"   • {key:<26}: {val:+.5f}")
+            print("=" * 65 + "\n")
+
+
+class CurriculumCallback(BaseCallback):
+    """
+    Monitors training metrics using SB3's native ep_info_buffer and 
+    dynamically advances curriculum phases across worker processes.
+    """
+
+    def __init__(self, reward_threshold: float = 100.0, verbose: int = 1):
+        super().__init__(verbose)
+        self.reward_threshold = reward_threshold
+
+    def _on_step(self) -> bool:
+        return True
+
+    def _on_rollout_end(self) -> None:
+        if self.training_env is None:
+            return
+
+        # SB3 natively records finished episode stats in self.model.ep_info_buffer
+        if hasattr(self.model, "ep_info_buffer") and len(self.model.ep_info_buffer) > 0:
+            recent_rewards = [ep_info["r"] for ep_info in self.model.ep_info_buffer]
+            mean_reward = float(np.mean(recent_rewards[-50:]))
+
+            try:
+                # Query current phase from worker 0
+                phases = self.training_env.get_attr("curriculum_phase")
+                current_phase = phases[0]
+
+                # Advance curriculum if mean reward exceeds threshold
+                if mean_reward >= self.reward_threshold:
+                    if hasattr(current_phase, "next"):
+                        next_phase = current_phase.next()
+                        if next_phase != current_phase:
+                            self.training_env.set_attr("curriculum_phase", next_phase)
+                            if self.verbose > 0:
+                                print(
+                                    f"\n[Curriculum] 🚀 Promoted workers to phase: {next_phase.name} "
+                                    f"(Mean Reward: {mean_reward:.2f})"
+                                )
+            except Exception:
+                pass  # Keep rollout loop running safely if phase check fails

ml/env.py

@@ -1,509 +1,456 @@
 """
-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
+from enum import IntEnum
 from typing import Optional, Tuple, Dict, Any, List
+from collections import defaultdict
 
 import gymnasium as gym
 from gymnasium import spaces
 import numpy as np
-import pybullet as p
 
 from config import cfg
+from simulation import Simulation
 from Robot import Robot, PyBulletBackend
-from ml.SimManager import SimManager
-from ml.MetricsOverlay import MetricsHUD, LeaderCrown
+from ml.MetricsOverlay import MetricsHUD
 
-# Color Palette RGBA for Terminated/Failed Robots
-COLOR_FAILED = [0.3, 0.3, 0.3, 0.6]   # Collapsed / Tilted Robot (Dark Semi-Transparent Gray)
+
+class CurriculumPhase(IntEnum):
+    STAND_ONLY = 0
+    FORWARD = 1
+    TURN_AND_DIRECTION = 2
+    OMNI_DIRECTION = 3
+    FULL_COMMAND = 4
 
 
 class JackBotEnv(gym.Env):
-    """Gymnasium environment wrapping JackBot hexapods with floating text & crown overlay."""
+    """Gymnasium environment wrapping JackBot hexapod simulation."""
 
     def __init__(
         self,
         use_gui: bool = True,
         random_command: bool = True,
-        robot_spacing: float = 0.5,
-        start_pose: str = "init_deg",
-        max_episode_steps: int = 5000,
+        max_episode_steps: int = 3000,
         urdf_path: str = cfg.urdf_path,
+        robot_mode: str = "direct",  # "direct", "residual", or "kinematics"
     ):
         super().__init__()
+        self.robot_mode = robot_mode
         self.use_gui = use_gui
         self.random_command = random_command
-        self.robot_spacing = robot_spacing
-        self.start_pose = start_pose
         self.max_episode_steps = max_episode_steps
         self.urdf_path = urdf_path
+        self.max_robot_speed = 0.6
 
-        self.episode_count = 0
         self.step_count = 0
         self.total_steps = 0
         self.cumulative_reward = 0.0
-        self.robot_rewards = [0.0]
-        self.failed_robots_mask = [False]
+        self.robot_reward = 0.0
+        self.is_failed = False
+
+        self.episode_count = 0
+        self.episode_height_sum = 0.0
+        self.episode_roll_sum = 0.0
+        self.episode_pitch_sum = 0.0
+        self._curriculum_advanced = False
         self._first_reset = True
 
-        # Initialize Simulation Manager
-        self.sim_manager = SimManager(use_gui=self.use_gui)
-        self.sim_manager.connect()
+        self.last_reward_components: Dict[str, float] = {}
+        self.episode_reward_components_sum: Dict[str, float] = defaultdict(float)
 
-        # Connect physics world
-        self.plane, self.pb_robots, self.robot_joint_indices = self.sim_manager.load_scene(
-            self.urdf_path, self.robot_spacing, self._robot_base_position
-        )
+        self.control_freq = 60
+        self.min_cmd_hold_steps = int(2.0 * self.control_freq)
+        self.max_cmd_hold_steps = int(6.0 * self.control_freq)
+        self.next_cmd_resample_step = 0
+        self.initial_stand_steps = 120
 
-        # Instantiate Robot Python wrappers per PyBullet body ID
-        self.robots = [
-            Robot(
-                backend_type=PyBulletBackend(self.sim_manager, body_id=pb_id),
-                start_pose=self.start_pose,
-                urdf_path=self.urdf_path
-            )
-            for pb_id in self.pb_robots
-        ]
+        # --- INITIALIZE PYBULLET SIMULATION ENGINE ---
+        self.sim = Simulation(urdf_path=self.urdf_path, use_gui=self.use_gui)
+        self.plane_id, self.pb_robot, self.revolute_joints = self.sim.load_scene()
+
+        # --- INITIALIZE ROBOT WITH BACKEND ---
+        self.backend = PyBulletBackend(self.sim)
+        self.robot = Robot(backend_type=self.backend, urdf_path=self.urdf_path, mode=self.robot_mode)
 
-        # Action (18 joint deltas per robot) & Observation (18 angles + 4 command dims per robot)
         action_dim = 18
-        obs_dim = 18 + 4
+        obs_dim = 18 + 3  # 18 Joint Angles + 3 Command Inputs [vx, vy, omega]
 
         self.action_space = spaces.Box(-1.0, 1.0, shape=(action_dim,), dtype=np.float32)
         self.observation_space = spaces.Box(-np.inf, np.inf, shape=(obs_dim,), dtype=np.float32)
+        self.min_joint_limits, self.max_joint_limits = self.sim._get_urdf_joint_limits()
 
-        self.commands = np.zeros((1, 4), dtype=np.float32)
+        self.command = np.zeros(3, dtype=np.float32)  # [vx, vy, omega]
         self.last_action = np.zeros(action_dim, dtype=np.float32)
-        self.target_height = 0.14
+        self.last_last_action = np.zeros(action_dim, dtype=np.float32)
+        self.target_height = 0.122
         self.collapse_height_fraction = 0.55
         self.tilt_failure_rad = 0.9
-        self.start_positions = [[0.0, 0.0, 0.0]]
-        self.max_distance_from_start = [0.0]
+
+        self.start_position = [0.0, 0.0, 0.0]
+        self.max_distance_from_start = 0.0
         self.max_survival_steps = 0
-        self.curriculum_phase = 0
-        self.curriculum_episode_limit = 150
+        self.default_joint_angles = np.zeros(18, dtype=np.float32)
+
+        self.curriculum_phase = CurriculumPhase.STAND_ONLY
         self.curriculum_stage_requirements = {
-            1: {"survival_steps": 200, "distance": 0.00, "stability_roll_pitch": 0.35},
-            2: {"survival_steps": 350, "distance": 0.50, "stability_roll_pitch": 0.30},
-            3: {"survival_steps": 550, "distance": 0.80, "stability_roll_pitch": 0.25},
+            CurriculumPhase.FORWARD: {"survival_steps": 300, "min_avg_height_ratio": 0.88, "max_avg_roll_pitch": 0.18},
+            CurriculumPhase.TURN_AND_DIRECTION: {"survival_steps": 500, "min_forward_distance": 2.5, "max_lateral_drift": 0.8, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25},
+            CurriculumPhase.OMNI_DIRECTION: {"survival_steps": 600, "min_distance": 5.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25},
+            CurriculumPhase.FULL_COMMAND: {"survival_steps": 750, "min_distance": 8.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.20},
         }
 
-        # Floating HUD & Leader Crown Visualizers
-        self.hud = MetricsHUD(physics_client_id=self.sim_manager.physics_client)
-        self.leader_crown = LeaderCrown(physics_client_id=self.sim_manager.physics_client)
+        self.hud = MetricsHUD(physics_client_id=self.sim.physics_client)
         self.last_time = time.time()
 
-    def _set_robot_color(self, pb_id: int, rgba: List[float]):
-        """Helper to change the visual color of a robot body and all its links."""
-        num_joints = p.getNumJoints(pb_id, physicsClientId=self.sim_manager.physics_client)
-        p.changeVisualShape(pb_id, -1, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client)
-        for j in range(num_joints):
-            p.changeVisualShape(pb_id, j, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client)
-
-    def _robot_base_position(self, robot_id: int, spacing: float = 0.5) -> list[float]:
-        return [0.0, 0.0, 0.14]
-
     def sample_command(self) -> np.ndarray:
-        """Curriculum command sampler with survival-gated difficulty progression."""
+        """Samples a command vector [vx, vy, omega] based on active curriculum phase."""
         phase = self.curriculum_phase
+        stand_probs = {
+            CurriculumPhase.STAND_ONLY: 1.0,
+            CurriculumPhase.FORWARD: 0.25,
+            CurriculumPhase.TURN_AND_DIRECTION: 0.20,
+            CurriculumPhase.OMNI_DIRECTION: 0.15,
+            CurriculumPhase.FULL_COMMAND: 0.15,
+        }
 
-        if phase == 0:
-            # Phase 1: Forward Walking Focus
-            vx = np.random.uniform(0.5, 1.0)
-            vy = 0.0
-            vz = 0.0
-            omega = 0.0
-        elif phase == 1:
-            # Phase 2: Forward/Backward + Turning
-            vx = np.random.uniform(-1.0, 1.0)
-            vy = 0.0
-            vz = 0.0
-            omega = np.random.uniform(-0.8, 0.8)
+        if np.random.random() < stand_probs.get(phase, 0.15):
+            return np.zeros(3, dtype=np.float32)
+
+        if phase == CurriculumPhase.FORWARD:
+            vx, vy, omega = np.random.uniform(0.15, 0.50), 0.0, 0.0
+        elif phase == CurriculumPhase.TURN_AND_DIRECTION:
+            vx, vy, omega = np.random.uniform(-0.8, 0.8), 0.0, np.random.uniform(-0.8, 0.8)
+        elif phase == CurriculumPhase.OMNI_DIRECTION:
+            vx, vy, omega = np.random.uniform(-0.8, 0.8), np.random.uniform(-0.5, 0.5), np.random.uniform(-0.8, 0.8)
         else:
-            # Phase 3: Full Omnidirectional Movement
-            vx = np.random.uniform(-1.0, 1.0)
-            vy = np.random.uniform(-0.5, 0.5)
-            vz = 0.0
-            omega = np.random.uniform(-1.0, 1.0)
+            vx, vy, omega = np.random.uniform(-1.0, 1.0), np.random.uniform(-1.0, 1.0), np.random.uniform(-1.0, 1.0)
 
-        return np.array([vx, vy, vz, omega], dtype=np.float32)
+        return np.array([vx, vy, omega], dtype=np.float32)
 
     def reset(self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None):
         super().reset(seed=seed)
         self.episode_count += 1
         self.step_count = 0
         self.cumulative_reward = 0.0
-        self.robot_rewards = [0.0]
-        self.failed_robots_mask = [False]
+        self.robot_reward = 0.0
+        self.is_failed = False
 
-        for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
-            spawn_pos = self._robot_base_position(idx, self.robot_spacing)
-            spawn_orn = [0, 0, 0, 1]
+        self.episode_height_sum = 0.0
+        self.episode_roll_sum = 0.0
+        self.episode_pitch_sum = 0.0
+        self.last_reward_components = {}
+        self.episode_reward_components_sum = defaultdict(float)
 
-            p.resetBasePositionAndOrientation(
-                pb_id, spawn_pos, spawn_orn, physicsClientId=self.sim_manager.physics_client
-            )
-            p.resetBaseVelocity(
-                pb_id, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 0],
-                physicsClientId=self.sim_manager.physics_client
-            )
+        spawn_pos = [0.0, 0.0, 0.15]
+        spawn_orn = [0.0, 0.0, 0.0, 1.0]
 
-            robot_obj.reset_to_init()
+        # 1. Reset base pose and velocities
+        self.sim.reset_robot_base(spawn_pos, spawn_orn)
 
-            if self.use_gui:
-                self._set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
+        # 2. Reset internal kinematics & hard reset joints in PyBullet
+        self.robot.reset_to_init()
 
-        self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32)
-        self.max_distance_from_start = [0.0]
+        action_dim = self.action_space.shape[0]
+        self.last_action = np.zeros(action_dim, dtype=np.float32)
+        self.last_last_action = np.zeros(action_dim, dtype=np.float32)
+        self.max_distance_from_start = 0.0
         self.max_survival_steps = 0
 
         if self._first_reset:
-            self.curriculum_phase = 0
-            self.curriculum_episode_limit = min(self.max_episode_steps, 150)
+            self.curriculum_phase = CurriculumPhase.STAND_ONLY
             self._first_reset = False
 
-        if self.random_command:
-            self.commands = np.stack([self.sample_command() for _ in range(1)])
-        else:
-            self.commands = np.zeros((1, 4), dtype=np.float32)
+        self.command = np.zeros(3, dtype=np.float32)
+        self.next_cmd_resample_step = self.initial_stand_steps
 
-        for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
-            pos, _ = p.getBasePositionAndOrientation(
-                pb_id, physicsClientId=self.sim_manager.physics_client
-            )
-            self.start_positions[idx] = [float(pos[0]), float(pos[1]), float(pos[2])]
+        pos, _ = self.sim.get_robot_pose()
+        self.start_position = list(pos)
 
-        for _ in range(200):
-            self.sim_manager.step()
-
-        self.target_height = self._measure_settled_height()
-        if self.target_height <= 0.0:
-            self.target_height = 0.14
+        # Drop settlement
+        self.target_height = self.sim.settle_and_measure_height(
+            target_angles=self.robot.current_rad, 
+            steps=300, 
+            fallback_height=0.122
+        )
+        self.default_joint_angles = self.sim.get_robot_joint_angles()
+        self.default_action = np.clip(self.default_joint_angles / (np.pi / 2.0), -1.0, 1.0)
 
         if self.use_gui:
             self.hud.reset()
-            self.leader_crown.reset()
             self._update_hud()
 
         return self._get_obs(), {}
 
-    def get_robot_velocities(self) -> list:
-            """Exposes velocities for the SB3 metrics callback."""
-            vels = []
-            for pb_id in self.pb_robots:
-                lin_v, ang_v = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client)
-                vels.append((lin_v, ang_v))
-            return vels
-
-    def get_robot_distance_metrics(self) -> list:
-            """Exposes distance-from-start metrics for logging without affecting reward."""
-            metrics = []
-            for idx, pb_id in enumerate(self.pb_robots):
-                pos, _ = p.getBasePositionAndOrientation(
-                    pb_id, physicsClientId=self.sim_manager.physics_client
-                )
-                start_x, start_y, _ = self.start_positions[idx]
-                dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
-                self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist)
-                metrics.append((dist, self.max_distance_from_start[idx]))
-
-            self.max_survival_steps = max(self.max_survival_steps, self.step_count)
-            return metrics
-
-    def get_current_robot_metrics(self) -> list:
-            """Returns current per-robot reward, distance, and velocity summaries for alive robots only."""
-            metrics = []
-            for idx, pb_id in enumerate(self.pb_robots):
-                if self.failed_robots_mask[idx]:
-                    continue
-
-                pos, _ = p.getBasePositionAndOrientation(
-                    pb_id, physicsClientId=self.sim_manager.physics_client
-                )
-                linear_vel, angular_vel = p.getBaseVelocity(
-                    pb_id, physicsClientId=self.sim_manager.physics_client
-                )
-                start_x, start_y, _ = self.start_positions[idx]
-                dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
-                speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)))
-                yaw_rate = float(abs(angular_vel[2]))
-                metrics.append({
-                    "reward": float(self.robot_rewards[idx]),
-                    "distance_from_start": dist,
-                    "speed": speed,
-                    "yaw_rate": yaw_rate,
-                    "alive": True,
-                    "survival_steps": int(self.step_count),
-                })
-
-            return metrics
-
-    def get_survival_steps(self) -> int:
-            """Returns the current survival length for the environment's current episode."""
-            return int(self.step_count)
-
     def _get_obs(self) -> np.ndarray:
-        obs_list = []
-        for idx, (pb_id, joint_indices) in enumerate(zip(self.pb_robots, self.robot_joint_indices)):
-            joint_states = p.getJointStates(
-                pb_id, 
-                joint_indices, 
-                physicsClientId=self.sim_manager.physics_client
-            )
-            joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32)
-            robot_obs = np.concatenate([joint_angles, self.commands[idx]])
-            obs_list.append(robot_obs)
+        # Read raw joint angles from backend
+        raw_angles = np.asarray(self.robot.backend.get_joint_angles(), dtype=np.float32).flatten()
 
-        return np.concatenate(obs_list).astype(np.float32)
+        min_lim = self.min_joint_limits.flatten()
+        max_lim = self.max_joint_limits.flatten()
 
-    def _measure_settled_height(self) -> float:
-        heights = []
-        for pb_id in self.pb_robots:
-            pos, _ = p.getBasePositionAndOrientation(
-                pb_id, physicsClientId=self.sim_manager.physics_client
-            )
-            heights.append(pos[2])
-        return float(np.mean(heights)) if heights else 0.14
+        # 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()
         self.step_count += 1
         self.total_steps += 1
-        previous_action = self.last_action.copy()
+
+        self.last_last_action = self.last_action.copy()
         self.last_action = action.copy()
+
+        # 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):
+            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 active [vx, vy, omega]
+        cmd_vx, cmd_vy, cmd_omega = self.command
+
+        # 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)]
+
+        # Direct Mode: Target joint scaling
+        action_flat = np.asarray(action, dtype=np.float32).flatten()
+        action_clipped = np.clip(action_flat, -1.0, 1.0)
+
+        if self.robot_mode == "direct":
+            # 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)
+            self.sim.apply_external_force(force=[random_force[0], random_force[1], 0.0])
+
+        self._update_robot_failure()
+        self._update_distance_metrics()
         self._update_curriculum()
 
-        # Resample commands every 300 steps during long episodes
-        if self.random_command and (self.step_count % 300 == 0 or self._curriculum_advanced):
-            self.commands = np.stack([self.sample_command() for _ in range(1)])
-
-        action_per_robot = action.reshape(1, 18)
-
-        for robot, act in zip(self.robots, action_per_robot):
-            robot.apply_rl_action(act)
-
-        self.sim_manager.step()
-
-        self._update_robot_failures()
-        self.get_robot_distance_metrics()
-
+        # Build next observation preserving active command
         obs = self._get_obs()
-        reward, per_robot_step_rewards = self._compute_reward(action, previous_action)
-        
+        reward = self._compute_reward(action_flat, previous_action)
+
         self.cumulative_reward += reward
-        for idx, r_step in enumerate(per_robot_step_rewards):
-            self.robot_rewards[idx] += r_step
+        self.robot_reward += reward
 
-        terminated = self._is_done()
-        truncated = self.step_count >= self.curriculum_episode_limit
+        terminated = self.is_failed
+        truncated = self.step_count >= self.max_episode_steps
+        info = {"reward_components": self.last_reward_components.copy()}
 
-        self._update_hud()
-        self._update_leader_visuals()
-        return obs, reward, terminated, truncated, {}
+        if self.step_count % 120 == 0 and self.use_gui:
+            self._update_hud()
 
-    def _phase_progress_ready(self, phase: int) -> bool:
-        if phase not in self.curriculum_stage_requirements:
+        if self.use_gui:
+            time.sleep(1.0 / self.control_freq)
+
+        return obs, reward, terminated, truncated, info
+
+    def _update_distance_metrics(self):
+        pos, _ = self.sim.get_robot_pose()
+        self.episode_height_sum += float(pos[2])
+        start_x, start_y, _ = self.start_position
+        dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
+        self.max_distance_from_start = max(self.max_distance_from_start, dist)
+        self.max_survival_steps = max(self.max_survival_steps, self.step_count)
+
+    def _phase_progress_ready(self, next_phase: CurriculumPhase) -> bool:
+        if next_phase not in self.curriculum_stage_requirements:
             return False
 
-        if not self.pb_robots or not self.max_distance_from_start:
-            return False
-
-        req = self.curriculum_stage_requirements[phase]
+        req = self.curriculum_stage_requirements[next_phase]
         survival_ok = self.max_survival_steps >= req["survival_steps"]
-        distance_ok = self.max_distance_from_start[0] >= req["distance"]
+        avg_roll = self.episode_roll_sum / max(1, self.step_count)
+        avg_pitch = self.episode_pitch_sum / max(1, self.step_count)
+        max_allowed_angle = req.get("max_avg_roll_pitch", 0.20)
+        stability_ok = (avg_roll <= max_allowed_angle) and (avg_pitch <= max_allowed_angle)
 
-        position, orientation = p.getBasePositionAndOrientation(
-            self.pb_robots[0], physicsClientId=self.sim_manager.physics_client
-        )
-        roll, pitch, _ = p.getEulerFromQuaternion(orientation)
-        stability_ok = abs(roll) <= req["stability_roll_pitch"] and abs(pitch) <= req["stability_roll_pitch"]
-        height_ok = position[2] >= max(0.09, self.target_height * 0.85)
+        avg_height = self.episode_height_sum / max(1, self.step_count)
+        required_min_avg_height = self.target_height * req.get("min_avg_height_ratio", 0.85)
+        height_ok = avg_height >= required_min_avg_height
 
-        return survival_ok and distance_ok and stability_ok and height_ok
+        pos, _ = self.sim.get_robot_pose()
+        dx, dy = pos[0] - self.start_position[0], pos[1] - self.start_position[1]
+        dist_2d = math.hypot(dx, dy)
+
+        distance_ok = True
+        if "min_forward_distance" in req:
+            distance_ok = dx >= req["min_forward_distance"]
+        elif "min_distance" in req:
+            distance_ok = dist_2d >= req["min_distance"]
+
+        drift_ok = abs(dy) <= req["max_lateral_drift"] if "max_lateral_drift" in req else True
+        return survival_ok and height_ok and stability_ok and distance_ok and drift_ok
 
     def _update_curriculum(self):
         self._curriculum_advanced = False
 
-        phase_labels = {
-            0: "stand-and-forward",
-            1: "turn-and-direction",
-            2: "omni-direction",
-            3: "full-command",
+        if self.curriculum_phase < CurriculumPhase.FORWARD and (self._phase_progress_ready(CurriculumPhase.FORWARD)):
+            self.curriculum_phase = CurriculumPhase.FORWARD
+            self._curriculum_advanced = True
+        elif self.curriculum_phase < CurriculumPhase.TURN_AND_DIRECTION and self._phase_progress_ready(CurriculumPhase.TURN_AND_DIRECTION):
+            self.curriculum_phase = CurriculumPhase.TURN_AND_DIRECTION
+            self._curriculum_advanced = True
+        elif self.curriculum_phase < CurriculumPhase.OMNI_DIRECTION and self._phase_progress_ready(CurriculumPhase.OMNI_DIRECTION):
+            self.curriculum_phase = CurriculumPhase.OMNI_DIRECTION
+            self._curriculum_advanced = True
+        elif self.curriculum_phase < CurriculumPhase.FULL_COMMAND and self._phase_progress_ready(CurriculumPhase.FULL_COMMAND):
+            self.curriculum_phase = CurriculumPhase.FULL_COMMAND
+            self._curriculum_advanced = True
+
+    def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> float:
+        pos, (roll, pitch, yaw) = self.sim.get_robot_pose_and_rpy()
+        linear_vel, angular_vel = self.sim.get_robot_velocity()
+        current_joints = self.sim.get_robot_joint_angles()
+
+        cmd_vx, cmd_vy, cmd_yaw = self.command
+        cmd_norm = math.hypot(cmd_vx, cmd_vy)
+
+        raw_speed = math.hypot(linear_vel[0], linear_vel[1])
+        # Forgiving zones: ignore noise below 0.04 m/s and 0.05 rad/s
+        filtered_vx, filtered_vy = (linear_vel[0], linear_vel[1]) if raw_speed >= 0.04 else (0.0, 0.0)
+        filtered_speed = raw_speed if raw_speed >= 0.04 else 0.0
+        raw_yaw_rate = abs(angular_vel[2])
+        filtered_yaw_rate = raw_yaw_rate if raw_yaw_rate >= 0.05 else 0.0
+
+        # --- 1. FIXED JITTER PENALTY ---
+        # Scaled way down (0.005) and capped so it can never dominate the reward
+        action_accel = action - 2.0 * self.last_action + self.last_last_action
+        raw_jitter = float(np.mean(np.square(action_accel)))
+        jitter_penalty = -0.005 * min(raw_jitter, 10.0)
+
+        # --- Base Components ---
+        height_error = pos[2] - self.target_height
+        r_height = math.exp(-150.0 * (height_error ** 2))
+        r_stability = math.exp(-25.0 * (roll**2 + pitch**2))
+        r_pose = math.exp(-2.0 * np.mean(np.square(current_joints - self.default_joint_angles)))
+        r_smoothness = math.exp(-0.1 * np.mean(np.square(action - previous_action)))
+
+        r_lin_vel, r_ang_vel, stillness_penalty = 0.0, 0.0, 0.0
+        stand_penalty = 0.0
+
+        # --- 2. COMMAND IS ZERO: STANDING MODE ---
+        if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05:
+            w_height, w_stability, w_pose, w_smoothness = 0.35, 0.35, 0.20, 0.10
+            base_reward = (w_height * r_height) + (w_stability * r_stability) + (w_pose * r_pose) + (w_smoothness * r_smoothness)
+
+            # Soft quadratic penalty on filtered speed (gives a forgiving gap near 0)
+            stand_penalty = -0.5 * filtered_speed - 0.1 * filtered_yaw_rate
+            total_reward = base_reward + stand_penalty
+
+        # --- 3. COMMAND IS NON-ZERO: WALKING MODE ---
+        else:
+            is_moving = (filtered_speed > 0.0) or (filtered_yaw_rate > 0.0)
+            target_vx, target_vy = cmd_vx * self.max_robot_speed, cmd_vy * self.max_robot_speed
+            target_speed = math.hypot(target_vx, target_vy)
+
+            if not is_moving:
+                # 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)
+                r_ang_vel = math.exp(-15.0 * ((filtered_yaw_rate - cmd_yaw)**2))
+
+                if target_speed > 0.08 and raw_speed < 0.03:
+                    r_lin_vel = 0.0
+                    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 = (
+                    (w_lin_vel * r_lin_vel) + (w_ang_vel * r_ang_vel) + (w_height * r_height)
+                    + (w_stability * r_stability) + (w_smoothness * r_smoothness) + stillness_penalty
+                )
+
+        step_reward = float(total_reward / 10.0)
+        alive_bonus = 0.01
+        final_reward = step_reward + jitter_penalty + alive_bonus
+
+        self.last_reward_components = {
+            "height": float(r_height),
+            "stability": float(r_stability),
+            "pose": float(r_pose),
+            "smoothness": float(r_smoothness),
+            "lin_vel": float(r_lin_vel),
+            "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():
+            self.episode_reward_components_sum[k] += v
 
-        if self.curriculum_phase < 1 and self._phase_progress_ready(1):
-            self.curriculum_phase = 1
-            self.curriculum_episode_limit = min(self.max_episode_steps, 400)
-            self._curriculum_advanced = True
-            print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
-        elif self.curriculum_phase < 2 and self._phase_progress_ready(2):
-            self.curriculum_phase = 2
-            self.curriculum_episode_limit = min(self.max_episode_steps, 700)
-            self._curriculum_advanced = True
-            print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
-        elif self.curriculum_phase < 3 and self._phase_progress_ready(3):
-            self.curriculum_phase = 3
-            self.curriculum_episode_limit = min(self.max_episode_steps, self.max_episode_steps)
-            self._curriculum_advanced = True
-            print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
+        return final_reward
+    
+    def get_reward_component_averages(self) -> Dict[str, float]:
+        steps = max(1, self.step_count)
+        return {k: v / steps for k, v in self.episode_reward_components_sum.items()}
 
-    def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> Tuple[float, list[float]]:
-        rewards = []
-        current_actions = action.reshape(1, 18)
-        previous_actions = previous_action.reshape(1, 18)
+    def get_current_robot_metrics(self) -> List[Dict[str, Any]]:
+        linear_vel, angular_vel = self.sim.get_robot_velocity()
+        speed = float(math.hypot(linear_vel[0], linear_vel[1]))
+        yaw_rate = float(abs(angular_vel[2]))
 
-        for idx, pb_id in enumerate(self.pb_robots):
-            pos, orientation = p.getBasePositionAndOrientation(
-                pb_id, physicsClientId=self.sim_manager.physics_client
-            )
-            linear_vel, angular_vel = p.getBaseVelocity(
-                pb_id, physicsClientId=self.sim_manager.physics_client
-            )
-            roll, pitch, _ = p.getEulerFromQuaternion(orientation)
-            command = self.commands[idx]
-
-            cmd_vx = command[0]
-            cmd_vy = command[1]
-            cmd_yaw = command[3]
-
-            # -------------------------------------------------------------
-            # 1. LINEAR VECTOR SPEED MAXIMIZATION (Magnitude + Direction)
-            # -------------------------------------------------------------
-            cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32)
-            cmd_norm = np.linalg.norm(cmd_dir)
-
-            if cmd_norm > 0.05:
-                # Normalize target direction vector
-                unit_cmd_dir = cmd_dir / cmd_norm
-                actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)
-
-                # Speed aligned with target direction (m/s)
-                aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir))
-
-                # Moderate movement reward to encourage directional motion
-                linear_speed_reward = 2.5 * aligned_speed
-
-                # Penalize sideways drift (perpendicular velocity to commanded direction)
-                perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir
-                drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel))
-            else:
-                # If no linear command given, penalize all horizontal movement
-                linear_speed_reward = 0.0
-                drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2)
-
-            # -------------------------------------------------------------
-            # 2. TURNING SPEED MAXIMIZATION
-            # -------------------------------------------------------------
-            actual_yaw_rate = angular_vel[2]  # rad/s in PyBullet Z-axis
-
-            if abs(cmd_yaw) > 0.05:
-                # Reward turning in the commanded direction, but less aggressively
-                turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw)
-            else:
-                # Penalize unwanted rotation when joystick turn is centered
-                turning_reward = -0.6 * (actual_yaw_rate ** 2)
-
-            # -------------------------------------------------------------
-            # 3. SMALL ALIVE BONUS + AGE-RAMPED STILLNESS PENALTY
-            # -------------------------------------------------------------
-            alive_reward = 0.02
-
-            age_ratio = min(1.0, self.step_count / max(1, self.curriculum_episode_limit))
-            still_penalty = 0.0
-            if (cmd_norm > 0.1 or abs(cmd_yaw) > 0.1) and (abs(linear_vel[0]) < 0.02 and abs(actual_yaw_rate) < 0.05):
-                still_penalty = 0.35 + 0.85 * age_ratio
-
-            # -------------------------------------------------------------
-            # 4. POSTURE & STABILITY PENALTIES
-            # -------------------------------------------------------------
-            height_penalty = 4.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0
-            stability_penalty = 1.5 * (roll**2 + pitch**2)
-
-            # Penalize only large control jumps; allow continuous low-amplitude motion
-            control_delta = np.abs(current_actions[idx] - previous_actions[idx])
-            large_delta_mask = control_delta > 0.12
-            large_delta_penalty = 0.002 * float(np.sum(np.square(control_delta[large_delta_mask]))) if np.any(large_delta_mask) else 0.0
-
-            # Combined Step Reward
-            r_step = (
-                alive_reward
-                + linear_speed_reward
-                + turning_reward
-                - drift_penalty
-                - still_penalty
-                - height_penalty
-                - stability_penalty
-                - large_delta_penalty
-            )
-            rewards.append(r_step)
-
-        return float(np.sum(rewards)), rewards
-
-    def _is_done(self) -> bool:
-        """Returns True when the robot has entered a failed state."""
-        return bool(self.failed_robots_mask[0]) if self.failed_robots_mask else False
+        metrics = {
+            "alive": not self.is_failed,
+            "phase_name": self.curriculum_phase.name,
+            "reward": float(self.cumulative_reward),
+            "speed": speed,
+            "yaw_rate": yaw_rate,
+            "distance_from_start": float(self.max_distance_from_start),
+            "survival_steps": int(self.max_survival_steps),
+        }
+        return [metrics]
 
     def _update_hud(self):
-        if not self.use_gui or not self.pb_robots:
+        if not self.use_gui:
             return
-
         now = time.time()
         fps = 1.0 / max(now - self.last_time, 1e-5)
         self.last_time = now
-
-        heights = []
-        rolls = []
-        pitches = []
-        for pb_id in self.pb_robots:
-            pos, orient = p.getBasePositionAndOrientation(pb_id, physicsClientId=self.sim_manager.physics_client)
-            roll, pitch, _ = p.getEulerFromQuaternion(orient)
-            heights.append(pos[2])
-            rolls.append(math.degrees(roll))
-            pitches.append(math.degrees(pitch))
-
-        avg_height = float(np.mean(heights))
-        avg_roll_pitch = (float(np.mean(rolls)), float(np.mean(pitches)))
-
+        pos, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy()
         self.hud.update(
-            episode=self.episode_count,
-            step=self.total_steps,
-            robot_rewards=self.robot_rewards,
-            cmd_vel=self.commands[0],
-            fps=fps,
-            avg_height=avg_height,
-            roll_pitch=avg_roll_pitch
+            episode=self.episode_count, step=self.total_steps, robot_rewards=[self.robot_reward],
+            cmd_vel=self.command, fps=fps, height=pos[2], roll_pitch=(math.degrees(roll), math.degrees(pitch))
         )
 
-    def _update_leader_visuals(self):
-        if not self.use_gui:
+    def _update_robot_failure(self):
+        if self.is_failed:
             return
 
-        best_idx = int(np.argmax(self.robot_rewards))
-        leader_pb_id = self.pb_robots[best_idx]
-        leader_pos, _ = p.getBasePositionAndOrientation(
-            leader_pb_id, physicsClientId=self.sim_manager.physics_client
-        )
-        self.leader_crown.update(leader_pos)
+        position, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy()
+        collapse_threshold = max(0.04, self.collapse_height_fraction * self.target_height)
+        is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad
+        is_collapsed = position[2] < collapse_threshold
 
-    def _update_robot_failures(self):
-        """Checks failure condition and colors failed robots dark gray."""
-        for idx, pb_id in enumerate(self.pb_robots):
-            if self.failed_robots_mask[idx]:
-                continue
-
-            position, orientation = p.getBasePositionAndOrientation(
-                pb_id, physicsClientId=self.sim_manager.physics_client
-            )
-            roll, pitch, _ = p.getEulerFromQuaternion(orientation)
-
-            collapse_threshold = max(0.06, self.collapse_height_fraction * self.target_height)
-            is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad
-            is_collapsed = position[2] < collapse_threshold
-
-            if is_tilted or is_collapsed:
-                self.failed_robots_mask[idx] = True
-                if self.use_gui:
-                    self._set_robot_color(pb_id, COLOR_FAILED)
+        if is_tilted or is_collapsed:
+            self.is_failed = True
 
     def close(self):
-        self.sim_manager.disconnect()
+        self.sim.disconnect()

ml/evaluate.py

@@ -1,92 +0,0 @@
-"""
-ml/evaluate.py - Evaluation routine for trained JackBot PPO policies.
-"""
-import argparse
-import json
-import time
-from pathlib import Path
-from typing import Optional, Dict, List, Any
-import numpy as np
-
-
-def evaluate(
-    model_path: str,
-    episodes: int = 5,
-    use_gui: bool = True,
-    robot_spacing: float = 0.5,
-    start_pose: str = "init_deg",
-    random_command: bool = False,
-    save_json: Optional[str] = None,
-) -> Dict[str, Any]:
-    try:
-        from stable_baselines3 import PPO
-    except ImportError as exc:
-        raise ImportError(
-            "stable-baselines3 is required for evaluation. Install with: pip install stable-baselines3"
-        ) from exc
-
-    from .env import JackBotEnv
-
-    print(f"[Eval] Loading policy model from: {model_path}")
-    # Force device="cpu" to prevent AMD ROCm/hipBLASLt matrix multiplication crashes
-    model = PPO.load(model_path, device="cpu")
-
-    # Initialize standard environment (returns single array of shape (352,))
-    env = JackBotEnv(
-        use_gui=use_gui,
-        random_command=random_command,
-        robot_spacing=robot_spacing,
-        start_pose=start_pose,
-    )
-
-    episode_rewards: List[float] = []
-    episode_lengths: List[int] = []
-
-    for ep in range(episodes):
-        obs, _ = env.reset()
-        done = False
-        total_reward = 0.0
-        steps = 0
-
-        print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} ---")
-
-        while not done:
-            action, _ = model.predict(obs, deterministic=True)
-            obs, reward, terminated, truncated, _ = env.step(action)
-            
-            done = terminated or truncated
-            total_reward += float(reward)
-            steps += 1
-
-            if use_gui:
-                time.sleep(1.0 / 240.0)
-
-        episode_rewards.append(total_reward)
-        episode_lengths.append(steps)
-        print(f"Episode {ep + 1} Finished: Total Reward = {total_reward:.2f} | Steps = {steps}")
-
-    env.close()
-
-    metrics = {
-        "model_path": str(model_path),
-        "episodes_evaluated": episodes,
-        "mean_reward": float(np.mean(episode_rewards)),
-        "std_reward": float(np.std(episode_rewards)),
-        "mean_episode_length": float(np.mean(episode_lengths)),
-        "raw_rewards": episode_rewards,
-    }
-
-    print("\n" + "=" * 50)
-    print(f"EVALUATION COMPLETE ({episodes} Episodes)")
-    print(f"Mean Reward:         {metrics['mean_reward']:.2f} ± {metrics['std_reward']:.2f}")
-    print(f"Mean Episode Length: {metrics['mean_episode_length']:.1f} steps")
-    print("=" * 50)
-
-    if save_json:
-        out_path = Path(save_json)
-        out_path.parent.mkdir(parents=True, exist_ok=True)
-        with open(out_path, "w") as f:
-            json.dump(metrics, f, indent=4)
-        print(f"[Eval] Saved evaluation metrics to: {out_path.resolve()}")
-
-    return metrics

ml/model.py

@@ -1,49 +0,0 @@
-import torch
-import torch.nn as nn
-from torch.distributions import Normal
-
-
-class ActorCritic(nn.Module):
-    def __init__(self, obs_dim: int, action_dim: int, hidden_sizes=(256, 256)):
-        super().__init__()
-        self.backbone = nn.Sequential(
-            nn.Linear(obs_dim, hidden_sizes[0]),
-            nn.ReLU(),
-            nn.Linear(hidden_sizes[0], hidden_sizes[1]),
-            nn.ReLU(),
-        )
-
-        self.mean_head = nn.Linear(hidden_sizes[1], action_dim)
-        self.value_head = nn.Linear(hidden_sizes[1], 1)
-        self.log_std = nn.Parameter(torch.zeros(action_dim, dtype=torch.float32))
-
-    def forward(self, obs: torch.Tensor):
-        x = self.backbone(obs)
-        mean = self.mean_head(x)
-        std = self.log_std.exp()
-        value = self.value_head(x).squeeze(-1)
-        return mean, std, value
-
-    def get_action(self, obs: torch.Tensor):
-        mean, std, value = self.forward(obs)
-        dist = Normal(mean, std)
-        action = dist.sample()
-        log_prob = dist.log_prob(action).sum(-1)
-        return action, log_prob, value
-
-    def evaluate_actions(self, obs: torch.Tensor, actions: torch.Tensor):
-        mean, std, value = self.forward(obs)
-        dist = Normal(mean, std)
-        log_prob = dist.log_prob(actions).sum(-1)
-        entropy = dist.entropy().sum(-1)
-        return value, log_prob, entropy
-
-    def save(self, path: str):
-        torch.save(self.state_dict(), path)
-
-    @classmethod
-    def load(cls, path: str, obs_dim: int, action_dim: int, hidden_sizes=(256, 256)):
-        model = cls(obs_dim, action_dim, hidden_sizes)
-        model.load_state_dict(torch.load(path, map_location=torch.device("cpu")))
-        model.eval()
-        return model

ml/pretrain_bc.py

@@ -0,0 +1,155 @@
+"""
+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
+import argparse
+from pathlib import Path
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, TensorDataset
+from stable_baselines3 import PPO
+from tqdm import tqdm  # <--- Progress Bar Support
+
+# Ensure project root is in sys.path
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+from ml.env import JackBotEnv, CurriculumPhase
+
+
+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})...")
+    
+    # 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 = []
+    actions = []
+
+    obs, _ = env.reset()
+    
+    # --- PROGRESS BAR: Data Collection ---
+    pbar = tqdm(range(num_samples), desc=" Collecting Data", unit="step")
+    for i in pbar:
+        # 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. Capture observation BEFORE stepping environment
+        current_obs = env._get_obs()
+
+        # 3. Step environment ONCE (updates kinematics solver, PyBullet physics, and computes IK)
+        obs, _, terminated, truncated, _ = env.step(np.zeros(18, dtype=np.float32))
+
+        # 4. Extract procedural IK joint targets computed during this step
+        target_ik_rad = env.robot.current_rad.data.flatten().copy()
+
+        # Step 5: Convert target radians directly to [-1, 1] relative to joint limits
+        min_lim = env.min_joint_limits.flatten()
+        max_lim = env.max_joint_limits.flatten()
+
+        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)
+
+        # Step 6: Store matching input (obs) and target ground truth (normalized_action)
+        observations.append(current_obs.copy())
+        actions.append(normalized_action.copy())
+
+        if use_gui:
+            time.sleep(1.0 / 60.0)
+
+        # 7. Handle episode boundaries using terminated and truncated
+        if terminated or truncated:
+            obs, _ = env.reset()
+
+    env.close()
+    print("[Pretrain] Data collection complete!\n")
+    return np.array(observations, dtype=np.float32), np.array(actions, dtype=np.float32)
+
+
+def pretrain_policy(
+    save_path: str = "ml/checkpoints/jackbot_kinematics_base.zip",
+    epochs: int = 15,
+    batch_size: int = 256,
+    num_samples: int = 100_000,
+    use_gui: bool = False
+):
+    # Collect dataset from Kinematics teacher
+    obs_data, action_data = collect_kinematics_dataset(num_samples=num_samples, use_gui=use_gui)
+
+    # Initialize Dummy Env & Fresh SB3 PPO Model
+    dummy_env = JackBotEnv(use_gui=False, robot_mode="direct")
+    model = PPO("MlpPolicy", dummy_env, learning_rate=5e-4, verbose=0, device="cpu")
+
+    # Extract PyTorch Policy Network & Optimizer
+    policy = model.policy
+    optimizer = torch.optim.Adam(policy.parameters(), lr=5e-4)
+    loss_fn = nn.MSELoss()
+
+    # Convert to PyTorch Dataloader
+    dataset = TensorDataset(torch.tensor(obs_data), torch.tensor(action_data))
+    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+    print(f"[Pretrain] Pre-training Policy Network ({epochs} Epochs)...")
+    policy.train()
+
+    # --- PROGRESS BAR: Epoch Training ---
+    for epoch in range(epochs):
+        epoch_loss = 0.0
+        
+        batch_pbar = tqdm(loader, desc=f" Epoch {epoch + 1:02d}/{epochs:02d}", leave=True, unit="batch")
+        for batch_obs, batch_actions in batch_pbar:
+            optimizer.zero_grad()
+            
+            distribution = policy.get_distribution(batch_obs)
+            predicted_actions = distribution.distribution.mean
+
+            loss = loss_fn(predicted_actions, batch_actions)
+            loss.backward()
+            optimizer.step()
+
+            current_loss = loss.item()
+            epoch_loss += current_loss * len(batch_obs)
+            
+            # Dynamic loss update in progress bar tail
+            batch_pbar.set_postfix({"loss": f"{current_loss:.6f}"})
+
+        avg_loss = epoch_loss / len(dataset)
+        tqdm.write(f"  └─ Epoch {epoch + 1:02d}/{epochs:02d} Complete | Mean MSE Loss: {avg_loss:.6f}")
+
+    # Save SB3 Model Checkpoint
+    out_file = Path(save_path)
+    out_file.parent.mkdir(parents=True, exist_ok=True)
+    model.save(out_file)
+    dummy_env.close()
+    print(f"\n[Pretrain] Successfully saved pre-trained base model to: {out_file.resolve()}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="JackBot Behavioral Cloning Pre-trainer")
+    parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering during collection")
+    parser.add_argument("--num-samples", type=int, default=100_000, help="Number of dataset samples to collect")
+    parser.add_argument("--epochs", type=int, default=15, help="Number of BC training epochs")
+    parser.add_argument("--save-path", type=str, default="ml/checkpoints/jackbot_kinematics_base.zip", help="Output path for pre-trained model .zip")
+    args = parser.parse_args()
+
+    pretrain_policy(
+        save_path=args.save_path,
+        epochs=args.epochs,
+        num_samples=args.num_samples,
+        use_gui=args.gui
+    )

ml/run_eval.py

@@ -1,41 +1,153 @@
-"""Run a trained policy in the PyBullet sim for quick inspection.
+"""
+ml/run_eval.py - Evaluate a saved JackBot policy across fixed command phases.
 
-Usage:
-    python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3 --gui
+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
+import json
+import time
 import sys
 from pathlib import Path
+import numpy as np
+from stable_baselines3 import PPO
 
-ROOT_DIR = Path(__file__).resolve().parent.parent
-if str(ROOT_DIR) not in sys.path:
-    sys.path.insert(0, str(ROOT_DIR))
-
-from ml.evaluate import evaluate
+# Ensure project root is in sys.path
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+from ml.env import JackBotEnv, CurriculumPhase
 
 
 def main():
-    parser = argparse.ArgumentParser(description="JackBot Policy Evaluator Wrapper")
-    parser.add_argument("--model", type=str, required=True, help="Path to the trained model file (.zip)")
-    parser.add_argument("--episodes", type=int, default=3, help="Number of evaluation episodes to run.")
-    parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during evaluation")
-    parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
-    parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
-    parser.add_argument("--random-command", action="store_true", help="Randomize command samples during evaluation")
-    parser.add_argument("--save-metrics", type=str, default=None, help="Optional path to save JSON metrics report")
+    parser = argparse.ArgumentParser(description="JackBot Phase-by-Phase Policy Evaluator")
+    parser.add_argument("--model", type=str, required=True, help="Path to trained model checkpoint (.zip)")
+    parser.add_argument("--episodes-per-phase", type=int, default=1, help="Number of test episodes per phase")
+    parser.add_argument("--max-steps-per-episode", type=int, default=600, help="Max simulation steps per episode")
+    parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
+    parser.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to save evaluation summary")
     args = parser.parse_args()
 
-    evaluate(
-        model_path=args.model,
-        episodes=args.episodes,
+    print(f"[Eval] Loading policy model from: {args.model}")
+    model = PPO.load(args.model, device="cpu")
+
+    # Instantiate environment with random_command disabled so our test command stays locked
+    env = JackBotEnv(
         use_gui=args.gui,
-        robot_spacing=args.robot_spacing,
-        start_pose=args.start_pose,
-        random_command=args.random_command,
-        save_json=args.save_metrics,
+        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.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)),
+        (CurriculumPhase.FULL_COMMAND,        "FULL OMNI COMBINATION",  np.array([0.3, 0.2, 0.3], dtype=np.float32)),
+    ]
+
+    phase_summary = []
+
+    try:
+        print("\n" + "=" * 75)
+        print(" STARTING MULTI-PHASE EVALUATION SUITE")
+        print("=" * 75)
+
+        for phase_enum, phase_name, test_cmd in phase_configs:
+            print(f"\n▶ Testing Phase [{phase_enum.value}]: {phase_enum.name} ({phase_name})")
+            print(f"  Target Command Vector [vx, vy, omega]: {test_cmd.tolist()}")
+
+            ep_rewards = []
+            ep_steps = []
+            ep_distances = []
+
+            for ep in range(args.episodes_per_phase):
+                obs, _ = env.reset()
+
+                # 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
+                total_reward = 0.0
+                steps = 0
+
+                while not done:
+                    # Predict deterministic action from policy
+                    action, _ = model.predict(obs, deterministic=True)
+                    
+                    # Step environment
+                    obs, reward, terminated, truncated, _ = env.step(action)
+                    done = terminated or truncated
+
+                    total_reward += float(reward)
+                    steps += 1
+
+                    if args.gui:
+                        time.sleep(1.0 / 60.0)
+
+                status_str = "FAILED (Collapsed)" if terminated else "SUCCESS (Completed)"
+                dist = float(env.max_distance_from_start)
+                print(
+                    f"   └─ Ep {ep + 1}/{args.episodes_per_phase}: {status_str:<19} | "
+                    f"Steps: {steps:<4} | Reward: {total_reward:+.2f} | Max Dist: {dist:.2f}m"
+                )
+
+                ep_rewards.append(total_reward)
+                ep_steps.append(steps)
+                ep_distances.append(dist)
+
+            phase_summary.append({
+                "phase_id": phase_enum.value,
+                "phase_name": phase_enum.name,
+                "label": phase_name,
+                "command": test_cmd.tolist(),
+                "mean_reward": float(np.mean(ep_rewards)),
+                "mean_steps": float(np.mean(ep_steps)),
+                "mean_distance": float(np.mean(ep_distances)),
+            })
+
+    finally:
+        env.close()
+
+    # Print Summary Table
+    print("\n" + "=" * 80)
+    print(" EVALUATION SUMMARY REPORT")
+    print("=" * 80)
+    print(f"{'Phase ID & Name':<25} | {'Label':<22} | {'Reward':<8} | {'Steps':<6} | {'Max Dist':<8}")
+    print("-" * 80)
+    for res in phase_summary:
+        phase_str = f"[{res['phase_id']}] {res['phase_name']}"
+        print(
+            f"{phase_str:<25} | "
+            f"{res['label']:<22} | "
+            f"{res['mean_reward']:<+8.2f} | "
+            f"{res['mean_steps']:<6.0f} | "
+            f"{res['mean_distance']:<8.2f}m"
+        )
+    print("=" * 80)
+
+    # Save JSON Report if requested
+    if args.save_metrics:
+        out_path = Path(args.save_metrics)
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(out_path, "w") as f:
+            json.dump({"model_path": str(args.model), "summary": phase_summary}, f, indent=4)
+        print(f"\n[Eval] Saved report to: {out_path.resolve()}")
+
 
 if __name__ == "__main__":
     main()

ml/run_eval_training.py

@@ -0,0 +1,78 @@
+"""
+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
+import sys
+from pathlib import Path
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+from ml.env import JackBotEnv, CurriculumPhase
+
+def evaluate_kinematics(episode_length: int = 1000):
+    env = JackBotEnv(
+        use_gui=True, 
+        random_command=False, 
+        max_episode_steps=episode_length,
+        robot_mode="kinematics"
+    )
+
+    print("\n" + "=" * 70)
+    print(" RUNNING MULTI-PHASE REWARD BENCHMARK (KINEMATICS MODE)")
+    print("=" * 70 + "\n")
+
+    phase_configs = [
+        (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)),
+        (CurriculumPhase.FULL_COMMAND,        "FULL OMNI COMBINATION",np.array([0.3, 0.2, 0.3], dtype=np.float32)),
+    ]
+
+    for phase_enum, label, cmd in phase_configs:
+        obs, _ = env.reset()
+        
+        env.curriculum_phase = phase_enum
+        env.command = cmd.copy()
+        
+        done = False
+        total_reward = 0.0
+        step_count = 0
+
+        while not done:
+            dummy_action = np.zeros(18, dtype=np.float32)
+            obs, reward, terminated, truncated, _ = env.step(dummy_action)
+            total_reward += reward
+            step_count += 1
+            done = terminated or truncated
+
+            time.sleep(1.0 / 60.0)
+
+        comp_averages = env.get_reward_component_averages()
+
+        print(f"\n--- Episode Stage: [{phase_enum.name}] ({label}) ---")
+        print(f"Command Applied: vx={cmd[0]:.2f}, vy={cmd[1]:.2f}, yaw={cmd[2]:.2f}")
+        print(f"Total Episode Reward: {total_reward:.4f}")
+        print("Component Step Averages:")
+        for name, value in comp_averages.items():
+            print(f"  • {name:<20}: {value:+.5f}")
+
+        metrics = env.get_current_robot_metrics()
+        dist = metrics[0]["distance_from_start"] if metrics else 0.0
+        speed = metrics[0]["speed"] if metrics else 0.0
+        avg_reward = total_reward / max(1, step_count)
+
+        print(f"  ├─ Average Reward / Step: {avg_reward:.4f}")
+        print(f"  ├─ Distance Travelled: {dist:.2f} m")
+        print(f"  ├─ Actual Avg Speed: {speed:.2f} m/s")
+        print(f"  └─ Steps Survived: {step_count} / {episode_length}")
+        print("-" * 70)
+
+    env.close()
+
+if __name__ == "__main__":
+    evaluate_kinematics()

ml/run_train.py

@@ -1,47 +1,161 @@
-"""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 --timesteps 500000 --model ml/checkpoints/ppo_joint_command
-
-This is a convenience wrapper around `ml.train.train` with friendly defaults
-for interactive experimentation.
+    python ml/run_train.py --total-timesteps 1500000 --gui
 """
 
 import argparse
+import os
+import re
 import sys
+import warnings
 from pathlib import Path
+from stable_baselines3 import PPO
+from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv
+from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
 
-ROOT_DIR = Path(__file__).resolve().parent.parent
-if str(ROOT_DIR) not in sys.path:
-    sys.path.insert(0, str(ROOT_DIR))
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+from ml.env import JackBotEnv
+from ml.callbacks import CurriculumCallback, RewardLoggerCallback
 
-from ml.train import train
+# Silence SB3's UserWarning about SubprocVecEnv vs DummyVecEnv
+warnings.filterwarnings("ignore", category=UserWarning, module="stable_baselines3")
+
+
+def get_next_run_number(save_dir: str) -> int:
+    """Scans the save directory for existing ppo<number> patterns and returns the next integer."""
+    if not os.path.exists(save_dir):
+        return 1
+    
+    existing_numbers = []
+    for item in os.listdir(save_dir):
+        # Match pattern ppo followed by numbers (e.g., ppo1, ppo_1, jackbot_ppo12)
+        matches = re.findall(r"ppo_?(\d+)", item, re.IGNORECASE)
+        for m in matches:
+            existing_numbers.append(int(m))
+            
+    return max(existing_numbers, default=0) + 1
+
+
+def make_env(rank: int, use_gui: bool = False, seed: int = 0):
+    """Utility helper to instantiate parallel JackBot environments."""
+    def _init():
+        env = JackBotEnv(
+            use_gui=use_gui,
+            random_command=True,
+        )
+        env.reset(seed=seed + rank)
+        return env
+    return _init
 
 
 def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--timesteps", type=int, default=500000, help="Total number of 'practice steps'.")
-    parser.add_argument("--model", type=str, default="ml/checkpoints/ppo_joint_command", help="Path to save the trained model")
-    parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducibility")
-    parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto'")
-    parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during training")
-    parser.add_argument("--num-workers", type=int, default=1, help="Number of training environment workers")
-    parser.add_argument("--robot-spacing", type=float, default=1.5, help="Spacing between robots in meters")
-    parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
+    parser = argparse.ArgumentParser(description="JackBot PPO Curriculum Trainer")
+    parser.add_argument("--num-workers", type=int, default=16, help="Number of parallel sub-process environments")
+    parser.add_argument("--total-timesteps", type=int, default=1_500_000, help="Total training timesteps")
+    parser.add_argument("--log-dir", type=str, default="ml/logs", help="Directory for TensorBoard logs")
+    parser.add_argument("--save-dir", type=str, default="ml/checkpoints", help="Directory for model checkpoints")
+    parser.add_argument("--save-freq", type=int, default=50_000, help="Checkpoint save frequency (steps)")
+    parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
+    parser.add_argument(
+        "--pretrained-model",
+        type=str,
+        default=None,
+        help="Path to pre-trained base model checkpoint (.zip) to start PPO training from"
+    )
     args = parser.parse_args()
 
-    Path(args.model).parent.mkdir(parents=True, exist_ok=True)
-    train(
-        total_timesteps=args.timesteps,
-        model_path=args.model,
-        seed=args.seed,
-        device=args.device,
-        use_gui=args.gui,
-        num_workers=args.num_workers,
-        robot_spacing=args.robot_spacing,
-        start_pose=args.start_pose,
+    os.makedirs(args.log_dir, exist_ok=True)
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    # Automatically determine next run number (e.g. ppo1, ppo2, ppo3...)
+    run_num = get_next_run_number(args.save_dir)
+    ppo_name = f"ppo{run_num}"
+
+    print(f"[Train] Initializing Run #{run_num} ('{ppo_name}') with {args.num_workers} parallel workers...")
+    
+    env_fns = [
+        make_env(rank=i, use_gui=(args.gui if i == 0 else False)) 
+        for i in range(args.num_workers)
+    ]
+    vec_env = SubprocVecEnv(env_fns)
+
+    # Initialize or Load PPO Policy Model
+    if args.pretrained_model and os.path.exists(args.pretrained_model):
+        print(f"[Train] Loading pre-trained base knowledge from: {args.pretrained_model}")
+        model = PPO.load(
+            args.pretrained_model,
+            env=vec_env,
+            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",
+        )
+    else:
+        print("[Train] No base model provided. Starting training from scratch...")
+        model = PPO(
+            policy="MlpPolicy",
+            env=vec_env,
+            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.0,
+            target_kl=0.05,
+            vf_coef=0.5,
+            max_grad_norm=0.5,
+            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),
+        save_path=args.save_dir,
+        name_prefix=f"jackbot_{ppo_name}",
     )
+    reward_logger_callback = RewardLoggerCallback(verbose=1)
+    curriculum_callback = CurriculumCallback()
+
+    eval_env = DummyVecEnv([lambda: JackBotEnv(use_gui=False, random_command=True)])
+    best_model_path = os.path.join(args.save_dir, f"best_model_{ppo_name}")
+
+    eval_callback = EvalCallback(
+        eval_env,
+        best_model_save_path=best_model_path,
+        log_path="ml/logs/results",
+        eval_freq=max(1, 50_000 // args.num_workers),
+        deterministic=True,
+        render=False,
+    )
+
+    print(f"[Train] Starting training for {args.total_timesteps} timesteps...")
+    try:
+        model.learn(
+            total_timesteps=args.total_timesteps,
+            callback=[checkpoint_callback, reward_logger_callback, curriculum_callback, eval_callback],
+            progress_bar=True,
+        )
+        final_model_path = os.path.join(args.save_dir, f"jackbot_{ppo_name}_final.zip")
+        model.save(final_model_path)
+        print(f"[Train] Training complete! Saved final model to {final_model_path}")
+    finally:
+        vec_env.close()
+        eval_env.close()
 
 
 if __name__ == "__main__":
-    main()
+    main()

ml/train.py

@@ -1,230 +0,0 @@
-import os
-import argparse
-from pathlib import Path
-
-import numpy as np
-from stable_baselines3.common.callbacks import BaseCallback
-from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv
-from .env import JackBotEnv
-
-
-class MilestoneCheckpointCallback(BaseCallback):
-    """
-    Saves a model checkpoint the FIRST time total_timesteps 
-    crosses every multiple of step_interval (e.g., 100,000).
-    """
-    def __init__(self, save_path: str, name_prefix: str = "ppo_jackbot", step_interval: int = 100_000, verbose: int = 1):
-        super().__init__(verbose)
-        self.save_path = save_path
-        self.name_prefix = name_prefix
-        self.step_interval = step_interval
-        self.last_milestone = 0
-        os.makedirs(self.save_path, exist_ok=True)
-
-    def _on_step(self) -> bool:
-        current_milestone = self.num_timesteps // self.step_interval
-        
-        if current_milestone > self.last_milestone:
-            self.last_milestone = current_milestone
-            milestone_step = current_milestone * self.step_interval
-            
-            save_file = os.path.join(
-                self.save_path, 
-                f"{self.name_prefix}_{milestone_step}_steps.zip"
-            )
-            self.model.save(save_file)
-            
-            if self.verbose > 0:
-                print(f"\n[Checkpoint] Saved milestone model at {self.num_timesteps} steps -> {save_file}\n")
-                
-        return True
-
-class JackBotMetricsCallback(BaseCallback):
-    """
-    Tracks the best current alive-robot performance for the most recent rollout,
-    instead of logging lifetime maxima from the entire training run.
-    """
-    def __init__(self, verbose=0):
-        super().__init__(verbose)
-        self.best_alive_speed = 0.0
-        self.best_alive_yaw_rate = 0.0
-        self.best_alive_distance = 0.0
-        self.best_alive_survival_steps = 0.0
-        self.best_alive_reward = -float('inf')
-
-    def _on_step(self) -> bool:
-        """Required by SB3 BaseCallback; no-op here because the rollout summary is emitted at rollout end."""
-        return True
-
-    def _on_rollout_end(self) -> bool:
-        """Executed right before PPO outputs the log table to console."""
-        try:
-            vec_env = self.training_env
-            alive_metrics = vec_env.env_method("get_current_robot_metrics")
-
-            self.best_alive_speed = 0.0
-            self.best_alive_yaw_rate = 0.0
-            self.best_alive_distance = 0.0
-            self.best_alive_survival_steps = 0.0
-            self.best_alive_reward = -float('inf')
-
-            for worker_res in alive_metrics:
-                for metrics in worker_res:
-                    if not metrics.get("alive", False):
-                        continue
-
-                    if metrics["reward"] > self.best_alive_reward:
-                        self.best_alive_reward = float(metrics["reward"])
-                    if metrics["speed"] > self.best_alive_speed:
-                        self.best_alive_speed = float(metrics["speed"])
-                    if metrics["yaw_rate"] > self.best_alive_yaw_rate:
-                        self.best_alive_yaw_rate = float(metrics["yaw_rate"])
-                    if metrics["distance_from_start"] > self.best_alive_distance:
-                        self.best_alive_distance = float(metrics["distance_from_start"])
-                    if metrics["survival_steps"] > self.best_alive_survival_steps:
-                        self.best_alive_survival_steps = float(metrics["survival_steps"])
-
-            self.logger.record("custom/best_alive_speed_mps", float(self.best_alive_speed))
-            self.logger.record("custom/best_alive_yaw_rate_rads", float(self.best_alive_yaw_rate))
-            self.logger.record("custom/best_alive_distance_from_start_m", float(self.best_alive_distance))
-            self.logger.record("custom/best_alive_survival_steps", float(self.best_alive_survival_steps))
-            self.logger.record("custom/best_alive_reward", float(self.best_alive_reward) if np.isfinite(self.best_alive_reward) else 0.0)
-
-            # Backward-compatible aliases so old dashboards keep a stable field name.
-            self.logger.record("custom/max_speed_mps", float(self.best_alive_speed))
-            self.logger.record("custom/max_yaw_rate_rads", float(self.best_alive_yaw_rate))
-            self.logger.record("custom/max_distance_from_start_m", float(self.best_alive_distance))
-            self.logger.record("custom/max_survival_steps", float(self.best_alive_survival_steps))
-            self.logger.record("custom/best_episode_reward", float(self.best_alive_reward) if np.isfinite(self.best_alive_reward) else 0.0)
-        except Exception:
-            pass
-
-        return True
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Train a joint-command policy for JackBot.")
-    parser.add_argument("--timesteps", type=int, default=500_000, help="Total training timesteps")
-    parser.add_argument("--model-path", type=str, default="ml/checkpoints/ppo_joint_command", help="Where to save the trained model")
-    parser.add_argument("--seed", type=int, default=0, help="Random seed")
-    parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto' to autodetect")
-    parser.add_argument("--gui", "--use-gui", dest="use_gui", action="store_true", help="Enable PyBullet GUI during training")
-    parser.add_argument("--num-workers", type=int, default=8, help="Number of parallel CPU worker processes")
-    parser.add_argument("--robot-spacing", type=float, default=3.0, help="Spacing between robots in meters")
-    parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
-    return parser.parse_args()
-
-
-def make_env(robot_spacing, start_pose, use_gui, rank, seed=0):
-    def _init():
-        env = JackBotEnv(
-            use_gui=use_gui if rank == 0 else False,  # Only rank 0 gets GUI if requested
-            random_command=True,
-            robot_spacing=robot_spacing,
-            start_pose=start_pose,
-        )
-        env.reset(seed=seed + rank)
-        return env
-    return _init
-
-
-def train(
-    total_timesteps: int,
-    model_path: str,
-    seed: int = 0,
-    device: str = "auto",
-    use_gui: bool = False,
-    num_workers: int = 8,
-    robot_spacing: float = 0.5,
-    start_pose: str = "init_deg",
-):
-    try:
-        from stable_baselines3 import PPO
-    except ImportError as exc:
-        raise ImportError("stable-baselines3 is required. Install with: pip install stable-baselines3") from exc
-
-    # Create multi-process vector environment
-    if num_workers > 1:
-        env_fns = [
-            make_env(robot_spacing, start_pose, use_gui, rank=i, seed=seed)
-            for i in range(num_workers)
-        ]
-        env = SubprocVecEnv(env_fns)
-    else:
-        env = DummyVecEnv([
-            make_env(robot_spacing, start_pose, use_gui, rank=0, seed=seed)
-        ])
-
-    def resolve_device(requested_device: str) -> str:
-        try:
-            import torch
-        except ImportError:
-            if requested_device != "cpu":
-                raise RuntimeError("PyTorch is not installed.")
-            return "cpu"
-
-        hip_supported = getattr(torch.version, "hip", None) is not None
-        cuda_available = torch.cuda.is_available()
-        hip_available = hip_supported and getattr(torch.backends, "hip", None) is not None and torch.backends.hip.is_available()
-
-        if requested_device == "auto":
-            return "cuda" if (hip_available or cuda_available) else "cpu"
-
-        if requested_device in {"cuda", "gpu", "hip"}:
-            if hip_available or cuda_available:
-                return "cuda"
-            raise RuntimeError(f"GPU requested ({requested_device}) but not available.")
-
-        return "cpu"
-
-    device = resolve_device(device)
-
-    model = PPO(
-        "MlpPolicy",
-        env,
-        verbose=1,
-        seed=seed,
-        learning_rate=3.5e-4,      # Cut LR in half (from 3e-4) to smooth out updates
-        n_steps=2048,              # Larger rollout buffer per env for stable gradients
-        batch_size=128,            # Larger minibatches reduce noise
-        n_epochs=10,               # Number of epoch updates per rollout
-        gamma=0.99,                # Discount factor
-        gae_lambda=0.95,           # GAE smoothing
-        clip_range=0.2,            # Standard PPO clipping
-        target_kl=0.03,           # EARLY STOPPING: Halts policy update if KL > 0.015!
-        ent_coef=0.03,            # Entropy coefficient to encourage exploration
-        vf_coef=0.5,
-        max_grad_norm=0.5,
-        device=device,
-        tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"),
-    )
-
-    save_dir = str(Path(model_path).parent)
-    model_prefix = Path(model_path).stem
-
-    milestone_cb = MilestoneCheckpointCallback(
-        save_path=save_dir,
-        name_prefix=model_prefix,
-        step_interval=100_000
-    )
-
-    metrics_callback = JackBotMetricsCallback()
-
-    model.learn(total_timesteps=total_timesteps, callback=[milestone_cb, metrics_callback])
-
-    Path(model_path).parent.mkdir(parents=True, exist_ok=True)
-    model.save(model_path)
-    env.close()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    train(
-        args.timesteps,
-        args.model_path,
-        seed=args.seed,
-        device=args.device,
-        use_gui=args.use_gui,
-        num_workers=args.num_workers,
-        robot_spacing=args.robot_spacing,
-        start_pose=args.start_pose,
-    )

requirements.txt

@@ -4,10 +4,12 @@ ikpy
 pybullet
 pyserial
 matplotlib
+dearpygui
 
 # Deep learning / RL
 torch
 stable-baselines3
-gymnasium[box2d]
+stable-baselines3[extra]
+gymnasium
 shimmy
 tensorboard

simulation.py

@@ -1,89 +1,216 @@
 """
-simulation.py - PyBullet Simulation Interface & Standalone Runner
+simulation.py - PyBullet Simulation Interface & Physics Engine
+Consolidates scene management, physics queries, motor control, and rendering.
 """
 import time
 import math
+from typing import List, Tuple, Optional, Union
 import numpy as np
 import pybullet as p
-import pybullet_data  # Added missing import
+import pybullet_data
 
 from config import cfg
 import DataTypes as dt
 
 
 class Simulation:
-    def __init__(self, urdf_path: str = cfg.urdf_path):
+    """Single PyBullet simulation manager providing getters/setters for JackBot."""
+
+    def __init__(self, urdf_path: str = cfg.urdf_path, use_gui: bool = True):
         self.urdf_path = urdf_path
-        
-        # Connect to PyBullet GUI
-        self.physicsClient = p.connect(p.GUI)
+        self.use_gui = use_gui
+        self.physics_client: Optional[int] = None
+        self.plane_id: Optional[int] = None
+        self.robot_id: Optional[int] = None
+        self.revolute_joints: List[int] = []
+
+        self.connect()
+
+    def connect(self) -> None:
+        """Establishes connection to PyBullet GUI or DIRECT mode."""
+        if self.physics_client is not None and p.isConnected(self.physics_client):
+            return
+
+        flags = p.GUI if self.use_gui else p.DIRECT
+        self.physics_client = p.connect(flags)
+
         p.setAdditionalSearchPath(pybullet_data.getDataPath())
-        p.setGravity(0, 0, -9.81)
-        
-        # Load plane and robot URDF
-        self.planeId = p.loadURDF("plane.urdf")
-        self.robot = p.loadURDF(self.urdf_path, [0, 0, 0.2])
+        p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client)
 
-        # Discover revolute joint indices dynamically
-        self.revolute_joints = []
-        for j in range(p.getNumJoints(self.robot)):
-            joint_info = p.getJointInfo(self.robot, j)
-            if joint_info[2] == p.JOINT_REVOLUTE:
-                self.revolute_joints.append(j)
-
-        self.set_all_joints_to_90()
-        p.resetDebugVisualizerCamera(
-            cameraDistance=1.0,
-            cameraYaw=50,
-            cameraPitch=-35,
-            cameraTargetPosition=[0, 0, 0],
-        )
-
-    def set_all_joints_to_90(self):
-        for joint_index in self.revolute_joints:
-            p.resetJointState(self.robot, joint_index, math.radians(90))
-
-    def updatePos(self, current_rad: dt.RadArray):
-        radflat = current_rad.data.flatten()
-        for joint_index, target_angle in zip(self.revolute_joints, radflat):
-            p.setJointMotorControl2(
-                bodyIndex=self.robot,
-                jointIndex=joint_index,
-                controlMode=p.POSITION_CONTROL,
-                targetPosition=target_angle,
-                force=500,
+        if self.use_gui:
+            p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.physics_client)
+            p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self.physics_client)
+            p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
+            p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
+            p.resetDebugVisualizerCamera(
+                cameraDistance=1.0, cameraYaw=50, cameraPitch=-35, cameraTargetPosition=[0, 0, 0],
+                physicsClientId=self.physics_client
             )
 
-    def step(self):
-        p.stepSimulation()
+    def load_scene(self, spawn_pos: Optional[List[float]] = None) -> Tuple[int, int, List[int]]:
+        """Loads plane and robot URDF, discovering revolute joint indices dynamically."""
+        if spawn_pos is None:
+            spawn_pos = [0.0, 0.0, 0.20]
 
-    def disconnect(self):
-        if p.isConnected(self.physicsClient):
-            p.disconnect(self.physicsClient)
+        self.plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
+        self.robot_id = p.loadURDF(self.urdf_path, spawn_pos, physicsClientId=self.physics_client)
 
-    def close(self):
-        """Cleanup wrapper for Robot backend compatibility."""
+        self.revolute_joints = []
+        for j in range(p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)):
+            info = p.getJointInfo(self.robot_id, j, physicsClientId=self.physics_client)
+            if info[2] == p.JOINT_REVOLUTE:
+                self.revolute_joints.append(j)
+
+        return self.plane_id, self.robot_id, self.revolute_joints
+
+    # --- ACTUATION (SETTERS) ---
+
+    def set_robot_joint_angles(
+        self, target_angles: Union[np.ndarray, List[float], dt.RadArray]
+    ) -> None:
+        """Applies motor torque to pull joints toward target position angles."""
+        if isinstance(target_angles, dt.RadArray):
+            radflat = target_angles.data.flatten()
+        elif isinstance(target_angles, np.ndarray):
+            radflat = target_angles.flatten()
+        else:
+            radflat = target_angles
+
+        for joint_index, target_angle in zip(self.revolute_joints, radflat):
+            p.setJointMotorControl2(
+                bodyIndex=self.robot_id,
+                jointIndex=joint_index,
+                controlMode=p.POSITION_CONTROL,
+                targetPosition=float(target_angle),
+                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()
+        for joint_index, target_angle in zip(self.revolute_joints, radflat):
+            p.resetJointState(
+                bodyUniqueId=self.robot_id,
+                jointIndex=joint_index,
+                targetValue=float(target_angle),
+                targetVelocity=0.0,
+                physicsClientId=self.physics_client
+            )
+
+    def reset_robot_base(
+        self,
+        pos: Optional[List[float]] = None,
+        orn: Optional[List[float]] = None,
+        linear_velocity: Optional[List[float]] = None,
+        angular_velocity: Optional[List[float]] = None
+    ) -> None:
+        """Resets root torso position, orientation quaternion, and clears base velocities."""
+        if pos is None:
+            pos = [0.0, 0.0, 0.20]
+        if orn is None:
+            orn = [0.0, 0.0, 0.0, 1.0]
+
+        lin_v = linear_velocity if linear_velocity is not None else [0.0, 0.0, 0.0]
+        ang_v = angular_velocity if angular_velocity is not None else [0.0, 0.0, 0.0]
+
+        p.resetBasePositionAndOrientation(self.robot_id, pos, orn, physicsClientId=self.physics_client)
+        p.resetBaseVelocity(self.robot_id, linearVelocity=lin_v, angularVelocity=ang_v, physicsClientId=self.physics_client)
+    
+    # --- TELEMETRY (GETTERS) ---
+
+    def get_robot_pose(self) -> Tuple[List[float], List[float]]:
+        pos, orn = p.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.physics_client)
+        return list(pos), list(orn)
+
+    def get_robot_rpy(self) -> Tuple[float, float, float]:
+        _, orn = p.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.physics_client)
+        roll, pitch, yaw = p.getEulerFromQuaternion(orn)
+        return float(roll), float(pitch), float(yaw)
+
+    def get_robot_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]:
+        pos, orn = p.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.physics_client)
+        roll, pitch, yaw = p.getEulerFromQuaternion(orn)
+        return list(pos), (float(roll), float(pitch), float(yaw))
+
+    def get_robot_velocity(self) -> Tuple[List[float], List[float]]:
+        lin_v, ang_v = p.getBaseVelocity(self.robot_id, physicsClientId=self.physics_client)
+        return list(lin_v), list(ang_v)
+
+    def get_robot_joint_angles(self) -> np.ndarray:
+        joint_states = p.getJointStates(self.robot_id, self.revolute_joints, physicsClientId=self.physics_client)
+        return np.array([state[0] for state in joint_states], dtype=np.float32)
+
+    def _get_urdf_joint_limits(self) -> Tuple[np.ndarray, np.ndarray]:
+        """Dynamically reads lower and upper limits for all revolute joints from PyBullet."""
+        lower_limits = []
+        upper_limits = []
+        
+        # Iterate through joints in PyBullet
+        for j_idx in range(p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)):
+            info = p.getJointInfo(self.robot_id, j_idx, physicsClientId=self.physics_client)
+            joint_type = info[2]
+            
+            # Only collect limits for revolute joints
+            if joint_type == p.JOINT_REVOLUTE:
+                lower_limits.append(info[8])  # Index 8 = jointLowerLimit
+                upper_limits.append(info[9])  # Index 9 = jointUpperLimit
+                
+        return np.array(lower_limits, dtype=np.float32), np.array(upper_limits, dtype=np.float32)
+
+    # --- SIMULATION LIFECYCLE CONTROLS ---
+
+    def step(self) -> None:
+        """Advances physics simulation by 1 time step."""
+        p.stepSimulation(physicsClientId=self.physics_client)
+
+    def settle_and_measure_height(
+        self, target_angles: Optional[np.ndarray] = None, steps: int = 100, fallback_height: float = 0.122
+    ) -> float:
+        """Settles the robot into the ground while actively holding target joint angles."""
+        for _ in range(steps):
+            if target_angles is not None:
+                self.set_robot_joint_angles(target_angles)
+            self.step()
+        pos, _ = self.get_robot_pose()
+        return pos[2] if pos[2] > 0.0 else fallback_height
+
+    def apply_external_force(
+        self, force: Union[List[float], np.ndarray], link_index: int = -1, position: Tuple[float, float, float] = (0.0, 0.0, 0.0)
+    ) -> None:
+        p.applyExternalForce(
+            objectUniqueId=self.robot_id,
+            linkIndex=link_index,
+            forceObj=list(force),
+            posObj=list(position),
+            flags=p.WORLD_FRAME,
+            physicsClientId=self.physics_client,
+        )
+
+    def disconnect(self) -> None:
+        if self.physics_client is not None and p.isConnected(self.physics_client):
+            p.disconnect(self.physics_client)
+            self.physics_client = None
+
+    def close(self) -> None:
         self.disconnect()
 
 
 if __name__ == "__main__":
     from Robot import Robot, PyBulletBackend
 
-    # 1. Initialize PyBullet simulation environment
-    sim_instance = Simulation()
-    backend = PyBulletBackend(sim_instance)
+    sim_instance = Simulation(use_gui=True)
+    sim_instance.load_scene()
 
-    # 2. Instantiate Robot with simulation backend
-    robot = Robot(backend_type=backend)
+    backend = PyBulletBackend(sim_instance)
+    robot = Robot(backend_type=backend, mode="kinematics")
     robot.reset_to_init()
 
-    # 3. Command forward movement [vx, vy, omega]
-    robot.vector_dirmov = [1.0, 0.0, 0.0]
+    # Forward velocity command
+    robot.vector_dirmov = [0.3, 0.0, 0.0]
 
-    # 4. Main test execution loop
     try:
         while True:
-            # Executes state machine logic
             robot.tick()
             time.sleep(1.0 / 60.0)
     except KeyboardInterrupt: