"""Run independent literature-informed rule-based VSL baselines.""" from __future__ import annotations import copy import os from typing import Type import matplotlib import numpy as np import yaml from tqdm import tqdm from agents.rule_vsl_agent import ( BaseRuleVSLAgent, BottleneckRuleVSLAgent, HarmonizationRuleVSLAgent, OccupancyRuleVSLAgent, ) from envs.edge_vsl_env import SUMOEdgeVSLEnvironment from envs.reward_system import REWARD_COMPONENT_COLUMNS, average_reward_components, init_reward_component_totals from utils.config import get_agent_config, get_training_config from utils.episode_artifacts import save_training_episode_artifacts from utils.logger import TrainingLogger from utils.plot import plot_training_curves from utils.run_dirs import resolve_run_dirs, write_shared_run_config from utils.seeding import derive_seed, resolve_base_seed, set_global_seed matplotlib.use("Agg") def train_sumo_rule_vsl_baseline( model_key: str, model_label: str, agent_cls: Type[BaseRuleVSLAgent], *, log_dir=None, checkpoint_dir=None, run_timestamp=None, ): with open("config_sumo_vsl.yaml", "r", encoding="utf-8") as f: config = yaml.safe_load(f) agent_config = get_agent_config(config, model_key) train_config = get_training_config(config) base_seed = resolve_base_seed(train_config) set_global_seed(base_seed) _, checkpoint_dir, log_dir = resolve_run_dirs( model_key, log_dir=log_dir, checkpoint_dir=checkpoint_dir, run_timestamp=run_timestamp, ) os.makedirs(checkpoint_dir, exist_ok=True) os.makedirs(log_dir, exist_ok=True) runtime_config = copy.deepcopy(config) runtime_config.setdefault("runtime", {})["output_dir"] = log_dir runtime_config["runtime"]["evaluation_mode"] = False write_shared_run_config( runtime_config, log_dir=log_dir, checkpoint_dir=checkpoint_dir, run_timestamp=run_timestamp, ) logger = TrainingLogger(log_dir, model_key) env = SUMOEdgeVSLEnvironment(runtime_config) agent = agent_cls.from_env(env, agent_config) num_episodes = train_config["num_episodes"] log_freq = train_config.get("log_freq", 10) episode_rewards = [] episode_throughputs = [] episode_mean_speeds = [] episode_speed_variance_norms = [] episode_ttc_risks = [] best_reward = -float("inf") print("=" * 70) print(f"{model_label} baseline - SUMO VSL environment") print("=" * 70) print(f" Controlled edges: {env.num_controlled_edges}") print(f" Actions per edge: {env.action_dim}") print(f" Episode steps: {env.episode_length}") print(f" Control interval: {env.control_interval}s") print(f" Global seed: {base_seed if base_seed is not None else 'None (random)'}") print(" Policy type: deterministic non-learning rule") print() try: for episode in range(1, num_episodes + 1): seed = derive_seed(base_seed, episode) state = env.reset(seed=seed) agent.reset_episode() episode_reward = 0.0 episode_throughput = 0.0 episode_speed = 0.0 episode_speed_variance_norm = 0.0 episode_reward_components = init_reward_component_totals() episode_ttc_risk = 0.0 done = False step = 0 pbar = tqdm(total=env.episode_length, desc=f"Ep {episode}/{num_episodes}", leave=False) while not done: action, _, _ = agent.select_action(state, deterministic=True) next_state, reward, done, info = env.step(action) episode_reward += reward episode_throughput += info["throughput"] episode_speed += info["mean_speed_kmh"] episode_speed_variance_norm += info["speed_variance_norm"] for column in REWARD_COMPONENT_COLUMNS: episode_reward_components[column] += float(info.get(column, 0.0)) episode_ttc_risk += float(info.get("ttc_risk_rate", 0.0)) state = next_state step += 1 pbar.set_postfix( r=f"{episode_reward:.1f}", tp=f"{info['throughput']:.0f}", v=f"{info['mean_speed_kmh']:.1f}", ) pbar.update(1) pbar.close() avg_tp = episode_throughput / max(step, 1) avg_speed = episode_speed / max(step, 1) avg_speed_variance_norm = episode_speed_variance_norm / max(step, 1) avg_reward_components = average_reward_components(episode_reward_components, step) episode_rewards.append(episode_reward) episode_throughputs.append(avg_tp) episode_mean_speeds.append(avg_speed) episode_speed_variance_norms.append(avg_speed_variance_norm) episode_ttc_risks.append(episode_ttc_risk) logger.log( episode, episode_reward, avg_tp, avg_speed, speed_variance_norm=avg_speed_variance_norm, reward_components=avg_reward_components, ttc_risk=episode_ttc_risk, ) episode_summary = { "episode": episode, "reward": float(episode_reward), "avg_throughput": float(avg_tp), "avg_mean_speed_kmh": float(avg_speed), "avg_speed_variance_norm": float(avg_speed_variance_norm), "ttc_risk": float(episode_ttc_risk), } for column, value in avg_reward_components.items(): episode_summary[f"avg_{column}"] = float(value) save_training_episode_artifacts( log_dir=log_dir, episode=episode, episode_metrics=env.episode_metrics, control_edges=env.control_edges, summary=episode_summary, ) if episode_reward > best_reward: best_reward = episode_reward if episode % log_freq == 0: recent_rewards = episode_rewards[-log_freq:] print(f"\nEpisode {episode}/{num_episodes}") print(f" Reward: {episode_reward:.2f} (Avg: {np.mean(recent_rewards):.2f})") print(f" Throughput: {avg_tp:.1f} veh/h") print(f" Mean Speed: {avg_speed:.1f} km/h") print(f" Normalized Speed Variance: {avg_speed_variance_norm:.4f}") print( " Reward Components: " + ", ".join( f"{column}={avg_reward_components[column]:.3f}" for column in REWARD_COMPONENT_COLUMNS ) ) finally: env.close() marker_path = os.path.join(checkpoint_dir, f"{model_key}_baseline.txt") with open(marker_path, "w", encoding="utf-8") as f: f.write(f"{model_label} is a deterministic rule-based baseline and has no trainable checkpoint.\n") plot_training_curves( episode_rewards, episode_throughputs, episode_mean_speeds, episode_speed_variance_norms, episode_ttc_risks, save_path=os.path.join(log_dir, "training_curves.png"), ) print("=" * 70) print(f"{model_label} run complete") print(f" Best reward: {best_reward:.2f}") print(f" Marker file: {marker_path}") print(f" Log dir: {log_dir}") print("=" * 70) def train_sumo_occ_rule_vsl(log_dir=None, checkpoint_dir=None, run_timestamp=None): return train_sumo_rule_vsl_baseline( "occ_rule_vsl", "Occ-Rule-VSL", OccupancyRuleVSLAgent, log_dir=log_dir, checkpoint_dir=checkpoint_dir, run_timestamp=run_timestamp, ) def train_sumo_bottleneck_rule_vsl(log_dir=None, checkpoint_dir=None, run_timestamp=None): return train_sumo_rule_vsl_baseline( "bottleneck_rule_vsl", "Bottleneck-Rule-VSL", BottleneckRuleVSLAgent, log_dir=log_dir, checkpoint_dir=checkpoint_dir, run_timestamp=run_timestamp, ) def train_sumo_harmonization_rule_vsl(log_dir=None, checkpoint_dir=None, run_timestamp=None): return train_sumo_rule_vsl_baseline( "harmonization_rule_vsl", "Harmonization-Rule-VSL", HarmonizationRuleVSLAgent, log_dir=log_dir, checkpoint_dir=checkpoint_dir, run_timestamp=run_timestamp, )