- The paper presents a novel integration of imitation and reinforcement learning via multi-mode trajectory probing in a differentiable 3D simulation environment.
- It employs flow matching with optimal transport and Sinkhorn regularization to generate diverse, precise future trajectories.
- Experimental results on nuScenes benchmarks show significant improvements in episode rewards, collision rates, and overall policy robustness.
GSDrive: Reinforcing Driving Policies via Multi-Mode Trajectory Probing in 3D Gaussian Splatting
Introduction
GSDrive addresses the challenge of robust end-to-end (E2E) policy optimization for autonomous driving by introducing a differentiable, physically grounded reward shaping mechanism. Traditional E2E approaches—combining imitation learning (IL) with reinforcement learning (RL)—lack sufficient reward signal attribution, rely on costly annotations, and often converge prematurely to suboptimal solutions due to sparse, delayed rewards. GSDrive proposes leveraging a 3D Gaussian Splatting (3DGS) environment for high-fidelity simulation and introduces multi-mode trajectory probing, facilitating bidirectional knowledge transfer between IL and RL and enabling dense, future-aware feedback during learning.
Architecture and Methodology
3D Gaussian Splatting-Based Environment
GSDrive utilizes VGGT (Visual Geometry Grounded Transformer) to reconstruct 3D driving scenes from multi-view images, transforming pixel features across spatially and temporally aligned camera viewpoints into physically consistent 3D Gaussian primitives. The primitives encode spatial centers, scaling, orientation (unit quaternions), opacity, and color, supervised by photometric, structural, and geometric losses against image and LiDAR data. The resulting 3DGS environment maintains spatial and temporal coherence for policy simulation, providing direct access to accurate agent-environment interaction signals.
Multi-Mode Trajectory Prediction with Flow Matching
A multi-modal trajectory head is trained via imitation learning, employing flow matching with an optimal transport (OT) objective. Clustered trajectories are interpolated towards ground-truth trajectories along cost-minimizing curves by OT-guided velocity fields, with Sinkhorn regularization ensuring computational tractability and entropy for large datasets. This mechanism allows explicit learning of diverse, physically plausible future paths, avoiding mode averaging and preserving sharp, decisive trajectories critical for driving tasks.
Trajectory predictions are transformed into discrete action anchors through an exponential kernel, producing action logits weighted by multi-mode classification probabilities and residual head outputs. The architecture enables hierarchical optimization, reinforcing informed exploration and policy stability.
Reinforcement Learning with Prospective Reward Probing
GSDrive defines RL as a Markov Decision Process in the reconstructed 3DGS environment. Actions correspond to trajectory endpoints, while state observations include rendered camera views, detection labels, and camera matrices. The main innovation is in reward construction:
- Base Rewards: Survival bonus, progress, collision penalties, and comfort metrics (jerk/acceleration).
- Trajectory Probing Rewards: Multiple candidate trajectories are simulated forward, each evaluated over a fixed probe horizon. The maximal reward from these probes is used to guide the current action update, effectively grounding reward signals in physically simulated interaction outcomes.
Policy optimization proceeds via a clipped surrogate objective (PPO style), value estimation, entropy regularization, and an adaptive KL divergence penalty maintained through exponential moving averages. This setup constrains the policy's deviation from demonstration-based trajectories, balancing exploration and safety.
Experimental Results
GSDrive demonstrates strong empirical performance on reconstructed nuScenes benchmarks. Closed-loop evaluations align with significant improvement across all metrics:
- Episode Reward (ER): Highest episodic reward (52.97) compared to prior RL-based approaches.
- Collision Rate (CR): Lowest collision rates (0.11), indicating effective preemptive avoidance.
- Driving Speed (DS), Lane Changes (LC), and Comfort Metrics: GSDrive achieves top scores, substantiating its ability to reconcile performance, safety, and comfort.
The ablation studies reveal the necessity of trajectory-constructed logits and physical probing rewards. Flow matching with classifier-free guidance (CFG) produces physically consistent, multi-modal trajectories, outperforming diffusion-based approaches (DDPM/DDIM) and pure flow matching models.
Qualitative assessments confirm GSDrive's superiority in handling out-of-distribution scenarios and complex interactions, with closed-loop stability and better anticipation of long-tail events.
Implications and Future Directions
GSDrive's approach fundamentally restructures reward assignment in RL for autonomous driving. By enabling dense, physically interpretable reward signals through multi-mode trajectory probing in a differentiable simulator, the method offers:
- Lifelong, Self-Supervised Policy Refinement: Unlabeled real-world observations can be continuously integrated, reducing dependence on manual annotation.
- Enhanced Robustness and Generalization: Dense prospective signals mitigate exposure bias and information bottlenecks inherent to standard RL, fostering safety-aware exploration and adaptation.
- Bridging IL and RL: GSDrive framework supports bidirectional knowledge transfer, allowing deviation from demonstrator behavior for enhanced safety and performance while retaining critical inductive priors.
Potential future developments include integrating more sophisticated world models, leveraging larger-scale reconstruction paradigms (e.g., video prior diffusion), enabling uncertainty quantification during probing, and scaling closed-loop evaluation to diverse geographic domains and sensor setups. The framework is extensible to embodied agents beyond autonomous vehicles where predictive, physically grounded reward shaping can significantly improve policy stability and safety.
Conclusion
GSDrive establishes a new paradigm for E2E driving policy optimization by integrating imitation priors with RL through multi-mode trajectory probing in a 3D Gaussian Splatting environment. Dense, future-aware reward feedback substantially elevates training stability and closed-loop robustness, as evidenced by strong empirical results. The method's ability to reconcile expert-driven inductive biases with exploration—without relying on sparse, catastrophic rewards—sets a firm foundation for continuous, self-supervised deployment and theoretical advances in physically grounded RL for embodied AI systems (2604.28111).
