init

.gitignore

@@ -0,0 +1,23 @@
+*~
+.hydra
+*pyc
+datasets
+results
+outputs
+/MUJOCO_LOG.TXT
+wandb
+experiments
+experiments_rw
+experiments_saved
+clip/
+gpt/
+bert/
+logs/
+
+*.mp4
+libero.egg-info/
+lotus.egg-info/
+.vscode/
+
+## random
+log.txt

README.md

@@ -0,0 +1,171 @@
+<div align="center">
+
+<b>**LOTUS: Continual Imitation Learning for Robot Manipulation Through Unsupervised Skill Discovery**</b>
+
+[[Website]](https://ut-austin-rpl.github.io/Lotus/)
+[[Paper]](https://arxiv.org/pdf/2311.02058)
+______________________________________________________________________
+</div>
+
+**LOTUS** is a continual imitation learning algorithm that empowers a physical robot to continuously and efficiently learn to solve new manipulation tasks throughout its lifespan. The core idea behind **LOTUS** is constructing an ever-growing skill library from a sequence of new tasks with a small number of corresponding task demonstrations. **LOTUS** starts with a continual skill discovery process using an open-vocabulary vision model, which extracts skills as recurring patterns presented in unstructured demonstrations. Continual skill discovery updates existing skills to avoid catastrophic forgetting of previous tasks and adds new skills to exhibit novel behaviors. **LOTUS** trains a meta-controller that flexibly composes various skills to tackle vision-based manipulation tasks in the lifelong learning process.
+
+---
+
+
+# Contents
+
+- [Installation](#Installation)
+- [Datasets](#Dataset)
+- [Getting Started](#Getting-Started)
+  - [Unsupervised Skill Discovery](#Unsupervised-Skill-Discovery)
+  - [Training](#Training)
+  - [Evaluation](#Evaluation)
+- [Acknowledgement](#Acknowledgement)
+- [Citation](#Citation)
+
+
+# Installtion
+First clone this repo, and then run the following commands in the given order to install the dependency for **LOTUS**.
+```
+conda create -n lotus python=3.9.19
+conda activate lotus
+cd Lotus
+pip install -r requirements.txt
+pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113
+pip install -e .
+```
+
+# Datasets
+We use high-quality human teleoperation demonstrations for the four task suites from [**LIBERO**] (https://github.com/Lifelong-Robot-Learning/LIBERO). To download the demonstration dataset, run:
+```python
+python libero_benchmark_scripts/download_libero_datasets.py
+```
+By default, the dataset will be stored under the ```LIBERO``` folder and all four datasets will be downloaded. To download a specific dataset, use
+```python
+python libero_benchmark_scripts/download_libero_datasets.py --datasets DATASET
+```
+where ```DATASET``` is chosen from `[libero_spatial, libero_object, libero_100, libero_goal`.
+
+
+# Getting Started
+
+In the following, we provide example scripts for unsupervised skill discovery, training and evaluation.
+
+## Unsupervised Skill Discovery
+
+The following is a example of unsupervised skill discovery on `libero_object` dataset.
+
+```shell
+cd lotus/skill_learning
+```
+
+### Encoding Representation
+```shell
+python multisensory_repr/dinov2_repr.py  --exp-name dinov2_libero_object_image_only --modality-str dinov2_agentview_eye_in_hand --feature-dim 1536
+```
+Output: 
+- `results/{exp_name}/repr/{DatasetCategoty}/{DatasetName}/embedding_{modality_str}_{feature_dim}.hdf5`
+
+### Hierarchical Agglomerative Clustering
+```shell
+python skill_discovery/hierarchical_agglomoration.py exp_name=dinov2_libero_object_image_only modality_str=dinov2_agentview_eye_in_hand repr.z_dim=1536 agglomoration.dist=cos agglomoration.footprint=global_pooling
+```
+Output: 
+- `results/{exp_name}/skill_classification/agglomoration_results/{DatasetCategoty}/{DatasetName}/{agglomoration.footprint}_{agglomoration.dist}_{modality_str}/{idx}.png`
+- `results/{exp_name}/skill_classification/trees/{DatasetCategoty}/{dataset_name}_trees_{modality_str}_{agglomoration.footprint}_{agglomoration.dist}.pkl`
+
+### Spectral Clustering
+```shell
+python skill_discovery/agglomoration_script.py exp_name=dinov2_libero_object_image_only modality_str=dinov2_agentview_eye_in_hand repr.z_dim=1536 agglomoration.segment_scale=1 agglomoration.min_len_thresh=30 agglomoration.K=2 agglomoration.scale=0.01 agglomoration.dist=cos
+```
+Output:
+- `results/{exp_name}_{current_task_num}/skill_data/{DatasetCategoty}/{DatasetCategoty}/{DatasetName}_subtasks_{modality_str}_{feature_dim}_mean_{agglomoration.dist}_concat_1_K{agglomoration.K}_{cfg.agglomoration.affinity}.hdf5`
+
+
+
+### Save Dinov2 Feature for Hierarchical Policy Training
+```shell
+cd lotus/skill_learning
+python multisensory_repr/save_dinov2_repr.py
+```
+Output: 
+- `../datasets/dinov2/{DatasetCategoty}/{DatasetName}.hdf5`
+
+## Training
+To start a lifelong learning experiment, please choose:
+- `BENCHMARK` from `[libero_object_exp6, ... ,LIBERO_KITCHEN_EXP50]` (Please see `lotus/libero/benchmark/__init__.py` for full registered benchmark name list)
+- `EXP_NAME`: experiment name
+- `SKILL_EXP_NAME`: experiment name in `Unsupervised Skill Discovery` (e.g., `dinov2_libero_object_image_only`)
+- `PRETRAIN_MODEL_PATH`: pretrain model path
+
+For single multitask policy training from scratch, run the following:
+```shell
+export CUDA_VISIBLE_DEVICES=GPU_ID && \
+export MUJOCO_EGL_DEVICE_ID=GPU_ID && \
+python lotus/lifelong/main_old.py seed=SEED \
+                               benchmark_name=BENCHMARK \
+                               policy=bc_transformer_policy \
+                               lifelong=multitask \
+                               exp={EXP_NAME}
+```
+For single multitask policy finetuning, run the following:
+```shell
+export CUDA_VISIBLE_DEVICES=GPU_ID && \
+export MUJOCO_EGL_DEVICE_ID=GPU_ID && \
+python lotus/lifelong/main_old.py seed=SEED \
+                               benchmark_name=BENCHMARK \
+                               policy=bc_transformer_policy \
+                               lifelong=multitask \
+                               exp={EXP_NAME} \
+                               pretrain_model_path={PRETRAIN_MODEL_PATH}
+```
+
+For hierarchical skill-based policy training from scratch, run the following:
+```shell
+export CUDA_VISIBLE_DEVICES=GPU_ID && \
+export MUJOCO_EGL_DEVICE_ID=GPU_ID && \
+python lotus/lifelong/main_old.py seed=SEED \
+                               benchmark_name=BENCHMARK \
+                               policy=bc_transformer_policy \
+                               lifelong=multitask_skill \
+                               skill_learning.exp_name={SKILL_EXP_NAME} \
+                               exp={EXP_NAME} \
+                               goal_modality=BUDS
+```
+For hierarchical skill-based policy finetuning, run the following:
+```shell
+export CUDA_VISIBLE_DEVICES=GPU_ID && \
+export MUJOCO_EGL_DEVICE_ID=GPU_ID && \
+python lotus/lifelong/main_old.py seed=SEED \
+                               benchmark_name=BENCHMARK \
+                               policy=bc_transformer_policy \
+                               lifelong=multitask_skill \
+                               exp={EXP_NAME} \
+                               pretrain_model_path={PRETRAIN_MODEL_PATH}
+```
+Currently, we haven't developed automated scripts for our method in the process of lifelong learning. At present, commands need to be run individually for different stages. Please see the scripts in `scripts/training` for full training command.
+
+
+## Evaluation
+
+By default the policies will be evaluated on the fly during training. If you have limited computing resource of GPUs, we offer an evaluation script for you to evaluate models separately. Please see `lotus/lifelong/evaluate_old.py` (for single multitask policy) and `lotus/lifelong/evaluate.py` (for hierarchical skill-based policy) for more details.
+
+# Acknowledgement
+The code base used in this project is sourced from these repository:
+
+[Lifelong-Robot-Learning/LIBERO](https://github.com/Lifelong-Robot-Learning/LIBERO)
+
+[UT-Austin-RPL/BUDS](https://github.com/UT-Austin-RPL/BUDS)
+
+
+# Citation
+If you find **LOTUS** to be useful in your own research, please consider citing our paper:
+
+```bibtex
+@article{wan2024lotus,
+  title={Lotus: Continual imitation learning for robot manipulation through unsupervised skill discovery},
+  author={Wan, Weikang and Zhu, Yifeng and Shah, Rutav and Zhu, Yuke},
+  booktitle={IEEE International Conference on Robotics and Automation (ICRA)},
+  year={2024}
+}
+```

libero_benchmark_scripts/check_task_suites.py

@@ -0,0 +1,117 @@
+"""
+This script is to test if users can successfully load all the environments, the benchmark initial states in their machines
+"""
+import os
+from termcolor import colored
+import cv2
+import h5py
+import subprocess
+import shutil
+import numpy as np
+
+from pathlib import Path
+
+# import init_path
+from lotus.libero import benchmark, get_libero_path
+
+
+# def render_task(task, bddl_file, init_states, demo_file):
+#     env_args = {
+#         "bddl_file_name": bddl_file,
+#         "camera_heights": 128,
+#         "camera_widths": 128
+#     }
+
+#     env = OffScreenRenderEnv(**env_args)
+#     env.reset()
+#     obs = env.set_init_state(init_states[0])
+#     for _ in range(5):
+#         obs, _, _, _ = env.step([0.] * 7)
+#     images = [obs["agentview_image"]]
+
+#     with h5py.File(demo_file, "r") as f:
+#         states = f["data/demo_0/states"][()]
+#         obs = env.set_init_state(states[-1])
+
+#     images.append(obs["agentview_image"])
+#     images = np.concatenate(images, axis=1)
+#     cv2.imwrite(f"benchmark_tasks/{task.problem}-{task.language}.png", images[::-1, :, ::-1])
+#     env.close()
+
+
+def main():
+
+    benchmark_root_path = get_libero_path("benchmark_root")
+    init_states_default_path = get_libero_path("init_states")
+    datasets_default_path = get_libero_path("datasets")
+    bddl_files_default_path = get_libero_path("bddl_files")
+
+    # Check all the files
+    task_tuples = []
+    demo_files = []
+    for benchmark_name in [
+        "libero_object",
+        "libero_goal",
+        "libero_spatial",
+        "libero_10",
+        "libero_90",
+    ]:
+        benchmark_instance = benchmark.get_benchmark_dict()[benchmark_name]()
+        num_tasks = benchmark_instance.get_num_tasks()
+        # see how many tasks involved in the benchmark
+        print(f"{num_tasks} tasks in the benchmark {benchmark_instance.name}: ")
+
+        # Check if all the task names and their bddl file names
+        task_names = benchmark_instance.get_task_names()
+        print("The benchmark contains the following tasks:")
+        for task_id in range(num_tasks):
+            task_name = task_names[task_id]
+            task = benchmark_instance.get_task(task_id)
+            bddl_file = os.path.join(
+                bddl_files_default_path, task.problem_folder, task.bddl_file
+            )
+            assert os.path.exists(bddl_file), f"{bddl_file} does not exist!"
+            init_states_path = os.path.join(
+                init_states_default_path, task.problem_folder, task.init_states_file
+            )
+            assert os.path.exists(
+                init_states_path
+            ), f"{init_states_path} does not exist!"
+            demo_file = os.path.join(
+                datasets_default_path,
+                benchmark_instance.get_task_demonstration(task_id),
+            )
+            assert os.path.exists(demo_file), f"{demo_file} does not exist!"
+            init_states = benchmark_instance.get_task_init_states(task_id)
+            task_tuples.append((benchmark_name, task_id, bddl_file, demo_file))
+            demo_files.append(demo_file)
+
+    print(colored("All the files exist!", "green"))
+    processes = []
+    if os.path.exists("benchmark_tasks"):
+        shutil.rmtree("benchmark_tasks")
+
+    for i in range(len(task_tuples)):
+        command = f"python benchmark_scripts/render_single_task.py --benchmark_name {task_tuples[i][0]} --task_id {task_tuples[i][1]} --bddl_file {task_tuples[i][2]} --demo_file {task_tuples[i][3]}"
+        p = subprocess.Popen(command, shell=True)
+        processes.append(p)
+        if i % 10 == 9:
+            for p in processes:
+                p.wait()
+            processes = []
+
+    count = len(list(Path("benchmark_tasks").glob("*.png")))
+    print(f"Expected 130 tasks, Rendered {count} tasks successfully.")
+    if count < 130:
+        print(colored("Some tasks failed to render!", "red"))
+        for demo_file in demo_files:
+            if not os.path.exists(
+                os.path.join(
+                    "benchmark_tasks", demo_file.split("/")[-1].replace(".hdf5", ".png")
+                )
+            ):
+                print(demo_file)
+
+
+if __name__ == "__main__":
+    main()

libero_benchmark_scripts/download_libero_datasets.py

@@ -0,0 +1,46 @@
+import init_path
+import argparse
+import os
+
+import lotus.libero.utils.download_utils as download_utils
+from lotus.libero import get_libero_path
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--download-dir",
+        type=str,
+        default=get_libero_path("datasets"),
+    )
+    parser.add_argument(
+        "--datasets",
+        type=str,
+        choices=["all", "libero_goal", "libero_spatial", "libero_object", "libero_100"],
+        default="all",
+    )
+    return parser.parse_args()
+
+
+def main():
+
+    args = parse_args()
+
+    # Ask users to specify the download directory of datasets
+    os.makedirs(args.download_dir, exist_ok=True)
+    print(f"Datasets downloaded to {args.download_dir}")
+    print(f"Downloading {args.datasets} datasets")
+
+    # If not, download
+    download_utils.libero_dataset_download(
+        download_dir=args.download_dir, datasets=args.datasets
+    )
+
+    # (TODO) If datasets exist, check if datasets are the same as benchmark
+
+    # Check if datasets exist first
+    download_utils.check_libero_dataset(download_dir=args.download_dir)
+
+
+if __name__ == "__main__":
+    main()

libero_benchmark_scripts/init_path.py

@@ -0,0 +1,8 @@
+import sys
+import os
+
+path = os.path.dirname(os.path.realpath(__file__))
+sys.path.insert(0, os.path.join(path, "../"))
+
+# import robosuite.utils.macros as macros
+# macros.IMAGE_CONVENTION = "opencv"