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
25 def __call__(self, done: np.ndarray, rewards: np.ndarray, values: np.ndarray) -> np.ndarray:
优势表
59 advantages = np.zeros((self.n_workers, self.worker_steps), dtype=np.float32)
60 last_advantage = 0
63 last_value = values[:, -1]
64
65 for t in reversed(range(self.worker_steps)):
如果剧集在步骤之后完成,请掩盖
67 mask = 1.0 - done[:, t]
68 last_value = last_value * mask
69 last_advantage = last_advantage * mask
71 delta = rewards[:, t] + self.gamma * last_value - values[:, t]
74 last_advantage = delta + self.gamma * self.lambda_ * last_advantage
77 advantages[:, t] = last_advantage
78
79 last_value = values[:, t]
82 return advantages