Reinforcement Learning with Neural Radiance Fields
Anonymous Author(s)
AffiliationAddress email
Abstract
It is a long-standing problem to find effective representations for training reinforcement learning (RL) agents.
This paper demonstrates that learning staterepresentations with supervision from Neural Radiance Fields (NeRFs) can improve the performance of RL compared to other learned representations or evenlow-dimensional, hand-engineered state information.
Specifically, we propose totrain an encoder that maps multiple image observations to a latent space describingthe objects in the scene.
The decoder built from a latent-conditioned NeRF servesas the supervision signal to learn the latent space.
An RL algorithm then operateson the learned latent space as its state representation.
We call this NeRF-RL.
Ourexperiments indicate that NeRF as supervision leads to a latent space better suitedfor the downstream RL tasks involving robotic object manipulations like hangingmugs on hooks, pushing objects, or opening doors
Introduction
The sample efficiency of reinforcement learning (RL) algorithms crucially depends on the representation of the underlying system state they operate on [1, 2, 3, 4, 5, 6, 7].
Sometimes, a low-dimensional(direct) representation of the state, such as the positions of the objects in the environment, is considered to make the resulting RL problem most efficient [2].
However, such low-dimensional, direct state representations can have several disadvantages.
On theone hand, a perception module, e.g., pose estimation, is necessary in the real world to obtain therepresentation from raw observations, which often is difficult to achieve in practice with sufficientrobustness.
On the other hand, if the goal is to learn policies that generalize over different objectshapes [8], using a low-dimensional state representation is often impractical.
Such scenarios, whilechallenging for RL, are common, e.g., in robotic manipulation tasks.
Therefore, there is a large history of approaches that consider RL directly from raw, high-dimensionalobservations like images (e.g., [9, 10]).
Typically, an encoder takes the high-dimensional input andmaps it to a low-dimensional latent representation of the state.
The RL algorithm (e.g., the Q-functionor the policy network) then operates on the latent vector as state input.
This way, no separateperception module is necessary, the framework can extract information from the raw observations thatare relevant for the task, and the RL agent, in principle, may generalize over challenging environments,such as if the object shapes are varied.
While these are advantages in principle, jointly trainingencoders capable of processing high-dimensional inputs from the RL signal alone is challenging.
Toaddress this, one approach is to pretrain the encoder on a different task, e.g., image reconstruction[1, 4, 11], multi-view consistency [6], or a time-constrastive task [3].
Alternatively, an auxiliary losson the latent encoding can be added during the RL procedure [5].
In both cases, the choice of the actual (auto-)encoder architecture and associated (auxiliary) lossfunction has a significant influence on the usefulness of the resulting latent space for the downstream
RL task.
Especially for image data, convolutional neural networks (CNNs) are commonly used for the encoder [12].
However, 2D CNNs have a 2D (equivariance) bias, while for many RL tasks, the3D structure of our world is essential.
Architectures like Vision Transformers [13, 14] may processimages with no such 2D bias, but they may require large scale data, which might be challenging in
RL applications.
Additionally, although multiple uncalibrated 2D image inputs can be used withgeneric image encoders [15], they do not benefit from 3D inductive bias, which may help for examplein resolving ambiguities in 2D images such as occlusions and object permanence.
Recently, Neural Radiance Fields (NeRFs)
[16] have shown great success in learning to representscenes with a neural network that enables to render the scene from novel viewpoints, and have sparkedbroad interest in computer vision [17].
NeRFs exhibit a strong 3D inductive bias, leading to betterscene reconstruction capabilities than methods composed of generic image encoders (e.g., [18]).
In the present work, we investigate whether incorporating these 3D inductive biases of NeRFs intolearning a state representation can benefit RL.
Specifically, we propose to train an encoder thatmaps multiple RGB image views of the scene to a latent representation through an auto-encoderstructure, where a (compositional) NeRF decoder provides the self-supervision signal using an imagereconstruction loss for each view.
In the experiments, we show for multiple environments that supervision from NeRF leads to alatent representation that makes the downstream RL procedure more sample efficient compared tosupervision via a 2D CNN decoder, a contrastive loss on the latent space, or even hand-engineered,perfect low-level state information given as keypoints.
Commonly, RL is trained on environmentswhere the objects have the same shape.
Our environments include hanging mugs on hooks, pushingobjects on a table, and a door opening scenario.
In all of these, the objects’ shapes are not fixed, andwe require the agent to generalize over all shapes from a distribution.
To summarize our main contributions: (i) we propose to train state representations for RL with NeRFsupervision, and (ii) we empirically demonstrate that an encoder trained with a latent-conditioned
NeRF decoder, especially with an object-compositional NeRF decoder, leads to increased RL performance relative to standard 2D CNN auto-encoders, contrastive learning, or expert keypoints.
Related Work
Neural Scene/Object Representations in Computer Vision, and Applications.
To our knowledge,the present work is the first to explore if neural scene representations like NeRFs can benefit RL.
Outside of RL, however, there has been a very active research field in the area of neural scenerepresentations, both in the representations themselves [19, 20, 21, 22] and their applications; see[23, 24, 17] for recent reviews.
Within the family of NeRFs and related methods, major thrusts ofresearch have included: improving modeling formulations [25, 26], modeling larger scenes [26, 27],addressing (re-)lighting [28, 29, 30], and an especially active area of research has been in improvingspeed, both of training and of inference-time rendering [31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41].
In our case, we are not constrained by inference-time computation issues, since we do not need torender images, and only have to run our latent-space encoder (with a runtime of approx.
7 ms on an
RTX3090).
Additionally of particular relevance, various methods have developed latent-conditioned[42, 43, 44] or compositional/object-oriented approaches for NeRFs [45, 46, 47, 48, 49, 50, 51, 52,53], although they, nor other NeRF-style methods to our knowledge, have been applied to RL.
Neuralscene representations have found application across many fields (i.e., augmented reality and medicalimaging [54]) and both NeRFs [55, 56, 57, 58] and other neural scene approaches [59, 60, 61, 62]have started to be used for various problems in robotics, including pose estimation [55], trajectoryplanning [56], visual foresight [11, 53], grasping [59, 57], and rearrangement tasks [60, 61, 58].
Learning State Representations for Reinforcement Learning.
One of the key enabling factorsfor the success of deep RL is its ability to find effective representations of the environment fromhigh-dimensional observation data [10, 63].
Extensive research has gone into investigating differentways to learn better state representations using various auxiliary objective functions.
Contrastivelearning is a common objective and has shown success in unsupervised representation learning incomputer vision applications [64, 65].
Researchers built upon this success and have shown suchlearning objectives can lead to better performance and sample efficiency in deep RL [66, 67], wherethe contrasting signals could come from time alignment [68, 3], camera viewpoints [69], and differentsensory modalities [70], with applications in real-world robotic tasks [6, 71].
Extensive effortshave investigated the role of representation learning in RL [72], provided a detailed analysis of the importance of different visual representation pretraining methods [73], and shown how we canimprove training stability in the face of multiple auxiliary losses [74].
There is also a range ofadditional explorations on pretraining methods with novel objective functions (e.g., bisimulationmetrics [75] and temporal cycle-consistency loss [76]) and less-explored data sources (e.g., in-thewild images [77] and action-free videos [78]).
Please check the survey for more related work in thisdirection [79].
Our method is different in that we explicitly utilize a decoder that includes strong 3Dinductive biases provided by NeRFs, which we empirically show improves RL for tasks that dependon the geometry of the objects.
Background
[Figure]
3.1 Reinforcement Learning
[Figure]
This work considers decision problems that can be described as discrete-time Markov Decision Processes (MDPs) M = ⟨S , A , T, γ, R, P 0 ⟩ .
S and A are the sets of all states and actions, respectively.
The transition probability (density) from s to s ′ using action a is T ( s ′ |s, a ) .
The agent receives a real-valued reward R ( s, a, s ′ ) after each step.
The discount factorγ ∈ [0 , 1) trades off immediate and future rewards.
P 0 : S → R +0 is the distribution of thestart state.
RL algorithms try to find the optimal policy π ∗ :
S ×
A → R + 0 , where π ∗ =argmax π (cid:80) ∞ t =0 γ t E s t +1 ∼ T ( ·| s t ,a t ) ,a t ∼ π ( ·| s t ) ,s 0 ∼ P
[ R ( s t , a t , s t +1 )] .
Importantly, in this work, we 108 consider RL problems where the state s encodes both the position and the shape of the objects in thescene.
We require the RL agent to generalize over all of these shapes at test time.
We can thereforethink of the state as a tuple s =
( s p , s s ) , where s p encodes positional information, and s s encodes theshapes involved.
We focus the experiments on sparse reward settings, meaning R ( s, a, s ′ )
=
R 0 >
0for s ′ ∈ S g and R ( s, a, s ′ )
=
for s ∈ S\S g , where the volume of S g ⊂ S is much smaller than thevolume of S .
The state space S usually is low-dimensional or a minimal description of the degreesof freedom of the system.
In this work, we consider that the RL algorithm has only access to a(high-dimensional) observation y ∈ Y of the scene (e.g., RGB images).
In particular, this meansthat the policy has observations as input a ∼ π ( · | y ) .
Since we assume that the underlying states =
( s p , s s ) is fully observable from y , we can treat y like a state for an MDP.
Reinforcement Learning with Learned Latent Scene Representations.
The general idea of RLwith learned latent scene representations is to learn an encoder Ω that maps an observation y ∈ Y toa k -dimensional latent vector z
= Ω( y ) ∈
Z ⊂ R k of the scene.
The actual RL components, e.g.,the Q-function or policy, then operate on z as its state description.
For a policy π , this means that theaction a ∼ π ( · | z ) =
π ( · | Ω( y )) is conditional on the latent vector z instead of the observation ydirectly.
The dimension k of the latent vector is typically (much) smaller than that of the observationspace Y , but larger than that of the state space S .
3.2 Neural Radiance Fields (NeRFs)
The general idea of NeRF, originally proposed by [16], is to learn a function f
= ( σ, c ) that predictsthe emitted RGB color value c ( x ) ∈ R 3 and volume density σ ( x ) ∈ R ≥ 0 at any 3D world coordinatex ∈ R 3 .
Based on f , an image from an arbitrary view and camera parameters can be rendered bycomputing the color C ( r ) ∈ R 3 of each pixel along its corresponding camera ray r ( α ) =
r (0) + αdthrough the volumetric rendering relation
C ( r ) =
(cid:90)
α f α n
T f ( r, α
[Equation]
σ ( r ( α )
[Equation]
c ( r ( α ))
d α with
T f ( r, α ) =
exp (cid:18) − (cid:90)
α α n σ ( r ( u ))
d u (cid:19) .
Here, r (0) ∈ R 3 is the camera origin, d ∈ R 3 the pixel dependent direction of the ray and α n , α f ∈ Rthe near and far bounds within which objects are expected, respectively.
The camera rays aredetermined from the camera matrix K (intrinsics and extrinsics) describing the desired view.
Learning State Representations for RL with NeRF Supervision
This section describes our proposed framework, in which we use a latent state space for RL that islearned from NeRF supervision.
For learning the latent space, we use an encoder-decoder where thedecoder is a latent-conditioned NeRF, which may either be a global [42, 43, 44] or a compositional
Observations
[Figure]
[Figure]
Encoder
LatentSpace
NeRF Decoder
RL Policy
Train offline, then freeze
Train online
Reconstructions renderings … 1 2 V multi-view
Latent-conditionedCompositional NeRF
Decoder *
NeRF decoder [53].
To our knowledge, no prior work has used such NeRF-derived supevision for
RL.
In Sec.
4.1 we describe this proposition, Sec.
4.2 provides an overview of the encoder-decodertraining, Sec. 4.3 and Sec.
4.4 introduce options for the NeRF decoder and encoder, respectively.
4.1 Using Latent-Conditioned NeRF for RL
We propose the state representation z on which an RL algorithm operates to be a latent vectorproduced by an encoder that maps images from multiple views to a latent z , which is trained with a(compositional) latent-conditioned NeRF decoder.
As will be verified in experiments, we hypothesizethat this framework is beneficial for the downstream RL task, as it produces latent vectors thatrepresent the actual 3D geometry of the objects in the scene, can handle multiple objects well, as wellas fuse multiple views in a consistent way to deal with occlusions by providing shape completion,all of which is relevant to solve tasks where the geometry is important.
There are two steps to ourframework, as shown in Fig. 1.
First, we train the encoder + decoder from a dataset collected byrandom interactions with the environment, i.e., we do not yet need a trained policy.
Second, wetake encoder trained in the first step, which we leave frozen, and use the latent space to train an RLpolicy.
Note that we investigate two variants of the auto-encoder framework, a global one, where thewhole scene is represented by one single latent vector, and a compositional one, where objects arerepresented by their own latent vector.
For the latter, objects are identified by masks in the views.
4.2 Overview: Auto-Encoder with Latent-Conditioned NeRF Decoder
Assume that an observation y = (cid:0) I 1: V , K 1: V , M 1: V (cid:1) of the scene consists of RGB images I i ∈
R 3 × h × w , i = 1 , . .
, V taken from V many camera views, their respective camera projection matrices
K i ∈ R 3 × 4 (including both intrinsics and extrinsics), and per-view image masks M 1: V .
For a global
NeRF decoder, these are global non-background masks M i tot ∈ { 0 , 1 } h × w , and for a compositional
NeRF decoder as in [53], these are sets of binary masks M ij ∈ { 0 , 1 } h × w that identify the objectsj = 1 , . . .
, m in the scene in view i .
The global case is equivalent to m = 1 , M ij =1 =
M i tot .
Theencoder Ω maps these posed image observations from the multiple views into a set of latent vectorsz 1: m , where each z j represents each object in the scene separately in the compositional case, or thesingle z 1 all objects in the scene.
This is achieved by querying Ω on the masks M 1: V j , i.e.,
[Equation]
for object j .
The supervision signal to train the encoder is the image reconstruction loss
[Equation]
[Equation]
on the input view i where the decoder D renders an image I =
D
( z 1: m , K ) for arbitrary
viewsspecified by the camera matrix K from the set of latent vectors z 1: m .
Both the encoder and decoderare trained end-to-end at the same time.
The target images for the decoder are the same in both the global and compositional case: the global-masked image
I
i ◦ M i tot ( ◦ is the element-wise product).
Inthe compositional case this can be computed with M i tot = (cid:87) mj =1 M ij .
By fusing the information frommultiple views of the objects into the latent vector from which the decoder has to be able to renderthe scene from multiple views, this auto-encoder framework can learn latent vectors that representthe 3D configurations (shape and pose) of the objects in the scene.
4.3 Latent-Conditioned NeRF Decoder Details
Global.
The original NeRF formulation [16] learns a fully connected network f that representsone single scene (Sec. 3.2).
In order to create a decoder from NeRFs within an auto-encoder tolearn a latent space, we condition the NeRF f ( · , z ) on the latent vector z ∈ R k [42, 43, 44].
Whileapproaches such as [42, 43, 44] use the latent code to represent factors such as lighting or categorylevel generalization, in our case the latent code is intended to represent the scene variation, i.e., shapeand configuration of objects, such that a downstream RL agent may use this as a state representation.
Compositional.
In the compositional case, the encoder produces a set of latent vectors z 1: m describingeach object j = 1 , . . .
, m individually, this leads to m many NeRFs ( σ j ( x ) , c j ( x )) =
f j ( x ) =
f
( x, z j ) , j = 1 , .
, m with their associated volume density σ j and color value c j .
Note thatwhile one could use different networks f j with their own network weights for each object, wehave a single network f for all objects.
This means that both the object’s pose as well as itsshape and type are represented through the latent code z j .
In order to force those conditioned
NeRFs to learn the 3D configuration of each object separately, we compose then into a global
NeRF model with the composition formulas (proposed e.g., by [80, 81]): σ ( x )
=
(cid:80) mj
=
σ j ( x ) ,c ( x ) = 1 σ ( x ) (cid:80) mj
=1 σ j ( x ) c
j
( x ) .
As this composition happens in 3D space, the latent vectors will belearned such that they correctly represent the actual shape and pose of the objects in the scene withrespect to the other objects, which we hypothesize may be useful for the downstream RL agent.
4.4 Encoder Details
The encoder Ω operates by fusing multiple views together to estimate the latent vector for the RL task.
Since the scientific question of this work is to investigate whether a decoder built from NeRFs to trainthe encoder end-to-end is beneficial for RL, we consider two different encoder architectures.
The firstone is a 2D CNN that averages feature encodings from the different views, where each encoding isadditionally conditioned on the camera matrix of that view.
The second one is based on a learned3D neural vector field that incorporates 3D biases by fusing the different camera views in 3D spacethrough 3D convolutions and camera projection.
This way, we are able to distinguish between theimportance of 3D priors incorporated into the encoder versus the decoder.
Per-image CNN Encoder (“Image encoder”).
For the global version, we utilize the network architecture from [11] as an encoder choice.
In order to work with multiple objects in the compositionalcase, we modify the architecture from [11] by taking the object masks into account as follows.
Foreach object j , the 2D CNN encoder computes
[Equation]
E CNN is a ResNet-18 [82] CNN feature extractor that determines a feature from the masked inputimage
I i ◦ M ij of object j for each view i , which is then concatenated with the (flattened) cameramatrix.
The output of the network g MLP is hence the encoding of each view, including the camerainformation, which is averaged and then processed with h MLP , to produce the final latent vector.
Notethat in the global case, we set m = 1 , M ij =
=
M i tot such that Ω CNN produces a single latent vector.
Neural Field 3D CNN Encoder (“Field encoder”).
Several authors [43] have considered to incorporate 3D biases into learning an encoder by computing pixel-aligned features from queried3D locations of the scene to fuse the information from the different camera views directly in 3Dspace.
We utilize the encoder architecture from [53], where the idea is to learn a neural vector fieldϕ (cid:2)
I 1: V , M 1: V j (cid:3) : R 3 → R E over 3D space, conditioned on the input views and masks.
The featuresof ϕ are computed from projecting the query point into the camera coordinate system from the respective view.
To turn ϕ into a latent vector, it is queried on a workspace set X h ∈ R
d
X ×
h
X × w X (a 3Dgrid) and then processed by a 3D convolutional network, i.e., z j = E 3D CNN (cid:0) ϕ (cid:2)
I 1: V , M 1: V j (cid:3) ( X h ) (cid:1) .
This method differs from [43, 83, 60] by computing a latent vector from the pixel-aligned features.
Baselines / Alternative State Representations
In this section, we briefly describe alternative ways of training an encoder for RL, which we willinvestigate in the experiments as baselines and ablations.
For details, refer to the appendix.
Conv.
Autoencoder.
This baseline uses a standard CNN decoder based on deconvolutions insteadof NeRF to reconstruct the image from the latent representation, similar to [1].
Therefore, with thisbaseline we investigate the influence of the NeRF decoder relative to CNN decoders.
We followthe architecture of [11] for the deconvolution part for the global case.
In the compositional case,we modify the architecture to be able to deal with a set of individual latent vectors instead of asingle, global one.
The image I =
D deconv ( g MLP ( 1 m (cid:80) mj
=1 z j ) , K ) is rendered from z 1: m by firstaveraging the latent vectors and then processing the averaged vector with a fully connected networkg MLP , leading to an aggregated feature.
This aggregated feature is concatenated with the (flattened)camera matrix K describing the desired view and then rendered into the image with D deconv .
In theexperiments, we utilize this decoder as the supervision signal to train the latent space produced bythe 2D CNN encoder from Sec. 4.4.
In the compositional version, the 2D CNN encoder (4) use thesame object masks as the compositional NeRF-RL variant.
Contrastive Learning.
As an alternative to learning an encoder via a reconstruction loss, the idea ofcontrastive learning [84] is to define a loss function directly on the latent space that tries to pull latentvectors describing the same configurations together (called positive samples) while ones representingdifferent system states apart (called negative samples).
A popular approach to achieve this is with the
InfoNCE loss [85, 64].
Let y i and ˜
y
i be two different observations of the same state.
Here, ˜ · denotes aperturbed/augmented version of the observation.
For a mini-batch of observations { ( y i , ˜ y i ) } ni =1 , afterencoding those into their respective latent vectors z i =
Ω
( y i ) , ˜ z
i = Ω(˜ y i ) with the encoder Ω , theloss for that batch would use ( z i , ˜ z i ) as a positive pair, and ( z i , ˜ z
̸ = i ) as a negative pair, or some similarvariation.
A crucial question in contrastive learning is how the observation y is perturbed/augmentedinto ˜ y to generate positive and negative training pairs, described in the following.
CURL.
In CURL [5], the input image is randomly cropped to generate y and ˜ y .
We closely follow thehyperparameters and design of [5].
CURL operates on a single input view and we carefully choose aview for this baseline from which the state of the environment can be inferred as best as possible.
Multi-View CURL.
This baseline investigates if the neural field 3D encoder (Sec. 4.4) can be trainedwith a contrastive loss.
As this encoder operates on multiple input views we double the amount ofavailable camera views.
Half of the views are the same as in the other experiments, the other half arecaptured from sightly perturbed camera angles.
We use the same loss as CURL, but with differentcontrastive pairs – rather than from augmentation, the contrastive style is taken from TCN [68]: thepositive pairs come from different views but at the same moment in time, while negative pairs comefrom different times.
Therefore, this baseline can be seen as a multi-view adaptation of CURL [5].
Direct State / Keypoint Representations.
Finally, we also consider a direct, low-dimensionalrepresentation of the state.
Since we are interested in generalizing over different object shapes,we consider multiple 3D keypoints that are attached at relevant locations of the objects by expertknowledge and observed with a perfect keypoint detector [8].
See Fig.
2b for a visualization of thosekeypoints.
The keypoints both provide information about object shape and its pose.
Furthermore, asseen in Fig.
2b, they have been chosen to reflect those locations in the environment relevant to solvethe task.
Additionally, we report results where the state is represented by the poses of the objects – asthis cannot represent object shape, in this case we use a constant object shape for training and test.
Experiments
We evaluate our proposed method on different environments where the geometry of the objects in thescene is important to solve the task successfully.
Please also refer to the video in the supplementarymaterial.
Commonly, RL is trained and evaluated on a single environment, where only the posesare changed, but the involved object shapes are kept constant.
Since latent-conditioned NeRFs havebeen shown to be capable of generalizing over geometry [43], we consider experiments where werequire the RL agent to generalize over object shapes within some distribution.
Answering thescientific question of this work requires environments with multi-view observations — and for thecompositional versions object masks as well.
These are not provided in standard RL benchmarks ,which is the reason for choosing the environments investigated in this work.
We use PPO [86] as the
RL algorithm and four camera views in all experiments.
Refer to the appendix for more details aboutour environments, parameter choices, network architectures, and training times.
6.1 Environments
Mug on Hook.
In this environment, adopted from [87] and visualized in Fig.
2b, the task is tohang a mug on a hook.
Both the mug and the hook shape are randomized.
The actions are small3D translations applied to the mug.
This environment is challenging as we require the RL agent togeneralize over mug and hook shapes and the tolerance between the handle opening and the hook isrelatively small.
Further, the agent receives a sparse reward only if the mug has been hung stably.
This reward is calculated by virtually simulating a mug drop after each action.
If the mug would notfall onto the ground from the current state, a reward of one is assigned, otherwise zero.
Planar Pushing.
The task in this environment, shown in Fig.
3b, is to push yellow box-shaped objectsinto the left region of the table and blue objects into the right region with the red pusher that canmove in the plane, i.e., the action is two dimensional.
This is the same environment as in [53] withthe same four different camera views.
Each run contains a single object on the table (plus the pusher).
If the box has been pushed inside its respective region, a sparse reward of one is received, otherwisezero.
The boxes in the environment have different sizes, two colors and are randomly initialized.
Inthis environment, we cannot use keypoints for the multi-shape setting, as the reward depends on theobject color; we evaluate the keypoints baseline only in the single shape case (Appendix).
Door Opening.
Fig.
4b shows the door environment, where the task is to open a sliding door withthe red end-effector that can be translated in 3 DoFs as the action.
To solve this task, the agent hasto push on the door handle.
As the handle position and size is randomized, the agent has to learnto interact with the handle geometry accordingly.
Interestingly, as can be seen in the video in thesupplementary material, the agent often chooses to push on the handle only at the beginning, as,afterwards, it is sufficient to push the door itself at its side.
The agent receives a sparse reward if thedoor has been opened sufficiently, otherwise, zero reward is assigned.
6.2 Results
Figs 2a, 3a, 4a show success rates (averaged over 6 independent experiment repetitions and over30 test rollouts per repetition per timestep) as a function of training steps.
Also shown are the 68%confidence intervals.
These success rates have been evaluated using randomized object shapes andinitial conditions, and therefore reflect the agent’s ability to generalize over these.
In all these experiments, a latent space trained with compositional NeRF supervision as the decoderconsistently outperformed all other learned representations, both in terms of sample efficiency andasymptotic performance.
Furthermore, our proposed framework with compositional NeRF evenoutperforms the expert keypoint representation.
For the door environment, the 3D neural field encoderplus NeRF decoder (NeRF-RL comp.
+ field) reaches nearly perfect success rates.
For the other twoenvironments, the compositional 2D CNN encoder plus NeRF decoder (NeRF-RL comp.
+ image)was slightly better than with the neural field encoder but not significantly.
This shows that the decoderbuilt from compositional NeRF is relevant for the performance, not so much the choice of the encoder.
Training the 3D neural field encoder with a contrastive loss as supervision signal for different cameraviews as positive/negative training pairs is not able to achieve significant learning progress in thesescenarios (Multi-CURL).
However, the other contrastive baseline, CURL, which has a differentencoder and uses image cropping as data augmentation instead of additional camera views, is able toachieve decent performance and sample efficiency on the door environment, but not for the pushingenvironment.
In the mug environment, CURL initially is able to make learning progress comparableto our framework, but never reaches a success rate above 59% and then becomes unstable.
Similarly,the global CNN autoencoder baseline shows decent learning progress initially on the mug and pushingscenario (not for the door), but then becomes unstable (mug) or never surpasses 50% success rate(pushing).
Such variations in performance or instable learning across the different environments havenot been observed with our method, which is stable in all cases.
The compositional variant (NeRF-RL comp.) of our framework achieves the highest performance.
Since the conv.
comp.
autoencoder baseline has worse performance than its global variant, composi- encoder decoder comp.
NeRF loss
NeRFcomp.+field 3D CNN comp.
3D NeRF ✓ ✓ image reconstr.
: L2 RL comp.+image 2D CNN comp.
3D NeRF ✓ ✓ image reconstr.
: L2 (ours) global+image 2D CNN global 3D NeRF ✗ ✓ image reconstr.
L
Conv.
Autoencoder, c 2D CNN comp.
2D CNN ✓ ✗ image reconstr.:
L2 Conv.
Autoencoder, g 2D CNN 2D CNN ✗ ✗ image reconstr.
: L2 CURL 2D CNN - ✗ ✗ contrast: InfoNCE Multi-CURL 3D CNN - ✓ ✗ contrast: InfoNCE
[Figure]
chosen by expert knowledge and perfect extraction
[Figure] [Figure]
[Figure] [Figure]
[Figure]
tionality alone is not the sole reason for the better performance of our state representation.
Indeed,the global NeRF-RL + image variant in the pushing env. is also better than all other baselines.
Discussion
Why NeRF provides better supervision.
The NeRF training objective (1) strongly forces eachf ( · , z j ) to represent each object in its actual 3D configuration and relative to other objects in the scene(compositional case), including their shape.
This implies that the latent vectors z j have to contain thisinformation, i.e., they are trained to determine the object type, shape and pose in the scene.
In theglobal case, z 1 has to represent the geometry of the whole secne.
As the tasks we consider requirepolicies to take the geometry of the objects into account, we hypothesize that a latent vector that iscapable of parameterizing a NeRF to reconstruct the scene in the 3D space has to contain enough ofthe relevant 3D information of the objects also for the policy to be successful.
Masks.
In order for the auto-encoder framework to be compositional, it requires object masks.
Webelieve that instance segmentation has reached a level of maturity [88] that this is a fair assumptionto make.
As we also utilize the individual masks for the compositional conv.
autoencoder and themulti-view CURL baseline, which do not show good performance, it indicates that the masks are notthe main reason that our state representation achieves higher performance.
This is further supportedby the fact that the global NeRF-RL variant which does not rely on individual object masks on thepushing scenario achieved a performance higher than all baselines, i.e., masks will increase theperformance of NeRF-RL as they enable the compositional version, but they do not seem essential.
Offline/Online.
In this work, we focused on pretraining the latent representation offline from adataset collected by random actions.
During RL, the encoder is fixed and only the policy networks arelearned.
This has the advantage that the same representation can be used for different RL tasks and thedataset to train the representation not necessarily has to come from the same distribution.
However,if a policy is needed to explore reasonable regions of the state space, collecting a dataset offline tolearn a latent space that covers the state space sufficiently might be more challenging for an offlineapproach.
This was not an issue for our experiments where data collection with random actions wassufficient.
Indeed, we show generalization over different starting states of the same environment andwith respect to different shapes (within distribution).
Future work could investigate NeRF supervisionin an online setup.
Note that the reconstruction loss via NeRF is computationally more demandingthan via a 2D CNN deconv.
decoder or a contrastive term, making NeRF supervision as an auxiliaryloss at each RL training step costly.
One potential solution for this is to apply the auxiliary loss not atevery RL training step, but with a lower frequency.
Regarding computational efficiency, this is wherecontrastive learning has an advantage over our proposed NeRF-based decoder, as the encoding with
CURL can be trained within half a day, whereas the NeRF auto-encoder took up to 2 days to train forour environments.
However, when using the encoder for RL, there is no difference in inference time.
Multi-View.
The auto-encoder framework we propose can fuse the information of multiple cameraviews into a latent vector describing an object in the scene.
This way, occlusions can be addressed andthe agent can gain a better 3D understanding of the scene from the different camera angles.
Havingaccess to multiple camera views and their camera matrices is an additional assumption we make,although we believe the capability to utilize this information is an advantage of our method.
Conclusion
In this work, we have proposed the idea to utilize Neural Radiance Fields (NeRFs) to train latentspaces for RL.
Our environments focus on tasks where the geometry of the objects in the scene isrelevant for successfully solving the tasks.
Training RL agents with the pretrained encoder that mapsmultiple views of the scene to a latent space consistently outperformed other ways of learning a staterepresentation and even keypoints chosen by expert knowledge.
Our results show that the 3D priorpresent in compositional NeRF as the decoder is more important than priors in the encoder.
Broader Impacts.
Our main contribution is a method to learn representations that improve theefficiency of vision-based RL, which could impact automation.
As such, our work inherits generalethical risks of AI, like the question of how to address the potential of increased automation in society.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[Yes]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] Seesupplementary material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] See appendix.(c)
Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See plots in main paper.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See appendix.4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a)
If your work uses existing assets, did you cite the creators?
[Yes] In the paper and inthe code.(b) Did you mention the license of the assets?
[Yes] In the code.(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]See supplementary material.(d)
Did you discuss whether and how consent was obtained from people whose data you’re
[Table]
using/curating?
[Yes] In the code.
[Table]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
[Table]
information or offensive content?
[N/A]
[Table]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
[Table]
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation?
[N/A]
[Table]