210 lines
7.8 KiB
Python
210 lines
7.8 KiB
Python
"""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()
|
|
|