Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,19 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/adept_models/kitchen/textures/marble1.png filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/adept_models/kitchen/textures/metal1.png filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/adept_models/kitchen/textures/tile1.png filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/adept_models/kitchen/textures/wood1.png filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/adept_models/scenes/textures/white_marble_tile.png filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/adept_models/scenes/textures/white_marble_tile2.png filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/third_party/franka/franka_panda.png filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/third_party/franka/meshes/visual/hand.stl filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/third_party/franka/meshes/visual/link0.stl filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/third_party/franka/meshes/visual/link1.stl filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/third_party/franka/meshes/visual/link2.stl filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/third_party/franka/meshes/visual/link3.stl filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/third_party/franka/meshes/visual/link4.stl filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/third_party/franka/meshes/visual/link5.stl filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/third_party/franka/meshes/visual/link6.stl filter=lfs diff=lfs merge=lfs -text
+uvd/envs/franka_kitchen/relay-policy-learning/third_party/franka/meshes/visual/link7.stl filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,169 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+.idea/
+
+*.pt
+*.pth
+*pl
+*.patch
+*used_configs
+.allenact_last_start_time_string
+*.lock
+*wandb

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 UVD
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,12 +1,163 @@
 ---
 title: UVD
-emoji: 🏢
-colorFrom: blue
-colorTo: green
-sdk: gradio
-sdk_version: 6.12.0
 app_file: app.py
-pinned: false
+sdk: gradio
+sdk_version: 6.11.0
 ---
+# Universal Visual Decomposer: <br>Long-Horizon Manipulation Made Easy
+
+<div align="center">
+
+[[Website]](https://zcczhang.github.io/UVD/)
+[[arXiv]](https://arxiv.org/abs/2310.08581)
+[[PDF]](https://zcczhang.github.io/UVD/assets/pdf/full_paper.pdf)
+[[Installation]](#Installation)
+[[Usage]](#Usage)
+[[BibTex]](#Citation)
+______________________________________________________________________
+
+
+
+https://github.com/zcczhang/UVD/assets/52727818/5555b99a-76eb-4d76-966f-787af763573a
+
+
+
+
+</div>
+
+# Installation
+
+- Follow the [instruction](https://github.com/openai/mujoco-py#install-mujoco) for installing `mujuco-py` and install the following apt packages if using Ubuntu:
+```commandline
+sudo apt install -y libosmesa6-dev libglfw3 patchelf libgl1 libglx-mesa0
+```
+- create conda env with Python==3.9
+```commandline
+conda create -n uvd python==3.9 -y && conda activate uvd
+```
+- Install any/all standalone visual foundation models from their repos separately *before* setup UVD, in case dependency conflicts, e.g.:
+<details><summary>
+<a href="https://github.com/facebookresearch/vip">VIP</a>
+</summary>
+<p>
+
+```commandline
+git clone https://github.com/facebookresearch/vip.git
+cd vip && pip install -e .
+python -c "from vip import load_vip; vip = load_vip()"
+```
+
+</p>
+</details>
+
+<details><summary>
+<a href="https://github.com/facebookresearch/r3m">R3M</a>
+</summary>
+<p>
+
+```commandline
+git clone https://github.com/facebookresearch/r3m.git
+cd r3m && pip install -e .
+python -c "from r3m import load_r3m; r3m = load_r3m('resnet50')"
+```
+
+</p>
+</details>
+
+<details><summary>
+<a href="https://github.com/penn-pal-lab/LIV">LIV (& CLIP)</a>
+</summary>
+<p>
+
+```commandline
+git clone https://github.com/penn-pal-lab/LIV.git
+cd LIV && pip install -e . && cd liv/models/clip && pip install -e .
+python -c "from liv import load_liv; liv = load_liv()"
+```
+
+</p>
+</details>
+
+
+<details><summary>
+<a href="https://github.com/facebookresearch/eai-vc">VC1</a>
+</summary>
+<p>
+
+```commandline
+git clone https://github.com/facebookresearch/eai-vc.git 
+cd eai-vc && pip install -e vc_models
+```
+
+</p>
+</details>
+
+<details><summary>
+<a href="https://github.com/facebookresearch/dinov2">DINOv2</a> and <a href="https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html">ResNet</a> pretrained with ImageNet-1k are directly loaded via <a href="https://pytorch.org/hub/">torch hub</a> and <a href="https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html">torchvision</a>.
+</summary></details>
+
+- Under *this* UVD repo directory, install other dependencies
+```commandline
+pip install -e .
+```
+
+# Usage
+
+We provide a simple API for decompose RGB videos:
+
+```python
+import torch
+import uvd
+
+# (N sub-goals, *video frame shape)
+subgoals = uvd.get_uvd_subgoals(
+    "/PATH/TO/VIDEO.*",   # video filename or (L, *video frame shape) video numpy array
+    preprocessor_name="vip",    # Literal["vip", "r3m", "liv", "clip", "vc1", "dinov2"]
+    device="cuda" if torch.cuda.is_available() else "cpu",  # device for loading frozen preprocessor
+    return_indices=False,   # True if only want the list of subgoal timesteps
+)
+```
+
+or run
+```commandline
+python demo.py
+```
+to host a Gradio demo locally with different choices of visual representations. 
+
+## Simulation Data
+
+We post-processed the data released from original [Relay-Policy-Learning](https://github.com/google-research/relay-policy-learning/tree/master) that keeps the successful trajectories only and adapt the control and observations used in our paper by:
+```commandline
+python datasets/data_gen.py raw_data_path=/PATH/TO/RAW_DATA
+```
+
+Also consider to force set `Builder = LinuxCPUExtensionBuilder` to `Builder = LinuxGPUExtensionBuilder` in `PATH/TO/CONDA/envs/uvd/lib/python3.9/site-packages/mujoco_py/builder.py` to enable (multi-)GPU acceleration.
+
+
+## Runtime Benchmark
+
+Since UVD's goal is to be an off-the-shelf method applying to *any* existing policy learning frameworks and models, across BC and RL, we provide minimal scripts for benchmarking the runtime showing negligible runtime under `./scripts` directory:
+```commandline
+python scripts/benchmark_decomp.py /PATH/TO/VIDEO
+```
+and passing `--preprocessor_name` with other preprocessors (default `vip`) and `--n` for the number of repeated iterations (default `100`).
+
+For inference or rollouts, we benchmark the runtime by 
+```commandline
+python scripts/benchmark_inference.py
+```
+and passing `--policy` for using MLP or causal GPT policy; `--preprocessor_name` with other preprocessors (default `vip`); `--use_uvd` as boolean arg for whether using UVD or no decomposition (i.e. final goal conditioned); and `--n` for the number of repeated iterations (default `100`). The default episode horizon is set to 300. We found that running in the terminal would be almost 2s slower every episode than directly running with python IDE (e.g. PyCharm, under the script directory and run as script instead of module), but the general trend that including UVD introduces negligible extra runtime still holds true. 
+
+# Citation
+If you find this project useful in your research, please consider citing:
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+```bibtex
+@inproceedings{zhang2024universal,
+  title={Universal visual decomposer: Long-horizon manipulation made easy},
+  author={Zhang, Zichen and Li, Yunshuang and Bastani, Osbert and Gupta, Abhishek and Jayaraman, Dinesh and Ma, Yecheng Jason and Weihs, Luca},
+  booktitle={2024 IEEE International Conference on Robotics and Automation (ICRA)},
+  pages={6973--6980},
+  year={2024},
+  organization={IEEE}
+}
+```

app.py

@@ -0,0 +1,34 @@
+import gradio as gr
+import torch
+import uvd
+import decord
+
+
+def proc_video(video, preprocessor_name):
+    frame_original_res = decord.VideoReader(video)[:].asnumpy()
+    indices = uvd.get_uvd_subgoals(
+        video, preprocessor_name.lower().replace("-", ""),
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        return_indices=True,
+    )
+    subgoals = frame_original_res[indices]
+    return [(img, f"No. {i+1} subgoal") for i, img in enumerate(subgoals)]
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        input_video = gr.Video(height=224, scale=3)
+        preprocessor_name = gr.Dropdown(
+            ["VIP", "R3M", "LIV", "CLIP", "DINO-v2", "VC-1", "ResNet"],
+            label="Preprocessor",
+            value="VIP",
+            scale=1,
+        )
+        output = gr.Gallery(label="UVD SubGoals", height=224, preview=True, scale=4)
+    with gr.Row():
+        submit = gr.Button("Submit")
+        clr = gr.ClearButton(components=[input_video, output])
+    submit.click(proc_video, inputs=[input_video, preprocessor_name], outputs=[output])
+
+
+demo.queue().launch(share=True, show_error=True)

datasets/.gitignore

@@ -0,0 +1,5 @@
+*
+!.gitignore
+!data_gen.py
+!generate_in_domain_vip_ft_data.py
+!data_gen.yaml

datasets/data_gen.py

@@ -0,0 +1,317 @@
+from __future__ import annotations
+
+import copy
+import glob
+import multiprocessing as mp
+import shutil
+from multiprocessing import get_logger
+from typing import Any
+
+import hydra
+import imageio
+import numpy as np
+import tqdm
+
+import uvd.utils as U
+from uvd.envs.franka_kitchen import KitchenBase
+
+
+def parse_task_name(task_name: str):
+    sub_tasks = task_name.split("-")[1].split(",")
+    parsed_tasks = []
+    for sub_task in sub_tasks:
+        if sub_task == "bottomknob":
+            parsed_tasks.append("bottom burner")
+        elif sub_task == "topknob":
+            parsed_tasks.append("top burner")
+        elif sub_task == "switch":
+            parsed_tasks.append("light switch")
+        elif sub_task == "slide":
+            parsed_tasks.append("slide cabinet")
+        elif sub_task == "hinge":
+            parsed_tasks.append("hinge cabinet")
+        else:
+            parsed_tasks.append(sub_task)
+    return "-".join(parsed_tasks)
+
+
+def _generate_one_task_demo(
+    *,
+    raw_data_path: str,
+    raw_task_dir: str,
+    # same seeds for diff tasks for now
+    seed_start: int,
+    output_path: str,
+    max_demos: int | None = None,
+    frame_height: int = 256,
+    frame_width: int = 256,
+    terminal_if_done: bool = True,
+    include_no_robot: bool = False,
+    max_reject_sampling: int = 10,
+    render: bool = False,
+    save_mp4: bool = False,
+    mp4_fps: int | None = None,
+    copy_raw_data: bool = False,
+    env: KitchenBase | None = None,
+    gpu_id,
+):
+    task_path = U.f_join(raw_data_path, raw_task_dir)
+    parsed_task_name = parse_task_name(raw_task_dir)
+    trajectories = glob.glob(U.f_join(task_path, "*.pkl"))
+    assert len(trajectories) > 0, task_path
+    # include space, connect goals with "-"
+    task_elements = parsed_task_name.split("-")
+    n_demos = min(max_demos or len(trajectories), len(trajectories))
+    saved_episode_dir = f"{parsed_task_name.replace(' ', '_')}"
+    episode_path = U.f_mkdir(output_path, saved_episode_dir)
+
+    close_env = False
+    if env is None:
+        env = KitchenBase(
+            task_elements=task_elements,
+            frame_height=frame_height,
+            frame_width=frame_width,
+            obs_keys=("rgb", "proprio"),
+            gpu_id=gpu_id,
+        )
+        close_env = True
+
+    num_demos_gen = 0
+    seed = seed_start - 1
+    traj_length_map = {}
+    unsuccessful_traj = []
+    for i, traj in tqdm.tqdm(
+        enumerate(trajectories[:n_demos]), total=n_demos, desc=parsed_task_name
+    ):
+        data = U.load_pickle(traj)
+        path = data["path"]
+        actions = path["actions"]
+        ctrls = data["ctrl"]
+        raw_data_length, action_dim = actions.shape
+        assert action_dim == env.action_space.shape[0]
+
+        data_dict, info = {}, {}  # placeholders
+        for _ in range(max_reject_sampling):
+            seed += 1
+            env.reset(init_qpos=data["qpos"][0], init_qvel=data["qvel"][0])
+            obs = env.reset(task_elements=task_elements, seed=seed)
+            obs_dict = env.obs_dict
+            # slightly difference than directly using
+            init_qpos = env.init_qpos.copy()
+            init_qvel = env.init_qvel.copy()
+            if render:
+                env.render()
+            data_dict: dict[str, Any] = {
+                "reset_kwargs": dict(
+                    task_elements=task_elements,
+                    init_qpos=init_qpos,
+                    init_qvel=init_qvel,
+                ),
+                "obs_full": [obs],  # {str: L+1, h, w, 3}
+                "actions": [],  # L, A
+                "rewards": [],  # L,
+                "dones": [],  # L,
+                "completed_tasks": [],  # L, x
+                "scores": [],  # L,
+                "seed": seed,
+                "oracle_milestones": [],  # n h, w, 3
+            }
+            if include_no_robot:
+                data_dict["no_robot_rgb"] = [
+                    env.render(mode="rgb_array", set_robot_alpha=0.0)
+                ]
+            for s in tqdm.trange(raw_data_length):
+                # Construct the action
+                ctrl = (ctrls[s] - obs_dict["qp"]) / (
+                    env.frame_skip * env.model.opt.timestep
+                )
+                act = (ctrl - env.act_mid) / env.act_amp
+                act = np.clip(act, -1.0, 1.0)
+                obs, reward, done, info = env.step(act)
+                if render:
+                    env.render()
+                obs_dict = info["obs_dict"]
+                # 1 if complete a new goal this step
+                score = info["score"]
+                # num goals achieved so far
+                completed_tasks = info["completed_tasks"]
+                if score == 1:
+                    data_dict["oracle_milestones"].append(obs["rgb"])
+
+                data_dict["obs_full"].append(obs)
+                data_dict["actions"].append(np.array(act))
+                data_dict["rewards"].append(float(reward))
+                data_dict["dones"].append(int(done))
+                data_dict["scores"].append(float(score))
+                data_dict["completed_tasks"].append(np.array(completed_tasks))
+
+                if include_no_robot:
+                    data_dict["no_robot_rgb"].append(
+                        env.render(mode="rgb_array", set_robot_alpha=0.0)
+                    )
+
+                if terminal_if_done and done:
+                    break
+            if len(env.tasks_to_complete) != 0:
+                get_logger().warning(
+                    f"reject sampling, with task left: {env.tasks_to_complete}, "
+                    f"distance left: {info['rewards']['distances_left']}"
+                )
+                continue
+            else:
+                break
+        if len(data_dict["dones"]) == 0 or data_dict["dones"][-1] != 1:
+            # raise ValueError(traj)
+            get_logger().warning(
+                f"unsuccessful data {traj}, distance left: {info['rewards']['distances_left']}"
+            )
+            debug_dir = U.f_mkdir(U.f_join(episode_path, "reject_sampling"))
+            # noinspection PyUnboundLocalVariable
+            imageio.imsave(U.f_join(debug_dir, f"{task_elements}_{i}.png"), obs["rgb"])
+            unsuccessful_traj.append(traj)
+            continue
+
+        traj_length = len(data_dict["dones"])
+        traj_length_map[i] = traj_length
+        # save accurate data
+        data_dict["obs_full"] = U.batch_observations(
+            data_dict["obs_full"], to_tensor=False
+        )
+        data_dict["actions"] = np.stack(data_dict["actions"])
+        data_dict["rewards"] = np.array(data_dict["rewards"])
+        data_dict["dones"] = np.array(data_dict["dones"])
+        data_dict["scores"] = np.array(data_dict["scores"])
+        data_dict["completed_tasks"] = np.array(data_dict["completed_tasks"])
+        data_dict["ctrls"] = ctrls[:traj_length]
+        data_dict["last_distances_to_goal"] = info["rewards"]["distances_left"]
+        data_dict["oracle_milestones"] = np.stack(data_dict["oracle_milestones"])
+        if include_no_robot:
+            data_dict["no_robot_rgb"] = np.stack(data_dict["no_robot_rgb"])
+            assert data_dict["obs_full"]["rgb"].shape == data_dict["no_robot_rgb"].shape
+
+        assert (
+            len(data_dict["obs_full"]["rgb"]) - 1
+            == len(data_dict["obs_full"]["proprio"]) - 1
+            == len(data_dict["actions"])
+            == len(data_dict["rewards"])
+            == len(data_dict["scores"])
+            == len(data_dict["completed_tasks"])
+            == len(data_dict["ctrls"])
+            == traj_length
+        )
+        assert len(data_dict["oracle_milestones"]) == len(task_elements)
+
+        num_demos_gen += 1
+        U.save_pickle(data_dict, U.f_join(episode_path, f"episode_{i}.pkl"))
+        if copy_raw_data:
+            out_raw_path = U.f_mkdir(output_path, "raw_data", saved_episode_dir)
+            shutil.copy(traj, U.f_join(out_raw_path, traj.split("/")[-1]))
+        if save_mp4:
+            video_dir = U.f_mkdir(U.f_join(episode_path, "videos"))
+            U.save_video(
+                video=data_dict["obs_full"]["rgb"],
+                fname=U.f_join(video_dir, f"episode_{i}.mp4"),
+                fps=mp4_fps,
+                compress=True,
+            )
+            if include_no_robot:
+                U.save_video(
+                    video=data_dict["no_robot_rgb"],
+                    fname=U.f_join(video_dir, f"episode_{i}_no_robot.mp4"),
+                    fps=mp4_fps,
+                    compress=True,
+                )
+    U.dump_json(
+        dict(
+            task_elements=task_elements,
+            seed_start=seed_start,
+            action_dim=env.action_space.shape[0],
+            frame_height=frame_height,
+            frame_width=frame_width,
+            num_demos=num_demos_gen,
+            trajectory_lengths=traj_length_map,
+            unsuccessful_traj=unsuccessful_traj,
+        ),
+        episode_path,
+        f"metadata.json",
+    )
+    if close_env:
+        env.close()
+
+
+def generate_one_task_demo(kwargs):
+    _generate_one_task_demo(**kwargs)
+
+
+@hydra.main(config_path=".", config_name="data_gen", version_base="1.1")
+def main(cfg):
+    if cfg.debug:
+        cfg = copy.deepcopy(cfg)
+        cfg.num_processes = 1
+        cfg.output_path += "_debug"
+        # cfg.render = True
+        cfg.max_demos = 3
+    U.ask_if_overwrite(cfg.output_path)
+    if cfg.dm_backend or (cfg.render and (not cfg.mp or cfg.num_processes == 1)):
+        import adept_envs.mujoco_env
+
+        adept_envs.mujoco_env.USE_DM_CONTROL = True
+
+    raw_data_path = cfg.raw_data_path
+    raw_task_dirs = U.f_listdir(raw_data_path)
+    assert raw_task_dirs
+    if cfg.mp:
+        num_processes = min(mp.cpu_count() - 2, cfg.num_processes or len(raw_task_dirs))
+        _, num_gpus = U.parse_gpu_devices(cfg.gpus)
+        with mp.Pool(num_processes) as pool:
+            pool.map(
+                generate_one_task_demo,
+                [
+                    dict(
+                        raw_data_path=cfg.raw_data_path,
+                        raw_task_dir=raw_task_dirs[i],
+                        seed_start=cfg.seed_start,
+                        output_path=cfg.output_path,
+                        max_demos=cfg.max_demos,
+                        frame_height=cfg.frame_height,
+                        frame_width=cfg.frame_width,
+                        terminal_if_done=cfg.terminal_if_done,
+                        max_reject_sampling=cfg.max_reject_sampling,
+                        include_no_robot=cfg.include_no_robot,
+                        render=cfg.render,
+                        save_mp4=cfg.save_mp4,
+                        mp4_fps=cfg.mp4_fps,
+                        copy_raw_data=cfg.copy_raw_data,
+                        env=None,
+                        gpu_id=i % num_gpus,
+                    )
+                    for i in range(len(raw_task_dirs))
+                ],
+            )
+    else:
+        env = KitchenBase(frame_height=cfg.frame_height, frame_width=cfg.frame_width)
+        for raw_task_dir in tqdm.tqdm(
+            raw_task_dirs, desc="Generate Franka Kitchen Demos"
+        ):
+            _generate_one_task_demo(
+                raw_data_path=cfg.raw_data_path,
+                raw_task_dir=raw_task_dir,
+                seed_start=cfg.seed_start,
+                output_path=cfg.output_path,
+                max_demos=cfg.max_demos,
+                frame_height=cfg.frame_height,
+                frame_width=cfg.frame_width,
+                terminal_if_done=cfg.terminal_if_done,
+                max_reject_sampling=cfg.max_reject_sampling,
+                include_no_robot=cfg.include_no_robot,
+                render=cfg.render,
+                save_mp4=cfg.save_mp4,
+                mp4_fps=cfg.mp4_fps,
+                copy_raw_data=cfg.copy_raw_data,
+                env=env,
+                gpu_id=-1,
+            )
+
+
+if __name__ == "__main__":
+    main()

datasets/data_gen.yaml

@@ -0,0 +1,29 @@
+raw_data_path: datasets/franka_kitchen_raw_dataset/full
+
+output_path: datasets/franka_kitchen_demos
+max_demos: null # null for all
+frame_height: 224
+frame_width: 224
+terminal_if_done: false
+max_reject_sampling: 1
+include_no_robot: false
+render: false
+save_mp4: true
+mp4_fps: null
+copy_raw_data: false
+
+debug: false
+
+seed_start: 42
+dm_backend: false
+mp: true
+num_processes: null
+gpus: null
+
+
+hydra:
+  job:
+    chdir: true
+  run:
+    dir: .
+  output_subdir: null

requirements.txt

@@ -0,0 +1,20 @@
+wheel==0.38.4
+setuptools==65.5.0
+torch==2.2.2
+torchvision==0.17.2
+cython<3
+dm_control==1.0.11
+einops==0.6.0
+gdown==4.7.1
+gym==0.25.2
+hydra-core==1.3.1
+mujoco-py<2.2,>=2.1
+numpy==1.23.5
+pytorch_lightning==2.0.0
+scikit-learn
+shimmy==0.2.0
+wandb==0.14.0
+termcolor
+decord==0.6.0
+gradio>=4.0.0
+allenact@git+https://github.com/allenai/allenact.git@usd#egg=allenact&subdirectory=allenact

scripts/auto_format.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+# Move to the directory containing the directory that this file is in
+cd "$( cd "$( dirname "${BASH_SOURCE[0]}/.." )" >/dev/null 2>&1 && pwd )" || exit
+
+echo RUNNING BLACK
+black . --exclude src --exclude external_projects
+echo BLACK DONE
+echo ""
+
+echo RUNNING DOCFORMATTER
+find . -name "*.py" | grep -v ^./src | grep -v ^./external_projects | grep -v used_configs | xargs docformatter --in-place -r
+echo DOCFORMATTER DONE
+
+echo ALL DONE

scripts/benchmark_decomp.py

@@ -0,0 +1,63 @@
+import argparse
+import warnings
+
+warnings.filterwarnings("ignore", category=UserWarning)
+
+import torch
+import tqdm
+
+import numpy as np
+import time
+
+from uvd.decomp.decomp import embedding_decomp, DEFAULT_DECOMP_KWARGS
+from uvd.models import Preprocessor, get_preprocessor
+import uvd.utils as U
+
+from decord import VideoReader
+
+
+def one_run(video_file: str, preprocessor: Preprocessor):
+    t = time.time()
+    vr = VideoReader(U.f_expand(video_file), height=224, width=224)
+    videos = vr[:].asnumpy()
+    load_v_t = time.time() - t
+
+    t = time.time()
+    embeddings = preprocessor.process(videos, return_numpy=True)
+    preprocess_t = time.time() - t
+
+    t = time.time()
+    _, decomp_meta = embedding_decomp(
+        embeddings=embeddings,
+        fill_embeddings=False,
+        return_intermediate_curves=False,
+        window_length=100,
+        **DEFAULT_DECOMP_KWARGS["embed"],
+    )
+    decomp_t = time.time() - t
+    return load_v_t, preprocess_t, decomp_t
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("video_file")
+    parser.add_argument("--preprocessor_name", default="vip")
+    parser.add_argument("--n", type=int, default=100)
+    args = parser.parse_args()
+
+    use_gpu = torch.cuda.is_available()
+    if not use_gpu:
+        print("NO GPU FOUND")
+    preprocessor = get_preprocessor(
+        args.preprocessor_name, device="cuda" if use_gpu else None
+    )
+    one_run(args.video_file, preprocessor)
+
+    benchmark_times = dict(load=[], preprocess=[], decomp=[])
+    for _ in tqdm.trange(args.n):
+        load_v_t, preprocess_t, decomp_t = one_run(args.video_file, preprocessor)
+        benchmark_times["load"].append(load_v_t)
+        benchmark_times["preprocess"].append(preprocess_t)
+        benchmark_times["decomp"].append(decomp_t)
+
+    print({k: (np.mean(v), np.std(v)) for k, v in benchmark_times.items()})

scripts/benchmark_inference.py

@@ -0,0 +1,163 @@
+import argparse
+import copy
+import time
+
+import gym
+import numpy as np
+import torch
+import yaml
+from omegaconf import DictConfig
+
+import uvd.utils as U
+from uvd.models.preprocessors import get_preprocessor
+from uvd.decomp.decomp import embedding_decomp, DEFAULT_DECOMP_KWARGS
+from uvd.envs.evaluator.inference_wrapper import InferenceWrapper
+from uvd.envs.franka_kitchen.franka_kitchen_base import KitchenBase
+
+MLP_CFG = """\
+policy:
+  _target_: uvd.models.policy.MLPPolicy
+  observation_space: ???
+  action_space: ???
+  preprocessor: ???
+  obs_encoder:
+    __target__: uvd.models.nn.MLP
+    hidden_dims: [1024, 512, 256]
+    activation: ReLU
+    normalization: false
+    input_normalization: BatchNorm1d
+    input_normalization_full_obs: false
+    proprio_output_dim: 512
+    proprio_add_layernorm: true
+    proprio_activation: Tanh
+    proprio_add_noise_eval: false
+    actor_act: Tanh
+  act_head:
+    __target__: uvd.models.distributions.DeterministicHead
+"""
+
+GPT_CFG = """\
+policy:
+  _target_: uvd.models.policy.GPTPolicy
+  observation_space: ???
+  action_space: ???
+  preprocessor: ???
+  use_kv_cache: true
+  max_seq_length: 10
+  obs_add: false
+  proprio_hidden_dim: 512
+  obs_encoder:
+    __target__: uvd.models.nn.GPT
+    use_wte: true
+  gpt_config:
+    block_size: 10
+    vocab_size: null
+    n_embd: 768
+    n_layer: 8
+    n_head: 8
+    dropout: 0.1
+    bias: false
+    use_llama_impl: true
+    position_embed: rotary
+  act_head:
+    __target__: uvd.models.distributions.DeterministicHead
+"""
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--policy", default="gpt")
+    parser.add_argument("--preprocessor_name", default="vip")
+    parser.add_argument("--use_uvd", action="store_true")
+    parser.add_argument("--n", type=int, default=100)
+    args = parser.parse_args()
+
+    use_gpu = torch.cuda.is_available()
+    if not use_gpu:
+        print("NO GPU FOUND")
+    preprocessor = get_preprocessor(
+        args.preprocessor_name, device="cuda" if use_gpu else None
+    )
+    policy_name = args.policy.lower()
+    assert policy_name in ["mlp", "gpt"]
+    is_causal = policy_name == "gpt"
+
+    env = KitchenBase(frame_height=224, frame_width=224)
+    env = InferenceWrapper(env, dummy_rtn=is_causal)
+    env.reset()
+
+    observation_space = gym.spaces.Dict(
+        rgb=gym.spaces.Box(-np.inf, np.inf, preprocessor.output_dim, np.float32),
+        proprio=gym.spaces.Box(-1, 1, (9,), np.float32),
+        milestones=gym.spaces.Box(
+            -np.inf, np.inf, (6,) + preprocessor.output_dim, np.float32
+        ),
+    )
+    action_space = env.action_space
+
+    cfg = yaml.safe_load(MLP_CFG if policy_name == "mlp" else GPT_CFG)
+    cfg = DictConfig(cfg)
+    policy = U.hydra_instantiate(
+        cfg.policy,
+        observation_space=observation_space,
+        action_space=action_space,
+        preprocessor=preprocessor,
+    )
+    policy = policy.to(preprocessor.device).eval()
+    U.debug_model_info(policy)
+    if is_causal:
+        assert policy.causal and policy.use_kv_cache
+
+    preprocessor = policy.preprocessor
+    # Or load FrankaKitchen dummy datas
+    dummy_data = np.random.random((300, 224, 224, 3)).astype(np.float32)
+    emb = preprocessor.process(dummy_data, return_numpy=True)
+    if args.use_uvd:
+        _, decomp_meta = embedding_decomp(
+            embeddings=emb,
+            fill_embeddings=False,
+            return_intermediate_curves=False,
+            **DEFAULT_DECOMP_KWARGS["embed"],
+        )
+        milestones = emb[decomp_meta.milestone_indices]  # nhw3
+    else:
+        milestones = emb[-1][None, ...]
+    env.milestones = milestones
+
+    MAX_HORIZON = 300
+    totals = []
+    for _ in range(args.n):
+        obs = env.reset()
+        if is_causal:
+            policy.reset_cache()
+
+        times = []
+        for st in range(MAX_HORIZON):
+            t = time.time()
+            obs = copy.deepcopy(obs)
+            batchify_obs = U.batch_observations([obs], device=policy.device)
+            if is_causal:
+                # B, T, ...
+                cur_milestone = env.current_milestone[None, None, ...]
+                for k in batchify_obs:
+                    batchify_obs[k] = batchify_obs[k][:, None, ...]
+            else:
+                # B, ...
+                cur_milestone = env.current_milestone[None, ...]
+            with torch.no_grad():
+                action, obs_embed, goal_embed = policy(
+                    batchify_obs,
+                    goal=torch.as_tensor(cur_milestone, device=policy.device),
+                    deterministic=True,
+                    return_embeddings=True,
+                    input_pos=torch.tensor([st], device=policy.device)
+                    if is_causal
+                    else None,
+                )
+            env.current_obs_embedding = obs_embed[0].cpu().numpy()
+            obs, r, done, info = env.step(action[0].cpu().numpy())
+            step_t = time.time() - t
+            times.append(step_t)
+        times = np.sum(times)
+        print(times)
+        totals.append(times)
+    print(np.mean(totals))

scripts/tsne_visualization.py

@@ -0,0 +1,110 @@
+import argparse
+import os.path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+import torch
+from sklearn.manifold import TSNE
+
+from uvd.decomp.decomp import (
+    embedding_decomp,
+)
+from uvd.models.preprocessors import *
+import uvd.utils as U
+
+from decord import VideoReader
+
+
+def vis_2d_tsne(embeddings: np.ndarray, labels: list):
+    tsne = TSNE(n_components=2)
+    tsne_result = tsne.fit_transform(embeddings)
+    tsne_result_df = pd.DataFrame(
+        {"tsne_1": tsne_result[:, 0], "tsne_2": tsne_result[:, 1], "label": labels}
+    )
+    fig, ax = plt.subplots(1)
+    sns.scatterplot(x="tsne_1", y="tsne_2", hue="label", data=tsne_result_df, ax=ax, s=120)
+    lim = (tsne_result.min() - 5, tsne_result.max() + 5)
+    ax.set_xlim(lim)
+    ax.set_ylim(lim)
+    ax.set_aspect("equal")
+    ax.set_title(f"{preprocessor.__class__.__name__}")
+    plt.show()
+
+
+def vis_3d_tsne(embeddings: np.ndarray, labels: list):
+    tsne = TSNE(n_components=3)
+    tsne_result = tsne.fit_transform(embeddings)
+    tsne_result_df = pd.DataFrame(
+        {
+            "tsne_1": tsne_result[:, 0],
+            "tsne_2": tsne_result[:, 1],
+            "tsne_3": tsne_result[:, 2],
+            "label": labels,
+        }
+    )
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection="3d")
+
+    palette = sns.color_palette("viridis", as_cmap=True)
+    unique_labels = tsne_result_df["label"].unique()
+    colors = palette(np.linspace(0, 1, len(unique_labels)))
+    color_dict = dict(zip(unique_labels, colors))
+
+    for label in unique_labels:
+        subset = tsne_result_df[tsne_result_df["label"] == label]
+        ax.scatter(
+            subset["tsne_1"],
+            subset["tsne_2"],
+            subset["tsne_3"],
+            c=[color_dict[label]],
+            label=label,
+            s=120,
+        )
+    ax.set_title(f"{preprocessor.__class__.__name__}")
+    plt.show()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--video_file",
+        default=U.f_join(
+         os.path.dirname(__file__), "examples/microwave-bottom_burner-light_switch-slide_cabinet.mp4"
+        )
+    )
+    parser.add_argument("--preprocessor_name", default="vip")
+    args = parser.parse_args()
+
+    use_gpu = torch.cuda.is_available()
+    if not use_gpu:
+        print("NO GPU FOUND")
+
+    frames = VideoReader(args.video_file, height=224, width=224)[:].asnumpy()
+    preprocessor = get_preprocessor(
+        args.preprocessor_name, device="cuda" if use_gpu else None
+    )
+    embeddings = preprocessor.process(frames, return_numpy=True)
+    _, decomp_meta = embedding_decomp(
+        embeddings=embeddings,
+        fill_embeddings=False,
+        return_intermediate_curves=False,
+        normalize_curve=False,
+        min_interval=20,
+        smooth_method="kernel",
+        gamma=0.1,
+    )
+    milestone_indices = decomp_meta.milestone_indices
+    milestone_rgbs = frames[milestone_indices]
+
+    labels = [
+        i
+        for i, count in enumerate(milestone_indices)
+        for _ in range(count - milestone_indices[i - 1] if i > 0 else count)
+    ]
+    labels = [labels[0]] + labels
+
+    vis_2d_tsne(embeddings, labels)
+    vis_3d_tsne(embeddings, labels)

setup.py

@@ -0,0 +1,49 @@
+import pathlib
+
+from setuptools import setup, find_packages
+
+PKG_NAME = "uvd"
+VERSION = "0.0.1"
+
+
+def _read_file(fname):
+    with pathlib.Path(fname).open() as fp:
+        return fp.read()
+
+
+def _read_install_requires():
+    with pathlib.Path("requirements.txt").open() as fp:
+        return [
+            line.strip()
+            for line in fp
+            if line.strip() and not line.startswith("#")
+        ]
+
+
+setup(
+    name=PKG_NAME,
+    version=VERSION,
+    author=f"{PKG_NAME} Developers",
+    # url='http://github.com/',
+    description="research project",
+    long_description=_read_file("README.md"),
+    long_description_content_type="text/markdown",
+    keywords=["Deep Learning", "Reinforcement Learning"],
+    license="MIT License",
+    packages=find_packages(include=f"{PKG_NAME}.*"),
+    include_package_data=True,
+    zip_safe=False,
+    entry_points={
+        "console_scripts": [
+            # 'cmd_tool=mylib.subpkg.module:main',
+        ]
+    },
+    install_requires=_read_install_requires(),
+    python_requires=">=3.9",
+    classifiers=[
+        "Development Status :: 3 - Alpha",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Environment :: Console",
+        "Programming Language :: Python :: 3",
+    ],
+)

uvd/__init__.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from typing import Literal
+
+import numpy as np
+import torch
+
+from .decomp import *
+from .models import *
+
+
+def get_uvd_subgoals(
+    frames: np.ndarray | str,
+    preprocessor_name: Literal["vip", "r3m", "liv", "clip", "vc1", "dinov2"] = "vip",
+    device: torch.device | str | None = "cuda",
+    return_indices: bool = False,
+) -> list | np.ndarray:
+    """Quick API for UVD decomposition."""
+    if isinstance(frames, str):
+        from decord import VideoReader
+
+        vr = VideoReader(frames, height=224, width=224)
+        frames = vr[:].asnumpy()
+    preprocessor = get_preprocessor(preprocessor_name, device=device)
+    rep = preprocessor.process(frames, return_numpy=True)
+    _, decomp_meta = decomp_trajectories("embed", rep)
+    indices = decomp_meta.milestone_indices
+    if return_indices:
+        return indices
+    return frames[indices]

uvd/data/__init__.py

@@ -0,0 +1,3 @@
+from .dataset_aug import *
+from .dataset_base import *
+from .franka_kitchen_datasets import *

uvd/data/dataset_aug.py

@@ -0,0 +1,310 @@
+from __future__ import annotations
+
+import random
+from typing import Literal
+
+import gym
+import numpy as np
+import torch
+
+import uvd.utils as U
+from uvd.decomp import decomp_trajectories
+from uvd.envs.franka_kitchen import KitchenBase
+from .franka_kitchen_datasets import (
+    FrankaKitchenDataset,
+)
+
+__all__ = ["DatasetWithAug", "MilestoneRandomSkipDataset"]
+
+
+class DatasetWithAug(FrankaKitchenDataset):
+    def __init__(
+        self,
+        dataset_path: str | list,
+        *,
+        specific_tasks: str | list | None,
+        obs_keys: str | tuple = ("rgb", "proprio"),
+        num_demos: int | None = None,
+        shuffle: bool = False,
+        replace_last_goal_as_last_frame: bool = True,
+        # if `preprocess`, the observation for training is directly embeddings
+        preprocess: bool = True,
+        preprocess_name: str | None = None,
+        preprocess_kwargs: dict | None = None,
+        decomp_method: Literal["oracle", "embed", "random", "equally", "near_future"]
+        | None = "embed",
+        decomp_kwargs: dict | None = None,
+        sequential: bool = True,
+        include_no_robot: bool = False,
+        del_preprocessor: bool = True,
+        device: int | str | torch.device | None = None,
+    ):
+        decomp_kwargs = decomp_kwargs or {}
+        decomp_kwargs.update(fill_embeddings=False)
+        self.idx_to_all_milestones = {}
+        self.max_num_milestones = -1
+        super().__init__(**U.prepare_locals_for_super(locals()))
+
+    def prepare_episode_data(
+        self,
+        episode_idx: int,
+        episode: str,
+        use_tensor: bool = True,
+        sequential: bool = True,
+    ):
+        self.idx_to_episode_path[episode_idx] = episode
+        data = U.load_pickle(episode)
+        actions = data["actions"]
+        episode_length = actions.shape[0]
+        full_obs_data = data["obs_full"]
+        rgb_traj = full_obs_data["rgb"][:-1, ...]  # rm goal
+        raw_rgb_no_robot = data.get("no_robot_rgb", None)
+        if self.include_no_robot:
+            assert raw_rgb_no_robot is not None
+            raw_rgb_no_robot = raw_rgb_no_robot[:-1]
+            assert raw_rgb_no_robot.shape == rgb_traj.shape
+        else:
+            raw_rgb_no_robot = None
+        assert rgb_traj.shape[0] == episode_length
+        self._idx_to_episode_length[episode_idx] = episode_length
+
+        prepared_data = dict()
+
+        if use_tensor:
+            actions = U.any_to_torch_tensor(actions, device="cpu", dtype=torch.float32)
+        else:
+            actions = U.any_to_numpy(actions, dtype="float32")
+        prepared_data["actions"] = actions
+
+        # embeddings are either fully preprocessed or transformed
+        if self.preprocess or (
+            self.decomp_method is not None and "embed" in self.decomp_method
+        ):
+            # preprocess with frozen visual backbones
+            embeddings = self.preprocessor.process(
+                raw_img_batch=rgb_traj, return_numpy=not use_tensor
+            )
+            # L, N
+            assert embeddings.shape[0] == episode_length, embeddings.shape
+        else:
+            # L, 3, H, W
+            embeddings = self.preprocessor.transform_images(rgb_traj)
+
+        obs = {}
+        for k, v in full_obs_data.items():
+            if k in self.obs_keys and k == "rgb":
+                # raw rgb, may use for rollout
+                obs[k] = rgb_traj  # v[:-1]
+            elif k in self.obs_keys:
+                obs[k] = (
+                    U.any_to_torch_tensor(
+                        v[:-1], dtype=torch.float32, device="cpu"  # self.device
+                    )
+                    if use_tensor
+                    else U.any_to_numpy(v[:-1], dtype="float32")
+                )
+
+        # obs = {k: v[:-1] for k, v in full_obs_data.items() if k in self.obs_keys}
+        prepared_data["obs"] = obs
+        rgb_no_robot = None
+        if raw_rgb_no_robot is not None:
+            if self.preprocess:
+                rgb_no_robot = self.preprocessor.process(
+                    raw_img_batch=raw_rgb_no_robot,
+                    return_numpy=not use_tensor,
+                    reconstruct_linear=True,
+                )
+            else:
+                rgb_no_robot = self.preprocessor.transform_images(raw_rgb_no_robot)
+            prepared_data["rgb_no_robot"] = rgb_no_robot
+
+        if self.decomp_method == "oracle":
+            num_goals_achieved = data["completed_tasks"]
+            _, decomp_meta = decomp_trajectories(
+                method_name=self.decomp_method,
+                embeddings=None,
+                goal_achieved_mask=num_goals_achieved,
+                **self.decomp_kwargs or {},
+            )
+        else:
+            decomp_kwargs = self.decomp_kwargs or {}
+            if (
+                self.decomp_method is not None
+                and "embed_no_robot" in self.decomp_method
+            ):
+                decomp_kwargs["no_robot_embeddings"] = rgb_no_robot
+            elif self.decomp_method == "embed2":
+                decomp_kwargs["no_robot_embeddings"] = embeddings
+                decomp_kwargs["task_name"] = str(data["reset_kwargs"]["task_elements"])
+            _, decomp_meta = decomp_trajectories(
+                method_name=self.decomp_method,
+                embeddings=embeddings,
+                **decomp_kwargs,
+            )
+
+        if self.preprocessor.use_language_goal:
+            raise NotImplementedError
+
+        if use_tensor:
+            embeddings = U.any_to_torch_tensor(embeddings, device="cpu")
+
+        milestone_indices = decomp_meta.milestone_indices
+        prepared_data["milestone_indices"] = milestone_indices
+        self.max_num_milestones = max(len(milestone_indices), self.max_num_milestones)
+
+        diffs = np.diff(milestone_indices, prepend=0)
+        milestone_step_mask = np.repeat(np.arange(len(diffs)), diffs)
+        milestone_step_mask = np.concatenate(
+            [np.array([0]), milestone_step_mask], dtype=np.int32
+        )
+        assert len(milestone_step_mask) == episode_length, (
+            len(milestone_step_mask),
+            episode_length,
+        )
+
+        milestone_embeddings = embeddings[milestone_indices]
+
+        if sequential:
+            if self.split_train_eval:
+                if episode.split("/")[-2] not in self.train_tasks:
+                    assert episode.split("/")[-2] in self.eval_tasks, episode.split(
+                        "/"
+                    )[-2]
+                    return prepared_data, episode_length
+
+            def maybe_tensor(x, dtype):
+                return (
+                    torch.tensor(x, device="cpu", dtype=dtype)
+                    if use_tensor
+                    else np.array(x, dtype=dtype)
+                )
+
+            for st in range(episode_length):
+                seq_data = dict(
+                    actions=actions[st],
+                    milestones=milestone_embeddings,
+                    milestone_indices=maybe_tensor(milestone_indices, dtype="int32"),
+                    cur_milestone_idx=maybe_tensor(
+                        milestone_step_mask[st], dtype="int32"
+                    ),
+                    episode_idx=maybe_tensor([episode_idx], dtype="int32"),
+                    timesteps=maybe_tensor([st], dtype="int32"),
+                )
+                obs_i = {
+                    k: v[st] if k != "rgb" else embeddings[st] for k, v in obs.items()
+                }
+                seq_data["obs"] = obs_i
+                self.sequence_data.append(seq_data)
+        else:
+            raise NotImplementedError
+
+        return prepared_data, episode_length
+
+    @property
+    def dataset_metadata(self) -> dict:
+        assert self.max_num_milestones > 0, self.max_num_milestones
+
+        dummy_env = KitchenBase(obs_keys=self.obs_keys)
+        action_space = dummy_env.action_space
+        observation_space = dummy_env.observation_space.spaces
+        dummy_env.close()
+        del dummy_env
+        if "rgb" in observation_space and self.preprocess:
+            observation_space["rgb"] = gym.spaces.Box(
+                low=-np.inf,
+                high=np.inf,
+                shape=self.preprocessor_output_dim,
+                dtype=np.float32,
+            )
+        observation_space = gym.spaces.Dict(observation_space)
+        metadata = dict(action_space=action_space, observation_space=observation_space)
+        metadata["observation_space"]["milestones"] = gym.spaces.Box(
+            low=-np.inf,
+            high=np.inf,
+            shape=(self.max_num_milestones,) + self.preprocessor_output_dim,
+            dtype=np.float32,
+        )
+        return metadata
+
+    def __getitem__(self, idx: int) -> dict:
+        data = self.sequence_data[idx]
+        milestones = data["milestones"]
+        milestone_indices = data["milestone_indices"]
+        num_milestones, d = milestones.shape
+        assert num_milestones <= self.max_num_milestones
+        pad_length = self.max_num_milestones - num_milestones
+        if pad_length != 0:
+            milestones = U.any_concat(
+                # assume not use tensor
+                [milestones, np.zeros((pad_length, d), dtype=milestones.dtype)],
+                dim=0,
+            )
+            milestone_indices = U.any_concat(
+                [milestone_indices, milestone_indices[-1][None] * pad_length]
+            )
+        # N, D
+        data["milestones"] = milestones  # .reshape((self.max_num_milestones * d,))
+        data["milestone_masks"] = np.array(
+            [1] * num_milestones + [0] * pad_length, dtype=np.bool_
+        )[
+            :, None
+        ]  # N, 1
+        data["milestone_indices"] = milestone_indices
+        return data
+
+
+class MilestoneRandomSkipDataset(DatasetWithAug):
+    def __init__(
+        self, *, skip_ratio: 0.1, min_skip_n_milestones: int | None = None, **kwargs
+    ):
+        self.skip_ratio = skip_ratio
+        self.min_skip_n_milestones = min_skip_n_milestones or 0
+        super().__init__(**kwargs)
+
+    def __getitem__(self, idx: int) -> dict:
+        data = self.sequence_data[idx]
+        milestones = data["milestones"]
+        milestone_indices = data["milestone_indices"].copy()
+        num_milestones = len(milestone_indices)
+        assert len(milestones) == num_milestones, (
+            len(milestones),
+            num_milestones,
+            milestone_indices,
+        )
+        cur_milestone_idx = data["cur_milestone_idx"].copy()
+
+        do_skip = (
+            len(milestone_indices) > self.min_skip_n_milestones
+            and random.random() > self.skip_ratio
+        )
+        if do_skip:
+            skip_back = (
+                cur_milestone_idx > 0 and random.random() > 0.5
+            ) or cur_milestone_idx == num_milestones - 1
+            if skip_back:
+                cur_milestone_idx -= 1
+            else:
+                cur_milestone_idx += 1
+            assert 0 <= cur_milestone_idx <= num_milestones - 1, (
+                cur_milestone_idx,
+                milestone_indices,
+            )
+
+        pad_length = self.max_num_milestones - num_milestones
+        if pad_length != 0:
+            milestone_indices = U.any_concat(
+                [data["milestone_indices"]]
+                + [data["milestone_indices"][-1][None]] * pad_length
+            )
+        else:
+            milestone_indices = data["milestone_indices"]
+
+        return dict(
+            obs=data["obs"],
+            milestones=milestones[cur_milestone_idx],
+            milestone_indices=milestone_indices,
+            cur_milestone_idx=cur_milestone_idx,
+            actions=data["actions"],
+            episode_idx=data["episode_idx"],
+            timesteps=data["timesteps"],
+        )

uvd/data/dataset_base.py

@@ -0,0 +1,19 @@
+import abc
+
+from torch.utils.data import Dataset
+
+__all__ = ["DatasetBase"]
+
+
+class DatasetBase(Dataset):
+    @abc.abstractproperty
+    def dataset_metadata(self) -> dict:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def __len__(self):
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def __getitem__(self, item: int):
+        raise NotImplementedError

uvd/data/franka_kitchen_datasets.py

@@ -0,0 +1,707 @@
+from __future__ import annotations
+
+import collections
+import copy
+import json
+import random
+from collections import OrderedDict
+from typing import Literal
+
+import gym
+import numpy as np
+import torch
+import tqdm
+import wandb
+from torch.distributions.utils import lazy_property
+from torch.utils.data import Dataset
+
+import uvd.utils as U
+from uvd.decomp import decomp_trajectories
+from uvd.envs.franka_kitchen import KitchenBase
+from uvd.envs.franka_kitchen.franka_kitchen_constants import ELEMENT_TO_IDX
+from uvd.models.preprocessors import get_preprocessor
+from .dataset_base import DatasetBase
+
+__all__ = [
+    "FrankaKitchenDataset",
+    "FrankaKitchenRolloutDataset",
+    "ALL_TASKS",
+    "task_elements_to_prompt",
+]
+
+
+ALL_TASKS = [
+    "bottom_burner-top_burner-light_switch-slide_cabinet",
+    "bottom_burner-top_burner-slide_cabinet-hinge_cabinet",
+    "kettle-bottom_burner-light_switch-hinge_cabinet",
+    "kettle-bottom_burner-light_switch-slide_cabinet",
+    "kettle-bottom_burner-slide_cabinet-hinge_cabinet",
+    "kettle-bottom_burner-top_burner-hinge_cabinet",
+    "kettle-bottom_burner-top_burner-light_switch",
+    "kettle-bottom_burner-top_burner-slide_cabinet",
+    "kettle-light_switch-slide_cabinet-hinge_cabinet",
+    "kettle-top_burner-light_switch-slide_cabinet",
+    "microwave-bottom_burner-light_switch-slide_cabinet",
+    "microwave-bottom_burner-slide_cabinet-hinge_cabinet",
+    "microwave-bottom_burner-top_burner-hinge_cabinet",
+    "microwave-bottom_burner-top_burner-light_switch",
+    "microwave-bottom_burner-top_burner-slide_cabinet",
+    "microwave-kettle-bottom_burner-hinge_cabinet",
+    "microwave-kettle-bottom_burner-slide_cabinet",
+    "microwave-kettle-light_switch-hinge_cabinet",
+    "microwave-kettle-light_switch-slide_cabinet",
+    "microwave-kettle-slide_cabinet-hinge_cabinet",
+    "microwave-kettle-top_burner-hinge_cabinet",
+    "microwave-kettle-top_burner-light_switch",
+    "microwave-light_switch-slide_cabinet-hinge_cabinet",
+    "microwave-top_burner-light_switch-hinge_cabinet",
+]
+
+
+class FrankaKitchenDataset(DatasetBase):
+    def __init__(
+        self,
+        dataset_path: str | list,
+        *,
+        specific_tasks: str | list | None,
+        obs_keys: str | tuple = ("rgb", "proprio"),
+        num_demos: int | None = None,
+        shuffle: bool = False,
+        replace_last_goal_as_last_frame: bool = True,
+        # if `preprocess`, the observation for training is directly embeddings
+        preprocess: bool = True,
+        preprocess_name: str | None = None,
+        preprocess_kwargs: dict | None = None,
+        decomp_method: Literal["oracle", "embed", "random", "equally", "near_future"]
+        | None = "embed",
+        decomp_kwargs: dict | None = None,
+        sequential: bool = True,
+        max_seq_length: int | None = None,
+        include_no_robot: bool = False,
+        del_preprocessor: bool = True,
+        device: int | str | torch.device | None = None,
+    ):
+        super().__init__()
+
+        self.split_train_eval = None
+        if "random_" in specific_tasks:
+            n_train = int(specific_tasks.split("_")[-1])
+            specific_tasks = ALL_TASKS
+            random.shuffle(specific_tasks)
+            self.split_train_eval = n_train
+
+        if isinstance(dataset_path, str):
+            assert U.f_exists(dataset_path), dataset_path
+            if specific_tasks is not None:
+                if isinstance(specific_tasks, str):
+                    specific_tasks = (specific_tasks,)
+                dataset_path = [U.f_join(dataset_path, t) for t in specific_tasks]
+            elif "metadata.json" in U.f_listdir(dataset_path):
+                # single task
+                dataset_path = (dataset_path,)
+            else:
+                # use all tasks
+                dataset_path = [
+                    U.f_join(dataset_path, d)
+                    for d in U.f_listdir(dataset_path)
+                    if "-" in d and "_" in d
+                ]
+
+        if self.split_train_eval is not None:
+            assert specific_tasks is not None
+            self.train_tasks = specific_tasks[: self.split_train_eval]
+            self.eval_tasks = specific_tasks[self.split_train_eval :]
+            U.rank_zero_print(
+                f"{len(self.train_tasks)} training tasks, {len(self.eval_tasks)} unseen tasks",
+                color="green",
+            )
+            if U.is_rank_zero() and wandb.run is not None:
+                table = wandb.Table(
+                    data=[[",\n".join(self.train_tasks), ",\n".join(self.eval_tasks)]],
+                    columns=["training tasks", "unseen tasks"],
+                )
+                wandb.log(dict(partition=table))
+
+        assert sum((U.f_exists(d) for d in dataset_path)) == len(dataset_path)
+        self.include_no_robot = include_no_robot and decomp_method is not None
+        if decomp_method == "embed_no_robot":
+            assert self.include_no_robot and preprocess
+        self.cache_scores = dict()
+
+        episode_for_task = collections.defaultdict(list)
+        for d in dataset_path:
+            for f in U.f_listdir(d):
+                if num_demos is not None and len(episode_for_task[d]) >= num_demos:
+                    break
+                if f.endswith(".pkl") and f.startswith("episode"):
+                    episode_for_task[d].append(U.f_join(d, f))
+
+        episodes = sorted([f for v in episode_for_task.values() for f in v])
+
+        U.rank_zero_print(
+            "# Train Demos:\n"
+            + json.dumps(
+                {
+                    k.split("/")[-1]: len(v)
+                    for k, v in episode_for_task.items()
+                    if k.split("/")[-1] in self.train_tasks
+                },
+                indent=2,
+            )
+            + "\n# Eval Demos:\n"
+            + json.dumps(
+                {
+                    k.split("/")[-1]: len(v)
+                    for k, v in episode_for_task.items()
+                    if k.split("/")[-1] in self.eval_tasks
+                },
+                indent=2,
+            )
+            + f"\nTotal: {len(episodes)}",
+            color="blue",
+        )
+
+        if shuffle:
+            random.shuffle(episodes)
+
+        self.obs_keys = (
+            tuple(obs_keys) if not isinstance(obs_keys, str) else (obs_keys,)
+        )
+        self.replace_last_goal_as_last_frame = replace_last_goal_as_last_frame
+        self.num_demos = num_demos
+        self.idx_to_episode_path = {}
+        self._idx_to_episode_length = {}
+
+        self.preprocess = preprocess
+        self.device = device
+        if preprocess or decomp_method == "embed":
+            assert preprocess_name is not None
+            self.preprocessor = get_preprocessor(
+                name=preprocess_name, device=device, **preprocess_kwargs or {}
+            )
+            self.preprocessor_output_dim = self.preprocessor.output_dim
+        else:
+            # do transform only, with self.preprocess is False
+            self.preprocessor = get_preprocessor(
+                name=preprocess_name, device=device, **preprocess_kwargs or {}
+            )
+        self.decomp_method = decomp_method
+        self.decomp_kwargs = decomp_kwargs
+        self.use_language_goal = self.preprocessor.use_language_goal
+
+        self.episode_data = {}
+        self.sequence_data = []
+        self.sequential = sequential
+        for episode_idx, episode in tqdm.tqdm(
+            enumerate(episodes),
+            total=len(episodes),
+            desc=f"load {self.__class__.__name__}",
+        ):
+            prepared_data, episode_length = self.prepare_episode_data(
+                episode_idx, episode, sequential=sequential, use_tensor=False  # True
+            )
+            self.episode_data[episode_idx] = prepared_data
+
+        assert len(self._idx_to_episode_length) == len(episodes)
+        self._max_episode_length = max(self._idx_to_episode_length.values())
+        self.max_seq_length = None
+        self.idx_partition = None
+        if max_seq_length is not None:
+            assert max_seq_length <= self._max_episode_length
+            if max_seq_length < self._max_episode_length:
+                self.max_seq_length = max_seq_length
+                self.idx_partition = [
+                    (ep_idx, (max(0, step_end - max_seq_length), step_end))
+                    for ep_idx in range(
+                        len(self.sequence_data)
+                    )  # only for training data
+                    for step_end in range(
+                        1,
+                        self._idx_to_episode_length[
+                            int(self.sequence_data[ep_idx]["episode_idx"])
+                        ]
+                        + 1,
+                    )
+                ]
+
+        if del_preprocessor:
+            # del after using
+            del self.preprocessor
+
+        U.rank_zero_print(
+            f"TOTAL TRAINING DATA: {len(self.sequence_data)}", color="blue"
+        )
+
+    @property
+    def max_episode_length(self) -> int:
+        return self._max_episode_length
+
+    @lazy_property
+    def dataset_metadata(self) -> dict:
+        dummy_env = KitchenBase(obs_keys=self.obs_keys)
+        action_space = dummy_env.action_space
+        observation_space = dummy_env.observation_space.spaces
+        dummy_env.close()
+        del dummy_env
+        if "rgb" in observation_space and self.preprocess:
+            observation_space["rgb"] = gym.spaces.Box(
+                low=-np.inf,
+                high=np.inf,
+                shape=(self.preprocessor_output_dim,)
+                if isinstance(self.preprocessor_output_dim, int)
+                else self.preprocessor_output_dim,
+                dtype=np.float32,
+            )
+        observation_space = gym.spaces.Dict(observation_space)
+        metadata = dict(action_space=action_space, observation_space=observation_space)
+        if not self.sequential:  # causal
+            metadata["max_episode_length"] = self.max_episode_length
+            if self.max_seq_length is not None:
+                metadata["max_seq_length"] = self.max_seq_length
+        return metadata
+
+    def prepare_episode_data(
+        self,
+        episode_idx: int,
+        episode: str,
+        use_tensor: bool = True,
+        sequential: bool = True,
+    ):
+        self.idx_to_episode_path[episode_idx] = episode
+        data = U.load_pickle(episode)
+        actions = data["actions"]
+        episode_length = actions.shape[0]
+        full_obs_data = data["obs_full"]
+        rgb_traj = full_obs_data["rgb"][:-1, ...]  # rm goal
+        task_elements = data["reset_kwargs"]["task_elements"]
+        raw_rgb_no_robot = data.get("no_robot_rgb", None)
+        if self.include_no_robot:
+            assert raw_rgb_no_robot is not None
+            raw_rgb_no_robot = raw_rgb_no_robot[:-1]
+            assert raw_rgb_no_robot.shape == rgb_traj.shape
+        else:
+            raw_rgb_no_robot = None
+        assert rgb_traj.shape[0] == episode_length
+        self._idx_to_episode_length[episode_idx] = episode_length
+
+        prepared_data = dict()
+
+        if use_tensor:
+            actions = U.any_to_torch_tensor(actions, device="cpu", dtype=torch.float32)
+        else:
+            actions = U.any_to_numpy(actions, dtype="float32")
+        prepared_data["actions"] = actions
+
+        # embeddings are either fully preprocessed or transformed
+        if self.preprocess or (
+            self.decomp_method is not None and "embed" in self.decomp_method
+        ):
+            # preprocess with frozen visual backbones
+            embeddings = self.preprocessor.process(
+                raw_img_batch=rgb_traj, return_numpy=not use_tensor
+            )
+            # L, N
+            assert embeddings.shape[0] == episode_length, embeddings.shape
+        else:
+            # L, 3, H, W
+            embeddings = self.preprocessor.transform_images(rgb_traj)
+
+        obs = {}
+        for k, v in full_obs_data.items():
+            if k in self.obs_keys and k == "rgb":
+                # raw rgb, may use for rollout
+                obs[k] = rgb_traj  # v[:-1]
+            elif k in self.obs_keys:
+                obs[k] = (
+                    U.any_to_torch_tensor(
+                        v[:-1], dtype=torch.float32, device="cpu"  # self.device
+                    )
+                    if use_tensor
+                    else U.any_to_numpy(v[:-1], dtype="float32")
+                )
+
+        # obs = {k: v[:-1] for k, v in full_obs_data.items() if k in self.obs_keys}
+        prepared_data["obs"] = obs
+        rgb_no_robot = None
+        if raw_rgb_no_robot is not None:
+            if self.preprocess:
+                rgb_no_robot = self.preprocessor.process(
+                    raw_img_batch=raw_rgb_no_robot,
+                    return_numpy=not use_tensor,
+                    reconstruct_linear=True,
+                )
+            else:
+                rgb_no_robot = self.preprocessor.transform_images(raw_rgb_no_robot)
+            prepared_data["rgb_no_robot"] = rgb_no_robot
+
+        if self.decomp_method == "oracle":
+            oracle_milestones = data["oracle_milestones"]
+            if self.replace_last_goal_as_last_frame:
+                oracle_milestones[-1, ...] = rgb_traj[-1]  # full_obs_data["rgb"][-1]
+            if self.preprocess:
+                oracle_milestone_embeddings = self.preprocessor.process(
+                    raw_img_batch=oracle_milestones, return_numpy=not use_tensor
+                )
+            else:
+                oracle_milestone_embeddings = self.preprocessor.transform_images(
+                    oracle_milestones
+                )
+            assert len(oracle_milestone_embeddings) == len(oracle_milestones)
+            num_goals_achieved = data["completed_tasks"]
+            milestone_embeddings, decomp_meta = decomp_trajectories(
+                method_name=self.decomp_method,
+                embeddings=oracle_milestone_embeddings,
+                goal_achieved_mask=num_goals_achieved,
+                **self.decomp_kwargs or {},
+            )
+        else:
+            decomp_kwargs = self.decomp_kwargs or {}
+            if (
+                self.decomp_method is not None
+                and "embed_no_robot" in self.decomp_method
+            ):
+                decomp_kwargs["no_robot_embeddings"] = rgb_no_robot
+            elif self.decomp_method == "embed2":
+                decomp_kwargs["no_robot_embeddings"] = embeddings
+                decomp_kwargs["task_name"] = str(data["reset_kwargs"]["task_elements"])
+            milestone_embeddings, decomp_meta = decomp_trajectories(
+                method_name=self.decomp_method,
+                embeddings=embeddings,
+                **decomp_kwargs,
+            )
+
+        if self.preprocessor.use_language_goal:
+            assert self.preprocess
+            prompt = task_elements_to_prompt(task_elements)
+            prompt_embed = self.preprocessor.encode_text(prompt)
+            milestone_embeddings = prompt_embed.repeat(episode_length, 1)
+            if not use_tensor:
+                milestone_embeddings = U.any_to_numpy(milestone_embeddings)
+            prepared_data["lang_embed"] = milestone_embeddings[-1]
+
+        if use_tensor:
+            embeddings = U.any_to_torch_tensor(embeddings, device="cpu")
+            milestone_embeddings = U.any_to_torch_tensor(
+                milestone_embeddings, device="cpu"
+            )
+        assert (
+            milestone_embeddings.shape == embeddings.shape
+        ), f"{milestone_embeddings.shape} != {embeddings.shape}"
+        prepared_data["milestone_indices"] = decomp_meta.milestone_indices
+
+        def maybe_tensor(x, dtype):
+            return (
+                torch.tensor(x, device="cpu", dtype=dtype)
+                if use_tensor
+                else np.array(x, dtype=dtype)
+            )
+
+        if sequential:
+            if self.split_train_eval:
+                if episode.split("/")[-2] not in self.train_tasks:
+                    assert episode.split("/")[-2] in self.eval_tasks, episode.split(
+                        "/"
+                    )[-2]
+                    return prepared_data, episode_length
+            for st in range(episode_length):
+                seq_data = dict(
+                    actions=actions[st],
+                    milestones=milestone_embeddings[st],
+                    episode_idx=maybe_tensor([episode_idx], dtype="int32"),
+                    timesteps=maybe_tensor([st], dtype="int32"),
+                )
+                obs_i = {
+                    k: v[st] if k != "rgb" else embeddings[st] for k, v in obs.items()
+                }
+                seq_data["obs"] = obs_i
+                self.sequence_data.append(seq_data)
+        else:
+            if self.split_train_eval:
+                if episode.split("/")[-2] not in self.train_tasks:
+                    assert episode.split("/")[-2] in self.eval_tasks, episode.split(
+                        "/"
+                    )[-2]
+                    return prepared_data, episode_length
+            self.sequence_data.append(
+                dict(
+                    actions=actions,
+                    obs={k: v if k != "rgb" else embeddings for k, v in obs.items()},
+                    milestones=milestone_embeddings,
+                    episode_idx=maybe_tensor([episode_idx], dtype="int32"),
+                    timesteps=maybe_tensor(list(range(episode_length)), dtype="int32"),
+                )
+            )
+        return prepared_data, episode_length
+
+    def __len__(self) -> int:
+        if self.max_seq_length is None:
+            return len(self.sequence_data)
+        else:
+            return len(self.idx_partition)
+
+    def __getitem__(self, idx: int) -> dict:
+        """Dict(obs, action, milestone, episode_idx, (maybe embedding))"""
+        if self.sequential:
+            return self.sequence_data[idx]
+
+        if self.max_seq_length is not None:
+            episode_idx, (start_idx, end_idx) = self.idx_partition[idx]
+            data = copy.deepcopy(self.sequence_data[episode_idx])
+            pad_len = self.max_seq_length - (end_idx - start_idx)
+            target_mask = np.array(
+                [1] * (end_idx - 1 - start_idx) + [0] * pad_len, dtype=np.int32
+            )
+            for k in data.keys():
+                val = data[k]
+                if k == "obs":
+                    for obs_k, v in val.items():
+                        if pad_len > 0:
+                            data[k][obs_k] = U.any_concat(
+                                [v[start_idx:end_idx, ...]]
+                                + [U.any_zeros_like(v[0])[None]] * pad_len
+                            )
+                        else:
+                            data[k][obs_k] = v[start_idx:end_idx, ...]
+                elif k != "episode_idx":
+                    if pad_len > 0:
+                        data[k] = U.any_concat(
+                            [val[start_idx:end_idx, ...]]
+                            + [U.any_zeros_like(val[0])[None]] * pad_len
+                        )
+                    else:
+                        data[k] = val[start_idx:end_idx]
+            data["target_mask"] = target_mask
+            return data
+
+        data = self.sequence_data[idx]  # seems no need to deepcopy st only pad once
+        actions = data["actions"]
+        ep_len = len(actions)
+        pad_len = self.max_episode_length - ep_len
+        if torch.is_tensor(actions):
+            target_mask = torch.tensor(
+                [1] * ep_len + [0] * pad_len, dtype=torch.int32, device=actions.device
+            )
+        else:
+            target_mask = np.array([1] * ep_len + [0] * pad_len, dtype=np.int32)
+        if pad_len > 0:
+            for k in data.keys():
+                val = data[k]
+                if k == "obs":
+                    for obs_k, v in val.items():
+                        if len(v) == ep_len:
+                            data[k][obs_k] = U.any_concat(
+                                [v] + [U.any_zeros_like(v[0])[None]] * pad_len
+                            )
+                elif len(val) == ep_len:
+                    data[k] = U.any_concat(
+                        [val] + [U.any_zeros_like(val[0])[None]] * pad_len
+                    )
+        data["target_mask"] = target_mask
+        return data
+
+
+class FrankaKitchenRolloutDataset(Dataset):
+    """Reuse training dataset for in-domain rollout evaluation."""
+
+    def __init__(
+        self,
+        train_dataset: FrankaKitchenDataset,
+        specific_tasks: str | list | None = None,
+        # has to be specific tasks
+        num_episodes_each_task: int | None = None,
+        num_tasks: int | None = None,
+        use_no_robot_rgb_milestones: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        assert "rgb" in train_dataset.obs_keys
+
+        self.is_hybrid = train_dataset.decomp_method == "embed_no_robot_extended"
+
+        self.max_num_milestones = -1
+        episode_data = copy.deepcopy(train_dataset.episode_data)
+        del train_dataset.episode_data
+        self.dataset_metadata = train_dataset.dataset_metadata
+
+        self._idx_to_episode_path = train_dataset.idx_to_episode_path
+        if specific_tasks is not None:
+            if isinstance(specific_tasks, str):
+                if "unseen" in specific_tasks:
+                    specific_tasks = train_dataset.eval_tasks
+                elif "all" in specific_tasks:
+                    specific_tasks = ALL_TASKS
+                else:
+                    specific_tasks = (specific_tasks,)
+            # assume all data in the same parent directory
+            ds_base_path = "/".join(self._idx_to_episode_path[0].split("/")[:-2])
+            spec_episode_paths = [U.f_join(ds_base_path, t) for t in specific_tasks]
+            assert sum([U.f_exists(p) for p in spec_episode_paths]) == len(
+                spec_episode_paths
+            ), spec_episode_paths
+
+            task_counter = {k: 1 for k in specific_tasks}
+            _idx_to_episode_path = {}
+            _n_episode_each_task = num_episodes_each_task or np.inf
+            for i, p in self._idx_to_episode_path.items():
+                task_name = p.split("/")[-2]
+                if (
+                    task_name not in specific_tasks
+                    or task_counter[task_name] > _n_episode_each_task
+                ):
+                    if task_name not in specific_tasks:
+                        U.rank_zero_print(
+                            f"WARNING: MISSING {task_name} in {specific_tasks}",
+                            color="red",
+                        )
+                    continue
+                task_counter[task_name] += 1
+                _idx_to_episode_path[i] = p
+            self._idx_to_episode_path = _idx_to_episode_path
+        use_no_robot_rgb_milestones = (
+            use_no_robot_rgb_milestones and train_dataset.decomp_method is not None
+        )
+        self.use_no_robot_rgb_milestones = use_no_robot_rgb_milestones
+
+        self.num_demos = len(self._idx_to_episode_path)
+        self.num_tasks = min(self.num_demos, num_tasks or self.num_demos)
+
+        self.episode_data = {}
+        idx = -1
+        for ds_idx in tqdm.tqdm(
+            episode_data.keys(),
+            desc=f"load {self.__class__.__name__}",
+            total=len(self._idx_to_episode_path),
+        ):
+            if ds_idx not in self._idx_to_episode_path:
+                continue
+            # else:
+            #     print(ds_idx)
+            idx += 1
+            prepared_data = episode_data[ds_idx]
+            milestone_indices = prepared_data["milestone_indices"]
+            rgb = prepared_data["obs"]["rgb"]
+            rgb = (
+                U.any_to_torch_tensor(rgb, device="cpu", copy=True)
+                if torch.is_tensor(rgb)
+                else rgb.copy()
+            )
+
+            data = U.load_pickle(self._idx_to_episode_path[ds_idx])
+            reset_kwargs = {
+                k: U.any_to_numpy(v, dtype="float32")[None]
+                if k != "task_elements"
+                else np.array([ELEMENT_TO_IDX[ele] for ele in v], dtype=np.uint8)[None]
+                for k, v in data["reset_kwargs"].items()
+            }
+
+            if train_dataset.use_language_goal:
+                milestones = prepared_data["lang_embed"][None]  # (1, embed_d)
+            else:
+                milestones = rgb[milestone_indices]  # use raw rgb anyway
+            if torch.is_tensor(milestones):
+                milestones = U.any_to_torch_tensor(milestones, device="cpu")
+            if rgb.ndim != milestones.ndim:
+                rgb_milestones = U.any_to_numpy(rgb)[milestone_indices]
+            else:
+                rgb_milestones = None
+            rgb_no_robot = prepared_data.get("rgb_no_robot", None)
+
+            self.max_num_milestones = max(len(milestones), self.max_num_milestones)
+            if use_no_robot_rgb_milestones and rgb_no_robot is None:
+                rgb_no_robot = data["no_robot_rgb"][:-1]
+
+            rgb_no_robot_milestones = (
+                rgb_no_robot[milestone_indices] if use_no_robot_rgb_milestones else None
+            )
+            if torch.is_tensor(rgb_no_robot_milestones):
+                rgb_no_robot_milestones = U.any_to_torch_tensor(
+                    rgb_no_robot_milestones, device="cpu"
+                )
+            if torch.is_tensor(milestones):
+                milestone_indices = torch.tensor(
+                    milestones, dtype=torch.int32, device=milestones.device
+                )
+            else:
+                milestone_indices = np.array(milestone_indices, dtype=np.int32)
+            self.episode_data[idx] = dict(
+                milestones=milestones,
+                reset_kwargs=reset_kwargs,
+                rgb_milestones=rgb_milestones,
+                rgb_no_robot_milestones=rgb_no_robot_milestones,
+                milestone_indices=milestone_indices,
+            )
+
+        del episode_data
+        if self.num_tasks != self.num_demos:
+            self.episode_data = [
+                self.episode_data[i % len(self.episode_data)]
+                for i in range(self.num_tasks)
+            ]
+
+    def __len__(self) -> int:
+        """Num tasks evaluated."""
+        return self.num_tasks
+
+    def __getitem__(
+        self, idx: int
+    ) -> OrderedDict[str, np.ndarray | dict[str, np.ndarray]]:
+        data = self.episode_data[idx]
+        milestones = data["milestones"]
+        milestone_indices = data.get("milestone_indices")
+        rgb_milestones = data["rgb_milestones"]
+        rgb_no_robot_milestones = data["rgb_no_robot_milestones"]
+        # assert milestones.shape[0] == rgb_milestones.shape[0]
+        num_milestones = len(milestones)
+        assert num_milestones <= self.max_num_milestones
+        if num_milestones < self.max_num_milestones:
+            # pad last
+            pad_length = self.max_num_milestones - num_milestones
+            milestones = U.any_concat(
+                [milestones] + [milestones[-1][None]] * pad_length
+            )
+            milestone_indices = U.any_concat(
+                [milestone_indices] + [milestone_indices[-1][None]] * pad_length
+            )
+            # if milestones.ndim != rgb_milestones.ndim:
+            if rgb_milestones is not None:
+                rgb_milestones = U.any_concat(
+                    [rgb_milestones] + [rgb_milestones[-1][None]] * pad_length
+                )
+            if rgb_no_robot_milestones is not None:
+                rgb_no_robot_milestones = U.any_concat(
+                    [rgb_no_robot_milestones]
+                    + [rgb_no_robot_milestones[-1][None]] * pad_length
+                )
+        rollout_data = OrderedDict(
+            milestones=milestones,
+            # rgb_milestones=rgb_milestones,
+            reset_kwargs=data["reset_kwargs"],
+            # rgb_no_robot_milestones=rgb_no_robot_milestones,
+            milestone_indices=milestone_indices,
+        )
+        if rgb_milestones is not None:
+            rollout_data["rgb_milestones"] = rgb_milestones
+        if rgb_no_robot_milestones is not None:
+            rollout_data["rgb_no_robot_milestones"] = rgb_no_robot_milestones
+        return rollout_data
+
+
+PROMPT_DICT = dict(
+    microwave="open the microwave",
+    kettle="move the kettle to the top left stove",
+    light_switch="turn on the light",
+    hinge_cabinet="open the left hinge cabinet",
+    slide_cabinet="open the right slide cabinet",
+    top_burner="turn on the top left burner",
+    bottom_burner="turn on the bottom left burner",
+)
+PROMPT_DICT.update({k.replace("_", " "): v for k, v in PROMPT_DICT.items()})
+
+
+def task_elements_to_prompt(task_elements: list) -> str | list[str]:
+    if isinstance(task_elements[0], str):
+        prompt = ", ".join([PROMPT_DICT[_] for _ in task_elements])
+        return prompt[0].capitalize() + prompt[1:]
+    # batch of task elements
+    assert isinstance(task_elements[0][0], str), task_elements
+    return [task_elements_to_prompt(ele) for ele in task_elements]

uvd/decomp/__init__.py

@@ -0,0 +1 @@
+from .decomp import decomp_trajectories

uvd/decomp/decomp.py

@@ -0,0 +1,636 @@
+from __future__ import annotations
+
+import random
+from typing import Literal, NamedTuple, Callable
+
+import numpy as np
+import torch
+import wandb
+from matplotlib import pyplot as plt
+from scipy.signal import medfilt
+from scipy.signal import savgol_filter, argrelextrema
+
+import uvd.utils as U
+from uvd.decomp.kernel_reg import KernelRegression
+
+
+def linear_random_skip(
+    cur_step: int, next_goal_step: int, ratio: float = 0.0, progress_lower: float = 0.3
+) -> bool:
+    """Progress toward next goal exceed `progress_lower`, linearly increasing
+    temperature for P(skip) <= ratio."""
+    if ratio == 0.0 or progress_lower == 1:
+        return False
+    assert cur_step <= next_goal_step, f"{cur_step} > {next_goal_step}"
+    # linearly increase until ratio
+    progress = cur_step / next_goal_step
+    temp = (
+        1.0
+        if progress < progress_lower
+        else ((progress - progress_lower) / (1 - progress_lower)) * ratio
+    )
+    if 0 < temp < random.random():
+        return True
+    return False
+
+
+class DecompMeta(NamedTuple):
+    milestone_indices: list
+    milestone_starts: list | None = None
+    iter_curves: list[np.ndarray] | None = None
+
+
+def _debug_plt(
+    xs: np.ndarray,
+    embed_distances: np.ndarray,
+    starts: np.ndarray | list,
+    ends: np.ndarray | list,
+    return_numpy: bool = False,
+):
+    fig = plt.figure()
+    for s in starts:
+        plt.axvline(x=s, linestyle="--")
+    for e in ends:
+        plt.axvline(x=e, linestyle="dotted")
+    plt.plot(xs, np.gradient(embed_distances), linewidth=1.5, label="1st derivative")
+    plt.plot(
+        xs,
+        np.gradient(np.gradient(embed_distances)),
+        linewidth=1.5,
+        label="2nd derivative",
+    )
+    plt.plot(xs, embed_distances, linewidth=1.5, label="embedding distance")
+    plt.legend()
+    if return_numpy:
+        # fig.canvas.draw()
+        return U.plt_to_numpy(fig)
+    else:
+        plt.show()
+
+
+def embed_decomp_no_robot(
+    embeddings: np.ndarray | torch.Tensor,
+    no_robot_embeddings: np.ndarray | torch.Tensor,
+    window_length: int = 10,
+    derivative_order: int = 1,
+    derivative_threshold: float = 1e-3,
+    threshold_subgoal_passing: float | None = None,
+    force_interleave: bool = False,
+    debug_plt: bool = False,
+    debug_plt_to_wandb: bool = False,
+    task_name: str | None = None,
+    fill_embeddings: bool = True,
+):
+    debug_plt = debug_plt and U.is_rank_zero()
+    debug_plt_to_wandb = debug_plt and debug_plt_to_wandb
+    if threshold_subgoal_passing is not None:
+        assert 0 < threshold_subgoal_passing <= 1.0, threshold_subgoal_passing
+    if isinstance(embeddings, torch.Tensor):
+        device = embeddings.device
+    else:
+        device = None
+    # clip based preprocessors would have bf16 by default
+    embeddings = U.any_to_numpy(embeddings, dtype="float32")
+    no_robot_embeddings = U.any_to_numpy(no_robot_embeddings, dtype="float32")
+    # L, N (though can be rgb for debugging as well)
+    traj_length = embeddings.shape[0]
+
+    debug_plt_figs = []
+    if fill_embeddings:
+        milestone_embeddings = []
+    else:
+        milestone_embeddings = None
+    # milestone indices, i.e. end of obj changing
+    subgoal_indices = []
+    # indices when obj changing starts, i.e. milestone indices for hand reaching
+    subgoal_starts = []
+    cur_subgoal_idx = traj_length - 1
+    # back to start
+    while cur_subgoal_idx > 15:
+        unnormalized_embed_distance = np.linalg.norm(
+            no_robot_embeddings[: cur_subgoal_idx + 1]
+            - no_robot_embeddings[cur_subgoal_idx],
+            axis=1,
+        )
+        cur_embed_distance = unnormalized_embed_distance / np.linalg.norm(
+            no_robot_embeddings[0] - no_robot_embeddings[cur_subgoal_idx]
+        )
+        cur_embed_distance = medfilt(cur_embed_distance, kernel_size=None)  # smooth
+
+        if derivative_order == 1:
+            slope = np.gradient(cur_embed_distance)
+        elif derivative_order == 2:
+            slope = np.gradient(np.gradient(cur_embed_distance))
+        else:
+            raise NotImplementedError(derivative_order)
+        valid_slope_indices = np.where(np.abs(slope) >= derivative_threshold)[0]
+        # Find the differences between consecutive valid_slope_indices
+        diffs = np.diff(valid_slope_indices)
+        # Find indices where the difference is greater than window_length + 1
+        break_indices = np.where(diffs > window_length + 1)[0]
+        # Extract start and end positions
+        start_positions = valid_slope_indices[np.concatenate([[0], break_indices + 1])]
+        end_positions = valid_slope_indices[
+            np.concatenate((break_indices, [len(valid_slope_indices) - 1]))
+        ]
+
+        # small tolerance
+        start_positions = [max(0, p - 1) for p in start_positions]
+        end_positions = [min(traj_length - 1, p + 1) for p in end_positions]
+
+        if debug_plt:
+            x = np.arange(0, len(cur_embed_distance))
+            fig = _debug_plt(
+                x,
+                cur_embed_distance,
+                start_positions,
+                end_positions,
+                return_numpy=debug_plt_to_wandb,
+            )
+            if debug_plt_to_wandb:
+                debug_plt_figs.append(fig)
+
+        subgoal_indices.append(cur_subgoal_idx)
+        subgoal_starts.append(start_positions[-1])
+        if len(end_positions) < 2 or end_positions[-2] < 15:
+            if threshold_subgoal_passing is not None:
+                cur_milestone_dist = None
+                if len(subgoal_indices) > 1:
+                    cur_milestone_dist = np.linalg.norm(
+                        no_robot_embeddings[cur_subgoal_idx]  # cur
+                        - no_robot_embeddings[subgoal_indices[-2]]  # prev
+                    )  # use the order from start to end
+                if fill_embeddings:
+                    for step in reversed(range(cur_subgoal_idx + 1)):
+                        if (
+                            len(subgoal_indices) == 1
+                            or unnormalized_embed_distance[step] / cur_milestone_dist
+                            > threshold_subgoal_passing
+                        ):
+                            milestone_embeddings.append(embeddings[cur_subgoal_idx])
+                        else:
+                            milestone_embeddings.append(embeddings[subgoal_indices[-2]])
+            break
+        if threshold_subgoal_passing is not None:
+            cur_milestone_dist = None
+            if len(subgoal_indices) > 1:
+                cur_milestone_dist = np.linalg.norm(
+                    no_robot_embeddings[cur_subgoal_idx]  # cur
+                    - no_robot_embeddings[subgoal_indices[-2]]  # prev
+                )
+            if fill_embeddings:
+                for step in reversed(range(end_positions[-2] + 1, cur_subgoal_idx + 1)):
+                    # not pass threshold or 1st iter with only last frame contained
+                    if (
+                        len(subgoal_indices) == 1
+                        or unnormalized_embed_distance[step] / cur_milestone_dist
+                        > threshold_subgoal_passing
+                    ):
+                        # use the real subgoal
+                        milestone_embeddings.append(embeddings[cur_subgoal_idx])
+                    else:
+                        # skip to the next subgoal if passing threshold
+                        milestone_embeddings.append(embeddings[subgoal_indices[-2]])
+        cur_subgoal_idx = end_positions[-2]
+
+    subgoal_starts = list(reversed(subgoal_starts))
+    subgoal_indices = list(reversed(subgoal_indices))
+
+    if len(debug_plt_figs) > 0 and wandb.run is not None:
+        debug_plt_figs = np.concatenate(debug_plt_figs, axis=1)
+        wandb.log(
+            {
+                f"decomp_curves/{task_name}": wandb.Image(
+                    debug_plt_figs,
+                    caption=f"starts: {subgoal_starts}, ends: {subgoal_indices}",
+                )
+            }
+        )
+
+    assert len(subgoal_starts) == len(
+        subgoal_indices
+    ), f"{subgoal_starts}, {subgoal_indices}"
+
+    if force_interleave:
+        _starts, _ends = [], []
+        for i in range(len(subgoal_starts)):
+            if subgoal_starts[i] < subgoal_indices[i]:
+                _starts.append(subgoal_starts[i])
+                _ends.append(subgoal_indices[i])
+            else:
+                U.get_logger().warning(
+                    f"{subgoal_starts} & {subgoal_indices} not interleaved"
+                )
+
+    if fill_embeddings:
+        if threshold_subgoal_passing is not None:
+            milestone_embeddings = np.stack(list(reversed(milestone_embeddings)))
+        else:
+            # slightly faster to do once here without threshold checking
+            milestone_embeddings = np.concatenate(
+                [embeddings[subgoal_indices[0], ...][None]]
+                + [
+                    np.full((end - start, *embeddings.shape[1:]), embeddings[end, ...])
+                    for start, end in zip([0] + subgoal_indices[:-1], subgoal_indices)
+                ],
+            )
+
+        if device is not None:
+            milestone_embeddings = U.any_to_torch_tensor(
+                milestone_embeddings, device=device
+            )
+    return milestone_embeddings, DecompMeta(
+        milestone_indices=subgoal_indices, milestone_starts=subgoal_starts
+    )
+
+
+def embed_decomp_no_robot_extended(
+    embeddings: np.ndarray | torch.Tensor,
+    no_robot_embeddings: np.ndarray | torch.Tensor,
+    threshold_subgoal_passing: float | None = None,
+    **kwargs,
+):
+    kwargs["fill_embeddings"] = False
+    _, decomp_meta = embed_decomp_no_robot(
+        embeddings,
+        no_robot_embeddings,
+        threshold_subgoal_passing=None,
+        **kwargs,
+    )
+    milestone_indices = decomp_meta.milestone_indices
+    milestone_starts = decomp_meta.milestone_starts
+    norm = (
+        np.linalg.norm if isinstance(embeddings[0], np.ndarray) else torch.linalg.norm
+    )
+    assert len(milestone_starts) == len(milestone_indices)
+
+    milestone_embeddings = []
+    hybrid_indices = list(sorted(milestone_starts + milestone_indices))
+    prev_idx = -1
+    s = -1
+    init_dist = None
+    for i, goal_idx in enumerate(hybrid_indices):
+        once_passed = False
+        for _ in range(goal_idx - prev_idx):
+            s += 1
+            if threshold_subgoal_passing is None:
+                milestone_embeddings.append(embeddings[goal_idx])
+            elif once_passed:
+                milestone_embeddings.append(embeddings[hybrid_indices[i + 1]])
+            else:
+                raw_cur_dist = float(norm(embeddings[s] - embeddings[goal_idx]))
+                if init_dist is None:
+                    assert s == 0, s
+                    cur_dist = 1.0
+                    init_dist = max(raw_cur_dist, 1e-7)
+                else:
+                    cur_dist = raw_cur_dist / init_dist
+                if (
+                    cur_dist <= threshold_subgoal_passing
+                    and i < len(hybrid_indices) - 1
+                ):
+                    init_dist = float(
+                        norm(embeddings[s] - embeddings[hybrid_indices[i + 1]])
+                    )
+                    init_dist = max(init_dist, 1e-7)
+                    once_passed = True
+                    milestone_embeddings.append(embeddings[hybrid_indices[i + 1]])
+                else:
+                    milestone_embeddings.append(embeddings[goal_idx])
+        prev_idx = goal_idx
+
+    milestone_embeddings = U.any_stack(milestone_embeddings)
+    U.assert_(milestone_embeddings.shape, embeddings.shape)
+    return milestone_embeddings, DecompMeta(milestone_indices=hybrid_indices)
+
+
+def get_hybrid_milestones(
+    start_embeddings: np.ndarray,  # w. robot
+    end_embeddings: np.ndarray,  # w.o robot
+    milestone_starts: list,
+    milestone_indices: list,
+) -> np.ndarray:
+    assert len(milestone_starts) == len(milestone_indices)
+    assert len(start_embeddings) == len(end_embeddings)
+    milestone_only = len(milestone_starts) == len(start_embeddings)
+    hybrid_milestones = np.empty(
+        (start_embeddings.shape[0] * 2, *start_embeddings.shape[1:]),
+        dtype=start_embeddings.dtype,
+    )
+    hybrid_milestones[::2] = (
+        start_embeddings if milestone_only else start_embeddings[milestone_starts]
+    )
+    hybrid_milestones[1::2] = (
+        end_embeddings if milestone_only else end_embeddings[milestone_indices]
+    )
+    return hybrid_milestones
+
+
+def embedding_decomp(
+    embeddings: np.ndarray | torch.Tensor,
+    normalize_curve: bool = True,
+    min_interval: int = 18,
+    window_length: int | None = None,
+    smooth_method: Literal["kernel", "savgol"] = "kernel",
+    extrema_comparator: Callable = np.greater,
+    fill_embeddings: bool = True,
+    return_intermediate_curves: bool = False,
+    **kwargs,
+) -> tuple[torch.Tensor | np.ndarray, DecompMeta]:
+    if torch.is_tensor(embeddings):
+        device = embeddings.device
+        embeddings = U.any_to_numpy(embeddings)
+    else:
+        device = None
+    # L, N
+    assert embeddings.ndim == 2, embeddings.shape
+    traj_length = embeddings.shape[0]
+
+    cur_goal_idx = traj_length - 1
+    goal_indices = [cur_goal_idx]
+    cur_embeddings = embeddings[
+        max(0, cur_goal_idx - (window_length or cur_goal_idx)) : cur_goal_idx + 1
+    ]
+    iterate_num = 0
+    iter_curves = [] if return_intermediate_curves else None
+    while cur_goal_idx > (window_length or min_interval):
+        iterate_num += 1
+        # get goal embedding
+        goal_embedding = cur_embeddings[-1]
+        distances = np.linalg.norm(cur_embeddings - goal_embedding, axis=1)
+        if normalize_curve:
+            distances = distances / np.linalg.norm(cur_embeddings[0] - goal_embedding)
+
+        x = np.arange(
+            max(0, cur_goal_idx - (window_length or cur_goal_idx)), cur_goal_idx + 1
+        )
+
+        if smooth_method == "kernel":
+            smooth_kwargs = dict(kernel="rbf", gamma=0.08)
+            smooth_kwargs.update(kwargs or {})
+            kr = KernelRegression(**smooth_kwargs)
+            kr.fit(x.reshape(-1, 1), distances)
+            distance_smoothed = kr.predict(x.reshape(-1, 1))
+        elif smooth_method == "savgol":
+            smooth_kwargs = dict(window_length=85, polyorder=2, mode="nearest")
+            smooth_kwargs.update(kwargs or {})
+            distance_smoothed = savgol_filter(distances, **smooth_kwargs)
+        elif smooth_method is None:
+            distance_smoothed = distances
+        else:
+            raise NotImplementedError(smooth_method)
+
+        if iter_curves is not None:
+            iter_curves.append(distance_smoothed)
+
+        extrema_indices = argrelextrema(distance_smoothed, extrema_comparator)[0]
+        x_extrema = x[extrema_indices]
+
+        update_goal = False
+        for i in range(len(x_extrema) - 1, -1, -1):
+            if cur_goal_idx < min_interval:
+                break
+            if (
+                cur_goal_idx - x_extrema[i] > min_interval
+                and x_extrema[i] > min_interval
+            ):
+                cur_goal_idx = x_extrema[i]
+                update_goal = True
+                goal_indices.append(cur_goal_idx)
+                break
+
+        if not update_goal or cur_goal_idx < min_interval:
+            break
+        cur_embeddings = embeddings[
+            max(0, cur_goal_idx - (window_length or cur_goal_idx)) : cur_goal_idx + 1
+        ]
+
+    goal_indices = goal_indices[::-1]
+    if fill_embeddings:
+        milestone_embeddings = np.concatenate(
+            [embeddings[goal_indices[0], ...][None]]
+            + [
+                np.full((end - start, *embeddings.shape[1:]), embeddings[end, ...])
+                for start, end in zip([0] + goal_indices[:-1], goal_indices)
+            ],
+        )
+        if device is not None:
+            milestone_embeddings = U.any_to_torch_tensor(
+                milestone_embeddings, device=device
+            )
+    else:
+        milestone_embeddings = None
+    return milestone_embeddings, DecompMeta(
+        milestone_indices=goal_indices, iter_curves=iter_curves
+    )
+
+
+def goal_idx_from_mask(goal_achieved_mask):
+    diff = np.diff(goal_achieved_mask)
+    goal_indices = np.where(diff != 0)[0] + 1
+    traj_length = goal_achieved_mask.shape[0]
+    goal_indices[-1] = traj_length - 1  # last
+    goal_indices = goal_indices.tolist()
+    return goal_indices
+
+
+def oracle_decomp(
+    embeddings: np.ndarray | torch.Tensor | None,
+    goal_achieved_mask: np.ndarray,
+    random_skip_ratio: float | None = None,
+    linearly_random_skip_lower: float | None = None,
+    fill_embeddings: bool = True,
+) -> tuple[torch.Tensor | np.ndarray, DecompMeta]:
+    """Note: embeddings here only has the oracle subgoals, not full trajectory"""
+    goal_indices = goal_idx_from_mask(goal_achieved_mask)
+    if not fill_embeddings:
+        return None, DecompMeta(milestone_indices=goal_indices)
+
+    traj_length = goal_achieved_mask.shape[0]
+    assert embeddings.shape[0] < traj_length, embeddings.shape
+    milestone_embeddings = (
+        torch.empty(
+            (goal_achieved_mask.shape[0], *embeddings.shape[1:]),
+            dtype=embeddings.dtype,
+            device=embeddings.device,
+        )
+        if not isinstance(embeddings, np.ndarray)
+        else np.empty(
+            (goal_achieved_mask.shape[0], *embeddings.shape[1:]),
+            dtype=embeddings.dtype,
+        )
+    )
+
+    assert len(goal_indices) == len(embeddings), goal_indices
+
+    last_embedding = embeddings[-1]
+    for i, idx in enumerate(goal_achieved_mask):
+        if idx >= embeddings.shape[0]:
+            # If the index in the mask is greater than the highest index in the embedding,
+            # just use the last row of the embedding
+            milestone_embeddings[i] = last_embedding
+        else:
+            skip = False
+            if linearly_random_skip_lower is not None:
+                skip = linear_random_skip(
+                    cur_step=i,
+                    next_goal_step=goal_indices[idx],
+                    ratio=random_skip_ratio,
+                    progress_lower=linearly_random_skip_lower,
+                )
+            elif (
+                random_skip_ratio is not None
+                and 0 < random_skip_ratio < random.random()
+            ):
+                skip = True
+            if skip:
+                milestone_embeddings[i] = embeddings[min(idx + 1, len(embeddings) - 1)]
+            else:
+                milestone_embeddings[i] = embeddings[idx]
+    return milestone_embeddings, DecompMeta(milestone_indices=goal_indices)
+
+
+def random_decomp(
+    embeddings: np.ndarray | torch.Tensor,
+    num_milestones: int | tuple[int, int],
+    fill_embeddings: bool = True,
+) -> tuple[torch.Tensor | np.ndarray, DecompMeta]:
+    if not isinstance(num_milestones, int):
+        assert len(num_milestones) == 2, num_milestones
+        # by randomly sample from lower and higher bound
+        num_milestones = random.randint(*num_milestones)
+    traj_length = embeddings.shape[0]
+    goal_indices = random.sample(range(traj_length), k=num_milestones)
+    goal_indices = list(sorted(goal_indices))
+    if fill_embeddings:
+        milestone_embeddings = (
+            torch.empty_like(
+                embeddings, dtype=embeddings.dtype, device=embeddings.device
+            )
+            if not isinstance(embeddings, np.ndarray)
+            else np.empty_like(embeddings, dtype=embeddings.dtype)
+        )
+        for i, goal_idx in enumerate(goal_indices):
+            milestone_embeddings[
+                (goal_indices[i - 1] + 1) if i != 0 else 0 : goal_idx + 1
+            ] = embeddings[goal_idx]
+    else:
+        milestone_embeddings = None
+    return milestone_embeddings, DecompMeta(milestone_indices=goal_indices)
+
+
+def equally_decomp(
+    embeddings: np.ndarray | torch.Tensor,
+    num_milestones: int | tuple[int, int],
+    fill_embeddings: bool = True,
+) -> tuple[torch.Tensor | np.ndarray, DecompMeta]:
+    if not isinstance(num_milestones, int):
+        assert len(num_milestones) == 2, num_milestones
+        # by randomly sample from lower and higher bound
+        num_milestones = random.randint(*num_milestones)
+    traj_length = embeddings.shape[0]
+    indices = np.linspace(0, traj_length - 1, num_milestones + 1, dtype=int)
+    if fill_embeddings:
+        milestone_embeddings = (
+            torch.empty_like(
+                embeddings, dtype=embeddings.dtype, device=embeddings.device
+            )
+            if not isinstance(embeddings, np.ndarray)
+            else np.empty_like(
+                embeddings,
+                dtype=embeddings.dtype,
+            )
+        )
+        for i, goal_idx in enumerate(indices[1:], start=1):
+            milestone_embeddings[
+                (indices[i - 1] + 1) if i != 1 else 0 : goal_idx + 1
+            ] = embeddings[goal_idx]
+    else:
+        milestone_embeddings = None
+    return milestone_embeddings, DecompMeta(milestone_indices=indices[1:].tolist())
+
+
+def near_future_decomp(
+    embeddings: np.ndarray | torch.Tensor, advance_steps: int, **kwargs
+) -> tuple[torch.Tensor | np.ndarray, DecompMeta]:
+    return equally_decomp(
+        embeddings, num_milestones=embeddings.shape[0] // advance_steps, **kwargs
+    )
+
+
+def no_decomp(
+    embeddings: np.ndarray | torch.Tensor, fill_embeddings: bool = True
+) -> tuple[torch.Tensor | np.ndarray, DecompMeta]:
+    """Only conditioned on final goal."""
+    if not fill_embeddings:
+        return None, DecompMeta(milestone_indices=[-1])
+    return embeddings[-1, ...].expand_as(embeddings).clone() if not isinstance(
+        embeddings, np.ndarray
+    ) else np.full(
+        embeddings.shape,
+        embeddings[-1, ...],
+        dtype=embeddings.dtype,
+    ), DecompMeta(
+        milestone_indices=[-1]
+    )
+
+
+def decomp_trajectories(
+    method_name: Literal[
+        "embed", "embed_no_robot", "oracle", "random", "equally", "near_future"
+    ]
+    | None,
+    embeddings: np.ndarray | torch.Tensor,
+    **kwargs,
+) -> tuple[torch.Tensor | np.ndarray, DecompMeta]:
+    assert embeddings.ndim == 2 or embeddings.ndim == 4, (
+        f"input embedding should be either 2 dimensional, "
+        f"with (L, feature_dim), or raw rgb with shape (L, H, W, 3), "
+        f"but get {embeddings.shape}"
+    )
+    if method_name is None:
+        return no_decomp(embeddings)
+    assert method_name in DEFAULT_DECOMP_KWARGS, method_name
+    method_kwargs = DEFAULT_DECOMP_KWARGS[method_name]
+    method_kwargs.update(kwargs)
+    if method_name == "embed":
+        return embedding_decomp(embeddings=embeddings, **method_kwargs)
+    elif method_name == "embed_no_robot":
+        return embed_decomp_no_robot(embeddings=embeddings, **method_kwargs)
+    elif method_name == "embed_no_robot_extended":
+        return embed_decomp_no_robot_extended(embeddings=embeddings, **method_kwargs)
+    elif method_name == "oracle":
+        return oracle_decomp(embeddings, **method_kwargs)
+    elif method_name == "random":
+        return random_decomp(embeddings, **method_kwargs)
+    elif method_name == "equally":
+        return equally_decomp(embeddings, **method_kwargs)
+    elif method_name == "near_future":
+        return near_future_decomp(embeddings, **method_kwargs)
+    raise NotImplementedError(method_name)
+
+
+DEFAULT_DECOMP_KWARGS = dict(
+    embed=dict(
+        normalize_curve=False,
+        min_interval=18,
+        smooth_method="kernel",
+        gamma=0.08,
+    ),
+    embed_no_robot=dict(
+        window_length=8,
+        derivative_order=1,
+        derivative_threshold=1e-3,
+        threshold_subgoal_passing=None,
+    ),
+    embed_no_robot_extended=dict(
+        window_length=3,
+        derivative_order=1,
+        derivative_threshold=1e-3,
+        threshold_subgoal_passing=None,
+    ),
+    oracle=dict(),
+    random=dict(num_milestones=(3, 6)),
+    equally=dict(num_milestones=(3, 6)),
+    near_future=dict(advance_steps=5),
+)

uvd/decomp/kernel_reg.py

@@ -0,0 +1,91 @@
+"""The :mod:`sklearn.kernel_regressor` module implements the Kernel
+Regressor."""
+# Author: Jan Hendrik Metzen <janmetzen@mailbox.de>
+#
+# License: BSD 3 clause
+
+import numpy as np
+from sklearn.base import BaseEstimator, RegressorMixin
+from sklearn.metrics.pairwise import pairwise_kernels
+
+
+class KernelRegression(BaseEstimator, RegressorMixin):
+    """Nadaraya-Watson kernel regression with automatic bandwidth selection.
+    This implements Nadaraya-Watson kernel regression with (optional) automatic
+    bandwith selection of the kernel via leave-one-out cross-validation. Kernel
+    regression is a simple non-parametric kernelized technique for learning a
+    non-linear relationship between input variable(s) and a target variable.
+
+    Parameters
+    ----------
+    kernel : string or callable, default="rbf"
+        Kernel map to be approximated. A callable should accept two arguments
+        and the keyword arguments passed to this object as kernel_params, and
+        should return a floating point number.
+    gamma : float, default=None
+        Gamma parameter for the RBF ("bandwidth"), polynomial,
+        exponential chi2 and sigmoid kernels. Interpretation of the default
+        value is left to the kernel; see the documentation for
+        sklearn.metrics.pairwise. Ignored by other kernels. If a sequence of
+        values is given, one of these values is selected which minimizes
+        the mean-squared-error of leave-one-out cross-validation.
+    See also
+    --------
+    sklearn.metrics.pairwise.kernel_metrics : List of built-in kernels.
+    """
+
+    def __init__(self, kernel="rbf", gamma=None):
+        self.kernel = kernel
+        self.gamma = gamma
+
+    def fit(self, X, y):
+        """Fit the model.
+
+        Parameters
+        ----------
+        X : array-like of shape = [n_samples, n_features]
+            The training input samples.
+        y : array-like, shape = [n_samples]
+            The target values
+        Returns
+        -------
+        self : object
+            Returns self.
+        """
+        self.X = X
+        self.y = y
+
+        if hasattr(self.gamma, "__iter__"):
+            self.gamma = self._optimize_gamma(self.gamma)
+
+        return self
+
+    def predict(self, X):
+        """Predict target values for X.
+
+        Parameters
+        ----------
+        X : array-like of shape = [n_samples, n_features]
+            The input samples.
+        Returns
+        -------
+        y : array of shape = [n_samples]
+            The predicted target value.
+        """
+        K = pairwise_kernels(self.X, X, metric=self.kernel, gamma=self.gamma)
+        return (K * self.y[:, None]).sum(axis=0) / K.sum(axis=0)
+
+    def _optimize_gamma(self, gamma_values):
+        # Select specific value of gamma from the range of given gamma_values
+        # by minimizing mean-squared error in leave-one-out cross validation
+        mse = np.empty_like(gamma_values, dtype=np.float)
+        for i, gamma in enumerate(gamma_values):
+            K = pairwise_kernels(self.X, self.X, metric=self.kernel, gamma=gamma)
+            np.fill_diagonal(K, 0)  # leave-one-out
+            Ky = K * self.y[:, np.newaxis]
+            y_pred = Ky.sum(axis=0) / K.sum(axis=0)
+            mse[i] = ((y_pred - self.y) ** 2).mean()
+        try:
+            return gamma_values[np.nanargmin(mse)]
+        except:
+            return 0

uvd/envs/evaluator/__init__.py

@@ -0,0 +1,3 @@
+from .evaluator import *
+from .inference_wrapper import *
+from .visualize_wrapper import *

uvd/envs/evaluator/evaluator.py

@@ -0,0 +1,571 @@
+from __future__ import annotations
+
+import copy
+
+import einops
+import gym
+import numpy as np
+import torch
+import torch.nn.functional as F
+import tqdm
+import uvd.utils as U
+import wandb
+from uvd.envs.franka_kitchen import KitchenBase
+from uvd.models.policy import PolicyBase
+
+from .inference_wrapper import InferenceWrapper
+from .vec_envs.vec_env import BaseVectorEnv, ShmemVectorEnv, SubprocVectorEnv
+from .visualize_wrapper import VisualizeWrapper
+from ...models import Preprocessor
+
+__all__ = ["VectorEnvEvaluator"]
+
+
+class VectorEnvEvaluator:
+    def __init__(
+        self,
+        *,
+        env_name: str,
+        reset_kwargs: dict,
+        num_tasks: int,
+        num_envs: int | None,
+        max_horizon: int = 400,
+        no_robot_rgb_milestones: np.ndarray | None = None,
+        use_no_robot_milestones: bool = False,
+        inference_kwargs: dict | None,
+        random_skip_inference: bool = False,
+        use_milestone_distances_normalization: bool = False,
+        decomp_method: str | None = None,
+        milestones: np.ndarray,
+        rgb_milestones: np.ndarray | None = None,
+        seed: int | None = None,
+        seed_list: list[int] | None = None,
+        save_video: bool = False,
+        save_video_kwargs: dict | None = None,
+        device: int | str | torch.device | None = None,
+        use_milestone_compressor: bool = False,
+        milestone_indices: np.ndarray | None = None,
+    ):
+        self.env_name = env_name
+        self.reset_kwargs = reset_kwargs
+        num_envs = min(num_envs or num_tasks, num_tasks, milestones.shape[0])
+        self.num_envs = num_envs
+        self.rollout_list = self.generate_rollout_list(num_envs, num_tasks)
+
+        inference_kwargs = inference_kwargs or {}
+        if not random_skip_inference:
+            inference_kwargs.update(random_skip_ratio=0.0)
+        # inference_kwargs["hybrid"] = decomp_method == "embed_no_robot_extended"
+        inference_kwargs["hybrid"] = False
+        if decomp_method == "embed_no_robot_extended":
+            use_no_robot_milestones = False
+        inference_kwargs["use_milestone_compressor"] = use_milestone_compressor
+        self.inference_kwargs = inference_kwargs
+        self.decomp_method = decomp_method
+        self.max_horizon = max_horizon
+        self.milestones = milestones
+        self.rgb_milestones = rgb_milestones
+        self.no_robot_rgb_milestones = no_robot_rgb_milestones
+        if not use_no_robot_milestones:
+            self.no_robot_rgb_milestones = None
+        elif self.no_robot_rgb_milestones is None:
+            U.rank_zero_print(
+                f"WARNING: set to use no robot milestones {use_no_robot_milestones} but there is no input",
+                color="red",
+            )
+        milestone_distances = None
+        if use_milestone_distances_normalization and milestones.shape[1] > 1:
+            if self.no_robot_rgb_milestones is not None:
+                # assume only using if with linear embedding
+                assert (
+                    self.no_robot_rgb_milestones.ndim == 3
+                ), self.no_robot_rgb_milestones.shape
+                # N_ENVS, N_GOALS, D
+                milestone_distances = np.linalg.norm(
+                    self.no_robot_rgb_milestones[:, :-1]
+                    - self.no_robot_rgb_milestones[:, 1:],
+                    axis=-1,
+                )
+            else:
+                assert self.milestones.ndim == 3, self.milestones.shape
+                milestone_distances = np.linalg.norm(
+                    self.milestones[:, :-1] - self.milestones[:, 1:], axis=-1
+                )
+            assert milestone_distances.shape == (
+                milestones.shape[0],
+                milestones.shape[1] - 1,
+            ), milestone_distances.shape
+        self.milestone_distances = milestone_distances
+
+        if seed_list is None:
+            _rng = np.random.default_rng(seed=seed)
+            seed_list = [_rng.integers(0, 2**31 - 1) for _ in range(num_tasks)]
+        assert len(seed_list) == num_tasks
+        self._num_tasks = num_tasks
+        self.seed_list = seed_list
+        self.save_video = save_video
+        if save_video:
+            assert save_video_kwargs is not None
+            self.video_fps = save_video_kwargs["fps"]
+            self.add_debug_text = save_video_kwargs["add_debug_text"]
+            self.wandb_logging = save_video_kwargs["wandb_logging"]
+            self.save_locally = save_video_kwargs["save_locally"]
+            if self.save_locally:
+                self.save_path = save_video_kwargs["save_path"]
+        self.device = device
+
+        self.batched_env: BaseVectorEnv | None = None
+        self.use_milestone_compressor = use_milestone_compressor
+        self.milestone_indices = milestone_indices  # B, N
+
+    @staticmethod
+    def confirm_embedding(
+        embedding: np.ndarray | torch.Tensor, preprocessor: Preprocessor | None
+    ) -> np.ndarray:
+        """make sure the embedding is linear for calculating embed distance during inference
+        Args:
+            embedding:
+                if embedding.ndim == 2: B, d already;
+                if embedding.ndim == 3: B, L, d for milestone embeddings;
+                if embedding.ndim == 4: B, H, W, 3 for raw rgb;
+                                        or without pooling, e.g. B, 2048, 7, 7
+                if embedding.ndim == 5: B, L, H, W, 3 for L milestones;
+            preprocessor: frozen preprocessor
+        Return:
+            embedding.ndim = 2 for single step obs, ndim=3 for milestone embeddings
+        """
+        if preprocessor is None:
+            return U.any_to_numpy(embedding)
+        preprocessor_output_dim = preprocessor.output_dim
+        if embedding.ndim not in [2, 3]:
+            input_dims = embedding.shape
+            if (
+                embedding.ndim in [4, 5]
+                and embedding.shape[-3:] != preprocessor_output_dim
+            ):
+                # raw rgb case (B, (L) H, W, 3)
+                if embedding.ndim == 5:
+                    embedding = einops.rearrange(embedding, "b l h w c -> (b l) h w c")
+                with torch.no_grad():
+                    embedding = preprocessor.process(
+                        embedding,
+                        # embedding.reshape((np.prod(input_dims[:-3]), *input_dims[-3:])),
+                        reconstruct_linear=True,
+                    )
+            # if preprocessor.remove_pool:
+            if (
+                embedding.ndim in [4, 5]
+                and embedding.shape[-3:] == preprocessor_output_dim
+            ):
+                U.rank_zero_print("DEPRECATED!", color="red")
+                # no pooling case, e.g. (B, (L), 2048, 7, 7)
+                embedding = embedding.reshape(
+                    (np.prod(input_dims[:-3]), *preprocessor.output_dim)
+                )
+                with torch.no_grad():
+                    embedding = F.adaptive_avg_pool2d(embedding, output_size=(1, 1))
+                    embedding = torch.flatten(embedding, 1)
+                    if preprocessor.preprocessor_fc is not None:
+                        embedding = preprocessor.preprocessor_fc(embedding)
+            embedding = embedding.reshape((*input_dims[:-3], embedding.shape[-1]))
+        if torch.is_tensor(embedding):
+            embedding = embedding.cpu().numpy()
+        assert embedding.ndim in [2, 3], embedding.shape
+        return embedding
+
+    def rollout(self, policy: PolicyBase, mode: str, epoch: int) -> dict:
+        use_kv_cache = hasattr(policy, "use_kv_cache") and policy.use_kv_cache
+        is_causal = hasattr(policy, "causal") and policy.causal
+        cache_obs = is_causal and not use_kv_cache
+        if cache_obs:  # means causal and not use kv cache
+            U.rank_zero_print(
+                "WARNING: `use_kv_cache=False` NOT RECOMMEND SINCE SOOOOOO INEFFICIENT!",
+                color="red",
+            )
+            self.inference_kwargs["cache_history"] = True
+        elif is_causal and use_kv_cache:
+            # keep the bs the same so not terminate the episode
+            self.inference_kwargs["dummy_rtn"] = True
+
+        if self.batched_env is None:
+            if cache_obs:  # dynamic leading dim obs for history
+                self.batched_env = SubprocVectorEnv(self._create_env_fns())
+            else:
+                self.batched_env = ShmemVectorEnv(self._create_env_fns())
+            self.batched_env.reset()
+
+        results = {k: None for k in range(self._num_tasks)}
+
+        seed_list = list(self.seed_list)
+        logging_videos = self.save_video
+        if logging_videos:
+            self.batched_env.set_env_attr("recording", True)
+        for i, env_ids in enumerate(self.rollout_list):
+            if logging_videos and i != 0 and not self.save_locally:
+                # logging wandb only 1st iter
+                self.batched_env.set_env_attr("recording", False)
+                logging_videos = False
+
+            if is_causal and use_kv_cache:
+                policy.reset_cache()
+
+            num_env_this_batch = len(env_ids)
+            self.batched_env.seed(
+                seed=[seed_list.pop() for _ in range(num_env_this_batch)]
+            )
+            self.prepare_states_before_reset(env_ids, global_idx=i)
+            if self.milestone_distances is not None:
+                milestone_distances_this_batch = self.milestone_distances[
+                    self.num_envs * i : self.num_envs * i + num_env_this_batch, ...
+                ].copy()
+                self.batched_env.set_env_attr(
+                    "milestone_distances",
+                    milestone_distances_this_batch,
+                    id=env_ids,
+                    diff_value=True,
+                )
+
+            if self.no_robot_rgb_milestones is not None:
+                no_robot_milestones_this_batch = self.no_robot_rgb_milestones[
+                    self.num_envs * i : self.num_envs * i + num_env_this_batch, ...
+                ].copy()
+                # num_env, num_milestone, (H, W, 3 or D)
+                no_robot_milestones_this_batch = self.confirm_embedding(
+                    no_robot_milestones_this_batch, preprocessor=policy.preprocessor
+                )
+
+                self.batched_env.set_env_attr(
+                    "no_robot_milestones",
+                    no_robot_milestones_this_batch,
+                    id=env_ids,
+                    diff_value=True,
+                )
+            # num_env_this_batch, num_goals, ...
+            milestones_this_batch = self.milestones[
+                self.num_envs * i : self.num_envs * i + num_env_this_batch, ...
+            ].copy()
+            # set milestones for this episode
+            self.batched_env.set_env_attr(
+                "milestones", milestones_this_batch, id=env_ids, diff_value=True
+            )
+
+            if self.milestone_indices is not None:
+                milestone_indices_this_batch = self.milestone_indices[
+                    self.num_envs * i : self.num_envs * i + num_env_this_batch, ...
+                ].copy()
+                self.batched_env.set_env_attr(
+                    "milestone_indices",
+                    milestone_indices_this_batch,
+                    id=env_ids,
+                    diff_value=True,
+                )
+
+            if milestones_this_batch.ndim != 3:
+                # num_env_this_batch, num_goals, 3, H, W
+                if i == 0:
+                    U.rank_zero_print(f"{milestones_this_batch.shape=}", color="blue")
+                self.batched_env.set_env_attr(
+                    "milestone_embeddings",
+                    self.confirm_embedding(
+                        milestones_this_batch.copy(), preprocessor=policy.preprocessor
+                    ),
+                    id=env_ids,
+                    diff_value=True,
+                )
+            if logging_videos:
+                rgb_milestones_this_batch = self.rgb_milestones[
+                    num_env_this_batch * i : num_env_this_batch * (i + 1), ...
+                ].copy()
+                self.batched_env.set_env_attr(
+                    "rgb_milestones",
+                    rgb_milestones_this_batch,
+                    id=env_ids,
+                    diff_value=True,
+                )
+            # reset: num_env_this_batch, h, w, 3 (or dict of obs)
+            obs = self.batched_env.reset(id=env_ids)
+            assert obs.shape[0] == num_env_this_batch, obs.shape
+
+            running_env_ids = np.arange(num_env_this_batch)
+
+            for st in tqdm.tqdm(
+                range(self.max_horizon),
+                initial=1,
+                desc=f"rank {U.get_local_rank()}: Rollout {i * self.num_envs + len(env_ids)}/{self._num_tasks}",
+                leave=False,
+            ):
+                self.batched_env.set_env_attr(
+                    "cur_milestone_idx", value=None, id=running_env_ids
+                )  # for random skipping, sample maybe next goal idx inside env every step
+
+                if cache_obs:
+                    current_milestone = self.batched_env.get_env_attr(
+                        "cached_prev_milestones", id=running_env_ids
+                    )
+                else:
+                    current_milestone = self.batched_env.get_env_attr(
+                        "current_milestone", id=running_env_ids
+                    )
+                current_milestone = U.any_to_numpy(current_milestone)
+                U.assert_(len(current_milestone), len(obs))
+
+                with torch.no_grad():
+                    # only batchify here, keep obs as list as length of running_env_ids anywhere else
+                    batchify_obs = U.batch_observations(obs, device=self.device)
+                    if cache_obs:
+                        b, t, *_ = current_milestone.shape
+                        if st == 0 and batchify_obs["rgb"].shape[:2] != (b, t):
+                            for k in batchify_obs:
+                                batchify_obs[k] = batchify_obs[k][
+                                    :, None, ...
+                                ]  # broadcast T dim
+                        assert batchify_obs["rgb"].shape[:2] == (b, t), (
+                            batchify_obs["rgb"].shape,
+                            b,
+                            t,
+                        )
+                    elif is_causal:
+                        # broadcast T dim, B H W 3 or B D
+                        assert current_milestone.ndim in [2, 4], current_milestone.shape
+                        current_milestone = current_milestone[:, None, ...]
+                        for k in batchify_obs:
+                            batchify_obs[k] = batchify_obs[k][:, None, ...]
+                    # L, n
+                    action, obs_embed, goal_embed = policy(
+                        batchify_obs,
+                        goal=current_milestone,
+                        deterministic=True,
+                        return_embeddings=True,
+                        timesteps=torch.tensor(
+                            [[st]], dtype=torch.int32, device=self.device
+                        ),
+                        input_pos=torch.tensor([st], device=self.device)
+                        if is_causal and use_kv_cache
+                        else None,
+                    )
+                    if action.ndim == 3:  # has T dim
+                        action = action[:, -1, :]
+                # switch milestones inside the inference wrapper based on current obs embedding
+                if self.no_robot_rgb_milestones is not None:
+                    no_robot_rgbs = self.batched_env.get_env_attr(
+                        "current_no_robot_frame", id=running_env_ids
+                    )
+                    cur_obs_embed = self.confirm_embedding(
+                        np.stack(no_robot_rgbs), preprocessor=policy.preprocessor
+                    )
+                else:
+                    cur_obs_embed = self.confirm_embedding(
+                        batchify_obs["rgb"] if obs_embed.ndim != 2 else obs_embed,
+                        preprocessor=policy.preprocessor,
+                    )
+                self.batched_env.set_env_attr(
+                    "current_obs_embedding",
+                    cur_obs_embed,
+                    id=running_env_ids,
+                    diff_value=True,
+                )
+
+                #  next_obs (n, dict/(256, 256, 3)) (n,) (n,)，list[dict]
+                obs, r, done, info = self.batched_env.step(
+                    action.cpu().numpy(), id=running_env_ids
+                )
+                prev_running_env_ids = copy.deepcopy(running_env_ids)
+                if np.any(done):
+                    if is_causal and use_kv_cache:
+                        assert np.all(done), done  # keep bs the same for kv cache
+                    # update running env id & collect results for envs that done
+                    terminated_env_local_idx = np.where(done)[0]
+                    masks = np.ones_like(running_env_ids, dtype=bool)
+                    masks[terminated_env_local_idx] = False
+                    running_env_ids = running_env_ids[masks]
+                    obs = obs[masks]
+                    terminated_local_ids = list(
+                        set(prev_running_env_ids) - set(running_env_ids)
+                    )
+                    metrics = self.batched_env.get_env_attr(
+                        "metrics", id=terminated_local_ids
+                    )
+                    try:
+                        task_names = self.batched_env.get_env_attr(
+                            "task_name", id=terminated_local_ids
+                        )
+                    except AttributeError:
+                        task_names = None
+                    for met_idx, local_id in enumerate(terminated_local_ids):
+                        assert (
+                            local_id not in running_env_ids
+                        ), f"{local_id=} is done but still in {running_env_ids=}"
+                        assert results[local_id + self.num_envs * i] is None, (
+                            f"{local_id + self.num_envs * i}-th task should be only rollout once, "
+                            f"{local_id=}, {running_env_ids=}, {prev_running_env_ids=}, {terminated_local_ids=}"
+                        )
+                        results[local_id + self.num_envs * i] = {
+                            (
+                                k
+                                if ("success" not in k and "completed_tasks" not in k)
+                                or task_names is None
+                                else f"{k}/{task_names[met_idx]}"
+                            ): v
+                            for k, v in metrics[met_idx].items()
+                        }
+
+                if np.all(done):
+                    assert results is not None
+                    break
+
+            # only save first num vec envs videos
+            if logging_videos:  # gather videos after all env done
+                # each video may have diff length among diff envs
+                all_frames = self.batched_env.get_env_attr("frames", id=env_ids)
+                try:
+                    task_names = self.batched_env.get_env_attr("task_name", id=env_ids)
+                except AttributeError:
+                    task_names = None
+                try:
+                    self.logging_videos(
+                        all_frames,
+                        task_names=task_names,
+                        global_task_ids=[
+                            self.num_envs * i + _ for _ in range(num_env_this_batch)
+                        ],
+                        mode=mode,
+                        global_step=epoch,
+                    )
+                except:
+                    pass
+
+        return self.process_results(results, mode=mode)
+
+    def prepare_states_before_reset(self, local_env_ids: list[int], global_idx: int):
+        if self.env_name == "franka_kitchen":
+            reset_states = [
+                dict(
+                    init_qpos=self.reset_kwargs["init_qpos"][
+                        env_id + self.num_envs * global_idx
+                    ],
+                    init_qvel=self.reset_kwargs["init_qvel"][
+                        env_id + self.num_envs * global_idx
+                    ],
+                    task_elements=self.reset_kwargs["task_elements"][
+                        env_id + self.num_envs * global_idx
+                    ],
+                )
+                for env_id in local_env_ids
+            ]
+            self.batched_env.set_env_attr(
+                "reset_states", reset_states, id=local_env_ids, diff_value=True
+            )
+        else:
+            raise NotImplementedError(self.env_name)
+
+    def process_results(self, results: dict, mode: str) -> dict:
+        res = {}
+        successes = []
+        completions = []
+        for result in results.values():
+            for k, v in result.items():
+                if "success" in k:
+                    successes.append(v)
+                elif "completed_tasks" in k:
+                    completions.append(v)
+                res.setdefault(k, []).append(v)
+        assert (
+            len(successes) == self._num_tasks
+        ), f"{len(successes)=} != {self._num_tasks}"
+        return {
+            f"{mode}/success": np.mean(successes),
+            f"{mode}/completed_tasks": np.mean(completions),
+            f"{mode}/num_tasks": float(self._num_tasks),
+            **{f"{mode}/{k}": np.mean(v) for k, v in res.items()},
+        }
+
+    def logging_videos(
+        self,
+        all_frames: list,
+        *,
+        task_names: str | None = None,
+        global_task_ids: list,
+        mode: str,
+        global_step: int,
+    ):
+        cur_rank = U.get_local_rank()
+        for i, ids in enumerate(global_task_ids):
+            # wandb channel-first logging: episode_length, h, w, 3 -> ..., 3, h, w
+            cur_video = U.any_to_numpy(all_frames[i])
+            if cur_video.ndim != 4 or cur_video.shape[-1] != 3:
+                U.rank_zero_print(
+                    f"cur_video has unexpected shape {cur_video.shape}", color="red"
+                )
+                continue
+            cur_video = cur_video.transpose([0, 3, 1, 2])
+            assert cur_video.ndim == 4 and cur_video.shape[1] == 3, cur_video.shape
+            if self.wandb_logging and cur_rank == 0:
+                log_name = U.f_join(
+                    f"{mode}", "rollout-videos", f"{task_names[i]}" or "", f"{ids}"
+                )
+                wandb.log(
+                    {
+                        log_name: wandb.Video(
+                            cur_video, fps=self.video_fps, format="mp4"
+                        )
+                    },
+                )
+            if self.save_locally:
+                save_path = U.f_mkdir(
+                    self.save_path,
+                    f"{mode}-rank_{cur_rank}"
+                    + f"/{task_names[i] if task_names is not None else ''}",
+                )
+                U.save_video(
+                    cur_video,
+                    U.f_join(save_path, f"env_{ids}_epoch_{global_step}.mp4"),
+                    fps=self.video_fps,
+                )
+
+    def _create_env_fns(self) -> list:
+        def _create_env(env_kwargs: dict) -> gym.Env:
+            if self.env_name == "franka_kitchen":
+                # don't reset before creating vector env
+                env = KitchenBase(frame_height=224, frame_width=224, **env_kwargs)
+                if self.decomp_method is not None:
+                    env.COMPLETE_IN_ANY_ORDER = True
+                    env.TERMINATE_ON_WRONG_COMPLETE = False
+                else:
+                    env.COMPLETE_IN_ANY_ORDER = False
+                    env.TERMINATE_ON_WRONG_COMPLETE = True
+            else:
+                raise NotImplementedError(self.env_name)
+            env = InferenceWrapper(env, **self.inference_kwargs)
+            if self.save_video:
+                env = VisualizeWrapper(
+                    env, add_goal=True, add_debug_text=self.add_debug_text
+                )
+                env.recording = True
+            return env
+
+        env_kwargs = dict(
+            max_horizon=self.max_horizon,
+            gpu_id=int(torch.device(self.device).index),
+        )
+
+        return [lambda: _create_env(env_kwargs) for _ in range(self.num_envs)]
+
+    def close(self):
+        self.batched_env.close()
+
+    @staticmethod
+    def generate_rollout_list(num_envs: int, num_tasks: int):
+        assert num_envs <= num_tasks
+        tasks = []
+        k = 0
+        for i in range(num_tasks):
+            sub_list = []
+            j = 0
+            while j < num_envs and k < num_tasks:
+                sub_list.append(j)
+                j += 1
+                k += 1
+            tasks.append(sub_list)
+            if k >= num_tasks:
+                break
+        return tasks

uvd/envs/evaluator/inference_wrapper.py

@@ -0,0 +1,435 @@
+from __future__ import annotations
+
+import random
+
+import gym
+import numpy as np
+
+import uvd.utils as U
+from uvd.envs.franka_kitchen import KitchenBase
+
+__all__ = ["InferenceWrapper"]
+
+
+class InferenceWrapper(gym.Wrapper):
+    """Set all milestones => reset => set current_embedding_distance => step =>
+    set current_embedding_distance => ...
+
+    current subgoal milestone should not be set but automatically be
+    changed after setting current embedding distance
+    """
+
+    def __init__(
+        self,
+        env: gym.Env,
+        thresh: float = 0.2,
+        hybrid: bool = False,
+        use_milestone_compressor: bool = False,
+        random_skip_ratio: float = 0.0,
+        delay_steps: int | None = 0,
+        cache_history: bool = False,
+        dummy_rtn: bool = False,
+    ):
+        super().__init__(env=env)
+        self._thresh = thresh if thresh is not None else np.inf
+        self._random_skip_ratio = random_skip_ratio
+        self._cur_milestone_idx = None
+        self._delay_steps = delay_steps or 0
+        self._n_steps_after_pass_thresh = -1
+
+        self.num_milestones = None
+        self._milestones = None
+        self._current_milestone = None
+        self._current_obs_embedding = None
+        self._achieved_goals = 0
+        self._current_embedding_distance = None
+        self._unnormalized_embedding_distance = None
+        self.milestone_distances = None
+
+        self._no_robot_milestones = None
+        self._current_no_robot_milestone = None
+
+        self._milestone_embeddings = None
+        self._current_milestone_embedding = None
+
+        self._subgoal_init_distance = None
+
+        self.hybrid_phase = 0 if hybrid else None
+        self._cur_rgb_frame = None
+        # U.rank_zero_print(f"{use_milestone_compressor=}", color="blue")
+        self.use_milestone_compressor = use_milestone_compressor
+        self._n_steps = 0
+        self._n_steps_this_milestone = 0
+        self._milestone_indices = None
+        self.cache_history = cache_history
+        if cache_history:
+            self.cached_obs = []
+            self.cached_prev_milestones = []
+        else:
+            self.cached_obs = None
+            self.cached_prev_milestones = None
+        self._has_dummy_reset = True
+
+        if dummy_rtn:
+            self.dummy_rtn = []
+        else:
+            self.dummy_rtn = None
+
+    def dummy_step(self) -> tuple:
+        self._n_steps += 1
+        return self.dummy_rtn
+
+    def step(self, action: np.ndarray) -> tuple:
+        if (
+            self.dummy_rtn is not None
+            and len(self.dummy_rtn) > 0
+            and self._n_steps + 1 >= self.env.max_horizon
+        ):
+            o, r, _, i = self.dummy_rtn
+            self.dummy_rtn = []
+            return o, r, True, i
+        elif self.dummy_rtn is not None and len(self.dummy_rtn) > 0:
+            return self.dummy_step()
+
+        o, r, d, i = super().step(action)
+        self._n_steps += 1
+        self._n_steps_this_milestone += 1
+        o = self.proc_obs(o)
+        self._cur_rgb_frame = o.get("rgb", None)
+        # switch delay triggered
+        if self._n_steps_after_pass_thresh >= 0:
+            self._n_steps_after_pass_thresh += 1
+
+        if d and self.dummy_rtn is not None:
+            self.dummy_rtn = (o, r, False, i)
+            d = self._n_steps >= self.env.max_horizon
+
+        if d:
+            # force set every episode
+            self.milestones = None
+            if self.dummy_rtn is not None:
+                self.dummy_rtn = []
+        return o, r, d, i
+
+    def reset(self, **kwargs):
+        obs = super().reset(**kwargs)
+        if self.cached_obs is not None:
+            self.cached_obs = []
+            self.cached_prev_milestones = []
+        obs = self.proc_obs(obs)
+        self._achieved_goals = 0
+        self._n_steps_after_pass_thresh = -1
+        self._subgoal_init_distance = None
+        self.hybrid_phase = 0 if self.hybrid_phase is not None else None
+        self._cur_rgb_frame = obs.get("rgb", None)
+        self._n_steps = 0
+        self._n_steps_this_milestone = 0
+        self._has_dummy_reset = False
+        return obs
+
+    def proc_obs(self, obs: dict):
+        if (
+            self.cached_obs is not None and not self._has_dummy_reset
+        ):  # in case dummy reset call
+            # this fn should be call only once for every step and reset
+            self.cached_obs.append(obs)
+            # steps, ...
+            obs = U.batch_observations(self.cached_obs, to_tensor=False)
+            self.cached_prev_milestones.append(self.current_milestone)
+        return obs
+
+    @property
+    def completed_goals(self):
+        return self._achieved_goals
+
+    @property
+    def milestones(self) -> np.ndarray:
+        assert (
+            self._milestones is not None
+        ), f"must set before using, step: {self.env.episode_length}"
+        return self._milestones
+
+    @milestones.setter
+    def milestones(self, milestones: np.ndarray | None):
+        """Set outside.
+
+        can be rgb or embedding
+        """
+        if milestones is not None:
+            self._milestones = milestones
+            self.num_milestones = len(milestones)
+            # set initial sub goal
+            self._current_milestone = self._milestones[0]
+        else:
+            self._milestones = None
+            self._current_milestone = None
+            self._subgoal_init_distance = None
+
+    @property
+    def milestone_embeddings(self) -> np.ndarray:
+        if self.milestones.ndim == 2:
+            # preprocessed embedding already
+            return self.milestones
+        assert self._milestone_embeddings is not None
+        return self._milestone_embeddings
+
+    @milestone_embeddings.setter
+    def milestone_embeddings(self, embeddings: np.ndarray | None):
+        if embeddings is not None:
+            assert embeddings.ndim == 2, embeddings.shape
+            self._milestone_embeddings = embeddings
+            self._current_milestone_embedding = self.milestone_embeddings[0]
+        else:
+            self._milestone_embeddings = self._current_milestone_embedding = None
+
+    @property
+    def no_robot_milestones(self) -> np.ndarray:
+        return self._no_robot_milestones
+
+    @no_robot_milestones.setter
+    def no_robot_milestones(self, milestones: np.ndarray | None):
+        """If set (outside), sub-goal switches will depend on this, while the
+        rollout inference is still based on the goal embed with arm."""
+        if milestones is not None:
+            self._no_robot_milestones = milestones.copy()
+            self._current_no_robot_milestone = milestones[0]
+        else:
+            self._no_robot_milestones = None
+            self._current_no_robot_milestone = None
+            self._subgoal_init_distance = None
+
+    @property
+    def cur_milestone_idx(self) -> int | None:
+        return self._cur_milestone_idx
+
+    @cur_milestone_idx.setter
+    def cur_milestone_idx(self, idx: int | None):
+        self._cur_milestone_idx = idx
+
+    def maybe_random_skip(self, cur, all_cur) -> np.ndarray:
+        if self._random_skip_ratio > 0:
+            if (
+                self._cur_milestone_idx is None
+                and self._random_skip_ratio > random.random()
+            ):
+                self.cur_milestone_idx = min(
+                    self._achieved_goals + 1, self.num_milestones - 1
+                )
+            elif self._cur_milestone_idx is None:
+                self.cur_milestone_idx = min(
+                    self._achieved_goals, self.num_milestones - 1
+                )
+            assert (
+                self.cur_milestone_idx
+                in [
+                    self._achieved_goals,
+                    self._achieved_goals + 1,
+                    self.num_milestones - 1,
+                ]
+                and self.cur_milestone_idx < self.num_milestones
+            ), f"WARN: {self.cur_milestone_idx}, {self._achieved_goals}"
+            return all_cur[self._cur_milestone_idx]
+        self._cur_milestone_idx = min(self._achieved_goals, self.num_milestones - 1)
+        return cur
+
+    @property
+    def current_milestone(self) -> np.ndarray:
+        if self.use_milestone_compressor:
+            return self.milestone_embeddings
+        # current sub-goal, switch automatically when set embedding distance
+        assert self._current_milestone is not None
+        return self.maybe_random_skip(self._current_milestone, self._milestones)
+
+    @property
+    def current_milestone_embedding(self) -> np.ndarray:
+        if self.current_milestone.ndim == 1:
+            return self.current_milestone
+        assert self._current_milestone_embedding is not None
+        return self.maybe_random_skip(
+            self._current_milestone_embedding, self._milestone_embeddings
+        )
+
+    def get_current_milestone_embedding_without_skip(self) -> np.ndarray:
+        if self.current_milestone.ndim == 1:
+            return self._current_milestone
+        assert self._current_milestone_embedding is not None
+        return self._current_milestone_embedding
+
+    @property
+    def current_no_robot_milestone(self) -> np.ndarray | None:
+        if self._current_no_robot_milestone is not None:
+            assert (
+                self._current_no_robot_milestone.ndim == 1
+            ), self.current_milestone.shape
+        # no need for random skip
+        return self._current_no_robot_milestone
+
+    @property
+    def current_embedding_distance(self) -> np.ndarray:
+        return self._current_embedding_distance
+
+    @current_embedding_distance.setter
+    def current_embedding_distance(self, dist: float | None):
+        if dist is None or self._milestones is None:
+            self._current_embedding_distance = dist
+            return
+        switch = False
+        if dist is not None and dist < self._thresh:
+            switch = True
+            if self._delay_steps > 0 and self._n_steps_after_pass_thresh == -1:
+                switch = False
+                # trigger the delay
+                self._n_steps_after_pass_thresh = 0
+        if 0 < self._delay_steps <= self._n_steps_after_pass_thresh:
+            assert self._delay_steps == self._n_steps_after_pass_thresh
+            switch = True
+            self._n_steps_after_pass_thresh = -1
+
+        current_milestone_budget = self.current_milestone_budget()
+        if switch and current_milestone_budget is not None:
+            if current_milestone_budget > self._n_steps_this_milestone + 3:
+                switch = False  # make sure not switch too fast by noise of the embed
+        elif not switch and current_milestone_budget is not None:
+            if current_milestone_budget < self._n_steps_this_milestone - 3:
+                switch = True  # in case not switch due to noise of the embed
+
+        if switch:
+            self._n_steps_this_milestone = 0
+            if self.hybrid_phase is not None:
+                self.hybrid_phase += 1
+            # switch milestones
+            self._achieved_goals += 1
+            self._achieved_goals = min(self._achieved_goals, self.num_milestones)
+            # self._achieved_goals = self.env.num_goals_achieved
+            self._cur_milestone_idx = min(
+                self._achieved_goals, self.num_milestones - 1
+            )  # next idx
+            self._current_milestone = self._milestones[self._cur_milestone_idx]
+            if self._current_milestone_embedding is not None:
+                # switch no next milestone embedding for conditioning
+                self._current_milestone_embedding = self.milestone_embeddings[
+                    self._cur_milestone_idx
+                ]
+            if self.current_no_robot_milestone is not None:
+                # switch to next milestone for calculating embed distance
+                assert len(self.milestones) == len(self.no_robot_milestones)
+                self._current_no_robot_milestone = self.no_robot_milestones[
+                    self._cur_milestone_idx
+                ]
+                if self.milestone_distances is not None:
+                    self._subgoal_init_distance = self.milestone_distances[
+                        min(self._achieved_goals - 1, self.num_milestones - 2)
+                    ]
+                else:
+                    # use no skipping milestone for distance check
+                    # self._subgoal_init_distance = np.linalg.norm(
+                    #     self.current_obs_embedding - self._current_no_robot_milestone
+                    # )
+                    self._subgoal_init_distance = None  # set next step back
+            else:
+                if self.milestone_distances is not None:
+                    # self._subgoal_init_distance = self.milestone_distances[
+                    #     min(self._achieved_goals - 1, self.num_milestones - 2)
+                    # ]
+                    self._subgoal_init_distance = None
+                else:
+                    # use no skipping milestone for distance check
+                    self._subgoal_init_distance = np.linalg.norm(
+                        self.current_obs_embedding - self._current_milestone_embedding
+                    )
+            self._current_embedding_distance = 1.0
+        else:
+            self._current_embedding_distance = dist
+
+    @property
+    def current_obs_embedding(self) -> np.ndarray:
+        assert self._current_obs_embedding is not None
+        return self._current_obs_embedding
+
+    @current_obs_embedding.setter
+    def current_obs_embedding(self, embedding: np.ndarray):
+        """Set outside, should be embedding already for calculating embed
+        distance."""
+        assert embedding.ndim == 1, embedding.shape
+        self._current_obs_embedding = embedding
+        if self.hybrid_phase is None:
+            # calculate embed dist automatically after setting embedding
+            # use no skipping milestone for distance check
+            cur_milestone = (
+                self.current_no_robot_milestone
+                if self.current_no_robot_milestone is not None
+                else self.get_current_milestone_embedding_without_skip()
+            )
+        else:
+            cur_milestone = (
+                self.get_current_milestone_embedding_without_skip()
+                if self.hybrid_phase % 2 == 0
+                else self.current_no_robot_milestone
+            )
+        assert (
+            cur_milestone.ndim == embedding.ndim == 1
+        ), f"{cur_milestone.shape} != {embedding.shape} != (n,)"
+        embed_distance = np.linalg.norm(embedding - cur_milestone)
+        self._unnormalized_embedding_distance = embed_distance
+
+        if self._subgoal_init_distance is None:
+            # first step after reset here
+            self._subgoal_init_distance = max(embed_distance, 1e-14)
+        # property setter
+        self.current_embedding_distance = embed_distance / self._subgoal_init_distance
+
+    @property
+    def reset_states(self) -> dict:
+        return self.env.reset_states
+
+    @reset_states.setter
+    def reset_states(self, state: dict):
+        self.env.reset_states = state
+
+    @property
+    def task_name(self) -> str:
+        return self.env.task_name
+
+    @property
+    def metrics(self) -> dict:
+        self.env: KitchenBase
+        metrics = self.env.metrics
+        achieved = (
+            self.completed_goals if not metrics["success"] else self.num_milestones
+        )
+        achieved_percentage = achieved / self.num_milestones
+        metrics["achieved_milestones"] = achieved_percentage
+        return metrics
+
+    @property
+    def current_no_robot_frame(self) -> np.ndarray:
+        if self.hybrid_phase is not None and self.hybrid_phase % 2 == 0:
+            # hack for hybrid return rgb again
+            assert self._cur_rgb_frame is not None
+            return self._cur_rgb_frame
+        else:
+            return self.env.current_no_robot_frame
+
+    @property
+    def current_rgb_frame(self) -> np.ndarray:
+        return self._cur_rgb_frame
+
+    @property
+    def milestone_indices(self):
+        return self._milestone_indices
+
+    @milestone_indices.setter
+    def milestone_indices(self, milestone_indices: np.ndarray | None):
+        self._milestone_indices = milestone_indices
+
+    def current_milestone_budget(self) -> int | None:
+        if self.cur_milestone_idx is None:
+            self.cur_milestone_idx = 0
+        if self.milestone_indices is None:
+            return None
+        return (
+            self.milestone_indices[self.cur_milestone_idx] + 1  # indices start from 0
+            if self.cur_milestone_idx == 0
+            else self.milestone_indices[self.cur_milestone_idx]
+            - self.milestone_indices[self.cur_milestone_idx - 1]
+        )

uvd/envs/evaluator/vec_envs/vec_env.py

@@ -0,0 +1,398 @@
+"""Modified from `tianshou`"""
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+import gym
+import numpy as np
+
+from .workers import EnvWorker, RayEnvWorker, SubprocEnvWorker
+
+GYM_RESERVED_KEYS = [
+    "metadata",
+    "reward_range",
+    "spec",
+    "action_space",
+    "observation_space",
+]
+
+
+class BaseVectorEnv(object):
+    """Base class for vectorized environments.
+
+    Usage:
+    ::
+
+        env_num = 8
+        envs = DummyVectorEnv([lambda: gym.make(task) for _ in range(env_num)])
+        assert len(envs) == env_num
+
+    It accepts a list of environment generators. In other words, an environment
+    generator ``efn`` of a specific task means that ``efn()`` returns the
+    environment of the given task, for example, ``gym.make(task)``.
+
+    All of the VectorEnv must inherit :class:`~tianshou.env.BaseVectorEnv`.
+    Here are some other usages:
+    ::
+
+        envs.seed(2)  # which is equal to the next line
+        envs.seed([2, 3, 4, 5, 6, 7, 8, 9])  # set specific seed for each env
+        obs = envs.reset()  # reset all environments
+        obs = envs.reset([0, 5, 7])  # reset 3 specific environments
+        obs, rew, done, info = envs.step([1] * 8)  # step synchronously
+        envs.render()  # render all environments
+        envs.close()  # close all environments
+
+    .. warning::
+
+        If you use your own environment, please make sure the ``seed`` method
+        is set up properly, e.g.,
+        ::
+
+            def seed(self, seed):
+                np.random.seed(seed)
+
+        Otherwise, the outputs of these envs may be the same with each other.
+
+    :param env_fns: a list of callable envs, ``env_fns[i]()`` generates the i-th env.
+    :param worker_fn: a callable worker, ``worker_fn(env_fns[i])`` generates a
+        worker which contains the i-th env.
+    :param int wait_num: use in asynchronous simulation if the time cost of
+        ``env.step`` varies with time and synchronously waiting for all
+        environments to finish a step is time-wasting. In that case, we can
+        return when ``wait_num`` environments finish a step and keep on
+        simulation in these environments. If ``None``, asynchronous simulation
+        is disabled; else, ``1 <= wait_num <= env_num``.
+    :param float timeout: use in asynchronous simulation same as above, in each
+        vectorized step it only deal with those environments spending time
+        within ``timeout`` seconds.
+    """
+
+    def __init__(
+        self,
+        env_fns: List[Callable[[], gym.Env]],
+        worker_fn: Callable[[Callable[[], gym.Env]], EnvWorker],
+        wait_num: Optional[int] = None,
+        timeout: Optional[float] = None,
+    ) -> None:
+        self._env_fns = env_fns
+        # A VectorEnv contains a pool of EnvWorkers, which corresponds to
+        # interact with the given envs (one worker <-> one env).
+        self.workers = [worker_fn(fn) for fn in env_fns]
+        self.worker_class = type(self.workers[0])
+        assert issubclass(self.worker_class, EnvWorker)
+        assert all([isinstance(w, self.worker_class) for w in self.workers])
+
+        self.env_num = len(env_fns)
+        self.wait_num = wait_num or len(env_fns)
+        assert (
+            1 <= self.wait_num <= len(env_fns)
+        ), f"wait_num should be in [1, {len(env_fns)}], but got {wait_num}"
+        self.timeout = timeout
+        assert (
+            self.timeout is None or self.timeout > 0
+        ), f"timeout is {timeout}, it should be positive if provided!"
+        self.is_async = self.wait_num != len(env_fns) or timeout is not None
+        self.waiting_conn: List[EnvWorker] = []
+        # environments in self.ready_id is actually ready
+        # but environments in self.waiting_id are just waiting when checked,
+        # and they may be ready now, but this is not known until we check it
+        # in the step() function
+        self.waiting_id: List[int] = []
+        # all environments are ready in the beginning
+        self.ready_id = list(range(self.env_num))
+        self.is_closed = False
+
+    def _assert_is_not_closed(self) -> None:
+        assert (
+            not self.is_closed
+        ), f"Methods of {self.__class__.__name__} cannot be called after close."
+
+    def __len__(self) -> int:
+        """Return len(self), which is the number of environments."""
+        return self.env_num
+
+    def __getattribute__(self, key: str) -> Any:
+        """Switch the attribute getter depending on the key.
+
+        Any class who inherits ``gym.Env`` will inherit some attributes,
+        like ``action_space``. However, we would like the attribute
+        lookup to go straight into the worker (in fact, this vector
+        env's action_space is always None).
+        """
+        if key in GYM_RESERVED_KEYS:  # reserved keys in gym.Env
+            return self.get_env_attr(key)
+        else:
+            return super().__getattribute__(key)
+
+    def get_env_attr(
+        self,
+        key: str,
+        id: Optional[Union[int, List[int], np.ndarray]] = None,
+    ) -> List[Any]:
+        """Get an attribute from the underlying environments.
+
+        If id is an int, retrieve the attribute denoted by key from the
+        environment underlying the worker at index id. The result is
+        returned as a list with one element. Otherwise, retrieve the
+        attribute for all workers at indices id and return a list that
+        is ordered correspondingly to id.
+
+        :param str key: The key of the desired attribute.
+        :param id: Indice(s) of the desired worker(s). Default to None
+            for all env_id.
+        :return list: The list of environment attributes.
+        """
+        self._assert_is_not_closed()
+        id = self._wrap_id(id)
+        if self.is_async:
+            self._assert_id(id)
+
+        return [self.workers[j].get_env_attr(key) for j in id]
+
+    def set_env_attr(
+        self,
+        key: str,
+        value: Any,
+        id: Optional[Union[int, List[int], np.ndarray]] = None,
+        diff_value: bool = False,
+    ) -> None:
+        """Set an attribute in the underlying environments.
+
+        If id is an int, set the attribute denoted by key from the
+        environment underlying the worker at index id to value.
+        Otherwise, set the attribute for all workers at indices id.
+
+        :param str key: The key of the desired attribute.
+        :param Any value: The new value of the attribute.
+        :param id: Indice(s) of the desired worker(s). Default to None
+            for all env_id.
+        """
+        self._assert_is_not_closed()
+        id = self._wrap_id(id)
+        if diff_value:
+            assert len(value) == len(id)
+        if self.is_async:
+            self._assert_id(id)
+        for i, j in enumerate(id):
+            if diff_value:
+                self.workers[j].set_env_attr(key, value[i])
+            else:
+                self.workers[j].set_env_attr(key, value)
+
+    def _wrap_id(
+        self,
+        id: Optional[Union[int, List[int], np.ndarray]] = None,
+    ) -> Union[List[int], np.ndarray]:
+        if id is None:
+            return list(range(self.env_num))
+        return [id] if np.isscalar(id) else id  # type: ignore
+
+    def _assert_id(self, id: Union[List[int], np.ndarray]) -> None:
+        for i in id:
+            assert (
+                i not in self.waiting_id
+            ), f"Cannot interact with environment {i} which is stepping now."
+            assert (
+                i in self.ready_id
+            ), f"Can only interact with ready environments {self.ready_id}."
+
+    # TODO: compatible issue with reset -> (obs, info)
+    def reset(
+        self, id: Optional[Union[int, List[int], np.ndarray]] = None
+    ) -> np.ndarray:
+        """Reset the state of some envs and return initial observations.
+
+        If id is None, reset the state of all the environments and
+        return initial observations, otherwise reset the specific
+        environments with the given id, either an int or a list.
+        """
+        self._assert_is_not_closed()
+        id = self._wrap_id(id)
+        if self.is_async:
+            self._assert_id(id)
+        # send(None) == reset() in worker
+        for i in id:
+            self.workers[i].send(None)
+        obs_list = [self.workers[i].recv() for i in id]
+        try:
+            obs = np.stack(obs_list)
+        except ValueError:  # different len(obs)
+            obs = np.array(obs_list, dtype=object)
+        return obs
+
+    def step(
+        self,
+        action: np.ndarray,
+        id: Optional[Union[int, List[int], np.ndarray]] = None,
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        """Run one timestep of some environments' dynamics.
+
+        If id is None, run one timestep of all the environments’ dynamics;
+        otherwise run one timestep for some environments with given id,  either
+        an int or a list. When the end of episode is reached, you are
+        responsible for calling reset(id) to reset this environment’s state.
+
+        Accept a batch of action and return a tuple (batch_obs, batch_rew,
+        batch_done, batch_info) in numpy format.
+
+        :param numpy.ndarray action: a batch of action provided by the agent.
+
+        :return: A tuple including four items:
+
+            * ``obs`` a numpy.ndarray, the agent's observation of current environments
+            * ``rew`` a numpy.ndarray, the amount of rewards returned after \
+                previous actions
+            * ``done`` a numpy.ndarray, whether these episodes have ended, in \
+                which case further step() calls will return undefined results
+            * ``info`` a numpy.ndarray, contains auxiliary diagnostic \
+                information (helpful for debugging, and sometimes learning)
+
+        For the async simulation:
+
+        Provide the given action to the environments. The action sequence
+        should correspond to the ``id`` argument, and the ``id`` argument
+        should be a subset of the ``env_id`` in the last returned ``info``
+        (initially they are env_ids of all the environments). If action is
+        None, fetch unfinished step() calls instead.
+        """
+        self._assert_is_not_closed()
+        id = self._wrap_id(id)
+        if not self.is_async:
+            assert len(action) == len(id)
+            for i, j in enumerate(id):
+                self.workers[j].send(action[i])
+            result = []
+            for j in id:
+                obs, rew, done, info = self.workers[j].recv()
+                info["env_id"] = j
+                result.append((obs, rew, done, info))
+        else:
+            if action is not None:
+                self._assert_id(id)
+                assert len(action) == len(id)
+                for act, env_id in zip(action, id):
+                    self.workers[env_id].send(act)
+                    self.waiting_conn.append(self.workers[env_id])
+                    self.waiting_id.append(env_id)
+                self.ready_id = [x for x in self.ready_id if x not in id]
+            ready_conns: List[EnvWorker] = []
+            while not ready_conns:
+                ready_conns = self.worker_class.wait(
+                    self.waiting_conn, self.wait_num, self.timeout
+                )
+            result = []
+            for conn in ready_conns:
+                waiting_index = self.waiting_conn.index(conn)
+                self.waiting_conn.pop(waiting_index)
+                env_id = self.waiting_id.pop(waiting_index)
+                obs, rew, done, info = conn.recv()
+                info["env_id"] = env_id
+                result.append((obs, rew, done, info))
+                self.ready_id.append(env_id)
+        obs_list, rew_list, done_list, info_list = zip(*result)
+        try:
+            obs_stack = np.stack(obs_list)
+        except ValueError:  # different len(obs)
+            obs_stack = np.array(obs_list, dtype=object)
+        rew_stack, done_stack, info_stack = map(
+            np.stack, [rew_list, done_list, info_list]
+        )
+        return obs_stack, rew_stack, done_stack, info_stack
+
+    def seed(
+        self,
+        seed: Optional[Union[int, List[int]]] = None,
+    ) -> List[Optional[List[int]]]:
+        """Set the seed for all environments.
+
+        Accept ``None``, an int (which will extend ``i`` to
+        ``[i, i + 1, i + 2, ...]``) or a list.
+
+        :return: The list of seeds used in this env's random number generators.
+            The first value in the list should be the "main" seed, or the value
+            which a reproducer pass to "seed".
+        """
+        self._assert_is_not_closed()
+        seed_list: Union[List[None], List[int]]
+        if seed is None:
+            seed_list = [seed] * self.env_num
+        elif isinstance(seed, int):
+            seed_list = [seed + i for i in range(self.env_num)]
+        else:
+            seed_list = seed
+        return [w.seed(s) for w, s in zip(self.workers, seed_list)]
+
+    def render(self, **kwargs: Any) -> List[Any]:
+        """Render all of the environments."""
+        self._assert_is_not_closed()
+        if self.is_async and len(self.waiting_id) > 0:
+            raise RuntimeError(
+                f"Environments {self.waiting_id} are still stepping, cannot "
+                "render them now."
+            )
+        return [w.render(**kwargs) for w in self.workers]
+
+    def close(self) -> None:
+        """Close all of the environments.
+
+        This function will be called only once (if not, it will be
+        called during garbage collected). This way, ``close`` of all
+        workers can be assured.
+        """
+        self._assert_is_not_closed()
+        for w in self.workers:
+            w.close()
+        self.is_closed = True
+
+
+class SubprocVectorEnv(BaseVectorEnv):
+    """Vectorized environment wrapper based on subprocess.
+
+    .. seealso::
+
+        Please refer to :class:`~tianshou.env.BaseVectorEnv` for other APIs' usage.
+    """
+
+    def __init__(self, env_fns: List[Callable[[], gym.Env]], **kwargs: Any) -> None:
+        def worker_fn(fn: Callable[[], gym.Env]) -> SubprocEnvWorker:
+            return SubprocEnvWorker(fn, share_memory=False)
+
+        super().__init__(env_fns, worker_fn, **kwargs)
+
+
+class ShmemVectorEnv(BaseVectorEnv):
+    """Optimized SubprocVectorEnv with shared buffers to exchange observations.
+
+    ShmemVectorEnv has exactly the same API as SubprocVectorEnv.
+
+    .. seealso::
+
+        Please refer to :class:`~tianshou.env.BaseVectorEnv` for other APIs' usage.
+    """
+
+    def __init__(self, env_fns: List[Callable[[], gym.Env]], **kwargs: Any) -> None:
+        def worker_fn(fn: Callable[[], gym.Env]) -> SubprocEnvWorker:
+            return SubprocEnvWorker(fn, share_memory=True)
+
+        super().__init__(env_fns, worker_fn, **kwargs)
+
+
+class RayVectorEnv(BaseVectorEnv):
+    """Vectorized environment wrapper based on ray.
+
+    This is a choice to run distributed environments in a cluster.
+
+    .. seealso::
+
+        Please refer to :class:`~tianshou.env.BaseVectorEnv` for other APIs' usage.
+    """
+
+    def __init__(self, env_fns: List[Callable[[], gym.Env]], **kwargs: Any) -> None:
+        try:
+            import ray
+        except ImportError as exception:
+            raise ImportError(
+                "Please install ray to support RayVectorEnv: pip install ray"
+            ) from exception
+        if not ray.is_initialized():
+            ray.init()
+        super().__init__(env_fns, RayEnvWorker, **kwargs)

uvd/envs/evaluator/vec_envs/workers.py

@@ -0,0 +1,409 @@
+import ctypes
+import time
+import warnings
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+from multiprocessing import Array, Pipe, connection
+from multiprocessing.context import Process
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+import cloudpickle
+import gym
+import numpy as np
+
+try:
+    import ray
+except ImportError:
+    pass
+
+
+warnings.simplefilter("once", DeprecationWarning)
+
+_NP_TO_CT = {
+    np.bool_: ctypes.c_bool,
+    np.uint8: ctypes.c_uint8,
+    np.uint16: ctypes.c_uint16,
+    np.uint32: ctypes.c_uint32,
+    np.uint64: ctypes.c_uint64,
+    np.int8: ctypes.c_int8,
+    np.int16: ctypes.c_int16,
+    np.int32: ctypes.c_int32,
+    np.int64: ctypes.c_int64,
+    np.float32: ctypes.c_float,
+    np.float64: ctypes.c_double,
+}
+
+
+def deprecation(msg: str) -> None:
+    """Deprecation warning wrapper."""
+    warnings.warn(msg, category=DeprecationWarning, stacklevel=2)
+
+
+class CloudpickleWrapper(object):
+    """A cloudpickle wrapper used in SubprocVectorEnv."""
+
+    def __init__(self, data: Any) -> None:
+        self.data = data
+
+    def __getstate__(self) -> str:
+        return cloudpickle.dumps(self.data)
+
+    def __setstate__(self, data: str) -> None:
+        self.data = cloudpickle.loads(data)
+
+
+class EnvWorker(ABC):
+    """An abstract worker for an environment."""
+
+    def __init__(self, env_fn: Callable[[], gym.Env]) -> None:
+        self._env_fn = env_fn
+        self.is_closed = False
+        self.result: Union[
+            Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], np.ndarray
+        ]
+        self.action_space = self.get_env_attr("action_space")  # noqa: B009
+        self.is_reset = False
+
+    @abstractmethod
+    def get_env_attr(self, key: str) -> Any:
+        pass
+
+    @abstractmethod
+    def set_env_attr(self, key: str, value: Any) -> None:
+        pass
+
+    def send(self, action: Optional[np.ndarray]) -> None:
+        """Send action signal to low-level worker.
+
+        When action is None, it indicates sending "reset" signal;
+        otherwise it indicates "step" signal. The paired return value
+        from "recv" function is determined by such kind of different
+        signal.
+        """
+        if hasattr(self, "send_action"):
+            deprecation(
+                "send_action will soon be deprecated. "
+                "Please use send and recv for your own EnvWorker."
+            )
+            if action is None:
+                self.is_reset = True
+                self.result = self.reset()
+            else:
+                self.is_reset = False
+                self.send_action(action)  # type: ignore
+
+    def recv(
+        self,
+    ) -> Union[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], np.ndarray]:
+        """Receive result from low-level worker.
+
+        If the last "send" function sends a NULL action, it only returns
+        a single observation; otherwise it returns a tuple of (obs, rew,
+        done, info).
+        """
+        if hasattr(self, "get_result"):
+            deprecation(
+                "get_result will soon be deprecated. "
+                "Please use send and recv for your own EnvWorker."
+            )
+            if not self.is_reset:
+                self.result = self.get_result()  # type: ignore
+        return self.result
+
+    def reset(self) -> np.ndarray:
+        self.send(None)
+        return self.recv()  # type: ignore
+
+    def step(
+        self, action: np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        """Perform one timestep of the environment's dynamic.
+
+        "send" and "recv" are coupled in sync simulation, so users only
+        call "step" function. But they can be called separately in async
+        simulation, i.e. someone calls "send" first, and calls "recv"
+        later.
+        """
+        self.send(action)
+        return self.recv()  # type: ignore
+
+    @staticmethod
+    def wait(
+        workers: List["EnvWorker"], wait_num: int, timeout: Optional[float] = None
+    ) -> List["EnvWorker"]:
+        """Given a list of workers, return those ready ones."""
+        raise NotImplementedError
+
+    def seed(self, seed: Optional[int] = None) -> Optional[List[int]]:
+        return self.action_space.seed(seed)  # issue 299
+
+    @abstractmethod
+    def render(self, **kwargs: Any) -> Any:
+        """Render the environment."""
+        pass
+
+    @abstractmethod
+    def close_env(self) -> None:
+        pass
+
+    def close(self) -> None:
+        if self.is_closed:
+            return None
+        self.is_closed = True
+        self.close_env()
+
+
+class ShArray:
+    """Wrapper of multiprocessing Array."""
+
+    def __init__(self, dtype: np.generic, shape: Tuple[int]) -> None:
+        self.arr = Array(_NP_TO_CT[dtype.type], int(np.prod(shape)))  # type: ignore
+        self.dtype = dtype
+        self.shape = shape
+
+    def save(self, ndarray: np.ndarray) -> None:
+        assert isinstance(ndarray, np.ndarray)
+        dst = self.arr.get_obj()
+        dst_np = np.frombuffer(dst, dtype=self.dtype).reshape(self.shape)
+        np.copyto(dst_np, ndarray)
+
+    def get(self) -> np.ndarray:
+        obj = self.arr.get_obj()
+        return np.frombuffer(obj, dtype=self.dtype).reshape(self.shape)
+
+
+def _setup_buf(space: gym.Space) -> Union[dict, tuple, ShArray]:
+    if isinstance(space, gym.spaces.Dict):
+        assert isinstance(space.spaces, OrderedDict)
+        return {k: _setup_buf(v) for k, v in space.spaces.items()}
+    elif isinstance(space, gym.spaces.Tuple):
+        assert isinstance(space.spaces, tuple)
+        return tuple([_setup_buf(t) for t in space.spaces])
+    else:
+        return ShArray(space.dtype, space.shape)  # type: ignore
+
+
+def _worker(
+    parent: connection.Connection,
+    p: connection.Connection,
+    env_fn_wrapper: CloudpickleWrapper,
+    obs_bufs: Optional[Union[dict, tuple, ShArray]] = None,
+) -> None:
+    def _encode_obs(
+        obs: Union[dict, tuple, np.ndarray], buffer: Union[dict, tuple, ShArray]
+    ) -> None:
+        if isinstance(obs, np.ndarray) and isinstance(buffer, ShArray):
+            buffer.save(obs)
+        elif isinstance(obs, tuple) and isinstance(buffer, tuple):
+            for o, b in zip(obs, buffer):
+                _encode_obs(o, b)
+        elif isinstance(obs, dict) and isinstance(buffer, dict):
+            for k in obs.keys():
+                _encode_obs(obs[k], buffer[k])
+        return None
+
+    parent.close()
+    env = env_fn_wrapper.data()
+    try:
+        while True:
+            try:
+                cmd, data = p.recv()
+            except EOFError:  # the pipe has been closed
+                p.close()
+                break
+            if cmd == "step":
+                if data is None:  # reset
+                    obs = env.reset()
+                else:
+                    obs, reward, done, info = env.step(data)
+                if obs_bufs is not None:
+                    _encode_obs(obs, obs_bufs)
+                    obs = None
+                if data is None:
+                    p.send(obs)
+                else:
+                    p.send((obs, reward, done, info))
+            elif cmd == "close":
+                p.send(env.close())
+                p.close()
+                break
+            elif cmd == "render":
+                p.send(env.render(**data) if hasattr(env, "render") else None)
+            elif cmd == "seed":
+                p.send(env.seed(data) if hasattr(env, "seed") else None)
+            elif cmd == "getattr":
+                p.send(getattr(env, data) if hasattr(env, data) else None)
+            elif cmd == "setattr":
+                setattr(env, data["key"], data["value"])
+            else:
+                p.close()
+                raise NotImplementedError
+    except KeyboardInterrupt:
+        p.close()
+
+
+class SubprocEnvWorker(EnvWorker):
+    """Subprocess worker used in SubprocVectorEnv and ShmemVectorEnv."""
+
+    def __init__(
+        self, env_fn: Callable[[], gym.Env], share_memory: bool = False
+    ) -> None:
+        self.parent_remote, self.child_remote = Pipe()
+        self.share_memory = share_memory
+        self.buffer: Optional[Union[dict, tuple, ShArray]] = None
+        if self.share_memory:
+            dummy = env_fn()
+            obs_space = dummy.observation_space
+            dummy.close()
+            del dummy
+            self.buffer = _setup_buf(obs_space)
+        args = (
+            self.parent_remote,
+            self.child_remote,
+            CloudpickleWrapper(env_fn),
+            self.buffer,
+        )
+        self.process = Process(target=_worker, args=args, daemon=True)
+        self.process.start()
+        self.child_remote.close()
+        self.is_reset = False
+        super().__init__(env_fn)
+
+    def get_env_attr(self, key: str) -> Any:
+        self.parent_remote.send(["getattr", key])
+        return self.parent_remote.recv()
+
+    def set_env_attr(self, key: str, value: Any) -> None:
+        self.parent_remote.send(["setattr", {"key": key, "value": value}])
+
+    def _decode_obs(self) -> Union[dict, tuple, np.ndarray]:
+        def decode_obs(
+            buffer: Optional[Union[dict, tuple, ShArray]]
+        ) -> Union[dict, tuple, np.ndarray]:
+            if isinstance(buffer, ShArray):
+                return buffer.get()
+            elif isinstance(buffer, tuple):
+                return tuple([decode_obs(b) for b in buffer])
+            elif isinstance(buffer, dict):
+                return {k: decode_obs(v) for k, v in buffer.items()}
+            else:
+                raise NotImplementedError
+
+        return decode_obs(self.buffer)
+
+    @staticmethod
+    def wait(  # type: ignore
+        workers: List["SubprocEnvWorker"],
+        wait_num: int,
+        timeout: Optional[float] = None,
+    ) -> List["SubprocEnvWorker"]:
+        remain_conns = conns = [x.parent_remote for x in workers]
+        ready_conns: List[connection.Connection] = []
+        remain_time, t1 = timeout, time.time()
+        while len(remain_conns) > 0 and len(ready_conns) < wait_num:
+            if timeout:
+                remain_time = timeout - (time.time() - t1)
+                if remain_time <= 0:
+                    break
+            # connection.wait hangs if the list is empty
+            new_ready_conns = connection.wait(remain_conns, timeout=remain_time)
+            ready_conns.extend(new_ready_conns)  # type: ignore
+            remain_conns = [conn for conn in remain_conns if conn not in ready_conns]
+        return [workers[conns.index(con)] for con in ready_conns]
+
+    def send(self, action: Optional[np.ndarray]) -> None:
+        self.parent_remote.send(["step", action])
+
+    def recv(
+        self,
+    ) -> Union[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], np.ndarray]:
+        result = self.parent_remote.recv()
+        if isinstance(result, tuple):
+            obs, rew, done, info = result
+            if self.share_memory:
+                obs = self._decode_obs()
+            return obs, rew, done, info
+        else:
+            obs = result
+            if self.share_memory:
+                obs = self._decode_obs()
+            return obs
+
+    def seed(self, seed: Optional[int] = None) -> Optional[List[int]]:
+        super().seed(seed)
+        self.parent_remote.send(["seed", seed])
+        return self.parent_remote.recv()
+
+    def render(self, **kwargs: Any) -> Any:
+        self.parent_remote.send(["render", kwargs])
+        return self.parent_remote.recv()
+
+    def close_env(self) -> None:
+        try:
+            self.parent_remote.send(["close", None])
+            # mp may be deleted so it may raise AttributeError
+            self.parent_remote.recv()
+            self.process.join()
+        except (BrokenPipeError, EOFError, AttributeError):
+            pass
+        # ensure the subproc is terminated
+        self.process.terminate()
+
+
+class _SetAttrWrapper(gym.Wrapper):
+    def set_env_attr(self, key: str, value: Any) -> None:
+        setattr(self.env, key, value)
+
+    def get_env_attr(self, key: str) -> Any:
+        return getattr(self.env, key)
+
+
+class RayEnvWorker(EnvWorker):
+    """Ray worker used in RayVectorEnv."""
+
+    def __init__(self, env_fn: Callable[[], gym.Env]) -> None:
+        self.env = (
+            ray.remote(_SetAttrWrapper)
+            .options(num_cpus=0)  # type: ignore
+            .remote(env_fn())
+        )
+        super().__init__(env_fn)
+
+    def get_env_attr(self, key: str) -> Any:
+        return ray.get(self.env.get_env_attr.remote(key))
+
+    def set_env_attr(self, key: str, value: Any) -> None:
+        ray.get(self.env.set_env_attr.remote(key, value))
+
+    def reset(self) -> Any:
+        return ray.get(self.env.reset.remote())
+
+    @staticmethod
+    def wait(  # type: ignore
+        workers: List["RayEnvWorker"], wait_num: int, timeout: Optional[float] = None
+    ) -> List["RayEnvWorker"]:
+        results = [x.result for x in workers]
+        ready_results, _ = ray.wait(results, num_returns=wait_num, timeout=timeout)
+        return [workers[results.index(result)] for result in ready_results]
+
+    def send(self, action: Optional[np.ndarray]) -> None:
+        # self.action is actually a handle
+        if action is None:
+            self.result = self.env.reset.remote()
+        else:
+            self.result = self.env.step.remote(action)
+
+    def recv(
+        self,
+    ) -> Union[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], np.ndarray]:
+        return ray.get(self.result)  # type: ignore
+
+    def seed(self, seed: Optional[int] = None) -> List[int]:
+        super().seed(seed)
+        return ray.get(self.env.seed.remote(seed))
+
+    def render(self, **kwargs: Any) -> Any:
+        return ray.get(self.env.render.remote(**kwargs))
+
+    def close_env(self) -> None:
+        ray.get(self.env.close.remote())

uvd/envs/evaluator/visualize_wrapper.py

@@ -0,0 +1,271 @@
+from __future__ import annotations
+
+from collections import OrderedDict
+
+import gym
+import numpy as np
+
+import uvd.utils as U
+from uvd.envs.evaluator.inference_wrapper import InferenceWrapper
+
+__all__ = ["VisualizeWrapper"]
+
+
+class VisualizeWrapper(gym.Wrapper):
+    def __init__(
+        self, env: InferenceWrapper, add_goal: bool = True, add_debug_text: bool = True
+    ):
+        super().__init__(env=env)
+        self.env: InferenceWrapper = env
+        self.add_goal = add_goal
+        self.add_debug_text = add_debug_text
+        self._frames = []
+        self._rgb_milestones = None
+        self._current_goal_image = None
+        self._stepwise_extra_debug_texts = None
+
+        self._is_init_reset = True
+        self._recording = True
+
+    @property
+    def recording(self) -> bool:
+        return self._recording
+
+    @recording.setter
+    def recording(self, val: bool):
+        self._recording = val
+
+    @property
+    def frames(self):
+        return self._frames
+
+    def clear(self):
+        self._frames = []
+
+    @property
+    def video_length(self):
+        return len(self._frames)
+
+    def reset(self, **kwargs) -> np.ndarray:
+        obs = super().reset(**kwargs)
+        self.clear()
+        if self._is_init_reset:
+            self._is_init_reset = False
+            return obs
+        self.add_frame(obs, info=None)
+        return obs
+
+    def step(self, action: np.ndarray) -> tuple[np.ndarray, float, bool, dict]:
+        o, r, d, i = super().step(action)
+        self.add_frame(o, i)
+        self.current_embedding_distance = None
+        self.stepwise_extra_debug_texts = None
+        return o, r, d, i
+
+    def add_frame(
+        self, obs: np.ndarray | OrderedDict[str, np.ndarray], info: dict | None
+    ):
+        if not self.recording:
+            return
+        if not isinstance(obs, np.ndarray):
+            assert isinstance(obs, (dict, OrderedDict)) and "rgb" in obs, obs
+            if self.env.cached_obs is not None:
+                rgb = obs["rgb"][-1, ...].copy()
+            else:
+                rgb = obs["rgb"].copy()
+        else:
+            rgb = obs.copy()
+        if self.add_goal:
+            rgb = np.concatenate([rgb, self.current_goal_image], axis=-2)
+        if self.add_debug_text:
+            if info is not None:
+                assert (
+                    self.current_embedding_distance is not None
+                ), f"should set before each step"
+                txts = dict(
+                    complete_milestones=self.completed_goals,
+                    num_milstones=self.env.num_milestones,
+                    current_embedding_distance=float(self.current_embedding_distance),
+                )
+                try:
+                    txts.update(
+                        current_sub_task=info["current_sub_task"],
+                        step=info["episode_length"],
+                        completed_tasks=info["completed_tasks"],
+                        # success=int(info["success"]),
+                    )
+                    if "distances_left" in info["rewards"]:
+                        distances_left = info["rewards"]["distances_left"]
+                        if isinstance(distances_left, dict):
+                            for ele, dist in distances_left.items():
+                                txts[f"low_dim_distances_left/{ele}"] = float(
+                                    np.sum(dist)
+                                )
+                        else:
+                            txts["low_dim_distances_left"] = float(
+                                np.sum(distances_left)
+                            )
+                except KeyError:
+                    pass
+            else:
+                try:
+                    current_sub_task = getattr(self.env, "current_subtask")
+                    txts = dict(current_sub_task=current_sub_task, step=0)
+                except AttributeError:
+                    txts = dict(step=0)
+            if self.stepwise_extra_debug_texts is not None:
+                txts.update(self.stepwise_extra_debug_texts)
+            rgb = U.debug_texts_to_frame(frame=rgb, debug_text=txts)
+        self._frames.append(rgb)
+
+    @property
+    def current_goal_image(self):
+        assert self._current_goal_image is not None
+        if self.cur_milestone_idx is not None:
+            return self.rgb_milestones[self.cur_milestone_idx]
+        return self._current_goal_image
+
+    @current_goal_image.setter
+    def current_goal_image(self, image: np.ndarray | None):
+        if image is not None:
+            assert image.ndim == 3, image.shape
+        self._current_goal_image = image
+
+    @property
+    def rgb_milestones(self) -> np.ndarray:
+        assert self._rgb_milestones is not None, f"must set before using"
+        return self._rgb_milestones
+
+    @rgb_milestones.setter
+    def rgb_milestones(self, rgbs: np.ndarray | None):
+        """Set outside."""
+        if rgbs is not None:
+            # num_goal, h, w, 3
+            assert (
+                rgbs.ndim == 4 and rgbs.shape[0] == self.milestones.shape[0]
+            ), rgbs.shape
+            self._rgb_milestones = rgbs.copy()
+            self._current_goal_image = self._rgb_milestones[0]
+        else:
+            self._rgb_milestones = None
+            self._current_goal_image = None
+
+    @property
+    def stepwise_extra_debug_texts(self) -> dict | None:
+        return self._stepwise_extra_debug_texts
+
+    @stepwise_extra_debug_texts.setter
+    def stepwise_extra_debug_texts(self, texts: dict | None):
+        self._stepwise_extra_debug_texts = texts
+
+    @property
+    def milestones(self) -> np.ndarray:
+        return self.env.milestones
+
+    @milestones.setter
+    def milestones(self, milestones: np.ndarray | None):
+        self.env.milestones = milestones
+
+    @property
+    def milestone_embeddings(self) -> np.ndarray:
+        return self.env.milestone_embeddings
+
+    @milestone_embeddings.setter
+    def milestone_embeddings(self, embeddings: np.ndarray | None):
+        self.env.milestone_embeddings = embeddings
+
+    @property
+    def no_robot_milestones(self) -> np.ndarray:
+        return self.env.no_robot_milestones
+
+    @no_robot_milestones.setter
+    def no_robot_milestones(self, milestones: np.ndarray | None):
+        self.env.no_robot_milestones = milestones
+
+    @property
+    def current_milestone(self) -> np.ndarray:
+        return self.env.current_milestone
+
+    @property
+    def current_no_robot_milestone(self) -> np.ndarray | None:
+        return self.env.no_robot_milestones
+
+    @property
+    def current_embedding_distance(self) -> np.ndarray:
+        return self.env.current_embedding_distance
+
+    @current_embedding_distance.setter
+    def current_embedding_distance(self, dist: float | None):
+        self.env.current_embedding_distance = dist
+
+    @property
+    def current_obs_embedding(self) -> np.ndarray:
+        return self.env.current_obs_embedding
+
+    @current_obs_embedding.setter
+    def current_obs_embedding(self, embedding: np.ndarray):
+        before_set_achieved = self.completed_goals
+        self.env.current_obs_embedding = embedding
+        after_set_achieved = self.completed_goals
+        # if after_set_achieved > before_set_achieved and self.recording:
+        if self.recording:
+            # switch rgb current goal image
+            self.current_goal_image = self.rgb_milestones[
+                min(after_set_achieved, len(self.rgb_milestones) - 1)
+            ]
+
+    @property
+    def cur_milestone_idx(self) -> int | None:
+        return self.env.cur_milestone_idx
+
+    @cur_milestone_idx.setter
+    def cur_milestone_idx(self, idx: int | None):
+        self.env.cur_milestone_idx = idx
+
+    @property
+    def completed_goals(self) -> int:
+        return self.env.completed_goals
+
+    @property
+    def reset_states(self) -> dict:
+        return self.env.reset_states
+
+    @reset_states.setter
+    def reset_states(self, state: dict):
+        self.env.reset_states = state
+
+    @property
+    def task_name(self) -> str:
+        return self.env.task_name
+
+    @property
+    def metrics(self) -> dict:
+        return self.env.metrics
+
+    @property
+    def milestone_distances(self) -> np.ndarray | None:
+        return self.env.milestone_distances
+
+    @milestone_distances.setter
+    def milestone_distances(self, milestone_distances: np.ndarray | None):
+        self.env.milestone_distances = milestone_distances
+
+    @property
+    def current_no_robot_frame(self) -> np.ndarray:
+        return self.env.current_no_robot_frame
+
+    @property
+    def current_rgb_frame(self) -> np.ndarray:
+        return self.env.current_rgb_frame
+
+    @property
+    def num_milestones(self):
+        return self.env.num_milestones
+
+    @property
+    def milestone_indices(self):
+        return self.env.milestone_indices
+
+    @milestone_indices.setter
+    def milestone_indices(self, milestone_indices: np.ndarray | None):
+        self.env.milestone_indices = milestone_indices

uvd/envs/franka_kitchen/__init__.py

@@ -0,0 +1,20 @@
+import os
+import sys
+
+os.environ[
+    "LD_LIBRARY_PATH"
+] = f":{os.environ['HOME']}/.mujoco/mujoco210/bin:/usr/lib/nvidia"
+
+# workaround to import adept envs
+ADEPT_DIR = os.path.join(
+    os.path.dirname(__file__), "relay-policy-learning", "adept_envs"
+)
+assert os.path.exists(ADEPT_DIR), ADEPT_DIR
+sys.path.append(ADEPT_DIR)
+
+from .franka_kitchen_base import *
+from .franka_kitchen_constants import *
+
+import adept_envs.mujoco_env
+
+adept_envs.mujoco_env.USE_DM_CONTROL = USE_DM_CONTROL

uvd/envs/franka_kitchen/franka_kitchen_base.py

@@ -0,0 +1,446 @@
+"""Environments using kitchen and Franka robot."""
+from __future__ import annotations
+
+import copy
+from collections import OrderedDict
+
+import gym
+import numpy as np
+
+
+from uvd.envs.franka_kitchen.franka_kitchen_constants import *
+from mujoco_py import MjRenderContextOffscreen
+
+from adept_envs.franka.kitchen_multitask_v0 import KitchenV0
+from adept_envs.simulation.renderer import RenderMode, DMRenderer
+
+__all__ = ["KitchenBase"]
+
+
+class KitchenBase(KitchenV0):
+    ALL_TASKS = FRANKA_KITCHEN_ALL_TASKS
+    REMOVE_TASKS_WHEN_COMPLETE = True
+    TERMINATE_ON_TASK_COMPLETE = True
+    TERMINATE_ON_WRONG_COMPLETE = False
+    COMPLETE_IN_ANY_ORDER = False
+
+    frame_height = 224
+    frame_width = 224
+
+    def __init__(
+        self,
+        *,
+        task_elements: list | None = None,
+        reward_config: dict | None = FRANKA_KITCHEN_REWARD_CONFIG,
+        goal_masking: bool = True,
+        obs_keys: str | tuple[str, ...] = ("rgb", "proprio"),
+        max_horizon: int = 1000,
+        frame_height: int = 224,
+        frame_width: int = 224,
+        robot_params: dict | None = None,
+        frame_skip: int = 40,
+        gpu_id: int = -1,
+    ):
+        task_elements = (
+            list(task_elements) if task_elements is not None else list(self.ALL_TASKS)
+        )
+        self._task_elements = [
+            e if isinstance(e, str) else IDX_TO_ELEMENT[e] for e in task_elements
+        ]
+        self.goal_masking = goal_masking
+        # workaround for vector env preventing render
+        self._init_step = True
+        super(KitchenBase, self).__init__(
+            robot_params=robot_params or {}, frame_skip=frame_skip
+        )
+        self._num_goals_achieved = 0
+        self.tasks_to_complete = copy.deepcopy(self._task_elements)
+        self.reward_config = reward_config
+        self.obs_keys = (
+            tuple(obs_keys) if not isinstance(obs_keys, str) else (obs_keys,)
+        )
+        dict_obs_space = OrderedDict()
+        for key in obs_keys:
+            if key == "rgb":
+                dict_obs_space[key] = gym.spaces.Box(
+                    low=0,
+                    high=255,
+                    shape=(frame_height, frame_width, 3),
+                    dtype=np.uint8,
+                )
+            elif key == "proprio":
+                dict_obs_space[key] = gym.spaces.Box(
+                    low=-np.inf,
+                    high=np.inf,
+                    shape=(self.obs_dict["qp"].shape[0],),
+                    dtype=np.float32,
+                )
+            else:
+                raise NotImplementedError(key)
+        self.observation_space = gym.spaces.Dict(dict_obs_space)
+        self.max_horizon = max_horizon
+        self.frame_height = frame_height
+        self.frame_width = frame_width
+        self.gpu_id = gpu_id
+        self._renderer_setup = False
+        self.env_info = None
+        self.episode_length = 0
+
+    @property
+    def task_elements(self) -> list[str]:
+        return self._task_elements
+
+    @property
+    def task_name(self) -> str:
+        return "-".join(self.task_elements).replace(" ", "_")
+
+    @property
+    def current_subtask(self) -> str:
+        if self.num_goals_achieved == len(self.task_elements):
+            return "null"  # success
+        return self.task_elements[self.num_goals_achieved]
+
+    @property
+    def metrics(self):
+        """Accessible metrics for vector env."""
+        assert self.env_info is not None, f"try to call metrics without stepping"
+        metrics = dict(
+            episode_length=self.episode_length,
+            success=int(self.env_info["success"]),
+            completed_tasks=self.env_info["completed_tasks"],
+        )
+        return metrics
+
+    def _get_task_goal(
+        self, task: list[str] | None = None, actually_return_goal: bool = True
+    ) -> np.ndarray:
+        if task is None:
+            task = self.task_elements
+        new_goal = np.zeros_like(self.goal)
+        if self.goal_masking and not actually_return_goal:
+            return new_goal
+        for element in task:
+            element_idx = OBS_ELEMENT_INDICES[element]
+            element_goal = OBS_ELEMENT_GOALS[element]
+            new_goal[element_idx] = element_goal
+        return new_goal
+
+    def reset(
+        self,
+        init_qpos: np.ndarray | None = None,
+        init_qvel: np.ndarray | None = None,
+        task_elements: list[str] | None = None,
+        seed: int | None = None,
+        **kwargs,
+    ) -> OrderedDict[str, np.ndarray]:
+        """Return visual obs."""
+        if seed is not None:
+            self.seed(seed)
+        # set initial states if not None
+        self.reset_states = dict(
+            init_qpos=init_qpos, init_qvel=init_qvel, task_elements=task_elements
+        )
+        assert len(self.task_elements) > 0
+        self.tasks_to_complete = copy.deepcopy(self.task_elements)
+        self._num_goals_achieved = 0
+
+        self.env_info = None
+        self.episode_length = 0
+
+        super().reset()
+        return self.get_dict_obs(**kwargs)
+
+    def get_dict_obs(self, **kwargs):
+        """Observation sensor for dict observation space."""
+        obs = OrderedDict()
+        for k in self.obs_keys:
+            if k == "rgb":
+                if "frame_height" in kwargs:
+                    self.frame_height = kwargs.pop("frame_height")
+                if "frame_width" in kwargs:
+                    self.frame_width = kwargs.pop("frame_width")
+                obs[k] = self.render(
+                    mode="rgb_array", height=self.frame_height, width=self.frame_width
+                )
+            elif k == "proprio":
+                obs[k] = self.obs_dict["qp"]
+            else:
+                raise NotImplementedError(k)
+        return obs
+
+    @property
+    def reset_states(self) -> dict:
+        return dict(
+            init_qpos=self.init_qpos,
+            init_qvel=self.init_qvel,
+            task_elements=self.task_elements,
+        )
+
+    @reset_states.setter
+    def reset_states(self, state: dict | None):
+        """Set state in vector env for convenience."""
+        if state.get("init_qpos", None) is not None:
+            assert (
+                state.get("init_qvel", None) is not None
+            ), "qpos and qvel set together"
+            qpos = self._squeeze_state_vec(state["init_qpos"].copy())
+            qvel = self._squeeze_state_vec(state["init_qvel"].copy())
+            self.set_state(qpos=qpos, qvel=qvel)
+            self.init_qpos = self.data.qpos.ravel().copy()
+            self.init_qvel = self.data.qvel.ravel().copy()
+        if state.get("task_elements", None) is not None:
+            self._task_elements = [
+                e if isinstance(e, str) else IDX_TO_ELEMENT[e]
+                for e in self._squeeze_state_vec(
+                    state["task_elements"].copy(), to_list=True
+                )
+            ]
+
+    def set_state(self, qpos: np.ndarray, qvel: np.ndarray):
+        if isinstance(self.sim_robot.renderer, DMRenderer):
+            assert qpos.shape == (self.model.nq,) and qvel.shape == (self.model.nv,)
+            state = np.concatenate([qpos, qvel])
+            self.sim.set_state(state)
+            self.sim.forward()
+        else:
+            super().set_state(qpos=qpos, qvel=qvel)
+
+    def state_vector(self) -> np.ndarray:
+        if isinstance(self.sim_robot.renderer, DMRenderer):
+            return self.sim.get_state()
+        else:
+            return super().state_vector()
+
+    @staticmethod
+    def _squeeze_state_vec(vec: np.ndarray, to_list: bool = False) -> np.ndarray | list:
+        if not isinstance(vec, np.ndarray):
+            return vec
+        if vec.ndim != 1:
+            vec = vec.reshape((vec.shape[-1],))
+        if to_list:
+            return vec.tolist()
+        return vec
+
+    def _get_reward_n_score(self, obs_dict: dict) -> tuple[dict, float]:
+        """Score here means whether complete a new goal this step."""
+        reward_dict = {"true_reward": 0.0}
+
+        next_q_obs = obs_dict["qp"]
+        next_obj_obs = obs_dict["obj_qp"]
+        next_goal = self._get_task_goal(
+            task=self.task_elements, actually_return_goal=True
+        )  # obs_dict['goal']
+        idx_offset = len(next_q_obs)
+        completions = []
+        all_completed_so_far = True
+        distances_dict = {}
+        # # for element in self.tasks_to_complete:
+        # make metrics consistent for sync in DDP rollout
+        for i, element in enumerate(self.task_elements):
+            element_idx = OBS_ELEMENT_INDICES[element]
+            # distance = np.linalg.norm(
+            #     next_obj_obs[..., element_idx - idx_offset] - next_goal[element_idx]
+            # )
+            distance = abs(
+                next_obj_obs[..., element_idx - idx_offset] - next_goal[element_idx]
+            )  # keep state dims
+            distances_dict[element] = distance
+            # complete = distance < BONUS_THRESH
+            complete = distance < OBS_ELEMENT_THRESH[element]  # diff criteria
+            if isinstance(complete, np.ndarray):
+                complete = np.all(complete)
+
+            if (
+                self.REMOVE_TASKS_WHEN_COMPLETE
+                and element not in self.tasks_to_complete
+            ):
+                # already achieved task(s), and
+                # edge case for knobs that distances increasing later though rendered states not changed
+                condition = True
+            else:
+                condition = (
+                    complete and all_completed_so_far
+                    if not self.COMPLETE_IN_ANY_ORDER
+                    else complete
+                )
+            if condition:
+                completions.append(element)
+            all_completed_so_far = all_completed_so_far and condition
+        prev_tasks_to_complete = list(self.tasks_to_complete)
+        if self.REMOVE_TASKS_WHEN_COMPLETE:
+            [
+                self.tasks_to_complete.remove(element)
+                for element in completions
+                if element in self.tasks_to_complete
+            ]
+            num_goals_achieved = len(self.task_elements) - len(self.tasks_to_complete)
+        else:
+            num_goals_achieved = len(completions)
+
+        complete_new_goal_this_step = False
+        if num_goals_achieved > self._num_goals_achieved:
+            self._num_goals_achieved += 1
+            # edge case for knobs
+            if num_goals_achieved != self._num_goals_achieved:
+                self._num_goals_achieved = len(self.task_elements) - len(
+                    self.tasks_to_complete
+                )
+
+            reward_dict["true_reward"] += self.reward_config["progress"]
+            complete_new_goal_this_step = True
+
+        if self.num_goals_achieved == len(self.task_elements):
+            reward_dict["true_reward"] += self.reward_config["terminal"]
+
+        reward_dict["distances_left"] = distances_dict
+        score = int(complete_new_goal_this_step)
+        return reward_dict, score
+
+    @property
+    def num_goals_achieved(self):
+        if self.REMOVE_TASKS_WHEN_COMPLETE:
+            assert self._num_goals_achieved == len(self.task_elements) - len(
+                self.tasks_to_complete
+            ), f"{self._num_goals_achieved}, {self.task_elements}, {self.tasks_to_complete}"
+        return self._num_goals_achieved
+
+    def step(
+        self, action: np.ndarray, **kwargs
+    ) -> tuple[OrderedDict[str, np.ndarray], float, bool, dict]:
+        action = np.clip(action, -1.0, 1.0)
+        if not self.initializing:
+            action = self.act_mid + action * self.act_amp  # mean center and scale
+        else:
+            self.goal = self._get_task_goal()  # update goal if init
+        self.robot.step(self, action, step_duration=self.skip * self.model.opt.timestep)
+
+        low_dim_obs = self._get_obs()
+        # workaround for vector env
+        if self._init_step:
+            self._init_step = False
+            return low_dim_obs, 0.0, False, {}
+
+        self.episode_length += 1
+
+        reward_dict, score = self._get_reward_n_score(self.obs_dict)
+        success = (
+            not self.tasks_to_complete
+            if self.REMOVE_TASKS_WHEN_COMPLETE
+            else self.num_goals_achieved == len(self.tasks_to_complete)
+        )
+
+        env_info = {
+            "time": self.obs_dict["t"],
+            "obs_dict": self.obs_dict,
+            "rewards": reward_dict,
+            "score": score,
+            "success": success,
+            # self.task_elements[:self.num_goals_achieved - 1]
+            "completed_tasks": self.num_goals_achieved,
+            "episode_length": self.episode_length,
+            "current_sub_task": self.current_subtask,
+        }
+        done = self.episode_length >= self.max_horizon
+        if self.TERMINATE_ON_TASK_COMPLETE and success:
+            done = True
+        if self.TERMINATE_ON_WRONG_COMPLETE:
+            all_goal = self._get_task_goal(task=self.ALL_TASKS)
+            for wrong_task in list(set(self.ALL_TASKS) - set(self.task_elements)):
+                element_idx = OBS_ELEMENT_INDICES[wrong_task]
+                distance = np.linalg.norm(
+                    low_dim_obs[..., element_idx] - all_goal[element_idx]
+                )
+                complete = distance < BONUS_THRESH
+                if complete:
+                    done = True
+                    break
+
+        obs = self.get_dict_obs(**kwargs)
+        env_info["done"] = done
+        self.env_info = env_info
+        return obs, reward_dict["true_reward"], done, env_info
+
+    def close(self):
+        super().close()
+        self.close_env()
+
+    def render(
+        self,
+        mode: str = "human",
+        height: int | None = None,
+        width: int | None = None,
+        camera_id: int = -1,
+        set_robot_alpha: float | None = None,
+    ) -> np.ndarray | None:
+        height = height or self.frame_height
+        width = width or self.frame_width
+        if not self._renderer_setup:
+            self.set_gpu_id(gpu_id=self.gpu_id)
+            camera_settings = dict(
+                distance=2.2, lookat=[-0.2, 0.5, 2.0], azimuth=70, elevation=-35
+            )
+            self.sim_robot.renderer._camera_settings = camera_settings
+        alpha_map = None
+        if set_robot_alpha is not None:
+            alpha_map = self.set_robot_alpha(alpha=set_robot_alpha)
+
+        height = height or self.frame_height
+        width = width or self.frame_width
+        if mode in ["rgb_array", "rgb"]:
+            frame = self.sim_robot.renderer.render_offscreen(
+                width=width, height=height, mode=RenderMode.RGB, camera_id=camera_id
+            )
+        elif mode == "human":
+            frame = self.sim_robot.renderer.render_to_window()
+        else:
+            frame = super(KitchenV0, self).render(
+                mode=mode, height=height, width=width, camera_id=camera_id
+            )
+        if set_robot_alpha is not None:
+            self.unset_robot_alpha(alpha_map)
+        return frame
+
+    def set_robot_alpha(self, alpha: float) -> dict[int, float]:
+        alpha_map = dict()
+        if not USE_DM_CONTROL:
+            for name in self.sim.model.site_names:
+                if "end_effector" in name:
+                    orig_val = self.sim.model.site_rgba[
+                        self.sim.model.site_name2id(name)
+                    ][-1]
+                    alpha_map[name] = max(orig_val, alpha_map.get(name, -1))
+                    self.sim.model.site_rgba[self.sim.model.site_name2id(name)][
+                        -1
+                    ] = alpha
+        for dof_id in self.sim.model.dof_jntid:
+            alpha_map[dof_id] = max(
+                self.sim.model.geom_rgba[dof_id][-1], alpha_map.get(dof_id, -1)
+            )
+            self.sim.model.geom_rgba[dof_id][-1] = alpha
+        for dof_id in self.sim.model.dof_bodyid:
+            alpha_map[dof_id] = max(
+                self.sim.model.geom_rgba[dof_id][-1], alpha_map.get(dof_id, -1)
+            )
+            self.sim.model.geom_rgba[dof_id][-1] = alpha
+        return alpha_map
+
+    def unset_robot_alpha(self, alpha_map: dict[int, float]):
+        for dof_id, alpha in alpha_map.items():
+            if isinstance(dof_id, str) and "end_effector" in dof_id:
+                self.sim.model.site_rgba[self.sim.model.site_name2id(dof_id)][
+                    -1
+                ] = alpha
+            else:
+                self.sim.model.geom_rgba[dof_id][-1] = alpha
+
+    @property
+    def current_no_robot_frame(self) -> np.ndarray:
+        return self.render(mode="rgb_array", set_robot_alpha=0.0)
+
+    def set_gpu_id(self, gpu_id: int):
+        self._renderer_setup = True
+        # if not USE_DM_CONTROL:
+        if not isinstance(self.sim_robot.renderer, DMRenderer):
+            self.sim_robot.renderer._offscreen_renderer = MjRenderContextOffscreen(
+                self.sim_robot.renderer._sim, device_id=gpu_id
+            )

uvd/envs/franka_kitchen/franka_kitchen_constants.py

@@ -0,0 +1,62 @@
+import numpy as np
+
+USE_DM_CONTROL = False
+BONUS_THRESH = 0.3
+
+FRANKA_KITCHEN_ALL_TASKS = [
+    "bottom burner",
+    "top burner",
+    "light switch",
+    "slide cabinet",
+    "hinge cabinet",
+    "microwave",
+    "kettle",
+]
+
+OBS_ELEMENT_INDICES = {
+    # rotation & opening for bottom left burner
+    "bottom burner": np.array([11, 12]),
+    # rotation & opening for top left burner
+    "top burner": np.array([15, 16]),
+    # joint angle & opening
+    "light switch": np.array([17, 18]),
+    # Translation of the slide cabinet joint
+    "slide cabinet": np.array([19]),
+    # Rotation of the joint in the (left, right) hinge cabinet
+    "hinge cabinet": np.array([20, 21]),
+    # Rotation of the joint in the microwave door
+    "microwave": np.array([22]),
+    # x, y, z, qx, qy, qz, qw
+    "kettle": np.array([23, 24, 25, 26, 27, 28, 29]),
+}
+
+OBS_ELEMENT_GOALS = {
+    "bottom burner": np.array([-0.88, -0.01]),
+    "top burner": np.array([-0.92, -0.01]),
+    "light switch": np.array([-0.69, -0.05]),
+    "slide cabinet": np.array([0.37]),
+    # right hinge, left should be (-1.45, 0)
+    "hinge cabinet": np.array([0.0, 1.45]),
+    "microwave": np.array([-0.75]),
+    "kettle": np.array([-0.23, 0.75, 1.62, 0.99, 0.0, 0.0, -0.06]),
+}
+
+OBS_ELEMENT_THRESH = {
+    "bottom burner": 0.31,
+    "top burner": 0.31,
+    "light switch": 0.3,  # 0.1
+    "slide cabinet": 0.2,
+    "hinge cabinet": 0.2,
+    "microwave": 0.2,
+    # "kettle": np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.2]),
+    "kettle": np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.3]),
+}
+
+ELEMENT_TO_IDX = {k: i for i, k in enumerate(OBS_ELEMENT_GOALS.keys())}
+IDX_TO_ELEMENT = {i: k for k, i in ELEMENT_TO_IDX.items()}
+
+FRANKA_KITCHEN_REWARD_CONFIG = {
+    "progress": 0.0,  # complete sub-goal
+    "terminal": 10.0,  # complete all goals
+    "intrinsic_weight": 1.0,  # intrinsic reward if using embedding-dist-diff reward
+}

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/.pylintrc

@@ -0,0 +1,433 @@
+[MASTER]
+
+# A comma-separated list of package or module names from where C extensions may
+# be loaded. Extensions are loading into the active Python interpreter and may
+# run arbitrary code.
+extension-pkg-whitelist=
+
+# Add files or directories to the blacklist. They should be base names, not
+# paths.
+ignore=CVS
+
+# Add files or directories matching the regex patterns to the blacklist. The
+# regex matches against base names, not paths.
+ignore-patterns=
+
+# Python code to execute, usually for sys.path manipulation such as
+# pygtk.require().
+#init-hook=
+
+# Use multiple processes to speed up Pylint. Specifying 0 will auto-detect the
+# number of processors available to use.
+jobs=1
+
+# Control the amount of potential inferred values when inferring a single
+# object. This can help the performance when dealing with large functions or
+# complex, nested conditions.
+limit-inference-results=100
+
+# List of plugins (as comma separated values of python modules names) to load,
+# usually to register additional checkers.
+load-plugins=
+
+# Pickle collected data for later comparisons.
+persistent=yes
+
+# Specify a configuration file.
+#rcfile=
+
+# When enabled, pylint would attempt to guess common misconfiguration and emit
+# user-friendly hints instead of false-positive error messages.
+suggestion-mode=yes
+
+# Allow loading of arbitrary C extensions. Extensions are imported into the
+# active Python interpreter and may run arbitrary code.
+unsafe-load-any-extension=no
+
+
+[MESSAGES CONTROL]
+
+# Only show warnings with the listed confidence levels. Leave empty to show
+# all. Valid levels: HIGH, INFERENCE, INFERENCE_FAILURE, UNDEFINED.
+confidence=
+
+# Disable the message, report, category or checker with the given id(s). You
+# can either give multiple identifiers separated by comma (,) or put this
+# option multiple times (only on the command line, not in the configuration
+# file where it should appear only once). You can also use "--disable=all" to
+# disable everything first and then reenable specific checks. For example, if
+# you want to run only the similarities checker, you can use "--disable=all
+# --enable=similarities". If you want to run only the classes checker, but have
+# no Warning level messages displayed, use "--disable=all --enable=classes
+# --disable=W".
+disable=relative-beyond-top-level
+
+
+[REPORTS]
+
+# Python expression which should return a note less than 10 (10 is the highest
+# note). You have access to the variables errors warning, statement which
+# respectively contain the number of errors / warnings messages and the total
+# number of statements analyzed. This is used by the global evaluation report
+# (RP0004).
+evaluation=10.0 - ((float(5 * error + warning + refactor + convention) / statement) * 10)
+
+# Template used to display messages. This is a python new-style format string
+# used to format the message information. See doc for all details.
+#msg-template=
+
+# Set the output format. Available formats are text, parseable, colorized, json
+# and msvs (visual studio). You can also give a reporter class, e.g.
+# mypackage.mymodule.MyReporterClass.
+output-format=text
+
+# Tells whether to display a full report or only the messages.
+reports=no
+
+# Activate the evaluation score.
+score=yes
+
+
+[REFACTORING]
+
+# Maximum number of nested blocks for function / method body
+max-nested-blocks=5
+
+# Complete name of functions that never returns. When checking for
+# inconsistent-return-statements if a never returning function is called then
+# it will be considered as an explicit return statement and no message will be
+# printed.
+never-returning-functions=sys.exit
+
+
+[LOGGING]
+
+# Format style used to check logging format string. `old` means using %
+# formatting, while `new` is for `{}` formatting.
+logging-format-style=old
+
+# Logging modules to check that the string format arguments are in logging
+# function parameter format.
+logging-modules=logging
+
+
+[VARIABLES]
+
+# List of additional names supposed to be defined in builtins. Remember that
+# you should avoid defining new builtins when possible.
+additional-builtins=
+
+# Tells whether unused global variables should be treated as a violation.
+allow-global-unused-variables=yes
+
+# List of strings which can identify a callback function by name. A callback
+# name must start or end with one of those strings.
+callbacks=cb_,
+          _cb
+
+# A regular expression matching the name of dummy variables (i.e. expected to
+# not be used).
+dummy-variables-rgx=_+$|(_[a-zA-Z0-9_]*[a-zA-Z0-9]+?$)|dummy|^ignored_|^unused_
+
+# Argument names that match this expression will be ignored. Default to name
+# with leading underscore.
+ignored-argument-names=_.*|^ignored_|^unused_
+
+# Tells whether we should check for unused import in __init__ files.
+init-import=no
+
+# List of qualified module names which can have objects that can redefine
+# builtins.
+redefining-builtins-modules=six.moves,past.builtins,future.builtins,builtins,io
+
+
+[FORMAT]
+
+# Expected format of line ending, e.g. empty (any line ending), LF or CRLF.
+expected-line-ending-format=
+
+# Regexp for a line that is allowed to be longer than the limit.
+ignore-long-lines=^\s*(# )?<?https?://\S+>?$
+
+# Number of spaces of indent required inside a hanging  or continued line.
+indent-after-paren=4
+
+# String used as indentation unit. This is usually "    " (4 spaces) or "\t" (1
+# tab).
+indent-string='    '
+
+# Maximum number of characters on a single line.
+max-line-length=80
+
+# Maximum number of lines in a module
+max-module-lines=99999
+
+# List of optional constructs for which whitespace checking is disabled. `dict-
+# separator` is used to allow tabulation in dicts, etc.: {1  : 1,\n222: 2}.
+# `trailing-comma` allows a space between comma and closing bracket: (a, ).
+# `empty-line` allows space-only lines.
+no-space-check=trailing-comma,
+               dict-separator
+
+# Allow the body of a class to be on the same line as the declaration if body
+# contains single statement.
+single-line-class-stmt=no
+
+# Allow the body of an if to be on the same line as the test if there is no
+# else.
+single-line-if-stmt=no
+
+
+[TYPECHECK]
+
+# List of decorators that produce context managers, such as
+# contextlib.contextmanager. Add to this list to register other decorators that
+# produce valid context managers.
+contextmanager-decorators=contextlib.contextmanager
+
+# List of members which are set dynamically and missed by pylint inference
+# system, and so shouldn't trigger E1101 when accessed. Python regular
+# expressions are accepted.
+generated-members=
+
+# Tells whether missing members accessed in mixin class should be ignored. A
+# mixin class is detected if its name ends with "mixin" (case insensitive).
+ignore-mixin-members=yes
+
+# Tells whether to warn about missing members when the owner of the attribute
+# is inferred to be None.
+ignore-none=yes
+
+# This flag controls whether pylint should warn about no-member and similar
+# checks whenever an opaque object is returned when inferring. The inference
+# can return multiple potential results while evaluating a Python object, but
+# some branches might not be evaluated, which results in partial inference. In
+# that case, it might be useful to still emit no-member and other checks for
+# the rest of the inferred objects.
+ignore-on-opaque-inference=yes
+
+# List of class names for which member attributes should not be checked (useful
+# for classes with dynamically set attributes). This supports the use of
+# qualified names.
+ignored-classes=optparse.Values,thread._local,_thread._local
+
+# List of module names for which member attributes should not be checked
+# (useful for modules/projects where namespaces are manipulated during runtime
+# and thus existing member attributes cannot be deduced by static analysis. It
+# supports qualified module names, as well as Unix pattern matching.
+ignored-modules=
+
+# Show a hint with possible names when a member name was not found. The aspect
+# of finding the hint is based on edit distance.
+missing-member-hint=yes
+
+# The minimum edit distance a name should have in order to be considered a
+# similar match for a missing member name.
+missing-member-hint-distance=1
+
+# The total number of similar names that should be taken in consideration when
+# showing a hint for a missing member.
+missing-member-max-choices=1
+
+
+[SIMILARITIES]
+
+# Ignore comments when computing similarities.
+ignore-comments=yes
+
+# Ignore docstrings when computing similarities.
+ignore-docstrings=yes
+
+# Ignore imports when computing similarities.
+ignore-imports=no
+
+# Minimum lines number of a similarity.
+min-similarity-lines=4
+
+
+[BASIC]
+
+# Naming style matching correct argument names
+argument-naming-style=snake_case
+
+# Regular expression matching correct argument names. Overrides argument-
+# naming-style
+argument-rgx=^[a-z][a-z0-9_]*$
+
+# Naming style matching correct attribute names
+attr-naming-style=snake_case
+
+# Regular expression matching correct attribute names. Overrides attr-naming-
+# style
+attr-rgx=^_{0,2}[a-z][a-z0-9_]*$
+
+# Bad variable names which should always be refused, separated by a comma
+bad-names=
+
+# Naming style matching correct class attribute names
+class-attribute-naming-style=any
+
+# Regular expression matching correct class attribute names. Overrides class-
+# attribute-naming-style
+class-attribute-rgx=^(_?[A-Z][A-Z0-9_]*|__[a-z0-9_]+__|_?[a-z][a-z0-9_]*)$
+
+# Naming style matching correct class names
+class-naming-style=PascalCase
+
+# Regular expression matching correct class names. Overrides class-naming-style
+class-rgx=^_?[A-Z][a-zA-Z0-9]*$
+
+# Naming style matching correct constant names
+const-naming-style=UPPER_CASE
+
+# Regular expression matching correct constant names. Overrides const-naming-
+# style
+const-rgx=^(_?[A-Z][A-Z0-9_]*|__[a-z0-9_]+__|_?[a-z][a-z0-9_]*)$
+
+# Minimum line length for functions/classes that require docstrings, shorter
+# ones are exempt.
+docstring-min-length=10
+
+# Naming style matching correct function names
+function-naming-style=snake_case
+
+# Regular expression matching correct function names. Overrides function-
+# naming-style
+function-rgx=^(?:(?P<exempt>setUp|tearDown|setUpModule|tearDownModule)|(?P<camel_case>_?[A-Z][a-zA-Z0-9]*)|(?P<snake_case>_?[a-z][a-z0-9_]*))$
+
+# Good variable names which should always be accepted, separated by a comma
+good-names=main,
+           _
+
+# Include a hint for the correct naming format with invalid-name
+include-naming-hint=no
+
+# Naming style matching correct inline iteration names
+inlinevar-naming-style=any
+
+# Regular expression matching correct inline iteration names. Overrides
+# inlinevar-naming-style
+inlinevar-rgx=^[a-z][a-z0-9_]*$
+
+# Naming style matching correct method names
+method-naming-style=snake_case
+
+# Regular expression matching correct method names. Overrides method-naming-
+# style
+method-rgx=(?x)^(?:(?P<exempt>_[a-z0-9_]+__|runTest|setUp|tearDown|setUpTestCase|tearDownTestCase|setupSelf|tearDownClass|setUpClass|(test|assert)_*[A-Z0-9][a-zA-Z0-9_]*|next)|(?P<camel_case>_{0,2}[A-Z][a-zA-Z0-9_]*)|(?P<snake_case>_{0,2}[a-z][a-z0-9_]*))$
+
+# Naming style matching correct module names
+module-naming-style=snake_case
+
+# Regular expression matching correct module names. Overrides module-naming-
+# style
+module-rgx=^(_?[a-z][a-z0-9_]*)|__init__|PRESUBMIT|PRESUBMIT_unittest$
+
+# Colon-delimited sets of names that determine each other's naming style when
+# the name regexes allow several styles.
+name-group=function:method
+
+# Regular expression which should only match function or class names that do
+# not require a docstring.
+no-docstring-rgx=(__.*__|main)
+
+# List of decorators that produce properties, such as abc.abstractproperty. Add
+# to this list to register other decorators that produce valid properties.
+property-classes=abc.abstractproperty,google3.pyglib.function_utils.cached.property
+
+# Naming style matching correct variable names
+variable-naming-style=snake_case
+
+# Regular expression matching correct variable names. Overrides variable-
+# naming-style
+variable-rgx=^[a-z][a-z0-9_]*$
+
+
+[SPELLING]
+
+# Limits count of emitted suggestions for spelling mistakes.
+max-spelling-suggestions=4
+
+# Spelling dictionary name. Available dictionaries: none. To make it working
+# install python-enchant package..
+spelling-dict=
+
+# List of comma separated words that should not be checked.
+spelling-ignore-words=
+
+# A path to a file that contains private dictionary; one word per line.
+spelling-private-dict-file=
+
+# Tells whether to store unknown words to indicated private dictionary in
+# --spelling-private-dict-file option instead of raising a message.
+spelling-store-unknown-words=no
+
+
+[MISCELLANEOUS]
+
+# List of note tags to take in consideration, separated by a comma.
+notes=FIXME,
+      XXX,
+      TODO
+
+
+[IMPORTS]
+
+# Allow wildcard imports from modules that define __all__.
+allow-wildcard-with-all=no
+
+# Analyse import fallback blocks. This can be used to support both Python 2 and
+# 3 compatible code, which means that the block might have code that exists
+# only in one or another interpreter, leading to false positives when analysed.
+analyse-fallback-blocks=no
+
+# Deprecated modules which should not be used, separated by a comma.
+deprecated-modules=optparse,tkinter.tix
+
+# Create a graph of external dependencies in the given file (report RP0402 must
+# not be disabled).
+ext-import-graph=
+
+# Create a graph of every (i.e. internal and external) dependencies in the
+# given file (report RP0402 must not be disabled).
+import-graph=
+
+# Create a graph of internal dependencies in the given file (report RP0402 must
+# not be disabled).
+int-import-graph=
+
+# Force import order to recognize a module as part of the standard
+# compatibility libraries.
+known-standard-library=
+
+# Force import order to recognize a module as part of a third party library.
+known-third-party=enchant
+
+
+[CLASSES]
+
+# List of method names used to declare (i.e. assign) instance attributes.
+defining-attr-methods=__init__,
+                      __new__,
+                      setUp
+
+# List of member names, which should be excluded from the protected access
+# warning.
+exclude-protected=_asdict,
+                  _fields,
+                  _replace,
+                  _source,
+                  _make
+
+# List of valid names for the first argument in a class method.
+valid-classmethod-first-arg=cls
+
+# List of valid names for the first argument in a metaclass class method.
+valid-metaclass-classmethod-first-arg=cls
+
+
+[EXCEPTIONS]
+
+# Exceptions that will emit a warning when being caught. Defaults to
+# "Exception".
+overgeneral-exceptions=Exception

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/.style.yapf

@@ -0,0 +1,323 @@
+[style]
+# Align closing bracket with visual indentation.
+align_closing_bracket_with_visual_indent=False
+
+# Allow dictionary keys to exist on multiple lines. For example:
+#
+#   x = {
+#       ('this is the first element of a tuple',
+#        'this is the second element of a tuple'):
+#            value,
+#   }
+allow_multiline_dictionary_keys=False
+
+# Allow lambdas to be formatted on more than one line.
+allow_multiline_lambdas=False
+
+# Allow splitting before a default / named assignment in an argument list.
+allow_split_before_default_or_named_assigns=True
+
+# Allow splits before the dictionary value.
+allow_split_before_dict_value=True
+
+#   Let spacing indicate operator precedence. For example:
+#
+#     a = 1 * 2 + 3 / 4
+#     b = 1 / 2 - 3 * 4
+#     c = (1 + 2) * (3 - 4)
+#     d = (1 - 2) / (3 + 4)
+#     e = 1 * 2 - 3
+#     f = 1 + 2 + 3 + 4
+#
+# will be formatted as follows to indicate precedence:
+#
+#     a = 1*2 + 3/4
+#     b = 1/2 - 3*4
+#     c = (1+2) * (3-4)
+#     d = (1-2) / (3+4)
+#     e = 1*2 - 3
+#     f = 1 + 2 + 3 + 4
+#
+arithmetic_precedence_indication=False
+
+# Number of blank lines surrounding top-level function and class
+# definitions.
+blank_lines_around_top_level_definition=2
+
+# Insert a blank line before a class-level docstring.
+blank_line_before_class_docstring=False
+
+# Insert a blank line before a module docstring.
+blank_line_before_module_docstring=False
+
+# Insert a blank line before a 'def' or 'class' immediately nested
+# within another 'def' or 'class'. For example:
+#
+#   class Foo:
+#                      # <------ this blank line
+#     def method():
+#       ...
+blank_line_before_nested_class_or_def=True
+
+# Do not split consecutive brackets. Only relevant when
+# dedent_closing_brackets is set. For example:
+#
+#    call_func_that_takes_a_dict(
+#        {
+#            'key1': 'value1',
+#            'key2': 'value2',
+#        }
+#    )
+#
+# would reformat to:
+#
+#    call_func_that_takes_a_dict({
+#        'key1': 'value1',
+#        'key2': 'value2',
+#    })
+coalesce_brackets=False
+
+# The column limit.
+column_limit=80
+
+# The style for continuation alignment. Possible values are:
+#
+# - SPACE: Use spaces for continuation alignment. This is default behavior.
+# - FIXED: Use fixed number (CONTINUATION_INDENT_WIDTH) of columns
+#   (ie: CONTINUATION_INDENT_WIDTH/INDENT_WIDTH tabs) for continuation
+#   alignment.
+# - LESS: Slightly left if cannot vertically align continuation lines with
+#   indent characters.
+# - VALIGN-RIGHT: Vertically align continuation lines with indent
+#   characters. Slightly right (one more indent character) if cannot
+#   vertically align continuation lines with indent characters.
+#
+# For options FIXED, and VALIGN-RIGHT are only available when USE_TABS is
+# enabled.
+continuation_align_style=SPACE
+
+# Indent width used for line continuations.
+continuation_indent_width=4
+
+# Put closing brackets on a separate line, dedented, if the bracketed
+# expression can't fit in a single line. Applies to all kinds of brackets,
+# including function definitions and calls. For example:
+#
+#   config = {
+#       'key1': 'value1',
+#       'key2': 'value2',
+#   }        # <--- this bracket is dedented and on a separate line
+#
+#   time_series = self.remote_client.query_entity_counters(
+#       entity='dev3246.region1',
+#       key='dns.query_latency_tcp',
+#       transform=Transformation.AVERAGE(window=timedelta(seconds=60)),
+#       start_ts=now()-timedelta(days=3),
+#       end_ts=now(),
+#   )        # <--- this bracket is dedented and on a separate line
+dedent_closing_brackets=False
+
+# Disable the heuristic which places each list element on a separate line
+# if the list is comma-terminated.
+disable_ending_comma_heuristic=False
+
+# Place each dictionary entry onto its own line.
+each_dict_entry_on_separate_line=True
+
+# The regex for an i18n comment. The presence of this comment stops
+# reformatting of that line, because the comments are required to be
+# next to the string they translate.
+i18n_comment=#\..*
+
+# The i18n function call names. The presence of this function stops
+# reformattting on that line, because the string it has cannot be moved
+# away from the i18n comment.
+i18n_function_call=N_, _
+
+# Indent blank lines.
+indent_blank_lines=False
+
+# Indent the dictionary value if it cannot fit on the same line as the
+# dictionary key. For example:
+#
+#   config = {
+#       'key1':
+#           'value1',
+#       'key2': value1 +
+#               value2,
+#   }
+indent_dictionary_value=False
+
+# The number of columns to use for indentation.
+indent_width=4
+
+# Join short lines into one line. E.g., single line 'if' statements.
+join_multiple_lines=True
+
+# Do not include spaces around selected binary operators. For example:
+#
+#   1 + 2 * 3 - 4 / 5
+#
+# will be formatted as follows when configured with "*,/":
+#
+#   1 + 2*3 - 4/5
+#
+no_spaces_around_selected_binary_operators=
+
+# Use spaces around default or named assigns.
+spaces_around_default_or_named_assign=False
+
+# Use spaces around the power operator.
+spaces_around_power_operator=False
+
+# The number of spaces required before a trailing comment.
+# This can be a single value (representing the number of spaces
+# before each trailing comment) or list of values (representing
+# alignment column values; trailing comments within a block will
+# be aligned to the first column value that is greater than the maximum
+# line length within the block). For example:
+#
+# With spaces_before_comment=5:
+#
+#   1 + 1 # Adding values
+#
+# will be formatted as:
+#
+#   1 + 1     # Adding values <-- 5 spaces between the end of the statement and comment
+#
+# With spaces_before_comment=15, 20:
+#
+#   1 + 1 # Adding values
+#   two + two # More adding
+#
+#   longer_statement # This is a longer statement
+#   short # This is a shorter statement
+#
+#   a_very_long_statement_that_extends_beyond_the_final_column # Comment
+#   short # This is a shorter statement
+#
+# will be formatted as:
+#
+#   1 + 1          # Adding values <-- end of line comments in block aligned to col 15
+#   two + two      # More adding
+#
+#   longer_statement    # This is a longer statement <-- end of line comments in block aligned to col 20
+#   short               # This is a shorter statement
+#
+#   a_very_long_statement_that_extends_beyond_the_final_column  # Comment <-- the end of line comments are aligned based on the line length
+#   short                                                       # This is a shorter statement
+#
+spaces_before_comment=2
+
+# Insert a space between the ending comma and closing bracket of a list,
+# etc.
+space_between_ending_comma_and_closing_bracket=False
+
+# Split before arguments
+split_all_comma_separated_values=False
+
+# Split before arguments if the argument list is terminated by a
+# comma.
+split_arguments_when_comma_terminated=False
+
+# Set to True to prefer splitting before '&', '|' or '^' rather than
+# after.
+split_before_bitwise_operator=False
+
+# Split before the closing bracket if a list or dict literal doesn't fit on
+# a single line.
+split_before_closing_bracket=True
+
+# Split before a dictionary or set generator (comp_for). For example, note
+# the split before the 'for':
+#
+#   foo = {
+#       variable: 'Hello world, have a nice day!'
+#       for variable in bar if variable != 42
+#   }
+split_before_dict_set_generator=False
+
+# Split before the '.' if we need to split a longer expression:
+#
+#   foo = ('This is a really long string: {}, {}, {}, {}'.format(a, b, c, d))
+#
+# would reformat to something like:
+#
+#   foo = ('This is a really long string: {}, {}, {}, {}'
+#          .format(a, b, c, d))
+split_before_dot=False
+
+# Split after the opening paren which surrounds an expression if it doesn't
+# fit on a single line.
+split_before_expression_after_opening_paren=False
+
+# If an argument / parameter list is going to be split, then split before
+# the first argument.
+split_before_first_argument=False
+
+# Set to True to prefer splitting before 'and' or 'or' rather than
+# after.
+split_before_logical_operator=False
+
+# Split named assignments onto individual lines.
+split_before_named_assigns=True
+
+# Set to True to split list comprehensions and generators that have
+# non-trivial expressions and multiple clauses before each of these
+# clauses. For example:
+#
+#   result = [
+#       a_long_var + 100 for a_long_var in xrange(1000)
+#       if a_long_var % 10]
+#
+# would reformat to something like:
+#
+#   result = [
+#       a_long_var + 100
+#       for a_long_var in xrange(1000)
+#       if a_long_var % 10]
+split_complex_comprehension=True
+
+# The penalty for splitting right after the opening bracket.
+split_penalty_after_opening_bracket=30
+
+# The penalty for splitting the line after a unary operator.
+split_penalty_after_unary_operator=10000
+
+# The penalty for splitting right before an if expression.
+split_penalty_before_if_expr=0
+
+# The penalty of splitting the line around the '&', '|', and '^'
+# operators.
+split_penalty_bitwise_operator=300
+
+# The penalty for splitting a list comprehension or generator
+# expression.
+split_penalty_comprehension=2100
+
+# The penalty for characters over the column limit.
+split_penalty_excess_character=7000
+
+# The penalty incurred by adding a line split to the unwrapped line. The
+# more line splits added the higher the penalty.
+split_penalty_for_added_line_split=30
+
+# The penalty of splitting a list of "import as" names. For example:
+#
+#   from a_very_long_or_indented_module_name_yada_yad import (long_argument_1,
+#                                                             long_argument_2,
+#                                                             long_argument_3)
+#
+# would reformat to something like:
+#
+#   from a_very_long_or_indented_module_name_yada_yad import (
+#       long_argument_1, long_argument_2, long_argument_3)
+split_penalty_import_names=0
+
+# The penalty of splitting the line around the 'and' and 'or'
+# operators.
+split_penalty_logical_operator=300
+
+# Use the Tab character for indentation.
+use_tabs=False
+

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/adept_envs/__init__.py

@@ -0,0 +1,15 @@
+#!/usr/bin/python
+#
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/adept_envs/base_robot.py

@@ -0,0 +1,153 @@
+#!/usr/bin/python
+#
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import deque
+
+import numpy as np
+
+
+class BaseRobot(object):
+    """Base class for all robot classes."""
+
+    def __init__(
+        self,
+        n_jnt,
+        n_obj,
+        pos_bounds=None,
+        vel_bounds=None,
+        calibration_path=None,
+        is_hardware=False,
+        device_name=None,
+        overlay=False,
+        calibration_mode=False,
+        observation_cache_maxsize=5,
+    ):
+        """Create a new robot.
+
+        Args:
+            n_jnt: The number of dofs in the robot.
+            n_obj: The number of dofs in the object.
+            pos_bounds: (n_jnt, 2)-shape matrix denoting the min and max joint
+                position for each joint.
+            vel_bounds: (n_jnt, 2)-shape matrix denoting the min and max joint
+                velocity for each joint.
+            calibration_path: File path to the calibration configuration file to
+                use.
+            is_hardware: Whether to run on hardware or not.
+            device_name: The device path for the robot hardware. Only required
+                in legacy mode.
+            overlay: Whether to show a simulation overlay of the hardware.
+            calibration_mode: Start with motors disengaged.
+        """
+
+        assert n_jnt > 0
+        assert n_obj >= 0
+
+        self._n_jnt = n_jnt
+        self._n_obj = n_obj
+        self._n_dofs = n_jnt + n_obj
+
+        self._pos_bounds = None
+        if pos_bounds is not None:
+            pos_bounds = np.array(pos_bounds, dtype=np.float32)
+            assert pos_bounds.shape == (self._n_dofs, 2)
+            for low, high in pos_bounds:
+                assert low < high
+            self._pos_bounds = pos_bounds
+        self._vel_bounds = None
+        if vel_bounds is not None:
+            vel_bounds = np.array(vel_bounds, dtype=np.float32)
+            assert vel_bounds.shape == (self._n_dofs, 2)
+            for low, high in vel_bounds:
+                assert low < high
+            self._vel_bounds = vel_bounds
+
+        self._is_hardware = is_hardware
+        self._device_name = device_name
+        self._calibration_path = calibration_path
+        self._overlay = overlay
+        self._calibration_mode = calibration_mode
+        self._observation_cache_maxsize = observation_cache_maxsize
+
+        # Gets updated
+        self._observation_cache = deque([], maxlen=self._observation_cache_maxsize)
+
+    @property
+    def n_jnt(self):
+        return self._n_jnt
+
+    @property
+    def n_obj(self):
+        return self._n_obj
+
+    @property
+    def n_dofs(self):
+        return self._n_dofs
+
+    @property
+    def pos_bounds(self):
+        return self._pos_bounds
+
+    @property
+    def vel_bounds(self):
+        return self._vel_bounds
+
+    @property
+    def is_hardware(self):
+        return self._is_hardware
+
+    @property
+    def device_name(self):
+        return self._device_name
+
+    @property
+    def calibration_path(self):
+        return self._calibration_path
+
+    @property
+    def overlay(self):
+        return self._overlay
+
+    @property
+    def has_obj(self):
+        return self._n_obj > 0
+
+    @property
+    def calibration_mode(self):
+        return self._calibration_mode
+
+    @property
+    def observation_cache_maxsize(self):
+        return self._observation_cache_maxsize
+
+    @property
+    def observation_cache(self):
+        return self._observation_cache
+
+    def clip_positions(self, positions):
+        """Clips the given joint positions to the position bounds.
+
+        Args:
+            positions: The joint positions.
+
+        Returns:
+            The bounded joint positions.
+        """
+        if self.pos_bounds is None:
+            return positions
+        assert len(positions) == self.n_jnt or len(positions) == self.n_dofs
+        pos_bounds = self.pos_bounds[: len(positions)]
+        return np.clip(positions, pos_bounds[:, 0], pos_bounds[:, 1])

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/adept_envs/franka/__init__.py

@@ -0,0 +1,24 @@
+#!/usr/bin/python
+#
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from gym.envs.registration import register
+
+# Relax the robot
+register(
+    id="kitchen_relax-v1",
+    entry_point="adept_envs.franka.kitchen_multitask_v0:KitchenTaskRelaxV1",
+    max_episode_steps=280,
+)

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/adept_envs/franka/assets/franka_kitchen_jntpos_act_ab.xml

@@ -0,0 +1,94 @@
+<!--Copyright 2020 Google LLC-->
+
+<!--Licensed under the Apache License, Version 2.0 (the "License");-->
+<!--you may not use this file except in compliance with the License.-->
+<!--You may obtain a copy of the License at-->
+
+    <!--https://www.apache.org/licenses/LICENSE-2.0-->
+
+<!--Unless required by applicable law or agreed to in writing, software-->
+<!--distributed under the License is distributed on an "AS IS" BASIS,-->
+<!--WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.-->
+<!--See the License for the specific language governing permissions and-->
+<!--limitations under the License.-->
+
+<mujoco model="franka_mocap_studyTable_buttons">
+
+    <size njmax='1000' nconmax='1000'/>
+
+    <include file="../../../../adept_models/scenes/basic_scene.xml"/>
+    <include file="../../../../third_party/franka/assets/assets.xml"/>
+    <include file="../../../../third_party/franka/assets/actuator0.xml"/>
+    <include file="../../../../adept_models/kitchen/assets/oven_asset.xml"/>
+    <include file="../../../../adept_models/kitchen/assets/counters_asset.xml"/>
+    <include file="../../../../adept_models/kitchen/assets/backwall_asset.xml"/>
+    <include file="../../../../adept_models/kitchen/assets/slidecabinet_asset.xml"/>
+    <include file="../../../../adept_models/kitchen/assets/hingecabinet_asset.xml"/>
+    <include file="../../../../adept_models/kitchen/assets/microwave_asset.xml"/>
+    <include file="../../../../adept_models/kitchen/assets/kettle_asset.xml"/>
+
+    <visual>
+    <global offwidth="2560" offheight="1920" />
+    <quality shadowsize="4096" offsamples="8" />
+    <map force="0.1" fogend="5" />
+    </visual>
+
+    <compiler inertiafromgeom='auto' inertiagrouprange='3 5' angle="radian"
+              meshdir="../../../../adept_models/kitchen"
+              texturedir="../../../../adept_models/kitchen"/>
+
+    <equality>
+        <weld body1="vive_controller" body2="world" solref="0.02 1" solimp=".7 .95 0.050"/>
+    </equality>
+
+    <worldbody>
+
+        <!-- Mocap -->
+        <body name="vive_controller" mocap="true" pos="-0.440 -0.092 2.026" euler="-1.57 0 -.785">
+            <geom type="box" group="2" pos='0 0 .142' size="0.02 0.10 0.03" contype="0" conaffinity="0" rgba=".9 .7 .95 0" euler="0 0 -.785"/>
+        </body>
+
+        <site name='target' pos='0 0 0' size='0.1' rgba='0 2 0 .2'/>
+        <camera name='left_cap' pos='-1.2 -0.5 1.8' quat='0.78 0.49 -0.22 -0.32' />
+        <camera name='right_cap' pos='1.2 -0.5 1.8' quat='0.76 0.5 0.21 0.35'/>
+
+        <!-- Robot -->
+        <body pos='0. 0 1.8' euler='0 0 1.57'>
+            <geom type='cylinder' size='.120 .90' pos='-.04 0 -0.90' class='panda_viz'/>
+            <include file="../../../../third_party/franka/assets/chain0.xml"/>
+        </body>
+
+        <body name='desk' pos='-0.1 0.75 0'>
+
+            <body name="counters1" pos="0 0 0" >
+                <include file="../../../../adept_models/kitchen/assets/counters_chain.xml"/>
+            </body>
+            <body name="oven" pos="0 0 0" >
+                <include file="../../../../adept_models/kitchen/assets/oven_chain.xml"/>
+            </body>
+            <body name="backwall" pos="0 0 0" >
+                <include file="../../../../adept_models/kitchen/assets/backwall_chain.xml"/>
+            </body>
+            <body name="slidecabinet" pos="0.4 0.3 2.6" >
+                <include file="../../../../adept_models/kitchen/assets/slidecabinet_chain.xml"/>
+            </body>
+            <body name="hingecabinet" pos="-0.504 0.28 2.6" >
+                <include file="../../../../adept_models/kitchen/assets/hingecabinet_chain.xml"/>
+            </body>
+            <body name="microwave" pos="-0.750 -0.025 1.6" euler="0 0 0.3">
+                <include file="../../../../adept_models/kitchen/assets/microwave_chain.xml"/>
+            </body>
+        </body>
+        <body name="kettle" pos="-0.269 0.35 1.626">
+            <freejoint/>
+            <include file="../../../../adept_models/kitchen/assets/kettle_chain.xml"/>
+        </body>
+
+    </worldbody>
+
+
+    <keyframe>
+        <key qpos='0.16 -1.76 1.84 -2.51 0.36 0.79 1.55 0.00 0.0 1.25561e-05 1.57437e-07 1.25561e-05 1.57437e-07 1.25561e-05 1.57437e-07 1.25561e-05 1.57437e-07 8.24417e-05 9.48283e-05 0 0 0 0 -0.269 0.35 1.61523 1 1.34939e-19 -3.51612e-05 -7.50168e-19'/>
+    </keyframe>
+
+</mujoco>

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/adept_envs/franka/kitchen_multitask_v0.py

@@ -0,0 +1,234 @@
+"""Kitchen environment for long horizon manipulation."""
+#!/usr/bin/python
+#
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import numpy as np
+from dm_control.mujoco import engine
+from gym import spaces
+
+from adept_envs import robot_env
+from adept_envs.utils.configurable import configurable
+
+
+@configurable(pickleable=True)
+class KitchenV0(robot_env.RobotEnv):
+    CALIBRATION_PATHS = {
+        "default": os.path.join(os.path.dirname(__file__), "robot/franka_config.xml")
+    }
+    # Converted to velocity actuation
+    ROBOTS = {"robot": "adept_envs.franka.robot.franka_robot:Robot_VelAct"}
+    MODEl = os.path.join(
+        os.path.dirname(__file__), "../franka/assets/franka_kitchen_jntpos_act_ab.xml"
+    )
+    N_DOF_ROBOT = 9
+    N_DOF_OBJECT = 21
+
+    def __init__(self, robot_params={}, frame_skip=40):
+        self.goal_concat = True
+        self.obs_dict = {}
+        self.robot_noise_ratio = 0.1  # 10% as per robot_config specs
+        self.goal = np.zeros((30,))
+
+        super().__init__(
+            self.MODEl,
+            robot=self.make_robot(
+                n_jnt=self.N_DOF_ROBOT,  # root+robot_jnts
+                n_obj=self.N_DOF_OBJECT,
+                **robot_params
+            ),
+            frame_skip=frame_skip,
+            camera_settings=dict(
+                distance=4.5,
+                azimuth=-66,
+                elevation=-65,
+            ),
+        )
+        self.init_qpos = self.sim.model.key_qpos[0].copy()
+
+        # For the microwave kettle slide hinge
+        self.init_qpos = np.array(
+            [
+                1.48388023e-01,
+                -1.76848573e00,
+                1.84390296e00,
+                -2.47685760e00,
+                2.60252026e-01,
+                7.12533105e-01,
+                1.59515394e00,
+                4.79267505e-02,
+                3.71350919e-02,
+                -2.66279850e-04,
+                -5.18043486e-05,
+                3.12877220e-05,
+                -4.51199853e-05,
+                -3.90842156e-06,
+                -4.22629655e-05,
+                6.28065475e-05,
+                4.04984708e-05,
+                4.62730939e-04,
+                -2.26906415e-04,
+                -4.65501369e-04,
+                -6.44129196e-03,
+                -1.77048263e-03,
+                1.08009684e-03,
+                -2.69397440e-01,
+                3.50383255e-01,
+                1.61944683e00,
+                1.00618764e00,
+                4.06395120e-03,
+                -6.62095997e-03,
+                -2.68278933e-04,
+            ]
+        )
+
+        self.init_qvel = self.sim.model.key_qvel[0].copy()
+
+        self.act_mid = np.zeros(self.N_DOF_ROBOT)
+        self.act_amp = 2.0 * np.ones(self.N_DOF_ROBOT)
+
+        act_lower = -1 * np.ones((self.N_DOF_ROBOT,))
+        act_upper = 1 * np.ones((self.N_DOF_ROBOT,))
+        self.action_space = spaces.Box(act_lower, act_upper)
+
+        obs_upper = 8.0 * np.ones(self.obs_dim)
+        obs_lower = -obs_upper
+        self.observation_space = spaces.Box(obs_lower, obs_upper)
+
+    def _get_reward_n_score(self, obs_dict):
+        raise NotImplementedError()
+
+    def step(self, a, b=None):
+        a = np.clip(a, -1.0, 1.0)
+
+        if not self.initializing:
+            a = self.act_mid + a * self.act_amp  # mean center and scale
+        else:
+            self.goal = self._get_task_goal()  # update goal if init
+
+        self.robot.step(self, a, step_duration=self.skip * self.model.opt.timestep)
+
+        # observations
+        obs = self._get_obs()
+
+        # rewards
+        reward_dict, score = self._get_reward_n_score(self.obs_dict)
+
+        # termination
+        done = False
+
+        # finalize step
+        env_info = {
+            "time": self.obs_dict["t"],
+            "obs_dict": self.obs_dict,
+            "rewards": reward_dict,
+            "score": score,
+            "images": np.asarray(self.render(mode="rgb_array")),
+        }
+        # self.render()
+        return obs, reward_dict["r_total"], done, env_info
+
+    def _get_obs(self):
+        t, qp, qv, obj_qp, obj_qv = self.robot.get_obs(
+            self, robot_noise_ratio=self.robot_noise_ratio
+        )
+
+        self.obs_dict = {}
+        self.obs_dict["t"] = t
+        self.obs_dict["qp"] = qp
+        self.obs_dict["qv"] = qv
+        self.obs_dict["obj_qp"] = obj_qp
+        self.obs_dict["obj_qv"] = obj_qv
+        self.obs_dict["goal"] = self.goal
+        if self.goal_concat:
+            return np.concatenate(
+                [self.obs_dict["qp"], self.obs_dict["obj_qp"], self.obs_dict["goal"]]
+            )
+
+    def reset_model(self):
+        reset_pos = self.init_qpos[:].copy()
+        reset_vel = self.init_qvel[:].copy()
+        self.robot.reset(self, reset_pos, reset_vel)
+        self.sim.forward()
+        self.goal = self._get_task_goal()  # sample a new goal on reset
+        return self._get_obs()
+
+    def evaluate_success(self, paths):
+        # score
+        mean_score_per_rollout = np.zeros(shape=len(paths))
+        for idx, path in enumerate(paths):
+            mean_score_per_rollout[idx] = np.mean(path["env_infos"]["score"])
+        mean_score = np.mean(mean_score_per_rollout)
+
+        # success percentage
+        num_success = 0
+        num_paths = len(paths)
+        for path in paths:
+            num_success += bool(path["env_infos"]["rewards"]["bonus"][-1])
+        success_percentage = num_success * 100.0 / num_paths
+
+        # fuse results
+        return np.sign(mean_score) * (
+            1e6 * round(success_percentage, 2) + abs(mean_score)
+        )
+
+    def close_env(self):
+        self.robot.close()
+
+    def set_goal(self, goal):
+        self.goal = goal
+
+    def _get_task_goal(self):
+        return self.goal
+
+    # Only include goal
+    @property
+    def goal_space(self):
+        len_obs = self.observation_space.low.shape[0]
+        env_lim = np.abs(self.observation_space.low[0])
+        return spaces.Box(
+            low=-env_lim, high=env_lim, shape=(len_obs // 2,), dtype=np.float32
+        )
+
+    def convert_to_active_observation(self, observation):
+        return observation
+
+
+class KitchenTaskRelaxV1(KitchenV0):
+    """Kitchen environment with proper camera and goal setup."""
+
+    def __init__(self):
+        super(KitchenTaskRelaxV1, self).__init__()
+
+    def _get_reward_n_score(self, obs_dict):
+        reward_dict = {}
+        reward_dict["true_reward"] = 0.0
+        reward_dict["bonus"] = 0.0
+        reward_dict["r_total"] = 0.0
+        score = 0.0
+        return reward_dict, score
+
+    def render(self, mode="human"):
+        if mode == "rgb_array":
+            camera = engine.MovableCamera(self.sim, 1920, 2560)
+            camera.set_pose(
+                distance=2.2, lookat=[-0.2, 0.5, 2.0], azimuth=70, elevation=-35
+            )
+            img = camera.render()
+            return img
+        else:
+            super(KitchenTaskRelaxV1, self).render()

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/adept_envs/franka/robot/franka_config.xml

@@ -0,0 +1,59 @@
+<!--Copyright 2020 Google LLC-->
+
+<!--Licensed under the Apache License, Version 2.0 (the "License");-->
+<!--you may not use this file except in compliance with the License.-->
+<!--You may obtain a copy of the License at-->
+
+    <!--https://www.apache.org/licenses/LICENSE-2.0-->
+
+<!--Unless required by applicable law or agreed to in writing, software-->
+<!--distributed under the License is distributed on an "AS IS" BASIS,-->
+<!--WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.-->
+<!--See the License for the specific language governing permissions and-->
+<!--limitations under the License.-->
+<config name='Franka'>
+
+	<!-- Franka -->
+	<qpos0 name='q0' mode='1' mj_dof='0' hardware_dof='0' scale='1' offset='0' pos_bound='-2.9 2.9' vel_bound='-10 10' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos1 name='q1' mode='1' mj_dof='1' hardware_dof='1' scale='1' offset='0' pos_bound='-1.8 1.8' vel_bound='-10 10' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos2 name='q2' mode='1' mj_dof='2' hardware_dof='2' scale='1' offset='0' pos_bound='-2.9 2.9' vel_bound='-10 10' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos3 name='q3' mode='1' mj_dof='3' hardware_dof='3' scale='1' offset='0' pos_bound='-3.1 0.0' vel_bound='-10 10' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos4 name='q4' mode='1' mj_dof='4' hardware_dof='4' scale='1' offset='0' pos_bound='-2.9 2.9' vel_bound='-10 10' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos5 name='q5' mode='1' mj_dof='5' hardware_dof='5' scale='1' offset='0' pos_bound='00.0 3.8' vel_bound='-10 10' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos6 name='q6' mode='1' mj_dof='6' hardware_dof='6' scale='1' offset='0' pos_bound='-2.9 2.9' vel_bound='-10 10' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos7 name='q7' mode='1' mj_dof='7' hardware_dof='7' scale='1' offset='0' pos_bound='00.0 0.04' vel_bound='-10 10' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos8 name='q8' mode='1' mj_dof='8' hardware_dof='8' scale='1' offset='0' pos_bound='0.0  0.04' vel_bound='-10 10' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+
+	<!-- Desk -->
+	<qpos9   name='deskSlideB' mode='1' mj_dof='9' hardware_dof='9' scale='1' offset='0' pos_bound='-.5 0.0' vel_bound='-5 5' pos_noise_amp='0.005' vel_noise_amp='0.005' />
+	<qpos10  name='deskSlideT' mode='1' mj_dof='10' hardware_dof='10' scale='1' offset='0' pos_bound='-.5 0.0' vel_bound='-5 5' pos_noise_amp='0.005' vel_noise_amp='0.005' />
+
+	<!-- Buttons -->
+	<qpos11  name='rBotton' mode='1' mj_dof='11' hardware_dof='11' scale='1' offset='0' pos_bound='-.005 0.0' vel_bound='-5 5' pos_noise_amp='0.0005' vel_noise_amp='0.005' />
+	<qpos12  name='gButton' mode='1' mj_dof='12' hardware_dof='12' scale='1' offset='0' pos_bound='-.005 0.0' vel_bound='-5 5' pos_noise_amp='0.0005' vel_noise_amp='0.005' />
+	<qpos13  name='bButton' mode='1' mj_dof='13' hardware_dof='13' scale='1' offset='0' pos_bound='-.005 0.0' vel_bound='-5 5' pos_noise_amp='0.0005' vel_noise_amp='0.005' />
+	<qpos14  name='rLight' mode='1' mj_dof='14' hardware_dof='14' scale='1' offset='0' pos_bound='-.005 0.0' vel_bound='-5 5' pos_noise_amp='0.0005' vel_noise_amp='0.005' />
+	<qpos15  name='bLight' mode='1' mj_dof='15' hardware_dof='15' scale='1' offset='0' pos_bound='-.005 0.0' vel_bound='-5 5' pos_noise_amp='0.0005' vel_noise_amp='0.005' />
+	<qpos16  name='gLight' mode='1' mj_dof='16' hardware_dof='16' scale='1' offset='0' pos_bound='-.005 0.0' vel_bound='-5 5' pos_noise_amp='0.0005' vel_noise_amp='0.005' />
+
+	<!-- Blocks -->
+	<qpos17 name='q17' mode='1' mj_dof='17' hardware_dof='17'  scale='1' offset='0' pos_bound='-1.5 1.5' vel_bound='-5 5' pos_noise_amp='0.005' vel_noise_amp='0.005' />
+	<qpos18 name='q18' mode='1' mj_dof='18' hardware_dof='18' scale='1' offset='0' pos_bound='-1.5 1.5' vel_bound='-5 5' pos_noise_amp='0.005' vel_noise_amp='0.005' />
+	<qpos19 name='q19' mode='1' mj_dof='19' hardware_dof='19' scale='1' offset='0' pos_bound='-1.5 1.5' vel_bound='-5 5' pos_noise_amp='0.005' vel_noise_amp='0.005' />
+	<qpos20 name='q20' mode='1' mj_dof='20' hardware_dof='20' scale='1' offset='0' pos_bound='-10.57 10.57' vel_bound='-.5 .5' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos21 name='q21' mode='1' mj_dof='21' hardware_dof='21' scale='1' offset='0' pos_bound='-10.57 10.57' vel_bound='-.5 .5' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos22 name='q22' mode='1' mj_dof='22' hardware_dof='22' scale='1' offset='0' pos_bound='-10.57 10.57' vel_bound='-.5 .5' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos23 name='q23' mode='1' mj_dof='23' hardware_dof='23' scale='1' offset='0' pos_bound='-1.5 1.5' vel_bound='-5 5' pos_noise_amp='0.005' vel_noise_amp='0.005' />
+	<qpos24 name='q24' mode='1' mj_dof='24' hardware_dof='24' scale='1' offset='0' pos_bound='-1.5 1.5' vel_bound='-5 5' pos_noise_amp='0.005' vel_noise_amp='0.005' />
+	<qpos25 name='q25' mode='1' mj_dof='25' hardware_dof='25' scale='1' offset='0' pos_bound='-1.5 1.5' vel_bound='-5 5' pos_noise_amp='0.005' vel_noise_amp='0.005' />
+	<qpos26 name='q26' mode='1' mj_dof='26' hardware_dof='26' scale='1' offset='0' pos_bound='-10.57 10.57' vel_bound='-.5 .5' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos27 name='q27' mode='1' mj_dof='27' hardware_dof='27' scale='1' offset='0' pos_bound='-10.57 10.57' vel_bound='-.5 .5' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos28 name='q28' mode='1' mj_dof='28' hardware_dof='28' scale='1' offset='0' pos_bound='-10.57 10.57' vel_bound='-.5 .5' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos29 name='q29' mode='1' mj_dof='29' hardware_dof='29' scale='1' offset='0' pos_bound='-1.5 1.5' vel_bound='-5 5' pos_noise_amp='0.005' vel_noise_amp='0.005' />
+	<qpos30 name='q30' mode='1' mj_dof='30' hardware_dof='30' scale='1' offset='0' pos_bound='-1.5 1.5' vel_bound='-5 5' pos_noise_amp='0.005' vel_noise_amp='0.005' />
+	<qpos31 name='q31' mode='1' mj_dof='31' hardware_dof='31' scale='1' offset='0' pos_bound='-1.5 1.5' vel_bound='-5 5' pos_noise_amp='0.005' vel_noise_amp='0.005' />
+	<qpos32 name='q32' mode='1' mj_dof='32' hardware_dof='32' scale='1' offset='0' pos_bound='-10.57 10.57' vel_bound='-.5 .5' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos33 name='q33' mode='1' mj_dof='33' hardware_dof='33' scale='1' offset='0' pos_bound='-10.57 10.57' vel_bound='-.5 .5' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+	<qpos34 name='q34' mode='1' mj_dof='34' hardware_dof='34' scale='1' offset='0' pos_bound='-10.57 10.57' vel_bound='-.5 .5' pos_noise_amp='0.1' vel_noise_amp='0.1' />
+
+</config>

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/adept_envs/franka/robot/franka_robot.py

@@ -0,0 +1,342 @@
+#!/usr/bin/python
+#
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+
+# obervations structure
+from collections import namedtuple
+
+import numpy as np
+from termcolor import cprint
+
+from adept_envs import base_robot
+from adept_envs.utils.config import get_config_root_node, read_config_from_node
+
+observation = namedtuple(
+    "observation", ["time", "qpos_robot", "qvel_robot", "qpos_object", "qvel_object"]
+)
+
+
+franka_interface = ""
+
+
+class Robot(base_robot.BaseRobot):
+    """Abstracts away the differences between the robot_simulation and
+    robot_hardware."""
+
+    def __init__(self, *args, **kwargs):
+        super(Robot, self).__init__(*args, **kwargs)
+        global franka_interface
+
+        # Read robot configurations
+        self._read_specs_from_config(robot_configs=self.calibration_path)
+
+        # Robot: Handware
+        if self.is_hardware:
+            if franka_interface is "":
+                raise NotImplementedError()
+                from handware.franka import franka
+
+                # initialize franka
+                self.franka_interface = franka()
+                franka_interface = self.franka_interface
+                cprint(
+                    "Initializing %s Hardware (Status:%d)"
+                    % (self.robot_name, self.franka.okay(self.robot_hardware_dof)),
+                    "white",
+                    "on_grey",
+                )
+            else:
+                self.franka_interface = franka_interface
+                cprint("Reusing previours Franka session", "white", "on_grey")
+
+        # Robot: Simulation
+        else:
+            self.robot_name = "Franka"
+            cprint("Initializing %s sim" % self.robot_name, "white", "on_grey")
+
+        # Robot's time
+        self.time_start = time.time()
+        self.time = time.time() - self.time_start
+        self.time_render = -1  # time of rendering
+
+    # read specs from the calibration file
+    def _read_specs_from_config(self, robot_configs):
+        root, root_name = get_config_root_node(config_file_name=robot_configs)
+        self.robot_name = root_name[0]
+        self.robot_mode = np.zeros(self.n_dofs, dtype=int)
+        self.robot_mj_dof = np.zeros(self.n_dofs, dtype=int)
+        self.robot_hardware_dof = np.zeros(self.n_dofs, dtype=int)
+        self.robot_scale = np.zeros(self.n_dofs, dtype=float)
+        self.robot_offset = np.zeros(self.n_dofs, dtype=float)
+        self.robot_pos_bound = np.zeros([self.n_dofs, 2], dtype=float)
+        self.robot_vel_bound = np.zeros([self.n_dofs, 2], dtype=float)
+        self.robot_pos_noise_amp = np.zeros(self.n_dofs, dtype=float)
+        self.robot_vel_noise_amp = np.zeros(self.n_dofs, dtype=float)
+
+        print("Reading configurations for %s" % self.robot_name)
+        for i in range(self.n_dofs):
+            self.robot_mode[i] = read_config_from_node(
+                root, "qpos" + str(i), "mode", int
+            )
+            self.robot_mj_dof[i] = read_config_from_node(
+                root, "qpos" + str(i), "mj_dof", int
+            )
+            self.robot_hardware_dof[i] = read_config_from_node(
+                root, "qpos" + str(i), "hardware_dof", int
+            )
+            self.robot_scale[i] = read_config_from_node(
+                root, "qpos" + str(i), "scale", float
+            )
+            self.robot_offset[i] = read_config_from_node(
+                root, "qpos" + str(i), "offset", float
+            )
+            self.robot_pos_bound[i] = read_config_from_node(
+                root, "qpos" + str(i), "pos_bound", float
+            )
+            self.robot_vel_bound[i] = read_config_from_node(
+                root, "qpos" + str(i), "vel_bound", float
+            )
+            self.robot_pos_noise_amp[i] = read_config_from_node(
+                root, "qpos" + str(i), "pos_noise_amp", float
+            )
+            self.robot_vel_noise_amp[i] = read_config_from_node(
+                root, "qpos" + str(i), "vel_noise_amp", float
+            )
+
+    # convert to hardware space
+    def _de_calib(self, qp_mj, qv_mj=None):
+        qp_ad = (qp_mj - self.robot_offset) / self.robot_scale
+        if qv_mj is not None:
+            qv_ad = qv_mj / self.robot_scale
+            return qp_ad, qv_ad
+        else:
+            return qp_ad
+
+    # convert to mujoco space
+    def _calib(self, qp_ad, qv_ad):
+        qp_mj = qp_ad * self.robot_scale + self.robot_offset
+        qv_mj = qv_ad * self.robot_scale
+        return qp_mj, qv_mj
+
+    # refresh the observation cache
+    def _observation_cache_refresh(self, env):
+        for _ in range(self.observation_cache_maxsize):
+            self.get_obs(env, sim_mimic_hardware=False)
+
+    # get past observation
+    def get_obs_from_cache(self, env, index=-1):
+        assert (index >= 0 and index < self.observation_cache_maxsize) or (
+            index < 0 and index >= -self.observation_cache_maxsize
+        ), (
+            "cache index out of bound. (cache size is %2d)"
+            % self.observation_cache_maxsize
+        )
+        obs = self.observation_cache[index]
+        if self.has_obj:
+            return (
+                obs.time,
+                obs.qpos_robot,
+                obs.qvel_robot,
+                obs.qpos_object,
+                obs.qvel_object,
+            )
+        else:
+            return obs.time, obs.qpos_robot, obs.qvel_robot
+
+    # get observation
+    def get_obs(
+        self, env, robot_noise_ratio=1, object_noise_ratio=1, sim_mimic_hardware=True
+    ):
+        if self.is_hardware:
+            raise NotImplementedError()
+
+        else:
+            # Gather simulated observation
+            qp = env.sim.data.qpos[: self.n_jnt].copy()
+            qv = env.sim.data.qvel[: self.n_jnt].copy()
+            if self.has_obj:
+                qp_obj = env.sim.data.qpos[-self.n_obj :].copy()
+                qv_obj = env.sim.data.qvel[-self.n_obj :].copy()
+            else:
+                qp_obj = None
+                qv_obj = None
+            self.time = env.sim.data.time
+
+            # Simulate observation noise
+            if not env.initializing:
+                qp += (
+                    robot_noise_ratio
+                    * self.robot_pos_noise_amp[: self.n_jnt]
+                    * env.np_random.uniform(low=-1.0, high=1.0, size=self.n_jnt)
+                )
+                qv += (
+                    robot_noise_ratio
+                    * self.robot_vel_noise_amp[: self.n_jnt]
+                    * env.np_random.uniform(low=-1.0, high=1.0, size=self.n_jnt)
+                )
+                if self.has_obj:
+                    qp_obj += (
+                        robot_noise_ratio
+                        * self.robot_pos_noise_amp[-self.n_obj :]
+                        * env.np_random.uniform(low=-1.0, high=1.0, size=self.n_obj)
+                    )
+                    qv_obj += (
+                        robot_noise_ratio
+                        * self.robot_vel_noise_amp[-self.n_obj :]
+                        * env.np_random.uniform(low=-1.0, high=1.0, size=self.n_obj)
+                    )
+
+        # cache observations
+        obs = observation(
+            time=self.time,
+            qpos_robot=qp,
+            qvel_robot=qv,
+            qpos_object=qp_obj,
+            qvel_object=qv_obj,
+        )
+        self.observation_cache.append(obs)
+
+        if self.has_obj:
+            return (
+                obs.time,
+                obs.qpos_robot,
+                obs.qvel_robot,
+                obs.qpos_object,
+                obs.qvel_object,
+            )
+        else:
+            return obs.time, obs.qpos_robot, obs.qvel_robot
+
+    # enforce position specs.
+    def ctrl_position_limits(self, ctrl_position):
+        ctrl_feasible_position = np.clip(
+            ctrl_position,
+            self.robot_pos_bound[: self.n_jnt, 0],
+            self.robot_pos_bound[: self.n_jnt, 1],
+        )
+        return ctrl_feasible_position
+
+    # step the robot env
+    def step(self, env, ctrl_desired, step_duration, sim_override=False):
+        # Populate observation cache during startup
+        if env.initializing:
+            self._observation_cache_refresh(env)
+
+        # enforce velocity limits
+        ctrl_feasible = self.ctrl_velocity_limits(ctrl_desired, step_duration)
+
+        # enforce position limits
+        ctrl_feasible = self.ctrl_position_limits(ctrl_feasible)
+
+        # Send controls to the robot
+        if self.is_hardware and (not sim_override):
+            raise NotImplementedError()
+        else:
+            env.do_simulation(
+                ctrl_feasible, int(step_duration / env.sim.model.opt.timestep)
+            )  # render is folded in here
+
+        # Update current robot state on the overlay
+        if self.overlay:
+            env.sim.data.qpos[self.n_jnt : 2 * self.n_jnt] = env.desired_pose.copy()
+            env.sim.forward()
+
+        # synchronize time
+        if self.is_hardware:
+            time_now = time.time() - self.time_start
+            time_left_in_step = step_duration - (time_now - self.time)
+            if time_left_in_step > 0.0001:
+                time.sleep(time_left_in_step)
+        return 1
+
+    def reset(
+        self,
+        env,
+        reset_pose,
+        reset_vel,
+        overlay_mimic_reset_pose=True,
+        sim_override=False,
+    ):
+        reset_pose = self.clip_positions(reset_pose)
+
+        if self.is_hardware:
+            raise NotImplementedError()
+        else:
+            env.sim.reset()
+            env.sim.data.qpos[: self.n_jnt] = reset_pose[: self.n_jnt].copy()
+            env.sim.data.qvel[: self.n_jnt] = reset_vel[: self.n_jnt].copy()
+            if self.has_obj:
+                env.sim.data.qpos[-self.n_obj :] = reset_pose[-self.n_obj :].copy()
+                env.sim.data.qvel[-self.n_obj :] = reset_vel[-self.n_obj :].copy()
+            env.sim.forward()
+
+        if self.overlay:
+            env.sim.data.qpos[self.n_jnt : 2 * self.n_jnt] = env.desired_pose[
+                : self.n_jnt
+            ].copy()
+            env.sim.forward()
+
+        # refresh observation cache before exit
+        self._observation_cache_refresh(env)
+
+    def close(self):
+        if self.is_hardware:
+            cprint(
+                "Closing Franka hardware... ", "white", "on_grey", end="", flush=True
+            )
+            status = 0
+            raise NotImplementedError()
+            cprint("Closed (Status: {})".format(status), "white", "on_grey", flush=True)
+        else:
+            cprint("Closing Franka sim", "white", "on_grey", flush=True)
+
+
+class Robot_PosAct(Robot):
+    # enforce velocity sepcs.
+    # ALERT: This depends on previous observation. This is not ideal as it breaks MDP addumptions. Be careful
+    def ctrl_velocity_limits(self, ctrl_position, step_duration):
+        last_obs = self.observation_cache[-1]
+        ctrl_desired_vel = (
+            ctrl_position - last_obs.qpos_robot[: self.n_jnt]
+        ) / step_duration
+
+        ctrl_feasible_vel = np.clip(
+            ctrl_desired_vel,
+            self.robot_vel_bound[: self.n_jnt, 0],
+            self.robot_vel_bound[: self.n_jnt, 1],
+        )
+        ctrl_feasible_position = (
+            last_obs.qpos_robot[: self.n_jnt] + ctrl_feasible_vel * step_duration
+        )
+        return ctrl_feasible_position
+
+
+class Robot_VelAct(Robot):
+    # enforce velocity sepcs.
+    # ALERT: This depends on previous observation. This is not ideal as it breaks MDP addumptions. Be careful
+    def ctrl_velocity_limits(self, ctrl_velocity, step_duration):
+        last_obs = self.observation_cache[-1]
+
+        ctrl_feasible_vel = np.clip(
+            ctrl_velocity,
+            self.robot_vel_bound[: self.n_jnt, 0],
+            self.robot_vel_bound[: self.n_jnt, 1],
+        )
+        ctrl_feasible_position = (
+            last_obs.qpos_robot[: self.n_jnt] + ctrl_feasible_vel * step_duration
+        )
+        return ctrl_feasible_position

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/adept_envs/mujoco_env.py

@@ -0,0 +1,222 @@
+"""Base environment for MuJoCo-based environments."""
+
+#!/usr/bin/python
+#
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import collections
+from collections import abc
+import os
+from typing import Dict, Optional
+
+import gym
+import numpy as np
+from gym import spaces
+from gym.utils import seeding
+
+from adept_envs.simulation.sim_robot import MujocoSimRobot, RenderMode
+
+# from simulation.renderer import RenderMode
+
+DEFAULT_RENDER_SIZE = 480
+
+USE_DM_CONTROL = True
+
+
+class MujocoEnv(gym.Env):
+    """Superclass for all MuJoCo environments."""
+
+    def __init__(
+        self,
+        model_path: str,
+        frame_skip: int,
+        camera_settings: Optional[Dict] = None,
+        use_dm_backend: Optional[bool] = None,
+    ):
+        """Initializes a new MuJoCo environment.
+
+        Args:
+            model_path: The path to the MuJoCo XML file.
+            frame_skip: The number of simulation steps per environment step. On
+              hardware this influences the duration of each environment step.
+            camera_settings: Settings to initialize the simulation camera. This
+              can contain the keys `distance`, `azimuth`, and `elevation`.
+            use_dm_backend: A boolean to switch between mujoco-py and dm_control.
+        """
+        self._seed()
+        if not os.path.isfile(model_path):
+            raise IOError(
+                "[MujocoEnv]: Model path does not exist: {}".format(model_path)
+            )
+        self.frame_skip = frame_skip
+
+        self.sim_robot = MujocoSimRobot(
+            model_path,
+            use_dm_backend=use_dm_backend or USE_DM_CONTROL,
+            camera_settings=camera_settings,
+        )
+        self.sim = self.sim_robot.sim
+        self.model = self.sim_robot.model
+        self.data = self.sim_robot.data
+
+        self.metadata = {
+            "render.modes": ["human", "rgb_array", "depth_array"],
+            "video.frames_per_second": int(np.round(1.0 / self.dt)),
+        }
+        self.mujoco_render_frames = False
+
+        self.init_qpos = self.data.qpos.ravel().copy()
+        self.init_qvel = self.data.qvel.ravel().copy()
+        observation, _reward, done, _info = self.step(np.zeros(self.model.nu))
+        assert not done
+
+        bounds = self.model.actuator_ctrlrange.copy()
+        act_upper = bounds[:, 1]
+        act_lower = bounds[:, 0]
+
+        # Define the action and observation spaces.
+        # HACK: MJRL is still using gym 0.9.x so we can't provide a dtype.
+        try:
+            self.action_space = spaces.Box(act_lower, act_upper, dtype=np.float32)
+            if isinstance(observation, abc.Mapping):
+                self.observation_space = spaces.Dict(
+                    {
+                        k: spaces.Box(-np.inf, np.inf, shape=v.shape, dtype=np.float32)
+                        for k, v in observation.items()
+                    }
+                )
+            else:
+                self.obs_dim = (
+                    np.sum([o.size for o in observation])
+                    if type(observation) is tuple
+                    else observation.size
+                )
+                self.observation_space = spaces.Box(
+                    -np.inf, np.inf, observation.shape, dtype=np.float32
+                )
+
+        except TypeError:
+            # Fallback case for gym 0.9.x
+            self.action_space = spaces.Box(act_lower, act_upper)
+            assert not isinstance(
+                observation, collections.Mapping
+            ), "gym 0.9.x does not support dictionary observation."
+            self.obs_dim = (
+                np.sum([o.size for o in observation])
+                if type(observation) is tuple
+                else observation.size
+            )
+            self.observation_space = spaces.Box(-np.inf, np.inf, observation.shape)
+
+    def seed(self, seed=None):  # Compatibility with new gym
+        return self._seed(seed)
+
+    def _seed(self, seed=None):
+        self.np_random, seed = seeding.np_random(seed)
+        return [seed]
+
+    # methods to override:
+    # ----------------------------
+
+    def reset_model(self):
+        """Reset the robot degrees of freedom (qpos and qvel).
+
+        Implement this in each subclass.
+        """
+        raise NotImplementedError
+
+    # -----------------------------
+
+    def reset(self):  # compatibility with new gym
+        return self._reset()
+
+    def _reset(self):
+        self.sim.reset()
+        self.sim.forward()
+        ob = self.reset_model()
+        return ob
+
+    def set_state(self, qpos, qvel):
+        assert qpos.shape == (self.model.nq,) and qvel.shape == (self.model.nv,)
+        state = self.sim.get_state()
+        for i in range(self.model.nq):
+            state.qpos[i] = qpos[i]
+        for i in range(self.model.nv):
+            state.qvel[i] = qvel[i]
+        self.sim.set_state(state)
+        self.sim.forward()
+
+    @property
+    def dt(self):
+        return self.model.opt.timestep * self.frame_skip
+
+    def do_simulation(self, ctrl, n_frames):
+        for i in range(self.model.nu):
+            self.sim.data.ctrl[i] = ctrl[i]
+
+        for _ in range(n_frames):
+            self.sim.step()
+
+            # TODO(michaelahn): Remove this; render should be called separately.
+            if self.mujoco_render_frames is True:
+                self.mj_render()
+
+    def render(
+        self,
+        mode="human",
+        width=DEFAULT_RENDER_SIZE,
+        height=DEFAULT_RENDER_SIZE,
+        camera_id=-1,
+    ):
+        """Renders the environment.
+
+        Args:
+            mode: The type of rendering to use.
+                - 'human': Renders to a graphical window.
+                - 'rgb_array': Returns the RGB image as an np.ndarray.
+                - 'depth_array': Returns the depth image as an np.ndarray.
+            width: The width of the rendered image. This only affects offscreen
+                rendering.
+            height: The height of the rendered image. This only affects
+                offscreen rendering.
+            camera_id: The ID of the camera to use. By default, this is the free
+                camera. If specified, only affects offscreen rendering.
+        """
+        if mode == "human":
+            self.sim_robot.renderer.render_to_window()
+        elif mode == "rgb_array":
+            assert width and height
+            return self.sim_robot.renderer.render_offscreen(
+                width, height, mode=RenderMode.RGB, camera_id=camera_id
+            )
+        elif mode == "depth_array":
+            assert width and height
+            return self.sim_robot.renderer.render_offscreen(
+                width, height, mode=RenderMode.DEPTH, camera_id=camera_id
+            )
+        else:
+            raise NotImplementedError(mode)
+
+    def close(self):
+        self.sim_robot.close()
+
+    def mj_render(self):
+        """Backwards compatibility with MJRL."""
+        self.render(mode="human")
+
+    def state_vector(self):
+        state = self.sim.get_state()
+        return np.concatenate([state.qpos.flat, state.qvel.flat])

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/adept_envs/robot_env.py

@@ -0,0 +1,178 @@
+"""Base class for robotics environments."""
+
+#!/usr/bin/python
+#
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Dict, Optional
+
+import numpy as np
+
+from adept_envs import mujoco_env
+from adept_envs.base_robot import BaseRobot
+from adept_envs.utils.configurable import import_class_from_path
+from adept_envs.utils.constants import MODELS_PATH
+
+
+class RobotEnv(mujoco_env.MujocoEnv):
+    """Base environment for all adept robots."""
+
+    # Mapping of robot name to fully qualified class path.
+    # e.g. 'robot': 'adept_envs.dclaw.robot.Robot'
+    # Subclasses should override this to specify the Robot classes they support.
+    ROBOTS = {}
+
+    # Mapping of device path to the calibration file to use. If the device path
+    # is not found, the 'default' key is used.
+    # This can be overriden by subclasses.
+    CALIBRATION_PATHS = {}
+
+    def __init__(
+        self,
+        model_path: str,
+        robot: BaseRobot,
+        frame_skip: int,
+        camera_settings: Optional[Dict] = None,
+    ):
+        """Initializes a robotics environment.
+
+        Args:
+            model_path: The path to the model to run. Relative paths will be
+              interpreted as relative to the 'adept_models' folder.
+            robot: The Robot object to use.
+            frame_skip: The number of simulation steps per environment step. On
+              hardware this influences the duration of each environment step.
+            camera_settings: Settings to initialize the simulation camera. This
+              can contain the keys `distance`, `azimuth`, and `elevation`.
+        """
+        self._robot = robot
+
+        # Initial pose for first step.
+        self.desired_pose = np.zeros(self.n_jnt)
+
+        if not model_path.startswith("/"):
+            model_path = os.path.abspath(os.path.join(MODELS_PATH, model_path))
+
+        self.remote_viz = None
+
+        try:
+            from adept_envs.utils.remote_viz import RemoteViz
+
+            self.remote_viz = RemoteViz(model_path)
+        except ImportError:
+            pass
+
+        self._initializing = True
+        super(RobotEnv, self).__init__(
+            model_path, frame_skip, camera_settings=camera_settings
+        )
+        self._initializing = False
+
+    @property
+    def robot(self):
+        return self._robot
+
+    @property
+    def n_jnt(self):
+        return self._robot.n_jnt
+
+    @property
+    def n_obj(self):
+        return self._robot.n_obj
+
+    @property
+    def skip(self):
+        """Alias for frame_skip.
+
+        Needed for MJRL.
+        """
+        return self.frame_skip
+
+    @property
+    def initializing(self):
+        return self._initializing
+
+    def close_env(self):
+        if self._robot is not None:
+            self._robot.close()
+
+    def make_robot(
+        self,
+        n_jnt,
+        n_obj=0,
+        is_hardware=False,
+        device_name=None,
+        legacy=False,
+        **kwargs
+    ):
+        """Creates a new robot for the environment.
+
+        Args:
+            n_jnt: The number of joints in the robot.
+            n_obj: The number of object joints in the robot environment.
+            is_hardware: Whether to run on hardware or not.
+            device_name: The device path for the robot hardware.
+            legacy: If true, runs using direct dynamixel communication rather
+              than DDS.
+            kwargs: See BaseRobot for other parameters.
+
+        Returns:
+            A Robot object.
+        """
+        if not self.ROBOTS:
+            raise NotImplementedError("Subclasses must override ROBOTS.")
+
+        if is_hardware and not device_name:
+            raise ValueError("Must provide device name if running on hardware.")
+
+        robot_name = "dds_robot" if not legacy and is_hardware else "robot"
+        if robot_name not in self.ROBOTS:
+            raise KeyError(
+                "Unsupported robot '{}', available: {}".format(
+                    robot_name, list(self.ROBOTS.keys())
+                )
+            )
+
+        cls = import_class_from_path(self.ROBOTS[robot_name])
+
+        calibration_path = None
+        if self.CALIBRATION_PATHS:
+            if not device_name:
+                calibration_name = "default"
+            elif device_name not in self.CALIBRATION_PATHS:
+                print(
+                    'Device "{}" not in CALIBRATION_PATHS; using default.'.format(
+                        device_name
+                    )
+                )
+                calibration_name = "default"
+            else:
+                calibration_name = device_name
+
+            calibration_path = self.CALIBRATION_PATHS[calibration_name]
+            if not os.path.isfile(calibration_path):
+                raise OSError(
+                    "Could not find calibration file at: {}".format(calibration_path)
+                )
+
+        return cls(
+            n_jnt,
+            n_obj,
+            is_hardware=is_hardware,
+            device_name=device_name,
+            calibration_path=calibration_path,
+            **kwargs
+        )

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/adept_envs/simulation/module.py

@@ -0,0 +1,135 @@
+#!/usr/bin/python
+#
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Module for caching Python modules related to simulation."""
+
+import sys
+
+_MUJOCO_PY_MODULE = None
+
+_DM_MUJOCO_MODULE = None
+_DM_VIEWER_MODULE = None
+_DM_RENDER_MODULE = None
+
+_GLFW_MODULE = None
+
+
+def get_mujoco_py():
+    """Returns the mujoco_py module."""
+    global _MUJOCO_PY_MODULE
+    if _MUJOCO_PY_MODULE:
+        return _MUJOCO_PY_MODULE
+    try:
+        import mujoco_py
+
+        # Override the warning function.
+        from mujoco_py.builder import cymj
+
+        cymj.set_warning_callback(_mj_warning_fn)
+    except ImportError:
+        print(
+            "Failed to import mujoco_py. Ensure that mujoco_py (using MuJoCo "
+            "v1.50) is installed.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    _MUJOCO_PY_MODULE = mujoco_py
+    return mujoco_py
+
+
+def get_mujoco_py_mjlib():
+    """Returns the mujoco_py mjlib module."""
+
+    class MjlibDelegate:
+        """Wrapper that forwards mjlib calls."""
+
+        def __init__(self, lib):
+            self._lib = lib
+
+        def __getattr__(self, name: str):
+            if name.startswith("mj"):
+                return getattr(self._lib, "_" + name)
+            raise AttributeError(name)
+
+    return MjlibDelegate(get_mujoco_py().cymj)
+
+
+def get_dm_mujoco():
+    """Returns the DM Control mujoco module."""
+    global _DM_MUJOCO_MODULE
+    if _DM_MUJOCO_MODULE:
+        return _DM_MUJOCO_MODULE
+    try:
+        from dm_control import mujoco
+    except ImportError:
+        print(
+            "Failed to import dm_control.mujoco. Ensure that dm_control (using "
+            "MuJoCo v2.00) is installed.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    _DM_MUJOCO_MODULE = mujoco
+    return mujoco
+
+
+def get_dm_viewer():
+    """Returns the DM Control viewer module."""
+    global _DM_VIEWER_MODULE
+    if _DM_VIEWER_MODULE:
+        return _DM_VIEWER_MODULE
+    try:
+        from dm_control import viewer
+    except ImportError:
+        print(
+            "Failed to import dm_control.viewer. Ensure that dm_control (using "
+            "MuJoCo v2.00) is installed.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    _DM_VIEWER_MODULE = viewer
+    return viewer
+
+
+def get_dm_render():
+    """Returns the DM Control render module."""
+    global _DM_RENDER_MODULE
+    if _DM_RENDER_MODULE:
+        return _DM_RENDER_MODULE
+    try:
+        try:
+            from dm_control import _render
+
+            render = _render
+        except ImportError:
+            print("Warning: DM Control is out of date.")
+            from dm_control import render
+    except ImportError:
+        print(
+            "Failed to import dm_control.render. Ensure that dm_control (using "
+            "MuJoCo v2.00) is installed.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    _DM_RENDER_MODULE = render
+    return render
+
+
+def _mj_warning_fn(warn_data: bytes):
+    """Warning function override for mujoco_py."""
+    print(
+        "WARNING: Mujoco simulation is unstable (has NaNs): {}".format(
+            warn_data.decode()
+        )
+    )

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/adept_envs/simulation/renderer.py

@@ -0,0 +1,304 @@
+#!/usr/bin/python
+#
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Module for viewing Physics objects in the DM Control viewer."""
+
+import abc
+import enum
+import sys
+from typing import Dict, Optional
+
+import numpy as np
+
+from adept_envs.simulation import module
+
+# Default window dimensions.
+DEFAULT_WINDOW_WIDTH = 1024
+DEFAULT_WINDOW_HEIGHT = 768
+
+DEFAULT_WINDOW_TITLE = "MuJoCo Viewer"
+
+_MAX_RENDERBUFFER_SIZE = 2048
+
+
+class RenderMode(enum.Enum):
+    """Rendering modes for offscreen rendering."""
+
+    RGB = 0
+    DEPTH = 1
+    SEGMENTATION = 2
+
+
+class Renderer(abc.ABC):
+    """Base interface for rendering simulations."""
+
+    def __init__(self, camera_settings: Optional[Dict] = None):
+        self._camera_settings = camera_settings
+
+    @abc.abstractmethod
+    def close(self):
+        """Cleans up any resources being used by the renderer."""
+
+    @abc.abstractmethod
+    def render_to_window(self):
+        """Renders the simulation to a window."""
+
+    @abc.abstractmethod
+    def render_offscreen(
+        self,
+        width: int,
+        height: int,
+        mode: RenderMode = RenderMode.RGB,
+        camera_id: int = -1,
+    ) -> np.ndarray:
+        """Renders the camera view as a NumPy array of pixels.
+
+        Args:
+            width: The viewport width (pixels).
+            height: The viewport height (pixels).
+            mode: The rendering mode.
+            camera_id: The ID of the camera to render from. By default, uses
+                the free camera.
+
+        Returns:
+            A NumPy array of the pixels.
+        """
+
+    def _update_camera(self, camera):
+        """Updates the given camera to move to the initial settings."""
+        if not self._camera_settings:
+            return
+        distance = self._camera_settings.get("distance")
+        azimuth = self._camera_settings.get("azimuth")
+        elevation = self._camera_settings.get("elevation")
+        lookat = self._camera_settings.get("lookat")
+
+        if distance is not None:
+            camera.distance = distance
+        if azimuth is not None:
+            camera.azimuth = azimuth
+        if elevation is not None:
+            camera.elevation = elevation
+        if lookat is not None:
+            camera.lookat[:] = lookat
+
+
+class MjPyRenderer(Renderer):
+    """Class for rendering mujoco_py simulations."""
+
+    def __init__(self, sim, **kwargs):
+        assert isinstance(
+            sim, module.get_mujoco_py().MjSim
+        ), "MjPyRenderer takes a mujoco_py MjSim object."
+        super().__init__(**kwargs)
+        self._sim = sim
+        self._onscreen_renderer = None
+        self._offscreen_renderer = None
+
+    def render_to_window(self):
+        """Renders the simulation to a window."""
+        if not self._onscreen_renderer:
+            self._onscreen_renderer = module.get_mujoco_py().MjViewer(self._sim)
+            self._update_camera(self._onscreen_renderer.cam)
+
+        self._onscreen_renderer.render()
+
+    def render_offscreen(
+        self,
+        width: int,
+        height: int,
+        mode: RenderMode = RenderMode.RGB,
+        camera_id: int = -1,
+    ) -> np.ndarray:
+        """Renders the camera view as a NumPy array of pixels.
+
+        Args:
+            width: The viewport width (pixels).
+            height: The viewport height (pixels).
+            mode: The rendering mode.
+            camera_id: The ID of the camera to render from. By default, uses
+                the free camera.
+
+        Returns:
+            A NumPy array of the pixels.
+        """
+        if not self._offscreen_renderer:
+            self._offscreen_renderer = module.get_mujoco_py().MjRenderContextOffscreen(
+                self._sim
+            )
+
+        # Update the camera configuration for the free-camera.
+        if camera_id == -1:
+            self._update_camera(self._offscreen_renderer.cam)
+
+        self._offscreen_renderer.render(width, height, camera_id)
+        if mode == RenderMode.RGB:
+            data = self._offscreen_renderer.read_pixels(width, height, depth=False)
+            # Original image is upside-down, so flip it
+            return data[::-1, :, :]
+        elif mode == RenderMode.DEPTH:
+            data = self._offscreen_renderer.read_pixels(width, height, depth=True)[1]
+            # Original image is upside-down, so flip it
+            return data[::-1, :]
+        else:
+            raise NotImplementedError(mode)
+
+    def close(self):
+        """Cleans up any resources being used by the renderer."""
+
+
+class DMRenderer(Renderer):
+    """Class for rendering DM Control Physics objects."""
+
+    def __init__(self, physics, **kwargs):
+        assert isinstance(
+            physics, module.get_dm_mujoco().Physics
+        ), "DMRenderer takes a DM Control Physics object."
+        super().__init__(**kwargs)
+        self._physics = physics
+        self._window = None
+
+        # Set the camera to lookat the center of the geoms. (mujoco_py does
+        # this automatically.
+        if "lookat" not in self._camera_settings:
+            self._camera_settings["lookat"] = [
+                np.median(self._physics.data.geom_xpos[:, i]) for i in range(3)
+            ]
+
+    def render_to_window(self):
+        """Renders the Physics object to a window.
+
+        The window continuously renders the Physics in a separate
+        thread.
+
+        This function is a no-op if the window was already created.
+        """
+        if not self._window:
+            self._window = DMRenderWindow()
+            self._window.load_model(self._physics)
+            self._update_camera(self._window.camera)
+        self._window.run_frame()
+
+    def render_offscreen(
+        self,
+        width: int,
+        height: int,
+        mode: RenderMode = RenderMode.RGB,
+        camera_id: int = -1,
+    ) -> np.ndarray:
+        """Renders the camera view as a NumPy array of pixels.
+
+        Args:
+            width: The viewport width (pixels).
+            height: The viewport height (pixels).
+            mode: The rendering mode.
+            camera_id: The ID of the camera to render from. By default, uses
+                the free camera.
+
+        Returns:
+            A NumPy array of the pixels.
+        """
+        mujoco = module.get_dm_mujoco()
+        # TODO(michaelahn): Consider caching the camera.
+        camera = mujoco.Camera(
+            physics=self._physics, height=height, width=width, camera_id=camera_id
+        )
+
+        # Update the camera configuration for the free-camera.
+        if camera_id == -1:
+            self._update_camera(
+                camera._render_camera,  # pylint: disable=protected-access
+            )
+
+        image = camera.render(
+            depth=(mode == RenderMode.DEPTH),
+            segmentation=(mode == RenderMode.SEGMENTATION),
+        )
+        camera._scene.free()  # pylint: disable=protected-access
+        return image
+
+    def close(self):
+        """Cleans up any resources being used by the renderer."""
+        if self._window:
+            self._window.close()
+            self._window = None
+
+
+class DMRenderWindow:
+    """Class that encapsulates a graphical window."""
+
+    def __init__(
+        self,
+        width: int = DEFAULT_WINDOW_WIDTH,
+        height: int = DEFAULT_WINDOW_HEIGHT,
+        title: str = DEFAULT_WINDOW_TITLE,
+    ):
+        """Creates a graphical render window.
+
+        Args:
+            width: The width of the window.
+            height: The height of the window.
+            title: The title of the window.
+        """
+        dmv = module.get_dm_viewer()
+        self._viewport = dmv.renderer.Viewport(width, height)
+        self._window = dmv.gui.RenderWindow(width, height, title)
+        self._viewer = dmv.viewer.Viewer(
+            self._viewport, self._window.mouse, self._window.keyboard
+        )
+        self._draw_surface = None
+        self._renderer = dmv.renderer.NullRenderer()
+
+    @property
+    def camera(self):
+        return self._viewer._camera._camera
+
+    def close(self):
+        self._viewer.deinitialize()
+        self._renderer.release()
+        self._draw_surface.free()
+        self._window.close()
+
+    def load_model(self, physics):
+        """Loads the given Physics object to render."""
+        self._viewer.deinitialize()
+
+        self._draw_surface = module.get_dm_render().Renderer(
+            max_width=_MAX_RENDERBUFFER_SIZE, max_height=_MAX_RENDERBUFFER_SIZE
+        )
+        self._renderer = module.get_dm_viewer().renderer.OffScreenRenderer(
+            physics.model, self._draw_surface
+        )
+
+        self._viewer.initialize(physics, self._renderer, touchpad=False)
+
+    def run_frame(self):
+        """Renders one frame of the simulation.
+
+        NOTE: This is extremely slow at the moment.
+        """
+        glfw = module.get_dm_viewer().gui.glfw_gui.glfw
+        glfw_window = self._window._context.window
+        if glfw.window_should_close(glfw_window):
+            sys.exit(0)
+
+        self._viewport.set_size(*self._window.shape)
+        self._viewer.render()
+        pixels = self._renderer.pixels
+
+        with self._window._context.make_current() as ctx:
+            ctx.call(self._window._update_gui_on_render_thread, glfw_window, pixels)
+        self._window._mouse.process_events()
+        self._window._keyboard.process_events()

uvd/envs/franka_kitchen/relay-policy-learning/adept_envs/adept_envs/simulation/sim_robot.py

@@ -0,0 +1,133 @@
+#!/usr/bin/python
+#
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Module for loading MuJoCo models."""
+
+import os
+from typing import Dict, Optional
+
+from adept_envs.simulation import module
+from adept_envs.simulation.renderer import DMRenderer, MjPyRenderer, RenderMode
+
+
+class MujocoSimRobot:
+    """Class that encapsulates a MuJoCo simulation.
+
+    This class exposes methods that are agnostic to the simulation backend.
+    Two backends are supported:
+    1. mujoco_py - MuJoCo v1.50
+    2. dm_control - MuJoCo v2.00
+    """
+
+    def __init__(
+        self,
+        model_file: str,
+        use_dm_backend: bool = False,
+        camera_settings: Optional[Dict] = None,
+    ):
+        """Initializes a new simulation.
+
+        Args:
+            model_file: The MuJoCo XML model file to load.
+            use_dm_backend: If True, uses DM Control's Physics (MuJoCo v2.0) as
+              the backend for the simulation. Otherwise, uses mujoco_py (MuJoCo
+              v1.5) as the backend.
+            camera_settings: Settings to initialize the renderer's camera. This
+              can contain the keys `distance`, `azimuth`, and `elevation`.
+        """
+        self._use_dm_backend = use_dm_backend
+
+        if not os.path.isfile(model_file):
+            raise ValueError(
+                "[MujocoSimRobot] Invalid model file path: {}".format(model_file)
+            )
+
+        if self._use_dm_backend:
+            dm_mujoco = module.get_dm_mujoco()
+            if model_file.endswith(".mjb"):
+                self.sim = dm_mujoco.Physics.from_binary_path(model_file)
+            else:
+                self.sim = dm_mujoco.Physics.from_xml_path(model_file)
+            self.model = self.sim.model
+            self._patch_mjlib_accessors(self.model, self.sim.data)
+            self.renderer = DMRenderer(self.sim, camera_settings=camera_settings)
+        else:  # Use mujoco_py
+            mujoco_py = module.get_mujoco_py()
+            self.model = mujoco_py.load_model_from_path(model_file)
+            self.sim = mujoco_py.MjSim(self.model)
+            self.renderer = MjPyRenderer(self.sim, camera_settings=camera_settings)
+
+        self.data = self.sim.data
+
+    def close(self):
+        """Cleans up any resources being used by the simulation."""
+        self.renderer.close()
+
+    def save_binary(self, path: str):
+        """Saves the loaded model to a binary .mjb file."""
+        if os.path.exists(path):
+            raise ValueError("[MujocoSimRobot] Path already exists: {}".format(path))
+        if not path.endswith(".mjb"):
+            path = path + ".mjb"
+        if self._use_dm_backend:
+            self.model.save_binary(path)
+        else:
+            with open(path, "wb") as f:
+                f.write(self.model.get_mjb())
+
+    def get_mjlib(self):
+        """Returns an object that exposes the low-level MuJoCo API."""
+        if self._use_dm_backend:
+            return module.get_dm_mujoco().wrapper.mjbindings.mjlib
+        else:
+            return module.get_mujoco_py_mjlib()
+
+    def _patch_mjlib_accessors(self, model, data):
+        """Adds accessors to the DM Control objects to support mujoco_py
+        API."""
+        assert self._use_dm_backend
+        mjlib = self.get_mjlib()
+
+        def name2id(type_name, name):
+            obj_id = mjlib.mj_name2id(
+                model.ptr, mjlib.mju_str2Type(type_name.encode()), name.encode()
+            )
+            if obj_id < 0:
+                raise ValueError('No {} with name "{}" exists.'.format(type_name, name))
+            return obj_id
+
+        if not hasattr(model, "body_name2id"):
+            model.body_name2id = lambda name: name2id("body", name)
+
+        if not hasattr(model, "geom_name2id"):
+            model.geom_name2id = lambda name: name2id("geom", name)
+
+        if not hasattr(model, "site_name2id"):
+            model.site_name2id = lambda name: name2id("site", name)
+
+        if not hasattr(model, "joint_name2id"):
+            model.joint_name2id = lambda name: name2id("joint", name)
+
+        if not hasattr(model, "actuator_name2id"):
+            model.actuator_name2id = lambda name: name2id("actuator", name)
+
+        if not hasattr(model, "camera_name2id"):
+            model.camera_name2id = lambda name: name2id("camera", name)
+
+        if not hasattr(data, "body_xpos"):
+            data.body_xpos = data.xpos
+
+        if not hasattr(data, "body_xquat"):
+            data.body_xquat = data.xquat