Update src/environments/visual_trading_env.py

src/environments/visual_trading_env.py

@@ -46,60 +46,87 @@ class VisualTradingEnvironment:
     
     def _get_visual_observation(self):
         """Generate visual representation of current market state"""
-        # Get recent price window
-        window_size = 50
-        start_idx = max(0, self.current_step - window_size)
-        end_idx = self.current_step + 1
-        
-        if end_idx > len(self.price_data):
-            end_idx = len(self.price_data)
-        
-        prices = self.price_data[start_idx:end_idx]
-        
-        # Create matplotlib figure
-        fig, ax = plt.subplots(figsize=(4.2, 4.2), dpi=20, facecolor='black')
-        ax.set_facecolor('black')
-        
-        # Plot price
-        ax.plot(prices, color='cyan', linewidth=1.5)
-        
-        # Remove axes for cleaner visual
-        ax.set_xticks([])
-        ax.set_yticks([])
-        ax.spines['top'].set_visible(False)
-        ax.spines['right'].set_visible(False)
-        ax.spines['bottom'].set_visible(False)
-        ax.spines['left'].set_visible(False)
-        
-        # Convert to numpy array
-        buf = io.BytesIO()
-        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, facecolor='black')
-        buf.seek(0)
-        img = Image.open(buf).convert('RGB')
-        img_array = np.array(img)
-        
-        plt.close(fig)
-        
-        # Create simple attention map (in real implementation, this would come from CNN)
-        attention_map = np.zeros((84, 84))
-        if len(prices) > 1:
-            recent_change = (prices[-1] - prices[-2]) / prices[-2]
-            intensity = min(255, abs(recent_change) * 5000)
-            
-            # Simple attention based on price movement
-            center_x, center_y = 42, 42
-            size = max(5, int(intensity / 50))
-            
-            for i in range(max(0, center_x-size), min(84, center_x+size)):
-                for j in range(max(0, center_y-size), min(84, center_y+size)):
-                    distance = np.sqrt((i-center_x)**2 + (j-center_y)**2)
-                    if distance <= size:
-                        attention_map[i, j] = max(attention_map[i, j], intensity * (1 - distance/size))
-        
-        # Combine RGB with attention map
-        visual_obs = np.concatenate([img_array, attention_map[:, :, np.newaxis]], axis=2)
-        
-        return visual_obs
+        try:
+            # Get recent price window
+            window_size = 50
+            start_idx = max(0, self.current_step - window_size)
+            end_idx = self.current_step + 1
+            
+            if end_idx > len(self.price_data):
+                end_idx = len(self.price_data)
+            
+            prices = self.price_data[start_idx:end_idx]
+            
+            # Create matplotlib figure with fixed size
+            fig, ax = plt.subplots(figsize=(4.2, 4.2), dpi=20, facecolor='black')
+            ax.set_facecolor('black')
+            
+            # Plot price
+            if len(prices) > 0:
+                ax.plot(prices, color='cyan', linewidth=1.5)
+            
+            # Remove axes for cleaner visual
+            ax.set_xticks([])
+            ax.set_yticks([])
+            ax.spines['top'].set_visible(False)
+            ax.spines['right'].set_visible(False)
+            ax.spines['bottom'].set_visible(False)
+            ax.spines['left'].set_visible(False)
+            
+            # Set fixed limits to ensure consistent size
+            ax.set_xlim(0, 50)
+            if len(prices) > 0:
+                price_min, price_max = min(prices), max(prices)
+                price_range = price_max - price_min
+                if price_range == 0:
+                    price_range = 1
+                ax.set_ylim(price_min - price_range * 0.1, price_max + price_range * 0.1)
+            else:
+                ax.set_ylim(0, 100)
+            
+            # Convert to numpy array with consistent size
+            buf = io.BytesIO()
+            plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, facecolor='black', dpi=20)
+            buf.seek(0)
+            img = Image.open(buf).convert('RGB')
+            
+            # Resize to consistent dimensions
+            img = img.resize((84, 84), Image.Resampling.LANCZOS)
+            img_array = np.array(img)
+            
+            plt.close(fig)
+            
+            # Create attention map with same dimensions
+            attention_map = np.zeros((84, 84))
+            if len(prices) > 1:
+                recent_change = (prices[-1] - prices[-2]) / prices[-2]
+                intensity = min(255, abs(recent_change) * 5000)
+                
+                # Simple attention based on price movement
+                center_x, center_y = 42, 42
+                size = max(5, int(intensity / 50))
+                
+                for i in range(max(0, center_x-size), min(84, center_x+size)):
+                    for j in range(max(0, center_y-size), min(84, center_y+size)):
+                        distance = np.sqrt((i-center_x)**2 + (j-center_y)**2)
+                        if distance <= size:
+                            attention_map[i, j] = max(attention_map[i, j], intensity * (1 - distance/size))
+            
+            # Ensure both arrays have same shape before concatenation
+            attention_map = attention_map.astype(np.uint8)
+            
+            # Combine RGB with attention map - now both are 84x84
+            visual_obs = np.concatenate([
+                img_array, 
+                attention_map[:, :, np.newaxis]  # Add channel dimension
+            ], axis=2)
+            
+            return visual_obs
+            
+        except Exception as e:
+            print(f"Error in visual observation: {e}")
+            # Return default observation in case of error
+            return np.zeros((84, 84, 4), dtype=np.uint8)
     
     def reset(self):
         """Reset environment to initial state"""
@@ -174,9 +201,11 @@ class VisualTradingEnvironment:
             'step': self.current_step
         }
         
-        return self._get_visual_observation(), reward, self.done, info
+        obs = self._get_visual_observation()
+        return obs, reward, self.done, info
     
     def get_price_history(self):
         """Get recent price history for visualization"""
         window_size = min(50, self.current_step)
-        return self.price_data[self.current_step - window_size:self.current_step].tolist()
+        start_idx = max(0, self.current_step - window_size)
+        return self.price_data[start_idx:self.current_step].tolist()