RL-AWB:基于深度强化学习的低光夜间场景自动白平衡校正 RL-AWB: Deep Reinforcement Learning for Auto White Balance Correction in Low-Light Night-time Scenes
首个将强化学习应用于夜间白平衡的框架,用5张训练图片实现跨传感器泛化。
前置知识
自动白平衡(Auto White Balance, AWB)
自动白平衡是相机图像信号处理(ISP)流水线中的核心模块,其目标是估计场景的光照色温并对图像进行色彩校正,使得白色物体在不同光照条件下看起来仍为中性白色。AWB 通常通过估计场景的光照颜色(illuminant)来对原始图像的 R、G、B 通道施加校正增益。传统方法包括灰度世界(Gray World)假设、最大 RGB 响应(Max-RGB)以及基于边缘和统计的方法。在日间场景中,这些方法表现良好;但在夜间低光场景中,传感器噪声大、光照复杂,传统统计假设失效,导致 AWB 性能严重下降。
本文的核心任务就是改进夜间场景的 AWB,理解这一基本模块是阅读本文的前提。
色彩恒常性(Color Constancy)
色彩恒常性是人类视觉系统的一种重要能力——无论光源色温如何变化,人眼都能感知物体的固有颜色。在计算摄影中,色彩恒常性算法试图模拟这种能力,通过估计场景光照来消除色偏。评估指标通常使用角度误差(angular error),即估计的光照方向与真实光照方向之间的夹角(单位:度)。角度越小,估计越准确。数学上,角度误差定义为 $E = \arccos\left(\frac{\langle\mathbf{e}_{\text{est}}, \mathbf{e}_{\text{gt}}\rangle}{\|\mathbf{e}_{\text{est}}\|\|\mathbf{e}_{\text{gt}}\|}\right)$,其中 $\mathbf{e}_{\text{est}}$ 和 $\mathbf{e}_{\text{gt}}$ 分别是估计和真实的光照向量。
角度误差是本文所有实验的核心评估指标,理解该指标才能解读实验结果的含义。
强化学习(Reinforcement Learning, RL)
强化学习是一种机器学习范式,智能体(agent)通过与环境(environment)交互来学习最优策略。在每一步,智能体观察环境状态(state),选择一个动作(action),获得奖励(reward),并转移到新状态。本文采用的 Soft Actor-Critic(SAC)算法是一种最大熵强化学习方法,其特点是在最大化累积奖励的同时最大化策略的熵,鼓励探索。SAC 使用两个 Q 值网络(twin critics)来减少过估计偏差,并支持自动熵调节。SAC 的目标函数为 $J(\pi) = \sum_t \mathbb{E}_{(s_t,a_t)\sim\rho_\pi}[r(s_t,a_t) + \alpha \mathcal{H}(\pi(\cdot|s_t))]$,其中 $\alpha$ 是温度参数,$\mathcal{H}$ 是策略的熵。
本文将 AWB 问题建模为序列决策问题并用 RL 求解,理解 RL 的基本概念(状态、动作、奖励、策略)是理解 RL-AWB 框架的关键。
灰度像素检测(Gray Pixel Detection)
灰度像素检测是统计色彩恒常性方法的核心技术,其基本假设是:场景中存在足够多的中性灰(achromatic)表面,它们在三个颜色通道上的响应应该相等或接近。通过识别这些灰度像素并计算它们的颜色比,可以估计场景的光照颜色。常用方法包括灰度世界假设(所有像素的平均即为光照)和更先进的显著灰度像素检测(如 MSGP、GI、RGP 等)。本文提出的 SGP-LRD 算法在此基础上引入了局部反射率差异和双重过滤机制来提升夜间场景的鲁棒性。
SGP-LRD 是本文提出的统计基础算法,理解灰度像素检测原理才能理解 RL-AWB 如何对该算法进行参数调优。
课程学习(Curriculum Learning)
课程学习是一种训练策略,灵感来自人类从易到难的学习过程。在 RL 中,课程学习通过逐步增加任务难度来提升训练效率和稳定性。本文采用两阶段课程学习:第一阶段(Stage 1)在单张图片上训练,让智能体学会基本的参数调整和收敛检测;第二阶段(Stage 2)在多张图片的课程池中循环训练,让智能体学习跨场景的泛化策略。课程池大小 $M=5$ 是通过消融实验确定的最优值。
课程学习是本文训练策略的核心组成部分,两阶段设计直接决定了方法的训练效率和泛化能力。
Minkowski 范数(Minkowski Norm)
Minkowski 范数是一种广义的向量范数,定义为 $\|x\|_p = (\sum_i |x_i|^p)^{1/p}$。在色彩恒常性中,Minkowski 范数的指数 $p$ 控制着光照估计时对不同像素的加权方式:较小的 $p$ 趋向均匀加权,较大的 $p$ 强调高置信度像素的贡献。$p=1$ 对应中位数,$p \to \infty$ 对应最大值。本文将 $p$ 作为 RL 智能体的一个可调参数,让智能体根据场景特性自动选择最优的 $p$ 值。
$p$ 是 RL-AWB 框架中两个关键可调参数之一,直接影响光照估计的精度。
跨传感器泛化(Cross-Sensor Generalization)
不同相机传感器具有不同的光谱响应特性(spectral response)和 ISP 配置,导致同一算法在不同传感器上的表现可能差异巨大。跨传感器泛化是指算法在某个传感器上训练后,能否直接部署到未见过的传感器上而无需重新训练或微调。这是色彩恒常性领域的核心挑战之一,因为标注数据通常只来自特定传感器。本文通过 (1) 传感器无关的统计算法和 (2) 课程学习暴露智能体于多样条件来实现跨传感器泛化。
跨传感器泛化是本文的核心目标和主要卖点之一,也是评判方法实用性的关键指标。
研究动机
夜间自动白平衡面临根本性的技术挑战,与日间场景形成鲜明对比。在低光环境下,传感器需要使用高 ISO 设置(本文 LEVI 数据集中 ISO 范围为 500 到 16,000),导致严重的色度噪声(chromatic noise)主导了信号,使得传统统计方法(如 Gray World、Max-RGB)依赖的颜色统计假设失效。具体而言,现有统计方法假设场景中存在足够多的灰度像素且其检测稳定,但在极端低光条件下这些假设被严重违反。深度学习方法虽然在日间场景有效,但需要大量标注的夜间训练数据,且在部署到未见过的相机传感器时泛化能力严重退化——如论文 Table 2 所示,FC4 在 NCC→LEVI 跨数据集评估中,中位角度误差飙升至 11.8°,PCC 更是达到 11.1°,GCC 的最差 25% 误差甚至高达 90.0°。此外,夜间场景对算法参数高度敏感,微小的参数变化就可能导致截然不同的光照估计结果。在跨传感器部署时,同一模型往往产生显著且不可预测的色偏,这在移动摄影、监控系统和汽车成像等实际应用中是不可接受的。
本文的目标是本文的目标是设计一个能够在夜间低光场景下实现鲁棒自动白平衡的框架,同时满足三个关键需求:第一,实现跨传感器泛化,即在一个传感器上训练后能直接部署到未见过的传感器上;第二,实现极高的数据效率,仅需极少的训练样本(目标是每个数据集仅需 5 张图片);第三,保持方法的可解释性,利用统计方法的物理意义而非纯粹的黑箱深度学习。具体性能目标是在 NCC 和 LEVI 两个夜间数据集上达到或超越现有统计方法和深度学习方法的精度,同时在跨数据集评估中保持稳定性能,避免现有学习方法普遍存在的严重性能退化。
与已有工作不同的是,本文的独特切入角度在于将自动白平衡问题从传统的回归问题(直接估计光照 RGB 值)重新建模为序列决策问题(sequential decision-making)。这一视角转换的核心洞察是:与其让神经网络直接预测光照颜色(如 FC4、C4 等方法),不如让 RL 智能体学会如何调整一个可解释的统计算法的参数。这种混合设计的巧妙之处在于:(1) 光照估计仍然由传感器无关的统计算法完成,天然具有跨传感器泛化能力;(2) RL 智能体只需学习参数调整策略,动作空间极小(仅 2 个连续参数),因此训练数据需求极低;(3) 整个过程保留了统计方法的可解释性。这是首次将强化学习应用于色彩恒常性领域,也是首次将 AWB 建模为每张图片的自适应参数优化问题。
核心方法
RL-AWB 框架的整体思路可以用一个直觉来理解:想象一位专业的 AWB 工程师在调试夜间照片的白平衡参数。他不会直接计算光照颜色,而是先用一个统计算法(SGP-LRD)得到一个初始估计,然后观察结果,逐步调整两个关键参数——灰度像素选择比例 $N$ 和 Minkowski 指数 $p$——直到获得满意的校正效果。RL-AWB 正是将这个人工调参过程自动化:SGP-LRD 算法充当'调参对象',SAC 强化学习智能体充当'专业工程师'。技术路线分为三个层次:底层是 SGP-LRD 统计算法,它通过显著灰度像素检测、局部反射率差异和双重过滤机制来估计场景光照;中层是 SAC 智能体,它观察图像的颜色统计特征和当前参数状态,输出参数调整量;顶层是两阶段课程学习,从单图训练到多图循环训练,逐步提升智能体的泛化能力。这种'统计引擎 + RL 调参器'的分层设计是整个框架的核心架构。
RL-AWB 的核心创新在于将 AWB 从'回归光照值'转变为'优化算法参数'。与现有深度学习方法(如 FC4 直接用 CNN 回归光照 RGB,C4/C5 使用级联网络,FFCC 使用全卷积色彩恒常性)的本质区别在于:这些方法直接用神经网络拟合从图像到光照的映射,因此对训练数据分布高度敏感,跨传感器时泛化能力差。RL-AWB 则只用神经网络学习参数调整策略,光照估计仍由统计算法完成。这带来三重优势:(1) 统计算法的传感器无关性天然保证了跨传感器泛化;(2) RL 的动作空间仅为 2 维($N$ 和 $p$ 的调整量),远小于直接回归 3 维光照值,因此训练效率极高;(3) SAC 算法的 off-policy 特性和经验回放机制使得即使在极少训练样本下也能高效学习。另一个关键创新是 SGP-LRD 算法本身,它通过局部窗口的重叠设计实现了隐式噪声过滤——传感器噪声产生的虚假像素只在有限窗口中出现且贡献很小,而真正的灰度区域在相邻窗口中呈现一致响应,通过空间冗余自然区分信号与噪声。
方法步骤详情
RL-AWB 的工作流程包含以下完整步骤:首先,对输入的原始 RAW 图像进行 RGB-uv 对数色度直方图计算(粒度 $m=60$,维度 $3m^2=10,800$),并进行 $\ell_1$ 归一化和逐元素平方根变换,得到光照相关特征 $s_{WB}$。同时,维护一个长度为 11 的历史向量 $h_s$,编码最近的参数值($N$ 和 $p$)和归一化时间步计数器,作为参数相关特征。完整状态为 $s_t = (s_{WB}, h_s)$。SAC 智能体的 Actor 网络采用双分支 MLP 编码器,将 $s_{WB}$ 和 $h_s$ 分别映射为 64 维嵌入 $z_{WB}$ 和 $z_{hist}$,融合后输出均值 $\mu$ 和对数方差 $\log\sigma^2 \in \mathbb{R}^2$,通过重参数化采样得到两个连续动作。动作经过 tanh 压缩到 $[-1, 1]$ 后重新缩放到有效范围($a^{(N)}_t \in [-0.6, 0.6]$,$a^{(p)}_t \in [-4, 4]$)。新参数 $\text{param}_{t+1} = \text{param}_t + a_t$ 被送入 SGP-LRD 算法,该算法依次执行:(1) 灰度性度量 $G(x,y)$ 计算,选取前 $N\%$ 像素;(2) 局部方差过滤,剔除方差过小的不可靠像素;(3) 颜色偏差过滤,剔除偏离场景主导色偏过远的像素;(4) 灰度像素置信度加权;(5) 逐像素局部反射率差异计算;(6) 基于 Minkowski 范数 $p$ 的光照估计。奖励设计采用相对误差改进 $R_{\text{err}} = (E_0 - E_t)/(E_0 + (E_0/c_1)^\alpha)$,加上动作成本惩罚 $R_{\text{act}}$ 和终止奖励 $R_\rho$。当连续三步估计稳定时终止,进入下一轮参数调整。课程学习分两阶段:Stage 1 在单张图片上训练直至收敛;Stage 2 在 5 张图片的课程池中循环训练,每张图片连续 5 个 episode 后切换。
技术新颖性
RL-AWB 的技术新颖性体现在多个层面。首先,这是强化学习在色彩恒常性领域的首次应用,开辟了一个全新的研究方向。此前 RL 在 ISP 中的应用(如 RL-Restore 用 DQN 进行图像恢复工具链选择,RL-AE 控制相机曝光)都聚焦于日间场景且动作空间面向相机设置,而 RL-AWB 面向夜间场景且动作空间面向统计算法参数。其次,SGP-LRD 算法本身具有技术新颖性:它将显著灰度像素检测与局部反射率差异(LRD)相结合,通过 3x3 窗口的重叠设计实现隐式噪声过滤和空间先验利用——这是对传统灰度像素检测方法的重要改进。双重过滤机制(局部方差过滤和颜色偏差过滤)专门针对夜间噪声大的特点设计。第三,奖励设计中的难度感知松弛机制(difficulty-aware relaxation)$R_{\text{step}} = R_{\text{err}} + (1 - E_0/c_2) \times R_{\text{act}}$ 是一个巧妙设计:初始误差大的图片自动减轻动作惩罚,避免在困难样本上过度约束探索。第四,LEVI 数据集是第一个多相机夜间色彩恒常性数据集(700 张图片,iPhone 16 Pro 和 Sony ILCE-6400 两个传感器),填补了该领域跨传感器评估基础设施的空白。
实验结果
实验结果全面验证了 RL-AWB 框架的有效性。在统计方法对比中(Table 1),SGP-LRD 在 NCC 数据集上取得最低中位误差 2.12°(对比次优 RGP 的 2.22°)和最低均值误差 3.11°(对比 RGP 的 3.33°),在 LEVI 数据集上同样领先(中位 3.08° vs RGP 3.21°)。这证明了 SGP-LRD 作为基础算法的优越性。在跨数据集评估中(Table 2),RL-AWB 展现出惊人的泛化能力:在 NCC→LEVI 方向上,中位误差仅 3.03°,远优于 C4 的 3.09°、FFCC 的 4.44°、FC4 的 11.8°;在 LEVI→NCC 方向上,中位误差 1.99°,大幅领先 C4 的 5.85° 和 FFCC 的 7.54°。最令人印象深刻的是,RL-AWB 仅用 5 张训练图片就达到了与全数据集训练的模型竞争的性能。在 Gehler-Shi 日间数据集上的跨场景评估(Table 3)中,RL-AWB 中位误差 2.24°,比 SGP-LRD 基础的 2.38° 降低了 5.9%,best-25% 误差从 0.51° 降至 0.46°(降低 9.8%),证明方法也能迁移到日间场景。消融实验(Table 4)显示课程池大小呈 U 形性能趋势:$M=5$ 是最优平衡点,过小($M=3$)缺乏多样性,过大($M=15$)减少每样本访问频率导致探索噪声过大。SGP-LRD 组件消融(Table 5)表明颜色过滤模块对性能至关重要,移除后 LEVI 中位误差从 3.08° 恶化至 3.25°,worst-25% 从 5.46° 恶化至 6.61°。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| NCC 夜间数据集统计方法对比 | 中位角度误差 (Med. Angular Error) | SGP-LRD: 2.12° | RGP: 2.22°, MSGP: 2.48°, BCC: 3.06°, GI: 3.13°, GE-1st: 4.14° | 比次优 RGP 降低 4.5%,比 MSGP 降低 14.5% |
| LEVI 夜间数据集统计方法对比 | 中位角度误差 (Med. Angular Error) | SGP-LRD: 3.08° | RGP: 3.21°, MSGP: 3.12°, GI: 3.10°, BCC: 4.23° | 比 RGP 降低 4.0%,比 BCC 降低 27.2% |
| NCC→LEVI 跨数据集评估 | 中位角度误差 (Med. Angular Error) | RL-AWB: 3.03° | C4: 3.09°, FFCC: 4.44°, C5: 9.12°, FC4: 11.8°, PCC: 11.1°, GCC: 28.1° | 比 C4 降低 1.9%,比 FC4 降低 74.3%,仅用 5 张训练图片 |
| LEVI→NCC 跨数据集评估 | 中位角度误差 (Med. Angular Error) | RL-AWB: 1.99° | C4: 5.85°, FFCC: 7.54°, C5: 4.47°, FC4: 13.2°, GCC: 9.77° | 比 C4 降低 66.0%,比 C5 降低 55.5% |
| Gehler-Shi 日间数据集泛化评估 | 中位角度误差 (Med. Angular Error) | RL-AWB: 2.24° | SGP-LRD: 2.38°, C5: 3.34°, C4: 5.62° | 比 SGP-LRD 降低 5.9%,比 C5 降低 32.9% |
| Gehler-Shi 日间数据集泛化评估 (B-25) | Best-25% 角度误差 | RL-AWB: 0.46° | SGP-LRD: 0.51°, GI: 0.49°, C5: 1.32°, C4: 2.43° | 比 SGP-LRD 降低 9.8%,比 GI 降低 6.1% |
局限与改进
论文坦诚地指出了 RL-AWB 的主要局限性。首先是过度校正问题(over-correction):对于初始估计误差已经较低的图片,RL 智能体可能会过度调整参数,反而导致输出质量下降。Fig. 7 展示了几个这样的失败案例,例如一张图片的误差从初始的 0.18° 上升到 RL 调整后的 7.05°。这种过度校正在初始误差较小的'容易'样本上尤为明显。其次,统计 AWB 方法对近灰度像素的依赖是一个根本性限制:如果场景中缺乏足够的中性灰表面,任何基于灰度像素假设的方法都无法准确估计光照。第三,当前方法假设光照在整个图像中是均匀的,对于突然的光照过渡(如锐利的阴影边界)可能无法正确处理。第四,RL-AWB 的推理时间约为每张图片 1.5 秒(每次调用约 360ms,平均 3 次迭代),虽然相比手动调参已大幅提速,但与直接前馈的深度学习方法(通常在毫秒级)相比仍有差距。第五,动作空间目前仅包含 2 个参数($N$ 和 $p$),虽然这简化了学习问题,但也限制了方法对更复杂 ISP 参数调优的潜力。此外,作者未提供模型参数量和内存占用等工程指标,对部署场景的适用性讨论不足。
独立分析的弱点
尽管 RL-AWB 取得了令人印象深刻的跨传感器泛化结果,但仍存在几个值得深入分析的弱点。第一,过度校正问题反映了奖励设计的缺陷:当前奖励函数仅优化角度误差的改善,缺乏对'何时应该停止调优'的显式建模。改进方向可以引入安全感知的奖励公式,当误差低于某个阈值时大幅增加停止奖励,或采用约束优化策略限制参数调整幅度。第二,两阶段课程学习的 Stage 1 使用固定图片可能导致智能体对特定图片过拟合,虽然 Stage 2 通过循环训练缓解了这一问题,但更先进的课程生成策略(如自动任务选择、反向课程生成)可能进一步提升样本效率。第三,当前的双分支 MLP 编码器结构相对简单,RGB-uv 直方图特征虽然是经典选择,但可能丢失了空间信息。引入轻量级的注意力机制或空间特征提取可能提升特征表达能力。第四,LEVI 数据集虽然开创性地提供了多传感器夜间数据,但仅包含 2 个传感器(iPhone 16 Pro 和 Sony ILCE-6400),跨传感器泛化的结论需要更多样化的传感器验证。第五,论文未讨论方法对不同分辨率和 ISO 范围的敏感性,而这些是实际部署中的关键变量。
未来方向
论文作者和本文的分析共同指出了多个有前景的未来研究方向。作者提出的两个方向尤为重要:(1) 安全感知奖励公式和约束优化策略,用于解决过度校正问题,例如可以设计一个'置信度门控'机制,当 RL 智能体对调整方向不确定时自动回退到基础 SGP-LRD 估计;(2) 层次化策略(hierarchical policies),将当前的两参数动作空间扩展到更多 ISP 模块参数(如去噪强度、锐化程度等),通过层次化结构高效协调多个可调参数。基于已有成果,还可以延伸以下方向:(3) 将 RL-AWB 框架推广到多光照场景(multi-illuminant),利用像素级参数调整而非全局参数;(4) 结合自监督或对比学习,进一步减少对标注数据的依赖,甚至实现零样本跨传感器部署;(5) 探索在线持续学习(online continual learning),使智能体在部署过程中通过用户反馈持续优化策略;(6) 将 SGP-LRD 算法与深度学习方法结合,例如用神经网络预测初始参数范围,再用 RL 精细调优;(7) 在移动设备上进行端到端的轻量化部署研究,优化推理速度和内存占用。
复现评估
从复现角度来看,RL-AWB 的复现条件总体较好但存在一些限制。有利方面:论文提供了详细的超参数设置(batch size 256, $\gamma=0.1$, $\tau=0.005$, learning rate $3 \times 10^{-4}$, 16 并行环境, 15,000 时间步),SGP-LRD 算法的数学描述足够详细(Eq. 1-5),SAC 框架有成熟的开源实现。LEVI 数据集论文中描述了采集细节(700 张图片,两个传感器,手动标注 Macbeth Color Checker),但未明确说明是否已公开发布。训练在 Intel Core i5-13600K CPU 上进行,推理在 NVIDIA RTX 3090 GPU 上进行,算力需求适中。复现的主要难点在于:(1) LEVI 数据集的可用性——如果未公开,跨传感器评估将无法复现;(2) SGP-LRD 算法的双重过滤参数(VarTh、ColorTh)的具体取值论文中未明确给出;(3) 论文提到使用了 16 个并行环境,但未详细说明环境实现细节。总体而言,核心算法和训练流程是可复现的,但完全复现所有实验结果可能需要等待代码和数据的正式开源。论文提供了项目主页链接:https://ntuneillee.github.io/research/rl-awb/
论文图表
这张图通过左右对比展示了传统 AWB 调参(上半部分)和本文提出的 RL-AWB(下半部分)的工作流程。传统流程中,AWB 专家手动调整参数,经过 N 步迭代优化;RL-AWB 中,RL 智能体通过观察 RGB-uv 直方图特征和状态信息,自动选择参数调整量,送入统计算法进行 N 步迭代,每步都通过奖励信号(角度误差变化)来指导优化方向。图中还展示了具体的误差从 1.29° 逐步降低到 0.28° 的过程。
这是理解论文核心思想的第一张图,直观地展示了 RL-AWB 与传统方法的本质区别——将人工调参过程自动化为 RL 决策过程。
展示了 RL-AWB 过度校正的失败案例。4 张图片的初始误差较低(0.18°, 0.32°, 0.48°, 0.70°),但经过 RL 调整后误差反而增大(5.24°, 5.69°, 6.65°, 7.05°)。视觉上可以看到 RL 调整后的图片出现了明显的色偏。
诚实地展示了方法的局限性,过度校正问题是 RL-AWB 最重要的待改进方向之一。