记忆
模型可以拥有记忆。只要你所在赛季的重量级允许,你的计算图就能在某一回合写出一个张量,由你的适配器在下一回合读回,并一直延续到对局结束,而运行器公开地携带这份记忆:它是你清单中声明的一个输出,会以该输出自己的名字,在你席位的下一次视图中交还给你。没有哪个节点保存它,也没有哪个节点隐藏它;而不声明记忆的模型,玩法和这一页出现之前完全一样。
两种记忆
清单可以在 policy 之外,再声明一个名为 memory 的输出、一个名为 ant_memory 的输出,两者都声明,或者都不声明:
| 输出 | 形状 | 用途 |
|---|---|---|
memory | 固定形状,或 [1, K, "H", "W"] | 地图层面:气味、见过食物的位置、敌人出现过的位置、已探索的区域 |
ant_memory | [1, "N", K] | 每只蚂蚁一行,按本回合 mine 的顺序:目标、角色、追逐目标已有多少回合 |
ant_memory 需要分清每只蚂蚁是谁,观测中的 ids 正是为此;memory 只需要地图本身。
"outputs": [
{ "name": "policy", "dtype": "f32", "shape": [1, 5, "H", "W"] },
{ "name": "memory", "dtype": "f32", "shape": [1, 1, "H", "W"] }
]
一个回合,两个步骤
每个回合,记忆都会依次经过两个步骤:
- 先是适配器,在模型之前。它读取上一回合的记忆和本回合的观测,更新它能更新的部分,再把输入交给模型:更新后的记忆,以及地图。
- 然后是模型。它算出走法,其
memory输出就是下一回合的记忆。在不需要改动的地方,这个输出就是通过Identity原样送回的输入,不花任何开销。
带到下一回合的记忆,就是模型输出的内容。适配器所做的更新只能通过模型输入才能到达那里,所以想保留适配器成果的计算图,要用 Identity 把那个输入原样送回。
适配器语言没有张量之间的算术运算(原因),所以这两个步骤各有所长。适配器负责选取、按 id 重新对齐和查表,计算图负责计算:
| 步骤 | 气味(memory) | 每只蚂蚁一个目标(ant_memory) |
|---|---|---|
| 适配器 | 把气味原样传递,并用 scatter 从 mine 构建本回合的沉积平面 | 按 id 把上一回合的行重新对齐到本回合的蚂蚁,并查出每个已存目标处的食物 |
| 模型 | Max(scent × 0.95, deposit):先衰减,再沉积 | 丢弃已到达、已消失或已过期的目标,挑选新目标,移动,并为每只蚂蚁输出一行 |
在适配器中读取记忆
第 0 回合的视图里没有 memory,所以适配器要自己构建零张量。此后它像读取其他字段一样读取这个键:
{ "name": "scent_in", "dtype": "f32", "shape": [1, 1, "H", "W"],
"adapter": { "if": [
{ "var": "memory" },
{ "tensor": [{ "var": "memory" }] },
{ "zeros": [{ "merge": [[1, 1], { "var": "size" }] }, "f32"] } ] } }
用 {"tensor": [{"var": "memory"}]} 读取它。无论在节点上还是在 tinybrains 里,这个键保存的都是你的计算图写出的张量,tensor 会按你输出声明的数据类型和形状把它原样传递。张量是不透明的:适配器能读到它的形状和数据类型,却读不到里面的元素,所以试图查看其内部值的程序在两边都一无所获。记忆唯一以 JSON 形式出现的地方,是你交给 tinybrains adapt 的观测文件,它在加载时会把记忆解码成张量。
从这里开始,适配器语言的其余部分照常适用。crop 和 concat 可以推移一段历史,gather 和 to_list 能把存下的值读回 JSON 用于比较,scatter 则把本回合的事实写到一个地图大小的平面上。
地图上的气味
气味轨迹是一种地图记忆。适配器像上面那样把上一回合的气味原样传递,再从 mine 构建一个沉积平面作为第二个输入:
{ "name": "deposit", "dtype": "f32", "shape": [1, 1, "H", "W"],
"adapter": { "reshape": [
{ "scatter": [{ "var": "mine" }, { "var": "size" }, "f32"] },
{ "merge": [[1, 1], { "var": "size" }] } ] } }
计算图让旧气味衰减,再把新气味叠加上去,并把结果作为 memory 输出:
memory = Max(scent_in × 0.95, deposit)
这不需要 id:每个回合里,格子还是那个格子。同样的模式也能记住见过的食物、见过的敌方蚁丘或已探索的陆地:用 Max 来记住,用衰减后的副本来遗忘。
每只蚂蚁一份记忆
ant_memory 为每只蚂蚁保留一行,而且各行按 id 跟随蚂蚁,从不按位置:mine 每回合都会重新排序,所以本回合的第三行通常和上一回合的第三行不是同一只蚂蚁。
适配器把本回合的 ids 作为一个输入交给模型,模型把每只蚂蚁的 id 写进它那一行的第 0 列。到了下一回合,适配器把上一回合的各行变成一张按 id 索引的表,再按本回合 ids 的顺序把这张表读回来:
prev = to_list(ant_memory) [P][K] as JSON
col_k = scatter(map(prev, [row.0, 0, row.k]), [S, 1]) one per column: row id → that column
table = concat(col_0 … col_K-1, axis 1) [S, K], row = id
rows = reshape(gather(table, ids, 0), [1, N, K]) this turn's ants, in mine order
每一步都是适配器语言已有的操作符:to_list、map、scatter、concat、gather 和 reshape。map 的主体只读取自己那一行,所以作用域规则不会造成麻烦。假设上一回合的蚂蚁是 7、3 和 12,本回合是 3、12 和 20:蚂蚁 7 已经死亡,蚂蚁 20 是新出生的。读回的各行依次是 3 的那一行、12 的那一行,然后是全零。
按 id 作键,让三种情况无需额外代码就能正常工作:
- 被判罚的回合。记忆停留在最后一次写入的值(见下文),而 id 不会变,所以各行仍能找到自己的蚂蚁。
- 死亡蚂蚁的那一行不会再被读回,并从下一次输出中消失。
- 新出生的蚂蚁读到的是全零的一行,模型按新蚂蚁的待遇处理它。
第 0 回合没有 ant_memory,适配器直接构建 [1, N, K] 的零张量。
表的高度 S 由你决定。遇到大于等于 S 的 id,gather 会让这次调用失败;而一个席位每回合每个蚁丘最多出生一只蚂蚁,所以只要 S 超过每席位蚁丘数 × (回合数 + 1),就永远不会触及。较小的 S 可以在两边都取 id % S,并接受两只活着的蚂蚁可能共用一行。这次连接的开销随 S 和蚂蚁数量增长;在 S = 4,096 时,它不到适配器预算的 5%。用 tinybrains adapt 测量你自己的开销。
保持目标直到到达或失效
把每只蚂蚁的目标存成一个绝对坐标的格子,ant_memory 的各列为 [id, has_target, tr, tc, turns]。
先是适配器,在模型之前:
- 像上面那样,按 id 把上一回合的行重新对齐到本回合的蚂蚁。
- 在每个已存目标处查出模型需要的信息。以食物为例,把食物平面展平,对每一行在
tr × W + tc处取出它的值,作为形状为[1, N]的第二个输入。目标是绝对坐标的格子,所以这只是对各行的一次map。
然后是模型:
- 保留蚂蚁的目标,除非已经到达、那里的食物已经消失,或者蚂蚁追逐它的回合数太多。这在计算图里就是一个
Where。 - 没有保留目标的地方,让策略头挑一个新的。
- 朝目标移动,并输出
ant_memory,每只蚂蚁一行。
被阻挡的移动不需要预测。引擎会结算每一次移动,下一回合蚂蚁的 id 就会告诉你它到了哪里。
记忆的携带
运行器为每场对局的每个席位保存每个记忆输出的最新值,并以同一个键交还:
| 时机 | 视图中的 memory(以及 ant_memory) |
|---|---|
| 第 0 回合 | 不存在。适配器自己构建零张量 |
| 一次有应答的调用之后 | 那次调用的同名输出 |
| 一次失败或超时的调用之后 | 不变:模型最后一次写入的记忆 |
| 新的对局 | 再次不存在 |
席位之间从不共享记忆,即使同一场对局里的两个席位是同一个模型。裁判只读取 policy,游戏规则从不读取记忆,所以回放里没有记忆:重新模拟一场对局只需要它的动作。
判罚不会改动记忆,所以一次调用失败的模型仍保有自己的笔记,可以从上次停下的地方继续。而记忆输入无法接受自己输出的模型,从第二回合起每回合都会失败;审核的往返检查(见下文)正是为了在对局暴露它之前先抓住它。
上限
重量级用两个数字来允许记忆,都由赛季设定:memory_flat_bytes 和 memory_cell_bytes。在一张有 rows × cols 个格子的地图上,你的两种记忆合计最多可占
cap = memory_flat_bytes + memory_cell_bytes × rows × cols
两者默认都是 0,即没有记忆。重量级一页说明了在哪里查看赛季的数字。
记忆的价格取决于它声明的形状。每个输出的开销是元素个数乘以数据类型的宽度:i8、u8 和 bool 为 1 字节,i16 和 u16 为 2 字节,f32、i32 和 u32 为 4 字节,f64、i64 和 u64 为 8 字节。没有命名维度的输出是固定开销。带有任何命名维度的输出按每格计价:其数字维度的乘积,乘以地图的格子数。这也涵盖了 ant_memory 的蚂蚁维度,因为一个席位的蚂蚁数永远不会超过地图的格子数。
memory最多可有两个命名维度,ant_memory最多一个。- 一个形状不能重复使用同一个名字。
审核会把两个输出合计,分别在赛季可用的最小地图(576 格,24 × 24)和最大地图(14,880 格)上计价。两边都随格子数线性增长,所以在两端都放得下的模型,在之间的每一张地图上也放得下,包括赛季之后新增的地图。对局时不会再检查上限。
举个例子:memory 声明为 f32[1, 2, "H", "W"] 是每格 8 字节,ant_memory 声明为 i8[1, "N", 5] 再加 5 字节。在 14,880 格时合计 193,440 字节,所以重量级的 memory_cell_bytes 需要达到 13 或更高,或者有一个在两端都能补足差额的固定配额。能存成 u8 或 i8 的就这么存,并在计算图内 Cast:每个值一字节,价格只有 f32 的四分之一。
记忆完全不计入你的大小。你的重量级仍由两个文件的字节数决定(S')。想要更多记忆的参赛者,就进入更大的重量级。记忆在每场对局开始时都是零,所以无法借它偷带权重。
审核检查什么
其中三项审核仅凭你的清单和重量级判定,不运行模型;第四项则通过让你的模型吃回自己的记忆来判定。每一项都是最终结论,所指出的问题也只有你能改变:
| 原因 | 时机 | 下一步 |
|---|---|---|
MEMORY_NOT_ALLOWED | 你声明了记忆输出,而你的重量级允许的是 0 和 0 | 删掉这个输出,或进入本赛季允许记忆的重量级 |
MEMORY_TOO_LARGE | 记忆的价格在 576 格或 14,880 格时超过了重量级的上限 | 缩小形状或数据类型,或进入更大的重量级 |
MEMORY_SHAPE | 命名维度过多、同一名字用了两次,或记忆输出缺少数据类型和形状 | 按两种记忆所示的方式声明 |
MEMORY_ROUND_TRIP | 喂回自己的记忆后,模型的调用失败了 | 检查每个记忆输入接受的正是其输出写出的内容 |
往返检查会把参考观测串联起来。审核按顺序运行它们,只要前后两个观测在同一张地图上且上一次调用有应答,每次调用就会拿到上一次写出的记忆。所以读错自己记忆的模型会在审核时失败,而不是在每场对局的第 1 回合被判罚。只在对局后期才显现的 bug 在这里抓不到:在你自己的机器上打一场对局来找它。
单一设备
每个运行器都在 CPU 上执行模型,这一点固定在它的配置里。一场对局的两个席位在同一个运行器上进行,所以没有哪个设备能偏袒某个席位;但记忆会把末位数字的差异带入之后的每一个回合。固定设备让带记忆的对局可以复现,而且对这种规模的计算图来说,这也是更快的选择。
测试记忆
tinybrains 按与运行器相同的规则携带记忆:打一场对局,你的席位每回合都会拿回自己的记忆。tinybrains check 会打印你的记忆开销,固定部分和每格部分各是多少,并接受 --memory-flat-bytes 和 --memory-cell-bytes 来对照某个重量级评判。check 也像审核那样串联运行参考观测。tinybrains adapt 接受带有 memory 或 ant_memory 的观测,既可以是嵌套数组,也可以是传输格式,所以训练器可以在设置了记忆的情况下比较自己的张量和适配器的张量(测试)。
携带规则的一个细节:如果一次调用返回的策略头裁判无法读取,那算一次判罚,但它写出的记忆会被保留。
训练器在 Python 里按同样的携带规则携带自己的记忆。如果记忆是观测的固定函数,比如上面的气味,用你已有的训练循环就能训练:训练器在 numpy 里算出同样的平面,再喂进去。由计算图计算的记忆则必须按回合顺序训练,因为每一回合的记忆都是计算图在上一回合的输出:这是一个用截断的时序反向传播训练的循环网络,在回合之间携带的是运行器将会携带的那个取整后的整数。
平台自己的基线两种都做了。ants/baselines 用 --memory 训练任意重量级,带两个用 Max 保留的平面:见过的食物和见过的敌方蚁丘;它的 train/seq.py 把同样两个平面训练成由优化器塑造的门控更新,加上 --ants 后还有每只蚂蚁一行、按 id 跟随蚂蚁的记忆,与上文完全一样。它的一致性测试通过 tinybrains adapt 运行每一个记忆适配器,包括按 id 的连接。各自携带了什么、又带来了多少收益,请读它的 README.md。