#  Make Encoder Great Again in 3D GAN Inversion through Geometry and
Occlusion-Aware Encoding

Ziyang Yuan1∗ Yiming Zhu1,2∗ Yu Li2† Hongyu Liu3 Chun Yuan1†
1Tsinghua Shenzhen International Graduate School 2International Digital
Economy Academy (IDEA)
3Hong Kong University of Science and Technology
https://eg3d-goae.github.io

###### Abstract

3D GAN inversion aims to achieve high reconstruction fidelity and reasonable
3D geometry simultaneously from a single image input. However, existing 3D GAN
inversion methods rely on time-consuming optimization for each individual
case. In this work, we introduce a novel encoder-based inversion framework
based on EG3D, one of the most widely-used 3D GAN models. We leverage the
inherent properties of EG3D’s latent space to design a discriminator and a
background depth regularization. This enables us to train a geometry-aware
encoder capable of converting the input image into corresponding latent code.
Additionally, we explore the feature space of EG3D and develop an adaptive
refinement stage that improves the representation ability of features in EG3D
to enhance the recovery of fine-grained textural details. Finally, we propose
an occlusion-aware fusion operation to prevent distortion in unobserved
regions. Our method achieves impressive results comparable to optimization-
based methods while operating up to 500 times faster. Our framework is well-
suited for applications such as semantic editing. ††$*$ Equal contributions;
${\dagger}$ Corresponding author.

Figure 1:  The inversion and editing results of our method on human portraits
and cat images. The novel view synthesized images demonstrate that our method
can achieve high-quality inversion and editing given a single image input.

##  1 Introduction

Recent advances in 3D-aware generative adversarial networks (GAN) (_e.g_. [5,
6, 15, 35]) have integrated neural radiance fields (NeRF) [30] into style-
based generation, resulting in remarkable success in generating highly
realistic 3D portraits. Similar to 2D GAN inversion techniques, 3D GAN
inversion also projects images into the latent space of pre-trained 3D GAN
models. However, in addition to pursuing faithful image reconstruction and
preserving fined details, 3D GAN inversion also aims for accurate 3D geometry,
which is critical for generating realistic and consistent multi-view images.
This ability to ensure the correctness of 3D geometry is a crucial advantage
of 3D GAN inversion over its 2D counterpart.

Some optimization-based methods, such as [23, 49, 51], first invert the image
to a latent code to obtain a coarse 3D shape. They then synthesize pseudo-
multi-view images with this shape using a warping process. These pseudo images
are used as supervision to optimize the parameters of the 3D GAN for more
precise 3D geometry and better detail preservation. While these iterative
optimization-based pipelines can prevent shape collapse and produce accurate
results, the optimization process needs to be performed for each individual
case, which is time-consuming and impractical for many applications.

Several encoder-based methods [24, 25] leverage a pre-trained 3D GAN to
generate proxy data by randomly sampling latent codes and camera poses. They
then employ these paired proxy data to train an encoder in a self-supervised
fashion and further refine features in 3D GAN. These encoder-based techniques
offer higher efficiency but lower high-fidelity reconstruction ability.
Moreover, due to a domain gap between the paired proxy data and real images,
the encoder may invert the in-the-wild images to unsuitable latent codes. To
address this issue, we study the properties of the latent space in 3D GAN and
develop a geometry-aware encoder that predicts the latent code incorporating
these properties. Furthermore, since the latent code can not preserve details
of input images well, we explore the characteristics of feature space in 3D
GAN to help our method improve the performance in reconstructing the original
appearance. The latent and feature spaces are widely discussed in 2D GAN
inversion [12, 26, 44].

This paper presents a novel encoder-based method for inverting 3D images
generated by EG3D [5], a well-known and high-quality 3D GAN. Our approach is
simple yet effective and is based on the discovery of a specific latent space
within EG3D, which we call the canonical latent space $\mathcal{W}_{c}$. We
find that latent codes belonging to this space can produce human faces with
superior 3D shape and texture quality following two principles. Firstly, the
input pose before the mapping network of EG3D must be a static canonical pose
throughout the generation process. Secondly, the depth of the background
should fall within a fixed range. Based on this, we propose using an encoder
to invert the image into this canonical latent space using a proposed latent
code discriminator and depth regularization. Although our encoder produces a
suitable latent code that preserves the reasonable 3D shape of the input
image, but the code is low dimension, making it challenging to capture texture
details. To address this, similar to the 2D GAN inversion [3, 47], we refine
the selected features in EG3D based on the residual between the input image
and reconstruction image with predicted latent code. Meanwhile, we also find
there is a canonical feature space $\mathcal{F}_{c}$ in EG3D, so we design an
alignment module to project the refined feature to the canonical feature
space. Moreover, we introduce an occlusion-aware fusion operation to avoid
texture and shape distortions in reconstructing invisible regions.

We conduct comprehensive experiments to evaluate the efficiency and
effectiveness of our method, and the results demonstrate that our method
achieves high-quality inversion comparable to optimization-based methods and
superior to other encoder-based methods. We also demonstrate our approach is
effective in 3D-aware editing applications. See Fig. 1 for some examples of
inversion and editing of human portraits and cat faces obtained by our method.
Our contribution can be summarised as following:

  * •

We conduct an exploration on the latent space and feature space of EG3D and
discover the presence of a canonical attribute.

  * •

Based on our analysis of the characteristics of the canonical latent space, we
propose a geometry-aware encoder by utilizing a canonical latent discriminator
and depth regularization.

  * •

We introduce an alignment module based on the canonical feature space and an
occlusion-aware method to refine the selected features of the generator and
supplement high-quality details.

  * •

The proposed method exhibits competitive performance with existing methods in
human portraits, both qualitatively and quantitatively. Additionally, our
method can also generalize to the cat faces and is effective in 3D editing
tasks.

##  2 Related Work

###  2.1 Generative 3D-Aware Image Synthesis.

Generative Adversarial Network [14] made significant advancements in
generating both 2D photorealistic images [21, 22] and 3D-Aware consistent
images. In 3D-aware image generation, some works build 3D GAN based on
explicit 3D representations, such as point-cloud [43], voxel [31, 32, 45], and
mesh [13, 16, 19]. However, these models often fail to generate high-
resolution 3D scenes with high-quality details.

Since NeRF [30] achieved outstanding multi-view image synthesis, a sequence of
works has begun to use implicit 3D representation and volume rendering for 3D
generation [5, 6, 10, 15, 33, 35]. Pi-GAN [6] is the first attempt at using
GAN to generate NeRF for 3D-aware image synthesis. After that, StyleNeRF [15]
proposes a view-consistent upsampling module to generate high-resolution
images. StyleSDF [35] combines an SDF-based volume renderer and a 2D CNN
network to obtain geometry information. GRAM [10] introduces 2D manifolds to
regulate point sampling strategy in volume rendering. EG3D [5] leverages
StyleGAN2 [22]’s structure and designs a computationally efficient 3D
representation called tri-plane. It can synthesize high-resolution view-
consistent images and produces high-quality 3D geometry. We explore EG3D’s
latent space and build our inversion encoder based on EG3D.

###  2.2 GAN Inversion

GAN Inversion is an inverse process of the generation that projects the given
image to corresponding latent code for faithful reconstruction from the source
image. Thus, the image editing process can be conducted by semantic editing on
the latent code [36, 40, 54]. In 2D GAN, inversion methods are mainly divided
into three categories. Optimization-based methods [1, 8, 48] optimize the
generator’s latent code or parameters both, achieve high reconstruction
quality, but are time-consuming. Encoder-based methods [2, 38, 44] train an
encoder to project input images to the latent codes. These methods can get
faster inference than optimization-based methods but unsatisfactory inversion
fidelity. Hybrid methods [4, 39, 54] use an encoder to get a latent as
initialization, then optimize the latent code and generator parameters, which
is still time-consuming.

Unlike 2D methods, 3D GAN Inversion should also consider multi-view-
consistence. Directly transferring 2D methods to the 3D inversion is prone to
suffering geometry collapse and artifacts. To solve it, Pose Opt. [23] jointly
optimizes camera pose, latent codes, and generator parameters. HFGI3D [49] is
an optimization-based method that extracts mesh to create pseudo-multi-view
supervision. Meanwhile, encoder methods are also gradually developing. IDE-3D
[42] uses semantic segmentation to re-train EG3D and learns an encoder for
inversion to edit images. NeRF-3DE [25] uses identity contrastive learning to
train a based encoder and then trains a refining encoder for better details.
E3DGE [24] trains a base encoder to learn coarse shape and texture and
supplements local details with pixel-aligned features. However, existing
encoder-based methods [24, 25] rely on synthetic paired data and may not
generalize well to real-world scenarios. None of the state-of-the-art methods
can handle time efficiency and 3D inversion quality simultaneously.

Figure 2: Recap of EG3D. $z,p$ denotes a Gaussian sample and a pose condition.
The relation of $w$, $F$, and tri-plane is illustrated.

##  3 EG3D Exploration

In this section, we first briefly review the EG3D model. Then, we explore and
analyze the properties of latent space and feature space in EG3D.

The pipeline of EG3D is shown in Fig 2. It first combines a Gaussian vector
$z$ and a pose condition $p$ as input and builds a mapper to generate a latent
code $w$, this process can be denoted as $w=mapping(z,p)$. Then this latent
code $w$ will be sent to each layer of a style-based hierarchical generator to
obtain the feature maps $F$. These feature maps are enlarged layer by layer.
The last layer’s feature map reshape to $tri\text{-}plane$, an efficient 3D
representation consists of three axis-aligned feature planes ($F_{xy}$,
$F_{xz}$, $F_{yz}$). With this efficient 3D representation, EG3D samples point
features from the tri-plane to synthesize multi-view images or depths by a
rendering decoder $R_{dec}$ with volume rendering [30]. Similar to the 2D GAN,
the latent code $w$ belongs to a latent space $\mathcal{W}$, and all feature
maps in the generator construct the $\mathcal{F}$ space.

Figure 3: Illustration of canonical Pose condition. We synthesize multi-view
images for three pose conditions $p$, including left, right, and canonical. As
shown in a $\sim$ c, the canonical pose condition in row L2 has no distortion
face shape. Moreover, depth maps in row L2 demonstrate that images on the
canonical pose condition have significant foreground and background
differences. In contrast, flattened shape depth maps in rows L1 and L3 have
foreground and background confusion. Columns g $\sim$ i reveal the tri-plane
is also related to the pose condition.

###  3.1 Canonical Latent Space $\mathcal{W}_{c}$

In order to investigate the properties of the latent space $\mathcal{W}$ and
its relationship with pose conditioning, we adopt a fixed $z$ and sample
various pose conditions to generate corresponding latent code $w$. Fig. 3
illustrates that when $w$ is conditioned on an unsuitable pose $p$ far away
from the canonical pose, the generated 3D shapes fail to generalize to other
views synthesis. Only when the pose condition is canonical can the synthesized
images maintain shape and texture consistency across multiple views.
Meanwhile, we find a clear depth gap between the face and background regions
when generating good results with the canonical pose condition. Under
canonical pose, the depth of the background is restricted to a fixed range,
while the face and background regions are confused for the other pose
conditions.

These observations lead us to identify a specific subspace of $\mathcal{W}$
intuitively, which we call it canonical latent space, denoted as
$\mathcal{W}_{c}$, that is responsible for ensuring reasonable 3D shape and
texture across multiple views. There are two principles of this canonical
latent space. Firstly, the input pose before the mapping network of EG3D must
be a static canonical pose. Secondly, the depth of the background should fall
within a fixed range to distinguish the face and background regions. Based on
these findings, we propose that a high-quality 3D-GAN inversion method should
first retrieve a $w$ latent code in $\mathcal{W}_{c}$ to ensure the best
possible image synthesis results.

###  3.2 Canonical Feature Space $\mathcal{F}_{c}$

To analyze the characteristics of $\mathcal{F}$ space, we generate the latent
codes with different pose conditions to synthesize multiple tri-planes. We
visualize the tri-plane as illustrated in Fig 3. The tri-plane also represents
a canonical view when the pose condition is canonical (_i.e_., the columns g
$\sim$ i in row L2). It indicates that there is also a specific subspace of
$\mathcal{F}$, we call it canonical feature space, denoted as
$\mathcal{F}_{c}$, and the canonical latent space $\mathcal{W}_{c}$ produces
the canonical feature space $\mathcal{F}_{c}$ with the generator layers. The
$\mathcal{F}_{c}$ can supply sufficient facial information for rendering
multi-view images.

##  4 Method

The pipeline of our method is shown in Fig. 4. We first project the input
image $I$ to the latent code $w^{+}\subset\mathbb{R}^{14\times 512}$ with our
encoder, where $w^{+}=(w_{0},w_{1},...,w_{13})$ is an extension of latent code
$w$ and can preserve the details better as mentioned in 2D GAN inversion [38,
44]. Then we calculate the residual between the $I$ and the generated image
$I_{w^{+}}$ based on $w^{+}$ to refine the features of the generator. In the
following Sec. 4.1 and Sec. 4.2, we will describe how we ensure the latent
code $w^{+}$ and refined feature belong to the canonical space, respectively.
Furthermore, we introduce an occlusion-aware mix tri-plane in Sec. 4.3 to
avoid the distortion of invisible regions reconstruction. In Sec. 4.4, we
exhibit how to perform 3D semantic editing with our inversion method.

###  4.1 Geometry-aware Encoder

We design a geometry-aware encoder to obtain the latent code $w^{+}$ given the
input image $I$ and corresponding rendering camera pose $c$. We generate the
inverted latent code $w^{+}$ according to the backbone’s pyramid feature
progressively. It is worth noting that we choose Swin-transformer [28] as the
backbone and further add attention modules at different scale feature layers
for different level latent codes (more details of encoder architecture are in
supplementary materials). After projecting the input image to the latent code
$w^{+}$, we can synthesize the $tri\text{-}plane_{w^{+}}$ through EG3D
generator $G$ and use rendering decoder $R_{dec}$ to get the reconstruction
image $I_{w^{+}}$ under the rendering camera pose $c$.

$\displaystyle w^{+}$ $\displaystyle=E(I),$ (1) $\displaystyle I_{w^{+}}$
$\displaystyle=R_{dec}(G(w^{+}),c)$

Figure 4: Overview of our method. Our framework could be divided into two
parts. (1) $\mathcal{W}$ space inversion. We design an encoder $E$ to invert
input image $I$ into $w^{+}$ latent codes. The $w^{+}$ latent codes are fed
into a pre-trained EG3D generator $G$ to get $tri$-$plane_{w^{+}}$ and
rendered into reconstruction image $I_{w^{+}}$. (2) Complement the
$\mathcal{F}$ space. We calculate the image residual $\Delta I$ between the
input image and its reconstruction and propose AFA module to refine the F
latent maps. The modified latent maps $F^{*}$ are transformed into
$tri$-$plane_{mix}$ by occlusion-aware mix and rendered into the fine detailed
inversion image $I_{F^{*}}$.

To make the inverted $w^{+}$ not deviate far from the canonical space
$\mathcal{W}_{c}$ discussed in Sec. 3.1, we use a canonical latent
discriminator [34] to discriminate between $w_{c}\in\mathcal{W}_{c}$ and
$w^{+}$. The canonical latent discriminator, denoted as
$\mathcal{D}_{\mathcal{W}_{c}}$, is trained in an adversarial scheme with the
non-saturating GAN loss [14] and $R_{1}$ regularization [29], the loss of
adversarial learning process can be written as below:

$\begin{gathered}\mathcal{L}^{\mathcal{D}_{\mathcal{W}_{c}}}_{\text{adv}}=-{\mathbb{E}}[\log\mathcal{D}_{\mathcal{W}_{c}}(w_{c})]-{\mathbb{E}}[\log(1-\mathcal{D}_{\mathcal{W}_{c}}(w_{i}))]\\\
+\frac{\gamma}{2}{\mathbb{E}}\left[\left\|\nabla\mathcal{D}_{\mathcal{W}_{c}}(w_{c})\right\|_{2}\right],\end{gathered}$
(2)
$\mathcal{L}_{\text{adv}}^{E}=-{\mathbb{E}}\left[\log{\mathcal{D}}_{\mathcal{W}_{c}}(w_{i})\right],$
(3)

where $\gamma$ is the weight of $R_{1}$ regularization, the
$\mathcal{L}^{{\mathcal{D}}_{\mathcal{W}_{c}}}_{\text{adv}}$ and
$\mathcal{L}_{\text{adv}}^{E}$ are the loss for the latent discriminator
$\mathcal{D}_{\mathcal{W}_{c}}$ and encoder respectively.
$\mathcal{D}_{\mathcal{W}_{c}}$ operates on $w_{c}$ and each entry $w_{i}$ of
$w^{+}$.

As discussed in Sec. 3.1, the background depth of an integrated face should be
limited in a range. To estimate the background depth, we first random sample
0.1 million front images with the canonical pose condition, then we obtain the
depth maps and background regions of these images. We use the parsing method
[55] to get a binary mask $M$, where $M=1$ indicates the background region and
0 otherwise. Finally, we calculate the average depth $D_{avg}$ of background
regions of these samples, and we set the $D_{avg}$ as a regularization to
constrain the background depth during inversion, the background regularization
can be written as below:

$\mathcal{L}_{BG}=\left\|D\odot M-D_{avg}\odot M\right\|_{2}.$ (4)

Here, $D$ and $M$ are the depth map and background mask inverted image under
the front-view pose.

In addition, we use reconstruction loss function $\mathcal{L}_{rec}$ follow
pSp [38] in the training stage. The $\mathcal{L}_{rec}$ is composited by
pixel-wise $\mathcal{L}_{2}$ loss, perceptual loss $\mathcal{L}_{LPIPS}$ [53],
and facial identity similarity loss $\mathcal{L}_{ID}$ with pre-trained
ArcFace network [9]:

$\displaystyle\mathcal{L}_{rec}$
$\displaystyle=\left\|I_{w^{+}}-I\right\|_{2}+\mathcal{L}_{LPIPS}(I_{w^{+}},I)+\mathcal{L}_{ID}(I_{w^{+}},I).$
(5)

In summary, the total loss for training our encoder can be written as:

$\mathcal{L}_{total}=\mathcal{L}^{\mathcal{D}}_{\text{adv}}+\mathcal{L}_{\text{adv}}^{E}+\mathcal{L}_{BG}+\mathcal{L}_{rec}.$
(6)

###  4.2 Adaptive Feature Alignment

It is non-trivial to restore fine-detailed texture with only a highly
compressed $w$ latent code. Thus, adjusting high-dimension feature maps in the
$\mathcal{F}$ space is natural to achieve high-fidelity inversion as mentioned
in related 2D GAN inversion methods [3, 12, 47]. Moreover, it is worth noting
that the feature maps $F$ generated with our inverted $w^{+}$ are in the
canonical feature space $\mathcal{F}_{c}$.

We calculate the residual $\Delta I$ between the $I$ and $I_{w^{+}}$ and hope
to refine the feature maps to eliminate this residual. However, directly
setting the residual as input and predicting the refined feature maps with CNN
to modify $F$ is invalid. This is because the residual is highly entangled
with the view of the input image, while $F$ is in $\mathcal{F}_{c}$. When the
view is not canonical, the refined feature maps with this residual can not
project to the $\mathcal{F}_{c}$ space simply. Thus, aligning the residual
feature maps to the $\mathcal{F}_{c}$ space is necessary to synthesize high-
quality images, as mentioned in Sec. 3.2.

According to the above analysis, we propose an adaptive feature alignment
(AFA) module that enables automatic canonical space alignment by cross-
attention mechanism, as shown in Fig 5. First, we use a CNN to capture the
feature $F_{\Delta I}$ of $\Delta I$, and we take $F_{\Delta I}$ as the key
$K$, value $V$, and the feature map $F$ as a query $Q$. Then, we utilize
cross-attention modules to align the $F_{\Delta I}$ with $F$ and output a
feature $F_{align}$. Since the $F$ is already in the canonical feature space
$\mathcal{F}_{c}$, and the feature $F_{align}$ can also have the properties of
$\mathcal{F}_{c}$ space. Finally, we use the $F_{align}$ to predict the scale
map $\gamma$ and translation map $\beta$ to refine the $F$ with Feature-wise
Linear Modulation (FiLM) [37]. The whole process of our AFA module can be
formulated as follows:

$\displaystyle F_{\Delta I}=CNN(\Delta I),$ (7) $\displaystyle
Q,K,V=W_{Q}F,W_{K}F_{\Delta I},W_{V}F_{\Delta I},$ $\displaystyle
F_{align}=Attention(Q,K,V),$
$\displaystyle\gamma,\beta=Conv_{\gamma}(F_{align}),Conv_{\beta}(F_{align}),$
$\displaystyle F^{*}=\gamma\odot F+\beta,$

where $W_{Q},W_{K},W_{V}$ are learnable projection heads and $d$ is the
dimension of feature, the $F^{*}$ is the final output of our AFA module. With
this adaptive alignment, we achieve a detailed complement of $\mathcal{F}_{c}$
space. Note that we adopt $\mathcal{L}_{total}$ and a $\mathcal{L}_{2}$
regularization loss on $\Delta F=F^{*}-F$ when training the AFA module.

Figure 5: The Architecture of AFA Module. We use a CNN to capture the feature
$F_{\Delta I}$ of residual image ${\Delta I}$, and align the $F_{\Delta I}$ to
$F$ with cross attention module. Finally, we refine the $F$ with the aligned
feature $F_{align}$.

###  4.3 Occlusion-aware Mix Tri-plane

The image residual $\Delta I$ only contains information in the visible area of
the input view. Using single-view residuals to modify the entire $F$ will lead
to distortion at the occlusion region. To solve this problem, we propose the
Occlusion-aware Mix Tri-plane.

The tri-plane contains all points in the 3D space which are divided into
visible points set and occlusion points set according to the input view. Here,
we define the 3D points viewed in the input image in tri-plane as the visible
points set $\mathcal{V}_{(x,y,z)}$, points in the occlusion regions as
occlusion points set $\mathcal{O}_{(x,y,z)}$. We can get the depth map $D$
under the input view, and the two points set can be defined by:

$\displaystyle\mathcal{V}_{(x,y,z)}=\\{[x,y,z]^{T}:$
$\displaystyle[x,y,z]^{T}=\pi D(u,v)\mathcal{K}^{-1}[u,v,1]^{T}\\},$ (8)
$\displaystyle\mathcal{O}_{(x,y,z)}$
$\displaystyle=\mathbb{R}^{3}\setminus\mathcal{V}_{(x,y,z)},$

where $[u,v,1]$ is the pixel homogeneous coordinate of the input image, $\pi$
is the extrinsic camera parameter matrix of the input view, $D(u,v)$ is the
depth of pixel $[u,v]$, $\mathcal{K}$ is the intrinsic camera parameter
matrix.

After determining the visible area, we can supplement the refined details of
the input view to the tri-plane. We define the occlusion-aware mix tri-plane
as follows:

$\displaystyle tri\text{-}plane_{w^{+}}$ $\displaystyle=G(w^{+}),$ (9)
$\displaystyle tri\text{-}plane_{F^{*}}$ $\displaystyle=G(F^{*}),$
$\displaystyle tri\text{-}plane_{mix}$
$\displaystyle=tri\text{-}plane_{F^{*}}|_{(x,y,z)\in\mathcal{V}}$
$\displaystyle+tri\text{-}plane_{w^{+}}|_{(x,y,z)\in\mathcal{O}},$

where $w^{+}$ is the latent code learned from the encoder, and $F^{*}$ is the
refined feature map. Finally, we use $tri\text{-}plane_{mix}$ with rendering
decoder to synthesize images under different views when inference.

###  4.4 Editing

To achieve 3D-consistent editing, we need to modify both $w^{+}$ and $F^{*}$.
We first get the edited latent code $\hat{w}^{+}$ through classic latent code
editing methods [36, 40]. Then for $F^{*}$, we follow FS [50] to extract the
features $F^{\hat{w}^{+}}$ and $F^{w}$ from the generator convolution layer
with input $\hat{w}^{+}$ and $w$, respectively. Finally, we calculate the
difference between these two features and add it to the $F^{*}$ to get the
edited features $\hat{F^{*}}$:

$\displaystyle\hat{F^{*}}=F^{*}+F^{\hat{w}^{+}}-F^{w}.$ (10)

With the modified $\hat{w}^{+}$ and $\hat{F^{*}}$, we use the method in
section 4.3 to get edited mix tri-plane, and use rendering decoder to
synthesize 3D-consistent edited images.

Figure 6: Visual comparison of inversion and novel-view synthesized results
between our method and others. Our method surpasses all the encoder-based
methods and is competitive with optimization methods in recovering texture
details, _i.e_. our method prevents the eye orientation from changing with the
camera view and artifacts in the novel view.

##  5 Experiments

Experiment setup. We evaluate our method on human faces and cat faces. For
human faces, we use FFHQ [1] dataset for training, and CelebA-HQ [20] and MEAD
[46] for evaluation. For the cat faces, we use AFHQ-CAT [7] for training and
evaluation. In addition, we employ the pre-trained EG3D generator and obtain
the corresponding camera parameters of each image following the official
implementation in EG3D.

Baselines. We compare our method with encoder (IDE-3D [42]) and optimization-
based (PTI [39] and Pose Opt. [23]) baselines in these two aspects. Moreover,
we build two encoder baselines imitating the state-of-the-art StyleGAN
inverters, pSp [38] and e4e [44], and present the model structure in the
supplementary.

###  5.1 Qualitative analysis

Source View Inversion Results. Fig. 6 shows the inversion performance
comparison between ours and other methods. Our method’s hair details and
facial similarity are obviously competitive with optimization methods and
surpass all the encoder-based methods.

Novel View Synthesis Results. The comparison results are shown in Fig. 6.
Inspecting the eye gaze direction and person’s identity preservation when
changing to novel views can reveal the differences. It can be seen that on
some extreme side-face camera poses, some optimization methods produce
artifacts and geometry distortions. Our method can keep synthesizing
3D-consistent images faithful to the person’s identity in the input.

Table 1:  Comparison with state-of-the-art methods on inversion performance.
The numbers in bold are the best performance of the optimization and encoder-
based methods, respectively. Metrics of Reconstruction Quality are calculated
under the input view. Metrics of Novel View Quality are the average of multi-
view. The yaw angle is the source view angle.

|  | Reconstruction Quality | Novel View Quality
---|---|---|---
Category | Method | Input View |  yaw=$\pm 60^{\circ}$ |  yaw=$0^{\circ}$
|  |  MSE $\downarrow$ |  LPIPS $\downarrow$ |  FID $\downarrow$ |  ID $\uparrow$ |  Geo. Err. $\downarrow$ |  Time (s)$\downarrow$ |  LPIPS $\downarrow$ |  FID $\downarrow$ |  ID $\uparrow$ |  Geo. Err. $\downarrow$ |  LPIPS $\downarrow$ |  FID $\downarrow$ |  ID $\uparrow$ |  Geo. Err. $\downarrow$
Optimization |  PTI [39] | 0.0163 | 0.0918 | 18.0 | 0.8417 | 0.0850 | 115.2 | 0.2926 | 69.8 | 0.621 | 0.1481 | 0.1925 | 53.5 | 0.718 | 0.1006
Pose Opt. [23] | 0.0164 | 0.0879 | 17.5 | 0.8390 | 0.0759 | 219.7 | 0.4477 | 132 | 0.321 | 0.1763 | 0.2436 | 60.2 | 0.605 | 0.1504
Encoder-based |  IDE-3D [42] | 0.1056 | 0.2806 | 43.6 | 0.7194 | 0.1835 | 0.0659 | 0.3317 | 79.0 | 0.505 | 0.2914 | 0.3074 | 63.7 | 0.513 | 0.2528
pSp-3D [38] | 0.0726 | 0.2082 | 31.2 | 0.7817 | 0.1846 | 0.1284 | 0.3223 | 65.7 | 0.591 | 0.1287 | 0.3132 | 65.2 | 0.671 | 0.1789
e4e-3D [44] | 0.0808 | 0.2281 | 33.0 | 0.7359 | 0.1161 | 0.1303 | 0.2495 | 57.7 | 0.588 | 0.1269 | 0.2394 | 53.8 | 0.653 | 0.1305
$\text{Ours-}{w^{+}}$ | 0.0697 | 0.1981 | 30.9 | 0.7873 | 0.1100 | 0.0963 | 0.2303 | 52.2 | 0.655 | 0.0943 | 0.2372 | 54.1 | 0.685 | 0.0955
Ours | 0.0301 | 0.1257 | 20.5 | 0.8752 | 0.0987 | 0.2262 | 0.2192 | 50.9 | 0.700 | 0.0950 | 0.2204 | 51.7 | 0.743 | 0.0944

###  5.2 Quantitative analysis

We calculate the mean square error (MSE), LPIPS [53], and FID [17] between the
source image and its inversion. We adopt a pre-trained face key point detector
parsenet [55] to calculate the identity similarity. To measure the geometry
properties, we followed [5, 41], utilizing a pre-trained 3D face
reconstruction model to extract a “pseudo” ground truth depth map from the
source image. Then, we obtain the depth of inversion and normalize all depth
maps to zero means and unit variance to calculate geometry error (Geo. Err.)
by using average $\mathcal{L}_{2}$ distance.

Source View Inversion Results. Table 3 presents the quantitative comparison of
the input view image reconstruction on CelebA-HQ test dataset. Our method is
two orders of magnitude faster than the optimized methods with similar
performance, which is notably better than other encoder-based methods.

Novel View Synthesis Results. In order to measure the multi-view consistency
and the synthesized quality, we evaluate our method and baselines with a
multi-view dataset, MEAD [46]. Each person in MEAD has five images with
increasing yaw angles. We select one angle as the input view, invert the input
image and synthesize the other four novel views. Then we use the same metrics
to calculate the difference between the synthesized images and the ground
truth. The quantitative comparison between the novel view synthesized results
and the ground truth in different views is shown in Table 3. Our method shows
similar performance with PTI when the input image is a front face (yaw =
$0^{\circ}$). When inverting side faces, the difficulty of inversion increases
drastically when the yaw angle increases. We calculated the metrics of the
largest yaw angle in MEAD (yaw = $\pm 60^{\circ}$), and our model has robust
performance when the test yaw angle varies. However, the performance of
optimization methods degrades remarkably at an extreme angle.

###  5.3 Ablation Analysis

We present ablation studies with qualitative analysis of our proposed method
below, and the corresponding numerical ablation results are in our
supplementary.

Geometry-aware Encoder. We conduct ablation experiments and make qualitative
comparisons to verify that our proposed canonical latent discriminator
($\mathcal{D}_{\mathcal{W}_{c}}$) and background regularization
($\mathcal{L}_{BG}$) contribute to restoring 3D geometry. Generally,
recovering the geometry from a side face input is more challenging and we test
such a case. As shown in Fig. 7, only if we adopt the canonical discriminator
and the background regularization could we obtain a satisfactory foreground
and background geometry.

Figure 7: Ablation of Geometry-aware Encoder. All images are synthesized only
with $w^{+}$. Without both modules, $w^{+}$ fails to generate integrated 3D
geometry and texture. The $L_{BG}$ helps distinguish foreground and
background. With $\mathcal{D}_{\mathcal{W}_{c}}$, we can get good inversion
under the input view but fail under the novel view. With both modules, $w^{+}$
can obtain an integrated side face of the woman, which is occlusion under the
input view.

Adaptive Feature Alignment. We proposed Adaptive Feature Alignment (AFA) to
refine the $F$ latent map and add texture details from the source view. Fig. 8
demonstrates that we can recover more details with AFA.

Occlusion-aware Mix Tri-plane. To verify the necessity of our proposed
occlusion-aware mix tri-plane, we ablate this operation and visualize the
novel view results in Fig. 9. The occlusion map shows the spatial position of
surface points in ray sampling, the visible region in the original view is
white, and the occlusion region is black. Without the mix tri-plane, the
occlusion area from the source view is easy to produce artifacts in novel
views.

###  5.4 Image Editing Results

Figure 8: Ablation of Adaptive Feature Alignment. We replace AFA with a CNN
network and obtain the results. Without feature alignment, it is difficult to
reconstruct the hair and facial details.  Figure 9: Ablation of Occlusion-
aware mix tri-plane. We visualize the occlusion regions
$\mathcal{O}_{(x,y,z)}$ of the input view in the last column (black points).
The green arrow points out $\mathcal{O}_{(x,y,z)}$ in the novel view, and it
is obvious that the occlusion-aware mix tri-plane can avoid artifacts. Figure
10: Editing result under novel views. We add three different attributes to
three faces. Our method generates both realistic editing and 3D-consistent
novel views.

In GAN inversion-based image manipulation, the editability and the
reconstruction trade-off widely exists in both 2D and 3D scenarios. We
validate the editing method mentioned in Sec. 4.4 and display the multi-view
edited results in Fig. 10. We choose three different edit directions, ’Young’
to reduce person’s age, ’Unhappy’ to change emotion, and ’Glasses’ to add
accessories. It is noticeable that our method not only restores the geometric
and texture information of the input image but also ensures the editing
consistency of multi-view. The edited results demonstrate that our inversion
method can be well applied to the facial image editing tasks.

##  6 Conclusion

In this work, we explore the intrinsic property of EG3D’s latent space and
find that the key to achieve high-quality shape and texture inversion is to
invert input images to the canonical space. In light of this, we design an
efficient geometry-aware encoder with an adaptive feature alignment module and
demonstrate their effectiveness. Besides, we propose an occlusion-aware mix
tri-plane to avoid invisible region distortion in novel views. Extensive
experiments demonstrate that our method has competitive inversion performance
with optimization-based methods, but operates two orders of magnitude faster.
In addition, our method is demonstrated to be a powerful tool for realistic
and 3D view-consistent editing of facial images.

## References

  * [1] Rameen Abdal, Yipeng Qin, and Peter Wonka.  Image2stylegan: How to embed images into the stylegan latent space?  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4432–4441, 2019.
  * [2] Yuval Alaluf, Or Patashnik, and Daniel Cohen-Or.  Restyle: A residual-based stylegan encoder via iterative refinement.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 6711–6720, 2021.
  * [3] Yuval Alaluf, Omer Tov, Ron Mokady, Rinon Gal, and Amit Bermano.  Hyperstyle: Stylegan inversion with hypernetworks for real image editing.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 18511–18521, 2022.
  * [4] Lucy Chai, Jun-Yan Zhu, Eli Shechtman, Phillip Isola, and Richard Zhang.  Ensembling with deep generative views.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14997–15007, 2021.
  * [5] Eric R Chan, Connor Z Lin, Matthew A Chan, Koki Nagano, Boxiao Pan, Shalini De Mello, Orazio Gallo, Leonidas J Guibas, Jonathan Tremblay, Sameh Khamis, et al.  Efficient geometry-aware 3d generative adversarial networks.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16123–16133, 2022.
  * [6] Eric R Chan, Marco Monteiro, Petr Kellnhofer, Jiajun Wu, and Gordon Wetzstein.  pi-gan: Periodic implicit generative adversarial networks for 3d-aware image synthesis.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5799–5809, 2021.
  * [7] Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo Ha.  Stargan v2: Diverse image synthesis for multiple domains.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8188–8197, 2020.
  * [8] Edo Collins, Raja Bala, Bob Price, and Sabine Susstrunk.  Editing in style: Uncovering the local semantics of gans.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5771–5780, 2020.
  * [9] Jiankang Deng, Jia Guo, Niannan Xue, and Stefanos Zafeiriou.  Arcface: Additive angular margin loss for deep face recognition.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 4690–4699, 2019.
  * [10] Yu Deng, Jiaolong Yang, Jianfeng Xiang, and Xin Tong.  Gram: Generative radiance manifolds for 3d-aware image generation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10673–10683, 2022.
  * [11] Yu Deng, Jiaolong Yang, Sicheng Xu, Dong Chen, Yunde Jia, and Xin Tong.  Accurate 3d face reconstruction with weakly-supervised learning: From single image to image set.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pages 0–0, 2019.
  * [12] Tan M Dinh, Anh Tuan Tran, Rang Nguyen, and Binh-Son Hua.  Hyperinverter: Improving stylegan inversion via hypernetwork.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11389–11398, 2022.
  * [13] Shubham Goel, Angjoo Kanazawa, and Jitendra Malik.  Shape and viewpoint without keypoints.  In European Conference on Computer Vision, 2020.
  * [14] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  In Advances in Neural Information Processing Systems, 2014.
  * [15] Jiatao Gu, Lingjie Liu, Peng Wang, and Christian Theobalt.  Stylenerf: A style-based 3d-aware generator for high-resolution image synthesis.  arXiv preprint arXiv:2110.08985, 2021.
  * [16] Paul Henderson, Vagia Tsiminaki, and Christoph H Lampert.  Leveraging 2d data to learn textured 3d mesh generation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2020.
  * [17] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  Gans trained by a two time-scale update rule converge to a local nash equilibrium.  Advances in Neural Information Processing Systems, 30, 2017.
  * [18] kairess.  Cat hipsterizer.  https://github.com/kairess/cat_hipsterizer/, 2018.
  * [19] Angjoo Kanazawa, Shubham Tulsiani, Alexei A. Efros, and Jitendra Malik.  Learning category-specific mesh reconstruction from image collections.  In European Conference on Computer Vision, 2018.
  * [20] Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen.  Progressive growing of GANs for improved quality, stability, and variation.  In International Conference On Learning Representations, 2018.
  * [21] Tero Karras, Samuli Laine, and Timo Aila.  A style-based generator architecture for generative adversarial networks.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2019.
  * [22] Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila.  Analyzing and improving the image quality of StyleGAN.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2020.
  * [23] Jaehoon Ko, Kyusun Cho, Daewon Choi, Kwangrok Ryoo, and Seungryong Kim.  3d gan inversion with pose optimization.  arXiv preprint arXiv:2210.07301, 2022.
  * [24] Yushi Lan, Xuyi Meng, Shuai Yang, Chen Change Loy, and Bo Dai.  Self-supervised geometry-aware encoder for style-based 3d gan inversion.  arXiv preprint arXiv:2212.07409, 2022.
  * [25] Yu-Jhe Li, Tao Xu, Bichen Wu, Ningyuan Zheng, Xiaoliang Dai, Albert Pumarola, Peizhao Zhang, Peter Vajda, and Kris Kitani.  3d-aware encoding for style-based neural radiance fields.  arXiv preprint arXiv:2211.06583, 2022.
  * [26] Hongyu Liu, Yibing Song, and Qifeng Chen.  Delving stylegan inversion for image editing: A foundation latent space viewpoint.  arXiv preprint arXiv:2211.11448, 2022.
  * [27] Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han.  On the variance of the adaptive learning rate and beyond.  arXiv preprint arXiv:1908.03265, 2019.
  * [28] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 10012–10022, 2021.
  * [29] Lars Mescheder, Andreas Geiger, and Sebastian Nowozin.  Which training methods for gans do actually converge?  In International Conference on Machine Learning, pages 3481–3490. PMLR, 2018.
  * [30] Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng.  Nerf: Representing scenes as neural radiance fields for view synthesis.  In European Conference on Computer Vision, 2020.
  * [31] Thu Nguyen-Phuoc, Chuan Li, Lucas Theis, Christian Richardt, and Yong-Liang Yang.  Hologan: Unsupervised learning of 3d representations from natural images.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 7588–7597, 2019.
  * [32] Thu H Nguyen-Phuoc, Christian Richardt, Long Mai, Yongliang Yang, and Niloy Mitra.  Blockgan: Learning 3d object-aware scene representations from unlabelled images.  Advances in Neural Information Processing Systems, 33:6767–6778, 2020.
  * [33] Michael Niemeyer and Andreas Geiger.  Giraffe: Representing scenes as compositional generative neural feature fields.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11453–11464, 2021.
  * [34] Yotam Nitzan, Amit Bermano, Yangyan Li, and Daniel Cohen-Or.  Face identity disentanglement via latent space mapping.  arXiv preprint arXiv:2005.07728, 2020.
  * [35] Roy Or-El, Xuan Luo, Mengyi Shan, Eli Shechtman, Jeong Joon Park, and Ira Kemelmacher-Shlizerman.  Stylesdf: High-resolution 3d-consistent image and geometry generation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 13503–13513, 2022.
  * [36] Or Patashnik, Zongze Wu, Eli Shechtman, Daniel Cohen-Or, and Dani Lischinski.  Styleclip: Text-driven manipulation of stylegan imagery.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 2085–2094, 2021.
  * [37] Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville.  Film: Visual reasoning with a general conditioning layer.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
  * [38] Elad Richardson, Yuval Alaluf, Or Patashnik, Yotam Nitzan, Yaniv Azar, Stav Shapiro, and Daniel Cohen-Or.  Encoding in style: a stylegan encoder for image-to-image translation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2287–2296, 2021.
  * [39] Daniel Roich, Ron Mokady, Amit H Bermano, and Daniel Cohen-Or.  Pivotal tuning for latent-based editing of real images.  ACM Transactions on Graphics (TOG), 42(1):1–13, 2022.
  * [40] Yujun Shen, Jinjin Gu, Xiaoou Tang, and Bolei Zhou.  Interpreting the latent space of gans for semantic face editing.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9243–9252, 2020.
  * [41] Yichun Shi, Divyansh Aggarwal, and Anil K Jain.  Lifting 2d stylegan for 3d-aware face generation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 6258–6266, 2021.
  * [42] Jingxiang Sun, Xuan Wang, Yichun Shi, Lizhen Wang, Jue Wang, and Yebin Liu.  Ide-3d: Interactive disentangled editing for high-resolution 3d-aware portrait synthesis.  arXiv preprint arXiv:2205.15517, 2022.
  * [43] Maxim Tatarchenko, Alexey Dosovitskiy, and Thomas Brox.  Multi-view 3D models from single images with a convolutional network.  In European Conference on Computer Vision, 2016.
  * [44] Omer Tov, Yuval Alaluf, Yotam Nitzan, Or Patashnik, and Daniel Cohen-Or.  Designing an encoder for stylegan image manipulation.  ACM Transactions on Graphics (TOG), 40(4):1–14, 2021.
  * [45] Shubham Tulsiani, Hao Su, Leonidas J. Guibas, Alexei A. Efros, and Jitendra Malik.  Learning shape abstractions by assembling volumetric primitives.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2017.
  * [46] Kaisiyuan Wang, Qianyi Wu, Linsen Song, Zhuoqian Yang, Wayne Wu, Chen Qian, Ran He, Yu Qiao, and Chen Change Loy.  Mead: A large-scale audio-visual dataset for emotional talking-face generation.  In European Conference on Computer Vision, pages 700–717. Springer, 2020.
  * [47] Tengfei Wang, Yong Zhang, Yanbo Fan, Jue Wang, and Qifeng Chen.  High-fidelity gan inversion for image attribute editing.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11379–11388, 2022.
  * [48] Zongze Wu, Dani Lischinski, and Eli Shechtman.  Stylespace analysis: Disentangled controls for stylegan image generation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12863–12872, 2021.
  * [49] Jiaxin Xie, Hao Ouyang, Jingtan Piao, Chenyang Lei, and Qifeng Chen.  High-fidelity 3d gan inversion by pseudo-multi-view optimization.  arXiv preprint arXiv:2211.15662, 2022.
  * [50] Xu Yao, Alasdair Newson, Yann Gousseau, and Pierre Hellier.  A style-based gan encoder for high fidelity reconstruction of images and videos.  European Conference on Computer Vision, 2022.
  * [51] Fei Yin, Yong Zhang, Xuan Wang, Tengfei Wang, Xiaoyu Li, Yuan Gong, Yanbo Fan, Xiaodong Cun, Ying Shan, Cengiz Oztireli, et al.  3d gan inversion with facial symmetry prior.  arXiv preprint arXiv:2211.16927, 2022.
  * [52] Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E Hinton.  Lookahead optimizer: k steps forward, 1 step back.  Advances in Neural Information Processing Systems, 32, 2019.
  * [53] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang.  The unreasonable effectiveness of deep features as a perceptual metric.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 586–595, 2018.
  * [54] Jiapeng Zhu, Yujun Shen, Deli Zhao, and Bolei Zhou.  In-domain gan inversion for real image editing.  In European Conference on Computer Vision, pages 592–608. Springer, 2020.
  * [55] zllrunning.  face-parsing.pytorch.  https://github.com/zllrunning/face-parsing.PyTorch/, 2019.

## Appendix

Figure 11: Encoder architecture comparison. We refer to pSp [38], e4e [44]’s
model structure, refine and build up EG3D inverters respectively (see in
(a),(b)). Our designed Swin-transformer based encoder is shown in (c).

##  Appendix A Implementation Details

Dataset Details. We conduct our experiments on two categories of data:
monocular images of human portraits and cat faces. We follow the method in
EG3D [5] to extract the camera parameters of these images with off-the-shelf
pose detectors [11, 18]. For human portraits, we use FFHQ [1] which contains
about 70,000 images to train our model. To evaluate our model’s performance of
input-view reconstruction, we randomly sample 1,500 images from CelebA-HQ [20]
test for quantitative evaluation. Additionally, we use a multi-view dataset
MEAD [46] to evaluate our model’s performance across novel views.
Specifically, we use five views (left$60^{\circ}$, left$30^{\circ}$, front,
right$30^{\circ}$, right$60^{\circ}$) images of 43 persons per frame. We
randomly sample 5 frames for each person. For cat faces, we use AFHQv2 Cats
[7] following EG3D. We split about 5,000 images into train, evaluation, and
test sets by 8:1:1 ratio.

Architecture of Geometry-aware Encoder. Our encoder uses Swin-transformer as
the backbone, and we further design attention modules at different scale
feature layers for different level latent codes. The encoder architecture is
shown in Fig. 11. We split the intermediate output of the Swin-transformer
into four levels “query, coarse, mid, fine” similar to the pyramid
architecture of CNN models. We use “query” to get $w_{0}$ and query
$Q_{coarse},Q_{mid},Q_{fine}$, and leverage “coarse”, “mid”, “fine” to obtain
keys and values $(K,V)_{coarse},(K,V)_{mid},(K,V)_{fine}$. Then the different
level queries with corresponding keys and values are sent into cross-attention
modules to yield different $w_{i}$. Finally, the final latent code $w^{+}$ is
obtained by:

$\displaystyle w^{+}=w_{avg}+w_{0}+(0,w_{1\sim 3},w_{3\sim 6},w_{7\sim 13}).$
(11)

Occlusion-aware Mix Tri-plane. As mentioned in Sec.4.3, We can get the visible
points set $\mathcal{V}_{(x,y,z)}$. Then we perform orthogonal projection that
projects these points to the three axis-aligned feature planes ($F_{xy}$,
$F_{xz}$, $F_{yz}$) of tri-plane to get three masks separately, denotes as
$tri\text{-}Mask$. The grid point in $tri\text{-}Mask$ is equal to 1 if the
corresponding 3D point is in $\mathcal{V}_{(x,y,z)}$, otherwise it is equal to
0. Finally, the mix tri-plane can be obtained by:

$\displaystyle tri\text{-}plane_{mix}$
$\displaystyle=tri\text{-}plane_{F^{*}}\odot tri\text{-}Mask$ (12)
$\displaystyle+tri\text{-}plane_{w^{+}}\odot(tri\text{-}I-tri\text{-}Mask),$

where $tri\text{-}I$ is the concatenated result of three all-one matrices,
which has the same dimension as $tri\text{-}Mask$.

Training Strategy We use a two-stage training strategy. In the first stage, we
only train our encoder model based on the loss in Sec. 4.1. After the loss is
converged, we freeze the encoder parameters, and train the Adaptive Feature
Alignment (AFA) module with occlusion-aware mix tri-plane. When training the
geometry-aware encoder, we follow e4e [44] to train different level modules
and output corresponding $w_{i}$ progressively. Only the latent code before
and at the current stage will be added to $w^{+}$. Different from e4e, we only
use 3 progressives stage (_i.e_. coarse, mid, fine), instead of 14 $w_{i}$
stages (_i.e_. $0,1,2,...,13$) in e4e. We train the canonical latent
discriminator $\mathcal{D}_{\mathcal{W}_{c}}$ at the beginning of each encoder
training iteration, and fix its parameter when training the encoder.

Losses. In the first training stage, we use the loss function mentioned in
Sec. 4.1. For the input view reconstruction, the loss $\mathcal{L}_{rec}$ is
computed by:

$\displaystyle\mathcal{L}_{rec}$
$\displaystyle=\lambda_{1}\left\|I_{w^{+}}-I\right\|_{2}+\lambda_{2}\mathcal{L}_{LPIPS}(I_{w^{+}},I)+\lambda_{3}\mathcal{L}_{ID}(I_{w^{+}},I),$
(13)

where $\lambda_{1}=1.0,\lambda_{2}=0.8,\lambda_{3}=0.25$ are constant weights.
The total loss for training our encoder and the canonical latent discriminator
can be written as:

$\mathcal{L}^{E}_{total}=\mathcal{L}^{\mathcal{D}}_{\text{adv}}+\lambda_{4}\mathcal{L}_{\text{adv}}^{E}+\lambda_{5}\mathcal{L}_{BG}+\mathcal{L}_{rec},$
(14)

where $R_{1}$ regularization weight $\gamma=10$,
$\lambda_{4}=0.05,\lambda_{5}=5.0$ are constant weights. Moreover, we apply
regularization to $\Delta w^{+}=(w_{1}-w_{0},w_{2}-w_{0},...,w_{13}-w_{0})$ by
$\mathcal{L}_{2}$ norm, _i.e_. $\mathcal{L}_{w^{+}}=\lambda_{6}\left\|\Delta
w^{+}\right\|_{2}$, where $\lambda_{6}=0.001$.

Once the loss in the first stage is converged, we freeze the encoder
parameters and drop the canonical latent discriminator and losses
$\mathcal{L}^{\mathcal{D}}_{\text{adv}},\mathcal{L}_{\text{adv}}^{E},\mathcal{L}_{w^{+}}$,
but add a $\mathcal{L}_{2}$ regularization loss on $\Delta F=F^{*}-F$ with
$\lambda_{7}=0.0001$ for training the AFA module. We apply the loss settings
to both human data and cat data.

For re-implemented pSp [38] and e4e [44], we train it with the same loss
weight in our settings, and train them until the losses are converged.

Experiment settings. For human data, we train the encoder with a batch size of
3 on one 3090 GPU for about 1,000,000 iterations. We use a discriminator
learning rate of 0.00002 with Adam optimizer and an encoder learning rate of
0.0001 with a ranger optimizer, which is a combination of Rectified Adam [27]
with the Lookahead technique [52]. Then we train our AFA module with a
learning rate of 0.000025 with Adam optimizer, which uses 800,000 iterations
of batch size 2. We only train the encoder 200,000 iterations and 150,000
iterations for cat data. We test our method and other methods with 3090 GPUs,
and test inference time with the same settings. We use images from CelebA-HQ
or AFHQv2 Cats to perform inversion.

IDE-3D. IDE-3D [42] uses semantic segmentation to re-train EG3D. It learns an
encoder to get a latent code and further optimize the generator parameters for
inversion and editing images. We compare our encoder with its encoder both
qualitatively and quantitatively.

Table 2: Ablation studies on Geometry-aware Encoder and Occlusion-aware Mix
Tri-plane. We evaluate the proposed canonical space $\mathcal{W}_{c}$, the
background regularization, and the occlusion-aware mix tri-plane on MEAD
dataset. The best performance on $\mathcal{W}$ space inversion and
$\mathcal{F}$ space inversion are in bold.

Method | Novel view quality
---|---
yaw$=\pm 60^{\circ}$ | yaw $=0^{\circ}$
PSNR $\uparrow$ | SSIM $\uparrow$ | MSE $\downarrow$ | LPIPS $\downarrow$ | $\mathrm{FID}\downarrow$ | $\mathrm{ID}\uparrow$ | Geo. Err. $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | MSE $\downarrow$ | LPIPS $\downarrow$ | $\mathrm{FID}\downarrow$ | $\mathrm{ID}\uparrow$ | Geo. Err. $\downarrow$
w/o $\mathcal{D}_{\mathcal{W}_{c}}$, w/o $\mathcal{L}_{BG}$ | 19.82 | 0.5730 | 0.0565 | 0.3069 | 62.1 | 0.6206 | 0.1422 | 19.83 | 0.5382 | 0.0540 | 0.2966 | 68.9 | 0.6842 | 0.1775
w/o $\mathcal{D}_{\mathcal{W}_{c}}$, w/ $\mathcal{L}_{BG}$ | 19.41 | 0.5576 | 0.0585 | 0.2883 | 59.0 | 0.6176 | 0.1016 | 20.76 | 0.5829 | 0.0417 | 0.2572 | 60.6 | 0.6864 | 0.1082
w/ $\mathcal{D}_{\mathcal{W}_{c}}$, w/o $\mathcal{L}_{BG}$ | 20.20 | 0.5999 | 0.0495 | 0.2599 | 55.5 | 0.6459 | 0.1414 | 20.67 | 0.5757 | 0.0428 | 0.2527 | 55.1 | 0.6844 | 0.1491
w/ $\mathcal{D}_{\mathcal{W}_{c}}$, w/ $\mathcal{L}_{BG}$ | 20.66 | 0.6211 | 0.0473 | 0.2203 | 52.2 | 0.6552 | 0.0943 | 21.12 | 0.5944 | 0.0383 | 0.2372 | 54.1 | 0.6848 | 0.0955
w/o occlusion-aware | 20.69 | 0.6176 | 0.0489 | 0.2327 | 52.5 | 0.6832 | 0.0983 | 20.26 | 0.6374 | 0.03202 | 0.2230 | 50.3 | 0.744 | 0.0954
Ours | 20.87 | 0.6299 | 0.0424 | 0.2192 | 50.9 | 0.700 | 0.0950 | 22.13 | 0.6502 | 0.0329 | 0.2204 | 51.7 | 0.743 | 0.0944

Table 3: Ablation study on AFA module. We replace the proposed adaptive
alignment module in AFA with CNN network and evaluate on CelebA-HQ dataset.

Method | PSNR $\uparrow$ | SSIM $\uparrow$ | MSE $\downarrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | ID $\uparrow$ | Geo. Err. $\downarrow$
---|---|---|---|---|---|---|---
w/ conv | 19.67 | 0.5962 | 0.0455 | 0.1766 | 25.1 | 0.8033 | 0.1094
w/ AFA | 21.84 | 0.7079 | 0.301 | 0.1242 | 18.1 | 0.8797 | 0.0984

PTI. PTI [39] is the most used optimization method in 2D GAN inversion. We
follow its official settings, and only replace the 2D GAN by the EG3D
generator. It learns the pivot latent code for 450 iterations, and finetunes
generator parameters for 350 iterations.

Pose Opt. Pose Opt. [23] jointly optimizes camera pose, latent codes, and
generator parameters for 3D GAN inversion. We follow its official settings and
learn the pivot latent code and camera pose for 400 iterations at the first
stage, and finetune generator parameters for 400 iterations at the second
stage. We also use its pre-trained encoder for pivot latent code and camera
pose initialization.

Editing. We perform InterfaceGAN [40] to get a semantic latent direction for
editing. First, we sample 500000 front faces whose corresponding latent codes
are conditioned by canonical poses, and sort them by attribution classifiers.
We choose the top 10000 and bottom 10000 samples according to their score of
classification, then we use SVM to get the direction. Finally, We use the
method in Sec. 4.4 to get 3D-consistent editing results.

##  Appendix B More Experimental Results

Ablation of Geometry-aware Encoder Designs. We test different ablation
settings of our encoder on MEAD, for novel view geometry and texture
evaluation. As shown in Table 2, the design of the canonical discriminator
($\mathcal{D}_{\mathcal{W}_{c}}$) and background depth regularization
($\mathcal{L}_{BG}$) is necessary for a good geometry inversion.

Ablation of Adaptive Feature Alignment. We evaluate Adaptive Feature Alignment
(AFA) module ablation on source view reconstruction performance on CelebA-HQ,
as shown in Table 3. Modified feature maps generated by only convolution
modules are hard to align to the facial region in canonical feature space
$\mathcal{F}_{c}$, whose reconstruction quality is inferior to our method.

Ablation of Occlusion-aware Mix Tri-plane. We evaluate occlusion-aware mix
tri-plane design on MEAD for novel view evaluation. As the distortion exists
in the occlusion part, our full models will perform better on novel view image
synthesis, as shown in Table 2.

Comparison with Encode-based Methods. We show more qualitative comparisons
with encode-based methods in Fig. 12. Our method significantly surpasses
others.

Comparison with Optimization-based Methods. We present more qualitative
results of comparison with optimization-based methods in Fig. 13. It is worth
noting that when the input image is a side face, optimization-based methods
tend to overfit the input image and are hard to synthesize novel-view images.
In some cases, Pose Opt. fails to converge to an accurate pose using
optimization and generates degradation results.

More Novel View Results with Our Method.  More novel view results with our
method can be found in Fig. 14, 15, 16. Our method can achieve high-quality
3D-consistent multi-view image synthesis.

More Results of Editing with Our Method. We present more 3D-consistent results
of human faces and cat faces in Fig. 17, 18, 19. Our method shows powerful
editing ability which can be used in real-world applications.

Figure 12: Comparison with encode-based methods. Figure 13: Comparison with
optimization-based methods. Figure 14: Multi-view human faces inversion
results of our method (part 1/2). Figure 15: Multi-view human faces inversion
results of our method (part 2/2). Figure 16: Multi-view cat faces inversion
results of our method. Figure 17: Editing results on human faces (part 1/2).
Figure 18: Editing results on human faces (part 2/2). Figure 19: Editing
results on cat faces.


# Dyna-T: Dyna-Q and Upper Confidence Bounds Applied to Trees

Tarek Faycal1 and Claudio Zito2∗ 1IRLab, School of Computer Science,
University of Birmingham, United Kingdom2Technology Innovation Institute, Abu
Dhabi, UAE∗Correspondent author<EMAIL_ADDRESS>
###### Abstract

In this work we present a preliminary investigation of a novel algorithm
called Dyna-T. In reinforcement learning (RL) a planning agent has its own
representation of the environment as a model. To discover an optimal policy to
interact with the environment, the agent collects experience in a trial and
error fashion. Experience can be used for learning a better model or improve
directly the value function and policy. Typically separated, Dyna-Q is an
hybrid approach which, at each iteration, exploits the real experience to
update the model as well as the value function, while planning its action
using simulated data from its model. However, the planning process is
computationally expensive and strongly depends on the dimensionality of the
state-action space. We propose to build a Upper Confidence Tree (UCT) on the
simulated experience and search for the best action to be selected during the
on-line learning process. We prove the effectiveness of our proposed method on
a set of preliminary tests on three testbed environments from Open AI. In
contrast to Dyna-Q, Dyna-T outperforms state-of-the-art RL agents in the
stochastic environments by choosing a more robust action selection strategy.

##  I INTRODUCTION

At the core of reinforcement learning (RL) there is the problem of how to
update an approximate value function such that a policy is found. Some methods
require a model of the environment to refine a policy by interacting with
simulated experience from the modelled environment. These methods are
typically referred to as planning methods. Other methods learns directly from
experience without the need of a model, and they are referred to as direct RL
methods. Within planning agents, it is still possible to use the experience
collected by interacting with the environments to improve the model, called
model-learning or indirect RL. Both direct and indirect approaches present
advantages and disadvantages. Indirect methods tend to exploit better a
limited amount of experience, thanks to their models, and thus achieve better
policies with few interactions. On the other hand, direct methods are
typically simpler and do not suffer any bias from the design of the model.

Dyna-Q is a form of hybrid method that uses the real experience generated by
acting on the environment for learning a model as well as a value function.
The value function is then refined by simulated experience generated by the
model. Acting, direct RL, model learning, and planning happen in the Dyna
agent at each interaction. The planning process is usually the computationally
intensive phase of the algorithm which strongly depends on the dimensionality
of the state-action space, as well as the complexity of the problem.

In this paper, we propose to improve the planning efficiency of Dyna-Q by
constructing a upper bound confidence tree (UCT) on sampled simulated
experience. We call this new algorithm Dyna-T.

UCT is a state-of-the-art approach to solving large state-space planning
problems. The algorithm estimates the utility of a state by building and
searching on a tree of simulated experience via a procedure called rollout.
The key aspect is that not all the branches of the tree are equally sampled so
to dedicate the most of the resources versus most promising paths. UCT
achieves that by treating each node of the tree as a multi-armed bandit
problem. Since random sampling is used to discover nodes with a good return,
UCT could take a long time to converge to good performance. Many enhancements
have been proposed to address this problem by either (i) tweaking the action
selection formula (e.g. which arm is pulled), or (ii) designing better-
informed rollout policies. Our proposed approach fall in the second category.
However, different than many of the previous attempts, we do not based our
rollout exploration on heuristic-based policies, but on a model-based RL
policy provided by the Dyna-Q algorithm.

##  II BACKGROUND

In this section, we briefly introduce the basic concepts about the nature of
RL and the specifics of designing agents for this machine learning paradigm.
We will focus on state-of-the-art approaches relative to this work and
highlight some of their strengths and drawbacks. A comprehensive review of
this topic is out of scope and we redirect the interested reader to a few
available sources, such as [6, 9, 1].

###  II-A Reinforcement Learning

A traditional reinforcement learning setting consists of two components: An
agent and an environment. The agent receives an observation from the
environment that contains the current state the agent is in, as well as a
reward signal that describes how favorable that current state is compared to
others. The agent would then map that state to an action using its policy
function returns an action to the environment that is intended to alter the
current state [6].

One iteration of this loop is a single step in the environment. The agent will
continue taking steps in the environment until a predetermined number is
reached or some environment-specific termination condition is met. This
interaction is termed episodic, because it has a clearly defined start and
end.

This way of learning poses a unique difficulty when compared to other branches
of machine learning. The simulation of experience might be computationally
expensive depending on environment complexity, added to the existing cost
associated with learning. Furthermore, the environments can vary greatly
making a general purpose reinforcement learning algorithm or agent difficult
to construct.

###  II-B Markov Decision Processes

One mathematical framework that can be used to model reinforcement learning
problems is a Markov Decision Process (MDP) [1]. An MDP consists of a set of
states $S$ that describe the environment, a set of actions $A$ available to
the agent, a reward function $R(s,s^{\prime})$, which returns the reward gain
from moving from state $s$ to state $s^{\prime}$, and a transition model
$P(s^{\prime}|s,a)$ that returns the probability of moving to state
$s^{\prime}$ given the current state $s$ and action $a$. The solution to this
problem comes in the form of a policy function $\pi(s)$ that returns an action
given a state. The optimal policy is the one that maximizes the rewards
collected by the agent [9]. The important assumption in this framework is the
process follows the Markov property, meaning the current state of the
environment must contain all information required by the agent to make a
decision regarding what action to take, which implies that the agent does not
need access to the history of this process. However, many real-world
applications in robotics are characterised by continuous state, action and
observational spaces and/or physical interactions with the environment
challenge the MDP formulation, although several attempts have been made in the
past[11, 14, 12, 5, 13].

###  II-C Active RL and the role of policy in exploration

There are two approaches to learning within the MDP framework: active and
passive. Passive RL evaluates a fixed policy over a certain number of trials
in the environment to derive the utility value associated with each state.
Utility is a concept borrowed from Economics and represents how good a state
is in the context of the problem and policy being evaluated. It is computed as
the expected sum of discounted rewards
$U^{\pi}(s)=\mathbb{E}[\sum\limits_{t=0}^{\infty}\gamma R(s_{t})]$. The $\pi$
symbol is the policy being evaluated, $\gamma\,\epsilon\,[0,1]$ is a discount
factor used to control the influence of future states on the value of the
current one, and $R(s_{t})$ is the reward gain in state $s$ at time $t$.

Active RL adopts a different approach, the agent must interact with the
environment and dynamically adjust its policy in order to arrive at an optimal
one. This introduces a new problem, known as the exploration/exploitation
dilemma. How much time should an agent spend exploring its environment, and
when should it stop and exploit the knowledge it has gathered so far to
maximize the reward? There is no single answer to this question because of the
wide range of possible environments. For example, in an environment with very
sparse rewards, the agent may need to spend a significant amount of time
exploring before finding any reward signal. Additionally, If it does come
across a solution early, how will the agent know to stop then? and how can it
guarantee that there aren’t any higher rewards to be gained from further
exploration? [9]. This work will focus on active RL given that it is a more
accurate representation of how biological agents learn, in addition to the
fact that some environments cannot be simulated due to high complexity or
dynamics that are not well understood.

One of the most commonly used strategy to guarantee exploration is $\epsilon-
greedy$. The agent always takes the best action according to the information
gathered so far, but with a small probability $\epsilon$ of choosing a random
action from those currently available. Presented by Sutton and Barto in 1998
[9], it has remained one of the most commonly used strategies in state of the
art research, albeit with some modifications such starting with $\epsilon=1$
and gradually reducing it to converge on the learned policy.

###  II-D Model-Free and Model-Based Learning

One approach to solving MDPs is framing it as an optimization problem.
Adaptive Dynamic Programming (ADP) was one of the first passive RL methods
utilized. The underlying principle is that the utility values of each state
obey the Bellman Equation [3] for a fixed policy $\pi$:

$U^{\pi}(s)=R(s)+\gamma\sum_{s^{\prime}}{P(s^{\prime}|s,\pi(s))}{U^{\pi}(s^{\prime})}$
(1)

Given a fixed policy, and the above condition, the agent can learn the
transition model $P(s^{\prime}|s,\pi(s))$ by going through a certain number of
episodes and collecting statistics about state transitions and observed
rewards. This information can be plugged in to Eq. 1 and solved for each
state, either directly or through iterative methods, giving us a utility
estimate for each state:

$U_{i+1}(s)\leftarrow
R(s)+\sum_{s^{\prime}}P(s^{\prime}|s,\pi_{i}(s))U_{i}(s^{\prime})$ (2)

Equation 2 shows a simplified iterative update to the utility where $i$ is the
current iteration of the algorithm. The above equation can be used as part of
a method called policy iteration [6] that starts with a random policy and
updates it after every run of utility evaluation to take the actions that
maximize the final reward gathered.

Such a method is called model-based because it explicitly learns the
transition model of the MDP. However, since the policy updates are based
solely on an estimated model, it may not capture the true model and may learn
a sub-optimal policy. Moreover, this approach will translate poorly into an
active RL setting because of how long it might take for the collected
statistics to reflect the true nature of the environment.

Model-free methods rely heavily on the Law of Large Numbers to update the
utility values without ever learning the model. This approach provides both
computational and memory advantages because the state transitions and model
does not need to be kept in memory and updated after each step in the
environment. Temporal Difference (TD) learning [7] is one of the most
prominent model-free approaches that uses the estimated value of the next
encountered state to update the value of the current one. Whenever a
transition from state $s$ to state $s^{\prime}$ occurs, the utility value of
state $s$ is updates as follows:

$U^{\pi}(s)\leftarrow U^{\pi}(s)+\alpha(R(s)+\gamma
U^{\pi}(s^{\prime})-U^{\pi}(s))$ (3)

The added parameter $\alpha\,\epsilon\,[0,1]$ is the learning rate, that value
of which ranges from The average value of $U^{\pi}(s)$ will converge after a
large number of episodes, taking longer than ADP but at a lower cost and with
added flexibility of application to domains where the model cannot be learned
[9]. TD learning can be seen as a first approximation of ADP because it does
not take into account all future states during each update (Recall the
summation over all possible $s^{\prime}$ in Eq. 2). This means that, given a
good exploration strategy and enough iteration, TD methods surpass ADP in
active RL. In both these approaches, updating value estimates is done using
other estimates; a process known as bootstrapping. A major advance in model-
free RL came in the form of Q-learning [10], where an agent would learn the
function $Q(s,a)$ that represents the quality of taking action $a$ in state
$s$. This is much more practical than methods based on utility because the
policy can be derived directly from the Q-function.

The RL algorithms implemented in this work rely heavily on these ideas, and
will be further elaborated on in Sec. V.

###  II-E Drawbacks of bootstrapping rewards

While methods that use bootstrapping can provide guarantees for a solution
(even an optimal one), it is conditioned on infinite exploration and assumes
that a reward signal will be encountered in the environment [9]. In
environments where rewards are very sparse, delayed in time compared to when
an action was taken, or when multiple actions need to be executed in a
specific order, bootstrapping cannot be guaranteed to succeed in a timely
manner. Of course, infinite exploration and learning using $\epsilon-greedy$
means that successful strategies will be learned at some point given enough
data [10], which is not practical in cases where simulations are
computationally expensive or time consuming.

###  II-F Q-Learning

Q-learning is a fairly flexible approach that can be applied to a host of
different problems. Q-learning is a TD method, with an update equation
identical to Eq. 3:

$Q(s,a)\leftarrow{}Q(s,a)+\alpha(R(s)+\gamma
max_{a^{\prime}}Q(s^{\prime},a^{\prime})-Q(s,a))$ (4)

This algorithm operates in an off-policy manner, meaning that the target value
used in the TD update in Eq. 4 is the best possible future value that can be
achieved in the next state, $max_{a^{\prime}}Q(s^{\prime},a^{\prime})$, and
not the value of the action actually taken. Standard Q-learning only looks
ahead one-step ahead to update the value of the current state, making no use
of longer sequences of steps to update multiple values. However, Q-learning is
guaranteed to converge given infinite exploration and learning making it a
good baseline for comparison.

###  II-G SARSA($\lambda$)

SARSA is another model-free algorithm that closely mirrors Q-learning, with a
slight modification. SARSA is on-policy, so the value update is not carried
out for the current state, until an actual action is taken in the next state.
This approach can lead to faster convergence but is highly dependent on the
policy used. Moreover, it is brittle when compared to Q-learning, which is
more flexible when applied to different types of environments and robust
across different hyperparameters and policies. The $\lambda$ portion of the
title refers to the eligibility traces mechanism [9] that allows for multi-
step look-ahead when updating the current state’s value. Effectively, this
maintains a memory of the sequence of state-action pairs that contribute to
the currently encountered rewards and updates the entire sequence in the
q-table.

Figure 1: Illustration of SARSA($\lambda$): Uses an additional table to store
the eligibility traces of state-action pairs in order execute multi-step
updates, as opposed to the one-step updates in our Q-Learning implementation

This is done by maintaining a separate table, refer to Fig. 1, where the
eligibility of states for an update is maintained. Whenever a state-action
pair is encountered, one is added to the corresponding table entry, and after
each update, the entire table is multiplied by $\lambda$ (a hyperparameter,
$\lambda\,\epsilon\,[0,1]$) which represents the decay of eligibility for each
state-action pair for the next update. Updates occur as follows:

$\Delta=R(s)+\gamma Q(s^{\prime},a^{\prime})-Q(s,a)$
$Q\leftarrow{}Q+\alpha*\Delta*Z$ (5)

Note the lack of a $max_{a^{\prime}}$ operator, which makes it on-policy. Also
note the fact that the entire Q-table is updated at each step which does lead
to increased computation when compared to Q-learning, especially with a high
$\lambda$. Setting $\lambda=0$ makes the algorithm one-step, as described in
3.1.1.

###  II-H Dyna-Q

The algorithms presented thus far do not use a model of the environment to
learn. Using one however, would facilitate planning via simulated experience.
One of the earliest methods to do so came with the Dyna architecture [8].

The model is built by adding newly encountered transitions to memory in a
table of states and actions. The transition is stored as a tuple
$(r,s^{\prime})$ where $r$ is the reward received from taking the chosen
action in the current state, and $s^{\prime}$ is the next state the agent ends
up in. The learning process is the same as Q-learning, with added planning
steps. After each transition and Q-table update, previously encountered state-
action pairs are sampled from the model and the stored tuple is used to
execute the Q-learning update again on the sampled state-action pair. The
added planning steps increase the speed of convergence.

Figure 2: Dyna-Q Architecture: Incorporates a model of the environment to
facilitate planning, which is done by sampling the stored transitions at every
step and executing the same update as Q-Learning

The original implementation presented in [8] assumes that the environment is
deterministic, which means that taking action $a$ in state $s$ will always
result in reward and next state $(r,s^{\prime})$. This is not a completely
accurate representation of most environments encountered in everyday life,
where a multitude of unconsidered factors can affect the consequences of agent
actions, whereas simulated environments tend to be rigorously controlled.
Moreover, this type of model will fail if the transitions are not the same
every time. So in order to try making a more robust model that can handle
stochastic environments, we implement both the original Dyna-Q in addition to
a modified version that builds a model of the environment in a similar fashion
to ADP, called Stochastic Dyna-Q. This is done by counting the number of times
state-action pairs are encountered, in addition to the number of times the
specific state-action pair leads to state $s^{\prime}$. Using this gathered
information, we can compute the expected reward of taking a certain action in
a state and use that for the planning portion of the algorithm. Planning
differs significantly in this case because the Q-table entry being updated
must reflect the new probabilities across all possible next states meaning
that the target value of the TD update will be:

$\sum\limits_{s^{\prime}}P(s^{\prime}|s,a)\;max_{a^{\prime}}Q(s^{\prime},a^{\prime}).$
(6)

##  III Dyna-T

While a stochastic model is more robust, it is affected by the same problems
as ADP, namely, the estimated model might take too long to accurately reflect
the environment, which significantly slows down learning or even deteriorate
it. To remedy this problem, we utilize the concept of Upper Confidence Bounds
applied to Trees [4]. This technique is used to improve Monte Carlo Tree
Search (MCTS) methods by computing a statistical interval whose upper bound
would represent the potential reward to be gained, while also emphasizing the
exploration of new state-action pairs. The algorithm computes the upper bounds
associated with a state-action pair with minimal increase in memory cost, but
with a slightly increased computational one. The upper bounds are computed in
a Dyna agent as follows:

$Q_{uct}(s,a)=Q(s,a)+c\sqrt{\frac{\ln n_{s}}{n_{s,a}}}$ (7)

Where $c$ is an experimentally determined exploration parameter, $n_{s}$ is
the number of times state $s$ has been encountered, and $n_{s,a}$ is the
number of times action $a$ was taken in state $s$. An important fact to note
is that this is not a modification to learning, but an improvement to the
policy. Dyna-T does not use an $\epsilon-greedy$ policy like the rest of the
algorithms presented so far, but chooses the action with the highest upper
bound, which is initialized to infinity for unexplored state-action pairs to
encourage exploratory behavior. This leads to an increase in the chances of
encountering a reward in the environment via more directed exploration when
compared to occasionally choosing random actions.

##  IV Test Environments

In this section we will introduce the environments used for our evaluation. We
focus on deterministic or stochastic single-agent environments with fully
observable states, discrete states and actions. All the environments are
available in the OpenAI Gym toolkit [2].

(a)

(b)

Figure 3: CliffWalking Results: Stochastic Dyna-Q and Dyna-T fail to converge
as quickly as the other algorithms.

(a)

(b)

Figure 4: NChain Results: Dyna-T converges on an optimal policy much quicker
in the stochastic environment, besting the rest of the algorithms tested.

(a)

(b)

Figure 5: FrozenLake Results: Dyna-T shows significantly better results on the
most complex of the three environments presented, indicating a more robust
action selection strategy.

###  IV-A Cliff Walking

Figure 6: The CliffWalking Environment: An agent must get from the start to
the goal in the shortest amount of steps since each additional step provides
-1 as a reward. The agent loses if it steps into the cliff.

The first environment considered is a deterministic one with a dense reward
distribution, presented in [9]. The agent must move from position S to G, and
receives a -1 reward at every white square. If the agent falls in the cliff,
it receives a reward of -100 and the episode is terminated. The goal square
gives a reward of 0 and also terminates the episode. The agent’s available
actions are Up, Down, Left and Right that move it one square in the desired
direction. The best performers are the ones that find the shortest path to the
goal, and minimize the negative reward.

###  IV-B N-Chain

Figure 7: The NChain Environment: The agent can only move forward and back,
and the optimal agent ignores the small reward of moving back and chooses to
go forward, despite the probability of ”slipping” and taking the opposite
action

The second environment is called N-Chain. There are only five states and two
actions, forward and back, and a small probability that the action chosen is
not the one actually taken in the environment. The agent starts in state
$S_{1}$. Moving forward yields a reward of 0 except in $S_{5}$ where the
reward for moving forward is 10. Moving back sends the agent to state $S_{1}$
and gives a reward of 1. Forward is represented by solid lines and backwards
with a dashed line, as shown in Fig. 7. Optimal behavior in this environment
means to forgo the immediate rewards gained by moving back and continuously go
forward to seek the larger reward available in the last state.

###  IV-C Frozen Lake

Figure 8: The FrozenLake Environment: The agent must navigate from start to
finish, the grey squares are holes that terminate the episode with a reward of
0, while the blue squares are frozen and can cause an agent to slip in a
random direction. The only positive reward available is 1, if the agent
reaches the goal.

The last environment combines elements from both of theones presented above
and it is shown in Fig. LABEL:fig:frozenlake. The agent must navigate from S
to G, and avoid holes, the grey squares, which terminate the episode with a
reward of $0$. The blue squares are frozen, meaning the agent can slip in a
random direction whenever it takes an action while it’s on a blue square. The
available actions are Up, Down, Left, and Right. The rewards are sparse when
compared to the previous environments because the agent only receives a reward
of $1$ if it manages to reach the goal state. Due to the nature of the
environment, for an algorithm to solve this task it must achieve an average
score of 0.74, at least, over $100$ episodes.

##  V Experiments

In this work, we assess the performance of Dyna-T against Q-Learning,
SARSA($\lambda$), Dyna-Q, and Stochastic Dyna-Q using as a benchmarking
platform the OpenAI Gym toolkit [2]. Our chosen experimental scenarios are
Cliff Walking, Frozen Lake, and NChain.

Q-Learning, SARSA($\lambda$), Dyna-Q, Stochastic Dyna-Q, and Dyna-UCT were
implemented and run on each of the three chosen environments. CliffWalking and
FrozenLake were tested for 2000 episodes, while NChain was tested for 500.
Their behavior with different hyperparameters was observed using parameter
evaluation scripts, that test each algorithm with different parameter values,
such as learning rate or planning steps. The primary measure of performance
was the reward collected in the environment, we also measure the duration
across iterations of each algorithm for further insight into their
practicality.

##  VI Results & Discussion

The algorithms all behave as expected in the deterministic CliffWalking
setting Fig. 3. SARSA($\lambda$) and Dyna-Q converge faster than one-step
Q-learning, while Stochastic Dyna-Q and Dyna-T fail to learn the appropriate
environment model quickly enough. Moreover, due to the added computational
cost of storing transitions and updating expected rewards at every step
Stochastic Dyna-Q and Dyna-T take a significantly longer time. This is also
partially due to the use of the Pandas package to handle the tabular
representation, storage, and retrieval operations of the algorithms. Pandas
tends to be slower than a numerical computation package like NumPy, but does
provide some advantages such as multi-indexing which was very helpful in the
implementation.

The balance shifts once the algorithms are tested in a stochastic environment.
In NChain, as shown in Fig. 4, one-step Q-learning and Dyna-Q fail to converge
on an optimal solution within the allotted timeframe while SARSA($\lambda$),
Stochastic Dyna-Q and Dyna-T perform much better. Dyna-T converges on the best
solution which is in line with our expectations.The time difference is
apparent, but less extreme due to the smaller state space and episode length.
Note the particularly poor performance of Dyna-Q, because it assumes a
deterministic environment, meaning that its model will store transitions that
are not guaranteed to occur every time which can inhibit learning.

FrozenLake is where we really start seeing the benefits of UCT, shown in Fig.
5. Dyna-T latches onto solutions much quicker than the rest of the algorithms.
While all of them are guaranteed to converge, Dyna-T will find the appropriate
solutions faster despite taking no random actions to explore the environment,
it is completely driven by the upper bounds computation.

##  VII Conclusion and Future Work

In this paper, we presented a novel algorithm for on-line model-based
reinforcement learning, called Dyna-T. We combined the benefits of two popular
technique: i) Dyna-Q and ii) UCT. The former is a hybrid RL approach that
combines planning, acting, model-learning, and direct-RL in a unique
framework, all occurring continually at each iteration of the algorithm.
However, the planning process can be computationally expensive and a drawback
for this approach. On the other hand, UCT is a state-of-the-art sampling-based
method for exploring large trees in an efficient way. We demonstrated that our
Dyna-T is efficient in exploring the simulated experience generated by the
Dyna-Q model leading to a more efficient action selection during the on-line
learning. By interacting with the environment with better actions, our model
converges quicker to the optimal solution than several state-of-the-art
methods in two testbed environments from Open AI. We also compared the
performance of our propose model in a deterministic scenario. Although Dyna-Q
is deterministic by nature, the combination with the UCT degraded its
performance in such scenario. Future work will focus on testing in larger
environments and to compare it against deep-learning techniques.

## References

  * [1] Richard Bellman.  A markovian decision process.  Journal of Mathematics and Mechanics, pp. 679–684, 1957.
  * [2] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba.  Openai gym, 2016.
  * [3] Avinash K Dixit.  Optimization in economic theory.  Oxford University Press on Demand, 1990.
  * [4] Levente Kocsis and Csaba Szepesvári.  Bandit based monte-carlo planning.  In European conference on machine learning, pp. 282–293. Springer, 2006.
  * [5] C. J. Rosales, F. Spinelli, M. Gabiccini, C. Zito, and J. L. Wyatt.  Gpatlasrrt: a local tactile exploration planner for recovering the shape of novel objects.  International Journal of Humanoid Robotics, Special Issue ’Tactile perception for manipulation: new progress and challenges’, 15, 2018\.
  * [6] Stuart J Russell and Peter Norvig.  Artificial intelligence: a modern approach.  Malaysia; Pearson Education Limited,, 2016.
  * [7] Richard S Sutton.  Learning to predict by the methods of temporal differences.  Machine learning, 3(1):9–44, 1988.
  * [8] Richard S Sutton.  Dyna, an integrated architecture for learning, planning, and reacting.  ACM SIGART Bulletin, 2(4):160–163, 1991.
  * [9] Richard S Sutton, Andrew G Barto, et al.  Reinforcement learning: An introduction.  MIT press, 1998.
  * [10] Christopher John Cornish Hellaby Watkins.  Learning from delayed rewards.  PhD thesis, King’s College, Cambridge, 1989.
  * [11] C. Zito.  Planning simultaneous perception and manipulation.  PhD thesis, University of Birmingham, 2016.
  * [12] C. Zito, M. Kopicki, R. Stolkin, C. Borst, F. Schmidt, M. A. Roa, and J.L. Wyatt.  Sequential re-planning for dextrous grasping under object-pose uncertainty.  In Workshop on Manipulation with Uncertain Models, Robotics: Science and Systems (RSS), 2013.
  * [13] C. Zito, V. Ortenzi, M. Adjigble, M. S. Kopicki, R. Stolkin, and J. L. Wyatt.  Hypothesis-based belief planning for dexterous grasping.  CoRR arXiv preprint, arXiv:1903.05517 [cs.RO] (cs.AI), 2019.
  * [14] C. Zito, R. Stolkin, M. Kopicki, M. Di Luca, and J.L. Wyatt.  Exploratory reach-to-grasp trajectories for uncertain object poses.  In Workshop of Beyond Robot Grasping, IEEE/RSJ Intelligent Robots and Systems (IROS), 2012.