The results in Flatland environment are presented in Figure REF . We use modified version of the winning solution from the RL track of Flatland 2020 challenge [1]} as our main baseline, which we will refer to as 1RL. Our only modification is disabling departure schedule, as it alters the underlying MDP of the environment. As there is no global reward in the environment, we do not compare with algorithms specific for solving credit assignment issue, e.g. COMA, QMIX and QTRAN. We also found it redundant to compare with MADDPG or communication methods such as G2A or CommNet as the winning solution employs PPO-based method with both augmented critics and communication between agents. In this environment, we use three different maps with the increasing size and number of agents. That way, we can progressively increase the difficulty of the environment and see how this affects the algorithms. Specifically, we use environments with 5, 10 and 15 agents, which turn traffic on the rail network more and more dense. This increase in difficulty affects the learning of world model as it must predict next states of all agents during training. Furthermore, during execution we provide agent messages only from its neighbours for both 1RL and MAMBA, thus limiting information about the environment only to local area of the map. Nevertheless, we can see that MAMBA can successfully achieve good performance in the 5 agents (Fig. REF  a) setting only in 300k samples while 1RL still needs more samples to improve performance. Naturally, performance decrease in 10 agents setting for both algorithms, but we can see that MAMBA still substantially outperforms 1RL method in 1 million samples (Fig. REF  b). However, we can see that sample efficiency starts to equalize compared to 5 agents setting. We hypothesize that it is easier for 1RL agents to learn on raw observations with communication than for MAMBA agents to learn from more noisy discrete latent space of the model. This issue of decaying sample efficiency is further exacerbated with the increased number of agents in the next setting. In 15 agents setting (Fig. REF  c) we can see that MAMBA is less sample efficient than 1RL but still achieves comparable performance in 1.5 million samples. We hypothesize that this decrease in sample efficiency results from the increased difficulty of the environment, which necessitates more expressive world model to approximate as well as more samples to learn.
