Controlling Neural Style Transfer with Deep Reinforcement Learning

1 Introduction

Neural style transfer (NST) refers to generation of a pastiche image combining the semantic content of one image (the content image) and the visual style of another image (the style image) using a deep neural network. NST can be used to create stylized non-photorealistic rendering of digital images with enriched expressiveness and artistic flavors.

Existing NST methods usually generate a stylized image with a one-step approach: a neural network is trained to minimize a loss function of the visual similarity between the content image and the stylized image and the style similarity between the style image and the stylized image [Cheng et al., 2021], and the trained Deep Learning (DL) model is run once to create a stylized image. This one-step approach has an apparent limitation: it is hard to determine a proper level of stylization to fit various flavors of different users since the ultimate metric of style transfer is very subjective. It is observed that generated stylized images by current NST methods tend to be under- or over-stylization [Cheng et al., 2021]. A remedy to under-stylization is to apply the DL model iteratively until a desired level of stylization is reached. However, this solution may suffer from high computational cost due to intrinsic complexity of one-step DL models. Other existing methods, like [Gatys et al., 2015b] and [Huang and Belongie, 2017], play a tradeoff between content and style by adjusting hyper-parameters. These methods are inefficient since there is no guarantee that a user can get the expected output via one-time adjusting.

∗Corresponding authors.

To address the aforementioned limitations of existing DLbased NST methods, we propose a novel framework in this paper, called RL-NST, based on a reinforcement-learning (RL) framework to progressively perform style translation as shown in Fig. 1. Given a content image, we consider the stylized content to be added progressively by the stylizer. More specifically, a stochastic actor in RL-NST first estimates a 2D Gaussian distribution to sample hidden actions, then uses the action to control the stylizer to generate an intermediate stylized image, which is in turn passed to the actor as the input for the next step. Our model also includes a critic to evaluate the latent action. The whole structure is shown in Fig. 2. Furthermore, by using a CNN+RNN architecture [Mirowski et al., 2017] for the actor and stylizer for both frame-wise and step-wise smoothing, our model can perform video NST tasks. To the best of our knowledge, this is the first work that successfully leverages RL for the NST scenario.

Our major contributions can be summarized as follows: 1) We propose the first reinforcement-learning-based NST method, RL-NST, that facilities step-wise style transfer. It provides more flexible control of the degree of stylization without any hyper-parameter adjustment during generation of stylized images. 2) Our RL-NST stylizes a content image progressively, with increased stylization along with more iterations (see Fig. 1). It leads to a lightweight NST model compared with existing one-step DL-based methods, making it computationally more efficient. 3) From our extensive experiments, our RL-NST demonstrates better effectiveness and robustness than existing state-of-the-art methods on both image and video NST tasks.

2 RelatedWork

Neural Style Transfer. Since the seminal work of Gatys et al. [Gatys et al., 2015b] that uses a neural network to produce striking stylized art images, many methods have been proposed to improve the quality and/or running efficiency of NST algorithms. Existing NST methods can be divided roughly into two groups: image-optimization-based [Gatys et al., 2015b; Li et al., 2017a; Risser et al., 2017] and model-optimization-based [Johnson et al., 2016; Ulyanov et al., 2016a; Chen et al., 2017; Huang and Belongie, 2017; Sheng et al., 2018; Park and Lee, 2019; An et al., 2021].

Reinforcement Learning. Reinforcement learning (RL) concerns how an agent takes actions in an environment to maximize its cumulative reward. Standard RL works well on tasks with finite and discrete action spaces. For realworld tasks with high dimensional continuous actions, such as robotic control [Tassa et al., 2018; Gu et al., 2017; Xiang et al., 2022], maximum entropy RL (MERL) and its variants, including Soft Q-learning [Haarnoja et al., 2017; Zhao et al., 2019] and SAC [Haarnoja et al., 2018; Hu et al., 2023], have been proven to have a stable and powerful performance. But they still have limitations to handle high dimensional continuous state and action spaces in image-to-image (I2I) transformation and synthesis. To address this problem, SAEC [Luo et al., 2021] extends the traditional MERL framework [Haarnoja et al., 2018] with an additional executor. It shows promising performance on several I2I tasks. However, its actions are 1D vectors, which do not preserve 2D spatial information of images.

3 Our RL-NST Framework

We formulate NST as a decision-making problem. Instead of directly converting a content image to a stylized image in a single step, we propose to use a lightweight model to perform the translation progressively, with more stylization added to the stylized image as the translation progresses.

Let (c, e) be a pair of content and style images from image domain X ∈ Rd , where d is the dimension. Our model consists of three components, as shown in Fig. 2: an actor πϕ parameterized by ϕ, a stylizer ηψ with model parameters ψ, and a critic Qθ parameterized by weights θ. The actor generates a latent action according to a stochastic policy, so as to capture the style transfer control. The critic evaluates the generated action. The stylizer leverages the latent action and image details to perform style transfer on the content image and generate a stylized image, referred to as a moving image, which updates the environment and is used subsequently as the content image in the next iteration. We use the pre-trained VGG [Simonyan and Zisserman, 2014] as our feature extraction network because it is widely used such as in [Chen et al., 2021; Lin et al., 2021]. Actor πϕ and stylizer ηψ are supervised jointly with content loss, style loss, total variation regularization, and compound temporal regularization (for video only) on the current content image (i.e., the moving image). Furthermore, the actor πϕ and critic Qθ form an actor-critic model.

3.1 RL-NST Settings

In forming our model, we define the infinite-horizon Markov decision process (MDP) as tuple (S,A,P, r, γ), where S is a set of states, A is the action space, and P : S × S ×A → [0,∞) represents the state transition probability of the next state st+1 given st ∈ S at time t and action a ∈ A, r : S ×A → R is the reward emitted from each transition, and γ ∈ [0, 1] is the reward discount factor. Specifically, (1) State. State st is the moving image, initialized by the content image. (2) Action. To extract high-level abstraction at of st from the actor, a stochastic latent action is modeled as at ∼ πϕ(at|st). In practice, it can be obtained by using a reparameterization trick [Kingma and Welling, 2013] at = fϕ(ϵt, st), where ϵt is an input noise vector sampled from a 2D Gaussian distribution. The moving image at time t, i.e., state image st+1, is created by the stylizer based on at and current state image st. (3) Reward. It is from the environment E, obtained by measuring the difference between current state st and the style image. The higher the difference is, the smaller the reward is.

3.2 Network Architecture

The Actor is a neural network model that consists of three convolutional layers and a residual layer. After each convolutional layer, there is an instance norm layer and a ReLU layer. In the residual layer, we use the residual block designed by He et al. [He et al., 2016]. The actor estimates a 2D Gaussian distribution for sampling our latent actions, which is forwarded to the stylizer to generate the moving image. The dimension of our latent action is 64×64, which is a 2D sample and able to preserve more spatial structure information of images. Notably, the actor in our method is to learn the latent actions, which are more compact than the image presentation. The learned actions are controlled by our framework to guide the stylizer to generate stylized images instead of the reconstruction.

The Stylizer has three up-sampling layers correspondingly. More importantly, by using 2D Gaussian sampling, our actorstylizer structure is a fully convolutional network (FCN), which can process images of any input size, instead of only accepting test images of the same size as training images. We also use three skip connections between the actor and the stylizer to stabilize the training process.

To handle video NST tasks, we expand the actor-stylizer to include RNN layers by following the work [Mirowski et al., 2017]. In particular, the ConvGRU layer [Shi et al., 2015] is used for the FCN structure (see the right bottom of Fig. 2). We add Step-wise GRU [Mirowski et al., 2017] to actor and Frame-wise GRU [Donahue et al., 2017] to stylizer. Specifically, for each frame, the hidden state of the Step-wise GRU at each step comes from the output of the Step-wise GRU at the previous step. The role of Step-wise GRU is to make the model maintain better content consistency. Furthermore, the hidden state of the Frame-wise GRU at each step is derived from the output of the Frame-wise GRU at the same step in the previous frame. Frame-wise GRU can make the model maintain better inter-frame consistency.

The Critic consists of seven convolutional layers and one fully-connected layer at the end. Since using standard zeropadded convolutions in style transfer leads to serious artifacts on the boundary of a generated image [Ulyanov et al., 2016b], we use reflection padding instead of zero padding for all the networks.

3.3 Model Training

Our RL-NST contains two learning procedures, namely style learning and step-wise learning.

Style Learning

To make the moving image not deviate from the content image, actor πϕ (encoder) and stylizer ηψ (decoder) are trained together to preserve the perceptual and semantic similarity with the content image. More specifically, the actor and the stylizer form a conditional generative process that translates state st to output moving image mt via the mapping mt = ηψ(πϕ(st)) at time t. Note that st is initialized to content image c and st+1 is equivalently mt. Inspired by [Johnson et al., 2016], we apply the content loss LCO , style loss LST , and total variation regularization LTV to optimize the model parameters of πϕ and ηψ in the image setting. These losses can better measure perceptual and semantic differences between the moving image and content image c. For the video setting, we add an additional loss named compound temporal regularization LCT , which can force the model to generate temporal consistent results under the compound transformation. More details about these losses are as follows.

Content Loss LCO . Following [Johnson et al., 2016], we use a pre-trained neural network F to extract the high-level feature representatives of mt and c. The reason for using this F is to encourage moving image mt to be perceptually similar to content image c but does not force them to match exactly. Denote F j (·) as the activations of the j-th layer of F . Suppose j-th layer is a convolutional layer, then the output of F j (·) will be a feature map with size C j × Hj × W j , where C j , Hj , and W j represent the number of channels, height, and width in the feature map of layer j, respectively. We apply the Euclidean distance, which is squared and normalized to design the content loss as follows,

Style Loss LST . To penalizemt when it deviates in content from c and in style from e, following [Gatys et al., 2015a], we define aGrammatrixGj (x) = ˜ F j (x)( ˜ F j (x))⊤ Cj Hj W j ∈ RCj ×Cj , where ˜ F j (·) is obtained by reshaping F j (·) into the shape C j × HjW j . The style loss can be defined as a squared Frobenius norm of the difference between the Gram matrices of mt and e. To preserve the spatial structure of images, we use a set of layers, J, instead of a single layer j. Thus, we define the style loss to be the sum of losses for each layer j ∈ J (J = 4 in our experiments):

Total Variation Regularization LTV . To ensure spatial smoothness in moving image m, we use a total variation regularizer LTV (mt), which has been widely used in existing works [Mahendran and Vedaldi, 2015; Johnson et al., 2016].

Compound Temporal Regularization LCT . Inspired by [Wang et al., 2020], we add a compound temporal regularization for video style transfer. Specifically, we first generate motions M(·) and then synthesize adjacent frames. With this approach, we do not need to estimate optical flow in the training process and we can guarantee the optical flows are absolutely accurate. Given noise △, to maintain temporal consistency, we can minimize the following loss

Summing up all the components, the final style learning loss is

where λ, β, and ζ are hyper-parameters to control the sensitivity of each term. For image style transfer, we use the first three terms. For video style transfer, we use all terms. Then we can update ϕ and ψ from the actor and stylizer by using the gradient descent method with a predefined learning rate η with the following steps:

Step-wise Learning

Our step-wise learning is based on the MERL framework [Haarnoja et al., 2018], where rewards and soft Q values are used to iteratively guide the stochastic policy improvement. Moreover, we focus on latent action a, and use it to estimate a soft state-action value to encourage high-level policies. Specifically, we concatenate at to the downsampled vector of the critic and output soft Q function Qθ(st, at), which is an estimation of the state value at time t. Because the critic is used to evaluate the actor, rewards rt and the soft Q values are used to iteratively guide the stochastic policy improvement by minimizing the soft Bellman residual:

where D is a replay pool and V¯ θ(st) = Eat∼πϕ [Q¯ θ(st, at) - α logπϕ(at|st)]. Note that we utilize the negative value of LST (st, e) for rt in practice.

The critic network Q¯ θ is used to stabilize the training, whose parameters ¯ θ are obtained by an exponential moving average of parameters of the critic network [Lillicrap et al., 2015]: ¯ θ → τθ + (1 - τ)¯ θ, with hyperparameter τ ∈ [0, 1]. To optimize JQ(θ), we use the gradient descent with respect to parameters θ as follows,

Algorithm 1: RL-NST

Input: c, e, and replay pool D

Init: ϕ, ψ, θ, ¯ θ, D ← ∅, η, ηQ, ηϕ, and environment E for each iteration do

where ηQ is a learning rate. Since the critic works on the actor, it will affect the actor's decisions. Therefore, the following objective can be applied to minimize the KL divergence between the policy and a Boltzmann distribution induced by the Q-function,

The last equation holds because at can be evaluated by fϕ(ϵt, st), as we discussed before.Note that hyperparameter α can be automatically adjusted by using the method proposed in [Haarnoja et al., 2018]. Similarly, we apply the gradient descent method with a learning rate ηϕ to optimize parameters as follows,

The pseudo-code of optimizing RL-NST is described in Algorithm 1. All parameters are optimized based on the samples from replay pool D.

4 Experiments

We have conducted a series of experiments to evaluate the effectiveness of RL-NST in realizing step-wise style transfer on both image and video NST tasks. Our code, a user study, and additional results with more detailed information can be found in the supplementary materials.

4.1 Experimental Settings

Datasets. (1) For image style transfer, we select style images from WikiArt [Phillips and Mackintosh, 2011] and use MS-COCO [Lin et al., 2014] as content images in which the training set includes 80K images and the test set includes 40K images. All training images are resized to 256×256. In the inference stage, our method is applicable for content images and style images of any size. (2) For video style transfer, we randomly collect 16 videos of different scenes from pexels[pex, 2022]. Then these videos are extracted into video frames and we obtain more than 2.5K frames. We regard these frames as the content images of training set. Note that the style images in the training set are also selected from WikiArt [Phillips and Mackintosh, 2011]. In addition, following [Wang et al., 2020], we use the training set of MPI Sintel dataset [Butler et al., 2012] as the test set, which contains 23 sequences with a total of 1K frames. Similarly, all training frames are resized to 256×256, we use the original frame size in testing.

Baselines and Evaluation Metrics. (1) For image style transfer, we choose the following eight classical and latest state-of-the-art style transfer methods as our baselines: Johnson et al. [Johnson et al., 2016], WCT [Li et al., 2017b], AdaIN [Huang and Belongie, 2017], SANet [Park and Lee, 2019], LapStyle [Lin et al., 2021], ArtFlow [An et al., 2021], IEContraAST [Chen et al., 2021], AdaAttN [Liu et al., 2021], and StyTR2 [Deng et al., 2021b]. Following StyTR2 [Deng et al., 2021b], we evaluate all the algorithms in terms of stylization effect, computing time, content loss, and style loss. (2) For video style transfer, we compare our method with the following four popular methods: Linear [Li et al., 2019], MCCNet [Deng et al., 2021a], ReReVST [Wang et al., 2020], and AdaAttN [Liu et al., 2021]. Following [Liu et al., 2021], we use temporal loss as the evaluation metric to compare the stability of stylized results. All these methods are performed using their public codes with the default settings.

Implementation Details. In the experiment, we set λ = 1e5, β = 1e - 7, ζ = 1e2 in Eq. (1), and η = 1e - 4 in Eq. (2). These settings yield nearly the best performance in our experiments. Following [Wang et al., 2020], in LCT , M(·) is implemented by warping with a random optical flow. Specifically, for a frame of size H×W, we first generate a Gaussian map (wavy twists) Mwt of shape H/100×W/100×2, mean 0, and standard deviation 0.001. Second, Mwt is resized to H ×W and blurred by a Gaussian filter of kernel size 100. Finally, we add two random values (translation motion) Mtm of range [-10,10] toMwt, and obtainM. In addition, random noise △ ∼ N (0, σ2 I), where σ ∼ U(0.001, 0.002).

4.2 Evaluations on Image RL-NST

Qualitative Comparison. Fig. 3 shows some stylized results of our RL-NST and the baseline methods. For content images with fine structures such as the forest image (Udnie style), all the baseline methods proposed to address content leakage, including ArtFlow, produce messy stylized images with a complete loss of content structure. Moreover, SANet has repeated texture patterns for all the cases, and most of its results are hard to generate sharp edges.

In contrast, our method can produce stable and diversified stylized results with good content structures. This may be attributed to our step-wise solution. More specifically, the content image is stylized progressively and hence smoothed stylization results are obtained. More importantly, as we mentioned before, despite that stylization is quite subjective, our step-wise method provides flexible control of different degrees of stylization to fit the need of different users.

Quantitative Results. To be consistent with all compared methods shown in Fig. 3, we compare our method with all baselines without caring which type (single or multiple styles) they are. The quantitative results are shown in Table 1. Our RL-NST (step=1) achieves better performance than the baseline methods in all evaluation metrics. Our method still has low content and style losses even if the step is equivalent to 10, which means our method is friendly to the user for choosing the results from specific steps accordingly. In addition, it is clear that our model has much fewer parameters and a faster speed. For example, the time cost and the parameter size of our method are 2/7 and 1/9 of Johnson et al., and 4/47 and 1/43 of LapStyle, respectively.

Ablation Study. (1) We study the effect of the RL model in our framework. As shown in Fig. 4, compared with the method that uses only Actor-Stylizer (AS), our method can generate more stable and clear stylized images at step 1. At step 10, AS loses the content information completely without the help of RL (Ours w/o RL), while our method can still produce amazing results. We also show the corresponding numerical comparison in Table 2. We can easily see that our method achieves the best performance consistently in both steps 1 and 10. This study indicates that RL can indeed improve the performance of DL-based NST models. (2) Since AdaIN and StyTR2 methods can adjust the hyperparameter 'alpha'∈ [0, 1] and the round of repetitive stylization to control the degree of stylization in the final results, respectively, we compare our method with them accordingly in Fig. 5. From the visualization results, we can see that the results of AdaIN are in the under-stylized state even if the style control hyperparameter is changed. Moreover, StyTR2 gets the results with small style changes and low quality after multiple rounds. However, our method not only ensures the gradual change in style, but also produces very smooth results.

4.3 Evaluations on Video RL-NST

Qualitative Comparison. We show the visualization results of our method compared with the four latest video style transfer methods in Fig. 6, wherein, for each method, the top portion shows the specific stylized results and the bottom portion is the heatmap of the differences in the adjacent frames of the input and stylized videos. Note that the adjacent frame indexes are the same for all methods. We can find that our method produces refined stylized results and our results are closest to the input frames. In particular, our method can highly promote the stability of video style transfer. The differences in our results are closest to the difference from input frames without reducing the effect of stylization. It is clear that MCCNet and ReReVST fail to keep the coherence of videos. In addition, Linear and AdaAttN also fail to keep the coherence in some regions that are close to the edge of objects such as the head and shoulder.

Quantitative Results. As shown in Table 3, we choose 23 different sequences from the MPI Sintel dataset [Butler et al., 2012] and eight different style images to calculate the average of temporal losses for comparison. It is clear that our method (step=1 and 5) outperforms the compared methods in all style settings. Our method still has a low temporal error even if step=10.

Ablation Study. We investigate the effect of the individual parts of the network structure group on the results, including the RL, Step-wise GRU, and Frame-wise GRU.

(1) As shown in Fig. 7 (first and second rows), our method generates more stable and clearer stylized results than the method using only Actor-Stylizer without RL. After step 5, AS no longer has the ability to keep the content information and style information, while our method with RL can still produce good results. (2) Similarly, we have compared our method with the results produced when Step-wise GRU is removed, and the results are shown in Fig. 7 (second and third rows). We can clearly see that most of the face details of the protagonist and dragon have been lost at step 10 when using our method without the Step-wise GRU. Also, the external details of the protagonist and dragon are completely lost in step 20. Our method with using Step-wise GRU, on the other hand, obtains very fine results even at step 20. (3) Table 4 shows the comparison of the temporal loss of our method with Frame-wise GRU (FWG) and without FWG. We find that the temporal loss is very low if we use FWG, which means the obtained final results are more consistent from frame to frame. The above experiments have shown that RL, Step-wise , and Frame-wise GRU all greatly improve the performance of the model.

5 Conclusion

In this paper, we propose a new RL-based framework called RL-NST for both image and video NST tasks. It achieves style transfer step by step for a flexible control of the level of stylization. Despite using a lightweight neural network model, RL-NST can handle an extremely high-dimensional latent action space (up to 64×64) and is thus capable of generating visually more satisfying artistic images than existing NST methods. Experimental results have demonstrated the effectiveness of the proposed method. The main goal of this work is to show the effectiveness of stylization-level controlling with our RL-based method and the superiority of our method in achieving the best NST quality. Therefore we use a single-style NST model. In the future, we would like to extend our model to multiple-style transfer tasks.

Acknowledgements

This work was supported in part by Sichuan province Key Technology Research and Development project under Grant No. 2023YFG0305, 2023YFG0124, CUIT Science and Technology Innovation Capacity Enhancement Program Project under Grant KYTD202206.

Contribution Statement

The contributions of Chengming Feng and Jing Hu to this paper were equal.