深度强化学习驱动下的股票交易策略：数据清洗、Rainbow DQN实现与回测全流程解析

import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import tushare as ts
import akshare as ak

class DataFetcher:
    """支持A股、港股、美股的数据聚合器"""
    def __init__(self, token=None):
        self.ts_pro = ts.pro_api(token) if token else None
        
    def fetch_daily(self, symbol, start, end, market='A'):
        if market == 'A':
            df = self.ts_pro.daily(ts_code=symbol, start_date=start, end_date=end)
            df = df.rename(columns={'trade_date': 'date', 'vol': 'volume'})
        elif market == 'US':
            df = ak.stock_us_hist(symbol=symbol, start_date=start, end_date=end, adjust='qfq')
        df['date'] = pd.to_datetime(df['date'])
        df = df.sort_values('date').set_index('date')
        # 统一列名：open, high, low, close, volume, amount
        return df[['open','high','low','close','volume','amount']]

from scipy.stats import mad

def winsorize_series(s, limits=(0.01, 0.99)):
    q_low, q_high = s.quantile(limits[0]), s.quantile(limits[1])
    return s.clip(q_low, q_high)

def clean_ohlcv(df, method='mad', threshold=5):
    for col in ['open','high','low','close']:
        if method == 'mad':
            med = df[col].median()
            mad_val = mad(df[col], nan_policy='omit')
            upper = med + threshold * mad_val
            lower = med - threshold * mad_val
            df[col] = df[col].clip(lower, upper)
        else:
            df[col] = winsorize_series(df[col])
    # 确保 high >= low, close 在合理范围
    df['high'] = df[['high','low']].max(axis=1)
    df['low'] = df[['high','low']].min(axis=1)
    df['close'] = df['close'].clip(df['low'], df['high'])
    return df

import talib

def add_technical_features(df):
    close = df['close'].values
    high = df['high'].values
    low = df['low'].values
    volume = df['volume'].values
    
    df['rsi_14'] = talib.RSI(close, timeperiod=14)
    df['macd'], df['macd_signal'], df['macd_hist'] = talib.MACD(close)
    df['bb_upper'], df['bb_middle'], df['bb_lower'] = talib.BBANDS(close)
    df['atr'] = talib.ATR(high, low, close, timeperiod=14)
    df['adx'] = talib.ADX(high, low, close, timeperiod=14)
    # 更多因子...
    
    # 标准化（滚动窗口，避免未来信息）
    lookback = 60
    for col in df.select_dtypes(include=[np.number]).columns:
        roll_mean = df[col].rolling(lookback).mean()
        roll_std = df[col].rolling(lookback).std()
        df[col] = (df[col] - roll_mean) / (roll_std + 1e-8)
    return df.fillna(0)

import gym
from gym import spaces

class TradingEnv(gym.Env):
    def __init__(self, df, feature_cols, history_len=30, initial_cash=1e6,
                 commission=0.0003, slippage=0.0001):
        super().__init__()
        self.df = df.values  # numpy for speed
        self.feature_cols = feature_cols
        self.history_len = history_len
        self.initial_cash = initial_cash
        self.commission = commission
        self.slippage = slippage
        
        self.action_space = spaces.Discrete(7)  # 0~6
        # 观测空间： (history_len, feature_dim)
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf,
                                            shape=(history_len, len(feature_cols)), dtype=np.float32)
        
    def reset(self):
        self.cash = self.initial_cash
        self.position = 0  # 股数
        self.step_index = self.history_len  # 从第history_len天开始
        self.done = False
        self.trade_log = []
        return self._get_obs()
    
    def _get_obs(self):
        start = self.step_index - self.history_len
        end = self.step_index
        return self.df[start:end, :][self.feature_cols].astype(np.float32)
    
    def step(self, action):
        current_price = self.df[self.step_index, self.df.columns.get_loc('close')]
        # 动作映射为仓位比例变化
        target_ratio = {0: 0, 1: 0.33, 2: 0.66, 3: 1.0, 4: -0.33, 5: -0.66, 6: -1.0}[action]
        current_value = self.cash + self.position * current_price
        target_position_value = current_value * (self.position * current_price / current_value + target_ratio)
        target_position = max(0, target_position_value // current_price)  # 只做多
        trade_volume = target_position - self.position
        
        # 计算成交价格（含滑点）
        exec_price = current_price * (1 + self.slippage if trade_volume > 0 else 1 - self.slippage)
        cost = abs(trade_volume) * exec_price * self.commission
        self.cash -= trade_volume * exec_price + cost
        self.position = target_position
        
        # 更新组合价值
        new_value = self.cash + self.position * current_price
        reward = (new_value - current_value) / current_value  # 简单收益率
        
        # 额外惩罚：回撤惩罚
        if hasattr(self, 'peak_value'):
            drawdown = (self.peak_value - new_value) / self.peak_value
            reward -= 0.5 * max(0, drawdown - 0.05)  # 超过5%回撤惩罚
        else:
            self.peak_value = new_value
        self.peak_value = max(self.peak_value, new_value)
        
        self.step_index += 1
        if self.step_index >= len(self.df) - 1:
            self.done = True
        return self._get_obs(), reward, self.done, {'value': new_value, 'position': self.position}

import torch
import torch.nn as nn
import torch.nn.functional as F

class NoisyLinear(nn.Module):
    def __init__(self, in_features, out_features, sigma_init=0.5):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight_mu = nn.Parameter(torch.empty(out_features, in_features))
        self.weight_sigma = nn.Parameter(torch.empty(out_features, in_features))
        self.bias_mu = nn.Parameter(torch.empty(out_features))
        self.bias_sigma = nn.Parameter(torch.empty(out_features))
        self.register_buffer('weight_epsilon', torch.empty(out_features, in_features))
        self.register_buffer('bias_epsilon', torch.empty(out_features))
        self.reset_parameters()
        self.sigma_init = sigma_init
        
    def reset_parameters(self):
        mu_range = 1 / self.in_features ** 0.5
        self.weight_mu.data.uniform_(-mu_range, mu_range)
        self.weight_sigma.data.fill_(self.sigma_init / self.in_features ** 0.5)
        self.bias_mu.data.uniform_(-mu_range, mu_range)
        self.bias_sigma.data.fill_(self.sigma_init / self.in_features ** 0.5)
    
    def forward(self, x):
        self.weight_epsilon.normal_()
        self.bias_epsilon.normal_()
        weight = self.weight_mu + self.weight_sigma * self.weight_epsilon
        bias = self.bias_mu + self.bias_sigma * self.bias_epsilon
        return F.linear(x, weight, bias)

class LSTM_Dueling_DQN(nn.Module):
    def __init__(self, feature_dim, history_len, n_actions, lstm_hidden=128):
        super().__init__()
        self.lstm = nn.LSTM(input_size=feature_dim, hidden_size=lstm_hidden,
                            num_layers=2, batch_first=True, dropout=0.2)
        # Dueling: value stream + advantage stream
        self.value_stream = nn.Sequential(
            NoisyLinear(lstm_hidden, 128),
            nn.ReLU(),
            NoisyLinear(128, 1)
        )
        self.advantage_stream = nn.Sequential(
            NoisyLinear(lstm_hidden, 128),
            nn.ReLU(),
            NoisyLinear(128, n_actions)
        )
        
    def forward(self, x):
        # x shape: (batch, history_len, feature_dim)
        lstm_out, (hn, cn) = self.lstm(x)
        # 取最后时刻的hidden state
        h = lstm_out[:, -1, :]  # (batch, lstm_hidden)
        value = self.value_stream(h)  # (batch, 1)
        advantage = self.advantage_stream(h)  # (batch, n_actions)
        # Q = V + (A - mean(A))
        q = value + advantage - advantage.mean(dim=1, keepdim=True)
        return q

class SumTree:
    def __init__(self, capacity):
        self.capacity = capacity
        self.tree = np.zeros(2 * capacity - 1)
        self.data = np.zeros(capacity, dtype=object)
        self.size = 0
        self.ptr = 0
        
    def add(self, priority, data):
        idx = self.ptr + self.capacity - 1
        self.data[self.ptr] = data
        self.update(idx, priority)
        self.ptr = (self.ptr + 1) % self.capacity
        self.size = min(self.size + 1, self.capacity)
    
    def update(self, idx, priority):
        delta = priority - self.tree[idx]
        self.tree[idx] = priority
        while idx != 0:
            idx = (idx - 1) // 2
            self.tree[idx] += delta
    
    def get_leaf(self, v):
        idx = 0
        while idx < self.capacity - 1:
            left = 2 * idx + 1
            right = left + 1
            if v <= self.tree[left]:
                idx = left
            else:
                v -= self.tree[left]
                idx = right
        data_idx = idx - self.capacity + 1
        return idx, self.tree[idx], self.data[data_idx]

class RainbowDQNAgent:
    def __init__(self, env, lr=1e-4, gamma=0.99, batch_size=64, 
                 replay_capacity=100000, per_alpha=0.6, per_beta=0.4):
        self.env = env
        self.gamma = gamma
        self.batch_size = batch_size
        self.per_alpha = per_alpha
        self.per_beta = per_beta
        self.n_actions = env.action_space.n
        self.feature_dim = env.observation_space.shape[1]
        self.history_len = env.observation_space.shape[0]
        
        self.q_net = LSTM_Dueling_DQN(self.feature_dim, self.history_len, self.n_actions)
        self.target_net = LSTM_Dueling_DQN(self.feature_dim, self.history_len, self.n_actions)
        self.target_net.load_state_dict(self.q_net.state_dict())
        self.optimizer = torch.optim.Adam(self.q_net.parameters(), lr=lr)
        self.memory = SumTree(replay_capacity)
        self.step_count = 0
        
    def select_action(self, state, epsilon=0.1):
        if np.random.rand() < epsilon:
            return np.random.randint(self.n_actions)
        state_t = torch.FloatTensor(state).unsqueeze(0)  # (1, H, F)
        with torch.no_grad():
            q_values = self.q_net(state_t)
        return q_values.argmax().item()
    
    def store_transition(self, state, action, reward, next_state, done):
        # 计算初始优先级（用当前网络估计的TD误差，但为了效率，先给最大优先级）
        priority = 1.0  # 后续在训练时更新
        self.memory.add(priority, (state, action, reward, next_state, done))
    
    def learn(self):
        if self.memory.size < self.batch_size:
            return
        # 采样
        idxs, priorities, batch_data = [], [], []
        segment = self.memory.total() / self.batch_size
        for i in range(self.batch_size):
            v = np.random.uniform(segment*i, segment*(i+1))
            idx, p, data = self.memory.get_leaf(v)
            idxs.append(idx)
            priorities.append(p)
            batch_data.append(data)
        # 解包
        states, actions, rewards, next_states, dones = map(np.array, zip(*batch_data))
        states = torch.FloatTensor(states)
        next_states = torch.FloatTensor(next_states)
        actions = torch.LongTensor(actions).unsqueeze(1)
        rewards = torch.FloatTensor(rewards).unsqueeze(1)
        dones = torch.BoolTensor(dones).unsqueeze(1)
        
        # 计算TD目标（Double DQN）
        with torch.no_grad():
            q_next = self.q_net(next_states)
            best_actions = q_next.argmax(dim=1, keepdim=True)
            q_target_next = self.target_net(next_states).gather(1, best_actions)
            target = rewards + self.gamma * q_target_next * (~dones)
        q_current = self.q_net(states).gather(1, actions)
        td_error = (target - q_current).abs().squeeze()
        
        # 更新优先级
        for i, idx in enumerate(idxs):
            priority = (td_error[i].item() + 1e-5) ** self.per_alpha
            self.memory.update(idx, priority)
        
        # Huber loss
        loss = F.smooth_l1_loss(q_current, target)
        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.q_net.parameters(), 10)
        self.optimizer.step()
        
        # 软更新target网络
        if self.step_count % 100 == 0:
            self.target_net.load_state_dict(self.q_net.state_dict())
        self.step_count += 1

def backtest(env, agent, n_episodes=1, render=False):
    total_rewards = []
    portfolio_values = []
    for ep in range(n_episodes):
        state = env.reset()
        done = False
        ep_vals = [env.initial_cash]
        while not done:
            action = agent.select_action(state, epsilon=0.0)  # greedy
            next_state, reward, done, info = env.step(action)
            state = next_state
            ep_vals.append(info['value'])
        total_rewards.append(np.sum(ep_vals) - env.initial_cash)
        portfolio_values.append(ep_vals)
    return portfolio_values, total_rewards

def evaluate_metrics(values, risk_free=0.03, periods_per_year=252):
    values = np.array(values)
    returns = (values[1:] / values[:-1] - 1)
    ann_return = returns.mean() * periods_per_year
    ann_vol = returns.std() * np.sqrt(periods_per_year)
    sharpe = (ann_return - risk_free) / ann_vol if ann_vol > 0 else 0
    cumulative = values / values[0]
    drawdown = 1 - cumulative / cumulative.cummax()
    max_drawdown = drawdown.max()
    calmar = ann_return / max_drawdown if max_drawdown > 0 else 0
    return {
        'annual_return': ann_return,
        'annual_volatility': ann_vol,
        'sharpe_ratio': sharpe,
        'max_drawdown': max_drawdown,
        'calmar_ratio': calmar
    }

if __name__ == '__main__':
    # 1. 数据获取
    fetcher = DataFetcher(token='你的token')
    df = fetcher.fetch_daily('000300.SH', '2015-01-01', '2024-12-31', market='A')
    df = clean_ohlcv(df)
    df = add_technical_features(df)
    
    # 2. 划分训练/验证/测试（时序切分）
    train_df = df.loc[:'2020-12-31']
    val_df = df.loc['2021-01-01':'2022-12-31']
    test_df = df.loc['2023-01-01':]
    
    feature_cols = [c for c in df.columns if c not in ['open','high','low','close','volume','amount']]
    # 保证特征列存在
    train_env = TradingEnv(train_df, feature_cols, history_len=30)
    test_env = TradingEnv(test_df, feature_cols, history_len=30)
    
    # 3. 训练智能体
    agent = RainbowDQNAgent(train_env, lr=1e-4, gamma=0.99)
    n_epochs = 50
    for epoch in range(n_epochs):
        state = train_env.reset()
        done = False
        while not done:
            epsilon = max(0.01, 0.3 - epoch * 0.005)  # 线性衰减
            action = agent.select_action(state, epsilon)
            next_state, reward, done, _ = train_env.step(action)
            agent.store_transition(state, action, reward, next_state, done)
            agent.learn()
            state = next_state
        # 每epoch结束在验证集评估
        if epoch % 5 == 0:
            val_values, _ = backtest(TradingEnv(val_df, feature_cols), agent, n_episodes=1)
            metrics = evaluate_metrics(val_values[0])
            print(f"Epoch {epoch}, Val Sharpe: {metrics['sharpe_ratio']:.3f}")
    
    # 4. 测试集最终评估
    test_values, _ = backtest(test_env, agent, n_episodes=1)
    test_metrics = evaluate_metrics(test_values[0])
    print("Test Metrics:", test_metrics)
    
    # 5. 保存模型
    torch.save(agent.q_net.state_dict(), 'rainbow_dqn_stock.pth')

from torch.distributions import Categorical

class PPOBuffer:
    def __init__(self, gamma, lam):
        self.gamma, self.lam = gamma, lam
        self.states, self.actions, self.rewards = [], [], []
        self.dones, self.logprobs, self.values = [], [], []
    
    def add(self, state, action, reward, done, logprob, value):
        self.states.append(state); self.actions.append(action); self.rewards.append(reward)
        self.dones.append(done); self.logprobs.append(logprob); self.values.append(value)
    
    def finish_trajectory(self, last_value=0):
        # GAE计算
        advantages = np.zeros(len(self.rewards), dtype=np.float32)
        last_advantage = 0
        last_value = last_value
        for t in reversed(range(len(self.rewards))):
            if t == len(self.rewards)-1:
                next_non_terminal = 1.0 - self.dones[-1]
                next_value = last_value
            else:
                next_non_terminal = 1.0 - self.dones[t+1]
                next_value = self.values[t+1]
            delta = self.rewards[t] + self.gamma * next_value * next_non_terminal - self.values[t]
            last_advantage = delta + self.gamma * self.lam * next_non_terminal * last_advantage
            advantages[t] = last_advantage
        returns = advantages + np.array(self.values)
        return advantages, returns

# PPO更新类似，但篇幅所限，此处仅展示GAE计算