Spaces:
Sleeping
Sleeping
File size: 4,533 Bytes
888cb3b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | """
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
# Allow running from any directory
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']}"
)
# Save trained model
agent.save(save_path)
# Save training log
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,
)
|