research/pcl_rl/optimizers.py
# Copyright 2017 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.
# ==============================================================================
"""Optimizers mostly for value estimate.
Gradient Descent optimizer
LBFGS optimizer
Best Fit optimizer
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from six.moves import xrange
import tensorflow as tf
import numpy as np
import scipy.optimize
def var_size(v):
return int(np.prod([int(d) for d in v.shape]))
def gradients(loss, var_list):
grads = tf.gradients(loss, var_list)
return [g if g is not None else tf.zeros(v.shape)
for g, v in zip(grads, var_list)]
def flatgrad(loss, var_list):
grads = gradients(loss, var_list)
return tf.concat([tf.reshape(grad, [-1])
for (v, grad) in zip(var_list, grads)
if grad is not None], 0)
def get_flat(var_list):
return tf.concat([tf.reshape(v, [-1]) for v in var_list], 0)
def set_from_flat(var_list, flat_theta):
assigns = []
shapes = [v.shape for v in var_list]
sizes = [var_size(v) for v in var_list]
start = 0
assigns = []
for (shape, size, v) in zip(shapes, sizes, var_list):
assigns.append(v.assign(
tf.reshape(flat_theta[start:start + size], shape)))
start += size
assert start == sum(sizes)
return tf.group(*assigns)
class LbfgsOptimization(object):
def __init__(self, max_iter=25, mix_frac=1.0):
self.max_iter = max_iter
self.mix_frac = mix_frac
def setup_placeholders(self):
self.flat_theta = tf.placeholder(tf.float32, [None], 'flat_theta')
self.intended_values = tf.placeholder(tf.float32, [None], 'intended_values')
def setup(self, var_list, values, targets, pads,
inputs, regression_weight):
self.setup_placeholders()
self.values = values
self.targets = targets
self.raw_loss = (tf.reduce_sum((1 - pads) * tf.square(values - self.intended_values))
/ tf.reduce_sum(1 - pads))
self.loss_flat_gradient = flatgrad(self.raw_loss, var_list)
self.flat_vars = get_flat(var_list)
self.set_vars = set_from_flat(var_list, self.flat_theta)
def optimize(self, sess, feed_dict):
old_theta = sess.run(self.flat_vars)
old_values, targets = sess.run([self.values, self.targets], feed_dict=feed_dict)
intended_values = targets * self.mix_frac + old_values * (1 - self.mix_frac)
feed_dict = dict(feed_dict)
feed_dict[self.intended_values] = intended_values
def calc_loss_and_grad(theta):
sess.run(self.set_vars, feed_dict={self.flat_theta: theta})
loss, grad = sess.run([self.raw_loss, self.loss_flat_gradient],
feed_dict=feed_dict)
grad = grad.astype('float64')
return loss, grad
theta, _, _ = scipy.optimize.fmin_l_bfgs_b(
calc_loss_and_grad, old_theta, maxiter=self.max_iter)
sess.run(self.set_vars, feed_dict={self.flat_theta: theta})
class GradOptimization(object):
def __init__(self, learning_rate=0.001, max_iter=25, mix_frac=1.0):
self.learning_rate = learning_rate
self.max_iter = max_iter
self.mix_frac = mix_frac
def get_optimizer(self):
return tf.train.AdamOptimizer(learning_rate=self.learning_rate,
epsilon=2e-4)
def setup_placeholders(self):
self.flat_theta = tf.placeholder(tf.float32, [None], 'flat_theta')
self.intended_values = tf.placeholder(tf.float32, [None], 'intended_values')
def setup(self, var_list, values, targets, pads,
inputs, regression_weight):
self.setup_placeholders()
self.values = values
self.targets = targets
self.raw_loss = (tf.reduce_sum((1 - pads) * tf.square(values - self.intended_values))
/ tf.reduce_sum(1 - pads))
opt = self.get_optimizer()
params = var_list
grads = tf.gradients(self.raw_loss, params)
self.gradient_ops = opt.apply_gradients(zip(grads, params))
def optimize(self, sess, feed_dict):
old_values, targets = sess.run([self.values, self.targets], feed_dict=feed_dict)
intended_values = targets * self.mix_frac + old_values * (1 - self.mix_frac)
feed_dict = dict(feed_dict)
feed_dict[self.intended_values] = intended_values
for _ in xrange(self.max_iter):
sess.run(self.gradient_ops, feed_dict=feed_dict)
class BestFitOptimization(object):
def __init__(self, mix_frac=1.0):
self.mix_frac = mix_frac
def setup_placeholders(self):
self.new_regression_weight = tf.placeholder(
tf.float32, self.regression_weight.shape)
def setup(self, var_list, values, targets, pads,
inputs, regression_weight):
self.values = values
self.targets = targets
self.inputs = inputs
self.regression_weight = regression_weight
self.setup_placeholders()
self.update_regression_weight = tf.assign(
self.regression_weight, self.new_regression_weight)
def optimize(self, sess, feed_dict):
reg_input, reg_weight, old_values, targets = sess.run(
[self.inputs, self.regression_weight, self.values, self.targets],
feed_dict=feed_dict)
intended_values = targets * self.mix_frac + old_values * (1 - self.mix_frac)
# taken from rllab
reg_coeff = 1e-5
for _ in range(5):
best_fit_weight = np.linalg.lstsq(
reg_input.T.dot(reg_input) +
reg_coeff * np.identity(reg_input.shape[1]),
reg_input.T.dot(intended_values))[0]
if not np.any(np.isnan(best_fit_weight)):
break
reg_coeff *= 10
if len(best_fit_weight.shape) == 1:
best_fit_weight = np.expand_dims(best_fit_weight, -1)
sess.run(self.update_regression_weight,
feed_dict={self.new_regression_weight: best_fit_weight})