train.py

"""
train.py
Standalone advanced training script for Drone Delivery RL.
Usage: python train.py --task easy_delivery --episodes 1000
"""
import os
import argparse
import random
import torch
import torch.nn as nn
import torch.optim as optim
from collections import deque

import sys
from pathlib import Path

# Add project root to sys.path
ROOT_DIR = Path(__file__).parent
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

from drone_env.rl.model import PathQNet, ACTIONS, ACTION2IDX, CELL2IDX
from drone_env.server.grid_world_environment import DroneDeliveryEnvironment
from drone_env.models import DroneAction

# ── Hyperparameters ───────────────────────────────────────────────────────────
GAMMA = 0.99
BATCH_SIZE = 64
LR = 1e-4
REPLAY_SIZE = 10000
TARGET_UPDATE = 10
EPS_START = 1.0
EPS_END = 0.05
EPS_DECAY = 0.995

# ── Replay Buffer ─────────────────────────────────────────────────────────────
class ReplayBuffer:
    def __init__(self, capacity: int):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size: int):
        return random.sample(self.buffer, batch_size)

    def __len__(self):
        return len(self.buffer)

# ── Helper to process observations ────────────────────────────────────────────
def obs_to_tensors(obs, device):
    """Converts a DroneObservation to the tensors required by PathQNet."""
    # 1. Grid (H*W)
    # Using cell_types instead of emojis for better robustness
    grid = obs.cell_types
    flat_grid = []
    for row in grid:
        for cell in row:
            # Map emoji back to int (simplified)
            idx = CELL2IDX.get(cell, 0) # default road
            flat_grid.append(idx)
    grid_t = torch.tensor([flat_grid], dtype=torch.long, device=device)

    # 2. Drone XY (normalized)
    drone_xy = torch.tensor([[obs.drone_x / obs.grid_width, obs.drone_y / obs.grid_height]], dtype=torch.float, device=device)

    # 3. Battery (0-1)
    battery = torch.tensor([[obs.battery]], dtype=torch.float, device=device)

    # 4. Target XY (normalized)
    if obs.current_target:
        tx, ty = obs.current_target
        target_xy = torch.tensor([[tx / obs.grid_width, ty / obs.grid_height]], dtype=torch.float, device=device)
    else:
        target_xy = torch.tensor([[0.0, 0.0]], dtype=torch.float, device=device)

    return grid_t, drone_xy, battery, target_xy

# ── Training Loop ─────────────────────────────────────────────────────────────
def train(task_name: str, episodes: int, device_name: str):
    device = torch.device(device_name)
    print(f"Starting training for {task_name} on {device}...")

    os.makedirs("data", exist_ok=True)
    with open("data/train.log", "a") as f:
        f.write(f"\n>>> Neural Training Engine Started [TASK: {task_name}]\n")
        
    env = DroneDeliveryEnvironment()
    policy_net = PathQNet(embed_dim=64).to(device)
    target_net = PathQNet(embed_dim=64).to(device)
    target_net.load_state_dict(policy_net.state_dict())
    
    optimizer = optim.Adam(policy_net.parameters(), lr=LR)
    memory = ReplayBuffer(REPLAY_SIZE)
    epsilon = EPS_START

    # Automatically resume if model exists
    task_short = task_name.split('_')[0]
    model_path = f"data/{task_short}.pth"
    if os.path.exists(model_path):
        print(f"Loading existing weights from {model_path}...")
        try:
            policy_net.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
            target_net.load_state_dict(policy_net.state_dict())
            epsilon = EPS_END # Start with low exploration if resuming
            with open("data/train.log", "a") as f:
                f.write(f"Resuming from checkpoint: {model_path}\n")
        except Exception as e:
            print(f"Failed to load checkpoint: {e}. Starting fresh.")

    for ep in range(episodes):
        obs = env.reset(DroneAction(task_name=task_name))
        state_t = obs_to_tensors(obs, device)
        total_reward = 0
        done = False

        while not done:
            # Select Action
            if random.random() < epsilon:
                action_idx = random.randint(0, len(ACTIONS) - 1)
            else:
                with torch.no_grad():
                    q_values = policy_net(*state_t)
                    action_idx = q_values.argmax().item()

            action_str = ACTIONS[action_idx]
            
            # Step
            next_obs = env.step(DroneAction(direction=action_str))
            reward = next_obs.reward_last
            done = next_obs.done
            total_reward += reward

            next_state_t = obs_to_tensors(next_obs, device)
            
            # Save to buffer
            memory.push(state_t, action_idx, reward, next_state_t, done)
            state_t = next_state_t

            # Optimize
            if len(memory) > BATCH_SIZE:
                batch = memory.sample(BATCH_SIZE)
                
                # Unzip batch
                states, actions, rewards, next_states, dones = zip(*batch)
                
                # Batch processing
                def cat_tensors(t_list):
                    return [torch.cat([t[i] for t in t_list]) for i in range(4)]
                
                b_grid, b_drone, b_bat, b_target = cat_tensors(states)
                bn_grid, bn_drone, bn_bat, bn_target = cat_tensors(next_states)
                
                b_actions = torch.tensor(actions, device=device).unsqueeze(1)
                b_rewards = torch.tensor(rewards, device=device, dtype=torch.float).unsqueeze(1)
                b_dones   = torch.tensor(dones, device=device, dtype=torch.float).unsqueeze(1)
                
                # 1. Current Q values
                current_q = policy_net(b_grid, b_drone, b_bat, b_target).gather(1, b_actions)
                
                # 2. Next Q values from target net
                with torch.no_grad():
                    next_q = target_net(bn_grid, bn_drone, bn_bat, bn_target).max(1)[0].unsqueeze(1)
                    expected_q = b_rewards + (GAMMA * next_q * (1 - b_dones))
                
                # 3. Loss & Step
                loss = nn.MSELoss()(current_q, expected_q)
                optimizer.zero_grad()
                loss.backward()
                # Gradient clipping
                torch.nn.utils.clip_grad_norm_(policy_net.parameters(), 1.0)
                optimizer.step()
        
        epsilon = max(EPS_END, epsilon * EPS_DECAY)
        
        if ep % TARGET_UPDATE == 0:
            target_net.load_state_dict(policy_net.state_dict())
            
        if ep % 50 == 0:
            log_msg = f"Episode {ep}/{episodes} | Avg Reward: {total_reward:.2f} | Epsilon: {epsilon:.2f}"
            print(log_msg)
            torch.save(policy_net.state_dict(), model_path)
            with open("data/train.log", "a") as f:
                f.write(log_msg + " (Periodic Save)\n")

    # Save
    task_short = task_name.split('_')[0] 
    model_path = f"data/{task_short}.pth"
    torch.save(policy_net.state_dict(), model_path)
    
    msg = f"Training complete. Model saved to {model_path}."
    print(msg)
    with open("data/train.log", "a") as f:
        f.write(msg + "\n")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--task", type=str, default="easy_delivery")
    parser.add_argument("--episodes", type=int, default=100)
    parser.add_argument("--gpu", action="store_true")
    args = parser.parse_args()

    device = "cuda" if args.gpu and torch.cuda.is_available() else "cpu"
    train(args.task, args.episodes, device)