Renton-Ren/multi-task-tcc-robosuite

XIRL

• Overview
• Setup
• Datasets
• Code Navigation
• Experiments: Reproducing Paper Results
• Extending XIRL
• Acknowledgments

Overview

Code release for our CoRL 2021 conference paper:

XIRL: Cross-embodiment Inverse Reinforcement Learning Kevin Zakka1,3, Andy Zeng1, Pete Florence1, Jonathan Tompson1, Jeannette Bohg2, and Debidatta Dwibedi1 Conference on Robot Learning (CoRL) 2021

¹Robotics at Google, ²Stanford University, ³UC Berkeley

---

This repository serves as a general-purpose library for (a) self-supervised pretraining on video data and (b) downstream reinforcement learning using the learned representations as reward functions. It also contains models, training scripts and config files for reproducing our results and as a reference for implementation details.

Our hope is that the code's modularity allows you to easily extend and build on top of our work. To aid in this effort, we're releasing two additional standalone libraries:

• x-magical: our Gym-like benchmark extension of MAGICAL geared towards cross-embodiment imitation.
• torchkit: a lightweight library containing useful PyTorch boilerplate utilities like logging and model checkpointing.

For the latest updates, see: x-irl.github.io

Setup

We use Python 3.8 and Miniconda for development. To create an environment and install dependencies, run the following steps:

# Clone and cd into xirl.
git clone git@github.com:google-research/google-research.git --depth=1
cd google-research/xirl

# Create and activate environment.
conda create -n xirl python=3.8
conda activate xirl

# Install dependencies.
pip install -r requirements.txt

Datasets

X-MAGICAL

Run the following bash script to download the demonstration dataset for the X-MAGICAL benchmark:

bash scripts/download_xmagical_dataset.sh

The dataset will be located in /tmp/xirl/datasets/xmagical. You are free to modify the save destination, just make sure you update config.data.root in the pretraining config file — see base_configs/pretrain.py.

X-REAL

Our real-world dataset X-REAL will be released as soon as it gets approval, stay tuned!

Code Navigation

At a high-level, our code relies on two important but generic python scripts: pretrain.py for pretraining and train_policy.py for reinforcement learning. We use ml_collections to parameterize these scripts with experiment-specific config files. All experiments must use config files that inherit from the base config files in base_configs/. Specifically, pretraining experiments must inherit from base_configs/pretrain.py and RL experiments must inherit from base_configs/rl.py.

The rest of the codebase is organized as follows:

• configs/ contains all config files used in our CoRL submission. They inherit from base_configs/.
• xirl/ is the core pretraining codebase.
• sac/ is the core Soft-Actor-Critic implementation adapted from pytorch_sac.
• scripts/ contains miscellaneous bash scripts.

Experiments: Reproducing Paper Results

Core Scripts

• Same-embodiment setting (Section 5.1)
  • Pretraining: python pretrain_xmagical_same_embodiment.py --help
  • RL: python rl_xmagical_learned_reward.py --help
• Cross-embodiment setting (Section 5.2)
  • Pretraining: python pretrain_xmagical_cross_embodiment.py --help
  • RL: python rl_xmagical_learned_reward.py --help
• RL with environment reward
  • python rl_xmagical_env_reward.py --help
• Interactive reward visualization (Section 5.4)
  • python interact_reward.py --help

Misc. Scripts

• Visualize dataloader for frame sampler debugging
  • python debug_dataset.py --help
• Compute goal embedding with a pretrained model
  • python compute_goal_embedding.py --help
• Quick n' dirty multi-GPU RL training
  • With environment reward: bash scripts/launch_rl_multi_gpu.sh

Extending XIRL

How can I implement my own self-supervised pretraining algorithm?

You'll want to inherit from xirl.trainers.base.Trainer and implement the __init__ and compute_loss methods. For reference, take a look at xirl/trainers/tcc.py to see how Temporal Cycle Consistency is implemented. Make sure to add your new algorithm to the TRAINERS dict in factory.py.

How do I modify the way frames are sampled in the dataloader?

Create your own sampler in xirl/frame_samplers.py and add it to the FRAME_SAMPLERS dict in factory.py.

How can I implement additional pretraining evaluation metrics?

You'll want to inherit from xirl.evaluators.base.Evaluator class and as you guessed, add it to the EVALUATORS dict in factory.py. See xirl/evaluators for our current list of qualitative and quantative evaluation metrics

Acknowledgments

Many people have contibuted one way or another in the making and shaping of this repository. In no particular order, we'd like to thank Alex Nichol, Nick Hynes, Brent Yi, Jimmy Wu and Sam Toyer for their fruitful back-and-forth discussions.

Citation

If you find this code useful, consider citing our work:

@inproceedings{zakka2021xirl,
  author    = {Zakka, Kevin and Zeng, Andy and Florence, Pete and Tompson, Jonathan and Bohg, Jeannette and Dwibedi, Debidatta},
  title     = {XIRL: Cross-embodiment Inverse Reinforcement Learning},
  booktitle = {Proceedings of the 5th Conference on Robot Learning (CoRL)},
  year      = {2021},
}