Visual Reinforcement Learning with Self-Supervised 3D Representations

Original PyTorch implementation of RL3D from

Visual Reinforcement Learning with Self-Supervised 3D Representations by

Yanjie Ze*, Nicklas Hansen*, Yinbo Chen, Mohit Jain, Xiaolong Wang

<p align="center"> <br><img src='media/rl3d.png' width="600"/><br> </p>

Method

RL3D is a framework for visual reinforcement learning (RL) using a pretrained 3D visual representation and jointly training with an auxiliary view synthesis task. RL3D could generate novel view synthesis for diverse RL tasks and achieve good sample efficiency, robustness to sim-to-real transfer, and generalization to unseen environments.

Instructions

Assuming that you already have MuJoCo installed, install dependencies using conda:

conda env create -f environment.yaml
conda activate rl3d

After installing dependencies, you can train an agent by using the provided script

bash scripts/train.sh

Evaluation videos and model weights can be saved with arguments save_video=1 and save_model=1. Refer to the arguments.py for a full list of options and default hyperparameters.

The training script supports both local logging as well as cloud-based logging with Weights & Biases. To use W&B, provide a key by setting the environment variable WANDB_API_KEY=<YOUR_KEY>, set use_wandb=1, and add your W&B project and entity details in the script.

Pretrained Model

Import the encoder in our pretrained 3D visual representation

import load_3d
encoder_2d = load_3d.visual_representation(ckpt_path="checkpoints/videoae_co3d.tar", use_3d=False)

You could set use_3d=True to use the 3D encoder. The encoder is trained on the CO3D dataset.

Tasks

We provide our xArm environments for sim-to-real experiments in our paper. The task name and the action space are listed below:

• reach (xyz)
• push (xy)
• pegbox (xyz)
• lift (xyzw)

License & Acknowledgements

RL3D is licensed under the MIT license. MuJoCo is licensed under the Apache 2.0 license.