# coding=utf-8

# Copyright 2020 The TF-Agents Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# https://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

"""Common utilities for TF-Agents.

Example of usage:

```python

from tf_agents.utils import eager_utils

@eager_utils.run_in_graph_and_eager_modes

def loss_fn(x, y):

v = tf.get_variable('v', initializer=tf.ones_initializer(), shape=())

return v + x - y

with tfe.graph_mode():

# loss and train_step are Tensors/Ops in the graph

loss_op = loss_fn(inputs, labels)

train_step_op = eager_utils.create_train_step(loss_op, optimizer)

# Compute the loss and apply gradients to the variables using the optimizer.

with tf.Session() as sess:

sess.run(tf.compat.v1.global_variables_initializer())

for _ in range(num_train_steps):

loss_value = sess.run(train_step_op)

with tfe.eager_mode():

# loss and train_step are lambda functions that can be called.

loss = loss_fn(inputs, labels)

train_step = eager_utils.create_train_step(loss, optimizer)

# Compute the loss and apply gradients to the variables using the optimizer.

for _ in range(num_train_steps):

loss_value = train_step()

```

"""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import copy

import functools

import inspect

from absl import logging

import numpy as np

import six

import tensorflow as tf # pylint: disable=g-explicit-tensorflow-version-import

from tf_agents.utils import common

from tensorflow.python.util import tf_decorator # pylint:disable=g-direct-tensorflow-import # TF internal

_USE_GLOBAL_STEP = 0

def has_self_cls_arg(func_or_method):

"""Checks if it is method which takes self/cls as the first argument."""

if isinstance(func_or_method, staticmethod):

return False

if inspect.ismethod(func_or_method):

return True

if isinstance(func_or_method, classmethod):

return True

if six.PY2:

arg_names = inspect.getargspec(func_or_method).args

else:

arg_names = list(inspect.signature(func_or_method).parameters)

if arg_names and arg_names[0] in ('self', 'cls'):

return True

return False

def is_unbound(method):

"""Checks if it is an unbounded method."""

return not (hasattr(method, '__self__') and method.__self__)

class Future(object):

"""Converts a function or class method call into a future callable."""

def __init__(self, func_or_method, *args, **kwargs):

self._func_or_method = func_or_method

self._args = args

self._kwargs = copy.copy(kwargs)

getargspec = inspect.getargspec if six.PY2 else inspect.getfullargspec

arg_names = getargspec(func_or_method).args

self._arg_names = arg_names

self._self = None

if has_self_cls_arg(func_or_method):

# Skip the first arg_name self/cls

self._arg_names = arg_names[1:]

if is_unbound(func_or_method):

# For unbound methods we require self/cls as args[0].

if not args:

raise ValueError(

'func_or_method is unbond, but not class/instance provided.'

)

else:

self._self = args[0]

self._args = args[1:]

def __call__(self, *args, **kwargs):

"""If *args/**kwargs are given they would replace those given at init.

Args:

*args: List of extra arguments.

**kwargs: Dict of extra keyword arguments.

Returns:

The result of func_or_method(*args, **kwargs).

"""

# By default use the init args.

call_args = args or self._args

call_kwargs = copy.copy(self._kwargs)

for arg_name in self._arg_names[: len(args)]:

# Remove any original kwargs replaced by new positional args.

call_kwargs.pop(arg_name, None)

call_kwargs.update(kwargs)

if self._self:

return self._func_or_method(self._self, *call_args, **call_kwargs)

return self._func_or_method(*call_args, **call_kwargs)

def future_in_eager_mode(func_or_method):

"""Decorator that allow a function/method to run in graph and in eager modes.

When applied in graph mode it calls the function and return its outputs.

When applied in eager mode it returns a lambda function that when called

returns the outputs.

```python

@eager_utils.future_in_eager_mode

def loss_fn(x):

v = tf.get_variable('v', initializer=tf.ones_initializer(), shape=())

return v + x

with context.graph_mode():

loss_op = loss_fn(inputs)

loss_value = sess.run(loss_op)

with context.eager_mode():

loss = loss_fn(inputs)

# Now loss is a Future callable.

loss_value = loss()

```

Args:

func_or_method: A function or method to decorate.

Returns:

Either the output ops of the function/method or a Future (lambda function).

"""

if not callable(func_or_method):

raise TypeError('func_or_method must be callable.')

def decorator(*args, **kwargs):

if tf.executing_eagerly():

return Future(func_or_method, *args, **kwargs)

else:

return func_or_method(*args, **kwargs)

return tf_decorator.make_decorator(func_or_method, decorator)

def add_variables_summaries(grads_and_vars, step):

"""Add summaries for variables.

Args:

grads_and_vars: A list of (gradient, variable) pairs.

step: Variable to use for summaries.

"""

with tf.name_scope('summarize_vars'):

for _, var in grads_and_vars:

if isinstance(var, tf.IndexedSlices):

var_values = var.values

else:

var_values = var

var_name = var.name.replace(':', '_')

tf.compat.v2.summary.histogram(

name=var_name + '_value', data=var_values, step=step

)

tf.compat.v2.summary.scalar(

name=var_name + '_value_norm',

data=tf.linalg.global_norm([var_values]),

step=step,

)

def add_gradients_summaries(grads_and_vars, step):

"""Add summaries to gradients.

Args:

grads_and_vars: A list of gradient to variable pairs (tuples).

step: Variable to use for summaries.

"""

with tf.name_scope('summarize_grads'):

for grad, var in grads_and_vars:

if grad is not None:

if isinstance(grad, tf.IndexedSlices):

grad_values = grad.values

else:

grad_values = grad

var_name = var.name.replace(':', '_')

tf.compat.v2.summary.histogram(

name=var_name + '_gradient', data=grad_values, step=step

)

tf.compat.v2.summary.scalar(

name=var_name + '_gradient_norm',

data=tf.linalg.global_norm([grad_values]),

step=step,

)

else:

logging.info('Var %s has no gradient', var.name)

def clip_gradient_norms(gradients_to_variables, max_norm):

"""Clips the gradients by the given value.

Args:

gradients_to_variables: A list of gradient to variable pairs (tuples).

max_norm: the maximum norm value.

Returns:

A list of clipped gradient to variable pairs.

"""

clipped_grads_and_vars = []

for grad, var in gradients_to_variables:

if grad is not None:

if isinstance(grad, tf.IndexedSlices):

tmp = tf.clip_by_norm(grad.values, max_norm)

grad = tf.IndexedSlices(tmp, grad.indices, grad.dense_shape)

else:

grad = tf.clip_by_norm(grad, max_norm)

clipped_grads_and_vars.append((grad, var))

return clipped_grads_and_vars

def clip_gradient_norms_fn(max_norm):

"""Returns a `transform_grads_fn` function for gradient clipping."""

def clip_norms(gradients_to_variables):

return clip_gradient_norms(gradients_to_variables, max_norm)

return clip_norms

def create_train_step(

loss,

optimizer,

global_step=_USE_GLOBAL_STEP,

total_loss_fn=None,

update_ops=None,

variables_to_train=None,

transform_grads_fn=None,

summarize_gradients=False,

gate_gradients=tf.compat.v1.train.Optimizer.GATE_OP,

aggregation_method=None,

check_numerics=True,

):

"""Creates a train_step that evaluates the gradients and returns the loss.

Args:

loss: A (possibly nested tuple of) `Tensor` or function representing the

loss.

optimizer: A tf.Optimizer to use for computing the gradients.

global_step: A `Tensor` representing the global step variable. If left as

`_USE_GLOBAL_STEP`, then tf.train.get_or_create_global_step() is used.

total_loss_fn: Function to call on loss value to access the final item to

minimize.

update_ops: An optional list of updates to execute. If `update_ops` is

`None`, then the update ops are set to the contents of the

`tf.GraphKeys.UPDATE_OPS` collection. If `update_ops` is not `None`, but

it doesn't contain all of the update ops in `tf.GraphKeys.UPDATE_OPS`, a

warning will be displayed.

variables_to_train: an optional list of variables to train. If None, it will

default to all tf.trainable_variables().

transform_grads_fn: A function which takes a single argument, a list of

gradient to variable pairs (tuples), performs any requested gradient

updates, such as gradient clipping or multipliers, and returns the updated

list.

summarize_gradients: Whether or not add summaries for each gradient.

gate_gradients: How to gate the computation of gradients. See tf.Optimizer.

aggregation_method: Specifies the method used to combine gradient terms.

Valid values are defined in the class `AggregationMethod`.

check_numerics: Whether or not we apply check_numerics.

Returns:

In graph mode: A (possibly nested tuple of) `Tensor` that when evaluated,

calculates the current loss, computes the gradients, applies the

optimizer, and returns the current loss.

In eager mode: A lambda function that when is called, calculates the loss,

then computes and applies the gradients and returns the original

loss values.

Raises:

ValueError: if loss is not callable.

RuntimeError: if resource variables are not enabled.

"""

if total_loss_fn is None:

total_loss_fn = lambda x: x

if not callable(total_loss_fn):

raise ValueError('`total_loss_fn` should be a function.')

if not common.resource_variables_enabled():

raise RuntimeError(common.MISSING_RESOURCE_VARIABLES_ERROR)

if not tf.executing_eagerly():

if callable(loss):

loss = loss()

if callable(variables_to_train):

variables_to_train = variables_to_train()

# Calculate loss first, then calculate train op, then return the original

# loss conditioned on executing the train op.

with tf.control_dependencies(tf.nest.flatten(loss)):

loss = tf.nest.map_structure(

lambda t: tf.identity(t, 'loss_pre_train'), loss

)

train_op = create_train_op(

total_loss_fn(loss),

optimizer,

global_step=global_step,

update_ops=update_ops,

variables_to_train=variables_to_train,

transform_grads_fn=transform_grads_fn,

summarize_gradients=summarize_gradients,

gate_gradients=gate_gradients,

aggregation_method=aggregation_method,

check_numerics=check_numerics,

)

with tf.control_dependencies([train_op]):

return tf.nest.map_structure(

lambda t: tf.identity(t, 'loss_post_train'), loss

)

if global_step is _USE_GLOBAL_STEP:

global_step = tf.compat.v1.train.get_or_create_global_step()

if not callable(loss):

raise ValueError('`loss` should be a function in eager mode.')

if not isinstance(loss, Future):

logging.warning('loss should be an instance of eager_utils.Future')

with tf.GradientTape() as tape:

loss_value = loss()

total_loss_value = total_loss_fn(loss_value)

if variables_to_train is None:

variables_to_train = tape.watched_variables()

elif callable(variables_to_train):

variables_to_train = variables_to_train()

variables_to_train = tf.nest.flatten(variables_to_train)

grads = tape.gradient(total_loss_value, variables_to_train)

grads_and_vars = list(zip(grads, variables_to_train))

if transform_grads_fn:

grads_and_vars = transform_grads_fn(grads_and_vars)

if summarize_gradients:

with tf.name_scope('summarize_grads'):

add_gradients_summaries(grads_and_vars, global_step)

if check_numerics:

with tf.name_scope('train_op'):

tf.debugging.check_numerics(total_loss_value, 'Loss is inf or nan')

optimizer.apply_gradients(grads_and_vars, global_step=global_step)

return loss_value

def create_train_op(

total_loss,

optimizer,

global_step=_USE_GLOBAL_STEP,

update_ops=None,

variables_to_train=None,

transform_grads_fn=None,

summarize_gradients=False,

gate_gradients=tf.compat.v1.train.Optimizer.GATE_OP,

aggregation_method=None,

check_numerics=True,

):

"""Creates an `Operation` that evaluates the gradients and returns the loss.

Args:

total_loss: A `Tensor` representing the total loss.

optimizer: A tf.Optimizer to use for computing the gradients.

global_step: A `Tensor` representing the global step variable. If left as

`_USE_GLOBAL_STEP`, then tf.train.get_or_create_global_step() is used.

update_ops: An optional list of updates to execute. If `update_ops` is

`None`, then the update ops are set to the contents of the

`tf.GraphKeys.UPDATE_OPS` collection. If `update_ops` is not `None`, but

it doesn't contain all of the update ops in `tf.GraphKeys.UPDATE_OPS`, a

warning will be displayed.

variables_to_train: an optional list of variables to train. If None, it will

default to all tf.trainable_variables().

transform_grads_fn: A function which takes a single argument, a list of

gradient to variable pairs (tuples), performs any requested gradient

updates, such as gradient clipping or multipliers, and returns the updated

list.

summarize_gradients: Whether or not add summaries for each gradient.

gate_gradients: How to gate the computation of gradients. See tf.Optimizer.

aggregation_method: Specifies the method used to combine gradient terms.

Valid values are defined in the class `AggregationMethod`.

check_numerics: Whether or not we apply check_numerics.

Returns:

A `Tensor` that when evaluated, computes the gradients and returns the total

loss value.

"""

if global_step is _USE_GLOBAL_STEP:

global_step = tf.compat.v1.train.get_or_create_global_step()

# Update ops use GraphKeys.UPDATE_OPS collection if update_ops is None.

global_update_ops = set(

tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)

)

if update_ops is None:

update_ops = global_update_ops

else:

update_ops = set(update_ops)

if not global_update_ops.issubset(update_ops):

tf.compat.v1.logging.warning(

'update_ops in create_train_op does not contain all the '

'update_ops in GraphKeys.UPDATE_OPS'

)

# Make sure update_ops are computed before total_loss.

if update_ops:

with tf.control_dependencies(update_ops):

barrier = tf.no_op(name='update_barrier')

with tf.control_dependencies([barrier]):

total_loss = tf.identity(total_loss)

if variables_to_train is None:

# Default to tf.trainable_variables()

variables_to_train = tf.compat.v1.trainable_variables()

else:

# Make sure that variables_to_train are in tf.trainable_variables()

trainable_variables = tf.compat.v1.trainable_variables()

for v in variables_to_train:

assert v.trainable or v in trainable_variables

assert variables_to_train

# Create the gradients. Note that apply_gradients adds the gradient

# computation to the current graph.

grads = optimizer.compute_gradients(

total_loss,

variables_to_train,

gate_gradients=gate_gradients,

aggregation_method=aggregation_method,

)

if transform_grads_fn:

grads = transform_grads_fn(grads)

# Summarize gradients.

if summarize_gradients:

with tf.name_scope('summarize_grads'):

add_gradients_summaries(grads, global_step)

# Create gradient updates.

grad_updates = optimizer.apply_gradients(grads, global_step=global_step)

with tf.name_scope('train_op'):

# Make sure total_loss is valid.

if check_numerics:

total_loss = tf.debugging.check_numerics(

total_loss, 'LossTensor is inf or nan'

)

# Ensure the train_tensor computes grad_updates.

with tf.control_dependencies([grad_updates]):

train_op = tf.identity(total_loss, name='train_op')

return train_op

def np_function(func=None, output_dtypes=None):

"""Decorator that allow a numpy function to be used in Eager and Graph modes.

Similar to `tf.py_func` and `tf.py_function` but it doesn't require defining

the inputs or the dtypes of the outputs a priori.

In Eager mode it would convert the tf.Tensors to np.arrays before passing to

`func` and then convert back the outputs from np.arrays to tf.Tensors.

In Graph mode it would create different tf.py_function for each combination

of dtype of the inputs and cache them for reuse.

NOTE: In Graph mode: if `output_dtypes` is not provided then `func` would

be called with `np.ones()` to infer the output dtypes, and therefore `func`

should be stateless.

```python

Instead of doing:

def sum(x):

return np.sum(x)

inputs = tf.constant([3, 4])

outputs = tf.py_function(sum, inputs, Tout=[tf.int64])

inputs = tf.constant([3., 4.])

outputs = tf.py_function(sum, inputs, Tout=[tf.float32])

Do:

@eager_utils.np_function

def sum(x):

return np.sum(x)

inputs = tf.constant([3, 4])

outputs = sum(inputs) # Infers that Tout is tf.int64

inputs = tf.constant([3., 4.])

outputs = sum(inputs) # Infers that Tout is tf.float32

# Output dtype is always float32 for valid input dtypes.

@eager_utils.np_function(output_dtypes=np.float32)

def mean(x):

return np.mean(x)

# Output dtype depends on the input dtype.

@eager_utils.np_function(output_dtypes=lambda x: (x, x))

def repeat(x):

return x, x

with context.graph_mode():

outputs = sum(tf.constant([3, 4]))

outputs2 = sum(tf.constant([3., 4.]))

sess.run(outputs) # np.array(7)

sess.run(outputs2) # np.array(7.)

with context.eager_mode():

inputs = tf.constant([3, 4])

outputs = sum(tf.constant([3, 4])) # tf.Tensor([7])

outputs = sum(tf.constant([3., 4.])) # tf.Tensor([7.])

```

Args:

func: A numpy function, that takes numpy arrays as inputs and return numpy

arrays as outputs.

output_dtypes: Optional list of dtypes or a function that maps input dtypes

to output dtypes. Examples: output_dtypes=[tf.float32],

output_dtypes=lambda x: x (outputs have the same dtype as inputs). If it

is not provided in Graph mode the `func` would be called to infer the

output dtypes.

Returns:

A wrapped function that can be used with TF code.

"""

def decorated(func):

"""Decorated func."""

dtype_map = {}

def wrapper(*args, **kwargs):

"""Wrapper to add nested input and outputs support."""

func_with_kwargs = functools.partial(func, **kwargs)

def func_flat_outputs(*args):

return tf.nest.flatten(func_with_kwargs(*args))

def compute_output_dtypes(*args):

"""Calls the func to compute output dtypes."""

result = func(*args, **kwargs)

return tf.nest.map_structure(lambda x: x.dtype, result)

if tf.executing_eagerly():

result = func_with_kwargs(

*tf.nest.map_structure(lambda x: x.numpy(), args)

)

convert = lambda x: x if x is None else tf.convert_to_tensor(value=x)

return tf.nest.map_structure(convert, result)

else:

input_dtypes = tuple([x.dtype for x in tf.nest.flatten(args)])

if input_dtypes not in dtype_map:

if output_dtypes is None:

dummy_args = tf.nest.map_structure(

lambda x: np.ones(x.shape, x.dtype.as_numpy_dtype), args

)

dtype_map[input_dtypes] = compute_output_dtypes(*dummy_args)

elif isinstance(output_dtypes, (list, tuple)):

# output_dtypes define the output dtypes.

dtype_map[input_dtypes] = output_dtypes

else:

try:

# See if output_dtypes define the output dtype directly.

tf.as_dtype(output_dtypes)

dtype_map[input_dtypes] = output_dtypes

except TypeError:

if callable(output_dtypes):

# output_dtypes is mapping from input_dtypes to output_dtypes.

dtype_map[input_dtypes] = output_dtypes(*input_dtypes)

else:

raise ValueError(

'output_dtypes not a list of dtypes or a callable.'

)

flat_output_dtypes = tf.nest.flatten(dtype_map[input_dtypes])

flat_outputs = tf.py_function(

func_flat_outputs, inp=args, Tout=flat_output_dtypes

)

return tf.nest.pack_sequence_as(dtype_map[input_dtypes], flat_outputs)

return tf_decorator.make_decorator(func, wrapper)

# This code path is for the `foo = np_function(foo, ...)` use case

if func is not None:

return decorated(func)

# This code path is for the decorator

# @np_function(...)

# def foo(...):

return decorated

def dataset_iterator(dataset):

"""Constructs a `Dataset` iterator.

The method used to construct the iterator is conditioned on whether Graph mode

is enabled. `dataset_iterator` and `get_next` are useful when we need to

construct an iterator and iterate through it inside a `tensorflow.function`.

Args:

dataset: a `tf.data.Dataset`.

Returns:

A `tf.data.Iterator` if Graph mode is enabled; a tf.data.EagerIterator if

in eager mode.

"""

if tf.executing_eagerly():

return iter(dataset)

try:

iterator = tf.compat.v1.data.make_one_shot_iterator(dataset)

except ValueError:

iterator = tf.compat.v1.data.make_initializable_iterator(dataset)

return iterator

def get_next(iterator):

"""Returns the next element in a `Dataset` iterator.

The syntax used to retrieve the next item is conditioned on whether Graph mode

is enabled. `dataset_iterator` and `get_next` are useful when we need to

construct an iterator and iterate through it inside a `tensorflow.function`.

Args:

iterator: a `tf.data.Iterator` if in Graph mode; a `tf.data.EagerIterator`

if in eager mode.

Returns:

A `tf.data.Iterator` if Graph mode is enabled; a Python iterator if in eager

mode.

"""

if tf.executing_eagerly():

return next(iterator)

return iterator.get_next()