diff --git a/README.md b/README.md index d16e43d479..f41d800a9f 100644 --- a/README.md +++ b/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 diff --git a/docs/g1_coke_pickup_pi05.md b/docs/g1_coke_pickup_pi05.md new file mode 100644 index 0000000000..c9b0605c02 --- /dev/null +++ b/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. diff --git a/docs/spark_gb10_pi05.md b/docs/spark_gb10_pi05.md new file mode 100644 index 0000000000..7de757eac3 --- /dev/null +++ b/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. diff --git a/scripts/docker/spark_gb10.Dockerfile b/scripts/docker/spark_gb10.Dockerfile new file mode 100644 index 0000000000..579f5426c6 --- /dev/null +++ b/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"] diff --git a/scripts/spark/build_image.sh b/scripts/spark/build_image.sh new file mode 100755 index 0000000000..971329cf1d --- /dev/null +++ b/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}}' diff --git a/scripts/spark/convert_g1_coke_recorder_rgbd.py b/scripts/spark/convert_g1_coke_recorder_rgbd.py new file mode 100755 index 0000000000..a84247c9c8 --- /dev/null +++ b/scripts/spark/convert_g1_coke_recorder_rgbd.py @@ -0,0 +1,557 @@ +#!/usr/bin/env python3 +"""Convert recorder RGB-D kinesthetic Coke demonstrations into LeRobot splits. + +The source recordings contain exact applied upper-body Teach commands, but the +Dex3 values are measured observations. The historical ``arm14`` contract +exports waist, both arms, and the observed hand. The ``left-arm7`` contract +strictly removes the right arm from both observation and action: its 17D state +is waist, left arm, and observed physical-left hand; its 7D action is the exact +applied left-arm target. Neither contract relabels observed hand motion as a +policy action. +""" + +from __future__ import annotations + +import argparse +from bisect import bisect_left +from dataclasses import dataclass +import hashlib +import itertools +import json +import math +from pathlib import Path +import shutil +from typing import Any + +from lerobot.common.datasets.lerobot_dataset import HF_LEROBOT_HOME +from lerobot.common.datasets.lerobot_dataset import LeRobotDataset +import numpy as np +from PIL import Image + +BATCH_SCHEMA = "wendy.g1.mujoco-trial-batch.v1" +TRIAL_SCHEMA_VERSION = 4 +OUTPUT_FPS = 15 +BODY_DIM = 17 +HAND_DIM = 7 +IMAGE_HEIGHT = 240 +IMAGE_WIDTH = 320 +# One source episode contains a single roughly 0.26 s camera dropout. Offline +# resampling may repeat its nearest real RGB-D pair across that gap, while the +# original camera-to-robot synchronization gate remains 20 ms and live +# inference retains its independent stale-frame rejection. +MAX_CAMERA_GRID_SKEW_S = 0.135 +MAX_ROBOT_CAMERA_SYNC_SKEW_MS = 20.0 +MAX_ABS_WAIST_ROLL_PITCH_RAD = math.radians(5.0) +NEAR_DEPTH_M = 0.20 +FAR_DEPTH_M = 2.00 +TASK = "Grasp the Coke can, lift it, and present it in front of the robot." +DEFAULT_HOLDOUT_NAMES = ("RCokeGrabbing4", "RCokeGrabbing9", "RCokeGrabbing15") + +BODY_JOINT_NAMES = ( + "waist_yaw_joint", + "waist_roll_joint", + "waist_pitch_joint", + "left_shoulder_pitch_joint", + "left_shoulder_roll_joint", + "left_shoulder_yaw_joint", + "left_elbow_joint", + "left_wrist_roll_joint", + "left_wrist_pitch_joint", + "left_wrist_yaw_joint", + "right_shoulder_pitch_joint", + "right_shoulder_roll_joint", + "right_shoulder_yaw_joint", + "right_elbow_joint", + "right_wrist_roll_joint", + "right_wrist_pitch_joint", + "right_wrist_yaw_joint", +) +LEFT_HAND_JOINT_NAMES = ( + "left_hand_thumb_0_joint", + "left_hand_thumb_1_joint", + "left_hand_thumb_2_joint", + "left_hand_middle_0_joint", + "left_hand_middle_1_joint", + "left_hand_index_0_joint", + "left_hand_index_1_joint", +) + + +@dataclass(frozen=True) +class TaskContract: + name: str + output_schema: str + state_body_indices: tuple[int, ...] + action_body_indices: tuple[int, ...] + state_names: tuple[str, ...] + action_names: tuple[str, ...] + + @property + def state_dim(self) -> int: + return len(self.state_names) + + @property + def action_dim(self) -> int: + return len(self.action_names) + + +ARM14_CONTRACT = TaskContract( + name="arm14", + output_schema="wendy.g1.coke-rgbd-arm14-conversion.v1", + state_body_indices=tuple(range(BODY_DIM)), + action_body_indices=tuple(range(3, BODY_DIM)), + state_names=BODY_JOINT_NAMES + LEFT_HAND_JOINT_NAMES, + action_names=BODY_JOINT_NAMES[3:], +) +LEFT_ARM7_CONTRACT = TaskContract( + name="left-arm7", + output_schema="wendy.g1.coke-rgbd-left-arm7-conversion.v1", + state_body_indices=tuple(range(10)), + action_body_indices=tuple(range(3, 10)), + state_names=BODY_JOINT_NAMES[:10] + LEFT_HAND_JOINT_NAMES, + action_names=BODY_JOINT_NAMES[3:10], +) +CONTRACTS = {contract.name: contract for contract in (ARM14_CONTRACT, LEFT_ARM7_CONTRACT)} + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(16 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _repo_id(value: str) -> str: + parts = value.split("/") + allowed = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-") + if len(parts) != 2 or not all(parts) or any(set(part) - allowed for part in parts): + raise argparse.ArgumentTypeError( + "repo id must be owner/dataset using letters, numbers, dot, dash, or underscore" + ) + return value + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--raw-dir", type=Path, required=True) + parser.add_argument("--train-repo-id", type=_repo_id, required=True) + parser.add_argument("--eval-repo-id", type=_repo_id, required=True) + parser.add_argument( + "--contract", + choices=tuple(CONTRACTS), + default=ARM14_CONTRACT.name, + help="Explicit state/action embodiment contract; use left-arm7 for the physical-left task", + ) + parser.add_argument( + "--holdout-name", + action="append", + dest="holdout_names", + help="Complete source episode name reserved for evaluation; repeat for multiple episodes", + ) + parser.add_argument("--overwrite", action="store_true") + return parser.parse_args() + + +def _finite_vector(value: Any, length: int, label: str) -> np.ndarray: + vector = np.asarray(value, dtype=np.float32) + if vector.shape != (length,) or not np.isfinite(vector).all(): + raise ValueError(f"{label} must be one finite {length}-value vector") + return vector + + +def _load_batch(raw_dir: Path) -> dict[str, Any]: + manifest_path = raw_dir / "manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + if manifest.get("schema") != BATCH_SCHEMA: + raise ValueError(f"unsupported recorder batch schema {manifest.get('schema')!r}") + if manifest.get("read_only_gets_only") is not True or manifest.get("robot_commands_sent") != 0: + raise ValueError("source batch does not preserve its read-only pull provenance") + trials = manifest.get("trials") + if not isinstance(trials, list) or not trials: + raise ValueError("recorder batch contains no trials") + return manifest + + +def _load_trial(raw_dir: Path, entry: dict[str, Any]) -> dict[str, Any]: + path = raw_dir / str(entry["file"]) + if _sha256(path) != entry.get("sha256"): + raise ValueError(f"{path} does not match its recorded SHA-256") + trial = json.loads(path.read_text(encoding="utf-8")) + name = trial.get("name") or trial.get("id") + if trial.get("schema_version") != TRIAL_SCHEMA_VERSION: + raise ValueError(f"{name} is not recorder schema v{TRIAL_SCHEMA_VERSION}") + if trial.get("id") != entry.get("id") or trial.get("name") != entry.get("name"): + raise ValueError(f"{path} identity does not match the batch manifest") + if trial.get("sample_hz") != 25.0 or trial.get("joint_count") != BODY_DIM: + raise ValueError(f"{name} does not use the expected 25 Hz 17D body contract") + + hand = trial.get("hand") + if not isinstance(hand, dict): + raise ValueError(f"{name} is missing its hand contract") + if ( + hand.get("semantic_hand") != "physical_left" + or hand.get("unitree_socket") != "left" + or hand.get("joint_count") != HAND_DIM + ): + raise ValueError(f"{name} does not use the physical-left Dex3 on the left socket") + if hand.get("policy_action_eligible") is not False: + raise ValueError(f"{name} unexpectedly marks measured hand pose as a policy action") + + frames = trial.get("frames") + if not isinstance(frames, list) or len(frames) != trial.get("frame_count"): + raise ValueError(f"{name} has a mismatched robot frame count") + timestamps = [float(frame["t_s"]) for frame in frames] + if any(not math.isfinite(value) for value in timestamps) or any( + right <= left for left, right in itertools.pairwise(timestamps) + ): + raise ValueError(f"{name} robot timestamps are not finite and strictly increasing") + for frame_index, frame in enumerate(frames): + body_state = _finite_vector( + frame.get("measured_q_rad"), BODY_DIM, f"{name} frame {frame_index} body state" + ) + if float(np.max(np.abs(body_state[1:3]))) > MAX_ABS_WAIST_ROLL_PITCH_RAD: + raise ValueError(f"{name} frame {frame_index} is not an upright waist observation") + hand_frame = frame.get("hand") + if not isinstance(hand_frame, dict): + raise ValueError(f"{name} frame {frame_index} is missing measured hand state") + _finite_vector( + hand_frame.get("measured_q_rad"), + HAND_DIM, + f"{name} frame {frame_index} hand state", + ) + teach = frame.get("teach_command") + if not isinstance(teach, dict) or teach.get("topic") != "rt/arm_sdk": + raise ValueError(f"{name} frame {frame_index} lacks an applied rt/arm_sdk Teach command") + _finite_vector(teach.get("q_rad"), BODY_DIM, f"{name} frame {frame_index} Teach command") + + video = trial.get("video") + depth = trial.get("depth") + if not isinstance(video, dict) or not isinstance(depth, dict): + raise ValueError(f"{name} does not contain RGB-D metadata") + if video.get("encoding") != "jpeg_sequence" or depth.get("encoding") != "png_z16_sequence": + raise ValueError(f"{name} uses an unsupported RGB-D encoding") + if depth.get("aligned_to") != "color": + raise ValueError(f"{name} depth is not aligned to RGB") + if float(video.get("maximum_sync_skew_ms", math.inf)) > MAX_ROBOT_CAMERA_SYNC_SKEW_MS: + raise ValueError(f"{name} RGB-to-robot synchronization exceeds 20 ms") + if float(depth.get("maximum_sync_skew_ms", math.inf)) > MAX_ROBOT_CAMERA_SYNC_SKEW_MS: + raise ValueError(f"{name} depth-to-robot synchronization exceeds 20 ms") + rgb_frames = video.get("frames") + depth_frames = depth.get("frames") + if ( + not isinstance(rgb_frames, list) + or not isinstance(depth_frames, list) + or len(rgb_frames) != len(depth_frames) + or len(rgb_frames) != video.get("frame_count") + or len(depth_frames) != depth.get("frame_count") + ): + raise ValueError(f"{name} RGB and depth frame counts do not match") + for pair_index, (rgb_frame, depth_frame) in enumerate(zip(rgb_frames, depth_frames, strict=True)): + if ( + abs(float(rgb_frame["t_s"]) - float(depth_frame["t_s"])) > 1.0e-9 + or rgb_frame.get("source_frame_id") != depth_frame.get("source_frame_id") + or rgb_frame.get("nearest_robot_frame_index") + != depth_frame.get("nearest_robot_frame_index") + ): + raise ValueError(f"{name} RGB-D pair {pair_index} is not aligned") + if not (raw_dir / str(rgb_frame["image"])).is_file(): + raise FileNotFoundError(raw_dir / str(rgb_frame["image"])) + if not (raw_dir / str(depth_frame["image"])).is_file(): + raise FileNotFoundError(raw_dir / str(depth_frame["image"])) + return trial + + +def _interpolate_robot( + trial: dict[str, Any], source_time_s: float, contract: TaskContract = ARM14_CONTRACT +) -> tuple[np.ndarray, np.ndarray]: + frames = trial["frames"] + timestamps = [float(frame["t_s"]) for frame in frames] + right_index = bisect_left(timestamps, source_time_s) + if right_index == 0: + left_index = right_index = 0 + ratio = 0.0 + elif right_index == len(frames): + left_index = right_index = len(frames) - 1 + ratio = 0.0 + else: + left_index = right_index - 1 + interval = timestamps[right_index] - timestamps[left_index] + ratio = (source_time_s - timestamps[left_index]) / interval + + left = frames[left_index] + right = frames[right_index] + + def lerp(left_value: Any, right_value: Any, length: int, label: str) -> np.ndarray: + left_vector = _finite_vector(left_value, length, label) + right_vector = _finite_vector(right_value, length, label) + return left_vector + (right_vector - left_vector) * np.float32(ratio) + + body = lerp(left["measured_q_rad"], right["measured_q_rad"], BODY_DIM, "body state") + hand = lerp( + left["hand"]["measured_q_rad"], + right["hand"]["measured_q_rad"], + HAND_DIM, + "hand state", + ) + applied_teach = lerp( + left["teach_command"]["q_rad"], + right["teach_command"]["q_rad"], + BODY_DIM, + "Teach command", + ) + state = np.concatenate((body[list(contract.state_body_indices)], hand)).astype( + np.float32, copy=False + ) + action = applied_teach[list(contract.action_body_indices)].astype(np.float32, copy=False) + if state.shape != (contract.state_dim,) or action.shape != (contract.action_dim,): + raise RuntimeError("internal G1 state/action slicing error") + return state, action + + +def _camera_pairs(trial: dict[str, Any]) -> list[dict[str, Any]]: + return [ + { + "t_s": float(rgb["t_s"]), + "rgb": rgb["image"], + "depth": depth["image"], + } + for rgb, depth in zip(trial["video"]["frames"], trial["depth"]["frames"], strict=True) + ] + + +def _nearest_camera_pair(pairs: list[dict[str, Any]], source_time_s: float) -> tuple[dict[str, Any], float]: + timestamps = [pair["t_s"] for pair in pairs] + right_index = bisect_left(timestamps, source_time_s) + candidates = [] + if right_index < len(pairs): + candidates.append(pairs[right_index]) + if right_index > 0: + candidates.append(pairs[right_index - 1]) + pair = min(candidates, key=lambda item: abs(item["t_s"] - source_time_s)) + skew_s = abs(pair["t_s"] - source_time_s) + if skew_s > MAX_CAMERA_GRID_SKEW_S: + raise ValueError( + f"nearest RGB-D frame is {skew_s * 1000.0:.1f} ms from the {OUTPUT_FPS} Hz grid" + ) + return pair, skew_s + + +def _load_rgb(path: Path) -> np.ndarray: + with Image.open(path) as image: + rgb = image.convert("RGB").resize( + (IMAGE_WIDTH, IMAGE_HEIGHT), + resample=Image.Resampling.BILINEAR, + ) + result = np.asarray(rgb, dtype=np.uint8) + if result.shape != (IMAGE_HEIGHT, IMAGE_WIDTH, 3): + raise ValueError(f"invalid RGB image {path}: {result.shape}") + return result + + +def _load_depth(path: Path, depth_scale_m: float) -> tuple[np.ndarray, float]: + with Image.open(path) as image: + raw = np.asarray(image) + if raw.shape != (480, 640) or raw.dtype != np.uint16: + raise ValueError(f"invalid Z16 depth image {path}: {raw.shape} {raw.dtype}") + depth_m = raw.astype(np.float32) * np.float32(depth_scale_m) + valid = np.isfinite(depth_m) & (depth_m > 0.0) + valid_fraction = float(np.mean(valid)) + if valid_fraction < 0.10: + raise ValueError(f"depth image {path} has only {valid_fraction:.3f} valid pixels") + clipped = np.clip(depth_m, NEAR_DEPTH_M, FAR_DEPTH_M) + inverse = (FAR_DEPTH_M - clipped) / (FAR_DEPTH_M - NEAR_DEPTH_M) + encoded = np.where(valid, np.rint(inverse * 255.0), 0.0).astype(np.uint8) + resized = Image.fromarray(encoded, mode="L").resize( + (IMAGE_WIDTH, IMAGE_HEIGHT), + resample=Image.Resampling.NEAREST, + ) + depth_gray = np.asarray(resized, dtype=np.uint8) + depth_rgb = np.repeat(depth_gray[..., None], 3, axis=-1) + return depth_rgb, valid_fraction + + +def _samples( + raw_dir: Path, trial: dict[str, Any], contract: TaskContract = ARM14_CONTRACT +) -> tuple[list[dict[str, Any]], dict[str, float]]: + pairs = _camera_pairs(trial) + first_time_s = pairs[0]["t_s"] + last_time_s = min(pairs[-1]["t_s"], float(trial["frames"][-1]["t_s"])) + count = math.floor((last_time_s - first_time_s) * OUTPUT_FPS + 1.0e-9) + 1 + if count < OUTPUT_FPS: + raise ValueError(f"{trial['name']} is shorter than one second after RGB-D alignment") + depth_scale_m = float(trial["depth"]["calibration"]["depth_scale_m"]) + if not 0.0 < depth_scale_m < 0.1: + raise ValueError(f"{trial['name']} has an invalid depth scale {depth_scale_m}") + + samples = [] + maximum_camera_grid_skew_s = 0.0 + minimum_depth_valid_fraction = 1.0 + for frame_index in range(count): + source_time_s = first_time_s + frame_index / OUTPUT_FPS + pair, camera_grid_skew_s = _nearest_camera_pair(pairs, source_time_s) + maximum_camera_grid_skew_s = max(maximum_camera_grid_skew_s, camera_grid_skew_s) + state, action = _interpolate_robot(trial, source_time_s, contract) + rgb = _load_rgb(raw_dir / pair["rgb"]) + depth, valid_fraction = _load_depth(raw_dir / pair["depth"], depth_scale_m) + minimum_depth_valid_fraction = min(minimum_depth_valid_fraction, valid_fraction) + samples.append( + { + "observation.images.head": rgb, + "observation.images.depth": depth, + "observation.state": state, + "action": action, + "task": TASK, + } + ) + return samples, { + "maximum_camera_grid_skew_ms": maximum_camera_grid_skew_s * 1000.0, + "minimum_depth_valid_fraction": minimum_depth_valid_fraction, + } + + +def _create_dataset(repo_id: str, contract: TaskContract, *, overwrite: bool) -> LeRobotDataset: + output_path = HF_LEROBOT_HOME / repo_id + if output_path.exists(): + if not overwrite: + raise FileExistsError(f"{output_path} exists; pass --overwrite to replace it") + shutil.rmtree(output_path) + return LeRobotDataset.create( + repo_id=repo_id, + robot_type="unitree_g1_upper_body_physical_left_dex3", + fps=OUTPUT_FPS, + features={ + "observation.images.head": { + "dtype": "image", + "shape": (IMAGE_HEIGHT, IMAGE_WIDTH, 3), + "names": ["height", "width", "channel"], + }, + "observation.images.depth": { + "dtype": "image", + "shape": (IMAGE_HEIGHT, IMAGE_WIDTH, 3), + "names": ["height", "width", "channel"], + }, + "observation.state": { + "dtype": "float32", + "shape": (contract.state_dim,), + "names": ["state"], + }, + "action": { + "dtype": "float32", + "shape": (contract.action_dim,), + "names": ["action"], + }, + }, + use_videos=True, + image_writer_threads=8, + image_writer_processes=0, + ) + + +def _write_conversion_metadata(repo_id: str, metadata: dict[str, Any]) -> None: + path = HF_LEROBOT_HOME / repo_id / "meta" / "wendy-rgbd-conversion.json" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def main() -> None: + args = parse_args() + contract = CONTRACTS[args.contract] + raw_dir = args.raw_dir.expanduser().resolve() + if args.train_repo_id == args.eval_repo_id: + raise ValueError("train and held-out evaluation repo ids must differ") + if contract is LEFT_ARM7_CONTRACT and any( + "left_only" not in repo_id for repo_id in (args.train_repo_id, args.eval_repo_id) + ): + raise ValueError("left-arm7 datasets must use distinct repo ids containing 'left_only'") + manifest = _load_batch(raw_dir) + holdout_names = set(args.holdout_names or DEFAULT_HOLDOUT_NAMES) + available_names = {str(entry["name"]) for entry in manifest["trials"]} + missing_holdouts = holdout_names - available_names + if missing_holdouts: + raise ValueError(f"holdout episodes are absent: {sorted(missing_holdouts)}") + if len(holdout_names) >= len(available_names): + raise ValueError("holdout split leaves no training episodes") + + train_dataset = _create_dataset(args.train_repo_id, contract, overwrite=args.overwrite) + eval_dataset = _create_dataset(args.eval_repo_id, contract, overwrite=args.overwrite) + split_records: dict[str, list[dict[str, Any]]] = {"train": [], "eval": []} + for entry in manifest["trials"]: + trial = _load_trial(raw_dir, entry) + split = "eval" if trial["name"] in holdout_names else "train" + samples, quality = _samples(raw_dir, trial, contract) + dataset = eval_dataset if split == "eval" else train_dataset + for sample in samples: + dataset.add_frame(sample) + dataset.save_episode() + split_records[split].append( + { + "id": trial["id"], + "name": trial["name"], + "source_json": entry["file"], + "source_json_sha256": entry["sha256"], + "frames": len(samples), + **quality, + } + ) + + metadata = { + "schema": contract.output_schema, + "contract": contract.name, + "source_batch_manifest_sha256": _sha256(raw_dir / "manifest.json"), + "fps": OUTPUT_FPS, + "task": TASK, + "image": { + "rgb": "320x240 uint8 RGB", + "depth": { + "representation": "three-channel uint8 inverse depth", + "near_m": NEAR_DEPTH_M, + "far_m": FAR_DEPTH_M, + "invalid_value": 0, + }, + }, + "state": { + "width": contract.state_dim, + "names": list(contract.state_names), + "source": ( + "measured waist and left arm plus measured physical-left Dex3" + if contract is LEFT_ARM7_CONTRACT + else "17D measured upper body plus 7D measured physical-left Dex3" + ), + "right_arm_included": contract is ARM14_CONTRACT, + }, + "action": { + "width": contract.action_dim, + "names": list(contract.action_names), + "source": ( + "applied rt/arm_sdk Teach q_rad indices 3 through 9" + if contract is LEFT_ARM7_CONTRACT + else "applied rt/arm_sdk Teach q_rad indices 3 through 16" + ), + "waist_excluded": True, + "right_arm_included": contract is ARM14_CONTRACT, + "measured_hand_as_action": False, + }, + "splits": split_records, + } + _write_conversion_metadata(args.train_repo_id, metadata) + _write_conversion_metadata(args.eval_repo_id, metadata) + print( + json.dumps( + { + "status": "passed", + "contract": contract.name, + "train_repo_id": args.train_repo_id, + "eval_repo_id": args.eval_repo_id, + "train_episodes": len(split_records["train"]), + "eval_episodes": len(split_records["eval"]), + "train_frames": sum(record["frames"] for record in split_records["train"]), + "eval_frames": sum(record["frames"] for record in split_records["eval"]), + }, + sort_keys=True, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/spark/convert_g1_coke_recorder_rgbd_test.py b/scripts/spark/convert_g1_coke_recorder_rgbd_test.py new file mode 100644 index 0000000000..8e58d59226 --- /dev/null +++ b/scripts/spark/convert_g1_coke_recorder_rgbd_test.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path +import sys + +import numpy as np + +SCRIPT = Path(__file__).with_name("convert_g1_coke_recorder_rgbd.py") +SPEC = importlib.util.spec_from_file_location("convert_g1_coke_recorder_rgbd", SCRIPT) +assert SPEC is not None +assert SPEC.loader is not None +converter = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = converter +SPEC.loader.exec_module(converter) + + +def _trial() -> dict: + def frame(t_s: float, offset: float) -> dict: + return { + "t_s": t_s, + "measured_q_rad": (np.arange(17, dtype=np.float32) + offset).tolist(), + "hand": { + "measured_q_rad": (np.arange(7, dtype=np.float32) + 20.0 + offset).tolist() + }, + "teach_command": { + "q_rad": (np.arange(17, dtype=np.float32) + 200.0 + offset).tolist() + }, + } + + return {"frames": [frame(0.0, 0.0), frame(1.0, 100.0)]} + + +def test_left_arm7_contract_has_no_right_arm_names() -> None: + contract = converter.LEFT_ARM7_CONTRACT + + assert contract.state_dim == 17 + assert contract.action_dim == 7 + assert all(not name.startswith("right_") for name in contract.state_names) + assert all(name.startswith("left_") for name in contract.action_names) + + +def test_left_arm7_slices_only_waist_left_arm_and_left_hand() -> None: + state, action = converter._interpolate_robot( # noqa: SLF001 + _trial(), 0.25, converter.LEFT_ARM7_CONTRACT + ) + + np.testing.assert_allclose(state[:10], np.arange(10, dtype=np.float32) + 25.0) + np.testing.assert_allclose(state[10:], np.arange(7, dtype=np.float32) + 45.0) + np.testing.assert_allclose(action, np.arange(3, 10, dtype=np.float32) + 225.0) + assert state.shape == (17,) + assert action.shape == (7,) + + +def test_arm14_contract_remains_backward_compatible() -> None: + state, action = converter._interpolate_robot(_trial(), 0.25) # noqa: SLF001 + + np.testing.assert_allclose(state[:17], np.arange(17, dtype=np.float32) + 25.0) + np.testing.assert_allclose(state[17:], np.arange(7, dtype=np.float32) + 45.0) + np.testing.assert_allclose(action, np.arange(3, 17, dtype=np.float32) + 225.0) + assert state.shape == (24,) + assert action.shape == (14,) diff --git a/scripts/spark/convert_g1_coke_sim_demos.py b/scripts/spark/convert_g1_coke_sim_demos.py new file mode 100755 index 0000000000..ae3fc85afa --- /dev/null +++ b/scripts/spark/convert_g1_coke_sim_demos.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +"""Convert validated Isaac G1 Coke-pickup demonstrations to LeRobot.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +import shutil + +from lerobot.common.datasets.lerobot_dataset import HF_LEROBOT_HOME +from lerobot.common.datasets.lerobot_dataset import LeRobotDataset +import numpy as np + +EXPECTED_FORMAT = "g1-coke-pickup-pi05-raw-v2" +EXPECTED_FPS = 50 +EXPECTED_STATE_DIM = 24 +EXPECTED_ACTION_DIM = 21 +MAX_ABS_WAIST_ROLL_PITCH_RAD = np.deg2rad(5.0) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(16 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _repo_id(value: str) -> str: + parts = value.split("/") + allowed = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-") + if len(parts) != 2 or not all(parts) or any(set(part) - allowed for part in parts): + raise argparse.ArgumentTypeError( + "repo id must be owner/dataset using letters, numbers, dot, dash, or underscore" + ) + return value + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--raw-dir", type=Path, required=True) + parser.add_argument("--repo-id", type=_repo_id, required=True) + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--push-to-hub", action="store_true") + parser.add_argument("--allow-failures", action="store_true") + return parser.parse_args() + + +def _load_manifest(raw_dir: Path) -> dict[str, object]: + manifest = json.loads((raw_dir / "manifest.json").read_text(encoding="utf-8")) + if manifest.get("format") != EXPECTED_FORMAT: + raise ValueError(f"unsupported raw format {manifest.get('format')!r}") + if manifest.get("fps") != EXPECTED_FPS: + raise ValueError(f"expected {EXPECTED_FPS} Hz demonstrations") + if manifest.get("state", {}).get("shape") != [EXPECTED_STATE_DIM]: + raise ValueError("manifest does not contain the 24-value G1 upper-body state") + if manifest.get("action", {}).get("shape") != [EXPECTED_ACTION_DIM]: + raise ValueError("manifest does not contain the 21-joint Coke task action") + episodes = manifest.get("episodes") + if not isinstance(episodes, list) or not episodes: + raise ValueError("manifest contains no episodes") + return manifest + + +def _load_episode(raw_dir: Path, entry: dict[str, object], image_shape: tuple[int, int, int]): + path = raw_dir / str(entry["path"]) + if _sha256(path) != entry.get("sha256"): + raise ValueError(f"{path} does not match its recorded SHA-256") + with np.load(path, allow_pickle=False) as archive: + image = archive["observation_images_head"] + state = archive["observation_state"] + action = archive["action"] + timestamp = archive["timestamp"] + frame_index = archive["frame_index"] + frames = len(image) + if frames < 1 or entry.get("frames") != frames: + raise ValueError(f"invalid or mismatched frame count in {path}") + if image.shape != (frames, *image_shape) or image.dtype != np.uint8: + raise ValueError(f"invalid RGB array in {path}: {image.shape} {image.dtype}") + if state.shape != (frames, EXPECTED_STATE_DIM) or state.dtype != np.float32: + raise ValueError(f"invalid state array in {path}: {state.shape} {state.dtype}") + if action.shape != (frames, EXPECTED_ACTION_DIM) or action.dtype != np.float32: + raise ValueError(f"invalid action array in {path}: {action.shape} {action.dtype}") + if not np.isfinite(state).all() or not np.isfinite(action).all() or np.max(np.abs(action)) > 2.0 * np.pi: + raise ValueError(f"non-finite or out-of-range values in {path}") + max_abs_waist_roll_pitch = float(np.max(np.abs(state[:, 1:3]))) + if max_abs_waist_roll_pitch > MAX_ABS_WAIST_ROLL_PITCH_RAD: + raise ValueError( + f"non-upright waist in {path}: {max_abs_waist_roll_pitch:.6f} rad exceeds " + f"{MAX_ABS_WAIST_ROLL_PITCH_RAD:.6f} rad" + ) + if not np.array_equal(frame_index, np.arange(frames, dtype=np.int64)): + raise ValueError(f"non-contiguous frame indices in {path}") + if not np.allclose(timestamp, np.arange(frames, dtype=np.float64) / EXPECTED_FPS, atol=1.0e-6): + raise ValueError(f"timestamps are not synchronized at {EXPECTED_FPS} Hz in {path}") + return image, state, action + + +def main() -> None: + args = parse_args() + raw_dir = args.raw_dir.expanduser().resolve() + manifest = _load_manifest(raw_dir) + image_shape = tuple(manifest["image"]["shape"]) + if image_shape != (240, 320, 3): + raise ValueError(f"expected rendered 240x320 RGB, got {image_shape}") + + output_path = HF_LEROBOT_HOME / args.repo_id + if output_path.exists(): + if not args.overwrite: + raise FileExistsError(f"{output_path} exists; pass --overwrite to replace it") + shutil.rmtree(output_path) + dataset = LeRobotDataset.create( + repo_id=args.repo_id, + robot_type="unitree_g1", + fps=EXPECTED_FPS, + features={ + "observation.images.head": { + "dtype": "image", + "shape": image_shape, + "names": ["height", "width", "channel"], + }, + "observation.state": {"dtype": "float32", "shape": (EXPECTED_STATE_DIM,), "names": ["state"]}, + "action": {"dtype": "float32", "shape": (EXPECTED_ACTION_DIM,), "names": ["action"]}, + }, + use_videos=True, + image_writer_threads=8, + image_writer_processes=4, + ) + + converted = 0 + for entry in manifest["episodes"]: + if not entry.get("success", False) and not args.allow_failures: + raise ValueError(f"{entry.get('path')} is unsuccessful; refusing conversion without --allow-failures") + image, state, action = _load_episode(raw_dir, entry, image_shape) + task = entry.get("task") or manifest["task"] + for frame_index in range(len(image)): + dataset.add_frame( + { + "observation.images.head": image[frame_index], + "observation.state": state[frame_index], + "action": action[frame_index], + "task": task, + } + ) + dataset.save_episode() + converted += 1 + + if args.push_to_hub: + dataset.push_to_hub( + tags=["unitree-g1", "coke-pickup", "isaaclab", "simulation", "pi0.5"], + private=True, + push_videos=True, + license="apache-2.0", + ) + print(json.dumps({"status": "passed", "episodes": converted, "output": str(output_path)})) + + +if __name__ == "__main__": + main() diff --git a/scripts/spark/convert_g1_sim_demos.py b/scripts/spark/convert_g1_sim_demos.py new file mode 100755 index 0000000000..7df5830730 --- /dev/null +++ b/scripts/spark/convert_g1_sim_demos.py @@ -0,0 +1,175 @@ +#!/usr/bin/env python3 +"""Convert synchronized G1 simulation shards into a LeRobot dataset.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +import shutil + +from lerobot.common.datasets.lerobot_dataset import HF_LEROBOT_HOME +from lerobot.common.datasets.lerobot_dataset import LeRobotDataset +import numpy as np + +EXPECTED_FORMAT = "g1-fruit-ninja-pi05-raw-v1" +EXPECTED_STATE_DIM = 29 +EXPECTED_ACTION_DIM = 21 + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(16 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--raw-dir", type=Path, required=True) + parser.add_argument("--repo-id", required=True, help="Hugging Face owner/dataset identifier") + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--push-to-hub", action="store_true") + parser.add_argument("--allow-failures", action="store_true") + return parser.parse_args() + + +def _validate_repo_id(repo_id: str) -> None: + parts = repo_id.split("/") + if len(parts) != 2 or not all(parts): + raise ValueError("--repo-id must be an owner/dataset identifier") + allowed = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-") + if any(set(part) - allowed for part in parts): + raise ValueError("--repo-id contains unsupported characters") + + +def _load_manifest(raw_dir: Path) -> dict: + manifest_path = raw_dir / "manifest.json" + if not manifest_path.is_file(): + raise FileNotFoundError(manifest_path) + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + if manifest.get("format") != EXPECTED_FORMAT: + raise ValueError(f"unsupported raw format: {manifest.get('format')!r}") + if manifest.get("fps") != 50: + raise ValueError(f"expected 50 Hz demonstrations, got {manifest.get('fps')!r}") + if manifest.get("state", {}).get("shape") != [EXPECTED_STATE_DIM]: + raise ValueError("manifest does not contain the 29-value G1 state contract") + if manifest.get("action", {}).get("shape") != [EXPECTED_ACTION_DIM]: + raise ValueError("manifest does not contain the 21-value G1 action contract") + episodes = manifest.get("episodes") + if not isinstance(episodes, list) or not episodes: + raise ValueError("manifest contains no recorded episodes") + return manifest + + +def _load_episode(raw_dir: Path, entry: dict, image_shape: tuple[int, int, int]): + path = raw_dir / entry["path"] + if not path.is_file(): + raise FileNotFoundError(path) + expected_sha256 = entry.get("sha256") + if not isinstance(expected_sha256, str) or _sha256(path) != expected_sha256: + raise ValueError(f"{path} does not match its recorded SHA-256") + with np.load(path, allow_pickle=False) as archive: + image = archive["observation_images_head"] + state = archive["observation_state"] + action = archive["action"] + timestamp = archive["timestamp"] + frame_index = archive["frame_index"] + + frames = len(image) + if frames < 1 or entry.get("frames") != frames: + raise ValueError(f"{path} has an invalid or mismatched frame count") + if image.shape != (frames, *image_shape) or image.dtype != np.uint8: + raise ValueError(f"{path} has invalid RGB shape or dtype: {image.shape} {image.dtype}") + if state.shape != (frames, EXPECTED_STATE_DIM) or state.dtype != np.float32: + raise ValueError(f"{path} has invalid state shape or dtype: {state.shape} {state.dtype}") + if action.shape != (frames, EXPECTED_ACTION_DIM) or action.dtype != np.float32: + raise ValueError(f"{path} has invalid action shape or dtype: {action.shape} {action.dtype}") + if timestamp.shape != (frames,) or frame_index.shape != (frames,): + raise ValueError(f"{path} has invalid timing arrays") + if not np.array_equal(frame_index, np.arange(frames)): + raise ValueError(f"{path} frame indices are not contiguous") + expected_timestamps = np.arange(frames, dtype=np.float64) / 50.0 + if not np.allclose(timestamp, expected_timestamps, atol=1e-6): + raise ValueError(f"{path} timestamps are not synchronized at 50 Hz") + if not np.isfinite(state).all() or not np.isfinite(action).all(): + raise ValueError(f"{path} contains non-finite state or action values") + return image, state, action + + +def main() -> None: + args = parse_args() + _validate_repo_id(args.repo_id) + manifest = _load_manifest(args.raw_dir) + image_shape = tuple(manifest["image"]["shape"]) + if len(image_shape) != 3 or image_shape[-1] != 3: + raise ValueError(f"invalid RGB shape in manifest: {image_shape}") + + output_path = HF_LEROBOT_HOME / args.repo_id + if output_path.exists(): + if not args.overwrite: + raise FileExistsError(f"{output_path} exists; pass --overwrite to replace it") + shutil.rmtree(output_path) + + dataset = LeRobotDataset.create( + repo_id=args.repo_id, + robot_type="unitree_g1", + fps=50, + features={ + "observation.images.head": { + "dtype": "image", + "shape": image_shape, + "names": ["height", "width", "channel"], + }, + "observation.state": { + "dtype": "float32", + "shape": (EXPECTED_STATE_DIM,), + "names": ["state"], + }, + "action": { + "dtype": "float32", + "shape": (EXPECTED_ACTION_DIM,), + "names": ["action"], + }, + }, + use_videos=True, + image_writer_threads=8, + image_writer_processes=4, + ) + + converted = 0 + for entry in manifest["episodes"]: + if not entry.get("success", False) and not args.allow_failures: + raise ValueError( + f"{entry.get('path')} is marked unsuccessful; refusing to convert without --allow-failures" + ) + image, state, action = _load_episode(args.raw_dir, entry, image_shape) + task = entry.get("task") or manifest["task"] + for frame_index in range(len(image)): + dataset.add_frame( + { + "observation.images.head": image[frame_index], + "observation.state": state[frame_index], + "action": action[frame_index], + "task": task, + } + ) + dataset.save_episode() + converted += 1 + + if converted == 0: + raise RuntimeError("no episodes were converted") + if args.push_to_hub: + dataset.push_to_hub( + tags=["unitree-g1", "fruit-ninja", "simulation", "pi0.5"], + private=True, + push_videos=True, + license="apache-2.0", + ) + print(json.dumps({"status": "passed", "episodes": converted, "output": str(output_path)})) + + +if __name__ == "__main__": + main() diff --git a/scripts/spark/evaluate_g1_coke_checkpoint_video.py b/scripts/spark/evaluate_g1_coke_checkpoint_video.py new file mode 100644 index 0000000000..60bab4d403 --- /dev/null +++ b/scripts/spark/evaluate_g1_coke_checkpoint_video.py @@ -0,0 +1,110 @@ +#!/usr/bin/env python3 +"""Render an offline RGB/action comparison video for a pi0.5 checkpoint.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import av +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +from openpi.policies import policy_config +from openpi.training import config as training_config + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint-dir", type=Path, required=True) + parser.add_argument("--episode", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--config-name", default="pi05_spark_g1_coke_pickup") + parser.add_argument("--prompt", default="pick up the Coke can and hold it upright") + parser.add_argument("--samples", type=int, default=24) + parser.add_argument("--denoise-steps", type=int, default=5) + args = parser.parse_args() + + if args.samples < 1 or args.denoise_steps < 1: + raise ValueError("samples and denoise steps must be positive") + with np.load(args.episode, allow_pickle=False) as archive: + rgb = np.asarray(archive["observation_images_head"], dtype=np.uint8) + state = np.asarray(archive["observation_state"], dtype=np.float32) + teacher = np.asarray(archive["action"], dtype=np.float32) + if rgb.ndim != 4 or state.shape != (len(rgb), 24) or teacher.shape != (len(rgb), 21): + raise ValueError("episode does not match the G1 Coke RGB/state/action contract") + + config = training_config.get_config(args.config_name) + policy = policy_config.create_trained_policy( + config, + args.checkpoint_dir, + default_prompt=args.prompt, + sample_kwargs={"num_steps": args.denoise_steps}, + pytorch_device="cuda:0", + ) + sample_ids = np.unique(np.linspace(0, len(rgb) - 1, min(args.samples, len(rgb)), dtype=np.int64)) + args.output.parent.mkdir(parents=True, exist_ok=True) + writer = av.open(str(args.output), mode="w") + stream = writer.add_stream("mpeg4", rate=4) + stream.width = int(rgb.shape[2]) + stream.height = int(rgb.shape[1]) + stream.pix_fmt = "yuv420p" + stream.bit_rate = 1_000_000 + try: + font = ImageFont.load_default(size=13) + except TypeError: + font = ImageFont.load_default() + + maes: list[float] = [] + inference_ms: list[float] = [] + try: + for frame_index in sample_ids: + result = policy.infer( + { + "head_image": rgb[frame_index], + "state": state[frame_index], + "prompt": args.prompt, + } + ) + actions = np.asarray(result["actions"], dtype=np.float32) + if actions.ndim != 2 or actions.shape[1] != 21 or not np.isfinite(actions).all(): + raise ValueError(f"checkpoint returned invalid actions {actions.shape}") + mae = float(np.mean(np.abs(actions[0] - teacher[frame_index]))) + maes.append(mae) + inference_ms.append(float(result.get("policy_timing", {}).get("infer_ms", np.nan))) + frame = Image.fromarray(rgb[frame_index], mode="RGB") + draw = ImageDraw.Draw(frame) + draw.rectangle((0, 0, frame.width, 51), fill=(0, 0, 0)) + draw.text((8, 5), "pi0.5 offline checkpoint evaluation", font=font, fill=(255, 255, 255)) + draw.text( + (8, 29), + f"sim frame {int(frame_index)} | action MAE {mae:.4f} rad", + font=font, + fill=(80, 220, 255), + ) + video_frame = av.VideoFrame.from_ndarray(np.asarray(frame), format="rgb24") + for packet in stream.encode(video_frame): + writer.mux(packet) + finally: + for packet in stream.encode(): + writer.mux(packet) + writer.close() + + report = { + "status": "passed", + "kind": "offline_sim_rgb_action_comparison", + "checkpoint": str(args.checkpoint_dir), + "episode": str(args.episode), + "video": str(args.output), + "samples": len(maes), + "action_mae_rad_mean": float(np.mean(maes)), + "action_mae_rad_max": float(np.max(maes)), + "inference_ms_mean": float(np.nanmean(inference_ms)), + } + args.output.with_suffix(".json").write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + print(json.dumps(report, sort_keys=True), flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/spark/evaluate_g1_coke_rgbd_checkpoint.py b/scripts/spark/evaluate_g1_coke_rgbd_checkpoint.py new file mode 100755 index 0000000000..c8d8fb17e2 --- /dev/null +++ b/scripts/spark/evaluate_g1_coke_rgbd_checkpoint.py @@ -0,0 +1,253 @@ +#!/usr/bin/env python3 +"""Evaluate a G1 RGB-D LoRA checkpoint on held-out demonstrations.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import av +from lerobot.common.datasets.lerobot_dataset import HF_LEROBOT_HOME +from lerobot.common.datasets.lerobot_dataset import LeRobotDataset +from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata +import numpy as np +from PIL import Image +from PIL import ImageDraw +from PIL import ImageFont + +from openpi.policies import policy_config +from openpi.training import config as training_config + +PROMPT = "grasp the Coke can, lift it, and present it" +ARM_MIN_RAD = np.asarray( + [ + -3.0892, + -1.5882, + -2.618, + -1.0472, + -1.97222, + -1.61443, + -1.61443, + -3.0892, + -2.2515, + -2.618, + -1.0472, + -1.97222, + -1.61443, + -1.61443, + ], + dtype=np.float32, +) +ARM_MAX_RAD = np.asarray( + [ + 2.6704, + 2.2515, + 2.618, + 2.0944, + 1.97222, + 1.61443, + 1.61443, + 2.6704, + 1.5882, + 2.618, + 2.0944, + 1.97222, + 1.61443, + 1.61443, + ], + dtype=np.float32, +) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(16 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _numpy(value) -> np.ndarray: + if hasattr(value, "detach"): + value = value.detach().cpu().numpy() + return np.asarray(value) + + +def _image(value) -> np.ndarray: + image = _numpy(value) + if image.ndim != 3: + raise ValueError(f"expected one image, got {image.shape}") + if image.shape[0] == 3: + image = np.moveaxis(image, 0, -1) + if np.issubdtype(image.dtype, np.floating): + image = np.clip(image * 255.0, 0, 255) + image = image.astype(np.uint8) + if image.shape[-1] != 3: + raise ValueError(f"expected three image channels, got {image.shape}") + return image + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint-dir", type=Path, required=True) + parser.add_argument("--repo-id", required=True) + parser.add_argument("--output-json", type=Path, required=True) + parser.add_argument("--output-video", type=Path, required=True) + parser.add_argument("--config-name", default="pi05_spark_g1_coke_rgbd_arm14") + parser.add_argument("--expected-contract", choices=("arm14", "left-arm7"), default="arm14") + parser.add_argument("--samples", type=int, default=32) + parser.add_argument("--denoise-steps", type=int, default=5) + args = parser.parse_args() + if args.samples < 1 or args.denoise_steps < 1: + raise ValueError("samples and denoise steps must be positive") + + conversion_path = HF_LEROBOT_HOME / args.repo_id / "meta" / "wendy-rgbd-conversion.json" + conversion = json.loads(conversion_path.read_text(encoding="utf-8")) + if conversion.get("contract") != args.expected_contract: + raise ValueError( + f"held-out dataset contract {conversion.get('contract')!r} does not match " + f"{args.expected_contract!r}" + ) + if args.expected_contract == "left-arm7": + if conversion.get("state", {}).get("right_arm_included") is not False: + raise ValueError("left-arm7 held-out state unexpectedly includes the right arm") + if conversion.get("action", {}).get("right_arm_included") is not False: + raise ValueError("left-arm7 held-out action unexpectedly includes the right arm") + + metadata = LeRobotDatasetMetadata(args.repo_id) + dataset = LeRobotDataset( + args.repo_id, + delta_timestamps={"action": [step / metadata.fps for step in range(10)]}, + ) + if len(dataset) < 1: + raise ValueError("held-out dataset is empty") + + config = training_config.get_config(args.config_name) + action_dim = int(config.data.task_action_dim) + expected_action_dim = 7 if args.expected_contract == "left-arm7" else 14 + if action_dim != expected_action_dim: + raise ValueError( + f"config {args.config_name!r} has {action_dim} actions, expected {expected_action_dim}" + ) + arm_min_rad = ARM_MIN_RAD[:action_dim] + arm_max_rad = ARM_MAX_RAD[:action_dim] + policy = policy_config.create_trained_policy( + config, + args.checkpoint_dir, + default_prompt=PROMPT, + sample_kwargs={"num_steps": args.denoise_steps}, + pytorch_device="cuda:0", + ) + sample_ids = np.unique(np.linspace(0, len(dataset) - 1, min(args.samples, len(dataset)), dtype=np.int64)) + noise_rng = np.random.default_rng(20260825) + maes: list[float] = [] + max_errors: list[float] = [] + inference_ms: list[float] = [] + temporal_steps: list[float] = [] + joint_error_sum = np.zeros(action_dim, dtype=np.float64) + joint_error_count = 0 + hard_limit_violations = 0 + + args.output_video.parent.mkdir(parents=True, exist_ok=True) + writer = av.open(str(args.output_video), mode="w") + stream = writer.add_stream("mpeg4", rate=4) + stream.width = 640 + stream.height = 240 + stream.pix_fmt = "yuv420p" + stream.bit_rate = 1_500_000 + try: + font = ImageFont.load_default(size=13) + except TypeError: + font = ImageFont.load_default() + + try: + for frame_index in sample_ids: + sample = dataset[int(frame_index)] + rgb = _image(sample["observation.images.head"]) + depth = _image(sample["observation.images.depth"]) + state = _numpy(sample["observation.state"]).astype(np.float32) + teacher = _numpy(sample["action"]).astype(np.float32) + result = policy.infer( + { + "head_image": rgb, + "depth_image": depth, + "state": state, + "prompt": PROMPT, + }, + noise=noise_rng.standard_normal((10, 32), dtype=np.float32), + ) + predicted = np.asarray(result["actions"], dtype=np.float32) + if predicted.shape != (10, action_dim) or not np.isfinite(predicted).all(): + raise ValueError(f"checkpoint returned invalid actions {predicted.shape}") + if teacher.shape != (10, action_dim) or not np.isfinite(teacher).all(): + raise ValueError(f"held-out sample has invalid actions {teacher.shape}") + + error = np.abs(predicted - teacher) + maes.append(float(np.mean(error))) + max_errors.append(float(np.max(error))) + joint_error_sum += np.sum(error, axis=0) + joint_error_count += error.shape[0] + temporal_steps.append(float(np.max(np.abs(np.diff(predicted, axis=0))))) + hard_limit_violations += int( + np.count_nonzero( + (predicted < arm_min_rad[None, :]) | (predicted > arm_max_rad[None, :]) + ) + ) + inference_ms.append(float(result.get("policy_timing", {}).get("infer_ms", np.nan))) + + canvas = Image.new("RGB", (640, 240)) + canvas.paste(Image.fromarray(rgb).resize((320, 240)), (0, 0)) + canvas.paste(Image.fromarray(depth).resize((320, 240)), (320, 0)) + draw = ImageDraw.Draw(canvas) + draw.rectangle((0, 0, 640, 43), fill=(0, 0, 0)) + draw.text((8, 5), "held-out RGB-D checkpoint evaluation", font=font, fill=(255, 255, 255)) + draw.text( + (8, 24), + f"frame {int(frame_index)} | 10-step action MAE {maes[-1]:.4f} rad", + font=font, + fill=(80, 220, 255), + ) + frame = av.VideoFrame.from_ndarray(np.asarray(canvas), format="rgb24") + for packet in stream.encode(frame): + writer.mux(packet) + finally: + for packet in stream.encode(): + writer.mux(packet) + writer.close() + + report = { + "status": "passed" if hard_limit_violations == 0 else "failed", + "kind": "g1_coke_rgbd_heldout_policy_evaluation_v1", + "config": args.config_name, + "contract": args.expected_contract, + "state_dim": int(config.data.state_dim), + "action_dim": action_dim, + "right_arm_inputs": args.expected_contract != "left-arm7", + "right_arm_actions": args.expected_contract != "left-arm7", + "checkpoint": str(args.checkpoint_dir), + "checkpoint_model_sha256": _sha256(args.checkpoint_dir / "model.safetensors"), + "repo_id": args.repo_id, + "dataset_frames": len(dataset), + "samples": len(maes), + "denoise_steps": args.denoise_steps, + "action_mae_rad_mean": float(np.mean(maes)), + "action_mae_rad_p95": float(np.quantile(maes, 0.95)), + "action_abs_error_rad_max": float(np.max(max_errors)), + "per_joint_action_mae_rad": (joint_error_sum / joint_error_count).tolist(), + "predicted_chunk_step_rad_max": float(np.max(temporal_steps)), + "hard_limit_violations": hard_limit_violations, + "inference_ms_mean": float(np.nanmean(inference_ms)), + "video": str(args.output_video), + "video_sha256": _sha256(args.output_video), + } + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(report, sort_keys=True), flush=True) + if report["status"] != "passed": + raise SystemExit(2) + + +if __name__ == "__main__": + main() diff --git a/scripts/spark/launch_g1_coke_48h_host.sh b/scripts/spark/launch_g1_coke_48h_host.sh new file mode 100644 index 0000000000..58215392b4 --- /dev/null +++ b/scripts/spark/launch_g1_coke_48h_host.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +set -euo pipefail + +image="${OPENPI_SPARK_IMAGE:-openpi-pi05-gb10:local}" +data_root="${OPENPI_SPARK_DATA_DIR:-/var/lib/openpi-spark}" +source_root="${OPENPI_SPARK_SOURCE_DIR:-${data_root}/openpi-source}" +raw_dataset_slug="${OPENPI_G1_COKE_RAW_DATASET_SLUG:-g1_coke_pickup_phase1_5cm_upright_v2}" +repo_id="${OPENPI_G1_COKE_DATASET_REPO_ID:-local/g1_coke_pickup_phase1_5cm_upright_v2}" +experiment="${OPENPI_G1_COKE_EXPERIMENT:-g1_coke_pickup_pi05_upright_48h_20260824}" +trainer_name="g1-coke-pickup-pi05-48h" +tensorboard_name="g1-coke-pickup-pi05-tensorboard" + +test -s "${data_root}/checkpoints/pi05_base_pytorch/model.safetensors" +test -s "${data_root}/lerobot/${repo_id}/meta/info.json" +test -s "${data_root}/raw/${raw_dataset_slug}/episode_000000.npz" +test -f "${source_root}/scripts/spark/train_g1_coke_48h.sh" +docker image inspect "${image}" >/dev/null +mkdir -p "${data_root}/reports" + +if docker container inspect "${trainer_name}" >/dev/null 2>&1; then + echo "${trainer_name} already exists; refusing to replace an existing training identity" >&2 + exit 3 +fi + +docker run --detach \ + --name "${trainer_name}" \ + --restart on-failure:20 \ + --gpus all \ + --ipc=host \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --volume "${data_root}:/openpi_assets" \ + --volume "${source_root}:/opt/openpi:ro" \ + --workdir /opt/openpi \ + --env "OPENPI_G1_COKE_DATASET_REPO_ID=${repo_id}" \ + --env "OPENPI_G1_COKE_EXPERIMENT=${experiment}" \ + --env "OPENPI_G1_COKE_RAW_DATASET_SLUG=${raw_dataset_slug}" \ + --env "OPENPI_G1_COKE_SESSION_SECONDS=172800" \ + --env "OPENPI_G1_COKE_VIDEO_INTERVAL_SECONDS=21600" \ + --env "OPENPI_SPARK_DATA_DIR=/openpi_assets" \ + --env "HF_LEROBOT_HOME=/openpi_assets/lerobot" \ + --env "PYTHONPATH=/opt/openpi/src" \ + --env "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True" \ + "${image}" \ + bash -lc "set -o pipefail; bash scripts/spark/train_g1_coke_48h.sh 2>&1 | tee -a /openpi_assets/reports/${experiment}.log" + +if ! docker container inspect "${tensorboard_name}" >/dev/null 2>&1; then + docker run --detach \ + --name "${tensorboard_name}" \ + --restart unless-stopped \ + --network host \ + --volume "${data_root}:/openpi_assets:ro" \ + "${image}" \ + tensorboard \ + --logdir "/openpi_assets/training/pi05_spark_g1_coke_pickup/${experiment}/tensorboard" \ + --bind_all \ + --port 6006 +fi + +docker inspect --format '{{.Name}} {{.State.Status}} {{.State.Pid}}' "${trainer_name}" "${tensorboard_name}" diff --git a/scripts/spark/launch_g1_coke_rgbd_4k_host.sh b/scripts/spark/launch_g1_coke_rgbd_4k_host.sh new file mode 100755 index 0000000000..2b479f6d46 --- /dev/null +++ b/scripts/spark/launch_g1_coke_rgbd_4k_host.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +set -euo pipefail + +image="${OPENPI_SPARK_IMAGE:-openpi-pi05-gb10:local}" +data_root="${OPENPI_SPARK_DATA_DIR:-/var/lib/openpi-spark}" +source_root="${OPENPI_SPARK_SOURCE_DIR:-${data_root}/openpi-source}" +train_repo_id="${OPENPI_G1_COKE_RGBD_DATASET_REPO_ID:-local/g1_coke_pickup_real_rgbd_left_rcoke_3_16_train_v1}" +experiment="${OPENPI_G1_COKE_RGBD_EXPERIMENT:-g1_coke_pickup_real_rgbd_left_arm14_pi05_4k_20260825}" +trainer_name="g1-coke-pickup-pi05-rgbd-arm14" +tensorboard_name="g1-coke-pickup-pi05-rgbd-tensorboard" + +test -s "${data_root}/checkpoints/pi05_base_pytorch/model.safetensors" +test -s "${data_root}/lerobot/${train_repo_id}/meta/info.json" +test -f "${source_root}/scripts/spark/train_g1_coke_rgbd_4k.sh" +docker image inspect "${image}" >/dev/null +mkdir -p "${data_root}/reports" + +if docker container inspect "${trainer_name}" >/dev/null 2>&1; then + echo "${trainer_name} already exists; refusing to replace an existing training identity" >&2 + exit 3 +fi + +docker run --detach \ + --name "${trainer_name}" \ + --restart on-failure:20 \ + --gpus all \ + --ipc=host \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --volume "${data_root}:/openpi_assets" \ + --volume "${source_root}:/opt/openpi:ro" \ + --workdir /opt/openpi \ + --env "OPENPI_G1_COKE_RGBD_DATASET_REPO_ID=${train_repo_id}" \ + --env "OPENPI_G1_COKE_RGBD_EXPERIMENT=${experiment}" \ + --env "OPENPI_SPARK_DATA_DIR=/openpi_assets" \ + --env "HF_LEROBOT_HOME=/openpi_assets/lerobot" \ + --env "PYTHONPATH=/opt/openpi/src" \ + --env "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True" \ + "${image}" \ + bash -lc "set -o pipefail; bash scripts/spark/train_g1_coke_rgbd_4k.sh 2>&1 | tee -a /openpi_assets/reports/${experiment}.log" + +if ! docker container inspect "${tensorboard_name}" >/dev/null 2>&1; then + docker run --detach \ + --name "${tensorboard_name}" \ + --restart unless-stopped \ + --network host \ + --volume "${data_root}:/openpi_assets:ro" \ + "${image}" \ + tensorboard \ + --logdir "/openpi_assets/training/pi05_spark_g1_coke_rgbd_arm14/${experiment}/tensorboard" \ + --bind_all \ + --port 6007 +fi + +docker inspect --format '{{.Name}} {{.State.Status}} {{.State.Pid}}' "${trainer_name}" "${tensorboard_name}" diff --git a/scripts/spark/prepare_g1_coke_dataset_host.sh b/scripts/spark/prepare_g1_coke_dataset_host.sh new file mode 100644 index 0000000000..d20dd6b61f --- /dev/null +++ b/scripts/spark/prepare_g1_coke_dataset_host.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +set -euo pipefail + +image="${OPENPI_SPARK_IMAGE:-openpi-pi05-gb10:local}" +data_root="${OPENPI_SPARK_DATA_DIR:-/var/lib/openpi-spark}" +source_root="${OPENPI_SPARK_SOURCE_DIR:-${data_root}/openpi-source}" +raw_dataset_slug="${OPENPI_G1_COKE_RAW_DATASET_SLUG:-g1_coke_pickup_phase1_5cm_upright_v2}" +repo_id="${OPENPI_G1_COKE_DATASET_REPO_ID:-local/g1_coke_pickup_phase1_5cm_upright_v2}" +raw_dir="${data_root}/raw/${raw_dataset_slug}" + +test -s "${raw_dir}/manifest.json" +test -f "${source_root}/scripts/spark/convert_g1_coke_sim_demos.py" +docker image inspect "${image}" >/dev/null + +docker run --rm \ + --ipc=host \ + --volume "${data_root}:/openpi_assets" \ + --volume "${source_root}:/opt/openpi:ro" \ + --workdir /opt/openpi \ + --env "HF_LEROBOT_HOME=/openpi_assets/lerobot" \ + --env "PYTHONPATH=/opt/openpi/src" \ + "${image}" \ + python3 scripts/spark/convert_g1_coke_sim_demos.py \ + --raw-dir "/openpi_assets/raw/${raw_dataset_slug}" \ + --repo-id "${repo_id}" \ + --overwrite + +test -s "${data_root}/lerobot/${repo_id}/meta/info.json" diff --git a/scripts/spark/prepare_g1_coke_rgbd_dataset_host.sh b/scripts/spark/prepare_g1_coke_rgbd_dataset_host.sh new file mode 100755 index 0000000000..df5e90842a --- /dev/null +++ b/scripts/spark/prepare_g1_coke_rgbd_dataset_host.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +set -euo pipefail + +image="${OPENPI_SPARK_IMAGE:-openpi-pi05-gb10:local}" +data_root="${OPENPI_SPARK_DATA_DIR:-/var/lib/openpi-spark}" +source_root="${OPENPI_SPARK_SOURCE_DIR:-${data_root}/openpi-source}" +raw_dataset_slug="${OPENPI_G1_COKE_RGBD_RAW_DATASET_SLUG:-g1_coke_pickup_real_rgbd_left_rcoke_3_16_v1}" +train_repo_id="${OPENPI_G1_COKE_RGBD_DATASET_REPO_ID:-local/g1_coke_pickup_real_rgbd_left_rcoke_3_16_train_v1}" +eval_repo_id="${OPENPI_G1_COKE_RGBD_EVAL_REPO_ID:-local/g1_coke_pickup_real_rgbd_left_rcoke_3_16_eval_v1}" +raw_dir="${data_root}/raw/${raw_dataset_slug}" + +test -s "${raw_dir}/manifest.json" +test -f "${source_root}/scripts/spark/convert_g1_coke_recorder_rgbd.py" +docker image inspect "${image}" >/dev/null + +docker run --rm \ + --ipc=host \ + --volume "${data_root}:/openpi_assets" \ + --volume "${source_root}:/opt/openpi:ro" \ + --workdir /opt/openpi \ + --env "HF_LEROBOT_HOME=/openpi_assets/lerobot" \ + --env "PYTHONPATH=/opt/openpi/src" \ + "${image}" \ + python3 scripts/spark/convert_g1_coke_recorder_rgbd.py \ + --raw-dir "/openpi_assets/raw/${raw_dataset_slug}" \ + --train-repo-id "${train_repo_id}" \ + --eval-repo-id "${eval_repo_id}" \ + --overwrite + +test -s "${data_root}/lerobot/${train_repo_id}/meta/info.json" +test -s "${data_root}/lerobot/${eval_repo_id}/meta/info.json" +test -s "${data_root}/lerobot/${train_repo_id}/meta/wendy-rgbd-conversion.json" diff --git a/scripts/spark/prepare_pi05.sh b/scripts/spark/prepare_pi05.sh new file mode 100755 index 0000000000..be5f32c59a --- /dev/null +++ b/scripts/spark/prepare_pi05.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +set -euo pipefail + +OPENPI_SPARK_IMAGE="${OPENPI_SPARK_IMAGE:-openpi-pi05-gb10:local}" +OPENPI_SPARK_DATA_DIR="${OPENPI_SPARK_DATA_DIR:-/var/lib/openpi-spark}" +OPENPI_SPARK_SOURCE_CHECKPOINT="${OPENPI_SPARK_SOURCE_CHECKPOINT:-gs://openpi-assets/checkpoints/pi05_base}" +OPENPI_SPARK_OUTPUT_CHECKPOINT="${OPENPI_SPARK_OUTPUT_CHECKPOINT:-/openpi_assets/checkpoints/pi05_base_pytorch}" + +mkdir -p "${OPENPI_SPARK_DATA_DIR}" + +if [[ -s "${OPENPI_SPARK_DATA_DIR}/checkpoints/pi05_base_pytorch/model.safetensors" ]]; then + echo "pi0.5 PyTorch checkpoint already exists; preserving it" + exit 0 +fi + +docker run --rm \ + --ipc=host \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --volume "${OPENPI_SPARK_DATA_DIR}:/openpi_assets" \ + --env JAX_PLATFORMS=cpu \ + --env OPENPI_DATA_HOME=/openpi_assets \ + --env "OPENPI_SPARK_SOURCE_CHECKPOINT=${OPENPI_SPARK_SOURCE_CHECKPOINT}" \ + --env "OPENPI_SPARK_OUTPUT_CHECKPOINT=${OPENPI_SPARK_OUTPUT_CHECKPOINT}" \ + "${OPENPI_SPARK_IMAGE}" \ + bash -lc ' + set -euo pipefail + source_checkpoint="$(python3 -c '\''import os; from openpi.shared.download import maybe_download; print(maybe_download(os.environ["OPENPI_SPARK_SOURCE_CHECKPOINT"]))'\'')" + python3 examples/convert_jax_model_to_pytorch.py \ + --checkpoint-dir "${source_checkpoint}" \ + --config-name pi05_libero \ + --output-path "${OPENPI_SPARK_OUTPUT_CHECKPOINT}" \ + --precision bfloat16 + ' + +test -s "${OPENPI_SPARK_DATA_DIR}/checkpoints/pi05_base_pytorch/model.safetensors" +sha256sum "${OPENPI_SPARK_DATA_DIR}/checkpoints/pi05_base_pytorch/model.safetensors" \ + | tee "${OPENPI_SPARK_DATA_DIR}/checkpoints/pi05_base_pytorch/model.safetensors.sha256" diff --git a/scripts/spark/qualify_g1_coke_rgbd_run.py b/scripts/spark/qualify_g1_coke_rgbd_run.py new file mode 100755 index 0000000000..f0655aaa49 --- /dev/null +++ b/scripts/spark/qualify_g1_coke_rgbd_run.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +"""Fail-closed qualification gate for a G1 Coke RGB-D training run.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +from pathlib import Path +import time + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(16 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument("--evaluation-report", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--minimum-reward-improvement", type=float, default=0.01) + parser.add_argument("--minimum-best-step", type=int, default=500) + parser.add_argument("--maximum-action-mae-rad", type=float, default=0.35) + parser.add_argument("--maximum-chunk-step-rad", type=float, default=0.45) + args = parser.parse_args() + + run_state = json.loads((args.run_dir / "run_state.json").read_text(encoding="utf-8")) + checkpoint_manifest = json.loads((args.run_dir / "reward_checkpoints.json").read_text(encoding="utf-8")) + evaluation = json.loads(args.evaluation_report.read_text(encoding="utf-8")) + records = [ + record + for record in checkpoint_manifest.get("records", []) + if (args.run_dir / str(record.get("step")) / "model.safetensors").is_file() + ] + if not records: + raise RuntimeError("qualification requires at least one complete reward checkpoint") + best = max(records, key=lambda record: float(record["reward"])) + best_dir = args.run_dir / str(best["step"]) + baseline = float(run_state.get("baseline_checkpoint_reward", math.nan)) + best_reward = float(best["reward"]) + action_mae = float(evaluation.get("action_mae_rad_mean", math.inf)) + chunk_step = float(evaluation.get("predicted_chunk_step_rad_max", math.inf)) + reasons: list[str] = [] + if run_state.get("state") != "TRAINING_STEPS_COMPLETE": + reasons.append("training_steps_incomplete") + if not run_state.get("checkpoint_eval_uses_heldout_data"): + reasons.append("checkpoint_evaluation_not_heldout") + if run_state.get("checkpoint_metric") != "heldout_offline_imitation_reward": + reasons.append("wrong_checkpoint_metric") + if not math.isfinite(baseline) or not math.isfinite(best_reward): + reasons.append("nonfinite_reward") + elif best_reward - baseline < args.minimum_reward_improvement: + reasons.append("insufficient_heldout_reward_improvement") + if int(best["step"]) < args.minimum_best_step: + reasons.append("best_checkpoint_too_early") + if evaluation.get("status") != "passed": + reasons.append("policy_evaluation_failed") + if Path(str(evaluation.get("checkpoint", ""))).resolve() != best_dir.resolve(): + reasons.append("evaluation_checkpoint_is_not_best") + if int(evaluation.get("hard_limit_violations", -1)) != 0: + reasons.append("predicted_arm_hard_limit_violation") + if not math.isfinite(action_mae) or action_mae > args.maximum_action_mae_rad: + reasons.append("heldout_action_mae_too_high") + if not math.isfinite(chunk_step) or chunk_step > args.maximum_chunk_step_rad: + reasons.append("predicted_chunk_step_too_large") + + report = { + "status": "qualified" if not reasons else "rejected", + "kind": "g1_coke_rgbd_qualification_v1", + "created_epoch": int(time.time()), + "run_dir": str(args.run_dir), + "checkpoint_eval_repo_id": run_state.get("checkpoint_eval_repo_id"), + "baseline_heldout_reward": baseline, + "best_heldout_reward": best_reward, + "heldout_reward_improvement": best_reward - baseline, + "best_checkpoint_step": int(best["step"]), + "best_checkpoint": str(best_dir), + "best_checkpoint_model_sha256": _sha256(best_dir / "model.safetensors"), + "evaluation_report": str(args.evaluation_report), + "evaluation_report_sha256": _sha256(args.evaluation_report), + "action_mae_rad_mean": action_mae, + "predicted_chunk_step_rad_max": chunk_step, + "rejection_reasons": reasons, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(report, sort_keys=True), flush=True) + if reasons: + raise SystemExit(2) + + +if __name__ == "__main__": + main() diff --git a/scripts/spark/qualify_g1_coke_rgbd_run_test.py b/scripts/spark/qualify_g1_coke_rgbd_run_test.py new file mode 100644 index 0000000000..873de302ec --- /dev/null +++ b/scripts/spark/qualify_g1_coke_rgbd_run_test.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import json +from pathlib import Path +import subprocess +import sys + +SCRIPT = Path(__file__).with_name("qualify_g1_coke_rgbd_run.py") + + +def _fixture(tmp_path: Path, *, best_reward: float = 0.80) -> tuple[Path, Path, Path]: + run_dir = tmp_path / "run" + checkpoint = run_dir / "800" + checkpoint.mkdir(parents=True) + (checkpoint / "model.safetensors").write_bytes(b"model") + (run_dir / "run_state.json").write_text( + json.dumps( + { + "state": "TRAINING_STEPS_COMPLETE", + "checkpoint_eval_uses_heldout_data": True, + "checkpoint_metric": "heldout_offline_imitation_reward", + "checkpoint_eval_repo_id": "local/eval", + "baseline_checkpoint_reward": 0.72, + } + ) + ) + (run_dir / "reward_checkpoints.json").write_text( + json.dumps({"records": [{"step": 800, "reward": best_reward}]}) + ) + evaluation = tmp_path / "evaluation.json" + evaluation.write_text( + json.dumps( + { + "status": "passed", + "checkpoint": str(checkpoint), + "hard_limit_violations": 0, + "action_mae_rad_mean": 0.10, + "predicted_chunk_step_rad_max": 0.20, + } + ) + ) + return run_dir, evaluation, tmp_path / "qualification.json" + + +def _run(run_dir: Path, evaluation: Path, output: Path) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [ + sys.executable, + str(SCRIPT), + "--run-dir", + str(run_dir), + "--evaluation-report", + str(evaluation), + "--output", + str(output), + ], + check=False, + capture_output=True, + text=True, + ) + + +def test_qualification_accepts_improving_heldout_checkpoint(tmp_path: Path) -> None: + run_dir, evaluation, output = _fixture(tmp_path) + + result = _run(run_dir, evaluation, output) + + assert result.returncode == 0, result.stderr + assert json.loads(output.read_text())["status"] == "qualified" + + +def test_qualification_rejects_nonimproving_checkpoint(tmp_path: Path) -> None: + run_dir, evaluation, output = _fixture(tmp_path, best_reward=0.725) + + result = _run(run_dir, evaluation, output) + + assert result.returncode == 2 + report = json.loads(output.read_text()) + assert report["status"] == "rejected" + assert "insufficient_heldout_reward_improvement" in report["rejection_reasons"] diff --git a/scripts/spark/recover_g1_coke_rgbd_qualification.py b/scripts/spark/recover_g1_coke_rgbd_qualification.py new file mode 100644 index 0000000000..f8bd31f09d --- /dev/null +++ b/scripts/spark/recover_g1_coke_rgbd_qualification.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Recover only the known 4k wrapper-state qualification mismatch.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +from pathlib import Path +import time + + +def _load(path: Path) -> dict: + return json.loads(path.read_text(encoding="utf-8")) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(16 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _atomic_write(path: Path, payload: dict) -> None: + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, sort_keys=True) + "\n", encoding="utf-8") + temporary.replace(path) + + +def _validate(run_dir: Path, session_path: Path, qualification_path: Path) -> tuple[dict, dict, dict]: + run_state_path = run_dir / "run_state.json" + session = _load(session_path) + run_state = _load(run_state_path) + qualification = _load(qualification_path) + + if session.get("state") != "SESSION_FAILED" or int(session.get("exit_code", -1)) != 2: + raise ValueError("session is not the known qualification rejection") + if int(session.get("deadline_epoch", 0)) <= int(time.time()) + 900: + raise ValueError("original supervised session deadline has expired or is too close") + if run_state.get("state") != "TRAINING_COMPLETE" or int(run_state.get("global_step", -1)) != 4000: + raise ValueError("4k training did not complete with the legacy wrapper state") + if not run_state.get("checkpoint_eval_uses_heldout_data"): + raise ValueError("checkpoint evaluation was not held out") + if run_state.get("checkpoint_metric") != "heldout_offline_imitation_reward": + raise ValueError("unexpected checkpoint metric") + if qualification.get("status") != "rejected" or qualification.get("rejection_reasons") != [ + "training_steps_incomplete" + ]: + raise ValueError("qualification did not fail solely on the known state mismatch") + if not math.isfinite(float(qualification.get("heldout_reward_improvement", math.nan))): + raise ValueError("qualification improvement is non-finite") + + best_checkpoint = Path(str(qualification.get("best_checkpoint", ""))) + model_path = best_checkpoint / "model.safetensors" + if not model_path.is_file() or _sha256(model_path) != qualification.get("best_checkpoint_model_sha256"): + raise ValueError("best checkpoint identity does not match qualification evidence") + evaluation_path = Path(str(qualification.get("evaluation_report", ""))) + if not evaluation_path.is_file() or _sha256(evaluation_path) != qualification.get("evaluation_report_sha256"): + raise ValueError("held-out evaluation identity does not match qualification evidence") + evaluation = _load(evaluation_path) + if evaluation.get("status") != "passed" or int(evaluation.get("hard_limit_violations", -1)) != 0: + raise ValueError("held-out policy evaluation did not pass") + if Path(str(evaluation.get("checkpoint", ""))).resolve() != best_checkpoint.resolve(): + raise ValueError("held-out evaluation did not use the best checkpoint") + return session, run_state, qualification + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument("--session-state", type=Path, required=True) + parser.add_argument("--qualification-report", type=Path, required=True) + mode = parser.add_mutually_exclusive_group(required=True) + mode.add_argument("--check", action="store_true") + mode.add_argument("--apply", action="store_true") + args = parser.parse_args() + + session, run_state, qualification = _validate( + args.run_dir, + args.session_state, + args.qualification_report, + ) + if args.apply: + recovered_epoch = int(time.time()) + run_state.update( + state="TRAINING_STEPS_COMPLETE", + qualification_recovery_epoch=recovered_epoch, + qualification_recovery_source="legacy_4k_wrapper_state_mismatch", + ) + session.update( + state="QUALIFICATION_RECOVERY_RUNNING", + qualification_recovery_epoch=recovered_epoch, + qualification_recovery_source="legacy_4k_wrapper_state_mismatch", + prior_failure={"failed_epoch": session.get("failed_epoch"), "exit_code": session.get("exit_code")}, + ) + session.pop("failed_epoch", None) + session.pop("exit_code", None) + _atomic_write(args.run_dir / "run_state.json", run_state) + _atomic_write(args.session_state, session) + print( + json.dumps( + { + "status": "recoverable" if args.check else "recovered", + "best_checkpoint": qualification["best_checkpoint"], + "deadline_epoch": session["deadline_epoch"], + }, + sort_keys=True, + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/spark/recover_g1_coke_rgbd_qualification_test.py b/scripts/spark/recover_g1_coke_rgbd_qualification_test.py new file mode 100644 index 0000000000..ee12b8ba70 --- /dev/null +++ b/scripts/spark/recover_g1_coke_rgbd_qualification_test.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +import subprocess +import sys +import time + +SCRIPT = Path(__file__).with_name("recover_g1_coke_rgbd_qualification.py") + + +def _fixture(tmp_path: Path, *, rejection_reasons: list[str] | None = None) -> tuple[Path, Path, Path]: + run_dir = tmp_path / "run" + checkpoint = run_dir / "1400" + checkpoint.mkdir(parents=True) + model = checkpoint / "model.safetensors" + model.write_bytes(b"model") + evaluation = tmp_path / "evaluation.json" + evaluation.write_text( + json.dumps({"status": "passed", "hard_limit_violations": 0, "checkpoint": str(checkpoint)}) + ) + run_state = { + "state": "TRAINING_COMPLETE", + "global_step": 4000, + "checkpoint_eval_uses_heldout_data": True, + "checkpoint_metric": "heldout_offline_imitation_reward", + } + (run_dir / "run_state.json").write_text(json.dumps(run_state)) + session = tmp_path / "session.json" + session.write_text( + json.dumps( + { + "state": "SESSION_FAILED", + "exit_code": 2, + "failed_epoch": int(time.time()), + "started_epoch": int(time.time()) - 100, + "deadline_epoch": int(time.time()) + 3600, + } + ) + ) + qualification = tmp_path / "qualification.json" + qualification.write_text( + json.dumps( + { + "status": "rejected", + "rejection_reasons": rejection_reasons or ["training_steps_incomplete"], + "heldout_reward_improvement": 0.12, + "best_checkpoint": str(checkpoint), + "best_checkpoint_model_sha256": hashlib.sha256(model.read_bytes()).hexdigest(), + "evaluation_report": str(evaluation), + "evaluation_report_sha256": hashlib.sha256(evaluation.read_bytes()).hexdigest(), + } + ) + ) + return run_dir, session, qualification + + +def _run(run_dir: Path, session: Path, qualification: Path, mode: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [ + sys.executable, + str(SCRIPT), + "--run-dir", + str(run_dir), + "--session-state", + str(session), + "--qualification-report", + str(qualification), + mode, + ], + check=False, + capture_output=True, + text=True, + ) + + +def test_recovers_only_known_wrapper_state_mismatch(tmp_path: Path) -> None: + run_dir, session, qualification = _fixture(tmp_path) + + check = _run(run_dir, session, qualification, "--check") + apply = _run(run_dir, session, qualification, "--apply") + + assert check.returncode == 0, check.stderr + assert apply.returncode == 0, apply.stderr + assert json.loads((run_dir / "run_state.json").read_text())["state"] == "TRAINING_STEPS_COMPLETE" + assert json.loads(session.read_text())["state"] == "QUALIFICATION_RECOVERY_RUNNING" + + +def test_refuses_any_additional_qualification_failure(tmp_path: Path) -> None: + run_dir, session, qualification = _fixture( + tmp_path, + rejection_reasons=["training_steps_incomplete", "heldout_action_mae_too_high"], + ) + + result = _run(run_dir, session, qualification, "--check") + + assert result.returncode != 0 + assert json.loads((run_dir / "run_state.json").read_text())["state"] == "TRAINING_COMPLETE" diff --git a/scripts/spark/run_smoke.sh b/scripts/spark/run_smoke.sh new file mode 100755 index 0000000000..781d4885c1 --- /dev/null +++ b/scripts/spark/run_smoke.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +set -euo pipefail + +OPENPI_SPARK_IMAGE="${OPENPI_SPARK_IMAGE:-openpi-pi05-gb10:local}" +OPENPI_SPARK_DATA_DIR="${OPENPI_SPARK_DATA_DIR:-/var/lib/openpi-spark}" +OPENPI_SPARK_REPORT_DIR="${OPENPI_SPARK_DATA_DIR}/reports" + +mkdir -p "${OPENPI_SPARK_REPORT_DIR}" + +docker run --rm \ + --gpus all \ + --ipc=host \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --volume "${OPENPI_SPARK_DATA_DIR}:/openpi_assets" \ + "${OPENPI_SPARK_IMAGE}" \ + python3 scripts/spark/verify_pi05.py \ + --checkpoint-dir /openpi_assets/checkpoints/pi05_base_pytorch \ + --output /openpi_assets/reports/pi05_inference.json + +docker run --rm \ + --gpus all \ + --ipc=host \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --volume "${OPENPI_SPARK_DATA_DIR}:/openpi_assets" \ + "${OPENPI_SPARK_IMAGE}" \ + python3 scripts/train_pytorch.py pi05_spark_smoke \ + 2>&1 | tee "${OPENPI_SPARK_REPORT_DIR}/pi05_train_smoke.log" + +test -s "${OPENPI_SPARK_REPORT_DIR}/pi05_inference.json" +test -s "${OPENPI_SPARK_REPORT_DIR}/pi05_train_smoke.log" +test -s "${OPENPI_SPARK_DATA_DIR}/training/pi05_spark_smoke/gb10_full_model_smoke/1/model.safetensors" + +echo "Inference report: ${OPENPI_SPARK_REPORT_DIR}/pi05_inference.json" +echo "Training log: ${OPENPI_SPARK_REPORT_DIR}/pi05_train_smoke.log" +echo "Fine-tuned checkpoint: ${OPENPI_SPARK_DATA_DIR}/training/pi05_spark_smoke/gb10_full_model_smoke/1" diff --git a/scripts/spark/serve_g1_coke_left_arm7.py b/scripts/spark/serve_g1_coke_left_arm7.py new file mode 100644 index 0000000000..a6cd44f00e --- /dev/null +++ b/scripts/spark/serve_g1_coke_left_arm7.py @@ -0,0 +1,127 @@ +#!/usr/bin/env python3 +"""Serve one hash-pinned physical-left G1 arm7 checkpoint for simulation.""" + +from __future__ import annotations + +import argparse +import hashlib +import logging +from pathlib import Path +import socket +from typing import Any + +POLICY_SERVER_SCHEMA = "wendy.g1.pi05.left-arm7-policy-server.v1" +POLICY_CONFIG_NAME = "pi05_spark_g1_coke_rgbd_left_arm7" +POLICY_PROMPT = "Grasp the Coke can, lift it, and present it in front of the robot." +MODEL_FILENAME = "model.safetensors" +STATE_NAMES = ( + "waist_yaw_joint", + "waist_roll_joint", + "waist_pitch_joint", + "left_shoulder_pitch_joint", + "left_shoulder_roll_joint", + "left_shoulder_yaw_joint", + "left_elbow_joint", + "left_wrist_roll_joint", + "left_wrist_pitch_joint", + "left_wrist_yaw_joint", + "left_hand_thumb_0_joint", + "left_hand_thumb_1_joint", + "left_hand_thumb_2_joint", + "left_hand_middle_0_joint", + "left_hand_middle_1_joint", + "left_hand_index_0_joint", + "left_hand_index_1_joint", +) +ACTION_NAMES = STATE_NAMES[3:10] + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(16 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def checkpoint_identity(checkpoint_dir: Path, expected_model_sha256: str) -> tuple[Path, str]: + resolved = checkpoint_dir.expanduser().resolve() + model_path = resolved / MODEL_FILENAME + if not model_path.is_file(): + raise FileNotFoundError(model_path) + actual = sha256(model_path) + if actual != expected_model_sha256: + raise ValueError( + f"checkpoint model hash mismatch: expected {expected_model_sha256}, got {actual}" + ) + return resolved, actual + + +def build_metadata(base: dict[str, Any], checkpoint_dir: Path, model_sha256: str) -> dict[str, Any]: + metadata = dict(base) + metadata.update( + { + "schema": POLICY_SERVER_SCHEMA, + "config": POLICY_CONFIG_NAME, + "checkpoint_dir": str(checkpoint_dir), + "checkpoint_model_sha256": model_sha256, + "state_dim": len(STATE_NAMES), + "action_dim": len(ACTION_NAMES), + "action_horizon": 10, + "state_names": list(STATE_NAMES), + "action_names": list(ACTION_NAMES), + "semantic_hand": "physical_left", + "right_arm_included": False, + "measured_hand_as_action": False, + "depth_representation": "three-channel uint8 inverse depth", + "prompt": POLICY_PROMPT, + "execution_scope": "software_only_simulation", + } + ) + return metadata + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint-dir", type=Path, required=True) + parser.add_argument("--expected-model-sha256", required=True) + parser.add_argument("--port", type=int, default=8000) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + checkpoint_dir, model_sha256 = checkpoint_identity( + args.checkpoint_dir, + args.expected_model_sha256, + ) + + from openpi.policies import policy_config + from openpi.serving import websocket_policy_server + from openpi.training import config as training_config + + policy = policy_config.create_trained_policy( + training_config.get_config(POLICY_CONFIG_NAME), + str(checkpoint_dir), + default_prompt=POLICY_PROMPT, + ) + metadata = build_metadata(policy.metadata, checkpoint_dir, model_sha256) + logging.info( + "Serving %s on %s:%d with model sha256 %s", + POLICY_CONFIG_NAME, + socket.gethostname(), + args.port, + model_sha256, + ) + websocket_policy_server.WebsocketPolicyServer( + policy=policy, + host="0.0.0.0", + port=args.port, + metadata=metadata, + ).serve_forever() + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, force=True) + main() + diff --git a/scripts/spark/serve_g1_coke_left_arm7_test.py b/scripts/spark/serve_g1_coke_left_arm7_test.py new file mode 100644 index 0000000000..012b382462 --- /dev/null +++ b/scripts/spark/serve_g1_coke_left_arm7_test.py @@ -0,0 +1,44 @@ +import importlib.util +from pathlib import Path +import tempfile + +import pytest + + +def _module(): + path = Path(__file__).with_name("serve_g1_coke_left_arm7.py") + spec = importlib.util.spec_from_file_location("serve_g1_coke_left_arm7", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_checkpoint_identity_and_metadata_are_exact() -> None: + module = _module() + with tempfile.TemporaryDirectory() as directory: + checkpoint = Path(directory) + model = checkpoint / module.MODEL_FILENAME + model.write_bytes(b"left-arm7") + expected = module.sha256(model) + resolved, actual = module.checkpoint_identity(checkpoint, expected) + metadata = module.build_metadata({}, resolved, actual) + assert metadata["schema"] == "wendy.g1.pi05.left-arm7-policy-server.v1" + assert metadata["config"] == "pi05_spark_g1_coke_rgbd_left_arm7" + assert metadata["state_dim"] == 17 + assert metadata["action_dim"] == 7 + assert metadata["action_horizon"] == 10 + assert metadata["semantic_hand"] == "physical_left" + assert metadata["right_arm_included"] is False + assert metadata["measured_hand_as_action"] is False + assert metadata["checkpoint_model_sha256"] == expected + assert all(not name.startswith("right_") for name in metadata["state_names"]) + + +def test_checkpoint_hash_mismatch_fails_closed() -> None: + module = _module() + with tempfile.TemporaryDirectory() as directory: + model = Path(directory) / module.MODEL_FILENAME + model.write_bytes(b"left-arm7") + with pytest.raises(ValueError, match="model hash mismatch"): + module.checkpoint_identity(Path(directory), "0" * 64) diff --git a/scripts/spark/train_g1.sh b/scripts/spark/train_g1.sh new file mode 100755 index 0000000000..191c20c9eb --- /dev/null +++ b/scripts/spark/train_g1.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +set -euo pipefail + +: "${OPENPI_G1_DATASET_REPO_ID:?Set OPENPI_G1_DATASET_REPO_ID to owner/dataset}" + +if [[ ! "${OPENPI_G1_DATASET_REPO_ID}" =~ ^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$ ]]; then + echo "OPENPI_G1_DATASET_REPO_ID must be an owner/dataset identifier" >&2 + exit 2 +fi + +OPENPI_SPARK_IMAGE="${OPENPI_SPARK_IMAGE:-openpi-pi05-gb10:local}" +OPENPI_SPARK_DATA_DIR="${OPENPI_SPARK_DATA_DIR:-/var/lib/openpi-spark}" +OPENPI_G1_NUM_TRAIN_STEPS="${OPENPI_G1_NUM_TRAIN_STEPS:-20000}" +OPENPI_G1_EXPERIMENT="${OPENPI_G1_EXPERIMENT:-fruit_ninja_pi05_$(date -u +%Y%m%dT%H%M%SZ)}" +OPENPI_G1_STATS_PATH="${OPENPI_SPARK_DATA_DIR}/assets/pi05_spark_g1_fruit_ninja/${OPENPI_G1_DATASET_REPO_ID}/norm_stats.json" +OPENPI_G1_REPORT_DIR="${OPENPI_SPARK_DATA_DIR}/reports" + +mkdir -p "${OPENPI_G1_REPORT_DIR}" + +if [[ ! -s "${OPENPI_G1_STATS_PATH}" ]]; then + docker run --rm \ + --ipc=host \ + --volume "${OPENPI_SPARK_DATA_DIR}:/openpi_assets" \ + --env "OPENPI_G1_DATASET_REPO_ID=${OPENPI_G1_DATASET_REPO_ID}" \ + "${OPENPI_SPARK_IMAGE}" \ + python3 scripts/compute_norm_stats.py pi05_spark_g1_fruit_ninja +fi + +test -s "${OPENPI_G1_STATS_PATH}" + +docker run --rm \ + --gpus all \ + --ipc=host \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --volume "${OPENPI_SPARK_DATA_DIR}:/openpi_assets" \ + --env "OPENPI_G1_DATASET_REPO_ID=${OPENPI_G1_DATASET_REPO_ID}" \ + "${OPENPI_SPARK_IMAGE}" \ + python3 scripts/train_pytorch.py pi05_spark_g1_fruit_ninja \ + --exp-name "${OPENPI_G1_EXPERIMENT}" \ + --num-train-steps "${OPENPI_G1_NUM_TRAIN_STEPS}" \ + 2>&1 | tee "${OPENPI_G1_REPORT_DIR}/${OPENPI_G1_EXPERIMENT}.log" + +echo "Training directory: ${OPENPI_SPARK_DATA_DIR}/training/pi05_spark_g1_fruit_ninja/${OPENPI_G1_EXPERIMENT}" +echo "Training log: ${OPENPI_G1_REPORT_DIR}/${OPENPI_G1_EXPERIMENT}.log" diff --git a/scripts/spark/train_g1_coke.sh b/scripts/spark/train_g1_coke.sh new file mode 100755 index 0000000000..5244ff64bd --- /dev/null +++ b/scripts/spark/train_g1_coke.sh @@ -0,0 +1,58 @@ +#!/usr/bin/env bash +set -euo pipefail + +: "${OPENPI_G1_COKE_DATASET_REPO_ID:?Set OPENPI_G1_COKE_DATASET_REPO_ID to owner/dataset}" + +if [[ ! "${OPENPI_G1_COKE_DATASET_REPO_ID}" =~ ^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$ ]]; then + echo "OPENPI_G1_COKE_DATASET_REPO_ID must be an owner/dataset identifier" >&2 + exit 2 +fi + +OPENPI_SPARK_IMAGE="${OPENPI_SPARK_IMAGE:-openpi-pi05-gb10:local}" +OPENPI_SPARK_DATA_DIR="${OPENPI_SPARK_DATA_DIR:-/var/lib/openpi-spark}" +OPENPI_G1_COKE_NUM_TRAIN_STEPS="${OPENPI_G1_COKE_NUM_TRAIN_STEPS:-10000000}" +OPENPI_G1_COKE_MAX_TRAIN_SECONDS="${OPENPI_G1_COKE_MAX_TRAIN_SECONDS:-172800}" +OPENPI_G1_COKE_EXPERIMENT="${OPENPI_G1_COKE_EXPERIMENT:-g1_coke_pickup_pi05_$(date -u +%Y%m%dT%H%M%SZ)}" +OPENPI_G1_COKE_STATS_PATH="${OPENPI_SPARK_DATA_DIR}/assets/pi05_spark_g1_coke_pickup/${OPENPI_G1_COKE_DATASET_REPO_ID}/norm_stats.json" +OPENPI_G1_COKE_REPORT_DIR="${OPENPI_SPARK_DATA_DIR}/reports" +OPENPI_G1_COKE_TRAINING_DIR="${OPENPI_SPARK_DATA_DIR}/training/pi05_spark_g1_coke_pickup/${OPENPI_G1_COKE_EXPERIMENT}" +OPENPI_LEROBOT_HOME="${OPENPI_SPARK_DATA_DIR}/lerobot" + +mkdir -p "${OPENPI_G1_COKE_REPORT_DIR}" + +if [[ ! -s "${OPENPI_G1_COKE_STATS_PATH}" ]]; then + docker run --rm \ + --ipc=host \ + --volume "${OPENPI_SPARK_DATA_DIR}:/openpi_assets" \ + --env "HF_LEROBOT_HOME=/openpi_assets/lerobot" \ + --env "OPENPI_G1_COKE_DATASET_REPO_ID=${OPENPI_G1_COKE_DATASET_REPO_ID}" \ + "${OPENPI_SPARK_IMAGE}" \ + python3 scripts/compute_norm_stats.py pi05_spark_g1_coke_pickup +fi + +test -s "${OPENPI_G1_COKE_STATS_PATH}" + +resume_args=() +if find "${OPENPI_G1_COKE_TRAINING_DIR}" -mindepth 1 -maxdepth 1 -type d -name '[0-9]*' -print -quit 2>/dev/null | grep -q .; then + resume_args+=(--resume) +fi + +docker run --rm \ + --gpus all \ + --ipc=host \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --volume "${OPENPI_SPARK_DATA_DIR}:/openpi_assets" \ + --env "HF_LEROBOT_HOME=/openpi_assets/lerobot" \ + --env "OPENPI_G1_COKE_DATASET_REPO_ID=${OPENPI_G1_COKE_DATASET_REPO_ID}" \ + "${OPENPI_SPARK_IMAGE}" \ + python3 scripts/train_pytorch.py pi05_spark_g1_coke_pickup \ + --exp-name "${OPENPI_G1_COKE_EXPERIMENT}" \ + --num-train-steps "${OPENPI_G1_COKE_NUM_TRAIN_STEPS}" \ + --max-train-seconds "${OPENPI_G1_COKE_MAX_TRAIN_SECONDS}" \ + "${resume_args[@]}" \ + 2>&1 | tee "${OPENPI_G1_COKE_REPORT_DIR}/${OPENPI_G1_COKE_EXPERIMENT}.log" + +echo "Training directory: ${OPENPI_SPARK_DATA_DIR}/training/pi05_spark_g1_coke_pickup/${OPENPI_G1_COKE_EXPERIMENT}" +echo "Training log: ${OPENPI_G1_COKE_REPORT_DIR}/${OPENPI_G1_COKE_EXPERIMENT}.log" +echo "TensorBoard log: ${OPENPI_G1_COKE_TRAINING_DIR}/tensorboard" diff --git a/scripts/spark/train_g1_coke_48h.sh b/scripts/spark/train_g1_coke_48h.sh new file mode 100644 index 0000000000..af8875d8d2 --- /dev/null +++ b/scripts/spark/train_g1_coke_48h.sh @@ -0,0 +1,67 @@ +#!/usr/bin/env bash +set -euo pipefail + +: "${OPENPI_G1_COKE_DATASET_REPO_ID:?Set OPENPI_G1_COKE_DATASET_REPO_ID}" + +experiment="${OPENPI_G1_COKE_EXPERIMENT:-g1_coke_pickup_pi05_48h_20260824}" +raw_dataset_slug="${OPENPI_G1_COKE_RAW_DATASET_SLUG:-g1_coke_pickup_phase1_5cm_upright_v2}" +session_seconds="${OPENPI_G1_COKE_SESSION_SECONDS:-172800}" +video_interval_seconds="${OPENPI_G1_COKE_VIDEO_INTERVAL_SECONDS:-21600}" +data_root="${OPENPI_SPARK_DATA_DIR:-/openpi_assets}" +run_dir="${data_root}/training/pi05_spark_g1_coke_pickup/${experiment}" +stats_path="${data_root}/assets/pi05_spark_g1_coke_pickup/${OPENPI_G1_COKE_DATASET_REPO_ID}/norm_stats.json" +episode="${OPENPI_G1_COKE_EVAL_EPISODE:-${data_root}/raw/${raw_dataset_slug}/episode_000000.npz}" +video_dir="${data_root}/videos/${experiment}" +session_path="${run_dir}/session_state.json" + +mkdir -p "${run_dir}" "${video_dir}" +if [[ ! -s "${stats_path}" ]]; then + python3 scripts/compute_norm_stats.py --config-name pi05_spark_g1_coke_pickup +fi +test -s "${stats_path}" +test -s "${episode}" + +if [[ -s "${session_path}" ]]; then + deadline_epoch="$(python3 -c 'import json,sys; print(int(json.load(open(sys.argv[1]))["deadline_epoch"]))' "${session_path}")" +else + started_epoch="$(date +%s)" + deadline_epoch="$((started_epoch + session_seconds))" + python3 -c 'import json,sys; json.dump({"state":"SESSION_RUNNING","started_epoch":int(sys.argv[2]),"deadline_epoch":int(sys.argv[3]),"video_interval_seconds":int(sys.argv[4])}, open(sys.argv[1],"w"), sort_keys=True)' "${session_path}" "${started_epoch}" "${deadline_epoch}" "${video_interval_seconds}" +fi + +while [[ "$(date +%s)" -lt "${deadline_epoch}" ]]; do + remaining="$((deadline_epoch - $(date +%s)))" + chunk_seconds="${video_interval_seconds}" + if [[ "${remaining}" -lt "${chunk_seconds}" ]]; then + chunk_seconds="${remaining}" + fi + resume_args=() + if find "${run_dir}" -mindepth 1 -maxdepth 1 -type d -name '[0-9]*' -print -quit | grep -q .; then + resume_args+=(--resume) + fi + python3 scripts/train_pytorch.py pi05_spark_g1_coke_pickup \ + --exp-name "${experiment}" \ + --num-train-steps 10000000 \ + --max-train-seconds "${chunk_seconds}" \ + "${resume_args[@]}" + + checkpoint="$(python3 -c ' +import json, pathlib, sys +root = pathlib.Path(sys.argv[1]) +manifest = json.loads((root / "reward_checkpoints.json").read_text()) +records = [r for r in manifest.get("records", []) if (root / str(r["step"])).is_dir()] +if not records: + raise SystemExit("no improving reward checkpoint is available for evaluation or resume") +best = max(records, key=lambda record: float(record["reward"])) +print(root / str(best["step"])) +' "${run_dir}")" + step="$(basename "${checkpoint}")" + python3 scripts/spark/evaluate_g1_coke_checkpoint_video.py \ + --checkpoint-dir "${checkpoint}" \ + --episode "${episode}" \ + --output "${video_dir}/checkpoint_${step}.mp4" \ + --samples 24 \ + --denoise-steps 5 +done + +python3 -c 'import json,sys,time; p=json.load(open(sys.argv[1])); p.update(state="SESSION_DEADLINE_COMPLETE", completed_epoch=int(time.time())); json.dump(p, open(sys.argv[1],"w"), sort_keys=True)' "${session_path}" diff --git a/scripts/spark/train_g1_coke_rgbd_4k.sh b/scripts/spark/train_g1_coke_rgbd_4k.sh new file mode 100755 index 0000000000..5f0b533aa9 --- /dev/null +++ b/scripts/spark/train_g1_coke_rgbd_4k.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +set -euo pipefail + +: "${OPENPI_G1_COKE_RGBD_DATASET_REPO_ID:?Set OPENPI_G1_COKE_RGBD_DATASET_REPO_ID}" +: "${OPENPI_G1_COKE_RGBD_EVAL_REPO_ID:=local/g1_coke_pickup_real_rgbd_left_rcoke_3_16_eval_v1}" +export OPENPI_G1_COKE_RGBD_EVAL_REPO_ID + +experiment="${OPENPI_G1_COKE_RGBD_EXPERIMENT:-g1_coke_pickup_real_rgbd_left_arm14_pi05_4k_20260825}" +data_root="${OPENPI_SPARK_DATA_DIR:-/openpi_assets}" +run_dir="${data_root}/training/pi05_spark_g1_coke_rgbd_arm14/${experiment}" +stats_path="${data_root}/assets/pi05_spark_g1_coke_rgbd_arm14/${OPENPI_G1_COKE_RGBD_DATASET_REPO_ID}/norm_stats.json" +eval_info_path="${data_root}/lerobot/${OPENPI_G1_COKE_RGBD_EVAL_REPO_ID}/meta/info.json" +run_state_path="${run_dir}/run_state.json" + +mkdir -p "${run_dir}" +test -s "${eval_info_path}" +python3 -c 'import json,sys,time; json.dump({"state":"PREPARING","started_epoch":int(time.time()),"config":"pi05_spark_g1_coke_rgbd_arm14","dataset_repo_id":sys.argv[2],"checkpoint_eval_repo_id":sys.argv[3],"target_steps":4000}, open(sys.argv[1],"w"), sort_keys=True)' "${run_state_path}" "${OPENPI_G1_COKE_RGBD_DATASET_REPO_ID}" "${OPENPI_G1_COKE_RGBD_EVAL_REPO_ID}" + +if [[ ! -s "${stats_path}" ]]; then + python3 scripts/compute_norm_stats.py --config-name pi05_spark_g1_coke_rgbd_arm14 +fi +test -s "${stats_path}" + +python3 -c 'import json,sys,time; p=json.load(open(sys.argv[1])); p.update(state="TRAINING_RUNNING",training_started_epoch=int(time.time())); json.dump(p, open(sys.argv[1],"w"), sort_keys=True)' "${run_state_path}" + +python3 scripts/train_pytorch.py pi05_spark_g1_coke_rgbd_arm14 \ + --exp-name "${experiment}" \ + --num-train-steps 4000 + +python3 -c 'import json,sys; p=json.load(open(sys.argv[1])); assert p.get("state") == "TRAINING_STEPS_COMPLETE", p.get("state")' "${run_state_path}" diff --git a/scripts/spark/train_g1_coke_rgbd_left_arm7_4k.sh b/scripts/spark/train_g1_coke_rgbd_left_arm7_4k.sh new file mode 100755 index 0000000000..6804893c4d --- /dev/null +++ b/scripts/spark/train_g1_coke_rgbd_left_arm7_4k.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash +set -euo pipefail + +: "${OPENPI_G1_COKE_LEFT_ONLY_DATASET_REPO_ID:?Set OPENPI_G1_COKE_LEFT_ONLY_DATASET_REPO_ID}" +: "${OPENPI_G1_COKE_LEFT_ONLY_EVAL_REPO_ID:=local/g1_coke_pickup_real_rgbd_left_only_rcoke_3_16_eval_v1}" +export OPENPI_G1_COKE_LEFT_ONLY_EVAL_REPO_ID + +config="pi05_spark_g1_coke_rgbd_left_arm7" +experiment="${OPENPI_G1_COKE_LEFT_ONLY_EXPERIMENT:-g1_coke_pickup_real_rgbd_left_arm7_upright_qualification_20260826}" +data_root="${OPENPI_SPARK_DATA_DIR:-/openpi_assets}" +run_dir="${data_root}/training/${config}/${experiment}" +stats_path="${data_root}/assets/${config}/${OPENPI_G1_COKE_LEFT_ONLY_DATASET_REPO_ID}/norm_stats.json" +train_info_path="${data_root}/lerobot/${OPENPI_G1_COKE_LEFT_ONLY_DATASET_REPO_ID}/meta/info.json" +train_contract_path="${data_root}/lerobot/${OPENPI_G1_COKE_LEFT_ONLY_DATASET_REPO_ID}/meta/wendy-rgbd-conversion.json" +eval_info_path="${data_root}/lerobot/${OPENPI_G1_COKE_LEFT_ONLY_EVAL_REPO_ID}/meta/info.json" +eval_contract_path="${data_root}/lerobot/${OPENPI_G1_COKE_LEFT_ONLY_EVAL_REPO_ID}/meta/wendy-rgbd-conversion.json" +run_state_path="${run_dir}/run_state.json" + +mkdir -p "${run_dir}" +test -s "${train_info_path}" +test -s "${eval_info_path}" +python3 - "${train_contract_path}" "${eval_contract_path}" <<'PY' +import json +import sys + +for path in sys.argv[1:]: + payload = json.load(open(path, encoding="utf-8")) + assert payload.get("contract") == "left-arm7", payload.get("contract") + assert payload.get("state", {}).get("width") == 17 + assert payload.get("state", {}).get("right_arm_included") is False + assert payload.get("action", {}).get("width") == 7 + assert payload.get("action", {}).get("right_arm_included") is False +PY + +python3 - "${run_state_path}" "${OPENPI_G1_COKE_LEFT_ONLY_DATASET_REPO_ID}" "${OPENPI_G1_COKE_LEFT_ONLY_EVAL_REPO_ID}" <<'PY' +import json +import sys +import time + +json.dump( + { + "state": "PREPARING", + "started_epoch": int(time.time()), + "config": "pi05_spark_g1_coke_rgbd_left_arm7", + "contract": "left-arm7", + "state_dim": 17, + "action_dim": 7, + "right_arm_inputs": False, + "right_arm_actions": False, + "dataset_repo_id": sys.argv[2], + "checkpoint_eval_repo_id": sys.argv[3], + "target_steps": 4000, + }, + open(sys.argv[1], "w", encoding="utf-8"), + sort_keys=True, +) +PY + +if [[ ! -s "${stats_path}" ]]; then + python3 scripts/compute_norm_stats.py --config-name "${config}" +fi +test -s "${stats_path}" + +python3 scripts/train_pytorch.py "${config}" \ + --exp-name "${experiment}" \ + --num-train-steps 4000 + +python3 -c 'import json,sys; p=json.load(open(sys.argv[1])); assert p.get("state") == "TRAINING_STEPS_COMPLETE", p.get("state")' "${run_state_path}" diff --git a/scripts/spark/train_g1_coke_rgbd_session.sh b/scripts/spark/train_g1_coke_rgbd_session.sh new file mode 100755 index 0000000000..81586611ef --- /dev/null +++ b/scripts/spark/train_g1_coke_rgbd_session.sh @@ -0,0 +1,130 @@ +#!/usr/bin/env bash +set -euo pipefail + +: "${OPENPI_G1_COKE_RGBD_DATASET_REPO_ID:?Set OPENPI_G1_COKE_RGBD_DATASET_REPO_ID}" +: "${OPENPI_G1_COKE_RGBD_EVAL_REPO_ID:?Set OPENPI_G1_COKE_RGBD_EVAL_REPO_ID}" + +experiment="${OPENPI_G1_COKE_RGBD_EXPERIMENT:-g1_coke_rgbd_arm14_augmented_7h_20260825}" +session_seconds="${OPENPI_G1_COKE_RGBD_SESSION_SECONDS:-25200}" +evaluation_interval_seconds="${OPENPI_G1_COKE_RGBD_EVAL_INTERVAL_SECONDS:-7200}" +data_root="${OPENPI_SPARK_DATA_DIR:-/openpi_assets}" +run_dir="${data_root}/training/pi05_spark_g1_coke_rgbd_arm14/${experiment}" +report_dir="${data_root}/reports/${experiment}" +video_dir="${data_root}/videos/${experiment}" +session_path="${run_dir}/session_state.json" +qualification_path="${report_dir}/qualification.json" + +if [[ "${session_seconds}" -lt 10800 ]]; then + echo "RGB-D session must allow at least three hours for qualification" >&2 + exit 2 +fi +if [[ "${evaluation_interval_seconds}" -lt 900 ]]; then + echo "RGB-D evaluation interval must be at least 15 minutes" >&2 + exit 2 +fi +mkdir -p "${run_dir}" "${report_dir}" "${video_dir}" + +recovery_mode=false +if python3 scripts/spark/recover_g1_coke_rgbd_qualification.py \ + --run-dir "${run_dir}" \ + --session-state "${session_path}" \ + --qualification-report "${qualification_path}" \ + --check >/dev/null 2>&1; then + recovery_mode=true + read -r started_epoch deadline_epoch < <( + python3 -c 'import json,sys; p=json.load(open(sys.argv[1])); print(p["started_epoch"], p["deadline_epoch"])' "${session_path}" + ) + python3 scripts/spark/recover_g1_coke_rgbd_qualification.py \ + --run-dir "${run_dir}" \ + --session-state "${session_path}" \ + --qualification-report "${qualification_path}" \ + --apply +else + started_epoch="$(date +%s)" + deadline_epoch="$((started_epoch + session_seconds))" + python3 -c ' +import json,sys +json.dump({"state":"QUALIFICATION_RUNNING","started_epoch":int(sys.argv[2]),"deadline_epoch":int(sys.argv[3]),"dataset_repo_id":sys.argv[4],"checkpoint_eval_repo_id":sys.argv[5]},open(sys.argv[1],"w"),sort_keys=True) +' "${session_path}" "${started_epoch}" "${deadline_epoch}" "${OPENPI_G1_COKE_RGBD_DATASET_REPO_ID}" "${OPENPI_G1_COKE_RGBD_EVAL_REPO_ID}" +fi + +on_error() { + local exit_code="$?" + trap - ERR + python3 -c ' +import json,sys,time +path=sys.argv[1]; payload=json.load(open(path)); payload.update(state="SESSION_FAILED",failed_epoch=int(time.time()),exit_code=int(sys.argv[2])); json.dump(payload,open(path,"w"),sort_keys=True) +' "${session_path}" "${exit_code}" || true + exec tail -f /dev/null +} +trap on_error ERR + +best_checkpoint() { + python3 -c ' +import json,pathlib,sys +root=pathlib.Path(sys.argv[1]) +manifest=json.loads((root/"reward_checkpoints.json").read_text()) +records=[record for record in manifest.get("records",[]) if (root/str(record["step"])/"model.safetensors").is_file()] +if not records: raise SystemExit("no complete reward checkpoint") +print(root/str(max(records,key=lambda record:float(record["reward"]))["step"])) +' "${run_dir}" +} + +evaluate_checkpoint() { + local checkpoint="$1" + local label="$2" + python3 scripts/spark/evaluate_g1_coke_rgbd_checkpoint.py \ + --checkpoint-dir "${checkpoint}" \ + --repo-id "${OPENPI_G1_COKE_RGBD_EVAL_REPO_ID}" \ + --output-json "${report_dir}/${label}.json" \ + --output-video "${video_dir}/${label}.mp4" \ + --samples 32 \ + --denoise-steps 5 +} + +if [[ "${recovery_mode}" == false ]]; then + bash scripts/spark/train_g1_coke_rgbd_4k.sh +fi +checkpoint="$(best_checkpoint)" +step="$(basename "${checkpoint}")" +evaluate_checkpoint "${checkpoint}" "qualification_checkpoint_${step}" +python3 scripts/spark/qualify_g1_coke_rgbd_run.py \ + --run-dir "${run_dir}" \ + --evaluation-report "${report_dir}/qualification_checkpoint_${step}.json" \ + --output "${qualification_path}" + +python3 -c ' +import json,sys,time +path=sys.argv[1]; payload=json.load(open(path)); payload.update(state="QUALIFIED_CONTINUATION_RUNNING",qualified_epoch=int(time.time()),qualification_report=sys.argv[2]); json.dump(payload,open(path,"w"),sort_keys=True) +' "${session_path}" "${qualification_path}" + +while [[ "$(date +%s)" -lt "${deadline_epoch}" ]]; do + remaining="$((deadline_epoch - $(date +%s)))" + chunk_seconds="${evaluation_interval_seconds}" + if [[ "${remaining}" -lt "${chunk_seconds}" ]]; then + chunk_seconds="${remaining}" + fi + python3 scripts/train_pytorch.py pi05_spark_g1_coke_rgbd_arm14 \ + --exp-name "${experiment}" \ + --num-train-steps 10000000 \ + --max-train-seconds "${chunk_seconds}" \ + --resume + checkpoint="$(best_checkpoint)" + step="$(basename "${checkpoint}")" + label="checkpoint_${step}_$(date +%s)" + evaluate_checkpoint "${checkpoint}" "${label}" + python3 -c ' +import json,sys,time +path=sys.argv[1]; payload=json.load(open(path)); payload.update(state="QUALIFIED_CONTINUATION_RUNNING",last_evaluation_epoch=int(time.time()),last_evaluation_report=sys.argv[2],best_checkpoint=sys.argv[3]); json.dump(payload,open(path,"w"),sort_keys=True) +' "${session_path}" "${report_dir}/${label}.json" "${checkpoint}" +done + +checkpoint="$(best_checkpoint)" +python3 -c ' +import json,sys,time +path=sys.argv[1]; payload=json.load(open(path)); payload.update(state="SESSION_COMPLETE",completed_epoch=int(time.time()),best_checkpoint=sys.argv[2]); json.dump(payload,open(path,"w"),sort_keys=True) +' "${session_path}" "${checkpoint}" + +# Wendy treats an exited one-shot service as unhealthy. Keep the completed +# session inspectable without restarting training or changing the checkpoint. +exec tail -f /dev/null diff --git a/scripts/spark/verify_pi05.py b/scripts/spark/verify_pi05.py new file mode 100755 index 0000000000..65799f82a0 --- /dev/null +++ b/scripts/spark/verify_pi05.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +"""Run a real pi0.5 checkpoint through GB10 GPU inference and emit evidence.""" + +import argparse +import dataclasses +import hashlib +import json +import math +import pathlib +import platform +import time + +import safetensors.torch +import torch + +from openpi.models import model as model_api +from openpi.models_pytorch import pi0_pytorch +from openpi.training import config as training_config + + +def _sha256(path: pathlib.Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as checkpoint_file: + while chunk := checkpoint_file.read(16 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _observation(config, device: torch.device) -> model_api.Observation: + images = {key: torch.zeros((1, 3, 224, 224), dtype=torch.float32, device=device) for key in model_api.IMAGE_KEYS} + image_masks = {key: torch.ones((1,), dtype=torch.bool, device=device) for key in model_api.IMAGE_KEYS} + tokenized_prompt = torch.ones((1, config.max_token_len), dtype=torch.int64, device=device) + tokenized_prompt_mask = torch.zeros((1, config.max_token_len), dtype=torch.bool, device=device) + tokenized_prompt_mask[:, :8] = True + return model_api.Observation( + images=images, + image_masks=image_masks, + state=torch.zeros((1, config.action_dim), dtype=torch.float32, device=device), + tokenized_prompt=tokenized_prompt, + tokenized_prompt_mask=tokenized_prompt_mask, + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint-dir", type=pathlib.Path, required=True) + parser.add_argument("--config-name", default="pi05_spark_smoke") + parser.add_argument("--num-steps", type=int, default=2) + parser.add_argument("--output", type=pathlib.Path) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is unavailable; refusing to report a CPU-only pi0.5 pass") + + device = torch.device("cuda:0") + capability = torch.cuda.get_device_capability(device) + if capability < (12, 1): + raise RuntimeError(f"Expected a GB10-class sm_121 GPU, found sm_{capability[0]}{capability[1]}") + + train_config = training_config.get_config(args.config_name) + model_config = dataclasses.replace(train_config.model, pytorch_compile_mode=None) + if not model_config.pi05: + raise ValueError(f"{args.config_name} is not a pi0.5 config") + + weight_path = args.checkpoint_dir / "model.safetensors" + if not weight_path.is_file(): + raise FileNotFoundError(weight_path) + + torch.manual_seed(0) + torch.cuda.reset_peak_memory_stats(device) + model = pi0_pytorch.PI0Pytorch(model_config) + safetensors.torch.load_model(model, weight_path) + model.to(device).eval() + + observation = _observation(model_config, device) + torch.cuda.synchronize(device) + started = time.perf_counter() + with torch.inference_mode(): + actions = model.sample_actions(device, observation, num_steps=args.num_steps) + torch.cuda.synchronize(device) + elapsed = time.perf_counter() - started + + actions_cpu = actions.float().cpu() + finite = bool(torch.isfinite(actions_cpu).all().item()) + if not finite: + raise RuntimeError("pi0.5 produced non-finite actions") + + report = { + "status": "passed", + "model": "pi0.5", + "config": args.config_name, + "checkpoint": { + "path": str(weight_path), + "bytes": weight_path.stat().st_size, + "sha256": _sha256(weight_path), + }, + "runtime": { + "architecture": platform.machine(), + "python": platform.python_version(), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(device), + "compute_capability": list(capability), + }, + "inference": { + "denoise_steps": args.num_steps, + "elapsed_seconds": elapsed, + "action_shape": list(actions_cpu.shape), + "all_finite": finite, + "action_mean": actions_cpu.mean().item(), + "action_std": actions_cpu.std().item(), + "action_l2": math.sqrt(torch.square(actions_cpu).sum().item()), + "peak_cuda_bytes": torch.cuda.max_memory_allocated(device), + }, + } + + rendered = json.dumps(report, indent=2, sort_keys=True) + print(rendered) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(f"{rendered}\n") + + +if __name__ == "__main__": + main() diff --git a/scripts/train_pytorch.py b/scripts/train_pytorch.py index c7ddd2b595..b3716f4ae0 100644 --- a/scripts/train_pytorch.py +++ b/scripts/train_pytorch.py @@ -25,7 +25,9 @@ import dataclasses import gc +import json import logging +import math import os import platform import shutil @@ -37,14 +39,19 @@ import torch import torch.distributed as dist import torch.nn.parallel +from torch.utils.tensorboard import SummaryWriter import tqdm import wandb import openpi.models.pi0_config +import openpi.models_pytorch.lora as _lora import openpi.models_pytorch.pi0_pytorch import openpi.shared.normalize as _normalize +from openpi.training.checkpoint_evaluation import CheckpointEvaluationPlan +from openpi.training.checkpoint_evaluation import resolve_checkpoint_evaluation import openpi.training.config as _config import openpi.training.data_loader as _data +from openpi.training.reward_checkpoints import RewardCheckpointStore def init_logging(): @@ -128,101 +135,117 @@ def build_datasets(config: _config.TrainConfig): return data_loader, data_loader.data_config() -def get_model_state_dict(model): - """Get state dict from model, handling DDP wrapper.""" - return ( - model.module.state_dict() - if isinstance(model, torch.nn.parallel.DistributedDataParallel) - else model.state_dict() +def build_reward_batches(config: _config.TrainConfig): + """Materialize a deterministic, episode-spanning scoring set.""" + loader = _data.create_data_loader(config, framework="pytorch", shuffle=False) + if not isinstance(loader, _data.DataLoaderImpl): + raise TypeError("Offline reward checkpointing requires the finite PyTorch dataset loader") + batches = loader.evenly_spaced_batches( + num_batches=config.checkpoint_reward_batches, + batch_size=config.batch_size, ) + if not batches: + raise RuntimeError("Cannot checkpoint by reward without evaluation batches") + return batches -def get_model_parameters(model): - """Get parameters from model, handling DDP wrapper.""" - return ( - model.module.parameters() - if isinstance(model, torch.nn.parallel.DistributedDataParallel) - else model.parameters() - ) - - -def save_checkpoint(model, optimizer, global_step, config, is_main, data_config): - """Save a checkpoint with model state, optimizer state, and metadata.""" - if not is_main: - return - - # Only save if it's time to save or if it's the final step - if (global_step % config.save_interval == 0 and global_step > 0) or global_step == config.num_train_steps - 1: - # Create temporary directory for atomic checkpoint saving - final_ckpt_dir = config.checkpoint_dir / f"{global_step}" - tmp_ckpt_dir = config.checkpoint_dir / f"tmp_{global_step}" +def unwrap_model(model): + return model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model - # Remove any existing temp directory and create new one - if tmp_ckpt_dir.exists(): - shutil.rmtree(tmp_ckpt_dir) - tmp_ckpt_dir.mkdir(parents=True, exist_ok=True) - - # Save model state using safetensors (handle shared tensors) - model_to_save = model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model - safetensors.torch.save_model(model_to_save, tmp_ckpt_dir / "model.safetensors") - - # Save optimizer state using PyTorch format - torch.save(optimizer.state_dict(), tmp_ckpt_dir / "optimizer.pt") - - # Save training metadata (avoid saving full config to prevent JAX/Flax compatibility issues) - metadata = { - "global_step": global_step, - "config": dataclasses.asdict(config), - "timestamp": time.time(), - } - torch.save(metadata, tmp_ckpt_dir / "metadata.pt") - # save norm stats - norm_stats = data_config.norm_stats - if norm_stats is not None and data_config.asset_id is not None: - _normalize.save(tmp_ckpt_dir / "assets" / data_config.asset_id, norm_stats) +def trainable_parameters(model): + return [parameter for parameter in model.parameters() if parameter.requires_grad] - # Atomically move temp directory to final location - if final_ckpt_dir.exists(): - shutil.rmtree(final_ckpt_dir) - tmp_ckpt_dir.rename(final_ckpt_dir) - logging.info(f"Saved checkpoint at step {global_step} -> {final_ckpt_dir}") - - # Log checkpoint to wandb - if config.wandb_enabled: - wandb.log({"checkpoint_step": global_step}, step=global_step) - - -def load_checkpoint(model, optimizer, checkpoint_dir, device): - """Load the latest checkpoint and return the global step.""" - checkpoint_steps = [ - int(d.name) - for d in checkpoint_dir.iterdir() - if d.is_dir() and d.name.isdigit() and not d.name.startswith("tmp_") - ] - - if not checkpoint_steps: - raise FileNotFoundError(f"No checkpoints found in {checkpoint_dir}") - - latest_step = max(checkpoint_steps) - ckpt_dir = checkpoint_dir / f"{latest_step}" +def evaluate_offline_reward(model, batches, device, *, seed: int) -> tuple[float, float]: + """Return deterministic imitation reward and loss; this is not task reward.""" + target = unwrap_model(model) + was_training = target.training + target.eval() + losses: list[float] = [] + fork_devices = [device] if device.type == "cuda" else [] + try: + with torch.no_grad(), torch.random.fork_rng(devices=fork_devices): + torch.manual_seed(seed) + if device.type == "cuda": + torch.cuda.manual_seed_all(seed) + for cpu_observation, cpu_actions in batches: + observation = jax.tree.map(lambda value: value.to(device), cpu_observation) + actions = cpu_actions.to(device=device, dtype=torch.float32) + batch_losses = target(observation, actions) + if isinstance(batch_losses, list | tuple): + batch_losses = torch.stack(batch_losses) + losses.append(float(batch_losses.float().mean().item())) + finally: + target.train(was_training) + mean_loss = sum(losses) / len(losses) + return math.exp(-mean_loss), mean_loss + + +def write_checkpoint_contents( + checkpoint_dir, + *, + model, + optimizer, + global_step, + reward, + reward_loss, + config, + data_config, + checkpoint_evaluation: CheckpointEvaluationPlan, +): + target = unwrap_model(model) + if config.pytorch_lora_rank is not None: + state = _lora.trainable_state_dict(target) + safetensors.torch.save_file(state, checkpoint_dir / "model.safetensors") + checkpoint_kind = "pytorch_lora_adapter" + else: + safetensors.torch.save_model(target, checkpoint_dir / "model.safetensors") + checkpoint_kind = "pytorch_full_model" + + torch.save(optimizer.state_dict(), checkpoint_dir / "optimizer.pt") + metadata = { + "global_step": global_step, + "config": dataclasses.asdict(config), + "timestamp": time.time(), + "checkpoint_kind": checkpoint_kind, + "checkpoint_metric": checkpoint_evaluation.metric_name, + "checkpoint_eval_repo_id": checkpoint_evaluation.source_repo_id, + "checkpoint_eval_uses_heldout_data": checkpoint_evaluation.uses_heldout_data, + "checkpoint_reward": reward, + "checkpoint_eval_loss": reward_loss, + "offline_imitation_loss": reward_loss, + } + torch.save(metadata, checkpoint_dir / "metadata.pt") + + norm_stats = data_config.norm_stats + if norm_stats is not None and data_config.asset_id is not None: + _normalize.save(checkpoint_dir / "assets" / data_config.asset_id, norm_stats) + + +def load_checkpoint(model, optimizer, checkpoint_store, device): + """Resume from the highest-reward checkpoint.""" + ckpt_dir = checkpoint_store.best_checkpoint_dir() + if ckpt_dir is None: + raise FileNotFoundError(f"No reward checkpoints found in {checkpoint_store.root}") + latest_step = int(ckpt_dir.name) - # Clear memory before loading checkpoints if torch.cuda.is_available(): torch.cuda.empty_cache() gc.collect() log_memory_usage(device, latest_step, "before_loading_checkpoint") try: - # Load model state with error handling - logging.info("Loading model state...") + logging.info("Loading highest-reward model state from %s", ckpt_dir) safetensors_path = ckpt_dir / "model.safetensors" - if safetensors_path.exists(): - model_to_load = model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model - safetensors.torch.load_model(model_to_load, safetensors_path, device=str(device)) - logging.info("Loaded model state from safetensors format") + target = unwrap_model(model) + state = safetensors.torch.load_file(safetensors_path, device=str(device)) + if any(name.endswith(("lora_a", "lora_b")) for name in state): + _lora.load_trainable_state_dict(target, state) + else: + safetensors.torch.load_model(target, safetensors_path, device=str(device)) + logging.info("Loaded reward-selected model state") else: raise FileNotFoundError(f"No model checkpoint found at {ckpt_dir}") @@ -255,7 +278,7 @@ def load_checkpoint(model, optimizer, checkpoint_dir, device): gc.collect() log_memory_usage(device, latest_step, "after_loading_metadata") - logging.info(f"Successfully loaded all checkpoint components from step {latest_step}") + logging.info(f"Successfully loaded reward checkpoint components from step {latest_step}") return global_step except RuntimeError as e: @@ -269,18 +292,6 @@ def load_checkpoint(model, optimizer, checkpoint_dir, device): "Out of memory while loading checkpoint. Try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True" ) from e raise - - -def get_latest_checkpoint_step(checkpoint_dir): - """Get the latest checkpoint step number from a checkpoint directory.""" - checkpoint_steps = [ - int(d.name) - for d in checkpoint_dir.iterdir() - if d.is_dir() and d.name.isdigit() and not d.name.startswith("tmp_") - ] - return max(checkpoint_steps) if checkpoint_steps else None - - def log_memory_usage(device, step, phase="unknown"): """Log detailed memory usage information.""" if not torch.cuda.is_available(): @@ -310,24 +321,15 @@ def train_loop(config: _config.TrainConfig): use_ddp, local_rank, device = setup_ddp() is_main = (not use_ddp) or (dist.get_rank() == 0) set_seed(config.seed, local_rank) + checkpoint_evaluation = resolve_checkpoint_evaluation(config) - # Initialize checkpoint directory and wandb + # Initialize checkpoint directory and wandb. resuming = False if config.resume: - # Find checkpoint directory based on experiment name exp_checkpoint_dir = config.checkpoint_dir - if exp_checkpoint_dir.exists(): - # Use validation to find the latest working checkpoint - latest_step = get_latest_checkpoint_step(exp_checkpoint_dir) - if latest_step is not None: - resuming = True - logging.info( - f"Resuming from experiment checkpoint directory: {exp_checkpoint_dir} at step {latest_step}" - ) - else: - raise FileNotFoundError(f"No valid checkpoints found in {exp_checkpoint_dir} for resume") - else: + if not exp_checkpoint_dir.exists(): raise FileNotFoundError(f"Experiment checkpoint directory {exp_checkpoint_dir} does not exist for resume") + resuming = True elif config.overwrite and config.checkpoint_dir.exists(): shutil.rmtree(config.checkpoint_dir) logging.info(f"Overwriting checkpoint directory: {config.checkpoint_dir}") @@ -339,9 +341,26 @@ def train_loop(config: _config.TrainConfig): exp_checkpoint_dir.mkdir(parents=True, exist_ok=True) logging.info(f"Created experiment checkpoint directory: {exp_checkpoint_dir}") else: - # For resume, checkpoint_dir is already set to the experiment directory logging.info(f"Using existing experiment checkpoint directory: {config.checkpoint_dir}") + checkpoint_store = RewardCheckpointStore( + config.checkpoint_dir, + metric_name=checkpoint_evaluation.metric_name, + minimum_reward=config.checkpoint_reward_min, + minimum_delta=config.checkpoint_reward_min_delta, + minimum_free_bytes=int(config.checkpoint_min_free_disk_gib * 1024**3), + minimum_free_fraction=config.checkpoint_min_free_disk_fraction, + ) + if resuming: + best = checkpoint_store.best_record() + if best is None: + raise FileNotFoundError(f"No reward checkpoint found in {config.checkpoint_dir}") + logging.info( + "Resuming from best reward checkpoint: step=%d reward=%.8f", + int(best["step"]), + float(best["reward"]), + ) + # Initialize wandb (only on main process) if is_main: init_wandb(config, resuming=resuming, enabled=config.wandb_enabled) @@ -357,6 +376,13 @@ def train_loop(config: _config.TrainConfig): # Pass the original batch size to data loader - it will handle DDP splitting internally loader, data_config = build_datasets(config) + reward_batches = build_reward_batches(checkpoint_evaluation.loader_config) + logging.info( + "Checkpoint evaluation dataset: repo_id=%s heldout=%s metric=%s", + checkpoint_evaluation.source_repo_id, + checkpoint_evaluation.uses_heldout_data, + checkpoint_evaluation.metric_name, + ) # Log sample images to wandb on first batch if is_main and config.wandb_enabled and not resuming: @@ -408,6 +434,37 @@ def train_loop(config: _config.TrainConfig): model = openpi.models_pytorch.pi0_pytorch.PI0Pytorch(model_cfg).to(device) + # Always load the immutable base weights before inserting LoRA wrappers. + if config.pytorch_weight_path is not None: + logging.info(f"Loading base weights from: {config.pytorch_weight_path}") + model_path = os.path.join(config.pytorch_weight_path, "model.safetensors") + safetensors.torch.load_model(model, model_path) + logging.info(f"Loaded PyTorch base weights from {config.pytorch_weight_path}") + + lora_modules: list[str] = [] + if config.pytorch_lora_rank is not None: + action_expert_rank = config.pytorch_lora_action_expert_rank or config.pytorch_lora_rank + lora_modules = _lora.apply_pi0_lora( + model, + paligemma_rank=config.pytorch_lora_rank, + action_expert_rank=action_expert_rank, + paligemma_alpha=config.pytorch_lora_alpha, + action_expert_alpha=config.pytorch_lora_action_expert_alpha, + dropout=config.pytorch_lora_dropout, + train_action_heads=config.pytorch_lora_train_action_heads, + ) + logging.info("Inserted LoRA into %d transformer linear layers", len(lora_modules)) + + total_parameters, trainable_parameter_count = _lora.parameter_counts(model) + if trainable_parameter_count == 0: + raise RuntimeError("Training configuration froze every model parameter") + logging.info( + "Parameters: total=%d trainable=%d trainable_fraction=%.6f", + total_parameters, + trainable_parameter_count, + trainable_parameter_count / total_parameters, + ) + if hasattr(model, "gradient_checkpointing_enable"): enable_gradient_checkpointing = True model.gradient_checkpointing_enable() @@ -433,21 +490,11 @@ def train_loop(config: _config.TrainConfig): model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[device.index] if device.type == "cuda" else None, - find_unused_parameters=True, # Disable for memory efficiency - gradient_as_bucket_view=True, # Enable for memory efficiency + find_unused_parameters=False, + gradient_as_bucket_view=True, static_graph=world_size >= 8, # Enable for 8+ GPUs ) - # Load weights from weight_loader if specified (for fine-tuning) - if config.pytorch_weight_path is not None: - logging.info(f"Loading weights from: {config.pytorch_weight_path}") - - model_path = os.path.join(config.pytorch_weight_path, "model.safetensors") - safetensors.torch.load_model( - (model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model), model_path - ) - logging.info(f"Loaded PyTorch weights from {config.pytorch_weight_path}") - # Optimizer + learning rate schedule from config warmup_steps = config.lr_schedule.warmup_steps peak_lr = config.lr_schedule.peak_lr @@ -455,8 +502,9 @@ def train_loop(config: _config.TrainConfig): end_lr = config.lr_schedule.decay_lr # Create optimizer with config parameters + optimized_parameters = trainable_parameters(model) optim = torch.optim.AdamW( - model.parameters(), + optimized_parameters, lr=peak_lr, betas=(config.optimizer.b1, config.optimizer.b2), eps=config.optimizer.eps, @@ -466,7 +514,7 @@ def train_loop(config: _config.TrainConfig): # Load checkpoint if resuming global_step = 0 if resuming: - global_step = load_checkpoint(model, optim, config.checkpoint_dir, device) + global_step = load_checkpoint(model, optim, checkpoint_store, device) logging.info(f"Resumed training from step {global_step}") def lr_schedule(step: int): @@ -480,6 +528,72 @@ def lr_schedule(step: int): return end_lr + (peak_lr - end_lr) * cos model.train() + training_started_at = time.time() + deadline_epoch = ( + training_started_at + config.max_train_seconds + if config.max_train_seconds is not None + else None + ) + tensorboard_dir = config.checkpoint_dir / "tensorboard" + tensorboard = SummaryWriter(log_dir=str(tensorboard_dir), purge_step=global_step) if is_main else None + state_path = config.checkpoint_dir / "run_state.json" + last_reward = None + last_reward_loss = None + last_reward_eval_step = None + + def write_run_state(state: str, **extra): + if not is_main: + return + payload = { + "state": state, + "pid": os.getpid(), + "global_step": global_step, + "training_started_epoch": training_started_at, + "deadline_epoch": deadline_epoch, + "tensorboard_dir": str(tensorboard_dir), + "checkpoint_metric": checkpoint_evaluation.metric_name, + "checkpoint_metric_is_task_reward": False, + "checkpoint_eval_repo_id": checkpoint_evaluation.source_repo_id, + "checkpoint_eval_uses_heldout_data": checkpoint_evaluation.uses_heldout_data, + "baseline_checkpoint_reward": baseline_reward, + "baseline_checkpoint_loss": baseline_reward_loss, + "last_offline_imitation_reward": last_reward, + "last_offline_imitation_loss": last_reward_loss, + "last_reward_eval_step": last_reward_eval_step, + "best_reward_checkpoint": checkpoint_store.best_record(), + "total_parameters": total_parameters, + "trainable_parameters": trainable_parameter_count, + **extra, + } + temporary = state_path.with_suffix(".tmp") + temporary.write_text(json.dumps(payload, sort_keys=True) + "\n", encoding="utf-8") + temporary.replace(state_path) + + baseline_reward, baseline_reward_loss = evaluate_offline_reward( + model, + reward_batches, + device, + seed=config.seed + 10_000, + ) + last_reward = baseline_reward + last_reward_loss = baseline_reward_loss + last_reward_eval_step = global_step + if checkpoint_store.best_record() is None: + checkpoint_store.minimum_reward = max( + checkpoint_store.minimum_reward, + baseline_reward + config.checkpoint_reward_min_delta, + ) + if is_main: + logging.info( + "Offline checkpoint baseline: reward=%.8f loss=%.8f (not simulator task reward)", + baseline_reward, + baseline_reward_loss, + ) + if tensorboard is not None: + tensorboard.add_scalar(checkpoint_evaluation.reward_log_key, baseline_reward, global_step) + tensorboard.add_scalar(checkpoint_evaluation.loss_log_key, baseline_reward_loss, global_step) + tensorboard.flush() + write_run_state("TRAINING_RUNNING") start_time = time.time() infos = [] # Collect stats over log interval if is_main: @@ -498,6 +612,68 @@ def lr_schedule(step: int): ) logging.info("EMA is not supported for PyTorch training") logging.info(f"Training precision: {model_cfg.dtype}") + logging.info( + "LoRA: modules=%d trainable=%d/%d (%.4f%%)", + len(lora_modules), + trainable_parameter_count, + total_parameters, + 100.0 * trainable_parameter_count / total_parameters, + ) + + def evaluate_and_maybe_checkpoint(): + nonlocal last_reward, last_reward_loss, last_reward_eval_step + reward, reward_loss = evaluate_offline_reward( + model, + reward_batches, + device, + seed=config.seed + 10_000, + ) + last_reward = reward + last_reward_loss = reward_loss + last_reward_eval_step = global_step + if not is_main: + return None + result = checkpoint_store.maybe_save( + step=global_step, + reward=reward, + write_checkpoint=lambda checkpoint_dir: write_checkpoint_contents( + checkpoint_dir, + model=model, + optimizer=optim, + global_step=global_step, + reward=reward, + reward_loss=reward_loss, + config=config, + data_config=data_config, + checkpoint_evaluation=checkpoint_evaluation, + ), + ) + logging.info( + "Offline checkpoint evaluation: step=%d reward=%.8f loss=%.8f saved=%s reason=%s", + global_step, + reward, + reward_loss, + result.saved, + result.reason, + ) + if tensorboard is not None: + tensorboard.add_scalar(checkpoint_evaluation.reward_log_key, reward, global_step) + tensorboard.add_scalar(checkpoint_evaluation.loss_log_key, reward_loss, global_step) + tensorboard.add_scalar("checkpoint/saved", float(result.saved), global_step) + if result.best_reward is not None: + tensorboard.add_scalar("checkpoint/best_reward", result.best_reward, global_step) + tensorboard.flush() + if config.wandb_enabled: + wandb.log( + { + checkpoint_evaluation.metric_name: reward, + checkpoint_evaluation.loss_log_key.removeprefix("eval/"): reward_loss, + "checkpoint_saved": int(result.saved), + }, + step=global_step, + ) + write_run_state("TRAINING_RUNNING") + return result # Training loop - iterate until we reach num_train_steps pbar = ( @@ -506,7 +682,8 @@ def lr_schedule(step: int): else None ) - while global_step < config.num_train_steps: + deadline_reached = False + while global_step < config.num_train_steps and not deadline_reached: # Set epoch for distributed training if use_ddp and hasattr(loader, "set_epoch"): loader.set_epoch(global_step // len(loader)) @@ -515,6 +692,9 @@ def lr_schedule(step: int): # Check if we've reached the target number of steps if global_step >= config.num_train_steps: break + if deadline_epoch is not None and time.time() >= deadline_epoch: + deadline_reached = True + break # The unified data loader returns (observation, actions) tuple observation = jax.tree.map(lambda x: x.to(device), observation) # noqa: PLW2901 @@ -543,7 +723,9 @@ def lr_schedule(step: int): log_memory_usage(device, global_step, "after_backward") # Gradient clipping - grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.optimizer.clip_gradient_norm) + grad_norm = torch.nn.utils.clip_grad_norm_( + optimized_parameters, max_norm=config.optimizer.clip_gradient_norm + ) # Optimizer step optim.step() @@ -597,12 +779,27 @@ def lr_schedule(step: int): log_payload["grad_norm"] = avg_grad_norm wandb.log(log_payload, step=global_step) + if tensorboard is not None: + tensorboard.add_scalar("train/loss", avg_loss, global_step) + tensorboard.add_scalar("train/learning_rate", avg_lr, global_step) + tensorboard.add_scalar("train/time_per_step_s", elapsed / max(1, len(infos)), global_step) + if avg_grad_norm is not None: + tensorboard.add_scalar("train/grad_norm", avg_grad_norm, global_step) + tensorboard.flush() + write_run_state( + "TRAINING_RUNNING", + loss=avg_loss, + learning_rate=avg_lr, + grad_norm=avg_grad_norm, + uptime_seconds=time.time() - training_started_at, + ) + start_time = time.time() infos = [] # Reset stats collection global_step += 1 - # Save checkpoint using the new mechanism - save_checkpoint(model, optim, global_step, config, is_main, data_config) + if global_step % config.save_interval == 0: + evaluate_and_maybe_checkpoint() # Update progress bar if pbar is not None: @@ -611,6 +808,13 @@ def lr_schedule(step: int): {"loss": f"{loss.item():.4f}", "lr": f"{optim.param_groups[0]['lr']:.2e}", "step": global_step} ) + if global_step > 0 and last_reward_eval_step != global_step: + evaluate_and_maybe_checkpoint() + write_run_state( + "TRAINING_DEADLINE_COMPLETE" if deadline_reached else "TRAINING_STEPS_COMPLETE", + uptime_seconds=time.time() - training_started_at, + ) + # Close progress bar if pbar is not None: pbar.close() @@ -618,6 +822,8 @@ def lr_schedule(step: int): # Finish wandb run if is_main and config.wandb_enabled: wandb.finish() + if tensorboard is not None: + tensorboard.close() cleanup_ddp() diff --git a/src/openpi/models/model.py b/src/openpi/models/model.py index 29618b4945..b24d687f58 100644 --- a/src/openpi/models/model.py +++ b/src/openpi/models/model.py @@ -17,6 +17,7 @@ import safetensors import torch +from openpi.models_pytorch import lora as _lora from openpi.models_pytorch import pi0_pytorch from openpi.shared import image_tools import openpi.shared.array_typing as at @@ -243,7 +244,25 @@ def load(self, params: at.Params, *, remove_extra_params: bool = True) -> "BaseM def load_pytorch(self, train_config, weight_path: str): logger.info(f"train_config: {train_config}") model = pi0_pytorch.PI0Pytorch(config=train_config.model) - safetensors.torch.load_model(model, weight_path) + if train_config.pytorch_lora_rank is None: + safetensors.torch.load_model(model, weight_path) + return model + + if train_config.pytorch_weight_path is None: + raise ValueError("A PyTorch base-weight path is required to load a LoRA checkpoint") + base_weight_path = pathlib.Path(train_config.pytorch_weight_path) / "model.safetensors" + safetensors.torch.load_model(model, base_weight_path) + _lora.apply_pi0_lora( + model, + paligemma_rank=train_config.pytorch_lora_rank, + action_expert_rank=train_config.pytorch_lora_action_expert_rank or train_config.pytorch_lora_rank, + paligemma_alpha=train_config.pytorch_lora_alpha, + action_expert_alpha=train_config.pytorch_lora_action_expert_alpha, + dropout=train_config.pytorch_lora_dropout, + train_action_heads=train_config.pytorch_lora_train_action_heads, + ) + adapter_state = safetensors.torch.load_file(weight_path, device="cpu") + _lora.load_trainable_state_dict(model, adapter_state) return model @abc.abstractmethod diff --git a/src/openpi/models_pytorch/lora.py b/src/openpi/models_pytorch/lora.py new file mode 100644 index 0000000000..a7803e3f15 --- /dev/null +++ b/src/openpi/models_pytorch/lora.py @@ -0,0 +1,167 @@ +"""Small, dependency-free LoRA utilities for the PyTorch pi0/pi0.5 model.""" + +from __future__ import annotations + +from collections.abc import Iterable +import math + +import torch +from torch import nn +import torch.nn.functional as F # noqa: N812 + + +class LoRALinear(nn.Module): + """A frozen linear layer with a trainable low-rank residual.""" + + def __init__(self, base: nn.Linear, *, rank: int, alpha: float, dropout: float = 0.0) -> None: + super().__init__() + if rank <= 0: + raise ValueError("LoRA rank must be positive") + if alpha <= 0: + raise ValueError("LoRA alpha must be positive") + if not 0.0 <= dropout < 1.0: + raise ValueError("LoRA dropout must be in [0, 1)") + + self.base = base + self.rank = rank + self.alpha = alpha + self.scaling = alpha / rank + self.dropout = nn.Dropout(dropout) if dropout else nn.Identity() + + for parameter in self.base.parameters(): + parameter.requires_grad_(requires_grad=False) + + factory_kwargs = {"device": base.weight.device, "dtype": base.weight.dtype} + self.lora_a = nn.Parameter(torch.empty(rank, base.in_features, **factory_kwargs)) + self.lora_b = nn.Parameter(torch.zeros(base.out_features, rank, **factory_kwargs)) + nn.init.kaiming_uniform_(self.lora_a, a=math.sqrt(5)) + + @property + def weight(self) -> torch.Tensor: + """Expose the base weight for existing dtype/device checks.""" + return self.base.weight + + @property + def bias(self) -> torch.Tensor | None: + return self.base.bias + + @property + def in_features(self) -> int: + return self.base.in_features + + @property + def out_features(self) -> int: + return self.base.out_features + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + base = self.base(inputs) + residual = F.linear(F.linear(self.dropout(inputs), self.lora_a), self.lora_b) + return base + residual * self.scaling + + +def _replace_linear_children( + module: nn.Module, + *, + rank: int, + alpha: float, + dropout: float, + target_names: frozenset[str], +) -> list[str]: + replaced: list[str] = [] + for name, child in list(module.named_children()): + if isinstance(child, LoRALinear): + continue + if isinstance(child, nn.Linear) and name in target_names: + setattr(module, name, LoRALinear(child, rank=rank, alpha=alpha, dropout=dropout)) + replaced.append(name) + continue + replaced.extend( + f"{name}.{nested_name}" + for nested_name in _replace_linear_children( + child, + rank=rank, + alpha=alpha, + dropout=dropout, + target_names=target_names, + ) + ) + return replaced + + +def apply_pi0_lora( + model: nn.Module, + *, + paligemma_rank: int, + action_expert_rank: int, + paligemma_alpha: float, + action_expert_alpha: float, + dropout: float = 0.0, + train_action_heads: bool = True, +) -> list[str]: + """Freeze pi0.5 and add LoRA to both language and action transformers.""" + for parameter in model.parameters(): + parameter.requires_grad_(requires_grad=False) + + target_names = frozenset({"q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"}) + roots = model.paligemma_with_expert + replaced = [ + f"paligemma.language_model.{name}" + for name in _replace_linear_children( + roots.paligemma.language_model, + rank=paligemma_rank, + alpha=paligemma_alpha, + dropout=dropout, + target_names=target_names, + ) + ] + replaced.extend( + f"gemma_expert.{name}" + for name in _replace_linear_children( + roots.gemma_expert, + rank=action_expert_rank, + alpha=action_expert_alpha, + dropout=dropout, + target_names=target_names, + ) + ) + if not replaced: + raise RuntimeError("No pi0.5 transformer linear layers matched the LoRA targets") + + if train_action_heads: + for head_name in ("action_in_proj", "action_out_proj", "time_mlp_in", "time_mlp_out"): + head = getattr(model, head_name, None) + if head is not None: + for parameter in head.parameters(): + parameter.requires_grad_(requires_grad=True) + + return replaced + + +def trainable_named_parameters(model: nn.Module) -> Iterable[tuple[str, nn.Parameter]]: + return ((name, parameter) for name, parameter in model.named_parameters() if parameter.requires_grad) + + +def trainable_state_dict(model: nn.Module) -> dict[str, torch.Tensor]: + """Return only adapter/action-head tensors, detached and safe to serialize.""" + return { + name: parameter.detach().cpu().contiguous() + for name, parameter in trainable_named_parameters(model) + } + + +def load_trainable_state_dict(model: nn.Module, state: dict[str, torch.Tensor]) -> None: + expected = {name for name, _ in trainable_named_parameters(model)} + received = set(state) + if missing := expected - received: + raise ValueError(f"Adapter checkpoint is missing trainable tensors: {sorted(missing)[:8]}") + if unexpected := received - expected: + raise ValueError(f"Adapter checkpoint has unexpected tensors: {sorted(unexpected)[:8]}") + incompatible = model.load_state_dict(state, strict=False) + if incompatible.unexpected_keys: + raise ValueError(f"Unexpected adapter keys: {incompatible.unexpected_keys[:8]}") + + +def parameter_counts(model: nn.Module) -> tuple[int, int]: + total = sum(parameter.numel() for parameter in model.parameters()) + trainable = sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad) + return total, trainable diff --git a/src/openpi/models_pytorch/lora_test.py b/src/openpi/models_pytorch/lora_test.py new file mode 100644 index 0000000000..d4de8f872a --- /dev/null +++ b/src/openpi/models_pytorch/lora_test.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import torch +from torch import nn + +from openpi.models_pytorch import lora + + +class _Transformer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.q_proj = nn.Linear(8, 8) + self.down_proj = nn.Linear(8, 8) + + +class _ToyPi0(nn.Module): + def __init__(self) -> None: + super().__init__() + self.paligemma_with_expert = nn.Module() + self.paligemma_with_expert.paligemma = nn.Module() + self.paligemma_with_expert.paligemma.language_model = _Transformer() + self.paligemma_with_expert.gemma_expert = _Transformer() + self.action_in_proj = nn.Linear(8, 8) + self.action_out_proj = nn.Linear(8, 8) + self.time_mlp_in = nn.Linear(8, 8) + self.time_mlp_out = nn.Linear(8, 8) + + +def test_lora_linear_starts_as_exact_base() -> None: + base = nn.Linear(8, 4) + inputs = torch.randn(3, 8) + expected = base(inputs) + + adapted = lora.LoRALinear(base, rank=2, alpha=2.0) + + torch.testing.assert_close(adapted(inputs), expected) + assert adapted.base.weight.requires_grad is False + assert adapted.lora_a.requires_grad is True + assert adapted.lora_b.requires_grad is True + + +def test_apply_pi0_lora_freezes_base_and_round_trips_trainable_state() -> None: + model = _ToyPi0() + replaced = lora.apply_pi0_lora( + model, + paligemma_rank=2, + action_expert_rank=4, + paligemma_alpha=2.0, + action_expert_alpha=4.0, + ) + + assert len(replaced) == 4 + assert model.paligemma_with_expert.paligemma.language_model.q_proj.base.weight.requires_grad is False + assert model.action_out_proj.weight.requires_grad is True + total, trainable = lora.parameter_counts(model) + assert 0 < trainable < total + + for _, parameter in lora.trainable_named_parameters(model): + parameter.data.uniform_(-0.1, 0.1) + state = lora.trainable_state_dict(model) + + restored = _ToyPi0() + lora.apply_pi0_lora( + restored, + paligemma_rank=2, + action_expert_rank=4, + paligemma_alpha=2.0, + action_expert_alpha=4.0, + ) + lora.load_trainable_state_dict(restored, state) + + restored_state = dict(lora.trainable_named_parameters(restored)) + for name, tensor in state.items(): + torch.testing.assert_close(restored_state[name], tensor) diff --git a/src/openpi/models_pytorch/preprocessing_pytorch.py b/src/openpi/models_pytorch/preprocessing_pytorch.py index 33c94a59b1..c323ce4dc9 100644 --- a/src/openpi/models_pytorch/preprocessing_pytorch.py +++ b/src/openpi/models_pytorch/preprocessing_pytorch.py @@ -141,9 +141,11 @@ def preprocess_observation_pytorch( # Back to [-1, 1] image = image * 2.0 - 1.0 - # Convert back to [B, C, H, W] format if it was originally channels-first - if is_channels_first: - image = image.permute(0, 3, 1, 2) # [B, H, W, C] -> [B, C, H, W] + # The Hugging Face SigLIP vision tower always consumes NCHW. Inputs may + # arrive in either layout (policy inputs are commonly NCHW, while the + # OpenPI fake/training dataset is NHWC), so normalize the output rather + # than restoring the input layout. + image = image.permute(0, 3, 1, 2) # [B, H, W, C] -> [B, C, H, W] out_images[key] = image diff --git a/src/openpi/models_pytorch/preprocessing_pytorch_test.py b/src/openpi/models_pytorch/preprocessing_pytorch_test.py new file mode 100644 index 0000000000..ee39dc3672 --- /dev/null +++ b/src/openpi/models_pytorch/preprocessing_pytorch_test.py @@ -0,0 +1,24 @@ +from types import SimpleNamespace + +import pytest +import torch + +from openpi.models_pytorch import preprocessing_pytorch + + +@pytest.mark.parametrize("image_shape", [(1, 224, 224, 3), (1, 3, 224, 224)]) +def test_preprocess_normalizes_images_to_nchw(image_shape): + image_keys = preprocessing_pytorch.IMAGE_KEYS + observation = SimpleNamespace( + images={key: torch.zeros(image_shape) for key in image_keys}, + image_masks={key: torch.ones((1,), dtype=torch.bool) for key in image_keys}, + state=torch.zeros((1, 32)), + tokenized_prompt=torch.ones((1, 8), dtype=torch.int64), + tokenized_prompt_mask=torch.ones((1, 8), dtype=torch.bool), + token_ar_mask=None, + token_loss_mask=None, + ) + + processed = preprocessing_pytorch.preprocess_observation_pytorch(observation) + + assert all(image.shape == (1, 3, 224, 224) for image in processed.images.values()) diff --git a/src/openpi/policies/g1_policy.py b/src/openpi/policies/g1_policy.py new file mode 100644 index 0000000000..d412abbbaf --- /dev/null +++ b/src/openpi/policies/g1_policy.py @@ -0,0 +1,175 @@ +import dataclasses + +import einops +import numpy as np + +from openpi import transforms +from openpi.models import model as _model + + +def _parse_image(image) -> np.ndarray: + image = np.asarray(image) + if np.issubdtype(image.dtype, np.floating): + image = np.clip(image * 255.0, 0, 255).astype(np.uint8) + if image.ndim != 3: + raise ValueError(f"Expected one HWC or CHW head image, got shape {image.shape}") + if image.shape[0] == 3: + image = einops.rearrange(image, "c h w -> h w c") + if image.shape[-1] != 3: + raise ValueError(f"Expected an RGB head image, got shape {image.shape}") + return image + + +def _translate_pair(rgb: np.ndarray, depth: np.ndarray, dx: int, dy: int) -> tuple[np.ndarray, np.ndarray]: + """Translate an aligned RGB-D pair with edge padding.""" + height, width = rgb.shape[:2] + pad_x, pad_y = abs(dx), abs(dy) + + def translate(image: np.ndarray) -> np.ndarray: + padded = np.pad(image, ((pad_y, pad_y), (pad_x, pad_x), (0, 0)), mode="edge") + start_x = pad_x - dx + start_y = pad_y - dy + return padded[start_y : start_y + height, start_x : start_x + width].copy() + + return translate(rgb), translate(depth) + + +@dataclasses.dataclass(frozen=True) +class G1RgbdAugment(transforms.DataTransformFn): + """Apply bounded real-camera perturbations to a parsed RGB-D pair.""" + + max_translation_px: int = 6 + max_brightness_delta: float = 12.0 + max_contrast_delta: float = 0.10 + max_channel_gain_delta: float = 0.05 + max_depth_noise_std: float = 3.0 + occlusion_probability: float = 0.25 + max_occlusion_fraction: float = 0.12 + + def __call__(self, data: dict) -> dict: + images = data.get("image") + masks = data.get("image_mask", {}) + if not isinstance(images, dict) or "base_0_rgb" not in images: + raise ValueError("G1 RGB-D augmentation requires parsed policy images") + if not bool(masks.get("left_wrist_0_rgb", False)): + raise ValueError("G1 RGB-D augmentation requires valid aligned depth") + + rgb = _parse_image(images["base_0_rgb"]) + depth = _parse_image(images["left_wrist_0_rgb"]) + if rgb.shape != depth.shape: + raise ValueError("G1 RGB-D augmentation requires matching RGB and depth shapes") + + dx = int(np.random.randint(-self.max_translation_px, self.max_translation_px + 1)) + dy = int(np.random.randint(-self.max_translation_px, self.max_translation_px + 1)) + rgb, depth = _translate_pair(rgb, depth, dx, dy) + + rgb_float = rgb.astype(np.float32) + contrast = 1.0 + float(np.random.uniform(-self.max_contrast_delta, self.max_contrast_delta)) + brightness = float(np.random.uniform(-self.max_brightness_delta, self.max_brightness_delta)) + channel_gain = np.random.uniform( + 1.0 - self.max_channel_gain_delta, + 1.0 + self.max_channel_gain_delta, + size=(1, 1, 3), + ).astype(np.float32) + rgb = np.clip( + (rgb_float - 127.5) * contrast * channel_gain + 127.5 + brightness, + 0, + 255, + ).astype(np.uint8) + + depth_float = depth[..., 0].astype(np.float32) + depth_noise_std = float(np.random.uniform(0.0, self.max_depth_noise_std)) + if depth_noise_std > 0.0: + depth_float += np.random.normal(0.0, depth_noise_std, size=depth_float.shape).astype(np.float32) + depth_plane = np.clip(depth_float, 0, 255).astype(np.uint8) + depth = np.repeat(depth_plane[..., None], 3, axis=-1) + + if float(np.random.random()) < self.occlusion_probability: + height, width = rgb.shape[:2] + min_fraction = min(0.04, self.max_occlusion_fraction) + fraction = float(np.random.uniform(min_fraction, self.max_occlusion_fraction)) + cutout_height = max(1, int(round(height * fraction))) + cutout_width = max(1, int(round(width * fraction))) + top = int(np.random.randint(0, height - cutout_height + 1)) + left = int(np.random.randint(0, width - cutout_width + 1)) + rgb_fill = np.median(rgb.reshape(-1, 3), axis=0).astype(np.uint8) + rgb[top : top + cutout_height, left : left + cutout_width] = rgb_fill + depth[top : top + cutout_height, left : left + cutout_width] = 0 + + images["base_0_rgb"] = rgb + images["left_wrist_0_rgb"] = depth + return data + + +@dataclasses.dataclass(frozen=True) +class G1Inputs(transforms.DataTransformFn): + """Map G1 RGB-D and low-dimensional state into pi0.5 image slots. + + The optional aligned depth image is an 8-bit, three-channel inverse-depth + visualization produced by the dataset/runtime adapter. It occupies a + pretrained auxiliary image slot so the pi0.5 checkpoint shape does not + change. Missing depth is accepted only for RGB-only configurations. + """ + + model_type: _model.ModelType + state_dim: int = 29 + task_action_dim: int = 21 + use_depth_image: bool = False + + def __call__(self, data: dict) -> dict: + state = np.asarray(data["state"], dtype=np.float32) + head_image = _parse_image(data["head_image"]) + depth_image = None + if self.use_depth_image: + if "depth_image" not in data: + raise ValueError("Expected an aligned G1 depth image") + depth_image = _parse_image(data["depth_image"]) + if depth_image.shape != head_image.shape: + raise ValueError( + "G1 RGB and aligned depth images must share one shape, got " + f"{head_image.shape} and {depth_image.shape}" + ) + + if state.shape != (self.state_dim,): + raise ValueError(f"Expected {self.state_dim} G1 state values, got shape {state.shape}") + if not np.isfinite(state).all(): + raise ValueError("G1 state contains non-finite values") + + if self.model_type not in (_model.ModelType.PI0, _model.ModelType.PI05): + raise ValueError(f"Unsupported G1 model type: {self.model_type}") + + inputs = { + "state": state, + "image": { + "base_0_rgb": head_image, + "left_wrist_0_rgb": depth_image if depth_image is not None else np.zeros_like(head_image), + "right_wrist_0_rgb": np.zeros_like(head_image), + }, + "image_mask": { + "base_0_rgb": np.True_, + "left_wrist_0_rgb": np.True_ if depth_image is not None else np.False_, + "right_wrist_0_rgb": np.False_, + }, + } + + if "actions" in data: + actions = np.asarray(data["actions"], dtype=np.float32) + if actions.shape[-1] != self.task_action_dim: + raise ValueError(f"Expected {self.task_action_dim} G1 task actions, got shape {actions.shape}") + if not np.isfinite(actions).all(): + raise ValueError("G1 actions contain non-finite values") + inputs["actions"] = actions + + if "prompt" in data: + prompt = data["prompt"] + inputs["prompt"] = prompt.decode("utf-8") if isinstance(prompt, bytes) else prompt + + return inputs + + +@dataclasses.dataclass(frozen=True) +class G1Outputs(transforms.DataTransformFn): + task_action_dim: int = 21 + + def __call__(self, data: dict) -> dict: + return {"actions": np.asarray(data["actions"])[..., : self.task_action_dim]} diff --git a/src/openpi/policies/g1_policy_test.py b/src/openpi/policies/g1_policy_test.py new file mode 100644 index 0000000000..bb24279d39 --- /dev/null +++ b/src/openpi/policies/g1_policy_test.py @@ -0,0 +1,184 @@ +import numpy as np +import pytest + +from openpi.models import model as _model +from openpi.policies import g1_policy + + +@pytest.mark.parametrize("image_shape", [(224, 224, 3), (3, 224, 224)]) +def test_g1_inputs_use_head_camera_and_mask_padding(image_shape): + transform = g1_policy.G1Inputs(model_type=_model.ModelType.PI05) + + result = transform( + { + "head_image": np.zeros(image_shape, dtype=np.uint8), + "state": np.zeros(29, dtype=np.float32), + "actions": np.zeros((10, 21), dtype=np.float32), + "prompt": "slice the fruit", + } + ) + + assert result["image"]["base_0_rgb"].shape == (224, 224, 3) + assert result["image_mask"] == { + "base_0_rgb": np.True_, + "left_wrist_0_rgb": np.False_, + "right_wrist_0_rgb": np.False_, + } + assert result["state"].shape == (29,) + assert result["actions"].shape == (10, 21) + + +def test_g1_inputs_use_aligned_depth_in_auxiliary_image_slot(): + transform = g1_policy.G1Inputs( + model_type=_model.ModelType.PI05, + state_dim=24, + task_action_dim=14, + use_depth_image=True, + ) + + depth = np.full((240, 320, 3), 73, dtype=np.uint8) + result = transform( + { + "head_image": np.zeros((240, 320, 3), dtype=np.uint8), + "depth_image": depth, + "state": np.zeros(24, dtype=np.float32), + "actions": np.zeros((10, 14), dtype=np.float32), + "prompt": "grasp the Coke can, lift it, and present it", + } + ) + + np.testing.assert_array_equal(result["image"]["left_wrist_0_rgb"], depth) + assert result["image_mask"] == { + "base_0_rgb": np.True_, + "left_wrist_0_rgb": np.True_, + "right_wrist_0_rgb": np.False_, + } + assert result["actions"].shape == (10, 14) + + +def test_g1_left_only_rgbd_contract_uses_17_state_and_7_actions(): + transform = g1_policy.G1Inputs( + model_type=_model.ModelType.PI05, + state_dim=17, + task_action_dim=7, + use_depth_image=True, + ) + + result = transform( + { + "head_image": np.zeros((240, 320, 3), dtype=np.uint8), + "depth_image": np.zeros((240, 320, 3), dtype=np.uint8), + "state": np.zeros(17, dtype=np.float32), + "actions": np.zeros((10, 7), dtype=np.float32), + "prompt": "grasp the Coke can, lift it, and present it", + } + ) + + assert result["state"].shape == (17,) + assert result["actions"].shape == (10, 7) + assert g1_policy.G1Outputs(task_action_dim=7)( + {"actions": np.zeros((10, 32), dtype=np.float32)} + )["actions"].shape == (10, 7) + + +def test_rgbd_augmentation_shares_spatial_translation() -> None: + feature = np.zeros((32, 48, 3), dtype=np.uint8) + feature[10:14, 20:24] = 255 + data = { + "image": {"base_0_rgb": feature.copy(), "left_wrist_0_rgb": feature.copy()}, + "image_mask": {"left_wrist_0_rgb": np.True_}, + } + transform = g1_policy.G1RgbdAugment( + max_translation_px=5, + max_brightness_delta=0.0, + max_contrast_delta=0.0, + max_channel_gain_delta=0.0, + max_depth_noise_std=0.0, + occlusion_probability=0.0, + ) + + np.random.seed(7) + result = transform(data) + + np.testing.assert_array_equal(result["image"]["base_0_rgb"], result["image"]["left_wrist_0_rgb"]) + assert result["image"]["base_0_rgb"].shape == feature.shape + assert not np.array_equal(result["image"]["base_0_rgb"], feature) + + +def test_rgbd_augmentation_is_seeded_and_preserves_depth_channels() -> None: + rgb = np.full((24, 32, 3), 120, dtype=np.uint8) + depth = np.full((24, 32, 3), 80, dtype=np.uint8) + transform = g1_policy.G1RgbdAugment(occlusion_probability=1.0) + + def apply_once() -> dict: + return transform( + { + "image": {"base_0_rgb": rgb.copy(), "left_wrist_0_rgb": depth.copy()}, + "image_mask": {"left_wrist_0_rgb": np.True_}, + } + ) + + np.random.seed(19) + first = apply_once() + np.random.seed(19) + second = apply_once() + + np.testing.assert_array_equal(first["image"]["base_0_rgb"], second["image"]["base_0_rgb"]) + np.testing.assert_array_equal(first["image"]["left_wrist_0_rgb"], second["image"]["left_wrist_0_rgb"]) + augmented_depth = first["image"]["left_wrist_0_rgb"] + np.testing.assert_array_equal(augmented_depth[..., 0], augmented_depth[..., 1]) + np.testing.assert_array_equal(augmented_depth[..., 1], augmented_depth[..., 2]) + + +def test_g1_rgbd_contract_requires_depth(): + transform = g1_policy.G1Inputs( + model_type=_model.ModelType.PI05, + state_dim=24, + task_action_dim=14, + use_depth_image=True, + ) + + with pytest.raises(ValueError, match="aligned G1 depth"): + transform( + { + "head_image": np.zeros((240, 320, 3), dtype=np.uint8), + "state": np.zeros(24, dtype=np.float32), + } + ) + + +def test_g1_outputs_drop_pi_padding(): + outputs = g1_policy.G1Outputs(task_action_dim=21) + + result = outputs({"actions": np.zeros((10, 32), dtype=np.float32)}) + + assert result["actions"].shape == (10, 21) + + +def test_g1_coke_contract_uses_24_state_and_21_actions(): + transform = g1_policy.G1Inputs(model_type=_model.ModelType.PI05, state_dim=24, task_action_dim=21) + + result = transform( + { + "head_image": np.zeros((240, 320, 3), dtype=np.uint8), + "state": np.zeros(24, dtype=np.float32), + "actions": np.zeros((10, 21), dtype=np.float32), + "prompt": "pick up the Coke can and hold it upright", + } + ) + + assert result["state"].shape == (24,) + assert result["actions"].shape == (10, 21) + assert g1_policy.G1Outputs(task_action_dim=21)({"actions": np.zeros((10, 32))})["actions"].shape == (10, 21) + + +def test_g1_contract_rejects_wrong_state_width(): + transform = g1_policy.G1Inputs(model_type=_model.ModelType.PI05, state_dim=24, task_action_dim=21) + + with pytest.raises(ValueError, match="Expected 24 G1 state"): + transform( + { + "head_image": np.zeros((240, 320, 3), dtype=np.uint8), + "state": np.zeros(29, dtype=np.float32), + } + ) diff --git a/src/openpi/training/checkpoint_evaluation.py b/src/openpi/training/checkpoint_evaluation.py new file mode 100644 index 0000000000..0a827a5c06 --- /dev/null +++ b/src/openpi/training/checkpoint_evaluation.py @@ -0,0 +1,62 @@ +"""Resolve the dataset and metric used for reward-gated checkpoints.""" + +from __future__ import annotations + +import dataclasses +from typing import Any + + +@dataclasses.dataclass(frozen=True) +class CheckpointEvaluationPlan: + """Immutable checkpoint-evaluation provenance and loader configuration.""" + + loader_config: Any + source_repo_id: str + metric_name: str + reward_log_key: str + loss_log_key: str + uses_heldout_data: bool + + +def resolve_checkpoint_evaluation(config: Any) -> CheckpointEvaluationPlan: + """Build a loader config for checkpoint scoring without mutating training data. + + When a held-out repo is configured, it reuses the training data factory and + training normalization assets, changing only the LeRobot repo selected by + the evaluation loader. This prevents checkpoint selection from fitting + normalization statistics on held-out examples. + """ + + training_repo_id = config.data.repo_id + if not isinstance(training_repo_id, str) or not training_repo_id: + raise ValueError("Checkpoint evaluation requires a concrete training repo id") + + heldout_repo_id = config.checkpoint_eval_repo_id + if heldout_repo_id is None: + return CheckpointEvaluationPlan( + loader_config=config, + source_repo_id=training_repo_id, + metric_name="offline_imitation_reward", + reward_log_key="eval/offline_imitation_reward", + loss_log_key="eval/offline_imitation_loss", + uses_heldout_data=False, + ) + if not isinstance(heldout_repo_id, str) or not heldout_repo_id: + raise ValueError("checkpoint_eval_repo_id must be a non-empty string when set") + if heldout_repo_id == training_repo_id: + raise ValueError("Held-out checkpoint evaluation repo must differ from the training repo") + + evaluation_assets = dataclasses.replace(config.data.assets, asset_id=training_repo_id) + evaluation_data = dataclasses.replace( + config.data, + repo_id=heldout_repo_id, + assets=evaluation_assets, + ) + return CheckpointEvaluationPlan( + loader_config=dataclasses.replace(config, data=evaluation_data), + source_repo_id=heldout_repo_id, + metric_name="heldout_offline_imitation_reward", + reward_log_key="eval/heldout_offline_imitation_reward", + loss_log_key="eval/heldout_offline_imitation_loss", + uses_heldout_data=True, + ) diff --git a/src/openpi/training/checkpoint_evaluation_test.py b/src/openpi/training/checkpoint_evaluation_test.py new file mode 100644 index 0000000000..c7960c5e35 --- /dev/null +++ b/src/openpi/training/checkpoint_evaluation_test.py @@ -0,0 +1,67 @@ +import dataclasses + +import pytest + +from openpi.training.checkpoint_evaluation import resolve_checkpoint_evaluation + + +@dataclasses.dataclass(frozen=True) +class _Assets: + assets_dir: str | None = None + asset_id: str | None = None + + +@dataclasses.dataclass(frozen=True) +class _Data: + repo_id: str + assets: _Assets = dataclasses.field(default_factory=_Assets) + use_depth_image: bool = True + + +@dataclasses.dataclass(frozen=True) +class _Config: + data: _Data + checkpoint_eval_repo_id: str | None = None + + +def test_default_checkpoint_evaluation_preserves_legacy_training_batches() -> None: + config = _Config(data=_Data(repo_id="local/train")) + + plan = resolve_checkpoint_evaluation(config) + + assert plan.loader_config is config + assert plan.source_repo_id == "local/train" + assert plan.metric_name == "offline_imitation_reward" + assert not plan.uses_heldout_data + + +def test_heldout_checkpoint_evaluation_changes_only_repo_and_reuses_training_norm_assets() -> None: + config = _Config( + data=_Data(repo_id="local/train", assets=_Assets(assets_dir="/assets")), + checkpoint_eval_repo_id="local/eval", + ) + + plan = resolve_checkpoint_evaluation(config) + + assert plan.loader_config is not config + assert plan.loader_config.data.repo_id == "local/eval" + assert plan.loader_config.data.assets == _Assets(assets_dir="/assets", asset_id="local/train") + assert plan.loader_config.data.use_depth_image + assert config.data.repo_id == "local/train" + assert config.data.assets.asset_id is None + assert plan.source_repo_id == "local/eval" + assert plan.metric_name == "heldout_offline_imitation_reward" + assert plan.reward_log_key == "eval/heldout_offline_imitation_reward" + assert plan.loss_log_key == "eval/heldout_offline_imitation_loss" + assert plan.uses_heldout_data + + +@pytest.mark.parametrize( + ("eval_repo_id", "message"), + [("", "non-empty string"), ("local/train", "must differ")], +) +def test_heldout_checkpoint_evaluation_rejects_empty_or_training_repo(eval_repo_id: str, message: str) -> None: + config = _Config(data=_Data(repo_id="local/train"), checkpoint_eval_repo_id=eval_repo_id) + + with pytest.raises(ValueError, match=message): + resolve_checkpoint_evaluation(config) diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py index 4ca47e1286..cebf8c8ab7 100644 --- a/src/openpi/training/config.py +++ b/src/openpi/training/config.py @@ -5,6 +5,7 @@ import dataclasses import difflib import logging +import os import pathlib from typing import Any, Literal, Protocol, TypeAlias @@ -19,6 +20,7 @@ import openpi.models.tokenizer as _tokenizer import openpi.policies.aloha_policy as aloha_policy import openpi.policies.droid_policy as droid_policy +import openpi.policies.g1_policy as g1_policy import openpi.policies.libero_policy as libero_policy import openpi.shared.download as _download import openpi.shared.normalize as _normalize @@ -79,6 +81,9 @@ class DataConfig: data_transforms: _transforms.Group = dataclasses.field(default_factory=_transforms.Group) # Model specific transforms. Will be applied after the data is normalized. model_transforms: _transforms.Group = dataclasses.field(default_factory=_transforms.Group) + # Stochastic transforms used only by shuffled training loaders. They are + # excluded from normalization, checkpoint evaluation, and policy inference. + training_transforms: _transforms.Group = dataclasses.field(default_factory=_transforms.Group) # If true, will use quantile normalization. Otherwise, normal z-score normalization will be used. use_quantile_norm: bool = False @@ -462,6 +467,54 @@ def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig ) +@dataclasses.dataclass(frozen=True) +class LeRobotG1DataConfig(DataConfigFactory): + """LeRobot contract for one G1 head camera and a task-level controller.""" + + state_dim: int = 29 + task_action_dim: int = 21 + use_depth_image: bool = False + augment_rgbd: bool = False + + @override + def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig: + repack_mapping = { + "head_image": "observation.images.head", + "state": "observation.state", + "actions": "action", + "prompt": "prompt", + } + if self.use_depth_image: + repack_mapping["depth_image"] = "observation.images.depth" + repack_transform = _transforms.Group( + inputs=[_transforms.RepackTransform(repack_mapping)] + ) + data_transforms = _transforms.Group( + inputs=[ + g1_policy.G1Inputs( + model_type=model_config.model_type, + state_dim=self.state_dim, + task_action_dim=self.task_action_dim, + use_depth_image=self.use_depth_image, + ) + ], + outputs=[g1_policy.G1Outputs(task_action_dim=self.task_action_dim)], + ) + training_transforms = _transforms.Group( + inputs=[g1_policy.G1RgbdAugment()] if self.augment_rgbd else [] + ) + model_transforms = ModelTransformFactory()(model_config) + + return dataclasses.replace( + self.create_base_config(assets_dirs, model_config), + repack_transforms=repack_transform, + data_transforms=data_transforms, + model_transforms=model_transforms, + training_transforms=training_transforms, + action_sequence_keys=("action",), + ) + + @dataclasses.dataclass(frozen=True) class TrainConfig: # Name of the config. Must be unique. Will be used to reference this config. @@ -485,6 +538,16 @@ class TrainConfig: # Precision for PyTorch training. pytorch_training_precision: Literal["bfloat16", "float32"] = "bfloat16" + # Optional dependency-free PyTorch LoRA. When enabled, the pretrained + # vision tower and base weights remain frozen; LoRA residuals are inserted + # into both transformers and the task action heads remain trainable. + pytorch_lora_rank: int | None = None + pytorch_lora_alpha: float = 16.0 + pytorch_lora_action_expert_rank: int | None = None + pytorch_lora_action_expert_alpha: float = 32.0 + pytorch_lora_dropout: float = 0.0 + pytorch_lora_train_action_heads: bool = True + lr_schedule: _optimizer.LRScheduleConfig = dataclasses.field(default_factory=_optimizer.CosineDecaySchedule) optimizer: _optimizer.OptimizerConfig = dataclasses.field(default_factory=_optimizer.AdamW) ema_decay: float | None = 0.99 @@ -509,6 +572,9 @@ class TrainConfig: num_workers: int = 2 # Number of train steps (batches) to run. num_train_steps: int = 30_000 + # Optional wall-clock limit for the actual optimization loop. A final + # checkpoint is written when this deadline is reached. + max_train_seconds: float | None = None # How often (in steps) to log training metrics. log_interval: int = 100 @@ -517,6 +583,19 @@ class TrainConfig: # If set, any existing checkpoints matching step % keep_period == 0 will not be deleted. keep_period: int | None = 5000 + # PyTorch reward-gated checkpointing. The current trainer defines the + # offline reward as exp(-deterministic imitation loss); it is explicitly + # not a simulator task reward. A checkpoint is written only when this + # reward improves by at least checkpoint_reward_min_delta. + checkpoint_reward_batches: int = 16 + checkpoint_reward_min: float = 0.0 + checkpoint_reward_min_delta: float = 1.0e-6 + # Optional separate LeRobot repo used only to rank reward-gated + # checkpoints. It is loaded with the training repo's normalization stats. + checkpoint_eval_repo_id: str | None = None + checkpoint_min_free_disk_gib: float = 256.0 + checkpoint_min_free_disk_fraction: float = 0.10 + # If true, will overwrite the checkpoint directory if it already exists. overwrite: bool = False # If true, will resume training from the last checkpoint. @@ -554,6 +633,24 @@ def trainable_filter(self) -> nnx.filterlib.Filter: def __post_init__(self) -> None: if self.resume and self.overwrite: raise ValueError("Cannot resume and overwrite at the same time.") + if self.max_train_seconds is not None and self.max_train_seconds <= 0: + raise ValueError("max_train_seconds must be positive when set.") + if self.pytorch_lora_rank is not None and self.pytorch_lora_rank <= 0: + raise ValueError("pytorch_lora_rank must be positive when set.") + if self.pytorch_lora_action_expert_rank is not None and self.pytorch_lora_action_expert_rank <= 0: + raise ValueError("pytorch_lora_action_expert_rank must be positive when set.") + if self.pytorch_lora_alpha <= 0 or self.pytorch_lora_action_expert_alpha <= 0: + raise ValueError("PyTorch LoRA alpha values must be positive.") + if not 0.0 <= self.pytorch_lora_dropout < 1.0: + raise ValueError("pytorch_lora_dropout must be in [0, 1).") + if self.checkpoint_reward_batches <= 0: + raise ValueError("checkpoint_reward_batches must be positive.") + if self.checkpoint_reward_min_delta < 0: + raise ValueError("checkpoint_reward_min_delta cannot be negative.") + if self.checkpoint_min_free_disk_gib < 0: + raise ValueError("checkpoint_min_free_disk_gib cannot be negative.") + if not 0.0 <= self.checkpoint_min_free_disk_fraction < 1.0: + raise ValueError("checkpoint_min_free_disk_fraction must be in [0, 1).") # Use `get_config` if you need to get a config by name in your code. @@ -916,6 +1013,193 @@ def __post_init__(self) -> None: num_train_steps=20_000, batch_size=32, ), + TrainConfig( + # Spark-ready full pi0.5 fine-tuning contract for the existing 21-D G1 + # Fruit Ninja task controller. Override data.repo_id with the recorded + # LeRobot dataset before computing stats or starting training. + name="pi05_spark_g1_fruit_ninja", + model=pi0_config.Pi0Config( + pi05=True, + action_dim=32, + action_horizon=10, + discrete_state_input=False, + pytorch_compile_mode=None, + ), + data=LeRobotG1DataConfig( + repo_id=os.getenv("OPENPI_G1_DATASET_REPO_ID", "your_hf_username/g1_fruit_ninja"), + state_dim=29, + task_action_dim=21, + base_config=DataConfig(prompt_from_task=True), + ), + pytorch_weight_path="/openpi_assets/checkpoints/pi05_base_pytorch", + pytorch_training_precision="bfloat16", + batch_size=1, + num_workers=4, + num_train_steps=20_000, + log_interval=10, + save_interval=1_000, + keep_period=5_000, + ema_decay=None, + checkpoint_base_dir="/openpi_assets/training", + assets_base_dir="/openpi_assets/assets", + exp_name="fruit_ninja_pi05", + wandb_enabled=False, + ), + TrainConfig( + # Kinesthetic RGB-D demonstrations recorded on the real G1. The + # observed 7D Dex3 pose remains in state, but is deliberately not used + # as an action label. Only the exact 14 arm targets accepted on + # rt/arm_sdk enter supervised action chunks; hand closure is owned by + # the deterministic grasp stage at runtime. + name="pi05_spark_g1_coke_rgbd_arm14", + model=pi0_config.Pi0Config( + pi05=True, + action_dim=32, + action_horizon=10, + discrete_state_input=False, + pytorch_compile_mode=None, + ), + data=LeRobotG1DataConfig( + repo_id=os.getenv( + "OPENPI_G1_COKE_RGBD_DATASET_REPO_ID", + "local/g1_coke_pickup_real_rgbd_left_rcoke_3_16_train_v1", + ), + state_dim=24, + task_action_dim=14, + use_depth_image=True, + augment_rgbd=True, + base_config=DataConfig(prompt_from_task=True), + ), + pytorch_weight_path="/openpi_assets/checkpoints/pi05_base_pytorch", + pytorch_training_precision="bfloat16", + pytorch_lora_rank=16, + pytorch_lora_alpha=16.0, + pytorch_lora_action_expert_rank=32, + pytorch_lora_action_expert_alpha=32.0, + pytorch_lora_dropout=0.0, + pytorch_lora_train_action_heads=True, + batch_size=4, + num_workers=0, + num_train_steps=4_000, + log_interval=10, + save_interval=200, + keep_period=None, + checkpoint_reward_batches=32, + checkpoint_reward_min=0.0, + checkpoint_reward_min_delta=1.0e-6, + checkpoint_eval_repo_id=os.getenv( + "OPENPI_G1_COKE_RGBD_EVAL_REPO_ID", + "local/g1_coke_pickup_real_rgbd_left_rcoke_3_16_eval_v1", + ), + checkpoint_min_free_disk_gib=256.0, + checkpoint_min_free_disk_fraction=0.10, + ema_decay=None, + checkpoint_base_dir="/openpi_assets/training", + assets_base_dir="/openpi_assets/assets", + exp_name="g1_coke_rgbd_arm14_pi05", + wandb_enabled=False, + ), + TrainConfig( + # Strict physical-left policy. The right arm is absent from both the + # observation and supervised target: RGB-D + waist3 + left-arm7 + + # measured physical-left Dex3 -> seven exact applied left-arm targets. + # The hand remains a deterministic staged controller until recordings + # contain its exact applied commands. + name="pi05_spark_g1_coke_rgbd_left_arm7", + model=pi0_config.Pi0Config( + pi05=True, + action_dim=32, + action_horizon=10, + discrete_state_input=False, + pytorch_compile_mode=None, + ), + data=LeRobotG1DataConfig( + repo_id=os.getenv( + "OPENPI_G1_COKE_LEFT_ONLY_DATASET_REPO_ID", + "local/g1_coke_pickup_real_rgbd_left_only_rcoke_3_16_train_v1", + ), + state_dim=17, + task_action_dim=7, + use_depth_image=True, + augment_rgbd=True, + base_config=DataConfig(prompt_from_task=True), + ), + pytorch_weight_path="/openpi_assets/checkpoints/pi05_base_pytorch", + pytorch_training_precision="bfloat16", + pytorch_lora_rank=16, + pytorch_lora_alpha=16.0, + pytorch_lora_action_expert_rank=32, + pytorch_lora_action_expert_alpha=32.0, + pytorch_lora_dropout=0.0, + pytorch_lora_train_action_heads=True, + batch_size=4, + num_workers=0, + num_train_steps=4_000, + log_interval=10, + save_interval=200, + keep_period=None, + checkpoint_reward_batches=32, + checkpoint_reward_min=0.0, + checkpoint_reward_min_delta=1.0e-6, + checkpoint_eval_repo_id=os.getenv( + "OPENPI_G1_COKE_LEFT_ONLY_EVAL_REPO_ID", + "local/g1_coke_pickup_real_rgbd_left_only_rcoke_3_16_eval_v1", + ), + checkpoint_min_free_disk_gib=256.0, + checkpoint_min_free_disk_fraction=0.10, + ema_decay=None, + checkpoint_base_dir="/openpi_assets/training", + assets_base_dir="/openpi_assets/assets", + exp_name="g1_coke_rgbd_left_arm7_pi05", + wandb_enabled=False, + ), + TrainConfig( + # Vision-language-action policy for the Isaac/real Coke pickup contract: + # one rendered/head RGB view, 24 upper-body positions, and 21 absolute + # arm/right-Dex3 joint targets. Keeping action_dim=32 + # preserves all pretrained pi0.5 checkpoint shapes. + name="pi05_spark_g1_coke_pickup", + model=pi0_config.Pi0Config( + pi05=True, + action_dim=32, + action_horizon=10, + discrete_state_input=False, + pytorch_compile_mode=None, + ), + data=LeRobotG1DataConfig( + repo_id=os.getenv("OPENPI_G1_COKE_DATASET_REPO_ID", "your_hf_username/g1_coke_pickup"), + state_dim=24, + task_action_dim=21, + base_config=DataConfig(prompt_from_task=True), + ), + pytorch_weight_path="/openpi_assets/checkpoints/pi05_base_pytorch", + pytorch_training_precision="bfloat16", + pytorch_lora_rank=16, + pytorch_lora_alpha=16.0, + pytorch_lora_action_expert_rank=32, + pytorch_lora_action_expert_alpha=32.0, + pytorch_lora_dropout=0.0, + pytorch_lora_train_action_heads=True, + batch_size=4, + # This bootstrap dataset contains a single recorded episode. Keep + # loading in-process so OpenCV is not imported concurrently by spawned + # workers during normalization on Spark. + num_workers=0, + num_train_steps=20_000, + log_interval=10, + save_interval=100, + keep_period=None, + checkpoint_reward_batches=8, + checkpoint_reward_min=0.0, + checkpoint_reward_min_delta=1.0e-6, + checkpoint_min_free_disk_gib=256.0, + checkpoint_min_free_disk_fraction=0.10, + ema_decay=None, + checkpoint_base_dir="/openpi_assets/training", + assets_base_dir="/openpi_assets/assets", + exp_name="g1_coke_pickup_pi05", + wandb_enabled=False, + ), # # ALOHA Sim configs. This config is used to demonstrate how to train on a simple simulated environment. # @@ -965,6 +1249,34 @@ def __post_init__(self) -> None: exp_name="debug_pi05", wandb_enabled=False, ), + # NVIDIA DGX Spark / GB10 compatibility smoke test. Unlike debug_pi05, this + # deliberately instantiates and updates the full pi0.5 model. The matching + # NVIDIA PyTorch runtime and converted checkpoint are provided by + # scripts/docker/spark_gb10.Dockerfile and scripts/spark/prepare_pi05.sh. + TrainConfig( + name="pi05_spark_smoke", + model=pi0_config.Pi0Config( + pi05=True, + action_dim=32, + action_horizon=10, + discrete_state_input=False, + pytorch_compile_mode=None, + ), + data=FakeDataConfig(), + pytorch_weight_path="/openpi_assets/checkpoints/pi05_base_pytorch", + pytorch_training_precision="bfloat16", + batch_size=1, + num_workers=0, + num_train_steps=1, + log_interval=1, + save_interval=1, + keep_period=None, + ema_decay=None, + checkpoint_base_dir="/openpi_assets/training", + overwrite=True, + exp_name="gb10_full_model_smoke", + wandb_enabled=False, + ), # RoboArena & PolaRiS configs. *roboarena_config.get_roboarena_configs(), *polaris_config.get_polaris_configs(), diff --git a/src/openpi/training/data_loader.py b/src/openpi/training/data_loader.py index e2ee7dd06b..568052eb93 100644 --- a/src/openpi/training/data_loader.py +++ b/src/openpi/training/data_loader.py @@ -169,7 +169,13 @@ def create_rlds_dataset( ) -def transform_dataset(dataset: Dataset, data_config: _config.DataConfig, *, skip_norm_stats: bool = False) -> Dataset: +def transform_dataset( + dataset: Dataset, + data_config: _config.DataConfig, + *, + skip_norm_stats: bool = False, + training: bool = False, +) -> Dataset: """Transform the dataset by applying the data transforms.""" norm_stats = {} if data_config.repo_id != "fake" and not skip_norm_stats: @@ -185,6 +191,7 @@ def transform_dataset(dataset: Dataset, data_config: _config.DataConfig, *, skip [ *data_config.repack_transforms.inputs, *data_config.data_transforms.inputs, + *(data_config.training_transforms.inputs if training else ()), _transforms.Normalize(norm_stats, use_quantiles=data_config.use_quantile_norm), *data_config.model_transforms.inputs, ], @@ -300,7 +307,7 @@ def create_torch_data_loader( seed: The seed to use for shuffling the data. """ dataset = create_torch_dataset(data_config, action_horizon, model_config) - dataset = transform_dataset(dataset, data_config, skip_norm_stats=skip_norm_stats) + dataset = transform_dataset(dataset, data_config, skip_norm_stats=skip_norm_stats, training=shuffle) # Use TorchDataLoader for both frameworks # For PyTorch DDP, create DistributedSampler and divide batch size by world size @@ -538,3 +545,20 @@ def data_config(self) -> _config.DataConfig: def __iter__(self): for batch in self._data_loader: yield _model.Observation.from_dict(batch), batch["actions"] + + def evenly_spaced_batches(self, *, num_batches: int, batch_size: int): + """Materialize deterministic random-access batches spanning the dataset.""" + if not isinstance(self._data_loader, TorchDataLoader): + raise TypeError("Evenly spaced evaluation batches require a TorchDataLoader") + dataset = self._data_loader.torch_loader.dataset + sample_count = min(len(dataset), num_batches * batch_size) + indices = np.linspace(0, len(dataset) - 1, sample_count, dtype=np.int64) + batches = [] + for start in range(0, sample_count, batch_size): + batch_indices = indices[start : start + batch_size] + if len(batch_indices) < batch_size: + break + batch = _collate_fn([dataset[int(index)] for index in batch_indices]) + batch = jax.tree.map(torch.as_tensor, batch) + batches.append((_model.Observation.from_dict(batch), batch["actions"])) + return batches diff --git a/src/openpi/training/data_loader_test.py b/src/openpi/training/data_loader_test.py index d15a73529e..82a39bba76 100644 --- a/src/openpi/training/data_loader_test.py +++ b/src/openpi/training/data_loader_test.py @@ -2,11 +2,26 @@ import jax +from openpi import transforms as _transforms from openpi.models import pi0_config from openpi.training import config as _config from openpi.training import data_loader as _data_loader +class _OneItemDataset: + def __getitem__(self, _index): + return {"marker": 0} + + def __len__(self): + return 1 + + +class _MarkTraining: + def __call__(self, data): + data["marker"] += 1 + return data + + def test_torch_data_loader(): config = pi0_config.Pi0Config(action_dim=24, action_horizon=50, max_token_len=48) dataset = _data_loader.FakeDataset(config, 16) @@ -23,6 +38,23 @@ def test_torch_data_loader(): assert all(x.shape[0] == 4 for x in jax.tree.leaves(batch)) +def test_training_transforms_are_excluded_from_evaluation() -> None: + data_config = _config.DataConfig( + repo_id="fake", + training_transforms=_transforms.Group(inputs=[_MarkTraining()]), + ) + + train = _data_loader.transform_dataset( + _OneItemDataset(), data_config, skip_norm_stats=True, training=True + ) + evaluation = _data_loader.transform_dataset( + _OneItemDataset(), data_config, skip_norm_stats=True, training=False + ) + + assert train[0]["marker"] == 1 + assert evaluation[0]["marker"] == 0 + + def test_torch_data_loader_infinite(): config = pi0_config.Pi0Config(action_dim=24, action_horizon=50, max_token_len=48) dataset = _data_loader.FakeDataset(config, 4) diff --git a/src/openpi/training/reward_checkpoints.py b/src/openpi/training/reward_checkpoints.py new file mode 100644 index 0000000000..7fa1b1012b --- /dev/null +++ b/src/openpi/training/reward_checkpoints.py @@ -0,0 +1,173 @@ +"""Atomic, reward-gated checkpoints with disk-pressure eviction.""" + +from __future__ import annotations + +from collections.abc import Callable +import dataclasses +import json +import logging +from pathlib import Path +import shutil +import time + + +@dataclasses.dataclass(frozen=True) +class SaveResult: + saved: bool + checkpoint_dir: Path | None + reward: float + best_reward: float | None + reason: str + evicted_steps: tuple[int, ...] = () + + +class RewardCheckpointStore: + """Keep only strict reward improvements and evict oldest under pressure.""" + + _SCHEMA_VERSION = 1 + + def __init__( + self, + root: Path, + *, + metric_name: str, + minimum_reward: float, + minimum_delta: float, + minimum_free_bytes: int, + minimum_free_fraction: float, + disk_usage: Callable[[Path], tuple[int, int, int]] = shutil.disk_usage, + ) -> None: + self.root = root + self.metric_name = metric_name + self.minimum_reward = minimum_reward + self.minimum_delta = minimum_delta + self.minimum_free_bytes = minimum_free_bytes + self.minimum_free_fraction = minimum_free_fraction + self._disk_usage = disk_usage + self.root.mkdir(parents=True, exist_ok=True) + + @property + def manifest_path(self) -> Path: + return self.root / "reward_checkpoints.json" + + def _load_manifest(self) -> dict: + if not self.manifest_path.exists(): + return { + "schema_version": self._SCHEMA_VERSION, + "metric_name": self.metric_name, + "records": [], + } + manifest = json.loads(self.manifest_path.read_text(encoding="utf-8")) + if manifest.get("schema_version") != self._SCHEMA_VERSION: + raise ValueError(f"Unsupported reward checkpoint manifest: {self.manifest_path}") + if manifest.get("metric_name") != self.metric_name: + raise ValueError( + f"Checkpoint metric changed from {manifest.get('metric_name')!r} to {self.metric_name!r}" + ) + manifest["records"] = [ + record for record in manifest.get("records", []) if (self.root / str(record["step"])).is_dir() + ] + return manifest + + def _write_manifest(self, manifest: dict) -> None: + temporary = self.manifest_path.with_suffix(".tmp") + temporary.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8") + temporary.replace(self.manifest_path) + + def best_record(self) -> dict | None: + records = self._load_manifest()["records"] + return max(records, key=lambda record: float(record["reward"])) if records else None + + def best_checkpoint_dir(self) -> Path | None: + best = self.best_record() + return self.root / str(best["step"]) if best is not None else None + + def _required_free_bytes(self) -> int: + usage = self._disk_usage(self.root) + return max(self.minimum_free_bytes, int(usage.total * self.minimum_free_fraction)) + + def _evict_under_pressure(self, manifest: dict, *, protected_step: int | None) -> tuple[int, ...]: + evicted: list[int] = [] + required = self._required_free_bytes() + while self._disk_usage(self.root).free < required: + candidates = sorted( + (record for record in manifest["records"] if int(record["step"]) != protected_step), + key=lambda record: (float(record["created_at"]), int(record["step"])), + ) + if not candidates: + raise RuntimeError( + f"Checkpoint filesystem has less than {required / 1024**3:.1f} GiB free and no older " + "reward checkpoint can be evicted safely" + ) + oldest = candidates[0] + step = int(oldest["step"]) + target = self.root / str(step) + if target.is_dir(): + shutil.rmtree(target) + manifest["records"].remove(oldest) + evicted.append(step) + logging.warning("Evicted oldest reward checkpoint at step %d because disk headroom is low", step) + return tuple(evicted) + + def maybe_save( + self, + *, + step: int, + reward: float, + write_checkpoint: Callable[[Path], None], + ) -> SaveResult: + manifest = self._load_manifest() + best = max((float(record["reward"]) for record in manifest["records"]), default=None) + if reward < self.minimum_reward: + return SaveResult( + saved=False, + checkpoint_dir=None, + reward=reward, + best_reward=best, + reason="below_minimum_reward", + ) + if best is not None and reward <= best + self.minimum_delta: + return SaveResult( + saved=False, + checkpoint_dir=None, + reward=reward, + best_reward=best, + reason="not_a_strict_improvement", + ) + + evicted = list(self._evict_under_pressure(manifest, protected_step=None)) + self._write_manifest(manifest) + + final_dir = self.root / str(step) + temporary_dir = self.root / f"tmp_{step}" + if temporary_dir.exists(): + shutil.rmtree(temporary_dir) + temporary_dir.mkdir(parents=True) + try: + write_checkpoint(temporary_dir) + if final_dir.exists(): + shutil.rmtree(final_dir) + temporary_dir.replace(final_dir) + except BaseException: + if temporary_dir.exists(): + shutil.rmtree(temporary_dir) + raise + + record = { + "step": step, + "reward": reward, + "created_at": time.time(), + "metric_name": self.metric_name, + } + manifest["records"].append(record) + evicted.extend(self._evict_under_pressure(manifest, protected_step=step)) + self._write_manifest(manifest) + logging.info("Saved new best reward checkpoint: step=%d reward=%.8f", step, reward) + return SaveResult( + saved=True, + checkpoint_dir=final_dir, + reward=reward, + best_reward=reward, + reason="strict_improvement", + evicted_steps=tuple(evicted), + ) diff --git a/src/openpi/training/reward_checkpoints_test.py b/src/openpi/training/reward_checkpoints_test.py new file mode 100644 index 0000000000..7f5071a590 --- /dev/null +++ b/src/openpi/training/reward_checkpoints_test.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from collections import namedtuple +import json +from pathlib import Path + +from openpi.training.reward_checkpoints import RewardCheckpointStore + + +def _writer(path: Path) -> None: + (path / "payload.bin").write_bytes(b"checkpoint") + + +def test_only_strict_reward_improvements_are_saved(tmp_path: Path) -> None: + store = RewardCheckpointStore( + tmp_path, + metric_name="reward", + minimum_reward=0.0, + minimum_delta=1.0e-6, + minimum_free_bytes=0, + minimum_free_fraction=0.0, + ) + + assert store.maybe_save(step=10, reward=0.2, write_checkpoint=_writer).saved + assert not store.maybe_save(step=20, reward=0.1, write_checkpoint=_writer).saved + assert not store.maybe_save(step=30, reward=0.2000005, write_checkpoint=_writer).saved + assert store.maybe_save(step=40, reward=0.3, write_checkpoint=_writer).saved + + assert (tmp_path / "10").is_dir() + assert not (tmp_path / "20").exists() + assert not (tmp_path / "30").exists() + assert (tmp_path / "40").is_dir() + assert store.best_checkpoint_dir() == tmp_path / "40" + + +def test_disk_pressure_evicts_oldest_reward_checkpoint(tmp_path: Path) -> None: + DiskUsage = namedtuple("DiskUsage", "total used free") + + def disk_usage(path: Path): + checkpoint_count = sum(child.is_dir() and child.name.isdigit() for child in path.iterdir()) + free = 0 if checkpoint_count > 1 else 100 + return DiskUsage(total=1000, used=1000 - free, free=free) + + store = RewardCheckpointStore( + tmp_path, + metric_name="reward", + minimum_reward=0.0, + minimum_delta=0.0, + minimum_free_bytes=50, + minimum_free_fraction=0.0, + disk_usage=disk_usage, + ) + + assert store.maybe_save(step=10, reward=0.2, write_checkpoint=_writer).saved + second = store.maybe_save(step=20, reward=0.3, write_checkpoint=_writer) + + assert second.saved + assert second.evicted_steps == (10,) + assert not (tmp_path / "10").exists() + assert (tmp_path / "20").is_dir() + manifest = json.loads((tmp_path / "reward_checkpoints.json").read_text()) + assert [record["step"] for record in manifest["records"]] == [20]