train/ppo_agent.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import numpy as np

class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, action_dim)
        )
        for layer in self.net:
            if isinstance(layer, nn.Linear):
                nn.init.xavier_uniform_(layer.weight)
                nn.init.constant_(layer.bias, 0)
    
    def forward(self, state):
        return self.net(state)

class ValueNetwork(nn.Module):
    def __init__(self, state_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 1)
        )
        for layer in self.net:
            if isinstance(layer, nn.Linear):
                nn.init.xavier_uniform_(layer.weight)
                nn.init.constant_(layer.bias, 0)
    
    def forward(self, state):
        return self.net(state)

class PPOAgent:
    def __init__(self, state_dim, action_dim, lr=3e-4, gamma=0.99, clip_eps=0.2):
        self.policy = PolicyNetwork(state_dim, action_dim)
        self.value = ValueNetwork(state_dim)
        self.optimizer = optim.Adam(
            list(self.policy.parameters()) + list(self.value.parameters()), 
            lr=lr
        )
        self.memory = []
        self.gamma = gamma
        self.clip_eps = clip_eps
        self.epsilon = 1.0
        self.epsilon_decay = 0.995
        self.epsilon_min = 0.1
        self.mini_batch_size = 256
        self.update_epochs = 4  

    def select_action(self, state):
        state = torch.FloatTensor(state).unsqueeze(0)
        
        # Add input validation
        if torch.isnan(state).any() or torch.isinf(state).any():
            print(f"Warning: Invalid state detected: {state}")
            action = np.random.randint(0, 8)
            log_prob = torch.log(torch.tensor(1/8))
            return action, log_prob
        
        with torch.no_grad():
            logits = self.policy(state)
            
            if torch.isnan(logits).any() or torch.isinf(logits).any():
                print(f"Warning: Invalid logits detected: {logits}")
                action = np.random.randint(0, 8)
                log_prob = torch.log(torch.tensor(1/8))
                return action, log_prob
            
            # Apply temperature scaling
            logits = logits / 1.0
            probs = torch.softmax(logits, dim=-1)
            
            # Add small epsilon to prevent zero probabilities
            probs = probs + 1e-8
            probs = probs / probs.sum(dim=-1, keepdim=True)
            
            dist = Categorical(probs)
            action = dist.sample()
            
        return action.item(), dist.log_prob(action)

    def store_experience(self, state, action, reward, next_state, done, log_prob):
        # Validate inputs before storing
        if not (torch.isnan(torch.FloatTensor(state)).any() or 
                torch.isnan(torch.FloatTensor(next_state)).any() or
                torch.isnan(log_prob).any() or
                np.isnan(reward)):
            self.memory.append((state, action, reward, next_state, done, log_prob))

    def calculate_returns_and_advantages_optimized(self, rewards, values, dones):
        """
        Optimized O(n) calculation of returns and advantages using GAE
        """
        returns = torch.zeros_like(rewards)
        advantages = torch.zeros_like(rewards)
        
        # GAE parameters
        gae_lambda = 0.95
        
        # Calculate advantages and returns in reverse order (O(n))
        last_gae = 0
        for i in reversed(range(len(rewards))):
            if i == len(rewards) - 1:
                next_value = 0  # Terminal state
            else:
                next_value = values[i + 1]
            
            # TD error
            delta = rewards[i] + self.gamma * next_value - values[i]
            
            # GAE advantage
            advantages[i] = last_gae = delta + self.gamma * gae_lambda * last_gae
            
            # Returns = advantages + values
            returns[i] = advantages[i] + values[i]
        
        return returns, advantages

    def update(self, epochs=None):
        if len(self.memory) < 32:  # Minimum batch size
            return None
        
        # Use instance variable if epochs not specified
        if epochs is None:
            epochs = self.update_epochs
            
        # Convert experience buffer to tensors
        states, actions, rewards, next_states, dones, old_log_probs = zip(*self.memory)
        
        states = torch.FloatTensor(np.array(states))
        actions = torch.LongTensor(actions)
        rewards = torch.FloatTensor(rewards)
        old_log_probs = torch.stack(old_log_probs)
        
        # Validate tensors
        if (torch.isnan(states).any() or torch.isnan(rewards).any() or 
            torch.isnan(old_log_probs).any()):
            print("Warning: NaN detected in experience buffer, skipping update")
            self.memory.clear()
            return None
        
        # Get values for all states at once
        with torch.no_grad():
            values = self.value(states).squeeze(-1)
        
        # OPTIMIZED: O(n) returns and advantages calculation
        returns, advantages = self.calculate_returns_and_advantages_optimized(
            rewards, values, torch.zeros_like(rewards)  # dones placeholder
        )
        
        # Normalize advantages
        if len(advantages) > 1 and advantages.std() > 1e-8:
            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        
        # Mini-batch training for large datasets
        dataset_size = len(states)
        indices = torch.randperm(dataset_size)
        
        total_policy_loss = 0
        total_value_loss = 0
        total_entropy = 0
        update_count = 0
        
        for epoch in range(epochs):
            # Process in mini-batches to avoid memory issues
            for start_idx in range(0, dataset_size, self.mini_batch_size):
                end_idx = min(start_idx + self.mini_batch_size, dataset_size)
                batch_indices = indices[start_idx:end_idx]
                
                # Get mini-batch
                batch_states = states[batch_indices]
                batch_actions = actions[batch_indices]
                batch_advantages = advantages[batch_indices]
                batch_returns = returns[batch_indices]
                batch_old_log_probs = old_log_probs[batch_indices]
                
                # Forward pass
                logits = self.policy(batch_states)
                probs = torch.softmax(logits, dim=-1) + 1e-8
                probs = probs / probs.sum(dim=-1, keepdim=True)
                
                dist = Categorical(probs)
                new_log_probs = dist.log_prob(batch_actions)
                
                # Policy loss
                ratio = torch.exp(new_log_probs - batch_old_log_probs)
                ratio = torch.clamp(ratio, 0.1, 10.0)
                
                surr1 = ratio * batch_advantages
                surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * batch_advantages
                policy_loss = -torch.min(surr1, surr2).mean()
                
                # Value loss
                current_values = self.value(batch_states).squeeze(-1)
                value_loss = nn.MSELoss()(current_values, batch_returns)
                
                # Entropy
                entropy = dist.entropy().mean()
                
                # Total loss
                total_loss = policy_loss + 0.5 * value_loss - 0.01 * entropy  # Reduced entropy coefficient
                
                # Check for NaN
                if torch.isnan(total_loss) or torch.isinf(total_loss):
                    print("Warning: NaN/Inf loss detected, skipping batch")
                    continue
                
                # Optimize
                self.optimizer.zero_grad()
                total_loss.backward()
                
                # Gradient clipping
                torch.nn.utils.clip_grad_norm_(self.policy.parameters(), 0.5)
                torch.nn.utils.clip_grad_norm_(self.value.parameters(), 0.5)
                
                self.optimizer.step()
                
                # Accumulate losses
                total_policy_loss += policy_loss.detach().item()
                total_value_loss += value_loss.detach().item()
                total_entropy += entropy.detach().item()
                update_count += 1
        
        # Clear memory
        self.memory.clear()
        
        if update_count == 0:
            return None
            
        return {
            'policy_loss': total_policy_loss / update_count,
            'value_loss': total_value_loss / update_count,
            'entropy': total_entropy / update_count
        }
    
    def save(self, filepath):
        """Save the model"""
        torch.save({
            'policy_state_dict': self.policy.state_dict(),
            'value_state_dict': self.value.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict()
        }, filepath)
    
    def load(self, filepath):
        """Load the model"""
        checkpoint = torch.load(filepath, map_location='cpu')
        self.policy.load_state_dict(checkpoint['policy_state_dict'])
        self.value.load_state_dict(checkpoint['value_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])