Add files using upload-large-folder tool

third_party/AnyBimanual/.gitignore

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

third_party/AnyBimanual/LICENSE

@@ -0,0 +1,402 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

third_party/AnyBimanual/README.md

@@ -0,0 +1,136 @@
+<!-- <p align="center">
+  <img src="docs/logo.png" alt="AnyBimanual Logo">
+</p> -->
+
+
+
+<div align="center">
+<img src="docs/logo.png" alt="AnyBimanual Logo" width="200">
+
+# AnyBimanual: Transferring Unimanual Policy for General Bimanual Manipulation
+
+[![python](https://img.shields.io/badge/-Python_3.8_-blue?logo=python&logoColor=white)](https://github.com/pre-commit/pre-commit)
+[![pytorch](https://img.shields.io/badge/PyTorch_2.0+-ee4c2c?logo=pytorch&logoColor=white)](https://pytorch.org/get-started/locally/)
+[![wandb](https://img.shields.io/badge/Wandb_0.14.0-yellow?logo=weightsandbiases&logoColor=white)](https://wandb.ai/site/)
+[![hydra](https://img.shields.io/badge/Config-Hydra_1.0.5-89b8cd)](https://hydra.cc/)
+[![license](https://img.shields.io/badge/License-MIT-gree.svg?labelColor=gray)](https://github.com/ashleve/lightning-hydra-template#license)
+
+[Guanxing Lu <sup>*</sup>](https://guanxinglu.github.io/), [Tengbo Yu <sup>*</sup>](https://github.com/Tengbo-Yu), [Haoyuan Deng](https://github.com/Denghaoyuan123?tab=repositories), [Season Si Chen](https://www.sigs.tsinghua.edu.cn/Chensi_en/main.htm), [Yansong Tang <sup>†</sup>](https://andytang15.github.io/), [Ziwei Wang](https://ziweiwangthu.github.io/)
+
+**[[Project Page](https://anybimanual.github.io/)] | [[Paper](https://arxiv.org/pdf/2412.06779)] | [[Real-World Codebase](https://github.com/Denghaoyuan123/Bimanual_ur5e_joystick_control)] | [[YouTube](https://www.youtube.com/watch?v=RFwLgtzrXuM)] | [[X](#)]**
+</div>
+<!-- ![](docs/pipeline.png) -->
+
+**AnyBimanual** is a training framework to transfer any pretrained unimanual robotic manipulation policy to multi-task bimanual manipulation policy with few bimanual demonstrations. We first introduce a **skill manager** to dynamically schedule the skill representations discovered from pretrained unimanual policy for bimanual manipulation tasks, which linearly combines skill primitives with task-oriented compensation to represent the bimanual manipulation instruction. To mitigate the observation discrepancy between unimanual and bimanual systems, we present a **visual aligner** to generate soft masks for visual embedding of the workspace, which aims to align visual input of unimanual policy model for each arm with those during pretraining stage. AnyBimanual shows superiority on **12** simulated tasks from **RLBench2** with a sizable **12.67\%** improvement in success rate over previous methods. Experiments on **9** real-world tasks further verify its practicality with an average success rate of **84.62\%**.
+
+
+https://github.com/user-attachments/assets/2cd23178-fde9-4af9-89c4-7cd703c88156
+
+
+🎉 **NEWS**: 
+
+- *Jun. 2025:* Our paper is accepted by ICCV2025!
+- *Dec. 2024:* Codebase for both simulated and real-world experiments is released!
+
+# 📝 TODO
+- [x] Release pretrained checkpoints.
+
+# 💻 Installation
+
+**NOTE**: AnyBimanual is mainly built upon the [Perceiver-Actor^2](https://github.com/markusgrotz/peract_bimanual) repo by Markus Grotz et al.
+
+See [INSTALL.md](docs/INSTALLATION.md) for installation instructions. 
+
+See [ERROR_CATCH.md](docs/ERROR_CATCH.md) for error catching.
+
+# 🛠️ Usage
+
+The following steps are structured in order.
+
+## 🗃️ Generate Demonstrations 
+
+Please checkout the website for [pre-generated RLBench
+demonstrations](https://bimanual.github.io). If you directly use these
+datasets, you don't need to run `tools/bimanual_data_generator.py` from
+RLBench. Using these datasets will also help reproducibility since each scene
+is randomly sampled in `data_generator_bimanual.py`.
+
+
+## 🚆 Training
+We use wandb to log some curves and visualizations. Login to wandb before running the scripts.
+```bash
+wandb login
+```
+To train our PerAct + AnyBimanual, run:
+```bash
+bash scripts/train.sh BIMANUAL_PERACT 0,1 12345 ${exp_name}
+```
+where the `exp_name` can be specified as you like.
+
+To train our PerAct-LF + AnyBimanual, run:
+```bash
+bash scripts/train.sh PERACT_BC 0,1 12345 ${exp_name}
+```
+
+To train our RVT-LF + AnyBimanual, run:
+```bash
+bash scripts/train.sh RVT 0,1 12345 ${exp_name}
+```
+
+Set the `augmentation_type` in the `scripts/train.sh` to choose whether to apply the augmentation methods mentioned in our paper or to use the original SE3 augmentation.
+
+## 🔬 Evaluation
+To evaluate the checkpoint in simulator, you can use:
+```bash
+bash scripts/eval.sh BIMANUAL_PERACT 0 ${exp_name}
+```
+
+# 🦾 Real Robot
+
+### 🎮 Prepare data in real world
+
+You can refer to [Demonstrations Collection by teleoperation](https://github.com/Denghaoyuan123/Bimanual_ur5e_joystick_control) to set up your device in the real world and collect **raw data**.
+
+Convert **raw data** into [**RLbench2**](https://github.com/markusgrotz/peract_bimanual) form, run:
+```bash
+python3 anybimanual_real_supply/data/preprocess_ntu_dualarm.py
+```  
+Keyframes selection, run:
+```bash
+python3 anybimanual_real_supply/data/auto_keyframe_mani.py
+```
+
+### 🎯 Finetune
+```bash
+bash scripts/train_real.sh BIMANUAL_PERACT 0,1 12345 ${exp_name}
+```
+
+### 🕹️ Evaluation on real robot
+Run model inference scripts to receive real-world observation to generate actions, here we give an example of the Agent Class.
+```bash
+python3 anybimanual_real_supply/eval_agent_on_robot.py
+```
+
+After receiving the action generated by the model, you can refer to [Bimanual_ur5e_action_control_for_IL](https://github.com/Denghaoyuan123/Bimanual_ur5e_action_control_for_IL) to drive dual_UR5e to perform the action.
+
+# 📍 Checkpoints
+Release the [checkpoints](https://cloud.tsinghua.edu.cn/d/dd54f729a902464cb2d1/).
+
+# 🏷️ License
+This repository is released under the MIT license.
+
+# 🙏 Acknowledgement
+
+Our code is built upon [Perceiver-Actor^2](https://github.com/markusgrotz/peract_bimanual), [SkillDiffuser](https://github.com/Liang-ZX/skilldiffuser), [PerAct](https://github.com/peract/peract), [RLBench](https://github.com/stepjam/RLBench), and [CLIP](https://github.com/openai/CLIP). We thank all these authors for their nicely open sourced code and their great contributions to the community.
+
+# 🔗 Citation
+If you find this repository helpful, please consider citing:
+
+```
+@article{lu2024anybimanual,
+  title={AnyBimanual: Transferring Unimanual Policy for General Bimanual Manipulation},
+  author={Lu, Guanxing and Yu, Tengbo and Deng, Haoyuan and Chen, Season Si and Tang, Yansong and Wang, Ziwei},
+  journal={arXiv preprint arXiv:2412.06779},
+  year={2024}
+}
+```

third_party/AnyBimanual/eval.py

@@ -0,0 +1,296 @@
+import gc
+import logging
+import os
+import sys
+
+import peract_config
+
+import hydra
+import numpy as np
+import torch
+import pandas as pd
+from omegaconf import DictConfig, OmegaConf, ListConfig
+from rlbench.action_modes.action_mode import BimanualMoveArmThenGripper
+from rlbench.action_modes.action_mode import BimanualJointPositionActionMode
+from rlbench.action_modes.arm_action_modes import BimanualEndEffectorPoseViaPlanning
+from rlbench.action_modes.arm_action_modes import BimanualJointPosition, JointPosition
+from rlbench.action_modes.gripper_action_modes import BimanualDiscrete
+from rlbench.action_modes.action_mode import MoveArmThenGripper
+from rlbench.action_modes.arm_action_modes import EndEffectorPoseViaPlanning
+from rlbench.action_modes.gripper_action_modes import Discrete
+
+from rlbench.backend import task as rlbench_task
+from rlbench.backend.utils import task_file_to_task_class
+from yarr.runners.independent_env_runner import IndependentEnvRunner
+from yarr.utils.stat_accumulator import SimpleAccumulator
+
+from helpers import utils
+from helpers import observation_utils
+
+from yarr.utils.rollout_generator import RolloutGenerator
+import torch.multiprocessing as mp
+
+from agents import agent_factory
+
+
+def eval_seed(
+    train_cfg, eval_cfg, logdir, env_device, multi_task, seed, env_config
+) -> None:
+    tasks = eval_cfg.rlbench.tasks
+    rg = RolloutGenerator()
+
+    train_cfg.method.robot_name = eval_cfg.method.robot_name
+
+    agent = agent_factory.create_agent(train_cfg)
+    stat_accum = SimpleAccumulator(eval_video_fps=30)
+
+    cwd = os.getcwd()
+    weightsdir = os.path.join(logdir, "weights")
+
+    env_runner = IndependentEnvRunner(
+        train_env=None,
+        agent=agent,
+        train_replay_buffer=None,
+        num_train_envs=0,
+        num_eval_envs=eval_cfg.framework.eval_envs,
+        rollout_episodes=99999,
+        eval_episodes=eval_cfg.framework.eval_episodes,
+        training_iterations=train_cfg.framework.training_iterations,
+        eval_from_eps_number=eval_cfg.framework.eval_from_eps_number,
+        episode_length=eval_cfg.rlbench.episode_length,
+        stat_accumulator=stat_accum,
+        weightsdir=weightsdir,
+        logdir=logdir,
+        env_device=env_device,
+        rollout_generator=rg,
+        num_eval_runs=len(tasks),
+        multi_task=multi_task,
+    )
+
+    env_runner._on_thread_start = peract_config.config_logging
+
+    manager = mp.Manager()
+    save_load_lock = manager.Lock()
+    writer_lock = manager.Lock()
+
+    # evaluate all checkpoints (0, 1000, ...) which don't have results, i.e. validation phase
+    if eval_cfg.framework.eval_type == "missing":
+        weight_folders = os.listdir(weightsdir)
+        weight_folders = sorted(map(int, weight_folders))
+
+        env_data_csv_file = os.path.join(logdir, "eval_data.csv")
+        if os.path.exists(env_data_csv_file):
+            env_dict = pd.read_csv(env_data_csv_file).to_dict()
+            evaluated_weights = sorted(map(int, list(env_dict["step"].values())))
+            weight_folders = [w for w in weight_folders if w not in evaluated_weights]
+
+        print("Missing weights: ", weight_folders)
+
+    # pick the best checkpoint from validation and evaluate, i.e. test phase
+    elif eval_cfg.framework.eval_type == "best":
+        env_data_csv_file = os.path.join(logdir, "eval_data.csv")
+        if os.path.exists(env_data_csv_file):
+            env_dict = pd.read_csv(env_data_csv_file).to_dict()
+            existing_weights = list(
+                map(int, sorted(os.listdir(os.path.join(logdir, "weights"))))
+            )
+            task_weights = {}
+            for task in tasks:
+                weights = list(env_dict["step"].values())
+
+                if len(tasks) > 1:
+                    task_score = list(env_dict["eval_envs/return/%s" % task].values())
+                else:
+                    task_score = list(env_dict["eval_envs/return"].values())
+
+                avail_weights, avail_task_scores = [], []
+                for step_idx, step in enumerate(weights):
+                    if step in existing_weights:
+                        avail_weights.append(step)
+                        avail_task_scores.append(task_score[step_idx])
+
+                assert len(avail_weights) == len(avail_task_scores)
+                best_weight = avail_weights[
+                    np.argwhere(avail_task_scores == np.amax(avail_task_scores))
+                    .flatten()
+                    .tolist()[-1]
+                ]
+                task_weights[task] = best_weight
+
+            weight_folders = [task_weights]
+            print("Best weights:", weight_folders)
+        else:
+            raise Exception("No existing eval_data.csv file found in %s" % logdir)
+
+    # evaluate only the last checkpoint
+    elif eval_cfg.framework.eval_type == "last":
+        weight_folders = os.listdir(weightsdir)
+        weight_folders = sorted(map(int, weight_folders))
+        weight_folders = [weight_folders[-1]]
+        print("Last weight:", weight_folders)
+
+    elif eval_cfg.framework.eval_type == "all":
+        weight_folders = os.listdir(weightsdir)
+        weight_folders = sorted(map(int, weight_folders))
+
+    # evaluate a specific checkpoint
+    elif type(eval_cfg.framework.eval_type) == int:
+        weight_folders = [int(eval_cfg.framework.eval_type)]
+        print("Weight:", weight_folders)
+
+    elif isinstance(eval_cfg.framework.eval_type, ListConfig):
+        weight_folders = [int(w) for w in eval_cfg.framework.eval_type]
+        print("Checking specified checkpoints:", weight_folders)
+    
+    else:
+        print(type(eval_cfg.framework.eval_type))
+        raise Exception("Unknown eval type")
+
+    if len(weight_folders) == 0:
+        logging.info(
+            "No weights to evaluate. Results are already available in eval_data.csv"
+        )
+        sys.exit(0)
+    # evaluate several checkpoints in parallel
+    # NOTE: in multi-task settings, each task is evaluated serially, which makes everything slow!
+    split_n = utils.split_list(weight_folders, eval_cfg.framework.eval_envs)
+    for split in split_n:
+        processes = []
+        for e_idx, weight in enumerate(split):
+            p = mp.Process(
+                target=env_runner.start,
+                args=(
+                    weight,
+                    save_load_lock,
+                    writer_lock,
+                    env_config,
+                    e_idx % torch.cuda.device_count(),
+                    eval_cfg.framework.eval_save_metrics,
+                    eval_cfg.cinematic_recorder,
+                ),
+            )
+            p.start()
+            processes.append(p)
+        for p in processes:
+            p.join()
+
+    del env_runner
+    del agent
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+@hydra.main(config_name="eval", config_path="conf")
+def main(eval_cfg: DictConfig) -> None:
+    logging.info("\n" + OmegaConf.to_yaml(eval_cfg))
+
+    start_seed = eval_cfg.framework.start_seed
+    logdir = os.path.join(
+        eval_cfg.framework.logdir,
+        eval_cfg.rlbench.task_name,
+        eval_cfg.method.name,
+        "seed%d" % start_seed,
+    )
+
+    train_config_path = os.path.join(logdir, "config.yaml")
+
+    if os.path.exists(train_config_path):
+        with open(train_config_path, "r") as f:
+            train_cfg = OmegaConf.load(f)
+    else:
+        raise Exception(f"Missing seed{start_seed}/config.yaml. Logdir is {logdir}")
+
+    # sanity checks
+    assert(train_cfg.method.name == eval_cfg.method.name)
+    assert(train_cfg.method.agent_type == eval_cfg.method.agent_type)
+    for task in eval_cfg.rlbench.tasks:
+        assert(task in train_cfg.rlbench.tasks)
+
+    env_device = utils.get_device(eval_cfg.framework.gpu)
+    logging.info("Using env device %s." % str(env_device))
+
+    gripper_mode = eval(eval_cfg.rlbench.gripper_mode)()
+    arm_action_mode = eval(eval_cfg.rlbench.arm_action_mode)()
+    action_mode = eval(eval_cfg.rlbench.action_mode)(arm_action_mode, gripper_mode)
+
+    
+    is_bimanual = eval_cfg.method.robot_name == "bimanual"
+
+    if is_bimanual:
+        # TODO: automate instantiation with eval
+        task_path = rlbench_task.BIMANUAL_TASKS_PATH
+    else:
+        task_path = rlbench_task.TASKS_PATH
+
+    task_files = [
+        t.replace(".py", "")
+        for t in os.listdir(task_path)
+        if t != "__init__.py" and t.endswith(".py")
+    ]
+    eval_cfg.rlbench.cameras = (
+        eval_cfg.rlbench.cameras
+        if isinstance(eval_cfg.rlbench.cameras, ListConfig)
+        else [eval_cfg.rlbench.cameras]
+    )
+    obs_config = observation_utils.create_obs_config(
+        eval_cfg.rlbench.cameras,
+        eval_cfg.rlbench.camera_resolution,
+        eval_cfg.method.name,
+        eval_cfg.method.robot_name
+    )
+
+    if eval_cfg.cinematic_recorder.enabled:
+        obs_config.record_gripper_closing = True
+
+    multi_task = len(eval_cfg.rlbench.tasks) > 1
+
+    tasks = eval_cfg.rlbench.tasks
+    task_classes = []
+    for task in tasks:
+        if task not in task_files:
+            raise ValueError('Task %s not recognised!.' % task)
+        task_classes.append(task_file_to_task_class(task, is_bimanual))
+
+
+    # single-task or multi-task
+    if multi_task:
+        env_config = (
+            task_classes,
+            obs_config,
+            action_mode,
+            eval_cfg.rlbench.demo_path,
+            eval_cfg.rlbench.episode_length,
+            eval_cfg.rlbench.headless,
+            eval_cfg.framework.eval_episodes,
+            train_cfg.rlbench.include_lang_goal_in_obs,
+            eval_cfg.rlbench.time_in_state,
+            eval_cfg.framework.record_every_n,
+        )
+    else:
+        env_config = (
+            task_classes[0],
+            obs_config,
+            action_mode,
+            eval_cfg.rlbench.demo_path,
+            eval_cfg.rlbench.episode_length,
+            eval_cfg.rlbench.headless,
+            train_cfg.rlbench.include_lang_goal_in_obs,
+            eval_cfg.rlbench.time_in_state,
+            eval_cfg.framework.record_every_n,
+        )
+
+    logging.info("Evaluating seed %d." % start_seed)
+    eval_seed(
+        train_cfg,
+        eval_cfg,
+        logdir,
+        env_device,
+        multi_task,
+        start_seed,
+        env_config,
+    )
+
+
+if __name__ == "__main__":
+    peract_config.on_init()
+    main()

third_party/AnyBimanual/peract_config.py

@@ -0,0 +1,34 @@
+"""
+System configuration for peract 
+"""
+import os
+import logging
+
+import torch.multiprocessing as mp
+
+
+def config_logging(logging_level=logging.INFO, reset=False):
+
+    if reset:
+        root = logging.getLogger()
+        list(map(root.removeHandler, root.handlers))
+        list(map(root.removeFilter, root.filters))
+        
+    from rich.logging import RichHandler
+    logging.basicConfig(level=logging_level, handlers=[RichHandler()])
+
+
+def on_init():
+
+    config_logging(logging.INFO)
+    
+    logging.debug("Configuring environment.")
+    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+    mp.set_start_method("spawn", force=True)
+    mp.set_sharing_strategy("file_system")
+
+
+def on_config(cfg):    
+    
+    os.environ["MASTER_ADDR"] = str(cfg.ddp.master_addr)
+    os.environ["MASTER_PORT"] = str(cfg.ddp.master_port)

third_party/AnyBimanual/pyproject.toml

@@ -0,0 +1,33 @@
+[tool.poetry]
+name = "peract_bimanual"
+version = "0.0.1"
+description = "A perceiver actor framework for bimanual manipulation tasks"
+authors = [ "Markus Grotz <grotz@uw.edu>",
+"Mohit Shridhar <mshr@cs.washington.edu>"]
+packages = [{include = "agents"}, {include = "helpers"}, {include = "voxel"}]
+
+
+readme = "README.md"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Framework :: Robot Framework "
+]
+
+[tool.poetry.dependencies]
+python = ">=3.8,<4.0"
+einops = "0.3.2"
+ftfy = "^6.1.1"
+hydra-core = "1.0.5"
+matplotlib = "^3.7.1"
+pandas = "1.4.1"
+regex = "^2023.6.3"
+tensorboard = "^2.13.0"
+perceiver-pytorch = "^0.8.7"
+
+
+[tool.poetry.extras]
+docs = ["sphinx"]
+
+[build-system]
+requires = ["setuptools", "wheel", "poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"

third_party/AnyBimanual/run_seed_fn.py

@@ -0,0 +1,243 @@
+import os
+import pickle
+import gc
+from typing import List
+import filecmp
+import hydra
+import numpy as np
+import torch
+from omegaconf import DictConfig
+
+from rlbench import CameraConfig, ObservationConfig
+from yarr.replay_buffer.wrappers.pytorch_replay_buffer import PyTorchReplayBuffer
+from yarr.runners.offline_train_runner import OfflineTrainRunner
+from yarr.utils.stat_accumulator import SimpleAccumulator
+from yarr.replay_buffer.task_uniform_replay_buffer import TaskUniformReplayBuffer
+
+from helpers.custom_rlbench_env import CustomRLBenchEnv, CustomMultiTaskRLBenchEnv
+import torch.distributed as dist
+from torch.utils.data import DataLoader, default_collate
+from torch.utils.data.distributed import DistributedSampler
+import random
+from agents import agent_factory
+from agents import replay_utils
+from typing import Tuple, Optional
+import peract_config
+from functools import partial
+import copy
+from tqdm import tqdm
+def run_seed(
+    rank,
+    cfg: DictConfig,
+    obs_config: ObservationConfig,
+    seed,
+    world_size,
+) -> None:
+    
+
+    peract_config.config_logging()
+    
+    dist.init_process_group("gloo", rank=rank, world_size=world_size)
+
+    tasks = cfg.rlbench.tasks
+    cams = cfg.rlbench.cameras
+
+    # task_folder = "debug" if len(tasks) > 1 else tasks[0] 
+    task_folder = cfg.replay.task_folder if len(tasks) > 1 else tasks[0] 
+    # task_folder = cfg.rlbench.task_name
+    # replay_path = os.path.join(
+    #     cfg.replay.path, task_folder, cfg.method.name, "seed%d" % seed
+    # )
+    replay_path = os.path.join(
+        cfg.replay.path, task_folder
+    )
+    # to do create agent
+    agent = agent_factory.create_agent(cfg)
+
+    if not agent:
+        print("Unable to create agent")
+        return
+
+    if cfg.method.name == "ARM":
+        raise NotImplementedError("ARM is not supported yet")
+    elif cfg.method.name == "BC_LANG":
+        from agents.baselines import bc_lang
+
+        assert cfg.ddp.num_devices == 1, "BC_LANG only supports single GPU training"
+        replay_buffer = bc_lang.launch_utils.create_replay(
+            cfg.replay.batch_size,
+            cfg.replay.timesteps,
+            cfg.replay.prioritisation,
+            cfg.replay.task_uniform,
+            replay_path if cfg.replay.use_disk else None,
+            cams,
+            cfg.rlbench.camera_resolution,
+        )
+
+        bc_lang.launch_utils.fill_multi_task_replay(
+            cfg,
+            obs_config,
+            rank,
+            replay_buffer,
+            tasks,
+            cfg.rlbench.demos,
+            cfg.method.demo_augmentation,
+            cfg.method.demo_augmentation_every_n,
+            cams,
+        )
+
+
+    elif cfg.method.name == "VIT_BC_LANG":
+        from agents.baselines import vit_bc_lang
+
+        assert cfg.ddp.num_devices == 1, "VIT_BC_LANG only supports single GPU training"
+        replay_buffer = vit_bc_lang.launch_utils.create_replay(
+            cfg.replay.batch_size,
+            cfg.replay.timesteps,
+            cfg.replay.prioritisation,
+            cfg.replay.task_uniform,
+            replay_path if cfg.replay.use_disk else None,
+            cams,
+            cfg.rlbench.camera_resolution,
+        )
+
+        vit_bc_lang.launch_utils.fill_multi_task_replay(
+            cfg,
+            obs_config,
+            rank,
+            replay_buffer,
+            tasks,
+            cfg.rlbench.demos,
+            cfg.method.demo_augmentation,
+            cfg.method.demo_augmentation_every_n,
+            cams,
+        )
+
+    elif cfg.method.name.startswith("ACT_BC_LANG"):
+        from agents import act_bc_lang
+
+        assert cfg.ddp.num_devices == 1, "ACT_BC_LANG only supports single GPU training"
+        replay_buffer = act_bc_lang.launch_utils.create_replay(
+            cfg.replay.batch_size,
+            cfg.replay.timesteps,
+            cfg.replay.prioritisation,
+            cfg.replay.task_uniform,
+            replay_path if cfg.replay.use_disk else None,
+            cams,
+            cfg.rlbench.camera_resolution,
+            replay_size=3e5,
+            prev_action_horizon=cfg.method.prev_action_horizon,
+            next_action_horizon=cfg.method.next_action_horizon
+        )
+
+        act_bc_lang.launch_utils.fill_multi_task_replay(
+            cfg,
+            obs_config,
+            rank,
+            replay_buffer,
+            tasks,
+            cfg.rlbench.demos,
+            cfg.method.demo_augmentation,
+            cfg.method.demo_augmentation_every_n,
+            cams,
+        )
+
+    elif cfg.method.name == "C2FARM_LINGUNET_BC":
+        from agents import c2farm_lingunet_bc
+
+        replay_buffer = c2farm_lingunet_bc.launch_utils.create_replay(
+            cfg.replay.batch_size,
+            cfg.replay.timesteps,
+            cfg.replay.prioritisation,
+            cfg.replay.task_uniform,
+            replay_path if cfg.replay.use_disk else None,
+            cams,
+            cfg.method.voxel_sizes,
+            cfg.rlbench.camera_resolution,
+        )
+
+        c2farm_lingunet_bc.launch_utils.fill_multi_task_replay(
+            cfg,
+            obs_config,
+            rank,
+            replay_buffer,
+            tasks,
+            cfg.rlbench.demos,
+            cfg.method.demo_augmentation,
+            cfg.method.demo_augmentation_every_n,
+            cams,
+            cfg.rlbench.scene_bounds,
+            cfg.method.voxel_sizes,
+            cfg.method.bounds_offset,
+            cfg.method.rotation_resolution,
+            cfg.method.crop_augmentation,
+            keypoint_method=cfg.method.keypoint_method,
+        )
+
+
+    elif cfg.method.name.startswith("BIMANUAL_PERACT") or cfg.method.name.startswith("RVT") or cfg.method.name.startswith("PERACT_BC") or cfg.method.name.startswith("DIFFUSER_ACTOR") or cfg.method.name.startswith("ANY_BIMANUAL"):
+        print(replay_path)
+        if os.path.exists(replay_path) and os.listdir(replay_path):
+            print("Replay files found. Loading...")
+            replay_buffer = replay_utils.create_replay(cfg, replay_path)
+            replay_files = [os.path.join(replay_path, f) for f in os.listdir(replay_path) if f.endswith('.replay')]
+            for replay_file in tqdm(replay_files, desc="Processing files"):
+                with open(replay_file, 'rb') as f:
+                    try:
+                        replay_data = pickle.load(f)
+                        replay_buffer.load_add(replay_data)
+                    except pickle.UnpicklingError as e:
+                        print(f"Error unpickling file {replay_file}: {e}")
+        else:
+            print("No replay files found. Creating replay...")
+            replay_buffer = replay_utils.create_replay(cfg, replay_path)
+            replay_utils.fill_multi_task_replay(
+                cfg,
+                obs_config,
+                rank,
+                replay_buffer,
+                tasks
+            )
+
+    elif cfg.method.name == "PERACT_RL":
+        raise NotImplementedError("PERACT_RL is not supported yet")
+    else:
+        raise ValueError("Method %s does not exists." % cfg.method.name)
+
+    wrapped_replay = PyTorchReplayBuffer(
+        replay_buffer, num_workers=cfg.framework.num_workers
+    )
+    stat_accum = SimpleAccumulator(eval_video_fps=30)
+
+    cwd = os.getcwd()
+    weightsdir = os.path.join(cwd, "seed%d" % seed, "weights")
+    logdir = os.path.join(cwd, "seed%d" % seed)
+
+    train_runner = OfflineTrainRunner(
+        agent=agent,
+        wrapped_replay_buffer=wrapped_replay,
+        train_device=rank,
+        stat_accumulator=stat_accum,
+        iterations=cfg.framework.training_iterations,
+        logdir=logdir,
+        logging_level=cfg.framework.logging_level,
+        log_freq=cfg.framework.log_freq,
+        weightsdir=weightsdir,
+        num_weights_to_keep=cfg.framework.num_weights_to_keep,
+        save_freq=cfg.framework.save_freq,
+        tensorboard_logging=cfg.framework.tensorboard_logging,
+        csv_logging=cfg.framework.csv_logging,
+        load_existing_weights=cfg.framework.load_existing_weights,
+        rank=rank,
+        world_size=world_size,
+        cfg=cfg
+    )
+
+    train_runner._on_thread_start = partial(peract_config.config_logging, cfg.framework.logging_level)
+    
+    train_runner.start()
+
+    del train_runner
+    del agent
+    gc.collect()
+    torch.cuda.empty_cache()

third_party/AnyBimanual/third_party/pytorch3d/.clang-format

@@ -0,0 +1,85 @@
+AccessModifierOffset: -1
+AlignAfterOpenBracket: AlwaysBreak
+AlignConsecutiveAssignments: false
+AlignConsecutiveDeclarations: false
+AlignEscapedNewlinesLeft: true
+AlignOperands:   false
+AlignTrailingComments: false
+AllowAllParametersOfDeclarationOnNextLine: false
+AllowShortBlocksOnASingleLine: false
+AllowShortCaseLabelsOnASingleLine: false
+AllowShortFunctionsOnASingleLine: Empty
+AllowShortIfStatementsOnASingleLine: false
+AllowShortLoopsOnASingleLine: false
+AlwaysBreakAfterReturnType: None
+AlwaysBreakBeforeMultilineStrings: true
+AlwaysBreakTemplateDeclarations: true
+BinPackArguments: false
+BinPackParameters: false
+BraceWrapping:
+  AfterClass:      false
+  AfterControlStatement: false
+  AfterEnum:       false
+  AfterFunction:   false
+  AfterNamespace:  false
+  AfterObjCDeclaration: false
+  AfterStruct:     false
+  AfterUnion:      false
+  BeforeCatch:     false
+  BeforeElse:      false
+  IndentBraces:    false
+BreakBeforeBinaryOperators: None
+BreakBeforeBraces: Attach
+BreakBeforeTernaryOperators: true
+BreakConstructorInitializersBeforeComma: false
+BreakAfterJavaFieldAnnotations: false
+BreakStringLiterals: false
+ColumnLimit:     80
+CommentPragmas:  '^ IWYU pragma:'
+ConstructorInitializerAllOnOneLineOrOnePerLine: true
+ConstructorInitializerIndentWidth: 4
+ContinuationIndentWidth: 4
+Cpp11BracedListStyle: true
+DerivePointerAlignment: false
+DisableFormat:   false
+ForEachMacros:   [ FOR_EACH, FOR_EACH_R, FOR_EACH_RANGE, ]
+IncludeCategories:
+  - Regex:           '^<.*\.h(pp)?>'
+    Priority:        1
+  - Regex:           '^<.*'
+    Priority:        2
+  - Regex:           '.*'
+    Priority:        3
+IndentCaseLabels: true
+IndentWidth:     2
+IndentWrappedFunctionNames: false
+KeepEmptyLinesAtTheStartOfBlocks: false
+MacroBlockBegin: ''
+MacroBlockEnd:   ''
+MaxEmptyLinesToKeep: 1
+NamespaceIndentation: None
+ObjCBlockIndentWidth: 2
+ObjCSpaceAfterProperty: false
+ObjCSpaceBeforeProtocolList: false
+PenaltyBreakBeforeFirstCallParameter: 1
+PenaltyBreakComment: 300
+PenaltyBreakFirstLessLess: 120
+PenaltyBreakString: 1000
+PenaltyExcessCharacter: 1000000
+PenaltyReturnTypeOnItsOwnLine: 200
+PointerAlignment: Left
+ReflowComments:  true
+SortIncludes:    true
+SpaceAfterCStyleCast: false
+SpaceBeforeAssignmentOperators: true
+SpaceBeforeParens: ControlStatements
+SpaceInEmptyParentheses: false
+SpacesBeforeTrailingComments: 1
+SpacesInAngles:  false
+SpacesInContainerLiterals: true
+SpacesInCStyleCastParentheses: false
+SpacesInParentheses: false
+SpacesInSquareBrackets: false
+Standard:        Cpp11
+TabWidth:        8
+UseTab:          Never

third_party/AnyBimanual/third_party/pytorch3d/.flake8

@@ -0,0 +1,9 @@
+[flake8]
+# B028 No explicit stacklevel argument found.
+# B907 'foo' is manually surrounded by quotes, consider using the `!r` conversion flag.
+# B905 `zip()` without an explicit `strict=` parameter.
+ignore = E203, E266, E501, W503, E221, B028, B905, B907
+max-line-length = 88
+max-complexity = 18
+select = B,C,E,F,W,T4,B9
+exclude = build,__init__.py

third_party/AnyBimanual/third_party/pytorch3d/.gitignore

@@ -0,0 +1,21 @@
+build/
+dist/
+*.egg-info/
+**/__pycache__/
+*-checkpoint.ipynb
+**/.ipynb_checkpoints
+**/.ipynb_checkpoints/**
+
+
+# Docusaurus site
+website/yarn.lock
+website/build/
+website/i18n/
+website/node_modules/*
+website/npm-debug.log
+
+## Generated for tutorials
+website/_tutorials/
+website/static/files/
+website/pages/tutorials/*
+!website/pages/tutorials/index.js

third_party/AnyBimanual/third_party/pytorch3d/INSTALL.md

@@ -0,0 +1,157 @@
+# Installation
+
+
+## Requirements
+
+### Core library
+
+The core library is written in PyTorch. Several components have underlying implementation in CUDA for improved performance. A subset of these components have CPU implementations in C++/PyTorch. It is advised to use PyTorch3D with GPU support in order to use all the features.
+
+- Linux or macOS or Windows
+- Python 3.8, 3.9 or 3.10
+- PyTorch 2.0.0, 2.0.1, 2.1.0, 2.1.1, 2.1.2, 2.2.0, 2.2.1, 2.2.2, 2.3.0 or 2.3.1.
+- torchvision that matches the PyTorch installation. You can install them together as explained at pytorch.org to make sure of this.
+- gcc & g++ ≥ 4.9
+- [ioPath](https://github.com/facebookresearch/iopath)
+- If CUDA is to be used, use a version which is supported by the corresponding pytorch version and at least version 9.2.
+- If CUDA older than 11.7 is to be used and you are building from source, the CUB library must be available. We recommend version 1.10.0.
+
+The runtime dependencies can be installed by running:
+```
+conda create -n pytorch3d python=3.9
+conda activate pytorch3d
+conda install pytorch=1.13.0 torchvision pytorch-cuda=11.6 -c pytorch -c nvidia
+conda install -c iopath iopath
+```
+
+For the CUB build time dependency, which you only need if you have CUDA older than 11.7, if you are using conda, you can continue with
+```
+conda install -c bottler nvidiacub
+```
+Otherwise download the CUB library from https://github.com/NVIDIA/cub/releases and unpack it to a folder of your choice.
+Define the environment variable CUB_HOME before building and point it to the directory that contains `CMakeLists.txt` for CUB.
+For example on Linux/Mac,
+```
+curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz
+tar xzf 1.10.0.tar.gz
+export CUB_HOME=$PWD/cub-1.10.0
+```
+
+### Tests/Linting and Demos
+
+For developing on top of PyTorch3D or contributing, you will need to run the linter and tests. If you want to run any of the notebook tutorials as `docs/tutorials` or the examples in `docs/examples` you will also need matplotlib and OpenCV.
+- scikit-image
+- black
+- usort
+- flake8
+- matplotlib
+- tdqm
+- jupyter
+- imageio
+- fvcore
+- plotly
+- opencv-python
+
+These can be installed by running:
+```
+# Demos and examples
+conda install jupyter
+pip install scikit-image matplotlib imageio plotly opencv-python
+
+# Tests/Linting
+conda install -c fvcore -c conda-forge fvcore
+pip install black usort flake8 flake8-bugbear flake8-comprehensions
+```
+
+## Installing prebuilt binaries for PyTorch3D
+After installing the above dependencies, run one of the following commands:
+
+### 1. Install with CUDA support from Anaconda Cloud, on Linux only
+
+```
+# Anaconda Cloud
+conda install pytorch3d -c pytorch3d
+```
+
+Or, to install a nightly (non-official, alpha) build:
+```
+# Anaconda Cloud
+conda install pytorch3d -c pytorch3d-nightly
+```
+
+### 2. Install wheels for Linux
+We have prebuilt wheels with CUDA for Linux for PyTorch 1.11.0, for each of the supported CUDA versions,
+for Python 3.8 and 3.9. This is for ease of use on Google Colab.
+These are installed in a special way.
+For example, to install for Python 3.8, PyTorch 1.11.0 and CUDA 11.3
+```
+pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu113_pyt1110/download.html
+```
+
+In general, from inside IPython, or in Google Colab or a jupyter notebook, you can install with
+```
+import sys
+import torch
+pyt_version_str=torch.__version__.split("+")[0].replace(".", "")
+version_str="".join([
+    f"py3{sys.version_info.minor}_cu",
+    torch.version.cuda.replace(".",""),
+    f"_pyt{pyt_version_str}"
+])
+!pip install iopath
+!pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html
+```
+
+## Building / installing from source.
+CUDA support will be included if CUDA is available in pytorch or if the environment variable
+`FORCE_CUDA` is set to `1`.
+
+### 1. Install from GitHub
+```
+pip install "git+https://github.com/facebookresearch/pytorch3d.git"
+```
+To install using the code of the released version instead of from the main branch, use the following instead.
+```
+pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"
+```
+
+For CUDA builds with versions earlier than CUDA 11, set `CUB_HOME` before building as described above.
+
+**Install from Github on macOS:**
+Some environment variables should be provided, like this.
+```
+MACOSX_DEPLOYMENT_TARGET=10.14 CC=clang CXX=clang++ pip install "git+https://github.com/facebookresearch/pytorch3d.git"
+```
+
+### 2. Install from a local clone
+```
+git clone https://github.com/facebookresearch/pytorch3d.git
+cd pytorch3d && pip install -e .
+```
+To rebuild after installing from a local clone run, `rm -rf build/ **/*.so` then `pip install -e .`. You often need to rebuild pytorch3d after reinstalling PyTorch. For CUDA builds with versions earlier than CUDA 11, set `CUB_HOME` before building as described above.
+
+**Install from local clone on macOS:**
+```
+MACOSX_DEPLOYMENT_TARGET=10.14 CC=clang CXX=clang++ pip install -e .
+```
+
+**Install from local clone on Windows:**
+
+Depending on the version of PyTorch, changes to some PyTorch headers may be needed before compilation. These are often discussed in issues in this repository.
+
+After any necessary patching, you can go to "x64 Native Tools Command Prompt for VS 2019" to compile and install
+```
+cd pytorch3d
+python3 setup.py install
+```
+
+After installing, you can run **unit tests**
+```
+python3 -m unittest discover -v -s tests -t .
+```
+
+# FAQ
+
+### Can I use Docker?
+
+We don't provide a docker file but see [#113](https://github.com/facebookresearch/pytorch3d/issues/113) for a docker file shared by a user (NOTE: this has not been tested by the PyTorch3D team).

third_party/AnyBimanual/third_party/pytorch3d/LICENSE

@@ -0,0 +1,30 @@
+BSD License
+
+For PyTorch3D software
+
+Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+ * Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+ * 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.
+
+ * Neither the name Meta 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.

third_party/AnyBimanual/third_party/pytorch3d/LICENSE-3RD-PARTY

@@ -0,0 +1,71 @@
+SRN license ( https://github.com/vsitzmann/scene-representation-networks/ ):
+
+MIT License
+
+Copyright (c) 2019 Vincent Sitzmann
+
+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.
+
+
+IDR license ( github.com/lioryariv/idr ):
+
+MIT License
+
+Copyright (c) 2020 Lior Yariv
+
+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.
+
+
+NeRF https://github.com/bmild/nerf/
+
+Copyright (c) 2020 bmild
+
+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.

third_party/AnyBimanual/third_party/pytorch3d/README.md

@@ -0,0 +1,183 @@
+<img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/pytorch3dlogo.png" width="900"/>
+
+[![CircleCI](https://circleci.com/gh/facebookresearch/pytorch3d.svg?style=svg)](https://circleci.com/gh/facebookresearch/pytorch3d)
+[![Anaconda-Server Badge](https://anaconda.org/pytorch3d/pytorch3d/badges/version.svg)](https://anaconda.org/pytorch3d/pytorch3d)
+
+# Introduction
+
+PyTorch3D provides efficient, reusable components for 3D Computer Vision research with [PyTorch](https://pytorch.org).
+
+Key features include:
+
+- Data structure for storing and manipulating triangle meshes
+- Efficient operations on triangle meshes (projective transformations, graph convolution, sampling, loss functions)
+- A differentiable mesh renderer
+- Implicitron, see [its README](projects/implicitron_trainer), a framework for new-view synthesis via implicit representations. ([blog post](https://ai.facebook.com/blog/implicitron-a-new-modular-extensible-framework-for-neural-implicit-representations-in-pytorch3d/))
+
+PyTorch3D is designed to integrate smoothly with deep learning methods for predicting and manipulating 3D data.
+For this reason, all operators in PyTorch3D:
+
+- Are implemented using PyTorch tensors
+- Can handle minibatches of hetereogenous data
+- Can be differentiated
+- Can utilize GPUs for acceleration
+
+Within FAIR, PyTorch3D has been used to power research projects such as [Mesh R-CNN](https://arxiv.org/abs/1906.02739).
+
+See our [blog post](https://ai.facebook.com/blog/-introducing-pytorch3d-an-open-source-library-for-3d-deep-learning/) to see more demos and learn about PyTorch3D.
+
+## Installation
+
+For detailed instructions refer to [INSTALL.md](INSTALL.md).
+
+## License
+
+PyTorch3D is released under the [BSD License](LICENSE).
+
+## Tutorials
+
+Get started with PyTorch3D by trying one of the tutorial notebooks.
+
+|<img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/dolphin_deform.gif" width="310"/>|<img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/bundle_adjust.gif" width="310"/>|
+|:-----------------------------------------------------------------------------------------------------------:|:--------------------------------------------------:|
+| [Deform a sphere mesh to dolphin](https://github.com/facebookresearch/pytorch3d/blob/main/docs/tutorials/deform_source_mesh_to_target_mesh.ipynb)| [Bundle adjustment](https://github.com/facebookresearch/pytorch3d/blob/main/docs/tutorials/bundle_adjustment.ipynb) |
+
+| <img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/render_textured_mesh.gif" width="310"/> | <img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/camera_position_teapot.gif" width="310" height="310"/>
+|:------------------------------------------------------------:|:--------------------------------------------------:|
+| [Render textured meshes](https://github.com/facebookresearch/pytorch3d/blob/main/docs/tutorials/render_textured_meshes.ipynb)| [Camera position optimization](https://github.com/facebookresearch/pytorch3d/blob/main/docs/tutorials/camera_position_optimization_with_differentiable_rendering.ipynb)|
+
+| <img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/pointcloud_render.png" width="310"/> | <img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/cow_deform.gif" width="310" height="310"/>
+|:------------------------------------------------------------:|:--------------------------------------------------:|
+| [Render textured pointclouds](https://github.com/facebookresearch/pytorch3d/blob/main/docs/tutorials/render_colored_points.ipynb)| [Fit a mesh with texture](https://github.com/facebookresearch/pytorch3d/blob/main/docs/tutorials/fit_textured_mesh.ipynb)|
+
+| <img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/densepose_render.png" width="310"/> | <img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/shapenet_render.png" width="310" height="310"/>
+|:------------------------------------------------------------:|:--------------------------------------------------:|
+| [Render DensePose data](https://github.com/facebookresearch/pytorch3d/blob/main/docs/tutorials/render_densepose.ipynb)| [Load & Render ShapeNet data](https://github.com/facebookresearch/pytorch3d/blob/main/docs/tutorials/dataloaders_ShapeNetCore_R2N2.ipynb)|
+
+| <img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/fit_textured_volume.gif" width="310"/> | <img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/fit_nerf.gif" width="310" height="310"/>
+|:------------------------------------------------------------:|:--------------------------------------------------:|
+| [Fit Textured Volume](https://github.com/facebookresearch/pytorch3d/blob/main/docs/tutorials/fit_textured_volume.ipynb)| [Fit A Simple Neural Radiance Field](https://github.com/facebookresearch/pytorch3d/blob/main/docs/tutorials/fit_simple_neural_radiance_field.ipynb)|
+
+| <img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/fit_textured_volume.gif" width="310"/> | <img src="https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/.github/implicitron_config.gif" width="310" height="310"/>
+|:------------------------------------------------------------:|:--------------------------------------------------:|
+| [Fit Textured Volume in Implicitron](https://github.com/facebookresearch/pytorch3d/blob/main/docs/tutorials/implicitron_volumes.ipynb)| [Implicitron Config System](https://github.com/facebookresearch/pytorch3d/blob/main/docs/tutorials/implicitron_config_system.ipynb)|
+
+
+
+
+
+## Documentation
+
+Learn more about the API by reading the PyTorch3D [documentation](https://pytorch3d.readthedocs.org/).
+
+We also have deep dive notes on several API components:
+
+- [Heterogeneous Batching](https://github.com/facebookresearch/pytorch3d/tree/main/docs/notes/batching.md)
+- [Mesh IO](https://github.com/facebookresearch/pytorch3d/tree/main/docs/notes/meshes_io.md)
+- [Differentiable Rendering](https://github.com/facebookresearch/pytorch3d/tree/main/docs/notes/renderer_getting_started.md)
+
+### Overview Video
+
+We have created a short (~14 min) video tutorial providing an overview of the PyTorch3D codebase including several code examples. Click on the image below to watch the video on YouTube:
+
+<a href="http://www.youtube.com/watch?v=Pph1r-x9nyY"><img src="http://img.youtube.com/vi/Pph1r-x9nyY/0.jpg" height="225" ></a>
+
+## Development
+
+We welcome new contributions to PyTorch3D and we will be actively maintaining this library! Please refer to [CONTRIBUTING.md](./.github/CONTRIBUTING.md) for full instructions on how to run the code, tests and linter, and submit your pull requests.
+
+## Development and Compatibility
+
+- `main` branch: actively developed, without any guarantee, Anything can be broken at any time
+  - REMARK: this includes nightly builds which are built from `main`
+  - HINT: the commit history can help locate regressions or changes
+- backward-compatibility between releases: no guarantee. Best efforts to communicate breaking changes and facilitate migration of code or data (incl. models).
+
+## Contributors
+
+PyTorch3D is written and maintained by the Facebook AI Research Computer Vision Team.
+
+In alphabetical order:
+
+* Amitav Baruah
+* Steve Branson
+* Krzysztof Chalupka
+* Jiali Duan
+* Luya Gao
+* Georgia Gkioxari
+* Taylor Gordon
+* Justin Johnson
+* Patrick Labatut
+* Christoph Lassner
+* Wan-Yen Lo
+* David Novotny
+* Nikhila Ravi
+* Jeremy Reizenstein
+* Dave Schnizlein
+* Roman Shapovalov
+* Olivia Wiles
+
+## Citation
+
+If you find PyTorch3D useful in your research, please cite our tech report:
+
+```bibtex
+@article{ravi2020pytorch3d,
+    author = {Nikhila Ravi and Jeremy Reizenstein and David Novotny and Taylor Gordon
+                  and Wan-Yen Lo and Justin Johnson and Georgia Gkioxari},
+    title = {Accelerating 3D Deep Learning with PyTorch3D},
+    journal = {arXiv:2007.08501},
+    year = {2020},
+}
+```
+
+If you are using the pulsar backend for sphere-rendering (the `PulsarPointRenderer` or `pytorch3d.renderer.points.pulsar.Renderer`), please cite the tech report:
+
+```bibtex
+@article{lassner2020pulsar,
+    author = {Christoph Lassner and Michael Zollh\"ofer},
+    title = {Pulsar: Efficient Sphere-based Neural Rendering},
+    journal = {arXiv:2004.07484},
+    year = {2020},
+}
+```
+
+## News
+
+Please see below for a timeline of the codebase updates in reverse chronological order. We are sharing updates on the releases as well as research projects which are built with PyTorch3D. The changelogs for the releases are available under [`Releases`](https://github.com/facebookresearch/pytorch3d/releases),  and the builds can be installed using `conda` as per the instructions in [INSTALL.md](INSTALL.md).
+
+**[Oct 31st 2023]:**   PyTorch3D [v0.7.5](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.5) released.
+
+**[May 10th 2023]:**   PyTorch3D [v0.7.4](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.4) released.
+
+**[Apr 5th 2023]:**   PyTorch3D [v0.7.3](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.3) released.
+
+**[Dec 19th 2022]:**   PyTorch3D [v0.7.2](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.2) released.
+
+**[Oct 23rd 2022]:**   PyTorch3D [v0.7.1](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.1) released.
+
+**[Aug 10th 2022]:**   PyTorch3D [v0.7.0](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.0) released with Implicitron and MeshRasterizerOpenGL.
+
+**[Apr 28th 2022]:**   PyTorch3D [v0.6.2](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.6.2) released
+
+**[Dec 16th 2021]:**   PyTorch3D [v0.6.1](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.6.1) released
+
+**[Oct 6th 2021]:**   PyTorch3D [v0.6.0](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.6.0) released
+
+**[Aug 5th 2021]:**   PyTorch3D [v0.5.0](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.5.0) released
+
+**[Feb 9th 2021]:** PyTorch3D [v0.4.0](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.4.0) released with support for implicit functions, volume rendering and a [reimplementation of NeRF](https://github.com/facebookresearch/pytorch3d/tree/main/projects/nerf).
+
+**[November 2nd 2020]:** PyTorch3D [v0.3.0](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.3.0) released, integrating the pulsar backend.
+
+**[Aug 28th 2020]:**   PyTorch3D [v0.2.5](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.2.5) released
+
+**[July 17th 2020]:**   PyTorch3D tech report published on ArXiv: https://arxiv.org/abs/2007.08501
+
+**[April 24th 2020]:**   PyTorch3D [v0.2.0](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.2.0) released
+
+**[March 25th 2020]:**   [SynSin](https://arxiv.org/abs/1912.08804) codebase released using PyTorch3D: https://github.com/facebookresearch/synsin
+
+**[March 8th 2020]:**   PyTorch3D [v0.1.1](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.1.1) bug fix release
+
+**[Jan 23rd 2020]:**   PyTorch3D [v0.1.0](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.1.0) released. [Mesh R-CNN](https://arxiv.org/abs/1906.02739) codebase released: https://github.com/facebookresearch/meshrcnn

third_party/AnyBimanual/third_party/pytorch3d/setup.cfg

@@ -0,0 +1,14 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+[isort]
+line_length = 88
+multi_line_output = 3
+include_trailing_comma = True
+force_grid_warp = 0
+default_section = THIRDPARTY
+lines_after_imports = 2
+combine_as_imports = True

third_party/AnyBimanual/third_party/pytorch3d/setup.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import glob
+import os
+import runpy
+import sys
+import warnings
+from typing import List, Optional
+
+import torch
+from setuptools import find_packages, setup
+from torch.utils.cpp_extension import CppExtension, CUDA_HOME, CUDAExtension
+
+
+def get_existing_ccbin(nvcc_args: List[str]) -> Optional[str]:
+    """
+    Given a list of nvcc arguments, return the compiler if specified.
+
+    Note from CUDA doc: Single value options and list options must have
+    arguments, which must follow the name of the option itself by either
+    one of more spaces or an equals character.
+    """
+    last_arg = None
+    for arg in reversed(nvcc_args):
+        if arg == "-ccbin":
+            return last_arg
+        if arg.startswith("-ccbin="):
+            return arg[7:]
+        last_arg = arg
+    return None
+
+
+def get_extensions():
+    no_extension = os.getenv("PYTORCH3D_NO_EXTENSION", "0") == "1"
+    if no_extension:
+        msg = "SKIPPING EXTENSION BUILD. PYTORCH3D WILL NOT WORK!"
+        print(msg, file=sys.stderr)
+        warnings.warn(msg)
+        return []
+
+    this_dir = os.path.dirname(os.path.abspath(__file__))
+    extensions_dir = os.path.join(this_dir, "pytorch3d", "csrc")
+    sources = glob.glob(os.path.join(extensions_dir, "**", "*.cpp"), recursive=True)
+    source_cuda = glob.glob(os.path.join(extensions_dir, "**", "*.cu"), recursive=True)
+    extension = CppExtension
+
+    extra_compile_args = {"cxx": ["-std=c++17"]}
+    define_macros = []
+    include_dirs = [extensions_dir]
+
+    force_cuda = os.getenv("FORCE_CUDA", "0") == "1"
+    force_no_cuda = os.getenv("PYTORCH3D_FORCE_NO_CUDA", "0") == "1"
+    if (
+        not force_no_cuda and torch.cuda.is_available() and CUDA_HOME is not None
+    ) or force_cuda:
+        extension = CUDAExtension
+        sources += source_cuda
+        define_macros += [("WITH_CUDA", None)]
+        # Thrust is only used for its tuple objects.
+        # With CUDA 11.0 we can't use the cudatoolkit's version of cub.
+        # We take the risk that CUB and Thrust are incompatible, because
+        # we aren't using parts of Thrust which actually use CUB.
+        define_macros += [("THRUST_IGNORE_CUB_VERSION_CHECK", None)]
+        cub_home = os.environ.get("CUB_HOME", None)
+        nvcc_args = [
+            "-DCUDA_HAS_FP16=1",
+            "-D__CUDA_NO_HALF_OPERATORS__",
+            "-D__CUDA_NO_HALF_CONVERSIONS__",
+            "-D__CUDA_NO_HALF2_OPERATORS__",
+        ]
+        if os.name != "nt":
+            nvcc_args.append("-std=c++17")
+        if cub_home is None:
+            prefix = os.environ.get("CONDA_PREFIX", None)
+            if prefix is not None and os.path.isdir(prefix + "/include/cub"):
+                cub_home = prefix + "/include"
+
+        if cub_home is None:
+            warnings.warn(
+                "The environment variable `CUB_HOME` was not found. "
+                "NVIDIA CUB is required for compilation and can be downloaded "
+                "from `https://github.com/NVIDIA/cub/releases`. You can unpack "
+                "it to a location of your choice and set the environment variable "
+                "`CUB_HOME` to the folder containing the `CMakeListst.txt` file."
+            )
+        else:
+            include_dirs.append(os.path.realpath(cub_home).replace("\\ ", " "))
+        nvcc_flags_env = os.getenv("NVCC_FLAGS", "")
+        if nvcc_flags_env != "":
+            nvcc_args.extend(nvcc_flags_env.split(" "))
+
+        # This is needed for pytorch 1.6 and earlier. See e.g.
+        # https://github.com/facebookresearch/pytorch3d/issues/436
+        # It is harmless after https://github.com/pytorch/pytorch/pull/47404 .
+        # But it can be problematic in torch 1.7.0 and 1.7.1
+        if torch.__version__[:4] != "1.7.":
+            CC = os.environ.get("CC", None)
+            if CC is not None:
+                existing_CC = get_existing_ccbin(nvcc_args)
+                if existing_CC is None:
+                    CC_arg = "-ccbin={}".format(CC)
+                    nvcc_args.append(CC_arg)
+                elif existing_CC != CC:
+                    msg = f"Inconsistent ccbins: {CC} and {existing_CC}"
+                    raise ValueError(msg)
+
+        extra_compile_args["nvcc"] = nvcc_args
+
+    sources = [os.path.join(extensions_dir, s) for s in sources]
+
+    ext_modules = [
+        extension(
+            "pytorch3d._C",
+            sources,
+            include_dirs=include_dirs,
+            define_macros=define_macros,
+            extra_compile_args=extra_compile_args,
+        )
+    ]
+
+    return ext_modules
+
+
+# Retrieve __version__ from the package.
+__version__ = runpy.run_path("pytorch3d/__init__.py")["__version__"]
+
+
+if os.getenv("PYTORCH3D_NO_NINJA", "0") == "1":
+
+    class BuildExtension(torch.utils.cpp_extension.BuildExtension):
+        def __init__(self, *args, **kwargs):
+            super().__init__(use_ninja=False, *args, **kwargs)
+
+else:
+    BuildExtension = torch.utils.cpp_extension.BuildExtension
+
+trainer = "pytorch3d.implicitron_trainer"
+
+setup(
+    name="pytorch3d",
+    version=__version__,
+    author="FAIR",
+    url="https://github.com/facebookresearch/pytorch3d",
+    description="PyTorch3D is FAIR's library of reusable components "
+    "for deep Learning with 3D data.",
+    packages=find_packages(
+        exclude=("configs", "tests", "tests.*", "docs.*", "projects.*")
+    )
+    + [trainer],
+    package_dir={trainer: "projects/implicitron_trainer"},
+    install_requires=["iopath"],
+    extras_require={
+        "all": ["matplotlib", "tqdm>4.29.0", "imageio", "ipywidgets"],
+        "dev": ["flake8", "usort"],
+        "implicitron": [
+            "hydra-core>=1.1",
+            "visdom",
+            "lpips",
+            "tqdm>4.29.0",
+            "matplotlib",
+            "accelerate",
+            "sqlalchemy>=2.0",
+        ],
+    },
+    entry_points={
+        "console_scripts": [
+            f"pytorch3d_implicitron_runner={trainer}.experiment:experiment",
+            f"pytorch3d_implicitron_visualizer={trainer}.visualize_reconstruction:main",
+        ]
+    },
+    ext_modules=get_extensions(),
+    cmdclass={"build_ext": BuildExtension},
+    package_data={
+        "": ["*.json"],
+    },
+)

third_party/AnyBimanual/third_party/pytorch3d/tests/test_shapenet_core.py

@@ -0,0 +1,297 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Sanity checks for loading ShapeNetCore.
+"""
+import os
+import unittest
+
+import numpy as np
+import torch
+from PIL import Image
+from pytorch3d.datasets import collate_batched_meshes, ShapeNetCore
+from pytorch3d.renderer import (
+    FoVPerspectiveCameras,
+    look_at_view_transform,
+    PointLights,
+    RasterizationSettings,
+)
+from torch.utils.data import DataLoader
+
+from .common_testing import get_tests_dir, load_rgb_image, TestCaseMixin
+
+
+# Set the SHAPENET_PATH to the local path to the dataset
+SHAPENET_PATH = None
+VERSION = 1
+# If DEBUG=True, save out images generated in the tests for debugging.
+# All saved images have prefix DEBUG_
+DEBUG = False
+DATA_DIR = get_tests_dir() / "data"
+
+
+class TestShapenetCore(TestCaseMixin, unittest.TestCase):
+    def setUp(self):
+        """
+        Check if the ShapeNet dataset is provided in the repo.
+        If not, download this separately and update the shapenet_path`
+        with the location of the dataset in order to run the tests.
+        """
+        if SHAPENET_PATH is None or not os.path.exists(SHAPENET_PATH):
+            url = "https://www.shapenet.org/"
+            msg = (
+                "ShapeNet data not found, download from %s, update "
+                "SHAPENET_PATH at the top of the file, and rerun."
+            )
+
+            self.skipTest(msg % url)
+
+    def test_load_shapenet_core(self):
+        """
+        Test loading both the entire ShapeNetCore dataset and a subset of the ShapeNetCore
+        dataset. Check the loaded datasets return items of the correct shapes and types.
+        """
+        # Try loading ShapeNetCore with an invalid version number and catch error.
+        with self.assertRaises(ValueError) as err:
+            ShapeNetCore(SHAPENET_PATH, version=3)
+        self.assertTrue("Version number must be either 1 or 2." in str(err.exception))
+
+        # Load ShapeNetCore without specifying any particular categories.
+        shapenet_dataset = ShapeNetCore(SHAPENET_PATH, version=VERSION)
+
+        # Count the number of grandchildren directories (which should be equal to
+        # the total number of objects in the dataset) by walking through the given
+        # directory.
+        wnsynset_list = [
+            wnsynset
+            for wnsynset in os.listdir(SHAPENET_PATH)
+            if os.path.isdir(os.path.join(SHAPENET_PATH, wnsynset))
+        ]
+        model_num_list = [
+            (len(next(os.walk(os.path.join(SHAPENET_PATH, wnsynset)))[1]))
+            for wnsynset in wnsynset_list
+        ]
+        # Check total number of objects in the dataset is correct.
+        self.assertEqual(len(shapenet_dataset), sum(model_num_list))
+
+        # Randomly retrieve an object from the dataset.
+        rand_obj = shapenet_dataset[torch.randint(len(shapenet_dataset), (1,))]
+        # Check that data types and shapes of items returned by __getitem__ are correct.
+        verts, faces = rand_obj["verts"], rand_obj["faces"]
+        self.assertTrue(verts.dtype == torch.float32)
+        self.assertTrue(faces.dtype == torch.int64)
+        self.assertEqual(verts.ndim, 2)
+        self.assertEqual(verts.shape[-1], 3)
+        self.assertEqual(faces.ndim, 2)
+        self.assertEqual(faces.shape[-1], 3)
+
+        # Load six categories from ShapeNetCore.
+        # Specify categories with a combination of offsets and labels.
+        shapenet_subset = ShapeNetCore(
+            SHAPENET_PATH,
+            synsets=[
+                "04330267",
+                "guitar",
+                "02801938",
+                "birdhouse",
+                "03991062",
+                "tower",
+            ],
+            version=1,
+        )
+        subset_offsets = [
+            "04330267",
+            "03467517",
+            "02801938",
+            "02843684",
+            "03991062",
+            "04460130",
+        ]
+        subset_model_nums = [
+            (len(next(os.walk(os.path.join(SHAPENET_PATH, offset)))[1]))
+            for offset in subset_offsets
+        ]
+        self.assertEqual(len(shapenet_subset), sum(subset_model_nums))
+
+    def test_collate_models(self):
+        """
+        Test collate_batched_meshes returns items of the correct shapes and types.
+        Check that when collate_batched_meshes is passed to Dataloader, batches of
+        the correct shapes and types are returned.
+        """
+        # Load ShapeNetCore without specifying any particular categories.
+        shapenet_dataset = ShapeNetCore(SHAPENET_PATH)
+        # Randomly retrieve several objects from the dataset.
+        rand_idxs = torch.randint(len(shapenet_dataset), (6,))
+        rand_objs = [shapenet_dataset[idx] for idx in rand_idxs]
+
+        # Collate the randomly selected objects
+        collated_meshes = collate_batched_meshes(rand_objs)
+        verts, faces = (collated_meshes["verts"], collated_meshes["faces"])
+        self.assertEqual(len(verts), 6)
+        self.assertEqual(len(faces), 6)
+
+        # Pass the custom collate_fn function to DataLoader and check elements
+        # in batch have the correct shape.
+        batch_size = 12
+        shapenet_core_loader = DataLoader(
+            shapenet_dataset, batch_size=batch_size, collate_fn=collate_batched_meshes
+        )
+        it = iter(shapenet_core_loader)
+        object_batch = next(it)
+        self.assertEqual(len(object_batch["synset_id"]), batch_size)
+        self.assertEqual(len(object_batch["model_id"]), batch_size)
+        self.assertEqual(len(object_batch["label"]), batch_size)
+        self.assertEqual(object_batch["mesh"].verts_padded().shape[0], batch_size)
+        self.assertEqual(object_batch["mesh"].faces_padded().shape[0], batch_size)
+
+    def test_catch_render_arg_errors(self):
+        """
+        Test rendering ShapeNetCore with invalid model_ids, categories or indices,
+        and catch corresponding errors.
+        """
+        # Load ShapeNetCore.
+        shapenet_dataset = ShapeNetCore(SHAPENET_PATH)
+
+        # Try loading with an invalid model_id and catch error.
+        with self.assertRaises(ValueError) as err:
+            shapenet_dataset.render(model_ids=["piano0"])
+        self.assertTrue("not found in the loaded dataset" in str(err.exception))
+
+        # Try loading with an index out of bounds and catch error.
+        with self.assertRaises(IndexError) as err:
+            shapenet_dataset.render(idxs=[100000])
+        self.assertTrue("are out of bounds" in str(err.exception))
+
+    def test_render_shapenet_core(self):
+        """
+        Test rendering objects from ShapeNetCore.
+        """
+        # Setup device and seed for random selections.
+        device = torch.device("cuda:0")
+        torch.manual_seed(39)
+
+        # Load category piano from ShapeNetCore.
+        piano_dataset = ShapeNetCore(SHAPENET_PATH, synsets=["piano"])
+
+        # Rendering settings.
+        R, T = look_at_view_transform(1.0, 1.0, 90)
+        cameras = FoVPerspectiveCameras(R=R, T=T, device=device)
+        raster_settings = RasterizationSettings(image_size=512)
+        lights = PointLights(
+            location=torch.tensor([0.0, 1.0, -2.0], device=device)[None],
+            # TODO: debug the source of the discrepancy in two images when rendering on GPU.
+            diffuse_color=((0, 0, 0),),
+            specular_color=((0, 0, 0),),
+            device=device,
+        )
+
+        # Render first three models in the piano category.
+        pianos = piano_dataset.render(
+            idxs=list(range(3)),
+            device=device,
+            cameras=cameras,
+            raster_settings=raster_settings,
+            lights=lights,
+        )
+        # Check that there are three images in the batch.
+        self.assertEqual(pianos.shape[0], 3)
+
+        # Compare the rendered models to the reference images.
+        for idx in range(3):
+            piano_rgb = pianos[idx, ..., :3].squeeze().cpu()
+            if DEBUG:
+                Image.fromarray((piano_rgb.numpy() * 255).astype(np.uint8)).save(
+                    DATA_DIR / ("DEBUG_shapenet_core_render_piano_by_idxs_%s.png" % idx)
+                )
+            image_ref = load_rgb_image(
+                "test_shapenet_core_render_piano_%s.png" % idx, DATA_DIR
+            )
+            self.assertClose(piano_rgb, image_ref, atol=0.05)
+
+        # Render the same piano models but by model_ids this time.
+        pianos_2 = piano_dataset.render(
+            model_ids=[
+                "13394ca47c89f91525a3aaf903a41c90",
+                "14755c2ee8e693aba508f621166382b0",
+                "156c4207af6d2c8f1fdc97905708b8ea",
+            ],
+            device=device,
+            cameras=cameras,
+            raster_settings=raster_settings,
+            lights=lights,
+        )
+
+        # Compare the rendered models to the reference images.
+        for idx in range(3):
+            piano_rgb_2 = pianos_2[idx, ..., :3].squeeze().cpu()
+            if DEBUG:
+                Image.fromarray((piano_rgb_2.numpy() * 255).astype(np.uint8)).save(
+                    DATA_DIR / ("DEBUG_shapenet_core_render_piano_by_ids_%s.png" % idx)
+                )
+            image_ref = load_rgb_image(
+                "test_shapenet_core_render_piano_%s.png" % idx, DATA_DIR
+            )
+            self.assertClose(piano_rgb_2, image_ref, atol=0.05)
+
+        #######################
+        # Render by categories
+        #######################
+
+        # Load ShapeNetCore.
+        shapenet_dataset = ShapeNetCore(SHAPENET_PATH)
+
+        # Render a mixture of categories and specify the number of models to be
+        # randomly sampled from each category.
+        mixed_objs = shapenet_dataset.render(
+            categories=["faucet", "chair"],
+            sample_nums=[2, 1],
+            device=device,
+            cameras=cameras,
+            raster_settings=raster_settings,
+            lights=lights,
+        )
+        # Compare the rendered models to the reference images.
+        for idx in range(3):
+            mixed_rgb = mixed_objs[idx, ..., :3].squeeze().cpu()
+            if DEBUG:
+                Image.fromarray((mixed_rgb.numpy() * 255).astype(np.uint8)).save(
+                    DATA_DIR
+                    / ("DEBUG_shapenet_core_render_mixed_by_categories_%s.png" % idx)
+                )
+            image_ref = load_rgb_image(
+                "test_shapenet_core_render_mixed_by_categories_%s.png" % idx, DATA_DIR
+            )
+            self.assertClose(mixed_rgb, image_ref, atol=0.05)
+
+        # Render a mixture of categories without specifying sample_nums.
+        mixed_objs_2 = shapenet_dataset.render(
+            categories=["faucet", "chair"],
+            device=device,
+            cameras=cameras,
+            raster_settings=raster_settings,
+            lights=lights,
+        )
+        # Compare the rendered models to the reference images.
+        for idx in range(2):
+            mixed_rgb_2 = mixed_objs_2[idx, ..., :3].squeeze().cpu()
+            if DEBUG:
+                Image.fromarray((mixed_rgb_2.numpy() * 255).astype(np.uint8)).save(
+                    DATA_DIR
+                    / ("DEBUG_shapenet_core_render_without_sample_nums_%s.png" % idx)
+                )
+            image_ref = load_rgb_image(
+                "test_shapenet_core_render_without_sample_nums_%s.png" % idx, DATA_DIR
+            )
+            self.assertClose(mixed_rgb_2, image_ref, atol=0.05)
+
+    def test_load_textures_false(self):
+        shapenet_dataset = ShapeNetCore(
+            SHAPENET_PATH, load_textures=False, version=VERSION
+        )
+        model = shapenet_dataset[0]
+        self.assertIsNone(model["textures"])

third_party/AnyBimanual/third_party/pytorch3d/tests/test_so3.py

@@ -0,0 +1,283 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import math
+import unittest
+from distutils.version import LooseVersion
+
+import numpy as np
+import torch
+from pytorch3d.transforms.so3 import (
+    hat,
+    so3_exp_map,
+    so3_log_map,
+    so3_relative_angle,
+    so3_rotation_angle,
+)
+
+from .common_testing import TestCaseMixin
+
+
+class TestSO3(TestCaseMixin, unittest.TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        torch.manual_seed(42)
+        np.random.seed(42)
+
+    @staticmethod
+    def init_log_rot(batch_size: int = 10):
+        """
+        Initialize a list of `batch_size` 3-dimensional vectors representing
+        randomly generated logarithms of rotation matrices.
+        """
+        device = torch.device("cuda:0")
+        log_rot = torch.randn((batch_size, 3), dtype=torch.float32, device=device)
+        return log_rot
+
+    @staticmethod
+    def init_rot(batch_size: int = 10):
+        """
+        Randomly generate a batch of `batch_size` 3x3 rotation matrices.
+        """
+        device = torch.device("cuda:0")
+
+        # TODO(dnovotny): replace with random_rotation from random_rotation.py
+        rot = []
+        for _ in range(batch_size):
+            r = torch.linalg.qr(torch.randn((3, 3), device=device))[0]
+            f = torch.randint(2, (3,), device=device, dtype=torch.float32)
+            if f.sum() % 2 == 0:
+                f = 1 - f
+            rot.append(r * (2 * f - 1).float())
+        rot = torch.stack(rot)
+
+        return rot
+
+    def test_determinant(self):
+        """
+        Tests whether the determinants of 3x3 rotation matrices produced
+        by `so3_exp_map` are (almost) equal to 1.
+        """
+        log_rot = TestSO3.init_log_rot(batch_size=30)
+        Rs = so3_exp_map(log_rot)
+        dets = torch.det(Rs)
+        self.assertClose(dets, torch.ones_like(dets), atol=1e-4)
+
+    def test_cross(self):
+        """
+        For a pair of randomly generated 3-dimensional vectors `a` and `b`,
+        tests whether a matrix product of `hat(a)` and `b` equals the result
+        of a cross product between `a` and `b`.
+        """
+        device = torch.device("cuda:0")
+        a, b = torch.randn((2, 100, 3), dtype=torch.float32, device=device)
+        hat_a = hat(a)
+        cross = torch.bmm(hat_a, b[:, :, None])[:, :, 0]
+        torch_cross = torch.cross(a, b, dim=1)
+        self.assertClose(torch_cross, cross, atol=1e-4)
+
+    def test_bad_so3_input_value_err(self):
+        """
+        Tests whether `so3_exp_map` and `so3_log_map` correctly return
+        a ValueError if called with an argument of incorrect shape or, in case
+        of `so3_exp_map`, unexpected trace.
+        """
+        device = torch.device("cuda:0")
+        log_rot = torch.randn(size=[5, 4], device=device)
+        with self.assertRaises(ValueError) as err:
+            so3_exp_map(log_rot)
+        self.assertTrue("Input tensor shape has to be Nx3." in str(err.exception))
+
+        rot = torch.randn(size=[5, 3, 5], device=device)
+        with self.assertRaises(ValueError) as err:
+            so3_log_map(rot)
+        self.assertTrue("Input has to be a batch of 3x3 Tensors." in str(err.exception))
+
+    def test_so3_exp_singularity(self, batch_size: int = 100):
+        """
+        Tests whether the `so3_exp_map` is robust to the input vectors
+        the norms of which are close to the numerically unstable region
+        (vectors with low l2-norms).
+        """
+        # generate random log-rotations with a tiny angle
+        log_rot = TestSO3.init_log_rot(batch_size=batch_size)
+        log_rot_small = log_rot * 1e-6
+        log_rot_small.requires_grad = True
+        R = so3_exp_map(log_rot_small)
+        # tests whether all outputs are finite
+        self.assertTrue(torch.isfinite(R).all())
+        # tests whether the gradient is not None and all finite
+        loss = R.sum()
+        loss.backward()
+        self.assertIsNotNone(log_rot_small.grad)
+        self.assertTrue(torch.isfinite(log_rot_small.grad).all())
+
+    def test_so3_log_singularity(self, batch_size: int = 100):
+        """
+        Tests whether the `so3_log_map` is robust to the input matrices
+        who's rotation angles are close to the numerically unstable region
+        (i.e. matrices with low rotation angles).
+        """
+        # generate random rotations with a tiny angle
+        device = torch.device("cuda:0")
+        identity = torch.eye(3, device=device)
+        rot180 = identity * torch.tensor([[1.0, -1.0, -1.0]], device=device)
+        r = [identity, rot180]
+        # add random rotations and random almost orthonormal matrices
+        r.extend(
+            [
+                torch.linalg.qr(identity + torch.randn_like(identity) * 1e-4)[0]
+                + float(i > batch_size // 2) * (0.5 - torch.rand_like(identity)) * 1e-3
+                # this adds random noise to the second half
+                # of the random orthogonal matrices to generate
+                # near-orthogonal matrices
+                for i in range(batch_size - 2)
+            ]
+        )
+        r = torch.stack(r)
+        r.requires_grad = True
+        # the log of the rotation matrix r
+        r_log = so3_log_map(r, cos_bound=1e-4, eps=1e-2)
+        # tests whether all outputs are finite
+        self.assertTrue(torch.isfinite(r_log).all())
+        # tests whether the gradient is not None and all finite
+        loss = r.sum()
+        loss.backward()
+        self.assertIsNotNone(r.grad)
+        self.assertTrue(torch.isfinite(r.grad).all())
+
+    def test_so3_log_to_exp_to_log_to_exp(self, batch_size: int = 100):
+        """
+        Check that
+        `so3_exp_map(so3_log_map(so3_exp_map(log_rot)))
+        == so3_exp_map(log_rot)`
+        for a randomly generated batch of rotation matrix logarithms `log_rot`.
+        Unlike `test_so3_log_to_exp_to_log`, this test checks the
+        correctness of converting a `log_rot` which contains values > math.pi.
+        """
+        log_rot = 2.0 * TestSO3.init_log_rot(batch_size=batch_size)
+        # check also the singular cases where rot. angle = {0, 2pi}
+        log_rot[:2] = 0
+        log_rot[1, 0] = 2.0 * math.pi - 1e-6
+        rot = so3_exp_map(log_rot, eps=1e-4)
+        rot_ = so3_exp_map(so3_log_map(rot, eps=1e-4, cos_bound=1e-6), eps=1e-6)
+        self.assertClose(rot, rot_, atol=0.01)
+        angles = so3_relative_angle(rot, rot_, cos_bound=1e-6)
+        self.assertClose(angles, torch.zeros_like(angles), atol=0.01)
+
+    def test_so3_log_to_exp_to_log(self, batch_size: int = 100):
+        """
+        Check that `so3_log_map(so3_exp_map(log_rot))==log_rot` for
+        a randomly generated batch of rotation matrix logarithms `log_rot`.
+        """
+        log_rot = TestSO3.init_log_rot(batch_size=batch_size)
+        # check also the singular cases where rot. angle = 0
+        log_rot[:1] = 0
+        log_rot_ = so3_log_map(so3_exp_map(log_rot))
+        self.assertClose(log_rot, log_rot_, atol=1e-4)
+
+    def test_so3_exp_to_log_to_exp(self, batch_size: int = 100):
+        """
+        Check that `so3_exp_map(so3_log_map(R))==R` for
+        a batch of randomly generated rotation matrices `R`.
+        """
+        rot = TestSO3.init_rot(batch_size=batch_size)
+        non_singular = (so3_rotation_angle(rot) - math.pi).abs() > 1e-2
+        rot = rot[non_singular]
+        rot_ = so3_exp_map(so3_log_map(rot, eps=1e-8, cos_bound=1e-8), eps=1e-8)
+        self.assertClose(rot_, rot, atol=0.1)
+        angles = so3_relative_angle(rot, rot_, cos_bound=1e-4)
+        self.assertClose(angles, torch.zeros_like(angles), atol=0.1)
+
+    def test_so3_cos_relative_angle(self, batch_size: int = 100):
+        """
+        Check that `so3_relative_angle(R1, R2, cos_angle=False).cos()`
+        is the same as `so3_relative_angle(R1, R2, cos_angle=True)` for
+        batches of randomly generated rotation matrices `R1` and `R2`.
+        """
+        rot1 = TestSO3.init_rot(batch_size=batch_size)
+        rot2 = TestSO3.init_rot(batch_size=batch_size)
+        angles = so3_relative_angle(rot1, rot2, cos_angle=False).cos()
+        angles_ = so3_relative_angle(rot1, rot2, cos_angle=True)
+        self.assertClose(angles, angles_, atol=1e-4)
+
+    def test_so3_cos_angle(self, batch_size: int = 100):
+        """
+        Check that `so3_rotation_angle(R, cos_angle=False).cos()`
+        is the same as `so3_rotation_angle(R, cos_angle=True)` for
+        a batch of randomly generated rotation matrices `R`.
+        """
+        rot = TestSO3.init_rot(batch_size=batch_size)
+        angles = so3_rotation_angle(rot, cos_angle=False).cos()
+        angles_ = so3_rotation_angle(rot, cos_angle=True)
+        self.assertClose(angles, angles_, atol=1e-4)
+
+    def test_so3_cos_bound(self, batch_size: int = 100):
+        """
+        Checks that for an identity rotation `R=I`, the so3_rotation_angle returns
+        non-finite gradients when `cos_bound=None` and finite gradients
+        for `cos_bound > 0.0`.
+        """
+        # generate random rotations with a tiny angle to generate cases
+        # with the gradient singularity
+        device = torch.device("cuda:0")
+        identity = torch.eye(3, device=device)
+        rot180 = identity * torch.tensor([[1.0, -1.0, -1.0]], device=device)
+        r = [identity, rot180]
+        r.extend(
+            [
+                torch.linalg.qr(identity + torch.randn_like(identity) * 1e-4)[0]
+                for _ in range(batch_size - 2)
+            ]
+        )
+        r = torch.stack(r)
+        r.requires_grad = True
+        for is_grad_finite in (True, False):
+            # clear the gradients and decide the cos_bound:
+            #     for is_grad_finite we run so3_rotation_angle with cos_bound
+            #     set to a small float, otherwise we set to 0.0
+            r.grad = None
+            cos_bound = 1e-4 if is_grad_finite else 0.0
+            # compute the angles of r
+            angles = so3_rotation_angle(r, cos_bound=cos_bound)
+            # tests whether all outputs are finite in both cases
+            self.assertTrue(torch.isfinite(angles).all())
+            # compute the gradients
+            loss = angles.sum()
+            loss.backward()
+            # tests whether the gradient is not None for both cases
+            self.assertIsNotNone(r.grad)
+            if is_grad_finite:
+                # all grad values have to be finite
+                self.assertTrue(torch.isfinite(r.grad).all())
+
+    @unittest.skipIf(LooseVersion(torch.__version__) < "1.9", "recent torchscript only")
+    def test_scriptable(self):
+        torch.jit.script(so3_exp_map)
+        torch.jit.script(so3_log_map)
+
+    @staticmethod
+    def so3_expmap(batch_size: int = 10):
+        log_rot = TestSO3.init_log_rot(batch_size=batch_size)
+        torch.cuda.synchronize()
+
+        def compute_rots():
+            so3_exp_map(log_rot)
+            torch.cuda.synchronize()
+
+        return compute_rots
+
+    @staticmethod
+    def so3_logmap(batch_size: int = 10):
+        log_rot = TestSO3.init_rot(batch_size=batch_size)
+        torch.cuda.synchronize()
+
+        def compute_logs():
+            so3_log_map(log_rot)
+            torch.cuda.synchronize()
+
+        return compute_logs

third_party/AnyBimanual/third_party/pytorch3d/tests/test_splatter_blend.py

@@ -0,0 +1,627 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from pytorch3d.common.compat import meshgrid_ij
+from pytorch3d.renderer.cameras import FoVPerspectiveCameras
+from pytorch3d.renderer.splatter_blend import (
+    _compute_occlusion_layers,
+    _compute_splatted_colors_and_weights,
+    _compute_splatting_colors_and_weights,
+    _get_splat_kernel_normalization,
+    _normalize_and_compose_all_layers,
+    _offset_splats,
+    _precompute,
+    _prepare_pixels_and_colors,
+)
+
+from .common_testing import TestCaseMixin
+
+offsets = torch.tensor(
+    [
+        [-1, -1],
+        [-1, 0],
+        [-1, 1],
+        [0, -1],
+        [0, 0],
+        [0, 1],
+        [1, -1],
+        [1, 0],
+        [1, 1],
+    ],
+    device=torch.device("cpu"),
+)
+
+
+def compute_splatting_colors_and_weights_naive(pixel_coords_screen, colors, sigma):
+    normalizer = float(_get_splat_kernel_normalization(offsets))
+    N, H, W, K, _ = colors.shape
+    splat_weights_and_colors = torch.zeros((N, H, W, K, 9, 5))
+    for n in range(N):
+        for h in range(H):
+            for w in range(W):
+                for k in range(K):
+                    q_xy = pixel_coords_screen[n, h, w, k]
+                    q_to_px_center = torch.floor(q_xy) - q_xy + 0.5
+                    color = colors[n, h, w, k]
+                    alpha = colors[n, h, w, k, 3:4]
+                    for d in range(9):
+                        dist_p_q = torch.sum((q_to_px_center + offsets[d]) ** 2)
+                        splat_weight = (
+                            alpha * torch.exp(-dist_p_q / (2 * sigma**2)) * normalizer
+                        )
+                        splat_color = splat_weight * color
+                        splat_weights_and_colors[n, h, w, k, d, :4] = splat_color
+                        splat_weights_and_colors[n, h, w, k, d, 4:5] = splat_weight
+    return splat_weights_and_colors
+
+
+class TestPrecompute(TestCaseMixin, unittest.TestCase):
+    def setUp(self):
+        self.results_cpu = _precompute((2, 3, 4, 5), torch.device("cpu"))
+        self.results1_cpu = _precompute((1, 1, 1, 1), torch.device("cpu"))
+
+    def test_offsets(self):
+        self.assertClose(self.results_cpu[2].shape, offsets.shape, atol=0)
+        self.assertClose(self.results_cpu[2], offsets, atol=0)
+
+        # Offsets should be independent of input_size.
+        self.assertClose(self.results_cpu[2], self.results1_cpu[2], atol=0)
+
+    def test_crops_h(self):
+        target_crops_h1 = torch.tensor(
+            [
+                # chennels being offset:
+                # R  G  B  A  W(eight)
+                [0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1],
+                [2, 2, 2, 2, 2],
+                [0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1],
+                [2, 2, 2, 2, 2],
+                [0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1],
+                [2, 2, 2, 2, 2],
+            ]
+            * 3,  # 3 because we're aiming at (N, H, W+2, K, 9, 5) with W=1.
+            device=torch.device("cpu"),
+        ).reshape(1, 1, 3, 1, 9, 5)
+        self.assertClose(self.results1_cpu[0], target_crops_h1, atol=0)
+
+        target_crops_h_base = target_crops_h1[0, 0, 0]
+        target_crops_h = torch.cat(
+            [target_crops_h_base, target_crops_h_base + 1, target_crops_h_base + 2],
+            dim=0,
+        )
+
+        # Check that we have the right shape, and (after broadcasting) it has the right
+        # values. These should be repeated (tiled) for each n and k.
+        self.assertClose(
+            self.results_cpu[0].shape, torch.tensor([2, 3, 6, 5, 9, 5]), atol=0
+        )
+        for n in range(2):
+            for w in range(6):
+                for k in range(5):
+                    self.assertClose(
+                        self.results_cpu[0][n, :, w, k],
+                        target_crops_h,
+                    )
+
+    def test_crops_w(self):
+        target_crops_w1 = torch.tensor(
+            [
+                # chennels being offset:
+                # R  G  B  A  W(eight)
+                [0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1],
+                [1, 1, 1, 1, 1],
+                [1, 1, 1, 1, 1],
+                [2, 2, 2, 2, 2],
+                [2, 2, 2, 2, 2],
+                [2, 2, 2, 2, 2],
+            ],
+            device=torch.device("cpu"),
+        ).reshape(1, 1, 1, 1, 9, 5)
+        self.assertClose(self.results1_cpu[1], target_crops_w1)
+
+        target_crops_w_base = target_crops_w1[0, 0, 0]
+        target_crops_w = torch.cat(
+            [
+                target_crops_w_base,
+                target_crops_w_base + 1,
+                target_crops_w_base + 2,
+                target_crops_w_base + 3,
+            ],
+            dim=0,
+        )  # Each w value needs an increment.
+
+        # Check that we have the right shape, and (after broadcasting) it has the right
+        # values. These should be repeated (tiled) for each n and k.
+        self.assertClose(self.results_cpu[1].shape, torch.tensor([2, 3, 4, 5, 9, 5]))
+        for n in range(2):
+            for h in range(3):
+                for k in range(5):
+                    self.assertClose(
+                        self.results_cpu[1][n, h, :, k],
+                        target_crops_w,
+                        atol=0,
+                    )
+
+
+class TestPreparPixelsAndColors(TestCaseMixin, unittest.TestCase):
+    def setUp(self):
+        self.device = torch.device("cpu")
+        N, H, W, K = 2, 3, 4, 5
+        self.pixel_coords_cameras = torch.randn(
+            (N, H, W, K, 3), device=self.device, requires_grad=True
+        )
+        self.colors_before = torch.rand((N, H, W, K, 3), device=self.device)
+        self.cameras = FoVPerspectiveCameras(device=self.device)
+        self.background_mask = torch.rand((N, H, W, K), device=self.device) < 0.5
+        self.pixel_coords_screen, self.colors_after = _prepare_pixels_and_colors(
+            self.pixel_coords_cameras,
+            self.colors_before,
+            self.cameras,
+            self.background_mask,
+        )
+
+    def test_background_z(self):
+        self.assertTrue(
+            torch.all(self.pixel_coords_screen[..., 2][self.background_mask] == 1.0)
+        )
+
+    def test_background_alpha(self):
+        self.assertTrue(
+            torch.all(self.colors_after[..., 3][self.background_mask] == 0.0)
+        )
+
+
+class TestGetSplatKernelNormalization(TestCaseMixin, unittest.TestCase):
+    def test_splat_kernel_normalization(self):
+        self.assertAlmostEqual(
+            float(_get_splat_kernel_normalization(offsets)), 0.6503, places=3
+        )
+        self.assertAlmostEqual(
+            float(_get_splat_kernel_normalization(offsets, 0.01)), 1.05, places=3
+        )
+        with self.assertRaisesRegex(ValueError, "Only positive standard deviations"):
+            _get_splat_kernel_normalization(offsets, 0)
+
+
+class TestComputeOcclusionLayers(TestCaseMixin, unittest.TestCase):
+    def test_single_layer(self):
+        # If there's only one layer, all splats must be on the surface level.
+        N, H, W, K = 2, 3, 4, 1
+        q_depth = torch.rand(N, H, W, K)
+        occlusion_layers = _compute_occlusion_layers(q_depth)
+        self.assertClose(occlusion_layers, torch.zeros(N, H, W, 9).long(), atol=0.0)
+
+    def test_all_equal(self):
+        # If all q-vals are equal, then all splats must be on the surface level.
+        N, H, W, K = 2, 3, 4, 5
+        q_depth = torch.ones((N, H, W, K)) * 0.1234
+        occlusion_layers = _compute_occlusion_layers(q_depth)
+        self.assertClose(occlusion_layers, torch.zeros(N, H, W, 9).long(), atol=0.0)
+
+    def test_mid_to_top_level_splatting(self):
+        # Check that occlusion buffers get accumulated as expected when the splatting
+        # and splatted pixels are co-surface on different intersection layers.
+        # This test will make best sense with accompanying Fig. 4 from "Differentiable
+        # Surface Rendering via Non-differentiable Sampling" by Cole et al.
+        for direction, offset in enumerate(offsets):
+            if direction == 4:
+                continue  # Skip self-splatting which is always co-surface.
+
+            depths = torch.zeros(1, 3, 3, 3)
+
+            # This is our q, the pixel splatted onto, in the center of the image.
+            depths[0, 1, 1] = torch.tensor([0.71, 0.8, 1.0])
+
+            # This is our p, the splatting pixel.
+            depths[0, offset[0] + 1, offset[1] + 1] = torch.tensor([0.5, 0.7, 0.9])
+
+            occlusion_layers = _compute_occlusion_layers(depths)
+
+            # Check that we computed that it is the middle layer of p that is co-
+            # surface with q. (1, 1) is the id of q in the depth array, and offset_id
+            # is the id of p's direction w.r.t. q.
+            psurfaceid_onto_q = occlusion_layers[0, 1, 1, direction]
+            self.assertEqual(int(psurfaceid_onto_q), 1)
+
+            # Conversely, if we swap p and q, we have a top-level splatting onto
+            # mid-level. offset + 1 is the id of p, and 8-offset_id is the id of
+            # q's direction w.r.t. p (e.g. if p is [-1, -1] w.r.t. q, then q is
+            # [1, 1] w.r.t. p; we use the ids of these two directions in the offsets
+            # array).
+            qsurfaceid_onto_p = occlusion_layers[
+                0, offset[0] + 1, offset[1] + 1, 8 - direction
+            ]
+            self.assertEqual(int(qsurfaceid_onto_p), -1)
+
+
+class TestComputeSplattingColorsAndWeights(TestCaseMixin, unittest.TestCase):
+    def setUp(self):
+        self.N, self.H, self.W, self.K = 2, 3, 4, 5
+        self.pixel_coords_screen = (
+            torch.stack(
+                meshgrid_ij(torch.arange(self.H), torch.arange(self.W)),
+                dim=-1,
+            )
+            .reshape(1, self.H, self.W, 1, 2)
+            .expand(self.N, self.H, self.W, self.K, 2)
+            .float()
+            + 0.5
+        )
+        self.colors = torch.ones((self.N, self.H, self.W, self.K, 4))
+
+    def test_all_equal(self):
+        # If all colors are equal and on a regular grid, all weights and reweighted
+        # colors should be equal given a specific splatting direction.
+        splatting_colors_and_weights = _compute_splatting_colors_and_weights(
+            self.pixel_coords_screen, self.colors * 0.2345, sigma=0.5, offsets=offsets
+        )
+
+        # Splatting directly to the top/bottom/left/right should have the same strenght.
+        non_diag_splats = splatting_colors_and_weights[
+            :, :, :, :, torch.tensor([1, 3, 5, 7])
+        ]
+
+        # Same for diagonal splats.
+        diag_splats = splatting_colors_and_weights[
+            :, :, :, :, torch.tensor([0, 2, 6, 8])
+        ]
+
+        # And for self-splats.
+        self_splats = splatting_colors_and_weights[:, :, :, :, torch.tensor([4])]
+
+        for splats in non_diag_splats, diag_splats, self_splats:
+            # Colors should be equal.
+            self.assertTrue(torch.all(splats[..., :4] == splats[0, 0, 0, 0, 0, 0]))
+
+            # Weights should be equal.
+            self.assertTrue(torch.all(splats[..., 4] == splats[0, 0, 0, 0, 0, 4]))
+
+        # Non-diagonal weights should be greater than diagonal weights.
+        self.assertGreater(
+            non_diag_splats[0, 0, 0, 0, 0, 0], diag_splats[0, 0, 0, 0, 0, 0]
+        )
+
+        # Self-splats should be strongest of all.
+        self.assertGreater(
+            self_splats[0, 0, 0, 0, 0, 0], non_diag_splats[0, 0, 0, 0, 0, 0]
+        )
+
+        # Splatting colors should be reweighted proportionally to their splat weights.
+        diag_self_color_ratio = (
+            diag_splats[0, 0, 0, 0, 0, 0] / self_splats[0, 0, 0, 0, 0, 0]
+        )
+        diag_self_weight_ratio = (
+            diag_splats[0, 0, 0, 0, 0, 4] / self_splats[0, 0, 0, 0, 0, 4]
+        )
+        self.assertEqual(diag_self_color_ratio, diag_self_weight_ratio)
+
+        non_diag_self_color_ratio = (
+            non_diag_splats[0, 0, 0, 0, 0, 0] / self_splats[0, 0, 0, 0, 0, 0]
+        )
+        non_diag_self_weight_ratio = (
+            non_diag_splats[0, 0, 0, 0, 0, 4] / self_splats[0, 0, 0, 0, 0, 4]
+        )
+        self.assertEqual(non_diag_self_color_ratio, non_diag_self_weight_ratio)
+
+    def test_zero_alpha_zero_weight(self):
+        # Pixels with zero alpha do no splatting, but should still be splatted on.
+        colors = self.colors.clone()
+        colors[0, 1, 1, 0, 3] = 0.0
+        splatting_colors_and_weights = _compute_splatting_colors_and_weights(
+            self.pixel_coords_screen, colors, sigma=0.5, offsets=offsets
+        )
+
+        # The transparent pixel should do no splatting.
+        self.assertTrue(torch.all(splatting_colors_and_weights[0, 1, 1, 0] == 0.0))
+
+        # Splatting *onto* the transparent pixel should be unaffected.
+        reference_weights_colors = splatting_colors_and_weights[0, 1, 1, 1]
+        for direction, offset in enumerate(offsets):
+            if direction == 4:
+                continue  # Ignore self-splats
+            # We invert the direction to get the right (h, w, d) coordinate of each
+            # pixel splatting *onto* the pixel with zero alpha.
+            self.assertClose(
+                splatting_colors_and_weights[
+                    0, 1 + offset[0], 1 + offset[1], 0, 8 - direction
+                ],
+                reference_weights_colors[8 - direction],
+                atol=0.001,
+            )
+
+    def test_random_inputs(self):
+        pixel_coords_screen = (
+            self.pixel_coords_screen
+            + torch.randn((self.N, self.H, self.W, self.K, 2)) * 0.1
+        )
+        colors = torch.rand((self.N, self.H, self.W, self.K, 4))
+        splatting_colors_and_weights = _compute_splatting_colors_and_weights(
+            pixel_coords_screen, colors, sigma=0.5, offsets=offsets
+        )
+        naive_colors_and_weights = compute_splatting_colors_and_weights_naive(
+            pixel_coords_screen, colors, sigma=0.5
+        )
+
+        self.assertClose(
+            splatting_colors_and_weights, naive_colors_and_weights, atol=0.01
+        )
+
+
+class TestOffsetSplats(TestCaseMixin, unittest.TestCase):
+    def test_offset(self):
+        device = torch.device("cuda:0")
+        N, H, W, K = 2, 3, 4, 5
+        colors_and_weights = torch.rand((N, H, W, K, 9, 5), device=device)
+        crop_ids_h, crop_ids_w, _ = _precompute((N, H, W, K), device=device)
+        offset_colors_and_weights = _offset_splats(
+            colors_and_weights, crop_ids_h, crop_ids_w
+        )
+
+        # Check each splatting direction individually, for clarity.
+        # offset_x, offset_y = (-1, -1)
+        direction = 0
+        self.assertClose(
+            offset_colors_and_weights[:, 1:, 1:, :, direction],
+            colors_and_weights[:, :-1, :-1, :, direction],
+            atol=0.001,
+        )
+        self.assertTrue(
+            torch.all(offset_colors_and_weights[:, 0, :, :, direction] == 0.0)
+        )
+        self.assertTrue(
+            torch.all(offset_colors_and_weights[:, :, 0, :, direction] == 0.0)
+        )
+
+        # offset_x, offset_y = (-1, 0)
+        direction = 1
+        self.assertClose(
+            offset_colors_and_weights[:, :, 1:, :, direction],
+            colors_and_weights[:, :, :-1, :, direction],
+            atol=0.001,
+        )
+        self.assertTrue(
+            torch.all(offset_colors_and_weights[:, :, 0, :, direction] == 0.0)
+        )
+
+        # offset_x, offset_y = (-1, 1)
+        direction = 2
+        self.assertClose(
+            offset_colors_and_weights[:, :-1, 1:, :, direction],
+            colors_and_weights[:, 1:, :-1, :, direction],
+            atol=0.001,
+        )
+        self.assertTrue(
+            torch.all(offset_colors_and_weights[:, -1, :, :, direction] == 0.0)
+        )
+        self.assertTrue(
+            torch.all(offset_colors_and_weights[:, :, 0, :, direction] == 0.0)
+        )
+
+        # offset_x, offset_y = (0, -1)
+        direction = 3
+        self.assertClose(
+            offset_colors_and_weights[:, 1:, :, :, direction],
+            colors_and_weights[:, :-1, :, :, direction],
+            atol=0.001,
+        )
+        self.assertTrue(
+            torch.all(offset_colors_and_weights[:, 0, :, :, direction] == 0.0)
+        )
+
+        # self-splat
+        direction = 4
+        self.assertClose(
+            offset_colors_and_weights[..., direction, :],
+            colors_and_weights[..., direction, :],
+            atol=0.001,
+        )
+
+        # offset_x, offset_y = (0, 1)
+        direction = 5
+        self.assertClose(
+            offset_colors_and_weights[:, :-1, :, :, direction],
+            colors_and_weights[:, 1:, :, :, direction],
+            atol=0.001,
+        )
+        self.assertTrue(
+            torch.all(offset_colors_and_weights[:, -1, :, :, direction] == 0.0)
+        )
+
+        # offset_x, offset_y = (1, -1)
+        direction = 6
+        self.assertClose(
+            offset_colors_and_weights[:, 1:, :-1, :, direction],
+            colors_and_weights[:, :-1, 1:, :, direction],
+            atol=0.001,
+        )
+        self.assertTrue(
+            torch.all(offset_colors_and_weights[:, 0, :, :, direction] == 0.0)
+        )
+        self.assertTrue(
+            torch.all(offset_colors_and_weights[:, :, -1, :, direction] == 0.0)
+        )
+
+        # offset_x, offset_y = (1, 0)
+        direction = 7
+        self.assertClose(
+            offset_colors_and_weights[:, :, :-1, :, direction],
+            colors_and_weights[:, :, 1:, :, direction],
+            atol=0.001,
+        )
+        self.assertTrue(
+            torch.all(offset_colors_and_weights[:, :, -1, :, direction] == 0.0)
+        )
+
+        # offset_x, offset_y = (1, 1)
+        direction = 8
+        self.assertClose(
+            offset_colors_and_weights[:, :-1, :-1, :, direction],
+            colors_and_weights[:, 1:, 1:, :, direction],
+            atol=0.001,
+        )
+        self.assertTrue(
+            torch.all(offset_colors_and_weights[:, -1, :, :, direction] == 0.0)
+        )
+        self.assertTrue(
+            torch.all(offset_colors_and_weights[:, :, -1, :, direction] == 0.0)
+        )
+
+
+class TestComputeSplattedColorsAndWeights(TestCaseMixin, unittest.TestCase):
+    def test_accumulation_background(self):
+        # Set occlusion_layers to all -1, so all splats are background splats.
+        splat_colors_and_weights = torch.rand((1, 1, 1, 3, 9, 5))
+        occlusion_layers = torch.zeros((1, 1, 1, 9)) - 1
+        splatted_colors, splatted_weights = _compute_splatted_colors_and_weights(
+            occlusion_layers, splat_colors_and_weights
+        )
+
+        # Foreground splats (there are none).
+        self.assertClose(
+            splatted_colors[0, 0, 0, :, 0],
+            torch.zeros((4)),
+            atol=0.001,
+        )
+
+        # Surface splats (there are none).
+        self.assertClose(
+            splatted_colors[0, 0, 0, :, 1],
+            torch.zeros((4)),
+            atol=0.001,
+        )
+
+        # Background splats.
+        self.assertClose(
+            splatted_colors[0, 0, 0, :, 2],
+            splat_colors_and_weights[0, 0, 0, :, :, :4].sum(dim=0).sum(dim=0),
+            atol=0.001,
+        )
+
+    def test_accumulation_middle(self):
+        # Set occlusion_layers to all 0, so top splats are co-surface with splatted
+        # pixels. Thus, the top splatting layer should be accumulated to surface, and
+        # all other layers to background.
+        splat_colors_and_weights = torch.rand((1, 1, 1, 3, 9, 5))
+        occlusion_layers = torch.zeros((1, 1, 1, 9))
+        splatted_colors, splatted_weights = _compute_splatted_colors_and_weights(
+            occlusion_layers, splat_colors_and_weights
+        )
+
+        # Foreground splats (there are none).
+        self.assertClose(
+            splatted_colors[0, 0, 0, :, 0],
+            torch.zeros((4)),
+            atol=0.001,
+        )
+
+        # Surface splats
+        self.assertClose(
+            splatted_colors[0, 0, 0, :, 1],
+            splat_colors_and_weights[0, 0, 0, 0, :, :4].sum(dim=0),
+            atol=0.001,
+        )
+
+        # Background splats
+        self.assertClose(
+            splatted_colors[0, 0, 0, :, 2],
+            splat_colors_and_weights[0, 0, 0, 1:, :, :4].sum(dim=0).sum(dim=0),
+            atol=0.001,
+        )
+
+    def test_accumulation_foreground(self):
+        # Set occlusion_layers to all 1. Then the top splatter is a foreground
+        # splatter, mid splatter is surface, and bottom splatter is background.
+        splat_colors_and_weights = torch.rand((1, 1, 1, 3, 9, 5))
+        occlusion_layers = torch.zeros((1, 1, 1, 9)) + 1
+        splatted_colors, splatted_weights = _compute_splatted_colors_and_weights(
+            occlusion_layers, splat_colors_and_weights
+        )
+
+        # Foreground splats
+        self.assertClose(
+            splatted_colors[0, 0, 0, :, 0],
+            splat_colors_and_weights[0, 0, 0, 0:1, :, :4].sum(dim=0).sum(dim=0),
+            atol=0.001,
+        )
+
+        # Surface splats
+        self.assertClose(
+            splatted_colors[0, 0, 0, :, 1],
+            splat_colors_and_weights[0, 0, 0, 1:2, :, :4].sum(dim=0).sum(dim=0),
+            atol=0.001,
+        )
+
+        # Background splats
+        self.assertClose(
+            splatted_colors[0, 0, 0, :, 2],
+            splat_colors_and_weights[0, 0, 0, 2:3, :, :4].sum(dim=0).sum(dim=0),
+            atol=0.001,
+        )
+
+
+class TestNormalizeAndComposeAllLayers(TestCaseMixin, unittest.TestCase):
+    def test_background_color(self):
+        # Background should always have alpha=0, and the chosen RGB.
+        N, H, W = 2, 3, 4
+        # Make a mask with background in the zeroth row of the first image.
+        bg_mask = torch.zeros([N, H, W, 1, 1])
+        bg_mask[0, :, 0] = 1
+
+        bg_color = torch.tensor([0.2, 0.3, 0.4])
+
+        color_layers = torch.rand((N, H, W, 4, 3)) * (1 - bg_mask)
+        color_weights = torch.rand((N, H, W, 1, 3)) * (1 - bg_mask)
+
+        colors = _normalize_and_compose_all_layers(
+            bg_color, color_layers, color_weights
+        )
+
+        # Background RGB should be .2, .3, .4, and alpha should be 0.
+        self.assertClose(
+            torch.masked_select(colors, bg_mask.bool()[..., 0]),
+            torch.tensor([0.2, 0.3, 0.4, 0, 0.2, 0.3, 0.4, 0, 0.2, 0.3, 0.4, 0.0]),
+            atol=0.001,
+        )
+
+    def test_compositing_opaque(self):
+        # When all colors are opaque, only the foreground layer should be visible.
+        N, H, W = 2, 3, 4
+        color_layers = torch.rand((N, H, W, 4, 3))
+        color_layers[..., 3, :] = 1.0
+        color_weights = torch.ones((N, H, W, 1, 3))
+
+        out_colors = _normalize_and_compose_all_layers(
+            torch.tensor([0.0, 0.0, 0.0]), color_layers, color_weights
+        )
+        self.assertClose(out_colors, color_layers[..., 0], atol=0.001)
+
+    def test_compositing_transparencies(self):
+        # When foreground layer is transparent and surface and bg are semi-transparent,
+        # we should return a  mix of the two latter.
+        N, H, W = 2, 3, 4
+        color_layers = torch.rand((N, H, W, 4, 3))
+        color_layers[..., 3, 0] = 0.1  # fg
+        color_layers[..., 3, 1] = 0.2  # surface
+        color_layers[..., 3, 2] = 0.3  # bg
+        color_weights = torch.ones((N, H, W, 1, 3))
+
+        out_colors = _normalize_and_compose_all_layers(
+            torch.tensor([0.0, 0.0, 0.0]), color_layers, color_weights
+        )
+        self.assertClose(
+            out_colors,
+            color_layers[..., 0]
+            + 0.9 * (color_layers[..., 1] + 0.8 * color_layers[..., 2]),
+        )

third_party/AnyBimanual/third_party/pytorch3d/tests/test_struct_utils.py

@@ -0,0 +1,227 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import unittest
+
+import torch
+from pytorch3d.structures import utils as struct_utils
+
+from .common_testing import TestCaseMixin
+
+
+class TestStructUtils(TestCaseMixin, unittest.TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        torch.manual_seed(43)
+
+    def _check_list_to_padded_slices(self, x, x_padded, ndim):
+        N = len(x)
+        for i in range(N):
+            slices = [i]
+            for dim in range(ndim):
+                if x[i].nelement() == 0 and x[i].ndim == 1:
+                    slice_ = slice(0, 0, 1)
+                else:
+                    slice_ = slice(0, x[i].shape[dim], 1)
+                slices.append(slice_)
+            if x[i].nelement() == 0 and x[i].ndim == 1:
+                x_correct = x[i].new_zeros(*[[0] * ndim])
+            else:
+                x_correct = x[i]
+            self.assertClose(x_padded[slices], x_correct)
+
+    def test_list_to_padded(self):
+        device = torch.device("cuda:0")
+        N = 5
+        K = 20
+        for ndim in [1, 2, 3, 4]:
+            x = []
+            for _ in range(N):
+                dims = torch.randint(K, size=(ndim,)).tolist()
+                x.append(torch.rand(dims, device=device))
+
+            # set 0th element to an empty 1D tensor
+            x[0] = torch.tensor([], dtype=x[0].dtype, device=device)
+
+            # set 1st element to an empty tensor with correct number of dims
+            x[1] = x[1].new_zeros(*[[0] * ndim])
+
+            pad_size = [K] * ndim
+            x_padded = struct_utils.list_to_padded(
+                x, pad_size=pad_size, pad_value=0.0, equisized=False
+            )
+
+            for dim in range(ndim):
+                self.assertEqual(x_padded.shape[dim + 1], K)
+
+            self._check_list_to_padded_slices(x, x_padded, ndim)
+
+            # check for no pad size (defaults to max dimension)
+            x_padded = struct_utils.list_to_padded(x, pad_value=0.0, equisized=False)
+            max_sizes = (
+                max(
+                    (0 if (y.nelement() == 0 and y.ndim == 1) else y.shape[dim])
+                    for y in x
+                )
+                for dim in range(ndim)
+            )
+            for dim, max_size in enumerate(max_sizes):
+                self.assertEqual(x_padded.shape[dim + 1], max_size)
+
+            self._check_list_to_padded_slices(x, x_padded, ndim)
+
+            # check for equisized
+            x = [torch.rand((K, *([10] * (ndim - 1))), device=device) for _ in range(N)]
+            x_padded = struct_utils.list_to_padded(x, equisized=True)
+            self.assertClose(x_padded, torch.stack(x, 0))
+
+        # catch ValueError for invalid dimensions
+        pad_size = [K] * (ndim + 1)
+        with self.assertRaisesRegex(ValueError, "Pad size must"):
+            struct_utils.list_to_padded(
+                x, pad_size=pad_size, pad_value=0.0, equisized=False
+            )
+
+        # invalid input tensor dimensions
+        x = []
+        ndim = 3
+        for _ in range(N):
+            dims = torch.randint(K, size=(ndim,)).tolist()
+            x.append(torch.rand(dims, device=device))
+        pad_size = [K] * 2
+        with self.assertRaisesRegex(ValueError, "Pad size must"):
+            x_padded = struct_utils.list_to_padded(
+                x, pad_size=pad_size, pad_value=0.0, equisized=False
+            )
+
+    def test_padded_to_list(self):
+        device = torch.device("cuda:0")
+        N = 5
+        K = 20
+        ndim = 2
+
+        for ndim in (2, 3, 4):
+
+            dims = [K] * ndim
+            x = torch.rand([N] + dims, device=device)
+
+            x_list = struct_utils.padded_to_list(x)
+            for i in range(N):
+                self.assertClose(x_list[i], x[i])
+
+            split_size = torch.randint(1, K, size=(N, ndim)).unbind(0)
+            x_list = struct_utils.padded_to_list(x, split_size)
+            for i in range(N):
+                slices = [i]
+                for dim in range(ndim):
+                    slices.append(slice(0, split_size[i][dim], 1))
+                self.assertClose(x_list[i], x[slices])
+
+            # split size is a list of ints
+            split_size = [int(z) for z in torch.randint(1, K, size=(N,)).unbind(0)]
+            x_list = struct_utils.padded_to_list(x, split_size)
+            for i in range(N):
+                self.assertClose(x_list[i], x[i][: split_size[i]])
+
+    def test_padded_to_packed(self):
+        device = torch.device("cuda:0")
+        N = 5
+        K = 20
+        ndim = 2
+        dims = [K] * ndim
+        x = torch.rand([N] + dims, device=device)
+
+        # Case 1: no split_size or pad_value provided
+        # Check output is just the flattened input.
+        x_packed = struct_utils.padded_to_packed(x)
+        self.assertTrue(x_packed.shape == (x.shape[0] * x.shape[1], x.shape[2]))
+        self.assertClose(x_packed, x.reshape(-1, K))
+
+        # Case 2: pad_value is provided.
+        # Check each section of the packed tensor matches the
+        # corresponding unpadded elements of the padded tensor.
+        # Check that only rows where all the values are padded
+        # are removed in the conversion to packed.
+        pad_value = -1
+        x_list = []
+        split_size = []
+        for _ in range(N):
+            dim = torch.randint(K, size=(1,)).item()
+            # Add some random values in the input which are the same as the pad_value.
+            # These should not be filtered out.
+            x_list.append(
+                torch.randint(low=pad_value, high=10, size=(dim, K), device=device)
+            )
+            split_size.append(dim)
+        x_padded = struct_utils.list_to_padded(x_list, pad_value=pad_value)
+        x_packed = struct_utils.padded_to_packed(x_padded, pad_value=pad_value)
+        curr = 0
+        for i in range(N):
+            self.assertClose(x_packed[curr : curr + split_size[i], ...], x_list[i])
+            self.assertClose(torch.cat(x_list), x_packed)
+            curr += split_size[i]
+
+        # Case 3: split_size is provided.
+        # Check each section of the packed tensor matches the corresponding
+        # unpadded elements.
+        x_packed = struct_utils.padded_to_packed(x_padded, split_size=split_size)
+        curr = 0
+        for i in range(N):
+            self.assertClose(x_packed[curr : curr + split_size[i], ...], x_list[i])
+            self.assertClose(torch.cat(x_list), x_packed)
+            curr += split_size[i]
+
+        # Case 4: split_size of the wrong shape is provided.
+        # Raise an error.
+        split_size = torch.randint(1, K, size=(2 * N,)).view(N, 2).unbind(0)
+        with self.assertRaisesRegex(ValueError, "1-dimensional"):
+            x_packed = struct_utils.padded_to_packed(x_padded, split_size=split_size)
+
+        split_size = torch.randint(1, K, size=(2 * N,)).view(N * 2).tolist()
+        with self.assertRaisesRegex(
+            ValueError, "same length as inputs first dimension"
+        ):
+            x_packed = struct_utils.padded_to_packed(x_padded, split_size=split_size)
+
+        # Case 5: both pad_value and split_size are provided.
+        # Raise an error.
+        with self.assertRaisesRegex(ValueError, "Only one of"):
+            x_packed = struct_utils.padded_to_packed(
+                x_padded, split_size=split_size, pad_value=-1
+            )
+
+        # Case 6: Input has more than 3 dims.
+        # Raise an error.
+        x = torch.rand((N, K, K, K, K), device=device)
+        split_size = torch.randint(1, K, size=(N,)).tolist()
+        with self.assertRaisesRegex(ValueError, "Supports only"):
+            struct_utils.padded_to_packed(x, split_size=split_size)
+
+    def test_list_to_packed(self):
+        device = torch.device("cuda:0")
+        N = 5
+        K = 20
+        x, x_dims = [], []
+        dim2 = torch.randint(K, size=(1,)).item()
+        for _ in range(N):
+            dim1 = torch.randint(K, size=(1,)).item()
+            x_dims.append(dim1)
+            x.append(torch.rand([dim1, dim2], device=device))
+
+        out = struct_utils.list_to_packed(x)
+        x_packed = out[0]
+        num_items = out[1]
+        item_packed_first_idx = out[2]
+        item_packed_to_list_idx = out[3]
+
+        cur = 0
+        for i in range(N):
+            self.assertTrue(num_items[i] == x_dims[i])
+            self.assertTrue(item_packed_first_idx[i] == cur)
+            self.assertTrue(item_packed_to_list_idx[cur : cur + x_dims[i]].eq(i).all())
+            self.assertClose(x_packed[cur : cur + x_dims[i]], x[i])
+            cur += x_dims[i]

third_party/AnyBimanual/third_party/pytorch3d/tests/test_subdivide_meshes.py

@@ -0,0 +1,234 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import unittest
+
+import torch
+from pytorch3d.ops.subdivide_meshes import SubdivideMeshes
+from pytorch3d.structures.meshes import Meshes
+from pytorch3d.utils.ico_sphere import ico_sphere
+
+from .common_testing import TestCaseMixin
+
+
+class TestSubdivideMeshes(TestCaseMixin, unittest.TestCase):
+    def simple_subdivide(self, with_init=False):
+        # Create a mesh with one face and check the subdivided mesh has
+        # 4 faces with the correct vertex coordinates.
+        device = torch.device("cuda:0")
+        verts = torch.tensor(
+            [[0.5, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
+            dtype=torch.float32,
+            device=device,
+            requires_grad=True,
+        )
+        faces = torch.tensor([[0, 1, 2]], dtype=torch.int64, device=device)
+        mesh = Meshes(verts=[verts], faces=[faces])
+        mesh_init = mesh.clone() if with_init else None
+        subdivide = SubdivideMeshes(meshes=mesh_init)
+        new_mesh = subdivide(mesh)
+
+        # Subdivided face:
+        #
+        #           v0
+        #           /\
+        #          /  \
+        #         / f0 \
+        #     v4 /______\ v3
+        #       /\      /\
+        #      /  \ f3 /  \
+        #     / f2 \  / f1 \
+        #    /______\/______\
+        #  v2       v5       v1
+        #
+        gt_subdivide_verts = torch.tensor(
+            [
+                [0.5, 1.0, 0.0],
+                [1.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0],
+                [0.75, 0.5, 0.0],
+                [0.25, 0.5, 0.0],
+                [0.5, 0.0, 0.0],
+            ],
+            dtype=torch.float32,
+            device=device,
+        )
+        gt_subdivide_faces = torch.tensor(
+            [[0, 3, 4], [1, 5, 3], [2, 4, 5], [5, 4, 3]],
+            dtype=torch.int64,
+            device=device,
+        )
+        new_verts, new_faces = new_mesh.get_mesh_verts_faces(0)
+        self.assertClose(new_verts, gt_subdivide_verts)
+        self.assertClose(new_faces, gt_subdivide_faces)
+        self.assertTrue(new_verts.requires_grad == verts.requires_grad)
+
+    def test_simple_subdivide(self):
+        self.simple_subdivide()
+
+    def test_simple_subdivide_with_init(self):
+        self.simple_subdivide(with_init=True)
+
+    def test_heterogeneous_meshes(self):
+        device = torch.device("cuda:0")
+        verts1 = torch.tensor(
+            [[0.5, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
+            dtype=torch.float32,
+            device=device,
+            requires_grad=True,
+        )
+        faces1 = torch.tensor([[0, 1, 2]], dtype=torch.int64, device=device)
+        verts2 = torch.tensor(
+            [[0.5, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.5, 1.0, 0.0]],
+            dtype=torch.float32,
+            device=device,
+            requires_grad=True,
+        )
+        faces2 = torch.tensor([[0, 1, 2], [0, 3, 1]], dtype=torch.int64, device=device)
+        faces3 = torch.tensor([[0, 1, 2], [0, 2, 3]], dtype=torch.int64, device=device)
+        mesh = Meshes(verts=[verts1, verts2, verts2], faces=[faces1, faces2, faces3])
+        subdivide = SubdivideMeshes()
+        new_mesh = subdivide(mesh.clone())
+
+        gt_subdivided_verts1 = torch.tensor(
+            [
+                [0.5, 1.0, 0.0],
+                [1.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0],
+                [0.75, 0.5, 0.0],
+                [0.25, 0.5, 0.0],
+                [0.5, 0.0, 0.0],
+            ],
+            dtype=torch.float32,
+            device=device,
+        )
+        gt_subdivided_faces1 = torch.tensor(
+            [[0, 3, 4], [1, 5, 3], [2, 4, 5], [5, 4, 3]],
+            dtype=torch.int64,
+            device=device,
+        )
+        # faces2:
+        #
+        #         v0 _______e2_______ v3
+        #           /\              /
+        #          /  \            /
+        #         /    \          /
+        #     e1 /      \ e0     / e4
+        #       /        \      /
+        #      /          \    /
+        #     /            \  /
+        #    /______________\/
+        #  v2       e3      v1
+        #
+        # Subdivided faces2:
+        #
+        #         v0 _______v6_______ v3
+        #           /\      /\      /
+        #          /  \ f1 /  \ f3 /
+        #         / f0 \  / f7 \  /
+        #     v5 /______v4______\/v8
+        #       /\      /\      /
+        #      /  \ f6 /  \ f5 /
+        #     / f4 \  / f2 \  /
+        #    /______\/______\/
+        #  v2       v7       v1
+        #
+        gt_subdivided_verts2 = torch.tensor(
+            [
+                [0.5, 1.0, 0.0],
+                [1.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0],
+                [1.5, 1.0, 0.0],
+                [0.75, 0.5, 0.0],
+                [0.25, 0.5, 0.0],
+                [1.0, 1.0, 0.0],
+                [0.5, 0.0, 0.0],
+                [1.25, 0.5, 0.0],
+            ],
+            dtype=torch.float32,
+            device=device,
+        )
+        gt_subdivided_faces2 = torch.tensor(
+            [
+                [0, 4, 5],
+                [0, 6, 4],
+                [1, 7, 4],
+                [3, 8, 6],
+                [2, 5, 7],
+                [1, 4, 8],
+                [7, 5, 4],
+                [8, 4, 6],
+            ],
+            dtype=torch.int64,
+            device=device,
+        )
+        gt_subdivided_verts3 = gt_subdivided_verts2.clone()
+        gt_subdivided_verts3[-1, :] = torch.tensor(
+            [0.75, 0.5, 0], dtype=torch.float32, device=device
+        )
+        gt_subdivided_faces3 = torch.tensor(
+            [
+                [0, 4, 5],
+                [0, 5, 6],
+                [1, 7, 4],
+                [2, 8, 5],
+                [2, 5, 7],
+                [3, 6, 8],
+                [7, 5, 4],
+                [8, 6, 5],
+            ],
+            dtype=torch.int64,
+            device=device,
+        )
+        new_mesh_verts1, new_mesh_faces1 = new_mesh.get_mesh_verts_faces(0)
+        new_mesh_verts2, new_mesh_faces2 = new_mesh.get_mesh_verts_faces(1)
+        new_mesh_verts3, new_mesh_faces3 = new_mesh.get_mesh_verts_faces(2)
+        self.assertClose(new_mesh_verts1, gt_subdivided_verts1)
+        self.assertClose(new_mesh_faces1, gt_subdivided_faces1)
+        self.assertClose(new_mesh_verts2, gt_subdivided_verts2)
+        self.assertClose(new_mesh_faces2, gt_subdivided_faces2)
+        self.assertClose(new_mesh_verts3, gt_subdivided_verts3)
+        self.assertClose(new_mesh_faces3, gt_subdivided_faces3)
+        self.assertTrue(new_mesh_verts1.requires_grad == verts1.requires_grad)
+        self.assertTrue(new_mesh_verts2.requires_grad == verts2.requires_grad)
+        self.assertTrue(new_mesh_verts3.requires_grad == verts2.requires_grad)
+
+    def test_subdivide_features(self):
+        device = torch.device("cuda:0")
+        mesh = ico_sphere(0, device)
+        N = 10
+        mesh = mesh.extend(N)
+        edges = mesh.edges_packed()
+        V = mesh.num_verts_per_mesh()[0]
+        D = 256
+        feats = torch.rand(
+            (N * V, D), dtype=torch.float32, device=device, requires_grad=True
+        )  # packed features
+        app_feats = feats[edges].mean(1)
+        subdivide = SubdivideMeshes()
+        new_mesh, new_feats = subdivide(mesh, feats)
+        gt_feats = torch.cat(
+            (feats.view(N, V, D), app_feats.view(N, -1, D)), dim=1
+        ).view(-1, D)
+        self.assertClose(new_feats, gt_feats)
+        self.assertTrue(new_feats.requires_grad == gt_feats.requires_grad)
+
+    @staticmethod
+    def subdivide_meshes_with_init(num_meshes: int = 10, same_topo: bool = False):
+        device = torch.device("cuda:0")
+        meshes = ico_sphere(0, device=device)
+        if num_meshes > 1:
+            meshes = meshes.extend(num_meshes)
+        meshes_init = meshes.clone() if same_topo else None
+        torch.cuda.synchronize()
+
+        def subdivide_meshes():
+            subdivide = SubdivideMeshes(meshes=meshes_init)
+            subdivide(meshes=meshes.clone())
+            torch.cuda.synchronize()
+
+        return subdivide_meshes

third_party/AnyBimanual/third_party/pytorch3d/tests/test_symeig3x3.py

@@ -0,0 +1,264 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import unittest
+
+import torch
+from pytorch3d.common.workaround import symeig3x3
+from pytorch3d.transforms.rotation_conversions import random_rotations
+
+from .common_testing import get_random_cuda_device, TestCaseMixin
+
+
+class TestSymEig3x3(TestCaseMixin, unittest.TestCase):
+    TEST_BATCH_SIZE = 1024
+
+    @staticmethod
+    def create_random_sym3x3(device, n):
+        random_3x3 = torch.randn((n, 3, 3), device=device)
+        random_3x3_T = torch.transpose(random_3x3, 1, 2)
+        random_sym_3x3 = (random_3x3 * random_3x3_T).contiguous()
+
+        return random_sym_3x3
+
+    @staticmethod
+    def create_diag_sym3x3(device, n, noise=0.0):
+        # Create purly diagonal matrices
+        random_diag_3x3 = torch.randn((n, 3), device=device).diag_embed()
+
+        # Make them 'almost' diagonal
+        random_diag_3x3 += noise * TestSymEig3x3.create_random_sym3x3(device, n)
+
+        return random_diag_3x3
+
+    def setUp(self) -> None:
+        super().setUp()
+        torch.manual_seed(42)
+
+        self._gpu = get_random_cuda_device()
+        self._cpu = torch.device("cpu")
+
+    def test_is_eigen_gpu(self):
+        test_input = self.create_random_sym3x3(self._gpu, n=self.TEST_BATCH_SIZE)
+
+        self._test_is_eigen(test_input)
+
+    def test_is_eigen_cpu(self):
+        test_input = self.create_random_sym3x3(self._cpu, n=self.TEST_BATCH_SIZE)
+
+        self._test_is_eigen(test_input)
+
+    def _test_is_eigen(self, test_input, atol=1e-04, rtol=1e-02):
+        """
+        Verify that values and vectors produced are really eigenvalues and eigenvectors
+        and can restore the original input matrix with good precision
+        """
+        eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
+
+        self.assertClose(
+            test_input,
+            eigenvectors @ eigenvalues.diag_embed() @ eigenvectors.transpose(-2, -1),
+            atol=atol,
+            rtol=rtol,
+        )
+
+    def test_eigenvectors_are_orthonormal_gpu(self):
+        test_input = self.create_random_sym3x3(self._gpu, n=self.TEST_BATCH_SIZE)
+
+        self._test_eigenvectors_are_orthonormal(test_input)
+
+    def test_eigenvectors_are_orthonormal_cpu(self):
+        test_input = self.create_random_sym3x3(self._cpu, n=self.TEST_BATCH_SIZE)
+
+        self._test_eigenvectors_are_orthonormal(test_input)
+
+    def _test_eigenvectors_are_orthonormal(self, test_input):
+        """
+        Verify that eigenvectors are an orthonormal set
+        """
+        eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
+
+        batched_eye = torch.zeros_like(test_input)
+        batched_eye[..., :, :] = torch.eye(3, device=batched_eye.device)
+
+        self.assertClose(
+            batched_eye, eigenvectors @ eigenvectors.transpose(-2, -1), atol=1e-06
+        )
+
+    def test_is_not_nan_or_inf_gpu(self):
+        test_input = self.create_random_sym3x3(self._gpu, n=self.TEST_BATCH_SIZE)
+
+        self._test_is_not_nan_or_inf(test_input)
+
+    def test_is_not_nan_or_inf_cpu(self):
+        test_input = self.create_random_sym3x3(self._cpu, n=self.TEST_BATCH_SIZE)
+
+        self._test_is_not_nan_or_inf(test_input)
+
+    def _test_is_not_nan_or_inf(self, test_input):
+        eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
+
+        self.assertTrue(torch.isfinite(eigenvalues).all())
+        self.assertTrue(torch.isfinite(eigenvectors).all())
+
+    def test_degenerate_inputs_gpu(self):
+        self._test_degenerate_inputs(self._gpu)
+
+    def test_degenerate_inputs_cpu(self):
+        self._test_degenerate_inputs(self._cpu)
+
+    def _test_degenerate_inputs(self, device):
+        """
+        Test degenerate case when input matrices are diagonal or near-diagonal
+        """
+
+        # Purely diagonal case
+        test_input = self.create_diag_sym3x3(device, self.TEST_BATCH_SIZE)
+
+        self._test_is_not_nan_or_inf(test_input)
+        self._test_is_eigen(test_input)
+        self._test_eigenvectors_are_orthonormal(test_input)
+
+        # Almost-diagonal case
+        test_input = self.create_diag_sym3x3(device, self.TEST_BATCH_SIZE, noise=1e-4)
+
+        self._test_is_not_nan_or_inf(test_input)
+        self._test_is_eigen(test_input)
+        self._test_eigenvectors_are_orthonormal(test_input)
+
+    def test_gradients_cpu(self):
+        self._test_gradients(self._cpu)
+
+    def test_gradients_gpu(self):
+        self._test_gradients(self._gpu)
+
+    def _test_gradients(self, device):
+        """
+        Tests if gradients pass though without any problems (infs, nans etc) and
+        also performs gradcheck (compares numerical and analytical gradients)
+        """
+        test_random_input = self.create_random_sym3x3(device, n=16)
+        test_diag_input = self.create_diag_sym3x3(device, n=16)
+        test_almost_diag_input = self.create_diag_sym3x3(device, n=16, noise=1e-4)
+
+        test_input = torch.cat(
+            (test_random_input, test_diag_input, test_almost_diag_input)
+        )
+        test_input.requires_grad = True
+
+        with torch.autograd.detect_anomaly():
+            eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
+
+            loss = eigenvalues.mean() + eigenvectors.mean()
+            loss.backward()
+
+        test_random_input.requires_grad = True
+        # Inputs are converted to double to increase the precision of gradcheck.
+        torch.autograd.gradcheck(
+            symeig3x3, test_random_input.double(), eps=1e-6, atol=1e-2, rtol=1e-2
+        )
+
+    def _test_eigenvalues_and_eigenvectors(
+        self, test_eigenvectors, test_eigenvalues, atol=1e-04, rtol=1e-04
+    ):
+        test_input = (
+            test_eigenvectors.transpose(-2, -1)
+            @ test_eigenvalues.diag_embed()
+            @ test_eigenvectors
+        )
+
+        test_eigenvalues_sorted, _ = torch.sort(test_eigenvalues, dim=-1)
+
+        eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
+
+        self.assertClose(
+            test_eigenvalues_sorted,
+            eigenvalues,
+            atol=atol,
+            rtol=rtol,
+        )
+
+        self._test_is_not_nan_or_inf(test_input)
+        self._test_is_eigen(test_input, atol=atol, rtol=rtol)
+        self._test_eigenvectors_are_orthonormal(test_input)
+
+    def test_degenerate_eigenvalues_gpu(self):
+        self._test_degenerate_eigenvalues(self._gpu)
+
+    def test_degenerate_eigenvalues_cpu(self):
+        self._test_degenerate_eigenvalues(self._cpu)
+
+    def _test_degenerate_eigenvalues(self, device):
+        """
+        Test degenerate eigenvalues like zero-valued and with 2-/3-multiplicity
+        """
+        # Error tolerances for degenerate values are increased as things might become
+        #  numerically unstable
+        deg_atol = 1e-3
+        deg_rtol = 1.0
+
+        # Construct random orthonormal sets
+        test_eigenvecs = random_rotations(n=self.TEST_BATCH_SIZE, device=device)
+
+        # Construct random eigenvalues
+        test_eigenvals = torch.randn(
+            (self.TEST_BATCH_SIZE, 3), device=test_eigenvecs.device
+        )
+        self._test_eigenvalues_and_eigenvectors(
+            test_eigenvecs, test_eigenvals, atol=deg_atol, rtol=deg_rtol
+        )
+
+        # First eigenvalue is always 0.0 here: [0.0 X Y]
+        test_eigenvals_with_zero = test_eigenvals.clone()
+        test_eigenvals_with_zero[..., 0] = 0.0
+        self._test_eigenvalues_and_eigenvectors(
+            test_eigenvecs, test_eigenvals_with_zero, atol=deg_atol, rtol=deg_rtol
+        )
+
+        # First two eigenvalues are always the same here: [X X Y]
+        test_eigenvals_with_multiplicity2 = test_eigenvals.clone()
+        test_eigenvals_with_multiplicity2[..., 1] = test_eigenvals_with_multiplicity2[
+            ..., 0
+        ]
+        self._test_eigenvalues_and_eigenvectors(
+            test_eigenvecs,
+            test_eigenvals_with_multiplicity2,
+            atol=deg_atol,
+            rtol=deg_rtol,
+        )
+
+        # All three eigenvalues are the same here: [X X X]
+        test_eigenvals_with_multiplicity3 = test_eigenvals_with_multiplicity2.clone()
+        test_eigenvals_with_multiplicity3[..., 2] = test_eigenvals_with_multiplicity2[
+            ..., 0
+        ]
+        self._test_eigenvalues_and_eigenvectors(
+            test_eigenvecs,
+            test_eigenvals_with_multiplicity3,
+            atol=deg_atol,
+            rtol=deg_rtol,
+        )
+
+    def test_more_dimensions(self):
+        """
+        Tests if function supports arbitrary leading dimensions
+        """
+        repeat = 4
+
+        test_input = self.create_random_sym3x3(self._cpu, n=16)
+        test_input_4d = test_input[None, ...].expand((repeat,) + test_input.shape)
+
+        eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
+        eigenvalues_4d, eigenvectors_4d = symeig3x3(test_input_4d, eigenvectors=True)
+
+        eigenvalues_4d_gt = eigenvalues[None, ...].expand((repeat,) + eigenvalues.shape)
+        eigenvectors_4d_gt = eigenvectors[None, ...].expand(
+            (repeat,) + eigenvectors.shape
+        )
+
+        self.assertClose(eigenvalues_4d_gt, eigenvalues_4d)
+        self.assertClose(eigenvectors_4d_gt, eigenvectors_4d)

third_party/AnyBimanual/third_party/pytorch3d/tests/test_texturing.py

@@ -0,0 +1,1325 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import unittest
+
+import torch
+import torch.nn.functional as F
+from pytorch3d.renderer.mesh.rasterizer import Fragments
+from pytorch3d.renderer.mesh.textures import (
+    _list_to_padded_wrapper,
+    TexturesAtlas,
+    TexturesUV,
+    TexturesVertex,
+)
+from pytorch3d.renderer.mesh.utils import (
+    pack_rectangles,
+    pack_unique_rectangles,
+    Rectangle,
+)
+from pytorch3d.structures import list_to_packed, Meshes, packed_to_list
+
+from .common_testing import TestCaseMixin
+from .test_meshes import init_mesh
+
+
+def tryindex(self, index, tex, meshes, source):
+    tex2 = tex[index]
+    meshes2 = meshes[index]
+    tex_from_meshes = meshes2.textures
+    for item in source:
+        basic = source[item][index]
+        from_texture = getattr(tex2, item + "_padded")()
+        from_meshes = getattr(tex_from_meshes, item + "_padded")()
+        if isinstance(index, int):
+            basic = basic[None]
+
+        if len(basic) == 0:
+            self.assertEqual(len(from_texture), 0)
+            self.assertEqual(len(from_meshes), 0)
+        else:
+            self.assertClose(basic, from_texture)
+            self.assertClose(basic, from_meshes)
+            self.assertEqual(from_texture.ndim, getattr(tex, item + "_padded")().ndim)
+            item_list = getattr(tex_from_meshes, item + "_list")()
+            self.assertEqual(basic.shape[0], len(item_list))
+            for i, elem in enumerate(item_list):
+                self.assertClose(elem, basic[i])
+
+
+class TestTexturesVertex(TestCaseMixin, unittest.TestCase):
+    def test_sample_vertex_textures(self):
+        """
+        This tests both interpolate_vertex_colors as well as
+        interpolate_face_attributes.
+        """
+        verts = torch.randn((4, 3), dtype=torch.float32)
+        faces = torch.tensor([[2, 1, 0], [3, 1, 0]], dtype=torch.int64)
+        vert_tex = torch.tensor(
+            [[0, 1, 0], [0, 1, 1], [1, 1, 0], [1, 1, 1]], dtype=torch.float32
+        )
+        verts_features = vert_tex
+        tex = TexturesVertex(verts_features=[verts_features])
+        mesh = Meshes(verts=[verts], faces=[faces], textures=tex)
+        pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
+        barycentric_coords = torch.tensor(
+            [[0.5, 0.3, 0.2], [0.3, 0.6, 0.1]], dtype=torch.float32
+        ).view(1, 1, 1, 2, -1)
+        expected_vals = torch.tensor(
+            [[0.5, 1.0, 0.3], [0.3, 1.0, 0.9]], dtype=torch.float32
+        ).view(1, 1, 1, 2, -1)
+        fragments = Fragments(
+            pix_to_face=pix_to_face,
+            bary_coords=barycentric_coords,
+            zbuf=torch.ones_like(pix_to_face),
+            dists=torch.ones_like(pix_to_face),
+        )
+        # sample_textures calls interpolate_vertex_colors
+        texels = mesh.sample_textures(fragments)
+        self.assertTrue(torch.allclose(texels, expected_vals[None, :]))
+
+    def test_sample_vertex_textures_grad(self):
+        verts = torch.randn((4, 3), dtype=torch.float32)
+        faces = torch.tensor([[2, 1, 0], [3, 1, 0]], dtype=torch.int64)
+        vert_tex = torch.tensor(
+            [[0, 1, 0], [0, 1, 1], [1, 1, 0], [1, 1, 1]],
+            dtype=torch.float32,
+            requires_grad=True,
+        )
+        verts_features = vert_tex
+        tex = TexturesVertex(verts_features=[verts_features])
+        mesh = Meshes(verts=[verts], faces=[faces], textures=tex)
+        pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
+        barycentric_coords = torch.tensor(
+            [[0.5, 0.3, 0.2], [0.3, 0.6, 0.1]], dtype=torch.float32
+        ).view(1, 1, 1, 2, -1)
+        fragments = Fragments(
+            pix_to_face=pix_to_face,
+            bary_coords=barycentric_coords,
+            zbuf=torch.ones_like(pix_to_face),
+            dists=torch.ones_like(pix_to_face),
+        )
+        grad_vert_tex = torch.tensor(
+            [[0.3, 0.3, 0.3], [0.9, 0.9, 0.9], [0.5, 0.5, 0.5], [0.3, 0.3, 0.3]],
+            dtype=torch.float32,
+        )
+        texels = mesh.sample_textures(fragments)
+        texels.sum().backward()
+        self.assertTrue(hasattr(vert_tex, "grad"))
+        self.assertTrue(torch.allclose(vert_tex.grad, grad_vert_tex[None, :]))
+
+    def test_textures_vertex_init_fail(self):
+        # Incorrect sized tensors
+        with self.assertRaisesRegex(ValueError, "verts_features"):
+            TexturesVertex(verts_features=torch.rand(size=(5, 10)))
+
+        # Not a list or a tensor
+        with self.assertRaisesRegex(ValueError, "verts_features"):
+            TexturesVertex(verts_features=(1, 1, 1))
+
+    def test_faces_verts_textures(self):
+        device = torch.device("cuda:0")
+        verts = torch.randn((2, 4, 3), dtype=torch.float32, device=device)
+        faces = torch.tensor(
+            [[[2, 1, 0], [3, 1, 0]], [[1, 3, 0], [2, 1, 3]]],
+            dtype=torch.int64,
+            device=device,
+        )
+
+        # define TexturesVertex
+        verts_texture = torch.rand(verts.shape, device=device)
+        textures = TexturesVertex(verts_features=verts_texture)
+
+        # compute packed faces
+        ff = faces.unbind(0)
+        faces_packed = torch.cat([ff[0], ff[1] + verts.shape[1]])
+
+        # face verts textures
+        faces_verts_texts = textures.faces_verts_textures_packed(faces_packed)
+
+        verts_texts_packed = torch.cat(verts_texture.unbind(0))
+        faces_verts_texts_packed = verts_texts_packed[faces_packed]
+
+        self.assertClose(faces_verts_texts_packed, faces_verts_texts)
+
+    def test_submeshes(self):
+        # define TexturesVertex
+        verts_features = torch.tensor(
+            [
+                [1, 0, 0],
+                [1, 0, 0],
+                [1, 0, 0],
+                [1, 0, 0],
+                [0, 1, 0],
+                [0, 1, 0],
+                [0, 1, 0],
+                [0, 1, 0],
+            ],
+            dtype=torch.float32,
+        )
+
+        textures = TexturesVertex(
+            verts_features=[verts_features, verts_features, verts_features]
+        )
+        subtextures = textures.submeshes(
+            [
+                [
+                    torch.LongTensor([0, 2, 3]),
+                    torch.LongTensor(list(range(8))),
+                ],
+                [],
+                [
+                    torch.LongTensor([4]),
+                ],
+            ],
+            None,
+        )
+
+        subtextures_features = subtextures.verts_features_list()
+
+        self.assertEqual(len(subtextures_features), 3)
+        self.assertTrue(
+            torch.equal(
+                subtextures_features[0],
+                torch.FloatTensor([[1, 0, 0], [1, 0, 0], [1, 0, 0]]),
+            )
+        )
+        self.assertTrue(torch.equal(subtextures_features[1], verts_features))
+        self.assertTrue(
+            torch.equal(subtextures_features[2], torch.FloatTensor([[0, 1, 0]]))
+        )
+
+    def test_clone(self):
+        tex = TexturesVertex(verts_features=torch.rand(size=(10, 100, 128)))
+        tex.verts_features_list()
+        tex_cloned = tex.clone()
+        self.assertSeparate(
+            tex._verts_features_padded, tex_cloned._verts_features_padded
+        )
+        self.assertClose(tex._verts_features_padded, tex_cloned._verts_features_padded)
+        self.assertSeparate(tex.valid, tex_cloned.valid)
+        self.assertTrue(tex.valid.eq(tex_cloned.valid).all())
+        for i in range(tex._N):
+            self.assertSeparate(
+                tex._verts_features_list[i], tex_cloned._verts_features_list[i]
+            )
+            self.assertClose(
+                tex._verts_features_list[i], tex_cloned._verts_features_list[i]
+            )
+
+    def test_detach(self):
+        tex = TexturesVertex(
+            verts_features=torch.rand(size=(10, 100, 128), requires_grad=True)
+        )
+        tex.verts_features_list()
+        tex_detached = tex.detach()
+        self.assertFalse(tex_detached._verts_features_padded.requires_grad)
+        self.assertClose(
+            tex_detached._verts_features_padded, tex._verts_features_padded
+        )
+        for i in range(tex._N):
+            self.assertClose(
+                tex._verts_features_list[i], tex_detached._verts_features_list[i]
+            )
+            self.assertFalse(tex_detached._verts_features_list[i].requires_grad)
+
+    def test_extend(self):
+        B = 10
+        mesh = init_mesh(B, 30, 50)
+        V = mesh._V
+        tex_uv = TexturesVertex(verts_features=torch.randn((B, V, 3)))
+        tex_mesh = Meshes(
+            verts=mesh.verts_padded(), faces=mesh.faces_padded(), textures=tex_uv
+        )
+        N = 20
+        new_mesh = tex_mesh.extend(N)
+
+        self.assertEqual(len(tex_mesh) * N, len(new_mesh))
+
+        tex_init = tex_mesh.textures
+        new_tex = new_mesh.textures
+
+        for i in range(len(tex_mesh)):
+            for n in range(N):
+                self.assertClose(
+                    tex_init.verts_features_list()[i],
+                    new_tex.verts_features_list()[i * N + n],
+                )
+                self.assertClose(
+                    tex_init._num_faces_per_mesh[i],
+                    new_tex._num_faces_per_mesh[i * N + n],
+                )
+
+        self.assertAllSeparate(
+            [tex_init.verts_features_padded(), new_tex.verts_features_padded()]
+        )
+
+        with self.assertRaises(ValueError):
+            tex_mesh.extend(N=-1)
+
+    def test_padded_to_packed(self):
+        # Case where each face in the mesh has 3 unique uv vertex indices
+        # - i.e. even if a vertex is shared between multiple faces it will
+        # have a unique uv coordinate for each face.
+        num_verts_per_mesh = [9, 6]
+        D = 10
+        verts_features_list = [torch.rand(v, D) for v in num_verts_per_mesh]
+        verts_features_packed = list_to_packed(verts_features_list)[0]
+        verts_features_list = packed_to_list(verts_features_packed, num_verts_per_mesh)
+        tex = TexturesVertex(verts_features=verts_features_list)
+
+        # This is set inside Meshes when textures is passed as an input.
+        # Here we set _num_faces_per_mesh and _num_verts_per_mesh explicity.
+        tex1 = tex.clone()
+        tex1._num_verts_per_mesh = num_verts_per_mesh
+        verts_packed = tex1.verts_features_packed()
+        verts_verts_list = tex1.verts_features_list()
+        verts_padded = tex1.verts_features_padded()
+
+        for f1, f2 in zip(verts_verts_list, verts_features_list):
+            self.assertTrue((f1 == f2).all().item())
+
+        self.assertTrue(verts_packed.shape == (sum(num_verts_per_mesh), D))
+        self.assertTrue(verts_padded.shape == (2, 9, D))
+
+        # Case where num_verts_per_mesh is not set and textures
+        # are initialized with a padded tensor.
+        tex2 = TexturesVertex(verts_features=verts_padded)
+        verts_packed = tex2.verts_features_packed()
+        verts_list = tex2.verts_features_list()
+
+        # Packed is just flattened padded as num_verts_per_mesh
+        # has not been provided.
+        self.assertTrue(verts_packed.shape == (9 * 2, D))
+
+        for i, (f1, f2) in enumerate(zip(verts_list, verts_features_list)):
+            n = num_verts_per_mesh[i]
+            self.assertTrue((f1[:n] == f2).all().item())
+
+    def test_getitem(self):
+        N = 5
+        V = 20
+        source = {"verts_features": torch.randn(size=(N, V, 128))}
+        tex = TexturesVertex(verts_features=source["verts_features"])
+
+        verts = torch.rand(size=(N, V, 3))
+        faces = torch.randint(size=(N, 10, 3), high=V)
+        meshes = Meshes(verts=verts, faces=faces, textures=tex)
+
+        tryindex(self, 2, tex, meshes, source)
+        tryindex(self, slice(0, 2, 1), tex, meshes, source)
+        index = torch.tensor([1, 0, 1, 0, 0], dtype=torch.bool)
+        tryindex(self, index, tex, meshes, source)
+        index = torch.tensor([0, 0, 0, 0, 0], dtype=torch.bool)
+        tryindex(self, index, tex, meshes, source)
+        index = torch.tensor([1, 2], dtype=torch.int64)
+        tryindex(self, index, tex, meshes, source)
+        tryindex(self, [2, 4], tex, meshes, source)
+
+    def test_sample_textures_error(self):
+        N = 5
+        V = 20
+        verts = torch.rand(size=(N, V, 3))
+        faces = torch.randint(size=(N, 10, 3), high=V)
+        tex = TexturesVertex(verts_features=torch.randn(size=(N, 10, 128)))
+
+        # Verts features have the wrong number of verts
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            Meshes(verts=verts, faces=faces, textures=tex)
+
+        # Verts features have the wrong batch dim
+        tex = TexturesVertex(verts_features=torch.randn(size=(1, V, 128)))
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            Meshes(verts=verts, faces=faces, textures=tex)
+
+        meshes = Meshes(verts=verts, faces=faces)
+        meshes.textures = tex
+
+        # Cannot use the texture attribute set on meshes for sampling
+        # textures if the dimensions don't match
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            meshes.sample_textures(None)
+
+
+class TestTexturesAtlas(TestCaseMixin, unittest.TestCase):
+    def test_sample_texture_atlas(self):
+        N, F, R = 1, 2, 2
+        verts = torch.randn((4, 3), dtype=torch.float32)
+        faces = torch.tensor([[2, 1, 0], [3, 1, 0]], dtype=torch.int64)
+        faces_atlas = torch.rand(size=(N, F, R, R, 3))
+        tex = TexturesAtlas(atlas=faces_atlas)
+        mesh = Meshes(verts=[verts], faces=[faces], textures=tex)
+        pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
+        barycentric_coords = torch.tensor(
+            [[0.5, 0.3, 0.2], [0.3, 0.6, 0.1]], dtype=torch.float32
+        ).view(1, 1, 1, 2, -1)
+        expected_vals = torch.tensor(
+            [[0.5, 1.0, 0.3], [0.3, 1.0, 0.9]], dtype=torch.float32
+        )
+        expected_vals = torch.zeros((1, 1, 1, 2, 3), dtype=torch.float32)
+        expected_vals[..., 0, :] = faces_atlas[0, 0, 0, 1, ...]
+        expected_vals[..., 1, :] = faces_atlas[0, 1, 1, 0, ...]
+
+        fragments = Fragments(
+            pix_to_face=pix_to_face,
+            bary_coords=barycentric_coords,
+            zbuf=torch.ones_like(pix_to_face),
+            dists=torch.ones_like(pix_to_face),
+        )
+        texels = mesh.textures.sample_textures(fragments)
+        self.assertTrue(torch.allclose(texels, expected_vals))
+
+    def test_textures_atlas_grad(self):
+        N, F, R = 1, 2, 2
+        verts = torch.randn((4, 3), dtype=torch.float32)
+        faces = torch.tensor([[2, 1, 0], [3, 1, 0]], dtype=torch.int64)
+        faces_atlas = torch.rand(size=(N, F, R, R, 3), requires_grad=True)
+        tex = TexturesAtlas(atlas=faces_atlas)
+        mesh = Meshes(verts=[verts], faces=[faces], textures=tex)
+        pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
+        barycentric_coords = torch.tensor(
+            [[0.5, 0.3, 0.2], [0.3, 0.6, 0.1]], dtype=torch.float32
+        ).view(1, 1, 1, 2, -1)
+        fragments = Fragments(
+            pix_to_face=pix_to_face,
+            bary_coords=barycentric_coords,
+            zbuf=torch.ones_like(pix_to_face),
+            dists=torch.ones_like(pix_to_face),
+        )
+        texels = mesh.textures.sample_textures(fragments)
+        grad_tex = torch.rand_like(texels)
+        grad_expected = torch.zeros_like(faces_atlas)
+        grad_expected[0, 0, 0, 1, :] = grad_tex[..., 0:1, :]
+        grad_expected[0, 1, 1, 0, :] = grad_tex[..., 1:2, :]
+        texels.backward(grad_tex)
+        self.assertTrue(hasattr(faces_atlas, "grad"))
+        self.assertTrue(torch.allclose(faces_atlas.grad, grad_expected))
+
+    def test_textures_atlas_init_fail(self):
+        # Incorrect sized tensors
+        with self.assertRaisesRegex(ValueError, "atlas"):
+            TexturesAtlas(atlas=torch.rand(size=(5, 10, 3)))
+
+        # Not a list or a tensor
+        with self.assertRaisesRegex(ValueError, "atlas"):
+            TexturesAtlas(atlas=(1, 1, 1))
+
+    def test_faces_verts_textures(self):
+        device = torch.device("cuda:0")
+        N, F, R = 2, 2, 8
+        num_faces = torch.randint(low=1, high=F, size=(N,))
+        faces_atlas = [
+            torch.rand(size=(num_faces[i].item(), R, R, 3), device=device)
+            for i in range(N)
+        ]
+        tex = TexturesAtlas(atlas=faces_atlas)
+
+        # faces_verts naive
+        faces_verts = []
+        for n in range(N):
+            ff = num_faces[n].item()
+            temp = torch.zeros(ff, 3, 3)
+            for f in range(ff):
+                t0 = faces_atlas[n][f, 0, -1]  # for v0, bary = (1, 0)
+                t1 = faces_atlas[n][f, -1, 0]  # for v1, bary = (0, 1)
+                t2 = faces_atlas[n][f, 0, 0]  # for v2, bary = (0, 0)
+                temp[f, 0] = t0
+                temp[f, 1] = t1
+                temp[f, 2] = t2
+            faces_verts.append(temp)
+        faces_verts = torch.cat(faces_verts, 0)
+
+        self.assertClose(faces_verts, tex.faces_verts_textures_packed().cpu())
+
+    def test_clone(self):
+        tex = TexturesAtlas(atlas=torch.rand(size=(1, 10, 2, 2, 3)))
+        tex.atlas_list()
+        tex_cloned = tex.clone()
+        self.assertSeparate(tex._atlas_padded, tex_cloned._atlas_padded)
+        self.assertClose(tex._atlas_padded, tex_cloned._atlas_padded)
+        self.assertSeparate(tex.valid, tex_cloned.valid)
+        self.assertTrue(tex.valid.eq(tex_cloned.valid).all())
+        for i in range(tex._N):
+            self.assertSeparate(tex._atlas_list[i], tex_cloned._atlas_list[i])
+            self.assertClose(tex._atlas_list[i], tex_cloned._atlas_list[i])
+
+    def test_detach(self):
+        tex = TexturesAtlas(atlas=torch.rand(size=(1, 10, 2, 2, 3), requires_grad=True))
+        tex.atlas_list()
+        tex_detached = tex.detach()
+        self.assertFalse(tex_detached._atlas_padded.requires_grad)
+        self.assertClose(tex_detached._atlas_padded, tex._atlas_padded)
+        for i in range(tex._N):
+            self.assertFalse(tex_detached._atlas_list[i].requires_grad)
+            self.assertClose(tex._atlas_list[i], tex_detached._atlas_list[i])
+
+    def test_extend(self):
+        B = 10
+        mesh = init_mesh(B, 30, 50)
+        F = mesh._F
+        tex_uv = TexturesAtlas(atlas=torch.randn((B, F, 2, 2, 3)))
+        tex_mesh = Meshes(
+            verts=mesh.verts_padded(), faces=mesh.faces_padded(), textures=tex_uv
+        )
+        N = 20
+        new_mesh = tex_mesh.extend(N)
+
+        self.assertEqual(len(tex_mesh) * N, len(new_mesh))
+
+        tex_init = tex_mesh.textures
+        new_tex = new_mesh.textures
+
+        for i in range(len(tex_mesh)):
+            for n in range(N):
+                self.assertClose(
+                    tex_init.atlas_list()[i], new_tex.atlas_list()[i * N + n]
+                )
+                self.assertClose(
+                    tex_init._num_faces_per_mesh[i],
+                    new_tex._num_faces_per_mesh[i * N + n],
+                )
+
+        self.assertAllSeparate([tex_init.atlas_padded(), new_tex.atlas_padded()])
+
+        with self.assertRaises(ValueError):
+            tex_mesh.extend(N=-1)
+
+    def test_padded_to_packed(self):
+        # Case where each face in the mesh has 3 unique uv vertex indices
+        # - i.e. even if a vertex is shared between multiple faces it will
+        # have a unique uv coordinate for each face.
+        R = 2
+        N = 20
+        num_faces_per_mesh = torch.randint(size=(N,), low=0, high=30)
+        atlas_list = [torch.rand(f, R, R, 3) for f in num_faces_per_mesh]
+        tex = TexturesAtlas(atlas=atlas_list)
+
+        # This is set inside Meshes when textures is passed as an input.
+        # Here we set _num_faces_per_mesh explicity.
+        tex1 = tex.clone()
+        tex1._num_faces_per_mesh = num_faces_per_mesh.tolist()
+        atlas_packed = tex1.atlas_packed()
+        atlas_list_new = tex1.atlas_list()
+        atlas_padded = tex1.atlas_padded()
+
+        for f1, f2 in zip(atlas_list_new, atlas_list):
+            self.assertTrue((f1 == f2).all().item())
+
+        sum_F = num_faces_per_mesh.sum()
+        max_F = num_faces_per_mesh.max().item()
+        self.assertTrue(atlas_packed.shape == (sum_F, R, R, 3))
+        self.assertTrue(atlas_padded.shape == (N, max_F, R, R, 3))
+
+        # Case where num_faces_per_mesh is not set and textures
+        # are initialized with a padded tensor.
+        atlas_list_padded = _list_to_padded_wrapper(atlas_list)
+        tex2 = TexturesAtlas(atlas=atlas_list_padded)
+        atlas_packed = tex2.atlas_packed()
+        atlas_list_new = tex2.atlas_list()
+
+        # Packed is just flattened padded as num_faces_per_mesh
+        # has not been provided.
+        self.assertTrue(atlas_packed.shape == (N * max_F, R, R, 3))
+
+        for i, (f1, f2) in enumerate(zip(atlas_list_new, atlas_list)):
+            n = num_faces_per_mesh[i]
+            self.assertTrue((f1[:n] == f2).all().item())
+
+    def test_getitem(self):
+        N = 5
+        V = 20
+        F = 10
+        source = {"atlas": torch.randn(size=(N, F, 4, 4, 3))}
+        tex = TexturesAtlas(atlas=source["atlas"])
+
+        verts = torch.rand(size=(N, V, 3))
+        faces = torch.randint(size=(N, F, 3), high=V)
+        meshes = Meshes(verts=verts, faces=faces, textures=tex)
+
+        tryindex(self, 2, tex, meshes, source)
+        tryindex(self, slice(0, 2, 1), tex, meshes, source)
+        index = torch.tensor([1, 0, 1, 0, 0], dtype=torch.bool)
+        tryindex(self, index, tex, meshes, source)
+        index = torch.tensor([0, 0, 0, 0, 0], dtype=torch.bool)
+        tryindex(self, index, tex, meshes, source)
+        index = torch.tensor([1, 2], dtype=torch.int64)
+        tryindex(self, index, tex, meshes, source)
+        tryindex(self, [2, 4], tex, meshes, source)
+
+    def test_sample_textures_error(self):
+        N = 1
+        V = 20
+        F = 10
+        verts = torch.rand(size=(5, V, 3))
+        faces = torch.randint(size=(5, F, 3), high=V)
+        meshes = Meshes(verts=verts, faces=faces)
+
+        # TexturesAtlas have the wrong batch dim
+        tex = TexturesAtlas(atlas=torch.randn(size=(1, F, 4, 4, 3)))
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            Meshes(verts=verts, faces=faces, textures=tex)
+
+        # TexturesAtlas have the wrong number of faces
+        tex = TexturesAtlas(atlas=torch.randn(size=(N, 15, 4, 4, 3)))
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            Meshes(verts=verts, faces=faces, textures=tex)
+
+        meshes = Meshes(verts=verts, faces=faces)
+        meshes.textures = tex
+
+        # Cannot use the texture attribute set on meshes for sampling
+        # textures if the dimensions don't match
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            meshes.sample_textures(None)
+
+    def test_submeshes(self):
+        N = 2
+        V = 5
+        F = 5
+        tex = TexturesAtlas(
+            atlas=torch.arange(N * F * 4 * 4 * 3, dtype=torch.float32).reshape(
+                N, F, 4, 4, 3
+            )
+        )
+
+        verts = torch.rand(size=(N, V, 3))
+        faces = torch.randint(size=(N, F, 3), high=V)
+        mesh = Meshes(verts=verts, faces=faces, textures=tex)
+
+        sub_faces = [
+            [torch.tensor([0, 2]), torch.tensor([1, 2])],
+            [],
+        ]
+        subtex = mesh.submeshes(sub_faces).textures
+        subtex_faces = subtex.atlas_list()
+
+        self.assertEqual(len(subtex_faces), 2)
+        self.assertClose(
+            subtex_faces[0].flatten().msort(),
+            torch.cat(
+                (
+                    torch.arange(4 * 4 * 3, dtype=torch.float32),
+                    torch.arange(96, 96 + 4 * 4 * 3, dtype=torch.float32),
+                ),
+                0,
+            ),
+        )
+
+
+class TestTexturesUV(TestCaseMixin, unittest.TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        torch.manual_seed(42)
+
+    def test_sample_textures_uv(self):
+        barycentric_coords = torch.tensor(
+            [[0.5, 0.3, 0.2], [0.3, 0.6, 0.1]], dtype=torch.float32
+        ).view(1, 1, 1, 2, -1)
+        dummy_verts = torch.zeros(4, 3)
+        vert_uvs = torch.tensor([[1, 0], [0, 1], [1, 1], [0, 0]], dtype=torch.float32)
+        face_uvs = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.int64)
+        interpolated_uvs = torch.tensor(
+            [[0.5 + 0.2, 0.3 + 0.2], [0.6, 0.3 + 0.6]], dtype=torch.float32
+        )
+
+        # Create a dummy texture map
+        H = 2
+        W = 2
+        x = torch.linspace(0, 1, W).view(1, W).expand(H, W)
+        y = torch.linspace(0, 1, H).view(H, 1).expand(H, W)
+        tex_map = torch.stack([x, y], dim=2).view(1, H, W, 2)
+        pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
+        fragments = Fragments(
+            pix_to_face=pix_to_face,
+            bary_coords=barycentric_coords,
+            zbuf=pix_to_face,
+            dists=pix_to_face,
+        )
+
+        for align_corners in [True, False]:
+            tex = TexturesUV(
+                maps=tex_map,
+                faces_uvs=[face_uvs],
+                verts_uvs=[vert_uvs],
+                align_corners=align_corners,
+            )
+            meshes = Meshes(verts=[dummy_verts], faces=[face_uvs], textures=tex)
+            mesh_textures = meshes.textures
+            texels = mesh_textures.sample_textures(fragments)
+
+            # Expected output
+            pixel_uvs = interpolated_uvs * 2.0 - 1.0
+            pixel_uvs = pixel_uvs.view(2, 1, 1, 2)
+            tex_map_ = torch.flip(tex_map, [1]).permute(0, 3, 1, 2)
+            tex_map_ = torch.cat([tex_map_, tex_map_], dim=0)
+            expected_out = F.grid_sample(
+                tex_map_, pixel_uvs, align_corners=align_corners, padding_mode="border"
+            )
+            self.assertTrue(torch.allclose(texels.squeeze(), expected_out.squeeze()))
+
+    def test_textures_uv_init_fail(self):
+        # Maps has wrong shape
+        with self.assertRaisesRegex(ValueError, "maps"):
+            TexturesUV(
+                maps=torch.ones((5, 16, 16, 3, 4)),
+                faces_uvs=torch.rand(size=(5, 10, 3)),
+                verts_uvs=torch.rand(size=(5, 15, 2)),
+            )
+
+        # faces_uvs has wrong shape
+        with self.assertRaisesRegex(ValueError, "faces_uvs"):
+            TexturesUV(
+                maps=torch.ones((5, 16, 16, 3)),
+                faces_uvs=torch.rand(size=(5, 10, 3, 3)),
+                verts_uvs=torch.rand(size=(5, 15, 2)),
+            )
+
+        # verts_uvs has wrong shape
+        with self.assertRaisesRegex(ValueError, "verts_uvs"):
+            TexturesUV(
+                maps=torch.ones((5, 16, 16, 3)),
+                faces_uvs=torch.rand(size=(5, 10, 3)),
+                verts_uvs=torch.rand(size=(5, 15, 2, 3)),
+            )
+
+        # verts has different batch dim to faces
+        with self.assertRaisesRegex(ValueError, "verts_uvs"):
+            TexturesUV(
+                maps=torch.ones((5, 16, 16, 3)),
+                faces_uvs=torch.rand(size=(5, 10, 3)),
+                verts_uvs=torch.rand(size=(8, 15, 2)),
+            )
+
+        # maps has different batch dim to faces
+        with self.assertRaisesRegex(ValueError, "maps"):
+            TexturesUV(
+                maps=torch.ones((8, 16, 16, 3)),
+                faces_uvs=torch.rand(size=(5, 10, 3)),
+                verts_uvs=torch.rand(size=(5, 15, 2)),
+            )
+
+        # verts on different device to faces
+        with self.assertRaisesRegex(ValueError, "verts_uvs"):
+            TexturesUV(
+                maps=torch.ones((5, 16, 16, 3)),
+                faces_uvs=torch.rand(size=(5, 10, 3)),
+                verts_uvs=torch.rand(size=(5, 15, 2, 3), device="cuda"),
+            )
+
+        # maps on different device to faces
+        with self.assertRaisesRegex(ValueError, "map"):
+            TexturesUV(
+                maps=torch.ones((5, 16, 16, 3), device="cuda"),
+                faces_uvs=torch.rand(size=(5, 10, 3)),
+                verts_uvs=torch.rand(size=(5, 15, 2)),
+            )
+
+        # maps ids are not none but maps doesn't have multiple map indices
+        with self.assertRaisesRegex(ValueError, "map"):
+            TexturesUV(
+                maps=torch.ones((5, 16, 16, 3)),
+                faces_uvs=torch.rand(size=(5, 10, 3)),
+                verts_uvs=torch.rand(size=(5, 15, 2)),
+                maps_ids=torch.randint(0, 1, (5, 10), dtype=torch.long),
+            )
+        # maps ids is none but maps have multiple map indices
+        with self.assertRaisesRegex(ValueError, "map"):
+            TexturesUV(
+                maps=torch.ones((5, 2, 16, 16, 3)),
+                faces_uvs=torch.rand(size=(5, 10, 3)),
+                verts_uvs=torch.rand(size=(5, 15, 2)),
+            )
+
+    def test_faces_verts_textures(self):
+        device = torch.device("cuda:0")
+        N, V, F, H, W = 2, 5, 12, 8, 8
+        vert_uvs = torch.rand((N, V, 2), dtype=torch.float32, device=device)
+        face_uvs = torch.randint(
+            high=V, size=(N, F, 3), dtype=torch.int64, device=device
+        )
+        maps = torch.rand((N, H, W, 3), dtype=torch.float32, device=device)
+
+        tex = TexturesUV(maps=maps, verts_uvs=vert_uvs, faces_uvs=face_uvs)
+
+        # naive faces_verts_textures
+        faces_verts_texs = []
+        for n in range(N):
+            temp = torch.zeros((F, 3, 3), device=device, dtype=torch.float32)
+            for f in range(F):
+                uv0 = vert_uvs[n, face_uvs[n, f, 0]]
+                uv1 = vert_uvs[n, face_uvs[n, f, 1]]
+                uv2 = vert_uvs[n, face_uvs[n, f, 2]]
+
+                idx = torch.stack((uv0, uv1, uv2), dim=0).view(1, 1, 3, 2)  # 1x1x3x2
+                idx = idx * 2.0 - 1.0
+                imap = maps[n].view(1, H, W, 3).permute(0, 3, 1, 2)  # 1x3xHxW
+                imap = torch.flip(imap, [2])
+
+                texts = torch.nn.functional.grid_sample(
+                    imap,
+                    idx,
+                    align_corners=tex.align_corners,
+                    padding_mode=tex.padding_mode,
+                )  # 1x3x1x3
+                temp[f] = texts[0, :, 0, :].permute(1, 0)
+            faces_verts_texs.append(temp)
+        faces_verts_texs = torch.cat(faces_verts_texs, 0)
+
+        self.assertClose(faces_verts_texs, tex.faces_verts_textures_packed())
+
+    def test_faces_verts_multiple_map_textures(self):
+        device = torch.device("cuda:0")
+        N, M, V, F, H, W = 2, 3, 5, 12, 8, 8
+        vert_uvs = torch.rand((N, V, 2), dtype=torch.float32, device=device)
+        face_uvs = torch.randint(
+            high=V, size=(N, F, 3), dtype=torch.int64, device=device
+        )
+        map_ids = torch.randint(0, M, (N, F), device=device)
+        maps = torch.rand((N, M, H, W, 3), dtype=torch.float32, device=device)
+
+        tex = TexturesUV(
+            maps=maps, verts_uvs=vert_uvs, faces_uvs=face_uvs, maps_ids=map_ids
+        )
+
+        # naive faces_verts_textures
+        faces_verts_texs = []
+        for n in range(N):
+            temp = torch.zeros((F, 3, 3), device=device, dtype=torch.float32)
+            for f in range(F):
+                uv0 = vert_uvs[n, face_uvs[n, f, 0]]
+                uv1 = vert_uvs[n, face_uvs[n, f, 1]]
+                uv2 = vert_uvs[n, face_uvs[n, f, 2]]
+                map_id = map_ids[n, f]
+
+                idx = torch.stack((uv0, uv1, uv2), dim=0).view(1, 1, 3, 2)  # 1x1x3x2
+                idx = idx * 2.0 - 1.0
+                imap = maps[n, map_id].view(1, H, W, 3).permute(0, 3, 1, 2)  # 1x3xHxW
+                imap = torch.flip(imap, [2])
+
+                texts = torch.nn.functional.grid_sample(
+                    imap,
+                    idx,
+                    align_corners=tex.align_corners,
+                    padding_mode=tex.padding_mode,
+                )  # 1x3x1x3
+                temp[f] = texts[0, :, 0, :].permute(1, 0)
+            faces_verts_texs.append(temp)
+        faces_verts_texs = torch.cat(faces_verts_texs, 0)
+
+        self.assertClose(faces_verts_texs, tex.faces_verts_textures_packed())
+
+    def test_clone(self):
+        tex = TexturesUV(
+            maps=torch.ones((5, 16, 16, 3)),
+            faces_uvs=torch.rand(size=(5, 10, 3)),
+            verts_uvs=torch.rand(size=(5, 15, 2)),
+        )
+        tex.faces_uvs_list()
+        tex.verts_uvs_list()
+        tex_cloned = tex.clone()
+        self.assertSeparate(tex._faces_uvs_padded, tex_cloned._faces_uvs_padded)
+        self.assertClose(tex._faces_uvs_padded, tex_cloned._faces_uvs_padded)
+        self.assertSeparate(tex._verts_uvs_padded, tex_cloned._verts_uvs_padded)
+        self.assertClose(tex._verts_uvs_padded, tex_cloned._verts_uvs_padded)
+        self.assertSeparate(tex._maps_padded, tex_cloned._maps_padded)
+        self.assertClose(tex._maps_padded, tex_cloned._maps_padded)
+        self.assertSeparate(tex.valid, tex_cloned.valid)
+        self.assertTrue(tex.valid.eq(tex_cloned.valid).all())
+        for i in range(tex._N):
+            self.assertSeparate(tex._faces_uvs_list[i], tex_cloned._faces_uvs_list[i])
+            self.assertClose(tex._faces_uvs_list[i], tex_cloned._faces_uvs_list[i])
+            self.assertSeparate(tex._verts_uvs_list[i], tex_cloned._verts_uvs_list[i])
+            self.assertClose(tex._verts_uvs_list[i], tex_cloned._verts_uvs_list[i])
+            # tex._maps_list is not use anywhere so it's not stored. We call it explicitly
+            self.assertSeparate(tex.maps_list()[i], tex_cloned.maps_list()[i])
+            self.assertClose(tex.maps_list()[i], tex_cloned.maps_list()[i])
+
+    def test_multiple_maps_clone(self):
+        tex = TexturesUV(
+            maps=torch.ones((5, 3, 16, 16, 3)),
+            faces_uvs=torch.rand(size=(5, 10, 3)),
+            verts_uvs=torch.rand(size=(5, 15, 2)),
+            maps_ids=torch.randint(0, 3, (5, 10)),
+        )
+        tex.faces_uvs_list()
+        tex.verts_uvs_list()
+        tex_cloned = tex.clone()
+        self.assertSeparate(tex._faces_uvs_padded, tex_cloned._faces_uvs_padded)
+        self.assertClose(tex._faces_uvs_padded, tex_cloned._faces_uvs_padded)
+        self.assertSeparate(tex._verts_uvs_padded, tex_cloned._verts_uvs_padded)
+        self.assertClose(tex._verts_uvs_padded, tex_cloned._verts_uvs_padded)
+        self.assertSeparate(tex._maps_padded, tex_cloned._maps_padded)
+        self.assertClose(tex._maps_padded, tex_cloned._maps_padded)
+        self.assertSeparate(tex.valid, tex_cloned.valid)
+        self.assertTrue(tex.valid.eq(tex_cloned.valid).all())
+        self.assertSeparate(tex._maps_ids_padded, tex_cloned._maps_ids_padded)
+        self.assertClose(tex._maps_ids_padded, tex_cloned._maps_ids_padded)
+        for i in range(tex._N):
+            self.assertSeparate(tex._faces_uvs_list[i], tex_cloned._faces_uvs_list[i])
+            self.assertClose(tex._faces_uvs_list[i], tex_cloned._faces_uvs_list[i])
+            self.assertSeparate(tex._verts_uvs_list[i], tex_cloned._verts_uvs_list[i])
+            self.assertClose(tex._verts_uvs_list[i], tex_cloned._verts_uvs_list[i])
+            # tex._maps_list is not use anywhere so it's not stored. We call it explicitly
+            self.assertSeparate(tex.maps_list()[i], tex_cloned.maps_list()[i])
+            self.assertClose(tex.maps_list()[i], tex_cloned.maps_list()[i])
+            self.assertSeparate(tex.maps_ids_list()[i], tex_cloned.maps_ids_list()[i])
+            self.assertClose(tex.maps_ids_list()[i], tex_cloned.maps_ids_list()[i])
+
+    def test_detach(self):
+        tex = TexturesUV(
+            maps=torch.ones((5, 16, 16, 3), requires_grad=True),
+            faces_uvs=torch.rand(size=(5, 10, 3)),
+            verts_uvs=torch.rand(size=(5, 15, 2)),
+        )
+        tex.faces_uvs_list()
+        tex.verts_uvs_list()
+        tex_detached = tex.detach()
+        self.assertFalse(tex_detached._maps_padded.requires_grad)
+        self.assertClose(tex._maps_padded, tex_detached._maps_padded)
+        self.assertFalse(tex_detached._verts_uvs_padded.requires_grad)
+        self.assertClose(tex._verts_uvs_padded, tex_detached._verts_uvs_padded)
+        self.assertFalse(tex_detached._faces_uvs_padded.requires_grad)
+        self.assertClose(tex._faces_uvs_padded, tex_detached._faces_uvs_padded)
+        for i in range(tex._N):
+            self.assertFalse(tex_detached._verts_uvs_list[i].requires_grad)
+            self.assertClose(tex._verts_uvs_list[i], tex_detached._verts_uvs_list[i])
+            self.assertFalse(tex_detached._faces_uvs_list[i].requires_grad)
+            self.assertClose(tex._faces_uvs_list[i], tex_detached._faces_uvs_list[i])
+            # tex._maps_list is not use anywhere so it's not stored. We call it explicitly
+            self.assertFalse(tex_detached.maps_list()[i].requires_grad)
+            self.assertClose(tex.maps_list()[i], tex_detached.maps_list()[i])
+
+    def test_multiple_maps_detach(self):
+        tex = TexturesUV(
+            maps=torch.ones((5, 3, 16, 16, 3), requires_grad=True),
+            faces_uvs=torch.rand(size=(5, 10, 3)),
+            verts_uvs=torch.rand(size=(5, 15, 2)),
+            maps_ids=torch.randint(0, 3, (5, 10)),
+        )
+        tex.faces_uvs_list()
+        tex.verts_uvs_list()
+        tex_detached = tex.detach()
+        self.assertFalse(tex_detached._maps_padded.requires_grad)
+        self.assertClose(tex._maps_padded, tex_detached._maps_padded)
+        self.assertFalse(tex_detached._verts_uvs_padded.requires_grad)
+        self.assertClose(tex._verts_uvs_padded, tex_detached._verts_uvs_padded)
+        self.assertFalse(tex_detached._faces_uvs_padded.requires_grad)
+        self.assertClose(tex._faces_uvs_padded, tex_detached._faces_uvs_padded)
+        self.assertFalse(tex_detached._maps_ids_padded.requires_grad)
+        self.assertClose(tex._maps_ids_padded, tex_detached._maps_ids_padded)
+        for i in range(tex._N):
+            self.assertFalse(tex_detached._verts_uvs_list[i].requires_grad)
+            self.assertClose(tex._verts_uvs_list[i], tex_detached._verts_uvs_list[i])
+            self.assertFalse(tex_detached._faces_uvs_list[i].requires_grad)
+            self.assertClose(tex._faces_uvs_list[i], tex_detached._faces_uvs_list[i])
+            # tex._maps_list is not use anywhere so it's not stored. We call it explicitly
+            self.assertFalse(tex_detached.maps_list()[i].requires_grad)
+            self.assertClose(tex.maps_list()[i], tex_detached.maps_list()[i])
+            self.assertFalse(tex_detached.maps_ids_list()[i].requires_grad)
+            self.assertClose(tex.maps_ids_list()[i], tex_detached.maps_ids_list()[i])
+
+    def test_extend(self):
+        B = 5
+        mesh = init_mesh(B, 30, 50)
+        V = mesh._V
+        num_faces = mesh.num_faces_per_mesh()
+        num_verts = mesh.num_verts_per_mesh()
+        faces_uvs_list = [torch.randint(size=(f, 3), low=0, high=V) for f in num_faces]
+        verts_uvs_list = [torch.rand(v, 2) for v in num_verts]
+        tex_uv = TexturesUV(
+            maps=torch.ones((B, 16, 16, 3)),
+            faces_uvs=faces_uvs_list,
+            verts_uvs=verts_uvs_list,
+        )
+        tex_mesh = Meshes(
+            verts=mesh.verts_list(), faces=mesh.faces_list(), textures=tex_uv
+        )
+        N = 2
+        new_mesh = tex_mesh.extend(N)
+
+        self.assertEqual(len(tex_mesh) * N, len(new_mesh))
+
+        tex_init = tex_mesh.textures
+        new_tex = new_mesh.textures
+
+        new_tex_num_verts = new_mesh.num_verts_per_mesh()
+        for i in range(len(tex_mesh)):
+            for n in range(N):
+                tex_nv = new_tex_num_verts[i * N + n]
+                self.assertClose(
+                    # The original textures were initialized using
+                    # verts uvs list
+                    tex_init.verts_uvs_list()[i],
+                    # In the new textures, the verts_uvs are initialized
+                    # from padded. The verts per mesh are not used to
+                    # convert from padded to list. See TexturesUV for an
+                    # explanation.
+                    new_tex.verts_uvs_list()[i * N + n][:tex_nv, ...],
+                )
+                self.assertClose(
+                    tex_init.faces_uvs_list()[i], new_tex.faces_uvs_list()[i * N + n]
+                )
+                self.assertClose(
+                    tex_init.maps_padded()[i, ...], new_tex.maps_padded()[i * N + n]
+                )
+                self.assertClose(
+                    tex_init._num_faces_per_mesh[i],
+                    new_tex._num_faces_per_mesh[i * N + n],
+                )
+
+        self.assertAllSeparate(
+            [
+                tex_init.faces_uvs_padded(),
+                new_tex.faces_uvs_padded(),
+                tex_init.verts_uvs_padded(),
+                new_tex.verts_uvs_padded(),
+                tex_init.maps_padded(),
+                new_tex.maps_padded(),
+            ]
+        )
+
+        with self.assertRaises(ValueError):
+            tex_mesh.extend(N=-1)
+
+    def test_padded_to_packed(self):
+        # Case where each face in the mesh has 3 unique uv vertex indices
+        # - i.e. even if a vertex is shared between multiple faces it will
+        # have a unique uv coordinate for each face.
+        N = 2
+        faces_uvs_list = [
+            torch.tensor([[0, 1, 2], [3, 5, 4], [7, 6, 8]]),
+            torch.tensor([[0, 1, 2], [3, 4, 5]]),
+        ]  # (N, 3, 3)
+        verts_uvs_list = [torch.ones(9, 2), torch.ones(6, 2)]
+        maps_ids_given_list = [torch.randint(0, 3, (3,)), torch.randint(0, 3, (2,))]
+
+        num_faces_per_mesh = [f.shape[0] for f in faces_uvs_list]
+        num_verts_per_mesh = [v.shape[0] for v in verts_uvs_list]
+        tex = TexturesUV(
+            maps=torch.ones((N, 3, 16, 16, 3)),
+            faces_uvs=faces_uvs_list,
+            verts_uvs=verts_uvs_list,
+            maps_ids=maps_ids_given_list,
+        )
+
+        # This is set inside Meshes when textures is passed as an input.
+        # Here we set _num_faces_per_mesh and _num_verts_per_mesh explicity.
+        tex1 = tex.clone()
+        tex1._num_faces_per_mesh = num_faces_per_mesh
+        tex1._num_verts_per_mesh = num_verts_per_mesh
+        verts_list = tex1.verts_uvs_list()
+        verts_padded = tex1.verts_uvs_padded()
+
+        faces_list = tex1.faces_uvs_list()
+        faces_padded = tex1.faces_uvs_padded()
+
+        maps_ids_list = tex1.maps_ids_list()
+        maps_ids_padded = tex1.maps_ids_padded()
+
+        for f1, f2 in zip(faces_list, faces_uvs_list):
+            self.assertTrue((f1 == f2).all().item())
+
+        for f1, f2 in zip(verts_list, verts_uvs_list):
+            self.assertTrue((f1 == f2).all().item())
+
+        for f1, f2 in zip(maps_ids_given_list, maps_ids_list):
+            self.assertTrue((f1 == f2).all().item())
+
+        self.assertTrue(faces_padded.shape == (2, 3, 3))
+        self.assertTrue(verts_padded.shape == (2, 9, 2))
+        self.assertTrue(maps_ids_padded.shape == (2, 3))
+
+        # Case where num_faces_per_mesh is not set and faces_verts_uvs
+        # are initialized with a padded tensor.
+        tex2 = TexturesUV(
+            maps=torch.ones((N, 3, 16, 16, 3)),
+            verts_uvs=verts_padded,
+            faces_uvs=faces_padded,
+            maps_ids=maps_ids_padded,
+        )
+        faces_list = tex2.faces_uvs_list()
+        verts_list = tex2.verts_uvs_list()
+        maps_ids_list = tex2.maps_ids_list()
+
+        for i, (f1, f2) in enumerate(zip(faces_list, faces_uvs_list)):
+            n = num_faces_per_mesh[i]
+            self.assertTrue((f1[:n] == f2).all().item())
+
+        for i, (f1, f2) in enumerate(zip(verts_list, verts_uvs_list)):
+            n = num_verts_per_mesh[i]
+            self.assertTrue((f1[:n] == f2).all().item())
+
+        for i, (f1, f2) in enumerate(zip(maps_ids_list, maps_ids_given_list)):
+            n = num_faces_per_mesh[i]
+            self.assertTrue((f1[:n] == f2).all().item())
+
+    def test_to(self):
+        tex = TexturesUV(
+            maps=torch.ones((5, 3, 16, 16, 3)),
+            faces_uvs=torch.randint(size=(5, 10, 3), high=15),
+            verts_uvs=torch.rand(size=(5, 15, 2)),
+            maps_ids=torch.randint(0, 3, (5, 10)),
+        )
+        device = torch.device("cuda:0")
+        tex = tex.to(device)
+        self.assertEqual(tex._faces_uvs_padded.device, device)
+        self.assertEqual(tex._verts_uvs_padded.device, device)
+        self.assertEqual(tex._maps_padded.device, device)
+        self.assertEqual(tex._maps_ids_padded.device, device)
+
+    def test_mesh_to(self):
+        tex_cpu = TexturesUV(
+            maps=torch.ones((5, 3, 16, 16, 3)),
+            faces_uvs=torch.randint(size=(5, 10, 3), high=15),
+            verts_uvs=torch.rand(size=(5, 15, 2)),
+            maps_ids=torch.randint(0, 3, (5, 10)),
+        )
+        verts = torch.rand(size=(5, 15, 3))
+        faces = torch.randint(size=(5, 10, 3), high=15)
+        mesh_cpu = Meshes(faces=faces, verts=verts, textures=tex_cpu)
+        cpu = torch.device("cpu")
+        device = torch.device("cuda:0")
+        tex = mesh_cpu.to(device).textures
+        self.assertEqual(tex._faces_uvs_padded.device, device)
+        self.assertEqual(tex._verts_uvs_padded.device, device)
+        self.assertEqual(tex._maps_padded.device, device)
+        self.assertEqual(tex._maps_ids_padded.device, device)
+        self.assertEqual(tex_cpu._verts_uvs_padded.device, cpu)
+        self.assertEqual(tex_cpu._maps_ids_padded.device, cpu)
+
+        self.assertEqual(tex_cpu.device, cpu)
+        self.assertEqual(tex.device, device)
+
+    def test_getitem(self):
+        N = 5
+        M = 3
+        V = 20
+        F = 10
+        source = {
+            "maps": torch.rand(size=(N, M, 1, 1, 3)),
+            "faces_uvs": torch.randint(size=(N, F, 3), high=V),
+            "verts_uvs": torch.randn(size=(N, V, 2)),
+            "maps_ids": torch.randint(0, M, (N, F)),
+        }
+        tex = TexturesUV(
+            maps=source["maps"],
+            faces_uvs=source["faces_uvs"],
+            verts_uvs=source["verts_uvs"],
+            maps_ids=source["maps_ids"],
+        )
+
+        verts = torch.rand(size=(N, V, 3))
+        faces = torch.randint(size=(N, F, 3), high=V)
+        meshes = Meshes(verts=verts, faces=faces, textures=tex)
+
+        tryindex(self, 2, tex, meshes, source)
+        tryindex(self, slice(0, 2, 1), tex, meshes, source)
+        index = torch.tensor([1, 0, 1, 0, 0], dtype=torch.bool)
+        tryindex(self, index, tex, meshes, source)
+        index = torch.tensor([0, 0, 0, 0, 0], dtype=torch.bool)
+        tryindex(self, index, tex, meshes, source)
+        index = torch.tensor([1, 2], dtype=torch.int64)
+        tryindex(self, index, tex, meshes, source)
+        tryindex(self, [2, 4], tex, meshes, source)
+
+    def test_centers_for_image(self):
+        maps = torch.rand(size=(1, 257, 129, 3))
+        verts_uvs = torch.FloatTensor([[[0.25, 0.125], [0.5, 0.625], [0.5, 0.5]]])
+        faces_uvs = torch.zeros(size=(1, 0, 3), dtype=torch.int64)
+        tex = TexturesUV(maps=maps, faces_uvs=faces_uvs, verts_uvs=verts_uvs)
+
+        expected = torch.FloatTensor([[32, 224], [64, 96], [64, 128]])
+        self.assertClose(tex.centers_for_image(0), expected)
+
+    def test_sample_textures_error(self):
+        N = 1
+        V = 20
+        F = 10
+        maps = torch.rand(size=(N, 1, 1, 3))
+        verts_uvs = torch.randn(size=(N, V, 2))
+        tex = TexturesUV(
+            maps=maps,
+            faces_uvs=torch.randint(size=(N, 15, 3), high=V),
+            verts_uvs=verts_uvs,
+        )
+        verts = torch.rand(size=(5, V, 3))
+        faces = torch.randint(size=(5, 10, 3), high=V)
+        meshes = Meshes(verts=verts, faces=faces)
+
+        # Wrong number of faces
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            Meshes(verts=verts, faces=faces, textures=tex)
+
+        # Wrong batch dim for faces
+        tex = TexturesUV(
+            maps=maps,
+            faces_uvs=torch.randint(size=(1, F, 3), high=V),
+            verts_uvs=verts_uvs,
+        )
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            Meshes(verts=verts, faces=faces, textures=tex)
+
+        # Wrong batch dim for verts_uvs is not necessary to check as
+        # there is already a check inside TexturesUV for a batch dim
+        # mismatch with faces_uvs
+
+        meshes = Meshes(verts=verts, faces=faces)
+        meshes.textures = tex
+
+        # Cannot use the texture attribute set on meshes for sampling
+        # textures if the dimensions don't match
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            meshes.sample_textures(None)
+
+    def test_submeshes(self):
+        N = 2
+        faces_uvs_list = [
+            torch.LongTensor([[0, 1, 2], [3, 5, 4], [7, 6, 8]]),
+            torch.LongTensor([[0, 1, 2], [3, 4, 5]]),
+        ]
+        verts_uvs_list = [
+            torch.arange(18, dtype=torch.float32).reshape(9, 2),
+            torch.ones(6, 2),
+        ]
+        tex = TexturesUV(
+            maps=torch.rand((N, 16, 16, 3)),
+            faces_uvs=faces_uvs_list,
+            verts_uvs=verts_uvs_list,
+        )
+
+        sub_faces = [
+            [torch.tensor([0, 1]), torch.tensor([1, 2])],
+            [],
+        ]
+
+        mesh = Meshes(
+            verts=[torch.rand(9, 3), torch.rand(6, 3)],
+            faces=faces_uvs_list,
+            textures=tex,
+        )
+        subtex = mesh.submeshes(sub_faces).textures
+        subtex_faces = subtex.faces_uvs_padded()
+        self.assertEqual(len(subtex_faces), 2)
+        self.assertClose(
+            subtex_faces[0],
+            torch.tensor([[0, 1, 2], [3, 5, 4]]),
+        )
+        self.assertClose(
+            subtex.verts_uvs_list()[0][subtex.faces_uvs_list()[0].flatten()]
+            .flatten()
+            .msort(),
+            torch.arange(12, dtype=torch.float32),
+        )
+        self.assertClose(
+            subtex.maps_padded(), tex.maps_padded()[:1].expand(2, -1, -1, -1)
+        )
+
+
+class TestRectanglePacking(TestCaseMixin, unittest.TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        torch.manual_seed(42)
+
+    def wrap_pack(self, sizes):
+        """
+        Call the pack_rectangles function, which we want to test,
+        and return its outputs.
+        Additionally makes some sanity checks on the output.
+        """
+        res = pack_rectangles(sizes)
+        total = res.total_size
+        self.assertGreaterEqual(total[0], 0)
+        self.assertGreaterEqual(total[1], 0)
+        mask = torch.zeros(total, dtype=torch.bool)
+        seen_x_bound = False
+        seen_y_bound = False
+        for (in_x, in_y), (out_x, out_y, flipped, is_first) in zip(
+            sizes, res.locations
+        ):
+            self.assertTrue(is_first)
+            self.assertGreaterEqual(out_x, 0)
+            self.assertGreaterEqual(out_y, 0)
+            placed_x, placed_y = (in_y, in_x) if flipped else (in_x, in_y)
+            upper_x = placed_x + out_x
+            upper_y = placed_y + out_y
+            self.assertGreaterEqual(total[0], upper_x)
+            if total[0] == upper_x:
+                seen_x_bound = True
+            self.assertGreaterEqual(total[1], upper_y)
+            if total[1] == upper_y:
+                seen_y_bound = True
+            already_taken = torch.sum(mask[out_x:upper_x, out_y:upper_y])
+            self.assertEqual(already_taken, 0)
+            mask[out_x:upper_x, out_y:upper_y] = 1
+        self.assertTrue(seen_x_bound)
+        self.assertTrue(seen_y_bound)
+
+        self.assertTrue(torch.all(torch.sum(mask, dim=0, dtype=torch.int32) > 0))
+        self.assertTrue(torch.all(torch.sum(mask, dim=1, dtype=torch.int32) > 0))
+        return res
+
+    def assert_bb(self, sizes, expected):
+        """
+        Apply the pack_rectangles function to sizes and verify the
+        bounding box dimensions are expected.
+        """
+        self.assertSetEqual(set(self.wrap_pack(sizes).total_size), expected)
+
+    def test_simple(self):
+        self.assert_bb([(3, 4), (4, 3)], {6, 4})
+        self.assert_bb([(2, 2), (2, 4), (2, 2)], {4})
+
+        # many squares
+        self.assert_bb([(2, 2)] * 9, {2, 18})
+
+        # One big square and many small ones.
+        self.assert_bb([(3, 3)] + [(1, 1)] * 2, {3, 4})
+        self.assert_bb([(3, 3)] + [(1, 1)] * 3, {3, 4})
+        self.assert_bb([(3, 3)] + [(1, 1)] * 4, {3, 5})
+        self.assert_bb([(3, 3)] + [(1, 1)] * 5, {3, 5})
+        self.assert_bb([(1, 1)] * 6 + [(3, 3)], {3, 5})
+        self.assert_bb([(3, 3)] + [(1, 1)] * 7, {3, 6})
+
+        # many identical rectangles
+        self.assert_bb([(7, 190)] * 4 + [(190, 7)] * 4, {190, 56})
+
+        # require placing the flipped version of a rectangle
+        self.assert_bb([(1, 100), (5, 96), (4, 5)], {100, 6})
+
+    def test_random(self):
+        for _ in range(5):
+            vals = torch.randint(size=(20, 2), low=1, high=18)
+            sizes = []
+            for j in range(vals.shape[0]):
+                sizes.append((int(vals[j, 0]), int(vals[j, 1])))
+            self.wrap_pack(sizes)
+
+    def test_all_identical(self):
+        sizes = [Rectangle(xsize=61, ysize=82, identifier=1729)] * 3
+        total_size, locations = pack_unique_rectangles(sizes)
+        self.assertEqual(total_size, (61, 82))
+        self.assertEqual(len(locations), 3)
+        for i, (x, y, is_flipped, is_first) in enumerate(locations):
+            self.assertEqual(x, 0)
+            self.assertEqual(y, 0)
+            self.assertFalse(is_flipped)
+            self.assertEqual(is_first, i == 0)
+
+    def test_one_different_id(self):
+        sizes = [Rectangle(xsize=61, ysize=82, identifier=220)] * 3
+        sizes.extend([Rectangle(xsize=61, ysize=82, identifier=284)] * 3)
+        total_size, locations = pack_unique_rectangles(sizes)
+        self.assertEqual(total_size, (82, 122))
+        self.assertEqual(len(locations), 6)
+        for i, (x, y, is_flipped, is_first) in enumerate(locations):
+            self.assertTrue(is_flipped)
+            self.assertEqual(is_first, i % 3 == 0)
+            self.assertEqual(x, 0)
+            if i < 3:
+                self.assertEqual(y, 61)
+            else:
+                self.assertEqual(y, 0)

third_party/AnyBimanual/third_party/pytorch3d/tests/test_transforms.py

@@ -0,0 +1,1350 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import os
+import unittest
+from unittest import mock
+
+import torch
+from pytorch3d.transforms import random_rotations
+from pytorch3d.transforms.se3 import se3_log_map
+from pytorch3d.transforms.so3 import so3_exp_map
+from pytorch3d.transforms.transform3d import (
+    Rotate,
+    RotateAxisAngle,
+    Scale,
+    Transform3d,
+    Translate,
+)
+
+from .common_testing import TestCaseMixin
+
+
+class TestTransform(TestCaseMixin, unittest.TestCase):
+    def setUp(self) -> None:
+        torch.manual_seed(42)
+
+    def test_to(self):
+        tr = Translate(torch.FloatTensor([[1.0, 2.0, 3.0]]))
+        R = torch.FloatTensor([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]])
+        R = Rotate(R)
+        t = Transform3d().compose(R, tr)
+
+        cpu_device = torch.device("cpu")
+
+        cpu_t = t.to("cpu")
+        self.assertEqual(cpu_device, cpu_t.device)
+        self.assertEqual(cpu_device, t.device)
+        self.assertEqual(torch.float32, cpu_t.dtype)
+        self.assertEqual(torch.float32, t.dtype)
+        self.assertIs(t, cpu_t)
+
+        cpu_t = t.to(cpu_device)
+        self.assertEqual(cpu_device, cpu_t.device)
+        self.assertEqual(cpu_device, t.device)
+        self.assertEqual(torch.float32, cpu_t.dtype)
+        self.assertEqual(torch.float32, t.dtype)
+        self.assertIs(t, cpu_t)
+
+        cpu_t = t.to(dtype=torch.float64, device=cpu_device)
+        self.assertEqual(cpu_device, cpu_t.device)
+        self.assertEqual(cpu_device, t.device)
+        self.assertEqual(torch.float64, cpu_t.dtype)
+        self.assertEqual(torch.float32, t.dtype)
+        self.assertIsNot(t, cpu_t)
+
+        cuda_device = torch.device("cuda:0")
+
+        cuda_t = t.to("cuda:0")
+        self.assertEqual(cuda_device, cuda_t.device)
+        self.assertEqual(cpu_device, t.device)
+        self.assertEqual(torch.float32, cuda_t.dtype)
+        self.assertEqual(torch.float32, t.dtype)
+        self.assertIsNot(t, cuda_t)
+
+        cuda_t = t.to(cuda_device)
+        self.assertEqual(cuda_device, cuda_t.device)
+        self.assertEqual(cpu_device, t.device)
+        self.assertEqual(torch.float32, cuda_t.dtype)
+        self.assertEqual(torch.float32, t.dtype)
+        self.assertIsNot(t, cuda_t)
+
+        cuda_t = t.to(dtype=torch.float64, device=cuda_device)
+        self.assertEqual(cuda_device, cuda_t.device)
+        self.assertEqual(cpu_device, t.device)
+        self.assertEqual(torch.float64, cuda_t.dtype)
+        self.assertEqual(torch.float32, t.dtype)
+        self.assertIsNot(t, cuda_t)
+
+        cpu_points = torch.rand(9, 3)
+        cuda_points = cpu_points.cuda()
+        for _ in range(3):
+            t = t.cpu()
+            t.transform_points(cpu_points)
+            t = t.cuda()
+            t.transform_points(cuda_points)
+            t = t.cuda()
+            t = t.cpu()
+
+    def test_dtype_propagation(self):
+        """
+        Check that a given dtype is correctly passed along to child
+        transformations.
+        """
+        # Use at least two dtypes so we avoid only testing on the
+        # default dtype.
+        for dtype in [torch.float32, torch.float64]:
+            R = torch.tensor(
+                [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]],
+                dtype=dtype,
+            )
+            tf = (
+                Transform3d(dtype=dtype)
+                .rotate(R)
+                .rotate_axis_angle(
+                    R[0],
+                    "X",
+                )
+                .translate(3, 2, 1)
+                .scale(0.5)
+            )
+
+            self.assertEqual(tf.dtype, dtype)
+            for inner_tf in tf._transforms:
+                self.assertEqual(inner_tf.dtype, dtype)
+
+            transformed = tf.transform_points(R)
+            self.assertEqual(transformed.dtype, dtype)
+
+    def test_clone(self):
+        """
+        Check that cloned transformations contain different _matrix objects.
+        Also, the clone of a composed translation and rotation has to be
+        the same as composition of clones of translation and rotation.
+        """
+        tr = Translate(torch.FloatTensor([[1.0, 2.0, 3.0]]))
+        R = torch.FloatTensor([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]])
+        R = Rotate(R)
+
+        # check that the _matrix property of clones of
+        # both transforms are different
+        for t in (R, tr):
+            self.assertTrue(t._matrix is not t.clone()._matrix)
+
+        # check that the _transforms lists of composition of R, tr contain
+        # different objects
+        t1 = Transform3d().compose(R, tr)
+        for t, t_clone in (t1._transforms, t1.clone()._transforms):
+            self.assertTrue(t is not t_clone)
+            self.assertTrue(t._matrix is not t_clone._matrix)
+
+        # check that all composed transforms are numerically equivalent
+        t2 = Transform3d().compose(R.clone(), tr.clone())
+        t3 = t1.clone()
+        for t_pair in ((t1, t2), (t1, t3), (t2, t3)):
+            matrix1 = t_pair[0].get_matrix()
+            matrix2 = t_pair[1].get_matrix()
+            self.assertTrue(torch.allclose(matrix1, matrix2))
+
+    def test_init_with_custom_matrix(self):
+        for matrix in (torch.randn(10, 4, 4), torch.randn(4, 4)):
+            t = Transform3d(matrix=matrix)
+            self.assertTrue(t.device == matrix.device)
+            self.assertTrue(t._matrix.dtype == matrix.dtype)
+            self.assertTrue(torch.allclose(t._matrix, matrix.view(t._matrix.shape)))
+
+    def test_init_with_custom_matrix_errors(self):
+        bad_shapes = [[10, 5, 4], [3, 4], [10, 4, 4, 1], [10, 4, 4, 2], [4, 4, 4, 3]]
+        for bad_shape in bad_shapes:
+            matrix = torch.randn(*bad_shape).float()
+            self.assertRaises(ValueError, Transform3d, matrix=matrix)
+
+    def test_get_se3(self):
+        N = 16
+        random_rotations(N)
+        tr = Translate(torch.rand((N, 3)))
+        R = Rotate(random_rotations(N))
+        transform = Transform3d().compose(R, tr)
+        se3_log = transform.get_se3_log()
+        gt_se3_log = se3_log_map(transform.get_matrix())
+        self.assertClose(se3_log, gt_se3_log)
+
+    def test_translate(self):
+        t = Transform3d().translate(1, 2, 3)
+        points = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.5, 0.5, 0.0]]).view(
+            1, 3, 3
+        )
+        normals = torch.tensor(
+            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [1.0, 1.0, 0.0]]
+        ).view(1, 3, 3)
+        points_out = t.transform_points(points)
+        normals_out = t.transform_normals(normals)
+        points_out_expected = torch.tensor(
+            [[2.0, 2.0, 3.0], [1.0, 3.0, 3.0], [1.5, 2.5, 3.0]]
+        ).view(1, 3, 3)
+        normals_out_expected = torch.tensor(
+            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [1.0, 1.0, 0.0]]
+        ).view(1, 3, 3)
+        self.assertTrue(torch.allclose(points_out, points_out_expected))
+        self.assertTrue(torch.allclose(normals_out, normals_out_expected))
+
+    @mock.patch.dict(os.environ, {"PYTORCH3D_CHECK_ROTATION_MATRICES": "1"}, clear=True)
+    def test_rotate_check_rot_valid_on(self):
+        R = so3_exp_map(torch.randn((1, 3)))
+        t = Transform3d().rotate(R)
+        points = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.5, 0.5, 0.0]]).view(
+            1, 3, 3
+        )
+        normals = torch.tensor(
+            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [1.0, 1.0, 0.0]]
+        ).view(1, 3, 3)
+        points_out = t.transform_points(points)
+        normals_out = t.transform_normals(normals)
+        points_out_expected = torch.bmm(points, R)
+        normals_out_expected = torch.bmm(normals, R)
+        self.assertTrue(torch.allclose(points_out, points_out_expected))
+        self.assertTrue(torch.allclose(normals_out, normals_out_expected))
+
+    @mock.patch.dict(os.environ, {"PYTORCH3D_CHECK_ROTATION_MATRICES": "0"}, clear=True)
+    def test_rotate_check_rot_valid_off(self):
+        R = so3_exp_map(torch.randn((1, 3)))
+        t = Transform3d().rotate(R)
+        points = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.5, 0.5, 0.0]]).view(
+            1, 3, 3
+        )
+        normals = torch.tensor(
+            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [1.0, 1.0, 0.0]]
+        ).view(1, 3, 3)
+        points_out = t.transform_points(points)
+        normals_out = t.transform_normals(normals)
+        points_out_expected = torch.bmm(points, R)
+        normals_out_expected = torch.bmm(normals, R)
+        self.assertTrue(torch.allclose(points_out, points_out_expected))
+        self.assertTrue(torch.allclose(normals_out, normals_out_expected))
+
+    def test_scale(self):
+        t = Transform3d().scale(2.0).scale(0.5, 0.25, 1.0)
+        points = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.5, 0.5, 0.0]]).view(
+            1, 3, 3
+        )
+        normals = torch.tensor(
+            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [1.0, 1.0, 0.0]]
+        ).view(1, 3, 3)
+        points_out = t.transform_points(points)
+        normals_out = t.transform_normals(normals)
+        points_out_expected = torch.tensor(
+            [[1.00, 0.00, 0.00], [0.00, 0.50, 0.00], [0.50, 0.25, 0.00]]
+        ).view(1, 3, 3)
+        normals_out_expected = torch.tensor(
+            [[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [1.0, 2.0, 0.0]]
+        ).view(1, 3, 3)
+        self.assertTrue(torch.allclose(points_out, points_out_expected))
+        self.assertTrue(torch.allclose(normals_out, normals_out_expected))
+
+    def test_scale_translate(self):
+        t = Transform3d().scale(2, 1, 3).translate(1, 2, 3)
+        points = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.5, 0.5, 0.0]]).view(
+            1, 3, 3
+        )
+        normals = torch.tensor(
+            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [1.0, 1.0, 0.0]]
+        ).view(1, 3, 3)
+        points_out = t.transform_points(points)
+        normals_out = t.transform_normals(normals)
+        points_out_expected = torch.tensor(
+            [[3.0, 2.0, 3.0], [1.0, 3.0, 3.0], [2.0, 2.5, 3.0]]
+        ).view(1, 3, 3)
+        normals_out_expected = torch.tensor(
+            [[0.5, 0.0, 0.0], [0.0, 1.0, 0.0], [0.5, 1.0, 0.0]]
+        ).view(1, 3, 3)
+        self.assertTrue(torch.allclose(points_out, points_out_expected))
+        self.assertTrue(torch.allclose(normals_out, normals_out_expected))
+
+    def test_rotate_axis_angle(self):
+        t = Transform3d().rotate_axis_angle(90.0, axis="Z")
+        points = torch.tensor([[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 1.0, 1.0]]).view(
+            1, 3, 3
+        )
+        normals = torch.tensor(
+            [[1.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 0.0, 0.0]]
+        ).view(1, 3, 3)
+        points_out = t.transform_points(points)
+        normals_out = t.transform_normals(normals)
+        points_out_expected = torch.tensor(
+            [[0.0, 0.0, 0.0], [-1.0, 0.0, 0.0], [-1.0, 0.0, 1.0]]
+        ).view(1, 3, 3)
+        normals_out_expected = torch.tensor(
+            [[0.0, 1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 1.0, 0.0]]
+        ).view(1, 3, 3)
+        self.assertTrue(torch.allclose(points_out, points_out_expected, atol=1e-7))
+        self.assertTrue(torch.allclose(normals_out, normals_out_expected, atol=1e-7))
+
+    def test_transform_points_fail(self):
+        t1 = Scale(0.1, 0.1, 0.1)
+        P = 7
+        with self.assertRaises(ValueError):
+            t1.transform_points(torch.randn(P))
+
+    def test_compose_fail(self):
+        # Only composing Transform3d objects is possible
+        t1 = Scale(0.1, 0.1, 0.1)
+        with self.assertRaises(ValueError):
+            t1.compose(torch.randn(100))
+
+    def test_transform_points_eps(self):
+        t1 = Transform3d()
+        persp_proj = [
+            [
+                [1.0, 0.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 1.0],
+                [0.0, 0.0, 1.0, 0.0],
+            ]
+        ]
+        t1._matrix = torch.FloatTensor(persp_proj)
+        points = torch.tensor(
+            [[0.0, 1.0, 0.0], [0.0, 0.0, 1e-5], [-1.0, 0.0, 1e-5]]
+        ).view(
+            1, 3, 3
+        )  # a set of points with z-coord very close to 0
+
+        proj = t1.transform_points(points)
+        proj_eps = t1.transform_points(points, eps=1e-4)
+
+        self.assertTrue(not bool(torch.isfinite(proj.sum())))
+        self.assertTrue(bool(torch.isfinite(proj_eps.sum())))
+
+    def test_inverse(self, batch_size=5):
+        device = torch.device("cuda:0")
+
+        # generate a random chain of transforms
+        for _ in range(10):  # 10 different tries
+
+            # list of transform matrices
+            ts = []
+
+            for i in range(10):
+                choice = float(torch.rand(1))
+                if choice <= 1.0 / 3.0:
+                    t_ = Translate(
+                        torch.randn(
+                            (batch_size, 3), dtype=torch.float32, device=device
+                        ),
+                        device=device,
+                    )
+                elif choice <= 2.0 / 3.0:
+                    t_ = Rotate(
+                        so3_exp_map(
+                            torch.randn(
+                                (batch_size, 3), dtype=torch.float32, device=device
+                            )
+                        ),
+                        device=device,
+                    )
+                else:
+                    rand_t = torch.randn(
+                        (batch_size, 3), dtype=torch.float32, device=device
+                    )
+                    rand_t = rand_t.sign() * torch.clamp(rand_t.abs(), 0.2)
+                    t_ = Scale(rand_t, device=device)
+                ts.append(t_._matrix.clone())
+
+                if i == 0:
+                    t = t_
+                else:
+                    t = t.compose(t_)
+
+            # generate the inverse transformation in several possible ways
+            m1 = t.inverse(invert_composed=True).get_matrix()
+            m2 = t.inverse(invert_composed=True)._matrix
+            m3 = t.inverse(invert_composed=False).get_matrix()
+            m4 = t.get_matrix().inverse()
+
+            # compute the inverse explicitly ...
+            m5 = torch.eye(4, dtype=torch.float32, device=device)
+            m5 = m5[None].repeat(batch_size, 1, 1)
+            for t_ in ts:
+                m5 = torch.bmm(torch.inverse(t_), m5)
+
+            # assert all same
+            for m in (m1, m2, m3, m4):
+                self.assertTrue(torch.allclose(m, m5, atol=1e-3))
+
+    def _check_indexed_transforms(self, t3d, t3d_selected, indices):
+        t3d_matrix = t3d.get_matrix()
+        t3d_selected_matrix = t3d_selected.get_matrix()
+        for order_index, selected_index in indices:
+            self.assertClose(
+                t3d_matrix[selected_index], t3d_selected_matrix[order_index]
+            )
+
+    def test_get_item(self, batch_size=5):
+        device = torch.device("cuda:0")
+
+        matrices = torch.randn(
+            size=[batch_size, 4, 4], device=device, dtype=torch.float32
+        )
+
+        # init the Transforms3D class
+        t3d = Transform3d(matrix=matrices)
+
+        # int index
+        index = 1
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), 1)
+        self._check_indexed_transforms(t3d, t3d_selected, [(0, 1)])
+
+        # negative int index
+        index = -1
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), 1)
+        self._check_indexed_transforms(t3d, t3d_selected, [(0, -1)])
+
+        # list index
+        index = [1, 2]
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), len(index))
+        self._check_indexed_transforms(t3d, t3d_selected, enumerate(index))
+
+        # empty list index
+        index = []
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), 0)
+        self.assertEqual(t3d_selected.get_matrix().nelement(), 0)
+
+        # slice index
+        index = slice(0, 2, 1)
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), 2)
+        self._check_indexed_transforms(t3d, t3d_selected, [(0, 0), (1, 1)])
+
+        # empty slice index
+        index = slice(0, 0, 1)
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), 0)
+        self.assertEqual(t3d_selected.get_matrix().nelement(), 0)
+
+        # bool tensor
+        index = (torch.rand(batch_size) > 0.5).to(device)
+        index[:2] = True  # make sure smth is selected
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), index.sum())
+        self._check_indexed_transforms(
+            t3d,
+            t3d_selected,
+            zip(
+                torch.arange(index.sum()),
+                torch.nonzero(index, as_tuple=False).squeeze(),
+            ),
+        )
+
+        # all false bool tensor
+        index = torch.zeros(batch_size).bool()
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), 0)
+        self.assertEqual(t3d_selected.get_matrix().nelement(), 0)
+
+        # int tensor
+        index = torch.tensor([1, 2], dtype=torch.int64, device=device)
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), index.numel())
+        self._check_indexed_transforms(t3d, t3d_selected, enumerate(index.tolist()))
+
+        # negative int tensor
+        index = -(torch.tensor([1, 2], dtype=torch.int64, device=device))
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), index.numel())
+        self._check_indexed_transforms(t3d, t3d_selected, enumerate(index.tolist()))
+
+        # invalid index
+        for invalid_index in (
+            torch.tensor([1, 0, 1], dtype=torch.float32, device=device),  # float tensor
+            1.2,  # float index
+        ):
+            with self.assertRaises(IndexError):
+                t3d_selected = t3d[invalid_index]
+
+    def test_stack(self):
+        rotations = random_rotations(3)
+        transform3 = Transform3d().rotate(rotations).translate(torch.full((3, 3), 0.3))
+        transform1 = Scale(37)
+        transform4 = transform1.stack(transform3)
+        self.assertEqual(len(transform1), 1)
+        self.assertEqual(len(transform3), 3)
+        self.assertEqual(len(transform4), 4)
+        self.assertClose(
+            transform4.get_matrix(),
+            torch.cat([transform1.get_matrix(), transform3.get_matrix()]),
+        )
+        points = torch.rand(4, 5, 3)
+        new_points_expect = torch.cat(
+            [
+                transform1.transform_points(points[:1]),
+                transform3.transform_points(points[1:]),
+            ]
+        )
+        new_points = transform4.transform_points(points)
+        self.assertClose(new_points, new_points_expect)
+
+
+class TestTranslate(unittest.TestCase):
+    def test_python_scalar(self):
+        t = Translate(0.2, 0.3, 0.4)
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0, 0.0, 0.0, 0],
+                    [0.0, 1.0, 0.0, 0],
+                    [0.0, 0.0, 1.0, 0],
+                    [0.2, 0.3, 0.4, 1],
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_torch_scalar(self):
+        x = torch.tensor(0.2)
+        y = torch.tensor(0.3)
+        z = torch.tensor(0.4)
+        t = Translate(x, y, z)
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0, 0.0, 0.0, 0],
+                    [0.0, 1.0, 0.0, 0],
+                    [0.0, 0.0, 1.0, 0],
+                    [0.2, 0.3, 0.4, 1],
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_mixed_scalars(self):
+        x = 0.2
+        y = torch.tensor(0.3)
+        z = 0.4
+        t = Translate(x, y, z)
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0, 0.0, 0.0, 0],
+                    [0.0, 1.0, 0.0, 0],
+                    [0.0, 0.0, 1.0, 0],
+                    [0.2, 0.3, 0.4, 1],
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_torch_scalar_grads(self):
+        # Make sure backprop works if we give torch scalars
+        x = torch.tensor(0.2, requires_grad=True)
+        y = torch.tensor(0.3, requires_grad=True)
+        z = torch.tensor(0.4)
+        t = Translate(x, y, z)
+        t._matrix.sum().backward()
+        self.assertTrue(hasattr(x, "grad"))
+        self.assertTrue(hasattr(y, "grad"))
+        self.assertTrue(torch.allclose(x.grad, x.new_ones(x.shape)))
+        self.assertTrue(torch.allclose(y.grad, y.new_ones(y.shape)))
+
+    def test_torch_vectors(self):
+        x = torch.tensor([0.2, 2.0])
+        y = torch.tensor([0.3, 3.0])
+        z = torch.tensor([0.4, 4.0])
+        t = Translate(x, y, z)
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0, 0.0, 0.0, 0],
+                    [0.0, 1.0, 0.0, 0],
+                    [0.0, 0.0, 1.0, 0],
+                    [0.2, 0.3, 0.4, 1],
+                ],
+                [
+                    [1.0, 0.0, 0.0, 0],
+                    [0.0, 1.0, 0.0, 0],
+                    [0.0, 0.0, 1.0, 0],
+                    [2.0, 3.0, 4.0, 1],
+                ],
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_vector_broadcast(self):
+        x = torch.tensor([0.2, 2.0])
+        y = torch.tensor([0.3, 3.0])
+        z = torch.tensor([0.4])
+        t = Translate(x, y, z)
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0, 0.0, 0.0, 0],
+                    [0.0, 1.0, 0.0, 0],
+                    [0.0, 0.0, 1.0, 0],
+                    [0.2, 0.3, 0.4, 1],
+                ],
+                [
+                    [1.0, 0.0, 0.0, 0],
+                    [0.0, 1.0, 0.0, 0],
+                    [0.0, 0.0, 1.0, 0],
+                    [2.0, 3.0, 0.4, 1],
+                ],
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_bad_broadcast(self):
+        x = torch.tensor([0.2, 2.0, 20.0])
+        y = torch.tensor([0.3, 3.0])
+        z = torch.tensor([0.4])
+        with self.assertRaises(ValueError):
+            Translate(x, y, z)
+
+    def test_mixed_broadcast(self):
+        x = 0.2
+        y = torch.tensor(0.3)
+        z = torch.tensor([0.4, 4.0])
+        t = Translate(x, y, z)
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0, 0.0, 0.0, 0],
+                    [0.0, 1.0, 0.0, 0],
+                    [0.0, 0.0, 1.0, 0],
+                    [0.2, 0.3, 0.4, 1],
+                ],
+                [
+                    [1.0, 0.0, 0.0, 0],
+                    [0.0, 1.0, 0.0, 0],
+                    [0.0, 0.0, 1.0, 0],
+                    [0.2, 0.3, 4.0, 1],
+                ],
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_mixed_broadcast_grad(self):
+        x = 0.2
+        y = torch.tensor(0.3, requires_grad=True)
+        z = torch.tensor([0.4, 4.0], requires_grad=True)
+        t = Translate(x, y, z)
+        t._matrix.sum().backward()
+        self.assertTrue(hasattr(y, "grad"))
+        self.assertTrue(hasattr(z, "grad"))
+        y_grad = torch.tensor(2.0)
+        z_grad = torch.tensor([1.0, 1.0])
+        self.assertEqual(y.grad.shape, y_grad.shape)
+        self.assertEqual(z.grad.shape, z_grad.shape)
+        self.assertTrue(torch.allclose(y.grad, y_grad))
+        self.assertTrue(torch.allclose(z.grad, z_grad))
+
+    def test_matrix(self):
+        xyz = torch.tensor([[0.2, 0.3, 0.4], [2.0, 3.0, 4.0]])
+        t = Translate(xyz)
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0, 0.0, 0.0, 0],
+                    [0.0, 1.0, 0.0, 0],
+                    [0.0, 0.0, 1.0, 0],
+                    [0.2, 0.3, 0.4, 1],
+                ],
+                [
+                    [1.0, 0.0, 0.0, 0],
+                    [0.0, 1.0, 0.0, 0],
+                    [0.0, 0.0, 1.0, 0],
+                    [2.0, 3.0, 4.0, 1],
+                ],
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_matrix_extra_args(self):
+        xyz = torch.tensor([[0.2, 0.3, 0.4], [2.0, 3.0, 4.0]])
+        with self.assertRaises(ValueError):
+            Translate(xyz, xyz[:, 1], xyz[:, 2])
+
+    def test_inverse(self):
+        xyz = torch.tensor([[0.2, 0.3, 0.4], [2.0, 3.0, 4.0]])
+        t = Translate(xyz)
+        im = t.inverse()._matrix
+        im_2 = t._matrix.inverse()
+        im_comp = t.get_matrix().inverse()
+        self.assertTrue(torch.allclose(im, im_comp))
+        self.assertTrue(torch.allclose(im, im_2))
+
+    def test_get_item(self, batch_size=5):
+        device = torch.device("cuda:0")
+        xyz = torch.randn(size=[batch_size, 3], device=device, dtype=torch.float32)
+        t3d = Translate(xyz)
+        index = 1
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), 1)
+        self.assertIsInstance(t3d_selected, Translate)
+
+
+class TestScale(unittest.TestCase):
+    def test_single_python_scalar(self):
+        t = Scale(0.1)
+        matrix = torch.tensor(
+            [
+                [
+                    [0.1, 0.0, 0.0, 0.0],
+                    [0.0, 0.1, 0.0, 0.0],
+                    [0.0, 0.0, 0.1, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_single_torch_scalar(self):
+        t = Scale(torch.tensor(0.1))
+        matrix = torch.tensor(
+            [
+                [
+                    [0.1, 0.0, 0.0, 0.0],
+                    [0.0, 0.1, 0.0, 0.0],
+                    [0.0, 0.0, 0.1, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_single_vector(self):
+        t = Scale(torch.tensor([0.1, 0.2]))
+        matrix = torch.tensor(
+            [
+                [
+                    [0.1, 0.0, 0.0, 0.0],
+                    [0.0, 0.1, 0.0, 0.0],
+                    [0.0, 0.0, 0.1, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ],
+                [
+                    [0.2, 0.0, 0.0, 0.0],
+                    [0.0, 0.2, 0.0, 0.0],
+                    [0.0, 0.0, 0.2, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ],
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_single_matrix(self):
+        xyz = torch.tensor([[0.1, 0.2, 0.3], [1.0, 2.0, 3.0]])
+        t = Scale(xyz)
+        matrix = torch.tensor(
+            [
+                [
+                    [0.1, 0.0, 0.0, 0.0],
+                    [0.0, 0.2, 0.0, 0.0],
+                    [0.0, 0.0, 0.3, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ],
+                [
+                    [1.0, 0.0, 0.0, 0.0],
+                    [0.0, 2.0, 0.0, 0.0],
+                    [0.0, 0.0, 3.0, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ],
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_three_python_scalar(self):
+        t = Scale(0.1, 0.2, 0.3)
+        matrix = torch.tensor(
+            [
+                [
+                    [0.1, 0.0, 0.0, 0.0],
+                    [0.0, 0.2, 0.0, 0.0],
+                    [0.0, 0.0, 0.3, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_three_torch_scalar(self):
+        t = Scale(torch.tensor(0.1), torch.tensor(0.2), torch.tensor(0.3))
+        matrix = torch.tensor(
+            [
+                [
+                    [0.1, 0.0, 0.0, 0.0],
+                    [0.0, 0.2, 0.0, 0.0],
+                    [0.0, 0.0, 0.3, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_three_mixed_scalar(self):
+        t = Scale(torch.tensor(0.1), 0.2, torch.tensor(0.3))
+        matrix = torch.tensor(
+            [
+                [
+                    [0.1, 0.0, 0.0, 0.0],
+                    [0.0, 0.2, 0.0, 0.0],
+                    [0.0, 0.0, 0.3, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_three_vector_broadcast(self):
+        x = torch.tensor([0.1])
+        y = torch.tensor([0.2, 2.0])
+        z = torch.tensor([0.3, 3.0])
+        t = Scale(x, y, z)
+        matrix = torch.tensor(
+            [
+                [
+                    [0.1, 0.0, 0.0, 0.0],
+                    [0.0, 0.2, 0.0, 0.0],
+                    [0.0, 0.0, 0.3, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ],
+                [
+                    [0.1, 0.0, 0.0, 0.0],
+                    [0.0, 2.0, 0.0, 0.0],
+                    [0.0, 0.0, 3.0, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ],
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_three_mixed_broadcast_grad(self):
+        x = 0.1
+        y = torch.tensor(0.2, requires_grad=True)
+        z = torch.tensor([0.3, 3.0], requires_grad=True)
+        t = Scale(x, y, z)
+        matrix = torch.tensor(
+            [
+                [
+                    [0.1, 0.0, 0.0, 0.0],
+                    [0.0, 0.2, 0.0, 0.0],
+                    [0.0, 0.0, 0.3, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ],
+                [
+                    [0.1, 0.0, 0.0, 0.0],
+                    [0.0, 0.2, 0.0, 0.0],
+                    [0.0, 0.0, 3.0, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ],
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+        t._matrix.sum().backward()
+        self.assertTrue(hasattr(y, "grad"))
+        self.assertTrue(hasattr(z, "grad"))
+        y_grad = torch.tensor(2.0)
+        z_grad = torch.tensor([1.0, 1.0])
+        self.assertTrue(torch.allclose(y.grad, y_grad))
+        self.assertTrue(torch.allclose(z.grad, z_grad))
+
+    def test_inverse(self):
+        x = torch.tensor([0.1])
+        y = torch.tensor([0.2, 2.0])
+        z = torch.tensor([0.3, 3.0])
+        t = Scale(x, y, z)
+        im = t.inverse()._matrix
+        im_2 = t._matrix.inverse()
+        im_comp = t.get_matrix().inverse()
+        self.assertTrue(torch.allclose(im, im_comp))
+        self.assertTrue(torch.allclose(im, im_2))
+
+    def test_get_item(self, batch_size=5):
+        device = torch.device("cuda:0")
+        s = torch.randn(size=[batch_size, 3], device=device, dtype=torch.float32)
+        t3d = Scale(s)
+        index = 1
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), 1)
+        self.assertIsInstance(t3d_selected, Scale)
+
+
+class TestTransformBroadcast(unittest.TestCase):
+    def test_broadcast_transform_points(self):
+        t1 = Scale(0.1, 0.1, 0.1)
+        N = 10
+        P = 7
+        M = 20
+        x = torch.tensor([0.2] * N)
+        y = torch.tensor([0.3] * N)
+        z = torch.tensor([0.4] * N)
+        tN = Translate(x, y, z)
+        p1 = t1.transform_points(torch.randn(P, 3))
+        self.assertTrue(p1.shape == (P, 3))
+        p2 = t1.transform_points(torch.randn(1, P, 3))
+        self.assertTrue(p2.shape == (1, P, 3))
+        p3 = t1.transform_points(torch.randn(M, P, 3))
+        self.assertTrue(p3.shape == (M, P, 3))
+        p4 = tN.transform_points(torch.randn(P, 3))
+        self.assertTrue(p4.shape == (N, P, 3))
+        p5 = tN.transform_points(torch.randn(1, P, 3))
+        self.assertTrue(p5.shape == (N, P, 3))
+
+    def test_broadcast_transform_normals(self):
+        t1 = Scale(0.1, 0.1, 0.1)
+        N = 10
+        P = 7
+        M = 20
+        x = torch.tensor([0.2] * N)
+        y = torch.tensor([0.3] * N)
+        z = torch.tensor([0.4] * N)
+        tN = Translate(x, y, z)
+        p1 = t1.transform_normals(torch.randn(P, 3))
+        self.assertTrue(p1.shape == (P, 3))
+        p2 = t1.transform_normals(torch.randn(1, P, 3))
+        self.assertTrue(p2.shape == (1, P, 3))
+        p3 = t1.transform_normals(torch.randn(M, P, 3))
+        self.assertTrue(p3.shape == (M, P, 3))
+        p4 = tN.transform_normals(torch.randn(P, 3))
+        self.assertTrue(p4.shape == (N, P, 3))
+        p5 = tN.transform_normals(torch.randn(1, P, 3))
+        self.assertTrue(p5.shape == (N, P, 3))
+
+    def test_broadcast_compose(self):
+        t1 = Scale(0.1, 0.1, 0.1)
+        N = 10
+        scale_n = torch.tensor([0.3] * N)
+        tN = Scale(scale_n)
+        t1N = t1.compose(tN)
+        self.assertTrue(t1._matrix.shape == (1, 4, 4))
+        self.assertTrue(tN._matrix.shape == (N, 4, 4))
+        self.assertTrue(t1N.get_matrix().shape == (N, 4, 4))
+        t11 = t1.compose(t1)
+        self.assertTrue(t11.get_matrix().shape == (1, 4, 4))
+
+    def test_broadcast_compose_fail(self):
+        # Cannot compose two transforms which have batch dimensions N and M
+        # other than the case where either N or M is 1
+        N = 10
+        M = 20
+        scale_n = torch.tensor([0.3] * N)
+        tN = Scale(scale_n)
+        x = torch.tensor([0.2] * M)
+        y = torch.tensor([0.3] * M)
+        z = torch.tensor([0.4] * M)
+        tM = Translate(x, y, z)
+        t = tN.compose(tM)
+        with self.assertRaises(ValueError):
+            t.get_matrix()
+
+    def test_multiple_broadcast_compose(self):
+        t1 = Scale(0.1, 0.1, 0.1)
+        t2 = Scale(0.2, 0.2, 0.2)
+        N = 10
+        scale_n = torch.tensor([0.3] * N)
+        tN = Scale(scale_n)
+        t1N2 = t1.compose(tN.compose(t2))
+        composed_mat = t1N2.get_matrix()
+        self.assertTrue(composed_mat.shape == (N, 4, 4))
+        expected_mat = torch.eye(3, dtype=torch.float32) * 0.3 * 0.2 * 0.1
+        self.assertTrue(torch.allclose(composed_mat[0, :3, :3], expected_mat))
+
+
+class TestRotate(unittest.TestCase):
+    def test_single_matrix(self):
+        R = torch.eye(3)
+        t = Rotate(R)
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0, 0.0, 0.0, 0.0],
+                    [0.0, 1.0, 0.0, 0.0],
+                    [0.0, 0.0, 1.0, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix))
+
+    def test_invalid_dimensions(self):
+        R = torch.eye(4)
+        with self.assertRaises(ValueError):
+            Rotate(R)
+
+    def test_inverse(self, batch_size=5):
+        device = torch.device("cuda:0")
+        log_rot = torch.randn((batch_size, 3), dtype=torch.float32, device=device)
+        R = so3_exp_map(log_rot)
+        t = Rotate(R)
+        im = t.inverse()._matrix
+        im_2 = t._matrix.inverse()
+        im_comp = t.get_matrix().inverse()
+        self.assertTrue(torch.allclose(im, im_comp, atol=1e-4))
+        self.assertTrue(torch.allclose(im, im_2, atol=1e-4))
+
+    def test_get_item(self, batch_size=5):
+        device = torch.device("cuda:0")
+        r = random_rotations(batch_size, dtype=torch.float32, device=device)
+        t3d = Rotate(r)
+        index = 1
+        t3d_selected = t3d[index]
+        self.assertEqual(len(t3d_selected), 1)
+        self.assertIsInstance(t3d_selected, Rotate)
+
+
+class TestRotateAxisAngle(unittest.TestCase):
+    def test_rotate_x_python_scalar(self):
+        t = RotateAxisAngle(angle=90, axis="X")
+        # fmt: off
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0,  0.0, 0.0, 0.0],  # noqa: E241, E201
+                    [0.0,  0.0, 1.0, 0.0],  # noqa: E241, E201
+                    [0.0, -1.0, 0.0, 0.0],  # noqa: E241, E201
+                    [0.0,  0.0, 0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        # fmt: on
+        points = torch.tensor([0.0, 1.0, 0.0])[None, None, :]  # (1, 1, 3)
+        transformed_points = t.transform_points(points)
+        expected_points = torch.tensor([0.0, 0.0, 1.0])
+        self.assertTrue(
+            torch.allclose(transformed_points.squeeze(), expected_points, atol=1e-7)
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix, atol=1e-7))
+
+    def test_rotate_x_torch_scalar(self):
+        angle = torch.tensor(90.0)
+        t = RotateAxisAngle(angle=angle, axis="X")
+        # fmt: off
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0,  0.0, 0.0, 0.0],  # noqa: E241, E201
+                    [0.0,  0.0, 1.0, 0.0],  # noqa: E241, E201
+                    [0.0, -1.0, 0.0, 0.0],  # noqa: E241, E201
+                    [0.0,  0.0, 0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        # fmt: on
+        points = torch.tensor([0.0, 1.0, 0.0])[None, None, :]  # (1, 1, 3)
+        transformed_points = t.transform_points(points)
+        expected_points = torch.tensor([0.0, 0.0, 1.0])
+        self.assertTrue(
+            torch.allclose(transformed_points.squeeze(), expected_points, atol=1e-7)
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix, atol=1e-7))
+
+    def test_rotate_x_torch_tensor(self):
+        angle = torch.tensor([0, 45.0, 90.0])  # (N)
+        t = RotateAxisAngle(angle=angle, axis="X")
+        r2_i = 1 / math.sqrt(2)
+        r2_2 = math.sqrt(2) / 2
+        # fmt: off
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0, 0.0, 0.0, 0.0],
+                    [0.0, 1.0, 0.0, 0.0],
+                    [0.0, 0.0, 1.0, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ],
+                [
+                    [1.0,   0.0,  0.0, 0.0],  # noqa: E241, E201
+                    [0.0,  r2_2, r2_i, 0.0],  # noqa: E241, E201
+                    [0.0, -r2_i, r2_2, 0.0],  # noqa: E241, E201
+                    [0.0,   0.0,  0.0, 1.0],  # noqa: E241, E201
+                ],
+                [
+                    [1.0,  0.0, 0.0,  0.0],   # noqa: E241, E201
+                    [0.0,  0.0, 1.0,  0.0],   # noqa: E241, E201
+                    [0.0, -1.0, 0.0,  0.0],   # noqa: E241, E201
+                    [0.0,  0.0, 0.0,  1.0],   # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        # fmt: on
+        self.assertTrue(torch.allclose(t._matrix, matrix, atol=1e-7))
+        angle = angle
+        t = RotateAxisAngle(angle=angle, axis="X")
+        self.assertTrue(torch.allclose(t._matrix, matrix, atol=1e-7))
+
+    def test_rotate_y_python_scalar(self):
+        t = RotateAxisAngle(angle=90, axis="Y")
+        # fmt: off
+        matrix = torch.tensor(
+            [
+                [
+                    [0.0, 0.0, -1.0, 0.0],  # noqa: E241, E201
+                    [0.0, 1.0,  0.0, 0.0],  # noqa: E241, E201
+                    [1.0, 0.0,  0.0, 0.0],  # noqa: E241, E201
+                    [0.0, 0.0,  0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        # fmt: on
+        points = torch.tensor([1.0, 0.0, 0.0])[None, None, :]  # (1, 1, 3)
+        transformed_points = t.transform_points(points)
+        expected_points = torch.tensor([0.0, 0.0, -1.0])
+        self.assertTrue(
+            torch.allclose(transformed_points.squeeze(), expected_points, atol=1e-7)
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix, atol=1e-7))
+
+    def test_rotate_y_torch_scalar(self):
+        """
+        Test rotation about Y axis. With a right hand coordinate system this
+        should result in a vector pointing along the x-axis being rotated to
+        point along the negative z axis.
+        """
+        angle = torch.tensor(90.0)
+        t = RotateAxisAngle(angle=angle, axis="Y")
+        # fmt: off
+        matrix = torch.tensor(
+            [
+                [
+                    [0.0, 0.0, -1.0, 0.0],  # noqa: E241, E201
+                    [0.0, 1.0,  0.0, 0.0],  # noqa: E241, E201
+                    [1.0, 0.0,  0.0, 0.0],  # noqa: E241, E201
+                    [0.0, 0.0,  0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        # fmt: on
+        points = torch.tensor([1.0, 0.0, 0.0])[None, None, :]  # (1, 1, 3)
+        transformed_points = t.transform_points(points)
+        expected_points = torch.tensor([0.0, 0.0, -1.0])
+        self.assertTrue(
+            torch.allclose(transformed_points.squeeze(), expected_points, atol=1e-7)
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix, atol=1e-7))
+
+    def test_rotate_y_torch_tensor(self):
+        angle = torch.tensor([0, 45.0, 90.0])
+        t = RotateAxisAngle(angle=angle, axis="Y")
+        r2_i = 1 / math.sqrt(2)
+        r2_2 = math.sqrt(2) / 2
+        # fmt: off
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0, 0.0, 0.0, 0.0],
+                    [0.0, 1.0, 0.0, 0.0],
+                    [0.0, 0.0, 1.0, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ],
+                [
+                    [r2_2,  0.0, -r2_i, 0.0],  # noqa: E241, E201
+                    [ 0.0,  1.0,   0.0, 0.0],  # noqa: E241, E201
+                    [r2_i,  0.0,  r2_2, 0.0],  # noqa: E241, E201
+                    [ 0.0,  0.0,   0.0, 1.0],  # noqa: E241, E201
+                ],
+                [
+                    [0.0,  0.0, -1.0, 0.0],  # noqa: E241, E201
+                    [0.0,  1.0,  0.0, 0.0],  # noqa: E241, E201
+                    [1.0,  0.0,  0.0, 0.0],  # noqa: E241, E201
+                    [0.0,  0.0,  0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        # fmt: on
+        self.assertTrue(torch.allclose(t._matrix, matrix, atol=1e-7))
+
+    def test_rotate_z_python_scalar(self):
+        t = RotateAxisAngle(angle=90, axis="Z")
+        # fmt: off
+        matrix = torch.tensor(
+            [
+                [
+                    [ 0.0, 1.0, 0.0, 0.0],  # noqa: E241, E201
+                    [-1.0, 0.0, 0.0, 0.0],  # noqa: E241, E201
+                    [ 0.0, 0.0, 1.0, 0.0],  # noqa: E241, E201
+                    [ 0.0, 0.0, 0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        # fmt: on
+        points = torch.tensor([1.0, 0.0, 0.0])[None, None, :]  # (1, 1, 3)
+        transformed_points = t.transform_points(points)
+        expected_points = torch.tensor([0.0, 1.0, 0.0])
+        self.assertTrue(
+            torch.allclose(transformed_points.squeeze(), expected_points, atol=1e-7)
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix, atol=1e-7))
+
+    def test_rotate_z_torch_scalar(self):
+        angle = torch.tensor(90.0)
+        t = RotateAxisAngle(angle=angle, axis="Z")
+        # fmt: off
+        matrix = torch.tensor(
+            [
+                [
+                    [ 0.0, 1.0, 0.0, 0.0],  # noqa: E241, E201
+                    [-1.0, 0.0, 0.0, 0.0],  # noqa: E241, E201
+                    [ 0.0, 0.0, 1.0, 0.0],  # noqa: E241, E201
+                    [ 0.0, 0.0, 0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        # fmt: on
+        points = torch.tensor([1.0, 0.0, 0.0])[None, None, :]  # (1, 1, 3)
+        transformed_points = t.transform_points(points)
+        expected_points = torch.tensor([0.0, 1.0, 0.0])
+        self.assertTrue(
+            torch.allclose(transformed_points.squeeze(), expected_points, atol=1e-7)
+        )
+        self.assertTrue(torch.allclose(t._matrix, matrix, atol=1e-7))
+
+    def test_rotate_z_torch_tensor(self):
+        angle = torch.tensor([0, 45.0, 90.0])
+        t = RotateAxisAngle(angle=angle, axis="Z")
+        r2_i = 1 / math.sqrt(2)
+        r2_2 = math.sqrt(2) / 2
+        # fmt: off
+        matrix = torch.tensor(
+            [
+                [
+                    [1.0, 0.0, 0.0, 0.0],
+                    [0.0, 1.0, 0.0, 0.0],
+                    [0.0, 0.0, 1.0, 0.0],
+                    [0.0, 0.0, 0.0, 1.0],
+                ],
+                [
+                    [ r2_2,   r2_i,  0.0, 0.0],  # noqa: E241, E201
+                    [-r2_i,   r2_2,  0.0, 0.0],  # noqa: E241, E201
+                    [  0.0,    0.0,  1.0, 0.0],  # noqa: E241, E201
+                    [  0.0,    0.0,  0.0, 1.0],  # noqa: E241, E201
+                ],
+                [
+                    [ 0.0,  1.0, 0.0, 0.0],  # noqa: E241, E201
+                    [-1.0,  0.0, 0.0, 0.0],  # noqa: E241, E201
+                    [ 0.0,  0.0, 1.0, 0.0],  # noqa: E241, E201
+                    [ 0.0,  0.0, 0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        # fmt: on
+        self.assertTrue(torch.allclose(t._matrix, matrix, atol=1e-7))
+
+    def test_rotate_compose_x_y_z(self):
+        angle = torch.tensor(90.0)
+        t1 = RotateAxisAngle(angle=angle, axis="X")
+        t2 = RotateAxisAngle(angle=angle, axis="Y")
+        t3 = RotateAxisAngle(angle=angle, axis="Z")
+        t = t1.compose(t2, t3)
+        # fmt: off
+        matrix1 = torch.tensor(
+            [
+                [
+                    [1.0,  0.0, 0.0, 0.0],  # noqa: E241, E201
+                    [0.0,  0.0, 1.0, 0.0],  # noqa: E241, E201
+                    [0.0, -1.0, 0.0, 0.0],  # noqa: E241, E201
+                    [0.0,  0.0, 0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        matrix2 = torch.tensor(
+            [
+                [
+                    [0.0, 0.0, -1.0, 0.0],  # noqa: E241, E201
+                    [0.0, 1.0,  0.0, 0.0],  # noqa: E241, E201
+                    [1.0, 0.0,  0.0, 0.0],  # noqa: E241, E201
+                    [0.0, 0.0,  0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        matrix3 = torch.tensor(
+            [
+                [
+                    [ 0.0, 1.0, 0.0, 0.0],  # noqa: E241, E201
+                    [-1.0, 0.0, 0.0, 0.0],  # noqa: E241, E201
+                    [ 0.0, 0.0, 1.0, 0.0],  # noqa: E241, E201
+                    [ 0.0, 0.0, 0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        # fmt: on
+        # order of transforms is t1 -> t2
+        matrix = torch.matmul(matrix1, torch.matmul(matrix2, matrix3))
+        composed_matrix = t.get_matrix()
+        self.assertTrue(torch.allclose(composed_matrix, matrix, atol=1e-7))
+
+    def test_rotate_angle_radians(self):
+        t = RotateAxisAngle(angle=math.pi / 2, degrees=False, axis="Z")
+        # fmt: off
+        matrix = torch.tensor(
+            [
+                [
+                    [ 0.0, 1.0, 0.0, 0.0],  # noqa: E241, E201
+                    [-1.0, 0.0, 0.0, 0.0],  # noqa: E241, E201
+                    [ 0.0, 0.0, 1.0, 0.0],  # noqa: E241, E201
+                    [ 0.0, 0.0, 0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        # fmt: on
+        self.assertTrue(torch.allclose(t._matrix, matrix, atol=1e-7))
+
+    def test_lower_case_axis(self):
+        t = RotateAxisAngle(angle=90.0, axis="z")
+        # fmt: off
+        matrix = torch.tensor(
+            [
+                [
+                    [ 0.0, 1.0, 0.0, 0.0],  # noqa: E241, E201
+                    [-1.0, 0.0, 0.0, 0.0],  # noqa: E241, E201
+                    [ 0.0, 0.0, 1.0, 0.0],  # noqa: E241, E201
+                    [ 0.0, 0.0, 0.0, 1.0],  # noqa: E241, E201
+                ]
+            ],
+            dtype=torch.float32,
+        )
+        # fmt: on
+        self.assertTrue(torch.allclose(t._matrix, matrix, atol=1e-7))
+
+    def test_axis_fail(self):
+        with self.assertRaises(ValueError):
+            RotateAxisAngle(angle=90.0, axis="P")
+
+    def test_rotate_angle_fail(self):
+        angle = torch.tensor([[0, 45.0, 90.0], [0, 45.0, 90.0]])
+        with self.assertRaises(ValueError):
+            RotateAxisAngle(angle=angle, axis="X")

third_party/AnyBimanual/third_party/pytorch3d/tests/test_vert_align.py

@@ -0,0 +1,194 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import unittest
+
+import torch
+import torch.nn.functional as F
+from pytorch3d.ops.vert_align import vert_align
+from pytorch3d.structures.meshes import Meshes
+from pytorch3d.structures.pointclouds import Pointclouds
+
+from .common_testing import TestCaseMixin
+
+
+class TestVertAlign(TestCaseMixin, unittest.TestCase):
+    @staticmethod
+    def vert_align_naive(
+        feats, verts, return_packed: bool = False, align_corners: bool = True
+    ):
+        """
+        Naive implementation of vert_align.
+        """
+        if torch.is_tensor(feats):
+            feats = [feats]
+        N = feats[0].shape[0]
+
+        out_feats = []
+        # sample every example in the batch separately
+        for i in range(N):
+            out_i_feats = []
+            for feat in feats:
+                feats_i = feat[i][None, :, :, :]  # (1, C, H, W)
+                if torch.is_tensor(verts):
+                    grid = verts[i][None, None, :, :2]  # (1, 1, V, 2)
+                elif hasattr(verts, "verts_list"):
+                    grid = verts.verts_list()[i][None, None, :, :2]  # (1, 1, V, 2)
+                elif hasattr(verts, "points_list"):
+                    grid = verts.points_list()[i][None, None, :, :2]  # (1, 1, V, 2)
+                else:
+                    raise ValueError("verts_or_meshes is invalid")
+                feat_sampled_i = F.grid_sample(
+                    feats_i,
+                    grid,
+                    mode="bilinear",
+                    padding_mode="zeros",
+                    align_corners=align_corners,
+                )  # (1, C, 1, V)
+                feat_sampled_i = feat_sampled_i.squeeze(2).squeeze(0)  # (C, V)
+                feat_sampled_i = feat_sampled_i.transpose(1, 0)  # (V, C)
+                out_i_feats.append(feat_sampled_i)
+            out_i_feats = torch.cat(out_i_feats, 1)  # (V, sum(C))
+            out_feats.append(out_i_feats)
+
+        if return_packed:
+            out_feats = torch.cat(out_feats, 0)  # (sum(V), sum(C))
+        else:
+            out_feats = torch.stack(out_feats, 0)  # (N, V, sum(C))
+        return out_feats
+
+    @staticmethod
+    def init_meshes(
+        num_meshes: int = 10, num_verts: int = 1000, num_faces: int = 3000
+    ) -> Meshes:
+        device = torch.device("cuda:0")
+        verts_list = []
+        faces_list = []
+        for _ in range(num_meshes):
+            verts = (
+                torch.rand((num_verts, 3), dtype=torch.float32, device=device) * 2.0
+                - 1.0
+            )  # verts in the space of [-1, 1]
+            faces = torch.randint(
+                num_verts, size=(num_faces, 3), dtype=torch.int64, device=device
+            )
+            verts_list.append(verts)
+            faces_list.append(faces)
+        meshes = Meshes(verts_list, faces_list)
+
+        return meshes
+
+    @staticmethod
+    def init_pointclouds(num_clouds: int = 10, num_points: int = 1000) -> Pointclouds:
+        device = torch.device("cuda:0")
+        points_list = []
+        for _ in range(num_clouds):
+            points = (
+                torch.rand((num_points, 3), dtype=torch.float32, device=device) * 2.0
+                - 1.0
+            )  # points in the space of [-1, 1]
+            points_list.append(points)
+        pointclouds = Pointclouds(points=points_list)
+
+        return pointclouds
+
+    @staticmethod
+    def init_feats(batch_size: int = 10, num_channels: int = 256, device: str = "cuda"):
+        H, W = [14, 28], [14, 28]
+        feats = []
+        for h, w in zip(H, W):
+            feats.append(torch.rand((batch_size, num_channels, h, w), device=device))
+        return feats
+
+    def test_vert_align_with_meshes(self):
+        """
+        Test vert align vs naive implementation with meshes.
+        """
+        meshes = TestVertAlign.init_meshes(10, 1000, 3000)
+        feats = TestVertAlign.init_feats(10, 256)
+
+        # feats in list
+        out = vert_align(feats, meshes, return_packed=True)
+        naive_out = TestVertAlign.vert_align_naive(feats, meshes, return_packed=True)
+        self.assertClose(out, naive_out)
+
+        # feats as tensor
+        out = vert_align(feats[0], meshes, return_packed=True)
+        naive_out = TestVertAlign.vert_align_naive(feats[0], meshes, return_packed=True)
+        self.assertClose(out, naive_out)
+
+    def test_vert_align_with_pointclouds(self):
+        """
+        Test vert align vs naive implementation with meshes.
+        """
+        pointclouds = TestVertAlign.init_pointclouds(10, 1000)
+        feats = TestVertAlign.init_feats(10, 256)
+
+        # feats in list
+        out = vert_align(feats, pointclouds, return_packed=True)
+        naive_out = TestVertAlign.vert_align_naive(
+            feats, pointclouds, return_packed=True
+        )
+        self.assertClose(out, naive_out)
+
+        # feats as tensor
+        out = vert_align(feats[0], pointclouds, return_packed=True)
+        naive_out = TestVertAlign.vert_align_naive(
+            feats[0], pointclouds, return_packed=True
+        )
+        self.assertClose(out, naive_out)
+
+    def test_vert_align_with_verts(self):
+        """
+        Test vert align vs naive implementation with verts as tensor.
+        """
+        feats = TestVertAlign.init_feats(10, 256)
+        verts = (
+            torch.rand((10, 100, 3), dtype=torch.float32, device=feats[0].device) * 2.0
+            - 1.0
+        )
+
+        # feats in list
+        out = vert_align(feats, verts, return_packed=True)
+        naive_out = TestVertAlign.vert_align_naive(feats, verts, return_packed=True)
+        self.assertClose(out, naive_out)
+
+        # feats as tensor
+        out = vert_align(feats[0], verts, return_packed=True)
+        naive_out = TestVertAlign.vert_align_naive(feats[0], verts, return_packed=True)
+        self.assertClose(out, naive_out)
+
+        out2 = vert_align(feats[0], verts, return_packed=True, align_corners=False)
+        naive_out2 = TestVertAlign.vert_align_naive(
+            feats[0], verts, return_packed=True, align_corners=False
+        )
+        self.assertFalse(torch.allclose(out, out2))
+        self.assertTrue(torch.allclose(out2, naive_out2))
+
+    @staticmethod
+    def vert_align_with_init(
+        num_meshes: int, num_verts: int, num_faces: int, device: str = "cpu"
+    ):
+        device = torch.device(device)
+        verts_list = []
+        faces_list = []
+        for _ in range(num_meshes):
+            verts = torch.rand((num_verts, 3), dtype=torch.float32, device=device)
+            faces = torch.randint(
+                num_verts, size=(num_faces, 3), dtype=torch.int64, device=device
+            )
+            verts_list.append(verts)
+            faces_list.append(faces)
+        meshes = Meshes(verts_list, faces_list)
+        feats = TestVertAlign.init_feats(num_meshes, device=device)
+        torch.cuda.synchronize()
+
+        def sample_features():
+            vert_align(feats, meshes, return_packed=True)
+            torch.cuda.synchronize()
+
+        return sample_features

third_party/AnyBimanual/third_party/pytorch3d/tests/test_vis.py

@@ -0,0 +1,74 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from pytorch3d.renderer import HeterogeneousRayBundle, PerspectiveCameras, RayBundle
+from pytorch3d.structures import Meshes, Pointclouds
+from pytorch3d.transforms import random_rotations
+
+# Some of these imports are only needed for testing code coverage
+from pytorch3d.vis import (  # noqa: F401
+    get_camera_wireframe,  # noqa: F401
+    plot_batch_individually,  # noqa: F401
+    plot_scene,
+    texturesuv_image_PIL,  # noqa: F401
+)
+
+
+class TestPlotlyVis(unittest.TestCase):
+    def test_plot_scene(
+        self,
+        B: int = 3,
+        n_rays: int = 128,
+        n_pts_per_ray: int = 32,
+        n_verts: int = 32,
+        n_edges: int = 64,
+        n_pts: int = 256,
+    ):
+        """
+        Tests plotting of all supported structures using plot_scene.
+        """
+        for device in ["cpu", "cuda:0"]:
+            plot_scene(
+                {
+                    "scene": {
+                        "ray_bundle": RayBundle(
+                            origins=torch.randn(B, n_rays, 3, device=device),
+                            xys=torch.randn(B, n_rays, 2, device=device),
+                            directions=torch.randn(B, n_rays, 3, device=device),
+                            lengths=torch.randn(
+                                B, n_rays, n_pts_per_ray, device=device
+                            ),
+                        ),
+                        "heterogeneous_ray_bundle": HeterogeneousRayBundle(
+                            origins=torch.randn(B * n_rays, 3, device=device),
+                            xys=torch.randn(B * n_rays, 2, device=device),
+                            directions=torch.randn(B * n_rays, 3, device=device),
+                            lengths=torch.randn(
+                                B * n_rays, n_pts_per_ray, device=device
+                            ),
+                            camera_ids=torch.randint(
+                                low=0, high=B, size=(B * n_rays,), device=device
+                            ),
+                        ),
+                        "camera": PerspectiveCameras(
+                            R=random_rotations(B, device=device),
+                            T=torch.randn(B, 3, device=device),
+                        ),
+                        "mesh": Meshes(
+                            verts=torch.randn(B, n_verts, 3, device=device),
+                            faces=torch.randint(
+                                low=0, high=n_verts, size=(B, n_edges, 3), device=device
+                            ),
+                        ),
+                        "point_clouds": Pointclouds(
+                            points=torch.randn(B, n_pts, 3, device=device),
+                        ),
+                    }
+                }
+            )

third_party/AnyBimanual/third_party/pytorch3d/tests/test_volumes.py

@@ -0,0 +1,987 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import copy
+import itertools
+import random
+import unittest
+
+import numpy as np
+import torch
+from pytorch3d.structures.volumes import VolumeLocator, Volumes
+from pytorch3d.transforms import Scale
+
+from .common_testing import TestCaseMixin
+
+
+class TestVolumes(TestCaseMixin, unittest.TestCase):
+    def setUp(self) -> None:
+        np.random.seed(42)
+        torch.manual_seed(42)
+        random.seed(42)
+
+    @staticmethod
+    def _random_volume_list(
+        num_volumes, min_size, max_size, num_channels, device, rand_sizes=None
+    ):
+        """
+        Init a list of `num_volumes` random tensors of size [num_channels, *rand_size].
+        If `rand_sizes` is None, rand_size is a 3D long vector sampled
+        from [min_size, max_size]. Otherwise, rand_size should be a list
+        [rand_size_1, rand_size_2, ..., rand_size_num_volumes] where each
+        `rand_size_i` denotes the size of the corresponding `i`-th tensor.
+        """
+        if rand_sizes is None:
+            rand_sizes = [
+                [random.randint(min_size, vs) for vs in max_size]
+                for _ in range(num_volumes)
+            ]
+
+        volume_list = [
+            torch.randn(
+                size=[num_channels, *rand_size], device=device, dtype=torch.float32
+            )
+            for rand_size in rand_sizes
+        ]
+
+        return volume_list, rand_sizes
+
+    def _check_indexed_volumes(self, v, selected, indices):
+        for selectedIdx, index in indices:
+            self.assertClose(selected.densities()[selectedIdx], v.densities()[index])
+            self.assertClose(
+                v.locator._local_to_world_transform.get_matrix()[index],
+                selected.locator._local_to_world_transform.get_matrix()[selectedIdx],
+            )
+            if selected.features() is not None:
+                self.assertClose(selected.features()[selectedIdx], v.features()[index])
+
+    def test_get_item(
+        self,
+        num_volumes=5,
+        num_channels=4,
+        volume_size=(10, 13, 8),
+        dtype=torch.float32,
+    ):
+
+        device = torch.device("cuda:0")
+
+        # make sure we have at least 3 volumes to prevent indexing crash
+        num_volumes = max(num_volumes, 3)
+
+        features = torch.randn(
+            size=[num_volumes, num_channels, *volume_size],
+            device=device,
+            dtype=torch.float32,
+        )
+        densities = torch.randn(
+            size=[num_volumes, 1, *volume_size], device=device, dtype=torch.float32
+        )
+
+        features_list, rand_sizes = TestVolumes._random_volume_list(
+            num_volumes, 3, volume_size, num_channels, device
+        )
+        densities_list, _ = TestVolumes._random_volume_list(
+            num_volumes, 3, volume_size, 1, device, rand_sizes=rand_sizes
+        )
+
+        volume_translation = -torch.randn(num_volumes, 3).type_as(features)
+        voxel_size = torch.rand(num_volumes, 1).type_as(features) + 0.5
+
+        for features_, densities_ in zip(
+            (None, features, features_list), (densities, densities, densities_list)
+        ):
+
+            # init the volume structure
+            v = Volumes(
+                features=features_,
+                densities=densities_,
+                volume_translation=volume_translation,
+                voxel_size=voxel_size,
+            )
+
+            # int index
+            index = 1
+            v_selected = v[index]
+            self.assertEqual(len(v_selected), 1)
+            self._check_indexed_volumes(v, v_selected, [(0, 1)])
+
+            # list index
+            index = [1, 2]
+            v_selected = v[index]
+            self.assertEqual(len(v_selected), len(index))
+            self._check_indexed_volumes(v, v_selected, enumerate(index))
+
+            # slice index
+            index = slice(0, 2, 1)
+            v_selected = v[0:2]
+            self.assertEqual(len(v_selected), 2)
+            self._check_indexed_volumes(v, v_selected, [(0, 0), (1, 1)])
+
+            # bool tensor
+            index = (torch.rand(num_volumes) > 0.5).to(device)
+            index[:2] = True  # make sure smth is selected
+            v_selected = v[index]
+            self.assertEqual(len(v_selected), index.sum())
+            self._check_indexed_volumes(
+                v,
+                v_selected,
+                zip(
+                    torch.arange(index.sum()),
+                    torch.nonzero(index, as_tuple=False).squeeze(),
+                ),
+            )
+
+            # int tensor
+            index = torch.tensor([1, 2], dtype=torch.int64, device=device)
+            v_selected = v[index]
+            self.assertEqual(len(v_selected), index.numel())
+            self._check_indexed_volumes(v, v_selected, enumerate(index.tolist()))
+
+            # invalid index
+            index = torch.tensor([1, 0, 1], dtype=torch.float32, device=device)
+            with self.assertRaises(IndexError):
+                v_selected = v[index]
+            index = 1.2  # floating point index
+            with self.assertRaises(IndexError):
+                v_selected = v[index]
+
+    def test_locator_init(self, batch_size=9, resolution=(3, 5, 7)):
+        with self.subTest("VolumeLocator init with all sizes equal"):
+            grid_sizes = [resolution for _ in range(batch_size)]
+            locator_tuple = VolumeLocator(
+                batch_size=batch_size, grid_sizes=resolution, device=torch.device("cpu")
+            )
+            locator_list = VolumeLocator(
+                batch_size=batch_size, grid_sizes=grid_sizes, device=torch.device("cpu")
+            )
+            locator_tensor = VolumeLocator(
+                batch_size=batch_size,
+                grid_sizes=torch.tensor(grid_sizes),
+                device=torch.device("cpu"),
+            )
+            expected_grid_sizes = torch.tensor(grid_sizes)
+            expected_resolution = resolution
+            assert torch.allclose(expected_grid_sizes, locator_tuple._grid_sizes)
+            assert torch.allclose(expected_grid_sizes, locator_list._grid_sizes)
+            assert torch.allclose(expected_grid_sizes, locator_tensor._grid_sizes)
+            self.assertEqual(expected_resolution, locator_tuple._resolution)
+            self.assertEqual(expected_resolution, locator_list._resolution)
+            self.assertEqual(expected_resolution, locator_tensor._resolution)
+
+        with self.subTest("VolumeLocator with different sizes in different grids"):
+            grid_sizes_list = [
+                torch.randint(low=1, high=42, size=(3,)) for _ in range(batch_size)
+            ]
+            grid_sizes_tensor = torch.cat([el[None] for el in grid_sizes_list])
+            locator_list = VolumeLocator(
+                batch_size=batch_size,
+                grid_sizes=grid_sizes_list,
+                device=torch.device("cpu"),
+            )
+            locator_tensor = VolumeLocator(
+                batch_size=batch_size,
+                grid_sizes=grid_sizes_tensor,
+                device=torch.device("cpu"),
+            )
+            expected_grid_sizes = grid_sizes_tensor
+            expected_resolution = tuple(torch.max(expected_grid_sizes, dim=0).values)
+            assert torch.allclose(expected_grid_sizes, locator_list._grid_sizes)
+            assert torch.allclose(expected_grid_sizes, locator_tensor._grid_sizes)
+            self.assertEqual(expected_resolution, locator_list._resolution)
+            self.assertEqual(expected_resolution, locator_tensor._resolution)
+
+    def test_coord_transforms(self, num_volumes=3, num_channels=4, dtype=torch.float32):
+        """
+        Test the correctness of the conversion between the internal
+        Transform3D Volumes.VolumeLocator._local_to_world_transform and the initialization
+        from the translation and voxel_size.
+        """
+
+        device = torch.device("cuda:0")
+
+        # try for 10 sets of different random sizes/centers/voxel_sizes
+        for _ in range(10):
+
+            size = torch.randint(high=10, size=(3,), low=3).tolist()
+
+            densities = torch.randn(
+                size=[num_volumes, num_channels, *size],
+                device=device,
+                dtype=torch.float32,
+            )
+
+            # init the transformation params
+            volume_translation = torch.randn(num_volumes, 3)
+            voxel_size = torch.rand(num_volumes, 3) * 3.0 + 0.5
+
+            # get the corresponding Transform3d object
+            local_offset = torch.tensor(list(size), dtype=torch.float32, device=device)[
+                [2, 1, 0]
+            ][None].repeat(num_volumes, 1)
+            local_to_world_transform = (
+                Scale(0.5 * local_offset - 0.5, device=device)
+                .scale(voxel_size)
+                .translate(-volume_translation)
+            )
+
+            # init the volume structures with the scale and translation,
+            # then get the coord grid in world coords
+            v_trans_vs = Volumes(
+                densities=densities,
+                voxel_size=voxel_size,
+                volume_translation=volume_translation,
+            )
+            grid_rot_trans_vs = v_trans_vs.get_coord_grid(world_coordinates=True)
+
+            # map the default local coords to the world coords
+            # with local_to_world_transform
+            v_default = Volumes(densities=densities)
+            grid_default_local = v_default.get_coord_grid(world_coordinates=False)
+            grid_default_world = local_to_world_transform.transform_points(
+                grid_default_local.view(num_volumes, -1, 3)
+            ).view(num_volumes, *size, 3)
+
+            # check that both grids are the same
+            self.assertClose(grid_rot_trans_vs, grid_default_world, atol=1e-5)
+
+            # check that the transformations are the same
+            self.assertClose(
+                v_trans_vs.get_local_to_world_coords_transform().get_matrix(),
+                local_to_world_transform.get_matrix(),
+                atol=1e-5,
+            )
+
+    def test_coord_grid_convention(
+        self, num_volumes=3, num_channels=4, dtype=torch.float32
+    ):
+        """
+        Check that for a trivial volume with spatial size DxHxW=5x7x5:
+        1) xyz_world=(0, 0, 0) lands right in the middle of the volume
+        with xyz_local=(0, 0, 0).
+        2) xyz_world=(-2, 3, 2) results in xyz_local=(-1, 1, -1).
+        3) The centeral voxel of the volume coordinate grid
+        has coords x_world=(0, 0, 0) and x_local=(0, 0, 0)
+        4) grid_sampler(world_coordinate_grid, local_coordinate_grid)
+        is the same as world_coordinate_grid itself. I.e. the local coordinate
+        grid matches the grid_sampler coordinate convention.
+        """
+
+        device = torch.device("cuda:0")
+
+        densities = torch.randn(
+            size=[num_volumes, num_channels, 5, 7, 5],
+            device=device,
+            dtype=torch.float32,
+        )
+        v_trivial = Volumes(densities=densities)
+
+        # check the case with x_world=(0,0,0)
+        pts_world = torch.zeros(num_volumes, 1, 3, device=device, dtype=torch.float32)
+        pts_local = v_trivial.world_to_local_coords(pts_world)
+        pts_local_expected = torch.zeros_like(pts_local)
+        self.assertClose(pts_local, pts_local_expected)
+
+        # check the case with x_world=(-2, 3, -2)
+        pts_world = torch.tensor([-2, 3, -2], device=device, dtype=torch.float32)[
+            None, None
+        ].repeat(num_volumes, 1, 1)
+        pts_local = v_trivial.world_to_local_coords(pts_world)
+        pts_local_expected = torch.tensor(
+            [-1, 1, -1], device=device, dtype=torch.float32
+        )[None, None].repeat(num_volumes, 1, 1)
+        self.assertClose(pts_local, pts_local_expected)
+
+        # check that the central voxel has coords x_world=(0, 0, 0) and x_local(0, 0, 0)
+        grid_world = v_trivial.get_coord_grid(world_coordinates=True)
+        grid_local = v_trivial.get_coord_grid(world_coordinates=False)
+        for grid in (grid_world, grid_local):
+            x0 = grid[0, :, :, 2, 0]
+            y0 = grid[0, :, 3, :, 1]
+            z0 = grid[0, 2, :, :, 2]
+            for coord_line in (x0, y0, z0):
+                self.assertClose(coord_line, torch.zeros_like(coord_line), atol=1e-7)
+
+        # resample grid_world using grid_sampler with local coords
+        # -> make sure the resampled version is the same as original
+        grid_world_resampled = torch.nn.functional.grid_sample(
+            grid_world.permute(0, 4, 1, 2, 3), grid_local, align_corners=True
+        ).permute(0, 2, 3, 4, 1)
+        self.assertClose(grid_world_resampled, grid_world, atol=1e-7)
+
+        for align_corners in [True, False]:
+            v_trivial = Volumes(densities=densities, align_corners=align_corners)
+
+            # check the case with x_world=(0,0,0)
+            pts_world = torch.zeros(
+                num_volumes, 1, 3, device=device, dtype=torch.float32
+            )
+            pts_local = v_trivial.world_to_local_coords(pts_world)
+            pts_local_expected = torch.zeros_like(pts_local)
+            self.assertClose(pts_local, pts_local_expected)
+
+            # check the case with x_world=(-2, 3, -2)
+            pts_world_tuple = [-2, 3, -2]
+            pts_world = torch.tensor(
+                pts_world_tuple, device=device, dtype=torch.float32
+            )[None, None].repeat(num_volumes, 1, 1)
+            pts_local = v_trivial.world_to_local_coords(pts_world)
+            pts_local_expected = torch.tensor(
+                [-1, 1, -1], device=device, dtype=torch.float32
+            )[None, None].repeat(num_volumes, 1, 1)
+            self.assertClose(pts_local, pts_local_expected)
+
+            # # check that the central voxel has coords x_world=(0, 0, 0) and x_local(0, 0, 0)
+            grid_world = v_trivial.get_coord_grid(world_coordinates=True)
+            grid_local = v_trivial.get_coord_grid(world_coordinates=False)
+            for grid in (grid_world, grid_local):
+                x0 = grid[0, :, :, 2, 0]
+                y0 = grid[0, :, 3, :, 1]
+                z0 = grid[0, 2, :, :, 2]
+                for coord_line in (x0, y0, z0):
+                    self.assertClose(
+                        coord_line, torch.zeros_like(coord_line), atol=1e-7
+                    )
+
+            # resample grid_world using grid_sampler with local coords
+            # -> make sure the resampled version is the same as original
+            grid_world_resampled = torch.nn.functional.grid_sample(
+                grid_world.permute(0, 4, 1, 2, 3),
+                grid_local,
+                align_corners=align_corners,
+            ).permute(0, 2, 3, 4, 1)
+            self.assertClose(grid_world_resampled, grid_world, atol=1e-7)
+
+    def test_coord_grid_convention_heterogeneous(
+        self, num_channels=4, dtype=torch.float32
+    ):
+        """
+        Check that for a list of 2 trivial volumes with
+        spatial sizes DxHxW=(5x7x5, 3x5x5):
+        1) xyz_world=(0, 0, 0) lands right in the middle of the volume
+        with xyz_local=(0, 0, 0).
+        2) xyz_world=((-2, 3, -2), (-2, -2,  1)) results
+        in xyz_local=((-1, 1, -1), (-1, -1, 1)).
+        3) The centeral voxel of the volume coordinate grid
+        has coords x_world=(0, 0, 0) and x_local=(0, 0, 0)
+        4) grid_sampler(world_coordinate_grid, local_coordinate_grid)
+        is the same as world_coordinate_grid itself. I.e. the local coordinate
+        grid matches the grid_sampler coordinate convention.
+        """
+
+        device = torch.device("cuda:0")
+
+        sizes = [(5, 7, 5), (3, 5, 5)]
+
+        densities_list = [
+            torch.randn(size=[num_channels, *size], device=device, dtype=torch.float32)
+            for size in sizes
+        ]
+
+        # init the volume
+        v_trivial = Volumes(densities=densities_list)
+
+        # check the border point locations
+        pts_world = torch.tensor(
+            [[-2.0, 3.0, -2.0], [-2.0, -2.0, 1.0]], device=device, dtype=torch.float32
+        )[:, None]
+        pts_local = v_trivial.world_to_local_coords(pts_world)
+        pts_local_expected = torch.tensor(
+            [[-1.0, 1.0, -1.0], [-1.0, -1.0, 1.0]], device=device, dtype=torch.float32
+        )[:, None]
+        self.assertClose(pts_local, pts_local_expected)
+
+        # check that the central voxel has coords x_world=(0, 0, 0) and x_local(0, 0, 0)
+        grid_world = v_trivial.get_coord_grid(world_coordinates=True)
+        grid_local = v_trivial.get_coord_grid(world_coordinates=False)
+        for grid in (grid_world, grid_local):
+            x0 = grid[0, :, :, 2, 0]
+            y0 = grid[0, :, 3, :, 1]
+            z0 = grid[0, 2, :, :, 2]
+            for coord_line in (x0, y0, z0):
+                self.assertClose(coord_line, torch.zeros_like(coord_line), atol=1e-7)
+            x0 = grid[1, :, :, 2, 0]
+            y0 = grid[1, :, 2, :, 1]
+            z0 = grid[1, 1, :, :, 2]
+            for coord_line in (x0, y0, z0):
+                self.assertClose(coord_line, torch.zeros_like(coord_line), atol=1e-7)
+
+        # resample grid_world using grid_sampler with local coords
+        # -> make sure the resampled version is the same as original
+        for grid_world_, grid_local_, size in zip(grid_world, grid_local, sizes):
+            grid_world_crop = grid_world_[: size[0], : size[1], : size[2], :][None]
+            grid_local_crop = grid_local_[: size[0], : size[1], : size[2], :][None]
+            grid_world_crop_resampled = torch.nn.functional.grid_sample(
+                grid_world_crop.permute(0, 4, 1, 2, 3),
+                grid_local_crop,
+                align_corners=True,
+            ).permute(0, 2, 3, 4, 1)
+            self.assertClose(grid_world_crop_resampled, grid_world_crop, atol=1e-7)
+
+    def test_coord_grid_transforms(
+        self, num_volumes=3, num_channels=4, dtype=torch.float32
+    ):
+        """
+        Test whether conversion between local-world coordinates of the
+        volume returns correct results.
+        """
+
+        device = torch.device("cuda:0")
+
+        # try for 10 sets of different random sizes/centers/voxel_sizes
+        for _ in range(10):
+
+            size = torch.randint(high=10, size=(3,), low=3).tolist()
+
+            center = torch.randn(num_volumes, 3, dtype=torch.float32, device=device)
+            voxel_size = torch.rand(1, dtype=torch.float32, device=device) * 5.0 + 0.5
+
+            for densities in (
+                torch.randn(
+                    size=[num_volumes, num_channels, *size],
+                    device=device,
+                    dtype=torch.float32,
+                ),
+                TestVolumes._random_volume_list(
+                    num_volumes, 3, size, num_channels, device, rand_sizes=None
+                )[0],
+            ):
+
+                # init the volume structure
+                v = Volumes(
+                    densities=densities,
+                    voxel_size=voxel_size,
+                    volume_translation=-center,
+                )
+
+                # get local coord grid
+                grid_local = v.get_coord_grid(world_coordinates=False)
+
+                # convert from world to local to world
+                grid_world = v.get_coord_grid(world_coordinates=True)
+                grid_local_2 = v.world_to_local_coords(grid_world)
+                grid_world_2 = v.local_to_world_coords(grid_local_2)
+
+                # assertions on shape and values of grid_world and grid_local
+                self.assertClose(grid_world, grid_world_2, atol=1e-5)
+                self.assertClose(grid_local, grid_local_2, atol=1e-5)
+
+                # check that the individual slices of the location grid have
+                # constant values along expected dimensions
+                for plane_dim in (1, 2, 3):
+                    for grid_plane in grid_world.split(1, dim=plane_dim):
+                        grid_coord_dim = {1: 2, 2: 1, 3: 0}[plane_dim]
+                        grid_coord_plane = grid_plane.squeeze()[..., grid_coord_dim]
+                        # check that all elements of grid_coord_plane are
+                        # the same for each batch element
+                        self.assertClose(
+                            grid_coord_plane.reshape(num_volumes, -1).max(dim=1).values,
+                            grid_coord_plane.reshape(num_volumes, -1).min(dim=1).values,
+                        )
+
+    def test_clone(
+        self, num_volumes=3, num_channels=4, size=(6, 8, 10), dtype=torch.float32
+    ):
+        """
+        Test cloning of a `Volumes` object
+        """
+
+        device = torch.device("cuda:0")
+
+        features = torch.randn(
+            size=[num_volumes, num_channels, *size], device=device, dtype=torch.float32
+        )
+        densities = torch.rand(
+            size=[num_volumes, 1, *size], device=device, dtype=torch.float32
+        )
+
+        for has_features in (True, False):
+            v = Volumes(
+                densities=densities, features=features if has_features else None
+            )
+            vnew = v.clone()
+            vnew._densities.data[0, 0, 0, 0, 0] += 1.0
+            self.assertNotAlmostEqual(
+                float(
+                    (vnew.densities()[0, 0, 0, 0, 0] - v.densities()[0, 0, 0, 0, 0])
+                    .abs()
+                    .max()
+                ),
+                0.0,
+            )
+
+            if has_features:
+                vnew._features.data[0, 0, 0, 0, 0] += 1.0
+                self.assertNotAlmostEqual(
+                    float(
+                        (vnew.features()[0, 0, 0, 0, 0] - v.features()[0, 0, 0, 0, 0])
+                        .abs()
+                        .max()
+                    ),
+                    0.0,
+                )
+
+    def _check_vars_on_device(self, v, desired_device):
+        for var_name, var in vars(v).items():
+            if var_name != "device":
+                if var is not None:
+                    self.assertTrue(
+                        var.device.type == desired_device.type,
+                        (var_name, var.device, desired_device),
+                    )
+            else:
+                self.assertTrue(var.type == desired_device.type)
+
+    def test_to(
+        self, num_volumes=3, num_channels=4, size=(6, 8, 10), dtype=torch.float32
+    ):
+        """
+        Test the moving of the volumes from/to gpu and cpu
+        """
+
+        features = torch.randn(
+            size=[num_volumes, num_channels, *size], dtype=torch.float32
+        )
+        densities = torch.rand(size=[num_volumes, 1, *size], dtype=dtype)
+        volumes = Volumes(densities=densities, features=features)
+        locator = VolumeLocator(
+            batch_size=5, grid_sizes=(3, 5, 7), device=volumes.device
+        )
+
+        for name, obj in (("VolumeLocator", locator), ("Volumes", volumes)):
+            with self.subTest(f"Moving {name} from/to gpu and cpu"):
+                # Test support for str and torch.device
+                cpu_device = torch.device("cpu")
+
+                converted_obj = obj.to("cpu")
+                self.assertEqual(cpu_device, converted_obj.device)
+                self.assertEqual(cpu_device, obj.device)
+                self.assertIs(obj, converted_obj)
+
+                converted_obj = obj.to(cpu_device)
+                self.assertEqual(cpu_device, converted_obj.device)
+                self.assertEqual(cpu_device, obj.device)
+                self.assertIs(obj, converted_obj)
+
+                cuda_device = torch.device("cuda:0")
+
+                converted_obj = obj.to("cuda:0")
+                self.assertEqual(cuda_device, converted_obj.device)
+                self.assertEqual(cpu_device, obj.device)
+                self.assertIsNot(obj, converted_obj)
+
+                converted_obj = obj.to(cuda_device)
+                self.assertEqual(cuda_device, converted_obj.device)
+                self.assertEqual(cpu_device, obj.device)
+                self.assertIsNot(obj, converted_obj)
+
+        with self.subTest("Test device placement of internal tensors of Volumes"):
+            features = features.to(cuda_device)
+            densities = features.to(cuda_device)
+
+            for features_ in (features, None):
+                volumes = Volumes(densities=densities, features=features_)
+
+                cpu_volumes = volumes.cpu()
+                cuda_volumes = cpu_volumes.cuda()
+                cuda_volumes2 = cuda_volumes.cuda()
+                cpu_volumes2 = cuda_volumes2.cpu()
+
+                for volumes1, volumes2 in itertools.combinations(
+                    (volumes, cpu_volumes, cpu_volumes2, cuda_volumes, cuda_volumes2), 2
+                ):
+                    if volumes1 is cuda_volumes and volumes2 is cuda_volumes2:
+                        # checks that we do not copy if the devices stay the same
+                        assert_fun = self.assertIs
+                    else:
+                        assert_fun = self.assertSeparate
+                    assert_fun(volumes1._densities, volumes2._densities)
+                    if features_ is not None:
+                        assert_fun(volumes1._features, volumes2._features)
+                    for volumes_ in (volumes1, volumes2):
+                        if volumes_ in (cpu_volumes, cpu_volumes2):
+                            self._check_vars_on_device(volumes_, cpu_device)
+                        else:
+                            self._check_vars_on_device(volumes_, cuda_device)
+
+        with self.subTest("Test device placement of internal tensors of VolumeLocator"):
+            for device1, device2 in itertools.combinations(
+                (torch.device("cpu"), torch.device("cuda:0")), 2
+            ):
+                locator = locator.to(device1)
+                locator = locator.to(device2)
+                self.assertEqual(locator._grid_sizes.device, device2)
+                self.assertEqual(locator._local_to_world_transform.device, device2)
+
+    def _check_padded(self, x_pad, x_list, grid_sizes):
+        """
+        Check that padded tensors x_pad are the same as x_list tensors.
+        """
+        num_volumes = len(x_list)
+        for i in range(num_volumes):
+            self.assertClose(
+                x_pad[i][:, : grid_sizes[i][0], : grid_sizes[i][1], : grid_sizes[i][2]],
+                x_list[i],
+            )
+
+    def test_feature_density_setters(self):
+        """
+        Tests getters and setters for padded/list representations.
+        """
+
+        device = torch.device("cuda:0")
+        diff_device = torch.device("cpu")
+
+        num_volumes = 30
+        num_channels = 4
+        K = 20
+
+        densities = []
+        features = []
+        grid_sizes = []
+        diff_grid_sizes = []
+
+        for _ in range(num_volumes):
+            grid_size = torch.randint(K - 1, size=(3,)).long() + 1
+            densities.append(
+                torch.rand((1, *grid_size), device=device, dtype=torch.float32)
+            )
+            features.append(
+                torch.rand(
+                    (num_channels, *grid_size), device=device, dtype=torch.float32
+                )
+            )
+            grid_sizes.append(grid_size)
+
+            diff_grid_size = (
+                copy.deepcopy(grid_size) + torch.randint(2, size=(3,)).long() + 1
+            )
+            diff_grid_sizes.append(diff_grid_size)
+        grid_sizes = torch.stack(grid_sizes).to(device)
+        diff_grid_sizes = torch.stack(diff_grid_sizes).to(device)
+
+        volumes = Volumes(densities=densities, features=features)
+        self.assertClose(volumes.get_grid_sizes(), grid_sizes)
+
+        # test the getters
+        features_padded = volumes.features()
+        densities_padded = volumes.densities()
+        features_list = volumes.features_list()
+        densities_list = volumes.densities_list()
+        for x_pad, x_list in zip(
+            (densities_padded, features_padded, densities_padded, features_padded),
+            (densities_list, features_list, densities, features),
+        ):
+            self._check_padded(x_pad, x_list, grid_sizes)
+
+        # test feature setters
+        features_new = [
+            torch.rand((num_channels, *grid_size), device=device, dtype=torch.float32)
+            for grid_size in grid_sizes
+        ]
+        volumes._set_features(features_new)
+        features_new_list = volumes.features_list()
+        features_new_padded = volumes.features()
+        for x_pad, x_list in zip(
+            (features_new_padded, features_new_padded),
+            (features_new, features_new_list),
+        ):
+            self._check_padded(x_pad, x_list, grid_sizes)
+
+        # wrong features to update
+        bad_features_new = [
+            [
+                torch.rand(
+                    (num_channels, *grid_size), device=diff_device, dtype=torch.float32
+                )
+                for grid_size in diff_grid_sizes
+            ],
+            torch.rand(
+                (num_volumes, num_channels, K + 1, K, K),
+                device=device,
+                dtype=torch.float32,
+            ),
+            None,
+        ]
+        for bad_features_new_ in bad_features_new:
+            with self.assertRaises(ValueError):
+                volumes._set_densities(bad_features_new_)
+
+        # test density setters
+        densities_new = [
+            torch.rand((1, *grid_size), device=device, dtype=torch.float32)
+            for grid_size in grid_sizes
+        ]
+        volumes._set_densities(densities_new)
+        densities_new_list = volumes.densities_list()
+        densities_new_padded = volumes.densities()
+        for x_pad, x_list in zip(
+            (densities_new_padded, densities_new_padded),
+            (densities_new, densities_new_list),
+        ):
+            self._check_padded(x_pad, x_list, grid_sizes)
+
+        # wrong densities to update
+        bad_densities_new = [
+            [
+                torch.rand((1, *grid_size), device=diff_device, dtype=torch.float32)
+                for grid_size in diff_grid_sizes
+            ],
+            torch.rand(
+                (num_volumes, 1, K + 1, K, K), device=device, dtype=torch.float32
+            ),
+            None,
+        ]
+        for bad_densities_new_ in bad_densities_new:
+            with self.assertRaises(ValueError):
+                volumes._set_densities(bad_densities_new_)
+
+        # test update_padded
+        volumes = Volumes(densities=densities, features=features)
+        volumes_updated = volumes.update_padded(
+            densities_new, new_features=features_new
+        )
+        densities_new_list = volumes_updated.densities_list()
+        densities_new_padded = volumes_updated.densities()
+        features_new_list = volumes_updated.features_list()
+        features_new_padded = volumes_updated.features()
+        for x_pad, x_list in zip(
+            (
+                densities_new_padded,
+                densities_new_padded,
+                features_new_padded,
+                features_new_padded,
+            ),
+            (densities_new, densities_new_list, features_new, features_new_list),
+        ):
+            self._check_padded(x_pad, x_list, grid_sizes)
+        self.assertIs(volumes.get_grid_sizes(), volumes_updated.get_grid_sizes())
+        self.assertIs(
+            volumes.get_local_to_world_coords_transform(),
+            volumes_updated.get_local_to_world_coords_transform(),
+        )
+        self.assertIs(volumes.device, volumes_updated.device)
+
+    def test_constructor_for_padded_lists(self):
+        """
+        Tests constructor for padded/list representations.
+        """
+
+        device = torch.device("cuda:0")
+        diff_device = torch.device("cpu")
+
+        num_volumes = 3
+        num_channels = 4
+        size = (6, 8, 10)
+        diff_size = (6, 8, 11)
+
+        # good ways to define densities
+        ok_densities = [
+            torch.randn(
+                size=[num_volumes, 1, *size], device=device, dtype=torch.float32
+            ).unbind(0),
+            torch.randn(
+                size=[num_volumes, 1, *size], device=device, dtype=torch.float32
+            ),
+        ]
+
+        # bad ways to define features
+        bad_features = [
+            torch.randn(
+                size=[num_volumes + 1, num_channels, *size],
+                device=device,
+                dtype=torch.float32,
+            ).unbind(
+                0
+            ),  # list with diff batch size
+            torch.randn(
+                size=[num_volumes + 1, num_channels, *size],
+                device=device,
+                dtype=torch.float32,
+            ),  # diff batch size
+            torch.randn(
+                size=[num_volumes, num_channels, *diff_size],
+                device=device,
+                dtype=torch.float32,
+            ).unbind(
+                0
+            ),  # list with different size
+            torch.randn(
+                size=[num_volumes, num_channels, *diff_size],
+                device=device,
+                dtype=torch.float32,
+            ),  # different size
+            torch.randn(
+                size=[num_volumes, num_channels, *size],
+                device=diff_device,
+                dtype=torch.float32,
+            ),  # different device
+            torch.randn(
+                size=[num_volumes, num_channels, *size],
+                device=diff_device,
+                dtype=torch.float32,
+            ).unbind(
+                0
+            ),  # list with different device
+        ]
+
+        # good ways to define features
+        ok_features = [
+            torch.randn(
+                size=[num_volumes, num_channels, *size],
+                device=device,
+                dtype=torch.float32,
+            ).unbind(
+                0
+            ),  # list of features of correct size
+            torch.randn(
+                size=[num_volumes, num_channels, *size],
+                device=device,
+                dtype=torch.float32,
+            ),
+        ]
+
+        for densities in ok_densities:
+            for features in bad_features:
+                self.assertRaises(
+                    ValueError, Volumes, densities=densities, features=features
+                )
+            for features in ok_features:
+                Volumes(densities=densities, features=features)
+
+    def test_constructor(
+        self, num_volumes=3, num_channels=4, size=(6, 8, 10), dtype=torch.float32
+    ):
+        """
+        Test different ways of calling the `Volumes` constructor
+        """
+
+        device = torch.device("cuda:0")
+
+        # all ways to define features
+        features = [
+            torch.randn(
+                size=[num_volumes, num_channels, *size],
+                device=device,
+                dtype=torch.float32,
+            ),  # padded tensor
+            torch.randn(
+                size=[num_volumes, num_channels, *size],
+                device=device,
+                dtype=torch.float32,
+            ).unbind(
+                0
+            ),  # list of features
+            None,  # no features
+        ]
+
+        # bad ways to define features
+        bad_features = [
+            torch.randn(
+                size=[num_volumes, num_channels, 2, *size],
+                device=device,
+                dtype=torch.float32,
+            ),  # 6 dims
+            torch.randn(
+                size=[num_volumes, *size], device=device, dtype=torch.float32
+            ),  # 4 dims
+            torch.randn(
+                size=[num_volumes, *size], device=device, dtype=torch.float32
+            ).unbind(
+                0
+            ),  # list of 4 dim tensors
+        ]
+
+        # all ways to define densities
+        densities = [
+            torch.randn(
+                size=[num_volumes, 1, *size], device=device, dtype=torch.float32
+            ),  # padded tensor
+            torch.randn(
+                size=[num_volumes, 1, *size], device=device, dtype=torch.float32
+            ).unbind(
+                0
+            ),  # list of densities
+        ]
+
+        # bad ways to define densities
+        bad_densities = [
+            None,  # omitted
+            torch.randn(
+                size=[num_volumes, 1, 1, *size], device=device, dtype=torch.float32
+            ),  # 6-dim tensor
+            torch.randn(
+                size=[num_volumes, 1, 1, *size], device=device, dtype=torch.float32
+            ).unbind(
+                0
+            ),  # list of 5-dim densities
+        ]
+
+        # all possible ways to define the voxels sizes
+        vox_sizes = [
+            torch.Tensor([1.0, 1.0, 1.0]),
+            [1.0, 1.0, 1.0],
+            torch.Tensor([1.0, 1.0, 1.0])[None].repeat(num_volumes, 1),
+            torch.Tensor([1.0])[None].repeat(num_volumes, 1),
+            1.0,
+            torch.Tensor([1.0]),
+        ]
+
+        # all possible ways to define the volume translations
+        vol_translations = [
+            torch.Tensor([1.0, 1.0, 1.0]),
+            [1.0, 1.0, 1.0],
+            torch.Tensor([1.0, 1.0, 1.0])[None].repeat(num_volumes, 1),
+        ]
+
+        # wrong ways to define voxel sizes
+        bad_vox_sizes = [
+            torch.Tensor([1.0, 1.0, 1.0, 1.0]),
+            [1.0, 1.0, 1.0, 1.0],
+            torch.Tensor([]),
+            None,
+        ]
+
+        # wrong ways to define the volume translations
+        bad_vol_translations = [
+            torch.Tensor([1.0, 1.0]),
+            [1.0, 1.0],
+            1.0,
+            torch.Tensor([1.0, 1.0, 1.0])[None].repeat(num_volumes + 1, 1),
+        ]
+
+        def zip_with_ok_indicator(good, bad):
+            return zip([*good, *bad], [*([True] * len(good)), *([False] * len(bad))])
+
+        for features_, features_ok in zip_with_ok_indicator(features, bad_features):
+            for densities_, densities_ok in zip_with_ok_indicator(
+                densities, bad_densities
+            ):
+                for vox_size, size_ok in zip_with_ok_indicator(
+                    vox_sizes, bad_vox_sizes
+                ):
+                    for vol_translation, trans_ok in zip_with_ok_indicator(
+                        vol_translations, bad_vol_translations
+                    ):
+                        if (
+                            size_ok and trans_ok and features_ok and densities_ok
+                        ):  # if all entries are good we check that this doesnt throw
+                            Volumes(
+                                features=features_,
+                                densities=densities_,
+                                voxel_size=vox_size,
+                                volume_translation=vol_translation,
+                            )
+
+                        else:  # otherwise we check for ValueError
+                            self.assertRaises(
+                                ValueError,
+                                Volumes,
+                                features=features_,
+                                densities=densities_,
+                                voxel_size=vox_size,
+                                volume_translation=vol_translation,
+                            )

third_party/AnyBimanual/third_party/pytorch3d/website/.dockerignore

@@ -0,0 +1,2 @@
+*/node_modules
+*.log

third_party/AnyBimanual/third_party/pytorch3d/website/.gitignore

@@ -0,0 +1,13 @@
+.DS_Store
+
+node_modules
+
+lib/core/metadata.js
+lib/core/MetadataBlog.js
+
+website/translated_docs
+website/build/
+website/yarn.lock
+website/node_modules
+website/i18n/*
+website/_tutorials/*

third_party/AnyBimanual/third_party/pytorch3d/website/README.md

@@ -0,0 +1,265 @@
+This website was created with [Docusaurus](https://docusaurus.io/).
+
+# Building the PyTorch3D website
+
+## Install
+
+1. Make sure all the dependencies for the website are installed:
+
+```sh
+# Install dependencies
+$ yarn
+
+or
+
+$ npm install docusaurus-init
+```
+
+2. Run your dev server:
+
+```sh
+# Start the site
+$ yarn start
+
+or
+$ ./node_modules/docusaurus/lib/start-server.js
+```
+
+## Build the tutorials
+
+We convert the ipython notebooks to html using `parse_tutorials.py` which is found in the scripts folder at the root of the PyTorch3D directory.
+
+Before running this script install the following dependencies:
+
+```
+pip install nbformat==4.4.0 nbconvert==5.3.1 ipywidgets==7.5.1 tornado==4.2 bs4
+```
+
+Install yarn:
+
+```
+brew install yarn
+
+# or
+
+curl -o- -L https://yarnpkg.com/install.sh | bash
+```
+
+Then run the build script:
+
+```
+bash scripts/build_website.sh
+```
+
+This will build the docusaurus website and run a script to parse the tutorials and generate:
+- `.html` files in the `website/_tutorials` folder
+- `.js` files in the `website/pages/tutorials` folder
+- `.py`/`.ipynb` files in the `website/static/files` folder
+
+
+TODO: Add support for latex in markdown in jupyter notebooks and embedded images.
+
+## Build and publish the website
+
+To update for a new version, you need to build the tutorials and the website and push to the gh-pages
+branch of `github.com/facebookresearch/pytorch3d`. The instructions in `scripts/publish_website.sh`
+bring it all together.
+
+## Add a new tutorial
+
+The tutorials to include in the website are listed in `website/tutorials.json`. If you create a new tutorial add an entry to the list in this file. This is needed in order to generate the sidebar for the tutorials page.
+
+
+## Edit the landing page
+
+To change the content of the landing page modify: `website/pages/en/index.js`.
+
+
+## Edit the tutorials page
+
+To change the content of the tutorials home page modify: `website/pages/tutorials/index.js`.
+
+
+---------------------------------------------------------
+
+## Docusaurus docs
+
+- [Get Started in 5 Minutes](#get-started-in-5-minutes)
+- [Directory Structure](#directory-structure)
+- [Editing Content](#editing-content)
+- [Adding Content](#adding-content)
+- [Full Documentation](#full-documentation)
+
+
+## Directory Structure
+
+Your project file structure should look something like this
+
+```
+my-docusaurus/
+  docs/
+    doc-1.md
+    doc-2.md
+    doc-3.md
+  website/
+    blog/
+      2016-3-11-oldest-post.md
+      2017-10-24-newest-post.md
+    core/
+    node_modules/
+    pages/
+    static/
+      css/
+      img/
+    package.json
+    sidebars.json
+    siteConfig.js
+```
+
+# Editing Content
+
+## Editing an existing docs page
+
+Edit docs by navigating to `docs/` and editing the corresponding document:
+
+`docs/doc-to-be-edited.md`
+
+```markdown
+---
+id: page-needs-edit
+title: This Doc Needs To Be Edited
+---
+
+Edit me...
+```
+
+For more information about docs, click [here](https://docusaurus.io/docs/en/navigation)
+
+## Editing an existing blog post
+
+Edit blog posts by navigating to `website/blog` and editing the corresponding post:
+
+`website/blog/post-to-be-edited.md`
+
+```markdown
+---
+id: post-needs-edit
+title: This Blog Post Needs To Be Edited
+---
+
+Edit me...
+```
+
+For more information about blog posts, click [here](https://docusaurus.io/docs/en/adding-blog)
+
+# Adding Content
+
+## Adding a new docs page to an existing sidebar
+
+1. Create the doc as a new markdown file in `/docs`, example `docs/newly-created-doc.md`:
+
+```md
+---
+id: newly-created-doc
+title: This Doc Needs To Be Edited
+---
+
+My new content here..
+```
+
+1. Refer to that doc's ID in an existing sidebar in `website/sidebars.json`:
+
+```javascript
+// Add newly-created-doc to the Getting Started category of docs
+{
+  "docs": {
+    "Getting Started": [
+      "quick-start",
+      "newly-created-doc" // new doc here
+    ],
+    ...
+  },
+  ...
+}
+```
+
+For more information about adding new docs, click [here](https://docusaurus.io/docs/en/navigation)
+
+## Adding a new blog post
+
+1. Make sure there is a header link to your blog in `website/siteConfig.js`:
+
+`website/siteConfig.js`
+
+```javascript
+headerLinks: [
+    ...
+    { blog: true, label: 'Blog' },
+    ...
+]
+```
+
+2. Create the blog post with the format `YYYY-MM-DD-My-Blog-Post-Title.md` in `website/blog`:
+
+`website/blog/2018-05-21-New-Blog-Post.md`
+
+```markdown
+---
+author: Frank Li
+authorURL: https://twitter.com/foobarbaz
+authorFBID: 503283835
+title: New Blog Post
+---
+
+Lorem Ipsum...
+```
+
+For more information about blog posts, click [here](https://docusaurus.io/docs/en/adding-blog)
+
+## Adding items to your site's top navigation bar
+
+1. Add links to docs, custom pages or external links by editing the headerLinks field of `website/siteConfig.js`:
+
+`website/siteConfig.js`
+
+```javascript
+{
+  headerLinks: [
+    ...
+    /* you can add docs */
+    { doc: 'my-examples', label: 'Examples' },
+    /* you can add custom pages */
+    { page: 'help', label: 'Help' },
+    /* you can add external links */
+    { href: 'https://github.com/facebook/docusaurus', label: 'GitHub' },
+    ...
+  ],
+  ...
+}
+```
+
+For more information about the navigation bar, click [here](https://docusaurus.io/docs/en/navigation)
+
+## Adding custom pages
+
+1. Docusaurus uses React components to build pages. The components are saved as .js files in `website/pages/en`:
+1. If you want your page to show up in your navigation header, you will need to update `website/siteConfig.js` to add to the `headerLinks` element:
+
+`website/siteConfig.js`
+
+```javascript
+{
+  headerLinks: [
+    ...
+    { page: 'my-new-custom-page', label: 'My New Custom Page' },
+    ...
+  ],
+  ...
+}
+```
+
+For more information about custom pages, click [here](https://docusaurus.io/docs/en/custom-pages).
+
+# Full Documentation
+
+Full documentation can be found on the [website](https://docusaurus.io/).

third_party/AnyBimanual/third_party/pytorch3d/website/core/Footer.js

@@ -0,0 +1,91 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+const PropTypes = require("prop-types");
+const React = require('react');
+
+function SocialFooter(props) {
+  const repoUrl = `https://github.com/${props.config.organizationName}/${props.config.projectName}`;
+  return (
+    <div className="footerSection">
+      <div className="social">
+        <a
+          className="github-button" // part of the https://buttons.github.io/buttons.js script in siteConfig.js
+          href={repoUrl}
+          data-count-href={`${repoUrl}/stargazers`}
+          data-show-count="true"
+          data-count-aria-label="# stargazers on GitHub"
+          aria-label="Star PyTorch3D on GitHub"
+        >
+          {props.config.projectName}
+        </a>
+      </div>
+    </div>
+  );
+}
+
+SocialFooter.propTypes = {
+  config: PropTypes.object
+};
+
+class Footer extends React.Component {
+  docUrl(doc, language) {
+    const baseUrl = this.props.config.baseUrl;
+    const docsUrl = this.props.config.docsUrl;
+    const docsPart = `${docsUrl ? `${docsUrl}/` : ''}`;
+    const langPart = `${language ? `${language}/` : ''}`;
+    return `${baseUrl}${docsPart}${langPart}${doc}`;
+  }
+
+  pageUrl(doc, language) {
+    const baseUrl = this.props.config.baseUrl;
+    return baseUrl + (language ? `${language}/` : '') + doc;
+  }
+
+  render() {
+    const repoUrl = `https://github.com/${this.props.config.organizationName}/${this.props.config.projectName}`;
+    return (
+      <footer className="nav-footer" id="footer">
+        <section className="sitemap">
+          <SocialFooter config={this.props.config} />
+        </section>
+
+        <a
+          href="https://opensource.facebook.com/"
+          target="_blank"
+          rel="noreferrer noopener"
+          className="fbOpenSource">
+          <img
+            src={`${this.props.config.baseUrl}img/oss_logo.png`}
+            alt="Facebook Open Source"
+            width="170"
+            height="45"
+          />
+        </a>
+        <section className="copyright">{this.props.config.copyright}
+          <br/>
+          Legal:
+          <a
+            href="https://opensource.facebook.com/legal/privacy/"
+            target="_blank"
+            rel="noreferrer noopener">
+            Privacy
+          </a>
+          <a
+            href="https://opensource.facebook.com/legal/terms/"
+            target="_blank"
+            rel="noreferrer noopener">
+            Terms
+          </a>
+        </section>
+      </footer>
+    );
+  }
+}
+
+module.exports = Footer;

third_party/AnyBimanual/third_party/pytorch3d/website/core/Tutorial.js

@@ -0,0 +1,100 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @format
+ */
+
+const React = require('react');
+
+const fs = require('fs-extra');
+const path = require('path');
+const CWD = process.cwd();
+
+const CompLibrary = require(`${CWD}/node_modules/docusaurus/lib/core/CompLibrary.js`);
+const Container = CompLibrary.Container;
+
+const TutorialSidebar = require(`${CWD}/core/TutorialSidebar.js`);
+
+function renderDownloadIcon() {
+  return (
+    <svg
+      aria-hidden="true"
+      focusable="false"
+      data-prefix="fas"
+      data-icon="file-download"
+      className="svg-inline--fa fa-file-download fa-w-12"
+      role="img"
+      xmlns="http://www.w3.org/2000/svg"
+      viewBox="0 0 384 512">
+      <path
+        fill="currentColor"
+        d="M224 136V0H24C10.7 0 0 10.7 0 24v464c0 13.3 10.7 24 24 24h336c13.3 0 24-10.7 24-24V160H248c-13.2 0-24-10.8-24-24zm76.45 211.36l-96.42 95.7c-6.65 6.61-17.39 6.61-24.04 0l-96.42-95.7C73.42 337.29 80.54 320 94.82 320H160v-80c0-8.84 7.16-16 16-16h32c8.84 0 16 7.16 16 16v80h65.18c14.28 0 21.4 17.29 11.27 27.36zM377 105L279.1 7c-4.5-4.5-10.6-7-17-7H256v128h128v-6.1c0-6.3-2.5-12.4-7-16.9z"
+      />
+    </svg>
+  );
+}
+
+class Tutorial extends React.Component {
+  render() {
+    const {baseUrl, tutorialID} = this.props;
+
+    const htmlFile = `${CWD}/_tutorials/${tutorialID}.html`;
+    const normalizedHtmlFile = path.normalize(htmlFile);
+
+    return (
+      <div className="docMainWrapper wrapper">
+        <TutorialSidebar currentTutorialID={tutorialID} />
+        <Container className="mainContainer">
+          <div className="tutorialButtonsWrapper">
+            <div className="tutorialButtonWrapper buttonWrapper">
+              <a
+                className="tutorialButton button"
+                download
+                href={`https://colab.research.google.com/github/facebookresearch/pytorch3d/blob/stable/docs/tutorials/${tutorialID}.ipynb`}
+                target="_blank">
+                <img
+                  className="colabButton"
+                  align="left"
+                  src={`${baseUrl}img/colab_icon.png`}
+                />
+                {'Run in Google Colab'}
+              </a>
+            </div>
+            <div className="tutorialButtonWrapper buttonWrapper">
+              <a
+                className="tutorialButton button"
+                download
+                href={`${baseUrl}files/${tutorialID}.ipynb`}
+                target="_blank">
+                {renderDownloadIcon()}
+                {'Download Tutorial Jupyter Notebook'}
+              </a>
+            </div>
+            <div className="tutorialButtonWrapper buttonWrapper">
+              <a
+                className="tutorialButton button"
+                download
+                href={`${baseUrl}files/${tutorialID}.py`}
+                target="_blank">
+                {renderDownloadIcon()}
+                {'Download Tutorial Source Code'}
+              </a>
+            </div>
+          </div>
+          <div
+            className="tutorialBody"
+            dangerouslySetInnerHTML={{
+              __html: fs.readFileSync(normalizedHtmlFile, {encoding: 'utf8'}),
+            }}
+          />
+        </Container>
+      </div>
+    );
+  }
+}
+
+module.exports = Tutorial;

third_party/AnyBimanual/third_party/pytorch3d/website/core/TutorialSidebar.js

@@ -0,0 +1,93 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @format
+ */
+
+const React = require('react');
+const fs = require('fs-extra');
+const path = require('path');
+const join = path.join;
+const CWD = process.cwd();
+
+const CompLibrary = require(join(
+  CWD,
+  '/node_modules/docusaurus/lib/core/CompLibrary.js',
+));
+const SideNav = require(join(
+  CWD,
+  '/node_modules/docusaurus/lib/core/nav/SideNav.js',
+));
+
+const Container = CompLibrary.Container;
+
+const OVERVIEW_ID = 'tutorial_overview';
+
+class TutorialSidebar extends React.Component {
+  render() {
+    const {currentTutorialID} = this.props;
+    const current = {
+      id: currentTutorialID || OVERVIEW_ID,
+    };
+
+    const toc = [
+      {
+        type: 'CATEGORY',
+        title: 'Tutorials',
+        children: [
+          {
+            type: 'LINK',
+            item: {
+              permalink: 'tutorials/',
+              id: OVERVIEW_ID,
+              title: 'Overview',
+            },
+          },
+        ],
+      },
+    ];
+
+    const jsonFile = join(CWD, 'tutorials.json');
+    const normJsonFile = path.normalize(jsonFile);
+    const json = JSON.parse(fs.readFileSync(normJsonFile, {encoding: 'utf8'}));
+
+    Object.keys(json).forEach(category => {
+      const categoryItems = json[category];
+      const items = [];
+      categoryItems.map(item => {
+        items.push({
+          type: 'LINK',
+          item: {
+            permalink: `tutorials/${item.id}`,
+            id: item.id,
+            title: item.title,
+          },
+        });
+      });
+
+      toc.push({
+        type: 'CATEGORY',
+        title: category,
+        children: items,
+      });
+    });
+
+    return (
+      <Container className="docsNavContainer" id="docsNav" wrapper={false}>
+        <SideNav
+          language={'tutorials'}
+          root={'tutorials'}
+          title="Tutorials"
+          contents={toc}
+          current={current}
+        />
+      </Container>
+    );
+  }
+}
+
+module.exports = TutorialSidebar;

third_party/AnyBimanual/third_party/pytorch3d/website/package.json

@@ -0,0 +1,14 @@
+{
+  "scripts": {
+    "examples": "docusaurus-examples",
+    "start": "docusaurus-start",
+    "build": "docusaurus-build",
+    "publish-gh-pages": "docusaurus-publish",
+    "write-translations": "docusaurus-write-translations",
+    "version": "docusaurus-version",
+    "rename-version": "docusaurus-rename-version"
+  },
+  "devDependencies": {
+    "docusaurus": "^1.14.4"
+  }
+}

third_party/AnyBimanual/third_party/pytorch3d/website/pages/en/help.js

@@ -0,0 +1,55 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+const React = require('react');
+
+const CompLibrary = require('../../core/CompLibrary.js');
+
+const Container = CompLibrary.Container;
+const GridBlock = CompLibrary.GridBlock;
+
+function Help(props) {
+  const {config: siteConfig, language = ''} = props;
+  const {baseUrl, docsUrl} = siteConfig;
+  const docsPart = `${docsUrl ? `${docsUrl}/` : ''}`;
+  const langPart = `${language ? `${language}/` : ''}`;
+  const docUrl = doc => `${baseUrl}${docsPart}${langPart}${doc}`;
+
+  const supportLinks = [
+    {
+      content: `Learn more using the [documentation on this site.](${docUrl(
+        'doc1.html',
+      )})`,
+      title: 'Browse Docs',
+    },
+    {
+      content: 'Ask questions about the documentation and project',
+      title: 'Join the community',
+    },
+    {
+      content: "Find out what's new with this project",
+      title: 'Stay up to date',
+    },
+  ];
+
+  return (
+    <div className="docMainWrapper wrapper">
+      <Container className="mainContainer documentContainer postContainer">
+        <div className="post">
+          <header className="postHeader">
+            <h1>Need help?</h1>
+          </header>
+          <p>This project is maintained by a dedicated group of people.</p>
+          <GridBlock contents={supportLinks} layout="threeColumn" />
+        </div>
+      </Container>
+    </div>
+  );
+}
+
+module.exports = Help;

third_party/AnyBimanual/third_party/pytorch3d/website/pages/en/index.js

@@ -0,0 +1,240 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+const React = require('react');
+
+const CompLibrary = require('../../core/CompLibrary.js');
+
+const MarkdownBlock = CompLibrary.MarkdownBlock; /* Used to read markdown */
+const Container = CompLibrary.Container;
+const GridBlock = CompLibrary.GridBlock;
+const bash = (...args) => `~~~bash\n${String.raw(...args)}\n~~~`;
+class HomeSplash extends React.Component {
+  render() {
+    const {siteConfig, language = ''} = this.props;
+    const {baseUrl, docsUrl} = siteConfig;
+    const docsPart = `${docsUrl ? `${docsUrl}/` : ''}`;
+    const langPart = `${language ? `${language}/` : ''}`;
+    const docUrl = doc => `${baseUrl}${docsPart}${langPart}${doc}`;
+
+    const SplashContainer = props => (
+      <div className="homeContainer">
+        <div className="homeSplashFade">
+          <div className="wrapper homeWrapper">{props.children}</div>
+        </div>
+      </div>
+    );
+
+    const Logo = props => (
+      <div className="splashLogo">
+        <img src={props.img_src} alt="Project Logo" />
+      </div>
+    );
+
+    const ProjectTitle = props => (
+      <h2 className="projectTitle">
+        <small>{props.tagline}</small>
+      </h2>
+    );
+
+    const PromoSection = props => (
+      <div className="section promoSection">
+        <div className="promoRow">
+          <div className="pluginRowBlock">{props.children}</div>
+        </div>
+      </div>
+    );
+
+    const Button = props => (
+      <div className="pluginWrapper buttonWrapper">
+        <a className="button" href={props.href} target={props.target}>
+          {props.children}
+        </a>
+      </div>
+    );
+
+    return (
+      <SplashContainer>
+        <Logo img_src={baseUrl + 'img/pytorch3dlogowhite.svg'} />
+        <div className="inner">
+          <ProjectTitle tagline={siteConfig.tagline} title={siteConfig.title} />
+          <PromoSection>
+            <Button href={docUrl('why_pytorch3d.html')}>Docs</Button>
+            <Button href={`${baseUrl}tutorials/`}>Tutorials</Button>
+            <Button href={'#quickstart'}>Get Started</Button>
+          </PromoSection>
+        </div>
+      </SplashContainer>
+    );
+  }
+}
+
+function SocialBanner() {
+  return (
+    <div className="socialBanner">
+      <div>
+        Support Ukraine 🇺🇦{' '}
+        <a href="https://opensource.fb.com/support-ukraine">
+          Help Provide Humanitarian Aid to Ukraine
+        </a>
+        .
+      </div>
+    </div>
+  );
+}
+
+class Index extends React.Component {
+  render() {
+    const {config: siteConfig, language = ''} = this.props;
+    const {baseUrl} = siteConfig;
+
+    const Block = props => (
+      <Container
+        padding={['bottom', 'top']}
+        id={props.id}
+        background={props.background}>
+        <GridBlock
+          align="center"
+          contents={props.children}
+          layout={props.layout}
+        />
+      </Container>
+    );
+
+    const Description = () => (
+      <Block background="light">
+        {[
+          {
+            content:
+              'This is another description of how this project is useful',
+            image: `${baseUrl}img/docusaurus.svg`,
+            imageAlign: 'right',
+            title: 'Description',
+          },
+        ]}
+      </Block>
+    );
+
+    const pre = '```';
+
+    const codeExample = `${pre}python
+from pytorch3d.utils import ico_sphere
+from pytorch3d.io import load_obj
+from pytorch3d.structures import Meshes
+from pytorch3d.ops import sample_points_from_meshes
+from pytorch3d.loss import chamfer_distance
+
+# Use an ico_sphere mesh and load a mesh from an .obj e.g. model.obj
+sphere_mesh = ico_sphere(level=3)
+verts, faces, _ = load_obj("model.obj")
+test_mesh = Meshes(verts=[verts], faces=[faces.verts_idx])
+
+# Differentiably sample 5k points from the surface of each mesh and then compute the loss.
+sample_sphere = sample_points_from_meshes(sphere_mesh, 5000)
+sample_test = sample_points_from_meshes(test_mesh, 5000)
+loss_chamfer, _ = chamfer_distance(sample_sphere, sample_test)
+    `;
+
+    const QuickStart = () => (
+      <div
+        className="productShowcaseSection"
+        id="quickstart"
+        style={{textAlign: 'center'}}>
+        <h2>Get Started</h2>
+        <Container>
+          <ol>
+            <li>
+              <strong>Install PyTorch3D  </strong> (following the instructions <a href="https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md">here</a>)
+            </li>
+            <li>
+              <strong>Try a few 3D operators  </strong>
+              e.g. compute the chamfer loss between two meshes:
+              <MarkdownBlock>{codeExample}</MarkdownBlock>
+            </li>
+          </ol>
+        </Container>
+      </div>
+    );
+
+    const Features = () => (
+      <div className="productShowcaseSection" style={{textAlign: 'center'}}>
+        <Block layout="fourColumn">
+          {[
+            {
+              content:
+                'Supports batching of 3D inputs of different sizes ' +
+                'such as meshes' ,
+              image: `${baseUrl}img/batching.svg`,
+              imageAlign: 'top',
+              title: 'Heterogeneous Batching',
+            },
+            {
+              content:
+                'Supports optimized implementations of ' +
+                'several common  functions for 3D data',
+              image: `${baseUrl}img/ops.png`,
+              imageAlign: 'top',
+              title: 'Fast 3D Operators',
+            },
+            {
+              content:
+                'Modular differentiable rendering API ' +
+                'with parallel implementations in ' +
+                'PyTorch, C++ and CUDA' ,
+              image: `${baseUrl}img/rendering.svg`,
+              imageAlign: 'top',
+              title: 'Differentiable Rendering',
+            },
+          ]}
+        </Block>
+      </div>
+    );
+
+    const Showcase = () => {
+      if ((siteConfig.users || []).length === 0) {
+        return null;
+      }
+
+      const showcase = siteConfig.users
+        .filter(user => user.pinned)
+        .map(user => (
+          <a href={user.infoLink} key={user.infoLink}>
+            <img src={user.image} alt={user.caption} title={user.caption} />
+          </a>
+        ));
+
+      const pageUrl = page => baseUrl + (language ? `${language}/` : '') + page;
+
+      return (
+        <div className="productShowcaseSection paddingBottom">
+          <h2>Who is Using This?</h2>
+          <p>This project is used by all these people</p>
+          <div className="logos">{showcase}</div>
+          <div className="more-users">
+            <a className="button" href={pageUrl('users.html')}>
+              More {siteConfig.title} Users
+            </a>
+          </div>
+        </div>
+      );
+    };
+
+    return (
+      <div>
+        <SocialBanner />
+        <HomeSplash siteConfig={siteConfig} language={language} />
+        <div className="landingPage mainContainer">
+          <Features />
+          <QuickStart />
+        </div>
+      </div>
+    );
+  }
+}
+
+module.exports = Index;

third_party/AnyBimanual/third_party/pytorch3d/website/pages/en/users.js

@@ -0,0 +1,49 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+const React = require('react');
+
+const CompLibrary = require('../../core/CompLibrary.js');
+
+const Container = CompLibrary.Container;
+
+class Users extends React.Component {
+  render() {
+    const {config: siteConfig} = this.props;
+    if ((siteConfig.users || []).length === 0) {
+      return null;
+    }
+
+    const editUrl = `${siteConfig.repoUrl}/edit/main/website/siteConfig.js`;
+    const showcase = siteConfig.users.map(user => (
+      <a href={user.infoLink} key={user.infoLink}>
+        <img src={user.image} alt={user.caption} title={user.caption} />
+      </a>
+    ));
+
+    return (
+      <div className="mainContainer">
+        <Container padding={['bottom', 'top']}>
+          <div className="showcaseSection">
+            <div className="prose">
+              <h1>Who is Using This?</h1>
+              <p>This project is used by many folks</p>
+            </div>
+            <div className="logos">{showcase}</div>
+            <p>Are you using this project?</p>
+            <a href={editUrl} className="button">
+              Add your company
+            </a>
+          </div>
+        </Container>
+      </div>
+    );
+  }
+}
+
+module.exports = Users;

third_party/AnyBimanual/third_party/pytorch3d/website/pages/tutorials/index.js

@@ -0,0 +1,83 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @format
+ */
+
+const React = require('react');
+
+const CWD = process.cwd();
+
+const CompLibrary = require(`${CWD}/node_modules/docusaurus/lib/core/CompLibrary.js`);
+const Container = CompLibrary.Container;
+const MarkdownBlock = CompLibrary.MarkdownBlock;
+
+const TutorialSidebar = require(`${CWD}/core/TutorialSidebar.js`);
+const bash = (...args) => `~~~bash\n${String.raw(...args)}\n~~~`;
+
+class TutorialHome extends React.Component {
+  render() {
+    return (
+      <div className="docMainWrapper wrapper">
+        <TutorialSidebar currentTutorialID={null} />
+        <Container className="mainContainer documentContainer postContainer">
+          <div className="post">
+            <header className="postHeader">
+              <h1 className="postHeaderTitle">
+                Welcome to the PyTorch3D Tutorials
+              </h1>
+            </header>
+            <p>
+              Here you can learn about the structure and applications of
+              PyTorch3D from examples which are in the form of ipython
+              notebooks.
+            </p>
+            <h3> Run interactively </h3>
+            <p>
+              At the top of each example you can find a button named{' '}
+              <strong>"Run in Google Colab"</strong> which will open the
+              notebook in{' '}
+              <a href="https://colab.research.google.com/notebooks/intro.ipynb">
+                {' '}
+                Google Colaboratory{' '}
+              </a>{' '}
+              where you can run the code directly in the browser with access to
+              GPU support - it looks like this:
+            </p>
+            <div className="tutorialButtonsWrapper">
+              <div className="tutorialButtonWrapper buttonWrapper">
+                <a className="tutorialButton button" target="_blank">
+                  <img
+                    className="colabButton"
+                    align="left"
+                    src="/img/colab_icon.png"
+                  />
+                  {'Run in Google Colab'}
+                </a>
+              </div>
+            </div>
+            <p>
+              {' '}
+              You can modify the code and experiment with varying different
+              settings. Remember to install the latest stable version of
+              PyTorch3D and its dependencies. Code to do this with pip is
+              provided in each notebook.{' '}
+            </p>
+            <h3> Run locally </h3>
+            <p>
+              {' '}
+              There is also a button to download the notebook and source code to
+              run it locally.{' '}
+            </p>
+          </div>
+        </Container>
+      </div>
+    );
+  }
+}
+
+module.exports = TutorialHome;

third_party/AnyBimanual/third_party/pytorch3d/website/sidebars.json

@@ -0,0 +1,9 @@
+{
+  "docs": {
+    "Introduction": ["why_pytorch3d"],
+    "Data": ["io", "meshes_io", "datasets", "batching"],
+    "Ops": ["cubify", "iou3d"],
+    "Visualization": ["visualization"],
+    "Renderer": ["renderer", "renderer_getting_started", "cameras"]
+  }
+}

third_party/AnyBimanual/third_party/pytorch3d/website/siteConfig.js

@@ -0,0 +1,91 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// @licenselint-loose-mode
+
+// See https://docusaurus.io/docs/site-config for all the possible
+// site configuration options.
+
+// List of projects/orgs using your project for the users page.
+const users = [
+  {
+    caption: 'User1',
+    // You will need to prepend the image path with your baseUrl
+    // if it is not '/', like: '/test-site/img/image.jpg'.
+    image: '/img/undraw_open_source.svg',
+    infoLink: 'https://www.facebook.com',
+    pinned: true,
+  },
+];
+
+const baseUrl = '/'
+
+const siteConfig = {
+  title: 'PyTorch3D', // Title for your website.
+  tagline: 'A library for deep learning with 3D data',
+  url: 'https://pytorch3d.org', // Your website URL
+  baseUrl: baseUrl, // Base URL for your project */
+  projectName: 'pytorch3d',
+  organizationName: 'facebookresearch',
+  customDocsPath: 'docs/notes',
+  headerLinks: [
+    {doc: 'why_pytorch3d', label: 'Docs'},
+    {page: 'tutorials', label: 'Tutorials'},
+    {href: "https://pytorch3d.readthedocs.io/", label: 'API'},
+    {href: "https://github.com/facebookresearch/pytorch3d", label: 'GitHub'},
+  ],
+
+  // If you have users set above, you add it here:
+  users,
+
+  /* path to images for header/footer */
+  headerIcon: 'img/pytorch3dfavicon.png',
+  footerIcon: 'img/pytorch3dfavicon.png',
+  favicon: 'img/pytorch3dfavicon.png',
+
+  /* Colors for website */
+  colors: {
+    primaryColor: '#812CE5',
+    secondaryColor: '#FFAF00',
+  },
+
+  // This copyright info is used in /core/Footer.js and blog RSS/Atom feeds.
+  copyright: `Copyright \u{00A9} ${new Date().getFullYear()} Meta Platforms, Inc`,
+
+  highlight: {
+    // Highlight.js theme to use for syntax highlighting in code blocks.
+    theme: 'default',
+  },
+
+  // Add custom scripts here that would be placed in <script> tags.
+  scripts: ['https://buttons.github.io/buttons.js'],
+
+  // On page navigation for the current documentation page.
+  onPageNav: 'separate',
+  // No .html extensions for paths.
+  cleanUrl: true,
+
+  // Open Graph and Twitter card images.
+  ogImage: 'img/pytorch3dlogoicon.svg',
+  twitterImage: 'img/pytorch3dlogoicon.svg',
+
+   // Google analytics
+   gaTrackingId: 'UA-157376881-1',
+
+  // For sites with a sizable amount of content, set collapsible to true.
+  // Expand/collapse the links and subcategories under categories.
+  // docsSideNavCollapsible: true,
+
+  // Show documentation's last contributor's name.
+  enableUpdateBy: true,
+
+  // Show documentation's last update time.
+  // enableUpdateTime: true,
+};
+
+module.exports = siteConfig;

third_party/AnyBimanual/third_party/pytorch3d/website/static/css/custom.css

@@ -0,0 +1,360 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+html body {
+  font-family: 'Montserrat', sans-serif;
+  overflow-x: hidden;
+}
+
+.fixedHeaderContainer {
+  background-color: #222222;
+}
+
+
+.fixedHeaderContainer header .headerTitleWithLogo {
+  display: block;
+  color: #ffffff;
+}
+
+.fixedHeaderContainer header .logo {
+  height: 50px;
+}
+
+.fixedHeaderContainer header a:nth-child(2) {
+  position: absolute;
+  right: 0px;
+}
+
+.fixedHeaderContainer header a:nth-child(2) h3 {
+  font-size: 14px;
+}
+
+.fixedHeaderContainer header a:nth-child(2) h3::before {
+  content: 'v: ';
+}
+
+.navigationSlider {
+  margin-right: 80px;
+}
+
+.navigationSlider .slidingNav ul {
+  background: #222222;
+}
+
+.navigationSlider .slidingNav ul li a {
+  color: #c7d4fd;
+}
+
+.navigationSlider .slidingNav ul li a:hover,
+.navigationSlider .slidingNav ul li a:focus {
+  color: #ffffff;
+  background-color: inherit;
+}
+
+.navigationSlider .slidingNav ul li.siteNavItemActive > a,
+.navigationSlider .slidingNav ul li.siteNavGroupActive > a {
+  background-color: inherit;
+}
+
+.homeContainer {
+  background: linear-gradient(
+    rgba(129, 44, 229, 1) 0%,
+    rgba(255, 175, 0, 1) 100%
+  );
+  padding: 25px 0px;
+}
+
+.splashLogo {
+  display: block;
+  margin: 0 auto;
+  width: 65%;
+}
+
+.projectTitle {
+  color: #ffffff;
+  font-variant: small-caps;
+  font-weight: 300;
+}
+
+.promoSection .button {
+  border: 2px solid #fff;
+  color: #ffffff;
+  font-size: 19px;
+  margin: 10px;
+}
+
+.promoSection .button:hover {
+  background: inherit;
+  border: 2px solid #ffffff;
+  color: #ffffff;
+}
+
+.landingPage {
+  padding: 0px;
+}
+
+
+.productShowcaseSection {
+   padding: 45px 20px 30px 20px;
+}
+
+div.productShowcaseSection {
+  color: #6c6c6c;
+  padding-top: 40px;
+}
+
+#quickstart {
+  padding-top: 80px;
+}
+
+.productShowcaseSection > h2 {
+  font-variant: small-caps;
+  font-weight: 360;
+  margin: 0px;
+  padding: 0px;
+  color: #5b1861;
+}
+
+.productShowcaseSection p {
+  font-weight: 360;
+}
+
+# Subtitles for key features
+.productShowcaseSection .blockContent > div span p {
+  font-size: 18px;
+}
+
+.productShowcaseSection div.container {
+  padding-bottom: 40px;
+  padding-top: 10px;
+  padding-left: 0px;
+  padding-right: 0px;
+}
+
+.productShowcaseSection img {
+  height: 100px;
+}
+
+.gridBlock .fourByGridBlock img {
+    max-width: 200%;
+}
+
+.productShowcaseSection li {
+  padding: 10px 0;
+}
+
+.productShowcaseSection pre {
+  margin: 10px 0;
+}
+
+.productShowcaseSection code {
+  background: #fff;
+}
+
+.container .wrapper .alignCenter h2 {
+  color: #222222;
+}
+
+div#quickstart {
+  background: #efefef;
+}
+
+div#quickstart ol {
+  margin-bottom: 0px;
+}
+
+.nav-footer {
+  background-color: #222222;
+}
+
+.nav-footer .sitemap a {
+  color: #c7d4fd;
+}
+
+.nav-footer .sitemap a:hover {
+  color: #ffffff;
+}
+
+.social {
+ text-align: center
+}
+
+a,
+p a {
+  color: #4872f9;
+}
+
+a:hover,
+p a:hover {
+  color: #4872f9;
+}
+
+.imageAlignTop .blockImage {
+    margin-bottom: 20px;
+    max-width: 200px;
+}
+
+/* Style tutorials */
+.tutorialBody {
+  margin-top: -20px;
+  color: #6c6c6c;
+}
+
+.tutorialBody h1 {
+  margin: 0px;
+}
+
+.tutorialBody h1,
+.tutorialBody h2,
+.tutorialBody h3 {
+  color: #222222;
+}
+
+.tutorialBody pre {
+  font-family: 'IBM Plex Mono', monospace;
+  font-size: 14px;
+  margin: 0px;
+}
+
+.tutorialBody .input_prompt,
+.tutorialBody .output_prompt {
+  color: darkred;
+  font-size: 12px;
+}
+
+.tutorialBody .highlight {
+  background: #f3f4f7;
+  padding: 10px 20px;
+  border: lightgray 1px solid;
+  border-radius: 3px;
+}
+
+.tutorialBody .cell {
+  margin: 20px;
+}
+
+.tutorialBody .output_stderr {
+  background-color: #fdede9;
+}
+
+.tutorialBody .anchor-link {
+  color: lightgray;
+}
+
+.tutorialBody iframe {
+  width: 100%;
+  height: 100vh;
+}
+
+.tutorialButtonWrapper,
+.tutorialRuntime {
+  margin: 20px;
+}
+
+.tutorialButtonWrapper {
+  float: left;
+  margin: 5px;
+}
+
+.colabButtonWrapper {
+  float: left;
+  margin: 5px;
+}
+
+.colabButtonWrapper img {
+ padding-right: 0.25em;
+}
+
+.colabButton {
+  width: 24px;
+}
+
+.tutorialButtonsWrapper {
+ display: flex;
+ align-items: center;
+ padding-bottom: 15px;
+}
+
+/* .tutorialButton {
+  color: #4872f9;
+  border: 1px solid #4872f9;
+}
+ */
+.tutorialButton svg {
+  height: 15px;
+  margin-right: 5px;
+}
+
+.tutorialButton:hover {
+  color: #4872f9;
+  background-color: inherit;
+}
+
+.wrapper {
+  max-width: 1400px;
+}
+
+
+@media only screen and (min-device-width: 360px) and (max-device-width: 736px) {
+  .fixedHeaderContainer header a:nth-child(2) {
+    position: absolute;
+    right: 150px;
+  }
+  .promoSection .button {
+    font-size: 12px;
+    margin: 3px;
+  }
+  .inner h2 {
+   margin-top: 0px;
+  }
+  .splashLogo {
+    width: 90%;
+  }
+  .headerTitleWithLogo {
+    display: block !important;
+  }
+  .blockContent > div span p {
+    margin-bottom: 30px
+  }
+  .productShowcaseSection div.container {
+    padding-top: 0px;
+  }
+  .productShowcaseSection > h2 {
+    padding-bottom: 20px;
+  }
+}
+
+@media only screen and (max-width: 1023px) {
+  .fixedHeaderContainer header a:nth-child(2) {
+    position: absolute;
+    right: 200px;
+  }
+}
+
+@media only screen and (min-width: 1024px) {
+}
+
+@media only screen and (min-width: 1400px) {
+}
+
+@media only screen and (min-width: 1500px) {
+}
+
+/* Social Banner */
+.socialBanner {
+  font-weight: bold;
+  font-size: 20px;
+  padding: 20px;
+  max-width: 768px;
+  margin: 0 auto;
+  text-align: center;
+}
+
+.socialBanner a {
+  text-decoration: underline;
+}

third_party/AnyBimanual/third_party/pytorch3d/website/static/css/pygments.css

@@ -0,0 +1,213 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+.highlight .hll {
+  background-color: #ffffcc;
+}
+.highlight .c {
+  color: #60a0b0;
+  font-style: italic;
+} /* Comment */
+.highlight .err {
+  border: 1px solid #ff0000;
+} /* Error */
+.highlight .k {
+  color: #007020;
+  font-weight: bold;
+} /* Keyword */
+.highlight .o {
+  color: #666666;
+} /* Operator */
+.highlight .cm {
+  color: #60a0b0;
+  font-style: italic;
+} /* Comment.Multiline */
+.highlight .cp {
+  color: #007020;
+} /* Comment.Preproc */
+.highlight .c1 {
+  color: #60a0b0;
+  font-style: italic;
+} /* Comment.Single */
+.highlight .cs {
+  color: #60a0b0;
+  background-color: #fff0f0;
+} /* Comment.Special */
+.highlight .gd {
+  color: #a00000;
+} /* Generic.Deleted */
+.highlight .ge {
+  font-style: italic;
+} /* Generic.Emph */
+.highlight .gr {
+  color: #ff0000;
+} /* Generic.Error */
+.highlight .gh {
+  color: #000080;
+  font-weight: bold;
+} /* Generic.Heading */
+.highlight .gi {
+  color: #00a000;
+} /* Generic.Inserted */
+.highlight .go {
+  color: #808080;
+} /* Generic.Output */
+.highlight .gp {
+  color: #c65d09;
+  font-weight: bold;
+} /* Generic.Prompt */
+.highlight .gs {
+  font-weight: bold;
+} /* Generic.Strong */
+.highlight .gu {
+  color: #800080;
+  font-weight: bold;
+} /* Generic.Subheading */
+.highlight .gt {
+  color: #0040d0;
+} /* Generic.Traceback */
+.highlight .kc {
+  color: #007020;
+  font-weight: bold;
+} /* Keyword.Constant */
+.highlight .kd {
+  color: #007020;
+  font-weight: bold;
+} /* Keyword.Declaration */
+.highlight .kn {
+  color: #007020;
+  font-weight: bold;
+} /* Keyword.Namespace */
+.highlight .kp {
+  color: #007020;
+} /* Keyword.Pseudo */
+.highlight .kr {
+  color: #007020;
+  font-weight: bold;
+} /* Keyword.Reserved */
+.highlight .kt {
+  color: #902000;
+} /* Keyword.Type */
+.highlight .m {
+  color: #40a070;
+} /* Literal.Number */
+.highlight .s {
+  color: #4070a0;
+} /* Literal.String */
+.highlight .na {
+  color: #4070a0;
+} /* Name.Attribute */
+.highlight .nb {
+  color: #007020;
+} /* Name.Builtin */
+.highlight .nc {
+  color: #0e84b5;
+  font-weight: bold;
+} /* Name.Class */
+.highlight .no {
+  color: #60add5;
+} /* Name.Constant */
+.highlight .nd {
+  color: #555555;
+  font-weight: bold;
+} /* Name.Decorator */
+.highlight .ni {
+  color: #d55537;
+  font-weight: bold;
+} /* Name.Entity */
+.highlight .ne {
+  color: #007020;
+} /* Name.Exception */
+.highlight .nf {
+  color: #06287e;
+} /* Name.Function */
+.highlight .nl {
+  color: #002070;
+  font-weight: bold;
+} /* Name.Label */
+.highlight .nn {
+  color: #0e84b5;
+  font-weight: bold;
+} /* Name.Namespace */
+.highlight .nt {
+  color: #062873;
+  font-weight: bold;
+} /* Name.Tag */
+.highlight .nv {
+  color: #bb60d5;
+} /* Name.Variable */
+.highlight .ow {
+  color: #007020;
+  font-weight: bold;
+} /* Operator.Word */
+.highlight .w {
+  color: #bbbbbb;
+} /* Text.Whitespace */
+.highlight .mf {
+  color: #40a070;
+} /* Literal.Number.Float */
+.highlight .mh {
+  color: #40a070;
+} /* Literal.Number.Hex */
+.highlight .mi {
+  color: #40a070;
+} /* Literal.Number.Integer */
+.highlight .mo {
+  color: #40a070;
+} /* Literal.Number.Oct */
+.highlight .sb {
+  color: #4070a0;
+} /* Literal.String.Backtick */
+.highlight .sc {
+  color: #4070a0;
+} /* Literal.String.Char */
+.highlight .sd {
+  color: #4070a0;
+  font-style: italic;
+} /* Literal.String.Doc */
+.highlight .s2 {
+  color: #4070a0;
+} /* Literal.String.Double */
+.highlight .se {
+  color: #4070a0;
+  font-weight: bold;
+} /* Literal.String.Escape */
+.highlight .sh {
+  color: #4070a0;
+} /* Literal.String.Heredoc */
+.highlight .si {
+  color: #70a0d0;
+  font-style: italic;
+} /* Literal.String.Interpol */
+.highlight .sx {
+  color: #c65d09;
+} /* Literal.String.Other */
+.highlight .sr {
+  color: #235388;
+} /* Literal.String.Regex */
+.highlight .s1 {
+  color: #4070a0;
+} /* Literal.String.Single */
+.highlight .ss {
+  color: #517918;
+} /* Literal.String.Symbol */
+.highlight .bp {
+  color: #007020;
+} /* Name.Builtin.Pseudo */
+.highlight .vc {
+  color: #bb60d5;
+} /* Name.Variable.Class */
+.highlight .vg {
+  color: #bb60d5;
+} /* Name.Variable.Global */
+.highlight .vi {
+  color: #bb60d5;
+} /* Name.Variable.Instance */
+.highlight .il {
+  color: #40a070;
+} /* Literal.Number.Integer.Long */