On GAE calculation math

[michael-lutz]

In losses.py, I noticed that the code includes the following step before returning value targets and advantages:

advantages = (rewards + discount * (1 - termination) * vs_t_plus_1 - values) * truncation_mask

From what I understand, compute_vs_minus_v_xs should return the standard GAE result. Why do we perform an additional TD computation at the end?

Second, I was hoping to ask why the value loss includes an extra 0.5 term:

v_loss = jnp.mean(v_error * v_error) * 0.5 * 0.5

Both these decisions seem non-standard. Did you find they improved performance empirically?

[ch33nchan]

the only logical explaination i could come up with was :

• Standard GAE implementation would use vs_minus_v_xs directly as advantages.However, Brax uses an additional TD computation to handle truncated episodes.I guess this approach is mainted for consistency with the rest of the codebase.
• The additional 0.5 factor (making it 0.5 * 0.5 * MSE) reduces the weight of the value loss relative to the policy loss, which is found to be empirical to improve performance in Brax environments

correct me if I'm wrong?

[btaba]

The extra 0.5 factor indeed just scales the value loss. I've futzed with a vf_coef in the past but it didn't seem to make a huge difference.

I'm not really sure about the deviation in the GAE implementation, but using vs_t_plus_1 (the value target) rather than values_t_plus_1 (the value) in the GAE calc. might contribute to more stability in the learning algo since vs_t_plus_1 should have lower variance compared to values_t_plus_1 (just a guess though).