Add files using upload-large-folder tool

README.md

@@ -14,6 +14,12 @@ This pass is a label study, not a policy study. No `pi0.5` integration is includ
 
 ## What is in this upload
 
+- `external/`
+  - Full local benchmark snapshots copied from the RunPod workspace.
+  - `external/rlbench/`: local RLBench tree used for this run.
+  - `external/pyrep/`: local PyRep tree used for this run.
+  - `external/peract_bimanual/`: local PerAct bimanual tree used for context.
+  - `external/yarr/`: local YARR tree used for context.
 - `code/rr_label_study/`
   - Core study code, including dense replay, visibility metrics, pregrasp/extraction oracles, keyframe extraction, intervention checks, and summary metric computation.
 - `code/scripts/`
@@ -30,6 +36,44 @@ This pass is a label study, not a policy study. No `pi0.5` integration is includ
 - `MANIFEST.txt`
   - Flat file listing of the uploaded bundle contents.
 
+## Repository map
+
+Relevant entry points and where to look:
+
+- Benchmark snapshots
+  - `external/README.md`
+  - `external/rlbench/README.md`
+  - `external/rlbench/rlbench/bimanual_tasks/`
+  - `external/rlbench/rlbench/action_modes/`
+  - `external/pyrep/README.md`
+  - `external/pyrep/pyrep/`
+  - `external/peract_bimanual/`
+  - `external/yarr/`
+- Study code
+  - `code/rr_label_study/oven_study.py`
+  - `code/scripts/run_oven_label_study.py`
+  - `code/scripts/launch_parallel_oven_label_study.py`
+  - `code/scripts/run_oven_single_frame.py`
+  - `code/scripts/repair_oven_episode_dense.py`
+- Final clean artifact
+  - `artifacts/results/oven_episode0_repaired_v1/episode0.dense.csv`
+  - `artifacts/results/oven_episode0_repaired_v1/episode0.keyframes.csv`
+  - `artifacts/results/oven_episode0_repaired_v1/episode0.metrics.json`
+  - `artifacts/results/oven_episode0_repaired_v1/summary.json`
+- Intermediate/debug artifacts
+  - `artifacts/results/oven_episode0_full*/`
+  - `artifacts/results/oven_to240_*/`
+  - `artifacts/results/oven_episode0_independent_v1/`
+  - `artifacts/results/parallel_smoke_2x10/`
+- Environment/repro
+  - `environment/system_info.txt`
+  - `environment/repo_revisions.txt`
+  - `environment/conda_env_rlbench.yml`
+  - `environment/pip_freeze_rlbench.txt`
+  - `environment/setup_same_hardware.sh`
+  - `environment/activate_rlbench_runtime.sh`
+  - `environment/dataset_notes.txt`
+
 ## Final validated artifact
 
 The clean single-episode artifact is:
@@ -122,11 +166,13 @@ The local run used:
 - `markusgrotz/peract_bimanual` at `bb0232a6ba3fe116566e9568f0c7af980ed6703d`
 - `markusgrotz/YARR` at `6822ff78602c77878b27d4cfe759ce029c67bffb`
 
+Those exact local source snapshots are also included under `external/`.
+
 ## Reproducing on the same hardware class
 
 1. Read `environment/dataset_notes.txt`.
 2. Run `environment/setup_same_hardware.sh /workspace`.
-3. Source `environment/activate_rlbench_runtime.sh /workspace`.
+   3. Source `environment/activate_rlbench_runtime.sh /workspace`.
 4. Run the dense study:
 
 ```bash
@@ -160,3 +206,7 @@ On a single deterministic episode, normalized time can become a degenerate perfe
 ## Dataset note
 
 The upstream RLBench demonstration dataset itself is not re-uploaded in this bundle. This repo contains the study code and all artifacts generated from the local run. The expected dataset path is documented in `environment/dataset_notes.txt`.
+
+The cloned benchmark code is included directly in this upload under `external/`.
+
+CoppeliaSim binaries are not included in this repo. The setup helpers expect a local extraction at `/workspace/coppelia_sim`.

external/README.md

@@ -0,0 +1,26 @@
+# External Benchmark Snapshots
+
+This directory contains the local benchmark/source trees copied from the RunPod workspace used for the study.
+
+Included trees:
+
+- `rlbench/`
+  - Source snapshot of `/workspace/rlbench`
+  - Upstream: `https://github.com/markusgrotz/RLBench.git`
+  - Commit: `8af748c51287989294e00c9c670e3330a0e35ed5`
+- `pyrep/`
+  - Source snapshot of `/workspace/pyrep`
+  - Upstream: `https://github.com/markusgrotz/PyRep.git`
+  - Commit: `b8bd1d7a3182adcd570d001649c0849047ebf197`
+- `peract_bimanual/`
+  - Source snapshot of `/workspace/peract_bimanual`
+  - Upstream: `https://github.com/markusgrotz/peract_bimanual.git`
+  - Commit: `bb0232a6ba3fe116566e9568f0c7af980ed6703d`
+- `yarr/`
+  - Source snapshot of `/workspace/yarr`
+  - Upstream: `https://github.com/markusgrotz/YARR.git`
+  - Commit: `6822ff78602c77878b27d4cfe759ce029c67bffb`
+
+These are source snapshots, not git clones with `.git/` metadata.
+
+See `../environment/repo_revisions.txt` for the recorded origin URLs and revisions.

external/peract_bimanual/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.

external/peract_bimanual/README.md

@@ -0,0 +1,300 @@
+# Perceiver-Actor^2: A Multi-Task Transformer for Bimanual Robotic Manipulation Tasks
+
+[![Code style](https://img.shields.io/badge/code%20style-black-black)](https://black.readthedocs.io/en/stable/)
+
+This work extends previous work [PerAct](https://peract.github.io) as well as
+[RLBench](https://sites.google.com/view/rlbench) for bimanual manipulation
+tasks.
+
+The repository and documentation are still work in progress.
+
+For the latest updates, see: [bimanual.github.io](https://bimanual.github.io)
+
+
+## Installation
+
+
+Please see [Installation](INSTALLATION.md) for further details.
+
+### Prerequisites
+
+The code PerAct^2 is built-off the [PerAct](https://peract.github.io) which itself is
+built on the [ARM repository](https://github.com/stepjam/ARM) by James et al.
+The prerequisites are the same as PerAct or ARM. 
+
+#### 1. Environment
+
+
+Install miniconda if not already present on the current system.
+You can use `scripts/install_conda.sh` for this step:
+
+```bash
+
+sudo apt install curl 
+
+curl -L -O https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
+chmod +x Miniconda3-latest-Linux-x86_64.sh 
+./Miniconda3-latest-Linux-x86_64.sh
+
+SHELL_NAME=`basename $SHELL`
+eval "$($HOME/miniconda3/bin/conda shell.${SHELL_NAME} hook)"
+conda init ${SHELL_NAME}
+conda install mamba -c conda-forge
+conda config --set auto_activate_base false
+```
+
+Next, create the rlbench environment and install the dependencies
+
+```bash
+conda create -n rlbench python=3.8
+conda activate rlbench
+conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
+```
+
+
+#### 2. Dependencies
+
+You need to setup  [RLBench](https://github.com/markusgrotz/rlbench/), [Pyrep](https://github.com/markusgrotz/Pyrep/), and [YARR](https://github.com/markusgrotz/YARR/).
+Please note that due to the bimanual functionallity the main repository does not work.
+You can use `scripts/install_dependencies.sh` to do so.
+See [Installation](INSTALLATION.md) for details.
+
+```bash
+./scripts/install_dependencies.sh
+```
+
+
+
+### Pre-Generated Datasets
+
+
+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
+
+
+#### Single-GPU Training
+
+To configure and train the model, follow these guidelines:
+
+- **General Parameters**: You can find and modify general parameters in the `conf/config.yaml` file. This file contains overall settings for the training environment, such as the number of cameras or the the tasks to use.
+
+- **Method-Specific Parameters**: For parameters specific to each method, refer to the corresponding files located in the `conf/method` directory. These files define configurations tailored to each method's requirements.
+
+
+
+When training adjust the `replay.batch_size` parameter to maximize the utilization of your GPU resources. Increasing this value can improve training efficiency based on the capacity of your available hardware.
+You can either modify the config files directly or you can pass parameters directly through the command line when running the training script. This allows for quick adjustments without editing configuration files:
+
+```bash
+python train.py replay.batch_size=3 method=BIMANUAL_PERACT
+```
+
+In this example, the command sets replay.batch_size to 3 and specifies the use of the BIMANUAL_PERACT method for training.
+Another important parameter to specify the tasks is `rlbench.task_name`, which sets the overall task, and `rlbench.tasks`, which is a list of tasks used for training. Note that these can be different for evaluation.
+A complete set of tasks is shown below:
+
+```yaml
+
+rlbench:
+  task_name: multi
+  tasks:
+  - bimanual_push_box
+  - bimanual_lift_ball
+  - bimanual_dual_push_buttons
+  - bimanual_pick_plate
+  - bimanual_put_item_in_drawer
+  - bimanual_put_bottle_in_fridge
+  - bimanual_handover_item
+  - bimanual_pick_laptop
+  - bimanual_straighten_rope
+  - bimanual_sweep_to_dustpan
+  - bimanual_lift_tray
+  - bimanual_handover_item_easy
+  - bimanual_take_tray_out_of_oven
+```
+
+
+#### Multi-GPU and Multi-Node Training
+
+This repository supports multi-GPU training and distributed training across multiple nodes using [PyTorch Distributed Data Parallel (DDP)](https://pytorch.org/docs/stable/notes/ddp.html). 
+Follow the instructions below to configure and run training across multiple GPUs and nodes.
+
+#### Multi-GPU Training on a Single Node
+
+To train using multiple GPUs on a single node, set the parameter `ddp.num_devices` to the number of GPUs available. For example, if you have 4 GPUs, you can start the training process as follows:
+
+```bash
+python train.py replay.batch_size=3 method=BIMANUAL_PERACT ddp.num_devices=4
+```
+
+This command will utilize 4 GPUs on the current node for training. Remember to set the `replay.batch_size`, which is per GPU.
+
+#### Multi-Node Training Across Different Nodes
+
+If you want to perform distributed training across multiple nodes, you need to set additional parameters: ddp.master_addr and ddp.master_port. These parameters should be configured as follows:
+
+`ddp.master_addr`: The IP address of the master node (usually the node where the training is initiated).
+`ddp.master_port`: A port number to be used for communication across nodes.
+
+Example Command:
+
+```bash
+python train.py replay.batch_size=3 method=BIMANUAL_PERACT ddp.num_devices=4 ddp.master_addr=192.168.1.1 ddp.master_port=29500
+```
+
+Note: Ensure that all nodes can communicate with each other through the specified IP and port, and that they have the same codebase, data access, and configurations for a successful distributed training run.
+
+
+
+### Evaluation
+
+
+Similar to training you can find general parameters in  `conf/eval.yaml` and method specific parameters in the `conf/method` directory.
+For each method, you have to set the execution mode in RLBench. For bimanual agents such as `BIMANUAL_PERACT` or `PERACT_BC` this is:
+
+```yaml
+rlbench:
+  gripper_mode: 'BimanualDiscrete'
+  arm_action_mode: 'BimanualEndEffectorPoseViaPlanning'
+  action_mode: 'BimanualMoveArmThenGripper'
+```
+
+
+To generate videos of the current evaluation you can set `cinematic_recorder.enabled` to `True`.
+It is recommended during evalution to disable the recorder, i.e. `cinematic_recorder.enabled=False`, as rendering the video increases the total evaluation time.
+
+
+## Acknowledgements
+
+This repository uses code from the following open-source projects:
+
+#### ARM 
+Original:  [https://github.com/stepjam/ARM](https://github.com/stepjam/ARM)  
+License: [ARM License](https://github.com/stepjam/ARM/LICENSE)    
+Changes: Data loading was modified for PerAct. Voxelization code was modified for DDP training.
+
+#### PerceiverIO
+Original: [https://github.com/lucidrains/perceiver-pytorch](https://github.com/lucidrains/perceiver-pytorch)   
+License: [MIT](https://github.com/lucidrains/perceiver-pytorch/blob/main/LICENSE)   
+Changes: PerceiverIO adapted for 6-DoF manipulation.
+
+#### ViT
+Original: [https://github.com/lucidrains/vit-pytorch](https://github.com/lucidrains/vit-pytorch)     
+License: [MIT](https://github.com/lucidrains/vit-pytorch/blob/main/LICENSE)   
+Changes: ViT adapted for baseline.   
+
+#### LAMB Optimizer
+Original: [https://github.com/cybertronai/pytorch-lamb](https://github.com/cybertronai/pytorch-lamb)   
+License: [MIT](https://github.com/cybertronai/pytorch-lamb/blob/master/LICENSE)   
+Changes: None.
+
+#### OpenAI CLIP
+Original: [https://github.com/openai/CLIP](https://github.com/openai/CLIP)  
+License: [MIT](https://github.com/openai/CLIP/blob/main/LICENSE)  
+Changes: Minor modifications to extract token and sentence features.
+
+Thanks for open-sourcing! 
+
+## Licenses
+- [PerAct License (Apache 2.0)](LICENSE) - Perceiver-Actor Transformer
+- [ARM License](ARM_LICENSE) - Voxelization and Data Preprocessing 
+- [YARR Licence (Apache 2.0)](https://github.com/stepjam/YARR/blob/main/LICENSE)
+- [RLBench Licence](https://github.com/stepjam/RLBench/blob/master/LICENSE)
+- [PyRep License (MIT)](https://github.com/stepjam/PyRep/blob/master/LICENSE)
+- [Perceiver PyTorch License (MIT)](https://github.com/lucidrains/perceiver-pytorch/blob/main/LICENSE)
+- [LAMB License (MIT)](https://github.com/cybertronai/pytorch-lamb/blob/master/LICENSE)
+- [CLIP License (MIT)](https://github.com/openai/CLIP/blob/main/LICENSE)
+
+## Release Notes
+
+**Update 2025-02-20**
+
+- Update instructions
+- Add missing dependency for install script
+- Add docker build file
+
+**Update 2024-11-06**
+
+- Regenerat and repack dataset. Closes #13. Task names are now more consistent. Dataset now includes waypoint information.
+
+
+**Update 2024-10-17**
+
+- Update Readme
+
+
+
+**Update 2024-07-10**
+
+- Initial release
+
+
+## Citations 
+
+
+**PerAct^2**
+```
+@misc{grotz2024peract2,
+      title={PerAct2: Benchmarking and Learning for Robotic Bimanual Manipulation Tasks}, 
+      author={Markus Grotz and Mohit Shridhar and Tamim Asfour and Dieter Fox},
+      year={2024},
+      eprint={2407.00278},
+      archivePrefix={arXiv},
+      primaryClass={cs.RO},
+      url={https://arxiv.org/abs/2407.00278}, 
+}
+```
+
+**PerAct**
+```
+@inproceedings{shridhar2022peract,
+  title     = {Perceiver-Actor: A Multi-Task Transformer for Robotic Manipulation},
+  author    = {Shridhar, Mohit and Manuelli, Lucas and Fox, Dieter},
+  booktitle = {Proceedings of the 6th Conference on Robot Learning (CoRL)},
+  year      = {2022},
+}
+```
+
+**C2FARM**
+```
+@inproceedings{james2022coarse,
+  title={Coarse-to-fine q-attention: Efficient learning for visual robotic manipulation via discretisation},
+  author={James, Stephen and Wada, Kentaro and Laidlow, Tristan and Davison, Andrew J},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={13739--13748},
+  year={2022}
+}
+```
+
+**PerceiverIO**
+```
+@article{jaegle2021perceiver,
+  title={Perceiver io: A general architecture for structured inputs \& outputs},
+  author={Jaegle, Andrew and Borgeaud, Sebastian and Alayrac, Jean-Baptiste and Doersch, Carl and Ionescu, Catalin and Ding, David and Koppula, Skanda and Zoran, Daniel and Brock, Andrew and Shelhamer, Evan and others},
+  journal={arXiv preprint arXiv:2107.14795},
+  year={2021}
+}
+```
+
+**RLBench**
+```
+@article{james2020rlbench,
+  title={Rlbench: The robot learning benchmark \& learning environment},
+  author={James, Stephen and Ma, Zicong and Arrojo, David Rovick and Davison, Andrew J},
+  journal={IEEE Robotics and Automation Letters},
+  volume={5},
+  number={2},
+  pages={3019--3026},
+  year={2020},
+  publisher={IEEE}
+}
+```
+
+## Questions or Issues?
+
+Please file an issue with the issue tracker.  

external/peract_bimanual/model-card.md

@@ -0,0 +1,47 @@
+# Model Card: Perceiver-Actor
+
+Following [Model Cards for Model Reporting (Mitchell et al.)](https://arxiv.org/abs/1810.03993) and [Lessons from Archives (Jo & Gebru)](https://arxiv.org/pdf/1912.10389.pdf) we provide additional information on PerAct.
+
+## Model Details
+
+
+### Overview
+- Developed by Shridhar et al. at University of Washington and NVIDIA. PerAct is an end-to-end behavior cloning agent that learns to perform a wide variety of language-conditioned manipulation tasks. PerAct uses a Transformer that exploits the 3D structure of _voxel patches_ to learn policies with just a few demonstrations per task.
+- Architecture: Transformer trained from scratch with end-to-end supervised learning.
+- Trained for 6-DoF manipulation tasks with objects that appear in tabletop scenes.
+
+### Model Date
+
+Nov 2022
+
+### Documents
+
+- [PerAct Paper](https://peract.github.io/paper/peract_corl2022.pdf)
+- [PerceiverIO Paper](https://arxiv.org/abs/2107.14795)
+- [C2FARM Paper](https://arxiv.org/abs/2106.12534)
+
+
+## Model Use
+
+- **Primary intended use case**: PerAct is intended for robotic manipulation research. We hope the benchmark and pre-trained models will enable researchers to study the capabilities of Transformers for end-to-end 6-DoF Manipulation. Specifically, we hope the setup serves a reproducible framework for evaluating  robustness and scaling capabilities of manipulation agents. 
+- **Primary intended users**: Robotics researchers. 
+- **Out-of-scope use cases**: Deployed use cases in real-world autonomous systems without human supervision during test-time is currently out-of-scope. Use cases that involve manipulating novel objects and observations with people, are not recommended for safety-critical systems. The agent is also intended to be trained and evaluated with English language instructions.
+
+## Data
+
+- Pre-training Data for CLIP's language encoder: See [OpenAI's Model Card](https://github.com/openai/CLIP/blob/main/model-card.md#data) for full details. **Note:** We do not use CLIP's vision encoders for any agents in the repo. 
+- Manipulation Data for PerAct: The agent was trained with expert demonstrations. In simulation, we use oracle agents and in real-world we use human demonstrations. Since the agent is used in few-shot settings with very limited data, the agent might exploit intended and un-intented biases in the training demonstrations. Currently, these biases are limited to just objects that appear on tabletops.
+
+
+## Limitations
+
+- Depends on a sampling-based motion planner.
+- Hard to extend to dexterous and continuous manipulation tasks.
+- Lacks memory to solve tasks with ordering and history-based sequencing. 
+- Exploits biases in training demonstrations.
+- Needs good hand-eye calibration.
+- Doesn't generalize to novel objects.
+- Struggles with grounding complex spatial relationships. 
+- Does not predict task completion.
+
+See Appendix L in the [paper](https://peract.github.io/paper/peract_corl2022.pdf) for an extended discussion.

external/peract_bimanual/peract_config.py

@@ -0,0 +1,32 @@
+"""
+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)

external/peract_bimanual/pyproject.toml

@@ -0,0 +1,35 @@
+[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"
+transformers = "^4.21"
+
+
+
+[tool.poetry.extras]
+docs = ["sphinx"]
+
+[build-system]
+requires = ["setuptools", "wheel", "poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"

external/peract_bimanual/run_seed_fn.py

@@ -0,0 +1,218 @@
+import os
+import pickle
+import gc
+from typing import List
+
+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 helpers.custom_rlbench_env import CustomRLBenchEnv, CustomMultiTaskRLBenchEnv
+import torch.distributed as dist
+
+from agents import agent_factory
+from agents import replay_utils
+
+import peract_config
+from functools import partial
+
+
+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 = "multi" if len(tasks) > 1 else tasks[0]
+    replay_path = os.path.join(
+        cfg.replay.path, task_folder, cfg.method.name, "seed%d" % seed
+    )
+
+    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")
+    ):
+        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,
+    )
+
+    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()

external/peract_bimanual/train.py

@@ -0,0 +1,116 @@
+from typing import List
+import logging
+import os
+import sys
+from datetime import datetime
+
+import peract_config
+
+import hydra
+from omegaconf import DictConfig, OmegaConf, ListConfig
+
+import run_seed_fn
+from helpers.observation_utils import create_obs_config
+
+import torch.multiprocessing as mp
+
+
+@hydra.main(config_name="config", config_path="conf")
+def main(cfg: DictConfig) -> None:
+    cfg_yaml = OmegaConf.to_yaml(cfg)
+    logging.info("\n" + cfg_yaml)
+
+    peract_config.on_config(cfg)
+
+    cfg.rlbench.cameras = (
+        cfg.rlbench.cameras
+        if isinstance(cfg.rlbench.cameras, ListConfig)
+        else [cfg.rlbench.cameras]
+    )
+
+    # sanity check if rgb is not used as camera name
+    for camera_name in cfg.rlbench.cameras:
+        assert "rgb" not in camera_name
+
+    obs_config = create_obs_config(
+        cfg.rlbench.cameras, cfg.rlbench.camera_resolution, cfg.method.name
+    )
+
+    cwd = os.getcwd()
+    logging.info("CWD:" + os.getcwd())
+
+    if cfg.framework.start_seed >= 0:
+        # seed specified
+        start_seed = cfg.framework.start_seed
+    elif (
+        cfg.framework.start_seed == -1
+        and len(list(filter(lambda x: "seed" in x, os.listdir(cwd)))) > 0
+    ):
+        # unspecified seed; use largest existing seed plus one
+        largest_seed = max(
+            [
+                int(n.replace("seed", ""))
+                for n in list(filter(lambda x: "seed" in x, os.listdir(cwd)))
+            ]
+        )
+        start_seed = largest_seed + 1
+    else:
+        # start with seed 0
+        start_seed = 0
+
+    seed_folder = os.path.join(os.getcwd(), "seed%d" % start_seed)
+    os.makedirs(seed_folder, exist_ok=True)
+
+    start_time = datetime.now()
+    with open(os.path.join(seed_folder, "config.yaml"), "w") as f:
+        f.write(cfg_yaml)
+
+    # check if previous checkpoints already exceed the number of desired training iterations
+    # if so, exit the script
+    latest_weight = 0
+    weights_folder = os.path.join(seed_folder, "weights")
+    if os.path.isdir(weights_folder) and len(os.listdir(weights_folder)) > 0:
+        weights = os.listdir(weights_folder)
+        latest_weight = sorted(map(int, weights))[-1]
+        if latest_weight >= cfg.framework.training_iterations:
+            logging.info(
+                "Agent was already trained for %d iterations. Exiting." % latest_weight
+            )
+            sys.exit(0)
+
+    with open(os.path.join(seed_folder, "training.log"), "a") as f:
+        f.write(
+            f"# Starting training from weights: {latest_weight} to {cfg.framework.training_iterations}"
+        )
+        f.write(f"# Training started on: {start_time.isoformat()}")
+        f.write(os.linesep)
+
+    # run train jobs with multiple seeds (sequentially)
+    for seed in range(start_seed, start_seed + cfg.framework.seeds):
+        logging.info("Starting seed %d." % seed)
+
+        world_size = cfg.ddp.num_devices
+        mp.spawn(
+            run_seed_fn.run_seed,
+            args=(
+                cfg,
+                obs_config,
+                seed,
+                world_size,
+            ),
+            nprocs=world_size,
+            join=True,
+        )
+
+    end_time = datetime.now()
+    duration = end_time - start_time
+    with open(os.path.join(seed_folder, "training.log"), "a") as f:
+        f.write(f"# Training finished on: {end_time.isoformat()}")
+        f.write(f"# Took {duration.total_seconds()}")
+        f.write(os.linesep)
+        f.write(os.linesep)
+
+
+if __name__ == "__main__":
+    peract_config.on_init()
+    main()

external/peract_bimanual/voxel/voxel_grid.py

@@ -0,0 +1,252 @@
+# Voxelizer modified from ARM for DDP training
+# Source: https://github.com/stepjam/ARM
+# License: https://github.com/stepjam/ARM/LICENSE
+
+from functools import reduce
+from operator import mul
+
+import torch
+from torch import nn
+
+MIN_DENOMINATOR = 1e-12
+INCLUDE_PER_VOXEL_COORD = False
+
+
+class VoxelGrid(nn.Module):
+    def __init__(
+        self,
+        coord_bounds,
+        voxel_size: int,
+        device,
+        batch_size,
+        feature_size,  # e.g. rgb or image features
+        max_num_coords: int,
+    ):
+        super(VoxelGrid, self).__init__()
+        self._device = device
+        self._voxel_size = voxel_size
+        self._voxel_shape = [voxel_size] * 3
+        self._voxel_d = float(self._voxel_shape[-1])
+        self._voxel_feature_size = 4 + feature_size
+        self._voxel_shape_spec = (
+            torch.tensor(
+                self._voxel_shape,
+            ).unsqueeze(0)
+            + 2
+        )  # +2 because we crop the edges.
+        self._coord_bounds = torch.tensor(
+            coord_bounds,
+            dtype=torch.float,
+        ).unsqueeze(0)
+        max_dims = self._voxel_shape_spec[0]
+        self._total_dims_list = torch.cat(
+            [
+                torch.tensor(
+                    [batch_size],
+                ),
+                max_dims,
+                torch.tensor(
+                    [4 + feature_size],
+                ),
+            ],
+            -1,
+        ).tolist()
+
+        self.register_buffer(
+            "_ones_max_coords", torch.ones((batch_size, max_num_coords, 1))
+        )
+        self._num_coords = max_num_coords
+
+        shape = self._total_dims_list
+        result_dim_sizes = torch.tensor(
+            [reduce(mul, shape[i + 1 :], 1) for i in range(len(shape) - 1)] + [1],
+        )
+        self.register_buffer("_result_dim_sizes", result_dim_sizes)
+        flat_result_size = reduce(mul, shape, 1)
+
+        self._initial_val = torch.tensor(0, dtype=torch.float)
+        flat_output = (
+            torch.ones(flat_result_size, dtype=torch.float) * self._initial_val
+        )
+        self.register_buffer("_flat_output", flat_output)
+
+        self.register_buffer("_arange_to_max_coords", torch.arange(4 + feature_size))
+        self._flat_zeros = torch.zeros(flat_result_size, dtype=torch.float)
+
+        self._const_1 = torch.tensor(
+            1.0,
+        )
+        self._batch_size = batch_size
+
+        # Coordinate Bounds:
+        bb_mins = self._coord_bounds[..., 0:3]
+        self.register_buffer("_bb_mins", bb_mins)
+        bb_maxs = self._coord_bounds[..., 3:6]
+        bb_ranges = bb_maxs - bb_mins
+        # get voxel dimensions. 'DIMS' mode
+        self._dims = dims = self._voxel_shape_spec.int()
+        dims_orig = self._voxel_shape_spec.int() - 2
+        self.register_buffer("_dims_orig", dims_orig)
+
+        # self._dims_m_one = (dims - 1).int()
+        dims_m_one = (dims - 1).int()
+        self.register_buffer("_dims_m_one", dims_m_one)
+
+        # BS x 1 x 3
+        res = bb_ranges / (dims_orig.float() + MIN_DENOMINATOR)
+        self._res_minis_2 = bb_ranges / (dims.float() - 2 + MIN_DENOMINATOR)
+        self.register_buffer("_res", res)
+
+        voxel_indicy_denmominator = res + MIN_DENOMINATOR
+        self.register_buffer("_voxel_indicy_denmominator", voxel_indicy_denmominator)
+
+        self.register_buffer("_dims_m_one_zeros", torch.zeros_like(dims_m_one))
+
+        batch_indices = torch.arange(self._batch_size, dtype=torch.int).view(
+            self._batch_size, 1, 1
+        )
+        self.register_buffer(
+            "_tiled_batch_indices", batch_indices.repeat([1, self._num_coords, 1])
+        )
+
+        w = self._voxel_shape[0] + 2
+        arange = torch.arange(
+            0,
+            w,
+            dtype=torch.float,
+        )
+        index_grid = (
+            torch.cat(
+                [
+                    arange.view(w, 1, 1, 1).repeat([1, w, w, 1]),
+                    arange.view(1, w, 1, 1).repeat([w, 1, w, 1]),
+                    arange.view(1, 1, w, 1).repeat([w, w, 1, 1]),
+                ],
+                dim=-1,
+            )
+            .unsqueeze(0)
+            .repeat([self._batch_size, 1, 1, 1, 1])
+        )
+        self.register_buffer("_index_grid", index_grid)
+
+    def _broadcast(self, src: torch.Tensor, other: torch.Tensor, dim: int):
+        if dim < 0:
+            dim = other.dim() + dim
+        if src.dim() == 1:
+            for _ in range(0, dim):
+                src = src.unsqueeze(0)
+        for _ in range(src.dim(), other.dim()):
+            src = src.unsqueeze(-1)
+        src = src.expand_as(other)
+        return src
+
+    def _scatter_mean(
+        self, src: torch.Tensor, index: torch.Tensor, out: torch.Tensor, dim: int = -1
+    ):
+        out = out.scatter_add_(dim, index, src)
+
+        index_dim = dim
+        if index_dim < 0:
+            index_dim = index_dim + src.dim()
+        if index.dim() <= index_dim:
+            index_dim = index.dim() - 1
+
+        ones = torch.ones(index.size(), dtype=src.dtype, device=src.device)
+        out_count = torch.zeros(out.size(), dtype=out.dtype, device=out.device)
+        out_count = out_count.scatter_add_(index_dim, index, ones)
+        out_count.clamp_(1)
+        count = self._broadcast(out_count, out, dim)
+        if torch.is_floating_point(out):
+            out.true_divide_(count)
+        else:
+            out.floor_divide_(count)
+        return out
+
+    def _scatter_nd(self, indices, updates):
+        indices_shape = indices.shape
+        num_index_dims = indices_shape[-1]
+        flat_updates = updates.view((-1,))
+        indices_scales = self._result_dim_sizes[0:num_index_dims].view(
+            [1] * (len(indices_shape) - 1) + [num_index_dims]
+        )
+        indices_for_flat_tiled = (
+            ((indices * indices_scales).sum(dim=-1, keepdims=True))
+            .view(-1, 1)
+            .repeat(*[1, self._voxel_feature_size])
+        )
+
+        implicit_indices = (
+            self._arange_to_max_coords[: self._voxel_feature_size]
+            .unsqueeze(0)
+            .repeat(*[indices_for_flat_tiled.shape[0], 1])
+        )
+        indices_for_flat = indices_for_flat_tiled + implicit_indices
+        flat_indices_for_flat = indices_for_flat.view((-1,)).long()
+
+        flat_scatter = self._scatter_mean(
+            flat_updates, flat_indices_for_flat, out=torch.zeros_like(self._flat_output)
+        )
+        return flat_scatter.view(self._total_dims_list)
+
+    def coords_to_bounding_voxel_grid(
+        self, coords, coord_features=None, coord_bounds=None
+    ):
+        voxel_indicy_denmominator = self._voxel_indicy_denmominator
+        res, bb_mins = self._res, self._bb_mins
+        if coord_bounds is not None:
+            bb_mins = coord_bounds[..., 0:3]
+            bb_maxs = coord_bounds[..., 3:6]
+            bb_ranges = bb_maxs - bb_mins
+            res = bb_ranges / (self._dims_orig.float() + MIN_DENOMINATOR)
+            voxel_indicy_denmominator = res + MIN_DENOMINATOR
+
+        bb_mins_shifted = bb_mins - res  # shift back by one
+        floor = torch.floor(
+            (coords - bb_mins_shifted.unsqueeze(1))
+            / voxel_indicy_denmominator.unsqueeze(1)
+        ).int()
+        voxel_indices = torch.min(floor, self._dims_m_one)
+        voxel_indices = torch.max(voxel_indices, self._dims_m_one_zeros)
+
+        # BS x NC x 3
+        voxel_values = coords
+        if coord_features is not None:
+            voxel_values = torch.cat([voxel_values, coord_features], -1)
+
+        _, num_coords, _ = voxel_indices.shape
+        # BS x N x (num_batch_dims + 2)
+        all_indices = torch.cat(
+            [self._tiled_batch_indices[:, :num_coords], voxel_indices], -1
+        )
+
+        # BS x N x 4
+        voxel_values_pruned_flat = torch.cat(
+            [voxel_values, self._ones_max_coords[:, :num_coords]], -1
+        )
+
+        # BS x x_max x y_max x z_max x 4
+        scattered = self._scatter_nd(
+            all_indices.view([-1, 1 + 3]),
+            voxel_values_pruned_flat.view(-1, self._voxel_feature_size),
+        )
+
+        vox = scattered[:, 1:-1, 1:-1, 1:-1]
+        if INCLUDE_PER_VOXEL_COORD:
+            res_expanded = res.unsqueeze(1).unsqueeze(1).unsqueeze(1)
+            res_centre = (res_expanded * self._index_grid) + res_expanded / 2.0
+            coord_positions = (
+                res_centre + bb_mins_shifted.unsqueeze(1).unsqueeze(1).unsqueeze(1)
+            )[:, 1:-1, 1:-1, 1:-1]
+            vox = torch.cat([vox[..., :-1], coord_positions, vox[..., -1:]], -1)
+
+        occupied = (vox[..., -1:] > 0).float()
+        vox = torch.cat([vox[..., :-1], occupied], -1)
+
+        return torch.cat(
+            [
+                vox[..., :-1],
+                self._index_grid[:, :-2, :-2, :-2] / self._voxel_d,
+                vox[..., -1:],
+            ],
+            -1,
+        )

external/yarr/.gitignore

@@ -0,0 +1,13 @@
+__pycache__
+venv
+.idea
+.bash_history
+.cache/
+.local/
+.python_history
+nvidia-persistenced/
+results/
+rlight.egg-info
+dist/
+build/
+yarr.egg-info

external/yarr/LICENSE

@@ -0,0 +1,201 @@
+                                 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.

external/yarr/README.md

@@ -0,0 +1,28 @@
+![Logo Missing](logo.png)
+
+**Note**: Pirate qualification not needed to use this library.
+
+YARR is **Y**et **A**nother **R**obotics and **R**einforcement learning framework for PyTorch.
+
+The framework allows for asynchronous training (i.e. agent and learner running in separate processes), which makes it suitable for robot learning.
+For an example of how to use this framework, see my [Attention-driven Robot Manipulation (ARM) repo](https://github.com/stepjam/ARM).
+
+This project is mostly intended for my personal use (Stephen James) and facilitate my research.
+
+## Modifcations
+
+This is my (Mohit Shridhar) fork of YARR. Honestly, I don't understand what exactly is happening in a lot of places, so there a lot of hacks to make it work for my purposes. If you are doing simple behavior cloning, you can probably write simpler training and evaluation routines, but YARR might be useful if you also want to do RL. Here is a quick summary of my modifcations:
+
+- Switched from randomly sampling evaluation episodes to deterministic reloading of val/test dataset episodes for one-to-one comparisons across models.
+- Separated training and evaluation routines. 
+- Task-uniform replay buffer for multi-task training. Each batch has a uniform distribution of tasks. 
+- Added cinematic recorder for rollouts.
+- Some other weird hacks to prevent memory leaks.
+
+## Install
+
+Ensure you have [PyTorch installed](https://pytorch.org/get-started/locally/).
+Then simply run:
+```bash
+python setup.py develop
+```

external/yarr/requirements.txt

@@ -0,0 +1,11 @@
+tensorboard
+moviepy
+natsort
+psutil
+timeout-decorator
+pyrender==0.1.45
+omegaconf
+hydra-core
+pandas==1.4.1
+opencv-python
+

external/yarr/setup.py

@@ -0,0 +1,37 @@
+import codecs
+import os
+
+import setuptools
+
+
+def read(rel_path):
+    here = os.path.abspath(os.path.dirname(__file__))
+    with codecs.open(os.path.join(here, rel_path), 'r') as fp:
+        return fp.read()
+
+
+def get_version(rel_path):
+    for line in read(rel_path).splitlines():
+        if line.startswith('__version__'):
+            delim = '"' if '"' in line else "'"
+            return line.split(delim)[1]
+    else:
+        raise RuntimeError("Unable to find version string.")
+
+
+def get_install_requires():
+    install_requires = []
+    with open('requirements.txt') as f:
+        for req in f:
+            install_requires.append(req.strip())
+    return install_requires
+
+
+setuptools.setup(
+    version=get_version("yarr/__init__.py"),
+    name='yarr',
+    author='Stephen James',
+    author_email='slj12@ic.ac.uk',
+    packages=setuptools.find_packages(),
+    install_requires=get_install_requires()
+)

external/yarr/yarr/__init__.py

@@ -0,0 +1 @@
+__version__ = '0.1'

external/yarr/yarr/agents/agent.py

@@ -0,0 +1,345 @@
+from abc import ABC, abstractmethod
+from typing import Any, List
+
+
+class Summary(object):
+    def __init__(self, name: str, value: Any):
+        self.name = name
+        self.value = value
+
+
+class ScalarSummary(Summary):
+    pass
+
+
+class HistogramSummary(Summary):
+    pass
+
+
+class ImageSummary(Summary):
+    pass
+
+
+class TextSummary(Summary):
+    pass
+
+
+class VideoSummary(Summary):
+    def __init__(self, name: str, value: Any, fps: int = 30):
+        super(VideoSummary, self).__init__(name, value)
+        self.fps = fps
+
+
+class ActResult(object):
+
+    def __init__(self, action: Any,
+                 observation_elements: dict = None,
+                 replay_elements: dict = None,
+                 info: dict = None):
+        self.action = action
+        self.observation_elements = observation_elements or {}
+        self.replay_elements = replay_elements or {}
+        self.info = info or {}
+
+
+class Agent(ABC):
+
+    @abstractmethod
+    def build(self, training: bool, device=None) -> None:
+        pass
+
+    @abstractmethod
+    def update(self, step: int, replay_sample: dict) -> dict:
+        pass
+
+    @abstractmethod
+    def act(self, step: int, observation: dict, deterministic: bool) -> ActResult:
+        # returns dict of values that get put in the replay.
+        # One of these must be 'action'.
+        pass
+
+    def reset(self) -> None:
+        pass
+
+    @abstractmethod
+    def update_summaries(self) -> List[Summary]:
+        pass
+
+    @abstractmethod
+    def act_summaries(self) -> List[Summary]:
+        pass
+
+    @abstractmethod
+    def load_weights(self, savedir: str) -> None:
+        pass
+
+    @abstractmethod
+    def save_weights(self, savedir: str) -> None:
+        pass
+
+
+class BimanualAgent(Agent):
+    """
+    
+    """
+
+    def __init__(self, right_agent: Agent, left_agent: Agent):
+        self.right_agent = right_agent
+        self.left_agent = left_agent
+        self._summaries = {}
+
+    def build(self, training: bool, device=None) -> None:
+        self.right_agent.build(training, device)
+        self.left_agent.build(training, device)
+
+    def update(self, step: int, replay_sample: dict) -> dict:
+        right_observation = {}
+        left_observation = {}
+
+        for k, v in replay_sample.items(): 
+            if "rgb" in k or "point_cloud" in k or "camera" in k:
+                right_observation[k] = v
+                left_observation[k] = v
+            elif "right_" in k :
+                right_observation[k[6:]] = v
+            elif "left_" in k:
+                left_observation[k[5:]] = v
+            else:
+                right_observation[k] = v
+                left_observation[k] = v
+
+        action = replay_sample["action"]
+        right_action, left_action = action.chunk(2, dim=2)
+        right_observation["action"] = right_action
+        left_observation["action"] = left_action
+
+        right_update_dict = self.right_agent.update(step, right_observation)
+        left_update_dict =  self.left_agent.update(step, left_observation)
+
+        total_losses = right_update_dict["total_losses"] + left_update_dict["total_losses"]
+        self._summaries.update({"total_losses": total_losses})
+        return self._summaries
+    
+
+    def act(self, step: int, observation: dict, deterministic: bool) -> ActResult:
+
+        observation_elements = {}
+        info = {}
+
+        right_observation = {}
+        left_observation = {}
+
+        for k, v in observation.items(): 
+            if "rgb" in k or "point_cloud" in k or "camera" in k:
+                right_observation[k] = v
+                left_observation[k] = v
+            elif "right_" in k :
+                right_observation[k[6:]] = v
+            elif "left_" in k:
+                left_observation[k[5:]] = v
+            else:
+                right_observation[k] = v
+                left_observation[k] = v
+
+        right_act_result = self.right_agent.act(step, right_observation, deterministic)
+        left_act_result = self.left_agent.act(step, left_observation, deterministic)
+
+        action = (*right_act_result.action, *left_act_result.action)
+
+        observation_elements.update(right_act_result.observation_elements)
+        observation_elements.update(left_act_result.observation_elements)
+
+        info.update(right_act_result.info)
+        info.update(left_act_result.info)
+
+        return ActResult(action, observation_elements=observation_elements, info=info)
+
+    def reset(self) -> None:
+        self.right_agent.reset()
+        self.left_agent.reset()
+
+    def update_summaries(self) -> List[Summary]:
+        summaries = []
+        for k, v in self._summaries.items():
+            summaries.append(ScalarSummary(f"{k}", v))
+
+        right_summaries = self.right_agent.update_summaries() 
+        left_summaries =  self.left_agent.update_summaries()
+        
+        for summary in right_summaries:
+            if not isinstance(summary, ImageSummary):
+                summary.name = f"agent_right/{summary.name}"
+
+        for summary in left_summaries:
+            if not isinstance(summary, ImageSummary):
+                summary.name = f"agent_left/{summary.name}"
+
+        return right_summaries + left_summaries + summaries
+
+
+    def act_summaries(self) -> List[Summary]:
+        right_summaries = self.right_agent.act_summaries() 
+        left_summaries =  self.left_agent.act_summaries()
+
+        for summary in right_summaries:
+            if not isinstance(summary, ImageSummary):
+                summary.name = f"agent_right/{summary.name}"
+
+        for summary in left_summaries:
+            if not isinstance(summary, ImageSummary):
+                summary.name = f"agent_left/{summary.name}"
+
+        return right_summaries + left_summaries
+
+    
+    def load_weights(self, savedir: str) -> None:
+        self.right_agent.load_weights(savedir)
+        self.left_agent.load_weights(savedir)
+
+    def save_weights(self, savedir: str) -> None:
+        self.right_agent.save_weights(savedir)
+        self.left_agent.save_weights(savedir)
+
+
+class LeaderFollowerAgent(Agent):
+
+    def __init__(self, leader_agent: Agent, follower_agent: Agent):
+        self.leader_agent = leader_agent
+        self.follower_agent = follower_agent
+        self._summaries = {}
+
+    def build(self, training: bool, device=None) -> None:
+        self.leader_agent.build(training, device)
+        self.follower_agent.build(training, device)
+
+    def update(self, step: int, replay_sample: dict) -> dict:
+
+        leader_observation = {}
+        follower_observation = {}
+
+
+        for k, v in replay_sample.items(): 
+            if "rgb" in k or "point_cloud" in k or "camera" in k:
+                leader_observation[k] = v
+                follower_observation[k] = v
+            elif "right_" in k :
+                leader_observation[k[6:]] = v
+            elif "left_" in k:
+                follower_observation[k[5:]] = v
+            else:
+                leader_observation[k] = v
+                follower_observation[k] = v
+
+        action = replay_sample["action"]
+        right_action, left_action = action.chunk(2, dim=2)
+        leader_observation["action"] = right_action
+        follower_observation["action"] = left_action
+
+        leader_update_dict = self.leader_agent.update(step, leader_observation)
+        import torch
+        follower_observation['low_dim_state'] = torch.cat([follower_observation['low_dim_state'],
+                                                           replay_sample["right_trans_action_indicies"], 
+                                                           replay_sample["right_rot_grip_action_indicies"], 
+                                                           replay_sample["right_ignore_collisions"]], dim=-1)
+
+        follower_update_dict = self.follower_agent.update(step, follower_observation)
+
+        total_losses = leader_update_dict["total_losses"] + follower_update_dict["total_losses"]
+        self._summaries.update({"total_losses": total_losses})
+        return self._summaries
+ 
+    def act(self, step: int, observation: dict, deterministic: bool) -> ActResult:
+
+        observation_elements = {}
+        info = {}
+
+        leader_observation = {}
+        follower_observation = {}
+
+        for k,v in observation.items():
+            if "right_" in k and not "rgb" in k and not "point_cloud" in k and not "camera" in k:
+                leader_observation[k[6:]] = v
+            elif "left_" in k and not "rgb" in k and not "point_cloud" in k and not "camera" in k:
+                follower_observation[k[5:]] = v
+            else:
+                leader_observation[k] = v
+                follower_observation[k] = v
+
+        right_act_result = self.leader_agent.act(step, leader_observation, deterministic)
+
+        right_observation_elements = right_act_result.observation_elements
+
+        import torch
+
+        device = follower_observation['low_dim_state'].device
+        if "trans_action_indicies" in right_observation_elements:
+            right_trans_action_indicies = torch.from_numpy(right_observation_elements["trans_action_indicies"]).unsqueeze(0).unsqueeze(0).to(device)
+            right_rot_grip_action_indicies = torch.from_numpy(right_observation_elements["rot_grip_action_indicies"]).unsqueeze(0).unsqueeze(0).to(device)
+            right_ignore_collisions = torch.from_numpy(right_act_result.action[-1:]).unsqueeze(0).unsqueeze(0).to(device)
+        else:
+            right_trans_action_indicies = torch.empty((1, 1, 3)).to(device)
+            right_rot_grip_action_indicies = torch.empty((1, 1, 4)).to(device)
+            right_ignore_collisions = torch.empty((1, 1, 1)).to(device)
+ 
+
+        follower_observation['low_dim_state'] = torch.cat([follower_observation['low_dim_state'],
+                                                           right_trans_action_indicies,
+                                                           right_rot_grip_action_indicies,
+                                                           right_ignore_collisions], dim=-1)  
+        
+        left_act_result = self.follower_agent.act(step, follower_observation, deterministic)
+
+        action = (*right_act_result.action, *left_act_result.action)
+        
+        observation_elements.update(right_act_result.observation_elements)
+        observation_elements.update(left_act_result.observation_elements)
+
+        info.update(right_act_result.info)
+        info.update(left_act_result.info)
+
+        return ActResult(action, observation_elements=observation_elements, info=info)
+    
+
+    def reset(self) -> None:
+        self.leader_agent.reset()
+        self.follower_agent.reset()
+
+    def update_summaries(self) -> List[Summary]:
+
+        summaries = []
+        for k, v in self._summaries.items():
+            summaries.append(ScalarSummary(f"{k}", v))
+
+        leader_summaries = self.leader_agent.update_summaries() 
+        follower_summaries =  self.follower_agent.update_summaries()
+
+        for summary in leader_summaries:
+            if not isinstance(summary, ImageSummary):
+                summary.name = f"agent_leader/{summary.name}"
+        for summary in follower_summaries:
+            if not isinstance(summary, ImageSummary):
+                summary.name = f"agent_follower/{summary.name}"
+
+        return leader_summaries + follower_summaries + summaries
+        
+
+    def act_summaries(self) -> List[Summary]:
+        leader_summaries = self.leader_agent.act_summaries() 
+        follower_summaries =  self.follower_agent.act_summaries()
+
+        for summary in leader_summaries:
+            if not isinstance(summary, ImageSummary):
+                summary.name = f"agent_leader/{summary.name}"
+        for summary in follower_summaries:
+            if not isinstance(summary, ImageSummary):
+                summary.name = f"agent_follower/{summary.name}"
+
+        return leader_summaries + follower_summaries
+
+    def load_weights(self, savedir: str) -> None:
+        self.leader_agent.load_weights(savedir)
+        self.follower_agent.load_weights(savedir)
+
+    def save_weights(self, savedir: str) -> None:
+        self.leader_agent.save_weights(savedir)
+        self.follower_agent.save_weights(savedir)

external/yarr/yarr/envs/env.py

@@ -0,0 +1,64 @@
+from abc import ABC, abstractmethod
+from typing import Any, List
+
+import numpy as np
+
+from yarr.utils.observation_type import ObservationElement
+from yarr.utils.transition import Transition
+
+
+class Env(ABC):
+
+    def __init__(self):
+        self._active_task_id = 0
+        self._eval_env = False
+
+    @property
+    def eval(self):
+        return self._eval_env
+
+    @eval.setter
+    def eval(self, eval):
+        self._eval_env = eval
+
+    @property
+    def active_task_id(self) -> int:
+        return self._active_task_id
+
+    @abstractmethod
+    def launch(self) -> None:
+        pass
+
+    def shutdown(self) -> None:
+        pass
+
+    @abstractmethod
+    def reset(self) -> dict:
+        pass
+
+    @abstractmethod
+    def step(self, action: np.ndarray) -> Transition:
+        pass
+
+    @property
+    @abstractmethod
+    def observation_elements(self) -> List[ObservationElement]:
+        pass
+
+    @property
+    @abstractmethod
+    def action_shape(self) -> tuple:
+        pass
+
+    @property
+    @abstractmethod
+    def env(self) -> Any:
+        pass
+
+
+class MultiTaskEnv(Env):
+
+    @property
+    @abstractmethod
+    def num_tasks(self) -> int:
+        pass

external/yarr/yarr/envs/rlbench_env.py

@@ -0,0 +1,332 @@
+import re
+from typing import Type, List
+
+import numpy as np
+try:
+    from rlbench import ObservationConfig, Environment, CameraConfig
+except (ModuleNotFoundError, ImportError) as e:
+    print("You need to install RLBench: 'https://github.com/stepjam/RLBench'")
+    raise e
+from rlbench.action_modes.action_mode import ActionMode
+from rlbench.backend.observation import BimanualObservation, Observation
+from rlbench.backend.task import Task
+from rlbench.backend.task import BimanualTask
+
+from helpers.clip.core.clip import tokenize
+
+from yarr.envs.env import Env, MultiTaskEnv
+from yarr.utils.observation_type import ObservationElement
+from yarr.utils.transition import Transition
+from yarr.utils.process_str import change_case
+
+import logging
+
+
+ROBOT_STATE_KEYS = ['joint_velocities', 'joint_positions', 'joint_forces',
+                        'gripper_open', 'gripper_pose',
+                        'gripper_joint_positions', 'gripper_touch_forces',
+                        'task_low_dim_state', 'misc', 'left', 'right']
+
+
+# ..todo:: possibly duplicated code.
+def _extract_obs_bimanual(obs: BimanualObservation, channels_last: bool, observation_config: ObservationConfig):
+    obs_dict = vars(obs)
+    obs_dict = {k: v for k, v in obs_dict.items() if v is not None}
+
+    right_robot_state = obs.get_low_dim_data(obs.right)
+    left_robot_state = obs.get_low_dim_data(obs.left)
+
+    obs_dict = {k: v for k, v in obs_dict.items()
+            if k not in ROBOT_STATE_KEYS}
+
+    if not channels_last:
+        # Swap channels from last dim to 1st dim
+        obs_dict = {k: np.transpose(v, [2, 0, 1]) if v.ndim == 3 else np.expand_dims(v, 0)
+                    for k, v in obs.perception_data.items() if v is not None}
+    else:
+        # Add extra dim to depth data
+        obs_dict = {k: v if v.ndim == 3 else np.expand_dims(v, -1)
+                    for k, v in obs.perception_data.items() if v is not None}
+        
+    if observation_config.robot_name == 'right':
+        obs_dict['low_dim_state'] = right_robot_state.astype(np.float32)
+        obs_dict['ignore_collisions'] = np.array([obs.right.ignore_collisions], dtype=np.float32)
+    elif observation_config.robot_name == 'left':
+        obs_dict['low_dim_state'] = left_robot_state.astype(np.float32)
+        obs_dict['ignore_collisions'] = np.array([obs.left.ignore_collisions], dtype=np.float32)
+    else:
+        obs_dict['right_low_dim_state'] = right_robot_state.astype(np.float32)
+        obs_dict['left_low_dim_state'] = left_robot_state.astype(np.float32)
+        obs_dict['right_ignore_collisions'] = np.array([obs.right.ignore_collisions], dtype=np.float32)
+        obs_dict['left_ignore_collisions'] = np.array([obs.left.ignore_collisions], dtype=np.float32)
+
+    for (k, v) in [(k, v) for k, v in obs_dict.items() if 'point_cloud' in k]:
+        # ..TODO:: 
+        obs_dict[k] = v.astype(np.float16)
+
+    for camera_name, config in observation_config.camera_configs.items():
+        if config.point_cloud:
+            obs_dict[f'{camera_name}_camera_extrinsics'] = obs.misc[f'{camera_name}_camera_extrinsics']
+            obs_dict[f'{camera_name}_camera_intrinsics'] = obs.misc[f'{camera_name}_camera_intrinsics']
+    return obs_dict
+    
+
+def _extract_obs_unimanual(obs: Observation, channels_last: bool, observation_config):
+    obs_dict = vars(obs)
+    obs_dict = {k: v for k, v in obs_dict.items() if v is not None}
+    robot_state = obs.get_low_dim_data()
+    # Remove all of the individual state elements
+    obs_dict = {k: v for k, v in obs_dict.items()
+                if k not in ROBOT_STATE_KEYS}
+    if not channels_last:
+        # Swap channels from last dim to 1st dim
+        obs_dict = {k: np.transpose(
+            v, [2, 0, 1]) if v.ndim == 3 else np.expand_dims(v, 0)
+                    for k, v in obs_dict.items()}
+    else:
+        # Add extra dim to depth data
+        obs_dict = {k: v if v.ndim == 3 else np.expand_dims(v, -1)
+                    for k, v in obs_dict.items()}
+    obs_dict['low_dim_state'] = np.array(robot_state, dtype=np.float32)
+    obs_dict['ignore_collisions'] = np.array([obs.ignore_collisions], dtype=np.float32)
+    for (k, v) in [(k, v) for k, v in obs_dict.items() if 'point_cloud' in k]:
+        obs_dict[k] = v.astype(np.float32)
+
+    for config, name in [
+        (observation_config.left_shoulder_camera, 'left_shoulder'),
+        (observation_config.right_shoulder_camera, 'right_shoulder'),
+        (observation_config.front_camera, 'front'),
+        (observation_config.wrist_camera, 'wrist'),
+        (observation_config.overhead_camera, 'overhead')]:
+        if config.point_cloud:
+            obs_dict['%s_camera_extrinsics' % name] = obs.misc['%s_camera_extrinsics' % name]
+            obs_dict['%s_camera_intrinsics' % name] = obs.misc['%s_camera_intrinsics' % name]
+    return obs_dict
+
+
+def _get_cam_observation_elements(camera: CameraConfig, prefix: str, channels_last):
+    elements = []
+    img_s = list(camera.image_size)
+    shape = img_s + [3] if channels_last else [3] + img_s
+    if camera.rgb:
+        elements.append(
+            ObservationElement('%s_rgb' % prefix, shape, np.uint8))
+    if camera.point_cloud:
+        elements.append(
+            ObservationElement('%s_point_cloud' % prefix, shape, np.float32))
+        elements.append(
+            ObservationElement('%s_camera_extrinsics' % prefix, (4, 4),
+                               np.float32))
+        elements.append(
+            ObservationElement('%s_camera_intrinsics' % prefix, (3, 3),
+                               np.float32))
+    if camera.depth:
+        shape = img_s + [1] if channels_last else [1] + img_s
+        elements.append(
+            ObservationElement('%s_depth' % prefix, shape, np.float32))
+    if camera.mask:
+        raise NotImplementedError()
+
+    return elements
+
+
+def _observation_elements(observation_config, channels_last) -> List[ObservationElement]:
+    elements = []
+    robot_state_len = 0
+    if observation_config.joint_velocities:
+        robot_state_len += 7
+    if observation_config.joint_positions:
+        robot_state_len += 7
+    if observation_config.joint_forces:
+        robot_state_len += 7
+    if observation_config.gripper_open:
+        robot_state_len += 1
+    if observation_config.gripper_pose:
+        robot_state_len += 7
+    if observation_config.gripper_joint_positions:
+        robot_state_len += 2
+    if observation_config.gripper_touch_forces:
+        robot_state_len += 2
+    if observation_config.task_low_dim_state:
+        raise NotImplementedError()
+    if robot_state_len > 0:
+        if observation_config.robot_name == 'bimanual':
+            elements.append(ObservationElement(
+                'right_low_dim_state', (robot_state_len,), np.float32))
+            elements.append(ObservationElement(
+                'left_low_dim_state', (robot_state_len,), np.float32))
+        elif observation_config.robot_name in ['unimanual', 'left', 'right']:            
+            elements.append(ObservationElement('low_dim_state', (robot_state_len,), np.float32))
+    elements.extend(_get_cam_observation_elements(
+        observation_config.left_shoulder_camera, 'left_shoulder', channels_last))
+    elements.extend(_get_cam_observation_elements(
+        observation_config.right_shoulder_camera, 'right_shoulder', channels_last))
+    elements.extend(_get_cam_observation_elements(
+        observation_config.front_camera, 'front', channels_last))
+    elements.extend(_get_cam_observation_elements(
+        observation_config.wrist_camera, 'wrist', channels_last))
+    return elements
+
+
+class RLBenchEnv(Env):
+
+    def __init__(self, task_class: Type[Task],
+                 observation_config: ObservationConfig,
+                 action_mode: ActionMode,
+                 dataset_root: str = '',
+                 channels_last=False,
+                 headless=True,
+                 include_lang_goal_in_obs=False):
+        super(RLBenchEnv, self).__init__()
+        self._task_class = task_class
+        self._observation_config = observation_config
+        self._channels_last = channels_last
+        self._include_lang_goal_in_obs = include_lang_goal_in_obs
+        if issubclass(task_class, BimanualTask):
+            robot_setup = "dual_panda"
+        else:
+            robot_setup = "panda"
+        self._rlbench_env = Environment(
+            action_mode=action_mode, obs_config=observation_config,
+            dataset_root=dataset_root, headless=headless, robot_setup=robot_setup)
+        self._task = None
+        self._lang_goal = 'unknown goal'
+
+
+    def extract_obs(self, obs: Observation):
+        if isinstance(obs, BimanualObservation):
+            extracted_obs = _extract_obs_bimanual(obs, self._channels_last, self._observation_config)
+        else:
+            extracted_obs = _extract_obs_unimanual(obs, self._channels_last, self._observation_config)
+        if self._include_lang_goal_in_obs:
+            extracted_obs['lang_goal_tokens'] = tokenize([self._lang_goal])[0].numpy()
+        return extracted_obs
+
+    def launch(self):
+        self._rlbench_env.launch()
+        self._task = self._rlbench_env.get_task(self._task_class)
+
+    def shutdown(self):
+        self._rlbench_env.shutdown()
+
+    def reset(self) -> dict:
+        descriptions, obs = self._task.reset()
+        self._lang_goal = descriptions[0] # first description variant
+        extracted_obs = self.extract_obs(obs)
+        return extracted_obs
+
+    def step(self, action: np.ndarray) -> Transition:
+        obs, reward, terminal = self._task.step(action)
+        obs = self.extract_obs(obs)
+        return Transition(obs, reward, terminal)
+
+    @property
+    def observation_elements(self) -> List[ObservationElement]:
+        return _observation_elements(self._observation_config, self._channels_last)
+
+    @property
+    def action_shape(self):
+        return (self._rlbench_env.action_size, )
+
+    @property
+    def env(self) -> Environment:
+        return self._rlbench_env
+
+
+class MultiTaskRLBenchEnv(MultiTaskEnv):
+
+    def __init__(self,
+                 task_classes: List[Type[Task]],
+                 observation_config: ObservationConfig,
+                 action_mode: ActionMode,
+                 dataset_root: str = '',
+                 channels_last=False,
+                 headless=True,
+                 swap_task_every: int = 1,
+                 include_lang_goal_in_obs=False):
+        super(MultiTaskRLBenchEnv, self).__init__()
+        self._task_classes = task_classes
+        self._observation_config = observation_config
+        self._channels_last = channels_last
+        self._include_lang_goal_in_obs = include_lang_goal_in_obs
+        if issubclass(task_classes[0], BimanualTask):
+            robot_setup = "dual_panda"
+        else:
+            robot_setup = "panda"
+        self._rlbench_env = Environment(
+            action_mode=action_mode, obs_config=observation_config,
+            dataset_root=dataset_root, headless=headless, robot_setup=robot_setup)
+        self._task = None
+        self._task_name = ''
+        self._lang_goal = 'unknown goal'
+        self._swap_task_every = swap_task_every
+        self._rlbench_env
+        self._episodes_this_task = 0
+        self._active_task_id = -1
+
+        self._task_name_to_idx = {change_case(tc.__name__):i for i, tc in enumerate(self._task_classes)}
+
+    def _set_new_task(self, shuffle=False):
+        if shuffle:
+            self._active_task_id = np.random.randint(0, len(self._task_classes))
+        else:
+            self._active_task_id = (self._active_task_id + 1) % len(self._task_classes)
+        task = self._task_classes[self._active_task_id]
+        self._task = self._rlbench_env.get_task(task)
+
+    def set_task(self, task_name: str):
+        self._active_task_id = self._task_name_to_idx[task_name]
+        task = self._task_classes[self._active_task_id]
+        self._task = self._rlbench_env.get_task(task)
+
+        descriptions, _ = self._task.reset()
+        self._lang_goal = descriptions[0] # first description variant
+
+    def extract_obs(self, obs: Observation):
+        if obs.is_bimanual:
+            extracted_obs = _extract_obs_bimanual(obs, self._channels_last, self._observation_config)
+        else:
+            extracted_obs = _extract_obs_unimanual(obs, self._channels_last, self._observation_config)
+        if self._include_lang_goal_in_obs:
+            extracted_obs['lang_goal_tokens'] = tokenize([self._lang_goal])[0].numpy()
+        return extracted_obs
+
+    def launch(self):
+        self._rlbench_env.launch()
+        self._set_new_task()
+
+    def shutdown(self):
+        self._rlbench_env.shutdown()
+
+    def reset(self) -> dict:
+        if self._episodes_this_task == self._swap_task_every:
+            self._set_new_task()
+            self._episodes_this_task = 0
+        self._episodes_this_task += 1
+
+        descriptions, obs = self._task.reset()
+        self._lang_goal = descriptions[0] # first description variant
+        extracted_obs = self.extract_obs(obs)
+
+        return extracted_obs
+
+    def step(self, action: np.ndarray) -> Transition:
+        obs, reward, terminal = self._task.step(action)
+        obs = self.extract_obs(obs)
+        return Transition(obs, reward, terminal)
+
+    @property
+    def observation_elements(self) -> List[ObservationElement]:
+        return _observation_elements(self._observation_config, self._channels_last)
+
+    @property
+    def action_shape(self):
+        return (self._rlbench_env.action_size, )
+
+    @property
+    def env(self) -> Environment:
+        return self._rlbench_env
+
+    @property
+    def num_tasks(self) -> int:
+        return len(self._task_classes)

external/yarr/yarr/replay_buffer/prioritized_replay_buffer.py

@@ -0,0 +1,217 @@
+"""An implementation of Prioritized Experience Replay (PER).
+
+This implementation is based on the paper "Prioritized Experience Replay"
+by Tom Schaul et al. (2015).
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from .uniform_replay_buffer import *
+from .sum_tree import *
+import numpy as np
+
+
+PRIORITY = 'priority'
+
+
+class PrioritizedReplayBuffer(UniformReplayBuffer):
+    """An out-of-graph Replay Buffer for Prioritized Experience Replay.
+
+    See uniform_replay_buffer.py for details.
+    """
+
+    def __init__(self, *args, **kwargs):
+        """Initializes OutOfGraphPrioritizedReplayBuffer."""
+        super(PrioritizedReplayBuffer, self).__init__(*args, **kwargs)
+        self._sum_tree = SumTree(self._replay_capacity)
+
+    def get_storage_signature(self) -> Tuple[List[ReplayElement],
+                                             List[ReplayElement]]:
+        """Returns a default list of elements to be stored in this replay memory.
+
+        Note - Derived classes may return a different signature.
+
+        Returns:
+          dict of ReplayElements defining the type of the contents stored.
+        """
+        storage_elements, obs_elements = super(
+            PrioritizedReplayBuffer, self).get_storage_signature()
+        storage_elements.append(ReplayElement(PRIORITY, (), np.float32),)
+
+        return storage_elements, obs_elements
+
+    def add(self, action, reward, terminal, timeout, priority=None, **kwargs):
+        kwargs['priority'] = priority
+        super(PrioritizedReplayBuffer, self).add(
+            action, reward, terminal, timeout, **kwargs)
+
+    def _add(self, kwargs: dict):
+        """Internal add method to add to the storage arrays.
+
+        Args:
+          kwargs: All the elements in a transition.
+        """
+        with self._lock:
+            cursor = self.cursor()
+            priority = kwargs[PRIORITY]
+            if priority is None:
+                priority = self._sum_tree.max_recorded_priority
+
+            if self._disk_saving:
+                term = self._store[TERMINAL]
+                term[cursor] = kwargs[TERMINAL]
+                self._store[TERMINAL] = term
+
+                with open(join(self._save_dir, '%d.replay' % cursor), 'wb') as f:
+                    pickle.dump(kwargs, f)
+                # If first add, then pad for correct wrapping
+                if self._add_count.value == 0:
+                    self._add_initial_to_disk(kwargs)
+            else:
+                for name, data in kwargs.items():
+                    item = self._store[name]
+                    item[cursor] = data
+                    self._store[name] = item
+
+
+            self._sum_tree.set(self.cursor(), priority)
+            self._add_count.value += 1
+            self.invalid_range = invalid_range(
+                self.cursor(), self._replay_capacity, self._timesteps,
+                self._update_horizon)
+
+    def add_final(self, **kwargs):
+        """Adds a transition to the replay memory.
+        Args:
+          **kwargs: The remaining args
+        """
+        # if self.is_empty() or self._store['terminal'][self.cursor() - 1] != 1:
+        #     raise ValueError('The previous transition was not terminal.')
+        self._check_add_types(kwargs, self._obs_signature)
+        transition = self._final_transition(kwargs)
+        for element_type in self._storage_signature:
+            # 0 priority for final observation.
+            if element_type.name == PRIORITY:
+                transition[element_type.name] = 0.0
+        self._add(transition)
+
+    def sample_index_batch(self, batch_size):
+        """Returns a batch of valid indices sampled as in Schaul et al. (2015).
+
+        Args:
+          batch_size: int, number of indices returned.
+
+        Returns:
+          list of ints, a batch of valid indices sampled uniformly.
+
+        Raises:
+          Exception: If the batch was not constructed after maximum number of tries.
+        """
+        # Sample stratified indices. Some of them might be invalid.
+        indices = self._sum_tree.stratified_sample(batch_size)
+        allowed_attempts = self._max_sample_attempts
+        for i in range(len(indices)):
+            if not self.is_valid_transition(indices[i]):
+                if allowed_attempts == 0:
+                    raise RuntimeError(
+                        'Max sample attempts: Tried {} times but only sampled {}'
+                        ' valid indices. Batch size is {}'.
+                            format(self._max_sample_attempts, i, batch_size))
+                index = indices[i]
+                while not self.is_valid_transition(
+                        index) and allowed_attempts > 0:
+                    # If index i is not valid keep sampling others. Note that this
+                    # is not stratified.
+                    index = self._sum_tree.sample()
+                    allowed_attempts -= 1
+                indices[i] = index
+        return indices
+
+    def sample_transition_batch(self, batch_size=None, indices=None,
+                                pack_in_dict=True):
+        """Returns a batch of transitions with extra storage and the priorities.
+
+        The extra storage are defined through the extra_storage_types constructor
+        argument.
+
+        When the transition is terminal next_state_batch has undefined contents.
+
+        Args:
+          batch_size: int, number of transitions returned. If None, the default
+            batch_size will be used.
+          indices: None or list of ints, the indices of every transition in the
+            batch. If None, sample the indices uniformly.
+
+        Returns:
+          transition_batch: tuple of np.arrays with the shape and type as in
+            get_transition_elements().
+        """
+        transition = super(
+            PrioritizedReplayBuffer, self).sample_transition_batch(
+            batch_size, indices, pack_in_dict=False)
+
+        transition_elements = self.get_transition_elements(batch_size)
+        transition_names = [e.name for e in transition_elements]
+        probabilities_index = transition_names.index('sampling_probabilities')
+        indices_index = transition_names.index('indices')
+        indices = transition[indices_index]
+        # The parent returned an empty array for the probabilities. Fill it with the
+        # contents of the sum tree.
+        transition[probabilities_index][:] = self.get_priority(indices)
+        batch_arrays = transition
+        if pack_in_dict:
+            batch_arrays = self.unpack_transition(transition,
+                                                  transition_elements)
+        return batch_arrays
+
+    def set_priority(self, indices, priorities):
+        """Sets the priority of the given elements according to Schaul et al.
+
+        Args:
+          indices: np.array with dtype int32, of indices in range
+            [0, replay_capacity).
+          priorities: float, the corresponding priorities.
+        """
+        assert indices.dtype == np.int32, ('Indices must be integers, '
+                                           'given: {}'.format(indices.dtype))
+        for index, priority in zip(indices, priorities):
+            self._sum_tree.set(index, priority)
+
+    def get_priority(self, indices):
+        """Fetches the priorities correspond to a batch of memory indices.
+
+        For any memory location not yet used, the corresponding priority is 0.
+
+        Args:
+          indices: np.array with dtype int32, of indices in range
+            [0, replay_capacity).
+
+        Returns:
+          priorities: float, the corresponding priorities.
+        """
+        assert indices.shape, 'Indices must be an array.'
+        assert indices.dtype == np.int32, ('Indices must be int32s, '
+                                           'given: {}'.format(indices.dtype))
+        batch_size = len(indices)
+        priority_batch = np.empty((batch_size), dtype=np.float32)
+        for i, memory_index in enumerate(indices):
+            priority_batch[i] = self._sum_tree.get(memory_index)
+        return priority_batch
+
+    def get_transition_elements(self, batch_size=None):
+        """Returns a 'type signature' for sample_transition_batch.
+
+        Args:
+          batch_size: int, number of transitions returned. If None, the default
+            batch_size will be used.
+        Returns:
+          signature: A namedtuple describing the method's return type signature.
+        """
+        parent_transition_type = (
+            super(PrioritizedReplayBuffer,
+                  self).get_transition_elements(batch_size))
+        probablilities_type = [
+            ReplayElement('sampling_probabilities', (batch_size,), np.float32)
+        ]
+        return parent_transition_type + probablilities_type

external/yarr/yarr/replay_buffer/replay_buffer.py

@@ -0,0 +1,71 @@
+from abc import ABC
+from typing import Tuple, List
+
+class ReplayElement(object):
+    def __init__(self, name, shape, type, is_observation=False):
+        self.name = name
+        self.shape = shape
+        self.type = type
+        self.is_observation = is_observation
+
+
+class ReplayBuffer(ABC):
+
+    def replay_capacity(self):
+        pass
+
+    def batch_size(self):
+        pass
+
+    def get_storage_signature(self) -> Tuple[List[ReplayElement],
+                                             List[ReplayElement]]:
+        pass
+
+    def add(self, action, reward, terminal, timeout, **kwargs):
+        pass
+
+    def add_final(self, **kwargs):
+        pass
+
+    def is_empty(self):
+        pass
+
+    def is_full(self):
+        pass
+
+    def cursor(self):
+        pass
+
+    def set_cursor(self):
+        pass
+
+    def get_range(self, array, start_index, end_index):
+        pass
+
+    def get_range_stack(self, array, start_index, end_index, terminals=None):
+        pass
+
+    def get_terminal_stack(self, index):
+        pass
+
+    def is_valid_transition(self, index):
+        pass
+
+    def sample_index_batch(self, batch_size):
+        pass
+
+    def unpack_transition(self, transition_tensors, transition_type):
+        pass
+
+    def sample_transition_batch(self, batch_size=None, indices=None,
+                                pack_in_dict=True):
+        pass
+
+    def get_transition_elements(self, batch_size=None):
+        pass
+
+    def shutdown(self):
+        pass
+
+    def using_disk(self):
+        pass

external/yarr/yarr/replay_buffer/sum_tree.py

@@ -0,0 +1,201 @@
+"""A sum tree data structure.
+
+Used for prioritized experience replay. See prioritized_replay_buffer.py
+and Schaul et al. (2015).
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import random
+
+import numpy as np
+
+
+class SumTree(object):
+    """A sum tree data structure for storing replay priorities.
+
+    A sum tree is a complete binary tree whose leaves contain values called
+    priorities. Internal nodes maintain the sum of the priorities of all leaf
+    nodes in their subtree.
+
+    For capacity = 4, the tree may look like this:
+
+                 +---+
+                 |2.5|
+                 +-+-+
+                   |
+           +-------+--------+
+           |                |
+         +-+-+            +-+-+
+         |1.5|            |1.0|
+         +-+-+            +-+-+
+           |                |
+      +----+----+      +----+----+
+      |         |      |         |
+    +-+-+     +-+-+  +-+-+     +-+-+
+    |0.5|     |1.0|  |0.5|     |0.5|
+    +---+     +---+  +---+     +---+
+
+    This is stored in a list of numpy arrays:
+    self.nodes = [ [2.5], [1.5, 1], [0.5, 1, 0.5, 0.5] ]
+
+    For conciseness, we allocate arrays as powers of two, and pad the excess
+    elements with zero values.
+
+    This is similar to the usual array-based representation of a complete binary
+    tree, but is a little more user-friendly.
+    """
+
+    def __init__(self, capacity, nodes=None):
+        """Creates the sum tree data structure for the given replay capacity.
+
+        Args:
+          capacity: int, the maximum number of elements that can be stored in this
+            data structure.
+          nodes: storage list for storing nodes
+
+        Raises:
+          ValueError: If requested capacity is not positive.
+        """
+        assert isinstance(capacity, int)
+        if capacity <= 0:
+            raise ValueError('Sum tree capacity should be positive. Got: {}'.
+                             format(capacity))
+
+        self.nodes = [] if nodes is None else nodes
+        tree_depth = int(math.ceil(np.log2(capacity)))
+        level_size = 1
+        for _ in range(tree_depth + 1):
+            nodes_at_this_depth = np.zeros(level_size)
+            self.nodes.append(nodes_at_this_depth)
+
+            level_size *= 2
+
+        self.max_recorded_priority = 1.0
+
+    def _total_priority(self):
+        """Returns the sum of all priorities stored in this sum tree.
+
+        Returns:
+          float, sum of priorities stored in this sum tree.
+        """
+        return self.nodes[0][0]
+
+    def sample(self, query_value=None):
+        """Samples an element from the sum tree.
+
+        Each element has probability p_i / sum_j p_j of being picked, where p_i is
+        the (positive) value associated with node i (possibly unnormalized).
+
+        Args:
+          query_value: float in [0, 1], used as the random value to select a
+          sample. If None, will select one randomly in [0, 1).
+
+        Returns:
+          int, a random element from the sum tree.
+
+        Raises:
+          Exception: If the sum tree is empty (i.e. its node values sum to 0), or if
+            the supplied query_value is larger than the total sum.
+        """
+        if self._total_priority() == 0.0:
+            raise Exception('Cannot sample from an empty sum tree.')
+
+        if query_value and (query_value < 0. or query_value > 1.):
+            raise ValueError('query_value must be in [0, 1].')
+
+        # Sample a value in range [0, R), where R is the value stored at the root.
+        query_value = random.random() if query_value is None else query_value
+        query_value *= self._total_priority()
+
+        # Now traverse the sum tree.
+        node_index = 0
+        for nodes_at_this_depth in self.nodes[1:]:
+            # Compute children of previous depth's node.
+            left_child = node_index * 2
+
+            left_sum = nodes_at_this_depth[left_child]
+            # Each subtree describes a range [0, a), where a is its value.
+            if query_value < left_sum:  # Recurse into left subtree.
+                node_index = left_child
+            else:  # Recurse into right subtree.
+                node_index = left_child + 1
+                # Adjust query to be relative to right subtree.
+                query_value -= left_sum
+
+        return node_index
+
+    def stratified_sample(self, batch_size):
+        """Performs stratified sampling using the sum tree.
+
+        Let R be the value at the root (total value of sum tree). This method will
+        divide [0, R) into batch_size segments, pick a random number from each of
+        those segments, and use that random number to sample from the sum_tree. This
+        is as specified in Schaul et al. (2015).
+
+        Args:
+          batch_size: int, the number of strata to use.
+        Returns:
+          list of batch_size elements sampled from the sum tree.
+
+        Raises:
+          Exception: If the sum tree is empty (i.e. its node values sum to 0).
+        """
+        if self._total_priority() == 0.0:
+            raise Exception('Cannot sample from an empty sum tree.')
+
+        bounds = np.linspace(0., 1., batch_size + 1)
+        assert len(bounds) == batch_size + 1
+        segments = [(bounds[i], bounds[i + 1]) for i in range(batch_size)]
+        # TODO removed for now
+        # query_values = [random.uniform(x[0], x[1]) for x in segments]
+        query_values = [random.uniform(0, 1) for x in segments]
+        return [self.sample(query_value=x) for x in query_values]
+
+    def get(self, node_index):
+        """Returns the value of the leaf node corresponding to the index.
+
+        Args:
+          node_index: The index of the leaf node.
+        Returns:
+          The value of the leaf node.
+        """
+        return self.nodes[-1][node_index]
+
+    def set(self, node_index, value):
+        """Sets the value of a leaf node and updates internal nodes accordingly.
+
+        This operation takes O(log(capacity)).
+        Args:
+          node_index: int, the index of the leaf node to be updated.
+          value: float, the value which we assign to the node. This value must be
+            nonnegative. Setting value = 0 will cause the element to never be
+            sampled.
+
+        Raises:
+          ValueError: If the given value is negative.
+        """
+        if value < 0.0:
+            raise ValueError('Sum tree values should be nonnegative. Got {}'.
+                             format(value))
+        self.max_recorded_priority = max(value, self.max_recorded_priority)
+
+        delta_value = value - self.nodes[-1][node_index]
+
+        # # Now traverse back the tree, adjusting all sums along the way.
+        # for nodes_at_this_depth in reversed(self.nodes):
+        #     # Note: Adding a delta leads to some tolerable numerical inaccuracies.
+        #     nodes_at_this_depth[node_index] += delta_value
+        #     self.nodes[]
+        #     node_index //= 2
+
+        for nodes_at_this_depth_idx in reversed(range(len(self.nodes))):
+            nodes_at_this_depth = self.nodes[nodes_at_this_depth_idx]
+            nodes_at_this_depth[node_index] += delta_value
+            self.nodes[nodes_at_this_depth_idx] = nodes_at_this_depth
+            node_index //= 2
+
+        assert node_index == 0, ('Sum tree traversal failed, final node index '
+                                 'is not 0.')

external/yarr/yarr/replay_buffer/task_uniform_replay_buffer.py

@@ -0,0 +1,182 @@
+import numpy as np
+import os
+from os.path import join
+import pickle
+import math
+from yarr.replay_buffer.uniform_replay_buffer import UniformReplayBuffer
+from yarr.replay_buffer.uniform_replay_buffer import invalid_range
+
+from yarr.replay_buffer.replay_buffer import ReplayBuffer, ReplayElement
+from yarr.utils.observation_type import ObservationElement
+
+ACTION = 'action'
+REWARD = 'reward'
+TERMINAL = 'terminal'
+TIMEOUT = 'timeout'
+INDICES = 'indices'
+TASK = 'task'
+
+
+class TaskUniformReplayBuffer(UniformReplayBuffer):
+    """
+    A uniform with uniform task sampling for each batch
+    """
+
+    def __init__(self, *args, **kwargs):
+        """Initializes OutOfGraphPrioritizedReplayBuffer."""
+        super(TaskUniformReplayBuffer, self).__init__(*args, **kwargs)
+        self._task_idxs = dict()
+
+    def _add(self, kwargs: dict):
+        """Internal add method to add to the storage arrays.
+
+        Args:
+          kwargs: All the elements in a transition.
+        """
+        with self._lock:
+            cursor = self.cursor()
+
+            if self._disk_saving:
+                term = self._store[TERMINAL]
+                term[cursor] = kwargs[TERMINAL]
+                self._store[TERMINAL] = term
+
+
+                ## reduce size
+                for k, v in kwargs.items():
+                    try:
+                        if 'float' in v.dtype.name and v.size > 100:
+                            v = v.astype(np.float16)
+                            kwargs[k] = v
+                    except:
+                        pass
+
+
+                with open(join(self._save_dir, '%d.replay' % cursor), 'wb') as f:
+                    pickle.dump(kwargs, f)
+                # If first add, then pad for correct wrapping
+                if self._add_count.value == 0:
+                    self._add_initial_to_disk(kwargs)
+            else:
+                for name, data in kwargs.items():
+                    item = self._store[name]
+                    item[cursor] = data
+                    self._store[name] = item
+            with self._add_count.get_lock():
+                task = kwargs[TASK]
+                if task not in self._task_idxs:
+                    self._task_idxs[task] = [cursor]
+                else:
+                    self._task_idxs[task] =  self._task_idxs[task] + [cursor]
+                self._add_count.value += 1
+
+            self.invalid_range = invalid_range(
+                self.cursor(), self._replay_capacity, self._timesteps,
+                self._update_horizon)
+
+    def sample_index_batch(self,
+                           batch_size):
+        """Returns a batch of valid indices sampled uniformly.
+
+        Args:
+          batch_size: int, number of indices returned.
+
+        Returns:
+          list of ints, a batch of valid indices sampled uniformly across tasks.
+
+        Raises:
+          RuntimeError: If the batch was not constructed after maximum number of
+            tries.
+        """
+        if self.is_full():
+            min_id = (self.cursor() - self._replay_capacity +
+                      self._timesteps - 1)
+            max_id = self.cursor() - self._update_horizon
+        else:
+            min_id = 0
+            max_id = self.cursor() - self._update_horizon
+            if max_id <= min_id:
+                raise RuntimeError(
+                    'Cannot sample a batch with fewer than stack size '
+                    '({}) + update_horizon ({}) transitions.'.
+                    format(self._timesteps, self._update_horizon))
+
+        tasks = list(self._task_idxs.keys())
+        attempt_count = 0
+        found_indicies = False
+
+        # uniform distribution of tasks
+        while not found_indicies and attempt_count < 1000:
+            # sample random tasks of batch_size length
+            sampled_tasks = list(np.random.choice(tasks, batch_size, replace=(batch_size > len(tasks))))
+            potential_indices = []
+            for task in sampled_tasks:
+                # DDP setting where each GPU only sees a fraction of the data
+                # reference: https://github.com/pytorch/pytorch/blob/master/torch/utils/data/distributed.py
+                task_data_size = len(self._task_idxs[task])
+                num_samples = math.ceil(task_data_size / self._num_replicas)
+                total_size = num_samples * self._num_replicas
+                task_indices = self._task_idxs[task][self._rank:total_size:self._num_replicas]
+
+                sampled_task_idx = np.random.choice(task_indices, 1)[0]
+                per_task_attempt_count = 0
+
+                # Argh.. this is slow
+                while not self.is_valid_transition(sampled_task_idx) and \
+                    per_task_attempt_count < self._max_sample_attempts:
+                    sampled_task_idx = np.random.choice(task_indices, 1)[0]
+                    per_task_attempt_count += 1
+
+                if not self.is_valid_transition(sampled_task_idx):
+                    attempt_count += 1
+                    continue
+                else:
+                    potential_indices.append(sampled_task_idx)
+            found_indicies = len(potential_indices) == batch_size
+        indices = potential_indices
+
+        if len(indices) != batch_size:
+            raise RuntimeError(
+                'Max sample attempts: Tried {} times but only sampled {}'
+                ' valid indices. Batch size is {}'.
+                    format(self._max_sample_attempts, len(indices), batch_size))
+
+        return indices
+
+    def get_transition_elements(self, batch_size=None):
+        """Returns a 'type signature' for sample_transition_batch.
+
+        Args:
+          batch_size: int, number of transitions returned. If None, the default
+            batch_size will be used.
+        Returns:
+          signature: A namedtuple describing the method's return type signature.
+        """
+        batch_size = self._batch_size if batch_size is None else batch_size
+
+        transition_elements = [
+            ReplayElement(ACTION, (batch_size, self._timesteps) + self._action_shape,
+                          self._action_dtype),
+            ReplayElement(REWARD, (batch_size, self._timesteps) + self._reward_shape,
+                          self._reward_dtype),
+            ReplayElement(TERMINAL, (batch_size, self._timesteps), np.int8),
+            ReplayElement(TIMEOUT, (batch_size, self._timesteps), bool),
+            ReplayElement(INDICES, (batch_size, self._timesteps), np.int32),
+        ]
+
+        for element in self._observation_elements:
+            transition_elements.append(ReplayElement(
+                element.name,
+                (batch_size, self._timesteps) + tuple(element.shape),
+                element.type, True))
+            transition_elements.append(ReplayElement(
+                element.name + '_tp1',
+                (batch_size, self._timesteps) + tuple(element.shape),
+                element.type, True))
+
+        for element in self._extra_replay_elements:
+            transition_elements.append(ReplayElement(
+                element.name,
+                (batch_size,) + tuple(element.shape),
+                element.type))
+        return transition_elements

external/yarr/yarr/replay_buffer/uniform_replay_buffer.py

@@ -0,0 +1,804 @@
+"""The standard DQN replay memory.
+
+This implementation is an out-of-graph replay memory + in-graph wrapper. It
+supports vanilla n-step updates of the form typically found in the literature,
+i.e. where rewards are accumulated for n steps and the intermediate trajectory
+is not exposed to the agent. This does not allow, for example, performing
+off-policy corrections.
+"""
+import ctypes
+import collections
+import concurrent.futures
+import os
+from os.path import join
+import pickle
+from typing import List, Tuple, Type
+import time
+import math
+# from threading import Lock
+import multiprocessing as mp
+from multiprocessing import Lock
+import numpy as np
+import logging
+
+from natsort import natsort
+
+from yarr.replay_buffer.replay_buffer import ReplayBuffer, ReplayElement
+from yarr.utils.observation_type import ObservationElement
+
+import torch.distributed as dist
+
+# Defines a type describing part of the tuple returned by the replay
+# memory. Each element of the tuple is a tensor of shape [batch, ...] where
+# ... is defined the 'shape' field of ReplayElement. The tensor type is
+# given by the 'type' field. The 'name' field is for convenience and ease of
+# debugging.
+
+
+# String constants for storage
+ACTION = 'action'
+REWARD = 'reward'
+TERMINAL = 'terminal'
+TIMEOUT = 'timeout'
+INDICES = 'indices'
+
+
+def invalid_range(cursor, replay_capacity, stack_size, update_horizon):
+    """Returns a array with the indices of cursor-related invalid transitions.
+
+    There are update_horizon + stack_size invalid indices:
+      - The update_horizon indices before the cursor, because we do not have a
+        valid N-step transition (including the next state).
+      - The stack_size indices on or immediately after the cursor.
+    If N = update_horizon, K = stack_size, and the cursor is at c, invalid
+    indices are:
+      c - N, c - N + 1, ..., c, c + 1, ..., c + K - 1.
+
+    It handles special cases in a circular buffer in the beginning and the end.
+
+    Args:
+      cursor: int, the position of the cursor.
+      replay_capacity: int, the size of the replay memory.
+      stack_size: int, the size of the stacks returned by the replay memory.
+      update_horizon: int, the agent's update horizon.
+    Returns:
+      np.array of size stack_size with the invalid indices.
+    """
+    assert cursor < replay_capacity
+    return np.array(
+        [(cursor - update_horizon + i) % replay_capacity
+         for i in range(stack_size + update_horizon)])
+
+
+class UniformReplayBuffer(ReplayBuffer):
+    """A simple out-of-graph Replay Buffer.
+
+    Stores transitions, state, action, reward, next_state, terminal (and any
+    extra contents specified) in a circular buffer and provides a uniform
+    transition sampling function.
+
+    When the states consist of stacks of observations storing the states is
+    inefficient. This class writes observations and constructs the stacked states
+    at sample time.
+
+    Attributes:
+      _add_count: int, counter of how many transitions have been added (including
+        the blank ones at the beginning of an episode).
+      invalid_range: np.array, an array with the indices of cursor-related invalid
+        transitions
+    """
+
+    def __init__(self,
+                 batch_size: int = 32,
+                 timesteps: int = 1,
+                 replay_capacity: int = int(1e6),
+                 update_horizon: int = 1,
+                 gamma: float = 0.99,
+                 max_sample_attempts: int = 10000,
+                 action_shape: tuple = (),
+                 action_dtype: Type[np.dtype] = np.float32,
+                 reward_shape: tuple = (),
+                 reward_dtype: Type[np.dtype] = np.float32,
+                 observation_elements: List[ObservationElement] = None,
+                 extra_replay_elements: List[ReplayElement] = None,
+                 save_dir: str = None,
+                 purge_replay_on_shutdown: bool = True,
+                 num_replicas: int = None,
+                 rank: int = None,
+                 ):
+        """Initializes OutOfGraphReplayBuffer.
+
+        Args:
+          batch_size: int.
+          timesteps: int, number of frames to use in state stack.
+          replay_capacity: int, number of transitions to keep in memory.
+          update_horizon: int, length of update ('n' in n-step update).
+          gamma: int, the discount factor.
+          max_sample_attempts: int, the maximum number of attempts allowed to
+            get a sample.
+          action_shape: tuple of ints, the shape for the action vector.
+            Empty tuple means the action is a scalar.
+          action_dtype: np.dtype, type of elements in the action.
+          reward_shape: tuple of ints, the shape of the reward vector.
+            Empty tuple means the reward is a scalar.
+          reward_dtype: np.dtype, type of elements in the reward.
+          observation_elements: list of ObservationElement defining the type of
+            the extra contents that will be stored and returned.
+          extra_storage_elements: list of ReplayElement defining the type of
+            the extra contents that will be stored and returned.
+
+        Raises:
+          ValueError: If replay_capacity is too small to hold at least one
+            transition.
+        """
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            self._num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            self._rank = dist.get_rank()
+        if self._rank >= self._num_replicas or self._rank < 0:
+            raise ValueError(
+                "Invalid rank {}, rank should be in the interval"
+                " [0, {}]".format(self._rank, self._num_replicas - 1))
+
+        if observation_elements is None:
+            observation_elements = []
+        if extra_replay_elements is None:
+            extra_replay_elements = []
+
+        if replay_capacity < update_horizon + timesteps:
+            raise ValueError('There is not enough capacity to cover '
+                             'update_horizon and stack_size.')
+
+        logging.info(
+            'Creating a %s replay memory with the following parameters:',
+            self.__class__.__name__)
+        logging.info('\t timesteps: %d', timesteps)
+        logging.info('\t replay_capacity: %d', replay_capacity)
+        logging.info('\t batch_size: %d', batch_size)
+        logging.info('\t update_horizon: %d', update_horizon)
+        logging.info('\t gamma: %f', gamma)
+
+        self._disk_saving = save_dir is not None
+        self._save_dir = save_dir
+        self._purge_replay_on_shutdown = purge_replay_on_shutdown
+        if self._disk_saving:
+            logging.info('\t saving to disk: %s', self._save_dir)
+            os.makedirs(save_dir, exist_ok=True)
+        else:
+            logging.info('\t saving to RAM')
+
+
+        self._action_shape = action_shape
+        self._action_dtype = action_dtype
+        self._reward_shape = reward_shape
+        self._reward_dtype = reward_dtype
+        self._timesteps = timesteps
+        self._replay_capacity = replay_capacity
+        self._batch_size = batch_size
+        self._update_horizon = update_horizon
+        self._gamma = gamma
+        self._max_sample_attempts = max_sample_attempts
+
+        self._observation_elements = observation_elements
+        self._extra_replay_elements = extra_replay_elements
+
+        self._storage_signature, self._obs_signature = self.get_storage_signature()
+        self._create_storage()
+
+        self._lock = Lock()
+        self._add_count = mp.Value('i', 0)
+
+        self._replay_capacity = replay_capacity
+
+        self.invalid_range = np.zeros((self._timesteps))
+
+        # When the horizon is > 1, we compute the sum of discounted rewards as a dot
+        # product using the precomputed vector <gamma^0, gamma^1, ..., gamma^{n-1}>.
+        self._cumulative_discount_vector = np.array(
+            [math.pow(self._gamma, n) for n in range(update_horizon)],
+            dtype=np.float32)
+
+    @property
+    def timesteps(self):
+        return self._timesteps
+
+    @property
+    def replay_capacity(self):
+        return self._replay_capacity
+
+    @property
+    def batch_size(self):
+        return self._batch_size
+
+    def _create_storage(self, store=None):
+        """Creates the numpy arrays used to store transitions.
+        """
+        self._store = {} if store is None else store
+        for storage_element in self._storage_signature:
+            array_shape = [self._replay_capacity] + list(storage_element.shape)
+            if storage_element.name == TERMINAL:
+                self._store[storage_element.name] = np.full(
+                    array_shape, -1, dtype=storage_element.type)
+            elif not self._disk_saving:
+                # If saving to disk, we don't need to store anything else.
+                self._store[storage_element.name] = np.empty(
+                    array_shape, dtype=storage_element.type)
+
+    def get_storage_signature(self) -> Tuple[List[ReplayElement],
+                                             List[ReplayElement]]:
+        """Returns a default list of elements to be stored in this replay memory.
+
+        Note - Derived classes may return a different signature.
+
+        Returns:
+          dict of ReplayElements defining the type of the contents stored.
+        """
+        storage_elements = [
+            ReplayElement(ACTION, self._action_shape, self._action_dtype),
+            ReplayElement(REWARD, self._reward_shape, self._reward_dtype),
+            ReplayElement(TERMINAL, (), np.int8),
+            ReplayElement(TIMEOUT, (), bool),
+        ]
+
+        obs_elements = []
+        for obs_element in self._observation_elements:
+            obs_elements.append(
+                ReplayElement(
+                    obs_element.name, obs_element.shape, obs_element.type))
+        storage_elements.extend(obs_elements)
+
+        for extra_replay_element in self._extra_replay_elements:
+            storage_elements.append(extra_replay_element)
+
+        return storage_elements, obs_elements
+
+    def add(self, action, reward, terminal, timeout, **kwargs):
+        """Adds a transition to the replay memory.
+
+        WE ONLY STORE THE TPS1s on the final frame
+
+        This function checks the types and handles the padding at the beginning of
+        an episode. Then it calls the _add function.
+
+        Since the next_observation in the transition will be the observation added
+        next there is no need to pass it.
+
+        If the replay memory is at capacity the oldest transition will be discarded.
+
+        Args:
+          action: int, the action in the transition.
+          reward: float, the reward received in the transition.
+          terminal: A uint8 acting as a boolean indicating whether the transition
+                    was terminal (1) or not (0).
+          **kwargs: The remaining args
+        """
+
+        # If previous transition was a terminal, then add_final wasn't called
+        # if not self.is_empty() and self._store['terminal'][self.cursor() - 1] == 1:
+        #     raise ValueError('The previous transition was a terminal, '
+        #                      'but add_final was not called.')
+
+        kwargs[ACTION] = action
+        kwargs[REWARD] = reward
+        kwargs[TERMINAL] = terminal
+        kwargs[TIMEOUT] = timeout
+        self._check_add_types(kwargs, self._storage_signature)
+        self._add(kwargs)
+
+    def add_final(self, **kwargs):
+        """Adds a transition to the replay memory.
+        Args:
+          **kwargs: The remaining args
+        """
+        # if self.is_empty() or self._store['terminal'][self.cursor() - 1] != 1:
+        #     raise ValueError('The previous transition was not terminal.')
+        self._check_add_types(kwargs, self._obs_signature)
+        transition = self._final_transition(kwargs)
+        self._add(transition)
+
+    def _final_transition(self, kwargs):
+        transition = {}
+        for element_type in self._storage_signature:
+            if element_type.name in kwargs:
+                transition[element_type.name] = kwargs[element_type.name]
+            elif element_type.name == TERMINAL:
+                # Used to check that user is correctly adding transitions
+                transition[element_type.name] = -1
+            else:
+                transition[element_type.name] = np.empty(
+                    element_type.shape, dtype=element_type.type)
+        return transition
+
+    def _add_initial_to_disk(self ,kwargs: dict):
+        for i in range(self._timesteps - 1):
+            with open(join(self._save_dir, '%d.replay' % (
+                    self._replay_capacity - 1 - i)), 'wb') as f:
+                pickle.dump(kwargs, f)
+
+    def _add(self, kwargs: dict):
+        """Internal add method to add to the storage arrays.
+
+        Args:
+          kwargs: All the elements in a transition.
+        """
+        with self._lock:
+            cursor = self.cursor()
+
+            if self._disk_saving:
+                term = self._store[TERMINAL]
+                term[cursor] = kwargs[TERMINAL]
+                self._store[TERMINAL] = term
+                with open(join(self._save_dir, '%d.replay' % cursor), 'wb') as f:
+                    pickle.dump(kwargs, f)
+                # If first add, then pad for correct wrapping
+                if self._add_count.value == 0:
+                    self._add_initial_to_disk(kwargs)
+            else:
+                for name, data in kwargs.items():
+                    item = self._store[name]
+                    item[cursor] = data
+                    self._store[name] = item
+            with self._add_count.get_lock():
+                self._add_count.value += 1
+            self.invalid_range = invalid_range(
+                self.cursor(), self._replay_capacity, self._timesteps,
+                self._update_horizon)
+
+    def _get_from_disk(self, start_index, end_index):
+        """Returns the range of array at the index handling wraparound if necessary.
+
+        Args:
+          start_index: int, index to the start of the range to be returned. Range
+            will wraparound if start_index is smaller than 0.
+          end_index: int, exclusive end index. Range will wraparound if end_index
+            exceeds replay_capacity.
+
+        Returns:
+          np.array, with shape [end_index - start_index, array.shape[1:]].
+        """
+        assert end_index > start_index, 'end_index must be larger than start_index'
+        assert end_index >= 0
+        assert start_index < self._replay_capacity
+        if not self.is_full():
+            assert end_index <= self.cursor(), (
+                'Index {} has not been added.'.format(start_index))
+
+        # Here we fake a mini store (buffer)
+        store = {store_element.name: {}
+                 for store_element in self._storage_signature}
+        if start_index % self._replay_capacity < end_index % self._replay_capacity:
+            for i in range(start_index, end_index):
+                with open(join(self._save_dir, '%d.replay' % i), 'rb') as f:
+                    d = pickle.load(f)
+                    for k, v in d.items():
+                        store[k][i] = v
+        else:
+            for i in range(end_index - start_index):
+                idx = (start_index + i) % self._replay_capacity
+                with open(join(self._save_dir, '%d.replay' % idx), 'rb') as f:
+                    d = pickle.load(f)
+                    for k, v in d.items():
+                        store[k][idx] = v
+        return store
+
+    def _check_add_types(self, kwargs, signature):
+        """Checks if args passed to the add method match those of the storage.
+
+        Args:
+          *args: Args whose types need to be validated.
+
+        Raises:
+          ValueError: If args have wrong shape or dtype.
+        """
+
+        if (len(kwargs)) != len(signature):
+            expected = str(natsort.natsorted([e.name for e in signature]))
+            actual = str(natsort.natsorted(list(kwargs.keys())))
+            error_list = '\nList of expected:\n{}\nList of actual:\n{}'.format(
+                expected, actual)
+            raise ValueError('Add expects {} elements, received {}.'.format(
+                len(signature), len(kwargs)) + error_list)
+
+        for store_element in signature:
+            arg_element = kwargs[store_element.name]
+            if isinstance(arg_element, np.ndarray):
+                arg_shape = arg_element.shape
+            elif isinstance(arg_element, tuple) or isinstance(arg_element, list):
+                # TODO: This is not efficient when arg_element is a list.
+                arg_shape = np.array(arg_element).shape
+            else:
+                # Assume it is scalar.
+                arg_shape = tuple()
+            store_element_shape = tuple(store_element.shape)
+            if arg_shape != store_element_shape:
+                raise ValueError('arg {} has shape {}, expected {}'.format(store_element.name,
+                    arg_shape, store_element_shape))
+
+    def is_empty(self):
+        """Is the Replay Buffer empty?"""
+        return self._add_count.value == 0
+
+    def is_full(self):
+        """Is the Replay Buffer full?"""
+        return self._add_count.value >= self._replay_capacity
+
+    def cursor(self):
+        """Index to the location where the next transition will be written."""
+        return self._add_count.value % self._replay_capacity
+
+    @property
+    def add_count(self):
+        return np.array(self._add_count.value)
+
+    @add_count.setter
+    def add_count(self, count):
+        if isinstance(count, int):
+            self._add_count = mp.Value('i', count)
+        else:
+            self._add_count = count
+
+
+    def get_range(self, array, start_index, end_index):
+        """Returns the range of array at the index handling wraparound if necessary.
+
+        Args:
+          array: np.array, the array to get the stack from.
+          start_index: int, index to the start of the range to be returned. Range
+            will wraparound if start_index is smaller than 0.
+          end_index: int, exclusive end index. Range will wraparound if end_index
+            exceeds replay_capacity.
+
+        Returns:
+          np.array, with shape [end_index - start_index, array.shape[1:]].
+        """
+        assert end_index > start_index, 'end_index must be larger than start_index'
+        assert end_index >= 0
+        assert start_index < self._replay_capacity
+        if not self.is_full():
+            assert end_index <= self.cursor(), (
+                'Index {} has not been added.'.format(start_index))
+
+        # Fast slice read when there is no wraparound.
+        if start_index % self._replay_capacity < end_index % self._replay_capacity:
+            return_array = np.array(
+                [array[i] for i in range(start_index, end_index)])
+        # Slow list read.
+        else:
+            indices = [(start_index + i) % self._replay_capacity
+                       for i in range(end_index - start_index)]
+            return_array = np.array([array[i] for i in indices])
+
+        return return_array
+
+    def get_range_stack(self, array, start_index, end_index, terminals=None):
+        """Returns the range of array at the index handling wraparound if necessary.
+
+        Args:
+          array: np.array, the array to get the stack from.
+          start_index: int, index to the start of the range to be returned. Range
+            will wraparound if start_index is smaller than 0.
+          end_index: int, exclusive end index. Range will wraparound if end_index
+            exceeds replay_capacity.
+
+        Returns:
+          np.array, with shape [end_index - start_index, array.shape[1:]].
+        """
+        return_array = np.array(self.get_range(array, start_index, end_index))
+        if terminals is None:
+            terminals = self.get_range(
+                self._store[TERMINAL], start_index, end_index)
+
+        terminals = terminals[:-1]
+
+        # Here we now check if we need to pad the front episodes
+        # If any have a terminal of -1, then we have spilled over
+        # into the the previous transition
+        if np.any(terminals == -1):
+            padding_item = return_array[-1]
+            _array = list(return_array)[:-1]
+            arr_len = len(_array)
+            pad_from_now = False
+            for i, (ar, term) in enumerate(
+                    zip(reversed(_array), reversed(terminals))):
+                if term == -1 or pad_from_now:
+                    # The first time we see a -1 term, means we have hit the
+                    # beginning of this episode, so pad from now.
+                    # pad_from_now needed because the next transition (reverse)
+                    # will not be a -1 terminal.
+                    pad_from_now = True
+                    return_array[arr_len - 1 - i] = padding_item
+                else:
+                    # After we hit out first -1 terminal, we never reassign.
+                    padding_item = ar
+
+        return return_array
+
+    def _get_element_stack(self, array, index, terminals=None):
+        state = self.get_range_stack(array,
+                                     index - self._timesteps + 1, index + 1,
+                                     terminals=terminals)
+        return state
+
+    def get_terminal_stack(self, index):
+        terminal_stack = self.get_range(self._store[TERMINAL],
+                              index - self._timesteps + 1,
+                              index + 1)
+        return terminal_stack
+
+    def is_valid_transition(self, index):
+        """Checks if the index contains a valid transition.
+
+        Checks for collisions with the end of episodes and the current position
+        of the cursor.
+
+        Args:
+          index: int, the index to the state in the transition.
+
+        Returns:
+          Is the index valid: Boolean.
+
+        """
+        # Check the index is in the valid range
+        if index < 0 or index >= self._replay_capacity:
+            return False
+        if not self.is_full():
+            # The indices and next_indices must be smaller than the cursor.
+            if index >= self.cursor() - self._update_horizon:
+                return False
+
+        # Skip transitions that straddle the cursor.
+        if index in set(self.invalid_range):
+            return False
+
+        term_stack = self.get_terminal_stack(index)
+        if term_stack[-1] == -1:
+            return False
+
+        return True
+
+    def _create_batch_arrays(self, batch_size):
+        """Create a tuple of arrays with the type of get_transition_elements.
+
+        When using the WrappedReplayBuffer with staging enabled it is important
+        to create new arrays every sample because StaginArea keeps a pointer to
+        the returned arrays.
+
+        Args:
+          batch_size: (int) number of transitions returned. If None the default
+            batch_size will be used.
+
+        Returns:
+          Tuple of np.arrays with the shape and type of get_transition_elements.
+        """
+        transition_elements = self.get_transition_elements(batch_size)
+        batch_arrays = []
+        for element in transition_elements:
+            batch_arrays.append(np.empty(element.shape, dtype=element.type))
+        return tuple(batch_arrays)
+
+    def sample_index_batch(self, batch_size):
+        """Returns a batch of valid indices sampled uniformly.
+
+        Args:
+          batch_size: int, number of indices returned.
+
+        Returns:
+          list of ints, a batch of valid indices sampled uniformly.
+
+        Raises:
+          RuntimeError: If the batch was not constructed after maximum number of
+            tries.
+        """
+        if self.is_full():
+            # add_count >= self._replay_capacity > self._stack_size
+            min_id = (self.cursor() - self._replay_capacity +
+                      self._timesteps - 1)
+            max_id = self.cursor() - self._update_horizon
+        else:
+            min_id = 0
+            max_id = self.cursor() - self._update_horizon
+            if max_id <= min_id:
+                raise RuntimeError(
+                    'Cannot sample a batch with fewer than stack size '
+                    '({}) + update_horizon ({}) transitions.'.
+                    format(self._timesteps, self._update_horizon))
+
+        indices = []
+        attempt_count = 0
+        while (len(indices) < batch_size and
+                       attempt_count < self._max_sample_attempts):
+            index = np.random.randint(min_id, max_id) % self._replay_capacity
+            if self.is_valid_transition(index):
+                indices.append(index)
+            else:
+                attempt_count += 1
+        if len(indices) != batch_size:
+            raise RuntimeError(
+                'Max sample attempts: Tried {} times but only sampled {}'
+                ' valid indices. Batch size is {}'.
+                    format(self._max_sample_attempts, len(indices), batch_size))
+
+        return indices
+
+    def unpack_transition(self, transition_tensors, transition_type):
+        """Unpacks the given transition into member variables.
+
+        Args:
+          transition_tensors: tuple of tf.Tensors.
+          transition_type: tuple of ReplayElements matching transition_tensors.
+        """
+        self.transition = collections.OrderedDict()
+        for element, element_type in zip(transition_tensors, transition_type):
+            self.transition[element_type.name] = element
+        return self.transition
+
+    def sample_transition_batch(self, batch_size=None, indices=None,
+                                pack_in_dict=True):
+        """Returns a batch of transitions (including any extra contents).
+
+        If get_transition_elements has been overridden and defines elements not
+        stored in self._store, an empty array will be returned and it will be
+        left to the child class to fill it. For example, for the child class
+        OutOfGraphPrioritizedReplayBuffer, the contents of the
+        sampling_probabilities are stored separately in a sum tree.
+
+        When the transition is terminal next_state_batch has undefined contents.
+
+        NOTE: This transition contains the indices of the sampled elements.
+        These are only valid during the call to sample_transition_batch,
+        i.e. they may  be used by subclasses of this replay buffer but may
+        point to different data as soon as sampling is done.
+
+        Args:
+          batch_size: int, number of transitions returned. If None, the default
+            batch_size will be used.
+          indices: None or list of ints, the indices of every transition in the
+            batch. If None, sample the indices uniformly.
+
+        Returns:
+          transition_batch: tuple of np.arrays with the shape and type as in
+            get_transition_elements().
+
+        Raises:
+          ValueError: If an element to be sampled is missing from the
+            replay buffer.
+        """
+        
+        if batch_size is None:
+            batch_size = self._batch_size
+        with self._lock:
+            if indices is None:
+                indices = self.sample_index_batch(batch_size)
+            assert len(indices) == batch_size
+
+            transition_elements = self.get_transition_elements(batch_size)
+            batch_arrays = self._create_batch_arrays(batch_size)
+
+            for batch_element, state_index in enumerate(indices):
+
+                if not self.is_valid_transition(state_index):
+                    raise ValueError('Invalid index %d.' % state_index)
+
+                trajectory_indices = [(state_index + j) % self._replay_capacity
+                                      for j in range(self._update_horizon)]
+                trajectory_terminals = self._store['terminal'][
+                    trajectory_indices]
+                is_terminal_transition = trajectory_terminals.any()
+                if not is_terminal_transition:
+                    trajectory_length = self._update_horizon
+                else:
+                    # np.argmax of a bool array returns index of the first True.
+                    trajectory_length = np.argmax(
+                        trajectory_terminals.astype(bool),
+                        0) + 1
+
+                next_state_index = state_index + trajectory_length
+
+                store = self._store
+                if self._disk_saving:
+                    store = self._get_from_disk(
+                        state_index - (self._timesteps - 1),
+                        next_state_index + 1)
+
+                trajectory_discount_vector = (
+                    self._cumulative_discount_vector[:trajectory_length])
+                trajectory_rewards = self.get_range(store['reward'],
+                                                    state_index,
+                                                    next_state_index)
+
+                terminal_stack = self.get_terminal_stack(state_index)
+                terminal_stack_tp1 = self.get_terminal_stack(
+                    next_state_index % self._replay_capacity)
+
+                # Fill the contents of each array in the sampled batch.
+                assert len(transition_elements) == len(batch_arrays)
+                for element_array, element in zip(batch_arrays,
+                                                  transition_elements):
+                    if element.is_observation:
+                        if element.name.endswith('tp1'):
+                            element_array[
+                                batch_element] = self._get_element_stack(
+                                store[element.name[:-4]],
+                                next_state_index % self._replay_capacity,
+                                terminal_stack_tp1)
+                        else:
+                            element_array[
+                                batch_element] = self._get_element_stack(
+                                store[element.name],
+                                state_index, terminal_stack)
+                    elif element.name == REWARD:
+                        # compute discounted sum of rewards in the trajectory.
+                        element_array[batch_element] = np.sum(
+                            trajectory_discount_vector * trajectory_rewards,
+                            axis=0)
+                    elif element.name == TERMINAL:
+                        element_array[batch_element] = is_terminal_transition
+                    elif element.name == INDICES:
+                        element_array[batch_element] = state_index
+                    elif element.name in store.keys():
+                        element_array[batch_element] = (
+                            store[element.name][state_index])
+
+        if pack_in_dict:
+            batch_arrays = self.unpack_transition(
+                batch_arrays, transition_elements)
+
+        # TODO(Mohit): proper fix to discard task names
+        if 'task' in batch_arrays:
+            del batch_arrays['task']
+        if 'task_tp1' in batch_arrays:
+            del batch_arrays['task_tp1']
+
+        return batch_arrays
+
+    def get_transition_elements(self, batch_size=None):
+        """Returns a 'type signature' for sample_transition_batch.
+
+        Args:
+          batch_size: int, number of transitions returned. If None, the default
+            batch_size will be used.
+        Returns:
+          signature: A namedtuple describing the method's return type signature.
+        """
+        batch_size = self._batch_size if batch_size is None else batch_size
+
+        transition_elements = [
+            ReplayElement(ACTION, (batch_size,) + self._action_shape,
+                          self._action_dtype),
+            ReplayElement(REWARD, (batch_size,) + self._reward_shape,
+                          self._reward_dtype),
+            ReplayElement(TERMINAL, (batch_size,), np.int8),
+            ReplayElement(TIMEOUT, (batch_size,), bool),
+            ReplayElement(INDICES, (batch_size,), np.int32),
+        ]
+
+        for element in self._observation_elements:
+            transition_elements.append(ReplayElement(
+                element.name,
+                (batch_size, self._timesteps) + tuple(element.shape),
+                element.type, True))
+            transition_elements.append(ReplayElement(
+                element.name + '_tp1',
+                (batch_size, self._timesteps) + tuple(element.shape),
+                element.type, True))
+
+        for element in self._extra_replay_elements:
+            transition_elements.append(ReplayElement(
+                element.name,
+                (batch_size,) + tuple(element.shape),
+                element.type))
+        return transition_elements
+
+    def shutdown(self):
+        if self._purge_replay_on_shutdown:
+            # Safely delete replay
+            logging.info('Clearing disk replay buffer.')
+            for f in [f for f in os.listdir(self._save_dir) if '.replay' in f]:
+                os.remove(join(self._save_dir, f))
+
+    def using_disk(self):
+        return self._disk_saving

external/yarr/yarr/replay_buffer/wrappers/__init__.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from yarr.replay_buffer.replay_buffer import ReplayBuffer
+
+
+class WrappedReplayBuffer(ABC):
+
+    def __init__(self, replay_buffer: ReplayBuffer):
+        """Initializes WrappedReplayBuffer.
+
+        Raises:
+          ValueError: If update_horizon is not positive.
+          ValueError: If discount factor is not in [0, 1].
+        """
+        self._replay_buffer = replay_buffer
+
+    @property
+    def replay_buffer(self):
+        return self._replay_buffer
+
+    @abstractmethod
+    def dataset(self) -> Any:
+        pass

external/yarr/yarr/replay_buffer/wrappers/pytorch_replay_buffer.py

@@ -0,0 +1,82 @@
+import time
+from threading import Thread
+
+from torch.utils.data import IterableDataset, DataLoader
+
+from yarr.replay_buffer.replay_buffer import ReplayBuffer
+from yarr.replay_buffer.wrappers import WrappedReplayBuffer
+
+
+class PyTorchIterableReplayDataset(IterableDataset):
+
+    def __init__(self, replay_buffer: ReplayBuffer):
+        self._replay_buffer = replay_buffer
+
+    def _generator(self):
+        while True:
+            yield self._replay_buffer.sample_transition_batch(pack_in_dict=True)
+
+    def __iter__(self):
+        return iter(self._generator())
+
+# class PyTorchIterableReplayDataset(IterableDataset):
+#
+#     BUFFER = 4
+#
+#     def __init__(self, replay_buffer: ReplayBuffer, num_workers: int):
+#         self._replay_buffer = replay_buffer
+#         self._num_wokers = num_workers
+#         self._samples = []
+#         self._lock = Lock()
+#
+#     def _run(self):
+#         while True:
+#             # Check if replay buffer is ig enough to be sampled
+#             while self._replay_buffer.add_count < self._replay_buffer.batch_size:
+#                 time.sleep(1.)
+#             s = self._replay_buffer.sample_transition_batch(pack_in_dict=True)
+#             while len(self._samples) >= PyTorchIterableReplayDataset.BUFFER:
+#                 time.sleep(0.25)
+#             with self._lock:
+#                 self._samples.append(s)
+#
+#     def _generator(self):
+#         ts = [Thread(
+#             target=self._run, args=()) for _ in range(self._num_wokers)]
+#         [t.start() for t in ts]
+#         while True:
+#             while len(self._samples) == 0:
+#                 time.sleep(0.1)
+#             with self._lock:
+#                 s = self._samples.pop(0)
+#             yield s
+#
+#     def __iter__(self):
+#         i = iter(self._generator())
+#         return i
+
+
+class PyTorchReplayBuffer(WrappedReplayBuffer):
+    """Wrapper of OutOfGraphReplayBuffer with an in graph sampling mechanism.
+
+    Usage:
+      To add a transition:  call the add function.
+
+      To sample a batch:    Construct operations that depend on any of the
+                            tensors is the transition dictionary. Every sess.run
+                            that requires any of these tensors will sample a new
+                            transition.
+    """
+
+    def __init__(self, replay_buffer: ReplayBuffer, num_workers: int = 2):
+        super(PyTorchReplayBuffer, self).__init__(replay_buffer)
+        self._num_workers = num_workers
+
+    def dataset(self, batch_size=None, drop_last=False) -> DataLoader:
+        # d = PyTorchIterableReplayDataset(self._replay_buffer)
+        d = PyTorchIterableReplayDataset(self._replay_buffer)
+
+        # Batch size None disables automatic batching
+        return DataLoader(d, batch_size=batch_size,
+                          drop_last=drop_last,
+                          num_workers=self._num_workers, pin_memory=True)

external/yarr/yarr/runners/_env_runner.py

@@ -0,0 +1,228 @@
+import copy
+import logging
+import os
+import time
+import pandas as pd
+
+from multiprocessing import Process, Manager
+from multiprocessing import get_start_method, set_start_method
+from typing import Any
+
+import numpy as np
+import torch
+from yarr.agents.agent import Agent
+from yarr.agents.agent import ScalarSummary
+from yarr.agents.agent import Summary
+from yarr.envs.env import Env
+from yarr.utils.rollout_generator import RolloutGenerator
+from yarr.utils.log_writer import LogWriter
+from yarr.utils.process_str import change_case
+from yarr.utils.video_utils import CircleCameraMotion, TaskRecorder
+
+from pyrep.objects.dummy import Dummy
+from pyrep.objects.vision_sensor import VisionSensor
+
+try:
+    if get_start_method() != 'spawn':
+        set_start_method('spawn', force=True)
+except RuntimeError:
+    pass
+
+
+class _EnvRunner(object):
+
+    def __init__(self,
+                 train_env: Env,
+                 eval_env: Env,
+                 agent: Agent,
+                 timesteps: int,
+                 train_envs: int,
+                 eval_envs: int,
+                 rollout_episodes: int,
+                 eval_episodes: int,
+                 training_iterations: int,
+                 eval_from_eps_number: int,
+                 episode_length: int,
+                 kill_signal: Any,
+                 step_signal: Any,
+                 num_eval_episodes_signal: Any,
+                 eval_epochs_signal: Any,
+                 eval_report_signal: Any,
+                 log_freq: int,
+                 rollout_generator: RolloutGenerator,
+                 save_load_lock,
+                 current_replay_ratio,
+                 target_replay_ratio,
+                 weightsdir: str = None,
+                 logdir: str = None,
+                 env_device: torch.device = None,
+                 previous_loaded_weight_folder: str = '',
+                 num_eval_runs: int = 1,
+                 ):
+        self._train_env = train_env
+        self._eval_env = eval_env
+        self._agent = agent
+        self._train_envs = train_envs
+        self._eval_envs = eval_envs
+        self._rollout_episodes = rollout_episodes
+        self._eval_episodes = eval_episodes
+        self._training_iterations = training_iterations
+        self._num_eval_runs = num_eval_runs
+        self._eval_from_eps_number = eval_from_eps_number
+        self._episode_length = episode_length
+        self._rollout_generator = rollout_generator
+        self._weightsdir = weightsdir
+        self._logdir = logdir
+        self._env_device = env_device
+        self._previous_loaded_weight_folder = previous_loaded_weight_folder
+
+        self._timesteps = timesteps
+
+        self._p_args = {}
+        self.p_failures = {}
+        manager = Manager()
+        self.write_lock = manager.Lock()
+        self.stored_transitions = manager.list()
+        self.agent_summaries = manager.list()
+        self._kill_signal = kill_signal
+        self._step_signal = step_signal
+        self._num_eval_episodes_signal = num_eval_episodes_signal
+        self._eval_epochs_signal = eval_epochs_signal
+        self._eval_report_signal = eval_report_signal
+        self._save_load_lock = save_load_lock
+        self._current_replay_ratio = current_replay_ratio
+        self._target_replay_ratio = target_replay_ratio
+        self._log_freq = log_freq
+
+        self._new_weights = False
+
+    def restart_process(self, name: str):
+        p = Process(target=self._run_env, args=self._p_args[name], name=name)
+        p.start()
+        return p
+
+    def spin_up_envs(self, name: str, num_envs: int, eval: bool):
+
+        ps = []
+        for i in range(num_envs):
+            n = name + str(i)
+            self._p_args[n] = (n, eval)
+            self.p_failures[n] = 0
+            p = Process(target=self._run_env, args=self._p_args[n], name=n)
+            p.start()
+            ps.append(p)
+        return ps
+
+    def _load_save(self):
+        if self._weightsdir is None:
+            logging.info("'weightsdir' was None, so not loading weights.")
+            return
+        while True:
+            weight_folders = []
+            with self._save_load_lock:
+                if os.path.exists(self._weightsdir):
+                    weight_folders = os.listdir(self._weightsdir)
+                if len(weight_folders) > 0:
+                    weight_folders = sorted(map(int, weight_folders))
+                    # Only load if there has been a new weight saving
+                    if self._previous_loaded_weight_folder != weight_folders[-1]:
+                        self._previous_loaded_weight_folder = weight_folders[-1]
+                        d = os.path.join(self._weightsdir, str(weight_folders[-1]))
+                        try:
+                            self._agent.load_weights(d)
+                        except FileNotFoundError:
+                            # Rare case when agent hasn't finished writing.
+                            time.sleep(1)
+                            self._agent.load_weights(d)
+                        print('Agent %s: Loaded weights: %s' % (self._name, d))
+                        self._new_weights = True
+                    else:
+                        self._new_weights = False
+                    break
+            print('Waiting for weights to become available.')
+            time.sleep(1)
+
+    def _get_type(self, x):
+        if x.dtype == np.float64:
+            return np.float32
+        return x.dtype
+
+    def _get_task_name(self):
+        if hasattr(self._eval_env, '_task_class'):
+            eval_task_name = change_case(self._eval_env._task_class.__name__)
+            multi_task = False
+        elif hasattr(self._eval_env, '_task_classes'):
+            if self._eval_env.active_task_id != -1:
+                task_id = (self._eval_env.active_task_id) % len(self._eval_env._task_classes)
+                eval_task_name = change_case(self._eval_env._task_classes[task_id].__name__)
+            else:
+                eval_task_name = ''
+            multi_task = True
+        else:
+            raise Exception('Neither task_class nor task_classes found in eval env')
+        return eval_task_name, multi_task
+
+    def _run_env(self, name: str, eval: bool):
+
+        self._name = name
+
+        self._agent = copy.deepcopy(self._agent)
+
+        self._agent.build(training=False, device=self._env_device)
+
+        logging.info('%s: Launching env.' % name)
+        np.random.seed()
+
+        logging.info('Agent information:')
+        logging.info(self._agent)
+
+        env = self._train_env
+        if eval:
+            env = self._eval_env
+        env.eval = eval
+        env.launch()
+        for ep in range(self._rollout_episodes):
+            self._load_save()
+            logging.debug('%s: Starting episode %d.' % (name, ep))
+            episode_rollout = []
+            generator = self._rollout_generator.generator(
+                self._step_signal, env, self._agent,
+                self._episode_length, self._timesteps,
+                eval, eval_demo_seed=eval_demo_seed,
+                record_enabled=rec_cfg.enabled)
+            try:
+                for replay_transition in generator:
+                    while True:
+                        if self._kill_signal.value:
+                            env.shutdown()
+                            return
+                        if (eval or self._target_replay_ratio is None or
+                                self._step_signal.value <= 0 or (
+                                        self._current_replay_ratio.value >
+                                        self._target_replay_ratio)):
+                            break
+                        time.sleep(1)
+                        logging.debug(
+                            'Agent. Waiting for replay_ratio %f to be more than %f' %
+                            (self._current_replay_ratio.value, self._target_replay_ratio))
+
+                    with self.write_lock:
+                        if len(self.agent_summaries) == 0:
+                            # Only store new summaries if the previous ones
+                            # have been popped by the main env runner.
+                            for s in self._agent.act_summaries():
+                                self.agent_summaries.append(s)
+                    episode_rollout.append(replay_transition)
+            except StopIteration as e:
+                continue
+            except Exception as e:
+                env.shutdown()
+                raise e
+
+            with self.write_lock:
+                for transition in episode_rollout:
+                    self.stored_transitions.append((name, transition, eval))
+        env.shutdown()
+
+    def kill(self):
+        self._kill_signal.value = True

external/yarr/yarr/runners/_independent_env_runner.py

@@ -0,0 +1,297 @@
+import copy
+import logging
+import os
+import time
+import pandas as pd
+
+from multiprocessing import Process, Manager
+from multiprocessing import get_start_method, set_start_method
+from typing import Any
+
+import numpy as np
+import torch
+from yarr.agents.agent import Agent
+from yarr.agents.agent import ScalarSummary
+from yarr.agents.agent import Summary
+from yarr.envs.env import Env
+from yarr.utils.rollout_generator import RolloutGenerator
+from yarr.utils.log_writer import LogWriter
+from yarr.utils.process_str import change_case
+from yarr.utils.video_utils import CircleCameraMotion, TaskRecorder
+
+from pyrep.objects.dummy import Dummy
+from pyrep.objects.vision_sensor import VisionSensor
+
+from yarr.runners._env_runner import _EnvRunner
+
+
+class _IndependentEnvRunner(_EnvRunner):
+
+    def __init__(self,
+                 train_env: Env,
+                 eval_env: Env,
+                 agent: Agent,
+                 timesteps: int,
+                 train_envs: int,
+                 eval_envs: int,
+                 rollout_episodes: int,
+                 eval_episodes: int,
+                 training_iterations: int,
+                 eval_from_eps_number: int,
+                 episode_length: int,
+                 kill_signal: Any,
+                 step_signal: Any,
+                 num_eval_episodes_signal: Any,
+                 eval_epochs_signal: Any,
+                 eval_report_signal: Any,
+                 log_freq: int,
+                 rollout_generator: RolloutGenerator,
+                 save_load_lock,
+                 current_replay_ratio,
+                 target_replay_ratio,
+                 weightsdir: str = None,
+                 logdir: str = None,
+                 env_device: torch.device = None,
+                 previous_loaded_weight_folder: str = '',
+                 num_eval_runs: int = 1,
+                 ):
+
+            super().__init__(train_env, eval_env, agent, timesteps,
+                             train_envs, eval_envs, rollout_episodes, eval_episodes,
+                             training_iterations, eval_from_eps_number, episode_length,
+                             kill_signal, step_signal, num_eval_episodes_signal,
+                             eval_epochs_signal, eval_report_signal, log_freq,
+                             rollout_generator, save_load_lock, current_replay_ratio,
+                             target_replay_ratio, weightsdir, logdir, env_device,
+                             previous_loaded_weight_folder, num_eval_runs)
+
+    def _load_save(self):
+        if self._weightsdir is None:
+            logging.info("'weightsdir' was None, so not loading weights.")
+            return
+        while True:
+            weight_folders = []
+            with self._save_load_lock:
+                if os.path.exists(self._weightsdir):
+                    weight_folders = os.listdir(self._weightsdir)
+                if len(weight_folders) > 0:
+                    weight_folders = sorted(map(int, weight_folders))
+                    # only load if there has been a new weight saving
+                    if self._previous_loaded_weight_folder != weight_folders[-1]:
+                        self._previous_loaded_weight_folder = weight_folders[-1]
+                        d = os.path.join(self._weightsdir, str(weight_folders[-1]))
+                        try:
+                            self._agent.load_weights(d)
+                        except FileNotFoundError:
+                            # rare case when agent hasn't finished writing.
+                            time.sleep(1)
+                            self._agent.load_weights(d)
+                        logging.info('Agent %s: Loaded weights: %s' % (self._name, d))
+                        self._new_weights = True
+                    else:
+                        self._new_weights = False
+                    break
+            logging.info('Waiting for weights to become available.')
+            time.sleep(1)
+
+    def _get_task_name(self):
+        if hasattr(self._eval_env, '_task_class'):
+            eval_task_name = change_case(self._eval_env._task_class.__name__)
+            multi_task = False
+        elif hasattr(self._eval_env, '_task_classes'):
+            if self._eval_env.active_task_id != -1:
+                task_id = (self._eval_env.active_task_id) % len(self._eval_env._task_classes)
+                eval_task_name = change_case(self._eval_env._task_classes[task_id].__name__)
+            else:
+                eval_task_name = ''
+            multi_task = True
+        else:
+            raise Exception('Neither task_class nor task_classes found in eval env')
+        return eval_task_name, multi_task
+
+    def _run_eval_independent(self, name: str,
+                              stats_accumulator,
+                              weight,
+                              writer_lock,
+                              eval=True,
+                              device_idx=0,
+                              save_metrics=True,
+                              cinematic_recorder_cfg=None):
+
+        self._name = name
+        self._save_metrics = save_metrics
+        self._is_test_set = type(weight) == dict
+
+        self._agent = copy.deepcopy(self._agent)
+
+        device = torch.device('cuda:%d' % device_idx) if torch.cuda.device_count() > 1 else torch.device('cuda:0')
+        with writer_lock: # hack to prevent multiple CLIP downloads ... argh should use a separate lock
+            self._agent.build(training=False, device=device)
+
+        logging.info('%s: Launching env.' % name)
+        np.random.seed()
+
+        logging.info('Agent information:')
+        logging.info(self._agent)
+
+        env = self._eval_env
+        env.eval = eval
+        env.launch()
+
+        # initialize cinematic recorder if specified
+        rec_cfg = cinematic_recorder_cfg
+        if rec_cfg.enabled:
+            cam_placeholder = Dummy('cam_cinematic_placeholder')
+            cam = VisionSensor.create(rec_cfg.camera_resolution)
+            cam.set_pose(cam_placeholder.get_pose())
+            cam.set_parent(cam_placeholder)
+
+            cam_motion = CircleCameraMotion(cam, Dummy('cam_cinematic_base'), rec_cfg.rotate_speed)
+            tr = TaskRecorder(env, cam_motion, fps=rec_cfg.fps)
+
+            env.env._action_mode.arm_action_mode.set_callable_each_step(tr.take_snap)
+
+        if not os.path.exists(self._weightsdir):
+            raise Exception('No weights directory found.')
+
+        # to save or not to save evaluation metrics (set as False for recording videos)
+        if self._save_metrics:
+            csv_file = 'eval_data.csv' if not self._is_test_set else 'test_data.csv'
+            writer = LogWriter(self._logdir, True, True,
+                               env_csv=csv_file)
+
+        # one weight for all tasks (used for validation)
+        if type(weight) == int:
+            logging.info('Evaluating weight %s' % weight)
+            weight_path = os.path.join(self._weightsdir, str(weight))
+            seed_path = self._weightsdir.replace('/weights', '')
+            self._agent.load_weights(weight_path)
+            weight_name = str(weight)
+
+        new_transitions = {'train_envs': 0, 'eval_envs': 0}
+        total_transitions = {'train_envs': 0, 'eval_envs': 0}
+        current_task_id = -1
+
+        for n_eval in range(self._num_eval_runs):
+            if rec_cfg.enabled:
+                tr._cam_motion.save_pose()
+
+            # best weight for each task (used for test evaluation)
+            if type(weight) == dict:
+                task_name = list(weight.keys())[n_eval]
+                task_weight = weight[task_name]
+                weight_path = os.path.join(self._weightsdir, str(task_weight))
+                seed_path = self._weightsdir.replace('/weights', '')
+                self._agent.load_weights(weight_path)
+                weight_name = str(task_weight)
+                print('Evaluating weight %s for %s' % (weight_name, task_name))
+
+            # evaluate on N tasks * M episodes per task = total eval episodes
+            for ep in range(self._eval_episodes):
+                eval_demo_seed = ep + self._eval_from_eps_number
+                logging.info('%s: Starting episode %d, seed %d.' % (name, ep, eval_demo_seed))
+
+                # the current task gets reset after every M episodes
+                episode_rollout = []
+                generator = self._rollout_generator.generator(
+                    self._step_signal, env, self._agent,
+                    self._episode_length, self._timesteps,
+                    eval, eval_demo_seed=eval_demo_seed,
+                    record_enabled=rec_cfg.enabled)
+                try:
+                    for replay_transition in generator:
+                        while True:
+                            if self._kill_signal.value:
+                                env.shutdown()
+                                return
+                            if (eval or self._target_replay_ratio is None or
+                                    self._step_signal.value <= 0 or (
+                                            self._current_replay_ratio.value >
+                                            self._target_replay_ratio)):
+                                break
+                            time.sleep(1)
+                            logging.debug(
+                                'Agent. Waiting for replay_ratio %f to be more than %f' %
+                                (self._current_replay_ratio.value, self._target_replay_ratio))
+
+                        with self.write_lock:
+                            if len(self.agent_summaries) == 0:
+                                # Only store new summaries if the previous ones
+                                # have been popped by the main env runner.
+                                for s in self._agent.act_summaries():
+                                    self.agent_summaries.append(s)
+                        episode_rollout.append(replay_transition)
+                except StopIteration as e:
+                    continue
+                except Exception as e:
+                    env.shutdown()
+                    raise e
+
+                with self.write_lock:
+                    for transition in episode_rollout:
+                        self.stored_transitions.append((name, transition, eval))
+
+                        new_transitions['eval_envs'] += 1
+                        total_transitions['eval_envs'] += 1
+                        stats_accumulator.step(transition, eval)
+                        current_task_id = transition.info['active_task_id']
+
+                self._num_eval_episodes_signal.value += 1
+
+                task_name, _ = self._get_task_name()
+                reward = episode_rollout[-1].reward
+                lang_goal = env._lang_goal
+                print(f"Evaluating {task_name} | Episode {ep} | Score: {reward} | Lang Goal: {lang_goal}")
+
+                # save recording
+                if rec_cfg.enabled:
+                    success = reward > 0.99
+                    record_file = os.path.join(seed_path, 'videos',
+                                               '%s_w%s_s%s_%s.mp4' % (task_name,
+                                                                      weight_name,
+                                                                      eval_demo_seed,
+                                                                      'succ' if success else 'fail'))
+
+                    lang_goal = self._eval_env._lang_goal
+
+                    tr.save(record_file, lang_goal, reward)
+                    tr._cam_motion.restore_pose()
+
+            # report summaries
+            summaries = []
+            summaries.extend(stats_accumulator.pop())
+
+            eval_task_name, multi_task = self._get_task_name()
+
+            if eval_task_name and multi_task:
+                for s in summaries:
+                    if 'eval' in s.name:
+                        s.name = '%s/%s' % (s.name, eval_task_name)
+
+            if len(summaries) > 0:
+                if multi_task:
+                    task_score = [s.value for s in summaries if f'eval_envs/return/{eval_task_name}' in s.name][0]
+                else:
+                    task_score = [s.value for s in summaries if f'eval_envs/return' in s.name][0]
+            else:
+                task_score = "unknown"
+
+            print(f"Finished {eval_task_name} | Final Score: {task_score}\n")
+
+            if self._save_metrics:
+                with writer_lock:
+                    writer.add_summaries(weight_name, summaries)
+
+            self._new_transitions = {'train_envs': 0, 'eval_envs': 0}
+            self.agent_summaries[:] = []
+            self.stored_transitions[:] = []
+
+        if self._save_metrics:
+            with writer_lock:
+                writer.end_iteration()
+
+        logging.info('Finished evaluation.')
+        env.shutdown()
+
+    def kill(self):
+        self._kill_signal.value = True

external/yarr/yarr/runners/env_runner.py

@@ -0,0 +1,224 @@
+import collections
+import logging
+import os
+import signal
+import time
+from multiprocessing import Value
+from threading import Thread
+from typing import List
+from typing import Union
+
+import numpy as np
+import torch
+from yarr.agents.agent import Agent
+from yarr.agents.agent import ScalarSummary
+from yarr.agents.agent import Summary
+from yarr.envs.env import Env
+from yarr.replay_buffer.replay_buffer import ReplayBuffer
+from yarr.runners._env_runner import _EnvRunner
+from yarr.utils.rollout_generator import RolloutGenerator
+from yarr.utils.stat_accumulator import StatAccumulator, SimpleAccumulator
+from yarr.utils.process_str import change_case
+from helpers.custom_rlbench_env import CustomRLBenchEnv, CustomMultiTaskRLBenchEnv
+
+class EnvRunner(object):
+
+    def __init__(self,
+                 train_env: Env,
+                 agent: Agent,
+                 train_replay_buffer: Union[ReplayBuffer, List[ReplayBuffer]],
+                 num_train_envs: int,
+                 num_eval_envs: int,
+                 rollout_episodes: int,
+                 eval_episodes: int,
+                 training_iterations: int,
+                 eval_from_eps_number: int,
+                 episode_length: int,
+                 eval_env: Union[Env, None] = None,
+                 eval_replay_buffer: Union[ReplayBuffer, List[ReplayBuffer], None] = None,
+                 stat_accumulator: Union[StatAccumulator, None] = None,
+                 rollout_generator: RolloutGenerator = None,
+                 weightsdir: str = None,
+                 logdir: str = None,
+                 max_fails: int = 10,
+                 num_eval_runs: int = 1,
+                 env_device: torch.device = None,
+                 multi_task: bool = False):
+        self._train_env = train_env
+        self._eval_env = eval_env if eval_env else train_env
+        self._agent = agent
+        self._train_envs = num_train_envs
+        self._eval_envs = num_eval_envs
+        self._train_replay_buffer = train_replay_buffer if isinstance(train_replay_buffer, list) else [train_replay_buffer]
+        self._timesteps = self._train_replay_buffer[0].timesteps if self._train_replay_buffer[0] is not None else 1
+
+        if eval_replay_buffer is not None:
+            eval_replay_buffer = eval_replay_buffer if isinstance(eval_replay_buffer, list) else [eval_replay_buffer]
+        self._eval_replay_buffer = eval_replay_buffer
+        self._rollout_episodes = rollout_episodes
+        self._eval_episodes = eval_episodes
+        self._num_eval_runs = num_eval_runs
+        self._training_iterations = training_iterations
+        self._eval_from_eps_number = eval_from_eps_number
+        self._episode_length = episode_length
+        self._stat_accumulator = stat_accumulator
+        self._rollout_generator = (
+            RolloutGenerator() if rollout_generator is None
+            else rollout_generator)
+        self._rollout_generator._env_device = env_device
+        self._weightsdir = weightsdir
+        self._logdir = logdir
+        self._max_fails = max_fails
+        self._env_device = env_device
+        self._previous_loaded_weight_folder = ''
+        self._p = None
+        self._kill_signal = Value('b', 0)
+        self._step_signal = Value('i', -1)
+        self._num_eval_episodes_signal = Value('i', 0)
+        self._eval_epochs_signal = Value('i', 0)
+        self._eval_report_signal = Value('b', 0)
+        self._new_transitions = {'train_envs': 0, 'eval_envs': 0}
+        self._total_transitions = {'train_envs': 0, 'eval_envs': 0}
+        self.log_freq = 1000  # Will get overridden later
+        self.target_replay_ratio = None  # Will get overridden later
+        self.current_replay_ratio = Value('f', -1)
+        self._current_task_id = -1
+        self._multi_task = multi_task
+
+    def summaries(self) -> List[Summary]:
+        summaries = []
+        if self._stat_accumulator is not None:
+            summaries.extend(self._stat_accumulator.pop())
+        for key, value in self._new_transitions.items():
+            summaries.append(ScalarSummary('%s/new_transitions' % key, value))
+        for key, value in self._total_transitions.items():
+            summaries.append(ScalarSummary('%s/total_transitions' % key, value))
+        self._new_transitions = {'train_envs': 0, 'eval_envs': 0}
+        summaries.extend(self._agent_summaries)
+
+        # add current task_name to eval summaries .... argh this should be inside a helper function
+        if hasattr(self._eval_env, '_task_class'):
+            eval_task_name = change_case(self._eval_env._task_class.__name__)
+        elif hasattr(self._eval_env, '_task_classes'):
+            if self._current_task_id != -1:
+                task_id = (self._current_task_id) % len(self._eval_env._task_classes)
+                eval_task_name = change_case(self._eval_env._task_classes[task_id].__name__)
+            else:
+                eval_task_name = ''
+        else:
+            raise Exception('Neither task_class nor task_classes found in eval env')
+
+        # multi-task summaries
+        if eval_task_name and self._multi_task:
+            for s in summaries:
+                if 'eval' in s.name:
+                    s.name = '%s/%s' % (s.name, eval_task_name)
+
+        return summaries
+
+    def _update(self):
+        # Move the stored transitions to the replay and accumulate statistics.
+        new_transitions = collections.defaultdict(int)
+        with self._internal_env_runner.write_lock:
+            self._agent_summaries = list(
+                self._internal_env_runner.agent_summaries)
+            if self._num_eval_episodes_signal.value % self._eval_episodes == 0 and self._num_eval_episodes_signal.value > 0:
+                self._internal_env_runner.agent_summaries[:] = []
+            for name, transition, eval in self._internal_env_runner.stored_transitions:
+                add_to_buffer = (not eval) or self._eval_replay_buffer is not None
+                if add_to_buffer:
+                    kwargs = dict(transition.observation)
+                    replay_index = transition.info["active_task_id"]
+                    rb = self._eval_replay_buffer[replay_index] if eval else self._train_replay_buffer[replay_index]
+                    rb.add(
+                        np.array(transition.action), transition.reward,
+                        transition.terminal,
+                        transition.timeout, **kwargs)
+                    if transition.terminal:
+                        rb.add_final(
+                            **transition.final_observation)
+                new_transitions[name] += 1
+                self._new_transitions[
+                    'eval_envs' if eval else 'train_envs'] += 1
+                self._total_transitions[
+                    'eval_envs' if eval else 'train_envs'] += 1
+                if self._stat_accumulator is not None:
+                    self._stat_accumulator.step(transition, eval)
+                self._current_task_id = transition.info["active_task_id"] if eval else -1
+            self._internal_env_runner.stored_transitions[:] = []  # Clear list
+        return new_transitions
+
+    def _run(self, save_load_lock):
+        self._internal_env_runner = _EnvRunner(
+            self._train_env, self._eval_env, self._agent, self._timesteps, self._train_envs,
+            self._eval_envs, self._rollout_episodes, self._eval_episodes,
+            self._training_iterations, self._eval_from_eps_number, self._episode_length, self._kill_signal,
+            self._step_signal, self._num_eval_episodes_signal,
+            self._eval_epochs_signal, self._eval_report_signal,
+            self.log_freq, self._rollout_generator, save_load_lock,
+            self.current_replay_ratio, self.target_replay_ratio,
+            self._weightsdir, self._logdir,
+            self._env_device, self._previous_loaded_weight_folder,
+            num_eval_runs=self._num_eval_runs)
+        training_envs = self._internal_env_runner.spin_up_envs('train_env', self._train_envs, False)
+        eval_envs = self._internal_env_runner.spin_up_envs('eval_env', self._eval_envs, True)
+        envs = training_envs + eval_envs
+        no_transitions = {env.name: 0 for env in envs}
+        while True:
+            for p in envs:
+                if p.exitcode is not None:
+                    envs.remove(p)
+                    if p.exitcode != 0:
+                        self._internal_env_runner.p_failures[p.name] += 1
+                        n_failures = self._internal_env_runner.p_failures[p.name]
+                        if n_failures > self._max_fails:
+                            logging.error('Env %s failed too many times (%d times > %d)' %
+                                          (p.name, n_failures, self._max_fails))
+                            raise RuntimeError('Too many process failures.')
+                        logging.warning('Env %s failed (%d times <= %d). restarting' %
+                                        (p.name, n_failures, self._max_fails))
+                        p = self._internal_env_runner.restart_process(p.name)
+                        envs.append(p)
+
+            if not self._kill_signal.value:
+                new_transitions = self._update()
+                for p in envs:
+                    if new_transitions[p.name] == 0:
+                        no_transitions[p.name] += 1
+                    else:
+                        no_transitions[p.name] = 0
+                    if no_transitions[p.name] > 1200: #600:  # 10min
+                        logging.warning("Env %s hangs, so restarting" % p.name)
+                        envs.remove(p)
+                        os.kill(p.pid, signal.SIGTERM)
+                        p = self._internal_env_runner.restart_process(p.name)
+                        envs.append(p)
+                        no_transitions[p.name] = 0
+
+            if len(envs) == 0:
+                break
+            time.sleep(1)
+
+    def start(self, save_load_lock):
+        self._p = Thread(target=self._run, args=(save_load_lock,), daemon=True)
+        self._p.name = 'EnvRunnerThread'
+        self._p.start()
+
+    def wait(self):
+        if self._p.is_alive():
+            self._p.join()
+
+    def stop(self):
+        if self._p.is_alive():
+            self._kill_signal.value = True
+            self._p.join()
+
+    def set_step(self, step):
+        self._step_signal.value = step
+
+    def set_eval_report(self, report):
+        self._eval_report_signal.value = report
+
+    def set_eval_epochs(self, epochs):
+        self._eval_epochs_signal.value = epochs
+

external/yarr/yarr/runners/independent_env_runner.py

@@ -0,0 +1,130 @@
+import numpy as np
+import torch
+from typing import List
+from typing import Union
+
+from yarr.agents.agent import Agent
+from yarr.envs.env import Env
+from yarr.replay_buffer.replay_buffer import ReplayBuffer
+from yarr.runners._independent_env_runner import _IndependentEnvRunner
+from yarr.utils.rollout_generator import RolloutGenerator
+from yarr.utils.stat_accumulator import StatAccumulator, SimpleAccumulator
+from yarr.agents.agent import Summary
+from helpers.custom_rlbench_env import CustomRLBenchEnv, CustomMultiTaskRLBenchEnv
+
+from yarr.runners.env_runner import EnvRunner
+
+
+class IndependentEnvRunner(EnvRunner):
+
+    def __init__(self,
+                 train_env: Env,
+                 agent: Agent,
+                 train_replay_buffer: Union[ReplayBuffer, List[ReplayBuffer]],
+                 num_train_envs: int,
+                 num_eval_envs: int,
+                 rollout_episodes: int,
+                 eval_episodes: int,
+                 training_iterations: int,
+                 eval_from_eps_number: int,
+                 episode_length: int,
+                 eval_env: Union[Env, None] = None,
+                 eval_replay_buffer: Union[ReplayBuffer, List[ReplayBuffer], None] = None,
+                 stat_accumulator: Union[StatAccumulator, None] = None,
+                 rollout_generator: RolloutGenerator = None,
+                 weightsdir: str = None,
+                 logdir: str = None,
+                 max_fails: int = 10,
+                 num_eval_runs: int = 1,
+                 env_device: torch.device = None,
+                 multi_task: bool = False):
+            super().__init__(train_env, agent, train_replay_buffer, num_train_envs, num_eval_envs,
+                            rollout_episodes, eval_episodes, training_iterations, eval_from_eps_number,
+                            episode_length, eval_env, eval_replay_buffer, stat_accumulator,
+                            rollout_generator, weightsdir, logdir, max_fails, num_eval_runs,
+                            env_device, multi_task)
+
+    def summaries(self) -> List[Summary]:
+        summaries = []
+        if self._stat_accumulator is not None:
+            summaries.extend(self._stat_accumulator.pop())
+        self._new_transitions = {'train_envs': 0, 'eval_envs': 0}
+        summaries.extend(self._agent_summaries)
+
+        # add current task_name to eval summaries .... argh this should be inside a helper function
+        if hasattr(self._eval_env, '_task_class'):
+            eval_task_name = change_case(self._eval_env._task_class.__name__)
+        elif hasattr(self._eval_env, '_task_classes'):
+            if self._current_task_id != -1:
+                task_id = (self._current_task_id) % len(self._eval_env._task_classes)
+                eval_task_name = change_case(self._eval_env._task_classes[task_id].__name__)
+            else:
+                eval_task_name = ''
+        else:
+            raise Exception('Neither task_class nor task_classes found in eval env')
+
+        # multi-task summaries
+        if eval_task_name and self._multi_task:
+            for s in summaries:
+                if 'eval' in s.name:
+                    s.name = '%s/%s' % (s.name, eval_task_name)
+
+        return summaries
+
+    # serialized evaluator for individual tasks
+    def start(self, weight,
+              save_load_lock, writer_lock,
+              env_config,
+              device_idx,
+              save_metrics,
+              cinematic_recorder_cfg):
+
+        if hasattr(self, "_on_thread_start"):
+            self._on_thread_start()
+            
+        multi_task = isinstance(env_config[0], list)
+        if multi_task:
+            eval_env = CustomMultiTaskRLBenchEnv(
+                task_classes=env_config[0],
+                observation_config=env_config[1],
+                action_mode=env_config[2],
+                dataset_root=env_config[3],
+                episode_length=env_config[4],
+                headless=env_config[5],
+                swap_task_every=env_config[6],
+                include_lang_goal_in_obs=env_config[7],
+                time_in_state=env_config[8],
+                record_every_n=env_config[9])
+        else:
+            eval_env = CustomRLBenchEnv(
+                task_class=env_config[0],
+                observation_config=env_config[1],
+                action_mode=env_config[2],
+                dataset_root=env_config[3],
+                episode_length=env_config[4],
+                headless=env_config[5],
+                include_lang_goal_in_obs=env_config[6],
+                time_in_state=env_config[7],
+                record_every_n=env_config[8])
+
+        self._internal_env_runner = _IndependentEnvRunner(
+            self._train_env, eval_env, self._agent, self._timesteps, self._train_envs,
+            self._eval_envs, self._rollout_episodes, self._eval_episodes,
+            self._training_iterations, self._eval_from_eps_number, self._episode_length, self._kill_signal,
+            self._step_signal, self._num_eval_episodes_signal,
+            self._eval_epochs_signal, self._eval_report_signal,
+            self.log_freq, self._rollout_generator, None,
+            self.current_replay_ratio, self.target_replay_ratio,
+            self._weightsdir, self._logdir,
+            self._env_device, self._previous_loaded_weight_folder,
+            num_eval_runs=self._num_eval_runs)
+
+        stat_accumulator = SimpleAccumulator(eval_video_fps=30)
+        self._internal_env_runner._run_eval_independent('eval_env',
+                                                        stat_accumulator,
+                                                        weight,
+                                                        writer_lock,
+                                                        True,
+                                                        device_idx,
+                                                        save_metrics,
+                                                        cinematic_recorder_cfg)

external/yarr/yarr/runners/offline_train_runner.py

@@ -0,0 +1,163 @@
+import copy
+import logging
+import os
+import shutil
+import time
+from typing import List
+from typing import Union
+
+import psutil
+import torch
+import pandas as pd
+from yarr.agents.agent import Agent
+from yarr.replay_buffer.wrappers.pytorch_replay_buffer import \
+    PyTorchReplayBuffer
+from yarr.utils.log_writer import LogWriter
+from yarr.utils.stat_accumulator import StatAccumulator
+
+
+class OfflineTrainRunner():
+
+    def __init__(self,
+                 agent: Agent,
+                 wrapped_replay_buffer: PyTorchReplayBuffer,
+                 train_device: torch.device,
+                 stat_accumulator: Union[StatAccumulator, None] = None,
+                 iterations: int = int(6e6),
+                 logdir: str = '/tmp/yarr/logs',
+                 logging_level: int = logging.INFO,
+                 log_freq: int = 10,
+                 weightsdir: str = '/tmp/yarr/weights',
+                 num_weights_to_keep: int = 60,
+                 save_freq: int = 100,
+                 tensorboard_logging: bool = True,
+                 csv_logging: bool = False,
+                 load_existing_weights: bool = True,
+                 rank: int = None,
+                 world_size: int = None):
+        self._agent = agent
+        self._wrapped_buffer = wrapped_replay_buffer
+        self._stat_accumulator = stat_accumulator
+        self._iterations = iterations
+        self._logdir = logdir
+        self._logging_level = logging_level
+        self._log_freq = log_freq
+        self._weightsdir = weightsdir
+        self._num_weights_to_keep = num_weights_to_keep
+        self._save_freq = save_freq
+
+        self._wrapped_buffer = wrapped_replay_buffer
+        self._train_device = train_device
+        self._tensorboard_logging = tensorboard_logging
+        self._csv_logging = csv_logging
+        self._load_existing_weights = load_existing_weights
+        self._rank = rank
+        self._world_size = world_size
+
+        self._writer = None
+        if logdir is None:
+            logging.info("'logdir' was None. No logging will take place.")
+        else:
+            self._writer = LogWriter(
+                self._logdir, tensorboard_logging, csv_logging)
+
+        if weightsdir is None:
+            logging.info(
+                "'weightsdir' was None. No weight saving will take place.")
+        else:
+            os.makedirs(self._weightsdir, exist_ok=True)
+
+    def _save_model(self, i):
+        d = os.path.join(self._weightsdir, str(i))
+        os.makedirs(d, exist_ok=True)
+        self._agent.save_weights(d)
+
+        # remove oldest save
+        prev_dir = os.path.join(self._weightsdir, str(
+            i - self._save_freq * self._num_weights_to_keep))
+        if os.path.exists(prev_dir):
+            shutil.rmtree(prev_dir)
+
+    def _step(self, i, sampled_batch):
+        update_dict = self._agent.update(i, sampled_batch)
+        total_losses = update_dict['total_losses']
+        return total_losses
+
+    def _get_resume_eval_epoch(self):
+        starting_epoch = 0
+        eval_csv_file = self._weightsdir.replace('weights', 'eval_data.csv') # TODO(mohit): check if it's supposed be 'env_data.csv'
+        if os.path.exists(eval_csv_file):
+             eval_dict = pd.read_csv(eval_csv_file).to_dict()
+             epochs = list(eval_dict['step'].values())
+             return epochs[-1] if len(epochs) > 0 else starting_epoch
+        else:
+            return starting_epoch
+
+    def start(self):
+
+        if hasattr(self, "_on_thread_start"):
+            self._on_thread_start()
+        else:
+            logging.getLogger().setLevel(self._logging_level)
+         
+
+        self._agent = copy.deepcopy(self._agent)
+        self._agent.build(training=True, device=self._train_device)
+
+        if self._weightsdir is not None:
+            existing_weights = sorted([int(f) for f in os.listdir(self._weightsdir)])
+            if (not self._load_existing_weights) or len(existing_weights) == 0:
+                # self._save_model(0)
+                start_iter = 0
+            else:
+                resume_iteration = existing_weights[-1]
+                self._agent.load_weights(os.path.join(self._weightsdir, str(resume_iteration)))
+                start_iter = resume_iteration + 1
+                if self._rank == 0:
+                    logging.info(f"Resuming training from iteration {resume_iteration} ...")
+
+        dataset = self._wrapped_buffer.dataset()
+        data_iter = iter(dataset)
+
+        process = psutil.Process(os.getpid())
+        num_cpu = psutil.cpu_count()
+
+        for i in range(start_iter, self._iterations):
+            log_iteration = i % self._log_freq == 0 and i > 0
+
+            if log_iteration:
+                process.cpu_percent(interval=None)
+
+            t = time.time()
+            sampled_batch = next(data_iter)
+            sample_time = time.time() - t
+
+            batch = {k: v.to(self._train_device) for k, v in sampled_batch.items() if type(v) == torch.Tensor}
+            t = time.time()
+            loss = self._step(i, batch)
+            step_time = time.time() - t
+
+            if self._rank == 0:
+                if log_iteration and self._writer is not None:
+                    agent_summaries = self._agent.update_summaries()
+                    self._writer.add_summaries(i, agent_summaries)
+
+                    self._writer.add_scalar(
+                        i, 'monitoring/memory_gb',
+                        process.memory_info().rss * 1e-9)
+                    self._writer.add_scalar(
+                        i, 'monitoring/cpu_percent',
+                        process.cpu_percent(interval=None) / num_cpu)
+
+                    logging.info(f"Train Step {i:06d} | Loss: {loss:0.5f} | Sample time: {sample_time:0.6f} | Step time: {step_time:0.4f}.")
+
+                self._writer.end_iteration()
+
+                if i % self._save_freq == 0 and self._weightsdir is not None:
+                    self._save_model(i)
+
+        if self._rank == 0 and self._writer is not None:
+            self._writer.close()
+            logging.info('Stopping envs ...')
+
+            self._wrapped_buffer.replay_buffer.shutdown()

external/yarr/yarr/runners/pytorch_train_runner.py

@@ -0,0 +1,308 @@
+import copy
+import logging
+import os
+import shutil
+import signal
+import sys
+import threading
+import time
+from multiprocessing import Lock
+from typing import Optional, List
+from typing import Union
+
+import gc
+import numpy as np
+import psutil
+import torch
+import pandas as pd
+from yarr.agents.agent import Agent
+from yarr.replay_buffer.wrappers.pytorch_replay_buffer import \
+    PyTorchReplayBuffer
+from yarr.runners.env_runner import EnvRunner
+from yarr.runners.train_runner import TrainRunner
+from yarr.utils.log_writer import LogWriter
+from yarr.utils.stat_accumulator import StatAccumulator
+from yarr.replay_buffer.prioritized_replay_buffer import PrioritizedReplayBuffer
+
+NUM_WEIGHTS_TO_KEEP = 60
+
+
+class PyTorchTrainRunner(TrainRunner):
+
+    def __init__(self,
+                 agent: Agent,
+                 env_runner: EnvRunner,
+                 wrapped_replay_buffer: Union[
+                     PyTorchReplayBuffer, List[PyTorchReplayBuffer]],
+                 train_device: torch.device,
+                 replay_buffer_sample_rates: List[float] = None,
+                 stat_accumulator: Union[StatAccumulator, None] = None,
+                 iterations: int = int(1e6),
+                 num_train_envs: int = 1,
+                 num_eval_envs: int = 1,
+                 eval_episodes: int = 10,
+                 logdir: str = '/tmp/yarr/logs',
+                 log_freq: int = 10,
+                 transitions_before_train: int = 1000,
+                 weightsdir: str = '/tmp/yarr/weights',
+                 save_freq: int = 100,
+                 replay_ratio: Optional[float] = None,
+                 tensorboard_logging: bool = True,
+                 csv_logging: bool = False,
+                 buffers_per_batch: int = -1,  # -1 = all
+                 load_existing_weights: bool = True):
+        super(PyTorchTrainRunner, self).__init__(
+            agent, env_runner, wrapped_replay_buffer,
+            stat_accumulator,
+            iterations, logdir, log_freq, transitions_before_train, weightsdir,
+            save_freq)
+
+        env_runner.log_freq = log_freq
+        env_runner.target_replay_ratio = replay_ratio
+        self._wrapped_buffer = wrapped_replay_buffer if isinstance(
+            wrapped_replay_buffer, list) else [wrapped_replay_buffer]
+        self._replay_buffer_sample_rates = (
+            [1.0] if replay_buffer_sample_rates is None else
+            replay_buffer_sample_rates)
+        if len(self._replay_buffer_sample_rates) != len(wrapped_replay_buffer):
+            logging.warning(
+                'Numbers of replay buffers differs from sampling rates. Setting as uniform sampling.')
+            self._replay_buffer_sample_rates = [1.0 / len(self._wrapped_buffer)] * len(self._wrapped_buffer)
+        if sum(self._replay_buffer_sample_rates) != 1:
+            raise ValueError('Sum of sampling rates should be 1.')
+
+        self._train_device = train_device
+        self._tensorboard_logging = tensorboard_logging
+        self._csv_logging = csv_logging
+        self._num_train_envs = num_train_envs
+        self._num_eval_envs = num_eval_envs
+        self._eval_episodes = eval_episodes
+        self._load_existing_weights = load_existing_weights
+
+        if replay_ratio is not None and replay_ratio < 0:
+            raise ValueError("max_replay_ratio must be positive.")
+        self._target_replay_ratio = replay_ratio
+
+        self._writer = None
+        if logdir is None:
+            logging.info("'logdir' was None. No logging will take place.")
+        else:
+            self._writer = LogWriter(
+                self._logdir, tensorboard_logging, csv_logging)
+        if weightsdir is None:
+            logging.info(
+                "'weightsdir' was None. No weight saving will take place.")
+        else:
+            os.makedirs(self._weightsdir, exist_ok=True)
+        self._buffers_per_batch = buffers_per_batch if buffers_per_batch > 0 else len(wrapped_replay_buffer)
+
+    def _save_model(self, i):
+        with self._save_load_lock:
+            d = os.path.join(self._weightsdir, str(i))
+            os.makedirs(d, exist_ok=True)
+            self._agent.save_weights(d)
+            # Remove oldest save
+            prev_dir = os.path.join(self._weightsdir, str(
+                i - self._save_freq * NUM_WEIGHTS_TO_KEEP))
+            if os.path.exists(prev_dir):
+                shutil.rmtree(prev_dir)
+
+    def _step(self, i, sampled_batch):
+        update_dict = self._agent.update(i, sampled_batch)
+        if "priority" in update_dict:
+            priority = update_dict['priority'].cpu().detach().numpy() if isinstance(update_dict['priority'], torch.Tensor) else np.numpy(update_dict['priority'])
+        else:
+            priority = None
+        indices = sampled_batch['indices'].cpu().detach().numpy()
+        acc_bs = 0
+        for wb_idx, wb in enumerate(self._wrapped_buffer):
+            bs = wb.replay_buffer.batch_size
+            if 'priority' in update_dict:
+                indices_ = indices[:, wb_idx]
+                if hasattr(wb, "replay_buffer"):
+                    if len(priority.shape) > 1:
+                        priority_ = priority[:, wb_idx]
+                    else:
+                        # legacy version
+                        priority_ = priority[acc_bs: acc_bs + bs]
+                    wb.replay_buffer.set_priority(indices_, priority_)
+            acc_bs += bs
+
+    def _signal_handler(self, sig, frame):
+        if threading.current_thread().name != 'MainThread':
+            return
+        logging.info('SIGINT captured. Shutting down.'
+                     'This may take a few seconds.')
+        self._env_runner.stop()
+        [r.replay_buffer.shutdown() for r in self._wrapped_buffer]
+        sys.exit(0)
+
+    def _get_add_counts(self):
+        return np.array([
+            r.replay_buffer.add_count for r in self._wrapped_buffer])
+
+    def _get_sum_add_counts(self):
+        return sum([
+            r.replay_buffer.add_count for r in self._wrapped_buffer])
+
+    def _get_resume_eval_epoch(self):
+        starting_epoch = 0
+        eval_csv_file = self._weightsdir.replace('weights', 'eval_data.csv') # TODO(mohit): check if it's supposed be 'env_data.csv'
+        if os.path.exists(eval_csv_file):
+             eval_dict = pd.read_csv(eval_csv_file).to_dict()
+             epochs = list(eval_dict['step'].values())
+             return epochs[-1] if len(epochs) > 0 else starting_epoch
+        else:
+            return starting_epoch
+
+    def start(self):
+
+        signal.signal(signal.SIGINT, self._signal_handler)
+
+        self._save_load_lock = Lock()
+
+        # Kick off the environments
+        self._env_runner.start(self._save_load_lock)
+
+        self._agent = copy.deepcopy(self._agent)
+        self._agent.build(training=True, device=self._train_device)
+
+        if self._weightsdir is not None:
+            existing_weights = sorted([int(f) for f in os.listdir(self._weightsdir)])
+            if (not self._load_existing_weights) or len(existing_weights) == 0:
+                self._save_model(0)
+                start_iter = 0
+            else:
+                resume_iteration = existing_weights[-1]
+                self._agent.load_weights(os.path.join(self._weightsdir, str(resume_iteration)))
+                start_iter = resume_iteration + 1
+                print(f"Resuming training from iteration {resume_iteration} ...")
+
+                if self._num_eval_envs > 0:
+                    eval_epoch = self._get_resume_eval_epoch()
+                    self._env_runner.set_eval_epochs(eval_epoch)
+                    self._writer.set_resumed_from_prev_run(True)
+                    print(f"Resuming evaluation from epoch {eval_epoch} ...")
+
+        while (np.any(self._get_add_counts() < self._transitions_before_train)):
+            time.sleep(1)
+            logging.info(
+                'Waiting for %d samples before training. Currently have %s.' %
+                (self._transitions_before_train, str(self._get_add_counts())))
+
+        datasets = [r.dataset() for r in self._wrapped_buffer]
+        data_iter = [iter(d) for d in datasets]
+
+        init_replay_size = self._get_sum_add_counts().astype(float)
+        batch_times_buffers_per_sample = sum([
+            r.replay_buffer.batch_size for r in self._wrapped_buffer[:self._buffers_per_batch]])
+        process = psutil.Process(os.getpid())
+        num_cpu = psutil.cpu_count()
+
+        for i in range(start_iter, self._iterations):
+            self._env_runner.set_step(i)
+
+            if self._num_train_envs > 0 or self._num_eval_envs == 0:
+                log_iteration = i % self._log_freq == 0 and i > 0
+            else:
+                num_eval_episodes = self._env_runner._num_eval_episodes_signal.value
+                log_iteration = self._env_runner._eval_report_signal.value and num_eval_episodes > 0
+
+            if log_iteration:
+                process.cpu_percent(interval=None)
+
+            def get_replay_ratio():
+                size_used = batch_times_buffers_per_sample * i
+                size_added = (
+                    self._get_sum_add_counts()
+                    - init_replay_size
+                )
+                replay_ratio = size_used / (size_added + 1e-6)
+                return replay_ratio
+
+            if self._target_replay_ratio is not None:
+                # wait for env_runner collecting enough samples
+                while True:
+                    replay_ratio = get_replay_ratio()
+                    self._env_runner.current_replay_ratio.value = replay_ratio
+                    if replay_ratio < self._target_replay_ratio:
+                        break
+                    time.sleep(1)
+                    logging.debug(
+                        'Waiting for replay_ratio %f to be less than %f.' %
+                        (replay_ratio, self._target_replay_ratio))
+                del replay_ratio
+
+            t = time.time()
+
+            sampled_task_ids = np.random.choice(
+                range(len(datasets)), self._buffers_per_batch, replace=False)
+            sampled_batch = [next(data_iter[j]) for j in sampled_task_ids]
+            result = {}
+            for key in sampled_batch[0]:
+                result[key] = torch.stack([d[key] for d in sampled_batch], 1)
+            sampled_batch = result
+            sample_time = time.time() - t
+
+            batch = {k: v.to(self._train_device) for k, v in sampled_batch.items()}
+            t = time.time()
+            self._step(i, batch)
+            step_time = time.time() - t
+
+            if log_iteration and self._writer is not None:
+                replay_ratio = get_replay_ratio()
+                logging.info('Train Step %d. Eval Epoch %d. Sample time: %s. Step time: %s. Replay ratio: %s.' % (
+                    i, self._env_runner._eval_epochs_signal.value, sample_time, step_time, replay_ratio))
+                agent_summaries = self._agent.update_summaries()
+                env_summaries = self._env_runner.summaries()
+
+                # agent summaries
+                self._writer.add_summaries(i, agent_summaries)
+
+                # env summaries
+                self._writer.add_summaries(self._env_runner._eval_epochs_signal.value, env_summaries)
+
+                for r_i, wrapped_buffer in enumerate(self._wrapped_buffer):
+                    self._writer.add_scalar(
+                        i, 'replay%d/add_count' % r_i,
+                        wrapped_buffer.replay_buffer.add_count)
+                    self._writer.add_scalar(
+                        i, 'replay%d/size' % r_i,
+                        wrapped_buffer.replay_buffer.replay_capacity
+                        if wrapped_buffer.replay_buffer.is_full()
+                        else wrapped_buffer.replay_buffer.add_count)
+
+                self._writer.add_scalar(
+                    i, 'replay/replay_ratio', replay_ratio)
+                self._writer.add_scalar(
+                    i, 'replay/update_to_insert_ratio',
+                    float(i) / float(
+                        self._get_sum_add_counts() -
+                        init_replay_size + 1e-6))
+
+                self._writer.add_scalar(
+                    i, 'monitoring/sample_time_per_item',
+                    sample_time / batch_times_buffers_per_sample)
+                self._writer.add_scalar(
+                    i, 'monitoring/train_time_per_item',
+                    step_time / batch_times_buffers_per_sample)
+                self._writer.add_scalar(
+                    i, 'monitoring/memory_gb',
+                    process.memory_info().rss * 1e-9)
+                self._writer.add_scalar(
+                    i, 'monitoring/cpu_percent',
+                    process.cpu_percent(interval=None) / num_cpu)
+
+                self._env_runner.set_eval_report(False)
+
+            self._writer.end_iteration()
+
+            if i % self._save_freq == 0 and self._weightsdir is not None:
+                self._save_model(i)
+
+        if self._writer is not None:
+            self._writer.close()
+
+        logging.info('Stopping envs ...')
+        self._env_runner.stop()
+        [r.replay_buffer.shutdown() for r in self._wrapped_buffer]

external/yarr/yarr/runners/train_runner.py

@@ -0,0 +1,37 @@
+from abc import abstractmethod, ABC
+from typing import Union, List
+
+from yarr.agents.agent import Agent
+from yarr.replay_buffer.wrappers import WrappedReplayBuffer
+from yarr.runners.env_runner import EnvRunner
+from yarr.utils.stat_accumulator import StatAccumulator
+
+
+class TrainRunner(ABC):
+
+    def __init__(self,
+                 agent: Agent,
+                 env_runner: EnvRunner,
+                 wrapped_replay_buffer: WrappedReplayBuffer,
+                 stat_accumulator: Union[StatAccumulator, None] = None,
+                 iterations: int = int(1e6),
+                 logdir: str = '/tmp/yarr/logs',
+                 log_freq: int = 500,
+                 transitions_before_train: int = 1000,
+                 weightsdir: str = '/tmp/yarr/weights',
+                 save_freq: int = 100,
+                 ):
+        self._agent = agent
+        self._env_runner = env_runner
+        self._wrapped_buffer = wrapped_replay_buffer
+        self._stat_accumulator = stat_accumulator
+        self._iterations = iterations
+        self._logdir = logdir
+        self._log_freq = log_freq
+        self._transitions_before_train = transitions_before_train
+        self._weightsdir = weightsdir
+        self._save_freq = save_freq
+
+    @abstractmethod
+    def start(self):
+        pass

external/yarr/yarr/utils/log_writer.py

@@ -0,0 +1,128 @@
+import csv
+import logging
+import os
+from collections import OrderedDict
+
+import numpy as np
+import torch
+from yarr.agents.agent import ScalarSummary, HistogramSummary, ImageSummary, \
+    VideoSummary, TextSummary
+from torch.utils.tensorboard import SummaryWriter
+
+
+class LogWriter(object):
+
+    def __init__(self,
+                 logdir: str,
+                 tensorboard_logging: bool,
+                 csv_logging: bool,
+                 train_csv: str = 'train_data.csv',
+                 env_csv: str = 'env_data.csv'):
+        self._tensorboard_logging = tensorboard_logging
+        self._csv_logging = csv_logging
+        os.makedirs(logdir, exist_ok=True)
+        if tensorboard_logging:
+            self._tf_writer = SummaryWriter(logdir)
+        if csv_logging:
+            self._train_prev_row_data = self._train_row_data = OrderedDict()
+            self._train_csv_file = os.path.join(logdir, train_csv)
+            self._env_prev_row_data = self._env_row_data = OrderedDict()
+            self._env_csv_file = os.path.join(logdir, env_csv)
+            self._train_field_names = None
+            self._env_field_names = None
+
+    def add_scalar(self, i, name, value):
+        if self._tensorboard_logging:
+            self._tf_writer.add_scalar(name, value, i)
+        if self._csv_logging:
+            if 'env' in name or 'eval' in name or 'test' in name:
+                if len(self._env_row_data) == 0:
+                    self._env_row_data['step'] = i
+                self._env_row_data[name] = value.item() if isinstance(
+                    value, torch.Tensor) else value
+            else:
+                if len(self._train_row_data) == 0:
+                    self._train_row_data['step'] = i
+                self._train_row_data[name] = value.item() if isinstance(
+                    value, torch.Tensor) else value
+
+    def add_summaries(self, i, summaries):
+        for summary in summaries:
+            try:
+                if isinstance(summary, ScalarSummary):
+                    self.add_scalar(i, summary.name, summary.value)
+                elif self._tensorboard_logging:
+                    if isinstance(summary, HistogramSummary):
+                        self._tf_writer.add_histogram(
+                            summary.name, summary.value, i)
+                    elif isinstance(summary, ImageSummary):
+                        # Only grab first item in batch
+                        v = (summary.value if summary.value.ndim == 3 else
+                             summary.value[0])
+                        self._tf_writer.add_image(summary.name, v, i)
+                    elif isinstance(summary, VideoSummary):
+                        # Only grab first item in batch
+                        v = (summary.value if summary.value.ndim == 5 else
+                             np.array([summary.value]))
+                        self._tf_writer.add_video(
+                            summary.name, v, i, fps=summary.fps)
+                    elif isinstance(summary, TextSummary):
+                        self._tf_writer.add_text(summary.name, summary.value, i)
+            except Exception as e:
+                logging.error('Error on summary: %s' % summary.name)
+                raise e
+
+    def end_iteration(self):
+        # write train data
+        if self._csv_logging and len(self._train_row_data) > 0:
+            should_write_train_header = not os.path.exists(self._train_csv_file)
+            with open(self._train_csv_file, mode='a+') as csv_f:
+                names = self._train_row_data.keys()
+                writer = csv.DictWriter(csv_f, fieldnames=names)
+                if should_write_train_header:
+                    if self._train_field_names is None:
+                        writer.writeheader()
+                    else:
+                        if not np.array_equal(self._train_field_names, self._train_row_data.keys()):
+                            # Special case when we are logging faster than new
+                            # summaries are coming in.
+                            missing_keys = list(set(self._train_field_names) - set(
+                                self._train_row_data.keys()))
+                            for mk in missing_keys:
+                                self._train_row_data[mk] = self._train_prev_row_data[mk]
+                self._train_field_names = names
+                try:
+                    writer.writerow(self._train_row_data)
+                except Exception as e:
+                    print(e)
+            self._train_prev_row_data = self._train_row_data
+            self._train_row_data = OrderedDict()
+
+        # write env data (also eval or test during evaluation)
+        if self._csv_logging and len(self._env_row_data) > 0:
+            should_write_env_header = not os.path.exists(self._env_csv_file)
+            with open(self._env_csv_file, mode='a+') as csv_f:
+                names = self._env_row_data.keys()
+                writer = csv.DictWriter(csv_f, fieldnames=names)
+                if should_write_env_header:
+                    if self._env_field_names is None:
+                        writer.writeheader()
+                    else:
+                        if not np.array_equal(self._env_field_names, self._env_row_data.keys()):
+                            # Special case when we are logging faster than new
+                            # summaries are coming in.
+                            missing_keys = list(set(self._env_field_names) - set(
+                                self._env_row_data.keys()))
+                            for mk in missing_keys:
+                                self._env_row_data[mk] = self._env_prev_row_data[mk]
+                self._env_field_names = names
+                try:
+                    writer.writerow(self._env_row_data)
+                except Exception as e:
+                    print(e)
+            self._env_prev_row_data = self._env_row_data
+            self._env_row_data = OrderedDict()
+
+    def close(self):
+        if self._tensorboard_logging:
+            self._tf_writer.close()

external/yarr/yarr/utils/multi_task_rollout_generator.py

@@ -0,0 +1,65 @@
+from multiprocessing import Value
+
+import numpy as np
+
+from yarr.agents.agent import Agent
+from yarr.envs.env import Env
+from yarr.envs.multi_task_env import MultiTaskEnv
+from yarr.utils.transition import ReplayTransition
+
+
+class RolloutGenerator(object):
+
+    def _get_type(self, x):
+        if x.dtype == np.float64:
+            return np.float32
+        return x.dtype
+
+    def generator(self, step_signal: Value, env: MultiTaskEnv, agent: Agent,
+                  episode_length: int, timesteps: int, eval: bool):
+        obs = env.reset()
+        agent.reset()
+        obs_history = {k: [np.array(v, dtype=self._get_type(v))] * timesteps for k, v in obs.items()}
+        for step in range(episode_length):
+
+            prepped_data = {k: np.array([v]) for k, v in obs_history.items()}
+
+            act_result = agent.act(step_signal.value, prepped_data,
+                                   deterministic=eval)
+
+            # Convert to np if not already
+            agent_obs_elems = {k: np.array(v) for k, v in
+                               act_result.observation_elements.items()}
+            agent_extra_elems = {k: np.array(v) for k, v in
+                                 act_result.replay_elements.items()}
+
+            transition = env.step(act_result)
+            timeout = False
+            if step == episode_length - 1:
+                # If last transition, and not terminal, then we timed out
+                timeout = not transition.terminal
+                if timeout:
+                    transition.terminal = True
+                    if "needs_reset" in transition.info:
+                        transition.info["needs_reset"] = True
+
+            obs.update(agent_obs_elems)
+            obs_tp1 = dict(transition.observation)
+
+            for k in obs_history.keys():
+                obs_history[k].append(transition.observation[k])
+                obs_history[k].pop(0)
+
+            transition.info["active_task_id"] = env.active_task_id
+
+            replay_transition = ReplayTransition(
+                obs, act_result.action, transition.reward,
+                transition.terminal,
+                timeout, obs_tp1, agent_extra_elems,
+                transition.info)
+
+            obs = transition.observation
+            yield replay_transition
+
+            if transition.info.get("needs_reset", transition.terminal):
+                return

external/yarr/yarr/utils/observation_type.py

@@ -0,0 +1,10 @@
+from typing import Type
+import numpy as np
+
+
+class ObservationElement(object):
+
+    def __init__(self, name: str, shape: tuple, type: Type[np.dtype]):
+        self.name = name
+        self.shape = shape
+        self.type = type

external/yarr/yarr/utils/process_str.py

@@ -0,0 +1,5 @@
+from functools import reduce
+
+
+def change_case(str):
+    return reduce(lambda x, y: x + ('_' if y.isupper() else '') + y, str).lower()

external/yarr/yarr/utils/rollout_generator.py

@@ -0,0 +1,89 @@
+from multiprocessing import Value
+
+import numpy as np
+import torch
+from yarr.agents.agent import Agent
+from yarr.envs.env import Env
+from yarr.utils.transition import ReplayTransition
+
+
+class RolloutGenerator(object):
+
+    def _get_type(self, x):
+        if x.dtype == np.float64:
+            return np.float32
+        return x.dtype
+
+    def generator(self, step_signal: Value, env: Env, agent: Agent,
+                  episode_length: int, timesteps: int,
+                  eval: bool, eval_demo_seed: int = 0,
+                  record_enabled: bool = False):
+
+        if eval:
+            obs = env.reset_to_demo(eval_demo_seed)
+        else:
+            obs = env.reset()
+
+        agent.reset()
+        obs_history = {k: [np.array(v, dtype=self._get_type(v))] * timesteps for k, v in obs.items()}
+        for step in range(episode_length):
+
+            prepped_data = {k:torch.tensor(np.array(v)[None], device=self._env_device) for k, v in obs_history.items()}
+
+            act_result = agent.act(step_signal.value, prepped_data,
+                                   deterministic=eval)
+
+            # Convert to np if not already
+            agent_obs_elems = {k: np.array(v) for k, v in
+                               act_result.observation_elements.items()}
+            extra_replay_elements = {k: np.array(v) for k, v in
+                                     act_result.replay_elements.items()}
+
+            transition = env.step(act_result)
+            obs_tp1 = dict(transition.observation)
+            timeout = False
+            if step == episode_length - 1:
+                # If last transition, and not terminal, then we timed out
+                timeout = not transition.terminal
+                if timeout:
+                    transition.terminal = True
+                    if "needs_reset" in transition.info:
+                        transition.info["needs_reset"] = True
+
+            obs_and_replay_elems = {}
+            obs_and_replay_elems.update(obs)
+            obs_and_replay_elems.update(agent_obs_elems)
+            obs_and_replay_elems.update(extra_replay_elements)
+
+            for k in obs_history.keys():
+                obs_history[k].append(transition.observation[k])
+                obs_history[k].pop(0)
+
+            transition.info["active_task_id"] = env.active_task_id
+
+            replay_transition = ReplayTransition(
+                obs_and_replay_elems, act_result.action, transition.reward,
+                transition.terminal, timeout, summaries=transition.summaries,
+                info=transition.info)
+
+            if transition.terminal or timeout:
+                # If the agent gives us observations then we need to call act
+                # one last time (i.e. acting in the terminal state).
+                if len(act_result.observation_elements) > 0:
+                    prepped_data = {k: torch.tensor([v], device=self._env_device) for k, v in obs_history.items()}
+                    act_result = agent.act(step_signal.value, prepped_data,
+                                           deterministic=eval)
+                    agent_obs_elems_tp1 = {k: np.array(v) for k, v in
+                                           act_result.observation_elements.items()}
+                    obs_tp1.update(agent_obs_elems_tp1)
+                replay_transition.final_observation = obs_tp1
+
+            if record_enabled and transition.terminal or timeout or step == episode_length - 1:
+                env.env._action_mode.arm_action_mode.record_end(env.env._scene,
+                                                                steps=60, step_scene=True)
+
+            obs = dict(transition.observation)
+            yield replay_transition
+
+            if transition.info.get("needs_reset", transition.terminal):
+                return

external/yarr/yarr/utils/stat_accumulator.py

@@ -0,0 +1,192 @@
+from multiprocessing import Lock
+from typing import List
+
+import numpy as np
+from yarr.agents.agent import Summary, ScalarSummary
+from yarr.utils.transition import ReplayTransition
+
+
+class StatAccumulator(object):
+
+    def step(self, transition: ReplayTransition, eval: bool):
+        pass
+
+    def pop(self) -> List[Summary]:
+        pass
+
+    def peak(self) -> List[Summary]:
+        pass
+
+    def reset(self) -> None:
+        pass
+
+
+class Metric(object):
+
+    def __init__(self):
+        self._previous = []
+        self._current = 0
+
+    def update(self, value):
+        self._current += value
+
+    def next(self):
+        self._previous.append(self._current)
+        self._current = 0
+
+    def reset(self):
+        self._previous.clear()
+
+    def min(self):
+        return np.min(self._previous)
+
+    def max(self):
+        return np.max(self._previous)
+
+    def mean(self):
+        return np.mean(self._previous)
+
+    def median(self):
+        return np.median(self._previous)
+
+    def std(self):
+        return np.std(self._previous)
+
+    def __len__(self):
+        return len(self._previous)
+
+    def __getitem__(self, i):
+        return self._previous[i]
+
+
+class _SimpleAccumulator(StatAccumulator):
+
+    def __init__(self, prefix, eval_video_fps: int = 30,
+                 mean_only: bool = True):
+        self._prefix = prefix
+        self._eval_video_fps = eval_video_fps
+        self._mean_only = mean_only
+        self._lock = Lock()
+        self._episode_returns = Metric()
+        self._episode_lengths = Metric()
+        self._summaries = []
+        self._transitions = 0
+
+    def _reset_data(self):
+        with self._lock:
+            self._episode_returns.reset()
+            self._episode_lengths.reset()
+            self._summaries.clear()
+
+    def step(self, transition: ReplayTransition, eval: bool):
+        with self._lock:
+            self._transitions += 1
+            self._episode_returns.update(transition.reward)
+            self._episode_lengths.update(1)
+            if transition.terminal:
+                self._episode_returns.next()
+                self._episode_lengths.next()
+            self._summaries.extend(list(transition.summaries))
+
+    def _get(self) -> List[Summary]:
+        sums = []
+
+        if self._mean_only:
+            stat_keys = ["mean"]
+        else:
+            stat_keys = ["min", "max", "mean", "median", "std"]
+        names = ["return", "length"]
+        metrics = [self._episode_returns, self._episode_lengths]
+        for name, metric in zip(names, metrics):
+            for stat_key in stat_keys:
+                if self._mean_only:
+                    assert stat_key == "mean"
+                    sum_name = '%s/%s' % (self._prefix, name)
+                else:
+                    sum_name = '%s/%s/%s' % (self._prefix, name, stat_key)
+                sums.append(
+                    ScalarSummary(sum_name, getattr(metric, stat_key)()))
+        sums.append(ScalarSummary(
+            '%s/total_transitions' % self._prefix, self._transitions))
+        sums.extend(self._summaries)
+        return sums
+
+    def pop(self) -> List[Summary]:
+        data = []
+        if len(self._episode_returns) > 1:
+            data = self._get()
+            self._reset_data()
+        return data
+
+    def peak(self) -> List[Summary]:
+        return self._get()
+    
+    def reset(self):
+        self._transitions = 0
+        self._reset_data()
+
+
+class SimpleAccumulator(StatAccumulator):
+
+    def __init__(self, eval_video_fps: int = 30, mean_only: bool = True):
+        self._train_acc = _SimpleAccumulator(
+            'train_envs', eval_video_fps, mean_only=mean_only)
+        self._eval_acc = _SimpleAccumulator(
+            'eval_envs', eval_video_fps, mean_only=mean_only)
+
+    def step(self, transition: ReplayTransition, eval: bool):
+        if eval:
+            self._eval_acc.step(transition, eval)
+        else:
+            self._train_acc.step(transition, eval)
+
+    def pop(self) -> List[Summary]:
+        return self._train_acc.pop() + self._eval_acc.pop()
+
+    def peak(self) -> List[Summary]:
+        return self._train_acc.peak() + self._eval_acc.peak()
+    
+    def reset(self) -> None:
+        self._train_acc.reset()
+        self._eval_acc.reset()
+
+
+class MultiTaskAccumulator(StatAccumulator):
+
+    def __init__(self, num_tasks,
+                 eval_video_fps: int = 30, mean_only: bool = True,
+                 train_prefix: str = 'train_task',
+                 eval_prefix: str = 'eval_task'):
+        self._train_accs = [_SimpleAccumulator(
+            '%s%d/envs' % (train_prefix, i), eval_video_fps, mean_only=mean_only)
+            for i in range(num_tasks)]
+        self._eval_accs = [_SimpleAccumulator(
+            '%s%d/envs' % (eval_prefix, i), eval_video_fps, mean_only=mean_only)
+            for i in range(num_tasks)]
+        self._train_accs_mean = _SimpleAccumulator(
+            '%s_summary/envs' % train_prefix, eval_video_fps,
+            mean_only=mean_only)
+
+    def step(self, transition: ReplayTransition, eval: bool):
+        replay_index = transition.info["active_task_id"]
+        if eval:
+            self._eval_accs[replay_index].step(transition, eval)
+        else:
+            self._train_accs[replay_index].step(transition, eval)
+            self._train_accs_mean.step(transition, eval)
+
+    def pop(self) -> List[Summary]:
+        combined = self._train_accs_mean.pop()
+        for acc in self._train_accs + self._eval_accs:
+            combined.extend(acc.pop())
+        return combined
+
+    def peak(self) -> List[Summary]:
+        combined = self._train_accs_mean.peak()
+        for acc in self._train_accs + self._eval_accs:
+            combined.extend(acc.peak())
+        return combined
+
+    def reset(self) -> None:
+        self._train_accs_mean.reset()
+        [acc.reset() for acc in self._train_accs + self._eval_accs]

external/yarr/yarr/utils/transition.py

@@ -0,0 +1,33 @@
+from typing import List
+
+import numpy as np
+from yarr.agents.agent import Summary
+
+
+class Transition(object):
+
+    def __init__(self, observation: dict, reward: float, terminal: bool,
+                 info: dict = None, summaries: List[Summary] = None):
+        self.observation = observation
+        self.reward = reward
+        self.terminal = terminal
+        self.info = info or {}
+        self.summaries = summaries or []
+
+
+class ReplayTransition(object):
+
+    def __init__(self, observation: dict, action: np.ndarray,
+                 reward: float, terminal: bool, timeout: bool,
+                 final_observation: dict = None,
+                 summaries: List[Summary] = None,
+                 info: dict = None):
+        self.observation = observation
+        self.action = action
+        self.reward = reward
+        self.terminal = terminal
+        self.timeout = timeout
+        # final only populated on last timestep
+        self.final_observation = final_observation
+        self.summaries = summaries or []
+        self.info = info

external/yarr/yarr/utils/video_utils.py

@@ -0,0 +1,80 @@
+import os
+import numpy as np
+from pyrep.objects.dummy import Dummy
+from pyrep.objects.vision_sensor import VisionSensor
+from rlbench import Environment
+from rlbench.backend.observation import Observation
+
+
+class CameraMotion(object):
+    def __init__(self, cam: VisionSensor):
+        self.cam = cam
+
+    def step(self):
+        raise NotImplementedError()
+
+    def save_pose(self):
+        self._prev_pose = self.cam.get_pose()
+
+    def restore_pose(self):
+        self.cam.set_pose(self._prev_pose)
+
+
+class CircleCameraMotion(CameraMotion):
+
+    def __init__(self, cam: VisionSensor, origin: Dummy,
+                 speed: float, init_rotation: float = np.deg2rad(180)):
+        super().__init__(cam)
+        self.origin = origin
+        self.speed = speed  # in radians
+        self.origin.rotate([0, 0, init_rotation])
+
+    def step(self):
+        self.origin.rotate([0, 0, self.speed])
+
+
+class TaskRecorder(object):
+
+    def __init__(self, env: Environment, cam_motion: CameraMotion, fps=30):
+        self._env = env
+        self._cam_motion = cam_motion
+        self._fps = fps
+        self._snaps = []
+        self._current_snaps = []
+
+    def take_snap(self, obs: Observation):
+        self._cam_motion.step()
+        self._current_snaps.append(
+            (self._cam_motion.cam.capture_rgb() * 255.).astype(np.uint8))
+
+    def save(self, path, lang_goal, reward):
+        print(f"Converting to video ... {path}")
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        # OpenCV QT version can conflict with PyRep, so import here
+        import cv2
+        image_size = self._cam_motion.cam.get_resolution()
+        video = cv2.VideoWriter(
+                path, cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), self._fps,
+                tuple(image_size))
+        
+        for image in self._current_snaps:
+            frame = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+            font = cv2.FONT_HERSHEY_DUPLEX
+            font_scale = (0.45 * image_size[0]) / 640
+            font_thickness = 2
+
+
+            if lang_goal:
+
+                lang_textsize = cv2.getTextSize(lang_goal, font, font_scale, font_thickness)[0]
+                lang_textX = (image_size[0] - lang_textsize[0]) // 2
+
+                frame = cv2.putText(frame, lang_goal, org=(lang_textX, image_size[1] - 35),
+                                    fontScale=font_scale, fontFace=font, color=(0, 0, 0),
+                                    thickness=font_thickness, lineType=cv2.LINE_AA)
+                
+
+            video.write(frame)
+        video.release()
+        self._current_snaps = []