GAE Bug for Envpool: Dummy Step Leak

Hello,

I've noticed a bug related to all PPO Atari scripts that use envpool.

Envpool has an odd handling of terminal states. The next transition after done=True is a "dummy" transition from the terminal state to the start state. That is, the stream of experience is:

State: $$S_{T-1} \rightarrow{} S_T \rightarrow{} S_0 \rightarrow{} S_1$$ Done: T -> F -> F -> F

The "Dummy" Step is from $$S_T\rightarrow{} S_0$$. The GAE calculation will set $$\delta = \gamma * V(S_0) * 1 - V_S(T)$$ for the dummy step. Since $$S_T$$ appears as an observation, the network is trained on this incorrect advantage.

Reproduction

Here is a notebook that visualizes breakout transitions, showing the dummy.

termination_error_reproduce.ipynb termination_error_reproduce.pdf

Fix:

In my own code, I handled this by making a wrapper around cleanRL that flags a transition as a dummy transition whether done=true on the last step. Then, the trace is cut (i.e. next value is zero-d out) whenever done=True or dummy=True:

def calculate_gaeE(traj_batch, γ, λ,):
    done = traj_batch.done
    is_dummy = traj_batch.info.get('is_dummy', jnp.zeros_like(done))
    cut_e_trace = done | is_dummy
    cut_e_mult = 1.0 - cut_e_trace.astype(jnp.float32)

    scan_inputs = (traj_batch, cut_e_mult)

    def _get_advantages(gae, inputs):
        transition, continue_e = inputs
        
        # --- Extrinsic ---
        delta = transition.reward + γ * transition.next_value * continue_e - transition.value
        gae = delta + (γ * λ * continue_e * gae)
        
        return gae, gae

    initial_acc = jnp.zeros_like(traj_batch.value[0])
    
    _, advantages = jax.lax.scan(
        _get_advantages, initial_acc, scan_inputs, reverse=True, unroll=16
    )
    
    return advantages, advantages + traj_batch.value