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]
multi-
[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 supervision 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]
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 them 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 followthe hyperparameters and design of [5].
CURL operates on a single input view and we choose a viewfor this baseline from which the state of the environment can be inferred as best as possible (Fig. 17).
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]
A Additional Experiments
[Figure]
A.1 Influence of Reconstruction Capabilities
In this section, we investigate for the mug environment how the reconstruction quality of the NeRF-RLauto-encoder influences the downstream RL performance.
The state representations of all experimentsreported in the paper have been trained for a maximum computational budget of 5 days.
The encodernetwork that has the lowest image reconstruction loss or contrastive loss is then chosen for the RLtraining procedure.
Now, for the NeRF-RL compositional + image variant, we deliberately choose3 more networks that are trained for fewer epochs and hence have a lower reconstruction quality.
Fig.
5 qualitatively shows the reconstruction quality for these networks.
As one can see in Fig.
5b,for a network trained on only 10 epochs, the reconstruction is still very rough and the handle isnot reconstructed properly yet.
After 30 epochs (Fig. 5c), the shape of the mug body and of thehook are reasonably reconstructed, but the handle is still reconstructed as a “mean” handle shape.
In
Fig.
5d (100 epochs), the handle is reasonably reconstructed, but still a bit imprecise.
Finally, Fig.
5eshows the network after 405 epochs, which is also the one used for the main experiment, where thereconstruction, especially with respect to the handle, has become quite precise.
[Figure] [Figure]
Fig.
6 compares the learning curves in the mug environment for NeRF-RL comp.
+ image with staterepresentation networks trained for different amounts of epochs.
As one can see, there is a correlationbetween asymptotic performance and reconstruction quality.
The better the reconstruction quality,the higher the performance.
An explanation for this is that the encoder network trained for only10 epochs (see Fig.
5b) is unable to reconstruct the mug handle precisely, hence the latent vectorlikely does not contain the necessary information about the handle shape either.
As the mug shapedistribution also contains handles that have larger openings, even without representing the handleshape precisely, an agent can have some success (red curve in Fig. 6).
However, it never reachesthe performance that is possible if the state representation is capable of precisely reconstructing thehandle shape (green curve), as for mugs with a small handle size, the policy needs this geometricalinformation in order to be successful.
A.2 Influence of Data Size for Encoder Pre-Training
In this section we investigate, for the mug and pushing environment, the effect on the RL performanceif the auto-encoder is trained with less data.
In Fig.
7 one can see that for the pushing scenario, having 10 times less data available (10,000instead of 100,000) to train the encoder offline does not affect the performance with statisticalsignificance.
Similarly, for the mug on hook scenario, as shown in Fig. 8, halving the dataset from10,000 to 5,000 samples does not affect performance either.
Note that, as described in Sec.
B.1, theparameter space from which the mug environment is sampled is 11 dimensional.
Despite neither5,000 (this experiment) nor 10,000 (original experiment) data points cover this space at all, theauto-encoder can be trained to learn a latent space that allows the RL policy to effectively generalizeto all possible shapes within the distribution.
Generally, the offline data required to train our latentstate representation is orders of magnitudes smaller than the number of samples during online RL.
A.3 Influence of Number of Views
In order for an MDP to be well-defined, the relevant state information (in the first-order Markov sense)has to be given to the policy.
If only one single view is used, and relevant parts of the environmentare, e.g., occluded, this assumption can be broken.
This problem becomes less severe if multipleviews of the scene are readily available, which is often the case in robotics applications.
In thesescenarios, NeRF-RL enables the fusion of multiple views into one single state representation.
On the other hand, it is useful to know to which extent NeRF-RL relies on having access to multipleviews.
In all experiments so far, we utilized 4 views of the environment as input to the encoder.
Ascan be seen from Fig.
7, the performance of NeRF-RL is largely unaffected by using two views oreven one single view in the pushing environment.
The single-view setup for the pushing environment
[Figure] [Figure]
[Figure]
uses a single top-down view.
Since in this environment, the movement of the object happens mostlyon the table, this setup represents a 2D task.
Therefore, these results show that NeRF-RL can also beapplied to scenes that are 2D in nature.
However, the focus of this work is on 3D scenes, where the3D geometry, as opposed to the pushing scenario, is crucial for the policy to be successful.
Therefore,in the more complicated mug environment, as shown in Fig. 8, the performance drops if we use 2instead of 4 views, but still leads to acceptable performance.
A.4 Generalization to More Objects (Pushing Scenario)
The compositional variants of the auto-encoder are able to generalize to different numbers of objectsin the scene without retraining.
Therefore, we can take a policy trained on a scenario with a singleobject and run it on scenes containing more than one object.
Fig.
9 reports learning curves for the pushing environment with two objects in the scene.
Theselearning curves show the evaluations on the two object scene environments for policies that weretrained for different amounts of time-steps on the original environment with only one object.
Ourapproach generalizes well in this case, while the (compositional) baselines fail.
In the supplementarymaterial, we show a video of a NeRF-RL policy solving a scene containing two objects.
A.5 Learning Curves
In this section, we report learning curves for all environments that include the global version of the
NeRF-RL framework (NeRF-RL global + image) as well as an extended keypoint version for themug environment.
As can be seen in Figures 10a and 10b (dark blue curve), although slightly worse than the compositional variant, the global version of NeRF-RL (NeRF-RL global + image) outperforms all otherbaselines where the encoder is not trained with NeRF supervision.
As discussed in Sec. 7, thisindicates that the masks are not the main reason for the improved performance of NeRF-RL.
Further,both the Conv.
Autoencoder, c and Multi-CURL baseline also use the masks in the same way as
NeRF-RL comp.
However, as shown in Fig. 10c, for the door environment, the global NeRF-RLversion does not learn anything significant.
This can be explained by the fact that for the doorenvironment, the auto-encoder NeRF-RL global + image did not learn a correct reconstruction of the pusher, therefore the latent code presumably does not represent the pusher’s position, which is anecessary prerequisite to solve the task.
Fig.
12 shows additional keypoints baselines (“keypoints+”, “keypoints larger network”, “keypoints++”) for the mug on hook environment.
For the “keypoints+” baseline, we choose a largernumber of keypoints, with additional ones located at the corner points of the mug handle, and onthe opposite side of the mug (visualized in Fig. 11).
For “keypoints++”, we additionally project the“keypoints+” locations into a total of 128 values with a fixed matrix, leading to the same input sizeto the policy/value network as for all other baselines that use a latent state representation.
Hence,in this case, the policy/value network has the same number of parameters as for NeRF-RL or theother latent space baselines.
Similarly, for “keypoints larger network”, we increase the number ofhidden units in the first layer of the policy/value networks such that these networks have the samenumber of parameters, despite the input dimension being smaller. Note that, as mentioned above, forthe other baselines, the policy/value networks have the same number of parameters as NeRF-RL. Theperformance with more keypoints (“keypoints+”, “keypoints++”) is slightly better than the “keypoints”baseline with less keypoints, but is still surpassed by compositional NeRF-RL.
The “keypoints largernetwork” baseline is able to reach similar asymptotic performance as NeRF-RL, but learns slowerthan NeRF-RL.
[Figure]
A.6 Mask Perturbations
For the compositional variants of the auto-encoders (NeRF-RL global + image, Conv.
Autoencoder,g), we assumed to have access to object masks, cf.
discussion in Sec.
In this experiment, we [Figure] investigate how the performance of an agent that was trained with perfect masks is deteriorated whenthe masks are perturbed.
This measures the robustness of the framework to mask perturbations, inscenarios where an agent is not allowed to adapt to these mask perturbations.
For each version of NeRF-RL, we train 6 independent agents on unperturbed input, create checkpointsat intervals of 50,000 steps, evaluate these checkpoints by running 30 test rollouts, and then for eachexperiment select the checkpoint with the highest average success rate on these test rollouts.
We thentest the success rate of these trained agents if we apply mask perturbations to the input before feedingthem into the encoder.
We randomly remove 2 (low perturbation) or 6 (high perturbation) patchesfrom the object masks of each object in each view.
As mentioned before, the agents are trained usingunperturbed masks only.
In the mug hanging environment (Fig. 13), we find that the performance of the NeRF-RL agents isvery robust even against quite strong perturbations in the input masks.
As can be seen in Fig.
13b, theperturbations remove part of the handle from the mug mask, but the model is still able to reconstructthe handle shape.
In box pushing (Fig. 14) and door opening (Fig. 15), the size of the red pusher in the image is smallcompared to the severeness of the mask perturbations.
Therefore, already the low perturbation veryoften removes the majority of the red pusher from the input to the encoder, hence the correspondinglatent vector representing the red pusher has nearly no information about its position in space.
Consequently, the agent’s performance deteriorates faster compared to the mug scenario.
Nevertheless, especially the NeRF-RL comp.
+ image variant can achieve reasonable performance even when thered pusher (and the box for the pushing scenario) is nearly invisible due to the mask perturbations.
[Figure]
[Figure] [Figure]
A.7 Single Shape
In this experiment, we consider the same environments as in Sec.
6.1 but with a fixed object shape.
This means that only the initial configuration is randomized in each run, but the same object shape isused throughout training and evaluation.
Hence, the agents are not required to generalize over objectshapes.
This allows us to consider low-dimensional pose-like state input (3D state) as a baseline.
In case of the mug environment, this is the 3D input of the position of the mug, for the pushingenvironment a 8D input (2D position of the pusher, 2D position of the box, 4D quaternion for rotationof box), and 4D for the door environment (3D position of the red pusher, 1D opening of the door).
As can be seen in Figures 16a and 16b, for the mug and pushing environment, even in the single shapecase, the NeRF-RL state representation outperforms both direct 3D state and keypoint representations,not only in terms of asymptotic performance, but especially in terms of sample efficiency.
Thekeypoint representation works better than the (lower dimensional) 3D state.
Fig.
16c shows that incase of the door environment, the keypoint representation leads to higher performance comparedto NeRF-RL comp.
+ field.
We hypothesize that when using a single shape, there is not enoughrandomness in the door environment for the agent to achieve a sufficient success signal by exploration(sparse reward setup), which can also be seen by the much larger confidence intervals.
Nevertheless,the overall best performance achieved in the single shape door case is lower than what our method(NeRF-RL comp.
+ field and NeRF-RL comp.
+ image) is able to achieve in the more challengingmulti-shape door environment task, see Fig. 4a.
[Figure] [Figure]
B Experiment/Environment Details
This section provides details about our environments, parameter choices and training times.
Pleaserefer to the provided code for details about the exact network architectures.
We use PPO [86] as the RL algorithm from the stable-baselines-3 [89] implementation for allexperiments.
The hyperparameters of the PPO algorithm are the default ones from [89] for allenvironments and all experiments.
The dimension of the latent vector is k = 64 for each object, or incase a global framework is used, it is 64 times the number of objects in the environment.
This way,the total dimensionality of the latent vector describing the scene is always the same, independentlyfrom whether a compositional- or global-style framework is used.
All comparison methods utilize thesame amount of training data, as well as, where applicable, use the same network sizes.
For CURL,we used the same architecture as proposed in the original publication, which resulted in NeRF-RLhaving larger network sizes for the mug and door environment, but a smaller network size in the boxscenario than the CURL baseline.
All experiments utilize four camera views as input to the encoder.
Please refer to the supplementaryvideo to see a visualization of these input views.
B.1 Mug on Hook
In this environment, adopted from [87] and visualized in Fig.
2b, the task is to hang a mug on a hook.
Both the mug and the hook shape are randomized.
More specifically, the height, radius, handle size,and handle position of the mug vary.
This leads to an 8 dimensional parameter space for the shape ofthe objects in the scene.
The actions are small 3D translations applied to the mug.
This environment is challenging as we require the RL agent to generalize over mug and hook shapesand the tolerance between the handle opening and the hook is relatively small.
Further, the agentreceives a sparse reward only if the mug has been hung stably.
This reward is calculated by virtuallysimulating a mug drop after each action.
If the mug would not fall onto the ground from the currentstate, a reward of one is assigned, otherwise zero.
To collect the data for training the state representations we uniformly sample the mug position withinthe workspace region and only add those samples to the dataset where the mug and the hook arenot in collision.
In total, the dataset contains 10,000 configurations.
As this includes not only theposition of the mug, but also variations in shape of the mug and the hook, this is not a dense coverageof the 11 dimensional (8 shape + 3 initial mug position) parameter space from which the scenesare sampled.
Note that as the training data for the auto-encoder is collected by randomly samplingmug configurations, there are no trajectories in the dataset.
Furthermore, it is much less likely tosample a mug configuration that is stable.
Hence the data distribution the RL agent generates onlinepotentially differs from the offline dataset to train the state representation.
Please refer to the videofor a visualization of the offline data generation procedure.
Fig.
2b shows the 4 camera poses (blue coordinate frames), which are also visualized in the video.
Inthe compositional case, the mug and the hook are represented by their own latent vector, in the globalcase both the mug and the hook have one single latent vector together.
The keypoint baseline utilizes 4 keypoints, three on the mug and one on the hook.
Fig.
2b visualizesthese keypoints as red dots, one keypoint on the mug is not visible (center of the mug).
As can beseen, those keypoints have been chosen to reflect the relevant geometry to solve the task, in particularthe center position of the mug and the hook position.
B.2 Planar Pushing
The task in this environment, shown in Fig.
3b, is to push yellow box-shaped objects into the leftregion of the table and blue objects into the right region, using the red pusher that can move in theplane, i.e., the action is two dimensional.
This is the same environment as in [53] with the same fourdifferent camera views.
Each scenario contains a single object on the table (plus the pusher).
Theboxes in the environment have different sizes, two different colors (blue, yellow) and are randomlyinitialized (shape, color, and pose).
The box and the pusher are represented by a separate latent vectorfor the compositional case, and a single latent vector for both in the global case.
If the box has beenpushed inside its respective region, a sparse reward of one is received, otherwise zero.
These regionsare shown only for visualization purposes in Fig.
3b and not part of the input views for the network.
We use the same dataset as in [53] to train the auto-encoder, which has been collected by applyingmovements of the red pusher in a randomly chosen direction, biased towards the object center untileither the pusher leaves the workspace in which case a new direction is chosen or the box is pushedfrom the table in which case the environment is reset with a new randomly chosen box shape, colorand initial configuration of the pusher and box.
Note that the offline training data in particular does notcontain trajectories that push boxes in goal directed ways.
Please refer to the video for a visualizationof the offline data generation procedure.
In this environment, we cannot use keypoints for the multi-shape setting, as the reward depends onthe object color.
Hence, we evaluate the keypoints baseline only in the single shape case, cf.
Sec.
A.7.
B.3 Door Opening
[Figure]
Fig.
4b shows the door environment, where the task is to open a sliding door with the red end-effectorthat can be translated in 3 DoFs as the action.
To solve this task, the agent has to push on the doorhandle.
As the handle position and size are randomized, the agent has to learn to interact with thehandle geometry accordingly.
Interestingly, as can be seen in the video in the supplementary material,the agent often chooses to push on the handle only at the beginning, as, afterwards, it is sufficientto push the door itself at its side.
The agent receives a sparse reward if the door has been openedsufficiently, otherwise, zero reward is assigned.
In total, the dataset to train the latent space contains 50,000 configurations, which is collected byrandomly sampling a handle configuration (size and position on the door), door opening, and theinitial height of the end-effector.
The end-effector then moves in a random direction until it eithercannot move anymore or it leaves the workspace, in which case a new door handle, door opening,and starting position of the end-effector is sampled.
Please refer to the video for a visualization ofthe offline data generation procedure.
These random directions are in a 2D plane of the sampledinitial pusher height.
Therefore, during RL, the agent has to apply actions to solve the task leading totrajectories that are not part of the offline dataset used to train the state representation.
Similar to the other environments, we use four different camera views.
There are two latent vectors inthe compositional case, one representing the red end-effector, and the other the yellow door includingthe blue handle.
The keypoint baseline utilizes 5 keypoints, three on the door handle, one on the doorand one on the red pusher.
B.4 Computational Resources
The majority of our experiments have been performed on a GPU server containing 7 NVIDIA
RTX3090 GPUs, 1 TB RAM and two AMD EPYC 7452 CPUs.
Training the NeRF-based autoencoder on a single such GPU takes 2-3 days.
Training the convolutional auto-encoder baselines on a single GPU takes 3-5 days.
Training the CURL baselines takes half a day on a single GPU.
After therepresentations have been learned, RL training on a single GPU takes 4 to 5 days in the mug hangingenvironment, roughly 1 day in the box pushing environment, and 1 to 2 days in the door openingenvironment.
The longer RL training time for the mug environment is mainly caused by an inefficientcollision check in the respective simulator.