Fully Convolutional Network with Multi-Step Reinforcement Learning

Fully Convolutional Network with Multi-Step Reinforcement Learningfor Image Processing（基于多步强化学习的全卷积网络图像处理 ）

0 Abstract

解决问题：图像处理的像素级奖励强化学习（reinforcement learning with pixel-wise rewards (pixelRL) for image processing）

Pixel IRL：每个像素都有一个代理，这个代理通过采用一个动作来更新像素值，

本实验提出一种PixelRL方法，该方法可应用于图像去噪、图像恢复和局部颜色增强。

1 Contribution

1、提出了一种新的问题设置:PixelRL用于图像处理，现有的多智能体RL技术无法应用。

2、提出reward map convolution，它是PixelRL的一种有效学习方法，能提高性能。

3、将PixelRL应用于图像去噪、图像恢复和局部颜色增强。

4、agents执行的动作是人类可理解的，与传统CNN不同。

2 Background

A3C（asynchronous advantage actor-critic）具有policy network 和 value network两个网络。将这两个网络的参数表示为${\theta }_p$和$ θ_v$，它们都是使用当前状态时间步长t时的$s^{(t)}$作为输入，thevaluenetworks输出$V(s^{{(t)})$：状态$s}{(t)}$的期待总奖励，显示当前状态有多好。

${\theta }_v$的梯度计算：

$R^{(t)}=r^{(t)}+\gamma r^{(t+1)}+{\gamma }^2{r}^{(t+2)}+\cdot \cdot \cdot +{\gamma }^{n-1}{r}^{(t+n-1)}+{\gamma }^{n}V(s^{(t+n)})$

KaTeX parse error: Unexpected character: '' at position 15: dθ_v = ∇_{θ_v}̲ \left(R^{(t)} …

${\gamma }^i$是第$i$次幂的discount factor $\gamma$ 。

策略网络输出采取$a^{(t)}\in A$的策略$\pi (a^{(t)}|s^{(t)})$（通过softmax的概率），策略网络的输出维数为$|A|$，${\theta }_p$的梯度：

$A(a^{(t)},s^{(t)})=R^{(t)}-V(s^{(t)})$

$d{\theta }_p=-{\nabla }_{{\theta }_p}log\pi (a^{(t)}|s^{(t)})A(a^{(t)},s^{(t)})$

$A(a^{(t)},s^{(t)})$是advantage，在$A(a^{(t)},s^{(t)})=R^{(t)}-V(s^{(t)})$中减去$V(s^{(t)})$来减少梯度的方差。

3 Reinforcement Learning with Pixel-wise Rewards (PixelRL)（用像素激励强化学习）

输入图像$I$拥有N个像素$(i=1,\dots \dots ,N)$，$I_i$表示输入图像$I$的第$i$个像素。

每个像素拥有agent，它的策略表示为${\pi }_i(a_i^{(t)}|s_i^{(t)})$，$a_i^{(t)}(\in A)$表示第$i$个agent在时间步长t时的动作，$s_i^{(t)}$表示时间步长t时的状态。

A是预定义的动作集，$s_i^{(0)}=I_i$，agents通过环境中获得下一个状态$s^{(t+1)}=\left(s_1^{(t+1)},\cdot \cdot \cdot ,s_N^{(t+1)}\right)$，

和奖励$r^{(t)}=(r_1^{(t)},\cdot \cdot \cdot ,r_N^{(t)})$。

问题的目标是学习最优策略$\pi =({\pi }_1，\cdot \cdot \cdot ，{\pi }_n)$，使所有像素点的总期望回报均值最大化“

${\pi }^{\ast }=\underset{\pi }{\mathrm{argmax}}{E}_{\pi }\left({\sum }_{t=0}^{\infty }{\gamma }^t{\overline{r}}^{(t)}\right)$

${\overline{r}}^{(t)}=\frac{1}{N}{\sum }_{i=1}^N{r}_i^{(t)}$

${\overline{r}}^{(t)}$是$r_i^{(t)}$的所有像素的奖励均值。

(1)第一种简单的解决方案是训练一个网络，输出Q-value或所有可能的行动集$a^{(t)}$的策略。 然而，由于最后一个全连接层的尺寸必须是$|A{|}^N$，这在计算上是不切实际的。

(2)第二种解决方案是将问题分解为N个独立的子问题，训练N个网络，我们训练第$i$个agent，使第

$i$个像素处的期望总回报最大化。

${\pi }_i^{\ast }=\underset{\pi i}{\mathrm{argmax}}{E}_{\pi i}\left({\sum }_{t=0}^{\infty }{\gamma }^t{r}_i^{(t)}\right)$

当像素数量大时，训练N个网络在计算上时不实际的，其只能处理固定大小的图像。因此本实验采用FCN代替N网络，通过使用FCN使得N个agent共享参数，并可以在GPU上并列化agent的计算，提高训练效率。

agent的数量N非常大，超过$10^5$。所以multi-agent RL无法直接应用于PixelRL 。

4 Reward Map Convolution （激励映射卷积）

当$R^{(t)}=r^{(t)}+\gamma r^{(t+1)}+{\gamma }^2{r}^{(t+2)}+\cdot \cdot \cdot +{\gamma }^{n-1}{r}^{(t+n-1)}+{\gamma }^{n}V(s^{(t+n)})$中$n=1$时，FCNs的接受域为1*1（即策略和值网络中的所有卷积滤波器都为1x1），N个子问题是完全独立的，与原A3C相比，计算两个网络的梯度：

$R_i^{(t)}=r_i^{(t)}+\gamma V(s_i^{(t+1)})$

$d{\theta }_v={\nabla }_{{\theta }_v}\frac{1}{N}{\sum }_{i=1}^N{\left(R_i^{(t)}-V(s_i^{(t)})\right)}^2$

$A(a_i^{(t)},s_i^{(t)})=R_i^{(t)}-V(s_i^{(t)})$

$d{\theta }_p=-{\nabla }_{{\theta }_p}\frac{1}{N}{\sum }_{i=1}^{N}log\pi \left(a_i^{(t)}|s_i^{(t)})A(a_i^{(t)},s_i^{(t)}\right)$

如$d{\theta }_v={\nabla }_{{\theta }_v}\frac{1}{N}{\sum }_{i=1}^N{\left(R_i^{(t)}-V(s_i^{(t)})\right)}^2$和$d{\theta }_p=-{\nabla }_{{\theta }_p}\frac{1}{N}{\sum }_{i=1}^{N}log\pi \left(a_i^{(t)}|s_i^{(t)})A(a_i^{(t)},s_i^{(t)}\right)$所示，每个网络参数的梯度为所有像素点上梯度的平均值。

本实验的网络架构具有很大的receptive field，策略和值网络需要观察第$i$个像素$s_i^{(t)}$和相邻像素，为了可以在第$i$个像素处输出策略$\pi$和值$V$，动作$a_i^{(t)}$会影响$s_i^{(i+1)}$，还影响下一个时间$N(i)$中的策略和值，$N(i)$是位于第$i$个像素中心的局部窗口。

将$R_i$代入$R_i^{(t)}=r_i^{(t)}+\gamma V(s_i^{(t+1)})$中

$R_i^{(t)}=r_i^{(t)}+\gamma {\sum }_{j\in N(i)}{w}_{i-j}V(s_j^{(t+1)})$

$w_{i-j}$是权值，表示在下一个时间步(t + 1)中考虑邻居像素值V的程度。

$w$是一个卷积滤波权值，可以与网络参数${\theta }_p$和${\theta }_v$同时学习。

$\gamma {\sum }_{j\in N(i)}{w}_{i-j}V(s_j^{(t+1)})$是一个2D卷积，因为每个像素$i$都有一个2D坐标$(i_x,i_y)$。

利用矩阵的形式，定义$n$步情况下的$R^{(t)}$：

$R^{(t)}=r^{(t)}+\gamma w\ast r^{(t+1)}+{\gamma }^2{w}^2\ast r^{(t+2)}+\cdot \cdot \cdot +{\gamma }^{n-1}{w}^{n-1}\ast r^{(t+n-1)}+{\gamma }^n{w}^n\ast V(s^{(t+n)})$

$w$的梯度计算：

$dw=-{\nabla }_w\frac{1}{N}{\sum }_{i=1}^{N}log\pi (a_i^{(t)}|s_i^{(t)})(R_i^{(t)}-V(s_i^{(t)}))+{\nabla }_w\frac{1}{N}{\sum }_{i=1}^N(R_i^{(t)}-V(s_i^{(t)}){)}^2$

$-{\nabla }_w\frac{1}{N}{\sum }_{i=1}^{N}log\pi (a_i^{(t)}|s_i^{(t)})(R_i^{(t)}-V(s_i^{(t)}))$鼓励更高的预期总回报

${\nabla }_w\frac{1}{N}{\sum }_{i=1}^N(R_i^{(t)}-V(s_i^{(t)}){)}^2$作为正则化，防止$R_i$不会通过卷积偏离预测$V(s_i^{(t)})$

5 Applications and Results

5.1 Image denoising

Method

输入图像$I(=s^{(0)})$为带噪声的灰度图像，agent通过执行动作迭代进行去噪。

reward $r_i^{(t)}$：$r_i^{(t)}=(I_i^{target}-s_i^{(t)}{)}^2-(I_i^{target}-s_i^{(t+1)}{)}^2$

$I_i^{target}$为原始干净图像的第i个像素值。

设定训练集数为30000集，每集长度$t_{max}$为5集。 我们将$w$的过滤大小设置为33×33，等于policy and value networks的receptive field大小。

Results

RMC：reward map convolution

aug.： 测试图像的数据增强

通过左右翻转和$90^{\circ }$、$180^{\circ }$、$270^{\circ }$旋转，将单个测试图像增强到8张图像。

在策略网络中加入convGRU后，PSNR提高了约0.1dB。 σ = 15时，RMC显著提高了PSNR，但σ = 25和50时，RMC提高不大。这是因为当噪声较强时，agent通过在自己的像素上去除噪声而不是考虑邻近像素，可以获得较大的奖励。

图1显示了在BSD68测试集上进行高斯去噪(σ = 50)时，本实验在每个时间步长所执行的动作数。

结果表明，agent首先采用强滤波器(box filter, bilateral filter σc = 1.0, and Gaussian filter σ = 1.5)去除噪声，然后通过其他动作调整像素值(像素值+=1和-=1)。

表3为Salt和Pepper去噪的比较。 观察到RMC显著地改善了性能。 此外，在噪声密度为0.5和0.9时，本文方法优于CNN (Zhang et al. 2017)。

由于原始像素值的信息丢失(即像素值被噪声改变为0或255)，CNN很难对噪声进行回归。本方法可以通过迭代滤波动作从相邻像素点预测出真实像素值。

5.2 Image Restoration

Method

将该方法应用于“blind”图像恢复中，不提供空白区域的掩模，该方法通过执行动作迭代地绘制空白区域。

使用与图像去噪相同的动作和奖励函数。

使用25,295幅图像和11,343个文档进行训练，在训练过程中，我们从25,295张图像中随机选择一张图像，从11,343个文档中随机选择一个文档，并将其覆盖在图像上，从而创建每个训练图像。

Results

表4为输出图像与ground truth图像的平均PSNR比较。

将Net-D、Net-E和CNN之间进行比较，发现RMC显著提高了绩效，本方法通过过滤动作迭代传播相邻像素值来预测真实像素值，规避了文本叠加会丢失原始像素值信息，回归比较困难的问题。

5.3 Local Color Enhancement

Method

使用了(Yan et al. 2016)创建的数据集，其中有70张列车图像和45张从Flicker下载的测试图像。使用Photoshop，所有的图像都由专业摄影师为三种不同风格的局部效果增强：Foreground Pop-Out, Local Xpro, and Watercolor.

给定一个输入图像I，所提出的方法通过执行一个动作来改变每个像素处的三个通道像素值。 将I输入到网络时，将RGB颜色值转换为CIELab颜色值。 将奖励函数定义为CIELab颜色空间中L2距离的减小：

$r_i^{(t)}=|I_i^{target}-s_i^{(t)}{|}_2-|I_i^{target}-s_i^{(t+1)}{|}_2$

所有超参数和设置与图像恢复时相同，除了episodes长度$t_{max}=10$。

Results

表6显示了45幅测试图像的平均L2误差的比较。在所有三种增强方式上，本文提出的方法都比DNN (Yan et al. 2016)取得了更好的结果，并且比pix2pix取得了类似或稍好的结果。观察到RMC改进了性能，尽管它们的改进程度取决于风格。 值得注意的是，现有使用深度RL的颜色增强方法 (Park et al. 2018; Hu et al. 2018)不能应用于这个局部增强应用，因为它们只能执行全局操作。

6 Conclusions

提出一种新的PixelRL方法，并应用于图像去噪、图像恢复和局部颜色增强。

提高PixelRL代理的性能，

该方法能最大化像素级奖励，可以最小化像素级的不可微目标函数。