RL-AWB: Deep Reinforcement Learning for Auto White Balance Correction in Low-Light Night-time Scenes

AI summary

Overview

Research area: Computational photography / camera image signal processing (ISP), specifically auto white balance (AWB) and color constancy for low-light nighttime scenes, combining classical statistical illuminant estimation with deep reinforcement learning.

Technical level: Advanced. The paper assumes familiarity with color constancy terminology (illuminant estimation, angular error, log-chrominance histograms), reinforcement learning (Soft Actor-Critic, off-policy updates, curriculum learning), and ISP pipeline concepts.

Scope: The paper introduces SGP-LRD (a nighttime statistical illuminant estimator), RL-AWB (a Soft Actor-Critic agent that tunes SGP-LRD's parameters per image without ground-truth illuminants), and LEVI (a 700-image, two-camera nighttime benchmark), and evaluates them on NCC, LEVI, and Gehler-Shi.

What This Paper Is About

Auto white balance estimates the color of the scene's light source and corrects color casts so white objects look neutral, but at night low illumination, high ISO settings, and severe chromatic noise break the statistical assumptions that daytime AWB relies on, producing unstable estimates that also shift unpredictably across different camera sensors. The authors build a nighttime-specific statistical estimator and then train a reinforcement learning agent to tune that estimator's two key hyperparameters on a per-image basis at inference time, without needing ground-truth illuminants or reference images.

Key Contributions

1. SGP-LRD (Salient Gray Pixels with Local Reflectance Differences): a nighttime-specific color constancy algorithm that the paper reports achieves state-of-the-art illuminant estimation on public nighttime benchmarks. It detects salient gray pixels, refines them with local variance and color-deviation filtering, weights them by a luminance-adaptive confidence measure, and estimates the illuminant from pixel-wise local reflectance differences under a Minkowski norm.
2. RL-AWB: described as the first deep reinforcement learning framework for color constancy. A Soft Actor-Critic (SAC) agent dynamically selects image-specific parameters (the gray-pixel sampling percentage N and the Minkowski order p) for SGP-LRD at inference time, trained with a two-stage curriculum and without requiring ground-truth illuminants or reference images.
3. LEVI dataset: described as the first multi-camera nighttime color constancy dataset, comprising 700 linear RAW images from two camera systems, intended to enable rigorous cross-sensor evaluation. Before this work, the paper states NCC (513 images, single camera) was the only public nighttime color constancy benchmark.
4. Data efficiency and cross-sensor generalization: using only 5 training images per dataset, RL-AWB is reported to generalize across unseen sensors better than fully supervised state-of-the-art methods.

Main Findings

• SGP-LRD beats statistical baselines on both nighttime benchmarks. On NCC, SGP-LRD records median 2.12°, mean 3.11°, trimean 2.29°, best-25% 0.68°, and worst-25% 7.22°, versus the next-best statistical method RGP (median 2.22°, mean 3.33°, trimean 2.44°, best-25% 0.68°, worst-25% 7.81°) and GE-1st (median 4.14°, worst-25% 10.87°). On LEVI, SGP-LRD records median 3.08°, mean 3.25°, trimean 3.07°, best-25% 1.40°, and worst-25% 5.46°.

• RL-AWB improves on SGP-LRD and is competitive with fully trained learned models despite minimal supervision. Trained with only 5 images per dataset, RL-AWB achieves 3.03° median, 3.24° mean, 3.04° trimean, 1.45° best-25%, and 5.36° worst-25° for NCC → LEVI; and 1.99° median, 3.12° mean, 2.25° trimean, 0.67° best-25%, and 7.39° worst-25° for LEVI → NCC. All learning-based baselines (FFCC, C4, C5, FC4, PCC, GCC, ePCC) were trained on the complete datasets using their official three-fold cross-validation protocols, and several degrade badly under cross-dataset shift — for example GCC reaches 28.1° median and 90.0° worst-25° on NCC → LEVI.

• Learning-based baselines degrade substantially under domain shift. The paper reports that for both NCC → LEVI and LEVI → NCC, the median and worst-25% errors of existing learned baselines increase markedly relative to their in-dataset results, which the authors attribute to differences in scene content and camera characteristics. RL-AWB avoids direct illuminant RGB regression and instead tunes SGP-LRD's control parameters per image, relying on the statistical model's robustness to distribution shifts.

• Daytime generalization after removing nighttime-specific filtering. On Gehler-Shi, with the local variance and color deviation filtering modules removed from SGP-LRD and RL-based parameter tuning applied, RL-AWB reaches median 2.24°, mean 3.50°, trimean 2.51°, best-25% 0.46°, and worst-25% 8.67°, compared with SGP-LRD at 2.38°, 3.64°, 2.64°, 0.51°, and 8.89°. The paper states this is a reduction of 5.9% in median angular error and 9.8% in best-25% angular error over SGP-LRD.

• Iterative RL beats one-step regressors. Against an MLP and a CNN trained to predict (N, p) in one step from the same RGB-uv histogram, RL-AWB records lower final median angular error in every tested transfer setting: RL-AWB 1.98/3.03/2.24 (N→N, N→L, N→G), 1.99/3.01/2.27 (L→N, L→L, L→G), and 1.99/3.12/2.16 (G→N, G→L, G→G), versus MLP 2.01/3.07/2.29, 2.25/3.02/2.39, 2.20/3.13/2.34 and CNN 2.06/3.05/2.34, 2.17/3.04/2.35, 2.07/3.17/2.30. The authors attribute this to the agent's ability to revise suboptimal intermediate decisions, which a one-step regressor cannot do. The paper notes that a dense grid search over 440 parameter configurations per image (11 values of N and 40 values of p) would require approximately 94 hours of computation and would need illuminant ground truth at test time.

• Curriculum pool size of 5 is the best trade-off. Varying Stage-2 pool size M ∈ {3, 5, 7, 9, 15} produces a U-shaped trend: M=3 gives NCC median 2.16°, mean 3.29°, worst-25% 7.69° and LEVI median 3.05°, mean 3.28°, worst-25% 5.55°; M=5 gives the best results (NCC 1.98/3.07/7.22, LEVI 3.01/3.22/5.32); larger pools degrade again. Small pools lack diversity, while large pools reduce per-sample visitation under a fixed replay budget and increase exploration noise.

• Both SGP-LRD filtering stages matter, with different roles. Removing noise filtering gives NCC 2.12/3.12/7.32 and LEVI 3.09/3.26/5.48; removing color filtering degrades more, to NCC 2.51/3.90/9.54 and LEVI 3.25/3.68/6.61; the full algorithm gives NCC 2.12/3.11/7.22 and LEVI 3.08/3.25/5.46. The paper notes the noise filter is a robustness safeguard against shot noise rather than an accuracy driver.

• Stated failure mode: over-correction. The agent can increase angular error relative to its initial estimate for images that already have low initial estimation error, producing slightly degraded outputs.

Methodology in Plain English

The method has two layers that work together.

The bottom layer, SGP-LRD, is a classical statistical estimator built for night images. It assumes that under a narrow-band spectral model, the local contrast of an image depends only on surface reflectance, not on the light source. It therefore finds pixels whose red, green, and blue local-contrast values point in the "gray direction" — the top N% ranked by this grayness measure — and then cleans that candidate set in two passes. A local variance filter removes pixels whose logarithmic RGB channels vary too little, since these tend to be pure sensor noise in dark regions; a color deviation filter removes pixels that stray too far from the scene's dominant color cast, using an adaptive threshold based on the minimum per-channel mean. Surviving pixels are the Salient Gray Pixels. Each is given a confidence weight that depends on its luminance relative to mean luminance, with the exponent chosen by the skewness of the luminance distribution (1.0 for high skewness, 2.0 for moderate, 4.0 for uniform illumination). For every pixel position, the algorithm then compares the local mean of the gray-pixel intensities in a 3×3 window against the local maximum for each channel, producing a normalized local reflectance difference. The illuminant is the ratio of the weighted accumulated gray-pixel intensity to the weighted accumulated reflectance difference, raised to 1/p — a Minkowski norm controlled by the parameter p.

The top layer, RL-AWB, addresses the fact that the best values of N% and p depend on the scene. This is framed as a sequential decision problem: a Soft Actor-Critic agent observes a state built from a log-chrominance (RGB-uv) histogram at granularity m = 60, ℓ1-normalized, square-rooted, and flattened, plus an 11-dimensional history vector of recent parameter values and a normalized timestep counter. Both actor and critic use two-branch MLP encoders that map inputs to 64-dimensional embeddings and fuse them; the actor outputs a mean and log-variance for two continuous actions, and twin critics take the minimum Q-value to reduce overestimation. Actions are relative adjustments squashed by tanh and rescaled (N adjustments to [-0.6, 0.6], p adjustments to [-4, 4]), applied to SGP-LRD, and the loop repeats. The reward centers on relative error improvement to normalize across images with different starting errors, adds a penalty for large moves that is relaxed for difficult images, and adds a terminal bonus from the set {+50, +30, +20, +10, -10} based on the final-to-initial error ratio. Training uses two curriculum stages: first a single fixed image to stabilize behavior and teach termination, then a cyclic pool of M = 5 images where the agent runs 5 consecutive episodes per image before cycling onward.

Implementation details include a 3×3 local window, a Laplacian-of-Gaussian operator with σ = 0.5 and a dynamically computed kernel of roughly 7×7, variance thresholds of 0.025 (LEVI) and 0.045 (NCC), color deviation thresholds of 0.3 (NCC) and 0.35 (LEVI), and a confidence threshold of 0.9, all chosen by grid search on a small validation set. Training used an Intel Core i5-13600K CPU with batch size 256, γ = 0.1, τ = 0.005, learning rate 3×10⁻⁴, 16 parallel environments over 15,000 timesteps, and updates starting after 100 initial steps. At inference the agent converges in an average of 3 iterations, taking approximately 1.1 s per image (about 360 ms per call) on an NVIDIA RTX 3080 GPU.

For data, the LEVI dataset contributes 700 linear RAW nighttime images from two systems — iPhone 16 Pro (images #1–370, 4320×2160, 12-bit) and Sony ILCE-6400 (images #371–700, 6000×4000, 14-bit) — with ISO ranging from 500 to 16,000. Each scene contains a Macbeth Color Checker with manual annotations, and ground-truth illuminants are the median RGB values of non-saturated achromatic patches, with all images black-level corrected and converted to linear RGB. For experiments, iPhone 16 Pro LEVI captures are resized by 0.25×, while Sony ILCE-6400 LEVI images and all NCC images are resized by 0.125×.

Why This Matters

The paper argues that robust cross-sensor nighttime AWB is critical for real-world deployment, and that existing paradigms each fail in a specific way: purely statistical methods lose reliability when noise dominates the signal, and purely learned methods are data-hungry and degrade badly across sensors.

Impact on research: The work positions reinforcement learning as a new paradigm for color constancy, distinct from regressing illuminants directly. It also contributes a cross-sensor nighttime benchmark where prior public nighttime data came from a single camera, enabling evaluation that the authors argue was previously impossible. The reported result that a 5-image-trained agent is competitive with or better than fully supervised baselines under domain shift is a notable claim about data efficiency in this domain.

Real-world applications (as named in the paper):

• Mobile photography, where nighttime color casts and cross-device sensor variation are common complaints.
• Surveillance systems operating in low-light conditions.
• Automotive imaging, where night scenes with mixed illumination are typical.
• General camera ISP tuning pipelines, where the paper notes that traditional AWB tuning for night scenes is difficult and time-consuming.

Industry relevance: The authors are affiliated with MediaTek Inc. and the work is framed around professional AWB tuning practice — the agent is described as mimicking AWB tuning experts. It targets two practical constraints that matter to camera makers: avoiding dependence on ground-truth illuminants at deployment, and reducing color shifts across sensors and ISP configurations. The stated inference cost of roughly 1.1 s per image on an RTX 3080 is a consideration for real-time deployment.

Future Directions

• Safety-aware rewards and constrained optimization to address the identified failure mode where the agent over-corrects images that already have low initial estimation error.
• Hierarchical policies to coordinate additional tunable ISP parameters beyond the current two-parameter action space of N and p.
• Broader cross-sensor validation, since LEVI covers two camera systems and the paper only evaluates Gehler-Shi for the daytime case (with the local variance and color deviation filters removed).
• Closing the gap to more accurate but impractical supervision, given that the grid-search procedure that could provide parameter ground truth requires illuminant ground truth at test time and roughly 94 hours of computation per 440 configurations.

Target Audience

This paper is best suited to researchers and engineers working on computational photography, camera ISP pipelines, and auto white balance; reinforcement learning practitioners interested in applied control of classical algorithms rather than end-to-end regression; and engineers at camera, smartphone, or automotive imaging companies who need nighttime color constancy that holds up across different sensors. Readers should have background in color constancy metrics (angular error) and in deep RL algorithms such as Soft Actor-Critic. The dataset contribution (LEVI) also makes it relevant to benchmark-focused researchers studying cross-sensor generalization.

Authors’ abstract

Nighttime color constancy still remains a challenging problem in computational photography due to low-light noise and complex illumination conditions. We present RL-AWB, a novel framework combining statistical methods with deep reinforcement learning for nighttime white balance. Our method begins with a statistical algorithm tailored for nighttime scenes, integrating salient gray pixel detection with novel illuminant estimation. Building on this foundation, we develop the first deep reinforcement learning approach for color constancy that leverages the statistical algorithm as its core, mimicking professional AWB tuning experts by dynamically determining image-specific parameters at inference time, without requiring ground-truth illuminants or reference images. To further facilitate cross-sensor evaluation, we introduce the first multi-sensor nighttime dataset. Experiment results demonstrate that our method achieves strong generalization capability across low-light and well-illuminated images. Project page: https://ntuneillee.github.io/research/rl-awb/