def update_weights(optimizer: tf.train.Optimizer, network: Network, batch,
weight_decay: float):
loss = 0
for image, actions, targets in batch:
# Initial step, from the real observation.
value, reward, policy_logits, hidden_state = network.initial_inference(
image)
predictions = [(1.0, value, reward, policy_logits)]
# Recurrent steps, from action and previous hidden state.
for action in actions:
value, reward, policy_logits, hidden_state = network.recurrent_inference(
hidden_state, action)
predictions.append((1.0 / len(actions), value, reward, policy_logits))
hidden_state = tf.scale_gradient(hidden_state, 0.5)
for prediction, target in zip(predictions, targets):
gradient_scale, value, reward, policy_logits = prediction
target_value, target_reward, target_policy = target
l = (
scalar_loss(value, target_value) +
scalar_loss(reward, target_reward) +
tf.nn.softmax_cross_entropy_with_logits(
logits=policy_logits, labels=target_policy))
loss += tf.scale_gradient(l, gradient_scale)
for weights in network.get_weights():
loss += weight_decay * tf.nn.l2_loss(weights)
optimizer.minimize(loss)
def update_weights(optimizer: tf.train.Optimizer, network: Network, batch,
weight_decay: float):
loss = 0
for image, (target_value, target_policy) in batch:
value, policy_logits = network.inference(image)
loss += (tf.losses.mean_squared_error(value, target_value) +
tf.nn.softmax_cross_entropy_with_logits(logits=policy_logits,
labels=target_policy))
for weights in network.get_weights():
loss += weight_decay * tf.nn.l2_loss(weights)
optimizer.minimize(loss)
def train_step(model: tf.keras.Model, optimizer: tf.train.Optimizer,
loss: loss, x: tf.Tensor, y: tf.Tensor):
"""Training operation. That is, we minimize the loss function here.
Arguments:
model {tf.keras.Model} -- Instance of tf.keras.Model
optimizer {tf.train.Optimizer} -- Optimizer to be used.
loss {loss} -- Loss function.
x {tf.Tensor} -- Input features.
y {tf.Tensor} -- Output labels.
"""
optimizer.minimize(loss=lambda: loss(model, x, y),
global_step=tf.train.get_or_create_global_step())
def dns_grad_op(loss, optimizer: tf.train.Optimizer, variables=None, global_step=None):
""" Create an operation the updates the weights by gradient descent.
In DNS, the weights are updated according to their derivative with respect to the masked
values, but the update is applied to the non-masked values, so that zeroed-out weights may
still change and in particular be spliced back in if necessary.
Parameters
----------
loss: A `tf.Tensor` representing the loss.
optimizer: The optimizer to use.
variables: The variables for which to create the gradient operation.
global_step: An optional global step to increment.
Returns
-------
train_op: An tensorflow op that when run updates the variables according to the gradient.
"""
if variables is None:
variables = tf.trainable_variables()
replaced = {}
wrt_variables = []
num_replaced = 0
for v in variables:
# look for variables having shadow values.
mvs = tf.get_collection(MASKED_WEIGHT_COLLECTION, v.op.name)
if len(mvs) == 0:
wrt_variables.append(v)
elif len(mvs) == 1:
num_replaced += 1
wrt_variables.append(mvs[0])
replaced[mvs[0]] = v
else:
raise ValueError('More than one masked weight for a given variable.')
tf.logging.info('Replaced {0} variables for Dynamic Network Surgery'.format(num_replaced))
grads_and_vars = optimizer.compute_gradients(loss, wrt_variables)
grads_and_vars = [(g, replaced.get(v, v)) for g, v in grads_and_vars]
train_op = optimizer.apply_gradients(grads_and_vars, global_step, 'dns_grad_op')
return train_op
def __init__(
self,
environment: ControlBenchmark,
experience_buffer: BaseExperienceBuffer,
tensorflow_session: tf.Session,
gamma: float,
layer_sizes: List[int],
layer_activations: List[str],
shared_layers: int,
tau: float,
optimizer: tf.train.Optimizer,
batch_size: int,
) -> None:
super().__init__(environment=environment,
experience_buffer=experience_buffer)
self.shared_layers = shared_layers
self.tensorflow_session = tensorflow_session
self.batch_size = batch_size
self.Q = NAFNetwork(layer_sizes=layer_sizes,
layer_activations=layer_activations,
shared_layers=shared_layers,
state_shape=environment.state_shape,
action_shape=environment.action_shape)
self.Q_lowpass = NAFNetwork(layer_sizes=layer_sizes,
layer_activations=layer_activations,
shared_layers=shared_layers,
state_shape=environment.state_shape,
action_shape=environment.action_shape)
self.Q_lowpass.model.set_weights(self.Q.model.get_weights())
self.observation_input = tf.keras.Input(
shape=self.environment.state_shape, name='state')
self.next_observation_input = tf.keras.Input(
shape=self.environment.state_shape, name='next_state')
self.action_input = tf.keras.Input(shape=self.environment.action_shape,
name='action_placeholder')
self.reward_input = tf.keras.Input(shape=(), name='reward')
self.terminal_input = tf.keras.Input(shape=(), name='terminal')
self.p_continue = gamma * (1 - self.terminal_input)
self.frozen_parameter_update_op = periodic_target_update(
target_variables=self.Q_lowpass.model.variables,
source_variables=self.Q.model.variables,
update_period=1,
tau=tau)
self.q_values_policy, self.mu_policy, _ = self.Q(
state_action=[self.observation_input, self.action_input])
_, _, self.vt_lowpass = self.Q_lowpass(
state_action=[self.next_observation_input, self.action_input])
# action is not actually used here to calculate the value
self.target = self.reward_input + self.p_continue * self.vt_lowpass
rl_loss = tf.reduce_mean(0.5 * (self.q_values_policy - self.target)**2)
self.train_op = optimizer.minimize(rl_loss)
self._initialize_tf_variables()
def train(self,
x_data: tf.Tensor,
y_data: tf.Tensor,
example_number: int,
epochs: int,
batch_size: int,
activation_function: ActivationFunction,
cost_function: CostFunction,
optimizer_function: tf.train.Optimizer,
) -> None:
assert x_data.shape[0] == y_data.shape[0] == example_number
x_dataset = tf.data.Dataset.from_tensor_slices(x_data)
y_dataset = tf.data.Dataset.from_tensor_slices(y_data)
assert x_dataset.output_shapes == self.x_size
assert y_dataset.output_shapes == self.y_size
assert self.is_initialized
self.activation_function = activation_function
batched_x, batched_y = (i.batch(batch_size)
for i in (x_dataset, y_dataset))
x_batch, y_batch = (
tf.placeholder('float', shape=(batch_size, self.x_size)),
tf.placeholder('float', shape=(batch_size, self.y_size)))
pred_y_batch = tf.map_fn(self.model, x_batch)
cost = tf.reduce_mean(cost_function(pred_y_batch, y_batch))
optimizer = optimizer_function.minimize(cost)
print('Training...')
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(epochs):
epoch_loss = 0
x_it, y_it = (b.make_one_shot_iterator()
for b in (batched_x, batched_y))
curr_x_init = x_it.get_next()
curr_y_init = y_it.get_next()
while True:
try:
curr_x = sess.run(curr_x_init)
curr_y = sess.run(curr_y_init)
_, c = sess.run([optimizer, cost], feed_dict={
x_batch: curr_x, y_batch: curr_y
})
epoch_loss += c
except tf.errors.OutOfRangeError:
break
print('Epoch', epoch + 1, 'out of', epochs, 'completed.',
'Current loss:', epoch_loss)
self.test_accuracy(sess, x_data, y_data, example_number)
def __init__(
self,
obs_spec: specs.Array,
action_spec: specs.DiscreteArray,
network: snt.AbstractModule,
optimizer: tf.train.Optimizer,
sequence_length: int,
td_lambda: float,
agent_discount: float,
seed: int,
):
"""A simple actor-critic agent."""
del action_spec # unused
tf.set_random_seed(seed)
self._sequence_length = sequence_length
self._count = 0
# Create the policy ops..
obs = tf.placeholder(shape=obs_spec.shape, dtype=obs_spec.dtype)
online_logits, _ = network(tf.expand_dims(obs, 0))
action = tf.squeeze(
tf.multinomial(online_logits, 1, output_dtype=tf.int32))
# Create placeholders and numpy arrays for learning from trajectories.
shapes = [obs_spec.shape, (), (), ()]
dtypes = [obs_spec.dtype, np.int32, np.float32, np.float32]
placeholders = [
tf.placeholder(shape=(self._sequence_length, 1) + shape,
dtype=dtype)
for shape, dtype in zip(shapes, dtypes)
]
observations, actions, rewards, discounts = placeholders
self.arrays = [
np.zeros(shape=(self._sequence_length, 1) + shape, dtype=dtype)
for shape, dtype in zip(shapes, dtypes)
]
# Build actor and critic losses.
logits, values = snt.BatchApply(network)(observations)
_, bootstrap_value = network(tf.expand_dims(obs, 0))
critic_loss, (advantages, _) = td_lambda_loss(
state_values=values,
rewards=rewards,
pcontinues=agent_discount * discounts,
bootstrap_value=bootstrap_value,
lambda_=td_lambda)
actor_loss = discrete_policy_gradient_loss(logits, actions, advantages)
train_op = optimizer.minimize(actor_loss + critic_loss)
# Create TF session and callables.
session = tf.Session()
self._policy_fn = session.make_callable(action, [obs])
self._update_fn = session.make_callable(train_op, placeholders + [obs])
session.run(tf.global_variables_initializer())
def ClippingOptimizer(opt: tf.train.Optimizer, low, high):
original = opt.apply_gradients
def apply_gradients(grads_and_vars, *a, **kw):
app = original(grads_and_vars, *a, **kw)
asg = [
v.assign_add(tf.maximum(high - v, 0) + tf.minimum(low - v, 0))
for g, v in grads_and_vars
]
return tf.group(app, *asg) # note that clipping is asynchronous here
opt.apply_gradients = apply_gradients
return opt
def train_model(optimizer: tf.train.Optimizer, loss: tf.Tensor):
"""Minimize the loss with respect to the model variables.
Args:
optimizer (tf.train.Optimizer):
loss (tf.Tensor): Loss value as defined by a loss function..
Returns:
An Operation that updates the variables in `var_list`
& also increments `global_step`.
"""
return optimizer.minimize(loss=loss,
global_step=tf.train.get_or_create_global_step())
def __distribute_training(self, iterator: tf.data.Iterator,
optimizer: tf.train.Optimizer) -> DistributedOps:
gpus_to_use = self.__get_gpu_to_use()
gradients, loss_operations = [], []
for gpu_id in gpus_to_use:
multi_gpu_operations = self.__place_operations(
target_gpu_id=gpu_id, iterator=iterator, optimizer=optimizer)
gradients.append(multi_gpu_operations.gradient)
loss_operations.append(multi_gpu_operations.loss_operation)
gradients = average_gradients(gradients)
loss_operation = average_loss(loss_operations)
training_step = optimizer.apply_gradients(gradients)
return loss_operation, training_step
def adversarial_train_op_func(
generator_loss: tf.Tensor,
discriminator_loss: tf.Tensor,
generator_weights: List[tf.Variable],
discriminator_weights: List[tf.Variable],
n_gen_steps: int = 1,
n_disc_steps: int = 5,
optimizer: tf.train.Optimizer = tf.train.RMSPropOptimizer(0.0005)
) -> tf.Operation:
"""
Build the adversarial train operation (n_disc_steps discriminator optimization steps
followed by n_gen_steps generator optimization steps).
Arguments:
generator_loss -- generator loss.
discriminator_loss -- discriminator loss.
generator_weights -- list of generator trainable weights.
discriminator_weights -- list of discriminator trainable weights.
n_gen_steps -- number of generator update steps per single train operation,
optional (default = 1).
n_disc_steps -- number of discriminator update steps per single train
operation, optional (default = 10).
optimizer -- optimizer to use, optional (default = tf.train.RMSPropOptimizer(0.001))
"""
disc_train_op = _op_repeat_n(
lambda: optimizer.minimize(discriminator_loss, var_list=discriminator_weights),
n_disc_steps
)
with tf.control_dependencies([disc_train_op]):
gen_train_op = _op_repeat_n(
lambda: optimizer.minimize(generator_loss, var_list=generator_weights),
n_gen_steps
)
return gen_train_op
def get_gradient_op(tensors: MDPTensors,
objective_initial_scales: SRLObjectives,
optimizer: tf.train.Optimizer,
gradient_clip: Optional[float], **kwargs):
objectives: SRLObjectives = SRLObjectives(
value_function=ValueFunction(tensors,
objective_initial_scales.value_function,
**kwargs),
reward_prediction=RewardPrediction(
tensors, objective_initial_scales.reward_prediction, **kwargs),
auto_encoding=AutoEncodingPrediction(
tensors, objective_initial_scales.auto_encoding, **kwargs),
forward_dynamics=ForwardDynamicsPrediction(
tensors, objective_initial_scales.forward_dynamics, **kwargs),
inverse_dynamics=InverseDynamicsPrediction(
tensors, objective_initial_scales.inverse_dynamics, **kwargs),
slowness=SlownessLoss(tensors, objective_initial_scales.slowness,
**kwargs),
diversity=DiversityLoss(tensors, objective_initial_scales.diversity,
**kwargs),
)
active_objectives = [
o for o in objectives
if o is not None and backend.get_value(o.scale) > 0
]
total_loss = backend.mean(
backend.stack([o.loss for o in active_objectives]))
if gradient_clip is not None:
gradients = optimizer.compute_gradients(total_loss)
for i, (grad, var) in enumerate(gradients):
if grad is not None:
gradients[i] = (tf.clip_by_norm(grad, gradient_clip), var)
return optimizer.apply_gradients(gradients)
else:
return optimizer.minimize(total_loss)
def feature_eval_setup(sess: Session,
X: Tensor,
Z: Tensor,
data_train: DataSet,
data_test: DataSet,
eval_fn: Callable[[Tensor, Tensor], Tensor],
eval_loss_fn: Callable[[Tensor, Tensor], Tensor],
supervise_net: Optional[Callable[[Tensor], Tensor]] = None,
optimizer: tf.train.Optimizer = (
tf.train.RMSPropOptimizer(learning_rate=1e-4)),
mb_size: Optional[int] = 128,
max_iter: int = 5000,
restart_training: bool = True
) -> Callable[[Session], Tuple[Number, Number]]:
with tf.variable_scope('feature_eval'):
if supervise_net is not None:
y_logits = supervise_net(Z)
else:
y_logits = dense_net(Z, [256, data_train.dim_y])
y_hat = tf.sigmoid(y_logits)
y = tf.placeholder(tf.float32, [None] + data_train.dim_Y)
eval_loss = tf.reduce_mean(eval_loss_fn(y_logits, y))
eval_result = eval_fn(y_hat, y)
vars_fteval = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='feature_eval')
train = optimizer.minimize(eval_loss, var_list=vars_fteval)
eval_vars_initializer = tf.variables_initializer(
tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='feature_eval'))
sess.run(eval_vars_initializer)
def feature_eval(_sess: Session) -> Tuple[Number, Number]:
if restart_training:
_sess.run(eval_vars_initializer)
for _ in tqdm(range(max_iter)):
if mb_size is not None:
_mb = data_train.sample(mb_size)
else:
_mb = data_train
data_feed = {X: _mb.x, y: _mb.y}
_sess.run(train, feed_dict=data_feed)
data_feed = {X: data_test.x, y: data_test.y}
val_eval_loss = _sess.run(eval_loss, feed_dict=data_feed)
val_eval = _sess.run(eval_result, feed_dict=data_feed)
return val_eval_loss, val_eval
return feature_eval
def ClippingOptimizer(opt: tf.train.Optimizer, low, high):
original = opt.apply_gradients
def apply_gradients(grads_and_vars, *a, **kw):
app = original(grads_and_vars, *a, **kw)
with tf.name_scope('clip'):
# clip = [v.assign_add(tf.maximum(high-v, 0)+tf.minimum(low-v, 0)) for g, v in grads_and_vars]
clip = [
v.assign(tf.clip_by_value(v, low, high))
for g, v in grads_and_vars
]
step = after(app, clip, name='step')
return step
opt.apply_gradients = apply_gradients
return opt
def train_op_with_clip_and_noise(
optimizer: tf.train.Optimizer,
grads_and_vars: _GRAD_AND_VARS_TYPE,
global_step: Optional[tf.Tensor] = None,
gradient_clip: Optional[float] = None,
gradient_noise_std: Optional[float] = None,
gradient_l2_norm: Optional[tf.Tensor] = None) -> tf.Operation:
"""
Produce train op for gradients and variables with
gradient clip and adding of gradient noise if they were provided inside of
optim config
Parameters
----------
optimizer
optimizer to use
grads_and_vars
list of (gradient, variable)
global_step
global step to use in the optimizer; Caution: provide global_step only
once, if you execute this method multiple times in one session
gradient_clip
gradient clip value
gradient_noise_std
standard deviation of the noise to add to gradients
gradient_l2_norm
gradient l2 norm used for the gradient clipping
Returns
-------
train_op
training operation, which can be used inside of session run
"""
if gradient_clip is not None:
grads_and_vars = clip_grads_and_vars(grads_and_vars, gradient_clip,
gradient_l2_norm)
if gradient_noise_std is not None:
grads_and_vars = add_noise_to_grads_and_vars(grads_and_vars,
gradient_noise_std)
train_op = optimizer.apply_gradients(grads_and_vars,
global_step=global_step)
return train_op
def train_step(model: tf.keras.Model, optimizer: tf.train.Optimizer,
loss_func: loss, inputs: tf.Tensor, labels: tf.Tensor, **kwargs):
"""Kicks off training for a given model.
Args:
model (tf.keras.Model):
optimizer (tf.train.Optimizer):
loss_func (loss): Loss function.
inputs (tf.Tensor): Dataset's input features.
labels (tf.Tensor): Dataset true labels.
Keyword Args:
sparse (bool): False if labels are not one-hot encoded.
Returns:
An Operation that updates the variables in `var_list`. If `global_step`
was not `None`, that operation also increments `global_step`.
"""
return optimizer.minimize(loss=lambda: loss_func(model, inputs, labels, **kwargs),
global_step=tf.train.get_or_create_global_step())
def __init__(
self,
obs_dim: int,
latent_dim: int,
encoder: Callable[[tf.Tensor, int], tf.Tensor],
decoder: Callable[[tf.Tensor], tf.Tensor],
decoder_loss: Callable[[tf.Tensor, tf.Tensor], tf.Tensor],
optimiser: tf.train.Optimizer,
seed: int,
) -> None:
super().__init__(obs_dim, latent_dim)
tf.set_random_seed(seed)
obs_pl = tf.placeholder(tf.float32, [None, obs_dim])
latent_pl = tf.placeholder(tf.float32, [None, latent_dim])
latent_dist_params = encoder(obs_pl, latent_dim)
latent = self._build_sampled_latent(latent_dist_params)
generated_dist_params = decoder(latent)
_fixed_prior = tf.random_normal([1, latent_dim])
_generated_dist_params = decoder(latent_pl)
loss = self._kl_divergence(latent_dist_params)
loss += decoder_loss(obs_pl, generated_dist_params)
loss *= np.log2(np.e)
train_op = optimiser.minimize(loss)
session = tf.Session()
self._prior_fn = session.make_callable(_fixed_prior)
self._posterior_fn = session.make_callable(latent_dist_params,
[obs_pl])
self._generator_fn = session.make_callable(_generated_dist_params,
[latent_pl])
self._loss_fn = session.make_callable(loss, [obs_pl])
self._train_fn = session.make_callable(train_op, [obs_pl])
session.run(tf.global_variables_initializer())
self.session = session
self.saver = tf.train.Saver()
def __init__(
self,
obs_spec: dm_env.specs.Array,
action_spec: dm_env.specs.BoundedArray,
ensemble: Sequence[snt.AbstractModule],
target_ensemble: Sequence[snt.AbstractModule],
batch_size: int,
agent_discount: float,
replay_capacity: int,
min_replay_size: int,
sgd_period: int,
target_update_period: int,
optimizer: tf.train.Optimizer,
mask_prob: float,
noise_scale: float,
epsilon_fn: Callable[[int], float] = lambda _: 0.,
seed: int = None,
):
"""Bootstrapped DQN with additive prior functions."""
# Dqn configurations.
self._ensemble = ensemble
self._target_ensemble = target_ensemble
self._num_actions = action_spec.maximum - action_spec.minimum + 1
self._batch_size = batch_size
self._sgd_period = sgd_period
self._target_update_period = target_update_period
self._min_replay_size = min_replay_size
self._epsilon_fn = epsilon_fn
self._replay = replay.Replay(capacity=replay_capacity)
self._mask_prob = mask_prob
self._noise_scale = noise_scale
self._rng = np.random.RandomState(seed)
tf.set_random_seed(seed)
self._total_steps = 0
self._total_episodes = 0
self._active_head = 0
self._num_ensemble = len(ensemble)
assert len(ensemble) == len(target_ensemble)
# Making the tensorflow graph
session = tf.Session()
# Placeholders = (obs, action, reward, discount, next_obs, mask, noise)
o_tm1 = tf.placeholder(shape=(None, ) + obs_spec.shape,
dtype=obs_spec.dtype)
a_tm1 = tf.placeholder(shape=(None, ), dtype=action_spec.dtype)
r_t = tf.placeholder(shape=(None, ), dtype=tf.float32)
d_t = tf.placeholder(shape=(None, ), dtype=tf.float32)
o_t = tf.placeholder(shape=(None, ) + obs_spec.shape,
dtype=obs_spec.dtype)
m_t = tf.placeholder(shape=(None, self._num_ensemble),
dtype=tf.float32)
z_t = tf.placeholder(shape=(None, self._num_ensemble),
dtype=tf.float32)
losses = []
value_fns = []
target_updates = []
for k in range(self._num_ensemble):
model = self._ensemble[k]
target_model = self._target_ensemble[k]
q_values = model(o_tm1)
train_value = batched_index(q_values, a_tm1)
target_value = tf.stop_gradient(
tf.reduce_max(target_model(o_t), axis=-1))
target_y = r_t + z_t[:, k] + agent_discount * d_t * target_value
loss = tf.square(train_value - target_y) * m_t[:, k]
value_fn = session.make_callable(q_values, [o_tm1])
target_update = update_target_variables(
target_variables=target_model.get_all_variables(),
source_variables=model.get_all_variables(),
)
losses.append(loss)
value_fns.append(value_fn)
target_updates.append(target_update)
sgd_op = optimizer.minimize(tf.stack(losses))
self._value_fns = value_fns
self._sgd_step = session.make_callable(
sgd_op, [o_tm1, a_tm1, r_t, d_t, o_t, m_t, z_t])
self._update_target_nets = session.make_callable(target_updates)
session.run(tf.global_variables_initializer())
def __init__(
self,
obs_spec: specs.Array,
action_spec: specs.DiscreteArray,
network: snt.RNNCore,
optimizer: tf.train.Optimizer,
sequence_length: int,
td_lambda: float,
agent_discount: float,
seed: int,
):
"""A recurrent actor-critic agent."""
del action_spec # unused
tf.set_random_seed(seed)
self._sequence_length = sequence_length
self._num_transitions_in_buffer = 0
# Create the policy ops.
obs = tf.placeholder(shape=(1,) + obs_spec.shape, dtype=obs_spec.dtype)
mask = tf.placeholder(shape=(1,), dtype=tf.float32)
state = self._placeholders_like(network.initial_state(batch_size=1))
(online_logits, _), next_state = network((obs, mask), state)
action = tf.squeeze(tf.multinomial(online_logits, 1, output_dtype=tf.int32))
# Create placeholders and numpy arrays for learning from trajectories.
shapes = [obs_spec.shape, (), (), (), ()]
dtypes = [obs_spec.dtype, np.int32, np.float32, np.float32, np.float32]
placeholders = [
tf.placeholder(shape=(self._sequence_length, 1) + shape, dtype=dtype)
for shape, dtype in zip(shapes, dtypes)]
observations, actions, rewards, discounts, masks = placeholders
# Build actor and critic losses.
(logits, values), final_state = tf.nn.dynamic_rnn(
network, (observations, tf.expand_dims(masks, -1)),
initial_state=state, dtype=tf.float32, time_major=True)
(_, bootstrap_value), _ = network((obs, mask), final_state)
values, bootstrap_value = tree.map_structure(
lambda t: tf.squeeze(t, axis=-1), (values, bootstrap_value))
critic_loss, (advantages, _) = td_lambda_loss(
state_values=values,
rewards=rewards,
pcontinues=agent_discount * discounts,
bootstrap_value=bootstrap_value,
lambda_=td_lambda)
actor_loss = discrete_policy_gradient_loss(logits, actions, advantages)
# Updates.
grads_and_vars = optimizer.compute_gradients(actor_loss + critic_loss)
grads, _ = tf.clip_by_global_norm([g for g, _ in grads_and_vars], 5.)
grads_and_vars = [(g, pair[1]) for g, pair in zip(grads, grads_and_vars)]
train_op = optimizer.apply_gradients(grads_and_vars)
# Create TF session and callables.
session = tf.Session()
self._reset_fn = session.make_callable(
network.initial_state(batch_size=1))
self._policy_fn = session.make_callable(
[action, next_state], [obs, mask, state])
self._update_fn = session.make_callable(
[train_op, final_state], placeholders + [obs, mask, state])
session.run(tf.global_variables_initializer())
# Initialize numpy buffers
self.state = self._reset_fn()
self.update_init_state = self._reset_fn()
self.arrays = [
np.zeros(shape=(self._sequence_length, 1) + shape, dtype=dtype)
for shape, dtype in zip(shapes, dtypes)]
def compute_adam_gradients(self, adam: tf.train.Optimizer, loss,
variables):
from tensorflow.python.training.optimizer import Optimizer
from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.training import distribute as distribute_lib
from tensorflow.python.training import distribution_strategy_context
from tensorflow.python.util import nest
def compute_gradients(optimizer,
loss,
var_list=None,
gate_gradients=Optimizer.GATE_OP,
aggregation_method=None,
colocate_gradients_with_ops=False,
grad_loss=None):
if callable(loss):
from tensorflow.python.eager import backprop
with backprop.GradientTape() as tape:
if var_list is not None:
tape.watch(var_list)
loss_value = loss()
# Scale loss if using a "mean" loss reduction and multiple towers.
# Have to be careful to call distribute_lib.get_loss_reduction()
# *after* loss() is evaluated, so we know what loss reduction it uses.
# TODO(josh11b): Test that we handle weight decay in a reasonable way.
if (distribute_lib.get_loss_reduction() ==
variable_scope.VariableAggregation.MEAN):
num_towers = distribution_strategy_context.get_distribution_strategy(
).num_towers
if num_towers > 1:
loss_value *= (1. / num_towers)
if var_list is None:
var_list = tape.watched_variables()
# TODO(jhseu): Figure out why GradientTape's gradients don't require loss
# to be executed.
with ops.control_dependencies([loss_value]):
grads = tape.gradient(loss_value, var_list, grad_loss)
return list(zip(grads, var_list))
# Non-callable/Tensor loss case
if context.executing_eagerly():
raise RuntimeError(
"`loss` passed to Optimizer.compute_gradients should "
"be a function when eager execution is enabled.")
# Scale loss if using a "mean" loss reduction and multiple towers.
if (distribute_lib.get_loss_reduction() ==
variable_scope.VariableAggregation.MEAN):
num_towers = distribution_strategy_context.get_distribution_strategy(
).num_towers
if num_towers > 1:
loss *= (1. / num_towers)
if gate_gradients not in [
Optimizer.GATE_NONE, Optimizer.GATE_OP,
Optimizer.GATE_GRAPH
]:
raise ValueError(
"gate_gradients must be one of: Optimizer.GATE_NONE, "
"Optimizer.GATE_OP, Optimizer.GATE_GRAPH. Not %s" %
gate_gradients)
optimizer._assert_valid_dtypes([loss])
if grad_loss is not None:
optimizer._assert_valid_dtypes([grad_loss])
if var_list is None:
var_list = (variables.trainable_variables() +
ops.get_collection(
ops.GraphKeys.TRAINABLE_RESOURCE_VARIABLES))
else:
var_list = nest.flatten(var_list)
# pylint: disable=protected-access
var_list += ops.get_collection(
ops.GraphKeys._STREAMING_MODEL_PORTS)
# pylint: enable=protected-access
from tensorflow.python.training.optimizer import _get_processor
processors = [_get_processor(v) for v in var_list]
if not var_list:
raise ValueError("No variables to optimize.")
var_refs = [p.target() for p in processors]
# original gradients computation
# grads = tf.gradients(
# loss, var_refs, grad_ys=grad_loss,
# gate_gradients=(gate_gradients == Optimizer.GATE_OP),
# aggregation_method=aggregation_method,
# colocate_gradients_with_ops=colocate_gradients_with_ops)
# using gradient check-pointing
from memory_saving_gradients import gradients
# setting outputs of different networks
tensors_to_checkpoint = self.get_tensors_to_checkpoint()
# just specifying memory as parameter fails
grads = gradients(
loss,
var_refs,
grad_ys=grad_loss,
gate_gradients=(gate_gradients == Optimizer.GATE_OP),
aggregation_method=aggregation_method,
colocate_gradients_with_ops=colocate_gradients_with_ops,
checkpoints='speed')
if gate_gradients == Optimizer.GATE_GRAPH:
grads = control_flow_ops.tuple(grads)
grads_and_vars = list(zip(grads, var_list))
optimizer._assert_valid_dtypes([
v for g, v in grads_and_vars
if g is not None and v.dtype != dtypes.resource
])
return grads_and_vars
# just copied so I can change gradients
# computed_gradients = compute_gradients(adam, loss, var_list=variables)
computed_gradients = adam.compute_gradients(
loss, var_list=variables) # original gradient
return computed_gradients
