TikTok GAE：TD 残差、反向递推与终止 Bootstrap

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

• 输入一段轨迹T奖励、T+1 value，terminals标真实终止；段尾非终止时提供bootstrap value。
• 多段独立轨迹需各自重置carry，不能让截断后的新episode接续旧carry；代码支持真终止mask。

解题路线

1. TD残差δ_t=r_t+γV(s_t+1)−V(s_t)，真终止下一value不应计。GAE优势A_t=Σ_l(γλ)^l δ_t+l直到轨迹边界，利用value作baseline。
2. 倒序递推A_t=δ_t+γλ·continuation·A_t+1，一次扫描避免每个t重新求后缀。return target常取A_t+V(s_t)，与裸奖励和不是同一个数。
3. λ=0只用一步TD，λ近1用更长残差序列，通常降低某些bootstrap偏差但增加方差；不能说λ=1在不完美value/截断下必无偏所有目标。
4. 真实终止continuation=0，既不bootstrap也不跨episode累积。时间限制截断可能仍应bootstrap，但残差递推不能跨到下一独立episode；实际批次可用分别的bootstrap与trace边界mask。
5. T+1 value提供最后一格，漏它会把非终止段尾误当0。padding与EOS在LLM rollout中需按奖励/终止语义设计，不能依字符串长度猜是否done。
6. advantage对actor通常固定/stop gradient，value训练另用returns；batch归一化优势是额外选择，单样本/近零方差需epsilon和mask。
7. 独立验证显式Σ残差与倒序实现一致，覆盖γ/λ为0/1、段中终止和非终止段尾，避免只用单条无终止轨迹。

手动走一遍

两步奖励[1,2]、values全0、γλ=0.5且末步终止：A1=2，A0=1+0.5×2=2。若第一步真终止，则A0只取1，不引用第二episode奖励。

Python 参考实现

下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。

代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。

from math import exp, log1p, isfinite


def generalized_advantage(rewards, values, terminals, gamma=.99, lam=.95):
    n = len(rewards)
    if len(values) != n + 1 or len(terminals) != n or not 0 <= gamma <= 1 or not 0 <= lam <= 1:
        raise ValueError('invalid trajectory shape or discount')
    if any(not isfinite(v) for v in rewards + values):
        raise ValueError('finite trajectory required')
    result, carry = [0.0] * n, 0.0
    for t in range(n - 1, -1, -1):
        continuation = 0.0 if terminals[t] else 1.0
        delta = rewards[t] + gamma * continuation * values[t + 1] - values[t]
        carry = delta + gamma * lam * continuation * carry
        result[t] = carry
    return result, [a + v for a, v in zip(result, values)]

复杂度

T步O(T)时间，输出优势/return O(T)，倒序carry O(1)工作；value模型推理成本另计。

易错点与追问

容易失分的边界

• 终止仍加下状态value。
• 截断和终止统一清所有bootstrap。
• 跨独立episode传播carry。

继续深挖

• 两个mask怎样区别？
• 优势归一化如何处理padding？
• λ改变哪些偏差来源？