# TMPC: Test-Time Alignment via Textual Model Predictive Control

Official implementation of **"Test-Time Alignment for Large Language Models via Textual Model Predictive Control"**, published at **ICLR 2026**.

[[Paper]](https://openreview.net/forum?id=DsS3xRPSs5)

TMPC is a lightweight, training-free framework for aligning LLMs at inference time. It adapts Model Predictive Control to text generation through two core principles:
1. **Hindsight Subgoal Identification** — analyzing high-reward intermediate outputs across multiple generation rollouts.
2. **Subgoal-Conditioned Re-Generation** — using identified subgoals to guide subsequent planning iterations.
---

## Installation

### 1. Create a Conda Environment

```bash
conda create -n TMPC python=3.10 -y
conda activate TMPC
```

### 2. Install PyTorch

Install PyTorch with the CUDA version matching your driver (check with `nvidia-smi`):

```bash
# For CUDA 13.0 (Blackwell GPUs, e.g., RTX PRO 6000)
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu130

# For CUDA 12.6
# pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu126
```

### 3. Install Python Dependencies

```bash
pip install uv
uv pip install -r requirements.txt
```

### 4. Install SpaCy Language Models

Install the SpaCy models for your source and target languages:

```bash
python -m spacy download en_core_web_sm   # English
python -m spacy download zh_core_web_sm   # Chinese
python -m spacy download ja_core_news_sm  # Japanese
python -m spacy download de_core_news_sm  # German
python -m spacy download ru_core_news_sm  # Russian
python -m spacy download es_core_news_sm  # Spanish
python -m spacy download ko_core_news_sm  # Korean
```

### 5. Setup LASER (Translation Task Only)

TMPC uses [LASER](https://github.com/facebookresearch/LASER) and [VecAlign](https://github.com/thompsonb/vecalign) for sentence alignment in the translation pipeline. VecAlign is bundled in this repository (`vecalign.py`, `overlap.py`, `dp_utils.py`, `dp_core.pyx`).

```bash
# Clone LASER
git clone https://github.com/facebookresearch/LASER.git /path/to/LASER
export LASER="/path/to/LASER"

# Install LASER external tools (moses tokenizer, sentencepiece, etc.)
cd $LASER && bash install_external_tools.sh

# Install LASER Python package and its runtime dependencies
pip install --no-deps fairseq
pip install hydra-core omegaconf bitarray unicategories
pip install -e $LASER --no-deps

# Download LASER2 model files
mkdir -p $LASER/models
wget -q https://dl.fbaipublicfiles.com/nllb/laser/laser2.pt -O $LASER/models/laser2.pt
wget -q https://dl.fbaipublicfiles.com/nllb/laser/laser2.spm -O $LASER/models/laser2.spm
```

Configure the LASER embed script to find the models:

```bash
# Edit $LASER/tasks/embed/embed.sh — set model_dir to:
model_dir="${LASER}/models"
```

Add `export LASER="/path/to/LASER"` to your `~/.bashrc` for persistence. The `$LASER` environment variable must be set when running WMT experiments.

### 6. Download Models from HuggingFace

TMPC requires one generation model and task-specific reward models:

```bash
# Generation model (used by all three tasks)
huggingface-cli download meta-llama/Meta-Llama-3.1-8B-Instruct

# Reward model for HH-RLHF task
huggingface-cli download rl-bandits-lab/hhrlhf_rm

# Reward model for WMT translation task
huggingface-cli download rl-bandits-lab/translation_rm

# COMET model for translation evaluation
huggingface-cli download Unbabel/wmt22-comet-da
```

> **Note:** Accessing `meta-llama/Meta-Llama-3.1-8B-Instruct` requires accepting the license on HuggingFace and authenticating via `huggingface-cli login`.

---

## Experiments

### HH-RLHF: Long-Form Response Generation

**Generate** responses with iterative TMPC refinement:

```bash
python run_hhrlhf.py \
    --input_file hhrlhf.csv \
    --output_folder tmpc_hhrlhf_output \
    --rm_path rl-bandits-lab/hhrlhf_rm \
    --max_iterations 3 \
    --buffer_size 3 \
    --threshold 4 \
    --cuda_num 0 \
    --start 0 \
    --end 1024
```

**Evaluate** the generated results (selects the best iteration per prompt):

```bash
python run_hhrlhf.py \
    --evaluate \
    --eval_input_folder tmpc_hhrlhf_output \
    --eval_it 3 \
    --eval_range 1024 \
    --eval_output_file hh_eval_result.csv
```

### WMT: Document-Level Machine Translation

**Generate** translations with TMPC buffer-based refinement (ensure `$LASER` is set):

```bash
export LASER="/path/to/LASER"
python run_wmt.py \
    --input_file wmt_zh_en.csv \
    --rm rl-bandits-lab/translation_rm \
    --src_language Chinese \
    --task_language English \
    --threshold 1.0 \
    --max_iterations 5 \
    --good_ref_contexts_num 5 \
    --cuda_num 0
```

**Merge** iteration results into a single CSV for evaluation:

```bash
python memory2csv.py \
    --num 4 \
    --input_csv wmt_zh_en.csv \
    --output_csv eval_zh_en.csv \
    --column_name TMPC
```

**Evaluate** with SEGALE context-level COMET scoring:

```bash
python segale_ctx.py \
    --file eval_zh_en.csv \
    --target_column TMPC \
    --save eval_zh_en \
    --src_language Chinese \
    --task_language English \
    --gpu_id 0
```

### MBPP: Program Synthesis

In this code, we use the **sanitized MBPP** dataset ([google-research/mbpp/sanitized-mbpp.json](https://github.com/google-research/google-research/blob/master/mbpp/sanitized-mbpp.json)).
The sanitized subset contains 427 problems; evaluation is filtered to Task IDs 11–510.

The paper reports results on both the sanitized subset and supplementary non-sanitized MBPP problems (Task IDs 11–510 not included in the sanitized set), where problem descriptions were augmented with function signature hints extracted from test cases.

**Generate** solutions with TMPC exploration:

```bash
python run_mbpp.py \
    --output_folder tmpc_mbpp_output \
    --max_iterations 5 \
    --num_candidates 3 \
    --cuda_num 0 \
    --load_in_4bit \
    --start 0 \
    --end 427
```

**Evaluate** pass rates (from generated files in folder):

```bash
python run_mbpp.py \
    --evaluate \
    --eval_folder tmpc_mbpp_output \
    --eval_max_iter 5 \
    --eval_output tmpc_results.csv \
    --eval_from_folder
```

---

## Running Full Experiments

We provide scripts under `scripts/` that run the full experiments:

```bash
# HH-RLHF
bash scripts/run_hhrlhf_full.sh

# WMT zh->en 
export LASER="/path/to/LASER"
bash scripts/run_wmt_zh_en.sh

# MBPP
bash scripts/run_mbpp_full.sh
```

Results are saved to the `results/` directory.

---

## VecAlign in the Translation Pipeline

The WMT translation task relies on [VecAlign](https://github.com/thompsonb/vecalign) for **sentence-level alignment** between source and translated texts. This alignment is central to both the generation loop and the evaluation stage.

### Generation (`run_wmt.py`)

At each TMPC iteration the LLM produces multiple translation rollouts for a given document. To score these translations at the sentence level, TMPC needs to know which source sentence corresponds to which translated sentence. The pipeline works as follows:

1. Source text and each translation rollout are segmented into sentences (SpaCy).
2. Sentences are written to temporary files; `overlap.py` generates n-gram overlaps and `$LASER/tasks/embed/embed.sh` produces LASER sentence embeddings.
3. `vecalign.py` runs an adaptive-penalty alignment search to find the best source↔translation sentence mapping.
4. Aligned (source, translation) pairs are sent to the Reward Model for scoring.
5. High-scoring pairs enter the **subgoal buffer**, enabling Hindsight Subgoal Identification for subsequent iterations.

Temporary files (`*.overlaps`, `*.emb`, `*.txt`) are created inside a `{src}_{tgt}_temp/` folder in the working directory and **automatically deleted** after each document is processed. To inspect intermediate alignment results, comment out the `shutil.rmtree(...)` line at the end of `generate_windows()` in `run_wmt.py`.

### Evaluation (`segale_ctx.py`)

Translation quality is evaluated using the [SEGALE](https://github.com/nvlabs/SEGALE) framework, introduced in *"Extending Automatic Machine Translation Evaluation to Book-Length Documents"* ([Wang et al., EMNLP 2025](https://aclanthology.org/2025.emnlp-main.1645/)). SEGALE extends sentence-level metrics (e.g., COMET) to document-level evaluation by first aligning source, reference, and MT sentences via VecAlign, then scoring aligned sliding windows with context. This provides more robust evaluation than naive sentence-level scoring, especially when over- or under-translation is present.

In our implementation (`segale_ctx.py`), the evaluation pipeline:

1. Aligns source↔reference and source↔MT sentences using VecAlign with adaptive penalty search.
2. Finds common alignments across both alignment sets.
3. Constructs sliding windows (window size = 3) from the aligned sentences.
4. Computes COMET scores (`Unbabel/wmt22-comet-da`) on each window with context.
5. Aggregates per-document scores into an overall SEGALE<sub>comet</sub> score.

All temporary files in this stage are created in the system `/tmp/` directory using Python's `tempfile` module and are cleaned up automatically.

### Bundled VecAlign Components

The following files are bundled from VecAlign and are used by both `run_wmt.py` and `segale_ctx.py`:

| File | Role |
|------|------|
| `vecalign.py` | Main alignment algorithm (dynamic programming with adaptive penalty search) |
| `overlap.py` | Generates n-gram overlap sequences for embedding |
| `score.py` | Scoring functions for alignment quality |
| `dp_utils.py` | Dynamic programming utilities |
| `dp_core.pyx` | Cython-optimized DP core (compiled at runtime via `pyximport`) |

---

## Repository Structure

```
TMPC/
├── run_hhrlhf.py           # TMPC for HH-RLHF long-form generation
├── run_wmt.py              # TMPC for WMT document-level translation
├── run_mbpp.py             # TMPC for MBPP program synthesis
├── segale_ctx.py           # SEGALE context-level COMET evaluation
├── memory2csv.py           # Merge translation iteration results to CSV
├── vecalign.py             # VecAlign sentence alignment
├── overlap.py              # Overlap generation for VecAlign
├── score.py                # VecAlign alignment scoring
├── dp_utils.py             # Dynamic programming utilities for VecAlign
├── dp_core.pyx             # Cython-optimized DP core for VecAlign
├── dp_core.*.so            # Compiled Cython module (platform-specific)
├── requirements.txt        # Python dependencies
├── hhrlhf.csv              # HH-RLHF dataset
├── sanitized-mbpp.json     # MBPP sanitized dataset (auto-downloaded if absent)
├── wmt_zh_en.csv           # WMT zh→en translation dataset
└── scripts/
    ├── run_hhrlhf_full.sh  # Full HH-RLHF experiment
    ├── run_wmt_zh_en.sh    # Full WMT zh→en experiment
    └── run_mbpp_full.sh    # Full MBPP experiment
```

---

## Citation

If you find this work useful, please cite:

```bibtex
@inproceedings{
wang2026testtime,
title={Test-Time Alignment for Large Language Models via Textual Model Predictive Control},
author={Kuang-Da Wang and Teng-Ruei Chen and Yu Heng Hung and Guo-Xun Ko and Shuoyang Ding and Yueh-Hua Wu and Yu-Chiang Frank Wang and Chao-Han Huck Yang and Wen-Chih Peng and Ping-Chun Hsieh},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=DsS3xRPSs5}
}
```
