1
M ULTI -ORFT: Stable Online Reinforcement Fine-Tuning for Multi-Agent Diffusion Planning in Cooperative Driving
arXiv:2604.11734v1 [cs.RO] 13 Apr 2026
Haojie Bai, Aimin Li, Member, IEEE, Ruoyu Yao, Xiongwei Zhao, Tingting Zhang, Member, IEEE, Xing Zhang, Caixiong Li, Lin Gao, Senior Member, IEEE, and Jun Ma, Senior Member, IEEE
✦
Abstract—Closed-loop cooperative driving requires planners that generate realistic multimodal multi-agent trajectories while improving safety and traffic efficiency. Existing diffusion planners can model multimodal behaviors from demonstrations, but they often exhibit weak scene consistency and remain poorly aligned with closed-loop objectives; meanwhile, stable online post-training in reactive multi-agent environments remains difficult. We present M ULTI -ORFT, which couples scene-conditioned diffusion pre-training with stable online reinforcement post-training. In pretraining, the planner uses inter-agent self-attention, cross-attention, and AdaLN-Zero-based scene conditioning to improve scene consistency and road adherence of joint trajectories. In post-training, we formulate a twolevel MDP that exposes step-wise reverse-kernel likelihoods for online optimization, and combine dense trajectory-level rewards with variancegated group-relative policy optimization (VG-GRPO) to stabilize training. On the WOMD closed-loop benchmark, M ULTI -ORFT reduces collision rate from 2.04% to 1.89% and off-road rate from 1.68% to 1.36%, while increasing average speed from 8.36 to 8.61 m/s relative to the pre-trained planner, and it outperforms strong open-source baselines including SMART-large, SMART-tiny-CLSFT, and VBD on the primary safety and efficiency metrics. These results show that coupling scene-consistent denoising with stable online diffusion-policy optimization improves the reliability of closed-loop cooperative driving. Index Terms—Autonomous driving, diffusion policy, multi-agent planning, reinforcement learning, closed-loop simulation
1
I NTRODUCTION
Advanced Connected Autonomous Vehicles (CAVs) have catalyzed the emergence of cooperative multi-agent driving paradigms, offering new avenues to address key transportation challenges and potentially improving both safety and • •
• •
•
This work is supported in part by the Natural Science Foundation of China under Grant No. 62171160. Haojie Bai, Xiongwei Zhao, Tingting Zhang, and Lin Gao are with the school of Information Science and Technology, Harbin Institute of Technology (Shenzhen), Shenzhen 518071, China (e-mail: [email protected], [email protected]). Aimin Li is with the Middle East Technology University (METU), Ankara, 06800, Turkiye (e-mail: [email protected]). Aimin Li contributes equally to this work. Ruoyu Yao and Jun Ma are with the Robotics and Autonomous Systems Thrust, Hong Kong University of Science and Technology (Guangzhou), Guangzhou 511453, China (e-mail: [email protected], [email protected]). Xing Zhang and Caixiong Li are with the School of Computer Science and Technology, Qinghai University, Xining, 810016, China
traffic efficiency [1], [2]. Central to this paradigm, multiagent behavior modeling provides the foundation for high-level cooperative driving in complex real-world environments. Nevertheless, real traffic behavior is inherently stochastic and multimodal: even within the same scene, multiple interaction behavior can be reasonable and effective. Generating multiagent driving behaviors that are both human-like and multimodal remains an open challenge, which in turn limits the deployment potential of cooperative autonomous driving systems [3]. Recent advances in diffusion models have introduced a powerful probabilistic paradigm for generating joint distributions over agent behaviors [4], [5]. By leveraging a forward diffusion and reverse denoising process and learning from large-scale human demonstrations via imitation learning, diffusion-based planners can capture dexterous behaviors and complex multimodal distributions. However, existing methods often struggle to balance multi-agent interaction modeling with scene-conditioned modeling, leading to joint trajectories that are inconsistent with scene constraints [6], [7]. Furthermore, these methods suffer from distribution shift and objective misalignment [8]. Behavior cloning mainly fits the data distribution and lacks the ability to explicitly enforce human preferences, expectations, or constraints, making it difficult to directly optimize safety and efficiency objectives in multi-agent systems. Moreover, safety-critical events such as collisions are exceedingly rare in human driving datasets, resulting in sparse supervision for learning safe interactions. Consequently, failures are more likely to occur under closedloop execution or out-of-distribution scenarios. Collectively, these issues limit the robustness and reliability of diffusionbased multi-agent planners in real-world deployment [9]. Reinforcement learning (RL) offers a compelling alternative for enhancing pretrained planning models by coupling sampling-based exploration with reward-driven policy optimization [10], [11]. Recent studies show that RL posttraining can improve human-likeness, planning quality, and trajectory stability [12]–[14]. In this paradigm, a pretrained trajectory planner serves as the actor, sampling diverse candidate futures that are scored by a reward function and iteratively refined through RL algorithms [15], [16]. However, most existing work focuses on offline post-training,
2
which is closer in spirit to reward-augmented supervised fine-tuning [17]. Without online interaction with reactive environments, constrained rollouts and limited exploration often yield only modest closed-loop gains [10], [14]. In general, achieving satisfactory closed-loop improvements with RL post-training primarily depends on (i) online interaction and (ii) well-shaped reward. (1) online interaction aligns naturally with closed-loop execution, enabling the policy to explore interaction scenarios beyond the pretraining distribution and to correct behaviors under reward guidance, thereby mitigating performance degradation from distribution shift; (2) well-shaped rewards provide stable, finegrained optimization signals for safety and efficiency-critical objectives while remaining compatible with realistic, humanlike driving. Despite these advantages, online RL in fully closed-loop settings inevitably faces interaction uncertainty and error accumulation, imposing stringent requirements on the stability and controllability of the post-training pipeline and remaining an open challenge. To address these limitations, we propose M ULTI -ORFT, a multi-agent trajectory diffusion planner and an online RL post-training framework tailored for closed-loop planning. Multi-ORFT combines condition-enhanced pre-training with stable online post-training to build a cooperative, sceneconsistent multi-agent planner that delivers substantial safety and efficiency gains. During pre-training, to better learn the scene-conditioned joint distribution of multi-agent trajectories, we design a multi-agent denoising network building on the Diffusion Transformer [18] that appropriately balances interaction modeling and scene conditioning. Specifically, we introduce AdaLN-Zero adaptive modulation mechanism in conjunction with cross-attention. Through scene-driven feature modulation and zero-initialization, the model strengthens scene consistency and adherence of multiagent trajectories while enhancing numerical stability during training, thereby improving conditioning utilization in multiagent generation beyond cross-attention alone. During posttraining, we develop a stable online RL post-training framework to strengthen cooperative behaviors. We formulate a two-level MDP to support online optimization, explicitly coupling the denoising chain with policy–environment interaction. To counteract training instability induced by closedloop interaction, we design dense trajectory-level rewards to characterize cooperative behaviors with safety and efficiency objectives, providing stable and fine-grained learning signals. We further propose a variance-gated group-relative optimization (VG-GRPO) strategy, which adaptively gates sampled groups and switches normalization schemes based on within-group reward variance, mitigating advantage collapse and gradient instability in standard GRPO [15] and improving the robustness of online training. Extensive experiments on the WOMD benchmark demonstrate that Multi-ORFT achieves substantial gains in closedloop planning performance, particularly on core safety and efficiency metrics. Further ablations show that, relative to supervised fine-tuning, preference optimization, and offline RL post-training, our online framework yields more consistent improvements in closed-loop safety and traffic efficiency. Our main contributions are summarized as follows: •
We propose a scene condition-enhanced multi-agent
•
•
trajectory diffusion model that integrates AdaLNZero adaptive modulation mechanism with crossattention in a cooperative design, simultaneously modeling inter-agent interactions and improving scene consistency and road adherence of multi-agent trajectories. We develop a stable online RL post-training pipeline to optimize multi-agent cooperative behaviors. To address instability from closed-loop interaction, we jointly design dense trajectory-level rewards and a variance-gated group-relative optimization (VGGRPO) strategy that mitigates advantage collapse and gradient instability, thereby significantly improving online training gains. Extensive experiments demonstrate Multi-ORFT’s superior closed-loop planning performance. The proposed method produces highly interactive and efficient cooperative trajectories that outperform imitation learning and other post-training methods, showing strong robustness and practical potential.
2
R ELATED W ORK
2.1
Multi-agent Behavior Modeling in Traffic Scenarios
Modeling the joint behaviors of multiple interacting agents is essential for advancing autonomous system, yet it remains challenging due to the inherent multimodality of vehicle trajectories [3], [19]. Typical distributional regression methods typically fit parametric continuous distributions such as Gaussian [20] , Laplace [21] to obtain a compact multimodal representation. Other studies incorporate goal anchors [22] or learnable intention queries [20], [23] into the decoding process to generate multi-modal future motions. However, these designs often increase model complexity and memory usage, limiting scalability. Recently, generative models, notably autoregressive Transformer and diffusion models, have advanced multiagent planning and simulation [11], [24], [25]. Autoregressive Transformer models cast multi-agent motion modeling as a next-token prediction task [26], [27], with methods such as SMART [24] learning categorical distributions over discrete motion tokens, and MotionLM [25] enabling weighted mode identification via pairwise sampling and a simple rollout aggregation. However, their discrete, sequential decoding can hinder temporal coherence and remain limited to partially joint dependencies. By contrast, diffusion models offer a compelling alternative, as they excel at modeling complex multimodal distributions while producing temporally consistent trajectories. They have been applied to decision-making, trajectory planning, and traffic simulation [28]–[30]. For example, CTG++ [31] employs a spatiotemporal Transformer that captures the evolving dynamics of multi-agent interactions, and VBD [6] combine diffusion trajectory model with behavior prediction to produce versatile traffic behaviors. Despite improved behavioral diversity, existing diffusion-based methods often struggle to balance agent interaction modeling with scene-conditioned consistency, leading to scene-inconsistent multi-agent trajectories [6]. Moreover, these methods are susceptible to closed-loop distribution shift and objective
3
misalignment, limiting robustness and reliability in realworld deployment [9]. In contrast, our method is built around this gap: we retain diffusion’s multimodal expressiveness while explicitly coupling scene-conditioned pre-training with RL posttraining to optimize both closed-loop safety and efficiency.
process gradually perturbs it into Gaussian noise through a Markov chain
q(u1:K | u0 ) :=
K Y
q(uk | uk−1 ),
(1)
k=1
where 2.2
Reinforcement Post-training for Driving Planners
Reinforcement learning (RL) has been shown to further enhance the capabilities of pretrained models by combining sampling-based exploration with reward-driven policy optimization [9], [32], [33]. Recent studies on reinforcement learning for driving planners generally focus on fine-tuning two major classes of generative model. The first line of work, analogous to RL fine-tuning for large language models, uses autoregressive generation to model each motion token as a continuous distribution and perform policy improvement [12], [34], [35]. However, this approach inherently suffers from conflicts between sequence-level and token-level objectives, and temporal instability induced by sequential decoding. In contrast, diffusion models provide an inherently temporally consistent decision process and produce diverse actions through probabilistic denoising in continuous space, making them particularly well-suited to RL’s exploration and exploitation paradigm. For example, TrajHF [13] proposes a human feedbackdriven RL fine-tuning scheme that aligns generative trajectory models with diverse human driving preferences. ReCogDrive [14] fine-tunes the diffusion planner using RL with a non-reactive simulator to generate safer and more stable trajectories. Nevertheless, existing methods are predominantly offline and do not interact with reactive environments; consequently, limited rollouts and exploration often yield marginal performance gains [11], [13], [14]. Due to the uncertainty and compounding errors introduced by environment interaction, RL training in a fully closed-loop manner remains an open challenge [10], [14]. Taken together, existing diffusion-based planners typically leave one of two gaps unresolved: weak scene-conditioned denoising or mostly offline post-training. Our method closes both by coupling a scene-conditioned denoising architecture with analytically tractable reversekernel online optimization, further stabilized by dense rewards and variance-gated updates.
3
P RELIMINARIES
This section summarizes the notation for scene-conditioned diffusion-based action generation and reinforcement learning used throughout the paper. 3.1
Diffusion Model and Diffusion Policies
Denoising Diffusion Probabilistic Models (DDPMs) [36] define a forward noising process together with a learned reverse denoising process for sample generation. In our setting, the clean sample u0 denotes a future action chunk, i.e., a control sequence to be generated for the controlled agents over the planning horizon, conditioned on the scene context c. Given a clean action chunk u0 , the forward diffusion
p q(uk | uk−1 ) = N uk ; 1 − βk uk−1 , βk I , k = 1, . . . , K. (2) Here, βk ∈ (0, 1) is a predefined noise schedule, αk = 1 − βk , Qk and ᾱk = i=1 αi . The noisy variable at step k admits the closed-form expression √ √ uk = ᾱk u0 + 1 − ᾱk ϵ, ϵ ∼ N (0, I). (3) The forward diffusion process is independent of the scene context c once u0 is given. Starting from Gaussian noise uK ∼ N (0, I), the reverse process generates actions conditioned on c:
pθ (u0:K | c) := p(uK )
K Y
pθ (uk−1 | uk , c),
(4)
k=1
where each reverse transition is parameterized as a Gaussian pθ (uk−1 | uk , c) := N uk−1 ; µ(uk , Dθ (uk , c, k)), σk2 I . (5) Here, uk denotes the noisy action variable at denoising step k , and uk−1 is its one-step denoised version. The reverse mean µ(·) is computed by the standard DDPM closed-form expression from the current noisy action chunk uk and the denoiser output Dθ (uk , c, k), while σk2 is determined by the fixed diffusion schedule. After K reverse steps, the final action chunk u0 is obtained. The corresponding conditioned diffusion policy πθ (u0 | c) is defined as marginal distribution induced on the final action chunk by the conditioned reverse denoising chain. 3.2
Markov Decision Process and Policy Optimization
A Markov decision process (MDP) is defined by a tuple M = (S, A, P0 , P, R), where S and A are the state and action spaces, P0 is the initial-state distribution, P is the transition kernel, and R is the immediate reward function. At time t, the agent observes a state st ∈ S , samples an action at ∈ A from the policy πθ (· | st ), receives reward R(st , at ), and transitions to st+1 ∼ P (· | st , at ). Starting from s0 ∼ P0 , the policy together with the environment dynamics induces a trajectory distribution. The RL objective is to maximize the expected discounted return "∞ # X t J (πθ ) = E(πθ ,P0 ,P ) γ R(st , at ) , (6) t=0
where γ ∈ (0, 1) is the discount factor. Policy-gradient methods optimize this objective using gradient estimators of the form "∞ # X ∇θ J (πθ ) = E(πθ ,P0 ,P ) ∇θ log πθ (at | st ) Gt , (7) t=0
where
Gt :=
X τ ≥t
γ τ −t R(sτ , aτ )
(8)
4
is the return-to-go from time t. In the MDP used later in this paper, the state corresponds to the current scene context and the action corresponds to the executable joint action chunk generated by the diffusion policy. The corresponding structured policy optimization will then be instantiated through the reverse denoising chain in Sec. 6.1.
4
P ROBLEM S TATEMENT
We consider multi-agent planning in interactive traffic scenes. A scene is represented as S = (x, u, c), where x ∈ RNa ×T ×Dx and u ∈ RNa ×T ×Du are state and control trajectories of Na agents over a finite horizon T , and c = (ca , cmp , ctl ) is scene context including: (i) agent history ca ; (ii) lane graph cmp ; and (iii) traffic light ctl . Dx , and Du the dimensional features. We aim to learn a scene-conditioned joint planning policy πθ (· | c) that produces interactionrealistic multi-vehicle behaviors while improving closed-loop safety and efficiency. To this end, leveraging diffusion models’ strong expressiveness for complex distributions, we first train a multiagent trajectory diffusion model via imitation learning to capture multimodal interactive behaviors, which jointly generate future trajectories for all vehicles conditioned on the scene context c. However, pretrained diffusion policies πθ often suffer from distribution shift and objective misalignment. A key question is: how can we preserve realistic interactive trajectories while promoting safety- and efficiencyoriented cooperative driving? We therefore propose M ULTI -ORFT, a cooperative, sceneconsistent multi-agent planner that post-trains the diffusion policy with stable online RL to improve closed-loop planning performance. Specifically, we sample rollouts x0 ∼ πθ (· | c) and optimize πθ by maximizing the expected cumulative reward J (πθ ).
5
M ULTI -AGENT D IFFUSION M ODEL P RETRAINING
This section describes the scene condition-enhanced multiagent diffusion planner, the pre-training half of M ULTI ORFT. It addresses the scene-consistency gap in prior diffusion planners, while Sec. 6 leverages the analytically tractable reverse kernel for stable online optimization. To better balance inter-agent interaction modeling with scene conditioning, we combine cross-attention with AdaLN-Zero modulation to improve the scene consistency and constraint adherence of joint trajectories. We first outline the planner architecture in Fig. 1, and then detail (i) a symmetric scene encoder for comprehensive scene conditioning and (ii) a multi-agent denoising decoder for scene-level interaction modeling. 5.1 Model Structure of the Multi-agent Diffusion Planner The planner Pθ ≜ (Eθ1 , Dθ2 ) is composed of a symmetric scene encoder Eθ1 and a multi-agent denoising decoder Dθ2 , as shown in Fig. 1: (i) Symmetric scene encoder Eθ1 : c 7→ c̃. The scene context is denoted by c = (ca , cmp , ctl ), where ca , cmp , and ctl represent agent history, lane-graph polylines, and traffic-light states, respectively. The encoder uses a querycentric attention Transformer to map the raw scene context
Algorithm 1: Pre-training of the multi-agent diffusion planner Input: Planner Pθ = (Eθ1 , Dθ2 ), dataset Dpre , diffusion steps K , dynamics f (·). Output: Trained parameters θ. 1 for each training iteration do 2 Sample expert trajectory xgt and scene context c from Dpre ; 3 Recover clean control chunk u0 ← f −1 (xgt ); 4 Sample diffusion step k ∼ U{1, . . . , K} and Gaussian noise ϵ ∼ N (0,√ I); √ 5 Form noisy control uk ← ᾱk u0 + 1 − ᾱk ϵ; 6 Predict denoised control û0 ← Dθ2 (uk , k, Eθ1 (c)); 7 Roll out denoised trajectory x̂0 ← f (û0 ); 8 Compute loss Lθ ← SL1 (x̂0 − xgt ); 9 Update θ = (θ1 , θ2 ) by backpropagation;
into latent scene tokens c̃. Specifically, scene elements are first transformed into local coordinate frames, and querycentric attention then combines token features with relative geometry to symmetrically capture pairwise relations among heterogeneous scene elements [20]. Detailed descriptions are provided in Sec. 5.2. (ii) Multi-agent denoising decoder Dθ2 : (c̃, uk , k) 7→ û0 , where û0 is the predicted clean action chunk u0 . Given the encoded scene context c̃, a noisy action chunk uk , and the denoising step index k , the decoder Dθ2 produces a denoising predicted output û0 . The decoder first models inter-agent dependencies through self-attention over agent trajectory tokens. It then injects scene information through cross-attention and further modulates the hidden states using AdaLN-Zero, allowing scene conditions to influence the denoising process at multiple layers. This conditioning strategy improves the consistency between generated trajectories and the surrounding scene context. Details of the decoder are given in Sec. 5.3. During pre-training, the scene encoder Eθ1 and the denoising decoder Dθ2 are optimized jointly in an endto-end manner. Algorithm 1 summarizes the loop: sample an expert trajectory and scene context, recover the clean control chunk through inverse dynamics, add noise at a random diffusion step, denoise it conditioned on the encoded scene, roll out the predicted controls through the vehicle dynamics, and optimize a Smooth-ℓ1 loss in trajectory space. The procedure clarifies that diffusion is defined in action space while supervision is imposed on the induced multiagent trajectories. 5.2
Symmetric Scene Context Modeling Encoder
Tokenization. The scene context c consists of three modalities: agent history ca ∈ RNa ×Th ×Da , lane graph polylines cmp ∈ RNm ×Mw ×Dp , and traffic-light states ctl ∈ RNt ×Dt , where Na , Nm , and Nt denote the numbers of agents, map polylines, and traffic lights, respectively; Th is the number of observed historical steps; Mw is the number of waypoints in each polyline; and Da , Dp , and Dt denote the corresponding feature dimensions. For each agent, the historical states are embedded by an MLP and fused with a learnable agenttype embedding to form an agent token. For each map polyline, the waypoint-level geometric attributes are first
5
Missing figure: Diffusion - vector.pdf
Fig. 1. Architecture of the multi-agent diffusion planner. A symmetric scene encoder models scene elements and their relations in local coordinates through query-centric self-attention. A denoising decoder predicts joint future plans via inter-agent self-attention, scene-conditioned generation, and AdaLN-Zero modulation.
encoded by an MLP, and then aggregated by max pooling along the waypoint dimension to obtain a polyline-level feature, which is further fused with discrete map attributes to form a polyline token. For each traffic light, the stoppoint coordinates and signal state are encoded by an MLP to produce a traffic-light token. In this way, heterogeneous scene elements are converted into a unified token representation. Local coordinates and query-centric attention. Before applying the Transformer encoder, positional attributes are transformed into local coordinate systems. Specifically, each agent is represented relative to its last observed state, and each polyline is represented relative to its first waypoint. This normalization reduces the burden of modeling absolute positions and emphasizes the relative geometric structure of the scene. The token set is then processed by L querycentric Transformer layers following [20]. For a query token (ℓ) i at layer ℓ, let Pi ∈ RD denote its embedding, and let Ω(i) denote the set of neighboring tokens associated with token i. For each neighbor j ∈ Ω(i), we express its pose in the local coordinate frame of token i and compute the relative geometric descriptor Rij = (∆xij , ∆yij , ∆ψij ). Let PE(Rij ) be a relative positional encoding and let [·, ·] denote the concatenation. The query-centric self-attention is (ℓ) ′(ℓ) Pi = MHSA Q : Pi , PE(Rii ) , (ℓ) (ℓ) K : Pj , PE(Rij ) j∈Ω(i) , V : Pj + PE(Rij ) j∈Ω(i) , (9) where MHSA(·query , ·key , ·value ) denotes multi-head selfattention operator. We apply the same computation to all tokens and produce the unified scene encoding, denoted by c̃ ∈ R(Na +Nm +Nt )×D . 5.3
Multi-agent trajectory denoising decoder
We diffuse joint control chunks in the action space of all vehicles and learn a denoiser Dθ2 conditioned on c̃ and diffusion step k . We embed noisy trajectories by rolling out controls through dynamics f (·) to obtain states x, which are then encoded with an MLP. This keeps kinematic structure explicit during denoising. The denoiser then applies Transformer blocks with inter-agent self-attention. Scene-conditioned modeling via cross attention and AdaLN-Zero. We propose a dual-path scene conditioning mechanism to model the scene-consistent joint trajectory distribution1 . As shown in Fig. 1, we couple (i) cross-attention, 1. Under high-dimensional inter-agent interaction modeling, relying on cross-attention alone can underutilize scene context and lead to scene-constraint violations [6].
which fuses the scene encoding as keys/values to update trajectory tokens, with (ii) AdaLN-Zero, which modulates trajectory tokens conditioned on the diffusion timestep and a dense road-feature representation. The former provides a direct information-injection pathway from the scene encoder, while the latter provides a flexible and stable constraintenforcement pathway that remains active throughout the denoising stack. Specifically, since road-graph embedding c̃mp can be sparse, we first densify it using a lightweight MLP-Mixer block that iteratively mixes information across the token and feature dimensions. ⊤ , c̃mp ← c̃mp + MLPtok c̃⊤ mp (10) c̃mp ← c̃mp + MLPch (c̃mp ) . AdaLN then modulates each trajectory-token feature x by AdaLN(x) = 1 + γ(k, c̃mp ) ⊙ x + β(k, c̃mp ), (11) where γ(·) and β(·) are the scale and shift factors regressed from the sum of the diffusion-step embedding and the dense road features. We further regress a gating factor α(k, c̃mp ) applied to the residual branch. We initialize all α to zero so that the full decoder starts as an identity function and gradually learns to introduce conditioning, which strengthens conditional modulation and improves training stability. Finally, the noise embedding is propagated to the final layer, which outputs the denoised action sequence u0 . We train the encoder Eθ1 and the denoising decoder Dθ2 by minimizing the loss:
Lθ = Eu0 , k, uk ∼pk (·|u0 ) [SL1 (x0 (Dθ2 (uk , k, Eθ1 (c))) − xgt )] , (12) where SL1 (·) is the Smooth-ℓ1 loss between the expert trajectory xgt and the trajectory x0 induced by executing the denoised controls u0 through the dynamics f (·).
6 P LANNING - ORIENTED O NLINE R EINFORCEMENT L EARNING A multi-agent diffusion planner trained solely via imitation learning can capture general driving capabilities, but it often fails under closed-loop execution due to distribution shift and the scarcity of safety-critical interaction data. Moreover, it fails to explicitly promote safety- and efficiency-oriented cooperative behaviors, relying solely on the BC loss. To overcome these limitations, we propose a stable online RL post-training framework (Fig. 3) that enhances both robustness and cooperation in closed-loop planning. To
6 •
Inner-MDP
ct
ct stK
ct
Outer-MDP
ct
utK −1 K −1 Pdp (⋅) 1 ut0 0 st st st
ct +1 PENV (⋅)
stK+1
RENV
(0, I ) K t
u
K −1 t
u
1 t
u
0 t
u
•
(0, I )
utK+1 Two-Layer MDP
support online optimization, we first formulate a two-level MDP. We then jointly design dense, well-shaped rewards and a variance-gated group-relative optimization (VG-GRPO) strategy to address instability during closed-loop training. Finally, we present the overall RL post-training pipeline. Two-level MDP Formulation
We now tailor the MDP to diffusion-policy optimization. Unlike standard policies that sample at in one shot, a diffusion policy produces the executed action u0t through a K -step stochastic denoising chain. Treating the entire denoising chain as a black-box sampler in a single-level MDP leads to poor credit assignment across denoising steps and highvariance gradients, which is exacerbated in multi-vehicle online rollouts due to severe non-stationarity from coupled agent interactions. We therefore model denoising as an inner MDP MDP with analytically tractable Gaussian likelihoods at each step, coupled with an outer environment-interaction MDP MENV that provides closed-loop rewards and constraints. Following [37], their composition yields the two-level MDP in Fig. 2. We use subscript t for the outer environment step and superscript k for the inner denoising step. •
•
•
•
State. s̄kt = (ct , ukt ), where ct is the scene observation at outer time step t and ukt is the intermediate denoising output at denoising step k . Action. ākt = uk−1 , sampled from the step-wise t reverse transition pθ (uk−1 | ukt , ct ); the final denoised t action is executed in the outer MDP. K Initial distribution. s̄K 0 = (c0 , u0 ) ∼ P̄0 , where c0 is drawn from the outer MDP’s initial distribution and uK 0 ∼ N (0, I). ′ Transition. The kernel P̄(s̄kt′ | s̄kt , ākt ) is defined as ( (ct , uk−1 ) (t′ , k ′ ) = (t, k − 1), k > 0, t k′ s̄t′ = K (ct+1 , ut+1 ) (t′ , k ′ ) = (t+1, K), k = 0, (13) For k > 0, transitions occur within the inner denoising chain: ct is fixed and the action is updated to uk−1 . For k = 0 we execute u0t in the environment to t obtain ct+1 and re-initialize uK t+1 ∼ N (0, I).
where RENV is specified in Sec. 6.2. Step-wise likelihood. In the inner MDP, the policy is identified with the reverse diffusion kernel. π̄θ (ākt | s̄kt ) can be evaluated by computing the diffusionpolicy likelihood of each denoising step along the sampled denoising chain.
π̄θ (ākt | s̄kt ) = pθ (uk−1 | ukt , ct ) t = N uk−1 ; µ ukt , Dθ2 ukt , k, ct , σk2 I t
Fig. 2. Two-level MDP formulation for reinforcement learning. We construct a two-level MDP consisting of an inner denoising MDP and an outer environment-interaction MDP, and use Gaussian likelihoods at each denoising step so that the denoising process can be optimized via policy gradients.
6.1
Reward. R̄(s̄kt , ākt ) is assigned only when denoising reaches k = 0: 0 k > 0, k k R̄(s̄t , āt ) = (14) RENV (ct , u0t ) k = 0,
•
(15)
where pθ (uk−1 | ukt , ct ) is a Gaussian transition with t mean µ(·) calculated from ukt and denoiser output k Dθ2 ut , k, ct , and variance σk2 specified by the fixed noise schedule. Hence, the step-wise log-likelihood is analytically tractable. Objective. The objective of the two-layer MDP is XX J¯ (π̄θ ) = Eπ̄θ ,P̄,P̄0 γ τ −t R̄ s̄kτ , ākτ . (16) t≥0 τ ≥t
6.2
Planning Reward Modeling
Rewards determine both the optimization direction and the driving preference. We use simple but fine-grained rulebased rewards that generalize across diverse large-scale scenarios and follow a safety-first principle, with explicit incentives for efficiency. The overall reward is defined as X RENV = wc Rcoll (j) + wo Roffroad (j) + we Reff (j). (17) j
where j is the step index within the reward horizon. Collision. We define the collision reward through SATbased oriented-box overlap checking, i.e., coll overlaps(·):
Rcoll (j) = 1[coll overlaps(j)].
(18)
Off-road. We define the off-road reward through signed distances of the four vehicle-box corners to lane boundaries, i.e., off road(·):
Roffroad (j) = 1[off road(j)].
(19)
Efficiency. We define the efficiency reward by measuring normalized progress along a road centerline: sj+1 − sj Reff (j) = max ,0 , (20) smax where sj is the arc-length projection of the vehicle position onto the centerline and smax normalizes the progress to [0, 1]. In multi-vehicle rollouts, we evaluate rewards at each step over the horizon and average them across the controlled vehicles. This formulation provides dense and well-shaped reward evaluations that enable finer-grained discrimination among sampled trajectories, thus yielding stable learning signals. To avoid trajectory dynamic-quality collapse when emphasizing safety, we use three built-in safeguards: dynamics-aware denoising with rollout through f (·) to ensure trajectory feasibility, KL anchoring to the
7
Policy Rollout Data
Closed-Loop
Scene Processing
Reference Model
Step
1.2
0.8
0.0
0.4
Gated Group Comp 1.3
0.4
KL
-1.3
-0.4
Policy Training
Fig. 3. Online RL post-training framework. The pipeline has three components: (1) policy rollout, where MADP samples multi-agent multimodal trajectories from updated scene context; (2) rule-based reward evaluation, where each trajectory group is scored by safety and efficiency metrics and converted to group-relative advantages; and (3) policy training, where the Actor network is optimized with policy loss and KL regularization.
pre-training prior, and short-horizon execute-then-replan in closed loop. 6.3
Variance-gated Group Relative Policy Optimization
With the two-level MDP (Sec. 6.1) and reward (Sec. 6.2) in place, we optimize the pretrained diffusion policy by maximizing the objective J¯(π̄θ ). We propose VG-GRPO for online refinement of multi-agent diffusion policy. Our design is built on the critic-free GRPO paradigm [15] and extends it with variance-gated advantages and denoisingaware optimization to stabilize training under multi-agent non-stationarity. This is a functional coupling rather than a loose combination: diffusion provides grouped multimodal candidates each step, and VG-GRPO directly consumes their relative ranking signal for online updates. At each outer environment step t, we fix the encoded context c̃t and sample a group of G candidate rollouts from the old policy π̄θold . Let i ∈ {1, . . . , G} index samples in the group. Each sample yields a scalar reward rt,i := R̄t,i , forming the within-group reward set {rt,1 , . . . , rt,G }. VG-GRPO objective. We maximize the following objective
J¯(θ) = E
+ min
G 1 X
G {āk t,i }i=1 ,s̄t ∼π̄θold G
ρkt,i , clip
−β DKL (π̄θ ||π̄ref )
i=1
ρkt,i , 1 − ε− , 1 + ε+
k γdenoise At,i
(21)
,
where ρkt,i is the importance ratio, computed from step-wise likelihoods under the two-level MDP in Sec. 6.1. π̄θ ākt,i | s̄kt , ρkt,i := (22) π̄θold ākt,i | s̄kt and [ε− , ε+ ] is the DAPO-style clipping ratio that stabilizes k training and encourages exploration, and γdenoise ∈ (0, 1] is a
denoising-step discount factor that downweights the policygradient contributions of noisier steps, thereby improving training stability [38]. To prevent mode collapse during exploration and preserve general driving capabilities, we regularize post-training with a KL term to the reference policy (set to the pretrained policy here), using the nonnegative unbiased estimator [39] and weighting it by β to control the regularization strength: π̄ref ākt,i | s̄kt π̄ref ākt,i | s̄kt − log − 1, DKL (π̄θ ||π̄ref ) = π̄θ ākt,i | s̄kt π̄θ ākt,i | s̄kt (23) Variance-gated advantages. Standard within-group normalization can induce training instability. In simple cases (e.g., short rollouts or near-stationary vehicles), sampled rewards become nearly identical; after GRPO’s mean–std normalization, the resulting advantages collapse toward zero, yielding vanishing policy gradients. This amplifies the sensitivity of minibatch gradients to noise and directly destabilizes training. To address this issue, we introduce a variance gated mechanism, inspired by [16] and adapted to fine-grained selection over group-based trajectory rollouts. Let σt ≜ std({rt,1 , . . . , rt,G }), the group relative advantage for sample i is defined as drop this group, rt,i − mean({rt,1 , . . . , rt,G }) , At,i = rt,i − mean({rt,1 , . . . , rt,G }) , std({rt,1 , . . . , rt,G })
σt ≤ std1 , std1 < σt ≤ std2 , σt > std2 .
(24) During policy sampling, we measure group diversity via the within-group standard deviation of rewards. (i) If the within-group standard deviation is near zero, rollouts are essentially identical, so we discard the group to avoid vanishing gradients. (ii) If it is small but nonzero, the differences may be noise-driven. Normalization would amplify this noise and corrupt the update direction, so we keep the raw differences to retain sample utility while reducing instability. (iii) If it is large, rollouts exhibit meaningful quality gaps and provide informative gradients; we then apply standard group-relative normalization to compute advantages. This variance gate retains informative samples and reduces gradient variance, thereby stabilizing typical GRPO for online RL training of multi-agent diffusion policies. 6.4
Reinforcement Learning Post-training Framework
We implement the online RL framework in Fig. 3 with three stages: rollout, reward evaluation, and policy update. In rollout, MADP repeatedly samples groups of multi-agent action chunks from the latest encoded scene context c̃t , which is refreshed after each executed action. In reward evaluation, rule-based metrics score collision, off-road behavior, and progress, producing a scalar reward for each sampled chunk. In policy update, grouped samples are used to optimize the diffusion-policy denoiser via VG-GRPO (Sec. 6.3) with KL regularization to the pre-training prior; the variance gate is critical for stabilizing multi-vehicle online refinement under non-stationary interactions. At each outer environment step t, the planner samples G candidate chunks by running the inner denoising chain
8
under the behavior policy pθold from the current context K s̄K t = (ct , ut ). Following [38], we use a rolling horizon: plan over Tp steps but execute only the first Ta steps. For group-relative optimization, we compute a scalar reward ri = RENV (ct , u0t,i ) for each sample i ∈ {1, . . . , G}. The environment executes the best-performing sample for interaction, while all group samples are retained to compute group-relative advantages and update the policy. We store grouped tuples in the rollout buffer Ditr : k G G {s̄kt,i }G (25) i=1 , {āt,i }i=1 , {rt,i }i=1 , k , where s̄kt,i = (ct , ukt,i ) and ākt,i = uk−1 follow the twot,i level MDP in Sec. 6.1. During updates, we sample minibatches from Ditr , compute variance-gated group-relative advantages (Eq. (24)), and optimize the planner parameters by maximizing the VG-GRPO objective (Eq. (21)). The complete post-training procedure and implementation details are summarized in Algorithm 2.
7
S IMULATION R ESULTS AND D ISCUSSION
This section reports closed-loop evaluation of the proposed planner. We first describe datasets, metrics, and implementation details, then present benchmark comparisons and ablation studies. The evaluation verifies three claims: closedloop safety and efficiency gains; robustness of the online post-training design choices; and the contribution of the AdaLN-Zero conditional module and VG-GRPO. 7.1
Algorithm 2: Closed-loop online RL post-training Input: Current policy πθ , dataset Dinter , group size G, denoise steps K , clip ε, KL weight β , thresholds stdlow , stdhigh . Output: Post-trained parameters θ. 1 for each iteration do 2 θold ← θ; Ditr ← ∅; 3 for each episode do 4 initialize outer state c1 from Dinter ; 5 for each outer step t do 6 initialize uK t,i ∼ N (0, I), i = 1, . . . , G; K 7 set s̄K = (c t , ut,i ), i = 1, . . . , G; t,i 8 for k = K, K − 1, . . . , 0 do 9 if k > 0 then 10 sample denoising actions ākt,i = uk−1 t,i ∼ pθold (· | ukt,i , ct ), i = 1, . . . , G; 11 set next inner states k−1 s̄k−1 t,i = (ct , ut,i ) and k R̄t,i = 0, i = 1, . . . , G; 12 else 0 13 compute terminal rewards R̄t,i ← 0 RENV (ct , ut,i ), i = 1, . . . , G; 0 14 select i⋆ = arg maxi R̄t,i and 0 execute ut,i⋆ in the simulator to obtain ct+1 ; 15
Setup
Dataset. For pre-training, we use the Waymo Open Motion Dataset (WOMD) [40], containing 486,995 training scenarios and 44,097 validation scenarios; each scenario covers 9 s of real traffic with all participant trajectories and map topology. We perform closed-loop evaluation and ablation studies across 41,590 Testing Interactive scenarios. For post-training, we uniformly sample 20,756 scenarios from the WOMD Validation Interactive split. We then evaluate the pre-trained trajectory diffuser on all sampled scenarios and partition them by performance score into three subsets: (i) a low-score set with 2,504 failure scenarios involving either inter-vehicle collisions or off-road events; (ii) a high-score set with 5,857 scenarios whose planned trajectories exhibit high quality; and (iii) a full set containing all available scenarios, which mixes the low-score, high-score, and regular scenarios. Implementation Details. Scene inputs include up to Nv = 32 agents, Nm = 256 polylines with Mw = 30 waypoints each, and Nt = 16 traffic lights. The planner generates a future control sequence of length Tp = 80 with a 0.1 s time step. We discard the entire history and condition on the current state. In closed-loop testing, we execute Ta = 10 steps. The scene encoder uses L = 6 query-centric Transformer layers with hidden size D = 256. The denoising decoder alternates two block types for three rounds (6 Transformer layers total). MLP-Mixer token/channel dimensions are 64/128. Pre-training uses the log noise schedule in [6] with ᾱmin = 10−9 , scaling 0.0031, and K = 20. We pre-train with AdamW (weight decay 0.01), learning rate 2×10−4 (3000-step warmup; ×0.02 decay every 3000 steps), gradient clipping 1.0, BF16 on 4×RTX 4090 for 30 epochs (global batch size 32).
store k G k G {s̄kt,i }G i=1 , {āt,i }i=1 , {R̄t,i }i=1 , k in
Ditr ; 16 17 18 19
20 21
for each update epoch do for each mini-batch do sample grouped tuples from Ditr ; compute variance-gated advantages At,i via Eq. (24); compute likelihood ratios ρkt,i ; update πθ by maximizing Eq. (21);
Post-training uses group size G = 10. Each rollout batch is optimized for 1 epoch with mini-batch size 16 for 10M fine-tuning steps. Following [9], we use 4 s reward horizon, reward weights (collision/off-road/progress) 8/1/4, KL weight β = 0.1, and learning rate 10−5 . For effective exploration, we use DAPO-style clipping [−0.15, 0.2] [16], and clamp a minimum sampling-time standard deviation of sam σmin = 0.2. For training stability, we clamp the per-step logprob likelihood standard deviation to σmin = 0.1. We set gated thresholds to std1 /std2 = 0.03/0.06 (10%/20% of the reward standard deviation based on), guided by signal-to-noise [41] and advantage-collapse analysis [42] and use γdenoise = 0.9. Post-training runs in BF16 on one NVIDIA 5090 GPU. Metrics. We report collision rate (CR), off-road rate (OR), average speed (AS), average displacement error (ADE), and kinematic infeasibility (Kin). We treat CR, OR, and AS as primary closed-loop objectives reflecting safety and traffic efficiency, and use ADE and Kin as secondary diagnostics
9
TABLE 1 Closed-loop benchmark on the WOMD testing interactive split (primary: CR, OR, AS; secondary: ADE, Kin) Method TrafficBotsV1.5 [43] SMART-large [24] VBD [6] SMART-tiny-CLSFT [44] M ULTI -ORFT
Primary closed-loop objectives
Secondary diagnostics
CR (%)↓
OR (%)↓
AS (m/s)↑ ADE (m)↓
Kin (%)↓
2.74±0.21 2.22±0.09 2.46±0.14 2.10±0.10 1.89±0.12
1.79±0.14 1.58±0.10 1.92±0.18 1.53±0.12 1.36±0.08
8.03±0.48 8.34±0.30 8.08±0.52 8.47±0.44 8.61±0.46
0.26±0.02 0.21±0.01 0.24±0.01 0.25±0.02 0.32±0.03
1.68±0.09 1.30±0.01 1.41±0.02 1.23±0.03 1.36±0.04
for imitation fidelity and physical feasibility. CR checks SAT-based oriented-box overlap; OR checks drivable-area boundary crossing; AS is mean per-step displacement; ADE is the ℓ2 position error to ground truth; and Kin counts violations of acceleration and curvature bounds with limits 6 m/s2 and 0.3 m−1 . For fairness, all methods are evaluated with the same simulator configuration, including scenario split, agent count, horizon, and replanning frequency; values with ± report mean and standard deviation over repeated closed-loop evaluations in the main comparison tables, while ablation tables report point estimates under the same protocol for compactness. 7.2
Evaluation
This subsection evaluates interactive trajectory generation using both benchmark tables and representative qualitative cases. On WOMD, we compare four strong open-source baselines: VBD [6], SMART-large [24], SMART-tiny-CLSFT [44], and TrafficBotsV1.5 [43]. They cover autoregressive, behaviorcloning, and diffusion paradigms. Each scenario runs for 8 s in closed loop with replanning at 1 Hz. Table 1 shows that M ULTI -ORFT leads on all three primary closed-loop objectives. Compared with the strong SMART-tiny-CLSFT baseline, which uses targeted closed-loop supervised finetuning, our method reduces CR from 2.10 to 1.89 (-10.0%) and OR from 1.53 to 1.36 (-11.1%), while increasing AS from 8.47 to 8.61 (+1.66%). Relative to the pretrained policy in Table 2, online post-training yields 7.4% lower CR, 19.0% lower OR, and 3.0% higher AS. This is consistent with the mechanism: VG-GRPO directly optimizes safety and efficiency rewards in closed loop, enabling better interaction outcomes under distribution shift. Post-training Methods. We further compare four posttraining strategies, SFT, DPO, offline RL, and online RL, as shown in Table 2. All methods improve over the pre-trainingonly policy in part, but with different objective alignment. SFT mainly improves imitation-oriented diagnostics (ADE). DPO and offline RL improve selected indicators but can degrade the balance between safety and efficiency. M ULTI ORFT with online RL provides the most consistent coimprovement on the primary closed-loop objectives CR, OR, and AS, reflecting the benefit of coupling diffusion group sampling with VG-GRPO. Online RL shows a mild ADE increase, indicating a distributional trade-off: closed-loop exploration can reduce one-step imitation accuracy while improving interaction robustness. Kin remains low across methods, all below 0.4%; compared with the pretrained policy, it increases by at
TABLE 2 Closed-loop comparison of post-training methods (primary: CR, OR, AS; secondary: ADE, Kin) Method Pre-trained only SFT DPO Offline RL M ULTI -ORFT (Online RL)
Primary closed-loop objectives
Secondary diagnostics
CR (%)↓
OR (%)↓
AS (m/s)↑ ADE (m)↓
Kin (%)↓
2.04±0.11 2.01±0.07 1.97±0.13 2.18±0.09 1.89±0.12
1.68±0.10 1.64±0.06 1.58±0.09 1.82±0.14 1.36±0.08
8.36±0.42 1.28±0.02 8.37±0.36 1.15±0.015 8.15±0.39 1.33±0.04 8.98±0.68 1.37±0.05 8.61±0.46 1.36±0.04
0.25±0.02 0.25±0.01 0.27±0.01 0.26±0.02 0.32±0.03
most 0.07 percentage points from 0.25 to 0.32, indicating no dynamic-quality collapse under safety-oriented optimization. Qualitative Visualization. Fig. 4 first illustrates the ability of the proposed planner to generate realistic and interactive trajectories in a real intersection scene involving about 20 vehicles. For clarity, each vehicle is indexed, and the colored traces behind the vehicles show their 1-second motion history. The timeordered rollout shows that the vehicles can pass through the scene in an orderly and safe manner under the planned trajectories. For example, when traversing a narrow road segment, Vehicles 0, 1, 20, and 28 adjust yielding and acceleration behaviors to maintain safe spacing and complete cooperative passing, indicating that the planner captures multi-vehicle coordination effectively. We further compare pre-training and RL post-training in an out-of-distribution interaction case. As shown in Fig. 5, in the pre-trained planner, the left-turn trajectory generated for Vehicle 3 does not adequately resolve its interaction conflict with Vehicle 20, leading to a collision at t = 6.0 s. After 1M RL steps, the planner learns a more conservative yielding strategy in the unprotected left-turn scenario, which avoids collision but reduces traffic efficiency. After 10M steps, the planner preserves safety while recovering more flexible and decisive interaction behaviors, and generates more efficient trajectories for moving vehicles such as Vehicles 0, 1, 2, 6, 10, and 18. These qualitative cases show that reward-driven RL post-training progressively reshapes the interaction strategy of the planner, improving the balance between safety and efficiency in closed-loop multi-agent driving. 7.3
Ablations
Unless otherwise stated, ablations use the same training pipeline and evaluation protocol as Section 7. Impact of AdaLN-Zero Module. AdaLN-Zero improves conditional generation quality. Table 3 shows gains on multiple metrics, especially off-road rate. The mechanism is direct: condition-driven modulation reshapes denoising features and strengthens road-geometry awareness, which is harder to achieve with cross-attention alone. Better boundary compliance in generation then transfers to improved closedloop stability. Fig. 6 provides a qualitative comparison on sharp rightturn cases. Without AdaLN-Zero, Vehicles 6 and 13 drift outward and hit the median boundary. With AdaLN-Zero, turn trajectories are more compact and maintain safer boundary margins. Impact of the Variance-gated Mechanism. Table 4 analyzes the effect of our variance-gated mechanism via ablations over several settings, including a non-gated baseline and
10
t = 4s
t = 2s
t = 6.0s
t = 8.0s
(a) Pretrain
Fig. 4. Closed-loop planning visualizations in real traffic scenes. We simulate the driving scenario in a closed-loop manner over an 8-second horizon, using Vehicle 1 as the centric viewpoint. The 1-second history trajectory is marked behind the vehicle.
Efficient
Safety
Efficient
Safety
Efficient
(c) RL 10M steps
(b) RL 1M steps
Safety
t = 4s
t = 0s
t = 6.0s
t = 8.0s
Fig. 5. Qualitative comparison between pre-training and RL post-training on an 8-second closed-loop rollout. TABLE 3 Ablation on AdaLN-Zero module Setting w/o AdaLN-Zero with AdaLN-Zero
CR (%)↓ OR (%)↓ AS (m/s)↑ ADE (m)↓ Kin (%)↓ 2.11 2.04
2.05 1.68
8.40 8.36
1.30 1.28
0.26 0.25
multiple threshold combinations. Without variance gating, post-training collapses at ∼0.5M steps, evidenced by per-
sistently large gradient-norm fluctuations and a sudden spike in the KL term, evaluating the last stable checkpoint also shows degraded performance across metrics. We then vary the gating thresholds. With std1 /std2 = 0.00/0.06, collapse is delayed to ∼2.0M steps and key closed-loop safety and efficiency metrics improve slightly, indicating that std2 contributes to training stability. In contrast, 0.03/0.06 and 0.03/0.09 prevent collapse completely, suggesting that stdlow is critical for stability by mitigating advantage degeneration
11
reactive closed-loop environments. For pre-training, M ULTI ORFT strengthens scene consistency and road adherence by combining inter-agent self-attention with dual-path scene conditioning. For post-training, M ULTI -ORFT achieves stable online learning by combining dense, well-shaped rewards with our proposed VG-GRPO, mitigating advantage collapse and gradient instability while strengthening closed-loop cooperative behaviors. Extensive experiments demonstrate superior closed-loop planning performance on key safety and efficiency metrics, outperforming state-of-the-art baselines on WOMD and across alternative post-training paradigms. (a) without AdaLN-Zero module (b) with AdaLN-Zero module Fig. 6. Qualitative comparison with and without the AdaLN-Zero module TABLE 4 Ablation of the variance-gated mechanism (performance is evaluated at the last checkpoint before collapse) Gating std1 /std2 Collapse step CR[%]↓ OR[%]↓ AS[m/s]↑ ADE[m]↓ Kin[%]↓ w/o gating 0.03/0.06 0.00/0.06 0.03/0.09
≈ 0.5M – ≈ 2.0M –
2.15 1.89 2.01 1.96
2.03 1.36 1.57 1.50
8.05 8.61 8.42 8.49
1.74 1.36 1.42 1.30
0.40 0.32 0.30 0.29
TABLE 5 Ablation on post-training data distribution Dataset type CR (%)↓ OR (%)↓ AS (m/s)↑ ADE (m)↓ Kin (%)↓ High-score Low-score Full
1.95 2.18 1.89
1.34 1.99 1.36
8.53 8.22 8.61
1.30 1.40 1.36
0.31 0.37 0.32
when within-group samples are nearly identical. Among stable settings, 0.03/0.06 achieves the best performance, a larger std2 appears overly conservative and weakens the effective learning signal, suggesting that moderate thresholds strike a better balance between stability and performance gains. Impact of Post-training Data Distribution. Table 5 analyzes how the composition of post-training scenes affects final performance. Training only on the high-score subset yields limited gains, because this subset contains many relatively simple scenes and thus provides weak optimization signals due to insufficient diversity among sampled trajectories. By contrast, training only on the lowscore subset substantially reduces collision rate and off-road rate, but also lowers average speed and degrades overall performance. This suggests that when post-training data are overly concentrated on hard failure cases, policy updates can be overly driven by local corrective objectives and fail to form a stable global optimization direction. The best overall result is obtained on the full dataset, which mixes high-score, low-score, and regular scenes. This observation indicates that a balanced post-training data distribution is important for stable reinforcement post-training: it maintains broader policy coverage and more stable sampling behavior, and thus provides more informative gradient signals for effective optimization and better generalization.
8
C ONCLUSION
In this paper, we introduced M ULTI -ORFT, a cooperative multi-agent planner that couples condition-enhanced diffusion pre-training with stable online RL post-training in
R EFERENCES J. Luo, T. Zhang, R. Hao, D. Li, C. Chen, Z. Na, and Q. Zhang, “Real-time cooperative vehicle coordination at unsignalized road intersections,” IEEE Trans. Intell. Transp. Syst., vol. 24, no. 5, pp. 5390–5405, 2023. [2] H. Bai, J. Luo, H. Li, X. Zhao, and Y. Wang, “A robust cooperative vehicle coordination framework for intersection crossing,” IEEE Trans. Veh. Technol., 2025. [3] N. Nayakanti, R. Al-Rfou, A. Zhou, K. Goel, K. S. Refaat, and B. Sapp, “Wayformer: Motion forecasting via simple and efficient attention networks,” arXiv preprint arXiv:2207.05844, 2022. [4] M. Janner, Y. Du, J. Tenenbaum, and S. Levine, “Planning with diffusion for flexible behavior synthesis,” in Proc. Int. Conf. Mach. Learn. (ICML). PMLR, 2022, pp. 9902–9915. [5] Y. Zheng, R. Liang, K. Zheng, J. Zheng, L. Mao, J. Li, W. Gu, R. Ai, S. E. Li, X. Zhan et al., “Diffusion-based planning for autonomous driving with flexible guidance,” arXiv preprint arXiv:2501.15564, 2025. [6] Z. Huang, Z. Zhang, A. Vaidya, Y. Chen, C. Lv, and J. F. Fisac, “Versatile behavior diffusion for generalized traffic agent simulation,” arXiv preprint arXiv:2404.02524, 2024. [7] Z. Huang, Z. Zhou, T. Cai, Y. Zhang, and J. Ma, “MDG: Masked denoising generation for multi-agent behavior modeling in traffic environments,” arXiv preprint arXiv:2511.17496, 2025. [8] Y. Lu, J. Fu, G. Tucker, X. Pan, E. Bronstein, R. Roelofs, B. Sapp, B. White, A. Faust, S. Whiteson et al., “Imitation is not enough: Robustifying imitation with reinforcement learning for challenging driving scenarios,” in Proc. IEEE/RSJ Int. Conf. Intell. Robots Syst. (IROS). IEEE, 2023, pp. 7553–7560. [9] H. Li, T. Li, J. Yang, H. Tian, C. Wang, L. Shi, M. Shang, Z. Lin, G. Wu, Z. Hao et al., “PlannerRFT: Reinforcing diffusion planners through closed-loop and sample-efficient fine-tuning,” arXiv preprint arXiv:2601.12901, 2026. [10] H. Gao, S. Chen, B. Jiang, B. Liao, Y. Shi, X. Guo, Y. Pu, H. Yin, X. Li, X. Zhang et al., “RAD: Training an end-to-end driving policy via large-scale 3DGS-based reinforcement learning,” arXiv preprint arXiv:2502.13144, 2025. [11] Z. Huang, X. Weng, M. Igl, Y. Chen, Y. Cao, B. Ivanovic, M. Pavone, and C. Lv, “Gen-drive: Enhancing diffusion generative driving policies with reward modeling and reinforcement learning finetuning,” in Proc. IEEE Int. Conf. Robot. Autom. (ICRA). IEEE, 2025, pp. 3445–3451. [12] Z. Peng, W. Luo, Y. Lu, T. Shen, C. Gulino, A. Seff, and J. Fu, “Improving agent behaviors with RL fine-tuning for autonomous driving,” in Proc. Eur. Conf. Comput. Vis. (ECCV). Springer, 2024, pp. 165–181. [13] D. Li, J. Ren, Y. Wang, X. Wen, P. Li, L. Xu, K. Zhan, Z. Xia, P. Jia, X. Lang et al., “Fine-tuning generative trajectory model with reinforcement learning from human feedback,” arXiv e-prints, 2025. [14] Y. Li, K. Xiong, X. Guo, F. Li, S. Yan, G. Xu, L. Zhou, L. Chen, H. Sun, B. Wang et al., “RecogDrive: A reinforced cognitive framework for end-to-end autonomous driving,” arXiv preprint arXiv:2506.08052, 2025. [15] D. Guo, D. Yang, H. Zhang, J. Song, P. Wang, Q. Zhu, R. Xu, R. Zhang, S. Ma, X. Bi et al., “DeepSeek-R1: Incentivizing reasoning capability in LLMs via reinforcement learning,” arXiv preprint arXiv:2501.12948, 2025. [16] Q. Yu, Z. Zhang, R. Zhu, Y. Yuan, X. Zuo, Y. Yue, W. Dai, T. Fan, G. Liu, L. Liu et al., “DAPO: An open-source LLM reinforcement learning system at scale,” arXiv preprint arXiv:2503.14476, 2025. [1]
12
[17] K. Zhang, Y. Zuo, B. He, Y. Sun, R. Liu, C. Jiang, Y. Fan, K. Tian, G. Jia, P. Li et al., “A survey of reinforcement learning for large reasoning models,” arXiv preprint arXiv:2509.08827, 2025. [18] W. Peebles and S. Xie, “Scalable diffusion models with transformers,” in Proc. IEEE/CVF Int. Conf. Comput. Vis. (ICCV), 2023, pp. 4195–4205. [19] S. Chen, B. Jiang, H. Gao, B. Liao, Q. Xu, Q. Zhang, C. Huang, W. Liu, and X. Wang, “VADv2: End-to-end vectorized autonomous driving via probabilistic planning,” arXiv preprint arXiv:2402.13243, 2024. [20] S. Shi, L. Jiang, D. Dai, and B. Schiele, “MTR++: Multi-agent motion prediction with symmetric scene modeling and guided intention querying,” IEEE Trans. Pattern Anal. Mach. Intell., vol. 46, no. 5, pp. 3955–3971, 2024. [21] Z. Zhou, Z. Wen, J. Wang, Y.-H. Li, and Y.-K. Huang, “QCNext: A next-generation framework for joint multi-agent trajectory prediction,” arXiv preprint arXiv:2306.10508, 2023. [22] J. Gu, C. Sun, and H. Zhao, “DenseTNT: End-to-end trajectory prediction from dense goal sets,” in Proc. IEEE/CVF Int. Conf. Comput. Vis. (ICCV), 2021, pp. 15 303–15 312. [23] B. Varadarajan, A. Hefny, A. Srivastava, K. S. Refaat, N. Nayakanti, A. Cornman, K. Chen, B. Douillard, C. P. Lam, D. Anguelov et al., “Multipath++: Efficient information fusion and trajectory aggregation for behavior prediction,” in Proc. IEEE Int. Conf. Robot. Autom. (ICRA). IEEE, 2022, pp. 7814–7821. [24] W. Wu, X. Feng, Z. Gao, and Y. Kan, “Smart: Scalable multi-agent real-time motion generation via next-token prediction,” Adv. Neural Inf. Process. Syst., vol. 37, pp. 114 048–114 071, 2024. [25] A. Seff, B. Cera, D. Chen, M. Ng, A. Zhou, N. Nayakanti, K. S. Refaat, R. Al-Rfou, and B. Sapp, “MotionLM: Multi-agent motion forecasting as language modeling,” in Proc. IEEE/CVF Int. Conf. Comput. Vis. (ICCV), 2023, pp. 8579–8590. [26] Z. Zhou, H. Haibo, X. Chen, J. Wang, N. Guan, K. Wu, Y.-H. Li, Y.-K. Huang, and C. J. Xue, “BehaviorGPT: Smart agent simulation for autonomous driving with next-patch prediction,” Adv. Neural Inf. Process. Syst., vol. 37, pp. 79 597–79 617, 2024. [27] Y. Yuan, X. Weng, Y. Ou, and K. M. Kitani, “Agentformer: Agentaware transformers for socio-temporal multi-agent forecasting,” in Proc. IEEE/CVF Int. Conf. Comput. Vis. (ICCV), 2021, pp. 9813–9823. [28] M. Jiang, Y. Bai, A. Cornman, C. Davis, X. Huang, H. Jeon, S. Kulshrestha, J. Lambert, S. Li, X. Zhou et al., “SceneDiffuser: Efficient and controllable driving simulation initialization and rollout,” Adv. Neural Inf. Process. Syst., vol. 37, pp. 55 729–55 760, 2024. [29] C. Jiang, A. Cornman, C. Park, B. Sapp, Y. Zhou, D. Anguelov et al., “MotionDiffuser: Controllable multi-agent motion prediction using diffusion,” in Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR), 2023, pp. 9644–9653. [30] Z. Zhong, D. Rempe, D. Xu, Y. Chen, S. Veer, T. Che, B. Ray, and M. Pavone, “Guided conditional diffusion for controllable traffic simulation,” in Proc. IEEE Int. Conf. Robot. Autom. (ICRA). IEEE, 2023, pp. 3560–3566.
[31] Z. Zhong, D. Rempe, Y. Chen, B. Ivanovic, Y. Cao, D. Xu, M. Pavone, and B. Ray, “Language-guided traffic simulation via scene-level diffusion,” in Proc. Conf. Robot Learn. (CoRL). PMLR, 2023, pp. 144–177. [32] D. Zhang, J. Liang, K. Guo, S. Lu, Q. Wang, R. Xiong, Z. Miao, and Y. Wang, “Carplanner: Consistent auto-regressive trajectory planning for large-scale reinforcement learning in autonomous driving,” in Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR), 2025, pp. 17 239–17 248. [33] B. Jiang, S. Chen, Q. Zhang, W. Liu, and X. Wang, “AlphaDrive: Unleashing the power of VLMs in autonomous driving via reinforcement learning and reasoning,” arXiv preprint arXiv:2503.07608, 2025. [34] Q. Li, X. Jia, S. Wang, and J. Yan, “Think2Drive: Efficient reinforcement learning by thinking with latent world model for autonomous driving (in CARLA-v2),” in Proc. Eur. Conf. Comput. Vis. (ECCV). Springer, 2024, pp. 142–158. [35] X. Tang, M. Kan, S. Shan, and X. Chen, “Plan-R1: Safe and feasible trajectory planning as language modeling,” arXiv preprint arXiv:2505.17659, 2025. [36] A. Q. Nichol and P. Dhariwal, “Improved denoising diffusion probabilistic models,” in Proc. Int. Conf. Mach. Learn. (ICML). PMLR, 2021, pp. 8162–8171. [37] K. Black, M. Janner, Y. Du, I. Kostrikov, and S. Levine, “Training diffusion models with reinforcement learning,” in Proc. ICML 2023 Workshop on Structured Probabilistic Inference and Generative Modeling, 2023. [38] A. Z. Ren, J. Lidard, L. L. Ankile, A. Simeonov, P. Agrawal, A. Majumdar, B. Burchfiel, H. Dai, and M. Simchowitz, “Diffusion policy policy optimization,” in Proc. CoRL 2024 Workshop on Mastering Robot Manipulation in a World of Abundant Data, 2024. [39] J. Schulman, “Approximating KL divergence,” http://joschu.net/blog/ kl-approx.html, 2020. [40] S. Ettinger, S. Cheng, B. Caine, C. Liu, H. Zhao, S. Pradhan, Y. Chai, B. Sapp, C. R. Qi, Y. Zhou et al., “Large scale interactive motion forecasting for autonomous driving: The Waymo Open Motion Dataset,” in Proc. IEEE/CVF Int. Conf. Comput. Vis. (ICCV), 2021, pp. 9710–9719. [41] J. Roberts and R. Tedrake, “Signal-to-noise ratio analysis of policy gradient algorithms,” Advances in neural information processing systems, vol. 21, 2008. [42] H. Zhong, J. Zhai, L. Song, J. Bian, Q. Liu, and T. Tan, “Rc-grpo: Reward-conditioned group relative policy optimization for multiturn tool calling agents,” arXiv preprint arXiv:2602.03025, 2026. [43] Z. Zhang, C. Sakaridis, and L. Van Gool, “TrafficBots V1.5: Traffic simulation via conditional VAEs and transformers with relative pose encoding,” arXiv preprint arXiv:2406.10898, 2024. [44] Z. Zhang, P. Karkus, M. Igl, W. Ding, Y. Chen, B. Ivanovic, and M. Pavone, “Closed-loop supervised fine-tuning of tokenized traffic models,” in Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR), 2025, pp. 5422–5432.