This repo contains an optimized version of flatland-rl's flatland.envs.observations.TreeObsForRailEnv. Tree-based observations allow RL models to learn much more quickly than the global observations do, but flatland's built-in TreeObsForRailEnv is kind of slow, so I wrote a faster version! This repo also contains an optimized version of https://gitlab.aicrowd.com/flatland/baselines/blob/master/utils/observation_utils.py, which flattens and normalizes the tree observations into 1D numpy arrays that can be passed to a feed-forward network.