research/pcl_rl/model.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.
# ==============================================================================
"""Model is responsible for setting up Tensorflow graph.
Creates policy and value networks. Also sets up all optimization
ops, including gradient ops, trust region ops, and value optimizers.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import tensorflow as tf
class Model(object):
def __init__(self, env_spec, global_step,
target_network_lag=0.95,
sample_from='online',
get_policy=None,
get_baseline=None,
get_objective=None,
get_trust_region_p_opt=None,
get_value_opt=None):
self.env_spec = env_spec
self.global_step = global_step
self.inc_global_step = self.global_step.assign_add(1)
self.target_network_lag = target_network_lag
self.sample_from = sample_from
self.policy = get_policy()
self.baseline = get_baseline()
self.objective = get_objective()
self.baseline.eps_lambda = self.objective.eps_lambda # TODO: do this better
self.trust_region_policy_opt = get_trust_region_p_opt()
self.value_opt = get_value_opt()
def setup_placeholders(self):
"""Create the Tensorflow placeholders."""
# summary placeholder
self.avg_episode_reward = tf.placeholder(
tf.float32, [], 'avg_episode_reward')
self.greedy_episode_reward = tf.placeholder(
tf.float32, [], 'greedy_episode_reward')
# sampling placeholders
self.internal_state = tf.placeholder(tf.float32,
[None, self.policy.rnn_state_dim],
'internal_state')
self.single_observation = []
for i, (obs_dim, obs_type) in enumerate(self.env_spec.obs_dims_and_types):
if self.env_spec.is_discrete(obs_type):
self.single_observation.append(
tf.placeholder(tf.int32, [None], 'obs%d' % i))
elif self.env_spec.is_box(obs_type):
self.single_observation.append(
tf.placeholder(tf.float32, [None, obs_dim], 'obs%d' % i))
else:
assert False
self.single_action = []
for i, (action_dim, action_type) in \
enumerate(self.env_spec.act_dims_and_types):
if self.env_spec.is_discrete(action_type):
self.single_action.append(
tf.placeholder(tf.int32, [None], 'act%d' % i))
elif self.env_spec.is_box(action_type):
self.single_action.append(
tf.placeholder(tf.float32, [None, action_dim], 'act%d' % i))
else:
assert False
# training placeholders
self.observations = []
for i, (obs_dim, obs_type) in enumerate(self.env_spec.obs_dims_and_types):
if self.env_spec.is_discrete(obs_type):
self.observations.append(
tf.placeholder(tf.int32, [None, None], 'all_obs%d' % i))
else:
self.observations.append(
tf.placeholder(tf.float32, [None, None, obs_dim], 'all_obs%d' % i))
self.actions = []
self.other_logits = []
for i, (action_dim, action_type) in \
enumerate(self.env_spec.act_dims_and_types):
if self.env_spec.is_discrete(action_type):
self.actions.append(
tf.placeholder(tf.int32, [None, None], 'all_act%d' % i))
if self.env_spec.is_box(action_type):
self.actions.append(
tf.placeholder(tf.float32, [None, None, action_dim],
'all_act%d' % i))
self.other_logits.append(
tf.placeholder(tf.float32, [None, None, None],
'other_logits%d' % i))
self.rewards = tf.placeholder(tf.float32, [None, None], 'rewards')
self.terminated = tf.placeholder(tf.float32, [None], 'terminated')
self.pads = tf.placeholder(tf.float32, [None, None], 'pads')
self.prev_log_probs = tf.placeholder(tf.float32, [None, None],
'prev_log_probs')
def setup(self, train=True):
"""Setup Tensorflow Graph."""
self.setup_placeholders()
tf.summary.scalar('avg_episode_reward', self.avg_episode_reward)
tf.summary.scalar('greedy_episode_reward', self.greedy_episode_reward)
with tf.variable_scope('model', reuse=None):
# policy network
with tf.variable_scope('policy_net'):
(self.policy_internal_states, self.logits, self.log_probs,
self.entropies, self.self_kls) = \
self.policy.multi_step(self.observations,
self.internal_state,
self.actions)
self.out_log_probs = sum(self.log_probs)
self.kl = self.policy.calculate_kl(self.other_logits, self.logits)
self.avg_kl = (tf.reduce_sum(sum(self.kl)[:-1] * (1 - self.pads)) /
tf.reduce_sum(1 - self.pads))
# value network
with tf.variable_scope('value_net'):
(self.values,
self.regression_input,
self.regression_weight) = self.baseline.get_values(
self.observations, self.actions,
self.policy_internal_states, self.logits)
# target policy network
with tf.variable_scope('target_policy_net'):
(self.target_policy_internal_states,
self.target_logits, self.target_log_probs,
_, _) = \
self.policy.multi_step(self.observations,
self.internal_state,
self.actions)
# target value network
with tf.variable_scope('target_value_net'):
(self.target_values, _, _) = self.baseline.get_values(
self.observations, self.actions,
self.target_policy_internal_states, self.target_logits)
# construct copy op online --> target
all_vars = tf.trainable_variables()
online_vars = [p for p in all_vars if
'/policy_net' in p.name or '/value_net' in p.name]
target_vars = [p for p in all_vars if
'target_policy_net' in p.name or 'target_value_net' in p.name]
online_vars.sort(key=lambda p: p.name)
target_vars.sort(key=lambda p: p.name)
aa = self.target_network_lag
self.copy_op = tf.group(*[
target_p.assign(aa * target_p + (1 - aa) * online_p)
for online_p, target_p in zip(online_vars, target_vars)])
if train:
# evaluate objective
(self.loss, self.raw_loss, self.regression_target,
self.gradient_ops, self.summary) = self.objective.get(
self.rewards, self.pads,
self.values[:-1, :],
self.values[-1, :] * (1 - self.terminated),
self.log_probs, self.prev_log_probs, self.target_log_probs,
self.entropies, self.logits, self.target_values[:-1, :],
self.target_values[-1, :] * (1 - self.terminated))
self.regression_target = tf.reshape(self.regression_target, [-1])
self.policy_vars = [
v for v in tf.trainable_variables()
if '/policy_net' in v.name]
self.value_vars = [
v for v in tf.trainable_variables()
if '/value_net' in v.name]
# trust region optimizer
if self.trust_region_policy_opt is not None:
with tf.variable_scope('trust_region_policy', reuse=None):
avg_self_kl = (
tf.reduce_sum(sum(self.self_kls) * (1 - self.pads)) /
tf.reduce_sum(1 - self.pads))
self.trust_region_policy_opt.setup(
self.policy_vars, self.raw_loss, avg_self_kl,
self.avg_kl)
# value optimizer
if self.value_opt is not None:
with tf.variable_scope('trust_region_value', reuse=None):
self.value_opt.setup(
self.value_vars,
tf.reshape(self.values[:-1, :], [-1]),
self.regression_target,
tf.reshape(self.pads, [-1]),
self.regression_input, self.regression_weight)
# we re-use variables for the sampling operations
with tf.variable_scope('model', reuse=True):
scope = ('target_policy_net' if self.sample_from == 'target'
else 'policy_net')
with tf.variable_scope(scope):
self.next_internal_state, self.sampled_actions = \
self.policy.sample_step(self.single_observation,
self.internal_state,
self.single_action)
self.greedy_next_internal_state, self.greedy_sampled_actions = \
self.policy.sample_step(self.single_observation,
self.internal_state,
self.single_action,
greedy=True)
def sample_step(self, sess,
single_observation, internal_state, single_action,
greedy=False):
"""Sample batch of steps from policy."""
if greedy:
outputs = [self.greedy_next_internal_state, self.greedy_sampled_actions]
else:
outputs = [self.next_internal_state, self.sampled_actions]
feed_dict = {self.internal_state: internal_state}
for action_place, action in zip(self.single_action, single_action):
feed_dict[action_place] = action
for obs_place, obs in zip(self.single_observation, single_observation):
feed_dict[obs_place] = obs
return sess.run(outputs, feed_dict=feed_dict)
def train_step(self, sess,
observations, internal_state, actions,
rewards, terminated, pads,
avg_episode_reward=0, greedy_episode_reward=0):
"""Train network using standard gradient descent."""
outputs = [self.raw_loss, self.gradient_ops, self.summary]
feed_dict = {self.internal_state: internal_state,
self.rewards: rewards,
self.terminated: terminated,
self.pads: pads,
self.avg_episode_reward: avg_episode_reward,
self.greedy_episode_reward: greedy_episode_reward}
time_len = None
for action_place, action in zip(self.actions, actions):
if time_len is None:
time_len = len(action)
assert time_len == len(action)
feed_dict[action_place] = action
for obs_place, obs in zip(self.observations, observations):
assert time_len == len(obs)
feed_dict[obs_place] = obs
assert len(rewards) == time_len - 1
return sess.run(outputs, feed_dict=feed_dict)
def trust_region_step(self, sess,
observations, internal_state, actions,
rewards, terminated, pads,
avg_episode_reward=0,
greedy_episode_reward=0):
"""Train policy using trust region step."""
feed_dict = {self.internal_state: internal_state,
self.rewards: rewards,
self.terminated: terminated,
self.pads: pads,
self.avg_episode_reward: avg_episode_reward,
self.greedy_episode_reward: greedy_episode_reward}
for action_place, action in zip(self.actions, actions):
feed_dict[action_place] = action
for obs_place, obs in zip(self.observations, observations):
feed_dict[obs_place] = obs
(prev_log_probs, prev_logits) = sess.run(
[self.out_log_probs, self.logits], feed_dict=feed_dict)
feed_dict[self.prev_log_probs] = prev_log_probs
for other_logit, prev_logit in zip(self.other_logits, prev_logits):
feed_dict[other_logit] = prev_logit
# fit policy
self.trust_region_policy_opt.optimize(sess, feed_dict)
ret = sess.run([self.raw_loss, self.summary], feed_dict=feed_dict)
ret = [ret[0], None, ret[1]]
return ret
def fit_values(self, sess,
observations, internal_state, actions,
rewards, terminated, pads):
"""Train value network using value-specific optimizer."""
feed_dict = {self.internal_state: internal_state,
self.rewards: rewards,
self.terminated: terminated,
self.pads: pads}
for action_place, action in zip(self.actions, actions):
feed_dict[action_place] = action
for obs_place, obs in zip(self.observations, observations):
feed_dict[obs_place] = obs
# fit values
if self.value_opt is None:
raise ValueError('Specific value optimizer does not exist')
self.value_opt.optimize(sess, feed_dict)