def step_fn(actor_output):
"""Per-replica StepFn."""
actor_output = tf.nest.pack_sequence_as(specs, actor_output)
(initial_agent_state, env_output, actor_agent_output, actor_action,
loss_type, info) = actor_output
with tf.GradientTape() as tape:
loss = loss_fns.compute_loss(
study_loss_types=study_loss_types,
current_batch_loss_type=loss_type,
agent=agent,
agent_state=initial_agent_state,
env_output=env_output,
actor_agent_output=actor_agent_output,
actor_action=actor_action,
num_steps=iterations)
grads = tape.gradient(loss, agent.trainable_variables)
if FLAGS.gradient_clip_norm > 0.:
for i, g in enumerate(grads):
if g is not None:
grads[i] = tf.clip_by_norm(g, FLAGS.gradient_clip_norm)
grad_norms = {}
for var, grad in zip(agent.trainable_variables, grads):
# For parameters which are initialized but not used for loss
# computation, gradient tape would return None.
if grad is not None:
grad_norms[var.name] = tf.norm(grad)
optimizer.apply_gradients(zip(grads, agent.trainable_variables))
return info, grad_norms
def _clip_gradients(self, grads):
clipped_grads = []
if self.clipnorm and self.clipnorm > 0:
for g in grads:
if g is None:
clipped_grads.append(g)
else:
clipped_grads.append(tf.clip_by_norm(g, self.clipnorm))
return clipped_grads
if self.global_clipnorm and self.global_clipnorm > 0:
return tf.clip_by_global_norm(grads, self.global_clipnorm)[0]
if self.clipvalue and self.clipvalue > 0:
for g in grads:
if g is None:
clipped_grads.append(g)
else:
clipped_grads.append(
tf.clip_by_value(
g,
clip_value_min=-self.clipvalue, # pylint: disable=invalid-unary-operand-type
clip_value_max=self.clipvalue))
return clipped_grads
return grads
def _clip_gradients(self, grads):
if not self._gradients_clip_option:
return grads
clipped_grads = []
if self._gradients_clip_option.clipnorm:
for g in grads:
if g is None:
clipped_grads.append(g)
else:
clipped_grads.append(
tf.clip_by_norm(g,
self._gradients_clip_option.clipnorm))
return clipped_grads
if self._gradients_clip_option.global_clipnorm:
return tf.clip_by_global_norm(
grads, self._gradients_clip_option.global_clipnorm)[0]
if self._gradients_clip_option.clipvalue:
for g in grads:
if g is None:
clipped_grads.append(g)
else:
clipped_grads.append(
tf.clip_by_value(
g,
clip_value_min=(
-self._gradients_clip_option.clipvalue),
clip_value_max=self._gradients_clip_option.
clipvalue))
return clipped_grads
return grads
def get_gradients(self, loss, params):
"""Returns gradients of `loss` with respect to `params`.
Args:
loss: Loss tensor.
params: List of variables.
Returns:
List of gradient tensors.
Raises:
ValueError: In case any gradient cannot be computed (e.g. if gradient
function not implemented).
"""
grads = backend.gradients(loss, params)
if any(g is None for g in grads):
raise ValueError('An operation has `None` for gradient. '
'Please make sure that all of your ops have a '
'gradient defined (i.e. are differentiable). '
'Common ops without gradient: '
'backend.argmax, backend.round, backend.eval.')
if hasattr(self, 'clipnorm'):
grads = [tf.clip_by_norm(g, self.clipnorm) for g in grads]
if hasattr(self, 'clipvalue'):
grads = [
tf.clip_by_value(g, -self.clipvalue, self.clipvalue)
for g in grads
]
return grads
def _clip_gradients(self, grads):
clipped_grads = []
if self.clipnorm and self.clipnorm > 0:
for g in grads:
if g is None:
clipped_grads.append(g)
else:
clipped_grads.append(tf.clip_by_norm(g, self.clipnorm))
return clipped_grads
if self.global_clipnorm and self.global_clipnorm > 0:
return tf.clip_by_global_norm(grads, self.global_clipnorm)[0]
if self.clipvalue and self.clipvalue > 0:
for g in grads:
if g is None:
clipped_grads.append(g)
else:
clipped_grads.append(
tf.clip_by_value(
g,
clip_value_min=-self.clipvalue,
clip_value_max=self.clipvalue,
))
return clipped_grads
return grads
def gradient_clipnorm_fn(grads_and_vars):
if isinstance(
tf.distribute.get_strategy(),
(tf.distribute.experimental.CentralStorageStrategy,
tf.compat.v1.distribute.experimental.CentralStorageStrategy)):
raise ValueError(
"`clipnorm` is not supported with `CenteralStorageStrategy`")
clipped_grads_and_vars = [(tf.clip_by_norm(g, clipnorm), v)
for g, v in grads_and_vars]
return clipped_grads_and_vars
def project(x, c):
"""Projects points to the Poincare ball.
Args:
x: Tensor of size B x dimension.
c: Tensor of size 1 representing the absolute hyperbolic curvature.
Returns:
Tensor of shape B x dimension where each row is a point that lies within
the Poincare ball.
"""
eps = BALL_EPS[x.dtype]
return tf.clip_by_norm(t=x, clip_norm=(1. - eps) / tf.sqrt(c), axes=[1])
