def step_logits(self, obs):
fetches = {
'a':
nest.map_structure_up_to(self._ac_structure, lambda head: head.sam,
self.net_out.self_fed_heads),
'logits':
nest.map_structure_up_to(self._ac_structure,
lambda head: head.logits,
self.net_out.self_fed_heads)
}
ret = self._forward(obs, fetches)
return _squeeze_batch_size_singleton_dim(ret['a']), ret['logits']
def setup_fetches(self, outputs):
def split_batch(template, tf_structure):
split_flatten = zip(*[
tf.split(t, self.batch_size)
for t in nest.flatten_up_to(template, tf_structure)
])
return [
nest.pack_sequence_as(template, flatten)
for flatten in split_flatten
]
if self.nc.use_self_fed_heads:
a = nest.map_structure_up_to(self._ac_structure,
lambda head: head.sam,
self.net_out.self_fed_heads)
neglogp = nest.map_structure_up_to(self._ac_structure,
lambda head: head.neglogp,
self.net_out.self_fed_heads)
flatparam = nest.map_structure_up_to(self._ac_structure,
lambda head: head.flatparam,
self.net_out.self_fed_heads)
self.all_outputs = {
'a':
split_batch(self._ac_structure, a),
'neglogp':
split_batch(self._ac_structure, neglogp),
'flatparam':
split_batch(self._ac_structure, flatparam),
'v':
tf.split(self.net_out.value_head, self.batch_size)
if self.net_out.value_head is not None else [[]] *
self.batch_size,
'state':
tf.split(self.net_out.S, self.batch_size)
if self.net_out.S is not None else [[]] * self.batch_size
}
else:
flatparam = nest.map_structure_up_to(self._ac_structure,
lambda head: head.flatparam,
self.net_out.outer_fed_heads)
self.all_outputs = {
'flatparam':
split_batch(self._ac_structure, flatparam),
'state':
tf.split(self.net_out.S, self.batch_size)
if self.net_out.S is not None else [[]] * self.batch_size
}
if self.nc.use_lstm and 'state' not in outputs:
outputs.append('state')
self.fetches = [
dict(zip(outputs, pred))
for pred in zip(*[self.all_outputs[o] for o in outputs])
]
def debatch_timestep(self, ts):
"""Debatches a single timestep.
Returns bs length of timesteps."""
traj_spec = self._traj_spec
def f(arr):
if arr is None:
return arr
l = np.split(arr, len(arr))
# remove the leading dimension
l = list(map(functools.partial(np.squeeze, axis=0), l))
return l
# split along the batch dimension
d = nest.map_structure_up_to(traj_spec, f, ts)
# determine the batch size
lens = [
len(v) for v in filter(lambda k: k is not None,
nest.flatten_up_to(traj_spec, d))
]
bs = lens[0]
assert all(x == bs for x in lens)
# Flatten and replicate by packing the sequence bs times.
d = nest.flatten_up_to(traj_spec, d)
l = []
for i in range(bs):
l.append(
nest.pack_sequence_as(
traj_spec, list(map(lambda k: k
if k is None else k[i], d))))
return l
def gym_ddpg_actor_test():
mycfg = {
'test': False,
'use_loss_type': 'none',
'use_value_head': False,
'n_v': 4,
'use_lstm': True,
'batch_size': 1,
'rollout_len': 1,
'nlstm': 64,
'hs_len': 64 * 2,
'lstm_layer_norm': True,
'weight_decay': 0.0005
}
ob_space = spaces.Box(shape=(11, ), dtype=np.float32, low=0, high=1)
ac_space = spaces.Box(shape=(2, ), low=-1.0, high=1.0, dtype=np.float32)
nc = net_config_cls(ob_space, ac_space, **mycfg)
inputs = net_inputs_placeholders_fun(nc)
out = net_build_fun(inputs, nc, scope='gym_ddpg')
sample = ob_space.sample()
sess = tf.Session()
tf.global_variables_initializer().run(session=sess)
feed_dict = {inputs.X: [sample]}
feed_dict[inputs.S] = np.zeros(shape=[1, nc.hs_len])
feed_dict[inputs.M] = np.zeros(shape=[1])
from tensorflow.contrib.framework import nest
import tpolicies.tp_utils as tp_utils
ac_structure = tp_utils.template_structure_from_gym_space(ac_space)
a = nest.map_structure_up_to(ac_structure, lambda head: head.sam,
out.self_fed_heads)
sam = sess.run(a, feed_dict=feed_dict)
print(sam)
pass
def forward_squeezed(self, obs):
if self.infserver_addr is None:
# prepare fetches dict
fetches = {
'a': nest.map_structure_up_to(self._ac_structure, lambda head: head.sam,
self.net_out.self_fed_heads),
'neglogp': nest.map_structure_up_to(self._ac_structure,
lambda head: head.neglogp,
self.net_out.self_fed_heads),
'v': self.net_out.value_head if self.net_out.value_head is not None else []
}
else:
fetches = None
ret = self._forward(obs, fetches=fetches)
ret['state'] = self._last_state
return ret.pop('a'), ret
def step(self, obs):
# prepare fetches dict
fetches = {
'a':
nest.map_structure_up_to(self._ac_structure, lambda head: head.sam,
self.net_out.self_fed_heads),
}
ret = self._forward(obs, fetches)
return _squeeze_batch_size_singleton_dim(ret['a'])
def _stack(trajs, traj_spec):
stacked_trajs = []
def f(spec, *l):
l = list(filter(lambda k: k is not None, l))
# copy leads to crazy cpu util
return np.stack(l, axis=0).astype(spec.dtype, copy=False)
return nest.map_structure_up_to(traj_spec, f, traj_spec, *trajs)
def batch(trajs, traj_spec):
batched_trajs = Trajectory._stack(trajs, traj_spec)
def f(l):
# make time major
return None if l is None else np.swapaxes(l, 0, 1)
batched_trajs = nest.map_structure_up_to(traj_spec, f, batched_trajs)
return batched_trajs
def step(self, obs):
if self.infserver_addr is None:
# prepare fetches dict
fetches = {
'a': nest.map_structure_up_to(self._ac_structure, lambda head: head.sam,
self.net_out.self_fed_heads),
}
else:
fetches = None
ret = self._forward(obs, fetches=fetches)
return ret['a']
def _stack_ts(self, timesteps):
"""Should be called after _make_step_spec."""
dict_tss = []
for ts in timesteps:
dict_tss.append(dict(ts._asdict()))
def f(spec, *l):
return np.stack(l, axis=0).astype(spec.dtype)
stacked_ts = nest.map_structure_up_to(self._step_spec, f,
self._step_spec, *dict_tss)
return TimeStep(**stacked_ts)
def forward_squeezed(self, obs):
# prepare fetches dict
fetches = {
'a':
nest.map_structure_up_to(self._ac_structure, lambda head: head.sam,
self.net_out.self_fed_heads),
'neglogp':
nest.map_structure_up_to(self._ac_structure,
lambda head: head.neglogp,
self.net_out.self_fed_heads),
# 'logits': nest.map_structure_up_to(self._ac_structure,
# lambda head: head.logits,
# self.net_out.self_fed_heads),
'v':
self.net_out.value_head
if self.net_out.value_head is not None else []
}
ret = self._forward(obs, fetches)
a = _squeeze_batch_size_singleton_dim(ret['a'])
v = _squeeze_batch_size_singleton_dim(ret['v'])
neglogp = _squeeze_batch_size_singleton_dim(ret['neglogp'])
# logits = _squeeze_batch_size_singleton_dim(ret['logits'])
return a, v, self._last_state, neglogp
def head_param(self, obs, action=None):
if self.infserver_addr is None:
if action is None:
assert self.net_out.self_fed_heads is not None
heads = self.net_out.self_fed_heads
else:
assert self.net_out.outer_fed_heads is not None
heads = self.net_out.outer_fed_heads
fetches = {'flatparam': nest.map_structure_up_to(
self._ac_structure, lambda head: head.flatparam, heads)}
else:
fetches = None
ret = self._forward(obs, fetches, action)
return ret['flatparam']
def logits(self, obs, action=None):
if action is None:
assert self.net_out.self_fed_heads is not None
heads = self.net_out.self_fed_heads
else:
assert self.net_out.outer_fed_heads is not None
heads = self.net_out.outer_fed_heads
fetches = {
'logits':
nest.map_structure_up_to(self._ac_structure,
lambda head: head.logits, heads)
}
ret = self._forward(obs, fetches, action)
return _squeeze_batch_size_singleton_dim(ret['logits'])
def conv_lstm_actor_test():
mycfg = {
'test': False,
'use_loss_type': 'none',
'use_value_head': False,
'n_v': 4,
'sync_statistics': None,
'use_lstm': True,
'batch_size': 1,
'rollout_len': 1,
'nlstm': 64,
'hs_len': 64 * 2,
'lstm_layer_norm': True,
'weight_decay': 0.0005
}
ob_space = spaces.Tuple([
spaces.Tuple([
spaces.Box(shape=(11, 11, 22), dtype=np.float32, low=0, high=1),
spaces.Box(shape=(2, ), dtype=np.int32, low=0, high=10),
spaces.Box(shape=[6], dtype=np.bool, low=0, high=1)
])
] * 2)
ac_space = spaces.Tuple([spaces.Discrete(n=6)] * 2)
nc = net_config_cls(ob_space, ac_space, **mycfg)
inputs = net_inputs_placeholders_fun(nc)
out = net_build_fun(inputs, nc, scope='conv_lstm')
sample = ob_space.sample()
sess = tf.Session()
tf.global_variables_initializer().run(session=sess)
feed_dict = {}
for s, input in zip(sample, inputs.X):
for x_np, x in zip(s, input):
feed_dict[x] = [x_np]
feed_dict[inputs.S] = np.zeros(shape=[1, nc.hs_len])
feed_dict[inputs.M] = np.zeros(shape=[1])
from tensorflow.contrib.framework import nest
import tpolicies.tp_utils as tp_utils
ac_structure = tp_utils.template_structure_from_gym_space(ac_space)
a = nest.map_structure_up_to(ac_structure, lambda head: head.sam,
out.self_fed_heads)
sam = sess.run(a, feed_dict=feed_dict)
print(sam)
pass
def debatch_and_stack(self):
"""Remove the leading batch dimension and then stack on timestamp.
Returns list of stacked timesteps for each batch."""
traj_spec = self._traj_spec
def f(arr):
if arr is None:
return arr
l = np.split(arr, len(arr))
# remove the leading dimension
l = list(map(functools.partial(np.squeeze, axis=0), l))
return l
l = []
for traj in self._trajs:
# split along the batch dimension
d = nest.map_structure_up_to(traj_spec, f, traj)
# determine the batch size
lens = [
len(v) for v in filter(lambda k: k is not None,
nest.flatten_up_to(traj_spec, d))
]
bs = lens[0]
assert all(x == bs for x in lens)
# Flatten and replicate by packing the sequence bs times.
d = nest.flatten_up_to(traj_spec, d)
if not l:
l = [[] for _ in range(bs)]
for i in range(bs):
l[i].append(
nest.pack_sequence_as(
traj_spec,
list(map(lambda k: k if k is None else k[i], d))))
return list(
map(
functools.partial(Trajectory._stack,
traj_spec=self._traj_spec), l))
def conv_lstm(inputs: ConvLstmInputs,
nc: ConvLstmConfig,
scope=None) -> ConvLstmOutputs:
"""create the whole net for conv-lstm"""
with tf.variable_scope(scope, default_name='pommerman') as sc:
# NOTE: use name_scope, in case multiple parameter-sharing nets are built
net_name_scope = tf.get_default_graph().get_name_scope()
endpoints_collections = net_name_scope + '_endpoints'
X = inputs.X
if nc.n_player == 1:
X = (X, )
ac_spaces = (nc.ac_space, )
else:
ac_spaces = tuple(nc.ac_space.spaces)
S = tf.split(inputs.S, nc.n_player, axis=1)
# make body
y = []
hs_new = []
heads = []
if nc.use_lstm and nc.n_player > 1:
nc.hs_len //= nc.n_player
nc.nlstm //= nc.n_player
for input, s, ac_space in zip(X, S, ac_spaces):
with tf.variable_scope('body', reuse=tf.AUTO_REUSE):
x = tfc_layers.conv2d(input[0],
nc.spa_ch_dim, [3, 3],
scope='conv0')
x = tfc_layers.conv2d(x, nc.spa_ch_dim, [5, 5], scope='conv1')
x = tfc_layers.conv2d(x,
nc.spa_ch_dim * 2, [3, 3],
scope='conv2')
x = tfc_layers.conv2d(x,
nc.spa_ch_dim * 2, [5, 5],
scope='conv3')
x = tfc_layers.conv2d(x,
nc.spa_ch_dim * 4, [3, 3],
scope='conv4')
pos = tf.to_int32(input[1])
ind = tf.concat(
[tf.expand_dims(tf.range(nc.batch_size), 1), pos], axis=1)
x = tf.gather_nd(x, ind)
if nc.use_lstm:
with tf.variable_scope('lstm_embed'):
x, hs = tp_layers.lstm_embed_block(
inputs_x=x,
inputs_hs=s,
inputs_mask=inputs.M,
nc=nc)
hs_new.append(hs)
y.append(x)
# make action head
with tf.variable_scope('action', reuse=tf.AUTO_REUSE):
head_logits = tfc_layers.fully_connected(x,
ac_space.n,
activation_fn=None,
normalizer_fn=None,
scope='logits')
if len(input) > 1:
head_logits = tp_ops.mask_logits(head_logits, input[2])
head = tp_layers.to_action_head(head_logits, CategoricalPdType)
heads.append(head)
if nc.use_lstm:
hs_new = tf.concat(hs_new, axis=1)
if nc.n_player > 1:
nc.hs_len *= nc.n_player
nc.nlstm *= nc.n_player
y = tf.concat(y, axis=1)
heads = tp_utils.pack_sequence_as_structure_like_gym_space(
nc.ac_space, heads)
if nc.n_player == 1:
heads = heads[0]
# make value head
vf = None
if nc.use_value_head:
assert nc.n_player == 2
with tf.variable_scope('vf'):
vf = tfc_layers.fully_connected(y, nc.spa_ch_dim * 4)
vf = tfc_layers.fully_connected(vf, nc.spa_ch_dim * 2)
vf = tfc_layers.fully_connected(vf,
nc.n_v,
activation_fn=None,
normalizer_fn=None)
# make loss
loss = None
if nc.use_loss_type in ['rl', 'rl_ppo', 'rl_vtrace']:
assert nc.n_player == 2
with tf.variable_scope('losses'):
# regularization loss
total_reg_loss = tf.losses.get_regularization_losses(
scope=sc.name)
# entropy loss
entropy_loss = nest.map_structure_up_to(
ac_spaces, lambda head: tf.reduce_mean(head.ent), heads)
# ppo loss
neglogp = nest.map_structure_up_to(
ac_spaces, lambda head, ac: head.pd.neglogp(ac), heads,
inputs.A)
loss_endpoints = {}
for k, v in enumerate(entropy_loss):
loss_endpoints['ent_' + str(k)] = v
if nc.use_loss_type == 'rl' or nc.use_loss_type == 'rl_ppo':
pg_loss, value_loss = tp_losses.ppo_loss(
neglogp=neglogp,
oldneglogp=inputs.neglogp,
vpred=vf,
R=inputs.R,
V=inputs.V,
masks=None,
reward_weights=nc.reward_weights,
adv_normalize=True,
sync_statistics=nc.sync_statistics)
elif nc.use_loss_type == 'rl_vtrace':
def _batch_to_TB(tsr):
return tf.transpose(
tf.reshape(tsr,
shape=(nc.nrollout, nc.rollout_len)))
lam = tf.convert_to_tensor(nc.lam, tf.float32)
vpred_list = [
_batch_to_TB(v) for v in tf.split(vf, nc.n_v, axis=1)
]
reward_list = [
_batch_to_TB(r)
for r in tf.split(inputs.r, nc.n_v, axis=1)
]
discounts = _batch_to_TB(inputs.discount)
value_loss = []
for values, rewards in zip(vpred_list, reward_list):
value_loss.append(
tp_losses.td_lambda(values,
rewards,
discounts,
lam=lam))
value_loss = tf.stack(value_loss)
neglogp_list = [
_batch_to_TB(neglogp)
for neglogp in nest.flatten(neglogp)
]
oldneglogp_list = [
_batch_to_TB(oldneglogp)
for oldneglogp in nest.flatten(inputs.neglogp)
]
shaped_values = tf.matmul(vf,
nc.reward_weights,
transpose_b=True)
shaped_rewards = tf.matmul(inputs.r,
nc.reward_weights,
transpose_b=True)
values = tf.transpose(
tf.reshape(shaped_values,
shape=(nc.nrollout, nc.rollout_len)))
rewards = tf.transpose(
tf.reshape(shaped_rewards,
shape=(nc.nrollout, nc.rollout_len)))
pg_loss = tf.reduce_sum([
tp_losses.vtrace_loss(neglogp, oldneglogp, None,
values, rewards, discounts, 1.0,
1.0) for oldneglogp, neglogp in
zip(oldneglogp_list, neglogp_list)
])
upgo_loss = tp_losses.upgo_loss(
tf.stack(neglogp_list, axis=-1),
tf.stack(oldneglogp_list, axis=-1), None,
vpred_list[0], reward_list[0], discounts)
loss_endpoints['upgo_loss'] = upgo_loss
loss_endpoints['pg_loss'] = pg_loss
if len(value_loss.shape) == 0:
loss_endpoints['value_loss'] = value_loss
else:
for i in range(value_loss.shape[0]):
loss_endpoints['value_loss_' + str(i)] = value_loss[i]
loss = ConvLstmLosses(total_reg_loss=total_reg_loss,
pg_loss=pg_loss,
value_loss=value_loss,
entropy_loss=entropy_loss,
loss_endpoints=loss_endpoints)
# collect vars, endpoints, etc.
trainable_vars = _make_vars(sc)
endpoints = OrderedDict() # TODO
return ConvLstmOutputs(self_fed_heads=heads,
outer_fed_heads=heads,
S=hs_new,
loss=loss,
vars=trainable_vars,
endpoints=endpoints,
value_head=vf)
def mnet_v6d6_loss(inputs: MNetV6Inputs,
outer_fed_heads,
value_head,
consts: MNetV6Consts,
nc: MNetV6Config,
net_level_scope: str,
structured_mw=None,
scope=None):
# regularization loss. Only `variable`s are involved, so it is safe to
# collect them using regular expression, e.g., 'mnet_v5.*', regardless
# of the current name_scope (e.g., 'mnet_v5_1', 'mnet_v5_2', ...)
total_reg_loss = tf.losses.get_regularization_loss(
scope='{}.*'.format(net_level_scope))
total_il_loss = None
pg_loss = None
value_loss = None
entropy_loss = None
distill_loss = None
loss_endpoints = {}
example_ac_sp = tp_utils.map_gym_space_to_structure(
lambda x: None, nc.ac_space)
with tf.variable_scope(scope, default_name='mnet_v6_losses'):
if nc.use_loss_type in ['il', 'rl', 'rl_ppo', 'rl_ppo2', 'rl_vtrace']:
# head masks and structure template
if structured_mw is None:
mw = _action_mask_weights(inputs_ab=inputs.A['A_AB'],
inputs_arg_mask=consts.arg_mask,
weights_include_ab=True)
structured_mw = tp_utils.pack_sequence_as_structure_like_gym_space(
nc.ac_space, mw)
outer_fed_head_pds = nest.map_structure_up_to(
example_ac_sp, lambda head: head.pd, outer_fed_heads)
if nc.use_loss_type == 'il':
# build imitation learning loss the cross entropy
total_il_loss, head_xe_loss = tp_losses.multi_head_neglogp_loss(
inputs_action_pds=outer_fed_head_pds,
inputs_action_labels=inputs.A,
inputs_mask_weights=structured_mw,
set_loss=nc.il_multi_label_loss,
)
assert type(head_xe_loss) == OrderedDict
loss_endpoints = head_xe_loss
elif nc.use_loss_type in ['rl', 'rl_ppo', 'rl_ppo2', 'rl_vtrace']:
# build rl losses
# the entropy regularizer
entropy_loss = nest.map_structure_up_to(
example_ac_sp,
lambda head, mask: tf.reduce_mean(head.ent * mask),
outer_fed_heads, structured_mw)
# distillation loss, i.e., the teacher-student KL regularizer
distill_loss = None
ab_distill_loss = None
if nc.distillation:
outer_fed_head_pds = nest.map_structure_up_to(
example_ac_sp, lambda head: head.pd, outer_fed_heads)
distill_loss = tp_losses.distill_loss(
student_pds=outer_fed_head_pds,
teacher_logits=inputs.logits,
masks=structured_mw)
ab_pd = outer_fed_head_pds['A_AB']
teacher_logit = inputs.logits['A_AB']
# TODO: this is from definition of position encoding, remove it?
first_4mins_mask = tf.cast(
inputs.X['X_VEC_GAME_PROG'][:, -1] >= np.cos(
60 * 4 * np.power(10000, -62 / 64)), tf.float32)
first_4mins_mask *= tf.cast((tf.reduce_sum(
inputs.X['Z_BUILD_ORDER'], axis=[1, 2]) > 0),
tf.float32)
ab_distill_loss = tp_losses.distill_loss(
ab_pd, teacher_logit, first_4mins_mask)
# the main policy gradient loss
outer_fed_head_neglogp = nest.map_structure_up_to(
example_ac_sp, lambda head, ac: head.pd.neglogp(ac),
outer_fed_heads, inputs.A)
loss_endpoints = {}
if nc.use_loss_type == 'rl' or nc.use_loss_type == 'rl_ppo':
# PPO loss
pg_loss, value_loss = tp_losses.ppo_loss(
outer_fed_head_neglogp,
inputs.neglogp,
value_head,
inputs.R,
inputs.V,
masks=structured_mw,
reward_weights=nc.reward_weights,
merge_pi=nc.merge_pi,
adv_normalize=nc.adv_normalize,
clip_range=nc.clip_range,
sync_statistics=nc.sync_statistics,
)
elif nc.use_loss_type in ['rl_ppo2', 'rl_vtrace']:
# Note: we need convert the shape (batch_size, ...) to the shape
# (T, B, ...) where T=nc.rollout_len, B=nc.nrollout, batch_size=B*T
# When computing ppo2-loss and value-loss, only T-1 time steps are
# used due to the value bootstrap at the tail. When doing so, the
# [:-1] indexing, leading to (T - 1, B, ...) tensor slice, makes life
# much easier
def _batch_to_TB(tsr):
return tf.transpose(
tf.reshape(tsr,
shape=(nc.nrollout, nc.rollout_len)))
# make the len=n_action_heads lists for action-head stuff
# for tensor entry, shape (batch_size, ...) -> shape (T, B, ...)
neglogp_list = [
_batch_to_TB(neglogp)
for neglogp in nest.flatten(outer_fed_head_neglogp)
]
oldneglogp_list = [
_batch_to_TB(oldneglogp)
for oldneglogp in nest.flatten(inputs.neglogp)
]
mask_list = [
_batch_to_TB(mw) for mw in nest.flatten(structured_mw)
]
# make the len=n_v lists for value-head stuff
# for tensor entry, shape (batch_size, ...) -> shape (T, B, ...)
# as aforementioned
vpred_list = [
_batch_to_TB(v)
for v in tf.split(value_head, nc.n_v, axis=1)
]
reward_list = [
_batch_to_TB(r)
for r in tf.split(inputs.r, nc.n_v, axis=1)
]
discounts = _batch_to_TB(inputs.discount)
# upgo_loss only use the win_loss, i.e, v[0]
upgo_loss = tp_losses.upgo_loss(
tf.stack(neglogp_list, axis=-1),
tf.stack(oldneglogp_list, axis=-1),
tf.stack(mask_list, axis=-1), vpred_list[0],
reward_list[0], discounts)
loss_endpoints['upgo_loss'] = upgo_loss
if nc.use_loss_type == 'rl_ppo2':
# PPO2 loss
# reward_weights size should be consistent with n_v
reward_weights = tf.squeeze(
tf.convert_to_tensor(nc.reward_weights,
tf.float32))
assert reward_weights.shape.as_list(
)[0] == len(reward_list), (
'For ppo2 loss, reward_weight size must be the same with number of'
' value head: each reward_weight element must correspond to one '
'value-head exactly.')
# lambda for td-lambda or lambda-return
assert nc.lam is not None, (
'building rl_ppo2, but lam for '
'lambda-return is None.')
lam = tf.convert_to_tensor(nc.lam, tf.float32)
# for each value-head, compute the corresponding policy gradient loss
# and the value loss
pg_loss, value_loss = [], []
for vpred, reward in zip(vpred_list, reward_list):
# compute the lambda-Return `R` in shape (T - 1, B)
# [:-1] means discarding the last one,
# [1:] means an off-one alignment.
# back_prop=False means R = tf.stop_gradient(R)
with tf.device("/cpu:0"):
R = multistep_forward_view(reward[:-1],
discounts[:-1],
vpred[1:],
lambda_=lam,
back_prop=False)
# compute the ppo2 loss using this value-head for each of the
# n_action_heads action-head; then reduce them
# [:-1] means discarding the last one and using only T - 1 time
# steps
_ploss = [
tp_losses.ppo2_loss(
neglogp[:-1],
oldneglogp[:-1],
tf.stop_gradient(vpred)[:-1],
R, # has been stop_gradient above
mask[:-1],
adv_normalize=nc.adv_normalize,
clip_range=nc.clip_range,
sync_statistics=nc.sync_statistics)
for neglogp, oldneglogp, mask in zip(
neglogp_list, oldneglogp_list, mask_list)
]
pg_loss.append(tf.reduce_sum(_ploss))
# compute the value loss for this value-head
value_loss.append(
tf.reduce_mean(0.5 *
tf.square(R - vpred[:-1])))
# element-wise times reward_weight and the pg_loss for that value-head
pg_loss = tf.stack(
pg_loss) * reward_weights # shape (n_v,)
# make the final pg_loss, value_loss in desired format
pg_loss = tf.reduce_sum(pg_loss)
value_loss = tf.stack(value_loss)
else:
# vtrace loss
# lambda for td-lambda or lambda-return
assert nc.lam is not None, (
'building rl_vtrace, but lam for '
'td-lambda is None.')
lam = tf.convert_to_tensor(nc.lam, tf.float32)
value_loss = []
for values, rewards in zip(vpred_list, reward_list):
value_loss.append(
tp_losses.td_lambda(values,
rewards,
discounts,
lam=lam))
shaped_values = tf.matmul(value_head,
nc.reward_weights,
transpose_b=True)
shaped_rewards = tf.matmul(inputs.r,
nc.reward_weights,
transpose_b=True)
values = tf.transpose(
tf.reshape(shaped_values,
shape=(nc.nrollout, nc.rollout_len)))
rewards = tf.transpose(
tf.reshape(shaped_rewards,
shape=(nc.nrollout, nc.rollout_len)))
pg_loss = tf.reduce_sum([
tp_losses.vtrace_loss(neglogp, oldneglogp, mask,
values, rewards, discounts,
1.0, 1.0)
for oldneglogp, neglogp, mask in zip(
oldneglogp_list, neglogp_list, mask_list)
])
value_loss = tf.stack(value_loss)
# TODO: maybe more rl endpoints
# policy gradient loss must be scalar
loss_endpoints['pg_loss'] = pg_loss
# value loss can be scalar or vector
if len(value_loss.shape) == 0:
loss_endpoints['value_loss'] = value_loss
else:
for i in range(value_loss.shape[0]):
loss_endpoints['value_loss_' + str(i)] = value_loss[i]
for k, v in entropy_loss.items():
loss_endpoints['ent_' + k] = v
if nc.distillation:
for k, v in distill_loss.items():
loss_endpoints['distill_' + k] = v
loss_endpoints['distill_ab_bf4mins'] = ab_distill_loss
else:
print('use_loss_type: {}. Nothing done.'.format(nc.use_loss_type))
pass
return MNetV6Losses(total_reg_loss=total_reg_loss,
total_il_loss=total_il_loss,
pg_loss=pg_loss,
value_loss=value_loss,
entropy_loss=entropy_loss,
distill_loss=distill_loss,
loss_endpoints=loss_endpoints)
def conv_lstm(inputs: ConvLstmInputs,
nc: ConvLstmConfig,
scope=None) -> ConvLstmOutputs:
"""create the whole net for conv-lstm"""
with tf.variable_scope(scope, default_name='pommerman') as sc:
# NOTE: use name_scope, in case multiple parameter-sharing nets are built
net_name_scope = tf.get_default_graph().get_name_scope()
endpoints_collections = net_name_scope + '_endpoints'
X = inputs.X
if nc.n_player == 1:
X = (X, )
ac_spaces = (nc.ac_space, )
else:
ac_spaces = tuple(nc.ac_space.spaces)
S = tf.split(inputs.S, nc.n_player, axis=1)
# make body
y = []
hs_new = []
heads = []
for input, s, ac_space in zip(X, S, ac_spaces):
with tf.variable_scope('body', reuse=tf.AUTO_REUSE):
x = tfc_layers.conv2d(input[0],
nc.spa_ch_dim, [3, 3],
scope='conv0')
x = tfc_layers.conv2d(x, nc.spa_ch_dim, [5, 5], scope='conv1')
x = tfc_layers.conv2d(x,
nc.spa_ch_dim * 2, [3, 3],
scope='conv2')
x = tfc_layers.conv2d(x,
nc.spa_ch_dim * 2, [5, 5],
scope='conv3')
x = tfc_layers.conv2d(x,
nc.spa_ch_dim * 4, [3, 3],
scope='conv4')
pos = tf.to_int32(input[1])
ind = tf.concat(
[tf.expand_dims(tf.range(nc.batch_size), 1), pos], axis=1)
x = tf.gather_nd(x, ind)
if nc.use_lstm:
with tf.variable_scope('lstm_embed'):
x, hs = _lstm_embed_block(inputs_x=x,
inputs_hs=s,
inputs_mask=inputs.M,
nc=nc)
hs_new.append(hs)
y.append(x)
# make action head
with tf.variable_scope('action', reuse=tf.AUTO_REUSE):
head_logits = tfc_layers.fully_connected(x,
ac_space.n,
activation_fn=None,
normalizer_fn=None,
scope='logits')
if len(input) > 1:
head_logits = tp_ops.mask_logits(head_logits, input[2])
head = tp_layers.to_action_head(head_logits, CategoricalPdType)
heads.append(head)
if nc.use_lstm:
hs_new = tf.concat(hs_new, axis=1)
y = tf.concat(y, axis=1)
heads = tp_utils.pack_sequence_as_structure_like_gym_space(
nc.ac_space, heads)
if nc.n_player == 1:
heads = heads[0]
# make value head
vf = None
if nc.use_value_head:
with tf.variable_scope('vf'):
vf = tfc_layers.fully_connected(y, nc.spa_ch_dim * 4)
vf = tfc_layers.fully_connected(vf, nc.spa_ch_dim * 2)
vf = tfc_layers.fully_connected(vf,
nc.n_v,
activation_fn=None,
normalizer_fn=None)
# make loss
loss = None
if nc.use_loss_type == 'rl':
# regularization loss
total_reg_loss = tf.losses.get_regularization_losses(scope=sc.name)
with tf.variable_scope('losses'):
# ppo loss
neglogp = nest.map_structure_up_to(
ac_spaces, lambda head, ac: head.pd.neglogp(ac), heads,
inputs.A)
ppo_loss, value_loss = tp_losses.ppo_loss(
neglogp=neglogp,
oldneglogp=inputs.neglogp,
vpred=vf,
R=inputs.R,
V=inputs.V,
masks=None,
reward_weights=None,
adv_normalize=True,
sync_statistics=nc.sync_statistics)
# entropy loss
entropy_loss = nest.map_structure_up_to(
ac_spaces, lambda head: tf.reduce_mean(head.ent), heads)
loss_endpoints = {}
loss = ConvLstmLosses(total_reg_loss=total_reg_loss,
pg_loss=ppo_loss,
value_loss=value_loss,
entropy_loss=entropy_loss,
loss_endpoints=loss_endpoints)
# collect vars, endpoints, etc.
trainable_vars = _make_vars(sc)
endpoints = OrderedDict() # TODO
return ConvLstmOutputs(self_fed_heads=heads,
outer_fed_heads=heads,
S=hs_new,
loss=loss,
vars=trainable_vars,
endpoints=endpoints,
value_head=vf)