Update app.py

app.py

@@ -1,30 +1,38 @@
 import gradio as gr
 import numpy as np
 import torch
-from pathlib import Path
-from typing import Dict, Tuple, Any
-from loguru import logger
-import yaml
-from gymnasium import spaces
+import random
+from collections import deque
+from typing import Tuple, Iterator, Any
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+import logging
+from datetime import datetime
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class TradingConfig:
-    def __init__(self):
-        self.initial_balance = 10000.0
+    """Configuration for trading environment"""
+    def __init__(self, initial_balance: float = 10000.0):
+        self.initial_balance = initial_balance
         self.max_steps = 1000
         self.transaction_cost = 0.001
-        self.risk_level = "Medium"
-        self.asset_type = "Crypto"
         self.learning_rate = 0.0001
         self.gamma = 0.99
         self.epsilon_start = 1.0
         self.epsilon_min = 0.01
-        self.epsilon_decay = 0.9995
+        self.epsilon_decay = 0.999
         self.batch_size = 32
-        self.memory_size = 10000
-        self.target_update = 100
+        self.memory_size = 5000  # Reduced for Spaces memory limits
+        self.target_update = 50
 
 class AdvancedTradingEnvironment:
-    def __init__(self, config):
+    """Simplified trading environment compatible with HF Spaces"""
+    
+    def __init__(self, config: TradingConfig):
+        self.config = config
         self.initial_balance = config.initial_balance
         self.balance = self.initial_balance
         self.position = 0.0
@@ -33,147 +41,218 @@ class AdvancedTradingEnvironment:
         self.max_steps = config.max_steps
         self.price_history = []
         self.sentiment_history = []
+        self.action_history = []
+        
+        # Initialize data
         self._initialize_data()
-        self.action_space = spaces.Discrete(4)
-        self.observation_space = spaces.Box(low=-2.0, high=2.0, shape=(12,), dtype=np.float32)
-
+        
+        # Spaces: Discrete actions (0=Hold, 1=Buy, 2=Sell, 3=Close)
+        self.action_space = gr.utils.Discrete(4)
+        self.observation_space = None  # For compatibility
+    
     def _initialize_data(self):
-        n_points = 100
+        """Initialize price and sentiment history"""
+        n_points = 50  # Reduced for faster init
         base_price = 100.0
+        
         for i in range(n_points):
-            price = base_price + np.sin(i * 0.1) * 10 + np.random.normal(0, 2)
-            self.price_history.append(max(10.0, price))
-            sentiment = 0.5 + np.random.normal(0, 0.1)
+            # Realistic price simulation
+            trend = np.sin(i * 0.2) * 5
+            noise = np.random.normal(0, 3)
+            price = max(50.0, base_price + trend + noise)
+            self.price_history.append(price)
+            
+            # Correlated sentiment
+            sentiment = 0.5 + 0.3 * np.tanh((price - base_price) / base_price) + np.random.normal(0, 0.1)
             self.sentiment_history.append(np.clip(sentiment, 0.0, 1.0))
+        
         self.current_price = self.price_history[-1]
-
-    def reset(self):
+    
+    def reset(self) -> Tuple[np.ndarray, dict]:
+        """Reset environment to initial state"""
         self.balance = self.initial_balance
         self.position = 0.0
         self.step_count = 0
-        self.price_history = [100.0 + np.random.normal(0, 5)]
-        self.sentiment_history = [0.5]
+        self.action_history = []
+        
+        # Reset price series
+        self._initialize_data()
+        
         obs = self._get_observation()
         info = self._get_info()
+        
+        logger.info(f"Environment reset: balance=${self.balance}, price=${self.current_price}")
         return obs, info
-
-    def step(self, action):
+    
+    def step(self, action: int) -> Tuple[np.ndarray, float, bool, bool, dict]:
+        """Execute one step"""
         self.step_count += 1
-        price_change = np.random.normal(0, 0.02)
-        self.current_price = max(10.0, self.current_price * (1 + price_change))
-        self.price_history.append(self.current_price)
+        self.action_history.append(action)
         
-        sentiment_change = np.random.normal(0, 0.05)
-        new_sentiment = np.clip(self.sentiment_history[-1] + sentiment_change, 0.0, 1.0)
-        self.sentiment_history.append(new_sentiment)
+        # Market evolution
+        self._market_step()
         
+        # Execute action
         reward = self._execute_action(action)
         
-        terminated = self.balance <= 0 or self.step_count >= self.max_steps
+        # Check termination
+        terminated = (self.balance <= 0 or self.step_count >= self.max_steps)
         truncated = False
         
         obs = self._get_observation()
         info = self._get_info()
         
         return obs, reward, terminated, truncated, info
-
-    def _execute_action(self, action):
+    
+    def _market_step(self):
+        """Simulate market movement"""
+        # Price evolution (geometric brownian motion approximation)
+        drift = 0.0001
+        volatility = 0.02
+        price_change = drift + volatility * np.random.normal()
+        self.current_price = max(10.0, self.current_price * np.exp(price_change))
+        
+        self.price_history.append(self.current_price)
+        if len(self.price_history) > 100:
+            self.price_history.pop(0)
+        
+        # Update sentiment
+        sentiment_change = 0.1 * price_change + np.random.normal(0, 0.05)
+        new_sentiment = np.clip(self.sentiment_history[-1] + sentiment_change, 0.0, 1.0)
+        self.sentiment_history.append(new_sentiment)
+        if len(self.sentiment_history) > 100:
+            self.sentiment_history.pop(0)
+    
+    def _execute_action(self, action: int) -> float:
+        """Execute trading action and return reward"""
         reward = 0.0
         prev_net_worth = self.balance + self.position * self.current_price
         
+        # Transaction cost
+        cost_factor = self.config.transaction_cost
+        
         if action == 1:  # Buy
-            trade_amount = min(self.balance * 0.2, self.balance)
-            cost = trade_amount
-            if cost <= self.balance:
-                self.position += trade_amount / self.current_price
-                self.balance -= cost
+            if self.balance > 100:  # Minimum trade size
+                trade_value = min(self.balance * 0.1, self.balance - 100)
+                cost = trade_value * (1 + cost_factor)
+                if cost <= self.balance:
+                    shares = trade_value / self.current_price
+                    self.position += shares
+                    self.balance -= cost
+                    reward -= cost_factor * trade_value
         
         elif action == 2:  # Sell
             if self.position > 0:
-                sell_amount = min(self.position * 0.2, self.position)
-                proceeds = sell_amount * self.current_price
-                self.position -= sell_amount
+                sell_shares = min(self.position * 0.1, self.position)
+                proceeds = sell_shares * self.current_price * (1 - cost_factor)
+                self.position -= sell_shares
                 self.balance += proceeds
+                reward -= cost_factor * sell_shares * self.current_price
         
-        elif action == 3:  # Close
+        elif action == 3:  # Close position
             if self.position > 0:
-                proceeds = self.position * self.current_price
+                proceeds = self.position * self.current_price * (1 - cost_factor)
                 self.balance += proceeds
                 self.position = 0
+                reward -= cost_factor * abs(self.position) * self.current_price * self.current_price
         
-        net_worth = self.balance + self.position * self.current_price
-        reward = (net_worth - prev_net_worth) / self.initial_balance * 100
+        # Portfolio reward
+        current_net_worth = self.balance + self.position * self.current_price
+        pnl = (current_net_worth - prev_net_worth) / self.initial_balance
+        reward += pnl * 100  # Scale reward
         
         return reward
-
-    def _get_observation(self):
-        recent_prices = self.price_history[-10:] if len(self.price_history) >= 10 else [self.current_price] * 10
-        recent_sentiments = self.sentiment_history[-10:] if len(self.sentiment_history) >= 10 else [0.5] * 10
+    
+    def _get_observation(self) -> np.ndarray:
+        """Get current state observation"""
+        # Use recent market data
+        recent_prices = self.price_history[-10:] if len(self.price_history) >= 10 else self.price_history
+        recent_sentiments = self.sentiment_history[-10:] if len(self.sentiment_history) >= 10 else self.sentiment_history
         
+        # Features
         features = [
-            self.balance / self.initial_balance,
-            self.position * self.current_price / self.initial_balance,
-            self.current_price / 100.0,
-            np.mean(recent_prices) / 100.0,
-            np.std(recent_prices) / 100.0,
-            np.mean(recent_sentiments),
-            np.std(recent_sentiments),
-            self.step_count / self.max_steps,
-            0.0, 0.0, 0.0, 0.0  # Padding
+            self.balance / self.initial_balance,  # Normalized balance
+            (self.position * self.current_price) / self.initial_balance,  # Position value
+            self.current_price / 200.0,  # Normalized price
+            np.mean(recent_prices) / 200.0,  # Short-term average
+            np.std(recent_prices) / 200.0 if len(recent_prices) > 1 else 0,  # Volatility
+            np.mean(recent_sentiments),  # Average sentiment
+            np.std(recent_sentiments) if len(recent_sentiments) > 1 else 0,  # Sentiment volatility
+            self.step_count / self.max_steps,  # Progress
+            len([a for a in self.action_history[-5:] if a == 1]) / 5.0,  # Recent buy ratio
+            len([a for a in self.action_history[-5:] if a == 2]) / 5.0,  # Recent sell ratio
+            np.mean(np.diff(recent_prices)) if len(recent_prices) > 1 else 0,  # Price momentum
+            0.0  # Padding
         ]
         
         return np.array(features[:12], dtype=np.float32)
-
-    def _get_info(self):
+    
+    def _get_info(self) -> dict:
+        """Get environment info"""
         net_worth = self.balance + self.position * self.current_price
-        return {'net_worth': net_worth}
+        return {
+            'net_worth': net_worth,
+            'balance': self.balance,
+            'position': self.position,
+            'current_price': self.current_price,
+            'step': self.step_count,
+            'pnl': (net_worth - self.initial_balance) / self.initial_balance * 100
+        }
 
 class DQNAgent:
-    def __init__(self, state_dim, action_dim, config, device='cpu'):
+    """Deep Q-Network agent"""
+    
+    def __init__(self, state_dim: int, action_dim: int, config: TradingConfig, device: str = 'cpu'):
+        self.state_dim = state_dim
+        self.action_dim = action_dim
         self.device = torch.device(device)
-        self.q_network = torch.nn.Sequential(
-            torch.nn.Linear(state_dim, 128),
-            torch.nn.ReLU(),
-            torch.nn.Linear(128, 128),
-            torch.nn.ReLU(),
-            torch.nn.Linear(128, action_dim)
-        ).to(self.device)
-        
-        self.target_network = torch.nn.Sequential(
-            torch.nn.Linear(state_dim, 128),
-            torch.nn.ReLU(),
-            torch.nn.Linear(128, 128),
-            torch.nn.ReLU(),
-            torch.nn.Linear(128, action_dim)
-        ).to(self.device)
+        self.config = config
         
+        # Neural networks
+        self.q_network = self._build_network().to(self.device)
+        self.target_network = self._build_network().to(self.device)
         self.target_network.load_state_dict(self.q_network.state_dict())
         
         self.optimizer = torch.optim.Adam(self.q_network.parameters(), lr=config.learning_rate)
         self.memory = deque(maxlen=config.memory_size)
-        self.gamma = config.gamma
         self.epsilon = config.epsilon_start
-        self.epsilon_min = config.epsilon_min
-        self.epsilon_decay = config.epsilon_decay
-        self.batch_size = config.batch_size
-        self.target_update = config.target_update
         self.steps = 0
-
-    def select_action(self, state, training=True):
-        state = torch.FloatTensor(state).unsqueeze(0).to(self.device)
+    
+    def _build_network(self) -> torch.nn.Module:
+        return torch.nn.Sequential(
+            torch.nn.Linear(self.state_dim, 128),
+            torch.nn.ReLU(),
+            torch.nn.Linear(128, 128),
+            torch.nn.ReLU(),
+            torch.nn.Linear(128, self.action_dim)
+        )
+    
+    def select_action(self, state: np.ndarray, training: bool = True) -> int:
+        """Epsilon-greedy action selection"""
         if training and random.random() < self.epsilon:
-            return random.randint(0, 3)
+            return random.randint(0, self.action_dim - 1)
+        
+        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
         with torch.no_grad():
-            return self.q_network(state).argmax(1).item()
-
+            q_values = self.q_network(state_tensor)
+            action = q_values.argmax().item()
+        
+        if training:
+            self.epsilon = max(self.config.epsilon_min, self.epsilon * self.config.epsilon_decay)
+        
+        return action
+    
     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:
+        """Store experience"""
+        self.memory.append((state, action, float(reward), next_state, done))
+    
+    def update(self) -> float:
+        """Train the network"""
+        if len(self.memory) < self.config.batch_size:
             return 0.0
         
-        batch = random.sample(self.memory, self.batch_size)
+        # Sample batch
+        batch = random.sample(self.memory, self.config.batch_size)
         states, actions, rewards, next_states, dones = zip(*batch)
         
         states = torch.FloatTensor(np.array(states)).to(self.device)
@@ -182,98 +261,275 @@ class DQNAgent:
         next_states = torch.FloatTensor(np.array(next_states)).to(self.device)
         dones = torch.FloatTensor(dones).to(self.device)
         
+        # Q-learning
         current_q = self.q_network(states).gather(1, actions.unsqueeze(1)).squeeze(1)
         next_q = self.target_network(next_states).max(1)[0]
-        target_q = rewards + self.gamma * next_q * (1 - dones)
+        target_q = rewards + self.config.gamma * next_q * (1 - dones)
         
         loss = torch.nn.MSELoss()(current_q, target_q)
         
         self.optimizer.zero_grad()
         loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.q_network.parameters(), 1.0)
         self.optimizer.step()
         
         self.steps += 1
-        if self.steps % self.target_update == 0:
-            self.target_network.load_state_dict(self.q_network.state_dict())
         
-        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)
+        # Update target network
+        if self.steps % self.config.target_update == 0:
+            self.target_network.load_state_dict(self.q_network.state_dict())
         
         return loss.item()
 
 class TradingDemo:
+    """Main demo class with proper state management"""
+    
     def __init__(self):
-        self.config = TradingConfig()
+        self.config = None
         self.env = None
         self.agent = None
         self.device = 'cpu'
-
-    def initialize(self, balance, risk, asset):
-        self.config.initial_balance = balance
-        self.config.risk_level = risk
-        self.config.asset_type = asset
-        self.env = AdvancedTradingEnvironment(self.config)
-        self.agent = DQNAgent(12, 4, self.config, self.device)
-        return "✅ Initialized!"
-
-    def train(self, episodes):
-        for ep in range(episodes):
+        self.is_initialized = False
+        self.training_history = []
+    
+    def initialize(self, balance: float) -> str:
+        """Initialize environment and agent"""
+        try:
+            self.config = TradingConfig(initial_balance=float(balance))
+            self.env = AdvancedTradingEnvironment(self.config)
+            state_dim = 12  # Fixed observation size
+            self.agent = DQNAgent(state_dim, 4, self.config, self.device)
+            self.is_initialized = True
+            self.training_history = []
+            
+            logger.info(f"Initialized: balance=${balance}")
+            return f"✅ Environment initialized! Balance: ${balance:,.2f}"
+            
+        except Exception as e:
+            logger.error(f"Initialization failed: {e}")
+            self.is_initialized = False
+            return f"❌ Initialization failed: {str(e)}"
+    
+    def train(self, episodes: int) -> Iterator[Tuple[str, go.Figure]]:
+        """Train the agent with progress updates"""
+        if not self.is_initialized or self.env is None or self.agent is None:
+            yield "❌ Please initialize first!", None
+            return
+        
+        try:
+            episodes = int(episodes)
+            episode_rewards = []
+            
+            for ep in range(episodes):
+                obs, _ = self.env.reset()
+                total_reward = 0
+                done = False
+                
+                while not done:
+                    action = self.agent.select_action(obs, training=True)
+                    next_obs, reward, done, _, info = self.env.step(action)
+                    
+                    self.agent.store_transition(obs, action, reward, next_obs, done)
+                    obs = next_obs
+                    total_reward += reward
+                
+                # Update agent
+                loss = self.agent.update()
+                episode_rewards.append(total_reward)
+                self.training_history.append({
+                    'episode': ep + 1,
+                    'reward': total_reward,
+                    'loss': loss
+                })
+                
+                # Progress update every 10 episodes
+                if (ep + 1) % 10 == 0 or ep == episodes - 1:
+                    avg_reward = np.mean(episode_rewards[-10:])
+                    progress = (ep + 1) / episodes * 100
+                    
+                    status = (f"📈 Training Progress: {ep+1}/{episodes} ({progress:.1f}%)\n"
+                            f"🎯 Episode Reward: {total_reward:.2f}\n"
+                            f"📊 Avg Reward (last 10): {avg_reward:.2f}\n"
+                            f"📉 Loss: {loss:.4f}")
+                    
+                    # Create progress chart
+                    chart = self._create_training_chart()
+                    yield status, chart
+            
+            final_status = f"✅ Training completed! Final Avg Reward: {np.mean(episode_rewards):.2f}"
+            yield final_status, self._create_training_chart()
+            
+        except Exception as e:
+            logger.error(f"Training error: {e}")
+            yield f"❌ Training failed: {str(e)}", None
+    
+    def simulate(self, steps: int) -> Tuple[str, go.Figure, go.Figure]:
+        """Run trading simulation"""
+        if not self.is_initialized or self.env is None or self.agent is None:
+            return "❌ Please initialize and train first!", None, None
+        
+        try:
+            steps = int(steps)
             obs, _ = self.env.reset()
-            total_reward = 0
-            done = False
-            while not done:
-                action = self.agent.select_action(obs)
+            
+            prices = []
+            actions = []
+            net_worths = []
+            rewards = []
+            
+            for step in range(steps):
+                action = self.agent.select_action(obs, training=False)
                 next_obs, reward, done, _, info = self.env.step(action)
-                self.agent.store_transition(obs, action, reward, next_obs, done)
+                
+                prices.append(self.env.current_price)
+                actions.append(action)
+                net_worths.append(info['net_worth'])
+                rewards.append(reward)
+                
                 obs = next_obs
-                total_reward += reward
-            self.agent.update()
-            yield f"Episode {ep+1}/{episodes} | Reward: {total_reward:.2f}", None
-        yield "✅ Training complete!", None
-
-    def simulate(self, steps):
-        obs, _ = self.env.reset()
-        prices = []
-        actions = []
-        net_worths = []
-        for _ in range(steps):
-            action = self.agent.select_action(obs, training=False)
-            next_obs, reward, done, _, info = self.env.step(action)
-            prices.append(self.env.current_price)
-            actions.append(action)
-            net_worths.append(info['net_worth'])
-            obs = next_obs
-            if done:
-                break
-        
-        import plotly.graph_objects as go
+                if done:
+                    break
+            
+            # Create charts
+            price_chart = self._create_price_chart(prices, actions)
+            performance_chart = self._create_performance_chart(net_worths, rewards)
+            
+            final_pnl = (net_worths[-1] - self.config.initial_balance) / self.config.initial_balance * 100
+            status = f"✅ Simulation completed! Steps: {len(prices)}, Final P&L: {final_pnl:+.2f}%"
+            
+            return status, price_chart, performance_chart
+            
+        except Exception as e:
+            logger.error(f"Simulation error: {e}")
+            return f"❌ Simulation failed: {str(e)}", None, None
+    
+    def _create_training_chart(self) -> go.Figure:
+        """Create training progress chart"""
+        if not self.training_history:
+            fig = go.Figure()
+            fig.update_layout(title="Training in progress...", height=400)
+            return fig
+        
+        episodes = [h['episode'] for h in self.training_history]
+        rewards = [h['reward'] for h in self.training_history]
+        losses = [h['loss'] for h in self.training_history]
+        
+        fig = make_subplots(rows=2, cols=1, subplot_titles=["Rewards", "Loss"])
+        
+        fig.add_trace(go.Scatter(x=episodes, y=rewards, mode='lines+markers', name='Reward'), row=1, col=1)
+        fig.add_trace(go.Scatter(x=episodes, y=losses, mode='lines', name='Loss'), row=2, col=1)
+        
+        fig.update_layout(height=500, title="Training Progress", showlegend=True)
+        return fig
+    
+    def _create_price_chart(self, prices: list, actions: list) -> go.Figure:
+        """Create price action chart"""
         fig = go.Figure()
-        fig.add_trace(go.Scatter(y=prices, mode='lines', name='Price'))
-        fig.add_trace(go.Scatter(y=net_worths, mode='lines', name='Net Worth'))
-        return "✅ Simulation complete!", fig
+        
+        # Price line
+        fig.add_trace(go.Scatter(x=list(range(len(prices))), y=prices, mode='lines', 
+                               name='Price', line=dict(color='blue', width=2)))
+        
+        # Action markers
+        buy_indices = [i for i, a in enumerate(actions) if a == 1]
+        sell_indices = [i for i, a in enumerate(actions) if a == 2]
+        
+        if buy_indices:
+            buy_prices = [prices[i] for i in buy_indices]
+            fig.add_trace(go.Scatter(x=buy_indices, y=buy_prices, mode='markers',
+                                   name='Buy', marker=dict(color='green', size=10, symbol='triangle-up')))
+        
+        if sell_indices:
+            sell_prices = [prices[i] for i in sell_indices]
+            fig.add_trace(go.Scatter(x=sell_indices, y=sell_prices, mode='markers',
+                                   name='Sell', marker=dict(color='red', size=10, symbol='triangle-down')))
+        
+        fig.update_layout(title="Price Action Simulation", height=400, showlegend=True)
+        return fig
+    
+    def _create_performance_chart(self, net_worths: list, rewards: list) -> go.Figure:
+        """Create performance dashboard"""
+        fig = make_subplots(rows=2, cols=1, subplot_titles=["Net Worth", "Rewards"])
+        
+        fig.add_trace(go.Scatter(x=list(range(len(net_worths))), y=net_worths, 
+                               mode='lines', name='Net Worth'), row=1, col=1)
+        fig.add_trace(go.Bar(x=list(range(len(rewards))), y=rewards, 
+                           name='Rewards'), row=2, col=1)
+        
+        fig.update_layout(title="Performance", height=500, showlegend=True)
+        return fig
 
+# Initialize demo
 demo = TradingDemo()
 
-with gr.Blocks() as interface:
-    gr.Markdown("# Trading AI Demo")
+# Create Gradio interface
+with gr.Blocks(theme=gr.themes.Soft(), title="🤖 AI Trading Demo") as interface:
+    gr.Markdown("# 🚀 AI Trading Assistant\nReinforcement Learning Trading Demo")
     
+    # Configuration
     with gr.Row():
-        balance = gr.Slider(1000, 50000, 10000, label="Balance")
-        risk = gr.Radio(["Low", "Medium", "High"], value="Medium", label="Risk")
-        asset = gr.Radio(["Crypto", "Stock", "Forex"], value="Crypto", label="Asset")
-        init_btn = gr.Button("Initialize")
+        with gr.Column(scale=1):
+            gr.Markdown("## ⚙️ Configuration")
+            balance = gr.Slider(1000, 50000, value=10000, step=1000, label="Initial Balance ($)")
+            init_btn = gr.Button("🚀 Initialize Environment", variant="primary")
+            status = gr.Textbox(label="Status", interactive=False, lines=2)
+    
+    # Training
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("## 🎓 Training")
+            episodes = gr.Slider(10, 100, value=30, step=5, label="Training Episodes")
+            train_btn = gr.Button("🤖 Start Training", variant="primary")
+        
+        with gr.Column(scale=2):
+            train_status = gr.Textbox(label="Training Progress", interactive=False, lines=3)
+            train_chart = gr.Plot(label="Training Metrics")
+    
+    # Simulation
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("## 🎯 Simulation")
+            sim_steps = gr.Slider(20, 200, value=50, step=10, label="Simulation Steps")
+            sim_btn = gr.Button("▶️ Run Simulation", variant="secondary")
+        
+        with gr.Column(scale=2):
+            sim_status = gr.Textbox(label="Simulation Status", interactive=False)
+            price_chart = gr.Plot(label="Price Chart")
+    
+    performance_chart = gr.Plot(label="Performance Dashboard")
+    
+    # Event handlers with proper error handling
+    init_btn.click(
+        demo.initialize,
+        inputs=[balance],
+        outputs=[status]
+    )
     
-    status = gr.Textbox(label="Status")
+    train_btn.click(
+        demo.train,
+        inputs=[episodes],
+        outputs=[train_status, train_chart]
+    )
     
-    episodes = gr.Number(value=50, label="Episodes")
-    train_btn = gr.Button("Train")
-    train_plot = gr.Plot()
+    sim_btn.click(
+        demo.simulate,
+        inputs=[sim_steps],
+        outputs=[sim_status, price_chart, performance_chart]
+    )
     
-    steps = gr.Number(value=100, label="Simulation Steps")
-    sim_btn = gr.Button("Simulate")
-    sim_plot = gr.Plot()
+    gr.Markdown("""
+    ## 📝 Usage Instructions:
+    1. **Initialize**: Set balance and click "Initialize Environment"
+    2. **Train**: Adjust episodes and click "Start Training"
+    3. **Simulate**: Set steps and click "Run Simulation"
     
-    init_btn.click(demo.initialize, [balance, risk, asset], status)
-    train_btn.click(demo.train, episodes, [status, train_plot])
-    sim_btn.click(demo.simulate, steps, [status, sim_plot])
+    **Note**: Training must complete before simulation!
+    """)
 
-interface.launch()
+if __name__ == "__main__":
+    interface.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        debug=True
+    )