# Test-Time Policy Adaptation for Enhanced Multi-Turn Interactions with LLMs

This repository contains the implementation and experimental code for the paper: **Test-Time Policy Adaptation for Enhanced Multi-Turn Interactions with LLMs**.


## Introduction to ROSA

Large Language Models (LLMs) employ multi-turn interaction as a fundamental paradigm for completing complex tasks. However, their performance often degrades in extended interactions, as they are typically trained on static, single-turn data, which hinders their ability to adapt to real-time user feedback. To address this limitation, we first propose a new paradigm: Test-Time Policy Adaptation for Multi-Turn Interactions (T$^2$PAM), which utilizes user feedback from the ongoing interaction as a reward signal to estimate a latent optimal policy aligned with user preferences, then updates a small subset of parameters to steer the model toward this policy, ultimately enabling efficient in-conversation self-correction. We then introduce Optimum-Referenced One-Step Adaptation (ROSA), a lightweight algorithm that operationalizes T$^2$PAM. ROSA guides the model parameters toward a theoretical optimal policy in a single, efficient update step, avoiding costly iterative gradient-based optimization and minimizing computational overhead. We provide a rigorous theoretical analysis guaranteeing that the policy of ROSA converges to the preference of user as the number of interactions increases. Extensive experiments on challenging benchmark demonstrate that ROSA achieves significant improvements in both task effectiveness and efficiency.

## File Structure

- **data**: Contains the testing data for ROSA.
- **direct_update_lmhead.py**: The test script of use the rule base reward to update LM Head of model.
- **direct_update_hidden_states.py**: The test script of use the rule base reward to update Hidden States of model. 
- **direct_update_lmhead.py**: The test script of use the model base reward to update LM Head of model.
- **run_direct.sh**: Startup Script.

## Getting Started

### Prerequisites

- Python 3.10.15
- torch==2.5.1
- transformers==4.49.0.dev0
- datasets==3.2.0

### Installation



```bash
 
conda create --name rosa python=3.10 pytorch torchvision pytorch-cuda -c nvidia -c pytorch -y
conda activate rosa
pip install -r requirements.txt
```
### Then, run the test script:
```bash
 bash run_direct.sh
```
