Update app.py

app.py

@@ -2,7 +2,7 @@
 Advanced AI Trading Demo - Hugging Face Spaces
 Deep Q-Network (DQN) Reinforcement Learning for Financial Trading Simulation
 
-Author: Your Name
+Author: AI Trading Team
 License: MIT
 """
 
@@ -20,6 +20,7 @@ import logging
 import os
 from datetime import datetime
 import json
+from dataclasses import dataclass  # ✅ Added missing import
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -152,7 +153,8 @@ class AdvancedTradingEnvironment:
             losses = -returns[returns < 0]
             avg_gain = np.mean(gains[-14:]) if len(gains) > 0 else 0.001
             avg_loss = np.mean(losses[-14:]) if len(losses) > 0 else 0.001
-            rsi = 100 - (100 / (1 + avg_gain / avg_loss))
+            rs = avg_gain / avg_loss if avg_loss != 0 else 100
+            rsi = 100 - (100 / (1 + rs))
         else:
             rsi = 50.0
             
@@ -178,6 +180,7 @@ class AdvancedTradingEnvironment:
     def reset(self) -> Tuple[np.ndarray, Dict]:
         """Reset environment to initial state."""
         self._reset_state()
+        self._initialize_market_data()  # ✅ Reinitialize market data on reset
         obs = self._get_observation()
         info = self._get_info()
         return obs, info
@@ -302,7 +305,7 @@ class AdvancedTradingEnvironment:
             np.mean(recent_prices) / 100.0,
             np.std(recent_prices) / 100.0,
             (self.current_price - np.min(recent_prices)) / 
-            (np.max(recent_prices) - np.min(recent_prices) + 1e-8) if len(set(recent_prices)) > 1 else 0.5
+            (np.max(recent_prices) - np.min(recent_prices) + 1e-8)
         ]
         
         # Portfolio features
@@ -324,9 +327,8 @@ class AdvancedTradingEnvironment:
         # Technical indicators
         technical_features = self._calculate_technical_indicators()
         
-        # Combine and truncate/pad to fixed size
-        all_features = (price_features + portfolio_features + sentiment_features + 
-                       technical_features[:6])  # Ensure exactly 15 features
+        # Combine all features (should be 4 + 3 + 3 + 6 = 16, take first 15)
+        all_features = price_features + portfolio_features + sentiment_features + technical_features
         observation = np.array(all_features[:15], dtype=np.float32)
         
         return observation
@@ -464,15 +466,21 @@ class DQNAgent:
             'config': self.config.__dict__
         }
         torch.save(checkpoint, path)
+        logger.info(f"Model saved to {path}")
     
     def load_checkpoint(self, path: str):
         """Load agent checkpoint."""
-        checkpoint = torch.load(path, map_location=self.device)
-        self.q_network.load_state_dict(checkpoint['q_network'])
-        self.target_network.load_state_dict(checkpoint['target_network'])
-        self.optimizer.load_state_dict(checkpoint['optimizer'])
-        self.epsilon = checkpoint['epsilon']
-        self.steps = checkpoint['steps']
+        if os.path.exists(path):
+            try:
+                checkpoint = torch.load(path, map_location=self.device)
+                self.q_network.load_state_dict(checkpoint['q_network'])
+                self.target_network.load_state_dict(checkpoint['target_network'])
+                self.optimizer.load_state_dict(checkpoint['optimizer'])
+                self.epsilon = checkpoint['epsilon']
+                self.steps = checkpoint['steps']
+                logger.info(f"Model loaded from {path}")
+            except Exception as e:
+                logger.warning(f"Failed to load model from {path}: {e}")
 
 # ---- 4. Main Trading Application ----
 class TradingDemo:
@@ -498,12 +506,23 @@ class TradingDemo:
     
     def load_model_if_exists(self):
         """Load existing model if available."""
+        self.agent = None  # Reset agent first
         if os.path.exists(self.model_path):
             try:
+                # Create agent first, then load
+                temp_config = TradingConfig()
+                temp_env = AdvancedTradingEnvironment(temp_config)
+                self.agent = DQNAgent(
+                    state_dim=temp_env.observation_space_dim,
+                    action_dim=temp_env.action_space,
+                    config=temp_config,
+                    device=self.device
+                )
                 self.agent.load_checkpoint(self.model_path)
-                logger.info("Loaded existing model checkpoint")
+                logger.info("✅ Loaded existing model checkpoint")
             except Exception as e:
                 logger.warning(f"Failed to load model: {e}")
+                self.agent = None
     
     def initialize(self, balance: float, risk: str, asset: str) -> str:
         """Initialize trading system with new parameters."""
@@ -556,7 +575,8 @@ class TradingDemo:
                 
                 while not done:
                     action = self.agent.select_action(obs)
-                    next_obs, reward, done, _, info = self.env.step(action)
+                    next_obs, reward, terminated, truncated, info = self.env.step(action)
+                    done = terminated or truncated
                     
                     self.agent.store_transition(obs, action, reward, next_obs, done)
                     loss = self.agent.update()
@@ -590,7 +610,8 @@ class TradingDemo:
                     yield progress, None
             
             # Save trained model
-            self.agent.save_checkpoint(self.model_path)
+            if self.agent:
+                self.agent.save_checkpoint(self.model_path)
             yield "✅ Training completed! Model saved.", self._create_training_plot()
             
         except Exception as e:
@@ -611,7 +632,7 @@ class TradingDemo:
             
             for step in range(steps):
                 action = self.agent.select_action(obs, training=False)
-                next_obs, _, done, _, info = self.env.step(action)
+                next_obs, _, terminated, truncated, info = self.env.step(action)
                 
                 prices.append(self.env.current_price)
                 actions.append(action)
@@ -620,7 +641,7 @@ class TradingDemo:
                 cash_balances.append(info['cash_balance'])
                 
                 obs = next_obs
-                if done:
+                if terminated or truncated:
                     break
             
             plot = self._create_simulation_plot(
@@ -630,11 +651,12 @@ class TradingDemo:
             final_return = ((net_worths[-1] - self.config.initial_balance) / 
                           self.config.initial_balance * 100)
             
+            last_action_name = DQNAgent.ACTION_NAMES.get(actions[-1], 'Unknown')
             result = (f"✅ Simulation completed!\n"
                      f"📈 Steps: {len(prices)}\n"
                      f"💰 Final Net Worth: ${net_worths[-1]:,.2f}\n"
                      f"📊 Total Return: {final_return:.2f}%\n"
-                     f"🎯 Final Action: {DQNAgent.ACTION_NAMES.get(actions[-1], 'Unknown')}")
+                     f"🎯 Final Action: {last_action_name}")
             
             return result, plot
             
@@ -652,9 +674,7 @@ class TradingDemo:
         fig = make_subplots(
             rows=2, cols=2,
             subplot_titles=('Episode Rewards', 'Training Loss', 'Epsilon Decay', 'Portfolio Performance'),
-            vertical_spacing=0.12,
-            specs=[[{"secondary_y": False}, {"secondary_y": False}],
-                   [{"secondary_y": True}, {"secondary_y": False}]]
+            vertical_spacing=0.12
         )
         
         # Rewards
@@ -692,16 +712,16 @@ class TradingDemo:
         # Portfolio performance
         returns = [(nw - self.config.initial_balance) / self.config.initial_balance * 100
                   for nw in self.training_history['net_worths']]
-        fig.add_trace(
-            go.Scatter(x=episodes, y=returns, mode='lines',
-                      name='Return %', line=dict(color='purple')),
-            row=2, col=2, secondary_y=True
-        )
         fig.add_trace(
             go.Scatter(x=episodes, y=self.training_history['net_worths'],
                       mode='lines', name='Net Worth',
                       line=dict(color='blue'), yaxis='y'),
-            row=2, col=2, secondary_y=False
+            row=2, col=2
+        )
+        fig.add_trace(
+            go.Scatter(x=episodes, y=returns, mode='lines',
+                      name='Return %', line=dict(color='purple'), yaxis='y2'),
+            row=2, col=2
         )
         
         fig.update_layout(
@@ -712,7 +732,7 @@ class TradingDemo:
         )
         
         fig.update_yaxes(title_text="Return (%)", secondary_y=True, row=2, col=2)
-        fig.update_yaxes(title_text="Net Worth ($)", secondary_y=False, row=2, col=2)
+        fig.update_yaxes(title_text="Net Worth ($)", row=2, col=2)
         
         return fig
     
@@ -726,9 +746,7 @@ class TradingDemo:
             rows=2, cols=2,
             subplot_titles=('Price Action & Trading Signals', 'Portfolio Performance',
                           'Portfolio Allocation', 'Action Distribution'),
-            vertical_spacing=0.12,
-            specs=[[{"secondary_y": False}, {"secondary_y": True}],
-                   [{"secondary_y": False}, {"secondary_y": False}]]
+            vertical_spacing=0.12
         )
         
         # Price and actions
@@ -748,7 +766,7 @@ class TradingDemo:
                 action_prices = [prices[i] for i in action_steps]
                 fig.add_trace(
                     go.Scatter(x=action_steps, y=action_prices, mode='markers',
-                              name=f'{name} Signals',
+                              name=f'{name}',
                               marker=dict(color=color, size=8, symbol='triangle-up')),
                     row=1, col=1
                 )
@@ -760,12 +778,12 @@ class TradingDemo:
         fig.add_trace(
             go.Scatter(x=steps, y=net_worths, mode='lines', name='Net Worth',
                       line=dict(color='purple', width=2)),
-            row=1, col=2, secondary_y=False
+            row=1, col=2
         )
         fig.add_trace(
             go.Scatter(x=steps, y=returns, mode='lines', name='Returns %',
-                      line=dict(color='orange', width=2)),
-            row=1, col=2, secondary_y=True
+                      line=dict(color='orange', width=2), yaxis='y2'),
+            row=1, col=2
         )
         
         # Portfolio composition
@@ -796,7 +814,7 @@ class TradingDemo:
         )
         
         fig.update_yaxes(title_text="Returns (%)", secondary_y=True, row=1, col=2)
-        fig.update_yaxes(title_text="Value ($)", secondary_y=False, row=1, col=2)
+        fig.update_yaxes(title_text="Value ($)", row=1, col=2)
         
         return fig
 
@@ -810,23 +828,15 @@ def create_interface() -> gr.Blocks:
         title="🤖 Advanced AI Trading Demo",
         css="""
         .gradio-container {max-width: 1400px !important;}
-        .status-box {background-color: #f0f9ff; padding: 1rem; border-radius: 8px;}
+        .status-box {background-color: #f0f9ff; padding: 1rem; border-radius: 8px; border-left: 4px solid #3b82f6;}
         """
     ) as interface:
         
         gr.Markdown("""
         # 🤖 Advanced AI Trading Demo
         **Deep Reinforcement Learning for Financial Markets**
-        
-        This demo showcases a **Deep Q-Network (DQN)** agent learning to trade in simulated financial markets. 
-        The agent uses technical indicators, sentiment analysis, and risk management to optimize trading strategies.
-        
-        **Key Features:**
-        - 📊 Multi-asset support (Crypto, Stocks, Forex)
-        - 🎯 Risk-adjusted position sizing
-        - 🧠 Deep Q-Network with experience replay
-        - 📈 Real-time training visualization
-        - 💾 Model persistence across sessions
+
+        This demo showcases a **Deep Q-Network (DQN)** agent learning to trade in simulated financial markets with realistic market dynamics, technical indicators, and risk management.
         """)
         
         # Configuration Row
@@ -836,20 +846,17 @@ def create_interface() -> gr.Blocks:
                 with gr.Group():
                     balance = gr.Slider(
                         1000, 50000, 10000, step=1000,
-                        label="💰 Initial Balance ($)",
-                        info="Starting capital for trading"
+                        label="💰 Initial Balance ($)"
                     )
                     risk = gr.Radio(
                         ["Low", "Medium", "High"], value="Medium",
-                        label="🎯 Risk Level",
-                        info="Affects position sizing and volatility"
+                        label="🎯 Risk Level"
                     )
                     asset = gr.Radio(
                         ["Crypto", "Stock", "Forex"], value="Crypto",
-                        label="📈 Asset Type",
-                        info="Different volatility characteristics"
+                        label="📈 Asset Type"
                     )
-                    init_btn = gr.Button("🚀 Initialize Trading System", variant="primary", size="lg")
+                    init_btn = gr.Button("🚀 Initialize Trading System", variant="primary")
             
             with gr.Column(scale=2):
                 gr.Markdown("## 📊 System Status")
@@ -864,8 +871,7 @@ def create_interface() -> gr.Blocks:
                 gr.Markdown("## 🏋️‍♂️ Train AI Agent")
                 with gr.Group():
                     episodes = gr.Number(
-                        value=100, label="🎯 Training Episodes",
-                        minimum=10, maximum=1000, precision=0
+                        value=50, label="🎯 Training Episodes", precision=0
                     )
                     train_btn = gr.Button("🎓 Start Training", variant="primary")
                     training_output = gr.Textbox(
@@ -877,8 +883,7 @@ def create_interface() -> gr.Blocks:
                 gr.Markdown("## ▶️ Test Trained Agent")
                 with gr.Group():
                     sim_steps = gr.Number(
-                        value=200, label="📊 Simulation Steps",
-                        minimum=50, maximum=1000, precision=0
+                        value=200, label="📊 Simulation Steps", precision=0
                     )
                     sim_btn = gr.Button("🎮 Run Simulation", variant="primary")
                     sim_output = gr.Textbox(
@@ -887,78 +892,66 @@ def create_interface() -> gr.Blocks:
                     sim_plot = gr.Plot(label="📊 Trading Results")
         
         # Event Handlers
+        def initialize_wrapper(balance, risk, asset):
+            return demo.initialize(balance, risk, asset)
+        
+        def simulate_wrapper(steps):
+            return demo.simulate(steps)
+        
+        def train_generator(episodes):
+            try:
+                for status_text, plot in demo.train(int(episodes)):
+                    yield status_text, plot
+            except Exception as e:
+                yield f"❌ Training error: {str(e)}", None
+        
         init_btn.click(
-            fn=demo.initialize,
+            fn=initialize_wrapper,
             inputs=[balance, risk, asset],
             outputs=status
         )
         
-        def train_generator(episodes):
-            for status_text, plot in demo.train(episodes):
-                yield status_text, plot
-        
-        train_btn.click(
+        train_btn.queue().click(
             fn=train_generator,
             inputs=episodes,
             outputs=[training_output, train_plot]
         )
         
         sim_btn.click(
-            fn=demo.simulate,
+            fn=simulate_wrapper,
             inputs=sim_steps,
             outputs=[sim_output, sim_plot]
         )
         
-        # Instructions
         gr.Markdown("""
-        ## 📖 How to Use
-        
-        ### 1. **Configuration**
-        - Set your **initial balance**, **risk level**, and **asset type**
-        - Risk level affects position sizing and market volatility
-        - Different assets have unique volatility characteristics
-        
-        ### 2. **Training**
-        - Click **Initialize** to setup the environment
-        - Start training with 100+ episodes for good results
-        - Monitor training progress through real-time charts
-        - Model is automatically saved for future use
-        
-        ### 3. **Simulation**
-        - Run simulations to test the trained agent's performance
-        - Observe trading decisions and portfolio evolution
-        - Analyze action distribution and performance metrics
+        ## 📖 Usage Instructions
+        1. **Configure** your trading parameters
+        2. **Initialize** the trading system
+        3. **Train** the AI agent (50+ episodes recommended)
+        4. **Simulate** trading with the trained agent
         
         ## 🎮 Trading Actions
-        | Action | ID | Description |
-        |--------|----|-------------|
-        | **Hold** | 0 | Maintain current position |
-        | **Buy** | 1 | Purchase assets (risk-adjusted size) |
-        | **Sell** | 2 | Sell portion of position |
-        | **Close** | 3 | Liquidate entire position |
-        
-        ## 🔬 Technical Details
-        - **Algorithm**: Deep Q-Network (DQN) with experience replay
-        - **State Space**: 15 features (price, technical indicators, portfolio)
-        - **Reward**: Risk-adjusted returns with drawdown penalties
-        - **Exploration**: Epsilon-greedy with decay
+        - **Hold (0)**: Maintain current position
+        - **Buy (1)**: Purchase assets (risk-adjusted)
+        - **Sell (2)**: Sell portion of position
+        - **Close (3)**: Liquidate entire position
         """)
     
     return interface
 
 # ---- 6. Hugging Face Spaces Entry Point ----
 if __name__ == "__main__":
-    # Required for Hugging Face Spaces
-    demo = create_interface()
-    
-    # Launch with Hugging Face Spaces configuration
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,  # Set to True for public sharing
-        show_error=True,
-        show_tips=False,
-        enable_queue=True,
-        max_threads=40,
-        debug=False
-    )
+    try:
+        interface = create_interface()
+        interface.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False,
+            show_error=True,
+            enable_queue=True,
+            max_threads=40,
+            debug=False
+        )
+    except Exception as e:
+        logger.error(f"Failed to launch application: {e}")
+        raise