第五轮:锁/续跑/超时/自动分支/R4 探针与 encode-once/控制台加固;阶段内断点续跑(gen 按条、prep 按 shard、训练与 APG 定时原子存档+体面停止);分歧以实现方意见为准

This commit is contained in:
Hark
2026-10-17 18:57:48 +00:00
parent e3bb78388d
commit a5e779c515
36 changed files with 1199 additions and 285 deletions
+16
View File
@@ -126,6 +126,22 @@ class Policy(nn.Module):
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)