X-Light: Cross-City Traffic Signal Control Using Transformer on Transformer as Meta Multi-Agent Reinforcement Learner

1 Introduction

An effective traffic signal control (TSC) system is the key to alleviating traffic congestion. In recent years, Reinforcement Learning (RL) has been widely used in the field of TSC. It can interact with the environment, explore, and exploit, which helps agents discover better policies without artificial priors and assumptions. Numerous studies [Wei et al., 2019a; Zheng et al., 2019; Chen et al., 2020; Wei et al., 2019b; Chu et al., 2019; Lu et al., 2023; Du et al., 2024] have demonstrated its noteworthy enhancements over conventional rulebased methods [Hunt et al., 1982; Lowrie, 1990; Roess et al., 2004; Smith et al., 2013].

This paper's extended version is at arxiv.org/abs/2404.12090, with code at github.com/jianghaoyuan1994/X-Light

Corresponding Author However, most of the existing methods are scenariospecific, meaning that the training and testing should be in the same scenario (A scenario is a simulation environment, e.g., a virtual city, with a set of intersections). When deploying on a new scenario, rebuilding the environment and re-training are needed, which involves significant costs. As a result, to the best of our knowledge, all the cities still use rule-based methods such as SCATS or SCOOT for a very fundamental reason: they can be easily reused in a new district/region/city.

This raised a critical problem that how to orchestrate the multiple intersections for various scenarios/cities with strong generalizability.

There are some solutions for single-agent settings: MetaLight [Zang et al., 2020] and GESA [Jiang et al., 2024] used one single agent to control all the intersections in a scenario and achieve transferability via gradient-based meta RL and multi-city co-training, respectively. Their drawbacks are obvious, i.e., neglecting the cooperation among the multiple intersections. Yet, learning a general Multi-Agent RL model for various scenarios is non-trivial, given various scenarios could have various road networks and traffic dynamics, rendering various environment states for each agent and collaboration patterns for multiple agents. To the best of our knowledge, only a few existing works target the same challenge: MetaVIM [Zhu et al., 2023] and MetaGAT [Lou et al., 2022]. They both utilized the Markov Decision Processes (MDPs) trajectories to help the agents learn and distinguish scenario context. However, they still display limitations, such as an unstable training process and large performance drops when encountering quite dissimilar scenarios (details in Sec. 2).

To enhance cooperation and generalizability, we will incorporate the full MDP trajectories, including the observations, rewards, and actions (o↪ a↪ r) of both the target intersection and its neighbors, into the method. Given the sequence nature of the trajectories from various MDPs, two natural questions are: Q1: Can Transformer utilize the o↪ a↪ r sequences of multiple intersections for better collaboration? Q2: Furthermore, can we Transformer learn the highlevel cross-scenario MDPs dynamics for better transferability. This forms our solution as in Figure 1: a Transformer on Transformer (TonT) model for Meta Multi-Agent Reinforcement Learners. The Lower Transformer extracts single-step trajectory information from the target intersection and its neighbor intersections to encourage cooperation. The Upper Transformer learns from the historical multi-scenario multi-modal MDPs distributions and makes actions.

The main contributions are three-fold:

• In the domain of TSC, we propose the first-ever Transformer-on-Transformer (TonT) framework for meta MARL, which solves both multi-intersection collaboration and cross-city transferability/generalizability.
• Specifically, the Lower Transformer aggregates the target and its neighbors' fine-grained o↪ a↪ r information and achieves better collaboration than other solutions such as the Graph Neural Network (GNN)-based methods [Wei et al., 2019b; Lou et al., 2022], which only aggregates traffic states o; The Upper Transformer, together with a dynamic prediction pretext task and multi-scenario co-training scheme, learns the scenarioagnostic decision process and achieves better cross-city decision. Additionally, a residual link is added before inputting to actor-critic for better decision.
• We conduct rigid experiments with various scenarios and zero-shot transfer to each, and our method is constantly the best performer. In non-transfer settings, we also achieve the most top one results.

2 RelatedWork

2.1 Meta Reinforcement Learning-based TSC

Meta RL-based methods are gaining attention in TSC since they can greatly reduce adaptation costs and have promising performance in new scenarios. GeneraLight [Zhang et al., 2020] and AttendLight [Oroojlooy et al., 2020] try to generalize TSC to different traffic flow patterns: they use GAN [Goodfellow et al., 2014] and two attention models to handle traffic diversity. However, it is not strictly adapting to a new scenario. MetaLight [Zang et al., 2020] can be seen as the pioneering work of training a TSC agent in multiple scenarios:

it uses gradient-based meta-reinforcement learning with base function reading each local data and updating local gradients, and meta-learner accumulating the memory and updating the global gradients. GESA [Jiang et al., 2024] instead proposed a general plug-in module (GPI) to make it possible to read various cities' maps without labels and further uses largescale scenarios to co-train an actor-critic-based agent. But all of them only considered single-agent settings and ignored the cooperation with others. However, to model the cooperative multi-agent under multiple various scenarios is quite technically challenging. Only a few models are proposed, yet their performance is quite unstable and limited.

MetaVIM [Zhu et al., 2023] uses context-based Meta RL to solve the generalization problem and uses intrinsic rewards to cooperate, but it is only trained in one scenario and the state does not contain neighbor information, leading to challenges in performance and adaptability when encountering an environment that is different from the training scenario. MetaGAT [Lou et al., 2022] extends MetaLight with a Graph Attention Network (GAT) to aggregate the neighbors' observations into the target's and Gated Recurrent unit (GRU) to learn the MDP dynamics. However, it has volatile fluctuations during the training process and suboptimal performance. Potential reasons are that (1) it is trained with one scenario after another, but we mix intersections from various scenarios within one batch; (2) Relying on GAT restricts them to only utilize observation information as node feature, overlooking crucial actions and rewards within the MDP. We instead use Lower Transformer to fully construct the relation among o↪ a↪ r from multiple intersections; (3) A hybrid model of combining GAT and GRU may be harder to train, and we are an elegant, Transformer-only model.

2.2 Transformers in RL and Other Fields

Offline RL unfolds the MDP process as a sequence of st↪ at↪ rt↪ st+1 etc. Thus, works such as Decision Transformer (DT) [Chen et al., 2021], PDiT [Mao et al., 2023], and Trajectory Transformer (TT) [Janner et al., 2021] use Transformer as the backbone, treating decision-making as a next-token prediction, i.e., to predict the next action at+1. TransformerLight [Wu et al., 2023] got inspired by DT and trained a TSC agent with offline-collected data. However, training with offline-collected data means a heavy workload to prepare high-quality data, and it cannot interact with the environment in real-time, let alone transfer to a new scenario. To increase transferability, PromptDT [Xu et al., 2022] proposes to add few-shot task-specific sequences as the prompt for quick adaptation for new tasks. However, offline data is still needed.

Transformer as Meta RL: TrMRL [Melo, 2022] is the pioneering work that proves Transformer can function as a meta-agent. Similar to the memory reinstatement mechanism, this agent establishes connections between the immediate working memories to construct an episodic memory across the transformer layers iteratively. Our work is the first one that applies a Transformer-based meta agent into the TSC domain; we further extend the original single-agent setting to multi-agent, with several technical challenges overcome by our final TonT solution: the Lower Transformer guides cooperation with neighbors utilizing o↪ a↪ r information, and the upper Transformer learns meta-features across tasks, showcasing remarkable generalization capabilities.

3 Methodology

This section delineates our proposed method, X-Light, a general cross-city multi-agent traffic signal control method. We begin by providing an overview of our method, followed by the introduction of the TonT Encoder module, which contains the Lower Transformer and the Upper Transformer. Finally, we will give detailed settings for training. The preliminary information is presented in the Appx. B.

3.1 Overview

As shown in Fig. 2, our method is trained using intersections across various scenarios within a batch. The i-th intersection and its neighbors Ni are selected, and their MDP trajectories (o↪ a↪ r) from time frame [t - K + 1↪ t] are sampled and fed into TonT Encoder module. It is worth noting that to allow the model to handle various intersections, we use the GPI module proposed in [Jiang et al., 2024], which maps various intersections' structure into a unified one, followed by an MLP to obtain o¶i↪Ni♦ t . More details are in the Appx. D.

As shown in Fig. 2, the TonT Encoder employs two types of transformers: the Lower Transformer and the Upper Transformer. The primary role of the Lower Transformer is to integrate the target and its neighbors' MPD information at time step t. It enhances agent collaboration compared with GNNbased collaboration by 6%-13%. The Upper Transformer utilizes historical trajectory information as context to infer the current task, thereby achieving improved transferability.

The output of the TonT Encoder is then utilized by both the Actor and Critic to output the policy ⇡ for executing the action and to estimate the state value for training.

Similar to [Wei et al., 2021], we choose five features from the observations as states: queue length, current phase, occupancy, flow, and the number of stopping cars. The action is to choose the eight pre-defined phases for the next time interval, as shown in Fig 2. At each time step t, the agent i can choose to execute action ait from available action set Ai in the next ∆t seconds. In our experiments, we set ∆t as 15 seconds. The reward r is defined as the weighted sum of queue length, wait time, delay time, and pressure. More details in Appx. D.

3.2 Lower Transformer

The Lower Transformer enhances collaboration between the target i and its n-nearest neighbors Ni, Ni = ¶N1↪ ▷ ▷ ▷ ↪ Nn♦. Existing multi-agent TSC methods only consider the states of neighboring intersections for cooperation, neglecting the complex and valuable interrelations among their observations, actions, and rewards. Considering the variations in road networks and traffic flow interconnections across diverse scenarios, one intersection's action will affect another's state and also vice versa. Thus, only relying on states is inadequate to simulate the neighbors' impacts on the target accurately. This can further result in instability within the learning process, especially during cross-scenario co-training, as we observed in MetaLight and MetaGAT's training process.

Thus, in the Lower Transformer, we utilize the full MDP features of the target and its neighbors to enhance the model's understanding of the cooperation. Since we are online RL, at time t, we can only observe immediate ot and the previous step's at-1↪ rt-1. Thus, the ith agent's MDP feature at t is:

Since original o↪ a↪ r have different dimensions, we first employ three trainable linear projections, i.e., Eo↪Ea↪Er, to map all of them to the same dimension of d. Then, the full MDP transition with neighbors concatenated is defined as:

If the number of neighbors is smaller than n, we use zeropadding and add a binary indicator embedding to all ot to indicate whether this neighbor exists.

Similar to ViT [Dosovitskiy et al., 2020], a learnable [class] token is prepended serving as the global image representation, we also prepend a trainable [decision] token qdecision 2 Rd , whose state at the output of the Lower Transformer is used as the intersection representation ct. We also add standard position embedding [Vaswani et al., 2017] to each input token to retain positional information. Thus, the input to the Lower Transformer zlower t↪i is:

where zlower t↪i ↪Elower pos 2 R(3(1+n)+1)⇥d and the attention in the Lower Transformer is in R(3(1+n)+1)⇥(3(1+n)+1) . Then, we feed zlower t↪i to the Lower transformer with 3 multi-head selfattention layers, and it outputs a capsulized intersection embedding for further decision-making, denoted as cit .

3.3 Upper Transformer

The Upper Transformer ensures that the model exhibits strong generalization in the presence of unseen intersections or scenarios. To achieve this goal, as shown in Fig. 2.(b2), we employ context-based meta RL [Duan et al., 2017] and a dynamic predictor within the Upper Transformer. ci[t-K+1:t] obtained from the Lower Transformer is first projected to dimension-d0 through a trainable projection denoted as E0 . Then, as the context-based meta RL, the Upper Transformer utilizes these embeddings from the last K time steps. Similarly, along with the positional embedding, the input to the Upper Transformer is zupper [t-K+1:t]↪i 2 RK⇥d0 :

The attention mechanism in the Upper Transformer (consisting of 3 multi-head self-attention layers) operates within the RK⇥K dimension, capturing the environmental dynamics related to the target intersection from historical features.

Each batch is a mix of intersections from various scenarios

Figure 2: Our method is co-trained with intersections' MDPs from various scenarios: (a) a GPI module unifying all the scenarios, (b) the proposed TonT Encoder, and (c) an actor-critic to make a decision. The TonT Encoder contains (b1) a Lower Transformer aggregating the o, a, and r among the target and its neighbors and (b2) an Upper Transformer learning general decisions from multi-scenario historical MDPs.

Dynamic Predictor: To enhance the agent's comprehension of the current task and the influence between intersections, we introduce a dynamic predictor between the Upper Transformer's each time-step output, enhancing the agent to learn the environment dynamics. The dynamic predictor is a pretext prediction task to conduct autoregression prediction, encouraging the Upper Transformer to capture the crossscenario dynamics. As shown in Fig. 4, this dynamic predictor can improve performance by 4.4%. Specifically, we feed the previous time step zoutput t-1↪i , concatenated with all the actions a¶i↪Ni♦ t-1 and rewards r¶i↪Ni♦ t-1 from the target and its neighbors, to predict the next zoutput t↪i :

A prediction loss based on mean squared error (MSE) is:

3.4 Actor-Critic

For the decision policy ⇡, we use the PPO method [Schulman et al., 2017], both Actor and Critic are two-layer MLPs. The actor receives the TonT Encoder's output and makes the action ati for the target intersection i.

Residual Link: However, before inputting into the actorcritic module, to avoid over-abstracting the intersection's embedding through the TonT Encoder, we directly add the observation oit of the target intersection i into the Upper Transformer's output zoutput t↪i . This ensures the embedding input into the actor has enough observational information from the target intersection.

The zoutput t↪i is the embedding after TonT Encoder's feature abstraction, which is particularly essential for scenarios with complex intercorrelations; oit instead is the direct selfobservation, which is proven as rather essential for simple scenarios, where focusing on yourself is enough. As shown in Fig.4, without residual link can lead to a 2% performance drop. Thus, our design strikes a good balance for scenarios in various complexity. Therefore, together with Eq. (8), the overall optimization objective can be formatted as:

where ↵↪ β↪ γ are tuning parameters. The Actor loss and Critic loss are the same as PPO.

3.5 Multi-scenario Co-Training

To further increase model generality, we employed the multiscenario co-train to increase data diversity. Unlike MetaGAT [Lou et al., 2022] or MetaLight [Zang et al., 2020], which utilizes the multi-scenario sequential training, i.e., within a batch, intersections are from the same scenario, and scenarios are read until previous one is all seen, we adopted multi-scenario co-training, i.e., within each batch, intersections are stochastically chosen from various scenarios. Given each intersection from each scenario has a distinct structure, we adopt a general preprocess module, GPI from GESA [Jiang et al., 2024], whose core idea is to map all various-structured intersections into a unified 4-leg one. This encourages the agent to learn generalist knowledge and also enables a more stable training process.

Lastly, the learning algorithm is shown in Algorithm 1.

Algorithm 1 X-Light training process

Input: A set of target intersections I from a set of multi-agent scenarios X; training episodes E; the number of neighbor n; the Upper Transformer input length K;

Initialize: buffer D; parameters ✓;

Output: Optimized policy ⇡✓

• 1: for episode=1, ▷ ▷ ▷ ↪ E do
• 2: clean buffer D ;;
• 3: for each scenario x do
• 4: Find nearest n neighbors N of each intersection;
• 5: for each time step t do
• 6: Get the last K transitions ¶mit-k+1 ♦xk=1↪▷▷▷↪K of each intersection i according to Eq. (1) and add these into D;
• 7: Get action ait according to Eq. (9) and take joint action ¶a1t ↪ ▷▷▷↪ anx t ♦;
• 8: Get the the next states oit+1 , and rewards rit ;
• 9: end for
• 10: end for
• 11: for each step in training steps do
• 12: sample minibatch data from D;
• 13: Get dynamic predictions ¶ˆzoutput t-k+1↪i ♦k=1↪▷▷▷↪K-1 according to Eq.(7);
• 14: Computer L according (10) and update parameter ✓;
• 15: end for
• 16: end for

4 Experiments

4.1 Datasets

We conducted experiments on the simulation of urban mobility (SUMO) as the simulator. The duration of each episode is 3600 seconds. Seven different scenarios [Ault and Sharon, 2021; Jiang et al., 2024] are employed in our co-training: five of them are reported in results, and another two scenarios (Fenglin and Nanshan) are only employed to increase the number of meta-training scenarios. To enable multi-scenario co-training, we utilize the GPI module in GESA to significantly reduce the need for manual labeling and unify the observation space and action space of all intersections. The main idea of GPI is to map all variously structured intersections into a unified 4-leg one using the relative angle and masking. For details, please refer to [Jiang et al., 2024]. Table 1 shows the properties of different scenarios. The details about these scenarios are in Appx. E.

4.2 Baselines

To comprehensively verify the effectiveness of our proposed method, we employed four types of methods for a comprehensive comparison:

Conventional methods:

• Fixed Time Control (FTC) [Roess et al., 2004] with random offset executes each phase within a loop, utilizing a pre-defined phase duration.
• MaxPressure [Kouvelas et al., 2014] is a powerful conventional method that consistently selects the phase with the highest pressure among all phases.

Multi-agent methods:

• MPLight [Chen et al., 2020] is based on a phase competition mechanism to select which phase to execute, and utilizes the concept of pressure as state and reward to coordinate multiple intersections.
• IPPO [Ault and Sharon, 2021] controls each intersection with an independent PPO agent, and the agents at each intersection have the same model architecture but different model parameters.
• rMAPPO [Yu et al., 2022a] is executed with independent PPO agents like IPPO, but it uses the overall information of all intersections to jointly optimize traffic efficiency. Furthermore, we use RNNs to introduce historical information into the agent.
• CoLight [Wei et al., 2019b] utilizes the GAT to aggregate information from neighboring intersections, enhancing cooperation between intersections.

Meta-Learning and Single-agent methods:

• MetaLight [Zang et al., 2020] using the meta-learning method to train multiple scenarios to increase the generalization of the model.
• GESA [Jiang et al., 2024] proposes a unified state and action space, and subsequently employs the GPI module for large-scale multi-scenario collaborative training, which improves performance and generality.

Meta-Learning and Multi-agent methods:

• MetaGAT [Lou et al., 2022] combines contextual metalearning based on GRU to improve generalization and GAT for cooperation. Yet, the GRU is not scenarioagnostic trained, and GAT only uses observations o.
• X-Light (ours) and X-LightGNN, which replaces the Lower Transformer with GNN [Hamilton et al., 2017], and the node feature of each intersection is obtained by concatenating o↪ a↪ r into a three times longer embedding, discarding the interrelations among them.

4.3 EvaluationMetrics

As same as most existing works [Lou et al., 2022; Oroojlooy et al., 2020; Chen et al., 2020], we use average trip time as a component of evaluation metrics, defined as the average time for each vehicle from entering the scenario to leaving the scenario. However, merely using the average trip time cannot accurately reflect the real traffic situation, e.g., in the cases of severe traffic congestion, new vehicles are prevented from entering the scenario, which leads to a relatively low average trip time though it is a severe traffic condition. Thus, we also add average delay time [Ault and Sharon, 2021] for evaluation, defined as the delay caused by signalized intersections and traffic congestion.

4.4 Results

Great Transferability When Handling A New Scenario

We use a zero-shot way to evaluate each model's performance in transferring to new scenarios. Only transferable models are selected. In Table 2, each column means that we use the other six scenarios during training and then directly transfer to this unseen scenario. Our model achieves the best transfer results in all scenarios. (1) Cooperation is needed: Unlike singleagent methods like MetaLight, cooperation enhances transferability. (2) TonT is better for Meta MARL: compared to the second-best MetaGAT, ours achieved a +7.91% improvement on average and +16.3% in Grid5⇥5, mainly because our unified Transformer on Transformer design captures collaborators' o↪ a↪ r interdependency for better local cooperation, and global cross-scenario dynamics via multi-scenario co-training, respectively. Yet, MetaGAT only focuses on the neighbors' o interdependency and lacks a scenario-agnostic training scheme. In some cases (Grid5⇥5, Ingostadt21), MetaGAT cannot beat MetaLight, which means that when cooperation is not as well designed as ours, it can even worsen the transferability when dealing with multi-scenarios. (3) Lower Transformer is better for cooperation: oursGNN flattens o↪ a↪ r, thus losing the o↪ a↪ r interdependency, further highlighting the necessity of the Lower Transformer.

Enhanced Performance in Non-transfer Setting

In this non-transfer setting, as shown in Table 3, the scenario in each column is present in training: specifically, all metalearning-based methods from MetaLight to Ours are trained with all seven scenarios, while other methods from FTC to CoLight are trained and tested on the same one scenario. Figure 3 further plots the top three methods' training process. (1)

TonT is also advantageous for non-transfer: our method achieves the best results in four scenarios, with a +6.20% gain on average and +9.96% in Grid5⇥5. The gain is relatively lower than the Transfer setting, which is reasonable given we are dedicated to transfer. Besides, we have tiny - 1.02% loss compared with the best GESA in the Grid4⇥4. A potential reason is the over-simplicity of the scenario. (2) Our multi-scenario co-training is more stable than MetaGAT: As shown in Figure 3, our training is much more stable than MetaGAT (both are multi-agent setting), due to the co-training design and strong scenario-agnostic trained Upper Transformer. GESA is the most stable given it is single-agent and also in a multi-scenario co-training scheme like ours.

Ablation Studies

The necessity of each component in model:

• w/o Lower Transformer replaces the Lower Transformer with MLP. No collaboration is formulated.
• w/o Upper Transformer replaces the Upper Transformer with GRU.

• w/o Residual Link removes the oti in the actor input.
• w/o Dynamic predictor removes the dynamic predictor in the Upper Transformer, thus without Lpred.

Figure 4 shows the trip time results, with other metrics in Appx. H. We can conclude that (1) the Lower Transformer, which promotes collaboration, is the most critical component. (2) In the Upper Transformer, the Transformer model is better than GRU; moreover, the dynamic predictor is also very necessary to ensure the Upper Transformer learns scenariodiagnostic dynamics. (3) As discussed before, the residual link can further boost performance, especially in easier scenarios such as Grid4⇥4 and Cologne8, where the direct selfobservation is more useful for decision in simple scenarios.

Numbers of co-trained scenarios: in Figure 5, we also conducted the impact of different numbers of co-training scenarios. We increase the number of co-trained scenarios from 1 to 7 and test in the two fixed scenarios (Ingolstadt21 and Arterial4⇥4). Our method consistently improves as the number of co-training scenarios increases.

We also conducted experiments on the Upper Transformer, exploring various historical lengths, with details in Appx. H.

Visualizing X-Light's Upper and Lower Transformers

This visualization offers great insights on our TonT. (1) Upper Transformer understands the scenario dynamics well: We first analyze the output embeddings from the Upper Transformer. In Figure 6, we visualize the model before and after training in five evaluation scenarios with the 3D t-SNE technique, and each scenario has a unique color. The Upper Transformer can map dynamics from the same scenario together, demonstrating the Upper Transformer's good understanding of various scenario dynamics. (2) Lower Transformer collaborates well by fully utilizing o↪ a↪ r: Furthermore, we visualize the attention weights of the Lower Transformers. In Figure 7, we visualize attention weights for two heads across three layers in the Cologne8. We find that (1) in the shallow Layer-1, attention is mainly to observations (green) of the target and its neighbors via different heads; in deeper Layer-2, more attention is to actions (blue); in deepest Layer-3, right the output being used for decision, all o↪ r↪ a, especially rewards (grey), have high attention (more cases in Appx. I); (2) This further justify the necessity of our Lower Transformer, with superiority of modeling o↪ r↪ ainterrelation over GNN-based collaborations.

Further Discussions: In Appx. A, we discuss the scalability, compatibility, assumptions for transferring, and theoretical guarantee of X-Light.

5 Conclusion

In this paper, we propose X-Light, an innovative framework that has two types of transformers with multi-scenario cotraining. This design enhances the agent's transferability and enables collaboration across multiple intersections simultaneously. We use a large number of scenarios and comparison methods to verify the effectiveness of our proposed method. The outcomes highlight not only the exceptional performance of our method but also its robust transferability.

Layer1 heads