"""No-control baseline runner for synchronized reward baselines.""" from __future__ import annotations import copy import os import matplotlib import numpy as np import yaml from tqdm import tqdm matplotlib.use("Agg") 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_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.reward_baseline import EpisodeBaselineWriter, resolve_baseline_dir from utils.run_dirs import resolve_run_dirs, write_shared_run_config from utils.seeding import derive_seed, resolve_base_seed, set_global_seed def _select_no_control_action(env: SUMOEdgeVSLEnvironment) -> np.ndarray: if env.num_controlled_edges <= 0: return np.zeros(0, dtype=np.int64) return np.full(env.num_controlled_edges, env.action_dim - 1, dtype=np.int64) def _baseline_row(episode: int, seed: int | None, reward: float, info: dict) -> dict: return { "episode": int(episode), "step": int(info.get("step", 0)), "seed": "" if seed is None else int(seed), "sim_time": float(info.get("sim_time", np.nan)), "reward": float(reward), "mean_speed_kmh": float(info.get("mean_speed_kmh", np.nan)), "num_vehicles": int(info.get("num_vehicles", 0)), "mainline_completed_count": int(info.get("mainline_completed_count", 0)), "mainline_interval_travel_time_mean_s": float( info.get("mainline_interval_travel_time_mean_s", np.nan) ), "mainline_travel_time_cumulative_mean_s": float( info.get("mainline_travel_time_cumulative_mean_s", np.nan) ), "ttc_risk_rate": float(info.get("ttc_risk_rate", np.nan)), } def train_sumo_no_control(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) train_config = get_training_config(config) base_seed = resolve_base_seed(train_config) set_global_seed(base_seed) resolved_run_timestamp, checkpoint_dir, log_dir = resolve_run_dirs( "no_control", 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 runtime_config["runtime"]["run_timestamp"] = resolved_run_timestamp reward_cfg = runtime_config.setdefault("environment", {}).setdefault("reward", {}) reward_cfg["mode"] = "absolute" reward_cfg["baseline_dir"] = resolve_baseline_dir(runtime_config, resolved_run_timestamp) baseline_dir = reward_cfg["baseline_dir"] write_shared_run_config( runtime_config, log_dir=log_dir, checkpoint_dir=checkpoint_dir, run_timestamp=run_timestamp, ) logger = TrainingLogger(log_dir, "no_control") env = SUMOEdgeVSLEnvironment(runtime_config) num_episodes = train_config["num_episodes"] log_freq = train_config.get("log_freq", 10) snapshot_interval = int(train_config.get("artifact_snapshot_interval", 50)) episode_rewards = [] episode_throughputs = [] episode_mean_speeds = [] episode_speed_variance_norms = [] episode_ttc_risks = [] print("=" * 70) print("NO_CONTROL baseline runner - synchronized reward baseline") print("=" * 70) print(f" Episode steps: {env.episode_length}") print(f" Baseline dir: {baseline_dir}") print(f" Global seed: {base_seed if base_seed is not None else 'None (random)'}") print() try: for episode in range(1, num_episodes + 1): seed = derive_seed(base_seed, episode) env.reset(seed=seed) baseline_writer = EpisodeBaselineWriter(baseline_dir=baseline_dir, episode=episode) episode_reward = 0.0 episode_throughput = 0.0 episode_speed = 0.0 episode_speed_variance_norm = 0.0 episode_ttc_risk = 0.0 episode_reward_components = init_reward_component_totals() done = False step = 0 pbar = tqdm(total=env.episode_length, desc=f"NO_CONTROL Ep {episode}/{num_episodes}", leave=False) while not done: action = _select_no_control_action(env) _, reward, done, info = env.step(action, apply_control=True) baseline_writer.append(_baseline_row(episode, seed, reward, info)) episode_reward += reward episode_throughput += info["throughput"] episode_speed += info["mean_speed_kmh"] episode_speed_variance_norm += info["speed_variance_norm"] episode_ttc_risk += float(info.get("ttc_risk_rate", 0.0)) for column in REWARD_COMPONENT_COLUMNS: episode_reward_components[column] += float(info.get(column, 0.0)) step += 1 pbar.set_postfix(r=f"{episode_reward:.1f}", 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": int(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), "baseline_dir": baseline_dir, } 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, snapshot_interval=snapshot_interval, ) if episode % log_freq == 0: recent_rewards = episode_rewards[-log_freq:] print(f"\nNO_CONTROL episode {episode}/{num_episodes}") print(f" Reward: {episode_reward:.2f} (Avg: {np.mean(recent_rewards):.2f})") print(f" Mean Speed: {avg_speed:.1f} km/h") finally: env.close() 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"), ) return { "model": "no_control", "log_dir": log_dir, "checkpoint_dir": checkpoint_dir, "baseline_dir": baseline_dir, } if __name__ == "__main__": train_sumo_no_control()