# Copyright 2017 reinforce.io. 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.

# ==============================================================================

from __future__ import absolute_import

from __future__ import print_function

from __future__ import division

import tensorflow as tf

from tensorforce.core.baselines import Baseline, AggregatedBaseline

from tensorforce.core.optimizers import Optimizer

from tensorforce.models import DistributionModel

class PGModel(DistributionModel):

"""

Base class for policy gradient models. It optionally defines a baseline

and handles its optimization. It implements the `tf_loss_per_instance` function, but requires

subclasses to implement `tf_pg_loss_per_instance`.

"""

def __init__(

self,

states,

actions,

scope,

device,

saver,

summarizer,

distributed,

batching_capacity,

variable_noise,

states_preprocessing,

actions_exploration,

reward_preprocessing,

update_mode,

memory,

optimizer,

discount,

network,

distributions,

entropy_regularization,

baseline_mode,

baseline,

baseline_optimizer,

gae_lambda

):

# Baseline mode

assert baseline_mode is None or baseline_mode in ('states', 'network')

self.baseline_mode = baseline_mode

self.baseline_spec = baseline

self.baseline_optimizer_spec = baseline_optimizer

# Generalized advantage function

assert gae_lambda is None or (0.0 <= gae_lambda <= 1.0 and self.baseline_mode is not None)

self.gae_lambda = gae_lambda

self.baseline = None

self.baseline_optimizer = None

self.fn_reward_estimation = None

super(PGModel, self).__init__(

states=states,

actions=actions,

scope=scope,

device=device,

saver=saver,

summarizer=summarizer,

distributed=distributed,

batching_capacity=batching_capacity,

variable_noise=variable_noise,

states_preprocessing=states_preprocessing,

actions_exploration=actions_exploration,

reward_preprocessing=reward_preprocessing,

update_mode=update_mode,

memory=memory,

optimizer=optimizer,

discount=discount,

network=network,

distributions=distributions,

entropy_regularization=entropy_regularization,

requires_deterministic=False

)

def as_local_model(self):

super(PGModel, self).as_local_model()

if self.baseline_optimizer_spec is not None:

self.baseline_optimizer_spec = dict(

type='global_optimizer',

optimizer=self.baseline_optimizer_spec

)

def initialize(self, custom_getter):

super(PGModel, self).initialize(custom_getter)

# Baseline

if self.baseline_spec is None:

assert self.baseline_mode is None

elif all(name in self.states_spec for name in self.baseline_spec):

# Implies AggregatedBaseline.

assert self.baseline_mode == 'states'

self.baseline = AggregatedBaseline(baselines=self.baseline_spec)

else:

assert self.baseline_mode is not None

self.baseline = Baseline.from_spec(

spec=self.baseline_spec,

kwargs=dict(

summary_labels=self.summary_labels

)

)

# Baseline optimizer

if self.baseline_optimizer_spec is not None:

assert self.baseline_mode is not None

self.baseline_optimizer = Optimizer.from_spec(spec=self.baseline_optimizer_spec)

# TODO: Baseline internal states !!! (see target_network q_model)

# Reward estimation

self.fn_reward_estimation = tf.make_template(

name_='reward-estimation',

func_=self.tf_reward_estimation,

custom_getter_=custom_getter

)

# Baseline loss

self.fn_baseline_loss = tf.make_template(

name_='baseline-loss',

func_=self.tf_baseline_loss,

custom_getter_=custom_getter

)

def tf_reward_estimation(self, states, internals, terminal, reward, update):

if self.baseline_mode is None:

return self.fn_discounted_cumulative_reward(terminal=terminal, reward=reward, discount=self.discount)

else:

if self.baseline_mode == 'states':

state_value = self.baseline.predict(

states=states,

internals=internals,

update=update

)

elif self.baseline_mode == 'network':

embedding = self.network.apply(

x=states,

internals=internals,

update=update

)

state_value = self.baseline.predict(

states=tf.stop_gradient(input=embedding),

internals=internals,

update=update

)

if self.gae_lambda is None:

reward = self.fn_discounted_cumulative_reward(

terminal=terminal,

reward=reward,

discount=self.discount

)

advantage = reward - state_value

else:

next_state_value = tf.concat(values=(state_value[1:], (0.0,)), axis=0)

zeros = tf.zeros_like(tensor=next_state_value)

next_state_value = tf.where(condition=terminal, x=zeros, y=next_state_value)

td_residual = reward + self.discount * next_state_value - state_value

gae_discount = self.discount * self.gae_lambda

advantage = self.fn_discounted_cumulative_reward(

terminal=terminal,

reward=td_residual,

discount=gae_discount

)

# Normalize advantage.

# mean, variance = tf.nn.moments(advantage, axes=[0], keep_dims=True)

# advantage = (advantage - mean) / tf.sqrt(x=tf.maximum(x=variance, y=util.epsilon))

return advantage

def tf_regularization_losses(self, states, internals, update):

losses = super(PGModel, self).tf_regularization_losses(

states=states,

internals=internals,

update=update

)

if self.baseline_mode is not None and self.baseline_optimizer is None:

baseline_regularization_loss = self.baseline.regularization_loss()

if baseline_regularization_loss is not None:

losses['baseline'] = baseline_regularization_loss

return losses

def tf_baseline_loss(self, states, internals, reward, update, reference=None):

"""

Creates the TensorFlow operations for calculating the baseline loss of a batch.

Args:

states: Dict of state tensors.

internals: List of prior internal state tensors.

reward: Reward tensor.

update: Boolean tensor indicating whether this call happens during an update.

reference: Optional reference tensor(s), in case of a comparative loss.

Returns:

Loss tensor.

"""

if self.baseline_mode == 'states':

loss = self.baseline.loss(

states=states,

internals=internals,

reward=reward,

update=update,

reference=reference

)

elif self.baseline_mode == 'network':

loss = self.baseline.loss(

states=self.network.apply(x=states, internals=internals, update=update),

internals=internals,

reward=reward,

update=update,

reference=reference

)

regularization_loss = self.baseline.regularization_loss()

if regularization_loss is not None:

loss += regularization_loss

return loss

def baseline_optimizer_arguments(self, states, internals, reward):

"""

Returns the baseline optimizer arguments including the time, the list of variables to

optimize, and various functions which the optimizer might require to perform an update

step.

Args:

states: Dict of state tensors.

internals: List of prior internal state tensors.

reward: Reward tensor.

Returns:

Baseline optimizer arguments as dict.

"""

arguments = dict(

time=self.global_timestep,

variables=self.baseline.get_variables(),

arguments=dict(

states=states,

internals=internals,

reward=reward,

update=tf.constant(value=True),

),

fn_reference=self.baseline.reference,

fn_loss=self.fn_baseline_loss,

# source_variables=self.network.get_variables()

)

if self.global_model is not None:

arguments['global_variables'] = self.global_model.baseline.get_variables()

return arguments

def tf_optimization(self, states, internals, actions, terminal, reward, next_states=None, next_internals=None):

assert next_states is None and next_internals is None # temporary

estimated_reward = self.fn_reward_estimation(

states=states,

internals=internals,

terminal=terminal,

reward=reward,

update=tf.constant(value=True)

)

if self.baseline_optimizer is not None:

estimated_reward = tf.stop_gradient(input=estimated_reward)

optimization = super(PGModel, self).tf_optimization(

states=states,

internals=internals,

actions=actions,

terminal=terminal,

reward=estimated_reward,

next_states=next_states,

next_internals=next_internals

)

if self.baseline_optimizer is not None:

cumulative_reward = self.fn_discounted_cumulative_reward(terminal=terminal, reward=reward, discount=self.discount)

arguments = self.baseline_optimizer_arguments(

states=states,

internals=internals,

reward=cumulative_reward,

)

baseline_optimization = self.baseline_optimizer.minimize(**arguments)

optimization = tf.group(optimization, baseline_optimization)

return optimization

def get_variables(self, include_submodules=False, include_nontrainable=False):

model_variables = super(PGModel, self).get_variables(

include_submodules=include_submodules,

include_nontrainable=include_nontrainable

)

if self.baseline_mode is not None and (include_submodules or self.baseline_optimizer is None):

baseline_variables = self.baseline.get_variables(include_nontrainable=include_nontrainable)

model_variables += baseline_variables

if include_nontrainable and self.baseline_optimizer is not None:

baseline_optimizer_variables = self.baseline_optimizer.get_variables()

# For some reason, some optimizer variables are only registered in the model.

for variable in baseline_optimizer_variables:

if variable in model_variables:

model_variables.remove(variable)

model_variables += baseline_optimizer_variables

return model_variables

def get_summaries(self):

if self.baseline_mode is None:

return super(PGModel, self).get_summaries()

else:

return super(PGModel, self).get_summaries() + self.baseline.get_summaries()