Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,96 @@ 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
+aloha_robot_project/Miniconda3-latest-Linux-x86_64.sh filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_A/meshes/light_link.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_A/textures/dark_wood.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_A/textures/dark_wood__black_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_A/textures/dark_wood__gray_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_A/textures/light_wood.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_A/textures/light_wood__black_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_A/textures/light_wood__gray_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_A/textures/wood.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_A/textures/wood__black_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_A/textures/wood__gray_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_B/meshes/light_link.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_B/textures/dark_wood.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_B/textures/dark_wood__black_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_B/textures/dark_wood__gray_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_B/textures/light_wood.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_B/textures/light_wood__black_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_B/textures/light_wood__gray_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_B/textures/wood.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_B/textures/wood__black_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_B/textures/wood__gray_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_C/meshes/light_link.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_C/textures/dark_wood.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_C/textures/dark_wood__black_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_C/textures/dark_wood__gray_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_C/textures/light_wood.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_C/textures/light_wood__black_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_C/textures/light_wood__gray_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_C/textures/wood.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_C/textures/wood__black_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_C/textures/wood__gray_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_D/meshes/light_link.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_D/textures/dark_wood.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_D/textures/dark_wood__black_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_D/textures/dark_wood__gray_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_D/textures/light_wood.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_D/textures/light_wood__black_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_D/textures/light_wood__gray_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_D/textures/wood.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_D/textures/wood__black_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/calvin_table_D/textures/wood__gray_handle.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/franka_panda/meshes/visual/Assem1.SLDASM filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/franka_panda/meshes/visual/FRANKA_Finger.SLDPRT filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/franka_panda/meshes/visual/digit.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/franka_panda/meshes/visual/digit_gel_only.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/data/franka_panda/meshes/visual/finger.SLDPRT filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/allegro_hand_description/meshes/base_link.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/allegro_hand_description/meshes/base_link_left.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/allegro_hand_description/meshes/digit.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/allegro_hand_description/meshes/link_12.0_left.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/allegro_hand_description/meshes/link_12.0_right.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_mp1/l6.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_mp3/l0.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_mp3/l1.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_pv/base.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_pv/head.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_pv/l0.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_pv/l1.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_pv/l2.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_pv/l3.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_pv/l4.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_pv/l5.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_pv/l6.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/sawyer_robot/sawyer_description/meshes/sawyer_pv/pedestal.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/wsg50/WSG-FMF.stl filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/examples/wsg50/digit.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_mp1/l6.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_mp3/l0.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_mp3/l1.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_pv/base.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_pv/head.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_pv/l0.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_pv/l1.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_pv/l2.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_pv/l3.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_pv/l4.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_pv/l5.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_pv/l6.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/sawyer_robot/sawyer_description/meshes/sawyer_pv/pedestal.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/wsg50/WSG-FMF.stl filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/grasp_stability/setup/wsg50/digit.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/experiments/rolling/setup/sensors/digit.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/meshes/digit.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/meshes/omnitact.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/meshes/omnitact_mod_pybullet.STL filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/website/static/img/demo_allegro.gif filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/website/static/img/demo_digit.gif filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/website/static/img/demo_grasp.gif filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/website/static/img/demo_omnitact.gif filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/website/static/img/demo_rolling.gif filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/calvin_env/tacto/website/static/img/demo_shadow.gif filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/media/sensors.png filter=lfs diff=lfs merge=lfs -text
+aloha_robot_project/calvin/media/teaser.png filter=lfs diff=lfs merge=lfs -text

aloha_robot_project/Miniconda3-latest-Linux-x86_64.sh

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e0b10e050e8928e2eb9aad2c522ee3b5d31d30048b8a9997663a8a460d538cef
+size 156772981

aloha_robot_project/calvin/.flake8

@@ -0,0 +1,15 @@
+[flake8]
+exclude = .git
+# Default is 79 in PEP 8
+max-line-length = 120
+select = E,F,W,C
+ignore=W503, # line break before binary operator, need for black
+       E203, # whitespace before ':'. Opposite convention enforced by black
+       E731, # do not assign a lambda expression, use a def
+       E722,
+       F401,
+       F841,
+       E402,  # module level import not at top of file
+       E741,  # ambiguous variable name
+       E501, # line too long. Handled by black
+       C406,  # Unnecessary list literal - rewrite as a dict literal

aloha_robot_project/calvin/.github/workflows/codeql.yml

@@ -0,0 +1,41 @@
+name: "CodeQL"
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+  schedule:
+    - cron: "33 22 * * 4"
+
+jobs:
+  analyze:
+    name: Analyze
+    runs-on: ubuntu-latest
+    permissions:
+      actions: read
+      contents: read
+      security-events: write
+
+    strategy:
+      fail-fast: false
+      matrix:
+        language: [ python ]
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+
+      - name: Initialize CodeQL
+        uses: github/codeql-action/init@v2
+        with:
+          languages: ${{ matrix.language }}
+          queries: +security-and-quality
+
+      - name: Autobuild
+        uses: github/codeql-action/autobuild@v2
+
+      - name: Perform CodeQL Analysis
+        uses: github/codeql-action/analyze@v2
+        with:
+          category: "/language:${{ matrix.language }}"

aloha_robot_project/calvin/.gitignore

@@ -0,0 +1,146 @@
+# 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/
+pip-wheel-metadata/
+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/
+
+# 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
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.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
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__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/
+
+# vscode
+.vscode
+
+*/runs/*
+
+dataset/task_D_D/*
+dataset/task_ABCD_D/*
+dataset/task_ABC_D/*
+
+.idea/*
+
+# Ignore venv created as part of installation instructions
+calvin_venv/*
+
+# Ignore datasets as downloaded according to instructions
+dataset/*

aloha_robot_project/calvin/.gitmodules

@@ -0,0 +1,3 @@
+[submodule "calvin_env"]
+	path = calvin_env
+	url = https://github.com/mees/calvin_env.git

aloha_robot_project/calvin/.pre-commit-config.yaml

@@ -0,0 +1,33 @@
+default_language_version:
+    python: python3.8
+repos:
+  - repo: https://github.com/psf/black
+    rev: 22.3.0
+    hooks:
+      - id: black
+        language_version: python3.8
+
+  - repo: https://gitlab.com/pycqa/flake8
+    rev: 3.8.4
+    hooks:
+      - id: flake8
+        additional_dependencies: [-e, "git+git://github.com/pycqa/pyflakes.git@c72d6cf#egg=pyflakes"]
+
+  - repo: https://github.com/pycqa/isort
+    rev: 5.7.0
+    hooks:
+      - id: isort
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v0.812
+    hooks:
+      - id: mypy
+        args: [--ignore-missing-imports, --warn-no-return, --warn-redundant-casts, --disallow-incomplete-defs]
+        additional_dependencies: [pytorch-lightning==1.4.9, torch==1.10.0, numpy]
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.0.1
+    hooks:
+      - id: check-yaml
+      - id: trailing-whitespace
+      - id: end-of-file-fixer

aloha_robot_project/calvin/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Oier Mees
+
+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.

aloha_robot_project/calvin/README.md

@@ -0,0 +1,377 @@
+# CALVIN
+[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![Language grade: Python](https://img.shields.io/lgtm/grade/python/g/mees/calvin.svg?logo=lgtm&logoWidth=18)](https://lgtm.com/projects/g/mees/calvin/context:python)
+[![Total alerts](https://img.shields.io/lgtm/alerts/g/mees/calvin.svg?logo=lgtm&logoWidth=18)](https://lgtm.com/projects/g/mees/calvin/alerts/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+
+[<b>CALVIN - A benchmark for Language-Conditioned Policy Learning for Long-Horizon Robot Manipulation Tasks</b>](https://arxiv.org/pdf/2112.03227.pdf)
+
+[Oier Mees](https://www.oiermees.com/), [Lukas Hermann](https://lukashermann.github.io/), [Erick Rosete](https://www.erickrosete.com/), [Wolfram Burgard](http://www2.informatik.uni-freiburg.de/~burgard)
+
+#### CALVIN won the 2022 IEEE Robotics and Automation Letters (RA-L) Best Paper Award!
+
+
+ We present **CALVIN** (**C**omposing **A**ctions from **L**anguage and **Vi**sio**n**), an open-source simulated benchmark to learn long-horizon language-conditioned tasks.
+Our aim is to make it possible to develop agents that can solve many robotic manipulation tasks over a long horizon, from onboard sensors, and specified only via human language. CALVIN tasks are more complex in terms of sequence length, action space, and language than existing vision-and-language task datasets and supports flexible specification of sensor
+suites.
+
+![](media/teaser.png)
+
+# :computer:  Quick Start
+To begin, clone this repository locally
+```bash
+git clone --recurse-submodules https://github.com/mees/calvin.git
+$ export CALVIN_ROOT=$(pwd)/calvin
+
+```
+Install requirements:
+```bash
+$ cd $CALVIN_ROOT
+$ conda create -n calvin_venv python=3.8  # or use virtualenv
+$ conda activate calvin_venv
+$ sh install.sh
+```
+If you encounter problems installing pyhash, you might have to downgrade setuptools to a version below 58.
+
+Download dataset (choose which split you want to download with the argument `D`, `ABC` or `ABCD`): \
+If you want to get started without downloading the whole dataset, use the argument `debug` to download a small debug dataset (1.3 GB).
+```bash
+$ cd $CALVIN_ROOT/dataset
+$ sh download_data.sh D | ABC | ABCD | debug
+```
+##	:weight_lifting_man: Train Baseline Agent
+Train baseline models:
+```bash
+$ cd $CALVIN_ROOT/calvin_models/calvin_agent
+$ python training.py datamodule.root_data_dir=/path/to/dataset/ datamodule/datasets=vision_lang_shm
+```
+The `vision_lang_shm` option loads the CALVIN dataset into shared memory at the beginning of the training,
+speeding up the data loading during training.
+The preparation of the shared memory cache will take some time
+(approx. 20 min at our SLURM cluster). \
+If you want to use the original data loader (e.g. for debugging) just override the command with `datamodule/datasets=vision_lang`. \
+For an additional speed up, you can disable the evaluation callbacks during training by adding `~callbacks/rollout` and `~callbacks/rollout_lh`
+
+You want to scale your training to a multi-gpu setup? Just specify the [number of GPUs](https://pytorch-lightning.readthedocs.io/en/latest/advanced/multi_gpu.html#select-gpu-devices) and DDP will automatically be used
+ for training thanks to [Pytorch Lightning](https://www.pytorchlightning.ai/).
+To train on all available GPUs:
+```bash
+$ python training.py trainer.gpus=-1
+```
+If you have access to a Slurm cluster, follow this [guide](https://github.com/mees/calvin/blob/main/slurm_scripts/README.md).
+
+You can use [Hydra's](https://hydra.cc/) flexible overriding system for changing hyperparameters.
+For example, to train a model with  rgb images from both static camera and the gripper camera with relative actions:
+```bash
+$ python training.py datamodule/observation_space=lang_rgb_static_gripper_rel_act model/perceptual_encoder=gripper_cam
+```
+To train a model with RGB-D from both cameras:
+```bash
+$ python training.py datamodule/observation_space=lang_rgbd_both model/perceptual_encoder=RGBD_both
+```
+To train a model with rgb images from the static camera and visual tactile observations with absolute actions:
+```bash
+$ python training.py datamodule/observation_space=lang_rgb_static_tactile_abs_act model/perceptual_encoder=static_RGB_tactile
+```
+
+To see all available hyperparameters:
+```console
+$ python training.py --help
+```
+To resume a training, just override the hydra working directory :
+```console
+$ python training.py hydra.run.dir=runs/my_dir
+```
+
+## :framed_picture: Sensory Observations
+ CALVIN  supports a range of sensors commonly utilized for visuomotor  control:
+1. **Static camera RGB images** - with shape `200x200x3`.
+2. **Static camera Depth maps** - with shape `200x200`.
+3. **Gripper camera RGB images** - with shape `84x84x3`.
+4. **Gripper camera Depth maps** - with shape `84x84`.
+5. **Tactile image** - with shape `120x160x6`.
+6. **Proprioceptive state** - EE position (3), EE orientation in euler angles (3), gripper width (1), joint positions (7), gripper action (1).
+
+<p align="center">
+<img src="media/sensors.png" alt="" width="50%">
+</p>
+
+## :joystick: Action Space
+In CALVIN, the  agent  must perform  closed-loop  continuous  control  to  follow  unconstrained  language  instructions  characterizing  complex  robot manipulation tasks, sending continuous actions to the robot at  30hz.
+In  order  to  give  researchers  and  practitioners  the freedom to experiment with different action spaces, CALVIN supports  the following actions spaces:
+1. **Absolute cartesian pose**  - EE position (3), EE orientation in euler angles (3),  gripper action (1).
+2. **Relative cartesian displacement**  - EE position (3), EE orientation in euler angles (3),  gripper action (1).
+3. **Joint action** -  Joint positions (7),  gripper action (1).
+
+For more information, please refer to this more detailed [README](https://github.com/mees/calvin/blob/main/dataset/README.md).
+
+## :muscle: Evaluation: The Calvin Challenge
+### Long-horizon Multi-task Language Control (LH-MTLC)
+The  aim  of  the  CALVIN  benchmark  is  to  evaluate  the learning  of  long-horizon  language-conditioned  continuous control  policies.  In  this  setting,  a  single  agent  must  solve complex  manipulation  tasks  by  understanding  a  series  of unconstrained  language  expressions  in  a  row,  e.g.,  “open the  drawer. . . pick  up  the  blue  block. . . now  push  the  block into the drawer. . . now open the sliding door”.
+We provide  an  evaluation  protocol  with  evaluation  modes  of varying  difficulty  by  choosing  different  combinations  of sensor  suites  and  amounts  of  training  environments.
+To avoid a biased initial position, the robot is reset to a neutral position before every multi-step sequence.
+
+To evaluate a trained calvin baseline agent, run the following command:
+
+```
+$ cd $CALVIN_ROOT/calvin_models/calvin_agent
+$ python evaluation/evaluate_policy.py --dataset_path <PATH/TO/DATASET> --train_folder <PATH/TO/TRAINING/FOLDER>
+```
+Optional arguments:
+
+- `--checkpoint <PATH/TO/CHECKPOINT>`: by default, the evaluation loads the last checkpoint in the training log directory.
+You can instead specify the path to another checkpoint by adding this to the evaluation command.
+- `--debug`: print debug information and visualize environment.
+
+If you want to evaluate your own model architecture on the CALVIN challenge, you can implement the `CustomModel` class in `evaluate_policy.py`
+as an interface to your agent. You need to implement the following methods:
+
+- \_\_init__():
+  gets called once at the beginning of the evaluation.
+- reset(): gets called at the beginning of each evaluation sequence.
+- step(obs, goal): gets called every step and returns the predicted action.
+
+Then evaluate the model by running:
+```
+$ python evaluation/evaluate_policy.py --dataset_path <PATH/TO/DATASET> --custom_model
+```
+
+You are also free to use your own language model instead of using the precomputed language embeddings provided by CALVIN.
+For this, implement `CustomLangEmbeddings` in `evaluate_policy.py` and add `--custom_lang_embeddings` to the evaluation command.
+
+### Multi-task Language Control (MTLC)
+Alternatively, you can evaluate the policy on single tasks and without resetting the robot to a neutral position.
+Note that this evaluation is currently only available for our baseline agent.
+```
+$ python evaluation/evaluate_policy_singlestep.py --dataset_path <PATH/TO/DATASET> --train_folder <PATH/TO/TRAINING/FOLDER> [--checkpoint <PATH/TO/CHECKPOINT>] [--debug]
+```
+
+### Pre-trained Model
+Download the [MCIL](http://calvin.cs.uni-freiburg.de/model_weights/D_D_static_rgb_baseline.zip) model checkpoint trained on the static camera rgb images on environment D.
+```
+$ wget http://calvin.cs.uni-freiburg.de/model_weights/D_D_static_rgb_baseline.zip
+$ unzip D_D_static_rgb_baseline.zip
+```
+## :speech_balloon: Relabeling Raw Language Annotations
+You want to try learning language conditioned policies in CALVIN with a new awesome language model?
+
+We provide an [example script](https://github.com/mees/calvin/blob/main/calvin_models/calvin_agent/utils/relabel_with_new_lang_model.py) to relabel the annotations with different language model provided in [SBert](https://www.sbert.net/docs/pretrained_models.html), such as the larger MPNet (paraphrase-mpnet-base-v2) or its corresponding multilingual model (paraphrase-multilingual-mpnet-base-v2).
+The supported options are "mini", "mpnet" and "multi". If you want to try different SBert models, just change the model name [here](https://github.com/mees/calvin/blob/main/calvin_models/calvin_agent/models/encoders/language_network.py#L18).
+```
+cd $CALVIN_ROOT/calvin_models/calvin_agent
+python utils/relabel_with_new_lang_model.py +path=$CALVIN_ROOT/dataset/task_D_D/ +name_folder=new_lang_model_folder model.nlp_model=mpnet
+```
+If you additionally want to sample different language annotations for each sequence (from the same task annotations) in the training split run the same command with the parameter `reannotate=true`.
+
+## :chart_with_upwards_trend: SOTA Models
+Open-source models that outperform the MCIL baselines from CALVIN:
+
+For a detailed overview of the evaluation performances, have a look at our **[LEADERBOARD](http://calvin.cs.uni-freiburg.de/)**.
+
+<br>
+<b> Grounding Language with Visual Affordances over Unstructured Data</b>
+<br>
+Oier Mees, Jessica Borja-Diaz, Wolfram Burgard
+<br>
+<a href="https://arxiv.org/pdf/2210.01911.pdf"> Paper</a>, <a href="https://github.com/mees/hulc2"> Code </a>
+
+<b> FLOWER: Democratizing Generalist Robot Policies with Efficient Vision-Language-Action Flow Policies </b>
+<br>
+Moritz Reuss, Hongyi Zhou, Marcel Rühle, Ömer Erdinç Yağmurlu, Fabian Otto, Rudolf Lioutikov
+<br>
+<a href="https://arxiv.org/pdf/2509.04996"> Paper</a>, <a href="https://intuitive-robots.github.io/flower_vla/"> Code </a>
+
+
+<b> Unified Vision-Language-Action Model </b>
+<br>
+Yuqi Wang, Xinghang Li, Wenxuan Wang, Junbo Zhang, Yingyan Li, Yuntao Chen, Xinlong Wang, Zhaoxiang Zhang
+<br>
+<a href="https://arxiv.org/pdf/2506.19850"> Paper</a>, <a href="https://robertwyq.github.io/univla.github.io/"> Code </a>
+
+<b> Predictive Inverse Dynamics Models are Scalable Learners for Robotic Manipulation </b>
+<br>
+Yang Tian, Sizhe Yang, Jia Zeng, Ping Wang, Dahua Lin, Hao Dong, Jiangmiao Pang
+<br>
+<a href="https://arxiv.org/pdf/2412.15109"> Paper</a>, <a href="https://github.com/OpenRobotLab/Seer/"> Code </a>
+
+<b> Diffusion Transformer Policy: Scaling Diffusion Transformer for Generalist Vision-Language-Action Learning </b>
+<br>
+Zhi Hou, Tianyi Zhang, Yuwen Xiong, Hengjun Pu, Chengyang Zhao, Ronglei Tong, Yu Qiao, Jifeng Dai, Yuntao Chen
+<br>
+<a href="https://arxiv.org/pdf/2410.15959"> Paper</a>, <a href="https://github.com/zhihou7/dit_policy_vla"> Code </a>
+
+<b> GR-MG: Leveraging Partially Annotated Data via Multi-Modal Goal Conditioned Policy </b>
+<br>
+Peiyan Li, Hongtao Wu, Yan Huang, Chilam Cheang, Liang Wang, Tao Kong
+<br>
+<a href="https://arxiv.org/pdf/2408.14368"> Paper</a>, <a href="https://github.com/bytedance/GR-MG/"> Code </a>
+
+<b> GHIL-Glue: Hierarchical Control with Filtered Subgoal Images </b>
+<br>
+Kyle B Hatch, Ashwin Balakrishna, Oier Mees, Suraj Nair, Seohong Park, Blake Wulfe, Masha Itkina, Benjamin Eysenbach, Sergey Levine, Thomas Kollar, Benjamin Burchfiel
+<br>
+<a href="https://arxiv.org/pdf/2410.20018"> Paper</a>, <a href="https://github.com/kyle-hatch-tri/ghil-glue"> Code </a>
+
+<b> Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers for Multitask Learning </b>
+<br>
+Moritz Reuss, Jyothish Pari, Pulkit Agrawal, Rudolf Lioutikov
+<br>
+<a href="https://arxiv.org/pdf/2412.12953"> Paper</a>, <a href="https://github.com/intuitive-robots/MoDE_Diffusion_Policy"> Code </a>
+
+<b> Incorporating Task Progress Knowledge for Subgoal Generation in Robotic Manipulation through Image Edits </b>
+<br>
+Xuhui Kang, Yen-Ling Kuo
+<br>
+<a href="https://arxiv.org/pdf/2410.11013"> Paper</a>, <a href="https://github.com/Shua-Kang/TaKSIE"> Code </a>
+
+<b> Closed-Loop Visuomotor Control with Generative Expectation for Robotic Manipulation </b>
+<br>
+Qingwen Bu, Jia Zeng, Li Chen, Yanchao Yang, Guyue Zhou, Junchi Yan, Ping Luo, Heming Cui, Yi Ma, Hongyang Li
+<br>
+<a href="https://arxiv.org/pdf/2409.09016"> Paper</a>, <a href="https://github.com/OpenDriveLab/CLOVER"> Code </a>
+
+<b> DeeR-VLA: Dynamic Inference of Multimodal Large Language Models for Efficient Robot Execution </b>
+<br>
+Yang Yue, Yulin Wang, Bingyi Kang, Yizeng Han, Shenzhi Wang, Shiji Song, Jiashi Feng, Gao Huang
+<br>
+<a href="https://arxiv.org/pdf/2411.02359"> Paper</a>, <a href="https://github.com/yueyang130/DeeR-VLA"> Code </a>
+
+<b> RoboUniView: Visual-Language Model with Unified View Representation for Robotic Manipulation </b>
+<br>
+Fanfan Liu, Feng Yan, Liming Zheng, Yiyang Huang, Chengjian Feng, Lin Ma
+<br>
+<a href="https://arxiv.org/pdf/2406.18977v2"> Paper</a>, <a href="https://github.com/liufanfanlff/RoboUniview"> Code </a>
+
+<b> Multimodal Diffusion Transformer: Learning Versatile Behavior from Multimodal Goals </b>
+<br>
+Moritz Reuss, Ömer Erdinç Yağmurlu, Fabian Wenzel, Rudolf Lioutikov
+<br>
+<a href="https://arxiv.org/pdf/2407.05996"> Paper</a>, <a href="https://github.com/intuitive-robots/mdt_policy"> Code </a>
+
+<b> 3D Diffuser Actor: Policy Diffusion with 3D Scene Representations</b>
+<br>
+Tsung-Wei Ke, Nikolaos Gkanatsios, Katerina Fragkiadaki
+<br>
+<a href="https://arxiv.org/pdf/2402.10885.pdf"> Paper</a>, <a href="https://github.com/nickgkan/3d_diffuser_actor"> Code </a>
+
+<b> Unleashing Large-Scale Video Generative Pre-training for Visual Robot Manipulation</b>
+<br>
+Hongtao Wu, Ya Jing, Chilam Cheang, Guangzeng Chen, Jiafeng Xu, Xinghang Li, Minghuan Liu, Hang Li, Tao Kong
+<br>
+<a href="https://arxiv.org/pdf/2312.13139.pdf"> Paper</a>, <a href="https://github.com/bytedance/GR-1"> Code </a>
+
+<b> Vision-Language Foundation Models as Effective Robot Imitators</b>
+<br>
+Xinghang Li, Minghuan Liu, Hanbo Zhang, Cunjun Yu, Jie Xu, Hongtao Wu, Chilam Cheang, Ya Jing, Weinan Zhang, Huaping Liu, Hang Li, and Tao Kong
+<br>
+<a href="https://arxiv.org/pdf/2311.01378.pdf"> Paper</a>, <a href="https://github.com/RoboFlamingo/RoboFlamingo"> Code </a>
+
+<b> Zero-Shot Robotic Manipulation With Pretrained Image-Editing Diffusion Models</b>
+<br>
+Kevin Black, Mitsuhiko Nakamoto, Pranav Atreya, Homer Walke, Chelsea Finn, Aviral Kumar, Sergey Levine
+<br>
+<a href="https://arxiv.org/pdf/2310.10639.pdf"> Paper</a>, <a href="https://github.com/kvablack/susie"> Code </a>
+
+<b> Language Control Diffusion: Efficiently Scaling through Space, Time, and Tasks</b>
+<br>
+Eddie Zhang, Yujie Lu, William Wang, Amy Zhang
+<br>
+<a href="https://arxiv.org/pdf/2210.15629.pdf"> Paper</a>, <a href="https://github.com/ezhang7423/language-control-diffusion"> Code </a>
+
+<b> What Matters in Language Conditioned Robotic Imitation Learning over Unstructured Data</b>
+<br>
+Oier Mees, Lukas Hermann, Wolfram Burgard
+<br>
+<a href="https://arxiv.org/pdf/2204.06252.pdf"> Paper</a>, <a href="https://github.com/lukashermann/hulc"> Code </a>
+
+<b> Language-Conditioned Imitation Learning with Base Skill Priors under Unstructured Data</b>
+<br>
+Hongkuan Zhou, Zhenshan Bing, Xiangtong Yao, Xiaojie Su, Chenguang Yang, Kai Huang, Alios Knoll
+<br>
+<a href="https://arxiv.org/pdf/2305.19075.pdf"> Paper</a>, <a href="https://github.com/hk-zh/spil"> Code
+
+Contact [Oier](https://www.oiermees.com/) to add your model here.
+
+## Reinforcement Learning with CALVIN
+Are you interested in trying  reinforcement learning agents for the different manipulation tasks in the CALVIN environment?
+We provide a [google colab](https://github.com/mees/calvin/blob/main/RL_with_CALVIN.ipynb) to showcase how to leverage the CALVIN task indicators to learn RL agents with a sparse reward.
+
+## FAQ
+
+#### Why do you use EGL rendering?
+We use EGL to move the bullet rendering from cpu (which is the default) to gpu, which is much faster.
+This way, we can also do rollouts during the training of the agent to track its performance.
+By changing from cpu to gpu, the rendered textures change slightly, so be aware of this if you plan on testing pretrained models.
+#### I am training with multiple GPUs and why am I get OOM errors during rollouts?
+PyBullet only recently added an option to select which GPU to use for rendering when using EGL (fix was commited in 3c4cb80
+on Oct 22, 2021, see [here](https://github.com/bulletphysics/bullet3/blob/master/examples/OpenGLWindow/EGLOpenGLWindow.cpp#L134).
+If you have an old version of PyBullet, there is no way to choose the GPU, which can lead to problems on cluster nodes with multiple GPUs, because all instances would be placed on the same GPU, slowing down the rendering and potentially leading to OOM erros.
+
+The fix introduced an environment variable EGL_VISIBLE_DEVICES (similar to CUDA_VISIBLE_DEVICES) which lets you specify the GPU device to render on.
+However, there is one catch: On some machines, the device ids of CUDA and EGL do not match (e.g. CUDA device 0 could be EGL device 3).
+We automatically handle this in our wrapper in calvin_env and find the corresponding egl device id, so you don't have to set EGL_VISIBLE_DEVICES yourself, see [here](https://github.com/mees/calvin_env/blob/main/calvin_env/envs/play_lmp_wrapper.py#L31).
+
+#### I am not interested in the manipulation tasks recorded, can I record different demonstration with teleop?
+Yes, although it is not documented right now, all the code to record data with a VR headset is present in
+calvin_env in [https://github.com/mees/calvin_env/blob/main/calvin_env/vrdatacollector.py](https://github.com/mees/calvin_env/blob/main/calvin_env/vrdatacollector.py)
+
+
+## Changelog
+### 24 Feb 2023
+- Wrong `scene_info.npy` in D dataset. Note that we have updated the corresponding checksum. Please replace as follows:
+```
+cd task_D_D
+wget http://calvin.cs.uni-freiburg.de/scene_info_fix/task_D_D_scene_info.zip
+unzip task_D_D_scene_info.zip && rm task_D_D_scene_info.zip
+```
+
+### 16 Sep 2022
+- **MAJOR BUG IN ABC and ABCD dataset:** If you downloaded these datasets before this date you have to do these fixes:
+   - Wrong language annotations in ABC and ABCD dataset. You can download the corrected language embeddings [here](https://github.com/mees/calvin/blob/main/dataset/README.md#language-embeddings).
+   - Bug in `calvin_env` that only affects the generation of language embeddings.
+   - Wrong `scene_info.npy` in ABC and ABCD dataset. Please replace as follows:
+```
+cd task_ABCD_D
+wget http://calvin.cs.uni-freiburg.de/scene_info_fix/task_ABCD_D_scene_info.zip
+unzip task_ABCD_D_scene_info.zip && rm task_ABCD_D_scene_info.zip
+```
+```
+cd task_ABC_D
+wget http://calvin.cs.uni-freiburg.de/scene_info_fix/task_ABC_D_scene_info.zip
+unzip task_ABC_D_scene_info.zip && rm task_ABC_D_scene_info.zip
+```
+- Added additional language embeddings to dataset.
+
+
+### 15 May 2022
+- Added shared memory dataset loader for faster training. Refactored data loading classes.
+
+### 7 Feb 2022
+- Minor changes to the distribution of tasks in the long-horizon multi-step sequences.
+- Changes to the task success criteria of pushing and lifting.
+- Set `use_nullspace: true` for robot in hydra cfg of dataset. If you downloaded one of the datasets prior to this date,
+edit this line in <PATH_TO_DATASET>/training/.hydra/merged_config.yaml and <PATH_TO_DATASET>/validation/.hydra/merged_config.yaml.
+- Renaming `model.decoder` to `model.action_decoder`.
+
+### 10 Jan 2022
+- Breaking change to evaluation, using different intitial states for environment.
+
+## Citation
+
+If you find the dataset or code useful, please cite:
+
+```bibtex
+@article{mees2022calvin,
+author = {Oier Mees and Lukas Hermann and Erick Rosete-Beas and Wolfram Burgard},
+title = {CALVIN: A Benchmark for Language-Conditioned Policy Learning for Long-Horizon Robot Manipulation Tasks},
+journal={IEEE Robotics and Automation Letters (RA-L)},
+volume={7},
+number={3},
+pages={7327-7334},
+year={2022}
+}
+```
+
+## License
+
+MIT License

aloha_robot_project/calvin/RL_with_CALVIN.ipynb

@@ -0,0 +1,325 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+  "colab": {
+   "name": "RL_with_CALVIN.ipynb",
+   "provenance": [],
+   "collapsed_sections": [],
+   "authorship_tag": "ABX9TyNwSDlhfoXLR2ORP+8/+3zP",
+   "include_colab_link": true
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "view-in-github",
+    "colab_type": "text"
+   },
+   "source": [
+    "<a href=\"https://colab.research.google.com/github/mees/calvin/blob/main/RL_with_CALVIN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "MDAC0SvBi1-9"
+   },
+   "source": [
+    "<h1>Reinforcement Learning with CALVIN</h1>\n",
+    "\n",
+    "The **CALVIN** simulated benchmark is perfectly suited for training agents with reinforcement learning, in this notebook we will demonstrate how to integrate your agents to these environments."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Gk7gEiqXizze"
+   },
+   "source": [
+    "## Installation\n",
+    "The first step is to install the CALVIN github repository such that we have access to the packages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "K1cSDZwLlAhR"
+   },
+   "source": [
+    "# Download repo\n",
+    "%mkdir /content/calvin\n",
+    "%cd /content/calvin\n",
+    "!git clone https://github.com/mees/calvin_env.git\n",
+    "%cd /content/calvin/calvin_env\n",
+    "!git clone https://github.com/lukashermann/tacto.git\n",
+    "# Install packages \n",
+    "%cd /content/calvin/calvin_env/tacto/\n",
+    "!pip3 install -e .\n",
+    "%cd /content/calvin/calvin_env\n",
+    "!pip3 install -e .\n",
+    "!pip3 install -U numpy"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "hPiKJ7mY58o3"
+   },
+   "source": [
+    "# Run this to check if the installation was succesful\n",
+    "from calvin_env.envs.play_table_env import PlayTableSimEnv"
+   ],
+   "execution_count": 2,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "JSmDkjlx0FOW"
+   },
+   "source": [
+    "## Loading the environment\n",
+    "After the installation has finished successfully, we can start using the environment for reinforcement Learning.\n",
+    "To be able to use the environment we need to have the appropriate configuration that define the desired features, for this example, we will load the static and gripper camera."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "vp1-mgIgekvY"
+   },
+   "source": [
+    "%cd /content/calvin\n",
+    "from hydra import initialize, compose\n",
+    "\n",
+    "with initialize(config_path=\"./calvin_env/conf/\"):\n",
+    "  cfg = compose(config_name=\"config_data_collection.yaml\", overrides=[\"cameras=static_and_gripper\"])\n",
+    "  cfg.env[\"use_egl\"] = False\n",
+    "  cfg.env[\"show_gui\"] = False\n",
+    "  cfg.env[\"use_vr\"] = False\n",
+    "  cfg.env[\"use_scene_info\"] = True\n",
+    "  print(cfg.env)"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "onN9sssj1JV6"
+   },
+   "source": [
+    "The environment has similar structure to traditional OpenAI Gym environments.\n",
+    "\n",
+    "*   We can restart the simulation with the *reset* function.\n",
+    "*   We can perform an action in the environment with the *step* function.\n",
+    "*   We can visualize images taken from the cameras in the environment by using the *render* function.\n",
+    "\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "MwIuxueazYOh"
+   },
+   "source": [
+    "import time\n",
+    "import hydra\n",
+    "import numpy as np\n",
+    "from google.colab.patches import cv2_imshow\n",
+    "\n",
+    "env = hydra.utils.instantiate(cfg.env)\n",
+    "observation = env.reset()\n",
+    "#The observation is given as a dictionary with different values\n",
+    "print(observation.keys())\n",
+    "for i in range(5):\n",
+    "  # The action consists in a pose displacement (position and orientation)\n",
+    "  action_displacement = np.random.uniform(low=-1, high=1, size=6)\n",
+    "  # And a binary gripper action, -1 for closing and 1 for oppening\n",
+    "  action_gripper = np.random.choice([-1, 1], size=1)\n",
+    "  action = np.concatenate((action_displacement, action_gripper), axis=-1)\n",
+    "  observation, reward, done, info = env.step(action)\n",
+    "  rgb = env.render(mode=\"rgb_array\")[:,:,::-1]\n",
+    "  cv2_imshow(rgb)"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "_0yYw5pHC0eE"
+   },
+   "source": [
+    "##  Custom environment for Reinforcement Learning\n",
+    "There are some aspects that needs to be defined to be able to use it for reinforcement learning, including:\n",
+    "\n",
+    "1.   Observation space\n",
+    "2.   Action space\n",
+    "3.   Reward function\n",
+    "\n",
+    "We are going to create a Custom environment that extends the **PlaytableSimEnv** to add these requirements. <br/>\n",
+    "The specific task that will be solved is called \"move_slider_left\", here you can find a [list of possible tasks](https://github.com/mees/calvin_env/blob/main/conf/tasks/new_playtable_tasks.yaml) that can be evaluated using CALVIN.\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "Bwj-5UQz2vyF"
+   },
+   "source": [
+    "from gym import spaces\n",
+    "from calvin_env.envs.play_table_env import PlayTableSimEnv\n",
+    "\n",
+    "class SlideEnv(PlayTableSimEnv):\n",
+    "    def __init__(self,\n",
+    "                 tasks: dict = {},\n",
+    "                 **kwargs):\n",
+    "        super(SlideEnv, self).__init__(**kwargs)\n",
+    "        # For this example we will modify the observation to\n",
+    "        # only retrieve the end effector pose\n",
+    "        self.action_space = spaces.Box(low=-1, high=1, shape=(7,))\n",
+    "        self.observation_space = spaces.Box(low=-1, high=1, shape=(7,))\n",
+    "        # We can use the task utility to know if the task was executed correctly\n",
+    "        self.tasks = hydra.utils.instantiate(tasks)\n",
+    "\n",
+    "    def reset(self):\n",
+    "        obs = super().reset()\n",
+    "        self.start_info = self.get_info()\n",
+    "        return obs\n",
+    "\n",
+    "    def get_obs(self):\n",
+    "        \"\"\"Overwrite robot obs to only retrieve end effector position\"\"\"\n",
+    "        robot_obs, robot_info = self.robot.get_observation()\n",
+    "        return robot_obs[:7]\n",
+    "\n",
+    "    def _success(self):\n",
+    "        \"\"\" Returns a boolean indicating if the task was performed correctly \"\"\"\n",
+    "        current_info = self.get_info()\n",
+    "        task_filter = [\"move_slider_left\"]\n",
+    "        task_info = self.tasks.get_task_info_for_set(self.start_info, current_info, task_filter)\n",
+    "        return 'move_slider_left' in task_info\n",
+    "\n",
+    "    def _reward(self):\n",
+    "        \"\"\" Returns the reward function that will be used \n",
+    "        for the RL algorithm \"\"\"\n",
+    "        reward = int(self._success()) * 10\n",
+    "        r_info = {'reward': reward}\n",
+    "        return reward, r_info\n",
+    "\n",
+    "    def _termination(self):\n",
+    "        \"\"\" Indicates if the robot has reached a terminal state \"\"\"\n",
+    "        success = self._success()\n",
+    "        done = success\n",
+    "        d_info = {'success': success}        \n",
+    "        return done, d_info\n",
+    "\n",
+    "    def step(self, action):\n",
+    "            \"\"\" Performing a relative action in the environment\n",
+    "                input:\n",
+    "                    action: 7 tuple containing\n",
+    "                            Position x, y, z. \n",
+    "                            Angle in rad x, y, z. \n",
+    "                            Gripper action\n",
+    "                            each value in range (-1, 1)\n",
+    "\n",
+    "                            OR\n",
+    "                            8 tuple containing\n",
+    "                            Relative Joint angles j1 - j7 (in rad)\n",
+    "                            Gripper action\n",
+    "                output:\n",
+    "                    observation, reward, done info\n",
+    "            \"\"\"\n",
+    "            # Transform gripper action to discrete space\n",
+    "            env_action = action.copy()\n",
+    "            env_action[-1] = (int(action[-1] >= 0) * 2) - 1\n",
+    "\n",
+    "            # for using actions in joint space\n",
+    "            if len(env_action) == 8:\n",
+    "                env_action = {\"action\": env_action, \"type\": \"joint_rel\"}\n",
+    "\n",
+    "            self.robot.apply_action(env_action)\n",
+    "            for i in range(self.action_repeat):\n",
+    "                self.p.stepSimulation(physicsClientId=self.cid)\n",
+    "            obs = self.get_obs()\n",
+    "            info = self.get_info()\n",
+    "            reward, r_info = self._reward()\n",
+    "            done, d_info = self._termination()\n",
+    "            info.update(r_info)\n",
+    "            info.update(d_info)\n",
+    "            return obs, reward, done, info"
+   ],
+   "execution_count": 5,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "IJo-HWRqcJHc"
+   },
+   "source": [
+    "# Training an RL agent\n",
+    "After generating the wrapper training a reinforcement learning agent is straightforward, for this example we will use stable baselines 3 agents"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "XOkD9S-iMVcj"
+   },
+   "source": [
+    "!pip3 install stable_baselines3"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "WxWRcJJFcpEF"
+   },
+   "source": [
+    "To train the agent we create an instance of our new environment and send it to the stable baselines agent to learn a policy.\n",
+    "\n",
+    "\n",
+    "> Note: the example uses Soft Actor Critic (SAC) which is one of the state of the art algorithm for off-policy RL.\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "3BUeaAnqMNLq"
+   },
+   "source": [
+    "import gym\n",
+    "import numpy as np\n",
+    "from stable_baselines3 import SAC\n",
+    "\n",
+    "new_env_cfg = {**cfg.env}\n",
+    "new_env_cfg[\"tasks\"] = cfg.tasks\n",
+    "new_env_cfg.pop('_target_', None)\n",
+    "new_env_cfg.pop('_recursive_', None)\n",
+    "env = SlideEnv(**new_env_cfg)\n",
+    "model = SAC(\"MlpPolicy\", env, verbose=1)\n",
+    "model.learn(total_timesteps=10000, log_interval=4)\n"
+   ],
+   "execution_count": null,
+   "outputs": []
+  }
+ ]
+}

aloha_robot_project/calvin/calvin_env/.flake8

@@ -0,0 +1,15 @@
+[flake8]
+exclude = .git
+# Default is 79 in PEP 8
+max-line-length = 120
+select = E,F,W,C
+ignore=W503, # line break before binary operator, need for black
+       E203, # whitespace before ':'. Opposite convention enforced by black
+       E731, # do not assign a lambda expression, use a def
+       E722,
+       F401,
+       F841,
+       E402,  # module level import not at top of file
+       E741,  # ambiguous variable name
+       E501, # line too long. Handled by black
+       C406,  # Unnecessary list literal - rewrite as a dict literal

aloha_robot_project/calvin/calvin_env/.gitignore

@@ -0,0 +1,161 @@
+# 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/
+pip-wheel-metadata/
+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/
+
+# 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
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.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
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__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/
+
+
+# NumPy
+*.npy
+*.npz
+
+# Pickles
+*.pickle
+
+# gzip
+*.gzip
+
+# Json
+#*.json
+
+# VSCode
+.vscode/*
+
+# Own playground area
+**/playground/*
+
+# tensorboard
+events.out*
+*.svg
+
+# pdfs
+*.pdf
+
+*.egg-info/*
+/.idea/
+
+egl_check/EGL_options.o

aloha_robot_project/calvin/calvin_env/.gitmodules

@@ -0,0 +1,3 @@
+[submodule "tacto"]
+	path = tacto
+	url = https://github.com/lukashermann/tacto.git

aloha_robot_project/calvin/calvin_env/.pre-commit-config.yaml

@@ -0,0 +1,32 @@
+default_language_version:
+    python: python3.8
+repos:
+  - repo: https://github.com/psf/black
+    rev: 22.1.0
+    hooks:
+      - id: black
+        language_version: python3.8
+
+  - repo: https://gitlab.com/pycqa/flake8
+    rev: 3.8.4
+    hooks:
+      - id: flake8
+        additional_dependencies: [-e, "git+git://github.com/pycqa/pyflakes.git@c72d6cf#egg=pyflakes"]
+
+  - repo: https://github.com/pre-commit/mirrors-isort
+    rev: v5.6.4
+    hooks:
+      - id: isort
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v0.790
+    hooks:
+      - id: mypy
+        args: [--ignore-missing-imports, --warn-no-return, --warn-redundant-casts, --disallow-incomplete-defs]
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v3.4.0
+    hooks:
+      - id: check-yaml
+      - id: trailing-whitespace
+      - id: end-of-file-fixer

aloha_robot_project/calvin/calvin_env/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Oier Mees
+
+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.

aloha_robot_project/calvin/calvin_env/README.md

@@ -0,0 +1,12 @@
+# calvin_env
+
+
+
+## Installation
+```bash
+git clone --recursive https://github.com/mees/calvin_env.git
+cd calvin_env/tacto
+pip install -e .
+cd ..
+pip install -e .
+```

aloha_robot_project/calvin/calvin_env/calvin_env/__init__.py

@@ -0,0 +1,10 @@
+"""'VR Data Collection and Rendering
+:copyright: 2019 by Oier Mees, Lukas Hermann, Wolfram Burgard
+:license: GPLv3, see LICENSE for more details.
+"""
+
+__version__ = "0.0.1"
+__project__ = "calvin_env"
+__author__ = "Oier Mees, Lukas Hermann"
+__license__ = "GPLv3"
+__email__ = "meeso@informatik.uni-freiburg.de, hermannl@informatik.uni-freiburg.de,"

aloha_robot_project/calvin/calvin_env/calvin_env/camera/camera.py

@@ -0,0 +1,91 @@
+import math
+
+import numpy as np
+
+
+class Camera:
+    def __init__(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def render(self):
+        raise NotImplementedError
+
+    def distance_map_to_point_cloud(self, distances, fov, width, height):
+        """Converts from a depth map to a point cloud.
+        Args:
+          distances: An numpy array which has the shape of (height, width) that
+            denotes a distance map. The unit is meter.
+          fov: The field of view of the camera in the vertical direction. The unit
+            is radian.
+          width: The width of the image resolution of the camera.
+          height: The height of the image resolution of the camera.
+        Returns:
+          point_cloud: The converted point cloud from the distance map. It is a numpy
+            array of shape (height, width, 3).
+        """
+        f = height / (2 * math.tan(fov / 2.0))
+        px = np.tile(np.arange(width), [height, 1])
+        x = (2 * (px + 0.5) - width) / f * distances / 2
+        py = np.tile(np.arange(height), [width, 1]).T
+        y = (2 * (py + 0.5) - height) / f * distances / 2
+        point_cloud = np.stack((x, y, distances), axis=-1)
+        return point_cloud
+
+    def z_buffer_to_real_distance(self, z_buffer, far, near):
+        """Function to transform depth buffer values to distances in camera space"""
+        return 1.0 * far * near / (far - (far - near) * z_buffer)
+
+    def process_rgbd(self, obs, nearval, farval):
+        (width, height, rgbPixels, depthPixels, segmentationMaskBuffer) = obs
+        rgb = np.reshape(rgbPixels, (height, width, 4))
+        rgb_img = rgb[:, :, :3]
+        depth_buffer = np.reshape(depthPixels, [height, width])
+        depth = self.z_buffer_to_real_distance(z_buffer=depth_buffer, far=farval, near=nearval)
+        return rgb_img, depth
+
+    # Reference: world2pixel
+    # https://github.com/bulletphysics/bullet3/issues/1952
+    def project(self, point):
+        """
+        Projects a world point in homogeneous coordinates to pixel coordinates
+        Args
+            point: np.array of len 4; indicates the desired point to project
+        Output
+            (x, y): tuple (u, v); pixel coordinates of the projected point
+        """
+
+        # reshape to get homogeneus transform
+        persp_m = np.array(self.projectionMatrix).reshape((4, 4)).T
+        view_m = np.array(self.viewMatrix).reshape((4, 4)).T
+
+        # Perspective proj matrix
+        world_pix_tran = persp_m @ view_m @ point
+        world_pix_tran = world_pix_tran / world_pix_tran[-1]  # divide by w
+        world_pix_tran[:3] = (world_pix_tran[:3] + 1) / 2
+        x, y = world_pix_tran[0] * self.width, (1 - world_pix_tran[1]) * self.height
+        x, y = np.floor(x).astype(int), np.floor(y).astype(int)
+        return (x, y)
+
+    def deproject(self, point, depth_img, homogeneous=False):
+        """
+        Deprojects a pixel point to 3D coordinates
+        Args
+            point: tuple (u, v); pixel coordinates of point to deproject
+            depth_img: np.array; depth image used as reference to generate 3D coordinates
+            homogeneous: bool; if true it returns the 3D point in homogeneous coordinates,
+                         else returns the world coordinates (x, y, z) position
+        Output
+            (x, y): np.array; world coordinates of the deprojected point
+        """
+        T_world_cam = np.linalg.inv(np.array(self.viewMatrix).reshape((4, 4)).T)
+
+        u, v = point
+        z = depth_img[v, u]
+        foc = self.height / (2 * np.tan(np.deg2rad(self.fov) / 2))
+        x = (u - self.width // 2) * z / foc
+        y = -(v - self.height // 2) * z / foc
+        z = -z
+        world_pos = T_world_cam @ np.array([x, y, z, 1])
+        if not homogeneous:
+            world_pos = world_pos[:3]
+        return world_pos

aloha_robot_project/calvin/calvin_env/calvin_env/camera/gripper_camera.py

@@ -0,0 +1,46 @@
+import numpy as np
+import pybullet as p
+
+from calvin_env.camera.camera import Camera
+
+
+class GripperCamera(Camera):
+    def __init__(self, fov, aspect, nearval, farval, width, height, robot_id, cid, name, objects=None):
+        self.cid = cid
+        self.robot_uid = robot_id
+        links = {
+            p.getJointInfo(self.robot_uid, i, physicsClientId=self.cid)[12].decode("utf-8"): i
+            for i in range(p.getNumJoints(self.robot_uid, physicsClientId=self.cid))
+        }
+        self.gripper_cam_link = links["gripper_cam"]
+        self.fov = fov
+        self.aspect = aspect
+        self.nearval = nearval
+        self.farval = farval
+        self.width = width
+        self.height = height
+
+        self.name = name
+
+    def render(self):
+        camera_ls = p.getLinkState(
+            bodyUniqueId=self.robot_uid, linkIndex=self.gripper_cam_link, physicsClientId=self.cid
+        )
+        camera_pos, camera_orn = camera_ls[:2]
+        cam_rot = p.getMatrixFromQuaternion(camera_orn)
+        cam_rot = np.array(cam_rot).reshape(3, 3)
+        cam_rot_y, cam_rot_z = cam_rot[:, 1], cam_rot[:, 2]
+        # camera: eye position, target position, up vector
+        self.view_matrix = p.computeViewMatrix(camera_pos, camera_pos + cam_rot_y, -cam_rot_z)
+        self.projection_matrix = p.computeProjectionMatrixFOV(
+            fov=self.fov, aspect=self.aspect, nearVal=self.nearval, farVal=self.farval
+        )
+        image = p.getCameraImage(
+            width=self.width,
+            height=self.height,
+            viewMatrix=self.view_matrix,
+            projectionMatrix=self.projection_matrix,
+            physicsClientId=self.cid,
+        )
+        rgb_img, depth_img = self.process_rgbd(image, self.nearval, self.farval)
+        return rgb_img, depth_img

aloha_robot_project/calvin/calvin_env/calvin_env/camera/static_camera.py

@@ -0,0 +1,72 @@
+import numpy as np
+import pybullet as p
+
+from calvin_env.camera.camera import Camera
+
+
+class StaticCamera(Camera):
+    def __init__(
+        self,
+        fov,
+        aspect,
+        nearval,
+        farval,
+        width,
+        height,
+        look_at,
+        look_from,
+        up_vector,
+        cid,
+        name,
+        robot_id=None,
+        objects=None,
+    ):
+        """
+        Initialize the camera
+        Args:
+            argument_group: initialize the camera and add needed arguments to argparse
+
+        Returns:
+            None
+        """
+        self.nearval = nearval
+        self.farval = farval
+        self.fov = fov
+        self.aspect = aspect
+        self.look_from = look_from
+        self.look_at = look_at
+        self.up_vector = up_vector
+        self.width = width
+        self.height = height
+        self.viewMatrix = p.computeViewMatrix(
+            cameraEyePosition=look_from, cameraTargetPosition=look_at, cameraUpVector=self.up_vector
+        )
+        self.projectionMatrix = p.computeProjectionMatrixFOV(
+            fov=fov, aspect=aspect, nearVal=self.nearval, farVal=self.farval
+        )
+        self.cid = cid
+        self.name = name
+
+    def set_position_from_gui(self):
+        info = p.getDebugVisualizerCamera(physicsClientId=self.cid)
+        look_at = np.array(info[-1])
+        dist = info[-2]
+        forward = np.array(info[5])
+        look_from = look_at - dist * forward
+        self.viewMatrix = p.computeViewMatrix(
+            cameraEyePosition=look_from, cameraTargetPosition=look_at, cameraUpVector=self.up_vector
+        )
+        look_from = [float(x) for x in look_from]
+        look_at = [float(x) for x in look_at]
+        return look_from, look_at
+
+    def render(self):
+        image = p.getCameraImage(
+            width=self.width,
+            height=self.height,
+            viewMatrix=self.viewMatrix,
+            projectionMatrix=self.projectionMatrix,
+            physicsClientId=self.cid,
+        )
+        rgb_img, depth_img = self.process_rgbd(image, self.nearval, self.farval)
+        return rgb_img, depth_img

aloha_robot_project/calvin/calvin_env/calvin_env/camera/tactile_sensor.py

@@ -0,0 +1,41 @@
+import os
+
+import numpy as np
+
+from calvin_env.camera.camera import Camera
+import tacto
+
+REPO_BASE = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+
+
+class TactileSensor(Camera):
+    def __init__(
+        self, width, height, digit_link_ids, visualize_gui, cid, name, config_path, robot_id=None, objects=None
+    ):
+        """
+        Initialize the camera
+        Args:
+            argument_group: initialize the camera and add needed arguments to argparse
+
+        Returns:
+            None
+        """
+        self.cid = cid
+        self.name = name
+        self.robot_uid = robot_id
+        self.digits = tacto.Sensor(
+            width=width, height=height, visualize_gui=visualize_gui, config_path=os.path.join(REPO_BASE, config_path)
+        )
+        self.digits.add_camera(robot_id, digit_link_ids)  # env.robot.digit_links()
+        for obj in objects:
+            # self.digits.add_body(obj)
+            self.digits.add_object(obj.file.as_posix(), obj.uid, obj.global_scaling)
+        self.visualize_gui = visualize_gui
+
+    def render(self):
+        rgb, depth = self.digits.render()
+        if self.visualize_gui:
+            self.digits.updateGUI(rgb, depth)
+        rgb = np.concatenate(rgb, axis=2)
+        depth = np.stack(depth, axis=2)
+        return rgb, depth

aloha_robot_project/calvin/calvin_env/calvin_env/datarenderer.py

@@ -0,0 +1,299 @@
+#!/usr/bin/python3
+
+from collections import deque
+import logging
+from multiprocessing import Process
+import os
+from pathlib import Path
+
+import hydra
+import numpy as np
+import omegaconf
+import pybullet as p
+from scipy.spatial.transform.rotation import Rotation as R
+
+from calvin_env.utils.utils import count_frames, get_episode_lengths, set_egl_device, to_relative_action
+
+# A logger for this file
+log = logging.getLogger(__name__)
+
+
+@hydra.main(config_path="../conf", config_name="config_rendering")
+def main(cfg):
+    log.info("pyBullet Data Renderer")
+    log.info("Determining maximum frame")
+
+    recording_dir = (Path(hydra.utils.get_original_cwd()) / cfg.load_dir).absolute()
+    max_frames = count_frames(recording_dir)
+    log.info(f"Found continuous interval of {max_frames} frames without gaps")
+
+    num_prev_rendered_episodes = num_previously_rendered_episodes()
+    if num_prev_rendered_episodes == 0:
+        playback_cfg = build_rendering_config(recording_dir, rendering_config=cfg)
+    else:
+        playback_cfg = load_rendering_config(cfg)
+
+    log.info("Initialization done!")
+    log.info(f"Starting {cfg.processes} processes")
+
+    if playback_cfg.set_static_cam:
+        playback_cfg = set_static_cams_from_gui(playback_cfg, recording_dir, max_frames)
+
+    if cfg.processes != 1 and playback_cfg.show_gui:
+        log.warning("Multiprocess rendering requires headless mode, setting cfg.show_gui = False")
+        playback_cfg.show_gui = False
+    # in order to distribute the rendering to multiple processes, predetermine the lengths of the
+    # (rendered) episodes and to which (recording) file ids the episode start and end correspond
+    # a rendered episode does not contain the done frame, thus length(render_episode) = length(recording_episode) -1
+    episode_lengths, render_start_end_ids = get_episode_lengths(cfg.load_dir, max_frames)
+    # episode_lengths = episode_lengths[:1]
+    # render_start_end_ids = render_start_end_ids[:1]
+
+    if cfg.processes > len(episode_lengths):
+        log.warning(f"Trying to use more processes ({cfg.processes}) than episodes ({len(episode_lengths)}).")
+        log.warning(f"Reducing number of processes to {len(episode_lengths)}.")
+        cfg.processes = len(episode_lengths)
+    # distribute the episodes equally to processes
+    split_indices = np.array_split(np.array(render_start_end_ids), cfg.processes, axis=0)
+    # every process renders the interval [proc_start_ids, proc_end_ids)
+    proc_start_ids = [split_indices[proc_num][0][0] for proc_num in range(cfg.processes)]
+    proc_end_ids = [split_indices[proc_num][-1][1] for proc_num in range(cfg.processes)]
+    # predetermine starting episode indices for multiple processes
+    proc_ep_ids = np.cumsum(
+        [0] + list(map(np.sum, np.array_split(np.array(episode_lengths), cfg.processes, axis=0)))[:-1]
+    )
+    proc_ep_ids += num_prev_rendered_episodes
+    if cfg.processes > 1:
+        processes = [
+            Process(
+                target=worker_run,
+                args=(
+                    recording_dir,
+                    playback_cfg,
+                    proc_num,
+                    proc_start_ids[proc_num],
+                    proc_end_ids[proc_num],
+                    proc_ep_ids[proc_num],
+                ),
+                name=f"Worker {proc_num}",
+            )
+            for proc_num in range(cfg.processes)
+        ]
+        deque(map(lambda proc: proc.start(), processes))
+        deque(map(lambda proc: proc.join(), processes))
+    else:
+        worker_run(recording_dir, playback_cfg, 0, 0, max_frames, num_prev_rendered_episodes)
+    save_ep_lens(episode_lengths, num_prev_rendered_episodes)
+
+    log.info("All workers done")
+
+
+def build_rendering_config(recording_dir, rendering_config):
+    merged_conf = omegaconf.OmegaConf.load(Path(recording_dir) / ".hydra" / "config.yaml")
+    merged_conf = omegaconf.OmegaConf.merge(merged_conf, rendering_config)
+    hydra.core.utils._save_config(merged_conf, "merged_config.yaml", Path(os.getcwd(), ".hydra"))
+    return merged_conf
+
+
+def load_rendering_config(rendering_config):
+    conf = omegaconf.OmegaConf.load(Path(os.getcwd()) / ".hydra" / "merged_config.yaml")
+    override_conf = omegaconf.OmegaConf.select(rendering_config, "scene")
+    omegaconf.OmegaConf.update(conf, "scene", override_conf, merge=False)
+    conf.set_static_cam = False
+    return conf
+
+
+def num_previously_rendered_episodes():
+    return len(list(Path(os.getcwd()).glob("*.npz")))
+
+
+def save_ep_lens(episode_lengths, num_prev_episodes):
+    if num_prev_episodes > 0:
+        previous_ep_lens = np.load("ep_lens.npy")
+        episode_lengths = np.concatenate((previous_ep_lens, episode_lengths))
+    np.save("ep_lens.npy", episode_lengths)
+    end_ids = np.cumsum(episode_lengths) - 1
+    start_ids = [0] + list(end_ids + 1)[:-1]
+    ep_start_end_ids = list(zip(start_ids, end_ids))
+    np.save("ep_start_end_ids.npy", ep_start_end_ids)
+
+
+def save_step(counter, rgbs, depths, actions, robot_obs, scene_obs, cam_names, **additional_infos):
+    rgb_entries = {f"rgb_{cam_name}": rgbs[f"rgb_{cam_name}"] for i, cam_name in enumerate(cam_names)}
+    depths_entries = {f"depth_{cam_name}": depths[f"depth_{cam_name}"] for i, cam_name in enumerate(cam_names)}
+    if actions[-1] == 0:
+        actions[-1] = -1
+    np.savez_compressed(
+        f"episode_{counter:07d}.npz",
+        actions=actions,
+        rel_actions=to_relative_action(actions, robot_obs),
+        robot_obs=robot_obs,
+        scene_obs=scene_obs,
+        **rgb_entries,
+        **depths_entries,
+        **additional_infos,
+    )
+
+
+def state_to_action(info):
+    """
+    save action as [tcp_pos, tcp_orn_quaternion, gripper_action]
+    """
+    tcp_pos = info["robot_info"]["tcp_pos"]
+    tcp_orn = info["robot_info"]["tcp_orn"]
+    gripper_action = info["robot_info"]["gripper_action"]
+    action = np.concatenate([tcp_pos, tcp_orn, [gripper_action]])
+    return action
+
+
+def set_static_cams_from_gui(cfg, load_dir, max_frames):
+    import cv2
+
+    assert cfg.env.show_gui
+    env = hydra.utils.instantiate(cfg.env)
+    env.reset()
+    frame = 0
+    log.info("--------------------------------------------------")
+    log.info("Use Debug GUI to change the position of the camera")
+    log.info("Use Render_view_window for keyboard input")
+    log.info("Press A or D to move through frames")
+    log.info("Press Q or E to skip through frames")
+    log.info("Press S to set camera position")
+    log.info("Press ENTER to save the set camera position")
+    log.info("Press ESC to skip setting position for current camera")
+    for cam_index, (cam_name, cam) in enumerate(cfg.cameras.items()):
+        if "static" in cam._target_:
+            # initialize variables
+            look_from = cam.look_from
+            look_at = cam.look_at
+            up_vector = cam.up_vector
+            fov = cam.fov
+            while True:
+                file_path = load_dir / f"{frame:012d}.pickle"
+                _, _, _ = env.reset_from_storage(file_path)
+                env.p.stepSimulation()
+                frame_rgbs, frame_depths = env.get_camera_obs()
+                rgb_static = frame_rgbs[cam_index]
+
+                cv2.imshow("Render_view_window", cv2.resize(rgb_static, (500, 500))[:, :, ::-1])
+                k = cv2.waitKey(10) % 256
+                if k == ord("a"):
+                    frame -= 1
+                    frame = np.clip(frame, 0, max_frames - 1)
+                if k == ord("d"):
+                    frame += 1
+                    frame = np.clip(frame, 0, max_frames - 1)
+                if k == ord("q"):
+                    frame -= 100
+                    frame = np.clip(frame, 0, max_frames - 1)
+                if k == ord("e"):
+                    frame += 100
+                    frame = np.clip(frame, 0, max_frames - 1)
+                if k == ord("z"):
+                    c = env.cameras[cam_index]
+                    c.fov -= 1
+                    c.projectionMatrix = p.computeProjectionMatrixFOV(
+                        fov=c.fov, aspect=c.aspect, nearVal=c.nearval, farVal=c.farval
+                    )
+                    print(c.fov)
+                    fov = c.fov
+                if k == ord("x"):
+                    c = env.cameras[cam_index]
+                    c.fov += 1
+                    c.projectionMatrix = p.computeProjectionMatrixFOV(
+                        fov=c.fov, aspect=c.aspect, nearVal=c.nearval, farVal=c.farval
+                    )
+                    print(c.fov)
+                    fov = c.fov
+                if k == ord("r"):
+                    c = env.cameras[cam_index]
+                    direction_vector = np.array(c.look_at) - np.array(c.look_from)
+                    c.up_vector = (
+                        R.from_rotvec(0.1 * direction_vector / np.linalg.norm(direction_vector)).as_matrix()
+                        @ c.up_vector
+                    )
+                    up_vector = c.up_vector
+                if k == ord("f"):
+                    c = env.cameras[cam_index]
+                    direction_vector = np.array(c.look_at) - np.array(c.look_from)
+                    c.up_vector = (
+                        R.from_rotvec(-0.1 * direction_vector / np.linalg.norm(direction_vector)).as_matrix()
+                        @ c.up_vector
+                    )
+                    up_vector = c.up_vector
+                if k == 13:  # Enter
+                    cam.look_from = look_from
+                    cam.look_at = look_at
+                    log.info(f"Set look_from of camera {cam_index} to {look_from}")
+                    log.info(f"Set look_at of camera {cam_index} to {look_at}")
+                    cam.up_vector = np.array(up_vector).tolist()
+                    log.info(f"Set up_vector of camera {cam_index} to {up_vector}")
+                    cam.fov = fov
+                    log.info(f"Set fov of camera {cam_index} to {fov}")
+                    break
+                if k == 27:  # ESC
+                    log.info(f"Do no change position of camera {cam_index}")
+                    break
+                # if k == ord("s"):
+                look_from, look_at = env.cameras[cam_index].set_position_from_gui()
+    hydra.core.utils._save_config(cfg, "merged_config.yaml", Path(os.getcwd(), ".hydra"))
+    env.close()
+    return cfg
+
+
+def worker_run(load_dir, rendering_cfg, proc_num, start_frame, stop_frame, episode_index):
+    log.info(f"[{proc_num}] Starting worker {proc_num}")
+    set_egl_device(0)
+    env = hydra.utils.instantiate(rendering_cfg.env)
+
+    log.info(f"[{proc_num}] Entering Loop")
+    frame_counter = 0
+    rgbs, depths, actions, robot_obs, scene_obs, = (
+        [],
+        [],
+        [],
+        [],
+        [],
+    )
+    for frame in range(start_frame, stop_frame):
+        file_path = load_dir / f"{frame:012d}.pickle"
+        state_ob, done, info = env.reset_from_storage(file_path)
+        action = state_to_action(info)
+        robot_obs.append(state_ob["robot_obs"])
+        scene_obs.append(state_ob["scene_obs"])
+
+        # action is robot state of next frame
+        if frame_counter > 0:
+            actions.append(action)
+        frame_rgbs, frame_depths = env.get_camera_obs()
+        rgbs.append(frame_rgbs)
+        depths.append(frame_depths)
+        # for terminal states save current robot state as action
+        frame_counter += 1
+        log.debug(f"episode counter {episode_index} frame counter {frame_counter} done {done}")
+
+        if frame_counter > 1:
+            save_step(
+                episode_index,
+                rgbs.pop(0),
+                depths.pop(0),
+                actions.pop(0),
+                robot_obs.pop(0),
+                scene_obs.pop(0),
+                cam_names=[cam.name for cam in env.cameras],
+            )
+            episode_index += 1
+            if done:
+                frame_counter = 0
+                rgbs, depths, actions, robot_obs, scene_obs = [], [], [], [], []
+
+        log.debug(f"[{proc_num}] Rendered frame {frame}")
+
+    assert done
+
+    env.close()
+    log.info(f"[{proc_num}] Finishing worker {proc_num}")
+
+
+if __name__ == "__main__":
+    main()

aloha_robot_project/calvin/calvin_env/calvin_env/envs/play_lmp_wrapper.py

@@ -0,0 +1,106 @@
+import logging
+import os
+from typing import Any, Dict, Tuple, Union
+
+from calvin_agent.datasets.utils.episode_utils import process_depth, process_rgb, process_state
+import gym
+import numpy as np
+import torch
+
+from calvin_env.envs.play_table_env import get_env
+from calvin_env.utils.utils import EglDeviceNotFoundError, get_egl_device_id
+
+logger = logging.getLogger(__name__)
+
+
+class PlayLMPWrapper(gym.Wrapper):
+    def __init__(self, dataset_loader, device, show_gui=False, **kwargs):
+        self.set_egl_device(device)
+        env = get_env(
+            dataset_loader.abs_datasets_dir, show_gui=show_gui, obs_space=dataset_loader.observation_space, **kwargs
+        )
+        super(PlayLMPWrapper, self).__init__(env)
+        self.observation_space_keys = dataset_loader.observation_space
+        self.transforms = dataset_loader.transforms
+        self.proprio_state = dataset_loader.proprio_state
+        self.device = device
+        self.relative_actions = "rel_actions" in self.observation_space_keys["actions"]
+        logger.info(f"Initialized PlayTableEnv for device {self.device}")
+
+    @staticmethod
+    def set_egl_device(device):
+        if "EGL_VISIBLE_DEVICES" in os.environ:
+            logger.warning("Environment variable EGL_VISIBLE_DEVICES is already set. Is this intended?")
+        cuda_id = device.index if device.type == "cuda" else 0
+        try:
+            egl_id = get_egl_device_id(cuda_id)
+        except EglDeviceNotFoundError:
+            logger.warning(
+                "Couldn't find correct EGL device. Setting EGL_VISIBLE_DEVICE=0. "
+                "When using DDP with many GPUs this can lead to OOM errors. "
+                "Did you install PyBullet correctly? Please refer to calvin env README"
+            )
+            egl_id = 0
+        os.environ["EGL_VISIBLE_DEVICES"] = str(egl_id)
+        logger.info(f"EGL_DEVICE_ID {egl_id} <==> CUDA_DEVICE_ID {cuda_id}")
+
+    def transform_observation(self, obs: Dict[str, Any]) -> Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]:
+        state_obs = process_state(obs, self.observation_space_keys, self.transforms, self.proprio_state)
+        rgb_obs = process_rgb(obs["rgb_obs"], self.observation_space_keys, self.transforms)
+        depth_obs = process_depth(obs["depth_obs"], self.observation_space_keys, self.transforms)
+
+        state_obs["robot_obs"] = state_obs["robot_obs"].to(self.device).unsqueeze(0)
+        rgb_obs.update({"rgb_obs": {k: v.to(self.device).unsqueeze(0) for k, v in rgb_obs["rgb_obs"].items()}})
+        depth_obs.update({"depth_obs": {k: v.to(self.device).unsqueeze(0) for k, v in depth_obs["depth_obs"].items()}})
+
+        obs_dict = {**rgb_obs, **state_obs, **depth_obs}
+        obs_dict["robot_obs_raw"] = torch.from_numpy(obs["robot_obs"]).to(self.device)
+        return obs_dict
+
+    def step(
+        self, action_tensor: torch.Tensor
+    ) -> Tuple[Dict[str, Union[torch.Tensor, Tuple[torch.Tensor, ...]]], int, bool, Dict]:
+        if self.relative_actions:
+            action = action_tensor.squeeze().cpu().detach().numpy()
+            assert len(action) == 7
+        else:
+            if action_tensor.shape[-1] == 7:
+                slice_ids = [3, 6]
+            elif action_tensor.shape[-1] == 8:
+                slice_ids = [3, 7]
+            else:
+                logger.error("actions are required to have length 8 (for euler angles) or 9 (for quaternions)")
+                raise NotImplementedError
+            action = np.split(action_tensor.squeeze().cpu().detach().numpy(), slice_ids)
+        action[-1] = 1 if action[-1] > 0 else -1
+        o, r, d, i = self.env.step(action)
+
+        obs = self.transform_observation(o)
+        return obs, r, d, i
+
+    def reset(
+        self,
+        reset_info: Dict[str, Any] = None,
+        batch_idx: int = 0,
+        seq_idx: int = 0,
+        scene_obs: Any = None,
+        robot_obs: Any = None,
+    ) -> Dict[str, Union[torch.Tensor, Tuple[torch.Tensor, ...]]]:
+        if reset_info is not None:
+            obs = self.env.reset(
+                robot_obs=reset_info["robot_obs"][batch_idx, seq_idx],
+                scene_obs=reset_info["scene_obs"][batch_idx, seq_idx],
+            )
+        elif scene_obs is not None or robot_obs is not None:
+            obs = self.env.reset(scene_obs=scene_obs, robot_obs=robot_obs)
+        else:
+            obs = self.env.reset()
+
+        return self.transform_observation(obs)
+
+    def get_info(self):
+        return self.env.get_info()
+
+    def get_obs(self):
+        obs = self.env.get_obs()
+        return self.transform_observation(obs)

aloha_robot_project/calvin/calvin_env/calvin_env/envs/play_table_env.py

@@ -0,0 +1,304 @@
+import logging
+from math import pi
+import os
+from pathlib import Path
+import pickle
+import pkgutil
+import re
+import sys
+import time
+
+import cv2
+import gym
+import gym.utils
+import gym.utils.seeding
+import hydra
+import numpy as np
+import pybullet as p
+import pybullet_utils.bullet_client as bc
+
+import calvin_env
+from calvin_env.utils.utils import FpsController, get_git_commit_hash
+
+# A logger for this file
+log = logging.getLogger(__name__)
+from rich.traceback import install
+
+install(show_locals=True)
+
+
+class PlayTableSimEnv(gym.Env):
+    def __init__(
+        self,
+        robot_cfg,
+        seed,
+        use_vr,
+        bullet_time_step,
+        cameras,
+        show_gui,
+        scene_cfg,
+        use_scene_info,
+        use_egl,
+        control_freq=30,
+    ):
+        self.p = p
+        # for calculation of FPS
+        self.t = time.time()
+        self.prev_time = time.time()
+        self.fps_controller = FpsController(bullet_time_step)
+        self.use_vr = use_vr
+        self.show_gui = show_gui
+        self.use_scene_info = use_scene_info
+        self.cid = -1
+        self.ownsPhysicsClient = False
+        self.use_egl = use_egl
+        self.control_freq = control_freq
+        self.action_repeat = int(bullet_time_step // control_freq)
+        render_width = max([cameras[cam].width for cam in cameras]) if cameras else None
+        render_height = max([cameras[cam].height for cam in cameras]) if cameras else None
+        self.initialize_bullet(bullet_time_step, render_width, render_height)
+        self.np_random = None
+        self.seed(seed)
+        self.robot = hydra.utils.instantiate(robot_cfg, cid=self.cid)
+        self.scene = hydra.utils.instantiate(scene_cfg, p=self.p, cid=self.cid, np_random=self.np_random)
+
+        # Load Env
+        self.load()
+
+        # init cameras after scene is loaded to have robot id available
+        self.cameras = [
+            hydra.utils.instantiate(
+                cameras[name], cid=self.cid, robot_id=self.robot.robot_uid, objects=self.scene.get_objects()
+            )
+            for name in cameras
+        ]
+        log.info(f"Using calvin_env with commit {get_git_commit_hash(Path(calvin_env.__file__))}.")
+
+    def __del__(self):
+        self.close()
+
+    # From pybullet gym_manipulator_envs code
+    # https://github.com/bulletphysics/bullet3/blob/master/examples/pybullet/gym/pybullet_envs/gym_manipulator_envs.py
+    def initialize_bullet(self, bullet_time_step, render_width, render_height):
+        if self.cid < 0:
+            self.ownsPhysicsClient = True
+            if self.use_vr:
+                self.p = bc.BulletClient(connection_mode=p.SHARED_MEMORY)
+                cid = self.p._client
+                if cid < 0:
+                    log.error("Failed to connect to SHARED_MEMORY bullet server.\n" " Is it running?")
+                    sys.exit(1)
+                self.p.setRealTimeSimulation(enableRealTimeSimulation=1, physicsClientId=cid)
+            elif self.show_gui:
+                self.p = bc.BulletClient(connection_mode=p.GUI)
+                cid = self.p._client
+                if cid < 0:
+                    log.error("Failed to connect to GUI.")
+            elif self.use_egl:
+                options = f"--width={render_width} --height={render_height}"
+                self.p = p
+                cid = self.p.connect(p.DIRECT, options=options)
+                p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=cid)
+                p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=cid)
+                p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=cid)
+                p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=cid)
+                egl = pkgutil.get_loader("eglRenderer")
+                log.info("Loading EGL plugin (may segfault on misconfigured systems)...")
+                if egl:
+                    plugin = p.loadPlugin(egl.get_filename(), "_eglRendererPlugin")
+                else:
+                    plugin = p.loadPlugin("eglRendererPlugin")
+                if plugin < 0:
+                    log.error("\nPlugin Failed to load!\n")
+                    sys.exit()
+                # set environment variable for tacto renderer
+                os.environ["PYOPENGL_PLATFORM"] = "egl"
+                log.info("Successfully loaded egl plugin")
+            else:
+                self.p = bc.BulletClient(connection_mode=p.DIRECT)
+                cid = self.p._client
+                if cid < 0:
+                    log.error("Failed to start DIRECT bullet mode.")
+            log.info(f"Connected to server with id: {cid}")
+
+            self.cid = cid
+            self.p.resetSimulation(physicsClientId=self.cid)
+            self.p.setPhysicsEngineParameter(deterministicOverlappingPairs=1, physicsClientId=self.cid)
+            self.p.configureDebugVisualizer(self.p.COV_ENABLE_GUI, 0)
+            log.info(f"Connected to server with id: {self.cid}")
+            self.p.setTimeStep(1.0 / bullet_time_step, physicsClientId=self.cid)
+            return cid
+
+    def load(self):
+        log.info("Resetting simulation")
+        self.p.resetSimulation(physicsClientId=self.cid)
+        log.info("Setting gravity")
+        self.p.setGravity(0, 0, -9.8, physicsClientId=self.cid)
+
+        self.robot.load()
+        self.scene.load()
+
+    def close(self):
+        if self.ownsPhysicsClient:
+            print("disconnecting id %d from server" % self.cid)
+            if self.cid >= 0 and self.p is not None:
+                try:
+                    self.p.disconnect(physicsClientId=self.cid)
+                except TypeError:
+                    pass
+
+        else:
+            print("does not own physics client id")
+
+    def render(self, mode="human"):
+        """render is gym compatibility function"""
+        rgb_obs, depth_obs = self.get_camera_obs()
+        if mode == "human":
+            if "rgb_static" in rgb_obs:
+                img = rgb_obs["rgb_static"][:, :, ::-1]
+                cv2.imshow("simulation cam", cv2.resize(img, (500, 500)))
+            if "rgb_gripper" in rgb_obs:
+                img2 = rgb_obs["rgb_gripper"][:, :, ::-1]
+                cv2.imshow("gripper cam", cv2.resize(img2, (500, 500)))
+            cv2.waitKey(1)
+        elif mode == "rgb_array":
+            assert "rgb_static" in rgb_obs, "Environment does not have static camera"
+            return rgb_obs["rgb_static"]
+        else:
+            raise NotImplementedError
+
+    def get_scene_info(self):
+        return self.scene.get_info()
+
+    def reset(self, robot_obs=None, scene_obs=None):
+        self.scene.reset(scene_obs)
+        self.robot.reset(robot_obs)
+        self.p.stepSimulation(physicsClientId=self.cid)
+        return self.get_obs()
+
+    def seed(self, seed=None):
+        self.np_random, seed = gym.utils.seeding.np_random(seed)
+        # self.robot.np_random = self.np_random  # use the same np_randomizer for robot as for env
+        return [seed]
+
+    def get_camera_obs(self):
+        assert self.cameras is not None
+        rgb_obs = {}
+        depth_obs = {}
+        for cam in self.cameras:
+            rgb, depth = cam.render()
+            rgb_obs[f"rgb_{cam.name}"] = rgb
+            depth_obs[f"depth_{cam.name}"] = depth
+        return rgb_obs, depth_obs
+
+    def get_obs(self):
+        """Collect camera, robot and scene observations."""
+        rgb_obs, depth_obs = self.get_camera_obs()
+        obs = {"rgb_obs": rgb_obs, "depth_obs": depth_obs}
+        obs.update(self.get_state_obs())
+        return obs
+
+    def get_state_obs(self):
+        """
+        Collect state observation dict
+        --state_obs
+            --robot_obs
+                --robot_state_full
+                    -- [tcp_pos, tcp_orn, gripper_opening_width]
+                --gripper_opening_width
+                --arm_joint_states
+                --gripper_action}
+            --scene_obs
+        """
+        robot_obs, robot_info = self.robot.get_observation()
+        scene_obs = self.scene.get_obs()
+        obs = {"robot_obs": robot_obs, "scene_obs": scene_obs}
+        return obs
+
+    def get_info(self):
+        _, robot_info = self.robot.get_observation()
+        info = {"robot_info": robot_info}
+        if self.use_scene_info:
+            info["scene_info"] = self.scene.get_info()
+        return info
+
+    def step(self, action):
+        # in vr mode real time simulation is enabled, thus p.stepSimulation() does not have to be called manually
+        if self.use_vr:
+            log.debug(f"SIM FPS: {(1 / (time.time() - self.t)):.0f}")
+            self.t = time.time()
+            current_time = time.time()
+            delta_t = current_time - self.prev_time
+            if delta_t >= (1.0 / self.control_freq):
+                log.debug(f"Act FPS: {1 / delta_t:.0f}")
+                self.prev_time = time.time()
+                self.robot.apply_action(action)
+            self.fps_controller.step()
+        # for RL call step simulation repeat
+        else:
+            self.robot.apply_action(action)
+            for i in range(self.action_repeat):
+                self.p.stepSimulation(physicsClientId=self.cid)
+        self.scene.step()
+        obs = self.get_obs()
+        info = self.get_info()
+        # obs, reward, done, info
+        return obs, 0, False, info
+
+    def reset_from_storage(self, filename):
+        """
+        Args:
+            filename: file to load from.
+        Returns:
+            observation
+        """
+        with open(filename, "rb") as file:
+            data = pickle.load(file)
+
+        self.robot.reset_from_storage(data["robot"])
+        self.scene.reset_from_storage(data["scene"])
+
+        self.p.stepSimulation(physicsClientId=self.cid)
+
+        return data["state_obs"], data["done"], data["info"]
+
+    def serialize(self):
+        data = {"time": time.time_ns() / (10**9), "robot": self.robot.serialize(), "scene": self.scene.serialize()}
+        return data
+
+
+def get_env(dataset_path, obs_space=None, show_gui=True, **kwargs):
+    from pathlib import Path
+
+    from omegaconf import OmegaConf
+
+    render_conf = OmegaConf.load(Path(dataset_path) / ".hydra" / "merged_config.yaml")
+
+    if obs_space is not None:
+        exclude_keys = set(render_conf.cameras.keys()) - {
+            re.split("_", key)[1] for key in obs_space["rgb_obs"] + obs_space["depth_obs"]
+        }
+        for k in exclude_keys:
+            del render_conf.cameras[k]
+    if "scene" in kwargs:
+        scene_cfg = OmegaConf.load(Path(calvin_env.__file__).parents[1] / "conf/scene" / f"{kwargs['scene']}.yaml")
+        OmegaConf.merge(render_conf, scene_cfg)
+    if not hydra.core.global_hydra.GlobalHydra.instance().is_initialized():
+        hydra.initialize(".")
+    env = hydra.utils.instantiate(render_conf.env, show_gui=show_gui, use_vr=False, use_scene_info=True)
+    return env
+
+
+@hydra.main(config_path="../../conf", config_name="config_data_collection")
+def run_env(cfg):
+    env = hydra.utils.instantiate(cfg.env, show_gui=True, use_vr=False, use_scene_info=True)
+
+    env.reset()
+    while True:
+        env.step(np.array((0.0, 0, 0, 0, 0, 1)))
+        # env.render()
+        time.sleep(0.01)
+
+
+if __name__ == "__main__":
+    run_env()

aloha_robot_project/calvin/calvin_env/calvin_env/envs/tasks.py

@@ -0,0 +1,306 @@
+from functools import partial
+from typing import Dict, Set
+
+import numpy as np
+from omegaconf import ListConfig
+from scipy.spatial.transform import Rotation as R
+
+
+class Tasks:
+    def __init__(self, tasks):
+        """
+        A task is defined as a specific change between the start_info and end_info dictionaries.
+        Use config file in conf/tasks/ to define tasks using the base task functions defined in this class
+        """
+        # register task functions from config file
+        self.tasks = {name: partial(getattr(self, args[0]), *args[1:]) for name, args in dict(tasks).items()}
+        # dictionary mapping from task name to task id
+        self.task_to_id = {name: i for i, name in enumerate(self.tasks.keys())}
+        # dictionary mapping from task id to task name
+        self.id_to_task = {i: name for i, name in enumerate(self.tasks.keys())}
+
+    def get_task_info(self, start_info: Dict, end_info: Dict) -> Set:
+        """
+        start_info: dict with scene info and robot info
+        end_info: dict with scene info and robot info
+        returns set with achieved tasks
+        """
+        # call functions that are registered in self.tasks
+        return {
+            task_name
+            for task_name, function in self.tasks.items()
+            if function(start_info=start_info, end_info=end_info)
+        }
+
+    def get_task_info_for_set(self, start_info: Dict, end_info: Dict, task_filter: Set) -> Set:
+        """
+        start_info: dict with scene info and robot info
+        end_info: dict with scene info and robot info
+        task_filter: set with task names to check
+        returns set with achieved tasks
+        """
+        # call functions that are registered in self.tasks
+        return {
+            task_name
+            for task_name, function in self.tasks.items()
+            if task_name in task_filter and function(start_info=start_info, end_info=end_info)
+        }
+
+    @property
+    def num_tasks(self):
+        return len(self.tasks)
+
+    @staticmethod
+    def rotate_object(
+        obj_name, z_degrees, x_y_threshold=30, z_threshold=180, movement_threshold=0.1, start_info=None, end_info=None
+    ):
+        """
+        Returns True if the object with obj_name was rotated more than z_degrees degrees around the z-axis while not
+        being rotated more than x_y_threshold degrees around the x or y axis.
+        z_degrees is negative for clockwise rotations and positive for counter-clockwise rotations.
+        """
+        obj_start_info = start_info["scene_info"]["movable_objects"][obj_name]
+        obj_end_info = end_info["scene_info"]["movable_objects"][obj_name]
+        start_orn = R.from_quat(obj_start_info["current_orn"])
+        end_orn = R.from_quat(obj_end_info["current_orn"])
+        rotation = end_orn * start_orn.inv()
+        x, y, z = rotation.as_euler("xyz", degrees=True)
+
+        start_pos = np.array(obj_start_info["current_pos"])
+        end_pos = np.array(obj_end_info["current_pos"])
+        pos_diff = end_pos - start_pos
+        if np.linalg.norm(pos_diff) > movement_threshold:
+            return False
+
+        end_contacts = set(c[2] for c in obj_end_info["contacts"])
+        robot_uid = {start_info["robot_info"]["uid"]}
+        if len(end_contacts - robot_uid) == 0:
+            return False
+
+        if z_degrees > 0:
+            return z_degrees < z < z_threshold and abs(x) < x_y_threshold and abs(y) < x_y_threshold
+        else:
+            return z_degrees > z > -z_threshold and abs(x) < x_y_threshold and abs(y) < x_y_threshold
+
+    @staticmethod
+    def push_object(obj_name, x_direction, y_direction, start_info, end_info):
+        """
+        Returns True if the object with 'obj_name' was moved more than 'x_direction' meters in x direction
+        (or 'y_direction' meters in y direction analogously).
+        Note that currently x and y pushes are mutually exclusive, meaning that one of the arguments has to be 0.
+        The sign matters, e.g. pushing an object to the right when facing the table coincides with a movement in
+        positive x-direction.
+        """
+        assert x_direction * y_direction == 0 and x_direction + y_direction != 0
+        obj_start_info = start_info["scene_info"]["movable_objects"][obj_name]
+        obj_end_info = end_info["scene_info"]["movable_objects"][obj_name]
+        start_pos = np.array(obj_start_info["current_pos"])
+        end_pos = np.array(obj_end_info["current_pos"])
+        pos_diff = end_pos - start_pos
+
+        robot_uid = start_info["robot_info"]["uid"]
+        # contacts excluding robot
+        start_contacts = set((c[2], c[4]) for c in obj_start_info["contacts"] if c[2] != robot_uid)
+        end_contacts = set((c[2], c[4]) for c in obj_end_info["contacts"] if c[2] != robot_uid)
+
+        # computing set difference to check if object had surface contact (excluding robot) at both times
+        surface_contact = len(start_contacts) > 0 and len(end_contacts) > 0 and start_contacts <= end_contacts
+        if not surface_contact:
+            return False
+
+        if x_direction > 0:
+            return pos_diff[0] > x_direction
+        elif x_direction < 0:
+            return pos_diff[0] < x_direction
+
+        if y_direction > 0:
+            return pos_diff[1] > y_direction
+        elif y_direction < 0:
+            return pos_diff[1] < y_direction
+
+    @staticmethod
+    def lift_object(obj_name, z_direction, surface_body=None, surface_link=None, start_info=None, end_info=None):
+        """
+        Returns True if the object with 'obj_name' was grasped by the robot and lifted more than 'z_direction' meters.
+        """
+        assert z_direction > 0
+        obj_start_info = start_info["scene_info"]["movable_objects"][obj_name]
+        obj_end_info = end_info["scene_info"]["movable_objects"][obj_name]
+
+        start_pos = np.array(obj_start_info["current_pos"])
+        end_pos = np.array(obj_end_info["current_pos"])
+        pos_diff = end_pos - start_pos
+        z_diff = pos_diff[2]
+
+        robot_uid = start_info["robot_info"]["uid"]
+        start_contacts = set(c[2] for c in obj_start_info["contacts"])
+        end_contacts = set(c[2] for c in obj_end_info["contacts"])
+
+        surface_criterion = True
+        if surface_body and surface_link is None:
+            surface_uid = start_info["scene_info"]["fixed_objects"][surface_body]["uid"]
+            surface_criterion = surface_uid in start_contacts
+        elif surface_body and surface_link:
+            surface_uid = start_info["scene_info"]["fixed_objects"][surface_body]["uid"]
+            surface_link_id = start_info["scene_info"]["fixed_objects"][surface_body]["links"][surface_link]
+            start_contacts_links = set((c[2], c[4]) for c in obj_start_info["contacts"])
+            surface_criterion = (surface_uid, surface_link_id) in start_contacts_links
+
+        return (
+            z_diff > z_direction
+            # and robot_uid not in start_contacts
+            and robot_uid in end_contacts
+            and len(end_contacts) == 1
+            and surface_criterion
+        )
+
+    @staticmethod
+    def place_object(dest_body, dest_link=None, start_info=None, end_info=None):
+        """
+        Returns True if the object that the robot has currently lifted is placed on the body 'dest_body'
+        (on 'dest_link' if provided).
+        The robot may not touch the object after placing.
+        """
+        robot_uid = start_info["robot_info"]["uid"]
+
+        robot_contacts_start = set(c[2] for c in start_info["robot_info"]["contacts"])
+        robot_contacts_end = set(c[2] for c in end_info["robot_info"]["contacts"])
+        if not len(robot_contacts_start) == 1:
+            return False
+        obj_uid = list(robot_contacts_start)[0]
+
+        if obj_uid in robot_contacts_end:
+            return False
+        _obj_name = [k for k, v in start_info["scene_info"]["movable_objects"].items() if v["uid"] == obj_uid]
+        if not len(_obj_name) == 1:
+            return False
+        obj_name = _obj_name[0]
+
+        dest_uid = end_info["scene_info"]["fixed_objects"][dest_body]["uid"]
+
+        object_contacts_start = set(c[2] for c in start_info["scene_info"]["movable_objects"][obj_name]["contacts"])
+        if dest_link is None:
+            object_contacts_end = set(c[2] for c in end_info["scene_info"]["movable_objects"][obj_name]["contacts"])
+            return (
+                robot_uid in object_contacts_start
+                and len(object_contacts_start) == 1
+                and dest_uid in object_contacts_end
+            )
+        else:
+            dest_link_id = end_info["scene_info"]["fixed_objects"][dest_body]["links"][dest_link]
+            end_contacts_links = set(
+                (c[2], c[4]) for c in end_info["scene_info"]["movable_objects"][obj_name]["contacts"]
+            )
+            return (
+                robot_uid in object_contacts_start
+                and len(object_contacts_start) == 1
+                and (dest_uid, dest_link_id) in end_contacts_links
+            )
+
+    @staticmethod
+    def push_object_into(obj_name, src_body, src_link, dest_body, dest_link, start_info=None, end_info=None):
+        """
+        obj_name is either a list of object names or a string
+        Returns True if the object / any of the objects changes contact from src_body to dest_body.
+        The robot may neither touch the object at start nor end.
+        """
+        if isinstance(obj_name, (list, ListConfig)):
+            return any(
+                Tasks.push_object_into(ob, src_body, src_link, dest_body, dest_link, start_info, end_info)
+                for ob in obj_name
+            )
+        robot_uid = start_info["robot_info"]["uid"]
+
+        src_uid = start_info["scene_info"]["fixed_objects"][src_body]["uid"]
+        src_link_id = start_info["scene_info"]["fixed_objects"][src_body]["links"][src_link]
+        dest_uid = end_info["scene_info"]["fixed_objects"][dest_body]["uid"]
+        dest_link_id = end_info["scene_info"]["fixed_objects"][dest_body]["links"][dest_link]
+
+        start_contacts = set((c[2], c[4]) for c in start_info["scene_info"]["movable_objects"][obj_name]["contacts"])
+        end_contacts = set((c[2], c[4]) for c in end_info["scene_info"]["movable_objects"][obj_name]["contacts"])
+        return (
+            robot_uid not in start_contacts | end_contacts
+            and len(start_contacts) == 1
+            and (src_uid, src_link_id) in start_contacts
+            and (dest_uid, dest_link_id) in end_contacts
+        )
+
+    @staticmethod
+    def move_door_abs(joint_name, start_threshold, end_threshold, start_info, end_info):
+        """
+        Returns True if the joint specified by 'obj_name' and 'joint_name' (e.g. a door or drawer)
+        is moved from at least 'start_threshold' to 'end_threshold'.
+        """
+        start_joint_state = start_info["scene_info"]["doors"][joint_name]["current_state"][0]
+        end_joint_state = end_info["scene_info"]["doors"][joint_name]["current_state"][0]
+
+        if start_threshold < end_threshold:
+            return start_joint_state < start_threshold < end_threshold < end_joint_state
+        elif start_threshold > end_threshold:
+            return start_joint_state > start_threshold > end_threshold > end_joint_state
+        else:
+            raise ValueError
+
+    @staticmethod
+    def move_door_rel(joint_name, threshold, start_info, end_info):
+        """
+        Returns True if the joint specified by 'obj_name' and 'joint_name' (e.g. a door or drawer)
+        is moved from at least 'start_threshold' to 'end_threshold'.
+        """
+        start_joint_state = start_info["scene_info"]["doors"][joint_name]["current_state"]
+        end_joint_state = end_info["scene_info"]["doors"][joint_name]["current_state"]
+
+        return (
+            0 < threshold < end_joint_state - start_joint_state or 0 > threshold > end_joint_state - start_joint_state
+        )
+
+    @staticmethod
+    def toggle_light(light_name, start_state, end_state, start_info, end_info):
+        return (
+            start_info["scene_info"]["lights"][light_name]["logical_state"] == start_state
+            and end_info["scene_info"]["lights"][light_name]["logical_state"] == end_state
+        )
+
+    @staticmethod
+    def stack_objects(max_vel=1, start_info=None, end_info=None):
+        obj_uids = set(obj["uid"] for obj in start_info["scene_info"]["movable_objects"].values())
+
+        for obj_name in start_info["scene_info"]["movable_objects"]:
+            obj_start_info = start_info["scene_info"]["movable_objects"][obj_name]
+            obj_end_info = end_info["scene_info"]["movable_objects"][obj_name]
+            obj_start_contacts = set(c[2] for c in obj_start_info["contacts"])
+            obj_end_contacts = set(c[2] for c in obj_end_info["contacts"])
+
+            if (
+                not len(obj_uids & obj_start_contacts)
+                and len(obj_uids & obj_end_contacts)
+                and not len(obj_end_contacts - obj_uids)
+            ):
+                # object velocity may not exceed max_vel for successful stack
+                if np.all(np.abs(obj_end_info["current_lin_vel"]) < max_vel) and np.all(
+                    np.abs(obj_end_info["current_ang_vel"]) < max_vel
+                ):
+                    return True
+        return False
+
+    @staticmethod
+    def unstack_objects(max_vel=1, start_info=None, end_info=None):
+        obj_uids = set(obj["uid"] for obj in start_info["scene_info"]["movable_objects"].values())
+
+        for obj_name in start_info["scene_info"]["movable_objects"]:
+            obj_start_info = start_info["scene_info"]["movable_objects"][obj_name]
+            obj_end_info = end_info["scene_info"]["movable_objects"][obj_name]
+            obj_start_contacts = set(c[2] for c in obj_start_info["contacts"])
+            obj_end_contacts = set(c[2] for c in obj_end_info["contacts"])
+
+            if (
+                len(obj_uids & obj_start_contacts)
+                and not len(obj_start_contacts - obj_uids)
+                and not len(obj_uids & obj_end_contacts)
+            ):
+                # object velocity may not exceed max_vel for successful stack
+                if np.all(np.abs(obj_start_info["current_lin_vel"]) < max_vel) and np.all(
+                    np.abs(obj_start_info["current_ang_vel"]) < max_vel
+                ):
+                    return True
+        return False

aloha_robot_project/calvin/calvin_env/calvin_env/io_utils/data_recorder.py

@@ -0,0 +1,135 @@
+import logging
+import multiprocessing as mp
+import os
+import pickle
+import time
+
+# A logger for this file
+log = logging.getLogger(__name__)
+
+
+class DataRecorder:
+    """
+    Collects frame information to file in output-dir with
+    filename: <TIMESTAMP>/<FRAME>.pickle
+    Saving facility with separate worker thread.
+    """
+
+    def __init__(self, env, record_fps, enable_tts):
+        """
+        Setup MultiprocessingStorage
+        """
+        self.env = env
+        self.queue = mp.Queue()
+        self.process = mp.Process(target=self.process_queue, name="MultiprocessingStorageWorker")
+        self.process.start()
+        self.running = True
+        self.save_frame_cnt = 0
+        log.info("Starting serialization worker process")
+        self.prev_time = time.time()
+        self.loop_time = 1.0 / record_fps
+        self._unsaved_vr_events = []
+        self.enable_tts = enable_tts
+        if enable_tts:
+            import pyttsx3
+
+            self.tts = pyttsx3.init()
+            self.tts.setProperty("rate", 175)
+        self.prev_done = False
+        self.current_episode_filenames = []
+
+    def step(self, prev_vr_event, state_obs, done, info):
+        self._unsaved_vr_events.extend(prev_vr_event)
+        current_time = time.time()
+        delta_t = current_time - self.prev_time
+        if delta_t >= self.loop_time or done:
+            log.debug(f"Record FPS: {1 / delta_t:.0f}")
+            self.prev_time = time.time()
+            file_path = f"{str(self.save_frame_cnt).zfill(12)}.pickle"
+            self.save(file_path, self._unsaved_vr_events, state_obs, done, info)
+            if self.prev_done and not done:
+                self.current_episode_filenames = []
+            self.current_episode_filenames.append(file_path)
+
+            # file_path_state = f"{str(self.save_frame_cnt).zfill(12)}_state.pickle"
+            # with open(file_path_state, 'wb') as file:
+            #     pickle.dump(state_obs, file)
+
+            self._unsaved_vr_events = []
+            self.save_frame_cnt += 1
+            self.prev_done = done
+
+    def save(self, filename, vr_events, state_obs, done, info):
+        """
+        Extract dataFrame from pybullet and enqueue for worker thread.
+
+        Args:
+            filename: path to file
+            vr_events: vrEvents to attach to data
+            state_obs: state observations
+            done: true if episode ends
+            info: info dict
+
+        Returns:
+            None
+        """
+        data = self.env.serialize()
+        data["vr_events"] = vr_events
+        data["state_obs"] = state_obs
+        data["done"] = done
+        data["info"] = info
+        self.queue.put((filename, data))
+
+    def delete_episode(self):
+        num_frames = len(self.current_episode_filenames)
+        if self.enable_tts:
+            self.tts.say(f"Deleting last episode with {num_frames} frames")
+            self.tts.runAndWait()
+        for filename in self.current_episode_filenames:
+            os.remove(filename)
+        if self.enable_tts:
+            self.tts.say("Finished deleting")
+            self.tts.runAndWait()
+        self.save_frame_cnt -= num_frames
+        self.current_episode_filenames = []
+
+    def process_queue(self):
+        """
+        Process function for queue.
+        Returns:
+            None
+        """
+        while True:
+            msg = self.queue.get()
+            if msg == "QUIT":
+                self.running = False
+                break
+            (filename, data) = msg
+            with open(filename, "wb") as file:
+                pickle.dump(data, file)
+
+    def close(self):
+        """
+        Tell Worker to shut down.
+        Returns:
+            None
+        """
+        if self.running:
+            self.queue.put("QUIT")
+            self.process.join()
+
+    def __enter__(self):
+        """
+            with ... as ... : logic
+        Returns:
+            None
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+            with ... as ... : logic
+        Returns:
+            None
+        """
+        self.close()

aloha_robot_project/calvin/calvin_env/calvin_env/io_utils/vr_input.py

@@ -0,0 +1,187 @@
+import collections
+import logging
+
+import numpy as np
+import pybullet as p
+import quaternion  # noqa
+
+import calvin_env.utils.utils as utils
+
+# A logger for this file
+log = logging.getLogger(__name__)
+
+
+# identical for both controllers
+RELEVANT_BUTTONS = {
+    "button_a": 7,
+    "button_b": 1,
+    "thumb_trigger": 2,  # on the left of the (right) controller; also 34 sent
+    "index_trigger": 33,
+}
+
+
+class ButtonState(collections.namedtuple("ButtonState", ("button_name", "is_down", "was_triggered", "was_released"))):
+    __slots__ = ()
+
+    @classmethod
+    def from_flag(cls, button_name, flag):
+        return cls(
+            button_name,
+            bool(flag & p.VR_BUTTON_IS_DOWN),
+            bool(flag & p.VR_BUTTON_WAS_TRIGGERED),
+            bool(flag & p.VR_BUTTON_WAS_RELEASED),
+        )
+
+
+class VrEvent(
+    collections.namedtuple(
+        "VrEvent",
+        [
+            "controller_id",
+            "position",
+            "orientation",
+            "analog_axis",
+            "n_button_events",
+            "n_move_events",
+            "buttons",
+            "device_type_flag",
+        ],
+    )
+):
+    __slots__ = ()
+
+    @property
+    def button_dicts(self):
+        return {
+            button_name: ButtonState.from_flag(button_name, self.buttons[button_idx])
+            for button_name, button_idx in RELEVANT_BUTTONS.items()
+        }
+        # overall 64 buttons according to OpenVR
+        # enumerate([f'button_{i}' for i in range(64)])}
+
+    @property
+    def device_type(self):
+        return {p.VR_DEVICE_HMD: "hmd", p.VR_DEVICE_CONTROLLER: "controller", p.VR_DEVICE_GENERIC_TRACKER: "generic"}[
+            self.device_type_flag
+        ]
+
+
+class VrInput:
+    """
+    pyBullet VR Data Collector
+    """
+
+    def __init__(self, vr_controller, limit_angle, visualize_vr_pos, reset_button_queue_len):
+        self.vr_controller_id = vr_controller.vr_controller_id
+        self.vr_controller = vr_controller
+        self.POSITION = vr_controller.POSITION
+        self.ORIENTATION = vr_controller.ORIENTATION
+        self.ANALOG = vr_controller.ANALOG
+        self.BUTTONS = vr_controller.BUTTONS
+        self.BUTTON_A = vr_controller.BUTTON_A
+        self.BUTTON_B = vr_controller.BUTTON_B
+        if limit_angle is not None:
+            self.limit_angle = limit_angle[0]
+            self.limit_vector = limit_angle[1:]
+        else:
+            self.limit_angle = None
+        self.gripper_position_offset = np.array(vr_controller.gripper_position_offset)
+        self.gripper_orientation_offset = quaternion.from_euler_angles(vr_controller.gripper_orientation_offset)
+        self.visualize_vr_pos = visualize_vr_pos
+        self.vr_pos_uid = None
+        if visualize_vr_pos:
+            self.vr_pos_uid = self.create_vr_pos_visualization_shape()
+        log.info("Disable Picking")
+        p.configureDebugVisualizer(p.COV_ENABLE_VR_PICKING, 0)
+        p.configureDebugVisualizer(p.COV_ENABLE_VR_RENDER_CONTROLLERS, 0)
+        self._prev_vr_events = None
+        self.prev_action = None
+        self._reset_button_pressed = False
+        self._start_button_pressed = False
+        self._prev_reset_button_pressed = False
+        self._prev_start_button_pressed = False
+        self.reset_button_press_counter = 0
+        self.reset_button_queue_len = reset_button_queue_len
+        # wait until first vr action event arrives
+        while self.prev_action is None:
+            _ = self.get_vr_action()
+
+    def create_vr_pos_visualization_shape(self):
+        visual_shape_id = p.createVisualShape(shapeType=p.GEOM_SPHERE, rgbaColor=[1, 0, 0, 1], radius=0.005)
+        return p.createMultiBody(baseMass=0, baseVisualShapeIndex=visual_shape_id, basePosition=[0, 0, 0])
+
+    def get_vr_action(self):
+        # self._reset_button_pressed = False
+        self._prev_reset_button_pressed = self._reset_button_pressed
+        self._prev_start_button_pressed = self._start_button_pressed
+        vr_events = p.getVREvents()
+        if vr_events != ():
+            self._prev_vr_events = vr_events
+            for event in vr_events:
+                # Only use one controller
+                # if event[0] == self.vr_controller_id:
+                action = (event[self.POSITION], event[self.ORIENTATION], event[self.ANALOG])
+
+                self._reset_button_pressed = event[self.BUTTONS][self.BUTTON_B] & p.VR_BUTTON_IS_DOWN > 0
+                self._start_button_pressed = event[self.BUTTONS][self.BUTTON_A] & p.VR_BUTTON_IS_DOWN > 0
+                robot_action = self.vr_action_to_robot_action(action)
+
+                self.update_reset_button_queue()
+                self.prev_action = robot_action
+                return robot_action
+
+        return self.prev_action
+
+    def vr_action_to_robot_action(self, action):
+        controller_pos, vr_controller_orientation, controller_analogue_axis = action
+        desired_ee_pos = controller_pos + self.gripper_position_offset
+        orientation = utils.xyzw_to_wxyz(vr_controller_orientation)
+        q1 = quaternion.from_float_array(orientation)
+        q2 = self.gripper_orientation_offset
+        q12 = q1 * q2
+        arr = quaternion.as_float_array(q12)
+        desired_ee_orn = utils.wxyz_to_xyzw(arr)
+
+        # v2 = quaternion.rotate_vectors(q12, [0, 0, 1])
+        # if self.limit_angle is not None:
+        #     if utils.angle_between(self.limit_vector, v2) <= self.limit_angle / 180 * math.pi:
+        #         self.prev_ee_orn = desired_ee_orn
+        #     else:
+        #         desired_ee_orn = self.prev_ee_orn
+
+        if controller_analogue_axis > 0.1:
+            gripper_action = -1
+        else:
+            gripper_action = 1
+
+        # change color of vr pos sphere when starting or ending recording
+        if self.visualize_vr_pos:
+            if self.start_button_pressed:
+                p.changeVisualShape(self.vr_pos_uid, -1, rgbaColor=[0, 1, 0, 1])
+            if self.reset_button_pressed:
+                p.changeVisualShape(self.vr_pos_uid, -1, rgbaColor=[1, 0, 0, 1])
+            p.resetBasePositionAndOrientation(self.vr_pos_uid, desired_ee_pos, desired_ee_orn)
+
+        return desired_ee_pos, desired_ee_orn, gripper_action
+
+    def update_reset_button_queue(self):
+        if self._reset_button_pressed:
+            self.reset_button_press_counter += 1
+        else:
+            self.reset_button_press_counter = 0
+
+    @property
+    def reset_button_hold(self):
+        return self.reset_button_press_counter >= self.reset_button_queue_len
+
+    @property
+    def reset_button_pressed(self):
+        return self._reset_button_pressed and not self._prev_reset_button_pressed
+
+    @property
+    def start_button_pressed(self):
+        return self._start_button_pressed and not self._prev_start_button_pressed
+
+    @property
+    def prev_vr_events(self):
+        return self._prev_vr_events

aloha_robot_project/calvin/calvin_env/calvin_env/robot/IKfast.py

@@ -0,0 +1,72 @@
+from ikfast_franka_panda import get_ik
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+
+
+class IKfast:
+    def __init__(
+        self,
+        robot_uid,
+        cid,
+        rp,
+        ll_real,
+        ul_real,
+        base_position,
+        base_orientation,
+        weights=[1, 1, 1, 1, 1, 1, 1],
+        num_angles=50,
+    ):
+        self.robot_uid = robot_uid
+        self.cid = cid
+        self.ll_real = ll_real
+        self.ul_real = ul_real
+        self.rp = rp
+        self.num_dof = len(self.ll_real)
+        self.weights = weights
+        self.num_angles = num_angles
+        T_world_robot = np.eye(4)
+        T_world_robot[:3, 3] = base_position
+        T_world_robot[:3, :3] = R.from_quat(base_orientation).as_matrix()
+        self.T_robot_world = np.linalg.inv(T_world_robot)
+
+    def world_to_robot(self, pos_w, orn_w):
+        """
+        pos, quat -> pos, Rot
+        """
+        pose_w = np.eye(4)
+        pose_w[:3, 3] = pos_w
+        pose_w[:3, :3] = R.from_quat(orn_w).as_matrix()
+        pose_r = self.T_robot_world @ pose_w
+        pos_r = list(pose_r[:3, 3])
+        orn_r = pose_r[:3, :3].tolist()
+        return pos_r, orn_r
+
+    def filter_solutions(self, sol):
+        test_sol = np.ones(self.num_dof) * 9999.0
+        for i in range(self.num_dof):
+            for add_ang in [-2.0 * np.pi, 0, 2.0 * np.pi]:
+                test_ang = sol[i] + add_ang
+                if self.ul_real[i] >= test_ang >= self.ll_real[i]:
+                    test_sol[i] = test_ang
+        if np.all(test_sol != 9999.0):
+            return test_sol
+        return None
+
+    def take_closest_sol(self, sols, last_q, weights):
+        best_sol_ind = np.argmin(np.sum((weights * (sols - np.array(last_q))) ** 2, 1))
+        return sols[best_sol_ind]
+
+    def get_ik_solution(self, target_pos, target_orn):
+        target_pos_robot, target_orn_robot = self.world_to_robot(target_pos, target_orn)
+        sols = []
+        feasible_sols = []
+        for q_6 in np.linspace(self.ll_real[-1], self.ul_real[-1], self.num_angles):
+            sols += get_ik(target_pos_robot, target_orn_robot, [q_6])
+        for sol in sols:
+            sol = self.filter_solutions(sol)
+            if sol is not None:
+                feasible_sols.append(sol)
+        if len(feasible_sols) < 1:
+            return None
+        best_sol = self.take_closest_sol(feasible_sols, self.rp[:7], self.weights)
+        return best_sol

aloha_robot_project/calvin/calvin_env/calvin_env/robot/mixed_ik.py

@@ -0,0 +1,118 @@
+import numpy as np
+import pybullet as p
+from scipy.spatial.transform import Rotation as R
+
+from calvin_env.utils.utils import angle_between_quaternions
+
+
+class MixedIK:
+    def __init__(
+        self,
+        robot_uid,
+        cid,
+        ll_real,
+        ul_real,
+        base_position,
+        base_orientation,
+        tcp_link_id,
+        ll,
+        ul,
+        jr,
+        rp,
+        use_ik_fast,
+        use_nullspace=True,
+        threshold_pos=0.01,
+        threshold_orn=0.05,
+        weights=(1, 1, 1, 1, 1, 1, 1),
+        num_angles=50,
+    ):
+        self.robot_uid = robot_uid
+        self.cid = cid
+        self.use_nullspace = use_nullspace
+        self.use_ik_fast = use_ik_fast
+        self.ik_fast = None
+        if self.use_ik_fast:
+            from ikfast_franka_panda import get_fk
+
+            from calvin_env.robot.IKfast import IKfast
+
+            self.get_fk = get_fk
+            self.ik_fast = IKfast(
+                robot_uid, cid, rp, ll_real, ul_real, base_position, base_orientation, weights, num_angles
+            )
+        self.tcp_link_id = tcp_link_id
+        self.ll = ll
+        self.ul = ul
+        self.jr = jr
+        self.ll_real = ll_real
+        self.ul_real = ul_real
+        self.rp = rp
+        self.num_dof = len(self.ll_real)
+        self.threshold_pos = threshold_pos
+        self.threshold_orn = threshold_orn
+        self.is_using_IK_fast = False
+
+    def get_bullet_ik(self, desired_ee_pos, desired_ee_orn):
+        if self.use_nullspace:
+            jnt_ps = p.calculateInverseKinematics(
+                self.robot_uid,
+                self.tcp_link_id,
+                desired_ee_pos,
+                desired_ee_orn,
+                self.ll,
+                self.ul,
+                self.jr,
+                self.rp,
+                physicsClientId=self.cid,
+            )
+        else:
+            jnt_ps = p.calculateInverseKinematics(
+                self.robot_uid, self.tcp_link_id, desired_ee_pos, desired_ee_orn, physicsClientId=self.cid
+            )
+        # clip joint positions outside the joint ranges
+        jnt_ps = np.clip(jnt_ps[: self.num_dof], self.ll_real, self.ul_real)
+        return jnt_ps
+
+    def robot_to_world(self, pos_r, orn_r):
+        """
+        pos, Rot -> pos, quat
+        """
+        pose_r = np.eye(4)
+        pose_r[:3, 3] = pos_r
+        pose_r[:3, :3] = orn_r
+        pose_w = np.linalg.inv(self.ik_fast.T_robot_world) @ pose_r
+        pos_r = pose_w[:3, 3]
+        orn_r = R.from_matrix(pose_w[:3, :3]).as_quat()
+        return pos_r, orn_r
+
+    def pose_within_threshold(self, target_pos, target_orn, q):
+        pos, orn = self.get_fk(q)
+        pos, orn = self.robot_to_world(pos, orn)
+        angular_diff = angle_between_quaternions(orn, target_orn)
+        threshold_pos_exceeded = np.linalg.norm(target_pos - pos) > self.threshold_pos
+        threshold_orn_exceeded = angular_diff > self.threshold_orn
+        return not (threshold_pos_exceeded or threshold_orn_exceeded)
+
+    def get_joint_states(self):
+        return list(zip(*p.getJointStates(self.robot_uid, range(self.num_dof))))[0]
+
+    def get_ik(self, target_pos, target_orn):
+        if self.is_using_IK_fast and not self.pose_within_threshold(target_pos, target_orn, self.get_joint_states()):
+            q_ik_fast = self.ik_fast.get_ik_solution(target_pos, target_orn)
+            if q_ik_fast is not None:
+                self.is_using_IK_fast = True
+                return q_ik_fast
+            else:
+                self.is_using_IK_fast = False
+                q_bullet = self.get_bullet_ik(target_pos, target_orn)
+                return q_bullet
+        self.is_using_IK_fast = False
+        q_bullet = self.get_bullet_ik(target_pos, target_orn)
+        if self.use_ik_fast and not self.pose_within_threshold(target_pos, target_orn, q_bullet):
+            q_ik_fast = self.ik_fast.get_ik_solution(target_pos, target_orn)
+            if q_ik_fast is not None:
+                self.is_using_IK_fast = True
+                return q_ik_fast
+            else:
+                return q_bullet
+        return q_bullet

aloha_robot_project/calvin/calvin_env/calvin_env/robot/robot.py

@@ -0,0 +1,412 @@
+import logging
+
+import numpy as np
+import pybullet as p
+
+from calvin_env.robot.mixed_ik import MixedIK
+
+# A logger for this file
+log = logging.getLogger(__name__)
+
+
+class Robot:
+    def __init__(
+        self,
+        filename,
+        base_position,
+        base_orientation,
+        initial_joint_positions,
+        max_joint_force,
+        gripper_force,
+        arm_joint_ids,
+        gripper_joint_ids,
+        gripper_joint_limits,
+        tcp_link_id,
+        end_effector_link_id,
+        cid,
+        use_nullspace,
+        max_velocity,
+        use_ik_fast,
+        euler_obs,
+        lower_joint_limits=(-2.8973, -1.7628, -2.8973, -3.0718, -2.8973, -0.0175, -2.8973),
+        upper_joint_limits=(2.8973, 1.7628, 2.8973, -0.0698, 2.8973, 3.7525, 2.8973),
+        max_rel_pos=0.02,
+        max_rel_orn=0.05,
+        magic_scaling_factor_pos=1,
+        magic_scaling_factor_orn=1,
+        use_target_pose=True,
+        **kwargs,
+    ):
+        log.info("Loading robot")
+        self.cid = cid
+        self.filename = filename
+        self.use_nullspace = use_nullspace
+        self.max_velocity = max_velocity
+        self.use_ik_fast = use_ik_fast
+        self.base_position = base_position
+        self.base_orientation = p.getQuaternionFromEuler(base_orientation)
+        self.arm_joint_ids = arm_joint_ids
+        self.initial_joint_positions = np.array(initial_joint_positions)
+        self.gripper_joint_ids = gripper_joint_ids
+        self.max_joint_force = max_joint_force
+        self.gripper_force = gripper_force
+        self.gripper_joint_limits = gripper_joint_limits
+        self.tcp_link_id = tcp_link_id
+        # Setup constraint
+        self.prev_ee_orn = p.getQuaternionFromEuler([0, 0, 0])
+        self.robot_uid = None
+        self.end_effector_link_id = end_effector_link_id
+        self.gripper_action = 1
+        self.ll = self.ul = self.jr = self.rp = None
+        self.ll_real = np.array(lower_joint_limits)
+        self.ul_real = np.array(upper_joint_limits)
+        self.mixed_ik = None
+        self.euler_obs = euler_obs
+        self.max_rel_pos = max_rel_pos
+        self.max_rel_orn = max_rel_orn
+        self.magic_scaling_factor_pos = magic_scaling_factor_pos
+        self.magic_scaling_factor_orn = magic_scaling_factor_orn
+        self.target_pos = None
+        self.target_orn = None
+        self.use_target_pose = use_target_pose
+        # self.reconfigure = False
+
+    def load(self):
+        self.robot_uid = p.loadURDF(
+            fileName=self.filename,
+            basePosition=self.base_position,
+            baseOrientation=self.base_orientation,
+            useFixedBase=True,
+            physicsClientId=self.cid,
+        )
+        self.add_base_cylinder()
+        # create a constraint to keep the fingers centered
+        c = p.createConstraint(
+            self.robot_uid,
+            self.gripper_joint_ids[0],
+            self.robot_uid,
+            self.gripper_joint_ids[1],
+            jointType=p.JOINT_GEAR,
+            jointAxis=[1, 0, 0],
+            parentFramePosition=[0, 0, 0],
+            childFramePosition=[0, 0, 0],
+            physicsClientId=self.cid,
+        )
+        p.changeConstraint(c, gearRatio=-1, erp=0.1, maxForce=50, physicsClientId=self.cid)
+        num_dof = p.computeDofCount(self.robot_uid)
+        # lower limits for null space (todo: set them to proper range)
+        self.ll = [-7] * num_dof
+        # upper limits for null space (todo: set them to proper range)
+        self.ul = [7] * num_dof
+        # joint ranges for null space (todo: set them to proper range)
+        self.jr = [7] * num_dof
+        # restposes for null space
+        self.rp = list(self.initial_joint_positions) + [self.gripper_joint_limits[1]] * 2
+        self.reset()
+        self.mixed_ik = MixedIK(
+            self.robot_uid,
+            self.cid,
+            self.ll_real,
+            self.ul_real,
+            self.base_position,
+            self.base_orientation,
+            self.tcp_link_id,
+            self.ll,
+            self.ul,
+            self.jr,
+            self.rp,
+            self.use_ik_fast,
+            threshold_pos=0.03,
+            threshold_orn=0.1,
+            weights=(10, 8, 6, 6, 2, 2, 1),
+            num_angles=30,
+        )
+
+    def add_base_cylinder(self):
+        """
+        TODO: this should happen in load(), but that would break compatibility with old recorded data
+        """
+        pos = self.base_position.copy()
+        pos[2] /= 2
+        angle = p.getEulerFromQuaternion(self.base_orientation)[2]
+        pos[0] -= np.cos(angle) * 0.05
+        pos[1] -= np.sin(angle) * 0.05
+        cylinder = p.createVisualShape(
+            shapeType=p.GEOM_CYLINDER,
+            rgbaColor=[1, 1, 1, 1],
+            radius=0.13,
+            length=self.base_position[2],
+            visualFramePosition=pos,
+        )
+        p.createMultiBody(baseVisualShapeIndex=cylinder)
+
+    def reset(self, robot_state=None):
+        if robot_state is None:
+            gripper_state = self.gripper_joint_limits[1]
+            joint_states = self.initial_joint_positions
+        else:
+            joint_indices = [i for i, x in enumerate(self.get_observation_labels()) if x.startswith("robot_joint")]
+            joint_states = robot_state[joint_indices]
+            gripper_state = robot_state[self.get_observation_labels().index("gripper_opening_width")] / 2
+
+        assert len(joint_states) == len(self.arm_joint_ids)
+        for i, _id in enumerate(self.arm_joint_ids):
+            p.resetJointState(self.robot_uid, _id, joint_states[i], physicsClientId=self.cid)
+            p.setJointMotorControl2(
+                bodyIndex=self.robot_uid,
+                jointIndex=_id,
+                controlMode=p.POSITION_CONTROL,
+                force=self.max_joint_force,
+                targetPosition=joint_states[i],
+                maxVelocity=self.max_velocity,
+                physicsClientId=self.cid,
+            )
+        for i in self.gripper_joint_ids:
+            p.resetJointState(self.robot_uid, i, gripper_state, physicsClientId=self.cid)
+            p.setJointMotorControl2(
+                bodyIndex=self.robot_uid,
+                jointIndex=i,
+                controlMode=p.POSITION_CONTROL,
+                force=self.gripper_force,
+                targetPosition=gripper_state,
+                maxVelocity=1,
+                physicsClientId=self.cid,
+            )
+        tcp_pos, tcp_orn = p.getLinkState(self.robot_uid, self.tcp_link_id, physicsClientId=self.cid)[:2]
+        if self.euler_obs:
+            tcp_orn = p.getEulerFromQuaternion(tcp_orn)
+        self.target_pos = np.array(tcp_pos)
+        self.target_orn = np.array(tcp_orn)
+
+    def get_observation(self):
+        """
+        returns:
+        - robot_state: ndarray (16,)
+            - tcp_pos: robot_state[:3]
+            - tcp_orn: robot_state[3:7] (quat) / [3:6] (euler)
+            - gripper_opening_width: robot_state[7:8] (quat) / [6:7] (euler)
+            - arm_joint_states: robot_state[8:15] (quat) / [7:14] (euler)
+            - gripper_action: robot_state[15:] (quat) / [14:] (euler)
+        - robot_info: Dict
+        """
+        tcp_pos, tcp_orn = p.getLinkState(self.robot_uid, self.tcp_link_id, physicsClientId=self.cid)[:2]
+        if self.euler_obs:
+            tcp_orn = p.getEulerFromQuaternion(tcp_orn)
+        gripper_opening_width = (
+            p.getJointState(self.robot_uid, self.gripper_joint_ids[0], physicsClientId=self.cid)[0]
+            + p.getJointState(self.robot_uid, self.gripper_joint_ids[1], physicsClientId=self.cid)[0]
+        )
+        arm_joint_states = []
+        for i in self.arm_joint_ids:
+            arm_joint_states.append(p.getJointState(self.robot_uid, i, physicsClientId=self.cid)[0])
+        robot_state = np.array([*tcp_pos, *tcp_orn, gripper_opening_width, *arm_joint_states, self.gripper_action])
+        robot_info = {
+            "tcp_pos": tcp_pos,
+            "tcp_orn": tcp_orn,
+            "gripper_opening_width": gripper_opening_width,
+            "arm_joint_states": arm_joint_states,
+            "gripper_action": self.gripper_action,
+            "uid": self.robot_uid,
+            "contacts": p.getContactPoints(bodyA=self.robot_uid, physicsClientId=self.cid),
+        }
+        return robot_state, robot_info
+
+    def get_observation_labels(self):
+        tcp_pos_labels = [f"tcp_pos_{ax}" for ax in ("x", "y", "z")]
+        if self.euler_obs:
+            tcp_orn_labels = [f"tcp_orn_{ax}" for ax in ("x", "y", "z")]
+        else:
+            tcp_orn_labels = [f"tcp_orn_{ax}" for ax in ("x", "y", "z", "w")]
+        return [
+            *tcp_pos_labels,
+            *tcp_orn_labels,
+            "gripper_opening_width",
+            *[f"robot_joint_{i}" for i in self.arm_joint_ids],
+            "gripper_action",
+        ]
+
+    def relative_to_absolute(self, action):
+        assert len(action) == 7
+        rel_pos, rel_orn, gripper = np.split(action, [3, 6])
+        rel_pos *= self.max_rel_pos * self.magic_scaling_factor_pos
+        rel_orn *= self.max_rel_orn * self.magic_scaling_factor_orn
+        if self.use_target_pose:
+            self.target_pos += rel_pos
+            self.target_orn += rel_orn
+            return self.target_pos, self.target_orn, gripper
+        else:
+            tcp_pos, tcp_orn = p.getLinkState(self.robot_uid, self.tcp_link_id, physicsClientId=self.cid)[:2]
+            tcp_orn = p.getEulerFromQuaternion(tcp_orn)
+            abs_pos = np.array(tcp_pos) + rel_pos
+            abs_orn = np.array(tcp_orn) + rel_orn
+            return abs_pos, abs_orn, gripper
+
+    def apply_action(self, action):
+        # cv2.imshow("win", np.zeros((300,300)))
+        # k = cv2.waitKey(1) % 255
+        # if k == ord('w'):
+        #     self.base_position[1] += 0.01
+        # elif k == ord('s'):
+        #     self.base_position[1] -= 0.01
+        # elif k == ord('d'):
+        #     self.base_position[0] += 0.01
+        # elif k == ord('a'):
+        #     self.base_position[0] -= 0.01
+        # elif k == ord('e'):
+        #     self.base_position[2] += 0.01
+        # elif k == ord('q'):
+        #     self.base_position[2] -= 0.01
+        # elif k == ord('r'):
+        #     self.initial_joint_positions[0] -= 0.1
+        # elif k == ord('f'):
+        #     self.initial_joint_positions[0] += 0.1
+        # elif k == ord('t'):
+        #     self.initial_joint_positions[1] -= 0.1
+        # elif k == ord('g'):
+        #     self.initial_joint_positions[1] += 0.1
+        # elif k == ord('y'):
+        #     self.initial_joint_positions[2] -= 0.1
+        # elif k == ord('h'):
+        #     self.initial_joint_positions[2] += 0.1
+        # elif k == ord('u'):
+        #     self.initial_joint_positions[3] -= 0.1
+        # elif k == ord('j'):
+        #     self.initial_joint_positions[3] += 0.1
+        # elif k == ord('i'):
+        #     self.initial_joint_positions[4] -= 0.1
+        # elif k == ord('k'):
+        #     self.initial_joint_positions[4] += 0.1
+        # elif k == ord('o'):
+        #     self.initial_joint_positions[5] -= 0.1
+        # elif k == ord('l'):
+        #     self.initial_joint_positions[5] += 0.1
+        # elif k == ord('p'):
+        #     self.initial_joint_positions[6] -= 0.1
+        # elif k == ord(';'):
+        #     self.initial_joint_positions[6] += 0.1
+        # elif k == ord('z'):
+        #     self.reconfigure = not self.reconfigure
+        #     print(f"{self.initial_joint_positions=}")
+        #     print(f"{self.base_position=}")
+        # if k != 254:
+        #     self.initial_joint_positions = np.clip(self.initial_joint_positions, self.ll_real, self.ul_real)
+        #     p.resetBasePositionAndOrientation(self.robot_uid, self.base_position, self.base_orientation, physicsClientId=self.cid)
+        #     self.rp = list(self.initial_joint_positions) + [self.gripper_joint_limits[1]] * 2
+        #     self.mixed_ik.rp = self.rp
+        #     for i, _id in enumerate(self.arm_joint_ids):
+        #         p.resetJointState(self.robot_uid, _id, self.initial_joint_positions[i], physicsClientId=self.cid)
+        #         p.setJointMotorControl2(
+        #             bodyIndex=self.robot_uid,
+        #             jointIndex=_id,
+        #             controlMode=p.POSITION_CONTROL,
+        #             force=self.max_joint_force,
+        #             targetPosition=self.initial_joint_positions[i],
+        #             maxVelocity=self.max_velocity,
+        #             physicsClientId=self.cid,
+        #         )
+        # if self.reconfigure:
+        #     return
+        #
+
+        if not len(action) == 3:
+            action = self.relative_to_absolute(action)
+        target_ee_pos, target_ee_orn, self.gripper_action = action
+
+        assert len(target_ee_pos) == 3
+        assert len(target_ee_orn) in (3, 4)
+        # automatically transform euler actions to quaternion
+        if len(target_ee_orn) == 3:
+            target_ee_orn = p.getQuaternionFromEuler(target_ee_orn)
+
+        if not isinstance(self.gripper_action, int) and len(self.gripper_action) == 1:
+            self.gripper_action = self.gripper_action[0]
+        assert self.gripper_action in (-1, 1)
+
+        # #
+        # cam_rot = p.getMatrixFromQuaternion(target_ee_orn)
+        # cam_rot = np.array(cam_rot).reshape(3, 3)
+        # cam_rot_x, cam_rot_y, cam_rot_z = cam_rot[:, 0], cam_rot[:, 1], cam_rot[:, 2]
+        # p.addUserDebugLine(target_ee_pos, target_ee_pos + cam_rot_x, lineWidth=3, lineColorRGB=[1,0,0])
+        # p.addUserDebugLine(target_ee_pos, target_ee_pos +cam_rot_y, lineWidth=3, lineColorRGB=[0,1,0])
+        # p.addUserDebugLine(target_ee_pos, target_ee_pos +cam_rot_z, lineWidth=3, lineColorRGB=[0,0,1])
+        #
+        # tcp_pos, tcp_orn = p.getLinkState(self.robotId, self.tcp_link_id)[:2]
+        # tcp_euler = p.getEulerFromQuaternion(tcp_orn)
+        # p.addUserDebugLine([0,0,0], target_ee_pos, lineWidth=8, lineColorRGB=[0,1,0])
+        # p.addUserDebugLine([0,0,0], p.getLinkState(self.robot_uid, 6)[4], lineWidth=3, lineColorRGB=[1,0,0])
+        # p.addUserDebugLine([0,0,0], p.getLinkState(self.robot_uid, 13)[4], lineWidth=3, lineColorRGB=[0,1,0])
+        # target_ee_pos, target_ee_orn = self.tcp_to_ee(target_ee_pos, target_ee_orn)
+        # p.addUserDebugLine([0,0,0], target_ee_pos, lineWidth=8, lineColorRGB=[1,0,0])
+        jnt_ps = self.mixed_ik.get_ik(target_ee_pos, target_ee_orn)
+        for i in range(self.end_effector_link_id):
+            # p.resetJointState(self.robot_uid, i, jnt_ps[i])
+            p.setJointMotorControl2(
+                bodyIndex=self.robot_uid,
+                jointIndex=i,
+                controlMode=p.POSITION_CONTROL,
+                force=self.max_joint_force,
+                targetPosition=jnt_ps[i],
+                maxVelocity=self.max_velocity,
+                physicsClientId=self.cid,
+            )
+
+        self.control_gripper(self.gripper_action)
+
+    def control_gripper(self, gripper_action):
+        if gripper_action == 1:
+            gripper_finger_position = self.gripper_joint_limits[1]
+            gripper_force = self.gripper_force / 100
+        else:
+            gripper_finger_position = self.gripper_joint_limits[0]
+            gripper_force = self.gripper_force
+        for id in self.gripper_joint_ids:
+            p.setJointMotorControl2(
+                bodyIndex=self.robot_uid,
+                jointIndex=id,
+                controlMode=p.POSITION_CONTROL,
+                targetPosition=gripper_finger_position,
+                force=gripper_force,
+                maxVelocity=1,
+                physicsClientId=self.cid,
+            )
+
+    def serialize(self):
+        return {
+            "uid": self.robot_uid,
+            "info": p.getBodyInfo(self.robot_uid, physicsClientId=self.cid),
+            "pose": p.getBasePositionAndOrientation(self.robot_uid, physicsClientId=self.cid),
+            "joints": p.getJointStates(
+                self.robot_uid,
+                list(range(p.getNumJoints(self.robot_uid, physicsClientId=self.cid))),
+                physicsClientId=self.cid,
+            ),
+            "gripper_action": self.gripper_action,
+        }
+
+    def reset_from_storage(self, data):
+        p.resetBasePositionAndOrientation(
+            bodyUniqueId=self.robot_uid, posObj=data["pose"][0], ornObj=data["pose"][1], physicsClientId=self.cid
+        )
+        num_joints = len(data["joints"])
+        assert num_joints == p.getNumJoints(self.robot_uid, physicsClientId=self.cid)
+        for i, (value, velocity, *_) in enumerate(data["joints"]):
+            p.resetJointState(
+                bodyUniqueId=self.robot_uid,
+                jointIndex=i,
+                targetValue=value,
+                targetVelocity=velocity,
+                physicsClientId=self.cid,
+            )
+            p.setJointMotorControl2(
+                bodyIndex=self.robot_uid,
+                jointIndex=i,
+                controlMode=p.POSITION_CONTROL,
+                force=self.max_joint_force,
+                targetPosition=value,
+                maxVelocity=self.max_velocity,
+                physicsClientId=self.cid,
+            )
+        self.control_gripper(data["gripper_action"])
+
+    def __str__(self):
+        return f"{self.filename} : {self.__dict__}"

aloha_robot_project/calvin/calvin_env/calvin_env/scene/objects/base_object.py

@@ -0,0 +1,13 @@
+class BaseObject:
+    def __init__(self, name, obj_cfg, p, cid, data_path, global_scaling):
+        self.p = p
+        self.cid = cid
+        self.name = name
+        self.file = data_path / obj_cfg["file"]
+        self.global_scaling = global_scaling
+
+    def reset(self, state):
+        pass
+
+    def get_info(self):
+        pass

aloha_robot_project/calvin/calvin_env/calvin_env/scene/objects/button.py

@@ -0,0 +1,76 @@
+from enum import Enum
+
+MAX_FORCE = 4
+
+
+class ButtonState(Enum):
+    ON = 1
+    OFF = 0
+
+
+class Button:
+    def __init__(self, name, cfg, uid, p, cid):
+        self.name = name
+        self.p = p
+        self.cid = cid
+        # get joint_index by name (to prevent index errors when additional joints are added)
+        joint_index = next(
+            i
+            for i in range(self.p.getNumJoints(uid, physicsClientId=self.cid))
+            if self.p.getJointInfo(uid, i, physicsClientId=self.cid)[1].decode("utf-8") == name
+        )
+        self.joint_index = joint_index
+        self.uid = uid
+        self.initial_state = cfg["initial_state"]
+        self.effect = cfg["effect"]
+        self.ll, self.ul = self.p.getJointInfo(uid, joint_index, physicsClientId=self.cid)[8:10]
+        self.trigger_threshold = (self.ll + self.ul) / 2
+        self.p.setJointMotorControl2(
+            self.uid,
+            self.joint_index,
+            controlMode=self.p.POSITION_CONTROL,
+            targetPosition=self.initial_state,
+            force=MAX_FORCE,
+            physicsClientId=self.cid,
+        )
+        self.state = ButtonState.OFF
+        self.prev_is_pressed = self._is_pressed
+        self.light = None
+
+    def reset(self, state=None):
+        _state = self.initial_state if state is None else state
+        self.p.resetJointState(
+            self.uid,
+            self.joint_index,
+            _state,
+            physicsClientId=self.cid,
+        )
+        self.state = ButtonState.OFF
+
+    def step(self):
+        if self.state == ButtonState.OFF and not self.prev_is_pressed and self._is_pressed:
+            self.state = ButtonState.ON
+            if self.light is not None:
+                self.light.turn_on()
+        elif self.state == ButtonState.ON and not self.prev_is_pressed and self._is_pressed:
+            self.state = ButtonState.OFF
+            if self.light is not None:
+                self.light.turn_off()
+        self.prev_is_pressed = self._is_pressed
+
+    @property
+    def _is_pressed(self):
+        if self.initial_state <= self.trigger_threshold:
+            return self.get_state() > self.trigger_threshold
+        elif self.initial_state > self.trigger_threshold:
+            return self.get_state() < self.trigger_threshold
+
+    def get_state(self):
+        """return button joint state"""
+        return float(self.p.getJointState(self.uid, self.joint_index, physicsClientId=self.cid)[0])
+
+    def get_info(self):
+        return {"joint_state": self.get_state(), "logical_state": self.state.value}
+
+    def add_effect(self, light):
+        self.light = light

aloha_robot_project/calvin/calvin_env/calvin_env/scene/objects/door.py

@@ -0,0 +1,40 @@
+MAX_FORCE = 4
+
+
+class Door:
+    def __init__(self, name, cfg, uid, p, cid):
+        self.name = name
+        self.p = p
+        self.cid = cid
+        # get joint_index by name (to prevent index errors when additional joints are added)
+        joint_index = next(
+            i
+            for i in range(self.p.getNumJoints(uid, physicsClientId=self.cid))
+            if self.p.getJointInfo(uid, i, physicsClientId=self.cid)[1].decode("utf-8") == name
+        )
+        self.joint_index = joint_index
+        self.uid = uid
+        self.initial_state = cfg["initial_state"]
+        self.p.setJointMotorControl2(
+            self.uid,
+            self.joint_index,
+            controlMode=p.VELOCITY_CONTROL,
+            force=MAX_FORCE,
+            physicsClientId=self.cid,
+        )
+
+    def reset(self, state=None):
+        _state = self.initial_state if state is None else state
+        self.p.resetJointState(
+            self.uid,
+            self.joint_index,
+            _state,
+            physicsClientId=self.cid,
+        )
+
+    def get_state(self):
+        joint_state = self.p.getJointState(self.uid, self.joint_index, physicsClientId=self.cid)
+        return float(joint_state[0])
+
+    def get_info(self):
+        return {"current_state": self.get_state()}

aloha_robot_project/calvin/calvin_env/calvin_env/scene/objects/fixed_object.py

@@ -0,0 +1,41 @@
+from calvin_env.scene.objects.base_object import BaseObject
+
+
+class FixedObject(BaseObject):
+    def __init__(self, name, obj_cfg, p, cid, data_path, global_scaling):
+        super().__init__(name, obj_cfg, p, cid, data_path, global_scaling)
+        self.initial_pos = obj_cfg["initial_pos"]
+        self.initial_orn = self.p.getQuaternionFromEuler(obj_cfg["initial_orn"])
+
+        self.uid = self.p.loadURDF(
+            self.file.as_posix(),
+            self.initial_pos,
+            self.initial_orn,
+            globalScaling=global_scaling,
+            physicsClientId=self.cid,
+        )
+        self.info_dict = {"uid": self.uid}
+        self.num_joints = self.p.getNumJoints(self.uid, physicsClientId=self.cid)
+        if self.num_joints > 0:
+            # save link names and ids in dictionary
+            links = {
+                self.p.getJointInfo(self.uid, i, physicsClientId=self.cid)[12].decode("utf-8"): i
+                for i in range(self.num_joints)
+            }
+            links["base_link"] = -1
+            self.info_dict["links"] = links
+
+    def reset(self, state=None):
+        pass
+
+    def get_info(self):
+        obj_info = {**self.info_dict, "contacts": self.p.getContactPoints(bodyA=self.uid, physicsClientId=self.cid)}
+        return obj_info
+
+    def serialize(self):
+        joints = (
+            self.p.getJointStates(self.uid, list(range(self.num_joints)), physicsClientId=self.cid)
+            if self.num_joints > 0
+            else ()
+        )
+        return {"uid": self.uid, "info": self.p.getBodyInfo(self.uid, physicsClientId=self.cid), "joints": joints}

aloha_robot_project/calvin/calvin_env/calvin_env/scene/objects/light.py

@@ -0,0 +1,52 @@
+from enum import Enum
+
+
+class LightState(Enum):
+    ON = 1
+    OFF = 0
+
+
+class Light:
+    def __init__(self, name, cfg, uid, p, cid):
+        self.name = name
+        self.uid = uid
+        self.p = p
+        self.cid = cid
+        self.link = cfg["link"]
+        self.link_id = next(
+            i
+            for i in range(self.p.getNumJoints(uid, physicsClientId=self.cid))
+            if self.p.getJointInfo(uid, i, physicsClientId=self.cid)[12].decode("utf-8") == self.link
+        )
+        self.color_on = cfg["color"]
+        self.color_off = [1, 1, 1, 1]
+        self.state = LightState.OFF
+
+    def reset(self, state=None):
+        if state is None:
+            self.turn_off()
+        else:
+            if state == LightState.ON.value:
+                self.turn_on()
+            elif state == LightState.OFF.value:
+                self.turn_off()
+            else:
+                print("Light state can be only 0 or 1.")
+                raise ValueError
+
+    def get_state(self):
+        return self.state.value
+
+    def get_info(self):
+        return {"logical_state": self.get_state()}
+
+    def turn_on(self):
+        self.state = LightState.ON
+        self.p.changeVisualShape(self.uid, self.link_id, rgbaColor=self.color_on, physicsClientId=self.cid)
+
+    def turn_off(self):
+        self.state = LightState.OFF
+        self.p.changeVisualShape(self.uid, self.link_id, rgbaColor=self.color_off, physicsClientId=self.cid)
+
+    def serialize(self):
+        return self.get_info()

aloha_robot_project/calvin/calvin_env/calvin_env/scene/objects/movable_object.py

@@ -0,0 +1,95 @@
+import logging
+
+import numpy as np
+from omegaconf.errors import ConfigKeyError
+
+from calvin_env.scene.objects.base_object import BaseObject
+
+log = logging.getLogger(__name__)
+
+
+class MovableObject(BaseObject):
+    def __init__(self, name, obj_cfg, p, cid, data_path, global_scaling, euler_obs, surfaces, np_random):
+        super().__init__(name, obj_cfg, p, cid, data_path, global_scaling)
+        self.initial_pos = obj_cfg["initial_pos"]
+        self.initial_orn = obj_cfg["initial_orn"]
+        if isinstance(self.initial_pos, list):
+            self.initial_pos = np.array(self.initial_pos)
+        if isinstance(self.initial_orn, list):
+            self.initial_orn = np.array(p.getQuaternionFromEuler(self.initial_orn))
+        self.euler_obs = euler_obs
+        self.surfaces = surfaces
+        self.np_random = np_random
+
+        initial_pos, initial_orn = self.sample_initial_pose()
+        self.uid = self.p.loadURDF(
+            self.file.as_posix(),
+            initial_pos,
+            initial_orn,
+            globalScaling=global_scaling,
+            physicsClientId=self.cid,
+        )
+
+    def reset(self, state=None):
+        if state is None:
+            initial_pos, initial_orn = self.sample_initial_pose()
+        else:
+            initial_pos, initial_orn = np.split(state, [3])
+            if len(initial_orn) == 3:
+                initial_orn = self.p.getQuaternionFromEuler(initial_orn)
+        self.p.resetBasePositionAndOrientation(
+            self.uid,
+            initial_pos,
+            initial_orn,
+            physicsClientId=self.cid,
+        )
+
+    def sample_initial_pose(self):
+        initial_pos = self.initial_pos
+        if isinstance(self.initial_pos, str):
+            if self.initial_pos == "any":
+                surface = self.np_random.choice(list(self.surfaces.keys()))
+                sampling_range = np.array(self.surfaces[surface])
+            else:
+                try:
+                    sampling_range = np.array(self.surfaces[self.initial_pos])
+                except ConfigKeyError:
+                    log.error(f"surface {self.initial_pos} not specified in scene config")
+                    raise KeyError
+            initial_pos = self.np_random.uniform(sampling_range[0], sampling_range[1])
+        initial_orn = self.initial_orn
+        if isinstance(self.initial_orn, str):
+            if self.initial_orn == "any":
+                initial_orn = np.array(
+                    self.p.getQuaternionFromEuler(self.np_random.uniform([0, 0, -np.pi], [0, 0, np.pi]))
+                )
+            else:
+                log.error("Only keyword 'any' supported at the moment")
+                raise ValueError
+        return initial_pos, initial_orn
+
+    def get_state(self):
+        pos, orn = self.p.getBasePositionAndOrientation(self.uid, physicsClientId=self.cid)
+        if self.euler_obs:
+            orn = self.p.getEulerFromQuaternion(orn)
+        return np.concatenate([pos, orn])
+
+    def get_info(self):
+        pos, orn = self.p.getBasePositionAndOrientation(self.uid, physicsClientId=self.cid)
+        lin_vel, ang_vel = self.p.getBaseVelocity(self.uid, physicsClientId=self.cid)
+        obj_info = {
+            "current_pos": pos,
+            "current_orn": orn,
+            "current_lin_vel": lin_vel,
+            "current_ang_vel": ang_vel,
+            "contacts": self.p.getContactPoints(bodyA=self.uid, physicsClientId=self.cid),
+            "uid": self.uid,
+        }
+        return obj_info
+
+    def serialize(self):
+        return {
+            "uid": self.uid,
+            "info": self.p.getBodyInfo(self.uid, physicsClientId=self.cid),
+            "pose": self.p.getBasePositionAndOrientation(self.uid, physicsClientId=self.cid),
+        }

aloha_robot_project/calvin/calvin_env/calvin_env/scene/objects/switch.py

@@ -0,0 +1,73 @@
+from enum import Enum
+
+MAX_FORCE = 4
+
+
+class ButtonState(Enum):
+    ON = 1
+    OFF = 0
+
+
+class Switch:
+    def __init__(self, name, cfg, uid, p, cid):
+        self.name = name
+        self.p = p
+        self.cid = cid
+        # get joint_index by name (to prevent index errors when additional joints are added)
+        joint_index = next(
+            i
+            for i in range(self.p.getNumJoints(uid, physicsClientId=self.cid))
+            if self.p.getJointInfo(uid, i, physicsClientId=self.cid)[1].decode("utf-8") == name
+        )
+        self.joint_index = joint_index
+        self.uid = uid
+        self.initial_state = cfg["initial_state"]
+        self.effect = cfg["effect"]
+        self.ll, self.ul = self.p.getJointInfo(uid, joint_index, physicsClientId=self.cid)[8:10]
+        self.trigger_threshold = (self.ll + self.ul) / 2
+        self.p.setJointMotorControl2(
+            self.uid,
+            self.joint_index,
+            controlMode=p.VELOCITY_CONTROL,
+            force=MAX_FORCE,
+            physicsClientId=self.cid,
+        )
+        self.state = ButtonState.OFF
+        self.light = None
+
+    def reset(self, state=None):
+        _state = self.initial_state if state is None else state
+        self.p.resetJointState(
+            self.uid,
+            self.joint_index,
+            _state,
+            physicsClientId=self.cid,
+        )
+        self.state = ButtonState.OFF
+
+    def step(self):
+        if self.is_pressed:
+            if self.light is not None and self.state == ButtonState.OFF:
+                self.light.turn_on()
+            self.state = ButtonState.ON
+        else:
+            if self.light is not None and self.state == ButtonState.ON:
+                self.light.turn_off()
+            self.state = ButtonState.OFF
+
+    @property
+    def is_pressed(self):
+        if self.initial_state <= self.trigger_threshold:
+            return self.get_state() > self.trigger_threshold
+        elif self.initial_state > self.trigger_threshold:
+            return self.get_state() < self.trigger_threshold
+
+    def get_state(self):
+        """return button joint state"""
+        return float(self.p.getJointState(self.uid, self.joint_index, physicsClientId=self.cid)[0])
+
+    def get_info(self):
+        return {"joint_state": self.get_state(), "logical_state": self.state.value}
+
+    def add_effect(self, light):
+        self.light = light

aloha_robot_project/calvin/calvin_env/calvin_env/scene/play_table_scene.py

@@ -0,0 +1,234 @@
+import itertools
+import logging
+import os
+from pathlib import Path
+
+import numpy as np
+
+# A logger for this file
+from omegaconf import OmegaConf
+
+from calvin_env.scene.objects.button import Button
+from calvin_env.scene.objects.door import Door
+from calvin_env.scene.objects.fixed_object import FixedObject
+from calvin_env.scene.objects.light import Light
+from calvin_env.scene.objects.movable_object import MovableObject
+from calvin_env.scene.objects.switch import Switch
+
+log = logging.getLogger(__name__)
+
+
+REPO_BASE = Path(__file__).parents[2]
+
+
+class PlayTableScene:
+    def __init__(self, objects, data_path, euler_obs, p, cid, global_scaling, surfaces, np_random, **kwargs):
+        self.p = p
+        self.cid = cid
+        self.global_scaling = global_scaling
+        self.euler_obs = euler_obs
+        self.surfaces = surfaces
+        self.np_random = np_random
+        if os.path.isabs(data_path):
+            self.data_path = Path(data_path)
+        else:
+            self.data_path = REPO_BASE / data_path
+        self.p.setAdditionalSearchPath(self.data_path.as_posix())
+        self.object_cfg = OmegaConf.to_container(objects)
+        self.fixed_objects, self.movable_objects = [], []
+        self.doors, self.buttons, self.switches, self.lights = [], [], [], []
+
+    def load(self):
+        self.fixed_objects, self.movable_objects = [], []
+        self.doors, self.buttons, self.switches, self.lights = [], [], [], []
+
+        for name, obj_cfg in self.object_cfg.get("movable_objects", {}).items():
+            self.movable_objects.append(
+                MovableObject(
+                    name,
+                    obj_cfg,
+                    self.p,
+                    self.cid,
+                    self.data_path,
+                    self.global_scaling,
+                    self.euler_obs,
+                    self.surfaces,
+                    self.np_random,
+                )
+            )
+
+        for name, obj_cfg in self.object_cfg.get("fixed_objects", {}).items():
+            fixed_obj = FixedObject(name, obj_cfg, self.p, self.cid, self.data_path, self.global_scaling)
+
+            if "joints" in obj_cfg:
+                for joint_name, cfg in obj_cfg["joints"].items():
+                    door = Door(joint_name, cfg, fixed_obj.uid, self.p, self.cid)
+                    self.doors.append(door)
+            self.fixed_objects.append(fixed_obj)
+
+            if "buttons" in obj_cfg:
+                for button_name, cfg in obj_cfg["buttons"].items():
+                    button = Button(button_name, cfg, fixed_obj.uid, self.p, self.cid)
+                    self.buttons.append(button)
+
+            if "switches" in obj_cfg:
+                for switch_name, cfg in obj_cfg["switches"].items():
+                    switch = Switch(switch_name, cfg, fixed_obj.uid, self.p, self.cid)
+                    self.switches.append(switch)
+
+            if "lights" in obj_cfg:
+                for light_name, cfg in obj_cfg["lights"].items():
+                    light = Light(light_name, cfg, fixed_obj.uid, self.p, self.cid)
+                    self.lights.append(light)
+
+        for light in self.lights:
+            for button_switch in itertools.chain(self.buttons, self.switches):
+                if button_switch.effect == light.name:
+                    button_switch.add_effect(light)
+
+        self.p.loadURDF(os.path.join(self.data_path, "plane/plane.urdf"), physicsClientId=self.cid)
+
+    def reset(self, scene_obs=None):
+        """Reset objects and doors to initial position."""
+        if scene_obs is None:
+            for obj in itertools.chain(self.doors, self.buttons, self.switches, self.lights):
+                obj.reset()
+            self.reset_movable_objects()
+        else:
+            door_info, button_info, switch_info, light_info, obj_info = self.parse_scene_obs(scene_obs)
+
+            for door, state in zip(self.doors, door_info):
+                door.reset(state)
+            for button, state in zip(self.buttons, button_info):
+                button.reset(state)
+            for switch, state in zip(self.switches, switch_info):
+                switch.reset(state)
+            for light, state in zip(self.lights, light_info):
+                light.reset(state)
+            for obj, state in zip(self.movable_objects, obj_info):
+                obj.reset(state)
+
+    def parse_scene_obs(self, scene_obs):
+        # an object pose is composed of position (3) and orientation (4 for quaternion)  / (3 for euler)
+        n_obj = len(self.movable_objects)
+        n_doors = len(self.doors)
+        n_buttons = len(self.buttons)
+        n_switches = len(self.switches)
+        n_lights = len(self.lights)
+
+        split_ids = np.cumsum([n_doors, n_buttons, n_switches, n_lights])
+        door_info, button_info, switch_info, light_info, obj_info = np.split(scene_obs, split_ids)
+        assert len(door_info) == n_doors
+        assert len(button_info) == n_buttons
+        assert len(switch_info) == n_switches
+        assert len(light_info) == n_lights
+        assert len(obj_info) // n_obj in [6, 7]  # depending on euler angles or quaternions
+
+        obj_info = np.split(obj_info, n_obj)
+
+        return door_info, button_info, switch_info, light_info, obj_info
+
+    def reset_movable_objects(self):
+        """reset movable objects such that there are no pairwise contacts"""
+        num_sampling_iterations = 1000
+        for i in range(num_sampling_iterations):
+            for obj in self.movable_objects:
+                obj.reset()
+            self.p.stepSimulation()
+            contact = False
+            for obj_a, obj_b in itertools.combinations(self.movable_objects, 2):
+                if np.any(len(self.p.getContactPoints(bodyA=obj_a.uid, bodyB=obj_b.uid, physicsClientId=self.cid))):
+                    contact = True
+                    break
+            if not contact:
+                return
+        log.error(f"Could not place objects in {num_sampling_iterations} iterations without contacts")
+        return
+
+    def step(self):
+        for button_switch in itertools.chain(self.buttons, self.switches):
+            button_switch.step()
+
+    def get_obs(self):
+        """Return state information of the doors, drawers and shelves."""
+        door_states = [door.get_state() for door in self.doors]
+        button_states = [button.get_state() for button in self.buttons]
+        switch_states = [switch.get_state() for switch in self.switches]
+        light_states = [light.get_state() for light in self.lights]
+        object_poses = list(itertools.chain(*[obj.get_state() for obj in self.movable_objects]))
+
+        return np.concatenate([door_states, button_states, switch_states, light_states, object_poses])
+
+    def get_info(self):
+        """
+        get dictionary of information about the objects in the scene
+        self.objects:
+            obj1:
+                joints:
+                    joint1:
+                        joint_index: int
+                        initial_state: float  # revolute
+                        current_state: float
+                    ...
+                current_pos: [x, y, z]
+                current_orn: [x, y, z, w]  # quaternion
+                contacts: output of pybullet getContactPoints(...)
+                links:  # key exists only if object has num_joints > 0
+                    link1: link_id  #  name: id
+            ...
+        """
+        info = {}
+        info["fixed_objects"] = {}
+        info["movable_objects"] = {}
+        info["doors"] = {}
+        info["buttons"] = {}
+        info["switches"] = {}
+        info["lights"] = {}
+
+        for obj in self.fixed_objects:
+            info["fixed_objects"][obj.name] = obj.get_info()
+        for obj in self.movable_objects:
+            info["movable_objects"][obj.name] = obj.get_info()
+        for obj in self.doors:
+            info["doors"][obj.name] = obj.get_info()
+        for obj in self.buttons:
+            info["buttons"][obj.name] = obj.get_info()
+        for obj in self.switches:
+            info["switches"][obj.name] = obj.get_info()
+        for obj in self.lights:
+            info["lights"][obj.name] = obj.get_info()
+        return info
+
+    def get_scene_obs_labels(self):
+        raise NotImplementedError
+
+    def get_objects(self):
+        return itertools.chain(self.fixed_objects, self.movable_objects)
+
+    def serialize(self):
+        data = {
+            "fixed_objects": [obj.serialize() for obj in self.fixed_objects],
+            "movable_objects": [obj.serialize() for obj in self.movable_objects],
+            "lights": [obj.serialize() for obj in self.lights],
+        }
+        return data
+
+    def reset_from_storage(self, data):
+        for fixed_obj in data["fixed_objects"]:
+            for i, (value, velocity, *_) in enumerate(fixed_obj["joints"]):
+                self.p.resetJointState(
+                    bodyUniqueId=fixed_obj["uid"],
+                    jointIndex=i,
+                    targetValue=value,
+                    targetVelocity=velocity,
+                    physicsClientId=self.cid,
+                )
+        for movable_obj in data["movable_objects"]:
+            self.p.resetBasePositionAndOrientation(
+                bodyUniqueId=movable_obj["uid"],
+                posObj=movable_obj["pose"][0],
+                ornObj=movable_obj["pose"][1],
+                physicsClientId=self.cid,
+            )
+        for light, state in zip(self.lights, data["lights"]):
+            light.reset(state["logical_state"])

aloha_robot_project/calvin/calvin_env/calvin_env/scripts/check_tasks.py

@@ -0,0 +1,103 @@
+from pathlib import Path
+
+import cv2
+import hydra
+import numpy as np
+import pybullet as p
+
+"""
+This script loads a rendered episode and replays it using the recorded actions.
+Optionally, gaussian noise can be added to the actions.
+"""
+
+
+def scale_depth(img):
+    img = np.clip(img, 0, 1.5)
+    img -= img.min()
+    img /= img.max()
+    return img
+
+
+def noise(action, pos_std=0.01, rot_std=1):
+    """
+    adds gaussian noise to position and orientation.
+    units are m for pos and degree for rot
+    """
+    pos, orn, gripper = action
+    rot_std = np.radians(rot_std)
+    pos_noise = np.random.normal(0, pos_std, 3)
+    rot_noise = p.getQuaternionFromEuler(np.random.normal(0, rot_std, 3))
+    pos, orn = p.multiplyTransforms(pos, orn, pos_noise, rot_noise)
+    return pos, orn, gripper
+
+
+@hydra.main(config_path="../../conf", config_name="config_data_collection")
+def run_env(cfg):
+    env = hydra.utils.instantiate(cfg.env, show_gui=True, use_vr=False, use_scene_info=True)
+
+    root_dir = Path(__file__).parents[3] / "dataset/task_D_D/training"
+
+    ep_start_end_ids = np.sort(np.load(root_dir / "ep_start_end_ids.npy"), axis=0)
+
+    seq_len = 32
+    tasks = hydra.utils.instantiate(cfg.tasks)
+    for s, e in ep_start_end_ids:
+        i = s
+        while 1:
+
+            file = root_dir / f"episode_{i:07}.npz"
+            data = np.load(file)
+            # gripper_img = data["rgb_gripper"]
+            # cv2.imshow("gripper", gripper_img[:, :, ::-1])
+            # static_img = data["rgb_static"]
+            # cv2.imshow("static", static_img[:, :, ::-1])
+            # print(data["robot_obs"])
+            # env.render()
+            env.reset(scene_obs=data["scene_obs"], robot_obs=data["robot_obs"])
+            start_info = env.get_info()
+            cv2.imshow("keylistener", np.zeros((300, 300)))
+            k = cv2.waitKey(0) % 256
+            if k == ord("a"):
+                i -= 1
+                i = np.clip(i, s, e - seq_len)
+            if k == ord("d"):
+                i += 1
+                i = np.clip(i, s, e - seq_len)
+            if k == ord("q"):
+                i -= 100
+                i = np.clip(i, s, e - seq_len)
+            if k == ord("z"):
+                i -= seq_len
+                i = np.clip(i, s, e - seq_len)
+            if k == ord("e"):
+                i += 100
+                i = np.clip(i, s, e - seq_len)
+            if k == ord("r"):
+                for j in range(i + 1, i + seq_len):
+                    file = root_dir / f"episode_{j:07d}.npz"
+                    data = np.load(file)
+                    env.reset(scene_obs=data["scene_obs"], robot_obs=data["robot_obs"])
+                    gripper_img = data["rgb_gripper"]
+                    cv2.imshow("gripper", gripper_img[:, :, ::-1])
+                    static_img = data["rgb_static"]
+                    cv2.imshow("static", static_img[:, :, ::-1])
+                    cv2.waitKey(100)
+                end_info = env.get_info()
+                task_info = tasks.get_task_info(start_info, end_info)
+                if len(task_info):
+                    print()
+                    print()
+                    print()
+                    print(task_info)
+                else:
+                    print()
+                    print()
+                    print()
+                    print("No task detected.")
+                i = j
+            if k == ord("n"):  # ESC
+                break
+
+
+if __name__ == "__main__":
+    run_env()

aloha_robot_project/calvin/calvin_env/calvin_env/scripts/convert_gripper_actions.py

@@ -0,0 +1,16 @@
+from pathlib import Path
+import sys
+
+import numpy as np
+from tqdm import tqdm
+
+path = Path(sys.argv[-1])
+
+for subdir in ["training", "validation"]:
+    for file in tqdm((path / subdir).glob("*.npz")):
+        data = np.load(file)
+        if data["rel_actions"][-1] == 0:
+            data = dict(data)
+            data["rel_actions"][-1] = -1
+            data["actions"][-1] = -1
+            np.savez(file, **data)

aloha_robot_project/calvin/calvin_env/calvin_env/scripts/dataset_to_euler.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+from shutil import copyfile, copytree
+
+import numpy as np
+import pybullet as p
+from tqdm import tqdm
+
+load_path = Path("/home/hermannl/phd/data/banana_dataset_01_29/validation")
+
+save_path = Path("/home/hermannl/phd/data/banana_dataset_01_29_euler/validation")
+save_path.mkdir(parents=True, exist_ok=True)
+
+for file in tqdm(load_path.glob("*.npz")):
+    data = np.load(file)
+    robot_obs = data["robot_obs"]
+    robot_obs_euler = np.concatenate([robot_obs[:3], p.getEulerFromQuaternion(robot_obs[3:7]), robot_obs[7:]])
+    scene_obs = data["scene_obs"]
+    scene_obs_euler = scene_obs[:3]
+    for i in range(6):
+        scene_obs_euler = np.append(scene_obs_euler, scene_obs[3 + i * 7 : 3 + i * 7 + 3])
+        scene_obs_euler = np.append(scene_obs_euler, p.getEulerFromQuaternion(scene_obs[3 + i * 7 + 3 : 3 + i * 7 + 7]))
+    actions = data["actions"]
+    actions_euler = np.concatenate([actions[:3], p.getEulerFromQuaternion(actions[3:7]), actions[7:]])
+    data_euler = dict(data.items())
+    data_euler["robot_obs"] = robot_obs_euler
+    data_euler["scene_obs"] = scene_obs_euler
+    data_euler["actions"] = actions_euler
+    np.savez(save_path / file.name, **data_euler)
+
+for file in set(load_path.glob("*")) - set(load_path.glob("*.npz")):
+    if file.is_dir():
+        copytree(file, save_path / file.name)
+    else:
+        copyfile(file, save_path / file.name)

aloha_robot_project/calvin/calvin_env/calvin_env/scripts/record_video_icra.py

@@ -0,0 +1,112 @@
+from copy import deepcopy
+from pathlib import Path
+import time
+
+import cv2
+import hydra
+import matplotlib.pyplot as plt
+import numpy as np
+import pybullet as p
+
+from calvin_env.utils import utils
+
+"""
+This script loads a rendered episode and replays it using the recorded actions.
+Optionally, gaussian noise can be added to the actions.
+"""
+
+
+def noise(action, pos_std=0.01, rot_std=1):
+    """
+    adds gaussian noise to position and orientation.
+    units are m for pos and degree for rot
+    """
+    pos, orn, gripper = action
+    rot_std = np.radians(rot_std)
+    pos_noise = np.random.normal(0, pos_std, 3)
+    rot_noise = p.getQuaternionFromEuler(np.random.normal(0, rot_std, 3))
+    pos, orn = p.multiplyTransforms(pos, orn, pos_noise, rot_noise)
+    return pos, orn, gripper
+
+
+def depth2rgb(img, minval=0, maxval=1):
+    img -= minval
+    img /= maxval - minval
+    img *= 255
+    img = np.clip(img, 0.0, 255.0)
+    img = img.astype(np.uint8)
+    img = np.tile(np.expand_dims(img, axis=2), 3)
+    return img
+
+
+@hydra.main(config_path="../../conf", config_name="config_data_collection")
+def run_env(cfg):
+    env = hydra.utils.instantiate(cfg.env, show_gui=False, use_vr=False, use_scene_info=True)
+
+    root_dir = Path("/home/hermannl/data/calvin_abcd_example")
+
+    ep_start_end_ids = [[100000, 104999]]
+    rel_actions = []
+    tasks = hydra.utils.instantiate(cfg.tasks)
+    prev_info = None
+    t1 = time.time()
+
+    video = cv2.VideoWriter(
+        "/home/hermannl/Documents/calvin/env_A_tactile.avi", cv2.VideoWriter_fourcc(*"XVID"), 30, (240, 160)
+    )
+    video_static = cv2.VideoWriter(
+        "/home/hermannl/Documents/calvin/env_A_static_depth.avi", cv2.VideoWriter_fourcc(*"XVID"), 30, (200, 200)
+    )
+    video_gripper = cv2.VideoWriter(
+        "/home/hermannl/Documents/calvin/env_A_gripper_depth.avi", cv2.VideoWriter_fourcc(*"XVID"), 30, (100, 100)
+    )
+    for s, e in ep_start_end_ids:
+        print("new_episode")
+        for i in range(s, e + 1):
+            file = root_dir / f"episode_{i:07d}.npz"
+            data = np.load(file)
+            # img = data["rgb_static"]
+            # cv2.imshow("win2", cv2.resize(img[:, :, ::-1], (500, 500)))
+            # cv2.waitKey(0)
+
+            # if (i - s) % 32 == 0:
+            #     print(f"reset {i}")
+            obs = env.reset(scene_obs=data["scene_obs"], robot_obs=data["robot_obs"])
+
+            im = obs["rgb_obs"][0]
+            im = np.concatenate([im[:, :, :3], im[:, :, 3:6]], axis=1)
+            cv2.imshow("im", im)
+            cv2.waitKey(1)
+            video.write(im[:, :, ::-1])
+            video_gripper.write(depth2rgb(obs["depth_obs"][1], minval=0.1, maxval=0.5)[:, :, ::-1])
+            video_static.write(depth2rgb(obs["depth_obs"][0], minval=3.5, maxval=5)[:, :, ::-1])
+            # action = data["rel_actions"]
+            action = np.split(data["actions"], [3, 6])
+            action = noise(action)
+
+            rel_actions.append(utils.to_relative_action(data["actions"], data["robot_obs"][:6]))
+            # action = utils.to_relative_action(data["actions"], data["robot_obs"], max_pos=0.04, max_orn=0.1)
+            # tcp_pos, tcp_orn = p.getLinkState(env.robot.robot_uid, env.robot.tcp_link_id, physicsClientId=env.cid)[:2]
+            # tcp_orn = p.getEulerFromQuaternion(tcp_orn)
+            # action2 = utils.to_relative_action(data["actions"], np.concatenate([tcp_pos, tcp_orn]))
+            o, _, _, info = env.step(action)
+            print(info["scene_info"]["lights"]["led"]["logical_state"])
+            if (i - s) % 32 != 0:
+                print(tasks.get_task_info(prev_info, info))
+            else:
+                prev_info = deepcopy(info)
+            time.sleep(0.01)
+    video.release()
+    video_static.release()
+    video_gripper.release()
+    print(time.time() - t1)
+    rel_actions = np.array(rel_actions)
+    for j in range(rel_actions.shape[1]):
+        plt.figure(j)
+        plt.hist(rel_actions[:, j], bins=10)
+        plt.plot()
+        plt.show()
+
+
+if __name__ == "__main__":
+    run_env()

aloha_robot_project/calvin/calvin_env/calvin_env/scripts/render_low_freq.py

@@ -0,0 +1,74 @@
+import argparse
+import itertools
+import os
+from pathlib import Path
+import shutil
+import sys
+
+import numpy as np
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+from calvin_env.utils import utils
+
+
+def pairwise(iterable):
+    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
+    a, b = itertools.tee(iterable)
+    next(b, None)
+    return zip(a, b)
+
+
+parser = argparse.ArgumentParser(description="convert dataset to 15 hz (leave one step out)")
+parser.add_argument("src", type=str)
+parser.add_argument("dest", type=str)
+parser.add_argument("--max_rel_pos", type=float, default=0.04)
+parser.add_argument("--max_rel_orn", type=float, default=0.1)
+args = parser.parse_args()
+src_path = Path(args.src)
+dest_path = Path(args.dest)
+
+os.mkdir(dest_path)
+os.mkdir(dest_path / "training")
+os.mkdir(dest_path / "validation")
+
+new_i = 0
+
+for subdir in ["training", "validation"]:
+    ep_lens = np.load(src_path / subdir / "ep_lens.npy")
+    ep_start_end_ids = np.load(src_path / subdir / "ep_start_end_ids.npy")
+
+    new_ep_lens = []
+    new_ep_start_end_ids = []
+
+    for start, end in tqdm(ep_start_end_ids):
+
+        ep_len = end - start + 1
+        for offset in (0, 1):
+            new_start = new_i
+            for old_i in range(start + offset, end + 1, 2):
+                old_data = np.load(src_path / subdir / f"episode_{old_i:06d}.npz")
+                data = dict(old_data)
+                if old_i < end:
+                    next_old_data = np.load(src_path / subdir / f"episode_{old_i + 1:06d}.npz")
+                    next_data = dict(next_old_data)
+                    data["actions"] = next_data["actions"]
+                    data["rel_actions"] = utils.to_relative_action(
+                        data["actions"], data["robot_obs"], max_pos=args.max_rel_pos, max_orn=args.max_rel_orn
+                    )
+                np.savez(dest_path / subdir / f"episode_{new_i:06d}.npz", **data)
+                new_i += 1
+            new_end = new_i - 1
+            new_ep_len = new_end - new_start + 1
+            new_ep_start_end_ids.append((new_start, new_end))
+            new_ep_lens.append(new_ep_len)
+    np.save(dest_path / subdir / "ep_lens.npy", new_ep_lens)
+    np.save(dest_path / subdir / "ep_start_end_ids.npy", new_ep_start_end_ids)
+    shutil.copy(src_path / subdir / "statistics.yaml", dest_path / subdir)
+    os.makedirs(dest_path / subdir / ".hydra")
+    shutil.copytree(src_path / subdir / ".hydra", dest_path / subdir / ".hydra", dirs_exist_ok=True)
+    cfg = OmegaConf.load(dest_path / subdir / ".hydra/merged_config.yaml")
+    cfg.robot.max_rel_pos = args.max_rel_pos
+    cfg.robot.max_rel_orn = args.max_rel_orn
+    cfg.env.control_freq = 15
+    OmegaConf.save(cfg, dest_path / subdir / ".hydra/merged_config.yaml")

aloha_robot_project/calvin/calvin_env/calvin_env/scripts/reset_env_rendered_episode.py

@@ -0,0 +1,84 @@
+from copy import deepcopy
+import glob
+import os
+from pathlib import Path
+import time
+
+import cv2
+import hydra
+import matplotlib.pyplot as plt
+import numpy as np
+import pybullet as p
+
+from calvin_env.envs.tasks import Tasks
+from calvin_env.utils import utils
+
+"""
+This script loads a rendered episode and replays it using the recorded actions.
+Optionally, gaussian noise can be added to the actions.
+"""
+
+
+def noise(action, pos_std=0.01, rot_std=1):
+    """
+    adds gaussian noise to position and orientation.
+    units are m for pos and degree for rot
+    """
+    pos, orn, gripper = action
+    rot_std = np.radians(rot_std)
+    pos_noise = np.random.normal(0, pos_std, 3)
+    rot_noise = p.getQuaternionFromEuler(np.random.normal(0, rot_std, 3))
+    pos, orn = p.multiplyTransforms(pos, orn, pos_noise, rot_noise)
+    return pos, orn, gripper
+
+
+@hydra.main(config_path="../../conf", config_name="config_data_collection")
+def run_env(cfg):
+    env = hydra.utils.instantiate(cfg.env, show_gui=False, use_vr=False, use_scene_info=True)
+
+    root_dir = Path("/tmp/test_render/2021-10-05/16-51-11")
+
+    ep_start_end_ids = np.sort(np.load(root_dir / "ep_start_end_ids.npy"), axis=0)
+    rel_actions = []
+    tasks = hydra.utils.instantiate(cfg.tasks)
+    prev_info = None
+    t1 = time.time()
+    for s, e in ep_start_end_ids:
+        print("new_episode")
+        for i in range(s, e + 1):
+            file = root_dir / f"episode_{i:07d}.npz"
+            data = np.load(file)
+            img = data["rgb_static"]
+            cv2.imshow("win2", cv2.resize(img[:, :, ::-1], (500, 500)))
+            cv2.waitKey(1)
+
+            if (i - s) % 32 == 0:
+                print(f"reset {i}")
+                env.reset(scene_obs=data["scene_obs"], robot_obs=data["robot_obs"])
+            action = data["rel_actions"]
+            # action = np.split(data["actions"], [3, 6])
+            # action = noise(action)
+
+            # rel_actions.append(create_relative_action(data["actions"], data["robot_obs"][:6]))
+            # action = utils.to_relative_action(data["actions"], data["robot_obs"], max_pos=0.04, max_orn=0.1)
+            # tcp_pos, tcp_orn = p.getLinkState(env.robot.robot_uid, env.robot.tcp_link_id, physicsClientId=env.cid)[:2]
+            # tcp_orn = p.getEulerFromQuaternion(tcp_orn)
+            # action2 = create_relative_action(data["actions"], np.concatenate([tcp_pos, tcp_orn]))
+            o, _, _, info = env.step(action)
+            print(info["scene_info"]["lights"]["led"]["logical_state"])
+            if (i - s) % 32 != 0:
+                print(tasks.get_task_info(prev_info, info))
+            else:
+                prev_info = deepcopy(info)
+            time.sleep(0.01)
+    print(time.time() - t1)
+    # rel_actions = np.array(rel_actions)
+    # for j in range(rel_actions.shape[1]):
+    #     plt.figure(j)
+    #     plt.hist(rel_actions[:, j], bins=10)
+    #     plt.plot()
+    #     plt.show()
+
+
+if __name__ == "__main__":
+    run_env()

aloha_robot_project/calvin/calvin_env/calvin_env/scripts/unnormalize_depth.py

@@ -0,0 +1,24 @@
+from pathlib import Path
+from shutil import copyfile, copytree
+
+import numpy as np
+from tqdm import tqdm
+
+load_path = Path("/home/meeso/expert_demos_03_10/training")
+
+save_path = Path("/home/meeso/expert_demos_03_10/training_unnormalized_depth")
+save_path.mkdir(parents=True, exist_ok=True)
+
+for file in tqdm(load_path.glob("*.npz")):
+    data = np.load(file)
+    corrected_data = dict(data.items())
+    corrected_data["depth_static"] = data["depth_static"] * 2.0
+    corrected_data["depth_gripper"] = data["depth_gripper"] * 2.0
+    np.savez(save_path / file.name, **corrected_data)
+
+
+for file in set(load_path.glob("*")) - set(load_path.glob("*.npz")):
+    if file.is_dir():
+        copytree(file, save_path / file.name)
+    else:
+        copyfile(file, save_path / file.name)

aloha_robot_project/calvin/calvin_env/calvin_env/utils/utils.py

@@ -0,0 +1,233 @@
+import itertools
+import logging
+import os
+from pathlib import Path
+import pickle
+import re
+import subprocess
+import time
+from typing import Union
+
+import git
+import numpy as np
+import quaternion
+
+# A logger for this file
+logger = logging.getLogger(__name__)
+
+
+def timeit(method):
+    def timed(*args, **kw):
+        ts = time.time()
+        result = method(*args, **kw)
+        te = time.time()
+        if "log_time" in kw:
+            name = kw.get("log_name", method.__name__.upper())
+            kw["log_time"][name] = int((te - ts) * 1000)
+        else:
+            print("%r  %2.2f ms" % (method.__name__, (te - ts) * 1000))
+        return result
+
+    return timed
+
+
+class FpsController:
+    def __init__(self, freq):
+        self.loop_time = (1.0 / freq) * 10**9
+        self.prev_time = time.time_ns()
+
+    def step(self):
+        current_time = time.time_ns()
+        delta_t = current_time - self.prev_time
+        if delta_t < self.loop_time:
+            nano_sleep(self.loop_time - delta_t)
+        self.prev_time = time.time_ns()
+
+
+def xyzw_to_wxyz(arr):
+    """
+    Convert quaternions from pyBullet to numpy.
+    """
+    return [arr[3], arr[0], arr[1], arr[2]]
+
+
+def wxyz_to_xyzw(arr):
+    """
+    Convert quaternions from numpy to pyBullet.
+    """
+    return [arr[1], arr[2], arr[3], arr[0]]
+
+
+def nano_sleep(time_ns):
+    """
+    Spinlock style sleep function. Burns cpu power on purpose
+    equivalent to time.sleep(time_ns / (10 ** 9)).
+
+    Should be more precise, especially on Windows.
+
+    Args:
+        time_ns: time to sleep in ns
+
+    Returns:
+
+    """
+    wait_until = time.time_ns() + time_ns
+    while time.time_ns() < wait_until:
+        pass
+
+
+def unit_vector(vector):
+    """Returns the unit vector of the vector."""
+    return vector / np.linalg.norm(vector)
+
+
+def angle_between_quaternions(q1, q2):
+    """
+    Returns the minimum rotation angle between to orientations expressed as quaternions
+    quaternions use X,Y,Z,W convention
+    """
+    q1 = xyzw_to_wxyz(q1)
+    q2 = xyzw_to_wxyz(q2)
+    q1 = quaternion.from_float_array(q1)
+    q2 = quaternion.from_float_array(q2)
+
+    theta = 2 * np.arcsin(np.linalg.norm((q1 * q2.conjugate()).vec))
+    return theta
+
+
+def angle_between(v1, v2):
+    """Returns the angle in radians between vectors 'v1' and 'v2'::
+
+    >>> angle_between((1, 0, 0), (0, 1, 0))
+    1.5707963267948966
+    >>> angle_between((1, 0, 0), (1, 0, 0))
+    0.0
+    >>> angle_between((1, 0, 0), (-1, 0, 0))
+    3.141592653589793
+    """
+    v1_u = unit_vector(v1)
+    v2_u = unit_vector(v2)
+    return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))
+
+
+def get_git_commit_hash(repo_path: Path) -> str:
+    repo = git.Repo(search_parent_directories=True, path=repo_path.parent)
+    assert repo, "not a repo"
+    changed_files = [item.a_path for item in repo.index.diff(None)]
+    if changed_files:
+        print("WARNING uncommitted modified files: {}".format(",".join(changed_files)))
+    return repo.head.object.hexsha
+
+
+class EglDeviceNotFoundError(Exception):
+    """Raised when EGL device cannot be found"""
+
+
+def get_egl_device_id(cuda_id: int) -> Union[int]:
+    """
+    >>> i = get_egl_device_id(0)
+    >>> isinstance(i, int)
+    True
+    """
+    assert isinstance(cuda_id, int), "cuda_id has to be integer"
+    dir_path = Path(__file__).absolute().parents[2] / "egl_check"
+    if not os.path.isfile(dir_path / "EGL_options.o"):
+        if os.environ.get("LOCAL_RANK", "0") == "0":
+            print("Building EGL_options.o")
+            subprocess.call(["bash", "build.sh"], cwd=dir_path)
+        else:
+            # In case EGL_options.o has to be built and multiprocessing is used, give rank 0 process time to build
+            time.sleep(5)
+    result = subprocess.run(["./EGL_options.o"], capture_output=True, cwd=dir_path)
+    n = int(result.stderr.decode("utf-8").split(" of ")[1].split(".")[0])
+    for egl_id in range(n):
+        my_env = os.environ.copy()
+        my_env["EGL_VISIBLE_DEVICE"] = str(egl_id)
+        result = subprocess.run(["./EGL_options.o"], capture_output=True, cwd=dir_path, env=my_env)
+        match = re.search(r"CUDA_DEVICE=[0-9]+", result.stdout.decode("utf-8"))
+        if match:
+            current_cuda_id = int(match[0].split("=")[1])
+            if cuda_id == current_cuda_id:
+                return egl_id
+    raise EglDeviceNotFoundError
+
+
+def angle_between_angles(a, b):
+    diff = b - a
+    return (diff + np.pi) % (2 * np.pi) - np.pi
+
+
+def to_relative_action(actions, robot_obs, max_pos=0.02, max_orn=0.05):
+    assert isinstance(actions, np.ndarray)
+    assert isinstance(robot_obs, np.ndarray)
+
+    rel_pos = actions[:3] - robot_obs[:3]
+    rel_pos = np.clip(rel_pos, -max_pos, max_pos) / max_pos
+
+    rel_orn = angle_between_angles(robot_obs[3:6], actions[3:6])
+    rel_orn = np.clip(rel_orn, -max_orn, max_orn) / max_orn
+
+    gripper = actions[-1:]
+    return np.concatenate([rel_pos, rel_orn, gripper])
+
+
+def set_egl_device(device):
+    assert "EGL_VISIBLE_DEVICES" not in os.environ, "Do not manually set EGL_VISIBLE_DEVICES"
+    try:
+        cuda_id = device.index if device.type == "cuda" else 0
+    except AttributeError:
+        cuda_id = device
+    try:
+        egl_id = get_egl_device_id(cuda_id)
+    except EglDeviceNotFoundError:
+        logger.warning(
+            "Couldn't find correct EGL device. Setting EGL_VISIBLE_DEVICE=0. "
+            "When using DDP with many GPUs this can lead to OOM errors. "
+            "Did you install PyBullet correctly? Please refer to VREnv README"
+        )
+        egl_id = 0
+    os.environ["EGL_VISIBLE_DEVICES"] = str(egl_id)
+    logger.info(f"EGL_DEVICE_ID {egl_id} <==> CUDA_DEVICE_ID {cuda_id}")
+
+
+def count_frames(directory):
+    """
+    counts the number of consecutive pickled frames in directory
+
+    Args:
+        directory: str of directory
+
+    Returns:
+         0 for none, otherwise >0
+    """
+
+    for i in itertools.count(start=0):
+        pickle_file = os.path.join(directory, f"{str(i).zfill(12)}.pickle")
+        if not os.path.isfile(pickle_file):
+            return i
+
+
+def get_episode_lengths(load_dir, num_frames):
+    episode_lengths = []
+    render_start_end_ids = [[0]]
+    i = 0
+    for frame in range(num_frames):
+        file_path = os.path.abspath(os.path.join(load_dir, f"{str(frame).zfill(12)}.pickle"))
+        with open(file_path, "rb") as file:
+            data = pickle.load(file)
+            done = data["done"]
+            if not done:
+                i += 1
+            else:
+                episode_lengths.append(i)
+                render_start_end_ids[-1].append(frame + 1)
+                render_start_end_ids.append([frame + 1])
+                i = 0
+    render_start_end_ids = render_start_end_ids[:-1]
+    return episode_lengths, render_start_end_ids
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

aloha_robot_project/calvin/calvin_env/calvin_env/vrdatacollector.py

@@ -0,0 +1,52 @@
+#!/usr/bin/python3
+from copy import deepcopy
+import logging
+import os
+import sys
+
+import hydra
+import pybullet as p
+import quaternion  # noqa
+
+from calvin_env.io_utils.data_recorder import DataRecorder
+from calvin_env.io_utils.vr_input import VrInput
+
+# A logger for this file
+log = logging.getLogger(__name__)
+
+
+@hydra.main(config_path="../conf", config_name="config_data_collection")
+def main(cfg):
+    # Load Scene
+    env = hydra.utils.instantiate(cfg.env)
+    vr_input = hydra.utils.instantiate(cfg.vr_input)
+
+    data_recorder = None
+    if cfg.recorder.record:
+        data_recorder = DataRecorder(env, cfg.recorder.record_fps, cfg.recorder.enable_tts)
+
+    log.info("Initialization done!")
+    log.info("Entering Loop")
+
+    record = False
+
+    while 1:
+        # get input events
+        action = vr_input.get_vr_action()
+        obs, _, _, info = env.step(action)
+        done = False
+        if vr_input.reset_button_pressed:
+            done = True
+        if vr_input.start_button_pressed:
+            record = True
+        if vr_input.reset_button_hold:
+            data_recorder.delete_episode()
+        if record and cfg.recorder.record:
+            data_recorder.step(vr_input.prev_vr_events, obs, done, info)
+        if done:
+            record = False
+            env.reset()
+
+
+if __name__ == "__main__":
+    main()

aloha_robot_project/calvin/calvin_env/conf/cameras/cameras/gripper.yaml

@@ -0,0 +1,8 @@
+_target_: calvin_env.camera.gripper_camera.GripperCamera
+name: gripper
+fov: 75
+aspect: 1
+nearval: 0.01
+farval: 2
+width: 84
+height: 84

aloha_robot_project/calvin/calvin_env/conf/cameras/cameras/opposing.yaml

@@ -0,0 +1,10 @@
+_target_: calvin_env.camera.static_camera.StaticCamera
+name: opposing
+fov: 75
+aspect: 1
+nearval: 0.01
+farval: 2
+width: 200
+height: 200
+look_at: [ 0.4, 0.5, 0.6 ]
+look_from: [ 0.4, 1.5, 0.9 ]

aloha_robot_project/calvin/calvin_env/conf/cameras/cameras/static.yaml

@@ -0,0 +1,11 @@
+_target_: calvin_env.camera.static_camera.StaticCamera
+name: static
+fov: 10
+aspect: 1
+nearval: 0.01
+farval: 10
+width: 200
+height: 200
+look_at: [ -0.026242351159453392, -0.0302329882979393, 0.3920000493526459]
+look_from: [ 2.871459009488717, -2.166602199425597, 2.555159848480571]
+up_vector: [ 0.4041403970338857, 0.22629790978217404, 0.8862616969685161]