This is a PyTorch implementation of paper Generalized Advantage Estimation.

You can find an experiment that uses it here.

`15import numpy as np`

`18class GAE:`

```
19 def __init__(self, n_workers: int, worker_steps: int, gamma: float, lambda_: float):
20 self.lambda_ = lambda_
21 self.gamma = gamma
22 self.worker_steps = worker_steps
23 self.n_workers = n_workers
```

$\hat{A_t^{(1)}}$ is high bias, low variance, whilst $\hat{A_t^{(\infty)}}$ is unbiased, high variance.

We take a weighted average of $\hat{A_t^{(k)}}$ to balance bias and variance. This is called Generalized Advantage Estimation. We set $w_k = \lambda^{k-1}$, this gives clean calculation for $\hat{A_t}$

`25 def __call__(self, done: np.ndarray, rewards: np.ndarray, values: np.ndarray) -> np.ndarray:`

advantages table

```
58 advantages = np.zeros((self.n_workers, self.worker_steps), dtype=np.float32)
59 last_advantage = 0
```

$V(s_{t+1})$

```
62 last_value = values[:, -1]
63
64 for t in reversed(range(self.worker_steps)):
```

mask if episode completed after step $t$

```
66 mask = 1.0 - done[:, t]
67 last_value = last_value * mask
68 last_advantage = last_advantage * mask
```

$\delta_t$

`70 delta = rewards[:, t] + self.gamma * last_value - values[:, t]`

$\hat{A_t} = \delta_t + \gamma \lambda \hat{A_{t+1}}$

`73 last_advantage = delta + self.gamma * self.lambda_ * last_advantage`

note that we are collecting in reverse order.
*My initial code was appending to a list and
I forgot to reverse it later.
It took me around 4 to 5 hours to find the bug.
The performance of the model was improving
slightly during initial runs,
probably because the samples are similar.*

```
82 advantages[:, t] = last_advantage
83
84 last_value = values[:, t]
85
86 return advantages
```