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])