README.md

---
base_model: meta-llama/Meta-Llama-3.1-8B-Instruct
library_name: peft
license: other
license_name: ncbi-public-domain
license_link: https://github.com/RAG-Gym/RAG-Gym/blob/main/LICENSE
pipeline_tag: text-classification
tags:
- rag-gym
- retrieval-augmented-generation
- agent
- peft
- lora
- process-supervision
- reward-model
---

# Direct-MedQA-PRM

This is the **PRM** process reward model (critic) for the **Direct** agent on **MedQA**, released as part of [RAG-Gym](https://github.com/RAG-Gym/RAG-Gym). It is a baseline agent used for comparison in RAG-Gym. The critic scores candidate actions so that an agent can select the best one at each step through critic-guided (Best-of-N) inference.

- **Base model:** [`meta-llama/Meta-Llama-3.1-8B-Instruct`](https://huggingface.co/meta-llama/Meta-Llama-3.1-8B-Instruct)
- **Agent architecture:** Direct (a baseline agent used for comparison in RAG-Gym)
- **Task / dataset:** MedQA (medical multiple-choice question answering in the USMLE style)
- **Training method:** process reward modeling (PRM)
- **Adapter:** PEFT LoRA (r=256, alpha=512), task type `SEQ_CLS`
- **Precision:** bf16

## Intended use

- Score and rank candidate actions generated by the Direct agent during search on MedQA-style tasks.
- Enable critic-guided (Best-of-N) inference within the RAG-Gym framework.

This model is intended for research on process-supervised retrieval-augmented generation. It is not intended for clinical decision-making or other high-stakes use.

## How to use

Load the LoRA adapter on top of the base model:

```python
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel

base = "meta-llama/Meta-Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(base)
model = AutoModelForSequenceClassification.from_pretrained(base, num_labels=1, torch_dtype=torch.bfloat16)
model = PeftModel.from_pretrained(model, "RAG-Gym/Direct-MedQA-PRM")
```

For full usage within the agent framework (environment setup, inference, and critic-guided Best-of-N selection), see the [RAG-Gym repository](https://github.com/RAG-Gym/RAG-Gym).

## Training

Fine-tuned from `meta-llama/Meta-Llama-3.1-8B-Instruct` on MedQA process-reward data using process reward modeling (PRM) with the TRL library. Adaptation uses LoRA (rank 256, alpha 512, task type `SEQ_CLS`) in bf16. The training code and full configuration are available in the [RAG-Gym repository](https://github.com/RAG-Gym/RAG-Gym).

## Limitations

The model is trained on MedQA-style data and is intended for research use within RAG-Gym. Outputs may be inaccurate or unsupported by retrieved evidence and should not be relied upon for high-stakes decisions. As an 8B-parameter model, it requires a CUDA-enabled GPU to run efficiently.

## License

Public Domain Notice (U.S. Government Work, NCBI). See the [LICENSE](https://github.com/RAG-Gym/RAG-Gym/blob/main/LICENSE) in the repository. Please cite the authors when using this material.

## Citation

```bibtex
@article{xiong2025raggym,
    title={RAG-Gym: Optimizing Reasoning and Search Agents with Process Supervision},
    author={Guangzhi Xiong and Qiao Jin and Xiao Wang and Yin Fang and Haolin Liu and Yifan Yang and Fangyuan Chen and Zhixing Song and Dengyu Wang and Minjia Zhang and Zhiyong Lu and Aidong Zhang},
    journal={arXiv preprint arXiv:2502.13957},
    year={2025}
}
```