Look Before You Restore: Learning to Predict Before Acting for Agentic Image Restoration

Xu Zhang; Jun Wan; Lefei Zhang; Liangpei Zhang

cs.CVOct 14, 2026

Look Before You Restore: Learning to Predict Before Acting for Agentic Image Restoration

Authors: Xu Zhang, Jun Wan, Lefei Zhang, Liangpei Zhang

Organizations: Wuhan University · Zhongnan University of Economics and Law

Abstract

Agents for mixed-degradation image restoration typically learn what a tool does to an image only after running it: failed calls waste computation and harmful outputs go unnoticed until too late. We propose PreActIR, an agentic restoration framework that predicts tool consequences before execution. For each candidate tool, a tool-conditioned world model forecasts changes in residual degradation, image quality, and fidelity, together with the risk of harm and the uncertainty of the prediction. The agent executes the top-ranked candidate, and an independent verifier commits or rolls back its output before the agent replans from the observed state. PreActIR also audits whether performance is limited by the selector or by the candidate space, and studies targeted tool adaptation as a controlled way to expand that space. In a controlled evaluation that fixes the candidate pool, verifier, and budget and employs reference-annotated degradation states, PreActIR uses 42.5% fewer real tool calls and 63.4% fewer rejected calls than reactive, execute-then-evaluate trial and error, while improving Y-PSNR by 0.92 dB.

Figures & tables

Figure 1: Motivation of PreActIR. Existing agentic restoration is largely reactive: tool effects are observed only after execution, while decisions remain restricted to the available action space. PreActIR predicts action consequences before execution.

Figure 2: Overview of PreActIR. The world model predicts action consequences before execution, while the predict–act–verify controller executes and verifies the selected action with commit or rollback. A tool-evolution loop performs candidate-frontier audit and targeted tool adaptation.

Figure 3: An illustrative predict–act–verify trajectory under haze and rain.

Group	Type	Method	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	MANIQA $\uparrow$	CLIP-IQA $\uparrow$	MUSIQ $\uparrow$
A	All-in-One	AirNet CVPR’22	19.13	0.602	0.428	0.258	0.393	42.46
		PromptIR NeurIPS’23	20.06	0.609	0.413	0.263	0.401	42.62
		MiOIR arXiv’24	20.84	0.656	0.372	0.245	0.393	47.82
		DA-CLIP ICLR’24	19.58	0.603	0.427	0.242	0.414	42.51
		InstructIR ECCV’24	18.03	0.575	0.443	0.266	0.353	45.77
		AutoDIR ECCV’24	19.64	0.629	0.397	0.250	0.377	47.01

Table 1: Quantitative comparison on the three compound-degradation subsets of MiO100. Best and second-best results are marked in bold and underlined, respectively.

Method	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	MANIQA $\uparrow$	CLIP-IQA $\uparrow$	MUSIQ $\uparrow$
AirNet	23.31	0.747	0.448	0.236	0.327	35.27
PromptIR	23.86	0.754	0.434	0.241	0.327	35.48
MiOIR	24.00	0.747	0.424	0.239	0.351	36.85
DA-CLIP	23.73	0.744	0.427	0.251	0.333	36.41
AutoDIR	20.75	0.664	0.479	0.247	0.315	42.68
AgenticIR	23.93	0.721	0.373	0.315	0.451	52.91

Table 2: Quantitative comparison on the paired real-world 100-image set used in MAIR.

Figure 4: Qualitative comparison on representative examples from MiO100 Groups A–C.

Figure 5: Qualitative comparison on a representative example from the paired real-world dataset.

Method	HR (%) $\uparrow$	Harm (%) $\downarrow$	Severe (%) $\downarrow$	$\Delta$ PSNR Y $\uparrow$	$\Delta$ SSIM Y $\uparrow$	$\Delta$ LPIPS $\downarrow$	$r_{\rm s}\uparrow$
Random	1.1	55.7	28.2	$-$ 0.10	0.015	$-$ 0.038	–
Tool prior	31.9	40.6	13.8	0.98	0.058	$-$ 0.067	0.284
Degradation routing	36.6	35.6	15.0	1.27	0.052	$-$ 0.079	0.378
Direct scorer	36.7	48.8	27.5	1.26	0.053	$-$ 0.056	0.247
World model (WM)	34.6	42.5	21.9	1.14	0.055	$-$ 0.069	0.309
WM-F	31.9	46.2	23.1	1.05	0.051	$-$ 0.057	0.314

Table 3: Open-loop top-1 selection on 160 OOF states with 2,147 executed candidates. HR: share of oracle headroom recovered. Harm: the selected action lowers Y-PSNR by more than 0.05 dB, Y-SSIM by more than 0.005, or mean proxy quality. Severe: Y-PSNR drop

\geq

1 dB or Y-SSIM drop

\geq

0.05.

r_{\rm s}

: mean within-state Spearman correlation of scores with realized

\Delta F_{Y}

. WM-F: fidelity head alone.

Figure 6: Risk information from pre-execution prediction; brackets and bands are 95% clean-identity cluster-bootstrap intervals. (a) ROC for detecting harmful actions among 2,147 candidates. Paired AUC gains of

p^{\rm harm}

: 0.118

[0.091,0.146]

over Direct and 0.073

[0.054,0.093]

over WM-F. (b) HR versus harmful submissions when the top action may be replaced by a no-op. Lines sweep an abstention threshold on predicted harm (WM-

\tau

) or gain (Direct-

\tau

), with bands shown where all draws have support; dots: top-1 results of Table 3 ; dashed lines: 5/10/15% harm.

Policy	Calls $\downarrow$	Rejects $\downarrow$	Time (s) $\downarrow$	PSNR Y $\uparrow$	SSIM Y $\uparrow$	LPIPS $\downarrow$	Harm (%) $\downarrow$
Reactive trial-and-error	2.97	2.36	39.6	21.16	0.676	0.422	15.6
Fidelity-head ranking	1.75	0.90	25.8	22.27	0.689	0.420	18.1
PreActIR (structured)	1.71	0.86	25.3	22.08	0.690	0.416	14.4

Table 4: Matched closed-loop comparison on the 160 OOF states with a shared candidate pool, verifier, and budget

B=4

. Calls and rejects are per image; time counts restoration-tool execution only. Rejects, PreActIR minus reactive:

-1.49

per image, CI

[-1.78,-1.21]

Figure 7: Group C frontier audit (post hoc, using test references). (a) Attainable mean Y-PSNR vs. MUSIQ for the fixed and expanded pools (dashed); gray: deployed fixed-pool system; red: final system, built on a different fixed-pool configuration (Sec. 4.5 ); hollow: baselines of Table 1 . (b) Per-combination Y-PSNR (LR: low resolution); oracle at the fixed system’s MUSIQ level.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Policy	PSNR ${}_{Y}\uparrow$	SSIM ${}_{Y}\uparrow$	LPIPS $\downarrow$	$\Delta F_{Y}\uparrow$	Harm (%) $\downarrow$	Regret $\downarrow$
Degradation routing	20.7006	0.6593	0.3974	1.1212	48.1	3.9618
Direct scorer	22.6932	0.6883	0.4174	3.4042	21.2	1.6788
World model ( $H=1$ )	22.5510	0.6958	0.4045	3.3364	15.6	1.7466
World model ( $H=2$ )	22.4985	0.6910	0.4094	3.2361	17.5	1.8469
w/o state rollout	22.5285	0.6938	0.4086	3.2939	16.2	1.7891
+ harm/severe penalties	22.5604	0.6909	0.4050	3.2974	18.8	1.7857

Appendix

Table 5: Controlled pre-action ranking ablation on the 160 OOF states. Every policy ranks the shared frozen pool of the closed-loop replay (Sec. 4.4 : annotated degradation states, 1,415 cached candidates) and executes exactly one real action per state;

H=2

uses only virtual lookahead. Harm rates are therefore not comparable with Table 3 , which uses estimated beliefs and the 2,147-candidate pool.

\Delta F_{Y}=\Delta\mathrm{PSNR}_{Y}+10\,\Delta\mathrm{SSIM}_{Y}

is the realized fidelity gain, and Regret is the oracle

\Delta F_{Y}

minus the realized one.

Figure 8: Additional qualitative comparisons on MiO100 Groups A–C.

Metric	World model $-$ Direct	WM-F $-$ Direct
HR	$-$ 2.1 [ $-$ 15.5, 11.8]	$-$ 4.8 [ $-$ 15.1, 6.0]
Harm	$-$ 6.2 [ $-$ 13.8, 1.2]	$-$ 2.5 [ $-$ 10.0, 5.0]
Severe	$-$ 5.6 [ $-$ 16.2, 4.4]	$-$ 4.4 [ $-$ 13.8, 5.0]
$\Delta$ PSNR Y (dB)	$-$ 0.13 [ $-$ 0.68, 0.41]	$-$ 0.21 [ $-$ 0.65, 0.23]
$\Delta$ SSIM Y	0.002 [ $-$ 0.013, 0.019]	$-$ 0.002 [ $-$ 0.017, 0.012]
$\Delta$ LPIPS	$-$ 0.013 [ $-$ 0.033, 0.006]	$-$ 0.001 [ $-$ 0.018, 0.018]

Appendix

Table 6: Paired differences to the Direct scorer with 95% clean-identity cluster-bootstrap intervals. Rates are in percentage points.

Utility	HR (%)	Harm (%)	Severe (%)	$\Delta$ PSNR Y	$\Delta$ SSIM Y	$\Delta$ LPIPS	$r_{\rm s}$
Full	34.6	42.5	21.9	1.14	0.055	$-$ 0.069	0.309
w/o $v$	$-$ 14.5	58.1	35.6	$-$ 0.66	$-$ 0.005	$-$ 0.019	$-$ 0.092
w/o $p^{\rm ben}$	32.4	46.2	21.9	1.08	0.050	$-$ 0.058	0.310
w/o $\xi$	36.1	41.9	20.6	1.17	0.060	$-$ 0.072	0.311
w/o $g$ , $q$ , $d$ , $u$ , or $C$	34.6	42.5	21.9	1.14	0.055	$-$ 0.069	0.309–0.312

Appendix

Table 7: Utility-term ablation (one term removed at a time, other coefficients frozen). Terms follow Eq. ( 8 ):

v

fidelity gain,

p^{\rm ben}

benefit probability,

\xi

non-target side effects,

g

target gain,

q

proxy-quality gain,

d

content damage,

u

uncertainty,

C

cost prior.

Figure 9: Supplementary open-loop diagnostics on the 160 OOF states. (a) Share of positive oracle headroom recovered by an ideal selector over the top-

k

ranked candidates; the dotted line marks 90%. (b) Reliability of the world model’s harm probability in ten equal-count bins with Wilson 95% intervals.

Budget	PSNR ${}_{Y}\uparrow$	SSIM ${}_{Y}\uparrow$	LPIPS $\downarrow$	Harm (%) $\downarrow$	No commit (%) $\downarrow$	Calls $\downarrow$	Time (s) $\downarrow$
$B=1$	21.88	0.683	0.423	8.8	31.2	1.00	18.5
$B=2$	22.02	0.687	0.419	11.2	21.2	1.31	21.0
$B=4$	22.08	0.690	0.416	14.4	15.6	1.71	25.3

Appendix

Table 8: Sensitivity of PreActIR to the real-execution budget on the 160 OOF states, with the candidate pool and verifier fixed. Harm counts no-commit outcomes as non-harmful; time counts restoration-tool execution only. The

B=4

row is the PreActIR row of Table 4 .

Figure 10: Quality versus real tool calls on the 160 OOF states with a shared candidate pool. Top: learned verifier; bottom: ground-truth verifier that commits exactly the candidates labeled beneficial (Appendix A.2 ). Points from left to right correspond to

B=1,2,3,4,6

; axes are shared across rows. (a,c) Y-PSNR gain over the input. (b,d) Share of states with no committed candidate within budget. Error bars are 95% clean-identity cluster-bootstrap intervals; the oracle order uses realized outcomes and is shown only as a reference.

Policy	Accepted	Beneficial accepted	Harmful accepted
Reactive	47/160	43/77 (55.8%)	4/78 (5.1%)
Routing	54/160	46/81 (56.8%)	8/74 (10.8%)
Fidelity-head	107/160	88/122 (72.1%)	19/34 (55.9%)
PreActIR	110/160	94/130 (72.3%)	16/26 (61.5%)
Oracle order	98/160	97/150 (64.7%)	1/7 (14.3%)

Appendix

Table 9: Learned-verifier decisions on each policy’s first choice at

B=1

, using the cached ground-truth labels.

Agent-Based Image Restoration

See all 12 papers

TIR-Agent: Training an Explorative and Efficient Agent for Image Restoration
Mar 29, 2026Guoli Jia, Yisheng Zhang, Haote Hu +11Tool-Using Vision-Language Agents Image Restoration
Tsinghua University, Beijing, China · Data Science & Artificial Intelligence Research Institute, China Unicom, Beijing, China · Hunan University, Changsha, China +2
SkillIR: Evolving Scene-Aware Skills for Agentic Image Restoration
Sep 18, 2026Jie Shao, Shengkai Hu, Xu Zhang +4Agent-Based Image Restoration Image Restoration
Zhongnan University of Economics and Law, Wuhan, China
OPERA: An Agent for Image Restoration with End-to-End Joint Planning-Execution Optimization
May 21, 2026Feng Zhu, Shuyang Xie, Yihan Zeng +2Agent-Based Image Restoration Image Restoration
Harbin Institute of Technology · Huawei Noah’s Ark Lab

Mar 29, 2026J/K move · Enter open · S save

TIR-Agent: Training an Explorative and Efficient Agent for Image Restoration

Guoli Jia, Yisheng Zhang, Haote Hu, Shanxu Zhao+10

Tsinghua University, Beijing, China · Data Science & Artificial Intelligence Research Institute, China Unicom, Beijing, China · Hunan University, Changsha, China+2

Tool-Using Vision-Language Agents Image Restoration RL for VLMs Agent-Based Image Restoration

Open paper AskPDF arXiv

Vision-language agents that orchestrate specialized tools for image restoration (IR) have emerged as a promising method, yet most existing frameworks operate in a training-free manner. They rely on heuristic task scheduling and exhaustive tool traversal, resulting in sub-optimal restoration paths and prohibitive computational cost. We argue that the core bottleneck lies in the absence of a learned policy to make decision, as a vision-language model cannot efficiently handle degradation-aware task ordering and tool composition. To this end, we propose TIR-Agent, a trainable image restoration agent that performs a direct tool-calling policy through a two-stage training pipeline of supervised fine-tuning (SFT) followed by reinforcement learning (RL). Two key designs underpin effective RL training: (i) a random perturbation strategy applied to the SFT data, which broadens the policy's exploration over task schedules and tool compositions, and (ii) a multi-dimensional adaptive reward mechanism that dynamically re-weights heterogeneous image quality metrics to mitigate reward hacking. To support high-throughput, asynchronous GPU-based tool invocation during training, we further develop a globally shared model-call pool. Experiments on both in-domain and out-of-domain degradations show that TIR-Agent outperforms 12 baselines, including 6 all-in-one models, 3 training-free agents, and 3 proprietary models, and achieves over 2.5

\times

inference speedup by eliminating redundant tool executions.

Abstract

Agent-Based Image Restoration

Explore similar work

TIR-Agent: Training an Explorative and Efficient Agent for Image Restoration

SkillIR: Evolving Scene-Aware Skills for Agentic Image Restoration

OPERA: An Agent for Image Restoration with End-to-End Joint Planning-Execution Optimization