| """ |
| train.py |
| ======== |
| Training script for the DQN agent on the Healthcare Routing environment. |
| |
| Run: |
| python train.py # default training |
| python train.py --episodes 500 # longer training |
| python train.py --render # enable console rendering |
| """ |
|
|
| import argparse |
| import sys |
| import os |
| import json |
| import numpy as np |
|
|
| |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "openenv_env")) |
| sys.path.insert(0, os.path.dirname(__file__)) |
|
|
| from healthcare_env import HealthcareRoutingEnv |
| from dqn_agent import DQNAgent |
|
|
|
|
| def train( |
| num_episodes: int = 300, |
| steps_per_ep: int = 100, |
| hidden_dim: int = 256, |
| lr: float = 1e-3, |
| gamma: float = 0.99, |
| render: bool = False, |
| save_path: str = "models/dqn_healthcare.pth", |
| log_interval: int = 20, |
| ): |
| os.makedirs(os.path.dirname(save_path) if os.path.dirname(save_path) else ".", exist_ok=True) |
|
|
| env = HealthcareRoutingEnv(render_mode="human" if render else None) |
| obs, _ = env.reset() |
|
|
| agent = DQNAgent( |
| obs_size = env.observation_space.shape[0], |
| action_size = env.action_space.n, |
| hidden_dim = hidden_dim, |
| lr = lr, |
| gamma = gamma, |
| ) |
|
|
| print(f"\n{'='*60}") |
| print(f" Healthcare Routing DQN Training") |
| print(f" Episodes: {num_episodes}") |
| print(f" Steps/ep: {steps_per_ep}") |
| print(f" Obs size: {env.observation_space.shape[0]}") |
| print(f" Actions: {env.action_space.n}") |
| print(f" Device: {agent.device}") |
| print(f"{'='*60}\n") |
|
|
| all_ep_rewards = [] |
|
|
| for episode in range(1, num_episodes + 1): |
| obs, _ = env.reset() |
| ep_reward = 0.0 |
| outcomes = {"success": 0, "no_bed": 0, "ambulance_busy": 0} |
|
|
| for step in range(steps_per_ep): |
| action = agent.select_action(obs) |
| next_obs, reward, terminated, truncated, info = env.step(action) |
|
|
| done = terminated or truncated |
| agent.store_transition(obs, action, reward, next_obs, done) |
| agent.update() |
|
|
| ep_reward += reward |
| obs = next_obs |
| outcomes[info.get("outcome", "success")] = outcomes.get(info.get("outcome", "success"), 0) + 1 |
|
|
| if render: |
| env.render() |
|
|
| if done: |
| break |
|
|
| all_ep_rewards.append(ep_reward) |
|
|
| if episode % log_interval == 0: |
| recent_mean = np.mean(all_ep_rewards[-log_interval:]) |
| stats = agent.get_stats() |
| print( |
| f"Episode {episode:4d}/{num_episodes} | " |
| f"Reward: {ep_reward:8.1f} | " |
| f"Mean({log_interval}): {recent_mean:8.1f} | " |
| f"ε={stats['epsilon']:.3f} | " |
| f"Loss={stats['mean_loss_100']:.4f} | " |
| f"✅{outcomes['success']} ❌{outcomes['no_bed']} 🚑{outcomes['ambulance_busy']}" |
| ) |
|
|
| |
| agent.save(save_path) |
|
|
| |
| log_path = save_path.replace(".pth", "_log.json") |
| with open(log_path, "w") as f: |
| json.dump({ |
| "episode_rewards": all_ep_rewards, |
| "final_stats": agent.get_stats(), |
| "config": { |
| "num_episodes": num_episodes, |
| "steps_per_ep": steps_per_ep, |
| "hidden_dim": hidden_dim, |
| "lr": lr, |
| "gamma": gamma, |
| } |
| }, f, indent=2) |
|
|
| print(f"\n[Train] Done! Final mean reward (last 50): {np.mean(all_ep_rewards[-50:]):.2f}") |
| print(f"[Train] Log saved → {log_path}") |
| return agent, all_ep_rewards |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="Train DQN on HealthcareRoutingEnv") |
| parser.add_argument("--episodes", type=int, default=300) |
| parser.add_argument("--steps", type=int, default=100) |
| parser.add_argument("--hidden", type=int, default=256) |
| parser.add_argument("--lr", type=float, default=1e-3) |
| parser.add_argument("--render", action="store_true") |
| parser.add_argument("--save", type=str, default="models/dqn_healthcare.pth") |
| args = parser.parse_args() |
|
|
| train( |
| num_episodes = args.episodes, |
| steps_per_ep = args.steps, |
| hidden_dim = args.hidden, |
| lr = args.lr, |
| render = args.render, |
| save_path = args.save, |
| ) |
|
|
