biggfish/src/ml/rl_agent.py.bak

"""
Reinforcement Learning Agent for Trading
DQN with experience replay and target network, implemented in PyTorch.
"""

import io
import random
import numpy as np
from collections import deque
from typing import Dict, List, Optional, Tuple
from loguru import logger

try:
    import torch
    import torch.nn as nn
    import torch.optim as optim
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    logger.warning("PyTorch not installed. RL agent will use random actions. "
                   "Install with: pip install torch")


class TradingNetwork:
    """Neural network for the RL agent (PyTorch or fallback)"""
    pass


if TORCH_AVAILABLE:
    class TradingNetwork(nn.Module):
        """MLP with 2 hidden layers for Q-value prediction"""

        def __init__(self, state_dim: int, action_dim: int, hidden_dim: int = 128):
            super().__init__()
            self.net = nn.Sequential(
                nn.Linear(state_dim, hidden_dim),
                nn.ReLU(),
                nn.Dropout(0.1),
                nn.Linear(hidden_dim, hidden_dim),
                nn.ReLU(),
                nn.Dropout(0.1),
                nn.Linear(hidden_dim, action_dim)
            )

        def forward(self, x):
            return self.net(x)


class RLAgent:
    """
    DQN-based RL agent for trading decisions.
    Falls back to random actions if PyTorch is not available.
    """

    def __init__(self, state_dim: int, action_dim: int = 7, config: Dict = None):
        config = config or {}
        self.config = config  # CRITICAL FIX: Save config for memory persistence
        self.state_dim = state_dim
        self.action_dim = action_dim

        # Hyperparameters
        self.gamma = config.get('gamma', 0.99)
        self.epsilon = config.get('epsilon_start', 1.0)
        self.epsilon_min = config.get('epsilon_min', 0.05)
        self.epsilon_decay = config.get('epsilon_decay', 0.9995)
        self.learning_rate = config.get('learning_rate', 0.0003)
        self.batch_size = config.get('batch_size', 64)
        self.memory_size = config.get('memory_size', 50000)
        self.target_update_freq = config.get('target_update_freq', 100)
        self.live_epsilon = config.get('live_epsilon', 0.1)

        # Experience replay buffer
        self.memory = deque(maxlen=self.memory_size)
        self.steps = 0
        self.training_losses = []

        # PyTorch setup
        self.use_torch = TORCH_AVAILABLE
        if self.use_torch:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            hidden_dim = config.get('hidden_dim', 128)
            self.policy_net = TradingNetwork(state_dim, action_dim, hidden_dim).to(self.device)
            self.target_net = TradingNetwork(state_dim, action_dim, hidden_dim).to(self.device)
            self.target_net.load_state_dict(self.policy_net.state_dict())
            self.target_net.eval()
            self.optimizer = optim.Adam(self.policy_net.parameters(), lr=self.learning_rate)
            logger.info(f"RL Agent initialized (PyTorch, device={self.device}, "
                       f"state_dim={state_dim}, action_dim={action_dim})")
        else:
            self.device = None
            self.policy_net = None
            self.target_net = None
            self.optimizer = None
            logger.info("RL Agent initialized (random mode - no PyTorch)")

    def select_action(self, state: np.ndarray, live_mode: bool = False) -> int:
        """
        Epsilon-greedy action selection.
        In live mode, uses live_epsilon instead of training epsilon.
        """
        eps = self.live_epsilon if live_mode else self.epsilon

        if random.random() < eps:
            return random.randint(0, self.action_dim - 1)

        if not self.use_torch:
            return random.randint(0, self.action_dim - 1)

        with torch.no_grad():
            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
            q_values = self.policy_net(state_tensor)
            return int(q_values.argmax(dim=1).item())

    def store_experience(self, state, action, reward, next_state, done):
        """Store transition in replay buffer"""
        self.memory.append((state, action, reward, next_state, done))

    def train_step(self) -> Optional[float]:
        """
        Sample mini-batch from replay buffer, compute DQN loss, update.
        Returns loss value or None if not enough samples.
        """
        if not self.use_torch:
            return None

        if len(self.memory) < self.batch_size:
            return None

        batch = random.sample(self.memory, self.batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)

        states = torch.FloatTensor(np.array(states)).to(self.device)
        actions = torch.LongTensor(actions).to(self.device)
        rewards = torch.FloatTensor(rewards).to(self.device)
        next_states = torch.FloatTensor(np.array(next_states)).to(self.device)
        dones = torch.BoolTensor(dones).to(self.device)

        # Current Q values
        current_q = self.policy_net(states).gather(1, actions.unsqueeze(1)).squeeze(1)

        # Target Q values
        with torch.no_grad():
            next_q = self.target_net(next_states).max(1)[0]
            next_q[dones] = 0.0
            target_q = rewards + self.gamma * next_q

        # Loss and backprop
        loss = nn.functional.smooth_l1_loss(current_q, target_q)
        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.policy_net.parameters(), 1.0)
        self.optimizer.step()

        self.steps += 1

        # Update target network
        if self.steps % self.target_update_freq == 0:
            self.target_net.load_state_dict(self.policy_net.state_dict())

        # Decay epsilon
        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)

        loss_val = loss.item()
        self.training_losses.append(loss_val)

        return loss_val

    def train_on_episode(self, env, candles_df, ga_signal_fn=None) -> Dict:
        """
        Train on a full episode (backtest run through candles).
        Returns training metrics.
        """
        state = env.reset(candles_df)
        if state is None:
            return {'avg_loss': 0, 'total_reward': 0, 'steps': 0}

        total_reward = 0
        total_loss = 0
        loss_count = 0
        step = 0
        done = False

        while not done:
            # Get GA signal if available
            ga_signal = None
            if ga_signal_fn and step < len(candles_df):
                ga_signal = ga_signal_fn(step)

            action = self.select_action(state)
            next_state, reward, done, info = env.step(action, ga_signal=ga_signal)

            self.store_experience(state, action, reward, next_state, done)

            loss = self.train_step()
            if loss is not None:
                total_loss += loss
                loss_count += 1

            total_reward += reward
            state = next_state
            step += 1

        avg_loss = total_loss / max(loss_count, 1)

        return {
            'avg_loss': round(avg_loss, 6),
            'total_reward': round(total_reward, 4),
            'steps': step,
            'epsilon': round(self.epsilon, 4),
            'total_trades': env.total_trades,
            'total_pnl': round(env.total_pnl, 4),
            'final_equity': round(env.equity_history[-1] if env.equity_history else 0, 2),
            'memory_size': len(self.memory),
        }

    def update_from_live_trade(self, state, action, reward, next_state):
        """
        Online learning: called after each live trade result.
        Stores experience and does one training step.
        """
        self.store_experience(state, action, reward, next_state, False)
        self.train_step()

    def save(self, store, epoch: int = None):
        """Save model checkpoint to DataStore"""
        if not self.use_torch:
            return

        if epoch is None:
            epoch = self.steps

        state_bytes = self._get_state_dict_bytes()
        metrics = {
            'epsilon': self.epsilon,
            'steps': self.steps,
            'memory_size': len(self.memory),
            'avg_loss': round(np.mean(self.training_losses[-100:]), 6)
                        if self.training_losses else 0,
        }
        store.save_model_checkpoint('rl_agent', epoch, state_bytes, metrics)
        
        # Also save experience replay memory
        import pickle
        from pathlib import Path
        db_path_str = str(store.db_path)
        memory_path = db_path_str.replace('.db', '_rl_memory.pkl')
        try:
            with open(memory_path, 'wb') as f:
                # Save memory as list to avoid deque pickle issues
                pickle.dump(list(self.memory), f)
            logger.debug(f"RL memory saved ({len(self.memory)} experiences)")
        except Exception as e:
            logger.warning(f"Failed to save RL memory: {e}")
        
        logger.info(f"RL model saved (epoch {epoch}, epsilon={self.epsilon:.4f})")

    def load(self, store) -> bool:
        """Load latest checkpoint from DataStore"""
        if not self.use_torch:
            return False

        checkpoint = store.load_latest_checkpoint('rl_agent')
        if checkpoint is None:
            logger.info("No RL checkpoint found, starting fresh")
            return False

        try:
            buffer = io.BytesIO(checkpoint['state_dict'])
            state_dict = torch.load(buffer, map_location=self.device, weights_only=True)

            # Check dimension compatibility before loading
            first_layer_key = 'net.0.weight'
            if first_layer_key in state_dict:
                saved_input_dim = state_dict[first_layer_key].shape[1]
                if saved_input_dim != self.state_dim:
                    logger.warning(f"RL checkpoint dimension mismatch (saved={saved_input_dim}, "
                                 f"current={self.state_dim}). Starting fresh with new architecture.")
                    return False

            self.policy_net.load_state_dict(state_dict)
            self.target_net.load_state_dict(state_dict)

            import json
            metrics = json.loads(checkpoint.get('metrics', '{}'))
            self.epsilon = metrics.get('epsilon', self.epsilon)
            self.steps = metrics.get('steps', self.steps)

            # Load experience replay memory
            import pickle
            from collections import deque
            from pathlib import Path
            db_path_str = str(store.db_path)
            memory_path = db_path_str.replace('.db', '_rl_memory.pkl')
            try:
                with open(memory_path, 'rb') as f:
                    saved_memory = pickle.load(f)
                    self.memory = deque(saved_memory, maxlen=self.config['memory_size'])
                logger.info(f"RL memory loaded ({len(self.memory)} experiences)")
            except FileNotFoundError:
                logger.debug("No saved RL memory found, starting with empty buffer")
            except Exception as e:
                logger.warning(f"Failed to load RL memory: {e}")

            logger.info(f"RL model loaded (epoch {checkpoint['epoch']}, "
                       f"epsilon={self.epsilon:.4f})")
            return True

        except Exception as e:
            logger.warning(f"Could not load RL checkpoint (likely dimension change): {e}. Starting fresh.")
            return False

    def _get_state_dict_bytes(self) -> bytes:
        """Serialize model state dict to bytes"""
        buffer = io.BytesIO()
        torch.save(self.policy_net.state_dict(), buffer)
        return buffer.getvalue()

    def get_stats(self) -> Dict:
        """Get current agent statistics"""
        return {
            'epsilon': round(self.epsilon, 4),
            'steps': self.steps,
            'memory_size': len(self.memory),
            'avg_loss': round(np.mean(self.training_losses[-100:]), 6)
                        if self.training_losses else 0,
            'device': str(self.device) if self.device else 'random',
            'use_torch': self.use_torch,
        }