agent.py

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
from collections import deque

# Import the model we discussed
class DuelingDQN(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(DuelingDQN, self).__init__()
        self.feature = nn.Sequential(
            nn.Linear(state_dim, 256),
            nn.ReLU()
        )
        self.advantage = nn.Sequential(
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, action_dim)
        )
        self.value = nn.Sequential(
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 1)
        )

    def forward(self, x):
        x = self.feature(x)
        advantage = self.advantage(x)
        value = self.value(x)
        return value + advantage - advantage.mean(dim=-1, keepdim=True)

class DQNAgent:
    def __init__(self, state_dim, action_dim, lr=1e-4, gamma=0.99, epsilon_start=0.5, epsilon_end=1e-5, epsilon_decay=0.9999):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.gamma = gamma
        self.epsilon = epsilon_start
        self.epsilon_end = epsilon_end
        self.epsilon_decay = epsilon_decay
        
        # Device configuration (MPS for Mac, CUDA for Windows/Linux, else CPU)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Networks: Policy and Target
        self.policy_net = DuelingDQN(state_dim, action_dim).to(self.device)
        self.target_net = DuelingDQN(state_dim, action_dim).to(self.device)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        
        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=lr)
        self.memory = deque(maxlen=50000)
        self.batch_size = 256

    def select_action(self, state, mask=None):
        """Selects action using epsilon-greedy policy with optional action masking."""
        if random.random() < self.epsilon:
            if mask is not None:
                # Only sample from valid actions (indices where mask is 1)
                valid_actions = np.where(mask == 1)[0]
                return random.choice(valid_actions)
            return random.randrange(self.action_dim)
        
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        with torch.no_grad():
            q_values = self.policy_net(state_tensor)
            
        if mask is not None:
            # Apply mask: set invalid actions to a very low value
            mask_tensor = torch.FloatTensor(mask).to(self.device)
            q_values = q_values + (1.0 - mask_tensor) * -1e9
            
        return q_values.argmax().item()

    def store_transition(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))

    def update(self):
        if len(self.memory) < self.batch_size:
            return None

        # Sample a batch
        batch = random.sample(self.memory, self.batch_size)
        state, action, reward, next_state, done = zip(*batch)

        state = torch.FloatTensor(np.array(state)).to(self.device)
        action = torch.LongTensor(action).unsqueeze(1).to(self.device)
        reward = torch.FloatTensor(reward).unsqueeze(1).to(self.device)
        next_state = torch.FloatTensor(np.array(next_state)).to(self.device)
        done = torch.FloatTensor(done).unsqueeze(1).to(self.device)

        # Current Q values
        current_q = self.policy_net(state).gather(1, action)

        # Target Q values (Double DQN logic: use policy_net to choose, target_net to evaluate)
        with torch.no_grad():
            next_actions = self.policy_net(next_state).argmax(1).unsqueeze(1)
            next_q = self.target_net(next_state).gather(1, next_actions)
            target_q = reward + (1 - done) * self.gamma * next_q

        loss = nn.MSELoss()(current_q, target_q)
        
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # Decay epsilon
        self.epsilon = max(self.epsilon_end, self.epsilon * self.epsilon_decay)
        
        return loss.item()

    def update_target_network(self):
        self.target_net.load_state_dict(self.policy_net.state_dict())