# Reinforcement Learning for Optimizing RAG for Domain Chatbots

arXiv: https://arxiv.org/abs/2401.06800 · HTML: https://mlsanity.com/paper/2401.06800/

## Key facts

- Cited by 26 papers in the map (data through 14 Oct 2026).
- Placed in RAG pipeline design and evaluation, one of 15 directions in Retrieval-augmented generation.
- Builds on 9 of its 11 references that are in the map.
- First posted on arXiv 10 Jan 2024.

## Placed in

- Retrieval-augmented generation / [RAG pipeline design and evaluation](https://mlsanity.com/rag/rag-pipeline-design-and-evaluation/)
- Deep reinforcement learning methods / [RL training of language and agent models](https://mlsanity.com/deeprl/rl-training-of-language-and-agent-models/)

## Builds on (most cited first, 9 of 9)

- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [G-Eval: NLG Evaluation using GPT-4 with Better Human Alignment](https://mlsanity.com/paper/2303.16634/) (2023)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [Enhancing Retrieval-Augmented Large Language Models with Iterative Retrieval-Generation Synergy](https://mlsanity.com/paper/2305.15294/) (2023)
- [Self-supervised Pretraining of Visual Features in the Wild](https://mlsanity.com/paper/2103.01988/) (2021)
- [A Survey on Retrieval-Augmented Text Generation](https://mlsanity.com/paper/2202.01110/) (2022)
- [Is GPT-4 a reliable rater? Evaluating Consistency in GPT-4 Text Ratings](https://mlsanity.com/paper/2308.02575/) (2023)

## Built on by (most cited first, 26 of 26)

- [Retrieval-Augmented Generation for AI-Generated Content: A Survey](https://mlsanity.com/paper/2402.19473/) (2024)
- [The Good and The Bad: Exploring Privacy Issues in Retrieval-Augmented Generation (RAG)](https://mlsanity.com/paper/2402.16893/) (2024)
- [Unveiling Privacy Risks in LLM Agent Memory](https://mlsanity.com/paper/2502.13172/) (2025)
- [A Survey of Graph Retrieval-Augmented Generation for Customized Large Language Models](https://mlsanity.com/paper/2501.13958/) (2025)
- [A Survey of Personalization: From RAG to Agent](https://mlsanity.com/paper/2504.10147/) (2025)
- [MA-RAG: Multi-Agent Retrieval-Augmented Generation via Collaborative Chain-of-Thought Reasoning](https://mlsanity.com/paper/2505.20096/) (2025)
- [DynamicRAG: Leveraging Outputs of Large Language Model as Feedback for Dynamic Reranking in Retrieval-Augmented Generation](https://mlsanity.com/paper/2505.07233/) (2025)
- [HyperMem: Hypergraph Memory for Long-Term Conversations](https://mlsanity.com/paper/2604.08256/) (2026)
- [Human-inspired Perspectives: A Survey on AI Long-term Memory](https://mlsanity.com/paper/2411.00489/) (2024)
- [Learning When to Retrieve, What to Rewrite, and How to Respond in Conversational QA](https://mlsanity.com/paper/2409.15515/) (2024)
- [Privacy-Preserving Federated Embedding Learning for Localized Retrieval-Augmented Generation](https://mlsanity.com/paper/2504.19101/) (2025)
- [End-to-End Optimization of LLM-Driven Multi-Agent Search Systems via Heterogeneous-Group-Based Reinforcement Learning](https://mlsanity.com/paper/2506.02718/) (2025)
- [Reward-RAG: Enhancing RAG with Reward Driven Supervision](https://mlsanity.com/paper/2410.03780/) (2024)
- [Self-Aligned Reward: Towards Effective and Efficient Reasoners](https://mlsanity.com/paper/2509.05489/) (2025)
- [Improving LLM Reliability with RAG in Religious Question-Answering: MufassirQAS](https://mlsanity.com/paper/2401.15378/) (2024)
- [Hybrid AI for Responsive Multi-Turn Online Conversations with Novel Dynamic Routing and Feedback Adaptation](https://mlsanity.com/paper/2506.02097/) (2025)
- [Token-Level Precise Attack on RAG: Searching for the Best Alternatives to Mislead Generation](https://mlsanity.com/paper/2508.03110/) (2025)
- [Improving Context Fidelity via Native Retrieval-Augmented Reasoning](https://mlsanity.com/paper/2509.13683/) (2025)
- [State Contamination in Memory-Augmented LLM Agents](https://mlsanity.com/paper/2605.16746/) (2026)
- [Reinforced Informativeness Optimization for Long-Form Retrieval-Augmented Generation](https://mlsanity.com/paper/2505.20825/) (2025)
- [Contextually Aware E-Commerce Product Question Answering using RAG](https://mlsanity.com/paper/2508.01990/) (2025)
- [Less Is More: Elevating RAG via Performance-Driven Context Compression](https://mlsanity.com/paper/2508.19282/) (2025)
- [MemPot: Defending Against Memory Extraction Attack with Optimized Honeypots](https://mlsanity.com/paper/2602.07517/) (2026)
- [NIFTY Financial News Headlines Dataset](https://mlsanity.com/paper/2405.09747/) (2024)
- [Optimizing Retrieval for RAG via Reinforcement Learning](https://mlsanity.com/paper/2510.24652/) (2025)
- [Reliable Conversational Agents under ASP Control that Understand Natural Language](https://mlsanity.com/paper/2502.09237/) (2025)

## Abstract (arXiv)

With the advent of Large Language Models (LLM), conversational assistants have become prevalent for domain use cases. LLMs acquire the ability to contextual question answering through training, and Retrieval Augmented Generation (RAG) further enables the bot to answer domain-specific questions. This paper describes a RAG-based approach for building a chatbot that answers user's queries using Frequently Asked Questions (FAQ) data. We train an in-house retrieval embedding model using infoNCE loss, and experimental results demonstrate that the in-house model works significantly better than the well-known general-purpose public embedding model, both in terms of retrieval accuracy and Out-of-Domain (OOD) query detection. As an LLM, we use an open API-based paid ChatGPT model. We noticed that a previously retrieved-context could be used to generate an answer for specific patterns/sequences of queries (e.g., follow-up queries). Hence, there is a scope to optimize the number of LLM tokens and cost. Assuming a fixed retrieval model and an LLM, we optimize the number of LLM tokens using Reinforcement Learning (RL). Specifically, we propose a policy-based model external to the RAG, which interacts with the RAG pipeline through policy actions and updates the policy to optimize the cost. The policy model can perform two actions: to fetch FAQ context or skip retrieval. We use the open API-based GPT-4 as the reward model. We then train a policy model using policy gradient on multiple training chat sessions. As a policy model, we experimented with a public gpt-2 model and an in-house BERT model. With the proposed RL-based optimization combined with similarity threshold, we are able to achieve significant cost savings while getting a slightly improved accuracy. Though we demonstrate results for the FAQ chatbot, the proposed RL approach is generic and can be experimented with any existing RAG pipeline.