SNeRL: Semantic-aware Neural Radiance Fields for Reinforcement Learning

To learn 3D-aware representation from a single view image, the previous methods exploit standard convolutional autoencoder architecture conditioned by the camera poses, which generates scenes from arbitrary views with either deterministic (Tatarchenko et al., 2016; Worrall et al., 2017) or stochastic (Eslami et al., 2018) latent vectors. Recently, neural radiance fields (NeRF) have achieved an exceptional progress in understanding 3D scenes and synthesizing novel views. Following, some approaches propose latent-conditioned NeRF (Martin-Brualla et al., 2021; Yu et al., 2021; Wang et al., 2021), but the major objective of the aforementioned methods is improving the quality of synthesized images rather than extracting time-variant latent vectors with 3D dynamic scene understanding from multi-view inputs. In this paper, we leverage the autoencoder with convolutional encoder and NeRF-style decoder (Li et al., 2022; Driess et al., 2022) so that the encoder can extract 3D-aware representation from multi-view inputs for RL downstream tasks.

The RL frameworks with image inputs typically have an encoder, which maps high-dimensional observations to a low-dimensional latent vector. RL agent is trained over the latent state space to maximize its objective functions, e.g., the total discounted reward for each episode. While a number of works have made significant advancements, it still remains a challenging open problem.

To address the sample inefficiency of image-based RL, prior works adopt various data-augmentation techniques (Laskin et al., 2020a; Yarats et al., 2021), contrastive learning with data augmentation (Laskin et al., 2020b; Schwarzer et al., 2020; Stooke et al., 2021; Liu & Abbeel, 2021; Zhan et al., 2022), representation learning from image reconstruction (Islam et al., 2022; Kulkarni et al., 2019), or task information reconstruction (Yang & Nachum, 2021; Yamada et al., 2022). Other approaches propose to capture the relations between multi-view data (Dwibedi et al., 2018; Kinose et al., 2022; Sermanet et al., 2018) or keypoints (Manuelli et al., 2020). There are also some approaches leveraging transition sequence data (Hansen et al., 2020; You et al., 2022), or pre-training with offline image-based RL (Wang et al., 2022). Unfortunately, these works have limited capability in learning 3D-structural information and could not obtain an intuitive understanding of the 3D environments that humans have because of the 2D bias that 2D convolutional neural networks have.

In recent, there have been attempts to learn the 3D structure of the real world (Li et al., 2022; Driess et al., 2022). Li et al. (2022) firstly proposes autoencoder with convolutional encoder and NeRF (Mildenhall et al., 2020) decoder to control the visuomotor with learned dynamics model and model-predictive control (MPC). Following, NeRF-RL (Driess et al., 2022) extends the prior study and firstly introduces NeRF-based architecture to the general model-free RL framework. However, they could not learn semantic features due to the limited RGB supervision with naïve NeRF. To learn object-centric representation only with RGB supervision, NeRF-RL presents compositional NeRF with object-individual masks, but requiring masks during the deployment of RL agents seems to be a strong assumption.

In this paper, we propose SNeRL which learns both geometric and semantic information with RGB, semantic, and distilled feature supervision for RL downstream tasks without any object masks during the inference phase.

The concept of neural radiance fields (NeRF) (Mildenhall et al., 2020) is to represent the 3D scene with learnable and continuous volumetric fields $f_{\theta}$. Specifically, at any 3D world coordinate $\mathbf{x}\in\mathbb{R}^{3}$ and unit viewing direction $\mathbf{d}\in\mathbb{R}^{3}$, $f_{\theta}$ estimates the differntiable volume density $\sigma$ and RGB color $\mathbf{c}$: $f_{\theta}(\mathbf{x},\mathbf{d})=(\sigma,\mathbf{c})$. Let the camera ray of the pixel in the camera coordinate be $\mathbf{r}=\mathbf{o}+t\mathbf{d}$, where $\mathbf{o}$ indicates the camera origin. The corresponding pixel value from an arbitrary view can be rendered through volumetric radiance fields as:

where $T(t)=\mathrm{exp}(-\int^{t}_{t_{n}}\sigma(s)ds)$ and $t_{n}$ and $t_{f}$ indicate pre-defined lower and upper bound of the depth respectively.

Then, $f_{\theta}$, which is usually formulated with MLP, is optimized through pixel-wise RGB supervision from multiple views as:

where $\mathbf{r}_{i,j}$ indicates ray $j$ from images of $i^{th}$ view. $\hat{C}$ and $C$ represents the rendered volumetric fields into 2D image and ground truth pixel value respectively.

We consider a finite-horizon Markov Decision Process (MDP) $\mathcal{M}=(\mathcal{O},\mathcal{A},\mathcal{T},\mathcal{R},\gamma)$, where $\mathcal{O}$ denotes the high-dimensional observation space (image pixels), $\mathcal{A}$ the action space, $\mathcal{T}(o^{\prime}|o,a)$ the transition dynamics $(o,o^{\prime}\in\mathcal{O},a\in\mathcal{A})$, $\mathcal{R}:\mathcal{O}\times\mathcal{A}\rightarrow\mathbb{R}$ the reward function, and $\gamma\in[0,1)$ the discount factor. Following the general idea of learning RL downstream tasks with pre-trained scene representations, we consider an encoder $\Omega:\mathcal{O}\rightarrow\mathcal{Z}$ that maps and high-dimensional observation $o\in\mathcal{O}$ to a low-dimensional latent state $z\in\mathcal{Z}$ on which an RL agent operates. To learn how to succeed in downstream tasks, the RL policy $\pi_{\theta}(a\in\mathcal{A}|z=\Omega(o))$ maximizes the total discounted reward $\sum_{t=0}^{H-1}=\gamma^{t}\mathcal{R}(o_{t},a_{t})$ of trajectories $\tau_{i}=(z_{0},o_{0},...,z_{H},o_{H})_{i}$.

Similar to Li et al. (2022), we adopt the multi-view encoder $\Omega$ which fuses the observations from multiple camera views together to learn a single latent vector $z$ for RL tasks. The encoder takes the pixel-level observations $o^{i}\in\mathbb{R}^{H\times W\times 3}$, and the corresponding camera projection matrices $K^{i}\in\mathbb{R}^{3\times 4}$ captured from $V$ different camera view as inputs, i.e., $i=1\cdots V$. To generate $z\in\mathcal{Z}$ from the inputs, a convolutional network $E_{\mathrm{CNN}}$ first extracts viewpoint-invariant features from each image. The features from different camera views are channel-wise concatenated with their corresponding (flattened) camera projection matrices to reflect the viewpoint information to the following feature vectors. Then, the concatenated vectors are passed through MLP layers, $g_{\mathrm{MLP}}$, to produce mid-level viewpoint-aware encodings. Lastly, the feature encodings from different camera views are averaged to generate a single encoding, and the averaged feature encoding is projected to the latent space $\mathcal{Z}$ with the latent encoder $h_{\mathrm{MLP}}$ as follows:

To inject 3D structural information into the latent vector $z$, we leverage a latent-conditioned NeRF architecture (Yu et al., 2021; Martin-Brualla et al., 2021; Wang et al., 2021) for the decoder. The difference between previous latent-conditioned NeRF and our proposed SNeRL is that the neural rendering function $f_{\theta}$ from SNeRL not only synthesizes novel views with RGB pixel value $\mathbf{c}$ but also with the semantic label $\mathbf{s}$ (Zhi et al., 2021; Fu et al., 2022; Kundu et al., 2022) and high-dimensional distilled features $\mathbf{f}$ from the large-scale teacher network (Kobayashi et al., 2022) as follows:

By estimating three different radiance fields (semantic, feature, and RGB), the latent vector $z$ is jointly optimized to learn the geometric and semantic representations of the 3D environment. Unlike RGB value $\mathbf{c}$ which is dependent on both the position $\mathbf{x}$ and the viewing direction $\mathbf{d}$, we formulate the semantic label and distilled feature to be invariant to the viewing direction $\mathbf{d}$ because the inherent properties of the scene or the object do not change according to the direction of the camera ray.

As SNeRL predicts three different fields, RGB, semantic, and distilled feature, by adding field-wise branches, they share the neural rendering function $f_{\theta}$ until estimating the density $\sigma$. It indicates that three radiance fields have the same accumulated transmittance $T(t)$ at depth $t\in[t_{n},t_{f}]$ along the ray $\mathbf{r}=\mathbf{o}+t\mathbf{d}$ as

For rendering the RGB field, we follow the same training framework as general latent-conditioned NeRF (Yu et al., 2021; Martin-Brualla et al., 2021; Wang et al., 2021), which optimizes the neural rendering function $f_{\theta}$ via pixel-wise RGB supervision. RGB supervision enables the encoder to extract geometric features from the observed environment by learning the RGB and density distribution in the 3-dimensional space. The rendered pixel value $\hat{C}(\mathbf{r})$ can be calculated as

and the loss function for RGB field, $\mathcal{L}_{\mathrm{RGB}}$, can be formulated with simple L2 loss between the rendered $\hat{C}(\mathbf{r})$ and the ground truth pixel colors $C(\mathbf{r})$,

where $\mathbf{r}_{i,j}$ indicates the camera ray $j$ from the observation $i$, $o^{i}$.

Unfortunately, optimizing the encoder only with only an RGB reconstruction is difficult to capture the semantic or object-centric properties of the 3D scene, which are crucial for downstream RL tasks. Therefore, we extend NeRF-based decoder by appending additional branches before injecting the viewing direction $\mathbf{d}$ into the rendering function, $f_{\theta}$ for semantic segmentation. The rendered semantic labels $\hat{S}(\mathbf{r})$ can be calculated as

and the loss function for semantic field, $\mathcal{L}_{seg}$, can be formulated with the standard cross entropy loss,

where $\hat{S}^{l}$ and $S^{l}$ denote the probability of the ray $j$ in observation $i$ belonging to the class $l$ and its corresponding ground-truth semantic labels, respectively.

To capture further fine-grained features that could not be fully expressed in semantic fields, SNeRL also synthesizes distilled feature fields (Kobayashi et al. (2022)) that predict the output of a pre-trained feature descriptor in a knowledge-distillation manner (Hinton et al., ). It is well known from prior literature (Caron et al., 2021) that Vision Transformer (ViT) (Dosovitskiy et al., 2020) trained in a self-supervised manner, e.g. DINO (Caron et al., 2021), can work as an excellent feature descriptor which explicitly represents the scene layouts such as object boundaries. Since the output of the ViT feature descriptor contains high-dimensional information with different values in all pixels depending on the geometric relationship or semantic meaning, the pre-trained ViT becomes a good feature descriptor with another advantage from the semantic label.

Therefore, we take advantage of such benefits to the NeRF-based decoder so that the latent vector $z$ learns high-level information by distilling the knowledge from ViT teacher network which cannot be learned via ground-truth semantic supervision. The distilled feature fields can be rendered as

The loss function for distilled feature field, $\mathcal{L}_{feat}$, is formulated by penalizing the difference between the rendered features $\hat{F}(\mathbf{r})$ and the outputs of ViT feature descriptor $F(o,\mathbf{r})$ as

Finally, the total loss function $\mathcal{L}$ for jointly optimizing the multi-view encoder and NeRF-based decoder can be formulated as the linear combination of aforementioned losses as:

where $\lambda_{sem}$ and $\lambda_{feat}$ are set to 0.004 and 0.04, respectively, to balance the losses (Zhi et al., 2021; Kobayashi et al., 2022). After training, the multi-view encoder $\Omega$ is exploited as a 3D structural and semantic feature extractor for any off-the-shelf downstream RL algorithms.

We additionally enforce the multi-view self-predictive loss to the latent vector $z$ to ensure that the encoder learns the viewpoint-invariant representation with observations from the same scene. The randomly sampled observations from two different camera pose, $o^{1}$ and $o^{2}$, are processed by the convolutional feature extractor, $E_{\mathrm{CNN}}$, and the weights of the feature extractor are shared between two inputs. A feature from one view, $z_{1}$, is mapped with a prediction network, $h_{pred}$, to match it to the feature from the other view, $z_{2}$. We formulate the self-predictive loss function $D$ with negative cosine similarity as follows:

where $p_{1}$ and $z_{2}$ indicate two output vectors, $p_{1}\triangleq h_{pred}(E_{\mathrm{CNN}}(o^{1}))$ and $z_{2}\triangleq E_{\mathrm{CNN}}(o^{2})$, respectively. We assume that $z_{2}$ is constant and the encoder $E_{\mathrm{CNN}}$ only receives the gradient from $p_{1}$ following Chen & He (2021).

The symmetrized auxiliary representation loss function can be formulated as follows:

Figure 3 shows the episode returns of SNeRL and baselines in 4 different visual control tasks. Thanks to the learned object-centric representation via semantic and distilled feature supervision, SNeRL consistently outperforms state-of-the-art visual RL methods and the prior 3D-aware RL method (NeRF-RL) in terms of data efficiency and performance.

Specifically, the contrastive baselines (CURL, CURL-multiview) and DrQ-v2 could not achieve high returns in the difficult environments (soccer and hammer) even though some of them succeed in the relatively easy environment (window). The results also show that pre-training CNN via naïve reconstruction loss (CNN-AE) with offline data could not succeed in the environments at all. These imply that extracting not only the 3D-aware geometric but also the object-centric and semantic information from multi-view observations is critical for RL performances.

Interestingly, we observe that pre-training a NeRF-based autoencoder only with RGB supervision (NeRF-RL) is not sufficient to learn the features for RL downstream tasks, and it can not outperform multi-view adaptation of the visual RL method with contrastive loss (CURL-multiview). These are contrary to the results reported in the prior work (Driess et al., 2022), which we analyze as follows: the environments we adopted in this work are more challenging compared to those of Driess et al. (2022), which consist of simple-shaped objects with primary colors. Therefore, it is relatively difficult to obtain semantic information simply using RGB supervision. Thus, leveraging a semantic-aware NeRF decoder is required to extract the features for better performances in RL downstream tasks in the practical use of 3D-aware RL, which is consistent with our analysis.

Even though we only leverage the convolutional encoder for downstream RL tasks, we compare the image rendering performance of NeRF-based decoders from SNeRL and NeRF-RL to explore the relationship between the synthesized image quality and RL performances. As NeRF originally aims to synthesize images of arbitrary camera views from a static scene, NeRF-RL, which also trains the volumetric field with a sole RGB supervision, cannot reconstruct dynamic objects, e.g. the robot arm, in input images without semantic information. On the other hand, SNeRL, which utilizes semantic labels and feature outputs from the ViT teacher network as additional supervision signals, not only achieves better RL performance but also well represents the dynamic scene and produces high-fidelity rendering outputs as shown in Figure 5.

In this section, we evaluate whether the learned representation via SNeRL can also be adopted in off-the-shelf model-based reinforcement learning algorithms, which train a world model to characterize the environment and conduct planning over the learned model. We adopt Dreamer (Hafner et al., 2019) as a downstream model-based RL agent and replace the encoder of the representation model with our pre-trained encoder. Refer to Appendix A for additional implementation details and an architectural overview.

Our results are shown in Figure 4(c). We observe that learning model-based RL with the pre-trained encoder of SNeRL outperforms the pre-trained weights of the prior 3D-aware RL method (NeRF-RL) and naïve CNN autoencoder. This empirical evidence is consistent with the case of model-free RL in section 5.1, indicating that the proposed method allows the encoder to learn representations that are important for general off-the-shelf RL agents.