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)