Update references to GCP assets

README.md

@@ -59,22 +59,22 @@ We provide multiple base VLA model checkpoints. These checkpoints have been pre-
 
 | Model        | Use Case    | Description                                                                                                 | Checkpoint Path                                |
 | ------------ | ----------- | ----------------------------------------------------------------------------------------------------------- | ---------------------------------------------- |
-| $\pi_0$      | Fine-Tuning | Base diffusion [π₀ model](https://www.physicalintelligence.company/blog/pi0) for fine-tuning                | `s3://openpi-assets/checkpoints/pi0_base`      |
-| $\pi_0$-FAST | Fine-Tuning | Base autoregressive [π₀-FAST model](https://www.physicalintelligence.company/research/fast) for fine-tuning | `s3://openpi-assets/checkpoints/pi0_fast_base` |
+| $\pi_0$      | Fine-Tuning | Base diffusion [π₀ 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` |
 
 ### 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 | `s3://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                                | `s3://openpi-assets/checkpoints/pi0_droid`            |
-| $\pi_0$-ALOHA-towel      | Inference   | $\pi_0$ model fine-tuned on internal ALOHA data, can fold diverse towels 0-shot on [ALOHA](https://tonyzhaozh.github.io/aloha/) robot platforms                                                          | `s3://openpi-assets/checkpoints/pi0_aloha_towel`      |
-| $\pi_0$-ALOHA-tupperware | Inference   | $\pi_0$ model fine-tuned on internal ALOHA data, can unpack food from a tupperware container                                                                                                             | `s3://openpi-assets/checkpoints/pi0_aloha_tupperware` |
-| $\pi_0$-ALOHA-pen-uncap  | Inference   | $\pi_0$ model fine-tuned on [public ALOHA data](https://dit-policy.github.io/), can uncap a pen                                                                                                          | `s3://openpi-assets/checkpoints/pi0_aloha_pen_uncap`  |
+| $\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 data, can fold diverse towels 0-shot on [ALOHA](https://tonyzhaozh.github.io/aloha/) robot platforms                                                          | `gs://openpi-assets/checkpoints/pi0_aloha_towel`      |
+| $\pi_0$-ALOHA-tupperware | Inference   | $\pi_0$ model fine-tuned on internal 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 data](https://dit-policy.github.io/), can uncap a pen                                                                                                          | `gs://openpi-assets/checkpoints/pi0_aloha_pen_uncap`  |
 
 
-By default, checkpoints are automatically downloaded from `s3://openpi-assets` and are cached in `~/.cache/openpi` when needed. You can overwrite the download path by setting the `OPENPI_DATA_HOME` environment variable.
+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.
 
 
 
@@ -88,7 +88,7 @@ from openpi.policies import policy_config
 from openpi.shared import download
 
 config = config.get_config("pi0_fast_droid")
-checkpoint_dir = download.maybe_download("s3://openpi-assets/checkpoints/pi0_fast_droid")
+checkpoint_dir = download.maybe_download("gs://openpi-assets/checkpoints/pi0_fast_droid")
 
 # Create a trained policy.
 policy = policy_config.create_trained_policy(config, checkpoint_dir)

docs/norm_stats.md

@@ -14,7 +14,7 @@ TrainConfig(
     data=LeRobotAlohaDataConfig(
         ...
         assets=AssetsConfig(
-            assets_dir="s3://openpi-assets/checkpoints/pi0_base/assets",
+            assets_dir="gs://openpi-assets/checkpoints/pi0_base/assets",
             asset_id="trossen",
         ),
     ),
@@ -30,7 +30,7 @@ For an example of a full training config that reloads normalization statistics,
 
 ## Provided Pre-training Normalization Statistics
 
-Below is a list of all the pre-training normalization statistics we provide. We provide them for both, the `pi0_base` and `pi0_fast_base` models. For `pi0_base`, set the `assets_dir` to `s3://openpi-assets/checkpoints/pi0_base/assets` and for `pi0_fast_base`, set the `assets_dir` to `s3://openpi-assets/checkpoints/pi0_fast_base/assets`.
+Below is a list of all the pre-training normalization statistics we provide. We provide them for both, the `pi0_base` and `pi0_fast_base` models. For `pi0_base`, set the `assets_dir` to `gs://openpi-assets/checkpoints/pi0_base/assets` and for `pi0_fast_base`, set the `assets_dir` to `gs://openpi-assets/checkpoints/pi0_fast_base/assets`.
 | Robot | Description | Asset ID |
 |-------|-------------|----------|
 | ALOHA | 6-DoF dual arm robot with parallel grippers | trossen |
@@ -59,7 +59,7 @@ Out of the box, both the `pi0_base` and `pi0_fast_base` use the following action
 
 The proprioceptive state uses the same definitions as the action space, except for the base x-y position (the last two dimensions) for mobile robots, which we don't include in the proprioceptive state.
 
-For 7-DoF robots (e.g. Franka), we use the first 7 dimensions of the action space for the joint actions, and the 8th dimension for the gripper action. 
+For 7-DoF robots (e.g. Franka), we use the first 7 dimensions of the action space for the joint actions, and the 8th dimension for the gripper action.
 
 General info for Pi robots:
 - Joint angles are expressed in radians, with position zero corresponding to the zero position reported by each robot's interface library, except for ALOHA, where the standard ALOHA code uses a slightly different convention (see the [ALOHA example code](../examples/aloha_real/README.md) for details).

docs/remote_inference.md

@@ -14,7 +14,7 @@ uv run scripts/serve_policy.py --env=[DROID | ALOHA | LIBERO]
 The `env` argument specifies which $\pi_0$ checkpoint should be loaded. Under the hood, this script will execute a command like the following, which you can use to start a policy server, e.g. for checkpoints you trained yourself (here an example for the DROID environment):
 
 ```bash
-uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_fast_droid --policy.dir=s3://openpi-assets/checkpoints/pi0_fast_droid
+uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_fast_droid --policy.dir=gs://openpi-assets/checkpoints/pi0_fast_droid
 ```
 
 This will start a policy server that will serve the policy specified by the `config` and `dir` arguments. The policy will be served on the specified port (default: 8000).

examples/aloha_real/README.md

@@ -53,74 +53,74 @@ While we’ve found the policies to work in unseen conditions across multiple AL
 
 ---
 
-### **Toast Task**  
+### **Toast Task**
 
-This task involves the robot taking two pieces of toast out of a toaster and placing them on a plate.  
+This task involves the robot taking two pieces of toast out of a toaster and placing them on a plate.
 
-- **Checkpoint path**: `s3://openpi-assets/checkpoints/pi0_base`
+- **Checkpoint path**: `gs://openpi-assets/checkpoints/pi0_base`
 - **Prompt**: "take the toast out of the toaster"
-- **Objects needed**: Two pieces of toast, a plate, and a standard toaster.  
-- **Object Distribution**:  
-  - Works on both real toast and rubber fake toast  
-  - Compatible with standard 2-slice toasters  
-  - Works with plates of varying colors  
+- **Objects needed**: Two pieces of toast, a plate, and a standard toaster.
+- **Object Distribution**:
+  - Works on both real toast and rubber fake toast
+  - Compatible with standard 2-slice toasters
+  - Works with plates of varying colors
 
 ### **Scene Setup Guidelines**
 <img width="500" alt="Screenshot 2025-01-31 at 10 06 02 PM" src="https://github.com/user-attachments/assets/3d043d95-9d1c-4dda-9991-e63cae61e02e" />
 
-- The toaster should be positioned in the top-left quadrant of the workspace.  
-- Both pieces of toast should start inside the toaster, with at least 1 cm of bread sticking out from the top.  
-- The plate should be placed roughly in the lower-center of the workspace.  
-- Works with both natural and synthetic lighting, but avoid making the scene too dark (e.g., don't place the setup inside an enclosed space or under a curtain).  
+- The toaster should be positioned in the top-left quadrant of the workspace.
+- Both pieces of toast should start inside the toaster, with at least 1 cm of bread sticking out from the top.
+- The plate should be placed roughly in the lower-center of the workspace.
+- Works with both natural and synthetic lighting, but avoid making the scene too dark (e.g., don't place the setup inside an enclosed space or under a curtain).
 
 
-### **Towel Task**  
+### **Towel Task**
 
 This task involves folding a small towel (e.g., roughly the size of a hand towel) into eighths.
 
-- **Checkpoint path**: `s3://openpi-assets/checkpoints/pi0_aloha_towel`
-- **Prompt**: "fold the towel"  
-- **Object Distribution**:  
-  - Works on towels of varying solid colors 
-  - Performance is worse on heavily textured or striped towels 
+- **Checkpoint path**: `gs://openpi-assets/checkpoints/pi0_aloha_towel`
+- **Prompt**: "fold the towel"
+- **Object Distribution**:
+  - Works on towels of varying solid colors
+  - Performance is worse on heavily textured or striped towels
 
-### **Scene Setup Guidelines**  
+### **Scene Setup Guidelines**
 <img width="500" alt="Screenshot 2025-01-31 at 10 01 15 PM" src="https://github.com/user-attachments/assets/9410090c-467d-4a9c-ac76-96e5b4d00943" />
 
-- The towel should be flattened and roughly centered on the table.  
-- Choose a towel that does not blend in with the table surface.  
+- The towel should be flattened and roughly centered on the table.
+- Choose a towel that does not blend in with the table surface.
 
 
-### **Tupperware Task**  
+### **Tupperware Task**
 
-This task involves opening a tupperware filled with food and pouring the contents onto a plate.  
+This task involves opening a tupperware filled with food and pouring the contents onto a plate.
 
-- **Checkpoint path**: `s3://openpi-assets/checkpoints/pi0_aloha_tupperware`
+- **Checkpoint path**: `gs://openpi-assets/checkpoints/pi0_aloha_tupperware`
 - **Prompt**: "open the tupperware and put the food on the plate"
-- **Objects needed**: Tupperware, food (or food-like items), and a plate.  
-- **Object Distribution**:  
-  - Works on various types of fake food (e.g., fake chicken nuggets, fries, and fried chicken).  
-  - Compatible with tupperware of different lid colors and shapes, with best performance on square tupperware with a corner flap (see images below).  
-  - The policy has seen plates of varying solid colors.  
+- **Objects needed**: Tupperware, food (or food-like items), and a plate.
+- **Object Distribution**:
+  - Works on various types of fake food (e.g., fake chicken nuggets, fries, and fried chicken).
+  - Compatible with tupperware of different lid colors and shapes, with best performance on square tupperware with a corner flap (see images below).
+  - The policy has seen plates of varying solid colors.
 
-### **Scene Setup Guidelines** 
+### **Scene Setup Guidelines**
 <img width="500" alt="Screenshot 2025-01-31 at 10 02 27 PM" src="https://github.com/user-attachments/assets/60fc1de0-2d64-4076-b903-f427e5e9d1bf" />
 
-- Best performance observed when both the tupperware and plate are roughly centered in the workspace.  
-- Positioning:  
-  - Tupperware should be on the left.  
-  - Plate should be on the right or bottom.  
-  - The tupperware flap should point toward the plate.  
+- Best performance observed when both the tupperware and plate are roughly centered in the workspace.
+- Positioning:
+  - Tupperware should be on the left.
+  - Plate should be on the right or bottom.
+  - The tupperware flap should point toward the plate.
 
 ## Training on your own Aloha dataset
 
-1. Convert the dataset to the LeRobot dataset v2.0 format. 
-    
-    We provide a script [convert_aloha_data_to_lerobot.py](./convert_aloha_data_to_lerobot.py) that converts the dataset to the LeRobot dataset v2.0 format. As an example we have converted the `aloha_pen_uncap_diverse_raw` dataset from the [BiPlay repo](https://huggingface.co/datasets/oier-mees/BiPlay/tree/main/aloha_pen_uncap_diverse_raw) and uploaded it to the HuggingFace Hub as [physical-intelligence/aloha_pen_uncap_diverse](https://huggingface.co/datasets/physical-intelligence/aloha_pen_uncap_diverse). 
+1. Convert the dataset to the LeRobot dataset v2.0 format.
 
+    We provide a script [convert_aloha_data_to_lerobot.py](./convert_aloha_data_to_lerobot.py) that converts the dataset to the LeRobot dataset v2.0 format. As an example we have converted the `aloha_pen_uncap_diverse_raw` dataset from the [BiPlay repo](https://huggingface.co/datasets/oier-mees/BiPlay/tree/main/aloha_pen_uncap_diverse_raw) and uploaded it to the HuggingFace Hub as [physical-intelligence/aloha_pen_uncap_diverse](https://huggingface.co/datasets/physical-intelligence/aloha_pen_uncap_diverse).
 
-2. Define a training config that uses the custom dataset. 
+
+2. Define a training config that uses the custom dataset.
 
     We provide the [pi0_aloha_pen_uncap config](../../src/openpi/training/config.py) as an example. You should refer to the root [README](../../README.md) for how to run training with the new config.
- 
+
 IMPORTANT: Our base checkpoint includes normalization stats from various common robot configurations. When fine-tuning a base checkpoint with a custom dataset from one of these configurations, we recommend using the corresponding normalization stats provided in the base checkpoint. In the example, this is done by specifying the trossen asset_id and a path to the pretrained checkpoint’s asset directory within the AssetsConfig.

examples/droid/README.md

@@ -11,7 +11,7 @@ Since the DROID control laptop does not have a powerful GPU, we will start a rem
 2. Start the OpenPI server via the following command:
 
 ```bash
-uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_fast_droid --policy.dir=s3://openpi-assets/checkpoints/pi0_fast_droid
+uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_fast_droid --policy.dir=gs://openpi-assets/checkpoints/pi0_fast_droid
 ```
 
 You can also run the equivalent command below:

examples/inference.ipynb

@@ -1,137 +1,137 @@
 {
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import dataclasses\n",
-    "\n",
-    "import jax\n",
-    "\n",
-    "from openpi.models import model as _model\n",
-    "from openpi.policies import droid_policy\n",
-    "from openpi.policies import policy_config as _policy_config\n",
-    "from openpi.shared import download\n",
-    "from openpi.training import config as _config\n",
-    "from openpi.training import data_loader as _data_loader"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Policy inference\n",
-    "\n",
-    "The following example shows how to create a policy from a checkpoint and run inference on a dummy example."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "config = _config.get_config(\"pi0_fast_droid\")\n",
-    "checkpoint_dir = download.maybe_download(\"s3://openpi-assets/checkpoints/pi0_fast_droid\")\n",
-    "\n",
-    "# Create a trained policy.\n",
-    "policy = _policy_config.create_trained_policy(config, checkpoint_dir)\n",
-    "\n",
-    "# Run inference on a dummy example. This example corresponds to observations produced by the DROID runtime.\n",
-    "example = droid_policy.make_droid_example()\n",
-    "result = policy.infer(example)\n",
-    "\n",
-    "# Delete the policy to free up memory.\n",
-    "del policy\n",
-    "\n",
-    "print(\"Actions shape:\", result[\"actions\"].shape)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Working with a live model\n",
-    "\n",
-    "\n",
-    "The following example shows how to create a live model from a checkpoint and compute training loss. First, we are going to demonstrate how to do it with fake data.\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "config = _config.get_config(\"pi0_aloha_sim\")\n",
-    "\n",
-    "checkpoint_dir = download.maybe_download(\"s3://openpi-assets/checkpoints/pi0_aloha_sim\")\n",
-    "key = jax.random.key(0)\n",
-    "\n",
-    "# Create a model from the checkpoint.\n",
-    "model = config.model.load(_model.restore_params(checkpoint_dir / \"params\"))\n",
-    "\n",
-    "# We can create fake observations and actions to test the model.\n",
-    "obs, act = config.model.fake_obs(), config.model.fake_act()\n",
-    "\n",
-    "# Sample actions from the model.\n",
-    "loss = model.compute_loss(key, obs, act)\n",
-    "print(\"Loss shape:\", loss.shape)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Now, we are going to create a data loader and use a real batch of training data to compute the loss."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Reduce the batch size to reduce memory usage.\n",
-    "config = dataclasses.replace(config, batch_size=2)\n",
-    "\n",
-    "# Load a single batch of data. This is the same data that will be used during training.\n",
-    "# NOTE: In order to make this example self-contained, we are skipping the normalization step\n",
-    "# since it requires the normalization statistics to be generated using `compute_norm_stats`.\n",
-    "loader = _data_loader.create_data_loader(config, num_batches=1, skip_norm_stats=True)\n",
-    "obs, act = next(iter(loader))\n",
-    "\n",
-    "# Sample actions from the model.\n",
-    "loss = model.compute_loss(key, obs, act)\n",
-    "\n",
-    "# Delete the model to free up memory.\n",
-    "del model\n",
-    "\n",
-    "print(\"Loss shape:\", loss.shape)"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": ".venv",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.9"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
+    "cells": [
+        {
+            "cell_type": "code",
+            "execution_count": 1,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "import dataclasses\n",
+                "\n",
+                "import jax\n",
+                "\n",
+                "from openpi.models import model as _model\n",
+                "from openpi.policies import droid_policy\n",
+                "from openpi.policies import policy_config as _policy_config\n",
+                "from openpi.shared import download\n",
+                "from openpi.training import config as _config\n",
+                "from openpi.training import data_loader as _data_loader"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "# Policy inference\n",
+                "\n",
+                "The following example shows how to create a policy from a checkpoint and run inference on a dummy example."
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "config = _config.get_config(\"pi0_fast_droid\")\n",
+                "checkpoint_dir = download.maybe_download(\"gs://openpi-assets/checkpoints/pi0_fast_droid\")\n",
+                "\n",
+                "# Create a trained policy.\n",
+                "policy = _policy_config.create_trained_policy(config, checkpoint_dir)\n",
+                "\n",
+                "# Run inference on a dummy example. This example corresponds to observations produced by the DROID runtime.\n",
+                "example = droid_policy.make_droid_example()\n",
+                "result = policy.infer(example)\n",
+                "\n",
+                "# Delete the policy to free up memory.\n",
+                "del policy\n",
+                "\n",
+                "print(\"Actions shape:\", result[\"actions\"].shape)"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "# Working with a live model\n",
+                "\n",
+                "\n",
+                "The following example shows how to create a live model from a checkpoint and compute training loss. First, we are going to demonstrate how to do it with fake data.\n"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "config = _config.get_config(\"pi0_aloha_sim\")\n",
+                "\n",
+                "checkpoint_dir = download.maybe_download(\"gs://openpi-assets/checkpoints/pi0_aloha_sim\")\n",
+                "key = jax.random.key(0)\n",
+                "\n",
+                "# Create a model from the checkpoint.\n",
+                "model = config.model.load(_model.restore_params(checkpoint_dir / \"params\"))\n",
+                "\n",
+                "# We can create fake observations and actions to test the model.\n",
+                "obs, act = config.model.fake_obs(), config.model.fake_act()\n",
+                "\n",
+                "# Sample actions from the model.\n",
+                "loss = model.compute_loss(key, obs, act)\n",
+                "print(\"Loss shape:\", loss.shape)"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "Now, we are going to create a data loader and use a real batch of training data to compute the loss."
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# Reduce the batch size to reduce memory usage.\n",
+                "config = dataclasses.replace(config, batch_size=2)\n",
+                "\n",
+                "# Load a single batch of data. This is the same data that will be used during training.\n",
+                "# NOTE: In order to make this example self-contained, we are skipping the normalization step\n",
+                "# since it requires the normalization statistics to be generated using `compute_norm_stats`.\n",
+                "loader = _data_loader.create_data_loader(config, num_batches=1, skip_norm_stats=True)\n",
+                "obs, act = next(iter(loader))\n",
+                "\n",
+                "# Sample actions from the model.\n",
+                "loss = model.compute_loss(key, obs, act)\n",
+                "\n",
+                "# Delete the model to free up memory.\n",
+                "del model\n",
+                "\n",
+                "print(\"Loss shape:\", loss.shape)"
+            ]
+        }
+    ],
+    "metadata": {
+        "kernelspec": {
+            "display_name": ".venv",
+            "language": "python",
+            "name": "python3"
+        },
+        "language_info": {
+            "codemirror_mode": {
+                "name": "ipython",
+                "version": 3
+            },
+            "file_extension": ".py",
+            "mimetype": "text/x-python",
+            "name": "python",
+            "nbconvert_exporter": "python",
+            "pygments_lexer": "ipython3",
+            "version": "3.11.9"
+        }
+    },
+    "nbformat": 4,
+    "nbformat_minor": 2
 }

examples/ur5/README.md

@@ -129,7 +129,7 @@ TrainConfig(
         # Reloading normalization stats can help transfer pre-trained models to new environments.
         # See the [norm_stats.md](../docs/norm_stats.md) file for more details.
         assets=AssetsConfig(
-            assets_dir="s3://openpi-assets/checkpoints/pi0_base/assets",
+            assets_dir="gs://openpi-assets/checkpoints/pi0_base/assets",
             asset_id="ur5e",
         ),
         base_config=DataConfig(
@@ -139,7 +139,7 @@ TrainConfig(
         ),
     ),
     # Load the pi0 base model checkpoint.
-    weight_loader=weight_loaders.CheckpointWeightLoader("s3://openpi-assets/checkpoints/pi0_base/params"),
+    weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"),
     num_train_steps=30_000,
 )
 ```

scripts/serve_policy.py

@@ -59,19 +59,19 @@ class Args:
 DEFAULT_CHECKPOINT: dict[EnvMode, Checkpoint] = {
     EnvMode.ALOHA: Checkpoint(
         config="pi0_aloha",
-        dir="s3://openpi-assets/checkpoints/pi0_base",
+        dir="gs://openpi-assets/checkpoints/pi0_base",
     ),
     EnvMode.ALOHA_SIM: Checkpoint(
         config="pi0_aloha_sim",
-        dir="s3://openpi-assets/checkpoints/pi0_aloha_sim",
+        dir="gs://openpi-assets/checkpoints/pi0_aloha_sim",
     ),
     EnvMode.DROID: Checkpoint(
         config="pi0_fast_droid",
-        dir="s3://openpi-assets/checkpoints/pi0_fast_droid",
+        dir="gs://openpi-assets/checkpoints/pi0_fast_droid",
     ),
     EnvMode.LIBERO: Checkpoint(
         config="pi0_fast_libero",
-        dir="s3://openpi-assets/checkpoints/pi0_fast_libero",
+        dir="gs://openpi-assets/checkpoints/pi0_fast_libero",
     ),
 }
 

src/openpi/models/model_test.py

@@ -84,7 +84,7 @@ def test_model_restore():
     obs, act = config.fake_obs(batch_size), config.fake_act(batch_size)
 
     model = config.load(
-        _model.restore_params(download.maybe_download("s3://openpi-assets/checkpoints/pi0_base/params"))
+        _model.restore_params(download.maybe_download("gs://openpi-assets/checkpoints/pi0_base/params"))
     )
 
     loss = model.compute_loss(key, obs, act)

src/openpi/policies/policy_test.py

@@ -9,7 +9,7 @@ from openpi.training import config as _config
 @pytest.mark.manual
 def test_infer():
     config = _config.get_config("pi0_aloha_sim")
-    policy = _policy_config.create_trained_policy(config, "s3://openpi-assets/checkpoints/pi0_aloha_sim")
+    policy = _policy_config.create_trained_policy(config, "gs://openpi-assets/checkpoints/pi0_aloha_sim")
 
     example = aloha_policy.make_aloha_example()
     result = policy.infer(example)
@@ -20,7 +20,7 @@ def test_infer():
 @pytest.mark.manual
 def test_broker():
     config = _config.get_config("pi0_aloha_sim")
-    policy = _policy_config.create_trained_policy(config, "s3://openpi-assets/checkpoints/pi0_aloha_sim")
+    policy = _policy_config.create_trained_policy(config, "gs://openpi-assets/checkpoints/pi0_aloha_sim")
 
     broker = action_chunk_broker.ActionChunkBroker(
         policy,

src/openpi/shared/download.py

@@ -10,15 +10,10 @@ import stat
 import time
 import urllib.parse
 
-import boto3
-import boto3.s3.transfer as s3_transfer
-import botocore
 import filelock
 import fsspec
 import fsspec.generic
-import s3transfer.futures as s3_transfer_futures
 import tqdm_loggable.auto as tqdm
-from types_boto3_s3.service_resource import ObjectSummary
 
 # Environment variable to control cache directory path, ~/.cache/openpi will be used by default.
 _OPENPI_DATA_HOME = "OPENPI_DATA_HOME"
@@ -92,22 +87,7 @@ def maybe_download(url: str, *, force_download: bool = False, **kwargs) -> pathl
             # Download the data to a local cache.
             logger.info(f"Downloading {url} to {local_path}")
             scratch_path = local_path.with_suffix(".partial")
-
-            if _is_openpi_url(url):
-                # Download without credentials.
-                _download_boto3(
-                    url,
-                    scratch_path,
-                    boto_session=boto3.Session(
-                        region_name="us-west-1",
-                    ),
-                    botocore_config=botocore.config.Config(signature_version=botocore.UNSIGNED),
-                )
-            elif url.startswith("s3://"):
-                # Download with default boto3 credentials.
-                _download_boto3(url, scratch_path)
-            else:
-                _download_fsspec(url, scratch_path, **kwargs)
+            _download_fsspec(url, scratch_path, **kwargs)
 
             shutil.move(scratch_path, local_path)
             _ensure_permissions(local_path)
@@ -126,7 +106,8 @@ def _download_fsspec(url: str, local_path: pathlib.Path, **kwargs) -> None:
     """Download a file from a remote filesystem to the local cache, and return the local path."""
     fs, _ = fsspec.core.url_to_fs(url, **kwargs)
     info = fs.info(url)
-    if is_dir := (info["type"] == "directory"):  # noqa: SIM108
+    # Folders are represented by 0-byte objects with a trailing forward slash.
+    if is_dir := (info["type"] == "directory" or (info["size"] == 0 and info["name"].endswith("/"))):
         total_size = fs.du(url)
     else:
         total_size = info["size"]
@@ -140,114 +121,6 @@ def _download_fsspec(url: str, local_path: pathlib.Path, **kwargs) -> None:
         pbar.update(total_size - pbar.n)
 
 
-def _download_boto3(
-    url: str,
-    local_path: pathlib.Path,
-    *,
-    boto_session: boto3.Session | None = None,
-    botocore_config: botocore.config.Config | None = None,
-    workers: int = 16,
-) -> None:
-    """Download a file from the OpenPI S3 bucket using boto3. This is a more performant version of download but can
-    only handle s3 urls. In openpi repo, this is mainly used to access assets in S3 with higher throughput.
-
-    Input:
-        url: URL to openpi checkpoint path.
-        local_path: local path to the downloaded file.
-        boto_session: Optional boto3 session, will create by default if not provided.
-        botocore_config: Optional botocore config.
-        workers: number of workers for downloading.
-    """
-
-    def validate_and_parse_url(maybe_s3_url: str) -> tuple[str, str]:
-        parsed = urllib.parse.urlparse(maybe_s3_url)
-        if parsed.scheme != "s3":
-            raise ValueError(f"URL must be an S3 URL (s3://), got: {maybe_s3_url}")
-        bucket_name = parsed.netloc
-        prefix = parsed.path.strip("/")
-        return bucket_name, prefix
-
-    bucket_name, prefix = validate_and_parse_url(url)
-    session = boto_session or boto3.Session()
-
-    s3api = session.resource("s3", config=botocore_config)
-    bucket = s3api.Bucket(bucket_name)
-
-    # Check if prefix points to an object and if not, assume that it's a directory and add a trailing slash.
-    try:
-        bucket.Object(prefix).load()
-    except botocore.exceptions.ClientError:
-        # Make sure to append a "/" to prevent getting objects from a different directory that shares the same prefix.
-        # For example, if we are downloading from s3://bucket/foo, we don't want to also download from s3://bucket/foobar.
-        if not prefix.endswith("/"):
-            prefix = prefix + "/"
-
-    # Get all candidate objects, filter out directories.
-    objects = [x for x in bucket.objects.filter(Prefix=prefix) if not x.key.endswith("/")]
-    if not objects:
-        raise FileNotFoundError(f"No objects found at {url}")
-
-    total_size = sum(obj.size for obj in objects)
-
-    s3t = _get_s3_transfer_manager(session, workers, botocore_config=botocore_config)
-
-    def transfer(
-        s3obj: ObjectSummary, dest_path: pathlib.Path, progress_func
-    ) -> s3_transfer_futures.TransferFuture | None:
-        if dest_path.exists():
-            dest_stat = dest_path.stat()
-            if s3obj.size == dest_stat.st_size:
-                progress_func(s3obj.size)
-                return None
-        dest_path.parent.mkdir(parents=True, exist_ok=True)
-        return s3t.download(
-            bucket_name,
-            s3obj.key,
-            str(dest_path),
-            subscribers=[
-                s3_transfer.ProgressCallbackInvoker(progress_func),
-            ],
-        )
-
-    try:
-        with tqdm.tqdm(total=total_size, unit="iB", unit_scale=True, unit_divisor=1024) as pbar:
-            if os.getenv("IS_DOCKER", "false").lower() == "true":
-                # tqdm is bugged when using docker-compose. See https://github.com/tqdm/tqdm/issues/771
-                def update_progress(size: int) -> None:
-                    pbar.update(size)
-                    print(pbar)
-            else:
-
-                def update_progress(size: int) -> None:
-                    pbar.update(size)
-
-            futures = []
-            for obj in objects:
-                relative_path = pathlib.Path(obj.key).relative_to(prefix)
-                dest_path = local_path / relative_path
-                if future := transfer(obj, dest_path, update_progress):
-                    futures.append(future)
-            for future in futures:
-                future.result()
-    finally:
-        s3t.shutdown()
-
-
-def _get_s3_transfer_manager(
-    session: boto3.Session, workers: int, botocore_config: botocore.config.Config | None = None
-) -> s3_transfer.TransferManager:
-    # Add a few extra connections to prevent exceeding the pool size.
-    config = botocore.config.Config(max_pool_connections=workers + 2)
-    if botocore_config is not None:
-        config = config.merge(botocore_config)
-    s3client = session.client("s3", config=config)
-    transfer_config = s3_transfer.TransferConfig(
-        use_threads=True,
-        max_concurrency=workers,
-    )
-    return s3_transfer.create_transfer_manager(s3client, transfer_config)
-
-
 def _set_permission(path: pathlib.Path, target_permission: int):
     """chmod requires executable permission to be set, so we skip if the permission is already match with the target."""
     if path.stat().st_mode & target_permission == target_permission:
@@ -295,11 +168,6 @@ def _ensure_permissions(path: pathlib.Path) -> None:
             _set_folder_permission(dir_path)
 
 
-def _is_openpi_url(url: str) -> bool:
-    """Check if the url is an OpenPI S3 bucket url."""
-    return url.startswith("s3://openpi-assets/")
-
-
 def _get_mtime(year: int, month: int, day: int) -> float:
     """Get the mtime of a given date at midnight UTC."""
     date = datetime.datetime(year, month, day, tzinfo=datetime.UTC)

src/openpi/shared/download_test.py

@@ -24,8 +24,8 @@ def test_download_local(tmp_path: pathlib.Path):
         download.maybe_download("bogus")
 
 
-def test_download_s3_dir():
-    remote_path = "s3://openpi-assets/testdata/random"
+def test_download_gs_dir():
+    remote_path = "gs://openpi-assets/testdata/random"
 
     local_path = download.maybe_download(remote_path)
     assert local_path.exists()
@@ -34,8 +34,8 @@ def test_download_s3_dir():
     assert new_local_path == local_path
 
 
-def test_download_s3():
-    remote_path = "s3://openpi-assets/testdata/random/random_512kb.bin"
+def test_download_gs():
+    remote_path = "gs://openpi-assets/testdata/random/random_512kb.bin"
 
     local_path = download.maybe_download(remote_path)
     assert local_path.exists()

src/openpi/training/config.py

@@ -44,7 +44,7 @@ class AssetsConfig:
 
     ```
     AssetsConfig(
-        assets_dir="s3://openpi-assets/checkpoints/pi0_base/assets",
+        assets_dir="gs://openpi-assets/checkpoints/pi0_base/assets",
         asset_id="trossen",
     )
     ```
@@ -556,7 +556,7 @@ _CONFIGS = [
         ),
         # Here you define which pre-trained checkpoint you want to load to initialize the model.
         # This should match the model config you chose above -- i.e. in this case we use the pi0 base model.
-        weight_loader=weight_loaders.CheckpointWeightLoader("s3://openpi-assets/checkpoints/pi0_base/params"),
+        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"),
         # Below you can define other hyperparameters like the learning rate, number of training steps, etc.
         # Check the base TrainConfig class for a full list of available hyperparameters.
         num_train_steps=30_000,
@@ -569,7 +569,7 @@ _CONFIGS = [
             repo_id="physical-intelligence/libero",
             base_config=DataConfig(prompt_from_task=True),
         ),
-        weight_loader=weight_loaders.CheckpointWeightLoader("s3://openpi-assets/checkpoints/pi0_base/params"),
+        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"),
         num_train_steps=30_000,
         # The freeze filter defines which parameters should be frozen during training.
         # We have a convenience function in the model config that returns the default freeze filter
@@ -599,7 +599,7 @@ _CONFIGS = [
             base_config=DataConfig(prompt_from_task=True),
         ),
         # Note that we load the pi0-FAST base model checkpoint here.
-        weight_loader=weight_loaders.CheckpointWeightLoader("s3://openpi-assets/checkpoints/pi0_fast_base/params"),
+        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_fast_base/params"),
         num_train_steps=30_000,
     ),
     TrainConfig(
@@ -613,7 +613,7 @@ _CONFIGS = [
             repo_id="physical-intelligence/libero",
             base_config=DataConfig(prompt_from_task=True),
         ),
-        weight_loader=weight_loaders.CheckpointWeightLoader("s3://openpi-assets/checkpoints/pi0_fast_base/params"),
+        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_fast_base/params"),
         num_train_steps=30_000,
         # Again, make sure to match the model config above when extracting the freeze filter
         # that specifies which parameters should be frozen during LoRA finetuning.
@@ -634,7 +634,7 @@ _CONFIGS = [
         data=LeRobotAlohaDataConfig(
             repo_id="physical-intelligence/aloha_pen_uncap_diverse",
             assets=AssetsConfig(
-                assets_dir="s3://openpi-assets/checkpoints/pi0_base/assets",
+                assets_dir="gs://openpi-assets/checkpoints/pi0_base/assets",
                 asset_id="trossen",
             ),
             default_prompt="uncap the pen",
@@ -654,7 +654,7 @@ _CONFIGS = [
                 ]
             ),
         ),
-        weight_loader=weight_loaders.CheckpointWeightLoader("s3://openpi-assets/checkpoints/pi0_base/params"),
+        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"),
         num_train_steps=20_000,
     ),
     #
@@ -698,7 +698,7 @@ _CONFIGS = [
             default_prompt="Transfer cube",
             use_delta_joint_actions=False,
         ),
-        weight_loader=weight_loaders.CheckpointWeightLoader("s3://openpi-assets/checkpoints/pi0_base/params"),
+        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"),
         num_train_steps=20_000,
     ),
     #

src/openpi/training/weight_loaders.py

@@ -42,7 +42,7 @@ class CheckpointWeightLoader(WeightLoader):
       trained checkpoints:
         example: "./checkpoints/<config>/<exp>/<step>/params"
       released checkpoints:
-        example: "s3://openpi-assets/checkpoints/<model>/params"
+        example: "gs://openpi-assets/checkpoints/<model>/params"
     """
 
     params_path: str