def test_generalized_advantage_estimation(self):
values = tf.constant([[2.] * 5], tf.float32)
step_types = tf.constant([[StepType.MID] * 5], tf.int64)
rewards = tf.constant([[3.] * 5], tf.float32)
discounts = tf.constant([[0.9] * 5], tf.float32)
td_lambda = 0.6 / 0.9
d = 2 * 0.9 + 1
expected = tf.constant([[((d * 0.6 + d) * 0.6 + d) * 0.6 + d,
(d * 0.6 + d) * 0.6 + d, d * 0.6 + d, d]],
dtype=tf.float32)
self.assertAllClose(
value_ops.generalized_advantage_estimation(rewards=rewards,
values=values,
step_types=step_types,
discounts=discounts,
td_lambda=td_lambda,
time_major=False),
expected)
# two episodes, and exceed by time limit (discount=1)
step_types = tf.constant([[
StepType.MID, StepType.MID, StepType.LAST, StepType.MID,
StepType.MID
]], tf.int32)
expected = tf.constant([[d * 0.6 + d, d, 0, d]], dtype=tf.float32)
self.assertAllClose(
value_ops.generalized_advantage_estimation(rewards=rewards,
values=values,
step_types=step_types,
discounts=discounts,
td_lambda=td_lambda,
time_major=False),
expected)
# two episodes, and end normal (discount=0)
step_types = tf.constant([[
StepType.MID, StepType.MID, StepType.LAST, StepType.MID,
StepType.MID
]], tf.int32)
discounts = tf.constant([[0.9, 0.9, 0.0, 0.9, 0.9]])
expected = tf.constant([[1 * 0.6 + d, 1, 0, d]], dtype=tf.float32)
self.assertAllClose(
value_ops.generalized_advantage_estimation(rewards=rewards,
values=values,
step_types=step_types,
discounts=discounts,
td_lambda=td_lambda,
time_major=False),
expected)
def _calc_returns_and_advantages(self, training_info, value):
if self._use_vtrace:
return value_ops.calc_vtrace_returns_and_advantages(
training_info, value, self._gamma, self._action_spec,
self._lambda, self._debug_summaries)
returns = value_ops.discounted_return(
rewards=training_info.reward,
values=value,
step_types=training_info.step_type,
discounts=training_info.discount * self._gamma)
returns = common.tensor_extend(returns, value[-1])
if not self._use_gae:
advantages = returns - value
else:
advantages = value_ops.generalized_advantage_estimation(
rewards=training_info.reward,
values=value,
step_types=training_info.step_type,
discounts=training_info.discount * self._gamma,
td_lambda=self._lambda)
advantages = common.tensor_extend_zero(advantages)
if self._use_td_lambda_return:
returns = advantages + value
return returns, advantages
def test_vtrace_returns_and_advantages_impl_on_policy_no_last_step(self):
"""Test vtrace_returns_and_advantages_impl on policy no last_step
in the middle of the trajectory.
"""
importance_ratio_clipped = tf.constant([[1.] * 5], tf.float32)
values = tf.constant([[2.] * 5], tf.float32)
step_types = tf.constant([[StepType.MID] * 5], tf.int64)
rewards = tf.constant([[3.] * 5], tf.float32)
discounts = tf.constant([[0.9] * 5], tf.float32)
td_lambda = 1.0
returns, advantages = value_ops.vtrace_returns_and_advantages_impl(
importance_ratio_clipped,
rewards,
values,
step_types,
discounts,
time_major=False)
sa_returns, sa_adv = vtrace_scalable_agent(importance_ratio_clipped,
discounts, rewards, values,
step_types)
self.assertAllClose(sa_adv,
advantages,
msg='advantages differ from scalable_agent')
self.assertAllClose(sa_returns,
returns,
msg='returns differ from scalable_agent')
expected_advantages = value_ops.generalized_advantage_estimation(
rewards=rewards,
values=values,
step_types=step_types,
discounts=discounts,
td_lambda=td_lambda,
time_major=False)
expected_advantages = tf.transpose(a=expected_advantages)
expected_advantages = common.tensor_extend_zero(expected_advantages)
expected_advantages = tf.transpose(a=expected_advantages)
self.assertAllClose(expected_advantages,
advantages,
msg='advantages differ from gold')
expected_returns = value_ops.discounted_return(rewards=rewards,
values=values,
step_types=step_types,
discounts=discounts,
time_major=False)
expected_returns = tf.transpose(a=expected_returns)
values = tf.transpose(a=values)
expected_returns = common.tensor_extend(expected_returns, values[-1])
expected_returns = tf.transpose(a=expected_returns)
self.assertAllClose(expected_returns,
returns,
msg='returns differ from gold')
def _check(self, rewards, values, step_types, discounts, td_lambda,
expected):
np.testing.assert_array_almost_equal(
value_ops.generalized_advantage_estimation(rewards=rewards,
values=values,
step_types=step_types,
discounts=discounts,
td_lambda=td_lambda,
time_major=False),
expected)
np.testing.assert_array_almost_equal(
value_ops.generalized_advantage_estimation(
rewards=torch.stack([rewards, 2 * rewards], dim=2),
values=torch.stack([values, 2 * values], dim=2),
step_types=step_types,
discounts=discounts,
td_lambda=td_lambda,
time_major=False),
torch.stack([expected, 2 * expected], dim=2),
decimal=5)
def preprocess_experience(self, exp: Experience):
"""Compute advantages and put it into exp.info."""
advantages = value_ops.generalized_advantage_estimation(
rewards=exp.reward,
values=exp.info.value,
step_types=exp.step_type,
discounts=exp.discount * self._loss._gamma,
td_lambda=self._loss._lambda,
time_major=False)
advantages = tf.concat([
advantages,
tf.zeros(shape=common.concat_shape(tf.shape(advantages)[:-1], [1]),
dtype=advantages.dtype)
],
axis=-1)
returns = exp.info.value + advantages
return exp._replace(info=PPOInfo(returns, advantages))
def test_vtrace_returns_and_advantages_impl_on_policy_has_last_step(self):
"""Test vtrace_returns_and_advantages_impl on policy has last_step
in the middle of the trajectory.
"""
importance_ratio_clipped = tf.constant([[1.] * 5], tf.float32)
values = tf.constant([[2., 2.1, 2.2, 2.3, 2.4]], tf.float32)
step_types = tf.constant([[
StepType.MID, StepType.MID, StepType.LAST, StepType.MID,
StepType.MID
]], tf.int32)
rewards = tf.constant([[3., 3.1, 3.2, 3.3, 3.4]], tf.float32)
discounts = tf.constant([[0.9, 0.9, 0.0, 0.9, 0.9]])
td_lambda = 1.0
returns, advantages = value_ops.vtrace_returns_and_advantages_impl(
importance_ratio_clipped,
rewards,
values,
step_types,
discounts,
time_major=False)
expected_advantages = value_ops.generalized_advantage_estimation(
rewards=rewards,
values=values,
step_types=step_types,
discounts=discounts,
td_lambda=td_lambda,
time_major=False)
expected_advantages = tf.transpose(a=expected_advantages)
expected_advantages = common.tensor_extend_zero(expected_advantages)
expected_advantages = tf.transpose(a=expected_advantages)
self.assertAllClose(expected_advantages,
advantages,
msg='advantages differ')
expected_returns = value_ops.discounted_return(rewards=rewards,
values=values,
step_types=step_types,
discounts=discounts,
time_major=False)
expected_returns = tf.transpose(a=expected_returns)
values = tf.transpose(a=values)
expected_returns = common.tensor_extend(expected_returns, values[-1])
expected_returns = tf.transpose(a=expected_returns)
self.assertAllClose(expected_returns, returns, msg='returns differ')
def preprocess_experience(self, exp: Experience):
"""Compute advantages and put it into exp.rollout_info."""
advantages = value_ops.generalized_advantage_estimation(
rewards=exp.reward,
values=exp.rollout_info.value,
step_types=exp.step_type,
discounts=exp.discount * self._loss._gamma,
td_lambda=self._loss._lambda,
time_major=False)
advantages = torch.cat([
advantages,
torch.zeros(*advantages.shape[:-1], 1, dtype=advantages.dtype)
],
dim=-1)
returns = exp.rollout_info.value + advantages
return exp._replace(rollout_info=PPOInfo(
exp.rollout_info.action_distribution, returns, advantages))
def preprocess_experience(self, exp: Experience):
"""Compute advantages and put it into exp.info."""
reward = exp.reward
if self._algorithm._reward_shaping_fn is not None:
reward = self._algorithm._reward_shaping_fn(reward)
reward = self._algorithm.calc_training_reward(reward, exp.info)
advantages = value_ops.generalized_advantage_estimation(
rewards=reward,
values=exp.info.value,
step_types=exp.step_type,
discounts=exp.discount * self._algorithm._loss._gamma,
td_lambda=self._algorithm._loss._lambda,
time_major=False)
advantages = tf.concat([
advantages,
tf.zeros(advantages.shape.as_list()[:-1] + [1],
dtype=advantages.dtype)
],
axis=-1)
returns = exp.info.value + advantages
return exp._replace(info=PPOInfo(returns, advantages))
def _calc_returns_and_advantages(self, experience, value):
returns = value_ops.discounted_return(rewards=experience.reward,
values=value,
step_types=experience.step_type,
discounts=experience.discount *
self._gamma)
returns = tensor_utils.tensor_extend(returns, value[-1])
if not self._use_gae:
advantages = returns - value
else:
advantages = value_ops.generalized_advantage_estimation(
rewards=experience.reward,
values=value,
step_types=experience.step_type,
discounts=experience.discount * self._gamma,
td_lambda=self._lambda)
advantages = tensor_utils.tensor_extend_zero(advantages)
if self._use_td_lambda_return:
returns = advantages + value
return returns, advantages
def forward(self, experience, value, target_value):
"""Cacluate the loss.
The first dimension of all the tensors is time dimension and the second
dimesion is the batch dimension.
Args:
experience (Experience): experience collected from ``unroll()`` or
a replay buffer. All tensors are time-major.
value (torch.Tensor): the time-major tensor for the value at each time
step. The loss is between this and the calculated return.
target_value (torch.Tensor): the time-major tensor for the value at
each time step. This is used to calculate return. ``target_value``
can be same as ``value``.
Returns:
LossInfo: with the ``extra`` field same as ``loss``.
"""
if self._lambda == 1.0:
returns = value_ops.discounted_return(
rewards=experience.reward,
values=target_value,
step_types=experience.step_type,
discounts=experience.discount * self._gamma)
elif self._lambda == 0.0:
returns = value_ops.one_step_discounted_return(
rewards=experience.reward,
values=target_value,
step_types=experience.step_type,
discounts=experience.discount * self._gamma)
else:
advantages = value_ops.generalized_advantage_estimation(
rewards=experience.reward,
values=target_value,
step_types=experience.step_type,
discounts=experience.discount * self._gamma,
td_lambda=self._lambda)
returns = advantages + target_value[:-1]
value = value[:-1]
if self._debug_summaries and alf.summary.should_record_summaries():
mask = experience.step_type[:-1] != StepType.LAST
with alf.summary.scope(self._name):
def _summarize(v, r, td, suffix):
alf.summary.scalar(
"explained_variance_of_return_by_value" + suffix,
tensor_utils.explained_variance(v, r, mask))
safe_mean_hist_summary('values' + suffix, v, mask)
safe_mean_hist_summary('returns' + suffix, r, mask)
safe_mean_hist_summary("td_error" + suffix, td, mask)
if value.ndim == 2:
_summarize(value, returns, returns - value, '')
else:
td = returns - value
for i in range(value.shape[2]):
suffix = '/' + str(i)
_summarize(value[..., i], returns[..., i], td[..., i],
suffix)
loss = self._td_error_loss_fn(returns.detach(), value)
if loss.ndim == 3:
# Multidimensional reward. Average over the critic loss for all dimensions
loss = loss.mean(dim=2)
# The shape of the loss expected by Algorith.update_with_gradient is
# [T, B], so we need to augment it with additional zeros.
loss = tensor_utils.tensor_extend_zero(loss)
return LossInfo(loss=loss, extra=loss)