Skip to content

README.md

Latest commit

 

History

History
328 lines (239 loc) · 11.1 KB

README.md

File metadata and controls

328 lines (239 loc) · 11.1 KB
Raw

CVPR 2026

MM-ReCoder: Advancing Chart-to-Code Generation with Reinforcement Learning and Self-Correction

Zitian Tang, Xu Zhang, Jianbo Yuan, Yang Zou, Varad Gunjal, Songyao Jiang, Davide Modolo

Project Page arXiv HuggingFace HuggingFace

This repository is the implementation of the paper MM-ReCoder: Advancing Chart-to-Code Generation with Reinforcement Learning and Self-Correction.

MM-ReCoder is a 7B vision-language model that converts chart images into the matplotlib code that reproduces them. It is trained with multi-turn reinforcement learning, and at inference time it renders its own code with a sandboxed matplotlib tool, inspects the result, and self-corrects across turns.

This repository is built on top of verl; MM-ReCoder–specific scripts live under examples/mmrecoder/.

⚙️ Installation

Tested on Linux x86_64 (glibc ≥ 2.28), Python 3.10, NVIDIA driver supporting the CUDA 12.4 runtime.

Create an Environment

git clone https://github.com/ZitianTang/MMReCoder.git
cd MMReCoder
conda create -n mmrecoder python=3.10 -y
conda activate mmrecoder

# Pulls PyTorch 2.6 (cu124), vLLM 0.8.5.post1, and the pinned ML / Hydra /
# Ray stack the codebase was developed against.
pip install -r requirements.txt

# Install this repo as an editable package without re-resolving deps.
pip install -e . --no-deps

Install Flash Attention

To install flash-attn, run

pip install flash-attn==2.7.4.post1 --no-build-isolation

If you find it very slow, an alternative way is to install a pre-built flash-attn wheel. We recommend picking a wheel compatible with your machine from flash-attention-prebuild-wheels. For our setup (torch 2.6 / cu124 / cp310):

pip install https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.0.8/flash_attn-2.7.4.post1+cu124torch2.6-cp310-cp310-linux_x86_64.whl

📊 Inference on Chart-to-code Benchmarks

Data Preparation

The pre-processed .parquet files used by the inference scripts live on HuggingFace under cwbc/MM-ReCoder-Data. Please pick the benchmark(s) you want to run.

ChartMimic — chart images from the ChartMimic ICLR release.

# Download the raw ChartMimic images from the upstream release.
mkdir -p data/chartmimic
wget -P data/chartmimic \
    https://huggingface.co/datasets/ChartMimic/ChartMimic/resolve/main/dataset-iclr.tar.gz
tar -xzvf data/chartmimic/dataset-iclr.tar.gz -C data/chartmimic

# Download our pre-processed parquet (prompts + image references).
huggingface-cli download cwbc/MM-ReCoder-Data chartmimic_direct_600.parquet \
    --repo-type dataset --local-dir data/chartmimic

Plot2Code — chart images from TencentARC/Plot2Code.

# Download the raw Plot2Code release and lift out the matplotlib images.
mkdir -p data/plot2code
huggingface-cli download TencentARC/Plot2Code \
    --repo-type dataset --local-dir data/plot2code/Plot2Code
cp -r data/plot2code/Plot2Code/data/python_matplotlib data/plot2code/

# Download our pre-processed parquet (prompts + image references).
huggingface-cli download cwbc/MM-ReCoder-Data plot2code.parquet \
    --repo-type dataset --local-dir data/plot2code

ChartX — chart images from InternScience/ChartVLM.

# Download the raw ChartX release from Google Drive.
mkdir -p data/chartx
cd data/chartx
gdown 1d6zyH3kIwgepTqR0fc67xzyUtblrvOIX
unzip ChartX.zip
cd ../..

# ChartX images are large; resize them in place before inference.
python examples/data_preprocess/chartx_resize.py

# Download our pre-processed parquet (prompts + image references).
huggingface-cli download cwbc/MM-ReCoder-Data chartx.parquet \
    --repo-type dataset --local-dir data/chartx

After these steps the layout should look like:

data/
├── chartmimic/
│   ├── chartmimic_direct_600.parquet
│   └── dataset/...
├── plot2code/
│   ├── plot2code.parquet
│   └── python_matplotlib/...
└── chartx/
    ├── chartx.parquet
    └── ChartX/...

Model Download

Before running inference, download the MM-ReCoder checkpoint from HuggingFace:

hf download cwbc/MM-ReCoder

This caches the 7B checkpoint under your local HuggingFace cache; the inference scripts will then load it by name (model.path=cwbc/MM-ReCoder) without re-downloading.

Run a Benchmark

All inference scripts live in examples/mmrecoder/inference/. Outputs are written to generations/mmrecoder_<benchmark>_<variant>.json as a list of {qid, input, outputs} records.

Script Benchmark Turns Default N_GPUS
chartmimic_1turn.sh ChartMimic 1 (no self-correction) 1
chartmimic_2turns.sh ChartMimic 2 8
chartmimic_4turns.sh ChartMimic 4 8
plot2code_4turns.sh Plot2Code 4 8
chartx_4turns.sh ChartX 4 8

The multi-turn scripts are sized for the paper's setup of 8 × 80 GB GPUs (e.g. H100 / A100 80 GB). They will run on as few as 2 × 80 GB GPUs — just edit N_GPUS at the top of the script. Throughput scales roughly linearly with the GPU count (vLLM is replicated per GPU).

To run, e.g., the 2-turn ChartMimic benchmark:

bash examples/mmrecoder/inference/chartmimic_2turns.sh

Single-turn Inference

chartmimic_1turn.sh is a one-shot, no-self-correction baseline: the model emits code once, with no rendering or self-correction step. It fits on a single 80 GB GPU by offloading FSDP parameters to CPU during vLLM init:

bash examples/mmrecoder/inference/chartmimic_1turn.sh

Visualize Generated Charts

To render the matplotlib code that the model emitted at each turn into PNG images, run:

bash scripts/visualization/visualize_generation.sh

The script reads a generation JSON (default generations/mmrecoder_chartmimic_2turns.json), extracts the python code block from every assistant turn across every rollout, and renders each block to {save_dir}/images/{qid}_{rollout_idx}_{turn_idx}.png in parallel via Ray. Edit GENERATION_JSON and SAVE_DIR at the top of the script to point at a different run.

🏋️ RL Training

MM-ReCoder is trained with multi-turn reinforcement learning (GRPO) on top of the SFT cold-start checkpoint cwbc/MM-ReCoder-SFT-Cold-Start. RL training runs in two stages:

  1. Stage 1 (stage1-shared-first-turn.sh) — multi-turn GRPO with a shared first turn across rollouts; warms up the self-correction behavior.
  2. Stage 2 (stage2-full-trajectory.sh) — multi-turn GRPO on the full trajectory; resumes from the stage 1 checkpoint.

Data Preparation

RL training uses Chart2Code-160k as the training data. Download the raw images from the upstream dataset and our pre-processed parquet:

# Download the raw Chart2Code-160k images from the upstream HuggingFace dataset.
mkdir -p data/chart2code_160k
wget -P data/chart2code_160k \
    https://huggingface.co/datasets/xxxllz/Chart2Code-160k/resolve/main/images.tar.gz
tar -xzvf data/chart2code_160k/images.tar.gz -C data/chart2code_160k

# Download our pre-processed parquet (prompts + image references).
huggingface-cli download cwbc/MM-ReCoder-Data chart2code_160k.parquet \
    --repo-type dataset --local-dir data/chart2code_160k

After this step the layout should look like:

data/chart2code_160k/
├── chart2code_160k.parquet
└── images/...

Launch the LLM-as-a-Judge Reward Server

Part of the reward signal during RL is computed by an LLM-as-a-judge served via a separate vLLM HTTP endpoint. Two options are provided under scripts/vllm_servers/:

# Option A (paper setup): Qwen2.5-VL-72B-Instruct as judge.
bash scripts/vllm_servers/vllm_server_72b.sh

# Option B (lighter): Qwen2.5-VL-7B-Instruct as judge,
bash scripts/vllm_servers/vllm_server_7b.sh

Both scripts expose the judge on port 18967 under served name judge. Launch the server on a separate node (or separate GPUs from the trainer), note its IP, and then set LLM_AS_A_JUDGE_BASE at the top of each training script to point at it:

# In examples/mmrecoder/train/*.sh, replace <VLLM_SERVER_IP>:
export LLM_AS_A_JUDGE_BASE="http://<VLLM_SERVER_IP>:18967/v1"

# Multiple judge servers can be load-balanced by comma-separating addresses:
export LLM_AS_A_JUDGE_BASE="http://host1:18967/v1,http://host2:18967/v1"

If you do not want to use the LLM-as-a-judge reward at all, you can skip the server entirely by deleting the following line from each training script:

+reward_model.reward_kwargs.reward_weights.chart_derendering.model_reward=0.1 \

The remaining rule-based reward components (format, color, text, layout, type) will still be used.

Multi-Turn RL Training

The two-stage multi-turn RL recipe used in the paper. Scripts live in examples/mmrecoder/train/:

# Stage 1: Shared-first-turn optimization.
bash examples/mmrecoder/train/stage1-shared-first-turn.sh

# Stage 2: Full-trajectory optimization (resumes from stage 1).
bash examples/mmrecoder/train/stage2-full-trajectory.sh

Single-Turn RL Training

We also provide a single-turn RL recipe that trains the model to emit code in one shot, with no self-correction:

bash examples/mmrecoder/train/single-turn-rl.sh

👍 Acknowledgements

This repository is built on top of verl and DeepEyes.

📑 Citation

@inproceedings{tang2026mmrecoder,
    title={MM-ReCoder: Advancing Chart-to-Code Generation with Reinforcement Learning and Self-Correction},
    author={Zitian Tang and Xu Zhang and Jianbo Yuan and Yang Zou and Varad Gunjal and Songyao Jiang and Davide Modolo},
    booktitle={CVPR},
    year={2026}
}

🔒 LICENSE

MM-ReCoder is released under the CC BY-NC 4.0 License.