TB-GAE: terminal feedback across turns

Let γ = 0.9, λ = 0.8, q = γλ = 0.72; only the terminal reward is 1.

Interaction

Turn 1

Turn 2

g₁

The 1st illustrative generated position in the turn-1 command. Real commands usually contain more tokens.

g₂

The last illustrative generated position in the turn-1 command. After the tool call finishes, the environment returns observations.

o₁ … o_L

The environment returns L observation tokens. They remain in the next turn’s model context; only the advantage recursion below skips these positions. The comparison uses L=3 for readability.

g₃

The first generated position of turn 2. The policy and Critic inputs still include the previous turn’s tool output.

g₄

The last generated position of turn 2. The example truly terminates the task here and applies the terminal reward once.

R = 1

Set terminal reward R=1. At true termination, the successor value and recursive advantage term are 0.

Observations stay in context; only GAE recursion skips them.

One run, three clocks for credit assignment

All-token chainBoth token types

g₁

Generated position g1. In the mixed-text chain, both generated and observed positions enter the advantage recursion. Even when observation positions are masked in the training loss, they still cause decay along this recursion.

g₂

Generated position g2. In the mixed-text chain, both generated and observed positions enter the advantage recursion. Even when observation positions are masked in the training loss, they still cause decay along this recursion.

g₃

Generated position g3. In the mixed-text chain, both generated and observed positions enter the advantage recursion. Even when observation positions are masked in the training loss, they still cause decay along this recursion.

g₄

Generated position g4. In the mixed-text chain, both generated and observed positions enter the advantage recursion. Even when observation positions are masked in the training loss, they still cause decay along this recursion.

R = 1

Reward is injected once at the final generated position, forming a terminal TD residual; recursion then proceeds backward through generated positions.

o₁

A tool-observation position. In the mixed-text chain it still enters GAE recursion, adding one decay step; the policy loss is 0 at this position.

o₂

A tool-observation position. In the mixed-text chain it still enters GAE recursion, adding one decay step; the policy loss is 0 at this position.

o₃

A tool-observation position. In the mixed-text chain it still enters GAE recursion, adding one decay step; the policy loss is 0 at this position.

Terminal residual weight: q⁶ ≈ 0.1393 (3 log tokens here)

SAO · Skip obs.Generated hops ×q

g₁

Generated position g1. This illustrative recursion skips only observations; every edge between generated positions still multiplies by q. It separates skipping observations from using a cross-turn clock.

g₂

Generated position g2. This illustrative recursion skips only observations; every edge between generated positions still multiplies by q. It separates skipping observations from using a cross-turn clock.

g₃

Generated position g3. This illustrative recursion skips only observations; every edge between generated positions still multiplies by q. It separates skipping observations from using a cross-turn clock.

g₄

Generated position g4. This illustrative recursion skips only observations; every edge between generated positions still multiplies by q. It separates skipping observations from using a cross-turn clock.

R = 1

Reward is injected once at the final generated position, forming a terminal TD residual; recursion then proceeds backward through generated positions.

Skip obs.

Terminal residual weight: q³ ≈ 0.3732

TB · Turn clockOnly across turns ×q

g₁

Generated position g1. TB-GAE retains this position’s own pre-token value baseline. Targets can be shared within a turn while advantages remain different.

g₂

Generated position g2. TB-GAE retains this position’s own pre-token value baseline. Targets can be shared within a turn while advantages remain different.

g₃

Generated position g3. TB-GAE retains this position’s own pre-token value baseline. Targets can be shared within a turn while advantages remain different.

g₄

Generated position g4. TB-GAE retains this position’s own pre-token value baseline. Targets can be shared within a turn while advantages remain different.

R = 1

Reward is injected once at the final generated position, forming a terminal TD residual; recursion then proceeds backward through generated positions.

Skip obs.

Terminal residual weight: q = 0.72; no within-turn decay

Shared turn targets; a distinct advantage for each prefix

Turn 1

g₁

g₂

0.20

0.40

0.81

0.61

0.41

At the next turn’s start, V=0.5 and A=0.5. The shared target is U₁=0.9×(0.5+0.8×0.5)=0.81. Each prefix subtracts its own V, giving A₁=0.81−0.2=0.61 and A₂=0.81−0.4=0.41.

Turn 2

g₃

g₄

0.50

0.80

1.00

0.50

0.20

Turn 2 ends in true termination, with only terminal reward 1. Both generated positions therefore share target 1; subtracting their own V gives advantages 0.5 and 0.2.

Pre-token value V

Shared target U

Adv. A = U − V

U₁ = 0.9 × (0.5 + 0.8 × 0.5) = 0.81

Above: terminal-residual propagation weights. Below: return targets and advantages.

With γ = λ = 1, all three chains have coefficient 1; RPR-1 uses standard GAE with this setting.