• Stars
    star
    147
  • Rank 242,560 (Top 5 %)
  • Language
    Python
  • License
    Other
  • Created 11 months ago
  • Updated 8 months ago

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

Official Code for RVT: Robotic View Transformer for 3D Object Manipulation

PWC

RVT: Robotic View Transformer for 3D Object Manipulation
Ankit Goyal, Jie Xu, Yijie Guo, Valts Blukis, Yu-Wei Chao, Dieter Fox
CoRL 2023 (Oral)

If you find our work useful, please consider citing:

@article{,
  title={RVT: Robotic View Transformer for 3D Object Manipulation},
  author={Goyal, Ankit and Xu, Jie and Guo, Yijie and Blukis, Valts and Chao, Yu-Wei and Fox, Dieter},
  journal={CoRL},
  year={2023}
}

Getting Started

Install RVT

  • Tested (Recommended) Versions: Python 3.8. We used CUDA 11.1.

  • Step 1 (Optional): We recommend using conda and creating a virtual environment.

conda create --name rvt python=3.8
conda activate rvt
  • Step 2: Install PyTorch. Make sure the PyTorch version is compatible with the CUDA version. One recommended version compatible with CUDA 11.1 and PyTorch3D can be installed with the following command. More instructions to install PyTorch can be found here.
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
  • Step 3: Install PyTorch3D. One recommended version that is compatible with the rest of the library can be installed as follows. Note that this might take some time. For more instructions visit here.
curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz
tar xzf 1.10.0.tar.gz
export CUB_HOME=$(pwd)/cub-1.10.0
pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'

Once you have downloaded CoppeliaSim, add the following to your ~/.bashrc file. (NOTE: the 'EDIT ME' in the first line)

export COPPELIASIM_ROOT=<EDIT ME>/PATH/TO/COPPELIASIM/INSTALL/DIR
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$COPPELIASIM_ROOT
export QT_QPA_PLATFORM_PLUGIN_PATH=$COPPELIASIM_ROOT
export DISLAY=:1.0

Remember to source your .bashrc (source ~/.bashrc) or .zshrc (source ~/.zshrc) after this.

  • Step 5: Clone the repository with the submodules using the following command.
git clone --recurse-submodules [email protected]:NVlabs/RVT.git && cd RVT && git submodule update --init

Now, locally install RVT and other libraries using the following command. Make sure you are in folder RVT.

pip install -e . 
pip install -e rvt/libs/PyRep 
pip install -e rvt/libs/RLBench 
pip install -e rvt/libs/YARR 
pip install -e rvt/libs/peract_colab
  • Step 6: Download dataset.
    • For experiments on RLBench, we use pre-generated dataset provided by PerAct. Please download and place them under RVT/rvt/data/xxx where xxx is either train, test, or val.

    • Additionally, we use the same dataloader as PerAct, which is based on YARR. YARR creates a replay buffer on the fly which can increase the startup time. We provide an option to directly load the replay buffer from the disk. We recommend using the pre-generated replay buffer (98 GB) as it reduces the startup time. You can either download replay.tar.xz which contains the replay buffer for all tasks or replay buffer for indidual tasks. After downloading, uncompress the replay buffer(s) (for example using the command tar -xf replay.tar.xz) and place it under RVT/rvt/replay/replay_xxx where xxx is either train or val. Note that is useful only if you want to train RVT from scratch and not needed if you want to evaluate the pre-trained model.

Using the library:

Training RVT

Default command

To train RVT on all RLBench tasks, use the following command (from folder RVT/rvt):

python train.py --exp_cfg_path configs/all.yaml --device 0,1,2,3,4,5,6,7

We use 8 V100 GPUs. Change the device flag depending on available compute.

More details about train.py
  • default parameters for an experiment are defined here.
  • default parameters for rvt are defined here.
  • the parameters in for experiment and rvt can be overwritten by two ways:
    • specifying the path of a yaml file
    • manually overwriting using a opts string of format <param1> <val1> <param2> <val2> ..
  • Manual overwriting has higher precedence over the yaml file.
python train.py --exp_cfg_opts <> --mvt_cfg_opts <> --exp_cfg_path <> --mvt_cfg_path <>

The following command overwrites the parameters for the experiment with the configs/all.yaml file. It also overwrites the bs parameters through the command line.

python train.py --exp_cfg_opts "bs 4" --exp_cfg_path configs/all.yaml --device 0

Evaluate on RLBench

Evaluate RVT on RLBench

Download the pretrained RVT model. Place the model (model_14.pth trained for 15 epochs or 100K steps) and the config files under the folder runs/rvt/. Run evaluation using (from folder RVT/rvt):

python eval.py --model-folder runs/rvt  --eval-datafolder ./data/test --tasks all --eval-episodes 25 --log-name test/1 --device 0 --headless --model-name model_14.pth
Evaluate the official PerAct model on RLBench

Download the officially released PerAct model. Put the downloaded policy under the runs folder with the recommended folder layout: runs/peract_official/seed0. Run the evaluation using:

python eval.py --eval-episodes 25 --peract_official --peract_model_dir runs/peract_official/seed0/weights/600000 --model-name QAttentionAgent_layer0.pt --headless --task all --eval-datafolder ./data/test --device 0 

Gotchas

  • If you get qt plugin error like qt.qpa.plugin: Could not load the Qt platform plugin "xcb", try uninstalling opencv-python and installing opencv-python-headless
pip uninstall opencv-python                                                                                         
pip install opencv-python-headless
  • If you have CUDA 11.7, an alternate installation strategy could be to use the following command for Step 2 and Step 3. Note that this is not heavily tested.
# Step 2:
pip install pytorch torchvision torchaudio
# Step 3:
pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'
  • If you are having issues running evaluation on a headless server, please refer to #2 (comment).

  • If you want to generate visualization videos, please refer to #5.

FAQ's

Q. What is the advantag of RVT over PerAct?

RVT is both faster to train and performs better than PerAct.

Q. What resources are required to train RVT?

For training on 18 RLBench tasks, with 100 demos per task, we use 8 V100 GPUs (16 GB memory each). The model trains in ~1 day.

Note that for fair comparison with PerAct, we used the same dataset, which means duplicate keyframes are loaded into the replay buffer. For other datasets, one could consider not doing so, which might further speed up training.

Q. Why do you use pe_fix=True in the rvt config?

For fair comparison with offical PerAct model, we use this setting. More detials about this can be found in PerAct code. For future, we recommend using pe_fix=False for language input.

Q. Why are the results for PerAct different from the PerAct paper?

In the PerAct paper, for each task, the best checkpoint is chosen based on the validation set performance. Hence, the model weights can be different for different tasks. We evaluate PerAct and RVT only on the final checkpoint, so that all tasks are strictly evaluated on the same model weights. Note that only the final model for PerAct has been released officially.

Q. Why is there a variance in performance on RLBench even when evaluting the same checkpoint?

We hypothesize that it is because of the sampling based planner used in RLBench, which could be the source of the randomization. Hence, we evaluate each checkpoint 5 times and report mean and variance.

Q. Why did you use a cosine decay learning rate scheduler instead of a fixed learning rate schedule as done in PerAct?

We found the cosine learning rate scheduler led to faster convergence for RVT. Training PerAct with our training hyper-parameters (cosine learning rate scheduler and same number of iterations) led to worse performance (in ~4 days of training time). Hence for Fig. 1, we used the official hyper-parameters for PerAct.

Q. For my use case, I want to render images at real camera locations (input camera poses) with PyTorch3D. Is it possible to do so and how can I do that?

Yes, it is possible to do so. A self-sufficient example is present here. Depending on your use case, the code may need be modified. Also note that 3D augmentation cannot be used while rendering images at real camera locations as it would change the pose of the camera with respect to the point cloud.

For questions and comments, please contact Ankit Goyal.

Acknowledgement

We sincerely thank the authors of the following repositories for sharing their code.

License

License Copyright © 2023, NVIDIA Corporation & affiliates. All rights reserved.

This work is made available under the Nvidia Source Code License. The pretrained RVT model is released under the CC-BY-NC-SA-4.0 license.

More Repositories

1

instant-ngp

Instant neural graphics primitives: lightning fast NeRF and more
Cuda
15,102
star
2

stylegan

StyleGAN - Official TensorFlow Implementation
Python
13,882
star
3

stylegan2

StyleGAN2 - Official TensorFlow Implementation
Python
10,740
star
4

SPADE

Semantic Image Synthesis with SPADE
Python
7,518
star
5

stylegan3

Official PyTorch implementation of StyleGAN3
Python
6,108
star
6

neuralangelo

Official implementation of "Neuralangelo: High-Fidelity Neural Surface Reconstruction" (CVPR 2023)
Python
4,125
star
7

imaginaire

NVIDIA's Deep Imagination Team's PyTorch Library
Python
3,941
star
8

stylegan2-ada-pytorch

StyleGAN2-ADA - Official PyTorch implementation
Python
3,866
star
9

ffhq-dataset

Flickr-Faces-HQ Dataset (FFHQ)
Python
3,483
star
10

tiny-cuda-nn

Lightning fast C++/CUDA neural network framework
C++
3,286
star
11

eg3d

Python
3,089
star
12

MUNIT

Multimodal Unsupervised Image-to-Image Translation
Python
2,564
star
13

SegFormer

Official PyTorch implementation of SegFormer
Python
2,252
star
14

nvdiffrec

Official code for the CVPR 2022 (oral) paper "Extracting Triangular 3D Models, Materials, and Lighting From Images".
Python
2,019
star
15

few-shot-vid2vid

Pytorch implementation for few-shot photorealistic video-to-video translation.
Python
1,780
star
16

stylegan2-ada

StyleGAN2 with adaptive discriminator augmentation (ADA) - Official TensorFlow implementation
Python
1,778
star
17

FUNIT

Translate images to unseen domains in the test time with few example images.
Python
1,545
star
18

PWC-Net

PWC-Net: CNNs for Optical Flow Using Pyramid, Warping, and Cost Volume, CVPR 2018 (Oral)
Python
1,512
star
19

noise2noise

Noise2Noise: Learning Image Restoration without Clean Data - Official TensorFlow implementation of the ICML 2018 paper
Python
1,356
star
20

alias-free-gan

Alias-Free GAN project website and code
1,320
star
21

prismer

The implementation of "Prismer: A Vision-Language Model with Multi-Task Experts".
Python
1,287
star
22

DG-Net

👫 Joint Discriminative and Generative Learning for Person Re-identification. CVPR'19 (Oral) 👫
Python
1,268
star
23

nvdiffrast

Nvdiffrast - Modular Primitives for High-Performance Differentiable Rendering
C++
1,137
star
24

edm

Elucidating the Design Space of Diffusion-Based Generative Models (EDM)
Python
1,014
star
25

Deep_Object_Pose

Deep Object Pose Estimation (DOPE) – ROS inference (CoRL 2018)
Python
955
star
26

VoxFormer

Official PyTorch implementation of VoxFormer [CVPR 2023 Highlight]
Python
937
star
27

NVAE

The Official PyTorch Implementation of "NVAE: A Deep Hierarchical Variational Autoencoder" (NeurIPS 2020 spotlight paper)
Python
889
star
28

BundleSDF

[CVPR 2023] BundleSDF: Neural 6-DoF Tracking and 3D Reconstruction of Unknown Objects
Python
842
star
29

ODISE

Official PyTorch implementation of ODISE: Open-Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models [CVPR 2023 Highlight]
Python
779
star
30

GroupViT

Official PyTorch implementation of GroupViT: Semantic Segmentation Emerges from Text Supervision, CVPR 2022.
Python
679
star
31

FasterViT

[ICLR 2024] Official PyTorch implementation of FasterViT: Fast Vision Transformers with Hierarchical Attention
Python
664
star
32

GA3C

Hybrid CPU/GPU implementation of the A3C algorithm for deep reinforcement learning.
Python
641
star
33

denoising-diffusion-gan

Tackling the Generative Learning Trilemma with Denoising Diffusion GANs https://arxiv.org/abs/2112.07804
Python
634
star
34

genvs

610
star
35

sionna

Sionna: An Open-Source Library for Next-Generation Physical Layer Research
Jupyter Notebook
580
star
36

curobo

CUDA Accelerated Robot Library
Python
545
star
37

FB-BEV

Official PyTorch implementation of FB-BEV & FB-OCC - Forward-backward view transformation for vision-centric autonomous driving perception
Python
518
star
38

Dancing2Music

Python
513
star
39

planercnn

PlaneRCNN detects and reconstructs piece-wise planar surfaces from a single RGB image
Python
502
star
40

pacnet

Pixel-Adaptive Convolutional Neural Networks (CVPR '19)
Python
490
star
41

CALM

Python
486
star
42

DeepInversion

Official PyTorch implementation of Dreaming to Distill: Data-free Knowledge Transfer via DeepInversion (CVPR 2020)
Python
474
star
43

EmerNeRF

PyTorch Implementation of EmerNeRF: Emergent Spatial-Temporal Scene Decomposition via Self-Supervision
Python
456
star
44

FAN

Official PyTorch implementation of Fully Attentional Networks
Python
454
star
45

FourCastNet

Initial public release of code, data, and model weights for FourCastNet
Python
421
star
46

GCVit

[ICML 2023] Official PyTorch implementation of Global Context Vision Transformers
Python
414
star
47

intrinsic3d

Intrinsic3D - High-Quality 3D Reconstruction by Joint Appearance and Geometry Optimization with Spatially-Varying Lighting (ICCV 2017)
C++
411
star
48

nvdiffmodeling

Differentiable rasterization applied to 3D model simplification tasks
Python
404
star
49

flip

A tool for visualizing and communicating the errors in rendered images.
C++
375
star
50

wetectron

Weakly-supervised object detection.
Python
355
star
51

FoundationPose

FoundationPose: Unified 6D Pose Estimation and Tracking of Novel Objects
JavaScript
349
star
52

nvdiffrecmc

Official code for the NeurIPS 2022 paper "Shape, Light, and Material Decomposition from Images using Monte Carlo Rendering and Denoising".
C
341
star
53

geomapnet

Geometry-Aware Learning of Maps for Camera Localization (CVPR2018)
Python
338
star
54

GLAMR

[CVPR 2022 Oral] Official PyTorch Implementation of "GLAMR: Global Occlusion-Aware Human Mesh Recovery with Dynamic Cameras”.
Python
329
star
55

LSGM

The Official PyTorch Implementation of "LSGM: Score-based Generative Modeling in Latent Space" (NeurIPS 2021)
Python
326
star
56

ssn_superpixels

Superpixel Sampling Networks (ECCV2018)
Python
323
star
57

DiffiT

Official Repository for DiffiT: Diffusion Vision Transformers for Image Generation
315
star
58

FreeSOLO

FreeSOLO for unsupervised instance segmentation, CVPR 2022
Python
307
star
59

long-video-gan

Official PyTorch implementation of LongVideoGAN
Python
297
star
60

neuralrgbd

Neural RGB→D Sensing: Per-pixel depth and its uncertainty estimation from a monocular RGB video
Python
294
star
61

selfsupervised-denoising

High-Quality Self-Supervised Deep Image Denoising - Official TensorFlow implementation of the NeurIPS 2019 paper
Python
293
star
62

Taylor_pruning

Pruning Neural Networks with Taylor criterion in Pytorch
Python
279
star
63

timeloop

Timeloop performs modeling, mapping and code-generation for tensor algebra workloads on various accelerator architectures.
C++
278
star
64

metfaces-dataset

Python
272
star
65

few_shot_gaze

Pytorch implementation and demo of FAZE: Few-Shot Adaptive Gaze Estimation (ICCV 2019, oral)
Python
272
star
66

splatnet

SPLATNet: Sparse Lattice Networks for Point Cloud Processing (CVPR2018)
Python
268
star
67

MinVIS

Python
261
star
68

edm2

Analyzing and Improving the Training Dynamics of Diffusion Models (EDM2)
Python
261
star
69

contact_graspnet

Efficient 6-DoF Grasp Generation in Cluttered Scenes
Python
260
star
70

CenterPose

Single-Stage Keypoint-based Category-level Object Pose Estimation from an RGB Image (ICRA 2022)
Python
251
star
71

trajdata

A unified interface to many trajectory forecasting datasets.
Python
245
star
72

STEP

STEP: Spatio-Temporal Progressive Learning for Video Action Detection. CVPR'19 (Oral)
Python
244
star
73

matchlib

SystemC/C++ library of commonly-used hardware functions and components for HLS.
C++
235
star
74

sim-web-visualizer

Web Based Visualizer for Simulation Environments
Python
231
star
75

SCOPS

SCOPS: Self-Supervised Co-Part Segmentation (CVPR'19)
Python
221
star
76

UMR

Self-supervised Single-view 3D Reconstruction
Python
221
star
77

DiffRL

[ICLR 2022] Accelerated Policy Learning with Parallel Differentiable Simulation
Python
220
star
78

cule

CuLE: A CUDA port of the Atari Learning Environment (ALE)
C++
216
star
79

SSV

Pytorch implementation of SSV: Self-Supervised Viewpoint Learning from Image Collections (CVPR 2020)
Python
214
star
80

DiffPure

A new adversarial purification method that uses the forward and reverse processes of diffusion models to remove adversarial perturbations.
Python
210
star
81

latentfusion

LatentFusion: End-to-End Differentiable Reconstruction and Rendering for Unseen Object Pose Estimation
Python
197
star
82

I2SB

Python
194
star
83

nvbio

NVBIO is a library of reusable components designed to accelerate bioinformatics applications using CUDA.
C++
193
star
84

6dof-graspnet

Implementation of 6-DoF GraspNet with tensorflow and python. This repo has been tested with python 2.7 and tensorflow 1.12.
Python
186
star
85

NVBit

183
star
86

AFNO-transformer

Adaptive FNO transformer - official Pytorch implementation
Python
174
star
87

UnseenObjectClustering

Learning RGB-D Feature Embeddings for Unseen Object Instance Segmentation
Python
166
star
88

AL-MDN

Official pytorch implementation of Active Learning for deep object detection via probabilistic modeling (ICCV 2021)
Python
159
star
89

fermat

Fermat is a high performance research oriented physically based rendering system, trying to produce beautiful pictures following the mathematician’s principle of least time
C++
158
star
90

PoseCNN-PyTorch

PyTorch implementation of the PoseCNN framework
C
156
star
91

mask-auto-labeler

Python
153
star
92

mimicgen_environments

This code corresponds to simulation environments used as part of the MimicGen project.
Python
153
star
93

Bi3D

Python
150
star
94

condensa

Programmable Neural Network Compression
Python
146
star
95

traffic-behavior-simulation

Python
145
star
96

learningrigidity

Learning Rigidity in Dynamic Scenes with a Moving Camera for 3D Motion Field Estimation (ECCV 2018)
Python
144
star
97

ocrodeg

document image degradation
Jupyter Notebook
142
star
98

ocropus3

Repository collecting all the submodules for the new PyTorch-based OCR System.
Shell
141
star
99

CGBN

CGBN: CUDA Accelerated Multiple Precision Arithmetic (Big Num) using Cooperative Groups
Cuda
139
star
100

PL4NN

Perceptual Losses for Neural Networks: Caffe implementation of loss layers based on perceptual image quality metrics.
Python
138
star