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

"""Utilities related to TensorSpec class."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import os

from typing import Union

import numpy as np

import tensorflow as tf

import tensorflow_probability as tfp

from tf_agents.specs import array_spec

from tf_agents.typing import types

from google.protobuf import text_format

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

from tensorflow.core.protobuf import struct_pb2 # TF internal

from tensorflow.python.framework import tensor_spec as ts # TF internal

from tensorflow.python.saved_model import nested_structure_coder # TF internal

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

tfd = tfp.distributions

TensorSpec = tf.TensorSpec

BoundedTensorSpec = ts.BoundedTensorSpec

def is_bounded(spec):

if isinstance(spec, (array_spec.BoundedArraySpec, BoundedTensorSpec)):

return True

elif hasattr(spec, "minimum") and hasattr(spec, "maximum"):

return hasattr(spec, "dtype") and hasattr(spec, "shape")

def is_discrete(spec):

if isinstance(spec, TensorSpec):

return spec.dtype.is_integer

else:

return array_spec.is_discrete(spec)

def is_continuous(spec):

if isinstance(spec, TensorSpec):

return spec.dtype.is_floating

else:

return array_spec.is_continuous(spec)

def from_spec(spec):

"""Maps the given spec into corresponding TensorSpecs keeping bounds."""

def _convert_to_tensor_spec(s):

# Need to check bounded first as non bounded specs are base class.

if isinstance(s, tf.TypeSpec):

return s

if is_bounded(s):

return BoundedTensorSpec.from_spec(s)

elif isinstance(s, array_spec.ArraySpec):

return TensorSpec.from_spec(s)

else:

raise ValueError(

"No known conversion from type `%s` to a TensorSpec. Saw:\n %s"

% (type(s), s)

)

return tf.nest.map_structure(_convert_to_tensor_spec, spec)

def to_array_spec(

tensor_spec: Union[types.NestedArraySpec, types.NestedTensorSpec]

) -> types.NestedArraySpec:

"""Converts TensorSpec into ArraySpec."""

def _convert(s):

if isinstance(s, array_spec.ArraySpec):

return s

if hasattr(s, "minimum") and hasattr(s, "maximum"):

return array_spec.BoundedArraySpec(

s.shape.as_list(),

s.dtype.as_numpy_dtype,

minimum=s.minimum,

maximum=s.maximum,

name=s.name,

)

else:

return array_spec.ArraySpec(

s.shape.as_list(), s.dtype.as_numpy_dtype, s.name

)

return tf.nest.map_structure(_convert, tensor_spec)

def to_nest_array_spec(

nest_array_spec: Union[types.NestedArraySpec, types.NestedTensorSpec]

) -> types.NestedArraySpec:

"""(Deprecated) Alias for `to_array_spec`."""

return to_array_spec(nest_array_spec)

def to_placeholder(spec, outer_dims=()):

"""Creates a placeholder from TensorSpec.

Args:

spec: instance of TensorSpec

outer_dims: optional leading dimensions of the placeholder.

Returns:

An instance of tf.placeholder.

"""

ph_shape = list(outer_dims) + spec.shape.as_list()

return tf.compat.v1.placeholder(spec.dtype, ph_shape, spec.name)

def to_placeholder_with_default(default, spec, outer_dims=()):

"""Creates a placeholder from TensorSpec.

Args:

default: A constant value of output type dtype.

spec: Instance of TensorSpec

outer_dims: Optional leading dimensions of the placeholder.

Returns:

An instance of tf.placeholder.

"""

ph_shape = list(outer_dims) + spec.shape.as_list()

return tf.compat.v1.placeholder_with_default(default, ph_shape, spec.name)

def to_nest_placeholder(

nested_tensor_specs, default=None, name_scope="", outer_dims=()

):

"""Converts a nest of TensorSpecs to a nest of matching placeholders.

Args:

nested_tensor_specs: A nest of tensor specs.

default: Optional constant value to set as a default for the placeholder.

name_scope: String name for the scope to create the placeholders in.

outer_dims: Optional leading dimensions for the placeholder.

Returns:

A nest of placeholders matching the given tensor spec.

Raises:

ValueError: If a default is provided outside of the allowed types, or if

default is a np.array that does not match the spec shape.

"""

if default is None:

to_ph = lambda spec: to_placeholder(spec, outer_dims=outer_dims)

else:

if not isinstance(default, (int, float, np.ndarray)):

raise ValueError(

"to_nest_placeholder default value must be an int, "

"float, or np.ndarray"

)

def to_ph(spec):

shape = list(outer_dims) + spec.shape.as_list()

if isinstance(default, np.ndarray) and list(default.shape) != shape:

raise ValueError(

"Shape mismatch between default value and spec. "

"Got {}, expected {}".format(default.shape, shape)

)

const = tf.constant(default, shape=shape, dtype=spec.dtype)

return to_placeholder_with_default(const, spec, outer_dims=outer_dims)

with tf.name_scope(name_scope):

return tf.nest.map_structure(to_ph, nested_tensor_specs)

def _random_uniform_int(shape, outer_dims, minval, maxval, dtype, seed=None):

"""Iterates over n-d tensor minval, maxval limits to sample uniformly."""

# maxval in BoundedTensorSpec is bound inclusive.

# tf.random_uniform is upper bound exclusive, +1 to fix the sampling

# behavior.

# However +1 could cause overflow, in such cases we use the original maxval.

maxval = np.broadcast_to(maxval, minval.shape).astype(dtype.as_numpy_dtype)

minval = np.broadcast_to(minval, maxval.shape).astype(dtype.as_numpy_dtype)

sampling_maxval = maxval

if dtype.is_integer:

sampling_maxval = np.where(maxval < dtype.max, maxval + 1, maxval)

if not np.all(shape[-len(minval.shape) :] == minval.shape):

raise ValueError(

"%s == shape[-%d:] != minval.shape == %s. shape == %s."

% (shape[len(minval.shape) :], len(minval.shape), minval.shape, shape)

)

# Example:

# minval = [1.0, 2.0]

# shape = [3, 2]

# outer_dims = [5]

# Sampling becomes:

# sample [5, 3] for minval 1.0

# sample [5, 3] for minval 2.0

# stack on innermost axis to get [5, 3, 2]

# reshape to get [5, 3, 2]

samples = []

shape = tf.convert_to_tensor(shape, dtype=tf.int32)

sample_shape = tf.concat((outer_dims, shape[: -len(minval.shape)]), axis=0)

full_shape = tf.concat((outer_dims, shape), axis=0)

for single_min, single_max in zip(minval.flat, sampling_maxval.flat):

samples.append(

tf.random.uniform(

shape=sample_shape,

minval=single_min,

maxval=single_max,

dtype=dtype,

seed=seed,

)

)

samples = tf.stack(samples, axis=-1)

samples = tf.reshape(samples, full_shape)

return samples

def sample_bounded_spec(spec, seed=None, outer_dims=None):

"""Samples uniformily the given bounded spec.

Args:

spec: A BoundedSpec to sample.

seed: A seed used for sampling ops

outer_dims: An optional `Tensor` specifying outer dimensions to add to the

spec shape before sampling.

Returns:

A Tensor sample of the requested spec.

"""

minval = spec.minimum

maxval = spec.maximum

dtype = tf.as_dtype(spec.dtype)

# To sample uint8 we will use int32 and cast later. This is needed for two

# reasons:

# - tf.random_uniform does not currently support uint8

# - if you want to sample [0, 255] range, there's no way to do this since

# tf.random_uniform has exclusive upper bound and 255 + 1 would overflow.

is_uint8 = dtype == tf.uint8

sampling_dtype = tf.int32 if is_uint8 else dtype

if dtype in [tf.float64, tf.float32]:

# Avoid under/over-flow as random_uniform can't sample over the full range

# for these types.

minval = np.maximum(dtype.min / 8, minval)

maxval = np.minimum(dtype.max / 8, maxval)

if outer_dims is None:

outer_dims = tf.constant([], dtype=tf.int32)

else:

outer_dims = tf.convert_to_tensor(outer_dims, dtype=tf.int32)

def _unique_vals(vals):

if vals.size > 0:

if vals.ndim > 0:

return np.all(vals == vals[0])

return True

if (minval.ndim != 0 or maxval.ndim != 0) and not (

_unique_vals(minval) and _unique_vals(maxval)

):

# tf.random_uniform can only handle minval/maxval 0-d tensors.

res = _random_uniform_int(

shape=spec.shape,

outer_dims=outer_dims,

minval=minval,

maxval=maxval,

dtype=sampling_dtype,

seed=seed,

)

else:

minval = minval.item(0) if minval.ndim != 0 else minval

maxval = maxval.item(0) if maxval.ndim != 0 else maxval

# BoundedTensorSpec are bounds inclusive.

# tf.random_uniform is upper bound exclusive, +1 to fix the sampling

# behavior.

# However +1 will cause overflow, in such cases we use the original maxval.

if sampling_dtype.is_integer and maxval < sampling_dtype.max:

maxval = maxval + 1

shape = tf.convert_to_tensor(spec.shape, dtype=tf.int32)

full_shape = tf.concat((outer_dims, shape), axis=0)

res = tf.random.uniform(

full_shape,

minval=minval,

maxval=maxval,

dtype=sampling_dtype,

seed=seed,

)

if is_uint8:

res = tf.cast(res, dtype=dtype)

return res

def sample_spec_nest(

structure, seed=None, outer_dims=(), minimum=None, maximum=None

):

"""Samples the given nest of specs.

Args:

structure: A nest of `TensorSpec`.

seed: A seed used for sampling ops

outer_dims: An optional `Tensor` specifying outer dimensions to add to the

spec shape before sampling.

minimum: An optional numeric value. If set, numeric specs within the nest

(both bounded and unbounded) will be restricted to this minimum.

maximum: Similar to the above but with maximums.

Returns:

A nest of sampled values following the ArraySpec definition.

Raises:

TypeError: If `spec` is an unknown type.

NotImplementedError: If `outer_dims` is not statically known but nest

contains a `SparseTensorSpec`.

"""

seed_stream = tfp.util.SeedStream(seed=seed, salt="sample_spec_nest")

def sample_fn(spec):

"""Return a composite tensor sample given `spec`.

Args:

spec: A TensorSpec, SparseTensorSpec, etc.

Returns:

A tensor or SparseTensor.

Raises:

NotImplementedError: If `outer_dims` is not statically known and a

SparseTensor is requested.

"""

if isinstance(spec, tf.SparseTensorSpec):

outer_shape = tf.get_static_value(outer_dims)

if outer_dims is not None and outer_shape is None:

raise NotImplementedError(

"outer_dims must be statically known, got: {}".format(outer_dims)

)

shape = tf.TensorShape(outer_shape or []).concatenate(spec.shape)

if shape.num_elements() == 0 or tf.compat.dimension_value(shape[0]) == 0:

return tf.SparseTensor(

indices=tf.zeros([0, shape.rank], dtype=tf.int64),

values=tf.zeros([0], dtype=spec.dtype),

dense_shape=shape,

)

indices_spec = BoundedTensorSpec(

dtype=tf.int64,

shape=[7, shape.rank],

minimum=[0] * shape.rank,

maximum=[x - 1 for x in shape.as_list()],

)

values_dtype = tf.int32 if spec.dtype == tf.string else spec.dtype

values_spec = BoundedTensorSpec(

dtype=values_dtype,

shape=[7],

minimum=0,

maximum=shape.as_list()[-1] - 1,

)

values_sample = sample_bounded_spec(values_spec, seed=seed_stream())

if spec.dtype == tf.string:

values_sample = tf.as_string(values_sample)

return tf.sparse.reorder(

tf.SparseTensor(

indices=sample_bounded_spec(indices_spec, seed=seed_stream()),

values=values_sample,

dense_shape=shape,

)

)

elif isinstance(spec, (TensorSpec, BoundedTensorSpec)):

if spec.dtype == tf.string:

sample_spec = BoundedTensorSpec(

spec.shape, tf.int32, minimum=0, maximum=10

)

return tf.as_string(

sample_bounded_spec(

sample_spec, outer_dims=outer_dims, seed=seed_stream()

)

)

elif spec.dtype == tf.bool:

sample_spec = BoundedTensorSpec(

spec.shape, tf.int32, minimum=0, maximum=1

)

return tf.cast(

sample_bounded_spec(

sample_spec, outer_dims=outer_dims, seed=seed_stream()

),

tf.bool,

)

else:

bounded_spec = BoundedTensorSpec.from_spec(spec)

spec_max = bounded_spec.maximum

if maximum is not None:

spec_max = min(maximum, spec_max)

spec_min = bounded_spec.minimum

if minimum is not None:

spec_min = max(minimum, spec_min)

bounded_spec = BoundedTensorSpec(

shape=bounded_spec.shape,

dtype=bounded_spec.dtype,

minimum=spec_min,

maximum=spec_max,

)

return sample_bounded_spec(

bounded_spec,

outer_dims=outer_dims,

seed=seed_stream(),

)

else:

raise TypeError("Spec type not supported: '{}'".format(spec))

return tf.nest.map_structure(sample_fn, structure)

def zero_spec_nest(specs, outer_dims=None):

"""Create zero tensors for a given spec.

Args:

specs: A nest of `TensorSpec`.

outer_dims: An optional list of constants or `Tensor` specifying outer

dimensions to add to the spec shape before sampling.

Returns:

A nest of zero tensors matching `specs`, with the optional outer

dimensions added.

Raises:

TypeError: If `specs` is an unknown type.

NotImplementedError: If `specs` contains non-dense tensor specs.

"""

def make_zero(spec):

if not isinstance(spec, TensorSpec):

raise NotImplementedError("Spec type not supported: '{}'".format(spec))

if outer_dims is None:

shape = spec.shape

else:

spec_shape = tf.convert_to_tensor(value=spec.shape, dtype=tf.int32)

shape = tf.concat((outer_dims, spec_shape), axis=0)

return tf.zeros(shape, spec.dtype)

if specs:

if outer_dims is None:

outer_dims = tf.constant([], dtype=tf.int32)

else:

outer_dims = tf.convert_to_tensor(outer_dims, dtype=tf.int32)

return tf.nest.map_structure(make_zero, specs)

def add_outer_dims_nest(specs, outer_dims):

"""Adds outer dimensions to the shape of input specs.

Args:

specs: Nested list/tuple/dict of TensorSpecs/ArraySpecs, describing the

shape of tensors.

outer_dims: a list or tuple, representing the outer shape to be added to the

TensorSpecs in specs.

Returns:

Nested TensorSpecs with outer dimensions added to the shape of input specs.

Raises:

ValueError: if any outer_dims is neither a list nor tuple.

"""

if not isinstance(outer_dims, (tuple, list)):

raise ValueError("outer_dims must be a tuple or list of dimensions")

def add_outer_dims(spec):

# TODO(b/187478998) Use spec.name when tf.SparseTensorSpec supports it.

name = getattr(spec, "name", None)

shape = outer_dims + spec.shape

if hasattr(spec, "minimum") and hasattr(spec, "maximum"):

return BoundedTensorSpec(

shape, spec.dtype, spec.minimum, spec.maximum, name

)

elif isinstance(specs, tf.SparseTensorSpec):

# TODO(b/187478998) Add name when tf.SparseTensorSpec supports it.

return tf.SparseTensorSpec(shape, spec.dtype)

return TensorSpec(shape, spec.dtype, name=name)

return tf.nest.map_structure(add_outer_dims, specs)

def add_outer_dim(specs, dim=None):

"""Adds an outer dimension to the shape of input specs.

Args:

specs: Nested list/tuple/dict of TensorSpecs/ArraySpecs, describing the

shape of tensors.

dim: Int, representing the outer dimension to be added to the TensorSpecs in

specs.

Returns:

Nested TensorSpecs with outer dimensions added to the shape of input specs.

"""

return add_outer_dims_nest(specs, outer_dims=(dim,))

def with_dtype(specs, dtype):

"""Updates dtypes of all specs in the input spec.

Args:

specs: Nested list/tuple/dict of TensorSpecs/ArraySpecs, describing the

shape of tensors.

dtype: dtype to update the specs to.

Returns:

Nested TensorSpecs with the udpated dtype.

"""

def update_dtype(spec):

if hasattr(spec, "minimum") and hasattr(spec, "maximum"):

return BoundedTensorSpec(

spec.shape, dtype, spec.minimum, spec.maximum, spec.name

)

return TensorSpec(spec.shape, dtype, name=spec.name)

return tf.nest.map_structure(update_dtype, specs)

def remove_outer_dims_nest(specs, num_outer_dims):

"""Removes the specified number of outer dimensions from the input spec nest.

Args:

specs: Nested list/tuple/dict of TensorSpecs/ArraySpecs, describing the

shape of tensors.

num_outer_dims: (int) Number of outer dimensions to remove.

Returns:

Nested TensorSpecs with outer dimensions removed from the input specs.

Raises:

Value error if a spec in the nest has shape rank less than `num_outer_dims`.

"""

def remove_outer_dims(spec):

"""Removes the num_outer_dims of a tensor spec."""

# TODO(b/187478998) Use spec.name when tf.SparseTensorSpec supports it.

name = getattr(spec, "name", None)

if len(spec.shape) < num_outer_dims:

raise ValueError(

"The shape of spec {} has rank lower than the specified "

"num_outer_dims {}".format(spec, num_outer_dims)

)

shape = list(spec.shape)[num_outer_dims:]

if hasattr(spec, "minimum") and hasattr(spec, "maximum"):

if isinstance(spec.minimum, (tuple, list)) and len(spec.minimum) == len(

spec.shape

):

minimum = spec.minimum[num_outer_dims:]

else:

minimum = spec.minimum

if isinstance(spec.maximum, (tuple, list)) and len(spec.maximum) == len(

spec.shape

):

maximum = spec.maximum[num_outer_dims:]

else:

maximum = spec.maximum

return BoundedTensorSpec(shape, spec.dtype, minimum, maximum, name)

elif isinstance(spec, tf.SparseTensorSpec):

# TODO(b/187478998) Add name when tf.SparseTensorSpec supports it.

return tf.SparseTensorSpec(shape, spec.dtype)

return TensorSpec(shape, spec.dtype, name=name)

return tf.nest.map_structure(remove_outer_dims, specs)

def to_proto(spec):

"""Encodes a nested spec into a struct_pb2.StructuredValue proto.

Args:

spec: Nested list/tuple or dict of TensorSpecs, describing the shape of the

non-batched Tensors.

Returns:

A `struct_pb2.StructuredValue` proto.

"""

# Make sure spec is a tensor_spec.

spec = from_spec(spec)

return nested_structure_coder.encode_structure(spec)

def from_proto(spec_proto):

"""Decodes a struct_pb2.StructuredValue proto into a nested spec."""

return nested_structure_coder.decode_proto(spec_proto)

def from_packed_proto(spec_packed_proto):

"""Decodes a packed Any proto containing the structured value for the spec."""

spec_proto = struct_pb2.StructuredValue()

spec_packed_proto.Unpack(spec_proto)

return from_proto(spec_proto)

def to_pbtxt_file(output_path, spec):

"""Saves a spec encoded as a struct_pb2.StructuredValue in a pbtxt file."""

spec_proto = to_proto(spec)

dir_path = os.path.split(output_path)[0]

tf.io.gfile.makedirs(dir_path)

with tf.io.gfile.GFile(output_path, "wb") as f:

f.write(text_format.MessageToString(spec_proto))

def from_pbtxt_file(spec_path):

"""Loads a spec encoded as a struct_pb2.StructuredValue from a pbtxt file."""

spec_proto = struct_pb2.StructuredValue()

with tf.io.gfile.GFile(spec_path, "rb") as f:

text_format.MergeLines(f, spec_proto)

return from_proto(spec_proto)