diff --git a/.gitignore b/.gitignore index 5a803b1..50d5978 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,9 @@ +# Folders generated by the repo +/datasets/ +training_results/ +paper/ + + # Mac OSX .DS_Store diff --git a/README.md b/README.md index 4f92f95..e443ed3 100644 --- a/README.md +++ b/README.md @@ -1,224 +1,46 @@ -# MimicGen Environments and Datasets +# MimicGen

- - +

-This repository contains the official release of simulation environments and datasets for the [CoRL 2023](https://www.corl2023.org/) paper "MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations". +This repository contains the official release of data generation code, simulation environments, and datasets for the [CoRL 2023](https://www.corl2023.org/) paper "MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations". -The datasets contain over 48,000 task demonstrations across 12 tasks. +The released datasets contain over 48,000 task demonstrations across 12 tasks and the MimicGen data generation tool can create as many as you'd like. Website: https://mimicgen.github.io Paper: https://arxiv.org/abs/2310.17596 +Documentation: https://mimicgen.github.io/docs/introduction/overview.html + For business inquiries, please submit this form: [NVIDIA Research Licensing](https://www.nvidia.com/en-us/research/inquiries/) ------- ## Latest Updates +- [07/09/2024] **v1.0.0**: Full code release, including data generation code - [04/04/2024] **v0.1.1**: Dataset license changed to [CC-BY 4.0](https://creativecommons.org/licenses/by/4.0/), which is less restrictive (see [License](#license)) - [09/28/2023] **v0.1.0**: Initial code and paper release ------- +## Useful Documentation Links -## Table of Contents - -- [Installation](#installation) -- [Downloading and Using Datasets](#downloading-and-using-datasets) -- [Reproducing Policy Learning Results](#reproducing-policy-learning-results) -- [Task Visualizations](#task-visualizations) -- [Data Generation Code](#data-generation-code) -- [Troubleshooting and Known Issues](#troubleshooting-and-known-issues) -- [License](#license) -- [Citation](#citation) - - -## Installation - -We recommend installing the repo into a new conda environment (it is called `mimicgen` in the example below): - -```sh -conda create -n mimicgen python=3.8 -conda activate mimicgen -``` - -You can install most of the dependencies by cloning the repository and then installing from source: - -```sh -cd -git clone https://github.com/NVlabs/mimicgen_environments.git -cd mimicgen_environments -pip install -e . -``` - -However, there are some additional dependencies that we list below. These are best installed from source: - -- [robosuite](https://robosuite.ai/) - - **Installation** - ```sh - cd - git clone https://github.com/ARISE-Initiative/robosuite.git - git checkout b9d8d3de5e3dfd1724f4a0e6555246c460407daa - cd robosuite - pip install -e . - ``` - - **Note**: the git checkout command corresponds to the commit we used for testing our policy learning results. In general the `master` branch (`v1.4+`) should be fine. - - For more detailed instructions, see [here](https://robosuite.ai/docs/installation.html) -- [robomimic](https://robomimic.github.io/) - - **Installation** - ```sh - cd - git clone https://github.com/ARISE-Initiative/robomimic.git - git checkout ab6c3dcb8506f7f06b43b41365e5b3288c858520 - cd robomimic - pip install -e . - ``` - - **Note**: the git checkout command corresponds to the commit we used for testing our policy learning results. In general the `master` branch (`v0.3+`) should be fine. - - For more detailed instructions, see [here](https://robomimic.github.io/docs/introduction/installation.html) -- [robosuite_task_zoo](https://github.com/ARISE-Initiative/robosuite-task-zoo) - - **Note**: This is optional and only needed for the Kitchen and Hammer Cleanup environments / datasets. - - **Installation** - ```sh - cd - git clone https://github.com/ARISE-Initiative/robosuite-task-zoo - git checkout 74eab7f88214c21ca1ae8617c2b2f8d19718a9ed - cd robosuite_task_zoo - pip install -e . - ``` - -Lastly, **please downgrade MuJoCo to 2.3.2**: -```sh -pip install mujoco==2.3.2 -``` - -**Note**: This MuJoCo version (`2.3.2`) is important -- in our testing, we found that other versions of MuJoCo could be problematic, especially for the Sawyer arm datasets (e.g. `2.3.5` causes problems with rendering and `2.3.7` changes the dynamics of the robot arm significantly from the collected datasets). - -### Test Your Installation - -The following script can be used to try random actions in a task. -```sh -cd mimicgen_envs/scripts -python demo_random_action.py -``` - -## Downloading and Using Datasets - -### Dataset Types - -As described in the paper, each task has a default reset distribution (D_0). Source human demonstrations (usually 10 demos) were collected on this distribution and MimicGen was subsequently used to generate large datasets (usually 1000 demos) across different task reset distributions (e.g. D_0, D_1, D_2), objects, and robots. - -The datasets are split into different types: - -- **source**: source human datasets used to generate all data -- this generally consists of 10 human demonstrations collected on the D_0 variant for each task. -- **core**: datasets generated with MimicGen for different task reset distributions. These correspond to the core set of results in Figure 4 of the paper. -- **object**: datasets generated with MimicGen for different objects. These correspond to the results in Appendix G of the paper. -- **robot**: datasets generated with MimicGen for different robots. These correspond to the results in Appendix F of the paper. -- **large_interpolation**: datasets generated with MimicGen using much larger interpolation segments. These correspond to the results in Appendix H in the paper. - -**Note 1**: All datasets are readily compatible with [robomimic](https://robomimic.github.io/) --- the structure is explained [here](https://robomimic.github.io/docs/datasets/overview.html#dataset-structure). This means that you can use robomimic to [visualize the data](https://robomimic.github.io/docs/tutorials/dataset_contents.html) or train models with different policy learning methods that we did not explore in our paper, such as [BC-Transformer](https://robomimic.github.io/docs/tutorials/training_transformers.html). - -**Note 2**: We found that the large_interpolation datasets pose a significant challenge for imitation learning, and have substantial room for improvement. - -### Dataset Statistics - -The datasets contain over 48,000 task demonstrations across 12 tasks. - -We provide more information on the amount of demonstrations for each dataset type: -- **source**: 120 human demonstrations across 12 tasks (10 per task) used to automatically generate the other datasets -- **core**: 26,000 task demonstrations across 12 tasks (26 task variants) -- **object**: 2000 task demonstrations on the Mug Cleanup task with different mugs -- **robot**: 16,000 task demonstrations across 4 different robot arms on 2 tasks (4 task variants) -- **large_interpolation**: 6000 task demonstrations across 6 tasks that pose significant challenges for modern imitation learning methods - -### Dataset Download - -#### Method 1: Using `download_datasets.py` (Recommended) - -`download_datasets.py` (located at `mimicgen_envs/scripts`) is a python script that provides a programmatic way of downloading the datasets. This is the preferred method, because this script also sets up a directory structure for the datasets that works out of the box with the code for reproducing policy learning results. - -A few examples of using this script are provided below: - -```sh -# default behavior - just download core square_d0 dataset -python download_datasets.py - -# download core datasets for square D0, D1, D2 and coffee D0, D1, D2 -python download_datasets.py --dataset_type core --tasks square_d0 square_d1 square_d2 coffee_d0 coffee_d1 coffee_d2 - -# download all core datasets, but do a dry run first to see what will be downloaded and where -python download_datasets.py --dataset_type core --tasks all --dry_run - -# download all source human datasets -python download_datasets.py --dataset_type source --tasks all -``` - -#### Method 2: Using Direct Download Links - -You can download the datasets manually through Google Drive. The folders each correspond to the dataset types described in [Dataset Types](#dataset-types). - -**Google Drive folder with all datasets:** [link](https://drive.google.com/drive/folders/14e9kkHGfApuQ709LBEbXrXVI1Lp5Ax7p?usp=drive_link) - -#### Method 3: Using Hugging Face - -You can download the datasets through Hugging Face. - -**Hugging Face dataset repository:** [link](https://huggingface.co/datasets/amandlek/mimicgen_datasets) - -## Reproducing Policy Learning Results - -After downloading the appropriate datasets you’re interested in using by running the `download_datasets.py` script, the `generate_training_configs.py` script (located at `mimicgen_envs/scripts`) can be used to generate all training config json files necessary to reproduce the experiments in the paper. A few examples are below. - -```sh -# Assume datasets already exist in mimicgen_envs/../datasets folder. Configs will be generated under mimicgen_envs/exps/paper, and training results will be at mimicgen_envs/../training_results after launching training runs. -python generate_training_configs.py - -# Alternatively, specify where datasets exist, and specify where configs should be generated. -python generate_training_configs.py --config_dir /tmp/configs --dataset_dir /tmp/datasets --output_dir /tmp/experiment_results -``` - -Then, to reproduce a specific set of training runs for different experiment groups (see [Dataset Types](#dataset-types)), we can simply navigate to the generated config directory, and copy training commands from the generated shell script there. As an example, we can reproduce the image training results on the Coffee D0 dataset, by looking for the correct set of commands in `mimicgen_envs/exps/paper/core.sh` and running them. The relevant section of the shell script is reproduced below. - -```sh -# task: coffee_d0 -# obs modality: image -python /path/to/robomimic/scripts/train.py --config /path/to/mimicgen_envs/exps/paper/core/coffee_d0/image/bc_rnn.json -``` - -**Note 1**: Another option is to directly run `robomimic/scripts/train.py` with any generated config jsons of interest -- the commands in the shell files do exactly this. - -**Note 2**: See the [robomimic documentation](https://robomimic.github.io/docs/introduction/getting_started.html) for more information on how training works. - -**Note 3**: In the MimicGen paper, we generated our datasets on versions of environments built on robosuite `v1.2`. Since then, we changed the environments and datasets (through postprocessing) to be based on robosuite `v1.4`. However, `v1.4` has some visual and dynamics differences from `v1.2`, so the learning results may not exactly match up with the ones we reported in the paper. In our testing on these released datasets, we were able to reproduce nearly all of our results, but within 10% of the performance reported in the paper. - - -## Task Visualizations - -We provide a convenience script to write videos for each task's reset distribution at `scripts/get_reset_videos.py`. Set the `OUTPUT_FOLDER` global variable to the folder where you want to write the videos, and set `DATASET_INFOS` appropriately if you would like to limit the environments visualized. Then run the script. - -The environments are also readily compatible with robosuite visualization scripts such as the [demo_random_action.py](https://github.com/ARISE-Initiative/robosuite/blob/b9d8d3de5e3dfd1724f4a0e6555246c460407daa/robosuite/demos/demo_random_action.py) script and the [make_reset_video.py](https://github.com/ARISE-Initiative/robosuite/blob/b9d8d3de5e3dfd1724f4a0e6555246c460407daa/robosuite/scripts/make_reset_video.py) script, but you will need to modify these files to add a `import mimicgen_envs` line to make sure that `robosuite` can find these environments. - - -**Note**: You can find task reset visualizations on the [website](https://mimicgen.github.io), but they may look a little different as they were generated with robosuite v1.2. - -## Data Generation Code - -If you are interested in the data generation code, please send an email to amandlekar@nvidia.com. Thanks! +Some helpful suggestions on useful documentation pages to view next: -## Troubleshooting and Known Issues +- [Getting Started](https://mimicgen.github.io/docs/tutorials/getting_started.html) +- [Launching Several Data Generation Runs](https://mimicgen.github.io/docs/tutorials/launching_several.html) +- [Reproducing Published Experiments and Results](https://mimicgen.github.io/docs/tutorials/reproducing_experiments.html) +- [Data Generation for Custom Environments](https://mimicgen.github.io/docs/tutorials/datagen_custom.html) +- [Overview of MimicGen Codebase](https://mimicgen.github.io/docs/modules/overview.html) -- If your robomimic training seems to be proceeding slowly (especially for image-based agents), it might be a problem with robomimic and more modern versions of PyTorch. We recommend PyTorch 1.12.1 (on Ubuntu, we used `conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch`). It is also a good idea to verify that the GPU is being utilized during training. -- In our testing on M1 macbook we ran into the following error when using `imageio-ffmpeg` installed through pip: `RuntimeError: No ffmpeg exe could be found. Install ffmpeg on your system, or set the IMAGEIO_FFMPEG_EXE environment variable.` Using `conda install imageio-ffmpeg` fixed this issue on our end. -- If you run into trouble with installing [egl_probe](https://github.com/StanfordVL/egl_probe) during robomimic installation (e.g. `ERROR: Failed building wheel for egl_probe`) you may need to make sure `cmake` is installed. A simple `pip install cmake` should work. -- If you run into other strange installation issues, one potential fix is to launch a new terminal, activate your conda environment, and try the install commands that are failing once again. One clue that the current terminal state is corrupt and this fix will help is if you see installations going into a different conda environment than the one you have active. -- If you run into rendering issues with the Sawyer robot arm, or have trouble reproducing our results, your MuJoCo version might be the issue. As noted in the [Installation](#installation) section, please use MuJoCo 2.3.2 (`pip install mujoco==2.3.2`). +## Troubleshooting -If you run into an error not documented above, please search through the [GitHub issues](https://github.com/NVlabs/mimicgen_environments/issues), and create a new one if you cannot find a fix. +Please see the [troubleshooting](https://mimicgen.github.io/docs/miscellaneous/troubleshooting.html) section for common fixes, or submit an issue on our github page. ## License -The code is released under the [NVIDIA Source Code License](https://github.com/NVlabs/mimicgen_environments/blob/main/LICENSE) and the datasets are released under [CC-BY 4.0](https://creativecommons.org/licenses/by/4.0/). +The code is released under the [NVIDIA Source Code License](https://github.com/NVlabs/mimicgen/blob/main/LICENSE) and the datasets are released under [CC-BY 4.0](https://creativecommons.org/licenses/by/4.0/). ## Citation diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 0000000..8508a32 --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,23 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = . +BUILDDIR = _build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +apidoc: + @sphinx-apidoc -T --force ../mimicgen -o api + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/_static/theme_overrides.css b/docs/_static/theme_overrides.css new file mode 100644 index 0000000..63ee6cc --- /dev/null +++ b/docs/_static/theme_overrides.css @@ -0,0 +1,13 @@ +/* override table width restrictions */ +@media screen and (min-width: 767px) { + + .wy-table-responsive table td { + /* !important prevents the common CSS stylesheets from overriding + this as on RTD they are loaded after this stylesheet */ + white-space: normal !important; + } + + .wy-table-responsive { + overflow: visible !important; + } +} diff --git a/docs/api/mimicgen.configs.rst b/docs/api/mimicgen.configs.rst new file mode 100644 index 0000000..631bfcc --- /dev/null +++ b/docs/api/mimicgen.configs.rst @@ -0,0 +1,37 @@ +mimicgen.configs package +======================== + +Submodules +---------- + +mimicgen.configs.config module +------------------------------ + +.. automodule:: mimicgen.configs.config + :members: + :undoc-members: + :show-inheritance: + +mimicgen.configs.robosuite module +--------------------------------- + +.. automodule:: mimicgen.configs.robosuite + :members: + :undoc-members: + :show-inheritance: + +mimicgen.configs.task\_spec module +---------------------------------- + +.. automodule:: mimicgen.configs.task_spec + :members: + :undoc-members: + :show-inheritance: + +Module contents +--------------- + +.. automodule:: mimicgen.configs + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/mimicgen.datagen.rst b/docs/api/mimicgen.datagen.rst new file mode 100644 index 0000000..6d58d07 --- /dev/null +++ b/docs/api/mimicgen.datagen.rst @@ -0,0 +1,45 @@ +mimicgen.datagen package +======================== + +Submodules +---------- + +mimicgen.datagen.data\_generator module +--------------------------------------- + +.. automodule:: mimicgen.datagen.data_generator + :members: + :undoc-members: + :show-inheritance: + +mimicgen.datagen.datagen\_info module +------------------------------------- + +.. automodule:: mimicgen.datagen.datagen_info + :members: + :undoc-members: + :show-inheritance: + +mimicgen.datagen.selection\_strategy module +------------------------------------------- + +.. automodule:: mimicgen.datagen.selection_strategy + :members: + :undoc-members: + :show-inheritance: + +mimicgen.datagen.waypoint module +-------------------------------- + +.. automodule:: mimicgen.datagen.waypoint + :members: + :undoc-members: + :show-inheritance: + +Module contents +--------------- + +.. automodule:: mimicgen.datagen + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/mimicgen.env_interfaces.rst b/docs/api/mimicgen.env_interfaces.rst new file mode 100644 index 0000000..3a153df --- /dev/null +++ b/docs/api/mimicgen.env_interfaces.rst @@ -0,0 +1,29 @@ +mimicgen.env\_interfaces package +================================ + +Submodules +---------- + +mimicgen.env\_interfaces.base module +------------------------------------ + +.. automodule:: mimicgen.env_interfaces.base + :members: + :undoc-members: + :show-inheritance: + +mimicgen.env\_interfaces.robosuite module +----------------------------------------- + +.. automodule:: mimicgen.env_interfaces.robosuite + :members: + :undoc-members: + :show-inheritance: + +Module contents +--------------- + +.. automodule:: mimicgen.env_interfaces + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/mimicgen.envs.robosuite.rst b/docs/api/mimicgen.envs.robosuite.rst new file mode 100644 index 0000000..8104943 --- /dev/null +++ b/docs/api/mimicgen.envs.robosuite.rst @@ -0,0 +1,93 @@ +mimicgen.envs.robosuite package +=============================== + +Submodules +---------- + +mimicgen.envs.robosuite.coffee module +------------------------------------- + +.. automodule:: mimicgen.envs.robosuite.coffee + :members: + :undoc-members: + :show-inheritance: + +mimicgen.envs.robosuite.hammer\_cleanup module +---------------------------------------------- + +.. automodule:: mimicgen.envs.robosuite.hammer_cleanup + :members: + :undoc-members: + :show-inheritance: + +mimicgen.envs.robosuite.kitchen module +-------------------------------------- + +.. automodule:: mimicgen.envs.robosuite.kitchen + :members: + :undoc-members: + :show-inheritance: + +mimicgen.envs.robosuite.mug\_cleanup module +------------------------------------------- + +.. automodule:: mimicgen.envs.robosuite.mug_cleanup + :members: + :undoc-members: + :show-inheritance: + +mimicgen.envs.robosuite.nut\_assembly module +-------------------------------------------- + +.. automodule:: mimicgen.envs.robosuite.nut_assembly + :members: + :undoc-members: + :show-inheritance: + +mimicgen.envs.robosuite.pick\_place module +------------------------------------------ + +.. automodule:: mimicgen.envs.robosuite.pick_place + :members: + :undoc-members: + :show-inheritance: + +mimicgen.envs.robosuite.single\_arm\_env\_mg module +--------------------------------------------------- + +.. automodule:: mimicgen.envs.robosuite.single_arm_env_mg + :members: + :undoc-members: + :show-inheritance: + +mimicgen.envs.robosuite.stack module +------------------------------------ + +.. automodule:: mimicgen.envs.robosuite.stack + :members: + :undoc-members: + :show-inheritance: + +mimicgen.envs.robosuite.threading module +---------------------------------------- + +.. automodule:: mimicgen.envs.robosuite.threading + :members: + :undoc-members: + :show-inheritance: + +mimicgen.envs.robosuite.three\_piece\_assembly module +----------------------------------------------------- + +.. automodule:: mimicgen.envs.robosuite.three_piece_assembly + :members: + :undoc-members: + :show-inheritance: + +Module contents +--------------- + +.. automodule:: mimicgen.envs.robosuite + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/mimicgen.envs.rst b/docs/api/mimicgen.envs.rst new file mode 100644 index 0000000..e7e805a --- /dev/null +++ b/docs/api/mimicgen.envs.rst @@ -0,0 +1,18 @@ +mimicgen.envs package +===================== + +Subpackages +----------- + +.. toctree:: + :maxdepth: 4 + + mimicgen.envs.robosuite + +Module contents +--------------- + +.. automodule:: mimicgen.envs + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/mimicgen.rst b/docs/api/mimicgen.rst new file mode 100644 index 0000000..5bce266 --- /dev/null +++ b/docs/api/mimicgen.rst @@ -0,0 +1,23 @@ +mimicgen package +================ + +Subpackages +----------- + +.. toctree:: + :maxdepth: 4 + + mimicgen.configs + mimicgen.datagen + mimicgen.env_interfaces + mimicgen.envs + mimicgen.scripts + mimicgen.utils + +Module contents +--------------- + +.. automodule:: mimicgen + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/mimicgen.scripts.rst b/docs/api/mimicgen.scripts.rst new file mode 100644 index 0000000..244657e --- /dev/null +++ b/docs/api/mimicgen.scripts.rst @@ -0,0 +1,133 @@ +mimicgen.scripts package +======================== + +Submodules +---------- + +mimicgen.scripts.annotate\_subtasks module +------------------------------------------ + +.. automodule:: mimicgen.scripts.annotate_subtasks + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.demo\_random\_action module +-------------------------------------------- + +.. automodule:: mimicgen.scripts.demo_random_action + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.download\_datasets module +------------------------------------------ + +.. automodule:: mimicgen.scripts.download_datasets + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.generate\_config\_templates module +--------------------------------------------------- + +.. automodule:: mimicgen.scripts.generate_config_templates + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.generate\_core\_configs module +----------------------------------------------- + +.. automodule:: mimicgen.scripts.generate_core_configs + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.generate\_core\_training\_configs module +--------------------------------------------------------- + +.. automodule:: mimicgen.scripts.generate_core_training_configs + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.generate\_dataset module +----------------------------------------- + +.. automodule:: mimicgen.scripts.generate_dataset + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.generate\_robot\_transfer\_configs module +---------------------------------------------------------- + +.. automodule:: mimicgen.scripts.generate_robot_transfer_configs + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.generate\_robot\_transfer\_training\_configs module +-------------------------------------------------------------------- + +.. automodule:: mimicgen.scripts.generate_robot_transfer_training_configs + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.generate\_training\_configs\_for\_public\_datasets module +-------------------------------------------------------------------------- + +.. automodule:: mimicgen.scripts.generate_training_configs_for_public_datasets + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.get\_reset\_videos module +------------------------------------------ + +.. automodule:: mimicgen.scripts.get_reset_videos + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.get\_source\_info module +----------------------------------------- + +.. automodule:: mimicgen.scripts.get_source_info + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.merge\_hdf5 module +----------------------------------- + +.. automodule:: mimicgen.scripts.merge_hdf5 + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.prepare\_src\_dataset module +--------------------------------------------- + +.. automodule:: mimicgen.scripts.prepare_src_dataset + :members: + :undoc-members: + :show-inheritance: + +mimicgen.scripts.visualize\_subtasks module +------------------------------------------- + +.. automodule:: mimicgen.scripts.visualize_subtasks + :members: + :undoc-members: + :show-inheritance: + +Module contents +--------------- + +.. automodule:: mimicgen.scripts + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/mimicgen.utils.rst b/docs/api/mimicgen.utils.rst new file mode 100644 index 0000000..93c064a --- /dev/null +++ b/docs/api/mimicgen.utils.rst @@ -0,0 +1,53 @@ +mimicgen.utils package +====================== + +Submodules +---------- + +mimicgen.utils.config\_utils module +----------------------------------- + +.. automodule:: mimicgen.utils.config_utils + :members: + :undoc-members: + :show-inheritance: + +mimicgen.utils.file\_utils module +--------------------------------- + +.. automodule:: mimicgen.utils.file_utils + :members: + :undoc-members: + :show-inheritance: + +mimicgen.utils.misc\_utils module +--------------------------------- + +.. automodule:: mimicgen.utils.misc_utils + :members: + :undoc-members: + :show-inheritance: + +mimicgen.utils.pose\_utils module +--------------------------------- + +.. automodule:: mimicgen.utils.pose_utils + :members: + :undoc-members: + :show-inheritance: + +mimicgen.utils.robomimic\_utils module +-------------------------------------- + +.. automodule:: mimicgen.utils.robomimic_utils + :members: + :undoc-members: + :show-inheritance: + +Module contents +--------------- + +.. automodule:: mimicgen.utils + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 0000000..74c2c52 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,177 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +# Configuration file for the Sphinx documentation builder. +# +# For the full list of built-in configuration values, see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +# -- Path setup -------------------------------------------------------------- + +# If extensions (or modules to document with autodoc) are in another directory, +# add these directories to sys.path here. If the directory is relative to the +# documentation root, use os.path.abspath to make it absolute, like shown here. + +import os +import sys +sys.path.insert(0, os.path.abspath('.')) + +import sphinx_book_theme +import mimicgen + + +# -- General configuration --------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration + +# If your documentation needs a minimal Sphinx version, state it here. +# +# needs_sphinx = '1.0' + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom +# ones. +extensions = [ + 'sphinx.ext.napoleon', + # 'sphinx_markdown_tables', + 'sphinx.ext.mathjax', + 'sphinx.ext.githubpages', + 'sphinx.ext.autodoc', + 'recommonmark', # use Sphinx-1.4 or newer + # 'nbsphinx', +] + + +# Sphinx-apidoc variables +apidoc_module_dir = "../mimicgen" +apidoc_output_dir = "reference" + + +# Add any paths that contain templates here, relative to this directory. +templates_path = ['_templates'] + +# source_parsers = { +# '.md': CommonMarkParser, +# } + +# The suffix(es) of source filenames. +# You can specify multiple suffix as a list of string: +source_suffix = ['.rst', '.md', '.ipynb'] + +# The master toctree document. +master_doc = 'index' + +# General information about the project. +project = 'mimicgen' +copyright = '2023-2024, NVIDIA' +author = 'Ajay Mandlekar' + +# The version info for the project you're documenting, acts as replacement for +# |version| and |release|, also used in various other places throughout the +# built documents. +# +# The short X.Y version. +version = (".").join(mimicgen.__version__.split(".")[:-1]) +# The full version, including alpha/beta/rc tags. +release = (".").join(mimicgen.__version__.split(".")[:-1]) + + +# The language for content autogenerated by Sphinx. Refer to documentation +# for a list of supported languages. +# +# This is also used if you do content translation via gettext catalogs. +# Usually you set "language" from the command line for these cases. +language = "en" + + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +# This pattern also affects html_static_path and html_extra_path. +exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] + +# The name of the Pygments (syntax highlighting) style to use. +pygments_style = 'sphinx' + +# If true, `todo` and `todoList` produce output, else they produce nothing. +todo_include_todos = False + + +# -- Options for HTML output ------------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +# +html_theme = 'sphinx_book_theme' +html_logo = "mimicgen_logo.png" + +# Theme options are theme-specific and customize the look and feel of a theme +# further. For a list of options available for each theme, see the +# documentation. +# +# html_theme_options = {} + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = ['_static'] + +# html_context = { +# 'css_files': [ +# '_static/theme_overrides.css', # override wide tables in RTD theme +# ], +# } + +# -- Options for HTMLHelp output ------------------------------------------ + +# Output file base name for HTML help builder. +htmlhelp_basename = 'mimicgendoc' + + +# -- Options for LaTeX output --------------------------------------------- + +latex_elements = { + # The paper size ('letterpaper' or 'a4paper'). + # + # 'papersize': 'letterpaper', + + # The font size ('10pt', '11pt' or '12pt'). + # + # 'pointsize': '10pt', + + # Additional stuff for the LaTeX preamble. + # + # 'preamble': '', + + # Latex figure (float) alignment + # + # 'figure_align': 'htbp', +} + +# Grouping the document tree into LaTeX files. List of tuples +# (source start file, target name, title, +# author, documentclass [howto, manual, or own class]). +latex_documents = [ + (master_doc, 'mimicgen.tex', u'mimicgen Documentation', author, 'manual'), +] + + +# -- Options for manual page output --------------------------------------- + +# One entry per manual page. List of tuples +# (source start file, name, description, authors, manual section). +man_pages = [ + (master_doc, 'mimicgen', u'mimicgen Documentation', + [author], 1) +] + + +# -- Options for Texinfo output ------------------------------------------- + +# Grouping the document tree into Texinfo files. List of tuples +# (source start file, target name, title, author, +# dir menu entry, description, category) +texinfo_documents = [ + (master_doc, 'mimicgen', u'mimicgen Documentation', + author, 'mimicgen', 'mimicgen', + 'Miscellaneous'), +] diff --git a/docs/datasets/mimicgen_corl_2023.md b/docs/datasets/mimicgen_corl_2023.md new file mode 100644 index 0000000..37db304 --- /dev/null +++ b/docs/datasets/mimicgen_corl_2023.md @@ -0,0 +1,104 @@ +# MimicGen (CoRL 2023) + +In this section, we provide an overview of the datasets we released with the MimicGen paper. These are the same exact datasets used in our paper (but postprocessed to support a higher version of robosuite). We also show how to easily reproduce policy learning results on this data. + +
+

Note

+ +If you would like to reproduce our data generation results in addition to the policy learning results, please see the [Reproducing Experiments](https://mimicgen.github.io/docs/tutorials/reproducing_experiments.html) tutorial for a comprehensive guide. + +
+ +## Downloading and Using Datasets + +### Dataset Types + +As described in the paper, each task has a default reset distribution (D_0). Source human demonstrations (usually 10 demos) were collected on this distribution and MimicGen was subsequently used to generate large datasets (usually 1000 demos) across different task reset distributions (e.g. D_0, D_1, D_2), objects, and robots. + +The datasets are split into different types: + +- **source**: source human datasets used to generate all data -- this generally consists of 10 human demonstrations collected on the D_0 variant for each task. +- **core**: datasets generated with MimicGen for different task reset distributions. These correspond to the core set of results in Figure 4 of the paper. +- **object**: datasets generated with MimicGen for different objects. These correspond to the results in Appendix G of the paper. +- **robot**: datasets generated with MimicGen for different robots. These correspond to the results in Appendix F of the paper. +- **large_interpolation**: datasets generated with MimicGen using much larger interpolation segments. These correspond to the results in Appendix H in the paper. + +
+

Note

+ +We found that the large_interpolation datasets pose a significant challenge for imitation learning, and have substantial room for improvement. + +
+ +### Dataset Statistics + +The datasets contain over 48,000 task demonstrations across 12 tasks. + +We provide more information on the amount of demonstrations for each dataset type: +- **source**: 120 human demonstrations across 12 tasks (10 per task) used to automatically generate the other datasets +- **core**: 26,000 task demonstrations across 12 tasks (26 task variants) +- **object**: 2000 task demonstrations on the Mug Cleanup task with different mugs +- **robot**: 16,000 task demonstrations across 4 different robot arms on 2 tasks (4 task variants) +- **large_interpolation**: 6000 task demonstrations across 6 tasks that pose significant challenges for modern imitation learning methods + +### Dataset Download + +#### Method 1: Using `download_datasets.py` (Recommended) + +`download_datasets.py` (located at `mimicgen/scripts`) is a python script that provides a programmatic way of downloading the datasets. This is the preferred method, because this script also sets up a directory structure for the datasets that works out of the box with the code for reproducing policy learning results. + +A few examples of using this script are provided below: + +```sh +# default behavior - just download core square_d0 dataset +python download_datasets.py + +# download core datasets for square D0, D1, D2 and coffee D0, D1, D2 +python download_datasets.py --dataset_type core --tasks square_d0 square_d1 square_d2 coffee_d0 coffee_d1 coffee_d2 + +# download all core datasets, but do a dry run first to see what will be downloaded and where +python download_datasets.py --dataset_type core --tasks all --dry_run + +# download all source human datasets +python download_datasets.py --dataset_type source --tasks all +``` + +#### Method 2: Using Direct Download Links + +You can download the datasets manually through Google Drive. The folders each correspond to the dataset types described in [Dataset Types](#dataset-types). + +**Google Drive folder with all datasets:** [link](https://drive.google.com/drive/folders/14e9kkHGfApuQ709LBEbXrXVI1Lp5Ax7p?usp=drive_link) + +#### Method 3: Using Hugging Face + +You can download the datasets through Hugging Face. + +**Hugging Face dataset repository:** [link](https://huggingface.co/datasets/amandlek/mimicgen_datasets) + + +## Reproducing Policy Learning Results + +After downloading the appropriate datasets you’re interested in using by running the `download_datasets.py` script, the `generate_training_configs_for_public_datasets.py` script (located at `mimicgen/scripts`) can be used to generate all training config json files necessary to reproduce the experiments in the paper. A few examples are below. + +```sh +# Assume datasets already exist in mimicgen/../datasets folder. Configs will be generated under mimicgen/exps/paper, and training results will be at mimicgen/../training_results after launching training runs. +python generate_training_configs_for_public_datasets.py + +# Alternatively, specify where datasets exist, and specify where configs should be generated. +python generate_training_configs_for_public_datasets.py --config_dir /tmp/configs --dataset_dir /tmp/datasets --output_dir /tmp/experiment_results +``` + +Then, to reproduce a specific set of training runs for different experiment groups (see [Dataset Types](#dataset-types)), we can simply navigate to the generated config directory, and copy training commands from the generated shell script there. As an example, we can reproduce the image training results on the Coffee D0 dataset, by looking for the correct set of commands in `mimicgen/exps/paper/core.sh` and running them. The relevant section of the shell script is reproduced below. + +```sh +# task: coffee_d0 +# obs modality: image +python /path/to/robomimic/scripts/train.py --config /path/to/mimicgen/exps/paper/core/coffee_d0/image/bc_rnn.json +``` + +
+

Note

+ +In the MimicGen paper, we generated our datasets on versions of environments built on robosuite `v1.2`. Since then, we changed the environments and datasets (through postprocessing) to be based on robosuite `v1.4`. However, `v1.4` has some visual and dynamics differences from `v1.2`, so the learning results may not exactly match up with the ones we reported in the paper. In our testing on these released datasets, we were able to reproduce nearly all of our results, but within 10% of the performance reported in the paper. + +
diff --git a/assets/mimicgen.gif b/docs/images/mimicgen.gif similarity index 100% rename from assets/mimicgen.gif rename to docs/images/mimicgen.gif diff --git a/assets/mosaic.gif b/docs/images/mosaic.gif similarity index 100% rename from assets/mosaic.gif rename to docs/images/mosaic.gif diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 index 0000000..23af397 --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,62 @@ +.. mimicgen documentation master file, created by + sphinx-quickstart on Thu Jul 4 12:05:47 2024. + You can adapt this file completely to your liking, but it should at least + contain the root `toctree` directive. + +Welcome to mimicgen's documentation! +==================================== + +.. toctree:: + :maxdepth: 1 + :caption: Introduction + + introduction/overview + introduction/installation + +.. toctree:: + :maxdepth: 1 + :caption: Datasets + + datasets/mimicgen_corl_2023 + +.. toctree:: + :maxdepth: 1 + :caption: Tutorials + + tutorials/getting_started + tutorials/reproducing_experiments + tutorials/launching_several + tutorials/datagen_custom + tutorials/task_visualizations + tutorials/subtask_termination_signals + tutorials/debugging_datagen + +.. toctree:: + :maxdepth: 1 + :caption: Modules + + modules/overview + modules/datagen + modules/task_spec + modules/env_interfaces + +.. toctree:: + :maxdepth: 1 + :caption: Miscellaneous + + miscellaneous/troubleshooting + +.. toctree:: + :maxdepth: 4 + :caption: Source API + + api/mimicgen + + + +Indices and tables +================== + +* :ref:`genindex` +* :ref:`modindex` +* :ref:`search` diff --git a/docs/introduction/installation.md b/docs/introduction/installation.md new file mode 100644 index 0000000..4601cea --- /dev/null +++ b/docs/introduction/installation.md @@ -0,0 +1,157 @@ +# Installation + +## Requirements + +- Mac OS X or Linux machine +- Python >= 3.6 (recommended 3.8.0) +- [conda](https://www.anaconda.com/products/individual) + - [virtualenv](https://virtualenv.pypa.io/en/latest/) is also an acceptable alternative, but we assume you have conda installed in our examples below + +## Install MimicGen + +We recommend installing the repo into a new conda environment (it is called `mimicgen` in the example below): + +```sh +conda create -n mimicgen python=3.8 +conda activate mimicgen +``` + +You can install most of the dependencies by cloning the repository and then installing from source: + +```sh +cd +git clone https://github.com/NVlabs/mimicgen.git +cd mimicgen +pip install -e . +``` + +However, there are some additional dependencies that we list below. + +### Additional Dependencies + +Most of these additional dependencies are best installed from source. + +#### robosuite + +
+

Note

+ +[robosuite](https://robosuite.ai/) is an optional dependency that is only needed if running the examples provided with this repository. The MimicGen source code does not rely on robosuite and can be used with other simulation frameworks. + +
+ +```sh +$ cd +$ git clone https://github.com/ARISE-Initiative/robosuite.git +$ cd robosuite +$ git checkout b9d8d3de5e3dfd1724f4a0e6555246c460407daa +$ pip install -e . +``` + +For more detailed instructions, see the [robosuite installation page](https://robosuite.ai/docs/installation.html). + +
+

Note

+ +The git checkout command corresponds to the commit we used for testing our policy learning results. In general the `master` branch (`v1.4+`) should be fine. + +
+ +#### robomimic + +[robomimic](https://robomimic.github.io/) is a required dependency that provides a standardized dataset format, wrappers around simulation environments, and policy learning utilities. + +```sh +$ cd +$ git clone https://github.com/ARISE-Initiative/robomimic.git +$ cd robomimic +$ git checkout d0b37cf214bd24fb590d182edb6384333f67b661 +$ pip install -e . +``` + +For more detailed instructions, see the [robomimic installation page](https://robomimic.github.io/docs/introduction/installation.html). + +
+

Note

+ +The git checkout command corresponds to the commit we used for testing our policy learning results. In general the `master` branch (`v0.3+`) should be fine, as long as it is after the above commit. + +
+ +#### robosuite_task_zoo + +
+

Note

+ +[robosuite_task_zoo](https://github.com/ARISE-Initiative/robosuite-task-zoo) is an optional dependency that is only needed if running the Kitchen and Hammer Cleanup environments and datasets provided with this repository. + +
+ +
+

Warning

+ +We recommend removing the dependencies in the `setup.py` file (the `install_requires` list) before installation, as it uses deprecated dependencies (such as mujoco-py). + +
+ +```sh +$ cd +$ git clone https://github.com/ARISE-Initiative/robosuite-task-zoo +$ cd robosuite-task-zoo +$ git checkout 74eab7f88214c21ca1ae8617c2b2f8d19718a9ed +# NOTE: should remove dependencies in setup.py in the "install_requires" before the last step +$ pip install -e . +``` + +#### mujoco + +If using robosuite, **please downgrade MuJoCo to 2.3.2**: + +```sh +$ pip install mujoco==2.3.2 +``` + +
+

Warning

+ +This MuJoCo version (`2.3.2`) can be important -- in our testing, we found that other versions of MuJoCo could be problematic, especially for the Sawyer arm datasets (e.g. `2.3.5` causes problems with rendering and `2.3.7` changes the dynamics of the robot arm significantly from the collected datasets). More modern versions of MuJoCo (e.g. `3.0`+) might be fine. + +
+ +#### pygame + +If you plan on using our subtask annotation interface (`scripts/annotate_subtasks.py`) you should also install pygame with `pip install pygame`. See the [Subtask Termination Signals](https://mimicgen.github.io/docs/tutorials/subtask_termination_signals.html) page for more information. + +## Test Your Installation + +The following script can be used to try random actions in one of our custom robosuite tasks. +```sh +$ cd mimicgen/scripts +$ python demo_random_action.py +``` + +
+

Note

+ +To test data generation please move on to the [Getting Started](https://mimicgen.github.io/docs/tutorials/getting_started.html) tutorial. + +
+ +## Install Documentation Dependencies + +If you plan to contribute to the repository and add new features, you must install the additional requirements required to build the documentation locally: + +```sh +$ pip install -r requirements-docs.txt +``` + +You can test generating the documentation and viewing it locally in a web browser: +```sh +$ cd /docs +$ make clean +$ make apidoc +$ make html +$ cp -r images _build/html/ +``` + +There should be a generated `_build` folder - navigate to `_build/html/` and open `index.html` in a web browser to view the documentation. diff --git a/docs/introduction/overview.md b/docs/introduction/overview.md new file mode 100644 index 0000000..f8fc9df --- /dev/null +++ b/docs/introduction/overview.md @@ -0,0 +1,51 @@ +# Overview + +

+ +

+ +This repository contains the official release of data generation code, simulation environments, and datasets for the [CoRL 2023](https://www.corl2023.org/) paper "MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations". + +Website: [https://mimicgen.github.io](https://mimicgen.github.io) + +Paper: [https://arxiv.org/abs/2310.17596](https://arxiv.org/abs/2310.17596) + +Documentation: [https://mimicgen.github.io/docs/introduction/overview.html](https://mimicgen.github.io/docs/introduction/overview.html) + +
+

Note

+ +For business inquiries, please submit this form: [NVIDIA Research Licensing](https://www.nvidia.com/en-us/research/inquiries/). + +
+ +## Useful Documentation Links + +Some helpful suggestions on useful documentation pages to view next: + +- [Getting Started](https://mimicgen.github.io/docs/tutorials/getting_started.html) +- [Launching Several Data Generation Runs](https://mimicgen.github.io/docs/tutorials/launching_several.html) +- [Reproducing Published Experiments and Results](https://mimicgen.github.io/docs/tutorials/reproducing_experiments.html) +- [Data Generation for Custom Environments](https://mimicgen.github.io/docs/tutorials/datagen_custom.html) +- [Overview of MimicGen Codebase](https://mimicgen.github.io/docs/modules/overview.html) + +## Troubleshooting + +Please see the [troubleshooting](https://mimicgen.github.io/docs/miscellaneous/troubleshooting.html) section for common fixes, or submit an issue on our github page. + +## License + +The code is released under the [NVIDIA Source Code License](https://github.com/NVlabs/mimicgen/blob/main/LICENSE) and the datasets are released under [CC-BY 4.0](https://creativecommons.org/licenses/by/4.0/). + +## Citation + +Please cite [the MimicGen paper](https://arxiv.org/abs/2310.17596) if you use this code in your work: + +```bibtex +@inproceedings{mandlekar2023mimicgen, + title={MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations}, + author={Mandlekar, Ajay and Nasiriany, Soroush and Wen, Bowen and Akinola, Iretiayo and Narang, Yashraj and Fan, Linxi and Zhu, Yuke and Fox, Dieter}, + booktitle={7th Annual Conference on Robot Learning}, + year={2023} +} +``` \ No newline at end of file diff --git a/docs/mimicgen_logo.png b/docs/mimicgen_logo.png new file mode 100644 index 0000000..5d035c4 Binary files /dev/null and b/docs/mimicgen_logo.png differ diff --git a/docs/miscellaneous/troubleshooting.md b/docs/miscellaneous/troubleshooting.md new file mode 100644 index 0000000..475e7a4 --- /dev/null +++ b/docs/miscellaneous/troubleshooting.md @@ -0,0 +1,14 @@ +# Troubleshooting and Known Issues + +If you run into an error not documented below, please search through the [GitHub issues](https://github.com/NVlabs/mimicgen/issues), and create a new one if you cannot find a fix. + +## Installation + +- In our testing on M1 macbook we ran into the following error: `RuntimeError: No ffmpeg exe could be found. Install ffmpeg on your system, or set the IMAGEIO_FFMPEG_EXE environment variable.` Using `conda install ffmpeg` fixed this issue on our end. +- If you run into trouble with installing [egl_probe](https://github.com/StanfordVL/egl_probe) during robomimic installation (e.g. `ERROR: Failed building wheel for egl_probe`) you may need to make sure `cmake` is installed. A simple `pip install cmake` should work. +- If you run into other strange installation issues, one potential fix is to launch a new terminal, activate your conda environment, and try the install commands that are failing once again. One clue that the current terminal state is corrupt and this fix will help is if you see installations going into a different conda environment than the one you have active. + +## Policy Learning + +- If your robomimic training seems to be proceeding slowly (especially for image-based agents), it might be a problem with robomimic and more modern versions of PyTorch. We recommend PyTorch 1.12.1 (on Ubuntu, we used `conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch`). It is also a good idea to verify that the GPU is being utilized during training. +- If you run into rendering issues with the Sawyer robot arm, or have trouble reproducing our results, your MuJoCo version might be the issue. As noted in the [Installation](https://mimicgen.github.io/docs/introduction/installation.html) section, please use MuJoCo 2.3.2 (`pip install mujoco==2.3.2`). diff --git a/docs/modules/datagen.md b/docs/modules/datagen.md new file mode 100644 index 0000000..9931bd1 --- /dev/null +++ b/docs/modules/datagen.md @@ -0,0 +1,316 @@ +# Data Generation Classes + +This section discusses the key classes related to data generation. + +## Data Generator + +The `DataGenerator` class (`datagen/data_generator.py`) is responsible for generating new demonstration trajectories. First, the internal `_load_dataset` method is used to parse the source dataset (using [Subtask Termination Signals](https://mimicgen.github.io/docs/tutorials/subtask_termination_signals.html)) into source subtask segments. Each segment is a sequence of [DatagenInfo](https://mimicgen.github.io/docs/modules/datagen.html#datagen-info) objects. Then, the `generate` method is called repeatedly (by the main script `scripts/generate_dataset.py`) to keep attempting to generate new trajectories using the source subtask segments. During each new attempt, the `select_source_demo` method is used to employ a [SelectionStrategy](https://mimicgen.github.io/docs/modules/datagen.html#selection-strategy) to pick a reference source subtask segment to transform. [WaypointTrajectory](https://mimicgen.github.io/docs/modules/datagen.html#waypoint) objects are used to transform and compose subtask segments together. + +## Datagen Info + +DatagenInfo objects keep track of important information used during data generation. These objects are added to source demonstrations with the `prepare_src_dataset.py` script, or provided directly by an [Environment Interface](https://mimicgen.github.io/docs/modules/env_interfaces.html) object. + +The structure of the object is below: + +```python +class DatagenInfo(object): + """ + Structure of information needed from an environment for data generation. To allow for + flexibility, not all information must be present. + """ + def __init__( + self, + eef_pose=None, + object_poses=None, + subtask_term_signals=None, + target_pose=None, + gripper_action=None, + ): + """ + Args: + eef_pose (np.array or None): robot end effector poses of shape [..., 4, 4] + object_poses (dict or None): dictionary mapping object name to object poses + of shape [..., 4, 4] + subtask_term_signals (dict or None): dictionary mapping subtask name to a binary + indicator (0 or 1) on whether subtask has been completed. Each value in the + dictionary could be an int, float, or np.array of shape [..., 1]. + target_pose (np.array or None): target end effector poses of shape [..., 4, 4] + gripper_action (np.array or None): gripper actions of shape [..., D] where D + is the dimension of the gripper actuation action for the robot arm + """ +``` + +## Selection Strategy + +
+

Note

+ +See Appendix N.3 in the MimicGen paper for a more thorough explanation of source subtask segment selection and some further intuition on when to use different settings. + +
+ +Each data generation attempt requires choosing one or more subtask segments from the source demonstrations to transform -- this is carried out by a SelectionStrategy instance: + +```python +@six.add_metaclass(MG_SelectionStrategyMeta) +class MG_SelectionStrategy(object): + """ + Defines methods and functions for selection strategies to implement. + """ + def __init__(self): + pass + + @property + @classmethod + def NAME(self): + """ + This name (str) will be used to register the selection strategy class in the global + registry. + """ + raise NotImplementedError + + @abc.abstractmethod + def select_source_demo( + self, + eef_pose, + object_pose, + src_subtask_datagen_infos, + ): + """ + Selects source demonstration index using the current robot pose, relevant object pose + for the current subtask, and relevant information from the source demonstrations for the + current subtask. + + Args: + eef_pose (np.array): current 4x4 eef pose + object_pose (np.array): current 4x4 object pose, for the object in this subtask + src_subtask_datagen_infos (list): DatagenInfo instance for the relevant subtask segment + in the source demonstrations + + Returns: + source_demo_ind (int): index of source demonstration - indicates which source subtask segment to use + """ + raise NotImplementedError +``` + +Every SelectionStrategy class must subclass this base class and implement the `NAME` and `select_source_demo` methods. The `NAME` field is used to register the SelectionStrategy class into the global registry, and `select_source_demo` implements the heuristic for selecting a source demonstration index. + +Each data generation config json specifies how source segment selection should be done during data generation. First, `config.experiment.generation.select_src_per_subtask` determines whether to select a different source demonstration for each subtask during data generation, or keep the same source demonstration as the one used for the first subtask. This corresponds to the `per-subtask` parameter described in the "Selection Frequency" paragraph of Appendix N.3 in the paper. + +The specific task config (`config.task.task_spec`), which corresponds to the [Task Spec](https://mimicgen.github.io/docs/modules/task_spec.html) object used in data generation, also specifies the selection strategy to use for each subtask via the `selection_strategy` parameter and the `selection_strategy_kwargs` parameter. The `selection_strategy` parameter corresponds to the `NAME` for the SelectionStrategy class, and the `selection_strategy_kwargs` correspond to any additional kwargs to specify when invoking the `select_source_demo` method. + +
+

Note

+ +Note that if `config.experiment.generation.select_src_per_subtask` is False, only the first subtask's selection strategy matters, since the selected source demonstration will be used for the remainder of the data generation attempt. + +
+ +As an example, the `NearestNeighborObjectStrategy` (see implementation below) can be specified by passing `nearest_neighbor_object` for the `selection_strategy` parameter and you can use the `selection_strategy_kwargs` parameter to specify a dictionary containing values for the `pos_weight`, `rot_weight`, and `nn_k` parameters. + +```python +class NearestNeighborObjectStrategy(MG_SelectionStrategy): + """ + Pick source demonstration to be the one with the closest object pose to the object + in the current scene. + """ + + # name for registering this class into registry + NAME = "nearest_neighbor_object" + + def select_source_demo( + self, + eef_pose, + object_pose, + src_subtask_datagen_infos, + pos_weight=1., + rot_weight=1., + nn_k=3, + ): + """ + Selects source demonstration index using the current robot pose, relevant object pose + for the current subtask, and relevant information from the source demonstrations for the + current subtask. + + Args: + eef_pose (np.array): current 4x4 eef pose + object_pose (np.array): current 4x4 object pose, for the object in this subtask + src_subtask_datagen_infos (list): DatagenInfo instance for the relevant subtask segment + in the source demonstrations + pos_weight (float): weight on position for minimizing pose distance + rot_weight (float): weight on rotation for minimizing pose distance + nn_k (int): pick source demo index uniformly at randomly from the top @nn_k nearest neighbors + + Returns: + source_demo_ind (int): index of source demonstration - indicates which source subtask segment to use + """ +``` + +## Waypoint + +### Waypoint Class Variants + +MimicGen uses a collection of convenience classes to represent waypoints and trajectories (`datagen/waypoint.py`). + +The `Waypoint` class represents a single 6-DoF target pose and the gripper action for that timestep: + +```python +class Waypoint(object): + """ + Represents a single desired 6-DoF waypoint, along with corresponding gripper actuation for this point. + """ + def __init__(self, pose, gripper_action, noise=None): +``` + +The `WaypointSequence` class represents a sequence of these `Waypoint` objects: + +```python +class WaypointSequence(object): + """ + Represents a sequence of Waypoint objects. + """ + def __init__(self, sequence=None): +``` + +It can easily be instantiated from a collection of poses (e.g. `WaypointSequence.from_poses`): + +```python + @classmethod + def from_poses(cls, poses, gripper_actions, action_noise): + """ + Instantiate a WaypointSequence object given a sequence of poses, + gripper actions, and action noise. + + Args: + poses (np.array): sequence of pose matrices of shape (T, 4, 4) + gripper_actions (np.array): sequence of gripper actions + that should be applied at each timestep of shape (T, D). + action_noise (float or np.array): sequence of action noise + magnitudes that should be applied at each timestep. If a + single float is provided, the noise magnitude will be + constant over the trajectory. + """ +``` + +Finally, the `WaypointTrajectory` class is a sequence of the `WaypointSequence` objects, and is a convenient way to represent 6-DoF trajectories and execute them: + +```python +class WaypointTrajectory(object): + """ + A sequence of WaypointSequence objects that corresponds to a full 6-DoF trajectory. + """ +``` + +`WaypointSequence` objects can be added directly to a `WaypointTrajectory` object: + +```python + def add_waypoint_sequence(self, sequence): + """ + Directly append sequence to list (no interpolation). + + Args: + sequence (WaypointSequence instance): sequence to add + """ +``` + +Interpolation segments can also be added easily using this helper method: + +```python + def add_waypoint_sequence_for_target_pose( + self, + pose, + gripper_action, + num_steps, + skip_interpolation=False, + action_noise=0., + ): + """ + Adds a new waypoint sequence corresponding to a desired target pose. A new WaypointSequence + will be constructed consisting of @num_steps intermediate Waypoint objects. These can either + be constructed with linear interpolation from the last waypoint (default) or be a + constant set of target poses (set @skip_interpolation to True). + + Args: + pose (np.array): 4x4 target pose + + gripper_action (np.array): value for gripper action + + num_steps (int): number of action steps when trying to reach this waypoint. Will + add intermediate linearly interpolated points between the last pose on this trajectory + and the target pose, so that the total number of steps is @num_steps. + + skip_interpolation (bool): if True, keep the target pose fixed and repeat it @num_steps + times instead of using linearly interpolated targets. + + action_noise (float): scale of random gaussian noise to add during action execution (e.g. + when @execute is called) + """ +``` + +The `merge` method is a thin wrapper around the above method, easily allowing for linear interpolation between two `WaypointTrajectory` objects: + +```python + def merge( + self, + other, + num_steps_interp=None, + num_steps_fixed=None, + action_noise=0., + ): + """ + Merge this trajectory with another (@other). + + Args: + other (WaypointTrajectory object): the other trajectory to merge into this one + + num_steps_interp (int or None): if not None, add a waypoint sequence that interpolates + between the end of the current trajectory and the start of @other + + num_steps_fixed (int or None): if not None, add a waypoint sequence that has constant + target poses corresponding to the first target pose in @other + + action_noise (float): noise to use during the interpolation segment + """ +``` + +Finally, the `execute` method makes it easy to execute the waypoint sequences in the simulation environment: + +```python + def execute( + self, + env, + env_interface, + render=False, + video_writer=None, + video_skip=5, + camera_names=None, + ): + """ + Main function to execute the trajectory. Will use env_interface.target_pose_to_action to + convert each target pose at each waypoint to an action command, and pass that along to + env.step. + + Args: + env (robomimic EnvBase instance): environment to use for executing trajectory + env_interface (MG_EnvInterface instance): environment interface for executing trajectory + render (bool): if True, render on-screen + video_writer (imageio writer): video writer + video_skip (int): determines rate at which environment frames are written to video + camera_names (list): determines which camera(s) are used for rendering. Pass more than + one to output a video with multiple camera views concatenated horizontally. + + Returns: + results (dict): dictionary with the following items for the executed trajectory: + states (list): simulator state at each timestep + observations (list): observation dictionary at each timestep + datagen_infos (list): datagen_info at each timestep + actions (list): action executed at each timestep + success (bool): whether the trajectory successfully solved the task or not + """ +``` + +### Waypoint Class Usage during Data Generation + +Each data generation attempt consists of executing a particular sequence of waypoints for each subtask. A `WaypointTrajectory` object is constructed and executed for each subtask in the `generate` method of the `DataGenerator`. + +For a given subtask, a `WaypointTrajectory` object is initialized with a single pose (usually the current robot end effector pose, or the last target pose from the previous subtask execution attempt). Next, a reference source subtask segment is selected, and then transformed using the `transform_source_data_segment_using_object_pose` method from `utils/pose_utils.py`. It is then merged into the trajectory object with linear interpolation using the `merge` method. Finally, the `execute` method is used to carry out the subtask. This process repeats for each subtask diff --git a/docs/modules/env_interfaces.md b/docs/modules/env_interfaces.md new file mode 100644 index 0000000..1c84247 --- /dev/null +++ b/docs/modules/env_interfaces.md @@ -0,0 +1,23 @@ +# Environment Interfaces + +Environment interface objects allow simulation environments to provide MimicGen with [DatagenInfo](https://mimicgen.github.io/docs/modules/datagen.html#datagen-info) class instances, which is a collection of information that MimicGen needs during data generation. + +These environment interface objects are used by `scripts/prepare_src_dataset.py` to add these DatagenInfo instances to the source datasets, along with metadata that stores which environment interface is to be used with this source dataset. They are also used during data generation to provide DatagenInfo instances that correspond to the current timestep of the environment being used for generation (via the `get_datagen_info` method). The environment interface objects are also used to go back and forth between environment actions (used by `env.step`) and target poses for the end effector controller (as described in MimicGen Appendix N.1). + +Every simulation framework should implement a base environment interface class that subclasses the `MG_EnvInterface` abstract base class in `env_interfaces/base.py`. For example, the `RobosuiteInterface` class in `env_interfaces/robosuite.py` is the base environment interface class for robosuite simulation environments. + +
+

Note

+ +The [Generating Data for New Simulation Frameworks](https://mimicgen.github.io/docs/tutorials/datagen_custom.html#generating-data-for-new-simulation-frameworks) tutorial provides a concrete example of how to implement a base environment interface class for a new simulation environment. + +
+ +Every simulation task should also implement a task-specific environment interface class that subclasses the corresponding base environment interface class for that simulation framework. + +
+

Note

+ +The [Generating Data for New Tasks](https://mimicgen.github.io/docs/tutorials/datagen_custom.html#generating-data-for-new-tasks) tutorial provides a concrete example of how to implement a task-specific environment interface class for new tasks. + +
diff --git a/docs/modules/overview.md b/docs/modules/overview.md new file mode 100644 index 0000000..b069ffa --- /dev/null +++ b/docs/modules/overview.md @@ -0,0 +1,32 @@ +# Codebase Overview + +## Codebase Structure + +We outline some important folders and files below. + +- `mimicgen/scripts`: utility scripts + - `generate_dataset.py`: main script for data generation +- `mimicgen/exps/templates`: collection of data generation config json templates for each task +- `mimicgen/configs`: implementation of data generation config classes + - `config.py`: base config class + - `task_spec.py`: [TaskSpec](https://mimicgen.github.io/docs/modules/task_spec.html) object for specifying sequence of subtasks for each task + - `robosuite.py`: robosuite-specific config classes +- `mimicgen/env_interfaces`: implementation of [Environment Interface](https://mimicgen.github.io/docs/modules/env_interfaces.html) classes that help simulation environments provide datagen info during data generation +- `mimicgen/datagen`: implementation of core [Data Generation](https://mimicgen.github.io/docs/modules/datagen.html) classes + - `data_generator.py`: [DataGenerator](https://mimicgen.github.io/docs/modules/datagen.html#data-generator) class used to generate new trajectories + - `datagen_info.py`: [DatagenInfo](https://mimicgen.github.io/docs/modules/datagen.html#datagen-info) class to group information from the sim environment needed during data generation + - `selection_strategy.py`: [SelectionStrategy](https://mimicgen.github.io/docs/modules/datagen.html#selection-strategy) classes that contain different heuristics for selecting source demos during each data generation trial + - `waypoint.py`: collection of [Waypoint](https://mimicgen.github.io/docs/modules/datagen.html#waypoint) classes to help end effector controllers execute waypoint targets and waypoint sequences +- `mimicgen/envs` and `mimicgen/models`: files containing collection of robosuite simulation environments and assets released with this project +- `mimicgen/utils`: collection of utility functions and classes +- `docs`: files related to documentation + +## Important Modules + +We provide some more guidance on some important modules and how they relate to one another. + +MimicGen starts with a handful of source demonstrations and generates new demonstrations automatically. MimicGen treats each task as a sequence of object-centric subtasks, and attempts to generate trajectories one subtask at a time. MimicGen must parse source demonstrations into contiguous subtask segments -- it uses [Subtask Termination Signals](https://mimicgen.github.io/docs/tutorials/subtask_termination_signals.html) to do this. It also requires object poses at the start of each subtask, both in the source demonstrations and in the current scene during data generation. Information on object poses, subtask termination signals, and other information needed at data generation time is collected into [DatagenInfo](https://mimicgen.github.io/docs/modules/datagen.html#datagen-info) objects, which are read from the source demonstrations, and also read from the current scene. This information is provided through [Environment Interface](https://mimicgen.github.io/docs/modules/env_interfaces.html) classes which connect underlying simulation environments to DatagenInfo objects. + +Data generation is carried out by the [DataGenerator](https://mimicgen.github.io/docs/modules/datagen.html#data-generator) class. Each data generation attempt requires choosing one or more subtask segments from the source demonstrations to transform -- this is carried out by a [SelectionStrategy](https://mimicgen.github.io/docs/modules/datagen.html#selection-strategy) instance. The transformation consists of keeping track of a collection of end effector target poses for a controller to execute -- this is managed by [Waypoint](https://mimicgen.github.io/docs/modules/datagen.html#waypoint) classes. + +The sequence of object-centric subtasks and other important data generation settings for each data generation run are communicated to MimicGen through the [TaskSpec object](https://mimicgen.github.io/docs/modules/task_spec.html), which is read as part of the MimicGen config. diff --git a/docs/modules/task_spec.md b/docs/modules/task_spec.md new file mode 100644 index 0000000..77e3be2 --- /dev/null +++ b/docs/modules/task_spec.md @@ -0,0 +1,78 @@ +# Task Spec + +The sequence of object-centric subtasks and other important data generation settings for each data generation run are communicated to MimicGen through the TaskSpec object (`configs/task_spec.py`). The TaskSpec object is instantiated from the MimicGen task config (`config.task.task_spec`) using the `TaskSpec.from_json` method. + +
+

Note

+ +See [this section](https://mimicgen.github.io/docs/tutorials/datagen_custom.html#step-2-implement-task-specific-config) of the Data Generation for Custom Environments tutorial for an example of how to implement a task config for a new task. + +
+ +We describe components of the TaskSpec object in more detail below. The TaskSpec is essentially list of dictionaries with each dictionary corresponding to exactly one subtask. The method below highlights the important settings that the TaskSpec object holds for each task: + +```python + def add_subtask( + self, + object_ref, + subtask_term_signal, + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0., + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ): + """ + Add subtask to this task spec. + + Args: + object_ref (str): each subtask involves manipulation with + respect to a single object frame. This string should + specify the object for this subtask. The name + should be consistent with the "datagen_info" from the + environment interface and dataset. + + subtask_term_signal (str or None): the "datagen_info" from the environment + and dataset includes binary indicators for each subtask + of the task at each timestep. This key should correspond + to the key in "datagen_info" that should be used to + infer when this subtask is finished (e.g. on a 0 to 1 + edge of the binary indicator). Should provide None for the final + subtask. + + subtask_term_offset_range (2-tuple): if provided, specifies time offsets to + be used during data generation when splitting a trajectory into + subtask segments. On each data generation attempt, an offset is sampled + and added to the boundary defined by @subtask_term_signal. + + selection_strategy (str): specifies how the source subtask segment should be + selected during data generation from the set of source human demos + + selection_strategy_kwargs (dict or None): optional keyword arguments for the selection + strategy function used + + action_noise (float): amount of action noise to apply during this subtask + + num_interpolation_steps (int): number of interpolation steps to bridge previous subtask segment + to this one + + num_fixed_steps (int): number of additional steps (with constant target pose of beginning of + this subtask segment) to add to give the robot time to reach the pose needed to carry + out this subtask segment + + apply_noise_during_interpolation (bool): if True, apply action noise during interpolation phase + leading up to this subtask, as well as during the execution of this subtask + """ +``` + +- The `object_ref` for each subtask determines the reference object frame for the motion in that subtask. The name here should be consistent with the `get_object_poses` method of the relevant task-specific [Environment Interface](https://mimicgen.github.io/docs/modules/env_interfaces.html)) object. + +- The `subtask_term_signal` and `subtask_term_offset_range` settings determine how subtask segments for this subtask is parsed from the source demonstrations -- see the [Subtask Termination Signals](https://mimicgen.github.io/docs/tutorials/subtask_termination_signals.html) page for more information. + +- The `selection_strategy` and `selection_strategy_kwargs` determine the [SelectionStrategy](https://mimicgen.github.io/docs/modules/datagen.html#selection-strategy) class used to select a source subtask segment at the start of each subtask during data generation. See Appendix N.3 in the MimicGen paper for more details. + +- The `action_noise` setting determines the magnitude of action noise added when executing actions during data generation. See Appendix N.4 of the MimicGen paper for more details. + +- The `num_interpolation_steps` and `num_fixed_steps` settings determines the number of interpolation waypoints in the interpolation segment that bridges this subtask segment and the previous subtask segment during data generation. See Appendix N.2 of the MimicGen paper for more details. diff --git a/docs/tutorials/datagen_custom.md b/docs/tutorials/datagen_custom.md new file mode 100644 index 0000000..177d5f6 --- /dev/null +++ b/docs/tutorials/datagen_custom.md @@ -0,0 +1,484 @@ +# Data Generation for Custom Environments + +In this section, we provide guidance on using MimicGen to generate data for custom tasks and simulation frameworks. + +
+

Note

+ +We recommend going through the [Getting Started](https://mimicgen.github.io/docs/tutorials/getting_started.html) tutorial first, so that you are familiar with the typical data generation workflow. We will refer back to the steps in the [Data Generation Workflow Overview](https://mimicgen.github.io/docs/tutorials/getting_started.html#overview-of-typical-data-generation-workflow) in this section. + +
+ +## Generating Data for New Tasks + +In this section, we will assume we are trying to generate data for a new task implemented in [robosuite](https://robosuite.ai/), and we will use the robosuite Stack Three task as a running example. The same instructions can be used for any task in any simulation framework, as long as an [Environment Interface](https://mimicgen.github.io/docs/modules/env_interfaces.html) base class already exists for the simulation framework. See the [Generating Data for New Simulation Frameworks](https://mimicgen.github.io/docs/tutorials/datagen_custom.html#generating-data-for-new-simulation-frameworks)) below for guidance on setting up a new simulation framework if this has not happened yet. + +### Step 1: Implement Task-Specific Environment Interface + +The first step is to subclass the appropriate base Environment Interface class -- for robosuite, this is the `RobosuiteInterface` class at the top of `env_interfaces/robosuite.py`. We create a new class as below: + +```python +class MG_StackThree(RobosuiteInterface): + """ + Corresponds to robosuite StackThree task and variants. + """ + pass +``` + +The `MG_EnvInterface` abstract base class in `env_interfaces/base.py` (which `RobosuiteInterface` inherits from) describes the methods that task-specific subclasses must implement. There are two important methods: + +```python + """ + These should be filled out by each simulation domain (e.g. nut assembly, coffee). + """ + @abc.abstractmethod + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + raise NotImplementedError + + @abc.abstractmethod + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + raise NotImplementedError +``` + +Recall that MimicGen generates data by composing object-centric subtask segments together (see the paper for more information). During data generation, MimicGen requires a way to observe the pose of the relevant object at the start of each subtask. The `get_object_poses` method will be used for this purpose - it should return a dictionary mapping object name to a pose matrix. + +The `RobosuiteInterface` base class offers a helper method `get_object_pose(self, obj_name, obj_type)` to make retrieving object poses from robosuite easily - we use it below to get the poses of each cube in the `StackThree` task. + +```python + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # three relevant objects - three cubes + return dict( + cubeA=self.get_object_pose(obj_name=self.env.cubeA.root_body, obj_type="body"), + cubeB=self.get_object_pose(obj_name=self.env.cubeB.root_body, obj_type="body"), + cubeC=self.get_object_pose(obj_name=self.env.cubeC.root_body, obj_type="body"), + ) +``` + +Next we need to implement `get_subtask_term_signals`. This function has only one purpose - it is used to provide [Subtask Termination Signals](https://mimicgen.github.io/docs/tutorials/subtask_termination_signals.html) for each timestep in the source demonstrations (this is part of what happens in `scripts/prepare_src_dataset.py`). These signals are used to determine where each subtask ends and the next one starts -- the first 0 to 1 transition in this signal during a source demonstration determines the end of the subtask. + +The StackThree tasks consists of 4 object-centric subtasks: + +``` +1. grasping cubeA (motion relative to cubeA) +2. placing cubeA on cubeB (motion relative to cubeB) +3. grasping cubeC (motion relative to cubeC) +4. placing cubeC on cubeA (motion relative to cubeA) +``` + +To define the end of subtask 1 and 3, we can just check for a successful grasp, and for the end of subtask 2, we can check for a placement (re-using part of the success check for the robosuite StackThree task): + +```python + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + # first subtask is grasping cubeA (motion relative to cubeA) + signals["grasp_1"] = int(self.env._check_grasp(gripper=self.env.robots[0].gripper, object_geoms=self.env.cubeA)) + + # second subtask is placing cubeA on cubeB (motion relative to cubeB) + signals["stack_1"] = int(self.env._check_cubeA_stacked()) + + # third subtask is grasping cubeC (motion relative to cubeC) + signals["grasp_2"] = int(self.env._check_grasp(gripper=self.env.robots[0].gripper, object_geoms=self.env.cubeC)) + + # final subtask is placing cubeC on cubeA (motion relative to cubeA) - but final subtask signal is not needed + return signals +``` + +
+

Warning

+ +The final subtask in a task never requires a subtask termination signal, since its end is determined by the end of the source demonstration. + +
+ +
+

Note

+ +Providing a proper implementation for the `get_subtask_term_signals` function is entirely optional. In most cases it is easy to specify heuristics to define these subtask boundaries as we did above, but sometimes you may want to just directly annotate the boundaries between subtasks. We provide an annotation script (`scripts/annotate_subtasks.py`) for this purpose. If you plan to do this, you can just return an empty dict for the `get_subtask_term_signals` function. + +
+ +### Step 2: Implement Task-Specific Config + +The next step is to implement a task-specific Config object that inherits from the `MG_Config` base class (`configs/config.py`). There are three things for the subclass to implement: + +```python + @property + @classmethod + def NAME(cls): + # must be specified by subclasses + raise NotImplementedError + + @property + @classmethod + def TYPE(cls): + # must be specified by subclasses + raise NotImplementedError + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. + """ + raise NotImplementedError +``` + +The `NAME` and `TYPE` are used to store the new subclass into the config registry, and chiefly determine where the auto-generated config template is stored in the repository (e.g. `mimicgen/exps/templates//.json`). + +The `task_config` is consistent with [TaskSpec objects](https://mimicgen.github.io/docs/modules/task_spec.html) -- there is an entry for each subtask which consists of the arguments to the `add_subtask` function in the TaskSpec object (`configs/task_spec.py`): + +```python + def add_subtask( + self, + object_ref, + subtask_term_signal, + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0., + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ): +``` + +We show the implementation for the StackThree config below: + +```python +class StackThree_Config(MG_Config): + """ + Corresponds to robosuite StackThree task and variants. + """ + NAME = "stack_three" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + self.task.task_spec.subtask_1 = dict( + object_ref="cubeA", + subtask_term_signal="grasp_1", + subtask_term_offset_range=(10, 20), + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_2 = dict( + object_ref="cubeB", + subtask_term_signal="stack_1", + subtask_term_offset_range=(10, 20), + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_3 = dict( + object_ref="cubeC", + subtask_term_signal="grasp_2", + subtask_term_offset_range=(10, 20), + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_4 = dict( + object_ref="cubeA", + subtask_term_signal=None, + subtask_term_offset_range=None, + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.do_not_lock_keys() +``` + +Notice that we set the `object_ref` for each subtask to be consistent with the object names in the `get_object_poses` method in the `MG_StackThree` environment interface we implemented. We also set the `subtask_term_signal` for each subtask to be consistent with the subtask signals in the `get_subtask_term_signals` method in the `MG_StackThree` class as well. Please see the [TaskSpec page](https://mimicgen.github.io/docs/modules/task_spec.html) for more information on the other settings. + +
+

Note

+ +If you used or plan to use `scripts/annotate_subtasks.py` to manually annotate the end of each subtask in the source demos, you should use signal names that are consistent with the `--signals` argument that you will pass to that script that give a name to each subtask. Internally, the annotations are stored as subtask termination signals with those names. + +
+ +
+

Note

+ +You should make sure that the config class you implemented is being imported somewhere in your codebase to make sure it gets registered in the config registry. In the MimicGen codebase, we do this in `mimicgen/configs/__init__.py`. + +
+ +Finally, run `scripts/generate_config_templates.py` to generate a config template for this new task. It should appear under `mimicgen/exps/templates//.json`. Ensure that the default settings look correct. These settings can be overridden using config generators ([see this tutorial](https://mimicgen.github.io/docs/tutorials/launching_several.html)). + +### Step 3: Execute Data Generation Workflow + +You are now all set to try data generation. You should be able to follow the steps documented in the [Data Generation Workflow Overview](https://mimicgen.github.io/docs/tutorials/getting_started.html#overview-of-typical-data-generation-workflow) with some minor changes. + +
+

Note

+ +Now that you have followed these steps, you can generate datasets for any other variants of this task, as long as the TaskSpec does not change -- e.g. different object placements, different robot arms, and different object instances. The [Reproducing Experiments](https://mimicgen.github.io/docs/tutorials/reproducing_experiments.html) tutorial provides examples of all three variations. + +
+ + +## Generating Data for New Simulation Frameworks + +
+

Note

+ +Before starting, you should ensure that a [robomimic environment wrapper](https://robomimic.github.io/docs/modules/environments.html) exists for the simulation framework you are using. See [this link](https://robomimic.github.io/docs/modules/environments.html#implement-an-environment-wrapper) for guidance on how to create one. The environment metadata in the source hdf5 should point to this environment wrapper. + +
+ +In this section, we will show how to apply MimicGen to new simulation frameworks. The key step is to implement an [Environment Interface](https://mimicgen.github.io/docs/modules/env_interfaces.html) base class for the simulation framework. We will use robosuite as a running example in this section. + +The `MG_EnvInterface` abstract base class in `env_interfaces/base.py` describes the methods that base subclasses for new simulators must implement. There are five important methods: + +```python + """ + These should be filled out by simulator subclasses (e.g. robosuite). + """ + @property + @classmethod + def INTERFACE_TYPE(self): + """ + Returns string corresponding to interface type. This is used to group + all subclasses together in the interface registry (for example, all robosuite + interfaces) and helps avoid name conflicts. + """ + raise NotImplementedError + + @abc.abstractmethod + def get_robot_eef_pose(self): + """ + Get current robot end effector pose. Should be the same frame as used by the robot end-effector controller. + + Returns: + pose (np.array): 4x4 eef pose matrix + """ + raise NotImplementedError + + @abc.abstractmethod + def target_pose_to_action(self, target_pose, relative=True): + """ + Takes a target pose for the end effector controller and returns an action + (usually a normalized delta pose action) to try and achieve that target pose. + + Args: + target_pose (np.array): 4x4 target eef pose + relative (bool): if True, use relative pose actions, else absolute pose actions + + Returns: + action (np.array): action compatible with env.step (minus gripper actuation) + """ + raise NotImplementedError + + @abc.abstractmethod + def action_to_target_pose(self, action, relative=True): + """ + Converts action (compatible with env.step) to a target pose for the end effector controller. + Inverse of @target_pose_to_action. Usually used to infer a sequence of target controller poses + from a demonstration trajectory using the recorded actions. + + Args: + action (np.array): environment action + relative (bool): if True, use relative pose actions, else absolute pose actions + + Returns: + target_pose (np.array): 4x4 target eef pose that @action corresponds to + """ + raise NotImplementedError + + @abc.abstractmethod + def action_to_gripper_action(self, action): + """ + Extracts the gripper actuation part of an action (compatible with env.step). + + Args: + action (np.array): environment action + + Returns: + gripper_action (np.array): subset of environment action for gripper actuation + """ + raise NotImplementedError +``` + +The `INTERFACE_TYPE` method is just used to make sure there are no class name conflicts in the environment interface registry. You should just make sure to choose a name unique to your simulation framework. For robosuite, we simply used `"robosuite"`. + +The `get_robot_eef_pose` method returns the current pose of the robot end effector and corresponds to the same frame used by the end effector controller in the simulation environment. In robosuite, the Operational Space Controller uses a specific MuJoCo site, so we return its pose: + +```python + def get_robot_eef_pose(self): + """ + Get current robot end effector pose. Should be the same frame as used by the robot end-effector controller. + + Returns: + pose (np.array): 4x4 eef pose matrix + """ + + # OSC control frame is a MuJoCo site - just retrieve its current pose + return self.get_object_pose( + obj_name=self.env.robots[0].controller.eef_name, + obj_type="site", + ) +``` + +The next two methods (`target_pose_to_action`, `action_to_target_pose`) are used to translate between simulator actions (e.g. those given to `env.step`) and absolute target poses for the end effector controller. This typically consists of simple scaling factors and transformations between different rotation conventions (as described in Appendix N.1 of the MimicGen paper): + +```python + def target_pose_to_action(self, target_pose, relative=True): + """ + Takes a target pose for the end effector controller and returns an action + (usually a normalized delta pose action) to try and achieve that target pose. + + Args: + target_pose (np.array): 4x4 target eef pose + relative (bool): if True, use relative pose actions, else absolute pose actions + + Returns: + action (np.array): action compatible with env.step (minus gripper actuation) + """ + + # version check for robosuite - must be v1.2+, so that we're using the correct controller convention + assert (robosuite.__version__.split(".")[0] == "1") + assert (robosuite.__version__.split(".")[1] >= "2") + + # target position and rotation + target_pos, target_rot = PoseUtils.unmake_pose(target_pose) + + # current position and rotation + curr_pose = self.get_robot_eef_pose() + curr_pos, curr_rot = PoseUtils.unmake_pose(curr_pose) + + # get maximum position and rotation action bounds + max_dpos = self.env.robots[0].controller.output_max[0] + max_drot = self.env.robots[0].controller.output_max[3] + + if relative: + # normalized delta position action + delta_position = target_pos - curr_pos + delta_position = np.clip(delta_position / max_dpos, -1., 1.) + + # normalized delta rotation action + delta_rot_mat = target_rot.dot(curr_rot.T) + delta_quat = T.mat2quat(delta_rot_mat) + delta_rotation = T.quat2axisangle(delta_quat) + delta_rotation = np.clip(delta_rotation / max_drot, -1., 1.) + return np.concatenate([delta_position, delta_rotation]) + + # absolute position and rotation action + target_quat = T.mat2quat(target_rot) + abs_rotation = T.quat2axisangle(target_quat) + return np.concatenate([target_pos, abs_rotation]) + + def action_to_target_pose(self, action, relative=True): + """ + Converts action (compatible with env.step) to a target pose for the end effector controller. + Inverse of @target_pose_to_action. Usually used to infer a sequence of target controller poses + from a demonstration trajectory using the recorded actions. + + Args: + action (np.array): environment action + relative (bool): if True, use relative pose actions, else absolute pose actions + + Returns: + target_pose (np.array): 4x4 target eef pose that @action corresponds to + """ + + # version check for robosuite - must be v1.2+, so that we're using the correct controller convention + assert (robosuite.__version__.split(".")[0] == "1") + assert (robosuite.__version__.split(".")[1] >= "2") + + if (not relative): + # convert absolute action to absolute pose + target_pos = action[:3] + target_quat = T.axisangle2quat(action[3:6]) + target_rot = T.quat2mat(target_quat) + else: + # get maximum position and rotation action bounds + max_dpos = self.env.robots[0].controller.output_max[0] + max_drot = self.env.robots[0].controller.output_max[3] + + # unscale actions + delta_position = action[:3] * max_dpos + delta_rotation = action[3:6] * max_drot + + # current position and rotation + curr_pose = self.get_robot_eef_pose() + curr_pos, curr_rot = PoseUtils.unmake_pose(curr_pose) + + # get pose target + target_pos = curr_pos + delta_position + delta_quat = T.axisangle2quat(delta_rotation) + delta_rot_mat = T.quat2mat(delta_quat) + target_rot = delta_rot_mat.dot(curr_rot) + + target_pose = PoseUtils.make_pose(target_pos, target_rot) + return target_pose +``` + +Finally, the `action_to_gripper_action` extracts the part of the simulator action that corresponds to gripper actuation: + +```python + def action_to_gripper_action(self, action): + """ + Extracts the gripper actuation part of an action (compatible with env.step). + + Args: + action (np.array): environment action + + Returns: + gripper_action (np.array): subset of environment action for gripper actuation + """ + + # last dimension is gripper action + return action[-1:] +``` + +Finally, you can follow the instructions in the [Generating Data for New Tasks](https://mimicgen.github.io/docs/tutorials/datagen_custom.html#generating-data-for-new-tasks) section to setup data generation for specific tasks in this simulation framework. diff --git a/docs/tutorials/debugging_datagen.md b/docs/tutorials/debugging_datagen.md new file mode 100644 index 0000000..6e40740 --- /dev/null +++ b/docs/tutorials/debugging_datagen.md @@ -0,0 +1,79 @@ +# Debugging Data Generation + +We provide some useful suggestions for debugging data generation runs. + +## Source Demo Validation + +### Get Source Dataset Information + +You can use the `get_source_info.py` script to validate whether the source demonstrations you are using have the expected [DatagenInfo](https://mimicgen.github.io/docs/modules/datagen.html#datagen-info) structure, and are using the correct [Environment Interface](https://mimicgen.github.io/docs/modules/env_interfaces.html): + +```sh +$ python mimicgen/scripts/get_source_info.py --dataset datasets/source/square.hdf5 +``` + +It will print out information that looks like the following. This is a good way to also validate the object poses and subtask termination signals present in the file: + +``` +Environment Interface: MG_Square +Environment Interface Type: robosuite + +Structure of datagen_info in episode demo_0: + eef_pose: shape (127, 4, 4) + gripper_action: shape (127, 1) + object_poses: + square_nut: shape (127, 4, 4) + square_peg: shape (127, 4, 4) + subtask_term_signals: + grasp: shape (127,) + target_pose: shape (127, 4, 4) +``` + +### Visualizing Subtasks in Source Dataset + +You can visualize each subtask segment in a source demonstration using the `visualize_subtasks.py` script. + +The script needs to be aware of the order of the subtask signals as well as the maximum termination offsets being used (see the [Subtask Termination Signals](https://mimicgen.github.io/docs/tutorials/subtask_termination_signals.html)) page for more information on offsets) -- this can be specified by providing a config json (`--config`) or by providing the sequence of signals (`--signals`) and offsets (`--offsets`) for all except the last subtask. The end of each subtask is defined as the timestep of the first 0 to 1 transition in the corresponding signal added to the maximum offset value. + +The script supports either on-screen rendering (`--render`) or off-screen rendering to a video (`--video_path`). If using on-screen rendering, the script will pause after every subtask segment. If using off-screen rendering, the video alternates between no borders and red borders to show each subtask segment. + +```sh +# render on-screen +$ python visualize_subtasks.py --dataset /path/to/demo.hdf5 --config /path/to/config.json --render + +# render to video +$ python visualize_subtasks.py --dataset /path/to/demo.hdf5 --config /path/to/config.json --video_path /path/to/video.mp4 + +# specify subtask information manually instead of using a config +$ python visualize_subtasks.py --dataset /path/to/demo.hdf5 --signals grasp_1 insert_1 grasp_2 --offsets 10 10 10 --render +``` + +The script is useful to tune the [Subtask Termination Signals](https://mimicgen.github.io/docs/tutorials/subtask_termination_signals.html) specified through the environment interface (the `get_subtask_term_signals` method) or the manual annotations provided through `scripts/annotate_subtasks.py` have defined that subtask segments properly, as well as the offsets you are using in the data generation config. + +
+

Warning

+ +When you change the `get_subtask_term_signals` method, you should re-run the `prepare_src_dataset.py` script on the source data to re-write the subtask termination signals to the dataset. + +
+ +## Visualization during Data Generation + +The main data generation script (`scripts/generate_dataset.py`) also has some useful features for debugging, including on-screen visualization (`--render`), off-screen rendering to video (`--video_path`), running a quick run for debugging (`--debug`), and pausing after each subtask execution (`--pause_subtask`): + +```sh +# run normal data generation +$ python generate_dataset.py --config /path/to/config.json + +# render all data generation attempts on-screen +$ python generate_dataset.py --config /path/to/config.json --render + +# render all data generation attempts to a video +$ python generate_dataset.py --config /path/to/config.json --video_path /path/to/video.mp4 + +# run a quick debug run +$ python generate_dataset.py --config /path/to/config.json --debug + +# pause after every subtask to debug data generation +$ python generate_dataset.py --config /path/to/config.json --render --pause_subtask +``` diff --git a/docs/tutorials/getting_started.md b/docs/tutorials/getting_started.md new file mode 100644 index 0000000..0300873 --- /dev/null +++ b/docs/tutorials/getting_started.md @@ -0,0 +1,142 @@ +# Getting Started + +
+

Note

+ +This section helps users get started with data generation. If you would just like to download our existing datasets and use them with policy learning methods please see the [Reproducing Experiments](https://mimicgen.github.io/docs/tutorials/reproducing_experiments.html) tutorial for a guide, or the [Datasets](https://mimicgen.github.io/docs/datasets/mimicgen_corl_2023.html) page to get details on the datasets. + +
+ + +## Quick Data Generation Run + +Let's run a quick data generation example. + +Before starting, make sure you are at the base repo path: +```sh +$ cd {/path/to/mimicgen} +``` + +### Step 1: Prepare source human dataset + +MimicGen requires a handful of human demonstrations to get started. + +Download the source demonstrations for the Square task below: +```sh +$ python mimicgen/scripts/download_datasets.py --dataset_type source --tasks square +``` + +This is a basic robomimic dataset collected via teleoperation in robosuite (see [here](https://robomimic.github.io/docs/datasets/robosuite.html)). We need to add in extra information to the hdf5 to make it compatible with MimicGen: +```sh +$ python mimicgen/scripts/prepare_src_dataset.py \ +--dataset datasets/source/square.hdf5 \ +--env_interface MG_Square \ +--env_interface_type robosuite +``` + +The `--env_interface` and `--env_interface_type` arguments allow the script to find the correct [Environment Interface](https://mimicgen.github.io/docs/modules/env_interfaces.html) class to extract [DatagenInfo objects](https://mimicgen.github.io/docs/modules/datagen.html#datagen-info) at each timestep. In general, each task needs to have an environment interface class to tell MimicGen how to retrieve object poses and other information needed during data generation. + +### Step 2: Prepare data generation config + +Each data generation run requires a config json (similar to robomimic [Configs](https://robomimic.github.io/docs/modules/configs.html)) that allows us to configure different settings. Template configs for each task are at `mimicgen/exps/templates` and are auto-generated (with `scripts/generate_config_templates.py`). The repository has easy ways to modify these templates to generate new config jsons. + +For now, we will use a script to produce experiment configs consistent with the MimicGen paper. Open `scripts/generate_core_configs.py` and set `NUM_TRAJ = 10` and `GUARANTEE = False` -- this means we will attempt to generate 10 new trajectories. + +
+

Warning

+ +If you do not edit `scripts/generate_core_configs.py` the default settings will run data generation until 1000 success trajectories have been collected. This is why it is important to set `NUM_TRAJ = 10` and `GUARANTEE = False` for a quick run. Alternatively, pass the `--debug` flag to the command in Step 3, which will run an even smaller data generation run. + +
+ +Next, run the script: + +```sh +$ python mimicgen/scripts/generate_core_configs.py +``` + +It generates a set of configs (and prints their paths) and also prints lines that correspond to data generation runs for each config. + +### Step 3: Run data generation and view outputs + +Next, we run data generation on the Square D1 task (this will take a couple minutes): +```sh +$ python mimicgen/scripts/generate_dataset.py \ +--config /tmp/core_configs/demo_src_square_task_D1.json \ +--auto-remove-exp +``` + +
+

Note

+ +If you run into a `RuntimeError: No ffmpeg exe could be found.` at the end of the script, this means rendering the dataset to video failed. We found that a simple `conda install ffmpeg` fixed the problem on our end (as documented on the [troubleshooting page](https://mimicgen.github.io/docs/miscellaneous/troubleshooting.html)). + +
+ +By default, the data generation folder can be found at `/tmp/core_datasets/square/demo_src_square_task_D1`. The contents of this folder are as follows: +``` +demo.hdf5 # generated hdf5 containing successful demonstrations +demo_failed.hdf5. # generated hdf5 containing failed demonstrations, up to a certain limit +important_stats.json # json file summarizing success rate and other statistics +log.txt # terminal output +logs/ # longer experiment runs will regularly sync progress jsons to this folder +mg_config.json # config used for this experiment +playback_demo_src_square_task_D1.mp4 # video that shows successful demonstrations, up to a certain limit +playback_demo_src_square_task_D1_failed.mp4. # video that shows failed demonstrations, up to a certain limit +``` + +
+

Note

+ +The generated `demo.hdf5` file is fully compatible with robomimic, which makes it easy to inspect the data (see [here](https://robomimic.github.io/docs/tutorials/dataset_contents.html)) or launch training jobs (we have a helper script at `scripts/generate_core_training_configs.py`). + +
+ +Next, we outline the typical MimicGen workflow. + + +## Overview of Typical Data Generation Workflow + +A typical application of MimicGen consists of the following steps. + + +### Step 1: Collect source demonstrations + +Collect human demonstrations for a task of interest. Typically, this is done using a teleoperation pipeline. For example, [this](https://robosuite.ai/docs/algorithms/demonstrations.html) is how robosuite demonstrations can be collected. Make sure you end up with an hdf5 [compatible with robomimic](https://robomimic.github.io/docs/datasets/mimicgen_corl_2023.html#dataset-structure) -- this typically involves a postprocessing script (for example [this one](https://robomimic.github.io/docs/datasets/robosuite.html#converting-robosuite-hdf5-datasets) for robosuite). + +
+

Note

+ +You should ensure that a [robomimic environment wrapper](https://robomimic.github.io/docs/modules/environments.html) exists for the simulation framework you are using. See [this link](https://robomimic.github.io/docs/modules/environments.html#implement-an-environment-wrapper) for guidance on how to create one. The environment metadata in the source hdf5 should point to this environment wrapper. + +
+ +### Step 2: Prepare source demonstrations with additional annotations + +Next, information must be added to the source hdf5 to make it compatible with MimicGen. This is typically done with `scripts/prepare_src_dataset.py` just like we did above. However, this requires an [Environment Interface](https://mimicgen.github.io/docs/modules/env_interfaces.html) class to be implemented for your simulation framework (usually a base class) and your task (a subclass of the base class). + +These classes typically specify how to translate between environment actions and target poses for the end effector controller in the environment (see the MimicGen paper for more details on why this is needed). Furthermore, for each task, they provide a dictionary that maps object name to object pose -- recall that MimicGen requires observing an object pose at the start of each object-centric subtask. The structure of the extracted information can be found [here](https://mimicgen.github.io/docs/modules/datagen.html#datagen-info). Finally, the class can optionally provide [subtask termination signals](https://mimicgen.github.io/docs/tutorials/subtask_termination_signals.html) that provide heuristics for splitting source demonstrations into subtask segments. As an example, the `MG_Square` environment interface senses when the nut has been grasped and provides this heuristic in the `grasp` subtask termination signal. + +Instead of using heuristics from the environment interface class, the source demonstrations can be manually segmented into subtasks using the annotation interface in `scripts/annotate_subtasks.py`. This step should be performed after running `scripts/prepare_src_dataset.py`. + +
+

Note

+ +See the [Data Generation for Custom Environments](https://mimicgen.github.io/docs/tutorials/datagen_custom.html) tutorial for a more comprehensive description of implementing environment interfaces for new simulators and new tasks. + +
+ +### Step 3: Run data generation + +Next, you can set up your MimicGen config and launch data generation. See the [Launching Several Data Generation Runs](https://mimicgen.github.io/docs/tutorials/launching_several.html) tutorial to see how data generation configs can be generated with ease. Once ready, run `scripts/generate_dataset.py`. + +
+

Note

+ +Data generation does not necessarily need to run on the exact same task. As described in the paper, you can run it on tasks with possibly different reset distributions, different robot arms, or different object instances. The [Reproducing Experiments](https://mimicgen.github.io/docs/tutorials/reproducing_experiments.html) tutorial provides examples of all three variations. + +
+ +### Step 4: Run policy learning on generated data + +You can now run any policy learning algorithm on the generated data to train an agent. A common choice is to run Behavioral Cloning. The generated data is compatible with [robomimic](https://robomimic.github.io/), which is an easy way to train an agent on generated datasets, and compare the performance of different learning methods. The [Reproducing Experiments](https://mimicgen.github.io/docs/tutorials/reproducing_experiments.html) tutorial shows examples of how to train agents on the generated data. diff --git a/docs/tutorials/launching_several.md b/docs/tutorials/launching_several.md new file mode 100644 index 0000000..a1d5d0a --- /dev/null +++ b/docs/tutorials/launching_several.md @@ -0,0 +1,18 @@ +# Launching Several Data Generation Runs + +MimicGen inherits the [Config](https://robomimic.github.io/docs/tutorials/configs.html) system from robomimic. Configs are specified as json files, and support both dictionary and "dot" syntax (e.g. `config.experiment.name` and `config["experiment"]["name"]`). + +MimicGen also uses the `ConfigGenerator` class from robomimic, and can use it to generate several config jsons efficiently. For a tutorial on how this generator works, please see the [tutorial from robomimic](https://robomimic.github.io/docs/tutorials/hyperparam_scan.html). In this repository, we grouped several related settings together into helper functions in `utils/config_utils.py` that operate over `ConfigGenerator` objects. + +We furthermore provide examples of how to create and use multiple `ConfigGenerator` objects in scripts such as `scripts/generate_core_configs.py` and `scripts/generate_core_training_configs.py`. These scripts support using multiple base configs (e.g. one per task, since the task spec for each task will be different), and a user can specify different parameter settings per base config (see the `make_generators` function in these files). There are additional settings that are specifed as global variables at the top of the files. These scripts print the file paths for all generated configs, and the commands to launch runs for each config. + +Users can easily modify the following files to generate large amounts of data generation and policy learning configs efficiently. + +Data Generation: +- `scripts/generate_core_configs.py` +- `scripts/generate_robot_transfer_configs.py` + + +Policy Learning: +- `scripts/generate_core_training_configs.py` +- `scripts/generate_robot_transfer_training_configs.py` diff --git a/docs/tutorials/reproducing_experiments.md b/docs/tutorials/reproducing_experiments.md new file mode 100644 index 0000000..05de63a --- /dev/null +++ b/docs/tutorials/reproducing_experiments.md @@ -0,0 +1,58 @@ +# Reproducing Published Experiments and Results + +There are two options for reproducing the set of results in the MimicGen paper. + +## Option 1: Download datasets and run policy learning + +You can directly download the datasets we generated and used in the MimicGen paper, and then subsequently run policy learning on the downloaded data. See the [Datasets](https://mimicgen.github.io/docs/datasets/mimicgen_corl_2023.html) page for exact instructions on how to do this. + +## Option 2: Run data generation and then policy learning + +In this section, we show how to run MimicGen on the source demonstrations we released to generate datasets equivalent to the ones we produced in our paper, and then subsequently train policies on the generated data. + +
+

Note

+ +We recommend going through the [Getting Started](https://mimicgen.github.io/docs/tutorials/getting_started.html) tutorial first, so that you are familiar with the way data generation works. + +
+ +The steps are very similar to the steps taken for the quick data generation run in the [Getting Started](https://mimicgen.github.io/docs/tutorials/getting_started.html) tutorial. We provide a brief outline of the steps and important changes below. + +Before starting, make sure you are at the base repo path: +```sh +$ cd {/path/to/mimicgen} +``` + +### Step 1: Prepare source human datasets + +Download all source demonstrations of interest. You can download all of them with the command below (optionally provide the `--download_dir` argument to set the download path): +```sh +$ python mimicgen/scripts/download_datasets.py --dataset_type source --tasks all +``` + +We need to prepare each one for data generation. The bash script `scripts/prepare_all_src_datasets.sh` outline the commands for each source dataset. We provide the command for Coffee below: +```sh +$ python mimicgen/scripts/prepare_src_dataset.py --dataset datasets/source/coffee.hdf5 --env_interface MG_Coffee --env_interface_type robosuite +``` + +### Step 2: Prepare data generation configs + +Open `scripts/generate_core_configs.py` and set `NUM_TRAJ = 1000` and `GUARANTEE = True` -- this means we will keep generating data until we generate 1000 successful trajectories. You can set additional parameters at the top of the file as well, e.g. in case you would like to change where data is generated. + +Next, run the script: +```sh +$ python mimicgen/scripts/generate_core_configs.py +``` + +The generated configs correspond to the **core** dataset type described on the [Datasets](https://mimicgen.github.io/docs/datasets/mimicgen_corl_2023.html) page, and the **object** dataset type as well (Mug Cleanup O1 and O2). + +If you would also like to generate configs for the **robot** dataset type (robot transfer experiments), you can follow the same steps above for `scripts/generate_robot_transfer_configs.py`. + +### Step 3: Run data generation + +The scripts above print lines that correspond to data generation runs for each config. You can pick and choose which ones you would like to run and then launch them with `scripts/generate_dataset.py`. + +### Step 4: Run policy training + +Finally, you can run policy training on the generated data -- to reproduce the paper results, you can run BC-RNN. To make this easy, we provide `scripts/generate_core_training_configs.py` and `scripts/generate_robot_transfer_training_configs.py`. As before, you can configure some settings in global variables at the top of the file, such as where to store policy training results. These scripts generate robomimic training configs that can be run with `scripts/train.py` in the robomimic repository. diff --git a/docs/tutorials/subtask_termination_signals.md b/docs/tutorials/subtask_termination_signals.md new file mode 100644 index 0000000..ba1fde7 --- /dev/null +++ b/docs/tutorials/subtask_termination_signals.md @@ -0,0 +1,75 @@ +# Subtask Termination Signals + +## What are these signals? + +MimicGen expects subtask termination signals to be present in each episode of the source demonstrations. Each signal is a flat numpy array with binary entries (e.g. 0 or 1). + +The [get_source_info.py script](https://mimicgen.github.io/docs/tutorials/debugging_datagen.html#get-source-dataset-information) can be used to print the signals present in a given source demonstration. For an episode (e.g. `demo_0`) they are expected to be at `f["data/demo_0/datagen_info/subtask_term_signals"]`. Under this hdf5 group, there will be one or more hdf5 datasets with the name of the subtask termination signal, and the corresponding flat numpy array. + +## How are they used? + +These signals are read from the source dataset by the `parse_source_dataset` function (see `mimicgen/utils/file_utils.py`) during data generation and used to split each source demonstration into contiguous subtask segments. Each subtask corresponds to a specific subtask termination signal, with the exception of the final subtask, which ends at the end of the source demonstration. This mapping between the subtask and the corresponding signal is specified through the [TaskSpec](https://mimicgen.github.io/docs/modules/task_spec.html) object that comes from the MimicGen config json. + +The end of each subtask is inferred from the first 0 to 1 transition in the corresponding signal. For example, let us consider the source demonstrations for the robosuite StackThree task. There are 4 subtasks with corresponding subtask termination signals: + +``` +1. (signal: grasp_1) grasping cubeA (motion relative to cubeA) +2. (signal: place_1) placing cubeA on cubeB (motion relative to cubeB) +3. (signal: grasp_2) grasping cubeC (motion relative to cubeC) +4. (signal: None) placing cubeC on cubeA (motion relative to cubeA) +``` + +For the first source demonstration, the first 0 to 1 transition for grasp_1 is 50, for place_1 is 94, and for grasp_2 is 148. This splits the first source demonstration into four subtask segments with start and end indices as follows: + +``` +1. (signal: grasp_1) [0, 50] +2. (signal: place_1) [50, 94] +3. (signal: grasp_2) [94, 148] +4. (signal: None) [148, 210] +``` + +These source subtask segments are subsequently transformed and stitched together through linear interpolation to generate data for new scenes. + +However, MimicGen also supports randomization of subtask boundaries -- this is where the `subtask_term_offset_range` parameter of the `TaskSpec` becomes relevant (see [TaskSpec](https://mimicgen.github.io/docs/modules/task_spec.html)). At the start of each data generation attempt, the subtask boundaries (indices 50, 94, and 148 above) can be randomized with an additive offset uniformly sampled in the given offset range bounds. + +The offset parameter can also be used to ensure that the end of each subtask happens at least N timesteps after the first 0 to 1 transition in the corresponding signal. For example, grasp_1 detects a successful grasp, but perhaps you would like the end of the subtask to be a little after grasping (e.g. when the lift begins). An easy way to do this is to specify an offset_range like (5, 10), so that the true subtask boundary will always occur 5 to 10 timesteps after the 0 to 1 transition. + +## How are they added to the source data? + +There are two ways that subtask termination signals can be added to a source dataset. The first is through the `prepare_src_dataset.py` script -- this will use the task-specific [Environment Interface](https://mimicgen.github.io/docs/modules/env_interfaces.html) (specifically, the `get_subtask_term_signals` method that is used to populate [DatagenInfo](https://mimicgen.github.io/docs/modules/datagen.html#datagen-info) objects) to get subtask termination signals for each timestep in each source demonstration. This relies on some heuristics per task. An example of this process can be found in the [Data Generation for Custom Environments](https://mimicgen.github.io/docs/tutorials/datagen_custom.html) tutorial. + +The second is to use manual human annotations for the end of each subtask in each source demonstration -- this can be done by using `scripts/annotate_subtasks.py`. + +
+

Note

+ +This script requires the `pygame` module to be installed (it will not be installed by default when installing the Mimicgen repo). + +
+ +The script plays back demonstrations (using visual observations and the pygame renderer) in order to allow a user to annotate portions of the demonstrations. This is useful to annotate the end of each object-centric subtask in each source demonstration used by MimicGen, as an alternative to implementing subtask termination signals directly in the simulation environment. Some example invocations of the script: + +```sh +# specify the sequence of signals that should be annotated and the dataset images to render on-screen +$ python annotate_subtasks.py --dataset /path/to/demo.hdf5 --signals grasp_1 insert_1 grasp_2 \ + --render_image_names agentview_image robot0_eye_in_hand_image + +# limit annotation to first 2 demos +$ python annotate_subtasks.py --dataset /path/to/demo.hdf5 --signals grasp_1 insert_1 grasp_2 \ + --render_image_names agentview_image robot0_eye_in_hand_image --n 2 + +# limit annotation to demo 2 and 3 +$ python annotate_subtasks.py --dataset /path/to/demo.hdf5 --signals grasp_1 insert_1 grasp_2 \ + --render_image_names agentview_image robot0_eye_in_hand_image --n 2 --start 1 + +# scale up dataset images when rendering to screen by factor of 10 +$ python annotate_subtasks.py --dataset /path/to/demo.hdf5 --signals grasp_1 insert_1 grasp_2 \ + --render_image_names agentview_image robot0_eye_in_hand_image --image_scale 10 +``` + +
+

Note

+ +We provide some utilities to debug your choice of subtask termination signals and offsets -- see the [Debugging Data Generation](https://mimicgen.github.io/docs/tutorials/debugging_datagen.html) tutorial for more information. + +
diff --git a/docs/tutorials/task_visualizations.md b/docs/tutorials/task_visualizations.md new file mode 100644 index 0000000..b815ac2 --- /dev/null +++ b/docs/tutorials/task_visualizations.md @@ -0,0 +1,13 @@ +# Task Visualizations + +We provide a convenience script to write videos for each task's reset distribution at `scripts/get_reset_videos.py` for the robosuite tasks we provide in this repository. Set the `OUTPUT_FOLDER` global variable to the folder where you want to write the videos, and set `DATASET_INFOS` appropriately if you would like to limit the environments visualized. Then run the script. + +The environments are also readily compatible with robosuite visualization scripts such as the [demo_random_action.py](https://github.com/ARISE-Initiative/robosuite/blob/b9d8d3de5e3dfd1724f4a0e6555246c460407daa/robosuite/demos/demo_random_action.py) script and the [make_reset_video.py](https://github.com/ARISE-Initiative/robosuite/blob/b9d8d3de5e3dfd1724f4a0e6555246c460407daa/robosuite/scripts/make_reset_video.py) script, but you will need to modify these files to add a `import mimicgen` line to make sure that `robosuite` can find these environments. + + +
+

Note

+ +You can find task reset visualizations on the [website](https://mimicgen.github.io), but they may look a little different as they were generated with robosuite v1.2. + +
\ No newline at end of file diff --git a/mimicgen_envs/__init__.py b/mimicgen/__init__.py similarity index 90% rename from mimicgen_envs/__init__.py rename to mimicgen/__init__.py index 593ef92..cda3f28 100644 --- a/mimicgen_envs/__init__.py +++ b/mimicgen/__init__.py @@ -1,24 +1,35 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. -__version__ = "0.1.1" +__version__ = "1.0.0" +# try to import all environment interfaces here +try: + from mimicgen.env_interfaces.robosuite import * +except ImportError as e: + print("WARNING: robosuite environment interfaces not imported...") + print("Got error: {}".format(e)) # import tasks to make sure they are added to robosuite task registry -from mimicgen_envs.envs.robosuite.threading import * -from mimicgen_envs.envs.robosuite.coffee import * -from mimicgen_envs.envs.robosuite.three_piece_assembly import * -from mimicgen_envs.envs.robosuite.mug_cleanup import * -from mimicgen_envs.envs.robosuite.stack import * -from mimicgen_envs.envs.robosuite.nut_assembly import * -from mimicgen_envs.envs.robosuite.pick_place import * +try: + from mimicgen.envs.robosuite.threading import * + from mimicgen.envs.robosuite.coffee import * + from mimicgen.envs.robosuite.three_piece_assembly import * + from mimicgen.envs.robosuite.mug_cleanup import * + from mimicgen.envs.robosuite.stack import * + from mimicgen.envs.robosuite.nut_assembly import * + from mimicgen.envs.robosuite.pick_place import * +except ImportError as e: + print("WARNING: robosuite environments not imported...") + print("Got error: {}".format(e)) try: - from mimicgen_envs.envs.robosuite.hammer_cleanup import * - from mimicgen_envs.envs.robosuite.kitchen import * -except ImportError: - print("WARNING: robosuite task zoo environments not imported since robosuite task zoo is not installed!") + from mimicgen.envs.robosuite.hammer_cleanup import * + from mimicgen.envs.robosuite.kitchen import * +except ImportError as e: + print("WARNING: robosuite task zoo environments not imported, possibly because robosuite_task_zoo is not installed...") + print("Got error: {}".format(e)) # stores released dataset links and rollout horizons in global dictionary. # Structure is given below for each type of dataset: diff --git a/mimicgen/configs/__init__.py b/mimicgen/configs/__init__.py new file mode 100644 index 0000000..51a51b5 --- /dev/null +++ b/mimicgen/configs/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +from mimicgen.configs.config import MG_Config, get_all_registered_configs, config_factory +from mimicgen.configs.task_spec import MG_TaskSpec +from mimicgen.configs.robosuite import * \ No newline at end of file diff --git a/mimicgen/configs/config.py b/mimicgen/configs/config.py new file mode 100644 index 0000000..3c9a34c --- /dev/null +++ b/mimicgen/configs/config.py @@ -0,0 +1,153 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Base MG_Config object for mimicgen data generation. +""" +import six +from copy import deepcopy + +import robomimic +from robomimic.config.config import Config + + +# global dictionary for remembering name - class mappings +REGISTERED_CONFIGS = {} + + +def get_all_registered_configs(): + """ + Give access to dictionary of all registered configs for external use. + """ + return deepcopy(REGISTERED_CONFIGS) + + +def config_factory(name, config_type, dic=None): + """ + Creates an instance of a config from the algo name. Optionally pass + a dictionary to instantiate the config from the dictionary. + """ + if (config_type not in REGISTERED_CONFIGS) or (name not in REGISTERED_CONFIGS[config_type]): + raise Exception("Config for name {} and type {} not found. Make sure it is a registered config among: {}".format( + name, config_type, ', '.join(REGISTERED_CONFIGS))) + return REGISTERED_CONFIGS[config_type][name](dict_to_load=dic) + + +class ConfigMeta(type): + """ + Define a metaclass for constructing a config class. + It registers configs into the global registry. + """ + def __new__(meta, name, bases, class_dict): + cls = super(ConfigMeta, meta).__new__(meta, name, bases, class_dict) + if cls.__name__ != "MG_Config": + if cls.TYPE not in REGISTERED_CONFIGS: + REGISTERED_CONFIGS[cls.TYPE] = dict() + REGISTERED_CONFIGS[cls.TYPE][cls.NAME] = cls + return cls + + +@six.add_metaclass(ConfigMeta) +class MG_Config(Config): + def __init__(self, dict_to_load=None): + if dict_to_load is not None: + super(MG_Config, self).__init__(dict_to_load) + return + + super(MG_Config, self).__init__() + + # store name class property in the config (must be implemented by subclasses) + self.name = type(self).NAME + self.type = type(self).TYPE + + self.experiment_config() + self.obs_config() + self.task_config() + + # After init, new keys cannot be added to the config, except under nested + # attributes that have called @do_not_lock_keys + self.lock_keys() + + @property + @classmethod + def NAME(cls): + # must be specified by subclasses + raise NotImplementedError + + @property + @classmethod + def TYPE(cls): + # must be specified by subclasses + raise NotImplementedError + + def experiment_config(self): + """ + This function populates the `config.experiment` attribute of the config, + which has general settings related to the dataset generation (e.g. + which environment, robot, and gripper to use for generation, how many + demonstrations to try collecting, etc). + """ + + # set the name of the experiment - which will be used to name the dataset folder that is generated + self.experiment.name = "demo" + + # settings related to source dataset + self.experiment.source.dataset_path = None # path to source hdf5 dataset + self.experiment.source.filter_key = None # filter key, to select a subset of trajectories in the source hdf5 dataset + self.experiment.source.n = None # if provided, use only the first @n trajectories in source hdf5 dataset + self.experiment.source.start = None # if provided, exclude the first @start trajectories in source hdf5 dataset + + # settings related to data generation + self.experiment.generation.path = None # path where new dataset folder will be created + self.experiment.generation.guarantee = False # whether to keep running data collection until we have @num_trials successful trajectories + self.experiment.generation.keep_failed = True # whether to keep failed trajectories as well + self.experiment.generation.num_trials = 10 # number of attempts to collect new data + + # if True, select a different source demonstration for each subtask during data generation, else + # keep the same one for the entire episode + self.experiment.generation.select_src_per_subtask = False + + # if True, each subtask segment will consist of the first robot pose and the target poses instead of just the target poses. + # Can sometimes help improve data generation quality as the interpolation segment will interpolate to where the robot + # started in the source segment instead of the first target pose. Note that the first subtask segment of each episode + # will always include the first robot pose, regardless of this argument. + self.experiment.generation.transform_first_robot_pose = False + + # if True, each interpolation segment will start from the last target pose in the previous subtask segment, instead + # of the current robot pose. Can sometimes improve data generation quality. + self.experiment.generation.interpolate_from_last_target_pose = True + + # settings related to task used for data generation + self.experiment.task.name = None # if provided, override the env name in env meta to collect data on a different environment from the one in source data + self.experiment.task.robot = None # if provided, override the robot name in env meta to collect data on a different robot from the one in source data + self.experiment.task.gripper = None # if provided, override the gripper in env meta to collect data on a different robot gripper from the one in source data + self.experiment.task.interface = None # if provided, override the environment interface class to use for this task to use a different one from the one in source data + self.experiment.task.interface_type = None # if provided, specify environment interface type (usually one per simulator) to use a different one from the one in source data + + # general settings + self.experiment.max_num_failures = 50 # maximum number of failure demos to save + self.experiment.render_video = True # whether to render some generated demos to video + self.experiment.num_demo_to_render = 50 # maxumum number of demos to render to video + self.experiment.num_fail_demo_to_render = 50 # maxumum number of failure demos to render to video + self.experiment.log_every_n_attempts = 50 # logs important info every N generation attempts + + # random seed for generation + self.experiment.seed = 1 + + def obs_config(self): + """ + This function populates the `config.obs` attribute of the config, which has + setings for which observations to collect during data generation. + """ + self.obs.collect_obs = True # whether to collect observations + self.obs.camera_names = [] # which cameras to render observations from + self.obs.camera_height = 84 # camera height + self.obs.camera_width = 84 # camera width + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. + """ + raise NotImplementedError diff --git a/mimicgen/configs/robosuite.py b/mimicgen/configs/robosuite.py new file mode 100644 index 0000000..b63a664 --- /dev/null +++ b/mimicgen/configs/robosuite.py @@ -0,0 +1,679 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Task configs for robosuite. + +See @Coffee_Config below for an explanation of each parameter. +""" +import mimicgen +from mimicgen.configs.config import MG_Config + + +class Coffee_Config(MG_Config): + """ + Corresponds to robosuite Coffee task and variants. + """ + NAME = "coffee" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + self.task.task_spec.subtask_1 = dict( + # Each subtask involves manipulation with respect to a single object frame. + # This string should specify the object for this subtask. The name should be + # consistent with the "datagen_info" from the environment interface and dataset. + object_ref="coffee_pod", + # The "datagen_info" from the environment and dataset includes binary indicators + # for each subtask of the task at each timestep. This key should correspond + # to the key in "datagen_info" that should be used to infer when this subtask + # is finished (e.g. on a 0 to 1 edge of the binary indicator). Should provide + # None for the final subtask. + subtask_term_signal="grasp", + # if not None, specifies time offsets to be used during data generation when splitting + # a trajectory into subtask segments. On each data generation attempt, an offset is sampled + # and added to the boundary defined by @subtask_term_signal. + subtask_term_offset_range=(5, 10), + # specifies how the source subtask segment should be selected during data generation + # from the set of source human demos + selection_strategy="random", + # optional keyword arguments for the selection strategy function used + selection_strategy_kwargs=None, + # amount of action noise to apply during this subtask + action_noise=0.05, + # number of interpolation steps to bridge previous subtask segment to this one + num_interpolation_steps=5, + # number of additional steps (with constant target pose of beginning of this subtask segment) to + # add to give the robot time to reach the pose needed to carry out this subtask segment + num_fixed_steps=0, + # if True, apply action noise during interpolation phase leading up to this subtask, as + # well as during the execution of this subtask + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_2 = dict( + object_ref="coffee_machine", + # end of final subtask does not need to be detected + subtask_term_signal=None, + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + # allow downstream code to completely replace the task spec from an external config + self.task.task_spec.do_not_lock_keys() + + +class Threading_Config(MG_Config): + """ + Corresponds to robosuite Threading task and variants. + """ + NAME = "threading" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + self.task.task_spec.subtask_1 = dict( + object_ref="needle", + subtask_term_signal="grasp", + subtask_term_offset_range=(5, 10), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_2 = dict( + object_ref="tripod", + subtask_term_signal=None, + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.do_not_lock_keys() + + +class ThreePieceAssembly_Config(MG_Config): + """ + Corresponds to robosuite ThreePieceAssembly task and variants. + """ + NAME = "three_piece_assembly" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + self.task.task_spec.subtask_1 = dict( + object_ref="piece_1", + subtask_term_signal="grasp_1", + subtask_term_offset_range=(5, 10), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_2 = dict( + object_ref="base", + subtask_term_signal="insert_1", + subtask_term_offset_range=(5, 10), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_3 = dict( + object_ref="piece_2", + subtask_term_signal="grasp_2", + subtask_term_offset_range=(5, 10), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_4 = dict( + object_ref="piece_1", + subtask_term_signal=None, + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.do_not_lock_keys() + + +class Square_Config(MG_Config): + """ + Corresponds to robosuite Square task and variants. + """ + NAME = "square" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + self.task.task_spec.subtask_1 = dict( + object_ref="square_nut", + subtask_term_signal="grasp", + subtask_term_offset_range=(10, 20), + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_2 = dict( + object_ref="square_peg", + subtask_term_signal=None, + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.do_not_lock_keys() + + +class Stack_Config(MG_Config): + """ + Corresponds to robosuite Stack task and variants. + """ + NAME = "stack" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + self.task.task_spec.subtask_1 = dict( + object_ref="cubeA", + subtask_term_signal="grasp", + subtask_term_offset_range=(10, 20), + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_2 = dict( + object_ref="cubeB", + subtask_term_signal=None, + subtask_term_offset_range=None, + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.do_not_lock_keys() + + +class StackThree_Config(MG_Config): + """ + Corresponds to robosuite StackThree task and variants. + """ + NAME = "stack_three" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + self.task.task_spec.subtask_1 = dict( + object_ref="cubeA", + subtask_term_signal="grasp_1", + subtask_term_offset_range=(10, 20), + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_2 = dict( + object_ref="cubeB", + subtask_term_signal="stack_1", + subtask_term_offset_range=(10, 20), + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_3 = dict( + object_ref="cubeC", + subtask_term_signal="grasp_2", + subtask_term_offset_range=(10, 20), + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_4 = dict( + object_ref="cubeA", + subtask_term_signal=None, + subtask_term_offset_range=None, + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.do_not_lock_keys() + + +class HammerCleanup_Config(MG_Config): + """ + Corresponds to robosuite HammerCleanup task and variants. + """ + NAME = "hammer_cleanup" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + self.task.task_spec.subtask_1 = dict( + object_ref="drawer", + subtask_term_signal="open", + subtask_term_offset_range=(10, 20), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_2 = dict( + object_ref="hammer", + subtask_term_signal="grasp", + subtask_term_offset_range=(10, 20), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_3 = dict( + object_ref="drawer", + subtask_term_signal=None, + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.do_not_lock_keys() + + +class MugCleanup_Config(MG_Config): + """ + Corresponds to robosuite MugCleanup task and variants. + """ + NAME = "mug_cleanup" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + self.task.task_spec.subtask_1 = dict( + object_ref="drawer", + subtask_term_signal="open", + subtask_term_offset_range=(10, 20), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_2 = dict( + object_ref="object", + subtask_term_signal="grasp", + subtask_term_offset_range=(10, 20), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_3 = dict( + object_ref="drawer", + subtask_term_signal=None, + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.do_not_lock_keys() + + +class NutAssembly_Config(MG_Config): + """ + Corresponds to robosuite NutAssembly task and variants. + """ + NAME = "nut_assembly" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + self.task.task_spec.subtask_1 = dict( + object_ref="square_nut", + subtask_term_signal="grasp_square_nut", + subtask_term_offset_range=(10, 20), + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_2 = dict( + object_ref="square_peg", + subtask_term_signal="insert_square_nut", + subtask_term_offset_range=(10, 20), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_3 = dict( + object_ref="round_nut", + subtask_term_signal="grasp_round_nut", + subtask_term_offset_range=(10, 20), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_4 = dict( + object_ref="round_peg", + subtask_term_signal=None, + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.do_not_lock_keys() + + +class PickPlace_Config(MG_Config): + """ + Corresponds to robosuite PickPlace task and variants. + """ + NAME = "pick_place" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + for i, obj in enumerate(["milk", "cereal", "bread", "can"]): + self.task.task_spec["subtask_{}".format(2 * i + 1)] = dict( + object_ref=obj, + subtask_term_signal="grasp_{}".format(obj), + subtask_term_offset_range=(10, 20), + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + # last subtask does not need subtask termination signal but all others do + self.task.task_spec["subtask_{}".format(2 * i + 2)] = dict( + object_ref=None, + subtask_term_signal=None if (obj == "can") else "place_{}".format(obj), + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.do_not_lock_keys() + + +class Kitchen_Config(MG_Config): + """ + Corresponds to robosuite Kitchen task and variants. + """ + NAME = "kitchen" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + self.task.task_spec.subtask_1 = dict( + object_ref="button", + subtask_term_signal="stove_on", + subtask_term_offset_range=(10, 20), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_2 = dict( + object_ref="pot", + subtask_term_signal="grasp_pot", + subtask_term_offset_range=(10, 20), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_3 = dict( + object_ref="stove", + subtask_term_signal="place_pot_on_stove", + subtask_term_offset_range=(10, 20), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_4 = dict( + object_ref="bread", + subtask_term_signal="grasp_bread", + subtask_term_offset_range=(10, 20), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_5 = dict( + object_ref="pot", + subtask_term_signal="place_bread_in_pot", + subtask_term_offset_range=(10, 20), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_6 = dict( + object_ref="serving_region", + subtask_term_signal="serve", + subtask_term_offset_range=(10, 20), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_7 = dict( + object_ref="button", + subtask_term_signal=None, + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.do_not_lock_keys() + + +class CoffeePreparation_Config(MG_Config): + """ + Corresponds to robosuite CoffeePreparation task and variants. + """ + NAME = "coffee_preparation" + TYPE = "robosuite" + + def task_config(self): + """ + This function populates the `config.task` attribute of the config, + which has settings for each object-centric subtask in a task. Each + dictionary should have kwargs for the @add_subtask method in the + @MG_TaskSpec object. + """ + self.task.task_spec.subtask_1 = dict( + object_ref="mug", + subtask_term_signal="mug_grasp", + subtask_term_offset_range=(5, 10), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_2 = dict( + object_ref="coffee_machine", + subtask_term_signal="mug_place", + subtask_term_offset_range=(5, 10), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_3 = dict( + object_ref="drawer", + subtask_term_signal="drawer_open", + subtask_term_offset_range=(5, 10), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_4 = dict( + object_ref="coffee_pod", + subtask_term_signal="pod_grasp", + subtask_term_offset_range=(5, 10), + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.subtask_5 = dict( + object_ref="coffee_machine", + subtask_term_signal=None, + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ) + self.task.task_spec.do_not_lock_keys() diff --git a/mimicgen/configs/task_spec.py b/mimicgen/configs/task_spec.py new file mode 100644 index 0000000..95e71ec --- /dev/null +++ b/mimicgen/configs/task_spec.py @@ -0,0 +1,139 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Defines task specification objects, which are used to store task-specific settings +for data generation. +""" +import json + +import mimicgen +from mimicgen.datagen.selection_strategy import assert_selection_strategy_exists + +class MG_TaskSpec: + """ + Stores task-specific settings for data generation. Each task is a sequence of + object-centric subtasks, and each subtask stores relevant settings used during + the data generation process. + """ + def __init__(self): + self.spec = [] + + def add_subtask( + self, + object_ref, + subtask_term_signal, + subtask_term_offset_range=None, + selection_strategy="random", + selection_strategy_kwargs=None, + action_noise=0., + num_interpolation_steps=5, + num_fixed_steps=0, + apply_noise_during_interpolation=False, + ): + """ + Add subtask to this task spec. + + Args: + object_ref (str): each subtask involves manipulation with + respect to a single object frame. This string should + specify the object for this subtask. The name + should be consistent with the "datagen_info" from the + environment interface and dataset. + + subtask_term_signal (str or None): the "datagen_info" from the environment + and dataset includes binary indicators for each subtask + of the task at each timestep. This key should correspond + to the key in "datagen_info" that should be used to + infer when this subtask is finished (e.g. on a 0 to 1 + edge of the binary indicator). Should provide None for the final + subtask. + + subtask_term_offset_range (2-tuple): if provided, specifies time offsets to + be used during data generation when splitting a trajectory into + subtask segments. On each data generation attempt, an offset is sampled + and added to the boundary defined by @subtask_term_signal. + + selection_strategy (str): specifies how the source subtask segment should be + selected during data generation from the set of source human demos + + selection_strategy_kwargs (dict or None): optional keyword arguments for the selection + strategy function used + + action_noise (float): amount of action noise to apply during this subtask + + num_interpolation_steps (int): number of interpolation steps to bridge previous subtask segment + to this one + + num_fixed_steps (int): number of additional steps (with constant target pose of beginning of + this subtask segment) to add to give the robot time to reach the pose needed to carry + out this subtask segment + + apply_noise_during_interpolation (bool): if True, apply action noise during interpolation phase + leading up to this subtask, as well as during the execution of this subtask + """ + if subtask_term_offset_range is None: + # corresponds to no offset + subtask_term_offset_range = (0, 0) + assert isinstance(subtask_term_offset_range, tuple) + assert len(subtask_term_offset_range) == 2 + assert subtask_term_offset_range[0] <= subtask_term_offset_range[1] + assert_selection_strategy_exists(selection_strategy) + self.spec.append(dict( + object_ref=object_ref, + subtask_term_signal=subtask_term_signal, + subtask_term_offset_range=subtask_term_offset_range, + selection_strategy=selection_strategy, + selection_strategy_kwargs=selection_strategy_kwargs, + action_noise=action_noise, + num_interpolation_steps=num_interpolation_steps, + num_fixed_steps=num_fixed_steps, + apply_noise_during_interpolation=apply_noise_during_interpolation, + )) + + @classmethod + def from_json(cls, json_string=None, json_dict=None): + """ + Instantiate a TaskSpec object from a json string. This should + be consistent with the output of @serialize. + + Args: + json_string (str): top-level of json has a key per subtask in-order (e.g. + "subtask_1", "subtask_2", "subtask_3") and under each subtask, there should + be an entry for each argument of @add_subtask + + json_dict (dict): optionally directly pass json dict + """ + if json_dict is None: + json_dict = json.loads(json_string) + task_spec = cls() + for subtask_name in json_dict: + if json_dict[subtask_name]["subtask_term_offset_range"] is not None: + json_dict[subtask_name]["subtask_term_offset_range"] = tuple(json_dict[subtask_name]["subtask_term_offset_range"]) + task_spec.add_subtask(**json_dict[subtask_name]) + return task_spec + + def serialize(self): + """ + Return a json string corresponding to this task spec object. Compatible with + @from_json classmethod. + """ + json_dict = dict() + for i, elem in enumerate(self.spec): + json_dict["subtask_{}".format(i + 1)] = elem + return json.dumps(json_dict, indent=4) + + def __len__(self): + return len(self.spec) + + def __getitem__(self, ind): + """Support list-like indexing""" + return self.spec[ind] + + def __iter__(self): + """Support list-like iteration.""" + return iter(self.spec) + + def __repr__(self): + return json.dumps(self.spec, indent=4) diff --git a/mimicgen_envs/utils/__init__.py b/mimicgen/datagen/__init__.py similarity index 57% rename from mimicgen_envs/utils/__init__.py rename to mimicgen/datagen/__init__.py index 981400f..8edfcfd 100644 --- a/mimicgen_envs/utils/__init__.py +++ b/mimicgen/datagen/__init__.py @@ -1,3 +1,3 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. \ No newline at end of file diff --git a/mimicgen/datagen/data_generator.py b/mimicgen/datagen/data_generator.py new file mode 100644 index 0000000..bc07081 --- /dev/null +++ b/mimicgen/datagen/data_generator.py @@ -0,0 +1,409 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Base class for data generator. +""" +import h5py +import sys +import numpy as np + +import mimicgen +import mimicgen.utils.pose_utils as PoseUtils +import mimicgen.utils.file_utils as MG_FileUtils + +from mimicgen.configs.task_spec import MG_TaskSpec +from mimicgen.datagen.datagen_info import DatagenInfo +from mimicgen.datagen.selection_strategy import make_selection_strategy +from mimicgen.datagen.waypoint import WaypointSequence, WaypointTrajectory + + +class DataGenerator(object): + """ + The main data generator object that loads a source dataset, parses it, and + generates new trajectories. + """ + def __init__( + self, + task_spec, + dataset_path, + demo_keys=None, + ): + """ + Args: + task_spec (MG_TaskSpec instance): task specification that will be + used to generate data + dataset_path (str): path to hdf5 dataset to use for generation + demo_keys (list of str): list of demonstration keys to use + in file. If not provided, all demonstration keys will be + used. + """ + assert isinstance(task_spec, MG_TaskSpec) + self.task_spec = task_spec + self.dataset_path = dataset_path + + # sanity check on task spec offset ranges - final subtask should not have any offset randomization + assert self.task_spec[-1]["subtask_term_offset_range"][0] == 0 + assert self.task_spec[-1]["subtask_term_offset_range"][1] == 0 + + # demonstration keys to use from hdf5 as source dataset + if demo_keys is None: + # get all demonstration keys from file + demo_keys = MG_FileUtils.get_all_demos_from_dataset(dataset_path=self.dataset) + self.demo_keys = demo_keys + + # parse source dataset + self._load_dataset(dataset_path=dataset_path, demo_keys=demo_keys) + + def _load_dataset(self, dataset_path, demo_keys): + """ + Load important information from a dataset into internal memory. + """ + print("\nDataGenerator: loading dataset at path {}...".format(dataset_path)) + self.src_dataset_infos, self.src_subtask_indices, self.subtask_names, _ = MG_FileUtils.parse_source_dataset( + dataset_path=dataset_path, + demo_keys=demo_keys, + task_spec=self.task_spec, + ) + print("\nDataGenerator: done loading\n") + + def __repr__(self): + """ + Pretty print this object. + """ + msg = str(self.__class__.__name__) + msg += " (\n\tdataset_path={}\n\tdemo_keys={}\n)".format( + self.dataset_path, + self.demo_keys, + ) + return msg + + def randomize_subtask_boundaries(self): + """ + Apply random offsets to sample subtask boundaries according to the task spec. + Recall that each demonstration is segmented into a set of subtask segments, and the + end index of each subtask can have a random offset. + """ + + # initial subtask start and end indices - shape (N, S, 2) + src_subtask_indices = np.array(self.src_subtask_indices) + + # for each subtask (except last one), sample all end offsets at once for each demonstration + # add them to subtask end indices, and then set them as the start indices of next subtask too + for i in range(src_subtask_indices.shape[1] - 1): + end_offsets = np.random.randint( + low=self.task_spec[i]["subtask_term_offset_range"][0], + high=self.task_spec[i]["subtask_term_offset_range"][1] + 1, + size=src_subtask_indices.shape[0] + ) + src_subtask_indices[:, i, 1] = src_subtask_indices[:, i, 1] + end_offsets + # don't forget to set these as start indices for next subtask too + src_subtask_indices[:, i + 1, 0] = src_subtask_indices[:, i, 1] + + # ensure non-empty subtasks + assert np.all((src_subtask_indices[:, :, 1] - src_subtask_indices[:, :, 0]) > 0), "got empty subtasks!" + + # ensure subtask indices increase (both starts and ends) + assert np.all((src_subtask_indices[:, 1:, :] - src_subtask_indices[:, :-1, :]) > 0), "subtask indices do not strictly increase" + + # ensure subtasks are in order + subtask_inds_flat = src_subtask_indices.reshape(src_subtask_indices.shape[0], -1) + assert np.all((subtask_inds_flat[:, 1:] - subtask_inds_flat[:, :-1]) >= 0), "subtask indices not in order" + + return src_subtask_indices + + def select_source_demo( + self, + eef_pose, + object_pose, + subtask_ind, + src_subtask_inds, + subtask_object_name, + selection_strategy_name, + selection_strategy_kwargs=None, + ): + """ + Helper method to run source subtask segment selection. + + Args: + eef_pose (np.array): current end effector pose + object_pose (np.array): current object pose for this subtask + subtask_ind (int): index of subtask + src_subtask_inds (np.array): start and end indices for subtask segment in source demonstrations of shape (N, 2) + subtask_object_name (str): name of reference object for this subtask + selection_strategy_name (str): name of selection strategy + selection_strategy_kwargs (dict): extra kwargs for running selection strategy + + Returns: + selected_src_demo_ind (int): selected source demo index + """ + if subtask_object_name is None: + # no reference object - only random selection is supported + assert selection_strategy_name == "random" + + # We need to collect the datagen info objects over the timesteps for the subtask segment in each source + # demo, so that it can be used by the selection strategy. + src_subtask_datagen_infos = [] + for i in range(len(self.demo_keys)): + # datagen info over all timesteps of the src trajectory + src_ep_datagen_info = self.src_dataset_infos[i] + + # time indices for subtask + subtask_start_ind = src_subtask_inds[i][0] + subtask_end_ind = src_subtask_inds[i][1] + + # get subtask segment using indices + src_subtask_datagen_infos.append(DatagenInfo( + eef_pose=src_ep_datagen_info.eef_pose[subtask_start_ind : subtask_end_ind], + # only include object pose for relevant object in subtask + object_poses={ subtask_object_name : src_ep_datagen_info.object_poses[subtask_object_name][subtask_start_ind : subtask_end_ind] } if (subtask_object_name is not None) else None, + # subtask termination signal is unused + subtask_term_signals=None, + target_pose=src_ep_datagen_info.target_pose[subtask_start_ind : subtask_end_ind], + gripper_action=src_ep_datagen_info.gripper_action[subtask_start_ind : subtask_end_ind], + )) + + # make selection strategy object + selection_strategy_obj = make_selection_strategy(selection_strategy_name) + + # run selection + if selection_strategy_kwargs is None: + selection_strategy_kwargs = dict() + selected_src_demo_ind = selection_strategy_obj.select_source_demo( + eef_pose=eef_pose, + object_pose=object_pose, + src_subtask_datagen_infos=src_subtask_datagen_infos, + **selection_strategy_kwargs, + ) + + return selected_src_demo_ind + + def generate( + self, + env, + env_interface, + select_src_per_subtask=False, + transform_first_robot_pose=False, + interpolate_from_last_target_pose=True, + render=False, + video_writer=None, + video_skip=5, + camera_names=None, + pause_subtask=False, + ): + """ + Attempt to generate a new demonstration. + + Args: + env (robomimic EnvBase instance): environment to use for data collection + + env_interface (MG_EnvInterface instance): environment interface for some data generation operations + + select_src_per_subtask (bool): if True, select a different source demonstration for each subtask + during data generation, else keep the same one for the entire episode + + transform_first_robot_pose (bool): if True, each subtask segment will consist of the first + robot pose and the target poses instead of just the target poses. Can sometimes help + improve data generation quality as the interpolation segment will interpolate to where + the robot started in the source segment instead of the first target pose. Note that the + first subtask segment of each episode will always include the first robot pose, regardless + of this argument. + + interpolate_from_last_target_pose (bool): if True, each interpolation segment will start from + the last target pose in the previous subtask segment, instead of the current robot pose. Can + sometimes improve data generation quality. + + render (bool): if True, render on-screen + + video_writer (imageio writer): video writer + + video_skip (int): determines rate at which environment frames are written to video + + camera_names (list): determines which camera(s) are used for rendering. Pass more than + one to output a video with multiple camera views concatenated horizontally. + + pause_subtask (bool): if True, pause after every subtask during generation, for + debugging. + + Returns: + results (dict): dictionary with the following items: + initial_state (dict): initial simulator state for the executed trajectory + states (list): simulator state at each timestep + observations (list): observation dictionary at each timestep + datagen_infos (list): datagen_info at each timestep + actions (np.array): action executed at each timestep + success (bool): whether the trajectory successfully solved the task or not + src_demo_inds (list): list of selected source demonstration indices for each subtask + src_demo_labels (np.array): same as @src_demo_inds, but repeated to have a label for each timestep of the trajectory + """ + + # sample new task instance + env.reset() + new_initial_state = env.get_state() + + # sample new subtask boundaries + all_subtask_inds = self.randomize_subtask_boundaries() # shape [N, S, 2], last dim is start and end action lengths + + # some state variables used during generation + selected_src_demo_ind = None + prev_executed_traj = None + + # save generated data in these variables + generated_states = [] + generated_obs = [] + generated_datagen_infos = [] + generated_actions = [] + generated_success = False + generated_src_demo_inds = [] # store selected src demo ind for each subtask in each trajectory + generated_src_demo_labels = [] # like @generated_src_demo_inds, but padded to align with size of @generated_actions + + for subtask_ind in range(len(self.task_spec)): + + # some things only happen on first subtask + is_first_subtask = (subtask_ind == 0) + + # get datagen info in current environment to get required info for selection (e.g. eef pose, object pose) + cur_datagen_info = env_interface.get_datagen_info() + + # name of object for this subtask + subtask_object_name = self.task_spec[subtask_ind]["object_ref"] + + # corresponding current object pose + cur_object_pose = cur_datagen_info.object_poses[subtask_object_name] if (subtask_object_name is not None) else None + + # We need source demonstration selection for the first subtask (always), and possibly for + # other subtasks if @select_src_per_subtask is set. + need_source_demo_selection = (is_first_subtask or select_src_per_subtask) + + # Run source demo selection or use selected demo from previous iteration + if need_source_demo_selection: + selected_src_demo_ind = self.select_source_demo( + eef_pose=cur_datagen_info.eef_pose, + object_pose=cur_object_pose, + subtask_ind=subtask_ind, + src_subtask_inds=all_subtask_inds[:, subtask_ind], + subtask_object_name=subtask_object_name, + selection_strategy_name=self.task_spec[subtask_ind]["selection_strategy"], + selection_strategy_kwargs=self.task_spec[subtask_ind]["selection_strategy_kwargs"], + ) + assert (selected_src_demo_ind is not None) + + # selected subtask segment time indices + selected_src_subtask_inds = all_subtask_inds[selected_src_demo_ind, subtask_ind] + + # get subtask segment, consisting of the sequence of robot eef poses, target poses, gripper actions + src_ep_datagen_info = self.src_dataset_infos[selected_src_demo_ind] + src_subtask_eef_poses = src_ep_datagen_info.eef_pose[selected_src_subtask_inds[0] : selected_src_subtask_inds[1]] + src_subtask_target_poses = src_ep_datagen_info.target_pose[selected_src_subtask_inds[0] : selected_src_subtask_inds[1]] + src_subtask_gripper_actions = src_ep_datagen_info.gripper_action[selected_src_subtask_inds[0] : selected_src_subtask_inds[1]] + + # get reference object pose from source demo + src_subtask_object_pose = src_ep_datagen_info.object_poses[subtask_object_name][selected_src_subtask_inds[0]] if (subtask_object_name is not None) else None + + if is_first_subtask or transform_first_robot_pose: + # Source segment consists of first robot eef pose and the target poses. This ensures that + # we will interpolate to the first robot eef pose in this source segment, instead of the + # first robot target pose. + src_eef_poses = np.concatenate([src_subtask_eef_poses[0:1], src_subtask_target_poses], axis=0) + else: + # Source segment consists of just the target poses. + src_eef_poses = np.array(src_subtask_target_poses) + + # account for extra timestep added to @src_eef_poses + src_subtask_gripper_actions = np.concatenate([src_subtask_gripper_actions[0:1], src_subtask_gripper_actions], axis=0) + + # Transform source demonstration segment using relevant object pose. + if subtask_object_name is not None: + transformed_eef_poses = PoseUtils.transform_source_data_segment_using_object_pose( + obj_pose=cur_object_pose, + src_eef_poses=src_eef_poses, + src_obj_pose=src_subtask_object_pose, + ) + else: + # skip transformation if no reference object is provided + transformed_eef_poses = src_eef_poses + + # We will construct a WaypointTrajectory instance to keep track of robot control targets + # that will be executed and then execute it. + traj_to_execute = WaypointTrajectory() + + if interpolate_from_last_target_pose and (not is_first_subtask): + # Interpolation segment will start from last target pose (which may not have been achieved). + assert prev_executed_traj is not None + last_waypoint = prev_executed_traj.last_waypoint + init_sequence = WaypointSequence(sequence=[last_waypoint]) + else: + # Interpolation segment will start from current robot eef pose. + init_sequence = WaypointSequence.from_poses( + poses=cur_datagen_info.eef_pose[None], + gripper_actions=src_subtask_gripper_actions[0:1], + action_noise=self.task_spec[subtask_ind]["action_noise"], + ) + traj_to_execute.add_waypoint_sequence(init_sequence) + + # Construct trajectory for the transformed segment. + transformed_seq = WaypointSequence.from_poses( + poses=transformed_eef_poses, + gripper_actions=src_subtask_gripper_actions, + action_noise=self.task_spec[subtask_ind]["action_noise"], + ) + transformed_traj = WaypointTrajectory() + transformed_traj.add_waypoint_sequence(transformed_seq) + + # Merge this trajectory into our trajectory using linear interpolation. + # Interpolation will happen from the initial pose (@init_sequence) to the first element of @transformed_seq. + traj_to_execute.merge( + transformed_traj, + num_steps_interp=self.task_spec[subtask_ind]["num_interpolation_steps"], + num_steps_fixed=self.task_spec[subtask_ind]["num_fixed_steps"], + action_noise=(float(self.task_spec[subtask_ind]["apply_noise_during_interpolation"]) * self.task_spec[subtask_ind]["action_noise"]), + ) + + # We initialized @traj_to_execute with a pose to allow @merge to handle linear interpolation + # for us. However, we can safely discard that first waypoint now, and just start by executing + # the rest of the trajectory (interpolation segment and transformed subtask segment). + traj_to_execute.pop_first() + + # Execute the trajectory and collect data. + exec_results = traj_to_execute.execute( + env=env, + env_interface=env_interface, + render=render, + video_writer=video_writer, + video_skip=video_skip, + camera_names=camera_names, + ) + + # check that trajectory is non-empty + if len(exec_results["states"]) > 0: + generated_states += exec_results["states"] + generated_obs += exec_results["observations"] + generated_datagen_infos += exec_results["datagen_infos"] + generated_actions.append(exec_results["actions"]) + generated_success = generated_success or exec_results["success"] + generated_src_demo_inds.append(selected_src_demo_ind) + generated_src_demo_labels.append(selected_src_demo_ind * np.ones((exec_results["actions"].shape[0], 1), dtype=int)) + + # remember last trajectory + prev_executed_traj = traj_to_execute + + if pause_subtask: + input("Pausing after subtask {} execution. Press any key to continue...".format(subtask_ind)) + + # merge numpy arrays + if len(generated_actions) > 0: + generated_actions = np.concatenate(generated_actions, axis=0) + generated_src_demo_labels = np.concatenate(generated_src_demo_labels, axis=0) + + results = dict( + initial_state=new_initial_state, + states=generated_states, + observations=generated_obs, + datagen_infos=generated_datagen_infos, + actions=generated_actions, + success=generated_success, + src_demo_inds=generated_src_demo_inds, + src_demo_labels=generated_src_demo_labels, + ) + return results diff --git a/mimicgen/datagen/datagen_info.py b/mimicgen/datagen/datagen_info.py new file mode 100644 index 0000000..a46b0b2 --- /dev/null +++ b/mimicgen/datagen/datagen_info.py @@ -0,0 +1,78 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Defines structure of information that is needed from an environment for data generation. +""" +import numpy as np +from copy import deepcopy + + +class DatagenInfo(object): + """ + Structure of information needed from an environment for data generation. To allow for + flexibility, not all information must be present. + """ + def __init__( + self, + eef_pose=None, + object_poses=None, + subtask_term_signals=None, + target_pose=None, + gripper_action=None, + ): + """ + Args: + eef_pose (np.array or None): robot end effector poses of shape [..., 4, 4] + object_poses (dict or None): dictionary mapping object name to object poses + of shape [..., 4, 4] + subtask_term_signals (dict or None): dictionary mapping subtask name to a binary + indicator (0 or 1) on whether subtask has been completed. Each value in the + dictionary could be an int, float, or np.array of shape [..., 1]. + target_pose (np.array or None): target end effector poses of shape [..., 4, 4] + gripper_action (np.array or None): gripper actions of shape [..., D] where D + is the dimension of the gripper actuation action for the robot arm + """ + self.eef_pose = None + if eef_pose is not None: + self.eef_pose = np.array(eef_pose) + + self.object_poses = None + if object_poses is not None: + self.object_poses = { k : np.array(object_poses[k]) for k in object_poses } + + self.subtask_term_signals = None + if subtask_term_signals is not None: + self.subtask_term_signals = dict() + for k in subtask_term_signals: + if isinstance(subtask_term_signals[k], float) or isinstance(subtask_term_signals[k], int): + self.subtask_term_signals[k] = subtask_term_signals[k] + else: + # only create numpy array if value is not a single value + self.subtask_term_signals[k] = np.array(subtask_term_signals[k]) + + self.target_pose = None + if target_pose is not None: + self.target_pose = np.array(target_pose) + + self.gripper_action = None + if gripper_action is not None: + self.gripper_action = np.array(gripper_action) + + def to_dict(self): + """ + Convert this instance to a dictionary containing the same information. + """ + ret = dict() + if self.eef_pose is not None: + ret["eef_pose"] = np.array(self.eef_pose) + if self.object_poses is not None: + ret["object_poses"] = deepcopy(self.object_poses) + if self.subtask_term_signals is not None: + ret["subtask_term_signals"] = deepcopy(self.subtask_term_signals) + if self.target_pose is not None: + ret["target_pose"] = np.array(self.target_pose) + if self.gripper_action is not None: + ret["gripper_action"] = np.array(self.gripper_action) + return ret diff --git a/mimicgen/datagen/selection_strategy.py b/mimicgen/datagen/selection_strategy.py new file mode 100644 index 0000000..68d81e9 --- /dev/null +++ b/mimicgen/datagen/selection_strategy.py @@ -0,0 +1,306 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Selection strategies used by MimicGen to select subtask segments from +source human demonstrations. +""" +import abc # for abstract base class definitions +import six # preserve metaclass compatibility between python 2 and 3 + +import numpy as np + +import mimicgen.utils.pose_utils as PoseUtils + + +# Global dictionary for remembering name to class mappings. +REGISTERED_SELECTION_STRATEGIES = {} + + +def make_selection_strategy(name, *args, **kwargs): + """ + Creates an instance of a selection strategy class, specified by @name, + which is used to look it up in the registry. + """ + assert_selection_strategy_exists(name) + return REGISTERED_SELECTION_STRATEGIES[name](*args, **kwargs) + + +def register_selection_strategy(cls): + """ + Register selection strategy class into global registry. + """ + ignore_classes = ["MG_SelectionStrategy"] + if cls.__name__ not in ignore_classes: + REGISTERED_SELECTION_STRATEGIES[cls.NAME] = cls + + +def assert_selection_strategy_exists(name): + """ + Allow easy way to check if selection strategy exists. + """ + if name not in REGISTERED_SELECTION_STRATEGIES: + raise Exception("assert_selection_strategy_exists: name {} not found. Make sure it is a registered selection strategy among {}".format(", ".join(REGISTERED_SELECTION_STRATEGIES))) + + +class MG_SelectionStrategyMeta(type): + """ + This metaclass adds selection strategy classes into the global registry. + """ + def __new__(meta, name, bases, class_dict): + cls = super(MG_SelectionStrategyMeta, meta).__new__(meta, name, bases, class_dict) + register_selection_strategy(cls) + return cls + + +@six.add_metaclass(MG_SelectionStrategyMeta) +class MG_SelectionStrategy(object): + """ + Defines methods and functions for selection strategies to implement. + """ + def __init__(self): + pass + + @property + @classmethod + def NAME(self): + """ + This name (str) will be used to register the selection strategy class in the global + registry. + """ + raise NotImplementedError + + @abc.abstractmethod + def select_source_demo( + self, + eef_pose, + object_pose, + src_subtask_datagen_infos, + ): + """ + Selects source demonstration index using the current robot pose, relevant object pose + for the current subtask, and relevant information from the source demonstrations for the + current subtask. + + Args: + eef_pose (np.array): current 4x4 eef pose + object_pose (np.array): current 4x4 object pose, for the object in this subtask + src_subtask_datagen_infos (list): DatagenInfo instance for the relevant subtask segment + in the source demonstrations + + Returns: + source_demo_ind (int): index of source demonstration - indicates which source subtask segment to use + """ + raise NotImplementedError + + +class RandomStrategy(MG_SelectionStrategy): + """ + Pick source demonstration randomly. + """ + + # name for registering this class into registry + NAME = "random" + + def select_source_demo( + self, + eef_pose, + object_pose, + src_subtask_datagen_infos, + ): + """ + Selects source demonstration index using the current robot pose, relevant object pose + for the current subtask, and relevant information from the source demonstrations for the + current subtask. + + Args: + eef_pose (np.array): current 4x4 eef pose + object_pose (np.array): current 4x4 object pose, for the object in this subtask + src_subtask_datagen_infos (list): DatagenInfo instance for the relevant subtask segment + in the source demonstrations + + Returns: + source_demo_ind (int): index of source demonstration - indicates which source subtask segment to use + """ + + # random selection + n_src_demo = len(src_subtask_datagen_infos) + return np.random.randint(0, n_src_demo) + + +class NearestNeighborObjectStrategy(MG_SelectionStrategy): + """ + Pick source demonstration to be the one with the closest object pose to the object + in the current scene. + """ + + # name for registering this class into registry + NAME = "nearest_neighbor_object" + + def select_source_demo( + self, + eef_pose, + object_pose, + src_subtask_datagen_infos, + pos_weight=1., + rot_weight=1., + nn_k=3, + ): + """ + Selects source demonstration index using the current robot pose, relevant object pose + for the current subtask, and relevant information from the source demonstrations for the + current subtask. + + Args: + eef_pose (np.array): current 4x4 eef pose + object_pose (np.array): current 4x4 object pose, for the object in this subtask + src_subtask_datagen_infos (list): DatagenInfo instance for the relevant subtask segment + in the source demonstrations + pos_weight (float): weight on position for minimizing pose distance + rot_weight (float): weight on rotation for minimizing pose distance + nn_k (int): pick source demo index uniformly at randomly from the top @nn_k nearest neighbors + + Returns: + source_demo_ind (int): index of source demonstration - indicates which source subtask segment to use + """ + + # collect object poses from start of subtask source segments into array of shape [N, 4, 4] + src_object_poses = [] + for di in src_subtask_datagen_infos: + src_obj_pose = list(di.object_poses.values()) + assert len(src_obj_pose) == 1 + # use object pose at start of subtask segment + src_object_poses.append(src_obj_pose[0][0]) + src_object_poses = np.array(src_object_poses) + + # split into positions and rotations + all_src_obj_pos, all_src_obj_rot = PoseUtils.unmake_pose(src_object_poses) + obj_pos, obj_rot = PoseUtils.unmake_pose(object_pose) + + # prepare for broadcasting + obj_pos = obj_pos.reshape(-1, 3) + obj_rot_T = obj_rot.T.reshape(-1, 3, 3) + + # pos dist is just L2 between positions + pos_dists = np.sqrt(((all_src_obj_pos - obj_pos) ** 2).sum(axis=-1)) + + # get angle (in axis-angle representation of delta rotation matrix) using the following formula + # (see http://www.boris-belousov.net/2016/12/01/quat-dist/) + + # batched matrix mult, [N, 3, 3] x [1, 3, 3] -> [N, 3, 3] + delta_R = np.matmul(all_src_obj_rot, obj_rot_T) + arc_cos_in = (np.trace(delta_R, axis1=-2, axis2=-1) - 1.) / 2. + arc_cos_in = np.clip(arc_cos_in, -1., 1.) # clip for numerical stability + rot_dists = np.arccos(arc_cos_in) + + # weight distances with coefficients + dists_to_minimize = pos_weight * pos_dists + rot_weight * rot_dists + + # clip top-k parameter to max possible value + nn_k = min(nn_k, len(dists_to_minimize)) + + # return one of the top-K nearest neighbors uniformly at random + rand_k = np.random.randint(0, nn_k) + top_k_neighbors_in_order = np.argsort(dists_to_minimize)[:nn_k] + return top_k_neighbors_in_order[rand_k] + + +class NearestNeighborRobotDistanceStrategy(MG_SelectionStrategy): + """ + Pick source demonstration to be the one that minimizes the distance the robot + end effector will need to travel from the current pose to the first pose + in the transformed segment. + """ + + # name for registering this class into registry + NAME = "nearest_neighbor_robot_distance" + + def select_source_demo( + self, + eef_pose, + object_pose, + src_subtask_datagen_infos, + pos_weight=1., + rot_weight=1., + nn_k=3, + ): + """ + Selects source demonstration index using the current robot pose, relevant object pose + for the current subtask, and relevant information from the source demonstrations for the + current subtask. + + Args: + eef_pose (np.array): current 4x4 eef pose + object_pose (np.array): current 4x4 object pose, for the object in this subtask + src_subtask_datagen_infos (list): DatagenInfo instance for the relevant subtask segment + in the source demonstrations + pos_weight (float): weight on position for minimizing pose distance + rot_weight (float): weight on rotation for minimizing pose distance + nn_k (int): pick source demo index uniformly at randomly from the top @nn_k nearest neighbors + + Returns: + source_demo_ind (int): index of source demonstration - indicates which source subtask segment to use + """ + + # collect eef and object poses from start of subtask source segments into arrays of shape [N, 4, 4] + src_eef_poses = [] + src_object_poses = [] + for di in src_subtask_datagen_infos: + # use eef pose at start of subtask segment + src_eef_poses.append(di.eef_pose[0]) + # use object pose at start of subtask segment + src_obj_pose = list(di.object_poses.values()) + assert len(src_obj_pose) == 1 + src_object_poses.append(src_obj_pose[0][0]) + src_eef_poses = np.array(src_eef_poses) + src_object_poses = np.array(src_object_poses) + + # Get source eef poses with respect to object frames. + # note: frame A is world, frame B is object + src_object_poses_inv = PoseUtils.pose_inv(src_object_poses) + src_eef_poses_in_obj = PoseUtils.pose_in_A_to_pose_in_B( + pose_in_A=src_eef_poses, + pose_A_in_B=src_object_poses_inv, + ) + + # Use this to find the first pose for the transformed subtask segment for each source demo. + # Note this is the same logic used in PoseUtils.transform_source_data_segment_using_object_pose + transformed_eef_poses = PoseUtils.pose_in_A_to_pose_in_B( + pose_in_A=src_eef_poses_in_obj, + pose_A_in_B=object_pose, + ) + + # split into positions and rotations + all_transformed_eef_pos, all_transformed_eef_rot = PoseUtils.unmake_pose(transformed_eef_poses) + eef_pos, eef_rot = PoseUtils.unmake_pose(eef_pose) + + # now measure distance from each of these transformed eef poses to our current eef pose + # and choose the source demo that minimizes this distance + + # prepare for broadcasting + eef_pos = eef_pos.reshape(-1, 3) + eef_rot_T = eef_rot.T.reshape(-1, 3, 3) + + # pos dist is just L2 between positions + pos_dists = np.sqrt(((all_transformed_eef_pos - eef_pos) ** 2).sum(axis=-1)) + + # get angle (in axis-angle representation of delta rotation matrix) using the following formula + # (see http://www.boris-belousov.net/2016/12/01/quat-dist/) + + # batched matrix mult, [N, 3, 3] x [1, 3, 3] -> [N, 3, 3] + delta_R = np.matmul(all_transformed_eef_rot, eef_rot_T) + arc_cos_in = (np.trace(delta_R, axis1=-2, axis2=-1) - 1.) / 2. + arc_cos_in = np.clip(arc_cos_in, -1., 1.) # clip for numerical stability + rot_dists = np.arccos(arc_cos_in) + + # weight distances with coefficients + dists_to_minimize = pos_weight * pos_dists + rot_weight * rot_dists + + # clip top-k parameter to max possible value + nn_k = min(nn_k, len(dists_to_minimize)) + + # return one of the top-K nearest neighbors uniformly at random + rand_k = np.random.randint(0, nn_k) + top_k_neighbors_in_order = np.argsort(dists_to_minimize)[:nn_k] + return top_k_neighbors_in_order[rand_k] diff --git a/mimicgen/datagen/waypoint.py b/mimicgen/datagen/waypoint.py new file mode 100644 index 0000000..fb1c534 --- /dev/null +++ b/mimicgen/datagen/waypoint.py @@ -0,0 +1,415 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +A collection of classes used to represent waypoints and trajectories. +""" +import json +import numpy as np +from copy import deepcopy + +import mimicgen +import mimicgen.utils.pose_utils as PoseUtils + + +class Waypoint(object): + """ + Represents a single desired 6-DoF waypoint, along with corresponding gripper actuation for this point. + """ + def __init__(self, pose, gripper_action, noise=None): + """ + Args: + pose (np.array): 4x4 pose target for robot controller + gripper_action (np.array): gripper action for robot controller + noise (float or None): action noise amplitude to apply during execution at this timestep + (for arm actions, not gripper actions) + """ + self.pose = np.array(pose) + self.gripper_action = np.array(gripper_action) + self.noise = noise + assert len(self.gripper_action.shape) == 1 + + +class WaypointSequence(object): + """ + Represents a sequence of Waypoint objects. + """ + def __init__(self, sequence=None): + """ + Args: + sequence (list or None): if provided, should be an list of Waypoint objects + """ + if sequence is None: + self.sequence = [] + else: + for waypoint in sequence: + assert isinstance(waypoint, Waypoint) + self.sequence = deepcopy(sequence) + + @classmethod + def from_poses(cls, poses, gripper_actions, action_noise): + """ + Instantiate a WaypointSequence object given a sequence of poses, + gripper actions, and action noise. + + Args: + poses (np.array): sequence of pose matrices of shape (T, 4, 4) + gripper_actions (np.array): sequence of gripper actions + that should be applied at each timestep of shape (T, D). + action_noise (float or np.array): sequence of action noise + magnitudes that should be applied at each timestep. If a + single float is provided, the noise magnitude will be + constant over the trajectory. + """ + assert isinstance(action_noise, float) or isinstance(action_noise, np.ndarray) + + # handle scalar to numpy array conversion + num_timesteps = poses.shape[0] + if isinstance(action_noise, float): + action_noise = action_noise * np.ones((num_timesteps, 1)) + action_noise = action_noise.reshape(-1, 1) + + # make WaypointSequence instance + sequence = [ + Waypoint( + pose=poses[t], + gripper_action=gripper_actions[t], + noise=action_noise[t, 0], + ) + for t in range(num_timesteps) + ] + return cls(sequence=sequence) + + def __len__(self): + # length of sequence + return len(self.sequence) + + def __getitem__(self, ind): + """ + Returns waypoint at index. + + Returns: + waypoint (Waypoint instance) + """ + return self.sequence[ind] + + def __add__(self, other): + """ + Defines addition (concatenation) of sequences + """ + return WaypointSequence(sequence=(self.sequence + other.sequence)) + + @property + def last_waypoint(self): + """ + Return last waypoint in sequence. + + Returns: + waypoint (Waypoint instance) + """ + return deepcopy(self.sequence[-1]) + + def split(self, ind): + """ + Splits this sequence into 2 pieces, the part up to time index @ind, and the + rest. Returns 2 WaypointSequence objects. + """ + seq_1 = self.sequence[:ind] + seq_2 = self.sequence[ind:] + return WaypointSequence(sequence=seq_1), WaypointSequence(sequence=seq_2) + + +class WaypointTrajectory(object): + """ + A sequence of WaypointSequence objects that corresponds to a full 6-DoF trajectory. + """ + def __init__(self): + self.waypoint_sequences = [] + + def __len__(self): + # sum up length of all waypoint sequences + return sum(len(s) for s in self.waypoint_sequences) + + def __getitem__(self, ind): + """ + Returns waypoint at time index. + + Returns: + waypoint (Waypoint instance) + """ + assert len(self.waypoint_sequences) > 0 + assert (ind >= 0) and (ind < len(self)) + + # find correct waypoint sequence we should index + end_ind = 0 + for seq_ind in range(len(self.waypoint_sequences)): + start_ind = end_ind + end_ind += len(self.waypoint_sequences[seq_ind]) + if (ind >= start_ind) and (ind < end_ind): + break + + # index within waypoint sequence + return self.waypoint_sequences[seq_ind][ind - start_ind] + + @property + def last_waypoint(self): + """ + Return last waypoint in sequence. + + Returns: + waypoint (Waypoint instance) + """ + return self.waypoint_sequences[-1].last_waypoint + + def add_waypoint_sequence(self, sequence): + """ + Directly append sequence to list (no interpolation). + + Args: + sequence (WaypointSequence instance): sequence to add + """ + assert isinstance(sequence, WaypointSequence) + self.waypoint_sequences.append(sequence) + + def add_waypoint_sequence_for_target_pose( + self, + pose, + gripper_action, + num_steps, + skip_interpolation=False, + action_noise=0., + ): + """ + Adds a new waypoint sequence corresponding to a desired target pose. A new WaypointSequence + will be constructed consisting of @num_steps intermediate Waypoint objects. These can either + be constructed with linear interpolation from the last waypoint (default) or be a + constant set of target poses (set @skip_interpolation to True). + + Args: + pose (np.array): 4x4 target pose + + gripper_action (np.array): value for gripper action + + num_steps (int): number of action steps when trying to reach this waypoint. Will + add intermediate linearly interpolated points between the last pose on this trajectory + and the target pose, so that the total number of steps is @num_steps. + + skip_interpolation (bool): if True, keep the target pose fixed and repeat it @num_steps + times instead of using linearly interpolated targets. + + action_noise (float): scale of random gaussian noise to add during action execution (e.g. + when @execute is called) + """ + if (len(self.waypoint_sequences) == 0): + assert skip_interpolation, "cannot interpolate since this is the first waypoint sequence" + + if skip_interpolation: + # repeat the target @num_steps times + assert num_steps is not None + poses = np.array([pose for _ in range(num_steps)]) + gripper_actions = np.array([[gripper] for _ in range(num_steps)]) + else: + # linearly interpolate between the last pose and the new waypoint + last_waypoint = self.last_waypoint + poses, num_steps_2 = PoseUtils.interpolate_poses( + pose_1=last_waypoint.pose, + pose_2=pose, + num_steps=num_steps, + ) + assert num_steps == num_steps_2 + gripper_actions = np.array([gripper_action for _ in range(num_steps + 2)]) + # make sure to skip the first element of the new path, which already exists on the current trajectory path + poses = poses[1:] + gripper_actions = gripper_actions[1:] + + # add waypoint sequence for this set of poses + sequence = WaypointSequence.from_poses( + poses=poses, + gripper_actions=gripper_actions, + action_noise=action_noise, + ) + self.add_waypoint_sequence(sequence) + + def pop_first(self): + """ + Removes first waypoint in first waypoint sequence and returns it. If the first waypoint + sequence is now empty, it is also removed. + + Returns: + waypoint (Waypoint instance) + """ + first, rest = self.waypoint_sequences[0].split(1) + if len(rest) == 0: + # remove empty waypoint sequence + self.waypoint_sequences = self.waypoint_sequences[1:] + else: + # update first waypoint sequence + self.waypoint_sequences[0] = rest + return first + + def merge( + self, + other, + num_steps_interp=None, + num_steps_fixed=None, + action_noise=0., + ): + """ + Merge this trajectory with another (@other). + + Args: + other (WaypointTrajectory object): the other trajectory to merge into this one + + num_steps_interp (int or None): if not None, add a waypoint sequence that interpolates + between the end of the current trajectory and the start of @other + + num_steps_fixed (int or None): if not None, add a waypoint sequence that has constant + target poses corresponding to the first target pose in @other + + action_noise (float): noise to use during the interpolation segment + """ + need_interp = (num_steps_interp is not None) and (num_steps_interp > 0) + need_fixed = (num_steps_fixed is not None) and (num_steps_fixed > 0) + use_interpolation_segment = (need_interp or need_fixed) + + if use_interpolation_segment: + # pop first element of other trajectory + other_first = other.pop_first() + + # Get first target pose of other trajectory. + # The interpolated segment will include this first element as its last point. + target_for_interpolation = other_first[0] + + if need_interp: + # interpolation segment + self.add_waypoint_sequence_for_target_pose( + pose=target_for_interpolation.pose, + gripper_action=target_for_interpolation.gripper_action, + num_steps=num_steps_interp, + action_noise=action_noise, + skip_interpolation=False, + ) + + if need_fixed: + # segment of constant target poses equal to @other's first target pose + + # account for the fact that we pop'd the first element of @other in anticipation of an interpolation segment + num_steps_fixed_to_use = num_steps_fixed if need_interp else (num_steps_fixed + 1) + self.add_waypoint_sequence_for_target_pose( + pose=target_for_interpolation.pose, + gripper_action=target_for_interpolation.gripper_action, + num_steps=num_steps_fixed_to_use, + action_noise=action_noise, + skip_interpolation=True, + ) + + # make sure to preserve noise from first element of other trajectory + self.waypoint_sequences[-1][-1].noise = target_for_interpolation.noise + + # concatenate the trajectories + self.waypoint_sequences += other.waypoint_sequences + + def execute( + self, + env, + env_interface, + render=False, + video_writer=None, + video_skip=5, + camera_names=None, + ): + """ + Main function to execute the trajectory. Will use env_interface.target_pose_to_action to + convert each target pose at each waypoint to an action command, and pass that along to + env.step. + + Args: + env (robomimic EnvBase instance): environment to use for executing trajectory + env_interface (MG_EnvInterface instance): environment interface for executing trajectory + render (bool): if True, render on-screen + video_writer (imageio writer): video writer + video_skip (int): determines rate at which environment frames are written to video + camera_names (list): determines which camera(s) are used for rendering. Pass more than + one to output a video with multiple camera views concatenated horizontally. + + Returns: + results (dict): dictionary with the following items for the executed trajectory: + states (list): simulator state at each timestep + observations (list): observation dictionary at each timestep + datagen_infos (list): datagen_info at each timestep + actions (list): action executed at each timestep + success (bool): whether the trajectory successfully solved the task or not + """ + write_video = (video_writer is not None) + video_count = 0 + + states = [] + actions = [] + observations = [] + datagen_infos = [] + success = { k: False for k in env.is_success() } # success metrics + + # iterate over waypoint sequences + for seq in self.waypoint_sequences: + + # iterate over waypoints in each sequence + for j in range(len(seq)): + + # on-screen render + if render: + env.render(mode="human", camera_name=camera_names[0]) + + # video render + if write_video: + if video_count % video_skip == 0: + video_img = [] + for cam_name in camera_names: + video_img.append(env.render(mode="rgb_array", height=512, width=512, camera_name=cam_name)) + video_img = np.concatenate(video_img, axis=1) # concatenate horizontally + video_writer.append_data(video_img) + video_count += 1 + + # current waypoint + waypoint = seq[j] + + # current state and obs + state = env.get_state()["states"] + obs = env.get_observation() + + # convert target pose to arm action + action_pose = env_interface.target_pose_to_action(target_pose=waypoint.pose) + + # maybe add noise to action + if waypoint.noise is not None: + action_pose += waypoint.noise * np.random.randn(*action_pose.shape) + action_pose = np.clip(action_pose, -1., 1.) + + # add in gripper action + play_action = np.concatenate([action_pose, waypoint.gripper_action], axis=0) + + # store datagen info too + datagen_info = env_interface.get_datagen_info(action=play_action) + + # step environment + env.step(play_action) + + # collect data + states.append(state) + play_action_record = play_action + actions.append(play_action_record) + observations.append(obs) + datagen_infos.append(datagen_info) + + cur_success_metrics = env.is_success() + for k in success: + success[k] = success[k] or cur_success_metrics[k] + + results = dict( + states=states, + observations=observations, + datagen_infos=datagen_infos, + actions=np.array(actions), + success=bool(success["task"]), + ) + return results diff --git a/mimicgen_envs/envs/__init__.py b/mimicgen/env_interfaces/__init__.py similarity index 57% rename from mimicgen_envs/envs/__init__.py rename to mimicgen/env_interfaces/__init__.py index 981400f..8edfcfd 100644 --- a/mimicgen_envs/envs/__init__.py +++ b/mimicgen/env_interfaces/__init__.py @@ -1,3 +1,3 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. \ No newline at end of file diff --git a/mimicgen/env_interfaces/base.py b/mimicgen/env_interfaces/base.py new file mode 100644 index 0000000..0cf638e --- /dev/null +++ b/mimicgen/env_interfaces/base.py @@ -0,0 +1,209 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Base class for environment interfaces used by MimicGen. Defines a set of +functions that should be implemented for every set of environments, and +a global registry. +""" +import abc # for abstract base class definitions +import six # preserve metaclass compatibility between python 2 and 3 + +import numpy as np + +from mimicgen.datagen.datagen_info import DatagenInfo + + +# Global dictionary for remembering name - class mappings. +# +# Organization: +# interface_type (str) +# class_name (str) +# class object +REGISTERED_ENV_INTERFACES = {} + + +def make_interface(name, interface_type, *args, **kwargs): + """ + Creates an instance of a env interface. Make sure to pass any other needed arguments. + """ + if interface_type not in REGISTERED_ENV_INTERFACES: + raise Exception("make_interface: interface type {} not found. Make sure it is a registered interface type among: {}".format(interface_type, ", ".join(REGISTERED_ENV_INTERFACES))) + if name not in REGISTERED_ENV_INTERFACES[interface_type]: + raise Exception("make_interface: interface name {} not found. Make sure it is a registered interface name among: {}".format(name, ', '.join(REGISTERED_ENV_INTERFACES[interface_type]))) + return REGISTERED_ENV_INTERFACES[interface_type][name](*args, **kwargs) + + +def register_env_interface(cls): + """ + Register environment interface class into global registry. + """ + ignore_classes = ["MG_EnvInterface"] + if cls.__name__ not in ignore_classes: + if cls.INTERFACE_TYPE not in REGISTERED_ENV_INTERFACES: + REGISTERED_ENV_INTERFACES[cls.INTERFACE_TYPE] = dict() + REGISTERED_ENV_INTERFACES[cls.INTERFACE_TYPE][cls.__name__] = cls + + +class MG_EnvInterfaceMeta(type): + """ + This metaclass adds env interface classes into the global registry. + """ + def __new__(meta, name, bases, class_dict): + cls = super(MG_EnvInterfaceMeta, meta).__new__(meta, name, bases, class_dict) + register_env_interface(cls) + return cls + + +@six.add_metaclass(MG_EnvInterfaceMeta) +class MG_EnvInterface(object): + """ + Environment interface API that MimicGen environment interfaces should conform to. + """ + def __init__(self, env): + """ + Args: + env: environment object + """ + self.env = env + self.interface_type = type(self).INTERFACE_TYPE + + def __repr__(self): + """ + Pretty-print env description. + """ + return self.__class__.__name__ + + """ + These should be filled out by simulator subclasses (e.g. robosuite). + """ + @property + @classmethod + def INTERFACE_TYPE(self): + """ + Returns string corresponding to interface type. This is used to group + all subclasses together in the interface registry (for example, all robosuite + interfaces) and helps avoid name conflicts. + """ + raise NotImplementedError + + @abc.abstractmethod + def get_robot_eef_pose(self): + """ + Get current robot end effector pose. Should be the same frame as used by the robot end-effector controller. + + Returns: + pose (np.array): 4x4 eef pose matrix + """ + raise NotImplementedError + + @abc.abstractmethod + def target_pose_to_action(self, target_pose, relative=True): + """ + Takes a target pose for the end effector controller and returns an action + (usually a normalized delta pose action) to try and achieve that target pose. + + Args: + target_pose (np.array): 4x4 target eef pose + relative (bool): if True, use relative pose actions, else absolute pose actions + + Returns: + action (np.array): action compatible with env.step (minus gripper actuation) + """ + raise NotImplementedError + + @abc.abstractmethod + def action_to_target_pose(self, action, relative=True): + """ + Converts action (compatible with env.step) to a target pose for the end effector controller. + Inverse of @target_pose_to_action. Usually used to infer a sequence of target controller poses + from a demonstration trajectory using the recorded actions. + + Args: + action (np.array): environment action + relative (bool): if True, use relative pose actions, else absolute pose actions + + Returns: + target_pose (np.array): 4x4 target eef pose that @action corresponds to + """ + raise NotImplementedError + + @abc.abstractmethod + def action_to_gripper_action(self, action): + """ + Extracts the gripper actuation part of an action (compatible with env.step). + + Args: + action (np.array): environment action + + Returns: + gripper_action (np.array): subset of environment action for gripper actuation + """ + raise NotImplementedError + + """ + These should be filled out by each simulation domain (e.g. nut assembly, coffee). + """ + @abc.abstractmethod + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + raise NotImplementedError + + @abc.abstractmethod + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + raise NotImplementedError + + """ + This method can be left as-is in most cases, as it calls other implemented methods to provide a + DatagenInfo object. + """ + def get_datagen_info(self, action=None): + """ + Get information needed for data generation, at the current + timestep of simulation. If @action is provided, it will be used to + compute the target eef pose for the controller, otherwise that + will be excluded. + + Returns: + datagen_info (DatagenInfo instance) + """ + + # current eef pose + eef_pose = self.get_robot_eef_pose() + + # object poses + object_poses = self.get_object_poses() + + # subtask termination signals + subtask_term_signals = self.get_subtask_term_signals() + + # these must be extracted from provided action + target_pose = None + gripper_action = None + if action is not None: + target_pose = self.action_to_target_pose(action=action, relative=True) + gripper_action = self.action_to_gripper_action(action=action) + + datagen_info = DatagenInfo( + eef_pose=eef_pose, + object_poses=object_poses, + subtask_term_signals=subtask_term_signals, + target_pose=target_pose, + gripper_action=gripper_action, + ) + return datagen_info diff --git a/mimicgen/env_interfaces/robosuite.py b/mimicgen/env_interfaces/robosuite.py new file mode 100644 index 0000000..a0d77fd --- /dev/null +++ b/mimicgen/env_interfaces/robosuite.py @@ -0,0 +1,749 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +MimicGen environment interface classes for basic robosuite environments. +""" +import numpy as np + +import robosuite +import robosuite.utils.transform_utils as T + +import mimicgen.utils.pose_utils as PoseUtils +from mimicgen.env_interfaces.base import MG_EnvInterface + + +class RobosuiteInterface(MG_EnvInterface): + """ + MimicGen environment interface base class for basic robosuite environments. + """ + + # Note: base simulator interface class must fill out interface type as a class property + INTERFACE_TYPE = "robosuite" + + def get_robot_eef_pose(self): + """ + Get current robot end effector pose. Should be the same frame as used by the robot end-effector controller. + + Returns: + pose (np.array): 4x4 eef pose matrix + """ + + # OSC control frame is a MuJoCo site - just retrieve its current pose + return self.get_object_pose( + obj_name=self.env.robots[0].controller.eef_name, + obj_type="site", + ) + + def target_pose_to_action(self, target_pose, relative=True): + """ + Takes a target pose for the end effector controller and returns an action + (usually a normalized delta pose action) to try and achieve that target pose. + + Args: + target_pose (np.array): 4x4 target eef pose + relative (bool): if True, use relative pose actions, else absolute pose actions + + Returns: + action (np.array): action compatible with env.step (minus gripper actuation) + """ + + # version check for robosuite - must be v1.2+, so that we're using the correct controller convention + assert (robosuite.__version__.split(".")[0] == "1") + assert (robosuite.__version__.split(".")[1] >= "2") + + # target position and rotation + target_pos, target_rot = PoseUtils.unmake_pose(target_pose) + + # current position and rotation + curr_pose = self.get_robot_eef_pose() + curr_pos, curr_rot = PoseUtils.unmake_pose(curr_pose) + + # get maximum position and rotation action bounds + max_dpos = self.env.robots[0].controller.output_max[0] + max_drot = self.env.robots[0].controller.output_max[3] + + if relative: + # normalized delta position action + delta_position = target_pos - curr_pos + delta_position = np.clip(delta_position / max_dpos, -1., 1.) + + # normalized delta rotation action + delta_rot_mat = target_rot.dot(curr_rot.T) + delta_quat = T.mat2quat(delta_rot_mat) + delta_rotation = T.quat2axisangle(delta_quat) + delta_rotation = np.clip(delta_rotation / max_drot, -1., 1.) + return np.concatenate([delta_position, delta_rotation]) + + # absolute position and rotation action + target_quat = T.mat2quat(target_rot) + abs_rotation = T.quat2axisangle(target_quat) + return np.concatenate([target_pos, abs_rotation]) + + def action_to_target_pose(self, action, relative=True): + """ + Converts action (compatible with env.step) to a target pose for the end effector controller. + Inverse of @target_pose_to_action. Usually used to infer a sequence of target controller poses + from a demonstration trajectory using the recorded actions. + + Args: + action (np.array): environment action + relative (bool): if True, use relative pose actions, else absolute pose actions + + Returns: + target_pose (np.array): 4x4 target eef pose that @action corresponds to + """ + + # version check for robosuite - must be v1.2+, so that we're using the correct controller convention + assert (robosuite.__version__.split(".")[0] == "1") + assert (robosuite.__version__.split(".")[1] >= "2") + + if (not relative): + # convert absolute action to absolute pose + target_pos = action[:3] + target_quat = T.axisangle2quat(action[3:6]) + target_rot = T.quat2mat(target_quat) + else: + # get maximum position and rotation action bounds + max_dpos = self.env.robots[0].controller.output_max[0] + max_drot = self.env.robots[0].controller.output_max[3] + + # unscale actions + delta_position = action[:3] * max_dpos + delta_rotation = action[3:6] * max_drot + + # current position and rotation + curr_pose = self.get_robot_eef_pose() + curr_pos, curr_rot = PoseUtils.unmake_pose(curr_pose) + + # get pose target + target_pos = curr_pos + delta_position + delta_quat = T.axisangle2quat(delta_rotation) + delta_rot_mat = T.quat2mat(delta_quat) + target_rot = delta_rot_mat.dot(curr_rot) + + target_pose = PoseUtils.make_pose(target_pos, target_rot) + return target_pose + + def action_to_gripper_action(self, action): + """ + Extracts the gripper actuation part of an action (compatible with env.step). + + Args: + action (np.array): environment action + + Returns: + gripper_action (np.array): subset of environment action for gripper actuation + """ + + # last dimension is gripper action + return action[-1:] + + # robosuite-specific helper method for getting object poses + def get_object_pose(self, obj_name, obj_type): + """ + Returns 4x4 object pose given the name of the object and the type. + + Args: + obj_name (str): name of object + obj_type (str): type of object - either "body", "geom", or "site" + + Returns: + obj_pose (np.array): 4x4 object pose + """ + assert obj_type in ["body", "geom", "site"] + + if obj_type == "body": + obj_id = self.env.sim.model.body_name2id(obj_name) + obj_pos = np.array(self.env.sim.data.body_xpos[obj_id]) + obj_rot = np.array(self.env.sim.data.body_xmat[obj_id].reshape(3, 3)) + elif obj_type == "geom": + obj_id = self.env.sim.model.geom_name2id(obj_name) + obj_pos = np.array(self.env.sim.data.geom_xpos[obj_id]) + obj_rot = np.array(self.env.sim.data.geom_xmat[obj_id].reshape(3, 3)) + elif obj_type == "site": + obj_id = self.env.sim.model.site_name2id(obj_name) + obj_pos = np.array(self.env.sim.data.site_xpos[obj_id]) + obj_rot = np.array(self.env.sim.data.site_xmat[obj_id].reshape(3, 3)) + + return PoseUtils.make_pose(obj_pos, obj_rot) + + +class MG_Coffee(RobosuiteInterface): + """ + Corresponds to robosuite Coffee task and variants. + """ + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # two relevant objects - coffee pod and coffee machine + return dict( + coffee_pod=self.get_object_pose(obj_name=self.env.coffee_pod.root_body, obj_type="body"), + coffee_machine=self.get_object_pose(obj_name=self.env.coffee_machine.root_body, obj_type="body"), + ) + + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + metrics = self.env._get_partial_task_metrics() + + # first subtask is grasping coffee pod (motion relative to pod) + signals["grasp"] = int(metrics["grasp"]) + + # final subtask is inserting pod into machine and closing the lid (motion relative to machine) - but final subtask signal is not needed + return signals + + +class MG_Threading(RobosuiteInterface): + """ + Corresponds to robosuite Threading task and variants. + """ + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # two relevant objects - needle and tripod + return dict( + needle=self.get_object_pose(obj_name=self.env.needle.root_body, obj_type="body"), + tripod=self.get_object_pose(obj_name=self.env.tripod.root_body, obj_type="body"), + ) + + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + # first subtask is grasping needle (motion relative to needle) + signals["grasp"] = int(self.env._check_grasp( + gripper=self.env.robots[0].gripper, + object_geoms=[g for g in self.env.needle.contact_geoms]) + ) + + # final subtask is inserting needle into tripod (motion relative to tripod) - but final subtask signal is not needed + return signals + + +class MG_ThreePieceAssembly(RobosuiteInterface): + """ + Corresponds to robosuite ThreePieceAssembly task and variants. + """ + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # three relevant objects - base piece, piece_1, piece_2 + return dict( + base=self.get_object_pose(obj_name=self.env.base.root_body, obj_type="body"), + piece_1=self.get_object_pose(obj_name=self.env.piece_1.root_body, obj_type="body"), + piece_2=self.get_object_pose(obj_name=self.env.piece_2.root_body, obj_type="body"), + ) + + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + metrics = self.env._get_partial_task_metrics() + + # first subtask is grasping piece_1 (motion relative to piece_1) + signals["grasp_1"] = int(self.env._check_grasp( + gripper=self.env.robots[0].gripper, + object_geoms=[g for g in self.env.piece_1.contact_geoms]) + ) + + # second subtask is inserting piece_1 into the base (motion relative to base) + signals["insert_1"] = int(metrics["first_piece_assembled"]) + + # third subtask is grasping piece_2 (motion relative to piece_2) + signals["grasp_2"] = int(self.env._check_grasp( + gripper=self.env.robots[0].gripper, + object_geoms=[g for g in self.env.piece_2.contact_geoms]) + ) + + # final subtask is inserting piece_2 into piece_1 (motion relative to piece_1) - but final subtask signal is not needed + return signals + + +class MG_Square(RobosuiteInterface): + """ + Corresponds to robosuite Square task and variants. + """ + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # two relevant objects - square nut and peg + return dict( + square_nut=self.get_object_pose(obj_name=self.env.nuts[self.env.nut_to_id["square"]].root_body, obj_type="body"), + square_peg=self.get_object_pose(obj_name="peg1", obj_type="body"), + ) + + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + # first subtask is grasping square nut (motion relative to square nut) + signals["grasp"] = int(self.env._check_grasp( + gripper=self.env.robots[0].gripper, + object_geoms=[g for g in self.env.nuts[self.env.nut_to_id["square"]].contact_geoms]) + ) + + # final subtask is inserting square nut onto square peg (motion relative to square peg) - but final subtask signal is not needed + return signals + + +class MG_Stack(RobosuiteInterface): + """ + Corresponds to robosuite Stack task and variants. + """ + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # two relevant objects - cubeA and cubeB + return dict( + cubeA=self.get_object_pose(obj_name=self.env.cubeA.root_body, obj_type="body"), + cubeB=self.get_object_pose(obj_name=self.env.cubeB.root_body, obj_type="body"), + ) + + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + # first subtask is grasping cubeA (motion relative to cubeA) + signals["grasp"] = int(self.env._check_grasp(gripper=self.env.robots[0].gripper, object_geoms=self.env.cubeA)) + + # final subtask is placing cubeA on cubeB (motion relative to cubeB) - but final subtask signal is not needed + return signals + + +class MG_StackThree(RobosuiteInterface): + """ + Corresponds to robosuite StackThree task and variants. + """ + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # three relevant objects - three cubes + return dict( + cubeA=self.get_object_pose(obj_name=self.env.cubeA.root_body, obj_type="body"), + cubeB=self.get_object_pose(obj_name=self.env.cubeB.root_body, obj_type="body"), + cubeC=self.get_object_pose(obj_name=self.env.cubeC.root_body, obj_type="body"), + ) + + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + # first subtask is grasping cubeA (motion relative to cubeA) + signals["grasp_1"] = int(self.env._check_grasp(gripper=self.env.robots[0].gripper, object_geoms=self.env.cubeA)) + + # second subtask is placing cubeA on cubeB (motion relative to cubeB) + signals["stack_1"] = int(self.env._check_cubeA_stacked()) + + # third subtask is grasping cubeC (motion relative to cubeC) + signals["grasp_2"] = int(self.env._check_grasp(gripper=self.env.robots[0].gripper, object_geoms=self.env.cubeC)) + + # final subtask is placing cubeC on cubeA (motion relative to cubeA) - but final subtask signal is not needed + return signals + + +class MG_HammerCleanup(RobosuiteInterface): + """ + Corresponds to robosuite HammerCleanup task and variants. + """ + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # two relevant objects - hammer and drawer + return dict( + hammer=self.get_object_pose(obj_name=self.env.sorting_object.root_body, obj_type="body"), + drawer=self.get_object_pose(obj_name=self.env.cabinet_object.root_body, obj_type="body"), + ) + + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + # first subtask is opening the drawer (motion relative to drawer) + # check that drawer is open enough and end effector is far enough from drawer after opening it + drawer_pos, _ = PoseUtils.unmake_pose(self.get_object_pose(obj_name="CabinetObject_drawer_link", obj_type="body")) + eef_pos, _ = PoseUtils.unmake_pose(self.get_robot_eef_pose()) + eef_drawer_dist = np.linalg.norm(eef_pos - drawer_pos) + signals["open"] = int( + (self.env.sim.data.qpos[self.env.cabinet_qpos_addrs] < -0.10) and (eef_drawer_dist > 0.24) + ) + + # second subtask is grasping the hammer (motion relative to hammer) + signals["grasp"] = int(self.env._check_grasp( + gripper=self.env.robots[0].gripper, + object_geoms=[g for g in self.env.sorting_object.contact_geoms] + )) + + # final subtask is placing the hammer into the drawer and closing the drawer (motion relative to drawer) - but final subtask signal not needed + return signals + + +class MG_MugCleanup(RobosuiteInterface): + """ + Corresponds to robosuite MugCleanup task and variants. + """ + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # two relevant objects - mug and drawer + return dict( + object=self.get_object_pose(obj_name=self.env.cleanup_object.root_body, obj_type="body"), + drawer=self.get_object_pose(obj_name=self.env.drawer.root_body, obj_type="body"), + ) + + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + # first subtask is opening the drawer (motion relative to drawer) + # check that drawer is open enough and end effector is far enough from drawer after opening it + drawer_pos, _ = PoseUtils.unmake_pose(self.get_object_pose(obj_name="DrawerObject_drawer_link", obj_type="body")) + eef_pos, _ = PoseUtils.unmake_pose(self.get_robot_eef_pose()) + eef_drawer_dist = np.linalg.norm(eef_pos - drawer_pos) + signals["open"] = int( + (self.env.sim.data.qpos[self.env.drawer_qpos_addr] < -0.10) and (eef_drawer_dist > 0.24) + ) + + # second subtask is grasping the mug (motion relative to mug) + signals["grasp"] = int(self.env._check_grasp_tolerant( + gripper=self.env.robots[0].gripper, + object_geoms=[g for g in self.env.cleanup_object.contact_geoms] + )) + + # final subtask is placing the mug into the drawer and closing the drawer (motion relative to drawer) - but final subtask signal not needed + return signals + + +class MG_NutAssembly(RobosuiteInterface): + """ + Corresponds to robosuite NutAssembly task and variants. + """ + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # four relevant objects - square and round nuts and pegs + return dict( + square_nut=self.get_object_pose(obj_name=self.env.nuts[self.env.nut_to_id["square"]].root_body, obj_type="body"), + round_nut=self.get_object_pose(obj_name=self.env.nuts[self.env.nut_to_id["round"]].root_body, obj_type="body"), + square_peg=self.get_object_pose(obj_name="peg1", obj_type="body"), + round_peg=self.get_object_pose(obj_name="peg2", obj_type="body"), + ) + + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + # checks which objects are on their correct pegs and records them in @self.objects_on_pegs + self.env._check_success() + + # first subtask is grasping square nut (motion relative to square nut) + signals["grasp_square_nut"] = int(self.env._check_grasp( + gripper=self.env.robots[0].gripper, + object_geoms=[g for g in self.env.nuts[self.env.nut_to_id["square"]].contact_geoms]) + ) + + # second subtask is inserting square nut onto square peg (motion relative to square peg) + signals["insert_square_nut"] = int(self.env.objects_on_pegs[self.env.nut_to_id["square"]]) + + # third subtask is grasping round nut (motion relative to round nut) + signals["grasp_round_nut"] = int(self.env._check_grasp( + gripper=self.env.robots[0].gripper, + object_geoms=[g for g in self.env.nuts[self.env.nut_to_id["round"]].contact_geoms]) + ) + + # final subtask is inserting round nut onto round peg (motion relative to round peg) - but final subtask signal is not needed + return signals + + +class MG_PickPlace(RobosuiteInterface): + """ + Corresponds to robosuite PickPlace task and variants. + """ + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # four relevant objects - milk, bread, cereal, can + object_poses = dict() + for obj_name in self.env.object_to_id: + obj = self.env.objects[self.env.object_to_id[obj_name]] + object_poses[obj_name] = self.get_object_pose(obj_name=obj.root_body, obj_type="body") + return object_poses + + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + # checks which objects are in their correct bins and records them in @self.objects_in_bins + self.env._check_success() + + object_names_in_order = ["milk", "cereal", "bread", "can"] + assert set(self.env.object_to_id.keys()) == set(object_names_in_order) + n_obj = len(object_names_in_order) + + # each subtask is a grasp and then a place + for i, obj_name in enumerate(object_names_in_order): + obj_id = self.env.object_to_id[obj_name] + + # first subtask for each object is grasping (motion relative to the object) + signals["grasp_{}".format(obj_name)] = int(self.env._check_grasp( + gripper=self.env.robots[0].gripper, + object_geoms=[g for g in self.env.objects[obj_id].contact_geoms]) + ) + + # skip final subtask - unneeded + if i < (n_obj - 1): + # second subtask for each object is placement into bin (motion relative to bin) + signals["place_{}".format(obj_name)] = int(self.env.objects_in_bins[obj_id]) + + return signals + + +class MG_Kitchen(RobosuiteInterface): + """ + Corresponds to robosuite Kitchen task and variants. + """ + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # five relevant objects - bread, pot, stove, button, and serving region + return dict( + bread=self.get_object_pose(obj_name=self.env.bread_ingredient.root_body, obj_type="body"), + pot=self.get_object_pose(obj_name=self.env.pot_object.root_body, obj_type="body"), + stove=self.get_object_pose(obj_name=self.env.stove_object_1.root_body, obj_type="body"), + button=self.get_object_pose(obj_name=self.env.button_object_1.root_body, obj_type="body"), + serving_region=self.get_object_pose(obj_name=self.env.serving_region.root_body, obj_type="body"), + ) + + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + # first subtask is to flip the switch to turn stove on (motion relative to button) + signals["stove_on"] = int(self.env.buttons_on[1]) + + # second subtask is to grasp pot (motion relative to pot) + grasped_pot = self.env._check_grasp( + gripper=self.env.robots[0].gripper, + object_geoms=[g for g in self.env.pot_object.contact_geoms] + ) + signals["grasp_pot"] = int(grasped_pot) + + # third subtask is to place pot on stove (motion relative to stove) + + # check for pot-stove contact and that hand is not grasping pot + pot_bottom_in_contact_with_stove = self.env.check_contact("PotObject_body_0", "Stove1_collision_burner") + signals["place_pot_on_stove"] = int(pot_bottom_in_contact_with_stove and not grasped_pot) + + # fourth subtask is to grasp bread (motion relative to bread) + signals["grasp_bread"] = int(self.env._check_grasp( + gripper=self.env.robots[0].gripper, + object_geoms=[g for g in self.env.bread_ingredient.contact_geoms] + )) + + # fifth subtask is to place bread in pot and grasp pot (motion relative to pot) + signals["place_bread_in_pot"] = int(self.env.check_contact(self.env.bread_ingredient, self.env.pot_object) and grasped_pot) + + # sixth subtask is to place pot in front of serving region and then push it into the serving region (motion relative to serving region) + pot_pos = self.env.sim.data.body_xpos[self.env.pot_object_id] + serving_region_pos = self.env.sim.data.body_xpos[self.env.serving_region_id] + dist_serving_pot = serving_region_pos - pot_pos + pot_in_serving_region = np.abs(dist_serving_pot[0]) < 0.05 and np.abs(dist_serving_pot[1]) < 0.10 and np.abs(dist_serving_pot[2]) < 0.05 + signals["serve"] = int(pot_in_serving_region) + + # final subtask is to turn off the stove (motion relative to button) - but final subtask signal not needed + return signals + + +class MG_CoffeePreparation(RobosuiteInterface): + """ + Corresponds to robosuite CoffeePreparation task and variants. + """ + def get_object_poses(self): + """ + Gets the pose of each object relevant to MimicGen data generation in the current scene. + + Returns: + object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 np.array) + """ + + # four relevant objects - coffee pod, coffee machine, drawer, and mug + return dict( + coffee_pod=self.get_object_pose(obj_name=self.env.coffee_pod.root_body, obj_type="body"), + coffee_machine=self.get_object_pose(obj_name=self.env.coffee_machine.root_body, obj_type="body"), + drawer=self.get_object_pose(obj_name=self.env.cabinet_object.root_body, obj_type="body"), + mug=self.get_object_pose(obj_name=self.env.mug.root_body, obj_type="body"), + ) + + def get_subtask_term_signals(self): + """ + Gets a dictionary of binary flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. MimicGen only uses this + when parsing source demonstrations at the start of data generation, and it only + uses the first 0 -> 1 transition in this signal to detect the end of a subtask. + + Returns: + subtask_term_signals (dict): dictionary that maps subtask name to termination flag (0 or 1) + """ + signals = dict() + + metrics = self.env._get_partial_task_metrics() + + # first subtask is grasping mug (motion relative to mug) + signals["mug_grasp"] = int(metrics["mug_grasp"]) + + # second subtask is placing the mug on the coffee machine base and then opening the lid (motion relative to coffee machine) + signals["mug_place"] = int(self.env._check_mug_placement() and (self.env.sim.data.qpos[self.env.hinge_qpos_addr] > 2.08)) + + # third subtask is opening the drawer (motion relative to drawer) + signals["drawer_open"] = int(self.env.sim.data.qpos[self.env.cabinet_qpos_addr] < -0.19) + + # fourth subtask is grasping the coffee pod (motion relative to coffee pod) + signals["pod_grasp"] = int(metrics["grasp"]) + + # final subtask is inserting pod into machine and closing the lid (motion relative to machine) - but final subtask signal is not needed + return signals diff --git a/mimicgen_envs/envs/robosuite/__init__.py b/mimicgen/envs/__init__.py similarity index 57% rename from mimicgen_envs/envs/robosuite/__init__.py rename to mimicgen/envs/__init__.py index 981400f..8edfcfd 100644 --- a/mimicgen_envs/envs/robosuite/__init__.py +++ b/mimicgen/envs/__init__.py @@ -1,3 +1,3 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. \ No newline at end of file diff --git a/mimicgen_envs/scripts/__init__.py b/mimicgen/envs/robosuite/__init__.py similarity index 57% rename from mimicgen_envs/scripts/__init__.py rename to mimicgen/envs/robosuite/__init__.py index 981400f..8edfcfd 100644 --- a/mimicgen_envs/scripts/__init__.py +++ b/mimicgen/envs/robosuite/__init__.py @@ -1,3 +1,3 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. \ No newline at end of file diff --git a/mimicgen_envs/envs/robosuite/coffee.py b/mimicgen/envs/robosuite/coffee.py similarity index 98% rename from mimicgen_envs/envs/robosuite/coffee.py rename to mimicgen/envs/robosuite/coffee.py index 3782f30..f7a72a7 100644 --- a/mimicgen_envs/envs/robosuite/coffee.py +++ b/mimicgen/envs/robosuite/coffee.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -17,9 +17,9 @@ from robosuite.utils.placement_samplers import SequentialCompositeSampler, UniformRandomSampler from robosuite.utils.observables import Observable, sensor -import mimicgen_envs -from mimicgen_envs.models.robosuite.objects import BlenderObject, CoffeeMachinePodObject, CoffeeMachineObject, LongDrawerObject, CupObject -from mimicgen_envs.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG +import mimicgen +from mimicgen.models.robosuite.objects import BlenderObject, CoffeeMachinePodObject, CoffeeMachineObject, LongDrawerObject, CupObject +from mimicgen.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG class Coffee(SingleArmEnv_MG): @@ -213,8 +213,6 @@ def reward(self, action=None): """ Reward function for the task. - Dense reward: TODO - The sparse reward only consists of the threading component. Note that the final reward is normalized and scaled by @@ -468,7 +466,6 @@ def _create_obj_sensors(self, obj_name, modality="object"): names (list): array of corresponding observable names """ - ### TODO: this was stolen from pick-place - do we want to move this into utils to share it? ### pf = self.robots[0].robot_model.naming_prefix @sensor(modality=modality) @@ -634,8 +631,6 @@ def _get_partial_task_metrics(self): lid_check = self._check_lid() pod_check = self._check_pod() - # TODO: should probably clean up redundant code below and in @_check_pod - # pod should be in pod holder pod_holder_pos = np.array(self.sim.data.body_xpos[self.obj_body_id["coffee_pod_holder"]]) pod_pos = np.array(self.sim.data.body_xpos[self.obj_body_id["coffee_pod"]]) @@ -858,7 +853,7 @@ def _get_mug_model(self): """ shapenet_id = "3143a4ac" # beige round mug, works well and matches color scheme of other assets shapenet_scale = 1.0 - base_mjcf_path = os.path.join(mimicgen_envs.__path__[0], "models/robosuite/assets/shapenet_core/mugs") + base_mjcf_path = os.path.join(mimicgen.__path__[0], "models/robosuite/assets/shapenet_core/mugs") mjcf_path = os.path.join(base_mjcf_path, "{}/model.xml".format(shapenet_id)) self.mug = BlenderObject( @@ -1250,7 +1245,7 @@ def _get_partial_task_metrics(self): # new task success includes mug placement metrics["task"] = metrics["task"] and metrics["mug_place"] - # TODO: decide if we want a check on drawer being closed here, to make the task even harder + # can have a check on drawer being closed here, to make the task even harder # print(self.sim.data.qpos[self.cabinet_qpos_addr]) return metrics diff --git a/mimicgen_envs/envs/robosuite/hammer_cleanup.py b/mimicgen/envs/robosuite/hammer_cleanup.py similarity index 98% rename from mimicgen_envs/envs/robosuite/hammer_cleanup.py rename to mimicgen/envs/robosuite/hammer_cleanup.py index f25867a..6fb99be 100644 --- a/mimicgen_envs/envs/robosuite/hammer_cleanup.py +++ b/mimicgen/envs/robosuite/hammer_cleanup.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -26,9 +26,9 @@ import robosuite_task_zoo from robosuite_task_zoo.environments.manipulation.hammer_place import HammerPlaceEnv -import mimicgen_envs -from mimicgen_envs.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG -from mimicgen_envs.models.robosuite.objects import DrawerObject +import mimicgen +from mimicgen.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG +from mimicgen.models.robosuite.objects import DrawerObject class HammerCleanup_D0(HammerPlaceEnv, SingleArmEnv_MG): diff --git a/mimicgen_envs/envs/robosuite/kitchen.py b/mimicgen/envs/robosuite/kitchen.py similarity index 98% rename from mimicgen_envs/envs/robosuite/kitchen.py rename to mimicgen/envs/robosuite/kitchen.py index 3d2b748..e08a0fa 100644 --- a/mimicgen_envs/envs/robosuite/kitchen.py +++ b/mimicgen/envs/robosuite/kitchen.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -26,8 +26,8 @@ from robosuite_task_zoo.environments.manipulation.kitchen import KitchenEnv from robosuite_task_zoo.models.kitchen import PotObject, StoveObject, ButtonObject, ServingRegionObject -import mimicgen_envs -from mimicgen_envs.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG +import mimicgen +from mimicgen.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG class StoveObjectNew(StoveObject): @@ -67,7 +67,7 @@ class ServingRegionObjectNew(MujocoXMLObject): def __init__(self, name, joints=None): # our custom serving region xml - turn site into visual-only geom so that it shows up on env reset (instead # of after first env step) - path_to_serving_region_xml = os.path.join(mimicgen_envs.__path__[0], "models/robosuite/assets/objects/serving_region.xml") + path_to_serving_region_xml = os.path.join(mimicgen.__path__[0], "models/robosuite/assets/objects/serving_region.xml") super().__init__(path_to_serving_region_xml, name=name, joints=None, obj_type="all", duplicate_collision_geoms=True) diff --git a/mimicgen_envs/envs/robosuite/mug_cleanup.py b/mimicgen/envs/robosuite/mug_cleanup.py similarity index 97% rename from mimicgen_envs/envs/robosuite/mug_cleanup.py rename to mimicgen/envs/robosuite/mug_cleanup.py index 3e84a88..0b9d602 100644 --- a/mimicgen_envs/envs/robosuite/mug_cleanup.py +++ b/mimicgen/envs/robosuite/mug_cleanup.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -22,9 +22,9 @@ from robosuite.utils.placement_samplers import SequentialCompositeSampler, UniformRandomSampler from robosuite.utils.observables import Observable, sensor -import mimicgen_envs -from mimicgen_envs.models.robosuite.objects import BlenderObject, DrawerObject, LongDrawerObject -from mimicgen_envs.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG +import mimicgen +from mimicgen.models.robosuite.objects import BlenderObject, DrawerObject, LongDrawerObject +from mimicgen.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG class MugCleanup(SingleArmEnv_MG): @@ -224,8 +224,6 @@ def reward(self, action=None): """ Reward function for the task. - Dense reward: TODO - The sparse reward only consists of the threading component. Note that the final reward is normalized and scaled by @@ -358,7 +356,7 @@ def _get_object_model(self): """ Allow subclasses to override which object to pack into drawer - should load into @self.cleanup_object. """ - base_mjcf_path = os.path.join(mimicgen_envs.__path__[0], "models/robosuite/assets/shapenet_core/mugs") + base_mjcf_path = os.path.join(mimicgen.__path__[0], "models/robosuite/assets/shapenet_core/mugs") mjcf_path = os.path.join(base_mjcf_path, "{}/model.xml".format(self._shapenet_id)) self.cleanup_object = BlenderObject( @@ -543,7 +541,6 @@ def _create_obj_sensors(self, obj_name, modality="object"): names (list): array of corresponding observable names """ - ### TODO: this was stolen from pick-place - do we want to move this into utils to share it? ### pf = self.robots[0].robot_model.naming_prefix @sensor(modality=modality) @@ -680,7 +677,7 @@ def __init__( ("e94e46bc", 0.8), # dark blue cylindrical mug ("fad118b3", 0.66666667), # tall green cylindrical mug ] - self._base_mjcf_path = os.path.join(mimicgen_envs.__path__[0], "models/robosuite/assets/shapenet_core/mugs") + self._base_mjcf_path = os.path.join(mimicgen.__path__[0], "models/robosuite/assets/shapenet_core/mugs") super(MugCleanup_O2, self).__init__(shapenet_id=None, shapenet_scale=None, **kwargs) def _get_object_model(self): diff --git a/mimicgen_envs/envs/robosuite/nut_assembly.py b/mimicgen/envs/robosuite/nut_assembly.py similarity index 99% rename from mimicgen_envs/envs/robosuite/nut_assembly.py rename to mimicgen/envs/robosuite/nut_assembly.py index d93425e..3f3c359 100644 --- a/mimicgen_envs/envs/robosuite/nut_assembly.py +++ b/mimicgen/envs/robosuite/nut_assembly.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -16,7 +16,7 @@ from robosuite.utils.mjcf_utils import array_to_string, string_to_array, find_elements from robosuite.utils import RandomizationError -from mimicgen_envs.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG +from mimicgen.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG class NutAssembly_D0(NutAssembly, SingleArmEnv_MG): diff --git a/mimicgen_envs/envs/robosuite/pick_place.py b/mimicgen/envs/robosuite/pick_place.py similarity index 89% rename from mimicgen_envs/envs/robosuite/pick_place.py rename to mimicgen/envs/robosuite/pick_place.py index 233cf9b..d9308df 100644 --- a/mimicgen_envs/envs/robosuite/pick_place.py +++ b/mimicgen/envs/robosuite/pick_place.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. diff --git a/mimicgen_envs/envs/robosuite/single_arm_env_mg.py b/mimicgen/envs/robosuite/single_arm_env_mg.py similarity index 93% rename from mimicgen_envs/envs/robosuite/single_arm_env_mg.py rename to mimicgen/envs/robosuite/single_arm_env_mg.py index 0d65052..38643f2 100644 --- a/mimicgen_envs/envs/robosuite/single_arm_env_mg.py +++ b/mimicgen/envs/robosuite/single_arm_env_mg.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -14,7 +14,7 @@ except ImportError: pass -import mimicgen_envs +import mimicgen class SingleArmEnv_MG(SingleArmEnv): @@ -57,11 +57,11 @@ def edit_model_xml(self, xml_str): new_path = "/".join(new_path_split) elem.set("file", new_path) - # replace all paths to mimicgen_envs assets - check_lst = [loc for loc, val in enumerate(old_path_split) if val == "mimicgen_envs"] + # replace all paths to mimicgen assets + check_lst = [loc for loc, val in enumerate(old_path_split) if val == "mimicgen"] if len(check_lst) > 0: ind = max(check_lst) # last occurrence index - new_path_split = os.path.split(mimicgen_envs.__file__)[0].split("/") + old_path_split[ind + 1 :] + new_path_split = os.path.split(mimicgen.__file__)[0].split("/") + old_path_split[ind + 1 :] new_path = "/".join(new_path_split) elem.set("file", new_path) diff --git a/mimicgen_envs/envs/robosuite/stack.py b/mimicgen/envs/robosuite/stack.py similarity index 99% rename from mimicgen_envs/envs/robosuite/stack.py rename to mimicgen/envs/robosuite/stack.py index ae9c585..6009e4f 100644 --- a/mimicgen_envs/envs/robosuite/stack.py +++ b/mimicgen/envs/robosuite/stack.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -16,7 +16,7 @@ from robosuite.utils.observables import Observable, sensor from robosuite.environments.manipulation.stack import Stack -from mimicgen_envs.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG +from mimicgen.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG class Stack_D0(Stack, SingleArmEnv_MG): diff --git a/mimicgen_envs/envs/robosuite/threading.py b/mimicgen/envs/robosuite/threading.py similarity index 98% rename from mimicgen_envs/envs/robosuite/threading.py rename to mimicgen/envs/robosuite/threading.py index c6c76c7..ee6f533 100644 --- a/mimicgen_envs/envs/robosuite/threading.py +++ b/mimicgen/envs/robosuite/threading.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -13,8 +13,8 @@ from robosuite.utils.placement_samplers import SequentialCompositeSampler, UniformRandomSampler from robosuite.utils.observables import Observable, sensor -from mimicgen_envs.models.robosuite.objects import NeedleObject, RingTripodObject -from mimicgen_envs.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG +from mimicgen.models.robosuite.objects import NeedleObject, RingTripodObject +from mimicgen.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG class Threading(SingleArmEnv_MG): @@ -208,8 +208,6 @@ def reward(self, action=None): """ Reward function for the task. - Dense reward: TODO - The sparse reward only consists of the threading component. Note that the final reward is normalized and scaled by @@ -423,7 +421,6 @@ def _create_obj_sensors(self, obj_name, modality="object"): names (list): array of corresponding observable names """ - ### TODO: this was stolen from pick-place - do we want to move this into utils to share it? ### pf = self.robots[0].robot_model.naming_prefix @sensor(modality=modality) diff --git a/mimicgen_envs/envs/robosuite/three_piece_assembly.py b/mimicgen/envs/robosuite/three_piece_assembly.py similarity index 99% rename from mimicgen_envs/envs/robosuite/three_piece_assembly.py rename to mimicgen/envs/robosuite/three_piece_assembly.py index 71111ed..82d1d21 100644 --- a/mimicgen_envs/envs/robosuite/three_piece_assembly.py +++ b/mimicgen/envs/robosuite/three_piece_assembly.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -16,8 +16,8 @@ from robosuite.utils.observables import Observable, sensor from robosuite.utils.mjcf_utils import CustomMaterial, find_elements, string_to_array -from mimicgen_envs.models.robosuite.objects import BoxPatternObject -from mimicgen_envs.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG +from mimicgen.models.robosuite.objects import BoxPatternObject +from mimicgen.envs.robosuite.single_arm_env_mg import SingleArmEnv_MG class ThreePieceAssembly(SingleArmEnv_MG): @@ -223,8 +223,6 @@ def reward(self, action=None): """ Reward function for the task. - Dense reward: TODO - The sparse reward only consists of the threading component. Note that the final reward is normalized and scaled by @@ -612,7 +610,6 @@ def _create_obj_sensors(self, obj_name, modality="object"): names (list): array of corresponding observable names """ - ### TODO: this was stolen from pick-place - do we want to move this into utils to share it? ### pf = self.robots[0].robot_model.naming_prefix @sensor(modality=modality) diff --git a/mimicgen/exps/templates/robosuite/coffee.json b/mimicgen/exps/templates/robosuite/coffee.json new file mode 100644 index 0000000..0086255 --- /dev/null +++ b/mimicgen/exps/templates/robosuite/coffee.json @@ -0,0 +1,70 @@ +{ + "name": "coffee", + "type": "robosuite", + "experiment": { + "name": "demo", + "source": { + "dataset_path": null, + "filter_key": null, + "n": null, + "start": null + }, + "generation": { + "path": null, + "guarantee": false, + "keep_failed": true, + "num_trials": 10, + "select_src_per_subtask": false, + "transform_first_robot_pose": false, + "interpolate_from_last_target_pose": true + }, + "task": { + "name": null, + "robot": null, + "gripper": null, + "interface": null, + "interface_type": null + }, + "max_num_failures": 50, + "render_video": true, + "num_demo_to_render": 50, + "num_fail_demo_to_render": 50, + "log_every_n_attempts": 50, + "seed": 1 + }, + "obs": { + "collect_obs": true, + "camera_names": [], + "camera_height": 84, + "camera_width": 84 + }, + "task": { + "task_spec": { + "subtask_1": { + "object_ref": "coffee_pod", + "subtask_term_signal": "grasp", + "subtask_term_offset_range": [ + 5, + 10 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_2": { + "object_ref": "coffee_machine", + "subtask_term_signal": null, + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + } + } + } +} \ No newline at end of file diff --git a/mimicgen/exps/templates/robosuite/coffee_preparation.json b/mimicgen/exps/templates/robosuite/coffee_preparation.json new file mode 100644 index 0000000..da4beb9 --- /dev/null +++ b/mimicgen/exps/templates/robosuite/coffee_preparation.json @@ -0,0 +1,112 @@ +{ + "name": "coffee_preparation", + "type": "robosuite", + "experiment": { + "name": "demo", + "source": { + "dataset_path": null, + "filter_key": null, + "n": null, + "start": null + }, + "generation": { + "path": null, + "guarantee": false, + "keep_failed": true, + "num_trials": 10, + "select_src_per_subtask": false, + "transform_first_robot_pose": false, + "interpolate_from_last_target_pose": true + }, + "task": { + "name": null, + "robot": null, + "gripper": null, + "interface": null, + "interface_type": null + }, + "max_num_failures": 50, + "render_video": true, + "num_demo_to_render": 50, + "num_fail_demo_to_render": 50, + "log_every_n_attempts": 50, + "seed": 1 + }, + "obs": { + "collect_obs": true, + "camera_names": [], + "camera_height": 84, + "camera_width": 84 + }, + "task": { + "task_spec": { + "subtask_1": { + "object_ref": "mug", + "subtask_term_signal": "mug_grasp", + "subtask_term_offset_range": [ + 5, + 10 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_2": { + "object_ref": "coffee_machine", + "subtask_term_signal": "mug_place", + "subtask_term_offset_range": [ + 5, + 10 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_3": { + "object_ref": "drawer", + "subtask_term_signal": "drawer_open", + "subtask_term_offset_range": [ + 5, + 10 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_4": { + "object_ref": "coffee_pod", + "subtask_term_signal": "pod_grasp", + "subtask_term_offset_range": [ + 5, + 10 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_5": { + "object_ref": "coffee_machine", + "subtask_term_signal": null, + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + } + } + } +} \ No newline at end of file diff --git a/mimicgen/exps/templates/robosuite/hammer_cleanup.json b/mimicgen/exps/templates/robosuite/hammer_cleanup.json new file mode 100644 index 0000000..f8a8feb --- /dev/null +++ b/mimicgen/exps/templates/robosuite/hammer_cleanup.json @@ -0,0 +1,84 @@ +{ + "name": "hammer_cleanup", + "type": "robosuite", + "experiment": { + "name": "demo", + "source": { + "dataset_path": null, + "filter_key": null, + "n": null, + "start": null + }, + "generation": { + "path": null, + "guarantee": false, + "keep_failed": true, + "num_trials": 10, + "select_src_per_subtask": false, + "transform_first_robot_pose": false, + "interpolate_from_last_target_pose": true + }, + "task": { + "name": null, + "robot": null, + "gripper": null, + "interface": null, + "interface_type": null + }, + "max_num_failures": 50, + "render_video": true, + "num_demo_to_render": 50, + "num_fail_demo_to_render": 50, + "log_every_n_attempts": 50, + "seed": 1 + }, + "obs": { + "collect_obs": true, + "camera_names": [], + "camera_height": 84, + "camera_width": 84 + }, + "task": { + "task_spec": { + "subtask_1": { + "object_ref": "drawer", + "subtask_term_signal": "open", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_2": { + "object_ref": "hammer", + "subtask_term_signal": "grasp", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_3": { + "object_ref": "drawer", + "subtask_term_signal": null, + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + } + } + } +} \ No newline at end of file diff --git a/mimicgen/exps/templates/robosuite/kitchen.json b/mimicgen/exps/templates/robosuite/kitchen.json new file mode 100644 index 0000000..f7fdb68 --- /dev/null +++ b/mimicgen/exps/templates/robosuite/kitchen.json @@ -0,0 +1,140 @@ +{ + "name": "kitchen", + "type": "robosuite", + "experiment": { + "name": "demo", + "source": { + "dataset_path": null, + "filter_key": null, + "n": null, + "start": null + }, + "generation": { + "path": null, + "guarantee": false, + "keep_failed": true, + "num_trials": 10, + "select_src_per_subtask": false, + "transform_first_robot_pose": false, + "interpolate_from_last_target_pose": true + }, + "task": { + "name": null, + "robot": null, + "gripper": null, + "interface": null, + "interface_type": null + }, + "max_num_failures": 50, + "render_video": true, + "num_demo_to_render": 50, + "num_fail_demo_to_render": 50, + "log_every_n_attempts": 50, + "seed": 1 + }, + "obs": { + "collect_obs": true, + "camera_names": [], + "camera_height": 84, + "camera_width": 84 + }, + "task": { + "task_spec": { + "subtask_1": { + "object_ref": "button", + "subtask_term_signal": "stove_on", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_2": { + "object_ref": "pot", + "subtask_term_signal": "grasp_pot", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_3": { + "object_ref": "stove", + "subtask_term_signal": "place_pot_on_stove", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_4": { + "object_ref": "bread", + "subtask_term_signal": "grasp_bread", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_5": { + "object_ref": "pot", + "subtask_term_signal": "place_bread_in_pot", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_6": { + "object_ref": "serving_region", + "subtask_term_signal": "serve", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_7": { + "object_ref": "button", + "subtask_term_signal": null, + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + } + } + } +} \ No newline at end of file diff --git a/mimicgen/exps/templates/robosuite/mug_cleanup.json b/mimicgen/exps/templates/robosuite/mug_cleanup.json new file mode 100644 index 0000000..e60976a --- /dev/null +++ b/mimicgen/exps/templates/robosuite/mug_cleanup.json @@ -0,0 +1,84 @@ +{ + "name": "mug_cleanup", + "type": "robosuite", + "experiment": { + "name": "demo", + "source": { + "dataset_path": null, + "filter_key": null, + "n": null, + "start": null + }, + "generation": { + "path": null, + "guarantee": false, + "keep_failed": true, + "num_trials": 10, + "select_src_per_subtask": false, + "transform_first_robot_pose": false, + "interpolate_from_last_target_pose": true + }, + "task": { + "name": null, + "robot": null, + "gripper": null, + "interface": null, + "interface_type": null + }, + "max_num_failures": 50, + "render_video": true, + "num_demo_to_render": 50, + "num_fail_demo_to_render": 50, + "log_every_n_attempts": 50, + "seed": 1 + }, + "obs": { + "collect_obs": true, + "camera_names": [], + "camera_height": 84, + "camera_width": 84 + }, + "task": { + "task_spec": { + "subtask_1": { + "object_ref": "drawer", + "subtask_term_signal": "open", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_2": { + "object_ref": "object", + "subtask_term_signal": "grasp", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_3": { + "object_ref": "drawer", + "subtask_term_signal": null, + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + } + } + } +} \ No newline at end of file diff --git a/mimicgen/exps/templates/robosuite/nut_assembly.json b/mimicgen/exps/templates/robosuite/nut_assembly.json new file mode 100644 index 0000000..922d8e5 --- /dev/null +++ b/mimicgen/exps/templates/robosuite/nut_assembly.json @@ -0,0 +1,100 @@ +{ + "name": "nut_assembly", + "type": "robosuite", + "experiment": { + "name": "demo", + "source": { + "dataset_path": null, + "filter_key": null, + "n": null, + "start": null + }, + "generation": { + "path": null, + "guarantee": false, + "keep_failed": true, + "num_trials": 10, + "select_src_per_subtask": false, + "transform_first_robot_pose": false, + "interpolate_from_last_target_pose": true + }, + "task": { + "name": null, + "robot": null, + "gripper": null, + "interface": null, + "interface_type": null + }, + "max_num_failures": 50, + "render_video": true, + "num_demo_to_render": 50, + "num_fail_demo_to_render": 50, + "log_every_n_attempts": 50, + "seed": 1 + }, + "obs": { + "collect_obs": true, + "camera_names": [], + "camera_height": 84, + "camera_width": 84 + }, + "task": { + "task_spec": { + "subtask_1": { + "object_ref": "square_nut", + "subtask_term_signal": "grasp_square_nut", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "nearest_neighbor_object", + "selection_strategy_kwargs": { + "nn_k": 3 + }, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_2": { + "object_ref": "square_peg", + "subtask_term_signal": "insert_square_nut", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_3": { + "object_ref": "round_nut", + "subtask_term_signal": "grasp_round_nut", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_4": { + "object_ref": "round_peg", + "subtask_term_signal": null, + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + } + } + } +} \ No newline at end of file diff --git a/mimicgen/exps/templates/robosuite/pick_place.json b/mimicgen/exps/templates/robosuite/pick_place.json new file mode 100644 index 0000000..0067a8f --- /dev/null +++ b/mimicgen/exps/templates/robosuite/pick_place.json @@ -0,0 +1,153 @@ +{ + "name": "pick_place", + "type": "robosuite", + "experiment": { + "name": "demo", + "source": { + "dataset_path": null, + "filter_key": null, + "n": null, + "start": null + }, + "generation": { + "path": null, + "guarantee": false, + "keep_failed": true, + "num_trials": 10, + "select_src_per_subtask": false, + "transform_first_robot_pose": false, + "interpolate_from_last_target_pose": true + }, + "task": { + "name": null, + "robot": null, + "gripper": null, + "interface": null, + "interface_type": null + }, + "max_num_failures": 50, + "render_video": true, + "num_demo_to_render": 50, + "num_fail_demo_to_render": 50, + "log_every_n_attempts": 50, + "seed": 1 + }, + "obs": { + "collect_obs": true, + "camera_names": [], + "camera_height": 84, + "camera_width": 84 + }, + "task": { + "task_spec": { + "subtask_1": { + "object_ref": "milk", + "subtask_term_signal": "grasp_milk", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "nearest_neighbor_object", + "selection_strategy_kwargs": { + "nn_k": 3 + }, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_2": { + "object_ref": null, + "subtask_term_signal": "place_milk", + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_3": { + "object_ref": "cereal", + "subtask_term_signal": "grasp_cereal", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "nearest_neighbor_object", + "selection_strategy_kwargs": { + "nn_k": 3 + }, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_4": { + "object_ref": null, + "subtask_term_signal": "place_cereal", + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_5": { + "object_ref": "bread", + "subtask_term_signal": "grasp_bread", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "nearest_neighbor_object", + "selection_strategy_kwargs": { + "nn_k": 3 + }, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_6": { + "object_ref": null, + "subtask_term_signal": "place_bread", + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_7": { + "object_ref": "can", + "subtask_term_signal": "grasp_can", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "nearest_neighbor_object", + "selection_strategy_kwargs": { + "nn_k": 3 + }, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_8": { + "object_ref": null, + "subtask_term_signal": null, + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + } + } + } +} \ No newline at end of file diff --git a/mimicgen/exps/templates/robosuite/square.json b/mimicgen/exps/templates/robosuite/square.json new file mode 100644 index 0000000..23dce5d --- /dev/null +++ b/mimicgen/exps/templates/robosuite/square.json @@ -0,0 +1,72 @@ +{ + "name": "square", + "type": "robosuite", + "experiment": { + "name": "demo", + "source": { + "dataset_path": null, + "filter_key": null, + "n": null, + "start": null + }, + "generation": { + "path": null, + "guarantee": false, + "keep_failed": true, + "num_trials": 10, + "select_src_per_subtask": false, + "transform_first_robot_pose": false, + "interpolate_from_last_target_pose": true + }, + "task": { + "name": null, + "robot": null, + "gripper": null, + "interface": null, + "interface_type": null + }, + "max_num_failures": 50, + "render_video": true, + "num_demo_to_render": 50, + "num_fail_demo_to_render": 50, + "log_every_n_attempts": 50, + "seed": 1 + }, + "obs": { + "collect_obs": true, + "camera_names": [], + "camera_height": 84, + "camera_width": 84 + }, + "task": { + "task_spec": { + "subtask_1": { + "object_ref": "square_nut", + "subtask_term_signal": "grasp", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "nearest_neighbor_object", + "selection_strategy_kwargs": { + "nn_k": 3 + }, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_2": { + "object_ref": "square_peg", + "subtask_term_signal": null, + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + } + } + } +} \ No newline at end of file diff --git a/mimicgen/exps/templates/robosuite/stack.json b/mimicgen/exps/templates/robosuite/stack.json new file mode 100644 index 0000000..2d42f08 --- /dev/null +++ b/mimicgen/exps/templates/robosuite/stack.json @@ -0,0 +1,74 @@ +{ + "name": "stack", + "type": "robosuite", + "experiment": { + "name": "demo", + "source": { + "dataset_path": null, + "filter_key": null, + "n": null, + "start": null + }, + "generation": { + "path": null, + "guarantee": false, + "keep_failed": true, + "num_trials": 10, + "select_src_per_subtask": false, + "transform_first_robot_pose": false, + "interpolate_from_last_target_pose": true + }, + "task": { + "name": null, + "robot": null, + "gripper": null, + "interface": null, + "interface_type": null + }, + "max_num_failures": 50, + "render_video": true, + "num_demo_to_render": 50, + "num_fail_demo_to_render": 50, + "log_every_n_attempts": 50, + "seed": 1 + }, + "obs": { + "collect_obs": true, + "camera_names": [], + "camera_height": 84, + "camera_width": 84 + }, + "task": { + "task_spec": { + "subtask_1": { + "object_ref": "cubeA", + "subtask_term_signal": "grasp", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "nearest_neighbor_object", + "selection_strategy_kwargs": { + "nn_k": 3 + }, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_2": { + "object_ref": "cubeB", + "subtask_term_signal": null, + "subtask_term_offset_range": null, + "selection_strategy": "nearest_neighbor_object", + "selection_strategy_kwargs": { + "nn_k": 3 + }, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + } + } + } +} \ No newline at end of file diff --git a/mimicgen/exps/templates/robosuite/stack_three.json b/mimicgen/exps/templates/robosuite/stack_three.json new file mode 100644 index 0000000..8d3b07c --- /dev/null +++ b/mimicgen/exps/templates/robosuite/stack_three.json @@ -0,0 +1,106 @@ +{ + "name": "stack_three", + "type": "robosuite", + "experiment": { + "name": "demo", + "source": { + "dataset_path": null, + "filter_key": null, + "n": null, + "start": null + }, + "generation": { + "path": null, + "guarantee": false, + "keep_failed": true, + "num_trials": 10, + "select_src_per_subtask": false, + "transform_first_robot_pose": false, + "interpolate_from_last_target_pose": true + }, + "task": { + "name": null, + "robot": null, + "gripper": null, + "interface": null, + "interface_type": null + }, + "max_num_failures": 50, + "render_video": true, + "num_demo_to_render": 50, + "num_fail_demo_to_render": 50, + "log_every_n_attempts": 50, + "seed": 1 + }, + "obs": { + "collect_obs": true, + "camera_names": [], + "camera_height": 84, + "camera_width": 84 + }, + "task": { + "task_spec": { + "subtask_1": { + "object_ref": "cubeA", + "subtask_term_signal": "grasp_1", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "nearest_neighbor_object", + "selection_strategy_kwargs": { + "nn_k": 3 + }, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_2": { + "object_ref": "cubeB", + "subtask_term_signal": "stack_1", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "nearest_neighbor_object", + "selection_strategy_kwargs": { + "nn_k": 3 + }, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_3": { + "object_ref": "cubeC", + "subtask_term_signal": "grasp_2", + "subtask_term_offset_range": [ + 10, + 20 + ], + "selection_strategy": "nearest_neighbor_object", + "selection_strategy_kwargs": { + "nn_k": 3 + }, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_4": { + "object_ref": "cubeA", + "subtask_term_signal": null, + "subtask_term_offset_range": null, + "selection_strategy": "nearest_neighbor_object", + "selection_strategy_kwargs": { + "nn_k": 3 + }, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + } + } + } +} \ No newline at end of file diff --git a/mimicgen/exps/templates/robosuite/threading.json b/mimicgen/exps/templates/robosuite/threading.json new file mode 100644 index 0000000..1511901 --- /dev/null +++ b/mimicgen/exps/templates/robosuite/threading.json @@ -0,0 +1,70 @@ +{ + "name": "threading", + "type": "robosuite", + "experiment": { + "name": "demo", + "source": { + "dataset_path": null, + "filter_key": null, + "n": null, + "start": null + }, + "generation": { + "path": null, + "guarantee": false, + "keep_failed": true, + "num_trials": 10, + "select_src_per_subtask": false, + "transform_first_robot_pose": false, + "interpolate_from_last_target_pose": true + }, + "task": { + "name": null, + "robot": null, + "gripper": null, + "interface": null, + "interface_type": null + }, + "max_num_failures": 50, + "render_video": true, + "num_demo_to_render": 50, + "num_fail_demo_to_render": 50, + "log_every_n_attempts": 50, + "seed": 1 + }, + "obs": { + "collect_obs": true, + "camera_names": [], + "camera_height": 84, + "camera_width": 84 + }, + "task": { + "task_spec": { + "subtask_1": { + "object_ref": "needle", + "subtask_term_signal": "grasp", + "subtask_term_offset_range": [ + 5, + 10 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_2": { + "object_ref": "tripod", + "subtask_term_signal": null, + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + } + } + } +} \ No newline at end of file diff --git a/mimicgen/exps/templates/robosuite/three_piece_assembly.json b/mimicgen/exps/templates/robosuite/three_piece_assembly.json new file mode 100644 index 0000000..2683598 --- /dev/null +++ b/mimicgen/exps/templates/robosuite/three_piece_assembly.json @@ -0,0 +1,98 @@ +{ + "name": "three_piece_assembly", + "type": "robosuite", + "experiment": { + "name": "demo", + "source": { + "dataset_path": null, + "filter_key": null, + "n": null, + "start": null + }, + "generation": { + "path": null, + "guarantee": false, + "keep_failed": true, + "num_trials": 10, + "select_src_per_subtask": false, + "transform_first_robot_pose": false, + "interpolate_from_last_target_pose": true + }, + "task": { + "name": null, + "robot": null, + "gripper": null, + "interface": null, + "interface_type": null + }, + "max_num_failures": 50, + "render_video": true, + "num_demo_to_render": 50, + "num_fail_demo_to_render": 50, + "log_every_n_attempts": 50, + "seed": 1 + }, + "obs": { + "collect_obs": true, + "camera_names": [], + "camera_height": 84, + "camera_width": 84 + }, + "task": { + "task_spec": { + "subtask_1": { + "object_ref": "piece_1", + "subtask_term_signal": "grasp_1", + "subtask_term_offset_range": [ + 5, + 10 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_2": { + "object_ref": "base", + "subtask_term_signal": "insert_1", + "subtask_term_offset_range": [ + 5, + 10 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_3": { + "object_ref": "piece_2", + "subtask_term_signal": "grasp_2", + "subtask_term_offset_range": [ + 5, + 10 + ], + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + }, + "subtask_4": { + "object_ref": "piece_1", + "subtask_term_signal": null, + "subtask_term_offset_range": null, + "selection_strategy": "random", + "selection_strategy_kwargs": null, + "action_noise": 0.05, + "num_interpolation_steps": 5, + "num_fixed_steps": 0, + "apply_noise_during_interpolation": false + } + } + } +} \ No newline at end of file diff --git a/mimicgen_envs/models/robosuite/assets/objects/coffee_base.xml b/mimicgen/models/robosuite/assets/objects/coffee_base.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/objects/coffee_base.xml rename to mimicgen/models/robosuite/assets/objects/coffee_base.xml diff --git a/mimicgen_envs/models/robosuite/assets/objects/coffee_body.xml b/mimicgen/models/robosuite/assets/objects/coffee_body.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/objects/coffee_body.xml rename to mimicgen/models/robosuite/assets/objects/coffee_body.xml diff --git a/mimicgen_envs/models/robosuite/assets/objects/coffee_lid.xml b/mimicgen/models/robosuite/assets/objects/coffee_lid.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/objects/coffee_lid.xml rename to mimicgen/models/robosuite/assets/objects/coffee_lid.xml diff --git a/mimicgen_envs/models/robosuite/assets/objects/coffee_pod.xml b/mimicgen/models/robosuite/assets/objects/coffee_pod.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/objects/coffee_pod.xml rename to mimicgen/models/robosuite/assets/objects/coffee_pod.xml diff --git a/mimicgen_envs/models/robosuite/assets/objects/drawer.xml b/mimicgen/models/robosuite/assets/objects/drawer.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/objects/drawer.xml rename to mimicgen/models/robosuite/assets/objects/drawer.xml diff --git a/mimicgen_envs/models/robosuite/assets/objects/drawer_long.xml b/mimicgen/models/robosuite/assets/objects/drawer_long.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/objects/drawer_long.xml rename to mimicgen/models/robosuite/assets/objects/drawer_long.xml diff --git a/mimicgen_envs/models/robosuite/assets/objects/meshes/coffee_base.stl b/mimicgen/models/robosuite/assets/objects/meshes/coffee_base.stl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/objects/meshes/coffee_base.stl rename to mimicgen/models/robosuite/assets/objects/meshes/coffee_base.stl diff --git a/mimicgen_envs/models/robosuite/assets/objects/meshes/coffee_body.stl b/mimicgen/models/robosuite/assets/objects/meshes/coffee_body.stl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/objects/meshes/coffee_body.stl rename to mimicgen/models/robosuite/assets/objects/meshes/coffee_body.stl diff --git a/mimicgen_envs/models/robosuite/assets/objects/meshes/coffee_lid.stl b/mimicgen/models/robosuite/assets/objects/meshes/coffee_lid.stl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/objects/meshes/coffee_lid.stl rename to mimicgen/models/robosuite/assets/objects/meshes/coffee_lid.stl diff --git a/mimicgen_envs/models/robosuite/assets/objects/meshes/coffee_pod.stl b/mimicgen/models/robosuite/assets/objects/meshes/coffee_pod.stl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/objects/meshes/coffee_pod.stl rename to mimicgen/models/robosuite/assets/objects/meshes/coffee_pod.stl diff --git a/mimicgen_envs/models/robosuite/assets/objects/serving_region.xml b/mimicgen/models/robosuite/assets/objects/serving_region.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/objects/serving_region.xml rename to mimicgen/models/robosuite/assets/objects/serving_region.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_10.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_10.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_10.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_10.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_11.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_11.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_11.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_11.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_12.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_12.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_12.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_12.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_13.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_13.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_13.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_13.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_14.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_14.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_14.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_14.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_15.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_15.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_15.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_15.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_16.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_16.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_16.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_16.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_17.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_17.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_17.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_17.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_18.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_18.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_18.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_18.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_19.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_19.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_19.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_19.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_20.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_20.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_20.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_20.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_21.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_21.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_21.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_21.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_22.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_22.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_22.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_22.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_23.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_23.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_23.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_23.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_24.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_24.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_24.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_24.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_25.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_25.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_25.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_25.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_26.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_26.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_26.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_26.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_27.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_27.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_27.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_27.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_28.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_28.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_28.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_28.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_29.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_29.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_29.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_29.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_30.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_30.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_30.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_30.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_31.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_31.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_31.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_31.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_5.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_5.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_5.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_5.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_6.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_6.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_6.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_6.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_7.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_7.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_7.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_7.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_8.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_8.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_8.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_8.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_9.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_9.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_9.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/collision/model_normalized_collision_9.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/meta.json b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/meta.json similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/meta.json rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/meta.json diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/model.xml b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/model.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/model.xml rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/model.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/material_0.mtl b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/material_0.mtl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/material_0.mtl rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/material_0.mtl diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/model_normalized_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/model_normalized_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/model_normalized_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/model_normalized_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/model_normalized_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/model_normalized_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/model_normalized_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/model_normalized_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/model_normalized_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/model_normalized_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/model_normalized_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/128ecbc1/visual/model_normalized_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_10.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_10.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_10.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_10.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_11.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_11.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_11.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_11.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_12.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_12.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_12.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_12.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_13.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_13.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_13.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_13.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_14.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_14.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_14.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_14.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_15.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_15.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_15.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_15.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_16.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_16.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_16.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_16.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_17.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_17.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_17.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_17.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_18.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_18.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_18.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_18.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_19.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_19.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_19.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_19.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_20.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_20.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_20.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_20.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_21.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_21.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_21.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_21.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_22.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_22.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_22.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_22.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_23.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_23.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_23.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_23.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_24.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_24.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_24.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_24.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_25.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_25.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_25.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_25.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_26.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_26.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_26.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_26.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_27.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_27.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_27.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_27.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_28.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_28.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_28.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_28.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_29.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_29.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_29.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_29.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_30.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_30.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_30.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_30.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_31.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_31.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_31.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_31.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_5.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_5.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_5.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_5.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_6.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_6.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_6.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_6.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_7.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_7.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_7.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_7.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_8.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_8.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_8.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_8.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_9.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_9.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_9.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/collision/model_normalized_collision_9.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/meta.json b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/meta.json similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/meta.json rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/meta.json diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/model.xml b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/model.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/model.xml rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/model.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/visual/material_0.mtl b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/visual/material_0.mtl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/visual/material_0.mtl rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/visual/material_0.mtl diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/visual/model_normalized_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/visual/model_normalized_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/3143a4ac/visual/model_normalized_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/3143a4ac/visual/model_normalized_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_10.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_10.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_10.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_10.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_11.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_11.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_11.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_11.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_12.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_12.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_12.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_12.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_13.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_13.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_13.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_13.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_14.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_14.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_14.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_14.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_15.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_15.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_15.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_15.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_16.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_16.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_16.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_16.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_17.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_17.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_17.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_17.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_18.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_18.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_18.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_18.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_19.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_19.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_19.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_19.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_20.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_20.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_20.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_20.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_21.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_21.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_21.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_21.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_22.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_22.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_22.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_22.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_23.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_23.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_23.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_23.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_24.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_24.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_24.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_24.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_25.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_25.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_25.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_25.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_26.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_26.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_26.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_26.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_27.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_27.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_27.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_27.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_28.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_28.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_28.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_28.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_29.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_29.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_29.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_29.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_30.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_30.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_30.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_30.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_31.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_31.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_31.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_31.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_5.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_5.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_5.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_5.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_6.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_6.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_6.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_6.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_7.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_7.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_7.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_7.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_8.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_8.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_8.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_8.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_9.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_9.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_9.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/collision/model_normalized_collision_9.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/meta.json b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/meta.json similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/meta.json rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/meta.json diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/model.xml b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/model.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/model.xml rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/model.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/material_0.jpeg b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/material_0.jpeg similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/material_0.jpeg rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/material_0.jpeg diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/material_0.mtl b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/material_0.mtl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/material_0.mtl rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/material_0.mtl diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/model_normalized_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture0.png b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture0.png similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture0.png rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture0.png diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture1.png b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture1.png similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture1.png rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture1.png diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture2.png b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture2.png similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture2.png rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture2.png diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture3.png b/mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture3.png similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture3.png rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/345d3e72/visual/texture3.png diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_10.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_10.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_10.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_10.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_11.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_11.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_11.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_11.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_12.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_12.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_12.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_12.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_13.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_13.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_13.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_13.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_14.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_14.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_14.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_14.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_15.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_15.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_15.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_15.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_16.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_16.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_16.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_16.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_17.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_17.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_17.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_17.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_18.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_18.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_18.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_18.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_19.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_19.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_19.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_19.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_20.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_20.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_20.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_20.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_21.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_21.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_21.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_21.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_22.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_22.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_22.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_22.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_23.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_23.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_23.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_23.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_24.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_24.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_24.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_24.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_25.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_25.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_25.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_25.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_26.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_26.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_26.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_26.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_27.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_27.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_27.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_27.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_28.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_28.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_28.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_28.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_29.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_29.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_29.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_29.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_30.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_30.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_30.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_30.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_31.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_31.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_31.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_31.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_5.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_5.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_5.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_5.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_6.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_6.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_6.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_6.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_7.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_7.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_7.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_7.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_8.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_8.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_8.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_8.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_9.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_9.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_9.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/collision/model_normalized_collision_9.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/meta.json b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/meta.json similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/meta.json rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/meta.json diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/model.xml b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/model.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/model.xml rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/model.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/material_0.mtl b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/material_0.mtl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/material_0.mtl rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/material_0.mtl diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/model_normalized_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/model_normalized_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/model_normalized_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/model_normalized_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/model_normalized_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/model_normalized_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/model_normalized_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/model_normalized_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/model_normalized_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/model_normalized_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/model_normalized_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/34ae0b61/visual/model_normalized_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_10.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_10.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_10.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_10.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_11.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_11.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_11.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_11.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_12.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_12.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_12.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_12.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_13.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_13.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_13.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_13.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_14.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_14.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_14.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_14.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_15.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_15.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_15.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_15.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_16.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_16.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_16.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_16.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_17.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_17.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_17.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_17.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_18.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_18.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_18.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_18.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_19.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_19.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_19.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_19.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_20.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_20.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_20.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_20.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_21.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_21.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_21.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_21.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_22.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_22.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_22.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_22.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_23.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_23.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_23.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_23.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_24.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_24.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_24.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_24.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_25.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_25.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_25.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_25.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_26.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_26.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_26.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_26.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_27.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_27.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_27.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_27.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_28.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_28.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_28.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_28.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_29.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_29.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_29.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_29.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_30.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_30.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_30.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_30.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_31.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_31.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_31.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_31.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_5.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_5.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_5.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_5.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_6.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_6.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_6.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_6.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_7.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_7.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_7.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_7.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_8.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_8.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_8.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_8.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_9.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_9.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_9.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/collision/model_normalized_collision_9.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/meta.json b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/meta.json similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/meta.json rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/meta.json diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/model.xml b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/model.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/model.xml rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/model.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/visual/material_0.mtl b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/visual/material_0.mtl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/visual/material_0.mtl rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/visual/material_0.mtl diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/visual/model_normalized_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/visual/model_normalized_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/visual/model_normalized_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/visual/model_normalized_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/visual/model_normalized_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/visual/model_normalized_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/48e260a6/visual/model_normalized_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/48e260a6/visual/model_normalized_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_10.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_10.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_10.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_10.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_11.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_11.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_11.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_11.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_12.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_12.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_12.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_12.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_13.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_13.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_13.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_13.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_14.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_14.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_14.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_14.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_15.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_15.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_15.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_15.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_16.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_16.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_16.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_16.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_17.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_17.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_17.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_17.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_18.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_18.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_18.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_18.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_19.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_19.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_19.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_19.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_20.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_20.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_20.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_20.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_21.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_21.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_21.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_21.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_22.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_22.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_22.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_22.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_23.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_23.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_23.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_23.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_24.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_24.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_24.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_24.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_25.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_25.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_25.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_25.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_26.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_26.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_26.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_26.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_27.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_27.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_27.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_27.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_28.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_28.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_28.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_28.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_29.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_29.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_29.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_29.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_30.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_30.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_30.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_30.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_31.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_31.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_31.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_31.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_5.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_5.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_5.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_5.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_6.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_6.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_6.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_6.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_7.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_7.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_7.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_7.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_8.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_8.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_8.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_8.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_9.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_9.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_9.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/collision/model_normalized_collision_9.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/meta.json b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/meta.json similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/meta.json rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/meta.json diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/model.xml b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/model.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/model.xml rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/model.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/material_0.mtl b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/material_0.mtl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/material_0.mtl rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/material_0.mtl diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/model_normalized_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/model_normalized_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/model_normalized_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/model_normalized_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/model_normalized_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/model_normalized_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/model_normalized_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/model_normalized_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/model_normalized_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/model_normalized_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/model_normalized_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/5fe74bab/visual/model_normalized_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_10.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_10.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_10.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_10.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_11.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_11.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_11.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_11.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_12.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_12.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_12.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_12.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_13.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_13.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_13.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_13.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_14.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_14.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_14.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_14.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_15.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_15.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_15.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_15.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_16.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_16.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_16.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_16.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_17.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_17.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_17.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_17.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_18.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_18.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_18.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_18.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_19.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_19.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_19.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_19.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_20.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_20.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_20.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_20.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_21.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_21.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_21.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_21.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_22.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_22.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_22.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_22.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_23.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_23.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_23.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_23.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_24.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_24.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_24.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_24.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_25.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_25.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_25.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_25.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_26.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_26.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_26.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_26.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_27.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_27.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_27.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_27.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_28.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_28.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_28.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_28.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_29.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_29.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_29.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_29.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_30.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_30.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_30.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_30.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_31.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_31.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_31.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_31.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_5.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_5.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_5.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_5.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_6.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_6.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_6.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_6.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_7.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_7.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_7.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_7.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_8.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_8.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_8.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_8.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_9.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_9.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_9.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/collision/model_normalized_collision_9.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/meta.json b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/meta.json similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/meta.json rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/meta.json diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/model.xml b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/model.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/model.xml rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/model.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/material_0.jpeg b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/material_0.jpeg similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/material_0.jpeg rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/material_0.jpeg diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/material_0.mtl b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/material_0.mtl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/material_0.mtl rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/material_0.mtl diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/model_normalized_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/texture0.png b/mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/texture0.png similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/texture0.png rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/8012f52d/visual/texture0.png diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_10.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_10.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_10.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_10.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_11.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_11.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_11.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_11.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_12.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_12.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_12.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_12.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_13.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_13.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_13.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_13.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_14.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_14.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_14.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_14.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_15.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_15.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_15.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_15.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_16.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_16.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_16.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_16.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_17.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_17.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_17.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_17.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_18.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_18.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_18.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_18.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_19.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_19.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_19.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_19.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_20.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_20.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_20.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_20.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_21.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_21.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_21.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_21.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_22.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_22.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_22.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_22.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_23.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_23.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_23.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_23.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_24.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_24.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_24.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_24.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_25.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_25.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_25.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_25.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_26.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_26.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_26.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_26.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_27.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_27.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_27.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_27.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_28.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_28.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_28.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_28.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_29.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_29.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_29.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_29.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_30.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_30.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_30.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_30.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_31.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_31.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_31.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_31.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_5.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_5.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_5.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_5.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_6.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_6.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_6.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_6.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_7.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_7.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_7.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_7.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_8.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_8.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_8.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_8.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_9.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_9.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_9.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/collision/model_normalized_collision_9.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/meta.json b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/meta.json similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/meta.json rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/meta.json diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/model.xml b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/model.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/model.xml rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/model.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/visual/material_0.mtl b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/visual/material_0.mtl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/visual/material_0.mtl rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/visual/material_0.mtl diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/visual/model_normalized_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/visual/model_normalized_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/visual/model_normalized_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/visual/model_normalized_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/visual/model_normalized_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/visual/model_normalized_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/visual/model_normalized_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/b4ae56d6/visual/model_normalized_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_10.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_10.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_10.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_10.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_11.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_11.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_11.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_11.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_12.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_12.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_12.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_12.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_13.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_13.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_13.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_13.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_14.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_14.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_14.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_14.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_15.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_15.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_15.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_15.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_16.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_16.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_16.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_16.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_17.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_17.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_17.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_17.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_18.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_18.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_18.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_18.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_19.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_19.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_19.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_19.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_20.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_20.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_20.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_20.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_21.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_21.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_21.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_21.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_22.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_22.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_22.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_22.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_23.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_23.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_23.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_23.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_24.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_24.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_24.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_24.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_25.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_25.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_25.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_25.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_26.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_26.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_26.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_26.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_27.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_27.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_27.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_27.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_28.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_28.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_28.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_28.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_29.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_29.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_29.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_29.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_30.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_30.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_30.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_30.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_31.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_31.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_31.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_31.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_5.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_5.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_5.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_5.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_6.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_6.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_6.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_6.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_7.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_7.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_7.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_7.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_8.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_8.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_8.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_8.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_9.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_9.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_9.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/collision/model_normalized_collision_9.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/meta.json b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/meta.json similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/meta.json rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/meta.json diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/model.xml b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/model.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/model.xml rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/model.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/material_0.jpeg b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/material_0.jpeg similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/material_0.jpeg rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/material_0.jpeg diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/material_0.mtl b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/material_0.mtl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/material_0.mtl rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/material_0.mtl diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/model_normalized_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/model_normalized_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/model_normalized_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/model_normalized_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/model_normalized_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/model_normalized_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/model_normalized_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/model_normalized_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/model_normalized_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/model_normalized_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/model_normalized_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/model_normalized_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/texture0.png b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/texture0.png similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/texture0.png rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/texture0.png diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/texture1.png b/mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/texture1.png similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/texture1.png rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/c2eacc52/visual/texture1.png diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_10.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_10.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_10.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_10.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_11.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_11.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_11.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_11.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_12.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_12.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_12.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_12.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_13.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_13.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_13.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_13.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_14.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_14.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_14.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_14.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_15.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_15.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_15.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_15.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_16.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_16.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_16.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_16.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_17.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_17.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_17.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_17.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_18.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_18.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_18.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_18.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_19.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_19.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_19.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_19.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_20.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_20.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_20.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_20.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_21.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_21.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_21.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_21.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_22.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_22.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_22.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_22.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_23.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_23.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_23.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_23.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_24.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_24.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_24.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_24.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_25.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_25.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_25.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_25.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_26.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_26.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_26.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_26.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_27.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_27.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_27.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_27.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_28.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_28.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_28.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_28.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_29.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_29.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_29.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_29.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_30.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_30.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_30.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_30.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_31.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_31.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_31.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_31.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_5.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_5.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_5.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_5.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_6.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_6.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_6.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_6.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_7.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_7.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_7.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_7.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_8.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_8.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_8.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_8.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_9.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_9.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_9.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/collision/model_normalized_collision_9.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/meta.json b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/meta.json similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/meta.json rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/meta.json diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/model.xml b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/model.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/model.xml rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/model.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/material_0.mtl b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/material_0.mtl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/material_0.mtl rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/material_0.mtl diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/d75af64a/visual/model_normalized_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_10.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_10.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_10.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_10.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_11.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_11.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_11.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_11.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_12.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_12.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_12.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_12.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_13.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_13.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_13.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_13.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_14.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_14.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_14.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_14.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_15.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_15.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_15.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_15.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_16.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_16.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_16.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_16.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_17.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_17.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_17.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_17.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_18.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_18.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_18.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_18.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_19.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_19.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_19.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_19.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_20.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_20.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_20.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_20.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_21.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_21.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_21.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_21.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_22.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_22.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_22.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_22.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_23.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_23.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_23.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_23.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_24.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_24.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_24.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_24.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_25.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_25.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_25.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_25.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_26.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_26.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_26.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_26.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_27.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_27.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_27.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_27.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_28.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_28.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_28.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_28.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_29.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_29.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_29.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_29.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_30.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_30.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_30.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_30.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_31.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_31.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_31.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_31.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_5.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_5.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_5.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_5.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_6.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_6.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_6.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_6.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_7.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_7.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_7.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_7.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_8.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_8.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_8.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_8.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_9.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_9.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_9.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/collision/model_normalized_collision_9.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/meta.json b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/meta.json similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/meta.json rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/meta.json diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/model.xml b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/model.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/model.xml rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/model.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/visual/material_0.mtl b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/visual/material_0.mtl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/visual/material_0.mtl rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/visual/material_0.mtl diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/visual/model_normalized_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/visual/model_normalized_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/visual/model_normalized_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/visual/model_normalized_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/visual/model_normalized_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/visual/model_normalized_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/e94e46bc/visual/model_normalized_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/e94e46bc/visual/model_normalized_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_1.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_1.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_1.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_1.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_10.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_10.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_10.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_10.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_11.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_11.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_11.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_11.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_12.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_12.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_12.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_12.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_13.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_13.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_13.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_13.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_14.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_14.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_14.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_14.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_15.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_15.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_15.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_15.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_16.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_16.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_16.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_16.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_17.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_17.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_17.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_17.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_18.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_18.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_18.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_18.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_19.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_19.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_19.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_19.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_2.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_2.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_2.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_2.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_20.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_20.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_20.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_20.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_21.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_21.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_21.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_21.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_22.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_22.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_22.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_22.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_23.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_23.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_23.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_23.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_24.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_24.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_24.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_24.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_25.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_25.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_25.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_25.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_26.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_26.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_26.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_26.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_27.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_27.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_27.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_27.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_28.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_28.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_28.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_28.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_29.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_29.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_29.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_29.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_3.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_3.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_3.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_3.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_30.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_30.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_30.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_30.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_31.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_31.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_31.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_31.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_4.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_4.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_4.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_4.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_5.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_5.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_5.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_5.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_6.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_6.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_6.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_6.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_7.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_7.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_7.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_7.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_8.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_8.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_8.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_8.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_9.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_9.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_9.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/collision/model_normalized_collision_9.obj diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/meta.json b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/meta.json similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/meta.json rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/meta.json diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/model.xml b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/model.xml similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/model.xml rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/model.xml diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/visual/material_0.mtl b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/visual/material_0.mtl similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/visual/material_0.mtl rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/visual/material_0.mtl diff --git a/mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/visual/model_normalized_0.obj b/mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/visual/model_normalized_0.obj similarity index 100% rename from mimicgen_envs/models/robosuite/assets/shapenet_core/mugs/fad118b3/visual/model_normalized_0.obj rename to mimicgen/models/robosuite/assets/shapenet_core/mugs/fad118b3/visual/model_normalized_0.obj diff --git a/mimicgen_envs/models/robosuite/assets/textures/ceramic.png b/mimicgen/models/robosuite/assets/textures/ceramic.png similarity index 100% rename from mimicgen_envs/models/robosuite/assets/textures/ceramic.png rename to mimicgen/models/robosuite/assets/textures/ceramic.png diff --git a/mimicgen_envs/models/robosuite/objects/__init__.py b/mimicgen/models/robosuite/objects/__init__.py similarity index 72% rename from mimicgen_envs/models/robosuite/objects/__init__.py rename to mimicgen/models/robosuite/objects/__init__.py index 7df5f0e..2b1a797 100644 --- a/mimicgen_envs/models/robosuite/objects/__init__.py +++ b/mimicgen/models/robosuite/objects/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. diff --git a/mimicgen_envs/models/robosuite/objects/composite/__init__.py b/mimicgen/models/robosuite/objects/composite/__init__.py similarity index 80% rename from mimicgen_envs/models/robosuite/objects/composite/__init__.py rename to mimicgen/models/robosuite/objects/composite/__init__.py index 6bbdf49..4542ffe 100644 --- a/mimicgen_envs/models/robosuite/objects/composite/__init__.py +++ b/mimicgen/models/robosuite/objects/composite/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. diff --git a/mimicgen_envs/models/robosuite/objects/composite/box_pattern_object.py b/mimicgen/models/robosuite/objects/composite/box_pattern_object.py similarity index 98% rename from mimicgen_envs/models/robosuite/objects/composite/box_pattern_object.py rename to mimicgen/models/robosuite/objects/composite/box_pattern_object.py index 15a37ee..114ac98 100644 --- a/mimicgen_envs/models/robosuite/objects/composite/box_pattern_object.py +++ b/mimicgen/models/robosuite/objects/composite/box_pattern_object.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. diff --git a/mimicgen_envs/models/robosuite/objects/composite/hollow_cylinder.py b/mimicgen/models/robosuite/objects/composite/hollow_cylinder.py similarity index 98% rename from mimicgen_envs/models/robosuite/objects/composite/hollow_cylinder.py rename to mimicgen/models/robosuite/objects/composite/hollow_cylinder.py index e819fc4..88de19c 100644 --- a/mimicgen_envs/models/robosuite/objects/composite/hollow_cylinder.py +++ b/mimicgen/models/robosuite/objects/composite/hollow_cylinder.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. diff --git a/mimicgen_envs/models/robosuite/objects/composite/needle.py b/mimicgen/models/robosuite/objects/composite/needle.py similarity index 95% rename from mimicgen_envs/models/robosuite/objects/composite/needle.py rename to mimicgen/models/robosuite/objects/composite/needle.py index 7777f9b..f50d9bc 100644 --- a/mimicgen_envs/models/robosuite/objects/composite/needle.py +++ b/mimicgen/models/robosuite/objects/composite/needle.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -22,8 +22,6 @@ def __init__( name, ): - ### TODO: make this object more general (with more args and configuration options) later ### - # Set object attributes self._name = name self.needle_mat_name = "darkwood_mat" diff --git a/mimicgen_envs/models/robosuite/objects/composite/ring_tripod.py b/mimicgen/models/robosuite/objects/composite/ring_tripod.py similarity index 97% rename from mimicgen_envs/models/robosuite/objects/composite/ring_tripod.py rename to mimicgen/models/robosuite/objects/composite/ring_tripod.py index 0644cf6..4c8d573 100644 --- a/mimicgen_envs/models/robosuite/objects/composite/ring_tripod.py +++ b/mimicgen/models/robosuite/objects/composite/ring_tripod.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -23,8 +23,6 @@ def __init__( name, ): - ### TODO: make this object more general (with more args and configuration options) later ### - # Set object attributes self._name = name self.tripod_mat_name = "lightwood_mat" diff --git a/mimicgen_envs/models/robosuite/objects/composite_body/__init__.py b/mimicgen/models/robosuite/objects/composite_body/__init__.py similarity index 71% rename from mimicgen_envs/models/robosuite/objects/composite_body/__init__.py rename to mimicgen/models/robosuite/objects/composite_body/__init__.py index 7a250c3..493c2c6 100644 --- a/mimicgen_envs/models/robosuite/objects/composite_body/__init__.py +++ b/mimicgen/models/robosuite/objects/composite_body/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. diff --git a/mimicgen_envs/models/robosuite/objects/composite_body/coffee_machine.py b/mimicgen/models/robosuite/objects/composite_body/coffee_machine.py similarity index 95% rename from mimicgen_envs/models/robosuite/objects/composite_body/coffee_machine.py rename to mimicgen/models/robosuite/objects/composite_body/coffee_machine.py index fa76ae2..48c7777 100644 --- a/mimicgen_envs/models/robosuite/objects/composite_body/coffee_machine.py +++ b/mimicgen/models/robosuite/objects/composite_body/coffee_machine.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -9,8 +9,8 @@ from robosuite.models.objects import CompositeBodyObject, BoxObject -from mimicgen_envs.models.robosuite.objects.composite_body.cup import CupObject -from mimicgen_envs.models.robosuite.objects.xml_objects import CoffeeMachineBodyObject, CoffeeMachineLidObject, CoffeeMachineBaseObject +from mimicgen.models.robosuite.objects.composite_body.cup import CupObject +from mimicgen.models.robosuite.objects.xml_objects import CoffeeMachineBodyObject, CoffeeMachineLidObject, CoffeeMachineBaseObject class CoffeeMachineObject(CompositeBodyObject): diff --git a/mimicgen_envs/models/robosuite/objects/composite_body/cup.py b/mimicgen/models/robosuite/objects/composite_body/cup.py similarity index 97% rename from mimicgen_envs/models/robosuite/objects/composite_body/cup.py rename to mimicgen/models/robosuite/objects/composite_body/cup.py index a5b9229..0b0d4c7 100644 --- a/mimicgen_envs/models/robosuite/objects/composite_body/cup.py +++ b/mimicgen/models/robosuite/objects/composite_body/cup.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -9,7 +9,7 @@ from robosuite.utils.mjcf_utils import RED, BLUE, CustomMaterial from robosuite.models.objects import CompositeBodyObject, BoxObject, CylinderObject -from mimicgen_envs.models.robosuite.objects import HollowCylinderObject +from mimicgen.models.robosuite.objects import HollowCylinderObject class CupObject(CompositeBodyObject): diff --git a/mimicgen_envs/models/robosuite/objects/xml_objects.py b/mimicgen/models/robosuite/objects/xml_objects.py similarity index 97% rename from mimicgen_envs/models/robosuite/objects/xml_objects.py rename to mimicgen/models/robosuite/objects/xml_objects.py index 702f45f..d317c3a 100644 --- a/mimicgen_envs/models/robosuite/objects/xml_objects.py +++ b/mimicgen/models/robosuite/objects/xml_objects.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -13,9 +13,9 @@ from robosuite.models.objects import MujocoXMLObject from robosuite.utils.mjcf_utils import string_to_array, array_to_string -import mimicgen_envs +import mimicgen -XML_ASSETS_BASE_PATH = os.path.join(mimicgen_envs.__path__[0], "models/robosuite/assets") +XML_ASSETS_BASE_PATH = os.path.join(mimicgen.__path__[0], "models/robosuite/assets") class BlenderObject(MujocoXMLObject): @@ -222,5 +222,4 @@ def top_offset(self): @property def horizontal_radius(self): - # TODO: this might need to change now return 0.15 diff --git a/mimicgen/scripts/__init__.py b/mimicgen/scripts/__init__.py new file mode 100644 index 0000000..8edfcfd --- /dev/null +++ b/mimicgen/scripts/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. \ No newline at end of file diff --git a/mimicgen/scripts/annotate_subtasks.py b/mimicgen/scripts/annotate_subtasks.py new file mode 100644 index 0000000..9847285 --- /dev/null +++ b/mimicgen/scripts/annotate_subtasks.py @@ -0,0 +1,414 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +A script to playback demonstrations (using visual observations and the pygame renderer) +in order to allow a user to annotate portions of the demonstrations. This is useful +to annotate the end of each object-centric subtask in each source demonstration used +by MimicGen, as an alternative to implementing subtask termination signals directly +in the simulation environment. + +Examples: + + # specify the sequence of signals that should be annotated and the dataset images to render on-screen + python annotate_subtasks.py --dataset /path/to/demo.hdf5 --signals grasp_1 insert_1 grasp_2 \ + --render_image_names agentview_image robot0_eye_in_hand_image + + # limit annotation to first 2 demos + python annotate_subtasks.py --dataset /path/to/demo.hdf5 --signals grasp_1 insert_1 grasp_2 \ + --render_image_names agentview_image robot0_eye_in_hand_image --n 2 + + # limit annotation to demo 2 and 3 + python annotate_subtasks.py --dataset /path/to/demo.hdf5 --signals grasp_1 insert_1 grasp_2 \ + --render_image_names agentview_image robot0_eye_in_hand_image --n 2 --start 1 + + # scale up dataset images when rendering to screen by factor of 10 + python annotate_subtasks.py --dataset /path/to/demo.hdf5 --signals grasp_1 insert_1 grasp_2 \ + --render_image_names agentview_image robot0_eye_in_hand_image --image_scale 10 + +""" + +import os +import sys +import h5py +import argparse +import imageio +import numpy as np + +# for rendering images on-screen +import cv2 +try: + import pygame +except ImportError as e: + print("Got error: {}".format(e)) + print("") + print("pygame is required. Please install with `pip install pygame`") + +import robomimic +from robomimic.utils.file_utils import get_env_metadata_from_dataset + +import mimicgen +import mimicgen.utils.file_utils as MG_FileUtils +import mimicgen.utils.misc_utils as MiscUtils + + +# scaling size for images when rendering to screen +# IMAGE_SCALE = 10 +IMAGE_SCALE = 5 +# IMAGE_SCALE = 1 + +# Grid of playback rates for the user to cycle through (e.g. 1hz, 5 hz, ...) +RATE_GRID = MiscUtils.Grid( + values=[1, 5, 10, 20, 40], + initial_ind=0, +) + + +def print_keyboard_commands(): + """ + Helper function to print keyboard annotation commands. + """ + def print_command(char, info): + char += " " * (11 - len(char)) + print("{}\t{}".format(char, info)) + + print("") + print_command("Keys", "Command") + print_command("up-down", "increase / decrease playback speed") + print_command("left-right", "seek left / right by N frames") + print_command("spacebar", "press and release to annotate the end of a subtask") + print_command("f", "next demo and save annotations") + print_command("r", "repeat demo and clear annotations") + print("") + + +def make_pygame_screen( + traj_grp, + image_names, + image_scale, +): + """ + Makes pygame screen. + + Args: + traj_grp (h5py.Group): group for a demonstration trajectory + image_names (list): list of image names that will be used for rendering + image_scale (int): scaling factor for the image to diplay in window + + Returns: + screen: pygame screen object + """ + # grab first image from all image modalities to infer size of window + im = [traj_grp["obs/{}".format(k)][0] for k in image_names] + frame = np.concatenate(im, axis=1) + width, height = frame.shape[:2] + width *= image_scale + height *= image_scale + screen = pygame.display.set_mode((height, width)) + return screen + + +def handle_pygame_events( + frame_ind, + subtask_signals, + subtask_ind, + rate_obj, + need_repeat, + annotation_done, + playback_rate_grid, +): + """ + Reads events from pygame window in order to provide the + following keyboard annotation functionality: + + up-down | increase / decrease playback speed + left-right | seek left / right by N frames + spacebar | press and release to annotate the end of a subtask + f | next demo and save annotations + r | repeat demo and clear annotations + + Args: + frame_ind (int): index of current frame in demonstration + subtask_signals (list): list of subtask termination signals that we will annotate + subtask_ind (int): index of current subtask (state variable) + rate_obj (Rate): rate object to maintain playback rate + need_repeat (bool): whether the demo should be repeated (state variable) + annotation_done (bool): whether user is done annotating this demo (state variable) + playback_rate_grid (Grid): grid object to easily toggle between different playback rates + + Returns: + subtask_end_ind (int or None): end index for current subtask, annotated by human, or None if no annotation + subtask_ind (int): possibly updated subtask index + need_repeat (bool): possibly updated value + annotation_done (bool): possibly updated value + seek (int): how much to seek forward or backward in demonstration (value read from user command) + """ + + subtask_end_ind = None + seek = 0 + for event in pygame.event.get(): + + if event.type == pygame.QUIT: + sys.exit() + + if event.type == pygame.KEYUP: + # print("released key {}".format(event.key)) + if event.key == pygame.K_SPACE: + # annotate end of current subtask and move on to next subtask + subtask_end_ind = frame_ind + 1 + print("") + print("*" * 50) + print("cmd: annotated end of subtask {} (signal {}) at index {}".format(subtask_ind, subtask_signals[subtask_ind], frame_ind + 1)) + print("*" * 50) + print("") + subtask_ind += 1 + elif event.key == pygame.K_UP: + # speed up traversal + rate_obj.update_hz(playback_rate_grid.next()) + print("cmd: playback rate increased to {} hz".format(rate_obj.hz)) + elif event.key == pygame.K_DOWN: + # slow down traversal + rate_obj.update_hz(playback_rate_grid.prev()) + print("cmd: playback rate decreased to {} hz".format(rate_obj.hz)) + elif event.key == pygame.K_LEFT: + # seek left + seek = -10 + print("cmd: seek {} frames".format(seek)) + elif event.key == pygame.K_RIGHT: + # seek right + seek = 10 + print("cmd: seek {} frames".format(seek)) + elif event.key == pygame.K_r: + # repeat demo + need_repeat = True + print("cmd: repeat demo") + elif event.key == pygame.K_f: + # next demo + annotation_done = True + print("cmd: next demo") + + return subtask_end_ind, subtask_ind, need_repeat, annotation_done, seek + + +def annotate_subtasks_in_trajectory( + ep, + traj_grp, + subtask_signals, + screen, + video_skip, + image_names, + playback_rate_grid, +): + """ + This function reads all "rgb" observations in the dataset trajectory and + writes them into a video. + + Args: + ep (str): name of hdf5 group for this demo + traj_grp (hdf5 file group): hdf5 group which corresponds to the dataset trajectory to annotate + subtask_signals (list): list of subtask termination signals that will be annotated + screen: pygame screen + video_skip (int): determines rate at which environment frames are written to video + image_names (list): determines which image observations are used for rendering. Pass more than + one to output a video with multiple image observations concatenated horizontally. + playback_rate_grid (Grid): grid object to easily toggle between different playback rates + """ + assert image_names is not None, "error: must specify at least one image observation to use in @image_names" + + traj_len = traj_grp["actions"].shape[0] + + rate_obj = MiscUtils.Rate(hz=playback_rate_grid.get()) + rate_measure = MiscUtils.RateMeasure(name="rate_measure") + + # repeat this demonstration until we have permission to move on + annotation_done = False + while not annotation_done: + print("Starting annotation for demo: {}".format(ep)) + print_keyboard_commands() + + need_repeat = False + subtask_end_inds = [] + + # keep looping through the video, reading user input from keyboard, until + # user indicates that demo is done being annotated + frame_ind = 0 + subtask_ind = 0 + should_add_border_to_frame = False + while (not need_repeat) and (not annotation_done): + + # maybe render frame to screen + if frame_ind % video_skip == 0: + # concatenate image obs together + im = [traj_grp["obs/{}".format(k)][frame_ind] for k in image_names] + frame = np.concatenate(im, axis=1) + # upscale frame to appropriate resolution + frame = cv2.resize(frame, + dsize=(frame.shape[1] * IMAGE_SCALE, frame.shape[0] * IMAGE_SCALE), + interpolation=cv2.INTER_CUBIC) + # maybe add red border + if should_add_border_to_frame: + frame = MiscUtils.add_red_border_to_frame(frame) + # write frame to window + frame = frame.transpose((1, 0, 2)) + pygame.pixelcopy.array_to_surface(screen, frame) + pygame.display.update() + + subtask_end_ind, subtask_ind, need_repeat, annotation_done, seek = handle_pygame_events( + frame_ind=frame_ind, + subtask_signals=subtask_signals, + subtask_ind=subtask_ind, + rate_obj=rate_obj, + need_repeat=need_repeat, + annotation_done=annotation_done, + playback_rate_grid=playback_rate_grid, + ) + + if subtask_end_ind is not None: + # store new annotation and toggle rendering of red border + subtask_end_inds.append(subtask_end_ind) + should_add_border_to_frame = (not should_add_border_to_frame) + + # try to enforce rate + rate_obj.sleep() + rate_measure.measure() + + # increment frame index appropriately (either by 1 or by seek amount), then + # clamp within bounds + mask = int(seek != 0) + frame_ind += (1 - mask) * 1 + mask * seek + frame_ind = max(min(frame_ind, traj_len - 1), 0) + + # if we don't need to repeat the demo, we're done + annotation_done = annotation_done or (not need_repeat) + + # check that we got the right number of annotations + if len(subtask_end_inds) != len(subtask_signals): + print("") + print("*" * 50) + print("Number of termination annotations {} does not match expected number {}...".format(len(subtask_end_inds), len(subtask_signals))) + print("Repeating annotation.") + print("*" * 50) + print("") + annotation_done = False + + # write subtask_termination_signals to hdf5 + assert len(subtask_end_inds) == len(subtask_signals) + + if "subtask_term_signals" in traj_grp["datagen_info"]: + del traj_grp["datagen_info"]["subtask_term_signals"] + for subtask_ind in range(len(subtask_signals)): + # subtask termination signal is 0 until subtask is complete, and 1 afterwards + subtask_signal_array = np.ones(traj_len, dtype=int) + subtask_signal_array[:subtask_end_inds[subtask_ind]] = 0 + traj_grp.create_dataset("datagen_info/subtask_term_signals/{}".format(subtask_signals[subtask_ind]), data=subtask_signal_array) + + # report rate measurements + print("\nFrame Rate (Hz) Statistics for demo {} annotation".format(ep)) + print(rate_measure) + + +def annotate_subtasks(args): + # Auto-fill camera rendering info if not specified + if args.render_image_names is None: + env_meta = get_env_metadata_from_dataset(dataset_path=args.dataset) + args.render_image_names = RobomimicUtils.get_default_env_cameras(env_meta=env_meta) + + # get demonstrations to annotate + dataset_path = args.dataset + demo_keys = MG_FileUtils.get_all_demos_from_dataset( + dataset_path=dataset_path, + filter_key=args.filter_key, + start=args.start, + n=args.n, + ) + + # Verify that the dataset has been processed and has datagen_info. + MG_FileUtils.get_env_interface_info_from_dataset( + dataset_path=dataset_path, + demo_keys=demo_keys, + ) + + # Open the file in read-write mode to add in annotations as subtask_term_signals in datagen_info. + f = h5py.File(dataset_path, "a") + + # make pygame screen first + screen = make_pygame_screen( + traj_grp=f["data/{}".format(demo_keys[0])], + image_names=args.render_image_names, + image_scale=args.image_scale, + ) + + for ind in range(len(demo_keys)): + ep = demo_keys[ind] + print("Annotating episode: {}".format(ep)) + + annotate_subtasks_in_trajectory( + ep=ep, + traj_grp=f["data/{}".format(ep)], + subtask_signals=args.signals, + screen=screen, + video_skip=args.video_skip, + image_names=args.render_image_names, + playback_rate_grid=RATE_GRID, + ) + + f.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--dataset", + type=str, + required=True, + help="path to hdf5 dataset which will be modified in-place", + ) + parser.add_argument( + "--signals", + type=str, + nargs='+', + required=True, + help="specify sequence of subtask termination signals for all except last subtask -- these will be written using the annotations", + ) + parser.add_argument( + "--filter_key", + type=str, + default=None, + help="(optional) filter key, to select a subset of trajectories in the file", + ) + parser.add_argument( + "--n", + type=int, + default=None, + help="(optional) stop after n trajectories are annotated", + ) + parser.add_argument( + "--start", + type=int, + default=None, + help="(optional) start after this many trajectories in the dataset", + ) + parser.add_argument( + "--video_skip", + type=int, + default=1, + help="(optional) render frames on-screen every n steps", + ) + parser.add_argument( + "--render_image_names", + type=str, + nargs='+', + default=None, + help="(optional) camera name(s) / image observation(s) to use for rendering on-screen or to video. Default is" + "None, which corresponds to a predefined camera for each env type", + ) + parser.add_argument( + "--image_scale", + type=int, + default=5, + help="(optional) scaling size for images when rendering to screen", + ) + + args = parser.parse_args() + annotate_subtasks(args) diff --git a/mimicgen_envs/scripts/demo_random_action.py b/mimicgen/scripts/demo_random_action.py similarity index 92% rename from mimicgen_envs/scripts/demo_random_action.py rename to mimicgen/scripts/demo_random_action.py index 4c0c1d9..666f244 100644 --- a/mimicgen_envs/scripts/demo_random_action.py +++ b/mimicgen/scripts/demo_random_action.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -28,14 +28,14 @@ def choose_mimicgen_environment(): robosuite_envs = set(suite.ALL_ENVIRONMENTS) # all environments including mimicgen environments - import mimicgen_envs + import mimicgen all_envs = set(suite.ALL_ENVIRONMENTS) # get only mimicgen envs - only_mimicgen_envs = sorted(all_envs - robosuite_envs) + only_mimicgen = sorted(all_envs - robosuite_envs) # keep only envs that correspond to the different reset distributions from the paper - envs = [x for x in only_mimicgen_envs if x[-1].isnumeric()] + envs = [x for x in only_mimicgen if x[-1].isnumeric()] # Select environment to run print("Here is a list of environments in the suite:\n") diff --git a/mimicgen_envs/scripts/download_datasets.py b/mimicgen/scripts/download_datasets.py similarity index 92% rename from mimicgen_envs/scripts/download_datasets.py rename to mimicgen/scripts/download_datasets.py index 26ce525..ec78c75 100644 --- a/mimicgen_envs/scripts/download_datasets.py +++ b/mimicgen/scripts/download_datasets.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -8,9 +8,9 @@ import os import argparse -import mimicgen_envs -import mimicgen_envs.utils.file_utils as FileUtils -from mimicgen_envs import DATASET_REGISTRY +import mimicgen +import mimicgen.utils.file_utils as FileUtils +from mimicgen import DATASET_REGISTRY if __name__ == "__main__": @@ -55,7 +55,7 @@ # set default base directory for downloads default_base_dir = args.download_dir if default_base_dir is None: - default_base_dir = os.path.join(mimicgen_envs.__path__[0], "../datasets") + default_base_dir = os.path.join(mimicgen.__path__[0], "../datasets") # load args download_dataset_type = args.dataset_type diff --git a/mimicgen/scripts/generate_config_templates.py b/mimicgen/scripts/generate_config_templates.py new file mode 100644 index 0000000..3653904 --- /dev/null +++ b/mimicgen/scripts/generate_config_templates.py @@ -0,0 +1,37 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Helpful script to generate example config files, one per config class. These should be re-generated +when new config options are added, or when default settings in the config classes are modified. +""" +import os +import json + +import mimicgen +from mimicgen.configs.config import get_all_registered_configs + + +def main(): + # store template config jsons in this directory + target_dir = os.path.join(mimicgen.__path__[0], "exps/templates/") + + # iterate through registered config classes + all_configs = get_all_registered_configs() + for config_type in all_configs: + # store config json by config type + target_type_dir = os.path.join(target_dir, config_type) + os.makedirs(target_type_dir, exist_ok=True) + for name in all_configs[config_type]: + # make config class to dump it to json + c = all_configs[config_type][name]() + assert name == c.name + assert config_type == c.type + # dump to json + json_path = os.path.join(target_type_dir, "{}.json".format(name)) + c.dump(filename=json_path) + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/mimicgen/scripts/generate_core_configs.py b/mimicgen/scripts/generate_core_configs.py new file mode 100644 index 0000000..f61acba --- /dev/null +++ b/mimicgen/scripts/generate_core_configs.py @@ -0,0 +1,377 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +We utilize robomimic's config generator class to easily generate data generation configs for our +core set of tasks in the paper. It can be modified easily to generate other configs. + +The global variables at the top of the file should be configured manually. + +See https://robomimic.github.io/docs/tutorials/hyperparam_scan.html for more info. +""" +import os +import json +import shutil + +import robomimic +from robomimic.utils.hyperparam_utils import ConfigGenerator + +import mimicgen +import mimicgen.utils.config_utils as ConfigUtils +from mimicgen.utils.file_utils import config_generator_to_script_lines + + +# set path to folder containing src datasets +SRC_DATA_DIR = os.path.join(mimicgen.__path__[0], "../datasets/source") + +# set base folder for where to copy each base config and generate new config files for data generation +CONFIG_DIR = "/tmp/core_configs" + +# set base folder for newly generated datasets +OUTPUT_FOLDER = "/tmp/core_datasets" + +# number of trajectories to generate (or attempt to generate) +NUM_TRAJ = 1000 + +# whether to guarantee that many successful trajectories (e.g. keep running until that many successes, or stop at that many attempts) +GUARANTEE = True + +# whether to run a quick debug run instead of full generation +DEBUG = False + +# camera settings for collecting observations +CAMERA_NAMES = ["agentview", "robot0_eye_in_hand"] +CAMERA_SIZE = (84, 84) + +# path to base config(s) +BASE_BASE_CONFIG_PATH = os.path.join(mimicgen.__path__[0], "exps/templates/robosuite") +BASE_CONFIGS = [ + os.path.join(BASE_BASE_CONFIG_PATH, "stack.json"), + os.path.join(BASE_BASE_CONFIG_PATH, "stack_three.json"), + os.path.join(BASE_BASE_CONFIG_PATH, "square.json"), + os.path.join(BASE_BASE_CONFIG_PATH, "threading.json"), + os.path.join(BASE_BASE_CONFIG_PATH, "three_piece_assembly.json"), + os.path.join(BASE_BASE_CONFIG_PATH, "coffee.json"), + os.path.join(BASE_BASE_CONFIG_PATH, "coffee_preparation.json"), + os.path.join(BASE_BASE_CONFIG_PATH, "nut_assembly.json"), + os.path.join(BASE_BASE_CONFIG_PATH, "pick_place.json"), + os.path.join(BASE_BASE_CONFIG_PATH, "hammer_cleanup.json"), + os.path.join(BASE_BASE_CONFIG_PATH, "mug_cleanup.json"), + os.path.join(BASE_BASE_CONFIG_PATH, "kitchen.json"), +] + + +def make_generators(base_configs): + """ + An easy way to make multiple config generators by using different + settings for each. + """ + all_settings = [ + # stack + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "stack.hdf5"), + dataset_name="stack", + generation_path="{}/stack".format(OUTPUT_FOLDER), + # task_interface="MG_Stack", + tasks=["Stack_D0", "Stack_D1"], + task_names=["D0", "D1"], + select_src_per_subtask=True, + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + subtask_term_offset_range=[[10, 20], None], + ), + # stack_three + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "stack_three.hdf5"), + dataset_name="stack_three", + generation_path="{}/stack_three".format(OUTPUT_FOLDER), + # task_interface="MG_StackThree", + tasks=["StackThree_D0", "StackThree_D1"], + task_names=["D0", "D1"], + select_src_per_subtask=True, + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + subtask_term_offset_range=[[10, 20], [10, 20], [10, 20], None], + ), + # square + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "square.hdf5"), + dataset_name="square", + generation_path="{}/square".format(OUTPUT_FOLDER), + # task_interface="MG_Square", + tasks=["Square_D0", "Square_D1", "Square_D2"], + task_names=["D0", "D1", "D2"], + select_src_per_subtask=False, + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + subtask_term_offset_range=[[10, 20], None], + ), + # threading + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "threading.hdf5"), + dataset_name="threading", + generation_path="{}/threading".format(OUTPUT_FOLDER), + # task_interface="MG_Threading", + tasks=["Threading_D0", "Threading_D1", "Threading_D2"], + task_names=["D0", "D1", "D2"], + select_src_per_subtask=False, + selection_strategy="random", + selection_strategy_kwargs=None, + subtask_term_offset_range=[[5, 10], None], + ), + # three_piece_assembly + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "three_piece_assembly.hdf5"), + dataset_name="three_piece_assembly", + generation_path="{}/three_piece_assembly".format(OUTPUT_FOLDER), + # task_interface="MG_ThreePieceAssembly", + tasks=["ThreePieceAssembly_D0", "ThreePieceAssembly_D1", "ThreePieceAssembly_D2"], + task_names=["D0", "D1", "D2"], + select_src_per_subtask=False, + selection_strategy="random", + selection_strategy_kwargs=None, + subtask_term_offset_range=[[5, 10], [5, 10], [5, 10], None], + ), + # coffee + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "coffee.hdf5"), + dataset_name="coffee", + generation_path="{}/coffee".format(OUTPUT_FOLDER), + # task_interface="MG_Coffee", + tasks=["Coffee_D0", "Coffee_D1", "Coffee_D2"], + task_names=["D0", "D1", "D2"], + select_src_per_subtask=False, + selection_strategy="random", + selection_strategy_kwargs=None, + subtask_term_offset_range=[[5, 10], None], + ), + # coffee_preparation + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "coffee_preparation.hdf5"), + dataset_name="coffee_preparation", + generation_path="{}/coffee_preparation".format(OUTPUT_FOLDER), + # task_interface="MG_CoffeePreparation", + tasks=["CoffeePreparation_D0", "CoffeePreparation_D1"], + task_names=["D0", "D1"], + select_src_per_subtask=False, + selection_strategy="random", + selection_strategy_kwargs=None, + subtask_term_offset_range=[[5, 10], [5, 10], [5, 10], [5, 10], None], + ), + # nut_assembly + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "nut_assembly.hdf5"), + dataset_name="nut_assembly", + generation_path="{}/nut_assembly".format(OUTPUT_FOLDER), + # task_interface="MG_NutAssembly", + tasks=["NutAssembly_D0"], + task_names=["D0"], + select_src_per_subtask=False, + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + subtask_term_offset_range=[[10, 20], [10, 20], [10, 20], None], + ), + # pick_place + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "pick_place.hdf5"), + dataset_name="pick_place", + generation_path="{}/pick_place".format(OUTPUT_FOLDER), + # task_interface="MG_PickPlace", + tasks=["PickPlace_D0"], + task_names=["D0"], + select_src_per_subtask=True, + # NOTE: selection strategy is set by default in the config template, and we will not change it + # selection_strategy="nearest_neighbor_object", + # selection_strategy_kwargs=dict(nn_k=3), + subtask_term_offset_range=[[10, 20], None, [10, 20], None, [10, 20], None, [10, 20], None], + ), + # hammer_cleanup + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "hammer_cleanup.hdf5"), + dataset_name="hammer_cleanup", + generation_path="{}/hammer_cleanup".format(OUTPUT_FOLDER), + # task_interface="MG_HammerCleanup", + tasks=["HammerCleanup_D0", "HammerCleanup_D1"], + task_names=["D0", "D1"], + select_src_per_subtask=False, + selection_strategy="random", + selection_strategy_kwargs=None, + subtask_term_offset_range=[[10, 20], [10, 20], None], + ), + # mug_cleanup + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "mug_cleanup.hdf5"), + dataset_name="mug_cleanup", + generation_path="{}/mug_cleanup".format(OUTPUT_FOLDER), + # task_interface="MG_MugCleanup", + tasks=["MugCleanup_D0", "MugCleanup_D1", "MugCleanup_O1", "MugCleanup_O2"], + task_names=["D0", "D1", "O1", "O2"], + select_src_per_subtask=False, + selection_strategy="random", + selection_strategy_kwargs=None, + subtask_term_offset_range=[[10, 20], [10, 20], None], + ), + # kitchen + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "kitchen.hdf5"), + dataset_name="kitchen", + generation_path="{}/kitchen".format(OUTPUT_FOLDER), + # task_interface="MG_Kitchen", + tasks=["Kitchen_D0", "Kitchen_D1"], + task_names=["D0", "D1"], + select_src_per_subtask=False, + selection_strategy="random", + selection_strategy_kwargs=None, + subtask_term_offset_range=[[10, 20], [10, 20], [10, 20], [10, 20], [10, 20], [10, 20], None], + ), + ] + + assert len(base_configs) == len(all_settings) + ret = [] + for conf, setting in zip(base_configs, all_settings): + ret.append(make_generator(os.path.expanduser(conf), setting)) + return ret + + +def make_generator(config_file, settings): + """ + Implement this function to setup your own hyperparameter scan. + Each config generator is created using a base config file (@config_file) + and a @settings dictionary that can be used to modify which parameters + are set. + """ + generator = ConfigGenerator( + base_config_file=config_file, + script_file="", # will be overriden in next step + ) + + # set basic settings + ConfigUtils.set_basic_settings( + generator=generator, + group=0, + source_dataset_path=settings["dataset_path"], + source_dataset_name=settings["dataset_name"], + generation_path=settings["generation_path"], + guarantee=GUARANTEE, + num_traj=NUM_TRAJ, + num_src_demos=10, + max_num_failures=25, + num_demo_to_render=10, + num_fail_demo_to_render=25, + verbose=False, + ) + + # set settings for subtasks + ConfigUtils.set_subtask_settings( + generator=generator, + group=0, + base_config_file=config_file, + select_src_per_subtask=settings["select_src_per_subtask"], + subtask_term_offset_range=settings["subtask_term_offset_range"], + selection_strategy=settings.get("selection_strategy", None), + selection_strategy_kwargs=settings.get("selection_strategy_kwargs", None), + # default settings: action noise 0.05, with 5 interpolation steps + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + verbose=False, + ) + + # optionally set env interface to use, and type + # generator.add_param( + # key="experiment.task.interface", + # name="", + # group=0, + # values=[settings["task_interface"]], + # ) + # generator.add_param( + # key="experiment.task.interface_type", + # name="", + # group=0, + # values=["robosuite"], + # ) + + # set task to generate data on + generator.add_param( + key="experiment.task.name", + name="task", + group=1, + values=settings["tasks"], + value_names=settings["task_names"], + ) + + # optionally set robot and gripper that will be used for data generation (robosuite-only) + if settings.get("robots", None) is not None: + generator.add_param( + key="experiment.task.robot", + name="r", + group=2, + values=settings["robots"], + ) + if settings.get("grippers", None) is not None: + generator.add_param( + key="experiment.task.gripper", + name="g", + group=2, + values=settings["grippers"], + ) + + # set observation collection settings + ConfigUtils.set_obs_settings( + generator=generator, + group=-1, + collect_obs=True, + camera_names=CAMERA_NAMES, + camera_height=CAMERA_SIZE[0], + camera_width=CAMERA_SIZE[1], + ) + + if DEBUG: + # set debug settings + ConfigUtils.set_debug_settings( + generator=generator, + group=-1, + ) + + # seed + generator.add_param( + key="experiment.seed", + name="", + group=1000000, + values=[1], + ) + + return generator + + +def main(): + + # make config generators + generators = make_generators(base_configs=BASE_CONFIGS) + + # maybe remove existing config directory + config_dir = CONFIG_DIR + if os.path.exists(config_dir): + ans = input("Non-empty dir at {} will be removed.\nContinue (y / n)? \n".format(config_dir)) + if ans != "y": + exit() + shutil.rmtree(config_dir) + + all_json_files, run_lines = config_generator_to_script_lines(generators, config_dir=config_dir) + + real_run_lines = [] + for line in run_lines: + line = line.strip().replace("train.py", "generate_dataset.py") + line += " --auto-remove-exp" + real_run_lines.append(line) + run_lines = real_run_lines + + print("configs") + print(json.dumps(all_json_files, indent=4)) + print("runs") + print(json.dumps(run_lines, indent=4)) + + +if __name__ == "__main__": + main() diff --git a/mimicgen/scripts/generate_core_training_configs.py b/mimicgen/scripts/generate_core_training_configs.py new file mode 100644 index 0000000..436b1d6 --- /dev/null +++ b/mimicgen/scripts/generate_core_training_configs.py @@ -0,0 +1,337 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +We utilize robomimic's config generator class to easily generate policy training configs for the +core set of experiments in the paper. It can be modified easily to generate other +training configs. + +See https://robomimic.github.io/docs/tutorials/hyperparam_scan.html for more info. +""" +import os +import json +import shutil +import argparse + +import robomimic +from robomimic.utils.hyperparam_utils import ConfigGenerator + +import mimicgen +import mimicgen.utils.config_utils as ConfigUtils +from mimicgen.utils.file_utils import config_generator_to_script_lines + + +# set path to folder with mimicgen generated datasets +DATASET_DIR = "/tmp/core_datasets" + +# set base folder for where to generate new config files for training runs +CONFIG_DIR = "/tmp/core_train_configs" + +# set base folder for training outputs (model checkpoints, videos, logs) +OUTPUT_DIR = "/tmp/core_training_results" + +# path to base config +BASE_CONFIG = os.path.join(robomimic.__path__[0], "exps/templates/bc.json") + + +def make_generators(base_config, dataset_dir, output_dir): + """ + An easy way to make multiple config generators by using different + settings for each. + """ + all_settings = [ + # stack + dict( + dataset_paths=[ + os.path.join(dataset_dir, "stack", "demo_src_stack_task_D0/demo.hdf5"), + os.path.join(dataset_dir, "stack", "demo_src_stack_task_D1/demo.hdf5"), + ], + dataset_names=[ + "stack_D0", + "stack_D1", + ], + horizon=400, + ), + # stack_three + dict( + dataset_paths=[ + os.path.join(dataset_dir, "stack_three", "demo_src_stack_three_task_D0/demo.hdf5"), + os.path.join(dataset_dir, "stack_three", "demo_src_stack_three_task_D1/demo.hdf5"), + ], + dataset_names=[ + "stack_three_D0", + "stack_three_D1", + ], + horizon=400, + ), + # square + dict( + dataset_paths=[ + os.path.join(dataset_dir, "square", "demo_src_square_task_D0/demo.hdf5"), + os.path.join(dataset_dir, "square", "demo_src_square_task_D1/demo.hdf5"), + os.path.join(dataset_dir, "square", "demo_src_square_task_D2/demo.hdf5"), + ], + dataset_names=[ + "square_D0", + "square_D1", + "square_D2", + ], + horizon=400, + ), + # threading + dict( + dataset_paths=[ + os.path.join(dataset_dir, "threading", "demo_src_threading_task_D0/demo.hdf5"), + os.path.join(dataset_dir, "threading", "demo_src_threading_task_D1/demo.hdf5"), + os.path.join(dataset_dir, "threading", "demo_src_threading_task_D2/demo.hdf5"), + ], + dataset_names=[ + "threading_D0", + "threading_D1", + "threading_D2", + ], + horizon=400, + ), + # three_piece_assembly + dict( + dataset_paths=[ + os.path.join(dataset_dir, "three_piece_assembly", "demo_src_three_piece_assembly_task_D0/demo.hdf5"), + os.path.join(dataset_dir, "three_piece_assembly", "demo_src_three_piece_assembly_task_D1/demo.hdf5"), + os.path.join(dataset_dir, "three_piece_assembly", "demo_src_three_piece_assembly_task_D2/demo.hdf5"), + ], + dataset_names=[ + "three_piece_assembly_D0", + "three_piece_assembly_D1", + "three_piece_assembly_D2", + ], + horizon=500, + ), + # coffee + dict( + dataset_paths=[ + os.path.join(dataset_dir, "coffee", "demo_src_coffee_task_D0/demo.hdf5"), + os.path.join(dataset_dir, "coffee", "demo_src_coffee_task_D1/demo.hdf5"), + os.path.join(dataset_dir, "coffee", "demo_src_coffee_task_D2/demo.hdf5"), + ], + dataset_names=[ + "coffee_D0", + "coffee_D1", + "coffee_D2", + ], + horizon=400, + ), + # coffee_preparation + dict( + dataset_paths=[ + os.path.join(dataset_dir, "coffee_preparation", "demo_src_coffee_preparation_task_D0/demo.hdf5"), + os.path.join(dataset_dir, "coffee_preparation", "demo_src_coffee_preparation_task_D1/demo.hdf5"), + ], + dataset_names=[ + "coffee_preparation_D0", + "coffee_preparation_D1", + ], + horizon=800, + ), + # nut_assembly + dict( + dataset_paths=[ + os.path.join(dataset_dir, "nut_assembly", "demo_src_nut_assembly_task_D0/demo.hdf5"), + ], + dataset_names=[ + "nut_assembly_D0", + ], + horizon=500, + ), + # pick_place + dict( + dataset_paths=[ + os.path.join(dataset_dir, "pick_place", "demo_src_pick_place_task_D0/demo.hdf5"), + ], + dataset_names=[ + "pick_place_D0", + ], + horizon=1000, + ), + # mug_cleanup + dict( + dataset_paths=[ + os.path.join(dataset_dir, "mug_cleanup", "demo_src_mug_cleanup_task_D0/demo.hdf5"), + os.path.join(dataset_dir, "mug_cleanup", "demo_src_mug_cleanup_task_D1/demo.hdf5"), + os.path.join(dataset_dir, "mug_cleanup", "demo_src_mug_cleanup_task_O1/demo.hdf5"), + os.path.join(dataset_dir, "mug_cleanup", "demo_src_mug_cleanup_task_O2/demo.hdf5"), + ], + dataset_names=[ + "mug_cleanup_D0", + "mug_cleanup_D1", + "mug_cleanup_O1", + "mug_cleanup_O2", + ], + horizon=500, + ), + # hammer_cleanup + dict( + dataset_paths=[ + os.path.join(dataset_dir, "hammer_cleanup", "demo_src_hammer_cleanup_task_D0/demo.hdf5"), + os.path.join(dataset_dir, "hammer_cleanup", "demo_src_hammer_cleanup_task_D1/demo.hdf5"), + ], + dataset_names=[ + "hammer_cleanup_D0", + "hammer_cleanup_D1", + ], + horizon=500, + ), + # kitchen + dict( + dataset_paths=[ + os.path.join(dataset_dir, "kitchen", "demo_src_kitchen_task_D0/demo.hdf5"), + os.path.join(dataset_dir, "kitchen", "demo_src_kitchen_task_D1/demo.hdf5"), + ], + dataset_names=[ + "kitchen_D0", + "kitchen_D1", + ], + horizon=800, + ), + ] + + ret = [] + for setting in all_settings: + for mod in ["low_dim", "image"]: + ret.append(make_gen(os.path.expanduser(base_config), setting, output_dir, mod)) + return ret + + +def make_gen(base_config, settings, output_dir, mod): + """ + Specify training configs to generate here. + """ + generator = ConfigGenerator( + base_config_file=base_config, + script_file="", # will be overriden in next step + base_exp_name="bc_rnn_{}".format(mod), + ) + + # set algo settings for bc-rnn + modality = mod + low_dim_keys = settings.get("low_dim_keys", None) + image_keys = settings.get("image_keys", None) + crop_size = settings.get("crop_size", None) + if modality == "low_dim": + if low_dim_keys is None: + low_dim_keys = ["robot0_eef_pos", "robot0_eef_quat", "robot0_gripper_qpos", "object"] + if modality == "image": + if low_dim_keys is None: + low_dim_keys = ["robot0_eef_pos", "robot0_eef_quat", "robot0_gripper_qpos"] + if image_keys is None: + image_keys = ["agentview_image", "robot0_eye_in_hand_image"] + if crop_size is None: + crop_size = [76, 76] + assert len(crop_size) == 2 + + ConfigUtils.set_learning_settings_for_bc_rnn( + generator=generator, + group=-1, + modality=modality, + seq_length=10, + low_dim_keys=low_dim_keys, + image_keys=image_keys, + crop_size=crop_size, + ) + + # set dataset + generator.add_param( + key="train.data", + name="ds", + group=0, + values=settings["dataset_paths"], + value_names=settings["dataset_names"], + ) + + # rollout settings + generator.add_param( + key="experiment.rollout.horizon", + name="", + group=1, + values=[settings["horizon"]], + ) + + # output path + generator.add_param( + key="train.output_dir", + name="", + group=-1, + values=[ + output_dir, + ], + ) + + # num data workers 4 by default (for both low-dim and image) and cache mode "low_dim" + generator.add_param( + key="train.num_data_workers", + name="", + group=-1, + values=[4], + ) + generator.add_param( + key="train.hdf5_cache_mode", + name="", + group=-1, + values=["low_dim"], + ) + + # seed + generator.add_param( + key="train.seed", + name="seed", + group=100000, + values=[101], + ) + + return generator + + +def main(args): + + # make config generators + generators = make_generators(base_config=BASE_CONFIG, dataset_dir=args.dataset_dir, output_dir=args.output_dir) + + if os.path.exists(args.config_dir): + ans = input("Non-empty dir at {} will be removed.\nContinue (y / n)? \n".format(args.config_dir)) + if ans != "y": + exit() + shutil.rmtree(args.config_dir) + + all_json_files, run_lines = config_generator_to_script_lines(generators, config_dir=args.config_dir) + + run_lines = [line.strip() for line in run_lines] + + print("configs") + print(json.dumps(all_json_files, indent=4)) + print("runs") + print(json.dumps(run_lines, indent=4)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--config_dir", + type=str, + default=os.path.expanduser(CONFIG_DIR), + help="set base folder for where to generate new config files for data generation", + ) + parser.add_argument( + "--dataset_dir", + type=str, + default=os.path.expanduser(DATASET_DIR), + help="set path to folder with datasets", + ) + parser.add_argument( + "--output_dir", + type=str, + default=os.path.expanduser(OUTPUT_DIR), + help="set base folder for where to generate new config files for data generation", + ) + + args = parser.parse_args() + main(args) diff --git a/mimicgen/scripts/generate_dataset.py b/mimicgen/scripts/generate_dataset.py new file mode 100644 index 0000000..65fa22b --- /dev/null +++ b/mimicgen/scripts/generate_dataset.py @@ -0,0 +1,655 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Main data generation script. + +Examples: + + # run normal data generation + python generate_dataset.py --config /path/to/config.json + + # render all data generation attempts on-screen + python generate_dataset.py --config /path/to/config.json --render + + # render all data generation attempts to a video + python generate_dataset.py --config /path/to/config.json --video_path /path/to/video.mp4 + + # run a quick debug run + python generate_dataset.py --config /path/to/config.json --debug + + # pause after every subtask to debug data generation + python generate_dataset.py --config /path/to/config.json --render --pause_subtask +""" + +import os +import shutil +import json +import time +import argparse +import traceback +import random +import imageio +import numpy as np +from copy import deepcopy + +import robomimic +from robomimic.utils.file_utils import get_env_metadata_from_dataset + +import mimicgen +import mimicgen.utils.file_utils as MG_FileUtils +import mimicgen.utils.robomimic_utils as RobomimicUtils + +from mimicgen.configs import config_factory, MG_TaskSpec +from mimicgen.datagen.data_generator import DataGenerator +from mimicgen.env_interfaces.base import make_interface + + +def get_important_stats( + new_dataset_folder_path, + num_success, + num_failures, + num_attempts, + num_problematic, + start_time=None, + ep_length_stats=None, +): + """ + Return a summary of important stats to write to json. + + Args: + new_dataset_folder_path (str): path to folder that will contain generated dataset + num_success (int): number of successful trajectories generated + num_failures (int): number of failed trajectories + num_attempts (int): number of total attempts + num_problematic (int): number of problematic trajectories that failed due + to a specific exception that was caught + start_time (float or None): starting time for this run from time.time() + ep_length_stats (dict or None): if provided, should have entries that summarize + the episode length statistics over the successfully generated trajectories + + Returns: + important_stats (dict): dictionary with useful summary of statistics + """ + important_stats = dict( + generation_path=new_dataset_folder_path, + success_rate=((100. * num_success) / num_attempts), + failure_rate=((100. * num_failures) / num_attempts), + num_success=num_success, + num_failures=num_failures, + num_attempts=num_attempts, + num_problematic=num_problematic, + ) + if (ep_length_stats is not None): + important_stats.update(ep_length_stats) + if start_time is not None: + # add in time taken + important_stats["time spent (hrs)"] = "{:.2f}".format((time.time() - start_time) / 3600.) + return important_stats + + +def generate_dataset( + mg_config, + auto_remove_exp=False, + render=False, + video_path=None, + video_skip=5, + render_image_names=None, + pause_subtask=False, +): + """ + Main function to collect a new dataset with MimicGen. + + Args: + mg_config (MG_Config instance): MimicGen config object + + auto_remove_exp (bool): if True, will remove generation folder if it exists, else + user will be prompted to decide whether to keep existing folder or not + + render (bool): if True, render each data generation attempt on-screen + + video_path (str or None): if provided, render the data generation attempts to the + provided video path + + video_skip (int): skip every nth frame when writing video + + render_image_names (list of str or None): if provided, specify camera names to + use during on-screen / off-screen rendering to override defaults + + pause_subtask (bool): if True, pause after every subtask during generation, for + debugging. + """ + + # time this run + start_time = time.time() + + # check some args + write_video = (video_path is not None) + assert not (render and write_video) # either on-screen or video but not both + if pause_subtask: + assert render, "should enable on-screen rendering for pausing to be useful" + + if write_video: + # debug video - use same cameras as observations + if len(mg_config.obs.camera_names) > 0: + assert render_image_names is None + render_image_names = list(mg_config.obs.camera_names) + + # path to source dataset + source_dataset_path = os.path.expandvars(os.path.expanduser(mg_config.experiment.source.dataset_path)) + + # get environment metadata from dataset + env_meta = get_env_metadata_from_dataset(dataset_path=source_dataset_path) + + # set seed for generation + random.seed(mg_config.experiment.seed) + np.random.seed(mg_config.experiment.seed) + + # create new folder for this data generation run + base_folder = os.path.expandvars(os.path.expanduser(mg_config.experiment.generation.path)) + new_dataset_folder_name = mg_config.experiment.name + new_dataset_folder_path = os.path.join( + base_folder, + new_dataset_folder_name, + ) + print("\nData will be generated at: {}".format(new_dataset_folder_path)) + + # ensure dataset folder does not exist, and make new folder + exist_ok = False + if os.path.exists(new_dataset_folder_path): + if not auto_remove_exp: + ans = input("\nWARNING: dataset folder ({}) already exists! \noverwrite? (y/n)\n".format(new_dataset_folder_path)) + else: + ans = "y" + if ans == "y": + print("Removed old results folder at {}".format(new_dataset_folder_path)) + shutil.rmtree(new_dataset_folder_path) + else: + print("Keeping old dataset folder. Note that individual files may still be overwritten.") + exist_ok = True + os.makedirs(new_dataset_folder_path, exist_ok=exist_ok) + + # log terminal output to text file + RobomimicUtils.make_print_logger(txt_file=os.path.join(new_dataset_folder_path, 'log.txt')) + + # save config to disk + MG_FileUtils.write_json( + json_dic=mg_config, + json_path=os.path.join(new_dataset_folder_path, "mg_config.json"), + ) + + print("\n============= Config =============") + print(mg_config) + print("") + + # some paths that we will create inside our new dataset folder + + # new dataset that will be generated + new_dataset_path = os.path.join(new_dataset_folder_path, "demo.hdf5") + + # tmp folder that will contain per-episode hdf5s that were successful (they will be merged later) + tmp_dataset_folder_path = os.path.join(new_dataset_folder_path, "tmp") + os.makedirs(tmp_dataset_folder_path, exist_ok=exist_ok) + + # folder containing logs + json_log_path = os.path.join(new_dataset_folder_path, "logs") + os.makedirs(json_log_path, exist_ok=exist_ok) + + if mg_config.experiment.generation.keep_failed: + # new dataset for failed trajectories, and tmp folder for per-episode hdf5s that failed + new_failed_dataset_path = os.path.join(new_dataset_folder_path, "demo_failed.hdf5") + tmp_dataset_failed_folder_path = os.path.join(new_dataset_folder_path, "tmp_failed") + os.makedirs(tmp_dataset_failed_folder_path, exist_ok=exist_ok) + + # get list of source demonstration keys from source hdf5 + all_demos = MG_FileUtils.get_all_demos_from_dataset( + dataset_path=source_dataset_path, + filter_key=mg_config.experiment.source.filter_key, + start=mg_config.experiment.source.start, + n=mg_config.experiment.source.n, + ) + + # prepare args for creating simulation environment + + # auto-fill camera rendering info if not specified + if (write_video or render) and (render_image_names is None): + render_image_names = RobomimicUtils.get_default_env_cameras(env_meta=env_meta) + if render: + # on-screen rendering can only support one camera + assert len(render_image_names) == 1 + + # env args: cameras to use come from debug camera video to write, or from observation collection + camera_names = (mg_config.obs.camera_names if not write_video else render_image_names) + + # env args: don't use image obs when writing debug video + use_image_obs = ((mg_config.obs.collect_obs and (len(mg_config.obs.camera_names) > 0)) if not write_video else False) + use_depth_obs = False + + # simulation environment + env = RobomimicUtils.create_env( + env_meta=env_meta, + env_class=None, + env_name=mg_config.experiment.task.name, + robot=mg_config.experiment.task.robot, + gripper=mg_config.experiment.task.gripper, + camera_names=camera_names, + camera_height=mg_config.obs.camera_height, + camera_width=mg_config.obs.camera_width, + render=render, + render_offscreen=write_video, + use_image_obs=use_image_obs, + use_depth_obs=use_depth_obs, + ) + print("\n==== Using environment with the following metadata ====") + print(json.dumps(env.serialize(), indent=4)) + print("") + + # get information necessary to create env interface + env_interface_name, env_interface_type = MG_FileUtils.get_env_interface_info_from_dataset( + dataset_path=source_dataset_path, + demo_keys=all_demos, + ) + # possibly override from config + if mg_config.experiment.task.interface is not None: + env_interface_name = mg_config.experiment.task.interface + if mg_config.experiment.task.interface_type is not None: + env_interface_type = mg_config.experiment.task.interface_type + + # create environment interface to use during data generation + env_interface = make_interface( + name=env_interface_name, + interface_type=env_interface_type, + # NOTE: env_interface takes underlying simulation environment, not robomimic wrapper + env=env.base_env, + ) + print("Created environment interface: {}".format(env_interface)) + + # make sure we except the same exceptions that we would normally except during policy rollouts + exceptions_to_except = env.rollout_exceptions + + # get task spec object from config + task_spec_json_string = mg_config.task.task_spec.dump() + task_spec = MG_TaskSpec.from_json(json_string=task_spec_json_string) + + # make data generator object + data_generator = DataGenerator( + task_spec=task_spec, + dataset_path=source_dataset_path, + demo_keys=all_demos, + ) + + print("\n==== Created Data Generator ====") + print(data_generator) + print("") + + # we might write a video to show the data generation attempts + video_writer = None + if write_video: + video_writer = imageio.get_writer(video_path, fps=20) + + # data generation statistics + num_success = 0 + num_failures = 0 + num_attempts = 0 + num_problematic = 0 + ep_lengths = [] # episode lengths for successfully generated data + selected_src_demo_inds_all = [] # selected source demo index in @all_demos for each trial + selected_src_demo_inds_succ = [] # selected source demo index in @all_demos for each successful trial + + # we will keep generating data until @num_trials successes (if @guarantee_success) else @num_trials attempts + num_trials = mg_config.experiment.generation.num_trials + guarantee_success = mg_config.experiment.generation.guarantee + + while True: + + # generate trajectory + try: + generated_traj = data_generator.generate( + env=env, + env_interface=env_interface, + select_src_per_subtask=mg_config.experiment.generation.select_src_per_subtask, + transform_first_robot_pose=mg_config.experiment.generation.transform_first_robot_pose, + interpolate_from_last_target_pose=mg_config.experiment.generation.interpolate_from_last_target_pose, + render=render, + video_writer=video_writer, + video_skip=video_skip, + camera_names=render_image_names, + pause_subtask=pause_subtask, + ) + except exceptions_to_except as e: + # problematic trajectory - do not have this count towards our total number of attempts, and re-try + print("") + print("*" * 50) + print("WARNING: got rollout exception {}".format(e)) + print("*" * 50) + print("") + num_problematic += 1 + continue + + # remember selection of source demos for each subtask + selected_src_demo_inds_all.append(generated_traj["src_demo_inds"]) + + # check if generated trajectory was successful + success = bool(generated_traj["success"]) + + if success: + num_success += 1 + + # store successful demonstration + ep_lengths.append(generated_traj["actions"].shape[0]) + MG_FileUtils.write_demo_to_hdf5( + folder=tmp_dataset_folder_path, + env=env, + initial_state=generated_traj["initial_state"], + states=generated_traj["states"], + observations=(generated_traj["observations"] if mg_config.obs.collect_obs else None), + datagen_info=generated_traj["datagen_infos"], + actions=generated_traj["actions"], + src_demo_inds=generated_traj["src_demo_inds"], + src_demo_labels=generated_traj["src_demo_labels"], + ) + selected_src_demo_inds_succ.append(generated_traj["src_demo_inds"]) + else: + num_failures += 1 + + # check if this failure should be kept + if mg_config.experiment.generation.keep_failed and \ + (mg_config.experiment.max_num_failures is None) or (num_failures <= mg_config.experiment.max_num_failures): + + # save failed trajectory in separate folder + MG_FileUtils.write_demo_to_hdf5( + folder=tmp_dataset_failed_folder_path, + env=env, + initial_state=generated_traj["initial_state"], + states=generated_traj["states"], + observations=(generated_traj["observations"] if mg_config.obs.collect_obs else None), + datagen_info=generated_traj["datagen_infos"], + actions=generated_traj["actions"], + src_demo_inds=generated_traj["src_demo_inds"], + src_demo_labels=generated_traj["src_demo_labels"], + ) + + num_attempts += 1 + print("") + print("*" * 50) + print("trial {} success: {}".format(num_attempts, success)) + print("have {} successes out of {} trials so far".format(num_success, num_attempts)) + print("have {} failures out of {} trials so far".format(num_failures, num_attempts)) + print("*" * 50) + + # regularly log progress to disk every so often + if (num_attempts % mg_config.experiment.log_every_n_attempts) == 0: + + # get summary stats + summary_stats = get_important_stats( + new_dataset_folder_path=new_dataset_folder_path, + num_success=num_success, + num_failures=num_failures, + num_attempts=num_attempts, + num_problematic=num_problematic, + start_time=start_time, + ep_length_stats=None, + ) + + # write stats to disk + max_digits = len(str(num_trials * 1000)) + 1 # assume we will never have lower than 0.1% data generation SR + json_file_path = os.path.join(json_log_path, "attempt_{}_succ_{}_rate_{}.json".format( + str(num_attempts).zfill(max_digits), # pad with leading zeros for ordered list of jsons in directory + num_success, + np.round((100. * num_success) / num_attempts, 2), + )) + MG_FileUtils.write_json(json_dic=summary_stats, json_path=json_file_path) + + # termination condition is on enough successes if @guarantee_success or enough attempts otherwise + check_val = num_success if guarantee_success else num_attempts + if check_val >= num_trials: + break + + if write_video: + video_writer.close() + + # merge all new created files + print("\nFinished data generation. Merging per-episode hdf5s together...\n") + MG_FileUtils.merge_all_hdf5( + folder=tmp_dataset_folder_path, + new_hdf5_path=new_dataset_path, + delete_folder=True, + ) + if mg_config.experiment.generation.keep_failed: + MG_FileUtils.merge_all_hdf5( + folder=tmp_dataset_failed_folder_path, + new_hdf5_path=new_failed_dataset_path, + delete_folder=True, + ) + + # get episode length statistics + ep_length_stats = None + if len(ep_lengths) > 0: + ep_lengths = np.array(ep_lengths) + ep_length_mean = float(np.mean(ep_lengths)) + ep_length_std = float(np.std(ep_lengths)) + ep_length_max = int(np.max(ep_lengths)) + ep_length_3std = int(np.ceil(ep_length_mean + 3. * ep_length_std)) + ep_length_stats = dict( + ep_length_mean=ep_length_mean, + ep_length_std=ep_length_std, + ep_length_max=ep_length_max, + ep_length_3std=ep_length_3std, + ) + + stats = get_important_stats( + new_dataset_folder_path=new_dataset_folder_path, + num_success=num_success, + num_failures=num_failures, + num_attempts=num_attempts, + num_problematic=num_problematic, + start_time=start_time, + ep_length_stats=ep_length_stats, + ) + print("\nStats Summary") + print(json.dumps(stats, indent=4)) + + # maybe render videos + if mg_config.experiment.render_video: + if (num_success > 0): + playback_video_path = os.path.join(new_dataset_folder_path, "playback_{}.mp4".format(new_dataset_folder_name)) + num_render = mg_config.experiment.num_demo_to_render + print("Rendering successful trajectories...") + RobomimicUtils.make_dataset_video( + dataset_path=new_dataset_path, + video_path=playback_video_path, + num_render=num_render, + ) + else: + print("\n" + "*" * 80) + print("\nWARNING: skipping dataset video creation since no successes") + print("\n" + "*" * 80 + "\n") + if mg_config.experiment.generation.keep_failed: + if (num_failures > 0): + playback_video_path = os.path.join(new_dataset_folder_path, "playback_{}_failed.mp4".format(new_dataset_folder_name)) + num_render = mg_config.experiment.num_fail_demo_to_render + print("Rendering failure trajectories...") + RobomimicUtils.make_dataset_video( + dataset_path=new_failed_dataset_path, + video_path=playback_video_path, + num_render=num_render, + ) + else: + print("\n" + "*" * 80) + print("\nWARNING: skipping dataset video creation since no failures") + print("\n" + "*" * 80 + "\n") + + # return some summary info + final_important_stats = get_important_stats( + new_dataset_folder_path=new_dataset_folder_path, + num_success=num_success, + num_failures=num_failures, + num_attempts=num_attempts, + num_problematic=num_problematic, + start_time=start_time, + ep_length_stats=ep_length_stats, + ) + + # write stats to disk + json_file_path = os.path.join(new_dataset_folder_path, "important_stats.json") + MG_FileUtils.write_json(json_dic=final_important_stats, json_path=json_file_path) + + # NOTE: we are not currently saving the choice of source human demonstrations for each trial, + # but you can do that if you wish -- the information is stored in @selected_src_demo_inds_all + # and @selected_src_demo_inds_succ + + return final_important_stats + + +def main(args): + + # load config object + with open(args.config, "r") as f: + ext_cfg = json.load(f) + # config generator from robomimic generates this part of config unused by MimicGen + if "meta" in ext_cfg: + del ext_cfg["meta"] + mg_config = config_factory(ext_cfg["name"], config_type=ext_cfg["type"]) + + # update config with external json - this will throw errors if + # the external config has keys not present in the base config + with mg_config.values_unlocked(): + mg_config.update(ext_cfg) + + # We assume that the external config specifies all subtasks, so + # delete any subtasks not in the external config. + source_subtasks = set(mg_config.task.task_spec.keys()) + new_subtasks = set(ext_cfg["task"]["task_spec"].keys()) + for subtask in (source_subtasks - new_subtasks): + print("deleting subtask {} in original config".format(subtask)) + del mg_config.task.task_spec[subtask] + + # maybe override some settings + if args.task_name is not None: + mg_config.experiment.task.name = args.task_name + + if args.source is not None: + mg_config.experiment.source.dataset_path = args.source + + if args.folder is not None: + mg_config.experiment.generation.path = args.folder + + if args.num_demos is not None: + mg_config.experiment.generation.num_trials = args.num_demos + + if args.seed is not None: + mg_config.experiment.seed = args.seed + + # maybe modify config for debugging purposes + if args.debug: + # shrink length of generation to test whether this run is likely to crash + mg_config.experiment.source.n = 3 + mg_config.experiment.generation.guarantee = False + mg_config.experiment.generation.num_trials = 2 + + # send output to a temporary directory + mg_config.experiment.generation.path = "/tmp/tmp_mimicgen" + + # catch error during generation and print it + res_str = "finished run successfully!" + important_stats = None + try: + important_stats = generate_dataset( + mg_config=mg_config, + auto_remove_exp=args.auto_remove_exp, + render=args.render, + video_path=args.video_path, + video_skip=args.video_skip, + render_image_names=args.render_image_names, + pause_subtask=args.pause_subtask, + ) + except Exception as e: + res_str = "run failed with error:\n{}\n\n{}".format(e, traceback.format_exc()) + print(res_str) + if important_stats is not None: + important_stats = json.dumps(important_stats, indent=4) + print("\nFinal Data Generation Stats") + print(important_stats) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--config", + type=str, + required=True, + help="path to MimicGen config json", + ) + parser.add_argument( + "--debug", + action='store_true', + help="set this flag to run a quick generation run for debugging purposes", + ) + parser.add_argument( + "--auto-remove-exp", + action='store_true', + help="force delete the experiment folder if it exists" + ) + parser.add_argument( + "--render", + action='store_true', + help="render each data generation attempt on-screen", + ) + parser.add_argument( + "--video_path", + type=str, + default=None, + help="if provided, render the data generation attempts to the provided video path", + ) + parser.add_argument( + "--video_skip", + type=int, + default=5, + help="skip every nth frame when writing video", + ) + parser.add_argument( + "--render_image_names", + type=str, + nargs='+', + default=None, + help="(optional) camera name(s) / image observation(s) to use for rendering on-screen or to video. Default is" + "None, which corresponds to a predefined camera for each env type", + ) + parser.add_argument( + "--pause_subtask", + action='store_true', + help="pause after every subtask during generation for debugging - only useful with render flag", + ) + parser.add_argument( + "--source", + type=str, + help="path to source dataset, to override the one in the config", + ) + parser.add_argument( + "--task_name", + type=str, + help="environment name to use for data generation, to override the one in the config", + default=None, + ) + parser.add_argument( + "--folder", + type=str, + help="folder that will be created with new data, to override the one in the config", + default=None, + ) + parser.add_argument( + "--num_demos", + type=int, + help="number of demos to generate, or attempt to generate, to override the one in the config", + default=None, + ) + parser.add_argument( + "--seed", + type=int, + help="seed, to override the one in the config", + default=None, + ) + + args = parser.parse_args() + main(args) diff --git a/mimicgen/scripts/generate_robot_transfer_configs.py b/mimicgen/scripts/generate_robot_transfer_configs.py new file mode 100644 index 0000000..b124cf7 --- /dev/null +++ b/mimicgen/scripts/generate_robot_transfer_configs.py @@ -0,0 +1,224 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +We utilize robomimic's config generator class to easily generate data generation configs for the +robot transfer set of experiments in the paper, where we use source data collected on the Panda arm +to generate demonstrations for other robot arms. It can be modified easily to generate other configs. + +The global variables at the top of the file should be configured manually. + +See https://robomimic.github.io/docs/tutorials/hyperparam_scan.html for more info. +""" +import os +import json +import shutil + +import robomimic +from robomimic.utils.hyperparam_utils import ConfigGenerator + +import mimicgen +import mimicgen.utils.config_utils as ConfigUtils +from mimicgen.utils.file_utils import config_generator_to_script_lines + + +# set path to folder containing src datasets +SRC_DATA_DIR = os.path.join(mimicgen.__path__[0], "../datasets/source") + +# set base folder for where to copy each base config and generate new config files for data generation +CONFIG_DIR = "/tmp/robot_configs" + +# set base folder for newly generated datasets +OUTPUT_FOLDER = "/tmp/robot_datasets" + +# number of trajectories to generate (or attempt to generate) +NUM_TRAJ = 1000 + +# whether to guarantee that many successful trajectories (e.g. keep running until that many successes, or stop at that many attempts) +GUARANTEE = True + +# whether to run a quick debug run instead of full generation +DEBUG = False + +# camera settings for collecting observations +CAMERA_NAMES = ["agentview", "robot0_eye_in_hand"] +CAMERA_SIZE = (84, 84) + +# path to base config(s) +BASE_BASE_CONFIG_PATH = os.path.join(mimicgen.__path__[0], "exps/templates/robosuite") +BASE_CONFIGS = [ + os.path.join(BASE_BASE_CONFIG_PATH, "square.json"), + os.path.join(BASE_BASE_CONFIG_PATH, "threading.json"), +] + + +def make_generators(base_configs): + """ + An easy way to make multiple config generators by using different + settings for each. + """ + all_settings = [ + # square + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "square.hdf5"), + dataset_name="square", + generation_path="{}/square".format(OUTPUT_FOLDER), + tasks=["Square_D0", "Square_D1"], + task_names=["D0", "D1"], + robots=["Sawyer", "IIWA", "UR5e"], + grippers=["RethinkGripper", "Robotiq85Gripper", "Robotiq85Gripper"], + select_src_per_subtask=False, + selection_strategy="nearest_neighbor_object", + selection_strategy_kwargs=dict(nn_k=3), + subtask_term_offset_range=[[10, 20], None], + ), + # threading + dict( + dataset_path=os.path.join(SRC_DATA_DIR, "threading.hdf5"), + dataset_name="threading", + generation_path="{}/threading".format(OUTPUT_FOLDER), + tasks=["Threading_D0", "Threading_D1"], + task_names=["D0", "D1"], + robots=["Sawyer", "IIWA", "UR5e"], + grippers=["RethinkGripper", "Robotiq85Gripper", "Robotiq85Gripper"], + select_src_per_subtask=False, + selection_strategy="random", + selection_strategy_kwargs=None, + subtask_term_offset_range=[[5, 10], None], + ), + ] + + assert len(base_configs) == len(all_settings) + ret = [] + for conf, setting in zip(base_configs, all_settings): + ret.append(make_generator(os.path.expanduser(conf), setting)) + return ret + + +def make_generator(config_file, settings): + """ + Implement this function to setup your own hyperparameter scan. + Each config generator is created using a base config file (@config_file) + and a @settings dictionary that can be used to modify which parameters + are set. + """ + generator = ConfigGenerator( + base_config_file=config_file, + script_file="", # will be overriden in next step + ) + + # set basic settings + ConfigUtils.set_basic_settings( + generator=generator, + group=0, + source_dataset_path=settings["dataset_path"], + source_dataset_name=settings["dataset_name"], + generation_path=settings["generation_path"], + guarantee=GUARANTEE, + num_traj=NUM_TRAJ, + num_src_demos=10, + max_num_failures=25, + num_demo_to_render=10, + num_fail_demo_to_render=25, + verbose=False, + ) + + # set settings for subtasks + ConfigUtils.set_subtask_settings( + generator=generator, + group=0, + base_config_file=config_file, + select_src_per_subtask=settings["select_src_per_subtask"], + subtask_term_offset_range=settings["subtask_term_offset_range"], + selection_strategy=settings.get("selection_strategy", None), + selection_strategy_kwargs=settings.get("selection_strategy_kwargs", None), + # default settings: action noise 0.05, with 5 interpolation steps + action_noise=0.05, + num_interpolation_steps=5, + num_fixed_steps=0, + verbose=False, + ) + + # set task to generate data on + generator.add_param( + key="experiment.task.name", + name="task", + group=1, + values=settings["tasks"], + value_names=settings["task_names"], + ) + + if settings.get("robots", None) is not None: + generator.add_param( + key="experiment.task.robot", + name="robot", + group=2, + values=settings["robots"], + ) + if settings.get("grippers", None) is not None: + generator.add_param( + key="experiment.task.gripper", + name="gripper", + group=2, + values=settings["grippers"], + ) + + # set observation collection settings + ConfigUtils.set_obs_settings( + generator=generator, + group=-1, + collect_obs=True, + camera_names=CAMERA_NAMES, + camera_height=CAMERA_SIZE[0], + camera_width=CAMERA_SIZE[1], + ) + + if DEBUG: + # set debug settings + ConfigUtils.set_debug_settings( + generator=generator, + group=-1, + ) + + # seed + generator.add_param( + key="experiment.seed", + name="", + group=1000000, + values=[1], + ) + + return generator + + +def main(): + + # make config generators + generators = make_generators(base_configs=BASE_CONFIGS) + + # maybe remove existing config directory + config_dir = CONFIG_DIR + if os.path.exists(config_dir): + ans = input("Non-empty dir at {} will be removed.\nContinue (y / n)? \n".format(config_dir)) + if ans != "y": + exit() + shutil.rmtree(config_dir) + + all_json_files, run_lines = config_generator_to_script_lines(generators, config_dir=config_dir) + + real_run_lines = [] + for line in run_lines: + line = line.strip().replace("train.py", "generate_dataset.py") + line += " --auto-remove-exp" + real_run_lines.append(line) + run_lines = real_run_lines + + print("configs") + print(json.dumps(all_json_files, indent=4)) + print("runs") + print(json.dumps(run_lines, indent=4)) + + +if __name__ == "__main__": + main() diff --git a/mimicgen/scripts/generate_robot_transfer_training_configs.py b/mimicgen/scripts/generate_robot_transfer_training_configs.py new file mode 100644 index 0000000..2656d5a --- /dev/null +++ b/mimicgen/scripts/generate_robot_transfer_training_configs.py @@ -0,0 +1,226 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +We utilize robomimic's config generator class to easily generate policy training configs for the +robot transfer set of experiments in the paper, where we use source data collected on the Panda arm +to generate demonstrations for other robot arms. It can be modified easily to generate other +training configs. + +See https://robomimic.github.io/docs/tutorials/hyperparam_scan.html for more info. +""" +import os +import json +import shutil +import argparse + +import robomimic +from robomimic.utils.hyperparam_utils import ConfigGenerator + +import mimicgen +import mimicgen.utils.config_utils as ConfigUtils +from mimicgen.utils.file_utils import config_generator_to_script_lines + + +# set path to folder with mimicgen generated datasets +DATASET_DIR = "/tmp/robot_datasets" + +# set base folder for where to generate new config files for training runs +CONFIG_DIR = "/tmp/robot_train_configs" + +# set base folder for training outputs (model checkpoints, videos, logs) +OUTPUT_DIR = "/tmp/robot_training_results" + +# path to base config +BASE_CONFIG = os.path.join(robomimic.__path__[0], "exps/templates/bc.json") + + +def make_generators(base_config, dataset_dir, output_dir): + """ + An easy way to make multiple config generators by using different + settings for each. + """ + all_settings = [ + # square + dict( + dataset_paths=[ + os.path.join(dataset_dir, "square", "demo_src_square_task_D0_robot_IIWA_gripper_Robotiq85Gripper/demo.hdf5"), + os.path.join(dataset_dir, "square", "demo_src_square_task_D0_robot_Sawyer_gripper_RethinkGripper/demo.hdf5"), + os.path.join(dataset_dir, "square", "demo_src_square_task_D0_robot_UR5e_gripper_Robotiq85Gripper/demo.hdf5"), + os.path.join(dataset_dir, "square", "demo_src_square_task_D1_robot_IIWA_gripper_Robotiq85Gripper/demo.hdf5"), + os.path.join(dataset_dir, "square", "demo_src_square_task_D1_robot_Sawyer_gripper_RethinkGripper/demo.hdf5"), + os.path.join(dataset_dir, "square", "demo_src_square_task_D1_robot_UR5e_gripper_Robotiq85Gripper/demo.hdf5"), + ], + dataset_names=[ + "square_D0_IIWA", + "square_D0_Sawyer", + "square_D0_UR5e", + "square_D1_IIWA", + "square_D1_Sawyer", + "square_D1_UR5e", + ], + horizon=400, + ), + # threading + dict( + dataset_paths=[ + os.path.join(dataset_dir, "threading", "demo_src_threading_task_D0_robot_IIWA_gripper_Robotiq85Gripper/demo.hdf5"), + os.path.join(dataset_dir, "threading", "demo_src_threading_task_D0_robot_Sawyer_gripper_RethinkGripper/demo.hdf5"), + os.path.join(dataset_dir, "threading", "demo_src_threading_task_D0_robot_UR5e_gripper_Robotiq85Gripper/demo.hdf5"), + os.path.join(dataset_dir, "threading", "demo_src_threading_task_D1_robot_IIWA_gripper_Robotiq85Gripper/demo.hdf5"), + os.path.join(dataset_dir, "threading", "demo_src_threading_task_D1_robot_Sawyer_gripper_RethinkGripper/demo.hdf5"), + os.path.join(dataset_dir, "threading", "demo_src_threading_task_D1_robot_UR5e_gripper_Robotiq85Gripper/demo.hdf5"), + ], + dataset_names=[ + "threading_D0_IIWA", + "threading_D0_Sawyer", + "threading_D0_UR5e", + "threading_D1_IIWA", + "threading_D1_Sawyer", + "threading_D1_UR5e", + ], + horizon=400, + ), + ] + + ret = [] + for setting in all_settings: + for mod in ["low_dim", "image"]: + ret.append(make_gen(os.path.expanduser(base_config), setting, output_dir, mod)) + return ret + + +def make_gen(base_config, settings, output_dir, mod): + """ + Specify training configs to generate here. + """ + generator = ConfigGenerator( + base_config_file=base_config, + script_file="", # will be overriden in next step + base_exp_name="bc_rnn_{}".format(mod), + ) + + # set algo settings for bc-rnn + modality = mod + low_dim_keys = settings.get("low_dim_keys", None) + image_keys = settings.get("image_keys", None) + crop_size = settings.get("crop_size", None) + if modality == "low_dim": + if low_dim_keys is None: + low_dim_keys = ["robot0_eef_pos", "robot0_eef_quat", "robot0_gripper_qpos", "object"] + if modality == "image": + if low_dim_keys is None: + low_dim_keys = ["robot0_eef_pos", "robot0_eef_quat", "robot0_gripper_qpos"] + if image_keys is None: + image_keys = ["agentview_image", "robot0_eye_in_hand_image"] + if crop_size is None: + crop_size = [76, 76] + assert len(crop_size) == 2 + + ConfigUtils.set_learning_settings_for_bc_rnn( + generator=generator, + group=-1, + modality=modality, + seq_length=10, + low_dim_keys=low_dim_keys, + image_keys=image_keys, + crop_size=crop_size, + ) + + # set dataset + generator.add_param( + key="train.data", + name="ds", + group=0, + values=settings["dataset_paths"], + value_names=settings["dataset_names"], + ) + + # rollout settings + generator.add_param( + key="experiment.rollout.horizon", + name="", + group=1, + values=[settings["horizon"]], + ) + + # output path + generator.add_param( + key="train.output_dir", + name="", + group=-1, + values=[ + output_dir, + ], + ) + + # num data workers 4 by default (for both low-dim and image) and cache mode "low_dim" + generator.add_param( + key="train.num_data_workers", + name="", + group=-1, + values=[4], + ) + generator.add_param( + key="train.hdf5_cache_mode", + name="", + group=-1, + values=["low_dim"], + ) + + # seed + generator.add_param( + key="train.seed", + name="seed", + group=100000, + values=[101], + ) + + return generator + + +def main(args): + + # make config generators + generators = make_generators(base_config=BASE_CONFIG, dataset_dir=args.dataset_dir, output_dir=args.output_dir) + + if os.path.exists(args.config_dir): + ans = input("Non-empty dir at {} will be removed.\nContinue (y / n)? \n".format(args.config_dir)) + if ans != "y": + exit() + shutil.rmtree(args.config_dir) + + all_json_files, run_lines = config_generator_to_script_lines(generators, config_dir=args.config_dir) + + run_lines = [line.strip() for line in run_lines] + + print("configs") + print(json.dumps(all_json_files, indent=4)) + print("runs") + print(json.dumps(run_lines, indent=4)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--config_dir", + type=str, + default=os.path.expanduser(CONFIG_DIR), + help="set base folder for where to generate new config files for data generation", + ) + parser.add_argument( + "--dataset_dir", + type=str, + default=os.path.expanduser(DATASET_DIR), + help="set path to folder with datasets", + ) + parser.add_argument( + "--output_dir", + type=str, + default=os.path.expanduser(OUTPUT_DIR), + help="set base folder for where to generate new config files for data generation", + ) + + args = parser.parse_args() + main(args) diff --git a/mimicgen_envs/scripts/generate_training_configs.py b/mimicgen/scripts/generate_training_configs_for_public_datasets.py similarity index 96% rename from mimicgen_envs/scripts/generate_training_configs.py rename to mimicgen/scripts/generate_training_configs_for_public_datasets.py index 00d1bc8..7f4bd16 100644 --- a/mimicgen_envs/scripts/generate_training_configs.py +++ b/mimicgen/scripts/generate_training_configs_for_public_datasets.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -12,8 +12,8 @@ from robomimic.config import config_factory, Config from robomimic.scripts.generate_paper_configs import modify_config_for_default_low_dim_exp, modify_config_for_default_image_exp -import mimicgen_envs -from mimicgen_envs import DATASET_REGISTRY +import mimicgen +from mimicgen import DATASET_REGISTRY def set_obs_config(config, obs_modality): @@ -230,7 +230,7 @@ def generate_all_configs( parser.add_argument( "--config_dir", type=str, - default=os.path.join(mimicgen_envs.__path__[0], "exps/paper"), + default=os.path.join(mimicgen.__path__[0], "exps/paper"), help="Directory where generated configs will be placed. Defaults to 'paper' subfolder in exps folder of repository", ) @@ -238,7 +238,7 @@ def generate_all_configs( parser.add_argument( "--dataset_dir", type=str, - default=os.path.join(mimicgen_envs.__path__[0], "../datasets"), + default=os.path.join(mimicgen.__path__[0], "../datasets"), help="Base dataset directory for released datasets. Defaults to datasets folder in repository.", ) @@ -246,7 +246,7 @@ def generate_all_configs( parser.add_argument( "--output_dir", type=str, - default=os.path.join(mimicgen_envs.__path__[0], "../training_results"), + default=os.path.join(mimicgen.__path__[0], "../training_results"), help="Base output directory for all training runs that will be written to generated configs. Defaults to training_results folder in repository", ) diff --git a/mimicgen_envs/scripts/get_reset_videos.py b/mimicgen/scripts/get_reset_videos.py similarity index 98% rename from mimicgen_envs/scripts/get_reset_videos.py rename to mimicgen/scripts/get_reset_videos.py index c56c2b9..571cbfc 100644 --- a/mimicgen_envs/scripts/get_reset_videos.py +++ b/mimicgen/scripts/get_reset_videos.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the NVIDIA Source Code License [see LICENSE for details]. @@ -15,7 +15,7 @@ import robosuite from robosuite.controllers import load_controller_config -import mimicgen_envs +import mimicgen # base output folder diff --git a/mimicgen/scripts/get_source_info.py b/mimicgen/scripts/get_source_info.py new file mode 100644 index 0000000..34e6df0 --- /dev/null +++ b/mimicgen/scripts/get_source_info.py @@ -0,0 +1,60 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Helper script to report source dataset information. It verifies that the dataset has a +"datagen_info" field for the first episode and prints its structure. +""" +import h5py +import argparse + +import mimicgen +import mimicgen.utils.file_utils as MG_FileUtils + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--dataset", + type=str, + help="path to hdf5 dataset", + required=True, + ) + args = parser.parse_args() + + dataset_path = args.dataset + + # get first demonstration + first_demo_key = MG_FileUtils.get_all_demos_from_dataset( + dataset_path=dataset_path, + filter_key=None, + start=None, + n=1, + )[0] + f = h5py.File(dataset_path, "r") + ep_grp = f["data/{}".format(first_demo_key)] + + # verify datagen info exists + assert "datagen_info" in ep_grp, "Could not find MimicGen metadata in dataset {}. Ensure you have run prepare_src_dataset.py on this hdf5".format(dataset_path) + + # environment interface information + env_interface_name = ep_grp["datagen_info"].attrs["env_interface_name"] + env_interface_type = ep_grp["datagen_info"].attrs["env_interface_type"] + + print("\nEnvironment Interface: {}".format(env_interface_name)) + print("Environment Interface Type: {}".format(env_interface_type)) + + # structure of datagen info + ep_datagen_info = ep_grp["datagen_info"] + + print("\nStructure of datagen_info in episode {}:".format(first_demo_key)) + for k in ep_datagen_info: + if k in ["object_poses", "subtask_term_signals"]: + print(" {}:".format(k)) + for k2 in ep_datagen_info[k]: + print(" {}: shape {}".format(k2, ep_datagen_info[k][k2].shape)) + else: + print(" {}: shape {}".format(k, ep_datagen_info[k].shape)) + print("") + + f.close() diff --git a/mimicgen/scripts/merge_hdf5.py b/mimicgen/scripts/merge_hdf5.py new file mode 100644 index 0000000..414e334 --- /dev/null +++ b/mimicgen/scripts/merge_hdf5.py @@ -0,0 +1,109 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Script to merge all hdf5s if scripts/generate_dataset.py is incomplete, +and doesn't make it to the line that merges all the hdf5s. +""" + +import os +import shutil +import json +import h5py +import argparse +import imageio + +import numpy as np + +import mimicgen +import mimicgen.utils.file_utils as MG_FileUtils +from mimicgen.configs import config_factory + + +def merge_hdf5s(args): + """ + Main function to collect a new dataset using trajectory transforms from + an existing dataset. + """ + have_config = (args.config is not None) + have_folder = (args.folder is not None) + assert have_config or have_folder + assert not (have_config and have_folder) + + folder_path = args.folder + if have_config: + # get folder path from config + + # load config object + with open(args.config, "r") as f: + ext_cfg = json.load(f) + # config generator from robomimic generates this part of config unused by MimicGen + if "meta" in ext_cfg: + del ext_cfg["meta"] + mg_config = config_factory(ext_cfg["name"], config_type=ext_cfg["type"]) + # update config with external json - this will throw errors if + # the external config has keys not present in the base config + with mg_config.values_unlocked(): + mg_config.update(ext_cfg) + + base_folder = os.path.expandvars(os.path.expanduser(mg_config.experiment.generation.path)) # path where new folder will be generated + new_dataset_folder_name = mg_config.experiment.name # name of folder to generate + folder_path = os.path.join( + base_folder, + new_dataset_folder_name, + ) + + path_to_hdf5s = os.path.join(folder_path, "tmp") + path_to_new_hdf5 = os.path.join(folder_path, "demo.hdf5") + path_to_failed_hdf5s = os.path.join(folder_path, "tmp_failed") + path_to_new_failed_hdf5 = os.path.join(folder_path, "demo_failed.hdf5") + + assert os.path.exists(path_to_hdf5s) + merge_failures = os.path.exists(path_to_failed_hdf5s) + + # merge all new created files + num_success = MG_FileUtils.merge_all_hdf5( + folder=path_to_hdf5s, + new_hdf5_path=path_to_new_hdf5, + delete_folder=args.delete, + dry_run=args.count, + ) + print("NUM SUCCESS: {}".format(num_success)) + if merge_failures: + num_failures = MG_FileUtils.merge_all_hdf5( + folder=path_to_failed_hdf5s, + new_hdf5_path=path_to_new_failed_hdf5, + delete_folder=args.delete, + dry_run=args.count, + ) + print("NUM FAILURE: {}".format(num_failures)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--config", + type=str, + default=None, + help="path to json for dataset generation, used to find dataset folder", + ) + parser.add_argument( + "--folder", + type=str, + default=None, + help="path to dataset folder that is generated by scripts/generate_dataset.py", + ) + parser.add_argument( + "--count", + action='store_true', + help="if provided just count the number of demos instead of merging all of them", + ) + parser.add_argument( + "--delete", + action='store_true', + help="if provided, delete the tmp directories instead of saving them", + ) + + args = parser.parse_args() + merge_hdf5s(args) diff --git a/mimicgen/scripts/prepare_all_src_datasets.sh b/mimicgen/scripts/prepare_all_src_datasets.sh new file mode 100644 index 0000000..199f46f --- /dev/null +++ b/mimicgen/scripts/prepare_all_src_datasets.sh @@ -0,0 +1,76 @@ +#!/bin/bash + +# Store commands used to prepare source datasets for each environment. + + +# coffee +python prepare_src_dataset.py \ + --dataset ../../datasets/source/coffee.hdf5 \ + --env_interface MG_Coffee \ + --env_interface_type robosuite + +# coffee_preparation +python prepare_src_dataset.py \ + --dataset ../../datasets/source/coffee_preparation.hdf5 \ + --env_interface MG_CoffeePreparation \ + --env_interface_type robosuite + +# hammer_cleanup +python prepare_src_dataset.py \ + --dataset ../../datasets/source/hammer_cleanup.hdf5 \ + --env_interface MG_HammerCleanup \ + --env_interface_type robosuite + +# kitchen +python prepare_src_dataset.py \ + --dataset ../../datasets/source/kitchen.hdf5 \ + --env_interface MG_Kitchen \ + --env_interface_type robosuite + +# mug_cleanup +python prepare_src_dataset.py \ + --dataset ../../datasets/source/mug_cleanup.hdf5 \ + --env_interface MG_MugCleanup \ + --env_interface_type robosuite + +# nut_assembly +python prepare_src_dataset.py \ + --dataset ../../datasets/source/nut_assembly.hdf5 \ + --env_interface MG_NutAssembly \ + --env_interface_type robosuite + +# pick_place +python prepare_src_dataset.py \ + --dataset ../../datasets/source/pick_place.hdf5 \ + --env_interface MG_PickPlace \ + --env_interface_type robosuite + +# square +python prepare_src_dataset.py \ + --dataset ../../datasets/source/square.hdf5 \ + --env_interface MG_Square \ + --env_interface_type robosuite + +# stack +python prepare_src_dataset.py \ + --dataset ../../datasets/source/stack.hdf5 \ + --env_interface MG_Stack \ + --env_interface_type robosuite + +# stack_three +python prepare_src_dataset.py \ + --dataset ../../datasets/source/stack_three.hdf5 \ + --env_interface MG_StackThree \ + --env_interface_type robosuite + +# threading +python prepare_src_dataset.py \ + --dataset ../../datasets/source/threading.hdf5 \ + --env_interface MG_Threading \ + --env_interface_type robosuite + +# three_piece_assembly +python prepare_src_dataset.py \ + --dataset ../../datasets/source/three_piece_assembly.hdf5 \ + --env_interface MG_ThreePieceAssembly \ + --env_interface_type robosuite diff --git a/mimicgen/scripts/prepare_src_dataset.py b/mimicgen/scripts/prepare_src_dataset.py new file mode 100644 index 0000000..7e8b339 --- /dev/null +++ b/mimicgen/scripts/prepare_src_dataset.py @@ -0,0 +1,262 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Script to extract information needed for data generation from low-dimensional simulation states +in a source dataset and add it to the source dataset. Basically a stripped down version of +dataset_states_to_obs.py script in the robomimic codebase, with a handful of modifications. + +Example usage: + + # prepare a source dataset collected on robosuite Stack task + python prepare_src_dataset.py --dataset /path/to/stack.hdf5 --env_interface MG_Stack --env_interface_type robosuite + + # prepare a source dataset collected on robosuite Square task, but only use first 10 demos, and write output to new hdf5 + python prepare_src_dataset.py --dataset /path/to/square.hdf5 --env_interface MG_Square --env_interface_type robosuite --n 10 --output /tmp/square_new.hdf5 +""" +import os +import shutil +import json +import h5py +import argparse +import numpy as np +from copy import deepcopy +from tqdm import tqdm + +import robomimic +import robomimic.utils.tensor_utils as TensorUtils +import robomimic.utils.file_utils as FileUtils +import robomimic.utils.env_utils as EnvUtils +from robomimic.envs.env_base import EnvBase + +import mimicgen +import mimicgen.utils.file_utils as MG_FileUtils +from mimicgen.env_interfaces.base import make_interface + + +def extract_datagen_info_from_trajectory( + env, + env_interface, + initial_state, + states, + actions, +): + """ + Helper function to extract observations, rewards, and dones along a trajectory using + the simulator environment. + + Args: + env (instance of robomimic EnvBase): environment + + env_interface (MG_EnvInterface instance): environment interface for some data generation operations + + initial_state (dict): initial simulation state to load + + states (np.array): array of simulation states to load to extract information + + actions (np.array): array of actions + + Returns: + datagen_infos (dict): the datagen info objects across all timesteps represented as a dictionary of + numpy arrays, for easy writes to an hdf5 + """ + assert isinstance(env, EnvBase) + assert len(states) == actions.shape[0] + + # load the initial state + env.reset() + env.reset_to(initial_state) + + all_datagen_infos = [] + traj_len = len(states) + for t in range(traj_len): + # reset to state + env.reset_to({"states" : states[t]}) + + # extract datagen info as a dictionary + datagen_info = env_interface.get_datagen_info(action=actions[t]).to_dict() + all_datagen_infos.append(datagen_info) + + # convert list of dict to dict of list for datagen info dictionaries (for convenient writes to hdf5 dataset) + all_datagen_infos = TensorUtils.list_of_flat_dict_to_dict_of_list(all_datagen_infos) + + for k in all_datagen_infos: + if k in ["object_poses", "subtask_term_signals"]: + # convert list of dict to dict of list again + all_datagen_infos[k] = TensorUtils.list_of_flat_dict_to_dict_of_list(all_datagen_infos[k]) + # list to numpy array + for k2 in all_datagen_infos[k]: + all_datagen_infos[k][k2] = np.array(all_datagen_infos[k][k2]) + else: + # list to numpy array + all_datagen_infos[k] = np.array(all_datagen_infos[k]) + + return all_datagen_infos + + +def prepare_src_dataset( + dataset_path, + env_interface_name, + env_interface_type, + filter_key=None, + n=None, + output_path=None, +): + """ + Adds DatagenInfo object instance for each timestep in each source demonstration trajectory + and stores it under the "datagen_info" key for each episode. Also store the @env_interface_name + and @env_interface_type used in the attribute of each key. This information is used during + MimicGen data generation. + + Args: + dataset_path (str): path to input hdf5 dataset, which will be modified in-place unless + @output_path is provided + + env_interface_name (str): name of environment interface class to use for this source dataset + + env_interface_type (str): type of environment interface to use for this source dataset + + filter_key (str or None): name of filter key + + n (int or None): if provided, stop after n trajectories are processed + + output_path (str or None): if provided, write a new hdf5 here instead of modifying the + original dataset in-place + """ + + # maybe write to new file instead of modifying existing file in-place + if output_path is not None: + shutil.copy(dataset_path, output_path) + dataset_path = output_path + + # create environment that was to collect source demonstrations + env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=dataset_path) + env = EnvUtils.create_env_for_data_processing( + env_meta=env_meta, + camera_names=[], + camera_height=84, + camera_width=84, + reward_shaping=False, + ) + print("") + print("==== Using environment with the following metadata ====") + print(json.dumps(env.serialize(), indent=4)) + print("") + + # create environment interface for us to grab relevant information from simulation at each timestep + env_interface = make_interface( + name=env_interface_name, + interface_type=env_interface_type, + # NOTE: env_interface takes underlying simulation environment, not robomimic wrapper + env=env.base_env, + ) + print("Created environment interface: {}".format(env_interface)) + print("") + + # some operations are env-type-specific + is_robosuite_env = EnvUtils.is_robosuite_env(env_meta) + + # get list of source demonstration keys from source hdf5 + demos = MG_FileUtils.get_all_demos_from_dataset( + dataset_path=dataset_path, + filter_key=filter_key, + start=None, + n=n, + ) + + print("File that will be modified with datagen info: {}".format(dataset_path)) + print("") + + # open file to modify it + f = h5py.File(dataset_path, "a") + + total_samples = 0 + for ind in tqdm(range(len(demos))): + ep = demos[ind] + ep_grp = f["data/{}".format(ep)] + + # prepare states to reload from + states = ep_grp["states"][()] + initial_state = dict(states=states[0]) + if is_robosuite_env: + initial_state["model"] = ep_grp.attrs["model_file"] + + # extract datagen info + actions = ep_grp["actions"][()] + datagen_info = extract_datagen_info_from_trajectory( + env=env, + env_interface=env_interface, + initial_state=initial_state, + states=states, + actions=actions, + ) + + # delete old dategen info if it already exists + if "datagen_info" in ep_grp: + del ep_grp["datagen_info"] + + for k in datagen_info: + if k in ["object_poses", "subtask_term_signals"]: + # handle dict + for k2 in datagen_info[k]: + ep_grp.create_dataset("datagen_info/{}/{}".format(k, k2), data=np.array(datagen_info[k][k2])) + else: + ep_grp.create_dataset("datagen_info/{}".format(k), data=np.array(datagen_info[k])) + + # remember the env interface used too + ep_grp["datagen_info"].attrs["env_interface_name"] = env_interface_name + ep_grp["datagen_info"].attrs["env_interface_type"] = env_interface_type + + print("Modified {} trajectories to include datagen info.".format(len(demos))) + f.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--dataset", + type=str, + required=True, + help="path to input hdf5 dataset, which will be modified in-place", + ) + parser.add_argument( + "--env_interface", + type=str, + required=True, + help="name of environment interface class to use for this source dataset", + ) + parser.add_argument( + "--env_interface_type", + type=str, + required=True, + help="type of environment interface to use for this source dataset", + ) + parser.add_argument( + "--n", + type=int, + default=None, + help="(optional) stop after n trajectories are processed", + ) + parser.add_argument( + "--filter_key", + type=str, + default=None, + help="(optional) name of filter key, to select a subset of demo keys in the source hdf5", + ) + parser.add_argument( + "--output", + type=str, + default=None, + help="(optional) path to output hdf5 dataset, instead of modifying existing dataset in-place", + ) + + args = parser.parse_args() + prepare_src_dataset( + dataset_path=args.dataset, + env_interface_name=args.env_interface, + env_interface_type=args.env_interface_type, + filter_key=args.filter_key, + n=args.n, + output_path=args.output, + ) diff --git a/mimicgen/scripts/visualize_subtasks.py b/mimicgen/scripts/visualize_subtasks.py new file mode 100644 index 0000000..5fbc0e5 --- /dev/null +++ b/mimicgen/scripts/visualize_subtasks.py @@ -0,0 +1,357 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +A script to visualize each subtask in a source demonstration. This is a useful way to +debug the subtask termination signals in a set of source demonstrations, as well as +the choice of maximum subtask termination offsets. + +Examples: + + # render on-screen + python visualize_subtasks.py --dataset /path/to/demo.hdf5 --config /path/to/config.json --render + + # render to video + python visualize_subtasks.py --dataset /path/to/demo.hdf5 --config /path/to/config.json --video_path /path/to/video.mp4 + + # specify subtask information manually instead of using a config + python visualize_subtasks.py --dataset /path/to/demo.hdf5 --signals grasp_1 insert_1 grasp_2 --offsets 10 10 10 --render + +""" + +import os +import sys +import json +import h5py +import argparse +import imageio +import numpy as np + +import robomimic +import robomimic.utils.obs_utils as ObsUtils +import robomimic.utils.env_utils as EnvUtils +from robomimic.envs.env_base import EnvBase +from robomimic.utils.file_utils import get_env_metadata_from_dataset + +import mimicgen +import mimicgen.utils.file_utils as MG_FileUtils +import mimicgen.utils.robomimic_utils as RobomimicUtils +from mimicgen.utils.misc_utils import add_red_border_to_frame +from mimicgen.configs import MG_TaskSpec + + +def visualize_subtasks_with_env( + env, + initial_state, + states, + subtask_end_indices, + render=False, + video_writer=None, + video_skip=5, + camera_names=None, +): + """ + Helper function to visualize each subtask in a trajectory using the simulator environment. + If using on-screen rendering, the script will pause for input at the end of each subtask. If + writing to a video, each subtask will toggle between having a red border around each + frame and no border in the video. + + Args: + env (instance of EnvBase): environment + initial_state (dict): initial simulation state to load + states (list): list of simulation states to load + subtask_end_indices (list): list containing the end index for each subtask + render (bool): if True, render on-screen + video_writer (imageio writer): video writer + video_skip (int): determines rate at which environment frames are written to video + camera_names (list): determines which camera(s) are used for rendering. Pass more than + one to output a video with multiple camera views concatenated horizontally. + """ + assert isinstance(env, EnvBase) + + write_video = (video_writer is not None) + video_count = 0 + assert not (render and write_video) + assert render or write_video + + # load the initial state + env.reset() + env.reset_to(initial_state) + traj_len = len(states) + + cur_subtask_ind = 0 + should_add_border_to_frame = False + for i in range(traj_len): + # reset to state + env.reset_to({"states" : states[i]}) + + # whether we are on last index of current subtask + is_last_subtask_ind = (i == subtask_end_indices[cur_subtask_ind] - 1) + + # on-screen render + if render: + env.render(mode="human", camera_name=camera_names[0]) + + if is_last_subtask_ind: + # pause on last index of current subtask + input("Pausing after subtask {} execution. Press any key to continue...".format(cur_subtask_ind)) + cur_subtask_ind += 1 + + # video render + if write_video: + if video_count % video_skip == 0: + video_img = [] + for cam_name in camera_names: + video_img.append(env.render(mode="rgb_array", height=512, width=512, camera_name=cam_name)) + video_img = np.concatenate(video_img, axis=1) # concatenate horizontally + if should_add_border_to_frame: + video_img = add_red_border_to_frame(video_img) + video_writer.append_data(video_img) + video_count += 1 + + if is_last_subtask_ind: + # toggle whether to add red border for next subtask + should_add_border_to_frame = (not should_add_border_to_frame) + cur_subtask_ind += 1 + + +def visualize_subtasks_with_obs( + traj_grp, + subtask_end_indices, + video_writer, + video_skip=5, + image_names=None, +): + """ + Helper function to visualize each subtask in a trajectory by writing image observations + to a video. Each subtask will toggle between having a red border around each + frame and no border in the video. + + Args: + traj_grp (hdf5 file group): hdf5 group which corresponds to the dataset trajectory to playback + subtask_end_indices (list): list containing the end index for each subtask + video_writer (imageio writer): video writer + video_skip (int): determines rate at which environment frames are written to video + image_names (list): determines which image observations are used for rendering. Pass more than + one to output a video with multiple image observations concatenated horizontally. + """ + assert image_names is not None, "error: must specify at least one image observation to use in @image_names" + assert (stage_signal is None) or (vis_stage_inds is None) + video_count = 0 + + traj_len = traj_grp["actions"].shape[0] + should_add_border_to_frame = False + cur_subtask_ind = 0 + for i in range(traj_len): + # whether we are on last index of current subtask + is_last_subtask_ind = (i == subtask_end_indices[cur_subtask_ind] - 1) + + if video_count % video_skip == 0: + # concatenate image obs together + im = [traj_grp["obs/{}".format(k)][i] for k in image_names] + frame = np.concatenate(im, axis=1) + if should_add_border_to_frame: + frame = add_red_border_to_frame(frame) + video_writer.append_data(frame) + video_count += 1 + + if is_last_subtask_ind: + # toggle whether to add red border for next subtask + should_add_border_to_frame = (not should_add_border_to_frame) + cur_subtask_ind += 1 + + +def visualize_subtasks(args): + # some arg checking + write_video = (args.video_path is not None) + + # either on-screen or video but not both + assert not (args.render and write_video) + + # either config or signals and offsets should be provided, but not both + assert (args.config is not None) or ((args.signals is not None) and (args.offsets is not None)) + assert (args.config is None) or ((args.signals is None) and (args.offsets is None)) + + # Auto-fill camera rendering info if not specified + if args.render_image_names is None: + env_meta = get_env_metadata_from_dataset(dataset_path=args.dataset) + args.render_image_names = RobomimicUtils.get_default_env_cameras(env_meta=env_meta) + + if args.render: + # on-screen rendering can only support one camera + assert len(args.render_image_names) == 1 + + if args.use_obs: + assert not args.render + else: + # create environment only if not playing back with observations + + # need to make sure ObsUtils knows which observations are images, but it doesn't matter + # for playback since observations are unused. Pass a dummy spec here. + dummy_spec = dict( + obs=dict( + low_dim=["robot0_eef_pos"], + rgb=[], + # image=[], + ), + ) + ObsUtils.initialize_obs_utils_with_obs_specs(obs_modality_specs=dummy_spec) + + env_meta = get_env_metadata_from_dataset(dataset_path=args.dataset) + env = EnvUtils.create_env_from_metadata(env_meta=env_meta, render=args.render, render_offscreen=write_video) + + # some operations for playback are env-type-specific + is_robosuite_env = EnvUtils.is_robosuite_env(env_meta) + + # get demonstrations to visualize subtasks for + dataset_path = args.dataset + demo_keys = MG_FileUtils.get_all_demos_from_dataset( + dataset_path=dataset_path, + filter_key=args.filter_key, + start=None, + n=args.n, + ) + + # we will parse the source dataset to get subtask boundaries using either the task spec in the + # provided config or the provided arguments + task_spec = None + subtask_term_signals = None + subtask_term_offset_ranges = None + if args.config is not None: + with open(args.config, 'r') as f_config: + mg_config = json.load(f_config) + task_spec = MG_TaskSpec.from_json(json_dict=mg_config["task"]["task_spec"]) + else: + subtask_term_signals = args.signals + [None] + subtask_term_offset_ranges = [(0, offset) for offset in args.offsets] + [None] + + # parse dataset to get subtask boundaries + _, subtask_indices, _, subtask_term_offset_ranges_ret = MG_FileUtils.parse_source_dataset( + dataset_path=dataset_path, + demo_keys=demo_keys, + task_spec=task_spec, + subtask_term_signals=subtask_term_signals, + subtask_term_offset_ranges=subtask_term_offset_ranges, + ) + + # apply maximum offset to each subtask boundary + offsets_to_apply = [x[1] for x in subtask_term_offset_ranges_ret] + offsets_to_apply[-1] = 0 + # subtask_indices is shape (N, S, 2) where N is num demos, S is num subtasks and each entry is 2-tuple of start and end + subtask_end_indices = subtask_indices[:, :, 1] + subtask_end_indices = subtask_end_indices + np.array(offsets_to_apply)[None] # offsets shape (1, S) + + f = h5py.File(args.dataset, "r") + + # maybe dump video + video_writer = None + if write_video: + video_writer = imageio.get_writer(args.video_path, fps=20) + + for ind in range(len(demo_keys)): + ep = demo_keys[ind] + print("Playing back episode: {}".format(ep)) + + if args.use_obs: + traj_grp = f["data/{}".format(ep)] + visualize_trajectory_with_obs( + traj_grp=traj_grp, + subtask_end_indices=subtask_end_indices[ind], + video_writer=video_writer, + video_skip=args.video_skip, + image_names=args.render_image_names, + ) + continue + + states = f["data/{}/states".format(ep)][()] + initial_state = dict(states=states[0]) + if is_robosuite_env: + initial_state["model"] = f["data/{}".format(ep)].attrs["model_file"] + visualize_subtasks_with_env( + env=env, + initial_state=initial_state, + states=states, + subtask_end_indices=subtask_end_indices[ind], + render=args.render, + video_writer=video_writer, + video_skip=args.video_skip, + camera_names=args.render_image_names, + ) + + f.close() + if write_video: + video_writer.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--dataset", + type=str, + help="path to hdf5 dataset", + required=True, + ) + parser.add_argument( + "--config", + type=str, + default=None, + help="use config to infer sequence of subtask signals and offsets", + ) + parser.add_argument( + "--signals", + type=str, + nargs='+', + default=None, + help="specify sequence of subtask termination signals for all except last subtask", + ) + parser.add_argument( + "--offsets", + type=int, + nargs='+', + default=None, + help="specify sequence of maximum subtask termination offsets for all except last subtask", + ) + parser.add_argument( + "--filter_key", + type=str, + default=None, + help="(optional) filter key, to select a subset of trajectories in the file", + ) + parser.add_argument( + "--n", + type=int, + default=None, + help="(optional) stop after n trajectories", + ) + parser.add_argument( + "--use-obs", + action='store_true', + help="visualize trajectories with dataset image observations instead of simulator", + ) + parser.add_argument( + "--render", + action='store_true', + help="on-screen rendering", + ) + parser.add_argument( + "--video_path", + type=str, + default=None, + help="(optional) render trajectories to this video file path", + ) + parser.add_argument( + "--video_skip", + type=int, + default=5, + help="render frames to video every n steps", + ) + parser.add_argument( + "--render_image_names", + type=str, + nargs='+', + default=None, + help="(optional) camera name(s) / image observation(s) to use for rendering on-screen or to video", + ) + args = parser.parse_args() + visualize_subtasks(args) diff --git a/mimicgen/utils/__init__.py b/mimicgen/utils/__init__.py new file mode 100644 index 0000000..8edfcfd --- /dev/null +++ b/mimicgen/utils/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. \ No newline at end of file diff --git a/mimicgen/utils/config_utils.py b/mimicgen/utils/config_utils.py new file mode 100644 index 0000000..a6ff70a --- /dev/null +++ b/mimicgen/utils/config_utils.py @@ -0,0 +1,500 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +A collection of utilities for working with config generators. These generators +are re-used from robomimic (https://robomimic.github.io/docs/tutorials/hyperparam_scan.html) +""" +import json +from collections.abc import Iterable + + +def set_debug_settings( + generator, + group, +): + """ + Sets config generator parameters for a quick debug run. + + Args: + generator (robomimic ConfigGenerator instance): config generator object + group (int): parameter group for these settings + """ + generator.add_param( + key="experiment.generation.guarantee", + name="", + group=group, + values=[False], + ) + generator.add_param( + key="experiment.generation.num_trials", + name="", + group=group, + values=[2], + ) + return generator + + +def set_basic_settings( + generator, + group, + source_dataset_path, + source_dataset_name, + generation_path, + guarantee, + num_traj, + num_src_demos=None, + max_num_failures=25, + num_demo_to_render=10, + num_fail_demo_to_render=25, + verbose=False, +): + """ + Sets config generator parameters for some basic data generation settings. + + Args: + generator (robomimic ConfigGenerator instance): config generator object + group (int): parameter group for these settings + source_dataset_path (str): path to source dataset + source_dataset_name (str): name to give source dataset in experiment name + generation_path (str): folder for generated data + guarantee (bool): whether to ensure @num_traj successes + num_traj (int): number of trajectories for generation + num_src_demos (int or None): number of source demos to take from @source_dataset_path + max_num_failures (int): max failures to keep + num_demo_to_render (int): max demos to render to video + num_fail_demo_to_render (int): max fail demos to render to video + verbose (bool): if True, make experiment name verbose using the passed settings + """ + + # set source dataset + generator.add_param( + key="experiment.source.dataset_path", + name="src" if source_dataset_name is not None else "", + group=group, + values=[source_dataset_path], + value_names=[source_dataset_name], + ) + + # set number of demos to use from source dataset + generator.add_param( + key="experiment.source.n", + name="n_src" if verbose else "", + group=group, + values=[num_src_demos], + ) + + # set generation settings + generator.add_param( + key="experiment.generation.path", + name="", + group=group, + values=[generation_path], + ) + generator.add_param( + key="experiment.generation.guarantee", + name="gt" if verbose else "", + group=group, + values=[guarantee], + value_names=["t" if guarantee else "f"], + ) + generator.add_param( + key="experiment.generation.num_trials", + name="nt" if verbose else "", + group=group, + values=[num_traj], + ) + generator.add_param( + key="experiment.max_num_failures", + name="", + group=group, + values=[max_num_failures], + ) + generator.add_param( + key="experiment.num_demo_to_render", + name="", + group=group, + values=[num_demo_to_render], + ) + generator.add_param( + key="experiment.num_fail_demo_to_render", + name="", + group=group, + values=[num_fail_demo_to_render], + ) + + return generator + + +def set_obs_settings( + generator, + group, + collect_obs, + camera_names, + camera_height, + camera_width, +): + """ + Sets config generator parameters for collecting observations. + """ + generator.add_param( + key="obs.collect_obs", + name="", + group=group, + values=[collect_obs], + ) + generator.add_param( + key="obs.camera_names", + name="", + group=group, + values=[camera_names], + ) + generator.add_param( + key="obs.camera_height", + name="", + group=group, + values=[camera_height], + ) + generator.add_param( + key="obs.camera_width", + name="", + group=group, + values=[camera_width], + ) + return generator + + +def set_subtask_settings( + generator, + group, + base_config_file, + select_src_per_subtask, + subtask_term_offset_range=None, + selection_strategy=None, + selection_strategy_kwargs=None, + action_noise=None, + num_interpolation_steps=None, + num_fixed_steps=None, + verbose=False, +): + """ + Sets config generator parameters for each subtask. + + Args: + generator (robomimic ConfigGenerator instance): config generator object + group (int): parameter group for these settings + base_config_file (str): path to base config file being used for generating configs + select_src_per_subtask (bool): whether to select src demo for each subtask + subtask_term_offset_range (list or None): if provided, should be list of 2-tuples, one + entry per subtask, with the last entry being None + selection_strategy (str or None): src demo selection strategy + selection_strategy_kwargs (dict or None): kwargs for selection strategy + action_noise (float or list or None): action noise for all subtasks + num_interpolation_steps (int or list or None): interpolation steps for all subtasks + num_fixed_steps (int or list or None): interpolation steps for all subtasks + verbose (bool): if True, make experiment name verbose using the passed settings + """ + + # get number of subtasks + with open(base_config_file, 'r') as f: + config = json.load(f) + num_subtasks = len(config["task"]["task_spec"]) + + # whether to select a different source demonstration for each subtask + generator.add_param( + key="experiment.generation.select_src_per_subtask", + name="select_src_per_subtask" if verbose else "", + group=group, + values=[select_src_per_subtask], + value_names=["t" if select_src_per_subtask else "f"], + ) + + # settings for each subtask + + # offset range + if subtask_term_offset_range is not None: + assert len(subtask_term_offset_range) == num_subtasks + for i in range(num_subtasks): + if (i == num_subtasks - 1): + assert subtask_term_offset_range[i] is None + else: + assert (subtask_term_offset_range[i] is None) or (len(subtask_term_offset_range[i]) == 2) + generator.add_param( + key="task.task_spec.subtask_{}.subtask_term_offset_range".format(i + 1), + name="offset" if (verbose and (i == 0)) else "", + group=group, + values=[subtask_term_offset_range[i]], + ) + + # selection strategy + if selection_strategy is not None: + for i in range(num_subtasks): + generator.add_param( + key="task.task_spec.subtask_{}.selection_strategy".format(i + 1), + name="ss" if (verbose and (i == 0)) else "", + group=group, + values=[selection_strategy], + ) + + # selection kwargs + if selection_strategy_kwargs is not None: + for i in range(num_subtasks): + generator.add_param( + key="task.task_spec.subtask_{}.selection_strategy_kwargs".format(i + 1), + name="", + group=group, + values=[selection_strategy_kwargs], + ) + + # action noise + if action_noise is not None: + if not isinstance(action_noise, Iterable): + action_noise = [action_noise for _ in range(num_subtasks)] + assert len(action_noise) == num_subtasks + for i in range(num_subtasks): + generator.add_param( + key="task.task_spec.subtask_{}.action_noise".format(i + 1), + name="noise" if (verbose and (i == 0)) else "", + group=group, + values=[action_noise[i]], + ) + + # interpolation + if num_interpolation_steps is not None: + if not isinstance(num_interpolation_steps, Iterable): + num_interpolation_steps = [num_interpolation_steps for _ in range(num_subtasks)] + assert len(num_interpolation_steps) == num_subtasks + for i in range(num_subtasks): + generator.add_param( + key="task.task_spec.subtask_{}.num_interpolation_steps".format(i + 1), + name="ni" if (verbose and (i == 0)) else "", + group=group, + values=[num_interpolation_steps[i]], + ) + if num_fixed_steps is not None: + if not isinstance(num_fixed_steps, Iterable): + num_fixed_steps = [num_fixed_steps for _ in range(num_subtasks)] + assert len(num_fixed_steps) == num_subtasks + for i in range(num_subtasks): + generator.add_param( + key="task.task_spec.subtask_{}.num_fixed_steps".format(i + 1), + name="ni" if (verbose and (i == 0)) else "", + group=group, + values=[num_fixed_steps[i]], + ) + + return generator + + +def set_learning_settings_for_bc_rnn( + generator, + group, + modality, + seq_length=10, + low_dim_keys=None, + image_keys=None, + crop_size=None, +): + """ + Sets config generator parameters for robomimic BC-RNN training runs. + + Args: + generator (robomimic ConfigGenerator instance): config generator object + group (int): parameter group for these settings + modality (str): whether this is a low-dim or image observation run + seq_length (int): BC-RNN context length + low_dim_keys (list or None): if provided, set low-dim observation keys, else use defaults + image_keys (list or None): if provided, set image observation keys, else use defaults + crop_size (tuple or None): if provided, size of crop to use for pixel shift augmentation + """ + supported_modalities = ["low_dim", "image"] + assert modality in supported_modalities, "got modality {} not in supported modalities {}".format(modality, supported_modalities) + + # setup RNN with GMM and desired seq_length + generator.add_param( + key="train.seq_length", + name="", + group=group, + values=[seq_length], + ) + generator.add_param( + key="algo.rnn.horizon", + name="", + group=group, + values=[seq_length], + ) + generator.add_param( + key="algo.rnn.enabled", + name="", + group=group, + values=[True], + ) + generator.add_param( + key="algo.gmm.enabled", + name="", + group=group, + values=[True], + ) + actor_layer_dims = [] + generator.add_param( + key="algo.actor_layer_dims", + name="", + group=group, + values=[actor_layer_dims], + ) + + # 4 data workers and low-dim cache mode seems to work well for both low-dim and image observations + generator.add_param( + key="train.num_data_workers", + name="", + group=group, + values=[4], + ) + generator.add_param( + key="train.hdf5_cache_mode", + name="", + group=group, + values=["low_dim"], + ) + + # modality-specific defaults + if modality == "image": + + # epoch settings + epoch_every_n_steps = 500 + validation_epoch_every_n_steps = 50 + eval_rate = 20 + + # learning settings + num_epochs = 600 + batch_size = 16 + policy_lr = 1e-4 + rnn_hidden_dim = 1000 + + # observation settings + if low_dim_keys is None: + low_dim_keys = ["robot0_eef_pos", "robot0_eef_quat", "robot0_gripper_qpos"] + if image_keys is None: + image_keys = ["agentview_image", "robot0_eye_in_hand_image"] + if crop_size is None: + crop_size = (76, 76) + + generator.add_param( + key="observation.encoder.rgb", + name="", + group=group, + values=[{ + "core_class": "VisualCore", + "core_kwargs": { + "feature_dimension": 64, + "flatten": True, + "backbone_class": "ResNet18Conv", + "backbone_kwargs": { + "pretrained": False, + "input_coord_conv": False, + }, + "pool_class": "SpatialSoftmax", + "pool_kwargs": { + "num_kp": 32, + "learnable_temperature": False, + "temperature": 1.0, + "noise_std": 0.0, + "output_variance": False, + }, + }, + "obs_randomizer_class": "CropRandomizer", + "obs_randomizer_kwargs": { + "crop_height": crop_size[0], + "crop_width": crop_size[1], + "num_crops": 1, + "pos_enc": False, + }, + }], + ) + + else: + + # epoch settings + epoch_every_n_steps = 100 + validation_epoch_every_n_steps = 10 + eval_rate = 50 + + # learning settings + num_epochs = 2000 + batch_size = 100 + policy_lr = 1e-3 + rnn_hidden_dim = 400 + + # observation settings + if low_dim_keys is None: + low_dim_keys = ["robot0_eef_pos", "robot0_eef_quat", "robot0_gripper_qpos", "object"] + if image_keys is None: + image_keys = [] + + generator.add_param( + key="observation.modalities.obs.low_dim", + name="", + group=group, + values=[low_dim_keys], + ) + generator.add_param( + key="observation.modalities.obs.rgb", + name="", + group=group, + values=[image_keys], + ) + + # epoch settings + generator.add_param( + key="experiment.epoch_every_n_steps", + name="", + group=group, + values=[epoch_every_n_steps], + ) + generator.add_param( + key="experiment.validation_epoch_every_n_steps", + name="", + group=group, + values=[validation_epoch_every_n_steps], + ) + generator.add_param( + key="experiment.save.every_n_epochs", + name="", + group=group, + values=[eval_rate], + ) + generator.add_param( + key="experiment.rollout.rate", + name="", + group=group, + values=[eval_rate], + ) + + # learning settings + generator.add_param( + key="train.num_epochs", + name="", + group=group, + values=[num_epochs], + ) + generator.add_param( + key="train.batch_size", + name="", + group=group, + values=[batch_size], + ) + generator.add_param( + key="algo.optim_params.policy.learning_rate.initial", + name="", + group=group, + values=[policy_lr], + ) + generator.add_param( + key="algo.rnn.hidden_dim", + name="", + group=group, + values=[rnn_hidden_dim], + ) + + return generator diff --git a/mimicgen/utils/file_utils.py b/mimicgen/utils/file_utils.py new file mode 100644 index 0000000..2dea426 --- /dev/null +++ b/mimicgen/utils/file_utils.py @@ -0,0 +1,475 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +A collection of utilities related to files. +""" +import os +import h5py +import json +import time +import datetime +import shutil +import shlex +import tempfile +import gdown +import numpy as np + +from glob import glob +from tqdm import tqdm + +import robomimic +import robomimic.utils.tensor_utils as TensorUtils +from robomimic.utils.file_utils import url_is_alive + +import mimicgen +from mimicgen.datagen.datagen_info import DatagenInfo + + +def write_json(json_dic, json_path): + """ + Write dictionary to json file. + """ + with open(json_path, 'w') as f: + # preserve original key ordering + json.dump(json_dic, f, sort_keys=False, indent=4) + + +def get_all_demos_from_dataset( + dataset_path, + filter_key=None, + start=None, + n=None, +): + """ + Helper function to get demonstration keys from robomimic hdf5 dataset. + + Args: + dataset_path (str): path to hdf5 dataset + filter_key (str or None): name of filter key + start (int or None): demonstration index to start from + n (int or None): number of consecutive demonstrations to retrieve + + Returns: + demo_keys (list): list of demonstration keys + """ + f = h5py.File(dataset_path, "r") + + # list of all demonstration episodes (sorted in increasing number order) + if filter_key is not None: + print("using filter key: {}".format(filter_key)) + demos = [elem.decode("utf-8") for elem in np.array(f["mask/{}".format(filter_key)])] + else: + demos = list(f["data"].keys()) + inds = np.argsort([int(elem[5:]) for elem in demos]) + demo_keys = [demos[i] for i in inds] + if start is not None: + demo_keys = demo_keys[start:] + if n is not None: + demo_keys = demo_keys[:n] + + f.close() + return demo_keys + + +def get_env_interface_info_from_dataset( + dataset_path, + demo_keys, +): + """ + Gets environment interface information from source dataset. + + Args: + dataset_path (str): path to hdf5 dataset + demo_keys (list): list of demonstration keys to extract info from + + Returns: + env_interface_name (str): name of environment interface class + env_interface_type (str): type of environment interface + """ + f = h5py.File(dataset_path, "r") + env_interface_names = [] + env_interface_types = [] + for ep in demo_keys: + datagen_info_key = "data/{}/datagen_info".format(ep) + assert datagen_info_key in f, "Could not find MimicGen metadata in dataset {}. Ensure you have run prepare_src_dataset.py on this hdf5".format(dataset_path) + env_interface_names.append(f[datagen_info_key].attrs["env_interface_name"]) + env_interface_types.append(f[datagen_info_key].attrs["env_interface_type"]) + f.close() + + # ensure all source demos are consistent + env_interface_name = env_interface_names[0] + env_interface_type = env_interface_types[0] + assert all(elem == env_interface_name for elem in env_interface_names) + assert all(elem == env_interface_type for elem in env_interface_types) + return env_interface_name, env_interface_type + + +def parse_source_dataset( + dataset_path, + demo_keys, + task_spec=None, + subtask_term_signals=None, + subtask_term_offset_ranges=None, +): + """ + Parses a source dataset to extract info needed for data generation (DatagenInfo instances) and + subtask indices that split each source dataset trajectory into contiguous subtask segments. + + Args: + dataset_path (str): path to hdf5 dataset + demo_keys (list): list of demo keys to use from dataset path + task_spec (MG_TaskSpec instance or None): task spec object, which will be used to + infer the sequence of subtask termination signals and offset ranges. + subtask_term_signals (list or None): sequence of subtask termination signals, which + should only be provided if not providing @task_spec. Should have an entry per subtask + and the last subtask entry should be None, since the final subtask ends when the + task ends. + subtask_term_offset_ranges (list or None): sequence of subtask termination offset ranges, which + should only be provided if not providing @task_spec. Should have an entry per subtask + and the last subtask entry should be None or (0, 0), since the final subtask ends when the + task ends. + + Returns: + + datagen_infos (list): list of DatagenInfo instances, one per source + demonstration. Each instance has entries with leading dimension [T, ...], + the length of the trajectory. + + subtask_indices (np.array): array of shape (N, S, 2) where N is the number of + demos and S is the number of subtasks for this task. Each entry is + a pair of integers that represents the index at which a subtask + segment starts and where it is completed. + + subtask_term_signals (list): sequence of subtask termination signals + + subtask_term_offset_ranges (list): sequence of subtask termination offset ranges + """ + + # should provide either task_spec or the subtask termination lists, but not both + assert (task_spec is not None) or ((subtask_term_signals is not None) and (subtask_term_offset_ranges is not None)) + assert (task_spec is None) or ((subtask_term_signals is None) and (subtask_term_offset_ranges is None)) + + if task_spec is not None: + subtask_term_signals = [subtask_spec["subtask_term_signal"] for subtask_spec in task_spec] + subtask_term_offset_ranges = [subtask_spec["subtask_term_offset_range"] for subtask_spec in task_spec] + + assert len(subtask_term_signals) == len(subtask_term_offset_ranges) + assert subtask_term_signals[-1] is None, "end of final subtask does not need to be detected" + assert (subtask_term_offset_ranges[-1] is None) or (subtask_term_offset_ranges[-1] == (0, 0)), "end of final subtask does not need to be detected" + subtask_term_offset_ranges[-1] = (0, 0) + + f = h5py.File(dataset_path, "r") + + datagen_infos = [] + subtask_indices = [] + for ind in tqdm(range(len(demo_keys))): + ep = demo_keys[ind] + ep_grp = f["data/{}".format(ep)] + + # extract datagen info + ep_datagen_info = ep_grp["datagen_info"] + ep_datagen_info_obj = DatagenInfo( + eef_pose=ep_datagen_info["eef_pose"][:], + object_poses={ k : ep_datagen_info["object_poses"][k][:] for k in ep_datagen_info["object_poses"] }, + subtask_term_signals={ k : ep_datagen_info["subtask_term_signals"][k][:] for k in ep_datagen_info["subtask_term_signals"] }, + target_pose=ep_datagen_info["target_pose"][:], + gripper_action=ep_datagen_info["gripper_action"][:], + ) + datagen_infos.append(ep_datagen_info_obj) + + # parse subtask indices using subtask termination signals + ep_subtask_indices = [] + prev_subtask_term_ind = 0 + for subtask_ind in range(len(subtask_term_signals)): + subtask_term_signal = subtask_term_signals[subtask_ind] + if subtask_term_signal is None: + # final subtask, finishes at end of demo + subtask_term_ind = ep_grp["actions"].shape[0] + else: + # trick to detect index where first 0 -> 1 transition occurs - this will be the end of the subtask + subtask_indicators = ep_datagen_info_obj.subtask_term_signals[subtask_term_signal] + diffs = subtask_indicators[1:] - subtask_indicators[:-1] + end_ind = int(diffs.nonzero()[0][0]) + 1 + subtask_term_ind = end_ind + 1 # increment to support indexing like demo[start:end] + ep_subtask_indices.append([prev_subtask_term_ind, subtask_term_ind]) + prev_subtask_term_ind = subtask_term_ind + + # run sanity check on subtask_term_offset_range in task spec to make sure we can never + # get an empty subtask in the worst case when sampling subtask bounds: + # + # end index of subtask i + max offset of subtask i < end index of subtask i + 1 + min offset of subtask i + 1 + # + assert len(ep_subtask_indices) == len(subtask_term_signals), "mismatch in length of extracted subtask info and number of subtasks" + for i in range(1, len(ep_subtask_indices)): + prev_max_offset_range = subtask_term_offset_ranges[i - 1][1] + assert ep_subtask_indices[i - 1][1] + prev_max_offset_range < ep_subtask_indices[i][1] + subtask_term_offset_ranges[i][0], \ + "subtask sanity check violation in demo key {} with subtask {} end ind {}, subtask {} max offset {}, subtask {} end ind {}, and subtask {} min offset {}".format( + demo_keys[ind], i - 1, ep_subtask_indices[i - 1][1], i - 1, prev_max_offset_range, i, ep_subtask_indices[i][1], i, subtask_term_offset_ranges[i][0]) + + subtask_indices.append(ep_subtask_indices) + f.close() + + # convert list of lists to array for easy indexing + subtask_indices = np.array(subtask_indices) + + return datagen_infos, subtask_indices, subtask_term_signals, subtask_term_offset_ranges + + +def write_demo_to_hdf5( + folder, + env, + initial_state, + states, + observations, + datagen_info, + actions, + src_demo_inds=None, + src_demo_labels=None, +): + """ + Helper function to write demonstration to an hdf5 file (robomimic format) in a folder. It will be + named using a timestamp. + + Args: + folder (str): folder to write hdf5 to + env (robomimic EnvBase instance): simulation environment + initial_state (dict): dictionary corresponding to initial simulator state (see robomimic dataset structure for more information) + states (list): list of simulator states + observations (list): list of observation dictionaries + datagen_info (list): list of DatagenInfo instances + actions (np.array): actions per timestep + src_demo_inds (list or None): if provided, list of selected source demonstration indices for each subtask + src_demo_labels (np.array or None): same as @src_demo_inds, but repeated to have a label for each timestep of the trajectory + """ + + # name hdf5 based on timestamp + timestamp = time.time() + time_str = datetime.datetime.fromtimestamp(timestamp).strftime('date_%m_%d_%Y_time_%H_%M_%S') + dataset_path = os.path.join(folder, "{}.hdf5".format(time_str)) + data_writer = h5py.File(dataset_path, "w") + data_grp = data_writer.create_group("data") + data_grp.attrs["timestamp"] = timestamp + data_grp.attrs["readable_timestamp"] = time_str + + # single episode + ep_data_grp = data_grp.create_group("demo_0") + + # write actions + ep_data_grp.create_dataset("actions", data=np.array(actions)) + + # write simulator states + if isinstance(states[0], dict): + states = TensorUtils.list_of_flat_dict_to_dict_of_list(states) + for k in states: + ep_data_grp.create_dataset("states/{}".format(k), data=np.array(states[k])) + else: + ep_data_grp.create_dataset("states", data=np.array(states)) + + # write observations + obs = TensorUtils.list_of_flat_dict_to_dict_of_list(observations) + for k in obs: + ep_data_grp.create_dataset("obs/{}".format(k), data=np.array(obs[k]), compression="gzip") + + # write datagen info + datagen_info = TensorUtils.list_of_flat_dict_to_dict_of_list([x.to_dict() for x in datagen_info]) + for k in datagen_info: + if k in ["object_poses", "subtask_term_signals"]: + # convert list of dict to dict of list again + datagen_info[k] = TensorUtils.list_of_flat_dict_to_dict_of_list(datagen_info[k]) + for k2 in datagen_info[k]: + datagen_info[k][k2] = np.array(datagen_info[k][k2]) + ep_data_grp.create_dataset("datagen_info/{}/{}".format(k, k2), data=np.array(datagen_info[k][k2])) + else: + ep_data_grp.create_dataset("datagen_info/{}".format(k), data=np.array(datagen_info[k])) + + # maybe write which source demonstrations generated this episode + if src_demo_inds is not None: + ep_data_grp.create_dataset("src_demo_inds", data=np.array(src_demo_inds)) + if src_demo_labels is not None: + ep_data_grp.create_dataset("src_demo_labels", data=np.array(src_demo_labels)) + + # episode metadata + if ("model" in initial_state) and (initial_state["model"] is not None): + # only for robosuite envs + ep_data_grp.attrs["model_file"] = initial_state["model"] # model xml for this episode + ep_data_grp.attrs["num_samples"] = actions.shape[0] # number of transitions in this episode + + # global metadata + data_grp.attrs["total"] = actions.shape[0] + data_grp.attrs["env_args"] = json.dumps(env.serialize(), indent=4) # environment info + data_writer.close() + + +def merge_all_hdf5( + folder, + new_hdf5_path, + delete_folder=False, + dry_run=False, + return_horizons=False, +): + """ + Helper function to take all hdf5s in @folder and merge them into a single one. + Returns the number of hdf5s that were merged. + """ + source_hdf5s = glob(os.path.join(folder, "*.hdf5")) + + # get all timestamps and sort files from lowest to highest + timestamps = [] + filtered_source_hdf5s = [] + for source_hdf5_path in source_hdf5s: + try: + f = h5py.File(source_hdf5_path, "r") + except Exception as e: + print("WARNING: problem with file {}".format(source_hdf5_path)) + print("Exception: {}".format(e)) + continue + filtered_source_hdf5s.append(source_hdf5_path) + timestamps.append(f["data"].attrs["timestamp"]) + f.close() + + assert len(timestamps) == len(filtered_source_hdf5s) + inds = np.argsort(timestamps) + sorted_hdf5s = [filtered_source_hdf5s[i] for i in inds] + + if dry_run: + if return_horizons: + horizons = [] + for source_hdf5_path in sorted_hdf5s: + with h5py.File(source_hdf5_path, "r") as f: + horizons.append(f["data"].attrs["total"]) + return len(sorted_hdf5s), horizons + return len(sorted_hdf5s) + + # write demos in order to new file + f_new = h5py.File(new_hdf5_path, "w") + f_new_grp = f_new.create_group("data") + + env_meta_str = None + total = 0 + if return_horizons: + horizons = [] + for i, source_hdf5_path in enumerate(sorted_hdf5s): + with h5py.File(source_hdf5_path, "r") as f: + # copy this episode over under a different name + demo_str = "demo_{}".format(i) + f.copy("data/demo_0", f_new_grp, name=demo_str) + if return_horizons: + horizons.append(f["data"].attrs["total"]) + total += f["data"].attrs["total"] + if env_meta_str is None: + env_meta_str = f["data"].attrs["env_args"] + + f_new["data"].attrs["total"] = total + f_new["data"].attrs["env_args"] = env_meta_str if env_meta_str is not None else "" + f_new.close() + + if delete_folder: + print("removing folder at path {}".format(folder)) + shutil.rmtree(folder) + + if return_horizons: + return len(sorted_hdf5s), horizons + return len(sorted_hdf5s) + + +def download_url_from_gdrive(url, download_dir, check_overwrite=True): + """ + Downloads a file at a URL from Google Drive. + + Example usage: + url = https://drive.google.com/file/d/1DABdqnBri6-l9UitjQV53uOq_84Dx7Xt/view?usp=drive_link + download_dir = "/tmp" + download_url_from_gdrive(url, download_dir, check_overwrite=True) + + Args: + url (str): url string + download_dir (str): path to directory where file should be downloaded + check_overwrite (bool): if True, will sanity check the download fpath to make sure a file of that name + doesn't already exist there + """ + assert url_is_alive(url), "@download_url_from_gdrive got unreachable url: {}".format(url) + + with tempfile.TemporaryDirectory() as td: + # HACK: Change directory to temp dir, download file there, and then move the file to desired directory. + # We do this because we do not know the name of the file beforehand. + cur_dir = os.getcwd() + os.chdir(td) + fpath = gdown.download(url, quiet=False, fuzzy=True) + fname = os.path.basename(fpath) + file_to_write = os.path.join(download_dir, fname) + if check_overwrite and os.path.exists(file_to_write): + user_response = input(f"Warning: file {file_to_write} already exists. Overwrite? y/n\n") + assert user_response.lower() in {"yes", "y"}, f"Did not receive confirmation. Aborting download." + shutil.move(fpath, file_to_write) + os.chdir(cur_dir) + + +def config_generator_to_script_lines(generator, config_dir): + """ + Takes a robomimic ConfigGenerator and uses it to + generate a set of training configs, and a set of bash command lines + that correspond to each training run (one per config). Note that + the generator's script_file will be overridden to be a temporary file that + will be removed from disk. + + Args: + generator (ConfigGenerator instance or list): generator(s) + to use for generating configs and training runs + + config_dir (str): path to directory where configs will be generated + + Returns: + config_files (list): a list of config files that were generated + + run_lines (list): a list of strings that are training commands, one per config + """ + + # make sure config dir exists + if not os.path.exists(config_dir): + os.makedirs(config_dir) + + # support one or more config generators + if not isinstance(generator, list): + generator = [generator] + + all_run_lines = [] + for gen in generator: + + # set new config directory by copying base config file from old location to new directory + base_config_file = gen.base_config_file + config_name = os.path.basename(base_config_file) + new_base_config_file = os.path.join(config_dir, config_name) + shutil.copyfile( + base_config_file, + new_base_config_file, + ) + gen.base_config_file = new_base_config_file + + # we'll write script file to a temp dir and parse it from there to get the training commands + with tempfile.TemporaryDirectory() as td: + gen.script_file = os.path.join(td, "tmp.sh") + + # generate configs + gen.generate() + + # collect training commands + with open(gen.script_file, "r") as f: + f_lines = f.readlines() + run_lines = [line for line in f_lines if line.startswith("python")] + all_run_lines += run_lines + + os.remove(gen.base_config_file) + + # get list of generated configs too + config_files = [] + config_file_dict = dict() + for line in all_run_lines: + cmd = shlex.split(line) + config_file_name = cmd[cmd.index("--config") + 1] + config_files.append(config_file_name) + assert config_file_name not in config_file_dict, "got duplicate config name {}".format(config_file_name) + config_file_dict[config_file_name] = 1 + + return config_files, all_run_lines diff --git a/mimicgen/utils/misc_utils.py b/mimicgen/utils/misc_utils.py new file mode 100644 index 0000000..2752fd8 --- /dev/null +++ b/mimicgen/utils/misc_utils.py @@ -0,0 +1,254 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +A collection of miscellaneous utilities. +""" +import time +import json +import numpy as np + +from collections import deque, OrderedDict +from contextlib import contextmanager + + +def add_red_border_to_frame(frame, ratio=0.02): + """ + Add a red border to an image frame. + """ + border_size_x = max(1, round(ratio * frame.shape[0])) + border_size_y = max(1, round(ratio * frame.shape[1])) + + frame[:border_size_x, :, :] = [255., 0., 0.] + frame[-border_size_x:, :, :] = [255., 0., 0.] + frame[:, :border_size_y, :] = [255., 0., 0.] + frame[:, -border_size_y:, :] = [255., 0., 0.] + return frame + + +class Grid(object): + """ + Keep track of a list of values, and point to a single value at a time. + """ + def __init__(self, values, initial_ind=0): + self.values = list(values) + self.ind = initial_ind + self.n = len(self.values) + + def get(self): + return self.values[self.ind] + + def next(self): + self.ind = min(self.ind + 1, self.n - 1) + return self.get() + + def prev(self): + self.ind = max(self.ind - 1, 0) + return self.get() + + +class Timer(object): + """ + A simple timer. + """ + def __init__(self, history=100, ignore_first=False): + """ + Args: + history (int): number of recent timesteps to record for reporting statistics + """ + self.total_time = 0. + self.calls = 0 + self.start_time = 0. + self.last_diff = 0. + self.average_time = 0. + self.min_diff = float("inf") + self.max_diff = 0. + self._measurements = deque(maxlen=history) + self._enabled = True + self.ignore_first = ignore_first + self._had_first = False + + def enable(self): + """ + Enable measurements with this timer. + """ + self._enabled = True + + def disable(self): + """ + Disable measurements with this timer. + """ + self._enabled = False + + def tic(self): + # using time.time instead of time.clock because time time.clock + # does not normalize for multithreading + self.start_time = time.time() + + def toc(self): + if self._enabled: + + if self.ignore_first and (self.start_time > 0. and not self._had_first): + self._had_first = True + return time.time() - self.start_time + + self.last_diff = time.time() - self.start_time + self.total_time += self.last_diff + self.calls += 1 + self.average_time = self.total_time / self.calls + self.min_diff = min(self.min_diff, self.last_diff) + self.max_diff = max(self.max_diff, self.last_diff) + self._measurements.append(self.last_diff) + last_diff = self.last_diff + return last_diff + + @contextmanager + def timed(self): + self.tic() + yield + self.toc() + + def report_stats(self, verbose=False): + stats = OrderedDict() + stats["global"] = OrderedDict( + mean=self.average_time, + min=self.min_diff, + max=self.max_diff, + num=self.calls, + ) + num = len(self._measurements) + stats["local"] = OrderedDict() + if num > 0: + stats["local"] = OrderedDict( + mean=np.mean(self._measurements), + std=np.std(self._measurements), + min=np.min(self._measurements), + max=np.max(self._measurements), + num=num, + ) + if verbose: + stats["local"]["values"] = list(self._measurements) + return stats + + +class Rate(object): + """ + Convenience class for enforcing rates in loops. Modeled after rospy.Rate. + + See http://docs.ros.org/en/jade/api/rospy/html/rospy.timer-pysrc.html#Rate.sleep + """ + def __init__(self, hz): + """ + Args: + hz (int): frequency to enforce + """ + self.update_hz(hz) + + def update_hz(self, hz): + """ + Update rate to enforce. + """ + self.hz = hz + self.last_time = time.time() + self.sleep_duration = (1. / hz) + + def _remaining(self, curr_time): + """ + Calculate time remaining for rate to sleep. + """ + assert curr_time >= self.last_time, "time moved backwards!" + elapsed = curr_time - self.last_time + return self.sleep_duration - elapsed + + def sleep(self): + """ + Attempt to sleep at the specified rate in hz, by taking the time + elapsed since the last call to this function into account. + """ + curr_time = time.time() + remaining = self._remaining(curr_time) + if remaining > 0: + time.sleep(remaining) + + # assume successful rate sleeping + self.last_time = self.last_time + self.sleep_duration + + # NOTE: this commented line is what we used to do, but this enforces a slower rate + # self.last_time = time.time() + + # detect time jumping forwards (e.g. loop is too slow) + if curr_time - self.last_time > self.sleep_duration * 2: + # we didn't sleep at all + self.last_time = curr_time + + +class RateMeasure(object): + """ + Measure approximate time intervals of code execution by calling @measure + """ + def __init__(self, name=None, history=100, freq_threshold=None): + self._timer = Timer(history=history, ignore_first=True) + self._timer.tic() + self.name = name + self.freq_threshold = freq_threshold + self._enabled = True + self._first = False + self.sum = 0. + self.calls = 0 + + def enable(self): + """ + Enable measurements. + """ + self._timer.enable() + self._enabled = True + + def disable(self): + """ + Disable measurements. + """ + self._timer.disable() + self._enabled = False + + def measure(self): + """ + Take a measurement of the time elapsed since the last @measure call + and also return the time elapsed. + """ + interval = self._timer.toc() + self._timer.tic() + self.sum += (1. / interval) + self.calls += 1 + if self._enabled and (self.freq_threshold is not None) and ((1. / interval) < self.freq_threshold): + print("WARNING: RateMeasure {} violated threshold {} hz with measurement {} hz".format(self.name, self.freq_threshold, (1. / interval))) + return (interval, True) + return (interval, False) + + def report_stats(self, verbose=False): + """ + Report statistics over measurements, converting timer measurements into frequencies. + """ + stats = self._timer.report_stats(verbose=verbose) + stats["name"] = self.name + if stats["global"]["num"] > 0: + stats["global"] = OrderedDict( + mean=(self.sum / float(self.calls)), + min=(1. / stats["global"]["max"]), + max=(1. / stats["global"]["min"]), + num=stats["global"]["num"], + ) + if len(stats["local"]) > 0: + measurements = [1. / x for x in self._timer._measurements] + stats["local"] = OrderedDict( + mean=np.mean(measurements), + std=np.std(measurements), + min=np.min(measurements), + max=np.max(measurements), + num=stats["local"]["num"], + ) + return stats + + def __str__(self): + stats = self.report_stats(verbose=False) + return json.dumps(stats, indent=4) diff --git a/mimicgen/utils/pose_utils.py b/mimicgen/utils/pose_utils.py new file mode 100644 index 0000000..c31c4d1 --- /dev/null +++ b/mimicgen/utils/pose_utils.py @@ -0,0 +1,297 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +A collection of utilities for working with poses. +""" + +import math +import collections +import numpy as np + +import robosuite +import robosuite.utils.transform_utils as T + + +def make_pose(pos, rot): + """ + Make homogenous pose matrices from a set of translation vectors and rotation matrices. + + Args: + pos (np.array): batch of position vectors with last dimension of 3 + rot (np.array): batch of rotation matrices with last 2 dimensions of (3, 3) + + Returns: + pose (np.array): batch of pose matrices with last 2 dimensions of (4, 4) + """ + assert pos.shape[:-1] == rot.shape[:-2] + assert pos.shape[-1] == rot.shape[-2] == rot.shape[-1] == 3 + pose = np.zeros(pos.shape[:-1] + (4, 4)) + pose[..., :3, :3] = rot + pose[..., :3, 3] = pos + pose[..., 3, 3] = 1. + return pose + + +def unmake_pose(pose): + """ + Split homogenous pose matrices back into translation vectors and rotation matrices. + + Args: + pose (np.array): batch of pose matrices with last 2 dimensions of (4, 4) + + Returns: + pos (np.array): batch of position vectors with last dimension of 3 + rot (np.array): batch of rotation matrices with last 2 dimensions of (3, 3) + """ + return pose[..., :3, 3], pose[..., :3, :3] + + +def pose_inv(pose): + """ + Computes the inverse of homogenous pose matrices. + + Note that the inverse of a pose matrix is the following: + [R t; 0 1]^-1 = [R.T -R.T*t; 0 1] + + Args: + pose (np.array): batch of pose matrices with last 2 dimensions of (4, 4) + + + Returns: + inv_pose (np.array): batch of inverse pose matrices with last 2 dimensions of (4, 4) + """ + num_axes = len(pose.shape) + assert num_axes >= 2 + + inv_pose = np.zeros_like(pose) + + # gymnastics to take transpose of last 2 dimensions + inv_pose[..., :3, :3] = np.transpose(pose[..., :3, :3], tuple(range(num_axes - 2)) + (num_axes - 1, num_axes - 2)) + + # note: numpy matmul wants shapes [..., 3, 3] x [..., 3, 1] -> [..., 3, 1] so we add a dimension and take it away after + inv_pose[..., :3, 3] = np.matmul(-inv_pose[..., :3, :3], pose[..., :3, 3:4])[..., 0] + inv_pose[..., 3, 3] = 1.0 + return inv_pose + + +def pose_in_A_to_pose_in_B(pose_in_A, pose_A_in_B): + """ + Converts homogenous matrices corresponding to a point C in frame A + to homogenous matrices corresponding to the same point C in frame B. + + Args: + pose_in_A (np.array): batch of homogenous matrices corresponding to the pose of C in frame A + pose_A_in_B (np.array): batch of homogenous matrices corresponding to the pose of A in frame B + + Returns: + pose_in_B (np.array): batch of homogenous matrices corresponding to the pose of C in frame B + """ + return np.matmul(pose_A_in_B, pose_in_A) + + +def quat2axisangle(quat): + """ + Converts (x, y, z, w) quaternion to axis-angle format. + Returns a unit vector direction and an angle. + + NOTE: this differs from robosuite's function because it returns + both axis and angle, not axis * angle. + """ + + # conversion from axis-angle to quaternion: + # qw = cos(theta / 2); qx, qy, qz = u * sin(theta / 2) + + # normalize qx, qy, qz by sqrt(qx^2 + qy^2 + qz^2) = sqrt(1 - qw^2) + # to extract the unit vector + + # clipping for scalar with if-else is orders of magnitude faster than numpy + if quat[3] > 1.: + quat[3] = 1. + elif quat[3] < -1.: + quat[3] = -1. + + den = np.sqrt(1. - quat[3] * quat[3]) + if math.isclose(den, 0.): + # This is (close to) a zero degree rotation, immediately return + return np.zeros(3), 0. + + return quat[:3] / den, 2. * math.acos(quat[3]) + + +def axisangle2quat(axis, angle): + """ + Converts axis-angle to (x, y, z, w) quat. + + NOTE: this differs from robosuite's function because it accepts + both axis and angle as arguments, not axis * angle. + """ + + # handle zero-rotation case + if math.isclose(angle, 0.): + return np.array([0., 0., 0., 1.]) + + # make sure that axis is a unit vector + assert math.isclose(np.linalg.norm(axis), 1., abs_tol=1e-3) + + q = np.zeros(4) + q[3] = np.cos(angle / 2.) + q[:3] = axis * np.sin(angle / 2.) + return q + + +def quat_slerp(q1, q2, tau): + """ + Adapted from robosuite. + """ + if tau == 0.0: + return q1 + elif tau == 1.0: + return q2 + d = np.dot(q1, q2) + if abs(abs(d) - 1.0) < np.finfo(float).eps * 4.: + return q1 + if d < 0.0: + # invert rotation + d = -d + q2 *= -1.0 + angle = math.acos(np.clip(d, -1, 1)) + if abs(angle) < np.finfo(float).eps * 4.: + return q1 + isin = 1.0 / math.sin(angle) + q1 = q1 * math.sin((1.0 - tau) * angle) * isin + q2 = q2 * math.sin(tau * angle) * isin + q1 = q1 + q2 + return q1 + + +def interpolate_rotations(R1, R2, num_steps, axis_angle=True): + """ + Interpolate between 2 rotation matrices. If @axis_angle, interpolate the axis-angle representation + of the delta rotation, else, use slerp. + + NOTE: I have verified empirically that both methods are essentially equivalent, so pick your favorite. + """ + if axis_angle: + # delta rotation expressed as axis-angle + delta_rot_mat = R2.dot(R1.T) + delta_quat = T.mat2quat(delta_rot_mat) + delta_axis, delta_angle = quat2axisangle(delta_quat) + + # fix the axis, and chunk the angle up into steps + rot_step_size = delta_angle / num_steps + + # convert into delta rotation matrices, and then convert to absolute rotations + if delta_angle < 0.05: + # small angle - don't bother with interpolation + rot_steps = np.array([R2 for _ in range(num_steps)]) + else: + delta_rot_steps = [T.quat2mat(axisangle2quat(delta_axis, i * rot_step_size)) for i in range(num_steps)] + rot_steps = np.array([delta_rot_steps[i].dot(R1) for i in range(num_steps)]) + else: + q1 = T.mat2quat(R1) + q2 = T.mat2quat(R2) + rot_steps = np.array([T.quat2mat(quat_slerp(q1, q2, tau=(float(i) / num_steps))) for i in range(num_steps)]) + + # add in endpoint + rot_steps = np.concatenate([rot_steps, R2[None]], axis=0) + + return rot_steps + + +def interpolate_poses(pose_1, pose_2, num_steps=None, step_size=None, perturb=False): + """ + Linear interpolation between two poses. + + Args: + pose_1 (np.array): 4x4 start pose + pose_2 (np.array): 4x4 end pose + num_steps (int): if provided, specifies the number of desired interpolated points (not excluding + the start and end points). Passing 0 corresponds to no interpolation, and passing None + means that @step_size must be provided to determine the number of interpolated points. + step_size (float): if provided, will be used to infer the number of steps, by taking the norm + of the delta position vector, and dividing it by the step size + perturb (bool): if True, randomly move all the interpolated position points in a uniform, non-overlapping grid. + + Returns: + pose_steps (np.array): array of shape (N + 2, 3) corresponding to the interpolated pose path, where N is @num_steps + num_steps (int): the number of interpolated points (N) in the path + """ + assert step_size is None or num_steps is None + + pos1, rot1 = unmake_pose(pose_1) + pos2, rot2 = unmake_pose(pose_2) + + if num_steps == 0: + # skip interpolation + return np.concatenate([pos1[None], pos2[None]], axis=0), np.concatenate([rot1[None], rot2[None]], axis=0), num_steps + + delta_pos = pos2 - pos1 + if num_steps is None: + assert np.linalg.norm(delta_pos) > 0 + num_steps = math.ceil(np.linalg.norm(delta_pos) / step_size) + + num_steps += 1 # include starting pose + assert num_steps >= 2 + + # linear interpolation of positions + pos_step_size = delta_pos / num_steps + grid = np.arange(num_steps).astype(np.float64) + if perturb: + # move the interpolation grid points by up to a half-size forward or backward + perturbations = np.random.uniform( + low=-0.5, + high=0.5, + size=(num_steps - 2,), + ) + grid[1:-1] += perturbations + pos_steps = np.array([pos1 + grid[i] * pos_step_size for i in range(num_steps)]) + + # add in endpoint + pos_steps = np.concatenate([pos_steps, pos2[None]], axis=0) + + # interpolate the rotations too + rot_steps = interpolate_rotations(R1=rot1, R2=rot2, num_steps=num_steps, axis_angle=True) + + pose_steps = make_pose(pos_steps, rot_steps) + return pose_steps, num_steps - 1 + + +def transform_source_data_segment_using_object_pose( + obj_pose, + src_eef_poses, + src_obj_pose, +): + """ + Transform a source data segment (object-centric subtask segment from source demonstration) such that + the relative poses between the target eef pose frame and the object frame are preserved. Recall that + each object-centric subtask segment corresponds to one object, and consists of a sequence of + target eef poses. + + Args: + obj_pose (np.array): 4x4 object pose in current scene + src_eef_poses (np.array): pose sequence (shape [T, 4, 4]) for the sequence of end effector control poses + from the source demonstration + src_obj_pose (np.array): 4x4 object pose from the source demonstration + + Returns: + transformed_eef_poses (np.array): transformed pose sequence (shape [T, 4, 4]) + """ + + # transform source end effector poses to be relative to source object frame + + # convert these eef poses from frame A (world frame) to frame B (source object frame) + src_eef_poses_rel_obj = pose_in_A_to_pose_in_B( + pose_in_A=src_eef_poses, + pose_A_in_B=pose_inv(src_obj_pose[None]), + ) + + # apply relative poses to current object frame to obtain new target eef poses + + # convert these eef poses from frame A (new object frame) to frame B (world frame) + transformed_eef_poses = pose_in_A_to_pose_in_B( + pose_in_A=src_eef_poses_rel_obj, + pose_A_in_B=obj_pose[None], + ) + return transformed_eef_poses diff --git a/mimicgen/utils/robomimic_utils.py b/mimicgen/utils/robomimic_utils.py new file mode 100644 index 0000000..c5baff6 --- /dev/null +++ b/mimicgen/utils/robomimic_utils.py @@ -0,0 +1,156 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the NVIDIA Source Code License [see LICENSE for details]. + +""" +Collection of utilities related to robomimic. +""" +import sys +import json +import traceback +import argparse +from copy import deepcopy + +import robomimic +from robomimic.utils.log_utils import PrintLogger +import robomimic.utils.env_utils as EnvUtils +from robomimic.scripts.playback_dataset import playback_dataset, DEFAULT_CAMERAS + + +def make_print_logger(txt_file): + """ + Makes a logger that mirrors stdout and stderr to a text file. + + Args: + txt_file (str): path to txt file to write + """ + logger = PrintLogger(txt_file) + sys.stdout = logger + sys.stderr = logger + return logger + + +def create_env( + env_meta, + env_name=None, + env_class=None, + robot=None, + gripper=None, + camera_names=None, + camera_height=84, + camera_width=84, + render=None, + render_offscreen=None, + use_image_obs=None, + use_depth_obs=None, +): + """ + Helper function to create the environment from dataset metadata and arguments. + + Args: + env_meta (dict): environment metadata compatible with robomimic, see + https://robomimic.github.io/docs/modules/environments.html + env_name (str or None): if provided, override environment name + in @env_meta + env_class (class or None): if provided, use this class instead of the + one inferred from @env_meta + robot (str or None): if provided, override the robot argument in + @env_meta. Currently only supported by robosuite environments. + gripper (str or None): if provided, override the gripper argument in + @env_meta. Currently only supported by robosuite environments. + camera_names (list of str or None): list of camera names that correspond to image observations + camera_height (int): camera height for all cameras + camera_width (int): camera width for all cameras + render (bool or None): optionally override rendering behavior + render_offscreen (bool or None): optionally override rendering behavior + use_image_obs (bool or None): optionally override rendering behavior + use_depth_obs (bool or None): optionally override rendering behavior + """ + env_meta = deepcopy(env_meta) + + # maybe override some settings in environment metadata + if env_name is not None: + env_meta["env_name"] = env_name + if robot is not None: + # for now, only support this argument for robosuite environments + assert EnvUtils.is_robosuite_env(env_meta) + assert robot in ["IIWA", "Sawyer", "UR5e", "Panda", "Jaco", "Kinova3"] + env_meta["env_kwargs"]["robots"] = [robot] + if gripper is not None: + # for now, only support this argument for robosuite environments + assert EnvUtils.is_robosuite_env(env_meta) + assert gripper in ["PandaGripper", "RethinkGripper", "Robotiq85Gripper", "Robotiq140Gripper"] + env_meta["env_kwargs"]["gripper_types"] = [gripper] + + if camera_names is None: + camera_names = [] + + # create environment + env = EnvUtils.create_env_for_data_processing( + env_meta=env_meta, + env_class=env_class, + camera_names=camera_names, + camera_height=camera_height, + camera_width=camera_width, + reward_shaping=False, + render=render, + render_offscreen=render_offscreen, + use_image_obs=use_image_obs, + use_depth_obs=use_depth_obs, + ) + + return env + + +def make_dataset_video( + dataset_path, + video_path, + num_render=None, + render_image_names=None, + use_obs=False, + video_skip=5, +): + """ + Helper function to set up args and call @playback_dataset from robomimic + to get video of generated dataset. + """ + print("\nmake_dataset_video(\n\tdataset_path={},\n\tvideo_path={},{}\n)".format( + dataset_path, + video_path, + "\n\tnum_render={},".format(num_render) if num_render is not None else "", + )) + playback_args = argparse.Namespace() + playback_args.dataset = dataset_path + playback_args.filter_key = None + playback_args.n = num_render + playback_args.use_obs = use_obs + playback_args.use_actions = False + playback_args.render = False + playback_args.video_path = video_path + playback_args.video_skip = video_skip + playback_args.render_image_names = render_image_names + if (render_image_names is None): + # default robosuite + playback_args.render_image_names = ["agentview"] + playback_args.render_depth_names = None + playback_args.first = False + + try: + playback_dataset(playback_args) + except Exception as e: + res_str = "playback failed with error:\n{}\n\n{}".format(e, traceback.format_exc()) + print(res_str) + + +def get_default_env_cameras(env_meta): + """ + Get the default set of cameras for a particular robomimic environment type. + + Args: + env_meta (dict): environment metadata compatible with robomimic, see + https://robomimic.github.io/docs/modules/environments.html + + Returns: + camera_names (list of str): list of camera names that correspond to image observations + """ + return DEFAULT_CAMERAS[EnvUtils.get_env_type(env_meta=env_meta)] diff --git a/mimicgen_envs/utils/file_utils.py b/mimicgen_envs/utils/file_utils.py deleted file mode 100644 index 4a65a80..0000000 --- a/mimicgen_envs/utils/file_utils.py +++ /dev/null @@ -1,45 +0,0 @@ -# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the NVIDIA Source Code License [see LICENSE for details]. - -""" -A collection of utility functions for working with files. -""" -import os -import shutil -import tempfile -import gdown - -from robomimic.utils.file_utils import url_is_alive - - -def download_url_from_gdrive(url, download_dir, check_overwrite=True): - """ - Downloads a file at a URL from Google Drive. - - Example usage: - url = https://drive.google.com/file/d/1DABdqnBri6-l9UitjQV53uOq_84Dx7Xt/view?usp=drive_link - download_dir = "/tmp" - download_url_from_gdrive(url, download_dir, check_overwrite=True) - - Args: - url (str): url string - download_dir (str): path to directory where file should be downloaded - check_overwrite (bool): if True, will sanity check the download fpath to make sure a file of that name - doesn't already exist there - """ - assert url_is_alive(url), "@download_url_from_gdrive got unreachable url: {}".format(url) - - with tempfile.TemporaryDirectory() as td: - # HACK: Change directory to temp dir, download file there, and then move the file to desired directory. - # We do this because we do not know the name of the file beforehand. - cur_dir = os.getcwd() - os.chdir(td) - fpath = gdown.download(url, quiet=False, fuzzy=True) - fname = os.path.basename(fpath) - file_to_write = os.path.join(download_dir, fname) - if check_overwrite and os.path.exists(file_to_write): - user_response = input(f"Warning: file {file_to_write} already exists. Overwrite? y/n\n") - assert user_response.lower() in {"yes", "y"}, f"Did not receive confirmation. Aborting download." - shutil.move(fpath, file_to_write) - os.chdir(cur_dir) diff --git a/requirements-docs.txt b/requirements-docs.txt new file mode 100644 index 0000000..3381b70 --- /dev/null +++ b/requirements-docs.txt @@ -0,0 +1,3 @@ +# requirements for building sphinx docs +sphinx_book_theme==0.3.3 +recommonmark==0.7.1 diff --git a/setup.py b/setup.py index 4bc0ecb..67bf26e 100644 --- a/setup.py +++ b/setup.py @@ -11,9 +11,9 @@ long_description = ''.join(lines) setup( - name="mimicgen_envs", + name="mimicgen", packages=[ - package for package in find_packages() if package.startswith("mimicgen_envs") + package for package in find_packages() if package.startswith("mimicgen") ], install_requires=[ "numpy>=1.13.3", @@ -23,16 +23,15 @@ "imageio-ffmpeg", "gdown", "chardet", - "mujoco==2.3.2", ], eager_resources=['*'], include_package_data=True, python_requires='>=3', description="MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations", author="Ajay Mandlekar", - url="https://gitlab-master.nvidia.com/srl/mimicgen_environments", - author_email="amandlek@cs.stanford.edu", - version="0.1.0", + url="https://github.com/NVlabs/mimicgen", + author_email="amandlekar@nvidia.com", + version="1.0.0", long_description=long_description, long_description_content_type='text/markdown' )