Codex/pi05 g1 coke rgbd heldout

README.md

@@ -123,6 +123,9 @@ We provide detailed step-by-step examples for running inference of our pre-train
 
 ## Fine-Tuning Base Models on Your Own Data
 
+For NVIDIA GB10 / DGX Spark inference and full pi0.5 fine-tuning, including the
+G1 Fruit Ninja LeRobot contract, see [docs/spark_gb10_pi05.md](docs/spark_gb10_pi05.md).
+
 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

docs/g1_coke_pickup_pi05.md

@@ -0,0 +1,42 @@
+# G1 Coke pickup on pi0.5
+
+This configuration keeps the pretrained pi0.5 tensor width at 32 while exposing the exact Coke task contract:
+
+- `head_image`: rendered Isaac or real G1 RGB, HWC `uint8`
+- `state`: 24 absolute upper-body joint positions
+- `actions`: 10-step chunks of seven normalized right-hand pose/grip commands
+- `prompt`: `pick up the Coke can and hold it upright`
+
+The first six outputs are relative right-hand pose commands. The seventh is binary Dex3 grip intent. pi0.5 does not directly command G1 joints, locomotion, or balance.
+
+## Convert successful Isaac demonstrations
+
+```bash
+python3 scripts/spark/convert_g1_coke_sim_demos.py \
+  --raw-dir /data/g1-coke-pi05-raw \
+  --repo-id YOUR_ORG/g1_coke_pickup \
+  --push-to-hub
+```
+
+Conversion verifies the manifest format, every shard SHA-256, 50 Hz timestamps, RGB/state/action shapes and dtypes, finite values, normalized action range, and successful episode flags. It fails closed on unsuccessful demonstrations unless `--allow-failures` is explicitly supplied for debugging.
+
+## Fine-tune on Spark
+
+```bash
+export OPENPI_G1_COKE_DATASET_REPO_ID=YOUR_ORG/g1_coke_pickup
+scripts/spark/train_g1_coke.sh
+```
+
+The launcher computes dataset-specific quantile normalization statistics once, then trains `pi05_spark_g1_coke_pickup` in bfloat16 from `/openpi_assets/checkpoints/pi05_base_pytorch`.
+
+## Serve the trained policy
+
+```bash
+python3 scripts/serve_policy.py \
+  --port 8000 \
+  policy:checkpoint \
+  --policy.config pi05_spark_g1_coke_pickup \
+  --policy.dir /openpi_assets/training/pi05_spark_g1_coke_pickup/EXPERIMENT/STEP
+```
+
+Use the Coke repository's `scripts/run_pi05_real.py` as the client. It is shadow-only unless explicitly enabled and leaves IK, joint limits, collision avoidance, controller ownership, and the hardware watchdog outside the VLA model.

docs/spark_gb10_pi05.md

@@ -0,0 +1,89 @@
+# pi0.5 on NVIDIA GB10 / DGX Spark
+
+This branch provides a working ARM64 CUDA 13 path for pi0.5 inference and full
+PyTorch fine-tuning. It keeps NVIDIA PyTorch 26.01 and TorchVision from the GB10
+base image, uses CPU JAX only to read the official Orbax checkpoint, and converts
+that checkpoint to SafeTensors before any GPU work.
+
+The upstream CUDA-12 JAX and PyTorch 2.7 wheels are intentionally not installed:
+they do not carry the GB10 `sm_121` kernel/runtime contract. MuJoCo-only ALOHA
+dependencies are also pruned because upstream MuJoCo 2.3.7 has no compatible
+Python 3.12 ARM64 wheel and is not needed to train or serve the G1 policy.
+
+## Reproduce the verified runtime
+
+Run these commands on the Spark host from this repository:
+
+```bash
+./scripts/spark/build_image.sh
+./scripts/spark/prepare_pi05.sh
+./scripts/spark/run_smoke.sh
+```
+
+The scripts preserve all model and report data under `/var/lib/openpi-spark`.
+The smoke test is deliberately a full pi0.5 model, not the small `debug_pi05`
+network. It runs denoising on the GPU, executes a backward and AdamW update, saves
+the new model and optimizer state, then leaves evidence at:
+
+- `reports/pi05_inference.json`
+- `reports/pi05_train_smoke.log`
+- `training/pi05_spark_smoke/gb10_full_model_smoke/1`
+
+## G1 Fruit Ninja fine-tuning contract
+
+pi0.5 does not learn from the PPO checkpoint or the AMP motion file directly. It
+needs time-aligned behavior demonstrations in LeRobot format. Each frame must
+contain these features:
+
+| LeRobot feature | Type and shape | Meaning |
+| --- | --- | --- |
+| `observation.images.head` | RGB image, `3xHxW` or `HxWx3` | G1 head RealSense color frame |
+| `observation.state` | `float32`, at most 32 values | fixed-order G1 proprioceptive state; 29 joint positions are the recommended starting contract |
+| `action` | `float32[21]` | existing Fruit Ninja task command: 3 walking, 1 chop-phase speed, 17 upper-body residuals |
+| task | string | language instruction such as `slice the fruit` |
+
+Frames, state, and the action actually applied to the controller must share the
+same timestamp. At 50 Hz, the configured 10-action horizon represents 0.2 s. Bad
+or aborted demonstrations should be excluded rather than labeled as successes.
+
+The G1 transform uses the head image as the base camera and supplies two masked
+zero wrist-camera slots. OpenPI pads the 29-D state and 21-D action to pi0.5's
+32-D internal width; policy output is cropped back to exactly 21 task actions.
+It does not bypass the existing frozen Unitree stabilizer, action scaling, safety
+supervisor, or operator arm/start gate.
+
+Once the demonstrations are uploaded as a LeRobot dataset, start a persistent
+full-model run with:
+
+```bash
+export OPENPI_G1_DATASET_REPO_ID=owner/dataset
+export OPENPI_G1_EXPERIMENT=fruit_ninja_pi05_v1
+export OPENPI_G1_NUM_TRAIN_STEPS=20000
+./scripts/spark/train_g1.sh
+```
+
+Simulation rollouts produced by `g1-fruit-ninja-mjwarp` can be validated and
+converted into this exact LeRobot contract before upload:
+
+```bash
+python scripts/spark/convert_g1_sim_demos.py \
+  --raw-dir /openpi_assets/demonstrations/g1-fruit-ninja \
+  --repo-id owner/g1-fruit-ninja-sim
+```
+
+The converter checks every episode's SHA-256, RGB, 29-joint state, 21-action
+command, frame count, contiguous frame index, finite values, and exact 50 Hz
+timestamps. It refuses failed episodes by default and creates private Hugging
+Face datasets when `--push-to-hub` is explicitly supplied.
+
+The launcher computes and persists normalization statistics before training. A
+new experiment name is required for each run so an older checkpoint is not
+silently replaced.
+
+## Hardware boundary
+
+A successful smoke test proves model execution and gradient updates on Spark
+48fd. It does not prove the policy can safely control the physical G1. Before any
+robot trial, validate the 21-D joint/order contract offline, replay output in
+simulation, enforce finite/range/rate checks, and require a supervised operator
+start with a zero-command fallback.

scripts/docker/spark_gb10.Dockerfile

@@ -0,0 +1,62 @@
+# syntax=docker/dockerfile:1.7
+
+# OpenPI pi0.5 runtime for NVIDIA GB10 / DGX Spark.
+#
+# The upstream Dockerfile installs PyTorch 2.7 and CUDA-12 JAX wheels. Those
+# wheels do not contain GB10 (sm_121) kernels. This image instead keeps the
+# NVIDIA 26.01 PyTorch + CUDA 13.1 stack that is validated on the target Spark,
+# while retaining CPU-only JAX for Orbax checkpoint conversion.
+ARG BASE_IMAGE=codex/mjlab-bench:20260809-egl
+FROM ${BASE_IMAGE}
+
+LABEL org.opencontainers.image.title="OpenPI pi0.5 for NVIDIA GB10"
+LABEL org.opencontainers.image.description="CUDA 13.1 PyTorch runtime with CPU JAX checkpoint conversion"
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV JAX_PLATFORMS=cpu
+ENV OPENPI_DATA_HOME=/openpi_assets
+ENV PYTHONUNBUFFERED=1
+ENV PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+ENV VIRTUAL_ENV=/opt/openpi-venv
+ENV PATH=/opt/openpi-venv/bin:${PATH}
+
+WORKDIR /opt/openpi
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends git-lfs \
+    && rm -rf /var/lib/apt/lists/* \
+    && python3 -m venv --system-site-packages "${VIRTUAL_ENV}"
+
+# Export the locked dependency graph, but prune packages that would replace the
+# NVIDIA-provided PyTorch/TorchVision pair or pull a parallel CUDA-12 runtime.
+COPY pyproject.toml uv.lock README.md LICENSE ./
+COPY packages/openpi-client/pyproject.toml packages/openpi-client/pyproject.toml
+COPY packages/openpi-client/src packages/openpi-client/src
+RUN uv export --quiet \
+        --frozen \
+        --no-dev \
+        --no-emit-project \
+        --no-emit-workspace \
+        --prune torch \
+        --prune torchvision \
+        --prune jax-cuda12-plugin \
+        --prune jax-cuda12-pjrt \
+        --prune gym-aloha \
+        --prune dm-control \
+        --prune labmaze \
+        --prune mujoco \
+        --output-file /tmp/spark-requirements.txt \
+    && ! grep -Eq '^(torch|torchvision|jax-cuda12|nvidia-.*-cu12)([= @;]|$)' /tmp/spark-requirements.txt \
+    && uv pip install --python "${VIRTUAL_ENV}/bin/python" --no-deps --requirement /tmp/spark-requirements.txt \
+    && rm /tmp/spark-requirements.txt
+
+COPY src src
+COPY scripts scripts
+COPY examples examples
+COPY packages packages
+
+RUN uv pip install --python "${VIRTUAL_ENV}/bin/python" --no-deps --editable packages/openpi-client --editable . \
+    && python3 -c "import pathlib, shutil, transformers; target = pathlib.Path(transformers.__file__).parent; [shutil.copy2(path, target / path.name) for path in pathlib.Path('src/openpi/models_pytorch/transformers_replace').glob('*') if path.is_file()]; [shutil.copytree(path, target / path.name, dirs_exist_ok=True) for path in pathlib.Path('src/openpi/models_pytorch/transformers_replace').glob('*') if path.is_dir()]" \
+    && JAX_PLATFORMS=cpu python3 -c "import jax, torch, transformers; assert jax.default_backend() == 'cpu'; assert torch.__version__.startswith('2.10.0a0'); assert transformers.__version__ == '4.53.2'; print('OpenPI GB10 dependency contract verified')"
+
+CMD ["bash"]

scripts/spark/build_image.sh

@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+OPENPI_SPARK_SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OPENPI_SPARK_REPO_DIR="$(cd "${OPENPI_SPARK_SCRIPT_DIR}/../.." && pwd)"
+OPENPI_SPARK_IMAGE="${OPENPI_SPARK_IMAGE:-openpi-pi05-gb10:local}"
+OPENPI_SPARK_BASE_IMAGE="${OPENPI_SPARK_BASE_IMAGE:-codex/mjlab-bench:20260809-egl}"
+
+docker image inspect "${OPENPI_SPARK_BASE_IMAGE}" >/dev/null
+
+docker build \
+    --progress=plain \
+    --build-arg "BASE_IMAGE=${OPENPI_SPARK_BASE_IMAGE}" \
+    --file "${OPENPI_SPARK_REPO_DIR}/scripts/docker/spark_gb10.Dockerfile" \
+    --tag "${OPENPI_SPARK_IMAGE}" \
+    "${OPENPI_SPARK_REPO_DIR}"
+
+docker image inspect "${OPENPI_SPARK_IMAGE}" \
+    --format 'built={{.Created}} image_id={{.Id}} size={{.Size}}'