Leverage Reinforcement Learning for building intelligent trains — Flatland NeurIPS Challenge 2020
By: Stevanoska Evgenija
Tags: Rllib, Reinforcement Learning, Sbb, Challenge, Optimization

SBB in collaboration with Deutsche Bahn (DB) and French railways (SNCF) organized another round of the flatland challenge. The goal is the same as in the previous challenge, we need to train intelligent trains that can successfully navigate the railway network, from an initial starting point to predefined target, without causing deadlocks and other delays in the network. (The blog post about last year’s challenge can be found here.) The evaluation is done on previously unseen set of environments. Additional complexity is introduced by trains malfunctions, the goal of which is to test the adaptability of the trained agents. A simple example of one environment and the movements of the trains is shown on Fig.1.

Fig.1. Visualization of the challenge. Shown are starting positions of the trains, their target stations and the railway structure.

The motivation of the challenge comes from the SBB’s need to increase its network capacity. Therefore, a mechanism for simulating and evaluating a change is needed. Self-adapting and intelligent trains are part of the plan of the challenge of increasing the railway network’s capacity.

New challenge

In this year’s challenge, the constraints for solving the environments was different. Whereas last year, when there was a set of environments that needed to be solved in 8 hours timeframe, now the evaluation is done on an infinite set of environments, and the goal is to score as many points as possible. The score of each environment is calculated based on the percentage of agents that successfully reached their targets, as well as total number of executed steps. The termination criteria for the evaluation were either low score (mean percentage of done agents in the last ten environments is lower than 25%) or 8 hours of evaluation time has passed. This means that the proposed solution needs to have two points of optimization: time and accuracy. This tested the ability of the reinforcement learning agents to generalize well in short amounts of time. The next sections give a detailed overview of our solution.

Solution

The challenge had a special Reinforcement Learning (RL) category, and this approach was preferable (and more challenging) over the algorithmic one. We competed in this category and managed to train a well performing agent. The main parts of our RL solutions are shown on Fig.2. We used the Flatland simulation environment as our training environment. Our intelligent agents represented the trains. We used the RLlib library for model implementation. The model took the state of the agent as an input and outputted the optimal action. In each episode, we collected SARSA pairs for each train, and provided them as the training set to the model (meaning there was one global model that learned the policy for all agents).

Fig.2. The basic elements of the RL algorithm

The advantages of the RL approach were manifold. Firstly, it can generalize well for new and unseen environments, and adapt to road blockages. Secondly, by modification to the reward function, we can hand-craft the desired behavior of the agents. And lastly, it is far more computationally efficient compared to a planning and algorithmic approach. At each timestamp, each agent observes just a small portion of the whole map and sends it through the model (Neural Network). Fig. 3 highlights the portion of the map visible to each agent.

Fig.3. Highlight of the portion of the map that each agents observers in its state vector

New observation (state of the agents)

The first choice we needed to make is how to represent the environment. Unlike last year where we used the pre-made observations modified with our custom features, this time we focused more on generating totally custom observation based on a graph. The first attempt was to build a cell-graph, meaning each cell in the map is represented as a node in the graph. An example of this representation is shown on Fig 4.

Fig.4. Cell graph representation of the environment

It is obvious that a lot of the cells have a degree of 2, and can be merged together, which means that this representation is not optimal. In our second attempt, we represented each switch as a node, and there is a edge between two nodes only if there is a directed path between the switches. This led us to the representation shown in Fig. 5.

Fig.5. Each node in the graph representation of the environment is a switch

But an important piece of information was missing from this representation and that is the direction of the agents.

Initially each rail in the environment is represented in the graph like (x, y, d) where x, y are the coordinates of the rail in the environment and d is the direction in which the agent will be if it is on that rail.

An edge between two nodes, (x1,y1,d1) and (x2,y2,d2) means that a train with position x1,y1 and a direction d1 can go to x2,y2 and its new direction would be d2. However, since for each rail there would be a minimum of 2 nodes but with different directions in the graph, we decided to shrink it in a such way that the nodes represent the switches and the edges are the segments between the switches, which we called segment graph. An example segment graph is shown on Fig.6. From there we generate a feature vector for each train in the environment.

Fig.6. Segment graph representation of the environment

The next choice we needed to make was how to represent each segment of the graph. The goal was to provide each agent with enough information, so it can make an intelligent and informed choice for the next action. We focused on features describing distance, deadlock and environment properties. We experimented with a lot of features, and choose the final set based on agents’ performances.

Many of the features that we included in our experiments proved to be more of noise to the model and led to reduced model performance. These features included statistical information about the environment (for example graph measurements that can be extracted from the representation of the environment, like betweenness and closeness), time tick info (the number of elapsed steps since the beginning of the episode), info about the next switch (usability and distance), info about the other agents (distance and malfunction), etc.

The most important features that were included in all of our well performing models were: information about the shortest path (distance, direction), some conflict info (for example deadlock, deadlock in the next whole segment until the next usable switch), the status of the agent and priority. Here, our most important feature is priority, which we will dive deeper into in the next section.

Priority

The concept of priority significantly improved the results in the last round, so we decided to include it in this round as well. The details can be found in the previous blog post (here). In short, each agent is assigned priority based on the conflicts it has. First, we build a conflict graph, (shown on Fig. 7). The nodes represent the agents, and there is an edge between two nodes if there is a conflict between the agents (assuming they follow their shortest path to the station). The conflicts are resolved by a graph coloring technique. Namely, in each step a node of the graph is chosen, and a priority of 1 is assigned. All of its neighbors are assigned a priority of 0.

Fig.7. Representation of the priority concept. On the right, the wait-for graph is shown, together with the color values for each node (1 or 0)

Two main asserts needs to be true for the conflict graph. First, no two neighbors can have a priority of 1, and second, there should be sanity checks implemented for an agent with priority of 1. For example, an agent whose path is blocked by another agent can’t have priority.

In this round we introduced another improvement to the priority concept. Namely, we added a different priority for ready-to-depart agents, allowing us to penalize this behavior differently in the reward function. We experimented with different strategies for assigning priority to ready-to-depart agents. The simplest idea is to limit the number of active agents on the map, meaning when more than x agents are active all ready-to-depart agents have a priority of 0. We upgraded this concept, and when calculating the priority of ready-to-depart agents, we added some logic for the free shortest path of each agent, for example information whether another agent with the same path is already on the map.

Reward structure

The rewarding is one of the most significant part of the RL algorithms. There are many examples where reward structures that make sense to humans are not appropriate for the RL algorithms; namely, the agents often learn how to leverage loopholes in the reward structure to optimize the smaller objectives and thus maximize the total reward at the end, without achieving the main goal. Think chess for example; if the reward is given for taking pieces, the agents may focus on taking the opponent’s pieces and still lose the game at the end.

We struggled with similar behavior from our agents. Upon visualization of the poor-performing agents, we identified examples when they deliberately moved to unwanted situations just so they can avoid them and get a positive reward. For example, if one of the choices on the switch has agents in the opposite direction, we award reward if the agent chooses the other option and doesn’t get into deadlock. The result was agents that deliberately got into “near deadlock” situations, so they can collect the “avoid deadlock” reward.

In our final reward setting, the only positive reward was getting into the target. We also included smaller and larger penalization for correcting unwanted behaviors. The biggest negative penalization was for getting into deadlock. With smaller penalization we correct not following priority, missing the shortest path, unwanted stopping etc.

Models and training

In this round we also experimented with different RL models, both on-policy (IMPALA, PPO, CCPPO) and off-policy models (APEX-DQN). For all the models we used the RLlib library for implementations. (https://docs.ray.io/en/master/rllib.html). RLlib offers scalability and unified API for different models, which together with easy configuration makes the experiments easier to implement. The best model, as well as the best parameters, were chosen experimentally. Our best results were achieved with APEX-DQN network.

Additionally, we adopted some of the baseline ideas, for example action masking. In each state, only a subset of actions are valid and available, due to the railway settings and the choice of switches. The action masking mechanism ensured that at each time stamp for each agent only valid actions can be chosen. Moreover, we decided to reduce the dimensionality of the action space. Namely, we removed the do nothing action which just copies the previous decision the train made.

We introduced the concept of incremental training, meaning the difficulty of the provided training environments changed over time. At the beginning, the training is done on simple environment with very few agents. The idea is for the agents to learn simple behaviors like following the shortest path to the stations. As the training progresses, so does the complexity of the training environments. The idea behind this is to allow the agents to experience more complex situations and thus learn more complex behaviors, like avoiding deadlocks, choosing less frequented paths etc.

Results and evaluation

As mentioned earlier, the agents need to solve as many episodes as possible. The complexity of the environments increases during the course of the evaluation, introducing the concept of levels. Each level consists of 10 environments to be solved with different malfunction frequencies. The level configurations can be found here.

Flatland was a really fun and exiting experience, but it was also challenging at times to find a good solution. We trained many agents with varying performances. Our best performing agent was able to get to environment 229, which means it solved environments from level 22, with 181 agents and 20 cities, which is really impressive! The total score we obtained was 181.48, and during the evaluation the average of 88.1% of the agents reached their targets successfully. The reason for termination was reached overall time limit for evaluation (which was 8 hours). Our best performing agent got us a second place in the competition in the RL category! A part of the evaluation of our winning agent is given on Figure 8.

Fig.8 Visual representation of the performance of our top agent

Although there is still a room for improvement, we are very satisfied with the score. The models showed some really intelligent behavior in terms of following priorities and avoiding deadlocks, which is another proof that a Reinforcement learning approach is suitable for this kind of tasks.