Robomimic

Highlights

- 🧠 New Algorithms (BC-Transformer, IQL)
- 🤖 Full compatibility with [robosuite v1.4](https://github.com/ARISE-Initiative/robosuite/releases/tag/v1.4.0) and DeepMind's MuJoCo bindings
- 👁️ The ability to use pre-trained image representations for policy learning
- 📈 Support for wandb logging
- Better data augmentation support
- Other quality of life improvements

New Algorithms

BC-Transformer

Transformer-based policies are becoming increasingly prevalent in the research community. We provide an implementation of BC-Transformer (Behavioral Cloning with a Transformer). It is a strong baseline for the robomimic datasets, often outperforming BC-RNN. We hope that this provides an easy way for our users to start using transformer-based policies.

We provide a [template configuration file](https://github.com/ARISE-Initiative/robomimic/tree/master/robomimic/config/default_templates/) for BC-Transformer that contains suggested parameters for the model. To learn how to train a transformer, either with our default configuration file or with custom parameters, please refer to our [tutorial](https://robomimic.github.io/docs/tutorials/training_transformers.html).

IQL

[Implicit Q-Learning (IQL)](https://arxiv.org/abs/2110.06169) is a popular offline RL method that has seen widespread use in the research community. We provide an implementation along with [benchmarking results](https://robomimic.github.io/docs/datasets/d4rl.html#d4rl-results) on [D4RL](https://arxiv.org/abs/2004.07219). Reproducing these results is easy -- see the instructions [here](https://robomimic.github.io/docs/datasets/d4rl.html#reproducing-d4rl-results).

Pre-trained Image Representations

We make it easy to use pre-trained image representations such as [R3M](https://arxiv.org/abs/2203.12601) and [MVP](https://arxiv.org/abs/2203.06173) for image-based agents in robomimic. See instructions [here](https://robomimic.github.io/docs/tutorials/pretrained_representations.html).

Wandb Logging

We now support logging with [wandb](https://wandb.ai/site) (in addition to [tensorboard](https://www.tensorflow.org/tensorboard)). Using it is simple. Follow [this guide](https://robomimic.github.io/docs/tutorials/viewing_results.html#viewing-wandb-results) for more details.

Better Data Augmentation Support

We now support observation randomizers, including crop randomizer, color randomizer, and gaussian randomizer. They will correctly detect the train/eval mode and change their behaviors accordingly.

- `CropRandomizer`: Take random crops of image-like observations.
- `ColorRandomizer`: Samples random color jitters in terms of brightness, contrast, saturation, and hue. Applies to RGB image observations.
- `GaussianRandomizer`: Add Gaussian noise to input. Applies to any observations.

For more details on how they work, please refer to the [documentation](https://robomimic.github.io/docs/modules/models.html?highlight=randomizer#randomizers)

Full compatibility with robosuite v1.4 and DeepMind's MuJoCo bindings

Our original [robomimic datasets](https://robomimic.github.io/docs/datasets/robomimic_v0.1.html) are now compatible with [robosuite v1.4](https://github.com/ARISE-Initiative/robosuite/releases/tag/v1.4.0) (the latest robosuite release) and the [DeepMind MuJoCo bindings](https://mujoco.readthedocs.io/en/latest/python.html#). The best part of this is that MuJoCo installation is much less painful now!

We highly recommend updating robosuite and switching from mujoco-py to DeepMind's MuJoCo bindings. All our [tests](https://github.com/ARISE-Initiative/robomimic/tree/master/tests) now assume robosuite v1.4 is installed, and our [dataset download script](https://robomimic.github.io/docs/datasets/robomimic_v0.1.html#downloading) download the versions compatible with robosuite v1.4 by default. You can still use the old versions by using the download script in robomimic v0.2, or using the download links in the [v0.2 documentation](https://robomimic.github.io/docs/v0.2/datasets/robomimic_v0.1.html).

Other Improvements

We outline other improvements here.
- Image datasets can use significantly less storage now -- we provide an option to use image compression and also exclude the `next_obs` key (which is unneeded for imitation learning algorithms) in our image extraction script. More details [here](https://robomimic.github.io/docs/datasets/robosuite.html#extracting-observations-from-mujoco-states).
- A new [colab notebook](https://robomimic.github.io/docs/introduction/getting_started.html#quickstart-colab-notebook) to get started easily
- Users must now specify train and validation hdf5 keys explicitly in the config when setting `config.experiment.validate = True`. This avoids any ambiguity in what data is being used for training and for validation, and offers more flexibility for using different training and validation splits. See more notes [here](https://robomimic.github.io/docs/datasets/overview.html#filter-keys).
- Moved some core observation processing models into `models/obs_core.py`
- Updated docs and tutorials

Contributor Spotlight

We would like to introduce the newest member of our robomimic core team [Matthew Bronars](https://github.com/MBronars), who helped with this release!

This release was a major team effort. Here is a breakdown of contributed features by team member.

- BC-Transformer (snasiriany, MBronars, amandlek)
- IQL (snasiriany)
- Pre-trained image representations (j96w, danfeiX, MBronars)
- Wandb logging (snasiriany)
- Better Data Augmentation Support (danfeiX, cremebrule)
- Compatibility with robosuite v1.4 (snasiriany)
- Colab notebook (danfeiX)
- Improvements to dataset storage and train-valid splits (amandlek)

Highlights
This release of robomimic brings integrated support for mobile manipulation datasets from the recent [MOMART](https://sites.google.com/view/il-for-mm/home) paper, and adds modular features for easily modifying and adding custom observation modalities and corresponding encoding networks.

MOMART Datasets
We have added integrated support for MOMART [datasets](https://sites.google.com/view/il-for-mm/datasets), a large-scale set of multi-stage, long-horizon mobile manipulation task demonstrations in a simulated kitchen environment collected in iGibson.

Using MOMART Datasets
Datasets can be easily downloaded using [download_momart_datasets.py](https://github.com/ARISE-Initiative/robomimic/tree/master/robomimic/scripts/download_momart_datasets.py).

For step-by-step instructions for setting up your machine environment to visualize and train with the MOMART datasets, please visit the [Getting Started](https://sites.google.com/view/il-for-mm/datasets#h.whukwluu16gm) page.

Modular Observation Modalities
We also introduce modular features for easily modifying and adding custom observation modalities and corresponding encoding networks. A **modality** corresponds to a group of specific observations that should be encoded the same way.

Default Modalities
robomimic natively supports the following modalities (expected size from a raw dataset shown, excluding the optional leading batch dimension):

- `rgb` (H, W, 3): Standard 3-channel color frames with values in range `[0, 255]`
- `depth` (H, W, 1): 1-channel frame with normalized values in range `[0, 1]`
- `low_dim` (N): low dimensional observations, e.g.: proprioception or object states
- `scan` (1, N): 1-channel, single-dimension data from a laser range scanner

We have default encoder networks which can be configured / modified by setting relevant parameters in your config, e.g.:

python
These keys should exist in your dataset
config.observation.modalities.obs.rgb = ["cam1", "cam2", "cam3"] Add camera observations to the RGB modality
config.observation.modalities.obs.low_dim = ["proprio", "object"] Add proprioception and object states to low dim modality
...

Now let's modify the default RGB encoder network and set the feature dimension size
config.observation.encoder.rgb.core_kwargs.feature_dimension = 128
...


To see the structure of the observation modalities and encoder parameters, please see the [base config](https://github.com/ARISE-Initiative/robomimic/blob/master/robomimic/config/base_config.py#L195) module.

Custom Modalities
You can also easily add your own modality and corresponding custom encoding network! Please see our example [add_new_modality.py](https://github.com/ARISE-Initiative/robomimic/tree/master/examples/add_new_modality.py).

Refactored Config Structure
With the introduction of modular modalities, our `Config` class structure has been modified slightly, and will likely cause breaking changes to any configs you have created using version 0.1.0. Below, we describe the exact changes in the config that need to be updated to match the current structure:

Observation Modalities
The `image` modality have been renamed to `rgb`. Thus, you would need to change your config in any places referencing `image` modality, e.g.:

python
Old format
config.observation.modalities.image.<etc>

New format
config.observation.modalities.rgb.<etc>


The `low_dim` modality remains unchanged. Note, however, that we have additionally added integrated support for both `depth` and `scan` modalities, and can be referenced in the same way, e.g.:

python
config.observation.modalities.depth.<etc>
config.observation.modalities.scan.<etc>


Observation Encoders / Randomizer Networks
We have modularized the encoder / randomizer arguments so that they are general, and are unique to each type of observation modality. All of the original arguments in v0.1.0 have been preserved, but are now re-formatted as follows:

python
OLD

Previously, a single set of arguments were specified, and was hardcoded to process image (rgb) observations

Assumes that you're using the VisualCore class, not general!
config.observation.encoder.visual_feature_dimension = 64
config.observation.encoder.visual_core = 'ResNet18Conv'
config.observation.encoder.visual_core_kwargs.pretrained = False
config.observation.encoder.visual_core_kwargs.input_coord_conv = False

For pooling, is hardcoded to use spatial softmax or not, not general!
config.observation.encoder.use_spatial_softmax = True
kwargs for spatial softmax layer
config.observation.encoder.spatial_softmax_kwargs.num_kp = 32
config.observation.encoder.spatial_softmax_kwargs.learnable_temperature = False

This is the first release of robomimic. This version should be used when trying to reproduce results from [the study](https://arxiv.org/abs/2108.03298).

Thankfully, the observation randomization network specifications were already modularized, but were hardcoded to process image (rgb) modality only. Thus, the only change we made is to allow the randomization kwargs to be specified per modality:

python
OLD
Previously, observation randomization was hardcoded for image / rgb modality
config.observation.encoder.obs_randomizer_class = None
config.observation.encoder.obs_randomizer_kwargs.crop_height = 76
config.observation.encoder.obs_randomizer_kwargs.crop_width = 76
config.observation.encoder.obs_randomizer_kwargs.num_crops = 1
config.observation.encoder.obs_randomizer_kwargs.pos_enc = False

NEW

Now, the randomization arguments are specified per modality. An example for RGB is shown below
config.observation.encoder.rgb.obs_randomizer_class = None
config.observation.encoder.rgb.obs_randomizer_kwargs.crop_height = 76
config.observation.encoder.rgb.obs_randomizer_kwargs.crop_width = 76
config.observation.encoder.rgb.obs_randomizer_kwargs.num_crops = 1
config.observation.encoder.rgb.obs_randomizer_kwargs.pos_enc = False


You can also view the [default configs](https://github.com/ARISE-Initiative/robomimic/tree/master/robomimic/exps/templates) and compare your config to these templates to view exact diffs in structure.