Stochastic Actor-Executor-Critic for Image-to-Image Translation

1 Introduction

Many computer vision problems, such as face inpainting, semantic segmentation, and realistic photo generated from sketch, can be defined as learning image-to-image translation (I2IT). Currently the most effective I2IT solution is based on a one-step framework that generates images in a single run of a deep learning (DL) model, such as VAE, U-Net, and conditional GANs. Directly learning I2IT with these DL models is challenging, due to the abundance of local minimums and poor generalization caused by overfitting. Although these problems could be potentially alleviated by using multi-scale models or a multi-stage pipelines, we are still left with models that have intrinsically high complexities, for which the optimal parameters (e.g. stage number and scale factor) have to be determined in a subjective and ad hoc manner.

∗Equal contribution.

†Corresponding authors.

To address these limitations of DL-based methods, we explore solving I2IT problems by leveraging the recent advances in deep reinforcement learning (DRL). The key idea is to decompose the monolithic learning process into small steps by a lighter-weight CNN, with the aim of progressively improving the quality of the model. Although recent works have successfully applied DRL to solve several visual tasks [Caicedo and Lazebnik, 2015], the action space in their tasks is usually discrete and can not be used for I2IT that requires continuous action spaces.

Currently, a promising direction for learning continuous actions is the maximum entropy reinforcement learning (MERL), which improves both exploration and robustness by maximizing a standard RL objective with an entropy term [Haarnoja et al., 2018]. Soft actor-critic (SAC) [Haarnoja et al., 2018] is an instance of MERL and has been applied to solve continuous action tasks. However, the main issue hindering the practical applicability of SAC on I2IT is its inability to handle high-dimensional states and actions effectively. Although recent work [Yarats et al., 2019] addresses this problem by combining SAC with a regularized autoencoder (RAE), RAE only provides an auxiliary loss for an endto-end RL training and is incapable for I2IT tasks.

Besides, a glaring issue is that the reward signal is so sparse that leads to an unstable training, and higher dimensional states such as pixels worsen this problem [Yarats et al., 2019]. Thus, training an image-based RL model requires much more exploration and exploitation. One solution to stablize training is to extract a lower dimensional visual representation with a separately pre-trained DNN model, and learn the value function and corresponding policy in the latent spaces [Nair et al., 2018]. However, this approach can not be trained from scratch, which could lead to inconsistent state representations with an optimal policy. Previous works like [Lee et al., 2020] have shown that it is beneficial to learn the image representations and continuous controls simultaneously.

In this paper, we propose a new DRL architecture, stochastic actor-executor critic (SAEC), to handle I2IT with very high dimensional continuous state and action spaces. SAEC is formed by three core deep neural networks: the actor, the critic, and the executor (see Fig. 1). High level actions are generated by the actor in a low dimensional latent space, and forwarded to the executor for image translation. The critic evaluates the latent actions to guide the policy update. These three neural networks form two different branches in SAEC. More specifically, the actor and the critic form the RL branch as an actor-critic model, while The actor and executor form the DL branch similar to autoencoder, with an important difference that skip connections are added between the actor and executor. The RL branch is trained based on the MERL framework, while the DL branch is trained by minimizing a reconstruction objective.

Our contributions can be summarized as follows:

• A new DRL framework stochastic actor-executor critic (SAEC) is proposed to handle very high dimensional state and action spaces in I2IT task.
• Even using a sparse reward signal, our SAEC model can be stably trained from scratch on I2IT problems by combining DL-based supervision.
• Our SAEC framework is flexible to incorporate many advanced DL methods for various complex I2IT applications. And this framework enables the agent to simultaneously learn feature representation, control, and image generation in high-dimensional continuous spaces.

2 Background

2.1 Image-to-Image Translation

Image-to-image translation (I2IT) is often addressed by learning a generative process G that maps state x to target y, G : x → y. The state should be consistent with the target and both of them are images, such as generating realistic photos from semantic segmentation labels [Isola et al., 2017], or synthesizing completed visually targets from images with missing regions [Pathak et al., 2016]. Autoencoder is leveraged in most research work to learn this process by minimizing the reconstruction error Lrec between the predicted image ˜y and the target y (see Figure 1). In addition, the generative adversarial network (GAN) is also vigorously studied in I2IT to synthesis realistic images [Isola et al., 2017]. Subsequent works enhance I2IT performance by using a coarse-to-fine deep learning framework [Yu et al., 2018] that recursively uses the previous stage's output as the input to the next stage (see the bottom of the Fig. 1(a)). In this way, I2IT task is transformed into a multi-stage, coarse-to-fine solution. Although iteration can be infinitely applied in theory, it is limited by the increasing model size and training instability.

2.2 Reinforcement Learning

Reinforcement learning (RL) is an infinite-horizon Markov decision process (MDP), defined by the tuple (X , A, P, r, γ). X is a set of states, A is action space, and P represents the state transition probability given x ∈ X and a, r is the reward emitted from each transition, and γ ∈ [0, 1] is the reward discount factor. The standard objective of RL is to maximize the expected sum of rewards. Maximum entropy RL (MERL) further adds an entropy term: ∑Tt=0 E(xt,at)∼ρπ [rt + αH(·|xt)], where α is a temperature parameter, ρπ denotes the marginal trajectory distribution induced by policy π. MERL model and has proven stable and powerful in low dimensional continuous action tasks, such as games and robotic controls [Haarnoja et al., 2018]. However, when facing complex visual problems such as I2IT, where observations and actions are high-dimensional, it remains a challenge for RL models [Lee et al., 2020].

3 Method

This work reformulates I2IT as a decision making problem. Instead of directly mapping the input to the target image in a single step, we propose to use a light-weight RL model that performs the translation progressively, where new details can be added as the translation progresses. The proposed model stochastic actor-executor-critic (SAEC) consists of three components: an actor (πφ), an executor (ηψ), and a critic (Qθ). Specifically, the actor πφ and the executor ηψ form a DL pathway that directs the agent learning with a DL-based objective. The same actor and the critic Qθ form a DRL pathway that works as an actor-critic model. In SAEC, the actor generates a latent action according to a stochastic policy, so as to capture the required image translation control. The critic evaluates the stochastic policy. The executor leverages the latent action and image details to perform image translation and the generated image is then applied to the environment as the new input image at next time.

3.1 I2IT as MDP

The key of our method is to formulate I2IT as a Markov decision process (MDP). As for state xt ∈ X, it is defined task specific. For realistic photo generation or segmentation, the state xt is the pair of the fixed source image If and the moving image Imt : xt = (If , Imt ). A new state is generated by warping the moving image with predicted image action ˜yt: xt+1 = (If , ˜yt ◦ Imt ). For face inpainting, the state is composed of the original image x and the synthesized missing part ˜yt. The next state is obtained by adding the new predicted image to the missing area: xt+1 = (˜yt ⊙ xt).

There are two types of actions in our action space: latent action zt ∈ Z and image action ˜yt ∈ Y. These two actions are intrinsically different, in that Z has a high-level of abstraction and its distribution is unknown. zt is sampled from the stochastic policy of the actor πφ, zt ∼ πφ(xt), while ˜yt is created by the executor based on zt and the input image. The reward function is flexible and can be any evaluation metric that measures the similarity between the predicted image ˜yt and the target y, such as the peak signal-to-noise ratio (PSNR) and the structural similarity index measure (SSIM).

The sequence of actions are implemented by the stochastic actor-executor-critic (SAEC) model. Specifically, at each time step t, the actor samples a latent action zt from state xt, then the executor generates image action ˜yt from the latent action zt and state xt. The change of the state will result in a reward rt, and the new state is fed back to the SAEC model to generate a new image action. Let Tt represents a single run from state xt to image action ˜yt via an actor and an executor: zt ∼ πφ(xt), ˜yt = ηψ(xt, zt). The whole process of I2IT in MDP can be formulated as follows:

where ◦ is the function composition operation, Tt is the tth translation step that predicts image from current state xt.

Algorithm 1 Stochastic Actor-Executor Critic

Input: Environment Env and initial parameters θ1, θ2 for

dynamically in experience replay. The objective of SAEC is to learn a policy πφ and an executor ηψ, by maximizing both the conditional likelihood and the expected sum of rewards with an entropy H.

Conditional Generative Model Learning

The actor and the executor form a conditional generative process that translates the current state xt to target y. Instead of learning the generative model alone like a VAE, we maximize

3.2 Policy with Skip Connections

As in an MERL model, the stochastic policy and maximum entropy in our SAEC improve the exploration for more diverse generation possibilities, which helps to prevent agent from producing a single type of plausible output during training (known as mode-collapse). In addition, one specific characteristic in the SAEC model is that skip connections are added from each down-sampling layer of the actor to the corresponding up-sampling layer of the executor, as shown in Fig. 1(b). In this way, a natural looking image is more likely to be reconstructed since the details of state xt can be passed to the executor by skip connections. Besides, since both zt and xt can be used by the executor to generate image action, over-exploration of the action space can be avoided in SAEC, where the variance is limited by these passed detail information. Furthermore, the skip connections also facilitate backpropagation of the DL gradients to the actor. Our experiments show that if the skip connections are removed, the details of the reconstructed image may lost because of down-sampling (max-pooling) and stochastic policies, and the training would be very unstable in such a situation. Combining skip connections and DL-based learning can avoid detail loss and thus stabilize training. (See experiments section for more details).

3.3 Stochastic Actor-Executor-Critic

To address I2IT, We propose to extend the maximum entropy RL (MERL) framework with an additional executor. Our proposed model has a similar structure with SAC [Haarnoja et al., 2018], but the latter cannot be applied to the I2IT tasks directly since the delayed and sparse rewards making it difficult to handle the inherent high-dimensional states and actions. To this end, we propose to train the actor and executor in a conventional DL-based manner, where the training data are collected from the agent-environmental interaction and change the conditional likelihood following the stochastic policy πφ:

The empirical objective is to minimize the reconstruction error over all samples:

where D is a replay pool, ‖·‖d denotes some distance measures, such as L1 and L2. To synthesis more realistic images, we also extend the actor-executor to a high quality generative model by jointly training an adversarial loss with discriminator D:

Our final DL guided objective can be expressed as

where λrec and λadv are used to balance the reconstruction and adversarial learning. In general, the DL pathway is very flexible and can readily leverage any other advanced DL losses or techniques into this framework.

MERL in I2IT

In RL pathway, the rewards and the soft Q values are used to iteratively guide the stochastic policy improve. Moreover, the latent-action zt, is used to estimate soft state-action value for encouraging high-level policies. Let actor and critic with parameters φ and θ be the function approximators for the policy πφ and the soft Q-function Qθ, respectively. The critic parameters θ are trained to minimize the soft Bellman residual:

where V¯ θ(xt) is the soft state value function which can be computed by

We use a target network Q¯ θ to stabilize training, whose parameters ¯ θ are obtained by an exponentially moving average of parameters of the critic network: ¯ θ ← τ θ + (1 - τ )¯ θ. The actor parameters φ are optimized to learn the policy towards the exponential of the soft Q-function:

In practice, we use two critics with the same structure but different parameters, (θ1, θ2), to mitigate positive bias in policy improvement and accelerate training. Moreover, we automatically tune the temperature hyperparameter α by

where ¯H represents a constant entropy that equals to the negative of the dimension of the latent action spaces.

The complete algorithm of SAEC is described in Algorithm 1. The training process alternates between learning I2IT from DL pathway and policy control from RL pathway. The joint training of DL and RL ensures a fast convergence since the supervised or unsupervised loss is a "highway" to facilitate back-propagation of the gradients to the actor and the executor. As the latent-action zt is concatenated into critic (shown in Figure 1(b)), the RL objective and DL objective can cooperate to train the stochastic actor fast and stable.

4 Experiments

In this section, we demonstrate experimentally the effectiveness and robustness of the SAEC framework on several different tasks of I2IT.

4.1 Face Inpainting

Settings In this experiment, we apply SAEC to the problem of face inpainting, which is to fill in pixels in the central area of a face image with synthesized contents that are semantically consistent with the original face and at the same time visually realistic. Celeba-HQ dataset is used in this study, of which 28, 000 images are used for training and 2, 000 images are used for testing. All images have a missing part of size of 64 × 64 cropped in the center. We compare the SAEC with several recent face inpainting methods, including CE [Pathak et al., 2016], CA [Yu et al., 2018], PEN [Zeng et al., 2019], PIC [Zheng et al., 2019] and RN [Yu et al., 2020].

The actor-executor for our methods uses the same network structure as the encoder-decoder for CE. The network structures of the discriminator also comes from CE but with a different SNGAN loss. Following the previous work [Pathak et al., 2016; Yu et al., 2018; Zheng et al., 2019], we use PSNR and SSIM as the evaluation metrics.

Results and Analysis As shown in Table 1, by using L1 and adversarial loss in the DL guided learning, our method achieves the best PSNR and SSIM scores comparing with the existing state-of-the-art methods. And as mentioned in the Method section, the reward function is very flexible for our RL framework. Both the PSNR and SSIM based reward are suitable for SAEC and can improve the performance on face inpainting. The qualitative comparison shown in Figure 2 illustrates that the SAEC gives obvious visual improvement for synthesizing realistic faces. The results of SAEC are very reasonable, and the generated faces are sharper and more natural. This may attribute to the high-level action zt, which focuses on learning the global semantic structure and then directs the executor with DL guided learning to further improve the local details of the generated image. We can also see that the synthesised images of the SAEC could have very different appearances from the ground truth, which indicates that although our training is based on paired images, the SAEC can successfully explore and exploit data for producing diverse results.

Ablation Study To illustrate the stability of training GANs in SAEC, we jointly use L1 and several advanced GAN techniques i.e. WGAN-GP [Gulrajani et al., 2017], RaGAN

[Jolicoeur-Martineau, 2018], and SNGAN [Miyato et al., 2018] for DL guided learning. To validate the effectiveness of RL pathway in SAEC for I2IT, we also separately train an actor-executor model (AE) by jointly optimizing the L1 and SNGAN loss. The results shown in Table 2 indicate that the proposed SAEC with different GANs are stable, and significantly improve the performance of training AE with SNGAN alone, which further demonstrates the contribution of the RL formulation.

4.2 Realistic Photo Translation

Settings In this section, we evaluate the SAEC on a set of general I2ITs. We select three datasets of realistic photo translation including:

• CMP Facades dataset for segmentation labels→images [Tyleˇcek and ˇS´ara, 2013].
• Cityscapes dataset for segmentation labels→images and images→labels [Cordts et al., 2016].
• Edge and shoes dataset for edges→shoes [Yu and Grauman, 2014].

In these tasks, SAEC uses the same network structure as in face inpainting experiment with PSNR reward and SNGAN loss. We compare our method with pix2pix [Isola et al., 2017] and PAN [Wang et al., 2018a]. We also compared with pix2pixHD [Wang et al., 2018b], DRPAN [Wang et al., 2019] and CHAN [Gao et al., 2021], which are designed for highquality I2IT. Moreover, we replace MERL with PPO [Schulman et al., 2017] in the RL pathway as SAEC∗ . We use PSNR, SSIM and LPIPS [Zhang et al., 2018] as the evaluation metrics.

Quantitative Evaluation The quantitative results are shown in Table 3. With a similar network structure, the proposed method significantly outperforms the pix2pix and PAN model on PSNR, SSIM and LPIPS over all datasets and tasks. The SAEC even achieves a comparable or better performance than the high-quality pix2pixHD and DRPAN model, which have much more complex architectures and training strategies. Moreover, using MERL instead of PPO obviously improves performance on most tasks. These experiments illustrate that the proposed SAEC is a robust and effective solution for I2IT.

Qualitative Evaluation The qualitative results of our SAEC with other I2IT methods on different tasks are shown in Figure 3. We observe that the pix2pix and PAN sometimes suffer from mode collapse and yield blurry outputs. The pix2pixHD is unstable in different datasets especially on Facades and Cityscapes. The DRPAN is more likely to produce blurred artifacts in several parts of the predicted image on Cityscapes. In contrast, the SAEC produces more stable and realistic results. Using stochastic policy and MERL helps to explore more possible solutions so as to seek out the best generation strategy by trial-and-error in the training steps, leading to a more robust agent for different datasets and tasks.

Comparison with other RL Algorithms The key components of SAEC are substituted by other structures or other state-of-the-art RL algorithms to test their importance. We use DDPG and PPO to demonstrate the effectiveness of stochastic policy and maximum entropy RL, respectively. The learning curves of different variants on four tasks are shown in Figure 4. By using the stochastic policy and the maximum entropy framework, the training is significantly improved.

5 Conclusion

In this paper, we present a new RL framework for highdimensional continuous control and apply it to I2IT problems. This framework contains a stochastic actor-executor critic (SAEC) structure that handles high dimensional policy control, high-level state representation, and low-level image generation simultaneously. We train SAEC by jointly optimizing the RL and the DL objective, which provides ubiquitous and instantaneous supervision for learning I2IT. Our empirical results on several different vision applications demonstrate the effectiveness and robustness of the proposed SAEC. We will subsequently extend this work to more computer vision tasks involving with high-dimensional continuous spaces.

Acknowledgements

This work was supported in part by the National Natural Science Foundation of China under Grant 61602065, Sichuan province Key Technology Research and Development project under Grant 2021YFG0038.