File size: 8,468 Bytes
f73f9b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 | import torch
import torch.nn as nn
from .encoder import LLMEncoder
from .mdp_head import MDPHead
from .decision_heads import DecisionHeads
from .domain_moe import DomainMoE
class SystemOneModel(nn.Module):
"""System One 决策链模型：完整组装。
架构：
1. LLM 编码器: 文本 → 语义嵌入（潜空间）
2. MDP Head: 隐式多步 Bellman 前瞻
3. 决策头族: 并行输出 Choice / Transition / Return / Confidence
关键设计：内部隐式多步前瞻，对外紧凑下一步输出。
+-- 2025-09-24: 逐动作架构重构 --+
所有 MLP 接受 [scene_emb ∥ action_vec[a]] 作为特征，对每个动作独立打分，
动作数 n 由运行时 action_vec 的形状决定，
不再依赖固定 num_actions 槽位，不再使用 one-hot 位置编码。
"""
def __init__(self, config):
super().__init__()
self.config = config
self.num_outcomes = config.get("num_outcomes", 2)
self.use_cross_attention = config.get("use_cross_attention", False)
self.encoder = LLMEncoder(config)
self.mdp_head = MDPHead(config)
self.decision_heads = DecisionHeads(config)
# 领域专家 MoE（默认禁用：config.domains 为空时自动关闭）
self.use_domain_experts = config.get("use_domain_experts", False)
self.domain_moe = DomainMoE(config)
def encode_options(self, actions_texts, device=None):
"""编码选项文本列表为动作语义嵌入（支持每样本不同动作数）。
Args:
actions_texts: list of list[str]，每个样本的选项文本列表
（每样本动作数可以不同）
device: 目标设备
Returns:
action_embeddings: (batch, batch_max_n, hidden_dim) 每个选项的语义嵌入
num_valid: (batch,) 每个样本的有效动作数
"""
batch_size = len(actions_texts)
num_valid = torch.tensor(
[len(row) for row in actions_texts], dtype=torch.long
)
batch_max_n = num_valid.max().item()
flat_texts = []
for row in actions_texts:
flat_texts.extend(row)
# padding 到 batch_max_n；padding 条目用占位符，编码后会被 mask
pad_count = batch_max_n - len(row)
flat_texts.extend(["无" for _ in range(pad_count)])
flat_emb = self.encoder(flat_texts)
action_embeddings = flat_emb.view(batch_size, batch_max_n, -1)
if device is not None:
action_embeddings = action_embeddings.to(device)
num_valid = num_valid.to(device)
return action_embeddings, num_valid
def _build_action_vec(self, embeddings, action_texts, num_valid_actions):
"""构建 action_vec: (batch, n, hidden_dim)。
语义模式：对各选项文本编码。
槽位模式：用可学习的位置嵌入表。
"""
if action_texts is not None:
action_embeddings, inferred_num_valid = self.encode_options(
action_texts, device=embeddings.device
)
if num_valid_actions is None:
num_valid_actions = inferred_num_valid
return action_embeddings, num_valid_actions
else:
max_n = self.mdp_head.max_actions
n = num_valid_actions.max().item() if num_valid_actions is not None else max_n
indices = (
torch.arange(max_n, device=embeddings.device)
.unsqueeze(0)
.expand(embeddings.shape[0], -1)
)
act_vec = self.mdp_head.action_embed_table(indices)
if n < max_n:
act_vec = act_vec[:, :n, :]
return act_vec, num_valid_actions
def _compute_domain_cond(self, scene_emb):
"""计算领域条件向量（领域 MoE 自动路由）。
无需外部传入领域标签——门控网络从场景语义自动判断激活哪些专家。
Args:
scene_emb: (batch, hidden_dim) 场景语义嵌入
Returns:
cond: (batch, hidden_dim) 领域条件向量；未启用时返回 None
"""
if not self.use_domain_experts or not self.domain_moe.enabled:
return None
return self.domain_moe(scene_emb)
def forward(
self,
texts,
action_texts=None,
return_mdp_details=False,
num_valid_actions=None,
):
"""完整前向传播。
Args:
texts: List[str] 输入文本
action_texts: list of list[str]，每个样本的选项文本（可选）
return_mdp_details: 是否返回 MDP 内部细节
num_valid_actions: (batch,) 有效动作数（用于 padding mask）
Returns:
result: dict 包含所有决策输出
"""
embeddings = self.encoder(texts)
action_vec, num_valid = self._build_action_vec(
embeddings, action_texts, num_valid_actions
)
# 领域条件向量由门控网络从场景语义自动路由，无需外部领域标签
cond = self._compute_domain_cond(embeddings)
q_summary, q_history, rewards, attended_emb = self.mdp_head(
embeddings, action_vec, num_valid, cond
)
choices, choice_probs, transition_probs, expected_returns, confidences = (
self.decision_heads(attended_emb, action_vec, num_valid, cond)
)
result = {
"choices": choices,
"choice_probs": choice_probs,
"transition_probs": transition_probs,
"expected_returns": expected_returns,
"q_values": q_summary,
"confidences": confidences.squeeze(-1),
"embeddings": embeddings,
}
if return_mdp_details:
result["q_history"] = q_history
result["rewards"] = rewards
return result
def generate_compact_output(self, texts, action_texts=None, num_valid_actions=None):
"""生成紧凑的下一步输出（符合输出契约）。
Args:
texts: List[str] 输入文本
action_texts: list of list[str]，每个样本的选项文本（可选）
num_valid_actions: (batch,) 有效动作数
Returns:
outputs: List[dict] 每个样本的紧凑输出
"""
embeddings = self.encoder(texts)
action_vec, num_valid = self._build_action_vec(
embeddings, action_texts, num_valid_actions
)
# 领域条件向量由门控网络从场景语义自动路由
cond = self._compute_domain_cond(embeddings)
q_summary, q_history, rewards, attended_emb = self.mdp_head(
embeddings, action_vec, num_valid, cond
)
choices, choice_probs, transition_probs, expected_returns, confidences = (
self.decision_heads(attended_emb, action_vec, num_valid, cond)
)
n_actions = action_vec.shape[1]
outputs = []
batch_size = len(texts)
for i in range(batch_size):
transition_dict = {}
for a in range(n_actions):
transition_dict[f"action_{a}"] = (
transition_probs[i, a].detach().cpu().tolist()
)
output = {
"answers": {
"q1": {
"type": "choice",
"choice": f"选项{choices[i].item() + 1}",
"prob_dist": transition_dict,
"expected_return": round(
expected_returns[i, choices[i]].item(), 4
),
"confidence": round(confidences[i].item(), 4),
}
}
}
outputs.append(output)
return outputs
def get_trainable_parameters(self):
"""获取可训练参数（用于优化器）。"""
trainable_params = []
for name, param in self.named_parameters():
if param.requires_grad:
trainable_params.append(param)
return trainable_params
def count_parameters(self):
"""统计参数数量。"""
total = sum(p.numel() for p in self.parameters())
trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
return total, trainable
|