176 lines
8.0 KiB
Python
176 lines
8.0 KiB
Python
"""策略网络（方案 §6.1）：条件特征 → 8 个条件 token → 4 层 encoder；16 个动作 token（噪声 z + 时间嵌入 (r,t)）
|
||
→ 4 层 decoder（自注意力 + 交叉注意力，非自回归一次前向）→ 每 token 6 维（平均速度场 u）。
|
||
价值头 / 子目标头只依赖观测（encoder 输出）。
|
||
注意力用显式 matmul+softmax 实现：torch.func.jvp（MeanFlow 训练）需要前向模式自动微分，SDPA 融合内核不支持。"""
|
||
import math
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
from ml.common.features import F_DIM, H
|
||
|
||
|
||
class Attn(nn.Module):
|
||
def __init__(self, d, nh):
|
||
super().__init__()
|
||
self.nh, self.dh = nh, d // nh
|
||
self.q, self.k, self.v, self.o = (nn.Linear(d, d) for _ in range(4))
|
||
|
||
def forward(self, x, mem=None):
|
||
mem = x if mem is None else mem
|
||
B, Lq, D = x.shape; Lk = mem.shape[1]
|
||
q = self.q(x).view(B, Lq, self.nh, self.dh).transpose(1, 2)
|
||
k = self.k(mem).view(B, Lk, self.nh, self.dh).transpose(1, 2)
|
||
v = self.v(mem).view(B, Lk, self.nh, self.dh).transpose(1, 2)
|
||
a = torch.softmax(q @ k.transpose(-1, -2) / math.sqrt(self.dh), dim=-1)
|
||
return self.o((a @ v).transpose(1, 2).reshape(B, Lq, D))
|
||
|
||
|
||
class FFN(nn.Module):
|
||
def __init__(self, d, f):
|
||
super().__init__(); self.a = nn.Linear(d, f); self.b = nn.Linear(f, d)
|
||
|
||
def forward(self, x):
|
||
return self.b(F.gelu(self.a(x)))
|
||
|
||
|
||
class EncBlock(nn.Module):
|
||
def __init__(self, d, nh, f):
|
||
super().__init__(); self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d); self.sa, self.ff = Attn(d, nh), FFN(d, f)
|
||
|
||
def forward(self, x):
|
||
x = x + self.sa(self.n1(x))
|
||
return x + self.ff(self.n2(x))
|
||
|
||
|
||
class DecBlock(nn.Module):
|
||
def __init__(self, d, nh, f):
|
||
super().__init__()
|
||
self.n1, self.n2, self.n3 = nn.LayerNorm(d), nn.LayerNorm(d), nn.LayerNorm(d)
|
||
self.sa, self.ca, self.ff = Attn(d, nh), Attn(d, nh), FFN(d, f)
|
||
|
||
def forward(self, x, mem):
|
||
x = x + self.sa(self.n1(x))
|
||
x = x + self.ca(self.n2(x), mem)
|
||
return x + self.ff(self.n3(x))
|
||
|
||
|
||
def time_emb(t, dim):
|
||
half = dim // 2
|
||
fr = torch.exp(-math.log(1000.0) * torch.arange(half, device=t.device, dtype=t.dtype) / half)
|
||
a = t[:, None] * 1000.0 * fr[None]
|
||
return torch.cat([torch.sin(a), torch.cos(a)], -1)
|
||
|
||
|
||
class Policy(nn.Module):
|
||
def __init__(self, d=192, nh=8, ff=768, n_enc=4, n_dec=4, n_cond=8, f_dim=F_DIM, horizon=H, head='meanflow', K=8):
|
||
super().__init__()
|
||
self.d, self.n_cond, self.h, self.head, self.K = d, n_cond, horizon, head, K
|
||
self.cond = nn.Sequential(nn.Linear(f_dim, 2 * d), nn.GELU(), nn.Linear(2 * d, n_cond * d))
|
||
self.cpos = nn.Parameter(torch.zeros(1, n_cond, d))
|
||
self.enc = nn.ModuleList(EncBlock(d, nh, ff) for _ in range(n_enc))
|
||
self.enc_norm = nn.LayerNorm(d)
|
||
self.ain = nn.Linear(6, d)
|
||
self.apos = nn.Parameter(torch.zeros(1, horizon, d))
|
||
self.temb = nn.Sequential(nn.Linear(2 * 64, d), nn.GELU(), nn.Linear(d, d))
|
||
self.dec = nn.ModuleList(DecBlock(d, nh, ff) for _ in range(n_dec))
|
||
self.dec_norm = nn.LayerNorm(d)
|
||
self.aout = nn.Linear(d, 6)
|
||
self.value = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, 1))
|
||
self.sub = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, 13))
|
||
if head == 'wta': # 预定回退：同骨干 WTA K 模式头（方案 §6.2）
|
||
self.mode_emb = nn.Parameter(torch.randn(1, K, 1, d) * 0.02)
|
||
self.mode_logit = nn.Linear(d, 1)
|
||
nn.init.normal_(self.cpos, std=0.02); nn.init.normal_(self.apos, std=0.02)
|
||
nn.init.zeros_(self.aout.weight); nn.init.zeros_(self.aout.bias)
|
||
|
||
def encode(self, feat):
|
||
x = self.cond(feat).view(-1, self.n_cond, self.d) + self.cpos
|
||
for b in self.enc: x = b(x)
|
||
return self.enc_norm(x)
|
||
|
||
def heads(self, mem):
|
||
p = mem.mean(1)
|
||
return self.value(p).squeeze(-1), self.sub(p)
|
||
|
||
def decode(self, mem, z, r, t):
|
||
"""平均速度场 u(z, r, t)；z [B,16,6]，r,t [B]"""
|
||
te = self.temb(torch.cat([time_emb(t, 64), time_emb(t - r, 64)], -1))
|
||
x = self.ain(z) + self.apos + te[:, None]
|
||
for b in self.dec: x = b(x, mem)
|
||
return self.aout(self.dec_norm(x))
|
||
|
||
def wta_modes(self, mem):
|
||
"""WTA：K 个模式一次解码 → act [B,K,16,6], logits [B,K]"""
|
||
B = mem.shape[0]
|
||
x = (self.ain.bias + self.apos)[:, None] + self.mode_emb # [1,K,16,d]
|
||
x = x.expand(B, -1, -1, -1).reshape(B * self.K, self.h, self.d)
|
||
m = mem.repeat_interleave(self.K, 0)
|
||
for b in self.dec: x = b(x, m)
|
||
x = self.dec_norm(x)
|
||
act = self.aout(x).view(B, self.K, self.h, 6)
|
||
return act, self.mode_logit(x.mean(1)).view(B, self.K)
|
||
|
||
def sample(self, feat, z):
|
||
"""部署接口（ONNX 同签名）：返回 (act [B,16,6] 归一化增量, value [B], sub [B,13])
|
||
meanflow：一步生成 x = z − u(z, 0, 1)；
|
||
wta：取第 0 行观测的 K 个模式，第 b 行输出模式 b mod K（部署时 8 行观测相同；前瞻只用各行的 value）。"""
|
||
mem = self.encode(feat)
|
||
v, s = self.heads(mem)
|
||
B = feat.shape[0]
|
||
if self.head == 'wta':
|
||
act, _ = self.wta_modes(mem[:1])
|
||
idx = torch.arange(B, device=feat.device) % self.K
|
||
return act[0, idx] + 0.0 * z, v, s # z 不参与计算，保留为输入以统一 ONNX 签名
|
||
one, zero = torch.ones(B, device=feat.device, dtype=feat.dtype), torch.zeros(B, device=feat.device, dtype=feat.dtype)
|
||
x = z - self.decode(mem, z, zero, one)
|
||
return x, v, s
|
||
|
||
def sample_shared(self, feat1, z):
|
||
"""部署接口（ONNX，2026-10-18 第五轮，R4 实测后改）：一次观测、K 条候选。
|
||
feat1 [1,F]，z [K,16,6] → act [K,16,6]，value [K]，sub [K,13]（value/sub 对 K 行相同，展开只为保持运行时读法）。
|
||
此前部署把同一观测复制 K=8 行喂 sample()，encoder 被重复算 8 次（约占总计算 29%）；这里只编码一次再展开给解码器。"""
|
||
mem = self.encode(feat1)
|
||
v, s = self.heads(mem)
|
||
K = z.shape[0]
|
||
if self.head == 'wta':
|
||
act, _ = self.wta_modes(mem)
|
||
idx = torch.arange(K, device=z.device) % self.K
|
||
return act[0, idx] + 0.0 * z, v.expand(K), s.expand(K, -1)
|
||
m = mem.expand(K, -1, -1)
|
||
one = torch.ones(K, device=z.device, dtype=z.dtype)
|
||
x = z - self.decode(m, z, torch.zeros_like(one), one)
|
||
return x, v.expand(K), s.expand(K, -1)
|
||
|
||
def sample_k(self, feat, K, z=None):
|
||
"""仿真 / APG 用：每个观测 K 条候选 → act [B,K,16,6], value [B], sub [B,13]"""
|
||
mem = self.encode(feat)
|
||
v, s = self.heads(mem)
|
||
B = feat.shape[0]
|
||
if self.head == 'wta':
|
||
act, logit = self.wta_modes(mem)
|
||
if K < self.K:
|
||
top = logit.topk(K, -1).indices
|
||
act = act.gather(1, top[:, :, None, None].expand(-1, -1, self.h, 6))
|
||
return act, v, s
|
||
z = torch.randn(B * K, self.h, 6, device=feat.device, dtype=feat.dtype) if z is None else z.reshape(B * K, self.h, 6)
|
||
m = mem.repeat_interleave(K, 0)
|
||
one = torch.ones(B * K, device=feat.device, dtype=feat.dtype)
|
||
x = z - self.decode(m, z, torch.zeros_like(one), one)
|
||
return x.view(B, K, self.h, 6), v, s
|
||
|
||
|
||
def build(cfg):
|
||
m = cfg.get('model', {})
|
||
return Policy(d=m.get('d', 192), nh=m.get('heads', 8), ff=m.get('ff', 768), n_enc=m.get('enc', 4), n_dec=m.get('dec', 4),
|
||
head=m.get('head', 'meanflow'), K=m.get('K', 8))
|
||
|
||
|
||
def wta_loss(model, mem, y, w_all=0.05):
|
||
"""WTA-L1 + 模式 CE；返回 (loss, 最优模式输出 [B,16,6])"""
|
||
act, logit = model.wta_modes(mem)
|
||
l1 = (act - y[:, None]).abs().mean(dim=(2, 3)) # [B,K]
|
||
best = l1.argmin(-1)
|
||
loss = l1.gather(1, best[:, None]).mean() + w_all * l1.mean() + 0.1 * F.cross_entropy(logit, best)
|
||
return loss, act[torch.arange(y.shape[0]), best]
|