def custom_loss(self, policy_loss, loss_inputs):
# create a new input reader per worker
reader = JsonReader(self.options["custom_options"]["input_files"])
input_ops = reader.tf_input_ops()
# define a secondary loss by building a graph copy with weight sharing
obs = tf.cast(input_ops["obs"], tf.float32)
logits, _ = self._build_layers_v2(
{"obs": restore_original_dimensions(obs, self.obs_space)},
self.num_outputs, self.options)
# You can also add self-supervised losses easily by referencing tensors
# created during _build_layers_v2(). For example, an autoencoder-style
# loss can be added as follows:
# ae_loss = squared_diff(
# loss_inputs["obs"], Decoder(self.fcnet.last_layer))
print("FYI: You can also use these tensors: {}, ".format(loss_inputs))
# compute the IL loss
action_dist = Categorical(logits, self.options)
self.policy_loss = policy_loss
self.imitation_loss = tf.reduce_mean(
-action_dist.logp(input_ops["actions"]))
total_loss = policy_loss + self.options["custom_options"]["lambda1"]\
* policy_loss + self.options["custom_options"]["lambda2"]\
* self.imitation_loss
return total_loss
def get_log_likelihood(policy, q_model, actions, input_dict, obs_space,
action_space, config):
# Action Q network.
q_vals = _compute_q_values(policy, q_model,
input_dict[SampleBatch.CUR_OBS], obs_space,
action_space)
q_vals = q_vals[0] if isinstance(q_vals, tuple) else q_vals
action_dist = Categorical(q_vals, q_model)
return action_dist.logp(actions)
def custom_loss(self, policy_loss, loss_inputs):
# create a new input reader per worker
reader = JsonReader(
self.model_config["custom_model_config"]["input_files"])
input_ops = reader.tf_input_ops(
self.model_config["custom_model_config"].get("expert_size", 1))
# define a secondary loss by building a graph copy with weight sharing
obs = restore_original_dimensions(
tf.cast(input_ops["obs"], tf.float32), self.obs_space)
logits, _ = self.forward({"obs": obs}, [], None)
# You can also add self-supervised losses easily by referencing tensors
# created during _build_layers_v2(). For example, an autoencoder-style
# loss can be added as follows:
# ae_loss = squared_diff(
# loss_inputs["obs"], Decoder(self.fcnet.last_layer))
# print("FYI: You can also use these tensors: {}, ".format(loss_inputs))
# compute the IL loss
self.policy_loss = policy_loss
(action_scores, model_logits,
dist) = self.get_q_value_distributions(logits)
model_logits = tf.squeeze(model_logits)
action_dist = Categorical(model_logits, self.model_config)
expert_logits = tf.cast(input_ops["actions"], tf.int32)
expert_action = tf.math.argmax(expert_logits)
expert_action_one_hot = tf.one_hot(expert_action, self.num_outputs)
model_action = action_dist.deterministic_sample()
model_action_one_hot = tf.one_hot(model_action, self.num_outputs)
model_expert = model_action_one_hot * expert_action_one_hot
imitation_loss = 0
loss_type = self.model_config["custom_model_config"].get("loss", "ce")
if loss_type == "ce":
imitation_loss = tf.reduce_mean(-action_dist.logp(expert_logits))
elif loss_type == "kl":
expert_dist = Categorical(tf.one_hot(expert_logits, \
self.num_outputs), self.model_config)
imitation_loss = tf.reduce_mean(-action_dist.kl(expert_dist))
elif loss_type == "dqfd":
max_value = float("-inf")
Q_select = model_logits # TODO: difference in action_scores,dist and logits
for a in range(self.num_outputs):
max_value = tf.maximum(
Q_select[a] + 0.8 * tf.cast(model_expert[a], tf.float32),
max_value)
imitation_loss = tf.reduce_mean(
1 * (max_value - Q_select[tf.cast(expert_action, tf.int32)]))
self.imitation_loss = imitation_loss
total_loss = self.model_config["custom_model_config"]["lambda1"] * policy_loss \
+ self.model_config["custom_model_config"]["lambda2"] \
* self.imitation_loss
return total_loss
def _postprocess_helper_tf(self, obs, next_obs, actions):
with (tf.GradientTape()
if self.framework != "tf" else NullContextManager()) as tape:
# Push both observations through feature net to get both phis.
phis, _ = self.model._curiosity_feature_net({
SampleBatch.OBS: tf.concat([obs, next_obs], axis=0)
})
phi, next_phi = tf.split(phis, 2)
# Predict next phi with forward model.
predicted_next_phi = self.model._curiosity_forward_fcnet(
tf.concat(
[phi, tf_one_hot(actions, self.action_space)], axis=-1))
# Forward loss term (predicted phi', given phi and action vs
# actually observed phi').
forward_l2_norm_sqared = 0.5 * tf.reduce_sum(
tf.square(predicted_next_phi - next_phi), axis=-1)
forward_loss = tf.reduce_mean(forward_l2_norm_sqared)
# Inverse loss term (prediced action that led from phi to phi' vs
# actual action taken).
phi_cat_next_phi = tf.concat([phi, next_phi], axis=-1)
dist_inputs = self.model._curiosity_inverse_fcnet(phi_cat_next_phi)
action_dist = Categorical(dist_inputs, self.model) if \
isinstance(self.action_space, Discrete) else \
MultiCategorical(
dist_inputs, self.model, self.action_space.nvec)
# Neg log(p); p=probability of observed action given the inverse-NN
# predicted action distribution.
inverse_loss = -action_dist.logp(actions)
inverse_loss = tf.reduce_mean(inverse_loss)
# Calculate the ICM loss.
loss = (1.0 - self.beta) * inverse_loss + self.beta * forward_loss
# Step the optimizer.
if self.framework != "tf":
grads = tape.gradient(loss, self._optimizer_var_list)
grads_and_vars = [(g, v)
for g, v in zip(grads, self._optimizer_var_list)
if g is not None]
update_op = self._optimizer.apply_gradients(grads_and_vars)
else:
update_op = self._optimizer.minimize(
loss, var_list=self._optimizer_var_list)
# Return the squared l2 norm and the optimizer update op.
return forward_l2_norm_sqared, update_op
def custom_loss(self, policy_loss, loss_inputs):
# Create a new input reader per worker.
reader = JsonReader(self.model_config["custom_model_config"]["input_files"])
input_ops = reader.tf_input_ops()
# Define a secondary loss by building a graph copy with weight sharing.
obs = restore_original_dimensions(
tf.cast(input_ops["obs"], tf.float32), self.obs_space
)
logits, _ = self.forward({"obs": obs}, [], None)
# You can also add self-supervised losses easily by referencing tensors
# created during _build_layers_v2(). For example, an autoencoder-style
# loss can be added as follows:
# ae_loss = squared_diff(
# loss_inputs["obs"], Decoder(self.fcnet.last_layer))
print("FYI: You can also use these tensors: {}, ".format(loss_inputs))
# Compute the IL loss.
action_dist = Categorical(logits, self.model_config)
self.policy_loss = policy_loss
self.imitation_loss = tf.reduce_mean(-action_dist.logp(input_ops["actions"]))
return policy_loss + 10 * self.imitation_loss
def _train(self):
import tensorflow as tf
policy = self.get_policy()
steps = 0
n_episodes = 1
for _ in range(n_episodes):
env = self.env._env.rail_env
obs = self.env.reset()
num_outputs = env.action_space[0]
n_agents = env.get_num_agents()
dispatcher = CellGraphDispatcher(env)
# TODO : Update max_steps as per latest version
# https://gitlab.aicrowd.com/flatland/flatland-examples/blob/master/reinforcement_learning/multi_agent_training.py
# max_steps = int(4 * 2 * (env.height + env.width + (n_agents / n_cities))) - 1
max_steps = int(4 * 2 * (20 + env.height + env.width))
episode_steps = 0
episode_max_steps = 0
episode_num_agents = 0
episode_score = 0
episode_done_agents = 0
done = {}
done["__all__"] = False
# TODO: Support for batch update
# batch_size = 2
# logits, _ = policy.model.forward({"obs": np.vstack([obs[a],obs[a]])}, [], None)
for step in range(max_steps):
action_dict = dispatcher.step(env._elapsed_steps)
with tf.GradientTape() as tape:
imitation_loss = 0
active_agents = 0
for a in range(n_agents):
if not done.get(a) and obs.get(a) is not None:
active_agents += 1
expert_action = action_dict[a].value
input_dict = {"obs": np.expand_dims(obs[a], 0)}
input_dict['obs_flat'] = input_dict['obs']
logits, _ = policy.model.forward(
input_dict, [], None)
model_logits = tf.squeeze(logits)
expert_logits = tf.cast(expert_action, tf.int32)
action_dist = Categorical(
logits, policy.model.model_config)
imitation_loss += tf.reduce_mean(-action_dist.logp(
tf.expand_dims(expert_logits, 0)))
imitation_loss = imitation_loss / max(active_agents, 1)
gradients = tape.gradient(imitation_loss,
policy.model.trainable_variables())
self.workers.local_worker().apply_gradients(gradients)
weights = ray.put(self.workers.local_worker().get_weights())
# print(self.workers.local_worker().get_weights()['default_policy'][0][:4])
for e in self.workers.remote_workers():
e.set_weights.remote(weights)
obs, all_rewards, done, info = self.env.step(action_dict)
steps += 1
for agent, agent_info in info.items():
if agent_info["agent_done"]:
episode_done_agents += 1
if done["__all__"]:
for agent, agent_info in info.items():
if episode_max_steps == 0:
episode_max_steps = agent_info["max_episode_steps"]
episode_num_agents = agent_info["num_agents"]
episode_steps = max(episode_steps,
agent_info["agent_step"])
episode_score += agent_info["agent_score"]
print(float(episode_done_agents) / episode_num_agents)
break
norm_factor = 1.0 / (episode_max_steps * episode_num_agents)
result = {
"expert_episode_reward_mean": episode_score,
"episode_reward_mean": episode_score,
"expert_episode_completion_mean":
float(episode_done_agents) / episode_num_agents,
"expert_episode_score_normalized": episode_score * norm_factor,
"episodes_this_iter": n_episodes,
"timesteps_this_iter": steps,
}
# Code taken from _train method of trainer_template.py - TODO: Not working
# res = self.collect_metrics()
# res = {}
# res.update(
# optimizer_steps_this_iter=steps,
# episode_reward_mean=episode_score,
# info=res.get("info", {}))
# res.update(expert_scores = result)
return result