def hx_plain():
"""Computes product of Hessian(f) and vector v.
Returns:
tf.Tensor: Symbolic result.
"""
with tf.name_scope('hx_plain',
values=[constraint_grads, params, xs]):
with tf.name_scope('hx_function',
values=[constraint_grads, xs]):
hx_f = tf.reduce_sum(
tf.stack([
tf.reduce_sum(g * x)
for g, x in zip(constraint_grads, xs)
])),
hx_plain_splits = tf.gradients(hx_f,
params,
name='gradients_hx_plain')
for idx, (hx,
param) in enumerate(zip(hx_plain_splits,
params)):
if hx is None:
hx_plain_splits[idx] = tf.zeros_like(param)
return tensor_utils.flatten_tensor_variables(
hx_plain_splits)
def get_opt_output():
"""Helper function to construct graph.
Returns:
list[tf.Tensor]: Loss and gradient tensor.
"""
with tf.name_scope('get_opt_output', values=[loss, params]):
flat_grad = tensor_utils.flatten_tensor_variables(
tf.gradients(loss, params))
return [
tf.cast(loss, tf.float64),
tf.cast(flat_grad, tf.float64)
]
def get_opt_output():
"""Helper function to construct graph.
Returns:
list[tf.Tensor]: Penalized loss and gradient tensor.
"""
with tf.name_scope('get_opt_output'):
grads = tf.gradients(penalized_loss, params)
for idx, (grad, param) in enumerate(zip(grads, params)):
if grad is None:
grads[idx] = tf.zeros_like(param)
flat_grad = tensor_utils.flatten_tensor_variables(grads)
return [
tf.cast(penalized_loss, tf.float64),
tf.cast(flat_grad, tf.float64),
]
def update_opt(
self,
loss,
target,
leq_constraint,
inputs,
extra_inputs=None,
name=None,
constraint_name='constraint',
):
"""Update the optimizer.
Build the functions for computing loss, gradient, and
the constraint value.
Args:
loss (tf.Tensor): Symbolic expression for the loss function.
target (metarl.tf.policies.Policy): A parameterized object to
optimize over.
leq_constraint (tuple[tf.Tensor, float]): A constraint provided
as a tuple (f, epsilon), of the form f(*inputs) <= epsilon.
inputs (list(tf.Tenosr)): A list of symbolic variables as inputs,
which could be subsampled if needed. It is assumed that the
first dimension of these inputs should correspond to the
number of data points.
extra_inputs (list[tf.Tenosr]): A list of symbolic variables as
extra inputs which should not be subsampled.
name (str): Name to be passed to tf.name_scope.
constraint_name (str): A constraint name for prupose of logging
and variable names.
"""
params = target.get_params()
ns_vals = [loss, target, leq_constraint, inputs, extra_inputs, params]
with tf.name_scope(name, 'ConjugateGradientOptimizer', ns_vals):
inputs = tuple(inputs)
if extra_inputs is None:
extra_inputs = tuple()
else:
extra_inputs = tuple(extra_inputs)
constraint_term, constraint_value = leq_constraint
with tf.name_scope('loss_gradients', values=[loss, params]):
grads = tf.gradients(loss, xs=params)
for idx, (grad, param) in enumerate(zip(grads, params)):
if grad is None:
grads[idx] = tf.zeros_like(param)
flat_grad = tensor_utils.flatten_tensor_variables(grads)
self._hvp_approach.update_hvp(f=constraint_term,
target=target,
inputs=inputs + extra_inputs,
reg_coeff=self._reg_coeff,
name='update_opt_' + constraint_name)
self._target = target
self._max_constraint_val = constraint_value
self._constraint_name = constraint_name
self._opt_fun = LazyDict(
f_loss=lambda: tensor_utils.compile_function(
inputs=inputs + extra_inputs,
outputs=loss,
log_name='f_loss',
),
f_grad=lambda: tensor_utils.compile_function(
inputs=inputs + extra_inputs,
outputs=flat_grad,
log_name='f_grad',
),
f_constraint=lambda: tensor_utils.compile_function(
inputs=inputs + extra_inputs,
outputs=constraint_term,
log_name='constraint',
),
f_loss_constraint=lambda: tensor_utils.compile_function(
inputs=inputs + extra_inputs,
outputs=[loss, constraint_term],
log_name='f_loss_constraint',
),
)
def update_hvp(self, f, target, inputs, reg_coeff, name=None):
"""Build the symbolic graph to compute the Hessian-vector product.
Args:
f (tf.Tensor): The function whose Hessian needs to be computed.
target (metarl.tf.policies.Policy): A parameterized object to
optimize over.
inputs (tuple[tf.Tensor]): The inputs for function f.
reg_coeff (float): A small value so that A -> A + reg*I.
name (str): Name to be used in tf.name_scope.
"""
self._target = target
self._reg_coeff = reg_coeff
params = target.get_params()
with tf.name_scope(name, 'FiniteDifferenceHvp',
[f, inputs, params, target]):
constraint_grads = tf.gradients(f,
xs=params,
name='gradients_constraint')
for idx, (grad, param) in enumerate(zip(constraint_grads, params)):
if grad is None:
constraint_grads[idx] = tf.zeros_like(param)
flat_grad = tensor_utils.flatten_tensor_variables(constraint_grads)
def f_hx_plain(*args):
"""Computes product of Hessian(f) and vector v.
Args:
args (tuple[numpy.ndarray]): Contains inputs of function f
, and vector v.
Returns:
tf.Tensor: Symbolic result.
"""
with tf.name_scope('f_hx_plain', values=[inputs,
self._target]):
inputs_ = args[:len(inputs)]
xs = args[len(inputs):]
flat_xs = np.concatenate(
[np.reshape(x, (-1, )) for x in xs])
param_val = self._target.get_param_values()
eps = np.cast['float32'](
self.base_eps / (np.linalg.norm(param_val) + 1e-8))
self._target.set_param_values(param_val + eps * flat_xs)
flat_grad_dvplus = self._hvp_fun['f_grad'](*inputs_)
self._target.set_param_values(param_val)
if self.symmetric:
self._target.set_param_values(param_val -
eps * flat_xs)
flat_grad_dvminus = self._hvp_fun['f_grad'](*inputs_)
hx = (flat_grad_dvplus - flat_grad_dvminus) / (2 * eps)
self._target.set_param_values(param_val)
else:
flat_grad = self._hvp_fun['f_grad'](*inputs_)
hx = (flat_grad_dvplus - flat_grad) / eps
return hx
self._hvp_fun = LazyDict(
f_grad=lambda: tensor_utils.compile_function(
inputs=inputs,
outputs=flat_grad,
log_name='f_grad',
),
f_hx_plain=lambda: f_hx_plain,
)