# Robust-U1: a multimodal model that restores a corrupted image before answering questions about it

> The ICML 2026 release from PolyX-Research adapts a 7B mixture-of-transformers VLM through three training stages so it can regenerate damaged visual content and then reason over what it recovered.

**PolyX-Research/Robust-U1** — 🚀🚀🚀 [ICML 2026] Official Implementation of Robust-U1: Can MLLMs Self-Recover Corrupted Visual Content for Robust Understanding?

- Repository: https://github.com/PolyX-Research/Robust-U1
- Website: https://huggingface.co/spaces/Jiaqi-hkust/Robust-U1
- Stars: 524 · Forks: 3
- Language: Python
- License: not declared
- Published: 2026-10-08 · Updated: 2026-10-08 · Language: en
- Canonical page: https://hysenlabs.com/projects/polyx-research-robust-u1

## Why recovery beats compensation for damaged images

The paper asks whether a multimodal large language model can recover corrupted visual content by itself, and the README states the two failures of existing work that make it ask. Feature-alignment methods are described as black-box alignment that lacks interpretability and never models the corruption process explicitly. Text-only compensation is the other failure: reasoning in text cannot bring back pixel-level detail that the image never delivered.

That framing decides the architecture. Instead of aligning features between a clean image and a degraded one, the model is trained to output the clean image, then to reason over its own output. A captioner working from a corrupted photo of a maths exercise has to guess digits it cannot see. A model that redraws the exercise first has something to read.

The repository is small and arranged around that pipeline. There are four directories, modeling, inference, rewards, and assets, plus demo.py and app.py at the root and a pyproject.toml alongside requirements.txt. It is not the training framework, and that is the single most important fact to know before you start.

## Three stages, one base model, two external frameworks

Training runs in three stages, and the model table shows what comes out of each one.

| Stage | Goal | Framework |
| --- | --- | --- |
| I. Visual Self-Recovery | Recover clean images from corrupted inputs, using SFT | MathCanvas |
| II. Visual Quality Alignment | Align recovery with pixel-level fidelity and semantics, using RL | Flow-GRPO |
| III. Multimodal Reasoning | Reason over corrupted and recovered images | MathCanvas |

Stage I adapts a base unified multimodal model, BAGEL-7B-MoT from ByteDance-Seed, to the recovery task. Stage II uses reinforcement learning to push the output toward structural fidelity and semantic agreement with the clean original. Stage III trains the model to reason over both the damaged input and the recovered image, which is what lets it be honest about what it recovered rather than only what it saw.

Four Hugging Face repositories correspond to the pipeline. BAGEL-7B-MoT supplies the starting weights, Jiaqi-hkust/`Robust-U1` is the final model for visual self-recovery and multimodal reasoning, `Robust-U1-SFT` is the Stage I checkpoint, and `Robust-U1-RL` is the Stage II checkpoint. There is no separate checkpoint for Stage III, so the final repository is the one that layers reasoning on top of reinforcement.

The two training frameworks are third-party. Stages I and III run on MathCanvas and Stage II runs on Flow-GRPO, so the `Robust-U1` contribution to the training side lives in the rewards directory, packaged so it can be registered directly in Flow-GRPO.

## An install that pulls in two different dependency lists

The install path is short and it installs more than you would expect from a repository with no releases.

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

The editable install pulls in a base set: numpy, Pillow, torch, pytorch-msssim, open-clip-torch, and safetensors. Three optional dependency groups carry the rest. The demo group adds accelerate, gradio, and huggingface-hub for the Gradio interface. The modeling group adds transformers, diffusers, einops, opencv-python-headless, sentencepiece, regex, and tqdm.

requirements.txt is a superset of pyproject.toml. The text file also pins torchvision, decord for video decoding, matplotlib, and further packages that the project file does not list. The README instructs installing both, so in practice you get the union rather than one authoritative list.

The project also declares two console entry points that mirror the two scripts at the repository root: `robust-u1-cli` points at the main function in demo, and `robust-u1-gradio` points at the main function in app. So the editable install gives you both commands, and if you would rather not rely on those shims you can call the scripts directly with python.

The Python requirement is not stated as a bare number in the README beyond the conda command, but pyproject.toml pins it: requires-python is 3.10 or newer, matching the conda environment the README creates.

## Running the CLI demo and the Gradio interface

Both entry points take a model path, which is the only argument that really matters. The CLI writes recovered images into an output directory and the Gradio app serves a browser interface on port 7860 by default.

```bash
export OUTPUT_DIR="./outputs"
```

The README writes the invocation across several lines with shell line continuations and an environment variable, so copying it literally means keeping the backslashes attached to the end of each line. Two environment variables carry the paths: one for the model directory and one for where recovered images land.



There is also a hosted demo on Hugging Face Spaces, linked from the README badges, which is the fastest way to see whether the recovery behaviour is what you are after before downloading several gigabytes of weights.

The CLI shape matters if you plan to batch. A model path and a directory of output images is a scriptable interface, and the recovered files are the useful artifact for any downstream comparison. The README does not describe batch input handling or a dataset flag, so looping over a corpus is something you would add yourself.

## Wiring the rewards into Flow-GRPO

Stage II is where this repository contributes most directly, and the README documents it as a small adapter registration rather than a training loop you run yourself. You clone Flow-GRPO, then in that project's rewards module you import the two `Robust-U1` names, register_flow_grpo_rewards and FLOW_GRPO_REFERENCE_REWARD_NAMES, and call the registration function with the framework's score_functions after it has built them.

```bash
git clone https://github.com/yifan123/flow_grpo.git
cd flow_grpo
```

Reference-based rewards in this design need clean target images, which makes sense for a restoration task: the recovered image gets scored against the original that was degraded. The pytorch-msssim dependency in both the base list and requirements.txt is the structural similarity side of that scoring, and open-clip-torch covers the semantic side.

For stages I and III the workflow is different and belongs to MathCanvas. You clone that repository, move into its BAGEL-Canvas directory, download BAGEL-7B-MoT, prepare your data, point data/dataset_info.py at your paths, and run the training scripts.

One instruction in that list is easy to misread. You run scripts/train/stage1.sh for the Stage I supervised fine-tuning, and then, after the Stage II reinforcement learning has finished, you run scripts/train/stage2.sh for the Stage III reasoning training. The script named stage2 is the third stage of the paper.

## A license field that does not match the badge

Three sources describe the license and one of them is empty. GitHub reports no license for this repository. pyproject.toml declares an MIT license in text form. The README carries an MIT badge pointing at opensource.org. There is no LICENSE file in the repository tree, which lists README.md, app.py, demo.py, pyproject.toml, requirements.txt, and the four directories.

So there are two facts to hold at once. The project presents itself as MIT in its own metadata and on its own README, and there is no license text in the repository and no license recorded on the hosting side. Both are consistent with an intention to be MIT and a failure to finish the paperwork. What you should not do is read the absence as permission.

If you need to settle it, read pyproject.toml for the declared intent, then check the repository root for a LICENSE file that is currently missing, and fall back to the OpenReview submission for whatever terms the paper itself carries.

Timeline is the last detail worth recording. The README news section notes acceptance to ICML 2026 on 2026-05-07 and the release of code, pretrained models, and the online demo on 2026-06-11, and the last push was 2026-06-12. There are no tagged releases, so there is no version line to pin and you reference a Hugging Face revision instead. The repository has 465 stars and a single open issue, which is what a paper drop looks like a few weeks after the fact.

## Conclusion

The idea is more transferable than the checkpoint. Teaching a vision model to regenerate pixels and then reason over the regeneration is something you could bolt onto an existing VLM, and this repository hands you the three pieces needed to try it: a supervised stage for recovery, a reinforcement stage for fidelity, and a reasoning stage that sees both the damaged and the restored image. Two things to check before committing time. The repository ships inference, modeling, and reward adapters, while the training scripts belong to MathCanvas and Flow-GRPO, so reproducing the paper means setting up someone else's framework. And the license is the least settled part, with GitHub reporting no license, pyproject declaring MIT, and the README displaying an MIT badge. Read pyproject.toml and confirm with the authors before building on it commercially.

## FAQ

### What does robust mean in machine learning?

In this project's usage it means holding up under real image degradations rather than clean inputs. The approach goes further than usual by having the model regenerate the corrupted visual content before answering questions about it, so the reasoning happens over recovered pixels rather than over a description of the damage.

### What base model does `Robust-U1` build on?

BAGEL-7B-MoT, a unified multimodal mixture-of-transformers model released by ByteDance-Seed on Hugging Face. It supplies the initial weights for the Stage I supervised fine-tuning, after which the pipeline adds the reinforcement alignment stage and the reasoning stage.

### What does the model actually restore?

Corrupted visual content rather than a text description of it. Stage I trains the model to recover a clean image from a degraded input, and Stage III then trains it to reason over both the corrupted and the recovered image, so it can use what it rebuilt and acknowledge what it could not.

### Can I train `Robust-U1` from this repository alone?

No, and the README is direct about it. This repository holds the inference, modeling, and reward adapters. Stages I and III run on MathCanvas and Stage II runs on Flow-GRPO, so a full reproduction means setting up one or both of those frameworks and preparing the corrupted-clean image pairs and reasoning annotations yourself.

## Sources

- [Issues](https://github.com/PolyX-Research/Robust-U1/issues)
- [PolyX-Research/Robust-U1 on GitHub](https://github.com/PolyX-Research/Robust-U1)
- [Project website](https://huggingface.co/spaces/Jiaqi-hkust/Robust-U1)
- [README](https://github.com/PolyX-Research/Robust-U1/blob/main/README.md)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/polyx-research-robust-u1