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

"""Keras layer to replace the Sequential Model object."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import copy

from typing import Any, List, Mapping, Optional, Sequence, Text, Tuple, Union

import tensorflow as tf

from tf_agents.keras_layers import rnn_wrapper

from tf_agents.networks import network

from tf_agents.typing import types

from tf_agents.utils import common

def _infer_state_specs(

layers: Sequence[tf.keras.layers.Layer],

) -> Tuple[types.NestedTensorSpec, List[bool]]:

"""Infer the state spec of a sequence of keras Layers and Networks.

Args:

layers: A list of Keras layers and Network.

Returns:

A tuple with `state_spec`, a tuple of the state specs of length

`len(layers)` and a list of bools indicating if the corresponding layer

has lists in it's state.

"""

state_specs = []

layer_state_is_list = []

for layer in layers:

spec = network.get_state_spec(layer)

if isinstance(spec, list):

layer_state_is_list.append(True)

state_specs.append(tuple(spec))

else:

state_specs.append(spec)

layer_state_is_list.append(False)

return tuple(state_specs), layer_state_is_list

class Sequential(network.Network):

"""The Sequential Network represents a sequence of Keras layers.

It is a TF-Agents network that should be used instead of

tf.keras.layers.Sequential. In contrast to keras Sequential, this layer can be

used as a pure Layer in tf.functions and when exporting SavedModels, without

having to pre-declare input and output shapes. In turn, this layer is usable

as a preprocessing layer for TF Agents Networks, and can be exported via

PolicySaver.

Stateful Keras layers (e.g. LSTMCell, RNN, LSTM, TF-Agents DynamicUnroll)

are all supported. The `state_spec` of `Sequential` is a tuple whose

length matches the number of stateful layers passed. If no stateful layers

or networks are passed to `Sequential` then `state_spec == ()`. Given that

the replay buffers do not support specs with lists due to tf.nest vs

tf.data.nest conflicts `Sequential` will also guarantee that all specs do not

contain lists.

Usage:

```python

c = Sequential([layer1, layer2, layer3])

output, next_state = c(inputs, state)

```

"""

def __init__(

self,

layers: Sequence[tf.keras.layers.Layer],

input_spec: Optional[types.NestedTensorSpec] = None,

name: Optional[Text] = None,

):

"""Create a Sequential Network.

Args:

layers: A list or tuple of layers to compose. Any layers that are

subclasses of `tf.keras.layers.{RNN,LSTM,GRU,...}` are wrapped in

`tf_agents.keras_layers.RNNWrapper`.

input_spec: (Optional.) A nest of `tf.TypeSpec` representing the input

observations to the first layer.

name: (Optional.) Network name.

Raises:

ValueError: If `layers` is empty.

ValueError: If `layers[0]` is a generic Keras layer (not a TF-Agents

network) and `input_spec is None`.

TypeError: If any of the layers are not instances of keras `Layer`.

"""

if not layers:

raise ValueError('`layers` must not be empty; saw: {}'.format(layers))

for layer in layers:

if not isinstance(layer, tf.keras.layers.Layer):

raise TypeError(

'Expected all layers to be instances of keras Layer, but saw'

": '{}'".format(layer)

)

layers = [

rnn_wrapper.RNNWrapper(layer)

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

else layer

for layer in layers

]

state_spec, self._layer_state_is_list = _infer_state_specs(layers)

# Now we remove all of the empty state specs so if there are no RNN layers,

# our state spec is empty. layer_has_state is a list of bools telling us

# which layers have a non-empty state and which don't.

flattened_specs = [tf.nest.flatten(s) for s in state_spec]

layer_has_state = [bool(fs) for fs in flattened_specs]

state_spec = tuple(

s for s, has_state in zip(state_spec, layer_has_state) if has_state

)

super(Sequential, self).__init__(

input_tensor_spec=input_spec, state_spec=state_spec, name=name

)

self._sequential_layers = layers

self._layer_has_state = layer_has_state

@property

def layers(self) -> List[tf.keras.layers.Layer]:

# Return a shallow copy so users don't modify the layers list.

return copy.copy(self._sequential_layers)

def copy(self, **kwargs) -> 'Sequential':

"""Make a copy of a `Sequential` instance.

**NOTE** A copy of a `Sequential` instance always performs a deep copy

of the underlying layers, so the new instance will not share weights

with the original - but it will start with the same weights.

Args:

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

overridden args include 'name'.

Returns:

A deep copy of this network.

Raises:

RuntimeError: If not `tf.executing_eagerly()`; as this is required to

be able to create deep copies of layers in `layers`.

"""

if not tf.executing_eagerly():

raise RuntimeError(

'Not executing eagerly - cannot make deep copies of `layers`.'

)

new_kwargs = dict(self._saved_kwargs, **kwargs)

if 'layers' not in kwargs:

new_layers = [copy.deepcopy(l) for l in self.layers]

new_kwargs['layers'] = new_layers

return type(self)(**new_kwargs)

def call(self, inputs, network_state=(), **kwargs):

if not tf.is_tensor(network_state) and not network_state:

network_state = ((),) * len(self.state_spec)

next_network_state = [()] * len(self.state_spec)

# Only Networks are expected to know about step_type; not Keras layers.

layer_kwargs = kwargs.copy()

layer_kwargs.pop('step_type', None)

stateful_layer_idx = 0

for i, layer in enumerate(self.layers):

if isinstance(layer, network.Network):

if self._layer_has_state[i]:

input_state = network_state[stateful_layer_idx]

if input_state is not None and self._layer_state_is_list[i]:

input_state = list(input_state)

inputs, next_state = layer(

inputs, network_state=network_state[stateful_layer_idx], **kwargs

)

if self._layer_state_is_list[i]:

next_network_state[stateful_layer_idx] = tuple(next_state)

else:

next_network_state[stateful_layer_idx] = next_state

stateful_layer_idx += 1

else:

inputs, _ = layer(inputs, **kwargs)

else:

# Generic Keras layer

if self._layer_has_state[i]:

# The layer maintains state. If a state was provided at input to

# `call`, then use it. Otherwise ask for an initial state.

maybe_network_state = network_state[stateful_layer_idx]

input_state = maybe_network_state

if maybe_network_state is None:

input_state = layer.get_initial_state(inputs)

elif (

common.safe_has_state(input_state)

and self._layer_state_is_list[i]

):

input_state = list(input_state)

outputs = layer(inputs, input_state, **layer_kwargs)

inputs, next_state = outputs

if self._layer_state_is_list[i]:

next_network_state[stateful_layer_idx] = tuple(next_state)

else:

next_network_state[stateful_layer_idx] = next_state

stateful_layer_idx += 1

else:

# Does not maintain state.

inputs = layer(inputs, **layer_kwargs)

return inputs, tuple(next_network_state)

def compute_output_shape(

self, input_shape: Union[List[int], Tuple[int], tf.TensorShape]

) -> tf.TensorShape:

output_shape = tf.TensorShape(input_shape)

for l in self._sequential_layers:

output_shape = l.compute_output_shape(output_shape)

return tf.TensorShape(output_shape)

def compute_output_signature(

self, input_signature: types.NestedSpec

) -> types.NestedSpec:

output_signature = input_signature

for l in self._sequential_layers:

output_signature = l.compute_output_signature(output_signature)

return output_signature

@property

def trainable_weights(self) -> List[tf.Variable]:

if not self.trainable:

return []

weights = {}

for l in self._sequential_layers:

for v in l.trainable_weights:

weights[id(v)] = v

return list(weights.values())

@property

def non_trainable_weights(self) -> List[tf.Variable]:

weights = {}

for l in self._sequential_layers:

for v in l.non_trainable_weights:

weights[id(v)] = v

return list(weights.values())

@property

def trainable(self) -> bool:

return any([l.trainable for l in self._sequential_layers])

@trainable.setter

def trainable(self, value: bool):

for l in self._sequential_layers:

l.trainable = value

def get_config(self) -> Mapping[int, Mapping[str, Any]]:

config = {}

for i, layer in enumerate(self._sequential_layers):

config[i] = {

'class_name': layer.__class__.__name__,

'config': copy.deepcopy(layer.get_config()),

}

return config

@classmethod

def from_config(cls, config, custom_objects=None) -> 'Sequential':

layers = [

tf.keras.layers.deserialize(conf, custom_objects=custom_objects)

for conf in config.values()

]

return cls(layers)

tf.keras.utils.get_custom_objects()['SequentialNetwork'] = Sequential