DiffLight: A Partial Rewards Conditioned Diffusion Model for Traffic Signal Control with Missing Data

Authors: hanyang  chen
yang  jiang
shengnan  guo
xiaowei  mao
youfang  lin
huaiyu  wan
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 17:12:48

Abstract: The application of reinforcement learning in traffic signal control (TSC) has been extensively researched and yielded notable achievements. However, most existing works for TSC assume that traffic data from all surrounding intersections is fully and continuously available through sensors. In real-world applications, this assumption often fails due to sensor malfunctions or data loss, making TSC with missing data a critical challenge. To meet the needs of practical applications, we introduce DiffLight, a novel conditional diffusion model for TSC under data-missing scenarios in the offline setting. Specifically, we integrate two essential sub-tasks, i.e., traffic data imputation and decision-making, by leveraging a Partial Rewards Conditioned Diffusion (PRCD) model to prevent missing rewards from interfering with the learning process. Meanwhile, to effectively capture the spatial-temporal dependencies among intersections, we design a Spatial-Temporal transFormer (STFormer) architecture. In addition, we propose a Diffusion Communication Mechanism (DCM) to promote better communication and control performance under data-missing scenarios. Extensive experiments on five datasets with various data-missing scenarios demonstrate that DiffLight is an effective controller to address TSC with missing data. The code of DiffLight is released at https://github.com/lokol5579/DiffLight-release.

Main paper: 
 introduction: 
 With the acceleration of urbanization, the surge in the number of vehicles in cities has led to increasingly severe traffic congestion and pollution problems [1] . Intersections, where traffic congestion often occurs, become a key in addressing these problems. To this end, solving the traffic signal control (TSC) problem is crucial for reducing traffic congestion in intersections by controlling traffic lights. Over the years, many approaches have been developed to tackle the TSC problem, which can be categorized into conventional approaches and reinforcement learning-based (RL-based) approaches. Conventional approaches, like Fixed-time [2] , SCOOT [3] and SCATS [4] , have been widely deployed in different cities. However, these approaches struggle to adapt to the inherent stochasticity and highly dynamic nature of real-time traffic conditions, limiting their effectiveness in responding to dynamic traffic demands.Recently, reinforcement learning (RL) is introduced into TSC to enable adaptive traffic signal control [5, 6, 7, 8, 9, 10, 11, 12, 13] . Unlike conventional approaches, RL-based approaches for TSC deploy a learnable agent at each intersection, allowing traffic signals to be adjusted dynamically based on real-time traffic conditions. However, most existing RL-based approaches for TSC assume that traffic There are several challenges that must be addressed to integrate the two sub-tasks mentioned above effectively. Firstly, in RL-based approaches for TSC, rewards which are typically vehicle queue length, are critical for the performance of controllers. However, due to the absence of traffic data, only partial rewards are available. A straightforward solution might be to fill in the missing rewards with padded values, which could confuse the imputed rewards with the actual ones, leading to a negative impact on performance. Secondly, relying solely on traffic data of the local intersection makes it challenging to capture the dynamic and spatial-temporal dependencies in the traffic network for traffic data imputation and decision-making tasks. The complexity of traffic flow often requires a broader context, as traffic data from the local intersection may not adequately reflect the behaviors and interactions occurring across the entire network. The absence of traffic data from neighboring intersections may further exacerbate this issue, hindering the ability to capture these dependencies and potentially leading to a decline in performance.To tackle these challenges, we introduce DiffLight, a novel conditional diffusion model for TSC with missing data. We propose a Partial Rewards Conditioned Diffusion (PRCD) model for both traffic data imputation and decision-making under data-missing scenarios to prevent missing rewards from interfering with the learning process. Meanwhile, to effectively capture the spatial-temporal dependencies among intersections, we design the noise model as a Spatial-Temporal transFormer (STFormer) architecture. In addition, we propose a Diffusion Communication Mechanism (DCM) to enable communication and promote the capture of spatio-temporal dependencies in the traffic network through the propagation of generated observations, facilitating better control in scenarios with missing data. Extensive experiments on five datasets with various data-missing scenarios are conducted to evaluate the effectiveness of DiffLight. The experimental results indicate that DiffLight is highly competitive for TSC with missing data. 
partially observable markov decision process: 
 We consider a partially observable Markov decision process (POMDP) in the offline setting, defined as a tuple ⟨S, A, P, R, Ω, O, γ⟩. S is the state space, and s t ∈ S denotes the state at time t. A is the set of available actions and a t ∈ A denotes the action of an agent at time t. The observation o t ∈ Ω observed by the agent is part of the state s t and can be derived from the function O(s t ). r t = R(s t ) is the immediate reward of an agent at time t. P and γ denote the transition probability function and the discount factor separately. The optimization objective is to learn a policy π for agents to maximize the expected return E st,at [R t ], where R t = t γ t r t . We formulate TSC with missing data as POMDP, and consider TSC in a traffic network with several intersections. Agents are deployed at each intersection of the traffic network. As illustrated in Figure 1 , for a fourway intersection, there are twelve traffic movements from the intersection's entrance lane l in to the departure lane l out , and four pairs of traffic movements without conflict comprise four traffic signal phases, i.e., A, B, C, D. For example, the traffic signal phase of the intersection in Figure 1 is phase-A which involves movement-2= {l 2 in → l 7 out , l 2 in → l 8 out , l 2 in → l 9 out } and movement-8 = {l 8 in → l 1 out , l 8 in → l 2 out , l 8 in → l 3 out }.In this paper, the observation o t contains the number of vehicles L num and queue length L queue in every entrance lane of the local intersection. Available actions are four phases. The immediate reward r t is defined as the sum of the queue length lin L queue . Due to the lack or error of sensors resulting in missing traffic data, o t and r t , which are derived from traffic data, could be missing in a particular missing pattern. We consider random missing and kriging missing patterns detailed in Appendix C.2. To simplify the problem, we assume that o t and r t in an intersection are missing simultaneously.Diffusion model. Diffusion models [15, 24, 25] , as a powerful generative model, provide a framework to model the data generative process as a discrete diffusion process. Diffusion models consist of two processes: the forward process and the reverse process. In this paper, the forward process is defined as q(x k |x k-1 ) := N ( 1 -β k x k-1 , β k I) by the Markov process, where β k is the variance of the noise at timestep k. We adopt DDIM sampler [24] to sample in the reverse process in order to accelerate sampling. DDIM sampler is parameterized with p θ (x k-1 |x k , x 0 ) := N (µ θ (x k , k), (σ k ) 2 I), which can be optimized by a simplified surrogate loss,EQUATIONThe reverse process begins by sampling an initial noise x K ∼ N (0, I). The estimated mean of Gaussian isµ θ (x k , k) = √ ᾱk-1 x0 + √ 1 -ᾱk-1 ϵ θ (x k , k) where x0 = 1 √ ᾱk (x k - √ 1 -ᾱk ϵ θ (x k , k)), α k = 1 -β k , ᾱk = K k=1α k and ϵ θ (x k , k) is a predictor used to estimate noise.Diffusing decision-making. Recently, many diffusion-based approaches have been proposed to address decision-making problems in RL. Among existing works, Diffuser [16] chooses to diffuse on an observation-action trajectory with returns as the condition to generate actions directly. Decision Diffuser [16] focuses solely on diffusing the observation trajectory conditioned on returns. By avoiding direct diffusion over actions, Decision Diffuser enhances performance in scenarios with discrete actions. In this paper, considering the discrete nature of actions in TSC, we focus on introducing the approach diffusing on the observation trajectory τ sampled from offline dataset D.We denote the k-step denoised output of the diffusion model as x k (τ ). The observation trajectory would be diffused to generate o t+1 . However, only diffusing on observation trajectory is not enough to make decisions. An inverse dynamics model f ϕ is adapted to generate the action a t that makes the observation transit from o t to o t+1 ,EQUATIONClassifier-free guidance. Classifier-free guidance [26] aims to learn the conditional distribution q(x(τ )|y(τ )). It learns both a conditional ϵ θ (x k (τ ), y(τ ), k) and an unconditional ϵ θ (x k (τ ), ϕ, k) for the noise. Then, the perturbed noiseϵ θ (x k (τ ), ϕ, k) + ω(ϵ θ (x k (τ ), y(τ ), k) -ϵ θ (x k (τ ), ϕ, k))can be used to generate samples, where ω is the guidance scale. 
methodology: 
 In order to effectively unify traffic data imputation and decision-making for TSC with missing data, we consider both of them as a conditional generative modeling problem via diffusion models,EQUATIONwhere p θ is a learnable model distribution to estimate the conditional data distribution of trajectory x 0 (τ ), conditioned on y(τ ). We construct our generative model according to the conditional diffusion process,EQUATIONwith conditions as,EQUATIONTo facilitate a better understanding of the symbol definitions, we categorize the trajectory segments into three types. The observed trajectory consists of data collected by sensors up to time t. The missing trajectory includes data that have not been collected during this period. The observable trajectory encompasses both the observed trajectory and the data that could potentially be collected in the future. In this context, r(τ ) is the observable reward trajectory, and τ ′ obs is the observed part of trajectory τ from the local intersection. τ ′ nei = ∪ N f nei (τ i ) is the observed observations from neighboring intersections, where τ i denotes the observation trajectory of the neighboring intersection i, N is total number of neighboring intersections, f nei (•) represents the observed observations from entrance lanes of all neighboring intersections that feed into the entrance lanes of the local intersection, as shown in Figure 1 . Due to discrete and high-frequent actions in TSC, we choose to diffuse solely on observations and utilize the inverse dynamic model to generate actions, which demonstrates a better performance proven in Appendix F.1 and F.6. We define the observation trajectory τ under data-missing scenarios as,EQUATIONwhere o t is the observation collected by sensors, ôt is the uncollected observation up to time t or a potentially collectible observation in the future, C is the length of historical observations and H is the horizon of future observations. It should be noted that we generate H-step future observations to enable effective long-horizon planning [16] .Figure 2 shows the overview of DiffLight, which consists of the Partial Rewards Conditioned Diffusion (PRCD), a noise model with a Spatial-Temporal transFormer structure (STFormer), and the Diffusion Communication Mechanism (DCM). We introduce each of them in the following sections.Under data-missing scenarios, sensors deployed to collect traffic data for rewards may malfunction or be lacking. The absence of partial rewards makes it challenging to calculate returns. Therefore, we adopt partial rewards as the condition instead of returns [16, 17, 19] , which can be expressed as,EQUATIONwhere r t is the collected reward or a potentially collectible reward in the future, and rt is the uncollected reward. There are two ways to handle the missing part of rewards: conditioning on the rewards with padded values or conditioning on the partial observable rewards directly. Padding a specific value in the missing part is a feasible way. However, padded values could confuse the imputed rewards with the actual ones, leading to a negative impact on performance proven in Section 4.3. For better decision-making, we choose to condition on only partial observable rewards. We assume that the observable part and the missing part of the trajectory are collected by real sensors and virtual sensors separately, and the distribution of traffic data collected by two kinds of sensors is independent. In this case, the distribution in Equation 3 can be factorized as,EQUATIONwhere τ obs is the observable part of the trajectory τ from the local intersection, and τ mis is the missing part. We parameterize Equation 8 as the same form of locally conditioned diffusion [27, 28] , propose partial rewards conditioned diffusion (PRCD) with classifier-free guidance [26] and introduce it into TSC with missing data. Given condition set c = {ϕ, r(τ )} and binary non-overlapping mask set m = {m mis , m obs }, PRCD assigns partial rewards to corresponding observation sub-trajectory masked by m, which can be formulated as,EQUATIONwhere εθ (x k (τ ), k, y ′ (τ ), c i ) could be expressed using classifier-free guidance,EQUATIONwhere c i ∈ c. We provide a derivation for the feasibility of PRCD in Appendix E.The noise model with a U-Net structure is widely applied in image generation [15, 24, 29, 30] , control [16, 17, 19] and other fields. However, it is hard to be applied to capture the spatial-temporal dependencies in TSC. The emergence of Transformer [31] and its applications on spatial-temporal modeling [32, 33, 34, 35] provide a promising solution to deal with it. In this section, we design a Spatial-Temporal transFormer (STFormer) structure to effectively model the spatial-temporal dependencies in TSC, which includes a data embedding layer, stacked L spatial-temporal encoder layers, and an output layer. Data embedding layer embeds different inputs into embeddings, including diffusion timestep, trajectory timestep, rewards, trajectory of the local intersection, and neighboring intersections. Spatial-Temporal Encoder layer (STE) is composed of Communication Cross-Attention module (CCA), Spatial Self-Attention module (SSA), and Temporal Self-Attention module (TSA).CCA is designed to capture the spatial-temporal dependencies between the local intersection and neighboring intersections. SSA and TSA are designed to capture the spatial dependencies and temporal dependencies at the local intersection separately. Output layer is used to convert the output of STE into the noise we desire to predict. We detail the structure of STFormer in Appendix B.2.Observations of neighboring intersections are crucial for TSC with missing data [14] . However, due to the possible absence of observations from neighboring intersections, the traffic signal could be controlled ineffectively. For instance, we assume an extreme situation where there is no available observation in both the local intersection and neighboring intersections. In this case, observation trajectories of intersections are all masked by noise, leading to difficulty in decision-making at the local intersection. Therefore, we propose a Diffusion Communication Mechanism (DCM) to disseminate observation information generated by the noise model in the reverse process among intersections. Formally, we formulate DCM as,EQUATIONThe reverse process begins by inputting original observations of neighboring intersections with missing data. During diffusing, we predict x0 (τ i ), which is the same in Section 2.3. With the help of DCM, generated observations of neighboring intersections could be spread from neighboring intersections for better decisions of the agent at the local intersection. Note that we train our model with ground-truth values collected from neighboring intersections and only use DCM in the inference process.Training process. Given an offline dataset D which consists of observation trajectories, rewards and actions, we train the reverse process p θ parameterized through the noise model ϵ θ , and the inverse dynamics model f ϕ in DiffLight with the following loss,EQUATIONFor each trajectory τ , we sample noise ϵ ∼ N (0, I) and a diffusion timestep k ∼ U(1, K), construct a masked noise array of observations x k (τ ) with τ ′ obs , and predict the noise εθ (x k (τ ), k, y ′ (τ ), r(τ ), m). Missing values in condition τ ′ nei are padded with zeros. It should be noted that we ignore the rewards condition r(τ ) with probability p and the inverse dynamics is trained with individual transitions without missing observation o or o ′ . For the training process of DiffLight, due to the inaccessibility of the ground-truth of missing data, we consider it self-supervised learning. In random missing, given a trajectory τ and conditions, we can separate the observed part into two parts and set one of them to miss. In kriging missing, we randomly mask the whole trajectories of one observed intersection. Then, we can train the noise model εθ by solving Equation 12.Inference process. DiffLight is deployed to every intersection in the traffic network in the inference process. Given rewards r(τ ), a C-step observed trajectory of local intersection τ ′ obs with missing data and trajectories of neighboring intersections τ ′ nei , the agent can impute the missing observations of local intersection, predict the observations in the future and generate next action with Equation 2, 9 and 10. In order to sample a high reward trajectory, the rewards from time t + 1 to t + H are considered in r(τ ), which is set to 1. 
experimental setup: 
 Experiment Settings We conduct our experiments on CityFlow [36] , a traffic simulator widely used in various RL-based methods. Similar to existing work in [13] , we set the phase number as four and the minimum action duration as 15 seconds.Datasets The datasets consist of two parts: offline datasets with missing data and real-world traffic flow datasets with traffic networks. We apply five real-world traffic flow datasets [9, 37] for comparison, including Hangzhou1 (D 1 HZ ), Hangzhou 2 (D 2 HZ ), Jinan 1 (D 1 JN ), Jinan 2 (D 2 JN )and Jinan 3 (D 3 JN ). Corresponding offline datasets are composed of training trajectories of three online methods. We detail the datasets in Appendix C.1. To simulate data-missing scenarios, we adopt masks of different missing rates in different missing patterns, including random missing (RM) and kriging missing (KM), to mask observations and rewards. We describe the details of two patterns and missing rates in Appendix C.2.Evaluation Metrics We use the average travel time (ATT) as the main metric for evaluation, which is widely used to evaluate the performance of TSC. It calculates the average time all the vehicles spend between entering and leaving the traffic network during simulation, which is formulated as,EQUATION)where N is the total number of vehicles entering the road network, t e i and t l i are the entering time and leaving time for the i-th vehicle respectively. The lower ATT indicates a better control performance.Compared Methods We compare our method with Behavior Cloning (BC) [38] and offline approaches, including CQL [39] , TD3+BC [40] , Decision Transformer (DT) [41] , Diffuser [16] , Decision Diffuser (DD) [17] . Similar to the existing work in [14] , we adopt store-and-forward method (SFM) [42] , a rule-based method that has generally more stable performances, to impute missing observations and rewards for these approaches. We detail these approaches and SFM in Appendix D.We train and test our method on all five datasets with different missing patterns and missing rates, and compare our method with all baselines. DiffLight performs the best on most of the datasets. We analyze experiments of different missing patterns below.Random missing. In random missing, we can see that DiffLight achieves optimal or sub-optimal performance compared with baselines in all datasets in Table 1 . Diffusion-based approaches, including Diffuser and DD, show a better performance than other baselines in most datasets. These diffusionbased approaches utilize a noise model to predict noise and generate actions or observations, which helps mitigate the disturbance caused by imputed observations and rewards during the diffusion process. Compared with diffusion-based approaches, the performance of other baselines without the diffusion process is disturbed by imputed observations and rewards more seriously. We provide the overall performance of DiffLight and baselines without missing data as well, which is detailed in Appendix F.1. In addition, we provide further experiments on the influence of unobserved locations of intersections in Appendix F.2, the limit of missing rates in Appendix F.3, and the scalability of the approach in Appendix F.4.We further evaluate the effectiveness of different parts in DiffLight with the following variants. ( We evaluate the generalization performance of DiffLight on Hangzhou 1 and Jinan 1 with different missing rates. We train our method in a specific missing rate and test it on the same dataset with other missing rates. To better compare the generalization performance among models trained in different missing rates, we formulate the relative generalization performance as,EQUATIONwhere P r is the relative generalization performance of the current missing rate, P g is the performance of the model trained in other missing rates, and P o is the performance of the model trained in the current missing rate. The results in Figure 3 demonstrate that the generalization performance of DiffLight is excellent in most situations. If we train DiffLight in a high missing rate and test it in a lower missing rate, the performance of DiffLight remains stable. In contrast, if we train DiffLight in a low missing rate and test it in a higher missing rate, the performance of DiffLight would decrease slightly. As data with a higher missing rate has more complex missing situations, making it difficult for models to handle these situations.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: See Section 4.1, Appendix B and Appendix C. 
related work: 
 Traffic Signal Control TSC approaches can be categorized into conventional and RL-based methods. Conventional approaches like Fixed-time [2] , SCOOT [3] and SCATS [4] have been widely deployed in different cities. In recent years, RL-based approaches for TSC get more attention. DQN algorithm is introduced into TSC in [5, 6, 7] for dynamic real-time control. Attention mechanisms are applied to promote inter-agent communication [9] and build universal models [10] . Max-pressure [8] and advanced-MP [13] are proposed to promote the performance of existing methods. TSC with missing data is considered with the help of the imputation model in the online setting. However, there is no existing work to solve this problem in the offline setting.There are various works applying the diffusion model to offline RL recently. Diffusion and deep q-learning are combined [18] , demonstrating the potential of diffusion in RL. The state-action trajectory is encoded in latent space [43] , enhancing credit assignment and reward propagation. In addition to methods relying on TD-learning, Diffuser [16] and Decision Diffuser [17] are proposed as planners to generate the trajectory with a conditional diffusion model. However, they are all studied under scenarios without missing data, while we model TSC with missing data.Traffic Data Imputation With the development of deep learning, RNN-based methods [44, 45, 46] show good performance for traffic data imputation. In subsequent studies, diffusion models are utilized to learn the complex distribution in traffic data [21, 22] . For TSC, store-and-forward method (SFM) [42] is proven to have a more stable performance than neural networks [14] . In this paper, we adopt SFM to impute observations and rewards for baselines.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes]Justification: We have submitted our data and code on: https://github.com/lokol5579/ DiffLight-release Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/public/ guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark). • The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/ guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc. • The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why. • At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable). • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion and limitation: 
 Conclusion In this paper, we introduce DiffLight, a novel conditional diffusion model designed for TSC in scenarios with missing data. Our approach centers on the Partial Rewards Conditioned Diffusion (PRCD) model, which addresses both traffic data imputation and decision-making in the presence of incomplete data. This model helps prevent missing rewards from disrupting the learning process. We address the challenge of capturing spatial-temporal dependencies across intersections by designing a Spatial-Temporal transFormer (STFormer) architecture as the noise model. Additionally, to enhance communication and control performance, we propose a Diffusion Communication Mechanism (DCM) that facilitates the propagation of generated observations. We conduct extensive experiments on different datasets and settings to demonstrate that DiffLight is an effective controller to address TSC with missing data.Limitation In this work, we only consider two missing patterns: random missing and kriging missing. While in the real world, the missing pattern in the traffic network is more complex. Meanwhile, we just adopt SFM which is similar to k-nearest neighbor (KNN) to impute the traffic data for baselines. In addition, our approach is conditioned on partial rewards instead of returns which could lead to the short-sightedness of agents. Future work could explore the influence on performance in more different missing patterns even mixed missing patterns, adopt more different imputation methods, and find out a more far-sighted method to control the traffic signals under data-missing scenarios.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: See Section 4.1 and Appendix B.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Our proposed method demonstrates effective abilities in TSC with missing data. It can handle different missing patterns and different missing rates when controlling traffic signals. Even in an intersection where there are no observed neighboring intersections, DiffLight can perform competitively against baselines with SFM. Moreover, the fast inference speed with fewer sampling steps and stable performance indicates that DiffLight can be deployed to achieve real-time control in the real world. However, a potential negative impact of this work is that bad decisions could lead to a collapse in the traffic network.In this section, we describe the details of hyperparameters,• We set the batch size as 64 and each sample contains the trajectory of the whole intersections in the traffic network. We train our model using Adam optimizer [47] with 2e -4 learning rate for 1.5e 5 train steps.• We train DiffLight on NVIDIA GeForce RTX A5000 for around 15 hours and test it on the same GPU.• We choose the probability p of removing the condition information to be 0.25 and guidance scale α = 1.2.• In DiffLight, we choose the length of historical observations C = 5 and the planning horizon of observation trajectory H = 3.• We use K = 100 for diffusion steps.STFormer is composed of a data embedding layer, stacked L spatial-temporal encoder layers, and an output layer. We detail each of them as follows.Data embedding layer. Spatial-temporal encoder layer. The spatial-temporal encoder layer, abbreviated as STE, is composed of Communication Cross-Attention module, Spatial Self-Attention module and Temporal Self-Attention module. We adopt the vanilla attention operator [31] in modules, represented as fAtt(Q, K, V ). The following slice notations are used to formulate attention modules. For the embedding of local intersection's trajectory ectr ∈ R T ×L×D where L is the number of entrance lanes, the t slice is e t:: ctr ∈ R L×D and the l slice is e :l: ctr ∈ R T ×D . For the embedding of neighboring intersections' partial trajectories entr ∈ R L×(T •L ′ )×D where L ′ is the number of neighboring intersections' entrance lanes taken into consideration, the l slice is e l:: ntr ∈ R (T •L ′ )×D . The Communication Cross-Attention module, abbreviated as CCA, is designed to capture the spatial-temporal dependencies between the local intersection and neighboring intersections. As illustrated in Figure 1 , e l:: ntr contains information of entrance lanes from neighboring intersections feeding into lane l. This module can be formulated as, fCCA(e :l:ctr , e l:: ntr ) := fAtt(e :l: ctr , e l:: ntr , e l:: ntr )e ′:l: ctr = fCCA(e :l: ctr , e l:: ntr ) + e :l: ctr (17) The Spatial Self-Attention module, abbreviated as SSA, and Temporal Self-Attention module, abbreviated as TSA, are designed to capture the spatial dependencies and temporal dependencies separately in the local intersection, which can be formulated as,fSSA(e ′h:: ctr ) := fAtt(e ′t:: ctr , e ′t:: ctr , e ′t:: ctr ), fTSA(e ′:l: ctr ) := fAtt(e ′:l: ctr , e ′:l: ctr , e ′:l: ctr )Therefore, the spatial-temporal encoder layer can be expressed as, fSTE(ectr, entr) := fMLP(fSSA(ectr) + fTSA(ectr)) + ectr (19) To simplify the expression, we omit the embedding of diffusion timestep, trajectory timestep and rewards in Equation 16, 17, 18 and 19. In the implementation, the embedding of diffusion timestep and trajectory timestep are added to every input of CCA, SSA and TSA, and the embedding of rewards is added to every input of SSA and TSA.Output layer. We use an MLP layer as the output layer to convert the output of the spatial-temporal encoder layers into the noise we desire to predict.We apply five real-world traffic flow datasets in three cities of different sizes including Hangzhou and Jinan:(1) Hangzhou datasets: the road network contains 16 (4×4) intersections with two traffic flow datasets, including Hangzhou 1 and Hangzhou 2 . (2) Jinan datasets: the road network contains 12 (3×4) intersections with three traffic flow datasets, including Jinan1, Jinan2 and Hangzhou 3 . All these datasets are accessible in https://traffic-signal-control.github.io/.We train AttendLight [10] , Efficient-CoLight [48] and Advanced-CoLight [13] in isolation for each dataset from scratch until convergence. Then we collect all transitions in the replay buffer for each dataset during training. We present the converged performance of three methods in Table 4 . In TSC with missing data, missing patterns have a significant impact on the performance of control. In this paper, as shown in Figure 4 , we conduct our experiments on random missing and kriging missing: (1) Random missing: traffic data collected by sensors for rewards and observations in every intersection is randomly masked with 10%∼50% probability. (2) Kriging missing: there is always no traffic data collected in certain intersections with 6.25%∼25.00% (1-intersection∼4-intersection) probability in Hangzhou In this section, we make a brief introduction to baseline approaches and the Store-and-forward method (SFM).Baseline approaches. We adopt six baseline approaches in experiments as follows:• BC is a type of imitation learning, where the agent learns to mimic the behavior of an expert demonstration. The agent is trained on a dataset of state-action pairs from the expert, and the goal is to learn a policy that can replicate the expert's actions given the same states.• CQL is an RL algorithm that aims to learn a conservative Q-function, which provides a lower bound on the true Q-values. This helps to address the issue of overestimation of Q-values, which can lead to poor performance in practice.• TD3+BC combines the TD3 algorithm with behavioral cloning. By incorporating BC into TD3, the agent can leverage expert demonstrations to accelerate learning and improve sample efficiency. TD3+BC offers the benefits of both TD3's stability and BC's ability to learn from expert demonstrations.• DT is a sequence-to-sequence model that casts reinforcement learning as a sequence modeling problem. It takes as input a sequence of past states, actions, and rewards, and it outputs a sequence of future actions that maximize the expected cumulative reward. We build DT based on the code https://github.com/kzl/decision-transformer/.• Diffuser is a diffusion-based approach for decision-making. Diffuser focuses on generating sequences of actions that lead to desirable outcomes by iteratively refining these sequences. We build Diffuser based on the code https://github.com/jannerm/diffuser.• DD is a diffusion-based approach for decision-making. DD diffuses over state trajectories and planning with an inverse dynamics model. We build DD based on the code https://github.com/ anuragajay/decision-diffuser.To avoid non-convergence caused by missing data, we train baselines on datasets without missing data and test them under data-missing scenarios with observations and rewards imputed by SFM.Store-and-forward method. Since baselines cannot handle the data-missing scenarios, we adopt a rulebased SFM to impute observations and rewards for baselines. It is proved that SFM has more stable performance compared to learning neural networks [14] . In this paper, we model current observation as:f (Vt-1, k) = Concat(∪ l i f ′ (li, k)) and f ′ (li, k) = 1 k l j o l j t-1, where V k t-1 is the intersection at time t, li ∈ V k t-1 is a lane and lj is the k's neighboring lane connected by traffic movements. We set k as 12 in this paper.To prove that the aim of partial rewards conditioned diffusion is the same as the goal in Equation 3, we assume that the observable part of the trajectory and the missing part of the trajectory are collected by real sensors and virtual sensors separately, and the distribution of traffic data collected by two kinds of sensors are independent. Thus, the distribution in Equation 3 can be factorized as follows,EQUATIONThe distribution without rewards condition p(x 0 (τ )|y ′ (τ )) can be regarded as the marginal distribution of that with rewards condition p(x 0 (τ )|y(τ )) = p(x 0 (τ )|r(τ ), y ′ (τ )),EQUATIONIn this case, we can adopt the same diffusion model with classifier-free guidance to model p θ (x 0 (τ obs )|r(τ ), y ′ (τ )) and p θ (x 0 (τmis)|y ′ (τ )).We train and test our method on all five datasets and compare our method with all baselines under no data-missing scenarios. DiffLight performs the best on over half of the datasets. Meanwhile, DD demonstrates a better performance than Diffuser, which shows that diffusing only on observations is a better choice in TSC. In previous experiments, the unobserved intersections in kriging missing are not adjacent. In this section, we study the influence of unobserved locations. We provide another mask of D 1 HZ with a missing rate of 25% in kriging missing, which contains a missing intersection where all neighboring intersections are missing. The performance of this experiment is shown in Table 6 . DiffLight still demonstrates the best performance. To further explore limits on missing proportions, we conduct experiments on the selected datasets in random missing with missing rates of 70% and 90%. In the experiment, DiffLight remains an acceptable performance at the missing rate of 70%. When the missing rate rises to 90%, the performance of DiffLight drops rapidly, which shows that the limit for the missing rate is around 70%. To further evaluate the efficacy and validate the performance of our approach, we conduct experiments on the New York dataset, which includes 48 intersections. In the experiment on the New York dataset, DiffLight achieves the best performance in most scenarios, demonstrating its ability to deal with complex traffic scenarios and control traffic signals in a larger-scale traffic network. In contrast, the performance of most baselines drops rapidly, due to the cumulative effect of errors in state imputation and decision-making at more intersections. To better clarify the core differences between DiffLight and MissLight [14] , we compare them from the following two aspects.MissLight is an online method with a state imputation model and a reward imputation model, which means that interaction with the environment is necessary. In the online setting, if the method is trained in the physical environment, safety problems must be taken into consideration. If the method is trained in a simulated environment, the difference between the physical environment and the simulated environment could affect the performance of the method to some extent when the online method is going to be employed in the physical environment. In contrast, our approach, DiffLight, is an offline method based on the diffusion model.In the offline setting, our method is trained using the collected dataset without interaction with the environment, which avoids the problems mentioned above.Model composition. MissLight is a two-stage method. In the first stage, state imputation and reward imputation models are used to fill in the missing data. In the second stage, the DQN algorithm is employed to complete the training process based on the imputed data. This approach suffers from the problem of error accumulation during the training process. However, our proposed DiffLight model, which incorporates both a diffusion model and an inverse dynamics model, can simultaneously train on missing data and collaboratively achieve traffic signal control with missing data.To better compare with MissLight, we implement the SDQN-SDQN (model-based) in [14] and replace the DQN algorithm with the CQL algorithm to adapt the offline setting. We replaced the DQN algorithm with different algorithms in the offline setting and imputed the states with the SFM model. Note that all the baselines in Section 4.2 were implemented with reference to the SDQN-SDQN (transferred) method in MissLight. To distinguish the baseline of CQL implemented in Section 4.2, the new baseline is named CQL (model-based). We provide the performance of CQL (model-based) in Table 14 and Table 15 . • The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.DiffLight achieves competitive performance compared with CQL (transferred) and CQL (model-based). The possible reason why DiffLight has better performance is that CQL (model-based) suffers from error accumulation caused by the reward imputation model while DiffLight can directly make decisions with Partial Rewards Conditioned Diffusion (PRCD). 


Decision: This paper tackles the practical application of traffic signal control under missing data, proposes a novel diffusion-based model architecture, and performs a sound evaluation that shows the proposed approach outperforms various relevant baselines and provides insights through ablation. Presentation is generally clear (though some typos and grammar errors must be fixed for the camera ready version). Moreover, as pointed by reviewers, the approach can be useful for other applications, making it relevant to the broader RL community.

The new experimental results (comparison with MissLight, results on larger-scale NY dataset, boundary conditions on missing rates...) and explanations (e.g., why kriging is more challenging) provided during the rebuttal period make the paper clearer and sounder and should be included in the camera ready version.

Reviews: review: paper_summary: paper_summary: This paper focuses on traffic signal control under the condition of missing data, presenting DiffLight, a conditional diffusion model that integrates traffic data imputation and decision-making tasks. It introduces a partial rewards conditioned diffusion method to handle missing rewards (PRCD), employs a spatial-temporal transformer for noise modeling to capture intersection dependencies (STFormer), and proposes a diffusion communication mechanism for coordinated intersection control (DCM). Experiments demonstrate the effectiveness of DiffLight over various baselines. Further, the three components of DiffLight are studied and shown to be individually effective.
strength_weakness: strength_weakness: [W1] A direct discussion or comparison to [14] would be appropriate as it addresses a substantially similar problem. While it is primarily an online method, it does have an offline stage drawing the work closer to the proposed method. However, by comparing the presented results in this paper and [14] it seems unlikely [14] would perform better.
questions: questions: **Minor comments**:

Line 50 "In addition, traffic data imputation for TSC with missing data is necessary." - somewhat repetitive

Line 121 "Diffuer"

Line 196 "in both current intersection" - in both "the" current intersection

Section 3.2 might be revised for clarity.

Consider including a requirements or setup file declaring the dependency versions used for the experiments in the paper. I had no issue reproducing some of the results, however dependencies may change substantially in the future.
limitations: limitations: Any potential negative societal impacts have been addressed.
paper_summary: paper_summary: The paper proposes a conditional diffusion framework to address the traffic signal control (TSC) problem under conditions of missing data. Utilizing Partial Rewards Conditioned Diffusion (PRCD) with classifier-free guidance, for both data imputation and decision making. The authors employ the DDIM sampling method and a spatial-temporal transformer noise predictor model to reconstruct the observation trajectory. The conditioning information is derived from partially observable rewards. A Diffusion Communication Mechanism (DCM) is introduced to exchange observation information between neighboring intersections. Finally, an inverse dynamic model is used to generate actions. Extensive experiments are conducted on datasets from multiple cities, with different missing patterns. The performance of the proposed approach is compared against multiple baseline methods.
strength_weakness: strength_weakness: 1. The selected datasets typically have few intersections and lack large-scale (e.g. thousands of intersections) datasets to validate scalability and cooperative performance.

2. The comparison benchmarks and metrics are limited, lacking some methods, such as RL-based methods and GNN-based methods. Additionally, there is a lack of comparison for certain metrics, such as queue length and throughput.

3. This paper contains many notations, formulas, and conclusions that are directly quoted without specific contextual explanations, making it confusing when reading for the first look.
questions: questions: What boundary conditions ensure the effectiveness of this method? 
Apart from assumptions such as missing data being independent of existing data, are there constraints such as limits on missing proportions and conflicts with neighboring traffic?
limitations: limitations: The authors discussed the limitations of the work in the Appendix. It is suggested that they put this part in the main body of the paper.
paper_summary: paper_summary: The paper introduces "DiffLight," a novel approach combining traffic data imputation and decision-making for Traffic Signal Control (TSC) under scenarios with missing data using a diffusion model framework. It employs partial rewards conditioned diffusion and a spatial-temporal transformer architecture to address challenges associated with incomplete data. The proposed model is tested extensively on existing datasets, demonstrating its efficacy in handling different missing data scenarios.
strength_weakness: strength_weakness: 1. There is a significant gap in the comparative analysis. The paper does not benchmark against relevant state-of-the-art methods like "MissLight" and other baseline diffusion models which are crucial for establishing the efficacy of the proposed approach.
2. In MissLight, they used $D_{HZ}$, and $D_{NY}$ but this paper used new dataset, which is not clear why.
questions: questions: 1. Can you explain why MissLight is neglected? If you have offline data, it is easy to employ an offline version of MissLight as well.
2. Why not compare on the existing dataset used in MissLight?
limitations: limitations: The paper discussed limitations.
paper_summary: paper_summary: This paper proposes a conditional diffusion model named DiffLight, which is able to unify traffic data imputation and decision-making for TSC when data is incomplete. Specifically, it proposes a partial reward conditional diffusion method to avoid the negative effects brought by the padded values of missing data. Additionally, a novel spatial-temporal transformer architecture and a diffusion communication mechanism are designed to capture spatiotemporal dependencies between intersections as well as enhance cooperative control among them. Extensive experiments on five datasets demonstrate the effectiveness of the DiffLight. Besides, the source code provided in this paper is accessible and available for use.
strength_weakness: strength_weakness: W1: The authors should further explain how the Diffusion Communication Mechanism (DCM) enables information communication among multiple intersections when data is missing at both the current intersection and neighboring intersections. This part is not clear enough.
W2: This paper only considers traffic signal control under two typical types of missing data patterns. While it is undeniable that random missing and Kriging missing are the most common patterns in real life, future research could explore traffic signal control under a wider variety of missing data patterns, including mixed missing data scenarios.
W3: There are some symbols in the paper that are not correctly defined. For example:
- At line 146, \(\tau\) is defined as the observation trajectory, but at line 152, it states \(\tau_{i}\) represents the neighboring intersection. I think here \(\tau_{i}\) should denote the observation trajectory of the neighboring intersection.
- In the experimental section, at lines 220 and 221, the evaluation metric (Average Travel Time, ATT) is defined. Is there a more objective formula definition and explanation for it?
questions: questions: Could you provide a specific example to explain the design of the DCM mechanism?
limitations: limitations: See weaknesses.

Comments: comment: We are truly grateful to the reviewer for the invaluable insights and detailed feedback. We will add the content of our discussion to the subsequent version.

To evaluate the performance of the inverse dynamics, we remove the inverse dynamics and extend the dimension of the noise model to generate both observations and actions.

**Random Missing:**

| Dataset                     | Rate   | w/o inverse dynamics | w/ inverse dynamics |
| --------------------------- | ------ | -------------------- | ------------------- |
| $\mathcal{D}_{\text{HZ}}^1$ | 50.00% | 572.61               | 303.91              |
| $\mathcal{D}_{\text{JN}}^1$ | 50.00% | 301.21               | 288.01              |

**Kriging Missing:**

| Dataset                     | Rate   | w/o inverse dynamics | w/ inverse dynamics |
| --------------------------- | ------ | -------------------- | ------------------- |
| $\mathcal{D}_{\text{HZ}}^1$ | 25.00% | 386.92               | 301.08              |
| $\mathcal{D}_{\text{JN}}^1$ | 25.00% | 395.46               | 334.12              |

In the ablation experiment of the inverse dynamics model, our approach with the inverse dynamics model achieves better performance.
We sincerely thank the reviewer for the feedback and appreciate the opportunity to clarify the question. 

We would like to point out that challenges posed by random missing and Kriging missing are inherently different. In Kriging missing, data collected at unobservable intersections is absent all the time. Therefore, decisions made by agents at unobservable intersections can only be made based on observations of neighboring intersections. In contrast, decisions can be made based on observations of current and neighboring intersections in random missing. Thus, compared with random missing scenarios, Kriging missing scenarios are more complex. Moreover, experiments in the literature [1, 2, 3, 4] on Kriging missing are often conducted with missing rates ranging from 0% to 25% while experiments in the literature [5, 6, 7] on random missing are often conducted with missing rates ranging from 0% to 50% or more. Furthermore, in practice, the performance of many baselines and the proposed method drops dramatically in Kriging missing with a missing rate of 25%, making it challenging to conduct further experiments with higher missing rates.

Reference:

[1] Wu, Yuankai, et al. "Inductive graph neural networks for spatiotemporal kriging." *Proceedings of the AAAI Conference on Artificial Intelligence*. Vol. 35. No. 5. 2021.

[2] Zheng, Chuanpan, et al. "Increase: Inductive graph representation learning for spatio-temporal kriging." *Proceedings of the ACM Web Conference 2023*. 2023.

[3] Mei, Hao, et al. "Uncertainty-aware Traffic Prediction under Missing Data." *2023 IEEE International Conference on Data Mining (ICDM)*. IEEE, 2023.

[4] Mei, Hao, et al. "Reinforcement learning approaches for traffic signal control under missing data." Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence. 2023.

[5] Tashiro, Yusuke, et al. "Csdi: Conditional score-based diffusion models for probabilistic time series imputation." *Advances in Neural Information Processing Systems* 34 (2021): 24804-24816.

[6] Liu, Mingzhe, et al. "Pristi: A conditional diffusion framework for spatiotemporal imputation." *2023 IEEE 39th International Conference on Data Engineering (ICDE)*. IEEE, 2023.

[7] Kollovieh, Marcel, et al. "Predict, refine, synthesize: Self-guiding diffusion models for probabilistic time series forecasting." *Advances in Neural Information Processing Systems* 36 (2023).
Thank you for taking the time to review and for your thoughtful feedback. We will take your suggestions into consideration in our future work.
Thank you for your response. Your suggestion is invaluable and we will follow your guidance to incorporate specific contextual explanations in the subsequent version.
Thank you for addressing my concerns! I will keep my score as is.
Thanks for the detailed responses. I'm generally fine with the results.  Thus, I will keep my score.
Thank you for the responses. I'm generally fine with the results. One question about the new results: why are the missing rates for the two experiments (random vs Kriging) totally different?

I will keep my score.
Thanks for the clarification on the comparisons to existing baselines. These new experimental results help us understand what actually helps under the missing data setting. From what I see, the offline RL method (CQL) definitely helps. I'm raising my score +1 and suggest you add the above results later in this paper.

As for the ablation study, it seems like we still miss some ablation on the Inverse Dynamics part. I understand it was originally part of the Decision Diffuser but I'm still curious to see its performance.