research/efficient-hrl/utils/utils.py
# Copyright 2018 The TensorFlow Authors All Rights Reserved.
#
# 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
#
# http://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.
# ==============================================================================
"""TensorFlow utility functions.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from copy import deepcopy
import tensorflow as tf
from tf_agents import specs
from tf_agents.utils import common
_tf_print_counts = dict()
_tf_print_running_sums = dict()
_tf_print_running_counts = dict()
_tf_print_ids = 0
def get_contextual_env_base(env_base, begin_ops=None, end_ops=None):
"""Wrap env_base with additional tf ops."""
# pylint: disable=protected-access
def init(self_, env_base):
self_._env_base = env_base
attribute_list = ["_render_mode", "_gym_env"]
for attribute in attribute_list:
if hasattr(env_base, attribute):
setattr(self_, attribute, getattr(env_base, attribute))
if hasattr(env_base, "physics"):
self_._physics = env_base.physics
elif hasattr(env_base, "gym"):
class Physics(object):
def render(self, *args, **kwargs):
return env_base.gym.render("rgb_array")
physics = Physics()
self_._physics = physics
self_.physics = physics
def set_sess(self_, sess):
self_._sess = sess
if hasattr(self_._env_base, "set_sess"):
self_._env_base.set_sess(sess)
def begin_episode(self_):
self_._env_base.reset()
if begin_ops is not None:
self_._sess.run(begin_ops)
def end_episode(self_):
self_._env_base.reset()
if end_ops is not None:
self_._sess.run(end_ops)
return type("ContextualEnvBase", (env_base.__class__,), dict(
__init__=init,
set_sess=set_sess,
begin_episode=begin_episode,
end_episode=end_episode,
))(env_base)
# pylint: enable=protected-access
def merge_specs(specs_):
"""Merge TensorSpecs.
Args:
specs_: List of TensorSpecs to be merged.
Returns:
a TensorSpec: a merged TensorSpec.
"""
shape = specs_[0].shape
dtype = specs_[0].dtype
name = specs_[0].name
for spec in specs_[1:]:
assert shape[1:] == spec.shape[1:], "incompatible shapes: %s, %s" % (
shape, spec.shape)
assert dtype == spec.dtype, "incompatible dtypes: %s, %s" % (
dtype, spec.dtype)
shape = merge_shapes((shape, spec.shape), axis=0)
return specs.TensorSpec(
shape=shape,
dtype=dtype,
name=name,
)
def merge_shapes(shapes, axis=0):
"""Merge TensorShapes.
Args:
shapes: List of TensorShapes to be merged.
axis: optional, the axis to merge shaped.
Returns:
a TensorShape: a merged TensorShape.
"""
assert len(shapes) > 1
dims = deepcopy(shapes[0].dims)
for shape in shapes[1:]:
assert shapes[0].ndims == shape.ndims
dims[axis] += shape.dims[axis]
return tf.TensorShape(dims=dims)
def get_all_vars(ignore_scopes=None):
"""Get all tf variables in scope.
Args:
ignore_scopes: A list of scope names to ignore.
Returns:
A list of all tf variables in scope.
"""
all_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
all_vars = [var for var in all_vars if ignore_scopes is None or not
any(var.name.startswith(scope) for scope in ignore_scopes)]
return all_vars
def clip(tensor, range_=None):
"""Return a tf op which clips tensor according to range_.
Args:
tensor: A Tensor to be clipped.
range_: None, or a tuple representing (minval, maxval)
Returns:
A clipped Tensor.
"""
if range_ is None:
return tf.identity(tensor)
elif isinstance(range_, (tuple, list)):
assert len(range_) == 2
return tf.clip_by_value(tensor, range_[0], range_[1])
else: raise NotImplementedError("Unacceptable range input: %r" % range_)
def clip_to_bounds(value, minimum, maximum):
"""Clips value to be between minimum and maximum.
Args:
value: (tensor) value to be clipped.
minimum: (numpy float array) minimum value to clip to.
maximum: (numpy float array) maximum value to clip to.
Returns:
clipped_value: (tensor) `value` clipped to between `minimum` and `maximum`.
"""
value = tf.minimum(value, maximum)
return tf.maximum(value, minimum)
clip_to_spec = common.clip_to_spec
def _clip_to_spec(value, spec):
"""Clips value to a given bounded tensor spec.
Args:
value: (tensor) value to be clipped.
spec: (BoundedTensorSpec) spec containing min. and max. values for clipping.
Returns:
clipped_value: (tensor) `value` clipped to be compatible with `spec`.
"""
return clip_to_bounds(value, spec.minimum, spec.maximum)
join_scope = common.join_scope
def _join_scope(parent_scope, child_scope):
"""Joins a parent and child scope using `/`, checking for empty/none.
Args:
parent_scope: (string) parent/prefix scope.
child_scope: (string) child/suffix scope.
Returns:
joined scope: (string) parent and child scopes joined by /.
"""
if not parent_scope:
return child_scope
if not child_scope:
return parent_scope
return '/'.join([parent_scope, child_scope])
def assign_vars(vars_, values):
"""Returns the update ops for assigning a list of vars.
Args:
vars_: A list of variables.
values: A list of tensors representing new values.
Returns:
A list of update ops for the variables.
"""
return [var.assign(value) for var, value in zip(vars_, values)]
def identity_vars(vars_):
"""Return the identity ops for a list of tensors.
Args:
vars_: A list of tensors.
Returns:
A list of identity ops.
"""
return [tf.identity(var) for var in vars_]
def tile(var, batch_size=1):
"""Return tiled tensor.
Args:
var: A tensor representing the state.
batch_size: Batch size.
Returns:
A tensor with shape [batch_size,] + var.shape.
"""
batch_var = tf.tile(
tf.expand_dims(var, 0),
(batch_size,) + (1,) * var.get_shape().ndims)
return batch_var
def batch_list(vars_list):
"""Batch a list of variables.
Args:
vars_list: A list of tensor variables.
Returns:
A list of tensor variables with additional first dimension.
"""
return [tf.expand_dims(var, 0) for var in vars_list]
def tf_print(op,
tensors,
message="",
first_n=-1,
name=None,
sub_messages=None,
print_freq=-1,
include_count=True):
"""tf.Print, but to stdout."""
# TODO(shanegu): `name` is deprecated. Remove from the rest of codes.
global _tf_print_ids
_tf_print_ids += 1
name = _tf_print_ids
_tf_print_counts[name] = 0
if print_freq > 0:
_tf_print_running_sums[name] = [0 for _ in tensors]
_tf_print_running_counts[name] = 0
def print_message(*xs):
"""print message fn."""
_tf_print_counts[name] += 1
if print_freq > 0:
for i, x in enumerate(xs):
_tf_print_running_sums[name][i] += x
_tf_print_running_counts[name] += 1
if (print_freq <= 0 or _tf_print_running_counts[name] >= print_freq) and (
first_n < 0 or _tf_print_counts[name] <= first_n):
for i, x in enumerate(xs):
if print_freq > 0:
del x
x = _tf_print_running_sums[name][i]/_tf_print_running_counts[name]
if sub_messages is None:
sub_message = str(i)
else:
sub_message = sub_messages[i]
log_message = "%s, %s" % (message, sub_message)
if include_count:
log_message += ", count=%d" % _tf_print_counts[name]
tf.logging.info("[%s]: %s" % (log_message, x))
if print_freq > 0:
for i, x in enumerate(xs):
_tf_print_running_sums[name][i] = 0
_tf_print_running_counts[name] = 0
return xs[0]
print_op = tf.py_func(print_message, tensors, tensors[0].dtype)
with tf.control_dependencies([print_op]):
op = tf.identity(op)
return op
periodically = common.periodically
def _periodically(body, period, name='periodically'):
"""Periodically performs a tensorflow op."""
if period is None or period == 0:
return tf.no_op()
if period < 0:
raise ValueError("period cannot be less than 0.")
if period == 1:
return body()
with tf.variable_scope(None, default_name=name):
counter = tf.get_variable(
"counter",
shape=[],
dtype=tf.int64,
trainable=False,
initializer=tf.constant_initializer(period, dtype=tf.int64))
def _wrapped_body():
with tf.control_dependencies([body()]):
return counter.assign(1)
update = tf.cond(
tf.equal(counter, period), _wrapped_body,
lambda: counter.assign_add(1))
return update
soft_variables_update = common.soft_variables_update