Robust RL with LLM-Driven Data Synthesis and Policy Adaptation for Autonomous Driving

TL;DR

Through fine-tuning via online environment interaction, RAPID reduces the forgetting of LLM knowledge while maintaining adaptability to different tasks, and demonstrates RAPID's capability to effectively integrate LLM knowledge into scaled-down RL policies in an efficient, adaptable, and robust way.

Abstract

The integration of Large Language Models (LLMs) into autonomous driving systems demonstrates strong common sense and reasoning abilities, effectively addressing the pitfalls of purely data-driven methods. Current LLM-based agents require lengthy inference times and face challenges in interacting with real-time autonomous driving environments. A key open question is whether we can effectively leverage the knowledge from LLMs to train an efficient and robust Reinforcement Learning (RL) agent. This paper introduces RAPID, a novel \underline{\textbf{R}}obust \underline{\textbf{A}}daptive \underline{\textbf{P}}olicy \underline{\textbf{I}}nfusion and \underline{\textbf{D}}istillation framework, which trains specialized mix-of-policy RL agents using data synthesized by an LLM-based driving agent and online adaptation. RAPID features three key designs: 1) utilization of offline data collected from an LLM agent to distil expert knowledge into RL policies for faster real-time inference; 2) introduction of robust distillation in RL to inherit both performance and robustness from LLM-based teacher; and 3) employment of a mix-of-policy approach for joint decision decoding with a policy adapter. Through fine-tuning via online environment interaction, RAPID reduces the forgetting of LLM knowledge while maintaining adaptability to different tasks. Extensive experiments demonstrate RAPID's capability to effectively integrate LLM knowledge into scaled-down RL policies in an efficient, adaptable, and robust way. Code and checkpoints will be made publicly available upon acceptance.

Figures (11)

• Figure 1: Cumulative reward on lane-3-density-2 with progressively increasing observation ($\mathbf{s}$) attack strength. The LLM-based agent $\pi_{\mathrm{LLM}}$ exhibits improved robustness compared with vanilla Offline DQN.
• Figure 2: Our proposed RAPID framework. Only modules tagged by * are trained. (a) Phase 1: Collect state-action rollouts from an environment and store them in a replay buffer. (b) Phase 2: Distill LLM knowledge into an offline policy using the collected data, the adapter policy is frozen and its output tokens are masked by zero gates. (c) Phase 3: Adapt the pre-trained model online by interacting with the environment, the LLM-distilled policy is frozen, and the zero gate is trained for progressive adaptation.
• Figure 3: Cumulative reward over different training ratios under offline training framework in highway-fast environment. The result is averaged over $5$ random trails.
• Figure 4: Performance of online adaptation (Phase 3) across three environments. Before the 5K-iteration mark, we pre-train the $\pi_{\mathrm{distil}}$ policy of RAPID using the method described in Phase 2. Note that $\pi_\mathrm{distil}$ keeps frozen in Phase 3. We report the average performance over $5$ random trails.
• Figure 5: Visualizing the contribution of $\pi_\mathrm{distil}$ and $\pi_\mathrm{adapt}$ to the final predicted action. The example is randomly sampled via interacting with lane-4-density-2.5 using $\Pi_\mathrm{MoP}$. More samples are demonstrated in Appendix. \ref{['app:visualization_policies']}.

Theorems & Definitions (1)

• Remark 1