macos brew

- Replace `use_pretrained_backbone` with `pretrained_backbone_weights`
- Bump diffusers' minimum version `0.26.3` -> `0.27.2`
- Add ignore flags in CI's pytest
- Change Box observation spaces in simulation environments
- Set `version_base="1.2"` in Hydra initializations
- Bump einops' minimum version `0.7.0` -> `0.8.0`

.dockerignore

@@ -0,0 +1,142 @@
+# Misc
+.git
+tmp
+wandb
+data
+outputs
+.vscode
+rl
+media
+
+
+# Logging
+logs
+
+# HPC
+nautilus/*.yaml
+*.key
+
+# Slurm
+sbatch*.sh
+
+# 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
+!tests/data
+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
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -0,0 +1,54 @@
+name: "\U0001F41B Bug Report"
+description: Submit a bug report to help us improve LeRobot
+body:
+  - type: markdown
+    attributes:
+      value: |
+        Thanks for taking the time to submit a bug report! 🐛
+        If this is not a bug related to the LeRobot library directly, but instead a general question about your code or the library specifically please use our [discord](https://discord.gg/s3KuuzsPFb).
+
+  - type: textarea
+    id: system-info
+    attributes:
+      label: System Info
+      description: If needed, you can share your lerobot configuration with us by running `python -m lerobot.scripts.display_sys_info` and copy-pasting its outputs below
+      render: Shell
+      placeholder: lerobot version, OS, python version, numpy version, torch version, and lerobot's configuration
+    validations:
+      required: true
+
+  - type: checkboxes
+    id: information-scripts-examples
+    attributes:
+      label: Information
+      description: 'The problem arises when using:'
+      options:
+        - label: "One of the scripts in the examples/ folder of LeRobot"
+        - label: "My own task or dataset (give details below)"
+
+  - type: textarea
+    id: reproduction
+    validations:
+      required: true
+    attributes:
+      label: Reproduction
+      description: |
+        If needed, provide a simple code sample that reproduces the problem you ran into. It can be a Colab link or just a code snippet.
+        Sharing error messages or stack traces could be useful as well!
+        Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
+        Try to avoid screenshots, as they are hard to read and don't allow copy-and-pasting.
+
+      placeholder: |
+        Steps to reproduce the behavior:
+
+          1.
+          2.
+          3.
+
+  - type: textarea
+    id: expected-behavior
+    validations:
+      required: true
+    attributes:
+      label: Expected behavior
+      description: "A clear and concise description of what you would expect to happen."

.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,32 @@
+# What does this PR do?
+
+Examples:
+- Fixes # (issue)
+- Adds new dataset
+- Optimizes something
+
+## How was it tested?
+
+Examples:
+- Added `test_something` in `tests/test_stuff.py`.
+- Added `new_feature` and checked that training converges with policy X on dataset/environment Y.
+- Optimized `some_function`, it now runs X times faster than previously.
+
+## How to checkout & try? (for the reviewer)
+
+Examples:
+```bash
+DATA_DIR=tests/data pytest -sx tests/test_stuff.py::test_something
+```
+```bash
+python lerobot/scripts/train.py --some.option=true
+```
+
+## Before submitting
+Please read the [contributor guideline](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md#submitting-a-pull-request-pr).
+
+
+## Who can review?
+
+Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
+members/contributors who may be interested in your PR. Try to avoid tagging more than 3 people.

.github/poetry/cpu/pyproject.toml

@@ -1,109 +0,0 @@
-[tool.poetry]
-name = "lerobot"
-version = "0.1.0"
-description = "Le robot is learning"
-authors = [
-    "Rémi Cadène <re.cadene@gmail.com>",
-    "Simon Alibert <alibert.sim@gmail.com>",
-]
-repository = "https://github.com/Cadene/lerobot"
-readme = "README.md"
-license = "MIT"
-classifiers=[
-    "Development Status :: 3 - Alpha",
-    "Intended Audience :: Developers",
-    "Topic :: Software Development :: Build Tools",
-    "License :: OSI Approved :: MIT License",
-    "Programming Language :: Python :: 3.10",
-]
-packages = [{include = "lerobot"}]
-
-
-[tool.poetry.dependencies]
-python = "^3.10"
-termcolor = "^2.4.0"
-omegaconf = "^2.3.0"
-dm-env = "^1.6"
-pandas = "^2.2.1"
-wandb = "^0.16.3"
-moviepy = "^1.0.3"
-imageio = {extras = ["pyav"], version = "^2.34.0"}
-gdown = "^5.1.0"
-hydra-core = "^1.3.2"
-einops = "^0.7.0"
-pygame = "^2.5.2"
-pymunk = "^6.6.0"
-zarr = "^2.17.0"
-shapely = "^2.0.3"
-scikit-image = "^0.22.0"
-numba = "^0.59.0"
-mpmath = "^1.3.0"
-torch = {version = "^2.2.1", source = "torch-cpu"}
-tensordict = {git = "https://github.com/pytorch/tensordict"}
-torchrl = {git = "https://github.com/pytorch/rl", rev = "13bef426dcfa5887c6e5034a6e9697993fa92c37"}
-mujoco = "^2.3.7"
-opencv-python = "^4.9.0.80"
-diffusers = "^0.26.3"
-torchvision = {version = "^0.17.1", source = "torch-cpu"}
-h5py = "^3.10.0"
-dm = "^1.3"
-dm-control = "1.0.14"
-robomimic = "0.2.0"
-huggingface-hub = "^0.21.4"
-gymnasium-robotics = "^1.2.4"
-gymnasium = "^0.29.1"
-cmake = "^3.29.0.1"
-
-
-[tool.poetry.group.dev.dependencies]
-pre-commit = "^3.6.2"
-debugpy = "^1.8.1"
-pytest = "^8.1.0"
-pytest-cov = "^5.0.0"
-
-
-[[tool.poetry.source]]
-name = "torch-cpu"
-url = "https://download.pytorch.org/whl/cpu"
-priority = "supplemental"
-
-
-[tool.ruff]
-line-length = 110
-target-version = "py310"
-exclude = [
-    ".bzr",
-    ".direnv",
-    ".eggs",
-    ".git",
-    ".git-rewrite",
-    ".hg",
-    ".mypy_cache",
-    ".nox",
-    ".pants.d",
-    ".pytype",
-    ".ruff_cache",
-    ".svn",
-    ".tox",
-    ".venv",
-    "__pypackages__",
-    "_build",
-    "buck-out",
-    "build",
-    "dist",
-    "node_modules",
-    "venv",
-]
-
-
-[tool.ruff.lint]
-select = ["E4", "E7", "E9", "F", "I", "N", "B", "C4", "SIM"]
-
-
-[tool.poetry-dynamic-versioning]
-enable = true
-
-
-[build-system]
-requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning>=1.0.0,<2.0.0"]
-build-backend = "poetry_dynamic_versioning.backend"

.github/scripts/dep_build.py

@@ -0,0 +1,30 @@
+PYPROJECT = "pyproject.toml"
+DEPS = {
+    "gym-pusht": '{ git = "git@github.com:huggingface/gym-pusht.git", optional = true}',
+    "gym-xarm": '{ git = "git@github.com:huggingface/gym-xarm.git", optional = true}',
+    "gym-aloha": '{ git = "git@github.com:huggingface/gym-aloha.git", optional = true}',
+}
+
+
+def update_envs_as_path_dependencies():
+    with open(PYPROJECT) as file:
+        lines = file.readlines()
+
+    new_lines = []
+    for line in lines:
+        if any(dep in line for dep in DEPS.values()):
+            for dep in DEPS:
+                if dep in line:
+                    new_line = f'{dep} = {{ path = "envs/{dep}/", optional = true}}\n'
+                    new_lines.append(new_line)
+                    break
+
+        else:
+            new_lines.append(line)
+
+    with open(PYPROJECT, "w") as file:
+        file.writelines(new_lines)
+
+
+if __name__ == "__main__":
+    update_envs_as_path_dependencies()

.github/workflows/build-docker-images.yml

@@ -0,0 +1,203 @@
+# Inspired by
+# https://github.com/huggingface/peft/blob/main/.github/workflows/build_docker_images.yml
+name: Builds
+
+on:
+  workflow_dispatch:
+  workflow_call:
+  schedule:
+    - cron: "0 1 * * *"
+
+env:
+  PYTHON_VERSION: "3.10"
+#   CI_SLACK_CHANNEL: ${{ secrets.CI_DOCKER_CHANNEL }}
+
+jobs:
+  latest-cpu:
+    name: "Build CPU"
+    runs-on: ubuntu-latest
+    steps:
+      - name: Cleanup disk
+        run: |
+          sudo df -h
+          # sudo ls -l /usr/local/lib/
+          # sudo ls -l /usr/share/
+          sudo du -sh /usr/local/lib/
+          sudo du -sh /usr/share/
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /usr/share/dotnet
+          sudo du -sh /usr/local/lib/
+          sudo du -sh /usr/share/
+          sudo df -h
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Check out code
+        uses: actions/checkout@v4
+
+      # HACK(aliberts): to be removed for release
+      # -----------------------------------------
+      - name: Checkout gym-aloha
+        uses: actions/checkout@v4
+        with:
+          repository: huggingface/gym-aloha
+          path: envs/gym-aloha
+          ssh-key: ${{ secrets.SSH_PRIVATE_KEY }}
+
+      - name: Checkout gym-xarm
+        uses: actions/checkout@v4
+        with:
+          repository: huggingface/gym-xarm
+          path: envs/gym-xarm
+          ssh-key: ${{ secrets.SSH_PRIVATE_KEY }}
+
+      - name: Checkout gym-pusht
+        uses: actions/checkout@v4
+        with:
+          repository: huggingface/gym-pusht
+          path: envs/gym-pusht
+          ssh-key: ${{ secrets.SSH_PRIVATE_KEY }}
+
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Change envs dependencies as local path
+        run: python .github/scripts/dep_build.py
+      # -----------------------------------------
+
+      - name: Login to DockerHub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push CPU
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          file: ./docker/lerobot-cpu/Dockerfile
+          push: true
+          tags: huggingface/lerobot-cpu
+          build-args: PYTHON_VERSION=${{ env.PYTHON_VERSION }}
+
+    #   - name: Post to a Slack channel
+    #     id: slack
+    #     #uses: slackapi/slack-github-action@v1.25.0
+    #     uses: slackapi/slack-github-action@6c661ce58804a1a20f6dc5fbee7f0381b469e001
+    #     with:
+    #       # Slack channel id, channel name, or user id to post message.
+    #       # See also: https://api.slack.com/methods/chat.postMessage#channels
+    #       channel-id: ${{ env.CI_SLACK_CHANNEL }}
+    #       # For posting a rich message using Block Kit
+    #       payload: |
+    #         {
+    #           "text": "lerobot-cpu Docker Image build result: ${{ job.status }}\n${{ github.event.pull_request.html_url || github.event.head_commit.url }}",
+    #           "blocks": [
+    #             {
+    #               "type": "section",
+    #               "text": {
+    #                 "type": "mrkdwn",
+    #                 "text": "lerobot-cpu Docker Image build result: ${{ job.status }}\n${{ github.event.pull_request.html_url || github.event.head_commit.url }}"
+    #               }
+    #             }
+    #           ]
+    #         }
+    #     env:
+    #       SLACK_BOT_TOKEN: ${{ secrets.SLACK_CIFEEDBACK_BOT_TOKEN }}
+
+  latest-cuda:
+    name: "Build GPU"
+    runs-on: ubuntu-latest
+    steps:
+      - name: Cleanup disk
+        run: |
+          sudo df -h
+          # sudo ls -l /usr/local/lib/
+          # sudo ls -l /usr/share/
+          sudo du -sh /usr/local/lib/
+          sudo du -sh /usr/share/
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /usr/share/dotnet
+          sudo du -sh /usr/local/lib/
+          sudo du -sh /usr/share/
+          sudo df -h
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Check out code
+        uses: actions/checkout@v4
+
+      # HACK(aliberts): to be removed for release
+      # -----------------------------------------
+      - name: Checkout gym-aloha
+        uses: actions/checkout@v4
+        with:
+          repository: huggingface/gym-aloha
+          path: envs/gym-aloha
+          ssh-key: ${{ secrets.SSH_PRIVATE_KEY }}
+
+      - name: Checkout gym-xarm
+        uses: actions/checkout@v4
+        with:
+          repository: huggingface/gym-xarm
+          path: envs/gym-xarm
+          ssh-key: ${{ secrets.SSH_PRIVATE_KEY }}
+
+      - name: Checkout gym-pusht
+        uses: actions/checkout@v4
+        with:
+          repository: huggingface/gym-pusht
+          path: envs/gym-pusht
+          ssh-key: ${{ secrets.SSH_PRIVATE_KEY }}
+
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Change envs dependencies as local path
+        run: python .github/scripts/dep_build.py
+      # -----------------------------------------
+
+      - name: Login to DockerHub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push GPU
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          file: ./docker/lerobot-gpu/Dockerfile
+          push: true
+          tags: huggingface/lerobot-gpu
+          build-args: PYTHON_VERSION=${{ env.PYTHON_VERSION }}
+
+      # - name: Post to a Slack channel
+      #   id: slack
+      #   #uses: slackapi/slack-github-action@v1.25.0
+      #   uses: slackapi/slack-github-action@6c661ce58804a1a20f6dc5fbee7f0381b469e001
+      #   with:
+      #     # Slack channel id, channel name, or user id to post message.
+      #     # See also: https://api.slack.com/methods/chat.postMessage#channels
+      #     channel-id: ${{ env.CI_SLACK_CHANNEL }}
+      #     # For posting a rich message using Block Kit
+      #     payload: |
+      #       {
+      #         "text": "lerobot-gpu Docker Image build result: ${{ job.status }}\n${{ github.event.pull_request.html_url || github.event.head_commit.url }}",
+      #         "blocks": [
+      #           {
+      #             "type": "section",
+      #             "text": {
+      #               "type": "mrkdwn",
+      #               "text": "lerobot-gpu Docker Image build result: ${{ job.status }}\n${{ github.event.pull_request.html_url || github.event.head_commit.url }}"
+      #             }
+      #           }
+      #         ]
+      #       }
+      #   env:
+      #     SLACK_BOT_TOKEN: ${{ secrets.SLACK_CIFEEDBACK_BOT_TOKEN }}

.github/workflows/nightly-tests.yml

@@ -0,0 +1,79 @@
+# Inspired by
+# https://github.com/huggingface/peft/blob/main/.github/workflows/nightly.yml
+name: Nightly
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 2 * * *"
+
+env:
+  DATA_DIR: tests/data
+  # SLACK_API_TOKEN: ${{ secrets.SLACK_API_TOKEN }}
+
+jobs:
+  run_all_tests_cpu:
+    name: "Test CPU"
+    strategy:
+      fail-fast: false
+    runs-on: ubuntu-latest
+    container:
+      image: huggingface/lerobot-cpu:latest
+      options: --shm-size "16gb"
+      credentials:
+        username: ${{ secrets.DOCKERHUB_USERNAME }}
+        password: ${{ secrets.DOCKERHUB_PASSWORD }}
+    defaults:
+      run:
+        shell: bash
+        working-directory: /lerobot
+    steps:
+      - name: Tests
+        env:
+          DATA_DIR: tests/data
+        run: pytest -v --cov=./lerobot --disable-warnings tests
+
+      - name: Tests end-to-end
+        env:
+          DATA_DIR: tests/data
+        run: make test-end-to-end
+
+
+  run_all_tests_single_gpu:
+    name: "Test GPU"
+    strategy:
+      fail-fast: false
+    runs-on: [single-gpu, nvidia-gpu, t4, ci]
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+      TEST_TYPE: "single_gpu"
+    container:
+      image: huggingface/lerobot-gpu:latest
+      options: --gpus all --shm-size "16gb"
+      credentials:
+        username: ${{ secrets.DOCKERHUB_USERNAME }}
+        password: ${{ secrets.DOCKERHUB_PASSWORD }}
+    defaults:
+      run:
+        shell: bash
+        working-directory: /lerobot
+    steps:
+      - name: Nvidia-smi
+        run: nvidia-smi
+
+      - name: Test
+        run: pytest -v --cov=./lerobot --cov-report=xml --disable-warnings tests
+      #   TODO(aliberts): Link with HF Codecov account
+      # - name: Upload coverage reports to Codecov with GitHub Action
+      #   uses: codecov/codecov-action@v4
+      #   with:
+      #     files: ./coverage.xml
+      #     verbose: true
+      - name: Tests end-to-end
+        run: make test-end-to-end
+
+    #   - name: Generate Report
+    #     if: always()
+    #     run: |
+    #       pip install slack_sdk tabulate
+    #       python scripts/log_reports.py >> $GITHUB_STEP_SUMMARY

.github/workflows/style.yml

@@ -0,0 +1,38 @@
+name: Style
+
+on:
+  workflow_dispatch:
+  workflow_call:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
+
+env:
+  PYTHON_VERSION: "3.10"
+
+jobs:
+  ruff_check:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ env.PYTHON_VERSION }}
+
+      - name: Get Ruff Version from pre-commit-config.yaml
+        id: get-ruff-version
+        run: |
+          RUFF_VERSION=$(awk '/repo: https:\/\/github.com\/astral-sh\/ruff-pre-commit/{flag=1;next}/rev:/{if(flag){print $2;exit}}' .pre-commit-config.yaml)
+          echo "RUFF_VERSION=${RUFF_VERSION}" >> $GITHUB_ENV
+
+      - name: Install Ruff
+        run: python -m pip install "ruff==${{ env.RUFF_VERSION }}"
+
+      - name: Run Ruff
+        run: ruff check .

.github/workflows/test-docker-build.yml

@@ -0,0 +1,109 @@
+# Inspired by
+# https://github.com/huggingface/peft/blob/main/.github/workflows/test-docker-build.yml
+name: Test Docker builds (PR)
+
+on:
+  pull_request:
+    branches:
+      - main
+    paths:
+      # Run only when DockerFile files are modified
+      - "docker/**"
+
+env:
+  PYTHON_VERSION: "3.10"
+
+jobs:
+  get_changed_files:
+    name: "Get all modified Dockerfiles"
+    runs-on: ubuntu-latest
+    outputs:
+      matrix: ${{ steps.set-matrix.outputs.matrix }}
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+
+      - name: Get changed files
+        id: changed-files
+        uses: tj-actions/changed-files@v44
+        with:
+          files: docker/**
+          json: "true"
+
+      - name: Run step if only the files listed above change
+        if: steps.changed-files.outputs.any_changed == 'true'
+        id: set-matrix
+        env:
+          ALL_CHANGED_FILES: ${{ steps.changed-files.outputs.all_changed_files }}
+        run: |
+          echo "matrix=${{ steps.changed-files.outputs.all_changed_files}}" >> $GITHUB_OUTPUT
+
+
+  build_modified_dockerfiles:
+    name: "Build all modified Docker images"
+    needs: get_changed_files
+    runs-on: ubuntu-latest
+    if: ${{ needs.get_changed_files.outputs.matrix }} != ''
+    strategy:
+      fail-fast: false
+      matrix:
+        docker-file: ${{ fromJson(needs.get_changed_files.outputs.matrix) }}
+    steps:
+      - name: Cleanup disk
+        run: |
+          sudo df -h
+          # sudo ls -l /usr/local/lib/
+          # sudo ls -l /usr/share/
+          sudo du -sh /usr/local/lib/
+          sudo du -sh /usr/share/
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /usr/share/dotnet
+          sudo du -sh /usr/local/lib/
+          sudo du -sh /usr/share/
+          sudo df -h
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Check out code
+        uses: actions/checkout@v4
+
+      # HACK(aliberts): to be removed for release
+      # -----------------------------------------
+      - name: Checkout gym-aloha
+        uses: actions/checkout@v4
+        with:
+          repository: huggingface/gym-aloha
+          path: envs/gym-aloha
+          ssh-key: ${{ secrets.SSH_PRIVATE_KEY }}
+
+      - name: Checkout gym-xarm
+        uses: actions/checkout@v4
+        with:
+          repository: huggingface/gym-xarm
+          path: envs/gym-xarm
+          ssh-key: ${{ secrets.SSH_PRIVATE_KEY }}
+
+      - name: Checkout gym-pusht
+        uses: actions/checkout@v4
+        with:
+          repository: huggingface/gym-pusht
+          path: envs/gym-pusht
+          ssh-key: ${{ secrets.SSH_PRIVATE_KEY }}
+
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Change envs dependencies as local path
+        run: python .github/scripts/dep_build.py
+      # -----------------------------------------
+
+      - name: Build Docker image
+        uses: docker/build-push-action@v5
+        with:
+          file: ${{ matrix.docker-file }}
+          context: .
+          push: False
+          build-args: PYTHON_VERSION=${{ env.PYTHON_VERSION }}

.github/workflows/test.yml

@@ -4,231 +4,71 @@ on:
   pull_request:
     branches:
       - main
-    types: [opened, synchronize, reopened, labeled]
+    paths:
+      - "lerobot/**"
+      - "tests/**"
+      - "examples/**"
+      - ".github/**"
+      - "poetry.lock"
   push:
     branches:
       - main
+    paths:
+      - "lerobot/**"
+      - "tests/**"
+      - "examples/**"
+      - ".github/**"
+      - "poetry.lock"
 
 jobs:
   tests:
-    if: |
-      ${{ github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'CI') }} ||
-      ${{ github.event_name == 'push' }}
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, macos-latest, macos-latest-large]
     env:
-      POETRY_VERSION: 1.8.2
       DATA_DIR: tests/data
       MUJOCO_GL: egl
     steps:
-      #----------------------------------------------
-      #       check-out repo and set-up python
-      #----------------------------------------------
-      - name: Check out repository
-        uses: actions/checkout@v4
+      - name: Add SSH key for installing envs
+        uses: webfactory/ssh-agent@v0.9.0
         with:
-          lfs: true
+          ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
 
-      - name: Set up python
-        id: setup-python
+      - uses: actions/checkout@v4
+
+      - name: Install EGL
+        run: |
+          if [[ "${{ matrix.os }}" == 'ubuntu-latest' ]]; then
+            sudo apt-get update && sudo apt-get install -y libegl1-mesa-dev
+          elif [[ "${{ matrix.os }}" == 'macos-latest' || "${{ matrix.os }}" == 'macos-latest-large' ]]; then
+            brew install mesa
+          fi
+
+      - name: Install poetry
+        run: |
+          pipx install poetry && poetry config virtualenvs.in-project true
+          echo "${{ github.workspace }}/.venv/bin" >> $GITHUB_PATH
+
+      - name: Set up Python 3.10
         uses: actions/setup-python@v5
         with:
-          python-version: '3.10'
+          python-version: "3.10"
+          cache: "poetry"
 
-      #----------------------------------------------
-      #         install & configure poetry
-      #----------------------------------------------
-      - name: Load cached Poetry installation
-        id: restore-poetry-cache
-        uses: actions/cache/restore@v3
-        with:
-          path: ~/.local
-          key: poetry-${{ env.POETRY_VERSION }}
-
-      - name: Install Poetry
-        if: steps.restore-poetry-cache.outputs.cache-hit != 'true'
-        uses: snok/install-poetry@v1
-        with:
-          version: ${{ env.POETRY_VERSION }}
-          virtualenvs-create: true
-          installer-parallel: true
-
-      - name: Save cached Poetry installation
-        if: |
-          steps.restore-poetry-cache.outputs.cache-hit != 'true' &&
-          github.ref_name == 'main'
-        id: save-poetry-cache
-        uses: actions/cache/save@v3
-        with:
-          path: ~/.local
-          key: poetry-${{ env.POETRY_VERSION }}
-
-      - name: Configure Poetry
-        run: poetry config virtualenvs.in-project true
-
-      #----------------------------------------------
-      #           install dependencies
-      #----------------------------------------------
-      # TODO(aliberts): move to gpu runners
-      - name: Select cpu dependencies  # HACK
-        run: cp -t . .github/poetry/cpu/pyproject.toml .github/poetry/cpu/poetry.lock
-
-      - name: Load cached venv
-        id: restore-dependencies-cache
-        uses: actions/cache/restore@v3
-        with:
-          path: .venv
-          key: venv-${{ steps.setup-python.outputs.python-version }}-${{ env.POETRY_VERSION }}-${{ hashFiles('**/poetry.lock') }}
-
-      - name: Install dependencies
-        if: steps.restore-dependencies-cache.outputs.cache-hit != 'true'
-        env:
-          TMPDIR: ~/tmp
-          TEMP: ~/tmp
-          TMP: ~/tmp
+      - name: Install poetry dependencies
         run: |
-          mkdir ~/tmp
-          poetry install --no-interaction --no-root
+          poetry install --all-extras
 
-      - name: Save cached venv
-        if: |
-            steps.restore-dependencies-cache.outputs.cache-hit != 'true' &&
-            github.ref_name == 'main'
-        id: save-dependencies-cache
-        uses: actions/cache/save@v3
-        with:
-          path: .venv
-          key: venv-${{ steps.setup-python.outputs.python-version }}-${{ env.POETRY_VERSION }}-${{ hashFiles('**/poetry.lock') }}
-
-      - name: Install libegl1-mesa-dev (to use MUJOCO_GL=egl)
-        run: sudo apt-get update && sudo apt-get install -y libegl1-mesa-dev
-
-      #----------------------------------------------
-      #             install project
-      #----------------------------------------------
-      - name: Install project
-        run: poetry install --no-interaction
-
-      #----------------------------------------------
-      #            run tests & coverage
-      #----------------------------------------------
-      - name: Run tests
-        env:
-          LEROBOT_TESTS_DEVICE: cpu
+      - name: Test with pytest
         run: |
-          source .venv/bin/activate
-          pytest --cov=./lerobot --cov-report=xml tests
+          pytest tests -v --cov=./lerobot --durations=0 \
+            -W ignore::DeprecationWarning:imageio_ffmpeg._utils:7 \
+            -W ignore::UserWarning:torch.utils.data.dataloader:558 \
+            -W ignore::UserWarning:gymnasium.utils.env_checker:247 \
+            && rm -rf tests/outputs outputs
 
-      #   TODO(aliberts): Link with HF Codecov account
-      # - name: Upload coverage reports to Codecov with GitHub Action
-      #   uses: codecov/codecov-action@v4
-      #   with:
-      #     files: ./coverage.xml
-      #     verbose: true
-
-      #----------------------------------------------
-      #            run end-to-end tests
-      #----------------------------------------------
-      - name: Test train ACT on ALOHA end-to-end
+      - name: Test end-to-end
         run: |
-          source .venv/bin/activate
-          python lerobot/scripts/train.py \
-            policy=act \
-            env=aloha \
-            wandb.enable=False \
-            offline_steps=2 \
-            online_steps=0 \
-            device=cpu \
-            save_model=true \
-            save_freq=2 \
-            horizon=20 \
-            policy.batch_size=2 \
-            hydra.run.dir=tests/outputs/act/
-
-      - name: Test eval ACT on ALOHA end-to-end
-        run: |
-          source .venv/bin/activate
-          python lerobot/scripts/eval.py \
-            --config tests/outputs/act/.hydra/config.yaml \
-            eval_episodes=1 \
-            env.episode_length=8 \
-            device=cpu \
-            policy.pretrained_model_path=tests/outputs/act/models/2.pt
-
-      # TODO(aliberts): This takes ~2mn to run, needs to be improved
-      # - name: Test eval ACT on ALOHA end-to-end (policy is None)
-      #   run: |
-      #     source .venv/bin/activate
-      #     python lerobot/scripts/eval.py \
-      #       --config lerobot/configs/default.yaml \
-      #       policy=act \
-      #       env=aloha \
-      #       eval_episodes=1 \
-      #       device=cpu
-
-      - name: Test train Diffusion on PushT end-to-end
-        run: |
-          source .venv/bin/activate
-          python lerobot/scripts/train.py \
-            policy=diffusion \
-            env=pusht \
-            wandb.enable=False \
-            offline_steps=2 \
-            online_steps=0 \
-            device=cpu \
-            save_model=true \
-            save_freq=2 \
-            hydra.run.dir=tests/outputs/diffusion/
-
-      - name: Test eval Diffusion on PushT end-to-end
-        run: |
-          source .venv/bin/activate
-          python lerobot/scripts/eval.py \
-            --config tests/outputs/diffusion/.hydra/config.yaml \
-            eval_episodes=1 \
-            env.episode_length=8 \
-            device=cpu \
-            policy.pretrained_model_path=tests/outputs/diffusion/models/2.pt
-
-      - name: Test eval Diffusion on PushT end-to-end (policy is None)
-        run: |
-          source .venv/bin/activate
-          python lerobot/scripts/eval.py \
-            --config lerobot/configs/default.yaml \
-            policy=diffusion  \
-            env=pusht \
-            eval_episodes=1 \
-            device=cpu
-
-      - name: Test train TDMPC on Simxarm end-to-end
-        run: |
-          source .venv/bin/activate
-          python lerobot/scripts/train.py \
-            policy=tdmpc \
-            env=simxarm \
-            wandb.enable=False \
-            offline_steps=1 \
-            online_steps=1 \
-            device=cpu \
-            save_model=true \
-            save_freq=2 \
-            hydra.run.dir=tests/outputs/tdmpc/
-
-      - name: Test eval TDMPC on Simxarm end-to-end
-        run: |
-          source .venv/bin/activate
-          python lerobot/scripts/eval.py \
-            --config tests/outputs/tdmpc/.hydra/config.yaml \
-            eval_episodes=1 \
-            env.episode_length=8 \
-            device=cpu \
-            policy.pretrained_model_path=tests/outputs/tdmpc/models/2.pt
-
-      - name: Test eval TDPMC on Simxarm end-to-end (policy is None)
-        run: |
-          source .venv/bin/activate
-          python lerobot/scripts/eval.py \
-            --config lerobot/configs/default.yaml \
-            policy=tdmpc  \
-            env=simxarm \
-            eval_episodes=1 \
-            device=cpu
+          make test-end-to-end \
+            && rm -rf outputs

.gitignore

@@ -11,6 +11,9 @@ rl
 nautilus/*.yaml
 *.key
 
+# Slurm
+sbatch*.sh
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

.pre-commit-config.yaml

@@ -1,9 +1,9 @@
-exclude: ^(data/|tests/)
+exclude: ^(tests/data)
 default_language_version:
     python: python3.10
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v4.6.0
     hooks:
       - id: check-added-large-files
       - id: debug-statements
@@ -18,7 +18,7 @@ repos:
     hooks:
     -   id: pyupgrade
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.3.4
+    rev: v0.4.2
     hooks:
       - id: ruff
         args: [--fix]

CODE_OF_CONDUCT.md

@@ -0,0 +1,133 @@
+
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, caste, color, religion, or sexual
+identity and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the overall
+  community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or advances of
+  any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email address,
+  without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official email address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at
+[feedback@huggingface.co](mailto:feedback@huggingface.co).
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series of
+actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or permanent
+ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior, harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within the
+community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.1, available at
+[https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1].
+
+Community Impact Guidelines were inspired by
+[Mozilla's code of conduct enforcement ladder][Mozilla CoC].
+
+For answers to common questions about this code of conduct, see the FAQ at
+[https://www.contributor-covenant.org/faq][FAQ]. Translations are available at
+[https://www.contributor-covenant.org/translations][translations].
+
+[homepage]: https://www.contributor-covenant.org
+[v2.1]: https://www.contributor-covenant.org/version/2/1/code_of_conduct.html
+[Mozilla CoC]: https://github.com/mozilla/diversity
+[FAQ]: https://www.contributor-covenant.org/faq
+[translations]: https://www.contributor-covenant.org/translations

CONTRIBUTING.md

@@ -0,0 +1,270 @@
+# How to contribute to 🤗 LeRobot?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter when it has
+helped you, or simply ⭐️ the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/lerobot/blob/main/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+Some of the ways you can contribute to 🤗 LeRobot:
+* Fixing outstanding issues with the existing code.
+* Implementing new models, datasets or simulation environments.
+* Contributing to the examples or to the documentation.
+* Submitting issues related to bugs or desired new features.
+
+Following the guides below, feel free to open issues and PRs and to coordinate your efforts with the community on our [Discord Channel](https://discord.gg/VjFz58wn3R). For specific inquiries, reach out to [Remi Cadene](remi.cadene@huggingface.co).
+
+If you are not sure how to contribute or want to know the next features we working on, look on this project page: [LeRobot TODO](https://github.com/orgs/huggingface/projects/46)
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 LeRobot library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s.
+* The full traceback if an exception is raised.
+* Attach any other additional information, like screenshots, you think may help.
+
+### Do you want a new feature?
+
+A good feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *paragraph* describing the feature.
+3. Provide a **code snippet** that demonstrates its future use.
+4. In case this is related to a paper, please attach a link.
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Adding new policies, datasets or environments
+
+Look at our implementations for [datasets](./lerobot/common/datasets/), [policies](./lerobot/common/policies/),
+environments ([aloha](https://github.com/huggingface/gym-aloha),
+[xarm](https://github.com/huggingface/gym-xarm),
+[pusht](https://github.com/huggingface/gym-pusht))
+and follow the same api design.
+
+When implementing a new dataset loadable with LeRobotDataset follow these steps:
+- Update `available_datasets_per_env` in `lerobot/__init__.py`
+
+When implementing a new environment (e.g. `gym_aloha`), follow these steps:
+- Update `available_tasks_per_env` and `available_datasets_per_env` in `lerobot/__init__.py`
+
+When implementing a new policy class (e.g. `DiffusionPolicy`) follow these steps:
+- Update `available_policies` and `available_policies_per_env`, in `lerobot/__init__.py`
+- Set the required `name` class attribute.
+- Update variables in `tests/test_available.py` by importing your new Policy class
+
+## Submitting a pull request (PR)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+🤗 LeRobot. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/lerobot) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote. The following command
+   assumes you have your public SSH key uploaded to GitHub. See the following guide for more
+   [information](https://docs.github.com/en/repositories/creating-and-managing-repositories/cloning-a-repository).
+
+   ```bash
+   git clone git@github.com:<your Github handle>/lerobot.git
+   cd lerobot
+   git remote add upstream https://github.com/huggingface/lerobot.git
+   ```
+
+3. Create a new branch to hold your development changes, and do this for every new PR you work on.
+
+   Start by synchronizing your `main` branch with the `upstream/main` branch (more details in the [GitHub Docs](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/syncing-a-fork)):
+
+   ```bash
+   git checkout main
+   git fetch upstream
+   git rebase upstream/main
+   ```
+
+   Once your `main` branch is synchronized, create a new branch from it:
+
+   ```bash
+   git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   🚨 **Do not** work on the `main` branch.
+
+4. for development, we use `poetry` instead of just `pip` to easily track our dependencies.
+   If you don't have it already, follow the [instructions](https://python-poetry.org/docs/#installation) to install it.
+
+   Set up a development environment with conda or miniconda:
+   ```bash
+   conda create -y -n lerobot-dev python=3.10 && conda activate lerobot-dev
+   ```
+
+   To develop on 🤗 LeRobot, you will at least need to install the `dev` and `test` extras dependencies along with the core library:
+   ```bash
+   poetry install --sync --extras "dev test"
+   ```
+
+   You can also install the project with all its dependencies (including environments):
+   ```bash
+   poetry install --sync --all-extras
+   ```
+
+   > **Note:** If you don't install simulation environments with `--all-extras`, the tests that require them will be skipped when running the pytest suite locally. However, they *will* be tested in the CI. In general, we advise you to install everything and test locally before pushing.
+
+   Whichever command you chose to install the project (e.g. `poetry install --sync --all-extras`), you should run it again when pulling code with an updated version of `pyproject.toml` and `poetry.lock` in order to synchronize your virtual environment with the new dependencies.
+
+   The equivalent of `pip install some-package`, would just be:
+   ```bash
+   poetry add some-package
+   ```
+
+   When making changes to the poetry sections of the `pyproject.toml`, you should run the following command to lock dependencies.
+   ```bash
+   poetry lock --no-update
+   ```
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see
+   below an explanation regarding the environment variable):
+
+   ```bash
+   pytest tests/<TEST_TO_RUN>.py
+   ```
+
+6. Follow our style.
+
+   `lerobot` relies on `ruff` to format its source code
+   consistently. Set up [`pre-commit`](https://pre-commit.com/) to run these checks
+   automatically as Git commit hooks.
+
+   Install `pre-commit` hooks:
+   ```bash
+   pre-commit install
+   ```
+
+   You can run these hooks whenever you need on staged files with:
+   ```bash
+   pre-commit
+   ```
+
+   Once you're happy with your changes, add changed files using `git add` and
+   make a commit with `git commit` to record your changes locally:
+
+   ```bash
+   git add modified_file.py
+   git commit
+   ```
+
+   Please write [good commit messages](https://chris.beams.io/posts/git-commit/).
+
+   It is a good idea to sync your copy of the code with the original
+   repository regularly. This way you can quickly account for changes:
+
+   ```bash
+   git fetch upstream
+   git rebase upstream/main
+   ```
+
+   Push the changes to your account using:
+
+   ```bash
+   git push -u origin a-descriptive-name-for-my-changes
+   ```
+
+6. Once you are satisfied (**and the checklist below is happy too**), go to the
+   webpage of your fork on GitHub. Click on 'Pull request' to send your changes
+   to the project maintainers for review.
+
+7. It's ok if maintainers ask you for changes. It happens to core contributors
+   too! So everyone can see the changes in the Pull request, work in your local
+   branch and push the changes to your fork. They will automatically appear in
+   the pull request.
+
+
+### Checklist
+
+1. The title of your pull request should be a summary of its contribution;
+2. If your pull request addresses an issue, please mention the issue number in
+   the pull request description to make sure they are linked (and people
+   consulting the issue know you are working on it);
+3. To indicate a work in progress please prefix the title with `[WIP]`, or preferably mark
+   the PR as a draft PR. These are useful to avoid duplicated work, and to differentiate
+   it from PRs ready to be merged;
+4. Make sure existing tests pass;
+<!-- 5. Add high-coverage tests. No quality testing = no merge.
+
+See an example of a good PR here: https://github.com/huggingface/lerobot/pull/ -->
+
+### Tests
+
+An extensive test suite is included to test the library behavior and several examples. Library tests can be found in the [tests folder](https://github.com/huggingface/lerobot/tree/main/tests).
+
+Install [git lfs](https://git-lfs.com/) to retrieve test artifacts (if you don't have it already).
+
+On Mac:
+```bash
+brew install git-lfs
+git lfs install
+```
+
+On Ubuntu:
+```bash
+sudo apt-get install git-lfs
+git lfs install
+```
+
+Pull artifacts if they're not in [tests/data](tests/data)
+```bash
+git lfs pull
+```
+
+We use `pytest` in order to run the tests. From the root of the
+repository, here's how to run tests with `pytest` for the library:
+
+```bash
+DATA_DIR="tests/data" python -m pytest -sv ./tests
+```
+
+
+You can specify a smaller set of tests in order to test only the feature
+you're working on.

Makefile

@@ -0,0 +1,95 @@
+.PHONY: tests
+
+PYTHON_PATH := $(shell which python)
+
+# If Poetry is installed, redefine PYTHON_PATH to use the Poetry-managed Python
+POETRY_CHECK := $(shell command -v poetry)
+ifneq ($(POETRY_CHECK),)
+    PYTHON_PATH := $(shell poetry run which python)
+endif
+
+export PATH := $(dir $(PYTHON_PATH)):$(PATH)
+
+
+build-cpu:
+	docker build -t lerobot:latest -f docker/lerobot-cpu/Dockerfile .
+
+build-gpu:
+	docker build -t lerobot:latest -f docker/lerobot-gpu/Dockerfile .
+
+test-end-to-end:
+	${MAKE} test-act-ete-train
+	${MAKE} test-act-ete-eval
+	${MAKE} test-diffusion-ete-train
+	${MAKE} test-diffusion-ete-eval
+	${MAKE} test-tdmpc-ete-train
+	${MAKE} test-tdmpc-ete-eval
+
+test-act-ete-train:
+	python lerobot/scripts/train.py \
+		policy=act \
+		env=aloha \
+		wandb.enable=False \
+		offline_steps=2 \
+		online_steps=0 \
+		eval_episodes=1 \
+		device=cpu \
+		save_model=true \
+		save_freq=2 \
+		policy.n_action_steps=20 \
+		policy.chunk_size=20 \
+		policy.batch_size=2 \
+		hydra.run.dir=tests/outputs/act/
+
+test-act-ete-eval:
+	python lerobot/scripts/eval.py \
+		--config tests/outputs/act/.hydra/config.yaml \
+		eval_episodes=1 \
+		env.episode_length=8 \
+		device=cpu \
+		policy.pretrained_model_path=tests/outputs/act/models/2.pt
+
+test-diffusion-ete-train:
+	python lerobot/scripts/train.py \
+		policy=diffusion \
+		env=pusht \
+		wandb.enable=False \
+		offline_steps=2 \
+		online_steps=0 \
+		eval_episodes=1 \
+		device=cpu \
+		save_model=true \
+		save_freq=2 \
+		policy.batch_size=2 \
+		hydra.run.dir=tests/outputs/diffusion/
+
+test-diffusion-ete-eval:
+	python lerobot/scripts/eval.py \
+		--config tests/outputs/diffusion/.hydra/config.yaml \
+		eval_episodes=1 \
+		env.episode_length=8 \
+		device=cpu \
+		policy.pretrained_model_path=tests/outputs/diffusion/models/2.pt
+
+test-tdmpc-ete-train:
+	python lerobot/scripts/train.py \
+		policy=tdmpc \
+		env=xarm \
+		wandb.enable=False \
+		offline_steps=1 \
+		online_steps=2 \
+		eval_episodes=1 \
+		env.episode_length=2 \
+		device=cpu \
+		save_model=true \
+		save_freq=2 \
+		policy.batch_size=2 \
+		hydra.run.dir=tests/outputs/tdmpc/
+
+test-tdmpc-ete-eval:
+	python lerobot/scripts/eval.py \
+		--config tests/outputs/tdmpc/.hydra/config.yaml \
+		eval_episodes=1 \
+		env.episode_length=8 \
+		device=cpu \
+		policy.pretrained_model_path=tests/outputs/tdmpc/models/2.pt

README.md

@@ -10,13 +10,14 @@
 
 <div align="center">
 
-[![Tests](https://github.com/huggingface/lerobot/actions/workflows/test.yml/badge.svg?branch=main)](https://github.com/huggingface/lerobot/actions/workflows/test.yml?query=branch%3Amain)
+[![Tests](https://github.com/huggingface/lerobot/actions/workflows/nightly-tests.yml/badge.svg?branch=main)](https://github.com/huggingface/lerobot/actions/workflows/nightly-tests.yml?query=branch%3Amain)
 [![Coverage](https://codecov.io/gh/huggingface/lerobot/branch/main/graph/badge.svg?token=TODO)](https://codecov.io/gh/huggingface/lerobot)
 [![Python versions](https://img.shields.io/pypi/pyversions/lerobot)](https://www.python.org/downloads/)
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/huggingface/lerobot/blob/main/LICENSE)
 [![Status](https://img.shields.io/pypi/status/lerobot)](https://pypi.org/project/lerobot/)
 [![Version](https://img.shields.io/pypi/v/lerobot)](https://pypi.org/project/lerobot/)
 [![Examples](https://img.shields.io/badge/Examples-green.svg)](https://github.com/huggingface/lerobot/tree/main/examples)
+[![Contributor Covenant](https://img.shields.io/badge/Contributor%20Covenant-v2.1%20adopted-ff69b4.svg)](https://github.com/huggingface/lerobot/blob/main/CODE_OF_CONDUCT.md)
 [![Discord](https://dcbadge.vercel.app/api/server/C5P34WJ68S?style=flat)](https://discord.gg/s3KuuzsPFb)
 
 </div>
@@ -62,19 +63,27 @@
 
 Download our source code:
 ```bash
-git clone https://github.com/huggingface/lerobot.git
-cd lerobot
+git clone https://github.com/huggingface/lerobot.git && cd lerobot
 ```
 
 Create a virtual environment with Python 3.10 and activate it, e.g. with [`miniconda`](https://docs.anaconda.com/free/miniconda/index.html):
 ```bash
-conda create -y -n lerobot python=3.10
-conda activate lerobot
+conda create -y -n lerobot python=3.10 && conda activate lerobot
 ```
 
-Then, install 🤗 LeRobot:
+Install 🤗 LeRobot:
 ```bash
-python -m pip install .
+pip install .
+```
+
+For simulations, 🤗 LeRobot comes with gymnasium environments that can be installed as extras:
+- [aloha](https://github.com/huggingface/gym-aloha)
+- [xarm](https://github.com/huggingface/gym-xarm)
+- [pusht](https://github.com/huggingface/gym-pusht)
+
+For instance, to install 🤗 LeRobot with aloha and pusht, use:
+```bash
+pip install ".[aloha, pusht]"
 ```
 
 To use [Weights and Biases](https://docs.wandb.ai/quickstart) for experiments tracking, log in with
@@ -89,11 +98,11 @@ wandb login
 ├── lerobot
 |   ├── configs          # contains hydra yaml files with all options that you can override in the command line
 |   |   ├── default.yaml   # selected by default, it loads pusht environment and diffusion policy
-|   |   ├── env            # various sim environments and their datasets: aloha.yaml, pusht.yaml, simxarm.yaml
+|   |   ├── env            # various sim environments and their datasets: aloha.yaml, pusht.yaml, xarm.yaml
 |   |   └── policy         # various policies: act.yaml, diffusion.yaml, tdmpc.yaml
 |   ├── common           # contains classes and utilities
-|   |   ├── datasets       # various datasets of human demonstrations: aloha, pusht, simxarm
-|   |   ├── envs           # various sim environments: aloha, pusht, simxarm
+|   |   ├── datasets       # various datasets of human demonstrations: aloha, pusht, xarm
+|   |   ├── envs           # various sim environments: aloha, pusht, xarm
 |   |   └── policies       # various policies: act, diffusion, tdmpc
 |   └── scripts                  # contains functions to execute via command line
 |       ├── visualize_dataset.py  # load a dataset and render its demonstrations
@@ -109,44 +118,19 @@ wandb login
 
 ### Visualize datasets
 
-You can import our dataset class, download the data from the HuggingFace hub and use our rendering utilities:
-```python
-""" Copy pasted from `examples/1_visualize_dataset.py` """
-import lerobot
-from lerobot.common.datasets.aloha import AlohaDataset
-from torchrl.data.replay_buffers import SamplerWithoutReplacement
-from lerobot.scripts.visualize_dataset import render_dataset
-
-print(lerobot.available_datasets)
-# >>> ['aloha_sim_insertion_human', 'aloha_sim_insertion_scripted', 'aloha_sim_transfer_cube_human', 'aloha_sim_transfer_cube_scripted', 'pusht', 'xarm_lift_medium']
-
-# we use this sampler to sample 1 frame after the other
-sampler = SamplerWithoutReplacement(shuffle=False)
-
-dataset = AlohaDataset("aloha_sim_transfer_cube_human", sampler=sampler)
-
-video_paths = render_dataset(
-    dataset,
-    out_dir="outputs/visualize_dataset/example",
-    max_num_samples=300,
-    fps=50,
-)
-print(video_paths)
-# >>> ['outputs/visualize_dataset/example/episode_0.mp4']
-```
+Check out [examples](./examples) to see how you can import our dataset class, download the data from the HuggingFace hub and use our rendering utilities.
 
 Or you can achieve the same result by executing our script from the command line:
 ```bash
 python lerobot/scripts/visualize_dataset.py \
-env=aloha \
-task=sim_sim_transfer_cube_human \
+env=pusht \
 hydra.run.dir=outputs/visualize_dataset/example
 # >>> ['outputs/visualize_dataset/example/episode_0.mp4']
 ```
 
 ### Evaluate a pretrained policy
 
-Check out [example 2](./examples/2_evaluate_pretrained_policy.py) to see how you can load a pretrained policy from HuggingFace hub, load up the corresponding environment and model, and run an evaluation.
+Check out [examples](./examples) to see how you can load a pretrained policy from HuggingFace hub, load up the corresponding environment and model, and run an evaluation.
 
 Or you can achieve the same result by executing our script from the command line:
 ```bash
@@ -169,94 +153,30 @@ See `python lerobot/scripts/eval.py --help` for more instructions.
 
 ### Train your own policy
 
-You can import our dataset, environment, policy classes, and use our training utilities (if some data is missing, it will be automatically downloaded from HuggingFace hub): check out [example 3](./examples/3_train_policy.py). After you run this, you may want to revisit [example 2](./examples/2_evaluate_pretrained_policy.py) to evaluate your training output!
+Check out [examples](./examples) to see how you can start training a model on a dataset, which will be automatically downloaded if needed.
 
 In general, you can use our training script to easily train any policy on any environment:
 ```bash
 python lerobot/scripts/train.py \
 env=aloha \
 task=sim_insertion \
-dataset_id=aloha_sim_insertion_scripted \
+repo_id=lerobot/aloha_sim_insertion_scripted \
 policy=act \
 hydra.run.dir=outputs/train/aloha_act
 ```
 
+After training, you may want to revisit model evaluation to change the evaluation settings. In fact, during training every checkpoint is already evaluated but on a low number of episodes for efficiency. Check out [example](./examples) to evaluate any model checkpoint on more episodes to increase statistical significance.
+
 ## Contribute
 
-Feel free to open issues and PRs, and to coordinate your efforts with the community on our [Discord Channel](https://discord.gg/VjFz58wn3R). For specific inquiries, reach out to [Remi Cadene](remi.cadene@huggingface.co).
-
-### TODO
-
-If you are not sure how to contribute or want to know the next features we working on, look on this project page: [LeRobot TODO](https://github.com/orgs/huggingface/projects/46)
-
-### Follow our style
-
-```bash
-# install if needed
-pre-commit install
-# apply style and linter checks before git commit
-pre-commit
-```
-
-### Add dependencies
-
-Instead of using `pip` directly, we use `poetry` for development purposes to easily track our dependencies.
-If you don't have it already, follow the [instructions](https://python-poetry.org/docs/#installation) to install it.
-
-Install the project with:
-```bash
-poetry install
-```
-
-Then, the equivalent of `pip install some-package`, would just be:
-```bash
-poetry add some-package
-```
-
-**NOTE:** Currently, to ensure the CI works properly, any new package must also be added in the CPU-only environment dedicated to the CI. To do this, you should create a separate environment and add the new package there as well. For example:
-```bash
-# Add the new package to your main poetry env
-poetry add some-package
-# Add the same package to the CPU-only env dedicated to CI
-conda create -y -n lerobot-ci python=3.10
-conda activate lerobot-ci
-cd .github/poetry/cpu
-poetry add some-package
-```
-
-### Run tests locally
-
-Install [git lfs](https://git-lfs.com/) to retrieve test artifacts (if you don't have it already).
-
-On Mac:
-```bash
-brew install git-lfs
-git lfs install
-```
-
-On Ubuntu:
-```bash
-sudo apt-get install git-lfs
-git lfs install
-```
-
-Pull artifacts if they're not in [tests/data](tests/data)
-```bash
-git lfs pull
-```
-
-When adding a new dataset, mock it with
-```bash
-python tests/scripts/mock_dataset.py --in-data-dir data/$DATASET --out-data-dir tests/data/$DATASET
-```
-
-Run tests
-```bash
-DATA_DIR="tests/data" pytest -sx tests
-```
+If you would like to contribute to 🤗 LeRobot, please check out our [contribution guide](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md).
 
 ### Add a new dataset
 
+```python
+# TODO(rcadene, AdilZouitine): rewrite this section
+```
+
 To add a dataset to the hub, first login and use a token generated from [huggingface settings](https://huggingface.co/settings/tokens) with write access:
 ```bash
 huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
@@ -271,7 +191,7 @@ HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli upload $HF_USER/$DATASET data/$DATAS
 
 You will need to set the corresponding version as a default argument in your dataset class:
 ```python
-  version: str | None = "v1.0",
+  version: str | None = "v1.1",
 ```
 See: [`lerobot/common/datasets/pusht.py`](https://github.com/Cadene/lerobot/blob/main/lerobot/common/datasets/pusht.py)
 
@@ -318,6 +238,10 @@ python tests/scripts/mock_dataset.py --in-data-dir data/$DATASET --out-data-dir
 
 ### Add a pretrained policy
 
+```python
+# TODO(rcadene, alexander-soare): rewrite this section
+```
+
 Once you have trained a policy you may upload it to the HuggingFace hub.
 
 Firstly, make sure you have a model repository set up on the hub. The hub ID looks like HF_USER/REPO_NAME.
@@ -326,15 +250,13 @@ Secondly, assuming you have trained a policy, you need:
 
 - `config.yaml` which you can get from the `.hydra` directory of your training output folder.
 - `model.pt` which should be one of the saved models in the `models` directory of your training output folder (they won't be named `model.pt` but you will need to choose one).
-- `stats.pth` which should point to the same file in the dataset directory (found in `data/{dataset_name}`).
 
 To upload these to the hub, prepare a folder with the following structure (you can use symlinks rather than copying):
 
 ```
 to_upload
     ├── config.yaml
-    ├── model.pt
-    └── stats.pth
+    └── model.pt
 ```
 
 With the folder prepared, run the following with a desired revision ID.

docker/lerobot-cpu/Dockerfile

@@ -0,0 +1,31 @@
+# Configure image
+ARG PYTHON_VERSION=3.10
+
+FROM python:${PYTHON_VERSION}-slim
+ARG PYTHON_VERSION
+ARG DEBIAN_FRONTEND=noninteractive
+
+# Install apt dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential cmake \
+    libglib2.0-0 libgl1-mesa-glx libegl1-mesa \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+# Create virtual environment
+RUN ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python
+RUN python -m venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+RUN echo "source /opt/venv/bin/activate" >> /root/.bashrc
+
+# Install LeRobot
+COPY . /lerobot
+WORKDIR /lerobot
+RUN pip install --upgrade --no-cache-dir pip
+RUN pip install --no-cache-dir ".[test, aloha, xarm, pusht]" \
+    --extra-index-url https://download.pytorch.org/whl/cpu
+
+# Set EGL as the rendering backend for MuJoCo
+ENV MUJOCO_GL="egl"
+
+# Execute in bash shell rather than python
+CMD ["/bin/bash"]

docker/lerobot-gpu/Dockerfile

@@ -0,0 +1,27 @@
+FROM nvidia/cuda:12.4.1-base-ubuntu22.04
+
+# Configure image
+ARG PYTHON_VERSION=3.10
+ARG DEBIAN_FRONTEND=noninteractive
+
+# Install apt dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential cmake \
+    libglib2.0-0 libgl1-mesa-glx libegl1-mesa \
+    python${PYTHON_VERSION} python${PYTHON_VERSION}-venv \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+# Create virtual environment
+RUN ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python
+RUN python -m venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+RUN echo "source /opt/venv/bin/activate" >> /root/.bashrc
+
+# Install LeRobot
+COPY . /lerobot
+WORKDIR /lerobot
+RUN pip install --upgrade --no-cache-dir pip
+RUN pip install --no-cache-dir ".[test, aloha, xarm, pusht]"
+
+# Set EGL as the rendering backend for MuJoCo
+ENV MUJOCO_GL="egl"

examples/1_load_lerobot_dataset.py

@@ -0,0 +1,91 @@
+"""
+This script demonstrates the use of `LeRobotDataset` class for handling and processing robotic datasets from Hugging Face.
+It illustrates how to load datasets, manipulate them, and apply transformations suitable for machine learning tasks in PyTorch.
+
+Features included in this script:
+- Loading a dataset and accessing its properties.
+- Filtering data by episode number.
+- Converting tensor data for visualization.
+- Saving video files from dataset frames.
+- Using advanced dataset features like timestamp-based frame selection.
+- Demonstrating compatibility with PyTorch DataLoader for batch processing.
+
+The script ends with examples of how to batch process data using PyTorch's DataLoader.
+"""
+
+from pathlib import Path
+
+import imageio
+import torch
+
+import lerobot
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+
+print("List of available datasets", lerobot.available_datasets)
+# # >>> ['lerobot/aloha_sim_insertion_human', 'lerobot/aloha_sim_insertion_scripted',
+# #     'lerobot/aloha_sim_transfer_cube_human', 'lerobot/aloha_sim_transfer_cube_scripted',
+# #     'lerobot/pusht', 'lerobot/xarm_lift_medium']
+
+repo_id = "lerobot/pusht"
+
+# You can easily load a dataset from a Hugging Face repositery
+dataset = LeRobotDataset(repo_id)
+
+# LeRobotDataset is actually a thin wrapper around an underlying Hugging Face dataset  (see https://huggingface.co/docs/datasets/index for more information).
+# TODO(rcadene): update to make the print pretty
+print(f"{dataset=}")
+print(f"{dataset.hf_dataset=}")
+
+# and provides additional utilities for robotics and compatibility with pytorch
+print(f"number of samples/frames: {dataset.num_samples=}")
+print(f"number of episodes: {dataset.num_episodes=}")
+print(f"average number of frames per episode: {dataset.num_samples / dataset.num_episodes:.3f}")
+print(f"frames per second used during data collection: {dataset.fps=}")
+print(f"keys to access images from cameras: {dataset.image_keys=}")
+
+# While the LeRobotDataset adds helpers for working within our library, we still expose the underling Hugging Face dataset.
+# It may be freely replaced or modified in place. Here we use the filtering to keep only frames from episode 5.
+# TODO(rcadene): remove this example of accessing hf_dataset
+dataset.hf_dataset = dataset.hf_dataset.filter(lambda frame: frame["episode_index"] == 5)
+
+# LeRobot datsets actually subclass PyTorch datasets. So you can do everything you know and love from working with the latter, for example: iterating through the dataset. Here we grab all the image frames.
+frames = [sample["observation.image"] for sample in dataset]
+
+# but frames are now float32 range [0,1] channel first (c,h,w) to follow pytorch convention,
+# to view them, we convert to uint8 range [0,255]
+frames = [(frame * 255).type(torch.uint8) for frame in frames]
+# and to channel last (h,w,c)
+frames = [frame.permute((1, 2, 0)).numpy() for frame in frames]
+
+# and finally save them to a mp4 video
+Path("outputs/examples/1_load_lerobot_dataset").mkdir(parents=True, exist_ok=True)
+imageio.mimsave("outputs/examples/1_load_lerobot_dataset/episode_5.mp4", frames, fps=dataset.fps)
+
+# For many machine learning applications we need to load histories of past observations, or trajectorys of future actions. Our datasets can load previous and future frames for each key/modality,
+# using timestamps differences with the current loaded frame. For instance:
+delta_timestamps = {
+    # loads 4 images: 1 second before current frame, 500 ms before, 200 ms before, and current frame
+    "observation.image": [-1, -0.5, -0.20, 0],
+    # loads 8 state vectors: 1.5 seconds before, 1 second before, ... 20 ms, 10 ms, and current frame
+    "observation.state": [-1.5, -1, -0.5, -0.20, -0.10, -0.02, -0.01, 0],
+    # loads 64 action vectors: current frame, 1 frame in the future, 2 frames, ... 63 frames in the future
+    "action": [t / dataset.fps for t in range(64)],
+}
+dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps)
+print(f"{dataset[0]['observation.image'].shape=}")  # (4,c,h,w)
+print(f"{dataset[0]['observation.state'].shape=}")  # (8,c)
+print(f"{dataset[0]['action'].shape=}")  # (64,c)
+
+# Finally, our datasets are fully compatible with PyTorch dataloaders and samplers
+# because they are just PyTorch datasets.
+dataloader = torch.utils.data.DataLoader(
+    dataset,
+    num_workers=0,
+    batch_size=32,
+    shuffle=True,
+)
+for batch in dataloader:
+    print(f"{batch['observation.image'].shape=}")  # (32,4,c,h,w)
+    print(f"{batch['observation.state'].shape=}")  # (32,8,c)
+    print(f"{batch['action'].shape=}")  # (32,64,c)
+    break

examples/1_visualize_dataset.py

@@ -1,24 +0,0 @@
-import os
-
-from torchrl.data.replay_buffers import SamplerWithoutReplacement
-
-import lerobot
-from lerobot.common.datasets.aloha import AlohaDataset
-from lerobot.scripts.visualize_dataset import render_dataset
-
-print(lerobot.available_datasets)
-# >>> ['aloha_sim_insertion_human', 'aloha_sim_insertion_scripted', 'aloha_sim_transfer_cube_human', 'aloha_sim_transfer_cube_scripted', 'pusht', 'xarm_lift_medium']
-
-# we use this sampler to sample 1 frame after the other
-sampler = SamplerWithoutReplacement(shuffle=False)
-
-dataset = AlohaDataset("aloha_sim_transfer_cube_human", sampler=sampler, root=os.environ.get("DATA_DIR"))
-
-video_paths = render_dataset(
-    dataset,
-    out_dir="outputs/visualize_dataset/example",
-    max_num_samples=300,
-    fps=50,
-)
-print(video_paths)
-# ['outputs/visualize_dataset/example/episode_0.mp4']

examples/2_evaluate_pretrained_policy.py

@@ -7,10 +7,11 @@ from pathlib import Path
 
 from huggingface_hub import snapshot_download
 
-from lerobot.common.utils import init_hydra_config
+from lerobot.common.utils.utils import init_hydra_config
 from lerobot.scripts.eval import eval
 
 # Get a pretrained policy from the hub.
+# TODO(alexander-soare): This no longer works until we upload a new model that uses the current configs.
 hub_id = "lerobot/diffusion_policy_pusht_image"
 folder = Path(snapshot_download(hub_id))
 # OR uncomment the following to evaluate a policy from the local outputs/train folder.
@@ -18,7 +19,6 @@ folder = Path(snapshot_download(hub_id))
 
 config_path = folder / "config.yaml"
 weights_path = folder / "model.pt"
-stats_path = folder / "stats.pth"  # normalization stats
 
 # Override some config parameters to do with evaluation.
 overrides = [
@@ -35,5 +35,4 @@ cfg = init_hydra_config(config_path, overrides)
 eval(
     cfg,
     out_dir=f"outputs/eval/example_{cfg.env.name}_{cfg.policy.name}",
-    stats_path=stats_path,
 )

examples/3_train_policy.py

@@ -9,47 +9,59 @@ from pathlib import Path
 
 import torch
 from omegaconf import OmegaConf
-from tqdm import trange
 
-from lerobot.common.datasets.factory import make_offline_buffer
-from lerobot.common.policies.diffusion.policy import DiffusionPolicy
-from lerobot.common.utils import init_hydra_config
+from lerobot.common.datasets.factory import make_dataset
+from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
+from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
+from lerobot.common.utils.utils import init_hydra_config
 
 output_directory = Path("outputs/train/example_pusht_diffusion")
 os.makedirs(output_directory, exist_ok=True)
 
-overrides = [
-    "env=pusht",
-    "policy=diffusion",
-    # Adjust as you prefer. 5000 steps are needed to get something worth evaluating.
-    "offline_steps=5000",
-    "log_freq=250",
-    "device=cuda",
-]
+# Number of offline training steps (we'll only do offline training for this example.
+# Adjust as you prefer. 5000 steps are needed to get something worth evaluating.
+training_steps = 5000
+device = torch.device("cuda")
+log_freq = 250
 
-cfg = init_hydra_config("lerobot/configs/default.yaml", overrides)
+# Set up the dataset.
+hydra_cfg = init_hydra_config("lerobot/configs/default.yaml", overrides=["env=pusht"])
+dataset = make_dataset(hydra_cfg)
 
-policy = DiffusionPolicy(
-    cfg=cfg.policy,
-    cfg_device=cfg.device,
-    cfg_noise_scheduler=cfg.noise_scheduler,
-    cfg_rgb_model=cfg.rgb_model,
-    cfg_obs_encoder=cfg.obs_encoder,
-    cfg_optimizer=cfg.optimizer,
-    cfg_ema=cfg.ema,
-    n_action_steps=cfg.n_action_steps + cfg.n_latency_steps,
-    **cfg.policy,
-)
+# Set up the the policy.
+# Policies are initialized with a configuration class, in this case `DiffusionConfig`.
+# For this example, no arguments need to be passed because the defaults are set up for PushT.
+# If you're doing something different, you will likely need to change at least some of the defaults.
+cfg = DiffusionConfig()
+# TODO(alexander-soare): Remove LR scheduler from the policy.
+policy = DiffusionPolicy(cfg, lr_scheduler_num_training_steps=training_steps, dataset_stats=dataset.stats)
 policy.train()
+policy.to(device)
 
-offline_buffer = make_offline_buffer(cfg)
+# Create dataloader for offline training.
+dataloader = torch.utils.data.DataLoader(
+    dataset,
+    num_workers=4,
+    batch_size=cfg.batch_size,
+    shuffle=True,
+    pin_memory=device != torch.device("cpu"),
+    drop_last=True,
+)
 
-for offline_step in trange(cfg.offline_steps):
-    train_info = policy.update(offline_buffer, offline_step)
-    if offline_step % cfg.log_freq == 0:
-        print(train_info)
+# Run training loop.
+step = 0
+done = False
+while not done:
+    for batch in dataloader:
+        batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
+        info = policy.update(batch)
+        if step % log_freq == 0:
+            print(f"step: {step} loss: {info['loss']:.3f} update_time: {info['update_s']:.3f} (seconds)")
+        step += 1
+        if step >= training_steps:
+            done = True
+            break
 
-# Save the policy, configuration, and normalization stats for later use.
+# Save the policy and configuration for later use.
 policy.save(output_directory / "model.pt")
-OmegaConf.save(cfg, output_directory / "config.yaml")
-torch.save(offline_buffer.transform[-1].stats, output_directory / "stats.pth")
+OmegaConf.save(hydra_cfg, output_directory / "config.yaml")

lerobot/__init__.py

@@ -7,53 +7,86 @@ Example:
         import lerobot
         print(lerobot.available_envs)
         print(lerobot.available_tasks_per_env)
-        print(lerobot.available_datasets_per_env)
         print(lerobot.available_datasets)
+        print(lerobot.available_datasets_per_env)
         print(lerobot.available_policies)
+        print(lerobot.available_policies_per_env)
     ```
 
-Note:
-    When implementing a concrete class (e.g. `AlohaDataset`, `PushtEnv`, `DiffusionPolicy`), you need to:
-        1. set the required class attributes:
-            - for classes inheriting from `AbstractDataset`: `available_datasets`
-            - for classes inheriting from `AbstractEnv`: `name`, `available_tasks`
-            - for classes inheriting from `AbstractPolicy`: `name`
-        2. update variables in `lerobot/__init__.py` (e.g. `available_envs`, `available_datasets_per_envs`, `available_policies`)
-        3. update variables in `tests/test_available.py` by importing your new class
+When implementing a new dataset loadable with LeRobotDataset follow these steps:
+- Update `available_datasets_per_env` in `lerobot/__init__.py`
+
+When implementing a new environment (e.g. `gym_aloha`), follow these steps:
+- Update `available_tasks_per_env` and `available_datasets_per_env` in `lerobot/__init__.py`
+
+When implementing a new policy class (e.g. `DiffusionPolicy`) follow these steps:
+- Update `available_policies` and `available_policies_per_env`, in `lerobot/__init__.py`
+- Set the required `name` class attribute.
+- Update variables in `tests/test_available.py` by importing your new Policy class
 """
 
+import itertools
+
 from lerobot.__version__ import __version__  # noqa: F401
 
-available_envs = [
-    "aloha",
-    "pusht",
-    "simxarm",
-]
-
 available_tasks_per_env = {
     "aloha": [
-        "sim_insertion",
-        "sim_transfer_cube",
+        "AlohaInsertion-v0",
+        "AlohaTransferCube-v0",
     ],
-    "pusht": ["pusht"],
-    "simxarm": ["lift"],
+    "pusht": ["PushT-v0"],
+    "xarm": ["XarmLift-v0"],
 }
+available_envs = list(available_tasks_per_env.keys())
 
 available_datasets_per_env = {
     "aloha": [
-        "aloha_sim_insertion_human",
-        "aloha_sim_insertion_scripted",
-        "aloha_sim_transfer_cube_human",
-        "aloha_sim_transfer_cube_scripted",
+        "lerobot/aloha_sim_insertion_human",
+        "lerobot/aloha_sim_insertion_scripted",
+        "lerobot/aloha_sim_transfer_cube_human",
+        "lerobot/aloha_sim_transfer_cube_scripted",
+    ],
+    "pusht": ["lerobot/pusht"],
+    "xarm": [
+        "lerobot/xarm_lift_medium",
+        "lerobot/xarm_lift_medium_replay",
+        "lerobot/xarm_push_medium",
+        "lerobot/xarm_push_medium_replay",
     ],
-    "pusht": ["pusht"],
-    "simxarm": ["xarm_lift_medium"],
 }
 
-available_datasets = [dataset for env in available_envs for dataset in available_datasets_per_env[env]]
+available_datasets_without_env = ["lerobot/umi_cup_in_the_wild"]
+
+available_datasets = list(
+    itertools.chain(*available_datasets_per_env.values(), available_datasets_without_env)
+)
+
+# TODO(rcadene, aliberts, alexander-soare): Add real-world env with a gym API
+available_datasets_without_env = ["lerobot/umi_cup_in_the_wild"]
+
+available_datasets = list(
+    itertools.chain(*available_datasets_per_env.values(), available_datasets_without_env)
+)
 
 available_policies = [
     "act",
     "diffusion",
     "tdmpc",
 ]
+
+available_policies_per_env = {
+    "aloha": ["act"],
+    "pusht": ["diffusion"],
+    "xarm": ["tdmpc"],
+}
+
+env_task_pairs = [(env, task) for env, tasks in available_tasks_per_env.items() for task in tasks]
+env_dataset_pairs = [
+    (env, dataset) for env, datasets in available_datasets_per_env.items() for dataset in datasets
+]
+env_dataset_policy_triplets = [
+    (env, dataset, policy)
+    for env, datasets in available_datasets_per_env.items()
+    for dataset in datasets
+    for policy in available_policies_per_env[env]
+]

lerobot/common/datasets/abstract.py

@@ -1,207 +0,0 @@
-import logging
-from pathlib import Path
-from typing import Callable
-
-import einops
-import torch
-import torchrl
-import tqdm
-from huggingface_hub import snapshot_download
-from tensordict import TensorDict
-from torchrl.data.replay_buffers.replay_buffers import TensorDictReplayBuffer
-from torchrl.data.replay_buffers.samplers import Sampler
-from torchrl.data.replay_buffers.storages import TensorStorage, _collate_id
-from torchrl.data.replay_buffers.writers import ImmutableDatasetWriter, Writer
-from torchrl.envs.transforms.transforms import Compose
-
-HF_USER = "lerobot"
-
-
-class AbstractDataset(TensorDictReplayBuffer):
-    """
-    AbstractDataset represents a dataset in the context of imitation learning or reinforcement learning.
-    This class is designed to be subclassed by concrete implementations that specify particular types of datasets.
-    These implementations can vary based on the source of the data, the environment the data pertains to,
-    or the specific kind of data manipulation applied.
-
-    Note:
-        - `TensorDictReplayBuffer` is the base class from which `AbstractDataset` inherits. It provides the foundational
-           functionality for storing and retrieving `TensorDict`-like data.
-        - `available_datasets` should be overridden by concrete subclasses to list the specific dataset variants supported.
-           It is expected that these variants correspond to a HuggingFace dataset on the hub.
-           For instance, the `AlohaDataset` which inherites from `AbstractDataset` has 4 available dataset variants:
-            - [aloha_sim_transfer_cube_scripted](https://huggingface.co/datasets/lerobot/aloha_sim_transfer_cube_scripted)
-            - [aloha_sim_insertion_scripted](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_scripted)
-            - [aloha_sim_transfer_cube_human](https://huggingface.co/datasets/lerobot/aloha_sim_transfer_cube_human)
-            - [aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human)
-        - When implementing a concrete class (e.g. `AlohaDataset`, `PushtEnv`, `DiffusionPolicy`), you need to:
-            1. set the required class attributes:
-                - for classes inheriting from `AbstractDataset`: `available_datasets`
-                - for classes inheriting from `AbstractEnv`: `name`, `available_tasks`
-                - for classes inheriting from `AbstractPolicy`: `name`
-            2. update variables in `lerobot/__init__.py` (e.g. `available_envs`, `available_datasets_per_envs`, `available_policies`)
-            3. update variables in `tests/test_available.py` by importing your new class
-    """
-
-    available_datasets: list[str] | None = None
-
-    def __init__(
-        self,
-        dataset_id: str,
-        version: str | None = None,
-        batch_size: int | None = None,
-        *,
-        shuffle: bool = True,
-        root: Path | None = None,
-        pin_memory: bool = False,
-        prefetch: int = None,
-        sampler: Sampler | None = None,
-        collate_fn: Callable | None = None,
-        writer: Writer | None = None,
-        transform: "torchrl.envs.Transform" = None,
-    ):
-        assert (
-            self.available_datasets is not None
-        ), "Subclasses of `AbstractDataset` should set the `available_datasets` class attribute."
-        assert (
-            dataset_id in self.available_datasets
-        ), f"The provided dataset ({dataset_id}) is not on the list of available datasets {self.available_datasets}."
-
-        self.dataset_id = dataset_id
-        self.version = version
-        self.shuffle = shuffle
-        self.root = root if root is None else Path(root)
-
-        if self.root is not None and self.version is not None:
-            logging.warning(
-                f"The version of the dataset ({self.version}) is not enforced when root is provided ({self.root})."
-            )
-
-        storage = self._download_or_load_dataset()
-
-        super().__init__(
-            storage=storage,
-            sampler=sampler,
-            writer=ImmutableDatasetWriter() if writer is None else writer,
-            collate_fn=_collate_id if collate_fn is None else collate_fn,
-            pin_memory=pin_memory,
-            prefetch=prefetch,
-            batch_size=batch_size,
-            transform=transform,
-        )
-
-    @property
-    def stats_patterns(self) -> dict:
-        return {
-            ("observation", "state"): "b c -> c",
-            ("observation", "image"): "b c h w -> c 1 1",
-            ("action",): "b c -> c",
-        }
-
-    @property
-    def image_keys(self) -> list:
-        return [("observation", "image")]
-
-    @property
-    def num_cameras(self) -> int:
-        return len(self.image_keys)
-
-    @property
-    def num_samples(self) -> int:
-        return len(self)
-
-    @property
-    def num_episodes(self) -> int:
-        return len(self._storage._storage["episode"].unique())
-
-    @property
-    def transform(self):
-        return self._transform
-
-    def set_transform(self, transform):
-        if not isinstance(transform, Compose):
-            # required since torchrl calls `len(self._transform)` downstream
-            if isinstance(transform, list):
-                self._transform = Compose(*transform)
-            else:
-                self._transform = Compose(transform)
-        else:
-            self._transform = transform
-
-    def compute_or_load_stats(self, num_batch=100, batch_size=32) -> TensorDict:
-        stats_path = self.data_dir / "stats.pth"
-        if stats_path.exists():
-            stats = torch.load(stats_path)
-        else:
-            logging.info(f"compute_stats and save to {stats_path}")
-            stats = self._compute_stats(num_batch, batch_size)
-            torch.save(stats, stats_path)
-        return stats
-
-    def _download_or_load_dataset(self) -> torch.StorageBase:
-        if self.root is None:
-            self.data_dir = Path(
-                snapshot_download(
-                    repo_id=f"{HF_USER}/{self.dataset_id}", repo_type="dataset", revision=self.version
-                )
-            )
-        else:
-            self.data_dir = self.root / self.dataset_id
-        return TensorStorage(TensorDict.load_memmap(self.data_dir / "replay_buffer"))
-
-    def _compute_stats(self, num_batch=100, batch_size=32):
-        rb = TensorDictReplayBuffer(
-            storage=self._storage,
-            batch_size=batch_size,
-            prefetch=True,
-        )
-
-        mean, std, max, min = {}, {}, {}, {}
-
-        # compute mean, min, max
-        for _ in tqdm.tqdm(range(num_batch)):
-            batch = rb.sample()
-            for key, pattern in self.stats_patterns.items():
-                batch[key] = batch[key].float()
-                if key not in mean:
-                    # first batch initialize mean, min, max
-                    mean[key] = einops.reduce(batch[key], pattern, "mean")
-                    max[key] = einops.reduce(batch[key], pattern, "max")
-                    min[key] = einops.reduce(batch[key], pattern, "min")
-                else:
-                    mean[key] += einops.reduce(batch[key], pattern, "mean")
-                    max[key] = torch.maximum(max[key], einops.reduce(batch[key], pattern, "max"))
-                    min[key] = torch.minimum(min[key], einops.reduce(batch[key], pattern, "min"))
-                batch = rb.sample()
-
-        for key in self.stats_patterns:
-            mean[key] /= num_batch
-
-        # compute std, min, max
-        for _ in tqdm.tqdm(range(num_batch)):
-            batch = rb.sample()
-            for key, pattern in self.stats_patterns.items():
-                batch[key] = batch[key].float()
-                batch_mean = einops.reduce(batch[key], pattern, "mean")
-                if key not in std:
-                    # first batch initialize std
-                    std[key] = (batch_mean - mean[key]) ** 2
-                else:
-                    std[key] += (batch_mean - mean[key]) ** 2
-                max[key] = torch.maximum(max[key], einops.reduce(batch[key], pattern, "max"))
-                min[key] = torch.minimum(min[key], einops.reduce(batch[key], pattern, "min"))
-
-        for key in self.stats_patterns:
-            std[key] = torch.sqrt(std[key] / num_batch)
-
-        stats = TensorDict({}, batch_size=[])
-        for key in self.stats_patterns:
-            stats[(*key, "mean")] = mean[key]
-            stats[(*key, "std")] = std[key]
-            stats[(*key, "max")] = max[key]
-            stats[(*key, "min")] = min[key]
-
-            if key[0] == "observation":
-                # use same stats for the next observations
-                stats[("next", *key)] = stats[key]
-        return stats

lerobot/common/datasets/aloha.py

@@ -1,185 +0,0 @@
-import logging
-from pathlib import Path
-from typing import Callable
-
-import einops
-import gdown
-import h5py
-import torch
-import torchrl
-import tqdm
-from tensordict import TensorDict
-from torchrl.data.replay_buffers.samplers import Sampler
-from torchrl.data.replay_buffers.storages import TensorStorage
-from torchrl.data.replay_buffers.writers import Writer
-
-from lerobot.common.datasets.abstract import AbstractDataset
-
-DATASET_IDS = [
-    "aloha_sim_insertion_human",
-    "aloha_sim_insertion_scripted",
-    "aloha_sim_transfer_cube_human",
-    "aloha_sim_transfer_cube_scripted",
-]
-
-FOLDER_URLS = {
-    "aloha_sim_insertion_human": "https://drive.google.com/drive/folders/1RgyD0JgTX30H4IM5XZn8I3zSV_mr8pyF",
-    "aloha_sim_insertion_scripted": "https://drive.google.com/drive/folders/1TsojQQSXtHEoGnqgJ3gmpPQR2DPLtS2N",
-    "aloha_sim_transfer_cube_human": "https://drive.google.com/drive/folders/1sc-E4QYW7A0o23m1u2VWNGVq5smAsfCo",
-    "aloha_sim_transfer_cube_scripted": "https://drive.google.com/drive/folders/1aRyoOhQwxhyt1J8XgEig4s6kzaw__LXj",
-}
-
-EP48_URLS = {
-    "aloha_sim_insertion_human": "https://drive.google.com/file/d/18Cudl6nikDtgRolea7je8iF_gGKzynOP/view?usp=drive_link",
-    "aloha_sim_insertion_scripted": "https://drive.google.com/file/d/1wfMSZ24oOh5KR_0aaP3Cnu_c4ZCveduB/view?usp=drive_link",
-    "aloha_sim_transfer_cube_human": "https://drive.google.com/file/d/18smMymtr8tIxaNUQ61gW6dG50pt3MvGq/view?usp=drive_link",
-    "aloha_sim_transfer_cube_scripted": "https://drive.google.com/file/d/1pnGIOd-E4-rhz2P3VxpknMKRZCoKt6eI/view?usp=drive_link",
-}
-
-EP49_URLS = {
-    "aloha_sim_insertion_human": "https://drive.google.com/file/d/1C1kZYyROzs-PrLc0SkDgUgMi4-L3lauE/view?usp=drive_link",
-    "aloha_sim_insertion_scripted": "https://drive.google.com/file/d/17EuCUWS6uCCr6yyNzpXdcdE-_TTNCKtf/view?usp=drive_link",
-    "aloha_sim_transfer_cube_human": "https://drive.google.com/file/d/1Nk7l53d9sJoGDBKAOnNrExX5nLacATc6/view?usp=drive_link",
-    "aloha_sim_transfer_cube_scripted": "https://drive.google.com/file/d/1GKReZHrXU73NMiC5zKCq_UtqPVtYq8eo/view?usp=drive_link",
-}
-
-NUM_EPISODES = {
-    "aloha_sim_insertion_human": 50,
-    "aloha_sim_insertion_scripted": 50,
-    "aloha_sim_transfer_cube_human": 50,
-    "aloha_sim_transfer_cube_scripted": 50,
-}
-
-EPISODE_LEN = {
-    "aloha_sim_insertion_human": 500,
-    "aloha_sim_insertion_scripted": 400,
-    "aloha_sim_transfer_cube_human": 400,
-    "aloha_sim_transfer_cube_scripted": 400,
-}
-
-CAMERAS = {
-    "aloha_sim_insertion_human": ["top"],
-    "aloha_sim_insertion_scripted": ["top"],
-    "aloha_sim_transfer_cube_human": ["top"],
-    "aloha_sim_transfer_cube_scripted": ["top"],
-}
-
-
-def download(data_dir, dataset_id):
-    assert dataset_id in DATASET_IDS
-    assert dataset_id in FOLDER_URLS
-    assert dataset_id in EP48_URLS
-    assert dataset_id in EP49_URLS
-
-    data_dir.mkdir(parents=True, exist_ok=True)
-
-    gdown.download_folder(FOLDER_URLS[dataset_id], output=str(data_dir))
-
-    # because of the 50 files limit per directory, two files episode 48 and 49 were missing
-    gdown.download(EP48_URLS[dataset_id], output=str(data_dir / "episode_48.hdf5"), fuzzy=True)
-    gdown.download(EP49_URLS[dataset_id], output=str(data_dir / "episode_49.hdf5"), fuzzy=True)
-
-
-class AlohaDataset(AbstractDataset):
-    available_datasets = DATASET_IDS
-
-    def __init__(
-        self,
-        dataset_id: str,
-        version: str | None = "v1.2",
-        batch_size: int | None = None,
-        *,
-        shuffle: bool = True,
-        root: Path | None = None,
-        pin_memory: bool = False,
-        prefetch: int = None,
-        sampler: Sampler | None = None,
-        collate_fn: Callable | None = None,
-        writer: Writer | None = None,
-        transform: "torchrl.envs.Transform" = None,
-    ):
-        super().__init__(
-            dataset_id,
-            version,
-            batch_size,
-            shuffle=shuffle,
-            root=root,
-            pin_memory=pin_memory,
-            prefetch=prefetch,
-            sampler=sampler,
-            collate_fn=collate_fn,
-            writer=writer,
-            transform=transform,
-        )
-
-    @property
-    def stats_patterns(self) -> dict:
-        d = {
-            ("observation", "state"): "b c -> c",
-            ("action",): "b c -> c",
-        }
-        for cam in CAMERAS[self.dataset_id]:
-            d[("observation", "image", cam)] = "b c h w -> c 1 1"
-        return d
-
-    @property
-    def image_keys(self) -> list:
-        return [("observation", "image", cam) for cam in CAMERAS[self.dataset_id]]
-
-    def _download_and_preproc_obsolete(self):
-        assert self.root is not None
-        raw_dir = self.root / f"{self.dataset_id}_raw"
-        if not raw_dir.is_dir():
-            download(raw_dir, self.dataset_id)
-
-        total_num_frames = 0
-        logging.info("Compute total number of frames to initialize offline buffer")
-        for ep_id in range(NUM_EPISODES[self.dataset_id]):
-            ep_path = raw_dir / f"episode_{ep_id}.hdf5"
-            with h5py.File(ep_path, "r") as ep:
-                total_num_frames += ep["/action"].shape[0] - 1
-        logging.info(f"{total_num_frames=}")
-
-        logging.info("Initialize and feed offline buffer")
-        idxtd = 0
-        for ep_id in tqdm.tqdm(range(NUM_EPISODES[self.dataset_id])):
-            ep_path = raw_dir / f"episode_{ep_id}.hdf5"
-            with h5py.File(ep_path, "r") as ep:
-                ep_num_frames = ep["/action"].shape[0]
-
-                # last step of demonstration is considered done
-                done = torch.zeros(ep_num_frames, 1, dtype=torch.bool)
-                done[-1] = True
-
-                state = torch.from_numpy(ep["/observations/qpos"][:])
-                action = torch.from_numpy(ep["/action"][:])
-
-                ep_td = TensorDict(
-                    {
-                        ("observation", "state"): state[:-1],
-                        "action": action[:-1],
-                        "episode": torch.tensor([ep_id] * (ep_num_frames - 1)),
-                        "frame_id": torch.arange(0, ep_num_frames - 1, 1),
-                        ("next", "observation", "state"): state[1:],
-                        # TODO: compute reward and success
-                        # ("next", "reward"): reward[1:],
-                        ("next", "done"): done[1:],
-                        # ("next", "success"): success[1:],
-                    },
-                    batch_size=ep_num_frames - 1,
-                )
-
-                for cam in CAMERAS[self.dataset_id]:
-                    image = torch.from_numpy(ep[f"/observations/images/{cam}"][:])
-                    image = einops.rearrange(image, "b h w c -> b c h w").contiguous()
-                    ep_td["observation", "image", cam] = image[:-1]
-                    ep_td["next", "observation", "image", cam] = image[1:]
-
-                if ep_id == 0:
-                    # hack to initialize tensordict data structure to store episodes
-                    td_data = ep_td[0].expand(total_num_frames).memmap_like(self.root / f"{self.dataset_id}")
-
-                td_data[idxtd : idxtd + len(ep_td)] = ep_td
-                idxtd = idxtd + len(ep_td)
-
-        return TensorStorage(td_data.lock_())

lerobot/common/datasets/factory.py

@@ -3,134 +3,42 @@ import os
 from pathlib import Path
 
 import torch
-from torchrl.data.replay_buffers import PrioritizedSliceSampler, SliceSampler
+from omegaconf import OmegaConf
 
-from lerobot.common.transforms import NormalizeTransform, Prod
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
 
-# DATA_DIR specifies to location where datasets are loaded. By default, DATA_DIR is None and
-# we load from `$HOME/.cache/huggingface/hub/datasets`. For our unit tests, we set `DATA_DIR=tests/data`
-# to load a subset of our datasets for faster continuous integration.
 DATA_DIR = Path(os.environ["DATA_DIR"]) if "DATA_DIR" in os.environ else None
 
 
-def make_offline_buffer(
+def make_dataset(
     cfg,
-    overwrite_sampler=None,
-    # set normalize=False to remove all transformations and keep images unnormalized in [0,255]
-    normalize=True,
-    overwrite_batch_size=None,
-    overwrite_prefetch=None,
-    stats_path=None,
+    split="train",
 ):
-    if cfg.policy.balanced_sampling:
-        assert cfg.online_steps > 0
-        batch_size = None
-        pin_memory = False
-        prefetch = None
-    else:
-        assert cfg.online_steps == 0
-        num_slices = cfg.policy.batch_size
-        batch_size = cfg.policy.horizon * num_slices
-        pin_memory = cfg.device == "cuda"
-        prefetch = cfg.prefetch
+    if cfg.env.name not in cfg.dataset.repo_id:
+        logging.warning(
+            f"There might be a mismatch between your training dataset ({cfg.dataset.repo_id=}) and your environment ({cfg.env.name=})."
+        )
 
-    if overwrite_batch_size is not None:
-        batch_size = overwrite_batch_size
+    delta_timestamps = cfg.policy.get("delta_timestamps")
+    if delta_timestamps is not None:
+        for key in delta_timestamps:
+            if isinstance(delta_timestamps[key], str):
+                delta_timestamps[key] = eval(delta_timestamps[key])
 
-    if overwrite_prefetch is not None:
-        prefetch = overwrite_prefetch
+    # TODO(rcadene): add data augmentations
 
-    if overwrite_sampler is None:
-        # TODO(rcadene): move batch_size outside
-        num_traj_per_batch = cfg.policy.batch_size  # // cfg.horizon
-        # TODO(rcadene): Sampler outputs a batch_size <= cfg.batch_size.
-        # We would need to add a transform to pad the tensordict to ensure batch_size == cfg.batch_size.
-
-        if cfg.offline_prioritized_sampler:
-            logging.info("use prioritized sampler for offline dataset")
-            sampler = PrioritizedSliceSampler(
-                max_capacity=100_000,
-                alpha=cfg.policy.per_alpha,
-                beta=cfg.policy.per_beta,
-                num_slices=num_traj_per_batch,
-                strict_length=False,
-            )
-        else:
-            logging.info("use simple sampler for offline dataset")
-            sampler = SliceSampler(
-                num_slices=num_traj_per_batch,
-                strict_length=False,
-            )
-    else:
-        sampler = overwrite_sampler
-
-    if cfg.env.name == "simxarm":
-        from lerobot.common.datasets.simxarm import SimxarmDataset
-
-        clsfunc = SimxarmDataset
-
-    elif cfg.env.name == "pusht":
-        from lerobot.common.datasets.pusht import PushtDataset
-
-        clsfunc = PushtDataset
-
-    elif cfg.env.name == "aloha":
-        from lerobot.common.datasets.aloha import AlohaDataset
-
-        clsfunc = AlohaDataset
-    else:
-        raise ValueError(cfg.env.name)
-
-    offline_buffer = clsfunc(
-        dataset_id=cfg.dataset_id,
-        sampler=sampler,
-        batch_size=batch_size,
+    dataset = LeRobotDataset(
+        cfg.dataset.repo_id,
+        split=split,
         root=DATA_DIR,
-        pin_memory=pin_memory,
-        prefetch=prefetch if isinstance(prefetch, int) else None,
+        delta_timestamps=delta_timestamps,
     )
 
-    if cfg.policy.name == "tdmpc":
-        img_keys = []
-        for key in offline_buffer.image_keys:
-            img_keys.append(("next", *key))
-        img_keys += offline_buffer.image_keys
-    else:
-        img_keys = offline_buffer.image_keys
+    if cfg.get("override_dataset_stats"):
+        for key, stats_dict in cfg.override_dataset_stats.items():
+            for stats_type, listconfig in stats_dict.items():
+                # example of stats_type: min, max, mean, std
+                stats = OmegaConf.to_container(listconfig, resolve=True)
+                dataset.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32)
 
-    if normalize:
-        transforms = [Prod(in_keys=img_keys, prod=1 / 255)]
-
-        # TODO(rcadene): make normalization strategy configurable between mean_std, min_max, manual_min_max,
-        # min_max_from_spec
-        stats = offline_buffer.compute_or_load_stats() if stats_path is None else torch.load(stats_path)
-
-        # we only normalize the state and action, since the images are usually normalized inside the model for
-        # now (except for tdmpc: see the following)
-        in_keys = [("observation", "state"), ("action")]
-
-        if cfg.policy.name == "tdmpc":
-            # TODO(rcadene): we add img_keys to the keys to normalize for tdmpc only, since diffusion and act policies normalize the image inside the model for now
-            in_keys += img_keys
-            # TODO(racdene): since we use next observations in tdmpc, we also add them to the normalization. We are wasting a bit of compute on this for now.
-            in_keys += [("next", *key) for key in img_keys]
-            in_keys.append(("next", "observation", "state"))
-
-        if cfg.policy.name == "diffusion" and cfg.env.name == "pusht":
-            # TODO(rcadene): we overwrite stats to have the same as pretrained model, but we should remove this
-            stats["observation", "state", "min"] = torch.tensor([13.456424, 32.938293], dtype=torch.float32)
-            stats["observation", "state", "max"] = torch.tensor([496.14618, 510.9579], dtype=torch.float32)
-            stats["action", "min"] = torch.tensor([12.0, 25.0], dtype=torch.float32)
-            stats["action", "max"] = torch.tensor([511.0, 511.0], dtype=torch.float32)
-
-        # TODO(rcadene): remove this and put it in config. Ideally we want to reproduce SOTA results just with mean_std
-        normalization_mode = "mean_std" if cfg.env.name == "aloha" else "min_max"
-        transforms.append(NormalizeTransform(stats, in_keys, mode=normalization_mode))
-
-        offline_buffer.set_transform(transforms)
-
-    if not overwrite_sampler:
-        index = torch.arange(0, offline_buffer.num_samples, 1)
-        sampler.extend(index)
-
-    return offline_buffer
+    return dataset

lerobot/common/datasets/lerobot_dataset.py

@@ -0,0 +1,72 @@
+from pathlib import Path
+
+import datasets
+import torch
+
+from lerobot.common.datasets.utils import (
+    load_episode_data_index,
+    load_hf_dataset,
+    load_info,
+    load_previous_and_future_frames,
+    load_stats,
+)
+
+
+class LeRobotDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        repo_id: str,
+        version: str | None = "v1.1",
+        root: Path | None = None,
+        split: str = "train",
+        transform: callable = None,
+        delta_timestamps: dict[list[float]] | None = None,
+    ):
+        super().__init__()
+        self.repo_id = repo_id
+        self.version = version
+        self.root = root
+        self.split = split
+        self.transform = transform
+        self.delta_timestamps = delta_timestamps
+        # load data from hub or locally when root is provided
+        self.hf_dataset = load_hf_dataset(repo_id, version, root, split)
+        self.episode_data_index = load_episode_data_index(repo_id, version, root)
+        self.stats = load_stats(repo_id, version, root)
+        self.info = load_info(repo_id, version, root)
+
+    @property
+    def fps(self) -> int:
+        return self.info["fps"]
+
+    @property
+    def image_keys(self) -> list[str]:
+        return [key for key, feats in self.hf_dataset.features.items() if isinstance(feats, datasets.Image)]
+
+    @property
+    def num_samples(self) -> int:
+        return len(self.hf_dataset)
+
+    @property
+    def num_episodes(self) -> int:
+        return len(self.hf_dataset.unique("episode_index"))
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, idx):
+        item = self.hf_dataset[idx]
+
+        if self.delta_timestamps is not None:
+            item = load_previous_and_future_frames(
+                item,
+                self.hf_dataset,
+                self.episode_data_index,
+                self.delta_timestamps,
+                tol=1 / self.fps - 1e-4,  # 1e-4 to account for possible numerical error
+            )
+
+        if self.transform is not None:
+            item = self.transform(item)
+
+        return item

lerobot/common/datasets/push_dataset_to_hub/_diffusion_policy_replay_buffer.py

@@ -1,3 +1,8 @@
+"""Helper code for loading PushT dataset from Diffusion Policy (https://diffusion-policy.cs.columbia.edu/)
+
+Copied from the original Diffusion Policy repository and used in our `download_and_upload_dataset.py` script.
+"""
+
 from __future__ import annotations
 
 import math

lerobot/common/datasets/push_dataset_to_hub/_download_raw.py

@@ -0,0 +1,179 @@
+"""
+This file contains all obsolete download scripts. They are centralized here to not have to load
+useless dependencies when using datasets.
+"""
+
+import io
+from pathlib import Path
+
+import tqdm
+
+
+def download_raw(root, dataset_id) -> Path:
+    if "pusht" in dataset_id:
+        return download_pusht(root=root, dataset_id=dataset_id)
+    elif "xarm" in dataset_id:
+        return download_xarm(root=root, dataset_id=dataset_id)
+    elif "aloha" in dataset_id:
+        return download_aloha(root=root, dataset_id=dataset_id)
+    elif "umi" in dataset_id:
+        return download_umi(root=root, dataset_id=dataset_id)
+    else:
+        raise ValueError(dataset_id)
+
+
+def download_and_extract_zip(url: str, destination_folder: Path) -> bool:
+    import zipfile
+
+    import requests
+
+    print(f"downloading from {url}")
+    response = requests.get(url, stream=True)
+    if response.status_code == 200:
+        total_size = int(response.headers.get("content-length", 0))
+        progress_bar = tqdm.tqdm(total=total_size, unit="B", unit_scale=True)
+
+        zip_file = io.BytesIO()
+        for chunk in response.iter_content(chunk_size=1024):
+            if chunk:
+                zip_file.write(chunk)
+                progress_bar.update(len(chunk))
+
+        progress_bar.close()
+
+        zip_file.seek(0)
+
+        with zipfile.ZipFile(zip_file, "r") as zip_ref:
+            zip_ref.extractall(destination_folder)
+        return True
+    else:
+        return False
+
+
+def download_pusht(root: str, dataset_id: str = "pusht", fps: int = 10) -> Path:
+    pusht_url = "https://diffusion-policy.cs.columbia.edu/data/training/pusht.zip"
+    pusht_zarr = Path("pusht/pusht_cchi_v7_replay.zarr")
+
+    root = Path(root)
+    raw_dir: Path = root / f"{dataset_id}_raw"
+    zarr_path: Path = (raw_dir / pusht_zarr).resolve()
+    if not zarr_path.is_dir():
+        raw_dir.mkdir(parents=True, exist_ok=True)
+        download_and_extract_zip(pusht_url, raw_dir)
+    return zarr_path
+
+
+def download_xarm(root: str, dataset_id: str, fps: int = 15) -> Path:
+    root = Path(root)
+    raw_dir: Path = root / "xarm_datasets_raw"
+    if not raw_dir.exists():
+        import zipfile
+
+        import gdown
+
+        raw_dir.mkdir(parents=True, exist_ok=True)
+        # from https://github.com/fyhMer/fowm/blob/main/scripts/download_datasets.py
+        url = "https://drive.google.com/uc?id=1nhxpykGtPDhmQKm-_B8zBSywVRdgeVya"
+        zip_path = raw_dir / "data.zip"
+        gdown.download(url, str(zip_path), quiet=False)
+        print("Extracting...")
+        with zipfile.ZipFile(str(zip_path), "r") as zip_f:
+            for member in zip_f.namelist():
+                if member.startswith("data/xarm") and member.endswith(".pkl"):
+                    print(member)
+                    zip_f.extract(member=member)
+        zip_path.unlink()
+
+    dataset_path: Path = root / f"{dataset_id}"
+    return dataset_path
+
+
+def download_aloha(root: str, dataset_id: str) -> Path:
+    folder_urls = {
+        "aloha_sim_insertion_human": "https://drive.google.com/drive/folders/1RgyD0JgTX30H4IM5XZn8I3zSV_mr8pyF",
+        "aloha_sim_insertion_scripted": "https://drive.google.com/drive/folders/1TsojQQSXtHEoGnqgJ3gmpPQR2DPLtS2N",
+        "aloha_sim_transfer_cube_human": "https://drive.google.com/drive/folders/1sc-E4QYW7A0o23m1u2VWNGVq5smAsfCo",
+        "aloha_sim_transfer_cube_scripted": "https://drive.google.com/drive/folders/1aRyoOhQwxhyt1J8XgEig4s6kzaw__LXj",
+    }
+
+    ep48_urls = {
+        "aloha_sim_insertion_human": "https://drive.google.com/file/d/18Cudl6nikDtgRolea7je8iF_gGKzynOP/view?usp=drive_link",
+        "aloha_sim_insertion_scripted": "https://drive.google.com/file/d/1wfMSZ24oOh5KR_0aaP3Cnu_c4ZCveduB/view?usp=drive_link",
+        "aloha_sim_transfer_cube_human": "https://drive.google.com/file/d/18smMymtr8tIxaNUQ61gW6dG50pt3MvGq/view?usp=drive_link",
+        "aloha_sim_transfer_cube_scripted": "https://drive.google.com/file/d/1pnGIOd-E4-rhz2P3VxpknMKRZCoKt6eI/view?usp=drive_link",
+    }
+
+    ep49_urls = {
+        "aloha_sim_insertion_human": "https://drive.google.com/file/d/1C1kZYyROzs-PrLc0SkDgUgMi4-L3lauE/view?usp=drive_link",
+        "aloha_sim_insertion_scripted": "https://drive.google.com/file/d/17EuCUWS6uCCr6yyNzpXdcdE-_TTNCKtf/view?usp=drive_link",
+        "aloha_sim_transfer_cube_human": "https://drive.google.com/file/d/1Nk7l53d9sJoGDBKAOnNrExX5nLacATc6/view?usp=drive_link",
+        "aloha_sim_transfer_cube_scripted": "https://drive.google.com/file/d/1GKReZHrXU73NMiC5zKCq_UtqPVtYq8eo/view?usp=drive_link",
+    }
+    num_episodes = {  # noqa: F841 # we keep this for reference
+        "aloha_sim_insertion_human": 50,
+        "aloha_sim_insertion_scripted": 50,
+        "aloha_sim_transfer_cube_human": 50,
+        "aloha_sim_transfer_cube_scripted": 50,
+    }
+
+    episode_len = {  # noqa: F841 # we keep this for reference
+        "aloha_sim_insertion_human": 500,
+        "aloha_sim_insertion_scripted": 400,
+        "aloha_sim_transfer_cube_human": 400,
+        "aloha_sim_transfer_cube_scripted": 400,
+    }
+
+    cameras = {  # noqa: F841 # we keep this for reference
+        "aloha_sim_insertion_human": ["top"],
+        "aloha_sim_insertion_scripted": ["top"],
+        "aloha_sim_transfer_cube_human": ["top"],
+        "aloha_sim_transfer_cube_scripted": ["top"],
+    }
+    root = Path(root)
+    raw_dir: Path = root / f"{dataset_id}_raw"
+    if not raw_dir.is_dir():
+        import gdown
+
+        assert dataset_id in folder_urls
+        assert dataset_id in ep48_urls
+        assert dataset_id in ep49_urls
+
+        raw_dir.mkdir(parents=True, exist_ok=True)
+
+        gdown.download_folder(folder_urls[dataset_id], output=str(raw_dir))
+
+        # because of the 50 files limit per directory, two files episode 48 and 49 were missing
+        gdown.download(ep48_urls[dataset_id], output=str(raw_dir / "episode_48.hdf5"), fuzzy=True)
+        gdown.download(ep49_urls[dataset_id], output=str(raw_dir / "episode_49.hdf5"), fuzzy=True)
+    return raw_dir
+
+
+def download_umi(root: str, dataset_id: str) -> Path:
+    url_cup_in_the_wild = "https://real.stanford.edu/umi/data/zarr_datasets/cup_in_the_wild.zarr.zip"
+    cup_in_the_wild_zarr = Path("umi/cup_in_the_wild/cup_in_the_wild.zarr")
+
+    root = Path(root)
+    raw_dir: Path = root / f"{dataset_id}_raw"
+    zarr_path: Path = (raw_dir / cup_in_the_wild_zarr).resolve()
+    if not zarr_path.is_dir():
+        raw_dir.mkdir(parents=True, exist_ok=True)
+        download_and_extract_zip(url_cup_in_the_wild, zarr_path)
+    return zarr_path
+
+
+if __name__ == "__main__":
+    root = "data"
+    dataset_ids = [
+        "pusht",
+        "xarm_lift_medium",
+        "xarm_lift_medium_replay",
+        "xarm_push_medium",
+        "xarm_push_medium_replay",
+        "aloha_sim_insertion_human",
+        "aloha_sim_insertion_scripted",
+        "aloha_sim_transfer_cube_human",
+        "aloha_sim_transfer_cube_scripted",
+        "umi_cup_in_the_wild",
+    ]
+    for dataset_id in dataset_ids:
+        download_raw(root=root, dataset_id=dataset_id)

lerobot/common/datasets/push_dataset_to_hub/_umi_imagecodecs_numcodecs.py

@@ -0,0 +1,311 @@
+# imagecodecs/numcodecs.py
+
+# Copyright (c) 2021-2022, Christoph Gohlke
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice,
+#    this list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+#    contributors may be used to endorse or promote products derived from
+#    this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+
+# Copied from: https://github.com/real-stanford/universal_manipulation_interface/blob/298776ce251f33b6b3185a98d6e7d1f9ad49168b/diffusion_policy/codecs/imagecodecs_numcodecs.py#L1
+"""Additional numcodecs implemented using imagecodecs."""
+
+__version__ = "2022.9.26"
+
+__all__ = ("register_codecs",)
+
+import imagecodecs
+import numpy
+from numcodecs.abc import Codec
+from numcodecs.registry import get_codec, register_codec
+
+# TODO (azouitine): Remove useless codecs
+
+
+def protective_squeeze(x: numpy.ndarray):
+    """
+    Squeeze dim only if it's not the last dim.
+    Image dim expected to be *, H, W, C
+    """
+    img_shape = x.shape[-3:]
+    if len(x.shape) > 3:
+        n_imgs = numpy.prod(x.shape[:-3])
+        if n_imgs > 1:
+            img_shape = (-1,) + img_shape
+    return x.reshape(img_shape)
+
+
+def get_default_image_compressor(**kwargs):
+    if imagecodecs.JPEGXL:
+        # has JPEGXL
+        this_kwargs = {
+            "effort": 3,
+            "distance": 0.3,
+            # bug in libjxl, invalid codestream for non-lossless
+            # when decoding speed > 1
+            "decodingspeed": 1,
+        }
+        this_kwargs.update(kwargs)
+        return JpegXl(**this_kwargs)
+    else:
+        this_kwargs = {"level": 50}
+        this_kwargs.update(kwargs)
+        return Jpeg2k(**this_kwargs)
+
+
+class Jpeg2k(Codec):
+    """JPEG 2000 codec for numcodecs."""
+
+    codec_id = "imagecodecs_jpeg2k"
+
+    def __init__(
+        self,
+        level=None,
+        codecformat=None,
+        colorspace=None,
+        tile=None,
+        reversible=None,
+        bitspersample=None,
+        resolutions=None,
+        numthreads=None,
+        verbose=0,
+    ):
+        self.level = level
+        self.codecformat = codecformat
+        self.colorspace = colorspace
+        self.tile = None if tile is None else tuple(tile)
+        self.reversible = reversible
+        self.bitspersample = bitspersample
+        self.resolutions = resolutions
+        self.numthreads = numthreads
+        self.verbose = verbose
+
+    def encode(self, buf):
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.jpeg2k_encode(
+            buf,
+            level=self.level,
+            codecformat=self.codecformat,
+            colorspace=self.colorspace,
+            tile=self.tile,
+            reversible=self.reversible,
+            bitspersample=self.bitspersample,
+            resolutions=self.resolutions,
+            numthreads=self.numthreads,
+            verbose=self.verbose,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.jpeg2k_decode(buf, verbose=self.verbose, numthreads=self.numthreads, out=out)
+
+
+class JpegXl(Codec):
+    """JPEG XL codec for numcodecs."""
+
+    codec_id = "imagecodecs_jpegxl"
+
+    def __init__(
+        self,
+        # encode
+        level=None,
+        effort=None,
+        distance=None,
+        lossless=None,
+        decodingspeed=None,
+        photometric=None,
+        planar=None,
+        usecontainer=None,
+        # decode
+        index=None,
+        keeporientation=None,
+        # both
+        numthreads=None,
+    ):
+        """
+        Return JPEG XL image from numpy array.
+        Float must be in nominal range 0..1.
+
+        Currently L, LA, RGB, RGBA images are supported in contig mode.
+        Extra channels are only supported for grayscale images in planar mode.
+
+        Parameters
+        ----------
+        level : Default to None, i.e. not overwriting lossess and decodingspeed options.
+            When < 0: Use lossless compression
+            When in [0,1,2,3,4]: Sets the decoding speed tier for the provided options.
+                Minimum is 0 (slowest to decode, best quality/density), and maximum
+                is 4 (fastest to decode, at the cost of some quality/density).
+        effort : Default to 3.
+            Sets encoder effort/speed level without affecting decoding speed.
+            Valid values are, from faster to slower speed: 1:lightning 2:thunder
+                3:falcon 4:cheetah 5:hare 6:wombat 7:squirrel 8:kitten 9:tortoise.
+            Speed: lightning, thunder, falcon, cheetah, hare, wombat, squirrel, kitten, tortoise
+            control the encoder effort in ascending order.
+            This also affects memory usage: using lower effort will typically reduce memory
+            consumption during encoding.
+            lightning and thunder are fast modes useful for lossless mode (modular).
+            falcon disables all of the following tools.
+            cheetah enables coefficient reordering, context clustering, and heuristics for selecting DCT sizes and quantization steps.
+            hare enables Gaborish filtering, chroma from luma, and an initial estimate of quantization steps.
+            wombat enables error diffusion quantization and full DCT size selection heuristics.
+            squirrel (default) enables dots, patches, and spline detection, and full context clustering.
+            kitten optimizes the adaptive quantization for a psychovisual metric.
+            tortoise enables a more thorough adaptive quantization search.
+        distance : Default to 1.0
+            Sets the distance level for lossy compression: target max butteraugli distance,
+            lower = higher quality. Range: 0 .. 15. 0.0 = mathematically lossless
+            (however, use JxlEncoderSetFrameLossless instead to use true lossless,
+            as setting distance to 0 alone is not the only requirement).
+            1.0 = visually lossless. Recommended range: 0.5 .. 3.0.
+        lossess : Default to False.
+            Use lossess encoding.
+        decodingspeed : Default to 0.
+            Duplicate to level. [0,4]
+        photometric : Return JxlColorSpace value.
+            Default logic is quite complicated but works most of the time.
+            Accepted value:
+                int: [-1,3]
+                str: ['RGB',
+                    'WHITEISZERO', 'MINISWHITE',
+                    'BLACKISZERO', 'MINISBLACK', 'GRAY',
+                    'XYB', 'KNOWN']
+        planar : Enable multi-channel mode.
+            Default to false.
+        usecontainer :
+            Forces the encoder to use the box-based container format (BMFF)
+            even when not necessary.
+            When using JxlEncoderUseBoxes, JxlEncoderStoreJPEGMetadata or
+            JxlEncoderSetCodestreamLevel with level 10, the encoder will
+            automatically also use the container format, it is not necessary
+            to use JxlEncoderUseContainer for those use cases.
+            By default this setting is disabled.
+        index : Selectively decode frames for animation.
+            Default to 0, decode all frames.
+            When set to > 0, decode that frame index only.
+        keeporientation :
+            Enables or disables preserving of as-in-bitstream pixeldata orientation.
+            Some images are encoded with an Orientation tag indicating that the
+            decoder must perform a rotation and/or mirroring to the encoded image data.
+
+            If skip_reorientation is JXL_FALSE (the default): the decoder will apply
+            the transformation from the orientation setting, hence rendering the image
+            according to its specified intent. When producing a JxlBasicInfo, the decoder
+            will always set the orientation field to JXL_ORIENT_IDENTITY (matching the
+            returned pixel data) and also align xsize and ysize so that they correspond
+            to the width and the height of the returned pixel data.
+
+            If skip_reorientation is JXL_TRUE: the decoder will skip applying the
+            transformation from the orientation setting, returning the image in
+            the as-in-bitstream pixeldata orientation. This may be faster to decode
+            since the decoder doesnt have to apply the transformation, but can
+            cause wrong display of the image if the orientation tag is not correctly
+            taken into account by the user.
+
+            By default, this option is disabled, and the returned pixel data is
+            re-oriented according to the images Orientation setting.
+        threads : Default to 1.
+            If <= 0, use all cores.
+            If > 32, clipped to 32.
+        """
+
+        self.level = level
+        self.effort = effort
+        self.distance = distance
+        self.lossless = bool(lossless)
+        self.decodingspeed = decodingspeed
+        self.photometric = photometric
+        self.planar = planar
+        self.usecontainer = usecontainer
+        self.index = index
+        self.keeporientation = keeporientation
+        self.numthreads = numthreads
+
+    def encode(self, buf):
+        # TODO: only squeeze all but last dim
+        buf = protective_squeeze(numpy.asarray(buf))
+        return imagecodecs.jpegxl_encode(
+            buf,
+            level=self.level,
+            effort=self.effort,
+            distance=self.distance,
+            lossless=self.lossless,
+            decodingspeed=self.decodingspeed,
+            photometric=self.photometric,
+            planar=self.planar,
+            usecontainer=self.usecontainer,
+            numthreads=self.numthreads,
+        )
+
+    def decode(self, buf, out=None):
+        return imagecodecs.jpegxl_decode(
+            buf,
+            index=self.index,
+            keeporientation=self.keeporientation,
+            numthreads=self.numthreads,
+            out=out,
+        )
+
+
+def _flat(out):
+    """Return numpy array as contiguous view of bytes if possible."""
+    if out is None:
+        return None
+    view = memoryview(out)
+    if view.readonly or not view.contiguous:
+        return None
+    return view.cast("B")
+
+
+def register_codecs(codecs=None, force=False, verbose=True):
+    """Register codecs in this module with numcodecs."""
+    for name, cls in globals().items():
+        if not hasattr(cls, "codec_id") or name == "Codec":
+            continue
+        if codecs is not None and cls.codec_id not in codecs:
+            continue
+        try:
+            try:  # noqa: SIM105
+                get_codec({"id": cls.codec_id})
+            except TypeError:
+                # registered, but failed
+                pass
+        except ValueError:
+            # not registered yet
+            pass
+        else:
+            if not force:
+                if verbose:
+                    log_warning(f"numcodec {cls.codec_id!r} already registered")
+                continue
+            if verbose:
+                log_warning(f"replacing registered numcodec {cls.codec_id!r}")
+        register_codec(cls)
+
+
+def log_warning(msg, *args, **kwargs):
+    """Log message with level WARNING."""
+    import logging
+
+    logging.getLogger(__name__).warning(msg, *args, **kwargs)

lerobot/common/datasets/push_dataset_to_hub/aloha_processor.py

@@ -0,0 +1,199 @@
+import re
+from pathlib import Path
+
+import h5py
+import torch
+import tqdm
+from datasets import Dataset, Features, Image, Sequence, Value
+from PIL import Image as PILImage
+
+from lerobot.common.datasets.push_dataset_to_hub.utils import concatenate_episodes
+from lerobot.common.datasets.utils import (
+    hf_transform_to_torch,
+)
+
+
+class AlohaProcessor:
+    """
+    Process HDF5 files formatted like in: https://github.com/tonyzhaozh/act
+
+    Attributes:
+        folder_path (Path): Path to the directory containing HDF5 files.
+        cameras (list[str]): List of camera identifiers to check in the files.
+        fps (int): Frames per second used in timestamp calculations.
+
+    Methods:
+        is_valid() -> bool:
+            Validates if each HDF5 file within the folder contains all required datasets.
+        preprocess() -> dict:
+            Processes the files and returns structured data suitable for further analysis.
+        to_hf_dataset(data_dict: dict) -> Dataset:
+            Converts processed data into a Hugging Face Dataset object.
+    """
+
+    def __init__(self, folder_path: Path, cameras: list[str] | None = None, fps: int | None = None):
+        """
+        Initializes the AlohaProcessor with a specified directory path containing HDF5 files,
+        an optional list of cameras, and a frame rate.
+
+        Args:
+            folder_path (Path): The directory path where HDF5 files are stored.
+            cameras (list[str] | None): Optional list of cameras to validate within the files. Defaults to ['top'] if None.
+            fps (int): Frame rate for the datasets, used in time calculations. Default is 50.
+
+        Examples:
+            >>> processor = AlohaProcessor(Path("path_to_hdf5_directory"), ["camera1", "camera2"])
+            >>> processor.is_valid()
+            True
+        """
+        self.folder_path = folder_path
+        if cameras is None:
+            cameras = ["top"]
+        self.cameras = cameras
+        if fps is None:
+            fps = 50
+        self._fps = fps
+
+    @property
+    def fps(self) -> int:
+        return self._fps
+
+    def is_valid(self) -> bool:
+        """
+        Validates the HDF5 files in the specified folder to ensure they contain the required datasets
+        for actions, positions, and images for each specified camera.
+
+        Returns:
+            bool: True if all files are valid HDF5 files with all required datasets, False otherwise.
+        """
+        hdf5_files: list[Path] = list(self.folder_path.glob("episode_*.hdf5"))
+        if len(hdf5_files) == 0:
+            return False
+        try:
+            hdf5_files = sorted(
+                hdf5_files, key=lambda x: int(re.search(r"episode_(\d+).hdf5", x.name).group(1))
+            )
+        except AttributeError:
+            # All file names must contain a numerical identifier matching 'episode_(\\d+).hdf5
+            return False
+
+        # Check if the sequence is consecutive eg episode_0, episode_1, episode_2, etc.
+        # If not, return False
+        previous_number = None
+        for file in hdf5_files:
+            current_number = int(re.search(r"episode_(\d+).hdf5", file.name).group(1))
+            if previous_number is not None and current_number - previous_number != 1:
+                return False
+            previous_number = current_number
+
+        for file in hdf5_files:
+            try:
+                with h5py.File(file, "r") as file:
+                    # Check for the expected datasets within the HDF5 file
+                    required_datasets = ["/action", "/observations/qpos"]
+                    # Add camera-specific image datasets to the required datasets
+                    camera_datasets = [f"/observations/images/{cam}" for cam in self.cameras]
+                    required_datasets.extend(camera_datasets)
+
+                    if not all(dataset in file for dataset in required_datasets):
+                        return False
+            except OSError:
+                return False
+        return True
+
+    def preprocess(self):
+        """
+        Collects episode data from the HDF5 file and returns it as an AlohaStep named tuple.
+
+        Returns:
+            AlohaStep: Named tuple containing episode data.
+
+        Raises:
+            ValueError: If the file is not valid.
+        """
+        if not self.is_valid():
+            raise ValueError("The HDF5 file is invalid or does not contain the required datasets.")
+
+        hdf5_files = list(self.folder_path.glob("*.hdf5"))
+        hdf5_files = sorted(hdf5_files, key=lambda x: int(re.search(r"episode_(\d+)", x.name).group(1)))
+        ep_dicts = []
+        episode_data_index = {"from": [], "to": []}
+
+        id_from = 0
+
+        for ep_path in tqdm.tqdm(hdf5_files):
+            with h5py.File(ep_path, "r") as ep:
+                ep_id = int(re.search(r"episode_(\d+)", ep_path.name).group(1))
+                num_frames = ep["/action"].shape[0]
+
+                # last step of demonstration is considered done
+                done = torch.zeros(num_frames, dtype=torch.bool)
+                done[-1] = True
+
+                state = torch.from_numpy(ep["/observations/qpos"][:])
+                action = torch.from_numpy(ep["/action"][:])
+
+                ep_dict = {}
+
+                for cam in self.cameras:
+                    image = torch.from_numpy(ep[f"/observations/images/{cam}"][:])  # b h w c
+                    ep_dict[f"observation.images.{cam}"] = [PILImage.fromarray(x.numpy()) for x in image]
+
+                ep_dict.update(
+                    {
+                        "observation.state": state,
+                        "action": action,
+                        "episode_index": torch.tensor([ep_id] * num_frames),
+                        "frame_index": torch.arange(0, num_frames, 1),
+                        "timestamp": torch.arange(0, num_frames, 1) / self.fps,
+                        # TODO(rcadene): compute reward and success
+                        # "next.reward": reward,
+                        "next.done": done,
+                        # "next.success": success,
+                    }
+                )
+
+                assert isinstance(ep_id, int)
+                ep_dicts.append(ep_dict)
+
+                episode_data_index["from"].append(id_from)
+                episode_data_index["to"].append(id_from + num_frames)
+
+            id_from += num_frames
+
+        data_dict = concatenate_episodes(ep_dicts)
+        return data_dict, episode_data_index
+
+    def to_hf_dataset(self, data_dict) -> Dataset:
+        """
+        Converts a dictionary of data into a Hugging Face Dataset object.
+
+        Args:
+            data_dict (dict): A dictionary containing the data to be converted.
+
+        Returns:
+            Dataset: The converted Hugging Face Dataset object.
+        """
+        image_features = {f"observation.images.{cam}": Image() for cam in self.cameras}
+        features = {
+            "observation.state": Sequence(
+                length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None)
+            ),
+            "action": Sequence(length=data_dict["action"].shape[1], feature=Value(dtype="float32", id=None)),
+            "episode_index": Value(dtype="int64", id=None),
+            "frame_index": Value(dtype="int64", id=None),
+            "timestamp": Value(dtype="float32", id=None),
+            # "next.reward": Value(dtype="float32", id=None),
+            "next.done": Value(dtype="bool", id=None),
+            # "next.success": Value(dtype="bool", id=None),
+            "index": Value(dtype="int64", id=None),
+        }
+        update_features = {**image_features, **features}
+        features = Features(update_features)
+        hf_dataset = Dataset.from_dict(data_dict, features=features)
+        hf_dataset.set_transform(hf_transform_to_torch)
+
+        return hf_dataset
+
+    def cleanup(self):
+        pass

lerobot/common/datasets/push_dataset_to_hub/pusht_processor.py

@@ -0,0 +1,180 @@
+from pathlib import Path
+
+import numpy as np
+import torch
+import tqdm
+import zarr
+from datasets import Dataset, Features, Image, Sequence, Value
+from PIL import Image as PILImage
+
+from lerobot.common.datasets.push_dataset_to_hub.utils import concatenate_episodes
+from lerobot.common.datasets.utils import (
+    hf_transform_to_torch,
+)
+
+
+class PushTProcessor:
+    """ Process zarr files formatted like in: https://github.com/real-stanford/diffusion_policy
+    """
+    def __init__(self, folder_path: Path, fps: int | None = None):
+        self.zarr_path = folder_path
+        if fps is None:
+            fps = 10
+        self._fps = fps
+
+    @property
+    def fps(self) -> int:
+        return self._fps
+
+    def is_valid(self):
+        try:
+            zarr_data = zarr.open(self.zarr_path, mode="r")
+        except Exception:
+            # TODO (azouitine): Handle the exception properly
+            return False
+        required_datasets = {
+            "data/action",
+            "data/img",
+            "data/keypoint",
+            "data/n_contacts",
+            "data/state",
+            "meta/episode_ends",
+        }
+        for dataset in required_datasets:
+            if dataset not in zarr_data:
+                return False
+        nb_frames = zarr_data["data/img"].shape[0]
+
+        required_datasets.remove("meta/episode_ends")
+
+        return all(nb_frames == zarr_data[dataset].shape[0] for dataset in required_datasets)
+
+    def preprocess(self):
+        try:
+            import pymunk
+            from gym_pusht.envs.pusht import PushTEnv, pymunk_to_shapely
+
+            from lerobot.common.datasets.push_dataset_to_hub._diffusion_policy_replay_buffer import (
+                ReplayBuffer as DiffusionPolicyReplayBuffer,
+            )
+        except ModuleNotFoundError as e:
+            print("`gym_pusht` is not installed. Please install it with `pip install 'lerobot[gym_pusht]'`")
+            raise e
+
+        # as define in env
+        success_threshold = 0.95  # 95% coverage,
+
+        dataset_dict = DiffusionPolicyReplayBuffer.copy_from_path(
+            self.zarr_path
+        )  # , keys=['img', 'state', 'action'])
+
+        episode_ids = torch.from_numpy(dataset_dict.get_episode_idxs())
+        num_episodes = dataset_dict.meta["episode_ends"].shape[0]
+        assert len(
+            {dataset_dict[key].shape[0] for key in dataset_dict.keys()}  # noqa: SIM118
+        ), "Some data type dont have the same number of total frames."
+
+        # TODO: verify that goal pose is expected to be fixed
+        goal_pos_angle = np.array([256, 256, np.pi / 4])  # x, y, theta (in radians)
+        goal_body = PushTEnv.get_goal_pose_body(goal_pos_angle)
+
+        imgs = torch.from_numpy(dataset_dict["img"])  # b h w c
+        states = torch.from_numpy(dataset_dict["state"])
+        actions = torch.from_numpy(dataset_dict["action"])
+
+        ep_dicts = []
+        episode_data_index = {"from": [], "to": []}
+
+        id_from = 0
+        for episode_id in tqdm.tqdm(range(num_episodes)):
+            id_to = dataset_dict.meta["episode_ends"][episode_id]
+
+            num_frames = id_to - id_from
+
+            assert (episode_ids[id_from:id_to] == episode_id).all()
+
+            image = imgs[id_from:id_to]
+            assert image.min() >= 0.0
+            assert image.max() <= 255.0
+            image = image.type(torch.uint8)
+
+            state = states[id_from:id_to]
+            agent_pos = state[:, :2]
+            block_pos = state[:, 2:4]
+            block_angle = state[:, 4]
+
+            reward = torch.zeros(num_frames)
+            success = torch.zeros(num_frames, dtype=torch.bool)
+            done = torch.zeros(num_frames, dtype=torch.bool)
+            for i in range(num_frames):
+                space = pymunk.Space()
+                space.gravity = 0, 0
+                space.damping = 0
+
+                # Add walls.
+                walls = [
+                    PushTEnv.add_segment(space, (5, 506), (5, 5), 2),
+                    PushTEnv.add_segment(space, (5, 5), (506, 5), 2),
+                    PushTEnv.add_segment(space, (506, 5), (506, 506), 2),
+                    PushTEnv.add_segment(space, (5, 506), (506, 506), 2),
+                ]
+                space.add(*walls)
+
+                block_body = PushTEnv.add_tee(space, block_pos[i].tolist(), block_angle[i].item())
+                goal_geom = pymunk_to_shapely(goal_body, block_body.shapes)
+                block_geom = pymunk_to_shapely(block_body, block_body.shapes)
+                intersection_area = goal_geom.intersection(block_geom).area
+                goal_area = goal_geom.area
+                coverage = intersection_area / goal_area
+                reward[i] = np.clip(coverage / success_threshold, 0, 1)
+                success[i] = coverage > success_threshold
+
+            # last step of demonstration is considered done
+            done[-1] = True
+
+            ep_dict = {
+                "observation.image": [PILImage.fromarray(x.numpy()) for x in image],
+                "observation.state": agent_pos,
+                "action": actions[id_from:id_to],
+                "episode_index": torch.tensor([episode_id] * num_frames, dtype=torch.int),
+                "frame_index": torch.arange(0, num_frames, 1),
+                "timestamp": torch.arange(0, num_frames, 1) / self.fps,
+                # "next.observation.image": image[1:],
+                # "next.observation.state": agent_pos[1:],
+                # TODO(rcadene): verify that reward and done are aligned with image and agent_pos
+                "next.reward": torch.cat([reward[1:], reward[[-1]]]),
+                "next.done": torch.cat([done[1:], done[[-1]]]),
+                "next.success": torch.cat([success[1:], success[[-1]]]),
+            }
+            ep_dicts.append(ep_dict)
+
+            episode_data_index["from"].append(id_from)
+            episode_data_index["to"].append(id_from + num_frames)
+
+            id_from += num_frames
+
+        data_dict = concatenate_episodes(ep_dicts)
+        return data_dict, episode_data_index
+
+    def to_hf_dataset(self, data_dict):
+        features = {
+            "observation.image": Image(),
+            "observation.state": Sequence(
+                length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None)
+            ),
+            "action": Sequence(length=data_dict["action"].shape[1], feature=Value(dtype="float32", id=None)),
+            "episode_index": Value(dtype="int64", id=None),
+            "frame_index": Value(dtype="int64", id=None),
+            "timestamp": Value(dtype="float32", id=None),
+            "next.reward": Value(dtype="float32", id=None),
+            "next.done": Value(dtype="bool", id=None),
+            "next.success": Value(dtype="bool", id=None),
+            "index": Value(dtype="int64", id=None),
+        }
+        features = Features(features)
+        hf_dataset = Dataset.from_dict(data_dict, features=features)
+        hf_dataset.set_transform(hf_transform_to_torch)
+        return hf_dataset
+
+    def cleanup(self):
+        pass

lerobot/common/datasets/push_dataset_to_hub/umi_processor.py

@@ -0,0 +1,280 @@
+import os
+import re
+import shutil
+from glob import glob
+
+import numpy as np
+import torch
+import tqdm
+import zarr
+from datasets import Dataset, Features, Image, Sequence, Value
+from PIL import Image as PILImage
+
+from lerobot.common.datasets.push_dataset_to_hub._umi_imagecodecs_numcodecs import register_codecs
+from lerobot.common.datasets.push_dataset_to_hub.utils import concatenate_episodes
+from lerobot.common.datasets.utils import (
+    hf_transform_to_torch,
+)
+
+
+class UmiProcessor:
+    """
+     Process UMI (Universal Manipulation Interface) data stored in Zarr format like in: https://github.com/real-stanford/universal_manipulation_interface
+
+    Attributes:
+        folder_path (str): The path to the folder containing Zarr datasets.
+        fps (int): Frames per second, used to calculate timestamps for frames.
+
+    """
+
+    def __init__(self, folder_path: str, fps: int | None = None):
+        self.zarr_path = folder_path
+        if fps is None:
+            # TODO (azouitine): Add reference to the paper
+            fps = 15
+        self._fps = fps
+        register_codecs()
+
+    @property
+    def fps(self) -> int:
+        return self._fps
+
+    def is_valid(self) -> bool:
+        """
+        Validates the Zarr folder to ensure it contains all required datasets with consistent frame counts.
+
+        Returns:
+            bool: True if all required datasets are present and have consistent frame counts, False otherwise.
+        """
+        # Check if the Zarr folder is valid
+        try:
+            zarr_data = zarr.open(self.zarr_path, mode="r")
+        except Exception:
+            # TODO (azouitine): Handle the exception properly
+            return False
+        required_datasets = {
+            "data/robot0_demo_end_pose",
+            "data/robot0_demo_start_pose",
+            "data/robot0_eef_pos",
+            "data/robot0_eef_rot_axis_angle",
+            "data/robot0_gripper_width",
+            "meta/episode_ends",
+            "data/camera0_rgb",
+        }
+        for dataset in required_datasets:
+            if dataset not in zarr_data:
+                return False
+        nb_frames = zarr_data["data/camera0_rgb"].shape[0]
+
+        required_datasets.remove("meta/episode_ends")
+
+        return all(nb_frames == zarr_data[dataset].shape[0] for dataset in required_datasets)
+
+    def preprocess(self):
+        """
+        Collects and processes all episodes from the Zarr dataset into structured data dictionaries.
+
+        Returns:
+            Tuple[Dict, Dict]: A tuple containing the structured episode data and episode index mappings.
+        """
+        zarr_data = zarr.open(self.zarr_path, mode="r")
+
+        # We process the image data separately because it is too large to fit in memory
+        end_pose = torch.from_numpy(zarr_data["data/robot0_demo_end_pose"][:])
+        start_pos = torch.from_numpy(zarr_data["data/robot0_demo_start_pose"][:])
+        eff_pos = torch.from_numpy(zarr_data["data/robot0_eef_pos"][:])
+        eff_rot_axis_angle = torch.from_numpy(zarr_data["data/robot0_eef_rot_axis_angle"][:])
+        gripper_width = torch.from_numpy(zarr_data["data/robot0_gripper_width"][:])
+
+        states_pos = torch.cat([eff_pos, eff_rot_axis_angle], dim=1)
+        states = torch.cat([states_pos, gripper_width], dim=1)
+
+        episode_ends = zarr_data["meta/episode_ends"][:]
+        num_episodes: int = episode_ends.shape[0]
+
+        episode_ids = torch.from_numpy(self.get_episode_idxs(episode_ends))
+
+        # We convert it in torch tensor later because the jit function does not support torch tensors
+        episode_ends = torch.from_numpy(episode_ends)
+
+        ep_dicts = []
+        episode_data_index = {"from": [], "to": []}
+        id_from = 0
+
+        for episode_id in tqdm.tqdm(range(num_episodes)):
+            id_to = episode_ends[episode_id]
+
+            num_frames = id_to - id_from
+
+            assert (
+                episode_ids[id_from:id_to] == episode_id
+            ).all(), f"episode_ids[{id_from}:{id_to}] != {episode_id}"
+
+            state = states[id_from:id_to]
+            ep_dict = {
+                # observation.image will be filled later
+                "observation.state": state,
+                "episode_index": torch.tensor([episode_id] * num_frames, dtype=torch.int),
+                "frame_index": torch.arange(0, num_frames, 1),
+                "timestamp": torch.arange(0, num_frames, 1) / self.fps,
+                "episode_data_index_from": torch.tensor([id_from] * num_frames),
+                "episode_data_index_to": torch.tensor([id_from + num_frames] * num_frames),
+                "end_pose": end_pose[id_from:id_to],
+                "start_pos": start_pos[id_from:id_to],
+                "gripper_width": gripper_width[id_from:id_to],
+            }
+            ep_dicts.append(ep_dict)
+            episode_data_index["from"].append(id_from)
+            episode_data_index["to"].append(id_from + num_frames)
+            id_from += num_frames
+
+        data_dict = concatenate_episodes(ep_dicts)
+
+        total_frames = id_from
+        data_dict["index"] = torch.arange(0, total_frames, 1)
+
+        print("Saving images to disk in temporary folder...")
+        # datasets.Image() can take a list of paths to images, so we save the images to a temporary folder
+        # to avoid loading them all in memory
+        _save_images_concurrently(
+            data=zarr_data, image_key="data/camera0_rgb", folder_path="tmp_umi_images", max_workers=12
+        )
+        print("Saving images to disk in temporary folder... Done")
+
+        # Sort files by number eg. 1.png, 2.png, 3.png, 9.png, 10.png instead of 1.png, 10.png, 2.png, 3.png, 9.png
+        # to correctly match the images with the data
+        images_path = sorted(
+            glob("tmp_umi_images/*"), key=lambda x: int(re.search(r"(\d+)\.png$", x).group(1))
+        )
+        data_dict["observation.image"] = images_path
+        print("Images saved to disk, do not forget to delete the folder tmp_umi_images/")
+
+        # Cleanup
+        return data_dict, episode_data_index
+
+    def to_hf_dataset(self, data_dict):
+        """
+        Converts the processed data dictionary into a Hugging Face dataset with defined features.
+
+        Args:
+            data_dict (Dict): The data dictionary containing tensors and episode information.
+
+        Returns:
+            Dataset: A Hugging Face dataset constructed from the provided data dictionary.
+        """
+        features = {
+            "observation.image": Image(),
+            "observation.state": Sequence(
+                length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None)
+            ),
+            "episode_index": Value(dtype="int64", id=None),
+            "frame_index": Value(dtype="int64", id=None),
+            "timestamp": Value(dtype="float32", id=None),
+            "index": Value(dtype="int64", id=None),
+            "episode_data_index_from": Value(dtype="int64", id=None),
+            "episode_data_index_to": Value(dtype="int64", id=None),
+            # `start_pos` and `end_pos` respectively represent the positions of the end-effector
+            # at the beginning and the end of the episode.
+            # `gripper_width` indicates the distance between the grippers, and this value is included
+            # in the state vector, which comprises the concatenation of the end-effector position
+            # and gripper width.
+            "end_pose": Sequence(
+                length=data_dict["end_pose"].shape[1], feature=Value(dtype="float32", id=None)
+            ),
+            "start_pos": Sequence(
+                length=data_dict["start_pos"].shape[1], feature=Value(dtype="float32", id=None)
+            ),
+            "gripper_width": Sequence(
+                length=data_dict["gripper_width"].shape[1], feature=Value(dtype="float32", id=None)
+            ),
+        }
+        features = Features(features)
+        hf_dataset = Dataset.from_dict(data_dict, features=features)
+        hf_dataset.set_transform(hf_transform_to_torch)
+
+        return hf_dataset
+
+    def cleanup(self):
+        # Cleanup
+        if os.path.exists("tmp_umi_images"):
+            print("Removing temporary images folder")
+            shutil.rmtree("tmp_umi_images")
+            print("Cleanup done")
+
+    @classmethod
+    def get_episode_idxs(cls, episode_ends: np.ndarray) -> np.ndarray:
+        # Optimized and simplified version of this function: https://github.com/real-stanford/universal_manipulation_interface/blob/298776ce251f33b6b3185a98d6e7d1f9ad49168b/diffusion_policy/common/replay_buffer.py#L374
+        from numba import jit
+
+        @jit(nopython=True)
+        def _get_episode_idxs(episode_ends):
+            result = np.zeros((episode_ends[-1],), dtype=np.int64)
+            start_idx = 0
+            for episode_number, end_idx in enumerate(episode_ends):
+                result[start_idx:end_idx] = episode_number
+                start_idx = end_idx
+            return result
+
+        return _get_episode_idxs(episode_ends)
+
+
+def _clear_folder(folder_path: str):
+    """
+    Clears all the content of the specified folder. Creates the folder if it does not exist.
+
+    Args:
+    folder_path (str): Path to the folder to clear.
+
+    Examples:
+    >>> import os
+    >>> os.makedirs('example_folder', exist_ok=True)
+    >>> with open('example_folder/temp_file.txt', 'w') as f:
+    ...     f.write('example')
+    >>> clear_folder('example_folder')
+    >>> os.listdir('example_folder')
+    []
+    """
+    if os.path.exists(folder_path):
+        for filename in os.listdir(folder_path):
+            file_path = os.path.join(folder_path, filename)
+            try:
+                if os.path.isfile(file_path) or os.path.islink(file_path):
+                    os.unlink(file_path)
+                elif os.path.isdir(file_path):
+                    shutil.rmtree(file_path)
+            except Exception as e:
+                print(f"Failed to delete {file_path}. Reason: {e}")
+    else:
+        os.makedirs(folder_path)
+
+
+def _save_image(img_array: np.array, i: int, folder_path: str):
+    """
+    Saves a single image to the specified folder.
+
+    Args:
+    img_array (ndarray): The numpy array of the image.
+    i (int): Index of the image, used for naming.
+    folder_path (str): Path to the folder where the image will be saved.
+    """
+    img = PILImage.fromarray(img_array)
+    img_format = "PNG" if img_array.dtype == np.uint8 else "JPEG"
+    img.save(os.path.join(folder_path, f"{i}.{img_format.lower()}"), quality=100)
+
+
+def _save_images_concurrently(data: dict, image_key: str, folder_path: str, max_workers: int = 4):
+    from concurrent.futures import ThreadPoolExecutor
+
+    """
+    Saves images from the zarr_data to the specified folder using multithreading.
+
+    Args:
+    zarr_data (dict): A dictionary containing image data in an array format.
+    folder_path (str): Path to the folder where images will be saved.
+    max_workers (int): The maximum number of threads to use for saving images.
+    """
+    num_images = len(data["data/camera0_rgb"])
+    _clear_folder(folder_path)  # Clear or create folder first
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        [executor.submit(_save_image, data[image_key][i], i, folder_path) for i in range(num_images)]

lerobot/common/datasets/push_dataset_to_hub/utils.py

@@ -0,0 +1,20 @@
+import torch
+
+
+def concatenate_episodes(ep_dicts):
+    data_dict = {}
+
+    keys = ep_dicts[0].keys()
+    for key in keys:
+        if torch.is_tensor(ep_dicts[0][key][0]):
+            data_dict[key] = torch.cat([ep_dict[key] for ep_dict in ep_dicts])
+        else:
+            if key not in data_dict:
+                data_dict[key] = []
+            for ep_dict in ep_dicts:
+                for x in ep_dict[key]:
+                    data_dict[key].append(x)
+
+    total_frames = data_dict["frame_index"].shape[0]
+    data_dict["index"] = torch.arange(0, total_frames, 1)
+    return data_dict

lerobot/common/datasets/push_dataset_to_hub/xarm_processor.py

@@ -0,0 +1,145 @@
+import pickle
+from pathlib import Path
+
+import einops
+import torch
+import tqdm
+from datasets import Dataset, Features, Image, Sequence, Value
+from PIL import Image as PILImage
+
+from lerobot.common.datasets.push_dataset_to_hub.utils import concatenate_episodes
+from lerobot.common.datasets.utils import (
+    hf_transform_to_torch,
+)
+
+
+class XarmProcessor:
+    """Process pickle files formatted like in: https://github.com/fyhMer/fowm"""
+
+    def __init__(self, folder_path: str, fps: int | None = None):
+        self.folder_path = Path(folder_path)
+        self.keys = {"actions", "rewards", "dones", "masks"}
+        self.nested_keys = {"observations": {"rgb", "state"}, "next_observations": {"rgb", "state"}}
+        if fps is None:
+            fps = 15
+        self._fps = fps
+
+    @property
+    def fps(self) -> int:
+        return self._fps
+
+    def is_valid(self) -> bool:
+        # get all .pkl files
+        xarm_files = list(self.folder_path.glob("*.pkl"))
+        if len(xarm_files) != 1:
+            return False
+
+        try:
+            with open(xarm_files[0], "rb") as f:
+                dataset_dict = pickle.load(f)
+        except Exception:
+            return False
+
+        if not isinstance(dataset_dict, dict):
+            return False
+
+        if not all(k in dataset_dict for k in self.keys):
+            return False
+
+        # Check for consistent lengths in nested keys
+        try:
+            expected_len = len(dataset_dict["actions"])
+            if any(len(dataset_dict[key]) != expected_len for key in self.keys if key in dataset_dict):
+                return False
+
+            for key, subkeys in self.nested_keys.items():
+                nested_dict = dataset_dict.get(key, {})
+                if any(
+                    len(nested_dict[subkey]) != expected_len for subkey in subkeys if subkey in nested_dict
+                ):
+                    return False
+        except KeyError:  # If any expected key or subkey is missing
+            return False
+
+        return True  # All checks passed
+
+    def preprocess(self):
+        if not self.is_valid():
+            raise ValueError("The Xarm file is invalid or does not contain the required datasets.")
+
+        xarm_files = list(self.folder_path.glob("*.pkl"))
+
+        with open(xarm_files[0], "rb") as f:
+            dataset_dict = pickle.load(f)
+        ep_dicts = []
+        episode_data_index = {"from": [], "to": []}
+
+        id_from = 0
+        id_to = 0
+        episode_id = 0
+        total_frames = dataset_dict["actions"].shape[0]
+        for i in tqdm.tqdm(range(total_frames)):
+            id_to += 1
+
+            if not dataset_dict["dones"][i]:
+                continue
+
+            num_frames = id_to - id_from
+
+            image = torch.tensor(dataset_dict["observations"]["rgb"][id_from:id_to])
+            image = einops.rearrange(image, "b c h w -> b h w c")
+            state = torch.tensor(dataset_dict["observations"]["state"][id_from:id_to])
+            action = torch.tensor(dataset_dict["actions"][id_from:id_to])
+            # TODO(rcadene): we have a missing last frame which is the observation when the env is done
+            # it is critical to have this frame for tdmpc to predict a "done observation/state"
+            # next_image = torch.tensor(dataset_dict["next_observations"]["rgb"][id_from:id_to])
+            # next_state = torch.tensor(dataset_dict["next_observations"]["state"][id_from:id_to])
+            next_reward = torch.tensor(dataset_dict["rewards"][id_from:id_to])
+            next_done = torch.tensor(dataset_dict["dones"][id_from:id_to])
+
+            ep_dict = {
+                "observation.image": [PILImage.fromarray(x.numpy()) for x in image],
+                "observation.state": state,
+                "action": action,
+                "episode_index": torch.tensor([episode_id] * num_frames, dtype=torch.int),
+                "frame_index": torch.arange(0, num_frames, 1),
+                "timestamp": torch.arange(0, num_frames, 1) / self.fps,
+                # "next.observation.image": next_image,
+                # "next.observation.state": next_state,
+                "next.reward": next_reward,
+                "next.done": next_done,
+            }
+            ep_dicts.append(ep_dict)
+
+            episode_data_index["from"].append(id_from)
+            episode_data_index["to"].append(id_from + num_frames)
+
+            id_from = id_to
+            episode_id += 1
+
+        data_dict = concatenate_episodes(ep_dicts)
+        return data_dict, episode_data_index
+
+    def to_hf_dataset(self, data_dict):
+        features = {
+            "observation.image": Image(),
+            "observation.state": Sequence(
+                length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None)
+            ),
+            "action": Sequence(length=data_dict["action"].shape[1], feature=Value(dtype="float32", id=None)),
+            "episode_index": Value(dtype="int64", id=None),
+            "frame_index": Value(dtype="int64", id=None),
+            "timestamp": Value(dtype="float32", id=None),
+            "next.reward": Value(dtype="float32", id=None),
+            "next.done": Value(dtype="bool", id=None),
+            #'next.success': Value(dtype='bool', id=None),
+            "index": Value(dtype="int64", id=None),
+        }
+        features = Features(features)
+        hf_dataset = Dataset.from_dict(data_dict, features=features)
+        hf_dataset.set_transform(hf_transform_to_torch)
+
+        return hf_dataset
+
+    def cleanup(self):
+        pass

lerobot/common/datasets/pusht.py

@@ -1,223 +0,0 @@
-from pathlib import Path
-from typing import Callable
-
-import einops
-import numpy as np
-import pygame
-import pymunk
-import torch
-import torchrl
-import tqdm
-from tensordict import TensorDict
-from torchrl.data.replay_buffers.samplers import Sampler
-from torchrl.data.replay_buffers.storages import TensorStorage
-from torchrl.data.replay_buffers.writers import Writer
-
-from lerobot.common.datasets.abstract import AbstractDataset
-from lerobot.common.datasets.utils import download_and_extract_zip
-from lerobot.common.envs.pusht.pusht_env import pymunk_to_shapely
-from lerobot.common.policies.diffusion.replay_buffer import ReplayBuffer as DiffusionPolicyReplayBuffer
-
-# as define in env
-SUCCESS_THRESHOLD = 0.95  # 95% coverage,
-
-PUSHT_URL = "https://diffusion-policy.cs.columbia.edu/data/training/pusht.zip"
-PUSHT_ZARR = Path("pusht/pusht_cchi_v7_replay.zarr")
-
-
-def get_goal_pose_body(pose):
-    mass = 1
-    inertia = pymunk.moment_for_box(mass, (50, 100))
-    body = pymunk.Body(mass, inertia)
-    # preserving the legacy assignment order for compatibility
-    # the order here doesn't matter somehow, maybe because CoM is aligned with body origin
-    body.position = pose[:2].tolist()
-    body.angle = pose[2]
-    return body
-
-
-def add_segment(space, a, b, radius):
-    shape = pymunk.Segment(space.static_body, a, b, radius)
-    shape.color = pygame.Color("LightGray")  # https://htmlcolorcodes.com/color-names
-    return shape
-
-
-def add_tee(
-    space,
-    position,
-    angle,
-    scale=30,
-    color="LightSlateGray",
-    mask=None,
-):
-    if mask is None:
-        mask = pymunk.ShapeFilter.ALL_MASKS()
-    mass = 1
-    length = 4
-    vertices1 = [
-        (-length * scale / 2, scale),
-        (length * scale / 2, scale),
-        (length * scale / 2, 0),
-        (-length * scale / 2, 0),
-    ]
-    inertia1 = pymunk.moment_for_poly(mass, vertices=vertices1)
-    vertices2 = [
-        (-scale / 2, scale),
-        (-scale / 2, length * scale),
-        (scale / 2, length * scale),
-        (scale / 2, scale),
-    ]
-    inertia2 = pymunk.moment_for_poly(mass, vertices=vertices1)
-    body = pymunk.Body(mass, inertia1 + inertia2)
-    shape1 = pymunk.Poly(body, vertices1)
-    shape2 = pymunk.Poly(body, vertices2)
-    shape1.color = pygame.Color(color)
-    shape2.color = pygame.Color(color)
-    shape1.filter = pymunk.ShapeFilter(mask=mask)
-    shape2.filter = pymunk.ShapeFilter(mask=mask)
-    body.center_of_gravity = (shape1.center_of_gravity + shape2.center_of_gravity) / 2
-    body.position = position
-    body.angle = angle
-    body.friction = 1
-    space.add(body, shape1, shape2)
-    return body
-
-
-class PushtDataset(AbstractDataset):
-    available_datasets = ["pusht"]
-
-    def __init__(
-        self,
-        dataset_id: str,
-        version: str | None = "v1.2",
-        batch_size: int | None = None,
-        *,
-        shuffle: bool = True,
-        root: Path | None = None,
-        pin_memory: bool = False,
-        prefetch: int = None,
-        sampler: Sampler | None = None,
-        collate_fn: Callable | None = None,
-        writer: Writer | None = None,
-        transform: "torchrl.envs.Transform" = None,
-    ):
-        super().__init__(
-            dataset_id,
-            version,
-            batch_size,
-            shuffle=shuffle,
-            root=root,
-            pin_memory=pin_memory,
-            prefetch=prefetch,
-            sampler=sampler,
-            collate_fn=collate_fn,
-            writer=writer,
-            transform=transform,
-        )
-
-    def _download_and_preproc_obsolete(self):
-        assert self.root is not None
-        raw_dir = self.root / f"{self.dataset_id}_raw"
-        zarr_path = (raw_dir / PUSHT_ZARR).resolve()
-        if not zarr_path.is_dir():
-            raw_dir.mkdir(parents=True, exist_ok=True)
-            download_and_extract_zip(PUSHT_URL, raw_dir)
-
-        # load
-        dataset_dict = DiffusionPolicyReplayBuffer.copy_from_path(
-            zarr_path
-        )  # , keys=['img', 'state', 'action'])
-
-        episode_ids = torch.from_numpy(dataset_dict.get_episode_idxs())
-        num_episodes = dataset_dict.meta["episode_ends"].shape[0]
-        total_frames = dataset_dict["action"].shape[0]
-        # to create test artifact
-        # num_episodes = 1
-        # total_frames = 50
-        assert len(
-            {dataset_dict[key].shape[0] for key in dataset_dict.keys()}  # noqa: SIM118
-        ), "Some data type dont have the same number of total frames."
-
-        # TODO: verify that goal pose is expected to be fixed
-        goal_pos_angle = np.array([256, 256, np.pi / 4])  # x, y, theta (in radians)
-        goal_body = get_goal_pose_body(goal_pos_angle)
-
-        imgs = torch.from_numpy(dataset_dict["img"])
-        imgs = einops.rearrange(imgs, "b h w c -> b c h w")
-        states = torch.from_numpy(dataset_dict["state"])
-        actions = torch.from_numpy(dataset_dict["action"])
-
-        idx0 = 0
-        idxtd = 0
-        for episode_id in tqdm.tqdm(range(num_episodes)):
-            idx1 = dataset_dict.meta["episode_ends"][episode_id]
-            # to create test artifact
-            # idx1 = 51
-
-            num_frames = idx1 - idx0
-
-            assert (episode_ids[idx0:idx1] == episode_id).all()
-
-            image = imgs[idx0:idx1]
-
-            state = states[idx0:idx1]
-            agent_pos = state[:, :2]
-            block_pos = state[:, 2:4]
-            block_angle = state[:, 4]
-
-            reward = torch.zeros(num_frames, 1)
-            success = torch.zeros(num_frames, 1, dtype=torch.bool)
-            done = torch.zeros(num_frames, 1, dtype=torch.bool)
-            for i in range(num_frames):
-                space = pymunk.Space()
-                space.gravity = 0, 0
-                space.damping = 0
-
-                # Add walls.
-                walls = [
-                    add_segment(space, (5, 506), (5, 5), 2),
-                    add_segment(space, (5, 5), (506, 5), 2),
-                    add_segment(space, (506, 5), (506, 506), 2),
-                    add_segment(space, (5, 506), (506, 506), 2),
-                ]
-                space.add(*walls)
-
-                block_body = add_tee(space, block_pos[i].tolist(), block_angle[i].item())
-                goal_geom = pymunk_to_shapely(goal_body, block_body.shapes)
-                block_geom = pymunk_to_shapely(block_body, block_body.shapes)
-                intersection_area = goal_geom.intersection(block_geom).area
-                goal_area = goal_geom.area
-                coverage = intersection_area / goal_area
-                reward[i] = np.clip(coverage / SUCCESS_THRESHOLD, 0, 1)
-                success[i] = coverage > SUCCESS_THRESHOLD
-
-            # last step of demonstration is considered done
-            done[-1] = True
-
-            ep_td = TensorDict(
-                {
-                    ("observation", "image"): image[:-1],
-                    ("observation", "state"): agent_pos[:-1],
-                    "action": actions[idx0:idx1][:-1],
-                    "episode": episode_ids[idx0:idx1][:-1],
-                    "frame_id": torch.arange(0, num_frames - 1, 1),
-                    ("next", "observation", "image"): image[1:],
-                    ("next", "observation", "state"): agent_pos[1:],
-                    # TODO: verify that reward and done are aligned with image and agent_pos
-                    ("next", "reward"): reward[1:],
-                    ("next", "done"): done[1:],
-                    ("next", "success"): success[1:],
-                },
-                batch_size=num_frames - 1,
-            )
-
-            if episode_id == 0:
-                # hack to initialize tensordict data structure to store episodes
-                td_data = ep_td[0].expand(total_frames).memmap_like(self.root / f"{self.dataset_id}")
-
-            td_data[idxtd : idxtd + len(ep_td)] = ep_td
-
-            idx0 = idx1
-            idxtd = idxtd + len(ep_td)
-
-        return TensorStorage(td_data.lock_())

lerobot/common/datasets/simxarm.py

@@ -1,127 +0,0 @@
-import pickle
-import zipfile
-from pathlib import Path
-from typing import Callable
-
-import torch
-import torchrl
-import tqdm
-from tensordict import TensorDict
-from torchrl.data.replay_buffers.samplers import (
-    Sampler,
-)
-from torchrl.data.replay_buffers.storages import TensorStorage
-from torchrl.data.replay_buffers.writers import Writer
-
-from lerobot.common.datasets.abstract import AbstractDataset
-
-
-def download():
-    raise NotImplementedError()
-    import gdown
-
-    url = "https://drive.google.com/uc?id=1nhxpykGtPDhmQKm-_B8zBSywVRdgeVya"
-    download_path = "data.zip"
-    gdown.download(url, download_path, quiet=False)
-    print("Extracting...")
-    with zipfile.ZipFile(download_path, "r") as zip_f:
-        for member in zip_f.namelist():
-            if member.startswith("data/xarm") and member.endswith(".pkl"):
-                print(member)
-                zip_f.extract(member=member)
-    Path(download_path).unlink()
-
-
-class SimxarmDataset(AbstractDataset):
-    available_datasets = [
-        "xarm_lift_medium",
-    ]
-
-    def __init__(
-        self,
-        dataset_id: str,
-        version: str | None = "v1.1",
-        batch_size: int | None = None,
-        *,
-        shuffle: bool = True,
-        root: Path | None = None,
-        pin_memory: bool = False,
-        prefetch: int = None,
-        sampler: Sampler | None = None,
-        collate_fn: Callable | None = None,
-        writer: Writer | None = None,
-        transform: "torchrl.envs.Transform" = None,
-    ):
-        super().__init__(
-            dataset_id,
-            version,
-            batch_size,
-            shuffle=shuffle,
-            root=root,
-            pin_memory=pin_memory,
-            prefetch=prefetch,
-            sampler=sampler,
-            collate_fn=collate_fn,
-            writer=writer,
-            transform=transform,
-        )
-
-    def _download_and_preproc_obsolete(self):
-        # assert self.root is not None
-        # TODO(rcadene): finish download
-        # download()
-
-        dataset_path = self.root / f"{self.dataset_id}" / "buffer.pkl"
-        print(f"Using offline dataset '{dataset_path}'")
-        with open(dataset_path, "rb") as f:
-            dataset_dict = pickle.load(f)
-
-        total_frames = dataset_dict["actions"].shape[0]
-
-        idx0 = 0
-        idx1 = 0
-        episode_id = 0
-        for i in tqdm.tqdm(range(total_frames)):
-            idx1 += 1
-
-            if not dataset_dict["dones"][i]:
-                continue
-
-            num_frames = idx1 - idx0
-
-            image = torch.tensor(dataset_dict["observations"]["rgb"][idx0:idx1])
-            state = torch.tensor(dataset_dict["observations"]["state"][idx0:idx1])
-            next_image = torch.tensor(dataset_dict["next_observations"]["rgb"][idx0:idx1])
-            next_state = torch.tensor(dataset_dict["next_observations"]["state"][idx0:idx1])
-            next_reward = torch.tensor(dataset_dict["rewards"][idx0:idx1])
-            next_done = torch.tensor(dataset_dict["dones"][idx0:idx1])
-
-            episode = TensorDict(
-                {
-                    ("observation", "image"): image,
-                    ("observation", "state"): state,
-                    "action": torch.tensor(dataset_dict["actions"][idx0:idx1]),
-                    "episode": torch.tensor([episode_id] * num_frames, dtype=torch.int),
-                    "frame_id": torch.arange(0, num_frames, 1),
-                    ("next", "observation", "image"): next_image,
-                    ("next", "observation", "state"): next_state,
-                    ("next", "reward"): next_reward,
-                    ("next", "done"): next_done,
-                },
-                batch_size=num_frames,
-            )
-
-            if episode_id == 0:
-                # hack to initialize tensordict data structure to store episodes
-                td_data = (
-                    episode[0]
-                    .expand(total_frames)
-                    .memmap_like(self.root / f"{self.dataset_id}" / "replay_buffer")
-                )
-
-            td_data[idx0:idx1] = episode
-
-            episode_id += 1
-            idx0 = idx1
-
-        return TensorStorage(td_data.lock_())

lerobot/common/datasets/utils.py

@@ -1,30 +1,358 @@
-import io
-import zipfile
+import json
+from copy import deepcopy
+from math import ceil
 from pathlib import Path
 
-import requests
+import datasets
+import einops
+import torch
 import tqdm
+from datasets import Image, load_dataset, load_from_disk
+from huggingface_hub import hf_hub_download
+from PIL import Image as PILImage
+from safetensors.torch import load_file
+from torchvision import transforms
 
 
-def download_and_extract_zip(url: str, destination_folder: Path) -> bool:
-    print(f"downloading from {url}")
-    response = requests.get(url, stream=True)
-    if response.status_code == 200:
-        total_size = int(response.headers.get("content-length", 0))
-        progress_bar = tqdm.tqdm(total=total_size, unit="B", unit_scale=True)
+def flatten_dict(d, parent_key="", sep="/"):
+    """Flatten a nested dictionary structure by collapsing nested keys into one key with a separator.
 
-        zip_file = io.BytesIO()
-        for chunk in response.iter_content(chunk_size=1024):
-            if chunk:
-                zip_file.write(chunk)
-                progress_bar.update(len(chunk))
+    For example:
+    ```
+    >>> dct = {"a": {"b": 1, "c": {"d": 2}}, "e": 3}`
+    >>> print(flatten_dict(dct))
+    {"a/b": 1, "a/c/d": 2, "e": 3}
+    """
+    items = []
+    for k, v in d.items():
+        new_key = f"{parent_key}{sep}{k}" if parent_key else k
+        if isinstance(v, dict):
+            items.extend(flatten_dict(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)
 
-        progress_bar.close()
 
-        zip_file.seek(0)
+def unflatten_dict(d, sep="/"):
+    outdict = {}
+    for key, value in d.items():
+        parts = key.split(sep)
+        d = outdict
+        for part in parts[:-1]:
+            if part not in d:
+                d[part] = {}
+            d = d[part]
+        d[parts[-1]] = value
+    return outdict
 
-        with zipfile.ZipFile(zip_file, "r") as zip_ref:
-            zip_ref.extractall(destination_folder)
-        return True
+
+def hf_transform_to_torch(items_dict):
+    """Get a transform function that convert items from Hugging Face dataset (pyarrow)
+    to torch tensors. Importantly, images are converted from PIL, which corresponds to
+    a channel last representation (h w c) of uint8 type, to a torch image representation
+    with channel first (c h w) of float32 type in range [0,1].
+    """
+    for key in items_dict:
+        first_item = items_dict[key][0]
+        if isinstance(first_item, PILImage.Image):
+            to_tensor = transforms.ToTensor()
+            items_dict[key] = [to_tensor(img) for img in items_dict[key]]
+        else:
+            items_dict[key] = [torch.tensor(x) for x in items_dict[key]]
+    return items_dict
+
+
+def load_hf_dataset(repo_id, version, root, split) -> datasets.Dataset:
+    """hf_dataset contains all the observations, states, actions, rewards, etc."""
+    if root is not None:
+        hf_dataset = load_from_disk(str(Path(root) / repo_id / split))
     else:
-        return False
+        hf_dataset = load_dataset(repo_id, revision=version, split=split)
+    hf_dataset.set_transform(hf_transform_to_torch)
+    return hf_dataset
+
+
+def load_episode_data_index(repo_id, version, root) -> dict[str, torch.Tensor]:
+    """episode_data_index contains the range of indices for each episode
+
+    Example:
+    ```python
+    from_id = episode_data_index["from"][episode_id].item()
+    to_id = episode_data_index["to"][episode_id].item()
+    episode_frames = [dataset[i] for i in range(from_id, to_id)]
+    ```
+    """
+    if root is not None:
+        path = Path(root) / repo_id / "meta_data" / "episode_data_index.safetensors"
+    else:
+        path = hf_hub_download(
+            repo_id, "meta_data/episode_data_index.safetensors", repo_type="dataset", revision=version
+        )
+
+    return load_file(path)
+
+
+def load_stats(repo_id, version, root) -> dict[str, dict[str, torch.Tensor]]:
+    """stats contains the statistics per modality computed over the full dataset, such as max, min, mean, std
+
+    Example:
+    ```python
+    normalized_action = (action - stats["action"]["mean"]) / stats["action"]["std"]
+    ```
+    """
+    if root is not None:
+        path = Path(root) / repo_id / "meta_data" / "stats.safetensors"
+    else:
+        path = hf_hub_download(repo_id, "meta_data/stats.safetensors", repo_type="dataset", revision=version)
+
+    stats = load_file(path)
+    return unflatten_dict(stats)
+
+
+def load_info(repo_id, version, root) -> dict:
+    """info contains useful information regarding the dataset that are not stored elsewhere
+
+    Example:
+    ```python
+    print("frame per second used to collect the video", info["fps"])
+    ```
+    """
+    if root is not None:
+        path = Path(root) / repo_id / "meta_data" / "info.json"
+    else:
+        path = hf_hub_download(repo_id, "meta_data/info.json", repo_type="dataset", revision=version)
+
+    with open(path) as f:
+        info = json.load(f)
+    return info
+
+
+def load_previous_and_future_frames(
+    item: dict[str, torch.Tensor],
+    hf_dataset: datasets.Dataset,
+    episode_data_index: dict[str, torch.Tensor],
+    delta_timestamps: dict[str, list[float]],
+    tol: float,
+) -> dict[torch.Tensor]:
+    """
+    Given a current item in the dataset containing a timestamp (e.g. 0.6 seconds), and a list of time differences of
+    some modalities (e.g. delta_timestamps={"observation.image": [-0.8, -0.2, 0, 0.2]}), this function computes for each
+    given modality a list of query timestamps (e.g. [-0.2, 0.4, 0.6, 0.8]) and loads the closest frames in the dataset.
+
+    Importantly, when no frame can be found around a query timestamp within a specified tolerance window, this function
+    raises an AssertionError. When a timestamp is queried before the first available timestamp of the episode or after
+    the last available timestamp, the violation of the tolerance doesnt raise an AssertionError, and the function
+    populates a boolean array indicating which frames are outside of the episode range. For instance, this boolean array
+    is useful during batched training to not supervise actions associated to timestamps coming after the end of the
+    episode, or to pad the observations in a specific way. Note that by default the observation frames before the start
+    of the episode are the same as the first frame of the episode.
+
+    Parameters:
+    - item (dict): A dictionary containing all the data related to a frame. It is the result of `dataset[idx]`. Each key
+      corresponds to a different modality (e.g., "timestamp", "observation.image", "action").
+    - hf_dataset (datasets.Dataset): A dictionary containing the full dataset. Each key corresponds to a different
+      modality (e.g., "timestamp", "observation.image", "action").
+    - episode_data_index (dict): A dictionary containing two keys ("from" and "to") associated to dataset indices.
+      They indicate the start index and end index of each episode in the dataset.
+    - delta_timestamps (dict): A dictionary containing lists of delta timestamps for each possible modality to be
+      retrieved. These deltas are added to the item timestamp to form the query timestamps.
+    - tol (float, optional): The tolerance level used to determine if a data point is close enough to the query
+      timestamp by asserting `tol > difference`. It is suggested to set `tol` to a smaller value than the
+      smallest expected inter-frame period, but large enough to account for jitter.
+
+    Returns:
+    - The same item with the queried frames for each modality specified in delta_timestamps, with an additional key for
+      each modality (e.g. "observation.image_is_pad").
+
+    Raises:
+    - AssertionError: If any of the frames unexpectedly violate the tolerance level. This could indicate synchronization
+      issues with timestamps during data collection.
+    """
+    # get indices of the frames associated to the episode, and their timestamps
+    ep_id = item["episode_index"].item()
+    ep_data_id_from = episode_data_index["from"][ep_id].item()
+    ep_data_id_to = episode_data_index["to"][ep_id].item()
+    ep_data_ids = torch.arange(ep_data_id_from, ep_data_id_to, 1)
+
+    # load timestamps
+    ep_timestamps = hf_dataset.select_columns("timestamp")[ep_data_id_from:ep_data_id_to]["timestamp"]
+    ep_timestamps = torch.stack(ep_timestamps)
+
+    # we make the assumption that the timestamps are sorted
+    ep_first_ts = ep_timestamps[0]
+    ep_last_ts = ep_timestamps[-1]
+    current_ts = item["timestamp"].item()
+
+    for key in delta_timestamps:
+        # get timestamps used as query to retrieve data of previous/future frames
+        delta_ts = delta_timestamps[key]
+        query_ts = current_ts + torch.tensor(delta_ts)
+
+        # compute distances between each query timestamp and all timestamps of all the frames belonging to the episode
+        dist = torch.cdist(query_ts[:, None], ep_timestamps[:, None], p=1)
+        min_, argmin_ = dist.min(1)
+
+        # TODO(rcadene): synchronize timestamps + interpolation if needed
+
+        is_pad = min_ > tol
+
+        # check violated query timestamps are all outside the episode range
+        assert ((query_ts[is_pad] < ep_first_ts) | (ep_last_ts < query_ts[is_pad])).all(), (
+            f"One or several timestamps unexpectedly violate the tolerance ({min_} > {tol=}) inside episode range."
+            "This might be due to synchronization issues with timestamps during data collection."
+        )
+
+        # get dataset indices corresponding to frames to be loaded
+        data_ids = ep_data_ids[argmin_]
+
+        # load frames modality
+        item[key] = hf_dataset.select_columns(key)[data_ids][key]
+        item[key] = torch.stack(item[key])
+        item[f"{key}_is_pad"] = is_pad
+
+    return item
+
+
+def get_stats_einops_patterns(hf_dataset):
+    """These einops patterns will be used to aggregate batches and compute statistics.
+
+    Note: We assume the images of `hf_dataset` are in channel first format
+    """
+
+    dataloader = torch.utils.data.DataLoader(
+        hf_dataset,
+        num_workers=0,
+        batch_size=2,
+        shuffle=False,
+    )
+    batch = next(iter(dataloader))
+
+    stats_patterns = {}
+    for key, feats_type in hf_dataset.features.items():
+        # sanity check that tensors are not float64
+        assert batch[key].dtype != torch.float64
+
+        if isinstance(feats_type, Image):
+            # sanity check that images are channel first
+            _, c, h, w = batch[key].shape
+            assert c < h and c < w, f"expect channel first images, but instead {batch[key].shape}"
+
+            # sanity check that images are float32 in range [0,1]
+            assert batch[key].dtype == torch.float32, f"expect torch.float32, but instead {batch[key].dtype=}"
+            assert batch[key].max() <= 1, f"expect pixels lower than 1, but instead {batch[key].max()=}"
+            assert batch[key].min() >= 0, f"expect pixels greater than 1, but instead {batch[key].min()=}"
+
+            stats_patterns[key] = "b c h w -> c 1 1"
+        elif batch[key].ndim == 2:
+            stats_patterns[key] = "b c -> c "
+        elif batch[key].ndim == 1:
+            stats_patterns[key] = "b -> 1"
+        else:
+            raise ValueError(f"{key}, {feats_type}, {batch[key].shape}")
+
+    return stats_patterns
+
+
+def compute_stats(hf_dataset, batch_size=32, max_num_samples=None):
+    if max_num_samples is None:
+        max_num_samples = len(hf_dataset)
+
+    stats_patterns = get_stats_einops_patterns(hf_dataset)
+
+    # mean and std will be computed incrementally while max and min will track the running value.
+    mean, std, max, min = {}, {}, {}, {}
+    for key in stats_patterns:
+        mean[key] = torch.tensor(0.0).float()
+        std[key] = torch.tensor(0.0).float()
+        max[key] = torch.tensor(-float("inf")).float()
+        min[key] = torch.tensor(float("inf")).float()
+
+    def create_seeded_dataloader(hf_dataset, batch_size, seed):
+        generator = torch.Generator()
+        generator.manual_seed(seed)
+        dataloader = torch.utils.data.DataLoader(
+            hf_dataset,
+            num_workers=4,
+            batch_size=batch_size,
+            shuffle=True,
+            drop_last=False,
+            generator=generator,
+        )
+        return dataloader
+
+    # Note: Due to be refactored soon. The point of storing `first_batch` is to make sure we don't get
+    # surprises when rerunning the sampler.
+    first_batch = None
+    running_item_count = 0  # for online mean computation
+    dataloader = create_seeded_dataloader(hf_dataset, batch_size, seed=1337)
+    for i, batch in enumerate(
+        tqdm.tqdm(dataloader, total=ceil(max_num_samples / batch_size), desc="Compute mean, min, max")
+    ):
+        this_batch_size = len(batch["index"])
+        running_item_count += this_batch_size
+        if first_batch is None:
+            first_batch = deepcopy(batch)
+        for key, pattern in stats_patterns.items():
+            batch[key] = batch[key].float()
+            # Numerically stable update step for mean computation.
+            batch_mean = einops.reduce(batch[key], pattern, "mean")
+            # Hint: to update the mean we need x̄ₙ = (Nₙ₋₁x̄ₙ₋₁ + Bₙxₙ) / Nₙ, where the subscript represents
+            # the update step, N is the running item count, B is this batch size, x̄ is the running mean,
+            # and x is the current batch mean. Some rearrangement is then required to avoid risking
+            # numerical overflow. Another hint: Nₙ₋₁ = Nₙ - Bₙ. Rearrangement yields
+            # x̄ₙ = x̄ₙ₋₁ + Bₙ * (xₙ - x̄ₙ₋₁) / Nₙ
+            mean[key] = mean[key] + this_batch_size * (batch_mean - mean[key]) / running_item_count
+            max[key] = torch.maximum(max[key], einops.reduce(batch[key], pattern, "max"))
+            min[key] = torch.minimum(min[key], einops.reduce(batch[key], pattern, "min"))
+
+        if i == ceil(max_num_samples / batch_size) - 1:
+            break
+
+    first_batch_ = None
+    running_item_count = 0  # for online std computation
+    dataloader = create_seeded_dataloader(hf_dataset, batch_size, seed=1337)
+    for i, batch in enumerate(
+        tqdm.tqdm(dataloader, total=ceil(max_num_samples / batch_size), desc="Compute std")
+    ):
+        this_batch_size = len(batch["index"])
+        running_item_count += this_batch_size
+        # Sanity check to make sure the batches are still in the same order as before.
+        if first_batch_ is None:
+            first_batch_ = deepcopy(batch)
+            for key in stats_patterns:
+                assert torch.equal(first_batch_[key], first_batch[key])
+        for key, pattern in stats_patterns.items():
+            batch[key] = batch[key].float()
+            # Numerically stable update step for mean computation (where the mean is over squared
+            # residuals).See notes in the mean computation loop above.
+            batch_std = einops.reduce((batch[key] - mean[key]) ** 2, pattern, "mean")
+            std[key] = std[key] + this_batch_size * (batch_std - std[key]) / running_item_count
+
+        if i == ceil(max_num_samples / batch_size) - 1:
+            break
+
+    for key in stats_patterns:
+        std[key] = torch.sqrt(std[key])
+
+    stats = {}
+    for key in stats_patterns:
+        stats[key] = {
+            "mean": mean[key],
+            "std": std[key],
+            "max": max[key],
+            "min": min[key],
+        }
+    return stats
+
+
+def cycle(iterable):
+    """The equivalent of itertools.cycle, but safe for Pytorch dataloaders.
+
+    See https://github.com/pytorch/pytorch/issues/23900 for information on why itertools.cycle is not safe.
+    """
+    iterator = iter(iterable)
+    while True:
+        try:
+            yield next(iterator)
+        except StopIteration:
+            iterator = iter(iterable)

lerobot/common/envs/abstract.py

@@ -1,92 +0,0 @@
-from collections import deque
-from typing import Optional
-
-from tensordict import TensorDict
-from torchrl.envs import EnvBase
-
-from lerobot.common.utils import set_global_seed
-
-
-class AbstractEnv(EnvBase):
-    """
-    Note:
-        When implementing a concrete class (e.g. `AlohaDataset`, `PushtEnv`, `DiffusionPolicy`), you need to:
-            1. set the required class attributes:
-                - for classes inheriting from `AbstractDataset`: `available_datasets`
-                - for classes inheriting from `AbstractEnv`: `name`, `available_tasks`
-                - for classes inheriting from `AbstractPolicy`: `name`
-            2. update variables in `lerobot/__init__.py` (e.g. `available_envs`, `available_datasets_per_envs`, `available_policies`)
-            3. update variables in `tests/test_available.py` by importing your new class
-    """
-
-    name: str | None = None  # same name should be used to instantiate the environment in factory.py
-    available_tasks: list[str] | None = None  # for instance: sim_insertion, sim_transfer_cube, pusht, lift
-
-    def __init__(
-        self,
-        task,
-        frame_skip: int = 1,
-        from_pixels: bool = False,
-        pixels_only: bool = False,
-        image_size=None,
-        seed=1337,
-        device="cpu",
-        num_prev_obs=1,
-        num_prev_action=0,
-    ):
-        super().__init__(device=device, batch_size=[])
-        assert self.name is not None, "Subclasses of `AbstractEnv` should set the `name` class attribute."
-        assert (
-            self.available_tasks is not None
-        ), "Subclasses of `AbstractEnv` should set the `available_tasks` class attribute."
-        assert (
-            task in self.available_tasks
-        ), f"The provided task ({task}) is not on the list of available tasks {self.available_tasks}."
-
-        self.task = task
-        self.frame_skip = frame_skip
-        self.from_pixels = from_pixels
-        self.pixels_only = pixels_only
-        self.image_size = image_size
-        self.num_prev_obs = num_prev_obs
-        self.num_prev_action = num_prev_action
-
-        if pixels_only:
-            assert from_pixels
-        if from_pixels:
-            assert image_size
-
-        self._make_env()
-        self._make_spec()
-
-        # self._next_seed will be used for the next reset. It is recommended that when self.set_seed is called
-        # you store the return value in self._next_seed (it will be a new randomly generated seed).
-        self._next_seed = seed
-        # Don't store the result of this in self._next_seed, as we want to make sure that the first time
-        # self._reset is called, we use seed.
-        self.set_seed(seed)
-
-        if self.num_prev_obs > 0:
-            self._prev_obs_image_queue = deque(maxlen=self.num_prev_obs)
-            self._prev_obs_state_queue = deque(maxlen=self.num_prev_obs)
-        if self.num_prev_action > 0:
-            raise NotImplementedError()
-            # self._prev_action_queue = deque(maxlen=self.num_prev_action)
-
-    def render(self, mode="rgb_array", width=640, height=480):
-        raise NotImplementedError("Abstract method")
-
-    def _reset(self, tensordict: Optional[TensorDict] = None):
-        raise NotImplementedError("Abstract method")
-
-    def _step(self, tensordict: TensorDict):
-        raise NotImplementedError("Abstract method")
-
-    def _make_env(self):
-        raise NotImplementedError("Abstract method")
-
-    def _make_spec(self):
-        raise NotImplementedError("Abstract method")
-
-    def _set_seed(self, seed: Optional[int]):
-        set_global_seed(seed)

lerobot/common/envs/aloha/assets/bimanual_viperx_end_effector_insertion.xml

@@ -1,59 +0,0 @@
-<mujoco>
-    <include file="scene.xml"/>
-    <include file="vx300s_dependencies.xml"/>
-
-    <equality>
-        <weld body1="mocap_left" body2="vx300s_left/gripper_link" solref="0.01 1" solimp=".25 .25 0.001" />
-        <weld body1="mocap_right" body2="vx300s_right/gripper_link" solref="0.01 1" solimp=".25 .25 0.001" />
-    </equality>
-
-
-    <worldbody>
-        <include file="vx300s_left.xml" />
-        <include file="vx300s_right.xml" />
-
-        <body mocap="true" name="mocap_left" pos="0.095 0.50 0.425">
-            <site pos="0 0 0" size="0.003 0.003 0.03" type="box" name="mocap_left_site1" rgba="1 0 0 1"/>
-            <site pos="0 0 0" size="0.003 0.03 0.003" type="box" name="mocap_left_site2" rgba="1 0 0 1"/>
-            <site pos="0 0 0" size="0.03 0.003 0.003" type="box" name="mocap_left_site3" rgba="1 0 0 1"/>
-        </body>
-        <body mocap="true" name="mocap_right" pos="-0.095 0.50 0.425">
-            <site pos="0 0 0" size="0.003 0.003 0.03" type="box" name="mocap_right_site1" rgba="1 0 0 1"/>
-            <site pos="0 0 0" size="0.003 0.03 0.003" type="box" name="mocap_right_site2" rgba="1 0 0 1"/>
-            <site pos="0 0 0" size="0.03 0.003 0.003" type="box" name="mocap_right_site3" rgba="1 0 0 1"/>
-        </body>
-
-        <body name="peg" pos="0.2 0.5 0.05">
-            <joint name="red_peg_joint" type="free" frictionloss="0.01" />
-            <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.06 0.01 0.01" type="box" name="red_peg" rgba="1 0 0 1" />
-        </body>
-
-        <body name="socket" pos="-0.2 0.5 0.05">
-            <joint name="blue_socket_joint" type="free" frictionloss="0.01" />
-            <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
-<!--            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.06 0.01 0.01" type="box" name="red_peg_ref" rgba="1 0 0 1" />-->
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 -0.02" size="0.06 0.018 0.002" type="box" name="socket-1" rgba="0 0 1 1" />
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 0.02" size="0.06 0.018 0.002" type="box" name="socket-2" rgba="0 0 1 1" />
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0.02 0" size="0.06 0.002 0.018" type="box" name="socket-3" rgba="0 0 1 1" />
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 -0.02 0" size="0.06 0.002 0.018" type="box" name="socket-4" rgba="0 0 1 1" />
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.04 0.01 0.01" type="box" name="pin" rgba="1 0 0 1" />
-        </body>
-
-    </worldbody>
-
-    <actuator>
-        <position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_left/left_finger" kp="200"  user="1"/>
-        <position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_left/right_finger" kp="200"  user="1"/>
-
-        <position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_right/left_finger" kp="200"  user="1"/>
-        <position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_right/right_finger" kp="200"  user="1"/>
-
-    </actuator>
-
-    <keyframe>
-        <key qpos="0 -0.96 1.16 0 -0.3 0 0.024 -0.024  0 -0.96 1.16 0 -0.3 0 0.024 -0.024  0.2 0.5 0.05 1 0 0 0  -0.2 0.5 0.05 1 0 0 0"/>
-    </keyframe>
-
-
-</mujoco>

lerobot/common/envs/aloha/assets/bimanual_viperx_end_effector_transfer_cube.xml

@@ -1,48 +0,0 @@
-<mujoco>
-    <include file="scene.xml"/>
-    <include file="vx300s_dependencies.xml"/>
-
-    <equality>
-        <weld body1="mocap_left" body2="vx300s_left/gripper_link" solref="0.01 1" solimp=".25 .25 0.001" />
-        <weld body1="mocap_right" body2="vx300s_right/gripper_link" solref="0.01 1" solimp=".25 .25 0.001" />
-    </equality>
-
-
-    <worldbody>
-        <include file="vx300s_left.xml" />
-        <include file="vx300s_right.xml" />
-
-        <body mocap="true" name="mocap_left" pos="0.095 0.50 0.425">
-            <site pos="0 0 0" size="0.003 0.003 0.03" type="box" name="mocap_left_site1" rgba="1 0 0 1"/>
-            <site pos="0 0 0" size="0.003 0.03 0.003" type="box" name="mocap_left_site2" rgba="1 0 0 1"/>
-            <site pos="0 0 0" size="0.03 0.003 0.003" type="box" name="mocap_left_site3" rgba="1 0 0 1"/>
-        </body>
-        <body mocap="true" name="mocap_right" pos="-0.095 0.50 0.425">
-            <site pos="0 0 0" size="0.003 0.003 0.03" type="box" name="mocap_right_site1" rgba="1 0 0 1"/>
-            <site pos="0 0 0" size="0.003 0.03 0.003" type="box" name="mocap_right_site2" rgba="1 0 0 1"/>
-            <site pos="0 0 0" size="0.03 0.003 0.003" type="box" name="mocap_right_site3" rgba="1 0 0 1"/>
-        </body>
-
-        <body name="box" pos="0.2 0.5 0.05">
-            <joint name="red_box_joint" type="free" frictionloss="0.01" />
-            <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.02 0.02 0.02" type="box" name="red_box" rgba="1 0 0 1" />
-        </body>
-
-    </worldbody>
-
-    <actuator>
-        <position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_left/left_finger" kp="200"  user="1"/>
-        <position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_left/right_finger" kp="200"  user="1"/>
-
-        <position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_right/left_finger" kp="200"  user="1"/>
-        <position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_right/right_finger" kp="200"  user="1"/>
-
-    </actuator>
-
-    <keyframe>
-        <key qpos="0 -0.96 1.16 0 -0.3 0 0.024 -0.024  0 -0.96 1.16 0 -0.3 0 0.024 -0.024  0.2 0.5 0.05 1 0 0 0"/>
-    </keyframe>
-
-
-</mujoco>

lerobot/common/envs/aloha/assets/bimanual_viperx_insertion.xml

@@ -1,53 +0,0 @@
-<mujoco>
-    <include file="scene.xml"/>
-    <include file="vx300s_dependencies.xml"/>
-    <worldbody>
-        <include file="vx300s_left.xml" />
-        <include file="vx300s_right.xml" />
-
-        <body name="peg" pos="0.2 0.5 0.05">
-            <joint name="red_peg_joint" type="free" frictionloss="0.01" />
-            <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.06 0.01 0.01" type="box" name="red_peg" rgba="1 0 0 1" />
-        </body>
-
-        <body name="socket" pos="-0.2 0.5 0.05">
-            <joint name="blue_socket_joint" type="free" frictionloss="0.01" />
-            <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
-<!--            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.06 0.01 0.01" type="box" name="red_peg_ref" rgba="1 0 0 1" />-->
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 -0.02" size="0.06 0.018 0.002" type="box" name="socket-1" rgba="0 0 1 1" />
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 0.02" size="0.06 0.018 0.002" type="box" name="socket-2" rgba="0 0 1 1" />
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0.02 0" size="0.06 0.002 0.018" type="box" name="socket-3" rgba="0 0 1 1" />
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 -0.02 0" size="0.06 0.002 0.018" type="box" name="socket-4" rgba="0 0 1 1" />
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.04 0.01 0.01" type="box" name="pin" rgba="1 0 0 1" />
-        </body>
-
-    </worldbody>
-
-    <actuator>
-        <position ctrllimited="true" ctrlrange="-3.14158 3.14158" joint="vx300s_left/waist" kp="800"  user="1" forcelimited="true" forcerange="-150 150"/>
-        <position ctrllimited="true" ctrlrange="-1.85005 1.25664" joint="vx300s_left/shoulder" kp="1600"  user="1" forcelimited="true" forcerange="-300 300"/>
-        <position ctrllimited="true" ctrlrange="-1.76278 1.6057" joint="vx300s_left/elbow" kp="800"  user="1" forcelimited="true" forcerange="-100 100"/>
-        <position ctrllimited="true" ctrlrange="-3.14158 3.14158" joint="vx300s_left/forearm_roll" kp="10"  user="1" forcelimited="true" forcerange="-100 100"/>
-        <position ctrllimited="true" ctrlrange="-1.8675 2.23402" joint="vx300s_left/wrist_angle" kp="50"  user="1"/>
-        <position ctrllimited="true" ctrlrange="-3.14158 3.14158" joint="vx300s_left/wrist_rotate" kp="20"  user="1"/>
-        <position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_left/left_finger" kp="200"  user="1"/>
-        <position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_left/right_finger" kp="200"  user="1"/>
-
-        <position ctrllimited="true" ctrlrange="-3.14158 3.14158" joint="vx300s_right/waist" kp="800"  user="1" forcelimited="true" forcerange="-150 150"/>
-        <position ctrllimited="true" ctrlrange="-1.85005 1.25664" joint="vx300s_right/shoulder" kp="1600"  user="1" forcelimited="true" forcerange="-300 300"/>
-        <position ctrllimited="true" ctrlrange="-1.76278 1.6057" joint="vx300s_right/elbow" kp="800"  user="1" forcelimited="true" forcerange="-100 100"/>
-        <position ctrllimited="true" ctrlrange="-3.14158 3.14158" joint="vx300s_right/forearm_roll" kp="10"  user="1" forcelimited="true" forcerange="-100 100"/>
-        <position ctrllimited="true" ctrlrange="-1.8675 2.23402" joint="vx300s_right/wrist_angle" kp="50"  user="1"/>
-        <position ctrllimited="true" ctrlrange="-3.14158 3.14158" joint="vx300s_right/wrist_rotate" kp="20"  user="1"/>
-        <position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_right/left_finger" kp="200"  user="1"/>
-        <position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_right/right_finger" kp="200"  user="1"/>
-
-    </actuator>
-
-    <keyframe>
-        <key qpos="0 -0.96 1.16 0 -0.3 0 0.024 -0.024  0 -0.96 1.16 0 -0.3 0 0.024 -0.024  0.2 0.5 0.05 1 0 0 0  -0.2 0.5 0.05 1 0 0 0"/>
-    </keyframe>
-
-
-</mujoco>

lerobot/common/envs/aloha/assets/bimanual_viperx_transfer_cube.xml

@@ -1,42 +0,0 @@
-<mujoco>
-    <include file="scene.xml"/>
-    <include file="vx300s_dependencies.xml"/>
-    <worldbody>
-        <include file="vx300s_left.xml" />
-        <include file="vx300s_right.xml" />
-
-        <body name="box" pos="0.2 0.5 0.05">
-            <joint name="red_box_joint" type="free" frictionloss="0.01" />
-            <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
-            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.02 0.02 0.02" type="box" name="red_box" rgba="1 0 0 1" />
-        </body>
-
-    </worldbody>
-
-    <actuator>
-        <position ctrllimited="true" ctrlrange="-3.14158 3.14158" joint="vx300s_left/waist" kp="800"  user="1" forcelimited="true" forcerange="-150 150"/>
-        <position ctrllimited="true" ctrlrange="-1.85005 1.25664" joint="vx300s_left/shoulder" kp="1600"  user="1" forcelimited="true" forcerange="-300 300"/>
-        <position ctrllimited="true" ctrlrange="-1.76278 1.6057" joint="vx300s_left/elbow" kp="800"  user="1" forcelimited="true" forcerange="-100 100"/>
-        <position ctrllimited="true" ctrlrange="-3.14158 3.14158" joint="vx300s_left/forearm_roll" kp="10"  user="1" forcelimited="true" forcerange="-100 100"/>
-        <position ctrllimited="true" ctrlrange="-1.8675 2.23402" joint="vx300s_left/wrist_angle" kp="50"  user="1"/>
-        <position ctrllimited="true" ctrlrange="-3.14158 3.14158" joint="vx300s_left/wrist_rotate" kp="20"  user="1"/>
-        <position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_left/left_finger" kp="200"  user="1"/>
-        <position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_left/right_finger" kp="200"  user="1"/>
-
-        <position ctrllimited="true" ctrlrange="-3.14158 3.14158" joint="vx300s_right/waist" kp="800"  user="1" forcelimited="true" forcerange="-150 150"/>
-        <position ctrllimited="true" ctrlrange="-1.85005 1.25664" joint="vx300s_right/shoulder" kp="1600"  user="1" forcelimited="true" forcerange="-300 300"/>
-        <position ctrllimited="true" ctrlrange="-1.76278 1.6057" joint="vx300s_right/elbow" kp="800"  user="1" forcelimited="true" forcerange="-100 100"/>
-        <position ctrllimited="true" ctrlrange="-3.14158 3.14158" joint="vx300s_right/forearm_roll" kp="10"  user="1" forcelimited="true" forcerange="-100 100"/>
-        <position ctrllimited="true" ctrlrange="-1.8675 2.23402" joint="vx300s_right/wrist_angle" kp="50"  user="1"/>
-        <position ctrllimited="true" ctrlrange="-3.14158 3.14158" joint="vx300s_right/wrist_rotate" kp="20"  user="1"/>
-        <position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_right/left_finger" kp="200"  user="1"/>
-        <position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_right/right_finger" kp="200"  user="1"/>
-
-    </actuator>
-
-    <keyframe>
-        <key qpos="0 -0.96 1.16 0 -0.3 0 0.024 -0.024  0 -0.96 1.16 0 -0.3 0 0.024 -0.024  0.2 0.5 0.05 1 0 0 0"/>
-    </keyframe>
-
-
-</mujoco>

lerobot/common/envs/aloha/assets/scene.xml

@@ -1,38 +0,0 @@
-<mujocoinclude>
-<!--    <option timestep='0.0025' iterations="50" tolerance="1e-10" solver="Newton" jacobian="dense" cone="elliptic"/>-->
-
-    <asset>
-        <mesh file="tabletop.stl" name="tabletop" scale="0.001 0.001 0.001"/>
-    </asset>
-
-    <visual>
-        <map fogstart="1.5" fogend="5" force="0.1" znear="0.1"/>
-        <quality shadowsize="4096" offsamples="4"/>
-        <headlight ambient="0.4 0.4 0.4"/>
-    </visual>
-
-    <worldbody>
-        <light castshadow="false" directional='true' diffuse='.3 .3 .3' specular='0.3 0.3 0.3' pos='-1 -1 1'
-               dir='1 1 -1'/>
-        <light directional='true' diffuse='.3 .3 .3' specular='0.3 0.3 0.3' pos='1 -1 1' dir='-1 1 -1'/>
-        <light castshadow="false" directional='true' diffuse='.3 .3 .3' specular='0.3 0.3 0.3' pos='0 1 1'
-               dir='0 -1 -1'/>
-
-        <body name="table" pos="0 .6 0">
-            <geom group="1" mesh="tabletop" pos="0 0 0" type="mesh" conaffinity="1" contype="1" name="table" rgba="0.2 0.2 0.2 1" />
-        </body>
-        <body name="midair" pos="0 .6 0.2">
-            <site pos="0 0 0" size="0.01" type="sphere" name="midair" rgba="1 0 0 0"/>
-        </body>
-
-        <camera name="left_pillar" pos="-0.5 0.2 0.6" fovy="78" mode="targetbody" target="table"/>
-        <camera name="right_pillar" pos="0.5 0.2 0.6" fovy="78" mode="targetbody" target="table"/>
-        <camera name="top" pos="0 0.6 0.8" fovy="78" mode="targetbody" target="table"/>
-        <camera name="angle" pos="0 0 0.6" fovy="78" mode="targetbody" target="table"/>
-        <camera name="front_close" pos="0 0.2 0.4" fovy="78" mode="targetbody" target="vx300s_left/camera_focus"/>
-
-    </worldbody>
-
-
-
-</mujocoinclude>

lerobot/common/envs/aloha/assets/tabletop.stl

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

lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_left.stl

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

lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_right.stl

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

lerobot/common/envs/aloha/assets/vx300s_10_gripper_finger.stl

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

lerobot/common/envs/aloha/assets/vx300s_11_ar_tag.stl

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

lerobot/common/envs/aloha/assets/vx300s_1_base.stl

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

lerobot/common/envs/aloha/assets/vx300s_2_shoulder.stl

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

lerobot/common/envs/aloha/assets/vx300s_3_upper_arm.stl

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

lerobot/common/envs/aloha/assets/vx300s_4_upper_forearm.stl

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

lerobot/common/envs/aloha/assets/vx300s_5_lower_forearm.stl

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

lerobot/common/envs/aloha/assets/vx300s_6_wrist.stl

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

lerobot/common/envs/aloha/assets/vx300s_7_gripper.stl

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

lerobot/common/envs/aloha/assets/vx300s_8_gripper_prop.stl

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

lerobot/common/envs/aloha/assets/vx300s_9_gripper_bar.stl

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

lerobot/common/envs/aloha/assets/vx300s_dependencies.xml

@@ -1,17 +0,0 @@
-<mujocoinclude>
-    <compiler angle="radian" inertiafromgeom="auto" inertiagrouprange="4 5"/>
-    <asset>
-        <mesh name="vx300s_1_base" file="vx300s_1_base.stl" scale="0.001 0.001 0.001" />
-        <mesh name="vx300s_2_shoulder" file="vx300s_2_shoulder.stl" scale="0.001 0.001 0.001" />
-        <mesh name="vx300s_3_upper_arm" file="vx300s_3_upper_arm.stl" scale="0.001 0.001 0.001" />
-        <mesh name="vx300s_4_upper_forearm" file="vx300s_4_upper_forearm.stl" scale="0.001 0.001 0.001" />
-        <mesh name="vx300s_5_lower_forearm" file="vx300s_5_lower_forearm.stl" scale="0.001 0.001 0.001" />
-        <mesh name="vx300s_6_wrist" file="vx300s_6_wrist.stl" scale="0.001 0.001 0.001" />
-        <mesh name="vx300s_7_gripper" file="vx300s_7_gripper.stl" scale="0.001 0.001 0.001" />
-        <mesh name="vx300s_8_gripper_prop" file="vx300s_8_gripper_prop.stl" scale="0.001 0.001 0.001" />
-        <mesh name="vx300s_9_gripper_bar" file="vx300s_9_gripper_bar.stl" scale="0.001 0.001 0.001" />
-        <mesh name="vx300s_10_gripper_finger_left" file="vx300s_10_custom_finger_left.stl" scale="0.001 0.001 0.001" />
-        <mesh name="vx300s_10_gripper_finger_right" file="vx300s_10_custom_finger_right.stl" scale="0.001 0.001 0.001" />
-    </asset>
-
-</mujocoinclude>

lerobot/common/envs/aloha/assets/vx300s_left.xml

@@ -1,59 +0,0 @@
-
-<mujocoinclude>
-    <body name="vx300s_left" pos="-0.469 0.5 0">
-        <geom quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_1_base" name="vx300s_left/1_base" contype="0" conaffinity="0"/>
-        <body name="vx300s_left/shoulder_link" pos="0 0 0.079">
-            <inertial pos="0.000259233 -3.3552e-06 0.0116129" quat="-0.476119 0.476083 0.52279 0.522826" mass="0.798614" diaginertia="0.00120156 0.00113744 0.0009388" />
-            <joint name="vx300s_left/waist" pos="0 0 0" axis="0 0 1" limited="true" range="-3.14158 3.14158" frictionloss="50" />
-            <geom pos="0 0 -0.003" quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_2_shoulder" name="vx300s_left/2_shoulder" />
-            <body name="vx300s_left/upper_arm_link" pos="0 0 0.04805">
-                <inertial pos="0.0206949 4e-10 0.226459" quat="0 0.0728458 0 0.997343" mass="0.792592" diaginertia="0.00911338 0.008925 0.000759317" />
-                <joint name="vx300s_left/shoulder" pos="0 0 0" axis="0 1 0" limited="true" range="-1.85005 1.25664" frictionloss="60" />
-                <geom quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_3_upper_arm" name="vx300s_left/3_upper_arm"/>
-                <body name="vx300s_left/upper_forearm_link" pos="0.05955 0 0.3">
-                    <inertial pos="0.105723 0 0" quat="-0.000621631 0.704724 0.0105292 0.709403" mass="0.322228" diaginertia="0.00144107 0.00134228 0.000152047" />
-                    <joint name="vx300s_left/elbow" pos="0 0 0" axis="0 1 0" limited="true" range="-1.76278 1.6057" frictionloss="60" />
-                    <geom type="mesh" mesh="vx300s_4_upper_forearm" name="vx300s_left/4_upper_forearm" />
-                    <body name="vx300s_left/lower_forearm_link" pos="0.2 0 0">
-                        <inertial pos="0.0513477 0.00680462 0" quat="-0.702604 -0.0796724 -0.702604 0.0796724" mass="0.414823" diaginertia="0.0005911 0.000546493 0.000155707" />
-                        <joint name="vx300s_left/forearm_roll" pos="0 0 0" axis="1 0 0" limited="true" range="-3.14158 3.14158" frictionloss="30" />
-                        <geom quat="0 1 0 0" type="mesh" mesh="vx300s_5_lower_forearm" name="vx300s_left/5_lower_forearm"/>
-                        <body name="vx300s_left/wrist_link" pos="0.1 0 0">
-                            <inertial pos="0.046743 -7.6652e-06 0.010565" quat="-0.00100191 0.544586 0.0026583 0.8387" mass="0.115395" diaginertia="5.45707e-05 4.63101e-05 4.32692e-05" />
-                            <joint name="vx300s_left/wrist_angle" pos="0 0 0" axis="0 1 0" limited="true" range="-1.8675 2.23402" frictionloss="30" />
-                            <geom quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_6_wrist" name="vx300s_left/6_wrist" />
-                            <body name="vx300s_left/gripper_link" pos="0.069744 0 0">
-                                <body name="vx300s_left/camera_focus" pos="0.15 0 0.01">
-                                    <site pos="0 0 0" size="0.01" type="sphere" name="left_cam_focus" rgba="0 0 1 0"/>
-                                </body>
-                                <site pos="0.15 0 0" size="0.003 0.003 0.03" type="box" name="cali_left_site1" rgba="0 0 1 0"/>
-                                <site pos="0.15 0 0" size="0.003 0.03 0.003" type="box" name="cali_left_site2" rgba="0 0 1 0"/>
-                                <site pos="0.15 0 0" size="0.03 0.003 0.003" type="box" name="cali_left_site3" rgba="0 0 1 0"/>
-                                <camera name="left_wrist" pos="-0.1 0 0.16" fovy="20" mode="targetbody" target="vx300s_left/camera_focus"/>
-                                <inertial pos="0.0395662 -2.56311e-07 0.00400649" quat="0.62033 0.619916 -0.339682 0.339869" mass="0.251652" diaginertia="0.000689546 0.000650316 0.000468142" />
-                                <joint name="vx300s_left/wrist_rotate" pos="0 0 0" axis="1 0 0" limited="true" range="-3.14158 3.14158" frictionloss="30" />
-                                <geom pos="-0.02 0 0" quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_7_gripper" name="vx300s_left/7_gripper" />
-                                <geom pos="-0.020175 0 0" quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_9_gripper_bar" name="vx300s_left/9_gripper_bar" />
-                                <body name="vx300s_left/gripper_prop_link" pos="0.0485 0 0">
-                                    <inertial pos="0.002378 2.85e-08 0" quat="0 0 0.897698 0.440611" mass="0.008009" diaginertia="4.2979e-06 2.8868e-06 1.5314e-06" />
-<!--                                    <joint name="vx300s_left/gripper" pos="0 0 0" axis="1 0 0" frictionloss="30" />-->
-                                    <geom pos="-0.0685 0 0" quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_8_gripper_prop" name="vx300s_left/8_gripper_prop" />
-                                </body>
-                                <body name="vx300s_left/left_finger_link" pos="0.0687 0 0">
-                                    <inertial pos="0.017344 -0.0060692 0" quat="0.449364 0.449364 -0.54596 -0.54596" mass="0.034796" diaginertia="2.48003e-05 1.417e-05 1.20797e-05" />
-                                    <joint name="vx300s_left/left_finger" pos="0 0 0" axis="0 1 0" type="slide" limited="true" range="0.021 0.057" frictionloss="30" />
-                                    <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0.005 -0.052 0" euler="3.14 1.57 0" type="mesh" mesh="vx300s_10_gripper_finger_left" name="vx300s_left/10_left_gripper_finger"/>
-                                </body>
-                                <body name="vx300s_left/right_finger_link" pos="0.0687 0 0">
-                                    <inertial pos="0.017344 0.0060692 0" quat="0.44937 -0.44937 0.545955 -0.545955" mass="0.034796" diaginertia="2.48002e-05 1.417e-05 1.20798e-05" />
-                                    <joint name="vx300s_left/right_finger" pos="0 0 0" axis="0 1 0" type="slide" limited="true" range="-0.057 -0.021" frictionloss="30" />
-                                    <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0.005 0.052 0" euler="3.14 1.57 0" type="mesh" mesh="vx300s_10_gripper_finger_right" name="vx300s_left/10_right_gripper_finger"/>
-                                </body>
-                            </body>
-                        </body>
-                    </body>
-                </body>
-            </body>
-        </body>
-    </body>
-</mujocoinclude>

lerobot/common/envs/aloha/assets/vx300s_right.xml

@@ -1,59 +0,0 @@
-
-<mujocoinclude>
-    <body name="vx300s_right" pos="0.469 0.5 0" euler="0 0 3.1416">
-        <geom quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_1_base" name="vx300s_right/1_base" contype="0" conaffinity="0"/>
-        <body name="vx300s_right/shoulder_link" pos="0 0 0.079">
-            <inertial pos="0.000259233 -3.3552e-06 0.0116129" quat="-0.476119 0.476083 0.52279 0.522826" mass="0.798614" diaginertia="0.00120156 0.00113744 0.0009388" />
-            <joint name="vx300s_right/waist" pos="0 0 0" axis="0 0 1" limited="true" range="-3.14158 3.14158" frictionloss="50" />
-            <geom pos="0 0 -0.003" quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_2_shoulder" name="vx300s_right/2_shoulder" />
-            <body name="vx300s_right/upper_arm_link" pos="0 0 0.04805">
-                <inertial pos="0.0206949 4e-10 0.226459" quat="0 0.0728458 0 0.997343" mass="0.792592" diaginertia="0.00911338 0.008925 0.000759317" />
-                <joint name="vx300s_right/shoulder" pos="0 0 0" axis="0 1 0" limited="true" range="-1.85005 1.25664" frictionloss="60" />
-                <geom quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_3_upper_arm" name="vx300s_right/3_upper_arm"/>
-                <body name="vx300s_right/upper_forearm_link" pos="0.05955 0 0.3">
-                    <inertial pos="0.105723 0 0" quat="-0.000621631 0.704724 0.0105292 0.709403" mass="0.322228" diaginertia="0.00144107 0.00134228 0.000152047" />
-                    <joint name="vx300s_right/elbow" pos="0 0 0" axis="0 1 0" limited="true" range="-1.76278 1.6057" frictionloss="60" />
-                    <geom type="mesh" mesh="vx300s_4_upper_forearm" name="vx300s_right/4_upper_forearm" />
-                    <body name="vx300s_right/lower_forearm_link" pos="0.2 0 0">
-                        <inertial pos="0.0513477 0.00680462 0" quat="-0.702604 -0.0796724 -0.702604 0.0796724" mass="0.414823" diaginertia="0.0005911 0.000546493 0.000155707" />
-                        <joint name="vx300s_right/forearm_roll" pos="0 0 0" axis="1 0 0" limited="true" range="-3.14158 3.14158" frictionloss="30" />
-                        <geom quat="0 1 0 0" type="mesh" mesh="vx300s_5_lower_forearm" name="vx300s_right/5_lower_forearm"/>
-                        <body name="vx300s_right/wrist_link" pos="0.1 0 0">
-                            <inertial pos="0.046743 -7.6652e-06 0.010565" quat="-0.00100191 0.544586 0.0026583 0.8387" mass="0.115395" diaginertia="5.45707e-05 4.63101e-05 4.32692e-05" />
-                            <joint name="vx300s_right/wrist_angle" pos="0 0 0" axis="0 1 0" limited="true" range="-1.8675 2.23402" frictionloss="30" />
-                            <geom quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_6_wrist" name="vx300s_right/6_wrist" />
-                            <body name="vx300s_right/gripper_link" pos="0.069744 0 0">
-                                <body name="vx300s_right/camera_focus" pos="0.15 0 0.01">
-                                    <site pos="0 0 0" size="0.01" type="sphere" name="right_cam_focus" rgba="0 0 1 0"/>
-                                </body>
-                                <site pos="0.15 0 0" size="0.003 0.003 0.03" type="box" name="cali_right_site1" rgba="0 0 1 0"/>
-                                <site pos="0.15 0 0" size="0.003 0.03 0.003" type="box" name="cali_right_site2" rgba="0 0 1 0"/>
-                                <site pos="0.15 0 0" size="0.03 0.003 0.003" type="box" name="cali_right_site3" rgba="0 0 1 0"/>
-                                <camera name="right_wrist" pos="-0.1 0 0.16" fovy="20" mode="targetbody" target="vx300s_right/camera_focus"/>
-                                <inertial pos="0.0395662 -2.56311e-07 0.00400649" quat="0.62033 0.619916 -0.339682 0.339869" mass="0.251652" diaginertia="0.000689546 0.000650316 0.000468142" />
-                                <joint name="vx300s_right/wrist_rotate" pos="0 0 0" axis="1 0 0" limited="true" range="-3.14158 3.14158" frictionloss="30" />
-                                <geom pos="-0.02 0 0" quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_7_gripper" name="vx300s_right/7_gripper" />
-                                <geom pos="-0.020175 0 0" quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_9_gripper_bar" name="vx300s_right/9_gripper_bar" />
-                                <body name="vx300s_right/gripper_prop_link" pos="0.0485 0 0">
-                                    <inertial pos="0.002378 2.85e-08 0" quat="0 0 0.897698 0.440611" mass="0.008009" diaginertia="4.2979e-06 2.8868e-06 1.5314e-06" />
-<!--                                    <joint name="vx300s_right/gripper" pos="0 0 0" axis="1 0 0" frictionloss="30" />-->
-                                    <geom pos="-0.0685 0 0" quat="0.707107 0 0 0.707107" type="mesh" mesh="vx300s_8_gripper_prop" name="vx300s_right/8_gripper_prop" />
-                                </body>
-                                <body name="vx300s_right/left_finger_link" pos="0.0687 0 0">
-                                    <inertial pos="0.017344 -0.0060692 0" quat="0.449364 0.449364 -0.54596 -0.54596" mass="0.034796" diaginertia="2.48003e-05 1.417e-05 1.20797e-05" />
-                                    <joint name="vx300s_right/left_finger" pos="0 0 0" axis="0 1 0" type="slide" limited="true" range="0.021 0.057" frictionloss="30" />
-                                    <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0.005 -0.052 0" euler="3.14 1.57 0" type="mesh" mesh="vx300s_10_gripper_finger_left" name="vx300s_right/10_left_gripper_finger"/>
-                                </body>
-                                <body name="vx300s_right/right_finger_link" pos="0.0687 0 0">
-                                    <inertial pos="0.017344 0.0060692 0" quat="0.44937 -0.44937 0.545955 -0.545955" mass="0.034796" diaginertia="2.48002e-05 1.417e-05 1.20798e-05" />
-                                    <joint name="vx300s_right/right_finger" pos="0 0 0" axis="0 1 0" type="slide" limited="true" range="-0.057 -0.021" frictionloss="30" />
-                                    <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0.005 0.052 0" euler="3.14 1.57 0" type="mesh" mesh="vx300s_10_gripper_finger_right" name="vx300s_right/10_right_gripper_finger"/>
-                                </body>
-                            </body>
-                        </body>
-                    </body>
-                </body>
-            </body>
-        </body>
-    </body>
-</mujocoinclude>

lerobot/common/envs/aloha/constants.py

@@ -1,163 +0,0 @@
-from pathlib import Path
-
-### Simulation envs fixed constants
-DT = 0.02  # 0.02 ms -> 1/0.2 = 50 hz
-FPS = 50
-
-
-JOINTS = [
-    # absolute joint position
-    "left_arm_waist",
-    "left_arm_shoulder",
-    "left_arm_elbow",
-    "left_arm_forearm_roll",
-    "left_arm_wrist_angle",
-    "left_arm_wrist_rotate",
-    # normalized gripper position 0: close, 1: open
-    "left_arm_gripper",
-    # absolute joint position
-    "right_arm_waist",
-    "right_arm_shoulder",
-    "right_arm_elbow",
-    "right_arm_forearm_roll",
-    "right_arm_wrist_angle",
-    "right_arm_wrist_rotate",
-    # normalized gripper position 0: close, 1: open
-    "right_arm_gripper",
-]
-
-ACTIONS = [
-    # position and quaternion for end effector
-    "left_arm_waist",
-    "left_arm_shoulder",
-    "left_arm_elbow",
-    "left_arm_forearm_roll",
-    "left_arm_wrist_angle",
-    "left_arm_wrist_rotate",
-    # normalized gripper position (0: close, 1: open)
-    "left_arm_gripper",
-    "right_arm_waist",
-    "right_arm_shoulder",
-    "right_arm_elbow",
-    "right_arm_forearm_roll",
-    "right_arm_wrist_angle",
-    "right_arm_wrist_rotate",
-    # normalized gripper position (0: close, 1: open)
-    "right_arm_gripper",
-]
-
-
-START_ARM_POSE = [
-    0,
-    -0.96,
-    1.16,
-    0,
-    -0.3,
-    0,
-    0.02239,
-    -0.02239,
-    0,
-    -0.96,
-    1.16,
-    0,
-    -0.3,
-    0,
-    0.02239,
-    -0.02239,
-]
-
-ASSETS_DIR = Path(__file__).parent.resolve() / "assets"  # note: absolute path
-
-# Left finger position limits (qpos[7]), right_finger = -1 * left_finger
-MASTER_GRIPPER_POSITION_OPEN = 0.02417
-MASTER_GRIPPER_POSITION_CLOSE = 0.01244
-PUPPET_GRIPPER_POSITION_OPEN = 0.05800
-PUPPET_GRIPPER_POSITION_CLOSE = 0.01844
-
-# Gripper joint limits (qpos[6])
-MASTER_GRIPPER_JOINT_OPEN = 0.3083
-MASTER_GRIPPER_JOINT_CLOSE = -0.6842
-PUPPET_GRIPPER_JOINT_OPEN = 1.4910
-PUPPET_GRIPPER_JOINT_CLOSE = -0.6213
-
-MASTER_GRIPPER_JOINT_MID = (MASTER_GRIPPER_JOINT_OPEN + MASTER_GRIPPER_JOINT_CLOSE) / 2
-
-############################ Helper functions ############################
-
-
-def normalize_master_gripper_position(x):
-    return (x - MASTER_GRIPPER_POSITION_CLOSE) / (
-        MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE
-    )
-
-
-def normalize_puppet_gripper_position(x):
-    return (x - PUPPET_GRIPPER_POSITION_CLOSE) / (
-        PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE
-    )
-
-
-def unnormalize_master_gripper_position(x):
-    return x * (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) + MASTER_GRIPPER_POSITION_CLOSE
-
-
-def unnormalize_puppet_gripper_position(x):
-    return x * (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) + PUPPET_GRIPPER_POSITION_CLOSE
-
-
-def convert_position_from_master_to_puppet(x):
-    return unnormalize_puppet_gripper_position(normalize_master_gripper_position(x))
-
-
-def normalizer_master_gripper_joint(x):
-    return (x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE)
-
-
-def normalize_puppet_gripper_joint(x):
-    return (x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE)
-
-
-def unnormalize_master_gripper_joint(x):
-    return x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
-
-
-def unnormalize_puppet_gripper_joint(x):
-    return x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
-
-
-def convert_join_from_master_to_puppet(x):
-    return unnormalize_puppet_gripper_joint(normalizer_master_gripper_joint(x))
-
-
-def normalize_master_gripper_velocity(x):
-    return x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE)
-
-
-def normalize_puppet_gripper_velocity(x):
-    return x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE)
-
-
-def convert_master_from_position_to_joint(x):
-    return (
-        normalize_master_gripper_position(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE)
-        + MASTER_GRIPPER_JOINT_CLOSE
-    )
-
-
-def convert_master_from_joint_to_position(x):
-    return unnormalize_master_gripper_position(
-        (x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE)
-    )
-
-
-def convert_puppet_from_position_to_join(x):
-    return (
-        normalize_puppet_gripper_position(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE)
-        + PUPPET_GRIPPER_JOINT_CLOSE
-    )
-
-
-def convert_puppet_from_joint_to_position(x):
-    return unnormalize_puppet_gripper_position(
-        (x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE)
-    )

lerobot/common/envs/aloha/env.py

@@ -1,298 +0,0 @@
-import importlib
-import logging
-from collections import deque
-from typing import Optional
-
-import einops
-import numpy as np
-import torch
-from dm_control import mujoco
-from dm_control.rl import control
-from tensordict import TensorDict
-from torchrl.data.tensor_specs import (
-    BoundedTensorSpec,
-    CompositeSpec,
-    DiscreteTensorSpec,
-    UnboundedContinuousTensorSpec,
-)
-
-from lerobot.common.envs.abstract import AbstractEnv
-from lerobot.common.envs.aloha.constants import (
-    ACTIONS,
-    ASSETS_DIR,
-    DT,
-    JOINTS,
-)
-from lerobot.common.envs.aloha.tasks.sim import BOX_POSE, InsertionTask, TransferCubeTask
-from lerobot.common.envs.aloha.tasks.sim_end_effector import (
-    InsertionEndEffectorTask,
-    TransferCubeEndEffectorTask,
-)
-from lerobot.common.envs.aloha.utils import sample_box_pose, sample_insertion_pose
-from lerobot.common.utils import set_global_seed
-
-_has_gym = importlib.util.find_spec("gymnasium") is not None
-
-
-class AlohaEnv(AbstractEnv):
-    name = "aloha"
-    available_tasks = ["sim_insertion", "sim_transfer_cube"]
-    _reset_warning_issued = False
-
-    def __init__(
-        self,
-        task,
-        frame_skip: int = 1,
-        from_pixels: bool = False,
-        pixels_only: bool = False,
-        image_size=None,
-        seed=1337,
-        device="cpu",
-        num_prev_obs=1,
-        num_prev_action=0,
-    ):
-        super().__init__(
-            task=task,
-            frame_skip=frame_skip,
-            from_pixels=from_pixels,
-            pixels_only=pixels_only,
-            image_size=image_size,
-            seed=seed,
-            device=device,
-            num_prev_obs=num_prev_obs,
-            num_prev_action=num_prev_action,
-        )
-
-    def _make_env(self):
-        if not _has_gym:
-            raise ImportError("Cannot import gymnasium.")
-
-        if not self.from_pixels:
-            raise NotImplementedError()
-
-        self._env = self._make_env_task(self.task)
-
-    def render(self, mode="rgb_array", width=640, height=480):
-        # TODO(rcadene): render and visualizer several cameras (e.g. angle, front_close)
-        image = self._env.physics.render(height=height, width=width, camera_id="top")
-        return image
-
-    def _make_env_task(self, task_name):
-        # time limit is controlled by StepCounter in env factory
-        time_limit = float("inf")
-
-        if "sim_transfer_cube" in task_name:
-            xml_path = ASSETS_DIR / "bimanual_viperx_transfer_cube.xml"
-            physics = mujoco.Physics.from_xml_path(str(xml_path))
-            task = TransferCubeTask(random=False)
-        elif "sim_insertion" in task_name:
-            xml_path = ASSETS_DIR / "bimanual_viperx_insertion.xml"
-            physics = mujoco.Physics.from_xml_path(str(xml_path))
-            task = InsertionTask(random=False)
-        elif "sim_end_effector_transfer_cube" in task_name:
-            raise NotImplementedError()
-            xml_path = ASSETS_DIR / "bimanual_viperx_end_effector_transfer_cube.xml"
-            physics = mujoco.Physics.from_xml_path(str(xml_path))
-            task = TransferCubeEndEffectorTask(random=False)
-        elif "sim_end_effector_insertion" in task_name:
-            raise NotImplementedError()
-            xml_path = ASSETS_DIR / "bimanual_viperx_end_effector_insertion.xml"
-            physics = mujoco.Physics.from_xml_path(str(xml_path))
-            task = InsertionEndEffectorTask(random=False)
-        else:
-            raise NotImplementedError(task_name)
-
-        env = control.Environment(
-            physics, task, time_limit, control_timestep=DT, n_sub_steps=None, flat_observation=False
-        )
-        return env
-
-    def _format_raw_obs(self, raw_obs):
-        if self.from_pixels:
-            image = torch.from_numpy(raw_obs["images"]["top"].copy())
-            image = einops.rearrange(image, "h w c -> c h w")
-            assert image.dtype == torch.uint8
-            obs = {"image": {"top": image}}
-
-            if not self.pixels_only:
-                obs["state"] = torch.from_numpy(raw_obs["qpos"]).type(torch.float32)
-        else:
-            # TODO(rcadene):
-            raise NotImplementedError()
-            # obs = {"state": torch.from_numpy(raw_obs["observation"]).type(torch.float32)}
-
-        return obs
-
-    def _reset(self, tensordict: Optional[TensorDict] = None):
-        if tensordict is not None and not AlohaEnv._reset_warning_issued:
-            logging.warning(f"{self.__class__.__name__}._reset ignores the provided tensordict.")
-            AlohaEnv._reset_warning_issued = True
-
-        # Seed the environment and update the seed to be used for the next reset.
-        self._next_seed = self.set_seed(self._next_seed)
-
-        # TODO(rcadene): do not use global variable for this
-        if "sim_transfer_cube" in self.task:
-            BOX_POSE[0] = sample_box_pose()  # used in sim reset
-        elif "sim_insertion" in self.task:
-            BOX_POSE[0] = np.concatenate(sample_insertion_pose())  # used in sim reset
-
-        raw_obs = self._env.reset()
-
-        obs = self._format_raw_obs(raw_obs.observation)
-
-        if self.num_prev_obs > 0:
-            stacked_obs = {}
-            if "image" in obs:
-                self._prev_obs_image_queue = deque(
-                    [obs["image"]["top"]] * (self.num_prev_obs + 1), maxlen=(self.num_prev_obs + 1)
-                )
-                stacked_obs["image"] = {"top": torch.stack(list(self._prev_obs_image_queue))}
-            if "state" in obs:
-                self._prev_obs_state_queue = deque(
-                    [obs["state"]] * (self.num_prev_obs + 1), maxlen=(self.num_prev_obs + 1)
-                )
-                stacked_obs["state"] = torch.stack(list(self._prev_obs_state_queue))
-            obs = stacked_obs
-
-        td = TensorDict(
-            {
-                "observation": TensorDict(obs, batch_size=[]),
-                "done": torch.tensor([False], dtype=torch.bool),
-            },
-            batch_size=[],
-        )
-
-        return td
-
-    def _step(self, tensordict: TensorDict):
-        td = tensordict
-        action = td["action"].numpy()
-        assert action.ndim == 1
-        # TODO(rcadene): add info["is_success"] and info["success"] ?
-
-        _, reward, _, raw_obs = self._env.step(action)
-
-        # TODO(rcadene): add an enum
-        success = done = reward == 4
-        obs = self._format_raw_obs(raw_obs)
-
-        if self.num_prev_obs > 0:
-            stacked_obs = {}
-            if "image" in obs:
-                self._prev_obs_image_queue.append(obs["image"]["top"])
-                stacked_obs["image"] = {"top": torch.stack(list(self._prev_obs_image_queue))}
-            if "state" in obs:
-                self._prev_obs_state_queue.append(obs["state"])
-                stacked_obs["state"] = torch.stack(list(self._prev_obs_state_queue))
-            obs = stacked_obs
-
-        td = TensorDict(
-            {
-                "observation": TensorDict(obs, batch_size=[]),
-                "reward": torch.tensor([reward], dtype=torch.float32),
-                # success and done are true when coverage > self.success_threshold in env
-                "done": torch.tensor([done], dtype=torch.bool),
-                "success": torch.tensor([success], dtype=torch.bool),
-            },
-            batch_size=[],
-        )
-        return td
-
-    def _make_spec(self):
-        obs = {}
-        from omegaconf import OmegaConf
-
-        if self.from_pixels:
-            if isinstance(self.image_size, int):
-                image_shape = (3, self.image_size, self.image_size)
-            elif OmegaConf.is_list(self.image_size) or isinstance(self.image_size, list):
-                assert len(self.image_size) == 3  # c h w
-                assert self.image_size[0] == 3  # c is RGB
-                image_shape = tuple(self.image_size)
-            else:
-                raise ValueError(self.image_size)
-            if self.num_prev_obs > 0:
-                image_shape = (self.num_prev_obs + 1, *image_shape)
-
-            obs["image"] = {
-                "top": BoundedTensorSpec(
-                    low=0,
-                    high=255,
-                    shape=image_shape,
-                    dtype=torch.uint8,
-                    device=self.device,
-                )
-            }
-            if not self.pixels_only:
-                state_shape = (len(JOINTS),)
-                if self.num_prev_obs > 0:
-                    state_shape = (self.num_prev_obs + 1, *state_shape)
-
-                obs["state"] = UnboundedContinuousTensorSpec(
-                    # TODO: add low and high bounds
-                    shape=state_shape,
-                    dtype=torch.float32,
-                    device=self.device,
-                )
-        else:
-            # TODO(rcadene): add observation_space achieved_goal and desired_goal?
-            state_shape = (len(JOINTS),)
-            if self.num_prev_obs > 0:
-                state_shape = (self.num_prev_obs + 1, *state_shape)
-
-            obs["state"] = UnboundedContinuousTensorSpec(
-                # TODO: add low and high bounds
-                shape=state_shape,
-                dtype=torch.float32,
-                device=self.device,
-            )
-        self.observation_spec = CompositeSpec({"observation": obs})
-
-        # TODO(rcadene): valid when controling end effector?
-        # action_space = self._env.action_spec()
-        # self.action_spec = BoundedTensorSpec(
-        #     low=action_space.minimum,
-        #     high=action_space.maximum,
-        #     shape=action_space.shape,
-        #     dtype=torch.float32,
-        #     device=self.device,
-        # )
-
-        # TODO(rcaene): add bounds (where are they????)
-        self.action_spec = BoundedTensorSpec(
-            shape=(len(ACTIONS)),
-            low=-1,
-            high=1,
-            dtype=torch.float32,
-            device=self.device,
-        )
-
-        self.reward_spec = UnboundedContinuousTensorSpec(
-            shape=(1,),
-            dtype=torch.float32,
-            device=self.device,
-        )
-
-        self.done_spec = CompositeSpec(
-            {
-                "done": DiscreteTensorSpec(
-                    2,
-                    shape=(1,),
-                    dtype=torch.bool,
-                    device=self.device,
-                ),
-                "success": DiscreteTensorSpec(
-                    2,
-                    shape=(1,),
-                    dtype=torch.bool,
-                    device=self.device,
-                ),
-            }
-        )
-
-    def _set_seed(self, seed: Optional[int]):
-        set_global_seed(seed)
-        # TODO(rcadene): seed the env
-        # self._env.seed(seed)
-        logging.warning("Aloha env is not seeded")

lerobot/common/envs/aloha/tasks/sim.py

@@ -1,219 +0,0 @@
-import collections
-
-import numpy as np
-from dm_control.suite import base
-
-from lerobot.common.envs.aloha.constants import (
-    START_ARM_POSE,
-    normalize_puppet_gripper_position,
-    normalize_puppet_gripper_velocity,
-    unnormalize_puppet_gripper_position,
-)
-
-BOX_POSE = [None]  # to be changed from outside
-
-"""
-Environment for simulated robot bi-manual manipulation, with joint position control
-Action space:      [left_arm_qpos (6),             # absolute joint position
-                    left_gripper_positions (1),    # normalized gripper position (0: close, 1: open)
-                    right_arm_qpos (6),            # absolute joint position
-                    right_gripper_positions (1),]  # normalized gripper position (0: close, 1: open)
-
-Observation space: {"qpos": Concat[ left_arm_qpos (6),         # absolute joint position
-                                    left_gripper_position (1),  # normalized gripper position (0: close, 1: open)
-                                    right_arm_qpos (6),         # absolute joint position
-                                    right_gripper_qpos (1)]     # normalized gripper position (0: close, 1: open)
-                    "qvel": Concat[ left_arm_qvel (6),         # absolute joint velocity (rad)
-                                    left_gripper_velocity (1),  # normalized gripper velocity (pos: opening, neg: closing)
-                                    right_arm_qvel (6),         # absolute joint velocity (rad)
-                                    right_gripper_qvel (1)]     # normalized gripper velocity (pos: opening, neg: closing)
-                    "images": {"main": (480x640x3)}        # h, w, c, dtype='uint8'
-"""
-
-
-class BimanualViperXTask(base.Task):
-    def __init__(self, random=None):
-        super().__init__(random=random)
-
-    def before_step(self, action, physics):
-        left_arm_action = action[:6]
-        right_arm_action = action[7 : 7 + 6]
-        normalized_left_gripper_action = action[6]
-        normalized_right_gripper_action = action[7 + 6]
-
-        left_gripper_action = unnormalize_puppet_gripper_position(normalized_left_gripper_action)
-        right_gripper_action = unnormalize_puppet_gripper_position(normalized_right_gripper_action)
-
-        full_left_gripper_action = [left_gripper_action, -left_gripper_action]
-        full_right_gripper_action = [right_gripper_action, -right_gripper_action]
-
-        env_action = np.concatenate(
-            [left_arm_action, full_left_gripper_action, right_arm_action, full_right_gripper_action]
-        )
-        super().before_step(env_action, physics)
-        return
-
-    def initialize_episode(self, physics):
-        """Sets the state of the environment at the start of each episode."""
-        super().initialize_episode(physics)
-
-    @staticmethod
-    def get_qpos(physics):
-        qpos_raw = physics.data.qpos.copy()
-        left_qpos_raw = qpos_raw[:8]
-        right_qpos_raw = qpos_raw[8:16]
-        left_arm_qpos = left_qpos_raw[:6]
-        right_arm_qpos = right_qpos_raw[:6]
-        left_gripper_qpos = [normalize_puppet_gripper_position(left_qpos_raw[6])]
-        right_gripper_qpos = [normalize_puppet_gripper_position(right_qpos_raw[6])]
-        return np.concatenate([left_arm_qpos, left_gripper_qpos, right_arm_qpos, right_gripper_qpos])
-
-    @staticmethod
-    def get_qvel(physics):
-        qvel_raw = physics.data.qvel.copy()
-        left_qvel_raw = qvel_raw[:8]
-        right_qvel_raw = qvel_raw[8:16]
-        left_arm_qvel = left_qvel_raw[:6]
-        right_arm_qvel = right_qvel_raw[:6]
-        left_gripper_qvel = [normalize_puppet_gripper_velocity(left_qvel_raw[6])]
-        right_gripper_qvel = [normalize_puppet_gripper_velocity(right_qvel_raw[6])]
-        return np.concatenate([left_arm_qvel, left_gripper_qvel, right_arm_qvel, right_gripper_qvel])
-
-    @staticmethod
-    def get_env_state(physics):
-        raise NotImplementedError
-
-    def get_observation(self, physics):
-        obs = collections.OrderedDict()
-        obs["qpos"] = self.get_qpos(physics)
-        obs["qvel"] = self.get_qvel(physics)
-        obs["env_state"] = self.get_env_state(physics)
-        obs["images"] = {}
-        obs["images"]["top"] = physics.render(height=480, width=640, camera_id="top")
-        obs["images"]["angle"] = physics.render(height=480, width=640, camera_id="angle")
-        obs["images"]["vis"] = physics.render(height=480, width=640, camera_id="front_close")
-
-        return obs
-
-    def get_reward(self, physics):
-        # return whether left gripper is holding the box
-        raise NotImplementedError
-
-
-class TransferCubeTask(BimanualViperXTask):
-    def __init__(self, random=None):
-        super().__init__(random=random)
-        self.max_reward = 4
-
-    def initialize_episode(self, physics):
-        """Sets the state of the environment at the start of each episode."""
-        # TODO Notice: this function does not randomize the env configuration. Instead, set BOX_POSE from outside
-        # reset qpos, control and box position
-        with physics.reset_context():
-            physics.named.data.qpos[:16] = START_ARM_POSE
-            np.copyto(physics.data.ctrl, START_ARM_POSE)
-            assert BOX_POSE[0] is not None
-            physics.named.data.qpos[-7:] = BOX_POSE[0]
-            # print(f"{BOX_POSE=}")
-        super().initialize_episode(physics)
-
-    @staticmethod
-    def get_env_state(physics):
-        env_state = physics.data.qpos.copy()[16:]
-        return env_state
-
-    def get_reward(self, physics):
-        # return whether left gripper is holding the box
-        all_contact_pairs = []
-        for i_contact in range(physics.data.ncon):
-            id_geom_1 = physics.data.contact[i_contact].geom1
-            id_geom_2 = physics.data.contact[i_contact].geom2
-            name_geom_1 = physics.model.id2name(id_geom_1, "geom")
-            name_geom_2 = physics.model.id2name(id_geom_2, "geom")
-            contact_pair = (name_geom_1, name_geom_2)
-            all_contact_pairs.append(contact_pair)
-
-        touch_left_gripper = ("red_box", "vx300s_left/10_left_gripper_finger") in all_contact_pairs
-        touch_right_gripper = ("red_box", "vx300s_right/10_right_gripper_finger") in all_contact_pairs
-        touch_table = ("red_box", "table") in all_contact_pairs
-
-        reward = 0
-        if touch_right_gripper:
-            reward = 1
-        if touch_right_gripper and not touch_table:  # lifted
-            reward = 2
-        if touch_left_gripper:  # attempted transfer
-            reward = 3
-        if touch_left_gripper and not touch_table:  # successful transfer
-            reward = 4
-        return reward
-
-
-class InsertionTask(BimanualViperXTask):
-    def __init__(self, random=None):
-        super().__init__(random=random)
-        self.max_reward = 4
-
-    def initialize_episode(self, physics):
-        """Sets the state of the environment at the start of each episode."""
-        # TODO Notice: this function does not randomize the env configuration. Instead, set BOX_POSE from outside
-        # reset qpos, control and box position
-        with physics.reset_context():
-            physics.named.data.qpos[:16] = START_ARM_POSE
-            np.copyto(physics.data.ctrl, START_ARM_POSE)
-            assert BOX_POSE[0] is not None
-            physics.named.data.qpos[-7 * 2 :] = BOX_POSE[0]  # two objects
-            # print(f"{BOX_POSE=}")
-        super().initialize_episode(physics)
-
-    @staticmethod
-    def get_env_state(physics):
-        env_state = physics.data.qpos.copy()[16:]
-        return env_state
-
-    def get_reward(self, physics):
-        # return whether peg touches the pin
-        all_contact_pairs = []
-        for i_contact in range(physics.data.ncon):
-            id_geom_1 = physics.data.contact[i_contact].geom1
-            id_geom_2 = physics.data.contact[i_contact].geom2
-            name_geom_1 = physics.model.id2name(id_geom_1, "geom")
-            name_geom_2 = physics.model.id2name(id_geom_2, "geom")
-            contact_pair = (name_geom_1, name_geom_2)
-            all_contact_pairs.append(contact_pair)
-
-        touch_right_gripper = ("red_peg", "vx300s_right/10_right_gripper_finger") in all_contact_pairs
-        touch_left_gripper = (
-            ("socket-1", "vx300s_left/10_left_gripper_finger") in all_contact_pairs
-            or ("socket-2", "vx300s_left/10_left_gripper_finger") in all_contact_pairs
-            or ("socket-3", "vx300s_left/10_left_gripper_finger") in all_contact_pairs
-            or ("socket-4", "vx300s_left/10_left_gripper_finger") in all_contact_pairs
-        )
-
-        peg_touch_table = ("red_peg", "table") in all_contact_pairs
-        socket_touch_table = (
-            ("socket-1", "table") in all_contact_pairs
-            or ("socket-2", "table") in all_contact_pairs
-            or ("socket-3", "table") in all_contact_pairs
-            or ("socket-4", "table") in all_contact_pairs
-        )
-        peg_touch_socket = (
-            ("red_peg", "socket-1") in all_contact_pairs
-            or ("red_peg", "socket-2") in all_contact_pairs
-            or ("red_peg", "socket-3") in all_contact_pairs
-            or ("red_peg", "socket-4") in all_contact_pairs
-        )
-        pin_touched = ("red_peg", "pin") in all_contact_pairs
-
-        reward = 0
-        if touch_left_gripper and touch_right_gripper:  # touch both
-            reward = 1
-        if (
-            touch_left_gripper and touch_right_gripper and (not peg_touch_table) and (not socket_touch_table)
-        ):  # grasp both
-            reward = 2
-        if peg_touch_socket and (not peg_touch_table) and (not socket_touch_table):  # peg and socket touching
-            reward = 3
-        if pin_touched:  # successful insertion
-            reward = 4
-        return reward

lerobot/common/envs/aloha/tasks/sim_end_effector.py

@@ -1,263 +0,0 @@
-import collections
-
-import numpy as np
-from dm_control.suite import base
-
-from lerobot.common.envs.aloha.constants import (
-    PUPPET_GRIPPER_POSITION_CLOSE,
-    START_ARM_POSE,
-    normalize_puppet_gripper_position,
-    normalize_puppet_gripper_velocity,
-    unnormalize_puppet_gripper_position,
-)
-from lerobot.common.envs.aloha.utils import sample_box_pose, sample_insertion_pose
-
-"""
-Environment for simulated robot bi-manual manipulation, with end-effector control.
-Action space:      [left_arm_pose (7),             # position and quaternion for end effector
-                    left_gripper_positions (1),    # normalized gripper position (0: close, 1: open)
-                    right_arm_pose (7),            # position and quaternion for end effector
-                    right_gripper_positions (1),]  # normalized gripper position (0: close, 1: open)
-
-Observation space: {"qpos": Concat[ left_arm_qpos (6),         # absolute joint position
-                                    left_gripper_position (1),  # normalized gripper position (0: close, 1: open)
-                                    right_arm_qpos (6),         # absolute joint position
-                                    right_gripper_qpos (1)]     # normalized gripper position (0: close, 1: open)
-                    "qvel": Concat[ left_arm_qvel (6),         # absolute joint velocity (rad)
-                                    left_gripper_velocity (1),  # normalized gripper velocity (pos: opening, neg: closing)
-                                    right_arm_qvel (6),         # absolute joint velocity (rad)
-                                    right_gripper_qvel (1)]     # normalized gripper velocity (pos: opening, neg: closing)
-                    "images": {"main": (480x640x3)}        # h, w, c, dtype='uint8'
-"""
-
-
-class BimanualViperXEndEffectorTask(base.Task):
-    def __init__(self, random=None):
-        super().__init__(random=random)
-
-    def before_step(self, action, physics):
-        a_len = len(action) // 2
-        action_left = action[:a_len]
-        action_right = action[a_len:]
-
-        # set mocap position and quat
-        # left
-        np.copyto(physics.data.mocap_pos[0], action_left[:3])
-        np.copyto(physics.data.mocap_quat[0], action_left[3:7])
-        # right
-        np.copyto(physics.data.mocap_pos[1], action_right[:3])
-        np.copyto(physics.data.mocap_quat[1], action_right[3:7])
-
-        # set gripper
-        g_left_ctrl = unnormalize_puppet_gripper_position(action_left[7])
-        g_right_ctrl = unnormalize_puppet_gripper_position(action_right[7])
-        np.copyto(physics.data.ctrl, np.array([g_left_ctrl, -g_left_ctrl, g_right_ctrl, -g_right_ctrl]))
-
-    def initialize_robots(self, physics):
-        # reset joint position
-        physics.named.data.qpos[:16] = START_ARM_POSE
-
-        # reset mocap to align with end effector
-        # to obtain these numbers:
-        # (1) make an ee_sim env and reset to the same start_pose
-        # (2) get env._physics.named.data.xpos['vx300s_left/gripper_link']
-        #     get env._physics.named.data.xquat['vx300s_left/gripper_link']
-        #     repeat the same for right side
-        np.copyto(physics.data.mocap_pos[0], [-0.31718881, 0.5, 0.29525084])
-        np.copyto(physics.data.mocap_quat[0], [1, 0, 0, 0])
-        # right
-        np.copyto(physics.data.mocap_pos[1], np.array([0.31718881, 0.49999888, 0.29525084]))
-        np.copyto(physics.data.mocap_quat[1], [1, 0, 0, 0])
-
-        # reset gripper control
-        close_gripper_control = np.array(
-            [
-                PUPPET_GRIPPER_POSITION_CLOSE,
-                -PUPPET_GRIPPER_POSITION_CLOSE,
-                PUPPET_GRIPPER_POSITION_CLOSE,
-                -PUPPET_GRIPPER_POSITION_CLOSE,
-            ]
-        )
-        np.copyto(physics.data.ctrl, close_gripper_control)
-
-    def initialize_episode(self, physics):
-        """Sets the state of the environment at the start of each episode."""
-        super().initialize_episode(physics)
-
-    @staticmethod
-    def get_qpos(physics):
-        qpos_raw = physics.data.qpos.copy()
-        left_qpos_raw = qpos_raw[:8]
-        right_qpos_raw = qpos_raw[8:16]
-        left_arm_qpos = left_qpos_raw[:6]
-        right_arm_qpos = right_qpos_raw[:6]
-        left_gripper_qpos = [normalize_puppet_gripper_position(left_qpos_raw[6])]
-        right_gripper_qpos = [normalize_puppet_gripper_position(right_qpos_raw[6])]
-        return np.concatenate([left_arm_qpos, left_gripper_qpos, right_arm_qpos, right_gripper_qpos])
-
-    @staticmethod
-    def get_qvel(physics):
-        qvel_raw = physics.data.qvel.copy()
-        left_qvel_raw = qvel_raw[:8]
-        right_qvel_raw = qvel_raw[8:16]
-        left_arm_qvel = left_qvel_raw[:6]
-        right_arm_qvel = right_qvel_raw[:6]
-        left_gripper_qvel = [normalize_puppet_gripper_velocity(left_qvel_raw[6])]
-        right_gripper_qvel = [normalize_puppet_gripper_velocity(right_qvel_raw[6])]
-        return np.concatenate([left_arm_qvel, left_gripper_qvel, right_arm_qvel, right_gripper_qvel])
-
-    @staticmethod
-    def get_env_state(physics):
-        raise NotImplementedError
-
-    def get_observation(self, physics):
-        # note: it is important to do .copy()
-        obs = collections.OrderedDict()
-        obs["qpos"] = self.get_qpos(physics)
-        obs["qvel"] = self.get_qvel(physics)
-        obs["env_state"] = self.get_env_state(physics)
-        obs["images"] = {}
-        obs["images"]["top"] = physics.render(height=480, width=640, camera_id="top")
-        obs["images"]["angle"] = physics.render(height=480, width=640, camera_id="angle")
-        obs["images"]["vis"] = physics.render(height=480, width=640, camera_id="front_close")
-        # used in scripted policy to obtain starting pose
-        obs["mocap_pose_left"] = np.concatenate(
-            [physics.data.mocap_pos[0], physics.data.mocap_quat[0]]
-        ).copy()
-        obs["mocap_pose_right"] = np.concatenate(
-            [physics.data.mocap_pos[1], physics.data.mocap_quat[1]]
-        ).copy()
-
-        # used when replaying joint trajectory
-        obs["gripper_ctrl"] = physics.data.ctrl.copy()
-        return obs
-
-    def get_reward(self, physics):
-        raise NotImplementedError
-
-
-class TransferCubeEndEffectorTask(BimanualViperXEndEffectorTask):
-    def __init__(self, random=None):
-        super().__init__(random=random)
-        self.max_reward = 4
-
-    def initialize_episode(self, physics):
-        """Sets the state of the environment at the start of each episode."""
-        self.initialize_robots(physics)
-        # randomize box position
-        cube_pose = sample_box_pose()
-        box_start_idx = physics.model.name2id("red_box_joint", "joint")
-        np.copyto(physics.data.qpos[box_start_idx : box_start_idx + 7], cube_pose)
-        # print(f"randomized cube position to {cube_position}")
-
-        super().initialize_episode(physics)
-
-    @staticmethod
-    def get_env_state(physics):
-        env_state = physics.data.qpos.copy()[16:]
-        return env_state
-
-    def get_reward(self, physics):
-        # return whether left gripper is holding the box
-        all_contact_pairs = []
-        for i_contact in range(physics.data.ncon):
-            id_geom_1 = physics.data.contact[i_contact].geom1
-            id_geom_2 = physics.data.contact[i_contact].geom2
-            name_geom_1 = physics.model.id2name(id_geom_1, "geom")
-            name_geom_2 = physics.model.id2name(id_geom_2, "geom")
-            contact_pair = (name_geom_1, name_geom_2)
-            all_contact_pairs.append(contact_pair)
-
-        touch_left_gripper = ("red_box", "vx300s_left/10_left_gripper_finger") in all_contact_pairs
-        touch_right_gripper = ("red_box", "vx300s_right/10_right_gripper_finger") in all_contact_pairs
-        touch_table = ("red_box", "table") in all_contact_pairs
-
-        reward = 0
-        if touch_right_gripper:
-            reward = 1
-        if touch_right_gripper and not touch_table:  # lifted
-            reward = 2
-        if touch_left_gripper:  # attempted transfer
-            reward = 3
-        if touch_left_gripper and not touch_table:  # successful transfer
-            reward = 4
-        return reward
-
-
-class InsertionEndEffectorTask(BimanualViperXEndEffectorTask):
-    def __init__(self, random=None):
-        super().__init__(random=random)
-        self.max_reward = 4
-
-    def initialize_episode(self, physics):
-        """Sets the state of the environment at the start of each episode."""
-        self.initialize_robots(physics)
-        # randomize peg and socket position
-        peg_pose, socket_pose = sample_insertion_pose()
-
-        def id2index(j_id):
-            return 16 + (j_id - 16) * 7  # first 16 is robot qpos, 7 is pose dim # hacky
-
-        peg_start_id = physics.model.name2id("red_peg_joint", "joint")
-        peg_start_idx = id2index(peg_start_id)
-        np.copyto(physics.data.qpos[peg_start_idx : peg_start_idx + 7], peg_pose)
-        # print(f"randomized cube position to {cube_position}")
-
-        socket_start_id = physics.model.name2id("blue_socket_joint", "joint")
-        socket_start_idx = id2index(socket_start_id)
-        np.copyto(physics.data.qpos[socket_start_idx : socket_start_idx + 7], socket_pose)
-        # print(f"randomized cube position to {cube_position}")
-
-        super().initialize_episode(physics)
-
-    @staticmethod
-    def get_env_state(physics):
-        env_state = physics.data.qpos.copy()[16:]
-        return env_state
-
-    def get_reward(self, physics):
-        # return whether peg touches the pin
-        all_contact_pairs = []
-        for i_contact in range(physics.data.ncon):
-            id_geom_1 = physics.data.contact[i_contact].geom1
-            id_geom_2 = physics.data.contact[i_contact].geom2
-            name_geom_1 = physics.model.id2name(id_geom_1, "geom")
-            name_geom_2 = physics.model.id2name(id_geom_2, "geom")
-            contact_pair = (name_geom_1, name_geom_2)
-            all_contact_pairs.append(contact_pair)
-
-        touch_right_gripper = ("red_peg", "vx300s_right/10_right_gripper_finger") in all_contact_pairs
-        touch_left_gripper = (
-            ("socket-1", "vx300s_left/10_left_gripper_finger") in all_contact_pairs
-            or ("socket-2", "vx300s_left/10_left_gripper_finger") in all_contact_pairs
-            or ("socket-3", "vx300s_left/10_left_gripper_finger") in all_contact_pairs
-            or ("socket-4", "vx300s_left/10_left_gripper_finger") in all_contact_pairs
-        )
-
-        peg_touch_table = ("red_peg", "table") in all_contact_pairs
-        socket_touch_table = (
-            ("socket-1", "table") in all_contact_pairs
-            or ("socket-2", "table") in all_contact_pairs
-            or ("socket-3", "table") in all_contact_pairs
-            or ("socket-4", "table") in all_contact_pairs
-        )
-        peg_touch_socket = (
-            ("red_peg", "socket-1") in all_contact_pairs
-            or ("red_peg", "socket-2") in all_contact_pairs
-            or ("red_peg", "socket-3") in all_contact_pairs
-            or ("red_peg", "socket-4") in all_contact_pairs
-        )
-        pin_touched = ("red_peg", "pin") in all_contact_pairs
-
-        reward = 0
-        if touch_left_gripper and touch_right_gripper:  # touch both
-            reward = 1
-        if (
-            touch_left_gripper and touch_right_gripper and (not peg_touch_table) and (not socket_touch_table)
-        ):  # grasp both
-            reward = 2
-        if peg_touch_socket and (not peg_touch_table) and (not socket_touch_table):  # peg and socket touching
-            reward = 3
-        if pin_touched:  # successful insertion
-            reward = 4
-        return reward

lerobot/common/envs/aloha/utils.py

@@ -1,39 +0,0 @@
-import numpy as np
-
-
-def sample_box_pose():
-    x_range = [0.0, 0.2]
-    y_range = [0.4, 0.6]
-    z_range = [0.05, 0.05]
-
-    ranges = np.vstack([x_range, y_range, z_range])
-    cube_position = np.random.uniform(ranges[:, 0], ranges[:, 1])
-
-    cube_quat = np.array([1, 0, 0, 0])
-    return np.concatenate([cube_position, cube_quat])
-
-
-def sample_insertion_pose():
-    # Peg
-    x_range = [0.1, 0.2]
-    y_range = [0.4, 0.6]
-    z_range = [0.05, 0.05]
-
-    ranges = np.vstack([x_range, y_range, z_range])
-    peg_position = np.random.uniform(ranges[:, 0], ranges[:, 1])
-
-    peg_quat = np.array([1, 0, 0, 0])
-    peg_pose = np.concatenate([peg_position, peg_quat])
-
-    # Socket
-    x_range = [-0.2, -0.1]
-    y_range = [0.4, 0.6]
-    z_range = [0.05, 0.05]
-
-    ranges = np.vstack([x_range, y_range, z_range])
-    socket_position = np.random.uniform(ranges[:, 0], ranges[:, 1])
-
-    socket_quat = np.array([1, 0, 0, 0])
-    socket_pose = np.concatenate([socket_position, socket_quat])
-
-    return peg_pose, socket_pose

lerobot/common/envs/factory.py

@@ -1,64 +1,43 @@
-from torchrl.envs import SerialEnv
-from torchrl.envs.transforms import Compose, StepCounter, Transform, TransformedEnv
+import importlib
+
+import gymnasium as gym
 
 
-def make_env(cfg, transform=None):
+def make_env(cfg, num_parallel_envs=0) -> gym.Env | gym.vector.SyncVectorEnv:
     """
-    Note: The returned environment is wrapped in a torchrl.SerialEnv with cfg.rollout_batch_size underlying
-    environments. The env therefore returns batches.`
+    Note: When `num_parallel_envs > 0`, this function returns a `SyncVectorEnv` which takes batched action as input and
+    returns batched observation, reward, terminated, truncated of `num_parallel_envs` items.
     """
-
     kwargs = {
-        "frame_skip": cfg.env.action_repeat,
-        "from_pixels": cfg.env.from_pixels,
-        "pixels_only": cfg.env.pixels_only,
-        "image_size": cfg.env.image_size,
-        "num_prev_obs": cfg.n_obs_steps - 1,
+        "obs_type": "pixels_agent_pos",
+        "render_mode": "rgb_array",
+        "max_episode_steps": cfg.env.episode_length,
+        "visualization_width": 384,
+        "visualization_height": 384,
     }
 
-    if cfg.env.name == "simxarm":
-        from lerobot.common.envs.simxarm.env import SimxarmEnv
+    package_name = f"gym_{cfg.env.name}"
 
-        kwargs["task"] = cfg.env.task
-        clsfunc = SimxarmEnv
-    elif cfg.env.name == "pusht":
-        from lerobot.common.envs.pusht.env import PushtEnv
+    try:
+        importlib.import_module(package_name)
+    except ModuleNotFoundError as e:
+        print(
+            f"{package_name} is not installed. Please install it with `pip install 'lerobot[{cfg.env.name}]'`"
+        )
+        raise e
 
-        # assert kwargs["seed"] > 200, "Seed 0-200 are used for the demonstration dataset, so we don't want to seed the eval env with this range."
+    gym_handle = f"{package_name}/{cfg.env.task}"
 
-        clsfunc = PushtEnv
-    elif cfg.env.name == "aloha":
-        from lerobot.common.envs.aloha.env import AlohaEnv
-
-        kwargs["task"] = cfg.env.task
-        clsfunc = AlohaEnv
+    if num_parallel_envs == 0:
+        # non-batched version of the env that returns an observation of shape (c)
+        env = gym.make(gym_handle, disable_env_checker=True, **kwargs)
     else:
-        raise ValueError(cfg.env.name)
+        # batched version of the env that returns an observation of shape (b, c)
+        env = gym.vector.SyncVectorEnv(
+            [
+                lambda: gym.make(gym_handle, disable_env_checker=True, **kwargs)
+                for _ in range(num_parallel_envs)
+            ]
+        )
 
-    def _make_env(seed):
-        nonlocal kwargs
-        kwargs["seed"] = seed
-        env = clsfunc(**kwargs)
-
-        # limit rollout to max_steps
-        env = TransformedEnv(env, StepCounter(max_steps=cfg.env.episode_length))
-
-        if transform is not None:
-            # useful to add normalization
-            if isinstance(transform, Compose):
-                for tf in transform:
-                    env.append_transform(tf.clone())
-            elif isinstance(transform, Transform):
-                env.append_transform(transform.clone())
-            else:
-                raise NotImplementedError()
-
-        return env
-
-    return SerialEnv(
-        cfg.rollout_batch_size,
-        create_env_fn=_make_env,
-        create_env_kwargs=[
-            {"seed": env_seed} for env_seed in range(cfg.seed, cfg.seed + cfg.rollout_batch_size)
-        ],
-    )
+    return env

lerobot/common/envs/pusht/env.py

@@ -1,245 +0,0 @@
-import importlib
-import logging
-from collections import deque
-from typing import Optional
-
-import cv2
-import numpy as np
-import torch
-from tensordict import TensorDict
-from torchrl.data.tensor_specs import (
-    BoundedTensorSpec,
-    CompositeSpec,
-    DiscreteTensorSpec,
-    UnboundedContinuousTensorSpec,
-)
-from torchrl.envs.libs.gym import _gym_to_torchrl_spec_transform
-
-from lerobot.common.envs.abstract import AbstractEnv
-from lerobot.common.utils import set_global_seed
-
-_has_gym = importlib.util.find_spec("gymnasium") is not None
-
-
-class PushtEnv(AbstractEnv):
-    name = "pusht"
-    available_tasks = ["pusht"]
-    _reset_warning_issued = False
-
-    def __init__(
-        self,
-        task="pusht",
-        frame_skip: int = 1,
-        from_pixels: bool = False,
-        pixels_only: bool = False,
-        image_size=None,
-        seed=1337,
-        device="cpu",
-        num_prev_obs=1,
-        num_prev_action=0,
-    ):
-        super().__init__(
-            task=task,
-            frame_skip=frame_skip,
-            from_pixels=from_pixels,
-            pixels_only=pixels_only,
-            image_size=image_size,
-            seed=seed,
-            device=device,
-            num_prev_obs=num_prev_obs,
-            num_prev_action=num_prev_action,
-        )
-
-    def _make_env(self):
-        if not _has_gym:
-            raise ImportError("Cannot import gymnasium.")
-
-        # TODO(rcadene) (PushTEnv is similar to PushTImageEnv, but without the image rendering, it's faster to iterate on)
-        # from lerobot.common.envs.pusht.pusht_env import PushTEnv
-
-        if not self.from_pixels:
-            raise NotImplementedError("Use PushTEnv, instead of PushTImageEnv")
-        from lerobot.common.envs.pusht.pusht_image_env import PushTImageEnv
-
-        self._env = PushTImageEnv(render_size=self.image_size)
-
-    def render(self, mode="rgb_array", width=96, height=96, with_marker=True):
-        """
-        with_marker adds a cursor showing the targeted action for the controller.
-        """
-        if width != height:
-            raise NotImplementedError()
-        tmp = self._env.render_size
-        if width != self._env.render_size:
-            self._env.render_cache = None
-            self._env.render_size = width
-        out = self._env.render(mode).copy()
-        if with_marker and self._env.latest_action is not None:
-            action = np.array(self._env.latest_action)
-            coord = (action / 512 * self._env.render_size).astype(np.int32)
-            marker_size = int(8 / 96 * self._env.render_size)
-            thickness = int(1 / 96 * self._env.render_size)
-            cv2.drawMarker(
-                out,
-                coord,
-                color=(255, 0, 0),
-                markerType=cv2.MARKER_CROSS,
-                markerSize=marker_size,
-                thickness=thickness,
-            )
-        self._env.render_size = tmp
-        return out
-
-    def _format_raw_obs(self, raw_obs):
-        if self.from_pixels:
-            image = torch.from_numpy(raw_obs["image"])
-            obs = {"image": image}
-
-            if not self.pixels_only:
-                obs["state"] = torch.from_numpy(raw_obs["agent_pos"]).type(torch.float32)
-        else:
-            # TODO:
-            obs = {"state": torch.from_numpy(raw_obs["observation"]).type(torch.float32)}
-
-        return obs
-
-    def _reset(self, tensordict: Optional[TensorDict] = None):
-        if tensordict is not None and not PushtEnv._reset_warning_issued:
-            logging.warning(f"{self.__class__.__name__}._reset ignores the provided tensordict.")
-            PushtEnv._reset_warning_issued = True
-
-        # Seed the environment and update the seed to be used for the next reset.
-        self._next_seed = self.set_seed(self._next_seed)
-        raw_obs = self._env.reset()
-
-        obs = self._format_raw_obs(raw_obs)
-
-        if self.num_prev_obs > 0:
-            stacked_obs = {}
-            if "image" in obs:
-                self._prev_obs_image_queue = deque(
-                    [obs["image"]] * (self.num_prev_obs + 1), maxlen=(self.num_prev_obs + 1)
-                )
-                stacked_obs["image"] = torch.stack(list(self._prev_obs_image_queue))
-            if "state" in obs:
-                self._prev_obs_state_queue = deque(
-                    [obs["state"]] * (self.num_prev_obs + 1), maxlen=(self.num_prev_obs + 1)
-                )
-                stacked_obs["state"] = torch.stack(list(self._prev_obs_state_queue))
-            obs = stacked_obs
-
-        td = TensorDict(
-            {
-                "observation": TensorDict(obs, batch_size=[]),
-                "done": torch.tensor([False], dtype=torch.bool),
-            },
-            batch_size=[],
-        )
-
-        return td
-
-    def _step(self, tensordict: TensorDict):
-        td = tensordict
-        action = td["action"].numpy()
-        assert action.ndim == 1
-        # TODO(rcadene): add info["is_success"] and info["success"] ?
-
-        raw_obs, reward, done, info = self._env.step(action)
-
-        obs = self._format_raw_obs(raw_obs)
-
-        if self.num_prev_obs > 0:
-            stacked_obs = {}
-            if "image" in obs:
-                self._prev_obs_image_queue.append(obs["image"])
-                stacked_obs["image"] = torch.stack(list(self._prev_obs_image_queue))
-            if "state" in obs:
-                self._prev_obs_state_queue.append(obs["state"])
-                stacked_obs["state"] = torch.stack(list(self._prev_obs_state_queue))
-            obs = stacked_obs
-
-        td = TensorDict(
-            {
-                "observation": TensorDict(obs, batch_size=[]),
-                "reward": torch.tensor([reward], dtype=torch.float32),
-                # success and done are true when coverage > self.success_threshold in env
-                "done": torch.tensor([done], dtype=torch.bool),
-                "success": torch.tensor([done], dtype=torch.bool),
-            },
-            batch_size=[],
-        )
-        return td
-
-    def _make_spec(self):
-        obs = {}
-        if self.from_pixels:
-            image_shape = (3, self.image_size, self.image_size)
-            if self.num_prev_obs > 0:
-                image_shape = (self.num_prev_obs + 1, *image_shape)
-
-            obs["image"] = BoundedTensorSpec(
-                low=0,
-                high=255,
-                shape=image_shape,
-                dtype=torch.uint8,
-                device=self.device,
-            )
-            if not self.pixels_only:
-                state_shape = self._env.observation_space["agent_pos"].shape
-                if self.num_prev_obs > 0:
-                    state_shape = (self.num_prev_obs + 1, *state_shape)
-
-                obs["state"] = BoundedTensorSpec(
-                    low=0,
-                    high=512,
-                    shape=state_shape,
-                    dtype=torch.float32,
-                    device=self.device,
-                )
-        else:
-            # TODO(rcadene): add observation_space achieved_goal and desired_goal?
-            state_shape = self._env.observation_space["observation"].shape
-            if self.num_prev_obs > 0:
-                state_shape = (self.num_prev_obs + 1, *state_shape)
-
-            obs["state"] = UnboundedContinuousTensorSpec(
-                # TODO:
-                shape=state_shape,
-                dtype=torch.float32,
-                device=self.device,
-            )
-        self.observation_spec = CompositeSpec({"observation": obs})
-
-        self.action_spec = _gym_to_torchrl_spec_transform(
-            self._env.action_space,
-            device=self.device,
-        )
-
-        self.reward_spec = UnboundedContinuousTensorSpec(
-            shape=(1,),
-            dtype=torch.float32,
-            device=self.device,
-        )
-
-        self.done_spec = CompositeSpec(
-            {
-                "done": DiscreteTensorSpec(
-                    2,
-                    shape=(1,),
-                    dtype=torch.bool,
-                    device=self.device,
-                ),
-                "success": DiscreteTensorSpec(
-                    2,
-                    shape=(1,),
-                    dtype=torch.bool,
-                    device=self.device,
-                ),
-            }
-        )
-
-    def _set_seed(self, seed: Optional[int]):
-        # Set global seed.
-        set_global_seed(seed)
-        # Set PushTImageEnv seed as it relies on it's own internal _seed attribute.
-        self._env.seed(seed)

lerobot/common/envs/pusht/pusht_env.py

@@ -1,378 +0,0 @@
-import collections
-
-import cv2
-import gymnasium as gym
-import numpy as np
-import pygame
-import pymunk
-import pymunk.pygame_util
-import shapely.geometry as sg
-import skimage.transform as st
-from gymnasium import spaces
-from pymunk.vec2d import Vec2d
-
-from lerobot.common.envs.pusht.pymunk_override import DrawOptions
-
-
-def pymunk_to_shapely(body, shapes):
-    geoms = []
-    for shape in shapes:
-        if isinstance(shape, pymunk.shapes.Poly):
-            verts = [body.local_to_world(v) for v in shape.get_vertices()]
-            verts += [verts[0]]
-            geoms.append(sg.Polygon(verts))
-        else:
-            raise RuntimeError(f"Unsupported shape type {type(shape)}")
-    geom = sg.MultiPolygon(geoms)
-    return geom
-
-
-class PushTEnv(gym.Env):
-    metadata = {"render.modes": ["human", "rgb_array"], "video.frames_per_second": 10}
-    reward_range = (0.0, 1.0)
-
-    def __init__(
-        self,
-        legacy=True,  # compatibility with original
-        block_cog=None,
-        damping=None,
-        render_action=True,
-        render_size=96,
-        reset_to_state=None,
-    ):
-        self._seed = None
-        self.seed()
-        self.window_size = ws = 512  # The size of the PyGame window
-        self.render_size = render_size
-        self.sim_hz = 100
-        # Local controller params.
-        self.k_p, self.k_v = 100, 20  # PD control.z
-        self.control_hz = self.metadata["video.frames_per_second"]
-        # legcay set_state for data compatibility
-        self.legacy = legacy
-
-        # agent_pos, block_pos, block_angle
-        self.observation_space = spaces.Box(
-            low=np.array([0, 0, 0, 0, 0], dtype=np.float64),
-            high=np.array([ws, ws, ws, ws, np.pi * 2], dtype=np.float64),
-            shape=(5,),
-            dtype=np.float64,
-        )
-
-        # positional goal for agent
-        self.action_space = spaces.Box(
-            low=np.array([0, 0], dtype=np.float64),
-            high=np.array([ws, ws], dtype=np.float64),
-            shape=(2,),
-            dtype=np.float64,
-        )
-
-        self.block_cog = block_cog
-        self.damping = damping
-        self.render_action = render_action
-
-        """
-        If human-rendering is used, `self.window` will be a reference
-        to the window that we draw to. `self.clock` will be a clock that is used
-        to ensure that the environment is rendered at the correct framerate in
-        human-mode. They will remain `None` until human-mode is used for the
-        first time.
-        """
-        self.window = None
-        self.clock = None
-        self.screen = None
-
-        self.space = None
-        self.teleop = None
-        self.render_buffer = None
-        self.latest_action = None
-        self.reset_to_state = reset_to_state
-
-    def reset(self):
-        seed = self._seed
-        self._setup()
-        if self.block_cog is not None:
-            self.block.center_of_gravity = self.block_cog
-        if self.damping is not None:
-            self.space.damping = self.damping
-
-        # use legacy RandomState for compatibility
-        state = self.reset_to_state
-        if state is None:
-            rs = np.random.RandomState(seed=seed)
-            state = np.array(
-                [
-                    rs.randint(50, 450),
-                    rs.randint(50, 450),
-                    rs.randint(100, 400),
-                    rs.randint(100, 400),
-                    rs.randn() * 2 * np.pi - np.pi,
-                ]
-            )
-        self._set_state(state)
-
-        observation = self._get_obs()
-        return observation
-
-    def step(self, action):
-        dt = 1.0 / self.sim_hz
-        self.n_contact_points = 0
-        n_steps = self.sim_hz // self.control_hz
-        if action is not None:
-            self.latest_action = action
-            for _ in range(n_steps):
-                # Step PD control.
-                # self.agent.velocity = self.k_p * (act - self.agent.position)    # P control works too.
-                acceleration = self.k_p * (action - self.agent.position) + self.k_v * (
-                    Vec2d(0, 0) - self.agent.velocity
-                )
-                self.agent.velocity += acceleration * dt
-
-                # Step physics.
-                self.space.step(dt)
-
-        # compute reward
-        goal_body = self._get_goal_pose_body(self.goal_pose)
-        goal_geom = pymunk_to_shapely(goal_body, self.block.shapes)
-        block_geom = pymunk_to_shapely(self.block, self.block.shapes)
-
-        intersection_area = goal_geom.intersection(block_geom).area
-        goal_area = goal_geom.area
-        coverage = intersection_area / goal_area
-        reward = np.clip(coverage / self.success_threshold, 0, 1)
-        done = coverage > self.success_threshold
-
-        observation = self._get_obs()
-        info = self._get_info()
-
-        return observation, reward, done, info
-
-    def render(self, mode):
-        return self._render_frame(mode)
-
-    def teleop_agent(self):
-        TeleopAgent = collections.namedtuple("TeleopAgent", ["act"])
-
-        def act(obs):
-            act = None
-            mouse_position = pymunk.pygame_util.from_pygame(Vec2d(*pygame.mouse.get_pos()), self.screen)
-            if self.teleop or (mouse_position - self.agent.position).length < 30:
-                self.teleop = True
-                act = mouse_position
-            return act
-
-        return TeleopAgent(act)
-
-    def _get_obs(self):
-        obs = np.array(
-            tuple(self.agent.position) + tuple(self.block.position) + (self.block.angle % (2 * np.pi),)
-        )
-        return obs
-
-    def _get_goal_pose_body(self, pose):
-        mass = 1
-        inertia = pymunk.moment_for_box(mass, (50, 100))
-        body = pymunk.Body(mass, inertia)
-        # preserving the legacy assignment order for compatibility
-        # the order here doesn't matter somehow, maybe because CoM is aligned with body origin
-        body.position = pose[:2].tolist()
-        body.angle = pose[2]
-        return body
-
-    def _get_info(self):
-        n_steps = self.sim_hz // self.control_hz
-        n_contact_points_per_step = int(np.ceil(self.n_contact_points / n_steps))
-        info = {
-            "pos_agent": np.array(self.agent.position),
-            "vel_agent": np.array(self.agent.velocity),
-            "block_pose": np.array(list(self.block.position) + [self.block.angle]),
-            "goal_pose": self.goal_pose,
-            "n_contacts": n_contact_points_per_step,
-        }
-        return info
-
-    def _render_frame(self, mode):
-        if self.window is None and mode == "human":
-            pygame.init()
-            pygame.display.init()
-            self.window = pygame.display.set_mode((self.window_size, self.window_size))
-        if self.clock is None and mode == "human":
-            self.clock = pygame.time.Clock()
-
-        canvas = pygame.Surface((self.window_size, self.window_size))
-        canvas.fill((255, 255, 255))
-        self.screen = canvas
-
-        draw_options = DrawOptions(canvas)
-
-        # Draw goal pose.
-        goal_body = self._get_goal_pose_body(self.goal_pose)
-        for shape in self.block.shapes:
-            goal_points = [
-                pymunk.pygame_util.to_pygame(goal_body.local_to_world(v), draw_options.surface)
-                for v in shape.get_vertices()
-            ]
-            goal_points += [goal_points[0]]
-            pygame.draw.polygon(canvas, self.goal_color, goal_points)
-
-        # Draw agent and block.
-        self.space.debug_draw(draw_options)
-
-        if mode == "human":
-            # The following line copies our drawings from `canvas` to the visible window
-            self.window.blit(canvas, canvas.get_rect())
-            pygame.event.pump()
-            pygame.display.update()
-
-            # the clock is already ticked during in step for "human"
-
-        img = np.transpose(np.array(pygame.surfarray.pixels3d(canvas)), axes=(1, 0, 2))
-        img = cv2.resize(img, (self.render_size, self.render_size))
-        if self.render_action and self.latest_action is not None:
-            action = np.array(self.latest_action)
-            coord = (action / 512 * 96).astype(np.int32)
-            marker_size = int(8 / 96 * self.render_size)
-            thickness = int(1 / 96 * self.render_size)
-            cv2.drawMarker(
-                img,
-                coord,
-                color=(255, 0, 0),
-                markerType=cv2.MARKER_CROSS,
-                markerSize=marker_size,
-                thickness=thickness,
-            )
-        return img
-
-    def close(self):
-        if self.window is not None:
-            pygame.display.quit()
-            pygame.quit()
-
-    def seed(self, seed=None):
-        if seed is None:
-            seed = np.random.randint(0, 25536)
-        self._seed = seed
-        self.np_random = np.random.default_rng(seed)
-
-    def _handle_collision(self, arbiter, space, data):
-        self.n_contact_points += len(arbiter.contact_point_set.points)
-
-    def _set_state(self, state):
-        if isinstance(state, np.ndarray):
-            state = state.tolist()
-        pos_agent = state[:2]
-        pos_block = state[2:4]
-        rot_block = state[4]
-        self.agent.position = pos_agent
-        # setting angle rotates with respect to center of mass
-        # therefore will modify the geometric position
-        # if not the same as CoM
-        # therefore should be modified first.
-        if self.legacy:
-            # for compatibility with legacy data
-            self.block.position = pos_block
-            self.block.angle = rot_block
-        else:
-            self.block.angle = rot_block
-            self.block.position = pos_block
-
-        # Run physics to take effect
-        self.space.step(1.0 / self.sim_hz)
-
-    def _set_state_local(self, state_local):
-        agent_pos_local = state_local[:2]
-        block_pose_local = state_local[2:]
-        tf_img_obj = st.AffineTransform(translation=self.goal_pose[:2], rotation=self.goal_pose[2])
-        tf_obj_new = st.AffineTransform(translation=block_pose_local[:2], rotation=block_pose_local[2])
-        tf_img_new = st.AffineTransform(matrix=tf_img_obj.params @ tf_obj_new.params)
-        agent_pos_new = tf_img_new(agent_pos_local)
-        new_state = np.array(list(agent_pos_new[0]) + list(tf_img_new.translation) + [tf_img_new.rotation])
-        self._set_state(new_state)
-        return new_state
-
-    def _setup(self):
-        self.space = pymunk.Space()
-        self.space.gravity = 0, 0
-        self.space.damping = 0
-        self.teleop = False
-        self.render_buffer = []
-
-        # Add walls.
-        walls = [
-            self._add_segment((5, 506), (5, 5), 2),
-            self._add_segment((5, 5), (506, 5), 2),
-            self._add_segment((506, 5), (506, 506), 2),
-            self._add_segment((5, 506), (506, 506), 2),
-        ]
-        self.space.add(*walls)
-
-        # Add agent, block, and goal zone.
-        self.agent = self.add_circle((256, 400), 15)
-        self.block = self.add_tee((256, 300), 0)
-        self.goal_color = pygame.Color("LightGreen")
-        self.goal_pose = np.array([256, 256, np.pi / 4])  # x, y, theta (in radians)
-
-        # Add collision handling
-        self.collision_handeler = self.space.add_collision_handler(0, 0)
-        self.collision_handeler.post_solve = self._handle_collision
-        self.n_contact_points = 0
-
-        self.max_score = 50 * 100
-        self.success_threshold = 0.95  # 95% coverage.
-
-    def _add_segment(self, a, b, radius):
-        shape = pymunk.Segment(self.space.static_body, a, b, radius)
-        shape.color = pygame.Color("LightGray")  # https://htmlcolorcodes.com/color-names
-        return shape
-
-    def add_circle(self, position, radius):
-        body = pymunk.Body(body_type=pymunk.Body.KINEMATIC)
-        body.position = position
-        body.friction = 1
-        shape = pymunk.Circle(body, radius)
-        shape.color = pygame.Color("RoyalBlue")
-        self.space.add(body, shape)
-        return body
-
-    def add_box(self, position, height, width):
-        mass = 1
-        inertia = pymunk.moment_for_box(mass, (height, width))
-        body = pymunk.Body(mass, inertia)
-        body.position = position
-        shape = pymunk.Poly.create_box(body, (height, width))
-        shape.color = pygame.Color("LightSlateGray")
-        self.space.add(body, shape)
-        return body
-
-    def add_tee(self, position, angle, scale=30, color="LightSlateGray", mask=None):
-        if mask is None:
-            mask = pymunk.ShapeFilter.ALL_MASKS()
-        mass = 1
-        length = 4
-        vertices1 = [
-            (-length * scale / 2, scale),
-            (length * scale / 2, scale),
-            (length * scale / 2, 0),
-            (-length * scale / 2, 0),
-        ]
-        inertia1 = pymunk.moment_for_poly(mass, vertices=vertices1)
-        vertices2 = [
-            (-scale / 2, scale),
-            (-scale / 2, length * scale),
-            (scale / 2, length * scale),
-            (scale / 2, scale),
-        ]
-        inertia2 = pymunk.moment_for_poly(mass, vertices=vertices1)
-        body = pymunk.Body(mass, inertia1 + inertia2)
-        shape1 = pymunk.Poly(body, vertices1)
-        shape2 = pymunk.Poly(body, vertices2)
-        shape1.color = pygame.Color(color)
-        shape2.color = pygame.Color(color)
-        shape1.filter = pymunk.ShapeFilter(mask=mask)
-        shape2.filter = pymunk.ShapeFilter(mask=mask)
-        body.center_of_gravity = (shape1.center_of_gravity + shape2.center_of_gravity) / 2
-        body.position = position
-        body.angle = angle
-        body.friction = 1
-        self.space.add(body, shape1, shape2)
-        return body

lerobot/common/envs/pusht/pusht_image_env.py

@@ -1,41 +0,0 @@
-import numpy as np
-from gymnasium import spaces
-
-from lerobot.common.envs.pusht.pusht_env import PushTEnv
-
-
-class PushTImageEnv(PushTEnv):
-    metadata = {"render.modes": ["rgb_array"], "video.frames_per_second": 10}
-
-    # Note: legacy defaults to True for compatibility with original
-    def __init__(self, legacy=True, block_cog=None, damping=None, render_size=96):
-        super().__init__(
-            legacy=legacy, block_cog=block_cog, damping=damping, render_size=render_size, render_action=False
-        )
-        ws = self.window_size
-        self.observation_space = spaces.Dict(
-            {
-                "image": spaces.Box(low=0, high=1, shape=(3, render_size, render_size), dtype=np.float32),
-                "agent_pos": spaces.Box(low=0, high=ws, shape=(2,), dtype=np.float32),
-            }
-        )
-        self.render_cache = None
-
-    def _get_obs(self):
-        img = super()._render_frame(mode="rgb_array")
-
-        agent_pos = np.array(self.agent.position)
-        img_obs = np.moveaxis(img, -1, 0)
-        obs = {"image": img_obs, "agent_pos": agent_pos}
-
-        self.render_cache = img
-
-        return obs
-
-    def render(self, mode):
-        assert mode == "rgb_array"
-
-        if self.render_cache is None:
-            self._get_obs()
-
-        return self.render_cache

lerobot/common/envs/pusht/pymunk_override.py

@@ -1,244 +0,0 @@
-# ----------------------------------------------------------------------------
-# pymunk
-# Copyright (c) 2007-2016 Victor Blomqvist
-#
-# 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.
-# ----------------------------------------------------------------------------
-
-"""This submodule contains helper functions to help with quick prototyping
-using pymunk together with pygame.
-
-Intended to help with debugging and prototyping, not for actual production use
-in a full application. The methods contained in this module is opinionated
-about your coordinate system and not in any way optimized.
-"""
-
-__docformat__ = "reStructuredText"
-
-__all__ = [
-    "DrawOptions",
-    "get_mouse_pos",
-    "to_pygame",
-    "from_pygame",
-    # "lighten",
-    "positive_y_is_up",
-]
-
-from typing import Sequence, Tuple
-
-import numpy as np
-import pygame
-import pymunk
-from pymunk.space_debug_draw_options import SpaceDebugColor
-from pymunk.vec2d import Vec2d
-
-positive_y_is_up: bool = False
-"""Make increasing values of y point upwards.
-
-When True::
-
-    y
-    ^
-    |      . (3, 3)
-    |
-    |   . (2, 2)
-    |
-    +------ > x
-
-When False::
-
-    +------ > x
-    |
-    |   . (2, 2)
-    |
-    |      . (3, 3)
-    v
-    y
-
-"""
-
-
-class DrawOptions(pymunk.SpaceDebugDrawOptions):
-    def __init__(self, surface: pygame.Surface) -> None:
-        """Draw a pymunk.Space on a pygame.Surface object.
-
-        Typical usage::
-
-        >>> import pymunk
-        >>> surface = pygame.Surface((10,10))
-        >>> space = pymunk.Space()
-        >>> options = pymunk.pygame_util.DrawOptions(surface)
-        >>> space.debug_draw(options)
-
-        You can control the color of a shape by setting shape.color to the color
-        you want it drawn in::
-
-        >>> c = pymunk.Circle(None, 10)
-        >>> c.color = pygame.Color("pink")
-
-        See pygame_util.demo.py for a full example
-
-        Since pygame uses a coordinate system where y points down (in contrast
-        to many other cases), you either have to make the physics simulation
-        with Pymunk also behave in that way, or flip everything when you draw.
-
-        The easiest is probably to just make the simulation behave the same
-        way as Pygame does. In that way all coordinates used are in the same
-        orientation and easy to reason about::
-
-        >>> space = pymunk.Space()
-        >>> space.gravity = (0, -1000)
-        >>> body = pymunk.Body()
-        >>> body.position = (0, 0) # will be positioned in the top left corner
-        >>> space.debug_draw(options)
-
-        To flip the drawing its possible to set the module property
-        :py:data:`positive_y_is_up` to True. Then the pygame drawing will flip
-        the simulation upside down before drawing::
-
-        >>> positive_y_is_up = True
-        >>> body = pymunk.Body()
-        >>> body.position = (0, 0)
-        >>> # Body will be position in bottom left corner
-
-        :Parameters:
-                surface : pygame.Surface
-                    Surface that the objects will be drawn on
-        """
-        self.surface = surface
-        super().__init__()
-
-    def draw_circle(
-        self,
-        pos: Vec2d,
-        angle: float,
-        radius: float,
-        outline_color: SpaceDebugColor,
-        fill_color: SpaceDebugColor,
-    ) -> None:
-        p = to_pygame(pos, self.surface)
-
-        pygame.draw.circle(self.surface, fill_color.as_int(), p, round(radius), 0)
-        pygame.draw.circle(self.surface, light_color(fill_color).as_int(), p, round(radius - 4), 0)
-
-        # circle_edge = pos + Vec2d(radius, 0).rotated(angle)
-        # p2 = to_pygame(circle_edge, self.surface)
-        # line_r = 2 if radius > 20 else 1
-        # pygame.draw.lines(self.surface, outline_color.as_int(), False, [p, p2], line_r)
-
-    def draw_segment(self, a: Vec2d, b: Vec2d, color: SpaceDebugColor) -> None:
-        p1 = to_pygame(a, self.surface)
-        p2 = to_pygame(b, self.surface)
-
-        pygame.draw.aalines(self.surface, color.as_int(), False, [p1, p2])
-
-    def draw_fat_segment(
-        self,
-        a: Tuple[float, float],
-        b: Tuple[float, float],
-        radius: float,
-        outline_color: SpaceDebugColor,
-        fill_color: SpaceDebugColor,
-    ) -> None:
-        p1 = to_pygame(a, self.surface)
-        p2 = to_pygame(b, self.surface)
-
-        r = round(max(1, radius * 2))
-        pygame.draw.lines(self.surface, fill_color.as_int(), False, [p1, p2], r)
-        if r > 2:
-            orthog = [abs(p2[1] - p1[1]), abs(p2[0] - p1[0])]
-            if orthog[0] == 0 and orthog[1] == 0:
-                return
-            scale = radius / (orthog[0] * orthog[0] + orthog[1] * orthog[1]) ** 0.5
-            orthog[0] = round(orthog[0] * scale)
-            orthog[1] = round(orthog[1] * scale)
-            points = [
-                (p1[0] - orthog[0], p1[1] - orthog[1]),
-                (p1[0] + orthog[0], p1[1] + orthog[1]),
-                (p2[0] + orthog[0], p2[1] + orthog[1]),
-                (p2[0] - orthog[0], p2[1] - orthog[1]),
-            ]
-            pygame.draw.polygon(self.surface, fill_color.as_int(), points)
-            pygame.draw.circle(
-                self.surface,
-                fill_color.as_int(),
-                (round(p1[0]), round(p1[1])),
-                round(radius),
-            )
-            pygame.draw.circle(
-                self.surface,
-                fill_color.as_int(),
-                (round(p2[0]), round(p2[1])),
-                round(radius),
-            )
-
-    def draw_polygon(
-        self,
-        verts: Sequence[Tuple[float, float]],
-        radius: float,
-        outline_color: SpaceDebugColor,
-        fill_color: SpaceDebugColor,
-    ) -> None:
-        ps = [to_pygame(v, self.surface) for v in verts]
-        ps += [ps[0]]
-
-        radius = 2
-        pygame.draw.polygon(self.surface, light_color(fill_color).as_int(), ps)
-
-        if radius > 0:
-            for i in range(len(verts)):
-                a = verts[i]
-                b = verts[(i + 1) % len(verts)]
-                self.draw_fat_segment(a, b, radius, fill_color, fill_color)
-
-    def draw_dot(self, size: float, pos: Tuple[float, float], color: SpaceDebugColor) -> None:
-        p = to_pygame(pos, self.surface)
-        pygame.draw.circle(self.surface, color.as_int(), p, round(size), 0)
-
-
-def get_mouse_pos(surface: pygame.Surface) -> Tuple[int, int]:
-    """Get position of the mouse pointer in pymunk coordinates."""
-    p = pygame.mouse.get_pos()
-    return from_pygame(p, surface)
-
-
-def to_pygame(p: Tuple[float, float], surface: pygame.Surface) -> Tuple[int, int]:
-    """Convenience method to convert pymunk coordinates to pygame surface
-    local coordinates.
-
-    Note that in case positive_y_is_up is False, this function won't actually do
-    anything except converting the point to integers.
-    """
-    if positive_y_is_up:
-        return round(p[0]), surface.get_height() - round(p[1])
-    else:
-        return round(p[0]), round(p[1])
-
-
-def from_pygame(p: Tuple[float, float], surface: pygame.Surface) -> Tuple[int, int]:
-    """Convenience method to convert pygame surface local coordinates to
-    pymunk coordinates
-    """
-    return to_pygame(p, surface)
-
-
-def light_color(color: SpaceDebugColor):
-    color = np.minimum(1.2 * np.float32([color.r, color.g, color.b, color.a]), np.float32([255]))
-    color = SpaceDebugColor(r=color[0], g=color[1], b=color[2], a=color[3])
-    return color

lerobot/common/envs/simxarm/env.py

@@ -1,237 +0,0 @@
-import importlib
-import logging
-from collections import deque
-from typing import Optional
-
-import einops
-import numpy as np
-import torch
-from tensordict import TensorDict
-from torchrl.data.tensor_specs import (
-    BoundedTensorSpec,
-    CompositeSpec,
-    DiscreteTensorSpec,
-    UnboundedContinuousTensorSpec,
-)
-from torchrl.envs.libs.gym import _gym_to_torchrl_spec_transform
-
-from lerobot.common.envs.abstract import AbstractEnv
-from lerobot.common.utils import set_global_seed
-
-MAX_NUM_ACTIONS = 4
-
-_has_gym = importlib.util.find_spec("gymnasium") is not None
-
-
-class SimxarmEnv(AbstractEnv):
-    name = "simxarm"
-    available_tasks = ["lift"]
-
-    def __init__(
-        self,
-        task,
-        frame_skip: int = 1,
-        from_pixels: bool = False,
-        pixels_only: bool = False,
-        image_size=None,
-        seed=1337,
-        device="cpu",
-        num_prev_obs=0,
-        num_prev_action=0,
-    ):
-        super().__init__(
-            task=task,
-            frame_skip=frame_skip,
-            from_pixels=from_pixels,
-            pixels_only=pixels_only,
-            image_size=image_size,
-            seed=seed,
-            device=device,
-            num_prev_obs=num_prev_obs,
-            num_prev_action=num_prev_action,
-        )
-
-    def _make_env(self):
-        if not _has_gym:
-            raise ImportError("Cannot import gymnasium.")
-
-        import gymnasium
-
-        from lerobot.common.envs.simxarm.simxarm import TASKS
-
-        if self.task not in TASKS:
-            raise ValueError(f"Unknown task {self.task}. Must be one of {list(TASKS.keys())}")
-
-        self._env = TASKS[self.task]["env"]()
-
-        num_actions = len(TASKS[self.task]["action_space"])
-        self._action_space = gymnasium.spaces.Box(low=-1.0, high=1.0, shape=(num_actions,))
-        self._action_padding = np.zeros((MAX_NUM_ACTIONS - num_actions), dtype=np.float32)
-        if "w" not in TASKS[self.task]["action_space"]:
-            self._action_padding[-1] = 1.0
-
-    def render(self, mode="rgb_array", width=384, height=384):
-        return self._env.render(mode, width=width, height=height)
-
-    def _format_raw_obs(self, raw_obs):
-        if self.from_pixels:
-            image = self.render(mode="rgb_array", width=self.image_size, height=self.image_size)
-            image = image.transpose(2, 0, 1)  # (H, W, C) -> (C, H, W)
-            image = torch.tensor(image.copy(), dtype=torch.uint8)
-
-            obs = {"image": image}
-
-            if not self.pixels_only:
-                obs["state"] = torch.tensor(self._env.robot_state, dtype=torch.float32)
-        else:
-            obs = {"state": torch.tensor(raw_obs["observation"], dtype=torch.float32)}
-
-        # obs = TensorDict(obs, batch_size=[])
-        return obs
-
-    def _reset(self, tensordict: Optional[TensorDict] = None):
-        td = tensordict
-        if td is None or td.is_empty():
-            raw_obs = self._env.reset()
-
-            obs = self._format_raw_obs(raw_obs)
-
-            if self.num_prev_obs > 0:
-                stacked_obs = {}
-                if "image" in obs:
-                    self._prev_obs_image_queue = deque(
-                        [obs["image"]] * (self.num_prev_obs + 1), maxlen=(self.num_prev_obs + 1)
-                    )
-                    stacked_obs["image"] = torch.stack(list(self._prev_obs_image_queue))
-                if "state" in obs:
-                    self._prev_obs_state_queue = deque(
-                        [obs["state"]] * (self.num_prev_obs + 1), maxlen=(self.num_prev_obs + 1)
-                    )
-                    stacked_obs["state"] = torch.stack(list(self._prev_obs_state_queue))
-                obs = stacked_obs
-
-            td = TensorDict(
-                {
-                    "observation": TensorDict(obs, batch_size=[]),
-                    "done": torch.tensor([False], dtype=torch.bool),
-                },
-                batch_size=[],
-            )
-        else:
-            raise NotImplementedError()
-
-        return td
-
-    def _step(self, tensordict: TensorDict):
-        td = tensordict
-        action = td["action"].numpy()
-        # step expects shape=(4,) so we pad if necessary
-        action = np.concatenate([action, self._action_padding])
-        # TODO(rcadene): add info["is_success"] and info["success"] ?
-        sum_reward = 0
-
-        if action.ndim == 1:
-            action = einops.repeat(action, "c -> t c", t=self.frame_skip)
-        else:
-            if self.frame_skip > 1:
-                raise NotImplementedError()
-
-        num_action_steps = action.shape[0]
-        for i in range(num_action_steps):
-            raw_obs, reward, done, info = self._env.step(action[i])
-            sum_reward += reward
-
-            obs = self._format_raw_obs(raw_obs)
-
-            if self.num_prev_obs > 0:
-                stacked_obs = {}
-                if "image" in obs:
-                    self._prev_obs_image_queue.append(obs["image"])
-                    stacked_obs["image"] = torch.stack(list(self._prev_obs_image_queue))
-                if "state" in obs:
-                    self._prev_obs_state_queue.append(obs["state"])
-                    stacked_obs["state"] = torch.stack(list(self._prev_obs_state_queue))
-                obs = stacked_obs
-
-        td = TensorDict(
-            {
-                "observation": self._format_raw_obs(raw_obs),
-                "reward": torch.tensor([sum_reward], dtype=torch.float32),
-                "done": torch.tensor([done], dtype=torch.bool),
-                "success": torch.tensor([info["success"]], dtype=torch.bool),
-            },
-            batch_size=[],
-        )
-        return td
-
-    def _make_spec(self):
-        obs = {}
-        if self.from_pixels:
-            image_shape = (3, self.image_size, self.image_size)
-            if self.num_prev_obs > 0:
-                image_shape = (self.num_prev_obs + 1, *image_shape)
-
-            obs["image"] = BoundedTensorSpec(
-                low=0,
-                high=255,
-                shape=image_shape,
-                dtype=torch.uint8,
-                device=self.device,
-            )
-            if not self.pixels_only:
-                state_shape = (len(self._env.robot_state),)
-                if self.num_prev_obs > 0:
-                    state_shape = (self.num_prev_obs + 1, *state_shape)
-
-                obs["state"] = UnboundedContinuousTensorSpec(
-                    shape=state_shape,
-                    dtype=torch.float32,
-                    device=self.device,
-                )
-        else:
-            # TODO(rcadene): add observation_space achieved_goal and desired_goal?
-            state_shape = self._env.observation_space["observation"].shape
-            if self.num_prev_obs > 0:
-                state_shape = (self.num_prev_obs + 1, *state_shape)
-
-            obs["state"] = UnboundedContinuousTensorSpec(
-                # TODO:
-                shape=state_shape,
-                dtype=torch.float32,
-                device=self.device,
-            )
-        self.observation_spec = CompositeSpec({"observation": obs})
-
-        self.action_spec = _gym_to_torchrl_spec_transform(
-            self._action_space,
-            device=self.device,
-        )
-
-        self.reward_spec = UnboundedContinuousTensorSpec(
-            shape=(1,),
-            dtype=torch.float32,
-            device=self.device,
-        )
-
-        self.done_spec = CompositeSpec(
-            {
-                "done": DiscreteTensorSpec(
-                    2,
-                    shape=(1,),
-                    dtype=torch.bool,
-                    device=self.device,
-                ),
-                "success": DiscreteTensorSpec(
-                    2,
-                    shape=(1,),
-                    dtype=torch.bool,
-                    device=self.device,
-                ),
-            }
-        )
-
-    def _set_seed(self, seed: Optional[int]):
-        set_global_seed(seed)
-        self._seed = seed
-        # TODO(aliberts): change self._reset so that it takes in a seed value
-        logging.warning("simxarm env is not properly seeded")

lerobot/common/envs/simxarm/simxarm/__init__.py

@@ -1,166 +0,0 @@
-from collections import OrderedDict, deque
-
-import gymnasium as gym
-import numpy as np
-from gymnasium.wrappers import TimeLimit
-
-from lerobot.common.envs.simxarm.simxarm.tasks.base import Base as Base
-from lerobot.common.envs.simxarm.simxarm.tasks.lift import Lift
-from lerobot.common.envs.simxarm.simxarm.tasks.peg_in_box import PegInBox
-from lerobot.common.envs.simxarm.simxarm.tasks.push import Push
-from lerobot.common.envs.simxarm.simxarm.tasks.reach import Reach
-
-TASKS = OrderedDict(
-    (
-        (
-            "reach",
-            {
-                "env": Reach,
-                "action_space": "xyz",
-                "episode_length": 50,
-                "description": "Reach a target location with the end effector",
-            },
-        ),
-        (
-            "push",
-            {
-                "env": Push,
-                "action_space": "xyz",
-                "episode_length": 50,
-                "description": "Push a cube to a target location",
-            },
-        ),
-        (
-            "peg_in_box",
-            {
-                "env": PegInBox,
-                "action_space": "xyz",
-                "episode_length": 50,
-                "description": "Insert a peg into a box",
-            },
-        ),
-        (
-            "lift",
-            {
-                "env": Lift,
-                "action_space": "xyzw",
-                "episode_length": 50,
-                "description": "Lift a cube above a height threshold",
-            },
-        ),
-    )
-)
-
-
-class SimXarmWrapper(gym.Wrapper):
-    """
-    A wrapper for the SimXarm environments. This wrapper is used to
-    convert the action and observation spaces to the correct format.
-    """
-
-    def __init__(self, env, task, obs_mode, image_size, action_repeat, frame_stack=1, channel_last=False):
-        super().__init__(env)
-        self._env = env
-        self.obs_mode = obs_mode
-        self.image_size = image_size
-        self.action_repeat = action_repeat
-        self.frame_stack = frame_stack
-        self._frames = deque([], maxlen=frame_stack)
-        self.channel_last = channel_last
-        self._max_episode_steps = task["episode_length"] // action_repeat
-
-        image_shape = (
-            (image_size, image_size, 3 * frame_stack)
-            if channel_last
-            else (3 * frame_stack, image_size, image_size)
-        )
-        if obs_mode == "state":
-            self.observation_space = env.observation_space["observation"]
-        elif obs_mode == "rgb":
-            self.observation_space = gym.spaces.Box(low=0, high=255, shape=image_shape, dtype=np.uint8)
-        elif obs_mode == "all":
-            self.observation_space = gym.spaces.Dict(
-                state=gym.spaces.Box(low=-np.inf, high=np.inf, shape=(4,), dtype=np.float32),
-                rgb=gym.spaces.Box(low=0, high=255, shape=image_shape, dtype=np.uint8),
-            )
-        else:
-            raise ValueError(f"Unknown obs_mode {obs_mode}. Must be one of [rgb, all, state]")
-        self.action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(len(task["action_space"]),))
-        self.action_padding = np.zeros(4 - len(task["action_space"]), dtype=np.float32)
-        if "w" not in task["action_space"]:
-            self.action_padding[-1] = 1.0
-
-    def _render_obs(self):
-        obs = self.render(mode="rgb_array", width=self.image_size, height=self.image_size)
-        if not self.channel_last:
-            obs = obs.transpose(2, 0, 1)
-        return obs.copy()
-
-    def _update_frames(self, reset=False):
-        pixels = self._render_obs()
-        self._frames.append(pixels)
-        if reset:
-            for _ in range(1, self.frame_stack):
-                self._frames.append(pixels)
-        assert len(self._frames) == self.frame_stack
-
-    def transform_obs(self, obs, reset=False):
-        if self.obs_mode == "state":
-            return obs["observation"]
-        elif self.obs_mode == "rgb":
-            self._update_frames(reset=reset)
-            rgb_obs = np.concatenate(list(self._frames), axis=-1 if self.channel_last else 0)
-            return rgb_obs
-        elif self.obs_mode == "all":
-            self._update_frames(reset=reset)
-            rgb_obs = np.concatenate(list(self._frames), axis=-1 if self.channel_last else 0)
-            return OrderedDict((("rgb", rgb_obs), ("state", self.robot_state)))
-        else:
-            raise ValueError(f"Unknown obs_mode {self.obs_mode}. Must be one of [rgb, all, state]")
-
-    def reset(self):
-        return self.transform_obs(self._env.reset(), reset=True)
-
-    def step(self, action):
-        action = np.concatenate([action, self.action_padding])
-        reward = 0.0
-        for _ in range(self.action_repeat):
-            obs, r, done, info = self._env.step(action)
-            reward += r
-        return self.transform_obs(obs), reward, done, info
-
-    def render(self, mode="rgb_array", width=384, height=384, **kwargs):
-        return self._env.render(mode, width=width, height=height)
-
-    @property
-    def state(self):
-        return self._env.robot_state
-
-
-def make(task, obs_mode="state", image_size=84, action_repeat=1, frame_stack=1, channel_last=False, seed=0):
-    """
-    Create a new environment.
-    Args:
-            task (str): The task to create an environment for. Must be one of:
-                    - 'reach'
-                    - 'push'
-                    - 'peg-in-box'
-                    - 'lift'
-            obs_mode (str): The observation mode to use. Must be one of:
-                    - 'state': Only state observations
-                    - 'rgb': RGB images
-                    - 'all': RGB images and state observations
-            image_size (int): The size of the image observations
-            action_repeat (int): The number of times to repeat the action
-            seed (int): The random seed to use
-    Returns:
-            gym.Env: The environment
-    """
-    if task not in TASKS:
-        raise ValueError(f"Unknown task {task}. Must be one of {list(TASKS.keys())}")
-    env = TASKS[task]["env"]()
-    env = TimeLimit(env, TASKS[task]["episode_length"])
-    env = SimXarmWrapper(env, TASKS[task], obs_mode, image_size, action_repeat, frame_stack, channel_last)
-    env.seed(seed)
-
-    return env

lerobot/common/envs/simxarm/simxarm/tasks/assets/lift.xml

@@ -1,53 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-
-<mujoco>
-	<compiler angle="radian" coordinate="local" meshdir="mesh" texturedir="texture"></compiler>
-	<size nconmax="2000" njmax="500"/>
-
-	<option timestep="0.002">
-		<flag warmstart="enable"></flag>
-	</option>
-
-	<include file="shared.xml"></include>
-
-	<worldbody>
-		<body name="floor0" pos="0 0 0">
-			<geom name="floorgeom0" pos="1.2 -2.0 0" size="20.0 20.0 1" type="plane" condim="3" material="floor_mat"></geom>
-		</body>
-
-		<include file="xarm.xml"></include>
-
-		<body pos="0.75 0 0.6325" name="pedestal0">
-			<geom name="pedestalgeom0" size="0.1 0.1 0.01" pos="0.32 0.27 0" type="box" mass="2000" material="pedestal_mat"></geom>
-			<site pos="0.30 0.30 0" size="0.075 0.075 0.002" type="box" name="robotmountsite0" rgba="0.55 0.54 0.53 1" />
-		</body>
-
-		<body pos="1.5 0.075 0.3425" name="table0">
-			<geom name="tablegeom0" size="0.3 0.6 0.2" pos="0 0 0" type="box" material="table_mat" density="2000" friction="1 1 1"></geom>
-		</body>
-
-		<body name="object" pos="1.405 0.3 0.58625">
-			<joint name="object_joint0" type="free" limited="false"></joint>
-			<geom size="0.035 0.035 0.035" type="box" name="object0" material="block_mat" density="50000" condim="4" friction="1 1 1" solimp="1 1 1" solref="0.02 1"></geom>
-			<site name="object_site" pos="0 0 0" size="0.035 0.035 0.035" rgba="1 0 0 0" type="box"></site>
-		</body>
-
-		<light directional="true" ambient="0.1 0.1 0.1" diffuse="0 0 0" specular="0 0 0" castshadow="false" pos="1.65 0 10" dir="-0.57 -0.57 -0.57" name="light0"></light>
-        <light directional="true" ambient="0.1 0.1 0.1" diffuse="0 0 0" specular="0 0 0" castshadow="false" pos="0 -4 4" dir="0 1 -0.1" name="light1"></light>
-        <light directional="true" ambient="0.05 0.05 0.05" diffuse="0 0 0" specular="0 0 0" castshadow="false" pos="2.13 1.6 2.5" name="light2"></light>
-        <light pos="0 0 2" dir="0.2 0.2 -0.8" directional="true"  diffuse="0.3 0.3 0.3"  castshadow="false" name="light3"></light>
-
-		<camera fovy="50" name="camera0" pos="0.9559 1.0 1.1" euler="-1.1 -0.6 3.4" />
-	</worldbody>
-
-	<equality>
-    	<connect body2="left_finger" body1="left_inner_knuckle" anchor="0.0 0.035 0.042" solimp="0.9 0.95 0.001 0.5 2" solref="0.0002 1.0" ></connect>
-    	<connect body2="right_finger" body1="right_inner_knuckle" anchor="0.0 -0.035 0.042" solimp="0.9 0.95 0.001 0.5 2" solref="0.0002 1.0" ></connect>
-        <joint joint1="left_inner_knuckle_joint" joint2="right_inner_knuckle_joint"></joint>
-    </equality>
-
-    <actuator>
-		<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="left_inner_knuckle_joint" gear="200.0"/>
-        <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="right_inner_knuckle_joint" gear="200.0"/>
-    </actuator>
-</mujoco>

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/base_link.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_inner.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_inner2.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_outer.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_finger.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_inner_knuckle.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_outer_knuckle.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link1.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link2.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link3.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link4.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link5.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link6.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link7.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link_base.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_finger.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_inner_knuckle.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_outer_knuckle.stl

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

lerobot/common/envs/simxarm/simxarm/tasks/assets/peg_in_box.xml

@@ -1,74 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-
-<mujoco>
-	<compiler angle="radian" coordinate="local" meshdir="mesh" texturedir="texture"></compiler>
-	<size nconmax="2000" njmax="500"/>
-
-	<option timestep="0.001">
-		<flag warmstart="enable"></flag>
-	</option>
-
-	<include file="shared.xml"></include>
-
-	<worldbody>
-		<body name="floor0" pos="0 0 0">
-			<geom name="floorgeom0" pos="1.2 -2.0 0" size="1.0 10.0 1" type="plane" condim="3" material="floor_mat"></geom>
-		</body>
-
-		<include file="xarm.xml"></include>
-
-		<body pos="0.75 0 0.6325" name="pedestal0">
-			<geom name="pedestalgeom0" size="0.1 0.1 0.01" pos="0.32 0.27 0" type="box" mass="2000" material="pedestal_mat"></geom>
-			<site pos="0.30 0.30 0" size="0.075 0.075 0.002" type="box" name="robotmountsite0" rgba="0.55 0.54 0.53 1" />
-		</body>
-
-		<body pos="1.5 0.075 0.3425" name="table0">
-			<geom name="tablegeom0" size="0.3 0.6 0.2" pos="0 0 0" type="box" material="table_mat" density="2000" friction="1 0.005 0.0002"></geom>
-		</body>
-
-		<body name="box0" pos="1.605 0.25 0.55">
-			<joint name="box_joint0" type="free" limited="false"></joint>
-			<site name="box_site" pos="0 0.075 -0.01" size="0.02" rgba="0 0 0 0" type="sphere"></site>
-			<geom name="box_side0" pos="0 0 0" size="0.065 0.002 0.04" type= "box" rgba="0.8 0.1 0.1 1" mass ="1" condim="4" />
-			<geom name="box_side1" pos="0 0.149 0" size="0.065 0.002 0.04" type="box" rgba="0.9 0.2 0.2 1" mass ="2" condim="4" />
-			<geom name="box_side2" pos="0.064 0.074 0" size="0.002 0.075 0.04" type="box" rgba="0.8 0.1 0.1 1" mass ="2" condim="4" />
-			<geom name="box_side3" pos="-0.064 0.074 0" size="0.002 0.075 0.04" type="box" rgba="0.9 0.2 0.2 1" mass ="2" condim="4" />
-			<geom name="box_side4" pos="-0 0.074 -0.038" size="0.065 0.075 0.002" type="box" rgba="0.5 0 0 1" mass ="2" condim="4"/>
-		</body>
-
-		<body name="object0" pos="1.4 0.25 0.65">
-			<joint name="object_joint0" type="free" limited="false"></joint>
-			<geom name="object_target0" type="cylinder" pos="0 0 -0.05" size="0.03 0.035" rgba="0.6 0.8 0.5 1" mass ="0.1" condim="3"  />
-			<site name="object_site" pos="0 0 -0.05" size="0.0325 0.0375" rgba="0 0 0 0" type="cylinder"></site>
-			<body name="B0" pos="0 0 0" euler="0 0 0 ">
-				<joint name="B0:joint" type="slide" limited="true" axis="0 0 1" damping="0.05" range="0.0001 0.0001001" solimpfriction="0.98 0.98 0.95" frictionloss="1"></joint>
-				<geom type="capsule" size="0.002 0.03" rgba="0 0 0 1" mass="0.001" condim="4"/>
-				<body name="B1" pos="0 0 0.04" euler="0 3.14 0 ">
-					<joint name="B1:joint1" type="hinge" axis="1 0 0" range="-0.1 0.1"  frictionloss="1"></joint>
-					<joint name="B1:joint2" type="hinge" axis="0 1 0" range="-0.1 0.1"  frictionloss="1"></joint>
-					<joint name="B1:joint3" type="hinge" axis="0 0 1" range="-0.1 0.1"  frictionloss="1"></joint>
-					<geom type="capsule" size="0.002 0.004" rgba="1 0 0 0" mass="0.001" condim="4"/>
-				</body>
-			</body>
-		</body>
-
-		<light directional="true" ambient="0.1 0.1 0.1" diffuse="0 0 0" specular="0 0 0" castshadow="false" pos="1.65 0 10" dir="-0.57 -0.57 -0.57" name="light0"></light>
-		<light directional="true" ambient="0.1 0.1 0.1" diffuse="0 0 0" specular="0 0 0" castshadow="false" pos="0 -4 4" dir="0 1 -0.1" name="light1"></light>
-		<light directional="true" ambient="0.05 0.05 0.05" diffuse="0 0 0" specular="0 0 0" castshadow="false" pos="2.13 1.6 2.5" name="light2"></light>
-		<light pos="0 0 2" dir="0.2 0.2 -0.8" directional="true"  diffuse="0.3 0.3 0.3"  castshadow="false" name="light3"></light>
-
-		<camera fovy="50" name="camera0" pos="0.9559 1.0 1.1" euler="-1.1 -0.6 3.4" />
-	</worldbody>
-
-	<equality>
-		<connect body2="left_finger" body1="left_inner_knuckle" anchor="0.0 0.035 0.042" solimp="0.9 0.95 0.001 0.5 2" solref="0.0002 1.0" ></connect>
-		<connect body2="right_finger" body1="right_inner_knuckle" anchor="0.0 -0.035 0.042" solimp="0.9 0.95 0.001 0.5 2" solref="0.0002 1.0" ></connect>
-		<weld body1="right_hand" body2="B1" solimp="0.99 0.99 0.99" solref="0.02 1"></weld>
-		<joint joint1="left_inner_knuckle_joint" joint2="right_inner_knuckle_joint"></joint>
-	</equality>
-
-	<actuator>
-		<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="left_inner_knuckle_joint" gear="200.0"/>
-		<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="right_inner_knuckle_joint" gear="200.0"/>
-	</actuator>
-</mujoco>

lerobot/common/envs/simxarm/simxarm/tasks/assets/push.xml

@@ -1,54 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-
-<mujoco>
-	<compiler angle="radian" coordinate="local" meshdir="mesh" texturedir="texture"></compiler>
-	<size nconmax="2000" njmax="500"/>
-
-	<option timestep="0.002">
-		<flag warmstart="enable"></flag>
-	</option>
-
-	<include file="shared.xml"></include>
-
-	<worldbody>
-		<body name="floor0" pos="0 0 0">
-			<geom name="floorgeom0" pos="1.2 -2.0 0" size="1.0 10.0 1" type="plane" condim="3" material="floor_mat"></geom>
-			<site name="target0" pos="1.565 0.3 0.545" size="0.0475 0.001" rgba="1 0 0 1" type="cylinder"></site>
-		</body>
-
-		<include file="xarm.xml"></include>
-
-		<body pos="0.75 0 0.6325" name="pedestal0">
-			<geom name="pedestalgeom0" size="0.1 0.1 0.01" pos="0.32 0.27 0" type="box" mass="2000" material="pedestal_mat"></geom>
-			<site pos="0.30 0.30 0" size="0.075 0.075 0.002" type="box" name="robotmountsite0" rgba="0.55 0.54 0.53 1" />
-		</body>
-
-		<body pos="1.5 0.075 0.3425" name="table0">
-			<geom name="tablegeom0" size="0.3 0.6 0.2" pos="0 0 0" type="box" material="table_mat" density="2000" friction="1 0.005 0.0002"></geom>
-		</body>
-
-		<body name="object" pos="1.655 0.3 0.68">
-			<joint name="object_joint0" type="free" limited="false"></joint>
-			<geom size="0.024 0.024 0.024" type="box" name="object" material="block_mat" density="50000" condim="4" friction="1 1 1" solimp="1 1 1" solref="0.02 1"></geom>
-			<site name="object_site" pos="0 0 0" size="0.024 0.024 0.024" rgba="0 0 0 0" type="box"></site>
-		</body>
-
-		<light directional="true" ambient="0.1 0.1 0.1" diffuse="0 0 0" specular="0 0 0" castshadow="false" pos="1.65 0 10" dir="-0.57 -0.57 -0.57" name="light0"></light>
-        <light directional="true" ambient="0.1 0.1 0.1" diffuse="0 0 0" specular="0 0 0" castshadow="false" pos="0 -4 4" dir="0 1 -0.1" name="light1"></light>
-        <light directional="true" ambient="0.05 0.05 0.05" diffuse="0 0 0" specular="0 0 0" castshadow="false" pos="2.13 1.6 2.5" name="light2"></light>
-        <light pos="0 0 2" dir="0.2 0.2 -0.8" directional="true"  diffuse="0.3 0.3 0.3"  castshadow="false" name="light3"></light>
-
-		<camera fovy="50" name="camera0" pos="0.9559 1.0 1.1" euler="-1.1 -0.6 3.4" />
-	</worldbody>
-
-	<equality>
-    	<connect body2="left_finger" body1="left_inner_knuckle" anchor="0.0 0.035 0.042" solimp="0.9 0.95 0.001 0.5 2" solref="0.0002 1.0" ></connect>
-    	<connect body2="right_finger" body1="right_inner_knuckle" anchor="0.0 -0.035 0.042" solimp="0.9 0.95 0.001 0.5 2" solref="0.0002 1.0" ></connect>
-        <joint joint1="left_inner_knuckle_joint" joint2="right_inner_knuckle_joint"></joint>
-    </equality>
-
-    <actuator>
-    	<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="left_inner_knuckle_joint" gear="200.0"/>
-        <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="right_inner_knuckle_joint" gear="200.0"/>
-    </actuator>
-</mujoco>

lerobot/common/envs/simxarm/simxarm/tasks/assets/reach.xml

@@ -1,48 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-
-<mujoco>
-	<compiler angle="radian" coordinate="local" meshdir="mesh" texturedir="texture"></compiler>
-	<size nconmax="2000" njmax="500"/>
-
-	<option timestep="0.002">
-		<flag warmstart="enable"></flag>
-	</option>
-
-	<include file="shared.xml"></include>
-
-	<worldbody>
-		<body name="floor0" pos="0 0 0">
-			<geom name="floorgeom0" pos="1.2 -2.0 0" size="1.0 10.0 1" type="plane" condim="3" material="floor_mat"></geom>
-			<site name="target0" pos="1.605 0.3 0.58" size="0.0475 0.001" rgba="1 0 0 1" type="cylinder"></site>
-		</body>
-
-		<include file="xarm.xml"></include>
-
-		<body pos="0.75 0 0.6325" name="pedestal0">
-			<geom name="pedestalgeom0" size="0.1 0.1 0.01" pos="0.32 0.27 0" type="box" mass="2000" material="pedestal_mat"></geom>
-			<site pos="0.30 0.30 0" size="0.075 0.075 0.002" type="box" name="robotmountsite0" rgba="0.55 0.54 0.53 1" />
-		</body>
-
-		<body pos="1.5 0.075 0.3425" name="table0">
-			<geom name="tablegeom0" size="0.3 0.6 0.2" pos="0 0 0" type="box" material="table_mat" density="2000" friction="1 0.005 0.0002"></geom>
-		</body>
-
-		<light directional="true" ambient="0.1 0.1 0.1" diffuse="0 0 0" specular="0 0 0" castshadow="false" pos="1.65 0 10" dir="-0.57 -0.57 -0.57" name="light0"></light>
-        <light directional="true" ambient="0.1 0.1 0.1" diffuse="0 0 0" specular="0 0 0" castshadow="false" pos="0 -4 4" dir="0 1 -0.1" name="light1"></light>
-        <light directional="true" ambient="0.05 0.05 0.05" diffuse="0 0 0" specular="0 0 0" castshadow="false" pos="2.13 1.6 2.5" name="light2"></light>
-        <light pos="0 0 2" dir="0.2 0.2 -0.8" directional="true"  diffuse="0.3 0.3 0.3"  castshadow="false" name="light3"></light>
-
-		<camera fovy="50" name="camera0" pos="0.9559 1.0 1.1" euler="-1.1 -0.6 3.4" />
-	</worldbody>
-
-	<equality>
-    	<connect body2="left_finger" body1="left_inner_knuckle" anchor="0.0 0.035 0.042" solimp="0.9 0.95 0.001 0.5 2" solref="0.0002 1.0" ></connect>
-    	<connect body2="right_finger" body1="right_inner_knuckle" anchor="0.0 -0.035 0.042" solimp="0.9 0.95 0.001 0.5 2" solref="0.0002 1.0" ></connect>
-        <joint joint1="left_inner_knuckle_joint" joint2="right_inner_knuckle_joint"></joint>
-    </equality>
-
-    <actuator>
-    	<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="left_inner_knuckle_joint" gear="200.0"/>
-        <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="right_inner_knuckle_joint" gear="200.0"/>
-    </actuator>
-</mujoco>

lerobot/common/envs/simxarm/simxarm/tasks/assets/shared.xml

@@ -1,51 +0,0 @@
-<mujoco>
-    <asset>
-        <texture type="skybox" builtin="gradient" rgb1="0.0 0.0 0.0" rgb2="0.0 0.0 0.0" width="32" height="32"></texture>
-        <material name="floor_mat" specular="0" shininess="0.0" reflectance="0" rgba="0.043 0.055 0.051 1"></material>
-
-        <material name="table_mat" specular="0.2" shininess="0.2" reflectance="0" rgba="1 1 1 1"></material>
-        <material name="pedestal_mat" specular="0.35" shininess="0.5" reflectance="0" rgba="0.705 0.585 0.405 1"></material>
-        <material name="block_mat" specular="0.5" shininess="0.9" reflectance="0.05" rgba="0.373 0.678 0.627 1"></material>
-
-        <material name="robot0:geomMat" shininess="0.03" specular="0.4"></material>
-        <material name="robot0:gripper_finger_mat" shininess="0.03" specular="0.4" reflectance="0"></material>
-        <material name="robot0:gripper_mat" shininess="0.03" specular="0.4" reflectance="0"></material>
-        <material name="background:gripper_mat" shininess="0.03" specular="0.4" reflectance="0"></material>
-        <material name="robot0:arm_mat" shininess="0.03" specular="0.4" reflectance="0"></material>
-        <material name="robot0:head_mat" shininess="0.03" specular="0.4" reflectance="0"></material>
-        <material name="robot0:torso_mat" shininess="0.03" specular="0.4" reflectance="0"></material>
-        <material name="robot0:base_mat" shininess="0.03" specular="0.4" reflectance="0"></material>
-
-        <mesh name="link_base" file="link_base.stl" />
-        <mesh name="link1" file="link1.stl" />
-        <mesh name="link2" file="link2.stl" />
-        <mesh name="link3" file="link3.stl" />
-        <mesh name="link4" file="link4.stl" />
-        <mesh name="link5" file="link5.stl" />
-        <mesh name="link6" file="link6.stl" />
-        <mesh name="link7" file="link7.stl" />
-        <mesh name="base_link" file="base_link.stl" />
-        <mesh name="left_outer_knuckle" file="left_outer_knuckle.stl" />
-        <mesh name="left_finger" file="left_finger.stl" />
-        <mesh name="left_inner_knuckle" file="left_inner_knuckle.stl" />
-        <mesh name="right_outer_knuckle" file="right_outer_knuckle.stl" />
-        <mesh name="right_finger" file="right_finger.stl" />
-        <mesh name="right_inner_knuckle" file="right_inner_knuckle.stl" />
-    </asset>
-
-    <equality>
-        <weld body1="robot0:mocap2" body2="link7" solimp="0.9 0.95 0.001" solref="0.02 1"></weld>
-    </equality>
-
-    <default>
-        <joint armature="1" damping="0.1" limited="true"/>
-         <default class="robot0:blue">
-            <geom rgba="0.086 0.506 0.767 1.0"></geom>
-        </default>
-
-        <default class="robot0:grey">
-            <geom rgba="0.356 0.361 0.376 1.0"></geom>
-        </default>
-    </default>
-
-</mujoco>

lerobot/common/envs/simxarm/simxarm/tasks/assets/xarm.xml

@@ -1,88 +0,0 @@
-<mujoco model="xarm7">
-    <body mocap="true" name="robot0:mocap2" pos="0 0 0">
-        <geom conaffinity="0" contype="0" pos="0 0 0" rgba="0 0.5 0 0" size="0.005 0.005 0.005" type="box"></geom>
-        <geom conaffinity="0" contype="0" pos="0 0 0" rgba="0.5 0 0 0" size="1 0.005 0.005" type="box"></geom>
-        <geom conaffinity="0" contype="0" pos="0 0 0" rgba="0 0 0.5 0" size="0.005 1 0.001" type="box"></geom>
-        <geom conaffinity="0" contype="0" pos="0 0 0" rgba="0.5 0.5 0 0" size="0.005 0.005 1" type="box"></geom>
-    </body>
-
-    <body name="link0" pos="1.09 0.28 0.655">
-        <geom name="bb" type="mesh" mesh="link_base" material="robot0:base_mat" rgba="1 1 1 1"/>
-        <body name="link1" pos="0 0 0.267">
-            <inertial pos="-0.0042142 0.02821 -0.0087788" quat="0.917781 -0.277115 0.0606681 0.277858" mass="0.42603" diaginertia="0.00144551 0.00137757 0.000823511" />
-            <joint name="joint1" pos="0 0 0" axis="0 0 1" limited="true" range="-6.28319 6.28319" damping="10" frictionloss="1" />
-            <geom name="j1" type="mesh" mesh="link1" material="robot0:arm_mat" rgba="1 1 1 1"/>
-            <body name="link2" pos="0 0 0" quat="0.707105 -0.707108 0 0">
-                <inertial pos="-3.3178e-05 -0.12849 0.026337" quat="0.447793 0.894132 -0.00224061 0.00218314" mass="0.56095" diaginertia="0.00319151 0.00311598 0.000980804" />
-                <joint name="joint2" pos="0 0 0" axis="0 0 1" limited="true" range="-2.059 2.0944" damping="10" frictionloss="1" />
-                <geom name="j2"  type="mesh" mesh="link2" material="robot0:head_mat" rgba="1 1 1 1"/>
-                <body name="link3" pos="0 -0.293 0" quat="0.707105 0.707108 0 0">
-                    <inertial pos="0.04223 -0.023258 -0.0096674" quat="0.883205 0.339803 0.323238 0.000542237" mass="0.44463" diaginertia="0.00133227 0.00119126 0.000780475" />
-                    <joint name="joint3" pos="0 0 0" axis="0 0 1" limited="true" range="-6.28319 6.28319" damping="5" frictionloss="1" />
-                    <geom name="j3" type="mesh" mesh="link3" material="robot0:gripper_mat" rgba="1 1 1 1"/>
-                    <body name="link4" pos="0.0525 0 0" quat="0.707105 0.707108 0 0">
-                        <inertial pos="0.067148 -0.10732 0.024479" quat="0.0654142 0.483317 -0.738663 0.465298" mass="0.52387" diaginertia="0.00288984 0.00282705 0.000894409" />
-                        <joint name="joint4" pos="0 0 0" axis="0 0 1" limited="true" range="-0.19198 3.927" damping="5" frictionloss="1" />
-                        <geom name="j4" type="mesh" mesh="link4" material="robot0:arm_mat" rgba="1 1 1 1"/>
-                        <body name="link5" pos="0.0775 -0.3425 0" quat="0.707105 0.707108 0 0">
-                            <inertial pos="-0.00023397 0.036705 -0.080064" quat="0.981064 -0.19003 0.00637998 0.0369004" mass="0.18554" diaginertia="0.00099553 0.000988613 0.000247126" />
-                            <joint name="joint5" pos="0 0 0" axis="0 0 1" limited="true" range="-6.28319 6.28319" damping="5" frictionloss="1" />
-                            <geom name="j5" type="mesh"  material="robot0:gripper_mat" rgba="1 1 1 1" mesh="link5" />
-                            <body name="link6" pos="0 0 0" quat="0.707105 0.707108 0 0">
-                                <inertial pos="0.058911 0.028469 0.0068428" quat="-0.188705 0.793535 0.166088 0.554173" mass="0.31344" diaginertia="0.000827892 0.000768871 0.000386708" />
-                                <joint name="joint6" pos="0 0 0" axis="0 0 1" limited="true" range="-1.69297 3.14159" damping="2" frictionloss="1" />
-                                <geom name="j6" type="mesh" material="robot0:gripper_mat" rgba="1 1 1 1" mesh="link6" />
-                                <body name="link7" pos="0.076 0.097 0" quat="0.707105 -0.707108 0 0">
-                                    <inertial pos="-0.000420033 -0.00287433 0.0257078" quat="0.999372 -0.0349129 -0.00605634 0.000551744" mass="0.85624" diaginertia="0.00137671 0.00118744 0.000514968" />
-                                    <joint name="joint7" pos="0 0 0" axis="0 0 1" limited="true" range="-6.28319 6.28319" damping="2" frictionloss="1" />
-                                    <geom name="j8" material="robot0:gripper_mat" type="mesh" rgba="0.753 0.753 0.753 1" mesh="link7" />
-                                    <geom name="j9" material="robot0:gripper_mat" type="mesh" rgba="1 1 1 1" mesh="base_link" />
-                                    <site name="grasp" pos="0 0 0.16" rgba="1 0 0 0" type="sphere" size="0.01" group="1"/>
-                                    <body name="left_outer_knuckle" pos="0 0.035 0.059098">
-                                        <inertial pos="0 0.021559 0.015181" quat="0.47789 0.87842 0 0" mass="0.033618" diaginertia="1.9111e-05 1.79089e-05 1.90167e-06" />
-                                        <joint name="drive_joint" pos="0 0 0" axis="1 0 0" limited="true" range="0 0.85" />
-                                        <geom type="mesh" rgba="0 0 0 1" conaffinity="1" contype="0"  mesh="left_outer_knuckle" />
-                                        <body name="left_finger" pos="0 0.035465 0.042039">
-                                            <inertial pos="0 -0.016413 0.029258" quat="0.697634 0.115353 -0.115353 0.697634" mass="0.048304" diaginertia="1.88037e-05 1.7493e-05 3.56792e-06" />
-                                            <joint name="left_finger_joint" pos="0 0 0" axis="-1 0 0" limited="true" range="0 0.85" />
-                                            <geom name="j10" material="robot0:gripper_finger_mat" type="mesh" rgba="0 0 0 1" conaffinity="3" contype="2"  mesh="left_finger" friction='1.5 1.5 1.5' solref='0.01 1' solimp='0.99 0.99 0.01'/>
-                                            <body name="right_hand" pos="0 -0.03 0.05" quat="-0.7071 0 0 0.7071">
-                                                <site name="ee" pos="0 0 0" rgba="0 0 1 0" type="sphere" group="1"/>
-                                                <site name="ee_x" pos="0 0 0" size="0.005 .1"  quat="0.707105 0.707108 0 0 " rgba="1 0 0 0" type="cylinder" group="1"/>
-                                                <site name="ee_z" pos="0 0 0" size="0.005 .1" quat="0.707105 0 0 0.707108" rgba="0 0 1 0" type="cylinder" group="1"/>
-                                                <site name="ee_y" pos="0 0 0" size="0.005 .1" quat="0.707105 0 0.707108 0 " rgba="0 1 0 0" type="cylinder" group="1"/>
-                                            </body>
-                                        </body>
-                                    </body>
-                                    <body name="left_inner_knuckle" pos="0 0.02 0.074098">
-                                        <inertial pos="1.86601e-06 0.0220468 0.0261335" quat="0.664139 -0.242732 0.242713 0.664146" mass="0.0230126" diaginertia="8.34216e-06 6.0949e-06 2.75601e-06" />
-                                        <joint name="left_inner_knuckle_joint" pos="0 0 0" axis="1 0 0" limited="true" range="0 0.85" />
-                                        <geom type="mesh" rgba="0 0 0 1" conaffinity="1" contype="0"  mesh="left_inner_knuckle" friction='1.5 1.5 1.5' solref='0.01 1' solimp='0.99 0.99 0.01'/>
-                                    </body>
-                                    <body name="right_outer_knuckle" pos="0 -0.035 0.059098">
-                                        <inertial pos="0 -0.021559 0.015181" quat="0.87842 0.47789 0 0" mass="0.033618" diaginertia="1.9111e-05 1.79089e-05 1.90167e-06" />
-                                        <joint name="right_outer_knuckle_joint" pos="0 0 0" axis="-1 0 0" limited="true" range="0 0.85" />
-                                        <geom type="mesh" rgba="0 0 0 1" conaffinity="1" contype="0"  mesh="right_outer_knuckle" />
-                                        <body name="right_finger" pos="0 -0.035465 0.042039">
-                                            <inertial pos="0 0.016413 0.029258" quat="0.697634 -0.115356 0.115356 0.697634" mass="0.048304" diaginertia="1.88038e-05 1.7493e-05 3.56779e-06" />
-                                            <joint name="right_finger_joint" pos="0 0 0" axis="1 0 0" limited="true" range="0 0.85" />
-                                            <geom name="j11" material="robot0:gripper_finger_mat" type="mesh" rgba="0 0 0 1" conaffinity="3" contype="2" mesh="right_finger" friction='1.5 1.5 1.5' solref='0.01 1' solimp='0.99 0.99 0.01'/>
-                                            <body name="left_hand" pos="0 0.03 0.05" quat="-0.7071 0 0 0.7071">
-                                                <site name="ee_2" pos="0 0 0" rgba="1 0 0 0" type="sphere" size="0.01" group="1"/>
-                                            </body>
-                                        </body>
-                                    </body>
-                                    <body name="right_inner_knuckle" pos="0 -0.02 0.074098">
-                                        <inertial pos="1.866e-06 -0.022047 0.026133" quat="0.66415 0.242702 -0.242721 0.664144" mass="0.023013" diaginertia="8.34209e-06 6.0949e-06 2.75601e-06" />
-                                        <joint name="right_inner_knuckle_joint" pos="0 0 0" axis="-1 0 0" limited="true" range="0 0.85" />
-                                        <geom type="mesh" rgba="0 0 0 1" conaffinity="1" contype="0" mesh="right_inner_knuckle" friction='1.5 1.5 1.5' solref='0.01 1' solimp='0.99 0.99 0.01'/>
-                                    </body>
-                                </body>
-                            </body>
-                        </body>
-                    </body>
-                </body>
-            </body>
-        </body>
-    </body>
-</mujoco>

lerobot/common/envs/simxarm/simxarm/tasks/base.py

@@ -1,145 +0,0 @@
-import os
-
-import mujoco
-import numpy as np
-from gymnasium_robotics.envs import robot_env
-
-from lerobot.common.envs.simxarm.simxarm.tasks import mocap
-
-
-class Base(robot_env.MujocoRobotEnv):
-    """
-    Superclass for all simxarm environments.
-    Args:
-            xml_name (str): name of the xml environment file
-            gripper_rotation (list): initial rotation of the gripper (given as a quaternion)
-    """
-
-    def __init__(self, xml_name, gripper_rotation=None):
-        if gripper_rotation is None:
-            gripper_rotation = [0, 1, 0, 0]
-        self.gripper_rotation = np.array(gripper_rotation, dtype=np.float32)
-        self.center_of_table = np.array([1.655, 0.3, 0.63625])
-        self.max_z = 1.2
-        self.min_z = 0.2
-        super().__init__(
-            model_path=os.path.join(os.path.dirname(__file__), "assets", xml_name + ".xml"),
-            n_substeps=20,
-            n_actions=4,
-            initial_qpos={},
-        )
-
-    @property
-    def dt(self):
-        return self.n_substeps * self.model.opt.timestep
-
-    @property
-    def eef(self):
-        return self._utils.get_site_xpos(self.model, self.data, "grasp")
-
-    @property
-    def obj(self):
-        return self._utils.get_site_xpos(self.model, self.data, "object_site")
-
-    @property
-    def robot_state(self):
-        gripper_angle = self._utils.get_joint_qpos(self.model, self.data, "right_outer_knuckle_joint")
-        return np.concatenate([self.eef, gripper_angle])
-
-    def is_success(self):
-        return NotImplementedError()
-
-    def get_reward(self):
-        raise NotImplementedError()
-
-    def _sample_goal(self):
-        raise NotImplementedError()
-
-    def get_obs(self):
-        return self._get_obs()
-
-    def _step_callback(self):
-        self._mujoco.mj_forward(self.model, self.data)
-
-    def _limit_gripper(self, gripper_pos, pos_ctrl):
-        if gripper_pos[0] > self.center_of_table[0] - 0.105 + 0.15:
-            pos_ctrl[0] = min(pos_ctrl[0], 0)
-        if gripper_pos[0] < self.center_of_table[0] - 0.105 - 0.3:
-            pos_ctrl[0] = max(pos_ctrl[0], 0)
-        if gripper_pos[1] > self.center_of_table[1] + 0.3:
-            pos_ctrl[1] = min(pos_ctrl[1], 0)
-        if gripper_pos[1] < self.center_of_table[1] - 0.3:
-            pos_ctrl[1] = max(pos_ctrl[1], 0)
-        if gripper_pos[2] > self.max_z:
-            pos_ctrl[2] = min(pos_ctrl[2], 0)
-        if gripper_pos[2] < self.min_z:
-            pos_ctrl[2] = max(pos_ctrl[2], 0)
-        return pos_ctrl
-
-    def _apply_action(self, action):
-        assert action.shape == (4,)
-        action = action.copy()
-        pos_ctrl, gripper_ctrl = action[:3], action[3]
-        pos_ctrl = self._limit_gripper(
-            self._utils.get_site_xpos(self.model, self.data, "grasp"), pos_ctrl
-        ) * (1 / self.n_substeps)
-        gripper_ctrl = np.array([gripper_ctrl, gripper_ctrl])
-        mocap.apply_action(
-            self.model,
-            self._model_names,
-            self.data,
-            np.concatenate([pos_ctrl, self.gripper_rotation, gripper_ctrl]),
-        )
-
-    def _render_callback(self):
-        self._mujoco.mj_forward(self.model, self.data)
-
-    def _reset_sim(self):
-        self.data.time = self.initial_time
-        self.data.qpos[:] = np.copy(self.initial_qpos)
-        self.data.qvel[:] = np.copy(self.initial_qvel)
-        self._sample_goal()
-        self._mujoco.mj_step(self.model, self.data, nstep=10)
-        return True
-
-    def _set_gripper(self, gripper_pos, gripper_rotation):
-        self._utils.set_mocap_pos(self.model, self.data, "robot0:mocap", gripper_pos)
-        self._utils.set_mocap_quat(self.model, self.data, "robot0:mocap", gripper_rotation)
-        self._utils.set_joint_qpos(self.model, self.data, "right_outer_knuckle_joint", 0)
-        self.data.qpos[10] = 0.0
-        self.data.qpos[12] = 0.0
-
-    def _env_setup(self, initial_qpos):
-        for name, value in initial_qpos.items():
-            self.data.set_joint_qpos(name, value)
-        mocap.reset(self.model, self.data)
-        mujoco.mj_forward(self.model, self.data)
-        self._sample_goal()
-        mujoco.mj_forward(self.model, self.data)
-
-    def reset(self):
-        self._reset_sim()
-        return self._get_obs()
-
-    def step(self, action):
-        assert action.shape == (4,)
-        assert self.action_space.contains(action), "{!r} ({}) invalid".format(action, type(action))
-        self._apply_action(action)
-        self._mujoco.mj_step(self.model, self.data, nstep=2)
-        self._step_callback()
-        obs = self._get_obs()
-        reward = self.get_reward()
-        done = False
-        info = {"is_success": self.is_success(), "success": self.is_success()}
-        return obs, reward, done, info
-
-    def render(self, mode="rgb_array", width=384, height=384):
-        self._render_callback()
-        # HACK
-        self.model.vis.global_.offwidth = width
-        self.model.vis.global_.offheight = height
-        return self.mujoco_renderer.render(mode)
-
-    def close(self):
-        if self.mujoco_renderer is not None:
-            self.mujoco_renderer.close()