Create src/agents/advanced_agent.py

src/agents/advanced_agent.py

@@ -0,0 +1,168 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import numpy as np
+from collections import deque
+import random
+from .visual_agent import VisualTradingAgent, SimpleTradingNetwork
+
+class AdvancedTradingAgent(VisualTradingAgent):
+    def __init__(self, state_dim, action_dim, learning_rate=0.001, use_sentiment=True):
+        super().__init__(state_dim, action_dim, learning_rate)
+        
+        self.use_sentiment = use_sentiment
+        self.sentiment_history = deque(maxlen=50)
+        
+        # Enhanced network architecture for sentiment analysis
+        if use_sentiment:
+            self.policy_net = EnhancedTradingNetwork(state_dim, action_dim)
+            self.policy_net = self.policy_net.to(self.device)
+            self.optimizer = optim.Adam(self.policy_net.parameters(), lr=learning_rate)
+    
+    def select_action(self, state, current_sentiment=0.5, sentiment_confidence=0.0):
+        """Select action with sentiment consideration"""
+        if random.random() < self.epsilon:
+            return random.randint(0, self.action_dim - 1)
+        
+        try:
+            state_normalized = state.astype(np.float32) / 255.0
+            state_tensor = torch.FloatTensor(state_normalized).unsqueeze(0).to(self.device)
+            
+            if self.use_sentiment:
+                # Add sentiment to the decision process
+                sentiment_tensor = torch.FloatTensor([current_sentiment, sentiment_confidence]).unsqueeze(0).to(self.device)
+                with torch.no_grad():
+                    q_values = self.policy_net(state_tensor, sentiment_tensor)
+            else:
+                with torch.no_grad():
+                    q_values = self.policy_net(state_tensor)
+            
+            return int(q_values.argmax().item())
+            
+        except Exception as e:
+            print(f"Error in advanced action selection: {e}")
+            return random.randint(0, self.action_dim - 1)
+    
+    def store_transition(self, state, action, reward, next_state, done, sentiment_data=None):
+        """Store experience with sentiment data"""
+        experience = (state, action, reward, next_state, done, sentiment_data)
+        self.memory.append(experience)
+    
+    def update(self):
+        """Update network with sentiment-enhanced learning"""
+        if len(self.memory) < self.batch_size:
+            return 0.0
+        
+        try:
+            batch = random.sample(self.memory, self.batch_size)
+            states, actions, rewards, next_states, dones, sentiment_data = zip(*batch)
+            
+            # Convert to tensors
+            states_tensor = torch.FloatTensor(np.array(states)).to(self.device) / 255.0
+            actions_tensor = torch.LongTensor(actions).to(self.device)
+            rewards_tensor = torch.FloatTensor(rewards).to(self.device)
+            next_states_tensor = torch.FloatTensor(np.array(next_states)).to(self.device) / 255.0
+            dones_tensor = torch.BoolTensor(dones).to(self.device)
+            
+            if self.use_sentiment and sentiment_data[0] is not None:
+                # Extract sentiment features
+                sentiment_features = []
+                for data in sentiment_data:
+                    if data:
+                        sentiment_features.append([data.get('sentiment', 0.5), data.get('confidence', 0.0)])
+                    else:
+                        sentiment_features.append([0.5, 0.0])
+                
+                sentiment_tensor = torch.FloatTensor(sentiment_features).to(self.device)
+                next_sentiment_tensor = sentiment_tensor  # Simplified
+                
+                # Current Q values with sentiment
+                current_q = self.policy_net(states_tensor, sentiment_tensor).gather(1, actions_tensor.unsqueeze(1))
+                
+                # Next Q values with sentiment
+                with torch.no_grad():
+                    next_q = self.policy_net(next_states_tensor, next_sentiment_tensor).max(1)[0]
+                    target_q = rewards_tensor + (self.gamma * next_q * ~dones_tensor)
+            else:
+                # Fallback to standard DQN
+                current_q = self.policy_net(states_tensor).gather(1, actions_tensor.unsqueeze(1))
+                
+                with torch.no_grad():
+                    next_q = self.policy_net(next_states_tensor).max(1)[0]
+                    target_q = rewards_tensor + (self.gamma * next_q * ~dones_tensor)
+            
+            # Compute loss
+            loss = nn.MSELoss()(current_q.squeeze(), target_q)
+            
+            # Optimize
+            self.optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(self.policy_net.parameters(), 1.0)
+            self.optimizer.step()
+            
+            # Update exploration
+            self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)
+            
+            return float(loss.item())
+            
+        except Exception as e:
+            print(f"Error in advanced update: {e}")
+            return 0.0
+
+class EnhancedTradingNetwork(nn.Module):
+    def __init__(self, state_dim, action_dim, sentiment_dim=2):
+        super(EnhancedTradingNetwork, self).__init__()
+        
+        # Visual processing branch (same as before)
+        self.visual_conv = nn.Sequential(
+            nn.Conv2d(4, 16, kernel_size=4, stride=2),
+            nn.ReLU(),
+            nn.Conv2d(16, 32, kernel_size=4, stride=2),
+            nn.ReLU(),
+            nn.Conv2d(32, 32, kernel_size=3, stride=1),
+            nn.ReLU(),
+            nn.AdaptiveAvgPool2d((8, 8))
+        )
+        
+        self.visual_fc = nn.Sequential(
+            nn.Linear(32 * 8 * 8, 256),
+            nn.ReLU(),
+            nn.Dropout(0.3)
+        )
+        
+        # Sentiment processing branch
+        self.sentiment_fc = nn.Sequential(
+            nn.Linear(sentiment_dim, 64),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(64, 32),
+            nn.ReLU()
+        )
+        
+        # Combined decision making
+        self.combined_fc = nn.Sequential(
+            nn.Linear(256 + 32, 128),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(128, 64),
+            nn.ReLU(),
+            nn.Linear(64, action_dim)
+        )
+        
+    def forward(self, x, sentiment=None):
+        # Visual processing
+        x = x.permute(0, 3, 1, 2)  # (batch, 84, 84, 4) -> (batch, 4, 84, 84)
+        visual_features = self.visual_conv(x)
+        visual_features = visual_features.view(visual_features.size(0), -1)
+        visual_features = self.visual_fc(visual_features)
+        
+        # Sentiment processing
+        if sentiment is not None:
+            sentiment_features = self.sentiment_fc(sentiment)
+            combined_features = torch.cat([visual_features, sentiment_features], dim=1)
+        else:
+            combined_features = visual_features
+        
+        # Final decision
+        q_values = self.combined_fc(combined_features)
+        return q_values