def _actor_loss_debug_summaries(self, actor_loss, actions, log_pi,
target_q_values, time_steps):
if self._debug_summaries:
common.generate_tensor_summaries('actor_loss', actor_loss,
self.train_step_counter)
try:
for name, action in nest_utils.flatten_with_joined_paths(
actions):
common.generate_tensor_summaries(name, action,
self.train_step_counter)
except ValueError:
pass # Guard against internal SAC variants that do not directly
# generate actions.
common.generate_tensor_summaries('log_pi', log_pi,
self.train_step_counter)
tf.compat.v2.summary.scalar(
name='entropy_avg',
data=-tf.reduce_mean(input_tensor=log_pi),
step=self.train_step_counter)
common.generate_tensor_summaries('target_q_values',
target_q_values,
self.train_step_counter)
batch_size = nest_utils.get_outer_shape(time_steps,
self._time_step_spec)[0]
policy_state = self._train_policy.get_initial_state(batch_size)
action_distribution = self._train_policy.distribution(
time_steps, policy_state).action
if isinstance(action_distribution, tfp.distributions.Normal):
common.generate_tensor_summaries('act_mean',
action_distribution.loc,
self.train_step_counter)
common.generate_tensor_summaries('act_stddev',
action_distribution.scale,
self.train_step_counter)
elif isinstance(action_distribution,
tfp.distributions.Categorical):
common.generate_tensor_summaries('act_mode',
action_distribution.mode(),
self.train_step_counter)
try:
for name, action_dist in nest_utils.flatten_with_joined_paths(
action_distribution):
common.generate_tensor_summaries('entropy_' + name,
action_dist.entropy(),
self.train_step_counter)
except NotImplementedError:
pass # Some distributions do not have an analytic entropy.
def __init__(self, time_step_spec=None, action_spec=None, batch_size=1):
"""Initialize instances of `BanditTFEnvironment`.
Args:
time_step_spec: A `TimeStep` namedtuple containing `TensorSpec`s
defining the tensors returned by
`step()` (step_type, reward, discount, and observation).
action_spec: A nest of BoundedTensorSpec representing the actions of the
environment.
batch_size: The batch size expected for the actions and observations.
"""
self._reset_called = tf.compat.v2.Variable(False,
trainable=False,
name='reset_called')
def _variable_from_spec(name, spec):
full_shape = [batch_size] + spec.shape.as_list()
if not name:
name = 'spec_var'
return common.create_variable(name,
0,
shape=full_shape,
dtype=spec.dtype)
paths_and_specs = nest_utils.flatten_with_joined_paths(time_step_spec)
variables = [
_variable_from_spec(path, spec) for path, spec in paths_and_specs
]
self._time_step_variables = tf.nest.pack_sequence_as(
time_step_spec, variables)
super(BanditTFEnvironment,
self).__init__(time_step_spec=time_step_spec,
action_spec=action_spec,
batch_size=batch_size)
def __init__(self,
time_step_spec,
action_spec,
batch_size=1,
episode_end_probability=0.1):
"""Initializes the environment.
Args:
time_step_spec: A `TimeStep` namedtuple containing `TensorSpec`s defining
the Tensors returned by `step()` (step_type, reward, discount, and
observation).
action_spec: A nest of BoundedTensorSpec representing the actions of the
environment.
batch_size: The batch size expected for the actions and observations.
episode_end_probability: Probability an episode will end when the
environment is stepped.
"""
super(RandomTFEnvironment, self).__init__(
time_step_spec, action_spec, batch_size=batch_size)
self._episode_end_probability = episode_end_probability
def _variable_from_spec(name, spec):
full_shape = [batch_size] + spec.shape.as_list()
if not name:
name = "spec_var"
return common.create_variable(name, 0, shape=full_shape, dtype=spec.dtype)
paths_and_specs = nest_utils.flatten_with_joined_paths(time_step_spec)
variables = [
_variable_from_spec(path, spec) for path, spec in paths_and_specs
]
self._time_step_variables = tf.nest.pack_sequence_as(
time_step_spec, variables)
def get_example_encoder(spec):
"""Get example encoder function for the given spec.
Given a spec, returns an example encoder function. The example encoder
function takes a nest of np.array feature values as input and returns a
TF Example proto.
Example:
spec = {
'lidar': array_spec.ArraySpec((900,), np.float32),
'joint_positions': {
'arm': array_spec.ArraySpec((7,), np.float32),
'hand': array_spec.BoundedArraySpec((3, 3), np.int32, -1, 1)
},
}
example_encoder = get_example_encoder(spec)
serialized = example_encoder({
'lidar': np.zeros((900,), np.float32),
'joint_positions': {
'arm': np.array([0.0, 1.57, 0.707, 0.2, 0.0, -1.57, 0.0],
np.float32),
'hand': np.ones((3, 3), np.int32)
},
})
The returned example encoder function requires that the feature nest passed
has the shape and exact dtype specified in the spec. For example, it is
an error to pass an array with np.float64 dtype where np.float32 is expected.
Args:
spec: list/tuple/nest of ArraySpecs describing a single example.
Returns:
Function
```python
encoder(features_nest of np.arrays) -> tf.train.Example
```
"""
# pylint: disable=g-complex-comprehension
feature_encoders = [
(path, _get_feature_encoder(spec.shape, spec.dtype))
for (path, spec) in nest_utils.flatten_with_joined_paths(spec)
]
# pylint: enable=g-complex-comprehension
def _example_encoder(features_nest):
flat_features = tf.nest.flatten(features_nest)
feature_dict = {
path: feature_encoder(feature)
for feature, (
path, feature_encoder) in zip(flat_features, feature_encoders)
}
return tf.train.Example(features=tf.train.Features(
feature=feature_dict))
return _example_encoder
def get_example_decoder(example_spec, batched=False, compress_image=False):
"""Get an example decoder function for a nested spec.
Given a spec, returns an example decoder function. The decoder function parses
string serialized example protos into tensors according to the given spec.
Args:
example_spec: list/tuple/nest of ArraySpecs describing a single example.
batched: Boolean indicating if the decoder will receive batches of
serialized data.
compress_image: Whether to decompress image. It is assumed that any uint8
tensor of rank 3 with shape (w,h,c) is an image.
If the tensor was compressed in the encoder, it needs to be decompressed.
Returns:
Function
```python
decoder(serialized_proto: tf.tensor[string]) -> example_spec nest of tensors
```
"""
features_dict = {}
parsers = []
for (path, spec) in nest_utils.flatten_with_joined_paths(example_spec):
feature, parser = _get_feature_parser(spec.shape, spec.dtype,
compress_image)
features_dict[path] = feature
parsers.append((path, parser))
def _example_decoder(serialized):
"""Parses string serialized example protos into tensors."""
if batched:
raw_features = tf.io.parse_example(serialized=serialized,
features=features_dict)
decoded_features = []
dtypes = [s.dtype for s in tf.nest.flatten(example_spec)]
if len(parsers) != len(dtypes):
raise ValueError(
'Decoding failed: Number of parsers (%d) does not match the number '
'of items in the example_spec (%d)' %
(len(parsers), len(dtypes)))
for (path, parser), dtype in zip(parsers, dtypes):
decoded_features.append(
tf.map_fn(parser, raw_features[path], dtype=dtype))
return tf.nest.pack_sequence_as(example_spec, decoded_features)
else:
raw_features = tf.io.parse_single_example(serialized=serialized,
features=features_dict)
return tf.nest.pack_sequence_as(
example_spec,
[parser(raw_features[path]) for path, parser in parsers])
return _example_decoder
def parameters_to_dict(value: Params) -> Mapping[Text, Any]:
"""Converts `value` to a nested `dict` (excluding all `type_` info).
Sub-dicts represent `Params` objects; keys represent flattened nest structures
in `value.params`.
Example:
```python
scale_matrix = tf.Variable([[1.0, 2.0], [-1.0, 0.0]])
d = tfp.distributions.MultivariateNormalDiag(
loc=[1.0, 1.0], scale_diag=[2.0, 3.0], validate_args=True)
b = tfp.bijectors.ScaleMatvecLinearOperator(
scale=tf.linalg.LinearOperatorFullMatrix(matrix=scale_matrix),
adjoint=True)
b_d = b(d)
p = utils.get_parameters(b_d)
params_dict = utils.parameters_to_dict(p)
```
results in the nested dictionary:
```python
{"bijector": {"adjoint": True,
"scale": {"matrix": scale_matrix}},
"distribution": {"validate_args": True,
# These are deeply nested because we passed lists
# intead of numpy arrays for `loc` and `scale_diag`.
"scale_diag:0": 2.0,
"scale_diag:1": 3.0,
"loc:0": 1.0,
"loc:1": 1.0}
}
```
The dictionary may then be modified or updated (e.g., in place), and converted
back to a `Params` object using `merge_to_parameters_from_dict`.
Args:
value: The (possibly recursively defined) `Params`.
Returns:
A `dict` mapping `value.params` to flattened key/value pairs. Any
sub-`Params` objects become nested dicts.
"""
convert = lambda p: parameters_to_dict(p) if isinstance(p, Params) else p
output_entries = {}
for k, v in value.params.items():
if tf.nest.is_nested(v):
flattened_params = nest_utils.flatten_with_joined_paths(v)
for (param_k, param_v) in flattened_params:
key = "{}:{}".format(k, param_k)
output_entries[key] = convert(param_v)
else:
output_entries[k] = convert(v)
return output_entries
def get_example_decoder(example_spec, batched=False):
"""Get an example decoder function for a nested spec.
Given a spec, returns an example decoder function. The decoder function parses
string serialized example protos into tensors according to the given spec.
Args:
example_spec: list/tuple/nest of ArraySpecs describing a single example.
batched: Boolean indicating if the decoder will receive batches of
serialized data.
Returns:
Function
```python
decoder(serialized_proto: tf.tensor[string]) -> example_spec nest of tensors
```
"""
features_dict = {}
parsers = []
for (path, spec) in nest_utils.flatten_with_joined_paths(example_spec):
feature, parser = _get_feature_parser(spec.shape, spec.dtype)
features_dict[path] = feature
parsers.append((path, parser))
def _example_decoder(serialized):
"""Parses string serialized example protos into tensors."""
if batched:
raw_features = tf.io.parse_example(serialized=serialized,
features=features_dict)
decoded_features = []
dtypes = [s.dtype for s in tf.nest.flatten(example_spec)]
for (path, parser), dtype in zip(parsers, dtypes):
decoded_features.append(
tf.map_fn(parser, raw_features[path], dtype=dtype))
return tf.nest.pack_sequence_as(example_spec, decoded_features)
else:
raw_features = tf.io.parse_single_example(serialized=serialized,
features=features_dict)
return tf.nest.pack_sequence_as(
example_spec,
[parser(raw_features[path]) for path, parser in parsers])
return _example_decoder
def _rename_spec_with_nest_paths(spec):
renamed_spec = [
tf.TensorSpec(shape=s.shape, name=path, dtype=s.dtype)
for path, s in nest_utils.flatten_with_joined_paths(spec)
]
return tf.nest.pack_sequence_as(spec, renamed_spec)
def parameters_to_dict(
value: Params,
tensors_only: bool = False) -> Mapping[Text, Any]:
"""Converts `value` to a nested `dict` (excluding all `type_` info).
Sub-dicts represent `Params` objects; keys represent flattened nest structures
in `value.params`.
Example:
```python
scale_matrix = tf.Variable([[1.0, 2.0], [-1.0, 0.0]])
d = tfp.distributions.MultivariateNormalDiag(
loc=[1.0, 1.0], scale_diag=[2.0, 3.0], validate_args=True)
b = tfp.bijectors.ScaleMatvecLinearOperator(
scale=tf.linalg.LinearOperatorFullMatrix(matrix=scale_matrix),
adjoint=True)
b_d = b(d)
p = utils.get_parameters(b_d)
params_dict = utils.parameters_to_dict(p, tensors_only)
```
results in the nested dictionary, if `tensors_only=False`:
```python
{
"bijector": {"adjoint": True,
"scale": {"matrix": scale_matrix}},
"distribution": {"validate_args": True,
# These are deeply nested because we passed lists
# intead of numpy arrays for `loc` and `scale_diag`.
"scale_diag:0": 2.0,
"scale_diag:1": 3.0,
"loc:0": 1.0,
"loc:1": 1.0}
}
```
results if `tensors_only=True`:
```python
{
"bijector": {"scale": {"matrix": scale_matrix}},
"distribution": {},
}
```
The dictionary may then be modified or updated (e.g., in place), and converted
back to a `Params` object using `merge_to_parameters_from_dict`.
Args:
value: The (possibly recursively defined) `Params`.
tensors_only: Whether to include all parameters or only
tensors and TypeSpecs. If `True`, then only
`tf.Tensor` and `tf.TypeSpec` leaf parameters are emitted.
Returns:
A `dict` mapping `value.params` to flattened key/value pairs. Any
sub-`Params` objects become nested dicts.
"""
output_entries = {}
# All tensors and TensorSpecs will be included for purposes of this test.
is_tensor_or_spec = (
lambda p: tf.is_tensor(p) or isinstance(p, tf.TypeSpec))
def convert(key, p):
if isinstance(p, Params):
output_entries[key] = parameters_to_dict(p, tensors_only)
elif not tensors_only or is_tensor_or_spec(p):
output_entries[key] = p
for k, v in value.params.items():
if tf.nest.is_nested(v):
flattened_params = nest_utils.flatten_with_joined_paths(v)
for (param_k, param_v) in flattened_params:
key = "{}:{}".format(k, param_k)
convert(key, param_v)
else:
convert(k, v)
return output_entries
def get_example_encoder(spec, compress_image=False, image_quality=95):
"""Get example encoder function for the given spec.
Given a spec, returns an example encoder function. The example encoder
function takes a nest of np.array feature values as input and returns a
TF Example proto.
Example:
spec = {
'lidar': array_spec.ArraySpec((900,), np.float32),
'joint_positions': {
'arm': array_spec.ArraySpec((7,), np.float32),
'hand': array_spec.BoundedArraySpec((3, 3), np.int32, -1, 1)
},
}
example_encoder = get_example_encoder(spec)
serialized = example_encoder({
'lidar': np.zeros((900,), np.float32),
'joint_positions': {
'arm': np.array([0.0, 1.57, 0.707, 0.2, 0.0, -1.57, 0.0],
np.float32),
'hand': np.ones((3, 3), np.int32)
},
})
The returned example encoder function requires that the feature nest passed
has the shape and exact dtype specified in the spec. For example, it is
an error to pass an array with np.float64 dtype where np.float32 is expected.
Args:
spec: list/tuple/nest of ArraySpecs describing a single example.
compress_image: Whether to compress image. It is assumed that any uint8
tensor of rank 3 with shape (w,h,c) is an image.
image_quality: An optional int. Defaults to 95. Quality of the compression
from 0 to 100 (higher is better and slower).
Returns:
Function
```python
encoder(features_nest of np.arrays) -> tf.train.Example
```
"""
# pylint: disable=g-complex-comprehension
feature_encoders = [
(path,
_get_feature_encoder(spec.shape, spec.dtype, compress_image,
image_quality))
for (path, spec) in nest_utils.flatten_with_joined_paths(spec)
]
# pylint: enable=g-complex-comprehension
def _example_encoder(features_nest):
flat_features = tf.nest.flatten(features_nest)
if len(flat_features) != len(feature_encoders):
raise ValueError(
'Encoding failed: The number of items to encode does not match the '
'number of encoders generated. Num features %d, num_encoders:%d'
% (len(flat_features), len(feature_encoders)))
feature_dict = {
path: feature_encoder(feature)
for feature, (
path, feature_encoder) in zip(flat_features, feature_encoders)
}
return tf.train.Example(features=tf.train.Features(
feature=feature_dict))
return _example_encoder