def __init__(self,
observation_space,
action_space,
config,
existing_inputs=None):
config = dict(ray.rllib.agents.impala.impala.DEFAULT_CONFIG, **config)
assert config["batch_mode"] == "truncate_episodes", \
"Must use `truncate_episodes` batch mode with V-trace."
self.config = config
self.sess = tf.get_default_session()
self.grads = None
if isinstance(action_space, gym.spaces.Discrete):
is_multidiscrete = False
output_hidden_shape = [action_space.n]
elif isinstance(action_space, gym.spaces.multi_discrete.MultiDiscrete):
is_multidiscrete = True
output_hidden_shape = action_space.nvec.astype(np.int32)
elif self.config["vtrace"]:
raise UnsupportedSpaceException(
"Action space {} is not supported for APPO + VTrace.",
format(action_space))
else:
is_multidiscrete = False
output_hidden_shape = 1
# Policy network model
dist_class, logit_dim = ModelCatalog.get_action_dist(
action_space, self.config["model"])
# Create input placeholders
if existing_inputs:
if self.config["vtrace"]:
actions, dones, behaviour_logits, rewards, observations, \
prev_actions, prev_rewards = existing_inputs[:7]
existing_state_in = existing_inputs[7:-1]
existing_seq_lens = existing_inputs[-1]
else:
actions, dones, behaviour_logits, rewards, observations, \
prev_actions, prev_rewards, adv_ph, value_targets = \
existing_inputs[:9]
existing_state_in = existing_inputs[9:-1]
existing_seq_lens = existing_inputs[-1]
else:
actions = ModelCatalog.get_action_placeholder(action_space)
dones = tf.placeholder(tf.bool, [None], name="dones")
rewards = tf.placeholder(tf.float32, [None], name="rewards")
behaviour_logits = tf.placeholder(tf.float32, [None, logit_dim],
name="behaviour_logits")
observations = tf.placeholder(tf.float32, [None] +
list(observation_space.shape))
existing_state_in = None
existing_seq_lens = None
if not self.config["vtrace"]:
adv_ph = tf.placeholder(tf.float32,
name="advantages",
shape=(None, ))
value_targets = tf.placeholder(tf.float32,
name="value_targets",
shape=(None, ))
self.observations = observations
# Unpack behaviour logits
unpacked_behaviour_logits = tf.split(behaviour_logits,
output_hidden_shape,
axis=1)
# Setup the policy
dist_class, logit_dim = ModelCatalog.get_action_dist(
action_space, self.config["model"])
prev_actions = ModelCatalog.get_action_placeholder(action_space)
prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward")
self.model = ModelCatalog.get_model(
{
"obs": observations,
"prev_actions": prev_actions,
"prev_rewards": prev_rewards,
"is_training": self._get_is_training_placeholder(),
},
observation_space,
logit_dim,
self.config["model"],
state_in=existing_state_in,
seq_lens=existing_seq_lens)
unpacked_outputs = tf.split(self.model.outputs,
output_hidden_shape,
axis=1)
dist_inputs = unpacked_outputs if is_multidiscrete else \
self.model.outputs
prev_dist_inputs = unpacked_behaviour_logits if is_multidiscrete else \
behaviour_logits
action_dist = dist_class(dist_inputs)
prev_action_dist = dist_class(prev_dist_inputs)
values = self.model.value_function()
self.value_function = values
self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
tf.get_variable_scope().name)
def make_time_major(tensor, drop_last=False):
"""Swaps batch and trajectory axis.
Args:
tensor: A tensor or list of tensors to reshape.
drop_last: A bool indicating whether to drop the last
trajectory item.
Returns:
res: A tensor with swapped axes or a list of tensors with
swapped axes.
"""
if isinstance(tensor, list):
return [make_time_major(t, drop_last) for t in tensor]
if self.model.state_init:
B = tf.shape(self.model.seq_lens)[0]
T = tf.shape(tensor)[0] // B
else:
# Important: chop the tensor into batches at known episode cut
# boundaries. TODO(ekl) this is kind of a hack
T = self.config["sample_batch_size"]
B = tf.shape(tensor)[0] // T
rs = tf.reshape(tensor,
tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0))
# swap B and T axes
res = tf.transpose(
rs,
[1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0]))))
if drop_last:
return res[:-1]
return res
if self.model.state_in:
max_seq_len = tf.reduce_max(self.model.seq_lens) - 1
mask = tf.sequence_mask(self.model.seq_lens, max_seq_len)
mask = tf.reshape(mask, [-1])
else:
mask = tf.ones_like(rewards)
# Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc.
if self.config["vtrace"]:
logger.info("Using V-Trace surrogate loss (vtrace=True)")
# Prepare actions for loss
loss_actions = actions if is_multidiscrete else tf.expand_dims(
actions, axis=1)
self.loss = VTraceSurrogateLoss(
actions=make_time_major(loss_actions, drop_last=True),
prev_actions_logp=make_time_major(
prev_action_dist.logp(actions), drop_last=True),
actions_logp=make_time_major(action_dist.logp(actions),
drop_last=True),
action_kl=prev_action_dist.kl(action_dist),
actions_entropy=make_time_major(action_dist.entropy(),
drop_last=True),
dones=make_time_major(dones, drop_last=True),
behaviour_logits=make_time_major(unpacked_behaviour_logits,
drop_last=True),
target_logits=make_time_major(unpacked_outputs,
drop_last=True),
discount=config["gamma"],
rewards=make_time_major(rewards, drop_last=True),
values=make_time_major(values, drop_last=True),
bootstrap_value=make_time_major(values)[-1],
valid_mask=make_time_major(mask, drop_last=True),
vf_loss_coeff=self.config["vf_loss_coeff"],
entropy_coeff=self.config["entropy_coeff"],
clip_rho_threshold=self.config["vtrace_clip_rho_threshold"],
clip_pg_rho_threshold=self.
config["vtrace_clip_pg_rho_threshold"],
clip_param=self.config["clip_param"])
else:
logger.info("Using PPO surrogate loss (vtrace=False)")
self.loss = PPOSurrogateLoss(
prev_actions_logp=make_time_major(
prev_action_dist.logp(actions)),
actions_logp=make_time_major(action_dist.logp(actions)),
action_kl=prev_action_dist.kl(action_dist),
actions_entropy=make_time_major(action_dist.entropy()),
values=make_time_major(values),
valid_mask=make_time_major(mask),
advantages=make_time_major(adv_ph),
value_targets=make_time_major(value_targets),
vf_loss_coeff=self.config["vf_loss_coeff"],
entropy_coeff=self.config["entropy_coeff"],
clip_param=self.config["clip_param"])
# KL divergence between worker and learner logits for debugging
model_dist = MultiCategorical(unpacked_outputs)
behaviour_dist = MultiCategorical(unpacked_behaviour_logits)
kls = model_dist.kl(behaviour_dist)
if len(kls) > 1:
self.KL_stats = {}
for i, kl in enumerate(kls):
self.KL_stats.update({
"mean_KL_{}".format(i):
tf.reduce_mean(kl),
"max_KL_{}".format(i):
tf.reduce_max(kl),
"median_KL_{}".format(i):
tf.contrib.distributions.percentile(kl, 50.0),
})
else:
self.KL_stats = {
"mean_KL": tf.reduce_mean(kls[0]),
"max_KL": tf.reduce_max(kls[0]),
"median_KL": tf.contrib.distributions.percentile(kls[0], 50.0),
}
# Initialize TFPolicyGraph
loss_in = [
("actions", actions),
("dones", dones),
("behaviour_logits", behaviour_logits),
("rewards", rewards),
("obs", observations),
("prev_actions", prev_actions),
("prev_rewards", prev_rewards),
]
if not self.config["vtrace"]:
loss_in.append(("advantages", adv_ph))
loss_in.append(("value_targets", value_targets))
LearningRateSchedule.__init__(self, self.config["lr"],
self.config["lr_schedule"])
TFPolicyGraph.__init__(
self,
observation_space,
action_space,
self.sess,
obs_input=observations,
action_sampler=action_dist.sample(),
action_prob=action_dist.sampled_action_prob(),
loss=self.loss.total_loss,
model=self.model,
loss_inputs=loss_in,
state_inputs=self.model.state_in,
state_outputs=self.model.state_out,
prev_action_input=prev_actions,
prev_reward_input=prev_rewards,
seq_lens=self.model.seq_lens,
max_seq_len=self.config["model"]["max_seq_len"],
batch_divisibility_req=self.config["sample_batch_size"])
self.sess.run(tf.global_variables_initializer())
values_batched = make_time_major(values,
drop_last=self.config["vtrace"])
self.stats_fetches = {
"stats":
dict(
{
"cur_lr":
tf.cast(self.cur_lr, tf.float64),
"policy_loss":
self.loss.pi_loss,
"entropy":
self.loss.entropy,
"grad_gnorm":
tf.global_norm(self._grads),
"var_gnorm":
tf.global_norm(self.var_list),
"vf_loss":
self.loss.vf_loss,
"vf_explained_var":
explained_variance(
tf.reshape(self.loss.value_targets, [-1]),
tf.reshape(values_batched, [-1])),
}, **self.KL_stats),
}
def __init__(self,
observation_space,
action_space,
config,
existing_inputs=None):
config = dict(ray.rllib.agents.impala.impala.DEFAULT_CONFIG, **config)
assert config["batch_mode"] == "truncate_episodes", \
"Must use `truncate_episodes` batch mode with V-trace."
self.config = config
self.sess = tf.get_default_session()
self.grads = None
if isinstance(action_space, gym.spaces.Discrete):
is_multidiscrete = False
output_hidden_shape = [action_space.n]
elif isinstance(action_space, gym.spaces.multi_discrete.MultiDiscrete):
is_multidiscrete = True
output_hidden_shape = action_space.nvec.astype(np.int32)
elif self.config["vtrace"]:
raise UnsupportedSpaceException(
"Action space {} is not supported for APPO + VTrace.",
format(action_space))
else:
is_multidiscrete = False
output_hidden_shape = 1
# Policy network model
dist_class, logit_dim = ModelCatalog.get_action_dist(
action_space, self.config["model"])
# Create input placeholders
if existing_inputs:
if self.config["vtrace"]:
actions, dones, behaviour_logits, rewards, observations, \
prev_actions, prev_rewards = existing_inputs[:7]
existing_state_in = existing_inputs[7:-1]
existing_seq_lens = existing_inputs[-1]
else:
actions, dones, behaviour_logits, rewards, observations, \
prev_actions, prev_rewards, adv_ph, value_targets = \
existing_inputs[:9]
existing_state_in = existing_inputs[9:-1]
existing_seq_lens = existing_inputs[-1]
else:
actions = ModelCatalog.get_action_placeholder(action_space)
dones = tf.placeholder(tf.bool, [None], name="dones")
rewards = tf.placeholder(tf.float32, [None], name="rewards")
behaviour_logits = tf.placeholder(
tf.float32, [None, logit_dim], name="behaviour_logits")
observations = tf.placeholder(
tf.float32, [None] + list(observation_space.shape))
existing_state_in = None
existing_seq_lens = None
if not self.config["vtrace"]:
adv_ph = tf.placeholder(
tf.float32, name="advantages", shape=(None, ))
value_targets = tf.placeholder(
tf.float32, name="value_targets", shape=(None, ))
self.observations = observations
# Unpack behaviour logits
unpacked_behaviour_logits = tf.split(
behaviour_logits, output_hidden_shape, axis=1)
# Setup the policy
dist_class, logit_dim = ModelCatalog.get_action_dist(
action_space, self.config["model"])
prev_actions = ModelCatalog.get_action_placeholder(action_space)
prev_rewards = tf.placeholder(tf.float32, [None], name="prev_reward")
self.model = ModelCatalog.get_model(
{
"obs": observations,
"prev_actions": prev_actions,
"prev_rewards": prev_rewards,
"is_training": self._get_is_training_placeholder(),
},
observation_space,
action_space,
logit_dim,
self.config["model"],
state_in=existing_state_in,
seq_lens=existing_seq_lens)
unpacked_outputs = tf.split(
self.model.outputs, output_hidden_shape, axis=1)
dist_inputs = unpacked_outputs if is_multidiscrete else \
self.model.outputs
prev_dist_inputs = unpacked_behaviour_logits if is_multidiscrete else \
behaviour_logits
action_dist = dist_class(dist_inputs)
prev_action_dist = dist_class(prev_dist_inputs)
values = self.model.value_function()
self.value_function = values
self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
tf.get_variable_scope().name)
def make_time_major(tensor, drop_last=False):
"""Swaps batch and trajectory axis.
Args:
tensor: A tensor or list of tensors to reshape.
drop_last: A bool indicating whether to drop the last
trajectory item.
Returns:
res: A tensor with swapped axes or a list of tensors with
swapped axes.
"""
if isinstance(tensor, list):
return [make_time_major(t, drop_last) for t in tensor]
if self.model.state_init:
B = tf.shape(self.model.seq_lens)[0]
T = tf.shape(tensor)[0] // B
else:
# Important: chop the tensor into batches at known episode cut
# boundaries. TODO(ekl) this is kind of a hack
T = self.config["sample_batch_size"]
B = tf.shape(tensor)[0] // T
rs = tf.reshape(tensor,
tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0))
# swap B and T axes
res = tf.transpose(
rs,
[1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0]))))
if drop_last:
return res[:-1]
return res
if self.model.state_in:
max_seq_len = tf.reduce_max(self.model.seq_lens) - 1
mask = tf.sequence_mask(self.model.seq_lens, max_seq_len)
mask = tf.reshape(mask, [-1])
else:
mask = tf.ones_like(rewards)
# Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc.
if self.config["vtrace"]:
logger.info("Using V-Trace surrogate loss (vtrace=True)")
# Prepare actions for loss
loss_actions = actions if is_multidiscrete else tf.expand_dims(
actions, axis=1)
self.loss = VTraceSurrogateLoss(
actions=make_time_major(loss_actions, drop_last=True),
prev_actions_logp=make_time_major(
prev_action_dist.logp(actions), drop_last=True),
actions_logp=make_time_major(
action_dist.logp(actions), drop_last=True),
action_kl=prev_action_dist.kl(action_dist),
actions_entropy=make_time_major(
action_dist.entropy(), drop_last=True),
dones=make_time_major(dones, drop_last=True),
behaviour_logits=make_time_major(
unpacked_behaviour_logits, drop_last=True),
target_logits=make_time_major(
unpacked_outputs, drop_last=True),
discount=config["gamma"],
rewards=make_time_major(rewards, drop_last=True),
values=make_time_major(values, drop_last=True),
bootstrap_value=make_time_major(values)[-1],
valid_mask=make_time_major(mask, drop_last=True),
vf_loss_coeff=self.config["vf_loss_coeff"],
entropy_coeff=self.config["entropy_coeff"],
clip_rho_threshold=self.config["vtrace_clip_rho_threshold"],
clip_pg_rho_threshold=self.config[
"vtrace_clip_pg_rho_threshold"],
clip_param=self.config["clip_param"])
else:
logger.info("Using PPO surrogate loss (vtrace=False)")
self.loss = PPOSurrogateLoss(
prev_actions_logp=make_time_major(
prev_action_dist.logp(actions)),
actions_logp=make_time_major(action_dist.logp(actions)),
action_kl=prev_action_dist.kl(action_dist),
actions_entropy=make_time_major(action_dist.entropy()),
values=make_time_major(values),
valid_mask=make_time_major(mask),
advantages=make_time_major(adv_ph),
value_targets=make_time_major(value_targets),
vf_loss_coeff=self.config["vf_loss_coeff"],
entropy_coeff=self.config["entropy_coeff"],
clip_param=self.config["clip_param"])
# KL divergence between worker and learner logits for debugging
model_dist = MultiCategorical(unpacked_outputs)
behaviour_dist = MultiCategorical(unpacked_behaviour_logits)
kls = model_dist.kl(behaviour_dist)
if len(kls) > 1:
self.KL_stats = {}
for i, kl in enumerate(kls):
self.KL_stats.update({
"mean_KL_{}".format(i): tf.reduce_mean(kl),
"max_KL_{}".format(i): tf.reduce_max(kl),
"median_KL_{}".format(i): tf.contrib.distributions.
percentile(kl, 50.0),
})
else:
self.KL_stats = {
"mean_KL": tf.reduce_mean(kls[0]),
"max_KL": tf.reduce_max(kls[0]),
"median_KL": tf.contrib.distributions.percentile(kls[0], 50.0),
}
# Initialize TFPolicyGraph
loss_in = [
("actions", actions),
("dones", dones),
("behaviour_logits", behaviour_logits),
("rewards", rewards),
("obs", observations),
("prev_actions", prev_actions),
("prev_rewards", prev_rewards),
]
if not self.config["vtrace"]:
loss_in.append(("advantages", adv_ph))
loss_in.append(("value_targets", value_targets))
LearningRateSchedule.__init__(self, self.config["lr"],
self.config["lr_schedule"])
TFPolicyGraph.__init__(
self,
observation_space,
action_space,
self.sess,
obs_input=observations,
action_sampler=action_dist.sample(),
action_prob=action_dist.sampled_action_prob(),
loss=self.loss.total_loss,
model=self.model,
loss_inputs=loss_in,
state_inputs=self.model.state_in,
state_outputs=self.model.state_out,
prev_action_input=prev_actions,
prev_reward_input=prev_rewards,
seq_lens=self.model.seq_lens,
max_seq_len=self.config["model"]["max_seq_len"],
batch_divisibility_req=self.config["sample_batch_size"])
self.sess.run(tf.global_variables_initializer())
values_batched = make_time_major(
values, drop_last=self.config["vtrace"])
self.stats_fetches = {
"stats": dict({
"cur_lr": tf.cast(self.cur_lr, tf.float64),
"policy_loss": self.loss.pi_loss,
"entropy": self.loss.entropy,
"grad_gnorm": tf.global_norm(self._grads),
"var_gnorm": tf.global_norm(self.var_list),
"vf_loss": self.loss.vf_loss,
"vf_explained_var": explained_variance(
tf.reshape(self.loss.value_targets, [-1]),
tf.reshape(values_batched, [-1])),
}, **self.KL_stats),
}
