def testGetTraverseShallowStructure(self):
scalar_traverse_input = [3, 4, (1, 2, [0]), [5, 6], {"a": (7, )}, []]
scalar_traverse_r = nest.get_traverse_shallow_structure(
lambda s: not isinstance(s, tuple), scalar_traverse_input)
self.assertEqual(scalar_traverse_r,
[True, True, False, [True, True], {
"a": False
}, []])
nest.assert_shallow_structure(scalar_traverse_r, scalar_traverse_input)
structure_traverse_input = [(1, [2]), ([1], 2)]
structure_traverse_r = nest.get_traverse_shallow_structure(
lambda s: (True, False)
if isinstance(s, tuple) else True, structure_traverse_input)
self.assertEqual(structure_traverse_r, [(True, False),
([True], False)])
nest.assert_shallow_structure(structure_traverse_r,
structure_traverse_input)
with self.assertRaisesRegexp(TypeError, "returned structure"):
nest.get_traverse_shallow_structure(lambda _: [True], 0)
with self.assertRaisesRegexp(TypeError, "returned a non-bool scalar"):
nest.get_traverse_shallow_structure(lambda _: 1, [1])
with self.assertRaisesRegexp(
TypeError, "didn't return a depth=1 structure of bools"):
nest.get_traverse_shallow_structure(lambda _: [1], [1])
def testNestAssertShallowStructureCompositeMismatch(self,
s1,
s2,
check_types=True):
with self.assertRaises(TypeError): # pylint: disable=g-error-prone-assert-raises
nest.assert_shallow_structure(
s1, s2, expand_composites=True, check_types=check_types)
def testAssertShallowStructure(self):
inp_ab = ["a", "b"]
inp_abc = ["a", "b", "c"]
expected_message = (
"The two structures don't have the same sequence length. Input "
"structure has length 2, while shallow structure has length 3.")
with self.assertRaisesRegexp(ValueError, expected_message):
nest.assert_shallow_structure(inp_abc, inp_ab)
inp_ab1 = [(1, 1), (2, 2)]
inp_ab2 = [[1, 1], [2, 2]]
expected_message = (
"The two structures don't have the same sequence type. Input structure "
"has type <(type|class) 'tuple'>, while shallow structure has type "
"<(type|class) 'list'>.")
with self.assertRaisesRegexp(TypeError, expected_message):
nest.assert_shallow_structure(inp_ab2, inp_ab1)
nest.assert_shallow_structure(inp_ab2, inp_ab1, check_types=False)
inp_ab1 = {"a": (1, 1), "b": {"c": (2, 2)}}
inp_ab2 = {"a": (1, 1), "b": {"d": (2, 2)}}
expected_message = (
r"The two structures don't have the same keys. Input "
r"structure has keys \['c'\], while shallow structure has "
r"keys \['d'\].")
with self.assertRaisesRegexp(ValueError, expected_message):
nest.assert_shallow_structure(inp_ab2, inp_ab1)
inp_ab = collections.OrderedDict([("a", 1), ("b", (2, 3))])
inp_ba = collections.OrderedDict([("b", (2, 3)), ("a", 1)])
nest.assert_shallow_structure(inp_ab, inp_ba)
def testPythonMapImpl(self):
t = data_structures._TupleWrapper((1, data_structures._TupleWrapper((2,))))
self.assertEqual(
(4, (5,)),
nest.map_structure_up_to((None, (None,)), lambda x: x + 3, t,
check_types=True))
nest.assert_shallow_structure((None, None), t)
def pbroadcast_fn(*args):
nest.assert_shallow_structure(args, in_axes)
nest.assert_shallow_structure(out_dtype, out_axes)
map_in_axes = nest.map_structure_up_to(args, canonicalize_axis_name,
in_axes)
map_out_axes = nest.map_structure_up_to(out_dtype, canonicalize_axis_name,
out_axes)
def _pbroadcast_input(out_axes, x, in_axes):
psum_axes = [
axis_name for axis_name in out_axes if axis_name not in in_axes
]
return pbroadcast(x, psum_axes)
def _flat_fn_index(i, *args):
out = fn(*args)
return tf.nest.flatten(out)[i]
def _flat_fn(*args):
outputs = []
for i, out_axis in enumerate(nest.flatten_up_to(out_dtype, map_out_axes)):
local_args = nest.map_structure_up_to(
args, functools.partial(_pbroadcast_input, out_axis), args,
map_in_axes)
outputs.append(_flat_fn_index(i, *local_args))
return tf.nest.pack_sequence_as(out_dtype, outputs)
return _flat_fn(*args)
def testGetTraverseShallowStructure(self):
scalar_traverse_input = [3, 4, (1, 2, [0]), [5, 6], {"a": (7,)}, []]
scalar_traverse_r = nest.get_traverse_shallow_structure(
lambda s: not isinstance(s, tuple),
scalar_traverse_input)
self.assertEqual(scalar_traverse_r,
[True, True, False, [True, True], {"a": False}, []])
nest.assert_shallow_structure(scalar_traverse_r,
scalar_traverse_input)
structure_traverse_input = [(1, [2]), ([1], 2)]
structure_traverse_r = nest.get_traverse_shallow_structure(
lambda s: (True, False) if isinstance(s, tuple) else True,
structure_traverse_input)
self.assertEqual(structure_traverse_r,
[(True, False), ([True], False)])
nest.assert_shallow_structure(structure_traverse_r,
structure_traverse_input)
with self.assertRaisesRegexp(TypeError, "returned structure"):
nest.get_traverse_shallow_structure(lambda _: [True], 0)
with self.assertRaisesRegexp(TypeError, "returned a non-bool scalar"):
nest.get_traverse_shallow_structure(lambda _: 1, [1])
with self.assertRaisesRegexp(
TypeError, "didn't return a depth=1 structure of bools"):
nest.get_traverse_shallow_structure(lambda _: [1], [1])
def testNestAssertShallowStructureCompositeMismatch(
self, s1, s2, check_types=True):
with self.assertRaises(TypeError): # pylint: disable=g-error-prone-assert-raises
nest.assert_shallow_structure(s1,
s2,
expand_composites=True,
check_types=check_types)
def python_function_wrapper(*py_args):
py_args, py_kwargs = nest.pack_sequence_as(args, py_args)
ret = func(*py_args, **py_kwargs)
# TODO(alextp): Catch Exceptions and improve msg, because tensorflow
# ist not able to preserve the traceback, i.e. the Exceptions does not
# contain any information where the Exception was raised.
nest.assert_shallow_structure(output_types, ret)
return nest.flatten(ret)
def python_function_wrapper(*py_args):
py_args, py_kwargs = nest.pack_sequence_as(args, py_args)
ret = func(*py_args, **py_kwargs)
# TODO(alextp): Catch Exceptions and improve msg, because tensorflow
# ist not able to preserve the traceback, i.e. the Exceptions does not
# contain any information where the Exception was raised.
nest.assert_shallow_structure(output_types, ret)
return nest.flatten(ret)
def testNestAssertSameStructureCompositeMismatch(self,
s1,
s2,
error=ValueError):
# s1 and s2 have the same structure if expand_composites=False; but
# different structures if expand_composites=True.
nest.assert_same_structure(s1, s2, expand_composites=False)
nest.assert_shallow_structure(s1, s2, expand_composites=False)
with self.assertRaises(error): # pylint: disable=g-error-prone-assert-raises
nest.assert_same_structure(s1, s2, expand_composites=True)
def testNestAssertSameStructureCompositeMismatch(self,
s1,
s2,
error=ValueError):
# s1 and s2 have the same structure if expand_composites=False; but
# different structures if expand_composites=True.
nest.assert_same_structure(s1, s2, expand_composites=False)
nest.assert_shallow_structure(s1, s2, expand_composites=False)
with self.assertRaises(error): # pylint: disable=g-error-prone-assert-raises
nest.assert_same_structure(s1, s2, expand_composites=True)
def tag_sharding_attribute_for_dequeued_tensors(dequeues, dims):
"""Tags appropriate XLA sharding attribute to the dequeued tensors.
Args:
dequeues: A list of dequeued tensors on TPU.
dims: A list of integer describes how the tensor is partitioned.
Returns:
The same dequeues with appropriate xla_sharding attribute.
"""
nest.assert_shallow_structure(dequeues, dims)
return nest.map_structure_up_to(
dequeues, _tag_sharding_attribute_for_dequeued_tensor, dequeues, dims)
def tag_sharding_attribute_for_dequeued_tensors(dequeues, dims):
"""Tags appropriate XLA sharding attribute to the dequeued tensors.
Args:
dequeues: A list of dequeued tensors on TPU.
dims: A list of integer describes how the tensor is partitioned.
Returns:
The same dequeues with appropriate xla_sharding attribute.
"""
nest.assert_shallow_structure(dequeues, dims)
return nest.map_structure_up_to(
dequeues, _tag_sharding_attribute_for_dequeued_tensor, dequeues, dims)
def testAssertShallowStructure(self):
inp_ab = ["a", "b"]
inp_abc = ["a", "b", "c"]
expected_message = (
"The two structures don't have the same sequence length. Input "
"structure has length 2, while shallow structure has length 3.")
with self.assertRaisesRegexp(ValueError, expected_message):
nest.assert_shallow_structure(inp_abc, inp_ab)
inp_ab1 = [(1, 1), (2, 2)]
inp_ab2 = [[1, 1], [2, 2]]
expected_message = (
"The two structures don't have the same sequence type. Input structure "
"has type <(type|class) 'tuple'>, while shallow structure has type "
"<(type|class) 'list'>.")
with self.assertRaisesRegexp(TypeError, expected_message):
nest.assert_shallow_structure(inp_ab2, inp_ab1)
nest.assert_shallow_structure(inp_ab2, inp_ab1, check_types=False)
inp_ab1 = {"a": (1, 1), "b": {"c": (2, 2)}}
inp_ab2 = {"a": (1, 1), "b": {"d": (2, 2)}}
expected_message = (
r"The two structures don't have the same keys. Input "
r"structure has keys \['c'\], while shallow structure has "
r"keys \['d'\].")
with self.assertRaisesRegexp(ValueError, expected_message):
nest.assert_shallow_structure(inp_ab2, inp_ab1)
def _psum_fn_fwd(*args):
nest.assert_shallow_structure(args, in_axes)
out_parts = fn(*args)
nest.assert_shallow_structure(out_parts, out_axes)
map_out_axes = nest.map_structure_up_to(out_parts, canonicalize_axis_name,
out_axes)
total_out_parts = nest.map_structure_up_to(
out_parts,
lambda out_part, axis_name: ( # pylint: disable=g-long-lambda
psum(out_part, axis_name=axis_name) if axis_name else out_part),
out_parts,
map_out_axes)
return total_out_parts, (args, out_parts)
def fold_in_axis_index(seed, axis_name=None):
"""Folds the active axis index into a seed according to its axis name."""
if axis_name is None:
return seed
nest.assert_shallow_structure(seed, axis_name)
axis_names = nest.map_structure_up_to(seed, canonicalize_named_axis,
axis_name)
def fold_in(seed, axes):
for name in axes:
axis_index = get_axis_index(name)
seed = samplers.fold_in(seed, tf.cast(axis_index, tf.int32))
return seed
return nest.map_structure_up_to(seed, fold_in, seed, axis_names)
def testAssertShallowStructure(self):
inp_ab = ["a", "b"]
inp_abc = ["a", "b", "c"]
expected_message = (
"The two structures don't have the same sequence length. Input "
"structure has length 2, while shallow structure has length 3.")
with self.assertRaisesRegexp(ValueError, expected_message):
nest.assert_shallow_structure(inp_abc, inp_ab)
inp_ab1 = [(1, 1), (2, 2)]
inp_ab2 = [[1, 1], [2, 2]]
expected_message = (
"The two structures don't have the same sequence type. Input structure "
"has type <(type|class) 'tuple'>, while shallow structure has type "
"<(type|class) 'list'>.")
with self.assertRaisesRegexp(TypeError, expected_message):
nest.assert_shallow_structure(inp_ab2, inp_ab1)
def _sharded_log_prob_parts_fwd(value):
nest.assert_shallow_structure(value, axis_names)
map_axes = nest.map_structure_up_to(value, canonicalize_axis_name,
axis_names)
log_prob_parts = log_prob_parts_fn(value)
nest.assert_same_structure(value, log_prob_parts)
nest.assert_shallow_structure(log_prob_parts, map_axes)
total_log_prob_parts = nest.map_structure_up_to(
log_prob_parts,
lambda log_prob_part, axis_name: ( # pylint: disable=g-long-lambda
psum(log_prob_part, axis_name=axis_name)
if axis_name else log_prob_part),
log_prob_parts,
map_axes)
return total_log_prob_parts, value
def testAssertShallowStructure(self):
inp_ab = ["a", "b"]
inp_abc = ["a", "b", "c"]
expected_message = (
"The two structures don't have the same sequence length. Input "
"structure has length 2, while shallow structure has length 3.")
with self.assertRaisesRegexp(ValueError, expected_message):
nest.assert_shallow_structure(inp_abc, inp_ab)
inp_ab1 = [(1, 1), (2, 2)]
inp_ab2 = [[1, 1], [2, 2]]
expected_message = (
"The two structures don't have the same sequence type. Input structure "
"has type <(type|class) 'tuple'>, while shallow structure has type "
"<(type|class) 'list'>.")
with self.assertRaisesRegexp(TypeError, expected_message):
nest.assert_shallow_structure(inp_ab2, inp_ab1)
def testNestAssertShallowStructure(self):
s1 = [[TestCompositeTensor(1, 2, 3)], 100, {'y': TestCompositeTensor(5, 6)}]
s2 = [[TestCompositeTensor(1, 2, 3)], 100, {
'y': TestCompositeTensor(TestCompositeTensor(4, 5), 6)
}]
nest.assert_shallow_structure(s1, s2, expand_composites=False)
nest.assert_shallow_structure(s1, s2, expand_composites=True)
nest.assert_shallow_structure(s2, s1, expand_composites=False)
with self.assertRaises(TypeError):
nest.assert_shallow_structure(s2, s1, expand_composites=True)
def testNestAssertShallowStructure(self):
s1 = [[TestCompositeTensor(1, 2, 3)], 100, {
'y': TestCompositeTensor(5, 6)
}]
s2 = [[TestCompositeTensor(1, 2, 3)], 100, {
'y': TestCompositeTensor(TestCompositeTensor(4, 5), 6)
}]
nest.assert_shallow_structure(s1, s2, expand_composites=False)
nest.assert_shallow_structure(s1, s2, expand_composites=True)
nest.assert_shallow_structure(s2, s1, expand_composites=False)
with self.assertRaises(TypeError):
nest.assert_shallow_structure(s2, s1, expand_composites=True)
def __init__(self,
input_tensor_spec,
preprocessing_layers=None,
preprocessing_combiner=None,
conv_layer_params=None,
fc_layer_params=None,
dropout_layer_params=None,
activation_fn=tf.keras.activations.relu,
weight_decay_params=None,
kernel_initializer=None,
batch_squash=True,
dtype=tf.float32,
name='EncodingNetwork',
conv_type=CONV_TYPE_2D):
"""Creates an instance of `EncodingNetwork`.
Network supports calls with shape outer_rank + input_tensor_spec.shape. Note
outer_rank must be at least 1.
For example an input tensor spec with shape `(2, 3)` will require
inputs with at least a batch size, the input shape is `(?, 2, 3)`.
Input preprocessing is possible via `preprocessing_layers` and
`preprocessing_combiner` Layers. If the `preprocessing_layers` nest is
shallower than `input_tensor_spec`, then the layers will get the subnests.
For example, if:
```python
input_tensor_spec = ([TensorSpec(3)] * 2, [TensorSpec(3)] * 5)
preprocessing_layers = (Layer1(), Layer2())
```
then preprocessing will call:
```python
preprocessed = [preprocessing_layers[0](observations[0]),
preprocessing_layers[1](obsrevations[1])]
```
However if
```python
preprocessing_layers = ([Layer1() for _ in range(2)],
[Layer2() for _ in range(5)])
```
then preprocessing will call:
```python
preprocessed = [
layer(obs) for layer, obs in zip(flatten(preprocessing_layers),
flatten(observations))
]
```
**NOTE** `preprocessing_layers` and `preprocessing_combiner` are not allowed
to have already been built. This ensures calls to `network.copy()` in the
future always have an unbuilt, fresh set of parameters. Furtheremore,
a shallow copy of the layers is always created by the Network, so the
layer objects passed to the network are never modified. For more details
of the semantics of `copy`, see the docstring of
`tf_agents.networks.Network.copy`.
Args:
input_tensor_spec: A nest of `tensor_spec.TensorSpec` representing the
input observations.
preprocessing_layers: (Optional.) A nest of `tf.keras.layers.Layer`
representing preprocessing for the different observations. All of these
layers must not be already built.
preprocessing_combiner: (Optional.) A keras layer that takes a flat list
of tensors and combines them. Good options include
`tf.keras.layers.Add` and `tf.keras.layers.Concatenate(axis=-1)`. This
layer must not be already built.
conv_layer_params: Optional list of convolution layers parameters, where
each item is either a length-three tuple indicating
`(filters, kernel_size, stride)` or a length-four tuple indicating
`(filters, kernel_size, stride, dilation_rate)`.
fc_layer_params: Optional list of fully_connected parameters, where each
item is the number of units in the layer.
dropout_layer_params: Optional list of dropout layer parameters, each item
is the fraction of input units to drop or a dictionary of parameters
according to the keras.Dropout documentation. The additional parameter
`permanent', if set to True, allows to apply dropout at inference for
approximated Bayesian inference. The dropout layers are interleaved with
the fully connected layers; there is a dropout layer after each fully
connected layer, except if the entry in the list is None. This list must
have the same length of fc_layer_params, or be None.
activation_fn: Activation function, e.g. tf.keras.activations.relu.
weight_decay_params: Optional list of weight decay parameters for the
fully connected layers.
kernel_initializer: Initializer to use for the kernels of the conv and
dense layers. If none is provided a default variance_scaling_initializer
batch_squash: If True the outer_ranks of the observation are squashed into
the batch dimension. This allow encoding networks to be used with
observations with shape [BxTx...].
dtype: The dtype to use by the convolution and fully connected layers.
name: A string representing name of the network.
conv_type: string, '1d' or '2d'. Convolution layers will be 1d or 2D
respectively
Raises:
ValueError: If any of `preprocessing_layers` is already built.
ValueError: If `preprocessing_combiner` is already built.
ValueError: If the number of dropout layer parameters does not match the
number of fully connected layer parameters.
ValueError: If conv_layer_params tuples do not have 3 or 4 elements each.
"""
if preprocessing_layers is None:
flat_preprocessing_layers = None
else:
flat_preprocessing_layers = [
_copy_layer(layer) for layer in tf.nest.flatten(preprocessing_layers)
]
# Assert shallow structure is the same. This verifies preprocessing
# layers can be applied on expected input nests.
input_nest = input_tensor_spec
# Given the flatten on preprocessing_layers above we need to make sure
# input_tensor_spec is a sequence for the shallow_structure check below
# to work.
if not nest.is_sequence(input_tensor_spec):
input_nest = [input_tensor_spec]
nest.assert_shallow_structure(preprocessing_layers, input_nest)
if (len(tf.nest.flatten(input_tensor_spec)) > 1 and
preprocessing_combiner is None):
raise ValueError(
'preprocessing_combiner layer is required when more than 1 '
'input_tensor_spec is provided.')
if preprocessing_combiner is not None:
preprocessing_combiner = _copy_layer(preprocessing_combiner)
if not kernel_initializer:
kernel_initializer = tf.compat.v1.variance_scaling_initializer(
scale=2.0, mode='fan_in', distribution='truncated_normal')
layers = []
if conv_layer_params:
if conv_type == '2d':
conv_layer_type = tf.keras.layers.Conv2D
elif conv_type == '1d':
conv_layer_type = tf.keras.layers.Conv1D
else:
raise ValueError('unsupported conv type of %s. Use 1d or 2d' % (
conv_type))
for config in conv_layer_params:
if len(config) == 4:
(filters, kernel_size, strides, dilation_rate) = config
elif len(config) == 3:
(filters, kernel_size, strides) = config
dilation_rate = (1, 1) if conv_type == '2d' else (1,)
else:
raise ValueError(
'only 3 or 4 elements permitted in conv_layer_params tuples')
layers.append(
conv_layer_type(
filters=filters,
kernel_size=kernel_size,
strides=strides,
dilation_rate=dilation_rate,
activation=activation_fn,
kernel_initializer=kernel_initializer,
dtype=dtype))
layers.append(tf.keras.layers.Flatten())
if fc_layer_params:
if dropout_layer_params is None:
dropout_layer_params = [None] * len(fc_layer_params)
else:
if len(dropout_layer_params) != len(fc_layer_params):
raise ValueError('Dropout and fully connected layer parameter lists'
'have different lengths (%d vs. %d.)' %
(len(dropout_layer_params), len(fc_layer_params)))
if weight_decay_params is None:
weight_decay_params = [None] * len(fc_layer_params)
else:
if len(weight_decay_params) != len(fc_layer_params):
raise ValueError('Weight decay and fully connected layer parameter '
'lists have different lengths (%d vs. %d.)' %
(len(weight_decay_params), len(fc_layer_params)))
for num_units, dropout_params, weight_decay in zip(
fc_layer_params, dropout_layer_params, weight_decay_params):
kernal_regularizer = None
if weight_decay is not None:
kernal_regularizer = tf.keras.regularizers.l2(weight_decay)
layers.append(
tf.keras.layers.Dense(
num_units,
activation=activation_fn,
kernel_initializer=kernel_initializer,
kernel_regularizer=kernal_regularizer,
dtype=dtype))
if not isinstance(dropout_params, dict):
dropout_params = {'rate': dropout_params} if dropout_params else None
if dropout_params is not None:
layers.append(utils.maybe_permanent_dropout(**dropout_params))
super(EncodingNetwork, self).__init__(
input_tensor_spec=input_tensor_spec, state_spec=(), name=name)
# Pull out the nest structure of the preprocessing layers. This avoids
# saving the original kwarg layers as a class attribute which Keras would
# then track.
self._preprocessing_nest = tf.nest.map_structure(lambda l: None,
preprocessing_layers)
self._flat_preprocessing_layers = flat_preprocessing_layers
self._preprocessing_combiner = preprocessing_combiner
self._postprocessing_layers = layers
self._batch_squash = batch_squash
self.built = True # Allow access to self.variables
def testNestAssertShallowStructure(self, s1, s2, expand_composites=True):
nest.assert_shallow_structure(s1,
s2,
expand_composites=expand_composites)
def testAssertShallowStructure(self):
inp_ab = ["a", "b"]
inp_abc = ["a", "b", "c"]
with self.assertRaisesWithLiteralMatch(
ValueError,
nest._INPUT_TREE_SMALLER_THAN_SHALLOW_TREE.format(
shallow_size=len(inp_abc), input_size=len(inp_ab))):
nest.assert_shallow_structure(shallow_tree=inp_abc,
input_tree=inp_ab)
inp_ab1 = [(1, 1), (2, 2)]
inp_ab2 = [[1, 1], [2, 2]]
with self.assertRaisesWithLiteralMatch(
TypeError,
nest._STRUCTURES_HAVE_MISMATCHING_TYPES.format(
shallow_type=type(inp_ab2[0]),
input_type=type(inp_ab1[0]))):
nest.assert_shallow_structure(inp_ab2, inp_ab1)
nest.assert_shallow_structure(inp_ab2, inp_ab1, check_types=False)
inp_ab1 = {"a": (1, 1), "b": {"c": (2, 2)}}
inp_ab2 = {"a": (1, 1), "b": {"d": (2, 2)}}
with self.assertRaisesWithLiteralMatch(
ValueError, nest._SHALLOW_TREE_HAS_INVALID_KEYS.format(["d"])):
nest.assert_shallow_structure(inp_ab2, inp_ab1)
inp_ab = collections.OrderedDict([("a", 1), ("b", (2, 3))])
inp_ba = collections.OrderedDict([("b", (2, 3)), ("a", 1)])
nest.assert_shallow_structure(inp_ab, inp_ba)
# This assertion is expected to pass: two namedtuples with the same
# name and field names are considered to be identical.
inp_shallow = NestTest.SameNameab(1, 2)
inp_deep = NestTest.SameNameab2(1, [1, 2, 3])
nest.assert_shallow_structure(inp_shallow, inp_deep, check_types=False)
nest.assert_shallow_structure(inp_shallow, inp_deep, check_types=True)
def _factored_surrogate_posterior( # pylint: disable=dangerous-default-value
event_shape=None,
bijector=None,
batch_shape=(),
base_distribution_cls=normal.Normal,
initial_parameters={'scale': 1e-2},
dtype=tf.float32,
validate_args=False,
name=None):
"""Builds a joint variational posterior that factors over model variables.
By default, this method creates an independent trainable Normal distribution
for each variable, transformed using a bijector (if provided) to
match the support of that variable. This makes extremely strong
assumptions about the posterior: that it is approximately normal (or
transformed normal), and that all model variables are independent.
Args:
event_shape: `Tensor` shape, or nested structure of `Tensor` shapes,
specifying the event shape(s) of the posterior variables.
bijector: Optional `tfb.Bijector` instance, or nested structure of such
instances, defining support(s) of the posterior variables. The structure
must match that of `event_shape` and may contain `None` values. A
posterior variable will be modeled as
`tfd.TransformedDistribution(underlying_dist, bijector)` if a
corresponding constraining bijector is specified, otherwise it is modeled
as supported on the unconstrained real line.
batch_shape: The `batch_shape` of the output distribution.
Default value: `()`.
base_distribution_cls: Subclass of `tfd.Distribution` that is instantiated
and optionally transformed by the bijector to define the component
distributions. May optionally be a structure of such subclasses
matching `event_shape`.
Default value: `tfd.Normal`.
initial_parameters: Optional `str : Tensor` dictionary specifying initial
values for some or all of the base distribution's trainable parameters,
or a Python `callable` with signature
`value = parameter_init_fn(parameter_name, shape, dtype, seed,
constraining_bijector)`, passed to `tfp.experimental.util.make_trainable`.
May optionally be a structure matching `event_shape` of such dictionaries
and/or callables. Dictionary entries that do not correspond to parameter
names are ignored.
Default value: `{'scale': 1e-2}` (ignored when `base_distribution` does
not have a `scale` parameter).
dtype: Optional float `dtype` for trainable parameters. May
optionally be a structure of such `dtype`s matching `event_shape`.
Default value: `tf.float32`.
validate_args: Python `bool`. Whether to validate input with asserts. This
imposes a runtime cost. If `validate_args` is `False`, and the inputs are
invalid, correct behavior is not guaranteed.
Default value: `False`.
name: Python `str` name prefixed to ops created by this function.
Default value: `None` (i.e., 'build_factored_surrogate_posterior').
Yields:
*parameters: sequence of `trainable_state_util.Parameter` namedtuples.
These are intended to be consumed by
`trainable_state_util.as_stateful_builder` and
`trainable_state_util.as_stateless_builder` to define stateful and
stateless variants respectively.
### Examples
Consider a Gamma model with unknown parameters, expressed as a joint
Distribution:
```python
Root = tfd.JointDistributionCoroutine.Root
def model_fn():
concentration = yield Root(tfd.Exponential(1.))
rate = yield Root(tfd.Exponential(1.))
y = yield tfd.Sample(tfd.Gamma(concentration=concentration, rate=rate),
sample_shape=4)
model = tfd.JointDistributionCoroutine(model_fn)
```
Let's use variational inference to approximate the posterior over the
data-generating parameters for some observed `y`. We'll build a
surrogate posterior distribution by specifying the shapes of the latent
`rate` and `concentration` parameters, and that both are constrained to
be positive.
```python
surrogate_posterior = tfp.experimental.vi.build_factored_surrogate_posterior(
event_shape=model.event_shape_tensor()[:-1], # Omit the observed `y`.
bijector=[tfb.Softplus(), # Rate is positive.
tfb.Softplus()]) # Concentration is positive.
```
This creates a trainable joint distribution, defined by variables in
`surrogate_posterior.trainable_variables`. We use `fit_surrogate_posterior`
to fit this distribution by minimizing a divergence to the true posterior.
```python
y = [0.2, 0.5, 0.3, 0.7]
losses = tfp.vi.fit_surrogate_posterior(
lambda rate, concentration: model.log_prob([rate, concentration, y]),
surrogate_posterior=surrogate_posterior,
num_steps=100,
optimizer=tf.optimizers.Adam(0.1),
sample_size=10)
# After optimization, samples from the surrogate will approximate
# samples from the true posterior.
samples = surrogate_posterior.sample(100)
posterior_mean = [tf.reduce_mean(x) for x in samples] # mean ~= [1.1, 2.1]
posterior_std = [tf.math.reduce_std(x) for x in samples] # std ~= [0.3, 0.8]
```
If we wanted to initialize the optimization at a specific location, we can
specify initial parameters when we build the surrogate posterior. Note that
these parameterize the distribution(s) over unconstrained values,
so we need to transform our desired constrained locations using the inverse
of the constraining bijector(s).
```python
surrogate_posterior = tfp.experimental.vi.build_factored_surrogate_posterior(
event_shape=tf.nest.map_fn(tf.shape, initial_loc),
bijector={'concentration': tfb.Softplus(), # Rate is positive.
'rate': tfb.Softplus()} # Concentration is positive.
initial_parameters={
'concentration': {'loc': tfb.Softplus().inverse(0.4), 'scale': 1e-2},
'rate': {'loc': tfb.Softplus().inverse(0.2), 'scale': 1e-2}})
```
"""
with tf.name_scope(name or 'build_factored_surrogate_posterior'):
# Convert event shapes to Tensors.
shallow_structure = _get_event_shape_shallow_structure(event_shape)
event_shape = nest.map_structure_up_to(
shallow_structure,
lambda s: tf.convert_to_tensor(s, dtype=tf.int32), event_shape)
if nest.is_nested(bijector):
event_space_bijector = joint_map.JointMap(
nest.map_structure(
lambda b: identity.Identity() if b is None else b,
nest_util.coerce_structure(event_shape, bijector)),
validate_args=validate_args)
else:
event_space_bijector = bijector
if event_space_bijector is None:
unconstrained_event_shape = event_shape
else:
unconstrained_event_shape = (
event_space_bijector.inverse_event_shape_tensor(event_shape))
unconstrained_batch_and_event_shape = tf.nest.map_structure(
lambda s: ps.concat([batch_shape, s], axis=0),
unconstrained_event_shape)
base_distribution_cls = nest_util.broadcast_structure(
event_shape, base_distribution_cls)
try:
# Check that we have initial parameters for each event part.
nest.assert_shallow_structure(event_shape, initial_parameters)
except (ValueError, TypeError):
# If not, broadcast the parameters to match the event structure.
# We do this manually rather than using `nest_util.broadcast_structure`
# because the initial parameters can themselves be structures (dicts).
initial_parameters = nest.map_structure(
lambda x: initial_parameters, event_shape)
unconstrained_trainable_distributions = yield from (
nest_util.map_structure_coroutine(
trainable._make_trainable, # pylint: disable=protected-access
cls=base_distribution_cls,
initial_parameters=initial_parameters,
batch_and_event_shape=unconstrained_batch_and_event_shape,
parameter_dtype=nest_util.broadcast_structure(
event_shape, dtype),
_up_to=event_shape))
unconstrained_trainable_distribution = (
joint_distribution_util.
independent_joint_distribution_from_structure(
unconstrained_trainable_distributions,
batch_ndims=ps.rank_from_shape(batch_shape),
validate_args=validate_args))
if event_space_bijector is None:
return unconstrained_trainable_distribution
return transformed_distribution.TransformedDistribution(
unconstrained_trainable_distribution, event_space_bijector)
def __init__(self,
input_tensor_spec,
preprocessing_layers=None,
preprocessing_combiner=None,
action_fc_layer_params=(200, ),
joint_fc_layer_params=(100, ),
lstm_size=(40, ),
output_fc_layer_params=(200, 100),
activation_fn=tf.keras.activations.relu,
name='MultiInputsCriticRnnNetwork'):
"""Creates an instance of `MultiInputsCriticRnnNetwork`.
Args:
input_tensor_spec: A tuple of (observation, action) each of type
`tensor_spec.TensorSpec` representing the inputs.
preprocessing_layers: (Optional.) A nest of `tf.keras.layers.Layer`
representing preprocessing for the different observations.
All of these layers must not be already built. For more details see
the documentation of `networks.EncodingNetwork`.
preprocessing_combiner: (Optional.) A keras layer that takes a flat list
of tensors and combines them. Good options include
`tf.keras.layers.Add` and `tf.keras.layers.Concatenate(axis=-1)`.
This layer must not be already built. For more details see
the documentation of `networks.EncodingNetwork`.
action_fc_layer_params: Optional list of parameters for a fully_connected
layer to apply to the actions, where each item is the number of units
in the layer.
joint_fc_layer_params: Optional list of parameters for a fully_connected
layer to apply after merging observations and actions, where each item
is the number of units in the layer.
lstm_size: An iterable of ints specifying the LSTM cell sizes to use.
output_fc_layer_params: Optional list of fully_connected parameters, where
each item is the number of units in the layer. This is applied after the
LSTM cell.
activation_fn: Activation function, e.g. tf.nn.relu, slim.leaky_relu, ...
name: A string representing name of the network.
Returns:
A tf.float32 Tensor of q-values.
Raises:
ValueError: If `observation_spec` or `action_spec` contains more than one
item.
"""
observation_spec, action_spec = input_tensor_spec
if len(tf.nest.flatten(action_spec)) > 1:
raise ValueError(
'Only a single action is supported by this network.')
if preprocessing_layers is None:
flat_preprocessing_layers = None
else:
flat_preprocessing_layers = [
_copy_layer(layer)
for layer in tf.nest.flatten(preprocessing_layers)
]
# Assert shallow structure is the same. This verifies preprocessing
# layers can be applied on expected input nests.
observation_nest = observation_spec
# Given the flatten on preprocessing_layers above we need to make sure
# input_tensor_spec is a sequence for the shallow_structure check below
# to work.
if not nest.is_sequence(observation_spec):
observation_nest = [observation_spec]
nest.assert_shallow_structure(preprocessing_layers,
observation_nest,
check_types=False)
if (len(tf.nest.flatten(observation_spec)) > 1
and preprocessing_combiner is None):
raise ValueError(
'preprocessing_combiner layer is required when more than 1 '
'observation_spec is provided.')
if preprocessing_combiner is not None:
preprocessing_combiner = _copy_layer(preprocessing_combiner)
action_layers = utils.mlp_layers(
None,
action_fc_layer_params,
activation_fn=activation_fn,
kernel_initializer=tf.compat.v1.keras.initializers.VarianceScaling(
scale=1. / 3., mode='fan_in', distribution='uniform'),
name='action_encoding')
joint_layers = utils.mlp_layers(
None,
joint_fc_layer_params,
activation_fn=activation_fn,
kernel_initializer=tf.compat.v1.keras.initializers.VarianceScaling(
scale=1. / 3., mode='fan_in', distribution='uniform'),
name='joint_mlp')
# Create RNN cell
if len(lstm_size) == 1:
cell = tf.keras.layers.LSTMCell(lstm_size[0])
else:
cell = tf.keras.layers.StackedRNNCells(
[tf.keras.layers.LSTMCell(size) for size in lstm_size])
counter = [-1]
def create_spec(size):
counter[0] += 1
return tensor_spec.TensorSpec(size,
dtype=tf.float32,
name='network_state_%d' % counter[0])
state_spec = tf.nest.map_structure(create_spec, cell.state_size)
output_layers = utils.mlp_layers(
fc_layer_params=output_fc_layer_params, name='output')
output_layers.append(
tf.keras.layers.Dense(
1,
activation=None,
kernel_initializer=tf.keras.initializers.RandomUniform(
minval=-0.003, maxval=0.003),
name='value'))
super(MultiInputsCriticRnnNetwork,
self).__init__(input_tensor_spec=input_tensor_spec,
state_spec=state_spec,
name=name)
self._action_layers = action_layers
self._joint_layers = joint_layers
self._dynamic_unroll = dynamic_unroll_layer.DynamicUnroll(cell)
self._output_layers = output_layers
self._preprocessing_nest = tf.nest.map_structure(
lambda l: None, preprocessing_layers)
self._flat_preprocessing_layers = flat_preprocessing_layers
self._preprocessing_combiner = preprocessing_combiner
def __init__(self,
input_tensor_spec,
output_tensor_spec,
preprocessing_layers=None,
preprocessing_combiner=None,
conv_layer_params=None,
input_fc_layer_params=(200, 100),
lstm_size=(40,),
output_fc_layer_params=(200, 100),
activation_fn=tf.keras.activations.relu,
name='MultiInputsActorRnnNetwork'):
"""Creates an instance of `MultiInputsActorRnnNetwork`.
Args:
input_tensor_spec: A nest of `tensor_spec.TensorSpec` representing the
input observations.
output_tensor_spec: A nest of `tensor_spec.BoundedTensorSpec` representing
the actions.
preprocessing_layers: (Optional.) A nest of `tf.keras.layers.Layer`
representing preprocessing for the different observations.
All of these layers must not be already built. For more details see
the documentation of `networks.EncodingNetwork`.
preprocessing_combiner: (Optional.) A keras layer that takes a flat list
of tensors and combines them. Good options include
`tf.keras.layers.Add` and `tf.keras.layers.Concatenate(axis=-1)`.
This layer must not be already built. For more details see
the documentation of `networks.EncodingNetwork`.
conv_layer_params: Optional list of convolution layers parameters, where
each item is a length-three tuple indicating (filters, kernel_size,
stride).
input_fc_layer_params: Optional list of fully_connected parameters, where
each item is the number of units in the layer. This is applied before
the LSTM cell.
lstm_size: An iterable of ints specifying the LSTM cell sizes to use.
output_fc_layer_params: Optional list of fully_connected parameters, where
each item is the number of units in the layer. This is applied after the
LSTM cell.
activation_fn: Activation function, e.g. tf.nn.relu, slim.leaky_relu, ...
name: A string representing name of the network.
Returns:
A nest of action tensors matching the action_spec.
Raises:
ValueError: If `input_tensor_spec` contains more than one observation.
"""
observation_spec = input_tensor_spec
if preprocessing_layers is None:
flat_preprocessing_layers = None
else:
flat_preprocessing_layers = [
_copy_layer(layer) for layer in tf.nest.flatten(preprocessing_layers)
]
# Assert shallow structure is the same. This verifies preprocessing
# layers can be applied on expected input nests.
observation_nest = observation_spec
# Given the flatten on preprocessing_layers above we need to make sure
# input_tensor_spec is a sequence for the shallow_structure check below
# to work.
if not nest.is_sequence(observation_spec):
observation_nest = [observation_spec]
nest.assert_shallow_structure(
preprocessing_layers, observation_nest, check_types=False)
if (len(tf.nest.flatten(observation_spec)) > 1 and
preprocessing_combiner is None):
raise ValueError(
'preprocessing_combiner layer is required when more than 1 '
'observation_spec is provided.')
if preprocessing_combiner is not None:
preprocessing_combiner = _copy_layer(preprocessing_combiner)
input_layers = utils.mlp_layers(
conv_layer_params,
input_fc_layer_params,
activation_fn=activation_fn,
kernel_initializer=tf.compat.v1.keras.initializers.glorot_uniform(),
name='input_mlp')
# Create RNN cell
if len(lstm_size) == 1:
cell = tf.keras.layers.LSTMCell(lstm_size[0])
else:
cell = tf.keras.layers.StackedRNNCells(
[tf.keras.layers.LSTMCell(size) for size in lstm_size])
state_spec = tf.nest.map_structure(
functools.partial(
tensor_spec.TensorSpec, dtype=tf.float32,
name='network_state_spec'), list(cell.state_size))
output_layers = utils.mlp_layers(fc_layer_params=output_fc_layer_params,
name='output')
flat_action_spec = tf.nest.flatten(output_tensor_spec)
action_layers = [
tf.keras.layers.Dense(
single_action_spec.shape.num_elements(),
activation=tf.keras.activations.tanh,
kernel_initializer=tf.keras.initializers.RandomUniform(
minval=-0.003, maxval=0.003),
name='action') for single_action_spec in flat_action_spec
]
super(MultiInputsActorRnnNetwork, self).__init__(
input_tensor_spec=input_tensor_spec,
state_spec=state_spec,
name=name)
self._output_tensor_spec = output_tensor_spec
self._flat_action_spec = flat_action_spec
self._conv_layer_params = conv_layer_params
self._input_layers = input_layers
self._dynamic_unroll = dynamic_unroll_layer.DynamicUnroll(cell)
self._output_layers = output_layers
self._action_layers = action_layers
self._preprocessing_nest = tf.nest.map_structure(lambda l: None,
preprocessing_layers)
self._flat_preprocessing_layers = flat_preprocessing_layers
self._preprocessing_combiner = preprocessing_combiner
def testAssertShallowStructure(self):
inp_ab = ["a", "b"]
inp_abc = ["a", "b", "c"]
expected_message = (
"The two structures don't have the same sequence length. Input "
"structure has length 2, while shallow structure has length 3.")
with self.assertRaisesRegexp(ValueError, expected_message):
nest.assert_shallow_structure(inp_abc, inp_ab)
inp_ab1 = [(1, 1), (2, 2)]
inp_ab2 = [[1, 1], [2, 2]]
expected_message = (
"The two structures don't have the same sequence type. Input structure "
"has type <(type|class) 'tuple'>, while shallow structure has type "
"<(type|class) 'list'>.")
with self.assertRaisesRegexp(TypeError, expected_message):
nest.assert_shallow_structure(inp_ab2, inp_ab1)
nest.assert_shallow_structure(inp_ab2, inp_ab1, check_types=False)
inp_ab1 = {"a": (1, 1), "b": {"c": (2, 2)}}
inp_ab2 = {"a": (1, 1), "b": {"d": (2, 2)}}
expected_message = (
r"The two structures don't have the same keys. Input "
r"structure has keys \['c'\], while shallow structure has "
r"keys \['d'\].")
with self.assertRaisesRegexp(ValueError, expected_message):
nest.assert_shallow_structure(inp_ab2, inp_ab1)
inp_ab = collections.OrderedDict([("a", 1), ("b", (2, 3))])
inp_ba = collections.OrderedDict([("b", (2, 3)), ("a", 1)])
nest.assert_shallow_structure(inp_ab, inp_ba)
# This assertion is expected to pass: two namedtuples with the same
# name and field names are considered to be identical.
inp_shallow = NestTest.SameNameab(1, 2)
inp_deep = NestTest.SameNameab2(1, [1, 2, 3])
nest.assert_shallow_structure(inp_shallow, inp_deep, check_types=False)
nest.assert_shallow_structure(inp_shallow, inp_deep, check_types=True)
def testNestAssertShallowStructure(self, s1, s2, expand_composites=True): nest.assert_shallow_structure(s1, s2, expand_composites=expand_composites)
def testAssertShallowStructure(self):
inp_ab = ["a", "b"]
inp_abc = ["a", "b", "c"]
expected_message = (
"The two structures don't have the same sequence length. Input "
"structure has length 2, while shallow structure has length 3.")
with self.assertRaisesRegexp(ValueError, expected_message):
nest.assert_shallow_structure(inp_abc, inp_ab)
inp_ab1 = [(1, 1), (2, 2)]
inp_ab2 = [[1, 1], [2, 2]]
expected_message = (
"The two structures don't have the same sequence type. Input structure "
"has type <(type|class) 'tuple'>, while shallow structure has type "
"<(type|class) 'list'>.")
with self.assertRaisesRegexp(TypeError, expected_message):
nest.assert_shallow_structure(inp_ab2, inp_ab1)
nest.assert_shallow_structure(inp_ab2, inp_ab1, check_types=False)
inp_ab1 = {"a": (1, 1), "b": {"c": (2, 2)}}
inp_ab2 = {"a": (1, 1), "b": {"d": (2, 2)}}
expected_message = (
r"The two structures don't have the same keys. Input "
r"structure has keys \['c'\], while shallow structure has "
r"keys \['d'\].")
with self.assertRaisesRegexp(ValueError, expected_message):
nest.assert_shallow_structure(inp_ab2, inp_ab1)
inp_ab = collections.OrderedDict([("a", 1), ("b", (2, 3))])
inp_ba = collections.OrderedDict([("b", (2, 3)), ("a", 1)])
nest.assert_shallow_structure(inp_ab, inp_ba)
# This assertion is expected to pass: two namedtuples with the same
# name and field names are considered to be identical.
inp_shallow = NestTest.SameNameab(1, 2)
inp_deep = NestTest.SameNameab2(1, [1, 2, 3])
nest.assert_shallow_structure(inp_shallow, inp_deep, check_types=False)
nest.assert_shallow_structure(inp_shallow, inp_deep, check_types=True)
def assert_same_shallow_tree(shallow, tree):
"""Asserts that `tree` has the same shallow structure as `shallow`."""
nest.assert_shallow_structure(shallow, tree)
