| """ |
| =================================================================================================== |
| SOTA PING PONG REINFORCEMENT LEARNING TRAINING SYSTEM |
| =================================================================================================== |
| A production-grade, CPU-optimized Reinforcement Learning framework for training a world-class |
| Ping Pong bot using Proximal Policy Optimization (PPO) with an adaptive multi-opponent curriculum: |
| - 50% Logic Engines (10% Easy, 10% Medium, 30% Hard) |
| - 16% Current Self-Play |
| - 3% Random Policy |
| - 25% Historical Self-Play (sampling checkpoints from 5, 10, 15, and 25 checkpoints ago) |
| - 3% Minimax Lookahead (depth = 2) |
| - 3% Minimax Lookahead (depth = 1) |
| |
| All configurable hyperparameters are exposed below at the top of the file. |
| =================================================================================================== |
| """ |
|
|
| from __future__ import annotations |
| import os |
| import sys |
| import math |
| import copy |
| import time |
| import random |
| import argparse |
| from dataclasses import dataclass, field |
| from typing import List, Tuple, Dict, Optional, Any |
| from collections import deque |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.optim as optim |
| from torch.distributions.categorical import Categorical |
| from PIL import Image, ImageDraw |
| import imageio |
|
|
|
|
| |
| |
| |
|
|
| @dataclass |
| class OpponentDistributionConfig: |
| """ |
| Opponent sampling probabilities across training episodes. |
| Empirically tuned via 4,200-game round-robin tournament. |
| Total must sum to 1.0 (100%). |
| """ |
| easy_logic: float = 0.05 |
| medium_logic: float = 0.15 |
| realistic_hard_logic: float = 0.18 |
| impossible_hard_logic: float = 0.03 |
| self_play: float = 0.16 |
| random: float = 0.03 |
| historical_self_play: float = 0.18 |
| minimax_depth_1: float = 0.10 |
| minimax_depth_2: float = 0.12 |
|
|
| |
| historical_lags: List[int] = field(default_factory=lambda: [5, 10, 15, 25, 35, 50, 65, 75]) |
| min_required_checkpoint_lag: int = 5 |
| min_historical_lag: int = 5 |
| max_historical_lag: int = 75 |
|
|
| def validate(self): |
| total = (self.easy_logic + self.medium_logic + self.realistic_hard_logic + |
| self.impossible_hard_logic + self.self_play + self.random + |
| self.historical_self_play + self.minimax_depth_2 + self.minimax_depth_1) |
| assert abs(total - 1.0) < 1e-5, f"Opponent probabilities must sum to 1.0, got {total:.4f}" |
|
|
|
|
| @dataclass |
| class ModelConfig: |
| """ |
| Actor-Critic Neural Network Architecture. |
| CPU-optimized: keeps total parameters well under the 100k limit (~41.1k params). |
| """ |
| obs_dim: int = 16 |
| action_dim: int = 3 |
| hidden_dims: List[int] = field(default_factory=lambda: [192, 192]) |
| activation: str = "tanh" |
| max_allowed_params: int = 100_000 |
|
|
|
|
| @dataclass |
| class PPOHyperparameters: |
| """ |
| Proximal Policy Optimization (PPO) training hyperparameters. |
| """ |
| learning_rate: float = 3.5e-4 |
| lr_annealing: bool = True |
| gamma: float = 0.99 |
| gae_lambda: float = 0.95 |
| clip_epsilon: float = 0.20 |
| value_coef: float = 0.50 |
| entropy_coef: float = 0.02 |
| clip_value_loss: bool = True |
| max_grad_norm: float = 0.75 |
| num_epochs: int = 4 |
| mini_batch_size: int = 64 |
| rollout_steps: int = 128 |
| num_envs: int = 12 |
|
|
|
|
| @dataclass |
| class PhysicsConfig: |
| """ |
| Ping Pong Game & Simulation Physics. |
| Coordinates are normalized to ego-centric coordinates in [0, 1]. |
| """ |
| table_width: float = 800.0 |
| table_height: float = 500.0 |
| paddle_height: float = 80.0 |
| paddle_width: float = 14.0 |
| paddle_speed: float = 8.0 |
| paddle_inertia: float = 0.70 |
| frame_skip: int = 3 |
| ball_radius: float = 8.0 |
| ball_speed_initial: float = 7.5 |
| ball_speed_max: float = 16.0 |
| ball_acceleration: float = 1.035 |
| max_rally_steps: int = 1500 |
|
|
|
|
| @dataclass |
| class RewardConfig: |
| """ |
| Reward shaping values for policy training. |
| """ |
| win_point: float = 3.0 |
| lose_point: float = -2.0 |
| paddle_hit: float = 0.20 |
| tracking_reward: float = 0.002 |
| edge_hit_bonus: float = 0.50 |
| smoothness_penalty: float = 0.005 |
| centering_reward: float = 0.001 |
| step_survival_penalty: float = 0.0000 |
|
|
|
|
| @dataclass |
| class TrainingConfig: |
| """ |
| Global training session execution settings. |
| """ |
| total_timesteps: int = 10_000_000 |
| checkpoint_interval_steps: int = 35_000 |
| eval_interval_steps: int = 500_000 |
| log_interval_updates: int = 13 |
| eval_episodes: int = 15 |
| save_dir: str = "./checkpoints_pong" |
| resume: bool = False |
| resume_checkpoint_path: Optional[str] = None |
| seed: int = 42 |
| device: str = "cpu" |
|
|
|
|
| @dataclass |
| class VideoConfig: |
| """ |
| Gameplay Video Recording Configuration. |
| Automatically records full gameplay matches at periodic SAVE steps or checkpoints. |
| """ |
| enabled: bool = True |
| save_video_every_checkpoint: bool = False |
| video_interval_steps: int = 100_000 |
| record_episodes: int = 1 |
| fps: int = 30 |
| video_format: str = "mp4" |
| video_dir: str = "./videos_pong" |
| width: int = 800 |
| height: int = 480 |
|
|
|
|
| |
| @dataclass |
| class Config: |
| opponents: OpponentDistributionConfig = field(default_factory=OpponentDistributionConfig) |
| model: ModelConfig = field(default_factory=ModelConfig) |
| ppo: PPOHyperparameters = field(default_factory=PPOHyperparameters) |
| physics: PhysicsConfig = field(default_factory=PhysicsConfig) |
| reward: RewardConfig = field(default_factory=RewardConfig) |
| training: TrainingConfig = field(default_factory=TrainingConfig) |
| video: VideoConfig = field(default_factory=VideoConfig) |
|
|
|
|
| CONFIG = Config() |
|
|
|
|
| |
| |
| |
|
|
| class PongEnv: |
| """ |
| Continuous 2D physics Ping Pong environment with continuous kinematics, |
| paddle deflection mechanics, edge spin modulation, and ego-centric observations. |
| |
| Coordinate System: |
| - Origin (0,0) at Top-Left. |
| - X in [0, table_width] (0 = Left/Ego, table_width = Right/Opponent). |
| - Y in [0, table_height] (0 = Top wall, table_height = Bottom wall). |
| |
| Actions: |
| - 0: STAY |
| - 1: MOVE UP |
| - 2: MOVE DOWN |
| """ |
| def __init__(self, physics: PhysicsConfig = CONFIG.physics, reward_cfg: RewardConfig = CONFIG.reward, seed: Optional[int] = None): |
| self.phys = physics |
| self.rew = reward_cfg |
| self.rng = random.Random(seed) |
| self.np_rng = np.random.RandomState(seed) |
| |
| |
| self.ball_x: float = 0.0 |
| self.ball_y: float = 0.0 |
| self.ball_vx: float = 0.0 |
| self.ball_vy: float = 0.0 |
| |
| self.ego_y: float = 0.0 |
| self.ego_vy: float = 0.0 |
| self.opp_y: float = 0.0 |
| self.opp_vy: float = 0.0 |
| |
| self.prev_ego_action: int = 0 |
| self.prev_opp_action: int = 0 |
|
|
| self.step_count: int = 0 |
| self.rally_count: int = 0 |
| self.reset() |
|
|
| def reset(self, serve_direction: Optional[int] = None) -> np.ndarray: |
| """ |
| Reset environment for a new point. |
| serve_direction: 1 (to right/opponent) or -1 (to left/ego). |
| """ |
| self.step_count = 0 |
| self.rally_count = 0 |
| self.prev_ego_action = 0 |
| self.prev_opp_action = 0 |
| |
| |
| self.ego_y = self.phys.table_height / 2.0 |
| self.ego_vy = 0.0 |
| self.opp_y = self.phys.table_height / 2.0 |
| self.opp_vy = 0.0 |
| |
| |
| self.ball_x = self.phys.table_width / 2.0 |
| self.ball_y = self.phys.table_height / 2.0 |
| |
| |
| if serve_direction is None: |
| direction = 1.0 if self.rng.random() > 0.5 else -1.0 |
| else: |
| direction = float(serve_direction) |
| |
| angle = self.rng.uniform(-math.pi / 4, math.pi / 4) |
| speed = self.phys.ball_speed_initial |
| self.ball_vx = direction * speed * math.cos(angle) |
| self.ball_vy = speed * math.sin(angle) |
| |
| return self.get_ego_observation() |
|
|
| def _get_action_velocity(self, action: int) -> float: |
| if action == 1: |
| return -self.phys.paddle_speed |
| elif action == 2: |
| return self.phys.paddle_speed |
| return 0.0 |
|
|
| def _physics_substep(self, ego_action: int, opp_action: int) -> Tuple[float, bool, Dict[str, Any]]: |
| """Single physics sub-step with Continuous Collision Detection (CCD) and smooth momentum.""" |
| sub_reward = 0.0 |
| done = False |
| info = { |
| "hit_ego": False, |
| "hit_opp": False, |
| "winner": None, |
| "rally_count": self.rally_count |
| } |
|
|
| |
| prev_ego_y = self.ego_y |
| prev_opp_y = self.opp_y |
|
|
| ego_target_v = self._get_action_velocity(ego_action) |
| opp_target_v = self._get_action_velocity(opp_action) |
|
|
| alpha = self.phys.paddle_inertia |
| self.ego_vy = alpha * self.ego_vy + (1.0 - alpha) * ego_target_v |
| self.opp_vy = alpha * self.opp_vy + (1.0 - alpha) * opp_target_v |
|
|
| half_h = self.phys.paddle_height / 2.0 |
| self.ego_y = float(np.clip(self.ego_y + self.ego_vy, half_h, self.phys.table_height - half_h)) |
| self.opp_y = float(np.clip(self.opp_y + self.opp_vy, half_h, self.phys.table_height - half_h)) |
|
|
| |
| prev_ball_x = self.ball_x |
| prev_ball_y = self.ball_y |
| r = self.phys.ball_radius |
|
|
| ego_paddle_x = self.phys.paddle_width |
| opp_paddle_x = self.phys.table_width - self.phys.paddle_width |
|
|
| ego_impact_plane = ego_paddle_x + r |
| opp_impact_plane = opp_paddle_x - r |
|
|
| next_ball_x = prev_ball_x + self.ball_vx |
| next_ball_y = prev_ball_y + self.ball_vy |
|
|
| |
| hit_occurred = False |
|
|
| |
| if self.ball_vx < 0 and prev_ball_x >= ego_impact_plane and next_ball_x <= ego_impact_plane: |
| t = (prev_ball_x - ego_impact_plane) / max(1e-6, -self.ball_vx) |
| t = float(np.clip(t, 0.0, 1.0)) |
|
|
| y_ball_at_impact = prev_ball_y + t * self.ball_vy |
| y_ego_at_impact = prev_ego_y + t * (self.ego_y - prev_ego_y) |
|
|
| if abs(y_ball_at_impact - y_ego_at_impact) <= (half_h + r * 0.6): |
| hit_occurred = True |
| self.rally_count += 1 |
| info["hit_ego"] = True |
| sub_reward += self.rew.paddle_hit |
|
|
| offset = float(np.clip((y_ball_at_impact - y_ego_at_impact) / half_h, -1.0, 1.0)) |
| |
| sub_reward += abs(offset) * 0.35 |
| if abs(offset) > 0.55: |
| sub_reward += self.rew.edge_hit_bonus |
|
|
| bounce_angle = offset * (math.pi / 3.0) |
| current_speed = math.hypot(self.ball_vx, self.ball_vy) |
| new_speed = min(current_speed * self.phys.ball_acceleration, self.phys.ball_speed_max) |
|
|
| new_vx = new_speed * math.cos(bounce_angle) |
| new_vy = new_speed * math.sin(bounce_angle) + 0.25 * self.ego_vy |
|
|
| |
| if self.opp_y < self.phys.table_height * 0.45 and new_vy > 2.0: |
| sub_reward += 0.25 |
| elif self.opp_y > self.phys.table_height * 0.55 and new_vy < -2.0: |
| sub_reward += 0.25 |
|
|
| rem_dt = 1.0 - t |
| self.ball_x = ego_impact_plane + rem_dt * new_vx |
| self.ball_y = y_ball_at_impact + rem_dt * new_vy |
| self.ball_vx = new_vx |
| self.ball_vy = new_vy |
|
|
| |
| elif self.ball_vx > 0 and prev_ball_x <= opp_impact_plane and next_ball_x >= opp_impact_plane: |
| t = (opp_impact_plane - prev_ball_x) / max(1e-6, self.ball_vx) |
| t = float(np.clip(t, 0.0, 1.0)) |
|
|
| y_ball_at_impact = prev_ball_y + t * self.ball_vy |
| y_opp_at_impact = prev_opp_y + t * (self.opp_y - prev_opp_y) |
|
|
| if abs(y_ball_at_impact - y_opp_at_impact) <= (half_h + r * 0.6): |
| hit_occurred = True |
| self.rally_count += 1 |
| info["hit_opp"] = True |
|
|
| offset = float(np.clip((y_ball_at_impact - y_opp_at_impact) / half_h, -1.0, 1.0)) |
| bounce_angle = offset * (math.pi / 3.0) |
| current_speed = math.hypot(self.ball_vx, self.ball_vy) |
| new_speed = min(current_speed * self.phys.ball_acceleration, self.phys.ball_speed_max) |
|
|
| new_vx = -new_speed * math.cos(bounce_angle) |
| new_vy = new_speed * math.sin(bounce_angle) + 0.25 * self.opp_vy |
|
|
| rem_dt = 1.0 - t |
| self.ball_x = opp_impact_plane + rem_dt * new_vx |
| self.ball_y = y_ball_at_impact + rem_dt * new_vy |
| self.ball_vx = new_vx |
| self.ball_vy = new_vy |
|
|
| if not hit_occurred: |
| self.ball_x = next_ball_x |
| self.ball_y = next_ball_y |
|
|
| |
| if self.ball_y - r <= 0: |
| self.ball_y = r + abs(r - self.ball_y) |
| self.ball_vy = abs(self.ball_vy) |
| elif self.ball_y + r >= self.phys.table_height: |
| self.ball_y = (self.phys.table_height - r) - abs(self.ball_y + r - self.phys.table_height) |
| self.ball_vy = -abs(self.ball_vy) |
|
|
| |
| if (ego_action == 1 and self.prev_ego_action == 2) or (ego_action == 2 and self.prev_ego_action == 1): |
| sub_reward -= self.rew.smoothness_penalty |
|
|
| self.prev_ego_action = ego_action |
| self.prev_opp_action = opp_action |
|
|
| |
| if self.ball_x < 0: |
| done = True |
| sub_reward += self.rew.lose_point |
| info["winner"] = "opponent" |
| elif self.ball_x > self.phys.table_width: |
| done = True |
| sub_reward += self.rew.win_point |
| info["winner"] = "ego" |
|
|
| |
| if self.ball_vx < 0 and not done: |
| dist_norm = abs(self.ball_y - self.ego_y) / self.phys.table_height |
| sub_reward += self.rew.tracking_reward * (1.0 - dist_norm) |
| |
| if dist_norm < 0.08 and ego_action == 0: |
| sub_reward += 0.001 |
| elif self.ball_vx > 0 and not done: |
| |
| center_dist = abs(self.ego_y - self.phys.table_height / 2.0) / (self.phys.table_height / 2.0) |
| sub_reward += self.rew.centering_reward * (1.0 - center_dist) |
|
|
| return sub_reward, done, info |
|
|
| def step(self, ego_action: int, opp_action: int) -> Tuple[np.ndarray, float, bool, Dict[str, Any]]: |
| """ |
| Execute one RL decision step with frame_skip sub-stepping for smooth motion. |
| Returns: (observation, ego_reward, done, info) |
| """ |
| self.step_count += 1 |
| total_reward = 0.0 |
| done = False |
| combined_info = { |
| "hit_ego": False, |
| "hit_opp": False, |
| "winner": None, |
| "rally_count": self.rally_count |
| } |
|
|
| |
| for _ in range(self.phys.frame_skip): |
| r, d, info = self._physics_substep(ego_action, opp_action) |
| total_reward += r |
| if info["hit_ego"]: |
| combined_info["hit_ego"] = True |
| if info["hit_opp"]: |
| combined_info["hit_opp"] = True |
| if d: |
| done = True |
| combined_info["winner"] = info["winner"] |
| break |
|
|
| if not done and self.step_count >= self.phys.max_rally_steps: |
| done = True |
| combined_info["winner"] = "draw" |
|
|
| combined_info["rally_count"] = self.rally_count |
| return self.get_ego_observation(), total_reward, done, combined_info |
|
|
| def calculate_intercept_y(self, target_x: float, ball_x: float, ball_y: float, ball_vx: float, ball_vy: float) -> float: |
| """Computes exact multi-bounce raycast intercept Y on the plane x = target_x.""" |
| if (target_x > ball_x and ball_vx <= 0) or (target_x < ball_x and ball_vx >= 0): |
| return self.phys.table_height / 2.0 |
| |
| bx, by = float(ball_x), float(ball_y) |
| bvx, bvy = float(ball_vx), float(ball_vy) |
| h = self.phys.table_height |
| r = self.phys.ball_radius |
| max_bounces = 10 |
| bounce = 0 |
|
|
| while bounce < max_bounces: |
| bounce += 1 |
| dt_x = (target_x - bx) / bvx if bvx != 0 else float('inf') |
| if dt_x <= 0: |
| break |
| if bvy > 0: |
| dt_y = (h - r - by) / bvy |
| elif bvy < 0: |
| dt_y = (r - by) / bvy |
| else: |
| dt_y = float('inf') |
|
|
| if dt_x <= dt_y: |
| by += bvy * dt_x |
| break |
| else: |
| bx += bvx * dt_y |
| by += bvy * dt_y |
| bvy = -bvy |
|
|
| return float(np.clip(by, r, h - r)) |
|
|
| def get_ego_observation(self) -> np.ndarray: |
| """ |
| 16-dim normalized state vector from Ego's perspective: |
| [rel_ball_y, rel_ball_x, ball_vx, ball_vy, ego_y, ego_vy, rel_opp_y, opp_vy, |
| ball_y, ball_x, rel_pred_y, pred_norm_y, opp_y_norm, opp_open_top, opp_open_bottom, speed_norm] |
| All values scaled to [-1, 1] or [0, 1]. |
| """ |
| w, h = self.phys.table_width, self.phys.table_height |
| v_max = self.phys.ball_speed_max |
| pv_max = self.phys.paddle_speed |
| half_h = self.phys.paddle_height / 2.0 |
| ego_x = self.phys.paddle_width |
|
|
| pred_intercept_y = self.calculate_intercept_y(ego_x, self.ball_x, self.ball_y, self.ball_vx, self.ball_vy) |
| rel_pred_y = (pred_intercept_y - self.ego_y) / h |
| pred_norm_y = pred_intercept_y / h |
|
|
| opp_y_norm = self.opp_y / h |
| opp_open_top = (self.opp_y - half_h) / h |
| opp_open_bottom = (h - (self.opp_y + half_h)) / h |
| speed_norm = math.hypot(self.ball_vx, self.ball_vy) / v_max |
|
|
| obs = np.array([ |
| (self.ball_y - self.ego_y) / h, |
| (self.ball_x - ego_x) / w, |
| self.ball_vx / v_max, |
| self.ball_vy / v_max, |
| self.ego_y / h, |
| self.ego_vy / pv_max, |
| (self.opp_y - self.ego_y) / h, |
| self.opp_vy / pv_max, |
| self.ball_y / h, |
| self.ball_x / w, |
| rel_pred_y, |
| pred_norm_y, |
| opp_y_norm, |
| opp_open_top, |
| opp_open_bottom, |
| speed_norm |
| ], dtype=np.float32) |
| return obs |
|
|
| def get_opp_observation(self) -> np.ndarray: |
| """ |
| 16-dim normalized state vector from Opponent's perspective (horizontally flipped). |
| Allows any model/agent to play on the right side seamlessly with zero modification. |
| """ |
| w, h = self.phys.table_width, self.phys.table_height |
| v_max = self.phys.ball_speed_max |
| pv_max = self.phys.paddle_speed |
| half_h = self.phys.paddle_height / 2.0 |
| opp_x = self.phys.table_width - self.phys.paddle_width |
|
|
| pred_intercept_y = self.calculate_intercept_y(opp_x, self.ball_x, self.ball_y, self.ball_vx, self.ball_vy) |
| rel_pred_y = (pred_intercept_y - self.opp_y) / h |
| pred_norm_y = pred_intercept_y / h |
|
|
| ego_y_norm = self.ego_y / h |
| ego_open_top = (self.ego_y - half_h) / h |
| ego_open_bottom = (h - (self.ego_y + half_h)) / h |
| speed_norm = math.hypot(self.ball_vx, self.ball_vy) / v_max |
|
|
| obs = np.array([ |
| (self.ball_y - self.opp_y) / h, |
| (opp_x - self.ball_x) / w, |
| -self.ball_vx / v_max, |
| self.ball_vy / v_max, |
| self.opp_y / h, |
| self.opp_vy / pv_max, |
| (self.ego_y - self.opp_y) / h, |
| self.ego_vy / pv_max, |
| self.ball_y / h, |
| (w - self.ball_x) / w, |
| rel_pred_y, |
| pred_norm_y, |
| ego_y_norm, |
| ego_open_top, |
| ego_open_bottom, |
| speed_norm |
| ], dtype=np.float32) |
| return obs |
|
|
| def clone(self) -> PongEnv: |
| """Deep copy environment state for tree search / minimax simulation.""" |
| env = PongEnv(self.phys, self.rew) |
| env.ball_x = self.ball_x |
| env.ball_y = self.ball_y |
| env.ball_vx = self.ball_vx |
| env.ball_vy = self.ball_vy |
| env.ego_y = self.ego_y |
| env.ego_vy = self.ego_vy |
| env.opp_y = self.opp_y |
| env.opp_vy = self.opp_vy |
| env.prev_ego_action = self.prev_ego_action |
| env.prev_opp_action = self.prev_opp_action |
| env.step_count = self.step_count |
| env.rally_count = self.rally_count |
| return env |
|
|
|
|
| |
| |
| |
|
|
| class OpponentPolicy: |
| """Base interface for all Pong opponent policies.""" |
| def act(self, env: PongEnv) -> int: |
| raise NotImplementedError |
|
|
|
|
| class RandomOpponent(OpponentPolicy): |
| """3% Random uniform baseline.""" |
| def __init__(self, seed: Optional[int] = None): |
| self.rng = random.Random(seed) |
| |
| def act(self, env: PongEnv) -> int: |
| return self.rng.choice([0, 1, 2]) |
|
|
|
|
| def smooth_aim_action(target_y: float, current_y: float, prev_action: int, deadzone: float = 8.0, exit_zone: float = 2.5) -> int: |
| """ |
| Hysteresis (Schmitt Trigger) controller to prevent discrete action chattering / jitter. |
| Maintains directional momentum until target is reached, preventing 60Hz oscillation. |
| """ |
| diff = target_y - current_y |
| if prev_action == 0: |
| if abs(diff) > deadzone: |
| return 1 if diff < 0 else 2 |
| return 0 |
| elif prev_action == 1: |
| if diff >= -exit_zone: |
| return 0 if abs(diff) <= deadzone else (1 if diff < 0 else 2) |
| return 1 |
| elif prev_action == 2: |
| if diff <= exit_zone: |
| return 0 if abs(diff) <= deadzone else (1 if diff < 0 else 2) |
| return 2 |
| return 0 |
|
|
|
|
| class EasyLogicOpponent(OpponentPolicy): |
| """ |
| 10% Easy Logic: |
| - High tracking deadzone (+/- 30px) |
| - Reaction latency (recalculates every 6 frames) |
| - Smooth hysteresis positioning |
| """ |
| def __init__(self, seed: Optional[int] = None): |
| self.rng = random.Random(seed) |
| self.latency_counter = 0 |
| self.target_y = 250.0 |
| self.prev_action = 0 |
|
|
| def act(self, env: PongEnv) -> int: |
| self.latency_counter += 1 |
| if self.latency_counter % 6 == 0: |
| noise = self.rng.uniform(-30.0, 30.0) |
| self.target_y = env.ball_y + noise |
| |
| action = smooth_aim_action(self.target_y, env.opp_y, self.prev_action, deadzone=30.0, exit_zone=10.0) |
| self.prev_action = action |
| return action |
|
|
|
|
| class MediumLogicOpponent(OpponentPolicy): |
| """ |
| 10% Medium Logic: |
| - Moderate deadzone (+/- 14px) |
| - Smooth tracking with linear trajectory extrapolation with hysteresis damping. |
| """ |
| def __init__(self): |
| self.prev_action = 0 |
|
|
| def act(self, env: PongEnv) -> int: |
| if env.ball_vx > 0: |
| time_to_reach = (env.phys.table_width - env.phys.paddle_width - env.ball_x) / max(1e-5, env.ball_vx) |
| predicted_y = env.ball_y + env.ball_vy * time_to_reach |
| target_y = float(np.clip(predicted_y, 0, env.phys.table_height)) |
| else: |
| target_y = env.phys.table_height / 2.0 |
|
|
| action = smooth_aim_action(target_y, env.opp_y, self.prev_action, deadzone=14.0, exit_zone=4.0) |
| self.prev_action = action |
| return action |
|
|
|
|
| class RealisticHardLogicOpponent(OpponentPolicy): |
| """ |
| Realistic Hard Logic (Human-ish Grandmaster Table Tennis Pro): |
| - When ball is on far side (X < 480px / 60%): Holds balanced athletic center stance while tracking ball elevation. |
| - When ball crosses into near zone (X >= 480px): Commits to multi-bounce raycast with realistic human perceptual variance (+/- 12px). |
| - Smooth fluid paddle control with realistic reaction window. |
| """ |
| def __init__(self, commit_x_ratio: float = 0.60, seed: Optional[int] = None): |
| self.commit_x_ratio = commit_x_ratio |
| self.prev_action = 0 |
| self.rng = random.Random(seed) |
| self.perceptual_noise = 0.0 |
|
|
| def predict_intercept_y(self, env: PongEnv) -> float: |
| if env.ball_vx <= 0: |
| self.perceptual_noise = self.rng.uniform(-12.0, 12.0) |
| return env.phys.table_height / 2.0 |
|
|
| if env.ball_x < env.phys.table_width * self.commit_x_ratio: |
| return 0.7 * (env.phys.table_height / 2.0) + 0.3 * env.ball_y |
|
|
| target_x = env.phys.table_width - env.phys.paddle_width |
| exact_y = env.calculate_intercept_y(target_x, env.ball_x, env.ball_y, env.ball_vx, env.ball_vy) |
| return float(np.clip(exact_y + self.perceptual_noise, env.phys.ball_radius, env.phys.table_height - env.phys.ball_radius)) |
|
|
| def act(self, env: PongEnv) -> int: |
| target_y = self.predict_intercept_y(env) |
| action = smooth_aim_action(target_y, env.opp_y, self.prev_action, deadzone=8.0, exit_zone=2.0) |
| self.prev_action = action |
| return action |
|
|
|
|
| class ImpossibleHardLogicOpponent(OpponentPolicy): |
| """ |
| Impossible Hard Logic (0ms Zero-Latency Mathematical Wall): |
| - Instant 0ms raycasting across entire table. |
| - Zero perception delay with smooth anti-chatter tracking. |
| """ |
| def __init__(self): |
| self.prev_action = 0 |
|
|
| def predict_intercept_y(self, env: PongEnv) -> float: |
| if env.ball_vx <= 0: |
| return env.phys.table_height / 2.0 |
| target_x = env.phys.table_width - env.phys.paddle_width |
| return env.calculate_intercept_y(target_x, env.ball_x, env.ball_y, env.ball_vx, env.ball_vy) |
|
|
| def act(self, env: PongEnv) -> int: |
| target_y = self.predict_intercept_y(env) |
| action = smooth_aim_action(target_y, env.opp_y, self.prev_action, deadzone=6.0, exit_zone=1.5) |
| self.prev_action = action |
| return action |
|
|
|
|
| |
| HardLogicOpponent = ImpossibleHardLogicOpponent |
|
|
|
|
| class MinimaxOpponent(OpponentPolicy): |
| """ |
| High-Speed Minimax Search Opponent with forward simulation rollouts. |
| Includes action inertia bias to prevent direction oscillation. |
| Uses ultra-fast float scalar simulation (0 allocations) for 3,500+ FPS. |
| """ |
| def __init__(self, depth: int = 1, horizon_steps: int = 3): |
| self.depth = depth |
| self.horizon_steps = horizon_steps |
| self.prev_action = 0 |
|
|
| def evaluate_state_fast(self, bx: float, by: float, bvx: float, bvy: float, ey: float, oy: float, w: float = 800.0) -> float: |
| if bx > w: |
| return -1000.0 |
| if bx < 0: |
| return 1000.0 |
| score = 0.0 |
| if bvx > 0: |
| score -= abs(by - oy) * 2.0 |
| if bvx < 0: |
| score += abs(by - ey) * 1.5 |
| return score |
|
|
| def simulate_fast(self, bx: float, by: float, bvx: float, bvy: float, ey: float, oy: float, |
| evy: float, ovy: float, opp_a: int, ego_a: int, |
| w: float = 800.0, h: float = 500.0, pw: float = 14.0, ph: float = 80.0, |
| r: float = 8.0, ps: float = 8.0, alpha: float = 0.70, b_acc: float = 1.035, v_max: float = 16.0): |
| half_h = ph / 2.0 |
| ego_target_v = -ps if ego_a == 1 else (ps if ego_a == 2 else 0.0) |
| opp_target_v = -ps if opp_a == 1 else (ps if opp_a == 2 else 0.0) |
|
|
| for _ in range(self.horizon_steps * 3): |
| evy = alpha * evy + (1.0 - alpha) * ego_target_v |
| ovy = alpha * ovy + (1.0 - alpha) * opp_target_v |
| ey = max(half_h, min(h - half_h, ey + evy)) |
| oy = max(half_h, min(h - half_h, oy + ovy)) |
|
|
| bx += bvx |
| by += bvy |
|
|
| |
| if by - r <= 0: |
| by = r + abs(r - by) |
| bvy = abs(bvy) |
| elif by + r >= h: |
| by = (h - r) - abs(by + r - h) |
| bvy = -abs(bvy) |
|
|
| |
| ego_front = pw + r |
| opp_front = w - pw - r |
| if bvx < 0 and bx <= ego_front: |
| if abs(by - ey) <= (half_h + r * 0.6): |
| offset = max(-1.0, min(1.0, (by - ey) / half_h)) |
| speed = min(math.hypot(bvx, bvy) * b_acc, v_max) |
| angle = offset * (math.pi / 3.0) |
| bvx = speed * math.cos(angle) |
| bvy = speed * math.sin(angle) + 0.25 * evy |
| bx = ego_front |
| elif bvx > 0 and bx >= opp_front: |
| if abs(by - oy) <= (half_h + r * 0.6): |
| offset = max(-1.0, min(1.0, (by - oy) / half_h)) |
| speed = min(math.hypot(bvx, bvy) * b_acc, v_max) |
| angle = offset * (math.pi / 3.0) |
| bvx = -speed * math.cos(angle) |
| bvy = speed * math.sin(angle) + 0.25 * ovy |
| bx = opp_front |
|
|
| if bx < 0 or bx > w: |
| break |
| return bx, by, bvx, bvy, ey, oy, evy, ovy |
|
|
| def _minimax(self, bx: float, by: float, bvx: float, bvy: float, ey: float, oy: float, |
| evy: float, ovy: float, depth: int, is_opp_turn: bool) -> Tuple[float, int]: |
| if depth == 0 or bx < 0 or bx > 800.0: |
| return self.evaluate_state_fast(bx, by, bvx, bvy, ey, oy), 0 |
|
|
| best_action = 0 |
| if is_opp_turn: |
| best_val = -float('inf') |
| for action in [0, 1, 2]: |
| ego_a = 1 if by < ey else (2 if by > ey else 0) |
| nbx, nby, nbvx, nbvy, ney, noy, nevy, novy = self.simulate_fast(bx, by, bvx, bvy, ey, oy, evy, ovy, action, ego_a) |
| val, _ = self._minimax(nbx, nby, nbvx, nbvy, ney, noy, nevy, novy, depth - 1, False) |
| if action == self.prev_action: |
| val += 1.5 |
| elif (action == 1 and self.prev_action == 2) or (action == 2 and self.prev_action == 1): |
| val -= 2.0 |
| if val > best_val: |
| best_val = val |
| best_action = action |
| return best_val, best_action |
| else: |
| best_val = float('inf') |
| for action in [0, 1, 2]: |
| opp_a = 1 if by < oy else (2 if by > oy else 0) |
| nbx, nby, nbvx, nbvy, ney, noy, nevy, novy = self.simulate_fast(bx, by, bvx, bvy, ey, oy, evy, ovy, opp_a, action) |
| val, _ = self._minimax(nbx, nby, nbvx, nbvy, ney, noy, nevy, novy, depth - 1, True) |
| if val < best_val: |
| best_val = val |
| best_action = action |
| return best_val, best_action |
|
|
| def act(self, env: PongEnv) -> int: |
| if env.ball_vx <= 0: |
| center_y = env.phys.table_height / 2.0 |
| if env.opp_y < center_y - 14.0: |
| action = 2 |
| elif env.opp_y > center_y + 14.0: |
| action = 1 |
| else: |
| action = 0 |
| self.prev_action = action |
| return action |
|
|
| _, action = self._minimax( |
| env.ball_x, env.ball_y, env.ball_vx, env.ball_vy, |
| env.ego_y, env.opp_y, env.ego_vy, env.opp_vy, |
| self.depth, True |
| ) |
| self.prev_action = action |
| return action |
|
|
|
|
| class NeuralOpponent(OpponentPolicy): |
| """ |
| Neural Policy Opponent used for Current Self-Play and Historical Checkpoints. |
| Observes game through horizontally flipped coordinate system. |
| """ |
| def __init__(self, model: nn.Module, device: str = "cpu"): |
| self.model = model |
| self.device = device |
|
|
| def act(self, env: PongEnv) -> int: |
| obs = env.get_opp_observation() |
| obs_tensor = torch.tensor(obs, dtype=torch.float32, device=self.device).unsqueeze(0) |
| with torch.no_grad(): |
| logits, _ = self.model(obs_tensor) |
| dist = Categorical(logits=logits) |
| action = dist.sample().item() |
| return action |
|
|
|
|
| class OpponentManager: |
| """ |
| Manages the multi-opponent pool, checkpoint history, and sampling distribution: |
| - 50% Logic (10% Easy, 10% Medium, 20% Realistic Hard, 10% Impossible Hard) |
| - 16% Current Self-Play |
| - 3% Random |
| - 25% Historical Self-Play (Lags: 5, 10, 15, 25) |
| - 3% Minimax Depth 2 |
| - 3% Minimax Depth 1 |
| """ |
| def __init__(self, config: OpponentDistributionConfig, current_model: nn.Module, device: str = "cpu"): |
| self.cfg = config |
| self.cfg.validate() |
| self.current_model = current_model |
| self.device = device |
| |
| |
| self.checkpoints: List[dict] = [] |
|
|
| def save_checkpoint(self, model: nn.Module): |
| """Register a new policy snapshot into historical buffer.""" |
| state_dict_clone = copy.deepcopy(model.state_dict()) |
| self.checkpoints.append(state_dict_clone) |
|
|
| def sample_opponent(self) -> Tuple[OpponentPolicy, str]: |
| """ |
| Sample an opponent following the configured probability distribution. |
| Returns a fresh independent instance to prevent state crosstalk across parallel environments. |
| """ |
| r = random.random() |
| c = self.cfg |
|
|
| |
| if r < c.easy_logic: |
| return EasyLogicOpponent(), "logic_easy" |
| r -= c.easy_logic |
|
|
| if r < c.medium_logic: |
| return MediumLogicOpponent(), "logic_medium" |
| r -= c.medium_logic |
|
|
| if r < c.realistic_hard_logic: |
| return RealisticHardLogicOpponent(), "logic_hard_realistic" |
| r -= c.realistic_hard_logic |
|
|
| if r < c.impossible_hard_logic: |
| return ImpossibleHardLogicOpponent(), "logic_hard_impossible" |
| r -= c.impossible_hard_logic |
|
|
| |
| if r < c.random: |
| return RandomOpponent(), "random" |
| r -= c.random |
|
|
| |
| if r < c.minimax_depth_1: |
| return MinimaxOpponent(depth=1), "minimax_d1" |
| r -= c.minimax_depth_1 |
|
|
| if r < c.minimax_depth_2: |
| return MinimaxOpponent(depth=2), "minimax_depth_2" |
| r -= c.minimax_depth_2 |
|
|
| |
| if r < c.self_play: |
| return NeuralOpponent(self.current_model, self.device), "self_play_current" |
| r -= c.self_play |
|
|
| |
| chosen_lag = random.choice(self.cfg.historical_lags) |
| num_checkpoints = len(self.checkpoints) |
| |
| if num_checkpoints >= chosen_lag: |
| target_idx = num_checkpoints - chosen_lag |
| hist_model = copy.deepcopy(self.current_model) |
| hist_model.load_state_dict(self.checkpoints[target_idx]) |
| hist_model.eval() |
| return NeuralOpponent(hist_model, self.device), f"historical_lag_{chosen_lag}" |
| elif num_checkpoints >= self.cfg.min_required_checkpoint_lag: |
| valid_lags = [l for l in self.cfg.historical_lags if l <= num_checkpoints] |
| fallback_lag = random.choice(valid_lags) |
| target_idx = num_checkpoints - fallback_lag |
| hist_model = copy.deepcopy(self.current_model) |
| hist_model.load_state_dict(self.checkpoints[target_idx]) |
| hist_model.eval() |
| return NeuralOpponent(hist_model, self.device), f"historical_lag_{fallback_lag}" |
| else: |
| |
| return RealisticHardLogicOpponent(), "historical_inactive_fallback" |
|
|
|
|
| |
| |
| |
|
|
| def layer_init(layer: nn.Linear, std: float = np.sqrt(2), bias_const: float = 0.0) -> nn.Linear: |
| """Orthogonal initialization for high-stability RL training.""" |
| nn.init.orthogonal_(layer.weight, std) |
| nn.init.constant_(layer.bias, bias_const) |
| return layer |
|
|
|
|
| class ActorCritic(nn.Module): |
| """ |
| Lightweight, SOTA Actor-Critic MLP architecture. |
| Designed for fast CPU cache residency and low latency forward passes. |
| |
| Total Parameters: ~18,180 parameters (well within the <100k constraint). |
| """ |
| def __init__(self, cfg: ModelConfig = CONFIG.model): |
| super().__init__() |
| self.cfg = cfg |
| |
| |
| if cfg.activation.lower() == "tanh": |
| act_cls = nn.Tanh |
| elif cfg.activation.lower() == "gelu": |
| act_cls = nn.GELU |
| else: |
| act_cls = nn.ReLU |
|
|
| |
| layers = [] |
| prev_dim = cfg.obs_dim |
| for hidden_dim in cfg.hidden_dims: |
| layers.append(layer_init(nn.Linear(prev_dim, hidden_dim))) |
| layers.append(act_cls()) |
| prev_dim = hidden_dim |
| |
| self.trunk = nn.Sequential(*layers) |
| |
| |
| self.actor = layer_init(nn.Linear(prev_dim, cfg.action_dim), std=0.01) |
| |
| |
| self.critic = layer_init(nn.Linear(prev_dim, 1), std=1.0) |
| |
| |
| total_params = sum(p.numel() for p in self.parameters() if p.requires_grad) |
| assert total_params <= cfg.max_allowed_params, ( |
| f"Model exceeds maximum parameter budget! ({total_params} > {cfg.max_allowed_params})" |
| ) |
|
|
| def get_value(self, x: torch.Tensor) -> torch.Tensor: |
| """Compute state value estimate V(s).""" |
| features = self.trunk(x) |
| return self.critic(features).squeeze(-1) |
|
|
| def get_action_and_value(self, x: torch.Tensor, action: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| """ |
| Evaluate policy and value heads for observation batch x. |
| Returns: (action, log_prob, entropy, state_value) |
| """ |
| features = self.trunk(x) |
| logits = self.actor(features) |
| dist = Categorical(logits=logits) |
| |
| if action is None: |
| action = dist.sample() |
| |
| return action, dist.log_prob(action), dist.entropy(), self.critic(features).squeeze(-1) |
|
|
| def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: |
| """Direct forward pass returning logits and value.""" |
| features = self.trunk(x) |
| return self.actor(features), self.critic(features).squeeze(-1) |
|
|
|
|
| |
| |
| |
|
|
| class VectorPongRolloutWorker: |
| """ |
| Manages parallel rollout environments on CPU with per-episode dynamic opponent sampling |
| and live rolling win/loss statistics tracking across all opponent categories. |
| """ |
| def __init__(self, num_envs: int, opp_manager: OpponentManager, seed: int = 42, history_window: int = 100): |
| self.num_envs = num_envs |
| self.opp_manager = opp_manager |
| self.envs = [PongEnv(seed=seed + i) for i in range(num_envs)] |
| self.opponents: List[OpponentPolicy] = [] |
| self.opp_names: List[str] = [] |
| |
| |
| self.history_window = history_window |
| self.category_keys = [ |
| "Easy Logic", "Medium Logic", "Realistic Hard", "Impossible Hard", |
| "Current Self-Play", "Historical Play", "Random Agent", |
| "Minimax Depth 1", "Minimax Depth 2" |
| ] |
| self.match_history: Dict[str, deque] = { |
| k: deque(maxlen=history_window) for k in self.category_keys |
| } |
| self.cumulative_stats: Dict[str, Dict[str, int]] = { |
| k: {"wins": 0, "draws": 0, "losses": 0, "total": 0} for k in self.category_keys |
| } |
| |
| |
| for env in self.envs: |
| opp, name = self.opp_manager.sample_opponent() |
| self.opponents.append(opp) |
| self.opp_names.append(name) |
| |
| self.obs = np.array([env.reset() for env in self.envs], dtype=np.float32) |
|
|
| def _map_category_name(self, raw_name: str) -> str: |
| if raw_name == "logic_easy": |
| return "Easy Logic" |
| elif raw_name == "logic_medium": |
| return "Medium Logic" |
| elif raw_name in ["logic_hard_realistic", "historical_inactive_fallback"]: |
| return "Realistic Hard" |
| elif raw_name in ["logic_hard_impossible", "logic_hard"]: |
| return "Impossible Hard" |
| elif raw_name == "self_play_current": |
| return "Current Self-Play" |
| elif raw_name.startswith("historical_lag_"): |
| return "Historical Play" |
| elif raw_name == "random": |
| return "Random Agent" |
| elif raw_name == "minimax_d1": |
| return "Minimax Depth 1" |
| elif raw_name in ["minimax_depth_2", "minimax_d2"]: |
| return "Minimax Depth 2" |
| return "Other" |
|
|
| def step(self, ego_actions: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, List[Dict[str, Any]]]: |
| """ |
| Advance all parallel environments by one step. |
| Automatically handles opponent actions, point completions, and win/loss logging. |
| """ |
| next_obs = np.zeros_like(self.obs) |
| rewards = np.zeros(self.num_envs, dtype=np.float32) |
| dones = np.zeros(self.num_envs, dtype=bool) |
| infos = [] |
|
|
| for i, (env, opp) in enumerate(zip(self.envs, self.opponents)): |
| opp_act = opp.act(env) |
| o, r, d, info = env.step(ego_actions[i], opp_act) |
| |
| rewards[i] = r |
| dones[i] = d |
| infos.append(info) |
| |
| if d: |
| |
| winner = info.get("winner") |
| cat = self._map_category_name(self.opp_names[i]) |
| if cat in self.match_history: |
| if winner == "ego": |
| self.match_history[cat].append("W") |
| self.cumulative_stats[cat]["wins"] += 1 |
| elif winner == "draw": |
| self.match_history[cat].append("D") |
| self.cumulative_stats[cat]["draws"] += 1 |
| elif winner == "opponent": |
| self.match_history[cat].append("L") |
| self.cumulative_stats[cat]["losses"] += 1 |
| self.cumulative_stats[cat]["total"] += 1 |
|
|
| |
| next_obs[i] = env.reset() |
| new_opp, new_name = self.opp_manager.sample_opponent() |
| self.opponents[i] = new_opp |
| self.opp_names[i] = new_name |
| else: |
| next_obs[i] = o |
|
|
| self.obs = next_obs |
| return next_obs, rewards, dones, infos |
|
|
| def get_live_match_stats(self) -> Dict[str, Dict[str, Any]]: |
| """ |
| Returns detailed live match breakdown per opponent category (both rolling window and lifetime): |
| {category: {win_rate, draw_rate, loss_rate, wins, draws, losses, total, cum_wins, cum_draws, cum_losses, cum_total, cum_win_rate}} |
| """ |
| stats = {} |
| for cat in self.category_keys: |
| history = self.match_history[cat] |
| total = len(history) |
| cum = self.cumulative_stats[cat] |
| cum_total = cum["total"] |
| cum_win_rate = (cum["wins"] / cum_total) if cum_total > 0 else 0.0 |
|
|
| if total > 0: |
| wins = sum(1 for x in history if x == "W" or x == 1) |
| draws = sum(1 for x in history if x == "D" or x == 0.5) |
| losses = sum(1 for x in history if x == "L" or x == 0) |
| stats[cat] = { |
| "win_rate": wins / total, |
| "draw_rate": draws / total, |
| "loss_rate": losses / total, |
| "wins": wins, |
| "draws": draws, |
| "losses": losses, |
| "total": total, |
| "cum_wins": cum["wins"], |
| "cum_draws": cum["draws"], |
| "cum_losses": cum["losses"], |
| "cum_total": cum_total, |
| "cum_win_rate": cum_win_rate |
| } |
| else: |
| stats[cat] = { |
| "win_rate": 0.0, |
| "draw_rate": 0.0, |
| "loss_rate": 0.0, |
| "wins": 0, |
| "draws": 0, |
| "losses": 0, |
| "total": 0, |
| "cum_wins": cum["wins"], |
| "cum_draws": cum["draws"], |
| "cum_losses": cum["losses"], |
| "cum_total": cum_total, |
| "cum_win_rate": cum_win_rate |
| } |
| return stats |
|
|
| def get_live_win_rates(self) -> Dict[str, Tuple[float, int]]: |
| """Backward compatible helper returning (win_rate, total_played).""" |
| stats = {} |
| for cat in self.category_keys: |
| history = self.match_history[cat] |
| total = len(history) |
| if total > 0: |
| wins = sum(1 for x in history if x == "W" or x == 1) |
| stats[cat] = (wins / total, total) |
| else: |
| stats[cat] = (0.0, 0) |
| return stats |
|
|
|
|
| |
| |
| |
|
|
| class PongRenderer: |
| """ |
| High-contrast 2D Game Renderer for Ping Pong Video Recording. |
| Draws table court, net, paddles, glowing ball, and real-time telemetry HUD overlay. |
| """ |
| def __init__(self, phys: PhysicsConfig, cfg: VideoConfig): |
| self.phys = phys |
| self.cfg = cfg |
| self.w = cfg.width |
| self.h = cfg.height |
| self.scale_x = cfg.width / phys.table_width |
| self.scale_y = cfg.height / phys.table_height |
|
|
| def render_frame( |
| self, |
| env: PongEnv, |
| step_idx: int, |
| ego_score: int, |
| opp_score: int, |
| opp_name: str, |
| global_step: int, |
| ego_act: int, |
| opp_act: int |
| ) -> np.ndarray: |
| img = Image.new("RGB", (self.w, self.h), color=(15, 23, 42)) |
| draw = ImageDraw.Draw(img) |
|
|
| |
| draw.rectangle([8, 8, self.w - 8, self.h - 8], outline=(51, 65, 85), width=3) |
| center_x = self.w // 2 |
| for y in range(16, self.h - 16, 24): |
| draw.line([(center_x, y), (center_x, y + 12)], fill=(71, 85, 105), width=2) |
|
|
| |
| |
| p_w = max(6, int(self.phys.paddle_width * self.scale_x)) |
| p_h = max(12, int(self.phys.paddle_height * self.scale_y)) |
| ego_x_px = int(self.phys.paddle_width * self.scale_x) |
| ego_y_px = int(env.ego_y * self.scale_y) |
| draw.rectangle( |
| [ego_x_px - p_w, ego_y_px - p_h // 2, ego_x_px, ego_y_px + p_h // 2], |
| fill=(56, 189, 248), |
| outline=(14, 165, 233), |
| width=1 |
| ) |
|
|
| |
| opp_x_px = int((self.phys.table_width - self.phys.paddle_width) * self.scale_x) |
| opp_y_px = int(env.opp_y * self.scale_y) |
| draw.rectangle( |
| [opp_x_px, opp_y_px - p_h // 2, opp_x_px + p_w, opp_y_px + p_h // 2], |
| fill=(251, 113, 133), |
| outline=(244, 63, 94), |
| width=1 |
| ) |
|
|
| |
| bx = int(env.ball_x * self.scale_x) |
| by = int(env.ball_y * self.scale_y) |
| br = max(4, int(self.phys.ball_radius * self.scale_x)) |
| draw.ellipse([bx - br - 2, by - br - 2, bx + br + 2, by + br + 2], fill=(254, 240, 138)) |
| draw.ellipse([bx - br, by - br, bx + br, by + br], fill=(255, 255, 255)) |
|
|
| |
| act_labels = ["STAY", "UP", "DOWN"] |
| ego_txt = f"AGENT [P1]: {ego_score} ({act_labels[ego_act]})" |
| opp_txt = f"{opp_name.upper()} [P2]: {opp_score} ({act_labels[opp_act]})" |
|
|
| |
| draw.text((24, 16), ego_txt, fill=(56, 189, 248)) |
|
|
| |
| draw.text((self.w - 360, 16), opp_txt, fill=(251, 113, 133)) |
|
|
| |
| draw.text((self.w // 2 - 15, 16), "VS", fill=(148, 163, 184)) |
|
|
| speed = math.hypot(env.ball_vx, env.ball_vy) |
| telemetry = ( |
| f"Step: {global_step:,} | Match: AGENT vs {opp_name} | " |
| f"Rally: {env.rally_count} hits | Ball Speed: {speed:.1f} px/f" |
| ) |
| draw.text((24, self.h - 28), telemetry, fill=(148, 163, 184)) |
|
|
| return np.array(img, dtype=np.uint8) |
|
|
|
|
| |
| |
| |
|
|
| class PPOTrainer: |
| """ |
| High-performance, stable PPO Training Engine with GAE, LR Annealing, and Opponent Tracking. |
| """ |
| def __init__(self, config: Config = CONFIG): |
| self.cfg = config |
| self.device = torch.device(config.training.device if torch.cuda.is_available() else "cpu") |
| |
| |
| torch.manual_seed(config.training.seed) |
| np.random.seed(config.training.seed) |
| random.seed(config.training.seed) |
| |
| |
| self.agent = ActorCritic(config.model).to(self.device) |
| self.optimizer = optim.AdamW( |
| self.agent.parameters(), |
| lr=config.ppo.learning_rate, |
| eps=1e-5, |
| weight_decay=1e-4 |
| ) |
| |
| |
| self.opp_manager = OpponentManager(config.opponents, self.agent, device=str(self.device)) |
| self.vector_worker = VectorPongRolloutWorker( |
| num_envs=config.ppo.num_envs, |
| opp_manager=self.opp_manager, |
| seed=config.training.seed |
| ) |
| |
| os.makedirs(config.training.save_dir, exist_ok=True) |
| |
| |
| param_count = sum(p.numel() for p in self.agent.parameters() if p.requires_grad) |
| print(f"[*] Initialized Actor-Critic with {param_count:,} trainable parameters on {self.device}.") |
|
|
| def evaluate_against_all_opponents(self, num_episodes: int = 15) -> Dict[str, Dict[str, Any]]: |
| """ |
| Benchmark current agent against each distinct opponent baseline. |
| Returns detailed stats dictionary {opponent: {win, draw, loss, wins, draws, losses, total}} |
| """ |
| self.agent.eval() |
| opponents = { |
| "Easy Logic": EasyLogicOpponent(), |
| "Medium Logic": MediumLogicOpponent(), |
| "Realistic Hard": RealisticHardLogicOpponent(), |
| "Impossible Hard": ImpossibleHardLogicOpponent(), |
| "Minimax Depth 1": MinimaxOpponent(depth=1), |
| "Minimax Depth 2": MinimaxOpponent(depth=2), |
| "Random Agent": RandomOpponent() |
| } |
| |
| results = {} |
| eval_env = PongEnv() |
| |
| for opp_name, opp in opponents.items(): |
| wins = 0 |
| draws = 0 |
| losses = 0 |
| for ep in range(num_episodes): |
| obs = eval_env.reset(serve_direction=1 if ep % 2 == 0 else -1) |
| done = False |
| while not done: |
| with torch.no_grad(): |
| obs_t = torch.tensor(obs, dtype=torch.float32, device=self.device).unsqueeze(0) |
| logits, _ = self.agent(obs_t) |
| dist = Categorical(logits=logits / 0.30) |
| ego_act = dist.sample().item() |
| |
| opp_act = opp.act(eval_env) |
| obs, _, done, info = eval_env.step(ego_act, opp_act) |
| |
| if done: |
| w = info.get("winner") |
| if w == "ego": |
| wins += 1 |
| elif w == "draw": |
| draws += 1 |
| else: |
| losses += 1 |
| |
| results[opp_name] = { |
| "win": wins / num_episodes, |
| "draw": draws / num_episodes, |
| "loss": losses / num_episodes, |
| "wins": wins, |
| "draws": draws, |
| "losses": losses, |
| "total": num_episodes |
| } |
| |
| self.agent.train() |
| return results |
|
|
| def find_latest_checkpoint(self) -> Optional[str]: |
| """Search save_dir for the most recent valid checkpoint state or model file.""" |
| save_dir = self.cfg.training.save_dir |
| if not os.path.exists(save_dir): |
| return None |
|
|
| |
| latest_state = os.path.join(save_dir, "pong_train_state_latest.pt") |
| if os.path.exists(latest_state): |
| return latest_state |
|
|
| |
| state_candidates = [] |
| for fname in os.listdir(save_dir): |
| if fname.startswith("pong_train_state_ckpt_") and fname.endswith(".pt"): |
| try: |
| num = int(fname.replace("pong_train_state_ckpt_", "").replace(".pt", "")) |
| state_candidates.append((num, os.path.join(save_dir, fname))) |
| except ValueError: |
| pass |
| if state_candidates: |
| state_candidates.sort(key=lambda x: x[0], reverse=True) |
| return state_candidates[0][1] |
|
|
| |
| model_candidates = [] |
| for fname in os.listdir(save_dir): |
| if fname.startswith("pong_model_ckpt_") and fname.endswith(".pt"): |
| try: |
| num = int(fname.replace("pong_model_ckpt_", "").replace(".pt", "")) |
| model_candidates.append((num, os.path.join(save_dir, fname))) |
| except ValueError: |
| pass |
| if model_candidates: |
| model_candidates.sort(key=lambda x: x[0], reverse=True) |
| return model_candidates[0][1] |
|
|
| return None |
|
|
| def load_checkpoint(self, checkpoint_path: str) -> Tuple[int, int, int]: |
| """ |
| Load complete training state or model weights from checkpoint. |
| Returns: (start_update, global_step, checkpoint_count) |
| """ |
| if not os.path.exists(checkpoint_path): |
| raise FileNotFoundError(f"Checkpoint not found at: {checkpoint_path}") |
|
|
| print(f"[*] Loading checkpoint from: {checkpoint_path} ...") |
| checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False) |
|
|
| if isinstance(checkpoint, dict) and "agent_state_dict" in checkpoint: |
| self.agent.load_state_dict(checkpoint["agent_state_dict"]) |
| if "optimizer_state_dict" in checkpoint: |
| self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) |
| if "opp_manager_checkpoints" in checkpoint: |
| self.opp_manager.checkpoints = checkpoint["opp_manager_checkpoints"] |
| if "match_history" in checkpoint: |
| for k, v in checkpoint["match_history"].items(): |
| if k in self.vector_worker.match_history: |
| converted = ["W" if x in (1, "W") else ("D" if x in (0.5, "D") else "L") for x in v] |
| self.vector_worker.match_history[k] = deque(converted, maxlen=self.vector_worker.history_window) |
| if "cumulative_stats" in checkpoint: |
| self.vector_worker.cumulative_stats = checkpoint["cumulative_stats"] |
| if "torch_rng" in checkpoint: |
| torch.set_rng_state(checkpoint["torch_rng"]) |
| if "numpy_rng" in checkpoint: |
| np.random.set_state(checkpoint["numpy_rng"]) |
| if "python_rng" in checkpoint: |
| random.setstate(checkpoint["python_rng"]) |
|
|
| start_update = checkpoint.get("update", 0) |
| global_step = checkpoint.get("global_step", 0) |
| checkpoint_count = checkpoint.get("checkpoint_count", 0) |
| print(f"[OK] Successfully resumed full training state from Step: {global_step:,} (Update: {start_update}, Checkpoints in pool: {len(self.opp_manager.checkpoints)})") |
| return start_update, global_step, checkpoint_count |
| elif isinstance(checkpoint, dict): |
| self.agent.load_state_dict(checkpoint) |
| print(f"[OK] Loaded model weights from checkpoint into agent.") |
| return 0, 0, 0 |
| else: |
| raise ValueError(f"Invalid checkpoint format in: {checkpoint_path}") |
|
|
| def record_gameplay_video(self, global_step: int, checkpoint_num: Optional[int] = None) -> Optional[str]: |
| """ |
| Record a gameplay match video of the current policy against Hard Logic and Minimax opponents. |
| Saves MP4/GIF to the configured video directory. |
| """ |
| if not self.cfg.video.enabled: |
| return None |
|
|
| os.makedirs(self.cfg.video.video_dir, exist_ok=True) |
| self.agent.eval() |
|
|
| renderer = PongRenderer(self.cfg.physics, self.cfg.video) |
| test_opponents = [ |
| ("Realistic Hard Pro", RealisticHardLogicOpponent()), |
| ("Medium Logic", MediumLogicOpponent()), |
| ("Minimax Depth 2", MinimaxOpponent(depth=2)), |
| ("Impossible Hard Wall", ImpossibleHardLogicOpponent()), |
| ("Self-Play Mirror", NeuralOpponent(self.agent, self.device)) |
| ] |
|
|
| frames: List[np.ndarray] = [] |
| env = PongEnv(self.cfg.physics, self.cfg.reward) |
|
|
| ego_score = 0 |
| opp_score = 0 |
|
|
| max_video_steps = 180 |
|
|
| for opp_name, opp in test_opponents: |
| for ep in range(self.cfg.video.record_episodes): |
| obs = env.reset(serve_direction=1 if ep % 2 == 0 else -1) |
| done = False |
| step_i = 0 |
|
|
| while not done and step_i < max_video_steps: |
| step_i += 1 |
| with torch.no_grad(): |
| obs_t = torch.tensor(obs, dtype=torch.float32, device=self.device).unsqueeze(0) |
| logits, _ = self.agent(obs_t) |
| ego_act = torch.argmax(logits, dim=-1).item() |
|
|
| opp_act = opp.act(env) |
|
|
| |
| for sub in range(env.phys.frame_skip): |
| frame = renderer.render_frame( |
| env, step_i, ego_score, opp_score, opp_name, global_step, ego_act, opp_act |
| ) |
| frames.append(frame) |
| _, d, info = env._physics_substep(ego_act, opp_act) |
| if d: |
| done = True |
| break |
|
|
| obs = env.get_ego_observation() |
|
|
| if done: |
| if info.get("winner") == "ego": |
| ego_score += 1 |
| elif info.get("winner") == "opponent": |
| opp_score += 1 |
|
|
| self.agent.train() |
|
|
| if not frames: |
| return None |
|
|
| ckpt_suffix = f"_ckpt_{checkpoint_num}" if checkpoint_num is not None else "" |
| filename = f"pong_gameplay_step_{global_step}{ckpt_suffix}.{self.cfg.video.video_format}" |
| filepath = os.path.join(self.cfg.video.video_dir, filename) |
|
|
| try: |
| imageio.mimsave(filepath, frames, fps=self.cfg.video.fps) |
| print(f"[+] Saved Gameplay Video at Step {global_step:,} -> {filepath}") |
| return filepath |
| except Exception as e: |
| gif_path = filepath.rsplit(".", 1)[0] + ".gif" |
| try: |
| imageio.mimsave(gif_path, frames, fps=self.cfg.video.fps) |
| print(f"[+] Saved Gameplay GIF at Step {global_step:,} -> {gif_path}") |
| return gif_path |
| except Exception as e2: |
| print(f"[!] Warning: Video export failed: {e2}") |
| return None |
|
|
| def train(self, resume_path: Optional[str] = None): |
| """Main PPO Training Loop with vectorized rollouts, GAE updates, and full resumability.""" |
| cfg = self.cfg |
| ppo = cfg.ppo |
| train_cfg = cfg.training |
| |
| total_steps = train_cfg.total_timesteps |
| num_envs = ppo.num_envs |
| rollout_steps = ppo.rollout_steps |
| batch_size = num_envs * rollout_steps |
| num_updates = total_steps // batch_size |
|
|
| start_update = 0 |
| global_step = 0 |
| checkpoint_count = 0 |
| last_video_step = 0 |
|
|
| |
| target_resume = resume_path or train_cfg.resume_checkpoint_path |
| if target_resume is None and train_cfg.resume: |
| target_resume = self.find_latest_checkpoint() |
|
|
| if target_resume: |
| start_update, global_step, checkpoint_count = self.load_checkpoint(target_resume) |
| |
| last_video_step = global_step |
| |
| |
| obs_buf = torch.zeros((rollout_steps, num_envs, cfg.model.obs_dim), dtype=torch.float32, device=self.device) |
| actions_buf = torch.zeros((rollout_steps, num_envs), dtype=torch.long, device=self.device) |
| logprobs_buf = torch.zeros((rollout_steps, num_envs), dtype=torch.float32, device=self.device) |
| rewards_buf = torch.zeros((rollout_steps, num_envs), dtype=torch.float32, device=self.device) |
| dones_buf = torch.zeros((rollout_steps, num_envs), dtype=torch.float32, device=self.device) |
| values_buf = torch.zeros((rollout_steps, num_envs), dtype=torch.float32, device=self.device) |
| |
| start_time = time.time() |
| |
| print("\n" + "="*80) |
| print(" STARTING SOTA PING PONG TRAINING LOOP") |
| print("="*80) |
| print(f"Total Target Timesteps: {total_steps:,} (Starting at: {global_step:,})") |
| print(f"Parallel CPU Envs : {num_envs}") |
| print(f"Rollout Length : {rollout_steps} steps (Batch: {batch_size} steps/update)") |
| print(f"Checkpoint Interval : Every {train_cfg.checkpoint_interval_steps:,} steps") |
| print(f"Evaluation Interval : Every {train_cfg.eval_interval_steps:,} steps") |
| print("="*80 + "\n") |
| |
| for update in range(start_update + 1, num_updates + 1): |
| |
| if ppo.lr_annealing: |
| frac = 1.0 - (update - 1.0) / num_updates |
| lr_now = frac * ppo.learning_rate |
| self.optimizer.param_groups[0]["lr"] = lr_now |
|
|
| |
| for step in range(rollout_steps): |
| global_step += num_envs |
| obs_tensor = torch.tensor(self.vector_worker.obs, dtype=torch.float32, device=self.device) |
| |
| with torch.no_grad(): |
| action, logprob, _, value = self.agent.get_action_and_value(obs_tensor) |
| |
| obs_buf[step] = obs_tensor |
| actions_buf[step] = action |
| logprobs_buf[step] = logprob |
| values_buf[step] = value |
| |
| |
| next_obs, rewards, dones, infos = self.vector_worker.step(action.cpu().numpy()) |
| rewards_buf[step] = torch.tensor(rewards, dtype=torch.float32, device=self.device) |
| dones_buf[step] = torch.tensor(dones, dtype=torch.float32, device=self.device) |
|
|
| |
| with torch.no_grad(): |
| next_obs_tensor = torch.tensor(self.vector_worker.obs, dtype=torch.float32, device=self.device) |
| next_value = self.agent.get_value(next_obs_tensor) |
| |
| advantages = torch.zeros_like(rewards_buf, device=self.device) |
| last_gae_lam = 0 |
| for t in reversed(range(rollout_steps)): |
| if t == rollout_steps - 1: |
| next_non_terminal = 1.0 - dones_buf[t] |
| next_val = next_value |
| else: |
| next_non_terminal = 1.0 - dones_buf[t + 1] |
| next_val = values_buf[t + 1] |
| |
| delta = rewards_buf[t] + ppo.gamma * next_val * next_non_terminal - values_buf[t] |
| advantages[t] = last_gae_lam = delta + ppo.gamma * ppo.gae_lambda * next_non_terminal * last_gae_lam |
| |
| returns = advantages + values_buf |
|
|
| |
| b_obs = obs_buf.reshape(-1, cfg.model.obs_dim) |
| b_actions = actions_buf.reshape(-1) |
| b_logprobs = logprobs_buf.reshape(-1) |
| b_advantages = advantages.reshape(-1) |
| b_returns = returns.reshape(-1) |
| b_values = values_buf.reshape(-1) |
|
|
| |
| b_advantages = (b_advantages - b_advantages.mean()) / (b_advantages.std() + 1e-8) |
|
|
| |
| b_inds = np.arange(batch_size) |
| clip_fracs = [] |
| |
| for epoch in range(ppo.num_epochs): |
| np.random.shuffle(b_inds) |
| for start in range(0, batch_size, ppo.mini_batch_size): |
| end = start + ppo.mini_batch_size |
| mb_inds = b_inds[start:end] |
| |
| _, newlogprob, entropy, newvalue = self.agent.get_action_and_value( |
| b_obs[mb_inds], b_actions[mb_inds] |
| ) |
| logratio = newlogprob - b_logprobs[mb_inds] |
| ratio = logratio.exp() |
|
|
| with torch.no_grad(): |
| clip_fracs.append(((ratio - 1.0).abs() > ppo.clip_epsilon).float().mean().item()) |
|
|
| mb_advantages = b_advantages[mb_inds] |
|
|
| |
| pg_loss1 = -mb_advantages * ratio |
| pg_loss2 = -mb_advantages * torch.clamp(ratio, 1.0 - ppo.clip_epsilon, 1.0 + ppo.clip_epsilon) |
| pg_loss = torch.max(pg_loss1, pg_loss2).mean() |
|
|
| |
| if ppo.clip_value_loss: |
| v_loss_unclipped = (newvalue - b_returns[mb_inds]) ** 2 |
| v_clipped = b_values[mb_inds] + torch.clamp( |
| newvalue - b_values[mb_inds], |
| -ppo.clip_epsilon, |
| ppo.clip_epsilon, |
| ) |
| v_loss_clipped = (v_clipped - b_returns[mb_inds]) ** 2 |
| v_loss_max = torch.max(v_loss_unclipped, v_loss_clipped) |
| v_loss = 0.5 * v_loss_max.mean() |
| else: |
| v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean() |
|
|
| |
| entropy_loss = entropy.mean() |
| |
| |
| loss = pg_loss - ppo.entropy_coef * entropy_loss + ppo.value_coef * v_loss |
|
|
| self.optimizer.zero_grad() |
| loss.backward() |
| nn.utils.clip_grad_norm_(self.agent.parameters(), ppo.max_grad_norm) |
| self.optimizer.step() |
|
|
| |
| if global_step >= (checkpoint_count + 1) * train_cfg.checkpoint_interval_steps: |
| checkpoint_count += 1 |
| self.opp_manager.save_checkpoint(self.agent) |
|
|
| |
| ckpt_path = os.path.join(train_cfg.save_dir, f"pong_model_ckpt_{checkpoint_count}.pt") |
| torch.save(self.agent.state_dict(), ckpt_path) |
|
|
| |
| full_state = { |
| "global_step": global_step, |
| "update": update, |
| "checkpoint_count": checkpoint_count, |
| "agent_state_dict": self.agent.state_dict(), |
| "optimizer_state_dict": self.optimizer.state_dict(), |
| "opp_manager_checkpoints": self.opp_manager.checkpoints, |
| "match_history": {k: list(v) for k, v in self.vector_worker.match_history.items()}, |
| "cumulative_stats": self.vector_worker.cumulative_stats, |
| "torch_rng": torch.get_rng_state(), |
| "numpy_rng": np.random.get_state(), |
| "python_rng": random.getstate(), |
| } |
| state_ckpt_path = os.path.join(train_cfg.save_dir, f"pong_train_state_ckpt_{checkpoint_count}.pt") |
| state_latest_path = os.path.join(train_cfg.save_dir, "pong_train_state_latest.pt") |
| torch.save(full_state, state_ckpt_path) |
| torch.save(full_state, state_latest_path) |
|
|
| print(f"[+] Saved Resumable Checkpoint #{checkpoint_count} at Step {global_step:,} -> {ckpt_path}") |
|
|
| |
| if cfg.video.enabled and cfg.video.save_video_every_checkpoint: |
| last_video_step = global_step |
| self.record_gameplay_video(global_step=global_step, checkpoint_num=checkpoint_count) |
|
|
| |
| if cfg.video.enabled and not (cfg.video.save_video_every_checkpoint and global_step >= (checkpoint_count) * train_cfg.checkpoint_interval_steps): |
| if (global_step - last_video_step) >= cfg.video.video_interval_steps: |
| last_video_step = global_step |
| self.record_gameplay_video(global_step=global_step, checkpoint_num=checkpoint_count) |
|
|
| |
| if update % train_cfg.log_interval_updates == 0 or update == num_updates: |
| fps = int(global_step / max(1e-5, (time.time() - start_time))) |
| mean_reward = rewards_buf.mean().item() |
| current_lr = self.optimizer.param_groups[0]["lr"] |
|
|
| print(f"\n[Step {global_step:08d} | Upd {update:04d}/{num_updates:04d} | FPS: {fps:4d} | LR: {current_lr:.2e} | " |
| f"Rew: {mean_reward:+.4f} | Ent: {entropy_loss.item():.4f} | PLoss: {pg_loss.item():+.4f} | VLoss: {v_loss.item():.4f}]") |
|
|
| |
| live_stats = self.vector_worker.get_live_match_stats() |
| col_items = [] |
| for cat_name, st in live_stats.items(): |
| if st["total"] > 0: |
| col_items.append( |
| f"{cat_name:17s}: {st['win_rate'] * 100:5.1f}% " |
| f"({st['wins']:2d}W/{st['draws']:2d}D/{st['losses']:2d}L | {st['total']:2d}p)" |
| ) |
| else: |
| col_items.append(f"{cat_name:17s}: N/A ( 0W/ 0D/ 0L | 0p)") |
|
|
| print(" >> Live Rolling Match Outcomes (Recent Rollout Matches):") |
| for j in range(0, len(col_items), 2): |
| chunk = " | ".join(col_items[j:j+2]) |
| print(f" * {chunk}") |
|
|
| if global_step % train_cfg.eval_interval_steps < batch_size or update == num_updates: |
| print("\n" + "="*58) |
| print(" --- MULTI-OPPONENT EVALUATION BENCHMARK ---") |
| print("="*58) |
| results = self.evaluate_against_all_opponents(num_episodes=train_cfg.eval_episodes) |
| for name, st in results.items(): |
| print(f" * vs {name:16s}: {st['win'] * 100:5.1f}% Win | {st['draw'] * 100:5.1f}% Draw | {st['loss'] * 100:5.1f}% Loss ({st['wins']}W / {st['draws']}D / {st['losses']}L)") |
| print("="*58 + "\n") |
|
|
| |
| final_model_path = os.path.join(train_cfg.save_dir, "pong_champion_final.pt") |
| torch.save(self.agent.state_dict(), final_model_path) |
| print(f"\n[OK] Training Complete! Final Champion Model saved to: {final_model_path}\n") |
|
|
| |
| if cfg.video.enabled: |
| self.record_gameplay_video(global_step=global_step, checkpoint_num=None) |
|
|
|
|
| |
| |
| |
|
|
| def run_self_tests(): |
| """ |
| Execute comprehensive automated test suite verifying physics, |
| opponents, model constraints, and PPO gradient flow. |
| """ |
| print("\n" + "="*80) |
| print(" RUNNING PING PONG AI SELF-TEST SUITE") |
| print("="*80) |
|
|
| |
| print("[1/7] Testing Model Architecture & Parameter Count Ceiling...") |
| model = ActorCritic(CONFIG.model) |
| total_params = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| print(f" Total parameters: {total_params:,} (Limit: {CONFIG.model.max_allowed_params:,})") |
| assert total_params <= CONFIG.model.max_allowed_params, "Param constraint exceeded!" |
| print(" -> PASSED: Parameter count constraint verified.") |
|
|
| |
| print("[2/7] Testing Continuous Collision Detection (CCD) & Observation Normalization...") |
| env = PongEnv() |
| obs = env.reset() |
| assert obs.shape == (CONFIG.model.obs_dim,), f"Expected obs shape ({CONFIG.model.obs_dim},), got {obs.shape}" |
| assert np.all(obs >= -1.5) and np.all(obs <= 1.5), "Observation normalization out of bounds!" |
|
|
| |
| env.ball_x = 25.0 |
| env.ball_y = 250.0 |
| env.ball_vx = -16.0 |
| env.ball_vy = 0.0 |
| env.ego_y = 250.0 |
| env.ego_vy = 0.0 |
| sub_r, d, info = env._physics_substep(ego_action=0, opp_action=0) |
| assert info["hit_ego"] is True, "High-speed ball failed to trigger CCD hit!" |
| assert env.ball_vx > 0, "Ball failed to bounce forward on CCD hit!" |
| assert env.ball_x >= env.phys.paddle_width + env.phys.ball_radius, "Ball tunneled behind paddle!" |
| |
| |
| next_obs, r, d, info = env.step(ego_action=1, opp_action=2) |
| assert next_obs.shape == (CONFIG.model.obs_dim,), "Step output shape mismatch!" |
| print(" -> PASSED: Continuous collision detection and physics dynamics verified.") |
|
|
| |
| print("[3/7] Testing all Opponent Strategies (including Realistic & Impossible Hard)...") |
| opponents = [ |
| ("Random", RandomOpponent()), |
| ("Easy Logic", EasyLogicOpponent()), |
| ("Medium Logic", MediumLogicOpponent()), |
| ("Realistic Hard", RealisticHardLogicOpponent()), |
| ("Impossible Hard", ImpossibleHardLogicOpponent()), |
| ("Minimax D1", MinimaxOpponent(depth=1)), |
| ("Minimax D2", MinimaxOpponent(depth=2)), |
| ("Neural Opponent", NeuralOpponent(model)), |
| ] |
| for name, opp in opponents: |
| act = opp.act(env) |
| assert act in [0, 1, 2], f"Opponent {name} produced invalid action: {act}" |
| print(f" - {name:20s}: Valid action generated ({act})") |
| print(" -> PASSED: All opponent strategies working as expected.") |
|
|
| |
| print("[4/7] Testing Historical Checkpoint Buffer & Lag Range [5 - 75]...") |
| opp_mgr = OpponentManager(CONFIG.opponents, model) |
| |
| sampled_opp, name = opp_mgr.sample_opponent() |
| assert sampled_opp is not None |
| |
| |
| for _ in range(80): |
| opp_mgr.save_checkpoint(model) |
| assert len(opp_mgr.checkpoints) == 80 |
| |
| |
| sampled_opp, name = opp_mgr.sample_opponent() |
| assert sampled_opp is not None |
| print(f" - Checkpoint buffer capacity: {len(opp_mgr.checkpoints)}, sampled: {name}") |
| print(" -> PASSED: Checkpoint sampling and fallback behavior verified.") |
|
|
| |
| print("[5/7] Testing PPO Forward Pass, Loss Computation & Gradient Step...") |
| dummy_obs = torch.randn((16, CONFIG.model.obs_dim)) |
| dummy_actions = torch.randint(0, 3, (16,)) |
| action, logprob, entropy, value = model.get_action_and_value(dummy_obs, dummy_actions) |
| |
| loss = -logprob.mean() + value.mean() |
| loss.backward() |
| optimizer = optim.Adam(model.parameters(), lr=1e-3) |
| optimizer.step() |
| print(" -> PASSED: Neural network gradient backward step verified.") |
|
|
| |
| print("[6/7] Testing 2D Canvas Frame Rendering & Video File Export...") |
| renderer = PongRenderer(CONFIG.physics, CONFIG.video) |
| sample_frame = renderer.render_frame( |
| env, step_idx=1, ego_score=0, opp_score=0, opp_name="Hard Logic", |
| global_step=1000, ego_act=1, opp_act=2 |
| ) |
| assert sample_frame.shape == (CONFIG.video.height, CONFIG.video.width, 3), "Frame dimension mismatch!" |
| os.makedirs("./videos_pong_test", exist_ok=True) |
| test_video_path = "./videos_pong_test/test_render_clip.mp4" |
| imageio.mimsave(test_video_path, [sample_frame] * 10, fps=CONFIG.video.fps) |
| assert os.path.exists(test_video_path), "Test video file was not created!" |
| os.remove(test_video_path) |
| os.rmdir("./videos_pong_test") |
| print(" -> PASSED: Video renderer and file export verified.") |
|
|
| |
| print("[7/7] Testing Full Training Checkpoint Save & Resume Fidelity...") |
| test_state_dir = "./checkpoints_pong_test" |
| os.makedirs(test_state_dir, exist_ok=True) |
| test_trainer = PPOTrainer(CONFIG) |
| test_trainer.cfg.training.save_dir = test_state_dir |
| test_ckpt_file = os.path.join(test_state_dir, "pong_train_state_latest.pt") |
| |
| |
| dummy_state = { |
| "global_step": 5000, |
| "update": 10, |
| "checkpoint_count": 2, |
| "agent_state_dict": test_trainer.agent.state_dict(), |
| "optimizer_state_dict": test_trainer.optimizer.state_dict(), |
| "opp_manager_checkpoints": test_trainer.opp_manager.checkpoints, |
| "match_history": {}, |
| } |
| torch.save(dummy_state, test_ckpt_file) |
| |
| |
| resumed_trainer = PPOTrainer(CONFIG) |
| upd, stp, cnt = resumed_trainer.load_checkpoint(test_ckpt_file) |
| assert upd == 10 and stp == 5000 and cnt == 2, "Resumed state metadata mismatch!" |
| os.remove(test_ckpt_file) |
| os.rmdir(test_state_dir) |
| print(" -> PASSED: Checkpoint state saving and resumability verified.") |
|
|
| print("\n" + "="*80) |
| print(" [OK] ALL 7 SELF-TESTS PASSED SUCCESSFULLY!") |
| print("="*80 + "\n") |
|
|
|
|
| |
| |
| |
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="SOTA Ping Pong Reinforcement Learning Training System") |
| parser.add_argument("--test", action="store_true", help="Run automated verification self-tests") |
| parser.add_argument("--resume", action="store_true", help="Auto-resume training from latest checkpoint") |
| parser.add_argument("--load-checkpoint", type=str, default=None, help="Resume training from specific checkpoint file") |
| parser.add_argument("--timesteps", type=int, default=None, help="Override total training timesteps") |
| parser.add_argument("--envs", type=int, default=None, help="Override number of parallel environments") |
| parser.add_argument("--eval-episodes", type=int, default=None, help="Override evaluation episodes") |
| parser.add_argument("--no-video", action="store_true", help="Disable gameplay video saving") |
| parser.add_argument("--video-interval", type=int, default=None, help="Override video save interval (steps)") |
| parser.add_argument("--checkpoint-interval", type=int, default=None, help="Override checkpoint save interval (steps)") |
| args = parser.parse_args() |
|
|
| if args.test: |
| run_self_tests() |
| return |
|
|
| |
| if args.timesteps is not None: |
| CONFIG.training.total_timesteps = args.timesteps |
| if args.envs is not None: |
| CONFIG.ppo.num_envs = args.envs |
| if args.eval_episodes is not None: |
| CONFIG.training.eval_episodes = args.eval_episodes |
| if args.no_video: |
| CONFIG.video.enabled = False |
| if args.video_interval is not None: |
| CONFIG.video.video_interval_steps = args.video_interval |
| if args.checkpoint_interval is not None: |
| CONFIG.training.checkpoint_interval_steps = args.checkpoint_interval |
| if args.resume: |
| CONFIG.training.resume = True |
| if args.load_checkpoint is not None: |
| CONFIG.training.resume_checkpoint_path = args.load_checkpoint |
|
|
| trainer = PPOTrainer(CONFIG) |
| trainer.train() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
