DualFlow/models/Dynamic_anomaly_diagnosis/rl_tracing.py

# -*- coding: utf-8 -*-
"""rl_tracing.py: 强化学习链路级异常溯源"""
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Categorical
import numpy as np
import pandas as pd
import os
from tqdm import tqdm
from config import config

# ----------------- 1. 环境 -----------------
class CausalTracingEnv:
    def __init__(self, causal_graph, window_scores, threshold_df, expert_knowledge=None):
        self.sensor_list = causal_graph['sensor_list']
        self.map = causal_graph['sensor_to_idx']
        self.idx_to_sensor = {v: k for k, v in self.map.items()}
        self.adj = causal_graph['adj_matrix']
        self.scores = window_scores
        
        self.expert_knowledge = expert_knowledge if expert_knowledge else {}
        self.num_sensors = len(self.sensor_list)
        
        # 解析属性
        self.node_props = {} 
        col_one_layer = self._find_col(threshold_df, config.KEYWORD_LAYER)
        col_device = self._find_col(threshold_df, config.KEYWORD_DEVICE)
        
        df_indexed = threshold_df.set_index('ID')
        dict_one = df_indexed[col_one_layer].to_dict() if col_one_layer else {}
        dict_dev = df_indexed[col_device].to_dict() if col_device else {}
        
        for name, idx in self.map.items():
            l_val = dict_one.get(name, -1)
            try: l_val = int(l_val)
            except: l_val = 0 
            d_val = dict_dev.get(name, None)
            d_val = str(d_val).strip() if pd.notna(d_val) and str(d_val).strip() != '' else None
            self.node_props[idx] = {'one_layer': l_val, 'device': d_val}

        self.current_window_idx = 0
        self.current_node_idx = 0
        self.prev_node_idx = 0
        self.trigger_node_idx = 0
        self.path = []
        self.current_expert_paths = []
        self.target_roots = set()
        
    def _find_col(self, df, keyword):
        if keyword in df.columns: return keyword
        for c in df.columns:
            if c.lower() == keyword.lower(): return c
        return None

    def reset(self, force_window_idx=None, force_trigger=None):
        if force_window_idx is not None:
            self.current_window_idx = force_window_idx
            t_name = force_trigger
        else:
            found = False
            for _ in range(100):
                w_idx = np.random.randint(len(self.scores))
                win_scores = self.scores[w_idx]
                candidates = []
                for t_name in config.TRIGGER_SENSORS:
                    if t_name in self.map:
                        idx = self.map[t_name]
                        if win_scores[idx] > config.TRIGGER_SCORE_THRESH:
                            candidates.append(t_name)
                if candidates:
                    self.current_window_idx = w_idx
                    t_name = np.random.choice(candidates)
                    found = True
                    break
            if not found:
                self.current_window_idx = np.random.randint(len(self.scores))
                t_name = list(self.map.keys())[0]

        self.current_node_idx = self.map.get(t_name, 0)
        self.trigger_node_idx = self.current_node_idx
        self.prev_node_idx = self.current_node_idx
        self.path = [self.current_node_idx]
        
        self.target_roots = set()
        self.current_expert_paths = []
        if self.current_node_idx in self.expert_knowledge:
            entry = self.expert_knowledge[self.current_node_idx]
            self.target_roots = entry['roots']
            self.current_expert_paths = entry['paths']
            
        return self._get_state()
    
    def _get_state(self):
        curr_score = self.scores[self.current_window_idx, self.current_node_idx]
        prev_score = self.scores[self.current_window_idx, self.prev_node_idx]
        curr_layer = self.node_props[self.current_node_idx]['one_layer'] / 20.0
        
        return (
            torch.LongTensor([self.current_node_idx, self.prev_node_idx, self.trigger_node_idx]), 
            torch.FloatTensor([curr_score, prev_score, curr_layer])
        )
    
    def get_valid_actions(self, curr_idx):
        neighbors = np.where(self.adj[curr_idx] == 1)[0]
        curr_props = self.node_props[curr_idx]
        curr_l, curr_d = curr_props['one_layer'], curr_props['device']
        valid = []
        for n in neighbors:
            if n in self.path: continue 
            tgt_props = self.node_props[n]
            tgt_l, tgt_d = tgt_props['one_layer'], tgt_props['device']
            if curr_l != 0 and tgt_l != 0:
                if not ((tgt_l == curr_l) or (tgt_l == curr_l - 1)): continue
            if (curr_d is not None) and (tgt_d is not None):
                if curr_d != tgt_d: continue
            valid.append(n)
        return np.array(valid)
    
    def step(self, action_idx):
        prev = self.current_node_idx
        self.prev_node_idx = prev
        self.current_node_idx = action_idx
        self.path.append(action_idx)
        
        score_curr = self.scores[self.current_window_idx, self.current_node_idx]
        
        reward = 0.0
        done = False
        
        # 奖励机制 (Imitation > Root > Gradient)
        in_expert_nodes = False
        for e_path in self.current_expert_paths:
            if action_idx in e_path:
                in_expert_nodes = True
                break
        
        if in_expert_nodes: reward += 2.0
        else: reward -= 0.2
            
        if action_idx in self.target_roots:
            reward += 10.0
            done = True
        
        score_prev = self.scores[self.current_window_idx, prev]
        diff = score_curr - score_prev
        if diff > 0: reward += diff * 3.0
        else: reward -= 0.5
            
        if len(self.path) >= config.MAX_PATH_LENGTH:
            done = True
            if action_idx not in self.target_roots: reward -= 5.0
            
        if score_curr < 0.15 and len(self.path) > 3:
            done = True
            reward -= 2.0
            
        return self._get_state(), reward, done, {}

# ----------------- 2. 网络 -----------------
class TargetDrivenActorCritic(nn.Module):
    def __init__(self, num_sensors, embedding_dim=64, hidden_dim=256):
        super().__init__()
        self.node_emb = nn.Embedding(num_sensors, embedding_dim)
        input_dim = (embedding_dim * 3) + 3 
        
        self.shared_net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU()
        )
        self.actor = nn.Linear(hidden_dim, num_sensors)
        self.critic = nn.Linear(hidden_dim, 1)
        
    def forward(self, int_data, float_data):
        curr_emb = self.node_emb(int_data[:, 0])
        prev_emb = self.node_emb(int_data[:, 1])
        trig_emb = self.node_emb(int_data[:, 2]) 
        x = torch.cat([curr_emb, prev_emb, trig_emb, float_data], dim=1)
        feat = self.shared_net(x)
        return self.actor(feat), self.critic(feat)

# ----------------- 3. 训练器 -----------------
class RLTrainer:
    def __init__(self, causal_graph, train_scores, threshold_df):
        self.sensor_map = causal_graph['sensor_to_idx']
        self.idx_to_sensor = {v: k for k, v in self.sensor_map.items()}
        self.threshold_df = threshold_df
        self.causal_graph = causal_graph
        self.expert_knowledge, self.bc_samples, _ = self._load_expert_data()
        
        self.env = CausalTracingEnv(causal_graph, train_scores, threshold_df, self.expert_knowledge)
        self.model = TargetDrivenActorCritic(self.env.num_sensors, config.EMBEDDING_DIM, config.HIDDEN_DIM)
        self.optimizer = optim.Adam(self.model.parameters(), lr=config.PPO_LR)
        
    def _load_expert_data(self):
        path = os.path.join(config.BASE_DIR, config.ABNORMAL_LINK_FILENAME)
        kb_data = {} 
        bc_data = [] 
        if not os.path.exists(path): return kb_data, bc_data, None
        
        df = pd.read_excel(path)
        for _, row in df.iterrows():
            link = str(row.get('Link Path', ''))
            if not link: continue
            nodes_str = [n.strip() for n in link.replace('→', '->').split('->')]
            path_nodes = nodes_str[::-1]
            
            ids = []
            valid = True
            for n in path_nodes:
                if n in self.sensor_map: ids.append(self.sensor_map[n])
                else: valid = False; break
            if not valid or len(ids)<2: continue
            
            trigger_id = ids[0]
            root_id = ids[-1]
            
            if trigger_id not in kb_data:
                kb_data[trigger_id] = {'paths': [], 'roots': set(), 'logic': row.get('Process Logic Basis', '')}
            kb_data[trigger_id]['paths'].append(ids)
            kb_data[trigger_id]['roots'].add(root_id)
            
            for i in range(len(ids) - 1):
                curr = ids[i]
                prev = ids[max(0, i-1)]
                nxt = ids[i+1]
                bc_data.append(((curr, prev, trigger_id), nxt))
                
        return kb_data, bc_data, df

    def pretrain_bc(self):
        if not self.bc_samples: return
        print(f"\n>>> [Step 3.1] 启动BC预训练 ({config.BC_EPOCHS}轮)...")
        states_int = torch.LongTensor([list(s) for s, a in self.bc_samples])
        actions = torch.LongTensor([a for s, a in self.bc_samples])
        states_float = torch.zeros((len(states_int), 3))
        states_float[:, 0] = 0.9 
        states_float[:, 1] = 0.8
        
        loss_fn = nn.CrossEntropyLoss()
        pbar = tqdm(range(config.BC_EPOCHS), desc="BC Training")
        for epoch in pbar:
            logits, _ = self.model(states_int, states_float)
            loss = loss_fn(logits, actions)
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            if epoch%100==0: pbar.set_postfix({'Loss': f"{loss.item():.4f}"})

    def train_ppo(self):
        print(f"\n>>> [Step 3.2] 启动PPO训练 ({config.RL_EPISODES}轮)...")
        pbar = tqdm(range(config.RL_EPISODES), desc="PPO Training")
        rewards_hist = []
        for _ in pbar:
            state_data = self.env.reset()
            done = False
            ep_r = 0
            b_int, b_float, b_act, b_lp, b_rew, b_mask = [], [], [], [], [], []
            
            while not done:
                s_int = state_data[0].unsqueeze(0)
                s_float = state_data[1].unsqueeze(0)
                valid = self.env.get_valid_actions(s_int[0, 0].item())
                if len(valid) == 0: break
                
                logits, _ = self.model(s_int, s_float)
                mask = torch.full_like(logits, -1e9)
                mask[0, valid] = 0
                dist = Categorical(F.softmax(logits+mask, dim=-1))
                action = dist.sample()
                
                next_s, r, done, _ = self.env.step(action.item())
                b_int.append(s_int); b_float.append(s_float)
                b_act.append(action); b_lp.append(dist.log_prob(action))
                b_rew.append(r); b_mask.append(1-done)
                state_data = next_s
                ep_r += r
            
            if len(b_rew) > 1:
                self._update_ppo(b_int, b_float, b_act, b_lp, b_rew, b_mask)
            
            rewards_hist.append(ep_r)
            if len(rewards_hist)>50: rewards_hist.pop(0)
            pbar.set_postfix({'AvgR': f"{np.mean(rewards_hist):.2f}"})

    def _update_ppo(self, b_int, b_float, b_act, b_lp, b_rew, b_mask):
        returns = []
        R = 0
        for r, m in zip(reversed(b_rew), reversed(b_mask)):
            R = r + config.PPO_GAMMA * R * m
            returns.insert(0, R)
        returns = torch.tensor(returns)
        
        if returns.numel() > 1 and returns.std() > 1e-5:
            returns = (returns - returns.mean()) / (returns.std() + 1e-5)
        elif returns.numel() > 1:
             returns = returns - returns.mean()
        
        s_int = torch.cat(b_int)
        s_float = torch.cat(b_float)
        act = torch.stack(b_act)
        old_lp = torch.stack(b_lp).detach()
        
        for _ in range(config.PPO_K_EPOCHS):
            logits, vals = self.model(s_int, s_float)
            dist = Categorical(logits=logits)
            new_lp = dist.log_prob(act)
            ratio = torch.exp(new_lp - old_lp)
            
            surr1 = ratio * returns
            surr2 = torch.clamp(ratio, 1-config.PPO_EPS_CLIP, 1+config.PPO_EPS_CLIP) * returns
            
            v_pred = vals.squeeze()
            if v_pred.shape != returns.shape:
                v_pred = v_pred.view(-1)
                returns = returns.view(-1)
                
            loss = -torch.min(surr1, surr2).mean() + 0.5 * F.mse_loss(v_pred, returns)
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

    def evaluate(self, test_scores):
        print("\n>>> [Step 4] 评估测试集...")
        self.model.eval()
        results = []
        
        cnt_detected = 0
        cnt_kb_covered = 0
        cnt_path_match = 0
        cnt_root_match = 0
        cnt_new = 0
        
        env = CausalTracingEnv(self.causal_graph, test_scores, self.threshold_df, self.expert_knowledge)
        
        for win_idx in range(len(test_scores)):
            scores = test_scores[win_idx]
            active = []
            for t_name in config.TRIGGER_SENSORS:
                if t_name in self.sensor_map:
                    idx = self.sensor_map[t_name]
                    if scores[idx] > config.TRIGGER_SCORE_THRESH:
                        active.append((t_name, idx))
            
            for t_name, t_idx in active:
                cnt_detected += 1
                state_data = env.reset(force_window_idx=win_idx, force_trigger=t_name)
                path_idxs = [t_idx]
                done = False
                
                while not done:
                    s_int = state_data[0].unsqueeze(0)
                    s_float = state_data[1].unsqueeze(0)
                    valid = env.get_valid_actions(path_idxs[-1])
                    if len(valid) == 0: break
                    
                    logits, _ = self.model(s_int, s_float)
                    mask = torch.full_like(logits, -1e9)
                    mask[0, valid] = 0
                    act = torch.argmax(logits + mask, dim=1).item()
                    state_data, _, done, _ = env.step(act)
                    path_idxs.append(act)
                    if len(path_idxs) >= config.MAX_PATH_LENGTH: done = True
                
                path_names = [self.idx_to_sensor[i] for i in path_idxs]
                root = path_names[-1]
                root_score = scores[self.sensor_map[root]]
                
                match_status = "未定义"
                logic = ""
                
                if t_idx in self.expert_knowledge:
                    cnt_kb_covered += 1
                    entry = self.expert_knowledge[t_idx]
                    logic = entry.get('logic', '')
                    
                    real_roots = [self.idx_to_sensor[r] for r in entry['roots']]
                    rm = False
                    for p_node in path_names:
                        if p_node in real_roots:
                            rm = True
                            break
                    
                    pm = False
                    path_set = set(path_idxs)
                    for exp_p in entry['paths']:
                        exp_set = set(exp_p)
                        intersection = len(path_set.intersection(exp_set))
                        union = len(path_set.union(exp_set))
                        if union > 0 and (intersection / union) >= 0.6:
                            pm = True
                            break
                    
                    if pm:
                        match_status = "路径吻合"
                        cnt_path_match += 1
                        cnt_root_match += 1
                    elif rm:
                        match_status = "仅根因吻合"
                        cnt_root_match += 1
                    else:
                        match_status = "不吻合"
                else:
                    match_status = "新链路"
                    cnt_new += 1
                
                results.append({
                    "窗口ID": win_idx,
                    "诱发变量": t_name,
                    "溯源路径": "->".join(path_names),
                    "根因变量": root,
                    "根因异常分": f"{root_score:.3f}",
                    "是否知识库": "是" if t_idx in self.expert_knowledge else "否",
                    "匹配情况": match_status,
                    "机理描述": logic
                })

        denom = max(cnt_kb_covered, 1)
        summary = [
            {"指标": "检测到的总异常样本数", "数值": cnt_detected},
            {"指标": "知识库覆盖的样本数", "数值": cnt_kb_covered},
            {"指标": "异常链路准确率", "数值": f"{cnt_path_match/denom:.2%}"},
            {"指标": "根因准确率", "数值": f"{cnt_root_match/denom:.2%}"},
            {"指标": "新发现异常模式数", "数值": cnt_new}
        ]
        
        save_path = os.path.join(config.RESULT_SAVE_DIR, config.TEST_RESULT_FILENAME)
        with pd.ExcelWriter(save_path, engine='openpyxl') as writer:
            pd.DataFrame(summary).to_excel(writer, sheet_name='Sheet1_概览指标', index=False)
            pd.DataFrame(results).to_excel(writer, sheet_name='Sheet2_测试集详情', index=False)
            
        print("\n" + "="*50)
        print(pd.DataFrame(summary).to_string(index=False))
        print(f"\n文件已保存: {save_path}")
        print("="*50)

    def save_model(self):
        path = os.path.join(config.MODEL_SAVE_DIR, "ppo_tracing_model.pth")
        torch.save(self.model.state_dict(), path)