def load_policy(filename, state_size=231, hidden_layer_size=256, seed=None):
# Training parameters
training_parameters = {
'buffer_size': int(1e5),
'batch_size': 32,
'update_every': 8,
'learning_rate': 0.5e-4,
'tau': 1e-3,
'gamma': 0.99,
'buffer_min_size': 0,
'hidden_size': hidden_layer_size,
'use_gpu': False
}
# The action space of flatland is 5 discrete actions
action_size = 5
# Create Double DQN Policy object by loading the network weights from file.
policy = DDDQNPolicy(state_size,
action_size,
Namespace(**training_parameters),
seed=seed)
policy.qnetwork_local = torch.load(filename)
return policy
def __init__(self, state_size, action_size, in_parameters=None):
print(">> MultiDecisionAgent")
super(MultiDecisionAgent, self).__init__()
self.state_size = state_size
self.action_size = action_size
self.in_parameters = in_parameters
self.memory = DummyMemory()
self.loss = 0
self.ppo_policy = PPOPolicy(state_size,
action_size,
use_replay_buffer=False,
in_parameters=in_parameters)
self.dddqn_policy = DDDQNPolicy(state_size, action_size, in_parameters)
self.policy_selector = PPOPolicy(state_size, 2)
def cartpole(use_dddqn=False):
eps = 1.0
eps_decay = 0.99
min_eps = 0.01
training_mode = True
env = gym.make("CartPole-v1")
observation_space = env.observation_space.shape[0]
action_space = env.action_space.n
if not use_dddqn:
policy = PPOPolicy(observation_space, action_space, False)
else:
policy = DDDQNPolicy(observation_space, action_space, dddqn_param)
episode = 0
checkpoint_interval = 20
scores_window = deque(maxlen=100)
writer = SummaryWriter()
while True:
episode += 1
state = env.reset()
policy.reset(env)
handle = 0
tot_reward = 0
policy.start_episode(train=training_mode)
while True:
# env.render()
policy.start_step(train=training_mode)
action = policy.act(handle, state, eps)
state_next, reward, terminal, info = env.step(action)
policy.end_step(train=training_mode)
tot_reward += reward
# reward = reward if not terminal else -reward
reward = 0 if not terminal else -1
policy.step(handle, state, action, reward, state_next, terminal)
state = np.copy(state_next)
if terminal:
break
policy.end_episode(train=training_mode)
eps = max(min_eps, eps * eps_decay)
scores_window.append(tot_reward)
if episode % checkpoint_interval == 0:
print(
'\rEpisode: {:5}\treward: {:7.3f}\t avg: {:7.3f}\t eps: {:5.3f}\t replay buffer: {}'
.format(episode, tot_reward, np.mean(scores_window), eps,
len(policy.memory)))
else:
print(
'\rEpisode: {:5}\treward: {:7.3f}\t avg: {:7.3f}\t eps: {:5.3f}\t replay buffer: {}'
.format(episode, tot_reward, np.mean(scores_window), eps,
len(policy.memory)),
end=" ")
writer.add_scalar("CartPole/value", tot_reward, episode)
writer.add_scalar("CartPole/smoothed_value", np.mean(scores_window),
episode)
writer.flush()
# This is the official formula used during evaluations
# See details in flatland.envs.schedule_generators.sparse_schedule_generator
max_steps = int(4 * 2 * (env.height + env.width + (n_agents / n_cities)))
action_count = [0] * action_size
action_dict = dict()
agent_obs = [None] * env.get_num_agents()
agent_prev_obs = [None] * env.get_num_agents()
agent_prev_action = [2] * env.get_num_agents()
update_values = False
smoothed_normalized_score = -1.0
smoothed_eval_normalized_score = -1.0
smoothed_completion = 0.0
smoothed_eval_completion = 0.0
policy = DDDQNPolicy(state_size, action_size, train_params)
def format_action_prob(action_probs):
action_probs = np.round(action_probs, 3)
actions = ["↻", "←", "↑", "→", "◼"]
buffer = ""
for action, action_prob in zip(actions, action_probs):
buffer += action + " " + "{:.3f}".format(action_prob) + " "
return buffer
def eval_policy(env, policy, n_eval_episodes, max_steps):
action_dict = dict()
def train_agent(seed, n_episodes, timed = False):
# Observation parameters
observation_radius = 10
# Exploration parameters
eps_start = 1.0
eps_end = 0.01
eps_decay = 0.997 # for 2500ts
# Setup the environment
env, max_steps, x_dim, y_dim, observation_tree_depth, _ = create_multi_agent_rail_env(seed, timed)
env.reset(True, True)
# Calculate the state size given the depth of the tree observation and the number of features
n_features_per_node = env.obs_builder.observation_dim
n_nodes = 0
for i in range(observation_tree_depth + 1):
n_nodes += np.power(4, i)
state_size = n_features_per_node * n_nodes
# The action space of flatland is 5 discrete actions
action_size = 5
action_dict = dict()
# And some variables to keep track of the progress
scores_window = deque(maxlen=100) # todo smooth when rendering instead
completion_window = deque(maxlen=100)
scores = []
completion = []
action_count = [0] * action_size
agent_obs = [None] * env.get_num_agents()
agent_prev_obs = [None] * env.get_num_agents()
agent_prev_action = [2] * env.get_num_agents()
update_values = False
# Training parameters
training_parameters = {
'buffer_size': int(1e5),
'batch_size': 32,
'update_every': 8,
'learning_rate': 0.5e-4,
'tau': 1e-3,
'gamma': 0.99,
'buffer_min_size': 0,
'hidden_size': 256,
'use_gpu': False
}
# Double Dueling DQN policy
policy = DDDQNPolicy(state_size, action_size, Namespace(**training_parameters), seed=(seed+1))
for episode_idx in range(n_episodes):
score = 0
# Reset environment
obs, info = env.reset(regenerate_rail=True, regenerate_schedule=True)
# Build agent specific observations
for agent in env.get_agent_handles():
if obs[agent]:
agent_obs[agent] = normalize_observation(obs[agent], observation_tree_depth, observation_radius=observation_radius)
agent_prev_obs[agent] = agent_obs[agent].copy()
# Run episode
for _ in range(max_steps - 1):
for agent in env.get_agent_handles():
if info['action_required'][agent]:
# If an action is required, we want to store the obs at that step as well as the action
update_values = True
action = policy.act(agent_obs[agent], eps=eps_start)
action_count[action] += 1
else:
update_values = False
action = 0
action_dict.update({agent: action})
# Environment step
next_obs, all_rewards, done, info = env.step(action_dict)
# Update replay buffer and train agent
for agent in range(env.get_num_agents()):
# Only update the values when we are done or when an action was taken and thus relevant information is present
if update_values or done[agent]:
policy.step(agent_prev_obs[agent], agent_prev_action[agent], all_rewards[agent], agent_obs[agent], done[agent])
agent_prev_obs[agent] = agent_obs[agent].copy()
agent_prev_action[agent] = action_dict[agent]
if next_obs[agent]:
agent_obs[agent] = normalize_observation(next_obs[agent], observation_tree_depth, observation_radius=observation_radius)
score += all_rewards[agent]
if done['__all__']:
break
# Epsilon decay
eps_start = max(eps_end, eps_decay * eps_start)
# Collection information about training
tasks_finished = np.sum([int(done[idx]) for idx in env.get_agent_handles()])
completion_window.append(tasks_finished / max(1, env.get_num_agents()))
scores_window.append(score / (max_steps * env.get_num_agents()))
completion.append((np.mean(completion_window)))
scores.append(np.mean(scores_window))
action_probs = action_count / np.sum(action_count)
if episode_idx % 100 == 0:
end = "\n"
action_count = [1] * action_size
else:
end = " "
print('\rTraining {} agents on {}x{}\t Episode {}\t Average Score: {:.3f}\tDones: {:.2f}%\tEpsilon: {:.2f} \t Action Probabilities: \t {}'.format(
env.get_num_agents(),
x_dim, y_dim,
episode_idx,
np.mean(scores_window),
100 * np.mean(completion_window),
eps_start,
action_probs
), end=end)
# Return trained policy.
return policy
def train_agent(train_params, train_env_params, eval_env_params, obs_params):
# Environment parameters
n_agents = train_env_params.n_agents
x_dim = train_env_params.x_dim
y_dim = train_env_params.y_dim
n_cities = train_env_params.n_cities
max_rails_between_cities = train_env_params.max_rails_between_cities
max_rails_in_city = train_env_params.max_rails_in_city
seed = train_env_params.seed
# Unique ID for this training
now = datetime.now()
training_id = now.strftime('%y%m%d%H%M%S')
# Observation parameters
observation_tree_depth = obs_params.observation_tree_depth
observation_radius = obs_params.observation_radius
observation_max_path_depth = obs_params.observation_max_path_depth
# Training parameters
eps_start = train_params.eps_start
eps_end = train_params.eps_end
eps_decay = train_params.eps_decay
n_episodes = train_params.n_episodes
checkpoint_interval = train_params.checkpoint_interval
n_eval_episodes = train_params.n_evaluation_episodes
restore_replay_buffer = train_params.restore_replay_buffer
save_replay_buffer = train_params.save_replay_buffer
# Set the seeds
random.seed(seed)
np.random.seed(seed)
# Observation builder
predictor = ShortestPathPredictorForRailEnv(observation_max_path_depth)
tree_observation = TreeObsForRailEnv(max_depth=observation_tree_depth,
predictor=predictor)
# Setup the environments
train_env = create_rail_env(train_env_params, tree_observation)
train_env.reset(regenerate_schedule=True, regenerate_rail=True)
eval_env = create_rail_env(eval_env_params, tree_observation)
eval_env.reset(regenerate_schedule=True, regenerate_rail=True)
# Setup renderer
if train_params.render:
env_renderer = RenderTool(train_env, gl="PGL")
# Calculate the state size given the depth of the tree observation and the number of features
n_features_per_node = train_env.obs_builder.observation_dim
n_nodes = sum([np.power(4, i) for i in range(observation_tree_depth + 1)])
state_size = n_features_per_node * n_nodes
# The action space of flatland is 5 discrete actions
action_size = 5
# Max number of steps per episode
# This is the official formula used during evaluations
# See details in flatland.envs.schedule_generators.sparse_schedule_generator
# max_steps = int(4 * 2 * (env.height + env.width + (n_agents / n_cities)))
max_steps = train_env._max_episode_steps
action_count = [0] * action_size
action_dict = dict()
agent_obs = [None] * n_agents
agent_prev_obs = [None] * n_agents
agent_prev_action = [2] * n_agents
update_values = [False] * n_agents
# Smoothed values used as target for hyperparameter tuning
smoothed_normalized_score = -1.0
smoothed_eval_normalized_score = -1.0
smoothed_completion = 0.0
smoothed_eval_completion = 0.0
# Double Dueling DQN policy
policy = DDDQNPolicy(state_size, action_size, train_params)
# Loads existing replay buffer
if restore_replay_buffer:
try:
policy.load_replay_buffer(restore_replay_buffer)
policy.test()
except RuntimeError as e:
print(
"\n🛑 Could't load replay buffer, were the experiences generated using the same tree depth?"
)
print(e)
exit(1)
print("\n💾 Replay buffer status: {}/{} experiences".format(
len(policy.memory.memory), train_params.buffer_size))
hdd = psutil.disk_usage('/')
if save_replay_buffer and (hdd.free / (2**30)) < 500.0:
print(
"⚠️ Careful! Saving replay buffers will quickly consume a lot of disk space. You have {:.2f}gb left."
.format(hdd.free / (2**30)))
# TensorBoard writer
writer = SummaryWriter()
writer.add_hparams(vars(train_params), {})
writer.add_hparams(vars(train_env_params), {})
writer.add_hparams(vars(obs_params), {})
training_timer = Timer()
training_timer.start()
print(
"\n🚉 Training {} trains on {}x{} grid for {} episodes, evaluating on {} episodes every {} episodes. Training id '{}'.\n"
.format(train_env.get_num_agents(), x_dim, y_dim, n_episodes,
n_eval_episodes, checkpoint_interval, training_id))
for episode_idx in range(n_episodes + 1):
step_timer = Timer()
reset_timer = Timer()
learn_timer = Timer()
preproc_timer = Timer()
inference_timer = Timer()
# Reset environment
reset_timer.start()
obs, info = train_env.reset(regenerate_rail=True,
regenerate_schedule=True)
reset_timer.end()
if train_params.render:
env_renderer.set_new_rail()
score = 0
nb_steps = 0
actions_taken = []
# Build initial agent-specific observations
for agent in train_env.get_agent_handles():
if obs[agent]:
agent_obs[agent] = normalize_observation(
obs[agent],
observation_tree_depth,
observation_radius=observation_radius)
agent_prev_obs[agent] = agent_obs[agent].copy()
# Run episode
for step in range(max_steps - 1):
inference_timer.start()
for agent in train_env.get_agent_handles():
if info['action_required'][agent]:
update_values[agent] = True
action = policy.act(agent_obs[agent], eps=eps_start)
action_count[action] += 1
actions_taken.append(action)
else:
# An action is not required if the train hasn't joined the railway network,
# if it already reached its target, or if is currently malfunctioning.
update_values[agent] = False
action = 0
action_dict.update({agent: action})
inference_timer.end()
# Environment step
step_timer.start()
next_obs, all_rewards, done, info = train_env.step(action_dict)
step_timer.end()
# Render an episode at some interval
if train_params.render and episode_idx % checkpoint_interval == 0:
env_renderer.render_env(show=True,
frames=False,
show_observations=False,
show_predictions=False)
# Update replay buffer and train agent
for agent in train_env.get_agent_handles():
if update_values[agent] or done['__all__']:
# Only learn from timesteps where somethings happened
learn_timer.start()
policy.step(agent_prev_obs[agent],
agent_prev_action[agent], all_rewards[agent],
agent_obs[agent], done[agent])
learn_timer.end()
agent_prev_obs[agent] = agent_obs[agent].copy()
agent_prev_action[agent] = action_dict[agent]
# Preprocess the new observations
if next_obs[agent]:
preproc_timer.start()
agent_obs[agent] = normalize_observation(
next_obs[agent],
observation_tree_depth,
observation_radius=observation_radius)
preproc_timer.end()
score += all_rewards[agent]
nb_steps = step
if done['__all__']:
break
# Epsilon decay
eps_start = max(eps_end, eps_decay * eps_start)
# Collect information about training
tasks_finished = sum(done[idx]
for idx in train_env.get_agent_handles())
completion = tasks_finished / max(1, train_env.get_num_agents())
normalized_score = score / (max_steps * train_env.get_num_agents())
action_probs = action_count / np.sum(action_count)
action_count = [1] * action_size
smoothing = 0.99
smoothed_normalized_score = smoothed_normalized_score * smoothing + normalized_score * (
1.0 - smoothing)
smoothed_completion = smoothed_completion * smoothing + completion * (
1.0 - smoothing)
# Print logs
if episode_idx % checkpoint_interval == 0:
torch.save(
policy.qnetwork_local, './checkpoints/' + training_id + '-' +
str(episode_idx) + '.pth')
if save_replay_buffer:
policy.save_replay_buffer('./replay_buffers/' + training_id +
'-' + str(episode_idx) + '.pkl')
if train_params.render:
env_renderer.close_window()
print('\r🚂 Episode {}'
'\t 🏆 Score: {:.3f}'
' Avg: {:.3f}'
'\t 💯 Done: {:.2f}%'
' Avg: {:.2f}%'
'\t 🎲 Epsilon: {:.3f} '
'\t 🔀 Action Probs: {}'.format(episode_idx, normalized_score,
smoothed_normalized_score,
100 * completion,
100 * smoothed_completion,
eps_start,
format_action_prob(action_probs)),
end=" ")
# Evaluate policy and log results at some interval
if episode_idx % checkpoint_interval == 0 and n_eval_episodes > 0:
scores, completions, nb_steps_eval = eval_policy(
eval_env, policy, train_params, obs_params)
writer.add_scalar("evaluation/scores_min", np.min(scores),
episode_idx)
writer.add_scalar("evaluation/scores_max", np.max(scores),
episode_idx)
writer.add_scalar("evaluation/scores_mean", np.mean(scores),
episode_idx)
writer.add_scalar("evaluation/scores_std", np.std(scores),
episode_idx)
writer.add_histogram("evaluation/scores", np.array(scores),
episode_idx)
writer.add_scalar("evaluation/completions_min",
np.min(completions), episode_idx)
writer.add_scalar("evaluation/completions_max",
np.max(completions), episode_idx)
writer.add_scalar("evaluation/completions_mean",
np.mean(completions), episode_idx)
writer.add_scalar("evaluation/completions_std",
np.std(completions), episode_idx)
writer.add_histogram("evaluation/completions",
np.array(completions), episode_idx)
writer.add_scalar("evaluation/nb_steps_min", np.min(nb_steps_eval),
episode_idx)
writer.add_scalar("evaluation/nb_steps_max", np.max(nb_steps_eval),
episode_idx)
writer.add_scalar("evaluation/nb_steps_mean",
np.mean(nb_steps_eval), episode_idx)
writer.add_scalar("evaluation/nb_steps_std", np.std(nb_steps_eval),
episode_idx)
writer.add_histogram("evaluation/nb_steps",
np.array(nb_steps_eval), episode_idx)
smoothing = 0.9
smoothed_eval_normalized_score = smoothed_eval_normalized_score * smoothing + np.mean(
scores) * (1.0 - smoothing)
smoothed_eval_completion = smoothed_eval_completion * smoothing + np.mean(
completions) * (1.0 - smoothing)
writer.add_scalar("evaluation/smoothed_score",
smoothed_eval_normalized_score, episode_idx)
writer.add_scalar("evaluation/smoothed_completion",
smoothed_eval_completion, episode_idx)
# Save logs to tensorboard
writer.add_scalar("training/score", normalized_score, episode_idx)
writer.add_scalar("training/smoothed_score", smoothed_normalized_score,
episode_idx)
writer.add_scalar("training/completion", np.mean(completion),
episode_idx)
writer.add_scalar("training/smoothed_completion",
np.mean(smoothed_completion), episode_idx)
writer.add_scalar("training/nb_steps", nb_steps, episode_idx)
writer.add_histogram("actions/distribution", np.array(actions_taken),
episode_idx)
writer.add_scalar("actions/nothing",
action_probs[RailEnvActions.DO_NOTHING], episode_idx)
writer.add_scalar("actions/left",
action_probs[RailEnvActions.MOVE_LEFT], episode_idx)
writer.add_scalar("actions/forward",
action_probs[RailEnvActions.MOVE_FORWARD],
episode_idx)
writer.add_scalar("actions/right",
action_probs[RailEnvActions.MOVE_RIGHT], episode_idx)
writer.add_scalar("actions/stop",
action_probs[RailEnvActions.STOP_MOVING],
episode_idx)
writer.add_scalar("training/epsilon", eps_start, episode_idx)
writer.add_scalar("training/buffer_size", len(policy.memory),
episode_idx)
writer.add_scalar("training/loss", policy.loss, episode_idx)
writer.add_scalar("timer/reset", reset_timer.get(), episode_idx)
writer.add_scalar("timer/step", step_timer.get(), episode_idx)
writer.add_scalar("timer/learn", learn_timer.get(), episode_idx)
writer.add_scalar("timer/preproc", preproc_timer.get(), episode_idx)
writer.add_scalar("timer/total", training_timer.get_current(),
episode_idx)
def eval_policy(env_params, checkpoint, n_eval_episodes, max_steps,
action_size, state_size, seed, render, allow_skipping,
allow_caching):
# Evaluation is faster on CPU (except if you use a really huge policy)
parameters = {'use_gpu': False}
policy = DDDQNPolicy(state_size,
action_size,
Namespace(**parameters),
evaluation_mode=True)
policy.qnetwork_local = torch.load(checkpoint)
env_params = Namespace(**env_params)
# Environment parameters
n_agents = env_params.n_agents
x_dim = env_params.x_dim
y_dim = env_params.y_dim
n_cities = env_params.n_cities
max_rails_between_cities = env_params.max_rails_between_cities
max_rails_in_city = env_params.max_rails_in_city
# Malfunction and speed profiles
# TODO pass these parameters properly from main!
malfunction_parameters = MalfunctionParameters(
malfunction_rate=1. / 2000, # Rate of malfunctions
min_duration=20, # Minimal duration
max_duration=50 # Max duration
)
# Only fast trains in Round 1
speed_profiles = {
1.: 1.0, # Fast passenger train
1. / 2.: 0.0, # Fast freight train
1. / 3.: 0.0, # Slow commuter train
1. / 4.: 0.0 # Slow freight train
}
# Observation parameters
observation_tree_depth = env_params.observation_tree_depth
observation_radius = env_params.observation_radius
observation_max_path_depth = env_params.observation_max_path_depth
# Observation builder
predictor = ShortestPathPredictorForRailEnv(observation_max_path_depth)
tree_observation = TreeObsForRailEnv(max_depth=observation_tree_depth,
predictor=predictor)
# Setup the environment
env = RailEnv(
width=x_dim,
height=y_dim,
rail_generator=sparse_rail_generator(
max_num_cities=n_cities,
grid_mode=False,
max_rails_between_cities=max_rails_between_cities,
max_rails_in_city=max_rails_in_city,
),
schedule_generator=sparse_schedule_generator(speed_profiles),
number_of_agents=n_agents,
malfunction_generator_and_process_data=malfunction_from_params(
malfunction_parameters),
obs_builder_object=tree_observation)
if render:
env_renderer = RenderTool(env, gl="PGL")
action_dict = dict()
scores = []
completions = []
nb_steps = []
inference_times = []
preproc_times = []
agent_times = []
step_times = []
for episode_idx in range(n_eval_episodes):
seed += 1
inference_timer = Timer()
preproc_timer = Timer()
agent_timer = Timer()
step_timer = Timer()
step_timer.start()
obs, info = env.reset(regenerate_rail=True,
regenerate_schedule=True,
random_seed=seed)
step_timer.end()
agent_obs = [None] * env.get_num_agents()
score = 0.0
if render:
env_renderer.set_new_rail()
final_step = 0
skipped = 0
nb_hit = 0
agent_last_obs = {}
agent_last_action = {}
for step in range(max_steps - 1):
if allow_skipping and check_if_all_blocked(env):
# FIXME why -1? bug where all agents are "done" after max_steps!
skipped = max_steps - step - 1
final_step = max_steps - 2
n_unfinished_agents = sum(not done[idx]
for idx in env.get_agent_handles())
score -= skipped * n_unfinished_agents
break
agent_timer.start()
for agent in env.get_agent_handles():
if obs[agent] and info['action_required'][agent]:
if agent in agent_last_obs and np.all(
agent_last_obs[agent] == obs[agent]):
nb_hit += 1
action = agent_last_action[agent]
else:
preproc_timer.start()
norm_obs = normalize_observation(
obs[agent],
tree_depth=observation_tree_depth,
observation_radius=observation_radius)
preproc_timer.end()
inference_timer.start()
action = policy.act(norm_obs, eps=0.0)
inference_timer.end()
action_dict.update({agent: action})
if allow_caching:
agent_last_obs[agent] = obs[agent]
agent_last_action[agent] = action
agent_timer.end()
step_timer.start()
obs, all_rewards, done, info = env.step(action_dict)
step_timer.end()
if render:
env_renderer.render_env(show=True,
frames=False,
show_observations=False,
show_predictions=False)
if step % 100 == 0:
print("{}/{}".format(step, max_steps - 1))
for agent in env.get_agent_handles():
score += all_rewards[agent]
final_step = step
if done['__all__']:
break
normalized_score = score / (max_steps * env.get_num_agents())
scores.append(normalized_score)
tasks_finished = sum(done[idx] for idx in env.get_agent_handles())
completion = tasks_finished / max(1, env.get_num_agents())
completions.append(completion)
nb_steps.append(final_step)
inference_times.append(inference_timer.get())
preproc_times.append(preproc_timer.get())
agent_times.append(agent_timer.get())
step_times.append(step_timer.get())
skipped_text = ""
if skipped > 0:
skipped_text = "\t⚡ Skipped {}".format(skipped)
hit_text = ""
if nb_hit > 0:
hit_text = "\t⚡ Hit {} ({:.1f}%)".format(nb_hit, (100 * nb_hit) /
(n_agents * final_step))
print(
"☑️ Score: {:.3f} \tDone: {:.1f}% \tNb steps: {:.3f} "
"\t🍭 Seed: {}"
"\t🚉 Env: {:.3f}s "
"\t🤖 Agent: {:.3f}s (per step: {:.3f}s) \t[preproc: {:.3f}s \tinfer: {:.3f}s]"
"{}{}".format(normalized_score, completion * 100.0, final_step,
seed, step_timer.get(), agent_timer.get(),
agent_timer.get() / final_step, preproc_timer.get(),
inference_timer.get(), skipped_text, hit_text))
return scores, completions, nb_steps, agent_times, step_times
if new:
tree_observation = TreeObsForRailEnv(max_depth=max_depth, predictor=predictor)
else:
tree_observation = TreeObsForRailEnv(max_depth=observation_tree_depth, predictor=predictor)
# Calculates state and action sizes
if new:
n_nodes = observation_tree_depth
state_size = (11 +1)* n_nodes-1
else:
n_nodes = sum([np.power(4, i) for i in range(observation_tree_depth + 1)])
state_size = tree_observation.observation_dim * n_nodes
action_size = 5
# Creates the policy. No GPU on evaluation server.
policy = DDDQNPolicy(state_size, action_size, Namespace(**{'use_gpu': False}), evaluation_mode=True)
if os.path.isfile(checkpoint):
policy.qnetwork_local = torch.load(checkpoint)
print("load checkpoint from %s"%(checkpoint))
else:
print("Checkpoint not found, using untrained policy! (path: {})".format(checkpoint))
#####################################################################
# Main evaluation loop
#####################################################################
evaluation_number = 0
while True:
evaluation_number += 1
def eval_policy(env_params, checkpoint, n_eval_episodes, max_steps, seed,
render):
# evaluation is faster on CPU, except if you have huge networks
parameters = {'use_gpu': False}
policy = DDDQNPolicy(state_size,
action_size,
Namespace(**parameters),
evaluation_mode=True)
policy.qnetwork_local = torch.load(checkpoint)
env_params = Namespace(**env_params)
# Environment parameters
n_agents = env_params.n_agents
x_dim = env_params.x_dim
y_dim = env_params.y_dim
n_cities = env_params.n_cities
max_rails_between_cities = env_params.max_rails_between_cities
max_rails_in_city = env_params.max_rails_in_city
# Observation parameters
observation_tree_depth = env_params.observation_tree_depth
observation_radius = env_params.observation_radius
observation_max_path_depth = env_params.observation_max_path_depth
# Malfunction and speed profiles
# TODO pass these parameters properly from main!
malfunction_parameters = MalfunctionParameters(
malfunction_rate=1. / 2000, # Rate of malfunctions
min_duration=20, # Minimal duration
max_duration=50 # Max duration
)
speed_profiles = {
1.: 1.0, # Fast passenger train
1. / 2.: 0.0, # Fast freight train
1. / 3.: 0.0, # Slow commuter train
1. / 4.: 0.0 # Slow freight train
}
# Observation builder
predictor = ShortestPathPredictorForRailEnv(observation_max_path_depth)
tree_observation = TreeObsForRailEnv(max_depth=observation_tree_depth,
predictor=predictor)
# Setup the environment
env = RailEnv(
width=x_dim,
height=y_dim,
rail_generator=sparse_rail_generator(
max_num_cities=n_cities,
grid_mode=False,
max_rails_between_cities=max_rails_between_cities,
max_rails_in_city=max_rails_in_city),
schedule_generator=sparse_schedule_generator(speed_profiles),
number_of_agents=n_agents,
malfunction_generator_and_process_data=malfunction_from_params(
malfunction_parameters),
obs_builder_object=tree_observation,
random_seed=seed)
env.reset(True, True)
if render:
env_renderer = RenderTool(env, gl="PGL")
action_dict = dict()
scores = []
completions = []
nb_steps = []
inference_times = []
preproc_times = []
agent_times = []
step_times = []
for episode_idx in range(n_eval_episodes):
inference_timer = Timer()
preproc_timer = Timer()
agent_timer = Timer()
step_timer = Timer()
agent_obs = [None] * env.get_num_agents()
score = 0.0
step_timer.start()
obs, info = env.reset(regenerate_rail=True, regenerate_schedule=True)
step_timer.end()
if render:
env_renderer.set_new_rail()
final_step = 0
for step in range(max_steps - 1):
agent_timer.start()
for agent in env.get_agent_handles():
if obs[agent]:
preproc_timer.start()
agent_obs[agent] = normalize_observation(
obs[agent],
tree_depth=observation_tree_depth,
observation_radius=observation_radius)
preproc_timer.end()
action = 0
if info['action_required'][agent]:
inference_timer.start()
action = policy.act(agent_obs[agent], eps=0.0)
inference_timer.end()
action_dict.update({agent: action})
agent_timer.end()
step_timer.start()
obs, all_rewards, done, info = env.step(action_dict)
step_timer.end()
if render:
env_renderer.render_env(show=True,
frames=False,
show_observations=False,
show_predictions=False)
for agent in env.get_agent_handles():
score += all_rewards[agent]
final_step = step
if done['__all__']:
break
normalized_score = score / (max_steps * env.get_num_agents())
scores.append(normalized_score)
tasks_finished = sum(done[idx] for idx in env.get_agent_handles())
completion = tasks_finished / max(1, env.get_num_agents())
completions.append(completion)
nb_steps.append(final_step)
inference_times.append(inference_timer.get())
preproc_times.append(preproc_timer.get())
agent_times.append(agent_timer.get())
step_times.append(step_timer.get())
print(
"☑️ Score: {:.3f} \tDone: {:.1f}% \tNb steps: {:.3f} "
"\t🚉 Env: {:.3f}s "
"\t🤖 Agent: {:.3f}s (per step: {:.3f}s) \t[preproc: {:.3f}s \tinfer: {:.3f}s]"
.format(normalized_score, completion * 100.0, final_step,
step_timer.get(), agent_timer.get(),
agent_timer.get() / final_step, preproc_timer.get(),
inference_timer.get()))
return scores, completions, nb_steps, agent_times, step_times
def train_agent(n_episodes):
# Environment parameters
n_agents = 1
x_dim = 25
y_dim = 25
n_cities = 4
max_rails_between_cities = 2
max_rails_in_city = 3
seed = 42
# Observation parameters
observation_tree_depth = 2
observation_radius = 10
# Exploration parameters
eps_start = 1.0
eps_end = 0.01
eps_decay = 0.997 # for 2500ts
# Set the seeds
random.seed(seed)
np.random.seed(seed)
# Observation builder
tree_observation = TreeObsForRailEnv(max_depth=observation_tree_depth)
# Setup the environment
env = RailEnv(width=x_dim,
height=y_dim,
rail_generator=sparse_rail_generator(
max_num_cities=n_cities,
seed=seed,
grid_mode=False,
max_rails_between_cities=max_rails_between_cities,
max_rails_in_city=max_rails_in_city),
schedule_generator=sparse_schedule_generator(),
number_of_agents=n_agents,
obs_builder_object=tree_observation)
env.reset(True, True)
# Calculate the state size given the depth of the tree observation and the number of features
n_features_per_node = env.obs_builder.observation_dim
n_nodes = 0
for i in range(observation_tree_depth + 1):
n_nodes += np.power(4, i)
state_size = n_features_per_node * n_nodes
# The action space of flatland is 5 discrete actions
action_size = 5
# Max number of steps per episode
# This is the official formula used during evaluations
max_steps = int(4 * 2 * (env.height + env.width + (n_agents / n_cities)))
action_dict = dict()
# And some variables to keep track of the progress
scores_window = deque(maxlen=100) # todo smooth when rendering instead
completion_window = deque(maxlen=100)
scores = []
completion = []
action_count = [0] * action_size
agent_obs = [None] * env.get_num_agents()
agent_prev_obs = [None] * env.get_num_agents()
agent_prev_action = [2] * env.get_num_agents()
update_values = False
# Training parameters
training_parameters = {
'buffer_size': int(1e5),
'batch_size': 32,
'update_every': 8,
'learning_rate': 0.5e-4,
'tau': 1e-3,
'gamma': 0.99,
'buffer_min_size': 0,
'hidden_size': 256,
'use_gpu': False
}
# Double Dueling DQN policy
policy = DDDQNPolicy(state_size, action_size,
Namespace(**training_parameters))
for episode_idx in range(n_episodes):
score = 0
# Reset environment
obs, info = env.reset(regenerate_rail=True, regenerate_schedule=True)
# Build agent specific observations
for agent in env.get_agent_handles():
if obs[agent]:
agent_obs[agent] = normalize_observation(
obs[agent],
observation_tree_depth,
observation_radius=observation_radius)
agent_prev_obs[agent] = agent_obs[agent].copy()
# Run episode
for step in range(max_steps - 1):
for agent in env.get_agent_handles():
if info['action_required'][agent]:
# If an action is required, we want to store the obs at that step as well as the action
update_values = True
action = policy.act(agent, agent_obs[agent], eps=eps_start)
action_count[action] += 1
else:
update_values = False
action = 0
action_dict.update({agent: action})
# Environment step
next_obs, all_rewards, done, info = env.step(action_dict)
# Update replay buffer and train agent
for agent in range(env.get_num_agents()):
# Only update the values when we are done or when an action was taken and thus relevant information is present
if update_values or done[agent]:
policy.step(agent, agent_prev_obs[agent],
agent_prev_action[agent], all_rewards[agent],
agent_obs[agent], done[agent])
agent_prev_obs[agent] = agent_obs[agent].copy()
agent_prev_action[agent] = action_dict[agent]
if next_obs[agent]:
agent_obs[agent] = normalize_observation(
next_obs[agent],
observation_tree_depth,
observation_radius=10)
score += all_rewards[agent]
if done['__all__']:
break
# Epsilon decay
eps_start = max(eps_end, eps_decay * eps_start)
# Collection information about training
tasks_finished = np.sum(
[int(done[idx]) for idx in env.get_agent_handles()])
completion_window.append(tasks_finished / max(1, env.get_num_agents()))
scores_window.append(score / (max_steps * env.get_num_agents()))
completion.append((np.mean(completion_window)))
scores.append(np.mean(scores_window))
action_probs = action_count / np.sum(action_count)
if episode_idx % 100 == 0:
end = "\n"
torch.save(policy.qnetwork_local,
'./checkpoints/single-' + str(episode_idx) + '.pth')
action_count = [1] * action_size
else:
end = " "
print(
'\rTraining {} agents on {}x{}\t Episode {}\t Average Score: {:.3f}\tDones: {:.2f}%\tEpsilon: {:.2f} \t Action Probabilities: \t {}'
.format(env.get_num_agents(), x_dim, y_dim, episode_idx,
np.mean(scores_window), 100 * np.mean(completion_window),
eps_start, action_probs),
end=end)
# Plot overall training progress at the end
plt.plot(scores)
plt.show()
plt.plot(completion)
plt.show()
break
print("Env Path : ", remote_client.current_env_path)
print("Env Creation Time : ", env_creation_time)
local_env = remote_client.env
nb_agents = len(local_env.agents)
max_nb_steps = local_env._max_episode_steps
tree_observation.set_env(local_env)
tree_observation.reset()
# Creates the policy. No GPU on evaluation server.
if load_policy == "DDDQN":
policy = DDDQNPolicy(state_size,
get_action_size(),
Namespace(**{'use_gpu': False}),
evaluation_mode=True)
elif load_policy == "PPO":
policy = PPOPolicy(state_size, get_action_size())
elif load_policy == "DeadLockAvoidance":
policy = DeadLockAvoidanceAgent(local_env,
get_action_size(),
enable_eps=False)
elif load_policy == "DeadLockAvoidanceWithDecision":
# inter_policy = PPOPolicy(state_size, get_action_size(), use_replay_buffer=False, in_parameters=train_params)
inter_policy = DDDQNPolicy(state_size,
get_action_size(),
Namespace(**{'use_gpu': False}),
evaluation_mode=True)
policy = DeadLockAvoidanceWithDecisionAgent(local_env, state_size,
get_action_size(),
def train_agent(env_params, train_params):
# Environment parameters
n_agents = env_params.n_agents
x_dim = env_params.x_dim
y_dim = env_params.y_dim
n_cities = env_params.n_cities
max_rails_between_cities = env_params.max_rails_between_cities
max_rails_in_city = env_params.max_rails_in_city
seed = env_params.seed
# Observation parameters
observation_tree_depth = env_params.observation_tree_depth
observation_radius = env_params.observation_radius
observation_max_path_depth = env_params.observation_max_path_depth
# Training parameters
eps_start = train_params.eps_start
eps_end = train_params.eps_end
eps_decay = train_params.eps_decay
n_episodes = train_params.n_episodes
checkpoint_interval = train_params.checkpoint_interval
n_eval_episodes = train_params.n_evaluation_episodes
# Set the seeds
random.seed(seed)
np.random.seed(seed)
# Break agents from time to time
malfunction_parameters = MalfunctionParameters(
malfunction_rate=1. / 10000, # Rate of malfunctions
min_duration=15, # Minimal duration
max_duration=50 # Max duration
)
# Observation builder
predictor = ShortestPathPredictorForRailEnv(observation_max_path_depth)
tree_observation = TreeObsForRailEnv(max_depth=observation_tree_depth, predictor=predictor)
# Fraction of train which each speed
speed_profiles = {
1.: 1.0, # Fast passenger train
1. / 2.: 0.0, # Fast freight train
1. / 3.: 0.0, # Slow commuter train
1. / 4.: 0.0 # Slow freight train
}
# Setup the environment
env = RailEnv(
width=x_dim,
height=y_dim,
rail_generator=sparse_rail_generator(
max_num_cities=n_cities,
grid_mode=False,
max_rails_between_cities=max_rails_between_cities,
max_rails_in_city=max_rails_in_city
),
schedule_generator=sparse_schedule_generator(speed_profiles),
number_of_agents=n_agents,
malfunction_generator_and_process_data=malfunction_from_params(malfunction_parameters),
obs_builder_object=tree_observation,
random_seed=seed
)
env.reset(regenerate_schedule=True, regenerate_rail=True)
# Setup renderer
if train_params.render:
env_renderer = RenderTool(env, gl="PGL")
# Calculate the state size given the depth of the tree observation and the number of features
n_features_per_node = env.obs_builder.observation_dim
n_nodes = 0
for i in range(observation_tree_depth + 1):
n_nodes += np.power(4, i)
state_size = n_features_per_node * n_nodes
# The action space of flatland is 5 discrete actions
action_size = 5
# Max number of steps per episode
# This is the official formula used during evaluations
# See details in flatland.envs.schedule_generators.sparse_schedule_generator
max_steps = int(4 * 2 * (env.height + env.width + (n_agents / n_cities)))
action_count = [0] * action_size
action_dict = dict()
agent_obs = [None] * env.get_num_agents()
agent_prev_obs = [None] * env.get_num_agents()
agent_prev_action = [2] * env.get_num_agents()
update_values = False
smoothed_normalized_score = -1.0
smoothed_eval_normalized_score = -1.0
smoothed_completion = 0.0
smoothed_eval_completion = 0.0
# Double Dueling DQN policy
policy = DDDQNPolicy(state_size, action_size, train_params)
# TensorBoard writer
writer = SummaryWriter()
writer.add_hparams(vars(train_params), {})
writer.add_hparams(vars(env_params), {})
training_timer = Timer()
training_timer.start()
print("\n🚉 Training {} trains on {}x{} grid for {} episodes, evaluating on {} episodes every {} episodes.\n"
.format(env.get_num_agents(), x_dim, y_dim, n_episodes, n_eval_episodes, checkpoint_interval))
for episode_idx in range(n_episodes + 1):
# Timers
step_timer = Timer()
reset_timer = Timer()
learn_timer = Timer()
preproc_timer = Timer()
# Reset environment
reset_timer.start()
obs, info = env.reset(regenerate_rail=True, regenerate_schedule=True)
reset_timer.end()
if train_params.render:
env_renderer.set_new_rail()
score = 0
nb_steps = 0
actions_taken = []
# Build agent specific observations
for agent in env.get_agent_handles():
if obs[agent]:
agent_obs[agent] = normalize_observation(obs[agent], observation_tree_depth, observation_radius=observation_radius)
agent_prev_obs[agent] = agent_obs[agent].copy()
# Run episode
for step in range(max_steps - 1):
for agent in env.get_agent_handles():
if info['action_required'][agent]:
# If an action is required, we want to store the obs at that step as well as the action
update_values = True
action = policy.act(agent_obs[agent], eps=eps_start)
action_count[action] += 1
actions_taken.append(action)
else:
update_values = False
action = 0
action_dict.update({agent: action})
# Environment step
step_timer.start()
next_obs, all_rewards, done, info = env.step(action_dict)
step_timer.end()
if train_params.render and episode_idx % checkpoint_interval == 0:
env_renderer.render_env(
show=True,
frames=False,
show_observations=False,
show_predictions=False
)
for agent in range(env.get_num_agents()):
# Update replay buffer and train agent
# Only update the values when we are done or when an action was taken and thus relevant information is present
if update_values or done[agent]:
learn_timer.start()
policy.step(agent_prev_obs[agent], agent_prev_action[agent], all_rewards[agent], agent_obs[agent], done[agent])
learn_timer.end()
agent_prev_obs[agent] = agent_obs[agent].copy()
agent_prev_action[agent] = action_dict[agent]
# Preprocess the new observations
if next_obs[agent]:
preproc_timer.start()
agent_obs[agent] = normalize_observation(next_obs[agent], observation_tree_depth, observation_radius=observation_radius)
preproc_timer.end()
score += all_rewards[agent]
nb_steps = step
if done['__all__']:
break
# Epsilon decay
eps_start = max(eps_end, eps_decay * eps_start)
# Collection information about training
tasks_finished = sum(done[idx] for idx in env.get_agent_handles())
completion = tasks_finished / max(1, env.get_num_agents())
normalized_score = score / (max_steps * env.get_num_agents())
action_probs = action_count / np.sum(action_count)
action_count = [1] * action_size
# Smoothed values for terminal display and for more stable hyper-parameter tuning
smoothing = 0.99
smoothed_normalized_score = smoothed_normalized_score * smoothing + normalized_score * (1.0 - smoothing)
smoothed_completion = smoothed_completion * smoothing + completion * (1.0 - smoothing)
# Print logs
if episode_idx % checkpoint_interval == 0:
torch.save(policy.qnetwork_local, './checkpoints/multi-' + str(episode_idx) + '.pth')
if train_params.render:
env_renderer.close_window()
print(
'\r🚂 Episode {}'
'\t 🏆 Score: {:.3f}'
' Avg: {:.3f}'
'\t 💯 Done: {:.2f}%'
' Avg: {:.2f}%'
'\t 🎲 Epsilon: {:.2f} '
'\t 🔀 Action Probs: {}'.format(
episode_idx,
normalized_score,
smoothed_normalized_score,
100 * completion,
100 * smoothed_completion,
eps_start,
format_action_prob(action_probs)
), end=" ")
# Evaluate policy
if episode_idx % train_params.checkpoint_interval == 0:
scores, completions, nb_steps_eval = eval_policy(env, policy, n_eval_episodes, max_steps)
writer.add_scalar("evaluation/scores_min", np.min(scores), episode_idx)
writer.add_scalar("evaluation/scores_max", np.max(scores), episode_idx)
writer.add_scalar("evaluation/scores_mean", np.mean(scores), episode_idx)
writer.add_scalar("evaluation/scores_std", np.std(scores), episode_idx)
writer.add_histogram("evaluation/scores", np.array(scores), episode_idx)
writer.add_scalar("evaluation/completions_min", np.min(completions), episode_idx)
writer.add_scalar("evaluation/completions_max", np.max(completions), episode_idx)
writer.add_scalar("evaluation/completions_mean", np.mean(completions), episode_idx)
writer.add_scalar("evaluation/completions_std", np.std(completions), episode_idx)
writer.add_histogram("evaluation/completions", np.array(completions), episode_idx)
writer.add_scalar("evaluation/nb_steps_min", np.min(nb_steps_eval), episode_idx)
writer.add_scalar("evaluation/nb_steps_max", np.max(nb_steps_eval), episode_idx)
writer.add_scalar("evaluation/nb_steps_mean", np.mean(nb_steps_eval), episode_idx)
writer.add_scalar("evaluation/nb_steps_std", np.std(nb_steps_eval), episode_idx)
writer.add_histogram("evaluation/nb_steps", np.array(nb_steps_eval), episode_idx)
smoothing = 0.9
smoothed_eval_normalized_score = smoothed_eval_normalized_score * smoothing + np.mean(scores) * (1.0 - smoothing)
smoothed_eval_completion = smoothed_eval_completion * smoothing + np.mean(completions) * (1.0 - smoothing)
writer.add_scalar("evaluation/smoothed_score", smoothed_eval_normalized_score, episode_idx)
writer.add_scalar("evaluation/smoothed_completion", smoothed_eval_completion, episode_idx)
# Save logs to tensorboard
writer.add_scalar("training/score", normalized_score, episode_idx)
writer.add_scalar("training/smoothed_score", smoothed_normalized_score, episode_idx)
writer.add_scalar("training/completion", np.mean(completion), episode_idx)
writer.add_scalar("training/smoothed_completion", np.mean(smoothed_completion), episode_idx)
writer.add_scalar("training/nb_steps", nb_steps, episode_idx)
writer.add_histogram("actions/distribution", np.array(actions_taken), episode_idx)
writer.add_scalar("actions/nothing", action_probs[RailEnvActions.DO_NOTHING], episode_idx)
writer.add_scalar("actions/left", action_probs[RailEnvActions.MOVE_LEFT], episode_idx)
writer.add_scalar("actions/forward", action_probs[RailEnvActions.MOVE_FORWARD], episode_idx)
writer.add_scalar("actions/right", action_probs[RailEnvActions.MOVE_RIGHT], episode_idx)
writer.add_scalar("actions/stop", action_probs[RailEnvActions.STOP_MOVING], episode_idx)
writer.add_scalar("training/epsilon", eps_start, episode_idx)
writer.add_scalar("training/buffer_size", len(policy.memory), episode_idx)
writer.add_scalar("training/loss", policy.loss, episode_idx)
writer.add_scalar("timer/reset", reset_timer.get(), episode_idx)
writer.add_scalar("timer/step", step_timer.get(), episode_idx)
writer.add_scalar("timer/learn", learn_timer.get(), episode_idx)
writer.add_scalar("timer/preproc", preproc_timer.get(), episode_idx)
writer.add_scalar("timer/total", training_timer.get_current(), episode_idx)
class MultiDecisionAgent(LearningPolicy):
def __init__(self, state_size, action_size, in_parameters=None):
print(">> MultiDecisionAgent")
super(MultiDecisionAgent, self).__init__()
self.state_size = state_size
self.action_size = action_size
self.in_parameters = in_parameters
self.memory = DummyMemory()
self.loss = 0
self.ppo_policy = PPOPolicy(state_size,
action_size,
use_replay_buffer=False,
in_parameters=in_parameters)
self.dddqn_policy = DDDQNPolicy(state_size, action_size, in_parameters)
self.policy_selector = PPOPolicy(state_size, 2)
def step(self, handle, state, action, reward, next_state, done):
self.ppo_policy.step(handle, state, action, reward, next_state, done)
self.dddqn_policy.step(handle, state, action, reward, next_state, done)
select = self.policy_selector.act(handle, state, 0.0)
self.policy_selector.step(handle, state, select, reward, next_state,
done)
def act(self, handle, state, eps=0.):
select = self.policy_selector.act(handle, state, eps)
if select == 0:
return self.dddqn_policy.act(handle, state, eps)
return self.policy_selector.act(handle, state, eps)
def save(self, filename):
self.ppo_policy.save(filename)
self.dddqn_policy.save(filename)
self.policy_selector.save(filename)
def load(self, filename):
self.ppo_policy.load(filename)
self.dddqn_policy.load(filename)
self.policy_selector.load(filename)
def start_step(self, train):
self.ppo_policy.start_step(train)
self.dddqn_policy.start_step(train)
self.policy_selector.start_step(train)
def end_step(self, train):
self.ppo_policy.end_step(train)
self.dddqn_policy.end_step(train)
self.policy_selector.end_step(train)
def start_episode(self, train):
self.ppo_policy.start_episode(train)
self.dddqn_policy.start_episode(train)
self.policy_selector.start_episode(train)
def end_episode(self, train):
self.ppo_policy.end_episode(train)
self.dddqn_policy.end_episode(train)
self.policy_selector.end_episode(train)
def load_replay_buffer(self, filename):
self.ppo_policy.load_replay_buffer(filename)
self.dddqn_policy.load_replay_buffer(filename)
self.policy_selector.load_replay_buffer(filename)
def test(self):
self.ppo_policy.test()
self.dddqn_policy.test()
self.policy_selector.test()
def reset(self, env: RailEnv):
self.ppo_policy.reset(env)
self.dddqn_policy.reset(env)
self.policy_selector.reset(env)
def clone(self):
multi_descision_agent = MultiDecisionAgent(self.state_size,
self.action_size,
self.in_parameters)
multi_descision_agent.ppo_policy = self.ppo_policy.clone()
multi_descision_agent.dddqn_policy = self.dddqn_policy.clone()
multi_descision_agent.policy_selector = self.policy_selector.clone()
return multi_descision_agent
