def _compute_advantage(self, rewards, valids, baselines): r"""Compute mean value of loss. Notes: P is the maximum episode length (self.max_episode_length) Args: rewards (torch.Tensor): Acquired rewards with shape :math:`(N, P)`. valids (list[int]): Numbers of valid steps in each episode baselines (torch.Tensor): Value function estimation at each step with shape :math:`(N, P)`. Returns: torch.Tensor: Calculated advantage values given rewards and baselines with shape :math:`(N \dot [T], )`. """ advantages = compute_advantages(self.discount, self._gae_lambda, self.max_episode_length, baselines, rewards) advantage_flat = torch.cat(filter_valids(advantages, valids)) if self._center_adv: means = advantage_flat.mean() variance = advantage_flat.var() advantage_flat = (advantage_flat - means) / (variance + 1e-8) if self._positive_adv: advantage_flat -= advantage_flat.min() return advantage_flat
def test_compute_advantages(self, num_eps, discount, gae_lambda, rewards_eps, baselines_eps): """Test compute_advantage function.""" def get_advantage(discount, gae_lambda, rewards, baselines): adv = torch.zeros(rewards.shape) for i in range(rewards.shape[0]): acc = 0 for j in range(rewards.shape[1]): acc = acc * discount * gae_lambda acc += rewards[i][-j - 1] - baselines[i][-j - 1] acc += discount * baselines[i][-j] if j else 0 adv[i][-j - 1] = acc return adv length = len(rewards_eps) rewards = torch.Tensor(stack(num_eps, rewards_eps)) baselines = torch.Tensor(stack(num_eps, baselines_eps)) expected_adv = get_advantage(discount, gae_lambda, rewards, baselines) computed_adv = compute_advantages(discount, gae_lambda, length, baselines, rewards) assert torch.allclose(expected_adv, computed_adv)