RB1001-08 : סימולציה פשוטה מאוד של נהיגה אוטונומית דו־ממדית RL PPO

מטרת התוכנית

זוהי סימולציה בסיסית של נהיגה אוטונומית דו־ממדית באמצעות למידת חיזוק מסוג PPO.

המודל לומד לשמור רכב וירטואלי קרוב למרכז כביש מתעקל. הוא אינו מקבל הוראות מתי לפנות, אלא לומד מתוך ניסוי, תגמול ועונש.

הקלט למודל

בכל צעד המודל מקבל מטריצה בגודל 50×39, הכוללת 1,950 ערכים:

0.0 – שטח ריק

0.5 – גבול הכביש

1.0 – הרכב

המטריצה המוצגת בצד שמאל היא אותו מידע שמקבל המודל, אך היא מוגדלת לצורך תצוגה.

המודל אינו מקבל ישירות את מרכז הכביש או את כיוון הפנייה. עליו להסיק אותם מתוך מיקום הרכב וצורת גבולות הכביש.

הפעולות האפשריות

למודל שלוש פעולות:

STRAIGHT – הפחתת המהירות הצידית לכיוון אפס.

LEFT – היגוי שמאלה.

RIGHT – היגוי ימינה.

המהירות קדימה קבועה. המודל שולט רק בהיגוי הצידי.

מערכת התגמולים

מטרת התגמול היא לעודד נסיעה במרכז הכביש ולהעניש התקרבות לקצה או יציאה מהכביש.

המשתנה d הוא המרחק המנורמל ממרכז הכביש:

d = 0 – הרכב נמצא במרכז הכביש.

d = 1 – קצה הרכב מגיע לקצה הכביש.

התגמול בכל צעד מחושב כך:

Reward = 0.02 + 0.08 × (1 − d²)

לכן:

במרכז הכביש מתקבל תגמול של 0.10+.

באמצע הדרך לקצה מתקבל תגמול של כ־0.08+.

קרוב לקצה מתקבל תגמול של כ־0.02+.

במרחק של עד 10 פיקסלים מקצה הכביש מתווסף עונש של 0.15−.

ביציאה מהכביש מתקבל עונש של 1.0− והפרק מסתיים.

לאחר השלמת מרחק של 2,800 יחידות מתקבל תגמול נוסף של 3.0+ והפרק מסתיים בהצלחה.

פרק אימון

Episode הוא ניסיון נהיגה אחד.

כל פרק מסתיים באחד המצבים הבאים:

יציאה מהכביש.

השלמת המסלול.

הגעה למגבלת הצעדים.

לאחר מכן מתחיל אוטומטית פרק חדש. נקודת ההתחלה לאורך הכביש משתנה כדי למנוע מהמודל לשנן מסלול התחלתי קבוע.

מה ה־Actor לומד?

ה־Actor לומד מדיניות בחירת פעולות.

עבור כל מטריצה הוא מחשב הסתברות לשלוש הפעולות: ישר, שמאלה וימינה.

לדוגמה:

ישר – 10%

שמאלה – 82%

ימינה – 8%

במהלך האימון הפעולה נדגמת מתוך התפלגות ההסתברויות. כך המודל משלב בין חקירת פעולות חדשות לבין שימוש בפעולות שכבר הוכיחו את עצמן.

במצב Play משתמשים ב־deterministic=True. במצב זה אין אימון או חקירה, והמודל בוחר בפעולה בעלת ההסתברות הגבוהה ביותר.

מה ה־Critic לומד?

ה־Critic מעריך את התגמול העתידי הצפוי מהמצב הנוכחי:

V(s) = Expected future reward

הוא אינו שולט ברכב. תפקידו להעריך אם המצב הנוכחי צפוי להוביל לנסיעה מוצלחת או להתרסקות.

ההפרש בין התחזית שלו לבין התוצאה בפועל משמש לשיפור ה־Actor וה־Critic.

תהליך הלמידה

המודל מקבל מטריצה.
ה־Actor בוחר פעולה.
הרכב משנה את מיקומו.
הסביבה מחשבת תגמול.
המצב, הפעולה והתגמול נשמרים.
לאחר 512 צעדים מתקבל Rollout.
PPO מחשב את יתרון הפעולות ביחס לתחזיות ה־Critic.
משקלי ה־Actor וה־Critic מתעדכנים במשך 10 Epochs.
מתחיל Rollout חדש.

מדוע משתמשים ב־PPO?

PPO הוא אלגוריתם Actor-Critic המגביל את גודל השינוי במדיניות בכל עדכון.

מנגנון זה מאפשר למודל להשתפר בהדרגה ומקטין את הסיכון שעדכון חד יפגע במדיניות שכבר נלמדה.

גרף התגמול

הגרף מציג את התגמולים של 240 הצעדים האחרונים:

קו ירוק – תגמול רגיל.

שיא ירוק סביב 3+ – השלמת המסלול.

שיא אדום בגובה 1− – יציאה מהכביש.

שיפור בלמידה אמור להתבטא בירידה במספר השיאים האדומים, בעלייה במספר השלמות המסלול ובתגמול רגיל הקרוב ל־0.10.

מדדי האימון

Actor Loss – מתאר את עדכון מדיניות הפעולות. אין לו ערך אידיאלי קבוע.

Value Loss – שגיאת תחזית ה־Critic ביחס לתגמול העתידי שהתקבל בפועל.

Explained Variance – מודד את איכות תחזיות ה־Critic. ערך הקרוב ל־1 טוב, וערך הקרוב ל־0 מצביע על תחזיות חלשות.

Entropy – מודדת את מידת האקראיות והחקירה במדיניות.

Approx KL – מודד את השינוי בין המדיניות הישנה לחדשה.

Clip Fraction – שיעור העדכונים שהוגבלו על ידי מנגנון ה־Clipping של PPO.

מבנה הרשת

הקוד משתמש ב־MlpPolicy.

המטריצה הדו־ממדית משוטחת לווקטור של 1,950 ערכים ומועברת לרשת עצבית מסוג MLP. לכן המודל אינו משתמש ב־CNN ואינו מקבל מראש הנחה מרחבית לגבי הקשר בין תאים סמוכים.

CNN עשויה לזהות עיקולים ומבנים מרחביים בצורה יעילה יותר.

מגבלת הקלט

המהירות הצידית של הרכב אינה כלולה בקלט.

לכן שתי מטריצות דומות יכולות לייצג רכב שנע שמאלה או ימינה. הפעולה המתאימה עשויה להיות שונה, אף שהקלט נראה זהה.

אפשר לשפר את המערכת באמצעות:

הוספת המהירות הצידית לקלט.

הוספת הפעולה הקודמת.

שימוש במספר מטריצות עוקבות.

שימוש ב־Recurrent PPO.

שימוש ב־CNN לניתוח המטריצה.

מגבלות הסימולציה

התוכנית אינה מערכת מלאה של רכב אוטונומי.

אין בה מצלמה אמיתית, מכשולים, תמרורים, תכנון מסלול, שליטה במהירות או פיזיקה מלאה של רכב.

ההגדרה המדויקת היא:

סימולציה בסיסית של שמירת רכב וירטואלי במרכז כביש דו־ממדי באמצעות PPO.

עיקרון מרכזי

המתכנת אינו מגדיר למודל מתי לפנות שמאלה או ימינה. הוא מגדיר את הקלט, הפעולות ומערכת התגמולים.

המודל לומד מדיניות נהיגה שמטרתה למקסם את סכום התגמולים העתידי.

R=0.02+0.08(1-d^2)
]

מצב	תגמול
מרכז הכביש	`+0.10`
באמצע הדרך לקצה	בערך `+0.08`
קרוב לקצה	בערך `+0.02`
עד 10 פיקסלים מהקצה	תוספת `−0.15`
יציאה מהכביש	`−1.0` וסיום הפרק
השלמת 2,800 יחידות	תוספת `+3.0` וסיום הפרק

התגמול בהשלמה הוא למעשה בין +3.02 ל־+3.10, לפי מיקום הרכב.

מה ה־Actor לומד?

ה־Actor לומד מדיניות:

[
\pi(a|s)
]

כלומר, מה ההסתברות לבחור בכל פעולה כאשר מתקבלת מטריצה מסוימת.

לדוגמה:

פעולה	הסתברות
ישר	10%
שמאלה	82%
ימינה	8%

במהלך האימון הפעולה נדגמת מההתפלגות, ולכן קיימת חקירה. במצב play משתמשים ב־deterministic=True, ולכן נבחרת הפעולה בעלת ההסתברות הגבוהה ביותר.

מה ה־Critic לומד?

ה־Critic מעריך:

[
V(s)=\text{התגמול העתידי הצפוי ממצב זה}
]

הוא מנסה לענות:

האם המצב הנוכחי צפוי להוביל לנסיעה מוצלחת או להתרסקות?

ה־Critic אינו נוהג ברכב. הוא מספק ל־Actor משוב לגבי איכות המצבים והפעולות.

כיצד PPO לומד?

המודל מקבל מטריצה.
ה־Actor בוחר פעולה.
הרכב זז.
הסביבה מחשבת תגמול.
התהליך חוזר במשך 512 צעדים.
PPO משווה בין התגמול שהתקבל לבין תחזית ה־Critic.
משקלי ה־Actor וה־Critic מתעדכנים במשך 10 epochs.
מתחיל Rollout חדש.

משמעות הגרף

קו ירוק: התגמול בכל צעד.
שיא ירוק סביב +3: השלמת המסלול.
שיא אדום בגובה −1: התרסקות.
הגרף מציג תגמול לכל צעד, לא את סכום התגמול של הפרק.

נקודה טכנית חשובה מאוד

למרות שהקלט הוא מטריצה דו־ממדית, הקוד משתמש ב:

"MlpPolicy"

1 2	"MlpPolicy"

לכן Stable-Baselines3 משטח אותה לווקטור של 1,950 מספרים. זו אינה רשת CNN.

המשמעות:

המודל יכול ללמוד.
אבל הוא אינו מקבל מראש את העיקרון שפיקסלים סמוכים קשורים זה לזה.
CNN עשויה ללמוד עיקולים ומבנה מרחבי ביעילות גבוהה יותר.

מגבלה נוספת

המהירות הצידית של הרכב אינה כלולה בקלט. שני מצבים עשויים להיראות כמעט זהים, אף שבאחד הרכב נע במהירות שמאלה ובשני ימינה.

לכן הסביבה אינה Markov מלאה מבחינת המודל. לשיפור עתידי אפשר:

להוסיף מהירות צידית כנתון.
לתת למודל מספר מטריצות עוקבות.
להשתמש ב־Recurrent PPO.

המודל אינו לומד את חוקי הנהיגה במילים. הוא לומד התאמה סטטיסטית בין מבנה הכביש, מיקום הרכב, הפעולה שנבחרה והתגמול העתידי שהתקבל.

קוד התכונה

"""
PPO road-centre driving trainer with a live Pygame display.

The PPO agent receives only the 50 x 39 AI matrix shown on the left:
    0.0 = empty
    0.5 = road boundary
    1.0 = car

Install:
    pip install pygame gymnasium stable-baselines3 numpy torch

Train and watch:
    python ppo_car_trainer.py --mode train --steps 500000

Continue training a saved model:
    python ppo_car_trainer.py --mode train --resume --steps 500000

Train without graphics:
    python ppo_car_trainer.py --mode train --headless --steps 500000

Watch the saved model without learning:
    python ppo_car_trainer.py --mode play

Live controls:
    SPACE = pause / continue
    Q     = save the PPO model immediately
    +     = faster display/training
    -     = slower display/training
    ESC   = stop and save
"""

from __future__ import annotations

import argparse
import math
import os
import time
from collections import deque
from collections.abc import Callable

import gymnasium as gym
import numpy as np
import pygame
from gymnasium import spaces


# Display and road geometry
MATRIX_WIDTH, MATRIX_HEIGHT = 50, 39
SENSOR_WIDTH, SENSOR_HEIGHT = 200, 156
SENSOR_X, SENSOR_Y = 10, 65
GAME_LEFT, GAME_WIDTH = 220, 900
WIDTH, HEIGHT = GAME_LEFT + GAME_WIDTH, 700
GRAPH_WIDTH, GRAPH_HEIGHT = 370, 180
GRAPH_X, GRAPH_Y = WIDTH - GRAPH_WIDTH - 15, 15

ROAD_WIDTH = 360
CAR_WIDTH, CAR_HEIGHT = 43, 77
CAR_Y = HEIGHT - 145

# Fixed simulation physics. PPO sees one action every 0.05 simulated seconds.
DT = 0.05
FORWARD_SPEED = 230.0
MAX_LATERAL_SPEED = 620.0
STEERING_RESPONSE = 12.0
EPISODE_DISTANCE = 2800.0
MAX_EPISODE_STEPS = 650

# Reward design agreed for road-centre driving.
FORWARD_REWARD = 0.02
CENTER_REWARD_SCALE = 0.08
NEAR_EDGE_PIXELS = 10.0
NEAR_EDGE_PENALTY = -0.15
CRASH_REWARD = -1.0
COMPLETE_REWARD = 3.0

# PPO parameters based on the supplied reference program.
LEARNING_RATE = 3e-4
ROLLOUT_STEPS = 512
BATCH_SIZE = 256
N_EPOCHS = 10
GAMMA = 0.995
GAE_LAMBDA = 0.95
ENTROPY_COEFFICIENT = 0.005
CLIP_RANGE = 0.2

MODEL_PATH = "ppo_road_center_car"


def road_center(world_y: float) -> float:
    """Return the road centre for straight, curved and linear sections."""
    section = world_y % 2400.0
    base = GAME_LEFT + GAME_WIDTH / 2

    def smoothstep(value: float) -> float:
        return value * value * (3.0 - 2.0 * value)

    if section < 350:
        offset = 0.0
    elif section < 700:
        t = (section - 350) / 350
        offset = 180 * smoothstep(t)
    elif section < 950:
        offset = 180.0
    elif section < 1250:
        t = (section - 950) / 300
        offset = 180 - 360 * t
    elif section < 1500:
        offset = -180.0
    elif section < 1850:
        t = (section - 1500) / 350
        offset = -180 + 180 * smoothstep(t)
    elif section < 2100:
        offset = 0.0
    elif section < 2250:
        t = (section - 2100) / 150
        offset = 120 * t
    else:
        t = (section - 2250) / 150
        offset = 120 * (1 - t)

    return base + offset


def build_road_points(scroll: float) -> tuple[list[tuple[int, int]], list[tuple[int, int]]]:
    left_points: list[tuple[int, int]] = []
    right_points: list[tuple[int, int]] = []
    for y in range(-20, HEIGHT + 21, 8):
        world_y = scroll + (HEIGHT - y)
        center = road_center(world_y)
        left_points.append((round(center - ROAD_WIDTH / 2), y))
        right_points.append((round(center + ROAD_WIDTH / 2), y))
    return left_points, right_points


def get_sensor_data(scroll: float, car_x: float) -> dict[str, object]:
    """Map the visible game into proportional 50 x 39 matrix coordinates."""
    scale_x = (MATRIX_WIDTH - 1) / GAME_WIDTH
    scale_y = (MATRIX_HEIGHT - 1) / HEIGHT
    left_edge: list[tuple[int, int]] = []
    right_edge: list[tuple[int, int]] = []

    for game_y in range(0, HEIGHT + 1, 5):
        world_y = scroll + (HEIGHT - game_y)
        center = road_center(world_y)
        sensor_y = round(game_y * scale_y)
        left_x = center - ROAD_WIDTH / 2
        right_x = center + ROAD_WIDTH / 2
        left_edge.append((round((left_x - GAME_LEFT) * scale_x), sensor_y))
        right_edge.append((round((right_x - GAME_LEFT) * scale_x), sensor_y))

    center_at_car = road_center(scroll + (HEIGHT - CAR_Y))
    road_left = center_at_car - ROAD_WIDTH / 2
    road_right = center_at_car + ROAD_WIDTH / 2
    car_left = car_x - CAR_WIDTH / 2
    car_right = car_x + CAR_WIDTH / 2

    return {
        "left_edge": left_edge,
        "right_edge": right_edge,
        "car_x": round((car_x - GAME_LEFT) * scale_x),
        "car_y": round(CAR_Y * scale_y),
        "road_center_at_car": center_at_car,
        "road_left_at_car": road_left,
        "road_right_at_car": road_right,
        "edge_clearance": min(car_left - road_left, road_right - car_right),
        "offset": car_x - center_at_car,
    }


def raster_line(
    matrix: np.ndarray,
    start: tuple[int, int],
    end: tuple[int, int],
    value: int,
) -> None:
    """Draw a line into the integer AI matrix with Bresenham's algorithm."""
    x1, y1 = start
    x2, y2 = end
    dx = abs(x2 - x1)
    dy = -abs(y2 - y1)
    step_x = 1 if x1 < x2 else -1
    step_y = 1 if y1 < y2 else -1
    error = dx + dy

    while True:
        if 0 <= x1 < MATRIX_WIDTH and 0 <= y1 < MATRIX_HEIGHT:
            matrix[y1, x1] = value
        if x1 == x2 and y1 == y2:
            break
        error2 = 2 * error
        if error2 >= dy:
            error += dy
            x1 += step_x
        if error2 <= dx:
            error += dx
            y1 += step_y


def make_ai_matrix(sensor: dict[str, object]) -> np.ndarray:
    """Create the exact matrix supplied to PPO: 0 empty, 1 edge, 2 car."""
    matrix = np.zeros((MATRIX_HEIGHT, MATRIX_WIDTH), dtype=np.uint8)

    for edge_name in ("left_edge", "right_edge"):
        points = sensor[edge_name]
        for point1, point2 in zip(points, points[1:]):
            raster_line(matrix, point1, point2, 1)

    # Integer approximation of the 20%-smaller car in the low-resolution view.
    car_w, car_h = 3, 6
    car_left = round(sensor["car_x"] - car_w / 2)
    car_top = round(sensor["car_y"] - car_h / 2)
    for y in range(car_top, car_top + car_h):
        for x in range(car_left, car_left + car_w):
            if 0 <= x < MATRIX_WIDTH and 0 <= y < MATRIX_HEIGHT:
                matrix[y, x] = 2
    return matrix


def draw_car(surface: pygame.Surface, x: float, y: float) -> None:
    """Draw the 20%-smaller gray car from above."""
    body = pygame.Rect(0, 0, CAR_WIDTH, CAR_HEIGHT)
    body.center = (round(x), round(y))

    for wheel_x in (body.left - 4, body.right - 3):
        pygame.draw.rect(surface, (25, 25, 28), (wheel_x, body.top + 12, 7, 18), border_radius=3)
        pygame.draw.rect(surface, (25, 25, 28), (wheel_x, body.bottom - 30, 7, 18), border_radius=3)

    pygame.draw.rect(surface, (105, 110, 116), body, border_radius=11)
    pygame.draw.rect(surface, (155, 160, 166), body.inflate(-10, -9), border_radius=8)
    pygame.draw.polygon(
        surface,
        (205, 215, 220),
        [
            (body.left + 9, body.top + 16),
            (body.right - 9, body.top + 16),
            (body.right - 7, body.top + 34),
            (body.left + 7, body.top + 34),
        ],
    )
    pygame.draw.rect(
        surface,
        (175, 185, 190),
        (body.left + 8, body.bottom - 31, body.width - 16, 18),
        border_radius=4,
    )
    pygame.draw.rect(surface, (255, 245, 155), (body.left + 5, body.top + 4, 9, 6), border_radius=2)
    pygame.draw.rect(surface, (255, 245, 155), (body.right - 14, body.top + 4, 9, 6), border_radius=2)


class RoadCenterEnv(gym.Env):
    """Gymnasium environment in which PPO keeps the car in the road centre."""

    metadata = {"render_modes": ["human"], "render_fps": 60}
    ACTION_NAMES = ("STRAIGHT", "LEFT", "RIGHT")

    def __init__(self, render_mode: str | None = None):
        super().__init__()
        self.render_mode = render_mode
        self.action_space = spaces.Discrete(3)
        self.observation_space = spaces.Box(
            low=0.0,
            high=1.0,
            shape=(MATRIX_HEIGHT, MATRIX_WIDTH),
            dtype=np.float32,
        )

        self.window: pygame.Surface | None = None
        self.clock: pygame.time.Clock | None = None
        self.font: pygame.font.Font | None = None
        self.small_font: pygame.font.Font | None = None
        self.save_handler: Callable[[], None] | None = None

        self.simulation_fps = self.metadata["render_fps"]
        self.paused = False
        self.save_message = ""
        self.save_message_until = 0.0

        self.episode_number = 0
        self.global_step = 0
        self.total_crashes = 0
        self.total_completions = 0
        self.last_episode_return = math.nan
        self.last_done_reason = ""
        self.last_action = 0
        self.last_reward = 0.0
        self.last_d = 0.0
        self.last_edge_clearance = 0.0
        self.reward_history: deque[tuple[float, bool, bool]] = deque(maxlen=240)
        self.training_metrics = {
            "critic_value": math.nan,
            "critic_loss": math.nan,
            "actor_loss": math.nan,
            "entropy_loss": math.nan,
            "approx_kl": math.nan,
            "policy_mode": "WAITING",
            "explore_percent": 0.0,
            "exploit_percent": 0.0,
        }

    def set_save_handler(self, handler: Callable[[], None]) -> None:
        self.save_handler = handler

    def set_training_metrics(self, metrics: dict[str, object]) -> None:
        self.training_metrics.update(metrics)

    def reset(self, *, seed=None, options=None):
        super().reset(seed=seed)
        self.episode_number += 1
        self.episode_step = 0
        self.episode_return = 0.0
        self.distance = 0.0

        # Start at varied road phases and small offsets to avoid memorisation.
        self.scroll = float(self.np_random.uniform(0.0, 2400.0))
        center = road_center(self.scroll + (HEIGHT - CAR_Y))
        start_offset = 0.0 if self.episode_number <= 100 else float(self.np_random.uniform(-35.0, 35.0))
        self.car_x = center + start_offset
        self.lateral_speed = 0.0
        self.last_action = 0
        self.last_reward = 0.0
        self.last_d = 0.0
        self.last_edge_clearance = ROAD_WIDTH / 2 - CAR_WIDTH / 2
        self.last_done_reason = ""
        return self._observation(), {}

    def _sensor(self) -> dict[str, object]:
        return get_sensor_data(self.scroll, self.car_x)

    def _observation(self) -> np.ndarray:
        # Values become 0.0 empty, 0.5 boundaries and 1.0 car.
        return make_ai_matrix(self._sensor()).astype(np.float32) / 2.0

    def step(self, action):
        action = int(action)
        self.last_action = action
        self.episode_step += 1
        self.global_step += 1

        steering = 0.0
        if action == 1:
            steering = -1.0
        elif action == 2:
            steering = 1.0

        target_lateral_speed = steering * MAX_LATERAL_SPEED
        response = min(1.0, STEERING_RESPONSE * DT)
        self.lateral_speed += (target_lateral_speed - self.lateral_speed) * response
        self.car_x += self.lateral_speed * DT
        self.scroll += FORWARD_SPEED * DT
        self.distance += FORWARD_SPEED * DT

        sensor = self._sensor()
        offset = float(sensor["offset"])
        usable_half_width = ROAD_WIDTH / 2 - CAR_WIDTH / 2
        d = abs(offset) / usable_half_width
        edge_clearance = float(sensor["edge_clearance"])

        # Dense road-centre reward: +0.10 at centre and +0.02 at the edge.
        clipped_d = min(d, 1.0)
        reward = FORWARD_REWARD + CENTER_REWARD_SCALE * (1.0 - clipped_d**2)

        terminated = False
        truncated = False
        crashed = edge_clearance < 0.0
        completed = self.distance >= EPISODE_DISTANCE

        if crashed:
            reward = CRASH_REWARD
            terminated = True
            self.total_crashes += 1
            self.last_done_reason = "OFF ROAD"
        else:
            if edge_clearance <= NEAR_EDGE_PIXELS:
                reward += NEAR_EDGE_PENALTY
            if completed:
                reward += COMPLETE_REWARD
                terminated = True
                self.total_completions += 1
                self.last_done_reason = "TRACK COMPLETE"

        if self.episode_step >= MAX_EPISODE_STEPS and not terminated:
            truncated = True
            self.last_done_reason = "TIME LIMIT"

        self.last_reward = float(reward)
        self.last_d = float(d)
        self.last_edge_clearance = edge_clearance
        self.episode_return += float(reward)
        self.reward_history.append((float(reward), crashed, completed and not crashed))

        if terminated or truncated:
            self.last_episode_return = self.episode_return

        info = {
            "crashed": crashed,
            "completed": completed,
            "d": d,
            "edge_clearance": edge_clearance,
            "episode_return": self.episode_return,
        }

        if self.render_mode == "human":
            self.render()
        return self._observation(), float(reward), terminated, truncated, info

    def _handle_events(self) -> None:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                raise KeyboardInterrupt
            if event.type != pygame.KEYDOWN:
                continue

            if event.key == pygame.K_ESCAPE:
                raise KeyboardInterrupt
            if event.key == pygame.K_SPACE:
                self.paused = not self.paused
            elif event.key == pygame.K_q:
                if self.save_handler is None:
                    self.save_message = "No active PPO model"
                else:
                    try:
                        self.save_handler()
                        self.save_message = "MODEL SAVED"
                    except Exception as error:
                        self.save_message = f"SAVE ERROR: {error}"
                self.save_message_until = time.monotonic() + 3.0
            elif event.key in (pygame.K_PLUS, pygame.K_KP_PLUS, pygame.K_EQUALS):
                self.simulation_fps = min(1920, self.simulation_fps * 2)
            elif event.key in (pygame.K_MINUS, pygame.K_KP_MINUS):
                self.simulation_fps = max(2, self.simulation_fps // 2)

    def _draw_world(self) -> None:
        assert self.window is not None
        self.window.fill((31, 34, 38))
        pygame.draw.rect(self.window, (73, 148, 63), (GAME_LEFT, 0, GAME_WIDTH, HEIGHT))
        left, right = build_road_points(self.scroll)
        shoulder = 18
        outer_left = [(x - shoulder, y) for x, y in left]
        outer_right = [(x + shoulder, y) for x, y in right]
        pygame.draw.polygon(self.window, (200, 190, 145), outer_left + list(reversed(outer_right)))
        pygame.draw.polygon(self.window, (61, 64, 68), left + list(reversed(right)))
        pygame.draw.lines(self.window, (245, 245, 245), False, left, 5)
        pygame.draw.lines(self.window, (245, 245, 245), False, right, 5)
        draw_car(self.window, self.car_x, CAR_Y)

    def _draw_matrix(self, sensor: dict[str, object], ai_matrix: np.ndarray) -> None:
        assert self.window is not None
        assert self.small_font is not None
        logical = pygame.Surface((MATRIX_WIDTH, MATRIX_HEIGHT))
        logical.fill((5, 7, 9))
        colors = {1: (255, 255, 255), 2: (155, 155, 155)}
        for y, x in np.argwhere(ai_matrix > 0):
            logical.set_at((int(x), int(y)), colors[int(ai_matrix[y, x])])

        enlarged = pygame.transform.scale(logical, (SENSOR_WIDTH, SENSOR_HEIGHT))
        self.window.blit(enlarged, (SENSOR_X, SENSOR_Y))
        pygame.draw.rect(
            self.window,
            (210, 210, 210),
            (SENSOR_X, SENSOR_Y, SENSOR_WIDTH, SENSOR_HEIGHT),
            2,
        )
        self.window.blit(
            self.small_font.render("PPO INPUT 50 x 39", True, (255, 255, 255)),
            (SENSOR_X, 15),
        )
        self.window.blit(
            self.small_font.render(
                f"dx:{float(sensor['offset']):+6.1f}  d:{self.last_d:4.2f}",
                True,
                (185, 190, 195),
            ),
            (SENSOR_X, 40),
        )

    @staticmethod
    def _metric(value: object) -> str:
        try:
            number = float(value)
        except (TypeError, ValueError):
            return "waiting"
        return "waiting" if not math.isfinite(number) else f"{number:+.6f}"

    def _draw_information(self) -> None:
        assert self.window is not None
        assert self.small_font is not None
        metrics = self.training_metrics
        last_return = "waiting" if not math.isfinite(self.last_episode_return) else f"{self.last_episode_return:+.3f}"
        status = "PAUSED" if self.paused else "TRAINING"
        lines = (
            f"Status: {status}",
            f"Episode: {self.episode_number}",
            f"Global step: {self.global_step}",
            f"Action: {self.ACTION_NAMES[self.last_action]}",
            f"Reward: {self.last_reward:+.4f}",
            f"Episode return: {self.episode_return:+.3f}",
            f"Previous return: {last_return}",
            f"Edge clearance: {self.last_edge_clearance:5.1f}px",
            f"Crashes / complete: {self.total_crashes}/{self.total_completions}",
            f"Critic V(s): {self._metric(metrics['critic_value'])}",
            f"Critic loss: {self._metric(metrics['critic_loss'])}",
            f"Actor loss: {self._metric(metrics['actor_loss'])}",
            f"Entropy loss: {self._metric(metrics['entropy_loss'])}",
            f"Approx KL: {self._metric(metrics['approx_kl'])}",
            f"Mode: {metrics['policy_mode']}",
            f"Explore: {float(metrics['explore_percent']):5.1f}%",
            f"Exploit: {float(metrics['exploit_percent']):5.1f}%",
            f"Speed: {self.simulation_fps} steps/s",
            "PPO lr:3e-4 roll:512",
            "batch:256 gamma:0.995",
            "SPACE pause   Q save",
            "+/- speed   ESC stop",
        )
        for index, line in enumerate(lines):
            color = (255, 225, 120) if index == 0 else (220, 225, 232)
            self.window.blit(
                self.small_font.render(line, True, color),
                (SENSOR_X, 242 + index * 20),
            )

        if self.save_message and time.monotonic() <= self.save_message_until:
            message = self.small_font.render(self.save_message, True, (100, 255, 130))
            self.window.blit(message, (GAME_LEFT + 20, HEIGHT - 40))

        # Small reward explanation on the road display.
        panel = pygame.Surface((310, 100), pygame.SRCALPHA)
        panel.fill((0, 0, 0, 150))
        self.window.blit(panel, (GAME_LEFT + 15, 15))
        reward_lines = (
            "Reward = 0.02 + 0.08(1-d^2)",
            "<=10 px from edge: -0.15",
            "Off road: -1 and end",
            "Complete: +3 and end",
        )
        for index, line in enumerate(reward_lines):
            self.window.blit(
                self.small_font.render(line, True, (245, 245, 245)),
                (GAME_LEFT + 28, 27 + index * 21),
            )

    def _draw_reward_graph(self) -> None:
        """Draw recent rewards: green positive peaks and red crash peaks."""
        assert self.window is not None
        assert self.small_font is not None

        panel = pygame.Surface((GRAPH_WIDTH, GRAPH_HEIGHT), pygame.SRCALPHA)
        panel.fill((3, 7, 8, 220))
        self.window.blit(panel, (GRAPH_X, GRAPH_Y))
        pygame.draw.rect(
            self.window,
            (130, 145, 150),
            (GRAPH_X, GRAPH_Y, GRAPH_WIDTH, GRAPH_HEIGHT),
            2,
        )

        plot_left = GRAPH_X + 38
        plot_right = GRAPH_X + GRAPH_WIDTH - 10
        plot_top = GRAPH_Y + 27
        plot_bottom = GRAPH_Y + GRAPH_HEIGHT - 20
        reward_min, reward_max = -1.1, 3.2

        def reward_y(value: float) -> int:
            clipped = max(reward_min, min(reward_max, value))
            fraction = (reward_max - clipped) / (reward_max - reward_min)
            return round(plot_top + fraction * (plot_bottom - plot_top))

        zero_y = reward_y(0.0)
        pygame.draw.line(
            self.window,
            (90, 100, 105),
            (plot_left, zero_y),
            (plot_right, zero_y),
            1,
        )
        for value, label in ((3.0, "+3"), (0.0, "0"), (-1.0, "-1")):
            y = reward_y(value)
            self.window.blit(
                self.small_font.render(label, True, (170, 180, 185)),
                (GRAPH_X + 7, y - 8),
            )

        title = self.small_font.render(
            "LIVE REWARD   GREEN=REWARD  RED=CRASH",
            True,
            (225, 235, 238),
        )
        self.window.blit(title, (GRAPH_X + 10, GRAPH_Y + 6))

        history = list(self.reward_history)
        if not history:
            return

        x_step = (plot_right - plot_left) / max(1, self.reward_history.maxlen - 1)
        start_x = plot_right - x_step * (len(history) - 1)
        previous_point: tuple[int, int] | None = None
        for index, (reward, crashed, completed) in enumerate(history):
            x = round(start_x + index * x_step)
            y = reward_y(reward)
            color = (255, 45, 45) if crashed else (45, 235, 95)
            point = (x, y)

            if previous_point is not None:
                pygame.draw.line(self.window, color, previous_point, point, 2)
            if crashed:
                pygame.draw.line(self.window, (255, 35, 35), (x, zero_y), point, 3)
                pygame.draw.circle(self.window, (255, 55, 55), point, 4)
            elif completed:
                pygame.draw.line(self.window, (35, 255, 90), (x, zero_y), point, 3)
                pygame.draw.circle(self.window, (60, 255, 110), point, 4)
            previous_point = point

    def _draw_frame(self) -> None:
        assert self.window is not None
        self._draw_world()
        sensor = self._sensor()
        ai_matrix = make_ai_matrix(sensor)
        self._draw_matrix(sensor, ai_matrix)
        self._draw_information()

        if self.paused:
            overlay = pygame.Surface((GAME_WIDTH, HEIGHT), pygame.SRCALPHA)
            overlay.fill((0, 0, 0, 105))
            self.window.blit(overlay, (GAME_LEFT, 0))
            assert self.font is not None
            pause_text = self.font.render("PAUSED", True, (255, 235, 100))
            self.window.blit(
                pause_text,
                pause_text.get_rect(center=(GAME_LEFT + GAME_WIDTH / 2, HEIGHT / 2)),
            )

        self._draw_reward_graph()

        pygame.display.flip()

    def render(self):
        if self.window is None:
            pygame.init()
            pygame.display.set_caption("PPO Road-Centre Car Trainer")
            self.window = pygame.display.set_mode((WIDTH, HEIGHT))
            self.clock = pygame.time.Clock()
            self.font = pygame.font.SysFont("arial", 48, bold=True)
            self.small_font = pygame.font.SysFont("consolas", 15, bold=True)

        assert self.clock is not None
        self._handle_events()
        self._draw_frame()

        while self.paused:
            self._handle_events()
            self._draw_frame()
            self.clock.tick(30)

        self.clock.tick(self.simulation_fps)

    def close(self):
        if self.window is not None:
            pygame.quit()
            self.window = None


def model_exists(model_path: str) -> bool:
    return os.path.exists(model_path) or os.path.exists(model_path + ".zip")


def train(total_steps: int, render_training: bool, model_path: str, resume: bool) -> None:
    from stable_baselines3 import PPO
    from stable_baselines3.common.callbacks import BaseCallback
    from stable_baselines3.common.monitor import Monitor

    class PygameMetricsCallback(BaseCallback):
        """Send current PPO values, losses and exploration data to Pygame."""

        def _on_step(self) -> bool:
            values = self.locals.get("values")
            raw_actions = self.locals.get("actions")
            obs_tensor = self.locals.get("obs_tensor")
            policy_mode = "EXPLORE"

            if obs_tensor is not None and raw_actions is not None:
                sampled = np.asarray(raw_actions).reshape(-1)
                deterministic = (
                    self.model.policy._predict(obs_tensor, deterministic=True)
                    .detach()
                    .cpu()
                    .numpy()
                    .reshape(-1)
                )
                if sampled.size and deterministic.size and sampled[0] == deterministic[0]:
                    policy_mode = "EXPLOIT"

            if not hasattr(self, "explore_count"):
                self.explore_count = 0
                self.exploit_count = 0
            if policy_mode == "EXPLOIT":
                self.exploit_count += 1
            else:
                self.explore_count += 1
            decisions = self.explore_count + self.exploit_count

            log = self.model.logger.name_to_value
            metrics = {
                "critic_value": float(values.flatten()[0].item()) if values is not None else math.nan,
                "critic_loss": float(log.get("train/value_loss", math.nan)),
                "actor_loss": float(log.get("train/policy_gradient_loss", math.nan)),
                "entropy_loss": float(log.get("train/entropy_loss", math.nan)),
                "approx_kl": float(log.get("train/approx_kl", math.nan)),
                "policy_mode": policy_mode,
                "explore_percent": 100.0 * self.explore_count / decisions,
                "exploit_percent": 100.0 * self.exploit_count / decisions,
            }
            self.training_env.env_method("set_training_metrics", metrics)
            return True

    render_mode = "human" if render_training else None
    base_env = RoadCenterEnv(render_mode=render_mode)
    env = Monitor(base_env)

    resumed = resume and model_exists(model_path)
    if resumed:
        model = PPO.load(model_path, env=env, device="auto")
        print(f"Continuing model: {model_path}.zip")
    else:
        if resume:
            print(f"Model {model_path}.zip not found; starting a new model.")
        model = PPO(
            "MlpPolicy",
            env,
            learning_rate=LEARNING_RATE,
            n_steps=ROLLOUT_STEPS,
            batch_size=BATCH_SIZE,
            n_epochs=N_EPOCHS,
            gamma=GAMMA,
            gae_lambda=GAE_LAMBDA,
            ent_coef=ENTROPY_COEFFICIENT,
            clip_range=CLIP_RANGE,
            policy_kwargs={"net_arch": dict(pi=[128, 128], vf=[128, 128])},
            seed=42,
            device="auto",
            verbose=1,
        )

    base_env.set_save_handler(lambda: model.save(model_path))
    callback = PygameMetricsCallback()

    try:
        model.learn(
            total_timesteps=total_steps,
            callback=callback,
            reset_num_timesteps=not resumed,
            progress_bar=False,
        )
    except KeyboardInterrupt:
        print("Training stopped by user.")
    finally:
        model.save(model_path)
        env.close()
        print(f"Saved model: {model_path}.zip")


def play(model_path: str) -> None:
    from stable_baselines3 import PPO

    if not model_exists(model_path):
        raise FileNotFoundError(f"Model '{model_path}.zip' was not found. Train first.")

    env = RoadCenterEnv(render_mode="human")
    model = PPO.load(model_path, env=env, device="auto")
    env.set_save_handler(lambda: model.save(model_path))
    env.set_training_metrics(
        {
            "policy_mode": "EXPLOIT / FROZEN",
            "explore_percent": 0.0,
            "exploit_percent": 100.0,
        }
    )
    observation, _ = env.reset()
    try:
        while True:
            action, _ = model.predict(observation, deterministic=True)
            observation, _, terminated, truncated, _ = env.step(action)
            if terminated or truncated:
                observation, _ = env.reset()
    except KeyboardInterrupt:
        pass
    finally:
        env.close()


def parse_args():
    parser = argparse.ArgumentParser(description="PPO learns to drive in the centre of a curved road")
    parser.add_argument("--mode", choices=("train", "play"), default="train")
    parser.add_argument("--steps", type=int, default=500_000, help="PPO training timesteps")
    parser.add_argument("--model", default=MODEL_PATH, help="model filename without .zip")
    parser.add_argument("--resume", action="store_true", help="continue an existing model")
    parser.add_argument(
        "--headless",
        action="store_false",
        dest="render_train",
        help="disable Pygame during training for maximum speed",
    )
    parser.set_defaults(render_train=True)
    return parser.parse_args()


if __name__ == "__main__":
    arguments = parse_args()
    if arguments.mode == "train":
        train(arguments.steps, arguments.render_train, arguments.model, arguments.resume)
    else:
        play(arguments.model)

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

511

512

513

514

515

516

517

518

519

520

521

522

523

524

525

526

527

528

529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

547

548

549

550

551

552

553

554

555

556

557

558

559

560

561

562

563

564

565

566

567

568

569

570

571

572

573

574

575

576

577

578

579

580

581

582

583

584

585

586

587

588

589

590

591

592

593

594

595

596

597

598

599

600

601

602

603

604

605

606

607

608

609

610

611

612

613

614

615

616

617

618

619

620

621

622

623

624

625

626

627

628

629

630

631

632

633

634

635

636

637

638

639

640

641

642

643

644

645

646

647

648

649

650

651

652

653

654

655

656

657

658

659

660

661

662

663

664

665

666

667

668

669

670

671

672

673

674

675

676

677

678

679

680

681

682

683

684

685

686

687

688

689

690

691

692

693

694

695

696

697

698

699

700

701

702

703

704

705

706

707

708

709

710

711

712

713

714

715

716

717

718

719

720

721

722

723

724

725

726

727

728

729

730

731

732

733

734

735

736

737

738

739

740

741

742

743

744

745

746

747

748

749

750

751

752

753

754

755

756

757

758

759

760

761

762

763

764

765

766

767

768

769

770

771

772

773

774

775

776

777

778

779

780

781

782

783

784

785

786

787

788

789

790

791

792

793

794

795

796

797

798

799

800

801

802

803

804

805

806

807

808

809

810

"""

PPO road-centre driving trainer with a live Pygame display.

The PPO agent receives only the 50 x 39 AI matrix shown on the left:

0.0 = empty

0.5 = road boundary

1.0 = car

Install:

pip install pygame gymnasium stable-baselines3 numpy torch

Train and watch:

python ppo_car_trainer.py --mode train --steps 500000

Continue training a saved model:

python ppo_car_trainer.py --mode train --resume --steps 500000

Train without graphics:

python ppo_car_trainer.py --mode train --headless --steps 500000

Watch the saved model without learning:

python ppo_car_trainer.py --mode play

Live controls:

SPACE = pause / continue

Q = save the PPO model immediately

+ = faster display/training

- = slower display/training

ESC = stop and save

"""

from __future__ import annotations

import argparse

import math

import os

import time

from collections import deque

from collections.abc import Callable

import gymnasium as gym

import numpy as np

import pygame

from gymnasium import spaces

# Display and road geometry

MATRIX_WIDTH, MATRIX_HEIGHT = 50, 39

SENSOR_WIDTH, SENSOR_HEIGHT = 200, 156

SENSOR_X, SENSOR_Y = 10, 65

GAME_LEFT, GAME_WIDTH = 220, 900

WIDTH, HEIGHT = GAME_LEFT + GAME_WIDTH, 700

GRAPH_WIDTH, GRAPH_HEIGHT = 370, 180

GRAPH_X, GRAPH_Y = WIDTH - GRAPH_WIDTH - 15, 15

ROAD_WIDTH = 360

CAR_WIDTH, CAR_HEIGHT = 43, 77

CAR_Y = HEIGHT - 145

# Fixed simulation physics. PPO sees one action every 0.05 simulated seconds.

DT = 0.05

FORWARD_SPEED = 230.0

MAX_LATERAL_SPEED = 620.0

STEERING_RESPONSE = 12.0

EPISODE_DISTANCE = 2800.0

MAX_EPISODE_STEPS = 650

# Reward design agreed for road-centre driving.

FORWARD_REWARD = 0.02

CENTER_REWARD_SCALE = 0.08

NEAR_EDGE_PIXELS = 10.0

NEAR_EDGE_PENALTY = -0.15

CRASH_REWARD = -1.0

COMPLETE_REWARD = 3.0

# PPO parameters based on the supplied reference program.

LEARNING_RATE = 3e-4

ROLLOUT_STEPS = 512

BATCH_SIZE = 256

N_EPOCHS = 10

GAMMA = 0.995

GAE_LAMBDA = 0.95

ENTROPY_COEFFICIENT = 0.005

CLIP_RANGE = 0.2

MODEL_PATH = "ppo_road_center_car"

def road_center(world_y: float) -> float:

"""Return the road centre for straight, curved and linear sections."""

section = world_y % 2400.0

base = GAME_LEFT + GAME_WIDTH / 2

def smoothstep(value: float) -> float:

return value * value * (3.0 - 2.0 * value)

if section < 350:

offset = 0.0

elif section < 700:

t = (section - 350) / 350

offset = 180 * smoothstep(t)

elif section < 950:

offset = 180.0

elif section < 1250:

t = (section - 950) / 300

offset = 180 - 360 * t

elif section < 1500:

offset = -180.0

elif section < 1850:

t = (section - 1500) / 350

offset = -180 + 180 * smoothstep(t)

elif section < 2100:

offset = 0.0

elif section < 2250:

t = (section - 2100) / 150

offset = 120 * t

else:

t = (section - 2250) / 150

offset = 120 * (1 - t)

return base + offset

def build_road_points(scroll: float) -> tuple[list[tuple[int, int]], list[tuple[int, int]]]:

left_points: list[tuple[int, int]] = []

right_points: list[tuple[int, int]] = []

for y in range(-20, HEIGHT + 21, 8):

world_y = scroll + (HEIGHT - y)

center = road_center(world_y)

left_points.append((round(center - ROAD_WIDTH / 2), y))

right_points.append((round(center + ROAD_WIDTH / 2), y))

return left_points, right_points

def get_sensor_data(scroll: float, car_x: float) -> dict[str, object]:

"""Map the visible game into proportional 50 x 39 matrix coordinates."""

scale_x = (MATRIX_WIDTH - 1) / GAME_WIDTH

scale_y = (MATRIX_HEIGHT - 1) / HEIGHT

left_edge: list[tuple[int, int]] = []

right_edge: list[tuple[int, int]] = []

for game_y in range(0, HEIGHT + 1, 5):

world_y = scroll + (HEIGHT - game_y)

center = road_center(world_y)

sensor_y = round(game_y * scale_y)

left_x = center - ROAD_WIDTH / 2

right_x = center + ROAD_WIDTH / 2

left_edge.append((round((left_x - GAME_LEFT) * scale_x), sensor_y))

right_edge.append((round((right_x - GAME_LEFT) * scale_x), sensor_y))

center_at_car = road_center(scroll + (HEIGHT - CAR_Y))

road_left = center_at_car - ROAD_WIDTH / 2

road_right = center_at_car + ROAD_WIDTH / 2

car_left = car_x - CAR_WIDTH / 2

car_right = car_x + CAR_WIDTH / 2

return {

"left_edge": left_edge,

"right_edge": right_edge,

"car_x": round((car_x - GAME_LEFT) * scale_x),

"car_y": round(CAR_Y * scale_y),

"road_center_at_car": center_at_car,

"road_left_at_car": road_left,

"road_right_at_car": road_right,

"edge_clearance": min(car_left - road_left, road_right - car_right),

"offset": car_x - center_at_car,

}

def raster_line(

matrix: np.ndarray,

start: tuple[int, int],

end: tuple[int, int],

value: int,

) -> None:

"""Draw a line into the integer AI matrix with Bresenham's algorithm."""

x1, y1 = start

x2, y2 = end

dx = abs(x2 - x1)

dy = -abs(y2 - y1)

step_x = 1 if x1 < x2 else -1

step_y = 1 if y1 < y2 else -1

error = dx + dy

while True:

if 0 <= x1 < MATRIX_WIDTH and 0 <= y1 < MATRIX_HEIGHT:

matrix[y1, x1] = value

if x1 == x2 and y1 == y2:

break

error2 = 2 * error

if error2 >= dy:

error += dy

x1 += step_x

if error2 <= dx:

error += dx

y1 += step_y

def make_ai_matrix(sensor: dict[str, object]) -> np.ndarray:

"""Create the exact matrix supplied to PPO: 0 empty, 1 edge, 2 car."""

matrix = np.zeros((MATRIX_HEIGHT, MATRIX_WIDTH), dtype=np.uint8)

for edge_name in ("left_edge", "right_edge"):

points = sensor[edge_name]

for point1, point2 in zip(points, points[1:]):

raster_line(matrix, point1, point2, 1)

# Integer approximation of the 20%-smaller car in the low-resolution view.

car_w, car_h = 3, 6

car_left = round(sensor["car_x"] - car_w / 2)

car_top = round(sensor["car_y"] - car_h / 2)

for y in range(car_top, car_top + car_h):

for x in range(car_left, car_left + car_w):

if 0 <= x < MATRIX_WIDTH and 0 <= y < MATRIX_HEIGHT:

matrix[y, x] = 2

return matrix

def draw_car(surface: pygame.Surface, x: float, y: float) -> None:

"""Draw the 20%-smaller gray car from above."""

body = pygame.Rect(0, 0, CAR_WIDTH, CAR_HEIGHT)

body.center = (round(x), round(y))

for wheel_x in (body.left - 4, body.right - 3):

pygame.draw.rect(surface, (25, 25, 28), (wheel_x, body.top + 12, 7, 18), border_radius=3)

pygame.draw.rect(surface, (25, 25, 28), (wheel_x, body.bottom - 30, 7, 18), border_radius=3)

pygame.draw.rect(surface, (105, 110, 116), body, border_radius=11)

pygame.draw.rect(surface, (155, 160, 166), body.inflate(-10, -9), border_radius=8)

pygame.draw.polygon(

surface,

(205, 215, 220),

[

(body.left + 9, body.top + 16),

(body.right - 9, body.top + 16),

(body.right - 7, body.top + 34),

(body.left + 7, body.top + 34),

)

pygame.draw.rect(

surface,

(175, 185, 190),

(body.left + 8, body.bottom - 31, body.width - 16, 18),

border_radius=4,

)

pygame.draw.rect(surface, (255, 245, 155), (body.left + 5, body.top + 4, 9, 6), border_radius=2)

pygame.draw.rect(surface, (255, 245, 155), (body.right - 14, body.top + 4, 9, 6), border_radius=2)

class RoadCenterEnv(gym.Env):

"""Gymnasium environment in which PPO keeps the car in the road centre."""

metadata = {"render_modes": ["human"], "render_fps": 60}

ACTION_NAMES = ("STRAIGHT", "LEFT", "RIGHT")

def __init__(self, render_mode: str | None = None):

super().__init__()

self.render_mode = render_mode

self.action_space = spaces.Discrete(3)

self.observation_space = spaces.Box(

low=0.0,

high=1.0,

shape=(MATRIX_HEIGHT, MATRIX_WIDTH),

dtype=np.float32,

)

self.window: pygame.Surface | None = None

self.clock: pygame.time.Clock | None = None

self.font: pygame.font.Font | None = None

self.small_font: pygame.font.Font | None = None

self.save_handler: Callable[[], None] | None = None

self.simulation_fps = self.metadata["render_fps"]

self.paused = False

self.save_message = ""

self.save_message_until = 0.0

self.episode_number = 0

self.global_step = 0

self.total_crashes = 0

self.total_completions = 0

self.last_episode_return = math.nan

self.last_done_reason = ""

self.last_action = 0

self.last_reward = 0.0

self.last_d = 0.0

self.last_edge_clearance = 0.0

self.reward_history: deque[tuple[float, bool, bool]] = deque(maxlen=240)

self.training_metrics = {

"critic_value": math.nan,

"critic_loss": math.nan,

"actor_loss": math.nan,

"entropy_loss": math.nan,

"approx_kl": math.nan,

"policy_mode": "WAITING",

"explore_percent": 0.0,

"exploit_percent": 0.0,

}

def set_save_handler(self, handler: Callable[[], None]) -> None:

self.save_handler = handler

def set_training_metrics(self, metrics: dict[str, object]) -> None:

self.training_metrics.update(metrics)

def reset(self, *, seed=None, options=None):

super().reset(seed=seed)

self.episode_number += 1

self.episode_step = 0

self.episode_return = 0.0

self.distance = 0.0

# Start at varied road phases and small offsets to avoid memorisation.

self.scroll = float(self.np_random.uniform(0.0, 2400.0))

center = road_center(self.scroll + (HEIGHT - CAR_Y))

start_offset = 0.0 if self.episode_number <= 100 else float(self.np_random.uniform(-35.0, 35.0))

self.car_x = center + start_offset

self.lateral_speed = 0.0

self.last_action = 0

self.last_reward = 0.0

self.last_d = 0.0

self.last_edge_clearance = ROAD_WIDTH / 2 - CAR_WIDTH / 2

self.last_done_reason = ""

return self._observation(), {}

def _sensor(self) -> dict[str, object]:

return get_sensor_data(self.scroll, self.car_x)

def _observation(self) -> np.ndarray:

# Values become 0.0 empty, 0.5 boundaries and 1.0 car.

return make_ai_matrix(self._sensor()).astype(np.float32) / 2.0

def step(self, action):

action = int(action)

self.last_action = action

self.episode_step += 1

self.global_step += 1

steering = 0.0

if action == 1:

steering = -1.0

elif action == 2:

steering = 1.0

target_lateral_speed = steering * MAX_LATERAL_SPEED

response = min(1.0, STEERING_RESPONSE * DT)

self.lateral_speed += (target_lateral_speed - self.lateral_speed) * response

self.car_x += self.lateral_speed * DT

self.scroll += FORWARD_SPEED * DT

self.distance += FORWARD_SPEED * DT

sensor = self._sensor()

offset = float(sensor["offset"])

usable_half_width = ROAD_WIDTH / 2 - CAR_WIDTH / 2

d = abs(offset) / usable_half_width

edge_clearance = float(sensor["edge_clearance"])

# Dense road-centre reward: +0.10 at centre and +0.02 at the edge.

clipped_d = min(d, 1.0)

reward = FORWARD_REWARD + CENTER_REWARD_SCALE * (1.0 - clipped_d**2)

terminated = False

truncated = False

crashed = edge_clearance < 0.0

completed = self.distance >= EPISODE_DISTANCE

if crashed:

reward = CRASH_REWARD

terminated = True

self.total_crashes += 1

self.last_done_reason = "OFF ROAD"

else:

if edge_clearance <= NEAR_EDGE_PIXELS:

reward += NEAR_EDGE_PENALTY

if completed:

reward += COMPLETE_REWARD

terminated = True

self.total_completions += 1

self.last_done_reason = "TRACK COMPLETE"

if self.episode_step >= MAX_EPISODE_STEPS and not terminated:

truncated = True

self.last_done_reason = "TIME LIMIT"

self.last_reward = float(reward)

self.last_d = float(d)

self.last_edge_clearance = edge_clearance

self.episode_return += float(reward)

self.reward_history.append((float(reward), crashed, completed and not crashed))

if terminated or truncated:

self.last_episode_return = self.episode_return

info = {

"crashed": crashed,

"completed": completed,

"d": d,

"edge_clearance": edge_clearance,

"episode_return": self.episode_return,

}

if self.render_mode == "human":

self.render()

return self._observation(), float(reward), terminated, truncated, info

def _handle_events(self) -> None:

for event in pygame.event.get():

if event.type == pygame.QUIT:

raise KeyboardInterrupt

if event.type != pygame.KEYDOWN:

continue

if event.key == pygame.K_ESCAPE:

raise KeyboardInterrupt

if event.key == pygame.K_SPACE:

self.paused = not self.paused

elif event.key == pygame.K_q:

if self.save_handler is None:

self.save_message = "No active PPO model"

else:

try:

self.save_handler()

self.save_message = "MODEL SAVED"

except Exception as error:

self.save_message = f"SAVE ERROR: {error}"

self.save_message_until = time.monotonic() + 3.0

elif event.key in (pygame.K_PLUS, pygame.K_KP_PLUS, pygame.K_EQUALS):

self.simulation_fps = min(1920, self.simulation_fps * 2)

elif event.key in (pygame.K_MINUS, pygame.K_KP_MINUS):

self.simulation_fps = max(2, self.simulation_fps // 2)

def _draw_world(self) -> None:

assert self.window is not None

self.window.fill((31, 34, 38))

pygame.draw.rect(self.window, (73, 148, 63), (GAME_LEFT, 0, GAME_WIDTH, HEIGHT))

left, right = build_road_points(self.scroll)

shoulder = 18

outer_left = [(x - shoulder, y) for x, y in left]

outer_right = [(x + shoulder, y) for x, y in right]

pygame.draw.polygon(self.window, (200, 190, 145), outer_left + list(reversed(outer_right)))

pygame.draw.polygon(self.window, (61, 64, 68), left + list(reversed(right)))

pygame.draw.lines(self.window, (245, 245, 245), False, left, 5)

pygame.draw.lines(self.window, (245, 245, 245), False, right, 5)

draw_car(self.window, self.car_x, CAR_Y)

def _draw_matrix(self, sensor: dict[str, object], ai_matrix: np.ndarray) -> None:

assert self.window is not None

assert self.small_font is not None

logical = pygame.Surface((MATRIX_WIDTH, MATRIX_HEIGHT))

logical.fill((5, 7, 9))

colors = {1: (255, 255, 255), 2: (155, 155, 155)}

for y, x in np.argwhere(ai_matrix > 0):

logical.set_at((int(x), int(y)), colors[int(ai_matrix[y, x])])

enlarged = pygame.transform.scale(logical, (SENSOR_WIDTH, SENSOR_HEIGHT))

self.window.blit(enlarged, (SENSOR_X, SENSOR_Y))

pygame.draw.rect(

self.window,

(210, 210, 210),

(SENSOR_X, SENSOR_Y, SENSOR_WIDTH, SENSOR_HEIGHT),

)

self.window.blit(

self.small_font.render("PPO INPUT 50 x 39", True, (255, 255, 255)),

(SENSOR_X, 15),

)

self.window.blit(

self.small_font.render(

f"dx:{float(sensor['offset']):+6.1f} d:{self.last_d:4.2f}",

True,

(185, 190, 195),

(SENSOR_X, 40),

)

@staticmethod

def _metric(value: object) -> str:

try:

number = float(value)

except (TypeError, ValueError):

return "waiting"

return "waiting" if not math.isfinite(number) else f"{number:+.6f}"

def _draw_information(self) -> None:

assert self.window is not None

assert self.small_font is not None

metrics = self.training_metrics

last_return = "waiting" if not math.isfinite(self.last_episode_return) else f"{self.last_episode_return:+.3f}"

status = "PAUSED" if self.paused else "TRAINING"

lines = (

f"Status: {status}",

f"Episode: {self.episode_number}",

f"Global step: {self.global_step}",

f"Action: {self.ACTION_NAMES[self.last_action]}",

f"Reward: {self.last_reward:+.4f}",

f"Episode return: {self.episode_return:+.3f}",

f"Previous return: {last_return}",

f"Edge clearance: {self.last_edge_clearance:5.1f}px",

f"Crashes / complete: {self.total_crashes}/{self.total_completions}",

f"Critic V(s): {self._metric(metrics['critic_value'])}",

f"Critic loss: {self._metric(metrics['critic_loss'])}",

f"Actor loss: {self._metric(metrics['actor_loss'])}",

f"Entropy loss: {self._metric(metrics['entropy_loss'])}",

f"Approx KL: {self._metric(metrics['approx_kl'])}",

f"Mode: {metrics['policy_mode']}",

f"Explore: {float(metrics['explore_percent']):5.1f}%",

f"Exploit: {float(metrics['exploit_percent']):5.1f}%",

f"Speed: {self.simulation_fps} steps/s",

"PPO lr:3e-4 roll:512",

"batch:256 gamma:0.995",

"SPACE pause Q save",

"+/- speed ESC stop",

)

for index, line in enumerate(lines):

color = (255, 225, 120) if index == 0 else (220, 225, 232)

self.window.blit(

self.small_font.render(line, True, color),

(SENSOR_X, 242 + index * 20),

)

if self.save_message and time.monotonic() <= self.save_message_until:

message = self.small_font.render(self.save_message, True, (100, 255, 130))

self.window.blit(message, (GAME_LEFT + 20, HEIGHT - 40))

# Small reward explanation on the road display.

panel = pygame.Surface((310, 100), pygame.SRCALPHA)

panel.fill((0, 0, 0, 150))

self.window.blit(panel, (GAME_LEFT + 15, 15))

reward_lines = (

"Reward = 0.02 + 0.08(1-d^2)",

"<=10 px from edge: -0.15",

"Off road: -1 and end",

"Complete: +3 and end",

)

for index, line in enumerate(reward_lines):

self.window.blit(

self.small_font.render(line, True, (245, 245, 245)),

(GAME_LEFT + 28, 27 + index * 21),

)

def _draw_reward_graph(self) -> None:

"""Draw recent rewards: green positive peaks and red crash peaks."""

assert self.window is not None

assert self.small_font is not None

panel = pygame.Surface((GRAPH_WIDTH, GRAPH_HEIGHT), pygame.SRCALPHA)

panel.fill((3, 7, 8, 220))

self.window.blit(panel, (GRAPH_X, GRAPH_Y))

pygame.draw.rect(

self.window,

(130, 145, 150),

(GRAPH_X, GRAPH_Y, GRAPH_WIDTH, GRAPH_HEIGHT),

)

plot_left = GRAPH_X + 38

plot_right = GRAPH_X + GRAPH_WIDTH - 10

plot_top = GRAPH_Y + 27

plot_bottom = GRAPH_Y + GRAPH_HEIGHT - 20

reward_min, reward_max = -1.1, 3.2

def reward_y(value: float) -> int:

clipped = max(reward_min, min(reward_max, value))

fraction = (reward_max - clipped) / (reward_max - reward_min)

return round(plot_top + fraction * (plot_bottom - plot_top))

zero_y = reward_y(0.0)

pygame.draw.line(

self.window,

(90, 100, 105),

(plot_left, zero_y),

(plot_right, zero_y),

)

for value, label in ((3.0, "+3"), (0.0, "0"), (-1.0, "-1")):

y = reward_y(value)

self.window.blit(

self.small_font.render(label, True, (170, 180, 185)),

(GRAPH_X + 7, y - 8),

)

title = self.small_font.render(

"LIVE REWARD GREEN=REWARD RED=CRASH",

True,

(225, 235, 238),

)

self.window.blit(title, (GRAPH_X + 10, GRAPH_Y + 6))

history = list(self.reward_history)

if not history:

return

x_step = (plot_right - plot_left) / max(1, self.reward_history.maxlen - 1)

start_x = plot_right - x_step * (len(history) - 1)

previous_point: tuple[int, int] | None = None

for index, (reward, crashed, completed) in enumerate(history):

x = round(start_x + index * x_step)

y = reward_y(reward)

color = (255, 45, 45) if crashed else (45, 235, 95)

point = (x, y)

if previous_point is not None:

pygame.draw.line(self.window, color, previous_point, point, 2)

if crashed:

pygame.draw.line(self.window, (255, 35, 35), (x, zero_y), point, 3)

pygame.draw.circle(self.window, (255, 55, 55), point, 4)

elif completed:

pygame.draw.line(self.window, (35, 255, 90), (x, zero_y), point, 3)

pygame.draw.circle(self.window, (60, 255, 110), point, 4)

previous_point = point

def _draw_frame(self) -> None:

assert self.window is not None

self._draw_world()

sensor = self._sensor()

ai_matrix = make_ai_matrix(sensor)

self._draw_matrix(sensor, ai_matrix)

self._draw_information()

if self.paused:

overlay = pygame.Surface((GAME_WIDTH, HEIGHT), pygame.SRCALPHA)

overlay.fill((0, 0, 0, 105))

self.window.blit(overlay, (GAME_LEFT, 0))

assert self.font is not None

pause_text = self.font.render("PAUSED", True, (255, 235, 100))

self.window.blit(

pause_text,

pause_text.get_rect(center=(GAME_LEFT + GAME_WIDTH / 2, HEIGHT / 2)),

)

self._draw_reward_graph()

pygame.display.flip()

def render(self):

if self.window is None:

pygame.init()

pygame.display.set_caption("PPO Road-Centre Car Trainer")

self.window = pygame.display.set_mode((WIDTH, HEIGHT))

self.clock = pygame.time.Clock()

self.font = pygame.font.SysFont("arial", 48, bold=True)

self.small_font = pygame.font.SysFont("consolas", 15, bold=True)

assert self.clock is not None

self._handle_events()

self._draw_frame()

while self.paused:

self._handle_events()

self._draw_frame()

self.clock.tick(30)

self.clock.tick(self.simulation_fps)

def close(self):

if self.window is not None:

pygame.quit()

self.window = None

def model_exists(model_path: str) -> bool:

return os.path.exists(model_path) or os.path.exists(model_path + ".zip")

def train(total_steps: int, render_training: bool, model_path: str, resume: bool) -> None:

from stable_baselines3 import PPO

from stable_baselines3.common.callbacks import BaseCallback

from stable_baselines3.common.monitor import Monitor

class PygameMetricsCallback(BaseCallback):

"""Send current PPO values, losses and exploration data to Pygame."""

def _on_step(self) -> bool:

values = self.locals.get("values")

raw_actions = self.locals.get("actions")

obs_tensor = self.locals.get("obs_tensor")

policy_mode = "EXPLORE"

if obs_tensor is not None and raw_actions is not None:

sampled = np.asarray(raw_actions).reshape(-1)

deterministic = (

self.model.policy._predict(obs_tensor, deterministic=True)

.detach()

.cpu()

.numpy()

.reshape(-1)

)

if sampled.size and deterministic.size and sampled[0] == deterministic[0]:

policy_mode = "EXPLOIT"

if not hasattr(self, "explore_count"):

self.explore_count = 0

self.exploit_count = 0

if policy_mode == "EXPLOIT":

self.exploit_count += 1

else:

self.explore_count += 1

decisions = self.explore_count + self.exploit_count

log = self.model.logger.name_to_value

metrics = {

"critic_value": float(values.flatten()[0].item()) if values is not None else math.nan,

"critic_loss": float(log.get("train/value_loss", math.nan)),

"actor_loss": float(log.get("train/policy_gradient_loss", math.nan)),

"entropy_loss": float(log.get("train/entropy_loss", math.nan)),

"approx_kl": float(log.get("train/approx_kl", math.nan)),

"policy_mode": policy_mode,

"explore_percent": 100.0 * self.explore_count / decisions,

"exploit_percent": 100.0 * self.exploit_count / decisions,

}

self.training_env.env_method("set_training_metrics", metrics)

return True

render_mode = "human" if render_training else None

base_env = RoadCenterEnv(render_mode=render_mode)

env = Monitor(base_env)

resumed = resume and model_exists(model_path)

if resumed:

model = PPO.load(model_path, env=env, device="auto")

print(f"Continuing model: {model_path}.zip")

else:

if resume:

print(f"Model {model_path}.zip not found; starting a new model.")

model = PPO(

"MlpPolicy",

env,

learning_rate=LEARNING_RATE,

n_steps=ROLLOUT_STEPS,

batch_size=BATCH_SIZE,

n_epochs=N_EPOCHS,

gamma=GAMMA,

gae_lambda=GAE_LAMBDA,

ent_coef=ENTROPY_COEFFICIENT,

clip_range=CLIP_RANGE,

policy_kwargs={"net_arch": dict(pi=[128, 128], vf=[128, 128])},

seed=42,

device="auto",

verbose=1,

)

base_env.set_save_handler(lambda: model.save(model_path))

callback = PygameMetricsCallback()

try:

model.learn(

total_timesteps=total_steps,

callback=callback,

reset_num_timesteps=not resumed,

progress_bar=False,

)

except KeyboardInterrupt:

print("Training stopped by user.")

finally:

model.save(model_path)

env.close()

print(f"Saved model: {model_path}.zip")

def play(model_path: str) -> None:

from stable_baselines3 import PPO

if not model_exists(model_path):

raise FileNotFoundError(f"Model '{model_path}.zip' was not found. Train first.")

env = RoadCenterEnv(render_mode="human")

model = PPO.load(model_path, env=env, device="auto")

env.set_save_handler(lambda: model.save(model_path))

env.set_training_metrics(

{

"policy_mode": "EXPLOIT / FROZEN",

"explore_percent": 0.0,

"exploit_percent": 100.0,

}

)

observation, _ = env.reset()

try:

while True:

action, _ = model.predict(observation, deterministic=True)

observation, _, terminated, truncated, _ = env.step(action)

if terminated or truncated:

observation, _ = env.reset()

except KeyboardInterrupt:

pass

finally:

env.close()

def parse_args():

parser = argparse.ArgumentParser(description="PPO learns to drive in the centre of a curved road")

parser.add_argument("--mode", choices=("train", "play"), default="train")

parser.add_argument("--steps", type=int, default=500_000, help="PPO training timesteps")

parser.add_argument("--model", default=MODEL_PATH, help="model filename without .zip")

parser.add_argument("--resume", action="store_true", help="continue an existing model")

parser.add_argument(

"--headless",

action="store_false",

dest="render_train",

help="disable Pygame during training for maximum speed",

)

parser.set_defaults(render_train=True)

return parser.parse_args()

if __name__ == "__main__":

arguments = parse_args()

if arguments.mode == "train":

train(arguments.steps, arguments.render_train, arguments.model, arguments.resume)

else:

play(arguments.model)

RB1001-08 : סימולציה פשוטה מאוד של נהיגה אוטונומית דו־ממדית RL PPO

מה ה־Actor לומד?

מה ה־Critic לומד?

כיצד PPO לומד?

משמעות הגרף

נקודה טכנית חשובה מאוד

מגבלה נוספת

אולי תאהב/י גם

בינה מאלכותית RB14-01 : קורס תרגול רשתוות נירונים

קורס אמבדד ומיקרובקרים : R-STM32-1-002 – התקנת סביבת פיתוח STM32

פיתוח מעגל מבוסס ESP32 מיקרובקר – פיתוח אב טיפוב