# 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.

"""Base extension to Keras network to simplify copy operations."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import abc

import typing

import six

import tensorflow as tf

import tensorflow_probability as tfp

from tf_agents.distributions import utils as distribution_utils

from tf_agents.networks import layer_utils

from tf_agents.specs import tensor_spec

from tf_agents.trajectories import time_step

from tf_agents.typing import types

from tf_agents.utils import common

from tf_agents.utils import nest_utils

from tf_agents.utils import object_identity

# pylint: disable=g-direct-tensorflow-import

from tensorflow.python.trackable import base # TF internal

from tensorflow.python.util import tf_decorator # TF internal

from tensorflow.python.util import tf_inspect # TF internal

# pylint: enable=g-direct-tensorflow-import

_ = tf.keras.layers # Force loading of keras layers.

class _NetworkMeta(abc.ABCMeta):

"""Meta class for Network object.

We mainly use this class to capture all args to `__init__` of all `Network`

instances, and store them in `instance._saved_kwargs`. This in turn is

used by the `instance.copy` method.

"""

def __new__(mcs, classname, baseclasses, attrs):

"""Control the creation of subclasses of the Network class.

Args:

classname: The name of the subclass being created.

baseclasses: A tuple of parent classes.

attrs: A dict mapping new attributes to their values.

Returns:

The class object.

Raises:

RuntimeError: if the class __init__ has *args in its signature.

"""

if baseclasses[0] == tf.keras.layers.Layer:

# This is just Network below. Return early.

return abc.ABCMeta.__new__(mcs, classname, baseclasses, attrs)

init = attrs.get("__init__", None)

if not init:

# This wrapper class does not define an __init__. When someone creates

# the object, the __init__ of its parent class will be called. We will

# call that __init__ instead separately since the parent class is also a

# subclass of Network. Here just create the class and return.

return abc.ABCMeta.__new__(mcs, classname, baseclasses, attrs)

arg_spec = tf_inspect.getargspec(init)

if arg_spec.varargs is not None:

raise RuntimeError(

"%s.__init__ function accepts *args. This is not allowed."

% classname

)

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

"""Captures init args and kwargs and stores them into `_saved_kwargs`."""

if len(args) > len(arg_spec.args) + 1:

# Error case: more inputs than args. Call init so that the appropriate

# error can be raised to the user.

init(self, *args, **kwargs)

# Convert to a canonical kwarg format.

kwargs = tf_inspect.getcallargs(init, self, *args, **kwargs)

kwargs.pop("self")

init(self, **kwargs)

# Avoid auto tracking which prevents keras from tracking layers that are

# passed as kwargs to the Network.

with base.no_automatic_dependency_tracking_scope(self):

setattr(self, "_saved_kwargs", kwargs)

attrs["__init__"] = tf_decorator.make_decorator(init, _capture_init)

return abc.ABCMeta.__new__(mcs, classname, baseclasses, attrs)

@six.add_metaclass(_NetworkMeta)

class Network(tf.keras.layers.Layer):

"""A class used to represent networks used by TF-Agents policies and agents.

The main differences between a TF-Agents Network and a Keras Layer include:

networks keep track of their underlying layers, explicitly represent RNN-like

state in inputs and outputs, and simplify variable creation and clone

operations.

When calling a network `net`, typically one passes data through it via:

```python

outputs, next_state = net(observation, network_state=...)

outputs, next_state = net(observation, step_type=..., network_state=...)

outputs, next_state = net(observation) # net.call must fill an empty state

outputs, next_state = net(observation, step_type=...)

outputs, next_state = net(

observation, step_type=..., network_state=..., learning=...)

```

etc.

To force construction of a network's variables:

```python

net.create_variables()

net.create_variables(input_tensor_spec=...) # To provide an input spec

net.create_variables(training=True) # Provide extra kwargs

net.create_variables(input_tensor_spec, training=True)

```

To create a copy of the network:

```python

cloned_net = net.copy()

cloned_net.variables # Raises ValueError: cloned net does not share weights.

cloned_net.create_variables(...)

cloned_net.variables # Now new variables have been created.

```

"""

# TODO(b/156314975): Rename input_tensor_spec to input_spec.

def __init__(self, input_tensor_spec=None, state_spec=(), name=None):

"""Creates an instance of `Network`.

Args:

input_tensor_spec: A nest of `tf.TypeSpec` representing the input

observations. Optional. If not provided, `create_variables()` will

fail unless a spec is provided.

state_spec: A nest of `tensor_spec.TensorSpec` representing the state

needed by the network. Default is `()`, which means no state.

name: (Optional.) A string representing the name of the network.

"""

# Disable autocast because it may convert bfloats to other types, breaking

# our spec checks.

super(Network, self).__init__(name=name, autocast=False)

common.check_tf1_allowed()

# Required for summary() to work.

self._is_graph_network = False

self._input_tensor_spec = (

tensor_spec.from_spec(input_tensor_spec)

if input_tensor_spec is not None

else None

)

# NOTE(ebrevdo): Would have preferred to call this output_tensor_spec, but

# looks like keras.Layer already reserves that one.

self._network_output_spec = None

self._state_spec = tensor_spec.from_spec(state_spec)

@property

def state_spec(self):

return self._state_spec

@property

def input_tensor_spec(self):

"""Returns the spec of the input to the network of type InputSpec."""

return self._input_tensor_spec

def create_variables(self, input_tensor_spec=None, **kwargs):

"""Force creation of the network's variables.

Return output specs.

Args:

input_tensor_spec: (Optional). Override or provide an input tensor spec

when creating variables.

**kwargs: Other arguments to `network.call()`, e.g. `training=True`.

Returns:

Output specs - a nested spec calculated from the outputs (excluding any

batch dimensions). If any of the output elements is a tfp `Distribution`,

the associated spec entry returned is a `DistributionSpec`.

Raises:

ValueError: If no `input_tensor_spec` is provided, and the network did

not provide one during construction.

"""

if self._network_output_spec is not None:

return self._network_output_spec

if self._input_tensor_spec is None:

self._input_tensor_spec = input_tensor_spec

input_tensor_spec = self._input_tensor_spec

if input_tensor_spec is None:

raise ValueError(

"Unable to create_variables: no input_tensor_spec provided, and "

"Network did not define one."

)

random_input = tensor_spec.sample_spec_nest(

input_tensor_spec, outer_dims=(1,)

)

initial_state = self.get_initial_state(batch_size=1)

step_type = tf.fill((1,), time_step.StepType.FIRST)

outputs = self.__call__(

random_input, step_type=step_type, network_state=initial_state, **kwargs

)

def _calc_unbatched_spec(x):

"""Build Network output spec by removing previously added batch dimension.

Args:

x: tfp.distributions.Distribution or Tensor.

Returns:

Specs without batch dimension representing x.

"""

if isinstance(x, tfp.distributions.Distribution):

parameters = distribution_utils.get_parameters(x)

parameter_specs = _convert_to_spec_and_remove_singleton_batch_dim(

parameters, outer_ndim=1

)

return distribution_utils.DistributionSpecV2(

event_shape=x.event_shape, dtype=x.dtype, parameters=parameter_specs

)

else:

return tensor_spec.remove_outer_dims_nest(

tf.type_spec_from_value(x), num_outer_dims=1

)

self._network_output_spec = tf.nest.map_structure(

_calc_unbatched_spec, outputs[0]

)

return self._network_output_spec

@property

def variables(self):

if not self.built:

raise ValueError(

f"Network `{self.name}` {type(self)} has not been built,"

" unable to access variables."

" Please call `create_variables` or apply the network first."

)

return super(Network, self).variables

@property

def trainable_variables(self):

if not self.built:

raise ValueError(

f"Network `{self.name}` {type(self)} has not been built,"

" unable to access trainable_variables."

" Please call `create_variables` or apply the network first."

)

return super(Network, self).trainable_variables

@property

def layers(self):

"""Get the list of all (nested) sub-layers used in this Network."""

return list(self._flatten_layers(include_self=False, recursive=False))

def get_layer(self, name=None, index=None):

"""Retrieves a layer based on either its name (unique) or index.

If `name` and `index` are both provided, `index` will take precedence.

Indices are based on order of horizontal graph traversal (bottom-up).

Args:

name: String, name of layer.

index: Integer, index of layer.

Returns:

A layer instance.

Raises:

ValueError: In case of invalid layer name or index.

"""

if index is not None and name is not None:

raise ValueError("Provide only a layer name or a layer index.")

if index is not None:

if len(self.layers) <= index:

raise ValueError(

"Was asked to retrieve layer at index "

+ str(index)

+ " but model only has "

+ str(len(self.layers))

+ " layers."

)

else:

return self.layers[index]

if name is not None:

for layer in self.layers:

if layer.name == name:

return layer

raise ValueError("No such layer: " + name + ".")

def summary(self, line_length=None, positions=None, print_fn=None):

"""Prints a string summary of the network.

Args:

line_length: Total length of printed lines (e.g. set this to adapt the

display to different terminal window sizes).

positions: Relative or absolute positions of log elements in each line.

If not provided, defaults to `[.33, .55, .67, 1.]`.

print_fn: Print function to use. Defaults to `print`. It will be called

on each line of the summary. You can set it to a custom function in

order to capture the string summary.

Raises:

ValueError: if `summary()` is called before the model is built.

"""

if not self.built:

raise ValueError(

"This model has not yet been built. "

"Build the model first by calling `build()` or "

"`__call__()` with some data, or `create_variables()`."

)

layer_utils.print_summary(

self, line_length=line_length, positions=positions, print_fn=print_fn

)

def copy(self, **kwargs):

"""Create a shallow copy of this network.

**NOTE** Network layer weights are *never* copied. This method recreates

the `Network` instance with the same arguments it was initialized with

(excepting any new kwargs).

Args:

**kwargs: Args to override when recreating this network. Commonly

overridden args include 'name'.

Returns:

A shallow copy of this network.

"""

return type(self)(**dict(self._saved_kwargs, **kwargs))

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

"""A wrapper around `Network.call`.

A typical `call` method in a class subclassing `Network` will have a

signature that accepts `inputs`, as well as other `*args` and `**kwargs`.

`call` can optionally also accept `step_type` and `network_state`

(if `state_spec != ()` is not trivial). e.g.:

```python

def call(self,

inputs,

step_type=None,

network_state=(),

training=False):

...

return outputs, new_network_state

```

We will validate the first argument (`inputs`)

against `self.input_tensor_spec` if one is available.

If a `network_state` kwarg is given it is also validated against

`self.state_spec`. Similarly, the return value of the `call` method is

expected to be a tuple/list with 2 values: `(output, new_state)`.

We validate `new_state` against `self.state_spec`.

If no `network_state` kwarg is given (or if empty `network_state = ()` is

given, it is up to `call` to assume a proper "empty" state, and to

emit an appropriate `output_state`.

Args:

inputs: The input to `self.call`, matching `self.input_tensor_spec`.

*args: Additional arguments to `self.call`.

**kwargs: Additional keyword arguments to `self.call`. These can include

`network_state` and `step_type`. `step_type` is required if the

network's `call` requires it. `network_state` is required if the

underlying network's `call` requires it.

Returns:

A tuple `(outputs, new_network_state)`.

"""

if self.input_tensor_spec is not None:

nest_utils.assert_matching_dtypes_and_inner_shapes(

inputs,

self.input_tensor_spec,

allow_extra_fields=True,

caller=self,

tensors_name="`inputs`",

specs_name="`input_tensor_spec`",

)

call_argspec = tf_inspect.getargspec(self.call)

# Convert *args, **kwargs to a canonical kwarg representation.

normalized_kwargs = tf_inspect.getcallargs(

self.call, inputs, *args, **kwargs

)

# TODO(b/156315434): Rename network_state to just state.

network_state = normalized_kwargs.get("network_state", None)

normalized_kwargs.pop("self", None)

if common.safe_has_state(network_state):

nest_utils.assert_matching_dtypes_and_inner_shapes(

network_state,

self.state_spec,

allow_extra_fields=True,

caller=self,

tensors_name="`network_state`",

specs_name="`state_spec`",

)

if "step_type" not in call_argspec.args and not call_argspec.keywords:

normalized_kwargs.pop("step_type", None)

# network_state can be a (), None, Tensor or NestedTensors.

if (

not tf.is_tensor(network_state)

and network_state in (None, ())

and "network_state" not in call_argspec.args

and not call_argspec.keywords

):

normalized_kwargs.pop("network_state", None)

outputs, new_state = super(Network, self).__call__(**normalized_kwargs) # pytype: disable=attribute-error # typed-keras

nest_utils.assert_matching_dtypes_and_inner_shapes(

new_state,

self.state_spec,

allow_extra_fields=True,

caller=self,

tensors_name="`new_state`",

specs_name="`state_spec`",

)

return outputs, new_state

def _check_trainable_weights_consistency(self):

"""Check trainable weights count consistency.

This method makes up for the missing method (b/143631010) of the same name

in `keras.Network`, which is needed when calling `Network.summary()`. This

method is a no op. If a Network wants to check the consistency of trainable

weights, see `keras.Model._check_trainable_weights_consistency` as a

reference.

"""

# TODO(b/143631010): If recognized and fixed, remove this entire method.

return

def get_initial_state(self, batch_size=None):

"""Returns an initial state usable by the network.

Args:

batch_size: Tensor or constant: size of the batch dimension. Can be None

in which case not dimensions gets added.

Returns:

A nested object of type `self.state_spec` containing properly

initialized Tensors.

"""

return self._get_initial_state(batch_size)

def _get_initial_state(self, batch_size):

"""Returns the initial state of the policy network.

Args:

batch_size: A constant or Tensor holding the batch size. Can be None, in

which case the state will not have a batch dimension added.

Returns:

A nest of zero tensors matching the spec of the policy network state.

"""

return tensor_spec.zero_spec_nest(

self._state_spec,

outer_dims=None if batch_size is None else [batch_size],

)

class DistributionNetwork(Network):

"""Base class for networks which generate Distributions as their output."""

def __init__(self, input_tensor_spec, state_spec, output_spec, name):

super(DistributionNetwork, self).__init__(

input_tensor_spec=input_tensor_spec, state_spec=state_spec, name=name

)

self._output_spec = output_spec

@property

def output_spec(self):

return self._output_spec

def _is_layer(obj):

"""Implicit check for Layer-like objects."""

# TODO(b/110718070): Replace with isinstance(obj, tf.keras.layers.Layer).

return hasattr(obj, "_is_layer") and not isinstance(obj, type)

def _filter_empty_layer_containers(layer_list):

"""Remove empty layer containers."""

existing = object_identity.ObjectIdentitySet()

to_visit = layer_list[::-1]

while to_visit:

obj = to_visit.pop()

if obj in existing:

continue

existing.add(obj)

if _is_layer(obj):

yield obj

else:

sub_layers = getattr(obj, "layers", None) or []

# Trackable data structures will not show up in ".layers" lists, but

# the layers they contain will.

to_visit.extend(sub_layers[::-1])

def _convert_to_spec_and_remove_singleton_batch_dim(

parameters: distribution_utils.Params, outer_ndim: int

) -> distribution_utils.Params:

"""Convert a `Params` object of tensors to one containing unbatched specs.

Note: The `Params` provided to this function are typically contain tensors

generated by Layers and therefore containing an outer singleton dimension.

Since TF-Agents specs exclude batch and time prefixes, here we need to

remove the singleton batch dimension from the specs created by these

input tensors.

Args:

parameters: Distribution parameters, including input tensors.

outer_ndim: Number of singleton outer dimensions expected in tensors found

in `parameters`.

Returns:

A `Params` object contanining `tf.TypeSpec` in place of tensors, with

up to `outer_ndim` outer singleton dimensions removed.

"""

def _maybe_convert_to_spec(p):

if isinstance(p, distribution_utils.Params):

return _convert_to_spec_and_remove_singleton_batch_dim(p, outer_ndim)

elif tf.is_tensor(p):

return tensor_spec.remove_outer_dims_nest(

tf.type_spec_from_value(p), num_outer_dims=outer_ndim

)

else:

return p

return distribution_utils.Params(

type_=parameters.type_,

params=tf.nest.map_structure(_maybe_convert_to_spec, parameters.params),

)

def create_variables(

module: typing.Union[Network, tf.keras.layers.Layer],

input_spec: typing.Optional[types.NestedTensorSpec] = None,

**kwargs: typing.Any,

) -> types.NestedTensorSpec:

"""Create variables in `module` given `input_spec`; return `output_spec`.

Here `module` can be a `tf_agents.networks.Network` or `Keras` layer.

Args:

module: The instance we would like to create layers on.

input_spec: The input spec (excluding batch dimensions).

**kwargs: Extra arguments to `module.__call__`, e.g. `training=True`.

Returns:

Output specs, a nested `tf.TypeSpec` describing the output signature.

If `module` returns a `tfp.Distribution`, then the associated nested

object is a `tf_agents.specs.DistributionSpecV2` (which is not a true

`tf.TypeSpec` but contains enough information to create a nested

`tf.TypeSpec` using `tf_agents.distributions.utils.parameters_to_dict`).

Raises:

ValueError: If `module` is a generic Keras layer but `input_spec is None`.

TypeError: If `module` is a `tf.keras.layers.{RNN,LSTM,GRU,...}`. These

must be wrapped in `tf_agents.keras_layers.RNNWrapper`.

"""

# NOTE(ebrevdo): As a side effect, for generic keras Layers (not Networks)

# this method stores new hidden properties in `module`:

# `_network_output_spec`, `_network_state_spec`,

# - which internal TF-Agents libraries make use of.

if isinstance(module, Network):

return module.create_variables(input_spec, **kwargs)

# Generic keras layer

if input_spec is None:

raise ValueError(

"Module is a Keras layer; an input_spec is required but saw "

"None: {}".format(module)

)

if isinstance(module, tf.keras.layers.RNN):

raise TypeError(

"Keras RNN layers (non-cell layers) must be wrapped in "

"tf_agents.keras_layers.RNNWrapper. Layer: {}".format(module)

)

maybe_spec = getattr(module, "_network_output_spec", None)

if maybe_spec is not None:

return maybe_spec

# Has state outputs.

recurrent_layer = getattr(module, "get_initial_state", None) is not None

# Required input rank

outer_ndim = _get_input_outer_ndim(module, input_spec)

random_input = tensor_spec.sample_spec_nest(

input_spec, outer_dims=(1,) * outer_ndim

)

if recurrent_layer:

state = module.get_initial_state(random_input)

def remove_singleton_batch_spec_dim(t):

# Convert tensor to its type-spec, and remove the batch dimension

# from the spec.

spec = tf.type_spec_from_value(t)

return nest_utils.remove_singleton_batch_spec_dim(spec, outer_ndim=1)

state_spec = tf.nest.map_structure(remove_singleton_batch_spec_dim, state)

outputs = module(random_input, state, **kwargs)

# tf.keras.layers.{LSTMCell, ...} return (output, [state1, state2,...]).

output = outputs[0]

else:

output = module(random_input, **kwargs)

state_spec = ()

def _calc_unbatched_spec(x):

if isinstance(x, tfp.distributions.Distribution):

parameters = distribution_utils.get_parameters(x)

parameter_specs = _convert_to_spec_and_remove_singleton_batch_dim(

parameters, outer_ndim=outer_ndim

)

return distribution_utils.DistributionSpecV2(

event_shape=x.event_shape, dtype=x.dtype, parameters=parameter_specs

)

else:

return tensor_spec.remove_outer_dims_nest(

tf.type_spec_from_value(x), num_outer_dims=outer_ndim

)

# pylint: disable=protected-access

module._network_output_spec = tf.nest.map_structure(

_calc_unbatched_spec, output

)

module._network_state_spec = state_spec

return module._network_output_spec

# pylint: enable=protected-access

def _get_input_outer_ndim(

layer: tf.keras.layers.Layer, input_spec: types.NestedTensorSpec

) -> int:

"""Calculate or guess the number of batch (outer) ndims in `layer`."""

if isinstance(layer, tf.keras.layers.RNN):

raise TypeError(

"Saw a tf.keras.layers.RNN layer nested inside e.g. a keras Sequential "

"layer. This is not directly supported. Please wrap your layer "

"inside a `tf_agents.keras_layers.RNNWrapper` or use "

"`tf_agents.networks.Sequential`. Layer: {}".format(layer)

)

if isinstance(layer, tf.keras.layers.TimeDistributed):

return 1 + _get_input_outer_ndim(layer.layer, input_spec)

if isinstance(layer, tf.keras.Sequential):

# We don't trust Sequential to give us the right thing if the first layer

# is e.g. a TimeDistributed.

return _get_input_outer_ndim(layer.layers[0], input_spec)

layer_input_spec = layer.input_spec

if layer_input_spec is None:

return 1

outer_ndim = layer_input_spec.ndim

if outer_ndim is None:

outer_ndim = layer_input_spec.min_ndim

if outer_ndim is None:

return 1

if input_spec:

input_spec = tf.nest.flatten(input_spec)[0]

if outer_ndim >= input_spec.shape.ndims:

# We can capture the "outer batch size" as the diff between the

# expected input rank and the rank of the non-batched spec passed in the

# input_spec.

return outer_ndim - input_spec.shape.ndims

# Empty input_spec.

return 1

def get_state_spec(layer: tf.keras.layers.Layer) -> types.NestedTensorSpec:

"""Extracts the state spec from a layer.

Args:

layer: The layer to extract from; can be a `Network`.

Returns:

The state spec.

Raises:

TypeError: If `layer` is a subclass of `tf.keras.layers.RNN` (it must

be wrapped by an `RNNWrapper` object).

ValueError: If `layer` is a Keras layer and `create_variables` has

not been called on it.

"""

if isinstance(layer, Network):

return layer.state_spec

if isinstance(layer, tf.keras.layers.RNN):

raise TypeError(

"RNN Layer must be wrapped inside "

"`tf_agents.keras_layers.RNNWrapper`: {}".format(layer)

)

initial_state = getattr(layer, "get_initial_state", None)

state_size = getattr(layer, "state_size", None)

if initial_state is not None and state_size is None:

raise ValueError(

"Layer lacks a `state_size` property. Unable to extract state "

"spec: {}".format(layer)

)

state_spec = ()

if state_size is not None:

state_spec = tf.nest.map_structure(

lambda s: tf.TensorSpec(dtype=layer.dtype, shape=s), state_size

)

return state_spec