Publish TaskAnchor source; model transfer in progress

README.md

@@ -1,3 +1,15 @@
 ---
 license: apache-2.0
+tags:
+- taskanchor
+- pi05
+- robomemarena
 ---
+
+# TaskAnchor for RoboMemArena
+
+The training, inference, and evaluation source is available in [code/](code/README.md).
+
+**Model upload is in progress.** The complete 50k checkpoint is not available yet; do not attempt model loading until `50000/params/` and the final release manifest have been published.
+
+Source code is Apache-2.0. The full pretrained-model release will retain applicable upstream Gemma terms. No benchmark results are published here.

code/.gitignore

@@ -0,0 +1,23 @@
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ruff_cache/
+.venv/
+.env
+.env.*
+!.env.example
+.cache/
+.wandb/
+wandb/
+logs/
+/checkpoints/
+/pretrained/
+/data/
+*.egg-info/
+build/
+dist/
+/reports/
+auth.json
+*token*.txt
+
+/local/

code/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.

code/LICENSES/LICENSE_GEMMA.txt

@@ -0,0 +1,77 @@
+Gemma Terms of Use 
+
+Last modified: February 21, 2024
+
+By using, reproducing, modifying, distributing, performing or displaying any portion or element of Gemma, Model Derivatives including via any Hosted Service, (each as defined below) (collectively, the "Gemma Services") or otherwise accepting the terms of this Agreement, you agree to be bound by this Agreement.
+
+Section 1: DEFINITIONS
+1.1 Definitions
+(a) "Agreement" or "Gemma Terms of Use" means these terms and conditions that govern the use, reproduction, Distribution or modification of the Gemma Services and any terms and conditions incorporated by reference.
+
+(b) "Distribution" or "Distribute" means any transmission, publication, or other sharing of Gemma or Model Derivatives to a third party, including by providing or making Gemma or its functionality available as a hosted service via API, web access, or any other electronic or remote means ("Hosted Service").
+
+(c) "Gemma" means the set of machine learning language models, trained model weights and parameters identified at ai.google.dev/gemma, regardless of the source that you obtained it from.
+
+(d) "Google" means Google LLC.
+
+(e) "Model Derivatives" means all (i) modifications to Gemma, (ii) works based on Gemma, or (iii) any other machine learning model which is created by transfer of patterns of the weights, parameters, operations, or Output of Gemma, to that model in order to cause that model to perform similarly to Gemma, including distillation methods that use intermediate data representations or methods based on the generation of synthetic data Outputs by Gemma for training that model. For clarity, Outputs are not deemed Model Derivatives.
+
+(f) "Output" means the information content output of Gemma or a Model Derivative that results from operating or otherwise using Gemma or the Model Derivative, including via a Hosted Service.
+
+1.2
+As used in this Agreement, "including" means "including without limitation".
+
+Section 2: ELIGIBILITY AND USAGE
+2.1 Eligibility
+You represent and warrant that you have the legal capacity to enter into this Agreement (including being of sufficient age of consent). If you are accessing or using any of the Gemma Services for or on behalf of a legal entity, (a) you are entering into this Agreement on behalf of yourself and that legal entity, (b) you represent and warrant that you have the authority to act on behalf of and bind that entity to this Agreement and (c) references to "you" or "your" in the remainder of this Agreement refers to both you (as an individual) and that entity.
+
+2.2 Use
+You may use, reproduce, modify, Distribute, perform or display any of the Gemma Services only in accordance with the terms of this Agreement, and must not violate (or encourage or permit anyone else to violate) any term of this Agreement.
+
+Section 3: DISTRIBUTION AND RESTRICTIONS
+3.1 Distribution and Redistribution
+You may reproduce or Distribute copies of Gemma or Model Derivatives if you meet all of the following conditions:
+
+You must include the use restrictions referenced in Section 3.2 as an enforceable provision in any agreement (e.g., license agreement, terms of use, etc.) governing the use and/or distribution of Gemma or Model Derivatives and you must provide notice to subsequent users you Distribute to that Gemma or Model Derivatives are subject to the use restrictions in Section 3.2.
+You must provide all third party recipients of Gemma or Model Derivatives a copy of this Agreement.
+You must cause any modified files to carry prominent notices stating that you modified the files.
+All Distributions (other than through a Hosted Service) must be accompanied by a "Notice" text file that contains the following notice: "Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms".
+You may add your own intellectual property statement to your modifications and, except as set forth in this Section, may provide additional or different terms and conditions for use, reproduction, or Distribution of your modifications, or for any such Model Derivatives as a whole, provided your use, reproduction, modification, Distribution, performance, and display of Gemma otherwise complies with the terms and conditions of this Agreement. Any additional or different terms and conditions you impose must not conflict with the terms of this Agreement.
+
+3.2 Use Restrictions
+You must not use any of the Gemma Services:
+
+for the restricted uses set forth in the Gemma Prohibited Use Policy at ai.google.dev/gemma/prohibited_use_policy ("Prohibited Use Policy"), which is hereby incorporated by reference into this Agreement; or
+in violation of applicable laws and regulations.
+To the maximum extent permitted by law, Google reserves the right to restrict (remotely or otherwise) usage of any of the Gemma Services that Google reasonably believes are in violation of this Agreement.
+
+3.3 Generated Output
+Google claims no rights in Outputs you generate using Gemma. You and your users are solely responsible for Outputs and their subsequent uses.
+
+Section 4: ADDITIONAL PROVISIONS
+4.1 Updates
+Google may update Gemma from time to time, and you must make reasonable efforts to use the latest version of Gemma.
+
+4.2 Trademarks
+Nothing in this Agreement grants you any rights to use Google's trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between you and Google. Google reserves any rights not expressly granted herein.
+
+4.3 DISCLAIMER OF WARRANTY
+UNLESS REQUIRED BY APPLICABLE LAW, THE GEMMA SERVICES, AND OUTPUTS, ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING 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, REPRODUCING, MODIFYING, PERFORMING, DISPLAYING OR OR DISTRIBUTING ANY OF THE GEMMA SERVICES OR OUTPUTS AND ASSUME ANY AND ALL RISKS ASSOCIATED WITH YOUR USE OR DISTRIBUTION OF ANY OF THE GEMMA SERVICES OR OUTPUTS AND YOUR EXERCISE OF RIGHTS AND PERMISSIONS UNDER THIS AGREEMENT.
+
+4.4 LIMITATION OF LIABILITY
+TO THE FULLEST EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), PRODUCT LIABILITY, CONTRACT, OR OTHERWISE, UNLESS REQUIRED BY APPLICABLE LAW, SHALL GOOGLE OR ITS AFFILIATES BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, CONSEQUENTIAL, OR PUNITIVE DAMAGES, OR LOST PROFITS OF ANY KIND ARISING FROM THIS AGREEMENT OR RELATED TO, ANY OF THE GEMMA SERVICES OR OUTPUTS EVEN IF GOOGLE OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+
+4.5 Term, Termination, and Survival
+The term of this Agreement will commence upon your acceptance of this Agreement (including acceptance by your use, modification, or Distribution, reproduction, performance or display of any portion or element of the Gemma Services) and will continue in full force and effect until terminated in accordance with the terms of this Agreement. Google may terminate this Agreement if you are in breach of any term of this Agreement. Upon termination of this Agreement, you must delete and cease use and Distribution of all copies of Gemma and Model Derivatives in your possession or control. Sections 1, 2.1, 3.3, 4.2 to 4.9 shall survive the termination of this Agreement.
+
+4.6 Governing Law and Jurisdiction
+This Agreement will be governed by the laws of the State of California without regard to choice of law principles. The UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The state and federal courts of Santa Clara County, California shall have exclusive jurisdiction of any dispute arising out of this Agreement.
+
+4.7 Severability
+If any provision of this Agreement is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein.
+
+4.8 Entire Agreement
+This Agreement states all the terms agreed between the parties and supersedes all other agreements between the parties as of the date of acceptance relating to its subject matter.
+
+4.9 No Waiver
+Google will not be treated as having waived any rights by not exercising (or delaying the exercise of) any rights under this Agreement.

code/NOTICE

@@ -0,0 +1,22 @@
+TaskAnchor RoboMemArena π0.5 release
+
+Source code in this package is distributed under Apache License 2.0, except
+where an individual file or upstream component states otherwise. Retain the
+existing copyright and license notices when redistributing it.
+
+This distribution includes modified OpenPI code:
+https://github.com/Physical-Intelligence/openpi
+Modifications add TaskAnchor, full-task history sampling, progress
+conditioning, RoboMemArena training/evaluation, and portable file locations.
+The original model computations in the selected release profile are retained.
+
+The accompanying full model includes pretrained π0.5/PaliGemma-derived
+components. The Apache source license does not override terms that apply
+to the underlying pretrained model. The included Gemma terms continue to
+apply to Gemma-derived material.
+
+Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms
+
+The official RoboMemArena simulator, assets, and training dataset are not
+redistributed in this source package; obtain them from their upstream
+locations and retain their respective terms.

code/README.md

@@ -0,0 +1,96 @@
+# TaskAnchor for RoboMemArena (π0.5, 50k)
+
+Training and evaluation source for the released **TaskAnchor–π0.5 50,000-step checkpoint** on RoboMemArena. The supported public profile is `robomemarena/pi05`. This package contains the native OpenPI backend, the TaskAnchor adapter, the benchmark interface, normalization assets, and evaluation tools.
+
+**Weights:** [henry13522/Taskanchor](https://huggingface.co/henry13522/Taskanchor). The complete source is also included there under `code/`; no separate GitHub repository is required to use this release.
+
+## Architecture and checkpoint contract
+
+The frozen SigLIP encoder produces 1152-dimensional tokens. TaskAnchor projects the main-camera tokens to a 512-dimensional latent space, applies the eight-layer, eight-head temporal adapter, and restores the visual residual to 1152 dimensions. The writeback is **`refined = current + offset`**, with no additional scalar gate. Adapter dropout is 0.10; the MLP ratio is 4. The native π0.5 backbone and action decoder remain present in the full checkpoint.
+
+The main camera and wrist camera are real inputs; the third native image slot is masked padding. Only the main camera receives the history-conditioned visual residual. Images are resized to 224 × 224. Each cached history frame contains 4 × 4 spatial tokens, bilinearly pooled from the native 16 × 16 SigLIP grid. The recorded protocol retains **a separate permanent frame0 plus up to 191 history frames**, sampled with stride 10. This is the 50k checkpoint's protocol, not the FIFO-only configuration used in some other experiments.
+
+The task-state head predicts a continuous scalar. Training targets are `round(subtask_id / num_subtasks, 2)`. At evaluation, the prediction is mapped to the nearest task-specific coordinate (for four stages: 0.00, 0.25, 0.50, 0.75) before conditioning the policy. These are predicted stages, not oracle labels. The model predicts 20 actions and executes 10 before replanning.
+
+There are **26,930,049 adapter parameters** and **3,380,363,921 parameters in the complete checkpoint**. The full tensor inventory and payload hashes are provided with the weights.
+
+## Download and run
+
+```bash
+hf download henry13522/Taskanchor --local-dir /path/to/Taskanchor
+cd /path/to/Taskanchor/code
+```
+
+Follow [REPRODUCING.md](docs/REPRODUCING.md) to install the model and simulator environments, pin the official benchmark, and run extraction, training, serving, or evaluation. [PORTABILITY.md](docs/PORTABILITY.md) lists the supported path overrides.
+
+After environment and benchmark setup, the evaluation entry point is shown below. The evaluator retains the original trial and seed schedule; a study requiring a different schedule must configure and record its evaluation protocol before collecting results.
+
+```bash
+export TASKANCHOR_TOKENIZER_PATH=/path/to/Taskanchor/tokenizer/paligemma_tokenizer.model
+export TASKANCHOR_PI05_PYTHON=/path/to/Taskanchor/code/pi05/.venv/bin/python
+export TASKANCHOR_ROBOMEM_EVAL_PYTHON=/path/to/simulator-env/bin/python
+export TASKANCHOR_LIBERO_CONFIG=/path/to/robomem_libero_config
+
+python -m taskanchor --benchmark robomemarena --model pi05 --stage eval \
+  --benchmark-root /path/to/RoboMemArena/evaluation_benchmark \
+  --python "$TASKANCHOR_PI05_PYTHON" -- \
+  --checkpoint /path/to/Taskanchor/50000 \
+  --output /path/to/reports/taskanchor_50k \
+  --workers 8 --base-port 8800
+```
+
+Keep `recipe.json` one level above `50000/`. Inference also needs `50000/assets/` and the supplied tokenizer. No training data, feature cache, base-model weight download, or optimizer state is required to load the complete released checkpoint.
+
+## Training settings
+
+The recorded run was configured for **55,000 total optimizer steps**, with 5,000 adapter-only steps followed by 50,000 joint steps. The **released 50k snapshot contains 5k adapter-only + 45k joint steps**. The visual encoder stays frozen throughout; the remaining backbone parameters participate after the warm-up. This is one training run with a freeze schedule, not the older four-stage pipeline.
+
+| Setting | Recorded value |
+| --- | --- |
+| Initialization | OpenPI `pi05_base` |
+| Global batch / devices | 256 / 8 GPUs, FSDP size 8 |
+| Optimizer | AdamW, β=(0.9, 0.95), gradient clipping 1.0 |
+| Adapter learning rate / weight decay | 2.5e-5 / 0.05 |
+| Backbone learning rate / weight decay | 2.5e-5 / 0 |
+| Image projection learning rate | 1e-5 |
+| Schedule | Adapter constant for the first 5k; cosine decay to 10% over the 50k joint phase |
+| Loss | `flow + 0.01 * (100 * MSE + 1 - Pearson)` |
+| Action prediction / execution | 20 / 10 |
+| History | stride 10, 191 frames + separate frame0, 16 tokens/frame |
+| Dataset | 2,600 full-task trajectories, 2,811,755 frames |
+| Seed / workers | 42 / 8 |
+
+The recorded run resumed at 10k after a data-loading change. Its initial recipe and saved checkpoint metadata retain this provenance; they are not claimed to be identical restart states. See `evidence/pipeline_restart_10000.json`. Training from scratch uses the same numerical settings, but exact retraining reproducibility is not asserted.
+
+## Source layout
+
+| Path | Responsibility |
+| --- | --- |
+| `taskanchor/pi05/taskanchor.py` | Adapter, task-state head, native π0.5 integration |
+| `taskanchor/pi05/ttt.py` | History-dependent temporary fast-weight computation |
+| `taskanchor/pi05/training.py` | Parameter groups, freeze schedule, optimization |
+| `taskanchor/data/` | Full-task indexing and historical token selection |
+| `taskanchor/__main__.py` | Native extraction/train/serve/evaluation launcher |
+| `taskanchor/paths.py` | Portable asset and data locations |
+| `pi05/src/openpi/` | Native model, transforms, checkpoint and serving utilities |
+| `pi05/scripts/*robomem*` | Feature extraction, training, inference, official evaluation bridge |
+| `pi05/packages/openpi-client/` | Policy client/server protocol support |
+| `tools/` | Benchmark setup and result comparison |
+| `assets/pi05_libero/` | The original normalization statistics |
+| `evidence/` | Configuration provenance and environment records |
+
+Some native OpenPI compatibility modules retain historical configuration names because they are imported by the backend. Only the RoboMemArena π0.5 entry points documented here are supported in this release.
+
+## Verification and scope
+
+During packaging, the complete checkpoint was restored on CPU: all 176 tensors are finite, and copied parameter/tokenizer/normalization files match the originals by SHA-256. Core adapter, TTT, history, and training implementations are retained without numerical changes. Packaging checks do not substitute for closed-loop benchmark evaluation. This repository does not publish a benchmark score. The packaged checkpoint omits optimizer state and therefore cannot be used as an exact training resume checkpoint.
+
+```bash
+python -m unittest discover -s tests -p test_portability.py
+```
+
+## License and attribution
+
+The code is distributed under [Apache-2.0](LICENSE), retaining upstream notices. The full checkpoint includes pretrained π0.5/PaliGemma-derived components and remains subject to upstream terms, including [Gemma Terms of Use](LICENSES/LICENSE_GEMMA.txt); Apache-2.0 for this source does not replace those weight terms. See [NOTICE](NOTICE).
+
+This work builds on [OpenPI](https://github.com/Physical-Intelligence/openpi) and [RoboMemArena](https://github.com/OpenHelix-Team/RoboMemArena). The benchmark, its assets, and the training dataset are obtained separately under their respective terms. Public model weights and source: [henry13522/Taskanchor](https://huggingface.co/henry13522/Taskanchor).

code/SOURCE_MANIFEST.json

@@ -0,0 +1,654 @@
+{
+  "files": [
+    {
+      "path": ".gitignore",
+      "size_bytes": 213,
+      "sha256": "281dcbe819391034b619056c0d6ae4fac7b69ba751600f6864d6ae035d54e967"
+    },
+    {
+      "path": "LICENSE",
+      "size_bytes": 11356,
+      "sha256": "43070e2d4e532684de521b885f385d0841030efa2b1a20bafb76133a5e1379c1"
+    },
+    {
+      "path": "LICENSES/LICENSE_GEMMA.txt",
+      "size_bytes": 8431,
+      "sha256": "3e2c24001f9ef57bf7ec959a3658fbb49cdad113cdf394c264da9d16f9bdd132"
+    },
+    {
+      "path": "NOTICE",
+      "size_bytes": 1114,
+      "sha256": "38e5a788a2ca6eb7d7fca20c19cbfdc12cc22cb54df32ecd5d82587dcce6b7c9"
+    },
+    {
+      "path": "README.md",
+      "size_bytes": 7644,
+      "sha256": "2224078f9713aedf4e746b16f538d3cb1882a53911d4f7d7d774a7a30e0a2ace"
+    },
+    {
+      "path": "assets/pi05_libero/norm_stats.json",
+      "size_bytes": 1874,
+      "sha256": "0c0d329a3345d2ea2e1348dac0edf4ae175085fac200cca1cab967aae1ae1767"
+    },
+    {
+      "path": "docs/PORTABILITY.md",
+      "size_bytes": 5477,
+      "sha256": "b9b39576f3d23c1edfd42f28f8d4c00a0e803e771b4856bd6bdf4cc7922765e2"
+    },
+    {
+      "path": "docs/PUBLISHING_ZH.md",
+      "size_bytes": 2346,
+      "sha256": "72579bf3257fb1ae7b6e72d58e3fb3e31584fb73ad23da93bfc31770f2279eb9"
+    },
+    {
+      "path": "docs/REPRODUCING.md",
+      "size_bytes": 12827,
+      "sha256": "7630f8f14f2d232ff49509fe54bca7dc33d20ed1fab3bcc22b7d0fca556692d2"
+    },
+    {
+      "path": "evidence/environment_pi05.json",
+      "size_bytes": 393,
+      "sha256": "463b318c2fcc9b4374bdf177e55211a134fad79be7fa4f8818d69237208a6a47"
+    },
+    {
+      "path": "evidence/packaging_checks.json",
+      "size_bytes": 465,
+      "sha256": "cdd294d1e4c23b38f42e37e2f15ca0f43d7600b7e6e76d4feebde61fa0f11de0"
+    },
+    {
+      "path": "evidence/pipeline_restart_10000.json",
+      "size_bytes": 1027,
+      "sha256": "55576107d7239b76ad608ea21a43307bdb72886b951ec2a214e006218a554984"
+    },
+    {
+      "path": "evidence/robomem50k_adapter_numerical_parity.json",
+      "size_bytes": 628,
+      "sha256": "d0067e72d6c768996c64d4ed66eb1d9351cdf94a232ce3985b8314fb5c76c14a"
+    },
+    {
+      "path": "evidence/task25_seed_candidates.json",
+      "size_bytes": 1930,
+      "sha256": "2c2c859aa60c0c58e04e9f8c7e7ec54cbb3a23a34d463d8432ab10080e7f4444"
+    },
+    {
+      "path": "pi05/LICENSE",
+      "size_bytes": 11356,
+      "sha256": "43070e2d4e532684de521b885f385d0841030efa2b1a20bafb76133a5e1379c1"
+    },
+    {
+      "path": "pi05/LICENSE_GEMMA.txt",
+      "size_bytes": 8431,
+      "sha256": "3e2c24001f9ef57bf7ec959a3658fbb49cdad113cdf394c264da9d16f9bdd132"
+    },
+    {
+      "path": "pi05/README.md",
+      "size_bytes": 23254,
+      "sha256": "567fc57766bd736664f50126cab8145af4071ca16cc5e238f2b11e369169c3b5"
+    },
+    {
+      "path": "pi05/packages/openpi-client/pyproject.toml",
+      "size_bytes": 388,
+      "sha256": "cacd55394787338bef9d6e7c2e4ecaefd0bc90b3362ac48fe20d9a7f201147d6"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/__init__.py",
+      "size_bytes": 22,
+      "sha256": "91447944015cec709e8aa7655f7e9d64e1e4508e7023a57fe3746911c0fc6fed"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/action_chunk_broker.py",
+      "size_bytes": 1404,
+      "sha256": "019b9d69cb26cc681331a2f95dbe60a2cc7b02d6d309c5d15b6457ab39798234"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/base_policy.py",
+      "size_bytes": 271,
+      "sha256": "88e98e0f06293f3c4767db675061664c4bba12f87242ba1e64377d1b8de8069e"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/image_tools.py",
+      "size_bytes": 2426,
+      "sha256": "d48b4bd7f44e79fe6db8a8e07c9161144fa250be686e1245014a8b47e6171977"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/image_tools_test.py",
+      "size_bytes": 1411,
+      "sha256": "4df4b0ba0b95ba5e47160039c3e837a64d50c802b419a588211dffbaa732589c"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/msgpack_numpy.py",
+      "size_bytes": 1944,
+      "sha256": "c04568948fcee52b691e3be4b6cffb759f7e79ad67530fcd5d23095a0d13c057"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/msgpack_numpy_test.py",
+      "size_bytes": 1643,
+      "sha256": "2c6dd236720dae716d7bdf842b7d506a165e9563bf9d4376c95289269adcb973"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/runtime/agent.py",
+      "size_bytes": 465,
+      "sha256": "b559885afd9645a91b9b6ee12a28513295474f987d11b16610fc0479dbe6de5b"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/runtime/agents/policy_agent.py",
+      "size_bytes": 509,
+      "sha256": "c4d63952d1954bd4b72a3122223df36794812fea74cf1c7023433b5a90473c1c"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/runtime/environment.py",
+      "size_bytes": 1044,
+      "sha256": "1c12e5054b3401dec5f1f8b71eef0f625cb60e2b50784dcf19afa32c77107a7e"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/runtime/runtime.py",
+      "size_bytes": 3049,
+      "sha256": "3daa242a9ce61ffb210789ef3ba6b1e3cade9de757fcd36af7b174c530345ea8"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/runtime/subscriber.py",
+      "size_bytes": 500,
+      "sha256": "d269a2af8ec4056beb310e5c3c97d4468af01549c7ef80dd905652d8ac64f85b"
+    },
+    {
+      "path": "pi05/packages/openpi-client/src/openpi_client/websocket_client_policy.py",
+      "size_bytes": 2131,
+      "sha256": "36557cb0b91ccf31cd4fb4b508306850d76ed0feb4028dac5182d0f5a5d88005"
+    },
+    {
+      "path": "pi05/pyproject.toml",
+      "size_bytes": 3110,
+      "sha256": "e1f457bd81622a47a4c4de17b92815977ac5fcd40116bec92398eb400bf76483"
+    },
+    {
+      "path": "pi05/scripts/__init__.py",
+      "size_bytes": 0,
+      "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+    },
+    {
+      "path": "pi05/scripts/check_robomem_memory_graph.py",
+      "size_bytes": 2521,
+      "sha256": "1100021fb2ae25cc85f3e1c0da2cd95fa78086da0a673fd606be8652adcab752"
+    },
+    {
+      "path": "pi05/scripts/evaluate_robomem_memory.py",
+      "size_bytes": 13126,
+      "sha256": "47fe434454e7db74d834016621f6ae61d12097ad270dd7fb155567604b8ab53d"
+    },
+    {
+      "path": "pi05/scripts/extract_robomem_memory.py",
+      "size_bytes": 7229,
+      "sha256": "4bf8a66484188ee321ee75b605e80cbfc55b9785c3332cf36dce15c45c186342"
+    },
+    {
+      "path": "pi05/scripts/robomem_eval_paths.py",
+      "size_bytes": 843,
+      "sha256": "f82cda8a7fecbe3def79e5c013cab15083d19e53634b80b41f984534c25d487e"
+    },
+    {
+      "path": "pi05/scripts/robomem_memory_eval_adapter.py",
+      "size_bytes": 4093,
+      "sha256": "771cc4492686bce2d81d6f320301df3dcd619e3e9eba8f61d70800f9274e84ce"
+    },
+    {
+      "path": "pi05/scripts/robomem_seed_search.py",
+      "size_bytes": 2742,
+      "sha256": "c48b3bfc0b30fc1ed7eb986925c42f4821fe52a856d9102a3181fd3e9dd6b2be"
+    },
+    {
+      "path": "pi05/scripts/run_robomem_eval_task.py",
+      "size_bytes": 4547,
+      "sha256": "18c698ec0bda3510816390778d1412a96a90f7d2bf098ee8224033a9ea1708f3"
+    },
+    {
+      "path": "pi05/scripts/run_robomem_task25_seed.py",
+      "size_bytes": 1225,
+      "sha256": "ca32a0bdb52376e0937bced4f87f41b87f18e5fb3f75cd6c975bc71dde166321"
+    },
+    {
+      "path": "pi05/scripts/serve_robomem_memory.py",
+      "size_bytes": 702,
+      "sha256": "e560e83a556a16b620c854140fb2a736ca1c003d8031a531801a71ef507af86e"
+    },
+    {
+      "path": "pi05/scripts/test_robomem_discrete_progress.py",
+      "size_bytes": 2915,
+      "sha256": "1f9e66c70775fac7933b7de3bc15f33aa434532fdad79165269f311c077b29a5"
+    },
+    {
+      "path": "pi05/scripts/test_robomem_memory.py",
+      "size_bytes": 8216,
+      "sha256": "e1b71aee89412650a99b4489961450c36d188ce654db182ce2605befcad8c893"
+    },
+    {
+      "path": "pi05/scripts/test_robomem_memory_inference.py",
+      "size_bytes": 2430,
+      "sha256": "d26aadc408ef8300a879c86cf1fe3fc32df7de3d56f1ff19c3ee7803a9d2ca57"
+    },
+    {
+      "path": "pi05/scripts/test_robomem_memory_loader.py",
+      "size_bytes": 3000,
+      "sha256": "80f421f909c0d3373171dccd013c46fbf8b2b9a67597b81a6198e6163cdd2ba7"
+    },
+    {
+      "path": "pi05/scripts/train_robomem_memory.py",
+      "size_bytes": 8168,
+      "sha256": "6e7d4a8ed7b6686437ed85101a0b90518fe555fd5dadb2d8d8d681d5beb25fb0"
+    },
+    {
+      "path": "pi05/src/openpi/__init__.py",
+      "size_bytes": 0,
+      "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+    },
+    {
+      "path": "pi05/src/openpi/conftest.py",
+      "size_bytes": 339,
+      "sha256": "626672e385fb276396c119e2220421065c58f8ec602401e8bf2ae509f9fffaf3"
+    },
+    {
+      "path": "pi05/src/openpi/models/__init__.py",
+      "size_bytes": 0,
+      "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+    },
+    {
+      "path": "pi05/src/openpi/models/gemma.py",
+      "size_bytes": 17146,
+      "sha256": "7e42ada4ae7e9995f0ef3e33a0c224758323c2015e2d9d72ec15726a8a001d2a"
+    },
+    {
+      "path": "pi05/src/openpi/models/gemma_fast.py",
+      "size_bytes": 15643,
+      "sha256": "563c6e1ffd8afa31c406bcb8894a605e87df88f6c075d398d97a57af6bc3aefa"
+    },
+    {
+      "path": "pi05/src/openpi/models/lora.py",
+      "size_bytes": 5342,
+      "sha256": "ced28716950c34d84a1a6d0b03b1d9550000732721df884818080b6e78cb3289"
+    },
+    {
+      "path": "pi05/src/openpi/models/lora_test.py",
+      "size_bytes": 3138,
+      "sha256": "a0f10525de7773310402b3ffa636e88ffff7035f3179b140aacc6809c122f085"
+    },
+    {
+      "path": "pi05/src/openpi/models/model.py",
+      "size_bytes": 15120,
+      "sha256": "386dd172e1f561ba31d54a0ec3eb92921a51f58e4ee520b12789c9b8e3a62d07"
+    },
+    {
+      "path": "pi05/src/openpi/models/model_test.py",
+      "size_bytes": 2966,
+      "sha256": "f8e71f195045d4e3ef1ddfafb2337fa52e5aa2f89bce1fde62f5cbfb67121f21"
+    },
+    {
+      "path": "pi05/src/openpi/models/pi0.py",
+      "size_bytes": 14318,
+      "sha256": "00df716074704c22ab4fccd88146bf73ed6ef7edcb4c9e6a9f92486ad030c89c"
+    },
+    {
+      "path": "pi05/src/openpi/models/pi0_config.py",
+      "size_bytes": 4435,
+      "sha256": "d321858c30b398126035ea2b6da09e1c422eeeebb5ebf8df631daf9dc77647fa"
+    },
+    {
+      "path": "pi05/src/openpi/models/pi0_fast.py",
+      "size_bytes": 13209,
+      "sha256": "627332a555046f2744ef144c45f8bd3361a9820cf68854a4528d0fdbba7e9aa3"
+    },
+    {
+      "path": "pi05/src/openpi/models/pi0_test.py",
+      "size_bytes": 1631,
+      "sha256": "5082873014a7fc45fb8414f517909e1c24cdafea980a1fb612e0aea6058ed0fe"
+    },
+    {
+      "path": "pi05/src/openpi/models/robomem_memory.py",
+      "size_bytes": 189,
+      "sha256": "cff45c0afe162f062ca54dfb52078483e64cb04db268746f5ea570e5268e1ca5"
+    },
+    {
+      "path": "pi05/src/openpi/models/siglip.py",
+      "size_bytes": 12471,
+      "sha256": "f129a7e1d5b8bf101d73f27ccf9b0005167c3e208ac934ffcdc36a886662d36f"
+    },
+    {
+      "path": "pi05/src/openpi/models/tokenizer.py",
+      "size_bytes": 17558,
+      "sha256": "0d604dfde16c3d00fe056ff7559efd4e5d95a042f4fe09e1c2ec7e77dfc47b0a"
+    },
+    {
+      "path": "pi05/src/openpi/models/tokenizer_test.py",
+      "size_bytes": 807,
+      "sha256": "5abd179ce883420e39d4d18b33d0123290af752da243003aa48a3bc06dbb9517"
+    },
+    {
+      "path": "pi05/src/openpi/models/utils/fsq_tokenizer.py",
+      "size_bytes": 15928,
+      "sha256": "88866638ee4775f46a222561a957bd35e20600fe38c5888b587016ad77b1065f"
+    },
+    {
+      "path": "pi05/src/openpi/models/vit.py",
+      "size_bytes": 10266,
+      "sha256": "d285d057636c3cba2d1fd3b37693b18075ad4b9c3b47bbde807c34c6c376ac93"
+    },
+    {
+      "path": "pi05/src/openpi/models_pytorch/gemma_pytorch.py",
+      "size_bytes": 13480,
+      "sha256": "225092e06509da0ae37bb9dca55c00d2804255bb1496f7c8700304b4c52a95eb"
+    },
+    {
+      "path": "pi05/src/openpi/models_pytorch/pi0_pytorch.py",
+      "size_bytes": 22558,
+      "sha256": "97dd290d699bb63719627e3933a414500b35ded4ef24c88faa0685f0acdc9e75"
+    },
+    {
+      "path": "pi05/src/openpi/models_pytorch/preprocessing_pytorch.py",
+      "size_bytes": 8098,
+      "sha256": "f8fb26173ecfe83b7f7832f8d6ca4104d377af08ebb7352dc68ad51b63a0eb43"
+    },
+    {
+      "path": "pi05/src/openpi/models_pytorch/transformers_replace/models/gemma/configuration_gemma.py",
+      "size_bytes": 8958,
+      "sha256": "26fd0d0f52730ab32b6f337eb73ae9a26cb71ae0915bc49eced1d0e2ed8d0ec6"
+    },
+    {
+      "path": "pi05/src/openpi/models_pytorch/transformers_replace/models/gemma/modeling_gemma.py",
+      "size_bytes": 36392,
+      "sha256": "17eeb54e277939e58cd6f8a92719a275e36baf0c9f463b2754227d42f87aedde"
+    },
+    {
+      "path": "pi05/src/openpi/models_pytorch/transformers_replace/models/paligemma/modeling_paligemma.py",
+      "size_bytes": 27688,
+      "sha256": "c945d330b829264f927f3ee5339f977238631772de07a28e855e265bb85fd53c"
+    },
+    {
+      "path": "pi05/src/openpi/models_pytorch/transformers_replace/models/siglip/check.py",
+      "size_bytes": 133,
+      "sha256": "466e8eb7887ac5ccc98118a3169298d6ab65a284d834bae74751e3854dbc4245"
+    },
+    {
+      "path": "pi05/src/openpi/models_pytorch/transformers_replace/models/siglip/modeling_siglip.py",
+      "size_bytes": 50231,
+      "sha256": "ef2e99500f263fdd78db9d4d9a255e29d844eb92f53c25ad640f8056700ef84b"
+    },
+    {
+      "path": "pi05/src/openpi/policies/aloha_policy.py",
+      "size_bytes": 7543,
+      "sha256": "50d7ad972d80142ecb0f03b0780dacfd3af212946f4c9aac31a204b02db4543d"
+    },
+    {
+      "path": "pi05/src/openpi/policies/droid_policy.py",
+      "size_bytes": 3186,
+      "sha256": "7d0893555964d485e231f70084125084a99495705ddd74b29461229d524fff26"
+    },
+    {
+      "path": "pi05/src/openpi/policies/libero_policy.py",
+      "size_bytes": 3461,
+      "sha256": "8606b38a3efeb8b97c1af18b33a873eb3c63f5359a661c9b4319cdb87fb66276"
+    },
+    {
+      "path": "pi05/src/openpi/policies/rmbench_memory_inference.py",
+      "size_bytes": 4643,
+      "sha256": "9335930d0da0efe1c64b6c06a58820c7f7eb522b21933a2cc4895508b5f65ff0"
+    },
+    {
+      "path": "pi05/src/openpi/policies/robomem_memory_inference.py",
+      "size_bytes": 12111,
+      "sha256": "3ff7d6ac0a2e9a54b011ecf557c4d2e4810c2ab144781b6ce95b1093103a02de"
+    },
+    {
+      "path": "pi05/src/openpi/policies/robomem_progress.py",
+      "size_bytes": 1371,
+      "sha256": "43b8afbdb6f96400f65b8bcb74e4b077856ede5e874d1513d6b16ef674ff98c3"
+    },
+    {
+      "path": "pi05/src/openpi/py.typed",
+      "size_bytes": 0,
+      "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+    },
+    {
+      "path": "pi05/src/openpi/rmbench_cover_blocks_10_full_prompts.py",
+      "size_bytes": 4524,
+      "sha256": "65e27c9cd1d89238e26005e10525cca8af2faec5837c0221dc3a9b02f2d8304f"
+    },
+    {
+      "path": "pi05/src/openpi/rmbench_eval_prompts.py",
+      "size_bytes": 1406,
+      "sha256": "c3df981e542ecb1215fb10ee4a56aa3fec7ea3af6ec998e096e0e29c5580d4f5"
+    },
+    {
+      "path": "pi05/src/openpi/rmbench_m1_5tasks_10_full_prompts.py",
+      "size_bytes": 14008,
+      "sha256": "488b3a0b0854a71d4a192bd9c02718ea7df307eb6d335276244178820cc0ae4c"
+    },
+    {
+      "path": "pi05/src/openpi/robomem_task6_100_instruction.py",
+      "size_bytes": 21358,
+      "sha256": "e6ae6b788d90b206187550b61d115d10a076cd4dde02afe2c2e5e5cc6dc95a97"
+    },
+    {
+      "path": "pi05/src/openpi/robomemarena_26tasks_10_full_prompts.py",
+      "size_bytes": 40410,
+      "sha256": "8f8c71573de9701321f62b9cbb334829b97199e6b01831f6d07fb990dcf92906"
+    },
+    {
+      "path": "pi05/src/openpi/robomemarena_26tasks_subtask_prompts.py",
+      "size_bytes": 5087,
+      "sha256": "8e4e2b3e7d4e88d927c5d4147ff41f91db423231072e302b4bc02fd77a45035e"
+    },
+    {
+      "path": "pi05/src/openpi/serving/websocket_policy_server.py",
+      "size_bytes": 3480,
+      "sha256": "001f06365933bd11386298905c9125ac28db9a8985f35d10ae090d174cc8e1ae"
+    },
+    {
+      "path": "pi05/src/openpi/shared/__init__.py",
+      "size_bytes": 0,
+      "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+    },
+    {
+      "path": "pi05/src/openpi/shared/array_typing.py",
+      "size_bytes": 3402,
+      "sha256": "6524945566eab3eb97b75a6280ba4724617b4d028c55f05172b283c3de68a026"
+    },
+    {
+      "path": "pi05/src/openpi/shared/download.py",
+      "size_bytes": 8862,
+      "sha256": "a12054a6ddaab4efa833c4ff49a1aba00c63759e4e88d6be21c1e0b3307fd9b2"
+    },
+    {
+      "path": "pi05/src/openpi/shared/download_test.py",
+      "size_bytes": 1490,
+      "sha256": "234ee70ab2501a1421f6c44cbd83c1e16747597023319ae1446bed9d44f17072"
+    },
+    {
+      "path": "pi05/src/openpi/shared/image_tools.py",
+      "size_bytes": 4665,
+      "sha256": "0f8fcb355288119e56acb56d550c1cc1f3a9edddceb0dc5b319f955943ef7767"
+    },
+    {
+      "path": "pi05/src/openpi/shared/image_tools_test.py",
+      "size_bytes": 1418,
+      "sha256": "842374579dc9f4c2907da5e8bb555efe5a3f4902b659720c8e90115863dd59fe"
+    },
+    {
+      "path": "pi05/src/openpi/shared/nnx_utils.py",
+      "size_bytes": 2844,
+      "sha256": "4da7111ba3011c4d1c381c9ed17257771df6a50ed89182647c3fb6f79208005e"
+    },
+    {
+      "path": "pi05/src/openpi/shared/normalize.py",
+      "size_bytes": 5529,
+      "sha256": "6c2cea4946fb07e51801530400d2b2fd94730e195cd290d6b6960114eca9739d"
+    },
+    {
+      "path": "pi05/src/openpi/shared/normalize_test.py",
+      "size_bytes": 1521,
+      "sha256": "fda27a21cf37f337b12bcab71229407c2a349c5965c33f278bdc468ed03f58af"
+    },
+    {
+      "path": "pi05/src/openpi/training/checkpoints.py",
+      "size_bytes": 6696,
+      "sha256": "ea420c6a816933ef96966a092801f388997800530aef7a29f7b42807cc3e5103"
+    },
+    {
+      "path": "pi05/src/openpi/training/config.py",
+      "size_bytes": 35031,
+      "sha256": "e1cde1a9906848460cd92a3c82155a83047610d983cb933e8f6f15bb93c1ac22"
+    },
+    {
+      "path": "pi05/src/openpi/training/data_loader.py",
+      "size_bytes": 23482,
+      "sha256": "7ccd627788901c9c6d1db222404f22e6cec76b54dab5a29586c21472262d945a"
+    },
+    {
+      "path": "pi05/src/openpi/training/data_loader_test.py",
+      "size_bytes": 2487,
+      "sha256": "a0cf1bc293e8b0eb9bc5516bc05d30d9940da06c94fcfa1addca7267f8358815"
+    },
+    {
+      "path": "pi05/src/openpi/training/droid_rlds_dataset.py",
+      "size_bytes": 10886,
+      "sha256": "1d3a057a040240218234d7dbc6680ee4fe8256f2058531a0f9fefd3f7ee995a4"
+    },
+    {
+      "path": "pi05/src/openpi/training/memory_training.py",
+      "size_bytes": 208,
+      "sha256": "bc17fe9f002a94ab38a09f46caeff75b5baeef2a4a6edd20b3381b611105b14b"
+    },
+    {
+      "path": "pi05/src/openpi/training/misc/polaris_config.py",
+      "size_bytes": 9624,
+      "sha256": "c2ab92efb5e0d675118ad812113fd1c3cb3a62cdab5d79066014266e82df12d8"
+    },
+    {
+      "path": "pi05/src/openpi/training/misc/roboarena_config.py",
+      "size_bytes": 4804,
+      "sha256": "914acdd74ccf90e6779c13a5093140595c7e4f687e9e0b167439446a81d74c01"
+    },
+    {
+      "path": "pi05/src/openpi/training/optimizer.py",
+      "size_bytes": 3202,
+      "sha256": "cfb5d6f6aec61fbe961287b4015c652260628a0ed4682a0dbdc271b8e65431fa"
+    },
+    {
+      "path": "pi05/src/openpi/training/rmbench_memory_data.py",
+      "size_bytes": 10093,
+      "sha256": "14923395d2d9395e98f3b5571b25cfc1a2d7ddb7f792521bb5a6620952403387"
+    },
+    {
+      "path": "pi05/src/openpi/training/robomem_memory_data.py",
+      "size_bytes": 7135,
+      "sha256": "6fd26b38cfb373630611e1cf1c54ef4af2b4992321c225ca1d73cb2d5830c0ec"
+    },
+    {
+      "path": "pi05/src/openpi/training/robomem_memory_loader.py",
+      "size_bytes": 3911,
+      "sha256": "84528f678e54805b2f3819ec596595086f115b1971ea893cb7c1ae3279ed41fc"
+    },
+    {
+      "path": "pi05/src/openpi/training/sharding.py",
+      "size_bytes": 4099,
+      "sha256": "c67c4112a126ab097571c8a882fd1ee3f106b15b6e608e7205ac7f29b44ca985"
+    },
+    {
+      "path": "pi05/src/openpi/training/utils.py",
+      "size_bytes": 1217,
+      "sha256": "b684e674343f86ef5e976f279a586d3fe2f885e901f396ac899ed136abc7eca6"
+    },
+    {
+      "path": "pi05/src/openpi/training/weight_loaders.py",
+      "size_bytes": 3738,
+      "sha256": "a54ac3222f81f8c6233241bc16b6f369c9416cdaa83f7058f96f7ddcfde1a3b4"
+    },
+    {
+      "path": "pi05/src/openpi/transforms.py",
+      "size_bytes": 46152,
+      "sha256": "ed9673a1f140f188cb6a0657f6938ce69ea6783d9f42c3ed3e6b6b46c1fc7339"
+    },
+    {
+      "path": "pi05/uv.lock",
+      "size_bytes": 795910,
+      "sha256": "793488b5a55bb87200db90a61fd0af51922b686d94e1da4f4c587ab119b37d74"
+    },
+    {
+      "path": "pyproject.toml",
+      "size_bytes": 365,
+      "sha256": "e21069c1bb24a3846579748d05fdbef49dda21c2df1b0d4af0fb028be475ce41"
+    },
+    {
+      "path": "taskanchor/__init__.py",
+      "size_bytes": 82,
+      "sha256": "efe209ecb3e663212784c22fb1891bf52261c241e069d6d90b5a85f2557a6956"
+    },
+    {
+      "path": "taskanchor/__main__.py",
+      "size_bytes": 4924,
+      "sha256": "efc77b76efe98c00c291de317a6e9822d58f908744b75df1a569a4aa2a08d690"
+    },
+    {
+      "path": "taskanchor/configs/__init__.py",
+      "size_bytes": 61,
+      "sha256": "a9ecb5d1640dd38980236cdb84878df8c498c3c1d05b291faea19d9e77428f5f"
+    },
+    {
+      "path": "taskanchor/configs/robomemarena_pi05.json",
+      "size_bytes": 634,
+      "sha256": "fd7d0e623817c60446b85250dd5fc3c6bb36fbc20bf057a30ddef2277d451099"
+    },
+    {
+      "path": "taskanchor/data/__init__.py",
+      "size_bytes": 58,
+      "sha256": "f40c70863322679971b4c9eaa17158d1ddee26fe194d3f0e9864906a494f5cb0"
+    },
+    {
+      "path": "taskanchor/data/history.py",
+      "size_bytes": 3157,
+      "sha256": "f8d4c155eaf1d636c7596795b33f28d9fad606f7b7d4b6c421e595ac4a274511"
+    },
+    {
+      "path": "taskanchor/data/robomemarena_metadata.py",
+      "size_bytes": 6932,
+      "sha256": "d8026e5cd225a0c22edc05004d921298ed8ce5d88f94c4c0110c6ba76e2dde48"
+    },
+    {
+      "path": "taskanchor/paths.py",
+      "size_bytes": 1190,
+      "sha256": "b829fba7bef5fa450a2621c6d4758e9dac951af5482ae3d4b34552cc9c6b2760"
+    },
+    {
+      "path": "taskanchor/pi05/__init__.py",
+      "size_bytes": 58,
+      "sha256": "4fe691e5c85a023c697680cb0f7faa1fdf4e5e21feb3bbc5c061403836d075af"
+    },
+    {
+      "path": "taskanchor/pi05/taskanchor.py",
+      "size_bytes": 13666,
+      "sha256": "6543c6114add32135786c20db971246b18e66a183f1bfe134e2bee7b5cf3992e"
+    },
+    {
+      "path": "taskanchor/pi05/training.py",
+      "size_bytes": 9393,
+      "sha256": "668be3f93ed82290d66319eb67ba77e85a34119b20685a36d4a20abced2f84f7"
+    },
+    {
+      "path": "taskanchor/pi05/ttt.py",
+      "size_bytes": 6106,
+      "sha256": "7cc97cd49dce55cee1bd60a95e07f095b47bad5d28f19d7d823ae1864f5159d0"
+    },
+    {
+      "path": "tests/test_portability.py",
+      "size_bytes": 3601,
+      "sha256": "236c1119e7d4868c26a04fcf8b650723e61bd191e4925b87e6e9ec65283bc685"
+    },
+    {
+      "path": "tools/compare_robomem_reports.py",
+      "size_bytes": 3875,
+      "sha256": "9398c81f25f9f7ab3eb3f50d792c28ebc0ab136fd8fc841969d5cc309d70945b"
+    },
+    {
+      "path": "tools/prepare_benchmark.py",
+      "size_bytes": 2786,
+      "sha256": "5eb7efdbb55d851c49b49d4c336d4675d20c894fdd93d70d637d60a8300049d9"
+    }
+  ]
+}

code/assets/pi05_libero/norm_stats.json

@@ -0,0 +1,84 @@
+{
+  "norm_stats": {
+    "state": {
+      "mean": [
+        0.048380594700574875,
+        0.11262303590774536,
+        0.6594910025596619,
+        3.370551347732544,
+        -0.30113735795021057,
+        -0.24474778771400452,
+        0.02991946041584015,
+        -0.029522070661187172
+      ],
+      "std": [
+        0.11915058642625809,
+        0.2559613585472107,
+        0.1345514953136444,
+        0.6972067952156067,
+        0.75095134973526,
+        0.7611529231071472,
+        0.011234981939196587,
+        0.01164152380079031
+      ],
+      "q01": [
+        -0.2685799194574356,
+        -0.33133132414221766,
+        0.4330671876549721,
+        2.4110534167289734,
+        -2.4456029770374297,
+        -2.43879587931633,
+        0.0009645143877714873,
+        -0.04009675695095211
+      ],
+      "q99": [
+        0.23303380120396622,
+        0.5012647942483425,
+        1.0027180651426315,
+        4.715920171499253,
+        0.03176737618446346,
+        1.2033550313949584,
+        0.04003552807592787,
+        0.00003360225446522036
+      ]
+    },
+    "actions": {
+      "mean": [
+        0.010760927572846413,
+        0.02506864257156849,
+        0.01419717539101839,
+        0.0020873930770903826,
+        0.0005215497221797705,
+        -0.00006810902414144948,
+        0.031027423217892647
+      ],
+      "std": [
+        0.20257413387298584,
+        0.363967627286911,
+        0.30624639987945557,
+        0.12264969944953918,
+        0.11749287694692612,
+        0.08579856157302856,
+        0.9992039203643799
+      ],
+      "q01": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -0.5,
+        -0.5001200099349021,
+        -0.5003199966907501,
+        -1.0
+      ],
+      "q99": [
+        0.8684000000000001,
+        0.9996,
+        0.9996,
+        0.5,
+        0.4999599966883661,
+        0.4997600099325181,
+        0.9997999999999998
+      ]
+    }
+  }
+}

code/docs/PORTABILITY.md

@@ -0,0 +1,108 @@
+# Paths, assets, and the evaluation protocol
+
+The release resolves paths relative to the checked-out source or explicit user
+arguments. These changes do not change model dimensions, parameter keys,
+normalization numbers, prompts, sampling, optimization, or scoring.
+
+## Training and feature extraction
+
+Run from the release root. Existing `--dataset`, `--base-params`, `--cache`, and
+`--checkpoint-root` arguments override these defaults; training also accepts
+`--assets-dir` (the parent of `pi05_libero/norm_stats.json`). Environment
+variables are read at process startup.
+
+| Environment variable | Default relative to the release root |
+| --- | --- |
+| `TASKANCHOR_ROBOMEM_DATASET` | `data/robomemarena_pi05_task_all_with_subtask_seed` |
+| `TASKANCHOR_BASE_PARAMS` | `pretrained/pi05_base/params` |
+| `TASKANCHOR_FEATURE_CACHE` | `data/robomem_siglip_base_main_spatial_v3` |
+| `TASKANCHOR_ASSETS_DIR` | `assets` |
+| `TASKANCHOR_CHECKPOINT_ROOT` | `checkpoints` |
+
+The packaged normalization file is the original `pi05_libero` file. Do not
+replace it with newly computed statistics when reproducing the provided model.
+During inference, normalization is loaded directly from the selected
+checkpoint's `assets/pi05_libero/norm_stats.json`; training-server paths are
+not consulted.
+
+Training requires the LeRobot Parquet dataset with embedded RGB `image` and
+`wrist_image`, `state`, `actions`, and metadata in `meta/episodes_stats.jsonl`.
+The metadata must preserve `task_id`, `subtask_id`, and `seed_rank`, so subtasks
+can be joined into the original full-task history. The existing loader does
+not require an additional predecoded observation cache. Feature extraction
+uses frozen base-model SigLIP weights and writes the main-camera 4×4 cache.
+The training dataset and feature cache are not required for evaluation.
+
+Feature-cache provenance intentionally retains its original strict check,
+including absolute dataset/base-parameter locations and metadata SHA-256.
+If moving a preexisting cache, extraction/training rejects a path-mismatched
+manifest. The supported portable route is to extract a new cache at the new
+location; do not silently rewrite provenance or relabel different features.
+Fresh training is not guaranteed to reproduce the checkpoint bit for bit,
+including the recorded restart at step 10,000.
+
+## Offline inference tokenizer
+
+`TASKANCHOR_TOKENIZER_PATH=/absolute/path/to/paligemma_tokenizer.model` selects
+the provided SentencePiece file. It changes the file location only. Without
+that setting, the original `gs://big_vision/paligemma_tokenizer.model` source is
+downloaded through OpenPI, whose cache can be relocated with `OPENPI_DATA_HOME`.
+The public full checkpoint contains the backbone and adapter; inference does
+not require downloading an additional `pi05_base` checkpoint.
+
+## Benchmark installation and paths
+
+Install RoboMemArena and its simulator dependencies separately, using commit
+`d09f63b671f362985e1e0d038d2014c2d2519488`. The release does not vendor the
+benchmark, MuJoCo/OSMesa libraries, or visual assets. Its evaluation process
+must import the pinned checkout's `libero_fork` and the official scorer.
+
+After checking out that commit, generate configuration without editing the
+official scorer:
+
+```bash
+python tools/prepare_benchmark.py \
+  --benchmark-root /path/to/RoboMemArena \
+  --config-dir /path/to/robomem_libero_config
+
+export TASKANCHOR_ROBOMEM_BENCHMARK=/path/to/RoboMemArena/evaluation_benchmark
+export TASKANCHOR_ROBOMEM_EVAL_PYTHON=/path/to/simulator-env/bin/python
+export TASKANCHOR_LIBERO_CONFIG=/path/to/robomem_libero_config
+export TASKANCHOR_PI05_PYTHON=/path/to/pi05-env/bin/python
+```
+
+The setup utility verifies the pinned Git revision and unchanged tracked
+benchmark files, then writes `OFFICIAL_SOURCE.json` and a local LIBERO
+`config.yaml`. It requires `libero_fork/libero/assets` and `bddl_files`.
+Evaluation uses seeded resets and does not read LIBERO stored `init_files`.
+Keep the simulator virtual environment's executable path; resolving its
+symlink to a base Python loses that environment's packages.
+
+OSMesa remains the renderer. If OSMesa libraries are installed outside system
+library paths, set `TASKANCHOR_RENDER_LIBS` to their directory. The original
+runtime used Mesa libraries under its `robomem_render_libs/usr/lib/x86_64-linux-gnu`
+directory; these operating-system binaries are not part of the code package.
+
+The original Task25 candidate list is retained in
+`evidence/task25_seed_candidates.json`. Only initialization failures before any
+policy action may advance to the next candidate. Policy failures are scored.
+Both main and wrist observations are still passed to the VLA; only the main
+view receives the TaskAnchor residual.
+
+## Preserved numerical settings
+
+The recorded 50k model uses separate permanent frame0, 191 history frames,
+16 tokens per frame, stride 10, action chunk 20, execution length 10, and
+nearest-stage predicted progress. Evaluation uses all 26 tasks and the original 2500-step budget.
+The evaluator retains the original trial and seed schedule; configure and record
+any different study protocol before collecting results. Paths do not select a
+different memory protocol or change these values.
+
+CPU portability checks can be run without JAX or simulator installation:
+
+```bash
+python -m unittest discover -s tests -p test_portability.py
+```
+
+These checks validate path relocation and benchmark setup only. They do not
+replace closed-loop benchmark evaluation.

code/docs/PUBLISHING_ZH.md

@@ -0,0 +1,46 @@
+# 发布和后续维护
+
+Hugging Face 目标：<https://huggingface.co/henry13522/Taskanchor>。
+代码许可证为 Apache-2.0；完整权重保留 π0.5 / PaliGemma / Gemma 上游条款。模型卡明确区分代码和权重许可。
+
+## 本次目录
+
+`model/` 是 Hugging Face 仓库的上传根目录：含 `50000/`、`recipe.json`、tokenizer、模型卡、校验清单和 `code/`。`code/` 同时作为独立 GitHub 源码仓库发布目录。HF 下载后直接按 `code/docs/REPRODUCING.md` 使用。
+
+只发布 50k 完整推理 checkpoint；不包含 optimizer/train_state、训练数据、特征缓存、机器环境、凭据或训练日志。不要将原始工作区整体推送到公开仓库。
+
+## Hugging Face 上传
+
+在终端安全登录有目标仓库写权限的 Hugging Face 账号：
+
+```bash
+hf auth login
+```
+
+使用官方 Hugging Face CLI 上传已检查的模型目录（目录内包含 `code/`）：
+
+```bash
+hf upload henry13522/Taskanchor /path/to/release/model .
+```
+
+上传使用 Hugging Face 官方 SDK/Xet，无需转换 Orbax checkpoint。若使用 Git 路线，应先安装 Git Xet，并保留发布包 `.gitattributes` 对**无扩展名** Orbax 分片的追踪规则，不能只追踪 `*.safetensors`。
+
+官方文档：[上传文件](https://huggingface.co/docs/huggingface_hub/guides/upload)、[Git Xet](https://huggingface.co/docs/hub/git-xet)。
+
+## GitHub 源码仓库
+
+账号/组织确定后，在 `code/` 内执行以下命令。将 `OWNER` 替换为实际账号；不要在 `model/` 内创建 GitHub 源码仓库。
+
+```bash
+cd /path/to/release/code
+git init -b main
+git add .
+git commit -m 'Release TaskAnchor RoboMemArena pi05 training and evaluation'
+gh repo create OWNER/TaskAnchor --public --source=. --remote=origin --push
+```
+
+在 GitHub About 中链接 Hugging Face 权重；在模型卡中添加 GitHub 仓库链接。后续对外报告结果时记录代码 commit、模型仓库 revision、benchmark commit 和完整评测命令。
+
+## 本次验证范围
+
+本次执行 CPU 权重恢复、有限性和文件校验、路径测试。公开包不提供历史小批次评测指标或复现结论；评测次数由正式研究协议确定；现有入口保留原试验次数和种子调度，数量不同时须先调整协议实现。原始记录仅保留在上传目录之外的本地审计目录。

code/docs/REPRODUCING.md

@@ -0,0 +1,283 @@
+# Reproducing the RoboMemArena 50k model
+
+This guide uses the downloaded Hugging Face repository, which contains both
+source and complete inference weights:
+
+```text
+Taskanchor/
+├── code/                         # source directory
+│   ├── taskanchor/
+│   ├── pi05/
+│   └── assets/pi05_libero/norm_stats.json
+├── recipe.json
+├── tokenizer/paligemma_tokenizer.model
+└── 50000/
+    ├── params/                   # complete Orbax parameter directory
+    └── assets/
+        ├── robomem_memory.json
+        └── pi05_libero/norm_stats.json
+```
+
+Set these absolute paths for the commands below:
+
+```bash
+export TASKANCHOR_MODEL=/absolute/path/to/Taskanchor
+export TASKANCHOR_CODE="$TASKANCHOR_MODEL/code"
+export TASKANCHOR_BENCHMARK_REPO=/absolute/path/to/RoboMemArena
+export TASKANCHOR_SIM_ENV=/absolute/path/to/robomem-sim-env
+cd "$TASKANCHOR_CODE"
+```
+
+## 1. Policy environment: Python 3.11 and JAX
+
+Use the checked-in `pi05/uv.lock`; do not replace it with an unpinned OpenPI
+installation. This environment runs feature extraction, training, model
+servers, and the evaluation scheduler. The simulator runs in a second
+environment. Install `uv` first if it is unavailable.
+
+```bash
+cd "$TASKANCHOR_CODE/pi05"
+uv python install 3.11
+uv sync --frozen --python 3.11
+uv pip install --python .venv/bin/python --no-deps -e "$TASKANCHOR_CODE"
+
+export TASKANCHOR_PI05_PYTHON="$TASKANCHOR_CODE/pi05/.venv/bin/python"
+export TASKANCHOR_TOKENIZER_PATH="$TASKANCHOR_MODEL/tokenizer/paligemma_tokenizer.model"
+cd "$TASKANCHOR_CODE"
+```
+
+The measured policy stack used Python 3.11.13, JAX/JAXlib 0.5.3, Flax 0.10.2,
+Optax 0.2.4, and Orbax 0.11.13. A compatible NVIDIA driver is required for the
+CUDA 12 JAX build. Do not force `JAX_PLATFORMS=cpu` for training or inference.
+The local tokenizer setting avoids a network download during model startup.
+The full 50k checkpoint already contains backbone weights.
+
+## 2. Simulator environment: Python 3.10 and MuJoCo
+
+Check out the exact benchmark revision. Use a fresh checkout or a clean
+worktree; `tools/prepare_benchmark.py` rejects modified tracked benchmark code.
+
+```bash
+git clone https://github.com/OpenHelix-Team/RoboMemArena.git "$TASKANCHOR_BENCHMARK_REPO"
+git -C "$TASKANCHOR_BENCHMARK_REPO" checkout d09f63b671f362985e1e0d038d2014c2d2519488
+
+uv python install 3.10
+uv venv --python 3.10 "$TASKANCHOR_SIM_ENV"
+uv pip install --python "$TASKANCHOR_SIM_ENV/bin/python" \
+  'numpy==1.26.4' 'mujoco==2.3.7' 'robosuite==1.4.1' \
+  'scipy==1.15.3' 'numba==0.68.0' 'llvmlite==0.50.0' \
+  'torch==2.4.1' 'torchvision==0.19.1' \
+  'opencv-python==4.11.0.86' 'Pillow==11.3.0' \
+  'imageio==2.34.2' 'imageio-ffmpeg==0.6.0' \
+  'PyOpenGL==3.1.10' 'PyYAML==6.0.3' 'h5py==3.16.0' \
+  'easydict==1.13' 'gym==0.26.2' 'bddl==3.6.0' \
+  'termcolor==3.3.0' 'cloudpickle==3.1.2' 'matplotlib==3.10.9' \
+  'websockets==16.1.1' 'msgpack==1.2.3' tqdm
+```
+
+These pins are read from the working evaluation environment. That environment
+inherited system packages; the clean-environment installation above has not
+been replayed as part of this release. Use the pinned benchmark's own setup
+instructions for any platform-specific missing dependency. Do not substitute
+a different MuJoCo/robosuite pair within a controlled evaluation comparison.
+
+The measured renderer is OSMesa. On Debian/Ubuntu, install the system OSMesa
+runtime if it is absent:
+
+```bash
+sudo apt-get update
+sudo apt-get install -y libosmesa6 libgl1 libegl1 ffmpeg
+```
+
+The benchmark must contain its meshes/textures and BDDL files. Follow its
+asset instructions if Git checkout alone does not supply these files.
+Create local LIBERO paths without importing its interactive setup:
+
+```bash
+"$TASKANCHOR_PI05_PYTHON" tools/prepare_benchmark.py \
+  --benchmark-root "$TASKANCHOR_BENCHMARK_REPO" \
+  --config-dir "$TASKANCHOR_CODE/local/robomem_libero_config"
+
+export TASKANCHOR_ROBOMEM_BENCHMARK="$TASKANCHOR_BENCHMARK_REPO/evaluation_benchmark"
+export TASKANCHOR_ROBOMEM_EVAL_PYTHON="$TASKANCHOR_SIM_ENV/bin/python"
+export TASKANCHOR_LIBERO_CONFIG="$TASKANCHOR_CODE/local/robomem_libero_config"
+```
+
+If OSMesa is outside the system linker paths, additionally set
+`TASKANCHOR_RENDER_LIBS=/absolute/path/to/osmesa/lib`. Keep the simulator
+virtualenv executable path as written; do not resolve its symlink to base
+Python. The launcher sets `MUJOCO_GL=osmesa`, `PYOPENGL_PLATFORM=osmesa`, and
+the local `libero_fork` import path for each simulator subprocess.
+
+## 3. Evaluate the released checkpoint
+
+The unified command starts its own policy servers and terminates only those
+servers when evaluation ends. Do not start a separate server first. The
+`--python` option below is the policy interpreter; the simulator interpreter
+comes from `TASKANCHOR_ROBOMEM_EVAL_PYTHON`.
+
+The evaluator preserves its original trial-count and seed schedule. A study
+requiring a different schedule must configure and validate that protocol before
+collecting results. No benchmark score is provided by this release.
+
+First inspect the command without allocating GPUs:
+
+```bash
+"$TASKANCHOR_PI05_PYTHON" -m taskanchor \
+  --benchmark robomemarena --model pi05 --stage eval \
+  --benchmark-root "$TASKANCHOR_ROBOMEM_BENCHMARK" \
+  --python "$TASKANCHOR_PI05_PYTHON" --dry-run -- \
+  --checkpoint "$TASKANCHOR_MODEL/50000" \
+  --output "$TASKANCHOR_CODE/reports/robomem50k_evaluation" \
+  --workers 8 --base-port 8800
+```
+
+Remove `--dry-run` to execute the supported task/trial configuration. The output directory
+must not already exist. `--workers 8` uses GPUs 0–7, with one model server per
+GPU. `--workers 1` uses GPU 0 and retains the same task/seed coverage, taking
+longer. The current scheduler assigns physical GPU indices itself; an outer
+`CUDA_VISIBLE_DEVICES` remapping is not an arbitrary GPU-selection interface.
+
+Inspect `final_report.json`, `status.json`, and
+per-task logs for complete coverage and runtime failures. Per-task folders
+contain benchmark summaries/videos and a `progress.jsonl` trace containing
+both raw and discretized predictions. Initialization failures are recorded;
+Task25 may advance its recorded candidate list only before any policy action.
+Failed policies are not discarded or retried as successful replacements.
+
+## 4. Exact evaluation contract
+
+| Setting | Released 50k protocol |
+| --- | --- |
+| Backbone | π0.5, Gemma 2B VLM / Gemma 300M action expert |
+| Observations | Main and wrist views; 224×224 model inputs |
+| Adapter insertion | Main view only; `current + offset`, 1152→512→1152 |
+| Temporal module | 8 layers, width 512, 8 heads; dropout 0.10, DropPath 0 |
+| History frame | 4×4 = 16 SigLIP tokens |
+| Memory protocol | Separate permanent frame0 + up to 191 history frames |
+| History stride | 10 environment steps |
+| Action chunk / replan | Predict 20; execute 10, then replan |
+| Progress | Continuous prediction mapped to nearest valid task stage |
+| Language | Fixed full-task rewrite 0 + `Progress:%.2f` |
+| Task coverage | Tasks 1–26; trial count and accepted seeds must be checked against the research protocol |
+| Episode budget | 2500 steps |
+| Rendering | OSMesa |
+
+There is no oracle progress input. For a four-stage task, conditioning values
+are 0.00, 0.25, 0.50, and 0.75. The separate frame0 is part of this model's
+training and evaluation contract; replacing it with FIFO frame0 changes the
+experiment even if the weight keys still load.
+
+This package provides runnable model and evaluation code. It does not publish
+benchmark scores or assert that a new evaluation with the sample size required
+by a paper has been completed. Record and report the actual task/seed coverage
+of each evaluation run.
+
+## 5. Extract training features
+
+Training starts from the converted LeRobot dataset, not the public 50k
+checkpoint. The dataset must contain embedded main/wrist RGB images, real
+states/actions, and the original subtask/seed metadata. This package does not
+include the raw-HDF5-to-LeRobot converter or the dataset. Preserve the original
+conversion and ordering when reproducing training. The reference manifest has
+2600 full-task groups, 2,811,755 frames, and a maximum 1918 frames per full task.
+Evaluation does not need any of these training files.
+
+Provide the original pretrained `pi05_base/params` and data locations:
+
+```bash
+export TASKANCHOR_ROBOMEM_DATASET=/absolute/path/to/robomemarena_pi05_task_all_with_subtask_seed
+export TASKANCHOR_BASE_PARAMS=/absolute/path/to/pi05_base/params
+export TASKANCHOR_FEATURE_CACHE=/absolute/path/to/new_robomem_spatial_cache
+export TASKANCHOR_CHECKPOINT_ROOT=/absolute/path/to/new_training_checkpoints
+
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+"$TASKANCHOR_PI05_PYTHON" -m taskanchor \
+  --benchmark robomemarena --model pi05 --stage extract \
+  --python "$TASKANCHOR_PI05_PYTHON" -- \
+  --dataset "$TASKANCHOR_ROBOMEM_DATASET" \
+  --base-params "$TASKANCHOR_BASE_PARAMS" \
+  --cache "$TASKANCHOR_FEATURE_CACHE" --batch-size 256 --workers 8
+```
+
+Extraction validates every completed feature file and writes `VERIFIED.json`.
+It stores FP16 values derived from BF16 main-camera SigLIP features, using the
+recorded 16×16→4×4 bilinear pooling without antialiasing. Existing cache
+provenance is checked strictly, including source paths; see [PORTABILITY](PORTABILITY.md)
+before relocating an old cache.
+
+## 6. Train one run with the 5k freeze prefix
+
+Use a new run name. The schedule is 55,000 total updates, with 5,000
+adapter-only updates followed by 50,000 joint updates. The released checkpoint
+is taken at update 50,000: 5,000 adapter-only + 45,000 joint updates. Setting
+`--steps 50000` would change the learning-rate schedule and is not the recorded
+experiment.
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+"$TASKANCHOR_PI05_PYTHON" -m taskanchor \
+  --benchmark robomemarena --model pi05 --stage train \
+  --python "$TASKANCHOR_PI05_PYTHON" -- \
+  --dataset "$TASKANCHOR_ROBOMEM_DATASET" \
+  --base-params "$TASKANCHOR_BASE_PARAMS" \
+  --cache "$TASKANCHOR_FEATURE_CACHE" \
+  --assets-dir "$TASKANCHOR_CODE/assets" \
+  --checkpoint-root "$TASKANCHOR_CHECKPOINT_ROOT" \
+  --run-name robomem50k_reproduction \
+  --steps 55000 --freeze-steps 5000 --batch-size 256 \
+  --action-horizon 20 --execution-horizon 10 \
+  --fsdp-devices 8 --workers 8 --save-interval 1000 --log-interval 20
+```
+
+The measured run restarted after a complete 10,000-step checkpoint to enable
+the shared-memory loader. To mirror that history, retain the **55k schedule**
+for the first leg, stop the first process after that checkpoint is complete,
+then rerun the exact command above with:
+
+```text
+--resume --resume-step 10000 --shared-memory-loader --device-prefetch 2 --prefetch-factor 2
+```
+
+Do not launch the resumed command while the original process is still running.
+The recorded restart arguments are retained in
+`evidence/pipeline_restart_10000.json`. The sampler restarts with seed
+`42 + restored_step`; it does not restore an exact input cursor. A fresh run
+therefore is not promised to reproduce identical parameter bytes.
+
+The training output is
+`$TASKANCHOR_CHECKPOINT_ROOT/pi05_robomemarena_memory_v5_main/robomem50k_reproduction/`.
+It contains the parent `recipe.json` and numbered checkpoints. Training saves
+every 1000 updates and retains 5000-step milestones under the original manager
+policy. Keep optimizer/train-state files for resuming; the public inference
+package intentionally omits them and cannot resume the original 50k optimizer.
+
+| Training setting | Value |
+| --- | --- |
+| Global batch / FSDP devices | 256 / 8 |
+| Visual encoder | Frozen throughout; image projection head joins after 5k |
+| Optimizer | AdamW, β=(0.9, 0.95), global gradient clip 1.0 |
+| Adapter learning rate | 2.5e-5 for first 5k; cosine decay afterward |
+| Backbone / image-head LR after unfreezing | 2.5e-5 / 1e-5, cosine decay |
+| Final cosine multiplier | 0.1 |
+| Weight decay | Adapter 0.05; backbone/image head 0 |
+| Progress target | `round(subtask_id / number_of_subtasks, 2)` |
+| Loss | Flow action loss + 0.01 × (100 × progress MSE + 1 − Pearson) |
+
+Frozen backbone Adam states do not advance during the first 5000 updates.
+Neither LoRA nor a separate planner is introduced. Numerical implementation
+and optimizer groups are in `taskanchor/pi05/taskanchor.py` and
+`taskanchor/pi05/training.py`.
+
+## 7. Validation scope
+
+```bash
+cd "$TASKANCHOR_CODE"
+python -m unittest discover -s tests -p test_portability.py
+```
+
+The packaging checks passed source parsing, CPU imports, path tests, and
+unchanged numerical source checks. Adapter/TTT/training-loop/history core
+files remain byte-identical to the previously tested release. Data sampling,
+online history, action sampling, and extraction computations are unchanged.
+No new training or GPU rollout was started during this packaging pass.

code/evidence/environment_pi05.json

@@ -0,0 +1,19 @@
+{
+  "python": "3.11.13",
+  "jax": "0.5.3",
+  "jaxlib": "0.5.3",
+  "flax": "0.10.2",
+  "optax": "0.2.4",
+  "orbax-checkpoint": "0.11.13",
+  "numpy": "1.26.4",
+  "torch": "2.7.1",
+  "transformers": "4.53.2",
+  "lerobot": "0.1.0",
+  "pyarrow": "20.0.0",
+  "Pillow": "11.2.1",
+  "sentencepiece": "0.2.0",
+  "ruff": "0.11.12",
+  "pytest": "8.3.5",
+  "openpi": "0.1.0",
+  "openpi-client": "0.1.0"
+}

code/evidence/packaging_checks.json

@@ -0,0 +1,17 @@
+{
+  "ast_files": 107,
+  "portability_tests": {
+    "passed": 3
+  },
+  "adapter_and_history_cpu_tests": {
+    "passed": 6
+  },
+  "no_gpu_or_rollouts_in_this_packaging_pass": true,
+  "numerical_model_logic_modified": false,
+  "changes": [
+    "portable asset and data paths",
+    "documentation and license notices",
+    "packaging"
+  ],
+  "training_data_requirement": "Previously converted LeRobot dataset with subtask/seed metadata; raw conversion not included."
+}

code/evidence/pipeline_restart_10000.json

@@ -0,0 +1,37 @@
+{
+  "command": [
+    "<OPENPI>/.venv/bin/python",
+    "-u",
+    "scripts/train_robomem_memory.py",
+    "--run-name",
+    "main_v5_spatial4x4_h191_a20_exec10_bs256_5k50k",
+    "--batch-size",
+    "256",
+    "--steps",
+    "55000",
+    "--freeze-steps",
+    "5000",
+    "--action-horizon",
+    "20",
+    "--execution-horizon",
+    "10",
+    "--fsdp-devices",
+    "8",
+    "--workers",
+    "8",
+    "--resume",
+    "--resume-step",
+    "10000",
+    "--shared-memory-loader",
+    "--device-prefetch",
+    "2",
+    "--prefetch-factor",
+    "2"
+  ],
+  "source_sha256": {
+    "scripts/train_robomem_memory.py": "218d6f44812d0043438b3c515bf8ca4d8e1fdd34a86dac1d6431398fb7aef9fb",
+    "src/openpi/training/robomem_memory_loader.py": "8924b7cb7146ae8f5a46230050f18938fc83338ab5382efb5cbe5e2ad6b9d038",
+    "scripts/test_robomem_memory_loader.py": "80f421f909c0d3373171dccd013c46fbf8b2b9a67597b81a6198e6163cdd2ba7",
+    "scripts/resume_robomem_pipeline_at_10k.py": "0d9dbfa45ce3a279878a203ff8ca147ff1189dfdaf166e137ac27459951165ea"
+  }
+}

code/evidence/robomem50k_adapter_numerical_parity.json

@@ -0,0 +1,24 @@
+{
+  "checkpoint": "<RUN>/50000/params",
+  "adapter_parameters": 26930049,
+  "adapter_arrays_sha256": "52573dc7d0ac6fb4373382bed7430ca1c0f30af1474cceab42c785ab6bae8a10",
+  "cases": [
+    {
+      "valid_history_frames": 0,
+      "refined_exact": true,
+      "progress_exact": true
+    },
+    {
+      "valid_history_frames": 17,
+      "refined_exact": true,
+      "progress_exact": true
+    },
+    {
+      "valid_history_frames": 191,
+      "refined_exact": true,
+      "progress_exact": true
+    }
+  ],
+  "device": "CPU",
+  "limitation": "Synthetic visual features; full VLA action inference and closed-loop episodes not rerun"
+}

code/evidence/task25_seed_candidates.json

@@ -0,0 +1,208 @@
+{
+  "candidates": [
+    50,
+    51,
+    52,
+    53,
+    54,
+    55,
+    56,
+    57,
+    58,
+    59,
+    60,
+    61,
+    62,
+    63,
+    64,
+    65,
+    66,
+    67,
+    68,
+    69,
+    70,
+    71,
+    72,
+    73,
+    74,
+    75,
+    76,
+    77,
+    78,
+    79,
+    80,
+    81,
+    82,
+    83,
+    84,
+    85,
+    86,
+    87,
+    88,
+    89,
+    90,
+    91,
+    92,
+    93,
+    94,
+    95,
+    96,
+    97,
+    98,
+    99,
+    216,
+    241,
+    247,
+    248,
+    250,
+    252,
+    253,
+    257,
+    258,
+    262,
+    265,
+    266,
+    267,
+    269,
+    271,
+    279,
+    280,
+    283,
+    289,
+    290,
+    293,
+    294,
+    295,
+    296,
+    300,
+    303,
+    308,
+    310,
+    313,
+    315,
+    316,
+    318,
+    322,
+    325,
+    326,
+    328,
+    329,
+    332,
+    333,
+    335,
+    336,
+    337,
+    338,
+    339,
+    343,
+    344,
+    345,
+    346,
+    350,
+    352,
+    360,
+    361,
+    362,
+    364,
+    367,
+    368,
+    369,
+    370,
+    371,
+    372,
+    373,
+    374,
+    375,
+    377,
+    378,
+    381,
+    382,
+    383,
+    384,
+    385,
+    386,
+    387,
+    388,
+    389,
+    390,
+    391,
+    392,
+    393,
+    394,
+    395,
+    396,
+    397,
+    398,
+    399,
+    401,
+    402,
+    403,
+    404,
+    405,
+    406,
+    407,
+    408,
+    409,
+    410,
+    411,
+    413,
+    414,
+    415,
+    416,
+    418,
+    419,
+    420,
+    421,
+    422,
+    423,
+    424,
+    425,
+    426,
+    427,
+    428,
+    429,
+    430,
+    431,
+    432,
+    433,
+    434,
+    435,
+    436,
+    437,
+    438,
+    440,
+    442,
+    446,
+    447,
+    448,
+    449,
+    450,
+    451,
+    453,
+    454,
+    455,
+    457,
+    458,
+    459,
+    460,
+    461,
+    462,
+    463,
+    465,
+    466,
+    468,
+    469,
+    470,
+    472,
+    473,
+    474,
+    475,
+    476,
+    477,
+    478
+  ],
+  "max_attempts": 200,
+  "selection": "ascending from 50, excluding all training seed_rank values",
+  "training_seed_min": 100,
+  "training_seed_max": 813
+}

code/pi05/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.

code/pi05/LICENSE_GEMMA.txt

@@ -0,0 +1,77 @@
+Gemma Terms of Use 
+
+Last modified: February 21, 2024
+
+By using, reproducing, modifying, distributing, performing or displaying any portion or element of Gemma, Model Derivatives including via any Hosted Service, (each as defined below) (collectively, the "Gemma Services") or otherwise accepting the terms of this Agreement, you agree to be bound by this Agreement.
+
+Section 1: DEFINITIONS
+1.1 Definitions
+(a) "Agreement" or "Gemma Terms of Use" means these terms and conditions that govern the use, reproduction, Distribution or modification of the Gemma Services and any terms and conditions incorporated by reference.
+
+(b) "Distribution" or "Distribute" means any transmission, publication, or other sharing of Gemma or Model Derivatives to a third party, including by providing or making Gemma or its functionality available as a hosted service via API, web access, or any other electronic or remote means ("Hosted Service").
+
+(c) "Gemma" means the set of machine learning language models, trained model weights and parameters identified at ai.google.dev/gemma, regardless of the source that you obtained it from.
+
+(d) "Google" means Google LLC.
+
+(e) "Model Derivatives" means all (i) modifications to Gemma, (ii) works based on Gemma, or (iii) any other machine learning model which is created by transfer of patterns of the weights, parameters, operations, or Output of Gemma, to that model in order to cause that model to perform similarly to Gemma, including distillation methods that use intermediate data representations or methods based on the generation of synthetic data Outputs by Gemma for training that model. For clarity, Outputs are not deemed Model Derivatives.
+
+(f) "Output" means the information content output of Gemma or a Model Derivative that results from operating or otherwise using Gemma or the Model Derivative, including via a Hosted Service.
+
+1.2
+As used in this Agreement, "including" means "including without limitation".
+
+Section 2: ELIGIBILITY AND USAGE
+2.1 Eligibility
+You represent and warrant that you have the legal capacity to enter into this Agreement (including being of sufficient age of consent). If you are accessing or using any of the Gemma Services for or on behalf of a legal entity, (a) you are entering into this Agreement on behalf of yourself and that legal entity, (b) you represent and warrant that you have the authority to act on behalf of and bind that entity to this Agreement and (c) references to "you" or "your" in the remainder of this Agreement refers to both you (as an individual) and that entity.
+
+2.2 Use
+You may use, reproduce, modify, Distribute, perform or display any of the Gemma Services only in accordance with the terms of this Agreement, and must not violate (or encourage or permit anyone else to violate) any term of this Agreement.
+
+Section 3: DISTRIBUTION AND RESTRICTIONS
+3.1 Distribution and Redistribution
+You may reproduce or Distribute copies of Gemma or Model Derivatives if you meet all of the following conditions:
+
+You must include the use restrictions referenced in Section 3.2 as an enforceable provision in any agreement (e.g., license agreement, terms of use, etc.) governing the use and/or distribution of Gemma or Model Derivatives and you must provide notice to subsequent users you Distribute to that Gemma or Model Derivatives are subject to the use restrictions in Section 3.2.
+You must provide all third party recipients of Gemma or Model Derivatives a copy of this Agreement.
+You must cause any modified files to carry prominent notices stating that you modified the files.
+All Distributions (other than through a Hosted Service) must be accompanied by a "Notice" text file that contains the following notice: "Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms".
+You may add your own intellectual property statement to your modifications and, except as set forth in this Section, may provide additional or different terms and conditions for use, reproduction, or Distribution of your modifications, or for any such Model Derivatives as a whole, provided your use, reproduction, modification, Distribution, performance, and display of Gemma otherwise complies with the terms and conditions of this Agreement. Any additional or different terms and conditions you impose must not conflict with the terms of this Agreement.
+
+3.2 Use Restrictions
+You must not use any of the Gemma Services:
+
+for the restricted uses set forth in the Gemma Prohibited Use Policy at ai.google.dev/gemma/prohibited_use_policy ("Prohibited Use Policy"), which is hereby incorporated by reference into this Agreement; or
+in violation of applicable laws and regulations.
+To the maximum extent permitted by law, Google reserves the right to restrict (remotely or otherwise) usage of any of the Gemma Services that Google reasonably believes are in violation of this Agreement.
+
+3.3 Generated Output
+Google claims no rights in Outputs you generate using Gemma. You and your users are solely responsible for Outputs and their subsequent uses.
+
+Section 4: ADDITIONAL PROVISIONS
+4.1 Updates
+Google may update Gemma from time to time, and you must make reasonable efforts to use the latest version of Gemma.
+
+4.2 Trademarks
+Nothing in this Agreement grants you any rights to use Google's trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between you and Google. Google reserves any rights not expressly granted herein.
+
+4.3 DISCLAIMER OF WARRANTY
+UNLESS REQUIRED BY APPLICABLE LAW, THE GEMMA SERVICES, AND OUTPUTS, ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING 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, REPRODUCING, MODIFYING, PERFORMING, DISPLAYING OR OR DISTRIBUTING ANY OF THE GEMMA SERVICES OR OUTPUTS AND ASSUME ANY AND ALL RISKS ASSOCIATED WITH YOUR USE OR DISTRIBUTION OF ANY OF THE GEMMA SERVICES OR OUTPUTS AND YOUR EXERCISE OF RIGHTS AND PERMISSIONS UNDER THIS AGREEMENT.
+
+4.4 LIMITATION OF LIABILITY
+TO THE FULLEST EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), PRODUCT LIABILITY, CONTRACT, OR OTHERWISE, UNLESS REQUIRED BY APPLICABLE LAW, SHALL GOOGLE OR ITS AFFILIATES BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, CONSEQUENTIAL, OR PUNITIVE DAMAGES, OR LOST PROFITS OF ANY KIND ARISING FROM THIS AGREEMENT OR RELATED TO, ANY OF THE GEMMA SERVICES OR OUTPUTS EVEN IF GOOGLE OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+
+4.5 Term, Termination, and Survival
+The term of this Agreement will commence upon your acceptance of this Agreement (including acceptance by your use, modification, or Distribution, reproduction, performance or display of any portion or element of the Gemma Services) and will continue in full force and effect until terminated in accordance with the terms of this Agreement. Google may terminate this Agreement if you are in breach of any term of this Agreement. Upon termination of this Agreement, you must delete and cease use and Distribution of all copies of Gemma and Model Derivatives in your possession or control. Sections 1, 2.1, 3.3, 4.2 to 4.9 shall survive the termination of this Agreement.
+
+4.6 Governing Law and Jurisdiction
+This Agreement will be governed by the laws of the State of California without regard to choice of law principles. The UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The state and federal courts of Santa Clara County, California shall have exclusive jurisdiction of any dispute arising out of this Agreement.
+
+4.7 Severability
+If any provision of this Agreement is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein.
+
+4.8 Entire Agreement
+This Agreement states all the terms agreed between the parties and supersedes all other agreements between the parties as of the date of acceptance relating to its subject matter.
+
+4.9 No Waiver
+Google will not be treated as having waived any rights by not exercising (or delaying the exercise of) any rights under this Agreement.

code/pi05/README.md

@@ -0,0 +1,323 @@
+# openpi
+
+openpi holds open-source models and packages for robotics, published by the [Physical Intelligence team](https://www.physicalintelligence.company/).
+
+Currently, this repo contains three types of models:
+- the [π₀ model](https://www.physicalintelligence.company/blog/pi0), a flow-based vision-language-action model (VLA).
+- the [π₀-FAST model](https://www.physicalintelligence.company/research/fast), an autoregressive VLA, based on the FAST action tokenizer.
+- the [π₀.₅ model](https://www.physicalintelligence.company/blog/pi05), an upgraded version of π₀ with better open-world generalization trained with [knowledge insulation](https://www.physicalintelligence.company/research/knowledge_insulation). Note that, in this repository, we currently only support the flow matching head for both $\pi_{0.5}$ training and inference.
+
+For all models, we provide _base model_ checkpoints, pre-trained on 10k+ hours of robot data, and examples for using them out of the box or fine-tuning them to your own datasets.
+
+This is an experiment: $\pi_0$ was developed for our own robots, which differ from the widely used platforms such as [ALOHA](https://tonyzhaozh.github.io/aloha/) and [DROID](https://droid-dataset.github.io/), and though we are optimistic that researchers and practitioners will be able to run creative new experiments adapting $\pi_0$ to their own platforms, we do not expect every such attempt to be successful. All this is to say: $\pi_0$ may or may not work for you, but you are welcome to try it and see!
+
+## Updates
+
+- [Sept 2025] We released PyTorch support in openpi.
+- [Sept 2025] We released pi05, an upgraded version of pi0 with better open-world generalization.
+- [Sept 2025]: We have added an [improved idle filter](examples/droid/README_train.md#data-filtering) for DROID training.
+- [Jun 2025]: We have added [instructions](examples/droid/README_train.md) for using `openpi` to train VLAs on the full [DROID dataset](https://droid-dataset.github.io/). This is an approximate open-source implementation of the training pipeline used to train pi0-FAST-DROID. 
+
+
+## Requirements
+
+To run the models in this repository, you will need an NVIDIA GPU with at least the following specifications. These estimations assume a single GPU, but you can also use multiple GPUs with model parallelism to reduce per-GPU memory requirements by configuring `fsdp_devices` in the training config. Please also note that the current training script does not yet support multi-node training.
+
+| Mode               | Memory Required | Example GPU        |
+| ------------------ | --------------- | ------------------ |
+| Inference          | > 8 GB          | RTX 4090           |
+| Fine-Tuning (LoRA) | > 22.5 GB       | RTX 4090           |
+| Fine-Tuning (Full) | > 70 GB         | A100 (80GB) / H100 |
+
+The repo has been tested with Ubuntu 22.04, we do not currently support other operating systems.
+
+## Installation
+
+When cloning this repo, make sure to update submodules:
+
+```bash
+git clone --recurse-submodules git@github.com:Physical-Intelligence/openpi.git
+
+# Or if you already cloned the repo:
+git submodule update --init --recursive
+```
+
+We use [uv](https://docs.astral.sh/uv/) to manage Python dependencies. See the [uv installation instructions](https://docs.astral.sh/uv/getting-started/installation/) to set it up. Once uv is installed, run the following to set up the environment:
+
+```bash
+GIT_LFS_SKIP_SMUDGE=1 uv sync
+GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .
+```
+
+NOTE: `GIT_LFS_SKIP_SMUDGE=1` is needed to pull LeRobot as a dependency.
+
+**Docker**: As an alternative to uv installation, we provide instructions for installing openpi using Docker. If you encounter issues with your system setup, consider using Docker to simplify installation. See [Docker Setup](docs/docker.md) for more details.
+
+
+
+
+## Model Checkpoints
+
+### Base Models
+We provide multiple base VLA model checkpoints. These checkpoints have been pre-trained on 10k+ hours of robot data, and can be used for fine-tuning.
+
+| Model        | Use Case    | Description                                                                                                 | Checkpoint Path                                |
+| ------------ | ----------- | ----------------------------------------------------------------------------------------------------------- | ---------------------------------------------- |
+| $\pi_0$      | Fine-Tuning | Base [π₀ model](https://www.physicalintelligence.company/blog/pi0) for fine-tuning                | `gs://openpi-assets/checkpoints/pi0_base`      |
+| $\pi_0$-FAST | Fine-Tuning | Base autoregressive [π₀-FAST model](https://www.physicalintelligence.company/research/fast) for fine-tuning | `gs://openpi-assets/checkpoints/pi0_fast_base` |
+| $\pi_{0.5}$    | Fine-Tuning | Base [π₀.₅ model](https://www.physicalintelligence.company/blog/pi05) for fine-tuning    | `gs://openpi-assets/checkpoints/pi05_base`      |
+
+### Fine-Tuned Models
+We also provide "expert" checkpoints for various robot platforms and tasks. These models are fine-tuned from the base models above and intended to run directly on the target robot. These may or may not work on your particular robot. Since these checkpoints were fine-tuned on relatively small datasets collected with more widely available robots, such as ALOHA and the DROID Franka setup, they might not generalize to your particular setup, though we found some of these, especially the DROID checkpoint, to generalize quite broadly in practice.
+
+| Model                    | Use Case    | Description                                                                                                                                                                                              | Checkpoint Path                                       |
+| ------------------------ | ----------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------- |
+| $\pi_0$-FAST-DROID       | Inference   | $\pi_0$-FAST model fine-tuned on the [DROID dataset](https://droid-dataset.github.io/): can perform a wide range of simple table-top manipulation tasks 0-shot in new scenes on the DROID robot platform | `gs://openpi-assets/checkpoints/pi0_fast_droid`       |
+| $\pi_0$-DROID            | Fine-Tuning | $\pi_0$ model fine-tuned on the [DROID dataset](https://droid-dataset.github.io/): faster inference than $\pi_0$-FAST-DROID, but may not follow language commands as well                                | `gs://openpi-assets/checkpoints/pi0_droid`            |
+| $\pi_0$-ALOHA-towel      | Inference   | $\pi_0$ model fine-tuned on internal [ALOHA](https://tonyzhaozh.github.io/aloha/) data: can fold diverse towels 0-shot on ALOHA robot platforms                                                          | `gs://openpi-assets/checkpoints/pi0_aloha_towel`      |
+| $\pi_0$-ALOHA-tupperware | Inference   | $\pi_0$ model fine-tuned on internal [ALOHA](https://tonyzhaozh.github.io/aloha/) data: can unpack food from a tupperware container                                                                                                             | `gs://openpi-assets/checkpoints/pi0_aloha_tupperware` |
+| $\pi_0$-ALOHA-pen-uncap  | Inference   | $\pi_0$ model fine-tuned on public [ALOHA](https://dit-policy.github.io/) data: can uncap a pen                                                                                                          | `gs://openpi-assets/checkpoints/pi0_aloha_pen_uncap`  |
+| $\pi_{0.5}$-LIBERO      | Inference   | $\pi_{0.5}$ model fine-tuned for the [LIBERO](https://libero-project.github.io/datasets) benchmark: gets state-of-the-art performance (see [LIBERO README](examples/libero/README.md)) | `gs://openpi-assets/checkpoints/pi05_libero`      |
+| $\pi_{0.5}$-DROID      | Inference / Fine-Tuning | $\pi_{0.5}$ model fine-tuned on the [DROID dataset](https://droid-dataset.github.io/) with [knowledge insulation](https://www.physicalintelligence.company/research/knowledge_insulation): fast inference and good language-following | `gs://openpi-assets/checkpoints/pi05_droid`      |
+
+
+By default, checkpoints are automatically downloaded from `gs://openpi-assets` and are cached in `~/.cache/openpi` when needed. You can overwrite the download path by setting the `OPENPI_DATA_HOME` environment variable.
+
+
+
+
+## Running Inference for a Pre-Trained Model
+
+Our pre-trained model checkpoints can be run with a few lines of code (here our $\pi_0$-FAST-DROID model):
+```python
+from openpi.training import config as _config
+from openpi.policies import policy_config
+from openpi.shared import download
+
+config = _config.get_config("pi05_droid")
+checkpoint_dir = download.maybe_download("gs://openpi-assets/checkpoints/pi05_droid")
+
+# Create a trained policy.
+policy = policy_config.create_trained_policy(config, checkpoint_dir)
+
+# Run inference on a dummy example.
+example = {
+    "observation/exterior_image_1_left": ...,
+    "observation/wrist_image_left": ...,
+    ...
+    "prompt": "pick up the fork"
+}
+action_chunk = policy.infer(example)["actions"]
+```
+You can also test this out in the [example notebook](examples/inference.ipynb).
+
+We provide detailed step-by-step examples for running inference of our pre-trained checkpoints on [DROID](examples/droid/README.md) and [ALOHA](examples/aloha_real/README.md) robots.
+
+**Remote Inference**: We provide [examples and code](docs/remote_inference.md) for running inference of our models **remotely**: the model can run on a different server and stream actions to the robot via a websocket connection. This makes it easy to use more powerful GPUs off-robot and keep robot and policy environments separate.
+
+**Test inference without a robot**: We provide a [script](examples/simple_client/README.md) for testing inference without a robot. This script will generate a random observation and run inference with the model. See [here](examples/simple_client/README.md) for more details.
+
+
+
+
+
+## Fine-Tuning Base Models on Your Own Data
+
+We will fine-tune the $\pi_{0.5}$ model on the [LIBERO dataset](https://libero-project.github.io/datasets) as a running example for how to fine-tune a base model on your own data. We will explain three steps:
+1. Convert your data to a LeRobot dataset (which we use for training)
+2. Defining training configs and running training
+3. Spinning up a policy server and running inference
+
+### 1. Convert your data to a LeRobot dataset
+
+We provide a minimal example script for converting LIBERO data to a LeRobot dataset in [`examples/libero/convert_libero_data_to_lerobot.py`](examples/libero/convert_libero_data_to_lerobot.py). You can easily modify it to convert your own data! You can download the raw LIBERO dataset from [here](https://huggingface.co/datasets/openvla/modified_libero_rlds), and run the script with:
+
+```bash
+uv run examples/libero/convert_libero_data_to_lerobot.py --data_dir /path/to/your/libero/data
+```
+
+**Note:** If you just want to fine-tune on LIBERO, you can skip this step, because our LIBERO fine-tuning configs point to a pre-converted LIBERO dataset. This step is merely an example that you can adapt to your own data.
+
+### 2. Defining training configs and running training
+
+To fine-tune a base model on your own data, you need to define configs for data processing and training. We provide example configs with detailed comments for LIBERO below, which you can modify for your own dataset:
+
+- [`LiberoInputs` and `LiberoOutputs`](src/openpi/policies/libero_policy.py): Defines the data mapping from the LIBERO environment to the model and vice versa. Will be used for both, training and inference.
+- [`LeRobotLiberoDataConfig`](src/openpi/training/config.py): Defines how to process raw LIBERO data from LeRobot dataset for training.
+- [`TrainConfig`](src/openpi/training/config.py): Defines fine-tuning hyperparameters, data config, and weight loader.
+
+We provide example fine-tuning configs for [π₀](src/openpi/training/config.py), [π₀-FAST](src/openpi/training/config.py), and [π₀.₅](src/openpi/training/config.py) on LIBERO data.
+
+Before we can run training, we need to compute the normalization statistics for the training data. Run the script below with the name of your training config:
+
+```bash
+uv run scripts/compute_norm_stats.py --config-name pi05_libero
+```
+
+Now we can kick off training with the following command (the `--overwrite` flag is used to overwrite existing checkpoints if you rerun fine-tuning with the same config):
+
+```bash
+XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py pi05_libero --exp-name=my_experiment --overwrite
+```
+
+The command will log training progress to the console and save checkpoints to the `checkpoints` directory. You can also monitor training progress on the Weights & Biases dashboard. For maximally using the GPU memory, set `XLA_PYTHON_CLIENT_MEM_FRACTION=0.9` before running training -- this enables JAX to use up to 90% of the GPU memory (vs. the default of 75%).
+
+**Note:** We provide functionality for *reloading* normalization statistics for state / action normalization from pre-training. This can be beneficial if you are fine-tuning to a new task on a robot that was part of our pre-training mixture. For more details on how to reload normalization statistics, see the [norm_stats.md](docs/norm_stats.md) file.
+
+### 3. Spinning up a policy server and running inference
+
+Once training is complete, we can run inference by spinning up a policy server and then querying it from a LIBERO evaluation script. Launching a model server is easy (we use the checkpoint for iteration 20,000 for this example, modify as needed):
+
+```bash
+uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi05_libero --policy.dir=checkpoints/pi05_libero/my_experiment/20000
+```
+
+This will spin up a server that listens on port 8000 and waits for observations to be sent to it. We can then run an evaluation script (or robot runtime) that queries the server.
+
+For running the LIBERO eval in particular, we provide (and recommend using) a Dockerized workflow that handles both the policy server and the evaluation script together. See the [LIBERO README](examples/libero/README.md) for more details.
+
+If you want to embed a policy server call in your own robot runtime, we have a minimal example of how to do so in the [remote inference docs](docs/remote_inference.md).
+
+
+
+### More Examples
+
+We provide more examples for how to fine-tune and run inference with our models on the ALOHA platform in the following READMEs:
+- [ALOHA Simulator](examples/aloha_sim)
+- [ALOHA Real](examples/aloha_real)
+- [UR5](examples/ur5)
+
+## PyTorch Support
+
+openpi now provides PyTorch implementations of π₀ and π₀.₅ models alongside the original JAX versions! The PyTorch implementation has been validated on the LIBERO benchmark (both inference and finetuning). A few features are currently not supported (this may change in the future):
+
+- The π₀-FAST model
+- Mixed precision training
+- FSDP (fully-sharded data parallelism) training
+- LoRA (low-rank adaptation) training
+- EMA (exponential moving average) weights during training
+
+### Setup
+1. Make sure that you have the latest version of all dependencies installed: `uv sync`
+
+2. Double check that you have transformers 4.53.2 installed: `uv pip show transformers`
+
+3. Apply the transformers library patches:
+   ```bash
+   cp -r ./src/openpi/models_pytorch/transformers_replace/* .venv/lib/python3.11/site-packages/transformers/
+   ```
+
+This overwrites several files in the transformers library with necessary model changes: 1) supporting AdaRMS, 2) correctly controlling the precision of activations, and 3) allowing the KV cache to be used without being updated.
+
+**WARNING**: With the default uv link mode (hardlink), this will permanently affect the transformers library in your uv cache, meaning the changes will survive reinstallations of transformers and could even propagate to other projects that use transformers. To fully undo this operation, you must run `uv cache clean transformers`.
+
+### Converting JAX Models to PyTorch
+
+To convert a JAX model checkpoint to PyTorch format:
+
+```bash
+uv run examples/convert_jax_model_to_pytorch.py \
+    --checkpoint_dir /path/to/jax/checkpoint \
+    --config_name <config name> \
+    --output_path /path/to/converted/pytorch/checkpoint
+```
+
+### Running Inference with PyTorch
+
+The PyTorch implementation uses the same API as the JAX version - you only need to change the checkpoint path to point to the converted PyTorch model:
+
+```python
+from openpi.training import config as _config
+from openpi.policies import policy_config
+from openpi.shared import download
+
+config = _config.get_config("pi05_droid")
+checkpoint_dir = "/path/to/converted/pytorch/checkpoint"
+
+# Create a trained policy (automatically detects PyTorch format)
+policy = policy_config.create_trained_policy(config, checkpoint_dir)
+
+# Run inference (same API as JAX)
+action_chunk = policy.infer(example)["actions"]
+```
+
+### Policy Server with PyTorch
+
+The policy server works identically with PyTorch models - just point to the converted checkpoint directory:
+
+```bash
+uv run scripts/serve_policy.py policy:checkpoint \
+    --policy.config=pi05_droid \
+    --policy.dir=/path/to/converted/pytorch/checkpoint
+```
+
+### Finetuning with PyTorch
+
+To finetune a model in PyTorch:
+
+1. Convert the JAX base model to PyTorch format:
+   ```bash
+   uv run examples/convert_jax_model_to_pytorch.py \
+       --config_name <config name> \
+       --checkpoint_dir /path/to/jax/base/model \
+       --output_path /path/to/pytorch/base/model
+   ```
+
+2. Specify the converted PyTorch model path in your config using `pytorch_weight_path`
+
+3. Launch training using one of these modes:
+
+```bash
+# Single GPU training:
+uv run scripts/train_pytorch.py <config_name> --exp_name <run_name> --save_interval <interval>
+
+# Example:
+uv run scripts/train_pytorch.py debug --exp_name pytorch_test
+uv run scripts/train_pytorch.py debug --exp_name pytorch_test --resume  # Resume from latest checkpoint
+
+# Multi-GPU training (single node):
+uv run torchrun --standalone --nnodes=1 --nproc_per_node=<num_gpus> scripts/train_pytorch.py <config_name> --exp_name <run_name>
+
+# Example:
+uv run torchrun --standalone --nnodes=1 --nproc_per_node=2 scripts/train_pytorch.py pi0_aloha_sim --exp_name pytorch_ddp_test
+uv run torchrun --standalone --nnodes=1 --nproc_per_node=2 scripts/train_pytorch.py pi0_aloha_sim --exp_name pytorch_ddp_test --resume
+
+# Multi-Node Training:
+uv run torchrun \
+    --nnodes=<num_nodes> \
+    --nproc_per_node=<gpus_per_node> \
+    --node_rank=<rank_of_node> \
+    --master_addr=<master_ip> \
+    --master_port=<port> \
+    scripts/train_pytorch.py <config_name> --exp_name=<run_name> --save_interval <interval>
+```
+
+### Precision Settings
+
+JAX and PyTorch implementations handle precision as follows:
+
+**JAX:**
+1. Inference: most weights and computations in bfloat16, with a few computations in float32 for stability
+2. Training: defaults to mixed precision: weights and gradients in float32, (most) activations and computations in bfloat16. You can change to full float32 training by setting `dtype` to float32 in the config.
+
+**PyTorch:**
+1. Inference: matches JAX -- most weights and computations in bfloat16, with a few weights converted to float32 for stability
+2. Training: supports either full bfloat16 (default) or full float32. You can change it by setting `pytorch_training_precision` in the config. bfloat16 uses less memory but exhibits higher losses compared to float32. Mixed precision is not yet supported.
+
+With torch.compile, inference speed is comparable between JAX and PyTorch.
+
+## Troubleshooting
+
+We will collect common issues and their solutions here. If you encounter an issue, please check here first. If you can't find a solution, please file an issue on the repo (see [here](CONTRIBUTING.md) for guidelines).
+
+| Issue                                     | Resolution                                                                                                                                                                                   |
+| ----------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `uv sync` fails with dependency conflicts | Try removing the virtual environment directory (`rm -rf .venv`) and running `uv sync` again. If issues persist, check that you have the latest version of `uv` installed (`uv self update`). |
+| Training runs out of GPU memory           | Make sure you set `XLA_PYTHON_CLIENT_MEM_FRACTION=0.9` (or higher) before running training to allow JAX to use more GPU memory. You can also use `--fsdp-devices <n>` where `<n>` is your number of GPUs, to enable [fully-sharded data parallelism](https://engineering.fb.com/2021/07/15/open-source/fsdp/), which reduces memory usage in exchange for slower training (the amount of slowdown depends on your particular setup). If you are still running out of memory, you may want to consider disabling EMA.        |
+| Policy server connection errors           | Check that the server is running and listening on the expected port. Verify network connectivity and firewall settings between client and server.                                            |
+| Missing norm stats error when training    | Run `scripts/compute_norm_stats.py` with your config name before starting training.                                                                                                          |
+| Dataset download fails                    | Check your internet connection. For HuggingFace datasets, ensure you're logged in (`huggingface-cli login`).                                                                                 |
+| CUDA/GPU errors                           | Verify NVIDIA drivers are installed correctly. For Docker, ensure nvidia-container-toolkit is installed. Check GPU compatibility. You do NOT need CUDA libraries installed at a system level --- they will be installed via uv. You may even want to try *uninstalling* system CUDA libraries if you run into CUDA issues, since system libraries can sometimes cause conflicts. |
+| Import errors when running examples       | Make sure you've installed all dependencies with `uv sync`. Some examples may have additional requirements listed in their READMEs.                    |
+| Action dimensions mismatch                | Verify your data processing transforms match the expected input/output dimensions of your robot. Check the action space definitions in your policy classes.                                  |
+| Diverging training loss                            | Check the `q01`, `q99`, and `std` values in `norm_stats.json` for your dataset. Certain dimensions that are rarely used can end up with very small `q01`, `q99`, or `std` values, leading to huge states and actions after normalization. You can manually adjust the norm stats as a workaround. |

code/pi05/packages/openpi-client/pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "openpi-client"
+version = "0.1.0"
+requires-python = ">=3.7"
+dependencies = [
+    "dm-tree>=0.1.8",
+    "msgpack>=1.0.5",
+    "numpy>=1.22.4,<2.0.0",
+    "pillow>=9.0.0",
+    "websockets>=11.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[dependency-groups]
+dev = ["pytest>=8.3.4"]
+
+[tool.ruff]
+line-length = 120
+target-version = "py37"

code/pi05/packages/openpi-client/src/openpi_client/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

code/pi05/packages/openpi-client/src/openpi_client/action_chunk_broker.py

@@ -0,0 +1,50 @@
+from typing import Dict
+
+import numpy as np
+import tree
+from typing_extensions import override
+
+from openpi_client import base_policy as _base_policy
+
+
+class ActionChunkBroker(_base_policy.BasePolicy):
+    """Wraps a policy to return action chunks one-at-a-time.
+
+    Assumes that the first dimension of all action fields is the chunk size.
+
+    A new inference call to the inner policy is only made when the current
+    list of chunks is exhausted.
+    """
+
+    def __init__(self, policy: _base_policy.BasePolicy, action_horizon: int):
+        self._policy = policy
+        self._action_horizon = action_horizon
+        self._cur_step: int = 0
+
+        self._last_results: Dict[str, np.ndarray] | None = None
+
+    @override
+    def infer(self, obs: Dict) -> Dict:  # noqa: UP006
+        if self._last_results is None:
+            self._last_results = self._policy.infer(obs)
+            self._cur_step = 0
+
+        def slicer(x):
+            if isinstance(x, np.ndarray):
+                return x[self._cur_step, ...]
+            else:
+                return x
+
+        results = tree.map_structure(slicer, self._last_results)
+        self._cur_step += 1
+
+        if self._cur_step >= self._action_horizon:
+            self._last_results = None
+
+        return results
+
+    @override
+    def reset(self) -> None:
+        self._policy.reset()
+        self._last_results = None
+        self._cur_step = 0

code/pi05/packages/openpi-client/src/openpi_client/base_policy.py

@@ -0,0 +1,12 @@
+import abc
+from typing import Dict
+
+
+class BasePolicy(abc.ABC):
+    @abc.abstractmethod
+    def infer(self, obs: Dict) -> Dict:
+        """Infer actions from observations."""
+
+    def reset(self) -> None:
+        """Reset the policy to its initial state."""
+        pass

code/pi05/packages/openpi-client/src/openpi_client/image_tools.py

@@ -0,0 +1,58 @@
+import numpy as np
+from PIL import Image
+
+
+def convert_to_uint8(img: np.ndarray) -> np.ndarray:
+    """Converts an image to uint8 if it is a float image.
+
+    This is important for reducing the size of the image when sending it over the network.
+    """
+    if np.issubdtype(img.dtype, np.floating):
+        img = (255 * img).astype(np.uint8)
+    return img
+
+
+def resize_with_pad(images: np.ndarray, height: int, width: int, method=Image.BILINEAR) -> np.ndarray:
+    """Replicates tf.image.resize_with_pad for multiple images using PIL. Resizes a batch of images to a target height.
+
+    Args:
+        images: A batch of images in [..., height, width, channel] format.
+        height: The target height of the image.
+        width: The target width of the image.
+        method: The interpolation method to use. Default is bilinear.
+
+    Returns:
+        The resized images in [..., height, width, channel].
+    """
+    # If the images are already the correct size, return them as is.
+    if images.shape[-3:-1] == (height, width):
+        return images
+
+    original_shape = images.shape
+
+    images = images.reshape(-1, *original_shape[-3:])
+    resized = np.stack([_resize_with_pad_pil(Image.fromarray(im), height, width, method=method) for im in images])
+    return resized.reshape(*original_shape[:-3], *resized.shape[-3:])
+
+
+def _resize_with_pad_pil(image: Image.Image, height: int, width: int, method: int) -> Image.Image:
+    """Replicates tf.image.resize_with_pad for one image using PIL. Resizes an image to a target height and
+    width without distortion by padding with zeros.
+
+    Unlike the jax version, note that PIL uses [width, height, channel] ordering instead of [batch, h, w, c].
+    """
+    cur_width, cur_height = image.size
+    if cur_width == width and cur_height == height:
+        return image  # No need to resize if the image is already the correct size.
+
+    ratio = max(cur_width / width, cur_height / height)
+    resized_height = int(cur_height / ratio)
+    resized_width = int(cur_width / ratio)
+    resized_image = image.resize((resized_width, resized_height), resample=method)
+
+    zero_image = Image.new(resized_image.mode, (width, height), 0)
+    pad_height = max(0, int((height - resized_height) / 2))
+    pad_width = max(0, int((width - resized_width) / 2))
+    zero_image.paste(resized_image, (pad_width, pad_height))
+    assert zero_image.size == (width, height)
+    return zero_image

code/pi05/packages/openpi-client/src/openpi_client/image_tools_test.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+import openpi_client.image_tools as image_tools
+
+
+def test_resize_with_pad_shapes():
+    # Test case 1: Resize image with larger dimensions
+    images = np.zeros((2, 10, 10, 3), dtype=np.uint8)  # Input images of shape (batch_size, height, width, channels)
+    height = 20
+    width = 20
+    resized_images = image_tools.resize_with_pad(images, height, width)
+    assert resized_images.shape == (2, height, width, 3)
+    assert np.all(resized_images == 0)
+
+    # Test case 2: Resize image with smaller dimensions
+    images = np.zeros((3, 30, 30, 3), dtype=np.uint8)
+    height = 15
+    width = 15
+    resized_images = image_tools.resize_with_pad(images, height, width)
+    assert resized_images.shape == (3, height, width, 3)
+    assert np.all(resized_images == 0)
+
+    # Test case 3: Resize image with the same dimensions
+    images = np.zeros((1, 50, 50, 3), dtype=np.uint8)
+    height = 50
+    width = 50
+    resized_images = image_tools.resize_with_pad(images, height, width)
+    assert resized_images.shape == (1, height, width, 3)
+    assert np.all(resized_images == 0)
+
+    # Test case 3: Resize image with odd-numbered padding
+    images = np.zeros((1, 256, 320, 3), dtype=np.uint8)
+    height = 60
+    width = 80
+    resized_images = image_tools.resize_with_pad(images, height, width)
+    assert resized_images.shape == (1, height, width, 3)
+    assert np.all(resized_images == 0)

code/pi05/packages/openpi-client/src/openpi_client/msgpack_numpy.py

@@ -0,0 +1,57 @@
+"""Adds NumPy array support to msgpack.
+
+msgpack is good for (de)serializing data over a network for multiple reasons:
+- msgpack is secure (as opposed to pickle/dill/etc which allow for arbitrary code execution)
+- msgpack is widely used and has good cross-language support
+- msgpack does not require a schema (as opposed to protobuf/flatbuffers/etc) which is convenient in dynamically typed
+    languages like Python and JavaScript
+- msgpack is fast and efficient (as opposed to readable formats like JSON/YAML/etc); I found that msgpack was ~4x faster
+    than pickle for serializing large arrays using the below strategy
+
+The code below is adapted from https://github.com/lebedov/msgpack-numpy. The reason not to use that library directly is
+that it falls back to pickle for object arrays.
+"""
+
+import functools
+
+import msgpack
+import numpy as np
+
+
+def pack_array(obj):
+    if (isinstance(obj, (np.ndarray, np.generic))) and obj.dtype.kind in ("V", "O", "c"):
+        raise ValueError(f"Unsupported dtype: {obj.dtype}")
+
+    if isinstance(obj, np.ndarray):
+        return {
+            b"__ndarray__": True,
+            b"data": obj.tobytes(),
+            b"dtype": obj.dtype.str,
+            b"shape": obj.shape,
+        }
+
+    if isinstance(obj, np.generic):
+        return {
+            b"__npgeneric__": True,
+            b"data": obj.item(),
+            b"dtype": obj.dtype.str,
+        }
+
+    return obj
+
+
+def unpack_array(obj):
+    if b"__ndarray__" in obj:
+        return np.ndarray(buffer=obj[b"data"], dtype=np.dtype(obj[b"dtype"]), shape=obj[b"shape"])
+
+    if b"__npgeneric__" in obj:
+        return np.dtype(obj[b"dtype"]).type(obj[b"data"])
+
+    return obj
+
+
+Packer = functools.partial(msgpack.Packer, default=pack_array)
+packb = functools.partial(msgpack.packb, default=pack_array)
+
+Unpacker = functools.partial(msgpack.Unpacker, object_hook=unpack_array)
+unpackb = functools.partial(msgpack.unpackb, object_hook=unpack_array)

code/pi05/packages/openpi-client/src/openpi_client/msgpack_numpy_test.py

@@ -0,0 +1,45 @@
+import numpy as np
+import pytest
+import tree
+
+from openpi_client import msgpack_numpy
+
+
+def _check(expected, actual):
+    if isinstance(expected, np.ndarray):
+        assert expected.shape == actual.shape
+        assert expected.dtype == actual.dtype
+        assert np.array_equal(expected, actual, equal_nan=expected.dtype.kind == "f")
+    else:
+        assert expected == actual
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        1,  # int
+        1.0,  # float
+        "hello",  # string
+        np.bool_(True),  # boolean scalar
+        np.array([1, 2, 3])[0],  # int scalar
+        np.str_("asdf"),  # string scalar
+        [1, 2, 3],  # list
+        {"key": "value"},  # dict
+        {"key": [1, 2, 3]},  # nested dict
+        np.array(1.0),  # 0D array
+        np.array([1, 2, 3], dtype=np.int32),  # 1D integer array
+        np.array(["asdf", "qwer"]),  # string array
+        np.array([True, False]),  # boolean array
+        np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32),  # 2D float array
+        np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dtype=np.int16),  # 3D integer array
+        np.array([np.nan, np.inf, -np.inf]),  # special float values
+        {"arr": np.array([1, 2, 3]), "nested": {"arr": np.array([4, 5, 6])}},  # nested dict with arrays
+        [np.array([1, 2]), np.array([3, 4])],  # list of arrays
+        np.zeros((3, 4, 5), dtype=np.float32),  # 3D zeros
+        np.ones((2, 3), dtype=np.float64),  # 2D ones with double precision
+    ],
+)
+def test_pack_unpack(data):
+    packed = msgpack_numpy.packb(data)
+    unpacked = msgpack_numpy.unpackb(packed)
+    tree.map_structure(_check, data, unpacked)

code/pi05/packages/openpi-client/src/openpi_client/runtime/agent.py

@@ -0,0 +1,17 @@
+import abc
+
+
+class Agent(abc.ABC):
+    """An Agent is the thing with agency, i.e. the entity that makes decisions.
+
+    Agents receive observations about the state of the world, and return actions
+    to take in response.
+    """
+
+    @abc.abstractmethod
+    def get_action(self, observation: dict) -> dict:
+        """Query the agent for the next action."""
+
+    @abc.abstractmethod
+    def reset(self) -> None:
+        """Reset the agent to its initial state."""

code/pi05/packages/openpi-client/src/openpi_client/runtime/agents/policy_agent.py

@@ -0,0 +1,18 @@
+from typing_extensions import override
+
+from openpi_client import base_policy as _base_policy
+from openpi_client.runtime import agent as _agent
+
+
+class PolicyAgent(_agent.Agent):
+    """An agent that uses a policy to determine actions."""
+
+    def __init__(self, policy: _base_policy.BasePolicy) -> None:
+        self._policy = policy
+
+    @override
+    def get_action(self, observation: dict) -> dict:
+        return self._policy.infer(observation)
+
+    def reset(self) -> None:
+        self._policy.reset()

code/pi05/packages/openpi-client/src/openpi_client/runtime/environment.py

@@ -0,0 +1,32 @@
+import abc
+
+
+class Environment(abc.ABC):
+    """An Environment represents the robot and the environment it inhabits.
+
+    The primary contract of environments is that they can be queried for observations
+    about their state, and have actions applied to them to change that state.
+    """
+
+    @abc.abstractmethod
+    def reset(self) -> None:
+        """Reset the environment to its initial state.
+
+        This will be called once before starting each episode.
+        """
+
+    @abc.abstractmethod
+    def is_episode_complete(self) -> bool:
+        """Allow the environment to signal that the episode is complete.
+
+        This will be called after each step. It should return `True` if the episode is
+        complete (either successfully or unsuccessfully), and `False` otherwise.
+        """
+
+    @abc.abstractmethod
+    def get_observation(self) -> dict:
+        """Query the environment for the current state."""
+
+    @abc.abstractmethod
+    def apply_action(self, action: dict) -> None:
+        """Take an action in the environment."""

code/pi05/packages/openpi-client/src/openpi_client/runtime/runtime.py

@@ -0,0 +1,92 @@
+import logging
+import threading
+import time
+
+from openpi_client.runtime import agent as _agent
+from openpi_client.runtime import environment as _environment
+from openpi_client.runtime import subscriber as _subscriber
+
+
+class Runtime:
+    """The core module orchestrating interactions between key components of the system."""
+
+    def __init__(
+        self,
+        environment: _environment.Environment,
+        agent: _agent.Agent,
+        subscribers: list[_subscriber.Subscriber],
+        max_hz: float = 0,
+        num_episodes: int = 1,
+        max_episode_steps: int = 0,
+    ) -> None:
+        self._environment = environment
+        self._agent = agent
+        self._subscribers = subscribers
+        self._max_hz = max_hz
+        self._num_episodes = num_episodes
+        self._max_episode_steps = max_episode_steps
+
+        self._in_episode = False
+        self._episode_steps = 0
+
+    def run(self) -> None:
+        """Runs the runtime loop continuously until stop() is called or the environment is done."""
+        for _ in range(self._num_episodes):
+            self._run_episode()
+
+        # Final reset, this is important for real environments to move the robot to its home position.
+        self._environment.reset()
+
+    def run_in_new_thread(self) -> threading.Thread:
+        """Runs the runtime loop in a new thread."""
+        thread = threading.Thread(target=self.run)
+        thread.start()
+        return thread
+
+    def mark_episode_complete(self) -> None:
+        """Marks the end of an episode."""
+        self._in_episode = False
+
+    def _run_episode(self) -> None:
+        """Runs a single episode."""
+        logging.info("Starting episode...")
+        self._environment.reset()
+        self._agent.reset()
+        for subscriber in self._subscribers:
+            subscriber.on_episode_start()
+
+        self._in_episode = True
+        self._episode_steps = 0
+        step_time = 1 / self._max_hz if self._max_hz > 0 else 0
+        last_step_time = time.time()
+
+        while self._in_episode:
+            self._step()
+            self._episode_steps += 1
+
+            # Sleep to maintain the desired frame rate
+            now = time.time()
+            dt = now - last_step_time
+            if dt < step_time:
+                time.sleep(step_time - dt)
+                last_step_time = time.time()
+            else:
+                last_step_time = now
+
+        logging.info("Episode completed.")
+        for subscriber in self._subscribers:
+            subscriber.on_episode_end()
+
+    def _step(self) -> None:
+        """A single step of the runtime loop."""
+        observation = self._environment.get_observation()
+        action = self._agent.get_action(observation)
+        self._environment.apply_action(action)
+
+        for subscriber in self._subscribers:
+            subscriber.on_step(observation, action)
+
+        if self._environment.is_episode_complete() or (
+            self._max_episode_steps > 0 and self._episode_steps >= self._max_episode_steps
+        ):
+            self.mark_episode_complete()

code/pi05/packages/openpi-client/src/openpi_client/runtime/subscriber.py

@@ -0,0 +1,20 @@
+import abc
+
+
+class Subscriber(abc.ABC):
+    """Subscribes to events in the runtime.
+
+    Subscribers can be used to save data, visualize, etc.
+    """
+
+    @abc.abstractmethod
+    def on_episode_start(self) -> None:
+        """Called when an episode starts."""
+
+    @abc.abstractmethod
+    def on_step(self, observation: dict, action: dict) -> None:
+        """Append a step to the episode."""
+
+    @abc.abstractmethod
+    def on_episode_end(self) -> None:
+        """Called when an episode ends."""

code/pi05/packages/openpi-client/src/openpi_client/websocket_client_policy.py

@@ -0,0 +1,58 @@
+import logging
+import time
+from typing import Dict, Optional, Tuple
+
+from typing_extensions import override
+import websockets.sync.client
+
+from openpi_client import base_policy as _base_policy
+from openpi_client import msgpack_numpy
+
+
+class WebsocketClientPolicy(_base_policy.BasePolicy):
+    """Implements the Policy interface by communicating with a server over websocket.
+
+    See WebsocketPolicyServer for a corresponding server implementation.
+    """
+
+    def __init__(self, host: str = "0.0.0.0", port: Optional[int] = None, api_key: Optional[str] = None) -> None:
+        if host.startswith("ws"):
+            self._uri = host
+        else:
+            self._uri = f"ws://{host}"
+        if port is not None:
+            self._uri += f":{port}"
+        self._packer = msgpack_numpy.Packer()
+        self._api_key = api_key
+        self._ws, self._server_metadata = self._wait_for_server()
+
+    def get_server_metadata(self) -> Dict:
+        return self._server_metadata
+
+    def _wait_for_server(self) -> Tuple[websockets.sync.client.ClientConnection, Dict]:
+        logging.info(f"Waiting for server at {self._uri}...")
+        while True:
+            try:
+                headers = {"Authorization": f"Api-Key {self._api_key}"} if self._api_key else None
+                conn = websockets.sync.client.connect(
+                    self._uri, compression=None, max_size=None, additional_headers=headers
+                )
+                metadata = msgpack_numpy.unpackb(conn.recv())
+                return conn, metadata
+            except ConnectionRefusedError:
+                logging.info("Still waiting for server...")
+                time.sleep(5)
+
+    @override
+    def infer(self, obs: Dict) -> Dict:  # noqa: UP006
+        data = self._packer.pack(obs)
+        self._ws.send(data)
+        response = self._ws.recv()
+        if isinstance(response, str):
+            # we're expecting bytes; if the server sends a string, it's an error.
+            raise RuntimeError(f"Error in inference server:\n{response}")
+        return msgpack_numpy.unpackb(response)
+
+    @override
+    def reset(self) -> None:
+        pass

code/pi05/pyproject.toml

@@ -0,0 +1,140 @@
+[project]
+name = "openpi"
+version = "0.1.0"
+description = "Physical Intelligence open source repo"
+readme = "README.md"
+requires-python = ">=3.11"
+license = { file = "LICENSE" }
+dependencies = [
+    "augmax>=0.3.4",
+    "dm-tree>=0.1.8",
+    "einops>=0.8.0",
+    "equinox>=0.11.8",
+    "flatbuffers>=24.3.25",
+    "flax==0.10.2",
+    "fsspec[gcs]>=2024.6.0",
+    "gym-aloha>=0.1.1",
+    "imageio>=2.36.1",
+    "jax[cuda12]==0.5.3",
+    "jaxtyping==0.2.36",
+    "lerobot",
+    "ml_collections==1.0.0",
+    "numpy>=1.22.4,<2.0.0",
+    "numpydantic>=1.6.6",
+    "opencv-python>=4.10.0.84",
+    "openpi-client",
+    "orbax-checkpoint==0.11.13",
+    "pillow>=11.0.0",
+    "sentencepiece>=0.2.0",
+    "torch==2.7.1",
+    "tqdm-loggable>=0.2",
+    "typing-extensions>=4.12.2",
+    "tyro>=0.9.5",
+    "wandb>=0.19.1",
+    "filelock>=3.16.1",
+    "beartype==0.19.0",
+    "chex==0.1.90",
+    "treescope>=0.1.7",
+    "transformers==4.53.2",
+    "rich>=14.0.0",
+    "polars>=1.30.0",
+]
+
+
+[project.urls]
+Repository = "https://github.com/Physical-Intelligence/openpi"
+
+[dependency-groups]
+dev = [
+    "pytest>=8.3.4",
+    "ruff>=0.8.6",
+    "pre-commit>=4.0.1",
+    "ipykernel>=6.29.5",
+    "ipywidgets>=8.1.5",
+    "matplotlib>=3.10.0",
+    "pynvml>=12.0.0",
+]
+rlds = [
+    # NOTE: tensorflow-cpu 2.15.0 only ships cp311 wheels on PyPI.
+    # Use `uv venv --python 3.11` before `uv sync --group rlds`.
+    "dlimp",
+    "tensorflow-cpu==2.15.0",
+    "tensorflow-datasets==4.9.9",
+]
+
+[tool.uv]
+override-dependencies = ["ml-dtypes==0.4.1", "tensorstore==0.1.74"]
+
+[tool.uv.sources]
+openpi-client = { workspace = true }
+lerobot = { git = "https://github.com/huggingface/lerobot", rev = "0cf864870cf29f4738d3ade893e6fd13fbd7cdb5" }
+dlimp = { git = "https://github.com/kvablack/dlimp", rev = "ad72ce3a9b414db2185bc0b38461d4101a65477a" }
+
+[tool.uv.workspace]
+members = ["packages/*"]
+
+[tool.ruff]
+line-length = 120
+target-version = "py311"
+extend-exclude = ["docker", "third_party", "src/openpi/models_pytorch/transformers_replace/*"]
+
+[tool.ruff.lint]
+# https://docs.astral.sh/ruff/rules/
+select = [
+    "B",
+    "C4",
+    "DTZ",
+    "E4",
+    "E7",
+    "E9",
+    "F",
+    "FBT",
+    "FURB",
+    "I",
+    "ICN",
+    "ISC",
+    "LOG",
+    "N",
+    "PD",
+    "PERF",
+    "PIE",
+    "PLC",
+    "PLE",
+    "PLR1",
+    "PLR5",
+    "PLW",
+    "PT",
+    "Q",
+    "RET",
+    "RUF",
+    "SIM",
+    "SLF",
+    "T10",
+    "T20",
+    "UP",
+    "W",
+]
+ignore = [
+    "F722",   # Conflicts with array typing.
+    "T201",   # We use print statements.
+    "PD008",  # Lots of false positives.
+    "ISC001", # Disabling to support ruff format.
+    "LOG015", # Use logger.info.
+]
+unfixable = [
+    "B905", # Fix defaults to strict=False, which is not what we want.
+]
+
+[tool.ruff.lint.isort]
+force-single-line = true
+force-sort-within-sections = true
+single-line-exclusions = ["collections.abc", "typing", "typing_extensions"]
+known-third-party = ["wandb"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.pytest.ini_options]
+markers = ["manual: should be run manually."]
+testpaths = ["src", "scripts", "packages"]

code/pi05/scripts/check_robomem_memory_graph.py

@@ -0,0 +1,59 @@
+"""Compile both phases without loading weights or executing training updates."""
+
+import json
+import logging
+from pathlib import Path
+import time
+
+import jax
+import numpy as np
+
+import train_robomem_memory as trainer
+from openpi.shared import array_typing as at
+from openpi.training import data_loader, sharding
+
+
+def main():
+    logging.basicConfig(level=logging.INFO, force=True)
+    args = trainer.build_argparser().parse_args()
+    cfg = trainer.recipe(args)
+    mesh = sharding.make_mesh(args.fsdp_devices)
+    state, state_sharding = trainer.init_state(cfg, args, mesh, True)
+    index = trainer.EpisodeIndex(args.dataset)
+    dataset = trainer.MemoryDataset(index, cfg.data.create(cfg.assets_dirs, cfg.model), cfg.model, args.cache)
+    ep = index.groups[sorted(index.groups)[0]][0]
+    pos = index.episodes.index(ep)
+    sample_index = int(index.ends[pos - 1]) if pos else 0
+    batch = data_loader._collate_fn([dataset[sample_index]] * args.batch_size)
+    batch = jax.tree.map(lambda a: jax.ShapeDtypeStruct(a.shape, a.dtype), batch)
+    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
+    replicated = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+    report = {}
+    with sharding.set_mesh(mesh):
+        for warmup in (True, False):
+            name = "adapter_only" if warmup else "joint_frozen_vision"
+            started = time.monotonic()
+            fn = jax.jit(
+                lambda s, b: trainer.train_step(s, b, warmup=warmup),
+                in_shardings=(state_sharding, data_sharding),
+                out_shardings=(state_sharding, replicated),
+                donate_argnums=(0,),
+            )
+            with at.disable_typechecking():
+                executable = fn.lower(state, batch).compile()
+            memory = executable.memory_analysis()
+            report[name] = dict(
+                seconds=time.monotonic() - started,
+                arguments_bytes=memory.argument_size_in_bytes,
+                output_bytes=memory.output_size_in_bytes,
+                temporary_bytes=memory.temp_size_in_bytes,
+                alias_bytes=memory.alias_size_in_bytes,
+            )
+            print(name, json.dumps(report[name]), flush=True)
+    report["action_horizon"] = cfg.model.action_horizon
+    report["execution_horizon"] = args.execution_horizon
+    Path(f"logs/robomem_memory_spatial_a{args.action_horizon}_compile_check.json").write_text(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()

code/pi05/scripts/evaluate_robomem_memory.py

@@ -0,0 +1,271 @@
+"""Run selected official tasks and persist coverage, exclusions, failures and CSR/TSR."""
+
+import argparse
+import concurrent.futures
+import hashlib
+import json
+import os
+from pathlib import Path
+import subprocess
+import threading
+import time
+import urllib.request
+
+from openpi_client import msgpack_numpy
+from openpi.policies.robomem_progress import PROGRESS_MODE
+from robomem_eval_paths import ROOT, BENCHMARK, POLICY_PYTHON
+from taskanchor.data.history import validate_robomem_recipe
+from run_robomem_eval_task import run_task
+from robomem_seed_search import run_task25
+import websockets.sync.client
+
+
+def verify_server_checkpoint(port, checkpoint):
+    """Reject a stale policy server before sending any evaluation observations."""
+    with websockets.sync.client.connect(
+        f"ws://127.0.0.1:{port}", compression=None, max_size=None, open_timeout=10
+    ) as connection:
+        metadata = msgpack_numpy.unpackb(connection.recv(timeout=10))
+    actual = metadata.get("checkpoint")
+    if not actual or Path(actual).resolve() != Path(checkpoint).resolve():
+        raise RuntimeError(f"Policy server {port} checkpoint mismatch: expected {checkpoint}, got {actual}")
+    return metadata
+
+
+def main(args):
+    tasks = [task for task in range(1, 27) if task not in args.skip_tasks]
+    if not tasks or args.trials < 1:
+        raise ValueError("Evaluation requires at least one task and one trial")
+    task25_candidates = json.loads(Path(args.task25_seeds).read_text())["candidates"] if args.task25_seeds else None
+    if task25_candidates is not None and (not task25_candidates or len(task25_candidates) > 200
+            or len(task25_candidates) != len(set(task25_candidates))):
+        raise ValueError("Task25 requires 1 to 200 distinct candidate seeds")
+    if args.wait_smoke:
+        smoke = Path(args.wait_smoke).resolve()
+        print(f"Waiting for closed-loop smoke result: {smoke}", flush=True)
+        deadline = time.monotonic() + 3600
+        while not (smoke / "aggregate.json").exists():
+            log = (smoke / "run.log").read_text() if (smoke / "run.log").exists() else ""
+            if "Traceback (most recent call last)" in log or "Episode failed:" in log:
+                raise RuntimeError("Closed-loop smoke encountered a runtime error; refusing full evaluation")
+            if time.monotonic() > deadline:
+                raise TimeoutError("Closed-loop smoke did not finish within one hour")
+            time.sleep(5)
+        if json.loads((smoke / "aggregate.json").read_text())["num_episodes"] != 1:
+            raise RuntimeError("Incomplete closed-loop smoke")
+        log = (smoke / "run.log").read_text()
+        if "Traceback (most recent call last)" in log or "Episode failed:" in log:
+            raise RuntimeError("Closed-loop smoke failed")
+        print("Closed-loop smoke passed; starting 26-task evaluation", flush=True)
+    checkpoint, output = Path(args.checkpoint).resolve(), Path(args.output).resolve()
+    recipe = json.loads((checkpoint.parent / "recipe.json").read_text())
+    validate_robomem_recipe(recipe)
+    action_horizon = recipe["model"]["action_horizon"]
+    replan_steps = recipe.get("execution_horizon", 5)
+    if not 0 < replan_steps <= action_horizon or recipe["history_stride"] % replan_steps:
+        raise ValueError("Checkpoint has incompatible execution and history settings")
+    allowed_trials = (3, 10) if args.reference_prefix else (10,)
+    if args.trials not in allowed_trials or args.skip_tasks:
+        raise ValueError("Use all26 tasks and10 trials, or --reference-prefix --trials 3 for seeds50–52")
+    output.mkdir(parents=True, exist_ok=False)
+    metadata = dict(
+        checkpoint=str(checkpoint),
+        tasks=tasks,
+        skipped_tasks=sorted(set(args.skip_tasks)),
+        task25_seed_candidates=task25_candidates,
+        seeds=list(range(50, 50 + args.trials)),
+        expected_episodes=len(tasks) * args.trials,
+        reference_prefix_check=args.reference_prefix,
+        action_horizon=action_horizon,
+        replan_steps=replan_steps,
+        history_stride=recipe["history_stride"],
+        max_steps=2500,
+        renderer="osmesa",
+        instruction="fixed full-task rewrite 0 from training instruction table",
+        progress=PROGRESS_MODE,
+        base_port=args.base_port,
+        launch_gpu0=args.launch_gpu0,
+        benchmark=json.loads((BENCHMARK.parent / "OFFICIAL_SOURCE.json").read_text()),
+        code_sha256={
+            str(p.relative_to(ROOT)): hashlib.sha256(p.read_bytes()).hexdigest()
+            for p in [
+                ROOT / "src/openpi/policies/robomem_memory_inference.py",
+                ROOT / "src/openpi/policies/robomem_progress.py",
+                ROOT / "src/openpi/robomemarena_26tasks_subtask_prompts.py",
+                ROOT / "scripts/robomem_memory_eval_adapter.py",
+                ROOT / "scripts/run_robomem_eval_task.py",
+                ROOT / "scripts/run_robomem_task25_seed.py",
+                ROOT / "scripts/robomem_seed_search.py",
+                ROOT / "scripts/robomem_eval_paths.py",
+                Path(__file__).resolve(),
+            ]
+        },
+    )
+    metadata["recipe_sha256"] = hashlib.sha256((checkpoint.parent / "recipe.json").read_bytes()).hexdigest()
+    metadata["norm_stats_sha256"] = hashlib.sha256(
+        (checkpoint / "assets/pi05_libero/norm_stats.json").read_bytes()
+    ).hexdigest()
+    metadata["taskanchor_code_sha256"] = {
+        rel: hashlib.sha256((ROOT.parent / rel).read_bytes()).hexdigest()
+        for rel in ("taskanchor/pi05/taskanchor.py", "taskanchor/pi05/ttt.py",
+                    "taskanchor/data/history.py")
+    }
+    (output / "evaluation_config.json").write_text(json.dumps(metadata, indent=2))
+    lock = threading.Lock()
+    completed, failures, running = {}, {}, {}
+
+    def write_status():
+        result = dict(
+            expected_tasks=len(tasks),
+            expected_episodes=len(tasks) * args.trials,
+            skipped_tasks=metadata["skipped_tasks"],
+            completed_tasks=sorted(completed),
+            failed_tasks=failures,
+            running_tasks=running,
+            completed_episodes=sum(r["num_episodes"] for r in completed.values()),
+        )
+        scored = [r for r in completed.values() if r["num_episodes"] > 0]
+        if scored:
+            result.update(
+                TSR=sum(r["TSR"] for r in scored) / len(scored),
+                CSR=sum(r["CSR"] for r in scored) / len(scored),
+            )
+        temp = output / "status.tmp.json"
+        temp.write_text(json.dumps(result, indent=2))
+        temp.replace(output / "status.json")
+
+    def worker(gpu):
+        assigned_tasks = [task for task in range(gpu + 1, 27, args.workers) if task in tasks]
+        if not assigned_tasks:
+            return
+        port, server = args.base_port + gpu, None
+        if gpu != 0 or args.launch_gpu0:
+            env = dict(
+                os.environ,
+                CUDA_VISIBLE_DEVICES=str(gpu),
+                JAX_PLATFORMS="cuda",
+                XLA_PYTHON_CLIENT_PREALLOCATE="false",
+                OMP_NUM_THREADS="1",
+                OPENBLAS_NUM_THREADS="1",
+                PYTHONPYCACHEPREFIX="/tmp/robomem_memory_pycache",
+            )
+            with (output / f"server_gpu{gpu}.log").open("w") as log:
+                server = subprocess.Popen(
+                    [
+                        str(POLICY_PYTHON),
+                        "-u",
+                        "scripts/serve_robomem_memory.py",
+                        "--checkpoint",
+                        str(checkpoint),
+                        "--port",
+                        str(port),
+                    ],
+                    cwd=ROOT,
+                    env=env,
+                    stdout=log,
+                    stderr=subprocess.STDOUT,
+                    stdin=subprocess.DEVNULL,
+                )
+            (output / f"server_gpu{gpu}.pid").write_text(str(server.pid))
+        try:
+            deadline = time.monotonic() + 300
+            while True:
+                try:
+                    with urllib.request.urlopen(f"http://127.0.0.1:{port}/healthz", timeout=2) as response:
+                        if response.status == 200:
+                            break
+                except OSError:
+                    if (server and server.poll() is not None) or time.monotonic() > deadline:
+                        raise RuntimeError(f"GPU {gpu} policy server unavailable")
+                    time.sleep(2)
+            server_metadata = verify_server_checkpoint(port, checkpoint)
+            if server_metadata["action_horizon"] != action_horizon:
+                raise RuntimeError("Server action horizon does not match the checkpoint recipe")
+            if server_metadata.get("progress") != PROGRESS_MODE:
+                raise RuntimeError("Server progress mode does not match evaluation protocol")
+            (output / f"server_gpu{gpu}_metadata.json").write_text(json.dumps(server_metadata, indent=2))
+            for task in assigned_tasks:
+                with lock:
+                    running[gpu] = task
+                    write_status()
+                print(f"START gpu={gpu} task={task}", flush=True)
+                try:
+                    if task == 25 and task25_candidates is not None:
+                        result = run_task25(gpu, port, args.trials, task25_candidates,
+                                            output / "task25", replan_steps)
+                    else:
+                        result = run_task(task, gpu, port, args.trials, 50, output / f"task{task:02d}",
+                                          replan_steps=replan_steps)
+                    with lock:
+                        completed[task] = result
+                    print(f"COMPLETE task={task} {json.dumps(result)}", flush=True)
+                except Exception as exc:
+                    with lock:
+                        failures[task] = str(exc)
+                    print(f"ERROR task={task}: {exc}", flush=True)
+                finally:
+                    with lock:
+                        running.pop(gpu, None)
+                        write_status()
+        except Exception as exc:
+            with lock:
+                for task in assigned_tasks:
+                    if task not in completed:
+                        failures[task] = str(exc)
+                write_status()
+        finally:
+            if server:
+                server.terminate()
+                try:
+                    server.wait(timeout=30)
+                except subprocess.TimeoutExpired:
+                    server.kill()
+                    server.wait()
+
+    write_status()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool:
+        futures = [pool.submit(worker, gpu) for gpu in range(args.workers)]
+        for future in futures:
+            future.result()
+    episodes, task_results = [], []
+    for task in sorted(completed):
+        rows = json.loads((output / f"task{task:02d}/summary.json").read_text())
+        task_results.extend(rows)
+        for row in rows:
+            episodes.extend(dict(task_id=task, **ep) for ep in row["episodes"])
+    scored = [r for r in completed.values() if r["num_episodes"] > 0]
+    report = dict(
+        metadata=metadata,
+        complete=set(completed) == set(tasks) and not failures,
+        failures=failures,
+        task_results=task_results,
+        episodes=episodes,
+        actual_episodes=len(episodes),
+        shortfalls={task: args.trials - r["num_episodes"] for task, r in completed.items()
+                    if r["num_episodes"] < args.trials},
+        metric_tasks=[task for task, r in sorted(completed.items()) if r["num_episodes"] > 0],
+        TSR=sum(r["TSR"] for r in scored) / len(scored) if scored else None,
+        CSR=sum(r["CSR"] for r in scored) / len(scored) if scored else None,
+    )
+    (output / "final_report.json").write_text(json.dumps(report, indent=2))
+    print(f"EVALUATION FINISHED complete={report['complete']} TSR={report['TSR']} CSR={report['CSR']}", flush=True)
+    if not report["complete"]:
+        raise RuntimeError("Some tasks failed; see final_report.json")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--workers", type=int, default=8, choices=range(1, 9))
+    parser.add_argument("--trials", type=int, default=10)
+    parser.add_argument("--reference-prefix", action="store_true",
+                        help="Allow3 trials per task to compare the seeds50–52 prefix of the recorded50k run")
+    parser.add_argument("--wait-smoke")
+    parser.add_argument("--base-port", type=int, default=8600)
+    parser.add_argument("--launch-gpu0", action="store_true")
+    parser.add_argument("--skip-tasks", type=int, nargs="*", default=[], choices=range(1, 27))
+    parser.add_argument("--task25-seeds", default=str(
+        ROOT.parent / "evidence/task25_seed_candidates.json"
+    ), help="Reference initialization seed list; policy failures are never retried")
+    main(parser.parse_args())

code/pi05/scripts/extract_robomem_memory.py

@@ -0,0 +1,169 @@
+"""Extract frozen pi05_base main-camera features before memory training."""
+
+import argparse
+import json
+import logging
+import os
+from pathlib import Path
+import time
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+
+from taskanchor.data.history import same_feature_provenance
+from openpi.models import siglip
+from openpi.models.model import restore_params
+from openpi_client.image_tools import resize_with_pad
+from openpi.training.robomem_memory_data import (
+    BASE_PARAMS,
+    CACHE_ROOT,
+    DATA_ROOT,
+    EpisodeIndex,
+    cache_file,
+    decode_image,
+    episode_path,
+    validate_cache,
+    validate_feature_file,
+)
+
+
+class ExtractionImages(torch.utils.data.IterableDataset):
+    def __init__(self, index, groups):
+        self.index, self.groups = index, groups
+
+    def __iter__(self):
+        torch.set_num_threads(1)
+        worker = torch.utils.data.get_worker_info()
+        groups = self.groups if worker is None else self.groups[worker.id :: worker.num_workers]
+        for task, seed in groups:
+            for ep in self.index.groups[(task, seed)]:
+                record = self.index.records[ep]
+                table = pq.read_table(episode_path(self.index.root, ep), columns=["image"])
+                for local in range(len(table)):
+                    yield dict(
+                        image=resize_with_pad(decode_image(table["image"][local].as_py()), 224, 224),
+                        task=task,
+                        seed=seed,
+                        frame=record.prefix_frames + local,
+                    )
+
+
+def main(args):
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s", force=True)
+    index = EpisodeIndex(args.dataset)
+    cache = Path(args.cache)
+    cache.mkdir(parents=True, exist_ok=True)
+    expected = index.manifest(args.base_params)
+    manifest = cache / "manifest.json"
+    if manifest.exists() and not same_feature_provenance(json.loads(manifest.read_text()), expected):
+        raise ValueError("Existing cache has different provenance")
+    if not manifest.exists():
+        manifest.write_text(json.dumps(expected, indent=2))
+    if args.verify_only:
+        validate_cache(index, cache, args.base_params)
+        logging.info("CACHE VERIFIED: %s", expected)
+        return
+    pending = []
+    for key in sorted(index.groups):
+        path = cache_file(cache, key)
+        if path.exists():
+            length = index.records[index.groups[key][0]].total_frames
+            validate_feature_file(path, length)
+        else:
+            pending.append(key)
+    if args.limit_tasks:
+        pending = pending[: args.limit_tasks]
+    if not pending:
+        validate_cache(index, cache, args.base_params)
+        return
+    devices = jax.devices()
+    if any(d.platform != "gpu" for d in devices):
+        raise RuntimeError("Extraction requires GPUs; refusing a silent CPU fallback")
+    if args.batch_size % len(devices):
+        raise ValueError("Extraction batch size must be divisible by GPU count")
+    logging.info("Loading pi05_base encoder; devices=%s pending_tasks=%d", devices, len(pending))
+    params = restore_params(args.base_params, restore_type=np.ndarray)["PaliGemma"]["img"]
+    mesh = jax.sharding.Mesh(np.asarray(devices), ("batch",))
+    replicated = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("batch"))
+    params = jax.device_put(jax.tree.map(lambda x: x.astype(jnp.bfloat16), params), replicated)
+    encoder = siglip.Module(num_classes=2048, variant="So400m/14", pool_type="none", scan=True, dtype_mm="bfloat16")
+
+    def encode(p, images):
+        images = images.astype(jnp.float32) / 255 * 2 - 1
+        tokens, _ = encoder.apply({"params": p}, images, train=False, return_encoded=True)
+        # Pool in FP32 then quantize to BF16; float16 storage preserves these values.
+        spatial = jax.image.resize(
+            tokens.reshape(-1, 16, 16, 1152).astype(jnp.float32),
+            (images.shape[0], 4, 4, 1152),
+            "linear",
+            antialias=False,
+        )
+        return spatial.reshape(-1, 16, 1152).astype(jnp.bfloat16).astype(jnp.float16)
+
+    encode = jax.jit(encode, in_shardings=(replicated, data_sharding), out_shardings=data_sharding)
+    loader = torch.utils.data.DataLoader(
+        ExtractionImages(index, pending),
+        batch_size=args.batch_size,
+        num_workers=args.workers,
+        multiprocessing_context="spawn" if args.workers else None,
+        prefetch_factor=2 if args.workers else None,
+    )
+    buffers, completed, frames = {}, 0, 0
+    started = time.monotonic()
+    for batch in loader:
+        images = batch["image"].numpy()
+        count = len(images)
+        if count < args.batch_size:
+            images = np.pad(images, ((0, args.batch_size - count), (0, 0), (0, 0), (0, 0)), mode="edge")
+        history = jax.device_get(encode(params, jax.device_put(images, data_sharding)))
+        for i in range(count):
+            key = int(batch["task"][i]), int(batch["seed"][i])
+            frame = int(batch["frame"][i])
+            if key not in buffers:
+                total = index.records[index.groups[key][0]].total_frames
+                buffers[key] = dict(history=np.empty((total, 16, 1152), np.float16), seen=np.zeros(total, bool))
+            buf = buffers[key]
+            if buf["seen"][frame]:
+                raise ValueError(f"Duplicate frame {key} {frame}")
+            buf["history"][frame] = history[i]
+            buf["seen"][frame] = True
+            if buf["seen"].all():
+                if not np.isfinite(buf["history"]).all():
+                    raise ValueError(f"Nonfinite encoder output {key}")
+                path = cache_file(cache, key)
+                temp = path.with_suffix(".tmp.npy")
+                np.save(temp, buf["history"], allow_pickle=False)
+                os.replace(temp, path)
+                del buffers[key]
+                completed += 1
+        frames += count
+        if frames == count or frames % (args.batch_size * 20) == 0:
+            logging.info(
+                "extracted_frames=%d tasks=%d/%d fps=%.1f",
+                frames,
+                completed,
+                len(pending),
+                frames / (time.monotonic() - started),
+            )
+    if buffers or completed != len(pending):
+        raise RuntimeError("Extraction ended with incomplete task buffers")
+    if not args.limit_tasks:
+        validate_cache(index, cache, args.base_params)
+        (cache / "VERIFIED.json").write_text(json.dumps(expected, indent=2))
+    logging.info("EXTRACTION COMPLETE: tasks=%d frames=%d", completed, frames)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", default=str(DATA_ROOT))
+    parser.add_argument("--base-params", default=str(BASE_PARAMS))
+    parser.add_argument("--cache", default=str(CACHE_ROOT))
+    parser.add_argument("--batch-size", type=int, default=256)
+    parser.add_argument("--workers", type=int, default=8)
+    parser.add_argument("--limit-tasks", type=int, default=0)
+    parser.add_argument("--verify-only", action="store_true")
+    main(parser.parse_args())

code/pi05/scripts/robomem_eval_paths.py

@@ -0,0 +1,20 @@
+"""Portable filesystem locations; numerical evaluation settings live elsewhere."""
+
+import os
+from pathlib import Path
+import sys
+
+ROOT = Path(__file__).resolve().parents[1]
+BENCHMARK = Path(os.environ.get(
+    "TASKANCHOR_ROBOMEM_BENCHMARK", ROOT.parent / "RoboMemArena/evaluation_benchmark"
+)).resolve()
+PYTHON = Path(os.environ.get(
+    "TASKANCHOR_ROBOMEM_EVAL_PYTHON", ROOT.parent / "robomem_eval_env/bin/python"
+)).absolute()  # Keep the virtualenv entry point; resolving its symlink loses its packages.
+POLICY_PYTHON = os.environ.get("TASKANCHOR_PI05_PYTHON", sys.executable)
+LIBERO_CONFIG = Path(os.environ.get(
+    "TASKANCHOR_LIBERO_CONFIG", ROOT / "data/robomem_eval_libero_config"
+)).resolve()
+RENDER_LIBS = Path(os.environ.get(
+    "TASKANCHOR_RENDER_LIBS", ROOT / "data/robomem_render_libs/usr/lib/x86_64-linux-gnu"
+)).resolve()

code/pi05/scripts/robomem_memory_eval_adapter.py

@@ -0,0 +1,93 @@
+"""RoboMemArena adapter: whole-task instruction, predicted progress, causal memory."""
+
+import json
+from pathlib import Path
+import sys
+import time
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / "packages/openpi-client/src"))
+sys.path.insert(0, str(ROOT / "src"))
+from openpi_client.websocket_client_policy import WebsocketClientPolicy
+from openpi.robomemarena_26tasks_10_full_prompts import ROBOMEMARENA_26_10_FULL_PROMPTS
+from openpi.policies.robomem_progress import PROGRESS_MODE, stage_values
+from policy_adapter import BasePolicyAdapter
+
+
+class MemoryEvalAdapter(BasePolicyAdapter):
+    def __init__(self, host="127.0.0.1", port=8600, replan_steps=None, trace_path=None, task_id=None):
+        self.client = WebsocketClientPolicy(host=host, port=int(port))
+        metadata = self.client.get_server_metadata()
+        if metadata.get("progress") != PROGRESS_MODE:
+            self.client._ws.close()
+            raise ValueError("Evaluation requires a nearest-stage progress policy server")
+        replan_steps = metadata.get("execution_horizon", 5) if replan_steps is None else replan_steps
+        if not 0 < replan_steps <= metadata["action_horizon"] or metadata["history_stride"] % replan_steps:
+            self.client._ws.close()
+            raise ValueError("Execution length must divide history stride and not exceed prediction horizon")
+        self.replan_steps = replan_steps
+        self.task_id = task_id
+        self.episode_id = -1
+        self.trace_path = Path(trace_path) if trace_path else None
+        self.frame = 0
+
+    def set_task_context(self, *, task_id=None, episode_id=None, **kwargs):
+        del kwargs
+        self.task_id, self.episode_id = task_id, episode_id
+
+    def reset(self):
+        self.frame = 0
+        self.episode_id += 1
+
+    def infer_actions(self, obs, prompt, resize_size):
+        del resize_size
+        task = self.task_id if self.task_id is not None else obs.get("task_id")
+        stage_values(task)
+        # A fixed full-task rewrite is the commanded instruction, never an oracle subtask.
+        if task is not None:
+            prompt = ROBOMEMARENA_26_10_FULL_PROMPTS[int(task)][0]
+        state = np.asarray(obs["observation/state"], np.float32).reshape(-1)
+        state = np.pad(state, (0, max(0, 8 - len(state))))[:8]
+        payload = {
+            "observation/image": obs["observation/image"],
+            "observation/wrist_image": obs["observation/wrist_image"],
+            "observation/state": state,
+            "prompt": prompt,
+            "task_id": int(task),
+            "environment_frame": self.frame,
+            "reset": self.frame == 0,
+        }
+        infer_started = time.perf_counter()
+        result = self.client.infer(payload)
+        round_trip_ms = (time.perf_counter() - infer_started) * 1000
+        if self.trace_path:
+            self.trace_path.parent.mkdir(parents=True, exist_ok=True)
+            with self.trace_path.open("a") as stream:
+                stream.write(
+                    json.dumps(
+                        dict(
+                            task_id=int(task) if task is not None else None,
+                            episode_id=self.episode_id,
+                            frame=self.frame,
+                            progress=float(result["progress"]),
+                            progress_raw=float(result["progress_raw"]),
+                            progress_stage=int(result["progress_stage"]),
+                            progress_num_stages=int(result["progress_num_stages"]),
+                            history_frames=int(result["history_frames"]),
+                            server_infer_ms=result.get("server_timing", {}).get("infer_ms"),
+                            round_trip_ms=round_trip_ms,
+                        )
+                    )
+                    + "\n"
+                )
+        self.frame += self.replan_steps
+        return np.asarray(result["actions"], np.float32)
+
+    def close(self):
+        self.client._ws.close()
+
+
+def build_adapter(**kwargs):
+    return MemoryEvalAdapter(**kwargs)

code/pi05/scripts/robomem_seed_search.py

@@ -0,0 +1,57 @@
+"""Bounded Task25 seed enumeration; unsuccessful policies remain scored episodes."""
+
+import json
+from pathlib import Path
+
+from run_robomem_eval_task import TaskInitializationError, run_task
+
+
+def run_task25(gpu, port, trials, candidates, output, replan_steps):
+    output = Path(output)
+    output.mkdir(parents=True, exist_ok=False)
+    if len(candidates) > 200 or len(candidates) != len(set(candidates)):
+        raise ValueError("At most 200 distinct seed candidates are allowed")
+    audit = dict(target_episodes=trials, max_seed_attempts=len(candidates), attempts=[], accepted_seeds=[])
+
+    def save():
+        temporary = output / "seed_search.tmp.json"
+        temporary.write_text(json.dumps(audit, indent=2))
+        temporary.replace(output / "seed_search.json")
+
+    rows, episodes = [], []
+    save()
+    for seed in candidates:
+        case = output / f"seed{seed:04d}"
+        try:
+            run_task(25, gpu, port, 1, seed, case, replan_steps=replan_steps, enumerate_seed=True)
+        except TaskInitializationError as exc:
+            audit["attempts"].append(dict(seed=seed, status="initialization_error", error=str(exc)))
+            save()
+            continue
+        except Exception as exc:
+            audit["attempts"].append(dict(seed=seed, status="runtime_error", error=str(exc)))
+            save()
+            raise
+        row = json.loads((case / "summary.json").read_text())[0]
+        if len(row["episodes"]) != 1 or row["episodes"][0]["seed"] != seed:
+            raise RuntimeError("Single-seed scorer returned inconsistent coverage")
+        ep = dict(row["episodes"][0], ep=len(episodes), video_dir=row["video_dir"])
+        episodes.append(ep)
+        rows.append(row)
+        audit["accepted_seeds"].append(seed)
+        audit["attempts"].append(dict(seed=seed, status="scored", TSR=ep["TSR"], CSR=ep["CSR"]))
+        save()
+        if len(episodes) == trials:
+            break
+    audit.update(actual_episodes=len(episodes), exhausted=len(audit["attempts"]) == len(candidates),
+                 shortfall=trials - len(episodes))
+    save()
+    tsr = sum(ep["TSR"] for ep in episodes) / len(episodes) if episodes else None
+    csr = sum(ep["CSR"] for ep in episodes) / len(episodes) if episodes else None
+    summary = dict(rows[0]) if rows else dict(task_id=25, task_key="task25")
+    summary.update(episodes=episodes, TSR=tsr, CSR=csr, seed_search=audit,
+                   video_dir=str(output), num_episodes=len(episodes))
+    aggregate = dict(num_tasks=1, num_episodes=len(episodes), TSR=tsr, CSR=csr, seed_search=audit)
+    (output / "summary.json").write_text(json.dumps([summary], indent=2))
+    (output / "aggregate.json").write_text(json.dumps(aggregate, indent=2))
+    return aggregate

code/pi05/scripts/run_robomem_eval_task.py

@@ -0,0 +1,103 @@
+"""Launch the unmodified official scorer for one task."""
+
+import argparse
+import json
+import os
+from pathlib import Path
+import subprocess
+
+from robomem_eval_paths import ROOT, BENCHMARK, PYTHON, LIBERO_CONFIG, RENDER_LIBS
+
+
+class TaskInitializationError(RuntimeError):
+    """An enumerated seed failed before any policy action was produced."""
+
+
+def run_task(task, gpu, port, trials, seed, output, replan_steps=5, enumerate_seed=False):
+    if replan_steps <= 0 or 10 % replan_steps:
+        raise ValueError("Replan steps must divide history stride 10")
+    output = Path(output).resolve()
+    output.mkdir(parents=True, exist_ok=True)
+    env = dict(
+        os.environ,
+        MUJOCO_GL="osmesa",
+        PYOPENGL_PLATFORM="osmesa",
+        MUJOCO_EGL_DEVICE_ID=str(gpu),
+        OMP_NUM_THREADS="1",
+        OPENBLAS_NUM_THREADS="1",
+        NUMBA_CACHE_DIR="/tmp/robomem_numba",
+        LIBERO_CONFIG_PATH=str(LIBERO_CONFIG),
+        PYTHONPATH=str(BENCHMARK / "libero_fork"),
+        LD_LIBRARY_PATH=str(RENDER_LIBS)
+        + ":"
+        + os.environ.get("LD_LIBRARY_PATH", ""),
+        LP_NUM_THREADS="2",
+    )
+    kwargs = dict(port=port, replan_steps=replan_steps, task_id=task, trace_path=str(output / "progress.jsonl"))
+    command = [
+        str(PYTHON),
+        "-u",
+        str(ROOT / "scripts/run_robomem_task25_seed.py") if enumerate_seed else "scripts/run_all_tasks1_26.py",
+        "--adapter-spec",
+        str(ROOT / "scripts/robomem_memory_eval_adapter.py") + ":build_adapter",
+        "--adapter-kwargs",
+        json.dumps(kwargs),
+        "--task-start",
+        str(task),
+        "--task-end",
+        str(task),
+        "--num-trials-per-task",
+        str(trials),
+        "--seed",
+        str(seed),
+        "--replan-steps",
+        str(replan_steps),
+        "--max-steps",
+        "2500",
+        "--out-root",
+        str(output),
+    ]
+    (output / "command.json").write_text(json.dumps(command, indent=2))
+    for attempt in range(1, 2 if enumerate_seed else 4):
+        with (output / "run.log").open("w", buffering=1) as log:
+            process = subprocess.run(command, cwd=BENCHMARK, env=env, stdout=log, stderr=subprocess.STDOUT)
+        log_text = (output / "run.log").read_text()
+        trace = output / "progress.jsonl"
+        if (enumerate_seed and process.returncode == 86 and "TASK25_INITIALIZATION_ERROR:" in log_text
+                and (not trace.exists() or trace.stat().st_size == 0)):
+            raise TaskInitializationError(f"Task25 seed {seed}: initialization failed; see {output / 'run.log'}")
+        # The official constructor samples before any episode seed is set. Its
+        # reset() already retries invalid placements, but its constructor does not.
+        # Retry only this startup failure, never an episode or a policy failure.
+        startup_placement_error = (
+            process.returncode != 0
+            and "robosuite.utils.errors.RandomizationError: Cannot place all objects" in log_text
+            and "env = OffScreenRenderEnv(" in log_text
+            and "Episode " not in log_text
+            and not (output / "progress.jsonl").exists()
+            and not (output / "summary.json").exists()
+        )
+        if not startup_placement_error or attempt == 3:
+            break
+        (output / f"startup_failure_attempt{attempt}.log").write_text(log_text)
+        print(f"Task {task}: constructor placement failed before episode 0; retry {attempt}/2", flush=True)
+    if process.returncode or "Traceback (most recent call last)" in log_text or "Episode failed:" in log_text:
+        raise RuntimeError(f"Task {task} evaluation encountered a runtime error; see {output / 'run.log'}")
+    result = json.loads((output / "aggregate.json").read_text())
+    if result["num_episodes"] != trials:
+        raise RuntimeError(f"Task {task}: incomplete episode coverage")
+    return result
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--task", type=int, required=True)
+    parser.add_argument("--gpu", type=int, default=0)
+    parser.add_argument("--port", type=int, default=8600)
+    parser.add_argument("--trials", type=int, default=1)
+    parser.add_argument("--seed", type=int, default=50)
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--replan-steps", type=int, default=5)
+    args = parser.parse_args()
+    print(json.dumps(run_task(args.task, args.gpu, args.port, args.trials, args.seed, args.output,
+                              replan_steps=args.replan_steps)), flush=True)

code/pi05/scripts/run_robomem_task25_seed.py

@@ -0,0 +1,36 @@
+"""Official scorer with deterministic, single-attempt Task25 initialization."""
+
+from pathlib import Path
+import sys
+
+import numpy as np
+
+from robomem_eval_paths import BENCHMARK
+sys.path.insert(0, str(BENCHMARK / "scripts"))
+import run_all_tasks1_26 as official
+
+
+def main():
+    args = official.build_argparser().parse_args()
+    from robosuite.utils.errors import RandomizationError
+    if args.task_start != 25 or args.task_end != 25 or args.num_trials_per_task != 1:
+        raise ValueError("This entry point is only for one Task25 seed")
+    np.random.seed(args.seed)  # Also seed the constructor, which precedes the episode loop.
+    base = official.ec._get_env_class()
+
+    class SingleAttemptEnvironment(base):
+        def reset(self):
+            # Original ControlEnv.reset retries placement errors indefinitely.
+            # Return the same observation on success; reject this seed on failure.
+            return self.env.reset()
+
+    official.ec._get_env_class = lambda: SingleAttemptEnvironment
+    try:
+        official.main()
+    except RandomizationError as exc:
+        print(f"TASK25_INITIALIZATION_ERROR: {exc}", flush=True)
+        raise SystemExit(86) from exc
+
+
+if __name__ == "__main__":
+    main()

code/pi05/scripts/serve_robomem_memory.py

@@ -0,0 +1,16 @@
+import argparse
+import logging
+
+from openpi.policies.robomem_memory_inference import RoboMemMemoryPolicy
+from openpi.serving.websocket_policy_server import WebsocketPolicyServer
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--port", type=int, default=8600)
+    args = parser.parse_args()
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s", force=True)
+    policy = RoboMemMemoryPolicy(args.checkpoint)
+    logging.info("Memory policy loaded: %s", policy.metadata)
+    WebsocketPolicyServer(policy, host="127.0.0.1", port=args.port, metadata=policy.metadata).serve_forever()

code/pi05/scripts/test_robomem_discrete_progress.py

@@ -0,0 +1,58 @@
+"""CPU checks for legal progress values and the action-conditioning path."""
+import unittest
+import numpy as np
+from openpi.policies.robomem_progress import discretize_progress, stage_values
+
+
+class ProgressTests(unittest.TestCase):
+    def test_four_stages_and_boundaries(self):
+        self.assertEqual(stage_values(1), (0.0, 0.25, 0.5, 0.75))
+        for raw, expected in [(0, 0), (.124, 0), (.125, 0), (.126, .25),
+                              (.375, .25), (.62, .5), (.625, .5), (.626, .75),
+                              (.86, .75), (1, .75), (-.1, 0), (1.1, .75)]:
+            self.assertEqual(discretize_progress(raw, 1)[0], expected)
+
+    def test_all_tasks_match_training_labels(self):
+        from openpi.robomemarena_26tasks_subtask_prompts import ROBOMEMARENA_26_SUBTASK_PROMPTS
+        for task, stages in ROBOMEMARENA_26_SUBTASK_PROMPTS.items():
+            expected = tuple(round(i / len(stages), 2) for i in range(len(stages)))
+            self.assertEqual(stage_values(task), expected)
+            for i, value in enumerate(expected):
+                self.assertEqual(discretize_progress(value, task), (value, i))
+            for value in np.linspace(0, 1, 101):
+                self.assertIn(discretize_progress(value, task)[0], expected)
+
+    def test_reject_invalid_inputs(self):
+        for task in [None, 0, 27, True, "1"]:
+            with self.assertRaises(ValueError):
+                discretize_progress(.5, task)
+        for raw in [float("nan"), float("inf"), -float("inf")]:
+            with self.assertRaises(FloatingPointError):
+                discretize_progress(raw, 1)
+
+    def test_action_sampler_receives_discrete_prompt(self):
+        from types import SimpleNamespace
+        import jax
+        from openpi.policies.robomem_memory_inference import RoboMemMemoryPolicy
+        policy = object.__new__(RoboMemMemoryPolicy)
+        policy.action_horizon = 20
+        policy.rng = jax.random.key(42)
+        policy._observation = lambda raw, prompt: SimpleNamespace(state=np.zeros((1, 8)), prompt=prompt)
+        policy.encode = lambda obs: (None, None, np.zeros((1, 16, 1152)))
+        policy.history = SimpleNamespace(add_and_read=lambda frame, spatial:
+            (np.zeros((16, 1152)), np.zeros((1, 16, 1152)), np.ones(1, bool)))
+        policy.refine = lambda *args: (None, np.array([.86]))
+        def sample(key, obs, refined, mask):
+            self.assertEqual(obs.prompt, "test task.Progress:0.75")
+            return np.zeros((1, 20, 7))
+        policy.sample = sample
+        policy.outputs = lambda data: data
+        result = policy.infer(dict(task_id=1, environment_frame=0, prompt="test task"))
+        self.assertEqual(result["progress"], .75)
+        self.assertEqual(result["progress_raw"], .86)
+        self.assertEqual(result["progress_stage"], 3)
+        self.assertEqual(result["actions"].shape, (20, 7))
+
+
+if __name__ == "__main__":
+    unittest.main()

code/pi05/scripts/test_robomem_memory.py

@@ -0,0 +1,163 @@
+import importlib.util
+from pathlib import Path
+import unittest
+import tempfile
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+from taskanchor.pi05.taskanchor import PrimaryMemoryAdapter, MemoryTTT
+from openpi.training.robomem_memory_data import (
+    EpisodeIndex,
+    MemoryDataset,
+    history_indices,
+    cache_file,
+    validate_feature_file,
+)
+
+spec = importlib.util.spec_from_file_location("memory_trainer", Path(__file__).with_name("train_robomem_memory.py"))
+trainer = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(trainer)
+
+
+class MemoryTests(unittest.TestCase):
+    def test_primary_residual_and_masked_history(self):
+        adapter = PrimaryMemoryAdapter(rngs=nnx.Rngs(0), input_dim=32, hidden_dim=32, depth=2, heads=4)
+        current = jax.random.normal(jax.random.key(1), (2, 3, 4, 32)).astype(jnp.bfloat16)
+        frame0 = jax.random.normal(jax.random.key(2), (2, 16, 32)).astype(jnp.bfloat16)
+        history = jax.random.normal(jax.random.key(3), (2, 5, 16, 32)).astype(jnp.bfloat16)
+        mask = jnp.array([[1, 1, 0, 0, 0]] * 2, bool)
+        cmask = jnp.ones((2, 3, 4), bool)
+        refined, progress = adapter(current, cmask, history, mask, frame0=frame0)
+        np.testing.assert_array_equal(refined[:, 1:], current[:, 1:])
+        self.assertGreater(float(jnp.linalg.norm((refined[:, 0] - current[:, 0]).astype(jnp.float32))), 0)
+        altered, p2 = adapter(current, cmask, history.at[:, 2:].set(100), mask, frame0=frame0)
+        np.testing.assert_array_equal(refined, altered)
+        np.testing.assert_array_equal(progress, p2)
+        grad = jax.grad(lambda h: adapter(current, cmask, h, mask, frame0=frame0)[0][:, 0, 0, 0].astype(jnp.float32).sum())(
+            history
+        )
+        self.assertGreater(float(jnp.linalg.norm(grad[:, :2].astype(jnp.float32))), 0)
+        self.assertTrue(np.all(np.linalg.norm(np.asarray(grad[:, :2], np.float32), axis=-1) > 0))
+        np.testing.assert_array_equal(grad[:, 2:], jnp.zeros_like(grad[:, 2:]))
+        masked, _ = adapter(current, cmask.at[:, 0].set(False), history, mask, frame0=frame0)
+        np.testing.assert_array_equal(masked, current)
+
+    def test_current_never_writes_fast_weights(self):
+        ttt = MemoryTTT(32, 4, nnx.Rngs(5))
+        context = jax.random.normal(jax.random.key(6), (1, 6, 32)).astype(jnp.bfloat16)
+        write = jnp.array([[0, 0, 1, 1, 1, 1]], bool)
+        positions = jnp.arange(6)[None].astype(jnp.float32)
+        original = ttt(context, write, positions, 2)
+        changed = ttt(context.at[:, 0].multiply(10), write, positions, 2)
+        np.testing.assert_array_equal(original[:, 1], changed[:, 1])
+
+    def test_history_strictly_past_and_across_subtasks(self):
+        for t, expected in [(0, []), (10, []), (11, [1]), (31, [1, 11, 21])]:
+            self.assertEqual(history_indices(t).tolist(), expected)
+        self.assertEqual(len(history_indices(1917)), 191)
+        index = EpisodeIndex()
+        self.assertEqual(len(index.groups), 2600)
+        self.assertEqual(index.max_task_frames, 1918)
+        self.assertEqual(index.history_frames, 191)
+        for episodes in index.groups.values():
+            previous_end = 0
+            for ep in episodes:
+                r = index.records[ep]
+                self.assertEqual(r.prefix_frames, previous_end)
+                previous_end += r.length
+                self.assertTrue(np.all(history_indices(r.prefix_frames) < r.prefix_frames))
+            self.assertEqual(previous_end, r.total_frames)
+            for t in (0, previous_end // 2, previous_end - 1):
+                expected = list(reversed(range(t - 10, 0, -10)))
+                self.assertEqual(history_indices(t, capacity=index.history_frames).tolist(), expected)
+
+    def test_spatial_cache_dataset_preserves_frames_and_patches(self):
+        index = EpisodeIndex()
+        key = sorted(index.groups)[0]
+        episodes = index.groups[key]
+        total = index.records[episodes[0]].total_frames
+        cfg = trainer.config.get_config("pi05_robomemarena_26tasks_subtask_only_full_sft_eval")
+        with tempfile.TemporaryDirectory() as directory:
+            path = cache_file(directory, key)
+            spatial = np.broadcast_to(np.arange(total, dtype=np.float16)[:, None, None], (total, 16, 1152)).copy()
+            spatial[:, :, 0] = np.arange(16, dtype=np.float16)
+            np.save(path, spatial)
+            validate_feature_file(path, total)
+            dataset = MemoryDataset(index, cfg.data.create(cfg.assets_dirs, cfg.model), cfg.model, directory)
+            ep = episodes[-1]
+            pos = index.episodes.index(ep)
+            sample = dataset[int(index.ends[pos]) - 1]
+            selected = history_indices(total - 1)
+            self.assertEqual(sample["history"].shape, (191, 16, 1152))
+            np.testing.assert_array_equal(sample["history_mask"], np.arange(191) < len(selected))
+            np.testing.assert_array_equal(sample["history"][: len(selected)], spatial[selected])
+            np.testing.assert_array_equal(sample["history"][len(selected) :], 0)
+            np.testing.assert_array_equal(sample["frame0"], spatial[0])
+            self.assertIsInstance(dataset.features[key], np.memmap)
+            np.save(path, spatial.mean(axis=1))
+            with self.assertRaisesRegex(ValueError, "shape/dtype"):
+                validate_feature_file(path, total)
+
+    def test_siglip_head_recomposition(self):
+        from openpi.models.siglip import Module
+
+        encoder = Module(num_classes=16, variant="mu/14", pool_type="none", scan=True, head_zeroinit=False)
+        images = jax.random.normal(jax.random.key(7), (1, 28, 28, 3))
+        params = encoder.init(jax.random.key(8), images, train=False)
+        original, _ = encoder.apply(params, images, train=False)
+        native, _ = encoder.apply(params, images, train=False, return_encoded=True)
+        recomposed = encoder.apply(params, native, method="project_encoded")
+        self.assertGreater(float(jnp.linalg.norm(original)), 0)
+        np.testing.assert_array_equal(original, recomposed)
+
+    def test_repeated_initialization_has_identical_static_graph(self):
+        def create():
+            return PrimaryMemoryAdapter(rngs=nnx.Rngs(0), input_dim=32, hidden_dim=32, depth=2, heads=4)
+
+        self.assertEqual(nnx.graphdef(create()), nnx.graphdef(create()))
+
+    def test_exact_optimizer_phase_boundary(self):
+        class Tiny(nnx.Module):
+            def __init__(self):
+                self.adapter = nnx.Dict(weight=nnx.Param(jnp.ones(2)))
+                self.backbone = nnx.Dict(
+                    weight=nnx.Param(jnp.ones(2)),
+                    PaliGemma=nnx.Dict(
+                        img=nnx.Dict(head=nnx.Dict(kernel=nnx.Param(jnp.ones(2))), encoder=nnx.Param(jnp.ones(2)))
+                    ),
+                )
+
+        params = nnx.state(Tiny()).filter(trainer.trainable)
+        self.assertNotIn("encoder", params["backbone"]["PaliGemma"]["img"])
+        tx = trainer.make_optimizer()
+        state = tx.init(params)
+        grads = jax.tree.map(jnp.ones_like, params)
+
+        @jax.jit
+        def update(p, s):
+            updates, s = tx.update(grads, s, p)
+            return optax.apply_updates(p, updates), s
+
+        initial = params.to_pure_dict()
+        for _ in range(5000):
+            params, state = update(params, state)
+        np.testing.assert_array_equal(params["backbone"]["weight"].value, initial["backbone"]["weight"])
+        np.testing.assert_array_equal(
+            params["backbone"]["PaliGemma"]["img"]["head"]["kernel"].value,
+            initial["backbone"]["PaliGemma"]["img"]["head"]["kernel"],
+        )
+        self.assertFalse(np.array_equal(params["adapter"]["weight"].value, initial["adapter"]["weight"]))
+        params, _ = update(params, state)
+        self.assertFalse(np.array_equal(params["backbone"]["weight"].value, initial["backbone"]["weight"]))
+        self.assertFalse(
+            np.array_equal(
+                params["backbone"]["PaliGemma"]["img"]["head"]["kernel"].value,
+                initial["backbone"]["PaliGemma"]["img"]["head"]["kernel"],
+            )
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

code/pi05/scripts/test_robomem_memory_inference.py

@@ -0,0 +1,54 @@
+import unittest
+import numpy as np
+from openpi.policies.robomem_memory_inference import FIFOHistory
+
+
+class OnlineHistoryTests(unittest.TestCase):
+    def test_environment_stride_and_rollover(self):
+        bank = FIFOHistory()
+        for t in range(0, 2051, 5):
+            spatial = np.full((16, 1152), t, np.float16)
+            spatial[:, 0] = np.arange(16)
+            memory, mask = bank.add_and_read(t, spatial)
+            expected = ([0] + list(reversed(range(t - 10, 0, -10))))[-191:]
+            np.testing.assert_array_equal(memory[mask, 0, 1], expected)
+            np.testing.assert_array_equal(memory[mask, :, 0], np.broadcast_to(np.arange(16), (len(expected), 16)))
+        self.assertNotIn(0, bank.frames)
+        self.assertEqual(int(mask.sum()), 191)
+        bank.reset()
+        self.assertFalse(bank.frames)
+        _, mask = bank.add_and_read(0, np.zeros((16, 1152)))
+        self.assertEqual(int(mask.sum()), 1)
+
+    def test_rmbench_fifo25_bootstrap_and_eviction(self):
+        bank = FIFOHistory(capacity=25, stride=30)
+        for t in range(0, 1201, 30):
+            memory, mask = bank.add_and_read(t, np.full((16, 1152), t, np.float16))
+            expected = [0] if t == 0 else list(range(0, t, 30))[-25:]
+            np.testing.assert_array_equal(memory[mask, 0, 0], expected)
+            self.assertEqual(memory.shape, (25, 16, 1152))
+            if t == 750:
+                self.assertEqual(memory[mask, 0, 0].tolist(), list(range(0, 750, 30)))
+            if t == 780:
+                self.assertEqual(memory[mask, 0, 0].tolist(), list(range(30, 780, 30)))
+                self.assertNotIn(0, bank.frames)
+        self.assertFalse(hasattr(bank, "frame0"))
+
+    def test_invalid_first_frame_does_not_mutate_bank(self):
+        bank = FIFOHistory(capacity=25, stride=30)
+        with self.assertRaisesRegex(ValueError, "frame zero"):
+            bank.add_and_read(30, np.zeros((16, 1152)))
+        self.assertEqual(bank.last_frame, -1)
+        self.assertFalse(bank.frames)
+
+    def test_missing_frames_and_duplicate_indices_rejected(self):
+        bank = FIFOHistory()
+        bank.add_and_read(0, np.zeros((16, 1152)))
+        with self.assertRaises(ValueError):
+            bank.add_and_read(0, np.zeros((16, 1152)))
+        with self.assertRaisesRegex(ValueError, "Missing historical"):
+            bank.add_and_read(25, np.zeros((16, 1152)))
+
+
+if __name__ == "__main__":
+    unittest.main()

code/pi05/scripts/test_robomem_memory_loader.py

@@ -0,0 +1,82 @@
+"""Verify batch equality, shared storage, epoch order, and producer failures."""
+import unittest
+import jax
+import numpy as np
+import torch
+from openpi.training.robomem_memory_loader import DevicePrefetch, make_host_loader
+
+
+class Samples:
+    def __len__(self):
+        return 12
+
+    def __getitem__(self, index):
+        return {"history": np.full((3, 4, 8), index, np.float16),
+                "mask": np.array([True, False, True]),
+                "data": {"ids": np.array([index, index + 1], np.int32)},
+                "progress": np.float32(index / 12)}
+
+
+class BrokenSamples(Samples):
+    def __getitem__(self, index):
+        raise ValueError("intentional producer failure")
+
+
+def numpy_collate(items):
+    return jax.tree.map(lambda *xs: np.stack([np.asarray(x) for x in xs]), *items)
+
+
+class LoaderTests(unittest.TestCase):
+    def setUp(self):
+        self.shard = jax.sharding.NamedSharding(
+            jax.sharding.Mesh(np.array(jax.devices()), ("batch",)), jax.sharding.PartitionSpec("batch")
+        )
+
+    def test_same_batches_two_epochs(self):
+        expected = torch.utils.data.DataLoader(
+            Samples(), batch_size=4, num_workers=2, shuffle=True,
+            generator=torch.Generator().manual_seed(7), drop_last=True,
+            persistent_workers=True, multiprocessing_context="spawn", collate_fn=numpy_collate,
+        )
+        source = make_host_loader(Samples(), 4, workers=2, seed=7)
+        prefetch = DevicePrefetch(source, self.shard, depth=2)
+        try:
+            for _ in range(2):
+                for host in expected:
+                    actual = next(prefetch)
+                    for a, b in zip(jax.tree.leaves(actual), jax.tree.leaves(host), strict=True):
+                        np.testing.assert_array_equal(np.asarray(a), np.asarray(b))
+                        self.assertEqual(a.dtype, np.asarray(b).dtype)
+        finally:
+            prefetch.close()
+            expected._iterator._shutdown_workers()
+
+    def test_worker_output_is_shared(self):
+        source = make_host_loader(Samples(), 4, workers=1, seed=7)
+        it = iter(source)
+        try:
+            for batch in it:
+                for value in jax.tree.leaves(batch):
+                    self.assertIsInstance(value, torch.Tensor)
+                    self.assertTrue(value.is_shared())
+        finally:
+            it._shutdown_workers()
+
+    def test_exception_reaches_trainer(self):
+        prefetch = DevicePrefetch(make_host_loader(BrokenSamples(), 4, workers=0, seed=7), self.shard)
+        try:
+            with self.assertRaisesRegex(ValueError, "intentional producer failure"):
+                next(prefetch)
+        finally:
+            prefetch.close()
+
+    def test_finite_iterator(self):
+        prefetch = DevicePrefetch(make_host_loader(Samples(), 4, workers=0, seed=7), self.shard, repeat=False)
+        try:
+            self.assertEqual(len(list(prefetch)), 3)
+        finally:
+            prefetch.close()
+
+
+if __name__ == "__main__":
+    unittest.main()

code/pi05/scripts/train_robomem_memory.py

@@ -0,0 +1,206 @@
+"""pi0.5: 5k adapter-only, then 50k joint updates with frozen SigLIP."""
+
+import argparse
+import dataclasses
+import json
+import logging
+import hashlib
+import shutil
+from pathlib import Path
+
+import jax
+from taskanchor.paths import ASSETS_DIR, CHECKPOINT_ROOT
+
+from openpi.training import checkpoints, config, sharding, weight_loaders
+from openpi.training.robomem_memory_data import (
+    BASE_PARAMS,
+    CACHE_ROOT,
+    DATA_ROOT,
+    EpisodeIndex,
+    MemoryDataset,
+    validate_cache,
+)
+
+
+from taskanchor.pi05.training import (
+    label as label,
+    trainable as trainable,
+    adapter_only as adapter_only,
+    make_optimizer as make_optimizer,
+    init_state,
+    train_step as train_step,
+    CheckpointData,
+    run_training_loop,
+)
+
+
+def recipe(args):
+    original = config.get_config("pi05_robomemarena_26tasks_subtask_only_full_sft_eval")
+    return dataclasses.replace(
+        original,
+        name="pi05_robomemarena_memory_v5_main",
+        exp_name=args.run_name,
+        model=dataclasses.replace(original.model, action_horizon=args.action_horizon),
+        data=dataclasses.replace(original.data, assets=dataclasses.replace(original.data.assets, assets_dir=args.assets_dir)),
+        batch_size=args.batch_size,
+        num_workers=args.workers,
+        num_train_steps=args.steps,
+        fsdp_devices=args.fsdp_devices,
+        ema_decay=None,
+        wandb_enabled=False,
+        save_interval=args.save_interval,
+        keep_period=5000,
+        checkpoint_base_dir=args.checkpoint_root,
+        weight_loader=weight_loaders.CheckpointWeightLoader(args.base_params),
+    )
+
+
+def main(args):
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s", force=True)
+    if not 0 < args.freeze_steps < args.steps:
+        raise ValueError("Require 0 < freeze_steps < steps")
+    if not 0 < args.execution_horizon <= args.action_horizon or 10 % args.execution_horizon:
+        raise ValueError("Execution horizon must divide history stride 10 and not exceed action horizon")
+    if any(d.platform != "gpu" for d in jax.devices()):
+        raise RuntimeError("Training requires GPUs; refusing a silent CPU fallback")
+    if args.batch_size % jax.device_count():
+        raise ValueError("Batch size must be divisible by GPU count")
+    index = EpisodeIndex(args.dataset)
+    logging.info("Verifying every cache file before training")
+    manifest = validate_cache(index, args.cache, args.base_params)
+    cfg = recipe(args)
+    dc = cfg.data.create(cfg.assets_dirs, cfg.model)
+    dataset = MemoryDataset(index, dc, cfg.model, args.cache)
+    mesh = sharding.make_mesh(args.fsdp_devices)
+    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
+    replicated = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+    manager, resume = checkpoints.initialize_checkpoint_dir(
+        cfg.checkpoint_dir, keep_period=5000, overwrite=False, resume=args.resume
+    )
+    metadata = dict(
+        metadata_filename="robomem_memory.json",
+        **vars(args),
+        cache_manifest=manifest,
+        model=dataclasses.asdict(cfg.model),
+        architecture="primary_camera_v5_spatial_history_offset_512_1152",
+        memory_layout="frame0_and_history",
+        frame0="separate",
+        history_stride=10,
+        history_capacity=index.history_frames,
+        max_task_frames=index.max_task_frames,
+        history_tokens_per_frame=16,
+        freeze_vision=True,
+        progress_target="round(subtask_id/num_subtasks,2)",
+        loss="flow+0.01*(100*MSE+1-Pearson)",
+    )
+    metadata_path = Path(cfg.checkpoint_dir) / "recipe.json"
+    if resume:
+        old = json.loads(metadata_path.read_text())
+        # Original50k recipes predate these explicit fields. Architecture identifies them.
+        from taskanchor.data.history import validate_robomem_recipe
+        validate_robomem_recipe(old)
+        old.setdefault("memory_layout", "frame0_and_history")
+        old.setdefault("frame0", "separate")
+        for key in (
+            "base_params",
+            "dataset",
+            "cache_manifest",
+            "freeze_steps",
+            "steps",
+            "batch_size",
+            "model",
+            "architecture",
+            "memory_layout",
+            "frame0",
+            "history_capacity",
+            "history_stride",
+            "history_tokens_per_frame",
+        ):
+            if old.get(key) != metadata[key]:
+                raise ValueError(f"Resume recipe changed: {key}")
+    else:
+        root = Path(__file__).resolve().parents[1]
+        snapshot = Path(cfg.checkpoint_dir) / "source_snapshot"
+        shutil.copytree(
+            root / "src", snapshot / "src", ignore=shutil.ignore_patterns("__pycache__", "*.pyc"), dirs_exist_ok=True
+        )
+        shutil.copytree(
+            root / "scripts",
+            snapshot / "scripts",
+            ignore=shutil.ignore_patterns("__pycache__", "*.pyc"),
+            dirs_exist_ok=True,
+        )
+        shutil.copytree(
+            root.parent / "taskanchor", snapshot / "taskanchor",
+            ignore=shutil.ignore_patterns("__pycache__", "*.pyc"), dirs_exist_ok=True,
+        )
+        metadata["source_sha256"] = {
+            str(path.relative_to(snapshot)): hashlib.sha256(path.read_bytes()).hexdigest()
+            for path in sorted(snapshot.rglob("*.py"))
+        }
+        metadata_path.write_text(json.dumps(metadata, indent=2))
+    logging.info(
+        "Spatial memory: max_task_frames=%d history_frames=%d tokens_per_frame=16 memory_tokens=%d",
+        index.max_task_frames,
+        index.history_frames,
+        (index.history_frames + 1) * 16,
+    )
+    logging.info("Initializing sharded pi0.5 + adapter")
+    state, state_sharding = init_state(cfg, args, mesh, resume)
+    checkpoint_data = CheckpointData(dc)
+    if resume:
+        state = checkpoints.restore_state(manager, state, checkpoint_data, step=args.resume_step)
+    step0 = int(state.step)
+    if args.resume_step is not None and (not resume or step0 != args.resume_step):
+        raise ValueError(f"Expected resume step {args.resume_step}, restored {step0}")
+    logging.info(
+        "Model initialized; start_step=%d total_steps=%d freeze_steps=%d", step0, args.steps, args.freeze_steps
+    )
+    run_training_loop(
+        args,
+        cfg,
+        dataset,
+        data_sharding,
+        replicated,
+        state_sharding,
+        state,
+        step0,
+        mesh,
+        manager,
+        checkpoint_data,
+        metadata,
+    )
+
+
+def build_argparser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", default=str(DATA_ROOT))
+    parser.add_argument("--base-params", default=str(BASE_PARAMS))
+    parser.add_argument("--cache", default=str(CACHE_ROOT))
+    parser.add_argument("--assets-dir", default=str(ASSETS_DIR), help="Directory containing pi05_libero/norm_stats.json")
+    parser.add_argument("--run-name", default="main_v5_spatial4x4_h191_a20_exec10_bs256_5k50k")
+    parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT))
+    parser.add_argument("--steps", type=int, default=55000)
+    parser.add_argument("--freeze-steps", type=int, default=5000)
+    parser.add_argument("--batch-size", type=int, default=256)
+    parser.add_argument("--action-horizon", type=int, default=20)
+    parser.add_argument(
+        "--execution-horizon",
+        type=int,
+        default=10,
+        help="Execution length for future deployment; training supervises the full action horizon",
+    )
+    parser.add_argument("--fsdp-devices", type=int, default=8)
+    parser.add_argument("--workers", type=int, default=8)
+    parser.add_argument("--save-interval", type=int, default=1000)
+    parser.add_argument("--log-interval", type=int, default=20)
+    parser.add_argument("--resume", action="store_true")
+    parser.add_argument("--resume-step", type=int, default=None)
+    parser.add_argument("--shared-memory-loader", action="store_true")
+    parser.add_argument("--prefetch-factor", type=int, default=2)
+    parser.add_argument("--device-prefetch", type=int, default=2)
+    return parser
+
+
+if __name__ == "__main__":
+    main(build_argparser().parse_args())